go-tunnel/cmd/server/main.go

package main

import (
	"context"
	"flag"
	"net/http"
	"strings"

	"cdr.dev/slog"
	"forge.cadoles.com/wpetit/go-tunnel"
	"github.com/pkg/errors"
	"gitlab.com/wpetit/goweb/logger"
)

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.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() {
		if err := server.Listen(ctx); err != nil {
			logger.Fatal(ctx, "error while listening", logger.E(err))
		}
	}()

	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 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)

			return nil, tunnel.ErrAbortProxy
		}

		clientID := subdomains[0]

		return registry.Get(clientID), nil
	})
}