Add --slug flag to reuse a specific subdomain across sessions

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

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)