// 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.h>
#include "pirclient.h"

#define PIRCLIENT_HDR_SIZE 13

#define PIRCLIENT_REQUEST_CREATE 0x41
#define PIRCLIENT_REQUEST_EXTRACT 0x42

#define PIRCLIENT_RESPONSE_CREATE 0xBF
#define PIRCLIENT_RESPONSE_EXTRACT 0xBE

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
PIRClient::mainloop()
{
    char header[PIRCLIENT_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, PIRCLIENT_HDR_SIZE);
        if (res <= 0) return; // stdin has reached EOF (or error); we
                              //  will terminate
        bodylen = ntohl(*(uint32_t*)(header+PIRCLIENT_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 PIRCLIENT_REQUEST_CREATE:
            if (bodylen >= 32) {
                string plainquery(body, 32);
                string params(body+32, bodylen-32);
                void *queryid = NULL;
                string pirquery;
                create(plainquery, params, queryid, pirquery);
                size_t pirquery_len = pirquery.length();
                header[8] = PIRCLIENT_RESPONSE_CREATE;
                *(uint32_t*)(header+PIRCLIENT_HDR_SIZE-4) =
                        htonl(8+pirquery_len);
                res = write_all(header, PIRCLIENT_HDR_SIZE);
                if (res <= 0) return;
                res = write_all((const char *)&queryid, sizeof(queryid));
                if (res <= 0) return;
                if (sizeof(queryid) < 8) {
                    res = write_all("\0\0\0\0\0\0\0\0", 8-sizeof(queryid));
                    if (res <= 0) return;
                }
                if (pirquery_len > 0) {
                    res = write_all(pirquery.c_str(), pirquery_len);
                    if (res <= 0) return;
                }
            }
            break;
        case PIRCLIENT_REQUEST_EXTRACT:
            if (bodylen >= 8) {
                void *queryid;
                memmove(&queryid, body, sizeof(queryid));
                string pirresponse(body+8, bodylen-8);
                string plainresponse;
                header[8] = PIRCLIENT_RESPONSE_EXTRACT;
                if (extract(queryid, pirresponse, plainresponse)) {
                    response_len = plainresponse.length();
                } else {
                    response_len = 0;
                }
                *(uint32_t*)(header+PIRCLIENT_HDR_SIZE-4) = htonl(response_len);
                res = write_all(header, PIRCLIENT_HDR_SIZE);
                if (res <= 0) return;
                if (response_len > 0) {
                    res = write_all(plainresponse.c_str(), response_len);
                    if (res <= 0) return;
                }
            }
            break;
        }

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