feat: add a timeout for handshake propose on server side

Fixes #582
blinklabs-io · Apr 10, 2024 · 85d6a5a · 85d6a5a
1 parent c051f8e
commit 85d6a5a
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/protocol/handshake/server.go b/protocol/handshake/server.go
@@ -37,6 +37,12 @@ func NewServer(protoOptions protocol.ProtocolOptions, cfg *Config) *Server {
 		Server:       s,
 		ConnectionId: protoOptions.ConnectionId,
 	}
+	// Update state map with timeout
+	stateMap := StateMap.Copy()
+	if entry, ok := stateMap[statePropose]; ok {
+		entry.Timeout = s.config.Timeout
+		stateMap[statePropose] = entry
+	}
 	protoConfig := protocol.ProtocolConfig{
 		Name:                ProtocolName,
 		ProtocolId:          ProtocolId,
@@ -46,7 +52,7 @@ func NewServer(protoOptions protocol.ProtocolOptions, cfg *Config) *Server {
 		Role:                protocol.ProtocolRoleServer,
 		MessageHandlerFunc:  s.handleMessage,
 		MessageFromCborFunc: NewMsgFromCbor,
-		StateMap:            StateMap,
+		StateMap:            stateMap,
 		InitialState:        statePropose,
 	}
 	s.Protocol = protocol.New(protoConfig)