blockchain sandbox: UTXO chain, p2p mining, attacker node

attacker/node.go

@@ -0,0 +1,206 @@
+// 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})
+		}
+	}
+}