From ad7385ff8192c3643828374a56f30042ec8a4412 Mon Sep 17 00:00:00 2001
From: Dave Mihalcik <dmihalcik@virtru.com>
Date: Mon, 31 Aug 2026 17:18:51 -0400
Subject: [PATCH] fix(sdk): fill each segment with io.ReadFull and size the
 buffer to the input

Two problems in CreateTDFContext's encrypt loop.

io.Reader.Read is permitted to return fewer bytes than the caller asked
for without erroring, and the loop treated that as fatal:
"io.ReadSeeker.Read size mismatch". A *bytes.Reader or *os.File on a
local disk rarely returns short, which is why this has held up, but any
wrapping ReadSeeker -- a decompressor, a network-backed store, an
instrumented reader -- can trigger it and there is nothing wrong with
the input when it does. io.ReadFull retries until the segment is full,
so the manual size check goes away with it. The new
Test_TDFCreateShortReads fails on main with exactly that error message.

The read buffer was also sized on defaultSegmentSize alone, which is
2 MiB, so encrypting a twelve-byte payload allocated 2 MiB to hold it.
Size it to min(segmentSize, inputSize) instead; the max(inputSize, 1)
keeps the empty-payload case, which still emits one empty segment, from
asking for a zero-length buffer.

Peeled out of the DSPX-2604 stack.

Signed-off-by: David Mihalcik <dmihalcik@virtru.com>
---
 sdk/tdf.go      | 17 +++++++++--------
 sdk/tdf_test.go | 40 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 49 insertions(+), 8 deletions(-)

diff --git a/sdk/tdf.go b/sdk/tdf.go
index ce42099603..1ff57f56a4 100644
--- a/sdk/tdf.go
+++ b/sdk/tdf.go
@@ -233,7 +233,10 @@ func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.R
 
 	var readPos int64
 	var aggregateHashBuilder strings.Builder
-	readBuf := bytes.NewBuffer(make([]byte, 0, tdfConfig.defaultSegmentSize))
+	// Only as large as the payload actually needs: the segment size defaults to
+	// 2 MiB, so sizing on it alone would allocate that much to encrypt a
+	// handful of bytes.
+	readBuf := make([]byte, min(segmentSize, max(inputSize, 1)))
 	segmentIndex := 0
 	for totalSegments != 0 { // adjust read size
 		readSize := segmentSize
@@ -241,16 +244,14 @@ func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.R
 			readSize = inputSize - readPos
 		}
 
-		n, err := reader.Read(readBuf.Bytes()[:readSize])
-		if err != nil {
+		// io.Reader.Read is free to return fewer bytes than asked for without
+		// erroring, so a bare Read would reject perfectly valid readers as a
+		// size mismatch. ReadFull retries until the segment is filled.
+		if _, err := io.ReadFull(reader, readBuf[:readSize]); err != nil {
 			return nil, fmt.Errorf("io.ReadSeeker.Read failed: %w", err)
 		}
 
-		if int64(n) != readSize {
-			return nil, errors.New("io.ReadSeeker.Read size mismatch")
-		}
-
-		cipherData, err := tdfObject.aesGcm.Encrypt(readBuf.Bytes()[:readSize])
+		cipherData, err := tdfObject.aesGcm.Encrypt(readBuf[:readSize])
 		if err != nil {
 			return nil, fmt.Errorf("io.ReadSeeker.Read failed: %w", err)
 		}
diff --git a/sdk/tdf_test.go b/sdk/tdf_test.go
index 2124c97003..83203a5e7c 100644
--- a/sdk/tdf_test.go
+++ b/sdk/tdf_test.go
@@ -1600,6 +1600,46 @@ func (s *TDFSuite) Test_TDFReader() { //nolint:gocognit // requires for testing
 	}
 }
 
+// shortReadSeeker hands back at most maxRead bytes per Read, which io.Reader
+// explicitly permits. A *bytes.Reader never does this, so nothing else in the
+// suite covers it.
+type shortReadSeeker struct {
+	io.ReadSeeker
+	maxRead int
+}
+
+func (s *shortReadSeeker) Read(p []byte) (int, error) {
+	if len(p) > s.maxRead {
+		p = p[:s.maxRead]
+	}
+	return s.ReadSeeker.Read(p)
+}
+
+// Test_TDFCreateShortReads pins that CreateTDF fills each segment rather than
+// treating a short read as a fatal size mismatch.
+func (s *TDFSuite) Test_TDFCreateShortReads() {
+	kasInfoList := []KASInfo{
+		{URL: s.kasTestURLLookup["http://localhost:65432/"]},
+	}
+
+	tdfBuf := bytes.Buffer{}
+	_, err := s.sdk.CreateTDF(
+		io.Writer(&tdfBuf),
+		&shortReadSeeker{ReadSeeker: bytes.NewReader([]byte(payload)), maxRead: 3},
+		WithKasInformation(kasInfoList...),
+		WithSegmentSize(7),
+	)
+	s.Require().NoError(err)
+
+	r, err := s.sdk.LoadTDF(bytes.NewReader(tdfBuf.Bytes()))
+	s.Require().NoError(err)
+
+	var out bytes.Buffer
+	_, err = r.WriteTo(&out)
+	s.Require().NoError(err)
+	s.Equal(payload, out.String())
+}
+
 func (s *TDFSuite) Test_TDFReaderFail() {
 	kasInfoList := []KASInfo{
 		{