lpra/internal/server/wsagent.go

package server

import (
	"context"
	"net/http"
	"sync"
	"time"

	"github.com/coder/websocket"
	"github.com/coder/websocket/wsjson"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// An agent's WebSocket carries what its long polls would: control messages
// pushed the moment they are queued, and the input and output of its
// terminal and http-forward sessions, without a request per keystroke. Tasks,
// file transfers and the speed test stay on plain requests.

// wsPing keeps the connection alive through proxies that drop quiet ones
// (Cloudflare after 100 s).
const wsPing = 20 * time.Second

// wsAgents are the agents connected by WebSocket now, for the panel to show.
type wsAgents struct {
	mu sync.Mutex
	m  map[string]*websocket.Conn
}

// set records the agent's connection, closing one it had before.
func (c *wsAgents) set(id string, conn *websocket.Conn) {
	c.mu.Lock()
	if c.m == nil {
		c.m = map[string]*websocket.Conn{}
	}
	old := c.m[id]
	c.m[id] = conn
	c.mu.Unlock()
	if old != nil {
		old.Close(websocket.StatusGoingAway, "replaced")
	}
}

func (c *wsAgents) drop(id string, conn *websocket.Conn) {
	c.mu.Lock()
	if c.m[id] == conn {
		delete(c.m, id)
	}
	c.mu.Unlock()
}

// On reports whether the agent is on its WebSocket now.
func (c *wsAgents) On(id string) bool {
	c.mu.Lock()
	defer c.mu.Unlock()
	return c.m[id] != nil
}

// transportOf is the agent's setting as stored, "auto" when unset.
func (s *Server) transportOf(agentID string) string {
	if a, err := s.db.GetAgent(agentID); err == nil && a != nil && a.Transport != "" {
		return a.Transport
	}
	return proto.TransportAuto
}

// GET /api/agent/ws - the agent's WebSocket. Refused (409) when the agent is
// set to long polling.
func (s *Server) handleAgentWS(w http.ResponseWriter, r *http.Request) {
	a := s.authenticateAgent(r)
	if a == nil {
		http.Error(w, "unauthorized", http.StatusUnauthorized)
		return
	}
	if a.Transport == proto.TransportLP {
		w.Header().Set(proto.TransportHeader, proto.TransportLP)
		http.Error(w, "long polling only", http.StatusConflict)
		return
	}
	if v := r.Header.Get("X-Agent-Version"); v != "" {
		_ = s.db.TouchAgentInfo(a.ID, v, firstNonEmpty(r.Header.Get("X-Agent-User"), a.RunUser))
	}
	conn, err := websocket.Accept(w, r, &websocket.AcceptOptions{CompressionMode: websocket.CompressionDisabled})
	if err != nil {
		return
	}
	conn.SetReadLimit(4 << 20)
	ctx, cancel := context.WithCancel(r.Context())
	defer cancel()
	s.wsAgents.set(a.ID, conn)
	defer s.wsAgents.drop(a.ID, conn)
	defer conn.CloseNow()

	send := func(m proto.WSMsg) error {
		wctx, wcancel := context.WithTimeout(ctx, 30*time.Second)
		defer wcancel()
		return wsjson.Write(wctx, conn, m)
	}
	cfg := func() proto.WSMsg {
		return proto.WSMsg{T: "cfg", Tunnel: s.tunnelAdvert(), Transport: s.transportOf(a.ID)}
	}
	if send(cfg()) != nil {
		return
	}

	// Control: what a poll would get, as soon as it is queued. An idle round
	// sends the settings, which doubles as a sign of life.
	go func() {
		defer cancel()
		for ctx.Err() == nil {
			msg := s.hub.WaitControl(ctx.Done(), a.ID, pollHold)
			if ctx.Err() != nil {
				return
			}
			_ = s.db.TouchAgent(a.ID)
			out := proto.WSMsg{T: "ctrl", Ctrl: &msg}
			if msg.Type == proto.CtrlIdle {
				out = cfg()
			}
			if send(out) != nil {
				return
			}
			if out.Transport == proto.TransportLP {
				conn.Close(websocket.StatusNormalClosure, "long polling only")
				return
			}
		}
	}()
	go func() {
		t := time.NewTicker(wsPing)
		defer t.Stop()
		for {
			select {
			case <-ctx.Done():
				return
			case <-t.C:
				pctx, pcancel := context.WithTimeout(ctx, 15*time.Second)
				err := conn.Ping(pctx)
				pcancel()
				if err != nil {
					cancel()
					return
				}
			}
		}
	}()

	subs := map[string]bool{}
	for {
		var m proto.WSMsg
		if err := wsjson.Read(ctx, conn, &m); err != nil {
			return
		}
		switch m.T {
		case "sub":
			if subs[m.SID] {
				continue
			}
			sess := s.hub.Session(m.SID)
			if sess == nil || sess.AgentID != a.ID {
				_ = send(proto.WSMsg{T: "gone", SID: m.SID})
				continue
			}
			subs[m.SID] = true
			go s.pumpSession(ctx, sess, send)
		case "out":
			sess := s.hub.Session(m.SID)
			if sess == nil || sess.AgentID != a.ID || m.Out == nil {
				_ = send(proto.WSMsg{T: "gone", SID: m.SID})
				continue
			}
			if !sess.PushOut(*m.Out) {
				_ = send(proto.WSMsg{T: "gone", SID: m.SID})
				continue
			}
			if m.Out.Exit {
				sess.close()
				time.AfterFunc(pollHold, func() { s.hub.CloseSession(sess.ID) })
			}
		}
	}
}

// pumpSession sends a session's input down the WebSocket as it comes. Input
// taken just as the connection ends goes back to the session, for the agent's
// long poll to collect.
func (s *Server) pumpSession(ctx context.Context, sess *Session, send func(proto.WSMsg) error) {
	for ctx.Err() == nil {
		frames, alive := sess.PollIn(pollHold)
		if !alive {
			_ = send(proto.WSMsg{T: "gone", SID: sess.ID})
			return
		}
		if len(frames) == 0 {
			continue
		}
		if ctx.Err() != nil || send(proto.WSMsg{T: "in", SID: sess.ID, Frames: frames}) != nil {
			sess.Unpoll(frames)
			return
		}
	}
}