Merge pull request #177 from cloudstruct/feat/handshake-timeout

feat: make handshake timeout configurable

Author: agaffney

protocol/handshake/client.go

@@ -11,9 +11,20 @@ type Client struct {
 }
 
 func NewClient(protoOptions protocol.ProtocolOptions, cfg *Config) *Client {
+	if cfg == nil {
+		tmpCfg := NewConfig()
+		cfg = &tmpCfg
+	}
 	c := &Client{
 		config: cfg,
 	}
+	// Update state map with timeout
+	stateMap := StateMap
+	if entry, ok := stateMap[STATE_CONFIRM]; ok {
+		entry.Timeout = c.config.Timeout
+		stateMap[STATE_CONFIRM] = entry
+	}
+	// Configure underlying Protocol
 	protoConfig := protocol.ProtocolConfig{
 		Name:                PROTOCOL_NAME,
 		ProtocolId:          PROTOCOL_ID,

protocol/handshake/handshake.go

@@ -31,8 +31,7 @@ var StateMap = protocol.StateMap{
 		},
 	},
 	STATE_CONFIRM: protocol.StateMapEntry{
-		Agency:  protocol.AGENCY_SERVER,
-		Timeout: 5 * time.Second,
+		Agency: protocol.AGENCY_SERVER,
 		Transitions: []protocol.StateTransition{
 			{
 				MsgType:  MESSAGE_TYPE_ACCEPT_VERSION,
@@ -59,6 +58,7 @@ type Config struct {
 	NetworkMagic     uint32
 	ClientFullDuplex bool
 	FinishedFunc     FinishedFunc
+	Timeout          time.Duration
 }
 
 type FinishedFunc func(uint16, bool) error
@@ -74,7 +74,9 @@ func New(protoOptions protocol.ProtocolOptions, cfg *Config) *Handshake {
 type HandshakeOptionFunc func(*Config)
 
 func NewConfig(options ...HandshakeOptionFunc) Config {
-	c := Config{}
+	c := Config{
+		Timeout: 5 * time.Second,
+	}
 	// Apply provided options functions
 	for _, option := range options {
 		option(&c)
@@ -105,3 +107,9 @@ func WithFinishedFunc(finishedFunc FinishedFunc) HandshakeOptionFunc {
 		c.FinishedFunc = finishedFunc
 	}
 }
+
+func WithTimeout(timeout time.Duration) HandshakeOptionFunc {
+	return func(c *Config) {
+		c.Timeout = timeout
+	}
+}