dmsgweb resolve specific address (#252)

* resolve a specific dmsg address locally * make format
skycoin · Feb 24, 2024 · 5632510 · 5632510
1 parent ecde4c5
commit 5632510
Showing 1 changed file with 58 additions and 49 deletions.
diff --git a/cmd/dmsgweb/commands/dmsgweb.go b/cmd/dmsgweb/commands/dmsgweb.go
@@ -72,6 +72,7 @@ var (
 	webPort            string
 	proxyPort          string
 	addProxy           string
+	resolveDmsgAddr    string
 	wg                 sync.WaitGroup
 )
 
@@ -81,6 +82,7 @@ func init() {
 	RootCmd.Flags().StringVarP(&proxyPort, "socks", "q", "4445", "port to serve the socks5 proxy")
 	RootCmd.Flags().StringVarP(&addProxy, "proxy", "r", "", "configure additional socks5 proxy for dmsgweb (i.e. 127.0.0.1:1080)")
 	RootCmd.Flags().StringVarP(&webPort, "port", "p", "8080", "port to serve the web application")
+	RootCmd.Flags().StringVarP(&resolveDmsgAddr, "resolve", "t", "", "resolve the specified dmsg address:port on the local port & disable proxy")
 	RootCmd.Flags().StringVarP(&dmsgDisc, "dmsg-disc", "d", "", "dmsg discovery url default:\n"+skyenv.DmsgDiscAddr)
 	RootCmd.Flags().IntVarP(&dmsgSessions, "sess", "e", 1, "number of dmsg servers to connect to")
 	RootCmd.Flags().StringVarP(&logLvl, "loglvl", "l", "", "[ debug | warn | error | fatal | panic | trace | info ]\033[0m")
@@ -149,72 +151,79 @@ var RootCmd = &cobra.Command{
 
 		httpC = http.Client{Transport: dmsghttp.MakeHTTPTransport(ctx, dmsgC)}
 
-		// Create a SOCKS5 server with custom name resolution
-		conf := &socks5.Config{
-			Resolver: &customResolver{},
-			Dial: func(ctx context.Context, network, addr string) (net.Conn, error) {
-				host, _, err := net.SplitHostPort(addr)
-				if err != nil {
-					return nil, err
-				}
-				regexPattern := `\` + filterDomainSuffix + `(:[0-9]+)?$`
-				match, _ := regexp.MatchString(regexPattern, host) //nolint:errcheck
-				if match {
-					port, ok := ctx.Value("port").(string)
-					if !ok {
-						port = webPort
+		if resolveDmsgAddr == "" {
+			// Create a SOCKS5 server with custom name resolution
+			conf := &socks5.Config{
+				Resolver: &customResolver{},
+				Dial: func(ctx context.Context, network, addr string) (net.Conn, error) {
+					host, _, err := net.SplitHostPort(addr)
+					if err != nil {
+						return nil, err
 					}
-					addr = "localhost:" + port
-				} else {
-					if addProxy != "" {
-						// Fallback to another SOCKS5 proxy
-						dialer, err := proxy.SOCKS5("tcp", addProxy, nil, proxy.Direct)
-						if err != nil {
-							return nil, err
+					regexPattern := `\` + filterDomainSuffix + `(:[0-9]+)?$`
+					match, _ := regexp.MatchString(regexPattern, host) //nolint:errcheck
+					if match {
+						port, ok := ctx.Value("port").(string)
+						if !ok {
+							port = webPort
+						}
+						addr = "localhost:" + port
+					} else {
+						if addProxy != "" {
+							// Fallback to another SOCKS5 proxy
+							dialer, err := proxy.SOCKS5("tcp", addProxy, nil, proxy.Direct)
+							if err != nil {
+								return nil, err
+							}
+							return dialer.Dial(network, addr)
 						}
-						return dialer.Dial(network, addr)
 					}
-				}
-				dmsgWebLog.Debug("Dialing address:", addr)
-				return net.Dial(network, addr)
-			},
-		}
-
-		// Start the SOCKS5 server
-		socksAddr := "127.0.0.1:" + proxyPort
-		log.Printf("SOCKS5 proxy server started on %s", socksAddr)
+					dmsgWebLog.Debug("Dialing address:", addr)
+					return net.Dial(network, addr)
+				},
+			}
 
-		server, err := socks5.New(conf)
-		if err != nil {
-			log.Fatalf("Failed to create SOCKS5 server: %v", err)
-		}
+			// Start the SOCKS5 server
+			socksAddr := "127.0.0.1:" + proxyPort
+			log.Printf("SOCKS5 proxy server started on %s", socksAddr)
 
-		wg.Add(1)
-		go func() {
-			dmsgWebLog.Debug("Serving SOCKS5 proxy on " + socksAddr)
-			err := server.ListenAndServe("tcp", socksAddr)
+			server, err := socks5.New(conf)
 			if err != nil {
-				log.Fatalf("Failed to start SOCKS5 server: %v", err)
+				log.Fatalf("Failed to create SOCKS5 server: %v", err)
 			}
-			defer server.Close()
-			dmsgWebLog.Debug("Stopped serving SOCKS5 proxy on " + socksAddr)
-		}()
 
+			wg.Add(1)
+			go func() {
+				dmsgWebLog.Debug("Serving SOCKS5 proxy on " + socksAddr)
+				err := server.ListenAndServe("tcp", socksAddr)
+				if err != nil {
+					log.Fatalf("Failed to start SOCKS5 server: %v", err)
+				}
+				defer server.Close()
+				dmsgWebLog.Debug("Stopped serving SOCKS5 proxy on " + socksAddr)
+			}()
+		}
 		r := gin.New()
 
 		r.Use(gin.Recovery())
 
 		r.Use(loggingMiddleware())
 
 		r.Any("/*path", func(c *gin.Context) {
-			hostParts := strings.Split(c.Request.Host, ":")
-			var dmsgp string
-			if len(hostParts) > 1 {
-				dmsgp = hostParts[1]
+			var urlStr string
+			if resolveDmsgAddr != "" {
+				urlStr = fmt.Sprintf("dmsg://%s%s", resolveDmsgAddr, c.Param("path"))
 			} else {
-				dmsgp = "80"
+
+				hostParts := strings.Split(c.Request.Host, ":")
+				var dmsgp string
+				if len(hostParts) > 1 {
+					dmsgp = hostParts[1]
+				} else {
+					dmsgp = "80"
+				}
+				urlStr = fmt.Sprintf("dmsg://%s:%s%s", strings.TrimRight(hostParts[0], filterDomainSuffix), dmsgp, c.Param("path"))
 			}
-			urlStr := fmt.Sprintf("dmsg://%s:%s%s", strings.TrimRight(hostParts[0], filterDomainSuffix), dmsgp, c.Param("path"))
 
 			maxSize := int64(1024)