124 lines
2.7 KiB
Go
124 lines
2.7 KiB
Go
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"))
|
||
}
|