Skip to content

Commit 02e1fe1

Browse files
LessUpCopilot
LessUp
and
Copilot
committed
fix: harden streaming go helpers
Co-authored-by: Copilot <223556219+Copilot@users.noreply.github.com>
1 parent fe47936 commit 02e1fe1

6 files changed

Lines changed: 177 additions & 4 deletions

File tree

‎algorithms/arithmetic/go/streaming.go‎

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -163,9 +163,14 @@ func (d *StreamingDecoder) Finish(out []byte) (int, error) {
163163
err := Decode(bytes.NewReader(d.inputBuf.Bytes()), &outBuf)
164164
if err != nil {
165165
d.state = codec.StateError
166+
errStr := err.Error()
166167
if err.Error() == "invalid input file format" {
167168
return 0, codec.ErrCorrupt
168169
}
170+
if errStr == "failed to read frequency table: unexpected EOF" ||
171+
errStr == "failed to read frequency table: EOF" {
172+
return 0, codec.ErrTruncated
173+
}
169174
return 0, err
170175
}
171176

‎algorithms/shared/go/codec/buffer.go‎

Lines changed: 22 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -39,10 +39,20 @@ func EncodeBuffer(encoder Encoder, input []byte) ([]byte, error) {
3939
if err != ErrBufTooSmall {
4040
break
4141
}
42+
totalWritten += n
43+
if totalWritten > MaxOutputSize {
44+
return nil, ErrSizeLimit
45+
}
4246
if len(outBuf) >= MaxOutputSize {
4347
return nil, ErrSizeLimit
4448
}
45-
outBuf = make([]byte, growBuffer(len(outBuf), MaxOutputSize))
49+
newSize := growBuffer(len(outBuf), MaxOutputSize)
50+
if newSize <= len(outBuf) {
51+
return nil, ErrSizeLimit
52+
}
53+
newBuf := make([]byte, newSize)
54+
copy(newBuf, outBuf[:totalWritten])
55+
outBuf = newBuf
4656
}
4757
if err != nil {
4858
return nil, err
@@ -87,10 +97,20 @@ func DecodeBuffer(decoder Decoder, input []byte) ([]byte, error) {
8797
if err != ErrBufTooSmall {
8898
break
8999
}
100+
totalWritten += n
101+
if totalWritten > MaxOutputSize {
102+
return nil, ErrSizeLimit
103+
}
90104
if len(outBuf) >= MaxOutputSize {
91105
return nil, ErrSizeLimit
92106
}
93-
outBuf = make([]byte, growBuffer(len(outBuf), MaxOutputSize))
107+
newSize := growBuffer(len(outBuf), MaxOutputSize)
108+
if newSize <= len(outBuf) {
109+
return nil, ErrSizeLimit
110+
}
111+
newBuf := make([]byte, newSize)
112+
copy(newBuf, outBuf[:totalWritten])
113+
outBuf = newBuf
94114
}
95115
if err != nil {
96116
return nil, err

‎algorithms/shared/go/codec/buffer_internal_test.go‎

Lines changed: 95 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,95 @@
1+
package codec
2+
3+
import (
4+
"errors"
5+
"testing"
6+
)
7+
8+
type processRetryEncoder struct {
9+
processCalls int
10+
}
11+
12+
func (e *processRetryEncoder) Process(in []byte, out []byte) (int, error) {
13+
e.processCalls++
14+
if e.processCalls == 1 {
15+
copy(out, []byte("abc"))
16+
return 3, ErrBufTooSmall
17+
}
18+
copy(out, []byte("def"))
19+
return 3, nil
20+
}
21+
22+
func (e *processRetryEncoder) Flush(out []byte) (int, error) { return 0, nil }
23+
func (e *processRetryEncoder) Finish(out []byte) (int, error) { return 0, nil }
24+
func (e *processRetryEncoder) Reset() {}
25+
func (e *processRetryEncoder) State() State { return StateStreaming }
26+
27+
type processRetryDecoder struct {
28+
processCalls int
29+
}
30+
31+
func (d *processRetryDecoder) Process(in []byte, out []byte) (int, error) {
32+
d.processCalls++
33+
if d.processCalls == 1 {
34+
copy(out, []byte("ghi"))
35+
return 3, ErrBufTooSmall
36+
}
37+
copy(out, []byte("jkl"))
38+
return 3, nil
39+
}
40+
41+
func (d *processRetryDecoder) Flush(out []byte) (int, error) { return 0, nil }
42+
func (d *processRetryDecoder) Finish(out []byte) (int, error) { return 0, nil }
43+
func (d *processRetryDecoder) Reset() {}
44+
func (d *processRetryDecoder) State() State { return StateStreaming }
45+
46+
func TestEncodeBuffer_PreservesOutputAcrossProcessRetry(t *testing.T) {
47+
out, err := EncodeBuffer(&processRetryEncoder{}, []byte("ignored"))
48+
if err != nil {
49+
t.Fatalf("EncodeBuffer() error = %v", err)
50+
}
51+
if string(out) != "abcdef" {
52+
t.Fatalf("EncodeBuffer() = %q, want %q", out, "abcdef")
53+
}
54+
}
55+
56+
func TestDecodeBuffer_PreservesOutputAcrossProcessRetry(t *testing.T) {
57+
out, err := DecodeBuffer(&processRetryDecoder{}, []byte("ignored"))
58+
if err != nil {
59+
t.Fatalf("DecodeBuffer() error = %v", err)
60+
}
61+
if string(out) != "ghijkl" {
62+
t.Fatalf("DecodeBuffer() = %q, want %q", out, "ghijkl")
63+
}
64+
}
65+
66+
type capRetryEncoder struct {
67+
processCalls int
68+
finishCalls int
69+
stopErr error
70+
}
71+
72+
func (e *capRetryEncoder) Process(in []byte, out []byte) (int, error) {
73+
e.processCalls++
74+
if e.processCalls > 2 {
75+
if e.stopErr == nil {
76+
e.stopErr = errors.New("unexpected extra process retry")
77+
}
78+
return 0, e.stopErr
79+
}
80+
return 0, ErrBufTooSmall
81+
}
82+
func (e *capRetryEncoder) Flush(out []byte) (int, error) { return 0, nil }
83+
func (e *capRetryEncoder) Reset() {}
84+
func (e *capRetryEncoder) State() State { return StateStreaming }
85+
86+
func (e *capRetryEncoder) Finish(out []byte) (int, error) {
87+
e.finishCalls++
88+
if e.finishCalls > 2 {
89+
if e.stopErr == nil {
90+
e.stopErr = errors.New("unexpected extra retry")
91+
}
92+
return 0, e.stopErr
93+
}
94+
return 0, ErrBufTooSmall
95+
}

‎algorithms/shared/go/codec/lifecycle_test.go‎

Lines changed: 15 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -589,6 +589,21 @@ func TestError_TruncatedFrame(t *testing.T) {
589589
}
590590
}
591591

592+
func TestError_ArithmeticTruncatedFrame(t *testing.T) {
593+
input := []byte("test data for arithmetic truncation")
594+
595+
encoded, err := codec.EncodeBuffer(arithmetic.NewStreamingEncoder(), input)
596+
if err != nil {
597+
t.Fatalf("EncodeBuffer() error = %v", err)
598+
}
599+
600+
truncated := encoded[:len(encoded)/2]
601+
_, err = codec.DecodeBuffer(arithmetic.NewStreamingDecoder(), truncated)
602+
if err != codec.ErrTruncated && err != codec.ErrCorrupt {
603+
t.Fatalf("DecodeBuffer() error = %v, want ErrTruncated or ErrCorrupt", err)
604+
}
605+
}
606+
592607
func TestDecodeBuffer_GrowsWithoutMaxAllocation(t *testing.T) {
593608
input := bytes.Repeat([]byte("abcd"), 1024)
594609
encoded, err := codec.EncodeBuffer(rle.NewStreamingEncoder(), input)

‎algorithms/shared/go/codec/writer.go‎

Lines changed: 10 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -31,7 +31,11 @@ func (we *WriterEncoder) Write(p []byte) (n int, err error) {
3131
if err := we.flushOutBuf(); err != nil {
3232
return 0, err
3333
}
34-
we.outBuf = make([]byte, growBuffer(len(we.outBuf), MaxOutputSize))
34+
newSize := growBuffer(len(we.outBuf), MaxOutputSize)
35+
if newSize <= len(we.outBuf) {
36+
return 0, ErrSizeLimit
37+
}
38+
we.outBuf = make([]byte, newSize)
3539
we.outBufOffset = 0
3640
written, err = we.encoder.Process(p, we.outBuf)
3741
}
@@ -59,7 +63,11 @@ func (we *WriterEncoder) Close() error {
5963
if err := we.flushOutBuf(); err != nil {
6064
return err
6165
}
62-
we.outBuf = make([]byte, growBuffer(len(we.outBuf), MaxOutputSize))
66+
newSize := growBuffer(len(we.outBuf), MaxOutputSize)
67+
if newSize <= len(we.outBuf) {
68+
return ErrSizeLimit
69+
}
70+
we.outBuf = make([]byte, newSize)
6371
we.outBufOffset = 0
6472
written, err = we.encoder.Finish(we.outBuf)
6573
}

‎algorithms/shared/go/codec/writer_internal_test.go‎

Lines changed: 30 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,30 @@
1+
package codec
2+
3+
import (
4+
"bytes"
5+
"testing"
6+
)
7+
8+
func TestWriterEncoder_ReturnsSizeLimitAtCap(t *testing.T) {
9+
stub := &capRetryEncoder{}
10+
var sink bytes.Buffer
11+
writer := NewWriterEncoder(stub, &sink)
12+
writer.outBuf = make([]byte, MaxOutputSize)
13+
14+
err := writer.Close()
15+
if err != ErrSizeLimit {
16+
t.Fatalf("Close() error = %v, want ErrSizeLimit", err)
17+
}
18+
}
19+
20+
func TestWriterEncoder_WriteReturnsSizeLimitAtCap(t *testing.T) {
21+
stub := &capRetryEncoder{}
22+
var sink bytes.Buffer
23+
writer := NewWriterEncoder(stub, &sink)
24+
writer.outBuf = make([]byte, MaxOutputSize)
25+
26+
_, err := writer.Write([]byte("payload"))
27+
if err != ErrSizeLimit {
28+
t.Fatalf("Write() error = %v, want ErrSizeLimit", err)
29+
}
30+
}

0 commit comments

Comments
 (0)