/*
 * Copyright (c) 2006, Stefan Walter
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 *     * Redistributions of source code must retain the above
 *       copyright notice, this list of conditions and the
 *       following disclaimer.
 *     * Redistributions in binary form must reproduce the
 *       above copyright notice, this list of conditions and
 *       the following disclaimer in the documentation and/or
 *       other materials provided with the distribution.
 *     * The names of contributors to this software may not be
 *       used to endorse or promote products derived from this
 *       software without specific prior written permission.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
 * "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
 * FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE
 * COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT,
 * INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING,
 * BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS
 * OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED
 * AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF
 * THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH
 * DAMAGE.
 *
 *
 * CONTRIBUTORS
 *  Stef Walter <stef@memberwebs.com>
 */

#include <sys/types.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <unistd.h>
#include <stdlib.h>
#include <stdio.h>
#include <err.h>
#include <errno.h>
#include <string.h>
#include <signal.h>

static void
usage ()
{
    fprintf (stderr, "usage: sockin [-t timeout] socket\n");
    fprintf (stderr, "       sockin [-t timeout] [-d] socket command arg ...\n");
    exit (2);
}

static int
write_all (int fd, void *buf, size_t len)
{
    int r;
    while (len > 0) {
        r = write(fd, buf, len);
        if(r >= 0) {
            buf += r;
            len -= r;
            continue;
        }
        return -1;
    }
    return 0;
}

static int
connect_sock (const char *path, int timeout)
{
    struct sockaddr_un addr;
    int len, sock;

    /* Connect to the socket */
    sock = socket (AF_UNIX, SOCK_STREAM, 0);
    if (sock < 0)
        err (1, "couldn't create socket");

    /* Try to connect once a second for timeout seconds */
    for(;;) {
        addr.sun_family = AF_UNIX;
        strlcpy (addr.sun_path, path, sizeof (addr.sun_path));
        len = sizeof(addr) - (sizeof(addr.sun_path) - strlen (addr.sun_path));\

        if (connect (sock, (struct sockaddr*)&addr, len) < 0) {
            if (errno == EAGAIN || errno == EINTR) {
                continue;
            } else if (timeout--) {
                sleep (1);
                continue;
            } else {
                err (1, "couldn't connect to socket: %s", path);
            }
        }

        break;
    }

    return sock;
}

int
main (int argc, char* argv[])
{
    int timeout = 0;
    int daemonize = 0;
    int sock, i, r, l;
    char buf[128];
    const char *path;
    char **args;
    char *t2;
    char ch;

    /* Parse the arguments nicely */
    while ((ch = getopt (argc, argv, "t:d")) != -1) {
        switch(ch) {

        /* Config directory */
        case 'd':
            daemonize = 1;
            break;

        /* Only print commands */
        case 't':
            timeout = strtol (optarg, &t2, 10);
            if (timeout < 0 || *t2)
                err (2, "invalid timeout: %s", optarg);
            break;

        /* Usage information */
        case '?':
        default:
            usage();
            break;
        }
    }

    argc -= optind;
    argv += optind;

    if (argc < 1)
        usage ();

    /* The socket path */
    path = argv[0];
    sock = connect_sock (path, timeout);

    argc--;
    argv++;

    /* Command mode */
    if (argc) {

        /* Allocate a new command line (no chance to free) */
        args = (char**)calloc (argc, sizeof (char*));
        if (!args)
            err (1, "out of memory");
        for (i = 0; i < argc; i++)
            args[i] = argv[i];

        if (daemonize) {
            if (daemon (1, 0) < 0)
                err (1, "couldn't daemonize process");
        }

        /* And duplicate our socket onto stdout */
        dup2 (sock, 1);

        execvp (args[0], args);
        err (1, "couldn't execute program: %s", args[0]);

    /* Pipe mode */
    } else {

        /* Handle some signals */
        signal (SIGPIPE, SIG_IGN);

        for (;;) {
            l = read (0, buf, sizeof (buf));

            if (l < 0)
                err (1, "couldn't read data");
            else if(l == 0)
                break;

            for (;;) {
                r = write_all (sock, buf, l);

                if (r < 0) {
                    if (errno != EPIPE)
                        err (1, "couldn't write data to socket: %s", argv[0]);

                    /* Try to reconnect when a broken pipe */
                    close (sock);
                    sock = connect_sock (path, timeout);
                    continue;
                }

                break;
            }
        }
    }

    close (sock);
    return 0;
}