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/db.go

@@ -118,6 +118,7 @@ CREATE TABLE IF NOT EXISTS forward_rules (
 		`ALTER TABLE forward_rules ADD COLUMN no_ufw INTEGER NOT NULL DEFAULT 0`,
 		`ALTER TABLE forward_rules ADD COLUMN label TEXT NOT NULL DEFAULT ''`,
 		`ALTER TABLE admins ADD COLUMN password_off INTEGER NOT NULL DEFAULT 0`,
+		`ALTER TABLE agents ADD COLUMN transport TEXT NOT NULL DEFAULT 'auto'`,
 	} {
 		if _, err := db.sql.Exec(alter); err != nil && !strings.Contains(err.Error(), "duplicate column") {
 			return err
@@ -354,6 +355,7 @@ type Agent struct {
 	Version   string
 	RunUser   string // OS user the agent process runs as (reported by the agent)
 	LoginUser string // OS user terminals are opened as; "" = RunUser
+	Transport string // auto, ws or lp: how the agent is to reach us
 	LastSeen  time.Time
 }
 
@@ -383,8 +385,8 @@ func (db *DB) GetAgent(id string) (*Agent, error) {
 	a := &Agent{}
 	var last int64
 	err := db.sql.QueryRow(
-		`SELECT id, name, hostname, os, arch, token_hash, version, run_user, login_user, last_seen FROM agents WHERE id = ?`, id).
-		Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.TokenHash, &a.Version, &a.RunUser, &a.LoginUser, &last)
+		`SELECT id, name, hostname, os, arch, token_hash, version, run_user, login_user, transport, last_seen FROM agents WHERE id = ?`, id).
+		Scan(&a.ID, &a.Name, &a.Hostname, &a.OS, &a.Arch, &a.TokenHash, &a.Version, &a.RunUser, &a.LoginUser, &a.Transport, &last)
 	if err == sql.ErrNoRows {
 		return nil, nil
 	}
@@ -415,6 +417,11 @@ func (db *DB) ListAgents() ([]Agent, error) {
 	return out, rows.Err()
 }
 
+func (db *DB) SetAgentTransport(id, transport string) error {
+	_, err := db.sql.Exec(`UPDATE agents SET transport = ? WHERE id = ?`, transport, id)
+	return err
+}
+
 func (db *DB) TouchAgent(id string) error {
 	_, err := db.sql.Exec(`UPDATE agents SET last_seen = ? WHERE id = ?`, time.Now().Unix(), id)
 	return err

internal/server/handlers_agent.go

@@ -53,8 +53,9 @@ func (s *Server) handlePoll(w http.ResponseWriter, r *http.Request) {
 		_ = s.db.TouchAgent(a.ID)
 	}
 	msg := s.hub.WaitControl(r.Context().Done(), a.ID, pollHold)
-	// Set after the hold so a port changed meanwhile reaches the agent now.
+	// Set after the hold so a setting changed meanwhile reaches the agent now.
 	s.setTunnelHeader(w)
+	w.Header().Set(proto.TransportHeader, s.transportOf(a.ID))
 	writeJSON(w, msg)
 }
 
@@ -71,11 +72,13 @@ func (s *Server) setTunnelHeader(w http.ResponseWriter) {
 // GET /api/agent/config - settings an agent picks up when it connects, without
 // waiting out a poll hold.
 func (s *Server) handleAgentConfig(w http.ResponseWriter, r *http.Request) {
-	if s.authenticateAgent(r) == nil {
+	a := s.authenticateAgent(r)
+	if a == nil {
 		http.Error(w, "unauthorized", http.StatusUnauthorized)
 		return
 	}
 	s.setTunnelHeader(w)
+	w.Header().Set(proto.TransportHeader, s.transportOf(a.ID))
 	w.WriteHeader(http.StatusNoContent)
 }
 

internal/server/handlers_panel.go

@@ -213,6 +213,9 @@ func (s *Server) handleAgentEdit(w http.ResponseWriter, r *http.Request) {
 		"Updated":       r.URL.Query().Get("updated") == "1",
 		"UserOK":        r.URL.Query().Get("uok"),
 		"UserErr":       r.URL.Query().Get("uerr"),
+		"Transport":     firstNonEmpty(a.Transport, proto.TransportAuto),
+		"OnWS":          s.wsAgents.On(a.ID),
+		"Online":        s.online(a),
 	}
 	// The user list lives on the machine, so it is only known while it is online.
 	if s.online(a) {
@@ -225,6 +228,21 @@ func (s *Server) handleAgentEdit(w http.ResponseWriter, r *http.Request) {
 	s.render(w, "edit.html", data)
 }
 
+// POST /agents/{id}/transport - how the agent is to reach the server: auto,
+// ws or lp. A connected agent picks it up within a poll round.
+func (s *Server) handleAgentTransport(w http.ResponseWriter, r *http.Request) {
+	if !checkCSRF(r) {
+		http.Error(w, "bad csrf", http.StatusForbidden)
+		return
+	}
+	id, t := r.PathValue("id"), r.FormValue("transport")
+	switch t {
+	case proto.TransportAuto, proto.TransportWS, proto.TransportLP:
+		_ = s.db.SetAgentTransport(id, t)
+	}
+	http.Redirect(w, r, "/agents/"+id+"/edit", http.StatusSeeOther)
+}
+
 // usersTaskTimeout bounds how long the edit page waits for the agent.
 const usersTaskTimeout = 8 * time.Second
 

internal/server/hub.go

@@ -275,6 +275,18 @@ func (s *Session) PushIn(f proto.InFrame) {
 	}
 }
 
+// Unpoll puts frames taken by PollIn back in front, for frames that could not
+// be delivered.
+func (s *Session) Unpoll(frames []proto.InFrame) {
+	s.mu.Lock()
+	s.inBuf = append(append([]proto.InFrame{}, frames...), s.inBuf...)
+	s.mu.Unlock()
+	select {
+	case s.inSignal <- struct{}{}:
+	default:
+	}
+}
+
 // PollIn drains queued downstream frames, blocking up to timeout for the first.
 func (s *Session) PollIn(timeout time.Duration) ([]proto.InFrame, bool) {
 	deadline := time.After(timeout)

internal/server/server.go

@@ -60,6 +60,7 @@ type Server struct {
 	meters     speedMeters       // byte counts of the speed tests under way
 	xfers      xfers             // files on their way between a browser and an agent
 	challenges challenges        // passkey challenges handed out
+	wsAgents   wsAgents          // agents on their WebSocket now
 	assets     map[string]string // static file -> hash of its content
 }
 
@@ -224,6 +225,7 @@ func (s *Server) Handler() http.Handler {
 	mux.Handle("POST /agents/{id}/rename", s.requireAdmin(s.handleAgentRename))
 	mux.Handle("POST /agents/{id}/update", s.requireAdmin(s.handleAgentUpdate))
 	mux.Handle("POST /agents/{id}/user", s.requireAdmin(s.handleAgentSetUser))
+	mux.Handle("POST /agents/{id}/transport", s.requireAdmin(s.handleAgentTransport))
 	mux.Handle("POST /agents/{id}/users/add", s.requireAdmin(s.handleAgentAddUser))
 	mux.Handle("POST /agents/{id}/delete", s.requireAdmin(s.handleAgentDelete))
 	mux.Handle("GET /settings", s.requireAdmin(s.handleSettings))
@@ -274,6 +276,7 @@ func (s *Server) Handler() http.Handler {
 	// Agent API (token/secret auth inside handlers).
 	mux.HandleFunc("POST /api/agent/enroll", s.handleEnroll)
 	mux.HandleFunc("GET /api/agent/poll", s.handlePoll)
+	mux.HandleFunc("GET /api/agent/ws", s.handleAgentWS)
 	mux.HandleFunc("GET /api/agent/config", s.handleAgentConfig)
 	mux.HandleFunc("GET /api/agent/term/{sid}/in", s.handleTermIn)
 	mux.HandleFunc("POST /api/agent/term/{sid}/out", s.handleTermOut)

internal/server/web/templates/edit.html

@@ -72,6 +72,25 @@
     {{end}}
   </div>
 
+  <div class="card">
+    <h2>Связь с сервером</h2>
+    <p class="muted">
+      Сейчас: {{if not .Online}}<span>офлайн</span>{{else if .OnWS}}<span style="color:var(--ok)">WebSocket</span>{{else}}<span>long polling</span>{{end}}
+    </p>
+    <form method="post" action="/agents/{{.Agent.ID}}/transport" class="row">
+      <input type="hidden" name="csrf" value="{{.CSRF}}">
+      <select name="transport">
+        <option value="auto" {{if eq .Transport "auto"}}selected{{end}}>Авто: WebSocket, а если не выходит - long polling</option>
+        <option value="ws" {{if eq .Transport "ws"}}selected{{end}}>Только WebSocket</option>
+        <option value="lp" {{if eq .Transport "lp"}}selected{{end}}>Только long polling</option>
+      </select>
+      <button type="submit">Сохранить</button>
+    </form>
+    <p class="muted" style="font-size:13px; margin-top:10px">WebSocket держит одно соединение: команды и терминал идут по нему без отдельного запроса на каждое нажатие.
+      В режиме «Авто» клиент, потеряв WebSocket, переходит на long polling и раз в минуту пробует вернуться.
+      Нужен клиент 0.1.0+ и прокси, пропускающий WebSocket (nginx: заголовки Upgrade и Connection).</p>
+  </div>
+
   <div class="card">
     <h2>Версия клиента</h2>
     {{if .Updated}}<div class="err" style="background:#1d3a24;border-color:#2b6b3a;color:#b3ffcb">Команда обновления отправлена. Клиент скачает свежий бинарник и перезапустится.</div>{{end}}

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
+		}
+	}
+}