feat: allowed hosts

Author: hugodutka

cmd/server/server.go

@@ -9,6 +9,7 @@ import (
 	"os"
 	"sort"
 	"strings"
+	"unicode"
 
 	"github.com/spf13/cobra"
 	"github.com/spf13/viper"
@@ -58,6 +59,26 @@ func parseAgentType(firstArg string, agentTypeVar string) (AgentType, error) {
 	return AgentTypeCustom, nil
 }
 
+// Validate allowed hosts don't contain whitespace or commas.
+// Viper/Cobra use different separators (space for env vars, comma for flags),
+// so these characters in origins likely indicate user error.
+func validateAllowedHosts(hosts []string) error {
+	if len(hosts) == 0 {
+		return fmt.Errorf("allowed hosts must not be empty")
+	}
+	for _, host := range hosts {
+		for _, r := range host {
+			if unicode.IsSpace(r) {
+				return fmt.Errorf("host '%s' contains whitespace characters, which are not allowed", host)
+			}
+			if strings.Contains(host, ",") {
+				return fmt.Errorf("host '%s' contains comma characters, which are not allowed", host)
+			}
+		}
+	}
+	return nil
+}
+
 func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) error {
 	agent := argsToPass[0]
 	agentTypeValue := viper.GetString(FlagType)
@@ -95,12 +116,16 @@ func runServer(ctx context.Context, logger *slog.Logger, argsToPass []string) er
 		}
 	}
 	port := viper.GetInt(FlagPort)
-	srv := httpapi.NewServer(ctx, httpapi.ServerConfig{
+	srv, err := httpapi.NewServer(ctx, httpapi.ServerConfig{
 		AgentType:    agentType,
 		Process:      process,
 		Port:         port,
 		ChatBasePath: viper.GetString(FlagChatBasePath),
+		AllowedHosts: viper.GetStringSlice(FlagAllowedHosts),
 	})
+	if err != nil {
+		return xerrors.Errorf("failed to create server: %w", err)
+	}
 	if printOpenAPI {
 		fmt.Println(srv.GetOpenAPI())
 		return nil
@@ -156,6 +181,7 @@ const (
 	FlagChatBasePath = "chat-base-path"
 	FlagTermWidth    = "term-width"
 	FlagTermHeight   = "term-height"
+	FlagAllowedHosts = "allowed-hosts"
 )
 
 func CreateServerCmd() *cobra.Command {
@@ -164,6 +190,13 @@ func CreateServerCmd() *cobra.Command {
 		Short: "Run the server",
 		Long:  fmt.Sprintf("Run the server with the specified agent (one of: %s)", strings.Join(agentNames, ", ")),
 		Args:  cobra.MinimumNArgs(1),
+		PreRunE: func(cmd *cobra.Command, args []string) error {
+			allowedOrigins := viper.GetStringSlice(FlagAllowedHosts)
+			if err := validateAllowedHosts(allowedOrigins); err != nil {
+				return err
+			}
+			return nil
+		},
 		Run: func(cmd *cobra.Command, args []string) {
 			logger := slog.New(slog.NewTextHandler(os.Stdout, nil))
 			ctx := logctx.WithLogger(context.Background(), logger)
@@ -181,6 +214,10 @@ func CreateServerCmd() *cobra.Command {
 		{FlagChatBasePath, "c", "/chat", "Base path for assets and routes used in the static files of the chat interface", "string"},
 		{FlagTermWidth, "W", uint16(80), "Width of the emulated terminal", "uint16"},
 		{FlagTermHeight, "H", uint16(1000), "Height of the emulated terminal", "uint16"},
+		// localhost:3284 is the default port for the server
+		// localhost:3000 is the default port for the chat interface during development
+		// localhost:3001 is used during development for the chat interface if 3000 is already in use
+		{FlagAllowedHosts, "a", []string{"localhost:3284", "localhost:3000", "localhost:3001"}, "HTTP allowed hosts. Use '*' for all, comma-separated list via flag, space-separated list via AGENTAPI_ALLOWED_HOSTS env var", "stringSlice"},
 	}
 
 	for _, spec := range flagSpecs {
@@ -193,6 +230,8 @@ func CreateServerCmd() *cobra.Command {
 			serverCmd.Flags().BoolP(spec.name, spec.shorthand, spec.defaultValue.(bool), spec.usage)
 		case "uint16":
 			serverCmd.Flags().Uint16P(spec.name, spec.shorthand, spec.defaultValue.(uint16), spec.usage)
+		case "stringSlice":
+			serverCmd.Flags().StringSliceP(spec.name, spec.shorthand, spec.defaultValue.([]string), spec.usage)
 		default:
 			panic(fmt.Sprintf("unknown flag type: %s", spec.flagType))
 		}

cmd/server/server_test.go

@@ -141,6 +141,7 @@ func TestServerCmd_AllArgs_Defaults(t *testing.T) {
 		{"chat-base-path default", FlagChatBasePath, "/chat", func() any { return viper.GetString(FlagChatBasePath) }},
 		{"term-width default", FlagTermWidth, uint16(80), func() any { return viper.GetUint16(FlagTermWidth) }},
 		{"term-height default", FlagTermHeight, uint16(1000), func() any { return viper.GetUint16(FlagTermHeight) }},
+		{"allowed-hosts default", FlagAllowedHosts, []string{"localhost:3284"}, func() any { return viper.GetStringSlice(FlagAllowedHosts) }},
 	}
 
 	for _, tt := range tests {
@@ -175,6 +176,7 @@ func TestServerCmd_AllEnvVars(t *testing.T) {
 		{"AGENTAPI_CHAT_BASE_PATH", "AGENTAPI_CHAT_BASE_PATH", "/api", "/api", func() any { return viper.GetString(FlagChatBasePath) }},
 		{"AGENTAPI_TERM_WIDTH", "AGENTAPI_TERM_WIDTH", "120", uint16(120), func() any { return viper.GetUint16(FlagTermWidth) }},
 		{"AGENTAPI_TERM_HEIGHT", "AGENTAPI_TERM_HEIGHT", "500", uint16(500), func() any { return viper.GetUint16(FlagTermHeight) }},
+		{"AGENTAPI_ALLOWED_HOSTS", "AGENTAPI_ALLOWED_HOSTS", "localhost:3284 localhost:3285", []string{"localhost:3284", "localhost:3285"}, func() any { return viper.GetStringSlice(FlagAllowedHosts) }},
 	}
 
 	for _, tt := range tests {
@@ -291,3 +293,119 @@ func TestMixed_ConfigurationScenarios(t *testing.T) {
 		assert.Equal(t, uint16(1000), viper.GetUint16(FlagTermHeight)) // default
 	})
 }
+
+func TestServerCmd_AllowedHosts(t *testing.T) {
+	tests := []struct {
+		name        string
+		env         map[string]string
+		args        []string
+		expectedErr string
+		expected    []string // only checked if expectedErr is empty
+	}{
+		// Environment variable scenarios (space-separated format)
+		{
+			name:     "env: single valid host",
+			env:      map[string]string{"AGENTAPI_ALLOWED_HOSTS": "localhost:3284"},
+			args:     []string{},
+			expected: []string{"localhost:3284"},
+		},
+		{
+			name:     "env: multiple valid hosts space-separated",
+			env:      map[string]string{"AGENTAPI_ALLOWED_HOSTS": "localhost:3284 example.com 192.168.1.1:8080"},
+			args:     []string{},
+			expected: []string{"localhost:3284", "example.com", "192.168.1.1:8080"},
+		},
+		{
+			name:     "env: host with tab",
+			env:      map[string]string{"AGENTAPI_ALLOWED_HOSTS": "localhost:3284\texample.com"},
+			args:     []string{},
+			expected: []string{"localhost:3284", "example.com"},
+		},
+		{
+			name:        "env: host with comma (invalid)",
+			env:         map[string]string{"AGENTAPI_ALLOWED_HOSTS": "localhost:3284,example.com"},
+			args:        []string{},
+			expectedErr: "contains comma characters",
+		},
+
+		// CLI flag scenarios (comma-separated format)
+		{
+			name:     "flag: single valid host",
+			args:     []string{"--allowed-hosts", "localhost:3284"},
+			expected: []string{"localhost:3284"},
+		},
+		{
+			name:     "flag: multiple valid hosts comma-separated",
+			args:     []string{"--allowed-hosts", "localhost:3284,example.com,192.168.1.1:8080"},
+			expected: []string{"localhost:3284", "example.com", "192.168.1.1:8080"},
+		},
+		{
+			name:     "flag: multiple valid hosts with multiple flags",
+			args:     []string{"--allowed-hosts", "localhost:3284", "--allowed-hosts", "example.com"},
+			expected: []string{"localhost:3284", "example.com"},
+		},
+		{
+			name:     "flag: host with newline",
+			args:     []string{"--allowed-hosts", "localhost:3284\n"},
+			expected: []string{"localhost:3284"},
+		},
+		{
+			name:        "flag: host with space in comma-separated list (invalid)",
+			args:        []string{"--allowed-hosts", "localhost:3284,example .com"},
+			expectedErr: "contains whitespace characters",
+		},
+
+		// Mixed scenarios (env + flag precedence)
+		{
+			name:     "mixed: flag overrides env",
+			env:      map[string]string{"AGENTAPI_ALLOWED_HOSTS": "localhost:8080"},
+			args:     []string{"--allowed-hosts", "override.com"},
+			expected: []string{"override.com"},
+		},
+		{
+			name:        "mixed: flag overrides env but flag is invalid",
+			env:         map[string]string{"AGENTAPI_ALLOWED_HOSTS": "localhost:8080"},
+			args:        []string{"--allowed-hosts", "invalid .com"},
+			expectedErr: "contains whitespace characters",
+		},
+
+		// Empty hosts are not allowed
+		{
+			name:        "empty host",
+			args:        []string{"--allowed-hosts", ""},
+			expectedErr: "allowed hosts must not be empty",
+		},
+
+		// Default behavior
+		{
+			name:     "default hosts when neither env nor flag provided",
+			args:     []string{},
+			expected: []string{"localhost:3284", "localhost:3000", "localhost:3001"},
+		},
+	}
+
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			isolateViper(t)
+
+			// Set environment variables if provided
+			for key, value := range tt.env {
+				t.Setenv(key, value)
+			}
+
+			serverCmd := CreateServerCmd()
+			// --print-openapi acts as an agent command that immediately exits
+			// use a 0 port to pick a random free port
+			serverCmd.SetArgs(append(tt.args, "--port", "0", "--", "sh", "-c", "echo ok"))
+			err := serverCmd.Execute()
+
+			if tt.expectedErr != "" {
+				require.Error(t, err)
+				assert.Contains(t, err.Error(), tt.expectedErr)
+			} else {
+				require.NoError(t, err)
+				assert.Equal(t, tt.expected, viper.GetStringSlice(FlagAllowedHosts))
+			}
+		})
+	}
+}

go.mod

@@ -54,6 +54,7 @@ require (
 	github.com/rogpeppe/go-internal v1.14.1 // indirect
 	github.com/spf13/afero v1.14.0
 	github.com/spf13/pflag v1.0.6 // indirect
+	github.com/unrolled/secure v1.17.0
 	golang.org/x/sync v0.12.0 // indirect
 	golang.org/x/sys v0.31.0 // indirect
 	golang.org/x/text v0.23.0 // indirect

go.sum

@@ -113,6 +113,8 @@ github.com/subosito/gotenv v1.6.0 h1:9NlTDc1FTs4qu0DDq7AEtTPNw6SVm7uBMsUCUjABIf8
 github.com/subosito/gotenv v1.6.0/go.mod h1:Dk4QP5c2W3ibzajGcXpNraDfq2IrhjMIvMSWPKKo0FU=
 github.com/tmaxmax/go-sse v0.10.0 h1:j9F93WB4Hxt8wUf6oGffMm4dutALvUPoDDxfuDQOSqA=
 github.com/tmaxmax/go-sse v0.10.0/go.mod h1:u/2kZQR1tyngo1lKaNCj1mJmhXGZWS1Zs5yiSOD+Eg8=
+github.com/unrolled/secure v1.17.0 h1:Io7ifFgo99Bnh0J7+Q+qcMzWM6kaDPCA5FroFZEdbWU=
+github.com/unrolled/secure v1.17.0/go.mod h1:BmF5hyM6tXczk3MpQkFf1hpKSRqCyhqcbiQtiAF7+40=
 go.uber.org/atomic v1.9.0 h1:ECmE8Bn/WFTYwEW/bpKD3M8VtR/zQVbavAoalC1PYyE=
 go.uber.org/atomic v1.9.0/go.mod h1:fEN4uk6kAWBTFdckzkM89CLk9XfWZrxpCo0nPH17wJc=
 go.uber.org/multierr v1.9.0 h1:7fIwc/ZtS0q++VgcfqFDxSBZVv/Xo49/SYnDFupUwlI=

lib/httpapi/server.go

@@ -7,6 +7,7 @@ import (
 	"log/slog"
 	"net/http"
 	"net/url"
+	"slices"
 	"strings"
 	"sync"
 	"time"
@@ -20,6 +21,7 @@ import (
 	"github.com/danielgtaylor/huma/v2/sse"
 	"github.com/go-chi/chi/v5"
 	"github.com/go-chi/cors"
+	"github.com/unrolled/secure"
 	"golang.org/x/xerrors"
 )
 
@@ -64,14 +66,55 @@ type ServerConfig struct {
 	Process      *termexec.Process
 	Port         int
 	ChatBasePath string
+	AllowedHosts []string
+}
+
+func parseAllowedHosts(hosts []string) ([]string, error) {
+	if slices.Contains(hosts, "*") {
+		return []string{}, nil
+	}
+	for _, host := range hosts {
+		if strings.Contains(host, "*") {
+			return nil, xerrors.Errorf("wildcard characters are not supported: %q", host)
+		}
+		if strings.Contains(host, "http://") || strings.Contains(host, "https://") {
+			return nil, xerrors.Errorf("host must not contain http:// or https://: %q", host)
+		}
+	}
+	return hosts, nil
+}
+
+func hostsToOrigins(hosts []string) []string {
+	if len(hosts) == 0 {
+		return []string{"*"}
+	}
+	origins := []string{}
+	for _, host := range hosts {
+		origins = append(origins, "http://"+host)
+		origins = append(origins, "https://"+host)
+	}
+	return origins
 }
 
 // NewServer creates a new server instance
-func NewServer(ctx context.Context, config ServerConfig) *Server {
+func NewServer(ctx context.Context, config ServerConfig) (*Server, error) {
 	router := chi.NewMux()
 
+	allowedHosts, err := parseAllowedHosts(config.AllowedHosts)
+	if err != nil {
+		return nil, xerrors.Errorf("failed to validate allowed hosts: %w", err)
+	}
+	secureMiddleware := secure.New(secure.Options{
+		AllowedHosts: allowedHosts,
+	})
+	badHostHandler := http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		http.Error(w, "Invalid host header. Allowed hosts: "+strings.Join(allowedHosts, ", "), http.StatusBadRequest)
+	})
+	secureMiddleware.SetBadHostHandler(badHostHandler)
+	router.Use(secureMiddleware.Handler)
+
 	corsMiddleware := cors.New(cors.Options{
-		AllowedOrigins:   []string{"*"},
+		AllowedOrigins:   hostsToOrigins(allowedHosts),
 		AllowedMethods:   []string{"GET", "POST", "PUT", "DELETE", "OPTIONS"},
 		AllowedHeaders:   []string{"Accept", "Authorization", "Content-Type", "X-CSRF-Token"},
 		ExposedHeaders:   []string{"Link"},
@@ -111,7 +154,7 @@ func NewServer(ctx context.Context, config ServerConfig) *Server {
 	// Register API routes
 	s.registerRoutes()
 
-	return s
+	return s, nil
 }
 
 // Handler returns the underlying chi.Router for testing purposes.