ayo/gocommon
8
0
Fork 1
Code Issues Pull Requests Projects Releases Wiki Activity

wshandler 추가

Browse Source
This commit is contained in:
mountain
2023-05-24 15:59:58 +09:00
parent e95efa06a6
commit abddf26379
3 changed files with 620 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

1
go.mod
View File

@ -7,6 +7,7 @@ replace repositories.action2quare.com/ayo/gocommon => ./
require (
github.com/go-redis/redis/v8 v8.11.5
github.com/golang-jwt/jwt v3.2.2+incompatible
github.com/gorilla/websocket v1.5.0
github.com/pires/go-proxyproto v0.7.0
go.mongodb.org/mongo-driver v1.11.6
golang.org/x/crypto v0.0.0-20220622213112-05595931fe9d

2
go.sum
View File

@ -14,6 +14,8 @@ github.com/golang/snappy v0.0.1 h1:Qgr9rKW7uDUkrbSmQeiDsGa8SjGyCOGtuasMWwvp2P4=
github.com/golang/snappy v0.0.1/go.mod h1:/XxbfmMg8lxefKM7IXC3fBNl/7bRcc72aCRzEWrmP2Q=
github.com/google/go-cmp v0.5.2 h1:X2ev0eStA3AbceY54o37/0PQ/UWqKEiiO2dKL5OPaFM=
github.com/google/go-cmp v0.5.2/go.mod h1:v8dTdLbMG2kIc/vJvl+f65V22dbkXbowE6jgT/gNBxE=
github.com/gorilla/websocket v1.5.0 h1:PPwGk2jz7EePpoHN/+ClbZu8SPxiqlu12wZP/3sWmnc=
github.com/gorilla/websocket v1.5.0/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/klauspost/compress v1.13.6 h1:P76CopJELS0TiO2mebmnzgWaajssP/EszplttgQxcgc=
github.com/klauspost/compress v1.13.6/go.mod h1:/3/Vjq9QcHkK5uEr5lBEmyoZ1iFhe47etQ6QUkpK6sk=
github.com/kr/pretty v0.1.0 h1:L/CwN0zerZDmRFUapSPitk6f+Q3+0za1rQkzVuMiMFI=

617
wshandler/wshandler.go Normal file
View File

@ -0,0 +1,617 @@
package wshandler
import (
"context"
"encoding/hex"
"encoding/json"
"fmt"
"io"
"net/http"
"os"
"strings"
"sync"
common "repositories.action2quare.com/ayo/gocommon"
"repositories.action2quare.com/ayo/gocommon/logger"
"github.com/go-redis/redis/v8"
"github.com/gorilla/websocket"
"go.mongodb.org/mongo-driver/bson"
"go.mongodb.org/mongo-driver/bson/primitive"
)
const (
connStateCachePrefix = "conn_state_"
connStateScript = `
local hosts = redis.call('keys',KEYS[1])
for index, key in ipairs(hosts) do
local ok = redis.call('hexists', key, KEYS[2])
if ok == 1 then
return redis.call('hget', key, KEYS[2])
end
end
return ""
`
)
var ConnStateCacheKey = func() string {
hn, _ := os.Hostname()
return connStateCachePrefix + hn
}()
type Richconn struct {
*websocket.Conn
lock sync.Mutex
alias primitive.ObjectID
tags []string
onClose map[string]func()
}
func (rc *Richconn) AddTag(name, val string) {
rc.lock.Lock()
defer rc.lock.Unlock()
prefix := name + "="
for i, tag := range rc.tags {
if strings.HasPrefix(tag, prefix) {
rc.tags[i] = prefix + val
return
}
}
rc.tags = append(rc.tags, prefix+val)
}
func (rc *Richconn) GetTag(name string) string {
rc.lock.Lock()
defer rc.lock.Unlock()
prefix := name + "="
for _, tag := range rc.tags {
if strings.HasPrefix(tag, prefix) {
return tag[len(prefix):]
}
}
return ""
}
func (rc *Richconn) RemoveTag(name string, val string) {
rc.lock.Lock()
defer rc.lock.Unlock()
whole := fmt.Sprintf("%s=%s", name, val)
for i, tag := range rc.tags {
if tag == whole {
if i == 0 && len(rc.tags) == 1 {
rc.tags = nil
} else {
lastidx := len(rc.tags) - 1
if i < lastidx {
rc.tags[i] = rc.tags[lastidx]
}
rc.tags = rc.tags[:lastidx]
}
return
}
}
}
func (rc *Richconn) RegistOnCloseFunc(name string, f func()) {
rc.lock.Lock()
defer rc.lock.Unlock()
if rc.onClose == nil {
f()
return
}
rc.onClose[name] = f
}
func (rc *Richconn) HasOnCloseFunc(name string) bool {
rc.lock.Lock()
defer rc.lock.Unlock()
if rc.onClose == nil {
return false
}
_, ok := rc.onClose[name]
return ok
}
func (rc *Richconn) UnregistOnCloseFunc(name string) (out func()) {
rc.lock.Lock()
defer rc.lock.Unlock()
if rc.onClose == nil {
return
}
out = rc.onClose[name]
delete(rc.onClose, name)
return
}
func (rc *Richconn) WriteBytes(data []byte) error {
rc.lock.Lock()
defer rc.lock.Unlock()
return rc.WriteMessage(websocket.TextMessage, data)
}
type DeliveryMessage struct {
Alias primitive.ObjectID
Body []byte
Command string
Conn *Richconn
}
func (dm *DeliveryMessage) Parse(out any) error {
return json.Unmarshal(dm.Body, out)
}
func (dm *DeliveryMessage) MarshalBinary() (data []byte, err error) {
return append(dm.Alias[:], dm.Body...), nil
}
func (dm *DeliveryMessage) UnmarshalBinary(data []byte) error {
copy(dm.Alias[:], data[:12])
dm.Body = data[12:]
return nil
}
type tagconn struct {
rc *Richconn
state string
}
type tagconnsmap = map[primitive.ObjectID]*tagconn
type tagconns struct {
sync.Mutex
tagconnsmap
}
type subhandler struct {
sync.Mutex
authCache *common.AuthCollection
conns map[primitive.ObjectID]*Richconn
aliases map[primitive.ObjectID]primitive.ObjectID
tags map[primitive.ObjectID]*tagconns
deliveryChan chan DeliveryMessage
url string
redisSync *redis.Client
}
// WebsocketHandler :
type WebsocketHandler struct {
authCaches map[string]*subhandler
RedisSync *redis.Client
}
type wsConfig struct {
SyncPipeline string `json:"ws_sync_pipeline"`
}
func NewWebsocketHandler(authglobal common.AuthCollectionGlobal) (wsh *WebsocketHandler) {
authCaches := make(map[string]*subhandler)
for _, region := range authglobal.Regions() {
sh := &subhandler{
authCache: authglobal.Get(region),
conns: make(map[primitive.ObjectID]*Richconn),
aliases: make(map[primitive.ObjectID]primitive.ObjectID),
tags: make(map[primitive.ObjectID]*tagconns),
deliveryChan: make(chan DeliveryMessage, 1000),
}
authCaches[region] = sh
}
var config wsConfig
common.LoadConfig(&config)
redisSync, err := common.NewRedisClient(config.SyncPipeline, 0)
if err != nil {
panic(err)
}
return &WebsocketHandler{
authCaches: authCaches,
RedisSync: redisSync,
}
}
func (ws *WebsocketHandler) Destructor() {
if ws.RedisSync != nil {
ws.RedisSync.Del(context.Background(), ConnStateCacheKey)
}
}
func (ws *WebsocketHandler) DeliveryChannel(region string) <-chan DeliveryMessage {
return ws.authCaches[region].deliveryChan
}
func (ws *WebsocketHandler) Conn(region string, alias primitive.ObjectID) *Richconn {
if sh := ws.authCaches[region]; sh != nil {
return sh.conns[alias]
}
return nil
}
func (ws *WebsocketHandler) JoinTag(region string, tag primitive.ObjectID, tid primitive.ObjectID, rc *Richconn, hint string) error {
if sh := ws.authCaches[region]; sh != nil {
sh.joinTag(tag, tid, rc, hint)
}
return nil
}
func (ws *WebsocketHandler) LeaveTag(region string, tag primitive.ObjectID, tid primitive.ObjectID) error {
if sh := ws.authCaches[region]; sh != nil {
sh.leaveTag(tag, tid)
}
return nil
}
func (ws *WebsocketHandler) SetStateInTag(region string, tag primitive.ObjectID, tid primitive.ObjectID, state string, hint string) error {
if sh := ws.authCaches[region]; sh != nil {
sh.setStateInTag(tag, tid, state, hint)
}
return nil
}
func (ws *WebsocketHandler) BroadcastRaw(region string, tag primitive.ObjectID, raw []byte) {
if sh := ws.authCaches[region]; sh != nil {
if cs := sh.cloneTag(tag); len(cs) > 0 {
go func(raw []byte) {
for _, c := range cs {
if c != nil {
c.WriteBytes(raw)
}
}
}(raw)
}
}
}
func (ws *WebsocketHandler) Broadcast(region string, tag primitive.ObjectID, doc bson.M) {
raw, _ := json.Marshal(doc)
ws.BroadcastRaw(region, tag, raw)
}
var onlineQueryScriptHash string
func (ws *WebsocketHandler) RegisterHandlers(ctx context.Context, serveMux *http.ServeMux, prefix string) error {
ws.RedisSync.Del(context.Background(), ConnStateCacheKey)
scriptHash, err := ws.RedisSync.ScriptLoad(context.Background(), connStateScript).Result()
if err != nil {
return err
}
onlineQueryScriptHash = scriptHash
for region, sh := range ws.authCaches {
if region == "default" {
region = ""
}
sh.url = common.MakeHttpHandlerPattern(prefix, region, "ws")
sh.redisSync = ws.RedisSync
if *common.NoSessionFlag {
serveMux.HandleFunc(sh.url, sh.upgrade_nosession)
} else {
serveMux.HandleFunc(sh.url, sh.upgrade)
}
}
return nil
}
func (sh *subhandler) cloneTag(tag primitive.ObjectID) (out []*Richconn) {
sh.Lock()
cs := sh.tags[tag]
sh.Unlock()
if cs == nil {
return nil
}
cs.Lock()
defer cs.Unlock()
out = make([]*Richconn, 0, len(cs.tagconnsmap))
for _, c := range cs.tagconnsmap {
out = append(out, c.rc)
}
return
}
func (sh *subhandler) joinTag(tag primitive.ObjectID, tid primitive.ObjectID, rc *Richconn, hint string) {
sh.Lock()
cs := sh.tags[tag]
if cs == nil {
cs = &tagconns{
tagconnsmap: make(map[primitive.ObjectID]*tagconn),
}
}
sh.Unlock()
cs.Lock()
states := make([]bson.M, 0, len(cs.tagconnsmap))
for tid, conn := range cs.tagconnsmap {
states = append(states, bson.M{
"_id": tid,
"_hint": hint,
"state": conn.state,
})
}
cs.tagconnsmap[tid] = &tagconn{rc: rc}
cs.Unlock()
sh.Lock()
sh.tags[tag] = cs
sh.Unlock()
if len(states) > 0 {
s, _ := json.Marshal(states)
rc.WriteBytes(s)
}
}
func (sh *subhandler) leaveTag(tag primitive.ObjectID, tid primitive.ObjectID) {
sh.Lock()
defer sh.Unlock()
cs := sh.tags[tag]
if cs == nil {
return
}
delete(cs.tagconnsmap, tid)
if len(cs.tagconnsmap) == 0 {
delete(sh.tags, tag)
} else {
sh.tags[tag] = cs
}
}
func (sh *subhandler) setStateInTag(tag primitive.ObjectID, tid primitive.ObjectID, state string, hint string) {
sh.Lock()
cs := sh.tags[tag]
sh.Unlock()
if cs == nil {
return
}
cs.Lock()
defer cs.Unlock()
if tagconn := cs.tagconnsmap[tid]; tagconn != nil {
tagconn.state = state
var clone []*Richconn
for _, c := range cs.tagconnsmap {
clone = append(clone, c.rc)
}
raw, _ := json.Marshal(map[string]any{
"_id": tid,
"_hint": hint,
"state": state,
})
go func(raw []byte) {
for _, c := range clone {
c.WriteBytes(raw)
}
}(raw)
}
}
func (wsh *WebsocketHandler) GetState(alias primitive.ObjectID) (string, error) {
state, err := wsh.RedisSync.EvalSha(context.Background(), onlineQueryScriptHash, []string{
connStateCachePrefix + "*", alias.Hex(),
}).Result()
if err != nil {
return "", err
}
return state.(string), nil
}
func (wsh *WebsocketHandler) IsOnline(alias primitive.ObjectID) (bool, error) {
state, err := wsh.GetState(alias)
if err != nil {
logger.Error("IsOnline failed. err :", err)
return false, err
}
return len(state) > 0, nil
}
func (sh *subhandler) closeConn(accid primitive.ObjectID) {
sh.Lock()
defer sh.Unlock()
if alias, ok := sh.aliases[accid]; ok {
if old := sh.conns[alias]; old != nil {
old.Close()
}
}
}
func (sh *subhandler) addConn(conn *Richconn, accid primitive.ObjectID) {
sh.Lock()
defer sh.Unlock()
sh.conns[conn.alias] = conn
sh.aliases[accid] = conn.alias
}
func upgrade_core(sh *subhandler, conn *websocket.Conn, initState string, accid primitive.ObjectID, alias primitive.ObjectID) {
sh.closeConn(accid)
newconn := sh.makeRichConn(alias, conn)
sh.addConn(newconn, accid)
sh.redisSync.HSet(context.Background(), ConnStateCacheKey, alias.Hex(), initState).Result()
go func(c *Richconn, accid primitive.ObjectID, deliveryChan chan<- DeliveryMessage) {
for {
mt, p, err := c.ReadMessage()
if err != nil {
c.Close()
break
}
switch mt {
case websocket.BinaryMessage:
msg := DeliveryMessage{
Alias: c.alias,
Body: p,
Conn: c,
}
deliveryChan <- msg
case websocket.TextMessage:
msg := string(p)
opcodes := strings.Split(msg, ";")
for _, opcode := range opcodes {
if strings.HasPrefix(opcode, "ps:") {
sh.redisSync.HSet(context.Background(), ConnStateCacheKey, alias.Hex(), opcode[3:]).Result()
} else if strings.HasPrefix(opcode, "cmd:") {
cmd := opcode[4:]
msg := DeliveryMessage{
Alias: c.alias,
Command: cmd,
Conn: c,
}
deliveryChan <- msg
}
}
}
}
sh.redisSync.HDel(context.Background(), ConnStateCacheKey, c.alias.Hex()).Result()
sh.Lock()
delete(sh.conns, c.alias)
delete(sh.aliases, accid)
sh.Unlock()
var funcs []func()
c.lock.Lock()
for _, f := range c.onClose {
funcs = append(funcs, f)
}
c.onClose = nil
c.lock.Unlock()
for _, f := range funcs {
f()
}
}(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
}
alias := accid
if v := r.Header.Get("AS-X-ALIAS"); len(v) > 0 {
alias = common.ParseObjectID(v)
}
initState := r.Header.Get("As-X-Tavern-InitialState")
if len(initState) == 0 {
initState = "online"
}
upgrade_core(sh, conn, initState, 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 := strings.Split(r.Header.Get("Authorization"), " ")
if len(auth) != 2 {
//TODO : 클라이언트는 BadRequest를 받으면 로그인 화면으로 돌아가야 한다.
w.WriteHeader(http.StatusBadRequest)
return
}
authtoken := auth[1]
accid, success := sh.authCache.IsValid(sk, authtoken)
if !success {
w.WriteHeader(http.StatusUnauthorized)
return
}
var upgrader = websocket.Upgrader{} // use default options
conn, err := upgrader.Upgrade(w, r, nil)
if err != nil {
w.WriteHeader(http.StatusInternalServerError)
return
}
alias := accid
if v := r.Header.Get("AS-X-ALIAS"); len(v) > 0 {
alias = common.ParseObjectID(v)
}
initState := r.Header.Get("As-X-Tavern-InitialState")
if len(initState) == 0 {
initState = "online"
}
upgrade_core(sh, conn, initState, accid, alias)
}
func (sh *subhandler) makeRichConn(alias primitive.ObjectID, conn *websocket.Conn) *Richconn {
rc := Richconn{
Conn: conn,
alias: alias,
onClose: make(map[string]func()),
}
return &rc
}