Add test udp server and client

Ref PAN-158

.github/workflows/test.yml

@@ -25,7 +25,7 @@ jobs:
         run: go build
 
       - name: Build Docker image
-        run: make build
+        run: make docker-build-release
 
       - name: Build binaries
         run: make go-build-release

clients/clients.go

@@ -46,6 +46,7 @@ type Target struct {
 type PortRange struct {
 	Min uint16 `json:"min"`
 	Max uint16 `json:"max"`
+	Protocol string `json:"protocol"` // "tcp" or "udp"
 }
 
 type Peer struct {
@@ -701,6 +702,7 @@ func (s *WireGuardService) ensureTargets(targets []Target) error {
 			portRanges = append(portRanges, netstack2.PortRange{
 				Min: pr.Min,
 				Max: pr.Max,
+				Protocol: pr.Protocol,
 			})
 		}
 

flake.nix

@@ -25,7 +25,7 @@
           inherit (pkgs) lib;
 
           # Update version when releasing
-          version = "1.6.1";
+          version = "1.7.0";
         in
         {
           default = self.packages.${system}.pangolin-newt;
@@ -35,7 +35,7 @@
             inherit version;
             src = pkgs.nix-gitignore.gitignoreSource [ ] ./.;
 
-            vendorHash = "sha256-krxkfH+4z0rgmFi8OyCzWIrxU5Rpb7gjtGcn3LnZCig=";
+            vendorHash = "sha256-5Xr6mwPtsqEliKeKv2rhhp6JC7u3coP4nnhIxGMqccU=";
 
             env = {
               CGO_ENABLED = 0;

healthcheck/healthcheck.go

@@ -58,7 +58,7 @@ type Target struct {
 	LastCheck  time.Time `json:"lastCheck"`
 	LastError  string    `json:"lastError,omitempty"`
 	CheckCount int       `json:"checkCount"`
-	ticker     *time.Ticker
+	timer      *time.Timer
 	ctx        context.Context
 	cancel     context.CancelFunc
 }
@@ -304,26 +304,26 @@ func (m *Monitor) monitorTarget(target *Target) {
 		go m.callback(m.GetTargets())
 	}
 
-	// Set up ticker based on current status
+	// Set up timer based on current status
 	interval := time.Duration(target.Config.Interval) * time.Second
 	if target.Status == StatusUnhealthy {
 		interval = time.Duration(target.Config.UnhealthyInterval) * time.Second
 	}
 
 	logger.Debug("Target %d: initial check interval set to %v", target.Config.ID, interval)
-	target.ticker = time.NewTicker(interval)
-	defer target.ticker.Stop()
+	target.timer = time.NewTimer(interval)
+	defer target.timer.Stop()
 
 	for {
 		select {
 		case <-target.ctx.Done():
 			logger.Info("Stopping health check monitoring for target %d", target.Config.ID)
 			return
-		case <-target.ticker.C:
+		case <-target.timer.C:
 			oldStatus := target.Status
 			m.performHealthCheck(target)
 
-			// Update ticker interval if status changed
+			// Update timer interval if status changed
 			newInterval := time.Duration(target.Config.Interval) * time.Second
 			if target.Status == StatusUnhealthy {
 				newInterval = time.Duration(target.Config.UnhealthyInterval) * time.Second
@@ -332,11 +332,12 @@ func (m *Monitor) monitorTarget(target *Target) {
 			if newInterval != interval {
 				logger.Debug("Target %d: updating check interval from %v to %v due to status change",
 					target.Config.ID, interval, newInterval)
-				target.ticker.Stop()
-				target.ticker = time.NewTicker(newInterval)
 				interval = newInterval
 			}
 
+			// Reset timer for next check with current interval
+			target.timer.Reset(interval)
+
 			// Notify callback if status changed
 			if oldStatus != target.Status && m.callback != nil {
 				logger.Info("Target %d status changed: %s -> %s",

main.go

@@ -388,6 +388,13 @@ func runNewtMain(ctx context.Context) {
 		tlsClientCAs = append(tlsClientCAs, tlsClientCAsFlag...)
 	}
 
+	if *version {
+		fmt.Println("Newt version " + newtVersion)
+		os.Exit(0)
+	} else {
+		logger.Info("Newt version %s", newtVersion)
+	}
+
 	logger.Init(nil)
 	loggerLevel := util.ParseLogLevel(logLevel)
 	logger.GetLogger().SetLevel(loggerLevel)
@@ -439,13 +446,6 @@ func runNewtMain(ctx context.Context) {
 		defer func() { _ = tel.Shutdown(context.Background()) }()
 	}
 
-	if *version {
-		fmt.Println("Newt version " + newtVersion)
-		os.Exit(0)
-	} else {
-		logger.Info("Newt version %s", newtVersion)
-	}
-
 	if err := updates.CheckForUpdate("fosrl", "newt", newtVersion); err != nil {
 		logger.Error("Error checking for updates: %v\n", err)
 	}

netstack2/proxy.go

@@ -22,10 +22,12 @@ import (
 	"gvisor.dev/gvisor/pkg/tcpip/transport/udp"
 )
 
-// PortRange represents an allowed range of ports (inclusive)
+// PortRange represents an allowed range of ports (inclusive) with optional protocol filtering
+// Protocol can be "tcp", "udp", or "" (empty string means both protocols)
 type PortRange struct {
-	Min uint16
-	Max uint16
+	Min      uint16
+	Max      uint16
+	Protocol string // "tcp", "udp", or "" for both
 }
 
 // SubnetRule represents a subnet with optional port restrictions and source address
@@ -97,14 +99,16 @@ func (sl *SubnetLookup) RemoveSubnet(sourcePrefix, destPrefix netip.Prefix) {
 	delete(sl.rules, key)
 }
 
-// Match checks if a source IP, destination IP, and port match any subnet rule
-// Returns the matched rule if BOTH:
+// Match checks if a source IP, destination IP, port, and protocol match any subnet rule
+// Returns the matched rule if ALL of these conditions are met:
 //   - The source IP is in the rule's source prefix
 //   - The destination IP is in the rule's destination prefix
 //   - The port is in an allowed range (or no port restrictions exist)
+//   - The protocol matches (or the port range allows both protocols)
 //
+// proto should be header.TCPProtocolNumber or header.UDPProtocolNumber
 // Returns nil if no rule matches
-func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16) *SubnetRule {
+func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16, proto tcpip.TransportProtocolNumber) *SubnetRule {
 	sl.mu.RLock()
 	defer sl.mu.RUnlock()
 
@@ -125,10 +129,20 @@ func (sl *SubnetLookup) Match(srcIP, dstIP netip.Addr, port uint16) *SubnetRule
 			return rule
 		}
 
-		// Check if port is in any of the allowed ranges
+		// Check if port and protocol are in any of the allowed ranges
 		for _, pr := range rule.PortRanges {
 			if port >= pr.Min && port <= pr.Max {
-				return rule
+				// Check protocol compatibility
+				if pr.Protocol == "" {
+					// Empty protocol means allow both TCP and UDP
+					return rule
+				}
+				// Check if the packet protocol matches the port range protocol
+				if (pr.Protocol == "tcp" && proto == header.TCPProtocolNumber) ||
+					(pr.Protocol == "udp" && proto == header.UDPProtocolNumber) {
+					return rule
+				}
+				// Port matches but protocol doesn't - continue checking other ranges
 			}
 		}
 	}
@@ -412,8 +426,8 @@ func (p *ProxyHandler) HandleIncomingPacket(packet []byte) bool {
 		dstPort = 0
 	}
 
-	// Check if the source IP, destination IP, and port match any subnet rule
-	matchedRule := p.subnetLookup.Match(srcAddr, dstAddr, dstPort)
+	// Check if the source IP, destination IP, port, and protocol match any subnet rule
+	matchedRule := p.subnetLookup.Match(srcAddr, dstAddr, dstPort, protocol)
 	if matchedRule != nil {
 		// Check if we need to perform DNAT
 		if matchedRule.RewriteTo != "" {

udp_client.py

@@ -0,0 +1,49 @@
+import socket
+import sys
+
+# Argument parsing: Check if IP and Port are provided
+if len(sys.argv) != 3:
+    print("Usage: python udp_client.py <HOST_IP> <HOST_PORT>")
+    # Example: python udp_client.py 127.0.0.1 12000
+    sys.exit(1)
+
+HOST = sys.argv[1]
+try:
+    PORT = int(sys.argv[2])
+except ValueError:
+    print("Error: HOST_PORT must be an integer.")
+    sys.exit(1)
+
+# The message to send to the server
+MESSAGE = "Hello UDP Server! How are you?"
+
+# Create a UDP socket
+try:
+    client_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+except socket.error as err:
+    print(f"Failed to create socket: {err}")
+    sys.exit()
+
+try:
+    print(f"Sending message to {HOST}:{PORT}...")
+
+    # Send the message (data must be encoded to bytes)
+    client_socket.sendto(MESSAGE.encode('utf-8'), (HOST, PORT))
+
+    # Wait for the server's response (buffer size 1024 bytes)
+    data, server_address = client_socket.recvfrom(1024)
+
+    # Decode and print the server's response
+    response = data.decode('utf-8')
+    print("-" * 30)
+    print(f"Received response from server {server_address[0]}:{server_address[1]}:")
+    print(f"-> Data: '{response}'")
+
+except socket.error as err:
+    print(f"Error during communication: {err}")
+
+finally:
+    # Close the socket
+    client_socket.close()
+    print("-" * 30)
+    print("Client finished and socket closed.")

udp_server.py

@@ -0,0 +1,58 @@
+import socket
+import sys
+
+# optionally take in some positional args for the port
+if len(sys.argv) > 1:
+    try:
+        PORT = int(sys.argv[1])
+    except ValueError:
+        print("Invalid port number. Using default port 12000.")
+        PORT = 12000
+else:
+    PORT = 12000
+
+# Define the server host and port
+HOST = '0.0.0.0'  # Standard loopback interface address (localhost)
+
+# Create a UDP socket
+try:
+    server_socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
+except socket.error as err:
+    print(f"Failed to create socket: {err}")
+    sys.exit()
+
+# Bind the socket to the address
+try:
+    server_socket.bind((HOST, PORT))
+    print(f"UDP Server listening on {HOST}:{PORT}")
+except socket.error as err:
+    print(f"Bind failed: {err}")
+    server_socket.close()
+    sys.exit()
+
+# Wait for and process incoming data
+while True:
+    try:
+        # Receive data and the client's address (buffer size 1024 bytes)
+        data, client_address = server_socket.recvfrom(1024)
+
+        # Decode the data and print the message
+        message = data.decode('utf-8')
+        print("-" * 30)
+        print(f"Received message from {client_address[0]}:{client_address[1]}:")
+        print(f"-> Data: '{message}'")
+
+        # Prepare the response message
+        response_message = f"Hello client! Server received: '{message.upper()}'"
+
+        # Send the response back to the client
+        server_socket.sendto(response_message.encode('utf-8'), client_address)
+        print(f"Sent response back to client.")
+
+    except Exception as e:
+        print(f"An error occurred: {e}")
+        break
+
+# Clean up (though usually unreachable in an infinite server loop)
+server_socket.close()
+print("Server stopped.")