text: /***********************************************************************
   *
   * Project:  PgSQL gzip/gunzip
   * Purpose:  Main file.
   *
   ***********************************************************************
   * Copyright 2019 Paul Ramsey <pramsey@cleverelephant.ca>
   *
   * Permission is hereby granted, free of charge, to any person obtaining a
   * copy of this software and associated documentation files (the
   * "Software"), to deal in the Software without restriction, including
   * without limitation the rights to use, copy, modify, merge, publish,
   * distribute, sublicense, and/or sell copies of the Software, and to
   * permit persons to whom the Software is furnished to do so, subject to
   * the following conditions:
   *
   * The above copyright notice and this permission notice shall be included
   * in all copies or substantial portions of the Software.
   *
   * THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
   * OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
   * MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
   * IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
   * CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
   * TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
   * SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
   *
   ***********************************************************************/
  
  /* Constants */
  #define ZCHUNK 262144 /* 256K */
  
  /* System */
  #include <assert.h>
  #include <stdio.h>
  #include <stdlib.h>
  #include <string.h>
  
  /* PostgreSQL */
  #include <postgres.h>
  #include <fmgr.h>
  #include <funcapi.h>
  
  /* LibZ */
  #include <zlib.h>
  
  /* Set up PgSQL */
  PG_MODULE_MAGIC;
  
  /**
  * Wrap palloc in a signature that matches what zalloc expects
  */
  static void*
  pg_gzip_alloc(void* opaque, unsigned int items, unsigned int itemsize)
  {
  	return palloc(items * itemsize);
  }
  
  /**
  * Wrap pfree in a signature that matches what zfree expects
  */
  static void
  pg_gzip_free(void* opaque, void* ptr)
  {
  	pfree(ptr);
  	return;
  }
  
  /* Zlib defines */
  #define WINDOW_BITS 15
  #define ENABLE_ZLIB_GZIP 32
  #define GZIP_ENCODING 16
  
  /**
  * gzip an uncompressed bytea
  */
  Datum pg_gzip(PG_FUNCTION_ARGS);
  PG_FUNCTION_INFO_V1(pg_gzip);
  Datum pg_gzip(PG_FUNCTION_ARGS)
  {
  	StringInfoData si;
  	int zs_rv;
  	z_stream zs;
  	uint8 out[ZCHUNK];
  	bytea* compressed;
  
  	bytea* uncompressed = PG_GETARG_BYTEA_P(0);
  	int32 compression_level = PG_GETARG_INT32(1);
  	uint8* in = (uint8*)(VARDATA(uncompressed));
  	size_t in_size = VARSIZE_ANY_EXHDR(uncompressed);
  
  	/* compression level -1 is default best effort (approx 6) */
  	/* level 0 is no compression, 1-9 are lowest to highest */
  	if (compression_level < -1 || compression_level > 9)
  		elog(ERROR, "invalid compression level: %d", compression_level);
  
  	/* Prepare the z_stream state */
  	zs.zalloc = pg_gzip_alloc;
  	zs.zfree = pg_gzip_free;
  	zs.opaque = Z_NULL;
  	zs.next_in = in;
  	zs.avail_in = in_size;
  
  	if (deflateInit2(&zs,
  	                 compression_level, Z_DEFLATED,
  	                 WINDOW_BITS|GZIP_ENCODING, /* Magic to initialize in gzip mode */
  	                 8, Z_DEFAULT_STRATEGY) != Z_OK)
  		elog(ERROR, "failed to deflateInit2");
  
  	zs.next_out = out;
  	zs.avail_out = ZCHUNK;
  
  	/* Compress until deflate stops returning output */
  	initStringInfo(&si);
  	zs_rv = Z_OK;
  	while (zs_rv == Z_OK)
  	{
  		if (zs.avail_out == 0)
  		{
  			/* build up output in stringinfo */
  			appendBinaryStringInfo(&si, (char*)out, ZCHUNK);
  			zs.avail_out = ZCHUNK;
  			zs.next_out = out;
  		}
  		zs_rv = deflate(&zs, Z_FINISH);
  	}
  	if (zs_rv != Z_STREAM_END)
  		elog(ERROR, "compression error: %s", zs.msg ? zs.msg : "");
  
  	appendBinaryStringInfo(&si, (char*)out, ZCHUNK - zs.avail_out);
  
  	/* Construct output bytea */
  	compressed = palloc(si.len + VARHDRSZ);
  	memcpy(VARDATA(compressed), si.data, si.len);
  	SET_VARSIZE(compressed, si.len + VARHDRSZ);
  	PG_FREE_IF_COPY(uncompressed, 0);
  	PG_RETURN_POINTER(compressed);
  }
  
  
  Datum pg_gunzip(PG_FUNCTION_ARGS);
  PG_FUNCTION_INFO_V1(pg_gunzip);
  Datum pg_gunzip(PG_FUNCTION_ARGS)
  {
  	StringInfoData si;
  	int zs_rv;
  	z_stream zs;
  	uint8 out[ZCHUNK];
  	bytea* uncompressed;
  
  	bytea* compressed = PG_GETARG_BYTEA_P(0);
  	uint8* in = (uint8*)(VARDATA(compressed));
  	size_t in_size = VARSIZE_ANY_EXHDR(compressed);
  
  	/* Prepare the z_stream state */
  	zs.zalloc = pg_gzip_alloc;
  	zs.zfree = pg_gzip_free;
  	zs.opaque = Z_NULL;
  	/* Magic to initialize in gzip mode */
  	if (inflateInit2(&zs, WINDOW_BITS|ENABLE_ZLIB_GZIP) != Z_OK)
  		elog(ERROR, "failed to inflateInit");
  
  	/* Point z_stream to input and output buffers */
  	zs.next_in = in;
  	zs.avail_in = in_size;
  	zs.next_out = out;
  	zs.avail_out = ZCHUNK;
  
  	/* Decompress until inflate stops returning output */
  	initStringInfo(&si);
  	zs_rv = Z_OK;
  	while (zs_rv == Z_OK)
  	{
  		if (zs.avail_out == 0)
  		{
  			/* build up output in stringinfo */
  			appendBinaryStringInfo(&si, (char*)out, ZCHUNK);
  			zs.avail_out = ZCHUNK;
  			zs.next_out = out;
  		}
  		zs_rv = inflate(&zs, Z_SYNC_FLUSH);
  	}
  
  	if (zs_rv != Z_STREAM_END)
  		elog(ERROR, "decompression error: %s", zs.msg ? zs.msg : "");
  
  	appendBinaryStringInfo(&si, (char*)out, ZCHUNK - zs.avail_out);
  
  	/* Construct output bytea */
  	uncompressed = palloc(si.len + VARHDRSZ);
  	memcpy(VARDATA(uncompressed), si.data, si.len);
  	SET_VARSIZE(uncompressed, si.len + VARHDRSZ);
  	PG_FREE_IF_COPY(compressed, 0);
  	PG_RETURN_POINTER(uncompressed);
  }
  