From bc5cdf120b1dbd7043231262a741aba9bfdd3306 Mon Sep 17 00:00:00 2001
From: bingoohuang <bingoo.huang@gmail.com>
Date: Thu, 16 Mar 2023 11:07:12 +0800
Subject: [PATCH] add allow_ips for frps

---
 cmd/frps/root.go     |  3 +++
 pkg/config/server.go |  4 ++++
 server/service.go    | 14 ++++++++++++++
 3 files changed, 21 insertions(+)

diff --git a/cmd/frps/root.go b/cmd/frps/root.go
index 955c90fb..abe27f0b 100644
--- a/cmd/frps/root.go
+++ b/cmd/frps/root.go
@@ -38,6 +38,7 @@ var (
 	showVersion bool
 
 	bindAddr             string
+	allowIps             []string
 	bindPort             int
 	bindUDPPort          int
 	kcpBindPort          int
@@ -69,6 +70,7 @@ func init() {
 	rootCmd.PersistentFlags().BoolVarP(&showVersion, "version", "v", false, "version of frps")
 
 	rootCmd.PersistentFlags().StringVarP(&bindAddr, "bind_addr", "", "0.0.0.0", "bind address")
+	rootCmd.PersistentFlags().StringArrayVarP(&allowIps, "allow_ip", "", nil, "ip whitelists")
 	rootCmd.PersistentFlags().IntVarP(&bindPort, "bind_port", "p", 7000, "bind port")
 	rootCmd.PersistentFlags().IntVarP(&bindUDPPort, "bind_udp_port", "", 0, "bind udp port")
 	rootCmd.PersistentFlags().IntVarP(&kcpBindPort, "kcp_bind_port", "", 0, "kcp bind udp port")
@@ -159,6 +161,7 @@ func parseServerCommonCfgFromCmd() (cfg config.ServerCommonConf, err error) {
 
 	cfg.BindAddr = bindAddr
 	cfg.BindPort = bindPort
+	cfg.AllowIps = allowIps
 	cfg.BindUDPPort = bindUDPPort
 	cfg.KCPBindPort = kcpBindPort
 	cfg.ProxyBindAddr = proxyBindAddr
diff --git a/pkg/config/server.go b/pkg/config/server.go
index cdc29daa..479adfc6 100644
--- a/pkg/config/server.go
+++ b/pkg/config/server.go
@@ -50,6 +50,10 @@ type ServerCommonConf struct {
 	// Set this value to 0 will disable this feature.
 	// By default, the value is 0.
 	QUICBindPort int `ini:"quic_bind_port" json:"quic_bind_port" validate:"gte=0,lte=65535"`
+
+	// AllowIps specifies the IP whitelists to limit clients.
+	AllowIps []string `ini:"allow_ips" json:"allow_ips"`
+
 	// QUIC protocol options
 	QUICKeepalivePeriod    int `ini:"quic_keepalive_period" json:"quic_keepalive_period" validate:"gte=0"`
 	QUICMaxIdleTimeout     int `ini:"quic_max_idle_timeout" json:"quic_max_idle_timeout" validate:"gte=0"`
diff --git a/server/service.go b/server/service.go
index e2f84940..bfa27a48 100644
--- a/server/service.go
+++ b/server/service.go
@@ -431,8 +431,22 @@ func (svr *Service) HandleListener(l net.Listener) {
 		}
 		log.Trace("check TLS connection success, isTLS: %v custom: %v", isTLS, custom)
 
+		allowIps := make(map[string]bool, len(svr.cfg.AllowIps))
+		for _, allowIP := range svr.cfg.AllowIps {
+			allowIps[allowIP] = true
+		}
 		// Start a new goroutine to handle connection.
 		go func(ctx context.Context, frpConn net.Conn) {
+			if len(allowIps) > 0 {
+				if addr, ok := frpConn.RemoteAddr().(*net.TCPAddr); ok {
+					if remoteIP := addr.IP.String(); !allowIps[remoteIP] {
+						log.Warn("Connection from %s is not allowed", remoteIP)
+						frpConn.Close()
+						return
+					}
+				}
+			}
+
 			if svr.cfg.TCPMux {
 				fmuxCfg := fmux.DefaultConfig()
 				fmuxCfg.KeepAliveInterval = time.Duration(svr.cfg.TCPMuxKeepaliveInterval) * time.Second