diff --git a/internal/io/limitWriter_test.go b/internal/io/limitWriter_test.go new file mode 100644 index 00000000..ef241c58 --- /dev/null +++ b/internal/io/limitWriter_test.go @@ -0,0 +1,45 @@ +package io + +import ( + "bytes" + "testing" +) + +func TestLimitWriter(t *testing.T) { + limit := int64(10) + longString := "1234567891011" + + tests := []struct { + input string + expected string + written int + }{ + {"hello", "hello", 5}, + {" world", " world", 6}, + {"!", "!", 1}, + {"1234567891011", "1234567891", 10}, + } + + for _, tt := range tests { + var buf bytes.Buffer + lw := NewLimitWriter(&buf, limit) + n, err := lw.Write([]byte(tt.input)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n != tt.written { + t.Errorf("expected %d bytes written, got %d", tt.written, n) + } + if buf.String() != tt.expected { + t.Errorf("expected buffer %q, got %q", tt.expected, buf.String()) + } + + n, err = lw.Write([]byte(longString)) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + if n == len(longString) { + t.Errorf("should not write more than the limit") + } + } +} diff --git a/internal/io/limitwriter.go b/internal/io/limitwriter.go new file mode 100644 index 00000000..102dee19 --- /dev/null +++ b/internal/io/limitwriter.go @@ -0,0 +1,45 @@ +// Copyright The Notary Project Authors. +// Licensed under the Apache License, Version 2.0 (the "License"); +// you may not use this file except in compliance with the License. +// You may obtain a copy of the License at +// +// http://www.apache.org/licenses/LICENSE-2.0 +// +// Unless required by applicable law or agreed to in writing, software +// distributed under the License is distributed on an "AS IS" BASIS, +// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +// See the License for the specific language governing permissions and +// limitations under the License. + +package io + +import "io" + +// LimitWriter is a writer that writes to an underlying writer up to a limit. +type LimitWriter struct { + w io.Writer + limit int64 + written int64 +} + +// NewLimitWriter returns a new LimitWriter that writes to w. +// +// parameters: +// w: the writer to write to +// limit: the maximum number of bytes to write +func NewLimitWriter(w io.Writer, limit int64) *LimitWriter { + return &LimitWriter{w: w, limit: limit} +} + +func (lw *LimitWriter) Write(p []byte) (int, error) { + remaining := lw.limit - lw.written + if remaining <= 0 { + return 0, nil + } + if int64(len(p)) > remaining { + p = p[:remaining] + } + n, err := lw.w.Write(p) + lw.written += int64(n) + return n, err +} diff --git a/plugin/plugin.go b/plugin/plugin.go index d7f0a932..77869b48 100644 --- a/plugin/plugin.go +++ b/plugin/plugin.go @@ -27,12 +27,18 @@ import ( "path/filepath" "strings" + "github.com/notaryproject/notation-go/internal/io" "github.com/notaryproject/notation-go/internal/slices" "github.com/notaryproject/notation-go/log" "github.com/notaryproject/notation-go/plugin/proto" "github.com/notaryproject/notation-plugin-framework-go/plugin" ) +const ( + // maxPluginOutputSize is the maximum size of the plugin output. + maxPluginOutputSize = 10 * 1024 * 1024 // 10MB +) + var executor commander = &execCommander{} // for unit test // GenericPlugin is the base requirement to be a plugin. @@ -218,12 +224,12 @@ func (c execCommander) Output(ctx context.Context, name string, command plugin.C var stdout, stderr bytes.Buffer cmd := exec.CommandContext(ctx, name, string(command)) cmd.Stdin = bytes.NewReader(req) - cmd.Stderr = &stderr - cmd.Stdout = &stdout + cmd.Stderr = io.NewLimitWriter(&stderr, maxPluginOutputSize) + cmd.Stdout = io.NewLimitWriter(&stdout, maxPluginOutputSize) err := cmd.Run() if err != nil { if errors.Is(ctx.Err(), context.DeadlineExceeded) { - return nil, stderr.Bytes(), fmt.Errorf("'%s %s' command execution timeout: %w", name, string(command), err); + return nil, stderr.Bytes(), fmt.Errorf("'%s %s' command execution timeout: %w", name, string(command), err) } return nil, stderr.Bytes(), err }