From 71eb81182952eab7155fb8bd25538f7716aa5033 Mon Sep 17 00:00:00 2001 From: Owen Date: Mon, 25 May 2026 21:58:27 -0700 Subject: [PATCH] Add sighup and block flag to block connections Former-commit-id: 694a7986f5b7c5f0610073b882e7a3203fe27108 --- clients.go | 7 +++++ clients/clients.go | 27 ++++++++++++++++++ main.go | 60 +++++++++++++++++++++++++++++++++++++++ netstack2/handlers.go | 18 ++++++++---- netstack2/http_handler.go | 5 ---- netstack2/proxy.go | 26 ++++++++++++++++- netstack2/tun.go | 5 ---- proxy/manager.go | 46 +++++++++++++++++++++++------- websocket/client.go | 5 ++++ websocket/types.go | 3 +- 10 files changed, 175 insertions(+), 27 deletions(-) diff --git a/clients.go b/clients.go index d650eeb..53ce63f 100644 --- a/clients.go +++ b/clients.go @@ -63,6 +63,13 @@ func closeClients() { } } +// setClientsBlocked enables or disables connection blocking on the WireGuard service. +func setClientsBlocked(v bool) { + if wgService != nil { + wgService.SetBlocked(v) + } +} + func clientsHandleNewtConnection(publicKey string, endpoint string, relayPort uint16) { if !ready { return diff --git a/clients/clients.go b/clients/clients.go index 3862160..cbd2816 100644 --- a/clients/clients.go +++ b/clients/clients.go @@ -13,6 +13,7 @@ import ( "strconv" "strings" "sync" + "sync/atomic" "time" "github.com/fosrl/newt/bind" @@ -114,6 +115,9 @@ type WireGuardService struct { netstackListener net.PacketConn netstackListenerMu sync.Mutex wgTesterServer *wgtester.Server + + // connection blocking: when true, all new incoming connections are dropped + blocked atomic.Bool } // generateChainId generates a random chain ID for deduplicating round-trip messages. @@ -286,6 +290,21 @@ func (s *WireGuardService) GetPublicKey() wgtypes.Key { return s.key.PublicKey() } +// SetBlocked enables or disables connection blocking for this WireGuard service. +// The state is persisted and applied immediately to any active proxy handler. +func (s *WireGuardService) SetBlocked(v bool) { + s.blocked.Store(v) + s.mu.Lock() + tnet := s.tnet + s.mu.Unlock() + if tnet == nil { + return + } + if ph := tnet.GetProxyHandler(); ph != nil { + ph.SetBlocked(v) + } +} + // SetOnNetstackReady sets a callback function to be called when the netstack interface is ready func (s *WireGuardService) SetOnNetstackReady(callback func(*netstack2.Net)) { s.onNetstackReady = callback @@ -891,6 +910,14 @@ func (s *WireGuardService) ensureWireguardInterface(wgconfig WgConfig) error { } // Note: we already unlocked above, so don't use defer unlock + + // Apply any pending blocked state to the newly created proxy handler + if s.blocked.Load() { + if ph := s.tnet.GetProxyHandler(); ph != nil { + ph.SetBlocked(true) + } + } + return nil } diff --git a/main.go b/main.go index 0c8e7dd..85d15e7 100644 --- a/main.go +++ b/main.go @@ -18,6 +18,7 @@ import ( "os/signal" "strconv" "strings" + "sync/atomic" "syscall" "time" @@ -736,11 +737,20 @@ func runNewtMain(ctx context.Context) { var connected bool var wgData WgData var dockerEventMonitor *docker.EventMonitor + var connectionBlocked atomic.Bool + var currentPM atomic.Pointer[proxy.ProxyManager] if !disableClients { setupClients(client) } + // Initialize connection blocked state from config + connectionBlocked.Store(client.GetConfig().Blocked) + if connectionBlocked.Load() { + logger.Info("Connection blocking is enabled (from config)") + setClientsBlocked(true) + } + // Initialize health check monitor with status change callback healthMonitor = healthcheck.NewMonitor(func(targets map[int]*healthcheck.Target) { logger.Debug("Health check status update for %d targets", len(targets)) @@ -780,6 +790,7 @@ func runNewtMain(ctx context.Context) { // Stop proxy manager if running if pm != nil { pm.Stop() + currentPM.Store(nil) pm = nil } @@ -950,6 +961,8 @@ persistent_keepalive_interval=5`, util.FixKey(privateKey.String()), util.FixKey( pm.SetUDPIdleTimeout(udpProxyIdleTimeout) // Set tunnel_id for metrics (WireGuard peer public key) pm.SetTunnelID(wgData.PublicKey) + pm.SetBlocked(connectionBlocked.Load()) + currentPM.Store(pm) connected = true @@ -1925,6 +1938,53 @@ persistent_keepalive_interval=5`, util.FixKey(privateKey.String()), util.FixKey( return nil }) + // Handle SIGHUP for config reload + sighupChan := make(chan os.Signal, 1) + signal.Notify(sighupChan, syscall.SIGHUP) + go func() { + defer signal.Stop(sighupChan) + for { + select { + case <-sighupChan: + logger.Info("SIGHUP received, reloading config...") + cfgPath := client.GetConfigFilePath() + data, err := os.ReadFile(cfgPath) + if err != nil { + logger.Error("Failed to read config file on SIGHUP: %v", err) + continue + } + var newCfg websocket.Config + if err := json.Unmarshal(data, &newCfg); err != nil { + logger.Error("Failed to parse config file on SIGHUP: %v", err) + continue + } + oldCfg := client.GetConfig() + // If credentials changed, exit so the supervisor can restart with new values + if newCfg.Endpoint != oldCfg.Endpoint || newCfg.ID != oldCfg.ID || newCfg.Secret != oldCfg.Secret { + logger.Info("Config credentials changed (endpoint/id/secret), exiting for supervisor restart...") + os.Exit(0) + } + // If blocked state changed, apply in-place without restart + if newCfg.Blocked != connectionBlocked.Load() { + connectionBlocked.Store(newCfg.Blocked) + if newCfg.Blocked { + logger.Info("Config reload: connection blocking enabled") + } else { + logger.Info("Config reload: connection blocking disabled") + } + if p := currentPM.Load(); p != nil { + p.SetBlocked(newCfg.Blocked) + } + setClientsBlocked(newCfg.Blocked) + } else { + logger.Info("Config reload: no relevant changes detected") + } + case <-ctx.Done(): + return + } + } + }() + // Connect to the WebSocket server if err := client.Connect(); err != nil { logger.Fatal("Failed to connect to server: %v", err) diff --git a/netstack2/handlers.go b/netstack2/handlers.go index e28a543..e6ea62e 100644 --- a/netstack2/handlers.go +++ b/netstack2/handlers.go @@ -1,8 +1,3 @@ -/* SPDX-License-Identifier: MIT - * - * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved. - */ - package netstack2 import ( @@ -144,6 +139,13 @@ func (h *TCPHandler) handleTCPConn(netstackConn *gonet.TCPConn, id stack.Transpo dstIP := id.LocalAddress.String() dstPort := id.LocalPort + // Drop connection if blocking is enabled + if h.proxyHandler != nil && h.proxyHandler.IsBlocked() { + logger.Debug("TCP Forwarder: connection blocked: %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort) + netstackConn.Close() + return + } + // For HTTP/HTTPS ports, look up the matching subnet rule. If the rule has // Protocol configured, hand the connection off to the HTTP handler which // takes full ownership of the lifecycle (the defer close must not be @@ -315,6 +317,12 @@ func (h *UDPHandler) handleUDPConn(netstackConn *gonet.UDPConn, id stack.Transpo logger.Info("UDP Forwarder: Handling connection %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort) + // Drop connection if blocking is enabled + if h.proxyHandler != nil && h.proxyHandler.IsBlocked() { + logger.Debug("UDP Forwarder: connection blocked: %s:%d -> %s:%d", srcIP, srcPort, dstIP, dstPort) + return + } + // Check if there's a destination rewrite for this connection (e.g., localhost targets) actualDstIP := dstIP if h.proxyHandler != nil { diff --git a/netstack2/http_handler.go b/netstack2/http_handler.go index ece82e9..7a4d69d 100644 --- a/netstack2/http_handler.go +++ b/netstack2/http_handler.go @@ -1,8 +1,3 @@ -/* SPDX-License-Identifier: MIT - * - * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved. - */ - package netstack2 import ( diff --git a/netstack2/proxy.go b/netstack2/proxy.go index 95fab6a..2df5c83 100644 --- a/netstack2/proxy.go +++ b/netstack2/proxy.go @@ -6,6 +6,7 @@ import ( "net" "net/netip" "sync" + "sync/atomic" "time" "github.com/fosrl/newt/logger" @@ -134,6 +135,7 @@ type ProxyHandler struct { notifiable channel.Notification // Notification handler for triggering reads accessLogger *AccessLogger // Access logger for tracking sessions httpRequestLogger *HTTPRequestLogger // HTTP request logger for proxied HTTP/HTTPS requests + blocked atomic.Bool // when true, all new connections are dropped } // ProxyHandlerOptions configures the proxy handler @@ -240,6 +242,28 @@ func (p *ProxyHandler) AddSubnetRule(rule SubnetRule) { p.subnetLookup.AddSubnet(rule) } +// SetBlocked enables or disables connection blocking on this proxy handler. +// When enabled, all new TCP/UDP connections from the tunnel are dropped immediately. +func (p *ProxyHandler) SetBlocked(v bool) { + if p == nil { + return + } + p.blocked.Store(v) + if v { + logger.Info("ProxyHandler: connection blocking enabled") + } else { + logger.Info("ProxyHandler: connection blocking disabled") + } +} + +// IsBlocked returns true if connection blocking is currently enabled. +func (p *ProxyHandler) IsBlocked() bool { + if p == nil { + return false + } + return p.blocked.Load() +} + // RemoveSubnetRule removes a subnet from the proxy handler func (p *ProxyHandler) RemoveSubnetRule(sourcePrefix, destPrefix netip.Prefix) { if p == nil || !p.enabled { @@ -612,7 +636,7 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool { } // logger.Debug("HandleIncomingPacket: No matching rule for %s -> %s (proto=%d, port=%d)", - // srcAddr, dstAddr, protocol, dstPort) + // srcAddr, dstAddr, protocol, dstPort) return false } diff --git a/netstack2/tun.go b/netstack2/tun.go index fae90dd..d104f2e 100644 --- a/netstack2/tun.go +++ b/netstack2/tun.go @@ -1,8 +1,3 @@ -/* SPDX-License-Identifier: MIT - * - * Copyright (C) 2017-2025 WireGuard LLC. All Rights Reserved. - */ - package netstack2 import ( diff --git a/proxy/manager.go b/proxy/manager.go index 0d1f750..b04be26 100644 --- a/proxy/manager.go +++ b/proxy/manager.go @@ -70,6 +70,9 @@ type ProxyManager struct { asyncBytes bool flushStop chan struct{} udpIdleTimeout time.Duration + + // connection blocking + blocked atomic.Bool } // tunnelEntry holds per-tunnel attributes and (optional) async counters. @@ -149,12 +152,12 @@ func classifyProxyError(err error) string { // NewProxyManager creates a new proxy manager instance func NewProxyManager(tnet *netstack.Net) *ProxyManager { return &ProxyManager{ - tnet: tnet, - tcpTargets: make(map[string]map[int]string), - udpTargets: make(map[string]map[int]string), - listeners: make([]*gonet.TCPListener, 0), - udpConns: make([]*gonet.UDPConn, 0), - tunnels: make(map[string]*tunnelEntry), + tnet: tnet, + tcpTargets: make(map[string]map[int]string), + udpTargets: make(map[string]map[int]string), + listeners: make([]*gonet.TCPListener, 0), + udpConns: make([]*gonet.UDPConn, 0), + tunnels: make(map[string]*tunnelEntry), udpIdleTimeout: defaultUDPIdleTimeout, } } @@ -229,10 +232,10 @@ func (pm *ProxyManager) ClearTunnelID() { // init function without tnet func NewProxyManagerWithoutTNet() *ProxyManager { return &ProxyManager{ - tcpTargets: make(map[string]map[int]string), - udpTargets: make(map[string]map[int]string), - listeners: make([]*gonet.TCPListener, 0), - udpConns: make([]*gonet.UDPConn, 0), + tcpTargets: make(map[string]map[int]string), + udpTargets: make(map[string]map[int]string), + listeners: make([]*gonet.TCPListener, 0), + udpConns: make([]*gonet.UDPConn, 0), udpIdleTimeout: defaultUDPIdleTimeout, } } @@ -244,6 +247,18 @@ func (pm *ProxyManager) SetTNet(tnet *netstack.Net) { pm.tnet = tnet } +// SetBlocked enables or disables connection blocking. +// When enabled, all new incoming TCP connections are immediately closed +// and all incoming UDP packets are silently dropped. +func (pm *ProxyManager) SetBlocked(v bool) { + pm.blocked.Store(v) + if v { + logger.Info("ProxyManager: connection blocking enabled, new connections will be dropped") + } else { + logger.Info("ProxyManager: connection blocking disabled, accepting connections") + } +} + // AddTarget adds as new target for proxying func (pm *ProxyManager) AddTarget(proto, listenIP string, port int, targetAddr string) error { pm.mutex.Lock() @@ -535,6 +550,12 @@ func (pm *ProxyManager) handleTCPProxy(listener net.Listener, targetAddr string) } tunnelID := pm.currentTunnelID + // Drop connection if blocking is enabled + if pm.blocked.Load() { + conn.Close() + logger.Debug("TCP proxy: connection dropped (blocking enabled)") + continue + } telemetry.IncProxyAccept(context.Background(), tunnelID, "tcp", "success", "") telemetry.IncProxyConnectionEvent(context.Background(), tunnelID, "tcp", telemetry.ProxyConnectionOpened) if tunnelID != "" { @@ -631,6 +652,11 @@ func (pm *ProxyManager) handleUDPProxy(conn *gonet.UDPConn, targetAddr string) { } clientKey := remoteAddr.String() + // Drop packet if blocking is enabled + if pm.blocked.Load() { + logger.Debug("UDP proxy: packet dropped (blocking enabled)") + continue + } // bytes from client -> target (direction=in) if pm.currentTunnelID != "" && n > 0 { if pm.asyncBytes { diff --git a/websocket/client.go b/websocket/client.go index 22187ac..0f4ea39 100644 --- a/websocket/client.go +++ b/websocket/client.go @@ -172,6 +172,11 @@ func (c *Client) GetConfig() *Config { return c.config } +// GetConfigFilePath returns the resolved path to the config file used by this client. +func (c *Client) GetConfigFilePath() string { + return getConfigPath(c.clientType, c.configFilePath) +} + func (c *Client) GetServerVersion() string { return c.serverVersion } diff --git a/websocket/types.go b/websocket/types.go index 195e06f..58b304e 100644 --- a/websocket/types.go +++ b/websocket/types.go @@ -7,6 +7,7 @@ type Config struct { TlsClientCert string `json:"tlsClientCert"` ProvisioningKey string `json:"provisioningKey,omitempty"` Name string `json:"name,omitempty"` + Blocked bool `json:"blocked,omitempty"` } type TokenResponse struct { @@ -31,4 +32,4 @@ type WSMessage struct { Type string `json:"type"` Data interface{} `json:"data"` ConfigVersion int64 `json:"configVersion,omitempty"` -} \ No newline at end of file +}