1// SPDX-License-Identifier: MIT
2
3package registryproxy
4
5import (
6 "bytes"
7 "context"
8 "crypto/sha256"
9 "fmt"
10 "io"
11 "net/http"
12 "net/http/httptest"
13 "os"
14 "path/filepath"
15 "strings"
16 "sync"
17 "sync/atomic"
18 "testing"
19 "time"
20)
21
22// TestMain fails the run if any store hit an accounting bug (unpinLocked of
23// an unpinned entry), whichever test caused it.
24func TestMain(m *testing.M) {
25 var bugs atomic.Int32
26 defaultBugHook = func(msg string) {
27 bugs.Add(1)
28 fmt.Fprintln(os.Stderr, "registry cache bug:", msg)
29 }
30 code := m.Run()
31 if n := bugs.Load(); n > 0 && code == 0 {
32 fmt.Fprintf(os.Stderr, "FAIL: %d registry cache accounting bugs\n", n)
33 code = 1
34 }
35 os.Exit(code)
36}
37
38const manifestBody = `{"schemaVersion":2}`
39const blobBody = "a layer of bytes"
40
41func sum(b string) string { return fmt.Sprintf("sha256:%x", sha256.Sum256([]byte(b))) }
42
43// storeCfg is a complete configuration for a test that opens a store directly.
44func storeCfg(dir string, max int64) storeConfig {
45 return storeConfig{dir: dir, max: max, now: time.Now, grace: time.Second, maxDownloads: 16, slotWait: time.Second}
46}
47
48func mustOpen(t *testing.T, c storeConfig) *store {
49 t.Helper()
50 s, err := openStore(c)
51 if err != nil {
52 t.Fatal(err)
53 }
54 return s
55}
56
57// put stores b as a blob.
58func put(t *testing.T, s *store, b string) {
59 t.Helper()
60 f, _, _, err := s.ensure(context.Background(), kindBlob, sum(b), func(_ context.Context, w io.Writer) (string, error) { _, err := io.WriteString(w, b); return "", err }, keepalive{})
61 if err != nil {
62 t.Fatal(err)
63 }
64 _ = f.Close()
65}
66
67// putManifest stores m as a manifest of media type media.
68func putManifest(t *testing.T, s *store, m, media string) {
69 t.Helper()
70 f, _, _, err := s.ensure(context.Background(), kindManifest, sum(m), func(_ context.Context, w io.Writer) (string, error) {
71 _, err := io.WriteString(w, m)
72 return media, err
73 }, keepalive{})
74 if err != nil {
75 t.Fatal(err)
76 }
77 _ = f.Close()
78}
79
80// accounted checks that the store's total is exactly the sum of its entries and
81// within Max, and returns the number of tag files on disk.
82func accounted(t *testing.T, s *store) int {
83 t.Helper()
84 s.mu.Lock()
85 defer s.mu.Unlock()
86 var sum int64
87 tags := 0
88 for k, e := range s.entries {
89 sum += e.size
90 if e.size < minCharge {
91 t.Errorf("%s charged %d, under the minimum", k, e.size)
92 }
93 if e.kind == kindTag {
94 tags++
95 if e.size != minCharge {
96 t.Errorf("tag %s charged %d", k, e.size)
97 }
98 }
99 }
100 if sum != s.total || s.total+s.pending > s.max {
101 t.Fatalf("accounting: sum=%d total=%d pending=%d max=%d", sum, s.total, s.pending, s.max)
102 }
103 files, err := os.ReadDir(filepath.Join(s.dir, kindTag.dir()))
104 if err != nil {
105 t.Fatal(err)
106 }
107 if len(files) != tags {
108 t.Fatalf("%d tag files, %d tag entries", len(files), tags)
109 }
110 return tags
111}
112
113// waitFor polls cond (under the store lock) until it holds.
114func waitFor(t *testing.T, s *store, what string, cond func() bool) {
115 t.Helper()
116 for end := time.Now().Add(5 * time.Second); ; {
117 s.mu.Lock()
118 ok := cond()
119 s.mu.Unlock()
120 if ok {
121 return
122 }
123 if time.Now().After(end) {
124 t.Fatalf("timed out waiting for %s", what)
125 }
126 time.Sleep(time.Millisecond)
127 }
128}
129
130// waiters is the number of requests waiting for the flight of key (-1: none).
131// The caller holds the store lock.
132func waiters(s *store, key string) int {
133 if fl, ok := s.flights[key]; ok {
134 return fl.waiters
135 }
136 return -1
137}
138
139// upstreamWith fakes Docker Hub: /token answers an anonymous token and h serves
140// everything else.
141func upstreamWith(t *testing.T, h http.HandlerFunc) *httptest.Server {
142 t.Helper()
143 srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
144 if r.URL.Path == "/token" {
145 _, _ = io.WriteString(w, `{"token":"public","expires_in":300}`)
146 return
147 }
148 h(w, r)
149 }))
150 t.Cleanup(srv.Close)
151 return srv
152}
153
154func proxyWith(t *testing.T, srv *httptest.Server, c Config) *Proxy {
155 t.Helper()
156 if c.Dir == "" {
157 c.Dir = t.TempDir()
158 }
159 if c.Max == 0 {
160 c.Max = 8 << 20
161 }
162 c.Upstream, c.AuthURL = srv.URL, srv.URL+"/token"
163 p, err := New(c)
164 if err != nil {
165 t.Fatal(err)
166 }
167 return p
168}
169
170// fakeRegistry is a Docker Hub with one manifest and one blob, served from a
171// CDN behind a redirect. It counts the requests it gets.
172type fakeRegistry struct {
173 server *httptest.Server
174 heads, gets, blobs, tokens atomic.Int32
175 status atomic.Int32
176 corrupt atomic.Bool
177}
178
179func fake(t *testing.T) *fakeRegistry {
180 t.Helper()
181 f := &fakeRegistry{}
182 cdn := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
183 if r.Header.Get("Authorization") != "" {
184 t.Error("bearer leaked to CDN")
185 }
186 w.Header().Set("Content-Length", fmt.Sprint(len(blobBody)))
187 if r.Method != http.MethodHead {
188 b := blobBody
189 if f.corrupt.Load() {
190 b = strings.Repeat("x", len(b))
191 }
192 _, _ = io.WriteString(w, b)
193 }
194 }))
195 t.Cleanup(cdn.Close)
196 f.server = httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
197 if r.URL.Path == "/token" {
198 f.tokens.Add(1)
199 if r.Header.Get("Authorization") != "" {
200 t.Error("guest credentials forwarded to auth")
201 }
202 if r.URL.Query().Get("scope") != "repository:library/alpine:pull" {
203 t.Error("wrong scope")
204 }
205 _, _ = io.WriteString(w, `{"token":"public","expires_in":300}`)
206 return
207 }
208 if r.Header.Get("Authorization") != "Bearer public" {
209 t.Errorf("wrong bearer %q", r.Header.Get("Authorization"))
210 }
211 if s := f.status.Load(); s != 0 {
212 w.WriteHeader(int(s))
213 return
214 }
215 if strings.Contains(r.URL.Path, "/manifests/") {
216 w.Header().Set("Content-Type", "application/vnd.oci.image.index.v1+json")
217 w.Header().Set("Docker-Content-Digest", sum(manifestBody))
218 w.Header().Set("Content-Length", fmt.Sprint(len(manifestBody)))
219 if r.Method == http.MethodHead {
220 f.heads.Add(1)
221 } else {
222 f.gets.Add(1)
223 _, _ = io.WriteString(w, manifestBody)
224 }
225 return
226 }
227 if strings.Contains(r.URL.Path, "/blobs/") {
228 if r.Method == http.MethodGet {
229 f.blobs.Add(1)
230 }
231 http.Redirect(w, r, cdn.URL+"/layer", http.StatusTemporaryRedirect)
232 return
233 }
234 http.NotFound(w, r)
235 }))
236 t.Cleanup(f.server.Close)
237 return f
238}
239
240func proxyFor(t *testing.T, f *fakeRegistry) *Proxy {
241 t.Helper()
242 p, err := New(Config{Dir: t.TempDir(), Max: 8 << 20, Upstream: f.server.URL, AuthURL: f.server.URL + "/token"})
243 if err != nil {
244 t.Fatal(err)
245 }
246 return p
247}
248
249// request serves one request with guest credentials that must never reach
250// upstream; headers are name, value pairs.
251func request(t *testing.T, p http.Handler, method, path string, headers ...string) *httptest.ResponseRecorder {
252 t.Helper()
253 r := httptest.NewRequest(method, path, nil)
254 r.Header.Set("Authorization", "Bearer guest-secret")
255 for i := 0; i < len(headers); i += 2 {
256 r.Header.Set(headers[i], headers[i+1])
257 }
258 w := httptest.NewRecorder()
259 p.ServeHTTP(w, r)
260 return w
261}
262
263// stallingUpstream serves a blob's first bytes and then holds the body until
264// the request is cancelled; it counts started and cancelled downloads.
265type stallingUpstream struct {
266 srv *httptest.Server
267 started, cancelled atomic.Int32
268}
269
270func newStallingUpstream(t *testing.T) *stallingUpstream {
271 t.Helper()
272 u := &stallingUpstream{}
273 u.srv = upstreamWith(t, func(w http.ResponseWriter, r *http.Request) {
274 u.started.Add(1)
275 w.Header().Set("Content-Length", "1000000")
276 _, _ = io.WriteString(w, "partial")
277 w.(http.Flusher).Flush()
278 select {
279 case <-r.Context().Done():
280 u.cancelled.Add(1)
281 case <-time.After(10 * time.Second):
282 }
283 })
284 return u
285}
286
287func blobPath(i int) string { return "/v2/library/alpine/blobs/" + sum(fmt.Sprint("blob-", i)) }
288
289// serveCancelled starts a request and cancels it once its download has begun.
290func serveCancelled(t *testing.T, p *Proxy, u *stallingUpstream, path string, started int32) {
291 t.Helper()
292 ctx, cancel := context.WithCancel(context.Background())
293 done := make(chan struct{})
294 go func() {
295 defer close(done)
296 p.ServeHTTP(httptest.NewRecorder(), httptest.NewRequest(http.MethodGet, path, nil).WithContext(ctx))
297 }()
298 waitFor(t, p.s, "the download to start", func() bool { return u.started.Load() >= started })
299 cancel()
300 <-done
301}
302
303// serveAsync serves path with ctx on its own goroutine, into w (a new
304// recorder when nil); the status code comes on the returned channel.
305func serveAsync(p *Proxy, ctx context.Context, path string, w *keepAliveRecorder) <-chan int {
306 if w == nil {
307 w = &keepAliveRecorder{ResponseRecorder: httptest.NewRecorder()}
308 }
309 code := make(chan int, 1)
310 go func() {
311 p.ServeHTTP(w, httptest.NewRequest(http.MethodGet, path, nil).WithContext(ctx))
312 code <- w.Code
313 }()
314 return code
315}
316
317type keepAliveRecorder struct {
318 *httptest.ResponseRecorder
319 calls atomic.Int32
320}
321
322func (k *keepAliveRecorder) KeepAlive() { k.calls.Add(1) }
323
324// waitKeptAlive waits until w has been kept alive at least once.
325func waitKeptAlive(t *testing.T, w *keepAliveRecorder) {
326 t.Helper()
327 for end := time.Now().Add(5 * time.Second); w.calls.Load() == 0; time.Sleep(time.Millisecond) {
328 if time.Now().After(end) {
329 t.Fatal("request never kept alive")
330 }
331 }
332}
333
334// blockingWriter holds a response in its first Write until release is closed.
335type blockingWriter struct {
336 *httptest.ResponseRecorder
337 started, release chan struct{}
338 once sync.Once
339}
340
341func (b *blockingWriter) Write(p []byte) (int, error) {
342 b.once.Do(func() { close(b.started) })
343 <-b.release
344 return b.ResponseRecorder.Write(p)
345}
346
347// slowBlob streams body in chunks, one every gap.
348func slowBlob(body string, chunks int, gap time.Duration) http.HandlerFunc {
349 return func(w http.ResponseWriter, r *http.Request) {
350 w.Header().Set("Content-Length", fmt.Sprint(len(body)))
351 size := len(body) / chunks
352 for i := 0; i < len(body); i += size {
353 _, _ = io.WriteString(w, body[i:min(i+size, len(body))])
354 w.(http.Flusher).Flush()
355 select {
356 case <-time.After(gap):
357 case <-r.Context().Done():
358 return
359 }
360 }
361 }
362}
363
364type syncBuffer struct {
365 mu sync.Mutex
366 b bytes.Buffer
367}
368
369func (s *syncBuffer) Write(p []byte) (int, error) {
370 s.mu.Lock()
371 defer s.mu.Unlock()
372 return s.b.Write(p)
373}
374
375func (s *syncBuffer) String() string { s.mu.Lock(); defer s.mu.Unlock(); return s.b.String() }