150 lines
3.7 KiB
Go
150 lines
3.7 KiB
Go
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"
|
|
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"
|
|
ogs "git.loafle.net/overflow/gateway/subscribe"
|
|
ogrs "git.loafle.net/overflow/gateway_rpc/servlet"
|
|
"git.loafle.net/overflow/member_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(ocpc.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(ocpc.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(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, 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)
|
|
}
|
|
|
|
func (s *ProbeServlets) OnDisconnect(servletCtx server.ServletCtx) {
|
|
s.RPCServlets.OnDisconnect(servletCtx)
|
|
}
|
|
|
|
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
|
|
}
|
|
}
|
|
|
|
}
|
|
}
|
|
}
|