Line data Source code
1 : #define _GNU_SOURCE
2 : #define _FILE_OFFSET_BITS 64
3 : #include <stdio.h>
4 : #include <stdlib.h>
5 : #include <zstd.h>
6 : #include "../../util/log/fd_log.h"
7 : #include "fd_libc_zstd.h"
8 :
9 : struct fd_zstd_rstream {
10 : FILE * file;
11 : ZSTD_DStream * dstream;
12 : ulong in_rd;
13 : uchar * in_buf;
14 : ulong in_buf_max;
15 : ZSTD_inBuffer input;
16 : uint eof : 1;
17 : };
18 :
19 : typedef struct fd_zstd_rstream fd_zstd_rstream_t;
20 :
21 : static ssize_t
22 : rstream_read( void * cookie,
23 : char * buf,
24 0 : size_t size ) {
25 0 : fd_zstd_rstream_t * zs = cookie;
26 0 : if( zs->eof ) return 0;
27 :
28 0 : ZSTD_outBuffer output = { buf, size, 0 };
29 0 : while( output.pos < output.size ) {
30 0 : if( zs->input.pos >= zs->input.size ) {
31 0 : size_t read_sz = fread( zs->in_buf, 1, zs->in_buf_max, zs->file );
32 0 : if( read_sz==0 ) {
33 0 : if( feof( zs->file ) ) break;
34 0 : return -1;
35 0 : }
36 0 : zs->input.src = zs->in_buf;
37 0 : zs->input.size = read_sz;
38 0 : zs->input.pos = 0;
39 0 : }
40 :
41 0 : size_t const ret = ZSTD_decompressStream( zs->dstream, &output, &zs->input );
42 0 : if( FD_UNLIKELY( ZSTD_isError( ret ) ) ) {
43 0 : FD_LOG_WARNING(( "ZSTD_decompressStream failed (%u-%s)", ZSTD_getErrorCode( ret ), ZSTD_getErrorName( ret ) ));
44 0 : return -1;
45 0 : }
46 :
47 0 : if( output.pos>0 && ret==0 ) break;
48 0 : }
49 :
50 0 : zs->eof = (!!feof( zs->file )) && (zs->input.pos >= zs->input.size);
51 0 : zs->in_rd += output.pos;
52 0 : return (ssize_t)output.pos;
53 0 : }
54 :
55 : static int
56 : rstream_seek( void * cookie,
57 : off64_t * pos,
58 0 : int w ) {
59 0 : fd_zstd_rstream_t * zs = cookie;
60 0 : if( FD_UNLIKELY( *pos ) ) {
61 0 : FD_LOG_WARNING(( "Invalid seek(%ld,%i) on fd_libc_zstd_rstream handle", *pos, w ));
62 0 : return -1;
63 0 : }
64 0 : if( w==SEEK_SET ) {
65 0 : if( FD_UNLIKELY( 0!=fseek( zs->file, 0L, SEEK_SET ) ) ) {
66 0 : return -1;
67 0 : }
68 0 : zs->input.src = zs->in_buf;
69 0 : zs->input.size = 0;
70 0 : zs->input.pos = 0;
71 0 : zs->in_rd = 0UL;
72 0 : zs->eof = 0;
73 0 : size_t const reset_ret = ZSTD_DCtx_reset( zs->dstream, ZSTD_reset_session_only );
74 0 : if( FD_UNLIKELY( ZSTD_isError( reset_ret ) ) ) return -1;
75 0 : *pos = 0L;
76 0 : } else if( w==SEEK_CUR ) {
77 0 : *pos = (long)zs->in_rd;
78 0 : } else {
79 0 : FD_LOG_CRIT(( "unsupported seek mode" ));
80 0 : }
81 0 : return 0;
82 0 : }
83 :
84 : static int
85 0 : rstream_close( void * cookie ) {
86 0 : fd_zstd_rstream_t * zs = cookie;
87 0 : int close_ret = fclose( zs->file );
88 0 : free( zs->in_buf );
89 0 : free( zs );
90 0 : return close_ret;
91 0 : }
92 :
93 : FILE *
94 : fd_zstd_rstream_open( FILE * file,
95 : ZSTD_DStream * dstream,
96 0 : ulong buf_sz ) {
97 0 : if( FD_UNLIKELY( !buf_sz ) ) return NULL;
98 :
99 0 : fd_zstd_rstream_t * zs = malloc( sizeof(fd_zstd_rstream_t) );
100 0 : if( FD_UNLIKELY( !zs ) ) return NULL;
101 :
102 0 : uchar * buf = malloc( buf_sz );
103 0 : if( FD_UNLIKELY( !buf ) ) {
104 0 : free( zs );
105 0 : return NULL;
106 0 : }
107 :
108 0 : zs->file = file;
109 0 : zs->dstream = dstream;
110 :
111 0 : size_t const init_ret = ZSTD_DCtx_reset( dstream, ZSTD_reset_session_only );
112 0 : if( FD_UNLIKELY( ZSTD_isError( init_ret ) ) ) {
113 0 : FD_LOG_WARNING(( "ZSTD_DCtx_reset failed: %s", ZSTD_getErrorName( init_ret ) ));
114 0 : free( zs );
115 0 : free( buf );
116 0 : return NULL;
117 0 : }
118 :
119 0 : zs->in_rd = 0UL;
120 0 : zs->in_buf = buf;
121 0 : zs->in_buf_max = buf_sz;
122 0 : zs->input.src = zs->in_buf;
123 0 : zs->input.size = 0;
124 0 : zs->input.pos = 0;
125 0 : zs->eof = 0;
126 :
127 0 : static cookie_io_functions_t const io_funcs = {
128 0 : .read = rstream_read,
129 0 : .write = NULL,
130 0 : .seek = rstream_seek,
131 0 : .close = rstream_close
132 0 : };
133 0 : FILE * ret = fopencookie( zs, "rb", io_funcs );
134 0 : if( FD_UNLIKELY( !ret ) ) {
135 0 : free( buf );
136 0 : free( zs );
137 0 : }
138 0 : return ret;
139 0 : }
140 :
141 : struct fd_zstd_wstream {
142 : FILE * file;
143 : ZSTD_CStream * cstream;
144 : uchar * out_buf;
145 : ulong out_buf_max;
146 : ulong wr_cnt;
147 : };
148 :
149 : typedef struct fd_zstd_wstream fd_zstd_wstream_t;
150 :
151 : static ssize_t
152 : wstream_write( void * cookie,
153 : char const * buf,
154 0 : size_t size ) {
155 0 : fd_zstd_wstream_t * zs = cookie;
156 0 : ZSTD_inBuffer input = { buf, size, 0 };
157 0 : while( input.pos < input.size ) {
158 0 : ZSTD_outBuffer output = { zs->out_buf, zs->out_buf_max, 0 };
159 0 : size_t const ret = ZSTD_compressStream( zs->cstream, &output, &input );
160 0 : if( FD_UNLIKELY( ZSTD_isError( ret ) ) ) {
161 0 : FD_LOG_WARNING(( "ZSTD_compressStream failed (%u-%s)", ZSTD_getErrorCode( ret ), ZSTD_getErrorName( ret ) ));
162 0 : return -1;
163 0 : }
164 0 : if( output.pos > 0 ) {
165 0 : size_t written = fwrite( zs->out_buf, 1, output.pos, zs->file );
166 0 : if( FD_UNLIKELY( written != output.pos ) ) return -1;
167 0 : }
168 0 : }
169 0 : zs->wr_cnt += size;
170 0 : return (ssize_t)size;
171 0 : }
172 :
173 : static int
174 0 : wstream_close( void * cookie ) {
175 0 : fd_zstd_wstream_t * zs = cookie;
176 0 : int ret_val = 0;
177 :
178 0 : size_t remaining;
179 0 : do {
180 0 : ZSTD_outBuffer output = { zs->out_buf, zs->out_buf_max, 0 };
181 0 : remaining = ZSTD_endStream( zs->cstream, &output );
182 0 : if( FD_UNLIKELY( ZSTD_isError( remaining ) ) ) {
183 0 : FD_LOG_WARNING(( "ZSTD_endStream failed (%u-%s)", ZSTD_getErrorCode( remaining ), ZSTD_getErrorName( remaining ) ));
184 0 : ret_val = -1;
185 0 : goto cleanup;
186 0 : }
187 0 : if( output.pos > 0 ) {
188 0 : size_t written = fwrite( zs->out_buf, 1, output.pos, zs->file );
189 0 : if( FD_UNLIKELY( written != output.pos ) ) {
190 0 : ret_val = -1;
191 0 : goto cleanup;
192 0 : }
193 0 : }
194 0 : } while( remaining > 0 );
195 :
196 0 : cleanup:
197 0 : ZSTD_freeCStream( zs->cstream );
198 0 : int close_ret = fclose( zs->file );
199 0 : if( FD_LIKELY( ret_val != -1 ) ) ret_val = close_ret;
200 0 : free( zs->out_buf );
201 0 : free( zs );
202 0 : return ret_val;
203 0 : }
204 :
205 : static int
206 : wstream_seek( void * cookie,
207 : off64_t * pos,
208 0 : int w ) {
209 0 : fd_zstd_wstream_t * zs = cookie;
210 0 : if( FD_UNLIKELY( !( w==SEEK_CUR && *pos==0 ) ) ) {
211 0 : FD_LOG_WARNING(( "Attempted to seek in fd_libc_zstd_wstream handle" ));
212 0 : return -1;
213 0 : }
214 0 : *pos = (long)zs->wr_cnt;
215 0 : return 0;
216 0 : }
217 :
218 : FILE *
219 : fd_zstd_wstream_open( FILE * file,
220 : int level,
221 0 : ulong buf_sz ) {
222 0 : if( FD_UNLIKELY( !buf_sz ) ) return NULL;
223 :
224 0 : fd_zstd_wstream_t * zs = malloc( sizeof(fd_zstd_wstream_t) );
225 0 : if( FD_UNLIKELY( !zs ) ) return NULL;
226 :
227 0 : uchar * buf = malloc( buf_sz );
228 0 : if( FD_UNLIKELY( !buf ) ) {
229 0 : free( zs );
230 0 : return NULL;
231 0 : }
232 :
233 0 : zs->file = file;
234 0 : zs->cstream = ZSTD_createCStream();
235 0 : zs->out_buf = buf;
236 0 : zs->out_buf_max = buf_sz;
237 0 : zs->wr_cnt = 0UL;
238 0 : if( FD_UNLIKELY( !zs->cstream ) ) {
239 0 : free( zs );
240 0 : free( buf );
241 0 : return NULL;
242 0 : }
243 :
244 0 : size_t const init_ret = ZSTD_initCStream( zs->cstream, level );
245 0 : if( FD_UNLIKELY( ZSTD_isError( init_ret ) ) ) {
246 0 : ZSTD_freeCStream( zs->cstream );
247 0 : free( zs );
248 0 : free( buf );
249 0 : return NULL;
250 0 : }
251 :
252 0 : static cookie_io_functions_t const io_funcs = {
253 0 : .read = NULL,
254 0 : .write = wstream_write,
255 0 : .seek = wstream_seek,
256 0 : .close = wstream_close
257 0 : };
258 0 : FILE * ret = fopencookie( zs, "wb", io_funcs );
259 0 : if( FD_UNLIKELY( !ret ) ) {
260 0 : ZSTD_freeCStream( zs->cstream );
261 0 : free( buf );
262 0 : free( zs );
263 0 : }
264 0 : return ret;
265 0 : }
|
