Skip to content

Commit

Permalink
s2: Document and test how to peek the stream for skippable blocks (#918)
Browse files Browse the repository at this point in the history
Co-authored-by: Klaus Post <klauspost@gmail.com>
  • Loading branch information
Jille and klauspost authored Jan 25, 2024
1 parent 3deb878 commit 6662a21
Show file tree
Hide file tree
Showing 2 changed files with 49 additions and 0 deletions.
4 changes: 4 additions & 0 deletions s2/reader.go
Original file line number Diff line number Diff line change
Expand Up @@ -104,6 +104,8 @@ func ReaderIgnoreStreamIdentifier() ReaderOption {
// For each chunk with the ID, the callback is called with the content.
// Any returned non-nil error will abort decompression.
// Only one callback per ID is supported, latest sent will be used.
// You can peek the stream, triggering the callback, by doing a Read with a 0
// byte buffer.
func ReaderSkippableCB(id uint8, fn func(r io.Reader) error) ReaderOption {
return func(r *Reader) error {
if id < 0x80 || id > 0xfd {
Expand Down Expand Up @@ -1053,6 +1055,8 @@ func (r *Reader) ReadByte() (byte, error) {
// Any returned non-nil error will abort decompression.
// Only one callback per ID is supported, latest sent will be used.
// Sending a nil function will disable previous callbacks.
// You can peek the stream, triggering the callback, by doing a Read with a 0
// byte buffer.
func (r *Reader) SkippableCB(id uint8, fn func(r io.Reader) error) error {
if id < 0x80 || id >= chunkTypePadding {
return fmt.Errorf("ReaderSkippableCB: Invalid id provided, must be 0x80-0xfe (inclusive)")
Expand Down
45 changes: 45 additions & 0 deletions s2/reader_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
// Copyright (c) 2019+ Klaus Post. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package s2

import (
"bytes"
"io"
"testing"
)

func TestLeadingSkippableBlock(t *testing.T) {
var buf bytes.Buffer
w := NewWriter(&buf)
if err := w.AddSkippableBlock(0x80, []byte("skippable block")); err != nil {
t.Fatalf("w.AddSkippableBlock: %v", err)
}
if _, err := w.Write([]byte("some data")); err != nil {
t.Fatalf("w.Write: %v", err)
}
if err := w.Close(); err != nil {
t.Fatalf("w.Close: %v", err)
}
r := NewReader(&buf)
var sb []byte
r.SkippableCB(0x80, func(sr io.Reader) error {
var err error
sb, err = io.ReadAll(sr)
return err
})
if _, err := r.Read([]byte{}); err != nil {
t.Errorf("empty read failed: %v", err)
}
if !bytes.Equal(sb, []byte("skippable block")) {
t.Errorf("didn't get correct data from skippable block: %q", string(sb))
}
data, err := io.ReadAll(r)
if err != nil {
t.Fatalf("r.Read: %v", err)
}
if !bytes.Equal(data, []byte("some data")) {
t.Errorf("didn't get correct compressed data: %q", string(data))
}
}

0 comments on commit 6662a21

Please sign in to comment.