修改校验的入口

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)
+	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 {

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)

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{