Revert "[dawn][wire] Handle command chunking in the wire in more robustly."

This reverts commit 11b115b26569c8aa46781b4dac589ce17f5d6c89.

Reason for revert: Fuzzer crash https://crbug.com/452575709

Original change's description:
> [dawn][wire] Handle command chunking in the wire in more robustly.
>
> - With the changes to make the wire more spontaneous, chunked
>   command serialization and handling need to be updated to work
>   in a multithreaded manner. Previously, the handler assumes that
>   once a chunked command is received, it will receive the entire
>   chunked command before it receives any other commands. This may
>   no longer be true, so we upgrade chunked commands into its own
>   command so with unique id's to ensure that we handle them
>   appropriately.
> - This change also unifies the command id enums in the wire to a
>   single enum since chunked command handling may occur on both
>   the server and the client (though in practice in Chromium,
>   only the server->client may be multithreaded for now.) This
>   allows us to handle commands in a unified way and also ensures
>   that the command ids are now unique both direction.
>
> Bug: 412761856
> Change-Id: Ia1eb34d57c2e2fe0457c70aae9348bc051fce573
> Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/266475
> Reviewed-by: Corentin Wallez <cwallez@chromium.org>
> Reviewed-by: Kai Ninomiya <kainino@chromium.org>
> Commit-Queue: Loko Kung <lokokung@google.com>

TBR=cwallez@chromium.org,kainino@chromium.org,dawn-scoped@luci-project-accounts.iam.gserviceaccount.com,lokokung@google.com

Fixed: 452575709
Bug: 412761856
Change-Id: I02915daf7e8d35729e3afdb9425ce6a51679958f
Reviewed-on: https://dawn-review.googlesource.com/c/dawn/+/267254
Bot-Commit: rubber-stamper@appspot.gserviceaccount.com <rubber-stamper@appspot.gserviceaccount.com>
Commit-Queue: rubber-stamper@appspot.gserviceaccount.com <rubber-stamper@appspot.gserviceaccount.com>
diff --git a/generator/dawn_json_generator.py b/generator/dawn_json_generator.py
index dd55164..4bd6757 100644
--- a/generator/dawn_json_generator.py
+++ b/generator/dawn_json_generator.py
@@ -701,7 +701,6 @@
 
     commands = []
     return_commands = []
-    special_commands = []
 
     wire_json['special items']['client_handwritten_commands'] += wire_json[
         'special items']['client_side_commands']
@@ -767,14 +766,9 @@
         return_commands.append(
             Command(name, linked_record_members(json_data, types)))
 
-    for (name, json_data) in wire_json['special commands'].items():
-        special_commands.append(
-            Command(name, linked_record_members(json_data, types)))
-
     wire_params['cmd_records'] = {
         'command': commands,
-        'return command': return_commands,
-        'special command': special_commands
+        'return command': return_commands
     }
 
     for commands in wire_params['cmd_records'].values():
diff --git a/generator/templates/dawn/wire/WireCmd.cpp b/generator/templates/dawn/wire/WireCmd.cpp
index 9714952..f12db7f 100644
--- a/generator/templates/dawn/wire/WireCmd.cpp
+++ b/generator/templates/dawn/wire/WireCmd.cpp
@@ -136,7 +136,7 @@
                       "Record must be at most one of is_cmd, extensible, and chained.");
         {% if is_cmd %}
             //* Start the transfer structure with the command ID, so that casting to WireCmd gives the ID.
-            WireCmd commandId;
+            {{Return}}WireCmd commandId;
         {% elif record.extensible %}
             WGPUBool hasNextInChain;
         {% elif record.chained %}
@@ -242,7 +242,7 @@
     ) {
         //* Handle special transfer members of methods.
         {% if is_cmd %}
-            transfer->commandId = WireCmd::{{Return}}{{name}};
+            transfer->commandId = {{Return}}WireCmd::{{name}};
         {% endif %}
 
         {% if record.extensible %}
@@ -348,7 +348,7 @@
             , const ObjectIdResolver& resolver
         {%- endif -%}) {
         {% if is_cmd %}
-            DAWN_ASSERT(transfer->commandId == WireCmd::{{Return}}{{name}});
+            DAWN_ASSERT(transfer->commandId == {{Return}}WireCmd::{{name}});
         {% endif %}
         {% if record.derived_method %}
             record->selfId = transfer->self;
@@ -690,12 +690,6 @@
     {% do sTypes.append(sType) %}
 {% endfor %}
 
-//* Output [de]serialization helpers for special commands
-{% for command in cmd_records["special command"] %}
-    {%- set name = command.name.CamelCase() -%}
-    {{write_record_serialization_helpers(command, name, command.members, is_cmd=True)}}
-{% endfor %}
-
 //* Output [de]serialization helpers for commands
 {% for command in cmd_records["command"] %}
     {%- set name = command.name.CamelCase() -%}
@@ -741,10 +735,6 @@
 
 }  // anonymous namespace
 
-{% for command in cmd_records["special command"] -%}
-    {{write_command_serialization_methods(command, False)}}
-{% endfor %}
-
 {% for command in cmd_records["command"] -%}
     {{write_command_serialization_methods(command, False)}}
 {% endfor %}
diff --git a/generator/templates/dawn/wire/WireCmd.h b/generator/templates/dawn/wire/WireCmd.h
index bd42e2a..ba1542a 100644
--- a/generator/templates/dawn/wire/WireCmd.h
+++ b/generator/templates/dawn/wire/WireCmd.h
@@ -63,20 +63,17 @@
             {% endfor %}
     };
 
+    //* Enum used as a prefix to each command on the wire format.
     enum class WireCmd : uint32_t {
-        // Enums used in special wire commands.
-        {% for command in cmd_records["special command"] %}
-            {{command.name.CamelCase()}},
-        {% endfor %}
-
-        // Enums used on the wire format.
         {% for command in cmd_records["command"] %}
             {{command.name.CamelCase()}},
         {% endfor %}
+    };
 
-        // Enums used on the return wire format.
+    //* Enum used as a prefix to each command on the return wire format.
+    enum class ReturnWireCmd : uint32_t {
         {% for command in cmd_records["return command"] %}
-            Return{{command.name.CamelCase()}},
+            {{command.name.CamelCase()}},
         {% endfor %}
     };
 
@@ -120,10 +117,6 @@
     };
 {% endmacro %}
 
-    {% for command in cmd_records["special command"] %}
-        {{write_command_struct(command, False)}}
-    {% endfor %}
-
     {% for command in cmd_records["command"] %}
         {{write_command_struct(command, False)}}
     {% endfor %}
diff --git a/generator/templates/dawn/wire/client/ClientHandlers.cpp b/generator/templates/dawn/wire/client/ClientHandlers.cpp
index 5c741fc..dc32d84 100644
--- a/generator/templates/dawn/wire/client/ClientHandlers.cpp
+++ b/generator/templates/dawn/wire/client/ClientHandlers.cpp
@@ -34,7 +34,7 @@
     {% for command in cmd_records["return command"] %}
         WireResult Client::Handle{{command.name.CamelCase()}}(DeserializeBuffer* deserializeBuffer) {
             Return{{command.name.CamelCase()}}Cmd cmd;
-            WIRE_TRY(cmd.Deserialize(deserializeBuffer, &mAllocator));
+            WIRE_TRY(cmd.Deserialize(deserializeBuffer, &mWireCommandAllocator));
 
             {% for member in command.members if member.handle_type %}
                 {% set Type = member.handle_type.name.CamelCase() %}
@@ -61,34 +61,37 @@
         }
     {% endfor %}
 
-    const volatile char* Client::HandleCommands(const volatile char* commands, size_t size) {
+    const volatile char* Client::HandleCommandsImpl(const volatile char* commands, size_t size) {
         DeserializeBuffer deserializeBuffer(commands, size);
 
-        while (deserializeBuffer.AvailableSize() >= sizeof(CmdHeader) + sizeof(WireCmd)) {
-            WireCmd cmdId = *static_cast<const volatile WireCmd*>(static_cast<const volatile void*>(
+        while (deserializeBuffer.AvailableSize() >= sizeof(CmdHeader) + sizeof(ReturnWireCmd)) {
+            // Start by chunked command handling, if it is done, then it means the whole buffer
+            // was consumed by it, so we return a pointer to the end of the commands.
+            switch (HandleChunkedCommands(deserializeBuffer.Buffer(), deserializeBuffer.AvailableSize())) {
+                case ChunkedCommandsResult::Consumed:
+                    return commands + size;
+                case ChunkedCommandsResult::Error:
+                    return nullptr;
+                case ChunkedCommandsResult::Passthrough:
+                    break;
+            }
+
+            ReturnWireCmd cmdId = *static_cast<const volatile ReturnWireCmd*>(static_cast<const volatile void*>(
                 deserializeBuffer.Buffer() + sizeof(CmdHeader)));
             WireResult result = WireResult::FatalError;
             switch (cmdId) {
-                {% for command in cmd_records["special command"] %}
-                    {% set Suffix = command.name.CamelCase() %}
-                    case WireCmd::{{Suffix}}:
-                        result = Handle{{Suffix}}(&deserializeBuffer);
-                        break;
-                {% endfor %}
                 {% for command in cmd_records["return command"] %}
                     {% set Suffix = command.name.CamelCase() %}
-                    case WireCmd::Return{{Suffix}}:
+                    case ReturnWireCmd::{{Suffix}}:
                         result = Handle{{Suffix}}(&deserializeBuffer);
                         break;
                 {% endfor %}
-                default:
-                    result = WireResult::FatalError;
             }
 
             if (result != WireResult::Success) {
                 return nullptr;
             }
-            mAllocator.Reset();
+            mWireCommandAllocator.Reset();
         }
 
         if (deserializeBuffer.AvailableSize() != 0) {
diff --git a/generator/templates/dawn/wire/client/ClientPrototypes.inc b/generator/templates/dawn/wire/client/ClientPrototypes.inc
index 576a0a3..12907e1 100644
--- a/generator/templates/dawn/wire/client/ClientPrototypes.inc
+++ b/generator/templates/dawn/wire/client/ClientPrototypes.inc
@@ -25,11 +25,14 @@
 //* OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 //* OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-// Command handlers & doers
+//* Return command handlers
 {% for command in cmd_records["return command"] %}
-    {% set Suffix = command.name.CamelCase() %}
-    WireResult Handle{{Suffix}}(DeserializeBuffer* deserializeBuffer);
-    WireResult Do{{Suffix}}(
+    WireResult Handle{{command.name.CamelCase()}}(DeserializeBuffer* deserializeBuffer);
+{% endfor %}
+
+//* Return command doers
+{% for command in cmd_records["return command"] %}
+    WireResult Do{{command.name.CamelCase()}}(
         {%- for member in command.members -%}
             {%- if member.handle_type -%}
                 {{as_wireType(member.handle_type)}} {{as_varName(member.name)}}
@@ -39,5 +42,4 @@
             {%- if not loop.last -%}, {% endif %}
         {%- endfor -%}
     );
-
 {% endfor %}
diff --git a/generator/templates/dawn/wire/server/ServerHandlers.cpp b/generator/templates/dawn/wire/server/ServerHandlers.cpp
index 57876a2..c954823 100644
--- a/generator/templates/dawn/wire/server/ServerHandlers.cpp
+++ b/generator/templates/dawn/wire/server/ServerHandlers.cpp
@@ -87,15 +87,26 @@
         }
     {% endfor %}
 
-    const volatile char* Server::HandleCommands(const volatile char* commands, size_t size) {
+    const volatile char* Server::HandleCommandsImpl(const volatile char* commands, size_t size) {
         DeserializeBuffer deserializeBuffer(commands, size);
 
         while (deserializeBuffer.AvailableSize() >= sizeof(CmdHeader) + sizeof(WireCmd)) {
+            // Start by chunked command handling, if it is done, then it means the whole buffer
+            // was consumed by it, so we return a pointer to the end of the commands.
+            switch (HandleChunkedCommands(deserializeBuffer.Buffer(), deserializeBuffer.AvailableSize())) {
+                case ChunkedCommandsResult::Consumed:
+                    return commands + size;
+                case ChunkedCommandsResult::Error:
+                    return nullptr;
+                case ChunkedCommandsResult::Passthrough:
+                    break;
+            }
+
             WireCmd cmdId = *static_cast<const volatile WireCmd*>(static_cast<const volatile void*>(
                 deserializeBuffer.Buffer() + sizeof(CmdHeader)));
             WireResult result;
             switch (cmdId) {
-                {% for command in cmd_records["special command"] + cmd_records["command"] %}
+                {% for command in cmd_records["command"] %}
                     case WireCmd::{{command.name.CamelCase()}}:
                         result = Handle{{command.name.CamelCase()}}(&deserializeBuffer);
                         break;
diff --git a/generator/templates/dawn/wire/server/ServerPrototypes.inc b/generator/templates/dawn/wire/server/ServerPrototypes.inc
index 7cdc7c3..f7404df 100644
--- a/generator/templates/dawn/wire/server/ServerPrototypes.inc
+++ b/generator/templates/dawn/wire/server/ServerPrototypes.inc
@@ -29,6 +29,7 @@
 {% for command in cmd_records["command"] %}
     {% set Suffix = command.name.CamelCase() %}
     WireResult Handle{{Suffix}}(DeserializeBuffer* deserializeBuffer);
+
     WireResult Do{{Suffix}}(
         {%- for member in command.members -%}
             {%- if member.is_return_value -%}
@@ -45,7 +46,6 @@
             {%- if not loop.last -%}, {% endif %}
         {%- endfor -%}
     );
-
 {% endfor %}
 
 {% for CommandName in server_custom_pre_handler_commands %}
diff --git a/src/dawn/dawn_wire.json b/src/dawn/dawn_wire.json
index 0fd2c01..ad848fd 100644
--- a/src/dawn/dawn_wire.json
+++ b/src/dawn/dawn_wire.json
@@ -221,14 +221,6 @@
             { "name": "features", "type": "feature name", "annotation": "const*", "length": "features count"}
         ]
     },
-    "special commands": {
-        "chunked command": [
-            {"name": "id", "type": "uint64_t"},
-            {"name": "size", "type": "uint64_t"},
-            {"name": "chunk data", "type": "char", "annotation": "const*", "length": "chunk size", "wire_is_data_only": true},
-            {"name": "chunk size", "type": "uint32_t"}
-        ]
-    },
     "special items": {
         "client_side_structures": [
             "FutureWaitInfo",
diff --git a/src/dawn/tests/unittests/wire/WireQueueTests.cpp b/src/dawn/tests/unittests/wire/WireQueueTests.cpp
index aab12c6..1751551 100644
--- a/src/dawn/tests/unittests/wire/WireQueueTests.cpp
+++ b/src/dawn/tests/unittests/wire/WireQueueTests.cpp
@@ -25,7 +25,6 @@
 // OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 // OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 
-#include <cstring>
 #include <memory>
 
 #include "dawn/common/StringViewUtils.h"
@@ -46,39 +45,6 @@
 using testing::Return;
 using testing::SizedString;
 
-class WireWriteBufferTests : public WireTest {};
-
-// Tests that commands are serialized properly when they are too large and need to be chunked.
-TEST_F(WireWriteBufferTests, WriteBufferChunkedCommands) {
-    WGPUBuffer apiBuffer = api.GetNewBuffer();
-
-    // In order to create a command larger than the the maximum that is allowed to be serialized at
-    // a time for a single command to force command chunking, use a value larger than the maximum
-    // allocation size.
-    static size_t kLargeAllocationSize = GetC2SMaxAllocationSize() + 16u;
-
-    wgpu::BufferDescriptor desc = {};
-    desc.usage = wgpu::BufferUsage::CopyDst | wgpu::BufferUsage::MapRead;
-    desc.size = kLargeAllocationSize;
-    wgpu::Buffer buffer = device.CreateBuffer(&desc);
-    EXPECT_CALL(api, DeviceCreateBuffer(apiDevice, _))
-        .WillOnce(Return(apiBuffer))
-        .RetiresOnSaturation();
-    FlushClient();
-
-    auto expected = std::make_unique<uint8_t[]>(kLargeAllocationSize);
-    std::memset(expected.get(), 0b10101010, kLargeAllocationSize);
-    queue.WriteBuffer(buffer, 0, expected.get(), kLargeAllocationSize);
-
-    EXPECT_CALL(
-        api, QueueWriteBuffer(apiQueue, apiBuffer, 0, MatchesLambda([&](void const* actual) {
-                                  return !std::memcmp(expected.get(), actual, kLargeAllocationSize);
-                              }),
-                              kLargeAllocationSize))
-        .Times(1);
-    FlushClient();
-}
-
 using WireQueueTestBase = WireFutureTest<wgpu::QueueWorkDoneCallback<void>*>;
 class WireQueueTests : public WireQueueTestBase {
   protected:
diff --git a/src/dawn/tests/unittests/wire/WireTest.cpp b/src/dawn/tests/unittests/wire/WireTest.cpp
index 41ef354..3d60204 100644
--- a/src/dawn/tests/unittests/wire/WireTest.cpp
+++ b/src/dawn/tests/unittests/wire/WireTest.cpp
@@ -260,10 +260,6 @@
     return mWireClient.get();
 }
 
-size_t WireTest::GetC2SMaxAllocationSize() {
-    return mC2sBuf->GetMaximumAllocationSize();
-}
-
 void WireTest::DeleteServer() {
     EXPECT_CALL(api, QueueRelease(apiQueue)).Times(1);
     EXPECT_CALL(api, DeviceRelease(apiDevice)).Times(1);
diff --git a/src/dawn/tests/unittests/wire/WireTest.h b/src/dawn/tests/unittests/wire/WireTest.h
index ab62ea6..41fdbac 100644
--- a/src/dawn/tests/unittests/wire/WireTest.h
+++ b/src/dawn/tests/unittests/wire/WireTest.h
@@ -180,8 +180,6 @@
     dawn::wire::WireServer* GetWireServer();
     dawn::wire::WireClient* GetWireClient();
 
-    size_t GetC2SMaxAllocationSize();
-
     void DeleteServer();
     void DeleteClient();
 
diff --git a/src/dawn/wire/ChunkedCommandHandler.cpp b/src/dawn/wire/ChunkedCommandHandler.cpp
index 5f23cb1..aa3de33 100644
--- a/src/dawn/wire/ChunkedCommandHandler.cpp
+++ b/src/dawn/wire/ChunkedCommandHandler.cpp
@@ -44,46 +44,56 @@
 
 ChunkedCommandHandler::~ChunkedCommandHandler() = default;
 
-WireResult ChunkedCommandHandler::HandleChunkedCommand(DeserializeBuffer* deserializeBuffer) {
-    ChunkedCommandCmd cmd;
-    WIRE_TRY(cmd.Deserialize(deserializeBuffer, &mAllocator));
+const volatile char* ChunkedCommandHandler::HandleCommands(const volatile char* commands,
+                                                           size_t size) {
+    if (mChunkedCommandRemainingSize > 0) {
+        // If there is a chunked command in flight, append the command data.
+        // We append at most |mChunkedCommandRemainingSize| which is enough to finish the
+        // in-flight chunked command, and then pass the rest along to a second call to
+        // |HandleCommandsImpl|.
+        size_t chunkSize = std::min(size, mChunkedCommandRemainingSize);
 
-    ChunkedCommand* chunkedCommand = nullptr;
-    if (auto it = mChunkedCommands.find(cmd.id); it != mChunkedCommands.end()) {
-        chunkedCommand = &(it->second);
-    } else {
-        ChunkedCommand newChunkedCommand = {};
-        newChunkedCommand.remainingSize = cmd.size;
-        newChunkedCommand.data.reset(AllocNoThrow<char>(cmd.size));
-        if (newChunkedCommand.data.get() == nullptr) {
-            return WireResult::FatalError;
-        }
+        memcpy(mChunkedCommandData.get() + mChunkedCommandPutOffset,
+               const_cast<const char*>(commands), chunkSize);
+        mChunkedCommandPutOffset += chunkSize;
+        mChunkedCommandRemainingSize -= chunkSize;
 
-        const auto& [newIt, inserted] =
-            mChunkedCommands.insert({cmd.id, std::move(newChunkedCommand)});
-        DAWN_ASSERT(inserted);
-        chunkedCommand = &(newIt->second);
-    }
-    DAWN_ASSERT(chunkedCommand);
+        commands += chunkSize;
+        size -= chunkSize;
 
-    DAWN_ASSERT(cmd.chunkSize <= chunkedCommand->remainingSize);
-    if (cmd.chunkSize > chunkedCommand->remainingSize) {
-        return WireResult::FatalError;
-    }
-    memcpy(chunkedCommand->data.get() + chunkedCommand->putOffset,
-           const_cast<const char*>(cmd.chunkData), cmd.chunkSize);
-    chunkedCommand->putOffset += cmd.chunkSize;
-    chunkedCommand->remainingSize -= cmd.chunkSize;
-
-    if (chunkedCommand->remainingSize == 0) {
-        ChunkedCommand fullCommand = std::move(*chunkedCommand);
-        mChunkedCommands.erase(cmd.id);
-        if (HandleCommands(fullCommand.data.get(), fullCommand.putOffset) == nullptr) {
-            return WireResult::FatalError;
+        if (mChunkedCommandRemainingSize == 0) {
+            // Once the chunked command is complete, pass the data to the command handler
+            // implemenation.
+            auto chunkedCommandData = std::move(mChunkedCommandData);
+            if (HandleCommandsImpl(chunkedCommandData.get(), mChunkedCommandPutOffset) == nullptr) {
+                // |HandleCommandsImpl| returns nullptr on error. Forward any errors
+                // out.
+                return nullptr;
+            }
         }
     }
 
-    return WireResult::Success;
+    return HandleCommandsImpl(commands, size);
+}
+
+ChunkedCommandHandler::ChunkedCommandsResult ChunkedCommandHandler::BeginChunkedCommandData(
+    const volatile char* commands,
+    size_t commandSize,
+    size_t initialSize) {
+    DAWN_ASSERT(!mChunkedCommandData);
+
+    // Reserve space for all the command data we're expecting, and copy the initial data
+    // to the start of the memory.
+    mChunkedCommandData.reset(AllocNoThrow<char>(commandSize));
+    if (!mChunkedCommandData) {
+        return ChunkedCommandsResult::Error;
+    }
+
+    memcpy(mChunkedCommandData.get(), const_cast<const char*>(commands), initialSize);
+    mChunkedCommandPutOffset = initialSize;
+    mChunkedCommandRemainingSize = commandSize - initialSize;
+
+    return ChunkedCommandsResult::Consumed;
 }
 
 }  // namespace dawn::wire
diff --git a/src/dawn/wire/ChunkedCommandHandler.h b/src/dawn/wire/ChunkedCommandHandler.h
index 050c813..d5ab08e 100644
--- a/src/dawn/wire/ChunkedCommandHandler.h
+++ b/src/dawn/wire/ChunkedCommandHandler.h
@@ -32,11 +32,9 @@
 #include <limits>
 #include <memory>
 
-#include "absl/container/flat_hash_map.h"
 #include "dawn/common/Assert.h"
 #include "dawn/wire/Wire.h"
 #include "dawn/wire/WireCmd_autogen.h"
-#include "dawn/wire/WireDeserializeAllocator.h"
 
 namespace dawn::wire {
 
@@ -45,21 +43,41 @@
     ChunkedCommandHandler();
     ~ChunkedCommandHandler() override;
 
-  protected:
-    WireResult HandleChunkedCommand(DeserializeBuffer* deserializeBuffer);
+    const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
 
-    WireDeserializeAllocator mAllocator;
+  protected:
+    enum class ChunkedCommandsResult {
+        Passthrough,
+        Consumed,
+        Error,
+    };
+
+    // Returns |true| if the commands were entirely consumed into the chunked command vector
+    // and should be handled later once we receive all the command data.
+    // Returns |false| if commands should be handled now immediately.
+    ChunkedCommandsResult HandleChunkedCommands(const volatile char* commands, size_t size) {
+        uint64_t commandSize64 = reinterpret_cast<const volatile CmdHeader*>(commands)->commandSize;
+
+        if (commandSize64 > std::numeric_limits<size_t>::max()) {
+            return ChunkedCommandsResult::Error;
+        }
+        size_t commandSize = static_cast<size_t>(commandSize64);
+        if (size < commandSize) {
+            return BeginChunkedCommandData(commands, commandSize, size);
+        }
+        return ChunkedCommandsResult::Passthrough;
+    }
 
   private:
-    // This map keeps track of all in-flight chunked commands. Note that because |HandleCommands|
-    // must be called in a thread-safe manner, we do not need to explicitly synchronize access to
-    // this map.
-    struct ChunkedCommand {
-        size_t remainingSize = 0;
-        size_t putOffset = 0;
-        std::unique_ptr<char[]> data = nullptr;
-    };
-    absl::flat_hash_map<uint64_t, ChunkedCommand> mChunkedCommands;
+    virtual const volatile char* HandleCommandsImpl(const volatile char* commands, size_t size) = 0;
+
+    ChunkedCommandsResult BeginChunkedCommandData(const volatile char* commands,
+                                                  size_t commandSize,
+                                                  size_t initialSize);
+
+    size_t mChunkedCommandRemainingSize = 0;
+    size_t mChunkedCommandPutOffset = 0;
+    std::unique_ptr<char[]> mChunkedCommandData;
 };
 
 }  // namespace dawn::wire
diff --git a/src/dawn/wire/ChunkedCommandSerializer.cpp b/src/dawn/wire/ChunkedCommandSerializer.cpp
index 51de9d5..f306495 100644
--- a/src/dawn/wire/ChunkedCommandSerializer.cpp
+++ b/src/dawn/wire/ChunkedCommandSerializer.cpp
@@ -32,8 +32,6 @@
 
 #include "dawn/wire/ChunkedCommandSerializer.h"
 
-#include "dawn/common/Assert.h"
-
 namespace dawn::wire {
 
 ChunkedCommandSerializer::ChunkedCommandSerializer(CommandSerializer* serializer)
@@ -41,35 +39,22 @@
     DAWN_ASSERT(mMaxAllocationSize > 0);
 }
 
-void ChunkedCommandSerializer::SetCommandSerializerForDisconnect(CommandSerializer* serializer) {
-    mSerializer = serializer;
-    mMaxAllocationSize = serializer->GetMaximumAllocationSize();
-    DAWN_ASSERT(mMaxAllocationSize > 0);
-}
-
 void ChunkedCommandSerializer::Flush() {
     mSerializer->Flush();
 }
 
 void ChunkedCommandSerializer::SerializeChunkedCommand(const char* allocatedBuffer,
-                                                       size_t totalSize) {
-    // Constant regarding the size of the WireChunkedCommandCmd that can be computed once.
-    static size_t kWireChunkedCmdPrefixSize = ChunkedCommandCmd{0, 0, nullptr, 0}.GetRequiredSize();
-
-    ChunkedCommandCmd cmd;
-    cmd.id = mNextChunkedCommandId++;
-    cmd.size = totalSize;
-
-    size_t remainingSize = totalSize;
+                                                       size_t remainingSize) {
     while (remainingSize > 0) {
-        cmd.chunkData = allocatedBuffer;
-        cmd.chunkSize = std::min(remainingSize, mMaxAllocationSize - kWireChunkedCmdPrefixSize);
-        DAWN_ASSERT(cmd.GetRequiredSize() <= mMaxAllocationSize);
+        size_t chunkSize = std::min(remainingSize, mMaxAllocationSize);
+        void* dst = mSerializer->GetCmdSpace(chunkSize);
+        if (dst == nullptr) {
+            return;
+        }
+        memcpy(dst, allocatedBuffer, chunkSize);
 
-        SerializeCommand(cmd);
-
-        allocatedBuffer += cmd.chunkSize;
-        remainingSize -= cmd.chunkSize;
+        allocatedBuffer += chunkSize;
+        remainingSize -= chunkSize;
     }
 }
 
diff --git a/src/dawn/wire/ChunkedCommandSerializer.h b/src/dawn/wire/ChunkedCommandSerializer.h
index bc15d50..6194315 100644
--- a/src/dawn/wire/ChunkedCommandSerializer.h
+++ b/src/dawn/wire/ChunkedCommandSerializer.h
@@ -29,7 +29,6 @@
 #define SRC_DAWN_WIRE_CHUNKEDCOMMANDSERIALIZER_H_
 
 #include <algorithm>
-#include <atomic>
 #include <cstring>
 #include <functional>
 #include <memory>
@@ -76,10 +75,6 @@
   public:
     explicit ChunkedCommandSerializer(CommandSerializer* serializer);
 
-    // This utility function is intended only for disconnect situations where we want the serializer
-    // to appear to keep working even though we are no longer serializing and flushing commands.
-    void SetCommandSerializerForDisconnect(CommandSerializer* serializer);
-
     template <typename Cmd>
     void SerializeCommand(const Cmd& cmd) {
         SerializeCommandImpl(
@@ -149,11 +144,10 @@
         SerializeChunkedCommand(cmdSpace.get(), requiredSize);
     }
 
-    void SerializeChunkedCommand(const char* allocatedBuffer, size_t totalSize);
+    void SerializeChunkedCommand(const char* allocatedBuffer, size_t remainingSize);
 
     raw_ptr<CommandSerializer> mSerializer;
     size_t mMaxAllocationSize;
-    std::atomic<uint64_t> mNextChunkedCommandId = 0;
 };
 
 }  // namespace dawn::wire
diff --git a/src/dawn/wire/client/Client.cpp b/src/dawn/wire/client/Client.cpp
index f3e12b1..6d7c7ad 100644
--- a/src/dawn/wire/client/Client.cpp
+++ b/src/dawn/wire/client/Client.cpp
@@ -174,7 +174,7 @@
 
 void Client::Disconnect() {
     mDisconnected = true;
-    mSerializer.SetCommandSerializerForDisconnect(NoopCommandSerializer::GetInstance());
+    mSerializer = ChunkedCommandSerializer(NoopCommandSerializer::GetInstance());
 
     // Transition all event managers to ClientDropped state.
     for (auto& [_, eventManager] : mEventManagers) {
diff --git a/src/dawn/wire/client/Client.h b/src/dawn/wire/client/Client.h
index 1406fa4..3732454 100644
--- a/src/dawn/wire/client/Client.h
+++ b/src/dawn/wire/client/Client.h
@@ -82,7 +82,7 @@
     }
 
     // ChunkedCommandHandler implementation
-    const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
+    const volatile char* HandleCommandsImpl(const volatile char* commands, size_t size) override;
 
     MemoryTransferService* GetMemoryTransferService() const { return mMemoryTransferService; }
 
@@ -143,6 +143,7 @@
 #include "dawn/wire/client/ClientPrototypes_autogen.inc"
 
     ChunkedCommandSerializer mSerializer;
+    WireDeserializeAllocator mWireCommandAllocator;
     PerObjectType<ObjectStore> mObjects;
     std::unique_ptr<MemoryTransferService> mOwnedMemoryTransferService = nullptr;
     raw_ptr<MemoryTransferService> mMemoryTransferService = nullptr;
diff --git a/src/dawn/wire/server/Server.h b/src/dawn/wire/server/Server.h
index 6ded567..a8896b4 100644
--- a/src/dawn/wire/server/Server.h
+++ b/src/dawn/wire/server/Server.h
@@ -173,7 +173,7 @@
     ~Server() override;
 
     // ChunkedCommandHandler implementation
-    const volatile char* HandleCommands(const volatile char* commands, size_t size) override;
+    const volatile char* HandleCommandsImpl(const volatile char* commands, size_t size) override;
 
     WireResult InjectBuffer(WGPUBuffer buffer, const Handle& handle, const Handle& deviceHandle);
     WireResult InjectTexture(WGPUTexture texture, const Handle& handle, const Handle& deviceHandle);
@@ -281,6 +281,7 @@
 
 #include "dawn/wire/server/ServerPrototypes_autogen.inc"
 
+    WireDeserializeAllocator mAllocator;
     MutexProtected<ChunkedCommandSerializer> mSerializer;
     std::unique_ptr<MemoryTransferService> mOwnedMemoryTransferService = nullptr;
     raw_ptr<MemoryTransferService> mMemoryTransferService = nullptr;