1 // Copyright (C) 2021 The Android Open Source Project
2 //
3 // Licensed under the Apache License, Version 2.0 (the "License");
4 // you may not use this file except in compliance with the License.
5 // You may obtain a copy of the License at
6 //
7 //      http://www.apache.org/licenses/LICENSE-2.0
8 //
9 // Unless required by applicable law or agreed to in writing, software
10 // distributed under the License is distributed on an "AS IS" BASIS,
11 // WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 // See the License for the specific language governing permissions and
13 // limitations under the License.
14 
15 #pragma once
16 
17 #include <linux/types.h>
18 #include <stdint.h>
19 #include <stdlib.h>
20 
21 #include <iostream>
22 
23 #include <libsnapshot/cow_reader.h>
24 
25 namespace android {
26 namespace snapshot {
27 
28 class BufferSink final {
29   public:
30     void Initialize(size_t size);
GetBufPtr()31     void* GetBufPtr() { return buffer_.get(); }
Clear()32     void Clear() { memset(GetBufPtr(), 0, buffer_size_); }
33     void* GetPayloadBuffer(size_t size);
34     void* GetBuffer(size_t requested, size_t* actual);
UpdateBufferOffset(size_t size)35     void UpdateBufferOffset(size_t size) { buffer_offset_ += size; }
36     struct dm_user_header* GetHeaderPtr();
ResetBufferOffset()37     void ResetBufferOffset() { buffer_offset_ = 0; }
38     void* GetPayloadBufPtr();
GetPayloadBytesWritten()39     loff_t GetPayloadBytesWritten() { return buffer_offset_; }
40 
41     // Same as calling GetPayloadBuffer and then UpdateBufferOffset.
42     //
43     // This is preferred over GetPayloadBuffer as it does not require a
44     // separate call to UpdateBufferOffset.
AcquireBuffer(size_t size)45     void* AcquireBuffer(size_t size) { return AcquireBuffer(size, size); }
46 
47     // Same as AcquireBuffer, but separates the requested size from the buffer
48     // offset. This is useful for a situation where a full run of data will be
49     // read, but only a partial amount will be returned.
50     //
51     // If size != to_write, the excess bytes may be reallocated by the next
52     // call to AcquireBuffer.
53     void* AcquireBuffer(size_t size, size_t to_write);
54 
55   private:
56     std::unique_ptr<uint8_t[]> buffer_;
57     loff_t buffer_offset_;
58     size_t buffer_size_;
59 };
60 
61 }  // namespace snapshot
62 }  // namespace android
63