cryptography rework and beginning of tcp-chat-protocol

dependencies/tlibc

@@ -1 +1 @@
-Subproject commit 0184d2e8c96d882b815eae05fa160a40e9f7faf2
+Subproject commit a0affaa6d0e4764c9a906bba516c1f7d1ded5405

src/client/ClientCredential.c

@@ -41,8 +41,8 @@ Result(ClientCredential*) ClientCredential_create(str username, str password){
     // lvl 1 hash - is used as AES key for user data
     hash_password(password_and_username, cred->aes_key.data, __PASSWORD_HASH_LVL_ITERATIONS);
 
-    DecryptorAES_construct(&cred->user_data_aes_dec, cred->aes_key);
+    AESBlockEncryptor_construct(&cred->user_data_aes_enc, cred->aes_key, AESBlockEncryptor_DEFAULT_CLASS);
-    EncryptorAES_construct(&cred->user_data_aes_enc, cred->aes_key);
+    AESBlockDecryptor_construct(&cred->user_data_aes_dec, cred->aes_key, AESBlockDecryptor_DEFAULT_CLASS);
 
     success = true;
     Return RESULT_VALUE(p, cred);

src/client/ServerConnection.c

@@ -1,4 +1,5 @@
 #include "client.h"
+#include "network/tcp-chat-protocol/v1.h"
 
 void ServerConnection_close(ServerConnection* conn){
     if(conn == NULL)
@@ -53,34 +54,70 @@ Result(ServerConnection*) ServerConnection_open(ClientCredential* client_credent
     );
 
     try_void(ServerLink_parse(server_link_cstr, &conn->server_end, &conn->server_pk));
+    RSAEncryptor_construct(&conn->rsa_enc, &conn->server_pk);
     
-    conn->session_key = Array_alloc_size(__AES_SESSION_KEY_SIZE);
+    conn->session_key = Array_alloc_size(AES_SESSION_KEY_SIZE);
     br_hmac_drbg_context key_rng = { .vtable = &br_hmac_drbg_vtable };
     rng_init_sha256_seedFromSystem(&key_rng.vtable);
     br_hmac_drbg_generate(&key_rng, conn->session_key.data, conn->session_key.size);
 
-    printf("connecting to server %s\n", server_link_cstr);
+    // connect to server address
-    try(Socket initial_socket, i, socket_open_TCP());
+    Socket _s;
-    try_void(socket_connect(initial_socket, conn->server_end));
+    try(_s, i, socket_open_TCP());
-    //TODO: add log
+    EncryptedSocketTCP_construct(&conn->system_socket, _s, conn->session_key);
-    EncryptorRSA_construct(&conn->rsa_enc, &conn->server_pk);
+    try_void(socket_connect(conn->system_socket.sock, conn->server_end));
-    // ClientHandshake client_handshake;
-    // Array(u8) encrypted_buf = Array_alloc_size(EncryptorAES_calcDstSize(sizeof(client_handshake)));
-    // u32 encrypted_size = 0;
-    // EncryptorRSA_encrypt(&conn->rsa_enc, struct_castTo_Array(&client_handshake), encrypted_buf, &encrypted_size);
-
-    // send handshake and session key to server
-    // request server info
-    // show server info
-    // save server info to user's db
-    // request log in
-    // if not registered, request registration and then log in
-
-    EncryptedSocket_construct(&conn->system_socket, initial_socket, conn->session_key);
     
-    try(Socket _s, i, socket_open_TCP());
+    Array(u8) encrypted_buf = Array_alloc_size(64*1024);
-    //try_void(socket_connect(_s, conn->server_end????));
+    Defer(free(encrypted_buf.data));
-    EncryptedSocket_construct(&conn->content_socket, _s, conn->session_key);
+    Array(u8) decrypted_buf = Array_alloc_size(64*1024);
+    Defer(free(decrypted_buf.data));
+    u32 encrypted_size = 0, decrypted_size = 0;
+
+    // send handshake to the server
+    ClientHandshake client_handshake;
+    try_void(ClientHandshake_tryConstruct(&client_handshake, conn->session_key));
+    // encryption by server public key
+    try(encrypted_size, u, RSAEncryptor_encrypt(&conn->rsa_enc, 
+        struct_castTo_Array(&client_handshake), 
+        encrypted_buf));
+    try_void(socket_send(conn->system_socket.sock, 
+        Array_construct_size(encrypted_buf.data, encrypted_size)));
+    
+    // receive server response
+    encrypted_size = sizeof(PacketHeader);
+    try(decrypted_size, u, EncryptedSocketTCP_recv(&conn->system_socket,
+        Array_construct_size(encrypted_buf.data, sizeof(PacketHeader)),
+        decrypted_buf,
+        SocketRecvFlag_WaitAll));
+    try_assert(decrypted_size == sizeof(PacketHeader));
+    PacketHeader* packet_header = decrypted_buf.data;
+    try_void(PacketHeader_validateMagic(packet_header));
+    
+    switch(packet_header->type){
+        case PacketType_ErrorMessage:
+            Array(u8) err_buf = Array_alloc_size(packet_header->content_size + 1);
+            bool err_msg_return = false;
+            Defer(
+                if(!err_msg_return)
+                    free(err_buf.data);
+            );
+            ((u8*)err_buf.data)[packet_header->content_size] = 0;
+            try_void(EncryptedSocketTCP_recv(&conn->system_socket, encrypted_buf, err_buf, SocketRecvFlag_WaitAll));
+            err_msg_return = true;
+            Return RESULT_ERROR((char*)err_buf.data, true);
+            break;
+        case PacketType_ClientHandshake:
+            //TODO: receive the rest of the struct
+            break;
+        default:
+            Return RESULT_ERROR_FMT("unexpected response type: %i", packet_header->type);
+    }
+    
+    
+    try(_s, i, socket_open_TCP());
+    EncryptedSocketTCP_construct(&conn->content_socket, _s, conn->session_key);
+    // TODO: how to connect the second socket to the server and associate it with current session?
+    //try_void(socket_connect(conn->content_socket.sock, conn->server_end????));
 
     success = true;
     Return RESULT_VALUE(p, conn);

src/client/client.c

@@ -102,12 +102,21 @@ static Result(void) commandExec(str command, bool* stop){
         );
     }
     else if (is_alias("j") || is_alias("join")){
+        ServerConnection_close(_server_connection);
+
         puts("Enter server address (ip:port): ");
         fgets(answer_buf, answer_buf_size, stdin);
         str new_server_link = str_from_cstr(answer_buf);
         str_trim(&new_server_link, true);
-        ServerConnection_close(_server_connection);
+
+        printf("connecting to server %s\n", new_server_link.data);
         try(_server_connection, p, ServerConnection_open(_client_credential, new_server_link.data));
+
+        // TODO: request server info
+        // show server info
+        // save server info to user's db
+        // request log in
+        // if not registered, request registration and then log in
     }
     else if(is_alias("c") || is_alias("connect")){
         // TODO: read saved servers from database

src/client/client.h

@@ -1,26 +1,28 @@
 #pragma once
-#include "cryptography/cryptography.h"
+#include "cryptography/AES.h"
-#include "network/EncryptedSocket.h"
+#include "cryptography/RSA.h"
+#include "network/encrypted_sockets.h"
 
 Result(void) client_run();
 
 typedef struct ClientCredential {
     str username;
     Array(u8) aes_key;
-    EncryptorAES user_data_aes_enc;
+    AESBlockEncryptor user_data_aes_enc;
-    DecryptorAES user_data_aes_dec;
+    AESBlockDecryptor user_data_aes_dec;
 } ClientCredential;
 
 Result(ClientCredential*) ClientCredential_create(str username, str password);
 void ClientCredential_free(ClientCredential* cred);
 
 typedef struct ServerConnection {
+    u64 session_id;
     EndpointIPv4 server_end;
     br_rsa_public_key server_pk;
-    EncryptorRSA rsa_enc;
+    RSAEncryptor rsa_enc;
     Array(u8) session_key;
-    EncryptedSocket system_socket;
+    EncryptedSocketTCP system_socket;
-    EncryptedSocket content_socket;
+    EncryptedSocketTCP content_socket;
 } ServerConnection;
 
 Result(ServerConnection*) ServerConnection_open(ClientCredential* client_credential, cstr server_link_cstr);

src/cryptography/AES.c

@@ -1,51 +1,56 @@
-#include "cryptography.h"
+#include "AES.h"
 #include <assert.h>
 
 // write data from src to array and increment array data pointer
-static inline void __Array_writeNext(Array(u8)* dst, void* src, size_t size){
+static inline void __Array_writeNext(Array(u8)* dst, u8* src, size_t size){
     memcpy(dst->data, src, size);
     dst->data = (u8*)dst->data + size;
     dst->size -= size;
 }
 
 // read data from array to dst and increment array data pointer
-static inline void __Array_readNext(void* dst, Array(u8)* src, size_t size){
+static inline void __Array_readNext(u8* dst, Array(u8)* src, size_t size){
     memcpy(dst, src->data, size);
     src->data = (u8*)src->data + size;
     src->size -= size;
 }
 
 
-void EncryptorAES_construct(EncryptorAES* ptr, Array(u8) key){
+
+void AESBlockEncryptor_construct(AESBlockEncryptor* ptr, Array(u8) key, const br_block_cbcenc_class* enc_class){
     assert(key.size == 16 || key.size == 24 || key.size == 32);
     
-    //TODO: use AES CTR encryption?
+    ptr->enc_class = enc_class;
-    br_aes_ct64_cbcenc_init(&ptr->enc_ctx, key.data, key.size);
+    ptr->enc_class->init((void*)ptr->enc_keys, key.data, key.size);
-
+    
     ptr->rng_ctx.vtable = &br_hmac_drbg_vtable;
     rng_init_sha256_seedFromSystem(&ptr->rng_ctx.vtable);
 }
 
-void EncryptorAES_encrypt(EncryptorAES* ptr, Array(u8) src, Array(u8) dst){
+Result(u32) AESBlockEncryptor_encrypt(AESBlockEncryptor* ptr, Array(u8) src, Array(u8) dst){
-    assert(dst.size >= EncryptorAES_calcDstSize(src.size));
+    Deferral(4);
+    u32 encrypted_size = AESBlockEncryptor_calcDstSize(src.size);
+    try_assert(dst.size >= encrypted_size);
     
     // generate random initial vector
-    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->iv, __AES_IV_SIZE);
+    br_hmac_drbg_generate(&ptr->rng_ctx, ptr->iv, __AES_BLOCK_IV_SIZE);
     // write IV to the beginning of dst
-    __Array_writeNext(&dst, ptr->iv, __AES_IV_SIZE);
+    __Array_writeNext(&dst, ptr->iv, __AES_BLOCK_IV_SIZE);
 
     const EncryptedBlockHeader header = { .padding_size = 16 - src.size % 16 };
     // write header to buffer
     memcpy(ptr->buf, &header, sizeof(header));
     // encrypt header
-    br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, sizeof(header));
+    ptr->enc_class->run((void*)ptr->enc_keys, ptr->iv, ptr->buf, sizeof(header));
     // write encrypted header to dst
     __Array_writeNext(&dst, ptr->buf, sizeof(header));
 
-    // encrypt full EncryptorAES buffers
+    // encrypt full AESBlockEncryptor buffers
     while(src.size > __AES_BUFFER_SIZE){
         __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
-        br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, __AES_BUFFER_SIZE);
+        ptr->enc_class->run((void*)ptr->enc_keys,
+            ptr->iv,
+            ptr->buf, __AES_BUFFER_SIZE);
         __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
     }
 
@@ -54,47 +59,152 @@ void EncryptorAES_encrypt(EncryptorAES* ptr, Array(u8) src, Array(u8) dst){
         memcpy(ptr->buf, src.data, src.size);
         u32 src_size_padded = src.size + header.padding_size;
         memset(ptr->buf + src.size, 0, header.padding_size);
-        br_aes_ct64_cbcenc_run(&ptr->enc_ctx, ptr->iv, ptr->buf, src_size_padded);
+        ptr->enc_class->run((void*)ptr->enc_keys,
+            ptr->iv,
+            ptr->buf, src_size_padded);
         memcpy(dst.data, ptr->buf, src_size_padded);
     }
+
+    Return RESULT_VALUE(u, encrypted_size);
 }
 
 
-void DecryptorAES_construct(DecryptorAES* ptr, Array(u8) key){
+
+void AESBlockDecryptor_construct(AESBlockDecryptor* ptr, Array(u8) key, const br_block_cbcdec_class* dec_class){
     assert(key.size == 16 || key.size == 24 || key.size == 32);
     
-    br_aes_ct64_cbcdec_init(&ptr->dec_ctx, key.data, key.size);
+    ptr->dec_class = dec_class;
+    ptr->dec_class->init((void*)ptr->dec_keys, key.data, key.size);
 }
 
-void DecryptorAES_decrypt(DecryptorAES* ptr, Array(u8) src, Array(u8) dst, u32* decrypted_size){
+Result(u32) AESBlockDecryptor_decrypt(AESBlockDecryptor* ptr, Array(u8) src, Array(u8) dst){
-    assert(src.size >= EncryptorAES_calcDstSize(0));
+    Deferral(4);
-    assert(src.size % 16 == 0 && "src must be array of 16-byte blocks");
+    try_assert(src.size >= AESBlockEncryptor_calcDstSize(0));
-    assert(dst.size >= src.size);
+    try_assert(src.size % 16 == 0 && "src must be array of 16-byte blocks");
+    try_assert(dst.size >= src.size);
     
     // read IV from the beginning of src
-    __Array_readNext(ptr->iv, &src, __AES_IV_SIZE);
+    __Array_readNext(ptr->iv, &src, __AES_BLOCK_IV_SIZE);
 
     EncryptedBlockHeader header;
     // read encrypted header from src
-    __Array_readNext(&header, &src, sizeof(header));
+    __Array_readNext((void*)&header, &src, sizeof(header));
     // decrypt header
-    br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, &header, sizeof(header));
+    ptr->dec_class->run((void*)ptr->dec_keys, ptr->iv, &header, sizeof(header));
     // size of decrypted data without padding
-    *decrypted_size = src.size - header.padding_size;
+    u32 decrypted_size = src.size - header.padding_size;
     const u32 src_size_padded = src.size;
-    src.size = *decrypted_size;
+    src.size = decrypted_size;
 
     // decrypt full buffers
     while(src.size > __AES_BUFFER_SIZE){
         __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
-        br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, ptr->buf, __AES_BUFFER_SIZE);
+        ptr->dec_class->run((void*)ptr->dec_keys, 
+            ptr->iv, 
+            ptr->buf, __AES_BUFFER_SIZE);
         __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
     }
     
     // decrypt buffer with remaining data
     if(src.size > 0){
         memcpy(ptr->buf, src.data, src.size);
-        br_aes_ct64_cbcdec_run(&ptr->dec_ctx, ptr->iv, ptr->buf, src_size_padded);
+        ptr->dec_class->run((void*)ptr->dec_keys, 
+            ptr->iv,
+            ptr->buf, src_size_padded);
         memcpy(dst.data, ptr->buf, src.size);
     }
+
+    Return RESULT_VALUE(u, decrypted_size);
+}
+
+
+
+void AESStreamEncryptor_construct(AESStreamEncryptor* ptr, Array(u8) key, const br_block_ctr_class* ctr_class){
+    assert(key.size == 16 || key.size == 24 || key.size == 32);
+    
+    ptr->ctr_class = ctr_class;
+    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.size);
+    
+    br_hmac_drbg_context rng_ctx;
+    rng_ctx.vtable = &br_hmac_drbg_vtable;
+    rng_init_sha256_seedFromSystem(&rng_ctx.vtable);
+    br_hmac_drbg_generate(&rng_ctx, ptr->iv, __AES_STREAM_IV_SIZE);
+
+    ptr->block_counter = 0;
+}
+
+Result(u32) AESStreamEncryptor_encrypt(AESStreamEncryptor* ptr, Array(u8) src, Array(u8) dst){
+    Deferral(4);
+    u32 encrypted_size = AESStreamEncryptor_calcDstSize(src.size);
+    try_assert(dst.size >= encrypted_size);
+    
+    // if it is the beginning of the stream, write IV
+    if(ptr->block_counter == 0){
+        __Array_writeNext(&dst, ptr->iv, __AES_STREAM_IV_SIZE);
+    }
+    
+    // encrypt full buffers
+    while(src.size > __AES_BUFFER_SIZE){
+        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
+        ptr->ctr_class->run((void*)ptr->ctr_keys, 
+            ptr->iv, ptr->block_counter,
+            ptr->buf, __AES_BUFFER_SIZE);
+        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
+    }
+
+    // encrypt remaining data
+    if(src.size > 0){
+        memcpy(ptr->buf, src.data, src.size);
+        ptr->ctr_class->run((void*)ptr->ctr_keys,
+            ptr->iv, ptr->block_counter,
+            ptr->buf, src.size);
+        memcpy(dst.data, ptr->buf, src.size);
+    }
+
+    ptr->block_counter++;
+    Return RESULT_VALUE(u, encrypted_size);
+}
+
+
+
+void AESStreamDecryptor_construct(AESStreamDecryptor* ptr, Array(u8) key, const br_block_ctr_class* ctr_class){
+    assert(key.size == 16 || key.size == 24 || key.size == 32);
+    
+    ptr->ctr_class = ctr_class;
+    ptr->ctr_class->init((void*)ptr->ctr_keys, key.data, key.size);
+
+    ptr->block_counter = 0;
+}
+
+Result(u32) AESStreamDecryptor_decrypt(AESStreamDecryptor* ptr, Array(u8) src, Array(u8) dst){
+    Deferral(4);
+    try_assert(src.size >= AESStreamEncryptor_calcDstSize(0));
+    u32 decrypted_size = __AESStreamDecryptor_calcDstSize(src.size);
+    try_assert(dst.size >= decrypted_size);
+
+    // if it is the beginning of the stream, read IV
+    if(ptr->block_counter == 0){
+        __Array_readNext(ptr->iv, &src, __AES_STREAM_IV_SIZE);
+    }
+
+    // decrypt full buffers
+    while(src.size > __AES_BUFFER_SIZE){
+        __Array_readNext(ptr->buf, &src, __AES_BUFFER_SIZE);
+        ptr->ctr_class->run((void*)ptr->ctr_keys, 
+            ptr->iv, ptr->block_counter,
+            ptr->buf, __AES_BUFFER_SIZE);
+        __Array_writeNext(&dst, ptr->buf, __AES_BUFFER_SIZE);
+    }
+
+    // decrypt remaining data
+    if(src.size > 0){
+        memcpy(ptr->buf, src.data, src.size);
+        ptr->ctr_class->run((void*)ptr->ctr_keys,
+            ptr->iv, ptr->block_counter,
+            ptr->buf, src.size);
+        memcpy(dst.data, ptr->buf, src.size);
+    }
+
+    ptr->block_counter++;
+    Return RESULT_VALUE(u, decrypted_size);
 }

src/cryptography/AES.h

@@ -0,0 +1,122 @@
+#pragma once
+#include "tlibc/collections/Array.h"
+#include "tlibc/errors.h"
+#include "bearssl_block.h"
+#include "cryptography.h"
+
+//////////////////////////////////////////////////////////////////////////////
+//                                                                          //
+//                                  AES.c                                   //
+//                                                                          //
+//////////////////////////////////////////////////////////////////////////////
+
+#define AESBlockEncryptor_DEFAULT_CLASS (&br_aes_ct64_cbcenc_vtable)
+#define AESBlockDecryptor_DEFAULT_CLASS (&br_aes_ct64_cbcdec_vtable)
+#define AESStream_DEFAULT_CLASS (&br_aes_ct64_ctr_vtable)
+
+
+//TODO: use PKS#7 instead of this garbage
+typedef struct EncryptedBlockHeader {
+    u8 padding_size;
+} __attribute__((aligned(16))) EncryptedBlockHeader;
+
+//////////////////////////////////////////////////////////////////////////////
+//                            AESBlockEncryptor                             //
+//////////////////////////////////////////////////////////////////////////////
+
+#define __AES_BLOCK_IV_SIZE 16
+// must be multiple of 16
+#define __AES_BUFFER_SIZE 512
+
+typedef struct AESBlockEncryptor {
+    const br_block_cbcenc_class* enc_class;
+    u8 enc_keys[sizeof(br_aes_ct64_cbcenc_keys)];
+    u8 buf[__AES_BUFFER_SIZE];
+    u8 iv[__AES_BLOCK_IV_SIZE];
+    br_hmac_drbg_context rng_ctx;
+} AESBlockEncryptor;
+
+/// @param key supported sizes: 16, 24, 32
+/// @param enc_class &br_aes_XXX_cbcenc_vtable
+void AESBlockEncryptor_construct(AESBlockEncryptor* ptr, Array(u8) key, const br_block_cbcenc_class* enc_class);
+
+/// @brief Encrypts a complete message. For part-by-part encryption use AESStreamEncryptor.
+/// @param src array of any size
+/// @param dst array of size >= AESBlockEncryptor_calcDstSize(src.size)
+/// @return size of encrypted data
+Result(u32) AESBlockEncryptor_encrypt(AESBlockEncryptor* ptr, Array(u8) src, Array(u8) dst);
+
+#define AESBlockEncryptor_calcDstSize(SRC_SIZE) (__AES_BLOCK_IV_SIZE + sizeof(EncryptedBlockHeader) + ALIGN_TO(SRC_SIZE, 16))
+
+//////////////////////////////////////////////////////////////////////////////
+//                            AESBlockDecryptor                             //
+//////////////////////////////////////////////////////////////////////////////
+
+typedef struct AESBlockDecryptor {
+    const br_block_cbcdec_class* dec_class;
+    u8 dec_keys[sizeof(br_aes_ct64_cbcdec_keys)];
+    u8 buf[__AES_BUFFER_SIZE];
+    u8 iv[__AES_BLOCK_IV_SIZE];
+} AESBlockDecryptor;
+
+/// @param key supported sizes: 16, 24, 32
+/// @param dec_class &br_aes_XXX_cbcdec_vtable
+void AESBlockDecryptor_construct(AESBlockDecryptor* ptr, Array(u8) key, const br_block_cbcdec_class* dec_class);
+
+/// @brief Decrypts a complete message. For part-by-part decryption use AESStreamEncryptor.
+/// @param src array of size at least AESBlockEncryptor_calcDstSize(0). Size must be multiple of 16.
+/// @param dst array of size >= src.size
+/// @return size of decrypted data
+Result(u32) AESBlockDecryptor_decrypt(AESBlockDecryptor* ptr, Array(u8) src, Array(u8) dst);
+
+//////////////////////////////////////////////////////////////////////////////
+//                            AESStreamEncryptor                            //
+//////////////////////////////////////////////////////////////////////////////
+
+#define __AES_STREAM_IV_SIZE 12
+
+typedef struct AESStreamEncryptor {
+    const br_block_ctr_class* ctr_class;
+    u8 ctr_keys[sizeof(br_aes_ct64_ctr_keys)];
+    u8 buf[__AES_BUFFER_SIZE];
+    u8 iv[__AES_STREAM_IV_SIZE];
+    u32 block_counter;
+} AESStreamEncryptor;
+
+/// @warning If you start from not first block in sequence (example: to continue a previously interrupted operation) set correct values for ptr->block_counter and ptr->iv or encryption/decryption will produce incorrect data.
+/// @param key supported sizes: 16, 24, 32
+/// @param dec_class &br_aes_XXX_ctr_vtable
+void AESStreamEncryptor_construct(AESStreamEncryptor* ptr, Array(u8) key, const br_block_ctr_class* ctr_class);
+
+#define AESStreamEncryptor_calcDstSize(src_size) (src_size + __AES_BLOCK_IV_SIZE)
+
+/// @brief If ptr->block_counter == 0, writes random IV to `dst`. After that writes encrypted data to dst.
+/// @param src array of any size
+/// @param dst array of size >= AESStreamEncryptor_calcDstSize(src.size)
+/// @return size of encrypted data
+Result(u32) AESStreamEncryptor_encrypt(AESStreamEncryptor* ptr, Array(u8) src, Array(u8) dst);
+
+//////////////////////////////////////////////////////////////////////////////
+//                            AESStreamDecryptor                            //
+//////////////////////////////////////////////////////////////////////////////
+
+typedef struct AESStreamDecryptor {
+    const br_block_ctr_class* ctr_class;
+    u8 ctr_keys[sizeof(br_aes_ct64_ctr_keys)];
+    u8 buf[__AES_BUFFER_SIZE];
+    u8 iv[__AES_STREAM_IV_SIZE];
+    u32 block_counter;
+} AESStreamDecryptor;
+
+/// @warning If you start from not first block in sequence (example: to continue a previously interrupted operation) set correct values for ptr->block_counter and ptr->iv or encryption/decryption will produce incorrect data.
+/// @param key supported sizes: 16, 24, 32
+/// @param dec_class &br_aes_XXX_ctr_vtable
+void AESStreamDecryptor_construct(AESStreamDecryptor* ptr, Array(u8) key, const br_block_ctr_class* ctr_class);
+
+/// @brief Reads IV from `src`, then decrypts data and writes it to dst
+/// @param src array of size at least AESStreamEncryptor_calcDstSize(0).
+/// @param dst array of size >= src.size
+/// @return size of decrypted data
+Result(u32) AESStreamDecryptor_decrypt(AESStreamDecryptor* ptr, Array(u8) src, Array(u8) dst);
+
+#define __AESStreamDecryptor_calcDstSize(src_size) (src_size - __AES_BLOCK_IV_SIZE)

src/cryptography/RSA.c

@@ -1,4 +1,4 @@
-#include "cryptography.h"
+#include "RSA.h"
 #include <assert.h>
 #include "bearssl_x509.h"
 #include "bearssl_pem.h"
@@ -173,33 +173,55 @@ Result(void) RSA_parsePrivateKey_base64(const str src, br_rsa_private_key* sk){
     Return RESULT_VOID;
 }
 
-void EncryptorRSA_construct(EncryptorRSA* ptr, const br_rsa_public_key* pk){
+void RSAEncryptor_construct(RSAEncryptor* ptr, const br_rsa_public_key* pk){
     ptr->pk = pk;
     ptr->rng.vtable = &br_hmac_drbg_vtable;
     rng_init_sha256_seedFromSystem(&ptr->rng.vtable);
 }
 
-void EncryptorRSA_encrypt(EncryptorRSA* ptr, Array(u8) src, Array(u8) dst, u32* encrypted_size){
+Result(u32) RSAEncryptor_encrypt(RSAEncryptor* ptr, Array(u8) src, Array(u8) dst){
+    const u32 max_src_size = RSAEncryptor_calcMaxSrcSize(ptr->pk->nlen * 8, 256);
+    if(src.size > max_src_size){  
+        return RESULT_ERROR_FMT("src.size (%u) must be <= %u (use RSAEncryptor_calcMaxSrcSize)", 
+            src.size, max_src_size);
+    }
+    if(dst.size < ptr->pk->nlen){
+        return RESULT_ERROR_FMT("dst.size (%u) must be >= %u (key length in bytes)", 
+            dst.size, (u32)ptr->pk->nlen);
+    }
     size_t sz = br_rsa_i31_oaep_encrypt(
         &ptr->rng.vtable, &br_sha256_vtable, 
         NULL, 0,
         ptr->pk,
         dst.data, dst.size,
         src.data, src.size);
-    *encrypted_size = sz;
+    
+    if(sz == 0){
+        return RESULT_ERROR("RSA encryption failed", false);
+    }
+    return RESULT_VALUE(u, sz);
 }
 
 
-void DecryptorRSA_construct(DecryptorRSA* ptr, const br_rsa_private_key* sk){
+void RSADecryptor_construct(RSADecryptor* ptr, const br_rsa_private_key* sk){
     ptr->sk = sk;
 }
 
-void DecryptorRSA_decrypt(DecryptorRSA* ptr, Array(u8) buf, u32* decrypted_size){
+Result(u32) RSADecryptor_decrypt(RSADecryptor* ptr, Array(u8) buf){
+    if(buf.size != ptr->sk->n_bitlen/8){
+        return RESULT_ERROR_FMT("buf.size (%u) must be == %u (key length in bytes)", 
+            buf.size, ptr->sk->n_bitlen/8);
+    }
+
     size_t sz = buf.size;
-    br_rsa_i31_oaep_decrypt(
+    size_t r = br_rsa_i31_oaep_decrypt(
         &br_sha256_vtable, 
         NULL, 0,
         ptr->sk,
         buf.data, &sz);
-    *decrypted_size = sz;
+    
+    if(r == 0){
+        return RESULT_ERROR("RSA encryption failed", false);
+    }
+    return RESULT_VALUE(u, sz);
 }

src/cryptography/RSA.h

@@ -0,0 +1,105 @@
+#pragma once
+#include "tlibc/collections/Array.h"
+#include "tlibc/errors.h"
+#include "bearssl_rsa.h"
+#include "cryptography.h"
+
+//////////////////////////////////////////////////////////////////////////////
+//                                                                          //
+//                                  RSA.c                                   //
+//                                                                          //
+//////////////////////////////////////////////////////////////////////////////
+
+#define RSA_DEFAULT_KEY_SIZE 3072
+
+/// @brief generate random key pair based on system time
+/// @param key_size size of public key in bits (2048/3072/4096)
+/// @param sk key for decryption
+/// @param pk key for encryption
+/// @param rng_vtable_ptr pointer to vtable field in prng context. The context must be initialized
+Result(void) RSA_generateKeyPair(u32 key_size, 
+    br_rsa_private_key* sk, br_rsa_public_key* pk,
+    const br_prng_class** rng_vtable_ptr);
+
+Result(void) RSA_generateKeyPairFromTime(u32 key_size, 
+    br_rsa_private_key* sk, br_rsa_public_key* pk);
+
+Result(void) RSA_generateKeyPairFromPassword(u32 key_size, 
+    br_rsa_private_key* sk, br_rsa_public_key* pk, str password);
+
+static inline void RSA_destroyPrivateKey(br_rsa_private_key* sk){
+    free(sk->p);
+}
+
+static inline void RSA_destroyPublicKey(br_rsa_public_key* sk){
+    free(sk->n);
+}
+
+/// @param sk some private key 
+/// @param pk out public key. WARNING: .n is allocated on heap
+Result(void) RSA_computePublicKey(const br_rsa_private_key* sk, br_rsa_public_key* pk);
+
+/// @brief Encode key data in human-readable format
+/// @param src some data
+/// @return heap-allocated string
+str RSA_serializePrivateKey_base64(const br_rsa_private_key* sk);
+
+/// @param src serialized private key format "RSA-Private-%SIZE%:%DATA_BASE64%"
+/// @param sk out private key. WARNING: .p is allocated on heap
+Result(void) RSA_parsePrivateKey_base64(const str src, br_rsa_private_key* sk);
+
+/// @brief Encode key data in human-readable format
+/// @param src some data
+/// @return heap-allocated string
+str RSA_serializePublicKey_base64(const br_rsa_public_key* sk);
+
+/// @param src serialized public key format "RSA-Public-%SIZE%:%DATA_BASE64%"
+/// @param sk out public key. WARNING: .p is allocated on heap
+Result(void) RSA_parsePublicKey_base64(const str src, br_rsa_public_key* sk);
+
+//////////////////////////////////////////////////////////////////////////////
+//                               RSAEncryptor                               //
+//////////////////////////////////////////////////////////////////////////////
+
+typedef struct RSAEncryptor {
+    const br_rsa_public_key* pk;
+    br_hmac_drbg_context rng;
+} RSAEncryptor;
+
+/// RSA OAEP encryption with SHA256 hashing algorithm
+void RSAEncryptor_construct(RSAEncryptor* ptr, const br_rsa_public_key* pk);
+
+/*
+https://crypto.stackexchange.com/a/42100
++---------+----------+-------------------------------------------+
+|         | overhead |         max message size (bytes)          |
+| Hash    | (bytes)  | RSA-1024 | RSA-2048 | RSA-3072 | RSA-4096 |
++---------+----------+----------+----------+----------+----------+
+| SHA-1   | 42       | 86       | 214      | 342      | 470      |
+| SHA-224 | 58       | 70       | 198      | 326      | 454      |
+| SHA-256 | 66       | 62       | 190      | 318      | 446      |
+| SHA-384 | 98       | 30       | 158      | 286      | 414      |
+| SHA-512 | 130      | N/A      | 126      | 254      | 382      |
++---------+----------+----------+----------+----------+----------+
+*/
+#define RSAEncryptor_calcMaxSrcSize(KEY_LEN_BITS, HASH_LEN_BITS) (KEY_LEN_BITS / 8 - 2 * HASH_LEN_BITS / 8 - 2)
+
+/// @param src buffer with size <= `RSAEncryptor_calcMaxSrcSize(key_size_bits, 256)`
+/// @param dst buffer with size >= key size in bytes
+/// @return size of encrypted data
+Result(u32) RSAEncryptor_encrypt(RSAEncryptor* ptr, Array(u8) src, Array(u8) dst);
+
+//////////////////////////////////////////////////////////////////////////////
+//                               RSADecryptor                               //
+//////////////////////////////////////////////////////////////////////////////
+
+typedef struct RSADecryptor {
+    const br_rsa_private_key* sk;
+} RSADecryptor;
+
+/// RSA OAEP encryption with SHA256 hashing algorithm
+void RSADecryptor_construct(RSADecryptor* ptr, const br_rsa_private_key* sk);
+
+/// @param buf buffer with size == key size in bytes
+/// @return size of decrypted data
+Result(u32) RSADecryptor_decrypt(RSADecryptor* ptr, Array(u8) buf);

src/cryptography/cryptography.h

@@ -1,14 +1,12 @@
 #pragma once
-#include "tlibc/std.h"
 #include "tlibc/collections/Array.h"
-#include "tlibc/string/str.h"
 #include "tlibc/errors.h"
-#include "bearssl_block.h"
 #include "bearssl_rand.h"
-#include "bearssl_rsa.h"
 
 //////////////////////////////////////////////////////////////////////////////
+//                                                                          //
 //                                  hash.c                                  //
+//                                                                          //
 //////////////////////////////////////////////////////////////////////////////
 
 /// @brief hashes password multiple times using its own hash as salt 
@@ -21,7 +19,9 @@ void hash_password(Array(u8) password, u8* out_buffer, i32 iterations);
 #define __PASSWORD_HASH_LVL_ITERATIONS 1e5
 
 //////////////////////////////////////////////////////////////////////////////
-//                                  rng.c                                  //
+//                                                                          //
+//                                  rng.c                                   //
+//                                                                          //
 //////////////////////////////////////////////////////////////////////////////
 
 /// @brief Initialize prng context with sha256 hashing algorithm 
@@ -42,120 +42,3 @@ void rng_init_sha256_seedFromSystem(const br_prng_class** rng_vtable_ptr);
 /// rng_init_sha256_seedFromTime(&rng_ctx.vtable);
 /// ```
 void rng_init_sha256_seedFromTime(const br_prng_class** rng_vtable_ptr);
-
-
-//////////////////////////////////////////////////////////////////////////////
-//                                  AES.c                                   //
-//////////////////////////////////////////////////////////////////////////////
-
-#define __AES_SESSION_KEY_SIZE 32
-#define __AES_DB_KEY_SIZE 32
-
-typedef struct EncryptedBlockHeader {
-    u8 padding_size;
-} __attribute__((aligned(16))) EncryptedBlockHeader;
-
-// must be multiple of 16
-#define __AES_IV_SIZE 16
-// must be multiple of 16
-#define __AES_BUFFER_SIZE 512
-
-typedef struct EncryptorAES {
-    br_aes_ct64_cbcenc_keys enc_ctx;
-    br_hmac_drbg_context rng_ctx;
-    u8 iv[__AES_IV_SIZE];
-    u8 buf[__AES_BUFFER_SIZE];
-} EncryptorAES;
-
-/// @param key Array<u8, 16 | 24 | 32>
-void EncryptorAES_construct(EncryptorAES* ptr, Array(u8) key);
-
-/// @brief Encrypts `src` and writes output to `dst`.
-/// @param src array of any size
-/// @param dst array of size >= EncryptorAES_calcDstSize(src.size)
-void EncryptorAES_encrypt(EncryptorAES* ptr, Array(u8) src, Array(u8) dst);
-
-#define EncryptorAES_calcDstSize(SRC_SIZE) (__AES_IV_SIZE + sizeof(EncryptedBlockHeader) + ALIGN_TO(SRC_SIZE, 16))
-
-
-typedef struct DecryptorAES {
-    br_aes_ct64_cbcdec_keys dec_ctx;
-    u8 iv[__AES_IV_SIZE];
-    u8 buf[__AES_BUFFER_SIZE];
-} DecryptorAES;
-
-/// @param key Array<u8, 16 | 24 | 32>
-void DecryptorAES_construct(DecryptorAES* ptr, Array(u8) key);
-
-/// @brief Decrypts `src` and writes output to `dst`.
-/// @param src array of size at least EncryptorAES_calcDstSize(0). Size must be multiple of 16.
-/// @param dst array of size >= src.size
-/// @param decrypted_size size of original data without padding added by EncryptorAES_encrypt
-void DecryptorAES_decrypt(DecryptorAES* ptr, Array(u8) src, Array(u8) dst, u32* decrypted_size);
-
-
-//////////////////////////////////////////////////////////////////////////////
-//                                  RSA.c                                   //
-//////////////////////////////////////////////////////////////////////////////
-
-#define __RSA_DEFAULT_KEY_SIZE 3072
-
-/// @brief generate random key pair based on system time
-/// @param key_size size of public key in bits (2048/3072/4096)
-/// @param sk key for decryption
-/// @param pk key for encryption
-/// @param rng_vtable_ptr pointer to vtable field in prng context. The context must be initialized
-Result(void) RSA_generateKeyPair(u32 key_size, 
-    br_rsa_private_key* sk, br_rsa_public_key* pk,
-    const br_prng_class** rng_vtable_ptr);
-
-Result(void) RSA_generateKeyPairFromTime(u32 key_size, 
-    br_rsa_private_key* sk, br_rsa_public_key* pk);
-
-Result(void) RSA_generateKeyPairFromPassword(u32 key_size, 
-    br_rsa_private_key* sk, br_rsa_public_key* pk, str password);
-
-static inline void RSA_destroyPrivateKey(br_rsa_private_key* sk){
-    free(sk->p);
-}
-
-static inline void RSA_destroyPublicKey(br_rsa_public_key* sk){
-    free(sk->n);
-}
-
-/// @param sk some private key 
-/// @param pk out public key. WARNING: .n is allocated on heap
-Result(void) RSA_computePublicKey(const br_rsa_private_key* sk, br_rsa_public_key* pk);
-
-/// @brief Encode key data in human-readable format
-/// @param src some data
-/// @return heap-allocated string
-str RSA_serializePrivateKey_base64(const br_rsa_private_key* sk);
-
-/// @param src serialized private key format "RSA-Private-%SIZE%:%DATA_BASE64%"
-/// @param sk out private key. WARNING: .p is allocated on heap
-Result(void) RSA_parsePrivateKey_base64(const str src, br_rsa_private_key* sk);
-
-/// @brief Encode key data in human-readable format
-/// @param src some data
-/// @return heap-allocated string
-str RSA_serializePublicKey_base64(const br_rsa_public_key* sk);
-
-/// @param src serialized public key format "RSA-Public-%SIZE%:%DATA_BASE64%"
-/// @param sk out public key. WARNING: .p is allocated on heap
-Result(void) RSA_parsePublicKey_base64(const str src, br_rsa_public_key* sk);
-
-
-typedef struct EncryptorRSA {
-    const br_rsa_public_key* pk;
-    br_hmac_drbg_context rng;
-} EncryptorRSA;
-void EncryptorRSA_construct(EncryptorRSA* ptr, const br_rsa_public_key* pk);
-void EncryptorRSA_encrypt(EncryptorRSA* ptr, Array(u8) src, Array(u8) dst, u32* encrypted_size);
-
-typedef struct DecryptorRSA {
-    const br_rsa_private_key* sk;
-} DecryptorRSA;
-
-void DecryptorRSA_construct(DecryptorRSA* ptr, const br_rsa_private_key* sk);
-void DecryptorRSA_decrypt(DecryptorRSA* ptr, Array(u8) buf, u32* decrypted_size);

src/db/idb.c

@@ -1,13 +1,9 @@
 #include "idb.h"
+#include "magic.h"
 #include "tlibc/filesystem.h"
 #include "tlibc/collections/HashMap.h"
 #include <pthread.h>
 
-typedef union Magic32 {
-    u32 n;
-    u8 bytes[4];
-} Magic32;
-
 typedef struct TableFileHeader {
     Magic32 magic;
     u16 version;

src/db/idb.h

@@ -3,6 +3,7 @@
 #include "tlibc/errors.h"
 
 #define IDB_VERSION 1
+#define IDB_AES_KEY_SIZE 32
 
 typedef struct IncrementalDB IncrementalDB;
 typedef struct Table Table;

src/magic.h

@@ -0,0 +1,12 @@
+#pragma once
+#include "tlibc/std.h"
+
+typedef union Magic32 {
+    u32 n;
+    u8 bytes[4];
+} Magic32;
+
+typedef union Magic64 {
+    u64 n;
+    u8 bytes[8];
+} Magic64;

src/main.c

@@ -20,17 +20,18 @@ int main(const int argc, cstr const* argv){
 
     ProgramMode mode = Client;
     cstr server_endpoint_cstr;
-    u32 key_size = 0;
+    u32 key_size = RSA_DEFAULT_KEY_SIZE;
 
     for(int argi = 1; argi < argc; argi++){
         str arg_str = str_from_cstr(argv[argi]);
         if(arg_is("-h") || arg_is("--help")){
             printf(
                 "USAGE:\n"
-                "no arguments                       Interactive client mode.\n"
+                "no arguments                   Interactive client mode.\n"
-                "-h, --help                         Show this message.\n"
+                "-h, --help                     Show this message.\n"
-                "-l, --listen [addr:port]           Start server.\n"
+                "-l, --listen [addr:port]       Start server.\n"
-                "--rsa-gen [size(2048/3072/4096)]   Generate RSA private and public keys based on stdin data (64Kb max)\n"
+                "--rsa-gen [size]               Generate RSA private and public keys based on stdin data (64Kb max)\n"
+                "    size: 2048 / 3072 (default) / 4096\n"
             );
             Return 0;
         }

src/network/EncryptedSocket.c

@@ -1,49 +0,0 @@
-#include "EncryptedSocket.h"
-
-void EncryptedSocket_construct(EncryptedSocket* ptr, Socket sock, Array(u8) aes_key)
-{
-    ptr->sock = sock;
-    EncryptorAES_construct(&ptr->enc, aes_key);
-    DecryptorAES_construct(&ptr->dec, aes_key);
-}
-
-
-Result(void) EncryptedSocket_send(EncryptedSocket* ptr,
-    Array(u8) decrypted_buf, Array(u8) encrypted_buf)
-{
-    Deferral(4);
-    EncryptorAES_encrypt(&ptr->enc, decrypted_buf, encrypted_buf);
-    try_void(socket_send(ptr->sock, encrypted_buf));
-    Return RESULT_VOID;
-}
-
-Result(void) EncryptedSocket_sendto(EncryptedSocket* ptr,
-    Array(u8) decrypted_buf, Array(u8) encrypted_buf,
-    EndpointIPv4 remote_end)
-{
-    Deferral(4);
-    EncryptorAES_encrypt(&ptr->enc, decrypted_buf, encrypted_buf);
-    try_void(socket_sendto(ptr->sock, encrypted_buf, remote_end));
-    Return RESULT_VOID;
-}
-
-Result(i32) EncryptedSocket_recv(EncryptedSocket* ptr,
-    Array(u8) encrypted_buf, Array(u8) decrypted_buf)
-{
-    Deferral(4);
-    try(i32 r, i, socket_recv(ptr->sock, encrypted_buf));
-    encrypted_buf.size = r;
-    DecryptorAES_decrypt(&ptr->dec, encrypted_buf, decrypted_buf, (u32*)&r);
-    Return RESULT_VALUE(i, r);
-}
-
-Result(i32) EncryptedSocket_recvfrom(EncryptedSocket* ptr,
-    Array(u8) encrypted_buf, Array(u8) decrypted_buf,
-    NULLABLE(EndpointIPv4*) remote_end)
-{
-    Deferral(4);
-    try(i32 r, i, socket_recvfrom(ptr->sock, encrypted_buf, remote_end));
-    encrypted_buf.size = r;
-    DecryptorAES_decrypt(&ptr->dec, encrypted_buf, decrypted_buf, (u32*)&r);
-    Return RESULT_VALUE(i, r);
-}

src/network/EncryptedSocket.h

@@ -1,25 +0,0 @@
-#pragma once
-#include "network/socket.h"
-#include "cryptography/cryptography.h"
-
-typedef struct EncryptedSocket {
-    Socket sock;
-    EncryptorAES enc;
-    DecryptorAES dec;
-} EncryptedSocket;
-
-void EncryptedSocket_construct(EncryptedSocket* ptr, Socket sock, Array(u8) aes_key);
-
-Result(void) EncryptedSocket_send(EncryptedSocket* ptr,
-    Array(u8) decrypted_buf, Array(u8) encrypted_buf);
-
-Result(void) EncryptedSocket_sendto(EncryptedSocket* ptr,
-    Array(u8) decrypted_buf, Array(u8) encrypted_buf,
-    EndpointIPv4 remote_end);
-
-Result(i32) EncryptedSocket_recv(EncryptedSocket* ptr,
-    Array(u8) encrypted_buf, Array(u8) decrypted_buf);
-
-Result(i32) EncryptedSocket_recvfrom(EncryptedSocket* ptr,
-    Array(u8) encrypted_buf, Array(u8) decrypted_buf,
-    NULLABLE(EndpointIPv4*) remote_end);

src/network/encrypted_sockets.c

@@ -0,0 +1,61 @@
+#include "encrypted_sockets.h"
+
+void EncryptedSocketTCP_construct(EncryptedSocketTCP* ptr, Socket sock, 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);
+}
+
+
+Result(void) EncryptedSocketTCP_send(EncryptedSocketTCP* ptr,
+    Array(u8) decrypted_buf, Array(u8) encrypted_buf)
+{
+    Deferral(4);
+    try(u32 encrypted_size, u, AESStreamEncryptor_encrypt(&ptr->enc, decrypted_buf, encrypted_buf));
+    encrypted_buf.size = encrypted_size;
+    try_void(socket_send(ptr->sock, encrypted_buf));
+    Return RESULT_VOID;
+}
+
+Result(i32) EncryptedSocketTCP_recv(EncryptedSocketTCP* ptr,
+    Array(u8) encrypted_buf, Array(u8) decrypted_buf,
+    SocketRecvFlag flags)
+{
+    Deferral(4);
+    try(i32 received_size, i, socket_recv(ptr->sock, encrypted_buf, flags));
+    encrypted_buf.size = received_size;
+    try(i32 decrypted_size, u, AESStreamDecryptor_decrypt(&ptr->dec, encrypted_buf, decrypted_buf));
+    Return RESULT_VALUE(i, decrypted_size);
+}
+
+
+
+void EncryptedSocketUDP_construct(EncryptedSocketUDP* ptr, Socket sock, 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);
+}
+
+Result(void) EncryptedSocketUDP_sendto(EncryptedSocketUDP* ptr,
+    Array(u8) decrypted_buf, Array(u8) encrypted_buf,
+    EndpointIPv4 remote_end)
+{
+    Deferral(4);
+    try(u32 encrypted_size, u, AESBlockEncryptor_encrypt(&ptr->enc, decrypted_buf, encrypted_buf));
+    encrypted_buf.size = encrypted_size;
+    try_void(socket_sendto(ptr->sock, encrypted_buf, remote_end));
+    Return RESULT_VOID;
+}
+
+Result(i32) EncryptedSocketUDP_recvfrom(EncryptedSocketUDP* ptr,
+    Array(u8) encrypted_buf, Array(u8) decrypted_buf,
+    SocketRecvFlag flags, NULLABLE(EndpointIPv4*) remote_end)
+{
+    Deferral(4);
+    try(i32 received_size, i, socket_recvfrom(ptr->sock, encrypted_buf, flags, remote_end));
+    encrypted_buf.size = received_size;
+    try(i32 decrypted_size, u, AESBlockDecryptor_decrypt(&ptr->dec, encrypted_buf, decrypted_buf));
+    Return RESULT_VALUE(i, decrypted_size);
+}

src/network/encrypted_sockets.h

@@ -0,0 +1,44 @@
+#pragma once
+#include "network/socket.h"
+#include "cryptography/AES.h"
+
+//////////////////////////////////////////////////////////////////////////////
+//                            EncryptedSocketTCP                            //
+//////////////////////////////////////////////////////////////////////////////
+
+typedef struct EncryptedSocketTCP {
+    Socket sock;
+    AESStreamEncryptor enc;
+    AESStreamDecryptor dec;
+} EncryptedSocketTCP;
+
+void EncryptedSocketTCP_construct(EncryptedSocketTCP* ptr, Socket sock, Array(u8) aes_key);
+
+Result(void) EncryptedSocketTCP_send(EncryptedSocketTCP* ptr,
+    Array(u8) decrypted_buf, Array(u8) encrypted_buf);
+
+Result(void) EncryptedSocketTCP_recv(EncryptedSocketTCP* ptr,
+    Array(u8) encrypted_buf, Array(u8) decrypted_buf,
+    SocketRecvFlag flags);
+
+
+//////////////////////////////////////////////////////////////////////////////
+//                            EncryptedSocketUDP                            //
+//////////////////////////////////////////////////////////////////////////////
+
+typedef struct EncryptedSocketUDP {
+    Socket sock;
+    AESBlockEncryptor enc;
+    AESBlockDecryptor dec;
+} EncryptedSocketUDP;
+
+void EncryptedSocketUDP_construct(EncryptedSocketUDP* ptr, Socket sock, Array(u8) aes_key);
+
+Result(void) EncryptedSocketUDP_sendto(EncryptedSocketUDP* ptr,
+    Array(u8) decrypted_buf, Array(u8) encrypted_buf,
+    EndpointIPv4 remote_end);
+
+Result(i32) EncryptedSocketUDP_recvfrom(EncryptedSocketUDP* ptr,
+    Array(u8) encrypted_buf, Array(u8) decrypted_buf,
+    SocketRecvFlag flags,
+    NULLABLE(EndpointIPv4*) remote_end);

src/network/socket.c

@@ -74,17 +74,26 @@ Result(void) socket_sendto(Socket s, Array(u8) buffer, EndpointIPv4 dst){
     return RESULT_VOID;
 }
 
-Result(i32) socket_recv(Socket s, Array(u8) buffer){
+static inline int SocketRecvFlags_toStd(SocketRecvFlag flags){
-    i32 r = recv(s, buffer.data, buffer.size, 0);
+    int f = 0;
+    if (flags & SocketRecvFlag_Peek)
+        f |= MSG_PEEK;
+    if (flags & SocketRecvFlag_WaitAll)
+        f |= MSG_WAITALL;
+    return f;
+}
+
+Result(i32) socket_recv(Socket s, Array(u8) buffer, SocketRecvFlag flags){
+    i32 r = recv(s, buffer.data, buffer.size, SocketRecvFlags_toStd(flags));
     if(r < 0)
         return RESULT_ERROR_SOCKET();
     return RESULT_VALUE(i, r);
 }
 
-Result(i32) socket_recvfrom(Socket s, Array(u8) buffer, NULLABLE(EndpointIPv4*) remote_end){
+Result(i32) socket_recvfrom(Socket s, Array(u8) buffer, SocketRecvFlag flags, NULLABLE(EndpointIPv4*) remote_end){
     struct sockaddr_in remote_addr = {0};
     i32 sockaddr_size = sizeof(remote_addr);
-    i32 r = recvfrom(s, buffer.data, buffer.size, 0, 
+    i32 r = recvfrom(s, buffer.data, buffer.size, SocketRecvFlags_toStd(flags), 
         (struct sockaddr*)&remote_addr, (void*)&sockaddr_size);
     if(r < 0)
         return RESULT_ERROR_SOCKET();

src/network/socket.h

@@ -2,6 +2,7 @@
 #include "endpoint.h"
 #include "tlibc/errors.h"
 #include "tlibc/collections/Array.h"
+#include "tlibc/time.h"
 
 typedef enum SocketShutdownType {
     SocketShutdownType_Receive = 0,
@@ -9,6 +10,12 @@ typedef enum SocketShutdownType {
     SocketShutdownType_Both = 2,
 } SocketShutdownType;
 
+typedef enum SocketRecvFlag {
+    SocketRecvFlag_None = 0,
+    SocketRecvFlag_Peek = 0b1 /* next recv call will read the same data */,
+    SocketRecvFlag_WaitAll = 0b10  /* waits until buffer is full */, 
+} SocketRecvFlag;
+
 typedef i64 Socket;
 
 Result(Socket)  socket_open_TCP();
@@ -20,5 +27,5 @@ Result(Socket)  socket_accept(Socket s, NULLABLE(EndpointIPv4*) remote_end);
 Result(void)    socket_connect(Socket s, EndpointIPv4 remote_end);
 Result(void)    socket_send(Socket s, Array(u8) buffer);
 Result(void)    socket_sendto(Socket s, Array(u8) buffer, EndpointIPv4 dst);
-Result(i32)     socket_recv(Socket s, Array(u8) buffer);
+Result(i32)     socket_recv(Socket s, Array(u8) buffer, SocketRecvFlag flags);
-Result(i32)     socket_recvfrom(Socket s, Array(u8) buffer, NULLABLE(EndpointIPv4*) remote_end);
+Result(i32)     socket_recvfrom(Socket s, Array(u8) buffer, SocketRecvFlag flags, NULLABLE(EndpointIPv4*) remote_end);

src/network/tcp-chat-protocol/constant.c

@@ -0,0 +1,17 @@
+#include "constant.h"
+
+const Magic64 PacketHeader_MAGIC = { .bytes = { 't', 'c', 'p', '-', 'c', 'h', 'a', 't' } };
+
+Result(void) PacketHeader_validateMagic(PacketHeader* ptr){
+    if (ptr->magic.n != PacketHeader_MAGIC.n){
+        return RESULT_ERROR("invalid packet magic", false);
+    }
+    return RESULT_VOID;
+}
+
+void PacketHeader_construct(PacketHeader* ptr, u8 protocol_version, u16 type, u64 content_size){
+    ptr->magic.n = PacketHeader_MAGIC.n;
+    ptr->protocol_version = protocol_version;
+    ptr->type = type;
+    ptr->content_size = content_size;
+}

src/network/tcp-chat-protocol/constant.h

@@ -0,0 +1,18 @@
+#pragma once
+#include "tlibc/errors.h"
+#include "magic.h"
+
+#define AES_SESSION_KEY_SIZE 32
+
+extern const Magic64 PacketHeader_MAGIC;
+
+typedef struct PacketHeader {
+    Magic64 magic;
+    u8 protocol_version;
+    u8 _reserved;
+    u16 type;
+    u64 content_size;
+} __attribute__((aligned(64))) PacketHeader;
+
+void PacketHeader_construct(PacketHeader* ptr, u8 protocol_version, u16 type, u64 content_size);
+Result(void) PacketHeader_validateMagic(PacketHeader* ptr);

src/network/tcp-chat-protocol/v1.c

@@ -0,0 +1,10 @@
+#include "v1.h"
+
+Result(void) ClientHandshake_tryConstruct(ClientHandshake* ptr, Array(u8) session_key){
+    Deferral(1);
+    try_assert(session_key.size == AES_SESSION_KEY_SIZE);
+    PacketHeader_construct(&ptr->header, PROTOCOL_VERSION, PacketType_ClientHandshake, session_key.size);
+    memcpy(ptr->session_key, session_key.data, session_key.size);
+    Return RESULT_VOID;
+}
+

src/network/tcp-chat-protocol/v1.h

@@ -0,0 +1,30 @@
+#pragma once
+#include "tlibc/errors.h"
+#include "network/tcp-chat-protocol/constant.h"
+
+#define PROTOCOL_VERSION 1 /* 1.0.0 */
+
+typedef enum PacketType {
+    PacketType_Invalid,
+    PacketType_ErrorMessage,
+    PacketType_ClientHandshake,
+    PacketType_ServerHandshake,
+} __attribute__((__packed__)) PacketType;
+
+typedef struct ErrorMessage {
+    PacketHeader header;
+    /* content stream of size `header.content_size` */
+} ErrorMessage;
+
+typedef struct ClientHandshake {
+    PacketHeader header;
+    u8 session_key[AES_SESSION_KEY_SIZE];
+} ClientHandshake;
+
+Result(void) ClientHandshake_tryConstruct(ClientHandshake* ptr, Array(u8) session_key);
+
+
+typedef struct ServerHandshake {
+    PacketHeader header;
+    u64 session_id;
+} ServerHandshake;

src/server/server.c

@@ -30,7 +30,7 @@ Result(void) server_run(cstr server_endpoint_str){
 
 static void* handle_connection(void* _args){
     Deferral(64);
-    ConnectionHandlerArgs* args = (ConnectionHandlerArgs*)_args;
+    //ConnectionHandlerArgs* args = (ConnectionHandlerArgs*)_args;
     // TODO: receive handshake and session key
     
     //ClientConnection conn;

src/server/server.h

@@ -1,12 +1,14 @@
 #pragma once
-#include "cryptography/cryptography.h"
+#include "cryptography/AES.h"
-#include "network/EncryptedSocket.h"
+#include "cryptography/RSA.h"
+#include "network/encrypted_sockets.h"
 
 Result(void) server_run(cstr server_endpoint_str);
 
 typedef struct ClientConnection {
+    u64 session_id;
     EndpointIPv4 client_end;
     Array(u8) session_key;
-    EncryptedSocket system_socket;
+    EncryptedSocketTCP system_socket;
-    EncryptedSocket content_socket;
+    EncryptedSocketTCP content_socket;
 } ClientConnection;