Add chainId based dedup

refactor(proxy): cleanup basics - constants, remove dead code, fix de…

Dockerfile

@@ -1,5 +1,5 @@
 # FROM golang:1.25-alpine AS builder
-FROM public.ecr.aws/docker/library/golang:1.26-alpine AS builder
+FROM public.ecr.aws/docker/library/golang:1.25-alpine AS builder
 
 # Install git and ca-certificates
 RUN apk --no-cache add ca-certificates git tzdata

clients/clients.go

@@ -2,6 +2,8 @@ package clients
 
 import (
 	"context"
+	"crypto/rand"
+	"encoding/hex"
 	"encoding/json"
 	"fmt"
 	"net"
@@ -34,6 +36,7 @@ type WgConfig struct {
 	IpAddress string   `json:"ipAddress"`
 	Peers     []Peer   `json:"peers"`
 	Targets   []Target `json:"targets"`
+	ChainId   string   `json:"chainId"`
 }
 
 type Target struct {
@@ -82,7 +85,8 @@ type WireGuardService struct {
 	host          string
 	serverPubKey  string
 	token         string
-	stopGetConfig func()
+	stopGetConfig        func()
+	pendingConfigChainId string
 	// Netstack fields
 	tun    tun.Device
 	tnet   *netstack2.Net
@@ -107,6 +111,13 @@ type WireGuardService struct {
 	wgTesterServer     *wgtester.Server
 }
 
+// generateChainId generates a random chain ID for deduplicating round-trip messages.
+func generateChainId() string {
+	b := make([]byte, 8)
+	_, _ = rand.Read(b)
+	return hex.EncodeToString(b)
+}
+
 func NewWireGuardService(interfaceName string, port uint16, mtu int, host string, newtId string, wsClient *websocket.Client, dns string, useNativeInterface bool) (*WireGuardService, error) {
 	key, err := wgtypes.GeneratePrivateKey()
 	if err != nil {
@@ -161,8 +172,9 @@ func NewWireGuardService(interfaceName string, port uint16, mtu int, host string
 		useNativeInterface: useNativeInterface,
 	}
 
-	// Create the holepunch manager
+	// Create the holepunch manager with ResolveDomain function
-	service.holePunchManager = holepunch.NewManager(sharedBind, newtId, "newt", key.PublicKey().String(), nil)
+	// We'll need to pass a domain resolver function
+	service.holePunchManager = holepunch.NewManager(sharedBind, newtId, "newt", key.PublicKey().String())
 
 	// Register websocket handlers
 	wsClient.RegisterHandler("newt/wg/receive-config", service.handleConfig)
@@ -441,9 +453,12 @@ func (s *WireGuardService) LoadRemoteConfig() error {
 		s.stopGetConfig()
 		s.stopGetConfig = nil
 	}
+	chainId := generateChainId()
+	s.pendingConfigChainId = chainId
 	s.stopGetConfig = s.client.SendMessageInterval("newt/wg/get-config", map[string]interface{}{
 		"publicKey": s.key.PublicKey().String(),
 		"port":      s.Port,
+		"chainId":   chainId,
 	}, 2*time.Second)
 
 	logger.Debug("Requesting WireGuard configuration from remote server")
@@ -468,6 +483,17 @@ func (s *WireGuardService) handleConfig(msg websocket.WSMessage) {
 		logger.Info("Error unmarshaling target data: %v", err)
 		return
 	}
+
+	// Deduplicate using chainId: discard responses that don't match the
+	// pending request, or that we have already processed.
+	if config.ChainId != "" {
+		if config.ChainId != s.pendingConfigChainId {
+			logger.Debug("Discarding duplicate/stale newt/wg/get-config response (chainId=%s, expected=%s)", config.ChainId, s.pendingConfigChainId)
+			return
+		}
+		s.pendingConfigChainId = "" // consume – further duplicates are rejected
+	}
+
 	s.config = config
 
 	if s.stopGetConfig != nil {

common.go

@@ -287,9 +287,12 @@ func startPingCheck(tnet *netstack.Net, serverIP string, client *websocket.Clien
 							}
 							stopFunc = client.SendMessageInterval("newt/ping/request", map[string]interface{}{}, 3*time.Second)
 							// Send registration message to the server for backward compatibility
+							bcChainId := generateChainId()
+							pendingRegisterChainId = bcChainId
 							err := client.SendMessage("newt/wg/register", map[string]interface{}{
 								"publicKey":           publicKey.String(),
 								"backwardsCompatible": true,
+								"chainId":             bcChainId,
 							})
 							if err != nil {
 								logger.Error("Failed to send registration message: %v", err)

flake.nix

@@ -35,7 +35,7 @@
             inherit version;
             src = pkgs.nix-gitignore.gitignoreSource [ ] ./.;
 
-            vendorHash = "sha256-vy6Dqjek7pLdASbCrM9snq5Dt9lbwNJ0AuQboy1JWNQ=";
+            vendorHash = "sha256-kmQM8Yy5TuOiNpMpUme/2gfE+vrhUK+0AphN+p71wGs=";
 
             nativeInstallCheckInputs = [ pkgs.versionCheckHook ];
 

healthcheck/healthcheck.go

@@ -5,7 +5,9 @@ import (
 	"crypto/tls"
 	"encoding/json"
 	"fmt"
+	"net"
 	"net/http"
+	"strconv"
 	"strings"
 	"sync"
 	"time"
@@ -365,11 +367,12 @@ func (m *Monitor) performHealthCheck(target *Target) {
 	target.LastCheck = time.Now()
 	target.LastError = ""
 
-	// Build URL
+	// Build URL (use net.JoinHostPort to properly handle IPv6 addresses with ports)
-	url := fmt.Sprintf("%s://%s", target.Config.Scheme, target.Config.Hostname)
+	host := target.Config.Hostname
 	if target.Config.Port > 0 {
-		url = fmt.Sprintf("%s:%d", url, target.Config.Port)
+		host = net.JoinHostPort(target.Config.Hostname, strconv.Itoa(target.Config.Port))
 	}
+	url := fmt.Sprintf("%s://%s", target.Config.Scheme, host)
 	if target.Config.Path != "" {
 		if !strings.HasPrefix(target.Config.Path, "/") {
 			url += "/"

holepunch/holepunch.go

@@ -27,17 +27,16 @@ type ExitNode struct {
 
 // Manager handles UDP hole punching operations
 type Manager struct {
-	mu          sync.Mutex
+	mu         sync.Mutex
-	running     bool
+	running    bool
-	stopChan    chan struct{}
+	stopChan   chan struct{}
-	sharedBind  *bind.SharedBind
+	sharedBind *bind.SharedBind
-	ID          string
+	ID         string
-	token       string
+	token      string
-	publicKey   string
+	publicKey  string
-	clientType  string
+	clientType string
-	exitNodes   map[string]ExitNode // key is endpoint
+	exitNodes  map[string]ExitNode // key is endpoint
-	updateChan  chan struct{}        // signals the goroutine to refresh exit nodes
+	updateChan chan struct{}       // signals the goroutine to refresh exit nodes
-	publicDNS []string
 
 	sendHolepunchInterval    time.Duration
 	sendHolepunchIntervalMin time.Duration
@@ -50,13 +49,12 @@ const defaultSendHolepunchIntervalMax = 60 * time.Second
 const defaultSendHolepunchIntervalMin = 1 * time.Second
 
 // NewManager creates a new hole punch manager
-func NewManager(sharedBind *bind.SharedBind, ID string, clientType string, publicKey string, publicDNS []string) *Manager {
+func NewManager(sharedBind *bind.SharedBind, ID string, clientType string, publicKey string) *Manager {
 	return &Manager{
 		sharedBind:               sharedBind,
 		ID:                       ID,
 		clientType:               clientType,
 		publicKey:                publicKey,
-		publicDNS:              publicDNS,
 		exitNodes:                make(map[string]ExitNode),
 		sendHolepunchInterval:    defaultSendHolepunchIntervalMin,
 		sendHolepunchIntervalMin: defaultSendHolepunchIntervalMin,
@@ -283,13 +281,7 @@ func (m *Manager) TriggerHolePunch() error {
 	// Send hole punch to all exit nodes
 	successCount := 0
 	for _, exitNode := range currentExitNodes {
-		var host string
+		host, err := util.ResolveDomain(exitNode.Endpoint)
-		var err error
-		if len(m.publicDNS) > 0 {
-			host, err = util.ResolveDomainUpstream(exitNode.Endpoint, m.publicDNS)
-		} else {
-			host, err = util.ResolveDomain(exitNode.Endpoint)
-		}
 		if err != nil {
 			logger.Warn("Failed to resolve endpoint %s: %v", exitNode.Endpoint, err)
 			continue
@@ -400,13 +392,7 @@ func (m *Manager) runMultipleExitNodes() {
 
 		var resolvedNodes []resolvedExitNode
 		for _, exitNode := range currentExitNodes {
-			var host string
+			host, err := util.ResolveDomain(exitNode.Endpoint)
-			var err error
-			if len(m.publicDNS) > 0 {
-				host, err = util.ResolveDomainUpstream(exitNode.Endpoint, m.publicDNS)
-			} else {
-				host, err = util.ResolveDomain(exitNode.Endpoint)
-			}
 			if err != nil {
 				logger.Warn("Failed to resolve endpoint %s: %v", exitNode.Endpoint, err)
 				continue

holepunch/tester.go

@@ -49,11 +49,10 @@ type cachedAddr struct {
 
 // HolepunchTester monitors holepunch connectivity using magic packets
 type HolepunchTester struct {
-	sharedBind  *bind.SharedBind
+	sharedBind *bind.SharedBind
-	publicDNS []string
+	mu         sync.RWMutex
-	mu          sync.RWMutex
+	running    bool
-	running     bool
+	stopChan   chan struct{}
-	stopChan    chan struct{}
 
 	// Pending requests waiting for responses (key: echo data as string)
 	pendingRequests sync.Map // map[string]*pendingRequest
@@ -85,10 +84,9 @@ type pendingRequest struct {
 }
 
 // NewHolepunchTester creates a new holepunch tester using the given SharedBind
-func NewHolepunchTester(sharedBind *bind.SharedBind, publicDNS []string) *HolepunchTester {
+func NewHolepunchTester(sharedBind *bind.SharedBind) *HolepunchTester {
 	return &HolepunchTester{
 		sharedBind:   sharedBind,
-		publicDNS:  publicDNS,
 		addrCache:    make(map[string]*cachedAddr),
 		addrCacheTTL: 5 * time.Minute, // Cache addresses for 5 minutes
 	}
@@ -171,13 +169,7 @@ func (t *HolepunchTester) resolveEndpoint(endpoint string) (*net.UDPAddr, error)
 	}
 
 	// Resolve the endpoint
-	var host string
+	host, err := util.ResolveDomain(endpoint)
-	var err error
-	if len(t.publicDNS) > 0 {
-		host, err = util.ResolveDomainUpstream(endpoint, t.publicDNS)
-	} else {
-		host, err = util.ResolveDomain(endpoint)
-	}
 	if err != nil {
 		host = endpoint
 	}

main.go

@@ -3,13 +3,16 @@ package main
 import (
 	"bytes"
 	"context"
+	"crypto/rand"
 	"crypto/tls"
+	"encoding/hex"
 	"encoding/json"
 	"errors"
 	"flag"
 	"fmt"
 	"net"
 	"net/http"
+	"net/http/pprof"
 	"net/netip"
 	"os"
 	"os/signal"
@@ -45,6 +48,7 @@ type WgData struct {
 	TunnelIP           string               `json:"tunnelIP"`
 	Targets            TargetsByType        `json:"targets"`
 	HealthCheckTargets []healthcheck.Config `json:"healthCheckTargets"`
+	ChainId            string               `json:"chainId"`
 }
 
 type TargetsByType struct {
@@ -127,6 +131,7 @@ var (
 	publicKey                          wgtypes.Key
 	pingStopChan                       chan struct{}
 	stopFunc                           func()
+	pendingRegisterChainId             string
 	healthFile                         string
 	useNativeInterface                 bool
 	authorizedKeysFile                 string
@@ -147,6 +152,7 @@ var (
 	adminAddr         string
 	region            string
 	metricsAsyncBytes bool
+	pprofEnabled      bool
 	blueprintFile     string
 	noCloud           bool
 
@@ -159,6 +165,13 @@ var (
 	tlsPrivateKey string
 )
 
+// generateChainId generates a random chain ID for deduplicating round-trip messages.
+func generateChainId() string {
+	b := make([]byte, 8)
+	_, _ = rand.Read(b)
+	return hex.EncodeToString(b)
+}
+
 func main() {
 	// Check for subcommands first (only principals exits early)
 	if len(os.Args) > 1 {
@@ -225,6 +238,7 @@ func runNewtMain(ctx context.Context) {
 	adminAddrEnv := os.Getenv("NEWT_ADMIN_ADDR")
 	regionEnv := os.Getenv("NEWT_REGION")
 	asyncBytesEnv := os.Getenv("NEWT_METRICS_ASYNC_BYTES")
+	pprofEnabledEnv := os.Getenv("NEWT_PPROF_ENABLED")
 
 	disableClientsEnv := os.Getenv("DISABLE_CLIENTS")
 	disableClients = disableClientsEnv == "true"
@@ -390,6 +404,14 @@ func runNewtMain(ctx context.Context) {
 			metricsAsyncBytes = v
 		}
 	}
+	// pprof debug endpoint toggle
+	if pprofEnabledEnv == "" {
+		flag.BoolVar(&pprofEnabled, "pprof", false, "Enable pprof debug endpoints on admin server")
+	} else {
+		if v, err := strconv.ParseBool(pprofEnabledEnv); err == nil {
+			pprofEnabled = v
+		}
+	}
 	// Optional region flag (resource attribute)
 	if regionEnv == "" {
 		flag.StringVar(&region, "region", "", "Optional region resource attribute (also NEWT_REGION)")
@@ -485,6 +507,14 @@ func runNewtMain(ctx context.Context) {
 		if tel.PrometheusHandler != nil {
 			mux.Handle("/metrics", tel.PrometheusHandler)
 		}
+		if pprofEnabled {
+			mux.HandleFunc("/debug/pprof/", pprof.Index)
+			mux.HandleFunc("/debug/pprof/cmdline", pprof.Cmdline)
+			mux.HandleFunc("/debug/pprof/profile", pprof.Profile)
+			mux.HandleFunc("/debug/pprof/symbol", pprof.Symbol)
+			mux.HandleFunc("/debug/pprof/trace", pprof.Trace)
+			logger.Info("pprof debugging enabled on %s/debug/pprof/", tcfg.AdminAddr)
+		}
 		admin := &http.Server{
 			Addr:              tcfg.AdminAddr,
 			Handler:           otelhttp.NewHandler(mux, "newt-admin"),
@@ -687,6 +717,24 @@ func runNewtMain(ctx context.Context) {
 		defer func() {
 			telemetry.IncSiteRegistration(ctx, regResult)
 		}()
+
+		// Deduplicate using chainId: if the server echoes back a chainId we have
+		// already consumed (or one that doesn't match our current pending request),
+		// throw the message away to avoid setting up the tunnel twice.
+		var chainData struct {
+			ChainId string `json:"chainId"`
+		}
+		if jsonBytes, err := json.Marshal(msg.Data); err == nil {
+			_ = json.Unmarshal(jsonBytes, &chainData)
+		}
+		if chainData.ChainId != "" {
+			if chainData.ChainId != pendingRegisterChainId {
+				logger.Debug("Discarding duplicate/stale newt/wg/connect (chainId=%s, expected=%s)", chainData.ChainId, pendingRegisterChainId)
+				return
+			}
+			pendingRegisterChainId = "" // consume – further duplicates with this id are rejected
+		}
+
 		if stopFunc != nil {
 			stopFunc()     // stop the ws from sending more requests
 			stopFunc = nil // reset stopFunc to nil to avoid double stopping
@@ -952,10 +1000,13 @@ persistent_keepalive_interval=5`, util.FixKey(privateKey.String()), util.FixKey(
 				},
 			}
 
+			chainId := generateChainId()
+			pendingRegisterChainId = chainId
 			stopFunc = client.SendMessageInterval(topicWGRegister, map[string]interface{}{
 				"publicKey":   publicKey.String(),
 				"pingResults": pingResults,
 				"newtVersion": newtVersion,
+				"chainId":     chainId,
 			}, 2*time.Second)
 
 			return
@@ -1055,10 +1106,13 @@ persistent_keepalive_interval=5`, util.FixKey(privateKey.String()), util.FixKey(
 		}
 
 		// Send the ping results to the cloud for selection
+		chainId := generateChainId()
+		pendingRegisterChainId = chainId
 		stopFunc = client.SendMessageInterval(topicWGRegister, map[string]interface{}{
 			"publicKey":   publicKey.String(),
 			"pingResults": pingResults,
 			"newtVersion": newtVersion,
+			"chainId":     chainId,
 		}, 2*time.Second)
 
 		logger.Debug("Sent exit node ping results to cloud for selection: pingResults=%+v", pingResults)
@@ -1721,10 +1775,13 @@ persistent_keepalive_interval=5`, util.FixKey(privateKey.String()), util.FixKey(
 		}
 
 		// Send registration message to the server for backward compatibility
+		bcChainId := generateChainId()
+		pendingRegisterChainId = bcChainId
 		err := client.SendMessage(topicWGRegister, map[string]interface{}{
 			"publicKey":           publicKey.String(),
 			"newtVersion":         newtVersion,
 			"backwardsCompatible": true,
+			"chainId":             bcChainId,
 		})
 
 		sendBlueprint(client)

proxy/manager.go

@@ -21,7 +21,10 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 )
 
-const errUnsupportedProtoFmt = "unsupported protocol: %s"
+const (
+	errUnsupportedProtoFmt = "unsupported protocol: %s"
+	maxUDPPacketSize       = 65507
+)
 
 // Target represents a proxy target with its address and port
 type Target struct {
@@ -105,13 +108,9 @@ func classifyProxyError(err error) string {
 	if errors.Is(err, net.ErrClosed) {
 		return "closed"
 	}
-	if ne, ok := err.(net.Error); ok {
+	var ne net.Error
-		if ne.Timeout() {
+	if errors.As(err, &ne) && ne.Timeout() {
-			return "timeout"
+		return "timeout"
-		}
-		if ne.Temporary() {
-			return "temporary"
-		}
 	}
 	msg := strings.ToLower(err.Error())
 	switch {
@@ -437,14 +436,6 @@ func (pm *ProxyManager) Stop() error {
 		pm.udpConns = append(pm.udpConns[:i], pm.udpConns[i+1:]...)
 	}
 
-	// // Clear the target maps
-	// for k := range pm.tcpTargets {
-	// 	delete(pm.tcpTargets, k)
-	// }
-	// for k := range pm.udpTargets {
-	// 	delete(pm.udpTargets, k)
-	// }
-
 	// Give active connections a chance to close gracefully
 	time.Sleep(100 * time.Millisecond)
 
@@ -498,7 +489,7 @@ func (pm *ProxyManager) handleTCPProxy(listener net.Listener, targetAddr string)
 			if !pm.running {
 				return
 			}
-			if ne, ok := err.(net.Error); ok && !ne.Temporary() {
+			if errors.Is(err, net.ErrClosed) {
 				logger.Info("TCP listener closed, stopping proxy handler for %v", listener.Addr())
 				return
 			}
@@ -564,7 +555,7 @@ func (pm *ProxyManager) handleTCPProxy(listener net.Listener, targetAddr string)
 }
 
 func (pm *ProxyManager) handleUDPProxy(conn *gonet.UDPConn, targetAddr string) {
-	buffer := make([]byte, 65507) // Max UDP packet size
+	buffer := make([]byte, maxUDPPacketSize) // Max UDP packet size
 	clientConns := make(map[string]*net.UDPConn)
 	var clientsMutex sync.RWMutex
 
@@ -583,7 +574,7 @@ func (pm *ProxyManager) handleUDPProxy(conn *gonet.UDPConn, targetAddr string) {
 			}
 
 			// Check for connection closed conditions
-			if err == io.EOF || strings.Contains(err.Error(), "use of closed network connection") {
+			if errors.Is(err, io.EOF) || errors.Is(err, net.ErrClosed) {
 				logger.Info("UDP connection closed, stopping proxy handler")
 
 				// Clean up existing client connections
@@ -662,10 +653,14 @@ func (pm *ProxyManager) handleUDPProxy(conn *gonet.UDPConn, targetAddr string) {
 					telemetry.IncProxyConnectionEvent(context.Background(), tunnelID, "udp", telemetry.ProxyConnectionClosed)
 				}()
 
-				buffer := make([]byte, 65507)
+				buffer := make([]byte, maxUDPPacketSize)
 				for {
 					n, _, err := targetConn.ReadFromUDP(buffer)
 					if err != nil {
+						// Connection closed is normal during cleanup
+						if errors.Is(err, net.ErrClosed) || errors.Is(err, io.EOF) {
+							return // defer will handle cleanup, result stays "success"
+						}
 						logger.Error("Error reading from target: %v", err)
 						result = "failure"
 						return // defer will handle cleanup

util/util.go

@@ -1,7 +1,6 @@
 package util
 
 import (
-	"context"
 	"encoding/base64"
 	"encoding/binary"
 	"encoding/hex"
@@ -15,99 +14,6 @@ import (
 	"golang.zx2c4.com/wireguard/device"
 )
 
-func ResolveDomainUpstream(domain string, publicDNS []string) (string, error) {
-	// trim whitespace
-	domain = strings.TrimSpace(domain)
-
-	// Remove any protocol prefix if present (do this first, before splitting host/port)
-	domain = strings.TrimPrefix(domain, "http://")
-	domain = strings.TrimPrefix(domain, "https://")
-
-	// if there are any trailing slashes, remove them
-	domain = strings.TrimSuffix(domain, "/")
-
-	// Check if there's a port in the domain
-	host, port, err := net.SplitHostPort(domain)
-	if err != nil {
-		// No port found, use the domain as is
-		host = domain
-		port = ""
-	}
-
-	// Check if host is already an IP address (IPv4 or IPv6)
-	// For IPv6, the host from SplitHostPort will already have brackets stripped
-	// but if there was no port, we need to handle bracketed IPv6 addresses
-	cleanHost := strings.TrimPrefix(strings.TrimSuffix(host, "]"), "[")
-	if ip := net.ParseIP(cleanHost); ip != nil {
-		// It's already an IP address, no need to resolve
-		ipAddr := ip.String()
-		if port != "" {
-			return net.JoinHostPort(ipAddr, port), nil
-		}
-		return ipAddr, nil
-	}
-
-	// Lookup IP addresses using the upstream DNS servers if provided
-	var ips []net.IP
-	if len(publicDNS) > 0 {
-		var lastErr error
-		for _, server := range publicDNS {
-			// Ensure the upstream DNS address has a port
-			dnsAddr := server
-			if _, _, err := net.SplitHostPort(dnsAddr); err != nil {
-				// No port specified, default to 53
-				dnsAddr = net.JoinHostPort(server, "53")
-			}
-
-			resolver := &net.Resolver{
-				PreferGo: true,
-				Dial: func(ctx context.Context, network, address string) (net.Conn, error) {
-					d := net.Dialer{}
-					return d.DialContext(ctx, "udp", dnsAddr)
-				},
-			}
-			ips, lastErr = resolver.LookupIP(context.Background(), "ip", host)
-			if lastErr == nil {
-				break
-			}
-		}
-		if lastErr != nil {
-			return "", fmt.Errorf("DNS lookup failed using all upstream servers: %v", lastErr)
-		}
-	} else {
-		ips, err = net.LookupIP(host)
-		if err != nil {
-			return "", fmt.Errorf("DNS lookup failed: %v", err)
-		}
-	}
-
-	if len(ips) == 0 {
-		return "", fmt.Errorf("no IP addresses found for domain %s", host)
-	}
-
-	// Get the first IPv4 address if available
-	var ipAddr string
-	for _, ip := range ips {
-		if ipv4 := ip.To4(); ipv4 != nil {
-			ipAddr = ipv4.String()
-			break
-		}
-	}
-
-	// If no IPv4 found, use the first IP (might be IPv6)
-	if ipAddr == "" {
-		ipAddr = ips[0].String()
-	}
-
-	// Add port back if it existed
-	if port != "" {
-		ipAddr = net.JoinHostPort(ipAddr, port)
-	}
-
-	return ipAddr, nil
-}
-
-
 func ResolveDomain(domain string) (string, error) {
 	// trim whitespace
 	domain = strings.TrimSpace(domain)