From 803b4049a6f48e6cd6e16e6246175f7caf688bfa Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Jose=20Luis=20Monta=C3=B1es=20Ojados?= Date: Mon, 9 Feb 2026 01:27:17 +0100 Subject: [PATCH] Add --slug flag to reuse a specific subdomain across sessions Co-Authored-By: Claude Opus 4.6 --- cmd/client/main.go | 3 ++- cmd/server/main.go | 54 +++++++++++++++++++++++++++------------ internal/tui/model.go | 4 +-- internal/tunnel/client.go | 38 +++++++++++++++++++-------- 4 files changed, 68 insertions(+), 31 deletions(-) diff --git a/cmd/client/main.go b/cmd/client/main.go index 23554d0..97aa011 100644 --- a/cmd/client/main.go +++ b/cmd/client/main.go @@ -37,6 +37,7 @@ func main() { tokenFlag := flag.String("token", "", "Authentication token (overrides config)") hostHeaderFlag := flag.String("host-header", "", "Custom Host header to send to local service") localHttpsFlag := flag.Bool("local-https", false, "Use HTTPS to connect to local service (implied if port is 443)") + slugFlag := flag.String("slug", "", "Request a specific subdomain slug (e.g., myapp)") flag.Parse() // Load config @@ -58,7 +59,7 @@ func main() { serverAddr = "localhost:2222" } - m := tui.InitialModel(*localPort, serverAddr, authToken, *hostHeaderFlag, *localHttpsFlag) + m := tui.InitialModel(*localPort, serverAddr, authToken, *hostHeaderFlag, *localHttpsFlag, *slugFlag) p := tea.NewProgram(m, tea.WithAltScreen()) if _, err := p.Run(); err != nil { diff --git a/cmd/server/main.go b/cmd/server/main.go index 32991cc..dd85aa2 100644 --- a/cmd/server/main.go +++ b/cmd/server/main.go @@ -8,6 +8,7 @@ import ( "math/rand" "net/http" "os" + "regexp" "strings" "sync" "time" @@ -44,6 +45,13 @@ var manager = &TunnelManager{ tunnels: make(map[string]*Tunnel), } +// requestedSlugs stores slugs requested by clients before tcpip-forward +var ( + requestedSlugs = make(map[*gossh.ServerConn]string) + requestedSlugsMu sync.Mutex + slugPattern = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{1,30}[a-z0-9]$`) +) + func (tm *TunnelManager) Register(id string, t *Tunnel) { tm.mu.Lock() defer tm.mu.Unlock() @@ -305,28 +313,22 @@ func main() { // This is where clients say "Please listen on port X" sshServer.RequestHandlers = map[string]ssh.RequestHandler{ "tcpip-forward": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) { - // Parse payload - // string address to bind (usually empty or 0.0.0.0 or 127.0.0.1) - // uint32 port number to bind + conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn) - // For Grokway, we ignore the requested port and assign a random subdomain/slug - // Or we use the requested port if valid? - // Let's assume we ignore it and generate a slug, - // OR we use the port as the "slug" if it's special. + // Check if client requested a specific slug + requestedSlugsMu.Lock() + slug, hasRequested := requestedSlugs[conn] + if hasRequested { + delete(requestedSlugs, conn) + } + requestedSlugsMu.Unlock() - // But wait, the client needs to know what we assigned! - // Standard SSH response to tcpip-forward contains the BOUND PORT. - // If we return a port, the client knows. - // If we want to return a URL, standard SSH doesn't have a field for that. - // But we can write it to the Session (stdout). - - // Let's accept any request and map it to a random slug. - slug := generateSlug(8) + if !hasRequested { + slug = generateSlug(8) + } log.Printf("Client requested forwarding. Assigning slug: %s", slug) - conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn) - manager.Register(slug, &Tunnel{ ID: slug, LocalPort: 80, // We assume client forwards to port 80 locally? No, @@ -351,6 +353,24 @@ func main() { "cancel-tcpip-forward": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) { return true, nil }, + "grokway-request-slug": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) { + slug := string(req.Payload) + if !slugPattern.MatchString(slug) { + log.Printf("Invalid slug format requested: %q", slug) + return false, []byte("invalid slug format (lowercase alphanumeric and hyphens, 3-32 chars)") + } + // Check if slug is already in use + if _, taken := manager.Get(slug); taken { + log.Printf("Slug %q already in use", slug) + return false, []byte("slug already in use") + } + conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn) + requestedSlugsMu.Lock() + requestedSlugs[conn] = slug + requestedSlugsMu.Unlock() + log.Printf("Client requested slug: %s", slug) + return true, nil + }, "grokway-whoami": func(ctx ssh.Context, srv *ssh.Server, req *gossh.Request) (bool, []byte) { conn := ctx.Value(ssh.ContextKeyConn).(*gossh.ServerConn) slug, ok := manager.FindSlugByConn(conn) diff --git a/internal/tui/model.go b/internal/tui/model.go index 59e67ba..ba1426c 100644 --- a/internal/tui/model.go +++ b/internal/tui/model.go @@ -68,8 +68,8 @@ type LogMsg string type MetricMsg int64 type ClearCopiedMsg struct{} -func InitialModel(localPort, serverAddr, authToken, hostHeader string, localHTTPS bool) Model { - c := tunnel.NewClient(serverAddr, localPort, authToken, hostHeader, localHTTPS) +func InitialModel(localPort, serverAddr, authToken, hostHeader string, localHTTPS bool, slug string) Model { + c := tunnel.NewClient(serverAddr, localPort, authToken, hostHeader, localHTTPS, slug) return Model{ Client: c, LogLines: []string{}, diff --git a/internal/tunnel/client.go b/internal/tunnel/client.go index 2a66a06..7c8bd50 100644 --- a/internal/tunnel/client.go +++ b/internal/tunnel/client.go @@ -21,26 +21,28 @@ func logToFile(msg string) { // Client handles the SSH connection and forwarding type Client struct { - ServerAddr string - LocalPort string - AuthToken string - HostHeader string // New field for custom Host header - LocalHTTPS bool // New field: connect to local service with HTTPS - SSHClient *ssh.Client - Listener net.Listener - Events chan string // Channel to send logs/events to TUI - Metrics chan int64 // Channel to send bytes transferred - PublicURL string // PublicURL is the URL accessible from the internet + ServerAddr string + LocalPort string + AuthToken string + HostHeader string // Custom Host header + LocalHTTPS bool // Connect to local service with HTTPS + Slug string // Requested slug (empty = server assigns random) + SSHClient *ssh.Client + Listener net.Listener + Events chan string // Channel to send logs/events to TUI + Metrics chan int64 // Channel to send bytes transferred + PublicURL string // PublicURL is the URL accessible from the internet stopKeepAlive chan struct{} // Signal to stop keepalive goroutine } -func NewClient(serverAddr, localPort, authToken, hostHeader string, localHTTPS bool) *Client { +func NewClient(serverAddr, localPort, authToken, hostHeader string, localHTTPS bool, slug string) *Client { return &Client{ ServerAddr: serverAddr, LocalPort: localPort, AuthToken: authToken, HostHeader: hostHeader, LocalHTTPS: localHTTPS, + Slug: slug, Events: make(chan string, 10), Metrics: make(chan int64, 10), } @@ -64,6 +66,20 @@ func (c *Client) connect() error { c.SSHClient = client c.Events <- "SSH Connected!" + // Request a specific slug if provided + if c.Slug != "" { + ok, reply, err := client.SendRequest("grokway-request-slug", true, []byte(c.Slug)) + if err != nil || !ok { + reason := string(reply) + if reason == "" && err != nil { + reason = err.Error() + } + c.Events <- fmt.Sprintf("Slug %q not available (%s), server will assign one", c.Slug, reason) + } else { + c.Events <- fmt.Sprintf("Slug %q reserved", c.Slug) + } + } + // Request remote listening (Reverse Forwarding) // Bind to 0.0.0.0 on server, random port (0) listener, err := client.Listen("tcp", "0.0.0.0:0")