修改bug

websocket/websocket.go

@@ -10,11 +10,11 @@ const (
 	connectionContextKey = "connection-context"
 )
 
-type HandleConnectFunc func(groupID string)
-type HandleDisconnectFunc func(groupID string)
-type HandleErrorFunc func(groupID string, err error)
-type HandleCloseFunc func(groupID string, i int, s string) error
-type HandlePongFunc func(groupID string)
+type HandleConnectFunc func(groupID string, context any)
+type HandleDisconnectFunc func(groupID string, context any)
+type HandleErrorFunc func(groupID string, err error, context any)
+type HandleCloseFunc func(groupID string, i int, s string, context any) error
+type HandlePongFunc func(groupID string, context any)
 type HandleMessageFunc func(groupID string, message []byte, context any)
 
 var managerInstance *Manager
@@ -49,7 +49,7 @@ type Manager struct {
 func (m *Manager) HandleConnect(handleConnectFunc HandleConnectFunc) {
 	m.melodyInstance.HandleConnect(func(session *melody.Session) {
 		if handleConnectFunc != nil {
-			handleConnectFunc(session.Keys[groupIDKey].(string))
+			handleConnectFunc(session.Keys[groupIDKey].(string), session.Keys[connectionContextKey])
 		}
 	})
 }
@@ -57,7 +57,7 @@ func (m *Manager) HandleConnect(handleConnectFunc HandleConnectFunc) {
 func (m *Manager) HandleDisconnect(handleDisconnectFunc HandleDisconnectFunc) {
 	m.melodyInstance.HandleDisconnect(func(session *melody.Session) {
 		if handleDisconnectFunc != nil {
-			handleDisconnectFunc(session.Keys[groupIDKey].(string))
+			handleDisconnectFunc(session.Keys[groupIDKey].(string), session.Keys[connectionContextKey])
 		}
 	})
 }
@@ -65,7 +65,7 @@ func (m *Manager) HandleDisconnect(handleDisconnectFunc HandleDisconnectFunc) {
 func (m *Manager) HandleError(handleErrorFunc HandleErrorFunc) {
 	m.melodyInstance.HandleError(func(session *melody.Session, err error) {
 		if handleErrorFunc != nil {
-			handleErrorFunc(session.Keys[groupIDKey].(string), err)
+			handleErrorFunc(session.Keys[groupIDKey].(string), err, session.Keys[connectionContextKey])
 		}
 	})
 }
@@ -73,7 +73,7 @@ func (m *Manager) HandleError(handleErrorFunc HandleErrorFunc) {
 func (m *Manager) HandleClose(handleCloseFunc HandleCloseFunc) {
 	m.melodyInstance.HandleClose(func(session *melody.Session, i int, s string) error {
 		if handleCloseFunc != nil {
-			err := handleCloseFunc(session.Keys[groupIDKey].(string), i, s)
+			err := handleCloseFunc(session.Keys[groupIDKey].(string), i, s, session.Keys[connectionContextKey])
 			if err != nil {
 				return err
 			}
@@ -86,7 +86,7 @@ func (m *Manager) HandleClose(handleCloseFunc HandleCloseFunc) {
 func (m *Manager) HandlePong(handlePongFunc HandlePongFunc) {
 	m.melodyInstance.HandlePong(func(session *melody.Session) {
 		if handlePongFunc != nil {
-			handlePongFunc(session.Keys[groupIDKey].(string))
+			handlePongFunc(session.Keys[groupIDKey].(string), session.Keys[connectionContextKey])
 		}
 	})
 }
@@ -128,8 +128,8 @@ func (m *Manager) BroadCast(groupID string, msg []byte) error {
 
 type ConnectionOption func(sessionMap map[string]any)
 
-func WithConnectionContext(context any) ConnectionOption {
+func WithConnectionContext(key string, context any) ConnectionOption {
 	return func(sessionMap map[string]any) {
-		sessionMap[connectionContextKey] = context
+		sessionMap[connectionContextKey+"::"+key] = context
 	}
 }