github.com/wangyougui/gf/v2@v2.6.5/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/wangyougui/gf.
     6  //
     7  
     8  package ghttp
     9  
    10  import (
    11  	"bufio"
    12  	"bytes"
    13  	"net"
    14  	"net/http"
    15  
    16  	"github.com/wangyougui/gf/v2/net/ghttp/internal/response"
    17  )
    18  
    19  // ResponseWriter is the custom writer for http response.
    20  type ResponseWriter struct {
    21  	Status int              // HTTP status.
    22  	writer *response.Writer // The underlying ResponseWriter.
    23  	buffer *bytes.Buffer    // The output buffer.
    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  	return w.writer.Hijack()
    50  }
    51  
    52  // Flush outputs the buffer to clients and clears the buffer.
    53  func (w *ResponseWriter) Flush() {
    54  	if w.writer.IsHijacked() {
    55  		return
    56  	}
    57  
    58  	if w.Status != 0 && !w.writer.IsHeaderWrote() {
    59  		w.writer.WriteHeader(w.Status)
    60  	}
    61  	// Default status text output.
    62  	if w.Status != http.StatusOK && w.buffer.Len() == 0 {
    63  		w.buffer.WriteString(http.StatusText(w.Status))
    64  	}
    65  	if w.buffer.Len() > 0 {
    66  		_, _ = w.writer.Write(w.buffer.Bytes())
    67  		w.buffer.Reset()
    68  		if flusher, ok := w.RawWriter().(http.Flusher); ok {
    69  			flusher.Flush()
    70  		}
    71  	}
    72  }