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