diff --git a/servlet/probe/gateway_rpc_handlers.go b/servlet/probe/gateway_rpc_handlers.go index 9c0c503..322811e 100644 --- a/servlet/probe/gateway_rpc_handlers.go +++ b/servlet/probe/gateway_rpc_handlers.go @@ -7,6 +7,7 @@ import ( cwf "git.loafle.net/commons_go/websocket_fasthttp" oocmp "git.loafle.net/overflow/overflow_commons_go/modules/probe" probeM "git.loafle.net/overflow/overflow_commons_go/modules/probe/model" + oogw "git.loafle.net/overflow/overflow_gateway_websocket" oogweg "git.loafle.net/overflow/overflow_gateway_websocket/external/grpc" "git.loafle.net/overflow/overflow_gateway_websocket/rpc" // jwt "github.com/dgrijalva/jwt-go" @@ -82,6 +83,9 @@ func (sh *GatewayRPCHandlers) Handshake(socketCTX cwf.SocketContext, ctx *fastht extensionsHeader = &fasthttp.ResponseHeader{} extensionsHeader.Add(oocmp.HTTPResponseHeaderKey_Probe_SetEncryptionKey, probe.EncryptionKey) + socketCTX.SetAttribute(oogw.ClientTypeKey, oogw.PROBE) + socketCTX.SetAttribute(oogw.TargetIDKey, probeKey) + return probeKey, extensionsHeader } diff --git a/servlet/probe/subscriber_handlers.go b/servlet/probe/subscriber_handlers.go index 9f0e568..56ce0c4 100644 --- a/servlet/probe/subscriber_handlers.go +++ b/servlet/probe/subscriber_handlers.go @@ -1,6 +1,9 @@ package probe import ( + "fmt" + + "git.loafle.net/commons_go/logging" "git.loafle.net/commons_go/websocket_fasthttp/websocket" oogws "git.loafle.net/overflow/overflow_gateway_websocket/servlet" oogwsub "git.loafle.net/overflow/overflow_gateway_websocket/subscribe" @@ -20,15 +23,18 @@ type SubscriberHandlers struct { } func (sh *SubscriberHandlers) OnSubscribe(channel string, message oos.SubscribeMessage) { - for _, uid := range message.Targets { - if soc := sh.GetSocket(uid); nil != soc { - soc.WriteMessage(websocket.TextMessage, message.Message) + switch message.TargetType { + case oos.PROBE: + for _, pid := range message.Targets { + if soc := sh.GetSocket(pid); nil != soc { + soc.WriteMessage(websocket.TextMessage, message.Message) + } } + default: + logging.Logger().Warn(fmt.Sprintf("[PROBE] SubscriberHandler: Unknown TargetType %s", message.TargetType)) } } func (sh *SubscriberHandlers) Validate() { sh.SubscriberHandlers.Validate() } - -