diff --git a/payload_consumer/block_extent_writer.cc b/payload_consumer/block_extent_writer.cc
index 055b485..456c491 100644
--- a/payload_consumer/block_extent_writer.cc
+++ b/payload_consumer/block_extent_writer.cc
@@ -16,9 +16,13 @@
 
 #include "update_engine/payload_consumer/block_extent_writer.h"
 
-#include <algorithm>
-#include <cstdint>
+#include <stdint.h>
 
+#include <algorithm>
+
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_ranges.h"
 #include "update_engine/update_metadata.pb.h"
 
 namespace chromeos_update_engine {
@@ -35,56 +39,65 @@
   return true;
 }
 
-size_t BlockExtentWriter::ConsumeWithBuffer(const uint8_t* data, size_t count) {
-  CHECK_LT(cur_extent_idx_, static_cast<size_t>(extents_.size()));
+bool BlockExtentWriter::WriteExtent(const void* bytes, const size_t count) {
   const auto& cur_extent = extents_[cur_extent_idx_];
-  const auto cur_extent_size = cur_extent.num_blocks() * block_size_;
+  const auto write_extent =
+      ExtentForRange(cur_extent.start_block() + offset_in_extent_ / block_size_,
+                     count / kBlockSize);
+  offset_in_extent_ += count;
+  if (offset_in_extent_ == cur_extent.num_blocks() * block_size_) {
+    NextExtent();
+  }
+  return WriteExtent(bytes, write_extent, block_size_);
+}
 
-  if (buffer_.empty() && count >= cur_extent_size) {
-    if (!WriteExtent(data, cur_extent, block_size_)) {
+size_t BlockExtentWriter::ConsumeWithBuffer(const uint8_t* const data,
+                                            const size_t count) {
+  if (cur_extent_idx_ >= static_cast<size_t>(extents_.size())) {
+    if (count > 0) {
+      LOG(ERROR) << "Exhausted all blocks, but still have " << count
+                 << " bytes pending for write";
+    }
+    return 0;
+  }
+  const auto& cur_extent = extents_[cur_extent_idx_];
+  const auto cur_extent_size =
+      static_cast<size_t>(cur_extent.num_blocks() * block_size_);
+
+  const auto write_size =
+      std::min(cur_extent_size - offset_in_extent_, BUFFER_SIZE);
+  if (buffer_.empty() && count >= write_size) {
+    if (!WriteExtent(data, write_size)) {
       LOG(ERROR) << "WriteExtent(" << cur_extent.start_block() << ", "
-                 << static_cast<const void*>(data) << ", " << cur_extent_size
+                 << static_cast<const void*>(data) << ", " << write_size
                  << ") failed.";
       // return value is expected to be greater than 0. Return 0 to signal error
       // condition
       return 0;
     }
-    if (!NextExtent()) {
-      if (count != cur_extent_size) {
-        LOG(ERROR) << "Exhausted all blocks, but still have "
-                   << count - cur_extent_size << " bytes left";
-        return 0;
-      }
-    }
-    return cur_extent_size;
+    return write_size;
   }
-  if (buffer_.size() >= cur_extent_size) {
+  if (buffer_.size() >= write_size) {
     LOG(ERROR)
-        << "Data left in buffer should never be >= cur_extent_size, otherwise "
+        << "Data left in buffer should never be >= write_size, otherwise "
            "we should have send that data to CowWriter. Buffer size: "
-        << buffer_.size() << " current extent size: " << cur_extent_size;
+        << buffer_.size() << " write_size: " << write_size;
   }
   const size_t bytes_to_copy =
-      std::min<size_t>(count, cur_extent_size - buffer_.size());
+      std::min<size_t>(count, write_size - buffer_.size());
   TEST_GT(bytes_to_copy, 0U);
 
   buffer_.insert(buffer_.end(), data, data + bytes_to_copy);
-  TEST_LE(buffer_.size(), cur_extent_size);
+  TEST_LE(buffer_.size(), write_size);
 
-  if (buffer_.size() == cur_extent_size) {
-    if (!WriteExtent(buffer_.data(), cur_extent, block_size_)) {
+  if (buffer_.size() == write_size) {
+    if (!WriteExtent(buffer_.data(), write_size)) {
       LOG(ERROR) << "WriteExtent(" << buffer_.data() << ", "
                  << cur_extent.start_block() << ", " << cur_extent.num_blocks()
                  << ") failed.";
       return 0;
     }
     buffer_.clear();
-    if (!NextExtent()) {
-      if (count != bytes_to_copy) {
-        LOG(ERROR) << "Exhausted all blocks, but still have "
-                   << count - bytes_to_copy << " bytes left";
-      }
-    }
   }
   return bytes_to_copy;
 }
@@ -111,6 +124,7 @@
 
 bool BlockExtentWriter::NextExtent() {
   cur_extent_idx_++;
+  offset_in_extent_ = 0;
   return cur_extent_idx_ < static_cast<size_t>(extents_.size());
 }
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/block_extent_writer.h b/payload_consumer/block_extent_writer.h
index 902e3e1..eeae36d 100644
--- a/payload_consumer/block_extent_writer.h
+++ b/payload_consumer/block_extent_writer.h
@@ -28,6 +28,7 @@
 // Cache data upto size of one extent before writing.
 class BlockExtentWriter : public chromeos_update_engine::ExtentWriter {
  public:
+  static constexpr size_t BUFFER_SIZE = 1024 * 1024;
   BlockExtentWriter() = default;
   ~BlockExtentWriter() = default;
   // Returns true on success.
@@ -44,15 +45,18 @@
   size_t BlockSize() const { return block_size_; }
 
  private:
+  bool WriteExtent(const void* bytes, size_t count);
   bool NextExtent();
-  [[nodiscard]] size_t ConsumeWithBuffer(const uint8_t* bytes, size_t count);
+  [[nodiscard]] size_t ConsumeWithBuffer(const uint8_t* const bytes,
+                                         const size_t count);
   // It's a non-owning pointer, because PartitionWriter owns the CowWruter. This
   // allows us to use a single instance of CowWriter for all operations applied
   // to the same partition.
   google::protobuf::RepeatedPtrField<Extent> extents_;
-  size_t cur_extent_idx_;
+  size_t cur_extent_idx_{};
   std::vector<uint8_t> buffer_;
-  size_t block_size_;
+  size_t block_size_{};
+  size_t offset_in_extent_{};
 };
 
 }  // namespace chromeos_update_engine
diff --git a/payload_consumer/block_extent_writer_unittest.cc b/payload_consumer/block_extent_writer_unittest.cc
new file mode 100644
index 0000000..558e46b
--- /dev/null
+++ b/payload_consumer/block_extent_writer_unittest.cc
@@ -0,0 +1,184 @@
+//
+// Copyright (C) 2023 The Android Open Source Project
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+
+#include "update_engine/payload_consumer/block_extent_writer.h"
+
+#include <fcntl.h>
+
+#include <algorithm>
+#include <memory>
+#include <string>
+#include <vector>
+
+#include <gtest/gtest.h>
+#include <gmock/gmock.h>
+
+#include "update_engine/common/test_utils.h"
+#include "update_engine/payload_generator/delta_diff_generator.h"
+#include "update_engine/payload_generator/extent_utils.h"
+#include "update_engine/common/utils.h"
+#include "update_engine/payload_generator/extent_ranges.h"
+
+using std::min;
+using std::string;
+using std::vector;
+using testing::_;
+using testing::Return;
+
+namespace chromeos_update_engine {
+
+class BlockExtentWriterTest : public ::testing::Test {
+ protected:
+  void SetUp() override {}
+  void TearDown() override {}
+};
+
+class MockBlockExtentWriter : public BlockExtentWriter {
+ public:
+  MOCK_METHOD(bool,
+              WriteExtent,
+              (const void*, const Extent&, size_t),
+              (override));
+};
+
+TEST_F(BlockExtentWriterTest, LongExtentTest) {
+  google::protobuf::RepeatedPtrField<Extent> extents;
+  *extents.Add() = ExtentForRange(0, 1);
+  *extents.Add() = ExtentForRange(2, 1);
+  *extents.Add() = ExtentForRange(4, 1);
+  // A single large extent which doesn't fit in 1 buffer
+  static constexpr auto BLOCKS_PER_BUFFER =
+      BlockExtentWriter::BUFFER_SIZE / kBlockSize;
+  *extents.Add() = ExtentForRange(10, BLOCKS_PER_BUFFER * 2);
+  MockBlockExtentWriter writer;
+  ASSERT_TRUE(writer.Init(extents, kBlockSize));
+  std::string buffer;
+  buffer.resize(BlockExtentWriter::BUFFER_SIZE * 2);
+  ON_CALL(writer, WriteExtent(_, _, _)).WillByDefault(Return(true));
+  EXPECT_CALL(writer,
+              WriteExtent(buffer.data(), ExtentForRange(0, 1), kBlockSize));
+  EXPECT_CALL(writer,
+              WriteExtent(static_cast<void*>(buffer.data() + kBlockSize),
+                          ExtentForRange(2, 1),
+                          kBlockSize));
+  EXPECT_CALL(writer,
+              WriteExtent(static_cast<void*>(buffer.data() + kBlockSize * 2),
+                          ExtentForRange(4, 1),
+                          kBlockSize));
+  // The last chunk should be split up into multiple chunks, each chunk is 1
+  // BUFFR_SIZE
+  EXPECT_CALL(writer,
+              WriteExtent(static_cast<void*>(buffer.data()),
+                          ExtentForRange(10, BLOCKS_PER_BUFFER),
+                          kBlockSize));
+  EXPECT_CALL(
+      writer,
+      WriteExtent(
+          static_cast<void*>(buffer.data() + BlockExtentWriter::BUFFER_SIZE),
+          ExtentForRange(10 + BLOCKS_PER_BUFFER, BLOCKS_PER_BUFFER),
+          kBlockSize));
+  ASSERT_TRUE(writer.Write(buffer.data(), kBlockSize * 3));
+  ASSERT_TRUE(writer.Write(buffer.data(), BlockExtentWriter::BUFFER_SIZE * 2));
+}
+
+TEST_F(BlockExtentWriterTest, LongExtentMultiCall) {
+  google::protobuf::RepeatedPtrField<Extent> extents;
+  static constexpr auto BLOCKS_PER_BUFFER =
+      BlockExtentWriter::BUFFER_SIZE / kBlockSize;
+  *extents.Add() = ExtentForRange(10, BLOCKS_PER_BUFFER * 5);
+  MockBlockExtentWriter writer;
+  ASSERT_TRUE(writer.Init(extents, kBlockSize));
+  std::string buffer;
+  buffer.resize(BlockExtentWriter::BUFFER_SIZE * 2);
+  ON_CALL(writer, WriteExtent(_, _, _)).WillByDefault(Return(true));
+  // The last chunk should be split up into multiple chunks, each chunk is 1
+  // BUFFR_SIZE
+  EXPECT_CALL(writer,
+              WriteExtent(static_cast<void*>(buffer.data()),
+                          ExtentForRange(10, BLOCKS_PER_BUFFER),
+                          kBlockSize));
+  EXPECT_CALL(
+      writer,
+      WriteExtent(static_cast<void*>(buffer.data()),
+                  ExtentForRange(10 + BLOCKS_PER_BUFFER, BLOCKS_PER_BUFFER),
+                  kBlockSize));
+  EXPECT_CALL(
+      writer,
+      WriteExtent(static_cast<void*>(buffer.data()),
+                  ExtentForRange(10 + BLOCKS_PER_BUFFER * 2, BLOCKS_PER_BUFFER),
+                  kBlockSize));
+  ASSERT_TRUE(writer.Write(buffer.data(), BlockExtentWriter::BUFFER_SIZE));
+  ASSERT_TRUE(writer.Write(buffer.data(), BlockExtentWriter::BUFFER_SIZE));
+  ASSERT_TRUE(writer.Write(buffer.data(), BlockExtentWriter::BUFFER_SIZE));
+}
+
+void FillArbitraryData(std::string* buffer) {
+  for (size_t i = 0; i < buffer->size(); i++) {
+    (*buffer)[i] = i;
+  }
+}
+
+TEST_F(BlockExtentWriterTest, SingleBufferMultiCall) {
+  google::protobuf::RepeatedPtrField<Extent> extents;
+  static constexpr auto BLOCKS_PER_BUFFER =
+      BlockExtentWriter::BUFFER_SIZE / kBlockSize;
+  *extents.Add() = ExtentForRange(10, BLOCKS_PER_BUFFER);
+  MockBlockExtentWriter writer;
+  ASSERT_TRUE(writer.Init(extents, kBlockSize));
+  std::string buffer;
+  buffer.resize(BlockExtentWriter::BUFFER_SIZE);
+  FillArbitraryData(&buffer);
+
+  ON_CALL(writer, WriteExtent(_, _, _)).WillByDefault(Return(true));
+  EXPECT_CALL(writer,
+              WriteExtent(_, ExtentForRange(10, BLOCKS_PER_BUFFER), kBlockSize))
+      .WillOnce([&buffer](const void* data, const Extent& extent, size_t) {
+        return memcmp(data, buffer.data(), buffer.size()) == 0;
+      });
+
+  ASSERT_TRUE(
+      writer.Write(buffer.data(), BlockExtentWriter::BUFFER_SIZE - kBlockSize));
+  ASSERT_TRUE(writer.Write(
+      buffer.data() + BlockExtentWriter::BUFFER_SIZE - kBlockSize, kBlockSize));
+}
+
+TEST_F(BlockExtentWriterTest, MultiBufferMultiCall) {
+  google::protobuf::RepeatedPtrField<Extent> extents;
+  static constexpr auto BLOCKS_PER_BUFFER =
+      BlockExtentWriter::BUFFER_SIZE / kBlockSize;
+  *extents.Add() = ExtentForRange(10, BLOCKS_PER_BUFFER + 1);
+  MockBlockExtentWriter writer;
+  ASSERT_TRUE(writer.Init(extents, kBlockSize));
+  std::string buffer;
+  buffer.resize(BlockExtentWriter::BUFFER_SIZE);
+  FillArbitraryData(&buffer);
+
+  ON_CALL(writer, WriteExtent(_, _, _)).WillByDefault(Return(true));
+  EXPECT_CALL(writer,
+              WriteExtent(_, ExtentForRange(10, BLOCKS_PER_BUFFER), kBlockSize))
+      .WillOnce([&buffer](const void* data, const Extent& extent, size_t) {
+        return memcmp(data, buffer.data(), extent.num_blocks() * kBlockSize) ==
+               0;
+      });
+  EXPECT_CALL(
+      writer,
+      WriteExtent(_, ExtentForRange(10 + BLOCKS_PER_BUFFER, 1), kBlockSize));
+
+  ASSERT_TRUE(writer.Write(buffer.data(), BlockExtentWriter::BUFFER_SIZE));
+  ASSERT_TRUE(writer.Write(buffer.data(), kBlockSize));
+}
+
+}  // namespace chromeos_update_engine
diff --git a/payload_consumer/bzip_extent_writer.cc b/payload_consumer/bzip_extent_writer.cc
index 9491964..bee9fcd 100644
--- a/payload_consumer/bzip_extent_writer.cc
+++ b/payload_consumer/bzip_extent_writer.cc
@@ -14,6 +14,7 @@
 // limitations under the License.
 //
 
+#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/bzip_extent_writer.h"
 
 using google::protobuf::RepeatedPtrField;
diff --git a/payload_consumer/bzip_extent_writer.h b/payload_consumer/bzip_extent_writer.h
index 38c041a..4926411 100644
--- a/payload_consumer/bzip_extent_writer.h
+++ b/payload_consumer/bzip_extent_writer.h
@@ -23,7 +23,6 @@
 
 #include <brillo/secure_blob.h>
 
-#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/extent_writer.h"
 
 // BzipExtentWriter is a concrete ExtentWriter subclass that bzip-decompresses
@@ -46,7 +45,7 @@
 
  private:
   std::unique_ptr<ExtentWriter> next_;  // The underlying ExtentWriter.
-  bz_stream stream_;                    // the libbz2 stream
+  bz_stream stream_{};                  // the libbz2 stream
   brillo::Blob input_buffer_;
 };
 
diff --git a/payload_consumer/extent_writer.h b/payload_consumer/extent_writer.h
index 8b1b532..5a4e3ee 100644
--- a/payload_consumer/extent_writer.h
+++ b/payload_consumer/extent_writer.h
@@ -23,7 +23,6 @@
 #include <base/logging.h>
 #include <brillo/secure_blob.h>
 
-#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/file_descriptor.h"
 #include "update_engine/update_metadata.pb.h"
 
diff --git a/payload_consumer/xz_extent_writer.cc b/payload_consumer/xz_extent_writer.cc
index 361e635..5c92d05 100644
--- a/payload_consumer/xz_extent_writer.cc
+++ b/payload_consumer/xz_extent_writer.cc
@@ -14,6 +14,7 @@
 // limitations under the License.
 //
 
+#include "update_engine/common/utils.h"
 #include "update_engine/payload_consumer/xz_extent_writer.h"
 
 using google::protobuf::RepeatedPtrField;
@@ -75,7 +76,7 @@
     count = input_buffer_.size();
   }
 
-  xz_buf request;
+  xz_buf request{};
   request.in = input;
   request.in_pos = 0;
   request.in_size = count;