// This Source Code Form is subject to the terms of the Mozilla Public License, v. 2.0. If a copy of the MPL was not distributed with this // file, You can obtain one at https://mozilla.org/MPL/2.0/. // Copyright 2025 WJQSERVER. All rights reserved. // All rights reserved by WJQSERVER, related rights can be exercised by the infinite-iroha organization. package touka import ( "bytes" "io" "net/http" "strings" ) // Event 代表一个服务器发送事件(SSE). type Event struct { // Event 是事件的名称. Event string // Data 是事件的内容, 可以是多行文本. Data string // Id 是事件的唯一标识符. Id string // Retry 是指定客户端在连接丢失后应等待多少毫秒后尝试重新连接. Retry string } // Render 将事件格式化并写入给定的 writer. // 通过逐行处理数据, 此方法可防止因数据中包含换行符而导致的CRLF注入问题. // 为了性能, 它使用 bytes.Buffer 并通过 WriteTo 直接写入, 以避免不必要的内存分配. func (e *Event) Render(w io.Writer) error { var buf bytes.Buffer if len(e.Id) > 0 { buf.WriteString("id: ") buf.WriteString(e.Id) buf.WriteString("\n") } if len(e.Event) > 0 { buf.WriteString("event: ") buf.WriteString(e.Event) buf.WriteString("\n") } if len(e.Data) > 0 { lines := strings.SplitSeq(e.Data, "\n") for line := range lines { buf.WriteString("data: ") buf.WriteString(line) buf.WriteString("\n") } } if len(e.Retry) > 0 { buf.WriteString("retry: ") buf.WriteString(e.Retry) buf.WriteString("\n") } // 每个事件都以一个额外的换行符结尾. buf.WriteString("\n") // 直接将 buffer 的内容写入 writer, 避免生成中间字符串. _, err := buf.WriteTo(w) return err } // EventStream 启动一个 SSE 事件流. // 这是推荐的、更简单安全的方式, 采用阻塞和回调的设计, 框架负责管理连接生命周期. // // 详细用法: // // r.GET("/sse/callback", func(c *touka.Context) { // // streamer 回调函数会在一个循环中被调用. // c.EventStream(func(w io.Writer) bool { // event := touka.Event{ // Event: "time-tick", // Data: time.Now().Format(time.RFC1123), // } // // if err := event.Render(w); err != nil { // // 发生写入错误, 停止发送. // return false // 返回 false 结束事件流. // } // // time.Sleep(2 * time.Second) // return true // 返回 true 继续事件流. // }) // // 当事件流结束后(例如客户端关闭页面), 这行代码会被执行. // fmt.Println("Client disconnected from /sse/callback") // }) func (c *Context) EventStream(streamer func(w io.Writer) bool) { // 为现代网络协议优化头部. c.Writer.Header().Set("Content-Type", "text/event-stream; charset=utf-8") c.Writer.Header().Set("Cache-Control", "no-cache, no-transform") c.Writer.Header().Del("Connection") c.Writer.Header().Del("Transfer-Encoding") c.Writer.WriteHeader(http.StatusOK) c.Writer.Flush() // 直接调用, ResponseWriter 接口保证了 Flush 方法的存在. for { select { case <-c.Request.Context().Done(): return default: if !streamer(c.Writer) { return } c.Writer.Flush() } } } // EventStreamChan 返回用于 SSE 事件流的 channel. // 这是为高级并发场景设计的、更灵活的API. // // 与 EventStream 回调模式类似, 此方法是阻塞的: handler 会在此方法中停留, // 直到事件 channel 被关闭 (close eventChan) 或客户端断开连接. // 这保证了 Context 不会在 SSE 流期间被 pool 回收. // // eventChan 必须在调用此方法之前创建, 以便调用者可以在独立的 goroutine 中发送事件. // 调用者必须在完成后 close(eventChan) 来结束流. // 生产者 goroutine 必须在 select 中监听 c.Request.Context().Done(), 否则在客户端断开时会产生 goroutine 泄漏. // // 详细用法: // // r.GET("/sse/channel", func(c *touka.Context) { // eventChan := make(chan touka.Event) // // // 在独立的 goroutine 中异步发送事件. // go func() { // defer close(eventChan) // 完成后关闭 channel 以结束事件流. // // for i := 1; i <= 5; i++ { // select { // case <-c.Request.Context().Done(): // return // 客户端已断开, 退出 goroutine. // default: // eventChan <- touka.Event{ // Id: fmt.Sprintf("%d", i), // Data: "hello from channel", // } // time.Sleep(2 * time.Second) // } // } // }() // // // 阻塞直到事件流结束. // c.EventStreamChan(eventChan) // }) func (c *Context) EventStreamChan(eventChan <-chan Event) { c.Writer.Header().Set("Content-Type", "text/event-stream; charset=utf-8") c.Writer.Header().Set("Cache-Control", "no-cache, no-transform") c.Writer.Header().Del("Connection") c.Writer.Header().Del("Transfer-Encoding") c.Writer.WriteHeader(http.StatusOK) c.Writer.Flush() // 捕获稳定的引用, 不持有 *Context 指针, 以免 Context 被 pool 回收后出现竞态. fl, _ := c.Writer.(http.Flusher) reqCtx := c.Request.Context() goroutineExited := make(chan struct{}) // 写入 goroutine: 从 eventChan 消费事件并写入响应. go func() { defer close(goroutineExited) for { select { case event, ok := <-eventChan: if !ok { return } if err := event.Render(c.Writer); err != nil { return } if fl != nil { fl.Flush() } case <-reqCtx.Done(): return } } }() // 阻塞直到: // 1. 写入 goroutine 退出 (eventChan 关闭或写入失败) // 2. 客户端断开连接 (reqCtx 取消) select { case <-goroutineExited: case <-reqCtx.Done(): } }