119 lines
2.3 KiB
Go
119 lines
2.3 KiB
Go
package tunnel
|
|
|
|
import (
|
|
"fmt"
|
|
"io"
|
|
"log"
|
|
"net"
|
|
"net/url"
|
|
"strings"
|
|
"sync"
|
|
"time"
|
|
|
|
"github.com/gorilla/websocket"
|
|
)
|
|
|
|
type Tunnel struct {
|
|
wsURL string
|
|
port int
|
|
sessionID string
|
|
ws *websocket.Conn
|
|
tcp net.Conn
|
|
stop chan struct{}
|
|
once sync.Once
|
|
}
|
|
|
|
// New returns a tunnel that pipes traffic between the Fleet Forge server
|
|
// (WebSocket) and the local VNC server (127.0.0.1:port). It does not connect
|
|
// until Run is called.
|
|
func New(serverURL, credential, sessionID string, vncPort int) *Tunnel {
|
|
wsURL := strings.Replace(serverURL, "https://", "wss://", 1)
|
|
wsURL = strings.Replace(wsURL, "http://", "ws://", 1)
|
|
wsURL = strings.TrimSuffix(wsURL, "/")
|
|
|
|
q := url.Values{}
|
|
q.Set("session_id", sessionID)
|
|
q.Set("token", credential)
|
|
|
|
full := wsURL + "/api/v1/remote/agent?" + q.Encode()
|
|
|
|
return &Tunnel{
|
|
wsURL: full,
|
|
port: vncPort,
|
|
sessionID: sessionID,
|
|
stop: make(chan struct{}),
|
|
}
|
|
}
|
|
|
|
// Run dials the server WebSocket and the local VNC port, then pipes bytes in
|
|
// both directions until either side closes or Stop is called.
|
|
func (t *Tunnel) Run() error {
|
|
ws, _, err := websocket.DefaultDialer.Dial(t.wsURL, nil)
|
|
if err != nil {
|
|
return fmt.Errorf("tunnel dial (server): %w", err)
|
|
}
|
|
defer ws.Close()
|
|
t.ws = ws
|
|
|
|
tcp, err := net.DialTimeout("tcp", fmt.Sprintf("127.0.0.1:%d", t.port), 10*time.Second)
|
|
if err != nil {
|
|
return fmt.Errorf("tunnel dial (local VNC): %w", err)
|
|
}
|
|
defer tcp.Close()
|
|
t.tcp = tcp
|
|
|
|
log.Printf("[tunnel] connected for session %s (vnc 127.0.0.1:%d)", t.sessionID, t.port)
|
|
|
|
done := make(chan error, 2)
|
|
|
|
go func() {
|
|
for {
|
|
_, data, err := ws.ReadMessage()
|
|
if err != nil {
|
|
done <- err
|
|
return
|
|
}
|
|
if _, err := tcp.Write(data); err != nil {
|
|
done <- err
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
go func() {
|
|
buf := make([]byte, 8192)
|
|
for {
|
|
n, err := tcp.Read(buf)
|
|
if n > 0 {
|
|
if err := ws.WriteMessage(websocket.BinaryMessage, buf[:n]); err != nil {
|
|
done <- err
|
|
return
|
|
}
|
|
}
|
|
if err != nil {
|
|
done <- err
|
|
return
|
|
}
|
|
}
|
|
}()
|
|
|
|
select {
|
|
case err := <-done:
|
|
return err
|
|
case <-t.stop:
|
|
return io.EOF
|
|
}
|
|
}
|
|
|
|
// Stop closes the underlying connections, causing Run to return.
|
|
func (t *Tunnel) Stop() {
|
|
t.once.Do(func() {
|
|
close(t.stop)
|
|
if t.ws != nil {
|
|
_ = t.ws.Close()
|
|
}
|
|
if t.tcp != nil {
|
|
_ = t.tcp.Close()
|
|
}
|
|
})
|
|
} |