Return an error instead of crashing when a Writer is used after Close

zstd_stream.go

@@ -72,6 +72,7 @@ import (
 
 var errShortRead = errors.New("short read")
 var errReaderClosed = errors.New("Reader is closed")
+var errWriterClosed = errors.New("Writer is closed")
 var ErrNoParallelSupport = errors.New("No parallel support")
 
 // Writer is an io.WriteCloser that zstd-compresses its input.
@@ -157,6 +158,9 @@ func (w *Writer) Write(p []byte) (int, error) {
 	if w.firstError != nil {
 		return 0, w.firstError
 	}
+	if w.ctx == nil {
+		return 0, errWriterClosed
+	}
 	if len(p) == 0 {
 		return 0, nil
 	}
@@ -220,6 +224,9 @@ func (w *Writer) Flush() error {
 	if w.firstError != nil {
 		return w.firstError
 	}
+	if w.ctx == nil {
+		return errWriterClosed
+	}
 
 	ret := 1 // So we loop at least once
 	for ret > 0 {
@@ -314,6 +321,9 @@ func (w *Writer) SetNbWorkers(n int) error {
 	if w.firstError != nil {
 		return w.firstError
 	}
+	if w.ctx == nil {
+		return errWriterClosed
+	}
 	if err := getError(int(C.ZSTD_CCtx_setParameter(w.ctx, C.ZSTD_c_nbWorkers, C.int(n)))); err != nil {
 		w.firstError = err
 		// First error case, a shared libary is used, and the library was compiled without parallel support

zstd_stream_test.go

@@ -535,6 +535,48 @@ func TestWriterDoubleCloseIsSafe(t *testing.T) {
 	}
 }
 
+func TestWriterUseAfterCloseReturnsError(t *testing.T) {
+	var buf bytes.Buffer
+	w := NewWriter(&buf)
+
+	_, err := w.Write([]byte("hello world"))
+	if err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+	failOnError(t, "Close failed", w.Close())
+
+	// The C context is freed by Close, so these must not reach the C library
+	if _, err := w.Write([]byte("hello again")); err != errWriterClosed {
+		t.Fatalf("Write() after Close() should return errWriterClosed, got: %v", err)
+	}
+	if err := w.Flush(); err != errWriterClosed {
+		t.Fatalf("Flush() after Close() should return errWriterClosed, got: %v", err)
+	}
+	if err := w.SetNbWorkers(2); err != errWriterClosed {
+		t.Fatalf("SetNbWorkers() after Close() should return errWriterClosed, got: %v", err)
+	}
+}
+
+func TestWriterWriteAfterFailedCloseReturnsError(t *testing.T) {
+	var bw bytes.Buffer
+	cw := &closeableWriter{w: &bw}
+	w := NewWriter(cw)
+
+	_, err := w.Write([]byte("hello world"))
+	if err != nil {
+		t.Fatalf("Write failed: %v", err)
+	}
+
+	cw.Close() // make the underlying writer fail
+	if err := w.Close(); err == nil {
+		t.Fatal("Close() should have returned an error from the underlying writer")
+	}
+
+	if _, err := w.Write([]byte("hello again")); err != errWriterClosed {
+		t.Fatalf("Write() after a failed Close() should return errWriterClosed, got: %v", err)
+	}
+}
+
 func TestWriterCloseNilsContextOnUnderlyingWriterError(t *testing.T) {
 	var bw bytes.Buffer
 	cw := &closeableWriter{w: &bw}