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

wshandler/wshandler.go

@@ -93,6 +93,7 @@ type websocketHandlerBase struct {
 	connInOutChan     chan *wsconn
 	deliveryChan      chan any
 	localDeliveryChan chan any
+	forceCloseChan    chan primitive.ObjectID
 	sendMsgChan       chan send_msg_queue_elem
 
 	connWaitGroup   sync.WaitGroup
@@ -154,7 +155,7 @@ func NewWebsocketHandler(consumer session.Consumer, redisUrl string) (*Websocket
 		}
 	}()
 
-	return &WebsocketHandler{
+	ws := &WebsocketHandler{
 		websocketHandlerBase: websocketHandlerBase{
 			redisMsgChanName:  fmt.Sprintf("_wsh_msg_%d", redisSync.Options().DB),
 			redisCmdChanName:  fmt.Sprintf("_wsh_cmd_%d", redisSync.Options().DB),
@@ -162,10 +163,13 @@ func NewWebsocketHandler(consumer session.Consumer, redisUrl string) (*Websocket
 			connInOutChan:     make(chan *wsconn),
 			deliveryChan:      make(chan any, 1000),
 			localDeliveryChan: make(chan any, 100),
+			forceCloseChan:    make(chan primitive.ObjectID),
 			sendMsgChan:       sendchan,
 			sessionConsumer:   consumer,
 		},
-	}, nil
+	}
+	consumer.RegisterOnSessionInvalidated(ws.onSessionInvalidated)
+	return ws, nil
 }
 
 func (ws *WebsocketHandler) Start(ctx context.Context) {
@@ -206,6 +210,10 @@ func (ws *WebsocketHandler) LeaveRoom(room string, accid primitive.ObjectID) {
 	}
 }
 
+func (ws *WebsocketHandler) onSessionInvalidated(accid primitive.ObjectID) {
+	ws.forceCloseChan <- accid
+}
+
 func (ws *WebsocketHandler) mainLoop(ctx context.Context) {
 	defer func() {
 		ws.connWaitGroup.Done()
@@ -358,6 +366,9 @@ func (ws *WebsocketHandler) mainLoop(ctx context.Context) {
 	}
 
 	// 유저에게서 온 메세지, 소켓 연결/해체 처리
+	unauthdata := []byte{0x03, 0xec}
+	unauthdata = append(unauthdata, []byte("unauthorized")...)
+
 	for {
 		buffer := bytes.NewBuffer(make([]byte, 0, 1024))
 		buffer.Reset()
@@ -442,6 +453,11 @@ func (ws *WebsocketHandler) mainLoop(ctx context.Context) {
 				logger.Println("ClientConnected :", c.sender.Alias)
 				go ws.ClientConnected(c)
 			}
+
+		case accid := <-ws.forceCloseChan:
+			if conn := entireConns[accid.Hex()]; conn != nil {
+				conn.WriteControl(websocket.CloseMessage, unauthdata, time.Time{})
+			}
 		}
 	}
 }
@@ -520,6 +536,18 @@ func (ws *WebsocketHandler) upgrade_nosession(w http.ResponseWriter, r *http.Req
 	raw := (*[12]byte)(temp)
 	accid := primitive.ObjectID(*raw)
 
+	sk := r.Header.Get("AS-X-SESSION")
+	authinfo, err := ws.sessionConsumer.Query(sk)
+	if err != nil {
+		w.WriteHeader(http.StatusBadRequest)
+		return
+	}
+
+	if authinfo.Account != accid {
+		w.WriteHeader(http.StatusUnauthorized)
+		return
+	}
+
 	var upgrader = websocket.Upgrader{} // use default options
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {
@@ -558,6 +586,11 @@ func (ws *WebsocketHandler) upgrade(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
+	if authinfo.Account.IsZero() {
+		w.WriteHeader(http.StatusUnauthorized)
+		return
+	}
+
 	var upgrader = websocket.Upgrader{} // use default options
 	conn, err := upgrader.Upgrade(w, r, nil)
 	if err != nil {