ing

servlet/rpc-servlet.go

@@ -2,11 +2,13 @@ package servlet
 
 import (
 	"context"
+	"sync"
 
 	logging "git.loafle.net/commons/logging-go"
 	crp "git.loafle.net/commons/rpc-go/protocol"
 	crpj "git.loafle.net/commons/rpc-go/protocol/json"
 	"git.loafle.net/commons/server-go"
+	css "git.loafle.net/commons/server-go/socket"
 	cssw "git.loafle.net/commons/server-go/socket/web"
 	og "git.loafle.net/overflow/gateway"
 	ogeg "git.loafle.net/overflow/gateway/external/grpc"
@@ -20,21 +22,39 @@ type RPCServlet interface {
 
 type RPCServlets struct {
 	cssw.Servlets
-}
 
-func init() {
-	// member RSA file read
+	UseSession bool
+
+	sessions sync.Map
 }
 
 func (s *RPCServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) {
-	// member auth token extraction
-	// auth token parse
-	// auth token valid
-	// servletCtx set member
-
 	return nil, nil
 }
 
+func (s *RPCServlets) OnConnect(servletCtx server.ServletCtx, conn css.Conn) {
+	s.Servlets.OnConnect(servletCtx, conn)
+
+	if s.UseSession {
+		sessionID := servletCtx.GetAttribute(og.SessionIDKey)
+		targetID := servletCtx.GetAttribute(og.SessionTargetIDKey)
+		if nil != sessionID && nil != targetID {
+			s.sessions.Store(sessionID.(string), RetainSession(targetID.(string), servletCtx))
+		}
+	}
+}
+
+func (s *RPCServlets) OnDisconnect(servletCtx server.ServletCtx) {
+	s.Servlets.OnDisconnect(servletCtx)
+
+	if s.UseSession {
+		sessionID := servletCtx.GetAttribute(og.SessionIDKey)
+		if nil != sessionID {
+			s.sessions.Delete(sessionID.(string))
+		}
+	}
+}
+
 func (s *RPCServlets) Handle(servletCtx server.ServletCtx,
 	stopChan <-chan struct{}, doneChan chan<- struct{},
 	readChan <-chan []byte, writeChan chan<- []byte) {
@@ -125,3 +145,42 @@ func (s *RPCServlets) writeError(src crp.ServerRequestCodec, writeChan chan<- []
 	}
 	writeChan <- buf
 }
+
+func (s *RPCServlets) GetSessions(sessionIDs []string) []*Session {
+	var sessions []*Session
+
+	if nil == sessionIDs || 0 == len(sessionIDs) {
+		return sessions
+	}
+
+	for _, sessionID := range sessionIDs {
+		session, ok := s.sessions.Load(sessionID)
+		if ok {
+			sessions = append(sessions, session.(*Session))
+		}
+	}
+
+	return sessions
+}
+
+func (s *RPCServlets) GetSessionsByTargetIDs(targetIDs []string) []*Session {
+	var sessions []*Session
+	if nil == targetIDs || 0 == len(targetIDs) {
+		return sessions
+	}
+
+	s.sessions.Range(func(k, v interface{}) bool {
+		session := v.(*Session)
+
+		for _, targetID := range targetIDs {
+			if session.TargetID == targetID {
+				sessions = append(sessions, session)
+				break
+			}
+		}
+
+		return true
+	})
+
+	return sessions
+}