Merge pull request #166 from cloudstruct/feat/protocol-state-transition-timeout

feat: support for timeouts on protocol state transitions

Author: agaffney

protocol/handshake/handshake.go

@@ -1,6 +1,8 @@
 package handshake
 
 import (
+	"time"
+
 	"github.com/cloudstruct/go-ouroboros-network/protocol"
 )
 
@@ -29,7 +31,8 @@ var StateMap = protocol.StateMap{
 		},
 	},
 	STATE_CONFIRM: protocol.StateMapEntry{
-		Agency: protocol.AGENCY_SERVER,
+		Agency:  protocol.AGENCY_SERVER,
+		Timeout: 5 * time.Second,
 		Transitions: []protocol.StateTransition{
 			{
 				MsgType:  MESSAGE_TYPE_ACCEPT_VERSION,

protocol/protocol.go

@@ -3,12 +3,14 @@ package protocol
 import (
 	"bytes"
 	"fmt"
-	"github.com/cloudstruct/go-ouroboros-network/muxer"
-	"github.com/cloudstruct/go-ouroboros-network/utils"
-	"github.com/fxamacker/cbor/v2"
 	"io"
 	"reflect"
 	"sync"
+	"time"
+
+	"github.com/cloudstruct/go-ouroboros-network/muxer"
+	"github.com/cloudstruct/go-ouroboros-network/utils"
+	"github.com/fxamacker/cbor/v2"
 )
 
 const (
@@ -17,17 +19,18 @@ const (
 )
 
 type Protocol struct {
-	config             ProtocolConfig
-	muxerSendChan      chan *muxer.Segment
-	muxerRecvChan      chan *muxer.Segment
-	state              State
-	stateMutex         sync.Mutex
-	recvBuffer         *bytes.Buffer
-	sendQueueChan      chan Message
-	sendStateQueueChan chan Message
-	recvReadyChan      chan bool
-	sendReadyChan      chan bool
-	doneChan           chan bool
+	config               ProtocolConfig
+	muxerSendChan        chan *muxer.Segment
+	muxerRecvChan        chan *muxer.Segment
+	state                State
+	stateMutex           sync.Mutex
+	recvBuffer           *bytes.Buffer
+	sendQueueChan        chan Message
+	sendStateQueueChan   chan Message
+	recvReadyChan        chan bool
+	sendReadyChan        chan bool
+	doneChan             chan bool
+	stateTransitionTimer *time.Timer
 }
 
 type ProtocolConfig struct {
@@ -316,6 +319,13 @@ func (p *Protocol) getNewState(msg Message) (State, error) {
 }
 
 func (p *Protocol) setState(state State) {
+	// Disable any previous state transition timer
+	if p.stateTransitionTimer != nil {
+		if !p.stateTransitionTimer.Stop() {
+			<-p.stateTransitionTimer.C
+		}
+		p.stateTransitionTimer = nil
+	}
 	// Set the new state
 	p.state = state
 	// Mark protocol as ready to send/receive based on role and agency of the new state
@@ -335,6 +345,15 @@ func (p *Protocol) setState(state State) {
 			p.recvReadyChan <- true
 		}
 	}
+	// Set timeout for state transition
+	if p.config.StateMap[p.state].Timeout > 0 {
+		p.stateTransitionTimer = time.AfterFunc(
+			p.config.StateMap[p.state].Timeout,
+			func() {
+				p.SendError(fmt.Errorf("%s: timeout waiting on transition from protocol state %s", p.config.Name, p.state))
+			},
+		)
+	}
 }
 
 func (p *Protocol) handleMessage(msg Message, isResponse bool) error {

protocol/state.go

@@ -1,5 +1,9 @@
 package protocol
 
+import (
+	"time"
+)
+
 const (
 	AGENCY_NONE   uint = 0
 	AGENCY_CLIENT uint = 1
@@ -33,6 +37,7 @@ type StateTransitionMatchFunc func(Message) bool
 type StateMapEntry struct {
 	Agency      uint
 	Transitions []StateTransition
+	Timeout     time.Duration
 }
 
 type StateMap map[State]StateMapEntry