source: trunk/instant_messenger/socket/BKP_20071105/BKP_20071026/BKP_20071019/proxy.c @ 151

#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <sys/socket.h>
#include <arpa/inet.h>
#include <netdb.h>
#include <unistd.h>
#include <signal.h>
#include <errno.h>
#include <sys/wait.h>
#include <sys/stat.h>
#include <stdarg.h>
#include <syslog.h>

/* Buffer size, in bytes, used for copying in the Cat() function. */
const size_t BUF_SIZE = 10240;

/* Typedef for a signal handler function. */
typedef void (*signal_handler_t) (int);

/* Prototypes for functions in this file. */
void ParseArgs (int argc, char *argv[], struct in_addr *ra,
                unsigned short *rp, struct in_addr *la, unsigned short *lp);
void Initialise (void);
void sig_child (int signo);
int CreateServerSocket (struct in_addr addr, unsigned short port);
void Daemonise (void);
void MainLoop (int listen_fd, struct in_addr rem_addr,
                unsigned short rem_port);
void Proxy (int server_fd, int client_fd);
int AcceptClientConnection (int listen_fd);
int ConnectToServer (struct in_addr addr, unsigned short port);
void Cat (int in_fd, int out_fd);
int NameToAddr (const char *name, struct in_addr *p_inaddr);
int NameToPort (const char *name, unsigned short *port, const char *proto);
#ifdef __GNUC__
void quit (const char *fmt, ...) __attribute__ ((format (printf, 1, 2)));
void pbomb (const char *fmt, ...) __attribute__ ((format (printf, 1, 2)));
void hbomb (const char *fmt, ...) __attribute__ ((format (printf, 1, 2)));
#else
void quit (const char *fmt, ...);
void pbomb (const char *fmt, ...);
void hbomb (const char *fmt, ...);
#endif
void set_signal_handler (int signum, signal_handler_t sig_handler_func);

/* Global variables */
char g_program_name[100];       /* Initialised from argv[0] in ParseArgs() */

        int
main (int argc, char *argv[])
{
        struct in_addr remote_addr, local_addr;
        unsigned short remote_port, local_port;
        int listen_fd;

        ParseArgs (argc, argv, &remote_addr, &remote_port, &local_addr,
                        &local_port);

        Initialise ();

        /* Create server socket before becoming a daemon so there is still a
         *      chance to print an error message. */
        listen_fd = CreateServerSocket (local_addr, local_port);
        if (listen_fd < 0)
                pbomb ("Unable to create server socket");

        Daemonise ();

        MainLoop (listen_fd, remote_addr, remote_port);       /* never returns */

        exit (EXIT_SUCCESS);
}

/* ParseArgs(): Parse the command line arguments to extract the remote and
 *  * local adresses and port numbers, ra, rp, la & lp.  Exit the program
 *   * gracefully upon error.
 *    */
        void
ParseArgs (int argc, char *argv[], struct in_addr *ra, unsigned short *rp,
                struct in_addr *la, unsigned short *lp)
{
        /*
         *    * argv[0] = program name
         *       * argv[1] = remote_addr
         *          * argv[2] = remote_port
         *             * argv[3] = local_addr (optional)
         *                * argv[4] = local_port (optional)
         *                   */

        char *p = strrchr (argv[0], '/');

        strncpy (g_program_name, (p == NULL) ? argv[0] : p + 1,
                        sizeof (g_program_name) - 1);

        if ((argc < 3) || (argc > 5))
        {
                fprintf (stderr,
                                "usage: %s remote_addr remote_port [local_addr] [local_port]\n",
                                argv[0]);
                exit (EXIT_FAILURE);
        }

        if (NameToAddr (argv[1], ra))
                hbomb ("Unable to resolve \"%s\" to an ip address", argv[1]);

        if (NameToPort (argv[2], rp, "tcp"))
                quit ("Unable to resolve \"%s\" to a port number", argv[2]);

        if (argc < 4)
                la->s_addr = htonl (INADDR_ANY);
        else if (NameToAddr (argv[3], la))
                hbomb ("Unable to resolve \"%s\" to an ip address", argv[3]);

        if (argc < 5)
                memcpy (lp, rp, sizeof (*lp));
        else if (NameToPort (argv[4], lp, "tcp"))
                quit ("Unable to resolve \"%s\" to a port number", argv[4]);
}

/* Initialise(): Setup syslog, signal handlers, and other intialisation.
 *  */
        void
Initialise (void)
{
        openlog (g_program_name, LOG_PID, LOG_USER);
        syslog (LOG_INFO, "%s started", g_program_name);

        chdir ("/");                  /* Change working directory to the root. */

        umask (0);                    /* Clear our file mode creation mask */

        set_signal_handler (SIGCHLD, sig_child);

        signal (SIGPIPE, SIG_IGN);
}

/* sig_child(): Handles SIGCHLD from exiting child processes.
 *  */
        void
sig_child (int signo)
{
        pid_t pid;

        (void) signo;                 /* suppress compiler warning */

        for (;;)
        {
                pid = waitpid (WAIT_ANY, NULL, WNOHANG);

                if (pid > 0)
                        syslog (LOG_INFO, "Caught SIGCHLD from pid %d", pid);
                else
                        break;
        }

        if ((pid < 0) && (errno != ECHILD))
                syslog (LOG_ERR, "waitpid(): %m"), exit (EXIT_FAILURE);

        return;
}

/* CreateServerSocket(): Create a socket, bind it to the specified address
 *  * and port, and set it to listen for client connections.  Returns <0 on
 *   * failure to bind, bombs on error otherwise, returns the fd of the new
 *    * socket on success.
 *     */
        int
CreateServerSocket (struct in_addr addr, unsigned short port)
{
        int err, fd;
        const int on = 1;
        struct sockaddr_in sa;

        /* Create a socket and get its descriptor. */
        fd = socket (AF_INET, SOCK_STREAM, 0);
        if (fd < 0)
                syslog (LOG_ERR, "socket(): %m"), exit (EXIT_FAILURE);

        /* Set SO_REUSEADDR socket option */
        if (setsockopt (fd, SOL_SOCKET, SO_REUSEADDR, &on, sizeof (on)) < 0)
                syslog (LOG_ERR, "setsockopt(fd%d, SO_REUSEADDR): %m", fd);

        /* Load a sa structure with the specified address and port */
        sa.sin_family = AF_INET;
        sa.sin_port = htons (port);
        sa.sin_addr = addr;
        memset (sa.sin_zero, 0, sizeof (sa.sin_zero));

        /* Bind our socket to the address and port specified */
        err = bind (fd, (struct sockaddr *) &sa, sizeof (sa));
        if (err < 0)
        {
                syslog (LOG_ERR, "bind(): %m");
                return err;
        }

        /* Tell socket to listen and queue up to 5 incoming connections. */
        if (listen (fd, 5) < 0)
                syslog (LOG_ERR, "listen(): %m"), exit (EXIT_FAILURE);

        return fd;
}

/* Daemonise(): Put the program in the background, set PPID=1, create a new
 *  * session and process group, without a controlling tty.
 *   */
        void
Daemonise (void)
{
        pid_t pid;

        /* Close stdin, stdout & stderr */
        /* TODO: open /dev/null and dup the fd to stdin, stdout & stderr
         *      close(STDIN_FILENO);
         *           close(STDOUT_FILENO);
         *                close(STDERR_FILENO);
         *                   */

        syslog (LOG_INFO, "%s daemonising", g_program_name);

        /* Fork the process to put it in the background. */

        pid = fork ();
        if (pid == -1)
                syslog (LOG_ERR, "fork(): %m"), exit (EXIT_FAILURE);

        /* parent terminates here, so shell thinks the command is done. */
        if (pid)
                exit (0);

        syslog (LOG_INFO, "%s in background", g_program_name);

        /* 1st child continues to run in the background with PPID=1 */

        /* Become leader of a new session and a new process group,
         *      with no controlling tty.  */
        setsid ();

        /* Fork again to guarantee the process will not be able to aquire a
         *      controlling tty. */

        /*signal (SIGHUP, SIG_IGN); *//* required according to Stevens' UNP2 p333 */

        pid = fork ();
        if (pid == -1)
                syslog (LOG_ERR, "fork(): %m"), exit (EXIT_FAILURE);

        if (pid)                      /* 1st child terminates */
                exit (0);

        /* 2nd child continues, no longer a session or group leader */

        syslog (LOG_INFO, "%s daemonised", g_program_name);
}

/* MainLoop(): Classic concurrent server model. Wait for a client to connect,
 *  * fork a child process to do the business with the client, parent process
 *   * continues to wait for the next connection.  This function does not return.
 *    */
        void
MainLoop (int listen_fd, struct in_addr rem_addr, unsigned short rem_port)
{
        int server_fd, client_fd;
        pid_t child_pid;

        for (;;)
        {
                client_fd = AcceptClientConnection (listen_fd);

                child_pid = fork ();
                if (child_pid == -1)
                        syslog (LOG_ERR, "fork(): %m"), exit (EXIT_FAILURE);

                if (child_pid)            /* parent */
                {
                        close (client_fd);

                        syslog (LOG_INFO, "Forked child pid %d to deal with client",
                                        child_pid);

                        continue;
                }
                else                      /* child */
                {
                        close (listen_fd);

                        server_fd = ConnectToServer (rem_addr, rem_port);

                        Proxy (server_fd, client_fd);

                        syslog (LOG_INFO, "exiting");
                        _exit (0);
                }
        }
}

/* Proxy(): Copies data between server_fd and client_fd in both directions
 *  * until one or other peer closes their connection.
 *   */
        void
Proxy (int server_fd, int client_fd)
{
        pid_t helper_pid;

        syslog (LOG_INFO, "Proxy(fd%d, fd%d)", server_fd, client_fd);

        signal (SIGCHLD, SIG_IGN);

        helper_pid = fork ();
        if (helper_pid == -1)
                syslog (LOG_ERR, "fork(): %m"), exit (EXIT_FAILURE);

        if (helper_pid)               /* parent */
        {
                syslog (LOG_INFO, "Forked child pid %d to help with proxying",
                                helper_pid);

                Cat (server_fd, client_fd);

                shutdown (client_fd, 1);
                close (server_fd);
                close (client_fd);

                wait (0);                 /* wait for helper to exit */
        }
        else                          /* child (helper) */
        {
                Cat (client_fd, server_fd);

                shutdown (server_fd, 1);

                syslog (LOG_INFO, "exiting");
                exit (0);                 /* helper exits here */
        }
}

/* AcceptClientConnection(): waits for a tcp connect to the socket listen_fd,
 *  * which must already be bound and set to listen on a local port.  Bombs on
 *   * error, returns the fd of the new socket on success.
 *    */
        int
AcceptClientConnection (int listen_fd)
{
        int newfd;
        struct sockaddr_in sa;
        socklen_t socklen;

        syslog (LOG_INFO, "AcceptClientConnection(fd%d)", listen_fd);

        /* Accept the connection and create a new socket for it. */
        socklen = sizeof (sa);
        memset (&sa, 0, socklen);
        do
        {
                newfd = accept (listen_fd, (struct sockaddr *) &sa, &socklen);
        }
        while ((newfd < 0) && (errno == EINTR));

        syslog (LOG_INFO, "Accepted client connection on new socket fd%d", newfd);

        if (newfd < 0)
                syslog (LOG_ERR, "accept(): %m"), exit (EXIT_FAILURE);

        if (socklen != sizeof (sa))
                syslog (LOG_ERR, "accept() screwed up!"), exit (EXIT_FAILURE);

        return (newfd);
}

/* ConnectToServer(): attempts a tcp connect to the server specified by addr
 *  * and port.  Bombs on failure to connect, returns the fd of the new socket
 *   * on success.
 *    */
        int
ConnectToServer (struct in_addr addr, unsigned short port)
{
        /* TODO: have a timeout for connect() - see Unix socket FAQ 6.2 */

        int fd, err;
        struct sockaddr_in sa;

        /* Create a socket and get its descriptor. */
        fd = socket (AF_INET, SOCK_STREAM, 0);
        if (fd < 0)
                syslog (LOG_ERR, "socket(): %m"), exit (EXIT_FAILURE);

        sa.sin_family = AF_INET;
        sa.sin_port = htons (port);
        sa.sin_addr = addr;
        memset (sa.sin_zero, 0, sizeof (sa.sin_zero));

        err = connect (fd, (struct sockaddr *) &sa, sizeof (sa));
        if (err < 0)
        {
                syslog (LOG_ERR, "Unable to connect socket fd%d to server: %m", fd);
                exit (EXIT_FAILURE);
        }

        syslog (LOG_INFO, "Connected socket fd%d to server", fd);

        return fd;
}

/* Cat(): read data from in_fd and write it to out_fd until the connection is
 *  * closed by one of the peers.  Data is copied using a dynamically allocated
 *   * buffer.
 *    */
        void
Cat (int in_fd, int out_fd)
{
        unsigned char *const buf = malloc (BUF_SIZE);
        int bytes_rcvd, bytes_sent = 0, i;

        syslog (LOG_INFO, "Cat(fd%d, fd%d)", in_fd, out_fd);

        if (buf == NULL)
                syslog (LOG_ERR, "malloc(): %m"), exit (EXIT_FAILURE);

        do
        {
                bytes_rcvd = recv (in_fd, buf, BUF_SIZE, 0);

                for (i = 0; i < bytes_rcvd; i += bytes_sent)
                {
                        bytes_sent = send (out_fd, buf + i, bytes_rcvd - i, 0);

                        if (bytes_sent < 0)
                                break;
                }

        }
        while ((bytes_rcvd > 0) && (bytes_sent > 0));

        if ((bytes_rcvd < 0) && (errno != ECONNRESET))
                syslog (LOG_ERR, "recv(): %m"), exit (EXIT_FAILURE);

        if ((bytes_sent < 0) && (errno != EPIPE))
                syslog (LOG_ERR, "send(): %m"), exit (EXIT_FAILURE);

        free (buf);
}

/* NameToAddress(): Convert name to an ip address.  Returns 0 on success,
 *  * -1 on failure.
 *   */
        int
NameToAddr (const char *name, struct in_addr *p_inaddr)
{
        struct hostent *he;

        /* First, attempt to convert from string ip format */
        /* TODO: use inet_aton() instead */
        p_inaddr->s_addr = inet_addr (name);
        if (p_inaddr->s_addr != -1U)  /* Success */
                return 0;

        /* Next, attempt to read from /etc/hosts or do a DNS lookup */
        he = gethostbyname (name);
        if (he != NULL)               /* Success */
        {
                memcpy (p_inaddr, he->h_addr, sizeof (struct in_addr));
                return 0;
        }

        return -1;                    /* Failed to resolve name to an ip address */
}

/* NameToPort(): Convert name to a port number. Name can either be a port
 *  * name (in which case proto must also be set to either "tcp" or "udp") or
 *   * name can be the ascii representation of the port number.  Returns 0 on
 *    * success, -1 on failure.
 *     */
        int
NameToPort (const char *name, unsigned short *port, const char *proto)
{
        unsigned long lport;
        char *errpos;
        struct servent *se;

        /* First, attempt to convert string to integer */
        lport = strtoul (name, &errpos, 0);
        if ((*errpos == 0) && (lport <= 65535))       /* Success */
        {
                *port = lport;
                return 0;
        }

        /* Next, attempt to read the string from /etc/services */
        se = getservbyname (name, proto);
        if (se != NULL)               /* Success */
        {
                *port = ntohs (se->s_port);
                return 0;
        }

        return -1;                    /* Failed to resolve port name to a number */
}

/* quit(), pbomb(), hbomb(): Print an error message to stderr and syslog, then
 *  * exit the program.  pbomb() and hbomb() additionally include the string
 *   * representation of errno and h_errno respectively.
 *    */
        void
quit (const char *fmt, ...)     /* quit with msg */
{
        va_list ap;

        fflush (stdout);

        fprintf (stderr, "%s: ", g_program_name);

        va_start (ap, fmt);
        vfprintf (stderr, fmt, ap);
        va_end (ap);

        fputc ('\n', stderr);

        syslog (LOG_ERR, "I quit!");

        exit (EXIT_FAILURE);
}

        void
pbomb (const char *fmt, ...)    /* bomb with perror */
{
        va_list ap;
        int errno_save = errno;
        char buf[100];

        fflush (stdout);

        fprintf (stderr, "%s: ", g_program_name);

        va_start (ap, fmt);
        vsnprintf (buf, sizeof (buf), fmt, ap);
        va_end (ap);

        errno = errno_save;
        perror (buf);

        syslog (LOG_ERR, "Bang!: %s: %m", buf);

        exit (EXIT_FAILURE);
}

        void
hbomb (const char *fmt, ...)    /* bomb with herror */
{
        va_list ap;
        int h_errno_save = h_errno;
        char buf[100];

        fflush (stdout);

        fprintf (stderr, "%s: ", g_program_name);

        va_start (ap, fmt);
        vsnprintf (buf, sizeof (buf), fmt, ap);
        va_end (ap);

        h_errno = h_errno_save;
        herror (buf);

        syslog (LOG_ERR, "Bang!: %s: %s", buf, hstrerror (h_errno));

        exit (EXIT_FAILURE);
}

/* set_signal_handler(): Sets a signal handler function.  Similar to signal()
 *  * but this method is more portable between platforms.
 *   */
        void
set_signal_handler (int signum, signal_handler_t sa_handler_func)
{
        struct sigaction act;

        act.sa_handler = sa_handler_func;
        sigemptyset (&(act.sa_mask));
        act.sa_flags = 0;

        if (sigaction (signum, &act, NULL) < 0)
        {
                syslog (LOG_ERR, "Error setting handler for signal %d: %m", signum);
                exit (EXIT_FAILURE);
        }
}