pgzstd/zstd.c at master · innogames/pgzstd · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
#include <postgres.h>
#if PG_VERSION_NUM >= 160000
#include <varatt.h>
#endif
#include <fmgr.h>
#include <utils/memutils.h>
#include <zstd.h>

#ifdef PG_MODULE_MAGIC
PG_MODULE_MAGIC;
#endif

#ifndef PG_INT32_MAX
#define PG_INT32_MAX 0x7fffffff
#endif

#define DEFAULT_COMPRESSION_LEVEL 3

static ZSTD_CCtx *cctx;
static ZSTD_DCtx *dctx;

void _PG_init(void);
Datum compress(PG_FUNCTION_ARGS);
Datum decompress(PG_FUNCTION_ARGS);
Datum length(PG_FUNCTION_ARGS);

void _PG_init(void)
{
    cctx = ZSTD_createCCtx();
    if (!cctx)
        elog(FATAL, "ZSTD_createCCtx failed");

    dctx = ZSTD_createDCtx();
    if (!dctx)
        elog(FATAL, "ZSTD_createDCtx failed");
}

PG_FUNCTION_INFO_V1(compress);
Datum compress(PG_FUNCTION_ARGS)
{
    bytea *in, *out;
    char *dict = NULL;
    int32 level;
    size_t in_len, out_len, dict_len = 0;

    if (PG_ARGISNULL(0))
        PG_RETURN_NULL();

    in = PG_GETARG_BYTEA_P(0);
    in_len = VARSIZE(in) - VARHDRSZ;

    level = PG_ARGISNULL(2) ? DEFAULT_COMPRESSION_LEVEL : PG_GETARG_INT32(2);

    if (!PG_ARGISNULL(1))
    {
        bytea *d = PG_GETARG_BYTEA_P(1);
        dict = VARDATA(d);
        dict_len = VARSIZE(d) - VARHDRSZ;
    }

    out_len = ZSTD_compressBound(in_len);
    out = palloc(out_len + VARHDRSZ);

    out_len = ZSTD_compress_usingDict(cctx, VARDATA(out), out_len, VARDATA(in), in_len, dict, dict_len, level);

    if (ZSTD_isError(out_len))
        elog(ERROR, "ZSTD_compress_usingDict failed: %s", ZSTD_getErrorName(out_len));

    out = repalloc(out, out_len + VARHDRSZ);

    SET_VARSIZE(out, out_len + VARHDRSZ);
    PG_RETURN_BYTEA_P(out);
}

PG_FUNCTION_INFO_V1(decompress);
Datum decompress(PG_FUNCTION_ARGS)
{
    bytea *in, *out;
    char *dict = NULL;
    size_t in_len, out_len, dict_len = 0;

    if (PG_ARGISNULL(0))
        PG_RETURN_NULL();

    in = PG_GETARG_BYTEA_P(0);
    in_len = VARSIZE(in) - VARHDRSZ;

    if (!PG_ARGISNULL(1))
    {
        bytea *d = PG_GETARG_BYTEA_P(1);
        dict = VARDATA(d);
        dict_len = VARSIZE(d) - VARHDRSZ;
    }

#ifdef ZSTD_CONTENTSIZE_UNKNOWN
    out_len = ZSTD_getFrameContentSize(VARDATA(in), in_len);
    if (out_len == ZSTD_CONTENTSIZE_UNKNOWN)
        elog(ERROR, "ZSTD_getFrameContentSize returned unknown");
    else if (out_len == ZSTD_CONTENTSIZE_ERROR)
        elog(ERROR, "ZSTD_getFrameContentSize failed");
#else
    out_len = ZSTD_getDecompressedSize(VARDATA(in), in_len);
#endif
    /*
     * The decompressed size is taken from the (untrusted) frame header, so
     * reject anything larger than a palloc can satisfy before attempting the
     * allocation. This turns a hostile/corrupt frame into a clear error rather
     * than an opaque "invalid memory alloc request size" failure.
     */
    if (out_len > MaxAllocSize - VARHDRSZ)
        elog(ERROR, "zstd frame claims decompressed size %zu, which exceeds the maximum of %zu",
             out_len, (size_t) (MaxAllocSize - VARHDRSZ));

    out = palloc(out_len + VARHDRSZ);

    out_len = ZSTD_decompress_usingDict(dctx, VARDATA(out), out_len, VARDATA(in), in_len, dict, dict_len);

    if (ZSTD_isError(out_len))
        elog(ERROR, "ZSTD_decompress_usingDict failed: %s", ZSTD_getErrorName(out_len));

    SET_VARSIZE(out, out_len + VARHDRSZ);
    PG_RETURN_BYTEA_P(out);
}

PG_FUNCTION_INFO_V1(length);
Datum length(PG_FUNCTION_ARGS)
{
    bytea *in;
    size_t in_len, out_len;

    if (PG_ARGISNULL(0))
        PG_RETURN_NULL();

    in = PG_GETARG_BYTEA_P(0);
    in_len = VARSIZE(in) - VARHDRSZ;

#ifdef ZSTD_CONTENTSIZE_UNKNOWN
    out_len = ZSTD_getFrameContentSize(VARDATA(in), in_len);
    if (out_len == ZSTD_CONTENTSIZE_UNKNOWN)
        PG_RETURN_NULL();
    else if (out_len == ZSTD_CONTENTSIZE_ERROR)
        elog(ERROR, "ZSTD_getFrameContentSize failed");
#else
    out_len = ZSTD_getDecompressedSize(VARDATA(in), in_len);
#endif
    if (out_len > PG_INT32_MAX)
        elog(ERROR, "ZSTD_getDecompressedSize returned value greater than PG_INT32_MAX");

    PG_RETURN_INT32((int32) out_len);
}