From 07fe1b5686ebed44ae56072a264be319c7c4f2cb Mon Sep 17 00:00:00 2001
From: alex89200 <zenderfile@yahoo.com>
Date: Thu, 16 Mar 2023 20:28:33 +0200
Subject: [PATCH] Added support for receiving responses from the broadcasted
 messages

---
 kcp2k/kcp2k/highlevel/KcpClient.cs | 26 ++++++++++++++++++++++++--
 kcp2k/kcp2k/highlevel/KcpConfig.cs |  9 ++++++++-
 kcp2k/kcp2k/highlevel/KcpServer.cs |  6 ++++--
 3 files changed, 36 insertions(+), 5 deletions(-)

diff --git a/kcp2k/kcp2k/highlevel/KcpClient.cs b/kcp2k/kcp2k/highlevel/KcpClient.cs
index 788d065c..76da9d3d 100644
--- a/kcp2k/kcp2k/highlevel/KcpClient.cs
+++ b/kcp2k/kcp2k/highlevel/KcpClient.cs
@@ -111,11 +111,22 @@ void OnDisconnectedWrap()
             // even a 1ms block per connection would stop us from scaling.
             socket.Blocking = false;
 
+            // set flag to be able to send/recv broadcast messages if needed
+            socket.EnableBroadcast = config.EnableBroadcast;
+
             // configure buffer sizes
             Common.ConfigureSocketBuffers(socket, config.RecvBufferSize, config.SendBufferSize);
 
             // bind to endpoint so we can use send/recv instead of sendto/recvfrom.
-            socket.Connect(remoteEndPoint);
+            // But we don't do that if we want to be able to receive answers to
+            // broadcast requests, we must use socket.SendTo instead of socket.Connect.
+            // Otherwise broadcast address will be bound to the socket and it will be
+            // trying to receive data from it instead of the real address of
+            // the client and nothing will be received.
+            if (!config.EnableBroadcast)
+            {
+                socket.Connect(remoteEndPoint);
+            }
 
             // client should send handshake to server as very first message
             peer.SendHandshake();
@@ -155,7 +166,18 @@ protected virtual void RawSend(ArraySegment<byte> data)
         {
             try
             {
-                socket.SendNonBlocking(data);
+                // if broadcast enabled, we use connectionless mode and must send
+                // data using the SendTo method. Note that using this method causes
+                // allocations (remoteEndPoint is re-serialized each time). Thus it
+                // is recommended to use broadcast-enabled client ONLY for discovery.
+                if (config.EnableBroadcast)
+                {
+                    socket.SendToNonBlocking(data, remoteEndPoint);
+                }
+                else
+                {
+                    socket.SendNonBlocking(data);
+                }
             }
             catch (SocketException e)
             {
diff --git a/kcp2k/kcp2k/highlevel/KcpConfig.cs b/kcp2k/kcp2k/highlevel/KcpConfig.cs
index 382d06bb..6935d9f3 100644
--- a/kcp2k/kcp2k/highlevel/KcpConfig.cs
+++ b/kcp2k/kcp2k/highlevel/KcpConfig.cs
@@ -63,6 +63,11 @@ public class KcpConfig
         // maximum retransmission attempts until dead_link
         public uint MaxRetransmits;
 
+        // Whether to enable broadcasting. If true, the client will be able
+        // to broadcast messages and the server will be able to receive them
+        // and reply.
+        public bool EnableBroadcast;
+
         // constructor /////////////////////////////////////////////////////////
         // constructor with defaults for convenience.
         // makes it easy to define "new KcpConfig(DualMode=false)" etc.
@@ -78,7 +83,8 @@ public KcpConfig(
             uint SendWindowSize    = Kcp.WND_SND,
             uint ReceiveWindowSize = Kcp.WND_RCV,
             int Timeout            = KcpPeer.DEFAULT_TIMEOUT,
-            uint MaxRetransmits    = Kcp.DEADLINK)
+            uint MaxRetransmits    = Kcp.DEADLINK,
+            bool EnableBroadcast   = false)
         {
             this.DualMode = DualMode;
             this.RecvBufferSize = RecvBufferSize;
@@ -92,6 +98,7 @@ public KcpConfig(
             this.ReceiveWindowSize = ReceiveWindowSize;
             this.Timeout = Timeout;
             this.MaxRetransmits = MaxRetransmits;
+            this.EnableBroadcast = EnableBroadcast;
         }
     }
 }
diff --git a/kcp2k/kcp2k/highlevel/KcpServer.cs b/kcp2k/kcp2k/highlevel/KcpServer.cs
index 64b990cd..37094362 100644
--- a/kcp2k/kcp2k/highlevel/KcpServer.cs
+++ b/kcp2k/kcp2k/highlevel/KcpServer.cs
@@ -63,7 +63,7 @@ public KcpServer(Action<int> OnConnected,
 
         public virtual bool IsActive() => socket != null;
 
-        static Socket CreateServerSocket(bool DualMode, ushort port)
+        static Socket CreateServerSocket(bool DualMode, ushort port, bool enableBroadcast)
         {
             if (DualMode)
             {
@@ -81,6 +81,7 @@ static Socket CreateServerSocket(bool DualMode, ushort port)
                 {
                     Log.Warning($"Failed to set Dual Mode, continuing with IPv6 without Dual Mode. Error: {e}");
                 }
+                socket.EnableBroadcast = enableBroadcast;
                 socket.Bind(new IPEndPoint(IPAddress.IPv6Any, port));
                 return socket;
             }
@@ -88,6 +89,7 @@ static Socket CreateServerSocket(bool DualMode, ushort port)
             {
                 // IPv4 socket @ "0.0.0.0" : port
                 Socket socket = new Socket(AddressFamily.InterNetwork, SocketType.Dgram, ProtocolType.Udp);
+                socket.EnableBroadcast = enableBroadcast;
                 socket.Bind(new IPEndPoint(IPAddress.Any, port));
                 return socket;
             }
@@ -103,7 +105,7 @@ public virtual void Start(ushort port)
             }
 
             // listen
-            socket = CreateServerSocket(config.DualMode, port);
+            socket = CreateServerSocket(config.DualMode, port, config.EnableBroadcast);
 
             // recv & send are called from main thread.
             // need to ensure this never blocks.