github.com/gogf/gf/v2@v2.7.4/net/ghttp/internal/response/response_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  package response
     8  
     9  import (
    10  	"bufio"
    11  	"net"
    12  	"net/http"
    13  )
    14  
    15  // Writer wraps http.ResponseWriter for extra features.
    16  type Writer struct {
    17  	http.ResponseWriter       // The underlying ResponseWriter.
    18  	hijacked            bool  // Mark this request is hijacked or not.
    19  	wroteHeader         bool  // Is header wrote or not, avoiding error: superfluous/multiple response.WriteHeader call.
    20  	bytesWritten        int64 // Bytes written to response.
    21  }
    22  
    23  // NewWriter creates and returns a new Writer.
    24  func NewWriter(writer http.ResponseWriter) *Writer {
    25  	return &Writer{
    26  		ResponseWriter: writer,
    27  	}
    28  }
    29  
    30  // WriteHeader implements the interface of http.ResponseWriter.WriteHeader.
    31  // Note that the underlying `WriteHeader` can only be called once in a http response.
    32  func (w *Writer) WriteHeader(status int) {
    33  	if w.wroteHeader {
    34  		return
    35  	}
    36  	w.ResponseWriter.WriteHeader(status)
    37  	w.wroteHeader = true
    38  }
    39  
    40  // BytesWritten returns the length that was written to response.
    41  func (w *Writer) BytesWritten() int64 {
    42  	return w.bytesWritten
    43  }
    44  
    45  // Write implements the interface function of http.ResponseWriter.Write.
    46  func (w *Writer) Write(data []byte) (int, error) {
    47  	n, err := w.ResponseWriter.Write(data)
    48  	w.bytesWritten += int64(n)
    49  	w.wroteHeader = true
    50  	return n, err
    51  }
    52  
    53  // Hijack implements the interface function of http.Hijacker.Hijack.
    54  func (w *Writer) Hijack() (conn net.Conn, writer *bufio.ReadWriter, err error) {
    55  	conn, writer, err = w.ResponseWriter.(http.Hijacker).Hijack()
    56  	w.hijacked = true
    57  	return
    58  }
    59  
    60  // IsHeaderWrote returns if the header status is written.
    61  func (w *Writer) IsHeaderWrote() bool {
    62  	return w.wroteHeader
    63  }
    64  
    65  // IsHijacked returns if the connection is hijacked.
    66  func (w *Writer) IsHijacked() bool {
    67  	return w.hijacked
    68  }
    69  
    70  // Flush sends any buffered data to the client.
    71  func (w *Writer) Flush() {
    72  	flusher, ok := w.ResponseWriter.(http.Flusher)
    73  	if ok {
    74  		flusher.Flush()
    75  		w.wroteHeader = true
    76  	}
    77  }