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             :   uint           mid_frame : 1;
      18             : };
      19             : 
      20             : typedef struct fd_zstd_rstream fd_zstd_rstream_t;
      21             : 
      22             : static ssize_t
      23             : rstream_read( void * cookie,
      24             :               char * buf,
      25           0 :               size_t size ) {
      26           0 :   fd_zstd_rstream_t * zs = cookie;
      27           0 :   if( zs->eof ) return 0;
      28             : 
      29           0 :   ZSTD_outBuffer output = { buf, size, 0 };
      30           0 :   while( output.pos < output.size ) {
      31           0 :     if( zs->input.pos >= zs->input.size ) {
      32           0 :       size_t read_sz = fread( zs->in_buf, 1, zs->in_buf_max, zs->file );
      33           0 :       if( read_sz==0 ) {
      34           0 :         if( !feof( zs->file ) ) return -1;
      35           0 :         if( FD_UNLIKELY( zs->mid_frame && !output.pos ) ) {
      36           0 :           FD_LOG_WARNING(( "truncated Zstandard frame" ));
      37           0 :           return -1;
      38           0 :         }
      39           0 :         break;
      40           0 :       }
      41           0 :       zs->input.src  = zs->in_buf;
      42           0 :       zs->input.size = read_sz;
      43           0 :       zs->input.pos  = 0;
      44           0 :     }
      45             : 
      46           0 :     size_t const ret = ZSTD_decompressStream( zs->dstream, &output, &zs->input );
      47           0 :     if( FD_UNLIKELY( ZSTD_isError( ret ) ) ) {
      48           0 :       FD_LOG_WARNING(( "ZSTD_decompressStream failed (%u-%s)", ZSTD_getErrorCode( ret ), ZSTD_getErrorName( ret ) ));
      49           0 :       return -1;
      50           0 :     }
      51           0 :     zs->mid_frame = ret!=0;
      52             : 
      53           0 :     if( output.pos>0 && ret==0 ) break;
      54           0 :   }
      55             : 
      56           0 :   zs->eof    = (!!feof( zs->file )) && (zs->input.pos >= zs->input.size) && !zs->mid_frame;
      57           0 :   zs->in_rd += output.pos;
      58           0 :   return (ssize_t)output.pos;
      59           0 : }
      60             : 
      61             : static int
      62             : rstream_seek( void * cookie,
      63             :               long * pos,
      64           0 :               int    w ) {
      65           0 :   fd_zstd_rstream_t * zs = cookie;
      66           0 :   if( FD_UNLIKELY( *pos ) ) {
      67           0 :     FD_LOG_WARNING(( "Invalid seek(%ld,%i) on fd_libc_zstd_rstream handle", *pos, w ));
      68           0 :     return -1;
      69           0 :   }
      70           0 :   if( w==SEEK_SET ) {
      71           0 :     if( FD_UNLIKELY( 0!=fseek( zs->file, 0L, SEEK_SET ) ) ) {
      72           0 :       return -1;
      73           0 :     }
      74           0 :     zs->input.src  = zs->in_buf;
      75           0 :     zs->input.size = 0;
      76           0 :     zs->input.pos  = 0;
      77           0 :     zs->in_rd      = 0UL;
      78           0 :     zs->eof        = 0;
      79           0 :     zs->mid_frame  = 0;
      80           0 :     size_t const reset_ret = ZSTD_DCtx_reset( zs->dstream, ZSTD_reset_session_only );
      81           0 :     if( FD_UNLIKELY( ZSTD_isError( reset_ret ) ) ) return -1;
      82           0 :     *pos = 0L;
      83           0 :   } else if( w==SEEK_CUR ) {
      84           0 :     *pos = (long)zs->in_rd;
      85           0 :   } else {
      86           0 :     FD_LOG_CRIT(( "unsupported seek mode" ));
      87           0 :   }
      88           0 :   return 0;
      89           0 : }
      90             : 
      91             : static int
      92           0 : rstream_close( void * cookie ) {
      93           0 :   fd_zstd_rstream_t * zs = cookie;
      94           0 :   int close_ret = fclose( zs->file );
      95           0 :   free( zs->in_buf );
      96           0 :   free( zs );
      97           0 :   return close_ret;
      98           0 : }
      99             : 
     100             : FILE *
     101             : fd_zstd_rstream_open( FILE *         file,
     102             :                       ZSTD_DStream * dstream,
     103           0 :                       ulong          buf_sz ) {
     104           0 :   if( FD_UNLIKELY( !buf_sz ) ) return NULL;
     105             : 
     106           0 :   fd_zstd_rstream_t * zs = malloc( sizeof(fd_zstd_rstream_t) );
     107           0 :   if( FD_UNLIKELY( !zs ) ) return NULL;
     108             : 
     109           0 :   uchar * buf = malloc( buf_sz );
     110           0 :   if( FD_UNLIKELY( !buf ) ) {
     111           0 :     free( zs );
     112           0 :     return NULL;
     113           0 :   }
     114             : 
     115           0 :   zs->file    = file;
     116           0 :   zs->dstream = dstream;
     117             : 
     118           0 :   size_t const init_ret = ZSTD_DCtx_reset( dstream, ZSTD_reset_session_only );
     119           0 :   if( FD_UNLIKELY( ZSTD_isError( init_ret ) ) ) {
     120           0 :     FD_LOG_WARNING(( "ZSTD_DCtx_reset failed: %s", ZSTD_getErrorName( init_ret ) ));
     121           0 :     free( zs );
     122           0 :     free( buf );
     123           0 :     return NULL;
     124           0 :   }
     125             : 
     126           0 :   zs->in_rd      = 0UL;
     127           0 :   zs->in_buf     = buf;
     128           0 :   zs->in_buf_max = buf_sz;
     129           0 :   zs->input.src  = zs->in_buf;
     130           0 :   zs->input.size = 0;
     131           0 :   zs->input.pos  = 0;
     132           0 :   zs->eof        = 0;
     133           0 :   zs->mid_frame  = 0;
     134             : 
     135           0 :   static cookie_io_functions_t const io_funcs = {
     136           0 :     .read  = rstream_read,
     137           0 :     .write = NULL,
     138           0 :     .seek  = rstream_seek,
     139           0 :     .close = rstream_close
     140           0 :   };
     141           0 :   FILE * ret = fopencookie( zs, "rb", io_funcs );
     142           0 :   if( FD_UNLIKELY( !ret ) ) {
     143           0 :     free( buf );
     144           0 :     free( zs );
     145           0 :   }
     146           0 :   return ret;
     147           0 : }