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/agent/client.go

@@ -18,6 +18,7 @@ import (
 	"path/filepath"
 	"runtime"
 	"sync"
+	"sync/atomic"
 	"time"
 
 	lpra "github.com/x-grk/longpolling-remote-access"
@@ -53,6 +54,8 @@ type Agent struct {
 	pollCount int
 
 	// Direct tunnel currently running (see setTunnelAddr).
+	transport atomic.Value // string: auto, ws or lp, as the server says
+	link      atomic.Pointer[wsLink]
 	tunMu     sync.Mutex
 	tunAddr   string
 	tunCancel context.CancelFunc
@@ -107,10 +110,22 @@ func (a *Agent) Run(ctx context.Context) error {
 		a.setTunnelAddr(ctx, a.cfg.TunnelAddr)
 	}
 
+	// The control channel: a WebSocket, or long polls. In "auto" a failed or
+	// dropped WebSocket gives way to long polling for a minute, then is tried
+	// again; "ws" keeps trying the WebSocket; "lp" only polls.
 	backoff := time.Second
 	for ctx.Err() == nil {
-		if err := a.pollOnce(ctx); err != nil {
-			log.Printf("poll: %v (retry in %s)", err, backoff)
+		if a.mode() == proto.TransportLP {
+			a.longPoll(ctx, 0)
+			continue
+		}
+		err := a.runWS(ctx)
+		if ctx.Err() != nil {
+			return nil
+		}
+		switch a.mode() {
+		case proto.TransportWS:
+			log.Printf("websocket: %v (retry in %s)", err, backoff)
 			select {
 			case <-time.After(backoff):
 			case <-ctx.Done():
@@ -120,10 +135,71 @@ func (a *Agent) Run(ctx context.Context) error {
 				backoff *= 2
 			}
 			continue
+		case proto.TransportAuto:
+			log.Printf("websocket: %v - long polling, websocket again in a minute", err)
+		}
+		backoff = time.Second
+		a.longPoll(ctx, time.Minute)
+	}
+	return nil
+}
+
+// mode is the transport the server set for this agent.
+func (a *Agent) mode() string {
+	if v, ok := a.transport.Load().(string); ok {
+		return v
+	}
+	return proto.TransportAuto
+}
+
+func (a *Agent) applyTransport(v string) {
+	switch v {
+	case proto.TransportAuto, proto.TransportWS, proto.TransportLP:
+		if v != a.mode() {
+			log.Printf("transport: %s", v)
+			a.transport.Store(v)
+		}
+	}
+}
+
+// longPoll runs control polls: for d, or with d zero until the transport is
+// no longer long polling.
+func (a *Agent) longPoll(ctx context.Context, d time.Duration) {
+	until := time.Now().Add(d)
+	backoff := time.Second
+	for ctx.Err() == nil {
+		if d > 0 && (time.Now().After(until) || a.mode() == proto.TransportWS) {
+			return
+		}
+		if d == 0 && a.mode() != proto.TransportLP {
+			return
+		}
+		// In "auto" the minute is a minute: a poll still held then is cut
+		// short (the server keeps its job for the next one).
+		reqCtx, cancel := ctx, context.CancelFunc(func() {})
+		if d > 0 {
+			reqCtx, cancel = context.WithDeadline(ctx, until)
+		}
+		err := a.pollOnce(ctx, reqCtx)
+		cut := reqCtx.Err() == context.DeadlineExceeded
+		cancel()
+		if err != nil && cut && ctx.Err() == nil {
+			return
+		}
+		if err != nil {
+			log.Printf("poll: %v (retry in %s)", err, backoff)
+			select {
+			case <-time.After(backoff):
+			case <-ctx.Done():
+				return
+			}
+			if backoff < 15*time.Second {
+				backoff *= 2
+			}
+			continue
 		}
 		backoff = time.Second
 	}
-	return nil
 }
 
 func (a *Agent) enroll(ctx context.Context) error {
@@ -157,8 +233,10 @@ func (a *Agent) enroll(ctx context.Context) error {
 }
 
 // pollOnce performs one control long-poll and dispatches any instruction.
-func (a *Agent) pollOnce(ctx context.Context) error {
-	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.Host+"/api/agent/poll", nil)
+// The poll itself runs on reqCtx, which may end early; what it starts runs on
+// ctx.
+func (a *Agent) pollOnce(ctx, reqCtx context.Context) error {
+	req, _ := http.NewRequestWithContext(reqCtx, http.MethodGet, a.cfg.Host+"/api/agent/poll", nil)
 	a.authHeaders(req)
 	// Report the version and our OS user occasionally (not every poll), so the
 	// panel can flag outdated agents without extra traffic.
@@ -183,7 +261,14 @@ func (a *Agent) pollOnce(ctx context.Context) error {
 		return err
 	}
 	a.applyTunnelHeader(ctx, resp.Header)
-	// Dispatch in the background so we keep polling for more work.
+	a.applyTransport(resp.Header.Get(proto.TransportHeader))
+	a.dispatch(ctx, msg)
+	return nil
+}
+
+// dispatch starts what a control message asks for, in the background, so the
+// control channel keeps listening.
+func (a *Agent) dispatch(ctx context.Context, msg proto.ControlMessage) {
 	switch msg.Type {
 	case proto.CtrlOpenTerminal:
 		go a.runTerminal(ctx, msg.SessionID, msg.User)
@@ -192,7 +277,6 @@ func (a *Agent) pollOnce(ctx context.Context) error {
 	case proto.CtrlTask:
 		go a.runTask(ctx, msg.TaskID)
 	}
-	return nil
 }
 
 // fetchConfig asks the server for the settings it hands to agents and reports
@@ -207,13 +291,17 @@ func (a *Agent) fetchConfig(ctx context.Context) bool {
 		return false
 	}
 	defer resp.Body.Close()
+	a.applyTransport(resp.Header.Get(proto.TransportHeader))
 	return a.applyTunnelHeader(ctx, resp.Header)
 }
 
 // applyTunnelHeader follows the tunnel address the server advertises as
 // "host:port"; an empty host means the host we already reach the server on.
 func (a *Agent) applyTunnelHeader(ctx context.Context, h http.Header) bool {
-	v := h.Get("X-Tunnel-Addr")
+	return a.applyTunnelAddr(ctx, h.Get("X-Tunnel-Addr"))
+}
+
+func (a *Agent) applyTunnelAddr(ctx context.Context, v string) bool {
 	if v == "" || a.cfg.TunnelPinned {
 		return false
 	}

internal/agent/session.go

@@ -54,6 +54,11 @@ func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWrite
 	ctx, cancel := context.WithCancel(parent)
 	defer cancel()
 	defer rw.Close()
+	defer func() {
+		if l := a.link.Load(); l != nil {
+			l.unsubscribe(sid)
+		}
+	}()
 	// A session closed from the panel (a tab reloaded, a window shut) ends
 	// the stream at once, not on its next output: an idle shell would live
 	// on until the agent stopped.
@@ -88,6 +93,21 @@ func (a *Agent) bridgeStream(parent context.Context, sid string, rw io.ReadWrite
 // pollIn performs one downstream long-poll and applies the frames.
 // Returns false when the session is gone and the loop should stop.
 func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize func(cols, rows uint16)) bool {
+	// On the WebSocket the input comes pushed; if it drops, by a long poll.
+	if l := a.link.Load(); l != nil {
+		if ch := l.subscribe(sid); ch != nil {
+			select {
+			case in := <-ch:
+				if in.gone {
+					return false
+				}
+				return applyFrames(in.frames, w, onResize)
+			case <-l.done:
+			case <-ctx.Done():
+				return false
+			}
+		}
+	}
 	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, a.cfg.Host+"/api/agent/term/"+sid+"/in", nil)
 	a.authHeaders(req)
 	resp, err := a.http.Do(req)
@@ -105,7 +125,13 @@ func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize fu
 	if err := json.NewDecoder(resp.Body).Decode(&batch); err != nil {
 		return true
 	}
-	for _, f := range batch.Frames {
+	return applyFrames(batch.Frames, w, onResize)
+}
+
+// applyFrames writes a session's input to its stream; false once the panel
+// closed it.
+func applyFrames(frames []proto.InFrame, w io.Writer, onResize func(cols, rows uint16)) bool {
+	for _, f := range frames {
 		switch f.Type {
 		case proto.FrameStdin:
 			if data, err := b64.DecodeString(f.Data); err == nil {
@@ -124,6 +150,11 @@ func (a *Agent) pollIn(ctx context.Context, sid string, w io.Writer, onResize fu
 
 // postOut sends one upstream message. Returns false if the session is gone.
 func (a *Agent) postOut(ctx context.Context, sid string, msg proto.OutMessage) bool {
+	if l := a.link.Load(); l != nil {
+		if l.send(proto.WSMsg{T: "out", SID: sid, Out: &msg}) == nil {
+			return true
+		}
+	}
 	body, _ := json.Marshal(msg)
 	req, _ := http.NewRequestWithContext(ctx, http.MethodPost, a.cfg.Host+"/api/agent/term/"+sid+"/out", bytes.NewReader(body))
 	a.authHeaders(req)

internal/agent/wslink.go

@@ -0,0 +1,172 @@
+package agent
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"log"
+	"net/http"
+	"strings"
+	"sync"
+	"time"
+
+	"github.com/coder/websocket"
+	"github.com/coder/websocket/wsjson"
+
+	"github.com/x-grk/longpolling-remote-access/internal/proto"
+)
+
+// The WebSocket to the server: control messages arrive on it as soon as they
+// are queued, and terminal and http-forward sessions send and take their
+// bytes over it instead of a request each. When it drops, sessions carry on
+// over plain requests.
+
+const wsPing = 20 * time.Second
+
+type wsLink struct {
+	c     *websocket.Conn
+	ctx   context.Context
+	mu    sync.Mutex
+	subs  map[string]chan wsIn
+	asked map[string]bool // sessions whose input was asked for
+	done  chan struct{}
+}
+
+// wsIn is what the server sent for a session: its input, or that it is over.
+type wsIn struct {
+	frames []proto.InFrame
+	gone   bool
+}
+
+func (l *wsLink) send(m proto.WSMsg) error {
+	ctx, cancel := context.WithTimeout(l.ctx, 30*time.Second)
+	defer cancel()
+	return wsjson.Write(ctx, l.c, m)
+}
+
+// subscribe asks the server for a session's input on this link, once, and
+// returns where it arrives; nil when the link cannot take it.
+func (l *wsLink) subscribe(sid string) chan wsIn {
+	l.mu.Lock()
+	ch := l.subs[sid]
+	if ch == nil {
+		ch = make(chan wsIn, 64)
+		l.subs[sid] = ch
+	}
+	fresh := !l.subscribed(sid)
+	l.mu.Unlock()
+	if fresh && l.send(proto.WSMsg{T: "sub", SID: sid}) != nil {
+		return nil
+	}
+	return ch
+}
+
+// subscribed marks a session as asked for; l.mu is held.
+func (l *wsLink) subscribed(sid string) bool {
+	if l.asked == nil {
+		l.asked = map[string]bool{}
+	}
+	if l.asked[sid] {
+		return true
+	}
+	l.asked[sid] = true
+	return false
+}
+
+func (l *wsLink) unsubscribe(sid string) {
+	l.mu.Lock()
+	delete(l.subs, sid)
+	delete(l.asked, sid)
+	l.mu.Unlock()
+}
+
+func (l *wsLink) deliver(sid string, in wsIn) {
+	l.mu.Lock()
+	ch := l.subs[sid]
+	l.mu.Unlock()
+	if ch == nil {
+		return
+	}
+	select {
+	case ch <- in:
+	case <-l.done:
+	}
+}
+
+// runWS holds the WebSocket until it ends, dispatching what comes on it.
+func (a *Agent) runWS(ctx context.Context) error {
+	u := a.cfg.Host + "/api/agent/ws"
+	u = "ws" + strings.TrimPrefix(u, "http")
+	hdr := http.Header{}
+	hdr.Set("X-Agent-Id", a.state.AgentID)
+	hdr.Set("Authorization", "Bearer "+a.state.Token)
+	hdr.Set("X-Agent-Version", Version)
+	hdr.Set("X-Agent-User", currentUser())
+	dctx, dcancel := context.WithTimeout(ctx, 15*time.Second)
+	c, resp, err := websocket.Dial(dctx, u, &websocket.DialOptions{
+		HTTPClient: a.stream, HTTPHeader: hdr, CompressionMode: websocket.CompressionDisabled,
+	})
+	dcancel()
+	if err != nil {
+		if resp != nil {
+			a.applyTransport(resp.Header.Get(proto.TransportHeader))
+			return fmt.Errorf("server answered %s", resp.Status)
+		}
+		return err
+	}
+	c.SetReadLimit(8 << 20)
+	lctx, lcancel := context.WithCancel(ctx)
+	l := &wsLink{c: c, ctx: lctx, subs: map[string]chan wsIn{}, done: make(chan struct{})}
+	a.link.Store(l)
+	defer func() {
+		a.link.CompareAndSwap(l, nil)
+		close(l.done)
+		lcancel()
+		c.CloseNow()
+	}()
+	log.Printf("connected by websocket")
+	go func() {
+		t := time.NewTicker(wsPing)
+		defer t.Stop()
+		for {
+			select {
+			case <-lctx.Done():
+				return
+			case <-t.C:
+				pctx, pcancel := context.WithTimeout(lctx, 15*time.Second)
+				err := c.Ping(pctx)
+				pcancel()
+				if err != nil {
+					lcancel()
+					return
+				}
+			}
+		}
+	}()
+	for {
+		var m proto.WSMsg
+		if err := wsjson.Read(lctx, c, &m); err != nil {
+			if errors.Is(err, context.Canceled) && ctx.Err() == nil {
+				return errors.New("no answer to ping")
+			}
+			return err
+		}
+		switch m.T {
+		case "ctrl":
+			if m.Ctrl != nil {
+				// Sessions live on the agent's context: they outlast the link.
+				a.dispatch(ctx, *m.Ctrl)
+			}
+		case "cfg":
+			a.applyTunnelAddr(ctx, m.Tunnel)
+			a.applyTransport(m.Transport)
+			if a.mode() == proto.TransportLP {
+				return errors.New("set to long polling")
+			}
+		case "in":
+			l.deliver(m.SID, wsIn{frames: m.Frames})
+		case "gone":
+			l.deliver(m.SID, wsIn{gone: true})
+		}
+	}
+}