Split udp/write operation into a 3-step beginPacket/write/endPacket

Author: cmaglie

network-api/network-api.go

@@ -47,7 +47,9 @@ func Register(router *msgpackrouter.Router) {
 	_ = router.RegisterMethod("tcp/connectSSL", tcpConnectSSL)
 
 	_ = router.RegisterMethod("udp/connect", udpConnect)
+	_ = router.RegisterMethod("udp/beginPacket", udpBeginPacket)
 	_ = router.RegisterMethod("udp/write", udpWrite)
+	_ = router.RegisterMethod("udp/endPacket", udpEndPacket)
 	_ = router.RegisterMethod("udp/awaitRead", udpAwaitRead)
 	_ = router.RegisterMethod("udp/read", udpRead)
 	_ = router.RegisterMethod("udp/close", udpClose)
@@ -58,6 +60,8 @@ var liveConnections = make(map[uint]net.Conn)
 var liveListeners = make(map[uint]net.Listener)
 var liveUdpConnections = make(map[uint]net.PacketConn)
 var udpReadBuffers = make(map[uint][]byte)
+var udpWriteTargets = make(map[uint]*net.UDPAddr)
+var udpWriteBuffers = make(map[uint][]byte)
 var nextConnectionID atomic.Uint32
 
 // takeLockAndGenerateNextID generates a new unique ID for a connection or listener.
@@ -375,9 +379,9 @@ func udpConnect(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (
 	return id, nil
 }
 
-func udpWrite(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (_result any, _err any) {
-	if len(params) != 4 {
-		return nil, []any{1, "Invalid number of parameters, expected udpConnId, dest address, dest port, payload"}
+func udpBeginPacket(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (_result any, _err any) {
+	if len(params) != 3 {
+		return nil, []any{1, "Invalid number of parameters, expected udpConnId, dest address, dest port"}
 	}
 	id, ok := msgpackrpc.ToUint(params[0])
 	if !ok {
@@ -391,9 +395,33 @@ func udpWrite(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (_r
 	if !ok {
 		return nil, []any{1, "Invalid parameter type, expected uint16 for server port"}
 	}
-	data, ok := params[3].([]byte)
+
+	lock.RLock()
+	defer lock.RUnlock()
+	if _, ok := liveUdpConnections[id]; !ok {
+		return nil, []any{2, fmt.Sprintf("UDP connection not found for ID: %d", id)}
+	}
+	targetAddr := net.JoinHostPort(targetIP, fmt.Sprintf("%d", targetPort))
+	addr, err := net.ResolveUDPAddr("udp", targetAddr) // TODO: This is inefficient, implement some caching
+	if err != nil {
+		return nil, []any{3, "Failed to resolve target address: " + err.Error()}
+	}
+	udpWriteTargets[id] = addr
+	udpWriteBuffers[id] = nil
+	return true, nil
+}
+
+func udpWrite(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (_result any, _err any) {
+	if len(params) != 2 {
+		return nil, []any{1, "Invalid number of parameters, expected expected udpConnId, payload"}
+	}
+	id, ok := msgpackrpc.ToUint(params[0])
+	if !ok {
+		return nil, []any{1, "Invalid parameter type, expected int for UDP connection ID"}
+	}
+	data, ok := params[1].([]byte)
 	if !ok {
-		if dataStr, ok := params[3].(string); ok {
+		if dataStr, ok := params[1].(string); ok {
 			data = []byte(dataStr)
 		} else {
 			// If data is not []byte or string, return an error
@@ -402,18 +430,45 @@ func udpWrite(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (_r
 	}
 
 	lock.RLock()
-	udpConn, ok := liveUdpConnections[id]
+	udpBuffer, ok := udpWriteBuffers[id]
+	if ok {
+		udpWriteBuffers[id] = append(udpBuffer, data...)
+	}
 	lock.RUnlock()
 	if !ok {
 		return nil, []any{2, fmt.Sprintf("UDP connection not found for ID: %d", id)}
 	}
+	return len(data), nil
+}
 
-	targetAddr := net.JoinHostPort(targetIP, fmt.Sprintf("%d", targetPort))
-	addr, err := net.ResolveUDPAddr("udp", targetAddr) // TODO: This is inefficient, implement some caching
-	if err != nil {
-		return nil, []any{3, "Failed to resolve target address: " + err.Error()}
+func udpEndPacket(ctx context.Context, rpc *msgpackrpc.Connection, params []any) (_result any, _err any) {
+	if len(params) != 1 {
+		return nil, []any{1, "Invalid number of parameters, expected expected udpConnId"}
 	}
-	if n, err := udpConn.WriteTo(data, addr); err != nil {
+	id, buffExists := msgpackrpc.ToUint(params[0])
+	if !buffExists {
+		return nil, []any{1, "Invalid parameter type, expected int for UDP connection ID"}
+	}
+
+	var udpBuffer []byte
+	var udpAddr *net.UDPAddr
+	lock.RLock()
+	udpConn, connExists := liveUdpConnections[id]
+	if connExists {
+		udpBuffer, buffExists = udpWriteBuffers[id]
+		udpAddr = udpWriteTargets[id]
+		delete(udpWriteBuffers, id)
+		delete(udpWriteTargets, id)
+	}
+	lock.RUnlock()
+	if !connExists {
+		return nil, []any{2, fmt.Sprintf("UDP connection not found for ID: %d", id)}
+	}
+	if !buffExists {
+		return nil, []any{3, fmt.Sprintf("No UDP packet begun for ID: %d", id)}
+	}
+
+	if n, err := udpConn.WriteTo(udpBuffer, udpAddr); err != nil {
 		return nil, []any{4, "Failed to write to UDP connection: " + err.Error()}
 	} else {
 		return n, nil

network-api/network-api_test.go

@@ -248,7 +248,13 @@ func TestUDPNetworkAPI(t *testing.T) {
 	require.NotEqual(t, conn1, conn2)
 
 	{
-		res, err := udpWrite(ctx, nil, []any{conn1, "127.0.0.1", 9900, []byte("Hello")})
+		res, err := udpBeginPacket(ctx, nil, []any{conn1, "127.0.0.1", 9900})
+		require.Nil(t, err)
+		require.True(t, res.(bool))
+		res, err = udpWrite(ctx, nil, []any{conn1, []byte("Hello")})
+		require.Nil(t, err)
+		require.Equal(t, 5, res)
+		res, err = udpEndPacket(ctx, nil, []any{conn1})
 		require.Nil(t, err)
 		require.Equal(t, 5, res)
 	}
@@ -262,12 +268,27 @@ func TestUDPNetworkAPI(t *testing.T) {
 		require.Equal(t, []uint8("Hello"), res2)
 	}
 	{
-		res, err := udpWrite(ctx, nil, []any{conn1, "127.0.0.1", 9900, []byte("One")})
+		res, err := udpBeginPacket(ctx, nil, []any{conn1, "127.0.0.1", 9900})
+		require.Nil(t, err)
+		require.True(t, res.(bool))
+		res, err = udpWrite(ctx, nil, []any{conn1, []byte("On")})
+		require.Nil(t, err)
+		require.Equal(t, 2, res)
+		res, err = udpWrite(ctx, nil, []any{conn1, []byte("e")})
+		require.Nil(t, err)
+		require.Equal(t, 1, res)
+		res, err = udpEndPacket(ctx, nil, []any{conn1})
 		require.Nil(t, err)
 		require.Equal(t, 3, res)
 	}
 	{
-		res, err := udpWrite(ctx, nil, []any{conn1, "127.0.0.1", 9900, []byte("Two")})
+		res, err := udpBeginPacket(ctx, nil, []any{conn1, "127.0.0.1", 9900})
+		require.Nil(t, err)
+		require.True(t, res.(bool))
+		res, err = udpWrite(ctx, nil, []any{conn1, []byte("Two")})
+		require.Nil(t, err)
+		require.Equal(t, 3, res)
+		res, err = udpEndPacket(ctx, nil, []any{conn1})
 		require.Nil(t, err)
 		require.Equal(t, 3, res)
 	}