1 // Copyright 2022 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 // https://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 #include <fstream>
16 #include <iostream>
17 #include <string>
18
19 #include "contrib/zstd/sandboxed.h"
20
21 constexpr size_t kFileMaxSize = 1024 * 1024 * 1024; // 1GB
22
GetStreamSize(std::ifstream & stream)23 std::streamsize GetStreamSize(std::ifstream& stream) {
24 stream.seekg(0, std::ios_base::end);
25 std::streamsize ssize = stream.tellg();
26 stream.seekg(0, std::ios_base::beg);
27
28 return ssize;
29 }
30
CompressInMemory(ZstdApi & api,std::ifstream & in_stream,std::ofstream & out_stream,int level)31 absl::Status CompressInMemory(ZstdApi& api, std::ifstream& in_stream,
32 std::ofstream& out_stream, int level) {
33 std::streamsize ssize = GetStreamSize(in_stream);
34 sapi::v::Array<uint8_t> inbuf(ssize);
35 in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), ssize);
36 if (in_stream.gcount() != ssize) {
37 return absl::UnavailableError("Unable to read file");
38 }
39
40 SAPI_ASSIGN_OR_RETURN(size_t size, api.ZSTD_compressBound(inbuf.GetSize()));
41 sapi::v::Array<uint8_t> outbuf(size);
42
43 SAPI_ASSIGN_OR_RETURN(
44 size_t outsize,
45 api.ZSTD_compress(outbuf.PtrAfter(), size, inbuf.PtrBefore(),
46 inbuf.GetSize(), level));
47 SAPI_ASSIGN_OR_RETURN(int iserr, api.ZSTD_isError(outsize));
48 if (iserr) {
49 return absl::UnavailableError("Unable to compress file");
50 }
51
52 out_stream.write(reinterpret_cast<char*>(outbuf.GetData()), outsize);
53 if (!out_stream.good()) {
54 return absl::UnavailableError("Unable to write file");
55 }
56
57 return absl::OkStatus();
58 }
59
DecompressInMemory(ZstdApi & api,std::ifstream & in_stream,std::ofstream & out_stream)60 absl::Status DecompressInMemory(ZstdApi& api, std::ifstream& in_stream,
61 std::ofstream& out_stream) {
62 int iserr;
63 std::streamsize ssize = GetStreamSize(in_stream);
64 sapi::v::Array<uint8_t> inbuf(ssize);
65 in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), ssize);
66 if (in_stream.gcount() != ssize) {
67 return absl::UnavailableError("Unable to read file");
68 }
69
70 SAPI_ASSIGN_OR_RETURN(size_t size, api.ZSTD_getFrameContentSize(
71 inbuf.PtrBefore(), inbuf.GetSize()));
72 if (size > kFileMaxSize) {
73 return absl::UnavailableError("File to large");
74 }
75 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(size));
76 if (iserr) {
77 return absl::UnavailableError("Unable to decompress file");
78 }
79 sapi::v::Array<uint8_t> outbuf(size);
80
81 SAPI_ASSIGN_OR_RETURN(size_t desize,
82 api.ZSTD_decompress(outbuf.PtrAfter(), size,
83 inbuf.PtrNone(), inbuf.GetSize()));
84 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(desize));
85 if (iserr) {
86 return absl::UnavailableError("Unable to decompress file");
87 }
88
89 out_stream.write(reinterpret_cast<char*>(outbuf.GetData()), desize);
90 if (!out_stream.good()) {
91 return absl::UnavailableError("Unable to write file");
92 }
93
94 return absl::OkStatus();
95 }
96
CompressStream(ZstdApi & api,std::ifstream & in_stream,std::ofstream & out_stream,int level)97 absl::Status CompressStream(ZstdApi& api, std::ifstream& in_stream,
98 std::ofstream& out_stream, int level) {
99 int iserr;
100
101 // Create necessary buffers.
102 SAPI_ASSIGN_OR_RETURN(size_t inbuf_size, api.ZSTD_CStreamInSize());
103 SAPI_ASSIGN_OR_RETURN(size_t outbuf_size, api.ZSTD_CStreamOutSize());
104 sapi::v::Array<uint8_t> inbuf(inbuf_size);
105 sapi::v::Array<uint8_t> outbuf(outbuf_size);
106
107 if (!api.GetSandbox()->Allocate(&inbuf).ok() ||
108 !api.GetSandbox()->Allocate(&outbuf).ok()) {
109 return absl::UnavailableError("Unable to allocate buffors");
110 }
111
112 // Create Zstd context.
113 SAPI_ASSIGN_OR_RETURN(ZSTD_CCtx * cctx, api.ZSTD_createCCtx());
114 sapi::v::RemotePtr rcctx(cctx);
115
116 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_CCtx_setParameter(
117 &rcctx, ZSTD_c_compressionLevel, level));
118 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr));
119 if (iserr) {
120 return absl::UnavailableError("Unable to set parameter");
121 }
122 SAPI_ASSIGN_OR_RETURN(
123 iserr, api.ZSTD_CCtx_setParameter(&rcctx, ZSTD_c_checksumFlag, 1));
124 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr));
125 if (iserr) {
126 return absl::UnavailableError("Unable to set parameter");
127 }
128
129 // Compress.
130 while (in_stream) {
131 in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), inbuf_size);
132
133 if (!api.GetSandbox()->TransferToSandboxee(&inbuf).ok()) {
134 return absl::UnavailableError("Unable to transfer data");
135 }
136
137 sapi::v::Struct<ZSTD_inBuffer_s> struct_in;
138 struct_in.mutable_data()->src = static_cast<uint8_t*>(inbuf.GetRemote());
139 struct_in.mutable_data()->pos = 0;
140 struct_in.mutable_data()->size = in_stream.gcount();
141
142 ZSTD_EndDirective mode = ZSTD_e_continue;
143 if (in_stream.gcount() < inbuf_size) {
144 mode = ZSTD_e_end;
145 }
146
147 bool isdone = false;
148 while (!isdone) {
149 sapi::v::Struct<ZSTD_outBuffer_s> struct_out;
150 struct_out.mutable_data()->dst =
151 static_cast<uint8_t*>(outbuf.GetRemote());
152 struct_out.mutable_data()->pos = 0;
153 struct_out.mutable_data()->size = outbuf.GetSize();
154
155 SAPI_ASSIGN_OR_RETURN(size_t remaining, api.ZSTD_compressStream2(
156 &rcctx, struct_out.PtrBoth(),
157 struct_in.PtrBoth(), mode));
158 SAPI_ASSIGN_OR_RETURN(int iserr, api.ZSTD_isError(remaining));
159 if (iserr) {
160 return absl::UnavailableError("Unable to decompress file");
161 }
162
163 if (!api.GetSandbox()->TransferFromSandboxee(&outbuf).ok()) {
164 return absl::UnavailableError("Unable to transfer data from");
165 }
166 out_stream.write(reinterpret_cast<char*>(outbuf.GetData()),
167 struct_out.mutable_data()->pos);
168 if (!out_stream.good()) {
169 return absl::UnavailableError("Unable to write file");
170 }
171
172 if (mode == ZSTD_e_continue) {
173 isdone = (struct_in.mutable_data()->pos == in_stream.gcount());
174 } else {
175 isdone = (remaining == 0);
176 }
177 }
178 }
179
180 api.ZSTD_freeDCtx(&rcctx).IgnoreError();
181
182 return absl::OkStatus();
183 }
184
DecompressStream(ZstdApi & api,std::ifstream & in_stream,std::ofstream & out_stream)185 absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_stream,
186 std::ofstream& out_stream) {
187 // Create necessary buffers.
188 SAPI_ASSIGN_OR_RETURN(size_t inbuf_size, api.ZSTD_CStreamInSize());
189 SAPI_ASSIGN_OR_RETURN(size_t outbuf_size, api.ZSTD_CStreamOutSize());
190 sapi::v::Array<uint8_t> inbuf(inbuf_size);
191 sapi::v::Array<uint8_t> outbuf(outbuf_size);
192
193 if (!api.GetSandbox()->Allocate(&inbuf).ok() ||
194 !api.GetSandbox()->Allocate(&outbuf).ok()) {
195 return absl::UnavailableError("Unable to allocate buffors");
196 }
197
198 // Create Zstd context.
199 SAPI_ASSIGN_OR_RETURN(ZSTD_DCtx * dctx, api.ZSTD_createDCtx());
200 sapi::v::RemotePtr rdctx(dctx);
201
202 // Decompress.
203 while (in_stream) {
204 in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), inbuf_size);
205
206 if (!api.GetSandbox()->TransferToSandboxee(&inbuf).ok()) {
207 return absl::UnavailableError("Unable to transfer data");
208 }
209
210 sapi::v::Struct<ZSTD_inBuffer_s> struct_in;
211 *struct_in.mutable_data() = {static_cast<uint8_t*>(inbuf.GetRemote()),
212 (size_t)in_stream.gcount(), 0};
213
214 bool isdone = false;
215 while (struct_in.mutable_data()->pos < in_stream.gcount()) {
216 sapi::v::Struct<ZSTD_outBuffer_s> struct_out;
217 *struct_out.mutable_data() = {static_cast<uint8_t*>(outbuf.GetRemote()),
218 (size_t)outbuf.GetSize(), 0};
219
220 SAPI_ASSIGN_OR_RETURN(
221 size_t ret, api.ZSTD_decompressStream(&rdctx, struct_out.PtrBoth(),
222 struct_in.PtrBoth()));
223 SAPI_ASSIGN_OR_RETURN(int iserr, api.ZSTD_isError(ret));
224 if (iserr) {
225 return absl::UnavailableError("Unable to decompress file");
226 }
227
228 if (!api.GetSandbox()->TransferFromSandboxee(&outbuf).ok()) {
229 return absl::UnavailableError("Unable to transfer data from");
230 }
231
232 out_stream.write(reinterpret_cast<char*>(outbuf.GetData()),
233 struct_out.mutable_data()->pos);
234 if (!out_stream.good()) {
235 return absl::UnavailableError("Unable to write file");
236 }
237 }
238 }
239
240 api.ZSTD_freeDCtx(&rdctx).IgnoreError();
241
242 return absl::OkStatus();
243 }
244
CompressInMemoryFD(ZstdApi & api,sapi::v::Fd & infd,sapi::v::Fd & outfd,int level)245 absl::Status CompressInMemoryFD(ZstdApi& api, sapi::v::Fd& infd,
246 sapi::v::Fd& outfd, int level) {
247 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&infd));
248 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&outfd));
249
250 SAPI_ASSIGN_OR_RETURN(
251 int iserr,
252 api.ZSTD_compress_fd(infd.GetRemoteFd(), outfd.GetRemoteFd(), 0));
253 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr));
254 if (iserr) {
255 return absl::UnavailableError("Unable to compress file");
256 }
257
258 infd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
259 outfd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
260
261 return absl::OkStatus();
262 }
263
DecompressInMemoryFD(ZstdApi & api,sapi::v::Fd & infd,sapi::v::Fd & outfd)264 absl::Status DecompressInMemoryFD(ZstdApi& api, sapi::v::Fd& infd,
265 sapi::v::Fd& outfd) {
266 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&infd));
267 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&outfd));
268
269 SAPI_ASSIGN_OR_RETURN(int iserr, api.ZSTD_decompress_fd(infd.GetRemoteFd(),
270 outfd.GetRemoteFd()));
271 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr));
272 if (iserr) {
273 return absl::UnavailableError("Unable to compress file");
274 }
275
276 infd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
277 outfd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
278
279 return absl::OkStatus();
280 }
281
CompressStreamFD(ZstdApi & api,sapi::v::Fd & infd,sapi::v::Fd & outfd,int level)282 absl::Status CompressStreamFD(ZstdApi& api, sapi::v::Fd& infd,
283 sapi::v::Fd& outfd, int level) {
284 SAPI_ASSIGN_OR_RETURN(ZSTD_CCtx * cctx, api.ZSTD_createCCtx());
285 sapi::v::RemotePtr rcctx(cctx);
286
287 int iserr;
288 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_CCtx_setParameter(
289 &rcctx, ZSTD_c_compressionLevel, level));
290 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr));
291 if (iserr) {
292 return absl::UnavailableError("Unable to set parameter l");
293 }
294 SAPI_ASSIGN_OR_RETURN(
295 iserr, api.ZSTD_CCtx_setParameter(&rcctx, ZSTD_c_checksumFlag, 1));
296 SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr));
297 if (iserr) {
298 return absl::UnavailableError("Unable to set parameter c");
299 }
300
301 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&infd));
302 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&outfd));
303
304 SAPI_ASSIGN_OR_RETURN(iserr,
305 api.ZSTD_compressStream_fd(&rcctx, infd.GetRemoteFd(),
306 outfd.GetRemoteFd()));
307 if (iserr) {
308 return absl::UnavailableError("Unable to compress");
309 }
310
311 infd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
312 outfd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
313
314 return absl::OkStatus();
315 }
316
DecompressStreamFD(ZstdApi & api,sapi::v::Fd & infd,sapi::v::Fd & outfd)317 absl::Status DecompressStreamFD(ZstdApi& api, sapi::v::Fd& infd,
318 sapi::v::Fd& outfd) {
319 SAPI_ASSIGN_OR_RETURN(ZSTD_DCtx * dctx, api.ZSTD_createDCtx());
320 sapi::v::RemotePtr rdctx(dctx);
321
322 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&infd));
323 SAPI_RETURN_IF_ERROR(api.GetSandbox()->TransferToSandboxee(&outfd));
324
325 SAPI_ASSIGN_OR_RETURN(int iserr,
326 api.ZSTD_decompressStream_fd(&rdctx, infd.GetRemoteFd(),
327 outfd.GetRemoteFd()));
328 if (iserr) {
329 return absl::UnavailableError("Unable to decompress");
330 }
331 infd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
332 outfd.CloseRemoteFd(api.GetSandbox()->rpc_channel()).IgnoreError();
333
334 return absl::OkStatus();
335 }
336