package servlet import ( "context" "encoding/json" "fmt" "github.com/valyala/fasthttp" "git.loafle.net/commons/logging-go" "git.loafle.net/commons/server-go" "git.loafle.net/commons/server-go/socket" occp "git.loafle.net/overflow/commons-go/config/probe" ocmp "git.loafle.net/overflow/commons-go/model/probe" 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/probe_gateway_rpc/subscribe" ) type ProbeServlet interface { ogrs.RPCServlet } type ProbeServlets struct { ogrs.RPCServlets } func (s *ProbeServlets) Init(serverCtx server.ServerCtx) error { if err := s.RPCServlets.Init(serverCtx); nil != err { return err } return nil } func (s *ProbeServlets) OnStart(serverCtx server.ServerCtx) error { if err := s.RPCServlets.OnStart(serverCtx); nil != err { return err } subscribeChan, err := subscribe.Subscriber.Subscribe(occp.HTTPEntry_Probe) if nil != err { return err } go s.handleSubscribe(serverCtx, subscribeChan) return nil } func (s *ProbeServlets) OnStop(serverCtx server.ServerCtx) { if err := subscribe.Subscriber.Unsubscribe(occp.HTTPEntry_Probe); nil != err { logging.Logger().Warn(err) } s.RPCServlets.OnStop(serverCtx) } func (s *ProbeServlets) Destroy(serverCtx server.ServerCtx) { s.RPCServlets.Destroy(serverCtx) } func (s *ProbeServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) { bMethod := ctx.Request.Header.Peek(occp.HTTPRequestHeaderKey_Probe_Method) if nil == bMethod { return nil, fmt.Errorf("Unexpected probe method: %v", bMethod) } method := string(bMethod) switch method { case occp.HTTPRequestHeaderValue_Probe_Method_Connect: default: return nil, fmt.Errorf("Unexpected probe method: %s", method) } bProbeKey := ctx.Request.Header.Peek(occp.HTTPRequestHeaderKey_Probe_ProbeKey) if nil == bProbeKey { return nil, fmt.Errorf("Unexpected probe key : %v", bProbeKey) } probeKey := string(bProbeKey) grpcCTX := context.Background() r, err := grpc.Exec(grpcCTX, "ProbeService.readByProbeKey", probeKey) if nil != err { return nil, fmt.Errorf("grpc call Error: %s", err.Error()) } probe := &ocmp.Probe{} err = json.Unmarshal([]byte(r), probe) if nil != err { return nil, fmt.Errorf("grpc result unMarshal Error: %s", err.Error()) } extHeader := &fasthttp.ResponseHeader{} extHeader.Add(occp.HTTPResponseHeaderKey_Probe_Method, occp.HTTPResponseHeaderValue_Probe_Method_EncryptionKey) extHeader.Add(occp.HTTPResponseHeaderKey_Probe_Method_Param, probe.EncryptionKey) servletCtx.SetAttribute(og.SessionIDKey, probeKey) servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE) servletCtx.SetAttribute(og.SessionTargetIDKey, probeKey) return extHeader, nil } func (s *ProbeServlets) OnConnect(servletCtx server.ServletCtx, conn socket.Conn) { s.RPCServlets.OnConnect(servletCtx, conn) _probeKey := servletCtx.GetAttribute(og.SessionIDKey) if nil != _probeKey { grpcCTX := context.Background() _, err := grpc.Exec(grpcCTX, "ProbeService.onConnect", _probeKey.(string), conn.RemoteAddr().String()) if nil != err { logging.Logger().Errorf("onConnect error %v", err) } } } func (s *ProbeServlets) OnDisconnect(servletCtx server.ServletCtx) { s.RPCServlets.OnDisconnect(servletCtx) _probeKey := servletCtx.GetAttribute(og.SessionIDKey) if nil != _probeKey { grpcCTX := context.Background() _, err := grpc.Exec(grpcCTX, "ProbeService.onDisconnect", _probeKey.(string)) if nil != err { logging.Logger().Errorf("onDisconnect error %v", err) } } } func (s *ProbeServlets) handleSubscribe(serverCtx server.ServerCtx, subscribeChan <-chan *ogs.Message) { var sessions []*ogrs.Session LOOP: for { select { case msg, ok := <-subscribeChan: if !ok { return } switch msg.TargetType { case ogs.PROBE: sessions = s.GetSessions(msg.Targets) default: logging.Logger().Warnf("Subscriber: Unknown TargetType %s", msg.TargetType) continue LOOP } if nil == sessions || 0 == len(sessions) { continue LOOP } for _, session := range sessions { _writeChan := session.ServletCtx.GetAttribute(og.SessionWriteChanKey) if nil != _writeChan { writeChan := _writeChan.(chan<- []byte) writeChan <- *msg.Message } } } } }