package servlet import ( "context" "encoding/base64" "encoding/json" "fmt" "sync" "github.com/valyala/fasthttp" "git.loafle.net/commons/logging-go" "git.loafle.net/commons/server-go" "git.loafle.net/commons/server-go/socket" ocnc "git.loafle.net/overflow/commons-go/noauthprobe/constants" ocnm "git.loafle.net/overflow/commons-go/noauthprobe/model" og "git.loafle.net/overflow/gateway" "git.loafle.net/overflow/gateway/external/grpc" ogs "git.loafle.net/overflow/gateway/subscribe" ogrs "git.loafle.net/overflow/gateway_rpc/servlet" "git.loafle.net/overflow/member_gateway_rpc/subscribe" ) type AuthServlet interface { ogrs.RPCServlet } type AuthServlets struct { ogrs.RPCServlets connections sync.Map } func (s *AuthServlets) Init(serverCtx server.ServerCtx) error { if err := s.RPCServlets.Init(serverCtx); nil != err { return err } return nil } func (s *AuthServlets) OnStart(serverCtx server.ServerCtx) error { if err := s.RPCServlets.OnStart(serverCtx); nil != err { return err } subscribeChan, err := subscribe.Subscriber.Subscribe(ocnc.HTTPEntry_Auth) if nil != err { return err } go s.handleSubscribe(serverCtx, subscribeChan) return nil } func (s *AuthServlets) OnStop(serverCtx server.ServerCtx) { if err := subscribe.Subscriber.Unsubscribe(ocnc.HTTPEntry_Auth); nil != err { logging.Logger().Warn(err) } s.RPCServlets.OnStop(serverCtx) } func (s *AuthServlets) Destroy(serverCtx server.ServerCtx) { s.RPCServlets.Destroy(serverCtx) } func (s *AuthServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) { bMethod := ctx.Request.Header.Peek(ocnc.HTTPRequestHeaderKey_NoAuthProbe_Method) if nil == bMethod { return nil, fmt.Errorf("Unexpected noauth probe method: %v", bMethod) } method := string(bMethod) switch method { case ocnc.HTTPRequestHeaderValue_NoAuthProbe_Method_Regist: bInfo := ctx.Request.Header.Peek(ocnc.HTTPRequestHeaderKey_NoAuthProbe_Info) if nil == bInfo { return nil, fmt.Errorf("Unexpected HTTPRequestHeaderKey NoAuthProbe Info") } rJSON := make([]byte, base64.StdEncoding.DecodedLen(len(bInfo))) _, err := base64.StdEncoding.Decode(rJSON, bInfo) if nil != err { return nil, fmt.Errorf("Base64 Encoding Error: %s", err.Error()) } grpcCTX := context.Background() r, err := grpc.Exec(grpcCTX, "NoAuthProbeService.regist", string(rJSON)) if nil != err { return nil, fmt.Errorf("grpc call Error: %s", err.Error()) } nap := &ocnm.NoAuthProbe{} err = json.Unmarshal([]byte(r), nap) if nil != err { return nil, fmt.Errorf("grpc result unMarshal Error: %s", err.Error()) } extHeader := &fasthttp.ResponseHeader{} extHeader.Add(ocnc.HTTPResponseHeaderKey_NoAuthProbe_SetTempProbeKey, nap.TempProbeKey) return extHeader, nil case ocnc.HTTPRequestHeaderValue_NoAuthProbe_Method_Connect: bTempProbeKey := ctx.Request.Header.Peek(ocnc.HTTPRequestHeaderKey_NoAuthProbe_TempProbeKey) if nil == bTempProbeKey { return nil, fmt.Errorf("Noauth probe temp key is not an existing key") } tempProbeKey := string(bTempProbeKey) grpcCTX := context.Background() _, err := grpc.Exec(grpcCTX, "NoAuthProbeService.readByTempKey", tempProbeKey) if nil != err { return nil, fmt.Errorf("grpc result error: %s", err.Error()) } servletCtx.SetAttribute(og.SessionIDKey, tempProbeKey) servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE) servletCtx.SetAttribute(og.SessionTargetIDKey, tempProbeKey) return nil, nil default: return nil, fmt.Errorf("Unexpected noauth probe httpRequestHeaderValue: %v", method) } } func (s *AuthServlets) OnConnect(servletCtx server.ServletCtx, conn socket.Conn) { s.RPCServlets.OnConnect(servletCtx, conn) sessionID := servletCtx.GetAttribute(og.SessionIDKey) targetID := servletCtx.GetAttribute(og.SessionTargetIDKey) if nil != sessionID && nil != targetID { s.connections.Store(sessionID.(string), retainConnection(targetID.(string), servletCtx)) } } func (s *AuthServlets) OnDisconnect(servletCtx server.ServletCtx) { s.RPCServlets.OnDisconnect(servletCtx) sessionID := servletCtx.GetAttribute(og.SessionIDKey) if nil != sessionID { s.connections.Delete(sessionID.(string)) } } func (s *AuthServlets) handleSubscribe(serverCtx server.ServerCtx, subscribeChan <-chan *ogs.Message) { for { select { case msg, ok := <-subscribeChan: if !ok { return } switch msg.TargetType { case ogs.PROBE: for _, targetID := range msg.Targets { _connections := s.getProbeConnections(targetID) if nil == _connections || 0 == len(_connections) { break } for _, _connection := range _connections { _writeChan := _connection.servletCtx.GetAttribute(og.SessionWriteChanKey) if nil != _writeChan { writeChan := _writeChan.(chan<- []byte) writeChan <- msg.Message } } } } } } } func (s *AuthServlets) getProbeConnections(targetID string) []*connection { var connections []*connection s.connections.Range(func(k, v interface{}) bool { _connection := v.(*connection) if _connection.targetID == targetID { connections = append(connections, _connection) } return true }) return connections } type connection struct { targetID string servletCtx server.ServletCtx } var connectionPool sync.Pool func retainConnection(targetID string, servletCtx server.ServletCtx) *connection { return nil } func releaseConnection(_connection *connection) { _connection.targetID = "" _connection.servletCtx = nil connectionPool.Put(_connection) }