github.com/gogf/gf@v1.16.9/net/ghttp/ghttp_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  
     8  package ghttp
     9  
    10  import (
    11  	"bufio"
    12  	"bytes"
    13  	"net"
    14  	"net/http"
    15  )
    16  
    17  // ResponseWriter is the custom writer for http response.
    18  type ResponseWriter struct {
    19  	Status      int                 // HTTP status.
    20  	writer      http.ResponseWriter // The underlying ResponseWriter.
    21  	buffer      *bytes.Buffer       // The output buffer.
    22  	hijacked    bool                // Mark this request is hijacked or not.
    23  	wroteHeader bool                // Is header wrote or not, avoiding error: superfluous/multiple response.WriteHeader call.
    24  }
    25  
    26  // RawWriter returns the underlying ResponseWriter.
    27  func (w *ResponseWriter) RawWriter() http.ResponseWriter {
    28  	return w.writer
    29  }
    30  
    31  // Header implements the interface function of http.ResponseWriter.Header.
    32  func (w *ResponseWriter) Header() http.Header {
    33  	return w.writer.Header()
    34  }
    35  
    36  // Write implements the interface function of http.ResponseWriter.Write.
    37  func (w *ResponseWriter) Write(data []byte) (int, error) {
    38  	w.buffer.Write(data)
    39  	return len(data), nil
    40  }
    41  
    42  // WriteHeader implements the interface of http.ResponseWriter.WriteHeader.
    43  func (w *ResponseWriter) WriteHeader(status int) {
    44  	w.Status = status
    45  }
    46  
    47  // Hijack implements the interface function of http.Hijacker.Hijack.
    48  func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
    49  	w.hijacked = true
    50  	return w.writer.(http.Hijacker).Hijack()
    51  }
    52  
    53  // OutputBuffer outputs the buffer to client and clears the buffer.
    54  func (w *ResponseWriter) Flush() {
    55  	if w.hijacked {
    56  		return
    57  	}
    58  	if w.Status != 0 && !w.wroteHeader {
    59  		w.wroteHeader = true
    60  		w.writer.WriteHeader(w.Status)
    61  	}
    62  	// Default status text output.
    63  	if w.Status != http.StatusOK && w.buffer.Len() == 0 {
    64  		w.buffer.WriteString(http.StatusText(w.Status))
    65  	}
    66  	if w.buffer.Len() > 0 {
    67  		w.writer.Write(w.buffer.Bytes())
    68  		w.buffer.Reset()
    69  	}
    70  }