feat(artwork): add teeReader to capture served artwork bytes

A wrapping io.ReadCloser that mirrors read bytes into a bounded buffer and, on a
fully-consumed Close, hands the captured bytes to a callback. Partial reads and
oversized streams are skipped so the callback only ever receives a complete,
bounded image — the exact bytes the client received. This is the capture side of
the served-bytes blurhash tee.
This commit is contained in:
Deluan 2026-07-17 19:36:31 -04:00
parent eb5ecabc5a
commit 2169938b30
2 changed files with 91 additions and 0 deletions

View File

@ -0,0 +1,48 @@
package artwork
import (
"bytes"
"io"
)
// teeReader mirrors bytes read from src into buf, and on Close invokes onComplete with the captured
// bytes only if the stream was fully consumed (EOF) and stayed within maxBytes. Partial reads and
// oversized streams are skipped, so a hash is only ever computed from a complete, bounded image.
type teeReader struct {
src io.ReadCloser
buf bytes.Buffer
maxBytes int
onComplete func(data []byte)
eof bool
over bool
done bool
}
func newTeeReader(src io.ReadCloser, maxBytes int, onComplete func(data []byte)) *teeReader {
return &teeReader{src: src, maxBytes: maxBytes, onComplete: onComplete}
}
func (t *teeReader) Read(p []byte) (int, error) {
n, err := t.src.Read(p)
if n > 0 && !t.over {
if t.buf.Len()+n > t.maxBytes {
t.over = true
t.buf.Reset()
} else {
t.buf.Write(p[:n])
}
}
if err == io.EOF {
t.eof = true
}
return n, err
}
func (t *teeReader) Close() error {
err := t.src.Close()
if !t.done && t.eof && !t.over && t.onComplete != nil {
t.done = true
t.onComplete(t.buf.Bytes())
}
return err
}

View File

@ -0,0 +1,43 @@
package artwork
import (
"bytes"
"io"
. "github.com/onsi/ginkgo/v2"
. "github.com/onsi/gomega"
)
var _ = Describe("teeReader", func() {
It("calls onComplete with the full bytes after a complete read+close", func() {
var got []byte
src := io.NopCloser(bytes.NewReader([]byte("hello world")))
tr := newTeeReader(src, 1024, func(data []byte) { got = data })
out, err := io.ReadAll(tr)
Expect(err).ToNot(HaveOccurred())
Expect(string(out)).To(Equal("hello world"))
Expect(tr.Close()).To(Succeed())
Expect(string(got)).To(Equal("hello world"))
})
It("does not call onComplete when the stream is not fully read", func() {
called := false
src := io.NopCloser(bytes.NewReader([]byte("hello world")))
tr := newTeeReader(src, 1024, func(data []byte) { called = true })
buf := make([]byte, 3)
_, err := tr.Read(buf) // partial read, then close without EOF
Expect(err).ToNot(HaveOccurred())
Expect(tr.Close()).To(Succeed())
Expect(called).To(BeFalse())
})
It("does not call onComplete when the data exceeds maxBytes", func() {
called := false
src := io.NopCloser(bytes.NewReader([]byte("hello world")))
tr := newTeeReader(src, 4, func(data []byte) { called = true })
_, err := io.ReadAll(tr)
Expect(err).ToNot(HaveOccurred())
Expect(tr.Close()).To(Succeed())
Expect(called).To(BeFalse())
})
})