From f4ae503982a5cc9de6638fca377fd8dac35e67c1 Mon Sep 17 00:00:00 2001
From: unknown <liuhe02@ESG.360ES.CN>
Date: Tue, 15 Oct 2019 13:59:08 +0800
Subject: [PATCH] module

---
 core/struct.go                            |  9 +++++++++
 server/common/common_inet_proxy_handle.go | 16 +++++-----------
 2 files changed, 14 insertions(+), 11 deletions(-)

diff --git a/core/struct.go b/core/struct.go
index 61b85b4..ecfafef 100644
--- a/core/struct.go
+++ b/core/struct.go
@@ -3,6 +3,7 @@ package core
 import (
 	"context"
 	"errors"
+	"github.com/cnlh/nps/bridge"
 	"net"
 )
 
@@ -67,3 +68,11 @@ func (npsPlugin *NpsPlugin) End(ctx context.Context, config map[string]string) e
 func (npsPlugin *NpsPlugin) GetClientConn(ctx context.Context) net.Conn {
 	return ctx.Value(CLIENT_CONNECTION).(net.Conn)
 }
+
+func (npsPlugin *NpsPlugin) GetBridge(ctx context.Context) *bridge.Bridge {
+	return ctx.Value(BRIDGE).(*bridge.Bridge)
+}
+
+func (npsPlugin *NpsPlugin) GetClientId(ctx context.Context) int {
+	return ctx.Value(CLIENT_ID).(int)
+}
diff --git a/server/common/common_inet_proxy_handle.go b/server/common/common_inet_proxy_handle.go
index 14fe1cd..de070a0 100644
--- a/server/common/common_inet_proxy_handle.go
+++ b/server/common/common_inet_proxy_handle.go
@@ -2,7 +2,6 @@ package common
 
 import (
 	"context"
-	"github.com/cnlh/nps/bridge"
 	"github.com/cnlh/nps/core"
 	"net"
 )
@@ -20,25 +19,20 @@ func (proxy *Proxy) GetConfigName() *core.NpsConfigs {
 func (proxy *Proxy) Run(ctx context.Context, config map[string]string) error {
 	proxy.clientConn = proxy.GetClientConn(ctx)
 	proxy.ctx = ctx
-	bg := ctx.Value(core.BRIDGE)
-	if bg == nil {
-		return core.BRIDGE_NOT_EXIST
-	}
+
 	clientCtxConn := ctx.Value(core.CLIENT_CONNECTION)
 	if clientCtxConn == nil {
 		return core.CLIENT_CONNECTION_NOT_EXIST
 	}
 
-	clientId := ctx.Value(core.CLIENT_ID)
-	if clientId == nil {
-		return core.CLIENT_ID_NOT_EXIST
-	}
+	clientId := proxy.GetClientId(ctx)
 
-	brg := bg.(*bridge.Bridge)
-	severConn, err := brg.GetConnByClientId(clientId.(int))
+	brg := proxy.GetBridge(ctx)
+	severConn, err := brg.GetConnByClientId(clientId)
 	if err != nil {
 		return err
 	}
+
 	go core.CopyBuffer(severConn, clientCtxConn.(net.Conn))
 	core.CopyBuffer(clientCtxConn.(net.Conn), severConn)
 	return nil