kpabe.cpp

#include <string>
#include <map>
#include <cmath>
#include <algorithm>
#include <numeric>
#include <functional>
#include <array>
#include <vector>

#include <mbedtls/cipher.h>
#include <mbedtls/md.h>
#include <pbc.h>

#include "kpabe.hpp"

//Debugging
#include <iostream>

using namespace std;

// For the encrypt/decrypt methods.
static const size_t AES_BLOCK_SIZE = 16;
static const size_t AES_KEY_SIZE = 32;

pairing_s pairing;
bool isInit = false;

pairing_ptr getPairing()
{
    if (!isInit)
    {
        pairing_init_set_str(&pairing, TYPE_A_PARAMS.c_str());
        isInit = true;
    }
    return &pairing;
}

void hashElement(element_t e, uint8_t *hashBuf)
{
    const int elementSize = element_length_in_bytes(e);
    uint8_t *elementBytes = new uint8_t[elementSize + 1];
    element_to_bytes(elementBytes, e);

    //TODO: use mbedtls_sha256
    auto mdInfo = mbedtls_md_info_from_type(MBEDTLS_MD_SHA256);
    mbedtls_md(mdInfo, elementBytes, elementSize, hashBuf);

    delete[] elementBytes;
}

/**
 * Common interface to for symmetric encryption and decryption.
 *
 * Uses AES-256-CBC and zero-filled IV.
 */
void mbedtlsSymCrypt(const uint8_t *input, size_t ilen, uint8_t *key, uint8_t *output, size_t *olen, mbedtls_operation_t mode)
{
    const auto cipherInfo = mbedtls_cipher_info_from_type(MBEDTLS_CIPHER_AES_256_CBC);
    mbedtls_cipher_context_t ctx;
    mbedtls_cipher_setup(&ctx, cipherInfo);
    mbedtls_cipher_setkey(&ctx, key, cipherInfo->key_bitlen, mode);
    array<uint8_t, 16> iv;
    iv.fill(0);
    mbedtls_cipher_crypt(&ctx, iv.data(), cipherInfo->iv_size, input, ilen, output, olen);
}

void symEncrypt(const uint8_t *input, size_t ilen, uint8_t *key, uint8_t *output, size_t *olen)
{
    mbedtlsSymCrypt(input, ilen, key, output, olen, MBEDTLS_ENCRYPT);
}

void symDecrypt(const uint8_t *input, size_t ilen, uint8_t *key, uint8_t *output, size_t *olen)
{
    mbedtlsSymCrypt(input, ilen, key, output, olen, MBEDTLS_DECRYPT);
}

// Node

Node::Node(const Node &other)
{
    attr = other.attr;
    type = other.type;
    children = other.children;
}

Node::Node(Node &&other) : attr(move(other.attr)),
                           type(other.type),
                           children(move(other.children))
{
}

Node::Node(int attr)
{
    this->attr = attr;
}

Node::Node(Type type, const vector<Node> &children)
{
    this->children = children;
    this->type = type;
}

//Fixed Copy Assignment
Node &Node::operator=(const Node &other)
{
    attr = other.attr;
    type = other.type;
    children = other.children;
}

//Move Assignment
Node &Node::operator=(Node &&other)
{
    //assert(this != &other);
    attr = move(other.attr);
    type = move(other.type);
    children = move(other.children);
    return *this;
}

void Node::addChild(const Node &node)
{
    children.push_back(node);
}

vector<int> Node::getLeafs() const
{
    vector<int> attrs;

    if (children.empty())
    {
        // Handles non-leaf node with one child
        attrs.push_back(attr);
    }
    else
    {
        for (const Node &child : children)
        {
            if (child.children.empty())
            {
                attrs.push_back(child.attr);
            }
            else
            {
                auto childAttrs = child.getLeafs();
                attrs.reserve(attrs.size() + childAttrs.size());
                attrs.insert(attrs.end(), childAttrs.begin(), childAttrs.end());
            }
        }
    }
    return attrs;
}

unsigned int Node::getThreshold() const
{
    return type == Type::OR ? 1 : static_cast<unsigned int>(children.size());
}

unsigned int Node::getPolyDegree() const
{
    return getThreshold() - 1;
}

vector<element_s> Node::splitShares(element_s &rootSecret)
{
    // Generate the coefficients for the polynomial.
    auto threshold = getThreshold();
    vector<element_s> coeff(threshold);

    element_init_same_as(&coeff[0], &rootSecret);
    element_set(&coeff[0], &rootSecret);

    // Generate random coefficients, except for q(0), which is set to the rootSecret.
    for (int i = 1; i <= getPolyDegree(); ++i)
    {
        element_init_same_as(&coeff[i], &rootSecret);
        element_random(&coeff[i]);
    }

    // Calculate the shares for each child.
    vector<element_s> shares(children.size());

    element_t temp;
    element_init_Zr(temp, getPairing());

    // The scheme decription defines an ordering on the children in a node (index(x)).
    // Here, we implicitly use a left to right order.
    for (int x = 1; x <= children.size(); ++x)
    {
        auto share = &shares[x - 1];
        element_init_same_as(share, &rootSecret);
        element_set0(share);
        // share = coeff[0] + coeff[1] * x + ... + coeff[threshold - 1] * x ^ (threshold - 1)
        for (int power = 0; power < coeff.size(); ++power)
        {
            element_set_si(temp, pow(x, power)); //TODO: handle pow
            element_mul(temp, temp, &coeff[power]);
            element_add(share, share, temp);
        }
    }

    element_clear(temp);
    for (element_s &c : coeff)
    {
        element_clear(&c);
    }

    return shares;
} //splitShares

vector<element_s> Node::getSecretShares(element_s &rootSecret)
{
    vector<element_s> shares;
    if (children.empty())
    {
        shares.push_back(rootSecret);
    }
    else
    {
        auto childSplits = splitShares(rootSecret);
        auto childSplitsIter = childSplits.begin();
        for (Node &child : children)
        {
            auto childShares = child.getSecretShares(*childSplitsIter++);
            shares.reserve(shares.size() + childShares.size());
            shares.insert(shares.end(), childShares.begin(), childShares.end());
        }
    }

    return shares;
}

vector<element_s> Node::recoverCoefficients()
{
    auto threshold = getThreshold();
    vector<element_s> coeff(threshold);

    element_t iVal, jVal, temp;
    element_init_Zr(iVal, getPairing());
    element_init_Zr(jVal, getPairing());
    element_init_Zr(temp, getPairing());

    for (int i = 1; i <= threshold; ++i)
    {
        element_set_si(iVal, i);
        element_s &result = coeff[i - 1];
        element_init_Zr(&result, getPairing());
        element_set1(&result);
        for (int j = 1; j <= threshold; ++j)
        {
            if (i == j)
            {
                continue;
            }
            // result *= (0 - j) / (i - j)
            element_set_si(jVal, -j);
            element_add(temp, iVal, jVal);
            element_div(temp, jVal, temp);
            element_mul(&result, &result, temp);
        }
    }

    element_clear(iVal);
    element_clear(jVal);
    element_clear(temp);

    return coeff;
}

vector<pair<int, element_s>>
Node::satisfyingAttributes(const vector<int> &attributes,
                           element_s &currentCoeff)
{
    vector<pair<int, element_s>> sat;

    if (children.empty())
    {
        if (find(attributes.begin(), attributes.end(), attr) != attributes.end())
        {
            sat.push_back({attr, currentCoeff});
        }
    }
    else
    {
        auto recCoeffs = recoverCoefficients();

        if (type == Type::AND)
        {
            bool allSatisfied = true;
            vector<pair<int, element_s>> totalChildSat;
            for (int i = 0; i < children.size(); ++i)
            {
                element_mul(&recCoeffs[i], &recCoeffs[i], &currentCoeff);
                auto childSat = children[i].satisfyingAttributes(attributes, recCoeffs[i]);
                if (childSat.empty())
                {
                    allSatisfied = false;
                    break;
                }
                totalChildSat.reserve(totalChildSat.size() + childSat.size());
                totalChildSat.insert(totalChildSat.end(), childSat.begin(), childSat.end());
            }
            if (allSatisfied)
            {
                sat = totalChildSat;
            }
        }
        else
        {
            auto &recCoeff0 = recCoeffs[0];
            element_mul(&recCoeff0, &recCoeff0, &currentCoeff);
            for (auto &child : children)
            {
                // TODO: Optimization -
                // Should return the shortest non-empty childSat instead of the first one.
                auto childSat = child.satisfyingAttributes(attributes, recCoeff0);
                if (!childSat.empty())
                {
                    sat = childSat;
                    break;
                }
            }
        }
    }

    return sat;
}

const vector<Node> &Node::getChildren() const
{
    return children;
}

// DecryptionKey

DecryptionKey::DecryptionKey(const Node &policy) : accessPolicy(policy) {}

// Algorithm Setup

void setup(const vector<int> &attributes,
           PublicParams &publicParams,
           PrivateParams &privateParams)
{
    element_init_Zr(&privateParams.mk, getPairing());
    element_random(&privateParams.mk);

    element_t g;
    element_init_G1(g, getPairing());
    element_random(g);

    // Generate a random public and private element for each attribute
    for (auto attr : attributes)
    {
        // private
        element_s &si = privateParams.Si[attr];
        element_init_Zr(&si, getPairing());
        element_random(&si);

        // public
        element_s &Pi = publicParams.Pi[attr];
        element_init_G1(&Pi, getPairing());
        element_pow_zn(&Pi, g, &si);
    }

    element_init_G1(&publicParams.pk, getPairing());
    element_pow_zn(&publicParams.pk, g, &privateParams.mk);
    element_clear(g);
}

/**
 * @brief An abstraction of createKey that allows different operation for hiding the
 *    secret shares.
 *
 * @param scramblingFunc A function that sets an element to the result of a function on
 *    a scambling key and a secret share. In the original paper the scrambling keys are
 *    the private keys and the function is division. The result of the scrambling is put
 *    in the first element, the shares in the second, the scramblng keys in the third.
 * @type scramblingFunc function<void (element_t, element_t, element_t)>
 */
DecryptionKey _keyGeneration(element_s &rootSecret,
                             map<int, element_s> &scramblingKeys,
                             function<void(element_t, element_t, element_t)> scramblingFunc,
                             Node &accessPolicy)
{
    auto leafs = accessPolicy.getLeafs();
    auto shares = accessPolicy.getSecretShares(rootSecret);

    DecryptionKey key(accessPolicy);
    auto attrIter = leafs.begin();
    auto sharesIter = shares.begin();
    // The below is: Du[attr] = shares[attr] / attributeSecrets[attr]
    for (; attrIter != leafs.end(); ++attrIter, ++sharesIter)
    {
        element_s &attrDi = key.Di[*attrIter];
        element_init_Zr(&attrDi, getPairing());
        scramblingFunc(&attrDi, &*sharesIter, &scramblingKeys[*attrIter]);
    }

    for (element_s &share : shares)
    {
        element_clear(&share);
    }

    return key;
}

DecryptionKey keyGeneration(PrivateParams &privateParams,
                            Node &accessPolicy)
{
    return _keyGeneration(privateParams.mk, privateParams.Si, element_div, accessPolicy);
}

Cw_t createSecret(PublicParams &params,
                  const vector<int> &attributes,
                  element_s &Cs)
{
    element_t k;
    element_init_Zr(k, getPairing());
    element_random(k);

    element_init_G1(&Cs, getPairing());
    element_pow_zn(&Cs, &params.pk, k);

    Cw_t Cw;
    for (auto attr : attributes)
    {
        element_s &i = Cw[attr];
        element_init_G1(&i, getPairing());
        element_pow_zn(&i, &params.Pi[attr], k);
    }
    element_clear(k);

    return Cw;
}

void recoverSecret(DecryptionKey &key,
                   Cw_t &Cw,
                   const vector<int> &attributes,
                   element_s &Cs)
{
    // Get attributes that can satisfy the policy (and their coefficients).
    element_t rootCoeff;
    element_init_Zr(rootCoeff, getPairing());
    element_set1(rootCoeff);
    auto attrs = key.accessPolicy.satisfyingAttributes(attributes, *rootCoeff);
    element_clear(rootCoeff);

    if (attrs.empty())
    {
        throw UnsatError();
        return;
    }

    element_t Zy;
    element_init_G1(&Cs, getPairing());
    element_init_G1(Zy, getPairing());
    bool pastFirst = false; // Is this the first "part" of the product

    // product = P(Ci ^ (Di * coeff(i)))
    // NOTE: attrCoeffPair is modified
    for (auto &attrCoeffPair : attrs)
    {
        element_mul(&attrCoeffPair.second, &key.Di[attrCoeffPair.first], &attrCoeffPair.second);
        element_pow_zn(Zy, &Cw[attrCoeffPair.first], &attrCoeffPair.second);

        if (pastFirst)
        {
            element_mul(&Cs, &Cs, Zy);
        }
        else
        {
            pastFirst = true;
            element_set(&Cs, Zy);
        }
    }

    for (auto &attrCoeffPair : attrs)
    {
        element_clear(&attrCoeffPair.second);
    }
    element_clear(Zy);
}

std::vector<uint8_t> encrypt(PublicParams &params,
                             const vector<int> &attributes,
                             const string &message,
                             Cw_t &Cw)
{
    element_s Cs;
    Cw = createSecret(params, attributes, Cs);

    // Use the key to encrypt the data using a symmetric cipher.
    size_t messageLen = message.size() + 1; // account for terminating byte
    size_t cipherMaxLen = messageLen + AES_BLOCK_SIZE;
    vector<uint8_t> ciphertext(cipherMaxLen);

    array<uint8_t, AES_KEY_SIZE> key;
    hashElement(&Cs, key.data());
    size_t clength = 0;
    symEncrypt((uint8_t *)message.c_str(), messageLen, key.data(), ciphertext.data(), &clength);
    ciphertext.resize(clength);

    element_clear(&Cs);

    return ciphertext;
}

string decrypt(DecryptionKey &key,
               Cw_t &Cw,
               const vector<int> &attributes,
               const vector<uint8_t> &ciphertext)
{
    element_s Cs;
    recoverSecret(key, Cw, attributes, Cs);
    vector<uint8_t> plaintext(ciphertext.size());
    size_t plaintextLen = 0;

    array<uint8_t, AES_KEY_SIZE> symKey;
    hashElement(&Cs, symKey.data());
    symDecrypt(ciphertext.data(), ciphertext.size(), symKey.data(), plaintext.data(), &plaintextLen);
    plaintext.resize(plaintextLen);
    string message((char *)plaintext.data());

    element_clear(&Cs);

    return message;
}