Merge pull request #74 from infinite-iroha/break/v1-feat-add-samesite

feat(cookie): add SameSite support to SetCookie method

context.go

@@ -417,6 +417,22 @@ func (c *Context) JSON(code int, obj any) {
 	}
 }
 
+// JSONBuf 先将 JSON 编码到 buffer, 成功后再写入状态码和响应体.
+// 与 JSON 相比，编码失败时可以正确返回 500 状态码，代价是多一次内存分配.
+func (c *Context) JSONBuf(code int, obj any) {
+	var buf bytes.Buffer
+	if err := json.MarshalWrite(&buf, obj); err != nil {
+		errMsg := fmt.Errorf("failed to marshal JSON: %w", err)
+		c.AddError(errMsg)
+		c.ErrorUseHandle(http.StatusInternalServerError, errMsg)
+		return
+	}
+
+	c.Writer.Header().Set("Content-Type", "application/json; charset=utf-8")
+	c.Writer.WriteHeader(code)
+	c.Writer.Write(buf.Bytes())
+}
+
 // GOB 向响应写入GOB数据
 // 设置 Content-Type 为 application/octet-stream
 func (c *Context) GOB(code int, obj any) {
@@ -431,6 +447,21 @@ func (c *Context) GOB(code int, obj any) {
 	}
 }
 
+// GOBBuf 先将 GOB 编码到 buffer, 成功后再写入状态码和响应体.
+func (c *Context) GOBBuf(code int, obj any) {
+	var buf bytes.Buffer
+	encoder := gob.NewEncoder(&buf)
+	if err := encoder.Encode(obj); err != nil {
+		errMsg := fmt.Errorf("failed to encode GOB: %w", err)
+		c.AddError(errMsg)
+		c.ErrorUseHandle(http.StatusInternalServerError, errMsg)
+		return
+	}
+	c.Writer.Header().Set("Content-Type", "application/octet-stream")
+	c.Writer.WriteHeader(code)
+	c.Writer.Write(buf.Bytes())
+}
+
 // WANF向响应写入WANF数据
 // 设置 application/vnd.wjqserver.wanf; charset=utf-8
 func (c *Context) WANF(code int, obj any) {
@@ -445,6 +476,21 @@ func (c *Context) WANF(code int, obj any) {
 	}
 }
 
+// WANFBuf 先将 WANF 编码到 buffer, 成功后再写入状态码和响应体.
+func (c *Context) WANFBuf(code int, obj any) {
+	var buf bytes.Buffer
+	encoder := wanf.NewStreamEncoder(&buf)
+	if err := encoder.Encode(obj); err != nil {
+		errMsg := fmt.Errorf("failed to encode WANF: %w", err)
+		c.AddError(errMsg)
+		c.ErrorUseHandle(http.StatusInternalServerError, errMsg)
+		return
+	}
+	c.Writer.Header().Set("Content-Type", "application/vnd.wjqserver.wanf; charset=utf-8")
+	c.Writer.WriteHeader(code)
+	c.Writer.Write(buf.Bytes())
+}
+
 // HTML 渲染 HTML 模板
 // 如果 Engine 配置了 HTMLRender，则使用它进行渲染
 // 否则，会进行简单的字符串输出
@@ -469,6 +515,36 @@ func (c *Context) HTML(code int, name string, obj any) {
 	c.Writer.Write(fmt.Appendf(nil, "<!-- HTML rendered for %s -->\n<pre>%v</pre>", name, obj))
 }
 
+// HTMLBuf 先将 HTML 模板渲染到 buffer, 成功后再写入状态码和响应体.
+// 如果模板渲染失败，则返回 500 错误且不写入任何内容.
+func (c *Context) HTMLBuf(code int, name string, obj any) {
+	if c.engine == nil || c.engine.HTMLRender == nil {
+		// 没有渲染器，回退到简单输出
+		c.HTML(code, name, obj)
+		return
+	}
+
+	if tpl, ok := c.engine.HTMLRender.(*template.Template); ok {
+		var buf bytes.Buffer
+		err := tpl.ExecuteTemplate(&buf, name, obj)
+		if err != nil {
+			// 渲染失败，记录错误并返回 500，不写入任何内容
+			errMsg := fmt.Errorf("failed to render HTML template '%s': %w", name, err)
+			c.AddError(errMsg)
+			c.ErrorUseHandle(http.StatusInternalServerError, errMsg)
+			return
+		}
+		// 渲染成功，写入响应
+		c.Writer.Header().Set("Content-Type", "text/html; charset=utf-8")
+		c.Writer.WriteHeader(code)
+		c.Writer.Write(buf.Bytes())
+		return
+	}
+
+	// 不支持的渲染器类型，回退到简单输出
+	c.HTML(code, name, obj)
+}
+
 // Redirect 执行 HTTP 重定向
 // code 应为 3xx 状态码 (如 http.StatusMovedPermanently, http.StatusFound)
 func (c *Context) Redirect(code int, location string) {
@@ -1084,17 +1160,25 @@ func (c *Context) SetSameSite(samesite http.SameSite) {
 }
 
 // SetCookie 设置一个 HTTP cookie
-func (c *Context) SetCookie(name, value string, maxAge int, path, domain string, secure, httpOnly bool) {
+// sameSite 参数是可选的，如果不提供则使用通过 SetSameSite 设置的值
+func (c *Context) SetCookie(name, value string, maxAge int, path, domain string, secure, httpOnly bool, sameSite ...http.SameSite) {
 	if path == "" {
 		path = "/"
 	}
+	site := c.sameSite
+	if len(sameSite) > 0 {
+		if len(sameSite) > 1 {
+			c.Warnf("SetCookie: only the first SameSite value will be used, got %d values", len(sameSite))
+		}
+		site = sameSite[0]
+	}
 	http.SetCookie(c.Writer, &http.Cookie{
 		Name:     name,
 		Value:    url.QueryEscape(value),
 		MaxAge:   maxAge,
 		Path:     path,
 		Domain:   domain,
-		SameSite: c.sameSite,
+		SameSite: site,
 		Secure:   secure,
 		HttpOnly: httpOnly,
 	})