github.com/xmidt-org/webpa-common@v1.11.9/xhttp/bufferedWriter.go (about)

     1  package xhttp
     2  
     3  import (
     4  	"bytes"
     5  	"errors"
     6  	"fmt"
     7  	"net/http"
     8  	"strconv"
     9  	"sync/atomic"
    10  )
    11  
    12  var ErrBufferedWriterClosed = errors.New("BufferedWriter has been closed")
    13  
    14  // BufferedWriter is a closeable http.ResponseWriter that holds all written response information in memory.
    15  // The zero value of this type is a fully usable writer.
    16  //
    17  // The http.ResponseWriter methods of this type are not safe for concurrent execution.  However, it is safe
    18  // to invoke Close concurrently with the other methods.
    19  //
    20  // Once closed, future Write and WriteTo calls will return errors.  This type is ideal for http.Handler code
    21  // that should be Buffered and optionally discarded depending on other logic.
    22  type BufferedWriter struct {
    23  	closed      uint32
    24  	wroteHeader bool
    25  	code        int
    26  	header      http.Header
    27  	buffer      bytes.Buffer
    28  }
    29  
    30  // Close closes this writer.  Once closed, this writer will reject writes with an error.
    31  // This method is idempotent, and will return an error if called more than once on a given writer instance.
    32  func (bw *BufferedWriter) Close() error {
    33  	if atomic.CompareAndSwapUint32(&bw.closed, 0, 1) {
    34  		return nil
    35  	}
    36  
    37  	return ErrBufferedWriterClosed
    38  }
    39  
    40  // Header returns the HTTP header to write to the response.  This method is unaffected by the close state.
    41  func (bw *BufferedWriter) Header() http.Header {
    42  	if bw.header == nil {
    43  		bw.header = make(http.Header)
    44  	}
    45  
    46  	return bw.header
    47  }
    48  
    49  // Write buffers content for later writing to a real http.ResponseWriter.  If this writer is closed,
    50  // this method returns a count of 0 with an error.
    51  func (bw *BufferedWriter) Write(p []byte) (int, error) {
    52  	if atomic.LoadUint32(&bw.closed) == 1 {
    53  		return 0, ErrBufferedWriterClosed
    54  	}
    55  
    56  	if !bw.wroteHeader {
    57  		bw.writeHeader(http.StatusOK)
    58  	}
    59  
    60  	return bw.buffer.Write(p)
    61  }
    62  
    63  // WriteHeader sets a status code and in most other ways behaves as the standard net/http ResponseWriter.
    64  // This method is idempotent.  Only the first invocation will have an effect.  If this writer is closed, this
    65  // method has no effect.
    66  func (bw *BufferedWriter) WriteHeader(code int) {
    67  	// mimic the current behavior of the stdlib
    68  	if code < 100 || code > 999 {
    69  		panic(fmt.Sprintf("Invalid WriteHeader code %v", code))
    70  	}
    71  
    72  	if atomic.LoadUint32(&bw.closed) == 1 || bw.wroteHeader {
    73  		return
    74  	}
    75  
    76  	bw.writeHeader(code)
    77  }
    78  
    79  func (bw *BufferedWriter) writeHeader(code int) {
    80  	bw.wroteHeader = true
    81  	bw.code = code
    82  }
    83  
    84  // WriteTo transfers this writer's state to the given ResponseWriter.  This method will only take effect once.
    85  // Once this method has been invoked successfully, this writer instance is closed and will reject future writes
    86  // (including this method).
    87  func (bw *BufferedWriter) WriteTo(response http.ResponseWriter) (int, error) {
    88  	if atomic.CompareAndSwapUint32(&bw.closed, 0, 1) {
    89  		destination := response.Header()
    90  		for k, v := range bw.header {
    91  			destination[http.CanonicalHeaderKey(k)] = v
    92  		}
    93  
    94  		contentLength := bw.buffer.Len()
    95  		if contentLength > 0 {
    96  			destination.Set("Content-Length", strconv.Itoa(contentLength))
    97  		}
    98  
    99  		code := bw.code
   100  		if code < 100 {
   101  			code = http.StatusOK
   102  		}
   103  
   104  		response.WriteHeader(code)
   105  		c, err := bw.buffer.WriteTo(response)
   106  		return int(c), err
   107  	}
   108  
   109  	return 0, ErrBufferedWriterClosed
   110  }