#include "encrypted_sockets.h"

void EncryptedSocketTCP_construct(EncryptedSocketTCP* ptr, 
    Socket sock, u32 crypto_buffer_size, Array(u8) aes_key)
{
    ptr->sock = sock;
    AESStreamEncryptor_construct(&ptr->enc, aes_key, AESStream_DEFAULT_CLASS);
    AESStreamDecryptor_construct(&ptr->dec, aes_key, AESStream_DEFAULT_CLASS);
    ptr->recv_buf = Array_alloc_size(crypto_buffer_size);
    ptr->send_buf = Array_alloc_size(crypto_buffer_size);
}

void EncryptedSocketTCP_destroy(EncryptedSocketTCP* ptr){
    socket_close(ptr->sock);
    free(ptr->recv_buf.data);
    free(ptr->send_buf.data);
}

Result(void) EncryptedSocketTCP_send(EncryptedSocketTCP* ptr, 
    Array(u8) buffer)
{
    Deferral(1);

    try(u32 encrypted_size, u,
        AESStreamEncryptor_encrypt(
            &ptr->enc,
            buffer,
            ptr->send_buf
        )
    );
    try_void(
        socket_send(
            ptr->sock,
            Array_sliceBefore(ptr->send_buf, encrypted_size)
        )
    );

    Return RESULT_VOID;
}

Result(u32) EncryptedSocketTCP_recv(EncryptedSocketTCP* ptr, 
    Array(u8) buffer, SocketRecvFlag flags)
{
    Deferral(1);

    u32 size_to_receive = buffer.size;
    if(ptr->dec.block_counter == 0){
        // There is some metadata at the beginning of AES stream
        size_to_receive = AESStreamEncryptor_calcDstSize(size_to_receive);
    }
    try(i32 received_size, i,
        socket_recv(
            ptr->sock,
            Array_sliceBefore(ptr->recv_buf, size_to_receive),
            flags
        )
    );
    try(u32 decrypted_size, u,
        AESStreamDecryptor_decrypt(
            &ptr->dec,
            Array_sliceBefore(ptr->recv_buf, received_size),
            buffer
        )
    );

    Return RESULT_VALUE(u, decrypted_size);
}



void EncryptedSocketUDP_construct(EncryptedSocketUDP* ptr, 
    Socket sock, u32 crypto_buffer_size, Array(u8) aes_key)
{
    ptr->sock = sock;
    AESBlockEncryptor_construct(&ptr->enc, aes_key, AESBlockEncryptor_DEFAULT_CLASS);
    AESBlockDecryptor_construct(&ptr->dec, aes_key, AESBlockDecryptor_DEFAULT_CLASS);
    ptr->recv_buf = Array_alloc_size(crypto_buffer_size);
    ptr->send_buf = Array_alloc_size(crypto_buffer_size);
}

void EncryptedSocketUDP_destroy(EncryptedSocketUDP* ptr){
    socket_close(ptr->sock);
    free(ptr->recv_buf.data);
    free(ptr->send_buf.data);
}

Result(void) EncryptedSocketUDP_sendto(EncryptedSocketUDP* ptr, 
    Array(u8) buffer, EndpointIPv4 remote_end)
{
    Deferral(1);

    try(u32 encrypted_size, u, 
        AESBlockEncryptor_encrypt(
            &ptr->enc,
            buffer,
            ptr->send_buf
        )
    );
    try_void(
        socket_sendto(
            ptr->sock, 
            Array_sliceBefore(ptr->send_buf, encrypted_size),
            remote_end
        )
    );

    Return RESULT_VOID;
}

Result(i32) EncryptedSocketUDP_recvfrom(EncryptedSocketUDP* ptr,
    Array(u8) buffer, SocketRecvFlag flags, NULLABLE(EndpointIPv4*) remote_end)
{
    Deferral(1);
    
    // There is some metadata at the start of each AES block
    u32 size_to_receive = AESBlockEncryptor_calcDstSize(buffer.size);
    try(i32 received_size, i,
        socket_recvfrom(
            ptr->sock,
            Array_sliceBefore(ptr->recv_buf, size_to_receive),
            flags,
            remote_end
        )
    );
    try(u32 decrypted_size, u,
        AESBlockDecryptor_decrypt(
            &ptr->dec,
            Array_sliceBefore(ptr->recv_buf, received_size),
            buffer
        )
    );

    Return RESULT_VALUE(u, decrypted_size);
}