zstrm.c

Go to the documentation of this file.

 
#include "config.h"
#include <limits.h>
#include <stdbool.h>
#include <string.h>
#include <time.h>
#include <zconf.h>
#include <zlib.h>
#include "mutt/lib.h"
#include "zstrm.h"
#include "connection.h"

struct ZstrmDirection
{
  z_stream z;          
  char *buf;           
  unsigned int len;    
  unsigned int pos;    
  bool conn_eof : 1;   
  bool stream_eof : 1; 
};

struct ZstrmContext
{
  struct ZstrmDirection read;  
  struct ZstrmDirection write; 
  struct Connection next_conn; 
};

static void *zstrm_malloc(void *opaque, unsigned int items, unsigned int size)
{
  return mutt_mem_calloc(items, size);
}

static void zstrm_free(void *opaque, void *address)
{
  FREE(&address);
}

static int zstrm_open(struct Connection *conn)
{
  return -1;
}

static int zstrm_close(struct Connection *conn)
{
  struct ZstrmContext *zctx = conn->sockdata;
 
  int rc = zctx->next_conn.close(&zctx->next_conn);
 
  double read_ratio = (zctx->read.z.total_in != 0) ?
                          (double) zctx->read.z.total_out /
                              (double) zctx->read.z.total_in :
                          0.0;
  double write_ratio = (zctx->write.z.total_out != 0) ?
                           (double) zctx->write.z.total_in /
                               (double) zctx->write.z.total_out :
                           0.0;
  mutt_debug(LL_DEBUG5, "read %lu->%lu (%.1fx) wrote %lu<-%lu (%.1fx)\n",
             zctx->read.z.total_in, zctx->read.z.total_out, read_ratio,
             zctx->write.z.total_in, zctx->write.z.total_out, write_ratio);
 
  // Restore the Connection's original functions
  conn->sockdata = zctx->next_conn.sockdata;
  conn->open = zctx->next_conn.open;
  conn->close = zctx->next_conn.close;
  conn->read = zctx->next_conn.read;
  conn->write = zctx->next_conn.write;
  conn->poll = zctx->next_conn.poll;
 
  inflateEnd(&zctx->read.z);
  deflateEnd(&zctx->write.z);
  FREE(&zctx->read.buf);
  FREE(&zctx->write.buf);
  FREE(&zctx);
 
  return rc;
}

static int zstrm_read(struct Connection *conn, char *buf, size_t len)
{
  struct ZstrmContext *zctx = conn->sockdata;
  int rc = 0;
  int zrc = 0;
 
retry:
  if (zctx->read.stream_eof)
    return 0;
 
  /* when avail_out was 0 on last call, we need to call inflate again, because
   * more data might be available using the current input, so avoid calling
   * read on the underlying stream in that case (for it might block) */
  if ((zctx->read.pos == 0) && !zctx->read.conn_eof)
  {
    rc = zctx->next_conn.read(&zctx->next_conn, zctx->read.buf, zctx->read.len);
    mutt_debug(LL_DEBUG5, "consuming data from next stream: %d bytes\n", rc);
    if (rc < 0)
      return rc;
    else if (rc == 0)
      zctx->read.conn_eof = true;
    else
      zctx->read.pos += rc;
  }
 
  zctx->read.z.avail_in = (uInt) zctx->read.pos;
  zctx->read.z.next_in = (Bytef *) zctx->read.buf;
  zctx->read.z.avail_out = (uInt) len;
  zctx->read.z.next_out = (Bytef *) buf;
 
  zrc = inflate(&zctx->read.z, Z_SYNC_FLUSH);
  mutt_debug(LL_DEBUG5, "rc=%d, consumed %u/%u bytes, produced %zu/%zu bytes\n",
             zrc, zctx->read.pos - zctx->read.z.avail_in, zctx->read.pos,
             len - zctx->read.z.avail_out, len);
 
  /* shift any remaining input data to the front of the buffer */
  if ((Bytef *) zctx->read.buf != zctx->read.z.next_in)
  {
    memmove(zctx->read.buf, zctx->read.z.next_in, zctx->read.z.avail_in);
    zctx->read.pos = zctx->read.z.avail_in;
  }
 
  switch (zrc)
  {
    case Z_OK:                            /* progress has been made */
      zrc = len - zctx->read.z.avail_out; /* "returned" bytes */
      if (zrc == 0)
      {
        /* there was progress, so must have been reading input */
        mutt_debug(LL_DEBUG5, "inflate just consumed\n");
        goto retry;
      }
      break;
 
    case Z_STREAM_END: /* everything flushed, nothing remaining */
      mutt_debug(LL_DEBUG5, "inflate returned Z_STREAM_END\n");
      zrc = len - zctx->read.z.avail_out; /* "returned" bytes */
      zctx->read.stream_eof = true;
      break;
 
    case Z_BUF_ERROR: /* no progress was possible */
      if (!zctx->read.conn_eof)
      {
        mutt_debug(LL_DEBUG5, "inflate returned Z_BUF_ERROR. retrying\n");
        goto retry;
      }
      zrc = 0;
      break;
 
    default:
      /* bail on other rcs, such as Z_DATA_ERROR, or Z_MEM_ERROR */
      mutt_debug(LL_DEBUG5, "inflate returned %d. aborting\n", zrc);
      zrc = -1;
      break;
  }
 
  return zrc;
}

static int zstrm_poll(struct Connection *conn, time_t wait_secs)
{
  struct ZstrmContext *zctx = conn->sockdata;
 
  mutt_debug(LL_DEBUG5, "%s\n",
             (zctx->read.z.avail_out == 0) || (zctx->read.pos > 0) ?
                 "last read wrote full buffer" :
                 "falling back on next stream");
  if ((zctx->read.z.avail_out == 0) || (zctx->read.pos > 0))
    return 1;
 
  return zctx->next_conn.poll(&zctx->next_conn, wait_secs);
}

static int zstrm_write(struct Connection *conn, const char *buf, size_t count)
{
  struct ZstrmContext *zctx = conn->sockdata;
  int rc;
 
  zctx->write.z.avail_in = (uInt) count;
  zctx->write.z.next_in = (Bytef *) buf;
  zctx->write.z.avail_out = (uInt) zctx->write.len;
  zctx->write.z.next_out = (Bytef *) zctx->write.buf;
 
  do
  {
    int zrc = deflate(&zctx->write.z, Z_PARTIAL_FLUSH);
    if (zrc == Z_OK)
    {
      /* push out produced data to the underlying stream */
      zctx->write.pos = zctx->write.len - zctx->write.z.avail_out;
      char *wbufp = zctx->write.buf;
      mutt_debug(LL_DEBUG5, "deflate consumed %zu/%zu bytes\n",
                 count - zctx->write.z.avail_in, count);
      while (zctx->write.pos > 0)
      {
        rc = zctx->next_conn.write(&zctx->next_conn, wbufp, zctx->write.pos);
        mutt_debug(LL_DEBUG5, "next stream wrote: %d bytes\n", rc);
        if (rc < 0)
          return -1; /* we can't recover from write failure */
 
        wbufp += rc;
        zctx->write.pos -= rc;
      }
 
      /* see if there's more for us to do, retry if the output buffer
       * was full (there may be something in zlib buffers), and retry
       * when there is still available input data */
      if ((zctx->write.z.avail_out != 0) && (zctx->write.z.avail_in == 0))
        break;
 
      zctx->write.z.avail_out = (uInt) zctx->write.len;
      zctx->write.z.next_out = (Bytef *) zctx->write.buf;
    }
    else
    {
      /* compression went wrong, but this is basically impossible
       * according to the docs */
      return -1;
    }
  } while (true);
 
  // Guard overflow when narrowing to int for Connection::write(); see #4940.
  if (count > INT_MAX)
  {
    mutt_debug(LL_DEBUG1, "zstrm_write: count %zu exceeds INT_MAX\n", count);
    return -1;
  }
  return (int) count;
}

void mutt_zstrm_wrap_conn(struct Connection *conn)
{
  struct ZstrmContext *zctx = MUTT_MEM_CALLOC(1, struct ZstrmContext);
 
  /* store wrapped stream as next stream */
  zctx->next_conn.fd = conn->fd;
  zctx->next_conn.sockdata = conn->sockdata;
  zctx->next_conn.open = conn->open;
  zctx->next_conn.close = conn->close;
  zctx->next_conn.read = conn->read;
  zctx->next_conn.write = conn->write;
  zctx->next_conn.poll = conn->poll;
 
  /* replace connection with our wrappers, where appropriate */
  conn->sockdata = zctx;
  conn->open = zstrm_open;
  conn->read = zstrm_read;
  conn->write = zstrm_write;
  conn->close = zstrm_close;
  conn->poll = zstrm_poll;
 
  /* allocate/setup (de)compression buffers */
  zctx->read.len = 8192;
  zctx->read.buf = MUTT_MEM_MALLOC(zctx->read.len, char);
  zctx->read.pos = 0;
  zctx->write.len = 8192;
  zctx->write.buf = MUTT_MEM_MALLOC(zctx->write.len, char);
  zctx->write.pos = 0;
 
  /* initialise zlib for inflate and deflate for RFC4978 */
  zctx->read.z.zalloc = zstrm_malloc;
  zctx->read.z.zfree = zstrm_free;
  zctx->read.z.opaque = NULL;
  zctx->read.z.avail_out = zctx->read.len;
  (void) inflateInit2(&zctx->read.z, -15);
  zctx->write.z.zalloc = zstrm_malloc;
  zctx->write.z.zfree = zstrm_free;
  zctx->write.z.opaque = NULL;
  zctx->write.z.avail_out = zctx->write.len;
  (void) deflateInit2(&zctx->write.z, Z_DEFAULT_COMPRESSION, Z_DEFLATED, -15, 8,
                      Z_DEFAULT_STRATEGY);
}