lpra/internal/agent/hosts.go

package agent

import (
	"context"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/x-grk/longpolling-remote-access/internal/proto"
)

// The server can be reached at two addresses: the main one (behind
// Cloudflare, say) and a direct one that goes to the server itself. The
// direct one comes from the installer or, for agents installed before it was
// set, from the server, and is kept in the state file. When the address in
// use has been out of reach for a while the agent tries the other; on the
// direct one it keeps checking the main one and goes back when it answers.

const (
	// switchAfter is how long the address in use may fail before the other
	// is tried.
	switchAfter = 30 * time.Second
	// mainCheck is how often the main address is checked from the direct one.
	mainCheck = 5 * time.Minute
)

type hostState struct {
	mu        sync.Mutex
	cur       string
	downSince time.Time
	changed   bool // the address changed: the WebSocket goes there at once
}

// host is the server address in use now.
func (a *Agent) host() string {
	a.hosts.mu.Lock()
	defer a.hosts.mu.Unlock()
	if a.hosts.cur == "" {
		return a.cfg.Host
	}
	return a.hosts.cur
}

// fallback is the direct address: the installer's, else the one the server
// told.
func (a *Agent) fallback() string {
	if a.cfg.Fallback != "" {
		return a.cfg.Fallback
	}
	a.stateMu.Lock()
	defer a.stateMu.Unlock()
	return a.state.Fallback
}

// learnFallback keeps the direct address the server tells.
func (a *Agent) learnFallback(v string) {
	v = strings.TrimRight(strings.TrimSpace(v), "/")
	if v == "" || a.cfg.Fallback != "" {
		return
	}
	a.stateMu.Lock()
	same := a.state.Fallback == v
	if !same {
		a.state.Fallback = v
	}
	a.stateMu.Unlock()
	if !same {
		log.Printf("direct address: %s", v)
		_ = a.saveState()
	}
}

// setHost moves to addr: requests go there from now on, and the WebSocket
// is dropped to reconnect there.
func (a *Agent) setHost(addr string) {
	a.hosts.mu.Lock()
	if a.hosts.cur == addr || (a.hosts.cur == "" && addr == a.cfg.Host) {
		a.hosts.mu.Unlock()
		return
	}
	a.hosts.cur, a.hosts.downSince, a.hosts.changed = addr, time.Time{}, true
	a.hosts.mu.Unlock()
	log.Printf("server address: %s", addr)
	if l := a.link.Load(); l != nil {
		l.close()
	}
}

// hostChanged reports, once, that the address changed.
func (a *Agent) hostChanged() bool {
	a.hosts.mu.Lock()
	defer a.hosts.mu.Unlock()
	c := a.hosts.changed
	a.hosts.changed = false
	return c
}

// reached notes the server answered; failed notes it did not, and after
// switchAfter of that moves to the other address.
func (a *Agent) reached() {
	a.hosts.mu.Lock()
	a.hosts.downSince = time.Time{}
	a.hosts.mu.Unlock()
}

func (a *Agent) failed() {
	fb := a.fallback()
	if fb == "" || fb == a.cfg.Host {
		return
	}
	a.hosts.mu.Lock()
	now := time.Now()
	if a.hosts.downSince.IsZero() {
		a.hosts.downSince = now
		a.hosts.mu.Unlock()
		return
	}
	long := now.Sub(a.hosts.downSince) >= switchAfter
	cur := a.hosts.cur
	a.hosts.mu.Unlock()
	if !long {
		return
	}
	next := fb
	if cur == fb {
		next = a.cfg.Host
	}
	log.Printf("server out of reach for %s, trying %s", switchAfter, next)
	a.setHost(next)
}

// watchMain moves back to the main address once it answers again.
func (a *Agent) watchMain(ctx context.Context) {
	t := time.NewTicker(mainCheck)
	defer t.Stop()
	for {
		select {
		case <-ctx.Done():
			return
		case <-t.C:
		}
		if a.host() == a.cfg.Host {
			continue
		}
		if a.reachableAt(ctx, a.cfg.Host) {
			log.Printf("main address answers again")
			a.setHost(a.cfg.Host)
		}
	}
}

// reachableAt reports whether the server answers a quick request at base.
func (a *Agent) reachableAt(ctx context.Context, base string) bool {
	ctx, cancel := context.WithTimeout(ctx, 10*time.Second)
	defer cancel()
	req, _ := http.NewRequestWithContext(ctx, http.MethodGet, base+"/api/agent/config", nil)
	a.authHeaders(req)
	resp, err := a.http.Do(req)
	if err != nil {
		return false
	}
	resp.Body.Close()
	if resp.StatusCode < 300 {
		a.learnFallback(resp.Header.Get(proto.FallbackHeader))
		return true
	}
	return false
}