1// SPDX-License-Identifier: AGPL-3.0-or-later
2
3package kvstore_test
4
5import (
6 "context"
7 "fmt"
8 "sync"
9 "testing"
10 "time"
11
12 "github.com/stretchr/testify/assert"
13 "github.com/stretchr/testify/require"
14
15 storagepb "bnema.dev/ermite/gen/ermite/storage"
16 "bnema.dev/ermite/internal/domain"
17)
18
19func newToken(t *testing.T, id, secret string) domain.Token {
20 t.Helper()
21 return domain.Token{
22 ID: id,
23 Label: "label " + id,
24 Hash: domain.HashTokenSecret(secret),
25 Scopes: []domain.Scope{domain.ScopeRepoRead, domain.ScopeRepoWrite},
26 Repos: []domain.RepoName{name(t, "a"), name(t, "b")},
27 ExpiresAt: created.Add(24 * time.Hour),
28 CreatedAt: created,
29 }
30}
31
32func TestTokensRoundTrip(t *testing.T) {
33 ctx := context.Background()
34 s := openStore(t, t.TempDir())
35 ts := s.Tokens()
36
37 tok := newToken(t, "t1", "secret-1")
38 require.NoError(t, ts.Create(ctx, tok))
39
40 got, err := ts.FindByHash(ctx, tok.Hash)
41 require.NoError(t, err)
42 assert.Equal(t, tok, got)
43
44 _, err = ts.FindByHash(ctx, domain.HashTokenSecret("other"))
45 assert.ErrorIs(t, err, domain.ErrNotFound)
46
47 // Minimal token: no repos, no expiry.
48 min := domain.Token{ID: "t2", Hash: domain.HashTokenSecret("secret-2"), Scopes: []domain.Scope{domain.ScopeRepoRead}, CreatedAt: created.Add(time.Minute)}
49 require.NoError(t, ts.Create(ctx, min))
50 got, err = ts.FindByHash(ctx, min.Hash)
51 require.NoError(t, err)
52 assert.Empty(t, got.Repos)
53 assert.True(t, got.ExpiresAt.IsZero())
54 assert.True(t, got.LastUsedAt.IsZero())
55
56 list, err := ts.List(ctx)
57 require.NoError(t, err)
58 require.Len(t, list, 2)
59 assert.Equal(t, "t1", list[0].ID)
60 assert.Equal(t, "t2", list[1].ID)
61}
62
63func TestTokensCreateRejects(t *testing.T) {
64 ctx := context.Background()
65 s := openStore(t, t.TempDir())
66 ts := s.Tokens()
67 require.NoError(t, ts.Create(ctx, newToken(t, "t1", "s1")))
68
69 assert.ErrorIs(t, ts.Create(ctx, newToken(t, "t1", "other")), domain.ErrAlreadyExists, "same id")
70 assert.ErrorIs(t, ts.Create(ctx, newToken(t, "t9", "s1")), domain.ErrAlreadyExists, "same hash")
71
72 // A failed create leaves no half-written index.
73 _, err := ts.FindByHash(ctx, domain.HashTokenSecret("other"))
74 assert.ErrorIs(t, err, domain.ErrNotFound)
75 list, err := ts.List(ctx)
76 require.NoError(t, err)
77 assert.Len(t, list, 1)
78
79 bad := newToken(t, "t3", "s3")
80 bad.Scopes = nil
81 assert.ErrorIs(t, ts.Create(ctx, bad), domain.ErrInvalidToken)
82 bad = newToken(t, "t3", "s3")
83 bad.Scopes = []domain.Scope{"nope"}
84 assert.ErrorIs(t, ts.Create(ctx, bad), domain.ErrInvalidScope)
85 bad = newToken(t, "a/b", "s3")
86 assert.ErrorIs(t, ts.Create(ctx, bad), domain.ErrInvalidToken)
87 bad = newToken(t, "", "s3")
88 assert.ErrorIs(t, ts.Create(ctx, bad), domain.ErrInvalidToken)
89}
90
91func TestTokensRevokeKeepsIndexConsistent(t *testing.T) {
92 ctx := context.Background()
93 dir := t.TempDir()
94 s := openStore(t, dir)
95 ts := s.Tokens()
96 a, b := newToken(t, "a", "sa"), newToken(t, "b", "sb")
97 require.NoError(t, ts.Create(ctx, a))
98 require.NoError(t, ts.Create(ctx, b))
99
100 require.NoError(t, ts.Revoke(ctx, "a"))
101 _, err := ts.FindByHash(ctx, a.Hash)
102 assert.ErrorIs(t, err, domain.ErrNotFound)
103 got, err := ts.FindByHash(ctx, b.Hash)
104 require.NoError(t, err)
105 assert.Equal(t, "b", got.ID)
106 assert.ErrorIs(t, ts.Revoke(ctx, "a"), domain.ErrNotFound)
107 assert.ErrorIs(t, ts.Revoke(ctx, "../x"), domain.ErrNotFound)
108
109 // The hash of a revoked token can be used again, and the index survives a restart.
110 require.NoError(t, ts.Create(ctx, newToken(t, "c", "sa")))
111 require.NoError(t, s.Close(ctx))
112 s2 := openStore(t, dir)
113 got, err = s2.Tokens().FindByHash(ctx, a.Hash)
114 require.NoError(t, err)
115 assert.Equal(t, "c", got.ID)
116 list, err := s2.Tokens().List(ctx)
117 require.NoError(t, err)
118 assert.Len(t, list, 2)
119}
120
121func TestTokensGet(t *testing.T) {
122 ctx := context.Background()
123 ts := openStore(t, t.TempDir()).Tokens()
124 tok := newToken(t, "t1", "s1")
125 require.NoError(t, ts.Create(ctx, tok))
126
127 got, err := ts.Get(ctx, "t1")
128 require.NoError(t, err)
129 assert.Equal(t, tok, got)
130
131 for _, id := range []string{"gone", "", "../t1"} {
132 _, err = ts.Get(ctx, id)
133 assert.ErrorIs(t, err, domain.ErrNotFound, id)
134 }
135 require.NoError(t, ts.Revoke(ctx, "t1"))
136 _, err = ts.Get(ctx, "t1")
137 assert.ErrorIs(t, err, domain.ErrNotFound)
138}
139
140func TestTokensTouch(t *testing.T) {
141 ctx := context.Background()
142 s := openStore(t, t.TempDir())
143 ts := s.Tokens()
144 tok := newToken(t, "t1", "s1")
145 require.NoError(t, ts.Create(ctx, tok))
146
147 at := created.Add(time.Hour)
148 require.NoError(t, ts.Touch(ctx, "t1", at))
149 got, err := ts.FindByHash(ctx, tok.Hash)
150 require.NoError(t, err)
151 assert.True(t, at.Equal(got.LastUsedAt))
152
153 // An older time does not move it back.
154 require.NoError(t, ts.Touch(ctx, "t1", created))
155 got, err = ts.FindByHash(ctx, tok.Hash)
156 require.NoError(t, err)
157 assert.True(t, at.Equal(got.LastUsedAt))
158
159 assert.ErrorIs(t, ts.Touch(ctx, "gone", at), domain.ErrNotFound)
160
161 // Everything else is untouched.
162 got.LastUsedAt = time.Time{}
163 assert.Equal(t, tok, got)
164}
165
166func TestTokensConcurrentTouchAndRevoke(t *testing.T) {
167 ctx := context.Background()
168 s := openStore(t, t.TempDir())
169 ts := s.Tokens()
170 tok := newToken(t, "t1", "s1")
171 require.NoError(t, ts.Create(ctx, tok))
172
173 var wg sync.WaitGroup
174 for i := range 8 {
175 wg.Add(1)
176 go func() {
177 defer wg.Done()
178 _ = ts.Touch(ctx, "t1", created.Add(time.Duration(i+1)*time.Second))
179 }()
180 }
181 wg.Add(1)
182 go func() { defer wg.Done(); _ = ts.Revoke(ctx, "t1") }()
183 wg.Wait()
184
185 // Whatever the interleaving, token and index are both there or both gone.
186 _, errTok := ts.FindByHash(ctx, tok.Hash)
187 list, err := ts.List(ctx)
188 require.NoError(t, err)
189 if errTok != nil {
190 assert.ErrorIs(t, errTok, domain.ErrNotFound, fmt.Sprint(errTok))
191 assert.Empty(t, list)
192 } else {
193 assert.Len(t, list, 1)
194 }
195}
196
197func TestTokensBranchPrefixes(t *testing.T) {
198 ctx := context.Background()
199
200 t.Run("round trip, order kept", func(t *testing.T) {
201 s := openStore(t, t.TempDir())
202 tok := newToken(t, "t1", "s1")
203 tok.BranchPrefixes = []string{"b/", "agent/claude/"}
204 require.NoError(t, s.Tokens().Create(ctx, tok))
205 got, err := s.Tokens().Get(ctx, "t1")
206 require.NoError(t, err)
207 assert.Equal(t, tok, got)
208 byHash, err := s.Tokens().FindByHash(ctx, tok.Hash)
209 require.NoError(t, err)
210 assert.Equal(t, []string{"b/", "agent/claude/"}, byHash.BranchPrefixes)
211 })
212
213 t.Run("a token stored before the field existed decodes with no prefixes", func(t *testing.T) {
214 dir := t.TempDir()
215 old := newToken(t, "old", "s-old")
216 putRaw(t, dir, "token/old", mustProto(t, &storagepb.Token{
217 Id: "old", Label: "old", Hash: old.Hash[:], Scopes: []string{"repo:write"},
218 }))
219 s := openStore(t, dir)
220 got, err := s.Tokens().Get(ctx, "old")
221 require.NoError(t, err)
222 assert.Empty(t, got.BranchPrefixes)
223 assert.Equal(t, []domain.Scope{domain.ScopeRepoWrite}, got.Scopes)
224 })
225
226 t.Run("an invalid prefix is not stored", func(t *testing.T) {
227 s := openStore(t, t.TempDir())
228 for _, bad := range []string{"", "agent", "/", "a b/"} {
229 tok := newToken(t, "t1", "s1")
230 tok.BranchPrefixes = []string{bad}
231 assert.ErrorIs(t, s.Tokens().Create(ctx, tok), domain.ErrInvalidToken, "%q", bad)
232 }
233 _, err := s.Tokens().Get(ctx, "t1")
234 assert.ErrorIs(t, err, domain.ErrNotFound)
235 })
236
237 t.Run("a stored prefix survives touch and stays limiting", func(t *testing.T) {
238 s := openStore(t, t.TempDir())
239 tok := newToken(t, "t1", "s1")
240 tok.BranchPrefixes = []string{"agent/"}
241 require.NoError(t, s.Tokens().Create(ctx, tok))
242 require.NoError(t, s.Tokens().Touch(ctx, "t1", created.Add(time.Hour)))
243 got, err := s.Tokens().Get(ctx, "t1")
244 require.NoError(t, err)
245 assert.Equal(t, []string{"agent/"}, got.BranchPrefixes)
246 })
247}