Skip to content

Commit

Permalink
Merge pull request #111 from mingmingtsao/fix-bitset-readfrom-incompl…
Browse files Browse the repository at this point in the history
…ete-buffer

fix `ReadFrom` generate corrupted bitset when reader incompletely fills buf
  • Loading branch information
lemire authored Sep 14, 2022
2 parents 55b2907 + 6079724 commit e160993
Show file tree
Hide file tree
Showing 2 changed files with 29 additions and 9 deletions.
2 changes: 1 addition & 1 deletion bitset.go
Original file line number Diff line number Diff line change
Expand Up @@ -945,7 +945,7 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) {
nWords := wordsNeeded(uint(length))
reader := bufio.NewReader(io.LimitReader(stream, 8*int64(nWords)))
for i := 0; i < nWords; i++ {
if _, err := reader.Read(item[:]); err != nil {
if _, err := io.ReadFull(reader, item[:]); err != nil {
if err == io.EOF {
err = io.ErrUnexpectedEOF
}
Expand Down
36 changes: 28 additions & 8 deletions bitset_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -1713,18 +1713,38 @@ func TestWriteTo(t *testing.T) {
}
}

type inCompleteRetBufReader struct {
returnEvery int64
reader io.Reader
offset int64
}

func (ir *inCompleteRetBufReader) Read(b []byte) (n int, err error) {
if ir.returnEvery > 0 {
maxRead := ir.returnEvery - (ir.offset % ir.returnEvery)
if len(b) > int(maxRead) {
b = b[:maxRead]
}
}
n, err = ir.reader.Read(b)
ir.offset += int64(n)
return
}

func TestReadFrom(t *testing.T) {
addBuf := []byte(`12345678`) // Bytes after stream
tests := []struct {
length uint
oneEvery uint
input string // base64+gzipped
wantErr error
length uint
oneEvery uint
input string // base64+gzipped
wantErr error
returnEvery int64
}{
{
length: 9585,
oneEvery: 97,
input: "H4sIAAAAAAAC/2IAA9VCCM3AyMDAwMSACVgYGBg4sIgLMDAwKGARd2BgYGjAFB41noDx6IAJajw64IAajw4UoMajg4ZR4/EaP5pQh1g+MDQyNjE1M7cABAAA//9W5OoOwAQAAA==",
length: 9585,
oneEvery: 97,
input: "H4sIAAAAAAAC/2IAA9VCCM3AyMDAwMSACVgYGBg4sIgLMDAwKGARd2BgYGjAFB41noDx6IAJajw64IAajw4UoMajg4ZR4/EaP5pQh1g+MDQyNjE1M7cABAAA//9W5OoOwAQAAA==",
returnEvery: 127,
},
{
length: 1337,
Expand Down Expand Up @@ -1764,7 +1784,7 @@ func TestReadFrom(t *testing.T) {
fatalErr(gz.Close())

bs := New(test.length)
_, err = bs.ReadFrom(&buf)
_, err = bs.ReadFrom(&inCompleteRetBufReader{returnEvery: test.returnEvery, reader: &buf})
if err != nil {
if errors.Is(err, test.wantErr) {
// Correct, nothing more we can test.
Expand Down

0 comments on commit e160993

Please sign in to comment.