Programování v UNIXu - 22.1.2013 - Kotal

#include <stdio.h>
#include <unistd.h>
#include <string.h>
#include <strings.h>
#include <stdlib.h>
#include <sys/types.h>
#include <sys/socket.h>
#include <netinet/in.h>
#include <libgen.h>
#include <err.h>
#include <sys/types.h>
#include <sys/wait.h>
#include <arpa/inet.h>
#include <time.h>
#include <sys/time.h>
#include <netdb.h>
#include <poll.h>

#include <signal.h>
#include <pthread.h>
#include <errno.h>
#include <semaphore.h>

#define KILL_MSG "killed\n"
#define TO_MSG "timeout\n"

#define BUFF_LEN 2048
#define MAX_ARGS 32
#define MAX_CONN 4

// @author Petr Bělohlávek

// GLOBAL vars =====================================

//  semaphor for conn. count
sem_t sem;
int error, server_socket, client_socket;


// my itoa =====================================
size_t
itoa(int i, char * b)
{
    char const digit[] = "0123456789";
    char * p = b;

    if (i < 0) {
        *p++ = '-';
        i *= -1;
    }

    int shifter = i;
    do { // Move to where representation ends
        ++p;
        shifter = shifter / 10;
    } while (shifter);
    *p = '\0';

    size_t s = p - b;

    do { // Move back, inserting digits as u go
        *--p = digit[i % 10];
        i = i / 10;
    } while (i);

    return (s);
}

// get time in ms =====================================
uint64_t
now()
{
    struct timeval tv;
    gettimeofday(&tv, NULL);
    double time_in_mill =  (tv.tv_sec) * 1000 + (tv.tv_usec) / 1000;
    return ((uint64_t) time_in_mill);
}

// parse by '|' into args =====================================
void
parse_args(char * arg_buffer, char ** args)
{
    // never free() args !!!
    char * token;
    char * str;
    char * saveptr;

    size_t i = 0;
    for (str = arg_buffer; i < MAX_ARGS; str = NULL) {
        token = strtok_r(str, "|", &saveptr);
        if (token == NULL) {
            break;
        }

        args[i++] = token;
    }
    args[i] = NULL;
}

// SIGINT handler =====================================
void
killed(int sig)
{
    write(1, KILL_MSG, strlen(KILL_MSG));
    kill(0, SIGKILL);

    close(client_socket);
    close(server_socket);

    sem_destroy(&sem);

    exit(0);
}

// fork and fork and meassure time =====================================
int
get_exec_time(char * arg_buffer, int timeout, int client_socket)
{
    char * args[MAX_ARGS];
    parse_args(arg_buffer, args);   // free args

    // arg test
    /*
    int j;
    for (j = 0; ; ++j) {
        if (args[j] == NULL) break;
        fprintf(stderr, "%s\n", args[j]);
    }
    */

    int t, rc;
    uint64_t start_time;

    int p[2];
    if (pipe(p) == -1) {
        err(1, "pipe");
    }

    struct pollfd ufds[1];
    ufds[0].fd = p[0];
    ufds[0].events = POLLIN;

    int pid1 = -1, pid2 = -1;

    //  other possibilities: alarm

    // ----------------------------------------------
    //  poll X handle
    switch (pid1 = fork()) {
        case -1:
            close(client_socket);
            err(1, "fork");

        // ----------------------------------------------
        case 0: // run command
            switch (pid2 = fork()) {
                case -1:
                    err(1, "fork");

            case 0: // run command
                close(0); close(1); close(2);
                close(p[0]); close(p[1]);
                close(client_socket);

                execv(args[0], args);
                err(1, "execv failed");

            default:    // parent
                wait(NULL);
                fprintf(stderr, "waited\n");
                write(p[1], ".", 1);
                close(p[0]); close(p[1]);

                close(client_socket);
                exit(0);
            }

            break;

        // ----------------------------------------------
        default:    //  meassure time

            start_time = now();
            rc = poll(ufds, 1, timeout);
            if (rc == -1) {
                close(client_socket);
                err(1, "poll");
            } else if (rc == 0) {
                fprintf(stderr, "timeout\n");
                t = -1;

            } else {
                fprintf(stderr, "stopped\n");
                t = now() - start_time;
            }

            // manage zombies
            if (pid1 != -1) {
                kill(pid1, SIGKILL);
                waitpid(pid1, NULL, 0);
            }
            if (pid2 != -1) {
                kill(pid2, SIGKILL);
                waitpid(pid2, NULL, 0);
            }
            close(p[0]); close(p[1]);
            // while (waitpid(-1, NULL, WNOHANG));
    }
    // ----------------------------------------------
    return (t);
}

int
main(int argc, char ** argv) {

    // TODO: check close() exitcode
    // let client request be valid

    //  sigint
    struct sigaction act;
    bzero(&act, sizeof (act));
    act.sa_handler = killed;
    if (sigaction(SIGINT, &act, NULL) == -1) {
        err(1, "sigaction");
    }

    if (argc != 3) {
        err(1, "usage: ./a.out <port> <timeout>");
    }

    char * portstr = argv[1];
    int timeout = 1000 * atoi(argv[2]); // in miliseconds for poll

    struct sockaddr_storage client;
    struct addrinfo *r, *rorig, hint;
    socklen_t sz;

    memset(&hint, 0, sizeof (hint));
    hint.ai_family = AF_UNSPEC;
    hint.ai_socktype = SOCK_STREAM;
    hint.ai_flags = AI_PASSIVE;

    //  get possible adresses
    if ((error = getaddrinfo(NULL, portstr, &hint, &r)) != 0) {
        errx(1, "%s", gai_strerror(error));
    }

    //  bind to any
    for (rorig = r; r != NULL; r = r->ai_next) {
        server_socket = socket(r->ai_family,
            r->ai_socktype, r->ai_protocol);

        if (!bind(server_socket, r->ai_addr, r->ai_addrlen)) {
            break;
        }
        close(server_socket);
    }

    //  free
    freeaddrinfo(rorig);

    //  if all tried and failed
    if (r == NULL) {
        close(server_socket);
        errx(1, "getaddrinfo");
    }

    //  set reuse addresses
    int optval = 1;
    if (setsockopt(server_socket, SOL_SOCKET, SO_REUSEADDR,
        &optval, sizeof (optval)) == -1) {

        close(server_socket);
        err(1, "setsockopt");
    }

    //  listen
    if (listen(server_socket, MAX_CONN) == -1) {
        close(server_socket);
        err(1, "listen");
    }

    //  init. semaphor
    if (sem_init(&sem, 0, MAX_CONN) == -1) {
        close(server_socket);
        err(1, "sem. init");
    }

    //  tmp vars
    char arg_buffer[BUFF_LEN], repl[BUFF_LEN];
    int n, t, s;

    //  accept every
    for (;;) {
        fprintf(stderr, "- waiting for free slot\n");
        //  don't care about another client for a while
        sem_wait(&sem);
        fprintf(stderr, "- ready\n");

        //  acccept conn.
        sz = sizeof (client);
        if ((client_socket = accept(server_socket,
            (struct sockaddr *) &client, &sz)) == -1) {

            close(client_socket);
            close(server_socket);
            err(1, "listen");
        }

        // create worker
        switch (fork()) {
            case -1:
                close(server_socket);
                err(1, "fork");
                //  may kill

            case 0:  // worker - handle this connection
                if ((n = read(client_socket,
                    &arg_buffer, BUFF_LEN)) == -1) {

                    close(client_socket);

                    err(1, "child read");
                }
                arg_buffer[n] = 0;
                fprintf(stderr, "accepted req.: %s\n",
                    arg_buffer);

                t = get_exec_time(arg_buffer,
                    timeout, client_socket);

                if (t == -1) { // timeout
                    if (write(client_socket, TO_MSG,
                        sizeof (TO_MSG)) == -1) {

                        close(client_socket);

                        err(1, "child write");
                    }
                } else {
                    s = itoa(t, repl);
                    if (write(client_socket,
                        repl, s) == -1) {

                        close(client_socket);

                        err(1, "child write");
                    }
                }

                close(client_socket);

                break;

            default:
                while (waitpid(-1, NULL, WNOHANG)) {
                    sem_post(&sem);
                }
                close(client_socket);
                break;
        }
    }

    // actually never reached
    close(server_socket);
    return (0);
}