protocol.cpp

#include "protocol.h"

#include <cerrno>
#include <sys/socket.h>
#include <sys/types.h>
#include <openssl/ssl.h>
#include <unistd.h>
#include <cassert>
#include "utils.h"
#include "tls_registry.h"


using namespace std;

static bool send_all(int fd, const void *buf, size_t len) {
    const uint8_t *charbuf = static_cast<const uint8_t *>(buf);
    
    SSL* ssl = tls_get_ssl(fd);
    if (ssl) {
        while(len){
            int n = SSL_write(ssl, charbuf, static_cast<int>(len));
            if(n<=0){
                int err = SSL_get_error(ssl, n);
                if (err == SSL_ERROR_WANT_READ || err == SSL_ERROR_WANT_WRITE)
                    continue;
                return false; // orderly shutdown
            }
            charbuf += n;
            len -= n;
        }
        return true;
    }
    
    while(len){
        ssize_t n = send(fd, charbuf, len, MSG_NOSIGNAL);
        if(n < 0){
            if(errno==EINTR || errno == EAGAIN || errno == EWOULDBLOCK) continue;
            else if(errno==EPIPE) return false; // peer closed conn
            else ERR_THROW("send");
        }
        charbuf += n;
        len -= static_cast<size_t>(n);
    }
    return true;
}


static bool recv_all(int fd, void *buf, size_t len) {
    uint8_t *charbuf = static_cast<uint8_t *>(buf);

    SSL* ssl = tls_get_ssl(fd);
    if (ssl) {
        while(len){
            int n = SSL_read(ssl, charbuf, static_cast<int>(len));
            if(n<=0){
                int err = SSL_get_error(ssl, n);
                if (err == SSL_ERROR_WANT_READ || err == SSL_ERROR_WANT_WRITE)
                    continue;
                return false; // orderly shutdown
            }
            charbuf += n;
            len -= n;
        }
        return true;
    }

    while(len){
        ssize_t n = recv(fd, charbuf, len, MSG_WAITALL);
        if(n < 0){
            if(errno==EINTR || errno == EAGAIN || errno == EWOULDBLOCK) continue;
            else ERR_THROW("recv");
        }else if (n == 0){
            return false; // peer closed conn
        }
        charbuf += n;
        len -= static_cast<size_t>(n);
    }
    return true;
}

bool send_packet(int fd, PacketHeader &h, const void *body, size_t body_len){
    h.len = htonl(h.len);

    // send header
    if(!send_all(fd, static_cast<const void *>(&h), sizeof(h))) return false;
    // send body (if needed)
    if(body_len > 0) 
        if(!send_all(fd, body, body_len)) return false;
    return true;
}

bool recv_packet(int fd, PacketHeader &h, void *(&body)){
    if (!recv_all(fd, &h, sizeof(h))) return false; // peer closed conn
    h.len = ntohl(h.len);

    void *body_tmp = nullptr;
    if(h.len > 0){
        if((body_tmp = malloc(h.len)) == NULL) ERR_THROW("malloc");
        if(!recv_all(fd, body_tmp, h.len)){
            free(body_tmp);
            return false;
        }
    }
    body = body_tmp;
    return true;
}

bool send_response(int fd, Cmd cmd, Status stat, const void *body, size_t body_len){
    // construct header
    assert(stat!=STATUS_REQ);
    PacketHeader h {
        .cmd = cmd,
        .status = stat,
        .len = static_cast<uint32_t>(body_len)
    };
    return send_packet(fd, h, body, body_len); 
}

bool send_request(int fd, Cmd cmd, const void *body, size_t body_len){
    // construct header
    PacketHeader h {
        .cmd = cmd,
        .status = STATUS_REQ,
        .len = static_cast<uint32_t>(body_len)
    };
    return send_packet(fd, h, body, body_len);
}

/* receives a response to a previous `cmd` request */
bool recv_response(int fd, Cmd cmd, Status &stat, void *(&body)){
    PacketHeader h;
    void *body_tmp;
    while(114514){
        if(!recv_packet(fd, h, body_tmp)) return false;
        if(h.cmd == cmd && h.status != STATUS_REQ) break;
        free(body_tmp);
    }
    stat = h.status;
    body = body_tmp;
    return true;
}

/* receives a request */
bool recv_request(int fd, Cmd &cmd, void *(&body)){
    PacketHeader h;
    void *body_tmp;
    while(114514){
        if(!recv_packet(fd, h, body_tmp)) return false;
        if(h.status == STATUS_REQ) break;
        free(body_tmp);
    }
    cmd = h.cmd;
    body = body_tmp;
    return true;
}