From ee2bca04dfdb879b379a67047d79af49dfc45e07 Mon Sep 17 00:00:00 2001
From: Jonathan Barrow <jonbarrow1998@gmail.com>
Date: Mon, 15 Jan 2024 17:55:35 -0500
Subject: [PATCH] prudp: add events back to PRUDPServer

---
 prudp_endpoint.go   |  3 +++
 prudp_server.go     | 38 +++++++++++++++++++++++++++++++-------
 server_interface.go |  1 +
 3 files changed, 35 insertions(+), 7 deletions(-)

diff --git a/prudp_endpoint.go b/prudp_endpoint.go
index 1502e41..e1e4ff0 100644
--- a/prudp_endpoint.go
+++ b/prudp_endpoint.go
@@ -58,6 +58,9 @@ func (pep *PRUDPEndPoint) emit(name string, packet PRUDPPacketInterface) {
 			go handler(packet)
 		}
 	}
+
+	// * propagate the event up to the PRUDP server
+	pep.Server.emit(name, packet)
 }
 
 func (pep *PRUDPEndPoint) emitConnectionEnded(connection *PRUDPConnection) {
diff --git a/prudp_server.go b/prudp_server.go
index 626ea6a..1176fe9 100644
--- a/prudp_server.go
+++ b/prudp_server.go
@@ -37,6 +37,7 @@ type PRUDPServer struct {
 	PRUDPv1ConnectionSignatureKey []byte
 	byteStreamSettings            *ByteStreamSettings
 	PRUDPV0Settings               *PRUDPV0Settings
+	packetEventHandlers           map[string][]func(packet PacketInterface)
 }
 
 // BindPRUDPEndPoint binds a provided PRUDPEndPoint to the server
@@ -470,15 +471,38 @@ func (s *PRUDPServer) SetByteStreamSettings(byteStreamSettings *ByteStreamSettin
 	s.byteStreamSettings = byteStreamSettings
 }
 
+// OnData adds an event handler which is fired when a new DATA packet is received
+func (s *PRUDPServer) OnData(handler func(packet PacketInterface)) {
+	s.on("data", handler)
+}
+
+func (s *PRUDPServer) on(name string, handler func(packet PacketInterface)) {
+	if _, ok := s.packetEventHandlers[name]; !ok {
+		s.packetEventHandlers[name] = make([]func(packet PacketInterface), 0)
+	}
+
+	s.packetEventHandlers[name] = append(s.packetEventHandlers[name], handler)
+}
+
+// emit emits an event to all relevant listeners. These events fire after the PRUDPEndPoint event handlers
+func (s *PRUDPServer) emit(name string, packet PRUDPPacketInterface) {
+	if handlers, ok := s.packetEventHandlers[name]; ok {
+		for _, handler := range handlers {
+			go handler(packet)
+		}
+	}
+}
+
 // NewPRUDPServer will return a new PRUDP server
 func NewPRUDPServer() *PRUDPServer {
 	return &PRUDPServer{
-		Endpoints:          NewMutexMap[uint8, *PRUDPEndPoint](),
-		Connections:        NewMutexMap[string, *SocketConnection](),
-		kerberosKeySize:    32,
-		FragmentSize:       1300,
-		pingTimeout:        time.Second * 15,
-		byteStreamSettings: NewByteStreamSettings(),
-		PRUDPV0Settings:    NewPRUDPV0Settings(),
+		Endpoints:           NewMutexMap[uint8, *PRUDPEndPoint](),
+		Connections:         NewMutexMap[string, *SocketConnection](),
+		kerberosKeySize:     32,
+		FragmentSize:        1300,
+		pingTimeout:         time.Second * 15,
+		byteStreamSettings:  NewByteStreamSettings(),
+		PRUDPV0Settings:     NewPRUDPV0Settings(),
+		packetEventHandlers: make(map[string][]func(PacketInterface)),
 	}
 }
diff --git a/server_interface.go b/server_interface.go
index bed89aa..a3ae1a2 100644
--- a/server_interface.go
+++ b/server_interface.go
@@ -16,6 +16,7 @@ type ServerInterface interface {
 	NATTraversalProtocolVersion() *LibraryVersion
 	SetDefaultLibraryVersion(version *LibraryVersion)
 	Send(packet PacketInterface)
+	OnData(handler func(packet PacketInterface))
 	PasswordFromPID(pid *types.PID) (string, uint32)
 	SetPasswordFromPIDFunction(handler func(pid *types.PID) (string, uint32))
 	ByteStreamSettings() *ByteStreamSettings