1.添加自定义配置项

1)在frps/root.go第64行添加变量
maxProxyPadCounts int64
2)在frps/root.go->init方法第96行指定变量缺省值
rootCmd.PersistentFlags().Int64VarP(&maxProxyPadCounts, "max_proxy_pad_counts", "", 10, "max proxy pad counts")
3)在frps/root.go->parseServerCommonCfgFromCmd方法第181行将缺省值赋值给cfg
cfg.MaxProxyPadCounts = maxProxyPadCounts

然后在frps.ini配置max_proxy_pad_counts字段就能指定代理的云机个数

2.修改tcp连接处理位置源码,根据自定义字段限制代理云机个数。

1)修改server/proxy/proxy.go,在tcp连接的地方+1

var tcp_count = 0

// startCommonTCPListenersHandler start a goroutine handler for each listener.
func (pxy *BaseProxy) startCommonTCPListenersHandler() {
        tcp_count = 0
        xl := xlog.FromContextSafe(pxy.ctx)
        for _, listener := range pxy.listeners {
                go func(l net.Listener) {
                        var tempDelay time.Duration // how long to sleep on accept failure

                        for {
                                // block
                                // if listener is closed, err returned
                                c, err := l.Accept()
                                if err != nil {
                                        if err, ok := err.(interface{ Temporary() bool }); ok && err.Temporary() {
                                                if tempDelay == 0 {
                                                        tempDelay = 5 * time.Millisecond
                                                } else {
                                                        tempDelay *= 2
                                                }
                                                if max := 1 * time.Second; tempDelay > max {
                                                        tempDelay = max
                                                }
                                                xl.Info("met temporary error: %s, sleep for %s ...", err, tempDelay)
                                                time.Sleep(tempDelay)
                                                continue
                                        }

                                        xl.Warn("listener is closed: %s", err)
                                        return
                               }

                                xl.Info("get a user connection [%s]", c.RemoteAddr().String())
                                if pxy.name == "tcp_remote" {
                                        tcp_count++
                                }
                                log.Info("proxy name:%s", pxy.name)
                                log.Info("max proxy pad counts:%d", pxy.serverCfg.MaxProxyPadCounts)
                                log.Info("cur tcp_remote counts:%d", tcp_count)
                                go pxy.handleUserTCPConnection(c)
                        }
                }(listener)
        }
}

2)修改server/proxy/proxy.go,超限则-1并主动断开连接。

func closeUsrConn(pxy *BaseProxy, userConn net.Conn) {
        xl := xlog.FromContextSafe(pxy.Context())
        xl.Info("closeUsrConn")
        tcp_count--
        userConn.Close()
}

// HandleUserTCPConnection is used for incoming user TCP connections.
func (pxy *BaseProxy) handleUserTCPConnection(userConn net.Conn) {
        xl := xlog.FromContextSafe(pxy.Context())
        defer closeUsrConn(pxy, userConn)
        if pxy.name == "tcp_remote" {
                if tcp_count > int(pxy.serverCfg.MaxProxyPadCounts) {
                        xl.Warn("tcp connect full,cur is:", tcp_count)
                        return
                }
        }

        serverCfg := pxy.serverCfg
        cfg := pxy.pxyConf.GetBaseConfig()
        // server plugin hook
        rc := pxy.GetResourceController()
        content := &plugin.NewUserConnContent{
                User:       pxy.GetUserInfo(),
                ProxyName:  pxy.GetName(),
                ProxyType:  cfg.ProxyType,
                RemoteAddr: userConn.RemoteAddr().String(),
        }
        _, err := rc.PluginManager.NewUserConn(content)
...
}
分类: frp

0 条评论

发表回复

您的电子邮箱地址不会被公开。