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      164625 :               size_t size ) {
      25      164625 :   fd_zstd_rstream_t * zs = cookie;
      26      164625 :   if( zs->eof ) return 0;
      27             : 
      28      134625 :   ZSTD_outBuffer output = { buf, size, 0 };
      29      239244 :   while( output.pos < output.size ) {
      30      134628 :     if( zs->input.pos >= zs->input.size ) {
      31       30009 :       size_t read_sz = fread( zs->in_buf, 1, zs->in_buf_max, zs->file );
      32       30009 :       if( read_sz==0 ) {
      33           6 :         if( feof( zs->file ) ) break;
      34           0 :         return -1;
      35           6 :       }
      36       30003 :       zs->input.src  = zs->in_buf;
      37       30003 :       zs->input.size = read_sz;
      38       30003 :       zs->input.pos  = 0;
      39       30003 :     }
      40             : 
      41      134622 :     size_t const ret = ZSTD_decompressStream( zs->dstream, &output, &zs->input );
      42      134622 :     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      134622 :     if( output.pos>0 && ret==0 ) break;
      48      134622 :   }
      49             : 
      50      134625 :   zs->eof    = (!!feof( zs->file )) && (zs->input.pos >= zs->input.size);
      51      134625 :   zs->in_rd += output.pos;
      52      134625 :   return (ssize_t)output.pos;
      53      134625 : }
      54             : 
      55             : static int
      56             : rstream_seek( void *    cookie,
      57             :               off64_t * pos,
      58      379725 :               int       w ) {
      59      379725 :   fd_zstd_rstream_t * zs = cookie;
      60      379725 :   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      379725 :   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 :     ZSTD_DCtx_reset( zs->dstream, ZSTD_reset_session_only );
      74           0 :     *pos = 0L;
      75      379725 :   } else if( w==SEEK_CUR ) {
      76      379725 :     *pos = (long)zs->in_rd;
      77      379725 :   } else {
      78           0 :     FD_LOG_CRIT(( "unsupported seek mode" ));
      79           0 :   }
      80      379725 :   return 0;
      81      379725 : }
      82             : 
      83             : static int
      84       30006 : rstream_close( void * cookie ) {
      85       30006 :   fd_zstd_rstream_t * zs = cookie;
      86       30006 :   int close_ret = fclose( zs->file );
      87       30006 :   free( zs->in_buf );
      88       30006 :   free( zs );
      89       30006 :   return close_ret;
      90       30006 : }
      91             : 
      92             : FILE *
      93             : fd_zstd_rstream_open( FILE *         file,
      94             :                       ZSTD_DStream * dstream,
      95       30006 :                       ulong          buf_sz ) {
      96       30006 :   fd_zstd_rstream_t * zs = malloc( sizeof(fd_zstd_rstream_t) );
      97       30006 :   if( FD_UNLIKELY( !zs ) ) return NULL;
      98             : 
      99       30006 :   uchar * buf = malloc( buf_sz );
     100       30006 :   if( FD_UNLIKELY( !buf ) ) {
     101           0 :     free( zs );
     102           0 :     return NULL;
     103           0 :   }
     104             : 
     105       30006 :   zs->file    = file;
     106       30006 :   zs->dstream = dstream;
     107             : 
     108       30006 :   size_t const init_ret = ZSTD_DCtx_reset( dstream, ZSTD_reset_session_only );
     109       30006 :   if( FD_UNLIKELY( ZSTD_isError( init_ret ) ) ) {
     110           0 :     FD_LOG_WARNING(( "ZSTD_DCtx_reset failed: %s", ZSTD_getErrorName( init_ret ) ));
     111           0 :     free( zs );
     112           0 :     return NULL;
     113           0 :   }
     114             : 
     115       30006 :   zs->in_rd      = 0UL;
     116       30006 :   zs->in_buf     = buf;
     117       30006 :   zs->in_buf_max = buf_sz;
     118       30006 :   zs->input.src  = zs->in_buf;
     119       30006 :   zs->input.size = 0;
     120       30006 :   zs->input.pos  = 0;
     121       30006 :   zs->eof        = 0;
     122             : 
     123       30006 :   static cookie_io_functions_t const io_funcs = {
     124       30006 :     .read  = rstream_read,
     125       30006 :     .write = NULL,
     126       30006 :     .seek  = rstream_seek,
     127       30006 :     .close = rstream_close
     128       30006 :   };
     129       30006 :   return fopencookie( zs, "rb", io_funcs );
     130       30006 : }
     131             : 
     132             : struct fd_zstd_wstream {
     133             :   FILE *         file;
     134             :   ZSTD_CStream * cstream;
     135             :   uchar *        out_buf;
     136             :   ulong          out_buf_max;
     137             :   ulong          wr_cnt;
     138             : };
     139             : 
     140             : typedef struct fd_zstd_wstream fd_zstd_wstream_t;
     141             : 
     142             : static ssize_t
     143             : wstream_write( void *       cookie,
     144             :                char const * buf,
     145           0 :                size_t       size ) {
     146           0 :   fd_zstd_wstream_t * zs = cookie;
     147           0 :   ZSTD_inBuffer input = { buf, size, 0 };
     148           0 :   while( input.pos < input.size ) {
     149           0 :     ZSTD_outBuffer output = { zs->out_buf, zs->out_buf_max, 0 };
     150           0 :     size_t const ret = ZSTD_compressStream( zs->cstream, &output, &input );
     151           0 :     if( FD_UNLIKELY( ZSTD_isError( ret ) ) ) {
     152           0 :       FD_LOG_WARNING(( "ZSTD_compressStream failed (%u-%s)", ZSTD_getErrorCode( ret ), ZSTD_getErrorName( ret ) ));
     153           0 :       return -1;
     154           0 :     }
     155           0 :     if( output.pos > 0 ) {
     156           0 :       size_t written = fwrite( zs->out_buf, 1, output.pos, zs->file );
     157           0 :       if( FD_UNLIKELY( written != output.pos ) ) return -1;
     158           0 :     }
     159           0 :   }
     160           0 :   zs->wr_cnt += size;
     161           0 :   return (ssize_t)size;
     162           0 : }
     163             : 
     164             : static int
     165           0 : wstream_close( void * cookie ) {
     166           0 :   fd_zstd_wstream_t * zs = cookie;
     167             : 
     168           0 :   ZSTD_outBuffer output = { zs->out_buf, zs->out_buf_max, 0 };
     169           0 :   size_t const ret = ZSTD_endStream( zs->cstream, &output );
     170           0 :   if( FD_UNLIKELY( ZSTD_isError( ret ) ) ) {
     171           0 :     FD_LOG_WARNING(( "ZSTD_endStream failed (%u-%s)", ZSTD_getErrorCode( ret ), ZSTD_getErrorName( ret ) ));
     172           0 :     return -1;
     173           0 :   }
     174           0 :   if( output.pos > 0 ) {
     175           0 :     size_t written = fwrite( zs->out_buf, 1, output.pos, zs->file );
     176           0 :     if( FD_UNLIKELY( written != output.pos ) ) return -1;
     177           0 :   }
     178             : 
     179           0 :   ZSTD_freeCStream( zs->cstream );
     180           0 :   int close_ret = fclose( zs->file );
     181           0 :   free( zs->out_buf );
     182           0 :   free( zs );
     183           0 :   return close_ret;
     184           0 : }
     185             : 
     186             : static int
     187             : wstream_seek( void *    cookie,
     188             :               off64_t * pos,
     189           0 :               int       w ) {
     190           0 :   fd_zstd_wstream_t * zs = cookie;
     191           0 :   if( FD_UNLIKELY( !( w==SEEK_CUR && *pos==0 ) ) ) {
     192           0 :     FD_LOG_WARNING(( "Attempted to seek in fd_libc_zstd_wstream handle" ));
     193           0 :     return -1;
     194           0 :   }
     195           0 :   *pos = (long)zs->wr_cnt;
     196           0 :   return 0;
     197           0 : }
     198             : 
     199             : FILE *
     200             : fd_zstd_wstream_open( FILE * file,
     201             :                       int    level,
     202           0 :                       ulong  buf_sz ) {
     203           0 :   fd_zstd_wstream_t * zs = malloc( sizeof(fd_zstd_wstream_t) );
     204           0 :   if( FD_UNLIKELY( !zs ) ) return NULL;
     205             : 
     206           0 :   uchar * buf = malloc( buf_sz );
     207           0 :   if( FD_UNLIKELY( !buf ) ) {
     208           0 :     free( zs );
     209           0 :     return NULL;
     210           0 :   }
     211             : 
     212           0 :   zs->file        = file;
     213           0 :   zs->cstream     = ZSTD_createCStream();
     214           0 :   zs->out_buf     = buf;
     215           0 :   zs->out_buf_max = buf_sz;
     216           0 :   zs->wr_cnt      = 0UL;
     217           0 :   if( FD_UNLIKELY( !zs->cstream ) ) {
     218           0 :     free( zs );
     219           0 :     return NULL;
     220           0 :   }
     221             : 
     222           0 :   size_t const init_ret = ZSTD_initCStream( zs->cstream, level );
     223           0 :   if( FD_UNLIKELY( ZSTD_isError( init_ret ) ) ) {
     224           0 :     ZSTD_freeCStream( zs->cstream );
     225           0 :     free( zs );
     226           0 :     return NULL;
     227           0 :   }
     228             : 
     229           0 :   static cookie_io_functions_t const io_funcs = {
     230             :     .read  = NULL,
     231           0 :     .write = wstream_write,
     232           0 :     .seek  = wstream_seek,
     233           0 :     .close = wstream_close
     234           0 :   };
     235           0 :   return fopencookie( zs, "wb", io_funcs );
     236           0 : }