diff --git a/README.md b/README.md index 8a226a9..83a6ee4 100644 --- a/README.md +++ b/README.md @@ -171,7 +171,9 @@ func min(a, b int) int { ### 内置 -Recovery `r.Use(touka.Recovery())` +Recovery `r.Use(Recovery())` + +Touka Gzip `r.Use(Gzip(-1))` ### fenthope diff --git a/context.go b/context.go index d9ac4a6..3172f6f 100644 --- a/context.go +++ b/context.go @@ -44,8 +44,8 @@ type Context struct { // 携带ctx以实现关闭逻辑 ctx context.Context - // HTTPClient 用于在此上下文中执行出站 HTTP 请求 - // 它由 Engine 提供 + // HTTPClient 用于在此上下文中执行出站 HTTP 请求。 + // 它由 Engine 提供。 HTTPClient *httpc.Client // 引用所属的 Engine 实例，方便访问 Engine 的配置（如 HTMLRender） @@ -56,8 +56,8 @@ type Context struct { // --- Context 相关方法实现 --- -// reset 重置 Context 对象以供复用 -// 每次从 sync.Pool 中获取 Context 后，都需要调用此方法进行初始化 +// reset 重置 Context 对象以供复用。 +// 每次从 sync.Pool 中获取 Context 后，都需要调用此方法进行初始化。 func (c *Context) reset(w http.ResponseWriter, req *http.Request) { if rw, ok := c.Writer.(*responseWriterImpl); ok && !rw.IsHijacked() { @@ -80,8 +80,8 @@ func (c *Context) reset(w http.ResponseWriter, req *http.Request) { // c.HTTPClient 和 c.engine 保持不变，它们引用 Engine 实例的成员 } -// Next 在处理链中执行下一个处理函数 -// 这是中间件模式的核心，允许请求依次经过多个处理函数 +// Next 在处理链中执行下一个处理函数。 +// 这是中间件模式的核心，允许请求依次经过多个处理函数。 func (c *Context) Next() { c.index++ for c.index < int8(len(c.handlers)) { @@ -90,25 +90,25 @@ func (c *Context) Next() { } } -// Abort 停止处理链的后续执行 -// 通常在中间件中，当遇到错误或需要提前终止请求时调用 +// Abort 停止处理链的后续执行。 +// 通常在中间件中，当遇到错误或需要提前终止请求时调用。 func (c *Context) Abort() { c.index = abortIndex // 将 index 设置为一个很大的值，使后续 Next() 调用跳过所有处理函数 } -// IsAborted 返回处理链是否已被中止 +// IsAborted 返回处理链是否已被中止。 func (c *Context) IsAborted() bool { return c.index >= abortIndex } -// AbortWithStatus 中止处理链并设置 HTTP 状态码 +// AbortWithStatus 中止处理链并设置 HTTP 状态码。 func (c *Context) AbortWithStatus(code int) { c.Writer.WriteHeader(code) // 设置响应状态码 c.Abort() // 中止处理链 } -// Set 将一个键值对存储到 Context 中 -// 这是一个线程安全的操作，用于在中间件之间传递数据 +// Set 将一个键值对存储到 Context 中。 +// 这是一个线程安全的操作，用于在中间件之间传递数据。 func (c *Context) Set(key string, value interface{}) { c.mu.Lock() // 加写锁 if c.Keys == nil { @@ -118,8 +118,8 @@ func (c *Context) Set(key string, value interface{}) { c.mu.Unlock() // 解写锁 } -// Get 从 Context 中获取一个值 -// 这是一个线程安全的操作 +// Get 从 Context 中获取一个值。 +// 这是一个线程安全的操作。 func (c *Context) Get(key string) (value interface{}, exists bool) { c.mu.RLock() // 加读锁 value, exists = c.Keys[key] @@ -127,8 +127,8 @@ func (c *Context) Get(key string) (value interface{}, exists bool) { return } -// MustGet 从 Context 中获取一个值，如果不存在则 panic -// 适用于确定值一定存在的场景 +// MustGet 从 Context 中获取一个值，如果不存在则 panic。 +// 适用于确定值一定存在的场景。 func (c *Context) MustGet(key string) interface{} { if value, exists := c.Get(key); exists { return value @@ -136,8 +136,8 @@ func (c *Context) MustGet(key string) interface{} { panic("Key \"" + key + "\" does not exist in context.") } -// Query 从 URL 查询参数中获取值 -// 懒加载解析查询参数，并进行缓存 +// Query 从 URL 查询参数中获取值。 +// 懒加载解析查询参数，并进行缓存。 func (c *Context) Query(key string) string { if c.queryCache == nil { c.queryCache = c.Request.URL.Query() // 首次访问时解析并缓存 @@ -145,7 +145,7 @@ func (c *Context) Query(key string) string { return c.queryCache.Get(key) } -// DefaultQuery 从 URL 查询参数中获取值，如果不存在则返回默认值 +// DefaultQuery 从 URL 查询参数中获取值，如果不存在则返回默认值。 func (c *Context) DefaultQuery(key, defaultValue string) string { if value := c.Query(key); value != "" { return value @@ -153,8 +153,8 @@ func (c *Context) DefaultQuery(key, defaultValue string) string { return defaultValue } -// PostForm 从 POST 请求体中获取表单值 -// 懒加载解析表单数据，并进行缓存 +// PostForm 从 POST 请求体中获取表单值。 +// 懒加载解析表单数据，并进行缓存。 func (c *Context) PostForm(key string) string { if c.formCache == nil { c.Request.ParseMultipartForm(defaultMemory) // 解析 multipart/form-data 或 application/x-www-form-urlencoded @@ -163,7 +163,7 @@ func (c *Context) PostForm(key string) string { return c.formCache.Get(key) } -// DefaultPostForm 从 POST 请求体中获取表单值，如果不存在则返回默认值 +// DefaultPostForm 从 POST 请求体中获取表单值，如果不存在则返回默认值。 func (c *Context) DefaultPostForm(key, defaultValue string) string { if value := c.PostForm(key); value != "" { return value @@ -171,8 +171,8 @@ func (c *Context) DefaultPostForm(key, defaultValue string) string { return defaultValue } -// Param 从 URL 路径参数中获取值 -// 例如，对于路由 /users/:id，c.Param("id") 可以获取 id 的值 +// Param 从 URL 路径参数中获取值。 +// 例如，对于路由 /users/:id，c.Param("id") 可以获取 id 的值。 func (c *Context) Param(key string) string { return c.Params.ByName(key) } @@ -184,14 +184,14 @@ func (c *Context) Raw(code int, contentType string, data []byte) { c.Writer.Write(data) } -// String 向响应写入格式化的字符串 +// String 向响应写入格式化的字符串。 func (c *Context) String(code int, format string, values ...interface{}) { c.Writer.WriteHeader(code) c.Writer.Write([]byte(fmt.Sprintf(format, values...))) } -// JSON 向响应写入 JSON 数据 -// 设置 Content-Type 为 application/json +// JSON 向响应写入 JSON 数据。 +// 设置 Content-Type 为 application/json。 func (c *Context) JSON(code int, obj interface{}) { c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8") c.Writer.WriteHeader(code) @@ -205,10 +205,10 @@ func (c *Context) JSON(code int, obj interface{}) { c.Writer.Write(jsonBytes) } -// HTML 渲染 HTML 模板 -// 如果 Engine 配置了 HTMLRender，则使用它进行渲染 -// 否则，会进行简单的字符串输出 -// 预留接口，可以扩展为支持多种模板引擎 +// HTML 渲染 HTML 模板。 +// 如果 Engine 配置了 HTMLRender，则使用它进行渲染。 +// 否则，会进行简单的字符串输出。 +// 预留接口，可以扩展为支持多种模板引擎。 func (c *Context) HTML(code int, name string, obj interface{}) { c.Writer.Header().Set("Content-Type", "text/html; charset=utf-8") c.Writer.WriteHeader(code) @@ -229,8 +229,8 @@ func (c *Context) HTML(code int, name string, obj interface{}) { c.Writer.Write([]byte(fmt.Sprintf("\n
%v
", name, obj))) } -// Redirect 执行 HTTP 重定向 -// code 应为 3xx 状态码 (如 http.StatusMovedPermanently, http.StatusFound) +// Redirect 执行 HTTP 重定向。 +// code 应为 3xx 状态码 (如 http.StatusMovedPermanently, http.StatusFound)。 func (c *Context) Redirect(code int, location string) { http.Redirect(c.Writer, c.Request, location, code) c.Abort() @@ -239,7 +239,7 @@ func (c *Context) Redirect(code int, location string) { } } -// ShouldBindJSON 尝试将请求体绑定到 JSON 对象 +// ShouldBindJSON 尝试将请求体绑定到 JSON 对象。 func (c *Context) ShouldBindJSON(obj interface{}) error { if c.Request.Body == nil { return errors.New("request body is empty") @@ -257,9 +257,9 @@ func (c *Context) ShouldBindJSON(obj interface{}) error { return nil } -// ShouldBind 尝试将请求体绑定到各种类型（JSON, Form, XML 等） -// 这是一个复杂的通用绑定接口，通常根据 Content-Type 或其他头部来判断绑定方式 -// 预留接口，可根据项目需求进行扩展 +// ShouldBind 尝试将请求体绑定到各种类型（JSON, Form, XML 等）。 +// 这是一个复杂的通用绑定接口，通常根据 Content-Type 或其他头部来判断绑定方式。 +// 预留接口，可根据项目需求进行扩展。 func (c *Context) ShouldBind(obj interface{}) error { // TODO: 完整的通用绑定逻辑 // 可以根据 c.Request.Header.Get("Content-Type") 来判断是 JSON, Form, XML 等 @@ -274,45 +274,45 @@ func (c *Context) ShouldBind(obj interface{}) error { return errors.New("generic binding not fully implemented yet, implement based on Content-Type") } -// AddError 添加一个错误到 Context -// 允许在处理请求过程中收集多个错误 +// AddError 添加一个错误到 Context。 +// 允许在处理请求过程中收集多个错误。 func (c *Context) AddError(err error) { c.Errors = append(c.Errors, err) } -// Errors 返回 Context 中收集的所有错误 +// Errors 返回 Context 中收集的所有错误。 func (c *Context) GetErrors() []error { return c.Errors } -// Client 返回 Engine 提供的 HTTPClient -// 方便在请求处理函数中进行出站 HTTP 请求 +// Client 返回 Engine 提供的 HTTPClient。 +// 方便在请求处理函数中进行出站 HTTP 请求。 func (c *Context) Client() *httpc.Client { return c.HTTPClient } -// Context() 返回请求的上下文，用于取消操作 -// 这是 Go 标准库的 `context.Context`，用于请求的取消和超时管理 +// Context() 返回请求的上下文，用于取消操作。 +// 这是 Go 标准库的 `context.Context`，用于请求的取消和超时管理。 func (c *Context) Context() context.Context { return c.ctx } // Done returns a channel that is closed when the request context is cancelled or times out. -// 继承自 `context.Context` +// 继承自 `context.Context`。 func (c *Context) Done() <-chan struct{} { return c.ctx.Done() } // Err returns the error, if any, that caused the context to be canceled or to // time out. -// 继承自 `context.Context` +// 继承自 `context.Context`。 func (c *Context) Err() error { return c.ctx.Err() } // Value returns the value associated with this context for key, or nil if no // value is associated with key. -// 可以用于从 Context 中获取与特定键关联的值，包括 Go 原生 Context 的值和 Touka Context 的 Keys +// 可以用于从 Context 中获取与特定键关联的值，包括 Go 原生 Context 的值和 Touka Context 的 Keys。 func (c *Context) Value(key interface{}) interface{} { if keyAsString, ok := key.(string); ok { if val, exists := c.Get(keyAsString); exists { @@ -322,18 +322,18 @@ func (c *Context) Value(key interface{}) interface{} { return c.ctx.Value(key) // 尝试从 Go 原生 Context 中获取值 } -// GetWriter 获得一个 io.Writer 接口，可以直接向响应体写入数据 -// 这对于需要自定义流式写入或与其他需要 io.Writer 的库集成非常有用 +// GetWriter 获得一个 io.Writer 接口，可以直接向响应体写入数据。 +// 这对于需要自定义流式写入或与其他需要 io.Writer 的库集成非常有用。 func (c *Context) GetWriter() io.Writer { return c.Writer // ResponseWriter 接口嵌入了 http.ResponseWriter，而 http.ResponseWriter 实现了 io.Writer } -// WriteStream 接受一个 io.Reader 并将其内容流式传输到响应体 -// 返回写入的字节数和可能遇到的错误 -// 该方法在开始写入之前，会确保设置 HTTP 状态码为 200 OK +// WriteStream 接受一个 io.Reader 并将其内容流式传输到响应体。 +// 返回写入的字节数和可能遇到的错误。 +// 该方法在开始写入之前，会确保设置 HTTP 状态码为 200 OK。 func (c *Context) WriteStream(reader io.Reader) (written int64, err error) { - // 确保在写入数据前设置状态码 - // WriteHeader 会在第一次写入时被 Write 方法隐式调用，但显式调用可以确保状态码的预期 + // 确保在写入数据前设置状态码。 + // WriteHeader 会在第一次写入时被 Write 方法隐式调用，但显式调用可以确保状态码的预期。 if !c.Writer.Written() { c.Writer.WriteHeader(http.StatusOK) // 默认 200 OK } @@ -346,14 +346,14 @@ func (c *Context) WriteStream(reader io.Reader) (written int64, err error) { } // GetReqBody 以获取一个 io.ReadCloser 接口，用于读取请求体 -// 注意：请求体只能读取一次 +// 注意：请求体只能读取一次。 func (c *Context) GetReqBody() io.ReadCloser { return c.Request.Body } // GetReqBodyFull -// GetReqBodyFull 读取并返回请求体的所有内容 -// 注意：请求体只能读取一次 +// GetReqBodyFull 读取并返回请求体的所有内容。 +// 注意：请求体只能读取一次。 func (c *Context) GetReqBodyFull() ([]byte, error) { if c.Request.Body == nil { return nil, nil @@ -367,9 +367,9 @@ func (c *Context) GetReqBodyFull() ([]byte, error) { return data, nil } -// RequestIP 返回客户端的 IP 地址 +// RequestIP 返回客户端的 IP 地址。 // 它会根据 Engine 的配置 (ForwardByClientIP) 尝试从 X-Forwarded-For 或 X-Real-IP 等头部获取， -// 否则回退到 Request.RemoteAddr +// 否则回退到 Request.RemoteAddr。 func (c *Context) RequestIP() string { if c.engine.ForwardByClientIP { for _, headerName := range c.engine.RemoteIPHeaders { @@ -409,60 +409,55 @@ func (c *Context) RequestIP() string { return "" } -// ClientIP 返回客户端的 IP 地址 -// 这是一个别名，与 RequestIP 功能相同 +// ClientIP 返回客户端的 IP 地址。 +// 这是一个别名，与 RequestIP 功能相同。 func (c *Context) ClientIP() string { return c.RequestIP() } -// ContentType 返回请求的 Content-Type 头部 +// ContentType 返回请求的 Content-Type 头部。 func (c *Context) ContentType() string { return c.GetReqHeader("Content-Type") } -// UserAgent 返回请求的 User-Agent 头部 +// UserAgent 返回请求的 User-Agent 头部。 func (c *Context) UserAgent() string { return c.GetReqHeader("User-Agent") } -// Status 设置响应状态码 +// Status 设置响应状态码。 func (c *Context) Status(code int) { c.Writer.WriteHeader(code) } -// File 将指定路径的文件作为响应发送 -// 它会设置 Content-Type 和 Content-Disposition 头部 +// File 将指定路径的文件作为响应发送。 +// 它会设置 Content-Type 和 Content-Disposition 头部。 func (c *Context) File(filepath string) { http.ServeFile(c.Writer, c.Request, filepath) c.Abort() // 发送文件后中止后续处理 } -// SetHeader 设置响应头部 +// SetHeader 设置响应头部。 func (c *Context) SetHeader(key, value string) { c.Writer.Header().Set(key, value) } -// AddHeader 添加响应头部 +// AddHeader 添加响应头部。 func (c *Context) AddHeader(key, value string) { c.Writer.Header().Add(key, value) } -// Header 作为SetHeader的别名 -func (c *Context) Header(key, value string) { - c.SetHeader(key, value) -} - -// DelHeader 删除响应头部 +// DelHeader 删除响应头部。 func (c *Context) DelHeader(key string) { c.Writer.Header().Del(key) } -// GetReqHeader 获取请求头部的值 +// GetReqHeader 获取请求头部的值。 func (c *Context) GetReqHeader(key string) string { return c.Request.Header.Get(key) } -// GetAllReqHeader 获取所有请求头部 +// GetAllReqHeader 获取所有请求头部。 func (c *Context) GetAllReqHeader() http.Header { return c.Request.Header } @@ -494,12 +489,12 @@ func (c *Context) GetLogger() *reco.Logger { return c.engine.LogReco } -// SetSameSite 设置响应的 SameSite cookie 属性 +// SetSameSite 设置响应的 SameSite cookie 属性。 func (c *Context) SetSameSite(samesite http.SameSite) { c.sameSite = samesite } -// SetCookie 设置一个 HTTP cookie +// SetCookie 设置一个 HTTP cookie。 func (c *Context) SetCookie(name, value string, maxAge int, path, domain string, secure, httpOnly bool) { if path == "" { path = "/" @@ -526,7 +521,7 @@ func (c *Context) SetCookieData(cookie *http.Cookie) { http.SetCookie(c.Writer, cookie) } -// GetCookie 获取指定名称的 cookie 值 +// GetCookie 获取指定名称的 cookie 值。 func (c *Context) GetCookie(name string) (string, error) { cookie, err := c.Request.Cookie(name) if err != nil { @@ -540,8 +535,8 @@ func (c *Context) GetCookie(name string) (string, error) { return value, nil } -// DeleteCookie 删除指定名称的 cookie -// 通过设置 MaxAge 为 -1 来删除 cookie +// DeleteCookie 删除指定名称的 cookie。 +// 通过设置 MaxAge 为 -1 来删除 cookie。 func (c *Context) DeleteCookie(name string) { c.SetCookie(name, "", -1, "/", "", false, false) // 设置 MaxAge 为 -1 删除 cookie } diff --git a/go.mod b/go.mod index 6c71120..dbeb743 100644 --- a/go.mod +++ b/go.mod @@ -4,9 +4,10 @@ go 1.24.4 require ( github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4 - github.com/WJQSERVER-STUDIO/httpc v0.6.0 + github.com/WJQSERVER-STUDIO/httpc v0.5.1 github.com/fenthope/reco v0.0.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 diff --git a/go.sum b/go.sum index e5d2a67..69661e4 100644 --- a/go.sum +++ b/go.sum @@ -1,10 +1,12 @@ github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4 h1:JLtFd00AdFg/TP+dtvIzLkdHwKUGPOAijN1sMtEYoFg= github.com/WJQSERVER-STUDIO/go-utils/copyb v0.0.4/go.mod h1:FZ6XE+4TKy4MOfX1xWKe6Rwsg0ucYFCdNh1KLvyKTfc= -github.com/WJQSERVER-STUDIO/httpc v0.6.0 h1:GXKc6BNGn5ALdDLkCsM+mP24jyi1S9QBLL2XQ1a2sPM= -github.com/WJQSERVER-STUDIO/httpc v0.6.0/go.mod h1:M7KNUZjjhCkzzcg9lBPs9YfkImI+7vqjAyjdA19+joE= +github.com/WJQSERVER-STUDIO/httpc v0.5.1 h1:+TKCPYBuj7PAHuiduGCGAqsHAa4QtsUfoVwRN777q64= +github.com/WJQSERVER-STUDIO/httpc v0.5.1/go.mod h1:M7KNUZjjhCkzzcg9lBPs9YfkImI+7vqjAyjdA19+joE= github.com/fenthope/reco v0.0.1 h1:GYcuXCEKYoctD0dFkiBC+t0RMTOyOiujBCin8bbLR3Y= github.com/fenthope/reco v0.0.1/go.mod h1:mDkGLHte5udWTIcjQTxrABRcf56SSdxBOCLgrRDwI/Y= 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= diff --git a/tgzip.go b/tgzip.go new file mode 100644 index 0000000..419cc6f --- /dev/null +++ b/tgzip.go @@ -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() 会刷新最终的压缩字节。 + } +} diff --git a/ws.go b/ws.go new file mode 100644 index 0000000..0ee0ee3 --- /dev/null +++ b/ws.go @@ -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/2，Upgrade() 方法不会发送 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 逻辑 + } +}