feat(sdk): accept io.Reader in CreateTDF and drop the 64 GB payload cap

sdk/tdf.go

@@ -11,6 +11,7 @@ import (
 	"hash/crc32"
 	"io"
 	"log/slog"
+	"math"
 	"net/http"
 	"strconv"
 	"strings"
@@ -31,7 +32,6 @@ import (
 
 const (
 	keyAccessSchemaVersion = "1.0"
-	maxFileSizeSupported   = 68719476736 // 64gb
 	defaultMimeType        = "application/octet-stream"
 	zip64MagicVal          = int64(^uint32(0))
 	tdfAsZip               = "zip"
@@ -53,6 +53,18 @@ const (
 	kAssertionHash         = "assertionHash"
 	hexSemverThreshold     = "4.3.0"
 	readActionName         = "read"
+
+	// maxPayloadSegments caps the segment count a declared input size may imply.
+	// The archive writer counts segments with an int, so the count has to fit one
+	// on every platform the SDK builds for. A payload that needs more segments than
+	// this — 4 PiB at the default segment size — is beyond what CreateTDF can write
+	// anyway, and is better refused than silently mis-sized.
+	//
+	// This bounds only a payload whose length resolves, declared through
+	// [WithInputSize] or recovered from a seekable reader. A stream that can be
+	// measured neither way is read to EOF with no segment ceiling; DSPX-4905 tracks
+	// giving it one.
+	maxPayloadSegments = math.MaxInt32
 )
 
 // Loads and reads ZTDF files
@@ -138,7 +150,15 @@ func (t TDFObject) Size() int64 {
 	return t.size
 }
 
-func (s SDK) CreateTDF(writer io.Writer, reader io.ReadSeeker, opts ...TDFOption) (*TDFObject, error) {
+// CreateTDF encrypts the payload read from reader and writes a TDF to writer. See
+// [SDK.CreateTDFContext] for how the payload length is resolved.
+//
+// Bytes are consumed from the reader's current position until the resolved length
+// is reached, or through EOF when the length cannot be resolved at all. A seekable
+// reader is not rewound first: a caller that has already advanced it — sniffing a
+// header for content-type detection, say — encrypts only what remains. Earlier
+// releases took an io.ReadSeeker and always rewound to byte 0.
+func (s SDK) CreateTDF(writer io.Writer, reader io.Reader, opts ...TDFOption) (*TDFObject, error) {
 	return s.CreateTDFContext(context.Background(), writer, reader, opts...)
 }
 
@@ -162,22 +182,30 @@ func uuidSplitIDGenerator() string {
 	return uuid.New().String()
 }
 
-// CreateTDFContext reads plain text from the given reader and saves it to the writer, subject to the given options
-func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.ReadSeeker, opts ...TDFOption) (*TDFObject, error) { //nolint:funlen, gocognit, lll // Better readability keeping it as is
-	inputSize, err := reader.Seek(0, io.SeekEnd)
-	if err != nil {
-		return nil, fmt.Errorf("readSeeker.Seek failed: %w", err)
-	}
-
-	if inputSize > maxFileSizeSupported {
-		return nil, errFileTooLarge
-	}
-
-	_, err = reader.Seek(0, io.SeekStart)
-	if err != nil {
-		return nil, fmt.Errorf("readSeeker.Seek failed: %w", err)
-	}
-
+// CreateTDFContext reads plain text from the given reader and saves it to the writer,
+// subject to the given options. Bytes are consumed from the reader's current position.
+//
+// A seekable reader is not rewound to byte 0 first — see [SDK.CreateTDF], which spells
+// out how that differs from the io.ReadSeeker signature this replaced.
+//
+// Knowing the length up front lets the archive stay in the compact ZIP32 layout when it
+// fits. The length comes from [WithInputSize] if given, otherwise from the reader when
+// it is seekable; a payload that can be measured neither way is written as ZIP64.
+//
+// A resolved length is exact, and it bounds the read no matter which of those two ways
+// produced it. Reading stops after that many bytes rather than at EOF, and a reader
+// that ends early fails the call. So a file that grows between the seek that measured
+// it and the read that follows contributes only its measured length — the later bytes
+// are left unread rather than silently appended — and one that is truncated in that
+// window fails instead of yielding a short TDF. Only a payload whose length resolves
+// neither way is read through to EOF.
+//
+// Memory grows with the segment count rather than the payload: one manifest segment and
+// one archive-writer entry apiece, all of them live until the archive is finalized. At
+// the 2 MiB default that is roughly 500k entries per terabyte, which is the practical
+// ceiling on a large stream now that the old 64 GB cap is gone. [WithSegmentSize] trades
+// that count against per-segment overhead.
+func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.Reader, opts ...TDFOption) (*TDFObject, error) { //nolint:funlen, gocognit, lll // Better readability keeping it as is
 	tdfConfig, err := newTDFConfig(opts...)
 	if err != nil {
 		return nil, fmt.Errorf("NewTDFConfig failed: %w", err)
@@ -200,60 +228,81 @@ func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.R
 	} else if segmentSize < minSegmentSize {
 		return nil, fmt.Errorf("segment size too small: %d", segmentSize)
 	}
-	totalSegments := inputSize / segmentSize
-	if inputSize%segmentSize != 0 {
-		totalSegments++
-	}
-
-	// empty payload we still want to create a payload
-	if totalSegments == 0 {
-		totalSegments = 1
-	}
 
 	encryptedSegmentSize := segmentSize + gcmIvSize + aesBlockSize
-	payloadSize := inputSize + (totalSegments * (gcmIvSize + aesBlockSize))
 
-	zipMode := zipstream.Zip64Auto
-	if payloadSize >= zip64MagicVal {
-		zipMode = zipstream.Zip64Always
-	}
+	// These describe a payload whose length cannot be known before it is read: no
+	// segment count can be declared up front, the archive has to assume it may need
+	// 64-bit offsets, and every read asks for a whole segment. A measurable payload
+	// replaces them below.
+	var (
+		declaredSize   int64
+		sizeIsDeclared bool
+	)
+	declaredSegments := 0
+	zipMode := zipstream.Zip64Always
+	readBufSize := segmentSize
+
+	if inputSize, err := resolveInputSize(tdfConfig, reader); err != nil {
+		return nil, err
+	} else if inputSize != inputSizeUnknown {
+		segments, err := segmentCount(inputSize, segmentSize)
+		if err != nil {
+			return nil, err
+		}
+		declaredSize, sizeIsDeclared, declaredSegments = inputSize, true, segments
+
+		// The ZIP64 choice is baked into the payload's local file header, which goes
+		// out ahead of the first segment, so it cannot be revisited once the archive
+		// has started. ZIP32 is only safe when the encrypted payload — the plaintext
+		// plus a per-segment IV and tag — stays under the 32-bit ceiling. Subtracting
+		// the overhead rather than adding it keeps the comparison honest for a size
+		// near math.MaxInt64, where the addition would wrap negative and read small.
+		encryptionOverhead := int64(segments) * (gcmIvSize + aesBlockSize)
+		if inputSize < zip64MagicVal-encryptionOverhead {
+			zipMode = zipstream.Zip64Auto
+		}
 
-	expectedSegments := int(totalSegments)
-	if expectedSegments < 1 {
-		expectedSegments = 1
+		// A known length doubles as a read limit: overrunning it would invalidate the
+		// ZIP64 choice made from it. The buffer is only as large as the payload
+		// actually needs, too — the segment size defaults to 2 MiB, so sizing on it
+		// alone would allocate that much to encrypt a handful of bytes. The buffer
+		// never shrinks to zero, so a read that comes back empty always means EOF.
+		reader = io.LimitReader(reader, inputSize)
+		readBufSize = max(1, min(segmentSize, inputSize))
 	}
 
-	archiveWriter := zipstream.NewSegmentTDFWriter(
-		expectedSegments,
-		zipstream.WithZip64Mode(zipMode),
-		zipstream.WithMaxSegments(expectedSegments),
-	)
+	archiveOpts := []zipstream.Option{zipstream.WithZip64Mode(zipMode)}
+	if declaredSegments > 0 {
+		archiveOpts = append(archiveOpts, zipstream.WithMaxSegments(declaredSegments))
+	}
+	archiveWriter := zipstream.NewSegmentTDFWriter(declaredSegments, archiveOpts...)
 
 	outputWriter := &countingWriter{writer: writer}
 
-	var readPos int64
 	var aggregateHashBuilder strings.Builder
-	// 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
-		}
-
+	var bytesRead int64
+	readBuf := make([]byte, readBufSize)
+	for segmentIndex := 0; ; segmentIndex++ {
 		// 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)
+		// erroring, so a bare Read would cut segments short at the whim of the
+		// reader. ReadFull retries until the segment is filled or the input runs
+		// out; a short final segment surfaces as io.ErrUnexpectedEOF.
+		n, readErr := io.ReadFull(reader, readBuf)
+		if readErr != nil && !errors.Is(readErr, io.EOF) && !errors.Is(readErr, io.ErrUnexpectedEOF) {
+			return nil, fmt.Errorf("io.Reader.Read failed: %w", readErr)
+		}
+		// A payload whose length is an exact multiple of the segment size reports
+		// EOF with nothing read. An empty payload still gets one empty segment.
+		if n == 0 && segmentIndex > 0 {
+			break
 		}
+		readSize := int64(n)
+		bytesRead += readSize
 
 		cipherData, err := tdfObject.aesGcm.Encrypt(readBuf[:readSize])
 		if err != nil {
-			return nil, fmt.Errorf("io.ReadSeeker.Read failed: %w", err)
+			return nil, fmt.Errorf("ocrypto.AesGcm.Encrypt failed: %w", err)
 		}
 
 		crc := crc32.ChecksumIEEE(cipherData)
@@ -289,9 +338,17 @@ func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.R
 
 		tdfObject.manifest.Segments = append(tdfObject.manifest.Segments, segmentInfo)
 
-		totalSegments--
-		readPos += readSize
-		segmentIndex++
+		if readErr != nil {
+			break
+		}
+	}
+
+	// A reader that ran dry before the declared length would otherwise yield a TDF
+	// that is silently short of the payload the caller asked to encrypt. The archive
+	// is already sized and partly written by now, so there is nothing to salvage —
+	// fail rather than hand back a truncated result that looks complete.
+	if sizeIsDeclared && bytesRead != declaredSize {
+		return nil, fmt.Errorf("%w: read %d of %d bytes", errInputShorterThanDeclared, bytesRead, declaredSize)
 	}
 
 	rootSignature, err := rootIntegrity([]byte(aggregateHashBuilder.String()), tdfObject.payloadKey[:],
@@ -369,6 +426,61 @@ func (s SDK) CreateTDFContext(ctx context.Context, writer io.Writer, reader io.R
 	return tdfObject, nil
 }
 
+// resolveInputSize reports the payload length in bytes, or inputSizeUnknown when it
+// cannot be established without consuming the reader. An explicit WithInputSize wins
+// over what the reader can report about itself.
+//
+// A reader may satisfy io.Seeker and still refuse to seek — os.Stdin on the end of a
+// pipe is the common case — so a failed probe is treated as an unmeasurable payload
+// rather than an error. Failing to restore the original position is different: the
+// cursor has already moved and the payload can no longer be read in full.
+func resolveInputSize(tdfConfig *TDFConfig, reader io.Reader) (int64, error) {
+	if tdfConfig.inputSize != inputSizeUnknown {
+		return tdfConfig.inputSize, nil
+	}
+
+	seeker, ok := reader.(io.Seeker)
+	if !ok {
+		return inputSizeUnknown, nil
+	}
+	start, err := seeker.Seek(0, io.SeekCurrent)
+	if err != nil {
+		return inputSizeUnknown, nil //nolint:nilerr // a reader that cannot seek is measured by reading it
+	}
+	end, err := seeker.Seek(0, io.SeekEnd)
+	if err != nil {
+		return inputSizeUnknown, nil //nolint:nilerr // a reader that cannot seek is measured by reading it
+	}
+	if _, err := seeker.Seek(start, io.SeekStart); err != nil {
+		return 0, fmt.Errorf("seeker.Seek failed to restore reader position: %w", err)
+	}
+	// Handle readers positioned at or past the end of the file: an empty payload, not a negative one.
+	return max(0, end-start), nil
+}
+
+// segmentCount returns the number of segments a payload of inputSize bytes occupies,
+// and refuses a size that would need more segments than maxPayloadSegments.
+//
+// The ceiling division is a quotient plus a remainder test rather than the usual
+// (inputSize + segmentSize - 1) / segmentSize: WithInputSize accepts any non-negative
+// int64, and that addition wraps negative for a size near math.MaxInt64 — which would
+// hand the archive writer a negative count for it to silently read as one segment.
+func segmentCount(inputSize, segmentSize int64) (int, error) {
+	if inputSize == 0 {
+		// An empty payload still gets one empty segment.
+		return 1, nil
+	}
+	count := inputSize / segmentSize
+	if inputSize%segmentSize != 0 {
+		count++
+	}
+	if count > maxPayloadSegments {
+		return 0, fmt.Errorf("%w: %d bytes in %d-byte segments needs %d segments, limit is %d",
+			errTooManySegments, inputSize, segmentSize, count, maxPayloadSegments)
+	}
+	return int(count), nil
+}
+
 // initKAOTemplate initializes the KAO template, from either the split plan, kaoTemplate, or autoconfigure based on tags.
 func (tdfConfig *TDFConfig) initKAOTemplate(ctx context.Context, s SDK) error {
 	// At most one of the following should be true:

sdk/tdf_config.go

@@ -21,6 +21,10 @@ const (
 	ECKeySize256           = 256
 	ECKeySize384           = 384
 	ECKeySize521           = 521
+
+	// inputSizeUnknown marks a payload whose length cannot be established
+	// before it is read.
+	inputSizeUnknown = -1
 )
 
 type TDFFormat = int
@@ -133,6 +137,7 @@ type TDFOption func(*TDFConfig) error
 type TDFConfig struct {
 	autoconfigure              bool
 	defaultSegmentSize         int64
+	inputSize                  int64
 	enableEncryption           bool
 	tdfFormat                  TDFFormat
 	metaData                   string
@@ -155,6 +160,7 @@ func newTDFConfig(opt ...TDFOption) (*TDFConfig, error) {
 	c := &TDFConfig{
 		autoconfigure:       true,
 		defaultSegmentSize:  defaultSegmentSize,
+		inputSize:           inputSizeUnknown,
 		enableEncryption:    true,
 		tdfFormat:           JSONFormat,
 		rootIntegrityAlg:    RootHS256,
@@ -264,6 +270,25 @@ func WithSegmentSize(size int64) TDFOption {
 	}
 }
 
+// WithInputSize declares the payload length, in bytes, for the reader passed to
+// CreateTDF. Supply it when the reader cannot report its own length — a pipe, a
+// network stream — but the length is known anyway: it lets the archive keep the
+// compact ZIP32 layout that an unmeasurable payload has to give up. When the reader
+// is seekable the length is recovered automatically and this option is unnecessary.
+//
+// The declared length is exact, not an upper bound: reading stops after size bytes
+// even if the reader has more to give, and a reader that reaches EOF first fails the
+// call rather than producing a TDF that is silently short of the payload.
+func WithInputSize(size int64) TDFOption {
+	return func(c *TDFConfig) error {
+		if size < 0 {
+			return fmt.Errorf("WithInputSize: size must not be negative, got %d", size)
+		}
+		c.inputSize = size
+		return nil
+	}
+}
+
 // WithDefaultAssertion returns an Option that adds a default assertion to the TDF.
 func WithSystemMetadataAssertion() TDFOption {
 	return func(c *TDFConfig) error {