server_test.go

// SPDX-FileCopyrightText: Copyright 2025 Stacklok, Inc.
// SPDX-License-Identifier: Apache-2.0

package authserver

import (
	"context"
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/x509"
	"encoding/json"
	"encoding/pem"
	"log/slog"
	"net/http"
	"net/http/httptest"
	"os"
	"path/filepath"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/go-jose/go-jose/v4"
	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"go.uber.org/mock/gomock"

	servercrypto "github.com/stacklok/toolhive/pkg/authserver/server/crypto"
	"github.com/stacklok/toolhive/pkg/authserver/server/keys"
	"github.com/stacklok/toolhive/pkg/authserver/server/registration"
	"github.com/stacklok/toolhive/pkg/authserver/storage"
	storagemocks "github.com/stacklok/toolhive/pkg/authserver/storage/mocks"
	"github.com/stacklok/toolhive/pkg/authserver/upstream"
	upstreammocks "github.com/stacklok/toolhive/pkg/authserver/upstream/mocks"
)

// validUpstreamConfig returns a valid upstream config for tests.
func validUpstreamConfig() *upstream.OAuth2Config {
	return &upstream.OAuth2Config{
		CommonOAuthConfig: upstream.CommonOAuthConfig{
			ClientID:    "test-client",
			RedirectURI: "https://example.com/callback",
		},
		AuthorizationEndpoint: "https://idp.example.com/auth",
		TokenEndpoint:         "https://idp.example.com/token",
	}
}

// validHMACSecret returns a valid HMAC secret for tests.
func validHMACSecret() []byte {
	secret := make([]byte, 32)
	_, _ = rand.Read(secret)
	return secret
}

func TestNew(t *testing.T) {
	t.Parallel()

	validKeyProvider := keys.NewGeneratingProvider(keys.DefaultAlgorithm)
	validHMAC := &servercrypto.HMACSecrets{Current: validHMACSecret()}
	validUpstreams := []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}}

	tests := []struct {
		name        string
		cfg         Config
		storageNil  bool
		wantErr     bool
		errContains string
	}{
		{
			name:        "nil storage returns error",
			cfg:         Config{},
			storageNil:  true,
			wantErr:     true,
			errContains: "invalid config",
		},
		{
			name:        "empty issuer returns error",
			cfg:         Config{},
			storageNil:  false,
			wantErr:     true,
			errContains: "issuer is required",
		},
		// Note: "missing HMAC secrets" no longer returns an error because
		// applyDefaults() auto-generates them when nil
		{
			name: "HMAC secret too short returns error",
			cfg: Config{
				Issuer:           "https://example.com",
				KeyProvider:      validKeyProvider,
				HMACSecrets:      &servercrypto.HMACSecrets{Current: []byte("short")},
				Upstreams:        validUpstreams,
				AllowedAudiences: []string{"https://mcp.example.com"},
			},
			storageNil:  false,
			wantErr:     true,
			errContains: "HMAC secret must be at least 32 bytes",
		},
		{
			name: "missing upstreams returns error",
			cfg: Config{
				Issuer:           "https://example.com",
				KeyProvider:      validKeyProvider,
				HMACSecrets:      validHMAC,
				AllowedAudiences: []string{"https://mcp.example.com"},
			},
			storageNil:  false,
			wantErr:     true,
			errContains: "at least one upstream is required",
		},
		{
			name: "missing allowed audiences returns error",
			cfg: Config{
				Issuer:      "https://example.com",
				KeyProvider: validKeyProvider,
				HMACSecrets: validHMAC,
				Upstreams:   validUpstreams,
			},
			storageNil:  false,
			wantErr:     true,
			errContains: "at least one allowed audience is required",
		},
	}

	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			t.Parallel()

			ctrl := gomock.NewController(t)
			defer ctrl.Finish()

			var stor *storagemocks.MockStorage
			if !tt.storageNil {
				stor = storagemocks.NewMockStorage(ctrl)
			}

			ctx := context.Background()
			_, err := New(ctx, tt.cfg, stor)

			if tt.wantErr {
				if err == nil {
					t.Errorf("New() error = nil, wantErr %v", tt.wantErr)
					return
				}
				if tt.errContains != "" && !strings.Contains(err.Error(), tt.errContains) {
					t.Errorf("New() error = %q, want error containing %q", err.Error(), tt.errContains)
				}
			} else {
				if err != nil {
					t.Errorf("New() unexpected error = %v", err)
				}
			}
		})
	}
}

// TestNewServer_Success tests the success path with mocked dependencies.
func TestNewServer_Success(t *testing.T) {
	t.Parallel()

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()

	mockUpstream := upstreammocks.NewMockOAuth2Provider(ctrl)

	// Use a real MemoryStorage rather than storagemocks.MockStorage: the
	// constructor type-asserts the storage to storage.DCRCredentialStore (per
	// the F6 design — Storage no longer embeds DCRCredentialStore), and
	// generated MockStorage does not implement DCRCredentialStore. This test
	// exercises the constructor flow, not specific storage method calls, so
	// a real MemoryStorage is sufficient and keeps the assertion path real.
	stor := storage.NewMemoryStorage()
	t.Cleanup(func() { _ = stor.Close() })

	// Create valid config
	cfg := Config{
		Issuer:           "https://example.com",
		KeyProvider:      keys.NewGeneratingProvider(keys.DefaultAlgorithm),
		HMACSecrets:      &servercrypto.HMACSecrets{Current: validHMACSecret()},
		Upstreams:        []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
		AllowedAudiences: []string{"https://mcp.example.com"},
	}

	// Create factory that returns our mock
	mockFactory := func(_ context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) {
		return mockUpstream, nil
	}

	// Call newServer with the mock factory
	ctx := context.Background()
	srv, err := newServer(ctx, cfg, stor, withUpstreamFactory(mockFactory))

	if err != nil {
		t.Fatalf("newServer() unexpected error: %v", err)
	}
	if srv == nil {
		t.Fatal("newServer() returned nil server")
	}
	if srv.Handler() == nil {
		t.Error("server.Handler() returned nil")
	}
	if srv.IDPTokenStorage() != stor {
		t.Error("server.IDPTokenStorage() did not return expected storage")
	}
}

// capturingSlogHandler records log records for assertions. slog's default
// handler is process-global, so tests using it must not run in parallel with
// other slog-capturing tests.
type capturingSlogHandler struct {
	sink *capturingSlogSink
	// attrs carries the slog.With(...) attributes in effect for this handler,
	// flattened into every rendered record by recordsContaining. Without this,
	// a secret leaked via slog.With("client_secret", s).Info(...) — the most
	// likely way one escapes during a refactor — would be invisible to the
	// leak-detection tests.
	attrs []slog.Attr
}

// capturingSlogSink is the shared record store behind a capturingSlogHandler
// and every handler derived from it via WithAttrs.
type capturingSlogSink struct {
	mu      sync.Mutex
	records []slog.Record
}

func newCapturingSlogHandler() *capturingSlogHandler {
	return &capturingSlogHandler{sink: &capturingSlogSink{}}
}

func (*capturingSlogHandler) Enabled(_ context.Context, _ slog.Level) bool { return true }

func (h *capturingSlogHandler) Handle(_ context.Context, r slog.Record) error {
	h.sink.mu.Lock()
	defer h.sink.mu.Unlock()
	// Flatten the With(...) attributes into the record so recordsContaining
	// sees them alongside per-call attrs.
	r.AddAttrs(h.attrs...)
	h.sink.records = append(h.sink.records, r)
	return nil
}

func (h *capturingSlogHandler) WithAttrs(attrs []slog.Attr) slog.Handler {
	return &capturingSlogHandler{sink: h.sink, attrs: append(h.attrs, attrs...)}
}
func (h *capturingSlogHandler) WithGroup(_ string) slog.Handler { return h }

func (h *capturingSlogHandler) messages(level slog.Level, containing string) []string {
	h.sink.mu.Lock()
	defer h.sink.mu.Unlock()
	var out []string
	for _, r := range h.sink.records {
		if r.Level == level && strings.Contains(r.Message, containing) {
			out = append(out, r.Message)
		}
	}
	return out
}

// recordsContaining returns every captured record — message plus all
// attribute values, at any level — that contains needle. Used by leak-detection
// tests that must prove a secret appears in ZERO log records, not just in
// zero messages.
func (h *capturingSlogHandler) recordsContaining(needle string) []string {
	h.sink.mu.Lock()
	defer h.sink.mu.Unlock()
	var out []string
	for _, r := range h.sink.records {
		var b strings.Builder
		b.WriteString(r.Message)
		r.Attrs(func(a slog.Attr) bool {
			b.WriteString(" ")
			b.WriteString(a.Value.String())
			return true
		})
		if strings.Contains(b.String(), needle) {
			out = append(out, b.String())
		}
	}
	return out
}

// TestNewServer_AllowConfidentialClientRegistration_Logs pins the startup logging
// contract: enabling the flag logs an Info naming the consequence when
// startup succeeds. Combining it with insecure_allow_http is rejected by
// Config.Validate (see ValidateConfidentialClientTransport) before this log
// line is ever reached, so that combination is covered by
// TestConfig_Validate_RejectsConfidentialClientOverInsecureHTTP instead.
//
//nolint:paralleltest // swaps the process-global slog default handler
func TestNewServer_AllowConfidentialClientRegistration_Logs(t *testing.T) {
	// Not parallel: swaps the process-global slog default handler.

	newCfg := func(allowConfidential bool) Config {
		return Config{
			Issuer:                              "https://example.com",
			KeyProvider:                         keys.NewGeneratingProvider(keys.DefaultAlgorithm),
			HMACSecrets:                         &servercrypto.HMACSecrets{Current: validHMACSecret()},
			Upstreams:                           []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
			AllowedAudiences:                    []string{"https://mcp.example.com"},
			AllowConfidentialClientRegistration: allowConfidential,
		}
	}

	tests := []struct {
		name              string
		allowConfidential bool
		wantInfo          bool
	}{
		{"flag off: no logs", false, false},
		{"flag on: Info naming the consequence", true, true},
	}
	for _, tt := range tests {
		t.Run(tt.name, func(t *testing.T) {
			capture := newCapturingSlogHandler()
			prev := slog.Default()
			slog.SetDefault(slog.New(capture))
			t.Cleanup(func() { slog.SetDefault(prev) })

			ctrl := gomock.NewController(t)
			defer ctrl.Finish()
			mockUpstream := upstreammocks.NewMockOAuth2Provider(ctrl)
			mockFactory := func(_ context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) {
				return mockUpstream, nil
			}
			stor := storage.NewMemoryStorage()
			t.Cleanup(func() { _ = stor.Close() })

			srv, err := newServer(context.Background(), newCfg(tt.allowConfidential), stor, withUpstreamFactory(mockFactory))
			require.NoError(t, err, "startup must succeed for every flag combination")
			require.NotNil(t, srv)

			// Filter to the flag's own log lines: other components (key
			// generation, baseline scopes) also log at Info during startup.
			infos := capture.messages(slog.LevelInfo, "client secrets")

			if tt.wantInfo {
				require.Len(t, infos, 1)
				assert.Contains(t, infos[0], "unauthenticated dynamic registration")
			} else {
				assert.Empty(t, infos)
			}
		})
	}
}

// TestConfig_Validate_RejectsConfidentialClientOverInsecureHTTP pins the
// rejection of allow_confidential_client_registration combined with insecure_allow_http:
// issuing client secrets over cleartext HTTP on an unauthenticated
// registration endpoint must fail loudly at config validation, not just log
// a warning.
func TestConfig_Validate_RejectsConfidentialClientOverInsecureHTTP(t *testing.T) {
	t.Parallel()

	cfg := Config{
		Issuer:                              "http://example.com",
		KeyProvider:                         keys.NewGeneratingProvider(keys.DefaultAlgorithm),
		HMACSecrets:                         &servercrypto.HMACSecrets{Current: validHMACSecret()},
		Upstreams:                           []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
		AllowedAudiences:                    []string{"https://mcp.example.com"},
		AllowConfidentialClientRegistration: true,
		InsecureAllowHTTP:                   true,
	}

	err := cfg.Validate()
	require.Error(t, err)
	assert.Contains(t, err.Error(), "allow_confidential_client_registration")
	assert.Contains(t, err.Error(), "insecure_allow_http")
}

func TestNewServer_CIMDEnabled_WrapsStorage(t *testing.T) {
	t.Parallel()

	mockUpstream := upstreammocks.NewMockOAuth2Provider(gomock.NewController(t))

	stor := storage.NewMemoryStorage()
	t.Cleanup(func() { _ = stor.Close() })

	cfg := Config{
		Issuer:               "https://example.com",
		KeyProvider:          keys.NewGeneratingProvider(keys.DefaultAlgorithm),
		HMACSecrets:          &servercrypto.HMACSecrets{Current: validHMACSecret()},
		Upstreams:            []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
		AllowedAudiences:     []string{"https://mcp.example.com"},
		CIMDEnabled:          true,
		CIMDCacheMaxSize:     16,
		CIMDCacheFallbackTTL: 5 * time.Minute,
	}

	mockFactory := func(_ context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) {
		return mockUpstream, nil
	}

	srv, err := newServer(context.Background(), cfg, stor, withUpstreamFactory(mockFactory))
	if err != nil {
		t.Fatalf("newServer() unexpected error: %v", err)
	}

	_, ok := srv.storage.(*storage.CIMDStorageDecorator)
	if !ok {
		t.Errorf("expected storage to be *storage.CIMDStorageDecorator when CIMDEnabled=true, got %T", srv.storage)
	}
}

// TestNewServer_UpstreamRefresherSharedInstance verifies the wiring this PR
// fixes: UpstreamTokenRefresher() must return the single refresher constructed
// in newServer rather than reallocating one per call. The pre-fix accessor
// rebuilt the refresher (and its singleflight.Group) on every call, so the
// handler chain-walk path and the runtime token-swap path ended up with
// independent groups and cross-path refresh deduplication was impossible.
// A regression that reintroduced per-call allocation would leave the
// refresher's own singleflight test green, so this asserts instance identity
// at the server boundary instead.
func TestNewServer_UpstreamRefresherSharedInstance(t *testing.T) {
	t.Parallel()

	mockUpstream := upstreammocks.NewMockOAuth2Provider(gomock.NewController(t))
	stor := storage.NewMemoryStorage()
	t.Cleanup(func() { _ = stor.Close() })

	cfg := Config{
		Issuer:           "https://example.com",
		KeyProvider:      keys.NewGeneratingProvider(keys.DefaultAlgorithm),
		HMACSecrets:      &servercrypto.HMACSecrets{Current: validHMACSecret()},
		Upstreams:        []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
		AllowedAudiences: []string{"https://mcp.example.com"},
	}
	mockFactory := func(_ context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) {
		return mockUpstream, nil
	}

	srv, err := newServer(context.Background(), cfg, stor, withUpstreamFactory(mockFactory))
	require.NoError(t, err)

	first := srv.UpstreamTokenRefresher()
	require.NotNil(t, first, "refresher must be non-nil when upstreams are configured")
	// Repeated calls must return the identical instance — i.e. the same
	// singleflight.Group — not a freshly allocated one.
	assert.Same(t, first, srv.UpstreamTokenRefresher(),
		"UpstreamTokenRefresher() must return the shared instance, not reallocate per call")
	// That instance must be the field stored on the server, which is the same
	// value wired into the handler via WithUpstreamRefresher in newServer.
	assert.Same(t, srv.upstreamRefresher, first,
		"accessor must return the stored instance shared with the handler")
}

// TestNewUpstreamTokenRefresher_NilWhenNoUpstreams verifies the true-nil
// interface contract: with no upstreams the constructor must return a nil
// interface value, not a typed nil (*upstreamTokenRefresher)(nil) wrapped in an
// interface, so that callers' `== nil` checks (runner, service, handler) work.
func TestNewUpstreamTokenRefresher_NilWhenNoUpstreams(t *testing.T) {
	t.Parallel()

	stor := storage.NewMemoryStorage()
	t.Cleanup(func() { _ = stor.Close() })

	refresher := newUpstreamTokenRefresher(nil, stor, 24*time.Hour)
	// Direct == nil comparison, not assert.Nil: testify's Nil also passes for a
	// typed nil pointer, which would hide exactly the bug this guards against.
	if refresher != nil {
		t.Fatalf("expected a true nil interface, got non-nil %T", refresher)
	}
}

func TestNewServer_RegistersDelegateClientsBeforeUpstreamConstruction(t *testing.T) {
	t.Parallel()

	ctx := t.Context()
	stor := storage.NewMemoryStorage()
	t.Cleanup(func() { _ = stor.Close() })
	cfg := Config{
		Issuer:           "https://example.com",
		KeyProvider:      keys.NewGeneratingProvider(keys.DefaultAlgorithm),
		HMACSecrets:      &servercrypto.HMACSecrets{Current: validHMACSecret()},
		Upstreams:        []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
		AllowedAudiences: []string{"https://mcp.example.com"},
		DelegateClients: []DelegateClient{{
			ClientID:     "delegate",
			ClientSecret: "delegate-secret-well-above-the-minimum-length",
			Scopes:       []string{"openid"}, Audiences: []string{"https://mcp.example.com"},
		}},
	}

	factory := func(ctx context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) {
		client, err := stor.GetClient(ctx, "delegate")
		require.NoError(t, err)
		assert.False(t, registration.DCRIssued(client))
		assert.False(t, client.IsPublic())
		return nil, assert.AnError
	}

	_, err := newServer(ctx, cfg, stor, withUpstreamFactory(factory))
	require.ErrorIs(t, err, assert.AnError)

	// Startup registration is an upsert. Replacing a same-ID DCR client makes
	// it permanent and unmarked rather than retaining DCR eviction semantics.
	dcrClient, err := registration.NewConfidentialPlain(registration.Config{ID: "delegate", Secret: "old-secret"})
	require.NoError(t, err)
	require.NoError(t, stor.RegisterClient(ctx, dcrClient))
	_, err = newServer(ctx, cfg, stor, withUpstreamFactory(factory))
	require.ErrorIs(t, err, assert.AnError)
	client, err := stor.GetClient(ctx, "delegate")
	require.NoError(t, err)
	assert.False(t, registration.DCRIssued(client))
}

func TestNewServer_JWKSIncludesFallbackKeys(t *testing.T) {
	t.Parallel()

	dir := t.TempDir()
	writePEM := func(key *ecdsa.PrivateKey, name string) string {
		der, err := x509.MarshalECPrivateKey(key)
		require.NoError(t, err)
		path := filepath.Join(dir, name)
		data := pem.EncodeToMemory(&pem.Block{Type: "EC PRIVATE KEY", Bytes: der})
		require.NoError(t, os.WriteFile(path, data, 0600))
		return name
	}
	k1, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err)
	k2, err := ecdsa.GenerateKey(elliptic.P256(), rand.Reader)
	require.NoError(t, err)
	signingFile := writePEM(k1, "signing.pem")
	fallbackFile := writePEM(k2, "fallback.pem")

	provider, err := keys.NewFileProvider(keys.Config{
		KeyDir:           dir,
		SigningKeyFile:   signingFile,
		FallbackKeyFiles: []string{fallbackFile},
	})
	require.NoError(t, err)

	ctrl := gomock.NewController(t)
	defer ctrl.Finish()
	mockUpstream := upstreammocks.NewMockOAuth2Provider(ctrl)
	stor := storage.NewMemoryStorage()
	t.Cleanup(func() { _ = stor.Close() })

	cfg := Config{
		Issuer:           "https://example.com",
		KeyProvider:      provider,
		HMACSecrets:      &servercrypto.HMACSecrets{Current: validHMACSecret()},
		Upstreams:        []UpstreamConfig{{Name: "default", Type: UpstreamProviderTypeOAuth2, OAuth2Config: validUpstreamConfig()}},
		AllowedAudiences: []string{"https://mcp.example.com"},
	}
	factory := func(_ context.Context, _ *UpstreamConfig) (upstream.OAuth2Provider, error) {
		return mockUpstream, nil
	}
	srv, err := newServer(context.Background(), cfg, stor, withUpstreamFactory(factory))
	require.NoError(t, err)

	// Hit the JWKS endpoint and verify both keys are published, primary first.
	req := httptest.NewRequest(http.MethodGet, "/.well-known/jwks.json", nil)
	w := httptest.NewRecorder()
	srv.Handler().ServeHTTP(w, req)
	require.Equal(t, http.StatusOK, w.Code)
	var jwks jose.JSONWebKeySet
	require.NoError(t, json.Unmarshal(w.Body.Bytes(), &jwks))
	require.Len(t, jwks.Keys, 2)
	pubKeys, err := provider.PublicKeys(context.Background())
	require.NoError(t, err)
	assert.Equal(t, pubKeys[0].KeyID, jwks.Keys[0].KeyID)
	assert.Equal(t, pubKeys[1].KeyID, jwks.Keys[1].KeyID)
}