dchain/consensus/pbft_test.go

package consensus_test

import (
	"encoding/json"
	"sync"
	"testing"
	"time"

	"go-blockchain/blockchain"
	"go-blockchain/consensus"
	"go-blockchain/identity"
)

// ─── helpers ─────────────────────────────────────────────────────────────────

func newID(t *testing.T) *identity.Identity {
	t.Helper()
	id, err := identity.Generate()
	if err != nil {
		t.Fatalf("identity.Generate: %v", err)
	}
	return id
}

func genesisFor(id *identity.Identity) *blockchain.Block {
	return blockchain.GenesisBlock(id.PubKeyHex(), id.PrivKey)
}

// network is a simple in-process message bus between engines.
// Each engine has a dedicated goroutine that delivers messages in FIFO order,
// which avoids the race where a PREPARE arrives before the PRE-PREPARE it
// depends on (which would cause the vote to be silently discarded).
type network struct {
	mu      sync.Mutex
	queues  []chan *blockchain.ConsensusMsg
	engines []*consensus.Engine
}

// addEngine registers an engine and starts its delivery goroutine.
func (n *network) addEngine(e *consensus.Engine) {
	ch := make(chan *blockchain.ConsensusMsg, 256)
	n.mu.Lock()
	n.engines = append(n.engines, e)
	n.queues = append(n.queues, ch)
	n.mu.Unlock()
	go func() {
		for msg := range ch {
			e.HandleMessage(msg)
		}
	}()
}

func (n *network) broadcast(msg *blockchain.ConsensusMsg) {
	n.mu.Lock()
	queues := make([]chan *blockchain.ConsensusMsg, len(n.queues))
	copy(queues, n.queues)
	n.mu.Unlock()
	for _, q := range queues {
		cp := *msg // copy to avoid concurrent signature-nil race in verifyMsgSig
		q <- &cp
	}
}

// committedBlocks collects blocks committed by an engine into a channel.
type committedBlocks struct {
	ch chan *blockchain.Block
}

func (cb *committedBlocks) onCommit(b *blockchain.Block) {
	cb.ch <- b
}

func newCommitted() *committedBlocks {
	return &committedBlocks{ch: make(chan *blockchain.Block, 16)}
}

func (cb *committedBlocks) waitOne(t *testing.T, timeout time.Duration) *blockchain.Block {
	t.Helper()
	select {
	case b := <-cb.ch:
		return b
	case <-time.After(timeout):
		t.Fatal("timed out waiting for committed block")
		return nil
	}
}

// ─── tests ───────────────────────────────────────────────────────────────────

// TestSingleValidatorCommit verifies that a single-validator network commits
// blocks immediately (f=0, quorum=1).
func TestSingleValidatorCommit(t *testing.T) {
	id := newID(t)
	genesis := genesisFor(id)

	committed := newCommitted()
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1, // seqNum = genesis+1
		committed.onCommit,
		net.broadcast,
	)
	net.addEngine(engine)

	// Propose block 1 from genesis.
	engine.Propose(genesis)

	b := committed.waitOne(t, 3*time.Second)
	if b.Index != 1 {
		t.Errorf("expected committed block index 1, got %d", b.Index)
	}
	if b.Validator != id.PubKeyHex() {
		t.Errorf("wrong validator in committed block")
	}
}

// TestSingleValidatorMultipleBlocks verifies sequential block commitment.
func TestSingleValidatorMultipleBlocks(t *testing.T) {
	id := newID(t)
	genesis := genesisFor(id)

	committed := newCommitted()
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		committed.onCommit,
		net.broadcast,
	)
	net.addEngine(engine)

	prev := genesis
	for i := uint64(1); i <= 3; i++ {
		engine.Propose(prev)
		b := committed.waitOne(t, 3*time.Second)
		if b.Index != i {
			t.Errorf("block %d: expected index %d, got %d", i, i, b.Index)
		}
		engine.SyncSeqNum(i + 1)
		prev = b
	}
}

// TestThreeValidatorCommit verifies a 3-node network reaches consensus.
// Messages are delivered synchronously through the in-process network bus.
// With f=1, quorum = ⌈2*3/3⌉ = 2.
func TestThreeValidatorCommit(t *testing.T) {
	ids := []*identity.Identity{newID(t), newID(t), newID(t)}
	valSet := []string{ids[0].PubKeyHex(), ids[1].PubKeyHex(), ids[2].PubKeyHex()}

	genesis := genesisFor(ids[0]) // block 0 signed by ids[0]

	committed := [3]*committedBlocks{newCommitted(), newCommitted(), newCommitted()}
	net := &network{}

	for i, id := range ids {
		idx := i
		engine := consensus.NewEngine(
			id, valSet, 1,
			func(b *blockchain.Block) { committed[idx].onCommit(b) },
			net.broadcast,
		)
		net.addEngine(engine)
	}

	// Leader for seqNum=1, view=0 is valSet[(1+0)%3] = ids[1].
	// Find and trigger the leader.
	for i, e := range net.engines {
		_ = i
		e.Propose(genesis)
	}

	// All three should commit the same block.
	timeout := 5 * time.Second
	var commitIdx [3]uint64
	for i := 0; i < 3; i++ {
		b := committed[i].waitOne(t, timeout)
		commitIdx[i] = b.Index
	}
	for i, idx := range commitIdx {
		if idx != 1 {
			t.Errorf("engine %d committed block at wrong index: got %d, want 1", i, idx)
		}
	}
}

// TestAddTransactionAndPropose verifies that pending transactions appear in committed blocks.
func TestAddTransactionAndPropose(t *testing.T) {
	id := newID(t)
	sender := newID(t)
	genesis := genesisFor(id)

	committed := newCommitted()
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		committed.onCommit,
		net.broadcast,
	)
	net.addEngine(engine)

	// Build a valid signed transaction.
	payload, _ := json.Marshal(blockchain.TransferPayload{})
	tx := &blockchain.Transaction{
		ID:        "test-tx-1",
		Type:      blockchain.EventTransfer,
		From:      sender.PubKeyHex(),
		To:        id.PubKeyHex(),
		Amount:    1000,
		Fee:       blockchain.MinFee,
		Payload:   payload,
		Timestamp: time.Now().UTC(),
	}
	// Sign with canonical bytes (matching validateTx).
	signData, _ := json.Marshal(struct {
		ID        string               `json:"id"`
		Type      blockchain.EventType `json:"type"`
		From      string               `json:"from"`
		To        string               `json:"to"`
		Amount    uint64               `json:"amount"`
		Fee       uint64               `json:"fee"`
		Payload   []byte               `json:"payload"`
		Timestamp time.Time            `json:"timestamp"`
	}{tx.ID, tx.Type, tx.From, tx.To, tx.Amount, tx.Fee, tx.Payload, tx.Timestamp})
	tx.Signature = sender.Sign(signData)

	if err := engine.AddTransaction(tx); err != nil {
		t.Fatalf("AddTransaction: %v", err)
	}

	engine.Propose(genesis)

	b := committed.waitOne(t, 3*time.Second)
	if len(b.Transactions) != 1 {
		t.Errorf("expected 1 transaction in committed block, got %d", len(b.Transactions))
	}
	if b.Transactions[0].ID != "test-tx-1" {
		t.Errorf("wrong transaction in committed block: %s", b.Transactions[0].ID)
	}
}

// TestDuplicateTransactionRejected verifies the mempool deduplicates by TX ID.
func TestDuplicateTransactionRejected(t *testing.T) {
	id := newID(t)
	sender := newID(t)
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		func(*blockchain.Block) {},
		net.broadcast,
	)

	payload, _ := json.Marshal(blockchain.TransferPayload{})
	tx := &blockchain.Transaction{
		ID:        "dup-tx",
		Type:      blockchain.EventTransfer,
		From:      sender.PubKeyHex(),
		To:        id.PubKeyHex(),
		Amount:    1000,
		Fee:       blockchain.MinFee,
		Payload:   payload,
		Timestamp: time.Now().UTC(),
	}
	signData, _ := json.Marshal(struct {
		ID        string               `json:"id"`
		Type      blockchain.EventType `json:"type"`
		From      string               `json:"from"`
		To        string               `json:"to"`
		Amount    uint64               `json:"amount"`
		Fee       uint64               `json:"fee"`
		Payload   []byte               `json:"payload"`
		Timestamp time.Time            `json:"timestamp"`
	}{tx.ID, tx.Type, tx.From, tx.To, tx.Amount, tx.Fee, tx.Payload, tx.Timestamp})
	tx.Signature = sender.Sign(signData)

	if err := engine.AddTransaction(tx); err != nil {
		t.Fatalf("first AddTransaction: %v", err)
	}
	if err := engine.AddTransaction(tx); err == nil {
		t.Fatal("expected duplicate transaction to be rejected, but it was accepted")
	}
}

// TestInvalidTxRejected verifies that transactions with bad signatures are rejected.
func TestInvalidTxRejected(t *testing.T) {
	id := newID(t)
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		func(*blockchain.Block) {},
		net.broadcast,
	)

	payload, _ := json.Marshal(blockchain.TransferPayload{})
	tx := &blockchain.Transaction{
		ID:        "bad-sig-tx",
		Type:      blockchain.EventTransfer,
		From:      id.PubKeyHex(),
		To:        id.PubKeyHex(),
		Amount:    1000,
		Fee:       blockchain.MinFee,
		Payload:   payload,
		Timestamp: time.Now().UTC(),
		Signature: []byte("not-a-real-signature"),
	}
	if err := engine.AddTransaction(tx); err == nil {
		t.Fatal("expected transaction with bad signature to be rejected")
	}
}

// TestFeeBelowMinimumRejected verifies that sub-minimum fees are rejected by the engine.
func TestFeeBelowMinimumRejected(t *testing.T) {
	id := newID(t)
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		func(*blockchain.Block) {},
		net.broadcast,
	)

	payload, _ := json.Marshal(blockchain.TransferPayload{})
	tx := &blockchain.Transaction{
		ID:        "low-fee-tx",
		Type:      blockchain.EventTransfer,
		From:      id.PubKeyHex(),
		To:        id.PubKeyHex(),
		Amount:    1000,
		Fee:       blockchain.MinFee - 1, // one µT below minimum
		Payload:   payload,
		Timestamp: time.Now().UTC(),
	}
	signData, _ := json.Marshal(struct {
		ID        string               `json:"id"`
		Type      blockchain.EventType `json:"type"`
		From      string               `json:"from"`
		To        string               `json:"to"`
		Amount    uint64               `json:"amount"`
		Fee       uint64               `json:"fee"`
		Payload   []byte               `json:"payload"`
		Timestamp time.Time            `json:"timestamp"`
	}{tx.ID, tx.Type, tx.From, tx.To, tx.Amount, tx.Fee, tx.Payload, tx.Timestamp})
	tx.Signature = id.Sign(signData)

	if err := engine.AddTransaction(tx); err == nil {
		t.Fatal("expected transaction with fee below MinFee to be rejected")
	}
}

// TestUpdateValidators verifies that UpdateValidators takes effect on the next round.
// We start with a 1-validator network (quorum=1), commit one block, then shrink
// the set back to the same single validator — confirming the hot-reload path runs
// without panicking and that the engine continues to commit blocks normally.
func TestUpdateValidators(t *testing.T) {
	id := newID(t)
	genesis := genesisFor(id)

	committed := newCommitted()
	net := &network{}

	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		committed.onCommit,
		net.broadcast,
	)
	net.engines = []*consensus.Engine{engine}

	// Block 1.
	engine.Propose(genesis)
	b1 := committed.waitOne(t, 3*time.Second)
	engine.SyncSeqNum(2)

	// Hot-reload: same single validator — ensures the method is exercised.
	engine.UpdateValidators([]string{id.PubKeyHex()})

	// Block 2 should still commit with the reloaded set.
	engine.Propose(b1)
	b2 := committed.waitOne(t, 3*time.Second)
	if b2.Index != 2 {
		t.Errorf("expected block index 2 after validator update, got %d", b2.Index)
	}
}

// TestSyncSeqNum verifies that SyncSeqNum advances the engine's expected block index.
func TestSyncSeqNum(t *testing.T) {
	id := newID(t)
	net := &network{}

	committed := newCommitted()
	engine := consensus.NewEngine(
		id,
		[]string{id.PubKeyHex()},
		1,
		committed.onCommit,
		net.broadcast,
	)
	net.engines = []*consensus.Engine{engine}

	// Simulate receiving a chain sync that jumps to block 5.
	engine.SyncSeqNum(5)

	genesis := genesisFor(id)
	// Build a fake block at index 5 to propose from.
	b5 := &blockchain.Block{
		Index:        4,
		Timestamp:    time.Now().UTC(),
		Transactions: []*blockchain.Transaction{},
		PrevHash:     genesis.Hash,
		Validator:    id.PubKeyHex(),
		TotalFees:    0,
	}
	b5.ComputeHash()
	b5.Sign(id.PrivKey)

	engine.Propose(b5)
	b := committed.waitOne(t, 3*time.Second)
	if b.Index != 5 {
		t.Errorf("expected committed block index 5, got %d", b.Index)
	}
}