package handler import ( "context" "io" "strings" "time" "github.com/docker/docker/api/types/container" "github.com/gofiber/fiber/v2" "github.com/gofiber/websocket/v2" "gitee.ltd/lxh/wechat-robot/internal/docker" ) // WebSocketContainerLogs 通过WebSocket推送容器日志 func WebSocketContainerLogs(c *websocket.Conn) { // 获取容器ID containerID := c.Params("id") if containerID == "" { c.Close() return } // 创建上下文,可以通过客户端关闭连接来取消 ctx, cancel := context.WithCancel(context.Background()) defer cancel() // 设置客户端关闭时取消上下文 go func() { defer cancel() for { if _, _, err := c.ReadMessage(); err != nil { // 客户端断开连接 return } } }() // 获取Docker客户端 cli := docker.GetClient() // 设置日志选项 options := container.LogsOptions{ ShowStdout: true, ShowStderr: true, Follow: true, // 持续跟踪日志 Tail: "100", // 初始返回100行 } // 获取容器日志流 logStream, err := cli.ContainerLogs(ctx, containerID, options) if err != nil { // 发送错误消息并关闭连接 c.WriteMessage(websocket.TextMessage, []byte("获取容器日志失败: "+err.Error())) c.Close() return } defer logStream.Close() // 创建缓冲区 buffer := make([]byte, 4096) // 持续读取日志并发送到WebSocket for { select { case <-ctx.Done(): // 上下文被取消,退出循环 return default: // 读取日志 n, err := logStream.Read(buffer) if err != nil { if err == io.EOF { // 日志读取完毕,等待新日志 time.Sleep(500 * time.Millisecond) continue } // 其他错误,关闭连接 return } if n > 0 { // 清理Docker日志头部的8个字节控制信息 cleanedLog := cleanDockerLogBuffer(buffer[:n]) // 发送日志到WebSocket if err := c.WriteMessage(websocket.TextMessage, cleanedLog); err != nil { // 发送失败,关闭连接 return } } } } } // ContainerLogsWebSocketMiddleware WebSocket中间件,用于升级HTTP连接到WebSocket func ContainerLogsWebSocketMiddleware() fiber.Handler { return websocket.New(WebSocketContainerLogs, websocket.Config{ ReadBufferSize: 1024, WriteBufferSize: 1024, }) } // cleanDockerLogBuffer 清理Docker日志缓冲区 func cleanDockerLogBuffer(buffer []byte) []byte { // 如果缓冲区太小,直接返回 if len(buffer) <= 8 { return buffer } // 分割日志为行 lines := strings.Split(string(buffer), "\n") for i, line := range lines { if len(line) > 8 { // 跳过每行前8个字节的头部信息 lines[i] = line[8:] } } // 合并并返回 return []byte(strings.Join(lines, "\n")) }