frpc: support multiple confs (fatedier#2873)

Author: fatedier

Makefile

@@ -16,6 +16,9 @@ file:
 fmt:
 	go fmt ./...
 
+vet:
+	go vet ./...
+
 frps:
 	env CGO_ENABLED=0 go build -trimpath -ldflags "$(LDFLAGS)" -o bin/frps ./cmd/frps
 
@@ -37,7 +40,7 @@ e2e:
 e2e-trace:
 	DEBUG=true LOG_LEVEL=trace ./hack/run-e2e.sh
 
-alltest: gotest e2e
+alltest: vet gotest e2e
 
 clean:
 	rm -f ./bin/frpc

client/service.go

@@ -19,9 +19,11 @@ import (
 	"crypto/tls"
 	"fmt"
 	"io"
+	"math/rand"
 	"net"
 	"runtime"
 	"strconv"
+	"strings"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -36,11 +38,17 @@ import (
 	"github.com/fatedier/frp/pkg/util/util"
 	"github.com/fatedier/frp/pkg/util/version"
 	"github.com/fatedier/frp/pkg/util/xlog"
+	"github.com/fatedier/golib/crypto"
 	libdial "github.com/fatedier/golib/net/dial"
 
 	fmux "github.com/hashicorp/yamux"
 )
 
+func init() {
+	crypto.DefaultSalt = "frp"
+	rand.Seed(time.Now().UnixNano())
+}
+
 // Service is a client service.
 type Service struct {
 	// uniq id got from frps, attach it in loginMsg
@@ -98,6 +106,21 @@ func (svr *Service) GetController() *Control {
 func (svr *Service) Run() error {
 	xl := xlog.FromContextSafe(svr.ctx)
 
+	// set custom DNSServer
+	if svr.cfg.DNSServer != "" {
+		dnsAddr := svr.cfg.DNSServer
+		if !strings.Contains(dnsAddr, ":") {
+			dnsAddr += ":53"
+		}
+		// Change default dns server for frpc
+		net.DefaultResolver = &net.Resolver{
+			PreferGo: true,
+			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
+				return net.Dial("udp", dnsAddr)
+			},
+		}
+	}
+
 	// login to frps
 	for {
 		conn, session, err := svr.login()

cmd/frpc/main.go

@@ -15,18 +15,10 @@
 package main
 
 import (
-	"math/rand"
-	"time"
-
 	_ "github.com/fatedier/frp/assets/frpc"
 	"github.com/fatedier/frp/cmd/frpc/sub"
-
-	"github.com/fatedier/golib/crypto"
 )
 
 func main() {
-	crypto.DefaultSalt = "frp"
-	rand.Seed(time.Now().UnixNano())
-
 	sub.Execute()
 }

cmd/frpc/sub/root.go

@@ -15,13 +15,14 @@
 package sub
 
 import (
-	"context"
 	"fmt"
+	"io/fs"
 	"net"
 	"os"
 	"os/signal"
+	"path/filepath"
 	"strconv"
-	"strings"
+	"sync"
 	"syscall"
 	"time"
 
@@ -41,6 +42,7 @@ const (
 
 var (
 	cfgFile     string
+	cfgDir      string
 	showVersion bool
 
 	serverAddr      string
@@ -72,15 +74,12 @@ var (
 	bindPort          int
 
 	tlsEnable bool
-
-	kcpDoneCh chan struct{}
 )
 
 func init() {
 	rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "./frpc.ini", "config file of frpc")
+	rootCmd.PersistentFlags().StringVarP(&cfgDir, "config_dir", "", "", "config directory, run one frpc service for each file in config directory")
 	rootCmd.PersistentFlags().BoolVarP(&showVersion, "version", "v", false, "version of frpc")
-
-	kcpDoneCh = make(chan struct{})
 }
 
 func RegisterCommonFlags(cmd *cobra.Command) {
@@ -104,6 +103,31 @@ var rootCmd = &cobra.Command{
 			return nil
 		}
 
+		// If cfgDir is not empty, run multiple frpc service for each config file in cfgDir.
+		// Note that it's only designed for testing. It's not guaranteed to be stable.
+		if cfgDir != "" {
+			var wg sync.WaitGroup
+			filepath.WalkDir(cfgDir, func(path string, d fs.DirEntry, err error) error {
+				if err != nil {
+					return nil
+				}
+				if d.IsDir() {
+					return nil
+				}
+				wg.Add(1)
+				go func() {
+					defer wg.Done()
+					err := runClient(path)
+					if err != nil {
+						fmt.Printf("frpc service error for config file [%s]\n", path)
+					}
+				}()
+				return nil
+			})
+			wg.Wait()
+			return nil
+		}
+
 		// Do not show command usage here.
 		err := runClient(cfgFile)
 		if err != nil {
@@ -120,12 +144,12 @@ func Execute() {
 	}
 }
 
-func handleSignal(svr *client.Service) {
-	ch := make(chan os.Signal)
+func handleSignal(svr *client.Service, doneCh chan struct{}) {
+	ch := make(chan os.Signal, 1)
 	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
 	<-ch
 	svr.GracefulClose(500 * time.Millisecond)
-	close(kcpDoneCh)
+	close(doneCh)
 }
 
 func parseClientCommonCfgFromCmd() (cfg config.ClientCommonConf, err error) {
@@ -182,28 +206,20 @@ func startService(
 	log.InitLog(cfg.LogWay, cfg.LogFile, cfg.LogLevel,
 		cfg.LogMaxDays, cfg.DisableLogColor)
 
-	if cfg.DNSServer != "" {
-		s := cfg.DNSServer
-		if !strings.Contains(s, ":") {
-			s += ":53"
-		}
-		// Change default dns server for frpc
-		net.DefaultResolver = &net.Resolver{
-			PreferGo: true,
-			Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
-				return net.Dial("udp", s)
-			},
-		}
+	if cfgFile != "" {
+		log.Trace("start frpc service for config file [%s]", cfgFile)
+		defer log.Trace("frpc service for config file [%s] stopped", cfgFile)
 	}
 	svr, errRet := client.NewService(cfg, pxyCfgs, visitorCfgs, cfgFile)
 	if errRet != nil {
 		err = errRet
 		return
 	}
 
+	kcpDoneCh := make(chan struct{})
 	// Capture the exit signal if we use kcp.
 	if cfg.Protocol == "kcp" {
-		go handleSignal(svr)
+		go handleSignal(svr, kcpDoneCh)
 	}
 
 	err = svr.Run()

test/e2e/plugin/server.go

@@ -150,7 +150,7 @@ var _ = Describe("[Feature: Server-Plugins]", func() {
 			type = tcp
 			local_port = {{ .%s }}
 			remote_port = 0
-			`, framework.TCPEchoServerPort, remotePort)
+			`, framework.TCPEchoServerPort)
 
 			f.RunProcesses([]string{serverConf}, []string{clientConf})