feat: 增强终端处理逻辑

This commit is contained in:
Akizon77
2025-11-04 19:11:24 +08:00
parent 384d3cf4bf
commit dbd8479981
3 changed files with 69 additions and 27 deletions

1
.gitignore vendored
View File

@@ -7,3 +7,4 @@ komari-agent
build/ build/
auto-discovery.json auto-discovery.json
net_static.json net_static.json
config.json

View File

@@ -3,6 +3,7 @@ package terminal
import ( import (
"encoding/json" "encoding/json"
"fmt" "fmt"
"time"
"github.com/gorilla/websocket" "github.com/gorilla/websocket"
@@ -40,41 +41,69 @@ func StartTerminal(conn *websocket.Conn) {
return return
} }
errChan := make(chan error, 1) errChan := make(chan error, 3) // 增加容量以容纳多个错误源
defer impl.term.Close() done := make(chan struct{})
defer func() {
gracefulShutdown(impl.term)
impl.term.Close()
conn.Close()
close(done)
}()
// 从 WebSocket 读取消息并写入终端 // 从 WebSocket 读取消息并写入终端
go handleWebSocketInput(conn, impl.term, errChan) go handleWebSocketInput(conn, impl.term, errChan, done)
// 从终端读取输出并写入 WebSocket // 从终端读取输出并写入 WebSocket
go handleTerminalOutput(conn, impl.term, errChan) go handleTerminalOutput(conn, impl.term, errChan, done)
// 错误处理和清理 // 等待终端进程结束或出现错误
go func() { select {
err := <-errChan case err := <-errChan:
if err != nil && conn != nil { // WebSocket 连接断开或出现错误
conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("Error: %v\r\n", err))) if err != nil {
conn.Close() conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("\r\nConnection error: %v\r\n", err)))
}
impl.term.Close()
}()
// 等待终端进程结束
if err := impl.term.Wait(); err != nil {
select {
case errChan <- err:
// 错误已发送
default:
// 错误通道已满或已关闭
conn.WriteMessage(websocket.TextMessage, []byte(fmt.Sprintf("Terminal exited with error: %v\r\n", err)))
} }
case <-done:
// 已经被其他地方关闭
} }
} }
// gracefulShutdown 尝试优雅地关闭终端
func gracefulShutdown(term Terminal) {
// Ctrl+C
for i := 0; i < 3; i++ {
if _, err := term.Write([]byte{3}); err != nil {
return
}
time.Sleep(50 * time.Millisecond)
}
time.Sleep(200 * time.Millisecond)
// Ctrl+D (EOF)
term.Write([]byte{4})
time.Sleep(100 * time.Millisecond)
term.Write([]byte("exit\n"))
time.Sleep(100 * time.Millisecond)
}
// handleWebSocketInput 处理 WebSocket 输入 // handleWebSocketInput 处理 WebSocket 输入
func handleWebSocketInput(conn *websocket.Conn, term Terminal, errChan chan<- error) { func handleWebSocketInput(conn *websocket.Conn, term Terminal, errChan chan<- error, done <-chan struct{}) {
for { for {
select {
case <-done:
return
default:
}
t, p, err := conn.ReadMessage() t, p, err := conn.ReadMessage()
if err != nil { if err != nil {
errChan <- err select {
case errChan <- err:
default:
}
return return
} }
if t == websocket.TextMessage { if t == websocket.TextMessage {
@@ -106,16 +135,28 @@ func handleWebSocketInput(conn *websocket.Conn, term Terminal, errChan chan<- er
} }
// handleTerminalOutput 处理终端输出 // handleTerminalOutput 处理终端输出
func handleTerminalOutput(conn *websocket.Conn, term Terminal, errChan chan<- error) { func handleTerminalOutput(conn *websocket.Conn, term Terminal, errChan chan<- error, done <-chan struct{}) {
buf := make([]byte, 4096) buf := make([]byte, 4096)
for { for {
select {
case <-done:
return
default:
}
n, err := term.Read(buf) n, err := term.Read(buf)
if err != nil { if err != nil {
errChan <- err select {
case errChan <- err:
default:
}
return return
} }
if err := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil { if err := conn.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
errChan <- err select {
case errChan <- err:
default:
}
return return
} }
} }

View File

@@ -29,7 +29,7 @@ func newTerminalImpl() (*terminalImpl, error) {
parts := strings.Split(line, ":") parts := strings.Split(line, ":")
if len(parts) >= 7 && parts[6] != "" { if len(parts) >= 7 && parts[6] != "" {
shell = parts[6] shell = parts[6]
log.Printf("Found shell from /etc/passwd: %s for user home: %s\n", shell, userHomeDir) //log.Printf("Found shell from /etc/passwd: %s for user home: %s\n", shell, userHomeDir)
break break
} }
} }