fix(cli): stream encrypt instead of buffering the whole payload

otdfctl/cmd/tdf/encrypt.go

@@ -1,8 +1,10 @@
 package tdf
 
 import (
+	"errors"
 	"io"
 	"log/slog"
+	"mime"
 	"os"
 	"path/filepath"
 	"strings"
@@ -11,8 +13,9 @@ import (
 	"github.com/opentdf/platform/lib/ocrypto"
 	"github.com/opentdf/platform/otdfctl/cmd/common"
 	"github.com/opentdf/platform/otdfctl/pkg/cli"
+	"github.com/opentdf/platform/otdfctl/pkg/handlers"
 	"github.com/opentdf/platform/otdfctl/pkg/man"
-	"github.com/opentdf/platform/otdfctl/pkg/utils"
+	"github.com/opentdf/platform/otdfctl/pkg/streamio"
 	"github.com/spf13/cobra"
 )
 
@@ -24,6 +27,39 @@ var (
 	EncryptCmd = &encryptDoc.Command
 )
 
+// detectMimeType sniffs the payload's type from its head and rewinds, so the
+// whole payload still reaches the encoder.
+//
+// Detection needs only the first megabyte, which is what mimetype is limited to
+// anyway, so this reads a bounded prefix rather than the whole payload.
+func detectMimeType(in io.ReadSeeker, fileExt string) (string, error) {
+	mimetype.SetLimit(Size1MB) // limit to 1MB
+
+	head := make([]byte, Size1MB)
+	// A payload shorter than the sniff window is the common case, not an error.
+	n, err := io.ReadFull(in, head)
+	if err != nil && !errors.Is(err, io.EOF) && !errors.Is(err, io.ErrUnexpectedEOF) {
+		return "", err
+	}
+	if _, err := in.Seek(0, io.SeekStart); err != nil {
+		return "", err
+	}
+
+	// defaults to application/octet-stream if nothing is recognized
+	detected := mimetype.Detect(head[:n]).String()
+	if detected == "application/octet-stream" && fileExt != "" {
+		// mime.TypeByExtension is the extension lookup. mimetype.Lookup takes a
+		// MIME type string, so passing it a bare extension always returned nil
+		// and dereferencing that panicked — which is what happened for any file
+		// whose contents were unrecognizable and whose name had an extension.
+		// An extension with no known type leaves octet-stream in place.
+		if byExt := mime.TypeByExtension("." + fileExt); byExt != "" {
+			detected = byExt
+		}
+	}
+	return detected, nil
+}
+
 func encryptRun(cmd *cobra.Command, args []string) {
 	c := cli.New(cmd, args, cli.WithPrintJSON())
 	h := common.NewHandler(c)
@@ -57,13 +93,16 @@ func encryptRun(cmd *cobra.Command, args []string) {
 		wrappingKeyAlgorithm = ocrypto.RSA2048Key
 	}
 
-	piped := readPipedStdin()
+	piped, hasPiped, err := streamio.PipeReader(os.Stdin)
+	if err != nil {
+		cli.ExitWithError("failed to scan bytes from stdin", err)
+	}
 
 	inputCount := 0
 	if filePath != "" {
 		inputCount++
 	}
-	if len(piped) > 0 {
+	if hasPiped {
 		inputCount++
 	}
 
@@ -76,71 +115,84 @@ func encryptRun(cmd *cobra.Command, args []string) {
 		cliExit("ONLY ONE")
 	}
 
-	// prefer filepath argument over stdin input
-	bytesSlice := piped
-	var err error
+	// The SDK seeks to the end of the payload to size it, so the input has to be
+	// seekable. A file already is; a pipe is spooled to disk, which trades the
+	// temporary file for the memory a whole-payload read used to cost.
+	var in io.ReadSeeker
+	var cleanup func()
 	if filePath != "" {
-		bytesSlice, err = utils.ReadBytesFromFile(filePath, MaxFileSize)
+		f, err := os.Open(filePath)
 		if err != nil {
 			cli.ExitWithError("Failed to read file:", err)
 		}
-	}
-
-	// auto-detect mime type if not provided
-	if fileMimeType == "" {
-		slog.Debug("detecting mime type of file")
-		// get the mime type of the file
-		mimetype.SetLimit(Size1MB) // limit to 1MB
-		m := mimetype.Detect(bytesSlice)
-		// default to application/octet-stream if no mime type is detected
-		fileMimeType = m.String()
-
-		if fileMimeType == "application/octet-stream" {
-			if fileExt != "" {
-				fileMimeType = mimetype.Lookup(fileExt).String()
-			}
+		in, cleanup = f, func() { f.Close() }
+	} else {
+		f, spoolCleanup, err := streamio.Spool(piped)
+		if err != nil {
+			cli.ExitWithError("Failed to read stdin:", err)
 		}
+		in, cleanup = f, spoolCleanup
 	}
-	slog.Debug("encrypting file",
-		slog.Int("file_len", len(bytesSlice)),
-		slog.String("mime_type", fileMimeType),
-	)
-
-	// Do the encryption
-	encrypted, err := h.EncryptBytes(
-		tdfType,
-		bytesSlice,
-		attrValues,
-		fileMimeType,
-		kasURLPath,
-		assertions,
-		wrappingKeyAlgorithm,
-		targetMode,
-	)
-	if err != nil {
-		cli.ExitWithError("Failed to encrypt", err)
-	}
-
-	// Find the destination as the output flag filename or stdout
-	var dest *os.File
+	// cli.ExitWithError calls os.Exit, which skips deferred functions, so every
+	// exit below goes through fail() to discard the spool and any partial output.
+	defer cleanup()
+
+	// Resolve the destination before encrypting, so the payload streams straight
+	// to it rather than accumulating in memory first.
+	var dest io.Writer
+	var tdfFile *streamio.OutputFile
 	if out != "" {
 		// make sure output ends in .tdf extension
 		if !strings.HasSuffix(out, ".tdf") {
 			out += ".tdf"
 		}
-		tdfFile, err := os.Create(out)
+		tdfFile, err = streamio.NewOutputFile(out)
 		if err != nil {
+			cleanup()
 			cli.ExitWithError("Failed to write encrypted file "+out, err)
 		}
-		defer tdfFile.Close()
+		defer tdfFile.Cleanup()
 		dest = tdfFile
 	} else {
 		dest = os.Stdout
 	}
 
-	_, e := io.Copy(dest, encrypted)
-	if e != nil {
-		cli.ExitWithError("Failed to write encrypted data to stdout", e)
+	fail := func(msg string, err error) {
+		if tdfFile != nil {
+			tdfFile.Cleanup()
+		}
+		cleanup()
+		cli.ExitWithError(msg, err)
+	}
+
+	// auto-detect mime type if not provided
+	if fileMimeType == "" {
+		slog.Debug("detecting mime type of file")
+		fileMimeType, err = detectMimeType(in, fileExt)
+		if err != nil {
+			fail("Failed to read file:", err)
+		}
+	}
+	slog.Debug("encrypting file", slog.String("mime_type", fileMimeType))
+
+	// Do the encryption
+	err = h.Encrypt(c.Context(), dest, in, handlers.EncryptOptions{
+		TDFType:              tdfType,
+		Attributes:           attrValues,
+		MimeType:             fileMimeType,
+		KASURLPath:           kasURLPath,
+		Assertions:           assertions,
+		WrappingKeyAlgorithm: wrappingKeyAlgorithm,
+		TargetMode:           targetMode,
+	})
+	if err != nil {
+		fail("Failed to encrypt", err)
+	}
+
+	if tdfFile != nil {
+		if err := tdfFile.Commit(); err != nil {
+			fail("Failed to write encrypted file "+out, err)
+		}
 	}
 }
 

otdfctl/cmd/tdf/encrypt_test.go

@@ -0,0 +1,67 @@
+package tdf
+
+import (
+	"bytes"
+	"io"
+	"strings"
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestDetectMimeTypeRewindsForTheEncoder(t *testing.T) {
+	for _, tc := range []struct {
+		name    string
+		content string
+		want    string
+	}{
+		{name: "text", content: "hello, world\n", want: "text/plain; charset=utf-8"},
+		{name: "json", content: `{"a":1}`, want: "application/json"},
+		// Larger than the sniff window, so a detector that consumed the reader
+		// instead of rewinding would truncate the payload.
+		{name: "larger than sniff window", content: strings.Repeat("a", Size1MB+7), want: "text/plain; charset=utf-8"},
+	} {
+		t.Run(tc.name, func(t *testing.T) {
+			in := strings.NewReader(tc.content)
+
+			got, err := detectMimeType(in, "")
+			require.NoError(t, err)
+			assert.Equal(t, tc.want, got)
+
+			// The whole payload must still reach the encoder.
+			all, err := io.ReadAll(in)
+			require.NoError(t, err)
+			assert.Len(t, all, len(tc.content))
+		})
+	}
+}
+
+func TestDetectMimeTypeFallsBackToExtension(t *testing.T) {
+	// Bytes mimetype cannot classify, so the extension decides. ".pdf" is in
+	// Go's builtin table, so this does not depend on the host's mime.types.
+	unrecognized := bytes.Repeat([]byte{0x01, 0x02, 0x03, 0x04}, 8)
+
+	got, err := detectMimeType(bytes.NewReader(unrecognized), "pdf")
+	require.NoError(t, err)
+	assert.Equal(t, "application/pdf", got)
+}
+
+func TestDetectMimeTypeUnknownExtensionStaysOctetStream(t *testing.T) {
+	unrecognized := bytes.Repeat([]byte{0x01, 0x02, 0x03, 0x04}, 8)
+
+	// The previous implementation called mimetype.Lookup(fileExt).String().
+	// Lookup takes a MIME type string, not an extension, so it returned nil for
+	// every extension and this path panicked.
+	got, err := detectMimeType(bytes.NewReader(unrecognized), "zzzznotathing")
+	require.NoError(t, err)
+	assert.Equal(t, "application/octet-stream", got)
+}
+
+func TestDetectMimeTypeEmptyPayload(t *testing.T) {
+	in := strings.NewReader("")
+
+	got, err := detectMimeType(in, "")
+	require.NoError(t, err)
+	assert.Equal(t, "text/plain", got)
+}

otdfctl/pkg/handlers/tdf.go

@@ -38,66 +38,72 @@ type TDFInspect struct {
 	UnencryptedMetadata []byte
 }
 
-func (h Handler) EncryptBytes(
-	tdfType string,
-	unencrypted []byte,
-	attrValues []string,
-	mimeType string,
-	kasURLPath string,
-	assertions string,
-	wrappingKeyAlgorithm ocrypto.KeyType,
-	targetMode string,
-) (*bytes.Buffer, error) {
-	var encrypted []byte
-	enc := bytes.NewBuffer(encrypted)
+// EncryptOptions carries the non-stream inputs to Encrypt.
+type EncryptOptions struct {
+	TDFType              string
+	Attributes           []string
+	MimeType             string
+	KASURLPath           string
+	Assertions           string
+	WrappingKeyAlgorithm ocrypto.KeyType
+	TargetMode           string
+}
 
-	switch tdfType {
+// Encrypt streams the plaintext from in to a TDF written to out. Memory use is
+// bounded by the SDK's segment size rather than by the payload length, so the
+// payload may be larger than RAM.
+//
+// in must be seekable: the SDK measures the payload by seeking to its end
+// before encrypting, and knowing the length up front is what lets it avoid
+// defaulting to ZIP64. A caller holding a pipe should spool it first.
+func (h Handler) Encrypt(ctx context.Context, out io.Writer, in io.ReadSeeker, o EncryptOptions) error {
+	switch o.TDFType {
 	// Encrypt the data as a ZTDF
 	case "", tdf.TypeTDF3, tdf.TypeZTDF:
 		opts := []sdk.TDFOption{
-			sdk.WithDataAttributes(attrValues...),
+			sdk.WithDataAttributes(o.Attributes...),
 			sdk.WithKasInformation(sdk.KASInfo{
-				URL: h.platformEndpoint + kasURLPath,
+				URL: h.platformEndpoint + o.KASURLPath,
 			}),
-			sdk.WithMimeType(mimeType),
-			sdk.WithWrappingKeyAlg(wrappingKeyAlgorithm), //nolint:staticcheck // SDK option is deprecated but no replacement is available in this SDK version.
+			sdk.WithMimeType(o.MimeType),
+			sdk.WithWrappingKeyAlg(o.WrappingKeyAlgorithm), //nolint:staticcheck // SDK option is deprecated but no replacement is available in this SDK version.
 		}
 
 		var assertionConfigs []sdk.AssertionConfig
 		//nolint:nestif // nested its mainly for error catching and handling case of string vs file
-		if assertions != "" {
-			err := json.Unmarshal([]byte(assertions), &assertionConfigs)
+		if o.Assertions != "" {
+			err := json.Unmarshal([]byte(o.Assertions), &assertionConfigs)
 			if err != nil {
 				// if unable to marshal to json, interpret as file string and try to read from file
-				assertionBytes, err := utils.ReadBytesFromFile(assertions, MaxAssertionsFileSize)
+				assertionBytes, err := utils.ReadBytesFromFile(o.Assertions, MaxAssertionsFileSize)
 				if err != nil {
-					return nil, fmt.Errorf("unable to read assertions file: %w", err)
+					return fmt.Errorf("unable to read assertions file: %w", err)
 				}
 				err = json.Unmarshal(assertionBytes, &assertionConfigs)
 				if err != nil {
-					return nil, fmt.Errorf("unable to unmarshal assertions json: %w", err)
+					return fmt.Errorf("unable to unmarshal assertions json: %w", err)
 				}
 			}
 			for i, config := range assertionConfigs {
 				if !config.SigningKey.IsEmpty() {
 					correctedKey, err := correctKeyType(config.SigningKey, false)
 					if err != nil {
-						return nil, fmt.Errorf("error with assertion signing key: %w", err)
+						return fmt.Errorf("error with assertion signing key: %w", err)
 					}
 					assertionConfigs[i].SigningKey.Key = correctedKey
 				}
 			}
 			opts = append(opts, sdk.WithAssertions(assertionConfigs...))
 		}
 
-		if targetMode != "" {
-			opts = append(opts, sdk.WithTargetMode(targetMode))
+		if o.TargetMode != "" {
+			opts = append(opts, sdk.WithTargetMode(o.TargetMode))
 		}
 
-		_, err := h.sdk.CreateTDF(enc, bytes.NewReader(unencrypted), opts...)
-		return enc, err
+		_, err := h.sdk.CreateTDFContext(ctx, out, in, opts...)
+		return err
 	default:
-		return nil, errors.New("unknown TDF type")
+		return errors.New("unknown TDF type")
 	}
 }