package servlet import ( "context" "crypto/rsa" "encoding/json" "fmt" "sync" "github.com/valyala/fasthttp" "git.loafle.net/commons/server-go" "git.loafle.net/commons/server-go/socket" ocpc "git.loafle.net/overflow/commons-go/probe/constants" ocpm "git.loafle.net/overflow/commons-go/probe/model" og "git.loafle.net/overflow/gateway" "git.loafle.net/overflow/gateway/external/grpc" ogrs "git.loafle.net/overflow/gateway_rpc/servlet" ) type DataServlet interface { ogrs.RPCServlet } type DataServlets struct { ogrs.RPCServlets VerifyKey *rsa.PublicKey SignKey *rsa.PrivateKey connections sync.Map } func (s *DataServlets) Init(serverCtx server.ServerCtx) error { if err := s.RPCServlets.Init(serverCtx); nil != err { return err } return nil } func (s *DataServlets) OnStart(serverCtx server.ServerCtx) error { if err := s.RPCServlets.OnStart(serverCtx); nil != err { return err } return nil } func (s *DataServlets) OnStop(serverCtx server.ServerCtx) { s.RPCServlets.OnStop(serverCtx) } func (s *DataServlets) Destroy(serverCtx server.ServerCtx) { s.RPCServlets.Destroy(serverCtx) } func (s *DataServlets) Handshake(servletCtx server.ServletCtx, ctx *fasthttp.RequestCtx) (*fasthttp.ResponseHeader, error) { // probe key extraction bMethod := ctx.Request.Header.Peek(ocpc.HTTPRequestHeaderKey_Probe_Method) if nil == bMethod { return nil, fmt.Errorf("Unexpected probe method: %v", bMethod) } method := string(bMethod) switch method { case ocpc.HTTPRequestHeaderValue_Probe_Method_Connect: default: return nil, fmt.Errorf("Unexpected probe method: %s", method) } bProbeKey := ctx.Request.Header.Peek(ocpc.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 := ocpm.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(ocpc.HTTPResponseHeaderKey_Probe_SetEncryptionKey, probe.EncryptionKey) servletCtx.SetAttribute(og.SessionIDKey, probe.ProbeKey) servletCtx.SetAttribute(og.SessionClientTypeKey, og.PROBE) servletCtx.SetAttribute(og.SessionTargetIDKey, probe.ProbeKey) return extHeader, nil } func (s *DataServlets) OnConnect(servletCtx server.ServletCtx, conn socket.Conn) { s.RPCServlets.OnConnect(servletCtx, conn) } func (s *DataServlets) OnDisconnect(servletCtx server.ServletCtx) { s.RPCServlets.OnDisconnect(servletCtx) }