diff --git a/core/artwork/blurhash_tee.go b/core/artwork/blurhash_tee.go new file mode 100644 index 000000000..f603f1f65 --- /dev/null +++ b/core/artwork/blurhash_tee.go @@ -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 +} diff --git a/core/artwork/blurhash_tee_internal_test.go b/core/artwork/blurhash_tee_internal_test.go new file mode 100644 index 000000000..595bbb099 --- /dev/null +++ b/core/artwork/blurhash_tee_internal_test.go @@ -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()) + }) +})