/*

 *    Copyright (c) 2020 Project CHIP Authors

 *    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.

 *    You may obtain a copy of the License at

 *

 *        http://www.apache.org/licenses/LICENSE-2.0

 *

 *    Unless required by applicable law or agreed to in writing, software

 *    distributed under the License 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

 *      This file implements the Matter Checkin protocol.

 */

#include <lib/core/CHIPCore.h>

#include <lib/core/CHIPEncoding.h>

#include <protocols/secure_channel/CheckinMessage.h>

#include <protocols/secure_channel/Constants.h>

using namespace chip::Crypto;

namespace chip {

namespace Protocols {

namespace SecureChannel {

CHIP_ERROR CheckinMessage::GenerateCheckinMessagePayload(const Crypto::Aes128KeyHandle & aes128KeyHandle,

                                                         const Crypto::Hmac128KeyHandle & hmacKeyHandle,

                                                         const CounterType & counter, const ByteSpan & appData,

                                                         MutableByteSpan & output)

{

    VerifyOrReturnError(output.size() >= (appData.size() + kMinPayloadSize), CHIP_ERROR_BUFFER_TOO_SMALL);

    size_t cursorIndex = 0;

    // Generate Nonce from Key and counter value

    {

        MutableByteSpan nonce = output.SubSpan(0, CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES);

        cursorIndex += nonce.size();

        Encoding::LittleEndian::BufferWriter writer(nonce);

        ReturnErrorOnFailure(GenerateCheckInMessageNonce(hmacKeyHandle, counter, writer));

    }

    // Encrypt Counter and Application Data

    {

        MutableByteSpan payloadByteSpan = output.SubSpan(cursorIndex, sizeof(CounterType) + appData.size());

        cursorIndex += payloadByteSpan.size();

        Encoding::LittleEndian::BufferWriter payloadWriter(payloadByteSpan);

        payloadWriter.EndianPut(counter, sizeof(counter));

        payloadWriter.Put(appData.data(), appData.size());

        VerifyOrReturnError(payloadWriter.Fit(), CHIP_ERROR_INTERNAL);

        MutableByteSpan mic = output.SubSpan(cursorIndex, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES);

        cursorIndex += mic.size();

        // Validate that the cursorIndex is within the available output space

        VerifyOrReturnError(cursorIndex <= output.size(), CHIP_ERROR_BUFFER_TOO_SMALL);

        // Validate that the cursorIndex matchs the message length

        VerifyOrReturnError(cursorIndex == appData.size() + kMinPayloadSize, CHIP_ERROR_INTERNAL);

        ReturnErrorOnFailure(Crypto::AES_CCM_encrypt(payloadByteSpan.data(), payloadByteSpan.size(), nullptr, 0, aes128KeyHandle,

                                                     output.data(), CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES, payloadByteSpan.data(),

                                                     mic.data(), mic.size()));

    }

    output.reduce_size(appData.size() + kMinPayloadSize);

    return CHIP_NO_ERROR;

}

CHIP_ERROR CheckinMessage::ParseCheckinMessagePayload(const Crypto::Aes128KeyHandle & aes128KeyHandle,

                                                      const Crypto::Hmac128KeyHandle & hmacKeyHandle, const ByteSpan & payload,

                                                      CounterType & counter, MutableByteSpan & appData)

{

    size_t appDataSize = GetAppDataSize(payload);

    VerifyOrReturnError(payload.size() >= kMinPayloadSize, CHIP_ERROR_INVALID_MESSAGE_LENGTH);

    // To prevent workbuffer usage, appData size needs to be large enough to hold both the appData and the counter

    VerifyOrReturnError(appData.size() >= sizeof(CounterType) + appDataSize, CHIP_ERROR_BUFFER_TOO_SMALL);

    // Decrypt received data

    {

        size_t cursorIndex = 0;

        ByteSpan nonce = payload.SubSpan(cursorIndex, CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES);

        cursorIndex += nonce.size();

        ByteSpan encryptedData = payload.SubSpan(cursorIndex, sizeof(CounterType) + appDataSize);

        cursorIndex += encryptedData.size();

        ByteSpan mic = payload.SubSpan(cursorIndex, CHIP_CRYPTO_AEAD_MIC_LENGTH_BYTES);

        cursorIndex += mic.size();

        // Return Invalid message length since the payload isn't the right size

        VerifyOrReturnError(cursorIndex == payload.size(), CHIP_ERROR_INVALID_MESSAGE_LENGTH);

        ReturnErrorOnFailure(Crypto::AES_CCM_decrypt(encryptedData.data(), encryptedData.size(), nullptr, 0, mic.data(), mic.size(),

                                                     aes128KeyHandle, nonce.data(), nonce.size(), appData.data()));

    }

    // Read decrypted counter and application data

    static_assert(sizeof(CounterType) == sizeof(uint32_t), "Expect counter to be 32 bits for correct decoding");

    CounterType tempCounter = Encoding::LittleEndian::Get32(appData.data());

    // Validate that the received nonce is correct

    {

        uint8_t calculatedNonceBuffer[CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES] = { 0 };

        Encoding::LittleEndian::BufferWriter writer(calculatedNonceBuffer, sizeof(calculatedNonceBuffer));

        ReturnErrorOnFailure(GenerateCheckInMessageNonce(hmacKeyHandle, tempCounter, writer));

        // Validate received nonce is the same as the calculated

        ByteSpan nonce = payload.SubSpan(0, CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES);

        VerifyOrReturnError(memcmp(nonce.data(), calculatedNonceBuffer, sizeof(calculatedNonceBuffer)) == 0, CHIP_ERROR_INTERNAL);

    }

    // We have successfully decrypted and validated Check-In message

    // Set output values

    counter = tempCounter;

    // Shift to remove the counter from the appData

    memmove(appData.data(), sizeof(CounterType) + appData.data(), appDataSize);

    appData.reduce_size(appDataSize);

    return CHIP_NO_ERROR;

}

CHIP_ERROR CheckinMessage::GenerateCheckInMessageNonce(const Crypto::Hmac128KeyHandle & hmacKeyHandle, CounterType counter,

                                                       Encoding::LittleEndian::BufferWriter & writer)

{

    VerifyOrReturnError(writer.Available() >= CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES, CHIP_ERROR_BUFFER_TOO_SMALL);

    uint8_t hashWorkBuffer[CHIP_CRYPTO_HASH_LEN_BYTES] = { 0 };

    uint8_t counterBuffer[sizeof(CounterType)];

    // validate that Check-In counter is a uint32_t

    static_assert(sizeof(CounterType) == sizeof(uint32_t), "Expect counter to be 32 bits for correct encoding");

    Encoding::LittleEndian::Put32(counterBuffer, counter);

    chip::Crypto::HMAC_sha shaHandler;

    ReturnErrorOnFailure(

        shaHandler.HMAC_SHA256(hmacKeyHandle, counterBuffer, sizeof(CounterType), hashWorkBuffer, CHIP_CRYPTO_HASH_LEN_BYTES));

    writer.Put(hashWorkBuffer, CHIP_CRYPTO_AEAD_NONCE_LENGTH_BYTES);

    VerifyOrReturnError(writer.Fit(), CHIP_ERROR_BUFFER_TOO_SMALL);

    return CHIP_NO_ERROR;

}

size_t CheckinMessage::GetAppDataSize(const ByteSpan & payload)

{

    return (payload.size() <= kMinPayloadSize) ? 0 : payload.size() - kMinPayloadSize;

}

} // namespace SecureChannel

} // namespace Protocols

} // namespace chip