153 lines
3.1 KiB
Go
153 lines
3.1 KiB
Go
package service
|
|
|
|
import (
|
|
"apigo.cc/go/cast"
|
|
"io"
|
|
"net/http"
|
|
"os"
|
|
)
|
|
|
|
// Response 封装 http.ResponseWriter
|
|
type Response struct {
|
|
Id string
|
|
Writer http.ResponseWriter
|
|
status int
|
|
outLen int
|
|
changed bool
|
|
headerWritten bool
|
|
dontLog200 bool
|
|
dontLogArgs []string
|
|
ProxyHeader *http.Header
|
|
}
|
|
|
|
// NewResponse 创建 Response 包装
|
|
func NewResponse(writer http.ResponseWriter) *Response {
|
|
return &Response{
|
|
Writer: writer,
|
|
status: http.StatusOK,
|
|
}
|
|
}
|
|
|
|
// Header 获取响应头部
|
|
func (r *Response) Header() http.Header {
|
|
r.changed = true
|
|
if r.ProxyHeader != nil {
|
|
return *r.ProxyHeader
|
|
}
|
|
return r.Writer.Header()
|
|
}
|
|
|
|
// Write 写入响应内容
|
|
func (r *Response) Write(bytes []byte) (int, error) {
|
|
r.checkWriteHeader()
|
|
r.changed = true
|
|
r.outLen += len(bytes)
|
|
if r.ProxyHeader != nil {
|
|
r.copyProxyHeader()
|
|
}
|
|
return r.Writer.Write(bytes)
|
|
}
|
|
|
|
// WriteString 写入字符串响应
|
|
func (r *Response) WriteString(s string) (int, error) {
|
|
return r.Write([]byte(s))
|
|
}
|
|
|
|
// WriteHeader 设置响应状态码
|
|
func (r *Response) WriteHeader(code int) {
|
|
r.changed = true
|
|
r.status = code
|
|
if r.ProxyHeader != nil && (r.status == http.StatusBadGateway || r.status == http.StatusServiceUnavailable || r.status == http.StatusGatewayTimeout) {
|
|
return
|
|
}
|
|
if r.ProxyHeader != nil {
|
|
r.copyProxyHeader()
|
|
}
|
|
}
|
|
|
|
func (r *Response) checkWriteHeader() {
|
|
if !r.headerWritten {
|
|
r.headerWritten = true
|
|
if r.status != http.StatusOK {
|
|
r.Writer.WriteHeader(r.status)
|
|
}
|
|
}
|
|
}
|
|
|
|
func (r *Response) copyProxyHeader() {
|
|
src := *r.ProxyHeader
|
|
dst := r.Writer.Header()
|
|
for k, vv := range src {
|
|
for _, v := range vv {
|
|
dst.Add(k, v)
|
|
}
|
|
}
|
|
r.ProxyHeader = nil
|
|
}
|
|
|
|
// Flush 刷新响应缓冲区
|
|
func (r *Response) Flush() {
|
|
if flusher, ok := r.Writer.(http.Flusher); ok {
|
|
flusher.Flush()
|
|
}
|
|
}
|
|
|
|
// GetStatusCode 获取当前状态码
|
|
func (r *Response) GetStatusCode() int {
|
|
return r.status
|
|
}
|
|
|
|
// DontLog200 标记不记录 200 状态码的日志
|
|
func (r *Response) DontLog200() {
|
|
r.dontLog200 = true
|
|
}
|
|
|
|
// Location 设置重定向地址
|
|
func (r *Response) Location(location string) {
|
|
r.WriteHeader(http.StatusFound)
|
|
r.Header().Set("Location", location)
|
|
}
|
|
|
|
// SendFile 发送文件
|
|
func (r *Response) SendFile(contentType, filename string) {
|
|
r.Header().Set("Content-Type", contentType)
|
|
// TODO: Integrate memory file support if needed
|
|
if fd, err := os.Open(filename); err == nil {
|
|
defer fd.Close()
|
|
_, _ = io.Copy(r, fd)
|
|
}
|
|
}
|
|
|
|
// DownloadFile 下载文件
|
|
func (r *Response) DownloadFile(contentType, filename string, data any) {
|
|
if contentType == "" {
|
|
contentType = "application/octet-stream"
|
|
}
|
|
r.Header().Set("Content-Type", contentType)
|
|
|
|
if filename != "" {
|
|
r.Header().Set("Content-Disposition", "attachment; filename="+filename)
|
|
}
|
|
|
|
var outBytes []byte
|
|
var reader io.Reader
|
|
|
|
switch v := data.(type) {
|
|
case []byte:
|
|
outBytes = v
|
|
case string:
|
|
outBytes = []byte(v)
|
|
case io.Reader:
|
|
reader = v
|
|
default:
|
|
outBytes, _ = cast.ToJSONBytes(data)
|
|
}
|
|
|
|
if outBytes != nil {
|
|
r.Header().Set("Content-Length", cast.String(len(outBytes)))
|
|
_, _ = r.Write(outBytes)
|
|
} else if reader != nil {
|
|
_, _ = io.Copy(r, reader)
|
|
}
|
|
}
|