Skip to content

Commit 38eb88e

Browse files
EmyrkEmyrk
committed
Support api keys in generator
1 parent 722194b commit 38eb88e

File tree

3 files changed

+93
-35
lines changed

3 files changed

+93
-35
lines changed

‎coderd/database/databasefake/generator.go

Lines changed: 72 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,10 +2,16 @@ package databasefake
22

33
import (
44
"context"
5+
"crypto/sha256"
56
"database/sql"
7+
"fmt"
68
"testing"
79
"time"
810

11+
"github.com/coder/coder/cryptorand"
12+
13+
"github.com/tabbed/pqtype"
14+
915
"github.com/coder/coder/coderd/database"
1016
"github.com/google/uuid"
1117
"github.com/moby/moby/pkg/namesgenerator"
@@ -52,10 +58,42 @@ func (g *Generator) PrimaryOrg(ctx context.Context) database.Organization {
5258
}
5359

5460
func populate[DBType any](ctx context.Context, g *Generator, name string, seed DBType) DBType {
61+
if name == "" {
62+
name = g.RandomName()
63+
}
64+
5565
out := g.Populate(ctx, map[string]interface{}{
5666
name: seed,
5767
})
58-
return out[name].(DBType)
68+
v, ok := out[name].(DBType)
69+
if !ok {
70+
panic("developer error, type mismatch")
71+
}
72+
return v
73+
}
74+
75+
func (g *Generator) RandomName() string {
76+
for {
77+
name := namesgenerator.GetRandomName(0)
78+
if _, ok := g.names[name]; !ok {
79+
return name
80+
}
81+
}
82+
}
83+
84+
func (g *Generator) APIKey(ctx context.Context, name string, seed database.APIKey) (key database.APIKey, token string) {
85+
if name == "" {
86+
name = g.RandomName()
87+
}
88+
89+
out := g.Populate(ctx, map[string]interface{}{
90+
name: seed,
91+
})
92+
key, keyOk := out[name].(database.APIKey)
93+
secret, secOk := out[name+"_secret"].(string)
94+
require.True(g.testT, keyOk && secOk, "APIKey & secret must be populated with the right type")
95+
96+
return key, fmt.Sprintf("%s-%s", key.ID, secret)
5997
}
6098

6199
func (g *Generator) WorkspaceResource(ctx context.Context, name string, seed database.WorkspaceResource) database.WorkspaceResource {
@@ -100,11 +138,34 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
100138

101139
for name, v := range seed {
102140
switch orig := v.(type) {
141+
case database.APIKey:
142+
id, _ := cryptorand.String(10)
143+
secret, _ := cryptorand.String(22)
144+
hashed := sha256.Sum256([]byte(secret))
145+
146+
key, err := db.InsertAPIKey(ctx, database.InsertAPIKeyParams{
147+
ID: takeFirst(orig.ID, id),
148+
LifetimeSeconds: takeFirst(orig.LifetimeSeconds, 3600),
149+
HashedSecret: takeFirstBytes(orig.HashedSecret, hashed[:]),
150+
IPAddress: pqtype.Inet{},
151+
UserID: takeFirst(orig.UserID, uuid.New()),
152+
LastUsed: takeFirstTime(orig.LastUsed, time.Now()),
153+
ExpiresAt: takeFirstTime(orig.ExpiresAt, time.Now().Add(time.Hour)),
154+
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
155+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
156+
LoginType: takeFirst(orig.LoginType, database.LoginTypePassword),
157+
Scope: takeFirst(orig.Scope, database.APIKeyScopeAll),
158+
})
159+
require.NoError(t, err, "insert api key")
160+
161+
seed[name] = key
162+
// Need to also save the secret
163+
seed[name+"_secret"] = secret
103164
case database.Template:
104165
template, err := db.InsertTemplate(ctx, database.InsertTemplateParams{
105166
ID: g.Lookup(name),
106167
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
107-
UpdatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
168+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
108169
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
109170
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
110171
Provisioner: takeFirst(orig.Provisioner, database.ProvisionerTypeEcho),
@@ -125,7 +186,7 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
125186
workspace, err := db.InsertWorkspace(ctx, database.InsertWorkspaceParams{
126187
ID: g.Lookup(name),
127188
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
128-
UpdatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
189+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
129190
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
130191
TemplateID: takeFirst(orig.TemplateID, uuid.New()),
131192
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
@@ -139,15 +200,15 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
139200
build, err := db.InsertWorkspaceBuild(ctx, database.InsertWorkspaceBuildParams{
140201
ID: g.Lookup(name),
141202
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
142-
UpdatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
203+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
143204
WorkspaceID: takeFirst(orig.WorkspaceID, uuid.New()),
144205
TemplateVersionID: takeFirst(orig.TemplateVersionID, uuid.New()),
145206
BuildNumber: takeFirst(orig.BuildNumber, 0),
146207
Transition: takeFirst(orig.Transition, database.WorkspaceTransitionStart),
147208
InitiatorID: takeFirst(orig.InitiatorID, uuid.New()),
148209
JobID: takeFirst(orig.JobID, uuid.New()),
149210
ProvisionerState: takeFirstBytes(orig.ProvisionerState, []byte{}),
150-
Deadline: takeFirstTime(orig.CreatedAt, time.Now().Add(time.Hour)),
211+
Deadline: takeFirstTime(orig.Deadline, time.Now().Add(time.Hour)),
151212
Reason: takeFirst(orig.Reason, database.BuildReasonInitiator),
152213
})
153214
require.NoError(t, err, "insert workspace build")
@@ -160,7 +221,7 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
160221
Username: takeFirst(orig.Username, namesgenerator.GetRandomName(1)),
161222
HashedPassword: takeFirstBytes(orig.HashedPassword, []byte{}),
162223
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
163-
UpdatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
224+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
164225
RBACRoles: []string{},
165226
LoginType: takeFirst(orig.LoginType, database.LoginTypePassword),
166227
})
@@ -172,9 +233,9 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
172233
org, err := db.InsertOrganization(ctx, database.InsertOrganizationParams{
173234
ID: g.Lookup(name),
174235
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
175-
Description: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
236+
Description: takeFirst(orig.Description, namesgenerator.GetRandomName(1)),
176237
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
177-
UpdatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
238+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
178239
})
179240
require.NoError(t, err, "insert organization")
180241

@@ -185,7 +246,7 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
185246
ID: g.Lookup(name),
186247
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
187248
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
188-
AvatarURL: takeFirst(orig.Name, "https://logo.example.com"),
249+
AvatarURL: takeFirst(orig.AvatarURL, "https://logo.example.com"),
189250
QuotaAllowance: takeFirst(orig.QuotaAllowance, 0),
190251
})
191252
require.NoError(t, err, "insert organization")
@@ -196,7 +257,7 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
196257
job, err := db.InsertProvisionerJob(ctx, database.InsertProvisionerJobParams{
197258
ID: g.Lookup(name),
198259
CreatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
199-
UpdatedAt: takeFirstTime(orig.CreatedAt, time.Now()),
260+
UpdatedAt: takeFirstTime(orig.UpdatedAt, time.Now()),
200261
OrganizationID: takeFirst(orig.OrganizationID, g.PrimaryOrg(ctx).ID),
201262
InitiatorID: takeFirst(orig.InitiatorID, uuid.New()),
202263
Provisioner: takeFirst(orig.Provisioner, database.ProvisionerTypeEcho),
@@ -220,7 +281,7 @@ func (g *Generator) Populate(ctx context.Context, seed map[string]interface{}) m
220281
Type: takeFirst(orig.Type, ""),
221282
Name: takeFirst(orig.Name, namesgenerator.GetRandomName(1)),
222283
Hide: takeFirst(orig.Hide, false),
223-
Icon: takeFirst(orig.Name, ""),
284+
Icon: takeFirst(orig.Icon, ""),
224285
InstanceType: sql.NullString{
225286
String: takeFirst(orig.InstanceType.String, ""),
226287
Valid: takeFirst(orig.InstanceType.Valid, false),

‎coderd/httpmw/apikey_test.go

Lines changed: 14 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -151,24 +151,21 @@ func TestAPIKey(t *testing.T) {
151151
t.Run("InvalidSecret", func(t *testing.T) {
152152
t.Parallel()
153153
var (
154-
db = databasefake.New()
155-
id, secret = randomAPIKeyParts()
156-
r = httptest.NewRequest("GET", "/", nil)
157-
rw = httptest.NewRecorder()
158-
user = createUser(r.Context(), t, db)
154+
db = databasefake.New()
155+
gen = databasefake.NewGenerator(t, db)
156+
r = httptest.NewRequest("GET", "/", nil)
157+
rw = httptest.NewRecorder()
158+
ctx = context.Background()
159+
user = gen.User(ctx, "", database.User{})
160+
161+
// Use a different secret so they don't match!
162+
hashed = sha256.Sum256([]byte("differentsecret"))
163+
_, token = gen.APIKey(ctx, "", database.APIKey{
164+
UserID: user.ID,
165+
HashedSecret: hashed[:],
166+
})
159167
)
160-
r.Header.Set(codersdk.SessionTokenHeader, fmt.Sprintf("%s-%s", id, secret))
161-
162-
// Use a different secret so they don't match!
163-
hashed := sha256.Sum256([]byte("differentsecret"))
164-
_, err := db.InsertAPIKey(r.Context(), database.InsertAPIKeyParams{
165-
ID: id,
166-
HashedSecret: hashed[:],
167-
UserID: user.ID,
168-
LoginType: database.LoginTypePassword,
169-
Scope: database.APIKeyScopeAll,
170-
})
171-
require.NoError(t, err)
168+
r.Header.Set(codersdk.SessionTokenHeader, token)
172169
httpmw.ExtractAPIKey(httpmw.ExtractAPIKeyConfig{
173170
DB: db,
174171
RedirectToLogin: false,

‎coderd/httpmw/ratelimit_test.go

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -94,9 +94,9 @@ func TestRateLimit(t *testing.T) {
9494
ctx := context.Background()
9595

9696
db := databasefake.New()
97-
98-
u := createUser(ctx, t, db)
99-
key := insertAPIKey(ctx, t, db, u.ID)
97+
gen := databasefake.NewGenerator(t, db)
98+
u := gen.User(ctx, "", database.User{})
99+
_, key := gen.APIKey(ctx, "", database.APIKey{UserID: u.ID})
100100

101101
rtr := chi.NewRouter()
102102
rtr.Use(httpmw.ExtractAPIKey(httpmw.ExtractAPIKeyConfig{
@@ -141,11 +141,11 @@ func TestRateLimit(t *testing.T) {
141141

142142
db := databasefake.New()
143143

144-
u := createUser(ctx, t, db, func(u *database.InsertUserParams) {
145-
u.RBACRoles = []string{rbac.RoleOwner()}
144+
gen := databasefake.NewGenerator(t, db)
145+
u := gen.User(ctx, "", database.User{
146+
RBACRoles: []string{rbac.RoleOwner()},
146147
})
147-
148-
key := insertAPIKey(ctx, t, db, u.ID)
148+
_, key := gen.APIKey(ctx, "", database.APIKey{UserID: u.ID})
149149

150150
rtr := chi.NewRouter()
151151
rtr.Use(httpmw.ExtractAPIKey(httpmw.ExtractAPIKeyConfig{

0 commit comments

Comments
 (0)