/* GIO - GLib Input, Output and Streaming Library
 *
 * Copyright (C) 2010 Collabora, Ltd.
 *
 * This library is free software; you can redistribute it and/or
 * modify it under the terms of the GNU Lesser General Public
 * License as published by the Free Software Foundation; either
 * version 2 of the License, or (at your option) any later version.
 *
 * This library 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
 * Lesser General Public License for more details.
 *
 * You should have received a copy of the GNU Lesser General
 * Public License along with this library; if not, see <http://www.gnu.org/licenses/>.
 *
 * Author: Nicolas Dufresne <nicolas.dufresne@collabora.co.uk>
 */

#include "config.h"

#include "gsocks4aproxy.h"

#include <string.h>

#include "giomodule.h"
#include "giomodule-priv.h"
#include "giostream.h"
#include "ginetaddress.h"
#include "ginputstream.h"
#include "glibintl.h"
#include "goutputstream.h"
#include "gproxy.h"
#include "gproxyaddress.h"
#include "gtask.h"

#define SOCKS4_VERSION		  4

#define SOCKS4_CMD_CONNECT	  1
#define SOCKS4_CMD_BIND		  2

#define SOCKS4_MAX_LEN		  255

#define SOCKS4_REP_VERSION	  0
#define SOCKS4_REP_GRANTED	  90
#define SOCKS4_REP_REJECTED       91
#define SOCKS4_REP_NO_IDENT       92
#define SOCKS4_REP_BAD_IDENT      93

static void g_socks4a_proxy_iface_init (GProxyInterface *proxy_iface);

#define g_socks4a_proxy_get_type _g_socks4a_proxy_get_type
G_DEFINE_TYPE_WITH_CODE (GSocks4aProxy, g_socks4a_proxy, G_TYPE_OBJECT,
			 G_IMPLEMENT_INTERFACE (G_TYPE_PROXY,
						g_socks4a_proxy_iface_init)
			 _g_io_modules_ensure_extension_points_registered ();
			 g_io_extension_point_implement (G_PROXY_EXTENSION_POINT_NAME,
							 g_define_type_id,
							 "socks4a",
							 0))

static void
g_socks4a_proxy_finalize (GObject *object)
{
  /* must chain up */
  G_OBJECT_CLASS (g_socks4a_proxy_parent_class)->finalize (object);
}

static void
g_socks4a_proxy_init (GSocks4aProxy *proxy)
{
  proxy->supports_hostname = TRUE;
}

/*                                                             |-> SOCKSv4a only
 * +----+----+----+----+----+----+----+----+----+----+....+----+------+....+------+
 * | VN | CD | DSTPORT |      DSTIP        | USERID       |NULL| HOST |    | NULL |
 * +----+----+----+----+----+----+----+----+----+----+....+----+------+....+------+
 *    1    1      2              4           variable       1    variable
 */
#define SOCKS4_CONN_MSG_LEN	    (9 + SOCKS4_MAX_LEN * 2)
static gint
set_connect_msg (guint8      *msg,
		 const gchar *hostname,
		 guint16      port,
		 const char  *username,
		 GError     **error)
{
  GInetAddress *addr;
  guint len = 0;
  gsize addr_len;
  gboolean is_ip;
  const gchar *ip;

  msg[len++] = SOCKS4_VERSION;
  msg[len++] = SOCKS4_CMD_CONNECT;

    {
      guint16 hp = g_htons (port);
      memcpy (msg + len, &hp, 2);
      len += 2;
    }

  is_ip = g_hostname_is_ip_address (hostname);

  if (is_ip)
    ip = hostname;
  else
    ip = "0.0.0.1";
    
  addr = g_inet_address_new_from_string (ip);
  addr_len = g_inet_address_get_native_size (addr);

  if (addr_len != 4)
    {
      g_set_error (error, G_IO_ERROR, G_IO_ERROR_PROXY_FAILED,
		  _("SOCKSv4 does not support IPv6 address “%s”"),
		  ip);
      g_object_unref (addr);
      return -1;
    }

  memcpy (msg + len, g_inet_address_to_bytes (addr), addr_len);
  len += addr_len;

  g_object_unref (addr);

  if (username)
    {
      gsize user_len = strlen (username);

      if (user_len > SOCKS4_MAX_LEN)
	{
	  g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_PROXY_FAILED,
			       _("Username is too long for SOCKSv4 protocol"));
	  return -1;
	}

      memcpy (msg + len, username, user_len);
      len += user_len;
    }

  msg[len++] = '\0';

  if (!is_ip)
    {
      gsize host_len = strlen (hostname);

      if (host_len > SOCKS4_MAX_LEN)
	{
	  g_set_error (error, G_IO_ERROR, G_IO_ERROR_PROXY_FAILED,
		       _("Hostname “%s” is too long for SOCKSv4 protocol"),
		       hostname);
	  return -1;
	}

      memcpy (msg + len, hostname, host_len);
      len += host_len;
      msg[len++] = '\0';
    }

  return len;
}

/*
 * +----+----+----+----+----+----+----+----+
 * | VN | CD | DSTPORT |      DSTIP        |
 * +----+----+----+----+----+----+----+----+
 *    1    1      2              4
 */
#define SOCKS4_CONN_REP_LEN	  8
static gboolean
parse_connect_reply (const guint8 *data, GError **error)
{
  if (data[0] != SOCKS4_REP_VERSION)
    {
      g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_PROXY_FAILED,
			   _("The server is not a SOCKSv4 proxy server."));
      return FALSE;
    }

  if (data[1] != SOCKS4_REP_GRANTED)
    {
      g_set_error_literal (error, G_IO_ERROR, G_IO_ERROR_PROXY_FAILED,
			   _("Connection through SOCKSv4 server was rejected"));
      return FALSE;
    }

  return TRUE;
}

static GIOStream *
g_socks4a_proxy_connect (GProxy            *proxy,
			 GIOStream         *io_stream,
			 GProxyAddress     *proxy_address,
			 GCancellable      *cancellable,
			 GError           **error)
{
  GInputStream *in;
  GOutputStream *out;
  const gchar *hostname;
  guint16 port;
  const gchar *username;

  hostname = g_proxy_address_get_destination_hostname (proxy_address);
  port = g_proxy_address_get_destination_port (proxy_address);
  username = g_proxy_address_get_username (proxy_address);

  in = g_io_stream_get_input_stream (io_stream);
  out = g_io_stream_get_output_stream (io_stream);

  /* Send SOCKS4 connection request */
    {
      guint8 msg[SOCKS4_CONN_MSG_LEN];
      gint len;
      
      len = set_connect_msg (msg, hostname, port, username, error);

      if (len < 0)
	goto error;

      if (!g_output_stream_write_all (out, msg, len, NULL,
				      cancellable, error))
	goto error;
    }

  /* Read SOCKS4 response */
    {
      guint8 data[SOCKS4_CONN_REP_LEN];

      if (!g_input_stream_read_all (in, data, SOCKS4_CONN_REP_LEN, NULL,
				    cancellable, error))
	goto error;

      if (!parse_connect_reply (data, error))
	goto error;
    }

  return g_object_ref (io_stream);

error:
  return NULL;
}


typedef struct
{
  GIOStream *io_stream;

  /* For connecting */
  guint8 *buffer;
  gssize length;
  gssize offset;

} ConnectAsyncData;

static void connect_msg_write_cb      (GObject          *source,
				       GAsyncResult     *result,
				       gpointer          user_data);
static void connect_reply_read_cb     (GObject          *source,
				       GAsyncResult     *result,
				       gpointer          user_data);

static void
free_connect_data (ConnectAsyncData *data)
{
  g_object_unref (data->io_stream);
  g_slice_free (ConnectAsyncData, data);
}

static void
do_read (GAsyncReadyCallback callback, GTask *task, ConnectAsyncData *data)
{
   GInputStream *in;
   in = g_io_stream_get_input_stream (data->io_stream);
   g_input_stream_read_async (in,
			      data->buffer + data->offset,
			      data->length - data->offset,
			      g_task_get_priority (task),
			      g_task_get_cancellable (task),
			      callback, task);
}

static void
do_write (GAsyncReadyCallback callback, GTask *task, ConnectAsyncData *data)
{
  GOutputStream *out;
  out = g_io_stream_get_output_stream (data->io_stream);
  g_output_stream_write_async (out,
			       data->buffer + data->offset,
			       data->length - data->offset,
			       g_task_get_priority (task),
			       g_task_get_cancellable (task),
			       callback, task);
}



static void
g_socks4a_proxy_connect_async (GProxy               *proxy,
			       GIOStream            *io_stream,
			       GProxyAddress        *proxy_address,
			       GCancellable         *cancellable,
			       GAsyncReadyCallback   callback,
			       gpointer              user_data)
{
  GError *error = NULL;
  GTask *task;
  ConnectAsyncData *data;
  const gchar *hostname;
  guint16 port;
  const gchar *username;

  data = g_slice_new0 (ConnectAsyncData);
  data->io_stream = g_object_ref (io_stream);

  task = g_task_new (proxy, cancellable, callback, user_data);
  g_task_set_source_tag (task, g_socks4a_proxy_connect_async);
  g_task_set_task_data (task, data, (GDestroyNotify) free_connect_data);

  hostname = g_proxy_address_get_destination_hostname (proxy_address);
  port = g_proxy_address_get_destination_port (proxy_address);
  username = g_proxy_address_get_username (proxy_address); 

  data->buffer = g_malloc0 (SOCKS4_CONN_MSG_LEN);
  data->length = set_connect_msg (data->buffer,
				  hostname, port, username,
				  &error);
  data->offset = 0;

  if (data->length < 0)
    {
      g_task_return_error (task, error);
      g_object_unref (task);
    }
  else
    {
      do_write (connect_msg_write_cb, task, data);
    }
}

static void
connect_msg_write_cb (GObject      *source,
		      GAsyncResult *result,
		      gpointer      user_data)
{
  GTask *task = user_data;
  ConnectAsyncData *data = g_task_get_task_data (task);
  GError *error = NULL;
  gssize written;

  written = g_output_stream_write_finish (G_OUTPUT_STREAM (source),
					  result, &error);
  
  if (written < 0)
    {
      g_task_return_error (task, error);
      g_object_unref (task);
      return;
    }

  data->offset += written;

  if (data->offset == data->length)
    {
      g_free (data->buffer);

      data->buffer = g_malloc0 (SOCKS4_CONN_REP_LEN);
      data->length = SOCKS4_CONN_REP_LEN;
      data->offset = 0;

      do_read (connect_reply_read_cb, task, data);
    }
  else
    {
      do_write (connect_msg_write_cb, task, data);
    }
}

static void
connect_reply_read_cb (GObject       *source,
		       GAsyncResult  *result,
		       gpointer       user_data)
{
  GTask *task = user_data;
  ConnectAsyncData *data = g_task_get_task_data (task);
  GError *error = NULL;
  gssize read;

  read = g_input_stream_read_finish (G_INPUT_STREAM (source),
				     result, &error);

  if (read < 0)
    {
      g_task_return_error (task, error);
      g_object_unref (task);
      return;
    }

  data->offset += read;

  if (data->offset == data->length)
    {
      if (!parse_connect_reply (data->buffer, &error))
	{
	  g_task_return_error (task, error);
	  g_object_unref (task);
	  return;
	}
      else
	{
	  g_task_return_pointer (task, g_object_ref (data->io_stream), g_object_unref);
	  g_object_unref (task);
	  return;
	}
    }
  else
    {
      do_read (connect_reply_read_cb, task, data);
    }
}

static GIOStream *g_socks4a_proxy_connect_finish (GProxy       *proxy,
						  GAsyncResult *result,
						  GError      **error);

static GIOStream *
g_socks4a_proxy_connect_finish (GProxy       *proxy,
			        GAsyncResult *result,
			        GError      **error)
{
  return g_task_propagate_pointer (G_TASK (result), error);
}

static gboolean
g_socks4a_proxy_supports_hostname (GProxy *proxy)
{
  return G_SOCKS4A_PROXY (proxy)->supports_hostname;
}

static void
g_socks4a_proxy_class_init (GSocks4aProxyClass *class)
{
  GObjectClass *object_class;

  object_class = (GObjectClass *) class;
  object_class->finalize = g_socks4a_proxy_finalize;
}

static void
g_socks4a_proxy_iface_init (GProxyInterface *proxy_iface)
{
  proxy_iface->connect  = g_socks4a_proxy_connect;
  proxy_iface->connect_async = g_socks4a_proxy_connect_async;
  proxy_iface->connect_finish = g_socks4a_proxy_connect_finish;
  proxy_iface->supports_hostname = g_socks4a_proxy_supports_hostname;
}