snapuserd: Refactor buffer allocation code

Since there is no change in the buffer allocation
code with respect to snapshots in user-space, refactor
it so that we have a common code base.

Bug: 193863397
Test: Build
Signed-off-by: Akilesh Kailash <akailash@google.com>
Change-Id: I3a7b3d319f6d1532b1bb63af9884b64e1346bb50
diff --git a/fs_mgr/libsnapshot/snapuserd/Android.bp b/fs_mgr/libsnapshot/snapuserd/Android.bp
index 5865507..837f33a 100644
--- a/fs_mgr/libsnapshot/snapuserd/Android.bp
+++ b/fs_mgr/libsnapshot/snapuserd/Android.bp
@@ -61,6 +61,7 @@
         "dm-snapshot-merge/snapuserd_worker.cpp",
         "dm-snapshot-merge/snapuserd_readahead.cpp",
         "snapuserd_daemon.cpp",
+	"snapuserd_buffer.cpp",
 	"user-space-merge/snapuserd_core.cpp",
 	"user-space-merge/snapuserd_dm_user.cpp",
 	"user-space-merge/snapuserd_merge.cpp",
@@ -109,6 +110,7 @@
         "dm-snapshot-merge/cow_snapuserd_test.cpp",
         "dm-snapshot-merge/snapuserd.cpp",
         "dm-snapshot-merge/snapuserd_worker.cpp",
+	"snapuserd_buffer.cpp",
     ],
     cflags: [
         "-Wall",
diff --git a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/cow_snapuserd_test.cpp b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/cow_snapuserd_test.cpp
index bff0a50..b86a802 100644
--- a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/cow_snapuserd_test.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/cow_snapuserd_test.cpp
@@ -33,6 +33,7 @@
 #include <libdm/dm.h>
 #include <libdm/loop_control.h>
 #include <libsnapshot/cow_writer.h>
+#include <snapuserd/snapuserd_buffer.h>
 #include <snapuserd/snapuserd_client.h>
 #include <storage_literals/storage_literals.h>
 
diff --git a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.h b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.h
index 6388a83..47b9b22 100644
--- a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.h
+++ b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd.h
@@ -42,6 +42,7 @@
 #include <libdm/dm.h>
 #include <libsnapshot/cow_reader.h>
 #include <libsnapshot/cow_writer.h>
+#include <snapuserd/snapuserd_buffer.h>
 #include <snapuserd/snapuserd_kernel.h>
 
 namespace android {
@@ -89,39 +90,6 @@
     READ_AHEAD_FAILURE,
 };
 
-class BufferSink : public IByteSink {
-  public:
-    void Initialize(size_t size);
-    void* GetBufPtr() { return buffer_.get(); }
-    void Clear() { memset(GetBufPtr(), 0, buffer_size_); }
-    void* GetPayloadBuffer(size_t size);
-    void* GetBuffer(size_t requested, size_t* actual) override;
-    void UpdateBufferOffset(size_t size) { buffer_offset_ += size; }
-    struct dm_user_header* GetHeaderPtr();
-    bool ReturnData(void*, size_t) override { return true; }
-    void ResetBufferOffset() { buffer_offset_ = 0; }
-    void* GetPayloadBufPtr();
-
-  private:
-    std::unique_ptr<uint8_t[]> buffer_;
-    loff_t buffer_offset_;
-    size_t buffer_size_;
-};
-
-class XorSink : public IByteSink {
-  public:
-    void Initialize(BufferSink* sink, size_t size);
-    void Reset();
-    void* GetBuffer(size_t requested, size_t* actual) override;
-    bool ReturnData(void* buffer, size_t len) override;
-
-  private:
-    BufferSink* bufsink_;
-    std::unique_ptr<uint8_t[]> buffer_;
-    size_t buffer_size_;
-    size_t returned_;
-};
-
 class Snapuserd;
 
 class ReadAheadThread {
diff --git a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp
index 5d184ad..0e9f0f1 100644
--- a/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/dm-snapshot-merge/snapuserd_worker.cpp
@@ -32,78 +32,6 @@
 #define SNAP_LOG(level) LOG(level) << misc_name_ << ": "
 #define SNAP_PLOG(level) PLOG(level) << misc_name_ << ": "
 
-void BufferSink::Initialize(size_t size) {
-    buffer_size_ = size;
-    buffer_offset_ = 0;
-    buffer_ = std::make_unique<uint8_t[]>(size);
-}
-
-void* BufferSink::GetPayloadBuffer(size_t size) {
-    if ((buffer_size_ - buffer_offset_) < size) return nullptr;
-
-    char* buffer = reinterpret_cast<char*>(GetBufPtr());
-    struct dm_user_message* msg = (struct dm_user_message*)(&(buffer[0]));
-    return (char*)msg->payload.buf + buffer_offset_;
-}
-
-void* BufferSink::GetBuffer(size_t requested, size_t* actual) {
-    void* buf = GetPayloadBuffer(requested);
-    if (!buf) {
-        *actual = 0;
-        return nullptr;
-    }
-    *actual = requested;
-    return buf;
-}
-
-struct dm_user_header* BufferSink::GetHeaderPtr() {
-    if (!(sizeof(struct dm_user_header) <= buffer_size_)) {
-        return nullptr;
-    }
-    char* buf = reinterpret_cast<char*>(GetBufPtr());
-    struct dm_user_header* header = (struct dm_user_header*)(&(buf[0]));
-    return header;
-}
-
-void* BufferSink::GetPayloadBufPtr() {
-    char* buffer = reinterpret_cast<char*>(GetBufPtr());
-    struct dm_user_message* msg = reinterpret_cast<struct dm_user_message*>(&(buffer[0]));
-    return msg->payload.buf;
-}
-
-void XorSink::Initialize(BufferSink* sink, size_t size) {
-    bufsink_ = sink;
-    buffer_size_ = size;
-    returned_ = 0;
-    buffer_ = std::make_unique<uint8_t[]>(size);
-}
-
-void XorSink::Reset() {
-    returned_ = 0;
-}
-
-void* XorSink::GetBuffer(size_t requested, size_t* actual) {
-    if (requested > buffer_size_) {
-        *actual = buffer_size_;
-    } else {
-        *actual = requested;
-    }
-    return buffer_.get();
-}
-
-bool XorSink::ReturnData(void* buffer, size_t len) {
-    uint8_t* xor_data = reinterpret_cast<uint8_t*>(buffer);
-    uint8_t* buff = reinterpret_cast<uint8_t*>(bufsink_->GetPayloadBuffer(len + returned_));
-    if (buff == nullptr) {
-        return false;
-    }
-    for (size_t i = 0; i < len; i++) {
-        buff[returned_ + i] ^= xor_data[i];
-    }
-    returned_ += len;
-    return true;
-}
-
 WorkerThread::WorkerThread(const std::string& cow_device, const std::string& backing_device,
                            const std::string& control_device, const std::string& misc_name,
                            std::shared_ptr<Snapuserd> snapuserd) {
diff --git a/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_buffer.h b/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_buffer.h
new file mode 100644
index 0000000..2e4cac6
--- /dev/null
+++ b/fs_mgr/libsnapshot/snapuserd/include/snapuserd/snapuserd_buffer.h
@@ -0,0 +1,62 @@
+// Copyright (C) 2021 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.
+
+#pragma once
+
+#include <linux/types.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+#include <iostream>
+
+#include <libsnapshot/cow_reader.h>
+
+namespace android {
+namespace snapshot {
+
+class BufferSink : public IByteSink {
+  public:
+    void Initialize(size_t size);
+    void* GetBufPtr() { return buffer_.get(); }
+    void Clear() { memset(GetBufPtr(), 0, buffer_size_); }
+    void* GetPayloadBuffer(size_t size);
+    void* GetBuffer(size_t requested, size_t* actual) override;
+    void UpdateBufferOffset(size_t size) { buffer_offset_ += size; }
+    struct dm_user_header* GetHeaderPtr();
+    bool ReturnData(void*, size_t) override { return true; }
+    void ResetBufferOffset() { buffer_offset_ = 0; }
+    void* GetPayloadBufPtr();
+
+  private:
+    std::unique_ptr<uint8_t[]> buffer_;
+    loff_t buffer_offset_;
+    size_t buffer_size_;
+};
+
+class XorSink : public IByteSink {
+  public:
+    void Initialize(BufferSink* sink, size_t size);
+    void Reset();
+    void* GetBuffer(size_t requested, size_t* actual) override;
+    bool ReturnData(void* buffer, size_t len) override;
+
+  private:
+    BufferSink* bufsink_;
+    std::unique_ptr<uint8_t[]> buffer_;
+    size_t buffer_size_;
+    size_t returned_;
+};
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/snapuserd/snapuserd_buffer.cpp b/fs_mgr/libsnapshot/snapuserd/snapuserd_buffer.cpp
new file mode 100644
index 0000000..ab763ab
--- /dev/null
+++ b/fs_mgr/libsnapshot/snapuserd/snapuserd_buffer.cpp
@@ -0,0 +1,96 @@
+/*
+ * Copyright (C) 2021 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 <snapuserd/snapuserd_buffer.h>
+#include <snapuserd/snapuserd_kernel.h>
+
+namespace android {
+namespace snapshot {
+
+void BufferSink::Initialize(size_t size) {
+    buffer_size_ = size;
+    buffer_offset_ = 0;
+    buffer_ = std::make_unique<uint8_t[]>(size);
+}
+
+void* BufferSink::GetPayloadBuffer(size_t size) {
+    if ((buffer_size_ - buffer_offset_) < size) return nullptr;
+
+    char* buffer = reinterpret_cast<char*>(GetBufPtr());
+    struct dm_user_message* msg = (struct dm_user_message*)(&(buffer[0]));
+    return (char*)msg->payload.buf + buffer_offset_;
+}
+
+void* BufferSink::GetBuffer(size_t requested, size_t* actual) {
+    void* buf = GetPayloadBuffer(requested);
+    if (!buf) {
+        *actual = 0;
+        return nullptr;
+    }
+    *actual = requested;
+    return buf;
+}
+
+struct dm_user_header* BufferSink::GetHeaderPtr() {
+    if (!(sizeof(struct dm_user_header) <= buffer_size_)) {
+        return nullptr;
+    }
+    char* buf = reinterpret_cast<char*>(GetBufPtr());
+    struct dm_user_header* header = (struct dm_user_header*)(&(buf[0]));
+    return header;
+}
+
+void* BufferSink::GetPayloadBufPtr() {
+    char* buffer = reinterpret_cast<char*>(GetBufPtr());
+    struct dm_user_message* msg = reinterpret_cast<struct dm_user_message*>(&(buffer[0]));
+    return msg->payload.buf;
+}
+
+void XorSink::Initialize(BufferSink* sink, size_t size) {
+    bufsink_ = sink;
+    buffer_size_ = size;
+    returned_ = 0;
+    buffer_ = std::make_unique<uint8_t[]>(size);
+}
+
+void XorSink::Reset() {
+    returned_ = 0;
+}
+
+void* XorSink::GetBuffer(size_t requested, size_t* actual) {
+    if (requested > buffer_size_) {
+        *actual = buffer_size_;
+    } else {
+        *actual = requested;
+    }
+    return buffer_.get();
+}
+
+bool XorSink::ReturnData(void* buffer, size_t len) {
+    uint8_t* xor_data = reinterpret_cast<uint8_t*>(buffer);
+    uint8_t* buff = reinterpret_cast<uint8_t*>(bufsink_->GetPayloadBuffer(len + returned_));
+    if (buff == nullptr) {
+        return false;
+    }
+    for (size_t i = 0; i < len; i++) {
+        buff[returned_ + i] ^= xor_data[i];
+    }
+    returned_ += len;
+    return true;
+}
+
+}  // namespace snapshot
+}  // namespace android
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
index e37be7b..c171eda 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_core.h
@@ -39,6 +39,7 @@
 #include <libdm/dm.h>
 #include <libsnapshot/cow_reader.h>
 #include <libsnapshot/cow_writer.h>
+#include <snapuserd/snapuserd_buffer.h>
 #include <snapuserd/snapuserd_kernel.h>
 
 namespace android {
@@ -64,39 +65,6 @@
     READ_AHEAD_FAILURE,
 };
 
-class Bufsink : public IByteSink {
-  public:
-    void Initialize(size_t size);
-    void* GetBufPtr() { return buffer_.get(); }
-    void Clear() { memset(GetBufPtr(), 0, buffer_size_); }
-    void* GetPayloadBuffer(size_t size);
-    void* GetBuffer(size_t requested, size_t* actual) override;
-    void UpdateBufferOffset(size_t size) { buffer_offset_ += size; }
-    struct dm_user_header* GetHeaderPtr();
-    bool ReturnData(void*, size_t) override { return true; }
-    void ResetBufferOffset() { buffer_offset_ = 0; }
-    void* GetPayloadBufPtr();
-
-  private:
-    std::unique_ptr<uint8_t[]> buffer_;
-    loff_t buffer_offset_;
-    size_t buffer_size_;
-};
-
-class XorBufSink : public IByteSink {
-  public:
-    void Initialize(Bufsink* sink, size_t size);
-    void Reset();
-    void* GetBuffer(size_t requested, size_t* actual) override;
-    bool ReturnData(void* buffer, size_t len) override;
-
-  private:
-    Bufsink* bufsink_;
-    std::unique_ptr<uint8_t[]> buffer_;
-    size_t buffer_size_;
-    size_t returned_;
-};
-
 class SnapshotHandler;
 
 class ReadAhead {
@@ -142,7 +110,7 @@
     std::unordered_set<uint64_t> dest_blocks_;
     std::unordered_set<uint64_t> source_blocks_;
     bool overlap_;
-    Bufsink bufsink_;
+    BufferSink bufsink_;
 };
 
 class Worker {
@@ -185,8 +153,8 @@
                      std::vector<const CowOperation*>* replace_zero_vec = nullptr);
 
     std::unique_ptr<CowReader> reader_;
-    Bufsink bufsink_;
-    XorBufSink xorsink_;
+    BufferSink bufsink_;
+    XorSink xorsink_;
 
     std::string cow_device_;
     std::string backing_store_device_;
diff --git a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
index 8036c81..18c7f2c 100644
--- a/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
+++ b/fs_mgr/libsnapshot/snapuserd/user-space-merge/snapuserd_dm_user.cpp
@@ -23,78 +23,6 @@
 using namespace android::dm;
 using android::base::unique_fd;
 
-void Bufsink::Initialize(size_t size) {
-    buffer_size_ = size;
-    buffer_offset_ = 0;
-    buffer_ = std::make_unique<uint8_t[]>(size);
-}
-
-void* Bufsink::GetPayloadBuffer(size_t size) {
-    if ((buffer_size_ - buffer_offset_) < size) return nullptr;
-
-    char* buffer = reinterpret_cast<char*>(GetBufPtr());
-    struct dm_user_message* msg = (struct dm_user_message*)(&(buffer[0]));
-    return (char*)msg->payload.buf + buffer_offset_;
-}
-
-void* Bufsink::GetBuffer(size_t requested, size_t* actual) {
-    void* buf = GetPayloadBuffer(requested);
-    if (!buf) {
-        *actual = 0;
-        return nullptr;
-    }
-    *actual = requested;
-    return buf;
-}
-
-struct dm_user_header* Bufsink::GetHeaderPtr() {
-    if (!(sizeof(struct dm_user_header) <= buffer_size_)) {
-        return nullptr;
-    }
-    char* buf = reinterpret_cast<char*>(GetBufPtr());
-    struct dm_user_header* header = (struct dm_user_header*)(&(buf[0]));
-    return header;
-}
-
-void* Bufsink::GetPayloadBufPtr() {
-    char* buffer = reinterpret_cast<char*>(GetBufPtr());
-    struct dm_user_message* msg = reinterpret_cast<struct dm_user_message*>(&(buffer[0]));
-    return msg->payload.buf;
-}
-
-void XorBufSink::Initialize(Bufsink* sink, size_t size) {
-    bufsink_ = sink;
-    buffer_size_ = size;
-    returned_ = 0;
-    buffer_ = std::make_unique<uint8_t[]>(size);
-}
-
-void XorBufSink::Reset() {
-    returned_ = 0;
-}
-
-void* XorBufSink::GetBuffer(size_t requested, size_t* actual) {
-    if (requested > buffer_size_) {
-        *actual = buffer_size_;
-    } else {
-        *actual = requested;
-    }
-    return buffer_.get();
-}
-
-bool XorBufSink::ReturnData(void* buffer, size_t len) {
-    uint8_t* xor_data = reinterpret_cast<uint8_t*>(buffer);
-    uint8_t* buff = reinterpret_cast<uint8_t*>(bufsink_->GetPayloadBuffer(len + returned_));
-    if (buff == nullptr) {
-        return false;
-    }
-    for (size_t i = 0; i < len; i++) {
-        buff[returned_ + i] ^= xor_data[i];
-    }
-    returned_ += len;
-    return true;
-}
-
 Worker::Worker(const std::string& cow_device, const std::string& backing_device,
                const std::string& control_device, const std::string& misc_name,
                const std::string& base_path_merge, std::shared_ptr<SnapshotHandler> snapuserd) {