protocol.c

#include "protocol.h"
#include <stdlib.h>
#include <string.h>
#include <unistd.h>

const uint8_t BYTE_MASK = 0xFF;
static const uint8_t SHIFT_BYTE1 = 8;
static const uint8_t SHIFT_BYTE2 = 16;
static const uint8_t SHIFT_BYTE3 = 24;

static const int zero = 0;
static const int one = 1;
static const int two = 2;
static const int three = 3;

// convert 32bit interger to little-endian
static void uint32_to_bytes(uint32_t value, unsigned char* bytes)
{
    bytes[zero] = value & BYTE_MASK;
    bytes[one] = (value >> SHIFT_BYTE1) & BYTE_MASK;
    bytes[two] = (value >> SHIFT_BYTE2) & BYTE_MASK;
    bytes[three] = (value >> SHIFT_BYTE3) & BYTE_MASK;
}

// convert back from little-endian to 32bit
static uint32_t bytes_to_uint32(const unsigned char* bytes)
{
    return (uint32_t)bytes[zero] | ((uint32_t)bytes[one] << SHIFT_BYTE1)
            | ((uint32_t)bytes[two] << SHIFT_BYTE2)
            | ((uint32_t)bytes[three] << SHIFT_BYTE3);
}

int send_message(FILE* stream, ProtocolMessage* message)
{
    unsigned char buffer[9];

    // send prefix
    uint32_to_bytes(message->prefix, buffer);
    if (fwrite(buffer, 1, 4, stream) != 4) {
        return -1;
    }

    // send operation
    buffer[0] = message->operation;
    if (fwrite(buffer, 1, 1, stream) != 1) {
        return -1;
    }

    // send image-1 size
    uint32_to_bytes(message->image1Size, buffer);
    if (fwrite(buffer, 1, 4, stream) != 4) {
        return -1;
    }

    // send image-1 data
    if (message->image1Size > 0 && message->image1Data) {
        if (fwrite(message->image1Data, 1, message->image1Size, stream)
                != message->image1Size) {
            return -1;
        }
    }

    // send image-2 for face replace
    if (message->operation == OP_FACE_REPLACE) {
        uint32_to_bytes(message->image2Size, buffer);
        if (fwrite(buffer, 1, 4, stream) != 4) {
            return -1;
        }

        if (message->image2Size > 0 && message->image2Data) {
            if (fwrite(message->image2Data, 1, message->image2Size, stream)
                    != message->image2Size) {
                return -1;
            }
        }
    }

    // need flush the picture to see result
    fflush(stream);
    return 0;
}

int receive_message(FILE* stream, ProtocolMessage* message)
{
    unsigned char buffer[9];

    memset(message, 0, sizeof(ProtocolMessage));

    // read prefix
    if (fread(buffer, 1, 4, stream) != 4) {
        return -1;
    }
    message->prefix = bytes_to_uint32(buffer);

    // read operation
    if (fread(buffer, 1, 1, stream) != 1) {
        return -1;
    }
    message->operation = buffer[0];

    // read image-1 size
    if (fread(buffer, 1, 4, stream) != 4) {
        return -1;
    }
    message->image1Size = bytes_to_uint32(buffer);

    // read image-1 data
    if (message->image1Size > 0) {
        message->image1Data = malloc(message->image1Size);
        if (!message->image1Data) {
            return -1;
        }
        if (fread(message->image1Data, 1, message->image1Size, stream)
                != message->image1Size) {
            free(message->image1Data);
            message->image1Data = NULL;
            return -1;
        }
    }

    // read image-1 for face replacement
    if (message->operation == OP_FACE_REPLACE) {
        if (fread(buffer, 1, 4, stream) != 4) {
            free_message(message);
            return -1;
        }
        message->image2Size = bytes_to_uint32(buffer);

        if (message->image2Size > 0) {
            message->image2Data = malloc(message->image2Size);
            if (!message->image2Data) {
                free_message(message);
                return -1;
            }
            if (fread(message->image2Data, 1, message->image2Size, stream)
                    != message->image2Size) {
                free_message(message);
                return -1;
            }
        }
    }

    return 0;
}

// free memory
void free_message(ProtocolMessage* message)
{
    if (message->image1Data) {
        free(message->image1Data);
        message->image1Data = NULL;
    }
    if (message->image2Data) {
        free(message->image2Data);
        message->image2Data = NULL;
    }
    message->image1Size = 0;
    message->image2Size = 0;
}

// read file into memory
unsigned char* read_file_data(const char* filename, uint32_t* size)
{
    FILE* file = fopen(filename, "rb");
    if (!file) {
        return NULL;
    }

    // get the file size
    if (fseek(file, 0, SEEK_END) != 0) {
        fclose(file);
        return NULL;
    }

    long fileSize = ftell(file);
    if (fileSize < 0 || fileSize > UINT32_MAX) {
        fclose(file);
        return NULL;
    }

    if (fseek(file, 0, SEEK_SET) != 0) {
        fclose(file);
        return NULL;
    }

    // allocate and read the data
    unsigned char* data = malloc(fileSize);
    if (!data) {
        fclose(file);
        return NULL;
    }

    if (fread(data, 1, fileSize, file) != (size_t)fileSize) {
        free(data);
        fclose(file);
        return NULL;
    }

    fclose(file);
    *size = (uint32_t)fileSize;
    return data;
}

// read stdin into memory
unsigned char* read_stdin_data(uint32_t* size)
{
    unsigned char* data = NULL;
    uint32_t capacity = 1024;
    uint32_t totalSize = 0;

    data = malloc(capacity);
    if (!data) {
        return NULL;
    }

    size_t bytesRead;
    while ((bytesRead = fread(data + totalSize, 1, capacity - totalSize, stdin))
            > 0) {
        totalSize += bytesRead;

        if (totalSize == capacity) {
            capacity *= 2;
            unsigned char* newData = realloc(data, capacity);
            if (!newData) {
                free(data);
                return NULL;
            }
            data = newData;
        }
    }

    if (totalSize == 0) {
        free(data);
        return NULL;
    }

    *size = totalSize;
    return data;
}

// write data to file
// REF: writeat.c
int write_file_data(const char* filename, unsigned char* data, uint32_t size)
{
    FILE* file = fopen(filename, "wb");
    if (!file) {
        return -1;
    }

    if (fwrite(data, 1, size, file) != size) {
        fclose(file);
        return -1;
    }

    fclose(file);
    return 0;
}

// write data to the stdout
int write_stdout_data(unsigned char* data, uint32_t size)
{
    if (fwrite(data, 1, size, stdout) != size) {
        return -1;
    }
    fflush(stdout);
    return 0;
}