/* proxyfwd.c -- Open a network connection through an http proxy server
   Written by Miles Bader <address@hidden>
*/

#include <stdio.h>
#include <stdlib.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <netdb.h>
#include <errno.h>
#include <error.h>
#include <string.h>
#include <sys/select.h>

#define CONNECT_STRING "CONNECT %s:%d HTTP/1.1\r\n\r\n"

#define BUF_LEN 8192

/* Read a line from socket SOCK.  Result does not include the
   terminating linefeed.  This is only used by the authentication
   protocol, which we call before we set up all the buffering stuff.
   It is possible it should use the buffers too, which would be faster
   (unlike the server, there isn't really a security issue in terms of
   separating authentication from the rest of the code).

   Space for the result is malloc'd and should be freed by the caller.

   Returns number of bytes read.  */
static int
recv_line (sock, host_name, resultp)
    int sock;
    const char *host_name;
    char **resultp;
{
    char *result;
    size_t input_index = 0;
    size_t result_size = 80;

    result = malloc (result_size);
    if (! result)
      error (20, 0, "Out of memory");

    while (1)
    {
	char ch;
	int n;
	n = recv (sock, &ch, 1, 0);
	if (n <= 0)
	    error (21, 0, "%s: %s", host_name, n == 0 ? "EOF" : strerror (errno));

	if (ch == '\012')
	    break;

	result[input_index++] = ch;
	while (input_index + 1 >= result_size)
	{
	    result_size *= 2;
	    result = (char *) realloc (result, result_size);
	    if (! result)
	      error (22, 0, "Out of memory");
	}
    }

    if (resultp)
	*resultp = result;

    /* Terminate it just for kicks, but we *can* deal with embedded NULs.  */
    result[input_index] = '\0';

    if (resultp == NULL)
	free (result);
    return input_index;
}

int main (argc, argv)
     int argc;
     char **argv;
{
  int sock;
  char *proxy_name, *host_name;
  int proxy_port, host_port;
  struct sockaddr_in proxy_sai;
  struct hostent *proxy_hostinfo;
  char in_buf[BUF_LEN], out_buf[BUF_LEN];
  char *in = in_buf, *out = out_buf;
  ssize_t in_buf_len = 0, out_buf_len = 0;
  int in_eof = 0, out_eof = 0;
  char *line_buf = 0;
  char http_id_buf[64], http_msg_buf[1024];
  int http_code;
  fd_set read_fds;
  int debug = 0;

  if (argc > 1 && strcmp (argv[1], "-D") == 0)
    {
      debug = 1;
      argv[1] = argv[0];
      argv++;
      argc--;
    }

  if (argc != 5 || !atoi(argv[2]) || !atoi(argv[4]))
    {
      fprintf (stderr, "Usage: %s [-D] PROXY_HOST PROXY_PORT HOST PORT\n", argv[0]);
      exit (1);
    }

  proxy_name = argv[1];
  proxy_port = atoi (argv[2]);
  host_name = argv[3];
  host_port = atoi (argv[4]);

  sock = socket (AF_INET, SOCK_STREAM, 0);
  if (sock == -1)
    error (2, errno, "cannot create socket");

  bzero (&proxy_sai, sizeof proxy_sai);

  proxy_sai.sin_family = AF_INET;
  proxy_sai.sin_port = htons (proxy_port);
  proxy_hostinfo = gethostbyname (proxy_name);
  if (proxy_hostinfo == NULL)
    error (3, 0, "%s: unknown host", proxy_name);
  proxy_sai.sin_addr = *(struct in_addr *) proxy_hostinfo->h_addr;

  if (connect (sock, (struct sockaddr *) &proxy_sai, sizeof (proxy_sai)) < 0)
    error (4, errno, "%s:%d", proxy_name, proxy_port);

  sprintf (out_buf, CONNECT_STRING, host_name, host_port);
  if (send (sock, out_buf, strlen (out_buf), 0) < 0)
    error (5, errno, "%s:%d", proxy_name, proxy_port);

  /* Wait for HTTP status code, bail out if you don't get back a 2xx code.*/
  recv_line (sock, proxy_name, &line_buf);
  sscanf (line_buf, "%63s %d %1023[^\n\r]", http_id_buf, &http_code, http_msg_buf);
  if (strncmp (http_id_buf, "HTTP/", 5) != 0)
    error (20, 0, "%s:%d: invalid http proxy id: %s",
	   proxy_name, proxy_port, http_id_buf);
  if ((http_code / 100) != 2)
    error (21, 0, "%s:%d: http proxy error %d: %s",
	   proxy_name, proxy_port, http_code, http_msg_buf);

  /* Skip through remaining part of MIME header, recv_line
     consumes the trailing \n */
  while (recv_line (sock, proxy_name, &line_buf) > 0)
    {
      if (line_buf[0] == '\r' || line_buf[0] == 0)
	{
	  free (line_buf);
	  break;
	}
      free (line_buf);
    }

  /* Now propagate I/O between stdin/stdout and the network socket.   */
  while (!in_eof || !out_eof)
    {
      FD_ZERO (&read_fds);
      if (!in_eof && in == in_buf)
	FD_SET (0, &read_fds);
      if (!out_eof && out == out_buf)
	FD_SET (sock, &read_fds);

      if (select (sock + 1, &read_fds, 0, 0, 0) < 0)
	error (4, errno, "select");

      if (!in_eof && in == in_buf && FD_ISSET (0, &read_fds))
	{
	  if (debug)
	    fprintf (stderr, "Reading stdin...\r\n");

	  in_buf_len = read (0, in_buf, sizeof in_buf);
	  if (in_buf_len < 0)
	    error (5, errno, "stdin");
	  if (in_buf_len == 0)
	    in_eof = 1;

	  if (debug)
	    fprintf (stderr, "Reading stdin...read %d bytes\r\n", in_buf_len);
	}

      if (!out_eof && out == out_buf && FD_ISSET (sock, &read_fds))
	{
	  if (debug)
	    fprintf (stderr, "Reading %s:%d...\r\n", proxy_name, proxy_port);

	  out_buf_len = read (sock, out_buf, sizeof out_buf);
	  if (out_buf_len < 0)
	    error (6, errno, "%s:%d", proxy_name, proxy_port);
	  if (out_buf_len == 0)
	    out_eof = 1;

	  if (debug)
	    fprintf (stderr, "Reading %s:%d...read %d bytes\r\n", proxy_name, proxy_port, out_buf_len);
	}

      while (in_buf_len > 0 || out_buf_len > 0)
	{
	  int wrote;

	  if (in < in_buf + in_buf_len)
	    {
	      if (debug)
		fprintf (stderr, "Writing %s:%d...\r\n", proxy_name, proxy_port);

	      wrote = write (sock, in, in_buf + in_buf_len - in);
	      if (wrote < 0)
		error (7, errno, "%s:%d", proxy_name, proxy_port);

	      in += wrote;
	      if (in == in_buf + in_buf_len)
		{
		  in = in_buf;
		  in_buf_len = 0;
		}

	      if (debug)
		fprintf (stderr, "Writing %s:%d...wrote %d bytes\r\n", proxy_name, proxy_port, wrote);
	    }

	  if (out < out_buf + out_buf_len)
	    {
	      if (debug)
		fprintf (stderr, "Writing stdout...\r\n");

	      wrote = write (1, out, out_buf + out_buf_len - out);
	      if (wrote < 0)
		error (8, errno, "stdout");

	      out += wrote;
	      if (out == out_buf + out_buf_len)
		{
		  out = out_buf;
		  out_buf_len = 0;
		}

	      if (debug)
		fprintf (stderr, "Writing stdout...wrote %d bytes\r\n", wrote);
	    }
	}

      if (in_eof)
	{
	  if (debug)
	    fprintf (stderr, "eof on stdin, forwarding to %s:%d...\r\n", proxy_name, proxy_port);
	  write (sock, in_buf, 0);
	}

      if (out_eof)
	{
	  if (debug)
	    fprintf (stderr, "eof on %s:%d, closing stdout...\r\n", proxy_name, proxy_port);
	  close (1);
	}
    }

  exit (0);
}