internal/pkg: expose response body

This uses the new measured reader provided by Cache. This should make httpArtifact zero-copy.

Signed-off-by: Ophestra <cat@gensokyo.uk>

1 files changed, 99 insertions, 3 deletions

diff --git a/internal/pkg/pkg.go b/internal/pkg/pkg.go
index d029eff9..b3b08510 100644
--- a/internal/pkg/pkg.go
+++ b/internal/pkg/pkg.go
@@ -215,6 +215,20 @@ func (f *FContext) GetArtifact(a Artifact) (
 	panic(InvalidLookupError(f.cache.Ident(a).Value()))
 }
 
+// RContext is passed to [FileArtifact.Cure] and provides helper methods useful
+// for curing the [FileArtifact].
+//
+// Methods of RContext are safe for concurrent use. RContext is valid
+// until [FileArtifact.Cure] returns.
+type RContext struct {
+	// Address of underlying [Cache], should be zeroed or made unusable after
+	// [FileArtifact.Cure] returns and must not be exposed directly.
+	cache *Cache
+}
+
+// Unwrap returns the underlying [context.Context].
+func (r *RContext) Unwrap() context.Context { return r.cache.ctx }
+
 // An Artifact is a read-only reference to a piece of data that may be created
 // deterministically but might not currently be available in memory or on the
 // filesystem.
@@ -323,7 +337,7 @@ type FileArtifact interface {
 	// Callers are responsible for closing the resulting [io.ReadCloser].
 	//
 	// Result must remain identical across multiple invocations.
-	Cure(ctx context.Context) (io.ReadCloser, error)
+	Cure(r *RContext) (io.ReadCloser, error)
 
 	Artifact
 }
@@ -952,7 +966,7 @@ func (c *Cache) openFile(f FileArtifact) (r io.ReadCloser, err error) {
 				}
 			}()
 		}
-		return f.Cure(c.ctx)
+		return f.Cure(&RContext{c})
 	}
 	return
 }
@@ -1229,6 +1243,88 @@ func (c *Cache) getWriter(w io.Writer) *bufio.Writer {
 	return bw
 }
 
+// measuredReader implements [io.ReadCloser] and measures the checksum during
+// Close. If the underlying reader is not read to EOF, Close blocks until all
+// remaining data is consumed and validated.
+type measuredReader struct {
+	// Underlying reader. Never exposed directly.
+	r io.ReadCloser
+	// For validating checksum. Never exposed directly.
+	h hash.Hash
+	// Buffers writes to h, initialised by [Cache]. Never exposed directly.
+	hbw *bufio.Writer
+	// Expected checksum, compared during Close.
+	want unique.Handle[Checksum]
+
+	// For accessing free lists.
+	c *Cache
+
+	// Set up via [io.TeeReader] by [Cache].
+	io.Reader
+}
+
+// Close reads the underlying [io.ReadCloser] to EOF, closes it and measures its
+// outcome. It returns a [ChecksumMismatchError] for an unexpected checksum.
+func (mr *measuredReader) Close() (err error) {
+	if mr.hbw == nil || mr.Reader == nil {
+		return os.ErrInvalid
+	}
+	err = mr.hbw.Flush()
+	mr.c.putWriter(mr.hbw)
+	mr.hbw, mr.Reader = nil, nil
+	if err != nil {
+		_ = mr.r.Close()
+		return
+	}
+	var n int64
+	if n, err = io.Copy(mr.h, mr.r); err != nil {
+		_ = mr.r.Close()
+		return
+	}
+
+	if n > 0 {
+		mr.c.msg.Verbosef("missed %d bytes on measured reader", n)
+	}
+
+	if err = mr.r.Close(); err != nil {
+		return
+	}
+
+	buf := mr.c.getIdentBuf()
+	mr.h.Sum(buf[:0])
+
+	if got := Checksum(buf[:]); got != mr.want.Value() {
+		err = &ChecksumMismatchError{
+			Got:  got,
+			Want: mr.want.Value(),
+		}
+	}
+
+	mr.c.putIdentBuf(buf)
+	return
+}
+
+// newMeasuredReader implements [RContext.NewMeasuredReader].
+func (c *Cache) newMeasuredReader(
+	r io.ReadCloser,
+	checksum unique.Handle[Checksum],
+) io.ReadCloser {
+	mr := measuredReader{r: r, h: sha512.New384(), want: checksum, c: c}
+	mr.hbw = c.getWriter(mr.h)
+	mr.Reader = io.TeeReader(r, mr.hbw)
+	return &mr
+}
+
+// NewMeasuredReader returns an [io.ReadCloser] implementing behaviour required
+// by [FileArtifact]. The resulting [io.ReadCloser] holds a buffer originating
+// from [Cache] and must be closed to return this buffer.
+func (r *RContext) NewMeasuredReader(
+	rc io.ReadCloser,
+	checksum unique.Handle[Checksum],
+) io.ReadCloser {
+	return r.cache.newMeasuredReader(rc, checksum)
+}
+
 // putWriter adds bw to bufioPool.
 func (c *Cache) putWriter(bw *bufio.Writer) { c.bufioPool.Put(bw) }
 
@@ -1363,7 +1459,7 @@ func (c *Cache) cure(a Artifact, curesExempt bool) (
 		if err = c.enterCure(curesExempt); err != nil {
 			return
 		}
-		r, err = f.Cure(c.ctx)
+		r, err = f.Cure(&RContext{c})
 		if err == nil {
 			if checksumPathname == nil || c.IsStrict() {
 				h := sha512.New384()