Merge pull request #106 from oshogbo:zstd_minor

PiperOrigin-RevId: 426328587
Change-Id: I22def7bb65ce0ec0f085bf3c44c7320026fb21d7
This commit is contained in:
Copybara-Service 2022-02-03 23:56:16 -08:00
commit c1fbdebeb8
4 changed files with 43 additions and 39 deletions

View File

@ -51,9 +51,11 @@ add_sapi_library(
ZSTD_createDCtx ZSTD_createDCtx
ZSTD_freeDCtx ZSTD_freeDCtx
ZSTD_createCCtx
ZSTD_freeCCtx
ZSTD_CCtx_setParameter ZSTD_CCtx_setParameter
ZSTD_compressBound
ZSTD_compress ZSTD_compress
ZSTD_compressStream ZSTD_compressStream
ZSTD_compressStream2 ZSTD_compressStream2

View File

@ -207,7 +207,7 @@ TEST(SandboxTest, CheckCompressStream) {
ASSERT_TRUE(outfile.is_open()); ASSERT_TRUE(outfile.is_open());
absl::Status status = CompressStream(api, infile, outfile, 0); absl::Status status = CompressStream(api, infile, outfile, 0);
ASSERT_THAT(status, IsOk()) << "Unable to compress file in memory"; ASSERT_THAT(status, IsOk()) << "Unable to compress stream";
infile.clear(); infile.clear();
@ -233,7 +233,7 @@ TEST(SandboxTest, CheckDecompressStream) {
ASSERT_TRUE(outfile.is_open()); ASSERT_TRUE(outfile.is_open());
absl::Status status = DecompressStream(api, infile, outfile); absl::Status status = DecompressStream(api, infile, outfile);
ASSERT_THAT(status, IsOk()) << "Unable to decompress file in memory"; ASSERT_THAT(status, IsOk()) << "Unable to decompress stream";
ASSERT_GT(outfile.tellp(), infile.tellg()); ASSERT_GT(outfile.tellp(), infile.tellg());
@ -266,7 +266,7 @@ TEST(SandboxTest, CheckCompressAndDecompressStream) {
ASSERT_TRUE(outmiddle.is_open()); ASSERT_TRUE(outmiddle.is_open());
status = CompressStream(api, infile, outmiddle, 0); status = CompressStream(api, infile, outmiddle, 0);
ASSERT_THAT(status, IsOk()) << "Unable to compress file in memory"; ASSERT_THAT(status, IsOk()) << "Unable to compress stream";
infile.clear(); infile.clear();
ASSERT_LT(outmiddle.tellp(), infile.tellg()); ASSERT_LT(outmiddle.tellp(), infile.tellg());

View File

@ -28,12 +28,12 @@ std::streamsize GetStreamSize(std::ifstream& stream) {
return ssize; return ssize;
} }
absl::Status CompressInMemory(ZstdApi& api, std::ifstream& in_file, absl::Status CompressInMemory(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_stream, int level) { std::ofstream& out_stream, int level) {
std::streamsize ssize = GetStreamSize(in_file); std::streamsize ssize = GetStreamSize(in_stream);
sapi::v::Array<uint8_t> inbuf(ssize); sapi::v::Array<uint8_t> inbuf(ssize);
in_file.read(reinterpret_cast<char*>(inbuf.GetData()), ssize); in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), ssize);
if (in_file.gcount() != ssize) { if (in_stream.gcount() != ssize) {
return absl::UnavailableError("Unable to read file"); return absl::UnavailableError("Unable to read file");
} }
@ -57,13 +57,13 @@ absl::Status CompressInMemory(ZstdApi& api, std::ifstream& in_file,
return absl::OkStatus(); return absl::OkStatus();
} }
absl::Status DecompressInMemory(ZstdApi& api, std::ifstream& in_file, absl::Status DecompressInMemory(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_stream) { std::ofstream& out_stream) {
int iserr; int iserr;
std::streamsize ssize = GetStreamSize(in_file); std::streamsize ssize = GetStreamSize(in_stream);
sapi::v::Array<uint8_t> inbuf(ssize); sapi::v::Array<uint8_t> inbuf(ssize);
in_file.read(reinterpret_cast<char*>(inbuf.GetData()), ssize); in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), ssize);
if (in_file.gcount() != ssize) { if (in_stream.gcount() != ssize) {
return absl::UnavailableError("Unable to read file"); return absl::UnavailableError("Unable to read file");
} }
@ -94,7 +94,7 @@ absl::Status DecompressInMemory(ZstdApi& api, std::ifstream& in_file,
return absl::OkStatus(); return absl::OkStatus();
} }
absl::Status CompressStream(ZstdApi& api, std::ifstream& in_file, absl::Status CompressStream(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_stream, int level) { std::ofstream& out_stream, int level) {
int iserr; int iserr;
@ -110,23 +110,25 @@ absl::Status CompressStream(ZstdApi& api, std::ifstream& in_file,
} }
// Create Zstd context. // Create Zstd context.
SAPI_ASSIGN_OR_RETURN(ZSTD_DCtx * dctx, api.ZSTD_createDCtx()); SAPI_ASSIGN_OR_RETURN(ZSTD_CCtx * cctx, api.ZSTD_createCCtx());
sapi::v::RemotePtr rdctx(dctx); sapi::v::RemotePtr rcctx(cctx);
SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_CCtx_setParameter( SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_CCtx_setParameter(
&rdctx, ZSTD_c_compressionLevel, level)); &rcctx, ZSTD_c_compressionLevel, level));
if (!iserr) { SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr))
if (iserr) {
return absl::UnavailableError("Unable to set parameter"); return absl::UnavailableError("Unable to set parameter");
} }
SAPI_ASSIGN_OR_RETURN( SAPI_ASSIGN_OR_RETURN(
iserr, api.ZSTD_CCtx_setParameter(&rdctx, ZSTD_c_checksumFlag, 1)); iserr, api.ZSTD_CCtx_setParameter(&rcctx, ZSTD_c_checksumFlag, 1));
if (!iserr) { SAPI_ASSIGN_OR_RETURN(iserr, api.ZSTD_isError(iserr))
if (iserr) {
return absl::UnavailableError("Unable to set parameter"); return absl::UnavailableError("Unable to set parameter");
} }
// Compress. // Compress.
while (in_file) { while (in_stream) {
in_file.read(reinterpret_cast<char*>(inbuf.GetData()), inbuf_size); in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), inbuf_size);
if (!api.GetSandbox()->TransferToSandboxee(&inbuf).ok()) { if (!api.GetSandbox()->TransferToSandboxee(&inbuf).ok()) {
return absl::UnavailableError("Unable to transfer data"); return absl::UnavailableError("Unable to transfer data");
@ -135,10 +137,10 @@ absl::Status CompressStream(ZstdApi& api, std::ifstream& in_file,
sapi::v::Struct<ZSTD_inBuffer_s> struct_in; sapi::v::Struct<ZSTD_inBuffer_s> struct_in;
struct_in.mutable_data()->src = static_cast<uint8_t*>(inbuf.GetRemote()); struct_in.mutable_data()->src = static_cast<uint8_t*>(inbuf.GetRemote());
struct_in.mutable_data()->pos = 0; struct_in.mutable_data()->pos = 0;
struct_in.mutable_data()->size = in_file.gcount(); struct_in.mutable_data()->size = in_stream.gcount();
ZSTD_EndDirective mode = ZSTD_e_continue; ZSTD_EndDirective mode = ZSTD_e_continue;
if (in_file.gcount() < inbuf_size) { if (in_stream.gcount() < inbuf_size) {
mode = ZSTD_e_end; mode = ZSTD_e_end;
} }
@ -151,7 +153,7 @@ absl::Status CompressStream(ZstdApi& api, std::ifstream& in_file,
struct_out.mutable_data()->size = outbuf.GetSize(); struct_out.mutable_data()->size = outbuf.GetSize();
SAPI_ASSIGN_OR_RETURN(size_t remaining, api.ZSTD_compressStream2( SAPI_ASSIGN_OR_RETURN(size_t remaining, api.ZSTD_compressStream2(
&rdctx, struct_out.PtrBoth(), &rcctx, struct_out.PtrBoth(),
struct_in.PtrBoth(), mode)); struct_in.PtrBoth(), mode));
SAPI_ASSIGN_OR_RETURN(int iserr, api.ZSTD_isError(remaining)) SAPI_ASSIGN_OR_RETURN(int iserr, api.ZSTD_isError(remaining))
if (iserr) { if (iserr) {
@ -168,19 +170,19 @@ absl::Status CompressStream(ZstdApi& api, std::ifstream& in_file,
} }
if (mode == ZSTD_e_continue) { if (mode == ZSTD_e_continue) {
isdone = (struct_in.mutable_data()->pos == in_file.gcount()); isdone = (struct_in.mutable_data()->pos == in_stream.gcount());
} else { } else {
isdone = (remaining == 0); isdone = (remaining == 0);
} }
} }
} }
api.ZSTD_freeDCtx(&rdctx).IgnoreError(); api.ZSTD_freeDCtx(&rcctx).IgnoreError();
return absl::OkStatus(); return absl::OkStatus();
} }
absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_file, absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_stream) { std::ofstream& out_stream) {
// Create necessary buffers. // Create necessary buffers.
SAPI_ASSIGN_OR_RETURN(size_t inbuf_size, api.ZSTD_CStreamInSize()); SAPI_ASSIGN_OR_RETURN(size_t inbuf_size, api.ZSTD_CStreamInSize());
@ -198,8 +200,8 @@ absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_file,
sapi::v::RemotePtr rdctx(dctx); sapi::v::RemotePtr rdctx(dctx);
// Decompress. // Decompress.
while (in_file) { while (in_stream) {
in_file.read(reinterpret_cast<char*>(inbuf.GetData()), inbuf_size); in_stream.read(reinterpret_cast<char*>(inbuf.GetData()), inbuf_size);
if (!api.GetSandbox()->TransferToSandboxee(&inbuf).ok()) { if (!api.GetSandbox()->TransferToSandboxee(&inbuf).ok()) {
return absl::UnavailableError("Unable to transfer data"); return absl::UnavailableError("Unable to transfer data");
@ -207,10 +209,10 @@ absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_file,
sapi::v::Struct<ZSTD_inBuffer_s> struct_in; sapi::v::Struct<ZSTD_inBuffer_s> struct_in;
*struct_in.mutable_data() = {static_cast<uint8_t*>(inbuf.GetRemote()), *struct_in.mutable_data() = {static_cast<uint8_t*>(inbuf.GetRemote()),
(size_t)in_file.gcount(), 0}; (size_t)in_stream.gcount(), 0};
bool isdone = false; bool isdone = false;
while (struct_in.mutable_data()->pos < in_file.gcount()) { while (struct_in.mutable_data()->pos < in_stream.gcount()) {
sapi::v::Struct<ZSTD_outBuffer_s> struct_out; sapi::v::Struct<ZSTD_outBuffer_s> struct_out;
*struct_out.mutable_data() = {static_cast<uint8_t*>(outbuf.GetRemote()), *struct_out.mutable_data() = {static_cast<uint8_t*>(outbuf.GetRemote()),
(size_t)outbuf.GetSize(), 0}; (size_t)outbuf.GetSize(), 0};

View File

@ -20,14 +20,14 @@
#include "contrib/zstd/sandboxed.h" #include "contrib/zstd/sandboxed.h"
absl::Status CompressInMemory(ZstdApi& api, std::ifstream& in_file, absl::Status CompressInMemory(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_file, int level); std::ofstream& out_stream, int level);
absl::Status DecompressInMemory(ZstdApi& api, std::ifstream& in_file, absl::Status DecompressInMemory(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_file); std::ofstream& out_stream);
absl::Status CompressStream(ZstdApi& api, std::ifstream& in_file, absl::Status CompressStream(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_file, int level); std::ofstream& out_stream, int level);
absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_file, absl::Status DecompressStream(ZstdApi& api, std::ifstream& in_stream,
std::ofstream& out_file); std::ofstream& out_stream);
#endif // CONTRIB_ZSTD_UTILS_UTILS_ZSTD_H_ #endif // CONTRIB_ZSTD_UTILS_UTILS_ZSTD_H_