Set to 127

Resolves #24

logger/logger.go

@@ -53,7 +53,23 @@ func (l *Logger) log(level LogLevel, format string, args ...interface{}) {
 	if level < l.level {
 		return
 	}
-	timestamp := time.Now().Format("2006/01/02 15:04:05")
+
+	// Get timezone from environment variable or use local timezone
+	timezone := os.Getenv("LOGGER_TIMEZONE")
+	var location *time.Location
+	var err error
+
+	if timezone != "" {
+		location, err = time.LoadLocation(timezone)
+		if err != nil {
+			// If invalid timezone, fall back to local
+			location = time.Local
+		}
+	} else {
+		location = time.Local
+	}
+
+	timestamp := time.Now().In(location).Format("2006/01/02 15:04:05")
 	message := fmt.Sprintf(format, args...)
 	l.logger.Printf("%s: %s %s", level.String(), timestamp, message)
 }

main.go

@@ -23,6 +23,7 @@ import (
 	"github.com/fosrl/newt/proxy"
 	"github.com/fosrl/newt/websocket"
 	"github.com/fosrl/newt/wg"
+	"github.com/fosrl/newt/wgtester"
 
 	"golang.org/x/net/icmp"
 	"golang.org/x/net/ipv4"
@@ -117,7 +118,12 @@ func ping(tnet *netstack.Net, dst string) error {
 }
 
 func startPingCheck(tnet *netstack.Net, serverIP string, stopChan chan struct{}) {
-	ticker := time.NewTicker(10 * time.Second)
+	initialInterval := 10 * time.Second
+	maxInterval := 60 * time.Second
+	currentInterval := initialInterval
+	consecutiveFailures := 0
+
+	ticker := time.NewTicker(currentInterval)
 	defer ticker.Stop()
 
 	go func() {
@@ -126,8 +132,34 @@ func startPingCheck(tnet *netstack.Net, serverIP string, stopChan chan struct{})
 			case <-ticker.C:
 				err := ping(tnet, serverIP)
 				if err != nil {
-					logger.Warn("Periodic ping failed: %v", err)
+					consecutiveFailures++
+					logger.Warn("Periodic ping failed (%d consecutive failures): %v",
+						consecutiveFailures, err)
 					logger.Warn("HINT: Do you have UDP port 51820 (or the port in config.yml) open on your Pangolin server?")
+
+					// Increase interval if we have consistent failures, with a maximum cap
+					if consecutiveFailures >= 3 && currentInterval < maxInterval {
+						// Increase by 50% each time, up to the maximum
+						currentInterval = time.Duration(float64(currentInterval) * 1.5)
+						if currentInterval > maxInterval {
+							currentInterval = maxInterval
+						}
+						ticker.Reset(currentInterval)
+						logger.Info("Increased ping check interval to %v due to consecutive failures",
+							currentInterval)
+					}
+				} else {
+					// On success, if we've backed off, gradually return to normal interval
+					if currentInterval > initialInterval {
+						currentInterval = time.Duration(float64(currentInterval) * 0.8)
+						if currentInterval < initialInterval {
+							currentInterval = initialInterval
+						}
+						ticker.Reset(currentInterval)
+						logger.Info("Decreased ping check interval to %v after successful ping",
+							currentInterval)
+					}
+					consecutiveFailures = 0
 				}
 			case <-stopChan:
 				logger.Info("Stopping ping check")
@@ -137,34 +169,97 @@ func startPingCheck(tnet *netstack.Net, serverIP string, stopChan chan struct{})
 	}()
 }
 
+// Function to track connection status and trigger reconnection as needed
+func monitorConnectionStatus(tnet *netstack.Net, serverIP string, client *websocket.Client) {
+	const checkInterval = 30 * time.Second
+	connectionLost := false
+	ticker := time.NewTicker(checkInterval)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-ticker.C:
+			// Try a ping to see if connection is alive
+			err := ping(tnet, serverIP)
+
+			if err != nil && !connectionLost {
+				// We just lost connection
+				connectionLost = true
+				logger.Warn("Connection to server lost. Continuous reconnection attempts will be made.")
+
+				// Notify the user they might need to check their network
+				logger.Warn("Please check your internet connection and ensure the Pangolin server is online.")
+				logger.Warn("Newt will continue reconnection attempts automatically when connectivity is restored.")
+			} else if err == nil && connectionLost {
+				// Connection has been restored
+				connectionLost = false
+				logger.Info("Connection to server restored!")
+
+				// Tell the server we're back
+				err := client.SendMessage("newt/wg/register", map[string]interface{}{
+					"publicKey": fmt.Sprintf("%s", privateKey.PublicKey()),
+				})
+
+				if err != nil {
+					logger.Error("Failed to send registration message after reconnection: %v", err)
+				} else {
+					logger.Info("Successfully re-registered with server after reconnection")
+				}
+			}
+		}
+	}
+}
+
 func pingWithRetry(tnet *netstack.Net, dst string) error {
 	const (
-		maxAttempts = 5
-		retryDelay  = 2 * time.Second
+		initialMaxAttempts = 15
+		initialRetryDelay  = 2 * time.Second
+		maxRetryDelay      = 60 * time.Second // Cap the maximum delay
 	)
 
-	var lastErr error
-	for attempt := 1; attempt <= maxAttempts; attempt++ {
-		logger.Info("Ping attempt %d of %d", attempt, maxAttempts)
-
-		if err := ping(tnet, dst); err != nil {
-			lastErr = err
-			logger.Warn("Ping attempt %d failed: %v", attempt, err)
-
-			if attempt < maxAttempts {
-				time.Sleep(retryDelay)
-				continue
-			}
-			return fmt.Errorf("all ping attempts failed after %d tries, last error: %w",
-				maxAttempts, lastErr)
-		}
+	attempt := 1
+	retryDelay := initialRetryDelay
 
+	// First try with the initial parameters
+	logger.Info("Ping attempt %d", attempt)
+	if err := ping(tnet, dst); err == nil {
 		// Successful ping
 		return nil
+	} else {
+		logger.Warn("Ping attempt %d failed: %v", attempt, err)
 	}
 
-	// This shouldn't be reached due to the return in the loop, but added for completeness
-	return fmt.Errorf("unexpected error: all ping attempts failed")
+	// Start a goroutine that will attempt pings indefinitely with increasing delays
+	go func() {
+		attempt = 2 // Continue from attempt 2
+
+		for {
+			logger.Info("Ping attempt %d", attempt)
+
+			if err := ping(tnet, dst); err != nil {
+				logger.Warn("Ping attempt %d failed: %v", attempt, err)
+
+				// Increase delay after certain thresholds but cap it
+				if attempt%5 == 0 && retryDelay < maxRetryDelay {
+					retryDelay = time.Duration(float64(retryDelay) * 1.5)
+					if retryDelay > maxRetryDelay {
+						retryDelay = maxRetryDelay
+					}
+					logger.Info("Increasing ping retry delay to %v", retryDelay)
+				}
+
+				time.Sleep(retryDelay)
+				attempt++
+			} else {
+				// Successful ping
+				logger.Info("Ping succeeded after %d attempts", attempt)
+				return
+			}
+		}
+	}()
+
+	// Return an error for the first batch of attempts (to maintain compatibility with existing code)
+	return fmt.Errorf("initial ping attempts failed, continuing in background")
 }
 
 func parseLogLevel(level string) logger.LogLevel {
@@ -263,6 +358,8 @@ var (
 	updownScript         string
 	interfaceName        string
 	generateAndSaveKeyTo string
+	rm                   bool
+	acceptClients        bool
 )
 
 func main() {
@@ -276,6 +373,8 @@ func main() {
 	updownScript = os.Getenv("UPDOWN_SCRIPT")
 	interfaceName = os.Getenv("INTERFACE")
 	generateAndSaveKeyTo = os.Getenv("GENERATE_AND_SAVE_KEY_TO")
+	rm = os.Getenv("RM") == "true"
+	acceptClients = os.Getenv("ACCEPT_CLIENTS") == "true"
 
 	if endpoint == "" {
 		flag.StringVar(&endpoint, "endpoint", "", "Endpoint of your pangolin server")
@@ -302,8 +401,10 @@ func main() {
 		flag.StringVar(&interfaceName, "interface", "wg1", "Name of the WireGuard interface")
 	}
 	if generateAndSaveKeyTo == "" {
-		flag.StringVar(&generateAndSaveKeyTo, "generateAndSaveKeyTo", "", "Path to save generated private key")
+		flag.StringVar(&generateAndSaveKeyTo, "generateAndSaveKeyTo", "/tmp/newtkey", "Path to save generated private key")
 	}
+	flag.BoolVar(&rm, "rm", false, "Remove the WireGuard interface")
+	flag.BoolVar(&acceptClients, "accept-clients", false, "Accept clients on the WireGuard interface")
 
 	// do a --version check
 	version := flag.Bool("version", false, "Print the version")
@@ -348,8 +449,9 @@ func main() {
 	var pm *proxy.ProxyManager
 	var connected bool
 	var wgData WgData
+	var wgTesterServer *wgtester.Server
 
-	if generateAndSaveKeyTo != "" {
+	if acceptClients {
 		// make sure we are running on linux
 		if runtime.GOOS != "linux" {
 			logger.Fatal("Tunnel management is only supported on Linux right now!")
@@ -370,7 +472,16 @@ func main() {
 		if err != nil {
 			logger.Fatal("Failed to create WireGuard service: %v", err)
 		}
-		defer wgService.Close()
+		defer wgService.Close(rm)
+
+		wgTesterServer = wgtester.NewServer("0.0.0.0", wgService.Port, id) // TODO: maybe make this the same ip of the wg server?
+		err := wgTesterServer.Start()
+		if err != nil {
+			logger.Error("Failed to start WireGuard tester server: %v", err)
+		} else {
+			// Make sure to stop the server on exit
+			defer wgTesterServer.Stop()
+		}
 	}
 
 	client.RegisterHandler("newt/terminate", func(msg websocket.WSMessage) {
@@ -393,13 +504,8 @@ func main() {
 
 		if connected {
 			logger.Info("Already connected! But I will send a ping anyway...")
-			// ping(tnet, wgData.ServerIP)
-			err = pingWithRetry(tnet, wgData.ServerIP)
-			if err != nil {
-				// Handle complete failure after all retries
-				logger.Warn("Failed to ping %s: %v", wgData.ServerIP, err)
-				logger.Warn("HINT: Do you have UDP port 51820 (or the port in config.yml) open on your Pangolin server?")
-			}
+			// Even if pingWithRetry returns an error, it will continue trying in the background
+			_ = pingWithRetry(tnet, wgData.ServerIP) // Ignoring initial error as pings will continue
 			return
 		}
 
@@ -458,18 +564,18 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey(
 		}
 
 		logger.Info("WireGuard device created. Lets ping the server now...")
-		// Ping to bring the tunnel up on the server side quickly
-		// ping(tnet, wgData.ServerIP)
-		err = pingWithRetry(tnet, wgData.ServerIP)
-		if err != nil {
-			// Handle complete failure after all retries
-			logger.Error("Failed to ping %s: %v", wgData.ServerIP, err)
-			fmt.Sprintf("%s", privateKey)
-		}
 
+		// Even if pingWithRetry returns an error, it will continue trying in the background
+		_ = pingWithRetry(tnet, wgData.ServerIP)
+
+		// Always mark as connected and start the proxy manager regardless of initial ping result
+		// as the pings will continue in the background
 		if !connected {
 			logger.Info("Starting ping check")
 			startPingCheck(tnet, wgData.ServerIP, pingStopChan)
+
+			// Start connection monitoring in a separate goroutine
+			go monitorConnectionStatus(tnet, wgData.ServerIP, client)
 		}
 
 		// Create proxy manager
@@ -619,7 +725,11 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey(
 	dev.Close()
 
 	if wgService != nil {
-		wgService.Close()
+		wgService.Close(rm)
+	}
+
+	if wgTesterServer != nil {
+		wgTesterServer.Stop()
 	}
 
 	if pm != nil {

network/network.go

@@ -193,10 +193,3 @@ func ParseResponse(response []byte) (net.IP, uint16) {
 	port := binary.BigEndian.Uint16(response[4:6])
 	return ip, port
 }
-
-func parseForBPF(response []byte) (srcIP net.IP, srcPort uint16, dstPort uint16) {
-	srcIP = net.IP(response[12:16])
-	srcPort = binary.BigEndian.Uint16(response[20:22])
-	dstPort = binary.BigEndian.Uint16(response[22:24])
-	return
-}

wg/wg.go

@@ -60,6 +60,7 @@ type WireGuardService struct {
 	host          string
 	serverPubKey  string
 	token         string
+	stopGetConfig chan struct{}
 }
 
 // Add this type definition
@@ -80,13 +81,20 @@ func NewFixedPortBind(port uint16) conn.Bind {
 	}
 }
 
+// find an available UDP port in the range [minPort, maxPort] and also the next port for the wgtester
 func FindAvailableUDPPort(minPort, maxPort uint16) (uint16, error) {
 	if maxPort < minPort {
 		return 0, fmt.Errorf("invalid port range: min=%d, max=%d", minPort, maxPort)
 	}
 
-	// Create a slice of all ports in the range
-	portRange := make([]uint16, maxPort-minPort+1)
+	// We need to check port+1 as well, so adjust the max port to avoid going out of range
+	adjustedMaxPort := maxPort - 1
+	if adjustedMaxPort < minPort {
+		return 0, fmt.Errorf("insufficient port range to find consecutive ports: min=%d, max=%d", minPort, maxPort)
+	}
+
+	// Create a slice of all ports in the range (excluding the last one)
+	portRange := make([]uint16, adjustedMaxPort-minPort+1)
 	for i := range portRange {
 		portRange[i] = minPort + uint16(i)
 	}
@@ -100,20 +108,35 @@ func FindAvailableUDPPort(minPort, maxPort uint16) (uint16, error) {
 
 	// Try each port in the randomized order
 	for _, port := range portRange {
-		addr := &net.UDPAddr{
+		// Check if port is available
+		addr1 := &net.UDPAddr{
 			IP:   net.ParseIP("127.0.0.1"),
 			Port: int(port),
 		}
-		conn, err := net.ListenUDP("udp", addr)
-		if err != nil {
+		conn1, err1 := net.ListenUDP("udp", addr1)
+		if err1 != nil {
 			continue // Port is in use or there was an error, try next port
 		}
-		_ = conn.SetDeadline(time.Now())
-		conn.Close()
+
+		// Check if port+1 is also available
+		addr2 := &net.UDPAddr{
+			IP:   net.ParseIP("127.0.0.1"),
+			Port: int(port + 1),
+		}
+		conn2, err2 := net.ListenUDP("udp", addr2)
+		if err2 != nil {
+			// The next port is not available, so close the first connection and try again
+			conn1.Close()
+			continue
+		}
+
+		// Both ports are available, close connections and return the first port
+		conn1.Close()
+		conn2.Close()
 		return port, nil
 	}
 
-	return 0, fmt.Errorf("no available UDP ports found in range %d-%d", minPort, maxPort)
+	return 0, fmt.Errorf("no available consecutive UDP ports found in range %d-%d", minPort, maxPort)
 }
 
 func NewWireGuardService(interfaceName string, mtu int, generateAndSaveKeyTo string, host string, newtId string, wsClient *websocket.Client) (*WireGuardService, error) {
@@ -146,12 +169,6 @@ func NewWireGuardService(interfaceName string, mtu int, generateAndSaveKeyTo str
 		}
 	}
 
-	port, err := FindAvailableUDPPort(49152, 65535)
-	if err != nil {
-		fmt.Printf("Error finding available port: %v\n", err)
-		return nil, err
-	}
-
 	service := &WireGuardService{
 		interfaceName: interfaceName,
 		mtu:           mtu,
@@ -159,25 +176,60 @@ func NewWireGuardService(interfaceName string, mtu int, generateAndSaveKeyTo str
 		wgClient:      wgClient,
 		key:           key,
 		newtId:        newtId,
-		lastReadings:  make(map[string]PeerReading),
-		Port:          port,
-		stopHolepunch: make(chan struct{}),
 		host:          host,
+		lastReadings:  make(map[string]PeerReading),
+		stopHolepunch: make(chan struct{}),
+		stopGetConfig: make(chan struct{}),
+	}
+
+	// Get the existing wireguard port (keep this part)
+	device, err := service.wgClient.Device(service.interfaceName)
+	if err == nil {
+		service.Port = uint16(device.ListenPort)
+		logger.Info("WireGuard interface %s already exists with port %d\n", service.interfaceName, service.Port)
+	} else {
+		service.Port, err = FindAvailableUDPPort(49152, 65535)
+		if err != nil {
+			fmt.Printf("Error finding available port: %v\n", err)
+			return nil, err
+		}
 	}
 
 	// Register websocket handlers
 	wsClient.RegisterHandler("newt/wg/receive-config", service.handleConfig)
 	wsClient.RegisterHandler("newt/wg/peer/add", service.handleAddPeer)
 	wsClient.RegisterHandler("newt/wg/peer/remove", service.handleRemovePeer)
+	wsClient.RegisterHandler("newt/wg/peer/update", service.handleUpdatePeer)
+
+	if err := service.sendUDPHolePunch(service.host + ":21820"); err != nil {
+		logger.Error("Failed to send UDP hole punch: %v", err)
+	}
+
+	// start the UDP holepunch
+	go service.keepSendingUDPHolePunch(service.host)
 
 	return service, nil
 }
 
-func (s *WireGuardService) Close() {
+func (s *WireGuardService) Close(rm bool) {
+	select {
+	case <-s.stopGetConfig:
+		// Already closed, do nothing
+	default:
+		close(s.stopGetConfig)
+	}
+
 	s.wgClient.Close()
 	// Remove the WireGuard interface
-	if err := s.removeInterface(); err != nil {
-		logger.Error("Failed to remove WireGuard interface: %v", err)
+	if rm {
+		if err := s.removeInterface(); err != nil {
+			logger.Error("Failed to remove WireGuard interface: %v", err)
+		}
+
+		// Remove the private key file
+		if err := os.Remove(s.key.String()); err != nil {
+			logger.Error("Failed to remove private key file: %v", err)
+		}
 	}
 }
 
@@ -190,24 +242,15 @@ func (s *WireGuardService) SetToken(token string) {
 }
 
 func (s *WireGuardService) LoadRemoteConfig() error {
-
-	// get the exising wireguard port
-	device, err := s.wgClient.Device(s.interfaceName)
-	if err == nil {
-		s.Port = uint16(device.ListenPort)
-		logger.Info("WireGuard interface %s already exists with port %d\n", s.interfaceName, s.Port)
-	}
-
-	err = s.client.SendMessage("newt/wg/get-config", map[string]interface{}{
-		"publicKey": fmt.Sprintf("%s", s.key.PublicKey().String()),
-		"port":      s.Port,
-	})
+	// Send the initial message
+	err := s.sendGetConfigMessage()
 	if err != nil {
-		logger.Error("Failed to send registration message: %v", err)
+		logger.Error("Failed to send initial get-config message: %v", err)
 		return err
 	}
 
-	logger.Info("Requesting WireGuard configuration from remote server")
+	// Start goroutine to periodically send the message until config is received
+	go s.keepSendingGetConfig()
 
 	go s.periodicBandwidthCheck()
 
@@ -231,6 +274,8 @@ func (s *WireGuardService) handleConfig(msg websocket.WSMessage) {
 	}
 	s.config = config
 
+	close(s.stopGetConfig)
+
 	// Ensure the WireGuard interface and peers are configured
 	if err := s.ensureWireguardInterface(config); err != nil {
 		logger.Error("Failed to ensure WireGuard interface: %v", err)
@@ -239,13 +284,6 @@ func (s *WireGuardService) handleConfig(msg websocket.WSMessage) {
 	if err := s.ensureWireguardPeers(config.Peers); err != nil {
 		logger.Error("Failed to ensure WireGuard peers: %v", err)
 	}
-
-	if err := s.sendUDPHolePunch(s.host + ":21820"); err != nil {
-		logger.Error("Failed to send UDP hole punch: %v", err)
-	}
-
-	// start the UDP holepunch
-	go s.keepSendingUDPHolePunch(s.host)
 }
 
 func (s *WireGuardService) ensureWireguardInterface(wgconfig WgConfig) error {
@@ -407,6 +445,7 @@ func (s *WireGuardService) ensureWireguardPeers(peers []Peer) error {
 }
 
 func (s *WireGuardService) handleAddPeer(msg websocket.WSMessage) {
+	logger.Info("Received message: %v", msg.Data)
 	var peer Peer
 
 	jsonData, err := json.Marshal(msg.Data)
@@ -450,8 +489,6 @@ func (s *WireGuardService) addPeer(peer Peer) error {
 			return fmt.Errorf("failed to resolve endpoint address: %w", err)
 		}
 
-		// make the endpoint localhost to test
-
 		peerConfig = wgtypes.PeerConfig{
 			PublicKey:                   pubKey,
 			AllowedIPs:                  allowedIPs,
@@ -481,6 +518,7 @@ func (s *WireGuardService) addPeer(peer Peer) error {
 }
 
 func (s *WireGuardService) handleRemovePeer(msg websocket.WSMessage) {
+	logger.Info("Received message: %v", msg.Data)
 	// parse the publicKey from the message which is json { "publicKey": "asdfasdfl;akjsdf" }
 	type RemoveRequest struct {
 		PublicKey string `json:"publicKey"`
@@ -527,6 +565,115 @@ func (s *WireGuardService) removePeer(publicKey string) error {
 	return nil
 }
 
+func (s *WireGuardService) handleUpdatePeer(msg websocket.WSMessage) {
+	logger.Info("Received message: %v", msg.Data)
+	// Define a struct to match the incoming message structure with optional fields
+	type UpdatePeerRequest struct {
+		PublicKey  string   `json:"publicKey"`
+		AllowedIPs []string `json:"allowedIps,omitempty"`
+		Endpoint   string   `json:"endpoint,omitempty"`
+	}
+	jsonData, err := json.Marshal(msg.Data)
+	if err != nil {
+		logger.Info("Error marshaling data: %v", err)
+		return
+	}
+	var request UpdatePeerRequest
+	if err := json.Unmarshal(jsonData, &request); err != nil {
+		logger.Info("Error unmarshaling peer data: %v", err)
+		return
+	}
+	// First, get the current peer configuration to preserve any unmodified fields
+	device, err := s.wgClient.Device(s.interfaceName)
+	if err != nil {
+		logger.Info("Error getting WireGuard device: %v", err)
+		return
+	}
+	pubKey, err := wgtypes.ParseKey(request.PublicKey)
+	if err != nil {
+		logger.Info("Error parsing public key: %v", err)
+		return
+	}
+	// Find the existing peer configuration
+	var currentPeer *wgtypes.Peer
+	for _, p := range device.Peers {
+		if p.PublicKey == pubKey {
+			currentPeer = &p
+			break
+		}
+	}
+	if currentPeer == nil {
+		logger.Info("Peer %s not found, cannot update", request.PublicKey)
+		return
+	}
+	// Create the update peer config
+	peerConfig := wgtypes.PeerConfig{
+		PublicKey:  pubKey,
+		UpdateOnly: true,
+	}
+	// Keep the default persistent keepalive of 1 second
+	keepalive := time.Second
+	peerConfig.PersistentKeepaliveInterval = &keepalive
+
+	// Handle Endpoint field special case
+	// If Endpoint is included in the request but empty, we want to remove the endpoint
+	// If Endpoint is not included, we don't modify it
+	endpointSpecified := false
+	for key := range msg.Data.(map[string]interface{}) {
+		if key == "endpoint" {
+			endpointSpecified = true
+			break
+		}
+	}
+
+	// Only update AllowedIPs if provided in the request
+	if request.AllowedIPs != nil && len(request.AllowedIPs) > 0 {
+		var allowedIPs []net.IPNet
+		for _, ipStr := range request.AllowedIPs {
+			_, ipNet, err := net.ParseCIDR(ipStr)
+			if err != nil {
+				logger.Info("Error parsing allowed IP %s: %v", ipStr, err)
+				return
+			}
+			allowedIPs = append(allowedIPs, *ipNet)
+		}
+		peerConfig.AllowedIPs = allowedIPs
+		peerConfig.ReplaceAllowedIPs = true
+		logger.Info("Updating AllowedIPs for peer %s", request.PublicKey)
+	} else if endpointSpecified && request.Endpoint == "" {
+		peerConfig.ReplaceAllowedIPs = false
+	}
+
+	if endpointSpecified {
+		if request.Endpoint != "" {
+			// Update to new endpoint
+			endpoint, err := net.ResolveUDPAddr("udp", request.Endpoint)
+			if err != nil {
+				logger.Info("Error resolving endpoint address %s: %v", request.Endpoint, err)
+				return
+			}
+			peerConfig.Endpoint = endpoint
+			logger.Info("Updating Endpoint for peer %s to %s", request.PublicKey, request.Endpoint)
+		} else {
+			// specify any address to listen for any incoming packets
+			peerConfig.Endpoint = &net.UDPAddr{
+				IP: net.IPv4(127, 0, 0, 1),
+			}
+			logger.Info("Removing Endpoint for peer %s", request.PublicKey)
+		}
+	}
+
+	// Apply the configuration update
+	config := wgtypes.Config{
+		Peers: []wgtypes.PeerConfig{peerConfig},
+	}
+	if err := s.wgClient.ConfigureDevice(s.interfaceName, config); err != nil {
+		logger.Info("Error updating peer configuration: %v", err)
+		return
+	}
+	logger.Info("Peer %s updated successfully", request.PublicKey)
+}
+
 func (s *WireGuardService) periodicBandwidthCheck() {
 	ticker := time.NewTicker(10 * time.Second)
 	defer ticker.Stop()
@@ -642,7 +789,8 @@ func (s *WireGuardService) reportPeerBandwidth() error {
 func (s *WireGuardService) sendUDPHolePunch(serverAddr string) error {
 
 	if s.serverPubKey == "" || s.token == "" {
-		return fmt.Errorf("server public key or token is not set")
+		logger.Debug("Server public key or token not set, skipping UDP hole punch")
+		return nil
 	}
 
 	// Parse server address
@@ -799,3 +947,33 @@ func (s *WireGuardService) removeInterface() error {
 
 	return nil
 }
+
+func (s *WireGuardService) sendGetConfigMessage() error {
+	err := s.client.SendMessage("newt/wg/get-config", map[string]interface{}{
+		"publicKey": fmt.Sprintf("%s", s.key.PublicKey().String()),
+		"port":      s.Port,
+	})
+	if err != nil {
+		logger.Error("Failed to send get-config message: %v", err)
+		return err
+	}
+	logger.Info("Requesting WireGuard configuration from remote server")
+	return nil
+}
+
+func (s *WireGuardService) keepSendingGetConfig() {
+	ticker := time.NewTicker(3 * time.Second)
+	defer ticker.Stop()
+
+	for {
+		select {
+		case <-s.stopGetConfig:
+			logger.Info("Stopping get-config messages")
+			return
+		case <-ticker.C:
+			if err := s.sendGetConfigMessage(); err != nil {
+				logger.Error("Failed to send periodic get-config: %v", err)
+			}
+		}
+	}
+}

wgtester/wgtester.go

@@ -0,0 +1,164 @@
+package wgtester
+
+import (
+	"encoding/binary"
+	"fmt"
+	"net"
+	"sync"
+	"time"
+
+	"github.com/fosrl/newt/logger"
+)
+
+const (
+	// Magic bytes to identify our packets
+	magicHeader uint32 = 0xDEADBEEF
+	// Request packet type
+	packetTypeRequest uint8 = 1
+	// Response packet type
+	packetTypeResponse uint8 = 2
+	// Packet format:
+	// - 4 bytes: magic header (0xDEADBEEF)
+	// - 1 byte: packet type (1 = request, 2 = response)
+	// - 8 bytes: timestamp (for round-trip timing)
+	packetSize = 13
+)
+
+// Server handles listening for connection check requests using UDP
+type Server struct {
+	conn         *net.UDPConn
+	serverAddr   string
+	serverPort   uint16
+	shutdownCh   chan struct{}
+	isRunning    bool
+	runningLock  sync.Mutex
+	newtID       string
+	outputPrefix string
+}
+
+// NewServer creates a new connection test server using UDP
+func NewServer(serverAddr string, serverPort uint16, newtID string) *Server {
+	return &Server{
+		serverAddr:   serverAddr,
+		serverPort:   serverPort + 1, // use the next port for the server
+		shutdownCh:   make(chan struct{}),
+		newtID:       newtID,
+		outputPrefix: "[WGTester] ",
+	}
+}
+
+// Start begins listening for connection test packets using UDP
+func (s *Server) Start() error {
+	s.runningLock.Lock()
+	defer s.runningLock.Unlock()
+
+	if s.isRunning {
+		return nil
+	}
+
+	//create the address to listen on
+	addr := net.JoinHostPort(s.serverAddr, fmt.Sprintf("%d", s.serverPort))
+
+	// Create UDP address to listen on
+	udpAddr, err := net.ResolveUDPAddr("udp", addr)
+	if err != nil {
+		return err
+	}
+
+	// Create UDP connection
+	conn, err := net.ListenUDP("udp", udpAddr)
+	if err != nil {
+		return err
+	}
+	s.conn = conn
+
+	s.isRunning = true
+	go s.handleConnections()
+
+	logger.Info("%sServer started on %s:%d", s.outputPrefix, s.serverAddr, s.serverPort)
+	return nil
+}
+
+// Stop shuts down the server
+func (s *Server) Stop() {
+	s.runningLock.Lock()
+	defer s.runningLock.Unlock()
+
+	if !s.isRunning {
+		return
+	}
+
+	close(s.shutdownCh)
+	if s.conn != nil {
+		s.conn.Close()
+	}
+	s.isRunning = false
+	logger.Info(s.outputPrefix + "Server stopped")
+}
+
+// handleConnections processes incoming packets
+func (s *Server) handleConnections() {
+	buffer := make([]byte, 2000) // Buffer large enough for any UDP packet
+
+	for {
+		select {
+		case <-s.shutdownCh:
+			return
+		default:
+			// Set read deadline to avoid blocking forever
+			err := s.conn.SetReadDeadline(time.Now().Add(1 * time.Second))
+			if err != nil {
+				logger.Error(s.outputPrefix+"Error setting read deadline: %v", err)
+				continue
+			}
+
+			// Read from UDP connection
+			n, addr, err := s.conn.ReadFromUDP(buffer)
+			if err != nil {
+				if netErr, ok := err.(net.Error); ok && netErr.Timeout() {
+					// Just a timeout, keep going
+					continue
+				}
+				logger.Error(s.outputPrefix+"Error reading from UDP: %v", err)
+				continue
+			}
+
+			// Process packet only if it meets minimum size requirements
+			if n < packetSize {
+				continue // Too small to be our packet
+			}
+
+			// Check magic header
+			magic := binary.BigEndian.Uint32(buffer[0:4])
+			if magic != magicHeader {
+				continue // Not our packet
+			}
+
+			// Check packet type
+			packetType := buffer[4]
+			if packetType != packetTypeRequest {
+				continue // Not a request packet
+			}
+
+			// Create response packet
+			responsePacket := make([]byte, packetSize)
+			// Copy the same magic header
+			binary.BigEndian.PutUint32(responsePacket[0:4], magicHeader)
+			// Change the packet type to response
+			responsePacket[4] = packetTypeResponse
+			// Copy the timestamp (for RTT calculation)
+			copy(responsePacket[5:13], buffer[5:13])
+
+			// Log response being sent for debugging
+			logger.Debug(s.outputPrefix+"Sending response to %s", addr.String())
+
+			// Send the response packet directly to the source address
+			_, err = s.conn.WriteToUDP(responsePacket, addr)
+			if err != nil {
+				logger.Error(s.outputPrefix+"Error sending response: %v", err)
+			} else {
+				logger.Debug(s.outputPrefix + "Response sent successfully")
+			}
+		}
+	}
+}