feat: agents connect by WebSocket, falling back to long polling and trying again each minute; auto, ws or lp per machine (0.1.0)

internal/server/wsagent.go

@@ -0,0 +1,202 @@
+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
+		}
+	}
+}