1// SPDX-License-Identifier: MIT
2
3package resultssvc
4
5import (
6 "bytes"
7 "context"
8 "crypto/sha256"
9 "encoding/hex"
10 "errors"
11 "fmt"
12 "net/http"
13 "net/http/httptest"
14 "sort"
15 "strings"
16 "sync"
17 "testing"
18 "time"
19
20 "bnema.dev/ermite-api/gen/ermite"
21 "bnema.dev/ermite-api/gen/ermite/ermiteconnect"
22 "connectrpc.com/connect"
23 "google.golang.org/protobuf/proto"
24 "google.golang.org/protobuf/types/known/timestamppb"
25)
26
27// fakeErmite is an in-memory JobArtifactService that follows the semantics of
28// proto/ermite/artifact.proto: job tokens bound to a (run, job, attempt),
29// pending, committed and finalized artifacts, blocks replaced on retry, commit
30// order, finalize replacing an earlier attempt's artifact of the same name,
31// lists by run, downloads from an offset in chunks of 1 MiB, and deletes. Its
32// hooks inject failures.
33type fakeErmite struct {
34 ermiteconnect.UnimplementedJobArtifactServiceHandler
35
36 mu sync.Mutex
37 jobs map[string]fakeJob // by job token
38 nextID int64
39 arts map[int64]*fakeArt
40 now time.Time
41
42 // maxArtifact, when > 0, is the most bytes of blocks one artifact may hold:
43 // beyond, UploadJobArtifactBlock fails with resource_exhausted.
44 maxArtifact int64
45 // failNext makes the next call of a procedure (by method name) fail.
46 failNext map[string]error
47 // cut, when set, is called before each chunk of a download stream with the
48 // offset the stream was opened at and the bytes it has sent; an error ends
49 // the stream.
50 cut func(open, sent int64) error
51 // blockStreams counts the UploadJobArtifactBlock calls that ended.
52 blockStreams int
53}
54
55type fakeJob struct {
56 run int64
57 job string
58 attempt int32
59}
60
61type artState int
62
63const (
64 pending artState = iota
65 committed
66 finalized
67)
68
69type fakeArt struct {
70 a *ermite.Artifact
71 state artState
72 blocks map[string][]byte
73 list [][]byte
74 content []byte
75}
76
77func newFakeErmite() *fakeErmite {
78 return &fakeErmite{jobs: map[string]fakeJob{}, arts: map[int64]*fakeArt{}, failNext: map[string]error{},
79 now: time.Date(2026, 10, 15, 12, 0, 0, 0, time.UTC)}
80}
81
82// serve starts the fake on an HTTP server and returns its URL.
83func (f *fakeErmite) serve(t *testing.T) string {
84 t.Helper()
85 path, h := ermiteconnect.NewJobArtifactServiceHandler(f)
86 mux := http.NewServeMux()
87 mux.Handle(path, h)
88 srv := httptest.NewServer(mux)
89 t.Cleanup(srv.Close)
90 return srv.URL
91}
92
93// addJob registers a job token.
94func (f *fakeErmite) addJob(token string, run int64, job string, attempt int32) {
95 f.mu.Lock()
96 defer f.mu.Unlock()
97 f.jobs[token] = fakeJob{run, job, attempt}
98}
99
100// bearer is an http.RoundTripper that adds a job token, as the runner's ermite
101// client does.
102type bearer struct {
103 token string
104 base http.RoundTripper
105}
106
107func (b bearer) RoundTrip(r *http.Request) (*http.Response, error) {
108 r = r.Clone(r.Context())
109 r.Header.Set("Authorization", "Bearer "+b.token)
110 return b.base.RoundTrip(r)
111}
112
113// client is a JobArtifactService client of the fake with a job token.
114func client(url, token string) ermiteconnect.JobArtifactServiceClient {
115 return ermiteconnect.NewJobArtifactServiceClient(&http.Client{Transport: bearer{token, http.DefaultTransport}}, url)
116}
117
118func (f *fakeErmite) streams() int {
119 f.mu.Lock()
120 defer f.mu.Unlock()
121 return f.blockStreams
122}
123
124// caller checks the token and the job id of a request. f.mu must be held.
125func (f *fakeErmite) caller(h http.Header, jobID, method string) (fakeJob, error) {
126 if err := f.failNext[method]; err != nil {
127 delete(f.failNext, method)
128 return fakeJob{}, err
129 }
130 tok, ok := strings.CutPrefix(h.Get("Authorization"), "Bearer ")
131 j, known := f.jobs[tok]
132 if !ok || !known {
133 return fakeJob{}, connect.NewError(connect.CodeUnauthenticated, errors.New("unknown job token"))
134 }
135 if jobID != j.job {
136 return fakeJob{}, connect.NewError(connect.CodePermissionDenied, errors.New("job_id is not the job of the token"))
137 }
138 return j, nil
139}
140
141func digestOf(b []byte) string {
142 s := sha256.Sum256(b)
143 return "sha256:" + hex.EncodeToString(s[:])
144}
145
146func (f *fakeErmite) CreateJobArtifact(_ context.Context, req *connect.Request[ermite.CreateJobArtifactRequest]) (*connect.Response[ermite.CreateJobArtifactResponse], error) {
147 f.mu.Lock()
148 defer f.mu.Unlock()
149 j, err := f.caller(req.Header(), req.Msg.GetJobId(), "CreateJobArtifact")
150 if err != nil {
151 return nil, err
152 }
153 name := req.Msg.GetName()
154 if name == "" || len(name) > 256 || strings.ContainsAny(name, "\":<>|*?\\/") {
155 return nil, connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("invalid artifact name %q", name))
156 }
157 for id, a := range f.arts {
158 if a.a.GetRunId() != j.run || a.a.GetName() != name {
159 continue
160 }
161 if a.state == finalized && a.a.GetAttempt() == j.attempt {
162 return nil, connect.NewError(connect.CodeAlreadyExists, errors.New(msgExists))
163 }
164 if a.state != finalized && a.a.GetJobId() == j.job {
165 delete(f.arts, id) // a new create is a new upload
166 }
167 }
168 mt := req.Msg.GetMediaType()
169 if mt == "" {
170 mt = "application/zip"
171 }
172 exp := f.now.Add(90 * 24 * time.Hour)
173 if e := req.Msg.GetExpiresAt(); e != nil && e.AsTime().After(f.now) {
174 exp = e.AsTime()
175 }
176 f.nextID++
177 f.arts[f.nextID] = &fakeArt{a: &ermite.Artifact{
178 Id: f.nextID, Repo: "bnema/app", RunId: j.run, Attempt: j.attempt, JobId: j.job, Name: name, MediaType: mt,
179 CreatedAt: timestamppb.New(f.now), ExpiresAt: timestamppb.New(exp),
180 }, blocks: map[string][]byte{}}
181 return connect.NewResponse(&ermite.CreateJobArtifactResponse{ArtifactId: f.nextID, ExpiresAt: timestamppb.New(exp)}), nil
182}
183
184func (f *fakeErmite) UploadJobArtifactBlock(_ context.Context, st *connect.ClientStream[ermite.UploadJobArtifactBlockRequest]) (*connect.Response[ermite.UploadJobArtifactBlockResponse], error) {
185 defer func() {
186 f.mu.Lock()
187 f.blockStreams++
188 f.mu.Unlock()
189 }()
190 if !st.Receive() {
191 if err := st.Err(); err != nil {
192 return nil, err
193 }
194 return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("empty stream"))
195 }
196 first := st.Msg()
197 f.mu.Lock()
198 j, err := f.caller(st.RequestHeader(), first.GetJobId(), "UploadJobArtifactBlock")
199 var art *fakeArt
200 var others int64
201 if err == nil {
202 art = f.arts[first.GetArtifactId()]
203 switch {
204 case art == nil || art.a.GetJobId() != j.job || art.state == finalized:
205 err = connect.NewError(connect.CodeNotFound, errors.New("not a pending artifact of this job"))
206 case art.state == committed:
207 err = connect.NewError(connect.CodeFailedPrecondition, errors.New("the artifact is committed"))
208 case len(first.GetBlockId()) == 0 || len(first.GetBlockId()) > 64:
209 err = connect.NewError(connect.CodeInvalidArgument, errors.New("block_id must be 1 to 64 bytes"))
210 }
211 if art != nil {
212 for id, b := range art.blocks {
213 if id != string(first.GetBlockId()) {
214 others += int64(len(b))
215 }
216 }
217 }
218 }
219 limit := f.maxArtifact
220 f.mu.Unlock()
221 if err != nil {
222 return nil, err
223 }
224 var data bytes.Buffer
225 msg := first
226 for {
227 if len(msg.GetData()) > chunkSize {
228 return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("data over 1 MiB"))
229 }
230 data.Write(msg.GetData())
231 if limit > 0 && others+int64(data.Len()) > limit {
232 return nil, connect.NewError(connect.CodeResourceExhausted, fmt.Errorf("the artifact is over the size of one artifact (%d bytes)", limit))
233 }
234 if !st.Receive() {
235 break
236 }
237 msg = st.Msg()
238 if msg.GetJobId() != "" || msg.GetArtifactId() != 0 || len(msg.GetBlockId()) != 0 {
239 return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("ids in a message after the first"))
240 }
241 }
242 if err := st.Err(); err != nil {
243 return nil, err // a broken stream stores nothing
244 }
245 f.mu.Lock()
246 defer f.mu.Unlock()
247 if art.state != pending || f.arts[art.a.GetId()] != art {
248 return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("the artifact is no longer pending"))
249 }
250 art.blocks[string(first.GetBlockId())] = data.Bytes()
251 return connect.NewResponse(&ermite.UploadJobArtifactBlockResponse{Size: int64(data.Len())}), nil
252}
253
254func sameList(a, b [][]byte) bool {
255 if len(a) != len(b) {
256 return false
257 }
258 for i := range a {
259 if !bytes.Equal(a[i], b[i]) {
260 return false
261 }
262 }
263 return true
264}
265
266func (f *fakeErmite) CommitJobArtifact(_ context.Context, req *connect.Request[ermite.CommitJobArtifactRequest]) (*connect.Response[ermite.CommitJobArtifactResponse], error) {
267 f.mu.Lock()
268 defer f.mu.Unlock()
269 j, err := f.caller(req.Header(), req.Msg.GetJobId(), "CommitJobArtifact")
270 if err != nil {
271 return nil, err
272 }
273 art := f.arts[req.Msg.GetArtifactId()]
274 if art == nil || art.a.GetJobId() != j.job || art.state == finalized {
275 return nil, connect.NewError(connect.CodeNotFound, errors.New("not a pending artifact of this job"))
276 }
277 ids := req.Msg.GetBlockIds()
278 if art.state == committed {
279 if !sameList(ids, art.list) {
280 return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("committed with another block list"))
281 }
282 return connect.NewResponse(&ermite.CommitJobArtifactResponse{Size: art.a.GetSize(), Digest: art.a.GetDigest()}), nil
283 }
284 if len(ids) > maxBlocks {
285 return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("more than 50,000 block ids"))
286 }
287 var content []byte
288 for _, id := range ids {
289 b, ok := art.blocks[string(id)]
290 if !ok {
291 return nil, connect.NewError(connect.CodeInvalidArgument, fmt.Errorf("block %x was never uploaded", id))
292 }
293 content = append(content, b...)
294 }
295 art.content, art.list, art.blocks, art.state = content, ids, nil, committed
296 art.a.Size, art.a.Digest = int64(len(content)), digestOf(content)
297 return connect.NewResponse(&ermite.CommitJobArtifactResponse{Size: art.a.GetSize(), Digest: art.a.GetDigest()}), nil
298}
299
300func (f *fakeErmite) FinalizeJobArtifact(_ context.Context, req *connect.Request[ermite.FinalizeJobArtifactRequest]) (*connect.Response[ermite.FinalizeJobArtifactResponse], error) {
301 f.mu.Lock()
302 defer f.mu.Unlock()
303 j, err := f.caller(req.Header(), req.Msg.GetJobId(), "FinalizeJobArtifact")
304 if err != nil {
305 return nil, err
306 }
307 var art *fakeArt
308 for _, a := range f.arts {
309 if a.a.GetJobId() == j.job && a.a.GetRunId() == j.run && a.a.GetName() == req.Msg.GetName() && (art == nil || a.a.GetId() > art.a.GetId()) {
310 art = a
311 }
312 }
313 switch {
314 case art == nil:
315 return nil, connect.NewError(connect.CodeNotFound, errors.New("no pending artifact of that name"))
316 case art.state == finalized:
317 return connect.NewResponse(&ermite.FinalizeJobArtifactResponse{Artifact: proto.CloneOf(art.a)}), nil
318 case art.state == pending:
319 return nil, connect.NewError(connect.CodeFailedPrecondition, errors.New("the artifact is not committed"))
320 case req.Msg.GetDigest() != "" && req.Msg.GetDigest() != art.a.GetDigest():
321 return nil, connect.NewError(connect.CodeInvalidArgument, errors.New("digest differs from the content"))
322 }
323 for id, a := range f.arts {
324 if a == art || a.state != finalized || a.a.GetRunId() != j.run || a.a.GetName() != art.a.GetName() {
325 continue
326 }
327 if a.a.GetAttempt() == j.attempt {
328 return nil, connect.NewError(connect.CodeAlreadyExists, errors.New(msgExists))
329 }
330 delete(f.arts, id) // a re-run replaces the earlier attempt's artifact
331 }
332 art.state = finalized
333 return connect.NewResponse(&ermite.FinalizeJobArtifactResponse{Artifact: proto.CloneOf(art.a)}), nil
334}
335
336func (f *fakeErmite) ListJobArtifacts(_ context.Context, req *connect.Request[ermite.ListJobArtifactsRequest]) (*connect.Response[ermite.ListJobArtifactsResponse], error) {
337 f.mu.Lock()
338 defer f.mu.Unlock()
339 j, err := f.caller(req.Header(), req.Msg.GetJobId(), "ListJobArtifacts")
340 if err != nil {
341 return nil, err
342 }
343 out := &ermite.ListJobArtifactsResponse{}
344 for _, a := range f.arts {
345 if a.state != finalized || a.a.GetRunId() != j.run ||
346 (req.Msg.GetName() != "" && a.a.GetName() != req.Msg.GetName()) ||
347 (req.Msg.GetArtifactId() != 0 && a.a.GetId() != req.Msg.GetArtifactId()) {
348 continue
349 }
350 out.Artifacts = append(out.Artifacts, proto.CloneOf(a.a))
351 }
352 sort.Slice(out.Artifacts, func(i, k int) bool { return out.Artifacts[i].GetId() < out.Artifacts[k].GetId() })
353 return connect.NewResponse(out), nil
354}
355
356func (f *fakeErmite) DownloadJobArtifact(_ context.Context, req *connect.Request[ermite.DownloadJobArtifactRequest], st *connect.ServerStream[ermite.DownloadJobArtifactResponse]) error {
357 f.mu.Lock()
358 j, err := f.caller(req.Header(), req.Msg.GetJobId(), "DownloadJobArtifact")
359 var art *fakeArt
360 if err == nil {
361 art = f.arts[req.Msg.GetArtifactId()]
362 if art == nil || art.state != finalized || art.a.GetRunId() != j.run {
363 err = connect.NewError(connect.CodeNotFound, errors.New("artifact not found"))
364 }
365 }
366 var meta *ermite.Artifact
367 var content []byte
368 if err == nil {
369 meta, content = proto.CloneOf(art.a), art.content
370 }
371 cut := f.cut
372 f.mu.Unlock()
373 if err != nil {
374 return err
375 }
376 off := req.Msg.GetOffset()
377 if off < 0 || off > int64(len(content)) {
378 return connect.NewError(connect.CodeOutOfRange, errors.New("offset beyond the size"))
379 }
380 if err := st.Send(&ermite.DownloadJobArtifactResponse{Artifact: meta}); err != nil {
381 return err
382 }
383 var sent int64
384 for p := content[off:]; len(p) > 0; {
385 if cut != nil {
386 if err := cut(off, sent); err != nil {
387 return err
388 }
389 }
390 n := min(len(p), chunkSize)
391 if err := st.Send(&ermite.DownloadJobArtifactResponse{Chunk: p[:n]}); err != nil {
392 return err
393 }
394 sent += int64(n)
395 p = p[n:]
396 }
397 return nil
398}
399
400func (f *fakeErmite) DeleteJobArtifact(_ context.Context, req *connect.Request[ermite.DeleteJobArtifactRequest]) (*connect.Response[ermite.DeleteJobArtifactResponse], error) {
401 f.mu.Lock()
402 defer f.mu.Unlock()
403 j, err := f.caller(req.Header(), req.Msg.GetJobId(), "DeleteJobArtifact")
404 if err != nil {
405 return nil, err
406 }
407 for id, a := range f.arts {
408 if a.state == finalized && a.a.GetRunId() == j.run && a.a.GetName() == req.Msg.GetName() {
409 delete(f.arts, id)
410 return connect.NewResponse(&ermite.DeleteJobArtifactResponse{ArtifactId: id}), nil
411 }
412 }
413 return nil, connect.NewError(connect.CodeNotFound, errors.New("artifact not found"))
414}
415
416// content returns the content of a finalized or committed artifact.
417func (f *fakeErmite) content(id int64) ([]byte, bool) {
418 f.mu.Lock()
419 defer f.mu.Unlock()
420 a := f.arts[id]
421 if a == nil || a.state == pending {
422 return nil, false
423 }
424 return a.content, true
425}
426
427// pendingBlocks returns the number of blocks a pending artifact holds.
428func (f *fakeErmite) pendingBlocks(id int64) int {
429 f.mu.Lock()
430 defer f.mu.Unlock()
431 if a := f.arts[id]; a != nil {
432 return len(a.blocks)
433 }
434 return -1
435}