zstd: Fix WriteTo error forwarding (#411)

Fix error forwarding when writing the final block.

Fixes #408

Author: klauspost

zstd/decoder.go

@@ -260,9 +260,10 @@ func (d *Decoder) WriteTo(w io.Writer) (int64, error) {
 		if len(d.current.b) > 0 {
 			n2, err2 := w.Write(d.current.b)
 			n += int64(n2)
-			if err2 != nil && d.current.err == nil {
+			if err2 != nil && (d.current.err == nil || d.current.err == io.EOF) {
 				d.current.err = err2
-				break
+			} else if n2 != len(d.current.b) {
+				d.current.err = io.ErrShortWrite
 			}
 		}
 		if d.current.err != nil {

zstd/decoder_test.go

@@ -165,6 +165,40 @@ func TestErrorReader(t *testing.T) {
 	}
 }
 
+type failingWriter struct {
+	err error
+}
+
+func (f failingWriter) Write(_ []byte) (n int, err error) {
+	return 0, f.err
+}
+
+func TestErrorWriter(t *testing.T) {
+	input := make([]byte, 100)
+	cmp := bytes.Buffer{}
+	w, err := NewWriter(&cmp)
+	if err != nil {
+		t.Fatal(err)
+	}
+	_, _ = rand.Read(input)
+	_, err = w.Write(input)
+	if err != nil {
+		t.Fatal(err)
+	}
+	err = w.Close()
+	if err != nil {
+		t.Fatal(err)
+	}
+	wantErr := fmt.Errorf("i'm a failure")
+	zr, err := NewReader(&cmp)
+	defer zr.Close()
+	out := failingWriter{err: wantErr}
+	_, err = zr.WriteTo(out)
+	if !errors.Is(err, wantErr) {
+		t.Errorf("error: wanted: %v, got: %v", wantErr, err)
+	}
+}
+
 func TestNewDecoder(t *testing.T) {
 	defer timeout(60 * time.Second)()
 	testDecoderFile(t, "testdata/decoder.zip")