embedded: add options for authn on the client

Author: ecordell

README.md

@@ -159,14 +159,19 @@ opts := proxy.NewOptions()
 opts.EmbeddedMode = true
 opts.SpiceDBOptions.SpiceDBEndpoint = "embedded://"
 
-// Customize authentication headers (optional)
-opts.Authentication.Embedded.UsernameHeaders = []string{"Custom-User"}
-opts.Authentication.Embedded.GroupHeaders = []string{"Custom-Groups"}
-opts.Authentication.Embedded.ExtraHeaderPrefixes = []string{"Custom-Extra-"}
-
-// Create and use the embedded proxy
-proxySrv, _ := proxy.NewServer(ctx, *opts)
-client := proxySrv.GetEmbeddedClient()
+// Complete configuration
+completedConfig, _ := opts.Complete(ctx)
+proxySrv, _ := proxy.NewServer(ctx, completedConfig)
+
+// Get client with automatic authentication headers
+client := proxySrv.GetEmbeddedClient(
+    proxy.WithUser("alice"),
+    proxy.WithGroups("developers", "admin"),
+    proxy.WithExtra("department", "engineering"),
+)
+
+// Or get a basic client without authentication
+basicClient := proxySrv.GetEmbeddedClient()
 ```
 
 See [docs/embedding.md](docs/embedding.md) for detailed usage examples.

docs/embedding.md

@@ -51,10 +51,15 @@ func main() {
     }
 
     // Get an HTTP client that connects directly to the embedded proxy
-    embeddedClient := proxySrv.GetEmbeddedClient()
+    // Use functional options to automatically add authentication headers
+    embeddedClient := proxySrv.GetEmbeddedClient(
+        proxy.WithUser("my-user"),
+        proxy.WithGroups("my-group", "admin"),
+        proxy.WithExtra("department", "engineering"),
+    )
 
     // Create a Kubernetes client that uses the embedded proxy
-    k8sClient := createKubernetesClient(embeddedClient, "my-user", []string{"my-group"})
+    k8sClient := createKubernetesClient(embeddedClient)
 
     // Use the client normally - all requests go through SpiceDB authorization
     pods, err := k8sClient.CoreV1().Pods("default").List(ctx, metav1.ListOptions{})
@@ -65,44 +70,19 @@ func main() {
     fmt.Printf("Found %d pods\n", len(pods.Items))
 }
 
-func createKubernetesClient(embeddedClient *http.Client, username string, groups []string) *kubernetes.Clientset {
+func createKubernetesClient(embeddedClient *http.Client) *kubernetes.Clientset {
     restConfig := &rest.Config{
         Host:      "http://embedded", // Special URL for embedded mode
         Transport: embeddedClient.Transport,
     }
 
-    // Add authentication headers
-    restConfig.WrapTransport = func(rt http.RoundTripper) http.RoundTripper {
-        return &authTransport{
-            username: username,
-            groups:   groups,
-            rt:       rt,
-        }
-    }
-
     k8sClient, err := kubernetes.NewForConfig(restConfig)
     if err != nil {
         panic(err)
     }
 
     return k8sClient
 }
-
-type authTransport struct {
-    username string
-    groups   []string
-    rt       http.RoundTripper
-}
-
-func (t *authTransport) RoundTrip(req *http.Request) (*http.Response, error) {
-    // Add authentication headers (using default header names)
-    // These can be customized using --embedded-* flags
-    req.Header.Set("X-Remote-User", t.username)
-    for _, group := range t.groups {
-        req.Header.Add("X-Remote-Group", group)
-    }
-    return t.base.RoundTrip(req)
-}
 ```
 
 ### Configuration Options
@@ -163,14 +143,73 @@ opts.EmbeddedMode = true
 opts.Authentication.Embedded.UsernameHeaders = []string{"Custom-User"}
 opts.Authentication.Embedded.GroupHeaders = []string{"Custom-Groups"}
 opts.Authentication.Embedded.ExtraHeaderPrefixes = []string{"Custom-Extra-"}
+
+// Complete and create the proxy server
+completedConfig, _ := opts.Complete(ctx)
+proxySrv, _ := proxy.NewServer(ctx, completedConfig)
+
+// The client will automatically use the custom header names
+embeddedClient := proxySrv.GetEmbeddedClient(
+    proxy.WithUser("alice"),
+    proxy.WithGroups("developers", "admin"),
+    proxy.WithExtra("department", "engineering"),
+)
+// Headers will be: Custom-User: alice, Custom-Groups: developers, etc.
 ```
 
-Then use custom headers in requests:
+This is similar to Kubernetes' request header authentication, but uses a separate dedicated `EmbeddedAuthentication` type for embedded mode and doesn't require client certificate configuration (the requests are trusted because the server is embedded).
+
+### Functional Options for GetEmbeddedClient
+
+The `GetEmbeddedClient()` method supports functional options that automatically add authentication headers based on your configured header names. This eliminates the need to manually add headers to each request:
+
+```go
+// Basic client without authentication
+client := proxySrv.GetEmbeddedClient()
+
+// Client with user authentication
+client := proxySrv.GetEmbeddedClient(
+    proxy.WithUser("alice"),
+)
+
+// Client with user and groups
+client := proxySrv.GetEmbeddedClient(
+    proxy.WithUser("alice"),
+    proxy.WithGroups("developers", "admin", "reviewers"),
+)
+
+// Client with user, groups, and extra attributes
+client := proxySrv.GetEmbeddedClient(
+    proxy.WithUser("alice"),
+    proxy.WithGroups("developers", "admin"),
+    proxy.WithExtra("department", "engineering"),
+    proxy.WithExtra("team", "platform"),
+    proxy.WithExtra("location", "remote"),
+)
 ```
-Custom-User: alice
-Custom-Groups: developers
-Custom-Groups: admin
-Custom-Extra-Department: engineering
+
+The functional options automatically use the header names you've configured in `opts.Authentication.Embedded`. For example, if you've configured custom header names:
+
+```go
+opts.Authentication.Embedded.UsernameHeaders = []string{"My-User"}
+opts.Authentication.Embedded.GroupHeaders = []string{"My-Groups"}
+opts.Authentication.Embedded.ExtraHeaderPrefixes = []string{"My-Extra-"}
+
+// This client will automatically add:
+// My-User: alice
+// My-Groups: developers
+// My-Groups: admin  
+// My-Extra-department: engineering
+client := proxySrv.GetEmbeddedClient(
+    proxy.WithUser("alice"),
+    proxy.WithGroups("developers", "admin"),
+    proxy.WithExtra("department", "engineering"),
+)
 ```
 
-This is similar to Kubernetes' request header authentication, but uses a separate dedicated `EmbeddedAuthentication` type for embedded mode and doesn't require client certificate configuration (the requests are trusted because the server is embedded).
+Available functional options:
+- `WithUser(username string)`: Sets the username
+- `WithGroups(groups ...string)`: Sets group memberships  
+- `WithExtra(key, value string)`: Sets extra user attributes (can be called multiple times)
+
+This approach provides a clean, type-safe way to configure authentication without manually managing headers.

pkg/proxy/embedded_test.go

@@ -19,6 +19,7 @@ import (
 )
 
 func TestEmbeddedMode(t *testing.T) {
+	defer require.NoError(t, logsv1.ResetForTest(utilfeature.DefaultFeatureGate))
 	ctx := t.Context()
 
 	opts := createEmbeddedTestOptions(t)
@@ -260,6 +261,144 @@ func TestEmbeddedModeDefaults(t *testing.T) {
 	require.NoError(t, err)
 }
 
+func TestEmbeddedClientFunctionalOptions(t *testing.T) {
+	defer require.NoError(t, logsv1.ResetForTest(utilfeature.DefaultFeatureGate))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	// Create one proxy server for all subtests to avoid logging config issues
+	opts := createEmbeddedTestOptions(t)
+	completedConfig, err := opts.Complete(ctx)
+	require.NoError(t, err)
+
+	proxySrv, err := NewServer(ctx, completedConfig)
+	require.NoError(t, err)
+
+	t.Run("basic client without options", func(t *testing.T) {
+		client := proxySrv.GetEmbeddedClient()
+		require.NotNil(t, client)
+
+		// Basic client should not add any authentication headers automatically
+		req, err := http.NewRequestWithContext(ctx, "GET", "http://embedded/healthz", nil)
+		require.NoError(t, err)
+
+		resp, err := client.Do(req)
+		require.NoError(t, err)
+		defer resp.Body.Close()
+
+		require.NotEqual(t, 0, resp.StatusCode)
+		_, err = io.ReadAll(resp.Body)
+		require.NoError(t, err)
+	})
+
+	t.Run("transport adds headers correctly", func(t *testing.T) {
+		// Test the authHeaderTransport directly
+		baseTransport := &testTransport{}
+		transport := &authHeaderTransport{
+			base:                baseTransport,
+			username:            "test-user",
+			groups:              []string{"developers", "admin"},
+			extra:               map[string]string{"department": "engineering", "location": "remote"},
+			usernameHeaders:     []string{"X-Remote-User"},
+			groupHeaders:        []string{"X-Remote-Group"},
+			extraHeaderPrefixes: []string{"X-Remote-Extra-"},
+		}
+
+		req, err := http.NewRequest("GET", "http://example.com/test", nil)
+		require.NoError(t, err)
+
+		_, err = transport.RoundTrip(req)
+		require.NoError(t, err)
+
+		// Check that baseTransport received the request with added headers
+		capturedReq := baseTransport.lastRequest
+		require.NotNil(t, capturedReq)
+
+		// Check username header
+		require.Equal(t, "test-user", capturedReq.Header.Get("X-Remote-User"))
+
+		// Check group headers
+		groups := capturedReq.Header.Values("X-Remote-Group")
+		require.Contains(t, groups, "developers")
+		require.Contains(t, groups, "admin")
+		require.Len(t, groups, 2)
+
+		// Check extra headers
+		require.Equal(t, "engineering", capturedReq.Header.Get("X-Remote-Extra-department"))
+		require.Equal(t, "remote", capturedReq.Header.Get("X-Remote-Extra-location"))
+	})
+
+	t.Run("functional options create correct transport", func(t *testing.T) {
+		client := proxySrv.GetEmbeddedClient(
+			WithUser("alice"),
+			WithGroups("security", "reviewers"),
+			WithExtra("team", "platform"),
+		)
+		require.NotNil(t, client)
+
+		// Check that the transport is wrapped
+		transport, ok := client.Transport.(*authHeaderTransport)
+		require.True(t, ok, "transport should be wrapped with authHeaderTransport")
+
+		// Check configuration
+		require.Equal(t, "alice", transport.username)
+		require.Equal(t, []string{"security", "reviewers"}, transport.groups)
+		require.Equal(t, "platform", transport.extra["team"])
+		require.Equal(t, []string{"X-Remote-User"}, transport.usernameHeaders)
+		require.Equal(t, []string{"X-Remote-Group"}, transport.groupHeaders)
+		require.Equal(t, []string{"X-Remote-Extra-"}, transport.extraHeaderPrefixes)
+	})
+
+}
+
+func TestEmbeddedClientCustomHeaderConfig(t *testing.T) {
+	defer require.NoError(t, logsv1.ResetForTest(utilfeature.DefaultFeatureGate))
+
+	ctx, cancel := context.WithCancel(context.Background())
+	t.Cleanup(cancel)
+
+	// Create proxy with custom header names
+	customOpts := createEmbeddedTestOptions(t)
+	customOpts.Authentication.Embedded.UsernameHeaders = []string{"Custom-User"}
+	customOpts.Authentication.Embedded.GroupHeaders = []string{"Custom-Groups"}
+	customOpts.Authentication.Embedded.ExtraHeaderPrefixes = []string{"Custom-Extra-"}
+
+	customCompletedConfig, err := customOpts.Complete(ctx)
+	require.NoError(t, err)
+
+	customProxySrv, err := NewServer(ctx, customCompletedConfig)
+	require.NoError(t, err)
+
+	client := customProxySrv.GetEmbeddedClient(
+		WithUser("charlie"),
+		WithGroups("security"),
+		WithExtra("team", "infrastructure"),
+	)
+	require.NotNil(t, client)
+
+	// Check that custom header names are used
+	transport, ok := client.Transport.(*authHeaderTransport)
+	require.True(t, ok)
+	require.Equal(t, []string{"Custom-User"}, transport.usernameHeaders)
+	require.Equal(t, []string{"Custom-Groups"}, transport.groupHeaders)
+	require.Equal(t, []string{"Custom-Extra-"}, transport.extraHeaderPrefixes)
+}
+
+// testTransport is a simple transport that captures the last request
+type testTransport struct {
+	lastRequest *http.Request
+}
+
+func (t *testTransport) RoundTrip(req *http.Request) (*http.Response, error) {
+	t.lastRequest = req
+	return &http.Response{
+		StatusCode: 200,
+		Body:       http.NoBody,
+		Header:     make(http.Header),
+	}, nil
+}
+
 // createEmbeddedTestOptions creates minimal options for embedded testing
 func createEmbeddedTestOptions(t *testing.T) *Options {
 	t.Helper()