From 580073da5aab33b5314cf7af3014d417fa3c2de9 Mon Sep 17 00:00:00 2001 From: Dave Mihalcik Date: Mon, 31 Aug 2026 21:28:23 -0400 Subject: [PATCH 1/2] chore(cli): move streaming IO helpers into pkg Adds otdfctl/pkg/streamio, holding the input and output plumbing that the streaming encrypt and decrypt work needs, and migrates `inspect` onto it so nothing is left calling the buffered helpers it supersedes. This is groundwork with one user-visible consequence: `inspect` no longer reads the whole TDF into memory. Everything else is a move. Why a new package rather than pkg/cli. The helpers in pkg/cli/pipe.go call ExitWithError -- which calls os.Exit -- from inside the read, so they cannot be used from anywhere that wants to handle the failure itself, and they read the entire input into memory. streamio returns errors and leaves the decision to exit with the command layer. What moved in: - PipeReader establishes whether stdin is a non-empty pipe with a one-byte Peek instead of a read, so the payload still reaches the caller. - Spool copies a pipe to a temporary file and rewinds it. A TDF's manifest sits at the end of the archive, so decrypt and inspect have to seek and cannot consume a pipe directly. - OpenSeekable resolves "file argument or piped stdin" to one seekable handle, reporting ErrNoInput for the shared "nothing to read" case. - OutputFile writes to a temporary sibling of the destination and renames it into place on Commit, so a failed run leaves no partial output. The temp file is a sibling so the rename stays atomic rather than degrading to a cross-filesystem copy. Per review feedback on #3921: - readPipedStdin now delegates its detection to streamio.PipeReader rather than answering "is there piped input?" a second way. Its read is still unbounded; the callers that must stop buffering are changed separately. - pkg/cli/pipe.go is deprecated rather than deleted, since the package is exported and may have callers outside this repository. Worth noting that ReadFromFile has no size cap at all -- not even the 10 GB the tdf commands apply -- which is its own argument for the notice. InspectTDF takes an io.ReadSeeker instead of a byte slice. GetTdfType already rewinds to the start, so the reader is positioned for LoadTDF. Because cli.ExitWithError calls os.Exit and skips deferred functions, inspectRun invokes cleanup explicitly on every exit path, including the successful one: piped input is spooled to disk and the temp file would otherwise survive. Signed-off-by: Dave Mihalcik --- otdfctl/cmd/tdf/inspect.go | 35 ++++- otdfctl/cmd/tdf/tdf.go | 26 ++-- otdfctl/pkg/cli/pipe.go | 47 ++++-- otdfctl/pkg/handlers/tdf.go | 13 +- otdfctl/pkg/streamio/input.go | 114 +++++++++++++++ otdfctl/pkg/streamio/input_test.go | 213 ++++++++++++++++++++++++++++ otdfctl/pkg/streamio/output.go | 86 +++++++++++ otdfctl/pkg/streamio/output_test.go | 148 +++++++++++++++++++ 8 files changed, 650 insertions(+), 32 deletions(-) create mode 100644 otdfctl/pkg/streamio/input.go create mode 100644 otdfctl/pkg/streamio/input_test.go create mode 100644 otdfctl/pkg/streamio/output.go create mode 100644 otdfctl/pkg/streamio/output_test.go diff --git a/otdfctl/cmd/tdf/inspect.go b/otdfctl/cmd/tdf/inspect.go index a35a3e0d04..7b8fd351de 100644 --- a/otdfctl/cmd/tdf/inspect.go +++ b/otdfctl/cmd/tdf/inspect.go @@ -2,11 +2,13 @@ package tdf import ( "errors" + "log/slog" "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/streamio" "github.com/opentdf/platform/sdk" "github.com/spf13/cobra" ) @@ -42,17 +44,36 @@ func inspectRun(cmd *cobra.Command, args []string) { h := common.NewHandler(c) defer h.Close() - data := cli.ReadFromArgsOrPipe(args, nil) - if len(data) == 0 { - c.ExitWithError("must provide ONE of the following: [file argument, stdin input]", errors.New("no input provided")) + var path string + if len(args) > 0 { + path = args[0] } + in, cleanup, err := streamio.OpenSeekable(path) + if err != nil { + if errors.Is(err, streamio.ErrNoInput) { + c.ExitWithError("must provide ONE of the following: [file argument, stdin input]", err) + } + c.ExitWithError("failed to read input", err) + } + // cli.ExitWithError calls os.Exit, which does not run deferred functions, so + // cleanup is also invoked explicitly before every exit below — including the + // successful one, since piped input is spooled to a temporary file. + defer cleanup() - result, errs := h.InspectTDF(data) + result, errs := h.InspectTDF(in) for _, err := range errs { - if errors.Is(err, handlers.ErrTDFInspectFailNotValidTDF) { + switch { + case errors.Is(err, handlers.ErrTDFInspectFailNotValidTDF): + cleanup() c.ExitWithError("not a valid TDF", err) - } else if errors.Is(err, handlers.ErrTDFInspectFailNotInspectable) { + case errors.Is(err, handlers.ErrTDFInspectFailNotInspectable): + cleanup() c.ExitWithError("failed to inspect TDF", err) + default: + // Attribute/metadata reads (e.g. a denied or unreachable KAS rewrap) + // are non-fatal by design so the manifest still prints, but must not + // vanish silently. + slog.Warn("inspect: partial result", slog.Any("error", err)) } } @@ -76,8 +97,10 @@ func inspectRun(cmd *cobra.Command, args []string) { Attributes: result.Attributes, } + cleanup() c.ExitWithJSON(m, cli.ExitCodeSuccess) } + cleanup() c.ExitWithError("failed to inspect TDF", nil) } diff --git a/otdfctl/cmd/tdf/tdf.go b/otdfctl/cmd/tdf/tdf.go index 7f727af92b..9f4395a782 100644 --- a/otdfctl/cmd/tdf/tdf.go +++ b/otdfctl/cmd/tdf/tdf.go @@ -5,6 +5,7 @@ import ( "os" "github.com/opentdf/platform/otdfctl/pkg/cli" + "github.com/opentdf/platform/otdfctl/pkg/streamio" ) const ( @@ -15,17 +16,24 @@ const ( GroupID = TDF ) +// readPipedStdin returns the whole of piped stdin, or nil when stdin is a +// terminal or an empty redirect. +// +// Detection is delegated to streamio.PipeReader so there is a single answer to +// "is there piped input?" across the CLI. The read itself is still unbounded; +// callers that must not hold the payload in memory should use +// streamio.OpenSeekable instead. func readPipedStdin() []byte { - stat, err := os.Stdin.Stat() + r, ok, err := streamio.PipeReader(os.Stdin) if err != nil { - cli.ExitWithError("Failed to read stat from stdin", err) + cli.ExitWithError("failed to scan bytes from stdin", err) } - if (stat.Mode() & os.ModeCharDevice) == 0 { - buf, err := io.ReadAll(os.Stdin) - if err != nil { - cli.ExitWithError("failed to scan bytes from stdin", err) - } - return buf + if !ok { + return nil } - return nil + buf, err := io.ReadAll(r) + if err != nil { + cli.ExitWithError("failed to scan bytes from stdin", err) + } + return buf } diff --git a/otdfctl/pkg/cli/pipe.go b/otdfctl/pkg/cli/pipe.go index 561c8ff19f..8d0472960d 100644 --- a/otdfctl/pkg/cli/pipe.go +++ b/otdfctl/pkg/cli/pipe.go @@ -3,8 +3,22 @@ package cli import ( "io" "os" + + "github.com/opentdf/platform/otdfctl/pkg/streamio" ) +// These wrappers read their whole input into memory and call ExitWithError — +// which calls os.Exit — from inside the read, so they cannot be used anywhere +// that wants to handle the failure itself. They delegate to pkg/streamio so +// there is one implementation of "find the file argument or piped stdin" to +// maintain. +// +// New code should call pkg/streamio directly, which streams and returns +// errors instead of buffering and exiting. + +// Deprecated: reads the entire input into memory and terminates the process on +// failure. Use streamio.OpenSeekable, which resolves the same "file argument or +// piped stdin" choice without buffering and returns an error. func ReadFromArgsOrPipe(args []string, pipe *os.File) []byte { if len(args) > 0 { return ReadFromFile(args[0]) @@ -15,31 +29,38 @@ func ReadFromArgsOrPipe(args []string, pipe *os.File) []byte { return ReadFromPipe(pipe) } +// Deprecated: reads the entire pipe into memory and terminates the process on +// failure. Use streamio.PipeReader, which reports whether input is present +// without consuming it, paired with io.ReadAll for the equivalent []byte. func ReadFromPipe(in *os.File) []byte { - stat, err := in.Stat() + r, ok, err := streamio.PipeReader(in) if err != nil { ExitWithError("failed to read stat from stdin", err) } - if (stat.Mode() & os.ModeCharDevice) == 0 { - buf, err := io.ReadAll(in) - if err != nil { - ExitWithError("failed to scan bytes from stdin", err) - } - return buf + if !ok { + return nil + } + buf, err := io.ReadAll(r) + if err != nil { + ExitWithError("failed to scan bytes from stdin", err) } - return nil + return buf } +// Deprecated: reads the entire file into memory with no size cap at all — not +// even the 10 GB one the tdf commands apply — and terminates the process on +// failure. Open the file and stream from it, or use utils.ReadBytesFromFile if +// a bounded in-memory read is genuinely wanted. func ReadFromFile(filePath string) []byte { - fileToEncrypt, err := os.Open(filePath) + f, err := os.Open(filePath) if err != nil { - ExitWithError("Failed to git open file at path: "+filePath, err) + ExitWithError("Failed to open file at path: "+filePath, err) } - defer fileToEncrypt.Close() + defer f.Close() - bytes, err := io.ReadAll(fileToEncrypt) + buf, err := io.ReadAll(f) if err != nil { ExitWithError("Failed to read bytes from file at path: "+filePath, err) } - return bytes + return buf } diff --git a/otdfctl/pkg/handlers/tdf.go b/otdfctl/pkg/handlers/tdf.go index d8a5b0d69f..73900bd79a 100644 --- a/otdfctl/pkg/handlers/tdf.go +++ b/otdfctl/pkg/handlers/tdf.go @@ -161,14 +161,19 @@ func (h Handler) DecryptBytes( return out, nil } -func (h Handler) InspectTDF(toInspect []byte) (TDFInspect, []error) { - b := bytes.NewReader(toInspect) - switch sdk.GetTdfType(b) { +// InspectTDF reads the manifest and attributes of a TDF. +// +// It takes an io.ReadSeeker rather than a byte slice because only the manifest +// at the end of the archive is needed; buffering the whole payload to reach it +// costs memory proportional to the file. GetTdfType rewinds to the start, so +// the reader is positioned for LoadTDF. +func (h Handler) InspectTDF(toInspect io.ReadSeeker) (TDFInspect, []error) { + switch sdk.GetTdfType(toInspect) { case sdk.Standard: // grouping errors so we don't impact the piping of the data errs := []error{} - tdfreader, err := h.sdk.LoadTDF(bytes.NewReader(toInspect)) + tdfreader, err := h.sdk.LoadTDF(toInspect) if err != nil { if strings.Contains(err.Error(), "zip: not a valid zip file") { return TDFInspect{}, []error{ErrTDFInspectFailNotInspectable} diff --git a/otdfctl/pkg/streamio/input.go b/otdfctl/pkg/streamio/input.go new file mode 100644 index 0000000000..dc7c29a874 --- /dev/null +++ b/otdfctl/pkg/streamio/input.go @@ -0,0 +1,114 @@ +// Package streamio provides the input and output plumbing shared by commands +// that move payloads too large to hold in memory. +// +// Everything here returns errors rather than terminating the process, so the +// decision to exit stays with the command layer. That is the difference from +// the older helpers in pkg/cli, which call cli.ExitWithError from inside the +// read and are therefore unusable from anywhere that wants to recover. +package streamio + +import ( + "bufio" + "errors" + "io" + "os" +) + +// PipeBufferSize is the window PipeReader buffers over a pipe — generous +// enough that a typical CLI payload is served from a single read. +const PipeBufferSize = 1024 * 1024 + +// ErrNoInput reports that a command was given neither a file argument nor a +// non-empty pipe. It is distinct from a failure to open a named file, which +// callers report differently. +var ErrNoInput = errors.New("no input provided") + +// PipeReader reports whether in is a pipe or redirect carrying at least one +// byte, and returns a reader over it. +// +// Presence is established with a one-byte Peek rather than a read, so the +// payload still reaches the caller intact and nothing is buffered beyond the +// reader's window. A terminal, or an empty redirect such as +// `otdfctl encrypt < /dev/null`, reports false — matching the behavior of a +// buffered implementation, which decides the same question by checking whether +// a full read came back empty. +func PipeReader(in *os.File) (*bufio.Reader, bool, error) { + stat, err := in.Stat() + if err != nil { + return nil, false, err + } + if (stat.Mode() & os.ModeCharDevice) != 0 { + return nil, false, nil + } + + r := bufio.NewReaderSize(in, PipeBufferSize) + if _, err := r.Peek(1); err != nil { + if errors.Is(err, io.EOF) { + return nil, false, nil + } + return nil, false, err + } + return r, true, nil +} + +// Spool copies r into a temporary file and rewinds it, giving a seekable view +// of a stream that has none. +// +// A TDF's manifest sits at the end of the archive, so any command that needs +// to seek it cannot consume a pipe directly. Spooling trades disk for the +// memory a whole-payload read would use, and needs a TMPDIR with room for the +// full TDF — it fails loudly if there isn't one. +// +// The returned cleanup must run on every path. cli.ExitWithError calls os.Exit +// and skips deferred functions, so deferring it alone is not enough. +func Spool(r io.Reader) (*os.File, func(), error) { + f, err := os.CreateTemp("", "otdfctl-spool-*.tdf") + if err != nil { + return nil, func() {}, err + } + cleanup := func() { + f.Close() + os.Remove(f.Name()) + } + if _, err := io.Copy(f, r); err != nil { + cleanup() + return nil, func() {}, err + } + if _, err := f.Seek(0, io.SeekStart); err != nil { + cleanup() + return nil, func() {}, err + } + return f, cleanup, nil +} + +// OpenSeekable resolves a command's input to something seekable: the named file +// when one is given, otherwise piped stdin spooled to disk. It returns +// ErrNoInput for the same "nothing to read" condition whether the file argument +// was absent or the pipe was empty. +// +// The returned cleanup must run on every path, per Spool. +func OpenSeekable(path string) (*os.File, func(), error) { + if path != "" { + f, err := os.Open(path) + if err != nil { + return nil, func() {}, err + } + if _, err := f.Seek(0, io.SeekCurrent); err == nil { + return f, func() { f.Close() }, nil + } + // Not seekable (a FIFO or /dev/fd/N): spool it like piped stdin so + // callers still get something they can seek. + spooled, cleanup, err := Spool(f) + f.Close() + return spooled, cleanup, err + } + + piped, ok, err := PipeReader(os.Stdin) + if err != nil { + return nil, func() {}, err + } + if !ok { + return nil, func() {}, ErrNoInput + } + return Spool(piped) +} diff --git a/otdfctl/pkg/streamio/input_test.go b/otdfctl/pkg/streamio/input_test.go new file mode 100644 index 0000000000..2c8f8ad3d5 --- /dev/null +++ b/otdfctl/pkg/streamio/input_test.go @@ -0,0 +1,213 @@ +package streamio + +import ( + "io" + "os" + "path/filepath" + "runtime" + "strings" + "syscall" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestPipeReader(t *testing.T) { + for _, tc := range []struct { + name string + content string + wantOK bool + }{ + {name: "with data", content: "hello world", wantOK: true}, + {name: "empty pipe reports absent", content: "", wantOK: false}, + } { + t.Run(tc.name, func(t *testing.T) { + r, w, err := os.Pipe() + require.NoError(t, err) + defer r.Close() + + go func() { + defer w.Close() + _, _ = io.WriteString(w, tc.content) + }() + + got, ok, err := PipeReader(r) + require.NoError(t, err) + require.Equal(t, tc.wantOK, ok) + if !tc.wantOK { + return + } + + // Peek must not consume: the whole payload is still readable. + all, err := io.ReadAll(got) + require.NoError(t, err) + assert.Equal(t, tc.content, string(all)) + }) + } +} + +func TestPipeReaderPreservesPayloadLargerThanBuffer(t *testing.T) { + content := strings.Repeat("a", PipeBufferSize*2+7) + + r, w, err := os.Pipe() + require.NoError(t, err) + defer r.Close() + + go func() { + defer w.Close() + _, _ = io.WriteString(w, content) + }() + + got, ok, err := PipeReader(r) + require.NoError(t, err) + require.True(t, ok) + + all, err := io.ReadAll(got) + require.NoError(t, err) + assert.Len(t, all, len(content)) +} + +func TestPipeReaderOnTerminalReportsAbsent(t *testing.T) { + // A regular file is not a char device, so use os.Stdin's actual mode only + // when it is one; otherwise this assertion is vacuous and we skip. + stat, err := os.Stdin.Stat() + require.NoError(t, err) + if (stat.Mode() & os.ModeCharDevice) == 0 { + t.Skip("stdin is not a terminal under this test runner") + } + _, ok, err := PipeReader(os.Stdin) + require.NoError(t, err) + assert.False(t, ok) +} + +func TestSpoolIsSeekableAndComplete(t *testing.T) { + // Larger than any plausible internal buffer, so a truncating copy shows up. + content := strings.Repeat("xyz", PipeBufferSize) + + f, cleanup, err := Spool(strings.NewReader(content)) + require.NoError(t, err) + defer cleanup() + + // The spool must be positioned at the head, not at the end of the copy. + pos, err := f.Seek(0, io.SeekCurrent) + require.NoError(t, err) + assert.Equal(t, int64(0), pos, "spool must be rewound for the caller") + + got, err := io.ReadAll(f) + require.NoError(t, err) + assert.Len(t, got, len(content)) + + // Seeking backwards is the whole reason for spooling: a TDF's manifest is at + // the end of the archive, so the reader has to be able to go back. + _, err = f.Seek(0, io.SeekStart) + require.NoError(t, err) + head := make([]byte, 3) + _, err = io.ReadFull(f, head) + require.NoError(t, err) + assert.Equal(t, "xyz", string(head)) +} + +func TestSpoolCleanupRemovesTheFile(t *testing.T) { + f, cleanup, err := Spool(strings.NewReader("payload")) + require.NoError(t, err) + + name := f.Name() + require.FileExists(t, name) + + cleanup() + assert.NoFileExists(t, name, "the spool must not outlive the command") +} + +func TestOpenSeekableReadsNamedFile(t *testing.T) { + path := filepath.Join(t.TempDir(), "in.tdf") + require.NoError(t, os.WriteFile(path, []byte("payload"), 0o600)) + + in, cleanup, err := OpenSeekable(path) + require.NoError(t, err) + defer cleanup() + + got, err := io.ReadAll(in) + require.NoError(t, err) + assert.Equal(t, "payload", string(got)) + + // A named file is opened directly, not copied through a spool. + assert.Equal(t, path, in.Name()) +} + +func TestOpenSeekableReportsMissingFile(t *testing.T) { + _, _, err := OpenSeekable(filepath.Join(t.TempDir(), "absent.tdf")) + require.Error(t, err) + // Distinguishable from "you gave me nothing", which callers report differently. + require.NotErrorIs(t, err, ErrNoInput) + require.ErrorIs(t, err, os.ErrNotExist) +} + +func TestOpenSeekableSpoolsNonSeekableNamedFile(t *testing.T) { + if runtime.GOOS == "windows" { + t.Skip("FIFOs are not portable to windows") + } + + path := filepath.Join(t.TempDir(), "in.fifo") + require.NoError(t, syscall.Mkfifo(path, 0o600)) + + content := "a named pipe cannot seek, so this has to be spooled" + go func() { + w, err := os.OpenFile(path, os.O_WRONLY, 0) + if err != nil { + return + } + defer w.Close() + _, _ = io.WriteString(w, content) + }() + + in, cleanup, err := OpenSeekable(path) + require.NoError(t, err) + defer cleanup() + + got, err := io.ReadAll(in) + require.NoError(t, err) + assert.Equal(t, content, string(got)) + + // A FIFO's own bytes were consumed into the spool; the caller's handle is a + // distinct, seekable temp file. + assert.NotEqual(t, path, in.Name()) + _, err = in.Seek(0, io.SeekStart) + require.NoError(t, err, "the whole point of spooling is that the result seeks") +} + +func TestOpenSeekableReadsFromStdinPipe(t *testing.T) { + r, w, err := os.Pipe() + require.NoError(t, err) + + origStdin := os.Stdin + os.Stdin = r + defer func() { os.Stdin = origStdin }() + + content := "piped stdin, spooled to disk" + go func() { + defer w.Close() + _, _ = io.WriteString(w, content) + }() + + in, cleanup, err := OpenSeekable("") + require.NoError(t, err) + defer cleanup() + + got, err := io.ReadAll(in) + require.NoError(t, err) + assert.Equal(t, content, string(got)) +} + +func TestOpenSeekableReturnsErrNoInputForEmptyStdinPipe(t *testing.T) { + r, w, err := os.Pipe() + require.NoError(t, err) + require.NoError(t, w.Close()) + + origStdin := os.Stdin + os.Stdin = r + defer func() { os.Stdin = origStdin }() + + _, _, err = OpenSeekable("") + require.ErrorIs(t, err, ErrNoInput) +} diff --git a/otdfctl/pkg/streamio/output.go b/otdfctl/pkg/streamio/output.go new file mode 100644 index 0000000000..7d3bb35d8b --- /dev/null +++ b/otdfctl/pkg/streamio/output.go @@ -0,0 +1,86 @@ +package streamio + +import ( + "errors" + "os" + "path/filepath" +) + +// ErrOutputFileFinished reports that Commit or Cleanup was called on an +// OutputFile that had already been committed or discarded. +var ErrOutputFileFinished = errors.New("streamio: output file already committed or discarded") + +// outputFileMode is the permission Commit applies to the destination. +// os.CreateTemp always creates the temp file with 0600; without an explicit +// Chmod that would leak onto the destination regardless of the caller's +// umask, so this matches the common default a plain os.Create would produce. +const outputFileMode = 0o644 + +// OutputFile writes to a temporary file alongside the destination and renames +// it into place only once the write has succeeded, so an interrupted or failed +// run leaves no partial output where a complete file is expected. +// +// Note that cli.ExitWithError calls os.Exit, which does not run deferred +// functions. Cleanup must therefore be called explicitly on every error path, +// not only via defer. +type OutputFile struct { + f *os.File + path string + finished bool +} + +// NewOutputFile creates the temporary file in the destination's own directory. +// A rename is only atomic within a single filesystem, so the temp file must +// live beside the destination rather than in a shared temp directory — +// Commit's os.Rename fails outright (EXDEV) if that invariant is broken. +func NewOutputFile(path string) (*OutputFile, error) { + f, err := os.CreateTemp(filepath.Dir(path), "."+filepath.Base(path)+".tmp-*") + if err != nil { + return nil, err + } + return &OutputFile{f: f, path: path}, nil +} + +func (o *OutputFile) Write(p []byte) (int, error) { return o.f.Write(p) } + +// Name reports the path of the temporary file currently being written, which is +// not the destination until Commit succeeds. +func (o *OutputFile) Name() string { return o.f.Name() } + +// Commit closes the temporary file and moves it onto the destination path. +// +// It returns ErrOutputFileFinished if called more than once, or after Cleanup — +// otherwise a second call would re-enter the close/rename-failure path against +// an already-closed or already-moved file and report a spurious error. +func (o *OutputFile) Commit() error { + if o.finished { + return ErrOutputFileFinished + } + o.finished = true + + if err := o.f.Chmod(outputFileMode); err != nil { + o.f.Close() + os.Remove(o.f.Name()) + return err + } + if err := o.f.Close(); err != nil { + os.Remove(o.f.Name()) + return err + } + if err := os.Rename(o.f.Name(), o.path); err != nil { + os.Remove(o.f.Name()) + return err + } + return nil +} + +// Cleanup discards the temporary file. It is a no-op after a successful Commit +// (or a prior Cleanup), so it is safe to both defer it and call it directly. +func (o *OutputFile) Cleanup() { + if o.finished { + return + } + o.finished = true + o.f.Close() + os.Remove(o.f.Name()) +} diff --git a/otdfctl/pkg/streamio/output_test.go b/otdfctl/pkg/streamio/output_test.go new file mode 100644 index 0000000000..e7e73cb643 --- /dev/null +++ b/otdfctl/pkg/streamio/output_test.go @@ -0,0 +1,148 @@ +package streamio + +import ( + "os" + "path/filepath" + "strings" + "testing" + + "github.com/stretchr/testify/assert" + "github.com/stretchr/testify/require" +) + +func TestOutputFileCommitRenamesIntoPlace(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + + _, err = o.Write([]byte("payload")) + require.NoError(t, err) + + // Nothing is visible at the destination until Commit. + _, err = os.Stat(dest) + require.ErrorIs(t, err, os.ErrNotExist, "destination must not exist before Commit") + + require.NoError(t, o.Commit()) + + got, err := os.ReadFile(dest) + require.NoError(t, err) + assert.Equal(t, "payload", string(got)) + assert.Empty(t, tempSiblings(t, dir, "out.tdf"), "temp file should be gone after Commit") +} + +func TestOutputFileCleanupLeavesNoPartialOutput(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + _, err = o.Write([]byte("partial")) + require.NoError(t, err) + + // Simulates the failure path: encryption died after some bytes were written. + o.Cleanup() + + _, err = os.Stat(dest) + require.ErrorIs(t, err, os.ErrNotExist, "a failed run must not leave a partial file") + assert.Empty(t, tempSiblings(t, dir, "out.tdf"), "a failed run must not leave a temp file") +} + +func TestOutputFileCleanupAfterCommitIsNoop(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + _, err = o.Write([]byte("payload")) + require.NoError(t, err) + require.NoError(t, o.Commit()) + + // Both deferred and explicit cleanup run on the success path. + o.Cleanup() + o.Cleanup() + + got, err := os.ReadFile(dest) + require.NoError(t, err) + assert.Equal(t, "payload", string(got), "Cleanup after Commit must not delete the output") +} + +func TestOutputFileTempIsSiblingOfDestination(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + defer o.Cleanup() + + // A rename is only atomic within one filesystem — a shared temp dir on a + // different filesystem would make Commit's os.Rename fail outright — so the + // temp file must live beside the destination. + assert.Equal(t, dir, filepath.Dir(o.Name())) +} + +func TestOutputFileCommitSetsReadableMode(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + _, err = o.Write([]byte("payload")) + require.NoError(t, err) + require.NoError(t, o.Commit()) + + info, err := os.Stat(dest) + require.NoError(t, err) + assert.Equal(t, os.FileMode(outputFileMode), info.Mode().Perm(), + "os.CreateTemp defaults to 0600; Commit must not leak that onto the destination") +} + +func TestOutputFileCommitAfterCommitReturnsError(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + _, err = o.Write([]byte("payload")) + require.NoError(t, err) + require.NoError(t, o.Commit()) + + require.ErrorIs(t, o.Commit(), ErrOutputFileFinished) + + got, err := os.ReadFile(dest) + require.NoError(t, err) + assert.Equal(t, "payload", string(got), "a redundant Commit must not disturb the already-committed output") +} + +func TestOutputFileCommitAfterCleanupReturnsError(t *testing.T) { + dir := t.TempDir() + dest := filepath.Join(dir, "out.tdf") + + o, err := NewOutputFile(dest) + require.NoError(t, err) + _, err = o.Write([]byte("payload")) + require.NoError(t, err) + + o.Cleanup() + require.ErrorIs(t, o.Commit(), ErrOutputFileFinished) + + _, err = os.Stat(dest) + require.ErrorIs(t, err, os.ErrNotExist, "Commit after Cleanup must not create the destination") +} + +// tempSiblings returns any leftover temp files NewOutputFile would have created +// for dest in dir. +func tempSiblings(t *testing.T, dir, dest string) []string { + t.Helper() + entries, err := os.ReadDir(dir) + require.NoError(t, err) + + var found []string + for _, e := range entries { + if strings.HasPrefix(e.Name(), "."+dest+".tmp-") { + found = append(found, e.Name()) + } + } + return found +} From c1ce2ced0dd5669f1ad81f3dd2a8fce789b69f44 Mon Sep 17 00:00:00 2001 From: Dave Mihalcik Date: Mon, 31 Aug 2026 21:33:23 -0400 Subject: [PATCH 2/2] 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") } }