세션 무효화시 접속 종료 처리

wshandler/wshandler_peer.go

@@ -22,10 +22,23 @@ type WebsocketPeerHandler interface {
 	RegisterHandlers(serveMux *http.ServeMux, prefix string) error
 }
 
+type peerCtorChannelValue struct {
+	accid primitive.ObjectID
+	conn  *websocket.Conn
+}
+
+type peerDtorChannelValue struct {
+	accid  primitive.ObjectID
+	closed bool
+}
+
 type websocketPeerHandler[T PeerInterface] struct {
 	methods         map[string]peerApiFuncType[T]
 	createPeer      func(primitive.ObjectID) T
 	sessionConsumer session.Consumer
+
+	peerCtorChannel chan peerCtorChannelValue
+	peerDtorChannel chan peerDtorChannelValue
 }
 
 type PeerInterface interface {
@@ -144,8 +157,11 @@ func NewWebsocketPeerHandler[T PeerInterface](consumer session.Consumer, creator
 		sessionConsumer: consumer,
 		methods:         methods,
 		createPeer:      creator,
+		peerCtorChannel: make(chan peerCtorChannelValue),
+		peerDtorChannel: make(chan peerDtorChannelValue),
 	}
 
+	consumer.RegisterOnSessionInvalidated(wsh.onSessionInvalidated)
 	return wsh
 }
 
@@ -155,18 +171,47 @@ func (ws *websocketPeerHandler[T]) RegisterHandlers(serveMux *http.ServeMux, pre
 	} else {
 		serveMux.HandleFunc(prefix, ws.upgrade)
 	}
+	go ws.sessionMonitoring()
 
 	return nil
 }
 
+func (ws *websocketPeerHandler[T]) onSessionInvalidated(accid primitive.ObjectID) {
+	ws.peerDtorChannel <- peerDtorChannelValue{
+		accid:  accid,
+		closed: false,
+	}
+}
+
+func (ws *websocketPeerHandler[T]) sessionMonitoring() {
+	all := make(map[primitive.ObjectID]*websocket.Conn)
+	unauthdata := []byte{0x03, 0xec}
+	unauthdata = append(unauthdata, []byte("unauthorized")...)
+	for {
+		select {
+		case estVal := <-ws.peerCtorChannel:
+			all[estVal.accid] = estVal.conn
+		case disVal := <-ws.peerDtorChannel:
+			if disVal.closed {
+				delete(all, disVal.accid)
+			} else if c := all[disVal.accid]; c != nil {
+				c.WriteControl(websocket.CloseMessage, unauthdata, time.Time{})
+				delete(all, disVal.accid)
+			}
+		}
+	}
+}
+
 func (ws *websocketPeerHandler[T]) upgrade_core(conn *websocket.Conn, accid primitive.ObjectID, nonce uint32) {
 	go func(c *websocket.Conn, accid primitive.ObjectID) {
 		peer := ws.createPeer(accid)
 		var closeReason string
 
 		peer.ClientConnected(conn)
+		ws.peerCtorChannel <- peerCtorChannelValue{accid: accid, conn: conn}
 
 		defer func() {
+			ws.peerDtorChannel <- peerDtorChannelValue{accid: accid, closed: true}
 			peer.ClientDisconnected(closeReason)
 		}()