source: trollforge/lastmeasure/server/lmserver.cpp @ 573

/* TODO:
 * Track referrers.
 */
#include <string>
#include <sstream>
#include <iostream>

#include <cstdio>
#include <cstring>
#include <cstdlib>

#include <arpa/inet.h> /* bind() */

#define PORT 31337
#define BACKLOG 1000
#define THREADS 10

#include "lmutil.h"
#include "lmresource.h"

class Server;

Server *g_server;

class Worker
{
public:
    Worker(Queue<BACKLOG> *queue) :
        m_ready(0),
        m_die(0),
        m_queue(queue)
    {
        pthread_attr_t attr;
        pthread_attr_init(&attr);
        pthread_attr_setdetachstate(&attr, PTHREAD_CREATE_JOINABLE);
        pthread_create(&m_thread, &attr, Helper, this);

        while (!m_ready)
            ; /* spinlock is OK */
    }

    ~Worker()
    {
        pthread_join(m_thread, NULL);
    }

    void HandleRequest(int client);

private:
    static void *Helper(void *data)
    {
        Worker *that = (Worker *)data;
        that->m_ready = 1; /* signal main thread spinlock */

        for (;;)
        {
            int client = that->m_queue->Pop();
            if (client < 0)
                return NULL;

            that->HandleRequest(client);
        }
        return NULL;
    }

    pthread_t m_thread;
    volatile int m_ready, m_die;
    Queue<BACKLOG> *m_queue;
};

class Server
{
public:
    Server(int port, int backlog, int nworkers) :
        m_fd(-1),
        m_port(port),
        m_nworkers(nworkers),
        m_requests(0),
        m_failures(0),
        m_starttime(time(NULL))
    {
        m_fd = socket(AF_INET, SOCK_STREAM, 0);
        if (m_fd < 0)
            Utils::OhShitShitShit("cannot create socket");

        int reuse = 1;
        int ret = setsockopt(m_fd, SOL_SOCKET, SO_REUSEADDR,
                             &reuse, (socklen_t)sizeof(reuse));
        if (ret < 0)
            Utils::OhShitShitShit("cannot set reuse flag on socket");

        struct sockaddr_in saddr;
        saddr.sin_family = AF_INET;
        saddr.sin_port = htons(m_port);
        saddr.sin_addr.s_addr = htonl(INADDR_ANY);
        ret = bind(m_fd, (struct sockaddr *)&saddr, sizeof(saddr));
        if (ret < 0)
            Utils::OhShitShitShit("cannot bind socket");

        ret = listen(m_fd, backlog);
        if (ret < 0)
            Utils::OhShitShitShit("cannot listen to socket");

        m_queue = new Queue<BACKLOG>();
        m_workers = new Worker *[m_nworkers];
        for (int i = 0; i < m_nworkers; i++)
            m_workers[i] = new Worker(m_queue);
    }

    ~Server()
    {
        close(m_fd);
        for (int i = 0; i < m_nworkers; i++)
            m_queue->Push(-1);
        for (int i = 0; i < m_nworkers; i++)
            delete m_workers[i];
        delete[] m_workers;
        delete m_queue;
    }

    void Run()
    {
        for (;;)
        {
            int client = GetNextRequest();
            if (client < 0)
                continue;
            m_queue->Push(client);
        }
    }

    int GetNextRequest()
    {
        struct sockaddr_in caddr;
        struct in_addr client_ip_addr;
        socklen_t addr_len;
        addr_len = sizeof(caddr);

        int ret = accept(m_fd, (struct sockaddr *)&caddr, &addr_len);
        if (ret < 0)
            m_failures++;
        else
            m_requests++;

        return ret;
    }

    std::string Status()
    {
        std::stringstream ret("");
        ret << "<p>Server status: " << m_requests << " requests served.</p>";
        ret << "<p>Uptime: " << (long)difftime(time(NULL), m_starttime) << " seconds.</p>";

        return ret.str();
    }

private:
    int m_fd, m_port, m_nworkers;
    uint64_t m_requests, m_failures;
    time_t m_starttime;
    Queue<BACKLOG> *m_queue;
    Worker **m_workers;
};

int main(void)
{
    g_server = new Server(PORT, BACKLOG, THREADS);

    Utils::Chroot();
    Utils::DropPrivileges();

    Resources::LoadAll();

    g_server->Run();

    delete g_server;

    return 0;
}

void Worker::HandleRequest(int client)
{
    enum { REQUEST_GET, REQUEST_POST, REQUEST_HEAD, REQUEST_UNKNOWN };

    char inbuf[BUFSIZ];
    int ret = recv(client, inbuf, BUFSIZ, 0);
    if (ret < 0)
    {
        close(client);
        return;
    }
    inbuf[ret] = '\0';

    int request = !strncmp(inbuf, "GET ", 4) ? REQUEST_GET
                : !strncmp(inbuf, "POST ", 5) ? REQUEST_POST
                : !strncmp(inbuf, "HEAD ", 5) ? REQUEST_HEAD
                : REQUEST_UNKNOWN;
    if (request == REQUEST_UNKNOWN)
    {
        close(client);
        return;
    }

    /* Find information in the request */
    /* TODO: Perhaps make `query' an array. Example:
     *       `?u=dongs&popup=1'
     *       * query[0] = "u=dongs";
     *       * query[1] = "popup=1";
     */
    char const *resource = "/", *query = "",
               *agent = "Unknown", *host = "127.0.0.1";

    char *resourceparser = strstr(inbuf, " ");
    char *queryparser = NULL;
    if (resourceparser)
    {
        resource = resourceparser + 1;
        queryparser = strstr(resourceparser + 1, "?");
        resourceparser = strstr(resourceparser + 1, " ");

        if (queryparser && resourceparser && queryparser < resourceparser)
        {
            query = queryparser + 1;
        }
    }

    char *agentparser = strstr(inbuf, "\r\nUser-Agent: ");
    if (agentparser)
    {
        agent = agentparser + 14;
        agentparser = strstr(agentparser + 14, "\r\n");
    }

    char *hostparser = strstr(inbuf, "\r\nHost: ");
    if (hostparser)
    {
        host = hostparser + 8;
        hostparser = strstr(hostparser + 8, "\r\n");
    }

    /* Put zeroes in the request for convenience */
    if (resourceparser)
        resourceparser[0] = '\0';
    if (queryparser)
        queryparser[0] = '\0';
    if (agentparser)
        agentparser[0] = '\0';
    if (hostparser)
        hostparser[0] = '\0';

    /* Now build the response */
    std::stringstream response("");
    Response data = Resources::Get(resource, query, host, agent);
    response << "HTTP/1.1 200 OK\r\n"
             << "Date: " << Utils::Time() << "\r\n"
             << "Server: " << Random::ServerName() << "\r\n"
             << "Content-Length: " << data.first.length() << "\r\n"
             //<< "Connection: Keep-Alive\r\n"
             << "Connection: Close\r\n"
             << "Content-Type: " << data.second << "\r\n"
             << "\r\n"
             << data.first;

    send(client, response.str().c_str(), response.str().length(), 0);

    /* FIXME: we could reuse the connection here */
    close(client);
}