From d9f2c85c275f8f8ba37f806ed4ed530fdab298bc Mon Sep 17 00:00:00 2001
From: AJ ONeal <coolaj86@gmail.com>
Date: Sat, 18 Jul 2020 00:18:09 -0600
Subject: [PATCH] move tcp muxing to own function

---
 cmd/telebit/telebit.go | 88 +++++++++++++++++++++++-------------------
 1 file changed, 49 insertions(+), 39 deletions(-)

diff --git a/cmd/telebit/telebit.go b/cmd/telebit/telebit.go
index 7cbbdb1..9bf65d6 100644
--- a/cmd/telebit/telebit.go
+++ b/cmd/telebit/telebit.go
@@ -287,45 +287,7 @@ func main() {
 		EnableTLSALPNChallenge: *enableTLSALPN01,
 	}
 
-	//mux := telebit.NewRouteMux(acme)
-	mux := telebit.NewRouteMux()
-
-	// Port forward without TerminatingTLS
-	for _, fwd := range portForwards {
-		msg := fmt.Sprintf("Fwd: %s %s", fwd.pattern, fwd.port)
-		fmt.Println(msg)
-		mux.ForwardTCP(fwd.pattern, "localhost:"+fwd.port, 120*time.Second, msg, "[Port Forward]")
-	}
-	// TODO close connection on invalid hostname
-	mux.HandleTCP("*", telebit.HandlerFunc(routeSubscribersAndClients), "[Tun => Remote Servers]")
-	mux.HandleTLS("*", acme, mux, "[Terminate TLS & Recurse]")
-
-	if 0 == len(*apiHostname) {
-		*apiHostname = os.Getenv("API_HOSTNAME")
-	}
-	if "" != *apiHostname {
-		listener := httpshim.NewListener()
-		go func() {
-			httpsrv.Serve(listener)
-		}()
-		fmt.Printf("Will respond to Websocket and API requests to %q\n", *apiHostname)
-		mux.HandleTCP(*apiHostname, telebit.HandlerFunc(func(client net.Conn) error {
-			if dbg.Debug {
-				fmt.Printf("[debug] Accepting API or WebSocket client %q\n", *apiHostname)
-			}
-			listener.Feed(client)
-			if dbg.Debug {
-				fmt.Printf("[debug] done with %q client\n", *apiHostname)
-			}
-			// nil now means handler in-progress (go routine)
-			// EOF now means handler finished
-			return nil
-		}), "[Admin API & Server Relays]")
-	}
-	for _, fwd := range forwards {
-		//mux.ForwardTCP("*", "localhost:"+fwd.port, 120*time.Second)
-		mux.ForwardTCP(fwd.pattern, "localhost:"+fwd.port, 120*time.Second, "[Servername Forward]")
-	}
+	mux := muxAll(portForwards, forwards, acme, apiHostname)
 
 	done := make(chan error)
 	if dbg.Debug {
@@ -396,6 +358,54 @@ func main() {
 	}
 }
 
+func muxAll(portForwards, forwards []Forward, acme *telebit.ACME, apiHostname *string) *telebit.RouteMux {
+	//mux := telebit.NewRouteMux(acme)
+	mux := telebit.NewRouteMux()
+
+	// Port forward without TerminatingTLS
+	for _, fwd := range portForwards {
+		msg := fmt.Sprintf("Fwd: %s %s", fwd.pattern, fwd.port)
+		fmt.Println(msg)
+		mux.ForwardTCP(fwd.pattern, "localhost:"+fwd.port, 120*time.Second, msg, "[Port Forward]")
+	}
+	// TODO close connection on invalid hostname
+	mux.HandleTCP("*", telebit.HandlerFunc(routeSubscribersAndClients), "[Tun => Remote Servers]")
+	mux.HandleTLS("*", acme, mux, "[Terminate TLS & Recurse]")
+	//mux.HandleTLSFunc(func (sni) bool {
+	//	// do whatever
+	//	return false
+	//}, acme, mux, "[Terminate TLS & Recurse]")
+
+	if 0 == len(*apiHostname) {
+		*apiHostname = os.Getenv("API_HOSTNAME")
+	}
+	if "" != *apiHostname {
+		listener := httpshim.NewListener()
+		go func() {
+			httpsrv.Serve(listener)
+		}()
+		fmt.Printf("Will respond to Websocket and API requests to %q\n", *apiHostname)
+		mux.HandleTCP(*apiHostname, telebit.HandlerFunc(func(client net.Conn) error {
+			if dbg.Debug {
+				fmt.Printf("[debug] Accepting API or WebSocket client %q\n", *apiHostname)
+			}
+			listener.Feed(client)
+			if dbg.Debug {
+				fmt.Printf("[debug] done with %q client\n", *apiHostname)
+			}
+			// nil now means handler in-progress (go routine)
+			// EOF now means handler finished
+			return nil
+		}), "[Admin API & Server Relays]")
+	}
+	for _, fwd := range forwards {
+		//mux.ForwardTCP("*", "localhost:"+fwd.port, 120*time.Second)
+		mux.ForwardTCP(fwd.pattern, "localhost:"+fwd.port, 120*time.Second, "[Servername Forward]")
+	}
+
+	return mux
+}
+
 func routeSubscribersAndClients(client net.Conn) error {
 	var wconn *telebit.ConnWrap
 	switch conn := client.(type) {