diff --git a/config.go b/config.go index 5959270..01fcad1 100644 --- a/config.go +++ b/config.go @@ -45,6 +45,7 @@ type OlmConfig struct { OverrideDNS bool `json:"overrideDNS"` TunnelDNS bool `json:"tunnelDNS"` DisableRelay bool `json:"disableRelay"` + WebSocketRelay bool `json:"websocketRelay"` // DoNotCreateNewClient bool `json:"doNotCreateNewClient"` // Parsed values (not in JSON) @@ -110,6 +111,7 @@ func DefaultConfig() *OlmConfig { config.sources["overrideDNS"] = string(SourceDefault) config.sources["tunnelDNS"] = string(SourceDefault) config.sources["disableRelay"] = string(SourceDefault) + config.sources["websocketRelay"] = string(SourceDefault) // config.sources["doNotCreateNewClient"] = string(SourceDefault) return config @@ -269,6 +271,10 @@ func loadConfigFromEnv(config *OlmConfig) { config.DisableRelay = true config.sources["disableRelay"] = string(SourceEnv) } + if val := os.Getenv("WEBSOCKET_RELAY"); val == "true" { + config.WebSocketRelay = true + config.sources["websocketRelay"] = string(SourceEnv) + } if val := os.Getenv("TUNNEL_DNS"); val == "true" { config.TunnelDNS = true config.sources["tunnelDNS"] = string(SourceEnv) @@ -303,6 +309,7 @@ func loadConfigFromCLI(config *OlmConfig, args []string) (bool, bool, error) { "disableHolepunch": config.DisableHolepunch, "overrideDNS": config.OverrideDNS, "disableRelay": config.DisableRelay, + "websocketRelay": config.WebSocketRelay, "tunnelDNS": config.TunnelDNS, // "doNotCreateNewClient": config.DoNotCreateNewClient, } @@ -327,6 +334,7 @@ func loadConfigFromCLI(config *OlmConfig, args []string) (bool, bool, error) { serviceFlags.BoolVar(&config.DisableHolepunch, "disable-holepunch", config.DisableHolepunch, "Disable hole punching") serviceFlags.BoolVar(&config.OverrideDNS, "override-dns", config.OverrideDNS, "When enabled, the client uses custom DNS servers to resolve internal resources and aliases. This overrides your system's default DNS settings. Queries that cannot be resolved as a Pangolin resource will be forwarded to your configured Upstream DNS Server. (default false)") serviceFlags.BoolVar(&config.DisableRelay, "disable-relay", config.DisableRelay, "Disable relay connections") + serviceFlags.BoolVar(&config.WebSocketRelay, "websocket-relay", config.WebSocketRelay, "Force WebSocket relay transport without waiting for fallback") serviceFlags.BoolVar(&config.TunnelDNS, "tunnel-dns", config.TunnelDNS, "When enabled, DNS queries are routed through the tunnel for remote resolution. To ensure queries are tunneled correctly, you must define the DNS server as a Pangolin resource and enter its address as an Upstream DNS Server. (default false)") // serviceFlags.BoolVar(&config.DoNotCreateNewClient, "do-not-create-new-client", config.DoNotCreateNewClient, "Do not create new client") @@ -403,6 +411,9 @@ func loadConfigFromCLI(config *OlmConfig, args []string) (bool, bool, error) { if config.DisableRelay != origValues["disableRelay"].(bool) { config.sources["disableRelay"] = string(SourceCLI) } + if config.WebSocketRelay != origValues["websocketRelay"].(bool) { + config.sources["websocketRelay"] = string(SourceCLI) + } if config.TunnelDNS != origValues["tunnelDNS"].(bool) { config.sources["tunnelDNS"] = string(SourceCLI) } @@ -530,6 +541,10 @@ func mergeConfigs(dest, src *OlmConfig) { dest.DisableRelay = src.DisableRelay dest.sources["disableRelay"] = string(SourceFile) } + if src.WebSocketRelay { + dest.WebSocketRelay = src.WebSocketRelay + dest.sources["websocketRelay"] = string(SourceFile) + } // if src.DoNotCreateNewClient { // dest.DoNotCreateNewClient = src.DoNotCreateNewClient // dest.sources["doNotCreateNewClient"] = string(SourceFile) @@ -621,6 +636,7 @@ func (c *OlmConfig) ShowConfig() { fmt.Printf(" override-dns = %v [%s]\n", c.OverrideDNS, getSource("overrideDNS")) fmt.Printf(" tunnel-dns = %v [%s]\n", c.TunnelDNS, getSource("tunnelDNS")) fmt.Printf(" disable-relay = %v [%s]\n", c.DisableRelay, getSource("disableRelay")) + fmt.Printf(" websocket-relay = %v [%s]\n", c.WebSocketRelay, getSource("websocketRelay")) // fmt.Printf(" do-not-create-new-client = %v [%s]\n", c.DoNotCreateNewClient, getSource("doNotCreateNewClient")) if c.TlsClientCert != "" { fmt.Printf(" tls-cert = %s [%s]\n", c.TlsClientCert, getSource("tlsClientCert")) diff --git a/main.go b/main.go index 7e2b9bf..8780b14 100644 --- a/main.go +++ b/main.go @@ -249,6 +249,7 @@ func runOlmMainWithArgs(ctx context.Context, cancel context.CancelFunc, signalCt OrgID: config.OrgID, OverrideDNS: config.OverrideDNS, DisableRelay: config.DisableRelay, + WebSocketRelay: config.WebSocketRelay, EnableUAPI: true, } go olm.StartTunnel(tunnelConfig) diff --git a/olm/connect.go b/olm/connect.go index 9f80ce3..0b2ccb8 100644 --- a/olm/connect.go +++ b/olm/connect.go @@ -56,6 +56,13 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) { logger.Info("Got new message. Closing existing tunnel!") o.dev.Close() } + o.relaySwitchMu.Lock() + if o.wsRelayBind != nil { + _ = o.wsRelayBind.Shutdown() + o.wsRelayBind = nil + } + o.wssRelayActive = false + o.relaySwitchMu.Unlock() jsonData, err := json.Marshal(msg.Data) if err != nil { @@ -96,8 +103,21 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) { o.middleDev = olmDevice.NewMiddleDevice(o.tdev) wgLogger := logger.GetLogger().GetWireGuardLogger("wireguard: ") - // Use filtered device instead of raw TUN device - o.dev = device.NewDevice(o.middleDev, o.sharedBind, (*device.Logger)(wgLogger)) + // Use filtered device instead of raw TUN device. + if o.tunnelConfig.WebSocketRelay { + relayURL, err := o.resolveRelayTunnelURL() + if err != nil { + logger.Error("Forced websocket relay mode requires relay URL: %v", err) + return + } + o.wsRelayBind = NewWebSocketRelayBind(relayURL, nil) + o.wssRelayActive = true + o.dev = device.NewDevice(o.middleDev, o.wsRelayBind, (*device.Logger)(wgLogger)) + logger.Info("Using websocket relay bind for forced relay mode") + } else { + o.dev = device.NewDevice(o.middleDev, o.sharedBind, (*device.Logger)(wgLogger)) + o.wssRelayActive = false + } if o.tunnelConfig.EnableUAPI { fileUAPI, err := func() (*os.File, error) { @@ -135,7 +155,12 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) { } if err = o.dev.Up(); err != nil { - logger.Error("Failed to bring up WireGuard device: %v", err) + logger.Error( + "Failed to bring up WireGuard device: %v (holepunch=%v)", + err, + o.tunnelConfig.Holepunch, + ) + return } // Extract interface IP (strip CIDR notation if present) @@ -246,6 +271,11 @@ func (o *Olm) handleConnect(msg websocket.WSMessage) { o.apiServer.SetRegistered(true) o.registered = true + if o.tunnelConfig.WebSocketRelay { + logger.Debug("Forced websocket relay mode enabled; runtime relay switch timer disabled") + } else { + logger.Debug("WebSocket relay forced mode disabled; keeping UDP transport until fallback triggers") + } // Start ping monitor now that we are registered and connected o.websocket.StartPingMonitor() diff --git a/olm/olm.go b/olm/olm.go index 07d144f..7441f80 100644 --- a/olm/olm.go +++ b/olm/olm.go @@ -8,6 +8,7 @@ import ( "net" "net/http" _ "net/http/pprof" + "net/url" "os" "sync" "time" @@ -41,6 +42,7 @@ type Olm struct { tdev tun.Device middleDev *olmDevice.MiddleDevice sharedBind *bind.SharedBind + wsRelayBind *WebSocketRelayBind dnsProxy *dns.DNSProxy apiServer *api.API @@ -68,13 +70,22 @@ type Olm struct { stopRegister func() updateRegister func(newData any) - stopPeerSends map[string]func() - stopPeerInits map[string]func() + stopPeerSends map[string]func() + stopPeerInits map[string]func() jitPendingSites map[int]string // siteId -> chainId for in-flight JIT requests - peerSendMu sync.Mutex + peerSendMu sync.Mutex // WaitGroup to track tunnel lifecycle tunnelWg sync.WaitGroup + + // Cached websocket relay URL for WSS / TCP fallback relay modes + relayTunnelURL string + relayURLMu sync.RWMutex + relayTimerMu sync.Mutex + relayTimers map[int]*time.Timer + relaySwitchMu sync.Mutex + relaySwitchTimer *time.Timer + wssRelayActive bool } // getPeerManager safely returns the current peerManager under a read-lock. @@ -138,6 +149,21 @@ func generateChainId() string { return hex.EncodeToString(b) } +func buildRelayTunnelURL(baseEndpointWss, token, userToken string) (string, error) { + u, err := url.Parse(baseEndpointWss) + if err != nil { + return "", fmt.Errorf("invalid relay endpoint wss URL: %w", err) + } + + q := u.Query() + q.Set("token", token) + if userToken != "" { + q.Set("userToken", userToken) + } + u.RawQuery = q.Encode() + return u.String(), nil +} + func Init(ctx context.Context, config OlmConfig) (*Olm, error) { logger.GetLogger().SetLevel(util.ParseLogLevel(config.LogLevel)) @@ -188,13 +214,14 @@ func Init(ctx context.Context, config OlmConfig) (*Olm, error) { apiServer.SetAgent(config.Agent) newOlm := &Olm{ - logFile: logFile, - olmCtx: ctx, - apiServer: apiServer, - olmConfig: config, + logFile: logFile, + olmCtx: ctx, + apiServer: apiServer, + olmConfig: config, stopPeerSends: make(map[string]func()), stopPeerInits: make(map[string]func()), jitPendingSites: make(map[int]string), + relayTimers: make(map[int]*time.Timer), } newOlm.registerAPICallbacks() @@ -457,10 +484,10 @@ func (o *Olm) StartTunnel(config TunnelConfig) { } if o.stopRegister == nil { - logger.Debug("Sending registration message to server with public key: %s and relay: %v", publicKey, !config.Holepunch) + logger.Debug("Sending registration message to server with public key: %s and relay: false", publicKey) o.stopRegister, o.updateRegister = o.websocket.SendMessageInterval("olm/wg/register", map[string]any{ "publicKey": publicKey.String(), - "relay": !config.Holepunch, + "relay": false, "olmVersion": o.olmConfig.Version, "olmAgent": o.olmConfig.Agent, "orgId": config.OrgID, @@ -479,9 +506,33 @@ func (o *Olm) StartTunnel(config TunnelConfig) { }) o.websocket.OnTokenUpdate(func(token string, exitNodes []websocket.ExitNode) { - // Check if tunnel is still running and hole punch manager exists - if !o.tunnelRunning || o.holePunchManager == nil { - logger.Debug("Tunnel stopped or hole punch manager nil, ignoring token update") + // Check if tunnel is still running + if !o.tunnelRunning { + logger.Debug("Tunnel stopped, ignoring token update") + return + } + + var relayEndpointWss string + for i := range exitNodes { + relayEndpointWss = exitNodes[i].RelayEndpointWss + if relayEndpointWss != "" { + break + } + } + if relayEndpointWss != "" { + relayURL, err := buildRelayTunnelURL(relayEndpointWss, token, userToken) + if err != nil { + logger.Warn("Failed to build relay tunnel URL: %v", err) + } else { + o.relayURLMu.Lock() + o.relayTunnelURL = relayURL + o.relayURLMu.Unlock() + logger.Debug("Cached websocket relay tunnel URL for WSS relay bind") + } + } + + if o.holePunchManager == nil { + logger.Debug("Hole punch manager nil, ignoring token update") return } @@ -563,6 +614,10 @@ func (o *Olm) Close() { o.stopRegister() o.stopRegister = nil } + o.cancelRelaySwitchTimer() + o.relaySwitchMu.Lock() + o.wssRelayActive = false + o.relaySwitchMu.Unlock() // Stop all pending peer init and send senders before closing websocket o.peerSendMu.Lock() @@ -652,10 +707,14 @@ func (o *Olm) Close() { // Now close WireGuard device - its TUN reader should have exited by now // This will call sharedBind.Close() which releases WireGuard's reference if o.dev != nil { + if o.wsRelayBind != nil { + _ = o.wsRelayBind.Shutdown() + } logger.Debug("Closing WireGuard device") o.dev.Close() o.dev = nil } + o.wsRelayBind = nil // Release the hole punch reference to the shared bind (WireGuard already // released its reference via dev.Close()) diff --git a/olm/peer.go b/olm/peer.go index 9c1c165..2ff9e7e 100644 --- a/olm/peer.go +++ b/olm/peer.go @@ -2,6 +2,7 @@ package olm import ( "encoding/json" + "fmt" "time" "github.com/fosrl/newt/holepunch" @@ -11,6 +12,70 @@ import ( "github.com/fosrl/olm/websocket" ) +func (o *Olm) cancelRelayTCPFallback(siteID int) { + o.relayTimerMu.Lock() + defer o.relayTimerMu.Unlock() + if t, ok := o.relayTimers[siteID]; ok { + t.Stop() + delete(o.relayTimers, siteID) + } +} + +func (o *Olm) scheduleRelayTCPFallback(siteID int, relayEndpointWss string) { + if relayEndpointWss == "" { + return + } + + o.relayTimerMu.Lock() + if t, ok := o.relayTimers[siteID]; ok { + t.Stop() + } + delay := 8 * time.Second + o.relayTimers[siteID] = time.AfterFunc(delay, func() { + o.relayTimerMu.Lock() + delete(o.relayTimers, siteID) + o.relayTimerMu.Unlock() + + if !o.tunnelRunning || o.peerManager == nil { + return + } + if o.peerManager.IsPeerConnected(siteID) { + logger.Info("Peer %d connected over UDP relay before websocket fallback timer", siteID) + return + } + if o.wsRelayBind != nil { + return + } + + logger.Warn("Peer %d still disconnected after UDP relay fallback delay (%s), switching to websocket relay transport", siteID, delay) + if err := o.switchToWebsocketRelayBind(relayEndpointWss); err != nil { + logger.Error("Failed to switch to websocket relay transport: %v", err) + } + }) + o.relayTimerMu.Unlock() +} + +func (o *Olm) switchToWebsocketRelayBind(relayEndpointWss string) error { + if o.websocket == nil { + return fmt.Errorf("websocket relay switch prerequisites missing") + } + + token, _ := o.websocket.GetTokenState() + if token == "" { + return fmt.Errorf("cached websocket token unavailable for relay switch") + } + + relayURL, err := buildRelayTunnelURL(relayEndpointWss, token, o.tunnelConfig.UserToken) + if err != nil { + return err + } + o.relayURLMu.Lock() + o.relayTunnelURL = relayURL + o.relayURLMu.Unlock() + + return o.switchToWebSocketRelayWithURL(relayURL) +} + func (o *Olm) handleWgPeerAdd(msg websocket.WSMessage) { logger.Debug("Received add-peer message: %v", msg.Data) @@ -25,7 +90,6 @@ func (o *Olm) handleWgPeerAdd(msg websocket.WSMessage) { logger.Debug("Ignoring add-peer message: peerManager is nil (shutdown in progress)") return } - jsonData, err := json.Marshal(msg.Data) if err != nil { logger.Error("Error marshaling data: %v", err) @@ -57,11 +121,26 @@ func (o *Olm) handleWgPeerAdd(msg websocket.WSMessage) { o.stopPeerSends = make(map[string]func()) o.peerSendMu.Unlock() } - if siteConfigMsg.PublicKey == "" { logger.Warn("Skipping add-peer for site %d (%s): no public key available (site may not be connected)", siteConfigMsg.SiteId, siteConfigMsg.Name) return } + if siteConfigMsg.ChainId != "" { + o.peerSendMu.Lock() + if stop, ok := o.stopPeerSends[siteConfigMsg.ChainId]; ok { + stop() + delete(o.stopPeerSends, siteConfigMsg.ChainId) + } + o.peerSendMu.Unlock() + } else { + // stop all of the stopPeerSends + o.peerSendMu.Lock() + for _, stop := range o.stopPeerSends { + stop() + } + o.stopPeerSends = make(map[string]func()) + o.peerSendMu.Unlock() + } _ = o.holePunchManager.TriggerHolePunch() // Trigger immediate hole punch attempt so that if the peer decides to relay we have already punched close to when we need it @@ -69,7 +148,6 @@ func (o *Olm) handleWgPeerAdd(msg websocket.WSMessage) { logger.Error("Failed to add peer: %v", err) return } - logger.Info("Successfully added peer for site %d", siteConfigMsg.SiteId) } @@ -226,6 +304,14 @@ func (o *Olm) handleWgPeerRelay(msg websocket.WSMessage) { // Update HTTP server to mark this peer as using relay o.apiServer.UpdatePeerRelayStatus(relayData.SiteId, relayData.RelayEndpoint, true) + if relayData.RelayEndpointWss != "" { + logger.Info("Received websocket relay endpoint for site %d: %s", relayData.SiteId, relayData.RelayEndpointWss) + pm.RelayPeerWss(relayData.SiteId, relayData.RelayEndpointWss) + if !o.tunnelConfig.WebSocketRelay { + o.scheduleRelayTCPFallback(relayData.SiteId, relayData.RelayEndpointWss) + } + } + pm.RelayPeer(relayData.SiteId, primaryRelay, relayData.RelayPort) } @@ -266,6 +352,7 @@ func (o *Olm) handleWgPeerUnrelay(msg websocket.WSMessage) { // Update HTTP server to mark this peer as using relay o.apiServer.UpdatePeerRelayStatus(relayData.SiteId, relayData.Endpoint, false) + o.cancelRelayTCPFallback(relayData.SiteId) pm.UnRelayPeer(relayData.SiteId, primaryRelay) } diff --git a/olm/relay_switch.go b/olm/relay_switch.go new file mode 100644 index 0000000..61e779d --- /dev/null +++ b/olm/relay_switch.go @@ -0,0 +1,208 @@ +package olm + +import ( + "fmt" + "time" + + "github.com/fosrl/newt/logger" + "github.com/fosrl/newt/util" + "github.com/fosrl/olm/peers" + wgConn "golang.zx2c4.com/wireguard/conn" + "golang.zx2c4.com/wireguard/device" +) + +func (o *Olm) getRelayTunnelURL() string { + o.relayURLMu.RLock() + defer o.relayURLMu.RUnlock() + return o.relayTunnelURL +} + +func (o *Olm) resolveRelayTunnelURL() (string, error) { + if cached := o.getRelayTunnelURL(); cached != "" { + return cached, nil + } + if o.websocket == nil { + return "", fmt.Errorf("websocket client is nil") + } + + token, exitNodes := o.websocket.GetTokenState() + if token == "" { + return "", fmt.Errorf("missing relay token") + } + + var relayEndpointWss string + for _, node := range exitNodes { + if node.RelayEndpointWss != "" { + relayEndpointWss = node.RelayEndpointWss + break + } + } + if relayEndpointWss == "" { + return "", fmt.Errorf("missing relay endpoint wss") + } + + relayURL, err := buildRelayTunnelURL(relayEndpointWss, token, o.websocket.GetUserToken()) + if err != nil { + return "", err + } + + o.relayURLMu.Lock() + o.relayTunnelURL = relayURL + o.relayURLMu.Unlock() + + return relayURL, nil +} + +func (o *Olm) requestRelayPreflight(pm *peers.PeerManager) { + monitor := pm.GetPeerMonitor() + if monitor == nil { + return + } + for _, peer := range pm.GetAllPeers() { + if monitor.IsPeerRelayed(peer.SiteId) { + continue + } + if err := monitor.RequestRelay(peer.SiteId); err != nil { + logger.Warn("Failed relay preflight for site %d: %v", peer.SiteId, err) + } + } + time.Sleep(300 * time.Millisecond) +} + +func (o *Olm) createWireGuardDevice(bind wgConn.Bind) (*device.Device, error) { + if o.middleDev == nil { + return nil, fmt.Errorf("middle device is nil") + } + wgLogger := logger.GetLogger().GetWireGuardLogger("wireguard: ") + dev := device.NewDevice(o.middleDev, bind, (*device.Logger)(wgLogger)) + if err := dev.IpcSet(fmt.Sprintf("private_key=%s", util.FixKey(o.privateKey.String()))); err != nil { + dev.Close() + return nil, fmt.Errorf("configure private key: %w", err) + } + if err := dev.Up(); err != nil { + dev.Close() + return nil, fmt.Errorf("bring up wireguard device: %w", err) + } + return dev, nil +} + +func (o *Olm) restoreSharedBindDevice(pm *peers.PeerManager) error { + if o.sharedBind == nil { + return fmt.Errorf("shared bind is nil") + } + dev, err := o.createWireGuardDevice(o.sharedBind) + if err != nil { + return err + } + if err := pm.ReplaceDevice(dev); err != nil { + dev.Close() + return err + } + o.dev = dev + o.wssRelayActive = false + return nil +} + +func (o *Olm) switchToWebSocketRelay() error { + o.relaySwitchMu.Lock() + defer o.relaySwitchMu.Unlock() + + return o.switchToWebSocketRelayLocked("") +} + +func (o *Olm) switchToWebSocketRelayWithURL(relayURL string) error { + o.relaySwitchMu.Lock() + defer o.relaySwitchMu.Unlock() + + return o.switchToWebSocketRelayLocked(relayURL) +} + +func (o *Olm) switchToWebSocketRelayLocked(prebuiltRelayURL string) error { + if o.wssRelayActive { + return nil + } + if o.wsRelayBind != nil { + return nil + } + if !o.tunnelRunning { + return fmt.Errorf("tunnel is not running") + } + + pm := o.getPeerManager() + if pm == nil { + return fmt.Errorf("peer manager is nil") + } + + relayURL := prebuiltRelayURL + if relayURL == "" { + var err error + relayURL, err = o.resolveRelayTunnelURL() + if err != nil { + return err + } + } else { + o.relayURLMu.Lock() + o.relayTunnelURL = relayURL + o.relayURLMu.Unlock() + } + o.requestRelayPreflight(pm) + + oldDev := o.dev + if oldDev != nil { + oldDev.Close() + } + + wsBind := NewWebSocketRelayBind(relayURL, nil) + newDev, err := o.createWireGuardDevice(wsBind) + if err != nil { + _ = wsBind.Shutdown() + if restoreErr := o.restoreSharedBindDevice(pm); restoreErr != nil { + return fmt.Errorf("switch to websocket relay failed: %w (restore failed: %v)", err, restoreErr) + } + return fmt.Errorf("switch to websocket relay failed: %w", err) + } + if err := pm.ReplaceDevice(newDev); err != nil { + newDev.Close() + _ = wsBind.Shutdown() + if restoreErr := o.restoreSharedBindDevice(pm); restoreErr != nil { + return fmt.Errorf("replace device failed: %w (restore failed: %v)", err, restoreErr) + } + return fmt.Errorf("replace device failed: %w", err) + } + + if o.wsRelayBind != nil { + _ = o.wsRelayBind.Shutdown() + } + o.dev = newDev + o.wsRelayBind = wsBind + o.wssRelayActive = true + if o.relaySwitchTimer != nil { + o.relaySwitchTimer.Stop() + o.relaySwitchTimer = nil + } + + logger.Info("Switched WireGuard transport to WSS relay bind") + return nil +} + +func (o *Olm) scheduleRelaySwitch(delay time.Duration) { + o.relaySwitchMu.Lock() + if o.relaySwitchTimer != nil { + o.relaySwitchTimer.Stop() + } + o.relaySwitchTimer = time.AfterFunc(delay, func() { + if err := o.switchToWebSocketRelay(); err != nil { + logger.Warn("WSS relay switch failed: %v", err) + } + }) + o.relaySwitchMu.Unlock() +} + +func (o *Olm) cancelRelaySwitchTimer() { + o.relaySwitchMu.Lock() + defer o.relaySwitchMu.Unlock() + if o.relaySwitchTimer != nil { + o.relaySwitchTimer.Stop() + o.relaySwitchTimer = nil + } +} diff --git a/olm/types.go b/olm/types.go index 9dd3189..169fc1d 100644 --- a/olm/types.go +++ b/olm/types.go @@ -87,4 +87,5 @@ type TunnelConfig struct { InitialPostures map[string]any DisableRelay bool + WebSocketRelay bool } diff --git a/olm/ws_relay_bind.go b/olm/ws_relay_bind.go new file mode 100644 index 0000000..81c6f52 --- /dev/null +++ b/olm/ws_relay_bind.go @@ -0,0 +1,261 @@ +package olm + +import ( + "fmt" + "net" + "net/http" + "net/netip" + "sync" + "time" + + "github.com/fosrl/newt/logger" + "github.com/gorilla/websocket" + wgConn "golang.zx2c4.com/wireguard/conn" +) + +// WebSocketRelayBind is a WireGuard bind implementation that transports packets +// over a WebSocket tunnel. +type WebSocketRelayBind struct { + url string + headers http.Header + + mu sync.RWMutex + conn *websocket.Conn + shutdown bool + port uint16 + recvCh chan []byte + shutdownCh chan struct{} + retryInterval time.Duration +} + +func NewWebSocketRelayBind(url string, headers http.Header) *WebSocketRelayBind { + return &WebSocketRelayBind{ + url: url, + headers: headers, + recvCh: make(chan []byte, 1024), + shutdownCh: make(chan struct{}), + // Keep retries short so WireGuard startup can recover quickly. + retryInterval: 500 * time.Millisecond, + } +} + +func (b *WebSocketRelayBind) connect() error { + for { + b.mu.RLock() + shutdown := b.shutdown + existingConn := b.conn + b.mu.RUnlock() + + if shutdown { + return net.ErrClosed + } + if existingConn != nil { + return nil + } + + dialer := websocket.Dialer{} + conn, _, err := dialer.Dial(b.url, b.headers) + if err != nil { + logger.Warn("WebSocket relay dial failed for %s: %v", b.url, err) + select { + case <-b.shutdownCh: + return net.ErrClosed + case <-time.After(b.retryInterval): + continue + } + } + + b.mu.Lock() + if b.shutdown { + b.mu.Unlock() + _ = conn.Close() + return net.ErrClosed + } + // If another goroutine won the race and already connected, use it. + if b.conn != nil { + b.mu.Unlock() + _ = conn.Close() + return nil + } + b.conn = conn + if conn.RemoteAddr() != nil { + if _, port, splitErr := net.SplitHostPort(conn.RemoteAddr().String()); splitErr == nil { + if parsed, lookupErr := net.LookupPort("tcp", port); lookupErr == nil { + b.port = uint16(parsed) + } + } + } + b.mu.Unlock() + + go b.readLoop(conn) + logger.Info("Connected WebSocket relay bind to %s", b.url) + return nil + } +} + +func (b *WebSocketRelayBind) readLoop(conn *websocket.Conn) { + for { + msgType, payload, err := conn.ReadMessage() + if err != nil { + select { + case <-b.shutdownCh: + return + default: + logger.Warn("WebSocket relay read failed: %v", err) + b.dropConn(conn) + return + } + } + if msgType != websocket.BinaryMessage { + continue + } + select { + case b.recvCh <- payload: + case <-b.shutdownCh: + return + } + } +} + +func (b *WebSocketRelayBind) dropConn(conn *websocket.Conn) { + b.mu.Lock() + defer b.mu.Unlock() + if b.conn == conn { + _ = b.conn.Close() + b.conn = nil + } +} + +func (b *WebSocketRelayBind) WriteToRelay(data []byte) error { + if len(data) == 0 { + return nil + } + + for { + b.mu.RLock() + conn := b.conn + shutdown := b.shutdown + b.mu.RUnlock() + if shutdown { + return net.ErrClosed + } + if conn == nil { + if err := b.connect(); err != nil { + return err + } + continue + } + + if err := conn.WriteMessage(websocket.BinaryMessage, data); err != nil { + logger.Warn("WebSocket relay write failed, reconnecting: %v", err) + b.dropConn(conn) + select { + case <-b.shutdownCh: + return net.ErrClosed + case <-time.After(b.retryInterval): + continue + } + } + return nil + } +} + +func (b *WebSocketRelayBind) Close() error { + b.mu.Lock() + defer b.mu.Unlock() + if b.conn != nil { + _ = b.conn.Close() + b.conn = nil + } + return nil +} + +// Shutdown permanently closes the bind and should only be called when OLM is +// shutting down, not during routine WireGuard bind updates. +func (b *WebSocketRelayBind) Shutdown() error { + b.mu.Lock() + defer b.mu.Unlock() + if b.shutdown { + return nil + } + b.shutdown = true + close(b.shutdownCh) + if b.conn != nil { + _ = b.conn.Close() + b.conn = nil + } + return nil +} + +func (b *WebSocketRelayBind) Open(_ uint16) ([]wgConn.ReceiveFunc, uint16, error) { + if err := b.connect(); err != nil { + return nil, 0, err + } + + recvFn := func(bufs [][]byte, sizes []int, eps []wgConn.Endpoint) (int, error) { + for { + select { + case <-b.shutdownCh: + return 0, net.ErrClosed + case packet := <-b.recvCh: + if len(bufs) == 0 { + return 0, nil + } + n := copy(bufs[0], packet) + sizes[0] = n + eps[0] = &wgConn.StdNetEndpoint{ + AddrPort: netip.AddrPortFrom(netip.IPv4Unspecified(), b.port), + } + return 1, nil + default: + b.mu.RLock() + shutdown := b.shutdown + hasConn := b.conn != nil + b.mu.RUnlock() + if shutdown { + return 0, net.ErrClosed + } + if !hasConn { + if err := b.connect(); err != nil { + return 0, err + } + continue + } + time.Sleep(10 * time.Millisecond) + } + } + } + + if b.port == 0 { + b.port = 443 + } + return []wgConn.ReceiveFunc{recvFn}, b.port, nil +} + +func (b *WebSocketRelayBind) Send(bufs [][]byte, _ wgConn.Endpoint) error { + for _, buf := range bufs { + if err := b.connect(); err != nil { + return fmt.Errorf("failed writing relay packet: %w", err) + } + if err := b.WriteToRelay(buf); err != nil { + return err + } + } + return nil +} + +func (b *WebSocketRelayBind) SetMark(_ uint32) error { + return nil +} + +func (b *WebSocketRelayBind) ParseEndpoint(s string) (wgConn.Endpoint, error) { + addrPort, err := netip.ParseAddrPort(s) + if err != nil { + return nil, err + } + return &wgConn.StdNetEndpoint{AddrPort: addrPort}, nil +} + +func (b *WebSocketRelayBind) BatchSize() int { + return 1 +} diff --git a/peers/manager.go b/peers/manager.go index 9cc1e75..b6e700d 100644 --- a/peers/manager.go +++ b/peers/manager.go @@ -107,10 +107,64 @@ func (pm *PeerManager) GetAllPeers() []SiteConfig { return peers } +// SetDevice swaps the WireGuard device reference used by peer operations. +func (pm *PeerManager) SetDevice(dev *device.Device) { + pm.mu.Lock() + defer pm.mu.Unlock() + pm.device = dev +} + +// IsPeerConnected reports the latest WireGuard connectivity status for a peer. +func (pm *PeerManager) IsPeerConnected(siteID int) bool { + if pm.peerMonitor == nil { + return false + } + return pm.peerMonitor.IsWGConnected(siteID) +} + +// IsPeerRelayed reports whether a peer is currently marked as relayed. +func (pm *PeerManager) IsPeerRelayed(siteID int) bool { + if pm.peerMonitor == nil { + return false + } + return pm.peerMonitor.IsPeerRelayed(siteID) +} + +// ReplaceDevice swaps the active WireGuard device and reapplies current peers. +func (pm *PeerManager) ReplaceDevice(dev *device.Device) error { + if dev == nil { + return fmt.Errorf("device is nil") + } + + pm.mu.Lock() + pm.device = dev + peersSnapshot := make([]SiteConfig, 0, len(pm.peers)) + for _, peer := range pm.peers { + peersSnapshot = append(peersSnapshot, peer) + } + monitor := pm.peerMonitor + privateKey := pm.privateKey + persistentKeepalive := pm.PersistentKeepalive + publicDNS := append([]string(nil), pm.publicDNS...) + pm.mu.Unlock() + + for _, peer := range peersSnapshot { + relayed := false + if monitor != nil { + relayed = monitor.IsPeerRelayed(peer.SiteId) + } + if err := ConfigurePeer(dev, peer, privateKey, relayed, persistentKeepalive, publicDNS); err != nil { + return fmt.Errorf("reconfigure peer %d: %w", peer.SiteId, err) + } + } + + return nil +} + func (pm *PeerManager) AddPeer(siteConfig SiteConfig) error { pm.mu.Lock() defer pm.mu.Unlock() - + for _, alias := range siteConfig.Aliases { address := net.ParseIP(alias.AliasAddress) if address == nil { @@ -118,7 +172,7 @@ func (pm *PeerManager) AddPeer(siteConfig SiteConfig) error { } pm.dnsProxy.AddDNSRecord(alias.Alias, address, siteConfig.SiteId) } - + if siteConfig.PublicKey == "" { logger.Debug("Skip adding site %d because no pub key", siteConfig.SiteId) return nil @@ -156,7 +210,7 @@ func (pm *PeerManager) AddPeer(siteConfig SiteConfig) error { if err := network.AddRoutes(siteConfig.RemoteSubnets, pm.interfaceName); err != nil { logger.Error("Failed to add routes for remote subnets: %v", err) } - + monitorAddress := strings.Split(siteConfig.ServerIP, "/")[0] monitorPeer := net.JoinHostPort(monitorAddress, strconv.Itoa(int(siteConfig.ServerPort+1))) // +1 for the monitor port @@ -183,7 +237,7 @@ func (pm *PeerManager) AddPeer(siteConfig SiteConfig) error { func (pm *PeerManager) UpdateAllPeersPersistentKeepalive(interval int) map[int]error { pm.mu.RLock() defer pm.mu.RUnlock() - + pm.PersistentKeepalive = interval errors := make(map[int]error) @@ -786,11 +840,23 @@ func (pm *PeerManager) RemoveAlias(siteId int, aliasName string) error { // RelayPeer handles failover to the relay server when a peer is disconnected func (pm *PeerManager) RelayPeer(siteId int, relayEndpoint string, relayPort uint16) { + // Check for IPv6 and format the endpoint correctly + formattedEndpoint := relayEndpoint + if strings.Contains(relayEndpoint, ":") { + formattedEndpoint = fmt.Sprintf("[%s]", relayEndpoint) + } + + if relayPort == 0 { + relayPort = 21820 // fall back to 21820 for backward compatibility + } + + relayHostPort := fmt.Sprintf("%s:%d", formattedEndpoint, relayPort) + pm.mu.Lock() peer, exists := pm.peers[siteId] if exists { - // Store the relay endpoint - peer.RelayEndpoint = relayEndpoint + // Persist host:port so ReplaceDevice reconfiguration keeps the same relay port. + peer.RelayEndpoint = relayHostPort pm.peers[siteId] = peer } pm.mu.Unlock() @@ -800,20 +866,10 @@ func (pm *PeerManager) RelayPeer(siteId int, relayEndpoint string, relayPort uin return } - // Check for IPv6 and format the endpoint correctly - formattedEndpoint := relayEndpoint - if strings.Contains(relayEndpoint, ":") { - formattedEndpoint = fmt.Sprintf("[%s]", relayEndpoint) - } - - if relayPort == 0 { - relayPort = 21820 // fall back to 21820 for backward compatibility - } - // Update only the endpoint for this peer (update_only preserves other settings) wgConfig := fmt.Sprintf(`public_key=%s update_only=true -endpoint=%s:%d`, util.FixKey(peer.PublicKey), formattedEndpoint, relayPort) +endpoint=%s`, util.FixKey(peer.PublicKey), relayHostPort) err := pm.device.IpcSet(wgConfig) if err != nil { @@ -829,6 +885,26 @@ endpoint=%s:%d`, util.FixKey(peer.PublicKey), formattedEndpoint, relayPort) logger.Info("Adjusted peer %d to point to relay!\n", siteId) } +// RelayPeerWss attempts to switch relay transport to WebSocket tunnel endpoint. +// Current implementation keeps the existing peer endpoint untouched and marks +// the peer relayed so the monitor/API reflect transport fallback intent. +func (pm *PeerManager) RelayPeerWss(siteId int, relayEndpointWss string) { + pm.mu.Lock() + _, exists := pm.peers[siteId] + pm.mu.Unlock() + + if !exists { + logger.Error("Cannot switch to websocket relay: peer with site ID %d not found", siteId) + return + } + + if pm.peerMonitor != nil { + pm.peerMonitor.MarkPeerRelayed(siteId, true) + } + + logger.Info("Requested websocket relay for peer %d via %s", siteId, relayEndpointWss) +} + // performRapidInitialTest performs a rapid holepunch test for a newly added peer. // If the test fails, it immediately requests relay to minimize connection delay. // This runs in a goroutine to avoid blocking AddPeer. diff --git a/peers/monitor/monitor.go b/peers/monitor/monitor.go index 56dcee4..5aea9b3 100644 --- a/peers/monitor/monitor.go +++ b/peers/monitor/monitor.go @@ -36,7 +36,7 @@ type PeerMonitor struct { timeout time.Duration maxAttempts int wsClient *websocket.Client - publicDNS []string + publicDNS []string // Relay sender tracking relaySends map[string]func() @@ -104,7 +104,7 @@ func NewPeerMonitor(wsClient *websocket.Client, middleDev *middleDevice.MiddleDe wsClient: wsClient, middleDev: middleDev, localIP: localIP, - publicDNS: publicDNS, + publicDNS: publicDNS, activePorts: make(map[uint16]bool), nsCtx: ctx, nsCancel: cancel, @@ -473,7 +473,6 @@ func (pm *PeerMonitor) CancelRelaySend(chainId string) { logger.Info("Cancelled all relay senders") return } - if stop, ok := pm.relaySends[chainId]; ok { stop() delete(pm.relaySends, chainId) @@ -719,6 +718,13 @@ func (pm *PeerMonitor) GetHolepunchStatus() map[int]bool { return status } +// IsWGConnected returns whether WireGuard is currently connected for a peer. +func (pm *PeerMonitor) IsWGConnected(siteID int) bool { + pm.mutex.Lock() + defer pm.mutex.Unlock() + return pm.wgConnectionStatus[siteID] +} + // Close stops monitoring and cleans up resources func (pm *PeerMonitor) Close() { // Stop holepunch monitor first (outside of mutex to avoid deadlock) diff --git a/peers/types.go b/peers/types.go index 9ef1462..a53cf09 100644 --- a/peers/types.go +++ b/peers/types.go @@ -34,6 +34,7 @@ type RelayPeerData struct { SiteId int `json:"siteId"` RelayEndpoint string `json:"relayEndpoint"` RelayPort uint16 `json:"relayPort"` + RelayEndpointWss string `json:"relayEndpointWss,omitempty"` } type UnRelayPeerData struct { diff --git a/websocket/client.go b/websocket/client.go index 3b4e894..0cea858 100644 --- a/websocket/client.go +++ b/websocket/client.go @@ -51,6 +51,7 @@ type TokenResponse struct { type ExitNode struct { Endpoint string `json:"endpoint"` RelayPort uint16 `json:"relayPort"` + RelayEndpointWss string `json:"relayEndpointWss,omitempty"` PublicKey string `json:"publicKey"` SiteIds []int `json:"siteIds"` } @@ -92,17 +93,17 @@ type Client struct { configNeedsSave bool // Flag to track if config needs to be saved configVersion int // Latest config version received from server configVersionMux sync.RWMutex - token string // Cached authentication token - exitNodes []ExitNode // Cached exit nodes from token response - tokenMux sync.RWMutex // Protects token and exitNodes - forceNewToken bool // Flag to force fetching a new token on next connection - processingMessage bool // Flag to track if a message is currently being processed - processingMux sync.RWMutex // Protects processingMessage - processingWg sync.WaitGroup // WaitGroup to wait for message processing to complete - getPingData func() map[string]any // Callback to get additional ping data - pingStarted bool // Flag to track if ping monitor has been started - pingStartedMux sync.Mutex // Protects pingStarted - pingDone chan struct{} // Channel to stop the ping monitor independently + token string // Cached authentication token + exitNodes []ExitNode // Cached exit nodes from token response + tokenMux sync.RWMutex // Protects token and exitNodes + forceNewToken bool // Flag to force fetching a new token on next connection + processingMessage bool // Flag to track if a message is currently being processed + processingMux sync.RWMutex // Protects processingMessage + processingWg sync.WaitGroup // WaitGroup to wait for message processing to complete + getPingData func() map[string]any // Callback to get additional ping data + pingStarted bool // Flag to track if ping monitor has been started + pingStartedMux sync.Mutex // Protects pingStarted + pingDone chan struct{} // Channel to stop the ping monitor independently } type ClientOption func(*Client) @@ -153,6 +154,32 @@ func (c *Client) OnTokenUpdate(callback func(token string, exitNodes []ExitNode) c.onTokenUpdate = callback } +// GetToken returns the currently cached websocket token. +func (c *Client) GetToken() string { + c.tokenMux.RLock() + defer c.tokenMux.RUnlock() + return c.token +} + +// GetUserToken returns the configured user token for websocket relay auth. +func (c *Client) GetUserToken() string { + if c.config == nil { + return "" + } + return c.config.UserToken +} + +// GetTokenState returns the currently cached token and exit nodes. +// A copy of exit nodes is returned to avoid exposing internal mutable state. +func (c *Client) GetTokenState() (string, []ExitNode) { + c.tokenMux.RLock() + defer c.tokenMux.RUnlock() + + exitNodesCopy := make([]ExitNode, len(c.exitNodes)) + copy(exitNodesCopy, c.exitNodes) + return c.token, exitNodesCopy +} + func (c *Client) OnAuthError(callback func(statusCode int, message string)) { c.onAuthError = callback } @@ -385,10 +412,10 @@ func (c *Client) getToken() (string, []ExitNode, error) { } tokenData := map[string]interface{}{ - "olmId": c.config.ID, - "secret": c.config.Secret, + "olmId": c.config.ID, + "secret": c.config.Secret, "userToken": c.config.UserToken, - "orgId": c.config.OrgID, + "orgId": c.config.OrgID, } jsonData, err := json.Marshal(tokenData) @@ -736,15 +763,15 @@ func (c *Client) pingMonitor() { func (c *Client) StartPingMonitor() { c.pingStartedMux.Lock() defer c.pingStartedMux.Unlock() - + if c.pingStarted { return } c.pingStarted = true - + // Create a new pingDone channel for this ping monitor instance c.pingDone = make(chan struct{}) - + // Send an initial ping immediately go func() { c.sendPing() @@ -756,11 +783,11 @@ func (c *Client) StartPingMonitor() { func (c *Client) stopPingMonitor() { c.pingStartedMux.Lock() defer c.pingStartedMux.Unlock() - + if !c.pingStarted { return } - + // Close the pingDone channel to stop the monitor close(c.pingDone) c.pingStarted = false