ensure fragment outgoing commands are allocated atomically with respect to failure

3 files changed, 82 insertions, 41 deletions

diff --git a/include/enet/enet.h b/include/enet/enet.h
index dca564c..7870723 100644
--- a/include/enet/enet.h
+++ b/include/enet/enet.h
@@ -478,6 +478,7 @@ ENET_API void                enet_peer_disconnect_later (ENetPeer *, enet_uint32
 ENET_API void                enet_peer_throttle_configure (ENetPeer *, enet_uint32, enet_uint32, enet_uint32);
 extern int                   enet_peer_throttle (ENetPeer *, enet_uint32);
 extern void                  enet_peer_reset_queues (ENetPeer *);
+extern void                  enet_peer_setup_outgoing_command (ENetPeer *, ENetOutgoingCommand *, ENetPacket *, enet_uint32, enet_uint16);
 extern ENetOutgoingCommand * enet_peer_queue_outgoing_command (ENetPeer *, const ENetProtocol *, ENetPacket *, enet_uint32, enet_uint16);
 extern ENetIncomingCommand * enet_peer_queue_incoming_command (ENetPeer *, const ENetProtocol *, ENetPacket *, enet_uint32);
 extern ENetAcknowledgement * enet_peer_queue_acknowledgement (ENetPeer *, const ENetProtocol *, enet_uint16);
diff --git a/peer.c b/peer.c
index d00eae7..2274f3e 100644
--- a/peer.c
+++ b/peer.c
@@ -115,9 +115,10 @@ enet_peer_send (ENetPeer * peer, enet_uint8 channelID, ENetPacket * packet)
       enet_uint32 fragmentCount = ENET_HOST_TO_NET_32 ((packet -> dataLength + fragmentLength - 1) / fragmentLength),
              fragmentNumber,
              fragmentOffset;
+      ENetList fragments;
+      ENetOutgoingCommand * fragment;
 
-      packet -> flags |= ENET_PACKET_FLAG_RELIABLE;
-      packet -> flags &= ~ENET_PACKET_FLAG_UNSEQUENCED;
+      enet_list_clear (& fragments);
 
       for (fragmentNumber = 0,
              fragmentOffset = 0;
@@ -128,17 +129,36 @@ enet_peer_send (ENetPeer * peer, enet_uint8 channelID, ENetPacket * packet)
          if (packet -> dataLength - fragmentOffset < fragmentLength)
            fragmentLength = packet -> dataLength - fragmentOffset;
 
-         command.header.command = ENET_PROTOCOL_COMMAND_SEND_FRAGMENT | ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE;
-         command.header.channelID = channelID;
-         command.sendFragment.startSequenceNumber = startSequenceNumber;
-         command.sendFragment.dataLength = ENET_HOST_TO_NET_16 (fragmentLength);
-         command.sendFragment.fragmentCount = fragmentCount;
-         command.sendFragment.fragmentNumber = ENET_HOST_TO_NET_32 (fragmentNumber);
-         command.sendFragment.totalLength = ENET_HOST_TO_NET_32 (packet -> dataLength);
-         command.sendFragment.fragmentOffset = ENET_NET_TO_HOST_32 (fragmentOffset);
-
-         if (enet_peer_queue_outgoing_command (peer, & command, packet, fragmentOffset, fragmentLength) == NULL)
-           return -1;
+         fragment = (ENetOutgoingCommand *) enet_malloc (sizeof (ENetOutgoingCommand));
+         if (fragment == NULL)
+         {
+            while (! enet_list_empty (& fragments))
+            {
+               fragment = (ENetOutgoingCommand *) enet_list_remove (enet_list_begin (& fragments));
+               
+               enet_free (fragment);
+            }
+            
+            return -1;
+         }
+                
+         fragment -> command.header.command = ENET_PROTOCOL_COMMAND_SEND_FRAGMENT | ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE;
+         fragment -> command.header.channelID = channelID;
+         fragment -> command.sendFragment.startSequenceNumber = startSequenceNumber;
+         fragment -> command.sendFragment.dataLength = ENET_HOST_TO_NET_16 (fragmentLength);
+         fragment -> command.sendFragment.fragmentCount = fragmentCount;
+         fragment -> command.sendFragment.fragmentNumber = ENET_HOST_TO_NET_32 (fragmentNumber);
+         fragment -> command.sendFragment.totalLength = ENET_HOST_TO_NET_32 (packet -> dataLength);
+         fragment -> command.sendFragment.fragmentOffset = ENET_NET_TO_HOST_32 (fragmentOffset);
+        
+         enet_list_insert (enet_list_end (& fragments), fragment);
+      }
+
+      while (! enet_list_empty (& fragments))
+      {
+         fragment = (ENetOutgoingCommand *) enet_list_remove (enet_list_begin (& fragments));
+ 
+         enet_peer_setup_outgoing_command (peer, fragment, packet, fragmentOffset, fragmentLength);
       }
 
       return 0;
@@ -146,9 +166,6 @@ enet_peer_send (ENetPeer * peer, enet_uint8 channelID, ENetPacket * packet)
 
    command.header.channelID = channelID;
 
-   if (! (packet -> flags & (ENET_PACKET_FLAG_RELIABLE | ENET_PACKET_FLAG_UNSEQUENCED)) && channel -> outgoingUnreliableSequenceNumber >= 0xFFFF)
-     packet -> flags |= ENET_PACKET_FLAG_RELIABLE;
-
    if (packet -> flags & ENET_PACKET_FLAG_RELIABLE)
    {
       command.header.command = ENET_PROTOCOL_COMMAND_SEND_RELIABLE | ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE;
@@ -161,6 +178,12 @@ enet_peer_send (ENetPeer * peer, enet_uint8 channelID, ENetPacket * packet)
       command.sendUnsequenced.unsequencedGroup = ENET_HOST_TO_NET_16 (peer -> outgoingUnsequencedGroup + 1);
       command.sendUnsequenced.dataLength = ENET_HOST_TO_NET_16 (packet -> dataLength);
    }
+   else 
+   if (channel -> outgoingUnreliableSequenceNumber >= 0xFFFF)
+   {
+      command.header.command = ENET_PROTOCOL_COMMAND_SEND_RELIABLE | ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE;
+      command.sendReliable.dataLength = ENET_HOST_TO_NET_16 (packet -> dataLength);
+   }
    else
    {
       command.header.command = ENET_PROTOCOL_COMMAND_SEND_UNRELIABLE;
@@ -503,17 +526,14 @@ enet_peer_queue_acknowledgement (ENetPeer * peer, const ENetProtocol * command,
     return acknowledgement;
 }
 
-ENetOutgoingCommand *
-enet_peer_queue_outgoing_command (ENetPeer * peer, const ENetProtocol * command, ENetPacket * packet, enet_uint32 offset, enet_uint16 length)
+void
+enet_peer_setup_outgoing_command (ENetPeer * peer, ENetOutgoingCommand * outgoingCommand, ENetPacket * packet, enet_uint32 offset, enet_uint16 length)
 {
-    ENetChannel * channel = & peer -> channels [command -> header.channelID];
-    ENetOutgoingCommand * outgoingCommand = (ENetOutgoingCommand *) enet_malloc (sizeof (ENetOutgoingCommand));
-    if (outgoingCommand == NULL)
-      return NULL;
+    ENetChannel * channel = & peer -> channels [outgoingCommand -> command.header.channelID];
 
-    peer -> outgoingDataTotal += enet_protocol_command_size (command -> header.command) + length;
+    peer -> outgoingDataTotal += enet_protocol_command_size (outgoingCommand -> command.header.command) + length;
 
-    if (command -> header.channelID == 0xFF)
+    if (outgoingCommand -> command.header.channelID == 0xFF)
     {
        ++ peer -> outgoingReliableSequenceNumber;
 
@@ -521,7 +541,7 @@ enet_peer_queue_outgoing_command (ENetPeer * peer, const ENetProtocol * command,
        outgoingCommand -> unreliableSequenceNumber = 0;
     }
     else
-    if (command -> header.command & ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE)
+    if (outgoingCommand -> command.header.command & ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE)
     {
        ++ channel -> outgoingReliableSequenceNumber;
        channel -> outgoingUnreliableSequenceNumber = 0;
@@ -530,7 +550,7 @@ enet_peer_queue_outgoing_command (ENetPeer * peer, const ENetProtocol * command,
        outgoingCommand -> unreliableSequenceNumber = 0;
     }
     else
-    if (command -> header.command & ENET_PROTOCOL_COMMAND_FLAG_UNSEQUENCED)
+    if (outgoingCommand -> command.header.command & ENET_PROTOCOL_COMMAND_FLAG_UNSEQUENCED)
     {
        ++ peer -> outgoingUnsequencedGroup;
 
@@ -552,16 +572,27 @@ enet_peer_queue_outgoing_command (ENetPeer * peer, const ENetProtocol * command,
     outgoingCommand -> fragmentOffset = offset;
     outgoingCommand -> fragmentLength = length;
     outgoingCommand -> packet = packet;
-    outgoingCommand -> command = * command;
     outgoingCommand -> command.header.reliableSequenceNumber = ENET_HOST_TO_NET_16 (outgoingCommand -> reliableSequenceNumber);
 
     if (packet != NULL)
       ++ packet -> referenceCount;
 
-    if (command -> header.command & ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE)
+    if (outgoingCommand -> command.header.command & ENET_PROTOCOL_COMMAND_FLAG_ACKNOWLEDGE)
       enet_list_insert (enet_list_end (& peer -> outgoingReliableCommands), outgoingCommand);
     else
       enet_list_insert (enet_list_end (& peer -> outgoingUnreliableCommands), outgoingCommand);
+}
+
+ENetOutgoingCommand *
+enet_peer_queue_outgoing_command (ENetPeer * peer, const ENetProtocol * command, ENetPacket * packet, enet_uint32 offset, enet_uint16 length)
+{
+    ENetOutgoingCommand * outgoingCommand = (ENetOutgoingCommand *) enet_malloc (sizeof (ENetOutgoingCommand));
+    if (outgoingCommand == NULL)
+      return NULL;
+
+    outgoingCommand -> command = * command;
+
+    enet_peer_setup_outgoing_command (peer, outgoingCommand, packet, offset, length);
 
     return outgoingCommand;
 }
@@ -569,6 +600,8 @@ enet_peer_queue_outgoing_command (ENetPeer * peer, const ENetProtocol * command,
 ENetIncomingCommand *
 enet_peer_queue_incoming_command (ENetPeer * peer, const ENetProtocol * command, ENetPacket * packet, enet_uint32 fragmentCount)
 {
+    static ENetIncomingCommand dummyCommand;
+
     ENetChannel * channel = & peer -> channels [command -> header.channelID];
     enet_uint32 unreliableSequenceNumber = 0, reliableSequenceNumber;
     enet_uint16 reliableWindow, currentWindow;
@@ -670,7 +703,7 @@ enet_peer_queue_incoming_command (ENetPeer * peer, const ENetProtocol * command,
 
     incomingCommand = (ENetIncomingCommand *) enet_malloc (sizeof (ENetIncomingCommand));
     if (incomingCommand == NULL)
-      goto freePacket;
+      goto notifyError;
 
     incomingCommand -> reliableSequenceNumber = command -> header.reliableSequenceNumber;
     incomingCommand -> unreliableSequenceNumber = unreliableSequenceNumber & 0xFFFF;
@@ -686,7 +719,8 @@ enet_peer_queue_incoming_command (ENetPeer * peer, const ENetProtocol * command,
        if (incomingCommand -> fragments == NULL)
        {
           enet_free (incomingCommand);
-          goto freePacket;
+
+          goto notifyError;
        }
        memset (incomingCommand -> fragments, 0, (fragmentCount + 31) / 32 * sizeof (enet_uint32));
     }
@@ -699,11 +733,17 @@ enet_peer_queue_incoming_command (ENetPeer * peer, const ENetProtocol * command,
     return incomingCommand;
 
 freePacket:
-    if (packet != NULL)
-    {
-       if (packet -> referenceCount == 0)
-         enet_packet_destroy (packet);
-    }
+    if (fragmentCount > 0)
+      goto notifyError;
+
+    if (packet != NULL && packet -> referenceCount == 0)
+      enet_packet_destroy (packet);
+
+    return & dummyCommand;
+
+notifyError:
+    if (packet != NULL && packet -> referenceCount == 0)
+      enet_packet_destroy (packet);
 
     return NULL;
 }
diff --git a/protocol.c b/protocol.c
index 2434e20..8518d1f 100644
--- a/protocol.c
+++ b/protocol.c
@@ -398,10 +398,10 @@ enet_protocol_handle_send_reliable (ENetHost * host, ENetPeer * peer, const ENet
     packet = enet_packet_create ((const enet_uint8 *) command + sizeof (ENetProtocolSendReliable),
                                  dataLength,
                                  ENET_PACKET_FLAG_RELIABLE);
-    if (packet == NULL)
+    if (packet == NULL ||
+        enet_peer_queue_incoming_command (peer, command, packet, 0) == NULL)
       return -1;
 
-    enet_peer_queue_incoming_command (peer, command, packet, 0);
     return 0;
 }
 
@@ -445,12 +445,12 @@ enet_protocol_handle_send_unsequenced (ENetHost * host, ENetPeer * peer, const E
     packet = enet_packet_create ((const enet_uint8 *) command + sizeof (ENetProtocolSendUnsequenced),
                                  dataLength,
                                  ENET_PACKET_FLAG_UNSEQUENCED);
-    if (packet == NULL)
+    if (packet == NULL ||
+        enet_peer_queue_incoming_command (peer, command, packet, 0) == NULL)
       return -1;
    
     peer -> unsequencedWindow [index / 32] |= 1 << (index % 32);
  
-    enet_peer_queue_incoming_command (peer, command, packet, 0);
     return 0;
 }
 
@@ -472,10 +472,10 @@ enet_protocol_handle_send_unreliable (ENetHost * host, ENetPeer * peer, const EN
     packet = enet_packet_create ((const enet_uint8 *) command + sizeof (ENetProtocolSendUnreliable),
                                  dataLength,
                                  0);
-    if (packet == NULL)
+    if (packet == NULL ||
+        enet_peer_queue_incoming_command (peer, command, packet, 0) == NULL)
       return -1;
 
-    enet_peer_queue_incoming_command (peer, command, packet, 0);
     return 0;
 }