/***************************************************************************
 *
 * Copyright (c) 2000 BalaBit IT Ltd, Budapest, Hungary
 *
 * This program is free software; you can redistribute it and/or modify
 * it under the terms of the GNU General Public License as published by
 * the Free Software Foundation; either version 2 of the License, or
 * (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU General Public License for more details.
 *
 * You should have received a copy of the GNU General Public License
 * along with this program; if not, write to the Free Software
 * Foundation, Inc., 675 Mass Ave, Cambridge, MA 02139, USA.
 *
 * $Id: plug.c,v 1.28.4.1 2001/05/11 14:19:11 bazsi Exp $
 *
 ***************************************************************************/

#include <zorp/stream.h>
#include <zorp/proxy.h>
#include <zorp/poll.h>
#include <zorp/policy.h>
#include <zorp/thread.h>
#include <zorp/log.h>
#include <zorp/registry.h>
#include <zorp/sockaddr.h>
#include <zorp/io.h>
#include <zorp/ctrack.h>

#include <sys/types.h>
#include <sys/time.h>
#include <unistd.h>
#include <errno.h>
#include <arpa/inet.h>

#include "plugtrack.h"
#include "plugudp.h"

#define PROXY_BUFSIZE    1500 /* 64k */

typedef struct _ZIOBuffer
{
  gchar buf[PROXY_BUFSIZE];
  guint ofs, end;
  guint packet_count, packet_bytes;
} ZIOBuffer;

#define EOF_CLIENT         0x0001
#define EOF_SERVER         0x0002
#define EOF_STACKED_CLIENT 0x0004
#define EOF_STACKED_SERVER 0x0008

typedef struct _PlugProxy
{
  ZProxy super;
  ZStackedProxy *stacked;
  ZIOBuffer buffers[EP_MAX];
  ZIOBuffer *downbufs[EP_MAX];
  ZPoll *poll;
  gint eofmask;
  gint timeout;
  guint global_packet_count;
  struct timeval started_time;
  gint packet_stats_interval;
  gboolean copy_to_server, copy_to_client;
} PlugProxy;

#define ENABLE_PKT_STATS 1

#define PLUG_DEBUG "plug.debug"

static void plug_proxy_free(ZProxy *s);

void
plug_config_set_defaults(PlugProxy *self)
{
  z_proxy_enter(self);

  self->copy_to_server = 1;
  self->copy_to_client = 1;
  self->timeout = -1;
  gettimeofday(&self->started_time, NULL);
  self->poll = z_poll_new();

  z_proxy_leave(self);
}

void
plug_register_vars(PlugProxy *self)
{
  z_proxy_enter(self);
  z_proxy_var_new(&self->super, "timeout", Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->timeout);
  z_proxy_var_new(&self->super, "copy_to_client", Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->copy_to_client);
  z_proxy_var_new(&self->super, "copy_to_server", Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->copy_to_server);
  z_proxy_var_new(&self->super, "packet_stats_interval", Z_VAR_GET | Z_VAR_SET_CONFIG | Z_VAR_TYPE_INT, &self->packet_stats_interval);
  z_proxy_leave(self);
}

static guint
plug_read_input(PlugProxy *self, ZStream *input, ZIOBuffer *buf)
{
  gint rc;
  gboolean called;

  z_proxy_enter(self);
  rc = z_stream_read(input, buf->buf, sizeof(buf->buf), &buf->end);
  switch (rc)
    {
    case G_IO_ERROR_NONE:

      buf->packet_bytes += buf->end;
      buf->packet_count++;
      self->global_packet_count++;
      if (self->packet_stats_interval &&
	  (self->global_packet_count % self->packet_stats_interval) == 0)
	{
	  int i;
	  ZPolicyObj *res;

	  z_policy_lock(self->super.thread);
	  for (i = EP_CLIENT; i < EP_MAX; i++)
	    {
	      ZIOBuffer *buf;
              struct timeval now, spent;
	      char *varname[EP_MAX] = { "bandwidth_to_client", "bandwidth_to_server" };

	      buf = &self->buffers[i];
      
              gettimeofday(&now, NULL);
              spent.tv_sec = now.tv_sec - self->started_time.tv_sec;
              spent.tv_usec = now.tv_usec - self->started_time.tv_usec;
              if (spent.tv_usec < -500000)
	        {
	          spent.tv_sec++;
		}
	      z_policy_setattr(self->super.handler, varname[i], z_policy_var_build("d", ((double) buf->packet_bytes) / spent.tv_sec));
	    }
	  res = z_policy_call(self->super.handler, "packetStats", 
	  		      z_policy_var_build("iiii", 
					         self->buffers[EP_CLIENT].packet_bytes, 
					         self->buffers[EP_CLIENT].packet_count, 
					         self->buffers[EP_SERVER].packet_bytes, 
					         self->buffers[EP_SERVER].packet_count),
			      &called);
	  if (res)
	    {
              if (!z_policy_var_parse(res, "i", &rc))
                rc = Z_DENY;
            }
          else
            if (called)
              rc = Z_DENY;
          z_policy_var_unref(res);
	  z_policy_unlock(self->super.thread);
	  if (rc != Z_ACCEPT)
	    {
	      z_proxy_leave(self);
	      return ST_CLOSE | ST_EOF;
	    }
	}
      if (buf->end == 0)
        {
          z_proxy_leave(self);
          return ST_CLOSE | ST_EOF;
        }
      z_proxy_leave(self);
      return ST_OK | ST_GOON;
    case G_IO_ERROR_AGAIN:
      z_proxy_leave(self);
      return ST_AGAIN | ST_GOON;
    default:
      z_proxy_leave(self);
      return ST_ERR | ST_CLOSE;
    }
  z_proxy_leave(self);
}

static guint
plug_write_output(PlugProxy *self, ZIOBuffer *buf, ZStream *output)
{
  int rc;
  guint bytes_written;
  
  z_proxy_enter(self);
  if (buf->ofs != buf->end)
    {
      /* buffer not empty */
      rc = z_stream_write(output, &buf->buf[buf->ofs], buf->end - buf->ofs, &bytes_written);
      switch (rc) 
        {
        case G_IO_ERROR_NONE:
          buf->ofs += bytes_written;
          break;
        case G_IO_ERROR_AGAIN:
          break;
        default:
          z_proxy_leave(self);
          return ST_ERR | ST_CLOSE;
        }
      if (buf->ofs != buf->end)
        {
          output->want_write = TRUE;
          z_proxy_leave(self);
          return ST_AGAIN | ST_GOON;
        }
      else
        {
          z_proxy_leave(self);
          return ST_OK | ST_GOON;
        }
    }
  z_proxy_leave(self);
  return ST_OK | ST_GOON;
}

static guint
plug_copy_data(PlugProxy *self, ZStream *from, ZStream *to, ZIOBuffer *buf)
{
  int rc;
  int pkt_count = 0;

  z_proxy_enter(self);
  from->want_read = FALSE;

  if (to)
    {
      to->want_write = FALSE;
      rc = plug_write_output(self, buf, to);
      if (ST_ERROR(rc) != ST_OK)
        {
          z_proxy_leave(self);
	  return rc;
	}
    }

  while (1) 
    {  
      buf->ofs = buf->end = 0;
      rc = plug_read_input(self, from, buf);
      if (ST_ERROR(rc) == ST_OK)
        {
	  if (to)
	    {
	      rc = plug_write_output(self, buf, to);
	      if (ST_ERROR(rc) == ST_AGAIN)
		break;
              else if (ST_FAILED(rc))
                {
                  z_message(0, "%s: Error writing %s (%m)", self->super.session_id, to->name);
                  z_proxy_leave(self);
                  return ST_QUIT | ST_OK;
                }
	    }
        }
      else if (ST_ERROR(rc) == ST_AGAIN)
        break;
      else if (ST_ERROR(rc) == ST_EOF)
        {
          z_proxy_leave(self);
          return ST_GOON | ST_EOF;
        }
      else if (ST_FAILED(rc))
        {
          z_message(0, "%s: Error reading %s (%s)", self->super.session_id, from->name, g_strerror(errno));
          return ST_QUIT | ST_OK;
        }
      pkt_count++;
    }

  if (buf->ofs == buf->end)
    {
      from->want_read = TRUE;
    }
  z_proxy_leave(self);
  return ST_OK | ST_GOON;
}

/* callbacks when no stacking is made */
static guint
plug_copy_client_to_server(ZStream *stream, GIOCondition cond)
{
  PlugProxy *self = (PlugProxy *) stream->user_data;
  gint ret;

  z_proxy_enter(self);
  if (self->copy_to_server)
    ret = plug_copy_data(self, self->super.endpoints[EP_CLIENT], self->super.endpoints[EP_SERVER], &self->buffers[EP_SERVER]);
  else
    ret = plug_copy_data(self, self->super.endpoints[EP_CLIENT], NULL, &self->buffers[EP_SERVER]);
  if (ST_ERROR(ret) == ST_EOF)
    {
      ret = ST_QUIT | ST_OK;
    }    
  z_proxy_leave(self);
  return ret;
}

static guint
plug_copy_server_to_client(ZStream *stream, GIOCondition cond)
{
  PlugProxy *self = (PlugProxy *) stream->user_data;
  gint ret;

  z_proxy_enter(self);
  if (self->copy_to_client)
    ret = plug_copy_data(self, self->super.endpoints[EP_SERVER], self->super.endpoints[EP_CLIENT], &self->buffers[EP_CLIENT]);
  else
    ret = plug_copy_data(self, self->super.endpoints[EP_SERVER], NULL, &self->buffers[EP_CLIENT]);

  if (ST_ERROR(ret) == ST_EOF)
    {
      ret = ST_QUIT | ST_OK;
    }
  z_proxy_leave(self);
  return ret;
}

/* callbacks when a stacked module exists */
static guint
plug_copy_client_to_down(ZStream *stream, GIOCondition cond)
{
  PlugProxy *self = (PlugProxy *) stream->user_data;
  gint ret;

  z_proxy_enter(self);
  if (self->copy_to_server)
    ret = plug_copy_data(self, self->super.endpoints[EP_CLIENT], self->stacked->downstreams[EP_CLIENT], self->downbufs[EP_CLIENT]);
  else
    ret = plug_copy_data(self, self->super.endpoints[EP_CLIENT], NULL, self->downbufs[EP_CLIENT]);

  if (ST_ERROR(ret) == ST_EOF)
    {
      z_stream_shutdown(self->super.endpoints[EP_CLIENT], SHUT_RD);
      z_stream_shutdown(self->stacked->downstreams[EP_CLIENT], SHUT_WR);
      self->eofmask |= EOF_CLIENT;
      ret = ST_GOON | ST_OK;
    }
  z_proxy_leave(self);
  return ret;
}

static guint
plug_copy_down_to_client(ZStream *stream, GIOCondition cond)
{
  PlugProxy *self = (PlugProxy *) stream->user_data;
  gint ret;

  z_proxy_enter(self);
  ret = plug_copy_data(self, self->stacked->downstreams[EP_CLIENT], self->super.endpoints[EP_CLIENT], &self->buffers[EP_CLIENT]);
  if (ST_ERROR(ret) == ST_EOF)
    {
      z_stream_shutdown(self->stacked->downstreams[EP_CLIENT], SHUT_RD);
      z_stream_shutdown(self->super.endpoints[EP_CLIENT], SHUT_WR);
      self->eofmask |= EOF_STACKED_CLIENT;
      if ((self->eofmask & (EOF_STACKED_CLIENT + EOF_STACKED_SERVER)) == (EOF_STACKED_CLIENT + EOF_STACKED_SERVER))
        ret = ST_QUIT | ST_OK;
      else
        ret = ST_GOON | ST_OK;
    }
  z_proxy_leave(self);
  return ret;
}

static guint
plug_copy_server_to_down(ZStream *stream, GIOCondition cond)
{
  PlugProxy *self = (PlugProxy *) stream->user_data;
  gint ret;

  z_proxy_enter(self);
  if (self->copy_to_client)
    ret = plug_copy_data(self, self->super.endpoints[EP_SERVER], self->stacked->downstreams[EP_SERVER], self->downbufs[EP_SERVER]);
  else
    ret = plug_copy_data(self, self->super.endpoints[EP_SERVER], NULL, self->downbufs[EP_SERVER]);

  if (ST_ERROR(ret) == ST_EOF)
    {
      z_stream_shutdown(self->super.endpoints[EP_SERVER], SHUT_RD);
      z_stream_shutdown(self->stacked->downstreams[EP_SERVER], SHUT_WR);
      self->eofmask |= EOF_SERVER;
      ret = ST_GOON | ST_OK;
    }
  z_proxy_leave(self);
  return ret;
}

static guint
plug_copy_down_to_server(ZStream *stream, GIOCondition cond)
{
  PlugProxy *self = (PlugProxy *) stream->user_data;
  gint ret;

  z_proxy_enter(self);
  ret = plug_copy_data(self, self->stacked->downstreams[EP_SERVER], self->super.endpoints[EP_SERVER], &self->buffers[EP_CLIENT]);
  if (ST_ERROR(ret) == ST_EOF)
    {
      z_stream_shutdown(self->stacked->downstreams[EP_SERVER], SHUT_RD);
      z_stream_shutdown(self->super.endpoints[EP_SERVER], SHUT_WR);
      self->eofmask |= EOF_STACKED_SERVER;
      if ((self->eofmask & (EOF_STACKED_CLIENT + EOF_STACKED_SERVER)) == (EOF_STACKED_CLIENT + EOF_STACKED_SERVER))
        ret = ST_QUIT | ST_OK;
      else
        ret = ST_GOON | ST_OK;
    }

  z_proxy_leave(self);
  return ret;
}

static gboolean
plug_init_streams(PlugProxy *self)
{
  z_proxy_enter(self);
  
  if (!self->super.endpoints[EP_CLIENT] ||
      !self->super.endpoints[EP_SERVER] ||
      !self->poll)
    {
      z_proxy_leave(self);
      return FALSE;
    }
    
  z_fd_set_nonblock(self->super.endpoints[EP_CLIENT]->fd, 1);
  self->super.endpoints[EP_CLIENT]->read = plug_copy_client_to_server;
  self->super.endpoints[EP_CLIENT]->write = plug_copy_server_to_client;
  self->super.endpoints[EP_CLIENT]->user_data = self;
  self->super.endpoints[EP_CLIENT]->want_read = TRUE;
  self->super.endpoints[EP_CLIENT]->timeout = -2;
  
  z_fd_set_nonblock(self->super.endpoints[EP_SERVER]->fd, 1);
  self->super.endpoints[EP_SERVER]->read = plug_copy_server_to_client;
  self->super.endpoints[EP_SERVER]->write = plug_copy_client_to_server;
  self->super.endpoints[EP_SERVER]->user_data = self;
  self->super.endpoints[EP_SERVER]->want_read = TRUE;
  self->super.endpoints[EP_SERVER]->timeout = -2;

  z_poll_add_stream(self->poll, self->super.endpoints[EP_CLIENT]);
  z_poll_add_stream(self->poll, self->super.endpoints[EP_SERVER]);
  
  z_proxy_leave(self);
  return TRUE;
}

static gboolean
plug_init_stacked_streams(PlugProxy *self)
{
  z_proxy_enter(self);
  
  if (self->stacked)
    {
      self->downbufs[EP_CLIENT] = g_new0(ZIOBuffer, 1);
      self->downbufs[EP_SERVER] = g_new0(ZIOBuffer, 1);

      self->super.endpoints[EP_CLIENT]->read = plug_copy_client_to_down;
      self->super.endpoints[EP_CLIENT]->write = plug_copy_down_to_client;

      self->super.endpoints[EP_SERVER]->read = plug_copy_server_to_down;
      self->super.endpoints[EP_SERVER]->write = plug_copy_down_to_server;

      self->stacked->downstreams[EP_CLIENT]->read = plug_copy_down_to_client;
      self->stacked->downstreams[EP_CLIENT]->write = plug_copy_client_to_down;
      self->stacked->downstreams[EP_CLIENT]->user_data = self;
      self->stacked->downstreams[EP_CLIENT]->want_read = TRUE;

      self->stacked->downstreams[EP_SERVER]->read = plug_copy_down_to_server;
      self->stacked->downstreams[EP_SERVER]->write = plug_copy_server_to_down;
      self->stacked->downstreams[EP_SERVER]->user_data = self;
      self->stacked->downstreams[EP_SERVER]->want_read = TRUE;

      z_poll_add_stream(self->poll, self->stacked->downstreams[EP_CLIENT]);
      z_poll_add_stream(self->poll, self->stacked->downstreams[EP_SERVER]);
    }
  z_proxy_leave(self);
  return TRUE;
}

static gboolean
plug_request_stack_event(PlugProxy *self)
{
  ZPolicyObj *res;
  gboolean called;
  gboolean rc = TRUE;

  z_proxy_enter(self);
  z_policy_lock(self->super.thread);
  res = z_policy_call(self->super.handler, "requestStack", NULL, &called);
  if (res)
    {
      if (res != z_policy_none)
        self->stacked = z_proxy_stack_proxy(&self->super, res);
    }
  else
    if (called)
      rc = FALSE;
  z_policy_var_unref(res);
  z_policy_unlock(self->super.thread);
  z_proxy_leave(self);
  return rc;
}

static gboolean
plug_main(PlugProxy *self)
{
  z_proxy_enter(self);
  
  ZPROXY_SET_STATE(self, PS_CONFIG);
  if (!z_proxy_config_event(&self->super))
    {
      z_proxy_leave(self);
      return FALSE;
    }

  ZPROXY_SET_STATE(self, PS_STARTING_UP);
  if (!z_proxy_startup_event(&self->super))
    {
      z_proxy_leave(self);
      return FALSE;
    }

  if (!z_proxy_connect_server_event(&self->super, NULL, 0))
    {
      z_proxy_leave(self);
      return FALSE;
    }

  if (!plug_init_streams(self))
    {
      z_proxy_leave(self);
      return FALSE;
    }

  if (!plug_request_stack_event(self))
    {
      z_proxy_leave(self);
      return FALSE;
    }
    
  if (!plug_init_stacked_streams(self))
    {
      z_proxy_leave(self);
      return FALSE;
    }
    
  ZPROXY_SET_STATE(self, PS_WORKING);

  
  while (!(self->super.flags & PF_QUIT) && z_poll_iter_timeout(self->poll, self->timeout))
    {
      ;
    }
 
  ZPROXY_SET_STATE(self, PS_SHUTTING_DOWN);
  z_proxy_shutdown_event(&self->super);
  
  z_proxy_leave(self);
  return TRUE;
}

static void
plug_thread(gpointer s)
{
  PlugProxy *self = (PlugProxy *) s;

  z_proxy_enter(self);

  plug_config_set_defaults(self);
  plug_register_vars(self);
  plug_main(self);

  z_proxy_destroy(&self->super);
  z_leave();
}

PlugProxy *
plug_init_proxy(gchar *session_id,
                ZStream *client,
                ZPolicyObj *handler)
{
  PlugProxy *self = g_new0(PlugProxy, 1);

  z_enter();
  z_proxy_init(&self->super, session_id, client, handler);
  
  self->super.free = plug_proxy_free;
  z_leave();
  return self;
}

ZProxy *
plug_proxy_new(gchar *session_id,
	       ZStream *client,
	       ZPolicyObj *handler)
{
  PlugProxy *self;
  
  z_enter();
  self = plug_init_proxy(session_id, client, handler);
  z_thread_new(session_id, plug_thread, self);
  z_leave();
  return &self->super;
}
    
static void
plug_proxy_free(ZProxy *s)
{
  PlugProxy *self = (PlugProxy *) s;
  
  z_proxy_enter(self);
  if (self->stacked)
    {
      g_free(self->downbufs[0]);
      g_free(self->downbufs[1]);
      z_stacked_proxy_destroy(self->stacked);
    }
  if (self->poll)
    {
      z_poll_unref(self->poll);
      self->poll = NULL;
    }

  z_proxy_leave(self);
}  

gint
zorp_module_init(void)
{
  z_registry_add("plug", ZR_SPROXY, plug_proxy_new);
#if ENABLE_CONNTRACK
  z_registry_add("plug", ZR_DPROXY, plug_packet_new);
  z_registry_add("plug", ZR_DTRACK, plug_tracker_new);
#endif
  return TRUE;
}
