Line data Source code
1 : /*-------------------------------------------------------------------------
2 : *
3 : * compress_zstd.c
4 : * Routines for archivers to write a Zstd compressed data stream.
5 : *
6 : * Portions Copyright (c) 1996-2026, PostgreSQL Global Development Group
7 : * Portions Copyright (c) 1994, Regents of the University of California
8 : *
9 : * IDENTIFICATION
10 : * src/bin/pg_dump/compress_zstd.c
11 : *
12 : *-------------------------------------------------------------------------
13 : */
14 :
15 : #include "postgres_fe.h"
16 : #include <unistd.h>
17 :
18 : #include "compress_zstd.h"
19 : #include "pg_backup_utils.h"
20 :
21 : #ifndef USE_ZSTD
22 :
23 : void
24 0 : InitCompressorZstd(CompressorState *cs, const pg_compress_specification compression_spec)
25 : {
26 0 : pg_fatal("this build does not support compression with %s", "ZSTD");
27 : }
28 :
29 : void
30 0 : InitCompressFileHandleZstd(CompressFileHandle *CFH, const pg_compress_specification compression_spec)
31 : {
32 0 : pg_fatal("this build does not support compression with %s", "ZSTD");
33 : }
34 :
35 : #else
36 :
37 : #include <zstd.h>
38 :
39 : typedef struct ZstdCompressorState
40 : {
41 : /* This is a normal file to which we read/write compressed data */
42 : FILE *fp;
43 :
44 : ZSTD_CStream *cstream;
45 : ZSTD_DStream *dstream;
46 : ZSTD_outBuffer output;
47 : ZSTD_inBuffer input;
48 : bool frame_finished;
49 :
50 : /* pointer to a static string like from strerror(), for Zstd_write() */
51 : const char *zstderror;
52 : } ZstdCompressorState;
53 :
54 : static ZSTD_CStream *_ZstdCStreamParams(pg_compress_specification compress);
55 : static void EndCompressorZstd(ArchiveHandle *AH, CompressorState *cs);
56 : static void WriteDataToArchiveZstd(ArchiveHandle *AH, CompressorState *cs,
57 : const void *data, size_t dLen);
58 : static void ReadDataFromArchiveZstd(ArchiveHandle *AH, CompressorState *cs);
59 :
60 : static void
61 : _Zstd_CCtx_setParam_or_die(ZSTD_CStream *cstream,
62 : ZSTD_cParameter param, int value, char *paramname)
63 : {
64 : size_t res;
65 :
66 : res = ZSTD_CCtx_setParameter(cstream, param, value);
67 : if (ZSTD_isError(res))
68 : pg_fatal("could not set compression parameter \"%s\": %s",
69 : paramname, ZSTD_getErrorName(res));
70 : }
71 :
72 : /* Return a compression stream with parameters set per argument */
73 : static ZSTD_CStream *
74 : _ZstdCStreamParams(pg_compress_specification compress)
75 : {
76 : ZSTD_CStream *cstream;
77 :
78 : cstream = ZSTD_createCStream();
79 : if (cstream == NULL)
80 : pg_fatal("could not initialize compression library");
81 :
82 : _Zstd_CCtx_setParam_or_die(cstream, ZSTD_c_compressionLevel,
83 : compress.level, "level");
84 :
85 : if (compress.options & PG_COMPRESSION_OPTION_LONG_DISTANCE)
86 : _Zstd_CCtx_setParam_or_die(cstream,
87 : ZSTD_c_enableLongDistanceMatching,
88 : compress.long_distance, "long");
89 :
90 : return cstream;
91 : }
92 :
93 : /* Helper function for WriteDataToArchiveZstd and EndCompressorZstd */
94 : static void
95 : _ZstdWriteCommon(ArchiveHandle *AH, CompressorState *cs, bool flush)
96 : {
97 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) cs->private_data;
98 : ZSTD_inBuffer *input = &zstdcs->input;
99 : ZSTD_outBuffer *output = &zstdcs->output;
100 :
101 : /* Loop while there's any input or until flushed */
102 : while (input->pos < input->size || flush)
103 : {
104 : size_t res;
105 :
106 : res = ZSTD_compressStream2(zstdcs->cstream, output,
107 : input, flush ? ZSTD_e_end : ZSTD_e_continue);
108 :
109 : if (ZSTD_isError(res))
110 : pg_fatal("could not compress data: %s", ZSTD_getErrorName(res));
111 :
112 : /* Dump output buffer if full, or if we're told to flush */
113 : if (output->pos >= output->size || flush)
114 : {
115 : cs->writeF(AH, output->dst, output->pos);
116 : output->pos = 0;
117 : }
118 :
119 : if (res == 0)
120 : break; /* End of frame or all input consumed */
121 : }
122 : }
123 :
124 : static void
125 : EndCompressorZstd(ArchiveHandle *AH, CompressorState *cs)
126 : {
127 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) cs->private_data;
128 :
129 : if (cs->readF != NULL)
130 : {
131 : Assert(zstdcs->cstream == NULL);
132 : ZSTD_freeDStream(zstdcs->dstream);
133 : pg_free(unconstify(void *, zstdcs->input.src));
134 : }
135 : else if (cs->writeF != NULL)
136 : {
137 : Assert(zstdcs->dstream == NULL);
138 : _ZstdWriteCommon(AH, cs, true);
139 : ZSTD_freeCStream(zstdcs->cstream);
140 : }
141 :
142 : /* output buffer may be allocated in either mode */
143 : pg_free(zstdcs->output.dst);
144 : pg_free(zstdcs);
145 : cs->private_data = NULL;
146 : }
147 :
148 : static void
149 : WriteDataToArchiveZstd(ArchiveHandle *AH, CompressorState *cs,
150 : const void *data, size_t dLen)
151 : {
152 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) cs->private_data;
153 :
154 : zstdcs->input.src = data;
155 : zstdcs->input.size = dLen;
156 : zstdcs->input.pos = 0;
157 :
158 : _ZstdWriteCommon(AH, cs, false);
159 : }
160 :
161 : static void
162 : ReadDataFromArchiveZstd(ArchiveHandle *AH, CompressorState *cs)
163 : {
164 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) cs->private_data;
165 : ZSTD_outBuffer *output = &zstdcs->output;
166 : ZSTD_inBuffer *input = &zstdcs->input;
167 : size_t input_allocated_size = ZSTD_DStreamInSize();
168 : size_t res;
169 : bool frame_finished = false;
170 :
171 : for (;;)
172 : {
173 : size_t cnt;
174 :
175 : /*
176 : * Read compressed data. Note that readF can resize the buffer; the
177 : * new size is tracked and used for future loops.
178 : */
179 : input->size = input_allocated_size;
180 : cnt = cs->readF(AH, (char **) unconstify(void **, &input->src), &input->size);
181 :
182 : /* ensure that readF didn't *shrink* the buffer */
183 : Assert(input->size >= input_allocated_size);
184 : input_allocated_size = input->size;
185 : input->size = cnt;
186 : input->pos = 0;
187 :
188 : if (cnt == 0)
189 : break;
190 :
191 : /* Now decompress */
192 : while (input->pos < input->size)
193 : {
194 : output->pos = 0;
195 : res = ZSTD_decompressStream(zstdcs->dstream, output, input);
196 : if (ZSTD_isError(res))
197 : pg_fatal("could not decompress data: %s", ZSTD_getErrorName(res));
198 : frame_finished = (res == 0);
199 :
200 : /*
201 : * then write the decompressed data to the output handle
202 : */
203 : ((char *) output->dst)[output->pos] = '\0';
204 : ahwrite(output->dst, 1, output->pos, AH);
205 :
206 : if (frame_finished)
207 : break; /* End of frame */
208 : }
209 : }
210 :
211 : if (!frame_finished)
212 : {
213 : ZSTD_inBuffer empty = {NULL, 0, 0};
214 :
215 : /*
216 : * A full output buffer with a positive return value might leave data
217 : * in zstd's internal buffers. Call the decompressor with empty input
218 : * until it has flushed that data.
219 : */
220 : while (!frame_finished)
221 : {
222 : output->pos = 0;
223 : res = ZSTD_decompressStream(zstdcs->dstream, output, &empty);
224 : if (ZSTD_isError(res))
225 : pg_fatal("could not decompress data: %s",
226 : ZSTD_getErrorName(res));
227 : frame_finished = (res == 0);
228 :
229 : if (output->pos == 0 && !frame_finished)
230 : pg_fatal("could not decompress data: compressed stream is incomplete");
231 :
232 : ((char *) output->dst)[output->pos] = '\0';
233 : ahwrite(output->dst, 1, output->pos, AH);
234 : }
235 : }
236 : }
237 :
238 : /* Public routine that supports Zstd compressed data I/O */
239 : void
240 : InitCompressorZstd(CompressorState *cs,
241 : const pg_compress_specification compression_spec)
242 : {
243 : ZstdCompressorState *zstdcs;
244 :
245 : cs->readData = ReadDataFromArchiveZstd;
246 : cs->writeData = WriteDataToArchiveZstd;
247 : cs->end = EndCompressorZstd;
248 :
249 : cs->compression_spec = compression_spec;
250 :
251 : zstdcs = pg_malloc0_object(ZstdCompressorState);
252 : cs->private_data = zstdcs;
253 :
254 : /* We expect that exactly one of readF/writeF is specified */
255 : Assert((cs->readF == NULL) != (cs->writeF == NULL));
256 :
257 : if (cs->readF != NULL)
258 : {
259 : zstdcs->dstream = ZSTD_createDStream();
260 : if (zstdcs->dstream == NULL)
261 : pg_fatal("could not initialize compression library");
262 :
263 : zstdcs->input.size = ZSTD_DStreamInSize();
264 : zstdcs->input.src = pg_malloc(zstdcs->input.size);
265 :
266 : /*
267 : * output.size is the buffer size we tell zstd it can output to.
268 : * Allocate an additional byte such that ReadDataFromArchiveZstd() can
269 : * call ahwrite() with a null-terminated string, which is an optimized
270 : * case in ExecuteSqlCommandBuf().
271 : */
272 : zstdcs->output.size = ZSTD_DStreamOutSize();
273 : zstdcs->output.dst = pg_malloc(zstdcs->output.size + 1);
274 : }
275 : else if (cs->writeF != NULL)
276 : {
277 : zstdcs->cstream = _ZstdCStreamParams(cs->compression_spec);
278 :
279 : zstdcs->output.size = ZSTD_CStreamOutSize();
280 : zstdcs->output.dst = pg_malloc(zstdcs->output.size);
281 : zstdcs->output.pos = 0;
282 : }
283 : }
284 :
285 : /*
286 : * Compressed stream API
287 : */
288 :
289 : static size_t
290 : Zstd_read_internal(void *ptr, size_t size, CompressFileHandle *CFH, bool exit_on_error)
291 : {
292 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) CFH->private_data;
293 : ZSTD_inBuffer *input = &zstdcs->input;
294 : ZSTD_outBuffer *output = &zstdcs->output;
295 : size_t input_allocated_size = ZSTD_DStreamInSize();
296 : size_t res,
297 : cnt;
298 :
299 : /*
300 : * If this is the first call to the reading function, initialize the
301 : * required datastructures.
302 : */
303 : if (zstdcs->dstream == NULL)
304 : {
305 : zstdcs->input.src = pg_malloc0(input_allocated_size);
306 : zstdcs->dstream = ZSTD_createDStream();
307 : if (zstdcs->dstream == NULL)
308 : {
309 : if (exit_on_error)
310 : pg_fatal("could not initialize compression library");
311 : return -1;
312 : }
313 : }
314 :
315 : output->size = size;
316 : output->dst = ptr;
317 : output->pos = 0;
318 :
319 : while (output->pos < output->size)
320 : {
321 : Assert(input->pos <= input->size);
322 : Assert(input->size <= input_allocated_size);
323 :
324 : /*
325 : * If the input is completely consumed, start back at the beginning
326 : */
327 : if (input->pos == input->size)
328 : {
329 : /* input->size is size produced by "fread" */
330 : input->size = 0;
331 : /* input->pos is position consumed by decompress */
332 : input->pos = 0;
333 : }
334 :
335 : /* read compressed data if we must produce more input */
336 : if (input->pos == input->size)
337 : {
338 : cnt = fread(unconstify(void *, input->src), 1, input_allocated_size, zstdcs->fp);
339 : if (ferror(zstdcs->fp))
340 : {
341 : if (exit_on_error)
342 : pg_fatal("could not read from input file: %m");
343 : return -1;
344 : }
345 :
346 : input->size = cnt;
347 :
348 : Assert(cnt <= input_allocated_size);
349 :
350 : /* If we have no more input to consume, verify the final frame. */
351 : if (cnt == 0)
352 : {
353 : ZSTD_inBuffer empty = {NULL, 0, 0};
354 :
355 : if (!zstdcs->frame_finished)
356 : {
357 : res = ZSTD_decompressStream(zstdcs->dstream, output, &empty);
358 : if (ZSTD_isError(res))
359 : {
360 : if (exit_on_error)
361 : pg_fatal("could not decompress data: %s", ZSTD_getErrorName(res));
362 : return -1;
363 : }
364 :
365 : zstdcs->frame_finished = (res == 0);
366 : if (!zstdcs->frame_finished && output->pos == 0)
367 : {
368 : zstdcs->zstderror = "compressed stream is incomplete";
369 : if (exit_on_error)
370 : pg_fatal("could not decompress data: %s", zstdcs->zstderror);
371 : return -1;
372 : }
373 : }
374 :
375 : break;
376 : }
377 : }
378 :
379 : while (input->pos < input->size)
380 : {
381 : /* now decompress */
382 : res = ZSTD_decompressStream(zstdcs->dstream, output, input);
383 :
384 : if (ZSTD_isError(res))
385 : {
386 : if (exit_on_error)
387 : pg_fatal("could not decompress data: %s", ZSTD_getErrorName(res));
388 : return -1;
389 : }
390 :
391 : zstdcs->frame_finished = (res == 0);
392 :
393 : if (output->pos == output->size)
394 : break; /* No more room for output */
395 :
396 : if (zstdcs->frame_finished)
397 : break; /* End of frame */
398 : }
399 : }
400 :
401 : return output->pos;
402 : }
403 :
404 : static void
405 : Zstd_write(const void *ptr, size_t size, CompressFileHandle *CFH)
406 : {
407 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) CFH->private_data;
408 : ZSTD_inBuffer *input = &zstdcs->input;
409 : ZSTD_outBuffer *output = &zstdcs->output;
410 : size_t res,
411 : cnt;
412 :
413 : input->src = ptr;
414 : input->size = size;
415 : input->pos = 0;
416 :
417 : if (zstdcs->cstream == NULL)
418 : {
419 : zstdcs->output.size = ZSTD_CStreamOutSize();
420 : zstdcs->output.dst = pg_malloc(zstdcs->output.size);
421 : zstdcs->output.pos = 0;
422 : zstdcs->cstream = _ZstdCStreamParams(CFH->compression_spec);
423 : if (zstdcs->cstream == NULL)
424 : pg_fatal("could not initialize compression library");
425 : }
426 :
427 : /* Consume all input, to be flushed later */
428 : while (input->pos < input->size)
429 : {
430 : res = ZSTD_compressStream2(zstdcs->cstream, output, input, ZSTD_e_continue);
431 : if (ZSTD_isError(res))
432 : pg_fatal("could not write to file: %s", ZSTD_getErrorName(res));
433 :
434 : /* Dump output buffer if full */
435 : if (output->pos >= output->size)
436 : {
437 : errno = 0;
438 : cnt = fwrite(output->dst, 1, output->pos, zstdcs->fp);
439 : if (cnt != output->pos)
440 : {
441 : errno = (errno) ? errno : ENOSPC;
442 : pg_fatal("could not write to file: %m");
443 : }
444 : output->pos = 0;
445 : }
446 : }
447 : }
448 :
449 : static int
450 : Zstd_getc(CompressFileHandle *CFH)
451 : {
452 : unsigned char ret;
453 :
454 : if (CFH->read_func(&ret, 1, CFH) != 1)
455 : pg_fatal("could not read from input file: end of file");
456 : return ret;
457 : }
458 :
459 : static char *
460 : Zstd_gets(char *buf, int len, CompressFileHandle *CFH)
461 : {
462 : int i;
463 :
464 : Assert(len > 0);
465 :
466 : /*
467 : * Read one byte at a time until newline or EOF. This is only used to read
468 : * the list of LOs, and the I/O is buffered anyway.
469 : */
470 : for (i = 0; i < len - 1; ++i)
471 : {
472 : if (Zstd_read_internal(&buf[i], 1, CFH, false) != 1)
473 : break;
474 : if (buf[i] == '\n')
475 : {
476 : ++i;
477 : break;
478 : }
479 : }
480 : buf[i] = '\0';
481 : return i > 0 ? buf : NULL;
482 : }
483 :
484 : static size_t
485 : Zstd_read(void *ptr, size_t size, CompressFileHandle *CFH)
486 : {
487 : return Zstd_read_internal(ptr, size, CFH, true);
488 : }
489 :
490 : static bool
491 : Zstd_close(CompressFileHandle *CFH)
492 : {
493 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) CFH->private_data;
494 : bool success = true;
495 :
496 : if (zstdcs->cstream)
497 : {
498 : size_t res,
499 : cnt;
500 : ZSTD_inBuffer *input = &zstdcs->input;
501 : ZSTD_outBuffer *output = &zstdcs->output;
502 :
503 : /* Loop until the compression buffers are fully consumed */
504 : for (;;)
505 : {
506 : res = ZSTD_compressStream2(zstdcs->cstream, output, input, ZSTD_e_end);
507 : if (ZSTD_isError(res))
508 : {
509 : zstdcs->zstderror = ZSTD_getErrorName(res);
510 : success = false;
511 : break;
512 : }
513 :
514 : errno = 0;
515 : cnt = fwrite(output->dst, 1, output->pos, zstdcs->fp);
516 : if (cnt != output->pos)
517 : {
518 : errno = (errno) ? errno : ENOSPC;
519 : zstdcs->zstderror = strerror(errno);
520 : success = false;
521 : break;
522 : }
523 : output->pos = 0;
524 :
525 : if (res == 0)
526 : break; /* End of frame */
527 : }
528 :
529 : ZSTD_freeCStream(zstdcs->cstream);
530 : pg_free(zstdcs->output.dst);
531 : }
532 :
533 : if (zstdcs->dstream)
534 : {
535 : if (!zstdcs->frame_finished)
536 : {
537 : errno = EIO;
538 : zstdcs->zstderror = "compressed stream is incomplete";
539 : success = false;
540 : }
541 : ZSTD_freeDStream(zstdcs->dstream);
542 : pg_free(unconstify(void *, zstdcs->input.src));
543 : }
544 :
545 : errno = 0;
546 : if (fclose(zstdcs->fp) != 0)
547 : {
548 : zstdcs->zstderror = strerror(errno);
549 : success = false;
550 : }
551 :
552 : pg_free(zstdcs);
553 : CFH->private_data = NULL;
554 : return success;
555 : }
556 :
557 : static bool
558 : Zstd_eof(CompressFileHandle *CFH)
559 : {
560 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) CFH->private_data;
561 :
562 : return feof(zstdcs->fp);
563 : }
564 :
565 : static bool
566 : Zstd_open(const char *path, int fd, const char *mode,
567 : CompressFileHandle *CFH)
568 : {
569 : FILE *fp;
570 : ZstdCompressorState *zstdcs;
571 :
572 : /*
573 : * Clear state storage to avoid having the fd point to non-NULL memory on
574 : * error return.
575 : */
576 : CFH->private_data = NULL;
577 :
578 : zstdcs = (ZstdCompressorState *) pg_malloc_extended(sizeof(*zstdcs),
579 : MCXT_ALLOC_NO_OOM | MCXT_ALLOC_ZERO);
580 : if (!zstdcs)
581 : {
582 : errno = ENOMEM;
583 : return false;
584 : }
585 :
586 : if (fd >= 0)
587 : {
588 : int dup_fd = dup(fd);
589 :
590 : if (dup_fd < 0)
591 : {
592 : pg_free(zstdcs);
593 : return false;
594 : }
595 : fp = fdopen(dup_fd, mode);
596 : if (fp == NULL)
597 : {
598 : close(dup_fd);
599 : pg_free(zstdcs);
600 : return false;
601 : }
602 : }
603 : else
604 : {
605 : fp = fopen(path, mode);
606 : if (fp == NULL)
607 : {
608 : pg_free(zstdcs);
609 : return false;
610 : }
611 : }
612 :
613 : zstdcs->fp = fp;
614 : CFH->private_data = zstdcs;
615 :
616 : return true;
617 : }
618 :
619 : static bool
620 : Zstd_open_write(const char *path, const char *mode, CompressFileHandle *CFH)
621 : {
622 : char fname[MAXPGPATH];
623 :
624 : sprintf(fname, "%s.zst", path);
625 : return CFH->open_func(fname, -1, mode, CFH);
626 : }
627 :
628 : static const char *
629 : Zstd_get_error(CompressFileHandle *CFH)
630 : {
631 : ZstdCompressorState *zstdcs = (ZstdCompressorState *) CFH->private_data;
632 :
633 : return zstdcs->zstderror;
634 : }
635 :
636 : void
637 : InitCompressFileHandleZstd(CompressFileHandle *CFH,
638 : const pg_compress_specification compression_spec)
639 : {
640 : CFH->open_func = Zstd_open;
641 : CFH->open_write_func = Zstd_open_write;
642 : CFH->read_func = Zstd_read;
643 : CFH->write_func = Zstd_write;
644 : CFH->gets_func = Zstd_gets;
645 : CFH->getc_func = Zstd_getc;
646 : CFH->close_func = Zstd_close;
647 : CFH->eof_func = Zstd_eof;
648 : CFH->get_error_func = Zstd_get_error;
649 :
650 : CFH->compression_spec = compression_spec;
651 :
652 : CFH->private_data = NULL;
653 : }
654 :
655 : #endif /* USE_ZSTD */
|
