feat: general protocol rewrite

cmd/server/main.go

@@ -2,12 +2,9 @@ package main
 
 import (
 	"context"
-	"net"
+	"flag"
 	"net/http"
-	"net/http/httputil"
-	"net/url"
 	"strings"
-	"time"
 
 	"cdr.dev/slog"
 	"forge.cadoles.com/wpetit/go-tunnel"
@@ -15,20 +12,61 @@ import (
 	"gitlab.com/wpetit/goweb/logger"
 )
 
-const sharedKey = "go-tunnel"
 const salt = "go-tunnel"
 
 var registry = NewRegistry()
 
 func main() {
+	var (
+		serverAddr = ":36543"
+		httpAddr   = ":3003"
+		sharedKey  = "go-tunnel"
+		targetURL  = "https://arcad.games"
+	)
+
+	flag.StringVar(&serverAddr, "server-addr", serverAddr, "server address")
+	flag.StringVar(&targetURL, "target-url", targetURL, "target url")
+	flag.StringVar(&httpAddr, "http-addr", httpAddr, "http server address")
+	flag.StringVar(&sharedKey, "shared-key", sharedKey, "shared key")
+
+	flag.Parse()
+
 	ctx := context.Background()
 
 	logger.SetLevel(slog.LevelDebug)
 
 	server := tunnel.NewServer(
-		tunnel.WithServerAESBlockCrypt(sharedKey, salt),
+		tunnel.WithServerAddress(serverAddr),
 		tunnel.WithServerOnClientAuth(registry.OnClientAuth),
 		tunnel.WithServerOnClientDisconnect(registry.OnClientDisconnect),
+		tunnel.WithServerOnClientAuth(func(ctx context.Context, remoteClient *tunnel.RemoteClient, credentials interface{}) (bool, error) {
+			remoteAddr := remoteClient.RemoteAddr().String()
+
+			ctx = logger.With(ctx, logger.F("remoteAddr", remoteAddr))
+
+			logger.Debug(ctx, "new client auth")
+
+			clientID, ok := credentials.(string)
+			if !ok {
+				logger.Debug(ctx, "client auth failed")
+
+				return false, nil
+			}
+
+			registry.Add(clientID, remoteAddr, remoteClient)
+
+			logger.Debug(ctx, "client auth success")
+
+			return true, nil
+		}),
+		tunnel.WithServerOnClientDisconnect(func(ctx context.Context, remoteClient *tunnel.RemoteClient) error {
+			remoteAddr := remoteClient.RemoteAddr().String()
+			ctx = logger.With(ctx, logger.F("remoteAddr", remoteAddr))
+			logger.Debug(ctx, "client disconnected")
+			registry.RemoveByRemoteAddr(remoteAddr)
+
+			return nil
+		}),
 	)
 
 	go func() {
@@ -37,51 +75,28 @@ func main() {
 		}
 	}()
 
-	if err := http.ListenAndServe(":3003", http.HandlerFunc(handleRequest)); err != nil {
+	handler, err := createProxyHandler(targetURL)
+	if err != nil {
+		logger.Fatal(ctx, "could not create proxy handler", logger.E(errors.WithStack(err)))
+	}
+
+	if err := http.ListenAndServe(httpAddr, handler); err != nil {
 		logger.Fatal(ctx, "error while listening", logger.E(err))
 	}
 }
 
-func handleRequest(w http.ResponseWriter, r *http.Request) {
-	subdomains := strings.SplitN(r.Host, ".", 2)
+func createProxyHandler(targetURL string) (http.Handler, error) {
+	return tunnel.ProxyHandler(targetURL, func(w http.ResponseWriter, r *http.Request) (*tunnel.RemoteClient, error) {
+		subdomains := strings.SplitN(r.Host, ".", 2)
 
-	if len(subdomains) < 2 {
-		http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
+		if len(subdomains) < 2 {
+			http.Error(w, http.StatusText(http.StatusNotFound), http.StatusNotFound)
 
-		return
-	}
+			return nil, tunnel.ErrAbortProxy
+		}
 
-	clientID := subdomains[0]
-	remoteClient := registry.Get(clientID)
+		clientID := subdomains[0]
 
-	if remoteClient == nil {
-		http.Error(w, http.StatusText(http.StatusBadGateway), http.StatusBadGateway)
-
-		return
-	}
-
-	target, err := url.Parse("https://arcad.games")
-	if err != nil {
-		logger.Fatal(r.Context(), "could not parse url", logger.E(err))
-	}
-
-	reverse := httputil.NewSingleHostReverseProxy(target)
-	reverse.Transport = &http.Transport{
-		Proxy: http.ProxyFromEnvironment,
-		DialContext: func(ctx context.Context, network string, addr string) (net.Conn, error) {
-			conn, err := remoteClient.Proxy(ctx, network, addr)
-			if err != nil {
-				return nil, errors.WithStack(err)
-			}
-
-			return conn, nil
-		},
-		ForceAttemptHTTP2:     true,
-		MaxIdleConns:          100,
-		IdleConnTimeout:       90 * time.Second,
-		TLSHandshakeTimeout:   10 * time.Second,
-		ExpectContinueTimeout: 1 * time.Second,
-	}
-
-	reverse.ServeHTTP(w, r)
+		return registry.Get(clientID), nil
+	})
 }