diff --git a/zstd_stream.go b/zstd_stream.go index 167c764..ccf5b38 100644 --- a/zstd_stream.go +++ b/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 diff --git a/zstd_stream_test.go b/zstd_stream_test.go index 049dd61..876aaa1 100644 --- a/zstd_stream_test.go +++ b/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}