mirror of
https://github.com/fankes/komari-agent.git
synced 2025-12-10 15:23:35 +08:00
feat: 增强终端处理逻辑
This commit is contained in:
1
.gitignore
vendored
1
.gitignore
vendored
@@ -7,3 +7,4 @@ komari-agent
|
|||||||
build/
|
build/
|
||||||
auto-discovery.json
|
auto-discovery.json
|
||||||
net_static.json
|
net_static.json
|
||||||
|
config.json
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
@@ -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
|
||||||
}
|
}
|
||||||
}
|
}
|
||||||
|
|||||||
Reference in New Issue
Block a user