mirror of
https://github.com/fankes/komari-agent.git
synced 2025-10-18 18:49:23 +08:00
feat: 添加安全的WebSocket连接管理,支持ping任务功能
This commit is contained in:
58
ws/safaConn.go
Normal file
58
ws/safaConn.go
Normal file
@@ -0,0 +1,58 @@
|
||||
package ws
|
||||
|
||||
import (
|
||||
"sync"
|
||||
"time"
|
||||
|
||||
"github.com/gorilla/websocket"
|
||||
)
|
||||
|
||||
type SafeConn struct {
|
||||
conn *websocket.Conn
|
||||
mu sync.Mutex
|
||||
}
|
||||
|
||||
func NewSafeConn(conn *websocket.Conn) *SafeConn {
|
||||
return &SafeConn{
|
||||
conn: conn,
|
||||
mu: sync.Mutex{},
|
||||
}
|
||||
}
|
||||
|
||||
func (sc *SafeConn) WriteMessage(messageType int, data []byte) error {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn.WriteMessage(messageType, data)
|
||||
}
|
||||
|
||||
func (sc *SafeConn) WriteJSON(v interface{}) error {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn.WriteJSON(v)
|
||||
}
|
||||
|
||||
func (sc *SafeConn) Close() error {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn.Close()
|
||||
}
|
||||
func (sc *SafeConn) ReadMessage() (int, []byte, error) {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn.ReadMessage()
|
||||
}
|
||||
func (sc *SafeConn) ReadJSON(v interface{}) error {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn.ReadJSON(v)
|
||||
}
|
||||
func (sc *SafeConn) SetReadDeadline(t time.Time) error {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn.SetReadDeadline(t)
|
||||
}
|
||||
func (sc *SafeConn) GetConn() *websocket.Conn {
|
||||
sc.mu.Lock()
|
||||
defer sc.mu.Unlock()
|
||||
return sc.conn
|
||||
}
|
Reference in New Issue
Block a user