blockchain/attacker/node.go

// Package attacker is a fork of the honest node.
// Modify this code to try to break the blockchain!
//
// You have full access to YOUR node — change any logic you want.
// But other nodes run the original honest code (p2p.Node).
// You can only communicate with them through messages (Broadcast).
//
// Can you steal coins? Double-spend? Rewrite history?
package attacker

import (
	"blockchain/chain"
	"blockchain/config"
	"blockchain/p2p"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"math/rand"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Node struct {
	Wallet    *chain.Wallet
	Chain     *chain.Blockchain
	Mempool   *chain.Mempool
	HashDelay time.Duration
	conn      chan p2p.Msg
	syncCh    chan p2p.Msg
	peers     map[p2p.Peer]*link
	mx        sync.Mutex
	seen      map[string]struct{}
}

type link struct {
	send    chan p2p.Msg
	latency time.Duration
}

func newlink(target p2p.Peer, latency time.Duration) *link {
	l := &link{
		send:    make(chan p2p.Msg, 50),
		latency: latency,
	}
	go func() {
		for msg := range l.send {
			jitter := time.Duration(rand.NormFloat64()*float64(l.latency)/4 + float64(l.latency))
			if jitter > 0 {
				time.Sleep(jitter)
			}
			target.Receive(msg)
		}
	}()
	return l
}

func NewNode() *Node {
	return &Node{
		Wallet:  chain.NewWallet(),
		Chain:   chain.NewBlockchain(),
		Mempool: chain.NewMempool(),
		conn:    make(chan p2p.Msg, 100),
		syncCh:  make(chan p2p.Msg, 10),
		peers:   make(map[p2p.Peer]*link),
		seen:    make(map[string]struct{}),
	}
}

// Receive implements p2p.Peer.
func (n *Node) Receive(msg p2p.Msg) {
	select {
	case n.conn <- msg:
	default:
	}
}

// ReceiveSync implements p2p.Peer.
func (n *Node) ReceiveSync(msg p2p.Msg) {
	select {
	case n.syncCh <- msg:
	default:
	}
}

func (n *Node) ShortID() string {
	return n.Wallet.ShortAddress()
}

func (n *Node) MarkSeen(hash string) bool {
	n.mx.Lock()
	defer n.mx.Unlock()
	if _, ok := n.seen[hash]; ok {
		return true
	}
	if len(n.seen) >= 10000 {
		n.seen = make(map[string]struct{})
	}
	n.seen[hash] = struct{}{}
	return false
}

func (n *Node) AddPeer(peer p2p.Peer) {
	if peer == n {
		return
	}
	n.mx.Lock()
	defer n.mx.Unlock()
	if _, exists := n.peers[peer]; exists {
		return
	}
	n.peers[peer] = newlink(peer, config.Default.P2P.NetDelay)
}

func (n *Node) Broadcast(msg p2p.Msg) {
	n.mx.Lock()
	links := make([]*link, 0, len(n.peers))
	for _, l := range n.peers {
		links = append(links, l)
	}
	n.mx.Unlock()

	for _, l := range links {
		select {
		case l.send <- msg:
		default:
		}
	}
}

func (n *Node) Mine(block *chain.Block) bool {
	prefix := strings.Repeat("0", config.Default.Chain.Difficulty)
	base := block.HashBase()
	for {
		data := base + strconv.Itoa(block.Nonce)
		h := sha256.Sum256([]byte(data))
		block.Hash = hex.EncodeToString(h[:])
		if strings.HasPrefix(block.Hash, prefix) {
			return true
		}
		block.Nonce++
		if n.HashDelay > 0 {
			time.Sleep(n.HashDelay)
		}
	}
}

// Run is the main loop — MODIFY THIS to implement your attack!
// Right now it behaves like an honest node.
func (n *Node) Run(stop chan struct{}) {
	// Start mining like an honest node
	go n.runMiner(stop)

	for {
		select {
		case msg := <-n.conn:
			switch m := msg.(type) {
			case p2p.BlockMsg:
				if n.MarkSeen(m.Block.Hash) {
					continue
				}
				if n.Chain.AddBlock(m.Block) {
					n.Mempool.RemoveMined(m.Block)
					fmt.Printf("[ATTACKER %s] Accepted block #%d\n", n.ShortID(), m.Block.Height)
					n.Broadcast(p2p.BlockMsg{Block: m.Block, From: n})
				}
			case p2p.TxMsg:
				if n.MarkSeen(m.Tx.TxID()) {
					continue
				}
				n.Mempool.Add(m.Tx)
				n.Broadcast(p2p.TxMsg{Tx: m.Tx, From: n})
			}
		case msg := <-n.syncCh:
			switch m := msg.(type) {
			case p2p.BlocksMsg:
				n.Chain.ReplaceBlocks(m.Blocks)
			}
		case <-stop:
			return
		}
	}
}

func (n *Node) runMiner(stop chan struct{}) {
	for {
		select {
		case <-stop:
			return
		default:
		}

		last := n.Chain.LastBlock()
		txs := n.Mempool.CollectBlock(n.Wallet.PublicKey)
		block := chain.NewBlock(last.Height+1, txs, last.Hash)
		n.Mine(block)

		if n.Chain.AddBlock(block) {
			n.Mempool.RemoveMined(block)
			fmt.Printf("[ATTACKER %s] Mined block #%d, hash: %s\n",
				n.ShortID(), block.Height, block.Hash[:16])
			n.Broadcast(p2p.BlockMsg{Block: block, From: n})
		}
	}
}