From 50dd46351db34f2a9817dbf049390c0171369608 Mon Sep 17 00:00:00 2001
From: Ananth <antsub@gmail.com>
Date: Sun, 14 Jul 2024 23:56:10 +0530
Subject: [PATCH] feat: Control URL flag (#5)

Use a different control server by setting the -control-url flag or the TS_CONTROL_URL environment variable.
---
 README.md |  6 +++++-
 main.go   | 64 ++++++++++++++++++++++++++++---------------------------
 2 files changed, 38 insertions(+), 32 deletions(-)

diff --git a/README.md b/README.md
index d3c3b44..35e2b6a 100644
--- a/README.md
+++ b/README.md
@@ -43,7 +43,11 @@ docker run -d \
     -v /path/to/state:/state \
     -e TSNS_HOSTNAME=<hostname> \
     -e TSNS_BACKEND=<backend> \
-    -e TSNS_FUNNEL=true \ # optional, enables Tailscale Funnel
+    # optional, enables Tailscale Funnel
+    # -e TSNS_FUNNEL=true \
+    # optional, set to your control URL
+    # leave empty for default
+    # -e TS_CONTROL_URL=https://your.tailscale.com \
     -e TS_AUTHKEY=<auth key> \
     ghcr.io/shayne/tsnet-serve:latest
 ```
diff --git a/main.go b/main.go
index bc813ae..4187c9a 100644
--- a/main.go
+++ b/main.go
@@ -30,6 +30,7 @@ var (
 	funnel     = flag.Bool("funnel", false, "enable funnel mode")
 	mountPath  = flag.String("mount-path", "/", "path to mount proxy on")
 	stateDir   = flag.String("state-dir", "/state", "directory to store state in")
+	controlURL = flag.String("control-url", "", "control URL to use, leave empty for default")
 	version    = flag.Bool("version", false, "print version and exit")
 )
 
@@ -51,36 +52,18 @@ func main() {
 		return
 	}
 
-	if os.Getenv("TSNS_HOSTNAME") != "" {
-		*hostname = os.Getenv("TSNS_HOSTNAME")
+	if h := os.Getenv("TSNS_HOSTNAME"); h != "" {
+		*hostname = h
 	}
-
-	if os.Getenv("TSNS_BACKEND") != "" {
-		*backend = os.Getenv("TSNS_BACKEND")
-	}
-
-	if os.Getenv("TSNS_LISTEN_PORT") != "" {
-		p, err := strconv.Atoi(os.Getenv("TSNS_LISTEN_PORT"))
-		if err != nil {
-			log.Fatalf("invalid TSNS_LISTEN_PORT: %v", err)
-		}
-		*listenPort = p
-	}
-
-	if os.Getenv("TSNS_MOUNT_PATH") != "" {
-		*mountPath = os.Getenv("TSNS_MOUNT_PATH")
-	}
-
-	if os.Getenv("TSNS_STATE_DIR") != "" {
-		*stateDir = os.Getenv("TSNS_STATE_DIR")
-	}
-
 	if *hostname == "" {
 		log.Fatal("hostname is required")
 	}
 
-	if *listenPort < 1 || *listenPort > 65535 {
-		log.Fatal("invalid port")
+	if b := os.Getenv("TSNS_BACKEND"); b != "" {
+		*backend = b
+	}
+	if *backend == "" {
+		log.Fatal("backend is required")
 	}
 
 	if fn := os.Getenv("TSNS_FUNNEL"); fn != "" {
@@ -91,16 +74,30 @@ func main() {
 		*funnel = ok
 	}
 
+	if p := os.Getenv("TSNS_LISTEN_PORT"); p != "" {
+		pr, err := strconv.Atoi(p)
+		if err != nil {
+			log.Fatalf("invalid TSNS_LISTEN_PORT: %v", err)
+		}
+		*listenPort = pr
+	}
+	if *listenPort < 1 || *listenPort > 65535 {
+		log.Fatal("invalid port")
+	}
 	if *funnel && *listenPort != 443 && *listenPort != 8443 && *listenPort != 10000 {
 		log.Fatal("funnel mode is only available on port 443, 8443, or 10000")
 	}
 	portStr := strconv.Itoa(*listenPort)
 
-	if *backend == "" {
-		log.Fatal("backend is required")
+	if p := os.Getenv("TSNS_MOUNT_PATH"); p != "" {
+		*mountPath = p
 	}
 
-	proxyTarget, err := expandProxyTarget(*backend)
+	if d := os.Getenv("TSNS_STATE_DIR"); d != "" {
+		*stateDir = d
+	}
+
+	proxyTarget, err := validateProxyTarget(*backend)
 	if err != nil {
 		log.Fatalf("invalid backend: %v", err)
 	}
@@ -117,9 +114,14 @@ func main() {
 		log.Fatalf("state directory is not writable")
 	}
 
+	if u := os.Getenv("TS_CONTROL_URL"); u != "" {
+		*controlURL = u
+	}
+
 	s := &tsnet.Server{
-		Hostname: *hostname,
-		Dir:      *stateDir,
+		Hostname:   *hostname,
+		Dir:        *stateDir,
+		ControlURL: *controlURL,
 	}
 	defer s.Close()
 
@@ -196,7 +198,7 @@ func startServer(ctx context.Context, s *tsnet.Server, portStr, proxyTarget stri
 	return lc.SetServeConfig(ctx, srvConfig)
 }
 
-func expandProxyTarget(source string) (*url.URL, error) {
+func validateProxyTarget(source string) (*url.URL, error) {
 	if !strings.Contains(source, "://") {
 		source = "http://" + source
 	}