1 // Copyright 2020 Google LLC
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 ///////////////////////////////////////////////////////////////////////////////
16
17 #include "tink/cc/python_output_stream.h"
18
19 #include <algorithm>
20 #include <memory>
21
22 #include "absl/memory/memory.h"
23 #include "absl/strings/string_view.h"
24 #include "tink/output_stream.h"
25 #include "tink/subtle/subtle_util.h"
26 #include "tink/util/status.h"
27 #include "tink/util/statusor.h"
28 #include "tink/cc/python_file_object_adapter.h"
29
30 namespace crypto {
31 namespace tink {
32
PythonOutputStream(std::shared_ptr<PythonFileObjectAdapter> adapter,int buffer_size)33 PythonOutputStream::PythonOutputStream(
34 std::shared_ptr<PythonFileObjectAdapter> adapter, int buffer_size) {
35 if (buffer_size <= 0) buffer_size = 128 * 1024; // 128 KB
36 adapter_ = adapter;
37 subtle::ResizeStringUninitialized(&buffer_, buffer_size);
38 is_first_call_ = true;
39 position_ = 0;
40 count_in_buffer_ = 0;
41 buffer_offset_ = 0;
42 status_ = util::OkStatus();
43 }
44
Next(void ** data)45 crypto::tink::util::StatusOr<int> PythonOutputStream::Next(void** data) {
46 if (!status_.ok()) return status_;
47
48 // This is the first call to Next(), so we return the whole buffer.
49 if (is_first_call_) {
50 is_first_call_ = false;
51 count_in_buffer_ = buffer_.size();
52 position_ = buffer_.size();
53 *data = &buffer_[0];
54 return buffer_.size();
55 }
56
57 // If some space was backed up, return it first.
58 if (count_in_buffer_ < buffer_.size()) {
59 int count_backedup = buffer_.size() - count_in_buffer_;
60 position_ += count_backedup;
61 buffer_offset_ = count_in_buffer_;
62 count_in_buffer_ = buffer_.size();
63 *data = &buffer_[buffer_offset_];
64 return count_backedup;
65 }
66
67 // Write the data from the buffer and return available space in the buffer.
68 // The available space might not span the entire buffer, as writing
69 // may succeed only for a prefix of the buffer -- in this case the data still
70 // to be written is shifted in the buffer and the remaining space is returned.
71 auto write_result = adapter_->Write(buffer_);
72 if (!write_result.ok()) return status_ = write_result.status();
73
74 // Some data was written, so we can return some portion of buffer_.
75 int written = write_result.value();
76 position_ += written;
77 count_in_buffer_ = buffer_.size();
78 buffer_offset_ = buffer_.size() - written;
79 if (written < buffer_.size()) {
80 // Only part of the data was written, shift the remaining data in buffer_.
81 // Using memmove, as source and destination may overlap.
82 std::memmove(&buffer_[0], &buffer_[written], buffer_offset_);
83 }
84 *data = &buffer_[buffer_offset_];
85 return written;
86 }
87
BackUp(int count)88 void PythonOutputStream::BackUp(int count) {
89 if (!status_.ok() || count < 1 || count_in_buffer_ == 0) return;
90 int actual_count = std::min(count, count_in_buffer_ - buffer_offset_);
91 count_in_buffer_ -= actual_count;
92 position_ -= actual_count;
93 }
94
~PythonOutputStream()95 PythonOutputStream::~PythonOutputStream() { Close().IgnoreError(); }
96
Close()97 util::Status PythonOutputStream::Close() {
98 if (!status_.ok()) return status_;
99 if (count_in_buffer_ > 0) {
100 // Try to write the remaining bytes.
101 int written = 0;
102 while (written < count_in_buffer_) {
103 auto write_result = adapter_->Write(absl::string_view(buffer_).substr(
104 written, count_in_buffer_ - written));
105 if (!write_result.ok()) return write_result.status();
106 written += write_result.value();
107 }
108 }
109 status_ = adapter_->Close();
110 if (!status_.ok()) return status_;
111 status_ =
112 util::Status(absl::StatusCode::kFailedPrecondition, "Stream closed");
113 return util::OkStatus();
114 }
115
Position() const116 int64_t PythonOutputStream::Position() const { return position_; }
117
118 } // namespace tink
119 } // namespace crypto
120