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 "contrib/zstd/wrapper/wrapper_zstd.h"
16
17 #include <errno.h>
18 #include <fcntl.h>
19 #include <unistd.h>
20
21 #include <cstdio>
22 #include <cstdlib>
23 #include <iostream>
24 #include <memory>
25
26 constexpr size_t kFileMaxSize = 1024 * 1024 * 1024; // 1GB
27
FDGetSize(int fd)28 off_t FDGetSize(int fd) {
29 off_t size = lseek(fd, 0, SEEK_END);
30 if (size < 0) {
31 return -1;
32 }
33 if (lseek(fd, 0, SEEK_SET) < 0) {
34 return -1;
35 }
36
37 return size;
38 }
39
ZSTD_compress_fd(int fdin,int fdout,int level)40 int ZSTD_compress_fd(int fdin, int fdout, int level) {
41 off_t sizein = FDGetSize(fdin);
42 if (sizein <= 0) {
43 return -1;
44 }
45
46 size_t sizeout = ZSTD_compressBound(sizein);
47
48 auto bufin = std::make_unique<int8_t[]>(sizein);
49 auto bufout = std::make_unique<int8_t[]>(sizeout);
50
51 if (read(fdin, bufin.get(), sizein) != sizein) {
52 return -1;
53 }
54
55 int retsize =
56 ZSTD_compress(bufout.get(), sizeout, bufin.get(), sizein, level);
57 if (ZSTD_isError(retsize)) {
58 return -1;
59 }
60
61 if (write(fdout, bufout.get(), retsize) != retsize) {
62 return -1;
63 }
64
65 return 0;
66 }
67
ZSTD_compressStream_fd(ZSTD_CCtx * cctx,int fdin,int fdout)68 int ZSTD_compressStream_fd(ZSTD_CCtx* cctx, int fdin, int fdout) {
69 size_t sizein = ZSTD_CStreamInSize();
70 size_t sizeout = ZSTD_CStreamOutSize();
71
72 auto bufin = std::make_unique<int8_t[]>(sizein);
73 auto bufout = std::make_unique<int8_t[]>(sizeout);
74
75 ssize_t size;
76 while ((size = read(fdin, bufin.get(), sizein)) > 0) {
77 ZSTD_inBuffer_s struct_in;
78 struct_in.src = bufin.get();
79 struct_in.pos = 0;
80 struct_in.size = size;
81
82 ZSTD_EndDirective mode = ZSTD_e_continue;
83 if (size < sizein) {
84 mode = ZSTD_e_end;
85 }
86
87 bool isdone = false;
88 while (!isdone) {
89 ZSTD_outBuffer_s struct_out;
90 struct_out.dst = bufout.get();
91 struct_out.pos = 0;
92 struct_out.size = sizeout;
93
94 size_t remaining =
95 ZSTD_compressStream2(cctx, &struct_out, &struct_in, mode);
96 if (ZSTD_isError(remaining)) {
97 return -1;
98 }
99 if (write(fdout, bufout.get(), struct_out.pos) != struct_out.pos) {
100 return -1;
101 }
102
103 if (mode == ZSTD_e_continue) {
104 isdone = (struct_in.pos == size);
105 } else {
106 isdone = (remaining == 0);
107 }
108 }
109 }
110
111 if (size != 0) {
112 return -1;
113 }
114
115 return 0;
116 }
117
ZSTD_decompress_fd(int fdin,int fdout)118 int ZSTD_decompress_fd(int fdin, int fdout) {
119 off_t sizein = FDGetSize(fdin);
120 if (sizein <= 0) {
121 return -1;
122 }
123 auto bufin = std::make_unique<int8_t[]>(sizein);
124 if (read(fdin, bufin.get(), sizein) != sizein) {
125 return -1;
126 }
127
128 size_t sizeout = ZSTD_getFrameContentSize(bufin.get(), sizein);
129 if (ZSTD_isError(sizeout) || sizeout > kFileMaxSize) {
130 return -1;
131 }
132
133 auto bufout = std::make_unique<int8_t[]>(sizeout);
134
135 size_t desize = ZSTD_decompress(bufout.get(), sizeout, bufin.get(), sizein);
136 if (ZSTD_isError(desize) || desize != sizeout) {
137 return -1;
138 }
139
140 if (write(fdout, bufout.get(), sizeout) != sizeout) {
141 return -1;
142 }
143
144 return 0;
145 }
146
ZSTD_decompressStream_fd(ZSTD_DCtx * dctx,int fdin,int fdout)147 int ZSTD_decompressStream_fd(ZSTD_DCtx* dctx, int fdin, int fdout) {
148 size_t sizein = ZSTD_CStreamInSize();
149 size_t sizeout = ZSTD_CStreamOutSize();
150
151 auto bufin = std::make_unique<int8_t[]>(sizein);
152 auto bufout = std::make_unique<int8_t[]>(sizeout);
153
154 ssize_t size;
155 while ((size = read(fdin, bufin.get(), sizein)) > 0) {
156 ZSTD_inBuffer_s struct_in;
157 struct_in.src = bufin.get();
158 struct_in.pos = 0;
159 struct_in.size = size;
160
161 while (struct_in.pos < size) {
162 ZSTD_outBuffer_s struct_out;
163 struct_out.dst = bufout.get();
164 struct_out.pos = 0;
165 struct_out.size = sizeout;
166
167 size_t ret = ZSTD_decompressStream(dctx, &struct_out, &struct_in);
168 if (ZSTD_isError(ret)) {
169 return -1;
170 }
171 if (write(fdout, bufout.get(), struct_out.pos) != struct_out.pos) {
172 return -1;
173 }
174 }
175 }
176
177 if (size != 0) {
178 return -1;
179 }
180
181 return 0;
182 }
183