From 18a2af47037874dfa76e99f2426609f715f89d1a Mon Sep 17 00:00:00 2001 From: fatedier Date: Mon, 28 Mar 2022 12:12:35 +0800 Subject: [PATCH] frpc: support multiple confs (#2873) --- Makefile | 5 +++- client/service.go | 23 +++++++++++++++ cmd/frpc/main.go | 8 ------ cmd/frpc/sub/root.go | 60 +++++++++++++++++++++++++-------------- test/e2e/plugin/server.go | 2 +- 5 files changed, 66 insertions(+), 32 deletions(-) diff --git a/Makefile b/Makefile index ae1a013..d3f09bd 100644 --- a/Makefile +++ b/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 diff --git a/client/service.go b/client/service.go index 11e369c..0eab403 100644 --- a/client/service.go +++ b/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() diff --git a/cmd/frpc/main.go b/cmd/frpc/main.go index 2c1e015..f9a6c25 100644 --- a/cmd/frpc/main.go +++ b/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() } diff --git a/cmd/frpc/sub/root.go b/cmd/frpc/sub/root.go index d03cdd1..997733c 100644 --- a/cmd/frpc/sub/root.go +++ b/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,18 +206,9 @@ 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 { @@ -201,9 +216,10 @@ func startService( 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() diff --git a/test/e2e/plugin/server.go b/test/e2e/plugin/server.go index b972f78..2af1eba 100644 --- a/test/e2e/plugin/server.go +++ b/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})