github.com/gogf/gf/v2@v2.7.4/net/ghttp/internal/response/response_buffer_writer.go (about)

     1  // Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
     2  //
     3  // This Source Code Form is subject to the terms of the MIT License.
     4  // If a copy of the MIT was not distributed with this file,
     5  // You can obtain one at https://github.com/gogf/gf.
     6  //
     7  
     8  package response
     9  
    10  import (
    11  	"bytes"
    12  	"net/http"
    13  )
    14  
    15  // BufferWriter is the custom writer for http response with buffer.
    16  type BufferWriter struct {
    17  	*Writer               // The underlying BufferWriter.
    18  	Status  int           // HTTP status.
    19  	buffer  *bytes.Buffer // The output buffer.
    20  }
    21  
    22  func NewBufferWriter(writer http.ResponseWriter) *BufferWriter {
    23  	return &BufferWriter{
    24  		Writer: NewWriter(writer),
    25  		buffer: bytes.NewBuffer(nil),
    26  	}
    27  }
    28  
    29  // RawWriter returns the underlying BufferWriter.
    30  func (w *BufferWriter) RawWriter() http.ResponseWriter {
    31  	return w.Writer
    32  }
    33  
    34  // Write implements the interface function of http.BufferWriter.Write.
    35  func (w *BufferWriter) Write(data []byte) (int, error) {
    36  	return w.buffer.Write(data)
    37  }
    38  
    39  // WriteString writes string content to internal buffer.
    40  func (w *BufferWriter) WriteString(data string) (int, error) {
    41  	return w.buffer.WriteString(data)
    42  }
    43  
    44  // Buffer returns the buffered content as []byte.
    45  func (w *BufferWriter) Buffer() []byte {
    46  	return w.buffer.Bytes()
    47  }
    48  
    49  // BufferString returns the buffered content as string.
    50  func (w *BufferWriter) BufferString() string {
    51  	return w.buffer.String()
    52  }
    53  
    54  // BufferLength returns the length of the buffered content.
    55  func (w *BufferWriter) BufferLength() int {
    56  	return w.buffer.Len()
    57  }
    58  
    59  // SetBuffer overwrites the buffer with `data`.
    60  func (w *BufferWriter) SetBuffer(data []byte) {
    61  	w.buffer.Reset()
    62  	w.buffer.Write(data)
    63  }
    64  
    65  // ClearBuffer clears the response buffer.
    66  func (w *BufferWriter) ClearBuffer() {
    67  	w.buffer.Reset()
    68  }
    69  
    70  // WriteHeader implements the interface of http.BufferWriter.WriteHeader.
    71  func (w *BufferWriter) WriteHeader(status int) {
    72  	w.Status = status
    73  }
    74  
    75  // Flush outputs the buffer to clients and clears the buffer.
    76  func (w *BufferWriter) Flush() {
    77  	if w.Writer.IsHijacked() {
    78  		return
    79  	}
    80  
    81  	if w.Status != 0 && !w.Writer.IsHeaderWrote() {
    82  		w.Writer.WriteHeader(w.Status)
    83  	}
    84  	// Default status text output.
    85  	if w.Status != http.StatusOK && w.buffer.Len() == 0 {
    86  		w.buffer.WriteString(http.StatusText(w.Status))
    87  	}
    88  	if w.buffer.Len() > 0 {
    89  		_, _ = w.Writer.Write(w.buffer.Bytes())
    90  		w.buffer.Reset()
    91  		if flusher, ok := w.RawWriter().(http.Flusher); ok {
    92  			flusher.Flush()
    93  		}
    94  	}
    95  }