hamlab/internal/nodes/nodes.go

// Package nodes is the hub side of multi-host: the list of managed Ham Labs, their
// health, adding one (pair with its password, receive a token, pin its certificate),
// removing one (revoke the token), and the reverse proxy that turns
// /api/nodes/<name>/… into /api/… on that node, WebSockets included.
package nodes

import (
	"bytes"
	"context"
	"crypto/tls"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"net"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"

	"hamlab/internal/config"
	"hamlab/internal/tlsauto"
)

var nameRe = regexp.MustCompile(`^[a-z0-9][a-z0-9-]{0,31}$`)

type Status struct {
	Name        string    `json:"name"`
	URL         string    `json:"url"`
	Fingerprint string    `json:"fingerprint,omitempty"`
	Added       time.Time `json:"added"`
	Online      bool      `json:"online"`
	Host        string    `json:"host,omitempty"`
	Version     string    `json:"version,omitempty"`
	Load        string    `json:"load,omitempty"`
	Uptime      float64   `json:"uptime,omitempty"`
	RebootNeed  bool      `json:"reboot_needed,omitempty"`
	LastSeen    time.Time `json:"last_seen,omitempty"`
	Error       string    `json:"error,omitempty"`
}

type Manager struct {
	cfg     *config.Config
	hubName string
	mu      sync.Mutex
	status  map[string]*Status
	proxies map[string]*httputil.ReverseProxy
	clients map[string]*http.Client
}

func New(cfg *config.Config, hubName string) *Manager {
	return &Manager{cfg: cfg, hubName: hubName, status: map[string]*Status{}, proxies: map[string]*httputil.ReverseProxy{}, clients: map[string]*http.Client{}}
}

func transportFor(n config.Node) *http.Transport {
	t := &http.Transport{
		DialContext:           (&net.Dialer{Timeout: 5 * time.Second}).DialContext,
		TLSHandshakeTimeout:   5 * time.Second,
		ResponseHeaderTimeout: 60 * time.Second,
		MaxIdleConnsPerHost:   8,
	}
	if strings.HasPrefix(n.URL, "https://") {
		if n.Fingerprint != "" {
			t.TLSClientConfig = tlsauto.PinnedConfig(n.Fingerprint)
		} else {
			t.TLSClientConfig = &tls.Config{InsecureSkipVerify: true}
		}
	}
	return t
}

func (m *Manager) client(n config.Node) *http.Client {
	m.mu.Lock()
	defer m.mu.Unlock()
	if c, ok := m.clients[n.Name]; ok {
		return c
	}
	c := &http.Client{Transport: transportFor(n), Timeout: 15 * time.Second}
	m.clients[n.Name] = c
	return c
}

func (m *Manager) node(name string) (config.Node, bool) {
	var found config.Node
	ok := false
	m.cfg.View(func(c *config.Config) {
		for _, n := range c.Nodes {
			if n.Name == name {
				found, ok = n, true
			}
		}
	})
	return found, ok
}

// ---- probing ----

func (m *Manager) probe(ctx context.Context, n config.Node) *Status {
	st := &Status{Name: n.Name, URL: n.URL, Fingerprint: n.Fingerprint, Added: n.Added}
	req, _ := http.NewRequestWithContext(ctx, "GET", strings.TrimRight(n.URL, "/")+"/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+n.Token)
	resp, err := m.client(n).Do(req)
	if err != nil {
		st.Error = shortErr(err)
		return st
	}
	defer resp.Body.Close()
	if resp.StatusCode != 200 {
		st.Error = fmt.Sprintf("HTTP %d", resp.StatusCode)
		if resp.StatusCode == 401 {
			st.Error = "token rejected (re-add the machine)"
		}
		return st
	}
	var me struct {
		Host, Version, Load string
		Uptime              float64
		RebootNeeded        bool `json:"reboot_needed"`
	}
	if json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&me) != nil {
		st.Error = "bad reply"
		return st
	}
	st.Online, st.Host, st.Version, st.Load, st.Uptime, st.RebootNeed, st.LastSeen = true, me.Host, me.Version, me.Load, me.Uptime, me.RebootNeeded, time.Now()
	return st
}

func shortErr(err error) string {
	s := err.Error()
	if i := strings.LastIndex(s, ": "); i >= 0 {
		s = s[i+2:]
	}
	return s
}

// Run refreshes every node's status periodically.
func (m *Manager) Run(ctx context.Context) {
	m.refresh(ctx)
	t := time.NewTicker(15 * time.Second)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
			m.refresh(ctx)
		}
	}
}

func (m *Manager) refresh(ctx context.Context) {
	var list []config.Node
	m.cfg.View(func(c *config.Config) { list = append(list, c.Nodes...) })
	var wg sync.WaitGroup
	for _, n := range list {
		wg.Add(1)
		go func(n config.Node) {
			defer wg.Done()
			pctx, cancel := context.WithTimeout(ctx, 8*time.Second)
			defer cancel()
			st := m.probe(pctx, n)
			m.mu.Lock()
			if prev, ok := m.status[n.Name]; ok && !st.Online && prev.Online {
				st.LastSeen = prev.LastSeen // keep the last good sighting
			}
			m.status[n.Name] = st
			m.mu.Unlock()
		}(n)
	}
	wg.Wait()
}

// List returns every managed node with its last known status.
func (m *Manager) List() []Status {
	var list []config.Node
	m.cfg.View(func(c *config.Config) { list = append(list, c.Nodes...) })
	m.mu.Lock()
	defer m.mu.Unlock()
	out := make([]Status, 0, len(list))
	for _, n := range list {
		if st, ok := m.status[n.Name]; ok {
			out = append(out, *st)
		} else {
			out = append(out, Status{Name: n.Name, URL: n.URL, Fingerprint: n.Fingerprint, Added: n.Added, Error: "not probed yet"})
		}
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Name < out[j].Name })
	return out
}

// ---- add / remove ----

// Add pairs with a node: pin its certificate (trust on first use), trade the node's
// password for a token named after this hub, and store the node. The password is
// used once and never kept.
func (m *Manager) Add(ctx context.Context, name, rawURL, username, password string) (Status, error) {
	if !nameRe.MatchString(name) {
		return Status{}, errors.New("name: lowercase letters, digits and dashes only")
	}
	if _, exists := m.node(name); exists {
		return Status{}, errors.New("a machine with that name already exists")
	}
	u, err := url.Parse(strings.TrimSpace(rawURL))
	if err != nil || (u.Scheme != "http" && u.Scheme != "https") || u.Host == "" {
		return Status{}, errors.New("address must be http(s)://host:port")
	}
	n := config.Node{Name: name, URL: u.Scheme + "://" + u.Host, Added: time.Now()}
	if u.Scheme == "https" {
		host := u.Host
		if _, _, err := net.SplitHostPort(host); err != nil {
			host += ":443"
		}
		fp, err := tlsauto.Peek(host)
		if err != nil {
			return Status{}, fmt.Errorf("connect: %s", shortErr(err))
		}
		n.Fingerprint = fp
	}
	body, _ := json.Marshal(map[string]string{"username": username, "password": password, "name": "hub:" + m.hubName})
	req, _ := http.NewRequestWithContext(ctx, "POST", n.URL+"/api/tokens", bytes.NewReader(body))
	req.Header.Set("Content-Type", "application/json")
	resp, err := (&http.Client{Transport: transportFor(n), Timeout: 15 * time.Second}).Do(req)
	if err != nil {
		return Status{}, fmt.Errorf("connect: %s", shortErr(err))
	}
	defer resp.Body.Close()
	var reply struct {
		Token string `json:"token"`
		Error string `json:"error"`
	}
	_ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&reply)
	if resp.StatusCode != 200 || reply.Token == "" {
		if reply.Error == "" {
			reply.Error = fmt.Sprintf("HTTP %d", resp.StatusCode)
		}
		return Status{}, errors.New("pairing refused: " + reply.Error)
	}
	n.Token = reply.Token
	if err := m.cfg.Update(func(c *config.Config) { c.Nodes = append(c.Nodes, n) }); err != nil {
		return Status{}, err
	}
	st := m.probe(ctx, n)
	m.mu.Lock()
	m.status[name] = st
	m.mu.Unlock()
	return *st, nil
}

// Remove revokes our token on the node (best effort) and forgets it.
func (m *Manager) Remove(ctx context.Context, name string) error {
	n, ok := m.node(name)
	if !ok {
		return errors.New("no such machine")
	}
	req, _ := http.NewRequestWithContext(ctx, "DELETE", n.URL+"/api/tokens/"+url.PathEscape("hub:"+m.hubName), nil)
	req.Header.Set("Authorization", "Bearer "+n.Token)
	if resp, err := m.client(n).Do(req); err == nil {
		resp.Body.Close()
	}
	err := m.cfg.Update(func(c *config.Config) {
		kept := c.Nodes[:0]
		for _, x := range c.Nodes {
			if x.Name != name {
				kept = append(kept, x)
			}
		}
		c.Nodes = kept
	})
	m.mu.Lock()
	delete(m.status, name)
	delete(m.proxies, name)
	delete(m.clients, name)
	m.mu.Unlock()
	return err
}

// ---- update from hub ----

type PushResult struct {
	Name    string `json:"name"`
	OK      bool   `json:"ok"`
	Version string `json:"version,omitempty"`
	Error   string `json:"error,omitempty"`
}

// Push uploads binPath (normally the hub's own executable) to the node's self-update
// endpoint. The node verifies it, installs it and restarts itself. Refused when the
// node's architecture differs from ours.
func (m *Manager) Push(ctx context.Context, name, binPath, ourArch string, force bool) PushResult {
	res := PushResult{Name: name}
	n, ok := m.node(name)
	if !ok {
		res.Error = "no such machine"
		return res
	}
	// arch check via /api/me
	req, _ := http.NewRequestWithContext(ctx, "GET", n.URL+"/api/me", nil)
	req.Header.Set("Authorization", "Bearer "+n.Token)
	resp, err := m.client(n).Do(req)
	if err != nil {
		res.Error = "unreachable: " + shortErr(err)
		return res
	}
	var me struct {
		Arch    string `json:"arch"`
		Version string `json:"version"`
	}
	_ = json.NewDecoder(io.LimitReader(resp.Body, 64<<10)).Decode(&me)
	resp.Body.Close()
	if me.Arch != "" && ourArch != "" && me.Arch != ourArch {
		// look for a sibling build: /usr/local/bin/hamlab-arm64 next to /usr/local/bin/hamlab
		alt := binPath + "-" + me.Arch
		if _, err := os.Stat(alt); err != nil {
			res.Error = fmt.Sprintf("node is %s but only a %s build is installed here (expected %s; run make deploy, which builds it)", me.Arch, ourArch, alt)
			return res
		}
		binPath = alt
	}
	f, err := os.Open(binPath)
	if err != nil {
		res.Error = err.Error()
		return res
	}
	defer f.Close()
	st, _ := f.Stat()
	target := n.URL + "/api/self-update"
	if force {
		target += "?force=1"
	}
	up, _ := http.NewRequestWithContext(ctx, "POST", target, f)
	up.ContentLength = st.Size()
	up.Header.Set("Authorization", "Bearer "+n.Token)
	up.Header.Set("Content-Type", "application/octet-stream")
	c := &http.Client{Transport: transportFor(n), Timeout: 5 * time.Minute}
	r2, err := c.Do(up)
	if err != nil {
		res.Error = "upload: " + shortErr(err)
		return res
	}
	defer r2.Body.Close()
	var reply struct {
		Version string `json:"version"`
		Error   string `json:"error"`
	}
	_ = json.NewDecoder(io.LimitReader(r2.Body, 64<<10)).Decode(&reply)
	if r2.StatusCode != 200 {
		if reply.Error == "" {
			reply.Error = fmt.Sprintf("HTTP %d", r2.StatusCode)
		}
		res.Error = reply.Error
		return res
	}
	res.OK, res.Version = true, reply.Version
	return res
}

// PushAll updates every node, in parallel.
func (m *Manager) PushAll(ctx context.Context, binPath, ourArch string, force bool) []PushResult {
	var list []config.Node
	m.cfg.View(func(c *config.Config) { list = append(list, c.Nodes...) })
	out := make([]PushResult, len(list))
	var wg sync.WaitGroup
	for i, n := range list {
		wg.Add(1)
		go func(i int, name string) {
			defer wg.Done()
			out[i] = m.Push(ctx, name, binPath, ourArch, force)
		}(i, n.Name)
	}
	wg.Wait()
	return out
}

// ---- proxy ----

// Proxy returns the handler for /api/nodes/<name>/<rest>: it rewrites to <node>/api/<rest>,
// swaps the hub's cookie for the node's bearer token, and drops Origin so the node's
// WebSocket origin check (meant for browsers) doesn't reject the hub.
// Do makes one request to a managed machine as this hub and returns the whole response.
func (m *Manager) Do(ctx context.Context, name, method, path string, body []byte, hdr http.Header) (int, []byte, error) {
	n, ok := m.node(name)
	if !ok {
		return 0, nil, errors.New("no such machine")
	}
	req, err := http.NewRequestWithContext(ctx, method, strings.TrimRight(n.URL, "/")+path, bytes.NewReader(body))
	if err != nil {
		return 0, nil, err
	}
	for k, v := range hdr {
		req.Header[k] = v
	}
	req.Header.Set("Authorization", "Bearer "+n.Token)
	resp, err := m.client(n).Do(req)
	if err != nil {
		return 0, nil, errors.New(shortErr(err))
	}
	defer resp.Body.Close()
	b, err := io.ReadAll(io.LimitReader(resp.Body, 8<<20))
	return resp.StatusCode, b, err
}

// URL is a managed machine's address, "" for one not managed here.
func (m *Manager) URL(name string) string {
	n, _ := m.node(name)
	return n.URL
}

// Names lists the managed machines.
func (m *Manager) Names() []string {
	var out []string
	m.cfg.View(func(c *config.Config) {
		for _, n := range c.Nodes {
			out = append(out, n.Name)
		}
	})
	return out
}

func (m *Manager) Proxy(name string) (http.Handler, bool) {
	n, ok := m.node(name)
	if !ok {
		return nil, false
	}
	m.mu.Lock()
	if p, ok := m.proxies[name]; ok {
		m.mu.Unlock()
		return p, true
	}
	m.mu.Unlock()
	target, _ := url.Parse(n.URL)
	prefix := "/api/nodes/" + name
	p := &httputil.ReverseProxy{
		Director: func(r *http.Request) {
			r.URL.Scheme, r.URL.Host = target.Scheme, target.Host
			r.URL.Path = "/api" + strings.TrimPrefix(r.URL.Path, prefix)
			r.URL.RawPath = ""
			r.Host = target.Host
			r.Header.Set("Authorization", "Bearer "+n.Token)
			r.Header.Del("Cookie")
			r.Header.Del("Origin")
			r.Header.Del("X-Forwarded-For")
			r.Header.Del("X-Forwarded-Proto")
		},
		Transport:     transportFor(n),
		FlushInterval: -1, // stream task output / log tails immediately
		// The node answers with its own policy headers and this hub adds the same ones to
		// everything it serves; relayed twice, Access-Control-Allow-Origin is invalid and a
		// browser on another origin refuses the response. The hub's copy is the one that counts.
		ModifyResponse: func(resp *http.Response) error {
			for _, h := range []string{"Access-Control-Allow-Origin", "X-Frame-Options", "Content-Security-Policy", "X-Content-Type-Options", "Referrer-Policy"} {
				resp.Header.Del(h)
			}
			return nil
		},
		ErrorHandler: func(w http.ResponseWriter, r *http.Request, err error) {
			w.Header().Set("Content-Type", "application/json")
			w.WriteHeader(http.StatusBadGateway)
			_ = json.NewEncoder(w).Encode(map[string]string{"error": name + " unreachable: " + shortErr(err)})
		},
	}
	m.mu.Lock()
	m.proxies[name] = p
	m.mu.Unlock()
	return p, true
}