Merge branch 'dev'

Might help #7

.github/workflows/cicd.yml

@@ -0,0 +1,61 @@
+name: CI/CD Pipeline
+
+on:
+    push:
+        tags:
+            - "*"
+
+jobs:
+    release:
+        name: Build and Release
+        runs-on: ubuntu-latest
+
+        steps:
+            - name: Checkout code
+              uses: actions/checkout@v3
+
+            - name: Set up QEMU
+              uses: docker/setup-qemu-action@v3
+
+            - name: Set up Docker Buildx
+              uses: docker/setup-buildx-action@v2
+
+            - name: Log in to Docker Hub
+              uses: docker/login-action@v2
+              with:
+                  username: ${{ secrets.DOCKER_HUB_USERNAME }}
+                  password: ${{ secrets.DOCKER_HUB_ACCESS_TOKEN }}
+
+            - name: Extract tag name
+              id: get-tag
+              run: echo "TAG=${GITHUB_REF#refs/tags/}" >> $GITHUB_ENV
+
+            - name: Install Go
+              uses: actions/setup-go@v4
+              with:
+                  go-version: 1.23.1
+
+            - name: Update version in main.go
+              run: |
+                  TAG=${{ env.TAG }}
+                  if [ -f main.go ]; then
+                    sed -i 's/Newt version replaceme/Newt version '"$TAG"'/' main.go
+                    echo "Updated main.go with version $TAG"
+                  else
+                    echo "main.go not found"
+                  fi
+
+            - name: Build and push Docker images
+              run: |
+                  TAG=${{ env.TAG }}
+                  make docker-build-release tag=$TAG
+
+            - name: Build binaries
+              run: |
+                  make go-build-release
+
+            - name: Upload artifacts from /bin
+              uses: actions/upload-artifact@v4
+              with:
+                  name: binaries
+                  path: bin/

.gitignore

@@ -1 +1,3 @@
-newt
+newt
+.DS_Store
+bin/

.go-version

@@ -0,0 +1 @@
+1.23.2

Dockerfile

@@ -15,19 +15,13 @@ COPY . .
 # Build the application
 RUN CGO_ENABLED=0 GOOS=linux go build -o /newt
 
-# Start a new stage from scratch
-FROM ubuntu:22.04 AS runner
+FROM alpine:3.19 AS runner
 
-RUN apt-get update && apt-get install ca-certificates -y  && rm -rf /var/lib/apt/lists/*
+RUN apk --no-cache add ca-certificates
 
-# Copy the pre-built binary file from the previous stage and the entrypoint script
 COPY --from=builder /newt /usr/local/bin/
 COPY entrypoint.sh /
 
 RUN chmod +x /entrypoint.sh
-
-# Copy the entrypoint script
 ENTRYPOINT ["/entrypoint.sh"]
-
-# Command to run the executable
 CMD ["newt"]

Makefile

@@ -1,6 +1,14 @@
 
 all: build push
 
+docker-build-release:
+	@if [ -z "$(tag)" ]; then \
+		echo "Error: tag is required. Usage: make build-all tag=<tag>"; \
+		exit 1; \
+	fi
+	docker buildx build --platform linux/arm/v7,linux/arm64,linux/amd64 -t fosrl/newt:latest -f Dockerfile --push .
+	docker buildx build --platform linux/arm/v7,linux/arm64,linux/amd64 -t fosrl/newt:$(tag) -f Dockerfile --push .
+
 build:
 	docker build -t fosrl/newt:latest .
 
@@ -11,7 +19,19 @@ test:
 	docker run fosrl/newt:latest
 
 local: 
-	 CGO_ENABLED=0 go build -o newt
+	CGO_ENABLED=0 go build -o newt
+
+go-build-release:
+	CGO_ENABLED=0 GOOS=linux GOARCH=arm64 go build -o bin/newt_linux_arm64
+	CGO_ENABLED=0 GOOS=linux GOARCH=arm GOARM=7 go build -o bin/newt_linux_arm32
+	CGO_ENABLED=0 GOOS=linux GOARCH=arm GOARM=6 go build -o bin/newt_linux_arm32v6
+	CGO_ENABLED=0 GOOS=linux GOARCH=amd64 go build -o bin/newt_linux_amd64
+	CGO_ENABLED=0 GOOS=linux GOARCH=riscv64 go build -o bin/newt_linux_riscv64
+	CGO_ENABLED=0 GOOS=darwin GOARCH=arm64 go build -o bin/newt_darwin_arm64
+	CGO_ENABLED=0 GOOS=darwin GOARCH=amd64 go build -o bin/newt_darwin_amd64
+	CGO_ENABLED=0 GOOS=windows GOARCH=amd64 go build -o bin/newt_windows_amd64.exe
+	CGO_ENABLED=0 GOOS=freebsd GOARCH=amd64 go build -o bin/newt_freebsd_amd64
+	CGO_ENABLED=0 GOOS=freebsd GOARCH=arm64 go build -o bin/newt_freebsd_arm64
 
 clean:
 	rm newt

README.md

@@ -19,7 +19,7 @@ _Sample output of a Newt container connected to Pangolin and hosting various res
 
 ### Registers with Pangolin
 
-Using the Newt ID and a secret the client will make HTTP requests to Pangolin to receive a session token. Using that token it will connect to a websocket and maintain that connection. Control messages will be sent over the websocket.
+Using the Newt ID and a secret, the client will make HTTP requests to Pangolin to receive a session token. Using that token, it will connect to a websocket and maintain that connection. Control messages will be sent over the websocket.
 
 ### Receives WireGuard Control Messages
 
@@ -27,7 +27,7 @@ When Newt receives WireGuard control messages, it will use the information encod
 
 ### Receives Proxy Control Messages
 
-When Newt receives WireGuard control messages, it will use the information encoded to crate local low level TCP and UDP proxies attached to the virtual tunnel in order to relay traffic to programmed targets.
+When Newt receives WireGuard control messages, it will use the information encoded to create a local low level TCP and UDP proxies attached to the virtual tunnel in order to relay traffic to programmed targets.
 
 ## CLI Args
 
@@ -36,7 +36,8 @@ When Newt receives WireGuard control messages, it will use the information encod
 - `secret`: A unique secret (not shared and kept private) used to authenticate the client ID with the websocket in order to receive commands. 
 - `dns`: DNS server to use to resolve the endpoint
 - `log-level` (optional): The log level to use. Default: INFO
-
+- `updown` (optional): A script to be called when targets are added or removed.
+ 
 Example:
 
 ```bash
@@ -74,6 +75,38 @@ services:
         - --endpoint https://example.com
 ```
 
+Finally a basic systemd service:
+
+```
+[Unit]
+Description=Newt VPN Client
+After=network.target
+
+[Service]
+ExecStart=/usr/local/bin/newt --id 31frd0uzbjvp721 --secret h51mmlknrvrwv8s4r1i210azhumt6isgbpyavxodibx1k2d6 --endpoint https://example.com
+Restart=always
+User=root
+
+[Install]
+WantedBy=multi-user.target
+```
+
+Make sure to `mv ./newt /usr/local/bin/newt`!
+
+### Updown
+
+You can pass in a updown script for Newt to call when it is adding or removing a target:
+
+`--updown "python3 test.py"`
+
+It will get called with args when a target is added: 
+`python3 test.py add tcp localhost:8556`
+`python3 test.py remove tcp localhost:8556`
+
+Returning a string from the script in the format of a target (`ip:dst` so `10.0.0.1:8080`) it will override the target and use this value instead to proxy.
+
+You can look at updown.py as a reference script to get started!
+
 ## Build
 
 ### Container 
@@ -98,4 +131,4 @@ Newt is dual licensed under the AGPLv3 and the Fossorial Commercial license. For
 
 ## Contributions
 
-Please see [CONTRIBUTIONS](./CONTRIBUTING.md) in the repository for guidelines and best practices.
+Please see [CONTRIBUTIONS](./CONTRIBUTING.md) in the repository for guidelines and best practices.

SECURITY.md

@@ -0,0 +1,14 @@
+# Security Policy
+
+If you discover a security vulnerability, please follow the steps below to responsibly disclose it to us:
+
+1. **Do not create a public GitHub issue or discussion post.** This could put the security of other users at risk.
+2. Send a detailed report to [security@fossorial.io](mailto:security@fossorial.io) or send a **private** message to a maintainer on [Discord](https://discord.gg/HCJR8Xhme4). Include:
+
+- Description and location of the vulnerability.
+- Potential impact of the vulnerability.
+- Steps to reproduce the vulnerability.
+- Potential solutions to fix the vulnerability.
+- Your name/handle and a link for recognition (optional).
+
+We aim to address the issue as soon as possible.

docker-compose.yml

@@ -6,4 +6,5 @@ services:
     environment:
       - PANGOLIN_ENDPOINT=https://example.com
       - NEWT_ID=2ix2t8xk22ubpfy 
-      - NEWT_SECRET=nnisrfsdfc7prqsp9ewo1dvtvci50j5uiqotez00dgap0ii2 
+      - NEWT_SECRET=nnisrfsdfc7prqsp9ewo1dvtvci50j5uiqotez00dgap0ii2 
+      - LOG_LEVEL=DEBUG

go.mod

@@ -4,16 +4,19 @@ go 1.23.1
 
 toolchain go1.23.2
 
-require golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
+require (
+	github.com/gorilla/websocket v1.5.3
+	golang.org/x/net v0.30.0
+	golang.zx2c4.com/wireguard v0.0.0-20231211153847-12269c276173
+	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6
+	gvisor.dev/gvisor v0.0.0-20230927004350-cbd86285d259
+)
 
 require (
 	github.com/google/btree v1.1.2 // indirect
-	github.com/gorilla/websocket v1.5.3 // indirect
+	github.com/google/go-cmp v0.6.0 // indirect
 	golang.org/x/crypto v0.28.0 // indirect
-	golang.org/x/net v0.30.0 // indirect
 	golang.org/x/sys v0.26.0 // indirect
 	golang.org/x/time v0.7.0 // indirect
 	golang.zx2c4.com/wintun v0.0.0-20230126152724-0fa3db229ce2 // indirect
-	golang.zx2c4.com/wireguard/wgctrl v0.0.0-20230429144221-925a1e7659e6 // indirect
-	gvisor.dev/gvisor v0.0.0-20230927004350-cbd86285d259 // indirect
 )

go.sum

@@ -1,5 +1,7 @@
 github.com/google/btree v1.1.2 h1:xf4v41cLI2Z6FxbKm+8Bu+m8ifhj15JuZ9sa0jZCMUU=
 github.com/google/btree v1.1.2/go.mod h1:qOPhT0dTNdNzV6Z/lhRX0YXUafgPLFUh+gZMl761Gm4=
+github.com/google/go-cmp v0.6.0 h1:ofyhxvXcZhMsU5ulbFiLKl/XBFqE1GSq7atu8tAmTRI=
+github.com/google/go-cmp v0.6.0/go.mod h1:17dUlkBOakJ0+DkrSSNjCkIjxS6bF9zb3elmeNGIjoY=
 github.com/gorilla/websocket v1.5.3 h1:saDtZ6Pbx/0u+bgYQ3q96pZgCzfhKXGPqt7kZ72aNNg=
 github.com/gorilla/websocket v1.5.3/go.mod h1:YR8l580nyteQvAITg2hZ9XVh4b55+EU/adAjf1fMHhE=
 golang.org/x/crypto v0.28.0 h1:GBDwsMXVQi34v5CCYUm2jkJvu4cbtru2U4TN2PSyQnw=

main.go

@@ -11,7 +11,9 @@ import (
 	"net"
 	"net/netip"
 	"os"
+	"os/exec"
 	"os/signal"
+	"strconv"
 	"strings"
 	"syscall"
 	"time"
@@ -123,6 +125,7 @@ func startPingCheck(tnet *netstack.Net, serverIP string, stopChan chan struct{})
 				err := ping(tnet, serverIP)
 				if err != nil {
 					logger.Warn("Periodic ping failed: %v", err)
+					logger.Warn("HINT: Do you have UDP port 51820 (or the port in config.yml) open on your Pangolin server?")
 				}
 			case <-stopChan:
 				logger.Info("Stopping ping check")
@@ -134,7 +137,7 @@ func startPingCheck(tnet *netstack.Net, serverIP string, stopChan chan struct{})
 
 func pingWithRetry(tnet *netstack.Net, dst string) error {
 	const (
-		maxAttempts = 5
+		maxAttempts = 15
 		retryDelay  = 2 * time.Second
 	)
 
@@ -242,23 +245,28 @@ func resolveDomain(domain string) (string, error) {
 	return ipAddr, nil
 }
 
-func main() {
-	var (
-		endpoint   string
-		id         string
-		secret     string
-		dns        string
-		privateKey wgtypes.Key
-		err        error
-		logLevel   string
-	)
+var (
+	endpoint     string
+	id           string
+	secret       string
+	mtu          string
+	mtuInt       int
+	dns          string
+	privateKey   wgtypes.Key
+	err          error
+	logLevel     string
+	updownScript string
+)
 
+func main() {
 	// if PANGOLIN_ENDPOINT, NEWT_ID, and NEWT_SECRET are set as environment variables, they will be used as default values
 	endpoint = os.Getenv("PANGOLIN_ENDPOINT")
 	id = os.Getenv("NEWT_ID")
 	secret = os.Getenv("NEWT_SECRET")
+	mtu = os.Getenv("MTU")
 	dns = os.Getenv("DNS")
 	logLevel = os.Getenv("LOG_LEVEL")
+	updownScript = os.Getenv("UPDOWN_SCRIPT")
 
 	if endpoint == "" {
 		flag.StringVar(&endpoint, "endpoint", "", "Endpoint of your pangolin server")
@@ -269,21 +277,37 @@ func main() {
 	if secret == "" {
 		flag.StringVar(&secret, "secret", "", "Newt secret")
 	}
+	if mtu == "" {
+		flag.StringVar(&mtu, "mtu", "1280", "MTU to use")
+	}
 	if dns == "" {
 		flag.StringVar(&dns, "dns", "8.8.8.8", "DNS server to use")
 	}
 	if logLevel == "" {
 		flag.StringVar(&logLevel, "log-level", "INFO", "Log level (DEBUG, INFO, WARN, ERROR, FATAL)")
 	}
+	if updownScript == "" {
+		flag.StringVar(&updownScript, "updown", "", "Path to updown script to be called when targets are added or removed")
+	}
+
+	// do a --version check
+	version := flag.Bool("version", false, "Print the version")
+
 	flag.Parse()
 
+	if *version {
+		fmt.Println("Newt version replaceme")
+		os.Exit(0)
+	}
+
 	logger.Init()
 	loggerLevel := parseLogLevel(logLevel)
 	logger.GetLogger().SetLevel(parseLogLevel(logLevel))
 
-	// Validate required fields
-	if endpoint == "" || id == "" || secret == "" {
-		logger.Fatal("endpoint, id, and secret are required either via CLI flags or environment variables")
+	// parse the mtu string into an int
+	mtuInt, err = strconv.Atoi(mtu)
+	if err != nil {
+		logger.Fatal("Failed to parse MTU: %v", err)
 	}
 
 	privateKey, err = wgtypes.GeneratePrivateKey()
@@ -333,7 +357,8 @@ func main() {
 			err = pingWithRetry(tnet, wgData.ServerIP)
 			if err != nil {
 				// Handle complete failure after all retries
-				logger.Error("Failed to ping %s: %v", wgData.ServerIP, err)
+				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?")
 			}
 			return
 		}
@@ -353,7 +378,7 @@ func main() {
 		tun, tnet, err = netstack.CreateNetTUN(
 			[]netip.Addr{netip.MustParseAddr(wgData.TunnelIP)},
 			[]netip.Addr{netip.MustParseAddr(dns)},
-			1420)
+			mtuInt)
 		if err != nil {
 			logger.Error("Failed to create TUN device: %v", err)
 		}
@@ -440,11 +465,6 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey(
 		if len(targetData.Targets) > 0 {
 			updateTargets(pm, "add", wgData.TunnelIP, "tcp", targetData)
 		}
-
-		err = pm.Start()
-		if err != nil {
-			logger.Error("Failed to start proxy manager: %v", err)
-		}
 	})
 
 	client.RegisterHandler("newt/udp/add", func(msg websocket.WSMessage) {
@@ -465,11 +485,6 @@ persistent_keepalive_interval=5`, fixKey(fmt.Sprintf("%s", privateKey)), fixKey(
 		if len(targetData.Targets) > 0 {
 			updateTargets(pm, "add", wgData.TunnelIP, "udp", targetData)
 		}
-
-		err = pm.Start()
-		if err != nil {
-			logger.Error("Failed to start proxy manager: %v", err)
-		}
 	})
 
 	client.RegisterHandler("newt/udp/remove", func(msg websocket.WSMessage) {
@@ -577,6 +592,18 @@ func updateTargets(pm *proxy.ProxyManager, action string, tunnelIP string, proto
 
 		if action == "add" {
 			target := parts[1] + ":" + parts[2]
+
+			// Call updown script if provided
+			processedTarget := target
+			if updownScript != "" {
+				newTarget, err := executeUpdownScript(action, proto, target)
+				if err != nil {
+					logger.Warn("Updown script error: %v", err)
+				} else if newTarget != "" {
+					processedTarget = newTarget
+				}
+			}
+
 			// Only remove the specific target if it exists
 			err := pm.RemoveTarget(proto, tunnelIP, port)
 			if err != nil {
@@ -587,10 +614,21 @@ func updateTargets(pm *proxy.ProxyManager, action string, tunnelIP string, proto
 			}
 
 			// Add the new target
-			pm.AddTarget(proto, tunnelIP, port, target)
+			pm.AddTarget(proto, tunnelIP, port, processedTarget)
 
 		} else if action == "remove" {
 			logger.Info("Removing target with port %d", port)
+
+			target := parts[1] + ":" + parts[2]
+
+			// Call updown script if provided
+			if updownScript != "" {
+				_, err := executeUpdownScript(action, proto, target)
+				if err != nil {
+					logger.Warn("Updown script error: %v", err)
+				}
+			}
+
 			err := pm.RemoveTarget(proto, tunnelIP, port)
 			if err != nil {
 				logger.Error("Failed to remove target: %v", err)
@@ -601,3 +639,45 @@ func updateTargets(pm *proxy.ProxyManager, action string, tunnelIP string, proto
 
 	return nil
 }
+
+func executeUpdownScript(action, proto, target string) (string, error) {
+	if updownScript == "" {
+		return target, nil
+	}
+
+	// Split the updownScript in case it contains spaces (like "/usr/bin/python3 script.py")
+	parts := strings.Fields(updownScript)
+	if len(parts) == 0 {
+		return target, fmt.Errorf("invalid updown script command")
+	}
+
+	var cmd *exec.Cmd
+	if len(parts) == 1 {
+		// If it's a single executable
+		logger.Info("Executing updown script: %s %s %s %s", updownScript, action, proto, target)
+		cmd = exec.Command(parts[0], action, proto, target)
+	} else {
+		// If it includes interpreter and script
+		args := append(parts[1:], action, proto, target)
+		logger.Info("Executing updown script: %s %s %s %s %s", parts[0], strings.Join(parts[1:], " "), action, proto, target)
+		cmd = exec.Command(parts[0], args...)
+	}
+
+	output, err := cmd.Output()
+	if err != nil {
+		if exitErr, ok := err.(*exec.ExitError); ok {
+			return "", fmt.Errorf("updown script execution failed (exit code %d): %s",
+				exitErr.ExitCode(), string(exitErr.Stderr))
+		}
+		return "", fmt.Errorf("updown script execution failed: %v", err)
+	}
+
+	// If the script returns a new target, use it
+	newTarget := strings.TrimSpace(string(output))
+	if newTarget != "" {
+		logger.Info("Updown script returned new target: %s", newTarget)
+		return newTarget, nil
+	}
+
+	return target, nil
+}

proxy/manager.go

@@ -9,326 +9,344 @@ import (
 	"time"
 
 	"github.com/fosrl/newt/logger"
-
 	"golang.zx2c4.com/wireguard/tun/netstack"
+	"gvisor.dev/gvisor/pkg/tcpip/adapters/gonet"
 )
 
+// Target represents a proxy target with its address and port
+type Target struct {
+	Address string
+	Port    int
+}
+
+// ProxyManager handles the creation and management of proxy connections
+type ProxyManager struct {
+	tnet       *netstack.Net
+	tcpTargets map[string]map[int]string // map[listenIP]map[port]targetAddress
+	udpTargets map[string]map[int]string
+	listeners  []*gonet.TCPListener
+	udpConns   []*gonet.UDPConn
+	running    bool
+	mutex      sync.RWMutex
+}
+
+// NewProxyManager creates a new proxy manager instance
 func NewProxyManager(tnet *netstack.Net) *ProxyManager {
 	return &ProxyManager{
-		tnet: tnet,
+		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),
 	}
 }
 
-func (pm *ProxyManager) AddTarget(protocol, listen string, port int, target string) {
-	pm.Lock()
-	defer pm.Unlock()
+// AddTarget adds as new target for proxying
+func (pm *ProxyManager) AddTarget(proto, listenIP string, port int, targetAddr string) error {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
 
-	logger.Info("Adding target: %s://%s:%d -> %s", protocol, listen, port, target)
-
-	newTarget := ProxyTarget{
-		Protocol: protocol,
-		Listen:   listen,
-		Port:     port,
-		Target:   target,
-		cancel:   make(chan struct{}),
-		done:     make(chan struct{}),
+	switch proto {
+	case "tcp":
+		if pm.tcpTargets[listenIP] == nil {
+			pm.tcpTargets[listenIP] = make(map[int]string)
+		}
+		pm.tcpTargets[listenIP][port] = targetAddr
+	case "udp":
+		if pm.udpTargets[listenIP] == nil {
+			pm.udpTargets[listenIP] = make(map[int]string)
+		}
+		pm.udpTargets[listenIP][port] = targetAddr
+	default:
+		return fmt.Errorf("unsupported protocol: %s", proto)
 	}
 
-	pm.targets = append(pm.targets, newTarget)
+	if pm.running {
+		return pm.startTarget(proto, listenIP, port, targetAddr)
+	} else {
+		logger.Debug("Not adding target because not running")
+	}
+	return nil
 }
 
-func (pm *ProxyManager) RemoveTarget(protocol, listen string, port int) error {
-	pm.Lock()
-	defer pm.Unlock()
+func (pm *ProxyManager) RemoveTarget(proto, listenIP string, port int) error {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
 
-	protocol = strings.ToLower(protocol)
-	if protocol != "tcp" && protocol != "udp" {
-		return fmt.Errorf("unsupported protocol: %s", protocol)
-	}
-
-	for i, target := range pm.targets {
-		if target.Listen == listen &&
-			target.Port == port &&
-			strings.ToLower(target.Protocol) == protocol {
-
-			// Signal the serving goroutine to stop
-			select {
-			case <-target.cancel:
-				// Channel is already closed, no need to close it again
-			default:
-				close(target.cancel)
-			}
-
-			// Close the appropriate listener/connection based on protocol
-			target.Lock()
-			switch protocol {
-			case "tcp":
-				if target.listener != nil {
-					select {
-					case <-target.cancel:
-						// Listener was already closed by Stop()
-					default:
-						target.listener.Close()
-					}
-				}
-			case "udp":
-				if target.udpConn != nil {
-					select {
-					case <-target.cancel:
-						// Connection was already closed by Stop()
-					default:
-						target.udpConn.Close()
-					}
+	switch proto {
+	case "tcp":
+		if targets, ok := pm.tcpTargets[listenIP]; ok {
+			delete(targets, port)
+			// Remove and close the corresponding TCP listener
+			for i, listener := range pm.listeners {
+				if addr, ok := listener.Addr().(*net.TCPAddr); ok && addr.Port == port {
+					listener.Close()
+					time.Sleep(50 * time.Millisecond)
+					// Remove from slice
+					pm.listeners = append(pm.listeners[:i], pm.listeners[i+1:]...)
+					break
 				}
 			}
-			target.Unlock()
-
-			// Wait for the target to fully stop
-			<-target.done
-
-			// Remove the target from the slice
-			pm.targets = append(pm.targets[:i], pm.targets[i+1:]...)
-			return nil
-		}
-	}
-
-	return fmt.Errorf("target not found for %s %s:%d", protocol, listen, port)
-}
-
-func (pm *ProxyManager) Start() error {
-	pm.RLock()
-	defer pm.RUnlock()
-
-	for i := range pm.targets {
-		target := &pm.targets[i]
-
-		target.Lock()
-		// If target is already running, skip it
-		if target.listener != nil || target.udpConn != nil {
-			target.Unlock()
-			continue
-		}
-
-		// Mark the target as starting by creating a nil listener/connection
-		// This prevents other goroutines from trying to start it
-		if strings.ToLower(target.Protocol) == "tcp" {
-			target.listener = nil
 		} else {
-			target.udpConn = nil
+			return fmt.Errorf("target not found: %s:%d", listenIP, port)
 		}
-		target.Unlock()
+	case "udp":
+		if targets, ok := pm.udpTargets[listenIP]; ok {
+			delete(targets, port)
+			// Remove and close the corresponding UDP connection
+			for i, conn := range pm.udpConns {
+				if addr, ok := conn.LocalAddr().(*net.UDPAddr); ok && addr.Port == port {
+					conn.Close()
+					time.Sleep(50 * time.Millisecond)
+					// Remove from slice
+					pm.udpConns = append(pm.udpConns[:i], pm.udpConns[i+1:]...)
+					break
+				}
+			}
+		} else {
+			return fmt.Errorf("target not found: %s:%d", listenIP, port)
+		}
+	default:
+		return fmt.Errorf("unsupported protocol: %s", proto)
+	}
+	return nil
+}
 
-		switch strings.ToLower(target.Protocol) {
-		case "tcp":
-			go pm.serveTCP(target)
-		case "udp":
-			go pm.serveUDP(target)
-		default:
-			return fmt.Errorf("unsupported protocol: %s", target.Protocol)
+// Start begins listening for all configured proxy targets
+func (pm *ProxyManager) Start() error {
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
+
+	if pm.running {
+		return nil
+	}
+
+	// Start TCP targets
+	for listenIP, targets := range pm.tcpTargets {
+		for port, targetAddr := range targets {
+			if err := pm.startTarget("tcp", listenIP, port, targetAddr); err != nil {
+				return fmt.Errorf("failed to start TCP target: %v", err)
+			}
 		}
 	}
+
+	// Start UDP targets
+	for listenIP, targets := range pm.udpTargets {
+		for port, targetAddr := range targets {
+			if err := pm.startTarget("udp", listenIP, port, targetAddr); err != nil {
+				return fmt.Errorf("failed to start UDP target: %v", err)
+			}
+		}
+	}
+
+	pm.running = true
 	return nil
 }
 
 func (pm *ProxyManager) Stop() error {
-	pm.Lock()
-	defer pm.Unlock()
+	pm.mutex.Lock()
+	defer pm.mutex.Unlock()
 
-	var wg sync.WaitGroup
-	for i := range pm.targets {
-		target := &pm.targets[i]
-		wg.Add(1)
-		go func(t *ProxyTarget) {
-			defer wg.Done()
-			close(t.cancel)
-			t.Lock()
-			if t.listener != nil {
-				t.listener.Close()
-			}
-			if t.udpConn != nil {
-				t.udpConn.Close()
-			}
-			t.Unlock()
-			// Wait for the target to fully stop
-			<-t.done
-		}(target)
+	if !pm.running {
+		return nil
 	}
-	wg.Wait()
+
+	// Set running to false first to signal handlers to stop
+	pm.running = false
+
+	// Close TCP listeners
+	for i := len(pm.listeners) - 1; i >= 0; i-- {
+		listener := pm.listeners[i]
+		if err := listener.Close(); err != nil {
+			logger.Error("Error closing TCP listener: %v", err)
+		}
+		// Remove from slice
+		pm.listeners = append(pm.listeners[:i], pm.listeners[i+1:]...)
+	}
+
+	// Close UDP connections
+	for i := len(pm.udpConns) - 1; i >= 0; i-- {
+		conn := pm.udpConns[i]
+		if err := conn.Close(); err != nil {
+			logger.Error("Error closing UDP connection: %v", err)
+		}
+		// Remove from slice
+		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)
+
 	return nil
 }
 
-func (pm *ProxyManager) serveTCP(target *ProxyTarget) {
-	defer close(target.done) // Signal that this target is fully stopped
+func (pm *ProxyManager) startTarget(proto, listenIP string, port int, targetAddr string) error {
+	switch proto {
+	case "tcp":
+		listener, err := pm.tnet.ListenTCP(&net.TCPAddr{Port: port})
+		if err != nil {
+			return fmt.Errorf("failed to create TCP listener: %v", err)
+		}
 
-	listener, err := pm.tnet.ListenTCP(&net.TCPAddr{
-		IP:   net.ParseIP(target.Listen),
-		Port: target.Port,
-	})
-	if err != nil {
-		logger.Info("Failed to start TCP listener for %s:%d: %v", target.Listen, target.Port, err)
-		return
+		pm.listeners = append(pm.listeners, listener)
+		go pm.handleTCPProxy(listener, targetAddr)
+
+	case "udp":
+		addr := &net.UDPAddr{Port: port}
+		conn, err := pm.tnet.ListenUDP(addr)
+		if err != nil {
+			return fmt.Errorf("failed to create UDP listener: %v", err)
+		}
+
+		pm.udpConns = append(pm.udpConns, conn)
+		go pm.handleUDPProxy(conn, targetAddr)
+
+	default:
+		return fmt.Errorf("unsupported protocol: %s", proto)
 	}
 
-	target.Lock()
-	target.listener = listener
-	target.Unlock()
+	logger.Info("Started %s proxy from %s:%d to %s", proto, listenIP, port, targetAddr)
 
-	defer listener.Close()
-	logger.Info("TCP proxy listening on %s", listener.Addr())
-
-	var activeConns sync.WaitGroup
-	acceptDone := make(chan struct{})
-
-	// Goroutine to handle shutdown signal
-	go func() {
-		<-target.cancel
-		close(acceptDone)
-		listener.Close()
-	}()
+	return nil
+}
 
+func (pm *ProxyManager) handleTCPProxy(listener net.Listener, targetAddr string) {
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
-			select {
-			case <-target.cancel:
-				// Wait for active connections to finish
-				activeConns.Wait()
+			// Check if we're shutting down or the listener was closed
+			if !pm.running {
 				return
-			default:
-				logger.Info("Failed to accept TCP connection: %v", err)
-				// Don't return here, try to accept new connections
-				time.Sleep(time.Second)
-				continue
 			}
+
+			// Check for specific network errors that indicate the listener is closed
+			if ne, ok := err.(net.Error); ok && !ne.Temporary() {
+				logger.Info("TCP listener closed, stopping proxy handler for %v", listener.Addr())
+				return
+			}
+
+			logger.Error("Error accepting TCP connection: %v", err)
+			// Don't hammer the CPU if we hit a temporary error
+			time.Sleep(100 * time.Millisecond)
+			continue
 		}
 
-		activeConns.Add(1)
 		go func() {
-			defer activeConns.Done()
-			pm.handleTCPConnection(conn, target.Target, acceptDone)
+			target, err := net.Dial("tcp", targetAddr)
+			if err != nil {
+				logger.Error("Error connecting to target: %v", err)
+				conn.Close()
+				return
+			}
+
+			// Create a WaitGroup to ensure both copy operations complete
+			var wg sync.WaitGroup
+			wg.Add(2)
+
+			go func() {
+				defer wg.Done()
+				io.Copy(target, conn)
+				target.Close()
+			}()
+
+			go func() {
+				defer wg.Done()
+				io.Copy(conn, target)
+				conn.Close()
+			}()
+
+			// Wait for both copies to complete
+			wg.Wait()
 		}()
 	}
 }
 
-func (pm *ProxyManager) handleTCPConnection(clientConn net.Conn, target string, done chan struct{}) {
-	defer clientConn.Close()
-
-	serverConn, err := net.Dial("tcp", target)
-	if err != nil {
-		logger.Info("Failed to connect to target %s: %v", target, err)
-		return
-	}
-	defer serverConn.Close()
-
-	var wg sync.WaitGroup
-	wg.Add(2)
-
-	// Client -> Server
-	go func() {
-		defer wg.Done()
-		select {
-		case <-done:
-			return
-		default:
-			io.Copy(serverConn, clientConn)
-		}
-	}()
-
-	// Server -> Client
-	go func() {
-		defer wg.Done()
-		select {
-		case <-done:
-			return
-		default:
-			io.Copy(clientConn, serverConn)
-		}
-	}()
-
-	wg.Wait()
-}
-
-func (pm *ProxyManager) serveUDP(target *ProxyTarget) {
-	defer close(target.done) // Signal that this target is fully stopped
-
-	addr := &net.UDPAddr{
-		IP:   net.ParseIP(target.Listen),
-		Port: target.Port,
-	}
-
-	conn, err := pm.tnet.ListenUDP(addr)
-	if err != nil {
-		logger.Info("Failed to start UDP listener for %s:%d: %v", target.Listen, target.Port, err)
-		return
-	}
-
-	target.Lock()
-	target.udpConn = conn
-	target.Unlock()
-
-	defer conn.Close()
-	logger.Info("UDP proxy listening on %s", conn.LocalAddr())
-
-	buffer := make([]byte, 65535)
-	var activeConns sync.WaitGroup
+func (pm *ProxyManager) handleUDPProxy(conn *gonet.UDPConn, targetAddr string) {
+	buffer := make([]byte, 65507) // Max UDP packet size
+	clientConns := make(map[string]*net.UDPConn)
+	var clientsMutex sync.RWMutex
 
 	for {
-		select {
-		case <-target.cancel:
-			activeConns.Wait() // Wait for all active UDP handlers to complete
-			return
-		default:
-			n, remoteAddr, err := conn.ReadFrom(buffer)
-			if err != nil {
-				select {
-				case <-target.cancel:
-					activeConns.Wait()
-					return
-				default:
-					logger.Info("Failed to read UDP packet: %v", err)
-					continue
-				}
+		n, remoteAddr, err := conn.ReadFrom(buffer)
+		if err != nil {
+			if !pm.running {
+				return
 			}
 
-			targetAddr, err := net.ResolveUDPAddr("udp", target.Target)
+			// Check for connection closed conditions
+			if err == io.EOF || strings.Contains(err.Error(), "use of closed network connection") {
+				logger.Info("UDP connection closed, stopping proxy handler")
+
+				// Clean up existing client connections
+				clientsMutex.Lock()
+				for _, targetConn := range clientConns {
+					targetConn.Close()
+				}
+				clientConns = nil
+				clientsMutex.Unlock()
+
+				return
+			}
+
+			logger.Error("Error reading UDP packet: %v", err)
+			continue
+		}
+
+		clientKey := remoteAddr.String()
+		clientsMutex.RLock()
+		targetConn, exists := clientConns[clientKey]
+		clientsMutex.RUnlock()
+
+		if !exists {
+			targetUDPAddr, err := net.ResolveUDPAddr("udp", targetAddr)
 			if err != nil {
-				logger.Info("Failed to resolve target address %s: %v", target.Target, err)
+				logger.Error("Error resolving target address: %v", err)
 				continue
 			}
 
-			activeConns.Add(1)
-			go func(data []byte, remote net.Addr) {
-				defer activeConns.Done()
-				targetConn, err := net.DialUDP("udp", nil, targetAddr)
-				if err != nil {
-					logger.Info("Failed to connect to target %s: %v", target.Target, err)
-					return
-				}
-				defer targetConn.Close()
+			targetConn, err = net.DialUDP("udp", nil, targetUDPAddr)
+			if err != nil {
+				logger.Error("Error connecting to target: %v", err)
+				continue
+			}
 
-				select {
-				case <-target.cancel:
-					return
-				default:
-					_, err = targetConn.Write(data)
+			clientsMutex.Lock()
+			clientConns[clientKey] = targetConn
+			clientsMutex.Unlock()
+
+			go func() {
+				buffer := make([]byte, 65507)
+				for {
+					n, _, err := targetConn.ReadFromUDP(buffer)
 					if err != nil {
-						logger.Info("Failed to write to target: %v", err)
+						logger.Error("Error reading from target: %v", err)
 						return
 					}
 
-					response := make([]byte, 65535)
-					n, err := targetConn.Read(response)
+					_, err = conn.WriteTo(buffer[:n], remoteAddr)
 					if err != nil {
-						logger.Info("Failed to read response from target: %v", err)
+						logger.Error("Error writing to client: %v", err)
 						return
 					}
-
-					_, err = conn.WriteTo(response[:n], remote)
-					if err != nil {
-						logger.Info("Failed to write response to client: %v", err)
-					}
 				}
-			}(buffer[:n], remoteAddr)
+			}()
+		}
+
+		_, err = targetConn.Write(buffer[:n])
+		if err != nil {
+			logger.Error("Error writing to target: %v", err)
+			targetConn.Close()
+			clientsMutex.Lock()
+			delete(clientConns, clientKey)
+			clientsMutex.Unlock()
 		}
 	}
 }

proxy/types.go

@@ -1,28 +0,0 @@
-package proxy
-
-import (
-	"log"
-	"net"
-	"sync"
-
-	"golang.zx2c4.com/wireguard/tun/netstack"
-)
-
-type ProxyTarget struct {
-	Protocol   string
-	Listen     string
-	Port       int
-	Target     string
-	cancel     chan struct{}  // Channel to signal shutdown
-	done       chan struct{}  // Channel to signal completion
-	listener   net.Listener   // For TCP
-	udpConn    net.PacketConn // For UDP
-	sync.Mutex                // Protect access to connection
-}
-
-type ProxyManager struct {
-	targets      []ProxyTarget
-	tnet         *netstack.Net
-	log          *log.Logger
-	sync.RWMutex // Protect access to targets slice
-}

updown.py

@@ -0,0 +1,77 @@
+"""
+Sample updown script for Newt proxy
+Usage: update.py <action> <protocol> <target>
+
+Parameters:
+- action: 'add' or 'remove'
+- protocol: 'tcp' or 'udp'
+- target: the target address in format 'host:port'
+
+If the action is 'add', the script can return a modified target that
+will be used instead of the original.
+"""
+
+import sys
+import logging
+import json
+from datetime import datetime
+
+# Configure logging
+LOG_FILE = "/tmp/newt-updown.log"
+logging.basicConfig(
+    filename=LOG_FILE,
+    level=logging.INFO,
+    format='%(asctime)s - %(levelname)s - %(message)s'
+)
+
+def log_event(action, protocol, target):
+    """Log each event to a file for auditing purposes"""
+    timestamp = datetime.now().isoformat()
+    event = {
+        "timestamp": timestamp,
+        "action": action,
+        "protocol": protocol,
+        "target": target
+    }
+    logging.info(json.dumps(event))
+
+def handle_add(protocol, target):
+    """Handle 'add' action"""
+    logging.info(f"Adding {protocol} target: {target}")
+    
+def handle_remove(protocol, target):
+    """Handle 'remove' action"""
+    logging.info(f"Removing {protocol} target: {target}")
+    # For remove action, no return value is expected or used
+    
+def main():
+    # Check arguments
+    if len(sys.argv) != 4:
+        logging.error(f"Invalid arguments: {sys.argv}")
+        sys.exit(1)
+    
+    action = sys.argv[1]
+    protocol = sys.argv[2]
+    target = sys.argv[3]
+    
+    # Log the event
+    log_event(action, protocol, target)
+    
+    # Handle the action
+    if action == "add":
+        new_target = handle_add(protocol, target)
+        # Print the new target to stdout (if empty, no change will be made)
+        if new_target and new_target != target:
+            print(new_target)
+    elif action == "remove":
+        handle_remove(protocol, target)
+    else:
+        logging.error(f"Unknown action: {action}")
+        sys.exit(1)
+
+if __name__ == "__main__":
+    try:
+        main()
+    except Exception as e:
+        logging.error(f"Unhandled exception: {e}")
+        sys.exit(1)

websocket/client.go

@@ -228,6 +228,10 @@ func (c *Client) getToken() (string, error) {
 
 	var tokenResp TokenResponse
 	if err := json.NewDecoder(resp.Body).Decode(&tokenResp); err != nil {
+		// print out the token response for debugging
+		buf := new(bytes.Buffer)
+		buf.ReadFrom(resp.Body)
+		logger.Info("Token response: %s", buf.String())
 		return "", fmt.Errorf("failed to decode token response: %w", err)
 	}
 
@@ -305,6 +309,10 @@ func (c *Client) establishConnection() error {
 	go c.readPump()
 
 	if c.onConnect != nil {
+		err := c.saveConfig()
+		if err != nil {
+			logger.Error("Failed to save config: %v", err)
+		}
 		if err := c.onConnect(); err != nil {
 			logger.Error("OnConnect callback failed: %v", err)
 		}