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") } }