f0dd810541
* removed the Async utility * now using Ping message to webclient for keep alive instaed of actual content * added remote shell to AgentInfo * retry of connections to the agent * better logging for SynchronizeStreams
86 lines
1.8 KiB
Go
86 lines
1.8 KiB
Go
package websocketutil
|
|
|
|
import (
|
|
"github.com/gorilla/websocket"
|
|
"net"
|
|
"time"
|
|
)
|
|
|
|
type WebSocketConn struct {
|
|
conn *websocket.Conn
|
|
buf []byte
|
|
text bool
|
|
}
|
|
|
|
func NewWebSocketConn(conn *websocket.Conn, text bool) *WebSocketConn {
|
|
return &WebSocketConn{conn: conn, text: text}
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) Read(p []byte) (n int, err error) {
|
|
if len(websocketConn.buf) == 0 {
|
|
_, message, err := websocketConn.conn.ReadMessage()
|
|
if err != nil {
|
|
return 0, err
|
|
}
|
|
websocketConn.buf = message
|
|
}
|
|
|
|
n = copy(p, websocketConn.buf)
|
|
websocketConn.buf = websocketConn.buf[n:]
|
|
|
|
return n, err
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) Write(p []byte) (n int, err error) {
|
|
messageType := websocket.BinaryMessage
|
|
switch {
|
|
case len(p) == 0:
|
|
messageType = websocket.PingMessage
|
|
case websocketConn.text:
|
|
messageType = websocket.TextMessage
|
|
}
|
|
err = websocketConn.conn.WriteMessage(messageType, p)
|
|
if err == nil {
|
|
n = len(p)
|
|
}
|
|
return n, err
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) Close() error {
|
|
return websocketConn.conn.Close()
|
|
}
|
|
|
|
func (WebSocketConn *WebSocketConn) LocalAddr() net.Addr {
|
|
return WebSocketConn.conn.LocalAddr()
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) RemoteAddr() net.Addr {
|
|
return websocketConn.conn.RemoteAddr()
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) SetDeadline(t time.Time) error {
|
|
return nil
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) SetReadDeadline(t time.Time) error {
|
|
return nil
|
|
}
|
|
|
|
func (websocketConn *WebSocketConn) SetWriteDeadline(t time.Time) error {
|
|
return nil
|
|
}
|
|
|
|
func ConnectWebSocket(conn net.Conn, urlStr string) (net.Conn, error) {
|
|
dialer := *websocket.DefaultDialer
|
|
dialer.NetDial = func(network, addr string) (net.Conn, error) {
|
|
return conn, nil
|
|
}
|
|
|
|
wsConn, _, err := dialer.Dial(urlStr, nil)
|
|
if err != nil {
|
|
return nil, err
|
|
}
|
|
|
|
return NewWebSocketConn(wsConn, false), nil
|
|
}
|