// Copyright 2024, The Android Open Source Project
//
// 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.

//! Helper wrapper around PostProcessor interface.

use android_hardware_security_keymint::aidl::android::hardware::security::keymint::Certificate::Certificate;
use android_security_postprocessor::aidl::android::security::postprocessor::{
    CertificateChain::CertificateChain,
    IKeystoreCertificatePostProcessor::IKeystoreCertificatePostProcessor,
};
use anyhow::{Context, Result};
use binder::{StatusCode, Strong};
use log::{error, info, warn};
use message_macro::source_location_msg;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::mpsc;
use std::thread;
use std::time::Duration;

/// Errors occurred during the interaction with Certificate Processor
#[derive(Debug, Clone, Copy, thiserror::Error, PartialEq, Eq)]
#[error("Binder transaction error {0:?}")]
pub struct Error(pub StatusCode);

static CERT_PROCESSOR_FAILURE: AtomicBool = AtomicBool::new(false);

fn send_certificate_chain_to_processor(
    attestation_chain: CertificateChain,
) -> Result<CertificateChain> {
    let cert_processing_server: Strong<dyn IKeystoreCertificatePostProcessor> = wait_for_interface(
        "rkp_cert_processor.service".to_string(),
    )
    .context(source_location_msg!("While trying to connect to the post processor service."))?;
    cert_processing_server
        .processKeystoreCertificates(&attestation_chain)
        .context(source_location_msg!("While trying to post process certificates."))
}

/// Processes the keystore certificates after the certificate chain has been generated by Keystore.
/// More details about this function provided in IKeystoreCertificatePostProcessor.aidl
pub fn process_certificate_chain(
    mut certificates: Vec<Certificate>,
    attestation_certs: Vec<u8>,
) -> Vec<Certificate> {
    // If no certificates are provided from keymint, return the original chain.
    if certificates.is_empty() {
        error!("No leaf certificate provided.");
        return vec![Certificate { encodedCertificate: attestation_certs }];
    }

    if certificates.len() > 1 {
        warn!("dropping {} unexpected extra certificates after the leaf", certificates.len() - 1);
    }

    let attestation_chain = CertificateChain {
        leafCertificate: certificates[0].encodedCertificate.clone(),
        remainingChain: attestation_certs.clone(),
    };
    let result = send_certificate_chain_to_processor(attestation_chain);
    match result {
        Ok(certificate_chain) => {
            info!("Post processing successful. Replacing certificates.");
            vec![
                Certificate { encodedCertificate: certificate_chain.leafCertificate },
                Certificate { encodedCertificate: certificate_chain.remainingChain },
            ]
        }
        Err(err) => {
            warn!("Failed to replace certificates ({err:#?}), falling back to original chain.");
            certificates.push(Certificate { encodedCertificate: attestation_certs });
            certificates
        }
    }
}

fn wait_for_interface(
    service_name: String,
) -> Result<Strong<dyn IKeystoreCertificatePostProcessor>> {
    if CERT_PROCESSOR_FAILURE.load(Ordering::Relaxed) {
        return Err(Error(StatusCode::INVALID_OPERATION).into());
    }
    let (sender, receiver) = mpsc::channel();
    let _t = thread::spawn(move || {
        if let Err(e) = sender.send(binder::wait_for_interface(&service_name)) {
            error!("failed to send result of wait_for_interface({service_name}), likely due to timeout: {e:?}");
        }
    });

    match receiver.recv_timeout(Duration::from_secs(5)) {
        Ok(service_binder) => Ok(service_binder?),
        Err(e) => {
            error!("Timed out while connecting to post processor service: {e:#?}");
            // Cert processor has failed. Retry only after reboot.
            CERT_PROCESSOR_FAILURE.store(true, Ordering::Relaxed);
            Err(e.into())
        }
    }
}