commit cc04de757c9087dab79b7c1fa48ad19f01fcdb82
author Kelvin Zhang <zhangkelvin@google.com> Tue Oct 05 14:43:02 2021 -0700
committer Kelvin Zhang <zhangkelvin@google.com> Thu Oct 14 13:38:44 2021 -0700
tree 5cf3a5865121ec14a48daa9e9984bf51e351e580
parent 1304fe7c95545b0f012b50929d77504a967d1f90

Make XzEtentWriter::Init() idempotent

xz_dec_init allocates memory, so avoid memory leak in calling Init()
    multiple times

Test: th
Change-Id: I2dc20f6f9127b6d7749369385f982a08c78e6e44

payload_consumer/xz_extent_writer.cc

diff --git a/payload_consumer/xz_extent_writer.cc b/payload_consumer/xz_extent_writer.cc
index a648351..361e635 100644
--- a/payload_consumer/xz_extent_writer.cc
+++ b/payload_consumer/xz_extent_writer.cc
@@ -53,13 +53,13 @@
 }  // namespace
 
 XzExtentWriter::~XzExtentWriter() {
-  xz_dec_end(stream_);
+  stream_.reset();
   TEST_AND_RETURN(input_buffer_.empty());
 }
 
 bool XzExtentWriter::Init(const RepeatedPtrField<Extent>& extents,
                           uint32_t block_size) {
-  stream_ = xz_dec_init(XZ_DYNALLOC, kXzMaxDictSize);
+  stream_.reset(xz_dec_init(XZ_DYNALLOC, kXzMaxDictSize));
   TEST_AND_RETURN_FALSE(stream_ != nullptr);
   return underlying_writer_->Init(extents, block_size);
 }
@@ -86,7 +86,7 @@
   for (;;) {
     request.out_pos = 0;
 
-    xz_ret ret = xz_dec_run(stream_, &request);
+    xz_ret ret = xz_dec_run(stream_.get(), &request);
     if (ret != XZ_OK && ret != XZ_STREAM_END) {
       LOG(ERROR) << "xz_dec_run returned " << XzErrorString(ret);
       return false;

payload_consumer/xz_extent_writer.h

diff --git a/payload_consumer/xz_extent_writer.h b/payload_consumer/xz_extent_writer.h
index 70338f2..caf34ec 100644
--- a/payload_consumer/xz_extent_writer.h
+++ b/payload_consumer/xz_extent_writer.h
@@ -34,6 +34,10 @@
 namespace chromeos_update_engine {
 
 class XzExtentWriter : public ExtentWriter {
+  struct xz_deleter {
+    constexpr void operator()(xz_dec* p) { xz_dec_end(p); }
+  };
+
  public:
   explicit XzExtentWriter(std::unique_ptr<ExtentWriter> underlying_writer)
       : underlying_writer_(std::move(underlying_writer)) {}
@@ -47,7 +51,7 @@
   // The underlying ExtentWriter.
   std::unique_ptr<ExtentWriter> underlying_writer_;
   // The opaque xz decompressor struct.
-  xz_dec* stream_{nullptr};
+  std::unique_ptr<xz_dec, xz_deleter> stream_{nullptr};
   brillo::Blob input_buffer_;
 
   DISALLOW_COPY_AND_ASSIGN(XzExtentWriter);