From 31b88638e03c2abbbf99c27b0c6430b215706251 Mon Sep 17 00:00:00 2001
From: Justin <justin@layerv.ai>
Date: Fri, 10 Jul 2026 00:02:48 -0600
Subject: [PATCH] fix(ac): recover UDP handler panics

---
 endpoints/ac/msghandler.go             |  5 ++-
 endpoints/ac/udp_handler_panic_test.go | 57 ++++++++++++++++++++++++++
 endpoints/ac/udpac.go                  | 26 +++++++++++-
 3 files changed, 84 insertions(+), 4 deletions(-)
 create mode 100644 endpoints/ac/udp_handler_panic_test.go

diff --git a/endpoints/ac/msghandler.go b/endpoints/ac/msghandler.go
index b9e67d543..710567da4 100644
--- a/endpoints/ac/msghandler.go
+++ b/endpoints/ac/msghandler.go
@@ -23,8 +23,6 @@ const (
 )
 
 func (a *UdpAC) HandleUdpACOperations(ppd *core.PacketParserData) (err error) {
-	defer a.wg.Done()
-
 	acId := a.config.ACId
 	dopMsg := &common.ServerACOpsMsg{}
 	artMsg := &common.ACOpsResultMsg{}
@@ -611,6 +609,7 @@ func (a *UdpAC) HandleAccessControl(au *common.AgentUser, srcAddrs []*common.Net
 
 func (a *UdpAC) tcpTempAccessHandler(listener *net.TCPListener, timeoutSec int, dstAddrs []*common.NetAddress, openTimeSec int) {
 	defer a.wg.Done()
+	defer a.recoverUDPHandler(core.NHP_ACC)
 	defer listener.Close()
 
 	// accept only the first incoming tcp connection
@@ -745,6 +744,7 @@ func (a *UdpAC) tcpTempAccessHandler(listener *net.TCPListener, timeoutSec int,
 
 func (a *UdpAC) udpTempAccessHandler(conn *net.UDPConn, timeoutSec int, dstAddrs []*common.NetAddress, openTimeSec int) {
 	defer a.wg.Done()
+	defer a.recoverUDPHandler(core.NHP_ACC)
 	defer conn.Close()
 	// listen to accept and handle only one incoming connection
 	startTime := time.Now()
@@ -908,6 +908,7 @@ func (a *UdpAC) udpTempAccessHandler(conn *net.UDPConn, timeoutSec int, dstAddrs
 }
 
 func (a *UdpAC) tempConnTerminator(conn net.Conn, ctx context.Context) {
+	defer a.recoverUDPHandler(core.NHP_ACC)
 	select {
 	case <-a.signals.stop:
 		conn.Close()
diff --git a/endpoints/ac/udp_handler_panic_test.go b/endpoints/ac/udp_handler_panic_test.go
new file mode 100644
index 000000000..251a9687a
--- /dev/null
+++ b/endpoints/ac/udp_handler_panic_test.go
@@ -0,0 +1,57 @@
+package ac
+
+import (
+	"sync/atomic"
+	"testing"
+	"time"
+
+	"github.com/OpenNHP/opennhp/nhp/core"
+)
+
+func TestRunUDPHandlerRecoversPanicAndReleasesWaitGroup(t *testing.T) {
+	a := &UdpAC{config: &Config{ACId: "test-ac"}}
+	a.wg.Add(1)
+
+	done := make(chan struct{})
+	go func() {
+		a.runUDPHandler(core.NHP_AOP, func() { panic("malformed packet") })
+		close(done)
+	}()
+
+	select {
+	case <-done:
+	case <-time.After(time.Second):
+		t.Fatal("panicking UDP handler did not return")
+	}
+
+	waitDone := make(chan struct{})
+	go func() {
+		a.wg.Wait()
+		close(waitDone)
+	}()
+	select {
+	case <-waitDone:
+	case <-time.After(time.Second):
+		t.Fatal("panicking UDP handler did not release UdpAC wait group")
+	}
+}
+
+func TestRunUDPHandlerHappyPath(t *testing.T) {
+	a := &UdpAC{}
+	a.wg.Add(1)
+	var called atomic.Bool
+
+	a.runUDPHandler(core.NHP_AOP, func() { called.Store(true) })
+	a.wg.Wait()
+
+	if !called.Load() {
+		t.Fatal("handler was not called")
+	}
+}
+
+func TestRecoverUDPHandlerNilConfigIsSafe(t *testing.T) {
+	a := &UdpAC{}
+	a.wg.Add(1)
+	a.runUDPHandler(core.NHP_AOP, func() { panic("nil config") })
+	a.wg.Wait()
+}
diff --git a/endpoints/ac/udpac.go b/endpoints/ac/udpac.go
index 9b8e15807..ebe8bfbe3 100644
--- a/endpoints/ac/udpac.go
+++ b/endpoints/ac/udpac.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net"
 	"path/filepath"
+	"runtime/debug"
 	"sync"
 	"sync/atomic"
 	"time"
@@ -523,15 +524,36 @@ func (a *UdpAC) recvMessageRoutine() {
 			switch ppd.HeaderType {
 			case core.NHP_AOP:
 				// deal with NHP_AOP message
+				p := ppd
 				a.wg.Add(1)
-				go func(p *core.PacketParserData) {
+				go a.runUDPHandler(p.HeaderType, func() {
 					_ = a.HandleUdpACOperations(p)
-				}(ppd)
+				})
 			}
 		}
 	}
 }
 
+// runUDPHandler contains panics from input-driven handler goroutines. A
+// malformed packet or an unexpected nil in one handler must drop that request,
+// not terminate the entire access-controller process.
+func (a *UdpAC) runUDPHandler(headerType int, handler func()) {
+	defer a.wg.Done()
+	defer a.recoverUDPHandler(headerType)
+	handler()
+}
+
+func (a *UdpAC) recoverUDPHandler(headerType int) {
+	if recovered := recover(); recovered != nil {
+		acID := "unknown"
+		if a != nil && a.config != nil && a.config.ACId != "" {
+			acID = a.config.ACId
+		}
+		log.Error("ac(%s)[%s] UDP handler panic recovered: %v\n%s",
+			acID, core.HeaderTypeToString(headerType), recovered, debug.Stack())
+	}
+}
+
 // keep interaction between ac and server in certain time interval to keep outwards ip path active
 func (a *UdpAC) maintainServerConnectionRoutine() {
 	defer a.wg.Done()