aes128_mmo.cuh

Go to the documentation of this file.

// SPDX-License-Identifier: Apache-2.0
#pragma once
#include <fss/prg.cuh>
#include <cuda_runtime.h>
#include <cuda/std/array>
#include <cuda/std/span>
#include <cassert>
#include <openssl/evp.h>
#include <openssl/aes.h>
#include <fss/util.cuh>
 
namespace fss::prg {
 
template <int mul>
class Aes128Mmo {
private:
    cuda::std::array<EVP_CIPHER_CTX *, mul> ctxs_;
 
public:
    Aes128Mmo(cuda::std::span<EVP_CIPHER_CTX *, mul> ctxs) {
        for (int i = 0; i < mul; ++i) ctxs_[i] = ctxs[i];
    }
 
    static cuda::std::array<EVP_CIPHER_CTX *, mul> CreateCtxs(const unsigned char *keys[mul]) {
        int ret;
        cuda::std::array<EVP_CIPHER_CTX *, mul> ctxs;
 
        for (int i = 0; i < mul; ++i) {
            ctxs[i] = EVP_CIPHER_CTX_new();
            assert(ctxs[i] != NULL);
 
            ret = EVP_EncryptInit_ex2(ctxs[i], EVP_aes_128_ecb(), keys[i], NULL, NULL);
            assert(ret == 1);
 
            ret = EVP_CIPHER_CTX_set_padding(ctxs[i], 0);
            assert(ret == 1);
        }
        return ctxs;
    }
 
    static void FreeCtxs(cuda::std::span<EVP_CIPHER_CTX *, mul> ctxs) {
        for (auto ctx : ctxs) {
            EVP_CIPHER_CTX_free(ctx);
        }
    }
 
    __host__ __device__ cuda::std::array<int4, mul> Gen(int4 seed) {
        cuda::std::array<int4, mul> out{};
 
#ifdef __CUDA_ARCH__
        assert(false && "Aes128Mmo is not supported on device side");
        __trap();
#else
        for (int i = 0; i < mul; ++i) {
            auto out_ptr = reinterpret_cast<unsigned char *>(&out[i]);
            auto seed_ptr = reinterpret_cast<const unsigned char *>(&seed);
            int cipher_len = 0;
            // Ctx does not change after block encryption because we use ECB, no padding, and AES_BLOCK_SIZE input size.
            int ret = EVP_EncryptUpdate(ctxs_[i], out_ptr, &cipher_len, seed_ptr, AES_BLOCK_SIZE);
            assert(ret == 1);
            assert(cipher_len == AES_BLOCK_SIZE);
 
            out[i] = fss::util::Xor(out[i], seed);
        }
#endif
 
        return out;
    }
};
static_assert(Prgable<Aes128Mmo<2>, 2> && Prgable<Aes128Mmo<4>, 4>);
 
}  // namespace fss::prg