#include "server.h"
#include "network/tcp-chat-protocol/v1.h"

void ClientConnection_close(ClientConnection* conn){
    if(conn == NULL)
        return;
    socket_close(conn->sock.sock);
    free(conn->session_key.data);
    free(conn);
}

Result(ClientConnection*) ClientConnection_accept(ServerCredentials* server_credentials, 
    Socket sock_tcp, EndpointIPv4 client_end, u64 session_id)
{
    Deferral(32);

    ClientConnection* conn = (ClientConnection*)malloc(sizeof(ClientConnection));
    memset(conn, 0, sizeof(*conn));
    bool success = false;
    Defer(
        if(!success)
            ClientConnection_close(conn);
    );

    conn->client_end = client_end;
    conn->session_id = session_id;
    conn->session_key = Array_alloc_size(AES_SESSION_KEY_SIZE);

    Array(u8) enc_buf = Array_alloc_size(8*1024);
    Defer(free(enc_buf.data));
    Array(u8) dec_buf = Array_alloc_size(8*1024);
    Defer(free(dec_buf.data));
    u32 enc_size = 0, dec_size = 0;
    
    // TODO: set socket timeout to 5 seconds

    // receive message encrypted by server public key
    try(enc_size, u, 
        socket_recv(
            sock_tcp,
            Array_sliceBefore(enc_buf, server_credentials->rsa_pk.nlen),
            SocketRecvFlag_WaitAll
        )
    );

    // decrypt the message using server private key
    RSADecryptor rsa_dec;
    RSADecryptor_construct(&rsa_dec, &server_credentials->rsa_sk);
    try(dec_size, u, 
        RSADecryptor_decrypt(
            &rsa_dec,
            Array_sliceBefore(enc_buf, enc_size),
            dec_buf
        )
    );
    
    // validate client handshake
    if(dec_size != sizeof(ClientHandshake)){
        Return RESULT_ERROR_FMT(
            "decrypted message (size: %u) is not a ClientHandshake (size: %u)",
            dec_size, (u32)sizeof(ClientHandshake) 
        );
    }
    ClientHandshake* client_handshake = dec_buf.data;
    try_void(PacketHeader_validateMagic(&client_handshake->header));
    if(client_handshake->header.type != PacketType_ClientHandshake){
        Return RESULT_ERROR_FMT(
            "received message of unexpected type: %u",
            client_handshake->header.type
        );
    }

    // use received session key
    memcpy(conn->session_key.data, client_handshake->session_key, conn->session_key.size);
    EncryptedSocketTCP_construct(&conn->sock, sock_tcp, conn->session_key);
    
    // construct ServerHandshake in dec_buf
    ServerHandshake_construct((ServerHandshake*)dec_buf.data, session_id);
    // send ServerHandshake over encrypted TCP socket
    try_void(
        EncryptedSocketTCP_send(
            &conn->sock,
            Array_sliceBefore(dec_buf, sizeof(ServerHandshake)),
            enc_buf
        )
    );

    success = true;
    Return RESULT_VALUE(p, conn);
}