zstd: Allow partial reads

zstd/decoder.go

@@ -124,7 +124,9 @@ func (d *Decoder) Read(p []byte) (int, error) {
 			if d.current.err != nil {
 				break
 			}
-			d.nextBlock()
+			if !d.nextBlock(n == 0) {
+				return n, nil
+			}
 		}
 	}
 	if len(d.current.b) > 0 {
@@ -252,7 +254,7 @@ func (d *Decoder) WriteTo(w io.Writer) (int64, error) {
 		if d.current.err != nil {
 			break
 		}
-		d.nextBlock()
+		d.nextBlock(true)
 	}
 	err := d.current.err
 	if err != nil {
@@ -329,7 +331,10 @@ func (d *Decoder) DecodeAll(input, dst []byte) ([]byte, error) {
 
 // nextBlock returns the next block.
 // If an error occurs d.err will be set.
-func (d *Decoder) nextBlock() {
+// Optionally the function can block for new output.
+// If non-blocking mode is used the returned boolean will be false
+// if no data was available without blocking.
+func (d *Decoder) nextBlock(blocking bool) (ok bool) {
 	if d.current.d != nil {
 		if debug {
 			printf("re-adding current decoder %p", d.current.d)
@@ -339,12 +344,22 @@ func (d *Decoder) nextBlock() {
 	}
 	if d.current.err != nil {
 		// Keep error state.
-		return
+		return blocking
+	}
+
+	if blocking {
+		d.current.decodeOutput = <-d.current.output
+	} else {
+		select {
+		case d.current.decodeOutput = <-d.current.output:
+		default:
+			return false
+		}
 	}
-	d.current.decodeOutput = <-d.current.output
 	if debug {
 		println("got", len(d.current.b), "bytes, error:", d.current.err)
 	}
+	return true
 }
 
 // Close will release all resources.

zstd/decoder_test.go

@@ -9,6 +9,7 @@ import (
 	"bytes"
 	"encoding/binary"
 	"encoding/hex"
+	"fmt"
 	"io"
 	"io/ioutil"
 	"log"
@@ -263,6 +264,83 @@ func TestNewDecoderSmallFile(t *testing.T) {
 	t.Logf("Decoded %d bytes with %f.2 MB/s", n, mbpersec)
 }
 
+type readAndBlock struct {
+	buf     []byte
+	unblock chan struct{}
+}
+
+func (r *readAndBlock) Read(p []byte) (int, error) {
+	n := copy(p, r.buf)
+	if n == 0 {
+		<-r.unblock
+		return 0, io.EOF
+	}
+	r.buf = r.buf[n:]
+	return n, nil
+}
+
+func TestNewDecoderFlushed(t *testing.T) {
+	if testing.Short() {
+		t.SkipNow()
+	}
+	file := "testdata/z000028.zst"
+	payload, err := ioutil.ReadFile(file)
+	if err != nil {
+		t.Fatal(err)
+	}
+	payload = append(payload, payload...) //2x
+	payload = append(payload, payload...) //4x
+	payload = append(payload, payload...) //8x
+	rng := rand.New(rand.NewSource(0x1337))
+	runs := 100
+	if testing.Short() {
+		runs = 5
+	}
+	enc, err := NewWriter(nil, WithWindowSize(128<<10))
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer enc.Close()
+	for i := 0; i < runs; i++ {
+		wantSize := rng.Intn(len(payload)-1) + 1
+		t.Run(fmt.Sprint("size-", wantSize), func(t *testing.T) {
+			var encoded bytes.Buffer
+			enc.Reset(&encoded)
+			_, err := enc.Write(payload[:wantSize])
+			if err != nil {
+				t.Fatal(err)
+			}
+			err = enc.Flush()
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			// We must be able to read back up until the flush...
+			r := readAndBlock{
+				buf:     encoded.Bytes(),
+				unblock: make(chan struct{}),
+			}
+			defer timeout(5 * time.Second)()
+			dec, err := NewReader(&r)
+			if err != nil {
+				t.Fatal(err)
+			}
+			defer dec.Close()
+			defer close(r.unblock)
+			readBack := 0
+			dst := make([]byte, 1024)
+			for readBack < wantSize {
+				// Read until we have enough.
+				n, err := dec.Read(dst)
+				if err != nil {
+					t.Fatal(err)
+				}
+				readBack += n
+			}
+		})
+	}
+}
+
 func TestDecoderRegression(t *testing.T) {
 	defer timeout(60 * time.Second)()
 	data, err := ioutil.ReadFile("testdata/regression.zip")