Merge pull request #4 from infinite-iroha/dev

0.0.4
This commit is contained in:
WJQSERVER 2025-05-30 21:43:57 +08:00 committed by GitHub
commit 7f32c15b4b
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
6 changed files with 589 additions and 13 deletions

69
ecw.go
View file

@ -1,6 +1,9 @@
package touka
import (
"bufio"
"errors"
"net"
"net/http"
"sync"
)
@ -44,9 +47,9 @@ func (ecw *errorCapturingResponseWriter) reset(w http.ResponseWriter, r *http.Re
// AcquireErrorCapturingResponseWriter 从对象池获取一个 errorCapturingResponseWriter 实例
// 必须在处理完成后调用 ReleaseErrorCapturingResponseWriter
func AcquireErrorCapturingResponseWriter(c *Context, eh ErrorHandler) *errorCapturingResponseWriter {
func AcquireErrorCapturingResponseWriter(c *Context) *errorCapturingResponseWriter {
ecw := errorResponseWriterPool.Get().(*errorCapturingResponseWriter)
ecw.reset(c.Writer, c.Request, c, eh) // 传入 Touka Context 的 Writer
ecw.reset(c.Writer, c.Request, c, c.engine.errorHandle.handler) // 传入 Touka Context 的 Writer
return ecw
}
@ -74,9 +77,9 @@ func (ecw *errorCapturingResponseWriter) WriteHeader(statusCode int) {
if ecw.responseStarted {
return // 响应已开始, 忽略后续的 WriteHeader 调用
}
ecw.statusCode = statusCode // 总是记录 FileServer 意图的状态码
ecw.statusCode = statusCode
if statusCode >= http.StatusBadRequest {
if ecw.Status() >= 400 {
ecw.capturedErrorSignal = true
// 是一个错误状态码 (>=400), 激活错误信号
// 不会将这个 WriteHeader 传递给原始的 w, 等待 processAfterFileServer 处理
@ -108,7 +111,7 @@ func (ecw *errorCapturingResponseWriter) Write(data []byte) (int, error) {
for k, v := range ecw.headerSnapshot {
ecw.w.Header()[k] = v // 直接赋值 []string, 保留所有值
}
ecw.w.WriteHeader(ecw.statusCode) // 发送实际的状态码 (可能是 200 或之前设置的 2xx)
ecw.w.WriteHeader(ecw.Status()) // 发送实际的状态码 (可能是 200 或之前设置的 2xx)
ecw.responseStarted = true
}
return ecw.w.Write(data) // 写入数据到原始 ResponseWriter
@ -133,7 +136,7 @@ func (ecw *errorCapturingResponseWriter) processAfterFileServer() {
ecw.ctx.Next()
} else {
// 调用用户自定义的 ErrorHandlerFunc, 由它负责完整的错误响应
ecw.errorHandlerFunc(ecw.ctx, ecw.statusCode)
ecw.errorHandlerFunc(ecw.ctx, ecw.Status())
ecw.ctx.Abort()
}
}
@ -141,3 +144,57 @@ func (ecw *errorCapturingResponseWriter) processAfterFileServer() {
// 如果 ecw.capturedErrorSignal && ecw.responseStarted, 表示在捕获错误信号之前,
// 成功路径的响应已经开始, 此时无法再进行错误处理覆盖
}
// Status 返回当前记录的状态码
func (ecw *errorCapturingResponseWriter) Status() int {
if ecw.statusCode == 0 && !ecw.responseStarted {
// 如果还没有显式设置状态码, 并且响应尚未开始,
// 则尝试从底层 ResponseWriter 获取状态码 (如果它实现了 Statuser)
if tw, ok := ecw.w.(ResponseWriter); ok {
return tw.Status()
}
// 否则, 默认返回 200 OK (Go HTTP server 的默认行为)
return http.StatusOK
}
return ecw.statusCode
}
// Size 返回已写入响应体的字节数
func (ecw *errorCapturingResponseWriter) Size() int {
// ecw 在捕获错误信号时会丢弃 FileServer 写入的数据, 所以 Size 应返回 0
if ecw.capturedErrorSignal {
return 0
}
// 否则, 尝试从底层 ResponseWriter 获取已写入的字节数
if tw, ok := ecw.w.(ResponseWriter); ok {
return tw.Size()
}
// 对于其他类型的 ResponseWriter, 无法可靠获取, 只能返回 0
return 0
}
// Written方式
func (ecw *errorCapturingResponseWriter) Written() bool {
// 如果响应已经通过这个包装器开始写入 (WriteHeader 或 Write 成功调用)
// 或者如果原始 ResponseWriter 已经标记为 Written (例如, 如果它是 touka.ResponseWriterImpl)
// 则认为响应已开始
if ecw.responseStarted {
return true
}
// 检查原始 ResponseWriter 是否已经写入
if tw, ok := ecw.w.(ResponseWriter); ok {
return tw.Written()
}
// 对于其他类型的 ResponseWriter, 无法可靠判断是否已写入, 只能依赖 responseStarted 标记
return false
}
// Hijack 实现 http.Hijacker 接口
// 它将 Hijack 调用委托给底层的 ResponseWriter
func (ecw *errorCapturingResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
hijacker, ok := ecw.w.(http.Hijacker)
if !ok {
return nil, nil, errors.New("the underlying ResponseWriter does not support the Hijacker interface")
}
return hijacker.Hijack()
}

View file

@ -2,6 +2,7 @@ package touka
import (
"context"
"log"
"reflect"
"runtime"
"strings"
@ -53,6 +54,7 @@ type Engine struct {
noRoute HandlerFunc
unMatchFS UnMatchFS // 未匹配下的处理
unMatchFileServer http.Handler // 处理handle
serverProtocols *http.Protocols //服务协议
Protocols ProtocolsConfig //协议版本配置
@ -73,6 +75,9 @@ func defaultErrorHandle(c *Context, code int) { // 检查客户端是否已断
return
default:
if c.Writer.Written() {
return
}
// 输出json 状态码与状态码对应描述
c.JSON(code, H{
"code": code,
@ -84,6 +89,22 @@ func defaultErrorHandle(c *Context, code int) { // 检查客户端是否已断
}
}
// 默认errorhandle包装 避免竞争意外问题, 保证稳定性
func defaultErrorWarp(handler ErrorHandler) ErrorHandler {
return func(c *Context, code int) {
select {
case <-c.Request.Context().Done():
return
default:
if c.Writer.Written() {
log.Printf("errpage: response already started for status %d, skipping error page rendering", code)
return
}
}
handler(c, code)
}
}
type UnMatchFS struct {
FSForUnmatched http.FileSystem
ServeUnmatchedAsFS bool
@ -146,7 +167,7 @@ func Default() *Engine {
// 设置自定义错误处理
func (engine *Engine) SetErrorHandler(handler ErrorHandler) {
engine.errorHandle.useDefault = false
engine.errorHandle.handler = handler
engine.errorHandle.handler = defaultErrorWarp(handler)
}
// 获取一个默认错误处理handle
@ -159,8 +180,10 @@ func (engine *Engine) SetUnMatchFS(fs http.FileSystem) {
if fs != nil {
engine.unMatchFS.FSForUnmatched = fs
engine.unMatchFS.ServeUnmatchedAsFS = true
engine.unMatchFileServer = http.FileServer(fs)
} else {
engine.unMatchFS.ServeUnmatchedAsFS = false
engine.unMatchFileServer = nil
}
}
@ -442,14 +465,20 @@ func (engine *Engine) handleRequest(c *Context) {
func unMatchFSHandle() HandlerFunc {
return func(c *Context) {
engine := c.engine
// 确保 engine.unMatchFileServer 存在
if !engine.unMatchFS.ServeUnmatchedAsFS || engine.unMatchFileServer == nil {
c.Next() // 如果未配置或 FileSystem 为 nil则继续处理链
return
}
if c.Request.Method == http.MethodGet || c.Request.Method == http.MethodHead {
// 使用 http.FileServer 处理未匹配的请求
fileServer := http.FileServer(engine.unMatchFS.FSForUnmatched)
//fileServer := http.FileServer(engine.unMatchFS.FSForUnmatched)
//ecw := newErrorCapturingResponseWriter(c, c.engine.errorHandle.handler)
ecw := AcquireErrorCapturingResponseWriter(c, c.engine.errorHandle.handler)
ecw := AcquireErrorCapturingResponseWriter(c)
defer ReleaseErrorCapturingResponseWriter(ecw)
fileServer.ServeHTTP(ecw, c.Request)
c.engine.unMatchFileServer.ServeHTTP(ecw, c.Request)
ecw.processAfterFileServer()
c.Abort()
return
} else {
if engine.noRoute == nil {
@ -726,7 +755,7 @@ func (engine *Engine) Static(relativePath, rootPath string) {
// 使用自定义的 ResponseWriter 包装器来捕获 FileServer 可能返回的错误状态码
// 这样我们可以在 FileServer 返回 404 或 403 时,使用 Engine 的 ErrorHandler 进行统一处理
ecw := AcquireErrorCapturingResponseWriter(c, c.engine.errorHandle.handler)
ecw := AcquireErrorCapturingResponseWriter(c)
defer ReleaseErrorCapturingResponseWriter(ecw)
//
@ -790,7 +819,7 @@ func (group *RouterGroup) Static(relativePath, rootPath string) {
// 使用自定义的 ResponseWriter 包装器来捕获 FileServer 可能返回的错误状态码
// 这样我们可以在 FileServer 返回 404 或 403 时,使用 Engine 的 ErrorHandler 进行统一处理
ecw := AcquireErrorCapturingResponseWriter(c, group.engine.errorHandle.handler)
ecw := AcquireErrorCapturingResponseWriter(c)
defer ReleaseErrorCapturingResponseWriter(ecw)
//

1
go.mod
View file

@ -6,6 +6,7 @@ require (
github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4
github.com/WJQSERVER-STUDIO/httpc v0.5.1
github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8
github.com/gorilla/websocket v1.5.3
)
require github.com/valyala/bytebufferpool v1.0.0 // indirect

2
go.sum
View file

@ -4,5 +4,7 @@ github.com/WJQSERVER-STUDIO/httpc v0.5.1 h1:+TKCPYBuj7PAHuiduGCGAqsHAa4QtsUfoVwR
github.com/WJQSERVER-STUDIO/httpc v0.5.1/go.mod h1:M7KNUZjjhCkzzcg9lBPs9YfkImI+7vqjAyjdA19+joE=
github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8 h1:o8UqXPI6SVwQt04RGsqKp3qqmbOfTNMqDrWsc4O47kk=
github.com/go-json-experiment/json v0.0.0-20250517221953-25912455fbc8/go.mod h1:TiCD2a1pcmjd7YnhGH0f/zKNcCD06B029pHhzV23c2M=
github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
github.com/valyala/bytebufferpool v1.0.0 h1:GqA5TC/0021Y/b9FG4Oi9Mr3q7XYx6KllzawFIhcdPw=
github.com/valyala/bytebufferpool v1.0.0/go.mod h1:6bBcMArwyJ5K/AmCkWv1jt77kVWyCJ6HpOuEn7z0Csc=

357
tgzip.go Normal file
View file

@ -0,0 +1,357 @@
package touka
import (
"bufio"
"compress/gzip"
"errors"
"io"
"net"
"net/http"
"strconv"
"strings"
"sync"
)
const (
headerAcceptEncoding = "Accept-Encoding" // 请求头部,客户端声明接受的编码
headerContentEncoding = "Content-Encoding" // 响应头部,服务器声明使用的编码
headerContentLength = "Content-Length" // 响应头部,内容长度
headerContentType = "Content-Type" // 响应头部,内容类型
headerVary = "Vary" // 响应头部,指示缓存行为
encodingGzip = "gzip" // Gzip 编码名称
)
var (
// 默认可压缩的 MIME 类型
defaultCompressibleTypes = []string{
"text/html", "text/css", "text/plain", "text/javascript",
"application/javascript", "application/x-javascript", "application/json",
"application/xml", "image/svg+xml",
}
)
// GzipOptions 用于配置 Gzip 中间件。
type GzipOptions struct {
// Level 设置 Gzip 压缩级别。
// 例如: gzip.DefaultCompression, gzip.BestSpeed, gzip.BestCompression。
Level int
// MinContentLength 是应用 Gzip 的最小内容长度。
// 如果响应的 Content-Length 小于此值,则不应用 Gzip。
// 默认为 0 (无最小长度限制)。
MinContentLength int64
// CompressibleTypes 是要压缩的 MIME 类型列表。
// 如果为空,将使用 defaultCompressibleTypes。
CompressibleTypes []string
// DecompressFn 是一个可选函数,用于解压缩请求体 (如果请求体是 gzipped)。
// 如果为 nil则禁用请求体解压缩。
// 注意: 本次实现主要关注响应压缩,请求解压可以作为扩展。
// DecompressFn func(c *Context)
}
// gzipResponseWriter 包装了 touka.ResponseWriter 以提供 Gzip 压缩功能。
type gzipResponseWriter struct {
ResponseWriter // 底层的 ResponseWriter (可能是 ecw 或 responseWriterImpl)
gzWriter *gzip.Writer // compress/gzip 的 writer
options *GzipOptions // Gzip 配置
wroteHeader bool // 标记 Header 是否已写入
doCompression bool // 标记是否执行压缩
statusCode int // 存储状态码,在实际写入底层 Writer 前使用
}
// --- 对象池 ---
var gzipResponseWriterPool = sync.Pool{
New: func() interface{} {
return &gzipResponseWriter{}
},
}
// gzip.Writer 实例的对象池。
// 注意: gzip.Writer.Reset() 不会改变压缩级别,所以对象池需要提供已正确初始化级别的 writer。
// 我们为每个可能的级别创建一个池。
var gzipWriterPools [gzip.BestCompression - gzip.BestSpeed + 2]*sync.Pool // 覆盖 -1 (Default) 到 9 (BestCompression)
func init() {
for i := gzip.BestSpeed; i <= gzip.BestCompression; i++ {
level := i // 捕获循环变量用于闭包
gzipWriterPools[level-gzip.BestSpeed] = &sync.Pool{
New: func() interface{} {
// 初始化时 writer 为 nil在 Reset 时设置
w, _ := gzip.NewWriterLevel(nil, level)
return w
},
}
}
// 为 gzip.DefaultCompression (-1) 映射一个索引
defaultLevelIndex := gzip.BestCompression - gzip.BestSpeed + 1
gzipWriterPools[defaultLevelIndex] = &sync.Pool{
New: func() interface{} {
w, _ := gzip.NewWriterLevel(nil, gzip.DefaultCompression)
return w
},
}
}
// 从对象池获取一个 gzip.Writer
func getGzipWriterFromPool(level int, underlyingWriter io.Writer) *gzip.Writer {
var poolIndex int
if level == gzip.DefaultCompression {
poolIndex = gzip.BestCompression - gzip.BestSpeed + 1
} else if level >= gzip.BestSpeed && level <= gzip.BestCompression {
poolIndex = level - gzip.BestSpeed
} else { // 无效级别,使用默认级别
poolIndex = gzip.BestCompression - gzip.BestSpeed + 1
level = gzip.DefaultCompression // 保证一致性
}
gz := gzipWriterPools[poolIndex].Get().(*gzip.Writer)
gz.Reset(underlyingWriter) // 重置并关联到底层的 io.Writer
return gz
}
// 将 gzip.Writer 返还给对象池
func putGzipWriterToPool(gz *gzip.Writer, level int) {
var poolIndex int
if level == gzip.DefaultCompression {
poolIndex = gzip.BestCompression - gzip.BestSpeed + 1
} else if level >= gzip.BestSpeed && level <= gzip.BestCompression {
poolIndex = level - gzip.BestSpeed
} else { // 不应该发生,如果 getGzipWriterFromPool 进行了标准化
poolIndex = gzip.BestCompression - gzip.BestSpeed + 1
}
gzipWriterPools[poolIndex].Put(gz)
}
// 从对象池获取一个 gzipResponseWriter
func acquireGzipResponseWriter(underlying ResponseWriter, opts *GzipOptions) *gzipResponseWriter {
gzw := gzipResponseWriterPool.Get().(*gzipResponseWriter)
gzw.ResponseWriter = underlying
gzw.options = opts
gzw.wroteHeader = false
gzw.doCompression = false
gzw.statusCode = 0 // 重置状态码
// gzWriter 将在 WriteHeader 中如果需要时获取
return gzw
}
// 将 gzipResponseWriter 返还给对象池
func releaseGzipResponseWriter(gzw *gzipResponseWriter) {
if gzw.gzWriter != nil {
// 确保它被关闭并返回到池中
_ = gzw.gzWriter.Close() // 关闭会 flush
putGzipWriterToPool(gzw.gzWriter, gzw.options.Level)
gzw.gzWriter = nil
}
gzw.ResponseWriter = nil // 断开引用
gzw.options = nil
gzipResponseWriterPool.Put(gzw)
}
// --- gzipResponseWriter 方法实现 ---
// Header 返回底层 ResponseWriter 的头部 map。
func (gzw *gzipResponseWriter) Header() http.Header {
return gzw.ResponseWriter.Header()
}
// WriteHeader 发送 HTTP 响应头部和指定的状态码。
// 在这里决定是否进行压缩。
func (gzw *gzipResponseWriter) WriteHeader(statusCode int) {
if gzw.wroteHeader {
return
}
gzw.wroteHeader = true
gzw.statusCode = statusCode // 存储状态码
// 在修改头部以进行压缩之前进行条件检查
// 1. 如果状态码是信息性(1xx)、重定向(3xx)、无内容(204)、重置内容(205)或未修改(304),则不压缩
if statusCode < http.StatusOK || statusCode == http.StatusNoContent || statusCode == http.StatusResetContent || statusCode == http.StatusNotModified {
gzw.ResponseWriter.WriteHeader(statusCode)
return
}
// 2. 如果响应已经被编码,则不压缩
if gzw.Header().Get(headerContentEncoding) != "" {
gzw.ResponseWriter.WriteHeader(statusCode)
return
}
// 3. 检查 Content-Type
contentType := strings.ToLower(strings.TrimSpace(strings.Split(gzw.Header().Get(headerContentType), ";")[0]))
compressibleTypes := gzw.options.CompressibleTypes
if len(compressibleTypes) == 0 {
compressibleTypes = defaultCompressibleTypes
}
isCompressible := false
for _, t := range compressibleTypes {
if strings.HasPrefix(contentType, t) { // 使用 HasPrefix 以匹配如 "text/html; charset=utf-8"
isCompressible = true
break
}
}
if !isCompressible {
gzw.ResponseWriter.WriteHeader(statusCode)
return
}
// 4. 检查 MinContentLength
if gzw.options.MinContentLength > 0 {
if clStr := gzw.Header().Get(headerContentLength); clStr != "" {
if cl, err := strconv.ParseInt(clStr, 10, 64); err == nil && cl < gzw.options.MinContentLength {
gzw.ResponseWriter.WriteHeader(statusCode)
return
}
}
// 如果未设置 Content-Length但设置了 MinContentLength我们可能仍会压缩。
// 这是一个权衡:可能会压缩小的动态内容。
}
// 所有检查通过,进行压缩
gzw.doCompression = true
gzw.Header().Set(headerContentEncoding, encodingGzip)
gzw.Header().Add(headerVary, headerAcceptEncoding) // 使用 Add 以避免覆盖其他 Vary 值
gzw.Header().Del(headerContentLength) // Gzip 会改变内容长度,所以删除它
// 从池中获取 gzWriter并将其 Reset 指向实际的底层 ResponseWriter
// 注意gzw.ResponseWriter 是被 Gzip 包装的 writer (例如,原始的 responseWriterImpl 或 ecw)
gzw.gzWriter = getGzipWriterFromPool(gzw.options.Level, gzw.ResponseWriter)
gzw.ResponseWriter.WriteHeader(statusCode) // 调用原始的 WriteHeader
}
// Write 将数据写入连接作为 HTTP 回复的一部分。
func (gzw *gzipResponseWriter) Write(data []byte) (int, error) {
if !gzw.wroteHeader {
// 如果在 WriteHeader 之前调用 Write根据 http.ResponseWriter 规范,
// 应写入 200 OK 头部。
gzw.WriteHeader(http.StatusOK)
}
if gzw.doCompression {
return gzw.gzWriter.Write(data)
}
return gzw.ResponseWriter.Write(data)
}
// Close 确保 gzip writer 被关闭并释放资源。
// 中间件应该在 c.Next() 之后调用它(通常在 defer 中)。
func (gzw *gzipResponseWriter) Close() error {
if gzw.gzWriter != nil {
err := gzw.gzWriter.Close() // Close 会 Flush
putGzipWriterToPool(gzw.gzWriter, gzw.options.Level)
gzw.gzWriter = nil // 标记为已返回
return err
}
return nil
}
// Flush 将所有缓冲数据发送到客户端。
// 实现 http.Flusher。
func (gzw *gzipResponseWriter) Flush() {
if gzw.doCompression && gzw.gzWriter != nil {
_ = gzw.gzWriter.Flush() // 确保 gzip writer 的缓冲被刷新
}
// 然后刷新底层的 writer (如果它支持)
if fl, ok := gzw.ResponseWriter.(http.Flusher); ok {
fl.Flush()
}
}
// Hijack 允许调用者接管连接。
// 实现 http.Hijacker。
func (gzw *gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
// 如果正在压缩hijack 的意义不大或不安全。
// 然而WriteHeader 应该会阻止对 101 状态码的压缩。
// 此调用必须转到实际的底层 ResponseWriter。
if hj, ok := gzw.ResponseWriter.(http.Hijacker); ok {
return hj.Hijack()
}
// 返回英文错误
return nil, nil, errors.New("touka.gzipResponseWriter: underlying ResponseWriter does not implement http.Hijacker")
}
// Status 返回已写入的 HTTP 状态码。委托给底层的 ResponseWriter。
// 这确保了与 ecw 或其他可能跟踪状态的包装器的兼容性。
func (gzw *gzipResponseWriter) Status() int {
if gzw.statusCode != 0 { // 如果我们在 WriteHeader 期间存储了它
return gzw.statusCode
}
return gzw.ResponseWriter.Status() // 委托
}
// Size 返回已写入的字节数。委托给底层的 ResponseWriter。
// 如果已压缩,这将是压缩后的大小 (由底层 writer 记录)。
func (gzw *gzipResponseWriter) Size() int {
return gzw.ResponseWriter.Size() // GzipResponseWriter 本身不直接跟踪大小,依赖底层
}
// Written 返回 WriteHeader 是否已被调用。委托给底层的 ResponseWriter。
func (gzw *gzipResponseWriter) Written() bool {
// 如果 gzw.wroteHeader 为 true说明 WriteHeader 至少被 gzw 处理过。
// 但最终是否写入底层取决于 gzw 的逻辑。
// 更可靠的是询问底层 writer。
return gzw.ResponseWriter.Written() // 委托
}
// --- Gzip 中间件 ---
// Gzip 返回一个使用 Gzip 压缩 HTTP 响应的中间件。
// 它会检查客户端的 "Accept-Encoding" 头部和响应的 "Content-Type"
// 来决定是否应用压缩。
// level 参数指定压缩级别 (例如 gzip.DefaultCompression)。
// opts 参数是可选的 GzipOptions。
func Gzip(level int, opts ...GzipOptions) HandlerFunc {
config := GzipOptions{ // 初始化默认配置
Level: level,
MinContentLength: 0, // 默认:无最小长度
CompressibleTypes: defaultCompressibleTypes,
}
if len(opts) > 0 { // 如果传入了 GzipOptions则覆盖默认值
opt := opts[0]
config.Level = opt.Level // 允许通过结构体覆盖级别
if opt.MinContentLength > 0 {
config.MinContentLength = opt.MinContentLength
}
if len(opt.CompressibleTypes) > 0 {
config.CompressibleTypes = opt.CompressibleTypes
}
}
// 验证级别
if config.Level < gzip.DefaultCompression || config.Level > gzip.BestCompression {
config.Level = gzip.DefaultCompression
}
return func(c *Context) {
// 1. 检查客户端是否接受 gzip
if !strings.Contains(c.Request.Header.Get(headerAcceptEncoding), encodingGzip) {
c.Next()
return
}
// 2. 包装 ResponseWriter
originalWriter := c.Writer
gzw := acquireGzipResponseWriter(originalWriter, &config)
c.Writer = gzw // 替换上下文的 writer
// defer 确保即使后续处理函数发生 panic也能进行清理
// 尽管恢复中间件应该自己处理 panic 响应。
defer func() {
// 必须关闭 gzip writer 以刷新其缓冲区。
// 这也会将 gzip.Writer 返回到其对象池。
if err := gzw.Close(); err != nil {
// 记录关闭 gzip writer 时的错误,但不应覆盖已发送的响应
// 通常这个错误不严重,因为数据可能已经大部分发送
// 使用英文记录日志
// log.Printf("Error closing gzip writer: %v", err)
c.AddError(err) // 可以选择将错误添加到 Context 中
}
// 恢复原始 writer 并将 gzipResponseWriter 返回到其对象池
c.Writer = originalWriter
releaseGzipResponseWriter(gzw)
}()
// 3. 调用链中的下一个处理函数
c.Next()
// c.Next() 执行完毕后,响应头部应该已经设置。
// gzw.WriteHeader 会被显式调用或通过第一次 Write 隐式调用。
// 如果 gzw.doCompression 为 true响应体已写入 gzw.gzWriter。
// defer 中的 gzw.Close() 会刷新最终的压缩字节。
}
}

130
ws.go Normal file
View file

@ -0,0 +1,130 @@
package touka
import (
"errors"
"log"
"net/http"
"github.com/gorilla/websocket"
)
// WebSocketHandler 是用户提供的用于处理 WebSocket 连接的函数类型。
// conn 是一个已经完成握手的 WebSocket 连接。
type WebSocketHandler func(c *Context, conn *websocket.Conn)
// WebSocketUpgradeOptions 用于配置 WebSocket 升级中间件。
type WebSocketUpgradeOptions struct {
// Upgrader 是 gorilla/websocket.Upgrader 的实例。
// 用户可以配置 ReadBufferSize, WriteBufferSize, CheckOrigin, Subprotocols 等。
// 如果为 nil将使用一个带有合理默认值的 Upgrader。
Upgrader *websocket.Upgrader
// Handler 是在 WebSocket 成功升级后调用的处理函数。
// 这个字段是必需的。
Handler WebSocketHandler
// OnError 是一个可选的回调函数,用于处理升级过程中发生的错误。
// 如果未提供,错误将导致一个标准的 HTTP 错误响应(例如 400 Bad Request
OnError func(c *Context, status int, err error)
}
// defaultWebSocketUpgrader 返回一个具有合理默认值的 websocket.Upgrader。
func defaultWebSocketUpgrader() *websocket.Upgrader {
return &websocket.Upgrader{
ReadBufferSize: 1024,
WriteBufferSize: 1024,
// CheckOrigin 应该由用户根据其安全需求来配置。
// 默认情况下,如果 Origin 头部存在且与 Host 头部不匹配,会拒绝连接。
// 对于开发,可以暂时设置为 func(r *http.Request) bool { return true }
// 但在生产环境中必须小心配置。
CheckOrigin: func(r *http.Request) bool {
// 简单的同源检查或允许所有 (根据需要调整)
// return r.Header.Get("Origin") == "" || strings.HasPrefix(r.Header.Get("Origin"), "http://"+r.Host) || strings.HasPrefix(r.Header.Get("Origin"), "https://"+r.Host)
return true // 示例:允许所有,生产环境请谨慎
},
}
}
// defaultWebSocketOnError 是默认的错误处理函数。
func defaultWebSocketOnError(c *Context, status int, err error) {
// 使用框架的错误处理机制或简单的字符串响应
// 确保不要写入一个已经开始的响应
if !c.Writer.Written() {
// 返回英文错误信息
errMsg := http.StatusText(status)
if err != nil {
errMsg = err.Error() // 可以考虑是否暴露详细错误
}
c.String(status, "%s", errMsg) // 或者 c.engine.errorHandle.handler(c, status)
}
c.Abort() // 总是中止
}
// WebSocketUpgrade 返回一个 WebSocket 升级中间件。
// 它能自动感知 HTTP/1.1 的 Upgrade 请求和 HTTP/2 的扩展 CONNECT 请求 (RFC 8441)。
func WebSocketUpgrade(opts WebSocketUpgradeOptions) HandlerFunc {
if opts.Handler == nil {
panic("touka: WebSocketUpgradeOptions.Handler cannot be nil")
}
upgrader := opts.Upgrader
if upgrader == nil {
upgrader = defaultWebSocketUpgrader()
}
onError := opts.OnError
if onError == nil {
onError = defaultWebSocketOnError
}
return func(c *Context) {
// 调试日志,查看请求详情
// reqBytes, _ := httputil.DumpRequest(c.Request, true)
// log.Printf("WebSocketUpgrade: Incoming request for path %s:\n%s", c.Request.URL.Path, string(reqBytes))
// log.Printf("Request Proto: %s, Method: %s", c.Request.Proto, c.Request.Method)
// 对于我们的目的,让 gorilla/websocket 的 Upgrade 方法去判断更佳,
// 它已经实现了 RFC 8441 的支持。
// 我们不再需要手动区分 HTTP/1.1 和 HTTP/2 的逻辑,
// gorilla/websocket.Upgrader.Upgrade 会自动处理。
// 它会检查请求是 HTTP/1.1 Upgrade 还是 HTTP/2 CONNECT with :protocol=websocket。
// 对于 HTTP/2Upgrade() 方法不会发送 101而是处理 CONNECT 的 200 OK。
// 它也不会调用 Hijack因为连接已经在 HTTP/2 流上。
conn, err := upgrader.Upgrade(c.Writer, c.Request, nil)
if err != nil {
// 升级失败。gorilla/websocket.Upgrade 会处理错误响应的发送。
// (对于 HTTP/1.1 会是 400/403 等;对于 HTTP/2 也是类似的非 2xx 响应)
var httpErr websocket.HandshakeError
statusCode := http.StatusBadRequest // 默认
if errors.As(err, &httpErr) {
// 尝试获取更具体的错误信息,但状态码可能不直接暴露
}
// 使用英文记录日志
log.Printf("WebSocket upgrade/handshake failed for %s (Proto: %s): %v", c.Request.RemoteAddr, c.Request.Proto, err)
onError(c, statusCode, err)
if !c.IsAborted() {
c.Abort()
}
return
}
// 升级/握手成功
// 使用英文记录日志
log.Printf("WebSocket connection established for %s (Proto: %s)", c.Request.RemoteAddr, c.Request.Proto)
if !c.IsAborted() {
c.Abort() // 确保 HTTP 处理链中止
}
defer func() {
// 使用英文记录日志
log.Printf("Closing WebSocket connection for %s", conn.RemoteAddr())
_ = conn.Close()
}()
opts.Handler(c, conn) // 执行用户定义的 WebSocket 逻辑
}
}