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{