coder/enterprise/coderd/aigatewaykeys_test.go

package coderd_test

import (
	"context"
	"encoding/json"
	"fmt"
	"io"
	"net/http"
	"strings"
	"testing"
	"time"

	"github.com/google/uuid"
	"github.com/stretchr/testify/require"
	"golang.org/x/xerrors"

	aibridgekeys "github.com/coder/coder/v2/coderd/aibridge/keys"
	"github.com/coder/coder/v2/coderd/audit"
	"github.com/coder/coder/v2/coderd/coderdtest"
	"github.com/coder/coder/v2/coderd/database"
	"github.com/coder/coder/v2/coderd/database/dbauthz"
	"github.com/coder/coder/v2/coderd/database/dbtestutil"
	"github.com/coder/coder/v2/codersdk"
	entaudit "github.com/coder/coder/v2/enterprise/audit"
	"github.com/coder/coder/v2/enterprise/audit/backends"
	"github.com/coder/coder/v2/enterprise/coderd/coderdenttest"
	"github.com/coder/coder/v2/enterprise/coderd/license"
	"github.com/coder/coder/v2/testutil"
)

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

	t.Run("CRUD", func(t *testing.T) {
		t.Parallel()

		ownerClient, _ := coderdenttest.New(t, aibridgeOpts(t))
		ctx := testutil.Context(t, testutil.WaitLong)

		//nolint:gocritic // Managing AI Gateway keys is owner-only.
		keys, err := ownerClient.ListAIGatewayKeys(ctx)
		require.NoError(t, err)
		require.Empty(t, keys)

		name := uniqueName(t, "happy")

		created, err := ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: name})
		require.NoError(t, err)
		require.NotEqual(t, uuid.Nil, created.ID)
		require.Equal(t, name, created.Name)
		require.Len(t, created.KeyPrefix, aibridgekeys.KeyPrefixLength)
		require.Len(t, created.Key, aibridgekeys.KeyLength)
		require.True(t, strings.HasPrefix(created.Key, created.KeyPrefix), "key must begin with key_prefix")
		require.WithinDuration(t, time.Now(), created.CreatedAt, time.Minute)

		keys, err = ownerClient.ListAIGatewayKeys(ctx)
		require.NoError(t, err)
		require.Len(t, keys, 1)
		require.Equal(t, created.ID, keys[0].ID)
		require.Equal(t, created.Name, keys[0].Name)
		require.Equal(t, created.KeyPrefix, keys[0].KeyPrefix)
		require.Nil(t, keys[0].LastHeartbeatAt)

		require.NoError(t, ownerClient.DeleteAIGatewayKey(ctx, created.ID))

		keys, err = ownerClient.ListAIGatewayKeys(ctx)
		require.NoError(t, err)
		require.Empty(t, keys)
	})

	t.Run("ListResponseDoesNotLeakSecrets", func(t *testing.T) {
		t.Parallel()

		ownerClient, _ := coderdenttest.New(t, aibridgeOpts(t))
		ctx := testutil.Context(t, testutil.WaitLong)

		//nolint:gocritic // Managing AI Gateway keys is owner-only.
		created, err := ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{
			Name: uniqueName(t, "leak"),
		})
		require.NoError(t, err)
		fullKey := created.Key

		resp, err := ownerClient.Request(ctx, http.MethodGet, "/api/v2/ai-gateway/keys", nil)
		require.NoError(t, err)
		t.Cleanup(func() { _ = resp.Body.Close() })
		require.Equal(t, http.StatusOK, resp.StatusCode)

		body, err := io.ReadAll(resp.Body)
		require.NoError(t, err)

		require.NotContains(t, string(body), fullKey, "LIST response leaked full key")
	})

	t.Run("CreateValidation", func(t *testing.T) {
		t.Parallel()

		ownerClient, _ := coderdenttest.New(t, aibridgeOpts(t))
		ctx := testutil.Context(t, testutil.WaitLong)

		// Empty name -> 400 (validate:"required" on request struct).
		//nolint:gocritic // Managing AI Gateway keys is owner-only.
		_, err := ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: ""})
		var sdkErr *codersdk.Error
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
		require.ErrorContains(t, err, "Validation failed")

		// >64 char name -> 400 (DB check constraint).
		longName := strings.Repeat("a", 65)
		_, err = ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: longName})
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
		require.ErrorContains(t, err, "Invalid key name")

		// Uppercase name -> 400 (DB check constraint rejects non-lowercase).
		_, err = ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: "UPPER-CASE"})
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
		require.ErrorContains(t, err, "Invalid key name")

		// Duplicate name -> 400.
		name := uniqueName(t, "dup")
		_, err = ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: name})
		require.NoError(t, err)
		_, err = ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: name})
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusBadRequest, sdkErr.StatusCode())
		require.ErrorContains(t, err, "must be unique")
	})

	t.Run("DeleteValidation", func(t *testing.T) {
		t.Parallel()

		ownerClient, _ := coderdenttest.New(t, aibridgeOpts(t))
		ctx := testutil.Context(t, testutil.WaitLong)

		// Invalid UUID -> 400 (raw request; SDK method accepts uuid.UUID).
		//nolint:gocritic // Managing AI Gateway keys is owner-only.
		resp, err := ownerClient.Request(ctx, http.MethodDelete, "/api/v2/ai-gateway/keys/not-a-uuid", nil)
		require.NoError(t, err)
		t.Cleanup(func() { _ = resp.Body.Close() })
		require.Equal(t, http.StatusBadRequest, resp.StatusCode)

		// Existing id -> 204.
		created, err := ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{
			Name: uniqueName(t, "del"),
		})
		require.NoError(t, err)
		// SDK returns no code on success, using raw request to check for 204.
		delResp, err := ownerClient.Request(ctx, http.MethodDelete, "/api/v2/ai-gateway/keys/"+created.ID.String(), nil)
		require.NoError(t, err)
		defer delResp.Body.Close()
		require.Equal(t, http.StatusNoContent, delResp.StatusCode)

		// Not existing id -> 404.
		err = ownerClient.DeleteAIGatewayKey(ctx, uuid.New())
		var sdkErr *codersdk.Error
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusNotFound, sdkErr.StatusCode())
	})

	t.Run("ReturnsForbiddenForNonOwners", func(t *testing.T) {
		t.Parallel()

		ownerClient, owner := coderdenttest.New(t, aibridgeOpts(t))
		ctx := testutil.Context(t, testutil.WaitLong)
		member, _ := coderdtest.CreateAnotherUser(t, ownerClient, owner.OrganizationID)

		_, err := member.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{
			Name: uniqueName(t, "denied"),
		})
		var sdkErr *codersdk.Error
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusForbidden, sdkErr.StatusCode())

		_, err = member.ListAIGatewayKeys(ctx)
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusForbidden, sdkErr.StatusCode())

		err = member.DeleteAIGatewayKey(ctx, uuid.New())
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusForbidden, sdkErr.StatusCode())
	})

	t.Run("LicenseEntitlement", func(t *testing.T) {
		t.Parallel()

		ownerClient, _ := coderdenttest.New(t, &coderdenttest.Options{
			LicenseOptions: &coderdenttest.LicenseOptions{
				Features: license.Features{},
			},
		})
		ctx := testutil.Context(t, testutil.WaitLong)

		//nolint:gocritic // Managing AI Gateway keys is owner-only.
		_, err := ownerClient.ListAIGatewayKeys(ctx)
		var sdkErr *codersdk.Error
		require.ErrorAs(t, err, &sdkErr)
		require.Equal(t, http.StatusForbidden, sdkErr.StatusCode())
		require.Contains(t, sdkErr.Message, "AI Gateway is a Premium feature")
	})
}

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

	db, ps := dbtestutil.NewDB(t)
	auditor := entaudit.NewAuditor(
		db,
		entaudit.DefaultFilter,
		backends.NewPostgres(db, true),
	)
	opts := aibridgeOpts(t)
	opts.AuditLogging = true
	opts.Options.Database = db
	opts.Options.Pubsub = ps
	opts.Options.Auditor = auditor
	opts.LicenseOptions.Features[codersdk.FeatureAuditLog] = 1

	ownerClient, _ := coderdenttest.New(t, opts)
	ctx, cancel := context.WithTimeout(context.Background(), testutil.WaitMedium)
	defer cancel()

	name := uniqueName(t, "audit")
	//nolint:gocritic // Managing AI Gateway coderd keys is owner-only.
	created, err := ownerClient.CreateAIGatewayKey(ctx, codersdk.CreateAIGatewayKeyRequest{Name: name})
	require.NoError(t, err)
	//nolint:gocritic // Managing AI Gateway coderd keys is owner-only.
	require.NoError(t, ownerClient.DeleteAIGatewayKey(ctx, created.ID))

	rows, err := db.GetAuditLogsOffset(
		dbauthz.AsSystemRestricted(ctx),
		database.GetAuditLogsOffsetParams{
			ResourceType: string(database.ResourceTypeAIGatewayKey),
			LimitOpt:     10,
		},
	)
	require.NoError(t, err)
	require.Len(t, rows, 2, "expected one create and one delete audit row")

	var createLog, deleteLog database.AuditLog
	for _, row := range rows {
		log := row.AuditLog
		switch log.Action {
		case database.AuditActionCreate:
			createLog = log
		case database.AuditActionDelete:
			deleteLog = log
		default:
			require.Failf(t, "unexpected audit action", "action: %s", log.Action)
		}
	}
	require.Equal(t, database.AuditActionCreate, createLog.Action)
	require.Equal(t, database.AuditActionDelete, deleteLog.Action)
	require.Equal(t, http.StatusCreated, int(createLog.StatusCode))
	require.Equal(t, http.StatusNoContent, int(deleteLog.StatusCode))

	for _, log := range []database.AuditLog{createLog, deleteLog} {
		require.Equal(t, database.ResourceTypeAIGatewayKey, log.ResourceType)
		require.Equal(t, created.ID, log.ResourceID)
		require.Equal(t, name, log.ResourceTarget)
	}

	var createDiff audit.Map
	require.NoError(t, json.Unmarshal(createLog.Diff, &createDiff))
	require.Contains(t, createDiff, "name")
	require.Equal(t, "", createDiff["name"].Old)
	require.Equal(t, name, createDiff["name"].New)
	require.Contains(t, createDiff, "secret_prefix")
	require.Equal(t, "", createDiff["secret_prefix"].Old)
	require.Equal(t, created.KeyPrefix, createDiff["secret_prefix"].New)
	require.NotContains(t, createDiff, "hashed_secret")

	var deleteDiff audit.Map
	require.NoError(t, json.Unmarshal(deleteLog.Diff, &deleteDiff))
	require.Contains(t, deleteDiff, "name")
	require.Equal(t, name, deleteDiff["name"].Old)
	require.Equal(t, "", deleteDiff["name"].New)
	require.NotContains(t, deleteDiff, "hashed_secret")
}

func uniqueName(t *testing.T, prefix string) string {
	t.Helper()
	return strings.ToLower(fmt.Sprintf("%s-%d", prefix, time.Now().UnixNano()))
}

// aiGatewayKeyErrorStore wraps a database.Store and forces specific
// methods to return errors, allowing tests to exercise error paths.
type aiGatewayKeyErrorStore struct {
	database.Store
	insertErr error
	listErr   error
	deleteErr error
}

func (s *aiGatewayKeyErrorStore) InsertAIGatewayKey(ctx context.Context, arg database.InsertAIGatewayKeyParams) (database.InsertAIGatewayKeyRow, error) {
	if s.insertErr != nil {
		return database.InsertAIGatewayKeyRow{}, s.insertErr
	}
	return s.Store.InsertAIGatewayKey(ctx, arg)
}

func (s *aiGatewayKeyErrorStore) ListAIGatewayKeys(ctx context.Context) ([]database.ListAIGatewayKeysRow, error) {
	if s.listErr != nil {
		return nil, s.listErr
	}
	return s.Store.ListAIGatewayKeys(ctx)
}

func (s *aiGatewayKeyErrorStore) DeleteAIGatewayKey(ctx context.Context, id uuid.UUID) (database.DeleteAIGatewayKeyRow, error) {
	if s.deleteErr != nil {
		return database.DeleteAIGatewayKeyRow{}, s.deleteErr
	}
	return s.Store.DeleteAIGatewayKey(ctx, id)
}

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

	dbErr := xerrors.New("internal db failure")

	tests := []struct {
		name       string
		errStore   aiGatewayKeyErrorStore
		method     string
		path       string
		body       any
		wantStatus int
		wantMsg    string
	}{
		{
			name:       "CreateDBError",
			errStore:   aiGatewayKeyErrorStore{insertErr: dbErr},
			method:     http.MethodPost,
			path:       "/api/v2/ai-gateway/keys",
			body:       codersdk.CreateAIGatewayKeyRequest{Name: "db-err-create"},
			wantStatus: http.StatusInternalServerError,
			wantMsg:    "Failed to create key. Please retry.",
		},
		{
			name:       "ListDBError",
			errStore:   aiGatewayKeyErrorStore{listErr: dbErr},
			method:     http.MethodGet,
			path:       "/api/v2/ai-gateway/keys",
			wantStatus: http.StatusInternalServerError,
			wantMsg:    "Failed to list keys.",
		},
		{
			name:       "DeleteDBError",
			errStore:   aiGatewayKeyErrorStore{deleteErr: dbErr},
			method:     http.MethodDelete,
			path:       "/api/v2/ai-gateway/keys/" + uuid.New().String(),
			wantStatus: http.StatusInternalServerError,
			wantMsg:    "Failed to delete key.",
		},
	}

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

			db, ps := dbtestutil.NewDB(t)
			errStore := tc.errStore
			errStore.Store = db

			opts := aibridgeOpts(t)
			opts.Options.Database = &errStore
			opts.Options.Pubsub = ps

			ownerClient, _ := coderdenttest.New(t, opts)
			ctx := testutil.Context(t, testutil.WaitLong)

			//nolint:gocritic // Managing AI Gateway keys is owner-only.
			resp, err := ownerClient.Request(ctx, tc.method, tc.path, tc.body)
			require.NoError(t, err)
			defer resp.Body.Close()

			require.Equal(t, tc.wantStatus, resp.StatusCode)

			var sdkResp codersdk.Response
			require.NoError(t, json.NewDecoder(resp.Body).Decode(&sdkResp))
			require.Equal(t, tc.wantMsg, sdkResp.Message)
			require.Empty(t, sdkResp.Detail, "response must not leak internal error details")
		})
	}
}