From c5189c090fc47503790bcb994232629897d61d7c Mon Sep 17 00:00:00 2001
From: eihrul <eihrul>
Date: Fri, 14 May 2010 18:20:20 +0000
Subject: [PATCH] added enet_host_channel_limit

---
 ChangeLog           |  2 ++
 host.c              | 17 +++++++++++++++++
 include/enet/enet.h |  3 +++
 protocol.c          | 10 +++++++++-
 4 files changed, 31 insertions(+), 1 deletion(-)

diff --git a/ChangeLog b/ChangeLog
index c9219d7..8520432 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -1,5 +1,7 @@
 ENet 1.2.2 (May 13, 2010):
 
+* added enet_host_channel_limit() for limiting the maximum number of
+channels allowed by connected peers
 * now uses dispatch queues for event dispatch rather than potentially
 unscalable array walking
 * added no_memory callback that is called when a malloc attempt fails,
diff --git a/host.c b/host.c
index 8d07e80..9edceb2 100644
--- a/host.c
+++ b/host.c
@@ -66,6 +66,7 @@ enet_host_create (const ENetAddress * address, size_t peerCount, enet_uint32 inc
     if (address != NULL)
       host -> address = * address;
 
+    host -> channelLimit = ENET_PROTOCOL_MAXIMUM_CHANNEL_COUNT;
     host -> incomingBandwidth = incomingBandwidth;
     host -> outgoingBandwidth = outgoingBandwidth;
     host -> bandwidthThrottleEpoch = 0;
@@ -208,6 +209,22 @@ enet_host_connect (ENetHost * host, const ENetAddress * address, size_t channelC
     return currentPeer;
 }
 
+/** Limits the maximum allowed channels of future incoming connections.
+    @param host host to limit
+    @param channelLimit the maximum number of channels allowed; if 0, then this is equivalent to ENET_PROTOCOL_MAXIMUM_CHANNEL_COUNT
+*/
+void
+enet_host_channel_limit (ENetHost * host, size_t channelLimit)
+{
+    if (! channelLimit || channelLimit > ENET_PROTOCOL_MAXIMUM_CHANNEL_COUNT)
+      channelLimit = ENET_PROTOCOL_MAXIMUM_CHANNEL_COUNT;
+    else
+    if (channelLimit < ENET_PROTOCOL_MINIMUM_CHANNEL_COUNT)
+      channelLimit = ENET_PROTOCOL_MINIMUM_CHANNEL_COUNT;
+
+    host -> channelLimit = channelLimit;
+}
+ 
 /** Queues a packet to be sent to all peers associated with the host.
     @param host host on which to broadcast the packet
     @param channelID channel on which to broadcast
diff --git a/include/enet/enet.h b/include/enet/enet.h
index 2ffcb3f..04b416e 100644
--- a/include/enet/enet.h
+++ b/include/enet/enet.h
@@ -294,6 +294,7 @@ typedef struct _ENetPeer
     @sa enet_host_service()
     @sa enet_host_flush()
     @sa enet_host_broadcast()
+    @sa enet_host_channel_limit()
     @sa enet_host_bandwidth_limit()
     @sa enet_host_bandwidth_throttle()
   */
@@ -308,6 +309,7 @@ typedef struct _ENetHost
    int                recalculateBandwidthLimits;
    ENetPeer *         peers;                       /**< array of peers allocated for this host */
    size_t             peerCount;                   /**< number of peers allocated for this host */
+   size_t             channelLimit;                /**< maximum number of channels allowed for connected peers */
    enet_uint32        serviceTime;
    ENetList           dispatchQueue;
    int                continueSending;
@@ -471,6 +473,7 @@ ENET_API int        enet_host_check_events (ENetHost *, ENetEvent *);
 ENET_API int        enet_host_service (ENetHost *, ENetEvent *, enet_uint32);
 ENET_API void       enet_host_flush (ENetHost *);
 ENET_API void       enet_host_broadcast (ENetHost *, enet_uint8, ENetPacket *);
+ENET_API void       enet_host_channel_limit (ENetHost *, size_t);
 ENET_API void       enet_host_bandwidth_limit (ENetHost *, enet_uint32, enet_uint32);
 extern   void       enet_host_bandwidth_throttle (ENetHost *);
 
diff --git a/protocol.c b/protocol.c
index 544feb2..1362179 100644
--- a/protocol.c
+++ b/protocol.c
@@ -292,6 +292,8 @@ enet_protocol_handle_connect (ENetHost * host, ENetProtocolHeader * header, ENet
     if (currentPeer >= & host -> peers [host -> peerCount])
       return NULL;
 
+    if (channelCount > host -> channelLimit)
+      channelCount = host -> channelLimit;
     currentPeer -> channels = (ENetChannel *) enet_malloc (channelCount * sizeof (ENetChannel));
     if (currentPeer -> channels == NULL)
       return NULL;
@@ -745,11 +747,14 @@ enet_protocol_handle_verify_connect (ENetHost * host, ENetEvent * event, ENetPee
 {
     enet_uint16 mtu;
     enet_uint32 windowSize;
+    size_t channelCount;
 
     if (peer -> state != ENET_PEER_STATE_CONNECTING)
       return 0;
 
-    if (ENET_NET_TO_HOST_32 (command -> verifyConnect.channelCount) != peer -> channelCount ||
+    channelCount = ENET_NET_TO_HOST_32 (command -> verifyConnect.channelCount);
+
+    if (channelCount < ENET_PROTOCOL_MINIMUM_CHANNEL_COUNT || channelCount > ENET_PROTOCOL_MAXIMUM_CHANNEL_COUNT ||
         ENET_NET_TO_HOST_32 (command -> verifyConnect.packetThrottleInterval) != peer -> packetThrottleInterval ||
         ENET_NET_TO_HOST_32 (command -> verifyConnect.packetThrottleAcceleration) != peer -> packetThrottleAcceleration ||
         ENET_NET_TO_HOST_32 (command -> verifyConnect.packetThrottleDeceleration) != peer -> packetThrottleDeceleration)
@@ -761,6 +766,9 @@ enet_protocol_handle_verify_connect (ENetHost * host, ENetEvent * event, ENetPee
 
     enet_protocol_remove_sent_reliable_command (peer, 1, 0xFF);
     
+    if (channelCount < peer -> channelCount)
+      peer -> channelCount = channelCount;
+
     peer -> outgoingPeerID = ENET_NET_TO_HOST_16 (command -> verifyConnect.outgoingPeerID);
 
     mtu = ENET_NET_TO_HOST_16 (command -> verifyConnect.mtu);