base: stream HTTP and websocket payloads into the handler's buffer

This makes it such that instead of the HttpServer providing the buffer
for HTTP/Websocket payloads, it instead asks the handler for where they
would like the payload to be stored.

The flow looks like this: socket read -> HTTPServer reads the
headers/framing -> once it knows the size it calls OnHttpRequestBody or
OnWebsocketPayload depending on the message type -> the handler provides
the storage bytes -> httpserver reads into those provided buffers.

The motivation for this is that it removes one copy allowing for
"zero-copy" parsing and tokenization etc.

The motivation for this is multi-threaded trace processor which requires
fast-handoffs of buffers between threads and reducing the latency of the
parse path (because we're going to be adding latency by a cross thread
hop). This change is one step on that journey.

Because a payload is now read across several socket reads, more than one
can be in flight at a time, so every buffer the handler hands out has to
belong to the connection asking for it. RPC payloads go into that
connection's Rpc::Stream, added in the previous change; the bodies of
the non-RPC endpoints go into a per-connection buffer alongside it.
Sharing either across connections lets one connection's payload be
dispatched into another's, or freed under it when the buffer grows.

The origin check moves above the point where the body is requested, so a
request that is about to be refused with a 403 is never handed a payload
sink to write into.

The reservation is held in the connection's state between the two calls,
so a peer that disappears mid-payload gives it back explicitly rather
than leaving the tokenizer believing a write is still in flight.

include/perfetto/ext/base/http/http_server.h

diff --git a/include/perfetto/ext/base/http/http_server.h b/include/perfetto/ext/base/http/http_server.h
index 93ad156..5a6aac2 100644
--- a/include/perfetto/ext/base/http/http_server.h
+++ b/include/perfetto/ext/base/http/http_server.h
@@ -42,8 +42,9 @@
 
   HttpServerConnection* conn;
 
-  // These StringViews point to memory in the rxbuf owned by |conn|. They are
-  // valid only within the OnHttpRequest() call.
+  // |method|, |uri| and |origin| point into the rxbuf owned by |conn|, |body|
+  // into the buffer returned by OnHttpRequestBody(). All are valid only within
+  // the OnHttpRequest() call.
   StringView method;
   StringView uri;
   StringView origin;
@@ -70,8 +71,8 @@
   // Note: message boundaries are not respected in case of fragmentation.
   // This websocket implementation preserves only the byte stream, but not the
   // atomicity of inbound messages (like SOCK_STREAM, unlike SOCK_DGRAM).
-  // Holds onto the connection's |rxbuf|. This is valid only within the scope
-  // of the OnWebsocketMessage() callback.
+  // Points into the buffer the handler returned from OnWebsocketPayload().
+  // This is valid only within the scope of the OnWebsocketMessage() callback.
   StringView data;
 
   // If false the payload contains binary data. If true it's supposed to contain
@@ -128,10 +129,20 @@
   friend class HttpServer;
 
   size_t rxbuf_avail() { return rxbuf.size() - rxbuf_used; }
+  void ClearPayload();
 
   std::unique_ptr<UnixSocket> sock;
+
+  // Holds request lines, headers and websocket frame headers only. Payloads
+  // are read straight into the buffer the handler supplies.
   PagedMemory rxbuf;
   size_t rxbuf_used = 0;
+
+  // Set while a payload is being read into the handler's buffer.
+  uint8_t* payload_ = nullptr;
+  size_t payload_size_ = 0;
+  size_t payload_used_ = 0;
+  uint8_t payload_mask_[4]{};  // Websocket masking key for |payload_|.
   bool is_websocket_ = false;
   bool headers_sent_ = false;
   size_t content_len_headers_ = 0;
@@ -150,6 +161,17 @@
 class HttpRequestHandler {
  public:
   virtual ~HttpRequestHandler();
+
+  // Called once the size of an inbound payload (an HTTP request body or a
+  // websocket frame payload) is known, to obtain the memory it is read into.
+  // The buffer must hold |size| bytes and stay valid until the matching
+  // OnHttpRequest()/OnWebsocketMessage() returns. Returning nullptr rejects the
+  // payload (413 for HTTP, close for websockets).
+  // Payloads on different connections can be in flight at the same time, so a
+  // handler serving more than one needs per-connection storage.
+  virtual uint8_t* OnHttpRequestBody(const HttpRequest&, size_t size) = 0;
+  virtual uint8_t* OnWebsocketPayload(HttpServerConnection*, size_t size) = 0;
+
   virtual void OnHttpRequest(const HttpRequest&) = 0;
   virtual void OnWebsocketMessage(const WebsocketMessage&);
   virtual void OnHttpConnectionClosed(HttpServerConnection*);

src/base/http/http_server.cc

diff --git a/src/base/http/http_server.cc b/src/base/http/http_server.cc
index 2b35925..f6c8c39 100644
--- a/src/base/http/http_server.cc
+++ b/src/base/http/http_server.cc
@@ -15,6 +15,8 @@
  */
 #include "perfetto/ext/base/http/http_server.h"
 
+#include <algorithm>
+
 #include <cinttypes>
 
 #include <cstddef>
@@ -41,7 +43,9 @@
 
 namespace {
 constexpr size_t kMaxPayloadSize = 64 * 1024 * 1024;
-constexpr size_t kMaxRequestSize = kMaxPayloadSize + 4096;
+// |rxbuf| holds only headers, so this also bounds how much of a payload can
+// arrive in the same read() as them and has to be moved to the handler.
+constexpr size_t kMaxHeadersSize = 64 * 1024;
 
 enum WebsocketOpcode : uint8_t {
   kOpcodeContinuation = 0x0,
@@ -142,22 +146,20 @@
   PERFETTO_CHECK(conn);
 
   char* rxbuf = reinterpret_cast<char*>(conn->rxbuf.Get());
+  // Reading and parsing interleave: it's the parser that discovers that the
+  // bytes still on the socket are a payload, not headers.
   for (;;) {
-    size_t avail = conn->rxbuf_avail();
-    PERFETTO_CHECK(avail <= kMaxRequestSize);
-    if (avail == 0) {
-      conn->SendResponseAndClose("413 Payload Too Large");
-      return;
+    if (conn->payload_used_ < conn->payload_size_) {
+      size_t rsize = sock->Receive(conn->payload_ + conn->payload_used_,
+                                   conn->payload_size_ - conn->payload_used_);
+      if (rsize == 0)
+        return;
+      conn->payload_used_ += rsize;
+      continue;
     }
-    size_t rsize = sock->Receive(&rxbuf[conn->rxbuf_used], avail);
-    conn->rxbuf_used += rsize;
-    if (rsize == 0 || conn->rxbuf_avail() == 0)
-      break;
-  }
 
-  // At this point |rxbuf| can contain a partial HTTP request, a full one or
-  // more (in case of HTTP Keepalive pipelining).
-  for (;;) {
+    // At this point |rxbuf| can contain a partial HTTP request, a full one or
+    // more (in case of HTTP Keepalive pipelining).
     size_t bytes_consumed;
 
     if (conn->is_websocket()) {
@@ -166,13 +168,33 @@
       bytes_consumed = ParseOneHttpRequest(conn);
     }
 
-    if (bytes_consumed == 0)
-      break;
-    memmove(rxbuf, &rxbuf[bytes_consumed], conn->rxbuf_used - bytes_consumed);
-    conn->rxbuf_used -= bytes_consumed;
+    if (bytes_consumed != 0) {
+      memmove(rxbuf, &rxbuf[bytes_consumed], conn->rxbuf_used - bytes_consumed);
+      conn->rxbuf_used -= bytes_consumed;
+      continue;
+    }
+    // The parse may have just installed a payload sink.
+    if (conn->payload_used_ < conn->payload_size_)
+      continue;
+
+    size_t avail = conn->rxbuf_avail();
+    if (avail == 0) {
+      conn->SendResponseAndClose("431 Request Header Fields Too Large");
+      return;
+    }
+    size_t rsize = sock->Receive(&rxbuf[conn->rxbuf_used], avail);
+    if (rsize == 0)
+      return;
+    conn->rxbuf_used += rsize;
   }
 }
 
+void HttpServerConnection::ClearPayload() {
+  payload_ = nullptr;
+  payload_size_ = 0;
+  payload_used_ = 0;
+}
+
 // Parses the HTTP request and invokes HandleRequest(). It returns the size of
 // the HTTP header + body that has been processed or 0 if there isn't enough
 // data for a full HTTP request in the buffer.
@@ -251,18 +273,50 @@
   PERFETTO_CHECK(buf_view.size() <= conn->rxbuf_used);
   const size_t headers_size = conn->rxbuf_used - buf_view.size();
 
-  if (body_size + headers_size >= kMaxRequestSize ||
-      body_size > kMaxPayloadSize) {
+  if (body_size > kMaxPayloadSize) {
     conn->SendResponseAndClose("413 Payload Too Large");
     return 0;
   }
 
   // If we can't read the full request return and try again next time with more
   // data.
-  if (!all_headers_received || buf_view.size() < body_size)
+  if (!all_headers_received)
     return 0;
 
-  http_req.body = buf_view.substr(0, body_size);
+  // CSRF defense: reject cross-origin requests from outside the allowlist.
+  // Before the body is requested, so a request about to be refused is never
+  // handed a payload sink. OPTIONS is the preflight itself, handled below.
+  if (http_req.method != "OPTIONS" && !http_req.origin.empty() &&
+      !IsOriginAllowed(http_req.origin)) {
+    conn->SendResponseAndClose("403 Forbidden", {}, "Origin not allowed");
+    return 0;
+  }
+
+  // Only what arrived in the same read() as the headers has to be moved; the
+  // rest is read from the socket into |dst|. The headers stay in |rxbuf| and
+  // are re-parsed on completion, rather than kept alive across the read.
+  if (body_size > 0 && conn->payload_ == nullptr) {
+    uint8_t* dst = req_handler_->OnHttpRequestBody(http_req, body_size);
+    if (dst == nullptr) {
+      conn->SendResponseAndClose("413 Payload Too Large");
+      return 0;
+    }
+    conn->payload_ = dst;
+    conn->payload_size_ = body_size;
+    conn->payload_used_ = std::min(buf_view.size(), body_size);
+    memcpy(dst, buf_view.data(), conn->payload_used_);
+    memmove(&rxbuf[headers_size], &rxbuf[headers_size + conn->payload_used_],
+            conn->rxbuf_used - headers_size - conn->payload_used_);
+    conn->rxbuf_used -= conn->payload_used_;
+  }
+  if (conn->payload_used_ < conn->payload_size_)
+    return 0;
+
+  http_req.body =
+      body_size == 0 ? buf_view.substr(0, 0)
+                     : StringView(reinterpret_cast<const char*>(conn->payload_),
+                                  body_size);
+  conn->ClearPayload();
 
   PERFETTO_LOG("[HTTP] %.*s %.*s [body=%zuB, origin=\"%.*s\"]",
                static_cast<int>(http_req.method.size()), http_req.method.data(),
@@ -272,9 +326,6 @@
 
   if (http_req.method == "OPTIONS") {
     HandleCorsPreflightRequest(http_req);
-  } else if (!http_req.origin.empty() && !IsOriginAllowed(http_req.origin)) {
-    // CSRF defense: reject cross-origin requests from outside the allowlist.
-    conn->SendResponseAndClose("403 Forbidden", {}, "Origin not allowed");
   } else {
     req_handler_->OnHttpRequest(http_req);
   }
@@ -286,7 +337,7 @@
   // Allow chaining multiple responses in the same HTTP-Keepalive connection.
   conn->headers_sent_ = false;
 
-  return headers_size + body_size;
+  return headers_size;
 }
 
 void HttpServer::HandleCorsPreflightRequest(const HttpRequest& req) {
@@ -447,9 +498,39 @@
   memcpy(mask, rd, sizeof(mask));
   rd += sizeof(mask);
 
-  if (avail() < payload_len)
-    return 0;  // Not enough data to read the payload.
-  uint8_t* const payload_start = rd;
+  // Control frames are capped at 125 bytes by RFC 6455 and are answered below
+  // rather than by the handler, so they stay in |rxbuf|.
+  const bool is_data_frame = opcode == kOpcodeBinary || opcode == kOpcodeText ||
+                             opcode == kOpcodeContinuation;
+  const size_t hdr_size = static_cast<size_t>(rd - rxbuf);
+  uint8_t* payload_start;
+  if (is_data_frame && payload_len > 0) {
+    if (conn->payload_ == nullptr) {
+      uint8_t* dst = req_handler_->OnWebsocketPayload(conn, payload_len);
+      if (dst == nullptr) {
+        PERFETTO_ELOG("[HTTP] Websocket payload rejected by the handler");
+        conn->Close();
+        return 0;
+      }
+      conn->payload_ = dst;
+      conn->payload_size_ = payload_len;
+      conn->payload_used_ = std::min(avail(), payload_len);
+      memcpy(dst, rd, conn->payload_used_);
+      memcpy(conn->payload_mask_, mask, sizeof(mask));
+      memmove(&rxbuf[hdr_size], &rxbuf[hdr_size + conn->payload_used_],
+              conn->rxbuf_used - hdr_size - conn->payload_used_);
+      conn->rxbuf_used -= conn->payload_used_;
+    }
+    if (conn->payload_used_ < conn->payload_size_)
+      return 0;
+    payload_start = conn->payload_;
+    memcpy(mask, conn->payload_mask_, sizeof(mask));
+    conn->ClearPayload();
+  } else {
+    if (avail() < payload_len)
+      return 0;  // Not enough data to read the payload.
+    payload_start = rd;
+  }
 
   // Unmask the payload, one 4-byte mask period per iteration.
   // Deliberately NOT written as the more natural `payload[i] ^= mask[i % 4]`:
@@ -490,7 +571,8 @@
   } else {
     PERFETTO_LOG("Unsupported WebSocket opcode: %d", opcode);
   }
-  return static_cast<size_t>(rd - rxbuf) + payload_len;
+  // A data frame's payload went to the handler, leaving only the header.
+  return hdr_size + (is_data_frame && payload_len > 0 ? 0 : payload_len);
 }
 
 void HttpServerConnection::SendResponseHeaders(
@@ -607,7 +689,7 @@
 }
 
 HttpServerConnection::HttpServerConnection(std::unique_ptr<UnixSocket> s)
-    : sock(std::move(s)), rxbuf(PagedMemory::Allocate(kMaxRequestSize)) {}
+    : sock(std::move(s)), rxbuf(PagedMemory::Allocate(kMaxHeadersSize)) {}
 
 HttpServerConnection::~HttpServerConnection() = default;
 

src/base/http/http_server_unittest.cc

diff --git a/src/base/http/http_server_unittest.cc b/src/base/http/http_server_unittest.cc
index 54a97bc..9b2340c 100644
--- a/src/base/http/http_server_unittest.cc
+++ b/src/base/http/http_server_unittest.cc
@@ -36,6 +36,19 @@
 
 class MockHttpHandler : public HttpRequestHandler {
  public:
+  // The tests drive one connection at a time, so one buffer is enough.
+  uint8_t* OnHttpRequestBody(const HttpRequest&, size_t size) override {
+    return Alloc(size);
+  }
+  uint8_t* OnWebsocketPayload(HttpServerConnection*, size_t size) override {
+    return Alloc(size);
+  }
+  uint8_t* Alloc(size_t size) {
+    payload_.reset(new uint8_t[size]);
+    return payload_.get();
+  }
+  std::unique_ptr<uint8_t[]> payload_;
+
   MOCK_METHOD(void, OnHttpRequest, (const HttpRequest&), (override));
   MOCK_METHOD(void,
               OnHttpConnectionClosed,

src/trace_processor/rpc/httpd.cc

diff --git a/src/trace_processor/rpc/httpd.cc b/src/trace_processor/rpc/httpd.cc
index c582889..dc56308 100644
--- a/src/trace_processor/rpc/httpd.cc
+++ b/src/trace_processor/rpc/httpd.cc
@@ -16,7 +16,6 @@
 
 #include <cstddef>
 #include <cstdint>
-#include <cstring>
 #include <initializer_list>
 #include <memory>
 #include <optional>
@@ -74,21 +73,32 @@
 
  private:
   // HttpRequestHandler implementation.
+  uint8_t* OnHttpRequestBody(const base::HttpRequest&, size_t size) override;
+  uint8_t* OnWebsocketPayload(base::HttpServerConnection*,
+                              size_t size) override;
   void OnHttpRequest(const base::HttpRequest&) override;
   void OnWebsocketMessage(const base::WebsocketMessage&) override;
   void OnHttpConnectionClosed(base::HttpServerConnection*) override;
 
-  // The RPC byte-pipe carried by |conn|, created on first use: most
-  // connections (the REST endpoints, /status polls) never carry one.
+  // The server can be part-way through a payload on several connections at
+  // once, so none of this can be shared.
+  struct ConnState {
+    // Created on first use: most connections never carry an RPC byte-pipe.
+    std::unique_ptr<Rpc::Stream> stream;
+    // Open from OnHttpRequestBody()/OnWebsocketPayload() until the matching
+    // OnHttpRequest()/OnWebsocketMessage(), spanning several socket reads.
+    Rpc::Stream::RequestHandle pending;
+    // Bodies of the non-RPC endpoints.
+    std::unique_ptr<uint8_t[]> payload;
+    size_t payload_size = 0;
+  };
+
   Rpc::Stream& GetRpcStream(base::HttpServerConnection* conn);
 
   static void ServeHelpPage(const base::HttpRequest&);
 
   Rpc& global_trace_processor_rpc_;
-  base::FlatHashMap<base::HttpServerConnection*,
-                    std::unique_ptr<Rpc::Stream>,
-                    ConnHasher>
-      streams_;
+  base::FlatHashMap<base::HttpServerConnection*, ConnState, ConnHasher> conns_;
   base::MaybeLockFreeTaskRunner task_runner_;
   base::HttpServer http_srv_;
   std::unique_ptr<IdleReaper> reaper_;
@@ -147,6 +157,28 @@
   task_runner_.Run();
 }
 
+uint8_t* Httpd::OnHttpRequestBody(const base::HttpRequest& req, size_t size) {
+  ConnState& state = conns_[req.conn];
+  if (req.uri == "/rpc") {
+    state.pending = GetRpcStream(req.conn).BeginRequest(size);
+    return state.pending.data();
+  }
+  if (size > state.payload_size) {
+    // Deliberately not value-initialized: |size| comes from Content-Length, so
+    // touching it here would make an unsent body resident.
+    state.payload.reset(new uint8_t[size]);
+    state.payload_size = size;
+  }
+  return state.payload.get();
+}
+
+uint8_t* Httpd::OnWebsocketPayload(base::HttpServerConnection* conn,
+                                   size_t size) {
+  ConnState& state = conns_[conn];
+  state.pending = GetRpcStream(conn).BeginRequest(size);
+  return state.pending.data();
+}
+
 void Httpd::OnHttpRequest(const base::HttpRequest& req) {
   if (reaper_)
     reaper_->OnActivity();
@@ -202,9 +234,7 @@
     conn.SendResponseHeaders("200 OK", chunked_headers,
                              base::HttpServerConnection::kOmitContentLength);
     if (!req.body.empty()) {
-      auto write = GetRpcStream(&conn).BeginRequest(req.body.size());
-      memcpy(write.data(), req.body.data(), req.body.size());
-      write.EndRequest(req.body.size());
+      conns_[&conn].pending.EndRequest(req.body.size());
     }
 
     // Terminate chunked stream.
@@ -332,13 +362,11 @@
     reaper_->OnActivity();
   if (msg.data.empty())
     return;
-  auto write = GetRpcStream(msg.conn).BeginRequest(msg.data.size());
-  memcpy(write.data(), msg.data.data(), msg.data.size());
-  write.EndRequest(msg.data.size());
+  conns_[msg.conn].pending.EndRequest(msg.data.size());
 }
 
 Rpc::Stream& Httpd::GetRpcStream(base::HttpServerConnection* conn) {
-  auto& stream = streams_[conn];
+  auto& stream = conns_[conn].stream;
   if (!stream) {
     stream = std::make_unique<Rpc::Stream>(
         global_trace_processor_rpc_, [conn](const void* data, uint32_t len) {
@@ -349,7 +377,12 @@
 }
 
 void Httpd::OnHttpConnectionClosed(base::HttpServerConnection* conn) {
-  streams_.Erase(conn);
+  // A peer that goes away mid-payload leaves a reservation open; the bytes it
+  // managed to send are an incomplete message and are dropped with it.
+  ConnState* state = conns_.Find(conn);
+  if (state && state->pending)
+    state->pending.AbortRequest();
+  conns_.Erase(conn);
 }
 
 }  // namespace

src/websocket_bridge/websocket_bridge.cc

diff --git a/src/websocket_bridge/websocket_bridge.cc b/src/websocket_bridge/websocket_bridge.cc
index 962b3b5..e1360bc 100644
--- a/src/websocket_bridge/websocket_bridge.cc
+++ b/src/websocket_bridge/websocket_bridge.cc
@@ -48,6 +48,25 @@
   void Main(int argc, char** argv);
 
   // base::HttpRequestHandler implementation.
+  // None of the bridge's HTTP endpoints take a body.
+  uint8_t* OnHttpRequestBody(const base::HttpRequest&, size_t) override {
+    return nullptr;
+  }
+  // A payload can span several reads, and payloads on different connections
+  // can be in flight at once, so it is buffered with the connection.
+  uint8_t* OnWebsocketPayload(base::HttpServerConnection* conn,
+                              size_t size) override {
+    auto it = conns_.find(conn);
+    PERFETTO_CHECK(it != conns_.end());
+    Conn& c = it->second;
+    if (size > c.payload_size) {
+      // Deliberately not value-initialized: |size| comes off the wire, so
+      // touching it here would make an unsent payload resident.
+      c.payload.reset(new uint8_t[size]);
+      c.payload_size = size;
+    }
+    return c.payload.get();
+  }
   void OnHttpRequest(const base::HttpRequest&) override;
   void OnWebsocketMessage(const base::WebsocketMessage&) override;
   void OnHttpConnectionClosed(base::HttpServerConnection*) override;
@@ -62,10 +81,15 @@
  private:
   base::HttpServerConnection* GetWebsocket(base::UnixSocket*);
 
+  struct Conn {
+    std::unique_ptr<base::UnixSocket> sock;
+    std::unique_ptr<uint8_t[]> payload;
+    size_t payload_size = 0;
+  };
+
   base::MaybeLockFreeTaskRunner task_runner_;
   std::vector<Endpoint> endpoints_;
-  std::map<base::HttpServerConnection*, std::unique_ptr<base::UnixSocket>>
-      conns_;
+  std::map<base::HttpServerConnection*, Conn> conns_;
 };
 
 void PrintUsage(char** argv) {
@@ -200,7 +224,7 @@
     sock_raw.SetBlocking(false);
 
     PERFETTO_DLOG("[WSBridge] Connected to %s", ep.endpoint);
-    conns_[req.conn] = base::UnixSocket::AdoptConnected(
+    conns_[req.conn].sock = base::UnixSocket::AdoptConnected(
         sock_raw.ReleaseFd(), this, &task_runner_, ep.family,
         base::SockType::kStream);
 
@@ -215,7 +239,7 @@
   auto it = conns_.find(msg.conn);
   PERFETTO_CHECK(it != conns_.end());
   // Pass through the websocket message onto the endpoint TCP socket.
-  base::UnixSocket& sock = *it->second;
+  base::UnixSocket& sock = *it->second.sock;
   sock.Send(msg.data.data(), msg.data.size());
 }
 
@@ -241,7 +265,7 @@
   auto it = conns_.find(websocket);
   if (it == conns_.end())
     return;  // Can happen if ADB closed first.
-  base::UnixSocket& sock = *it->second;
+  base::UnixSocket& sock = *it->second.sock;
   sock.Shutdown(/*notify=*/true);
   conns_.erase(websocket);
 }
@@ -258,7 +282,7 @@
 
 base::HttpServerConnection* WSBridge::GetWebsocket(base::UnixSocket* sock) {
   for (const auto& it : conns_) {
-    if (it.second.get() == sock) {
+    if (it.second.sock.get() == sock) {
       return it.first;
     }
   }