net/ghttp/ghttp_response_writer.go - master - githubmirror/GoFrame

net/ghttp/ghttp_response_writer.go
82 строки2 KB
John Guo
fix issue #2334 when accessing static files with cache time (#2366)
22 дек 2022, 05:25
22 дек 2022, 05:25a853984
// Copyright GoFrame Author(https://goframe.org). All Rights Reserved.
//
// This Source Code Form is subject to the terms of the MIT License.
// If a copy of the MIT was not distributed with this file,
// You can obtain one at https://github.com/gogf/gf.
//

package ghttp

import (
	"bufio"
	"bytes"
	"net"
	"net/http"
)

// ResponseWriter is the custom writer for http response.
type ResponseWriter struct {
	Status      int                 // HTTP status.
	writer      http.ResponseWriter // The underlying ResponseWriter.
	buffer      *bytes.Buffer       // The output buffer.
	hijacked    bool                // Mark this request is hijacked or not.
	wroteHeader bool                // Is header wrote or not, avoiding error: superfluous/multiple response.WriteHeader call.
}

// RawWriter returns the underlying ResponseWriter.
func (w *ResponseWriter) RawWriter() http.ResponseWriter {
	return w.writer
}

// Header implements the interface function of http.ResponseWriter.Header.
func (w *ResponseWriter) Header() http.Header {
	return w.writer.Header()
}

// Write implements the interface function of http.ResponseWriter.Write.
func (w *ResponseWriter) Write(data []byte) (int, error) {
	w.buffer.Write(data)
	return len(data), nil
}

// WriteHeader implements the interface of http.ResponseWriter.WriteHeader.
func (w *ResponseWriter) WriteHeader(status int) {
	w.Status = status
}

// Hijack implements the interface function of http.Hijacker.Hijack.
func (w *ResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	w.hijacked = true
	return w.writer.(http.Hijacker).Hijack()
}

// Flush outputs the buffer to clients and clears the buffer.
func (w *ResponseWriter) Flush() {
	if w.hijacked {
		return
	}

	if w.Status != 0 && !w.isHeaderWritten() {
		w.wroteHeader = true
		w.writer.WriteHeader(w.Status)
	}
	// Default status text output.
	if w.Status != http.StatusOK && w.buffer.Len() == 0 {
		w.buffer.WriteString(http.StatusText(w.Status))
	}
	if w.buffer.Len() > 0 {
		_, _ = w.writer.Write(w.buffer.Bytes())
		w.buffer.Reset()
	}
}

// isHeaderWrote checks and returns whether the header is written.
func (w *ResponseWriter) isHeaderWritten() bool {
	if w.wroteHeader {
		return true
	}
	if _, ok := w.writer.Header()[responseHeaderContentLength]; ok {
		return true
	}
	return false
}