rebrand: frp -> kanhole (kanhole server, kanholec client)

cmd/kanholec/sub/root.go

@@ -0,0 +1,264 @@
+// Copyright 2018 fatedier, fatedier@gmail.com
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package sub
+
+import (
+	"context"
+	"fmt"
+	"io"
+	"io/fs"
+	"net/http"
+	"os"
+	"os/signal"
+	"path/filepath"
+	"strings"
+	"sync"
+	"syscall"
+	"time"
+
+	"github.com/spf13/cobra"
+
+	"kanhole/client"
+	"kanhole/pkg/config"
+	"kanhole/pkg/config/source"
+	v1 "kanhole/pkg/config/v1"
+	"kanhole/pkg/config/v1/validation"
+	"kanhole/pkg/policy/featuregate"
+	"kanhole/pkg/policy/security"
+	"kanhole/pkg/util/log"
+	"kanhole/pkg/util/version"
+)
+
+var (
+	cfgFile          string
+	cfgDir           string
+	showVersion      bool
+	strictConfigMode bool
+	allowUnsafe      []string
+	serverConfigURL  string
+)
+
+func init() {
+	rootCmd.PersistentFlags().StringVarP(&cfgFile, "config", "c", "./kanholec.toml", "config file of kanholec")
+	rootCmd.PersistentFlags().StringVarP(&cfgDir, "config_dir", "", "", "config directory, run one kanholec service for each file in config directory")
+	rootCmd.PersistentFlags().BoolVarP(&showVersion, "version", "v", false, "version of kanholec")
+	rootCmd.PersistentFlags().BoolVarP(&strictConfigMode, "strict_config", "", true, "strict config parsing mode, unknown fields will cause an errors")
+	rootCmd.PersistentFlags().StringVarP(&serverConfigURL, "server-config", "", "", "fetch config from frps server URL (auto-reloads on change)")
+
+	rootCmd.PersistentFlags().StringSliceVarP(&allowUnsafe, "allow-unsafe", "", []string{},
+		fmt.Sprintf("allowed unsafe features, one or more of: %s", strings.Join(security.ClientUnsafeFeatures, ", ")))
+}
+
+var rootCmd = &cobra.Command{
+	Use: "kanholec",
+	Short: "kanholec is the client component of kanhole",
+	RunE: func(cmd *cobra.Command, args []string) error {
+		if showVersion {
+			fmt.Println(version.Full())
+			return nil
+		}
+
+		unsafeFeatures := security.NewUnsafeFeatures(allowUnsafe)
+
+		// If cfgDir is not empty, run multiple kanholec service for each config file in cfgDir.
+		// Note that it's only designed for testing. It's not guaranteed to be stable.
+		if cfgDir != "" {
+			_ = runMultipleClients(cfgDir, unsafeFeatures)
+			return nil
+		}
+
+		// If server-config is set, use it instead of local config file.
+		if serverConfigURL != "" {
+			err := runClientWithServerConfig(serverConfigURL, unsafeFeatures)
+			if err != nil {
+				fmt.Println(err)
+				os.Exit(1)
+			}
+			return nil
+		}
+
+		// Do not show command usage here.
+		err := runClient(cfgFile, unsafeFeatures)
+		if err != nil {
+			fmt.Println(err)
+			os.Exit(1)
+		}
+		return nil
+	},
+}
+
+func runMultipleClients(cfgDir string, unsafeFeatures *security.UnsafeFeatures) error {
+	var wg sync.WaitGroup
+	err := filepath.WalkDir(cfgDir, func(path string, d fs.DirEntry, err error) error {
+		if err != nil || d.IsDir() {
+			return nil
+		}
+		wg.Add(1)
+		time.Sleep(time.Millisecond)
+		go func() {
+			defer wg.Done()
+			err := runClient(path, unsafeFeatures)
+			if err != nil {
+				fmt.Printf("kanholec service error for config file [%s]\n", path)
+			}
+		}()
+		return nil
+	})
+	wg.Wait()
+	return err
+}
+
+func Execute() {
+	rootCmd.SetGlobalNormalizationFunc(config.WordSepNormalizeFunc)
+	if err := rootCmd.Execute(); err != nil {
+		os.Exit(1)
+	}
+}
+
+func handleTermSignal(svr *client.Service) {
+	ch := make(chan os.Signal, 1)
+	signal.Notify(ch, syscall.SIGINT, syscall.SIGTERM)
+	<-ch
+	svr.GracefulClose(500 * time.Millisecond)
+}
+
+func runClientWithServerConfig(url string, unsafeFeatures *security.UnsafeFeatures) error {
+	log.Infof("fetching config from server: %s", url)
+	resp, err := http.Get(url)
+	if err != nil {
+		return fmt.Errorf("failed to fetch config from server: %w", err)
+	}
+	defer resp.Body.Close()
+
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return fmt.Errorf("failed to read config from server: %w", err)
+	}
+
+	allCfg := v1.ClientConfig{}
+	if err := config.LoadConfigure(body, &allCfg, strictConfigMode, "toml"); err != nil {
+		return fmt.Errorf("failed to parse config from server: %w", err)
+	}
+
+	result := &config.ClientConfigLoadResult{
+		Common:  &allCfg.ClientCommonConfig,
+		Proxies: make([]v1.ProxyConfigurer, 0),
+	}
+	for _, c := range allCfg.Proxies {
+		result.Proxies = append(result.Proxies, c.ProxyConfigurer)
+	}
+
+	result.Common.ConfigURL = url
+
+	return runClientWithAggregator(result, unsafeFeatures, "")
+}
+
+func runClient(cfgFilePath string, unsafeFeatures *security.UnsafeFeatures) error {
+	// Load configuration
+	result, err := config.LoadClientConfigResult(cfgFilePath, strictConfigMode)
+	if err != nil {
+		return err
+	}
+	if result.IsLegacyFormat {
+		fmt.Printf("WARNING: ini format is deprecated and the support will be removed in the future, " +
+			"please use yaml/json/toml format instead!\n")
+	}
+
+	if len(result.Common.FeatureGates) > 0 {
+		if err := featuregate.SetFromMap(result.Common.FeatureGates); err != nil {
+			return err
+		}
+	}
+
+	return runClientWithAggregator(result, unsafeFeatures, cfgFilePath)
+}
+
+// runClientWithAggregator runs the client using the internal source aggregator.
+func runClientWithAggregator(result *config.ClientConfigLoadResult, unsafeFeatures *security.UnsafeFeatures, cfgFilePath string) error {
+	configSource := source.NewConfigSource()
+	if err := configSource.ReplaceAll(result.Proxies, result.Visitors); err != nil {
+		return fmt.Errorf("failed to set config source: %w", err)
+	}
+
+	var storeSource *source.StoreSource
+
+	if result.Common.Store.IsEnabled() {
+		storePath := result.Common.Store.Path
+		if storePath != "" && cfgFilePath != "" && !filepath.IsAbs(storePath) {
+			storePath = filepath.Join(filepath.Dir(cfgFilePath), storePath)
+		}
+
+		s, err := source.NewStoreSource(source.StoreSourceConfig{
+			Path: storePath,
+		})
+		if err != nil {
+			return fmt.Errorf("failed to create store source: %w", err)
+		}
+		storeSource = s
+	}
+
+	aggregator := source.NewAggregator(configSource)
+	if storeSource != nil {
+		aggregator.SetStoreSource(storeSource)
+	}
+
+	proxyCfgs, visitorCfgs, err := aggregator.Load()
+	if err != nil {
+		return fmt.Errorf("failed to load config from sources: %w", err)
+	}
+
+	proxyCfgs, visitorCfgs = config.FilterClientConfigurers(result.Common, proxyCfgs, visitorCfgs)
+	proxyCfgs = config.CompleteProxyConfigurers(proxyCfgs)
+	visitorCfgs = config.CompleteVisitorConfigurers(visitorCfgs)
+
+	warning, err := validation.ValidateAllClientConfig(result.Common, proxyCfgs, visitorCfgs, unsafeFeatures)
+	if warning != nil {
+		fmt.Printf("WARNING: %v\n", warning)
+	}
+	if err != nil {
+		return err
+	}
+
+	return startServiceWithAggregator(result.Common, aggregator, unsafeFeatures, cfgFilePath)
+}
+
+func startServiceWithAggregator(
+	cfg *v1.ClientCommonConfig,
+	aggregator *source.Aggregator,
+	unsafeFeatures *security.UnsafeFeatures,
+	cfgFile string,
+) error {
+	log.InitLogger(cfg.Log.To, cfg.Log.Level, int(cfg.Log.MaxDays), cfg.Log.DisablePrintColor)
+
+	if cfgFile != "" {
+		log.Infof("start kanholec service for config file [%s] with aggregated configuration", cfgFile)
+		defer log.Infof("kanholec service for config file [%s] stopped", cfgFile)
+	}
+	svr, err := client.NewService(client.ServiceOptions{
+		Common:                 cfg,
+		ConfigSourceAggregator: aggregator,
+		UnsafeFeatures:         unsafeFeatures,
+		ConfigFilePath:         cfgFile,
+	})
+	if err != nil {
+		return err
+	}
+
+	shouldGracefulClose := cfg.Transport.Protocol == "kcp" || cfg.Transport.Protocol == "quic"
+	if shouldGracefulClose {
+		go handleTermSignal(svr)
+	}
+	return svr.Run(context.Background())
+}