From 0280eae5213eae1ea1b2716b28a62ff79a91be90 Mon Sep 17 00:00:00 2001
From: afocus <wangyinliang@parkingwang.com>
Date: Wed, 24 Jan 2018 11:45:42 +0800
Subject: [PATCH] =?UTF-8?q?=E4=BF=AE=E6=94=B9=E6=A0=A1=E9=AA=8C=E7=9A=84?=
 =?UTF-8?q?=E5=85=A5=E5=8F=A3?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 server/irain.go   | 9 +++++++--
 server/proxy.go   | 8 +++++++-
 server/service.go | 9 ---------
 3 files changed, 14 insertions(+), 12 deletions(-)

diff --git a/server/irain.go b/server/irain.go
index 095f9bbd..11c5fe75 100644
--- a/server/irain.go
+++ b/server/irain.go
@@ -5,8 +5,10 @@ import (
 	"encoding/json"
 	"fmt"
 	"github.com/julienschmidt/httprouter"
+	"net"
 	"net/http"
 	"strconv"
+	"strings"
 	"sync"
 	"time"
 )
@@ -22,7 +24,9 @@ func IRainRespone(w http.ResponseWriter, code int, msg string) {
 		Message: msg,
 	}
 	b, _ := json.Marshal(ret)
-	w.WriteHeader(code)
+	if code != 0 {
+		w.WriteHeader(code)
+	}
 	w.Write(b)
 }
 
@@ -47,7 +51,8 @@ func (p *IRainIPPool) Put(ip string) {
 	p.list[ip] = time.Now().Add(time.Minute * 30)
 }
 
-func (p *IRainIPPool) Check(ip string) bool {
+func (p *IRainIPPool) Check(addr net.Addr) bool {
+	ip := strings.Split(addr.String(), ":")[0]
 	p.mux.RLock()
 	defer p.mux.RUnlock()
 	if v, ok := p.list[ip]; ok {
diff --git a/server/proxy.go b/server/proxy.go
index 554e8181..09f3721e 100644
--- a/server/proxy.go
+++ b/server/proxy.go
@@ -110,7 +110,13 @@ func (pxy *BaseProxy) startListenHandler(p Proxy, handler func(Proxy, frpNet.Con
 					pxy.Info("listener is closed")
 					return
 				}
-				pxy.Debug("get a user connection [%s]", c.RemoteAddr().String())
+				useraddr := c.RemoteAddr().String()
+				pxy.Debug("get a user connection [%s]", useraddr)
+				if !globalIRainIPPool.Check(c.RemoteAddr()) {
+					c.Close()
+					pxy.Warn("user connection not auth [%s]", useraddr)
+					return
+				}
 				go handler(p, c)
 			}
 		}(listener)
diff --git a/server/service.go b/server/service.go
index b59e8a8a..e976658a 100644
--- a/server/service.go
+++ b/server/service.go
@@ -214,15 +214,6 @@ func (svr *Service) HandleListener(l frpNet.Listener) {
 				case *msg.NewWorkConn:
 					svr.RegisterWorkConn(conn, m)
 				case *msg.NewVisitorConn:
-
-					// irain
-					// 检测是否允许访问者ip
-					if globalIRainIPPool.Check(conn.RemoteAddr().String()) {
-						log.Warn("visitor[%s] not allow access: %s", conn.RemoteAddr())
-						conn.Close()
-						return
-					}
-
 					if err = svr.RegisterVisitorConn(conn, m); err != nil {
 						conn.Warn("%v", err)
 						msg.WriteMsg(conn, &msg.NewVisitorConnResp{