1 /*
2 * Copyright (c) 2017-present, Facebook, Inc.
3 * All rights reserved.
4 *
5 * This source code is licensed under both the BSD-style license (found in the
6 * LICENSE file in the root directory of this source tree) and the GPLv2 (found
7 * in the COPYING file in the root directory of this source tree).
8 */
9
10 #include <stdlib.h> /* malloc, free */
11 #include <limits.h> /* UINT_MAX */
12 #include <assert.h>
13
14 #define XXH_STATIC_LINKING_ONLY
15 #define XXH_NAMESPACE ZSTD_
16 #include "xxhash.h"
17
18 #define ZSTD_STATIC_LINKING_ONLY
19 #include "zstd.h"
20 #include "zstd_errors.h"
21 #include "mem.h"
22 #include "zstd_seekable.h"
23
24 #define CHECK_Z(f) { size_t const ret = (f); if (ret != 0) return ret; }
25
26 #undef ERROR
27 #define ERROR(name) ((size_t)-ZSTD_error_##name)
28
29 #undef MIN
30 #undef MAX
31 #define MIN(a, b) ((a) < (b) ? (a) : (b))
32 #define MAX(a, b) ((a) > (b) ? (a) : (b))
33
34 typedef struct {
35 U32 cSize;
36 U32 dSize;
37 U32 checksum;
38 } framelogEntry_t;
39
40 struct ZSTD_frameLog_s {
41 framelogEntry_t* entries;
42 U32 size;
43 U32 capacity;
44
45 int checksumFlag;
46
47 /* for use when streaming out the seek table */
48 U32 seekTablePos;
49 U32 seekTableIndex;
50 } framelog_t;
51
52 struct ZSTD_seekable_CStream_s {
53 ZSTD_CStream* cstream;
54 ZSTD_frameLog framelog;
55
56 U32 frameCSize;
57 U32 frameDSize;
58
59 XXH64_state_t xxhState;
60
61 U32 maxFrameSize;
62
63 int writingSeekTable;
64 };
65
ZSTD_seekable_frameLog_allocVec(ZSTD_frameLog * fl)66 size_t ZSTD_seekable_frameLog_allocVec(ZSTD_frameLog* fl)
67 {
68 /* allocate some initial space */
69 size_t const FRAMELOG_STARTING_CAPACITY = 16;
70 fl->entries = (framelogEntry_t*)malloc(
71 sizeof(framelogEntry_t) * FRAMELOG_STARTING_CAPACITY);
72 if (fl->entries == NULL) return ERROR(memory_allocation);
73 fl->capacity = FRAMELOG_STARTING_CAPACITY;
74
75 return 0;
76 }
77
ZSTD_seekable_frameLog_freeVec(ZSTD_frameLog * fl)78 size_t ZSTD_seekable_frameLog_freeVec(ZSTD_frameLog* fl)
79 {
80 if (fl != NULL) free(fl->entries);
81 return 0;
82 }
83
ZSTD_seekable_createFrameLog(int checksumFlag)84 ZSTD_frameLog* ZSTD_seekable_createFrameLog(int checksumFlag)
85 {
86 ZSTD_frameLog* fl = malloc(sizeof(ZSTD_frameLog));
87 if (fl == NULL) return NULL;
88
89 if (ZSTD_isError(ZSTD_seekable_frameLog_allocVec(fl))) {
90 free(fl);
91 return NULL;
92 }
93
94 fl->checksumFlag = checksumFlag;
95 fl->seekTablePos = 0;
96 fl->seekTableIndex = 0;
97 fl->size = 0;
98
99 return fl;
100 }
101
ZSTD_seekable_freeFrameLog(ZSTD_frameLog * fl)102 size_t ZSTD_seekable_freeFrameLog(ZSTD_frameLog* fl)
103 {
104 ZSTD_seekable_frameLog_freeVec(fl);
105 free(fl);
106 return 0;
107 }
108
ZSTD_seekable_createCStream()109 ZSTD_seekable_CStream* ZSTD_seekable_createCStream()
110 {
111 ZSTD_seekable_CStream* zcs = malloc(sizeof(ZSTD_seekable_CStream));
112
113 if (zcs == NULL) return NULL;
114
115 memset(zcs, 0, sizeof(*zcs));
116
117 zcs->cstream = ZSTD_createCStream();
118 if (zcs->cstream == NULL) goto failed1;
119
120 if (ZSTD_isError(ZSTD_seekable_frameLog_allocVec(&zcs->framelog))) goto failed2;
121
122 return zcs;
123
124 failed2:
125 ZSTD_freeCStream(zcs->cstream);
126 failed1:
127 free(zcs);
128 return NULL;
129 }
130
ZSTD_seekable_freeCStream(ZSTD_seekable_CStream * zcs)131 size_t ZSTD_seekable_freeCStream(ZSTD_seekable_CStream* zcs)
132 {
133 if (zcs == NULL) return 0; /* support free on null */
134 ZSTD_freeCStream(zcs->cstream);
135 ZSTD_seekable_frameLog_freeVec(&zcs->framelog);
136 free(zcs);
137
138 return 0;
139 }
140
ZSTD_seekable_initCStream(ZSTD_seekable_CStream * zcs,int compressionLevel,int checksumFlag,unsigned maxFrameSize)141 size_t ZSTD_seekable_initCStream(ZSTD_seekable_CStream* zcs,
142 int compressionLevel,
143 int checksumFlag,
144 unsigned maxFrameSize)
145 {
146 zcs->framelog.size = 0;
147 zcs->frameCSize = 0;
148 zcs->frameDSize = 0;
149
150 /* make sure maxFrameSize has a reasonable value */
151 if (maxFrameSize > ZSTD_SEEKABLE_MAX_FRAME_DECOMPRESSED_SIZE) {
152 return ERROR(frameParameter_unsupported);
153 }
154
155 zcs->maxFrameSize = maxFrameSize
156 ? maxFrameSize
157 : ZSTD_SEEKABLE_MAX_FRAME_DECOMPRESSED_SIZE;
158
159 zcs->framelog.checksumFlag = checksumFlag;
160 if (zcs->framelog.checksumFlag) {
161 XXH64_reset(&zcs->xxhState, 0);
162 }
163
164 zcs->framelog.seekTablePos = 0;
165 zcs->framelog.seekTableIndex = 0;
166 zcs->writingSeekTable = 0;
167
168 return ZSTD_initCStream(zcs->cstream, compressionLevel);
169 }
170
ZSTD_seekable_logFrame(ZSTD_frameLog * fl,unsigned compressedSize,unsigned decompressedSize,unsigned checksum)171 size_t ZSTD_seekable_logFrame(ZSTD_frameLog* fl,
172 unsigned compressedSize,
173 unsigned decompressedSize,
174 unsigned checksum)
175 {
176 if (fl->size == ZSTD_SEEKABLE_MAXFRAMES)
177 return ERROR(frameIndex_tooLarge);
178
179 /* grow the buffer if required */
180 if (fl->size == fl->capacity) {
181 /* exponential size increase for constant amortized runtime */
182 size_t const newCapacity = fl->capacity * 2;
183 framelogEntry_t* const newEntries = realloc(fl->entries,
184 sizeof(framelogEntry_t) * newCapacity);
185
186 if (newEntries == NULL) return ERROR(memory_allocation);
187
188 fl->entries = newEntries;
189 assert(newCapacity <= UINT_MAX);
190 fl->capacity = (U32)newCapacity;
191 }
192
193 fl->entries[fl->size] = (framelogEntry_t){
194 compressedSize, decompressedSize, checksum
195 };
196 fl->size++;
197
198 return 0;
199 }
200
ZSTD_seekable_endFrame(ZSTD_seekable_CStream * zcs,ZSTD_outBuffer * output)201 size_t ZSTD_seekable_endFrame(ZSTD_seekable_CStream* zcs, ZSTD_outBuffer* output)
202 {
203 size_t const prevOutPos = output->pos;
204 /* end the frame */
205 size_t ret = ZSTD_endStream(zcs->cstream, output);
206
207 zcs->frameCSize += output->pos - prevOutPos;
208
209 /* need to flush before doing the rest */
210 if (ret) return ret;
211
212 /* frame done */
213
214 /* store the frame data for later */
215 ret = ZSTD_seekable_logFrame(
216 &zcs->framelog, zcs->frameCSize, zcs->frameDSize,
217 zcs->framelog.checksumFlag
218 ? XXH64_digest(&zcs->xxhState) & 0xFFFFFFFFU
219 : 0);
220 if (ret) return ret;
221
222 /* reset for the next frame */
223 zcs->frameCSize = 0;
224 zcs->frameDSize = 0;
225
226 ZSTD_resetCStream(zcs->cstream, 0);
227 if (zcs->framelog.checksumFlag)
228 XXH64_reset(&zcs->xxhState, 0);
229
230 return 0;
231 }
232
ZSTD_seekable_compressStream(ZSTD_seekable_CStream * zcs,ZSTD_outBuffer * output,ZSTD_inBuffer * input)233 size_t ZSTD_seekable_compressStream(ZSTD_seekable_CStream* zcs, ZSTD_outBuffer* output, ZSTD_inBuffer* input)
234 {
235 const BYTE* const inBase = (const BYTE*) input->src + input->pos;
236 size_t inLen = input->size - input->pos;
237
238 inLen = MIN(inLen, (size_t)(zcs->maxFrameSize - zcs->frameDSize));
239
240 /* if we haven't finished flushing the last frame, don't start writing a new one */
241 if (inLen > 0) {
242 ZSTD_inBuffer inTmp = { inBase, inLen, 0 };
243 size_t const prevOutPos = output->pos;
244
245 size_t const ret = ZSTD_compressStream(zcs->cstream, output, &inTmp);
246
247 if (zcs->framelog.checksumFlag) {
248 XXH64_update(&zcs->xxhState, inBase, inTmp.pos);
249 }
250
251 zcs->frameCSize += output->pos - prevOutPos;
252 zcs->frameDSize += inTmp.pos;
253
254 input->pos += inTmp.pos;
255
256 if (ZSTD_isError(ret)) return ret;
257 }
258
259 if (zcs->maxFrameSize == zcs->frameDSize) {
260 /* log the frame and start over */
261 size_t const ret = ZSTD_seekable_endFrame(zcs, output);
262 if (ZSTD_isError(ret)) return ret;
263
264 /* get the client ready for the next frame */
265 return (size_t)zcs->maxFrameSize;
266 }
267
268 return (size_t)(zcs->maxFrameSize - zcs->frameDSize);
269 }
270
ZSTD_seekable_seekTableSize(const ZSTD_frameLog * fl)271 static inline size_t ZSTD_seekable_seekTableSize(const ZSTD_frameLog* fl)
272 {
273 size_t const sizePerFrame = 8 + (fl->checksumFlag?4:0);
274 size_t const seekTableLen = ZSTD_SKIPPABLEHEADERSIZE +
275 sizePerFrame * fl->size +
276 ZSTD_seekTableFooterSize;
277
278 return seekTableLen;
279 }
280
ZSTD_stwrite32(ZSTD_frameLog * fl,ZSTD_outBuffer * output,U32 const value,U32 const offset)281 static inline size_t ZSTD_stwrite32(ZSTD_frameLog* fl,
282 ZSTD_outBuffer* output, U32 const value,
283 U32 const offset)
284 {
285 if (fl->seekTablePos < offset + 4) {
286 BYTE tmp[4]; /* so that we can work with buffers too small to write a whole word to */
287 size_t const lenWrite =
288 MIN(output->size - output->pos, offset + 4 - fl->seekTablePos);
289 MEM_writeLE32(tmp, value);
290 memcpy((BYTE*)output->dst + output->pos,
291 tmp + (fl->seekTablePos - offset), lenWrite);
292 output->pos += lenWrite;
293 fl->seekTablePos += lenWrite;
294
295 if (lenWrite < 4) return ZSTD_seekable_seekTableSize(fl) - fl->seekTablePos;
296 }
297 return 0;
298 }
299
ZSTD_seekable_writeSeekTable(ZSTD_frameLog * fl,ZSTD_outBuffer * output)300 size_t ZSTD_seekable_writeSeekTable(ZSTD_frameLog* fl, ZSTD_outBuffer* output)
301 {
302 /* seekTableIndex: the current index in the table and
303 * seekTableSize: the amount of the table written so far
304 *
305 * This function is written this way so that if it has to return early
306 * because of a small buffer, it can keep going where it left off.
307 */
308
309 size_t const sizePerFrame = 8 + (fl->checksumFlag?4:0);
310 size_t const seekTableLen = ZSTD_seekable_seekTableSize(fl);
311
312 CHECK_Z(ZSTD_stwrite32(fl, output, ZSTD_MAGIC_SKIPPABLE_START | 0xE, 0));
313 assert(seekTableLen <= (size_t)UINT_MAX);
314 CHECK_Z(ZSTD_stwrite32(fl, output, (U32)seekTableLen - ZSTD_SKIPPABLEHEADERSIZE, 4));
315
316 while (fl->seekTableIndex < fl->size) {
317 unsigned long long const start = ZSTD_SKIPPABLEHEADERSIZE + sizePerFrame * fl->seekTableIndex;
318 assert(start + 8 <= UINT_MAX);
319 CHECK_Z(ZSTD_stwrite32(fl, output,
320 fl->entries[fl->seekTableIndex].cSize,
321 (U32)start + 0));
322
323 CHECK_Z(ZSTD_stwrite32(fl, output,
324 fl->entries[fl->seekTableIndex].dSize,
325 (U32)start + 4));
326
327 if (fl->checksumFlag) {
328 CHECK_Z(ZSTD_stwrite32(
329 fl, output, fl->entries[fl->seekTableIndex].checksum,
330 (U32)start + 8));
331 }
332
333 fl->seekTableIndex++;
334 }
335
336 assert(seekTableLen <= UINT_MAX);
337 CHECK_Z(ZSTD_stwrite32(fl, output, fl->size,
338 (U32)seekTableLen - ZSTD_seekTableFooterSize));
339
340 if (output->size - output->pos < 1) return seekTableLen - fl->seekTablePos;
341 if (fl->seekTablePos < seekTableLen - 4) {
342 BYTE sfd = 0;
343 sfd |= (fl->checksumFlag) << 7;
344
345 ((BYTE*)output->dst)[output->pos] = sfd;
346 output->pos++;
347 fl->seekTablePos++;
348 }
349
350 CHECK_Z(ZSTD_stwrite32(fl, output, ZSTD_SEEKABLE_MAGICNUMBER,
351 (U32)seekTableLen - 4));
352
353 if (fl->seekTablePos != seekTableLen) return ERROR(GENERIC);
354 return 0;
355 }
356
ZSTD_seekable_endStream(ZSTD_seekable_CStream * zcs,ZSTD_outBuffer * output)357 size_t ZSTD_seekable_endStream(ZSTD_seekable_CStream* zcs, ZSTD_outBuffer* output)
358 {
359 if (!zcs->writingSeekTable && zcs->frameDSize) {
360 const size_t endFrame = ZSTD_seekable_endFrame(zcs, output);
361 if (ZSTD_isError(endFrame)) return endFrame;
362 /* return an accurate size hint */
363 if (endFrame) return endFrame + ZSTD_seekable_seekTableSize(&zcs->framelog);
364 }
365
366 zcs->writingSeekTable = 1;
367
368 return ZSTD_seekable_writeSeekTable(&zcs->framelog, output);
369 }
370