Supporting socksd fallback for web2

cmd/client_debug/client_debug.go

@@ -45,7 +45,7 @@ func main() {
 	}
 	var wg sync.WaitGroup
 	config = parseFlag()
-	prox, _ := url.Parse(fmt.Sprintf("socks5://%s", config.SocksServerAddr))
+	prox, _ := url.Parse(fmt.Sprintf("socks5://%s", config.SocksServerAddr()))
 	proxyTransport.Proxy = http.ProxyURL(prox)
 	log.Printf("Start to connect %d times", config.Conn)
 	wg.Add(config.Conn)

cmd/client_debug/flag.go

@@ -13,7 +13,6 @@ type Config struct {
 	Target          string
 	Conn            int
 	EnableTransport bool
-	SocksServerAddr string
 	SocksServerHost string
 	SocksServerPort int
 }
@@ -29,6 +28,9 @@ func parseFlag() *Config {
 	flag.IntVar(&cfg.Conn, "conn", 100, "total connection concurrently")
 
 	flag.Parse()
-	cfg.SocksServerAddr = fmt.Sprintf("%s:%d", cfg.SocksServerHost, cfg.SocksServerPort)
 	return cfg
 }
+
+func (cfg *Config) SocksServerAddr() string {
+	return fmt.Sprintf("%s:%d", cfg.SocksServerHost, cfg.SocksServerPort)
+}

cmd/diode/diode.go

@@ -25,29 +25,25 @@ import (
 )
 
 var (
-	version string = "development"
+	version     string = "development"
+	socksServer *rpc.Server
+	proxyServer *rpc.ProxyServer
+	pool        *rpc.DataPool
 )
 
 func init() {
 	config.ParseFlag()
 }
 
 func main() {
-	var socksServer *rpc.Server
-	var proxyServer *rpc.ProxyServer
 	var err error
-	var pool *rpc.DataPool
 
 	if version != "development" {
 		doUpdate()
 	}
 
 	cfg := config.AppConfig
-	if len(cfg.PublishedPorts) > 0 {
-		pool = rpc.NewPoolWithPublishedPorts(cfg.PublishedPorts)
-	} else {
-		pool = rpc.NewPool()
-	}
+	pool = rpc.NewPool()
 
 	printLabel("Diode Client version", version)
 
@@ -178,59 +174,62 @@ func main() {
 		sig := <-sigChan
 		switch sig {
 		case syscall.SIGINT:
-			closeDiode(client, socksServer, proxyServer, cfg)
+			os.Exit(0)
 		}
 	}()
 
-	socksConfig := &rpc.Config{
-		Addr:            cfg.SocksServerAddr,
+	socksServer = client.NewSocksServer(pool)
+	proxyServer = rpc.NewProxyServer(socksServer)
+
+	processConfig(cfg)
+	// start
+	client.Wait()
+	closeDiode(client, cfg)
+}
+
+func processConfig(cfg *config.Config) {
+	if len(cfg.PublishedPorts) > 0 {
+		pool.SetPublishedPorts(cfg.PublishedPorts)
+	}
+
+	socksServer.SetConfig(&rpc.Config{
+		Addr:            cfg.SocksServerAddr(),
 		FleetAddr:       cfg.FleetAddr,
 		Blacklists:      cfg.Blacklists,
 		Whitelists:      cfg.Whitelists,
 		EnableProxy:     cfg.EnableProxyServer,
-		ProxyServerAddr: cfg.ProxyServerAddr,
-	}
-	socksServer = client.NewSocksServer(socksConfig, pool)
+		ProxyServerAddr: cfg.ProxyServerAddr(),
+		Fallback:        cfg.SocksFallback,
+	})
 
 	if cfg.EnableSocksServer {
 		// start socks server
 		if err := socksServer.Start(); err != nil {
 			cfg.Logger.Error(err.Error(), "module", "main")
-			return
 		}
+	} else {
+		socksServer.Stop()
 	}
+
 	if cfg.EnableProxyServer {
-		proxyConfig := rpc.ProxyConfig{
+		proxyServer.SetConfig(rpc.ProxyConfig{
 			EnableProxy:      cfg.EnableProxyServer,
 			EnableSProxy:     cfg.EnableSProxyServer,
-			ProxyServerAddr:  cfg.ProxyServerAddr,
-			SProxyServerAddr: cfg.SProxyServerAddr,
+			ProxyServerAddr:  cfg.ProxyServerAddr(),
+			SProxyServerAddr: cfg.SProxyServerAddr(),
 			CertPath:         cfg.SProxyServerCertPath,
 			PrivPath:         cfg.SProxyServerPrivPath,
 			AllowRedirect:    cfg.AllowRedirectToSProxy,
-		}
+		})
 		// Start proxy server
-		if proxyServer, err = rpc.NewProxyServer(socksServer, proxyConfig); err != nil {
-			cfg.Logger.Error(err.Error(), "module", "main")
-			return
-		}
 		if err := proxyServer.Start(); err != nil {
 			cfg.Logger.Error(err.Error(), "module", "main")
-			return
-		}
-	}
-
-	for _, bind := range cfg.Binds {
-		err = socksServer.StartBind(bind)
-		if err != nil {
-			cfg.Logger.Error(err.Error(), "module", "main")
-			return
 		}
+	} else {
+		proxyServer.Stop()
 	}
 
-	// start
-	client.Wait()
-	closeDiode(client, socksServer, proxyServer, cfg)
+	socksServer.SetBinds(cfg.Binds)
 }
 
 func doConfig(cfg *config.Config) {
@@ -511,7 +510,7 @@ func connect(c chan *rpc.RPCClient, host string, cfg *config.Config, wg *sync.Wa
 	}
 }
 
-func closeDiode(client *rpc.RPCClient, socksServer *rpc.Server, proxyServer *rpc.ProxyServer, cfg *config.Config) {
+func closeDiode(client *rpc.RPCClient, cfg *config.Config) {
 	if client.Started() {
 		client.Close()
 	}

config/command_flag.go

@@ -73,6 +73,9 @@ func wrapPublishCommandFlag(cfg *Config) {
 	publishCommandFlag.Flag.Var(&cfg.PublicPublishedPorts, "public", "expose ports to public users, so that user could connect to")
 	publishCommandFlag.Flag.Var(&cfg.ProtectedPublishedPorts, "protected", "expose ports to protected users (in fleet contract), so that user could connect to")
 	publishCommandFlag.Flag.Var(&cfg.PrivatePublishedPorts, "private", "expose ports to private users, so that user could connect to")
+	publishCommandFlag.Flag.StringVar(&cfg.SocksServerHost, "proxy_host", "127.0.0.1", "host of socksd proxy server")
+	publishCommandFlag.Flag.IntVar(&cfg.SocksServerPort, "proxy_port", 1080, "port of socksd proxy server")
+	publishCommandFlag.Flag.BoolVar(&cfg.EnableSocksServer, "socksd", false, "enable socksd proxy server")
 	publishCommandFlag.Flag.Usage = func() {
 		printUsage(publishCommandFlag)
 	}
@@ -81,6 +84,7 @@ func wrapPublishCommandFlag(cfg *Config) {
 func wrapSocksdCommandFlag(cfg *Config) {
 	socksdCommandFlag.Flag.StringVar(&cfg.SocksServerHost, "socksd_host", "127.0.0.1", "host of socks server listening to")
 	socksdCommandFlag.Flag.IntVar(&cfg.SocksServerPort, "socksd_port", 1080, "port of socks server listening to")
+	socksdCommandFlag.Flag.StringVar(&cfg.SocksFallback, "fallback", "localhost", "how to resolve web2 addresses")
 	socksdCommandFlag.Flag.Usage = func() {
 		printUsage(socksdCommandFlag)
 	}

config/flag.go

@@ -72,10 +72,8 @@ type Config struct {
 	Command                 string           `yaml:"-" json:"-"`
 	FleetAddr               Address          `yaml:"-" json:"-"`
 	RegistryAddr            Address          `yaml:"-" json:"-"`
-	ProxyServerAddr         string           `yaml:"-" json:"-"`
 	ProxyServerHost         string           `yaml:"-" json:"-"`
 	ProxyServerPort         int              `yaml:"-" json:"-"`
-	SProxyServerAddr        string           `yaml:"-" json:"-"`
 	SProxyServerHost        string           `yaml:"-" json:"-"`
 	SProxyServerPort        int              `yaml:"-" json:"-"`
 	SProxyServerCertPath    string           `yaml:"-" json:"-"`
@@ -84,9 +82,9 @@ type Config struct {
 	EnableProxyServer       bool             `yaml:"-" json:"-"`
 	EnableSProxyServer      bool             `yaml:"-" json:"-"`
 	EnableSocksServer       bool             `yaml:"-" json:"-"`
-	SocksServerAddr         string           `yaml:"-" json:"-"`
 	SocksServerHost         string           `yaml:"-" json:"-"`
 	SocksServerPort         int              `yaml:"-" json:"-"`
+	SocksFallback           string           `yaml:"-" json:"-"`
 	ConfigUnsafe            bool             `yaml:"-" json:"-"`
 	ConfigList              bool             `yaml:"-" json:"-"`
 	ConfigDelete            stringValues     `yaml:"-" json:"-"`
@@ -429,15 +427,9 @@ func ParseFlag() {
 	case "socksd":
 		commandFlag.Parse(args[1:])
 		cfg.EnableSocksServer = true
-		cfg.SocksServerAddr = fmt.Sprintf("%s:%d", cfg.SocksServerHost, cfg.SocksServerPort)
 	case "httpd":
 		commandFlag.Parse(args[1:])
 		cfg.EnableProxyServer = true
-		cfg.SocksServerAddr = fmt.Sprintf("%s:%d", cfg.SocksServerHost, cfg.SocksServerPort)
-		cfg.ProxyServerAddr = fmt.Sprintf("%s:%d", cfg.ProxyServerHost, cfg.ProxyServerPort)
-		if cfg.EnableSProxyServer {
-			cfg.SProxyServerAddr = fmt.Sprintf("%s:%d", cfg.SProxyServerHost, cfg.SProxyServerPort)
-		}
 	case "publish":
 		commandFlag.Parse(args[1:])
 		publishedPorts := make(map[int]*Port)
@@ -552,3 +544,15 @@ func ParseFlag() {
 	AppConfig = cfg
 	// return cfg
 }
+
+func (cfg *Config) SocksServerAddr() string {
+	return fmt.Sprintf("%s:%d", cfg.SocksServerHost, cfg.SocksServerPort)
+}
+
+func (cfg *Config) ProxyServerAddr() string {
+	return fmt.Sprintf("%s:%d", cfg.ProxyServerHost, cfg.ProxyServerPort)
+}
+
+func (cfg *Config) SProxyServerAddr() string {
+	return fmt.Sprintf("%s:%d", cfg.SProxyServerHost, cfg.SProxyServerPort)
+}

rpc/datapool.go

@@ -27,14 +27,6 @@ func NewPool() *DataPool {
 	}
 }
 
-func NewPoolWithPublishedPorts(publishedPorts map[int]*config.Port) *DataPool {
-	return &DataPool{
-		memoryCache:    cache.New(5*time.Minute, 10*time.Minute),
-		devices:        make(map[string]*ConnectedDevice),
-		publishedPorts: publishedPorts,
-	}
-}
-
 func (p *DataPool) GetCacheDNS(key string) (dns Address, ok bool) {
 	p.rm.RLock()
 	defer p.rm.RUnlock()
@@ -123,12 +115,8 @@ func (p *DataPool) GetPublishedPort(port int) *config.Port {
 	return p.publishedPorts[port]
 }
 
-func (p *DataPool) SetPublishedPort(port int, publishedPort *config.Port) {
+func (p *DataPool) SetPublishedPorts(ports map[int]*config.Port) {
 	p.rm.Lock()
 	defer p.rm.Unlock()
-	if publishedPort == nil {
-		delete(p.publishedPorts, port)
-	} else {
-		p.publishedPorts[port] = publishedPort
-	}
+	p.publishedPorts = ports
 }

rpc/proxy.go

@@ -181,25 +181,41 @@ func (proxyServer *ProxyServer) pipeProxy(w http.ResponseWriter, r *http.Request
 	connDevice.Close()
 }
 
-func NewProxyServer(socksServer *Server, config ProxyConfig) (*ProxyServer, error) {
-	if socksServer == nil {
-		return nil, fmt.Errorf("socks server should not be nil")
-	}
-	if !socksServer.Started() {
-		return nil, fmt.Errorf("should start socks server first")
+func NewProxyServer(socksServer *Server) *ProxyServer {
+	proxyServer := &ProxyServer{
+		socksServer: socksServer,
 	}
+	return proxyServer
+}
+
+func (proxyServer *ProxyServer) SetConfig(config ProxyConfig) error {
 	if config.AllowRedirect && !config.EnableSProxy {
-		return nil, fmt.Errorf("wrong parameters, need started httpsd server for http redirect")
+		return fmt.Errorf("wrong parameters, need started httpsd server for http redirect")
 	}
-	proxyServer := &ProxyServer{
-		socksServer: socksServer,
-		Config:      config,
+	proxyServer.Config = config
+	return nil
+}
+
+func (proxyServer *ProxyServer) Stop() {
+	if proxyServer.httpServer != nil {
+		proxyServer.httpServer.Close()
+		proxyServer.httpServer = nil
+	}
+	if proxyServer.httpsServer != nil {
+		proxyServer.httpsServer.Close()
+		proxyServer.httpsServer = nil
 	}
-	return proxyServer, nil
 }
 
 func (proxyServer *ProxyServer) Start() error {
 	// start httpd proxy server
+	if proxyServer.socksServer == nil || !proxyServer.socksServer.Started() {
+		return fmt.Errorf("should start socks server first")
+	}
+	if proxyServer.started {
+		return nil
+	}
+	proxyServer.started = true
 	if proxyServer.Config.EnableProxy {
 		proxyServer.socksServer.Client.Info("Start httpd server %s", proxyServer.Config.ProxyServerAddr)
 		prox, _ := url.Parse(fmt.Sprintf("socks5://%s", proxyServer.socksServer.Config.Addr))
@@ -223,8 +239,12 @@ func (proxyServer *ProxyServer) Start() error {
 				proxyServer.httpServer = nil
 				proxyServer.socksServer.Client.Error("cannot start http proxy: %v", err)
 			}
-			proxyServer.started = true
 		}()
+	} else {
+		if proxyServer.httpServer != nil {
+			proxyServer.httpServer.Close()
+			proxyServer.httpServer = nil
+		}
 	}
 
 	// start httpsd proxy server
@@ -240,8 +260,12 @@ func (proxyServer *ProxyServer) Start() error {
 				proxyServer.httpsServer = nil
 				proxyServer.socksServer.Client.Error("cannot start https proxy: %v", err)
 			}
-			proxyServer.started = true
 		}()
+	} else {
+		if proxyServer.httpsServer != nil {
+			proxyServer.httpsServer.Close()
+			proxyServer.httpsServer = nil
+		}
 	}
 	return nil
 }