feat(security): add auth token flag to protect self-hosted server

internal/cli/config_test.go

@@ -1 +1,92 @@
 package cli
+
+import (
+	"os"
+	"strings"
+	"testing"
+
+	"github.com/kernelshard/expose/internal/config"
+)
+
+func TestConfigCmd_Structure(t *testing.T) {
+	cmd := newConfigCmd()
+	if cmd.Name() != "config" {
+		t.Errorf("expected command name 'config', got %q", cmd.Name())
+	}
+
+	subCommands := cmd.Commands()
+	if len(subCommands) != 2 {
+		t.Fatalf("expected 2 subcommands, got %d", len(subCommands))
+	}
+}
+
+func TestConfigListCmd(t *testing.T) {
+	t.Run("fails when config does not exist", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		prevDir, _ := os.Getwd()
+		_ = os.Chdir(tmpDir)
+		defer os.Chdir(prevDir)
+
+		cmd := newConfigListCmd()
+		err := cmd.RunE(cmd, nil)
+		if err == nil || !strings.Contains(err.Error(), "config not found") {
+			t.Fatalf("expected 'config not found' error, got: %v", err)
+		}
+	})
+
+	t.Run("succeeds when config exists", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		prevDir, _ := os.Getwd()
+		_ = os.Chdir(tmpDir)
+		defer os.Chdir(prevDir)
+
+		_, err := config.Init()
+		if err != nil {
+			t.Fatalf("failed to init config: %v", err)
+		}
+
+		cmd := newConfigListCmd()
+		if err := cmd.RunE(cmd, nil); err != nil {
+			t.Fatalf("expected config list to succeed, got: %v", err)
+		}
+	})
+}
+
+func TestConfigGetCmd(t *testing.T) {
+	t.Run("fails when config does not exist", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		prevDir, _ := os.Getwd()
+		_ = os.Chdir(tmpDir)
+		defer os.Chdir(prevDir)
+
+		cmd := newConfigGetCmd()
+		err := cmd.RunE(cmd, []string{"port"})
+		if err == nil || !strings.Contains(err.Error(), "config not found") {
+			t.Fatalf("expected 'config not found' error, got: %v", err)
+		}
+	})
+
+	t.Run("succeeds for valid key and fails for invalid key", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		prevDir, _ := os.Getwd()
+		_ = os.Chdir(tmpDir)
+		defer os.Chdir(prevDir)
+
+		_, err := config.Init()
+		if err != nil {
+			t.Fatalf("failed to init config: %v", err)
+		}
+
+		cmd := newConfigGetCmd()
+
+		// Valid key
+		if err := cmd.RunE(cmd, []string{"port"}); err != nil {
+			t.Fatalf("expected get 'port' to succeed, got: %v", err)
+		}
+
+		// Invalid key
+		if err := cmd.RunE(cmd, []string{"nonexistent"}); err == nil {
+			t.Fatal("expected get 'nonexistent' to fail, got nil")
+		}
+	})
+}

internal/cli/init_test.go

@@ -0,0 +1,36 @@
+package cli
+
+import (
+	"os"
+	"testing"
+)
+
+func TestInitCmd(t *testing.T) {
+	cmd := newInitCmd()
+	if cmd == nil {
+		t.Fatal("expected newInitCmd to return non-nil command")
+	}
+	if cmd.Name() != "init" {
+		t.Errorf("expected command name 'init', got %q", cmd.Name())
+	}
+
+	t.Run("successfully creates config and rejects duplicate", func(t *testing.T) {
+		tmpDir := t.TempDir()
+		prevDir, _ := os.Getwd()
+		_ = os.Chdir(tmpDir)
+		defer os.Chdir(prevDir)
+
+		// First execution creates .expose.yml
+		err := cmd.RunE(cmd, nil)
+		if err != nil {
+			t.Fatalf("expected init to succeed, got: %v", err)
+		}
+
+		// Second execution should fail because config already exists
+		err = cmd.RunE(cmd, nil)
+		if err == nil {
+			t.Fatal("expected second init to fail with 'config already exists', got nil")
+		}
+	})
+}
+

internal/cli/root_test.go

@@ -0,0 +1,13 @@
+package cli
+
+import (
+	"testing"
+)
+
+func TestExecute(t *testing.T) {
+	rootCmd.SetArgs([]string{"--help"})
+	if err := Execute(); err != nil {
+		t.Fatalf("expected Execute() with --help to succeed, got: %v", err)
+	}
+}
+

internal/cli/server.go

@@ -31,6 +31,7 @@ func newServerCmd() *cobra.Command {
 	cmd.Flags().IntP("control-port", "c", 7890, "Port for client connections")
 	cmd.Flags().IntP("public-port", "p", 8080, "Port for public HTTP traffic")
 	cmd.Flags().StringP("domain", "d", "localhost", "Domain for public URLs")
+	cmd.Flags().StringP("token", "t", "", "Auth token to protect the server (or EXPOSE_SERVER_TOKEN env)")
 	return cmd
 }
 
@@ -39,8 +40,12 @@ func runServerCmd(cmd *cobra.Command, _ []string) error {
 	domain, _ := cmd.Flags().GetString("domain")
 	controlPort, _ := cmd.Flags().GetInt("control-port")
 	publicPort, _ := cmd.Flags().GetInt("public-port")
+	token, _ := cmd.Flags().GetString("token")
+	if token == "" {
+		token = os.Getenv("EXPOSE_SERVER_TOKEN")
+	}
 
-	srv := server.NewServer(domain, controlPort, publicPort)
+	srv := server.NewServer(domain, controlPort, publicPort).WithToken(token)
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 

internal/cli/server_test.go

@@ -0,0 +1,38 @@
+package cli
+
+import (
+	"testing"
+)
+
+func TestServerCmd(t *testing.T) {
+	cmd := newServerCmd()
+
+	if cmd == nil {
+		t.Fatal("newServerCmd returned nil")
+	}
+
+	if cmd.Name() != "server" {
+		t.Errorf("expected command name 'server', got '%s'", cmd.Name())
+	}
+
+	flags := []struct {
+		name      string
+		shorthand string
+	}{
+		{"control-port", "c"},
+		{"public-port", "p"},
+		{"domain", "d"},
+		{"token", "t"},
+	}
+
+	for _, f := range flags {
+		flag := cmd.Flags().Lookup(f.name)
+		if flag == nil {
+			t.Errorf("flag %q not defined", f.name)
+			continue
+		}
+		if flag.Shorthand != f.shorthand {
+			t.Errorf("flag %q: expected shorthand %q, got %q", f.name, f.shorthand, flag.Shorthand)
+		}
+	}
+}

internal/cli/tunnel.go

@@ -43,6 +43,9 @@ func newTunnelCmd() *cobra.Command {
 	// subdomain flag for custom subdomain request
 	cmd.Flags().StringP("subdomain", "d", "", "Request a custom subdomain (e.g. --subdomain my-app)")
 
+	// token flag for self-hosted server auth
+	cmd.Flags().StringP("token", "t", "", "Auth token for self-hosted server (or EXPOSE_SERVER_TOKEN env)")
+
 	return cmd
 }
 
@@ -109,13 +112,17 @@ func runTunnelCmd(cmd *cobra.Command, args []string) error {
 	}
 
 	subdomain, _ := cmd.Flags().GetString("subdomain")
+	token, _ := cmd.Flags().GetString("token")
+	if token == "" {
+		token = os.Getenv("EXPOSE_SERVER_TOKEN")
+	}
 
 	var svc *tunnel.Service
 	switch providerName {
 	case "cloudflare":
 		svc = tunnel.NewService(provider.NewCloudFlare())
 	case "selfhosted":
-		svc = tunnel.NewService(provider.NewSelfHosted(serverAddr, subdomain))
+		svc = tunnel.NewService(provider.NewSelfHosted(serverAddr, subdomain).WithToken(token))
 	default:
 		svc = tunnel.NewService(provider.NewLocalTunnel(nil))
 	}

internal/cli/tunnel_test.go

@@ -1,6 +1,7 @@
 package cli
 
 import (
+	"net"
 	"os"
 	"testing"
 )
@@ -33,6 +34,14 @@ func TestTunnelCmd(t *testing.T) {
 	if subdomainFlag.Shorthand != "d" {
 		t.Errorf("expected shorthand 'd', got %s", subdomainFlag.Shorthand)
 	}
+
+	tokenFlag := cmd.Flags().Lookup("token")
+	if tokenFlag == nil {
+		t.Error("token flag not defined")
+	}
+	if tokenFlag.Shorthand != "t" {
+		t.Errorf("expected shorthand 't', got %s", tokenFlag.Shorthand)
+	}
 }
 
 func TestResolvePort(t *testing.T) {
@@ -102,3 +111,21 @@ func TestResolvePort(t *testing.T) {
 		}
 	})
 }
+
+func TestIsPortListening(t *testing.T) {
+	ln, err := net.Listen("tcp", "127.0.0.1:0")
+	if err != nil {
+		t.Fatalf("failed to listen: %v", err)
+	}
+	defer ln.Close()
+
+	port := ln.Addr().(*net.TCPAddr).Port
+	if !isPortListening(port) {
+		t.Errorf("expected port %d to be listening", port)
+	}
+
+	_ = ln.Close()
+	if isPortListening(port) {
+		t.Errorf("expected closed port %d to not be listening", port)
+	}
+}

internal/provider/selfhosted.go

@@ -5,6 +5,7 @@ import (
 	"encoding/json"
 	"fmt"
 	"net"
+	"strings"
 	"sync"
 	"time"
 
@@ -21,6 +22,7 @@ const (
 type SelfHosted struct {
 	serverAddr     string // e.g: tunnel.mysite.com:7890
 	subdomain      string
+	token          string
 	conn           net.Conn
 	publicURL      string
 	connected      bool
@@ -39,6 +41,11 @@ func NewSelfHosted(serverAddr string, subdomain string) *SelfHosted {
 	}
 }
 
+func (s *SelfHosted) WithToken(token string) *SelfHosted {
+	s.token = strings.TrimSpace(token)
+	return s
+}
+
 // Name returns the name of the provider.
 func (s *SelfHosted) Name() string {
 	return "selfhosted"
@@ -84,7 +91,9 @@ func (s *SelfHosted) Connect(ctx context.Context, localPort int) (string, error)
 	req := protocol.TunnelRequest{
 		Type:      protocol.TypeControl,
 		Subdomain: s.subdomain,
+		Token:     s.token,
 	}
+
 	if err := json.NewEncoder(conn).Encode(req); err != nil {
 		conn.Close()
 		return "", err

internal/provider/selfhosted_test.go

@@ -361,3 +361,49 @@ func TestSelfHosted_StartHeartBeat_Cancelled(t *testing.T) {
 	cancel()
 	p.startHeartBeat(ctx)
 }
+
+func TestSelfHosted_Connect_TokenAuth_Success(t *testing.T) {
+	const SecretToken = "secret123"
+	srv := server.NewServer("localhost", 0, 0).WithToken(SecretToken)
+	ctx := t.Context()
+	go func() {
+		_ = srv.Start(ctx)
+	}()
+	<-srv.Ready()
+
+	controlAddr := srv.ControlAddr()
+
+	client := NewSelfHosted(controlAddr, "example").WithToken(SecretToken)
+	url, err := client.Connect(ctx, 3000)
+	if err != nil {
+		t.Fatalf("Connect failed: %v", err)
+	}
+	defer client.Close()
+
+	if url == "" {
+		t.Fatal("expected public URL, got empty string")
+	}
+	if !client.IsConnected() {
+		t.Fatal("expected IsConnected to be true")
+	}
+}
+
+func TestSelfHosted_Connect_TokenAuth_Rejected(t *testing.T) {
+	srv := server.NewServer("localhost", 0, 0).WithToken("secret123")
+	ctx := t.Context()
+	go func() {
+		_ = srv.Start(ctx)
+	}()
+	<-srv.Ready()
+
+	controlAddr := srv.ControlAddr()
+
+	client := NewSelfHosted(controlAddr, "example").WithToken("wrong")
+	_, err := client.Connect(ctx, 3000)
+	if err == nil {
+		t.Fatalf("expected error from Connect() with invalid token, got nil")
+	}
+	if !strings.Contains(err.Error(), protocol.ErrUnauthorized) {
+		t.Errorf("expected '%s' error, got: %v", protocol.ErrUnauthorized, err)
+	}
+}

internal/server/protocol/protocol.go

@@ -1,16 +1,18 @@
 package protocol
 
 const (
-	TypeControl = "control"
-	TypeData    = "data"
+	TypeControl     = "control"
+	TypeData        = "data"
+	ErrUnauthorized = "unauthorized: invalid or missing auth token"
 )
 
 // TunnelRequest is sent by the client when connecting to the server.
 // Type determines the purpose: TypeControl for registration, TypeData for proxying.
 type TunnelRequest struct {
 	// Subdomain is the desired subdomain. If empty, server will assign a random one.
 	Subdomain string `json:"subdomain,omitempty"`
-	Type      string `json:"type,omitempty"` // "control" or "data"
+	Type      string `json:"type"`            // "control" or "data"
+	Token     string `json:"token,omitempty"` // optional auth token for self-hosted server
 }
 
 // TunnelResponse is sent by the server after registration.