lpra/internal/agent/wslink.go

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{}
	stop  context.CancelFunc
}

// close drops the link: the agent then connects again, to the address in
// use by then.
func (l *wsLink) close() { l.stop() }

// 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.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{}), stop: lcancel}
	a.link.Store(l)
	defer func() {
		a.link.CompareAndSwap(l, nil)
		close(l.done)
		lcancel()
		c.CloseNow()
	}()
	log.Printf("connected by websocket")
	a.reached()
	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.learnFallback(m.Fallback)
			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})
		}
	}
}