573 lines
14 KiB
Go
573 lines
14 KiB
Go
package wshandler
|
|
|
|
import (
|
|
"context"
|
|
"encoding/hex"
|
|
"encoding/json"
|
|
"fmt"
|
|
"io"
|
|
"net/http"
|
|
"strings"
|
|
"sync"
|
|
|
|
"go.mongodb.org/mongo-driver/bson/primitive"
|
|
"repositories.action2quare.com/ayo/gocommon"
|
|
"repositories.action2quare.com/ayo/gocommon/flagx"
|
|
"repositories.action2quare.com/ayo/gocommon/logger"
|
|
|
|
"github.com/go-redis/redis/v8"
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
var noSessionFlag = flagx.Bool("nosession", false, "nosession=[true|false]")
|
|
|
|
type wsconn struct {
|
|
*websocket.Conn
|
|
sender *Sender
|
|
joinedRooms []*room
|
|
}
|
|
|
|
func (conn *wsconn) popRoom(r *room) int {
|
|
for i, jr := range conn.joinedRooms {
|
|
if jr == r {
|
|
conn.joinedRooms = append(conn.joinedRooms[:i], conn.joinedRooms[i+1:]...)
|
|
break
|
|
}
|
|
}
|
|
return len(conn.joinedRooms)
|
|
}
|
|
|
|
func (conn *wsconn) pushRoom(r *room) {
|
|
conn.joinedRooms = append(conn.joinedRooms, r)
|
|
}
|
|
|
|
type UpstreamMessage struct {
|
|
Alias string
|
|
Accid primitive.ObjectID
|
|
Target string
|
|
Body any
|
|
Tag []string
|
|
}
|
|
|
|
type DownstreamMessage struct {
|
|
Alias string `json:",omitempty"`
|
|
Body any `json:",omitempty"`
|
|
Tag []string `json:",omitempty"`
|
|
}
|
|
|
|
type commandType string
|
|
|
|
const (
|
|
commandType_JoinRoom = commandType("join_room")
|
|
commandType_LeaveRoom = commandType("leave_room")
|
|
)
|
|
|
|
type commandMessage struct {
|
|
Cmd commandType
|
|
Args []any
|
|
}
|
|
|
|
type WebSocketMessageType int
|
|
|
|
const (
|
|
TextMessage = WebSocketMessageType(websocket.TextMessage)
|
|
BinaryMessage = WebSocketMessageType(websocket.BinaryMessage)
|
|
CloseMessage = WebSocketMessageType(websocket.CloseMessage)
|
|
PingMessage = WebSocketMessageType(websocket.PingMessage)
|
|
PongMessage = WebSocketMessageType(websocket.PongMessage)
|
|
Connected = WebSocketMessageType(100)
|
|
Disconnected = WebSocketMessageType(101)
|
|
)
|
|
|
|
type Sender struct {
|
|
Region string
|
|
Accid primitive.ObjectID
|
|
Alias string
|
|
}
|
|
|
|
type WebSocketMessageReceiver func(sender *Sender, messageType WebSocketMessageType, body io.Reader)
|
|
|
|
type subhandler struct {
|
|
redisMsgChanName string
|
|
redisCmdChanName string
|
|
redisSync *redis.Client
|
|
connInOutChan chan *wsconn
|
|
deliveryChan chan any
|
|
localDeliveryChan chan any
|
|
callReceiver WebSocketMessageReceiver
|
|
connWaitGroup sync.WaitGroup
|
|
region string
|
|
receiverChain []WebSocketMessageReceiver
|
|
}
|
|
|
|
// WebsocketHandler :
|
|
type WebsocketHandler struct {
|
|
subhandlers map[string]*subhandler
|
|
}
|
|
|
|
type wsConfig struct {
|
|
gocommon.RegionStorageConfig
|
|
Maingate string `json:"maingate_service_url"`
|
|
}
|
|
|
|
var config wsConfig
|
|
|
|
func init() {
|
|
gocommon.LoadConfig(&config)
|
|
}
|
|
|
|
func NewWebsocketHandler() (*WebsocketHandler, error) {
|
|
// decoder := func(r io.Reader) *T {
|
|
// if r == nil {
|
|
// // 접속이 끊겼을 때.
|
|
// return nil
|
|
// }
|
|
// var m T
|
|
// dec := json.NewDecoder(r)
|
|
// if err := dec.Decode(&m); err != nil {
|
|
// logger.Println(err)
|
|
// }
|
|
|
|
// // decoding 실패하더라도 빈 *T를 내보냄
|
|
// return &m
|
|
// }
|
|
|
|
subhandlers := make(map[string]*subhandler)
|
|
for region, cfg := range config.RegionStorage {
|
|
redisSync, err := gocommon.NewRedisClient(cfg.Redis["wshandler"])
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
sh := &subhandler{
|
|
redisMsgChanName: fmt.Sprintf("_wsh_msg_%s", region),
|
|
redisCmdChanName: fmt.Sprintf("_wsh_cmd_%s", region),
|
|
redisSync: redisSync,
|
|
connInOutChan: make(chan *wsconn),
|
|
deliveryChan: make(chan any, 1000),
|
|
localDeliveryChan: make(chan any, 100),
|
|
region: region,
|
|
}
|
|
|
|
subhandlers[region] = sh
|
|
}
|
|
|
|
return &WebsocketHandler{
|
|
subhandlers: subhandlers,
|
|
}, nil
|
|
}
|
|
|
|
func (ws *WebsocketHandler) RegisterReceiver(region string, receiver WebSocketMessageReceiver) {
|
|
if sh := ws.subhandlers[region]; sh != nil {
|
|
sh.receiverChain = append(sh.receiverChain, receiver)
|
|
}
|
|
}
|
|
|
|
func (ws *WebsocketHandler) Start(ctx context.Context) {
|
|
for _, sh := range ws.subhandlers {
|
|
chain := sh.receiverChain
|
|
if len(chain) == 0 {
|
|
sh.callReceiver = func(sender *Sender, messageType WebSocketMessageType, body io.Reader) {}
|
|
} else if len(chain) == 1 {
|
|
sh.callReceiver = chain[0]
|
|
} else {
|
|
sh.callReceiver = func(sender *Sender, messageType WebSocketMessageType, body io.Reader) {
|
|
for _, r := range chain {
|
|
r(sender, messageType, body)
|
|
}
|
|
}
|
|
}
|
|
|
|
go sh.mainLoop(ctx)
|
|
}
|
|
}
|
|
|
|
func (ws *WebsocketHandler) Cleanup() {
|
|
for _, sh := range ws.subhandlers {
|
|
sh.connWaitGroup.Wait()
|
|
}
|
|
}
|
|
|
|
func (ws *WebsocketHandler) RegisterHandlers(serveMux *http.ServeMux, prefix string) error {
|
|
for region, sh := range ws.subhandlers {
|
|
if region == "default" {
|
|
region = ""
|
|
}
|
|
url := gocommon.MakeHttpHandlerPattern(prefix, region, "ws")
|
|
if *noSessionFlag {
|
|
serveMux.HandleFunc(url, sh.upgrade_nosession)
|
|
} else {
|
|
serveMux.HandleFunc(url, sh.upgrade)
|
|
}
|
|
}
|
|
|
|
return nil
|
|
}
|
|
|
|
func (ws *WebsocketHandler) GetState(region string, accid primitive.ObjectID) string {
|
|
if sh := ws.subhandlers[region]; sh != nil {
|
|
state, _ := sh.redisSync.Get(context.Background(), accid.Hex()).Result()
|
|
return state
|
|
}
|
|
return ""
|
|
}
|
|
|
|
func (ws *WebsocketHandler) SetState(region string, accid primitive.ObjectID, state string) {
|
|
if sh := ws.subhandlers[region]; sh != nil {
|
|
sh.redisSync.SetArgs(context.Background(), accid.Hex(), state, redis.SetArgs{Mode: "XX"}).Result()
|
|
}
|
|
}
|
|
|
|
func (ws *WebsocketHandler) SendUpstreamMessage(region string, msg *UpstreamMessage) {
|
|
sh := ws.subhandlers[region]
|
|
if sh != nil {
|
|
sh.localDeliveryChan <- msg
|
|
}
|
|
}
|
|
|
|
func (ws *WebsocketHandler) EnterRoom(region string, room string, accid primitive.ObjectID) {
|
|
sh := ws.subhandlers[region]
|
|
if sh != nil {
|
|
sh.localDeliveryChan <- &commandMessage{
|
|
Cmd: commandType_JoinRoom,
|
|
Args: []any{room, accid},
|
|
}
|
|
}
|
|
}
|
|
|
|
func (ws *WebsocketHandler) LeaveRoom(region string, room string, accid primitive.ObjectID) {
|
|
sh := ws.subhandlers[region]
|
|
if sh != nil {
|
|
sh.localDeliveryChan <- &commandMessage{
|
|
Cmd: commandType_LeaveRoom,
|
|
Args: []any{room, accid},
|
|
}
|
|
}
|
|
}
|
|
|
|
func (sh *subhandler) mainLoop(ctx context.Context) {
|
|
defer func() {
|
|
s := recover()
|
|
if s != nil {
|
|
logger.Error(s)
|
|
}
|
|
}()
|
|
|
|
// redis channel에서 유저가 보낸 메시지를 읽는 go rountine
|
|
go func() {
|
|
var pubsub *redis.PubSub
|
|
for {
|
|
if pubsub == nil {
|
|
pubsub = sh.redisSync.Subscribe(ctx, sh.redisMsgChanName, sh.redisCmdChanName)
|
|
}
|
|
|
|
raw, err := pubsub.ReceiveMessage(ctx)
|
|
if err == nil {
|
|
if raw.Channel == sh.redisMsgChanName {
|
|
var msg UpstreamMessage
|
|
if err := json.Unmarshal([]byte(raw.Payload), &msg); err == nil {
|
|
sh.deliveryChan <- &msg
|
|
} else {
|
|
logger.Println("decode UpstreamMessage failed :", err)
|
|
}
|
|
} else if raw.Channel == sh.redisCmdChanName {
|
|
var cmd commandMessage
|
|
if err := json.Unmarshal([]byte(raw.Payload), &cmd); err == nil {
|
|
sh.deliveryChan <- &cmd
|
|
} else {
|
|
logger.Println("decode UpstreamMessage failed :", err)
|
|
}
|
|
}
|
|
} else {
|
|
logger.Println("pubsub.ReceiveMessage failed :", err)
|
|
pubsub.Close()
|
|
pubsub = nil
|
|
|
|
if ctx.Err() != nil {
|
|
break
|
|
}
|
|
}
|
|
}
|
|
}()
|
|
|
|
entireConns := make(map[string]*wsconn)
|
|
rooms := make(map[string]*room)
|
|
roomDestroyChan := make(chan string, 1000)
|
|
findRoom := func(name string, create bool) *room {
|
|
room := rooms[name]
|
|
if room == nil && create {
|
|
room = makeRoom(name, roomDestroyChan)
|
|
rooms[name] = room
|
|
room.start(ctx)
|
|
}
|
|
return room
|
|
}
|
|
|
|
// 유저에게서 온 메세지, 소켓 연결/해체 처리
|
|
for {
|
|
select {
|
|
case destroyedRoom := <-roomDestroyChan:
|
|
delete(rooms, destroyedRoom)
|
|
|
|
case usermsg := <-sh.localDeliveryChan:
|
|
// 로컬에 connection이 있는지 먼저 확인해 보기 위한 채널
|
|
// 없으면 publish한다.
|
|
switch usermsg := usermsg.(type) {
|
|
case *UpstreamMessage:
|
|
target := usermsg.Target
|
|
if target[0] == '@' {
|
|
accid := target[1:]
|
|
conn := entireConns[accid]
|
|
if conn != nil {
|
|
// 이 경우 아니면 publish 해야 함
|
|
ds, _ := json.Marshal(DownstreamMessage{
|
|
Alias: usermsg.Alias,
|
|
Body: usermsg.Body,
|
|
Tag: usermsg.Tag,
|
|
})
|
|
|
|
conn.WriteMessage(websocket.TextMessage, ds)
|
|
break
|
|
}
|
|
}
|
|
if bt, err := json.Marshal(usermsg); err == nil {
|
|
sh.redisSync.Publish(context.Background(), sh.redisMsgChanName, bt).Result()
|
|
}
|
|
|
|
case *commandMessage:
|
|
if usermsg.Cmd == commandType_JoinRoom && len(usermsg.Args) == 2 {
|
|
roomName := usermsg.Args[0].(string)
|
|
accid := usermsg.Args[1].(primitive.ObjectID)
|
|
conn := entireConns[accid.Hex()]
|
|
if conn != nil {
|
|
conn.pushRoom(findRoom(roomName, true).in(conn))
|
|
break
|
|
}
|
|
} else if usermsg.Cmd == commandType_LeaveRoom && len(usermsg.Args) == 2 {
|
|
roomName := usermsg.Args[0].(string)
|
|
accid := usermsg.Args[1].(primitive.ObjectID)
|
|
conn := entireConns[accid.Hex()]
|
|
if conn != nil {
|
|
if room := findRoom(roomName, false); room != nil {
|
|
if conn.popRoom(room.out(conn)) == 0 {
|
|
conn.Close()
|
|
}
|
|
break
|
|
}
|
|
}
|
|
}
|
|
|
|
// 위에서 break 안걸리면 나한테 없으므로 publish를 해야 함. 그러면 다른 호스트가 deliveryChan으로 받는다
|
|
if bt, err := json.Marshal(usermsg); err == nil {
|
|
sh.redisSync.Publish(context.Background(), sh.redisCmdChanName, bt).Result()
|
|
}
|
|
}
|
|
|
|
case usermsg := <-sh.deliveryChan:
|
|
switch usermsg := usermsg.(type) {
|
|
case *UpstreamMessage:
|
|
target := usermsg.Target
|
|
if target[0] == '#' {
|
|
// 룸에 브로드캐스팅
|
|
roomName := target[1:]
|
|
if room := findRoom(roomName, false); room != nil {
|
|
room.broadcast(usermsg)
|
|
}
|
|
} else if target[0] == '@' {
|
|
accid := target[1:]
|
|
conn := entireConns[accid]
|
|
if conn != nil {
|
|
ds, _ := json.Marshal(DownstreamMessage{
|
|
Alias: usermsg.Alias,
|
|
Body: usermsg.Body,
|
|
Tag: usermsg.Tag,
|
|
})
|
|
conn.WriteMessage(websocket.TextMessage, ds)
|
|
}
|
|
}
|
|
|
|
case *commandMessage:
|
|
if usermsg.Cmd == commandType_JoinRoom && len(usermsg.Args) == 2 {
|
|
roomName := usermsg.Args[0].(string)
|
|
accid := usermsg.Args[1].(primitive.ObjectID)
|
|
conn := entireConns[accid.Hex()]
|
|
if conn != nil {
|
|
conn.pushRoom(findRoom(roomName, true).in(conn))
|
|
}
|
|
} else if usermsg.Cmd == commandType_LeaveRoom && len(usermsg.Args) == 2 {
|
|
roomName := usermsg.Args[0].(string)
|
|
accid := usermsg.Args[1].(primitive.ObjectID)
|
|
conn := entireConns[accid.Hex()]
|
|
if conn != nil {
|
|
if room := findRoom(roomName, false); room != nil {
|
|
if conn.popRoom(room.out(conn)) == 0 {
|
|
conn.Close()
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
default:
|
|
logger.Println("usermsg is unknown type")
|
|
}
|
|
|
|
case c := <-sh.connInOutChan:
|
|
if c.Conn == nil {
|
|
delete(entireConns, c.sender.Accid.Hex())
|
|
for _, room := range c.joinedRooms {
|
|
room.out(c)
|
|
}
|
|
c.joinedRooms = nil
|
|
sh.callReceiver(c.sender, Disconnected, nil)
|
|
} else {
|
|
entireConns[c.sender.Accid.Hex()] = c
|
|
sh.callReceiver(c.sender, Connected, nil)
|
|
}
|
|
}
|
|
}
|
|
}
|
|
|
|
func upgrade_core(sh *subhandler, conn *websocket.Conn, accid primitive.ObjectID, alias string) {
|
|
newconn := &wsconn{
|
|
Conn: conn,
|
|
sender: &Sender{
|
|
Region: sh.region,
|
|
Alias: alias,
|
|
Accid: accid,
|
|
},
|
|
}
|
|
sh.connInOutChan <- newconn
|
|
|
|
sh.connWaitGroup.Add(1)
|
|
go func(c *wsconn, accid primitive.ObjectID, deliveryChan chan<- any) {
|
|
sh.redisSync.Set(context.Background(), accid.Hex(), "online", 0)
|
|
for {
|
|
messageType, r, err := c.NextReader()
|
|
if err != nil {
|
|
c.Close()
|
|
break
|
|
}
|
|
|
|
if messageType == websocket.CloseMessage {
|
|
sh.callReceiver(c.sender, CloseMessage, r)
|
|
break
|
|
}
|
|
|
|
if messageType == websocket.TextMessage {
|
|
// 유저가 직접 보낸 메시지
|
|
sh.callReceiver(c.sender, TextMessage, r)
|
|
} else if messageType == websocket.BinaryMessage {
|
|
sh.callReceiver(c.sender, BinaryMessage, r)
|
|
}
|
|
}
|
|
sh.redisSync.Del(context.Background(), accid.Hex())
|
|
sh.connWaitGroup.Done()
|
|
|
|
c.Conn = nil
|
|
sh.connInOutChan <- c
|
|
}(newconn, accid, sh.deliveryChan)
|
|
}
|
|
|
|
func (sh *subhandler) upgrade_nosession(w http.ResponseWriter, r *http.Request) {
|
|
// 클라이언트 접속
|
|
defer func() {
|
|
s := recover()
|
|
if s != nil {
|
|
logger.Error(s)
|
|
}
|
|
io.Copy(io.Discard, r.Body)
|
|
r.Body.Close()
|
|
}()
|
|
|
|
auth := strings.Split(r.Header.Get("Authorization"), " ")
|
|
if len(auth) != 2 {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
if auth[0] != "Editor" {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
return
|
|
}
|
|
temp, err := hex.DecodeString(auth[1])
|
|
if err != nil {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
if len(temp) != len(primitive.NilObjectID) {
|
|
w.WriteHeader(http.StatusBadRequest)
|
|
return
|
|
}
|
|
|
|
raw := (*[12]byte)(temp)
|
|
accid := primitive.ObjectID(*raw)
|
|
|
|
var upgrader = websocket.Upgrader{} // use default options
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
var alias string
|
|
if v := r.Header.Get("AS-X-ALIAS"); len(v) > 0 {
|
|
alias = v
|
|
} else {
|
|
alias = accid.Hex()
|
|
}
|
|
|
|
upgrade_core(sh, conn, accid, alias)
|
|
}
|
|
|
|
func (sh *subhandler) upgrade(w http.ResponseWriter, r *http.Request) {
|
|
// 클라이언트 접속
|
|
defer func() {
|
|
s := recover()
|
|
if s != nil {
|
|
logger.Error(s)
|
|
}
|
|
io.Copy(io.Discard, r.Body)
|
|
r.Body.Close()
|
|
}()
|
|
|
|
sk := r.Header.Get("AS-X-SESSION")
|
|
auth := r.Header.Get("Authorization")
|
|
|
|
req, _ := http.NewRequest("GET", fmt.Sprintf("%s/query?sk=%s", config.Maingate, sk), nil)
|
|
req.Header.Add("Authorization", auth)
|
|
|
|
client := http.Client{}
|
|
resp, err := client.Do(req)
|
|
if err != nil {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
logger.Error("authorize query failed :", err)
|
|
return
|
|
}
|
|
defer resp.Body.Close()
|
|
|
|
var authinfo gocommon.Authinfo
|
|
dec := json.NewDecoder(resp.Body)
|
|
if err = dec.Decode(&authinfo); err != nil {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
logger.Error("authorize query failed :", err)
|
|
return
|
|
}
|
|
|
|
var upgrader = websocket.Upgrader{} // use default options
|
|
conn, err := upgrader.Upgrade(w, r, nil)
|
|
if err != nil {
|
|
w.WriteHeader(http.StatusInternalServerError)
|
|
return
|
|
}
|
|
|
|
var alias string
|
|
if v := r.Header.Get("AS-X-ALIAS"); len(v) > 0 {
|
|
alias = v
|
|
} else {
|
|
alias = authinfo.Accid.Hex()
|
|
}
|
|
|
|
upgrade_core(sh, conn, authinfo.Accid, alias)
|
|
}
|