Fix cors serving (#156)

Author: djthorpe

pkg/cmd/server.go

@@ -114,6 +114,7 @@ func (s *RunServer) Run(ctx server.Cmd) error {
 	if err != nil {
 		return fmt.Errorf("router: %w", err)
 	}
+	srv.SetHandler(router)
 
 	// Register routes
 	for _, fn := range s.register {

pkg/httpserver/httpserver.go

@@ -100,6 +100,17 @@ func (server *server) Router() *http.ServeMux {
 	return server.mux
 }
 
+// SetHandler replaces the server's active HTTP handler. This is used when a
+// higher-level router wraps the underlying ServeMux with middleware such as
+// CORS, CSRF protection, tracing, or authentication.
+func (server *server) SetHandler(handler http.Handler) {
+	if handler == nil {
+		server.http.Handler = server.mux
+		return
+	}
+	server.http.Handler = handler
+}
+
 // Addr returns the listen address. After [Listen] has been called this
 // returns the actual bound address (which may differ from the configured
 // address when an ephemeral port is used).

pkg/httpserver/httpserver_test.go

@@ -203,6 +203,40 @@ func Test_Options_004(t *testing.T) {
 	assert.NotNil(s)
 }
 
+func Test_Server_UsesCustomHandler(t *testing.T) {
+	assert := assert.New(t)
+
+	s, err := httpserver.New(":0", nil)
+	if !assert.NoError(err) {
+		return
+	}
+
+	s.SetHandler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(http.StatusAccepted)
+	}))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	defer cancel()
+
+	if !assert.NoError(s.Listen()) {
+		return
+	}
+
+	done := make(chan error, 1)
+	go func() {
+		done <- s.Run(ctx)
+	}()
+
+	resp, err := http.Get(s.URL().String())
+	if assert.NoError(err) {
+		defer resp.Body.Close()
+		assert.Equal(http.StatusAccepted, resp.StatusCode)
+	}
+
+	cancel()
+	assert.NoError(<-done)
+}
+
 ///////////////////////////////////////////////////////////////////////////////
 // TESTS - TLS CONFIG