/* Copyright (c) 2002-2018 Dovecot authors, see the included COPYING file */

#include "lib.h"

#include "buffer.h"

#include "str.h"

#include "iostream-pump.h"

#include "istream.h"

#include "ostream.h"

#include <unistd.h>

#undef iostream_pump_set_completion_callback

struct iostream_pump {

  int refcount;

  struct istream *input;

  struct ostream *output;

  struct io *io;

  iostream_pump_callback_t *callback;

  void *context;

  bool waiting_output;

  bool completed;

};

static void iostream_pump_copy(struct iostream_pump *pump)

{

  enum ostream_send_istream_result res;

  size_t old_size;

  o_stream_cork(pump->output);

  old_size = o_stream_get_max_buffer_size(pump->output);

  o_stream_set_max_buffer_size(pump->output,

    I_MIN(IO_BLOCK_SIZE,

          o_stream_get_max_buffer_size(pump->output)));

  res = o_stream_send_istream(pump->output, pump->input);

  o_stream_set_max_buffer_size(pump->output, old_size);

  o_stream_uncork(pump->output);

  switch(res) {

  case OSTREAM_SEND_ISTREAM_RESULT_ERROR_INPUT:

    io_remove(&pump->io);

    pump->callback(IOSTREAM_PUMP_STATUS_INPUT_ERROR,

             pump->context);

    return;

  case OSTREAM_SEND_ISTREAM_RESULT_ERROR_OUTPUT:

    io_remove(&pump->io);

    pump->callback(IOSTREAM_PUMP_STATUS_OUTPUT_ERROR,

             pump->context);

    return;

  case OSTREAM_SEND_ISTREAM_RESULT_WAIT_OUTPUT:

    i_assert(!pump->output->blocking);

    pump->waiting_output = TRUE;

    io_remove(&pump->io);

    return;

  case OSTREAM_SEND_ISTREAM_RESULT_FINISHED:

    pump->waiting_output = FALSE;

    io_remove(&pump->io);

    /* flush it */

    switch (o_stream_finish(pump->output)) {

    case -1:

      pump->callback(IOSTREAM_PUMP_STATUS_OUTPUT_ERROR,

               pump->context);

      break;

    case 0:

      pump->waiting_output = TRUE;

      pump->completed = TRUE;

      break;

    default:

      pump->callback(IOSTREAM_PUMP_STATUS_INPUT_EOF,

               pump->context);

      break;

    }

    return;

  case OSTREAM_SEND_ISTREAM_RESULT_WAIT_INPUT:

    i_assert(!pump->input->blocking);

    pump->waiting_output = FALSE;

    return;

  }

  i_unreached();

}

static int iostream_pump_flush(struct iostream_pump *pump)

{

  int ret;

  if ((ret = o_stream_flush(pump->output)) <= 0) {

    if (ret < 0) {

      pump->callback(IOSTREAM_PUMP_STATUS_OUTPUT_ERROR,

               pump->context);

    }

    return ret;

  }

  pump->waiting_output = FALSE;

  if (pump->completed) {

    pump->callback(IOSTREAM_PUMP_STATUS_INPUT_EOF, pump->context);

    return 1;

  }

  if (pump->input->blocking)

    iostream_pump_copy(pump);

  else if (pump->io == NULL) {

    pump->io = io_add_istream(pump->input,

            iostream_pump_copy, pump);

    io_set_pending(pump->io);

  }

  return ret;

}

struct iostream_pump *

iostream_pump_create(struct istream *input, struct ostream *output)

{

  struct iostream_pump *pump;

  i_assert(input != NULL &&

     output != NULL);

  i_assert(!input->blocking || !output->blocking);

  /* ref streams */

  i_stream_ref(input);

  o_stream_ref(output);

  /* create pump */

  pump = i_new(struct iostream_pump, 1);

  pump->refcount = 1;

  pump->input = input;

  pump->output = output;

  return pump;

}

void iostream_pump_start(struct iostream_pump *pump)

{

  i_assert(pump != NULL);

  i_assert(pump->callback != NULL);

  /* add flush handler */

  if (!pump->output->blocking) {

    o_stream_set_flush_callback(pump->output,

              iostream_pump_flush, pump);

  }

  /* make IO objects */

  if (pump->input->blocking) {

    i_assert(!pump->output->blocking);

    o_stream_set_flush_pending(pump->output, TRUE);

  } else {

    pump->io = io_add_istream(pump->input,

            iostream_pump_copy, pump);

    io_set_pending(pump->io);

  }

}

struct istream *iostream_pump_get_input(struct iostream_pump *pump)

{

  i_assert(pump != NULL);

  return pump->input;

}

struct ostream *iostream_pump_get_output(struct iostream_pump *pump)

{

  i_assert(pump != NULL);

  return pump->output;

}

void iostream_pump_set_completion_callback(struct iostream_pump *pump,

             iostream_pump_callback_t *callback,

             void *context)

{

  i_assert(pump != NULL);

  pump->callback = callback;

  pump->context = context;

}

void iostream_pump_ref(struct iostream_pump *pump)

{

  i_assert(pump != NULL);

  i_assert(pump->refcount > 0);

  pump->refcount++;

}

void iostream_pump_unref(struct iostream_pump **_pump)

{

  i_assert(_pump != NULL);

  struct iostream_pump *pump = *_pump;

  if (pump == NULL)

    return;

  i_assert(pump->refcount > 0);

  *_pump = NULL;

  if (--pump->refcount > 0)

    return;

  iostream_pump_stop(pump);

  o_stream_unref(&pump->output);

  i_stream_unref(&pump->input);

  i_free(pump);

}

void iostream_pump_destroy(struct iostream_pump **_pump)

{

  i_assert(_pump != NULL);

  struct iostream_pump *pump = *_pump;

  if (pump == NULL)

    return;

  *_pump = NULL;

  iostream_pump_stop(pump);

  o_stream_unref(&pump->output);

  i_stream_unref(&pump->input);

  iostream_pump_unref(&pump);

}

void iostream_pump_stop(struct iostream_pump *pump)

{

  i_assert(pump != NULL);

  if (pump->output != NULL)

    o_stream_unset_flush_callback(pump->output);

  io_remove(&pump->io);

}

bool iostream_pump_is_waiting_output(struct iostream_pump *pump)

{

  return pump->waiting_output;

}

void iostream_pump_switch_ioloop_to(struct iostream_pump *pump,

            struct ioloop *ioloop)

{

  i_assert(pump != NULL);

  if (pump->io != NULL)

    pump->io = io_loop_move_io_to(ioloop, &pump->io);

  o_stream_switch_ioloop_to(pump->output, ioloop);

  i_stream_switch_ioloop_to(pump->input, ioloop);

}

void iostream_pump_switch_ioloop(struct iostream_pump *pump)

{

  iostream_pump_switch_ioloop_to(pump, current_ioloop);

}