diff --git a/streampush/README.md b/streampush/README.md new file mode 100644 index 0000000..fd298e9 --- /dev/null +++ b/streampush/README.md @@ -0,0 +1,96 @@ +# streampush + +Push a container image whose single layer is **generated as it uploads** and +**never buffered**. The goal: stream a 200GB layer to a registry using a small, +constant amount of memory, with the layer digest finalized only once the upload +completes. + +``` +go run . --size 200GB --data random --compression none --registry stream +``` + +## How it works + +The whole pipeline is single-pass and back-pressured by the network: + +1. **Generator** (`gen.go`) — a `genReader` produces exactly `--size` bytes on + demand into each `Read` buffer (zeros or a fast allocation-free xorshift + PRNG). It never allocates per-read and never holds more than one buffer. +2. **`stream.Layer`** (go-containerregistry's `pkg/v1/stream`) — wraps the + generator, gzips it on the fly through an `io.Pipe`, and computes the + compressed digest, the uncompressed `diffID`, and the size **during the + single streaming pass**. Until the stream is fully consumed, `Digest()`, + `DiffID()` and `Size()` return `stream.ErrNotComputed`. +3. **`remote.Write`** — the pusher sees `stream.ErrNotComputed`, so it uploads + the layer first via one chunked `PATCH` (no `Content-Length`, no retry + buffer), then derives the config + manifest from the now-known digests and + `PUT`s them. The layer body is sent straight from the gzip pipe to the socket. + +So memory stays flat regardless of layer size. A 4 GiB and a 32 GiB push use the +same peak heap: + +``` +=== 4GiB === elapsed: 5.3s peak Go heap during push: 50.53 MiB +=== 32GiB === elapsed: 42.7s peak Go heap during push: 57.51 MiB +``` + +(The peak heap includes the in-process registry, since `--registry stream` runs +in the same process.) + +## About `pkg/registry` + +The prompt asked to use go-containerregistry's `pkg/registry` as a Docker Hub +stand-in *if it supports incremental uploads*. It speaks the incremental upload +**protocol** (`POST` → `PATCH` → `PUT`), and you can select it with +`--registry pkg` for small images. But it is **not** a non-buffering target: + +- Its `PATCH` handler does `io.Copy(&bytes.Buffer{}, req.Body)`, reading the + entire chunk into memory (`pkg/registry/blobs.go`). +- Its default blob handler stores each blob in a `map[string][]byte`. + +A 200GB layer would therefore require ~200GB of RAM on the server. `streampush` +prints a warning and will likely OOM if you point `--registry pkg` at a +multi-gigabyte `--size`. + +To actually receive a 200GB layer, this repo includes **`sinkRegistry`** +(`registry.go`): a minimal OCI distribution registry that streams the +`PATCH`/`PUT` body through a `sha256` hasher with a fixed staging buffer, +retaining only a bounded prefix (`--retain-below`, default 32 MiB) so the config +and manifest stay pullable while the giant layer is digested-and-dropped. This +is the default (`--registry stream`). + +## Flags + +| flag | default | meaning | +| --- | --- | --- | +| `--size` | `200GB` | uncompressed layer size (`200GB`, `512MiB`, raw bytes, ...) | +| `--registry` | `stream` | `stream` (non-buffering sink), `pkg` (buffers!), or a `host[:port]` to push to a real registry | +| `--data` | `zero` | `zero` (fast/compressible) or `random` (incompressible, CPU heavy) | +| `--compression` | `none` | `none`, `speed`, `default`, `best`, or `0`–`9` | +| `--retain-below` | `32MiB` | sink retains blob content below this size | +| `--repo` / `--tag` | `streampush/big` / `latest` | image coordinates | +| `--insecure` | `false` | use http for an explicit `--registry` host | +| `--verbose` | `false` | log registry + push activity | + +### Wire size vs. generated size + +`stream.Layer` always gzips. To put real bytes on the wire: + +- `--data zero --compression none` — fast; ~200GB of (gzip-wrapped, undeflated) + zeros actually traverse the socket. Good for plumbing/throughput tests. +- `--data random --compression none` — incompressible bytes, on-wire size tracks + generated size, low CPU. +- `--data random --compression best` — realistic CPU cost; the digest is of the + compressed stream. + +### Pushing to a real registry + +``` +# e.g. a local registry:2 container, or Docker Hub (with auth configured) +go run . --size 1GiB --data random --registry localhost:5000 --insecure +go run . --size 1GiB --data random --registry index.docker.io --repo youruser/streamtest +``` + +`remote.Write` uses the default keychain, so `docker login` credentials apply. +Note that real registries enforce blob/layer size limits; Docker Hub will not +accept a 200GB layer. diff --git a/streampush/gen.go b/streampush/gen.go new file mode 100644 index 0000000..cb12bca --- /dev/null +++ b/streampush/gen.go @@ -0,0 +1,131 @@ +package main + +import ( + "fmt" + "io" + "strconv" + "strings" + "sync/atomic" +) + +// genReader produces exactly `remaining` bytes on demand, filling each Read +// buffer with `fill` and incrementing `counter` atomically. It never allocates +// per Read and never buffers content, so it can back an arbitrarily large layer +// with constant memory. +type genReader struct { + remaining int64 + fill func([]byte) + counter *int64 +} + +func (g *genReader) Read(p []byte) (int, error) { + if g.remaining <= 0 { + return 0, io.EOF + } + n := len(p) + if int64(n) > g.remaining { + n = int(g.remaining) + } + g.fill(p[:n]) + g.remaining -= int64(n) + atomic.AddInt64(g.counter, int64(n)) + return n, nil +} + +// Close implements io.Closer so genReader satisfies io.ReadCloser, which is +// what stream.NewLayer expects. +func (g *genReader) Close() error { return nil } + +// fillFunc returns a buffer-filling function for the requested data mode. +// +// - "zero": write zero bytes. Trivial CPU; gzip crushes it (use with +// --compression none to still put real bytes on the wire). +// - "random": fill with a fast, allocation-free xorshift PRNG. Incompressible, +// so the on-wire size tracks the generated size even under gzip, at the cost +// of more CPU. +func fillFunc(mode string) (func([]byte), error) { + switch strings.ToLower(mode) { + case "zero", "zeros", "0": + return func(b []byte) { clear(b) }, nil + case "random", "rand": + // One state per returned closure; Reads are sequential within a layer. + state := uint64(0x9e3779b97f4a7c15) + return func(b []byte) { + i := 0 + for ; i+8 <= len(b); i += 8 { + state ^= state << 13 + state ^= state >> 7 + state ^= state << 17 + b[i+0] = byte(state) + b[i+1] = byte(state >> 8) + b[i+2] = byte(state >> 16) + b[i+3] = byte(state >> 24) + b[i+4] = byte(state >> 32) + b[i+5] = byte(state >> 40) + b[i+6] = byte(state >> 48) + b[i+7] = byte(state >> 56) + } + if i < len(b) { + state ^= state << 13 + state ^= state >> 7 + state ^= state << 17 + for j := 0; i < len(b); i, j = i+1, j+1 { + b[i] = byte(state >> (8 * j)) + } + } + }, nil + default: + return nil, fmt.Errorf("unknown data mode %q (want zero or random)", mode) + } +} + +// parseSize parses human-friendly sizes: a bare integer (bytes), or a number +// with a decimal (KB/MB/GB/TB) or binary (KiB/MiB/GiB/TiB) unit suffix. +func parseSize(s string) (int64, error) { + s = strings.TrimSpace(s) + if s == "" { + return 0, fmt.Errorf("empty size") + } + upper := strings.ToUpper(s) + + type unit struct { + suffix string + mult int64 + } + // Order matters: check longer/binary suffixes before shorter decimal ones. + units := []unit{ + {"TIB", 1 << 40}, {"GIB", 1 << 30}, {"MIB", 1 << 20}, {"KIB", 1 << 10}, + {"TB", 1e12}, {"GB", 1e9}, {"MB", 1e6}, {"KB", 1e3}, + {"T", 1 << 40}, {"G", 1 << 30}, {"M", 1 << 20}, {"K", 1 << 10}, + {"B", 1}, + } + for _, u := range units { + if strings.HasSuffix(upper, u.suffix) { + num := strings.TrimSpace(upper[:len(upper)-len(u.suffix)]) + f, err := strconv.ParseFloat(num, 64) + if err != nil { + return 0, fmt.Errorf("parse %q: %w", s, err) + } + return int64(f * float64(u.mult)), nil + } + } + n, err := strconv.ParseInt(upper, 10, 64) + if err != nil { + return 0, fmt.Errorf("parse %q: %w", s, err) + } + return n, nil +} + +// humanBytes formats a byte count using binary (IEC) units. +func humanBytes(n int64) string { + const unit = 1024 + if n < unit { + return fmt.Sprintf("%d B", n) + } + div, exp := int64(unit), 0 + for x := n / unit; x >= unit; x /= unit { + div *= unit + exp++ + } + return fmt.Sprintf("%.2f %ciB", float64(n)/float64(div), "KMGTPE"[exp]) +} diff --git a/streampush/go.mod b/streampush/go.mod new file mode 100644 index 0000000..c1cd4d5 --- /dev/null +++ b/streampush/go.mod @@ -0,0 +1,17 @@ +module github.com/imjasonh/terraform-playground/streampush + +go 1.26.3 + +require github.com/google/go-containerregistry v0.21.6 + +require ( + github.com/docker/cli v29.4.3+incompatible // indirect + github.com/docker/docker-credential-helpers v0.9.3 // indirect + github.com/klauspost/compress v1.18.6 // indirect + github.com/opencontainers/go-digest v1.0.0 // indirect + github.com/opencontainers/image-spec v1.1.1 // indirect + github.com/sirupsen/logrus v1.9.4 // indirect + golang.org/x/sync v0.20.0 // indirect + golang.org/x/sys v0.44.0 // indirect + gotest.tools/v3 v3.5.2 // indirect +) diff --git a/streampush/go.sum b/streampush/go.sum new file mode 100644 index 0000000..ac9f39a --- /dev/null +++ b/streampush/go.sum @@ -0,0 +1,34 @@ +github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c= +github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38= +github.com/docker/cli v29.4.3+incompatible h1:u+UliYm2J/rYrIh2FqHQg32neRG8GjbvNuwQRTzGspU= +github.com/docker/cli v29.4.3+incompatible/go.mod h1:JLrzqnKDaYBop7H2jaqPtU4hHvMKP+vjCwu2uszcLI8= +github.com/docker/docker-credential-helpers v0.9.3 h1:gAm/VtF9wgqJMoxzT3Gj5p4AqIjCBS4wrsOh9yRqcz8= +github.com/docker/docker-credential-helpers v0.9.3/go.mod h1:x+4Gbw9aGmChi3qTLZj8Dfn0TD20M/fuWy0E5+WDeCo= +github.com/google/go-cmp v0.7.0 h1:wk8382ETsv4JYUZwIsn6YpYiWiBsYLSJiTsyBybVuN8= +github.com/google/go-cmp v0.7.0/go.mod h1:pXiqmnSA92OHEEa9HXL2W4E7lf9JzCmGVUdgjX3N/iU= +github.com/google/go-containerregistry v0.21.6 h1:T+yqQIlJXKrM98Om4DlW3GoWQAmhZuLMwoDOvVrtiUM= +github.com/google/go-containerregistry v0.21.6/go.mod h1:U7MMSBIJynke2MVQrQk19NP9k/uQsGz/h0amIFSHMbo= +github.com/klauspost/compress v1.18.6 h1:2jupLlAwFm95+YDR+NwD2MEfFO9d4z4Prjl1XXDjuao= +github.com/klauspost/compress v1.18.6/go.mod h1:cwPg85FWrGar70rWktvGQj8/hthj3wpl0PGDogxkrSQ= +github.com/opencontainers/go-digest v1.0.0 h1:apOUWs51W5PlhuyGyz9FCeeBIOUDA/6nW8Oi/yOhh5U= +github.com/opencontainers/go-digest v1.0.0/go.mod h1:0JzlMkj0TRzQZfJkVvzbP0HBR3IKzErnv2BNG4W4MAM= +github.com/opencontainers/image-spec v1.1.1 h1:y0fUlFfIZhPF1W537XOLg0/fcx6zcHCJwooC2xJA040= +github.com/opencontainers/image-spec v1.1.1/go.mod h1:qpqAh3Dmcf36wStyyWU+kCeDgrGnAve2nCC8+7h8Q0M= +github.com/pmezard/go-difflib v1.0.0 h1:4DBwDE0NGyQoBHbLQYPwSUPoCMWR5BEzIk/f1lZbAQM= +github.com/pmezard/go-difflib v1.0.0/go.mod h1:iKH77koFhYxTK1pcRnkKkqfTogsbg7gZNVY4sRDYZ/4= +github.com/sirupsen/logrus v1.9.4 h1:TsZE7l11zFCLZnZ+teH4Umoq5BhEIfIzfRDZ1Uzql2w= +github.com/sirupsen/logrus v1.9.4/go.mod h1:ftWc9WdOfJ0a92nsE2jF5u5ZwH8Bv2zdeOC42RjbV2g= +github.com/stretchr/testify v1.10.0 h1:Xv5erBjTwe/5IxqUQTdXv5kgmIvbHo3QQyRwhJsOfJA= +github.com/stretchr/testify v1.10.0/go.mod h1:r2ic/lqez/lEtzL7wO/rwa5dbSLXVDPFyf8C91i36aY= +golang.org/x/mod v0.36.0 h1:JJjpVx6myfUsUdAzZuOSTTmRE0PfZeNWzzvKrP7amb4= +golang.org/x/mod v0.36.0/go.mod h1:moc6ELqsWcOw5Ef3xVprK5ul/MvtVvkIXLziUOICjUQ= +golang.org/x/sync v0.20.0 h1:e0PTpb7pjO8GAtTs2dQ6jYa5BWYlMuX047Dco/pItO4= +golang.org/x/sync v0.20.0/go.mod h1:9xrNwdLfx4jkKbNva9FpL6vEN7evnE43NNNJQ2LF3+0= +golang.org/x/sys v0.44.0 h1:ildZl3J4uzeKP07r2F++Op7E9B29JRUy+a27EibtBTQ= +golang.org/x/sys v0.44.0/go.mod h1:4GL1E5IUh+htKOUEOaiffhrAeqysfVGipDYzABqnCmw= +golang.org/x/tools v0.45.0 h1:18qN3FAooORvApf5XjCXgsuayZOEtXf6JK18I3+ONa8= +golang.org/x/tools v0.45.0/go.mod h1:LuUGqqaXcXMEFEruIVJVm5mgDD8vww/z/SR1gQ4uE/0= +gopkg.in/yaml.v3 v3.0.1 h1:fxVm/GzAzEWqLHuvctI91KS9hhNmmWOoWu0XTYJS7CA= +gopkg.in/yaml.v3 v3.0.1/go.mod h1:K4uyk7z7BCEPqu6E+C64Yfv1cQ7kz7rIZviUmN+EgEM= +gotest.tools/v3 v3.5.2 h1:7koQfIKdy+I8UTetycgUqXWSDwpgv193Ka+qRsmBY8Q= +gotest.tools/v3 v3.5.2/go.mod h1:LtdLGcnqToBH83WByAAi/wiwSFCArdFIUV/xxN4pcjA= diff --git a/streampush/main.go b/streampush/main.go new file mode 100644 index 0000000..bd61899 --- /dev/null +++ b/streampush/main.go @@ -0,0 +1,236 @@ +// Command streampush pushes a container image whose single layer is generated +// on the fly as it uploads. The layer content is never buffered: bytes are +// produced by a generator, gzipped by go-containerregistry's stream.Layer, and +// streamed to the registry in a single chunked PATCH. The layer's digest and +// diffID are only known once the upload completes. +// +// This lets you push, e.g., a 200GB layer using a tiny, constant amount of +// memory on both the client and (with the default streaming registry) the +// server. +package main + +import ( + "compress/gzip" + "context" + "flag" + "fmt" + "log" + "net/http/httptest" + "os" + "runtime" + "strconv" + "strings" + "sync/atomic" + "time" + + "github.com/google/go-containerregistry/pkg/name" + "github.com/google/go-containerregistry/pkg/registry" + v1 "github.com/google/go-containerregistry/pkg/v1" + "github.com/google/go-containerregistry/pkg/v1/empty" + "github.com/google/go-containerregistry/pkg/v1/mutate" + "github.com/google/go-containerregistry/pkg/v1/remote" + "github.com/google/go-containerregistry/pkg/v1/stream" +) + +func main() { + var ( + sizeStr = flag.String("size", "200GB", "uncompressed layer size to generate (e.g. 200GB, 512MiB, 1073741824)") + repo = flag.String("repo", "streampush/big", "repository name to push to") + tag = flag.String("tag", "latest", "tag to push") + registryArg = flag.String("registry", "stream", `target registry: "stream" (built-in non-buffering sink), "pkg" (go-containerregistry in-memory registry; buffers!), or an explicit host[:port] to push to a real registry`) + insecure = flag.Bool("insecure", false, "use http (and skip TLS) when pushing to an explicit --registry host") + dataMode = flag.String("data", "zero", `layer byte pattern: "zero" (fast, compressible) or "random" (incompressible, CPU heavy)`) + compArg = flag.String("compression", "none", `gzip level: "none", "speed", "default", "best", or 0-9`) + smallLimit = flag.String("retain-below", "32MiB", "sink registry retains blob content below this size (so config/manifest stay pullable)") + verbose = flag.Bool("verbose", false, "verbose registry + push logging") + ) + flag.Parse() + + size, err := parseSize(*sizeStr) + if err != nil { + fatalf("invalid --size: %v", err) + } + retainBelow, err := parseSize(*smallLimit) + if err != nil { + fatalf("invalid --retain-below: %v", err) + } + level, err := parseCompression(*compArg) + if err != nil { + fatalf("invalid --compression: %v", err) + } + fill, err := fillFunc(*dataMode) + if err != nil { + fatalf("invalid --data: %v", err) + } + + ctx := context.Background() + + // Resolve the target registry host, standing up an in-process server if needed. + host, nameOpts, cleanup := resolveRegistry(*registryArg, *insecure, retainBelow, *verbose, size) + defer cleanup() + + ref, err := name.ParseReference(fmt.Sprintf("%s/%s:%s", host, *repo, *tag), nameOpts...) + if err != nil { + fatalf("parse reference: %v", err) + } + + // The generator: produces `size` bytes on demand, counting as it goes. + var generated int64 + gen := &genReader{remaining: size, fill: fill, counter: &generated} + + // stream.Layer reads the generator exactly once, gzipping on the fly and + // computing the compressed digest, uncompressed diffID, and size during the + // single streaming pass. None of it is buffered. + layer := stream.NewLayer(gen, stream.WithCompressionLevel(level)) + + img, err := mutate.AppendLayers(empty.Image, layer) + if err != nil { + fatalf("append layer: %v", err) + } + + fmt.Printf("streampush: pushing %s\n", ref) + fmt.Printf(" layer size (uncompressed): %s (%d bytes)\n", humanBytes(size), size) + fmt.Printf(" data=%s compression=%s registry=%s\n", *dataMode, *compArg, *registryArg) + fmt.Println(" content is generated as it uploads; digest is finalized at the end.") + fmt.Println() + + // remote.WithProgress reports compressed bytes accepted by the registry. + progressCh := make(chan v1.Update, 64) + var pushed int64 + go func() { + for u := range progressCh { + atomic.StoreInt64(&pushed, u.Complete) + } + }() + + start := time.Now() + stopReport := make(chan struct{}) + var peakHeap int64 + go reportProgress(&generated, &pushed, &peakHeap, start, stopReport) + + writeOpts := []remote.Option{ + remote.WithContext(ctx), + remote.WithProgress(progressCh), + } + + err = remote.Write(ref, img, writeOpts...) + close(stopReport) + if err != nil { + fatalf("push failed: %v", err) + } + + // Ensure we have at least one heap sample even for very fast pushes. + var ms runtime.MemStats + runtime.ReadMemStats(&ms) + if h := int64(ms.HeapAlloc); h > atomic.LoadInt64(&peakHeap) { + atomic.StoreInt64(&peakHeap, h) + } + + // All values are now computed because the stream was fully consumed. + digest, _ := layer.Digest() + diffID, _ := layer.DiffID() + compressedSize, _ := layer.Size() + manifestDigest, _ := img.Digest() + elapsed := time.Since(start) + + fmt.Println() + fmt.Println("push complete.") + fmt.Printf(" generated (uncompressed): %s\n", humanBytes(atomic.LoadInt64(&generated))) + fmt.Printf(" uploaded (compressed): %s\n", humanBytes(compressedSize)) + fmt.Printf(" layer digest (blob): %s\n", digest) + fmt.Printf(" layer diffID (tar): %s\n", diffID) + fmt.Printf(" manifest: %s@%s\n", ref.Context().Name(), manifestDigest) + fmt.Printf(" elapsed: %s (%s/s generated)\n", elapsed.Round(time.Millisecond), + humanBytes(int64(float64(atomic.LoadInt64(&generated))/elapsed.Seconds()))) + fmt.Printf(" peak Go heap during push: %s (constant regardless of layer size)\n", + humanBytes(atomic.LoadInt64(&peakHeap))) +} + +// resolveRegistry returns the registry host to target, name parse options, and +// a cleanup function. For "stream" and "pkg" it stands up an in-process httptest +// server; for an explicit host it pushes to that registry directly. +func resolveRegistry(arg string, insecure bool, retainBelow int64, verbose bool, size int64) (string, []name.Option, func()) { + switch arg { + case "stream": + srv := httptest.NewServer(newSinkRegistry(retainBelow, verbose)) + host := strings.TrimPrefix(srv.URL, "http://") + return host, []name.Option{name.Insecure}, srv.Close + + case "pkg": + if size > 1<<30 { + fmt.Fprintf(os.Stderr, + "WARNING: --registry pkg buffers the entire layer in memory; %s will likely OOM.\n", + humanBytes(size)) + } + opts := []registry.Option{} + if verbose { + opts = append(opts, registry.Logger(log.New(os.Stderr, "[pkg-registry] ", log.LstdFlags))) + } + srv := httptest.NewServer(registry.New(opts...)) + host := strings.TrimPrefix(srv.URL, "http://") + return host, []name.Option{name.Insecure}, srv.Close + + default: + // Treat as an explicit registry host. + nameOpts := []name.Option{} + if insecure { + nameOpts = append(nameOpts, name.Insecure) + } + return arg, nameOpts, func() {} + } +} + +func reportProgress(generated, pushed, peakHeap *int64, start time.Time, stop <-chan struct{}) { + ticker := time.NewTicker(2 * time.Second) + defer ticker.Stop() + sample := time.NewTicker(250 * time.Millisecond) + defer sample.Stop() + var ms runtime.MemStats + for { + select { + case <-stop: + return + case <-sample.C: + runtime.ReadMemStats(&ms) + if h := int64(ms.HeapAlloc); h > atomic.LoadInt64(peakHeap) { + atomic.StoreInt64(peakHeap, h) + } + case <-ticker.C: + elapsed := time.Since(start).Seconds() + g := atomic.LoadInt64(generated) + p := atomic.LoadInt64(pushed) + rate := int64(0) + if elapsed > 0 { + rate = int64(float64(g) / elapsed) + } + fmt.Printf(" ... generated %s | uploaded %s | %s/s | heap %s\n", + humanBytes(g), humanBytes(p), humanBytes(rate), humanBytes(atomic.LoadInt64(peakHeap))) + } + } +} + +func parseCompression(s string) (int, error) { + switch strings.ToLower(s) { + case "none", "no", "store": + return gzip.NoCompression, nil + case "speed", "fast": + return gzip.BestSpeed, nil + case "default", "": + return gzip.DefaultCompression, nil + case "best", "max": + return gzip.BestCompression, nil + } + n, err := strconv.Atoi(s) + if err != nil { + return 0, fmt.Errorf("not a level: %q", s) + } + if n < gzip.HuffmanOnly || n > gzip.BestCompression { + return 0, fmt.Errorf("level %d out of range", n) + } + return n, nil +} + +func fatalf(format string, args ...any) { + fmt.Fprintf(os.Stderr, "streampush: "+format+"\n", args...) + os.Exit(1) +} diff --git a/streampush/registry.go b/streampush/registry.go new file mode 100644 index 0000000..a952a45 --- /dev/null +++ b/streampush/registry.go @@ -0,0 +1,359 @@ +package main + +import ( + "bytes" + "crypto/sha256" + "encoding/hex" + "fmt" + "hash" + "io" + "log" + "math/rand" + "net/http" + "strconv" + "strings" + "sync" +) + +// sinkRegistry is a minimal OCI distribution registry that supports the +// incremental ("chunked") blob upload protocol WITHOUT ever buffering the full +// blob in memory. +// +// The default in-memory registry shipped with go-containerregistry +// (pkg/registry) speaks the same protocol, but its PATCH handler does +// `io.Copy(&bytes.Buffer{}, req.Body)` and its blob handler stores the whole +// blob in a map[string][]byte. That makes it impossible to receive a 200GB +// layer without 200GB of RAM. +// +// sinkRegistry instead streams the PATCH/PUT body straight through a sha256 +// hasher into io.Discard, keeping only the running digest and byte count. Blobs +// smaller than smallBlobLimit (config, manifests) are retained so the resulting +// image is still describable/pullable; larger blobs are digested and dropped. +type sinkRegistry struct { + mu sync.Mutex + blobSizes map[string]int64 // digest -> size (everything we've committed) + smallBlobs map[string][]byte // digest -> content (only blobs <= smallBlobLimit) + uploads map[string]*upload + manifests map[string][]byte // repo+"/"+reference -> raw manifest + manifestCT map[string]string // repo+"/"+reference -> content type + + smallBlobLimit int64 + verbose bool +} + +type upload struct { + h hash.Hash + buf *bytes.Buffer // retained prefix, capped at limit + limit int64 + n int64 +} + +func newUpload(limit int64) *upload { + return &upload{h: sha256.New(), buf: &bytes.Buffer{}, limit: limit} +} + +// Write streams bytes through the digest while retaining at most `limit` bytes +// of prefix. This is the crux of the non-buffering design: io.Copy hands us the +// whole body, but we never hold more than `limit` of it in memory. +func (u *upload) Write(p []byte) (int, error) { + u.h.Write(p) + if room := u.limit - int64(u.buf.Len()); room > 0 { + if int64(len(p)) <= room { + u.buf.Write(p) + } else { + u.buf.Write(p[:room]) + } + } + u.n += int64(len(p)) + return len(p), nil +} + +// retained returns the full blob content if it fit entirely within the cap, and +// false if it was too large to keep (and was therefore digested-and-dropped). +func (u *upload) retained() ([]byte, bool) { + if u.n <= u.limit { + return u.buf.Bytes(), true + } + return nil, false +} + +func newSinkRegistry(smallBlobLimit int64, verbose bool) *sinkRegistry { + return &sinkRegistry{ + blobSizes: map[string]int64{}, + smallBlobs: map[string][]byte{}, + uploads: map[string]*upload{}, + manifests: map[string][]byte{}, + manifestCT: map[string]string{}, + smallBlobLimit: smallBlobLimit, + verbose: verbose, + } +} + +func (s *sinkRegistry) logf(format string, args ...any) { + if s.verbose { + log.Printf("[sink-registry] "+format, args...) + } +} + +func writeErr(w http.ResponseWriter, status int, code, msg string) { + w.Header().Set("Content-Type", "application/json") + w.WriteHeader(status) + fmt.Fprintf(w, `{"errors":[{"code":%q,"message":%q}]}`, code, msg) +} + +func (s *sinkRegistry) ServeHTTP(w http.ResponseWriter, r *http.Request) { + w.Header().Set("Docker-Distribution-Api-Version", "registry/2.0") + + if r.URL.Path == "/v2" || r.URL.Path == "/v2/" { + w.WriteHeader(http.StatusOK) + return + } + + parts := strings.Split(strings.Trim(r.URL.Path, "/"), "/") + if len(parts) < 3 || parts[0] != "v2" { + writeErr(w, http.StatusNotFound, "NAME_UNKNOWN", "unsupported path") + return + } + + // Find the "blobs" or "manifests" marker. + marker, idx := "", -1 + for i := len(parts) - 1; i >= 1; i-- { + if parts[i] == "blobs" || parts[i] == "manifests" { + marker, idx = parts[i], i + break + } + } + if idx == -1 { + writeErr(w, http.StatusNotFound, "NAME_UNKNOWN", "expected blobs or manifests in path") + return + } + repo := strings.Join(parts[1:idx], "/") + rest := parts[idx+1:] + + switch marker { + case "blobs": + s.handleBlobs(w, r, repo, rest) + case "manifests": + if len(rest) != 1 { + writeErr(w, http.StatusNotFound, "MANIFEST_UNKNOWN", "bad manifest path") + return + } + s.handleManifests(w, r, repo, rest[0]) + } +} + +func (s *sinkRegistry) handleBlobs(w http.ResponseWriter, r *http.Request, repo string, rest []string) { + // Upload sub-resource: /v2//blobs/uploads[/] + if len(rest) >= 1 && rest[0] == "uploads" { + switch r.Method { + case http.MethodPost: + s.startUpload(w, r, repo) + case http.MethodPatch: + if len(rest) < 2 { + writeErr(w, http.StatusBadRequest, "BLOB_UPLOAD_INVALID", "missing upload id") + return + } + s.patchUpload(w, r, repo, rest[1]) + case http.MethodPut: + if len(rest) < 2 { + writeErr(w, http.StatusBadRequest, "BLOB_UPLOAD_INVALID", "missing upload id") + return + } + s.putUpload(w, r, repo, rest[1]) + default: + writeErr(w, http.StatusMethodNotAllowed, "UNSUPPORTED", "method not allowed on uploads") + } + return + } + + // Blob by digest: /v2//blobs/ + if len(rest) != 1 { + writeErr(w, http.StatusNotFound, "BLOB_UNKNOWN", "bad blob path") + return + } + digest := rest[0] + switch r.Method { + case http.MethodHead, http.MethodGet: + s.getBlob(w, r, digest) + default: + writeErr(w, http.StatusMethodNotAllowed, "UNSUPPORTED", "method not allowed on blob") + } +} + +func (s *sinkRegistry) startUpload(w http.ResponseWriter, r *http.Request, repo string) { + // Monolithic upload: POST .../uploads/?digest=sha256:... + if digest := r.URL.Query().Get("digest"); digest != "" { + u := newUpload(s.smallBlobLimit) + s.consume(u, r.Body) + s.finalizeBlob(w, u, digest) + return + } + + id := strconv.FormatInt(rand.Int63(), 16) + strconv.FormatInt(rand.Int63(), 16) + s.mu.Lock() + s.uploads[id] = newUpload(s.smallBlobLimit) + s.mu.Unlock() + s.logf("POST start upload repo=%s id=%s", repo, id) + + loc := fmt.Sprintf("/v2/%s/blobs/uploads/%s", repo, id) + w.Header().Set("Location", loc) + w.Header().Set("Range", "0-0") + w.Header().Set("Docker-Upload-Uuid", id) + w.WriteHeader(http.StatusAccepted) +} + +func (s *sinkRegistry) patchUpload(w http.ResponseWriter, r *http.Request, repo, id string) { + s.mu.Lock() + u := s.uploads[id] + s.mu.Unlock() + if u == nil { + writeErr(w, http.StatusNotFound, "BLOB_UPLOAD_UNKNOWN", "no such upload") + return + } + + // Stream the body straight through the hasher; never hold the whole blob. + n := s.consume(u, r.Body) + s.logf("PATCH upload id=%s chunk=%d total=%d", id, n, u.n) + + loc := fmt.Sprintf("/v2/%s/blobs/uploads/%s", repo, id) + w.Header().Set("Location", loc) + w.Header().Set("Range", fmt.Sprintf("0-%d", maxInt64(u.n-1, 0))) + w.Header().Set("Docker-Upload-Uuid", id) + w.WriteHeader(http.StatusAccepted) +} + +func (s *sinkRegistry) putUpload(w http.ResponseWriter, r *http.Request, repo, id string) { + digest := r.URL.Query().Get("digest") + if digest == "" { + writeErr(w, http.StatusBadRequest, "DIGEST_INVALID", "missing digest on commit") + return + } + + s.mu.Lock() + u := s.uploads[id] + delete(s.uploads, id) + s.mu.Unlock() + if u == nil { + writeErr(w, http.StatusNotFound, "BLOB_UPLOAD_UNKNOWN", "no such upload") + return + } + + // Any trailing bytes on the PUT are also streamed through. + s.consume(u, r.Body) + s.logf("PUT commit id=%s digest=%s size=%d", id, digest, u.n) + s.finalizeBlob(w, u, digest) +} + +// consume streams rc through the upload (digesting + capped retention) using a +// fixed-size staging buffer. Memory stays constant no matter how big rc is. +func (s *sinkRegistry) consume(u *upload, rc io.Reader) int64 { + before := u.n + // io.Copy uses a 32KiB staging buffer; u.Write caps what it retains. + _, _ = io.Copy(u, rc) + return u.n - before +} + +func (s *sinkRegistry) finalizeBlob(w http.ResponseWriter, u *upload, expected string) { + got := "sha256:" + hex.EncodeToString(u.h.Sum(nil)) + if expected != got { + writeErr(w, http.StatusBadRequest, "DIGEST_INVALID", + fmt.Sprintf("digest mismatch: client said %s, computed %s", expected, got)) + return + } + + s.mu.Lock() + s.blobSizes[got] = u.n + if content, ok := u.retained(); ok { + s.smallBlobs[got] = append([]byte(nil), content...) + } + s.mu.Unlock() + + w.Header().Set("Docker-Content-Digest", got) + w.WriteHeader(http.StatusCreated) +} + +func (s *sinkRegistry) getBlob(w http.ResponseWriter, r *http.Request, digest string) { + s.mu.Lock() + size, known := s.blobSizes[digest] + content, haveContent := s.smallBlobs[digest] + s.mu.Unlock() + + if !known { + writeErr(w, http.StatusNotFound, "BLOB_UNKNOWN", "blob not found") + return + } + + w.Header().Set("Docker-Content-Digest", digest) + w.Header().Set("Content-Length", strconv.FormatInt(size, 10)) + if r.Method == http.MethodHead { + w.WriteHeader(http.StatusOK) + return + } + if !haveContent { + // We digested-and-dropped this (large) blob; we can't replay it. + writeErr(w, http.StatusNotFound, "BLOB_UNKNOWN", + "blob content was discarded by the streaming sink (too large to retain)") + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write(content) +} + +func (s *sinkRegistry) handleManifests(w http.ResponseWriter, r *http.Request, repo, ref string) { + key := repo + "/" + ref + switch r.Method { + case http.MethodPut: + body, err := io.ReadAll(r.Body) + if err != nil { + writeErr(w, http.StatusBadRequest, "MANIFEST_INVALID", err.Error()) + return + } + sum := sha256.Sum256(body) + digest := "sha256:" + hex.EncodeToString(sum[:]) + ct := r.Header.Get("Content-Type") + + s.mu.Lock() + s.manifests[key] = body + s.manifests[repo+"/"+digest] = body + s.manifestCT[key] = ct + s.manifestCT[repo+"/"+digest] = ct + s.mu.Unlock() + s.logf("PUT manifest repo=%s ref=%s digest=%s size=%d", repo, ref, digest, len(body)) + + w.Header().Set("Docker-Content-Digest", digest) + w.WriteHeader(http.StatusCreated) + + case http.MethodHead, http.MethodGet: + s.mu.Lock() + body, ok := s.manifests[key] + ct := s.manifestCT[key] + s.mu.Unlock() + if !ok { + writeErr(w, http.StatusNotFound, "MANIFEST_UNKNOWN", "manifest not found") + return + } + sum := sha256.Sum256(body) + if ct == "" { + ct = "application/vnd.docker.distribution.manifest.v2+json" + } + w.Header().Set("Content-Type", ct) + w.Header().Set("Docker-Content-Digest", "sha256:"+hex.EncodeToString(sum[:])) + w.Header().Set("Content-Length", strconv.Itoa(len(body))) + if r.Method == http.MethodHead { + w.WriteHeader(http.StatusOK) + return + } + w.WriteHeader(http.StatusOK) + _, _ = w.Write(body) + + default: + writeErr(w, http.StatusMethodNotAllowed, "UNSUPPORTED", "method not allowed on manifest") + } +} + +func maxInt64(a, b int64) int64 { + if a > b { + return a + } + return b +}