Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
feat: Add support for receiving gzip event payloads
- Loading branch information
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,109 @@ | ||
| package util | ||
|
|
||
| // Taken from event-recorder reader.go | ||
|
|
||
| import ( | ||
| "compress/gzip" | ||
| "errors" | ||
| "io" | ||
| ) | ||
|
|
||
| // PayloadReader is an implementation of io.Reader that reads bytes off the request body | ||
| // optionally decompresses them, and has a limit attached. | ||
| // If the limit is reached, an error will be returned and the underlying stream will be closed. | ||
| // | ||
| // Note: limit is applied to *both* compressed and uncompressed number of bytes. This | ||
| // protects us from potential zipbombs | ||
| type PayloadReader struct { | ||
| IsGzipped bool | ||
| MaxBytes int64 | ||
|
|
||
| uncompressedBytesRead int64 | ||
|
|
||
| wrappedBaseStream *byteCountingReader | ||
| stream *io.LimitedReader | ||
| } | ||
|
|
||
| // NewReader creates a new reader | ||
| func NewReader(r io.ReadCloser, isGzipped bool, maxBytes int64) (io.ReadCloser, error) { | ||
| // If this isn't compressed, and we don't want to limit the size, just | ||
| // return the original reader | ||
| if !isGzipped && maxBytes == 0 { | ||
| return r, nil | ||
| } | ||
|
|
||
| var baseStream = &byteCountingReader{ | ||
| bytesRead: 0, | ||
| baseStream: r, | ||
| } | ||
| var s io.Reader = baseStream | ||
|
|
||
| if isGzipped { | ||
| var err error | ||
| gzipReader, err := gzip.NewReader(s) | ||
| if err != nil { | ||
| return nil, err | ||
| } | ||
|
|
||
| // If we don't want to limit the size, just return the gzip reader | ||
| if maxBytes == 0 { | ||
| return gzipReader, nil | ||
| } | ||
|
|
||
| s = gzipReader | ||
| } | ||
| stream := io.LimitReader(s, maxBytes).(*io.LimitedReader) | ||
|
|
||
| payloadReader := &PayloadReader{ | ||
| IsGzipped: isGzipped, | ||
| MaxBytes: maxBytes, | ||
| wrappedBaseStream: baseStream, | ||
| stream: stream, | ||
| } | ||
| return payloadReader, nil | ||
| } | ||
|
|
||
| // GetBytesRead returns the total number of bytes read off the original stream | ||
| func (pr *PayloadReader) GetBytesRead() int64 { | ||
| return pr.wrappedBaseStream.bytesRead | ||
| } | ||
|
|
||
| // GetUncompressedBytesRead Total number of Bytes in the uncompressed stream. | ||
| // | ||
| // GetBytesRead and GetUncompressedBytesRead will return the same value if the stream is uncompressed. | ||
| func (pr *PayloadReader) GetUncompressedBytesRead() int64 { | ||
| return pr.uncompressedBytesRead | ||
| } | ||
|
|
||
| func (pr *PayloadReader) Read(p []byte) (int, error) { | ||
| n, err := pr.stream.Read(p) | ||
| pr.uncompressedBytesRead += int64(n) | ||
| if err != nil { | ||
| return n, err | ||
| } | ||
| if pr.stream.N <= 0 { | ||
| pr.Close() | ||
|
Check failure on line 85 in internal/util/reader.go
|
||
| return n, errors.New("Max bytes exceeded") | ||
|
Check failure on line 86 in internal/util/reader.go
|
||
| } | ||
| return n, err | ||
| } | ||
|
|
||
| func (pr *PayloadReader) Close() error { | ||
| c, ok := pr.wrappedBaseStream.baseStream.(io.Closer) | ||
| if ok { | ||
| return c.Close() | ||
| } | ||
| return nil | ||
| } | ||
|
|
||
| // a simple reader decorator that keeps a running total of bytes | ||
| type byteCountingReader struct { | ||
| baseStream io.Reader | ||
| bytesRead int64 | ||
| } | ||
|
|
||
| func (bt *byteCountingReader) Read(p []byte) (int, error) { | ||
| n, err := bt.baseStream.Read(p) | ||
| bt.bytesRead += int64(n) | ||
| return n, err | ||
| } | ||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,77 @@ | ||
| package util | ||
|
|
||
| import ( | ||
| "bytes" | ||
| "compress/gzip" | ||
| "io" | ||
| "testing" | ||
|
|
||
| "github.com/stretchr/testify/assert" | ||
| ) | ||
|
|
||
| func TestUncompressed(t *testing.T) { | ||
| // make sure the bytes read are correct for non-gzipped streams | ||
| nonZipBytes := []byte("this is a test") | ||
|
|
||
| reader, _ := NewReader(io.NopCloser(bytes.NewReader(nonZipBytes)), false, 1000) | ||
| payloadReader, _ := reader.(*PayloadReader) | ||
| readBytes, _ := io.ReadAll(reader) | ||
|
|
||
| assert.Equal(t, int64(len(nonZipBytes)), payloadReader.GetBytesRead()) | ||
| assert.Equal(t, int64(len(nonZipBytes)), payloadReader.GetUncompressedBytesRead()) | ||
| assert.Equal(t, nonZipBytes, readBytes) | ||
| } | ||
|
|
||
| func TestPayloadBytesTracking(t *testing.T) { | ||
| // build a string | ||
| var ret string | ||
| for i := 0; i < 500; i++ { | ||
| ret += "00" | ||
| } | ||
|
|
||
| // make sure the bytes read are correct for gzipped streams | ||
| nonZipBytes := []byte(ret) | ||
|
|
||
| var b bytes.Buffer | ||
| w := gzip.NewWriter(&b) | ||
| w.Write(nonZipBytes) | ||
| w.Close() | ||
| zipBytes := b.Bytes() | ||
|
|
||
| reader, _ := NewReader(io.NopCloser(bytes.NewReader(zipBytes)), true, 10000) | ||
| payloadReader, _ := reader.(*PayloadReader) | ||
| readBytes, _ := io.ReadAll(reader) | ||
|
|
||
| assert.Equal(t, int64(len(zipBytes)), payloadReader.GetBytesRead()) // compressed is 32 | ||
| assert.Equal(t, int64(len(nonZipBytes)), payloadReader.GetUncompressedBytesRead()) // uncompressed is 1000 | ||
| assert.Equal(t, nonZipBytes, readBytes) | ||
| } | ||
|
|
||
| func TestZipBombing(t *testing.T) { | ||
| // build a string | ||
| var ret string | ||
| for i := 0; i < 500; i++ { | ||
| ret += "00" | ||
| } | ||
|
|
||
| nonZipBytes := []byte(ret) | ||
|
|
||
| var b bytes.Buffer | ||
| w := gzip.NewWriter(&b) | ||
| w.Write(nonZipBytes) | ||
| w.Close() | ||
| zipBytes := b.Bytes() | ||
|
|
||
| maxBytes := int64(100) | ||
| reader, _ := NewReader(io.NopCloser(bytes.NewReader(zipBytes)), true, maxBytes) | ||
| payloadReader, _ := reader.(*PayloadReader) | ||
| bytesRead, err := io.ReadAll(reader) | ||
|
|
||
| // we should have an error because the uncompressed stream is larger than maxBytes | ||
| assert.Error(t, err) | ||
| assert.Equal(t, int64(len(bytesRead)), maxBytes) | ||
|
|
||
| assert.Equal(t, int64(len(zipBytes)), payloadReader.GetBytesRead()) | ||
| assert.Equal(t, maxBytes, payloadReader.GetUncompressedBytesRead()) | ||
|
|
||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters