Support connecting dynamically to gerbil

go.mod

@@ -32,4 +32,4 @@ require (
 )
 
 // To be used ONLY for local development
-// replace github.com/fosrl/newt => ../newt
+replace github.com/fosrl/newt => ../newt

go.sum

@@ -1,7 +1,5 @@
 github.com/Microsoft/go-winio v0.6.2 h1:F2VQgta7ecxGYO8k3ZZz3RS8fVIXVxONVUPlNERoyfY=
 github.com/Microsoft/go-winio v0.6.2/go.mod h1:yd8OoFMLzJbo9gZq8j5qaps8bJ9aShtEA8Ipt1oGCvU=
-github.com/fosrl/newt v1.15.0 h1:WpL0whZM1FMjUe2Vy5jSH1bgbxm1O9k1qCyF/mqZT+s=
-github.com/fosrl/newt v1.15.0/go.mod h1:l6kWoZPSaXT+ZRUjiyPgwflRqZWYaXpUj9oQ0sOPh4o=
 github.com/godbus/dbus/v5 v5.2.2 h1:TUR3TgtSVDmjiXOgAAyaZbYmIeP3DPkld3jgKGV8mXQ=
 github.com/godbus/dbus/v5 v5.2.2/go.mod h1:3AAv2+hPq5rdnr5txxxRwiGjPXamgoIHgz9FPBfOp3c=
 github.com/google/btree v1.1.3 h1:CVpQJjYgC4VbzxeGVHfvZrv1ctoYCAI8vbl07Fcxlyg=

olm/connect.go

@@ -262,6 +262,14 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) {
 		network.SetDNSServers([]string{o.dnsProxy.GetProxyIP().String()})
 	}
 
+	if wgData.ExitNode != nil && wgData.ExitNode.Connect {
+		if err := o.connectExitNode(*wgData.ExitNode); err != nil {
+			logger.Error("Failed to connect to exit node: %v", err)
+		}
+	} else {
+		logger.Debug("No exit node to connect to (not provided, or connect flag is false)")
+	}
+
 	o.apiServer.SetRegistered(true)
 
 	o.registered = true

olm/exitnode.go

@@ -0,0 +1,182 @@
+package olm
+
+import (
+	"encoding/json"
+	"fmt"
+	"strings"
+
+	"github.com/fosrl/newt/logger"
+	"github.com/fosrl/newt/network"
+	"github.com/fosrl/newt/util"
+	"github.com/fosrl/olm/peers"
+	"github.com/fosrl/olm/websocket"
+)
+
+// connectExitNode configures a WireGuard peer connection to an exit node, on the
+// same interface and WireGuard device already used for site peers. The exit node
+// lives in a different address space than the site tunnel, so a secondary address
+// (ExitNodeConfig.TunnelIP) is added to the interface for it - the exit node's own
+// WireGuard peer entry only accepts traffic sourced from that address. Nothing here
+// is persisted; it's purely in-memory WireGuard/routing state, same as site peers.
+func (o *Olm) connectExitNode(cfg ExitNodeConfig) error {
+	if !o.tunnelRunning {
+		return fmt.Errorf("tunnel not running")
+	}
+	if cfg.PublicKey == "" || cfg.Endpoint == "" || cfg.ServerIP == "" || cfg.TunnelIP == "" {
+		return fmt.Errorf("incomplete exit node configuration")
+	}
+
+	o.exitNodeMu.Lock()
+	defer o.exitNodeMu.Unlock()
+
+	dev := o.dev
+	if dev == nil {
+		return fmt.Errorf("wireguard device not initialized")
+	}
+
+	if o.exitNode != nil && o.exitNode.PublicKey != cfg.PublicKey {
+		logger.Info("Switching exit nodes, removing previous exit node peer")
+		if err := o.removeExitNodePeerLocked(); err != nil {
+			logger.Warn("Failed to remove previous exit node peer: %v", err)
+		}
+	}
+
+	endpoint := cfg.Endpoint
+	if !strings.Contains(endpoint, ":") {
+		relayPort := cfg.RelayPort
+		if relayPort == 0 {
+			relayPort = 21820
+		}
+		endpoint = fmt.Sprintf("%s:%d", endpoint, relayPort)
+	}
+
+	resolvedEndpoint, err := util.ResolveDomain(endpoint)
+	if err != nil {
+		return fmt.Errorf("failed to resolve exit node endpoint: %w", err)
+	}
+
+	persistentKeepalive := 0
+	if pm := o.getPeerManager(); pm != nil {
+		persistentKeepalive = pm.PersistentKeepalive
+	}
+
+	allowedIP := strings.Split(cfg.ServerIP, "/")[0] + "/32"
+
+	wgConfig := fmt.Sprintf(`public_key=%s
+allowed_ip=%s
+endpoint=%s
+persistent_keepalive_interval=%d`, util.FixKey(cfg.PublicKey), allowedIP, resolvedEndpoint, persistentKeepalive)
+
+	if err := dev.IpcSet(wgConfig); err != nil {
+		return fmt.Errorf("failed to configure exit node peer: %w", err)
+	}
+
+	interfaceName := o.tunnelConfig.InterfaceName
+	tunnelIP := cfg.TunnelIP
+	if !strings.Contains(tunnelIP, "/") {
+		tunnelIP += "/32"
+	}
+	if err := network.AddSecondaryAddress(interfaceName, tunnelIP); err != nil {
+		logger.Warn("Failed to add secondary address %s for exit node: %v", tunnelIP, err)
+	}
+
+	if err := network.AddRouteForServerIP(cfg.ServerIP, interfaceName); err != nil {
+		logger.Warn("Failed to add route for exit node server IP: %v", err)
+	}
+
+	cfgCopy := cfg
+	o.exitNode = &cfgCopy
+
+	logger.Info("Connected to exit node at %s", resolvedEndpoint)
+	return nil
+}
+
+// disconnectExitNode tears down the current exit node peer connection, if any.
+func (o *Olm) disconnectExitNode() error {
+	o.exitNodeMu.Lock()
+	defer o.exitNodeMu.Unlock()
+
+	return o.removeExitNodePeerLocked()
+}
+
+// removeExitNodePeerLocked removes the current exit node peer, its secondary
+// interface address, and its server IP route. Must be called with exitNodeMu held.
+func (o *Olm) removeExitNodePeerLocked() error {
+	if o.exitNode == nil {
+		return nil
+	}
+	cfg := o.exitNode
+	o.exitNode = nil
+
+	if o.dev != nil {
+		if err := peers.RemovePeer(o.dev, 0, cfg.PublicKey); err != nil {
+			logger.Warn("Failed to remove exit node peer: %v", err)
+		}
+	}
+
+	interfaceName := o.tunnelConfig.InterfaceName
+	if err := network.RemoveRouteForServerIP(cfg.ServerIP, interfaceName); err != nil {
+		logger.Warn("Failed to remove route for exit node server IP: %v", err)
+	}
+
+	tunnelIP := cfg.TunnelIP
+	if !strings.Contains(tunnelIP, "/") {
+		tunnelIP += "/32"
+	}
+	if err := network.RemoveSecondaryAddress(interfaceName, tunnelIP); err != nil {
+		logger.Warn("Failed to remove secondary address %s for exit node: %v", tunnelIP, err)
+	}
+
+	logger.Info("Disconnected from exit node")
+	return nil
+}
+
+// handleExitNodeConnect handles a server-initiated request to connect to (or switch to)
+// an exit node, delivered as a full ExitNodeConfig payload.
+func (o *Olm) handleExitNodeConnect(msg websocket.WSMessage) {
+	logger.Debug("Received exit node connect message: %v", msg.Data)
+
+	if !o.tunnelRunning {
+		logger.Debug("Tunnel stopped, ignoring exit node connect message")
+		return
+	}
+
+	jsonData, err := json.Marshal(msg.Data)
+	if err != nil {
+		logger.Error("Error marshaling exit node connect data: %v", err)
+		return
+	}
+
+	var cfg ExitNodeConfig
+	if err := json.Unmarshal(jsonData, &cfg); err != nil {
+		logger.Error("Error unmarshaling exit node connect data: %v", err)
+		return
+	}
+
+	if !cfg.Connect {
+		logger.Debug("Exit node connect message has connect=false, disconnecting instead")
+		if err := o.disconnectExitNode(); err != nil {
+			logger.Error("Failed to disconnect from exit node: %v", err)
+		}
+		return
+	}
+
+	if err := o.connectExitNode(cfg); err != nil {
+		logger.Error("Failed to connect to exit node: %v", err)
+	}
+}
+
+// handleExitNodeDisconnect handles a server-initiated request to disconnect from the
+// currently connected exit node.
+func (o *Olm) handleExitNodeDisconnect(msg websocket.WSMessage) {
+	logger.Debug("Received exit node disconnect message: %v", msg.Data)
+
+	if !o.tunnelRunning {
+		logger.Debug("Tunnel stopped, ignoring exit node disconnect message")
+		return
+	}
+
+	if err := o.disconnectExitNode(); err != nil {
+		logger.Error("Failed to disconnect from exit node: %v", err)
+	}
+}

olm/olm.go

@@ -57,6 +57,11 @@ type Olm struct {
 	holePunchManager   *holepunch.Manager
 	peerManager        *peers.PeerManager
 	peerManagerMu      sync.RWMutex
+
+	// exitNode tracks the currently connected exit node peer, if any. It lives on a
+	// secondary address on the same interface/WireGuard device as the site peers.
+	exitNode   *ExitNodeConfig
+	exitNodeMu sync.Mutex
 	// Power mode management
 	currentPowerMode string
 	powerModeMu      sync.Mutex
@@ -557,6 +562,10 @@ func (o *Olm) StartTunnel(config TunnelConfig) {
 	o.websocket.RegisterHandler("olm/wg/peer/chain/cancel", o.handleCancelChain)
 	o.websocket.RegisterHandler("olm/sync", o.handleSync)
 
+	// Handlers for the server to direct connecting/disconnecting an exit node after registration
+	o.websocket.RegisterHandler("olm/wg/exitnode/connect", o.handleExitNodeConnect)
+	o.websocket.RegisterHandler("olm/wg/exitnode/disconnect", o.handleExitNodeDisconnect)
+
 	o.websocket.RegisterHandler("olm/ping/exitNodes", func(msg websocket.WSMessage) {
 		logger.Debug("Received exit node ping request")
 
@@ -827,6 +836,13 @@ func (o *Olm) Close() {
 	}
 	o.peerManagerMu.Unlock()
 
+	// The WireGuard device and TUN interface are being torn down below, which takes
+	// the exit node peer and its secondary address with them - just clear the
+	// in-memory record so a stale config isn't reused on the next connect.
+	o.exitNodeMu.Lock()
+	o.exitNode = nil
+	o.exitNodeMu.Unlock()
+
 	if o.uapiListener != nil {
 		_ = o.uapiListener.Close()
 		o.uapiListener = nil

olm/types.go

@@ -10,6 +10,23 @@ type WgData struct {
 	Sites         []peers.SiteConfig `json:"sites"`
 	TunnelIP      string             `json:"tunnelIP"`
 	UtilitySubnet string             `json:"utilitySubnet"` // this is for things like the DNS server, and alias addresses
+	ExitNode      *ExitNodeConfig    `json:"exitNode,omitempty"`
+}
+
+// ExitNodeConfig describes an exit node the olm client can connect to for
+// resources (e.g. inference) hosted on that node, separate from the site
+// peers. It lives in a different address space than the site tunnel - the
+// client is assigned TunnelIP (within the exit node's subnet) to reach the
+// node at ServerIP. It arrives on the initial "olm/wg/connect" message and can
+// also be sent later via "olm/wg/exitnode/connect" / "olm/wg/exitnode/disconnect"
+// so the server can direct a client to connect/disconnect after registration.
+type ExitNodeConfig struct {
+	Connect   bool   `json:"connect"`
+	Endpoint  string `json:"endpoint"`
+	RelayPort uint16 `json:"relayPort"`
+	PublicKey string `json:"publicKey"`
+	ServerIP  string `json:"serverIP"`
+	TunnelIP  string `json:"tunnelIP"`
 }
 
 type SyncData struct {