/* ====================================================================
 * Copyright (c) 2008 The OpenSSL Project. All rights reserved.
 *
 * Rights for redistribution and usage in source and binary
 * forms are granted according to the OpenSSL license.
 */

#include <openssl/crypto.h>
#include "modes_lcl.h"
#include <string.h>

#ifndef MODES_DEBUG
# ifndef NDEBUG
#  define NDEBUG
# endif
#endif
#include <assert.h>

/*
 * Trouble with Ciphertext Stealing, CTS, mode is that there is no
 * common official specification, but couple of cipher/application
 * specific ones: RFC2040 and RFC3962. Then there is 'Proposal to
 * Extend CBC Mode By "Ciphertext Stealing"' at NIST site, which
 * deviates from mentioned RFCs. Most notably it allows input to be
 * of block length and it doesn't flip the order of the last two
 * blocks. CTS is being discussed even in ECB context, but it's not
 * adopted for any known application. This implementation provides
 * two interfaces: one compliant with above mentioned RFCs and one
 * compliant with the NIST proposal, both extending CBC mode.
 */

size_t CRYPTO_cts128_encrypt_block(const unsigned char *in,
                                   unsigned char *out, size_t len,
                                   const void *key, unsigned char ivec[16],
                                   block128_f block)
{
    size_t residue, n;

    assert(in && out && key && ivec);

    if (len <= 16)
        return 0;

    if ((residue = len % 16) == 0)
        residue = 16;

    len -= residue;

    CRYPTO_cbc128_encrypt(in, out, len, key, ivec, block);

    in += len;
    out += len;

    for (n = 0; n < residue; ++n)
        ivec[n] ^= in[n];
    (*block) (ivec, ivec, key);
    memcpy(out, out - 16, residue);
    memcpy(out - 16, ivec, 16);

    return len + residue;
}

size_t CRYPTO_nistcts128_encrypt_block(const unsigned char *in,
                                       unsigned char *out, size_t len,
                                       const void *key,
                                       unsigned char ivec[16],
                                       block128_f block)
{
    size_t residue, n;

    assert(in && out && key && ivec);

    if (len < 16)
        return 0;

    residue = len % 16;

    len -= residue;

    CRYPTO_cbc128_encrypt(in, out, len, key, ivec, block);

    if (residue == 0)
        return len;

    in += len;
    out += len;

    for (n = 0; n < residue; ++n)
        ivec[n] ^= in[n];
    (*block) (ivec, ivec, key);
    memcpy(out - 16 + residue, ivec, 16);

    return len + residue;
}

size_t CRYPTO_cts128_encrypt(const unsigned char *in, unsigned char *out,
                             size_t len, const void *key,
                             unsigned char ivec[16], cbc128_f cbc)
{
    size_t residue;
    union {
        size_t align;
        unsigned char c[16];
    } tmp;

    assert(in && out && key && ivec);

    if (len <= 16)
        return 0;

    if ((residue = len % 16) == 0)
        residue = 16;

    len -= residue;

    (*cbc) (in, out, len, key, ivec, 1);

    in += len;
    out += len;

#if defined(CBC_HANDLES_TRUNCATED_IO)
    memcpy(tmp.c, out - 16, 16);
    (*cbc) (in, out - 16, residue, key, ivec, 1);
    memcpy(out, tmp.c, residue);
#else
    memset(tmp.c, 0, sizeof(tmp));
    memcpy(tmp.c, in, residue);
    memcpy(out, out - 16, residue);
    (*cbc) (tmp.c, out - 16, 16, key, ivec, 1);
#endif
    return len + residue;
}

size_t CRYPTO_nistcts128_encrypt(const unsigned char *in, unsigned char *out,
                                 size_t len, const void *key,
                                 unsigned char ivec[16], cbc128_f cbc)
{
    size_t residue;
    union {
        size_t align;
        unsigned char c[16];
    } tmp;

    assert(in && out && key && ivec);

    if (len < 16)
        return 0;

    residue = len % 16;

    len -= residue;

    (*cbc) (in, out, len, key, ivec, 1);

    if (residue == 0)
        return len;

    in += len;
    out += len;

#if defined(CBC_HANDLES_TRUNCATED_IO)
    (*cbc) (in, out - 16 + residue, residue, key, ivec, 1);
#else
    memset(tmp.c, 0, sizeof(tmp));
    memcpy(tmp.c, in, residue);
    (*cbc) (tmp.c, out - 16 + residue, 16, key, ivec, 1);
#endif
    return len + residue;
}

size_t CRYPTO_cts128_decrypt_block(const unsigned char *in,
                                   unsigned char *out, size_t len,
                                   const void *key, unsigned char ivec[16],
                                   block128_f block)
{
    size_t residue, n;
    union {
        size_t align;
        unsigned char c[32];
    } tmp;

    assert(in && out && key && ivec);

    if (len <= 16)
        return 0;

    if ((residue = len % 16) == 0)
        residue = 16;

    len -= 16 + residue;

    if (len) {
        CRYPTO_cbc128_decrypt(in, out, len, key, ivec, block);
        in += len;
        out += len;
    }

    (*block) (in, tmp.c + 16, key);

    memcpy(tmp.c, tmp.c + 16, 16);
    memcpy(tmp.c, in + 16, residue);
    (*block) (tmp.c, tmp.c, key);

    for (n = 0; n < 16; ++n) {
        unsigned char c = in[n];
        out[n] = tmp.c[n] ^ ivec[n];
        ivec[n] = c;
    }
    for (residue += 16; n < residue; ++n)
        out[n] = tmp.c[n] ^ in[n];

    return 16 + len + residue;
}

size_t CRYPTO_nistcts128_decrypt_block(const unsigned char *in,
                                       unsigned char *out, size_t len,
                                       const void *key,
                                       unsigned char ivec[16],
                                       block128_f block)
{
    size_t residue, n;
    union {
        size_t align;
        unsigned char c[32];
    } tmp;

    assert(in && out && key && ivec);

    if (len < 16)
        return 0;

    residue = len % 16;

    if (residue == 0) {
        CRYPTO_cbc128_decrypt(in, out, len, key, ivec, block);
        return len;
    }

    len -= 16 + residue;

    if (len) {
        CRYPTO_cbc128_decrypt(in, out, len, key, ivec, block);
        in += len;
        out += len;
    }

    (*block) (in + residue, tmp.c + 16, key);

    memcpy(tmp.c, tmp.c + 16, 16);
    memcpy(tmp.c, in, residue);
    (*block) (tmp.c, tmp.c, key);

    for (n = 0; n < 16; ++n) {
        unsigned char c = in[n];
        out[n] = tmp.c[n] ^ ivec[n];
        ivec[n] = in[n + residue];
        tmp.c[n] = c;
    }
    for (residue += 16; n < residue; ++n)
        out[n] = tmp.c[n] ^ tmp.c[n - 16];

    return 16 + len + residue;
}

size_t CRYPTO_cts128_decrypt(const unsigned char *in, unsigned char *out,
                             size_t len, const void *key,
                             unsigned char ivec[16], cbc128_f cbc)
{
    size_t residue;
    union {
        size_t align;
        unsigned char c[32];
    } tmp;

    assert(in && out && key && ivec);

    if (len <= 16)
        return 0;

    if ((residue = len % 16) == 0)
        residue = 16;

    len -= 16 + residue;

    if (len) {
        (*cbc) (in, out, len, key, ivec, 0);
        in += len;
        out += len;
    }

    memset(tmp.c, 0, sizeof(tmp));
    /*
     * this places in[16] at &tmp.c[16] and decrypted block at &tmp.c[0]
     */
    (*cbc) (in, tmp.c, 16, key, tmp.c + 16, 0);

    memcpy(tmp.c, in + 16, residue);
#if defined(CBC_HANDLES_TRUNCATED_IO)
    (*cbc) (tmp.c, out, 16 + residue, key, ivec, 0);
#else
    (*cbc) (tmp.c, tmp.c, 32, key, ivec, 0);
    memcpy(out, tmp.c, 16 + residue);
#endif
    return 16 + len + residue;
}

size_t CRYPTO_nistcts128_decrypt(const unsigned char *in, unsigned char *out,
                                 size_t len, const void *key,
                                 unsigned char ivec[16], cbc128_f cbc)
{
    size_t residue;
    union {
        size_t align;
        unsigned char c[32];
    } tmp;

    assert(in && out && key && ivec);

    if (len < 16)
        return 0;

    residue = len % 16;

    if (residue == 0) {
        (*cbc) (in, out, len, key, ivec, 0);
        return len;
    }

    len -= 16 + residue;

    if (len) {
        (*cbc) (in, out, len, key, ivec, 0);
        in += len;
        out += len;
    }

    memset(tmp.c, 0, sizeof(tmp));
    /*
     * this places in[16] at &tmp.c[16] and decrypted block at &tmp.c[0]
     */
    (*cbc) (in + residue, tmp.c, 16, key, tmp.c + 16, 0);

    memcpy(tmp.c, in, residue);
#if defined(CBC_HANDLES_TRUNCATED_IO)
    (*cbc) (tmp.c, out, 16 + residue, key, ivec, 0);
#else
    (*cbc) (tmp.c, tmp.c, 32, key, ivec, 0);
    memcpy(out, tmp.c, 16 + residue);
#endif
    return 16 + len + residue;
}

#if defined(SELFTEST)
# include <stdio.h>
# include <openssl/aes.h>

/* test vectors from RFC 3962 */
static const unsigned char test_key[16] = "chicken teriyaki";
static const unsigned char test_input[64] =
    "I would like the" " General Gau's C"
    "hicken, please, " "and wonton soup.";
static const unsigned char test_iv[16] =
    { 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0 };

static const unsigned char vector_17[17] = {
    0xc6, 0x35, 0x35, 0x68, 0xf2, 0xbf, 0x8c, 0xb4,
    0xd8, 0xa5, 0x80, 0x36, 0x2d, 0xa7, 0xff, 0x7f,
    0x97
};

static const unsigned char vector_31[31] = {
    0xfc, 0x00, 0x78, 0x3e, 0x0e, 0xfd, 0xb2, 0xc1,
    0xd4, 0x45, 0xd4, 0xc8, 0xef, 0xf7, 0xed, 0x22,
    0x97, 0x68, 0x72, 0x68, 0xd6, 0xec, 0xcc, 0xc0,
    0xc0, 0x7b, 0x25, 0xe2, 0x5e, 0xcf, 0xe5
};

static const unsigned char vector_32[32] = {
    0x39, 0x31, 0x25, 0x23, 0xa7, 0x86, 0x62, 0xd5,
    0xbe, 0x7f, 0xcb, 0xcc, 0x98, 0xeb, 0xf5, 0xa8,
    0x97, 0x68, 0x72, 0x68, 0xd6, 0xec, 0xcc, 0xc0,
    0xc0, 0x7b, 0x25, 0xe2, 0x5e, 0xcf, 0xe5, 0x84
};

static const unsigned char vector_47[47] = {
    0x97, 0x68, 0x72, 0x68, 0xd6, 0xec, 0xcc, 0xc0,
    0xc0, 0x7b, 0x25, 0xe2, 0x5e, 0xcf, 0xe5, 0x84,
    0xb3, 0xff, 0xfd, 0x94, 0x0c, 0x16, 0xa1, 0x8c,
    0x1b, 0x55, 0x49, 0xd2, 0xf8, 0x38, 0x02, 0x9e,
    0x39, 0x31, 0x25, 0x23, 0xa7, 0x86, 0x62, 0xd5,
    0xbe, 0x7f, 0xcb, 0xcc, 0x98, 0xeb, 0xf5
};

static const unsigned char vector_48[48] = {
    0x97, 0x68, 0x72, 0x68, 0xd6, 0xec, 0xcc, 0xc0,
    0xc0, 0x7b, 0x25, 0xe2, 0x5e, 0xcf, 0xe5, 0x84,
    0x9d, 0xad, 0x8b, 0xbb, 0x96, 0xc4, 0xcd, 0xc0,
    0x3b, 0xc1, 0x03, 0xe1, 0xa1, 0x94, 0xbb, 0xd8,
    0x39, 0x31, 0x25, 0x23, 0xa7, 0x86, 0x62, 0xd5,
    0xbe, 0x7f, 0xcb, 0xcc, 0x98, 0xeb, 0xf5, 0xa8
};

static const unsigned char vector_64[64] = {
    0x97, 0x68, 0x72, 0x68, 0xd6, 0xec, 0xcc, 0xc0,
    0xc0, 0x7b, 0x25, 0xe2, 0x5e, 0xcf, 0xe5, 0x84,
    0x39, 0x31, 0x25, 0x23, 0xa7, 0x86, 0x62, 0xd5,
    0xbe, 0x7f, 0xcb, 0xcc, 0x98, 0xeb, 0xf5, 0xa8,
    0x48, 0x07, 0xef, 0xe8, 0x36, 0xee, 0x89, 0xa5,
    0x26, 0x73, 0x0d, 0xbc, 0x2f, 0x7b, 0xc8, 0x40,
    0x9d, 0xad, 0x8b, 0xbb, 0x96, 0xc4, 0xcd, 0xc0,
    0x3b, 0xc1, 0x03, 0xe1, 0xa1, 0x94, 0xbb, 0xd8
};

static AES_KEY encks, decks;

void test_vector(const unsigned char *vector, size_t len)
{
    unsigned char iv[sizeof(test_iv)];
    unsigned char cleartext[64], ciphertext[64];
    size_t tail;

    printf("vector_%d\n", len);
    fflush(stdout);

    if ((tail = len % 16) == 0)
        tail = 16;
    tail += 16;

    /* test block-based encryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_cts128_encrypt_block(test_input, ciphertext, len, &encks, iv,
                                (block128_f) AES_encrypt);
    if (memcmp(ciphertext, vector, len))
        fprintf(stderr, "output_%d mismatch\n", len), exit(1);
    if (memcmp(iv, vector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(1);

    /* test block-based decryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_cts128_decrypt_block(ciphertext, cleartext, len, &decks, iv,
                                (block128_f) AES_decrypt);
    if (memcmp(cleartext, test_input, len))
        fprintf(stderr, "input_%d mismatch\n", len), exit(2);
    if (memcmp(iv, vector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(2);

    /* test streamed encryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_cts128_encrypt(test_input, ciphertext, len, &encks, iv,
                          (cbc128_f) AES_cbc_encrypt);
    if (memcmp(ciphertext, vector, len))
        fprintf(stderr, "output_%d mismatch\n", len), exit(3);
    if (memcmp(iv, vector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(3);

    /* test streamed decryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_cts128_decrypt(ciphertext, cleartext, len, &decks, iv,
                          (cbc128_f) AES_cbc_encrypt);
    if (memcmp(cleartext, test_input, len))
        fprintf(stderr, "input_%d mismatch\n", len), exit(4);
    if (memcmp(iv, vector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(4);
}

void test_nistvector(const unsigned char *vector, size_t len)
{
    unsigned char iv[sizeof(test_iv)];
    unsigned char cleartext[64], ciphertext[64], nistvector[64];
    size_t tail;

    printf("nistvector_%d\n", len);
    fflush(stdout);

    if ((tail = len % 16) == 0)
        tail = 16;

    len -= 16 + tail;
    memcpy(nistvector, vector, len);
    /* flip two last blocks */
    memcpy(nistvector + len, vector + len + 16, tail);
    memcpy(nistvector + len + tail, vector + len, 16);
    len += 16 + tail;
    tail = 16;

    /* test block-based encryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_nistcts128_encrypt_block(test_input, ciphertext, len, &encks, iv,
                                    (block128_f) AES_encrypt);
    if (memcmp(ciphertext, nistvector, len))
        fprintf(stderr, "output_%d mismatch\n", len), exit(1);
    if (memcmp(iv, nistvector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(1);

    /* test block-based decryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_nistcts128_decrypt_block(ciphertext, cleartext, len, &decks, iv,
                                    (block128_f) AES_decrypt);
    if (memcmp(cleartext, test_input, len))
        fprintf(stderr, "input_%d mismatch\n", len), exit(2);
    if (memcmp(iv, nistvector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(2);

    /* test streamed encryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_nistcts128_encrypt(test_input, ciphertext, len, &encks, iv,
                              (cbc128_f) AES_cbc_encrypt);
    if (memcmp(ciphertext, nistvector, len))
        fprintf(stderr, "output_%d mismatch\n", len), exit(3);
    if (memcmp(iv, nistvector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(3);

    /* test streamed decryption */
    memcpy(iv, test_iv, sizeof(test_iv));
    CRYPTO_nistcts128_decrypt(ciphertext, cleartext, len, &decks, iv,
                              (cbc128_f) AES_cbc_encrypt);
    if (memcmp(cleartext, test_input, len))
        fprintf(stderr, "input_%d mismatch\n", len), exit(4);
    if (memcmp(iv, nistvector + len - tail, sizeof(iv)))
        fprintf(stderr, "iv_%d mismatch\n", len), exit(4);
}

int main()
{
    AES_set_encrypt_key(test_key, 128, &encks);
    AES_set_decrypt_key(test_key, 128, &decks);

    test_vector(vector_17, sizeof(vector_17));
    test_vector(vector_31, sizeof(vector_31));
    test_vector(vector_32, sizeof(vector_32));
    test_vector(vector_47, sizeof(vector_47));
    test_vector(vector_48, sizeof(vector_48));
    test_vector(vector_64, sizeof(vector_64));

    test_nistvector(vector_17, sizeof(vector_17));
    test_nistvector(vector_31, sizeof(vector_31));
    test_nistvector(vector_32, sizeof(vector_32));
    test_nistvector(vector_47, sizeof(vector_47));
    test_nistvector(vector_48, sizeof(vector_48));
    test_nistvector(vector_64, sizeof(vector_64));

    return 0;
}
#endif