feat: no error return on connection close

Signed-off-by: Jenita <[email protected]>

Author: Jenita

protocol/blockfetch/blockfetch.go

@@ -121,24 +121,17 @@ func New(protoOptions protocol.ProtocolOptions, cfg *Config) *BlockFetch {
 	return b
 }
 
-func (b *BlockFetch) IsDone() bool {
-	if b.Client != nil && b.Client.IsDone() {
-		return true
-	}
-	if b.Server != nil && b.Server.IsDone() {
-		return true
-	}
-	return false
-}
-
 func (b *BlockFetch) HandleConnectionError(err error) error {
 	if err == nil {
 		return nil
 	}
+	// Check if protocol is done or if it's a normal connection closure
+	if b.Client.IsDone() || b.Server.IsDone() {
+		return nil
+	}
+
 	if errors.Is(err, io.EOF) || isConnectionReset(err) {
-		if b.IsDone() {
-			return nil
-		}
+		return err
 	}
 	return err
 }

protocol/blockfetch/blockfetch_test.go

@@ -38,14 +38,33 @@ func (a testAddr) String() string  { return "test-addr" }
 // testConn implements net.Conn for testing with buffered writes
 type testConn struct {
 	writeChan chan []byte
+	closed    bool
+	closeChan chan struct{}
+}
+
+func newTestConn() *testConn {
+	return &testConn{
+		writeChan: make(chan []byte, 100),
+		closeChan: make(chan struct{}),
+	}
 }
 
 func (c *testConn) Read(b []byte) (n int, err error) { return 0, nil }
 func (c *testConn) Write(b []byte) (n int, err error) {
-	c.writeChan <- b
-	return len(b), nil
+	select {
+	case c.writeChan <- b:
+		return len(b), nil
+	case <-c.closeChan:
+		return 0, io.EOF
+	}
+}
+func (c *testConn) Close() error {
+	if !c.closed {
+		close(c.closeChan)
+		c.closed = true
+	}
+	return nil
 }
-func (c *testConn) Close() error                       { return nil }
 func (c *testConn) LocalAddr() net.Addr                { return testAddr{} }
 func (c *testConn) RemoteAddr() net.Addr               { return testAddr{} }
 func (c *testConn) SetDeadline(t time.Time) error      { return nil }
@@ -54,6 +73,12 @@ func (c *testConn) SetWriteDeadline(t time.Time) error { return nil }
 
 func getTestProtocolOptions(conn net.Conn) protocol.ProtocolOptions {
 	mux := muxer.New(conn)
+	go mux.Start()
+	go func() {
+		<-conn.(*testConn).closeChan
+		mux.Stop()
+	}()
+
 	return protocol.ProtocolOptions{
 		ConnectionId: connection.ConnectionId{
 			LocalAddr:  testAddr{},
@@ -65,7 +90,8 @@ func getTestProtocolOptions(conn net.Conn) protocol.ProtocolOptions {
 }
 
 func TestNewBlockFetch(t *testing.T) {
-	conn := &testConn{writeChan: make(chan []byte, 1)}
+	conn := newTestConn()
+	defer conn.Close()
 	cfg := NewConfig()
 	bf := New(getTestProtocolOptions(conn), &cfg)
 	assert.NotNil(t, bf.Client)
@@ -92,10 +118,17 @@ func TestConfigOptions(t *testing.T) {
 }
 
 func TestConnectionErrorHandling(t *testing.T) {
-	conn := &testConn{writeChan: make(chan []byte, 1)}
+	conn := newTestConn()
+	defer conn.Close()
 	cfg := NewConfig()
 	bf := New(getTestProtocolOptions(conn), &cfg)
 
+	// Start protocols
+	bf.Client.Start()
+	defer bf.Client.Stop()
+	bf.Server.Start()
+	defer bf.Server.Stop()
+
 	t.Run("Non-EOF error when not done", func(t *testing.T) {
 		err := bf.HandleConnectionError(errors.New("test error"))
 		assert.Error(t, err)
@@ -110,10 +143,23 @@ func TestConnectionErrorHandling(t *testing.T) {
 		err := bf.HandleConnectionError(errors.New("connection reset by peer"))
 		assert.Error(t, err)
 	})
+
+	t.Run("EOF error when done", func(t *testing.T) {
+		// Send done message to properly transition to done state
+		err := bf.Client.SendMessage(NewMsgClientDone())
+		assert.NoError(t, err)
+
+		// Wait for state transition
+		time.Sleep(100 * time.Millisecond)
+
+		err = bf.HandleConnectionError(io.EOF)
+		assert.NoError(t, err, "expected no error when protocol is in done state")
+	})
 }
 
 func TestCallbackRegistration(t *testing.T) {
-	conn := &testConn{writeChan: make(chan []byte, 1)}
+	conn := newTestConn()
+	defer conn.Close()
 
 	t.Run("Block callback registration", func(t *testing.T) {
 		blockFunc := func(ctx CallbackContext, slot uint, block ledger.Block) error {
@@ -137,16 +183,17 @@ func TestCallbackRegistration(t *testing.T) {
 }
 
 func TestClientMessageSending(t *testing.T) {
-	conn := &testConn{writeChan: make(chan []byte, 1)}
+	conn := newTestConn()
+	defer conn.Close()
 	cfg := NewConfig()
 	client := NewClient(getTestProtocolOptions(conn), &cfg)
 
 	t.Run("Client can send messages", func(t *testing.T) {
-		// Start the client protocol
 		client.Start()
+		defer client.Stop()
 
 		// Send a done message
-		err := client.Protocol.SendMessage(NewMsgClientDone())
+		err := client.SendMessage(NewMsgClientDone())
 		assert.NoError(t, err)
 
 		// Verify message was written to connection
@@ -158,14 +205,3 @@ func TestClientMessageSending(t *testing.T) {
 		}
 	})
 }
-
-func TestServerMessageHandling(t *testing.T) {
-	conn := &testConn{writeChan: make(chan []byte, 1)}
-	cfg := NewConfig()
-	server := NewServer(getTestProtocolOptions(conn), &cfg)
-
-	t.Run("Server can be started", func(t *testing.T) {
-		server.Start()
-
-	})
-}

protocol/blockfetch/client.go

@@ -35,8 +35,6 @@ type Client struct {
 	blockUseCallback     bool
 	onceStart            sync.Once
 	onceStop             sync.Once
-	currentState         protocol.State
-	stateMutex           sync.Mutex
 }
 
 func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
@@ -48,7 +46,6 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 		config:               cfg,
 		blockChan:            make(chan ledger.Block),
 		startBatchResultChan: make(chan error),
-		currentState:         StateIdle,
 	}
 	c.callbackContext = CallbackContext{
 		Client:       c,
@@ -85,18 +82,6 @@ func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
 	return c
 }
 
-func (c *Client) IsDone() bool {
-	c.stateMutex.Lock()
-	defer c.stateMutex.Unlock()
-	return c.currentState.Id == StateDone.Id
-}
-
-func (c *Client) setState(newState protocol.State) {
-	c.stateMutex.Lock()
-	defer c.stateMutex.Unlock()
-	c.currentState = newState
-}
-
 func (c *Client) Start() {
 	c.onceStart.Do(func() {
 		c.Protocol.Logger().
@@ -124,12 +109,12 @@ func (c *Client) Stop() error {
 				"protocol", ProtocolName,
 				"connection_id", c.callbackContext.ConnectionId.String(),
 			)
-		msg := NewMsgClientDone()
-		if sendErr := c.SendMessage(msg); sendErr != nil {
-			err = sendErr
-			return
+		if !c.IsDone() {
+			msg := NewMsgClientDone()
+			if err = c.SendMessage(msg); err != nil {
+				return
+			}
 		}
-		c.setState(StateDone)
 	})
 	return err
 }
@@ -216,7 +201,7 @@ func (c *Client) messageHandler(msg protocol.Message) error {
 	case MessageTypeBatchDone:
 		err = c.handleBatchDone()
 	case MessageTypeClientDone:
-		c.setState(StateDone)
+		return nil
 	default:
 		err = fmt.Errorf(
 			"%s: received unexpected message type %d",

protocol/blockfetch/server.go

@@ -17,7 +17,6 @@ package blockfetch
 import (
 	"errors"
 	"fmt"
-	"sync"
 
 	"github.com/blinklabs-io/gouroboros/cbor"
 	"github.com/blinklabs-io/gouroboros/protocol"
@@ -28,16 +27,13 @@ type Server struct {
 	config          *Config
 	callbackContext CallbackContext
 	protoOptions    protocol.ProtocolOptions
-	currentState    protocol.State
-	stateMutex      sync.Mutex
 }
 
 func NewServer(protoOptions protocol.ProtocolOptions, cfg *Config) *Server {
 	s := &Server{
 		config: cfg,
 		// Save this for re-use later
 		protoOptions: protoOptions,
-		currentState: StateIdle,
 	}
 	s.callbackContext = CallbackContext{
 		Server:       s,
@@ -47,18 +43,6 @@ func NewServer(protoOptions protocol.ProtocolOptions, cfg *Config) *Server {
 	return s
 }
 
-func (s *Server) IsDone() bool {
-	s.stateMutex.Lock()
-	defer s.stateMutex.Unlock()
-	return s.currentState.Id == StateDone.Id
-}
-
-func (s *Server) setState(newState protocol.State) {
-	s.stateMutex.Lock()
-	defer s.stateMutex.Unlock()
-	s.currentState = newState
-}
-
 func (s *Server) initProtocol() {
 	protoConfig := protocol.ProtocolConfig{
 		Name:                ProtocolName,
@@ -142,8 +126,8 @@ func (s *Server) messageHandler(msg protocol.Message) error {
 	case MessageTypeRequestRange:
 		err = s.handleRequestRange(msg)
 	case MessageTypeClientDone:
-		s.setState(StateDone)
-		err = s.handleClientDone()
+		// State handled automatically by base protocol
+		return s.handleClientDone()
 	default:
 		err = fmt.Errorf(
 			"%s: received unexpected message type %d",

protocol/chainsync/chainsync.go

@@ -195,13 +195,10 @@ var PipelineIsNotEmpty = func(context any, msg protocol.Message) bool {
 
 // ChainSync is a wrapper object that holds the client and server instances
 type ChainSync struct {
-	Client       *Client
-	Server       *Server
-	stateMutex   sync.Mutex
-	currentState protocol.State
+	Client *Client
+	Server *Server
 }
 
-// Config is used to configure the ChainSync protocol instance
 type Config struct {
 	RollBackwardFunc   RollBackwardFunc
 	RollForwardFunc    RollForwardFunc
@@ -335,26 +332,20 @@ func WithRecvQueueSize(size int) ChainSyncOptionFunc {
 	}
 }
 
-// HandleConnectionError handles connection errors and determines if they should be ignored
 func (c *ChainSync) HandleConnectionError(err error) error {
 	if err == nil {
 		return nil
 	}
+	if c.Client.IsDone() || c.Server.IsDone() {
+		return nil
+	}
+
 	if errors.Is(err, io.EOF) || isConnectionReset(err) {
-		if c.IsDone() {
-			return nil
-		}
+		return err
 	}
 	return err
 }
 
-// IsDone returns true if the protocol is in the Done state
-func (c *ChainSync) IsDone() bool {
-	c.stateMutex.Lock()
-	defer c.stateMutex.Unlock()
-	return c.currentState.Id == stateDone.Id
-}
-
 func isConnectionReset(err error) bool {
 	return strings.Contains(err.Error(), "connection reset") ||
 		strings.Contains(err.Error(), "broken pipe")