Implement P2P GetBlockBodies Handler (#1359)

adapters/core2p2p/class.go

@@ -0,0 +1,30 @@
+package core2p2p
+
+import (
+	"fmt"
+
+	"github.com/NethermindEth/juno/core"
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/p2p/starknet/spec"
+)
+
+func AdaptClass(class core.Class, compiledHash *felt.Felt) *spec.Class {
+	if class == nil {
+		return nil
+	}
+
+	switch v := class.(type) {
+	case *core.Cairo0Class:
+		return &spec.Class{
+			CompiledHash: AdaptHash(compiledHash),
+			Definition:   []byte(v.Program),
+		}
+	case *core.Cairo1Class:
+		return &spec.Class{
+			CompiledHash: AdaptHash(compiledHash),
+			Definition:   v.Compiled,
+		}
+	default:
+		panic(fmt.Errorf("unsupported cairo class %T (version=%d)", v, class.Version()))
+	}
+}

adapters/core2p2p/state.go

@@ -0,0 +1,26 @@
+package core2p2p
+
+import (
+	"github.com/NethermindEth/juno/core"
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/p2p/starknet/spec"
+	"github.com/NethermindEth/juno/utils"
+)
+
+func AdaptStateDiff(addr, classHash, nonce *felt.Felt, diff []core.StorageDiff) *spec.StateDiff_ContractDiff {
+	return &spec.StateDiff_ContractDiff{
+		Address:   AdaptAddress(addr),
+		Nonce:     AdaptFelt(nonce),
+		ClassHash: AdaptFelt(classHash),
+		Values:    AdaptStorageDiff(diff),
+	}
+}
+
+func AdaptStorageDiff(diff []core.StorageDiff) []*spec.ContractStoredValue {
+	return utils.Map(diff, func(item core.StorageDiff) *spec.ContractStoredValue {
+		return &spec.ContractStoredValue{
+			Key:   AdaptFelt(item.Key),
+			Value: AdaptFelt(item.Value),
+		}
+	})
+}

adapters/p2p2core/felt.go

@@ -12,3 +12,11 @@ func AdaptHash(h *spec.Hash) *felt.Felt {
 
 	return new(felt.Felt).SetBytes(h.Elements)
 }
+
+func AdaptAddress(h *spec.Address) *felt.Felt {
+	if h == nil {
+		return nil
+	}
+
+	return new(felt.Felt).SetBytes(h.Elements)
+}

p2p/starknet/block_bodies.go

@@ -0,0 +1,247 @@
+package starknet
+
+import (
+	"crypto/rand"
+	"slices"
+
+	"github.com/NethermindEth/juno/adapters/core2p2p"
+	"github.com/NethermindEth/juno/blockchain"
+	"github.com/NethermindEth/juno/core"
+	"github.com/NethermindEth/juno/core/felt"
+	"github.com/NethermindEth/juno/p2p/starknet/spec"
+	"github.com/NethermindEth/juno/utils"
+	"google.golang.org/protobuf/proto"
+)
+
+type blockBodyStep int
+
+const (
+	_ blockBodyStep = iota
+
+	sendDiff // initial
+	sendClasses
+	sendProof
+	sendBlockFin
+	terminal // final step
+)
+
+type blockBodyIterator struct {
+	log         utils.Logger
+	stateReader core.StateReader
+	stateCloser func() error
+
+	step        blockBodyStep
+	header      *core.Header
+	stateUpdate *core.StateUpdate
+}
+
+func newBlockBodyIterator(bcReader blockchain.Reader, header *core.Header, log utils.Logger) (*blockBodyIterator, error) {
+	stateUpdate, err := bcReader.StateUpdateByNumber(header.Number)
+	if err != nil {
+		return nil, err
+	}
+
+	stateReader, closer, err := bcReader.StateAtBlockNumber(header.Number)
+	if err != nil {
+		return nil, err
+	}
+
+	return &blockBodyIterator{
+		step:        sendDiff,
+		header:      header,
+		log:         log,
+		stateReader: stateReader,
+		stateCloser: closer,
+		stateUpdate: stateUpdate,
+	}, nil
+}
+
+func (b *blockBodyIterator) hasNext() bool {
+	return slices.Contains([]blockBodyStep{
+		sendDiff,
+		sendClasses,
+		sendProof,
+		sendBlockFin,
+	}, b.step)
+}
+
+// Either BlockBodiesResponse_Diff, *_Classes, *_Proof, *_Fin
+func (b *blockBodyIterator) next() (msg proto.Message, valid bool) {
+	switch b.step {
+	case sendDiff:
+		msg, valid = b.diff()
+		b.step = sendClasses
+	case sendClasses:
+		msg, valid = b.classes()
+		b.step = sendProof
+	case sendProof:
+		msg, valid = b.proof()
+		b.step = sendBlockFin
+	case sendBlockFin:
+		// fin changes step to terminal internally
+		msg, valid = b.fin()
+	case terminal:
+		panic("next called on terminal step")
+	default:
+		b.log.Errorw("Unknown step in blockBodyIterator", "step", b.step)
+	}
+
+	return
+}
+
+func (b *blockBodyIterator) classes() (proto.Message, bool) {
+	classesM := make(map[felt.Felt]*spec.Class)
+
+	stateDiff := b.stateUpdate.StateDiff
+
+	for _, hash := range stateDiff.DeclaredV0Classes {
+		cls, err := b.stateReader.Class(hash)
+		if err != nil {
+			return b.fin()
+		}
+
+		classesM[*hash] = core2p2p.AdaptClass(cls.Class, hash)
+	}
+	for _, class := range stateDiff.DeclaredV1Classes {
+		cls, err := b.stateReader.Class(class.ClassHash)
+		if err != nil {
+			return b.fin()
+		}
+
+		classesM[*class.ClassHash] = core2p2p.AdaptClass(cls.Class, cls.Class.(*core.Cairo1Class).Hash())
+	}
+	for _, class := range stateDiff.DeployedContracts {
+		if _, ok := classesM[*class.ClassHash]; ok {
+			// Skip if the class was declared and deployed in the same block. Otherwise, there will be duplicate classes.
+			continue
+		}
+		cls, err := b.stateReader.Class(class.ClassHash)
+		if err != nil {
+			return b.fin()
+		}
+
+		var compiledHash *felt.Felt
+		switch cairoClass := cls.Class.(type) {
+		case *core.Cairo0Class:
+			compiledHash = class.ClassHash
+		case *core.Cairo1Class:
+			compiledHash = cairoClass.Hash()
+		default:
+			b.log.Errorw("Unknown cairo class", "cairoClass", cairoClass)
+			return b.fin()
+		}
+
+		classesM[*compiledHash] = core2p2p.AdaptClass(cls.Class, compiledHash)
+	}
+
+	var classes []*spec.Class
+	for _, c := range classesM {
+		c := c
+		classes = append(classes, c)
+	}
+
+	return &spec.BlockBodiesResponse{
+		Id: core2p2p.AdaptBlockID(b.header),
+		BodyMessage: &spec.BlockBodiesResponse_Classes{
+			Classes: &spec.Classes{
+				Domain:  0,
+				Classes: classes,
+			},
+		},
+	}, true
+}
+
+type contractDiff struct {
+	address      *felt.Felt
+	classHash    *felt.Felt
+	storageDiffs []core.StorageDiff
+	nonce        *felt.Felt
+}
+
+func (b *blockBodyIterator) diff() (proto.Message, bool) {
+	var err error
+	diff := b.stateUpdate.StateDiff
+
+	modifiedContracts := make(map[felt.Felt]*contractDiff)
+	initContractDiff := func(addr *felt.Felt) (*contractDiff, error) {
+		var cHash *felt.Felt
+		cHash, err = b.stateReader.ContractClassHash(addr)
+		if err != nil {
+			return nil, err
+		}
+		return &contractDiff{address: addr, classHash: cHash}, nil
+	}
+
+	for addr, n := range diff.Nonces {
+		addr := addr // copy
+		cDiff, ok := modifiedContracts[addr]
+		if !ok {
+			cDiff, err = initContractDiff(&addr)
+			if err != nil {
+				b.log.Errorw("Failed to get class hash", "err", err)
+				return b.fin()
+			}
+			modifiedContracts[addr] = cDiff
+		}
+		cDiff.nonce = n
+	}
+
+	for addr, sDiff := range diff.StorageDiffs {
+		addr := addr // copy
+		cDiff, ok := modifiedContracts[addr]
+		if !ok {
+			cDiff, err = initContractDiff(&addr)
+			if err != nil {
+				b.log.Errorw("Failed to get class hash", "err", err)
+				return b.fin()
+			}
+			modifiedContracts[addr] = cDiff
+		}
+		cDiff.storageDiffs = sDiff
+	}
+
+	var contractDiffs []*spec.StateDiff_ContractDiff
+	for _, c := range modifiedContracts {
+		contractDiffs = append(contractDiffs, core2p2p.AdaptStateDiff(c.address, c.classHash, c.nonce, c.storageDiffs))
+	}
+
+	return &spec.BlockBodiesResponse{
+		Id: core2p2p.AdaptBlockID(b.header),
+		BodyMessage: &spec.BlockBodiesResponse_Diff{
+			Diff: &spec.StateDiff{
+				Domain:        0,
+				ContractDiffs: contractDiffs,
+			},
+		},
+	}, true
+}
+
+func (b *blockBodyIterator) fin() (proto.Message, bool) {
+	b.step = terminal
+	if err := b.stateCloser(); err != nil {
+		b.log.Errorw("Call to state closer failed", "err", err)
+	}
+	return &spec.BlockBodiesResponse{
+		Id:          core2p2p.AdaptBlockID(b.header),
+		BodyMessage: &spec.BlockBodiesResponse_Fin{},
+	}, true
+}
+
+func (b *blockBodyIterator) proof() (proto.Message, bool) {
+	// proof size is currently 142K
+	proof := make([]byte, 142*1024) //nolint:gomnd
+	_, err := rand.Read(proof)
+	if err != nil {
+		b.log.Errorw("Failed to generate rand proof", "err", err)
+		return b.fin()
+	}
+
+	return &spec.BlockBodiesResponse{
+		Id: core2p2p.AdaptBlockID(b.header),
+		BodyMessage: &spec.BlockBodiesResponse_Proof{
+			Proof: &spec.BlockProof{
+				Proof: proof,
+			},
+		},
+	}, true
+}

p2p/starknet/handlers.go

@@ -150,18 +150,40 @@ func (h *Handler) onBlockHeadersRequest(req *spec.BlockHeadersRequest) (Stream[p
 }
 
 func (h *Handler) onBlockBodiesRequest(req *spec.BlockBodiesRequest) (Stream[proto.Message], error) {
-	// todo: read from bcReader and adapt to p2p type
-	count := uint64(0)
+	it, err := h.newIterator(req.Iteration)
+	if err != nil {
+		return nil, err
+	}
+
+	fin := h.newFin(&spec.BlockBodiesResponse{
+		BodyMessage: &spec.BlockBodiesResponse_Fin{},
+	})
+
+	var bodyIterator *blockBodyIterator
 	return func() (proto.Message, bool) {
-		if count > 3 {
-			return nil, false
+		// bodyIterator is nil only during the first iteration
+		if bodyIterator != nil && bodyIterator.hasNext() {
+			return bodyIterator.next()
 		}
-		count++
-		return &spec.BlockBodiesResponse{
-			Id: &spec.BlockID{
-				Number: count - 1,
-			},
-		}, true
+
+		if !it.Valid() {
+			return fin()
+		}
+
+		header, err := it.Header()
+		if err != nil {
+			h.log.Errorw("Failed to fetch block header", "err", err)
+			return fin()
+		}
+		it.Next()
+
+		bodyIterator, err = newBlockBodyIterator(h.bcReader, header, h.log)
+		if err != nil {
+			h.log.Errorw("Failed to create block body iterator", "err", err)
+			return fin()
+		}
+		// no need to call hasNext since it's first iteration over a block
+		return bodyIterator.next()
 	}, nil
 }
 

p2p/starknet/iterator_test.go

@@ -1,7 +1,6 @@
 package starknet
 
 import (
-	"fmt"
 	"testing"
 
 	"github.com/NethermindEth/juno/core"
@@ -126,7 +125,6 @@ func TestIterator(t *testing.T) {
 
 			var i int
 			for it.Valid() {
-				fmt.Println("Block")
 				block, err := it.Block()
 				if err != nil {
 					assert.Equal(t, err, db.ErrKeyNotFound)