fix(sdk): fill each segment with io.ReadFull and size the buffer to the input

sdk/tdf.go

@@ -233,24 +233,25 @@ 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
 		if (inputSize - readPos) < segmentSize {
 			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)
 		}

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{
 		{