pirserver/pirserver.cc

// Implementation of the main loop of the pirserver, responsible for the
// communication with the tor process.  All the actual private lookup
// work is done by an appropriate subclass of PIRServer.

#include <stdlib.h>
#include <unistd.h>
#include <arpa/inet.h>
#include <string>
#include "pirserver.h"

#define PIRSERVER_HDR_SIZE 13

#define PIRSERVER_REQUEST_PARAMS 0x01
#define PIRSERVER_REQUEST_STORE 0x02
#define PIRSERVER_REQUEST_LOOKUP 0x03

#define PIRSERVER_RESPONSE_PARAMS 0xFF
#define PIRSERVER_RESPONSE_LOOKUP_SUCCESS 0xFE
#define PIRSERVER_RESPONSE_LOOKUP_FAILURE 0xFD

static int
read_all(char *buf, size_t len)
{
    int tot_read = 0;
    while(len > 0) {
        int res = read(0, buf, len);
        if (res <= 0) return res;
        buf += res;
        len -= res;
        tot_read += res;
    }
    return tot_read;
}

static int
write_all(const char *buf, size_t len)
{
    int tot_written = 0;
    while(len > 0) {
        int res = write(1, buf, len);
        if (res <= 0) return res;
        buf += res;
        len -= res;
        tot_written += res;
    }
    return tot_written;
}

void
PIRServer::mainloop()
{
    char header[PIRSERVER_HDR_SIZE];
    size_t bodylen = 0;
    char *body = NULL;
    string query, response;
    size_t response_len;
    while(1) {
        // Read the request from stdin
        int res = read_all(header, PIRSERVER_HDR_SIZE);
        if (res <= 0) return; // stdin has reached EOF (or error); we
                              //  will terminate
        bodylen = ntohl(*(uint32_t*)(header+PIRSERVER_HDR_SIZE-4));
        if (bodylen > 0) {
            body = (char *)malloc(bodylen);
            res = read_all(body, bodylen);
            if (res <= 0) return;
        }

        // We have a complete request.  Dispatch it.
        switch(header[8]) {
        case PIRSERVER_REQUEST_PARAMS:
            get_params(response);
            response_len = response.length();
            header[8] = PIRSERVER_RESPONSE_PARAMS;
            *(uint32_t*)(header+PIRSERVER_HDR_SIZE-4) = htonl(response_len);
            res = write_all(header, PIRSERVER_HDR_SIZE);
            if (res <= 0) return;
            if (response_len > 0) {
                res = write_all(response.c_str(), response_len);
                if (res <= 0) return;
            }
            break;
        case PIRSERVER_REQUEST_STORE:
            if (bodylen >= 32) {
                string key(body, 32);
                string value(body+32, bodylen-32);
                store(key, value);
            }
            break;
        case PIRSERVER_REQUEST_LOOKUP:
            query.assign(body, bodylen);
            if (lookup(query, response)) {
                response_len = response.length();
                header[8] = PIRSERVER_RESPONSE_LOOKUP_SUCCESS;
            } else {
                response_len = 0;
                header[8] = PIRSERVER_RESPONSE_LOOKUP_FAILURE;
            }
            *(uint32_t*)(header+PIRSERVER_HDR_SIZE-4) = htonl(response_len);
            res = write_all(header, PIRSERVER_HDR_SIZE);
            if (res <= 0) return;
            if (response_len > 0) {
                res = write_all(response.c_str(), response_len);
                if (res <= 0) return;
            }
            break;
        }

        // Clean up for the next request.
        free(body);
        body = NULL;
        bodylen = 0;
    }
}