From eec983c8b766e9ec7169fa4ead67f9c16f156ec6 Mon Sep 17 00:00:00 2001
From: Unknown <fred.nicolson@outlook.com>
Date: Sat, 3 Jun 2017 14:39:49 +0100
Subject: [PATCH] Added fr::Socket::set_max_packet_size

Can be used to limit the size of packets being received, to prevent malicious attacks.
---
 include/frnetlib/Socket.h | 18 ++++++++++++++++++
 src/Socket.cpp            | 17 +++++++++++++++--
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/include/frnetlib/Socket.h b/include/frnetlib/Socket.h
index 5b5e1d9..78dc2fa 100644
--- a/include/frnetlib/Socket.h
+++ b/include/frnetlib/Socket.h
@@ -26,6 +26,7 @@ namespace fr
             ConnectionFailed = 7,
             HandshakeFailed = 8,
             VerificationFailed = 9,
+            MaxPacketSizeExceeded = 10,
         };
 
         enum IP
@@ -158,6 +159,22 @@ namespace fr
          * @param version Should IPv4, IPv6 be used, or any?
          */
         void set_inet_version(IP version);
+
+        /*!
+         * Sets the maximum fr::Packet size that may be received by the socket.
+         *
+         * If a client attempts to send a packet larger than sz bytes, then
+         * the client will be disconnected and an fr::Socket::MaxPacketSizeExceeded
+         * will be returned. Pass '0' to indicate no limit. The default value is 0.
+         *
+         * This should be used to prevent potential abuse, as a client could say that
+         * it's going to send a 200GiB packet, which would cause the Socket to try and
+         * allocate that much memory to accommodate the data, which is most likely not
+         * desirable.
+         *
+         * @param sz The maximum number of bytes that may be received in an fr::Packet
+         */
+        void set_max_packet_size(uint32_t sz);
     protected:
 
         /*!
@@ -171,6 +188,7 @@ namespace fr
         std::mutex outbound_mutex;
         std::mutex inbound_mutex;
         int ai_family;
+        uint32_t max_packet_size;
         #ifdef _WIN32
                 static WSADATA wsaData;
         #endif // _WIN32
diff --git a/src/Socket.cpp b/src/Socket.cpp
index e044686..354f40f 100644
--- a/src/Socket.cpp
+++ b/src/Socket.cpp
@@ -16,7 +16,8 @@ namespace fr
 
     Socket::Socket() noexcept
     : is_blocking(true),
-      ai_family(AF_UNSPEC)
+      ai_family(AF_UNSPEC),
+      max_packet_size(0)
     {
             if(instance_count == 0)
             {
@@ -68,6 +69,13 @@ namespace fr
             return status;
         packet_length = ntohl(packet_length);
 
+        //Check that packet_length doesn't exceed the limit, if any
+        if(max_packet_size && max_packet_size > max_packet_size)
+        {
+            close_socket();
+            return fr::Socket::MaxPacketSizeExceeded;
+        }
+
         //Now we've got the length, read the rest of the data in
         packet.buffer.resize(packet_length + PACKET_HEADER_LENGTH);
         status = receive_all(&packet.buffer[PACKET_HEADER_LENGTH], packet_length);
@@ -83,7 +91,7 @@ namespace fr
         if(!connected())
             return Socket::Disconnected;
 
-		int32_t bytes_remaining = buffer_size;
+		int32_t bytes_remaining = (int32_t) buffer_size;
         size_t bytes_read = 0;
         while(bytes_remaining > 0)
         {
@@ -128,4 +136,9 @@ namespace fr
                 throw std::logic_error("Unknown Socket::IP value passed to set_inet_version()");
         }
     }
+
+    void Socket::set_max_packet_size(uint32_t sz)
+    {
+        max_packet_size = sz;
+    }
 }
\ No newline at end of file