////////////////////////////////////////////////////////////////
// Copyright 2019 Galois, Inc. All Rights Reserved
//
// Licensed under the Apache License, Version 2.0 (the "License").
// You may not use this file except in compliance with the License.
// A copy of the License is located at
//
//  http://aws.amazon.com/apache2.0
//
// or in the "license" file accompanying this file. This file is distributed
// on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
// express or implied. See the License for the specific language governing
// permissions and limitations under the License.
//
////////////////////////////////////////////////////////////////
//
// This file describes the correspondence between the Cryptol
// specification of DRBG in DRBG.cry and the C implementation in
// crypto/s2n_drbg.c, allowing SAW to prove them equivalent.
//
////////////////////////////////////////////////////////////////

import "DRBG.cry";

////////////////////////////////////////////////////////////////////////////////
// Generic utility functions
////////////////////////////////////////////////////////////////////////////////

let alloc_init ty v = do {
    p <- crucible_alloc ty;
    crucible_points_to p v;
    return p;
};

let alloc_init_readonly ty v = do {
    p <- crucible_alloc_readonly ty;
    crucible_points_to p v;
    return p;
};

let ptr_to_fresh n ty = do {
    x <- crucible_fresh_var n ty;
    p <- alloc_init ty (crucible_term x);
    return (x, p);
};

let ptr_to_fresh_readonly n ty = do {
    x <- crucible_fresh_var n ty;
    p <- alloc_init_readonly ty (crucible_term x);
    return (x, p);
};

let blocksize = 16; // blocklen / 8
let keysize = 16; // keylen / 8
let seedsize = 32;

let bool = llvm_int 1;
let i8 = llvm_int 8;
let i32 = llvm_int 32;
let i64 = llvm_int 64;
let tm = crucible_term;

let bytes_type n = llvm_array n i8;
let alloc_bytes n = crucible_alloc (bytes_type n);

////////////////////////////////////////////////////////////////////////////////
// Convenient Cryptol definitions
////////////////////////////////////////////////////////////////////////////////

let {{

max_seq_number : [blocksize*8]
max_seq_number = ~zero

drbg_generate_seedlen : s2n_drbg -> [seedlen] -> Bit -> ([seedlen], s2n_drbg)
drbg_generate_seedlen = drbg_generate `{n=seedlen,blocks=2}

encrypt_128 : [keysize][8] -> [blocksize][8] -> [blocksize][8]
encrypt_128 key msg =
  split (block_encrypt (join key) (join msg))

mode_128 = 0
mode_256 = 1

}};

// A symbolic variable representing the entropy returned by the system.
fake_entropy <- fresh_symbolic "fake_entropy" {| [seedsize*8] |};

////////////////////////////////////////////////////////////////////////////////
// Types specific to s2n
////////////////////////////////////////////////////////////////////////////////

let ctx_type = bytes_type keysize;

let blob_type = llvm_struct "struct.s2n_blob";
let alloc_blob n = do {
    p <- crucible_alloc_readonly blob_type;
    datap <- alloc_bytes n;
    crucible_points_to (crucible_field p "data") datap;
    crucible_points_to (crucible_field p "size") (tm {{ `n : [32] }});
    crucible_points_to (crucible_field p "allocated") (tm {{ 0 : [32] }});
    crucible_points_to (crucible_field p "growable") (tm {{ 0 : [8] }});
    return (p, datap);
};

let alloc_blob_readonly n = do {
    p <- crucible_alloc_readonly blob_type;
    datap <- crucible_alloc_readonly (llvm_array n i8);
    crucible_points_to (crucible_field p "data") datap;
    crucible_points_to (crucible_field p "size") (tm {{ `n : [32] }});
    return (p, datap);
};

let drbg_state n = do {
    state <- crucible_alloc (llvm_struct "struct.s2n_drbg");
    (key, keyp) <- ptr_to_fresh "key" ctx_type;
    bytes_used <- crucible_fresh_var (str_concat n "->bytes_used") i64;
    mixes <- crucible_fresh_var (str_concat n "->mixes") i64;
    v <- crucible_fresh_var (str_concat n "->v") (bytes_type blocksize);
    crucible_points_to (crucible_field state "bytes_used") (tm bytes_used);
    crucible_points_to (crucible_field state "mixes") (tm mixes);
    crucible_points_to (crucible_field state "ctx") keyp;
    crucible_points_to (crucible_field state "v") (tm v);
    return ( state
           , keyp
           , {{ { bytes_used = bytes_used
               , ctx = { key = join key }
               , v = join v
               } }}
           );
};

let ensure_drbg_state p keyp s = do {
    crucible_points_to (crucible_field p "bytes_used") (tm {{ s.bytes_used }});
    crucible_points_to (crucible_field p "ctx") keyp;
    crucible_points_to_untyped keyp (tm {{
        split s.ctx.key : [keysize][8]
    }});
    crucible_points_to (crucible_field p "v") (tm {{
        split s.v : [blocksize][8]
    }});
    mixes <- crucible_fresh_var ("mixes'") i64;
    crucible_points_to (crucible_field p "mixes") (tm mixes);
};

////////////////////////////////////////////////////////////////////////////////
// Assumed specifications
////////////////////////////////////////////////////////////////////////////////

let getenv_spec = do {
    p <- crucible_fresh_pointer i8;
    crucible_execute_func [p];
    crucible_return crucible_null;
};

let aes_128_ecb_spec = do {
    crucible_execute_func [];
    r <- crucible_fresh_pointer ctx_type;
    crucible_return r;
};

let cipher_new_spec = do {
    crucible_execute_func [];
    r <- crucible_alloc ctx_type;
    crucible_return r;
};

let cipher_init_spec = do {
    ctx <- crucible_alloc ctx_type;
    crucible_execute_func [ctx];
    key <- crucible_fresh_var "key" ctx_type;
    crucible_points_to ctx (crucible_term key);
};

let cipher_free_spec = do {
    ctx <- crucible_alloc ctx_type;
    crucible_execute_func [ctx];
};

let cipher_cleanup_spec = do {
    ctx <- crucible_alloc ctx_type;
    crucible_execute_func [ctx];
    crucible_points_to ctx (tm {{ zero : [keysize][8] }});
    crucible_return (tm {{ 1 : [32] }});
};

let cipher_key_length_spec = do {
    ctx <- crucible_alloc_readonly ctx_type;
    crucible_execute_func [ctx];
    // Specialized to AES-128 for now
    crucible_return (tm {{ 16 : [32] }});
};

let encryptInit_spec = do {
    ctx <- crucible_alloc ctx_type;
    st <- crucible_fresh_pointer ctx_type;
    (key, keyp) <- ptr_to_fresh_readonly "key" ctx_type;
    crucible_execute_func [ctx, st, crucible_null, keyp, crucible_null];
    crucible_points_to ctx (tm key);
    crucible_return (tm {{ 1 : [32] }});
};

let encryptInit_nokey_spec = do {
    ctx <- crucible_alloc ctx_type;
    st <- crucible_fresh_pointer ctx_type;
    crucible_execute_func [ctx, st, crucible_null, crucible_null, crucible_null];
    crucible_return (tm {{ 1 : [32] }});
};

let encryptUpdate_spec n = do {
    // the first argument of `EVP_EncryptUpdate` is not `const`,
    // but it is constant in the DRBG cryptol specification.
    (key, keyp) <- ptr_to_fresh_readonly "key" ctx_type;
    outp <- alloc_bytes n;
    lenp <- alloc_init i32 (tm {{ `n : [32] }});
    (msg, msgp) <- ptr_to_fresh_readonly "msg" (bytes_type n);
    crucible_execute_func [keyp, outp, lenp, msgp, tm {{ `blocksize : [32] }} ];
    crucible_points_to outp (tm {{ encrypt_128 key msg }});
    crucible_points_to lenp (tm {{ `n : [32] }});
    crucible_return (tm {{ 1 : [32] }});
};

let get_public_random_spec = do {
    (p, datap) <- alloc_blob seedsize;
    crucible_execute_func [p];
    crucible_points_to datap (tm {{ split fake_entropy : [seedsize][8] }});
    crucible_return (tm {{ 0 : [32] }});
};

let supports_rdrand_spec = do {
    crucible_execute_func [];
    r <- crucible_fresh_var "supports_rdrand" bool;
    crucible_return (tm r);
};

////////////////////////////////////////////////////////////////////////////////
// Specifications to be verified
////////////////////////////////////////////////////////////////////////////////

let block_encrypt_spec = do {
    (key, keyp) <- ptr_to_fresh_readonly "ctx" ctx_type;
    (msg, msgp) <- ptr_to_fresh_readonly "msg" (bytes_type blocksize);
    outp <- alloc_bytes blocksize;
    crucible_execute_func [keyp, msgp, outp];
    crucible_points_to outp (tm {{ encrypt_128 key msg }});
    crucible_return (tm {{ 0 : [32] }});
};

let blob_zero_spec n = do {
    (p, datap) <- alloc_blob n;
    crucible_execute_func [p];
    crucible_points_to datap (tm {{ zero : [n][8] }});
    crucible_return (tm {{ 0 : [32] }});
};

let increment_drbg_counter_spec = do {
    (p, datap) <- alloc_blob blocksize;
    v <- crucible_fresh_var "v" (bytes_type blocksize);
    crucible_points_to datap (tm v);
    crucible_execute_func [p];
    let res = {{ split ((join v) + 1) : [blocksize][8] }};
    crucible_points_to datap (tm res);
    crucible_return (tm {{ 0 : [32] }});
};

let bits_spec n = do {
    (sp, keyp, s) <- drbg_state "drbg";
    (outp, datap) <- alloc_blob n;
    crucible_execute_func [sp, outp];
    let res = {{ drbg_generate_internal `{n=n*8} s }};
    crucible_points_to datap (tm {{ split res.0 : [n][8] }});
    ensure_drbg_state sp keyp {{ res.1 }};
    crucible_return (tm {{ 0 : [32] }});
};

let update_spec n = do {
    (sp, keyp, s) <- drbg_state "drbg";
    (providedp, datap) <- alloc_blob_readonly n;
    data <- crucible_fresh_var "data" (bytes_type n);
    crucible_points_to datap (tm data);
    crucible_execute_func [sp, providedp];
    ensure_drbg_state sp keyp {{ drbg_update (join data) s }};
    crucible_return (tm {{ 0 : [32] }});
};

let seed_spec n = do {
    (sp, keyp, s) <- drbg_state "drbg";
    (psp, datap) <- alloc_blob_readonly n;
    data <- crucible_fresh_var "data" (bytes_type n);
    crucible_points_to datap (tm data);
    crucible_execute_func [sp, psp];
    ensure_drbg_state sp keyp {{ drbg_reseed s fake_entropy (join data) }};
    crucible_return (tm {{ 0 : [32] }});
};

let mix_spec n = do {
    (crucible_alloc_global "ignore_prediction_resistance_for_testing");
    (crucible_points_to (crucible_global "ignore_prediction_resistance_for_testing") (crucible_term {{ 0 : [8] }}) );
    (state_ptr, key_ptr, state) <- drbg_state "drbg";
    (personalization_string_ptr, data_ptr) <- alloc_blob_readonly n;
    data <- crucible_fresh_var "data" (bytes_type n);
    crucible_points_to data_ptr (tm data);
    crucible_execute_func [state_ptr, personalization_string_ptr];
    (crucible_points_to (crucible_global "ignore_prediction_resistance_for_testing") (crucible_term {{ 0 : [8] }}) );
    ensure_drbg_state state_ptr key_ptr {{ drbg_mix state fake_entropy (join data) }};
    crucible_return (tm {{ 0 : [32] }});
};

let instantiate_spec n = do {
    (sp, keyp, s) <- drbg_state "drbg";
    (psp, datap) <- alloc_blob_readonly n;
    data <- crucible_fresh_var "data" (bytes_type n);
    crucible_points_to datap (tm data);

    crucible_execute_func [sp, psp, tm {{ mode_128 : [32] }}];

    newkeyp <- crucible_alloc ctx_type;
    crucible_points_to (crucible_field sp "ctx") newkeyp;
    ensure_drbg_state sp newkeyp {{ drbg_instantiate fake_entropy (join data) }};
    crucible_return (tm {{ 0 : [32] }});
};

let generate_spec = do {
    (crucible_alloc_global "ignore_prediction_resistance_for_testing");
    (crucible_points_to (crucible_global "ignore_prediction_resistance_for_testing") (crucible_term {{ 0 : [8] }}) );
    (sp, keyp, s) <- drbg_state "drbg";
    (psp, datap) <- alloc_blob seedsize;
    crucible_execute_func [sp, psp];
    (crucible_points_to (crucible_global "ignore_prediction_resistance_for_testing") (crucible_term {{ 0 : [8] }}) );
    let {{ res = drbg_generate_seedlen s fake_entropy True }};
    crucible_points_to datap (tm {{ split res.0 : [seedsize][8] }});
    ensure_drbg_state sp keyp {{ res.1 }};
    crucible_return (tm {{ 0 : [32] }});
};

let bytes_used_spec = do {
    (sp, keyp, s) <- drbg_state "drbg";
    bytes_used <- alloc_init i64 (tm {{ 0 : [64] }});
    crucible_execute_func [sp, bytes_used];
    crucible_points_to bytes_used (tm {{ s.bytes_used }});
    crucible_return (tm {{ 0 : [32] }});
};

////////////////////////////////////////////////////////////////////////////////
// Bitcode processing
////////////////////////////////////////////////////////////////////////////////

m <- llvm_load_module "../../bitcode/all_llvm.bc";

////////////////////////////////////////////////////////////////////////////////
// Assumptions about external functions
////////////////////////////////////////////////////////////////////////////////

getenv_ov <- crucible_llvm_unsafe_assume_spec m "getenv" getenv_spec;

aes_128_ecb_ov <- crucible_llvm_unsafe_assume_spec m "EVP_aes_128_ecb" aes_128_ecb_spec;
cipher_new_ov <- crucible_llvm_unsafe_assume_spec m "EVP_CIPHER_CTX_new" cipher_new_spec;

cipher_free_ov <- crucible_llvm_unsafe_assume_spec m "EVP_CIPHER_CTX_free" cipher_free_spec;

cipher_cleanup_ov <- crucible_llvm_unsafe_assume_spec m "EVP_CIPHER_CTX_reset" cipher_cleanup_spec;

cipher_key_length_ov <- crucible_llvm_unsafe_assume_spec m "EVP_CIPHER_CTX_key_length" cipher_key_length_spec;

encryptInit_ov <- crucible_llvm_unsafe_assume_spec m "EVP_EncryptInit_ex"
    encryptInit_spec;

encryptInit_nokey_ov <- crucible_llvm_unsafe_assume_spec m "EVP_EncryptInit_ex"
    encryptInit_nokey_spec;

encryptUpdate_ov <- crucible_llvm_unsafe_assume_spec m "EVP_EncryptUpdate"
    (encryptUpdate_spec 16);

supports_rdrand_ov <- crucible_llvm_unsafe_assume_spec m "s2n_cpu_supports_rdrand" supports_rdrand_spec;

get_public_random_ov <- crucible_llvm_unsafe_assume_spec m "s2n_get_public_random_data" get_public_random_spec;

get_seed_entropy_ov <- crucible_llvm_unsafe_assume_spec m "s2n_get_seed_entropy" get_public_random_spec;

get_mix_entropy_ov <- crucible_llvm_unsafe_assume_spec m "s2n_get_mix_entropy" get_public_random_spec;

////////////////////////////////////////////////////////////////////////////////
// Proofs about internal functions
////////////////////////////////////////////////////////////////////////////////

zero_ov_block <- crucible_llvm_verify m "s2n_blob_zero" [] false (blob_zero_spec blocksize) (w4_unint_yices []);
zero_ov_seed <- crucible_llvm_verify m "s2n_blob_zero" [] false (blob_zero_spec seedsize) (w4_unint_yices []);
zero_ov_drbg <- crucible_llvm_verify m "s2n_blob_zero" [] false (blob_zero_spec 48) (w4_unint_yices []);

inc_ov <- crucible_llvm_verify m "s2n_increment_drbg_counter" [] false increment_drbg_counter_spec (w4_unint_yices []);

crucible_llvm_verify m "s2n_drbg_bytes_used" [] false bytes_used_spec (w4_unint_yices []);

blk_enc_ov <- crucible_llvm_verify m "s2n_drbg_block_encrypt" [encryptUpdate_ov] false block_encrypt_spec (w4_unint_yices ["block_encrypt"]);

bits_ov <- crucible_llvm_verify m "s2n_drbg_bits" [inc_ov, encryptUpdate_ov, blk_enc_ov] false (bits_spec seedsize) (w4_unint_yices ["block_encrypt"]);

update_ov <- crucible_llvm_verify m "s2n_drbg_update" [bits_ov, encryptInit_ov, aes_128_ecb_ov, cipher_key_length_ov] false (update_spec seedsize) (w4_unint_yices ["block_encrypt"]);

seed_ov <- crucible_llvm_verify m "s2n_drbg_seed" [get_public_random_ov, get_seed_entropy_ov, update_ov, cipher_key_length_ov] false (seed_spec seedsize) (w4_unint_yices ["block_encrypt"]);

mix_ov <- crucible_llvm_verify m "s2n_drbg_mix" [get_public_random_ov, get_mix_entropy_ov, update_ov, cipher_key_length_ov] false (mix_spec seedsize) (w4_unint_yices ["block_encrypt"]);

instantiate_ov <- crucible_llvm_verify m "s2n_drbg_instantiate" [zero_ov_block, zero_ov_seed, cipher_new_ov, cipher_cleanup_ov, encryptInit_ov, encryptInit_nokey_ov, aes_128_ecb_ov, cipher_key_length_ov, seed_ov, supports_rdrand_ov] false(instantiate_spec seedsize) (w4_unint_yices ["block_encrypt"]);

generate_ov <- crucible_llvm_verify m "s2n_drbg_generate" [getenv_ov, mix_ov, bits_ov, update_ov, cipher_key_length_ov] false generate_spec (w4_unint_yices ["block_encrypt"]);