Merge pull request #34 from muir/deferred-body

Add Body() to deferred writer

Author: muir

.github/workflows/codecov.yml

@@ -1,6 +1,6 @@
 name: Test and coverage
 
-on: [push, pull_request]
+on: [push]
 
 permissions:  # added using https://github.com/step-security/secure-workflows
   contents: read

.github/workflows/go.yml

@@ -7,7 +7,7 @@ jobs:
   test:
     strategy:
       matrix:
-        go-version: [1.16.x, 1.17.x, 1.18.x]
+        go-version: [1.16.x, 1.17.x, 1.18.x, 1.19.x]
         os: [ubuntu-latest, macos-latest, windows-latest]
     runs-on: ${{ matrix.os }}
     steps:

.github/workflows/golangci-lint.yml

@@ -1,5 +1,5 @@
 name: golangci-lint
-on: [push ] 
+on: [push] 
 permissions:  # added using https://github.com/step-security/secure-workflows
   contents: read
 

deferred.go

@@ -17,6 +17,7 @@ type DeferredWriter struct {
 	buffer      []byte
 	status      int
 	resetHeader http.Header
+	flushed     bool
 }
 
 // NewDeferredWriter returns a DeferredWriter based on a
@@ -85,7 +86,13 @@ func (w *DeferredWriter) PreserveHeader() {
 // switches to passthrough mode: all future calls to Write(),
 // Header(), etc are passed through to the http.ResponseWriter that
 // was used to initialize the DeferredWrited.
+//
+// Any writes made before the call to UnderlyingWriter are discarded.
+// Call Flush() first to preserve writes.
 func (w *DeferredWriter) UnderlyingWriter() http.ResponseWriter {
+	if w.passthrough {
+		return w.base
+	}
 	w.passthrough = true
 	h := w.base.Header()
 	for k := range h {
@@ -112,7 +119,7 @@ func (w *DeferredWriter) Flush() error {
 	if w.passthrough {
 		return errors.New("Attempt flush deferred writer that is not deferred")
 	}
-	w.passthrough = true
+	w.flushed = true
 	base := w.UnderlyingWriter()
 	if w.status != 0 {
 		base.WriteHeader(w.status)
@@ -147,3 +154,14 @@ func (w *DeferredWriter) FlushIfNotFlushed() error {
 func (w *DeferredWriter) Done() bool {
 	return w.passthrough
 }
+
+// Body returns the internal buffer used by DeferredWriter.  Do not modify it.
+// It also returns the status code (if set).
+// If UnderlyingWriter() has been called, then Body() will return an error since
+// the underlying buffer does not represent what has been written.
+func (w *DeferredWriter) Body() ([]byte, int, error) {
+	if w.passthrough && !w.flushed {
+		return nil, 0, errors.New("unable to provide body because DeferredWriter is operating in passthrough mode")
+	}
+	return w.buffer, w.status, nil
+}

deferred_test.go

@@ -42,8 +42,24 @@ func (w *testResponseWriter) Write(b []byte) (int, error) {
 
 func TestUnderlyingWriter(t *testing.T) {
 	tw := &testResponseWriter{header: make(http.Header)}
+	tw.header.Set("X", "Y")
 	w, _ := nvelope.NewDeferredWriter(tw)
+	tw.header.Set("Foo", "bar")
+	tw.header.Set("A", "B")
+	w.Header().Set("Baz", "bap")
+	w.Header().Set("A", "C")
+	_, _ = w.Write([]byte("howdy"))
+	_, _, err := w.Body()
+	assert.NoError(t, err, "body before Underlying")
+	assert.Equal(t, tw, w.UnderlyingWriter())
+	_, _, err = w.Body()
+	assert.Error(t, err, "body after Underlying")
 	assert.Equal(t, tw, w.UnderlyingWriter())
+	assert.Equal(t, []byte(nil), tw.buffer, "underlying buffer after two calls")
+	assert.Empty(t, tw.header["Foo"], "Foo")
+	assert.Equal(t, tw.header.Get("Baz"), "bap", "Baz")
+	assert.Equal(t, tw.header.Get("A"), "C", "A")
+	assert.Equal(t, tw.header.Get("X"), "Y", "X")
 }
 
 func TestFlush(t *testing.T) {
@@ -59,6 +75,10 @@ func TestFlush(t *testing.T) {
 	assert.Equal(t, "", tw.Header().Get("c"), "original header untouched with existing key")
 	assert.Equal(t, "d", w.Header().Get("c"), "new header override works though")
 	w.WriteHeader(303)
+	body, code, err := w.Body()
+	assert.NoError(t, err, "body")
+	assert.Equal(t, 303, code, "body code")
+	assert.Equal(t, []byte("howdy"), body, code, "body")
 	assert.Equal(t, 0, tw.code, "code not written before flush")
 	assert.False(t, w.Done(), "done before flush")
 	require.NoError(t, w.Flush(), "flush")
@@ -67,6 +87,10 @@ func TestFlush(t *testing.T) {
 	assert.Equal(t, 303, tw.code, "code written after flush")
 	assert.Equal(t, "d", tw.Header().Get("c"), "new header written - c")
 	assert.Equal(t, "d", tw.Header().Get("a"), "new header written - a")
+	body, code, err = w.Body()
+	assert.NoError(t, err, "body")
+	assert.Equal(t, 303, code, "body code")
+	assert.Equal(t, []byte("howdy"), body, code, "body")
 }
 
 func TestReset(t *testing.T) {