/*-------------------------------------------------------------------------

 *

 * basebackup_gzip.c

 *    Basebackup sink implementing gzip compression.

 *

 * Portions Copyright (c) 2010-2026, PostgreSQL Global Development Group

 *

 * IDENTIFICATION

 *    src/backend/backup/basebackup_gzip.c

 *

 *-------------------------------------------------------------------------

 */

#include "postgres.h"

#ifdef HAVE_LIBZ

#include <zlib.h>

#endif

#include "backup/basebackup_sink.h"

#ifdef HAVE_LIBZ

typedef struct bbsink_gzip

{

  /* Common information for all types of sink. */

  bbsink    base;

  /* Compression level. */

  int     compresslevel;

  /* Compressed data stream. */

  z_stream  zstream;

  /* Number of bytes staged in output buffer. */

  size_t    bytes_written;

  /* Has the zstream been initialized? */

  bool    zstream_initialized;

} bbsink_gzip;

static void bbsink_gzip_begin_backup(bbsink *sink);

static void bbsink_gzip_begin_archive(bbsink *sink, const char *archive_name);

static void bbsink_gzip_archive_contents(bbsink *sink, size_t len);

static void bbsink_gzip_manifest_contents(bbsink *sink, size_t len);

static void bbsink_gzip_end_archive(bbsink *sink);

static void bbsink_gzip_cleanup(bbsink *sink);

static void *gzip_palloc(void *opaque, unsigned items, unsigned size);

static void gzip_pfree(void *opaque, void *address);

static const bbsink_ops bbsink_gzip_ops = {

  .begin_backup = bbsink_gzip_begin_backup,

  .begin_archive = bbsink_gzip_begin_archive,

  .archive_contents = bbsink_gzip_archive_contents,

  .end_archive = bbsink_gzip_end_archive,

  .begin_manifest = bbsink_forward_begin_manifest,

  .manifest_contents = bbsink_gzip_manifest_contents,

  .end_manifest = bbsink_forward_end_manifest,

  .end_backup = bbsink_forward_end_backup,

  .cleanup = bbsink_gzip_cleanup

};

#endif

/*

 * Create a new basebackup sink that performs gzip compression.

 */

bbsink *

bbsink_gzip_new(bbsink *next, pg_compress_specification *compress)

{

#ifndef HAVE_LIBZ

  ereport(ERROR,

      (errcode(ERRCODE_FEATURE_NOT_SUPPORTED),

       errmsg("gzip compression is not supported by this build")));

  return NULL;        /* keep compiler quiet */

#else

  bbsink_gzip *sink;

  int     compresslevel;

  Assert(next != NULL);

  compresslevel = compress->level;

  Assert((compresslevel >= 1 && compresslevel <= 9) ||

       compresslevel == Z_DEFAULT_COMPRESSION);

  sink = palloc0_object(bbsink_gzip);

  *((const bbsink_ops **) &sink->base.bbs_ops) = &bbsink_gzip_ops;

  sink->base.bbs_next = next;

  sink->compresslevel = compresslevel;

  return &sink->base;

#endif

}

#ifdef HAVE_LIBZ

/*

 * Begin backup.

 */

static void

bbsink_gzip_begin_backup(bbsink *sink)

{

  /*

   * We need our own buffer, because we're going to pass different data to

   * the next sink than what gets passed to us.

   */

  sink->bbs_buffer = palloc(sink->bbs_buffer_length);

  /*

   * Since deflate() doesn't require the output buffer to be of any

   * particular size, we can just make it the same size as the input buffer.

   */

  bbsink_begin_backup(sink->bbs_next, sink->bbs_state,

            sink->bbs_buffer_length);

}

/*

 * Prepare to compress the next archive.

 */

static void

bbsink_gzip_begin_archive(bbsink *sink, const char *archive_name)

{

  bbsink_gzip *mysink = (bbsink_gzip *) sink;

  char     *gz_archive_name;

  z_stream   *zs = &mysink->zstream;

  /* Initialize compressor object. */

  memset(zs, 0, sizeof(z_stream));

  zs->zalloc = gzip_palloc;

  zs->zfree = gzip_pfree;

  zs->next_out = (uint8 *) sink->bbs_next->bbs_buffer;

  zs->avail_out = sink->bbs_next->bbs_buffer_length;

  /*

   * We need to use deflateInit2() rather than deflateInit() here so that we

   * can request a gzip header rather than a zlib header. Otherwise, we want

   * to supply the same values that would have been used by default if we

   * had just called deflateInit().

   *

   * Per the documentation for deflateInit2, the third argument must be

   * Z_DEFLATED; the fourth argument is the number of "window bits", by

   * default 15, but adding 16 gets you a gzip header rather than a zlib

   * header; the fifth argument controls memory usage, and 8 is the default;

   * and likewise Z_DEFAULT_STRATEGY is the default for the sixth argument.

   */

  if (deflateInit2(zs, mysink->compresslevel, Z_DEFLATED, 15 + 16, 8,

           Z_DEFAULT_STRATEGY) != Z_OK)

    ereport(ERROR,

        errcode(ERRCODE_INTERNAL_ERROR),

        errmsg("could not initialize compression library"));

  mysink->zstream_initialized = true;

  /*

   * Add ".gz" to the archive name. Note that the pg_basebackup -z produces

   * archives named ".tar.gz" rather than ".tgz", so we match that here.

   */

  gz_archive_name = psprintf("%s.gz", archive_name);

  Assert(sink->bbs_next != NULL);

  bbsink_begin_archive(sink->bbs_next, gz_archive_name);

  pfree(gz_archive_name);

}

/*

 * Compress the input data to the output buffer until we run out of input

 * data. Each time the output buffer fills up, invoke the archive_contents()

 * method for then next sink.

 *

 * Note that since we're compressing the input, it may very commonly happen

 * that we consume all the input data without filling the output buffer. In

 * that case, the compressed representation of the current input data won't

 * actually be sent to the next bbsink until a later call to this function,

 * or perhaps even not until bbsink_gzip_end_archive() is invoked.

 */

static void

bbsink_gzip_archive_contents(bbsink *sink, size_t len)

{

  bbsink_gzip *mysink = (bbsink_gzip *) sink;

  z_stream   *zs = &mysink->zstream;

  /* Compress data from input buffer. */

  zs->next_in = (uint8 *) mysink->base.bbs_buffer;

  zs->avail_in = len;

  while (zs->avail_in > 0)

  {

    int     res;

    /* Write output data into unused portion of output buffer. */

    Assert(mysink->bytes_written < mysink->base.bbs_next->bbs_buffer_length);

    zs->next_out = (uint8 *)

      mysink->base.bbs_next->bbs_buffer + mysink->bytes_written;

    zs->avail_out =

      mysink->base.bbs_next->bbs_buffer_length - mysink->bytes_written;

    /*

     * Try to compress. Note that this will update zs->next_in and

     * zs->avail_in according to how much input data was consumed, and

     * zs->next_out and zs->avail_out according to how many output bytes

     * were produced.

     *

     * According to the zlib documentation, Z_STREAM_ERROR should only

     * occur if we've made a programming error, or if say there's been a

     * memory clobber; we use elog() rather than Assert() here out of an

     * abundance of caution.

     */

    res = deflate(zs, Z_NO_FLUSH);

    if (res == Z_STREAM_ERROR)

      elog(ERROR, "could not compress data: %s", zs->msg);

    /* Update our notion of how many bytes we've written. */

    mysink->bytes_written =

      mysink->base.bbs_next->bbs_buffer_length - zs->avail_out;

    /*

     * If the output buffer is full, it's time for the next sink to

     * process the contents.

     */

    if (mysink->bytes_written >= mysink->base.bbs_next->bbs_buffer_length)

    {

      bbsink_archive_contents(sink->bbs_next, mysink->bytes_written);

      mysink->bytes_written = 0;

    }

  }

}

/*

 * There might be some data inside zlib's internal buffers; we need to get

 * that flushed out and forwarded to the successor sink as archive content.

 *

 * Then we can end processing for this archive.

 */

static void

bbsink_gzip_end_archive(bbsink *sink)

{

  bbsink_gzip *mysink = (bbsink_gzip *) sink;

  z_stream   *zs = &mysink->zstream;

  /* There is no more data available. */

  zs->next_in = (uint8 *) mysink->base.bbs_buffer;

  zs->avail_in = 0;

  while (1)

  {

    int     res;

    /* Write output data into unused portion of output buffer. */

    Assert(mysink->bytes_written < mysink->base.bbs_next->bbs_buffer_length);

    zs->next_out = (uint8 *)

      mysink->base.bbs_next->bbs_buffer + mysink->bytes_written;

    zs->avail_out =

      mysink->base.bbs_next->bbs_buffer_length - mysink->bytes_written;

    /*

     * As bbsink_gzip_archive_contents, but pass Z_FINISH since there is

     * no more input.

     */

    res = deflate(zs, Z_FINISH);

    if (res == Z_STREAM_ERROR)

      elog(ERROR, "could not compress data: %s", zs->msg);

    /* Update our notion of how many bytes we've written. */

    mysink->bytes_written =

      mysink->base.bbs_next->bbs_buffer_length - zs->avail_out;

    /*

     * Apparently we had no data in the output buffer and deflate() was

     * not able to add any. We must be done.

     */

    if (mysink->bytes_written == 0)

      break;

    /* Send whatever accumulated output bytes we have. */

    bbsink_archive_contents(sink->bbs_next, mysink->bytes_written);

    mysink->bytes_written = 0;

  }

  /* Release the compression resources. */

  deflateEnd(zs);

  mysink->zstream_initialized = false;

  /* Must also pass on the information that this archive has ended. */

  bbsink_forward_end_archive(sink);

}

/*

 * Manifest contents are not compressed, but we do need to copy them into

 * the successor sink's buffer, because we have our own.

 */

static void

bbsink_gzip_manifest_contents(bbsink *sink, size_t len)

{

  memcpy(sink->bbs_next->bbs_buffer, sink->bbs_buffer, len);

  bbsink_manifest_contents(sink->bbs_next, len);

}

/*

 * Wrapper function to adjust the signature of palloc to match what libz

 * expects.

 */

static void *

gzip_palloc(void *opaque, unsigned items, unsigned size)

{

  return palloc(items * size);

}

/*

 * Wrapper function to adjust the signature of pfree to match what libz

 * expects.

 */

static void

gzip_pfree(void *opaque, void *address)

{

  pfree(address);

}

/*

 * In case the backup fails, make sure we free the compression context by

 * calling deflateEnd() if needed to avoid a resource leak.

 */

static void

bbsink_gzip_cleanup(bbsink *sink)

{

  bbsink_gzip *mysink = (bbsink_gzip *) sink;

  if (mysink->zstream_initialized)

  {

    deflateEnd(&mysink->zstream);

    mysink->zstream_initialized = false;

  }

}

#endif