diff --git a/lib/include/chiaki/common.h b/lib/include/chiaki/common.h index 29c21ff..44bd7b4 100644 --- a/lib/include/chiaki/common.h +++ b/lib/include/chiaki/common.h @@ -18,6 +18,8 @@ #ifndef CHIAKI_COMMON_H #define CHIAKI_COMMON_H +#include + #ifdef __cplusplus extern "C" { #endif @@ -49,6 +51,8 @@ typedef enum CHIAKI_EXPORT const char *chiaki_error_string(ChiakiErrorCode code); +void *chiaki_aligned_alloc(size_t alignment, size_t size); + #ifdef __cplusplus } #endif diff --git a/lib/include/chiaki/gkcrypt.h b/lib/include/chiaki/gkcrypt.h index b1443b9..a2c76e7 100644 --- a/lib/include/chiaki/gkcrypt.h +++ b/lib/include/chiaki/gkcrypt.h @@ -20,6 +20,7 @@ #include "common.h" #include "log.h" +#include "thread.h" #include #include @@ -29,13 +30,25 @@ extern "C" { #endif #define CHIAKI_GKCRYPT_BLOCK_SIZE 0x10 +#define CHIAKI_GKCRYPT_KEY_BUF_BLOCKS_DEFAULT 0x20 // 2MB #define CHIAKI_GKCRYPT_GMAC_SIZE 4 #define CHIAKI_GKCRYPT_GMAC_KEY_REFRESH_KEY_POS 45000 #define CHIAKI_GKCRYPT_GMAC_KEY_REFRESH_IV_OFFSET 44910 typedef struct chiaki_gkcrypt_t { - uint8_t *key_buf; + uint8_t index; + + uint8_t *key_buf; // circular buffer of the ctr mode key stream size_t key_buf_size; + size_t key_buf_populated; // size of key_buf that is already populated (on startup) + size_t key_buf_key_pos_min; // minimal key pos currently in key_buf + size_t key_buf_start_offset; // offset in key_buf of the minimal key pos + size_t last_key_pos; // last key pos that has been requested + bool key_buf_thread_stop; + ChiakiMutex key_buf_mutex; + ChiakiCond key_buf_cond; + ChiakiThread key_buf_thread; + uint8_t iv[CHIAKI_GKCRYPT_BLOCK_SIZE]; uint8_t key_base[CHIAKI_GKCRYPT_BLOCK_SIZE]; uint8_t key_gmac_base[CHIAKI_GKCRYPT_BLOCK_SIZE]; @@ -46,9 +59,14 @@ typedef struct chiaki_gkcrypt_t { struct chiaki_session_t; -CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_init(ChiakiGKCrypt *gkcrypt, ChiakiLog *log, size_t key_buf_blocks, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret); +/** + * @param key_buf_chunks if > 0, use a thread to generate the ctr mode key stream + */ +CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_init(ChiakiGKCrypt *gkcrypt, ChiakiLog *log, size_t key_buf_chunks, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret); + CHIAKI_EXPORT void chiaki_gkcrypt_fini(ChiakiGKCrypt *gkcrypt); CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_gen_key_stream(ChiakiGKCrypt *gkcrypt, size_t key_pos, uint8_t *buf, size_t buf_size); +CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_get_key_stream(ChiakiGKCrypt *gkcrypt, size_t key_pos, uint8_t *buf, size_t buf_size); CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_decrypt(ChiakiGKCrypt *gkcrypt, size_t key_pos, uint8_t *buf, size_t buf_size); static inline ChiakiErrorCode chiaki_gkcrypt_encrypt(ChiakiGKCrypt *gkcrypt, size_t key_pos, uint8_t *buf, size_t buf_size) { return chiaki_gkcrypt_decrypt(gkcrypt, key_pos, buf, buf_size); } CHIAKI_EXPORT void chiaki_gkcrypt_gen_gmac_key(uint64_t index, const uint8_t *key_base, const uint8_t *iv, uint8_t *key_out); @@ -56,12 +74,12 @@ CHIAKI_EXPORT void chiaki_gkcrypt_gen_new_gmac_key(ChiakiGKCrypt *gkcrypt, uint6 CHIAKI_EXPORT void chiaki_gkcrypt_gen_tmp_gmac_key(ChiakiGKCrypt *gkcrypt, uint64_t index, uint8_t *key_out); CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_gmac(ChiakiGKCrypt *gkcrypt, size_t key_pos, const uint8_t *buf, size_t buf_size, uint8_t *gmac_out); -static inline ChiakiGKCrypt *chiaki_gkcrypt_new(ChiakiLog *log, size_t key_buf_blocks, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret) +static inline ChiakiGKCrypt *chiaki_gkcrypt_new(ChiakiLog *log, size_t key_buf_chunks, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret) { ChiakiGKCrypt *gkcrypt = CHIAKI_NEW(ChiakiGKCrypt); if(!gkcrypt) return NULL; - ChiakiErrorCode err = chiaki_gkcrypt_init(gkcrypt, log, key_buf_blocks, index, handshake_key, ecdh_secret); + ChiakiErrorCode err = chiaki_gkcrypt_init(gkcrypt, log, key_buf_chunks, index, handshake_key, ecdh_secret); if(err != CHIAKI_ERR_SUCCESS) { free(gkcrypt); diff --git a/lib/src/common.c b/lib/src/common.c index 40d8c7f..b7dfe5e 100644 --- a/lib/src/common.c +++ b/lib/src/common.c @@ -17,6 +17,8 @@ #include +#include + CHIAKI_EXPORT const char *chiaki_error_string(ChiakiErrorCode code) { @@ -56,3 +58,8 @@ CHIAKI_EXPORT const char *chiaki_error_string(ChiakiErrorCode code) return "Unknown"; } } + +void *chiaki_aligned_alloc(size_t alignment, size_t size) +{ + return aligned_alloc(alignment, size); +} diff --git a/lib/src/gkcrypt.c b/lib/src/gkcrypt.c index 370c333..c43c47c 100644 --- a/lib/src/gkcrypt.c +++ b/lib/src/gkcrypt.c @@ -28,35 +28,91 @@ #include "utils.h" +#define KEY_BUF_CHUNK_SIZE 0x1000 + + static ChiakiErrorCode gkcrypt_gen_key_iv(ChiakiGKCrypt *gkcrypt, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret); +static void *gkcrypt_thread_func(void *user); -CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_init(ChiakiGKCrypt *gkcrypt, ChiakiLog *log, size_t key_buf_blocks, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret) +CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_init(ChiakiGKCrypt *gkcrypt, ChiakiLog *log, size_t key_buf_chunks, uint8_t index, const uint8_t *handshake_key, const uint8_t *ecdh_secret) { gkcrypt->log = log; - gkcrypt->key_buf_size = key_buf_blocks * CHIAKI_GKCRYPT_BLOCK_SIZE; - gkcrypt->key_buf = malloc(gkcrypt->key_buf_size); - if(!gkcrypt->key_buf) - return CHIAKI_ERR_MEMORY; + gkcrypt->index = index; - ChiakiErrorCode err = gkcrypt_gen_key_iv(gkcrypt, index, handshake_key, ecdh_secret); + gkcrypt->key_buf_size = key_buf_chunks * KEY_BUF_CHUNK_SIZE; + gkcrypt->key_buf_populated = 0; + gkcrypt->key_buf_key_pos_min = 0; + gkcrypt->key_buf_start_offset = 0; + gkcrypt->last_key_pos = 0; + gkcrypt->key_buf_thread_stop = false; + + ChiakiErrorCode err; + if(gkcrypt->key_buf_size) + { + gkcrypt->key_buf = chiaki_aligned_alloc(KEY_BUF_CHUNK_SIZE, gkcrypt->key_buf_size); + if(!gkcrypt->key_buf) + { + err = CHIAKI_ERR_MEMORY; + goto error; + } + + err = chiaki_mutex_init(&gkcrypt->key_buf_mutex, false); + if(err != CHIAKI_ERR_SUCCESS) + goto error_key_buf; + + err = chiaki_cond_init(&gkcrypt->key_buf_cond); + if(err != CHIAKI_ERR_SUCCESS) + goto error_key_buf_mutex; + } + else + { + gkcrypt->key_buf = NULL; + } + + err = gkcrypt_gen_key_iv(gkcrypt, index, handshake_key, ecdh_secret); if(err != CHIAKI_ERR_SUCCESS) { CHIAKI_LOGE(gkcrypt->log, "GKCrypt failed to generate key and IV"); - free(gkcrypt->key_buf); - return CHIAKI_ERR_UNKNOWN; + goto error_key_buf_cond; } chiaki_gkcrypt_gen_gmac_key(0, gkcrypt->key_base, gkcrypt->iv, gkcrypt->key_gmac_base); gkcrypt->key_gmac_index_current = 0; memcpy(gkcrypt->key_gmac_current, gkcrypt->key_gmac_base, sizeof(gkcrypt->key_gmac_current)); + if(gkcrypt->key_buf) + { + err = chiaki_thread_create(&gkcrypt->key_buf_thread, gkcrypt_thread_func, gkcrypt); + if(err != CHIAKI_ERR_SUCCESS) + goto error_key_buf_cond; + } + return CHIAKI_ERR_SUCCESS; + +error_key_buf_cond: + if(gkcrypt->key_buf) + chiaki_cond_fini(&gkcrypt->key_buf_cond); +error_key_buf_mutex: + if(gkcrypt->key_buf) + chiaki_mutex_fini(&gkcrypt->key_buf_mutex); +error_key_buf: + free(gkcrypt->key_buf); +error: + return err; } CHIAKI_EXPORT void chiaki_gkcrypt_fini(ChiakiGKCrypt *gkcrypt) { - free(gkcrypt->key_buf); + if(gkcrypt->key_buf) + { + chiaki_mutex_lock(&gkcrypt->key_buf_mutex); + gkcrypt->key_buf_thread_stop = true; + chiaki_mutex_unlock(&gkcrypt->key_buf_mutex); + chiaki_cond_signal(&gkcrypt->key_buf_cond); + chiaki_thread_join(&gkcrypt->key_buf_thread, NULL); + free(gkcrypt->key_buf); + } } @@ -162,6 +218,53 @@ CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_gen_key_stream(ChiakiGKCrypt *gkcry return CHIAKI_ERR_SUCCESS; } +static bool gkcrypt_key_buf_should_generate(ChiakiGKCrypt *gkcrypt) +{ + return gkcrypt->last_key_pos > gkcrypt->key_buf_key_pos_min + gkcrypt->key_buf_populated / 2; +} + +CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_get_key_stream(ChiakiGKCrypt *gkcrypt, size_t key_pos, uint8_t *buf, size_t buf_size) +{ + if(!gkcrypt->key_buf) + return chiaki_gkcrypt_gen_key_stream(gkcrypt, key_pos, buf, buf_size); + + chiaki_mutex_lock(&gkcrypt->key_buf_mutex); + + if(key_pos + buf_size > gkcrypt->last_key_pos) + gkcrypt->last_key_pos = key_pos + buf_size; + bool signal = gkcrypt_key_buf_should_generate(gkcrypt); + + ChiakiErrorCode err; + if(key_pos < gkcrypt->key_buf_key_pos_min + || key_pos + buf_size >= gkcrypt->key_buf_key_pos_min + gkcrypt->key_buf_populated) + { + CHIAKI_LOGW(gkcrypt->log, "Requested key stream for key pos %#llx on GKCrypt %d, but it's not in the buffer", (int)key_pos, gkcrypt->index); + chiaki_mutex_unlock(&gkcrypt->key_buf_mutex); + err = chiaki_gkcrypt_gen_key_stream(gkcrypt, key_pos, buf, buf_size); + } + else + { + size_t offset_in_buf = key_pos - gkcrypt->key_buf_key_pos_min + gkcrypt->key_buf_start_offset; + offset_in_buf %= gkcrypt->key_buf_size; + size_t end = offset_in_buf + buf_size; + if(end > gkcrypt->key_buf_size) + { + size_t excess = end - gkcrypt->key_buf_size; + memcpy(buf, gkcrypt->key_buf + offset_in_buf, buf_size - excess); + memcpy(buf + (buf_size - excess), gkcrypt->key_buf, excess); + } + else + memcpy(buf, gkcrypt->key_buf + offset_in_buf, buf_size); + err = CHIAKI_ERR_SUCCESS; + chiaki_mutex_unlock(&gkcrypt->key_buf_mutex); + } + + if(signal) + chiaki_cond_signal(&gkcrypt->key_buf_cond); + + return err; +} + CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_decrypt(ChiakiGKCrypt *gkcrypt, size_t key_pos, uint8_t *buf, size_t buf_size) { size_t padding_pre = key_pos % CHIAKI_GKCRYPT_BLOCK_SIZE; @@ -171,7 +274,7 @@ CHIAKI_EXPORT ChiakiErrorCode chiaki_gkcrypt_decrypt(ChiakiGKCrypt *gkcrypt, siz if(!key_stream) return CHIAKI_ERR_MEMORY; - ChiakiErrorCode err = chiaki_gkcrypt_gen_key_stream(gkcrypt, key_pos - padding_pre, key_stream, full_size); + ChiakiErrorCode err = chiaki_gkcrypt_get_key_stream(gkcrypt, key_pos - padding_pre, key_stream, full_size); if(err != CHIAKI_ERR_SUCCESS) { free(key_stream); @@ -253,4 +356,90 @@ fail_cipher: EVP_CIPHER_CTX_free(ctx); fail: return ret; -} \ No newline at end of file +} + +static bool key_buf_mutex_pred(void *user) +{ + ChiakiGKCrypt *gkcrypt = user; + if(gkcrypt->key_buf_thread_stop) + return true; + + if(gkcrypt->key_buf_populated < gkcrypt->key_buf_size) + return true; + + if(gkcrypt_key_buf_should_generate(gkcrypt)) + return true; + + return false; +} + +static ChiakiErrorCode gkcrypt_generate_next_chunk(ChiakiGKCrypt *gkcrypt) +{ + assert(gkcrypt->key_buf_populated + KEY_BUF_CHUNK_SIZE <= gkcrypt->key_buf_size); + size_t buf_offset = (gkcrypt->key_buf_start_offset + gkcrypt->key_buf_populated) % gkcrypt->key_buf_size; + size_t key_pos = gkcrypt->key_buf_key_pos_min + gkcrypt->key_buf_populated; + uint8_t *buf_start = gkcrypt->key_buf + buf_offset; + + chiaki_mutex_unlock(&gkcrypt->key_buf_mutex); + + ChiakiErrorCode err = chiaki_gkcrypt_gen_key_stream(gkcrypt, key_pos, buf_start, KEY_BUF_CHUNK_SIZE); + if(err != CHIAKI_ERR_SUCCESS) + CHIAKI_LOGE(gkcrypt->log, "GKCrypt failed to generate key stream chunk"); + + chiaki_mutex_lock(&gkcrypt->key_buf_mutex); + + if(err == CHIAKI_ERR_SUCCESS) + gkcrypt->key_buf_populated += KEY_BUF_CHUNK_SIZE; + + return err; +} + +static void *gkcrypt_thread_func(void *user) +{ + ChiakiGKCrypt *gkcrypt = user; + + CHIAKI_LOGV(gkcrypt->log, "GKCrypt %d thread starting", (int)gkcrypt->index); + + ChiakiErrorCode err = chiaki_mutex_lock(&gkcrypt->key_buf_mutex); + assert(err == CHIAKI_ERR_SUCCESS); + + while(1) + { + err = chiaki_cond_wait_pred(&gkcrypt->key_buf_cond, &gkcrypt->key_buf_mutex, key_buf_mutex_pred, gkcrypt); + if(gkcrypt->key_buf_thread_stop || err != CHIAKI_ERR_SUCCESS) + break; + + CHIAKI_LOGV(gkcrypt->log, "GKCrypt %d key buf size %#llx, start offset: %#llx, populated: %#llx, min key pos: %#llx, last key pos: %#llx, generating next chunk", + (int)gkcrypt->index, + (unsigned long long)gkcrypt->key_buf_size, + (unsigned long long)gkcrypt->key_buf_start_offset, + (unsigned long long)gkcrypt->key_buf_populated, + (unsigned long long)gkcrypt->key_buf_key_pos_min, + (unsigned long long)gkcrypt->last_key_pos); + + if(gkcrypt->last_key_pos > gkcrypt->key_buf_key_pos_min + gkcrypt->key_buf_populated) + { + // skip ahead if the last key pos is already beyond our buffer + size_t key_pos = (gkcrypt->last_key_pos / KEY_BUF_CHUNK_SIZE) * KEY_BUF_CHUNK_SIZE; + CHIAKI_LOGW(gkcrypt->log, "Already requested a higher key pos than in the buffer, skipping ahead from min %#llx to %#llx", + (unsigned long long)gkcrypt->key_buf_key_pos_min, + (unsigned long long)key_pos); + gkcrypt->key_buf_key_pos_min = key_pos; + gkcrypt->key_buf_start_offset = 0; + gkcrypt->key_buf_populated = 0; + } + else if(gkcrypt->key_buf_populated == gkcrypt->key_buf_size) + { + gkcrypt->key_buf_start_offset = (gkcrypt->key_buf_start_offset + KEY_BUF_CHUNK_SIZE) % gkcrypt->key_buf_size; + gkcrypt->key_buf_key_pos_min += KEY_BUF_CHUNK_SIZE; + gkcrypt->key_buf_populated -= KEY_BUF_CHUNK_SIZE; + } + + err = gkcrypt_generate_next_chunk(gkcrypt); + if(err != CHIAKI_ERR_SUCCESS) + break; + } + + chiaki_mutex_unlock(&gkcrypt->key_buf_mutex); + return NULL; +} diff --git a/lib/src/streamconnection.c b/lib/src/streamconnection.c index 72d8cd6..8007f42 100644 --- a/lib/src/streamconnection.c +++ b/lib/src/streamconnection.c @@ -396,13 +396,13 @@ static ChiakiErrorCode stream_connection_init_crypt(ChiakiStreamConnection *stre { ChiakiSession *session = stream_connection->session; - stream_connection->gkcrypt_local = chiaki_gkcrypt_new(stream_connection->log, 0 /* TODO */, 2, session->handshake_key, stream_connection->ecdh_secret); + stream_connection->gkcrypt_local = chiaki_gkcrypt_new(stream_connection->log, CHIAKI_GKCRYPT_KEY_BUF_BLOCKS_DEFAULT, 2, session->handshake_key, stream_connection->ecdh_secret); if(!stream_connection->gkcrypt_local) { CHIAKI_LOGE(stream_connection->log, "StreamConnection failed to initialize local GKCrypt with index 2"); return CHIAKI_ERR_UNKNOWN; } - stream_connection->gkcrypt_remote = chiaki_gkcrypt_new(stream_connection->log, 0 /* TODO */, 3, session->handshake_key, stream_connection->ecdh_secret); + stream_connection->gkcrypt_remote = chiaki_gkcrypt_new(stream_connection->log, CHIAKI_GKCRYPT_KEY_BUF_BLOCKS_DEFAULT, 3, session->handshake_key, stream_connection->ecdh_secret); if(!stream_connection->gkcrypt_remote) { CHIAKI_LOGE(stream_connection->log, "StreamConnection failed to initialize remote GKCrypt with index 3");