From 3805c7b9be052bc52deb79d4950fec4915610083 Mon Sep 17 00:00:00 2001
From: Brian Ting <brianting@google.com>
Date: Mon, 3 Jun 2024 16:02:57 -0700
Subject: [PATCH] Increase WebDriverServer max buffer size.

Resolves timeout issues when screenshots exceed the default 1mb buffer
size and are unable to be sent.

b/344007302

Change-Id: I8b7d7972fe0dbdf449b1c6ddc02ae50e397ec6a2
---
 cobalt/webdriver/server.cc | 5 ++++-
 cobalt/webdriver/server.h  | 3 +++
 2 files changed, 7 insertions(+), 1 deletion(-)

diff --git a/cobalt/webdriver/server.cc b/cobalt/webdriver/server.cc
index 51e5709c89bd..b1bb2071b89b 100644
--- a/cobalt/webdriver/server.cc
+++ b/cobalt/webdriver/server.cc
@@ -76,7 +76,10 @@ class ResponseHandlerImpl : public WebDriverServer::ResponseHandler {
   ResponseHandlerImpl(net::HttpServer* server, int connection_id)
       : task_runner_(base::SequencedTaskRunner::GetCurrentDefault()),
         server_(server),
-        connection_id_(connection_id) {}
+        connection_id_(connection_id) {
+    server_->SetReceiveBufferSize(connection_id_, kMaxRecieveBufferSize);
+    server_->SetSendBufferSize(connection_id_, kMaxSendBufferSize);
+  }
 
   // https://www.selenium.dev/documentation/legacy/json_wire_protocol/#responses
   void Success(std::unique_ptr<base::Value> value) override {
diff --git a/cobalt/webdriver/server.h b/cobalt/webdriver/server.h
index 7d850aa3e6a7..6575a1ce882f 100644
--- a/cobalt/webdriver/server.h
+++ b/cobalt/webdriver/server.h
@@ -54,6 +54,9 @@ class WebDriverServer : public net::HttpServer::Delegate {
   // client.
   class ResponseHandler {
    public:
+    static const int kMaxRecieveBufferSize = 2 * 1024 * 1024;  // 2 Mbytes.
+    static const int kMaxSendBufferSize = 8 * 1024 * 1024;     // 8 Mbytes.
+
     // Called after a successful WebDriver command.
     // https://www.selenium.dev/documentation/legacy/json_wire_protocol/#responses
     virtual void Success(std::unique_ptr<base::Value>) = 0;