修改websocket

websocket/websocket.go

@@ -6,7 +6,8 @@ import (
 )
 
 const (
-	groupIDKey = "group-id"
+	groupIDKey           = "group-id"
+	connectionContextKey = "connection-context"
 )
 
 type HandleConnectFunc func(groupID string)
@@ -90,8 +91,16 @@ func (m *Manager) HandlePong(handlePongFunc HandlePongFunc) {
 	})
 }
 
-func (m *Manager) HandleRequest(groupID string, w http.ResponseWriter, r *http.Request) error {
-	err := m.melodyInstance.HandleRequestWithKeys(w, r, map[string]interface{}{groupIDKey: groupID})
+func (m *Manager) HandleRequest(groupID string, w http.ResponseWriter, r *http.Request, opts ...ConnectionOption) error {
+	sessionMap := map[string]interface{}{
+		groupIDKey: groupID,
+	}
+
+	for _, opt := range opts {
+		opt(sessionMap)
+	}
+
+	err := m.melodyInstance.HandleRequestWithKeys(w, r, sessionMap)
 	if err != nil {
 		return err
 	}
@@ -116,3 +125,11 @@ func (m *Manager) BroadCast(groupID string, msg []byte) error {
 		return true
 	})
 }
+
+type ConnectionOption func(sessionMap map[string]any)
+
+func WithConnectionContext(context any) ConnectionOption {
+	return func(sessionMap map[string]any) {
+		sessionMap[connectionContextKey] = context
+	}
+}