/*
 * Copyright 2018 Amazon.com, Inc. or its affiliates. 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.
 */

/* 
* File containing transition functions. 
* Functions that directly map to functions in the AWS Encryption C SDK end with Model. 
* These are functions called in the loop of aws_cryptosdk_session_process in session.c (lines 180 - 243)
*/

module transition_functions where

import type_def
import helper_functions

// Cryptol translation of function try_write_header in session_encrypt.c (line 219 - 236)

// Success criteria : there is enough room in the buffer to write the header 

// Success case : state is updated
//                output size estimate is the size of the header
//                output buffer len has decreased by header size
//                output buffer ptr has increased by header size 

// Failure case : output size estimate is the size of the header 


tryWriteHeaderModel : state_t -> state_t
tryWriteHeaderModel state_old = 
    if state_old.output_buffer.len >= state_old.sess.header_size then 
        {
        sess = ( updateOutputInputEstimates 
                    ( advanceStateOnly state_old.sess ST_ENCRYPT_BODY ) 
                    state_old.sess.header_size
                    state_old.sess.input_size_estimate ), 
        output_buffer = {
                        ptr = state_old.output_buffer.ptr + state_old.sess.header_size,
                        len = state_old.output_buffer.len - state_old.sess.header_size
                        },  
        input_buffer = state_old.input_buffer
        }
    else 
        {
        sess = updateOutputInputEstimates 
            state_old.sess 
            state_old.sess.header_size 
            state_old.sess.input_size_estimate,
        output_buffer = state_old.output_buffer,
        input_buffer = state_old.input_buffer
        }


// Cryptol translation of function try_parse_header in session_decrypt.c (line 156 - 191)

// Success criteria : the full header can be read from the input buffer

// Success case : state is updated
//                input buffer len has decreased by header size
//                input buffer ptr has increased by header size 

// Failure case : input size estimate is increased by 128 
//                if this increase results in overflow: set to size_t - 1 
getHeaderSize : state_t -> size_bits_t
getHeaderSize state_old = 
    if ( state_old.sess.header.aad_len_total > 0 ) then 
        ( HEADER_FIELD_OVERHEAD_WITH_AAD + MESSAGE_ID_LEN 
            + state_old.sess.header.aad_len_total + state_old.sess.header.edk_len_total )
    else ( HEADER_FIELD_OVERHEAD + MESSAGE_ID_LEN 
            + state_old.sess.header.aad_len_total + state_old.sess.header.edk_len_total )

updateInputBuffer : state_t -> size_bits_t
updateInputBuffer state_old = 
    if ( state_old.sess.input_size_estimate < state_old.input_buffer.len ) then 
        if ( (state_old.input_buffer.len + DEFAULT_SIZE_INCREMENT) < state_old.input_buffer.len ) then 
             ( SIZE_T_MAX - 1 )
        else  (state_old.input_buffer.len + DEFAULT_SIZE_INCREMENT)
    else state_old.sess.input_size_estimate

tryParseHeaderModel : state_t -> state_t 
tryParseHeaderModel state_old = 
    if (state_old.input_buffer.len < ( getHeaderSize state_old ) ) then 
        {
        sess = updateOutputInputEstimates 
                state_old.sess 
                0 
                ( updateInputBuffer state_old ),
        output_buffer = state_old.output_buffer, //UNCHANGED
        input_buffer = state_old.input_buffer    //UNCHANGED 
        }

    else 
        {
        sess = advanceStateOnly state_old.sess ST_UNWRAP_KEY,
        output_buffer = state_old.output_buffer, //UNCHANGED 
        input_buffer = { 
                        ptr = state_old.input_buffer.ptr + ( getHeaderSize state_old ), 
                        len = state_old.output_buffer.ptr - ( getHeaderSize state_old )
                        }
        }
    

// Cryptol translation of function try_encrypt_body in session_encrypt.c (line 238 - 333)

// Next state criteria : the precise_size is known 
//                       we are encrypting a final frame (if message is framed)

// Other success criteria : the output buffer is long enough to write the ciphertext
//                          the input buffer contains all the plaintext                           

// Next state case : state is updated


// All success case : input buffer len has decreased by plaintext size
//                      input buffer ptr has increased by plaintext size
//                      output buffer len has decreased by ciphertext size
//                      output buffer ptr has increased by ciphertext size
//                      input size estimate is the plaintext size 
//                      output size estimate is the ciphertext size 

// Failure case : (if framed or precise size known) : input size estimate is the plaintext size 
//                                                    output size estimate is the ciphertext size 

// Failure case : (if unframed and precise size unknown) : input size estimate is 0
//                                                      output size estimate is 0
encryptHelper : state_t -> size_bits_t -> size_bits_t -> state_size_t -> state_t 
encryptHelper state_old size_in size_out new_state = 
    if ( state_old.input_buffer.len < size_in ) || ( state_old.output_buffer.len < size_out ) then 
        { 
        sess = updateOutputInputEstimates state_old.sess size_out size_in,
        output_buffer = state_old.output_buffer, //UNCHANGED 
        input_buffer = state_old.input_buffer    //UNCHANGED 
        }
    else 
        { 
        sess = updateOutputInputEstimates (advanceStateOnly (incrementFrameSeqno state_old.sess) new_state) size_out size_in,
                output_buffer = { 
                                ptr = state_old.output_buffer.ptr + size_out,
                                len = state_old.output_buffer.len - size_out
                                },
                input_buffer = { 
                                ptr = state_old.input_buffer.ptr + size_in,
                                len = state_old.input_buffer.len - size_in
                                }       
        }

encryptBodyFramed : state_t -> state_t
encryptBodyFramed state_old = 
    if ( state_old.sess.precise_size_known == 1 ) 
        && (state_old.sess.precise_size - state_old.sess.data_so_far < state_old.sess.frame_size) then 
            encryptHelper 
                state_old 
                (state_old.sess.precise_size - state_old.sess.data_so_far) 
                (FRAME_FIELD_OVERHEAD_FINAL_FRAME + state_old.sess.alg_prop.iv_len + 4 + state_old.sess.precise_size 
                    - state_old.sess.data_so_far + state_old.sess.alg_prop.tag_len) 
                ST_WRITE_TRAILER
    else encryptHelper 
            state_old 
            state_old.sess.frame_size 
            (FRAME_FIELD_OVERHEAD + state_old.sess.alg_prop.iv_len + state_old.sess.frame_size + state_old.sess.alg_prop.tag_len) 
            ST_ENCRYPT_BODY

encryptBodyUnframed : state_t -> state_t
encryptBodyUnframed state_old = 
    if ( state_old.sess.precise_size_known == 1 ) then 
        encryptHelper 
            state_old 
            state_old.sess.precise_size 
            (state_old.sess.alg_prop.iv_len + FRAME_FIELD_OVERHEAD_UNFRAMED + state_old.sess.precise_size + state_old.sess.alg_prop.tag_len) 
            ST_WRITE_TRAILER
    else { 
        sess = updateOutputInputEstimates state_old.sess 0 0,
        output_buffer = state_old.output_buffer,
        input_buffer = state_old.input_buffer 
        }


tryEncryptBodyModel : state_t -> state_t
tryEncryptBodyModel state_old = 
    if state_old.sess.frame_size > 0 then encryptBodyFramed state_old
    else encryptBodyUnframed state_old
        

// Partial Cryptol translation of function try_decrypt_body in session_decrypt.c (line 193 - 252)
decryptHelper : state_t -> size_bits_t -> size_bits_t -> state_size_t -> state_t
decryptHelper state_old size_in size_out new_state = 
    if ( state_old.input_buffer.len < size_in ) || ( state_old.output_buffer.len < size_out ) then 
        { 
        sess = updateOutputInputEstimates state_old.sess size_out size_in, 
        output_buffer = state_old.output_buffer,
        input_buffer = state_old.input_buffer
        }
    else 
        { 
        sess = updateOutputInputEstimates (advanceStateOnly state_old.sess new_state) size_out size_in,
        output_buffer = { 
                        ptr = state_old.output_buffer.ptr + size_out,
                        len = state_old.output_buffer.len - size_out
                        },
        input_buffer = { 
                        ptr = state_old.input_buffer.ptr + size_in,
                        len = state_old.input_buffer.len - size_in
                        }       
        }

//TODO: Size of the unframed data is stored in the ciphertext (input_buffer). What is the best way to model this?
//TODO: Should not be using state_old.sess.precise_size. This is only temporary. 
decryptBodyUnframed : state_t -> state_t 
decryptBodyUnframed state_old = decryptHelper state_old ( state_old.sess.alg_prop.iv_len + 8 + 
state_old.sess.precise_size + state_old.sess.alg_prop.tag_len)   state_old.sess.precise_size ST_CHECK_TRAILER

//TODO: Whether the frame is final as well as the size of the final frame is stored in the ciphertext (input_buffer). 
decryptBodyFramed : state_t -> state_t 
decryptBodyFramed state_old = state_old 

tryDecryptBodyModel : state_t -> state_t 
tryDecryptBodyModel state_old = 
    if ( state_old.sess.frame_size > 0 ) then decryptBodyUnframed state_old
    else decryptBodyUnframed state_old