Merge pull request #36 from primevprotocol/iowar/optimization.0

optimization: repeated handshake creation removed

pkg/p2p/libp2p/internal/handshake/handshake.go

@@ -31,6 +31,7 @@ type Service struct {
 	signer        signer.Signer
 	register      register.Register
 	minimumStake  *big.Int
+	handshakeReq  *HandshakeReq
 	getEthAddress func(core.PeerID) (common.Address, error)
 }
 
@@ -43,8 +44,8 @@ func New(
 	register register.Register,
 	minimumStake *big.Int,
 	getEthAddress func(core.PeerID) (common.Address, error),
-) *Service {
-	return &Service{
+) (*Service, error) {
+	s := &Service{
 		privKey:       privKey,
 		ethAddress:    ethAddress,
 		peerType:      peerType,
@@ -54,6 +55,13 @@ func New(
 		minimumStake:  minimumStake,
 		getEthAddress: getEthAddress,
 	}
+
+	err := s.setHandshakeReq()
+	if err != nil {
+		return nil, err
+	}
+
+	return s, nil
 }
 
 func ProtocolID() protocol.ID {
@@ -123,6 +131,22 @@ func (h *Service) createSignature() ([]byte, error) {
 	return sig, nil
 }
 
+func (h *Service) setHandshakeReq() error {
+	sig, err := h.createSignature()
+	if err != nil {
+		return err
+	}
+
+	req := &HandshakeReq{
+		PeerType: h.peerType.String(),
+		Token:    h.passcode,
+		Sig:      sig,
+	}
+
+	h.handshakeReq = req
+	return nil
+}
+
 func (h *Service) verifyResp(resp *HandshakeResp) error {
 	if !bytes.Equal(resp.ObservedAddress.Bytes(), h.ethAddress.Bytes()) {
 		return errors.New("observed address mismatch")
@@ -152,11 +176,6 @@ func (h *Service) Handle(
 		return p2p.Peer{}, err
 	}
 
-	sig, err := h.createSignature()
-	if err != nil {
-		return p2p.Peer{}, err
-	}
-
 	resp := &HandshakeResp{
 		ObservedAddress: ethAddress,
 		PeerType:        req.PeerType,
@@ -168,12 +187,7 @@ func (h *Service) Handle(
 
 	ar, aw := msgpack.NewReaderWriter[HandshakeResp, HandshakeReq](stream)
 
-	err = aw.WriteMsg(ctx, &HandshakeReq{
-		PeerType: h.peerType.String(),
-		Token:    h.passcode,
-		Sig:      sig,
-	},
-	)
+	err = aw.WriteMsg(ctx, h.handshakeReq)
 	if err != nil {
 		return p2p.Peer{}, err
 	}
@@ -199,20 +213,9 @@ func (h *Service) Handshake(
 	stream p2p.Stream,
 ) (p2p.Peer, error) {
 
-	sig, err := h.createSignature()
-	if err != nil {
-		return p2p.Peer{}, err
-	}
-
-	req := &HandshakeReq{
-		PeerType: h.peerType.String(),
-		Token:    h.passcode,
-		Sig:      sig,
-	}
-
 	r, w := msgpack.NewReaderWriter[HandshakeResp, HandshakeReq](stream)
 
-	if err := w.WriteMsg(ctx, req); err != nil {
+	if err := w.WriteMsg(ctx, h.handshakeReq); err != nil {
 		return p2p.Peer{}, err
 	}
 

pkg/p2p/libp2p/internal/handshake/handshake_test.go

@@ -52,7 +52,7 @@ func TestHandshake(t *testing.T) {
 			t.Fatal(err)
 		}
 
-		hs1 := handshake.New(
+		hs1, err := handshake.New(
 			privKey1,
 			common.HexToAddress("0x1"),
 			p2p.PeerTypeBuilder,
@@ -64,7 +64,11 @@ func TestHandshake(t *testing.T) {
 				return common.HexToAddress("0x2"), nil
 			},
 		)
-		hs2 := handshake.New(
+		if err != nil {
+			t.Fatal(err)
+		}
+
+		hs2, err := handshake.New(
 			privKey2,
 			common.HexToAddress("0x2"),
 			p2p.PeerTypeBuilder,
@@ -76,6 +80,9 @@ func TestHandshake(t *testing.T) {
 				return common.HexToAddress("0x1"), nil
 			},
 		)
+		if err != nil {
+			t.Fatal(err)
+		}
 
 		out, in := p2ptest.NewDuplexStream()
 

pkg/p2p/libp2p/libp2p.go

@@ -113,7 +113,7 @@ func New(opts *Options) (*Service, error) {
 		return nil, err
 	}
 
-	hsSvc := handshake.New(
+	hsSvc, err := handshake.New(
 		opts.PrivKey,
 		ethAddress,
 		opts.PeerType,
@@ -123,6 +123,9 @@ func New(opts *Options) (*Service, error) {
 		opts.MinimumStake,
 		GetEthAddressFromPeerID,
 	)
+	if err != nil {
+		return nil, err
+	}
 
 	baseCtx, baseCtxCancel := context.WithCancel(context.Background())