From f5f56d292445df320b4d53760e2147ea84e8a15e Mon Sep 17 00:00:00 2001
From: crusader <crusader@loafle.com>
Date: Mon, 19 Mar 2018 17:05:43 +0900
Subject: [PATCH] ing

---
 servlet/probe/gateway_rpc_handlers.go |  4 ++++
 servlet/probe/subscriber_handlers.go  | 16 +++++++++++-----
 2 files changed, 15 insertions(+), 5 deletions(-)

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()
 }
-
-