moonlight-common-c/src/PlatformCrypto.c

#include "Limelight-internal.h"

#ifdef USE_MBEDTLS
#include <mbedtls/entropy.h>
#include <mbedtls/ctr_drbg.h>

mbedtls_entropy_context EntropyContext;
mbedtls_ctr_drbg_context CtrDrbgContext;
bool RandomStateInitialized = false;
#else
#include <openssl/evp.h>
#include <openssl/rand.h>

#if OPENSSL_VERSION_NUMBER < 0x10100000L
#define EVP_CIPHER_CTX_reset(x) EVP_CIPHER_CTX_cleanup(x); EVP_CIPHER_CTX_init(x)
#endif

static int addPkcs7PaddingInPlace(unsigned char* plaintext, int plaintextLen) {
    int paddedLength = ROUND_TO_PKCS7_PADDED_LEN(plaintextLen);
    unsigned char paddingByte = (unsigned char)(16 - (plaintextLen % 16));

    memset(&plaintext[plaintextLen], paddingByte, paddedLength - plaintextLen);

    return paddedLength;
}
#endif

// For CBC modes, inputData buffer must be allocated with length rounded up to next multiple of 16 and inputData buffer may be modified!
bool PltEncryptMessage(PPLT_CRYPTO_CONTEXT ctx, int algorithm,
                       unsigned char* key, int keyLength,
                       unsigned char* iv, int ivLength,
                       unsigned char* tag, int tagLength,
                       unsigned char* inputData, int inputDataLength,
                       unsigned char* outputData, int* outputDataLength) {
#ifdef USE_MBEDTLS
    mbedtls_cipher_mode_t cipherMode;
    size_t outLength;

    switch (algorithm) {
    case ALGORITHM_AES_CBC:
        LC_ASSERT(tag == NULL);
        LC_ASSERT(tagLength == 0);
        cipherMode = MBEDTLS_MODE_CBC;
        break;
    case ALGORITHM_AES_GCM:
        LC_ASSERT(tag != NULL);
        LC_ASSERT(tagLength > 0);
        cipherMode = MBEDTLS_MODE_GCM;
        break;
    default:
        LC_ASSERT(false);
        return false;
    }

    if (!ctx->initialized) {
        if (mbedtls_cipher_setup(&ctx->ctx, mbedtls_cipher_info_from_values(MBEDTLS_CIPHER_ID_AES, keyLength * 8, cipherMode)) != 0) {
            return false;
        }

        if (mbedtls_cipher_setkey(&ctx->ctx, key, keyLength * 8, MBEDTLS_ENCRYPT) != 0) {
            return false;
        }

        ctx->initialized = true;
    }

    outLength = *outputDataLength;

    if (tag != NULL) {
        if (mbedtls_cipher_auth_encrypt(&ctx->ctx, iv, ivLength, NULL, 0, inputData, inputDataLength, outputData, &outLength, tag, tagLength) != 0) {
            return false;
        }
    }
    else {
        if (mbedtls_cipher_crypt(&ctx->ctx, iv, ivLength, inputData, inputDataLength, outputData, &outLength) != 0) {
            return false;
        }
    }

    *outputDataLength = outLength;
    return true;
#else
    bool ret = false;
    const EVP_CIPHER* cipher;

    switch (algorithm) {
    case ALGORITHM_AES_CBC:
        LC_ASSERT(keyLength == 16);
        LC_ASSERT(tag == NULL);
        LC_ASSERT(tagLength == 0);
        cipher = EVP_aes_128_cbc();
        break;
    case ALGORITHM_AES_GCM:
        LC_ASSERT(keyLength == 16);
        LC_ASSERT(tag != NULL);
        LC_ASSERT(tagLength > 0);
        cipher = EVP_aes_128_gcm();
        break;
    default:
        LC_ASSERT(false);
        return false;
    }

    if (algorithm == ALGORITHM_AES_GCM) {
        if (EVP_EncryptInit_ex(ctx->ctx, cipher, NULL, NULL, NULL) != 1) {
            goto cleanup;
        }

        if (EVP_CIPHER_CTX_ctrl(ctx->ctx, EVP_CTRL_GCM_SET_IVLEN, ivLength, NULL) != 1) {
            goto cleanup;
        }

        if (EVP_EncryptInit_ex(ctx->ctx, NULL, NULL, key, iv) != 1) {
            goto cleanup;
        }
    }
    else {
        if (!ctx->initialized) {
            if (EVP_EncryptInit_ex(ctx->ctx, cipher, NULL, key, iv) != 1) {
                goto cleanup;
            }

            ctx->initialized = true;
        }

        inputDataLength = addPkcs7PaddingInPlace(inputData, inputDataLength);
    }

    if (EVP_EncryptUpdate(ctx->ctx, outputData, outputDataLength, inputData, inputDataLength) != 1) {
        goto cleanup;
    }

    if (algorithm == ALGORITHM_AES_GCM) {
        int len;

        // GCM encryption won't ever fill ciphertext here but we have to call it anyway
        if (EVP_EncryptFinal_ex(ctx->ctx, outputData, &len) != 1) {
            goto cleanup;
        }
        LC_ASSERT(len == 0);

        if (EVP_CIPHER_CTX_ctrl(ctx->ctx, EVP_CTRL_GCM_GET_TAG, tagLength, tag) != 1) {
            goto cleanup;
        }
    }

    ret = true;

cleanup:
    if (algorithm == ALGORITHM_AES_GCM) {
        EVP_CIPHER_CTX_reset(ctx->ctx);
    }
    return ret;
#endif
}

bool PltDecryptMessage(PPLT_CRYPTO_CONTEXT ctx, int algorithm,
                       unsigned char* key, int keyLength,
                       unsigned char* iv, int ivLength,
                       unsigned char* tag, int tagLength,
                       unsigned char* inputData, int inputDataLength,
                       unsigned char* outputData, int* outputDataLength) {
#ifdef USE_MBEDTLS
    mbedtls_cipher_mode_t cipherMode;
    size_t outLength;

    switch (algorithm) {
    case ALGORITHM_AES_CBC:
        LC_ASSERT(tag == NULL);
        LC_ASSERT(tagLength == 0);
        cipherMode = MBEDTLS_MODE_CBC;
        break;
    case ALGORITHM_AES_GCM:
        LC_ASSERT(tag != NULL);
        LC_ASSERT(tagLength > 0);
        cipherMode = MBEDTLS_MODE_GCM;
        break;
    default:
        LC_ASSERT(false);
        return false;
    }

    if (!ctx->initialized) {
        if (mbedtls_cipher_setup(&ctx->ctx, mbedtls_cipher_info_from_values(MBEDTLS_CIPHER_ID_AES, keyLength * 8, cipherMode)) != 0) {
            return false;
        }

        if (mbedtls_cipher_setkey(&ctx->ctx, key, keyLength * 8, MBEDTLS_DECRYPT) != 0) {
            return false;
        }

        ctx->initialized = true;
    }

    outLength = *outputDataLength;

    if (tag != NULL) {
        if (mbedtls_cipher_auth_decrypt(&ctx->ctx, iv, ivLength, NULL, 0, inputData, inputDataLength, outputData, &outLength, tag, tagLength) != 0) {
            return false;
        }
    }
    else {
        if (mbedtls_cipher_crypt(&ctx->ctx, iv, ivLength, inputData, inputDataLength, outputData, &outLength) != 0) {
            return false;
        }
    }

    *outputDataLength = outLength;
    return true;
#else
    bool ret = false;
    const EVP_CIPHER* cipher;

    switch (algorithm) {
    case ALGORITHM_AES_CBC:
        LC_ASSERT(keyLength == 16);
        LC_ASSERT(tag == NULL);
        LC_ASSERT(tagLength == 0);
        cipher = EVP_aes_128_cbc();
        break;
    case ALGORITHM_AES_GCM:
        LC_ASSERT(keyLength == 16);
        LC_ASSERT(tag != NULL);
        LC_ASSERT(tagLength > 0);
        cipher = EVP_aes_128_gcm();
        break;
    default:
        LC_ASSERT(false);
        return false;
    }

    if (algorithm == ALGORITHM_AES_GCM) {
        if (EVP_DecryptInit_ex(ctx->ctx, cipher, NULL, NULL, NULL) != 1) {
            goto cleanup;
        }

        if (EVP_CIPHER_CTX_ctrl(ctx->ctx, EVP_CTRL_GCM_SET_IVLEN, ivLength, NULL) != 1) {
            goto cleanup;
        }

        if (EVP_DecryptInit_ex(ctx->ctx, NULL, NULL, key, iv) != 1) {
            goto cleanup;
        }
    }
    else {
        if (!ctx->initialized) {
            if (EVP_DecryptInit_ex(ctx->ctx, cipher, NULL, key, iv) != 1) {
                goto cleanup;
            }

            ctx->initialized = true;
        }
    }

    if (EVP_DecryptUpdate(ctx->ctx, outputData, outputDataLength, inputData, inputDataLength) != 1) {
        goto cleanup;
    }

    if (algorithm == ALGORITHM_AES_GCM) {
        int len;

        // Set the GCM tag before calling EVP_DecryptFinal_ex()
        if (EVP_CIPHER_CTX_ctrl(ctx->ctx, EVP_CTRL_GCM_SET_TAG, tagLength, tag) != 1) {
            goto cleanup;
        }

        // GCM will never have additional plaintext here, but we need to call it to
        // ensure that the GCM authentication tag is correct for this data.
        if (EVP_DecryptFinal_ex(ctx->ctx, outputData, &len) != 1) {
            goto cleanup;
        }
        LC_ASSERT(len == 0);
    }

    ret = true;

cleanup:
    if (algorithm == ALGORITHM_AES_GCM) {
        EVP_CIPHER_CTX_reset(ctx->ctx);
    }
    return ret;
#endif
}

PPLT_CRYPTO_CONTEXT PltCreateCryptoContext(void) {
    PPLT_CRYPTO_CONTEXT ctx = malloc(sizeof(*ctx));
    if (!ctx) {
        return NULL;
    }

    ctx->initialized = false;

#ifdef USE_MBEDTLS
    mbedtls_cipher_init(&ctx->ctx);
#else
    ctx->ctx = EVP_CIPHER_CTX_new();
    if (!ctx->ctx) {
        free(ctx);
        return NULL;
    }
#endif

    return ctx;
}

void PltDestroyCryptoContext(PPLT_CRYPTO_CONTEXT ctx) {
#ifdef USE_MBEDTLS
    mbedtls_cipher_free(&ctx->ctx);
#else
    EVP_CIPHER_CTX_free(ctx->ctx);
#endif
    free(ctx);
}

void PltGenerateRandomData(unsigned char* data, int length) {
#ifdef USE_MBEDTLS
    // FIXME: This is not thread safe...
    if (!RandomStateInitialized) {
        mbedtls_entropy_init(&EntropyContext);
        mbedtls_ctr_drbg_init(&CtrDrbgContext);
        if (mbedtls_ctr_drbg_seed(&CtrDrbgContext, mbedtls_entropy_func, &EntropyContext, NULL, 0) != 0) {
            // Nothing we can really do here...
            Limelog("Seeding MbedTLS random number generator failed!\n");
            LC_ASSERT(false);
            return;
        }

        RandomStateInitialized = true;
    }

    mbedtls_ctr_drbg_random(&CtrDrbgContext, data, length);
#else
    RAND_bytes(data, length);
#endif
}