Improve error reporting on oversized messages

src/enclave/rpcendpoint.h

@@ -47,7 +47,7 @@ namespace enclave
       return std::make_pair(actor, method);
     }
 
-    bool handle_data(const std::vector<uint8_t>& data)
+    bool handle_data(const std::vector<uint8_t>& data) override
     {
       LOG_DEBUG_FMT(
         "Entered handle_data, session {} with {} bytes",
@@ -129,5 +129,20 @@ namespace enclave
 
       return true;
     }
+
+    std::vector<uint8_t> oversized_message_error(
+      size_t msg_size, size_t max_msg_size) override
+    {
+      return jsonrpc::pack(
+        jsonrpc::error_response(
+          0,
+          jsonrpc::StandardErrorCodes::PARSE_ERROR,
+          fmt::format(
+            "Requested message ({} bytes) is too large. Maximum allowed is {} "
+            "bytes. Closing connection.",
+            msg_size,
+            max_msg_size)),
+        jsonrpc::Pack::Text);
+    }
   };
 }

src/enclave/tlsendpoint.h

@@ -17,7 +17,6 @@ namespace enclave
     std::unique_ptr<ringbuffer::AbstractWriter> to_host;
     size_t session_id;
 
-  private:
     enum Status
     {
       handshake,
@@ -27,6 +26,24 @@ namespace enclave
       error
     };
 
+    Status get_status() const
+    {
+      return status;
+    }
+
+    virtual std::vector<uint8_t> oversized_message_error(
+      size_t msg_size, size_t max_msg_size)
+    {
+      const auto s = fmt::format(
+        "Requested message ({} bytes) is too large. Maximum allowed is {} "
+        "bytes. Closing connection.",
+        msg_size,
+        max_msg_size);
+      const auto data = (const uint8_t*)s.data();
+      return std::vector<uint8_t>(data, data + s.size());
+    }
+
+  private:
     std::vector<uint8_t> pending_write;
     std::vector<uint8_t> pending_read;
     // Decrypted data, read through mbedtls

src/enclave/tlsframedendpoint.h

@@ -12,6 +12,8 @@ namespace enclave
     uint32_t msg_size;
     size_t count;
 
+    static constexpr size_t max_msg_size = 2 * 1024 * 1024;
+
   public:
     FramedTLSEndpoint(
       size_t session_id,
@@ -24,6 +26,16 @@ namespace enclave
 
     void recv(const uint8_t* data, size_t size)
     {
+      const auto status = get_status();
+      if (status >= closed)
+      {
+        LOG_INFO_FMT(
+          "Received additional data for {}, ignoring due to status {}",
+          session_id,
+          status);
+        return;
+      }
+
       recv_buffered(data, size);
 
       while (true)
@@ -43,8 +55,14 @@ namespace enclave
 
         // Arbitrary limit on RPC size to stop a client from requesting
         // a very large allocation.
-        if (msg_size > 1 << 21)
+        if (msg_size > max_msg_size)
         {
+          LOG_FAIL_FMT(
+            "Received oversized message request ({} bytes) - closing session "
+            "{}",
+            msg_size,
+            session_id);
+          send(oversized_message_error(msg_size, max_msg_size));
           close();
           return;
         }

src/host/tcp.h

@@ -31,7 +31,7 @@ namespace asynchost
     friend class close_ptr<TCPImpl>;
 
     static constexpr int backlog = 128;
-    static constexpr size_t read_size = 1024;
+    static constexpr size_t max_read_size = 16384;
 
     enum Status
     {
@@ -456,15 +456,17 @@ namespace asynchost
       return true;
     }
 
-    static void on_alloc(uv_handle_t* handle, size_t, uv_buf_t* buf)
+    static void on_alloc(
+      uv_handle_t* handle, size_t suggested_size, uv_buf_t* buf)
     {
-      static_cast<TCPImpl*>(handle->data)->on_alloc(buf);
+      static_cast<TCPImpl*>(handle->data)->on_alloc(suggested_size, buf);
     }
 
-    void on_alloc(uv_buf_t* buf)
+    void on_alloc(size_t suggested_size, uv_buf_t* buf)
     {
-      buf->base = new char[read_size];
-      buf->len = read_size;
+      auto alloc_size = std::min(suggested_size, max_read_size);
+      buf->base = new char[alloc_size];
+      buf->len = alloc_size;
     }
 
     void on_free(const uv_buf_t* buf)

tests/infra/jsonrpc.py

@@ -123,6 +123,15 @@ def from_json(self, data):
         self._from_parsed(parsed)
 
 
+def human_readable_size(n):
+    suffixes = ("B", "KB", "MB", "GB")
+    i = 0
+    while n >= 1024 and i < len(suffixes) - 1:
+        n /= 1024.0
+        i += 1
+    return f"{n:,.2f} {suffixes[i]}"
+
+
 class FramedTLSClient:
     def __init__(self, host, port, server_hostname, cert=None, key=None, cafile=None):
         self.host = host
@@ -159,11 +168,13 @@ def connect(self):
         self.conn.connect((self.host, self.port))
 
     def send(self, msg):
+        LOG.trace(f"Sending {human_readable_size(len(msg))} message")
         frame = struct.pack("<I", len(msg)) + msg
         self.conn.sendall(frame)
 
     def _read(self):
         (size,) = struct.unpack("<I", self.conn.recv(4))
+        LOG.trace(f"Reading {human_readable_size(size)} response")
         data = self.conn.recv(size)
         while len(data) < size:
             data += self.conn.recv(size - len(data))