From 06947a4d5df9a578c4448ba44c3983df13b32d86 Mon Sep 17 00:00:00 2001 From: Dave Mihalcik Date: Mon, 31 Aug 2026 21:33:23 -0400 Subject: [PATCH] fix(cli): stream encrypt instead of buffering the whole payload `otdfctl encrypt` read the entire plaintext into memory, handed the slice to EncryptBytes, which accumulated the entire ciphertext in a bytes.Buffer, and only then copied it to the destination. Peak RSS was therefore roughly twice the payload, and a 10 GB cap existed solely to keep that from taking the process down. The payload now streams from the input to the destination. The SDK already takes an io.ReadSeeker and writes incrementally, so this is a matter of not getting in its way: the file argument is opened directly and passed through, and piped stdin is spooled to a temporary file because the SDK seeks to the end of the payload to size it. Spooling trades disk for the memory a whole-payload read cost; once CreateTDF accepts a plain io.Reader the spool goes away. Handler.EncryptBytes becomes Handler.Encrypt(ctx, out, in, EncryptOptions). The options struct replaces eight positional parameters, and passing the command's context through means an interrupted encrypt is now cancellable. MIME detection no longer needs the payload in memory either. detectMimeType reads the first megabyte -- which is all mimetype inspects, given SetLimit -- and rewinds, so the whole payload still reaches the encoder. Two behaviors worth calling out: Fixes a panic. The extension fallback called mimetype.Lookup(fileExt), but Lookup takes a MIME type string, not an extension: it parses its argument with mime.ParseMediaType, which fails on a bare extension, so it returned nil for every extension and .String() on that nil dereferenced. Any file whose contents mimetype could not classify and whose name had an extension crashed the CLI. The stdlib mime.TypeByExtension is the extension lookup, and an unknown extension now leaves application/octet-stream in place. Output to a file is atomic. Encryption writes to a temporary sibling of the destination and renames it into place only on success, so a failed run no longer leaves a truncated .tdf where a complete one is expected. Output to stdout cannot offer this; a failure part-way through has already written bytes downstream. Because cli.ExitWithError calls os.Exit and skips deferred functions, every exit path discards the temporary output and the spool explicitly. Signed-off-by: Dave Mihalcik --- otdfctl/cmd/tdf/encrypt.go | 150 +++++++++++++++++++++----------- otdfctl/cmd/tdf/encrypt_test.go | 67 ++++++++++++++ otdfctl/pkg/handlers/tdf.go | 62 +++++++------ 3 files changed, 202 insertions(+), 77 deletions(-) create mode 100644 otdfctl/cmd/tdf/encrypt_test.go diff --git a/otdfctl/cmd/tdf/encrypt.go b/otdfctl/cmd/tdf/encrypt.go index 3935825f28..9358338a76 100644 --- a/otdfctl/cmd/tdf/encrypt.go +++ b/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) + } } } diff --git a/otdfctl/cmd/tdf/encrypt_test.go b/otdfctl/cmd/tdf/encrypt_test.go new file mode 100644 index 0000000000..8c22bea7ed --- /dev/null +++ b/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) +} diff --git a/otdfctl/pkg/handlers/tdf.go b/otdfctl/pkg/handlers/tdf.go index 73900bd79a..f4bfb25e27 100644 --- a/otdfctl/pkg/handlers/tdf.go +++ b/otdfctl/pkg/handlers/tdf.go @@ -38,51 +38,57 @@ 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 } @@ -90,14 +96,14 @@ func (h Handler) EncryptBytes( 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") } }