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

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

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 {

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)

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{},

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