mirror of
https://github.com/coder/coder.git
synced 2026-09-24 15:04:27 +08:00
Paweł Banaszewski
96e3a64b12
Adds create, list and delete endpoints for AI Gateway keys. Those keys are used to authenticate into Coderd. All endpoints require Owner permission.
388 lines
13 KiB
Go
388 lines
13 KiB
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].LastUsedAt)
|
|
|
|
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/aibridge/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/aibridge/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/aibridge/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/aibridge/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/aibridge/keys",
|
|
wantStatus: http.StatusInternalServerError,
|
|
wantMsg: "Failed to list keys.",
|
|
},
|
|
{
|
|
name: "DeleteDBError",
|
|
errStore: aiGatewayKeyErrorStore{deleteErr: dbErr},
|
|
method: http.MethodDelete,
|
|
path: "/api/v2/aibridge/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")
|
|
})
|
|
}
|
|
}
|