sequoia-cert-store 0.7.3

use std::collections::BTreeMap;
use std::collections::btree_map::Entry;
use std::sync::Arc;
use std::sync::RwLock;

use smallvec::SmallVec;
use smallvec::smallvec;

use anyhow::Context;

use sequoia_openpgp as openpgp;
use openpgp::cert::Cert;
use openpgp::cert::raw::RawCert;
use openpgp::cert::raw::RawCertParser;
use openpgp::Fingerprint;
use openpgp::KeyID;
use openpgp::KeyHandle;
use openpgp::packet::UserID;
use openpgp::parse::Parse;
use openpgp::Result;

use crate::LazyCert;
use crate::store::MergeCerts;
use crate::store::Store;
use crate::store::StoreError;
use crate::store::StoreUpdate;
use crate::store::UserIDIndex;
use crate::store::UserIDQueryParams;

const TRACE: bool = false;

/// Manages a slice of bytes, `RawCert`s, `Cert`s, or `LazyCert`s.
///
/// `Cert`s implements `StoreUpdate`, but it does not write the
/// updates to disk; they are only updated in memory.
pub struct Certs<'a> {
    inner: RwLock<CertsInner<'a>>,
}
assert_send_and_sync!(Certs<'_>);

struct CertsInner<'a> {
    // Indexed by primary key fingerprint.
    certs: BTreeMap<Fingerprint, Arc<LazyCert<'a>>>,
    // Indexed by a key's KeyID (primary key or subkey) and maps to
    // the primary key.
    keys: BTreeMap<KeyID, SmallVec<[Fingerprint; 1]>>,

    userid_index: UserIDIndex,
}

impl<'a> Certs<'a>
{
    /// Returns an empty `Certs` store.
    ///
    /// This is useful as a placeholder.  But, certificates can also
    /// be added to it using the [`StoreUpdate`] interface.
    pub fn empty() -> Self {
        Certs {
            inner: RwLock::new(CertsInner {
                certs: Default::default(),
                keys: Default::default(),
                userid_index: UserIDIndex::new(),
            }),
        }
    }

    /// Returns a new `Certs`.
    pub fn from_bytes(bytes: &'a [u8]) -> Result<Self> {
        tracer!(TRACE, "Certs::from_bytes");

        let raw_certs = RawCertParser::from_bytes(bytes)?
            .filter_map(|c| {
                match c {
                    Ok(c) => Some(c),
                    Err(err) => {
                        t!("Parsing raw certificate: {}", err);
                        None
                    }
                }
            });
        Self::from_certs(raw_certs)
    }

    /// Returns a new `Certs`.
    pub fn from_certs<I>(certs: impl IntoIterator<Item=I>)
        -> Result<Self>
        where I: Into<LazyCert<'a>>
    {
        let r = Self::empty();
        for cert in certs {
            r.update(Arc::new(cert.into())).expect("implementation doesn't fail")
        }

        Ok(r)
    }
}

impl<'a> Store<'a> for CertsInner<'a>
{
    fn lookup_by_cert(&self, kh: &KeyHandle) -> Result<Vec<Arc<LazyCert<'a>>>> {
        tracer!(TRACE, "Certs::lookup_by_cert");

        match kh {
            KeyHandle::Fingerprint(fpr) => {
                self.lookup_by_cert_fpr(fpr).map(|c| vec![ c ])
            }
            KeyHandle::KeyID(keyid) => {
                let certs: Vec<Arc<LazyCert>> = self.keys.get(keyid)
                    .ok_or_else(|| {
                        anyhow::Error::from(
                            StoreError::NotFound(kh.clone()))
                    })?
                    .iter()
                    .filter_map(|fpr| self.certs.get(fpr).cloned())
                    // Check the constraints before we convert the
                    // rawcert to a cert.
                    .filter(|cert| cert.key_handle().aliases(kh))
                    .collect();

                if certs.is_empty() {
                    Err(StoreError::NotFound(kh.clone()).into())
                } else {
                    Ok(certs)
                }
            }
        }
    }

    fn lookup_by_cert_fpr(&self, fingerprint: &Fingerprint) -> Result<Arc<LazyCert<'a>>> {
        tracer!(TRACE, "Certs::lookup_by_cert_fpr");

        if let Some(cert) = self.certs.get(fingerprint).cloned() {
            Ok(cert)
        } else {
            Err(StoreError::NotFound(
                KeyHandle::from(fingerprint.clone())).into())
        }
    }

    fn lookup_by_cert_or_subkey(&self, kh: &KeyHandle) -> Result<Vec<Arc<LazyCert<'a>>>> {
        tracer!(TRACE, "Certs::lookup_by_cert_or_subkey");

        let keyid = KeyID::from(kh);
        let certs: Vec<Arc<LazyCert<'a>>> = self.keys.get(&keyid)
            .ok_or_else(|| {
                anyhow::Error::from(
                    StoreError::NotFound(kh.clone()))
            })?
            .iter()
            .filter_map(|fpr| self.certs.get(fpr))
            // Check the constraints before we convert the rawcert to a
            // cert.
            .filter(|cert| {
                cert.keys().any(|k| k.key_handle().aliases(kh))
            })
            .cloned()
            .collect();

        if certs.is_empty() {
            Err(StoreError::NotFound(kh.clone()).into())
        } else {
            Ok(certs)
        }
    }

    fn select_userid(&self, params: &UserIDQueryParams, pattern: &str)
        -> Result<Vec<Arc<LazyCert<'a>>>>
    {
        tracer!(TRACE, "Certs::select_userid");
        t!("params: {:?}, pattern: {:?}", params, pattern);

        let matches = self.userid_index.select_userid(params, pattern)?;

        let matches = matches
            .into_iter()
            .map(|fpr| {
                self.lookup_by_cert_fpr(&fpr).expect("indexed")
            })
            .collect();

        Ok(matches)
    }

    fn fingerprints<'b>(&'b self) -> Box<dyn Iterator<Item=Fingerprint> + 'b> {
        Box::new(self.certs.keys().cloned())
    }

    fn certs<'b>(&'b self) -> Box<dyn Iterator<Item=Arc<LazyCert<'a>>> + 'b>
        where 'a: 'b
    {
        Box::new(self.certs.values().cloned())
    }

    fn prefetch_all(&self) {
        // Because the interior mutability is implemented via Certs,
        // not CertInner, and we need a mutable reference here, but we
        // only have a normal, non-mutable reference, the prefetch
        // implementation is dispatched via `Certs`.
    }

    fn prefetch_some(&self, _certs: &[KeyHandle]) {
        // Because the interior mutability is implemented via Certs,
        // not CertInner, and we need a mutable reference here, but we
        // only have a normal, non-mutable reference, the prefetch
        // implementation is dispatched via `Certs`.
    }
}

impl<'a> Store<'a> for Certs<'a>
{
    fn lookup_by_cert(&self, kh: &KeyHandle) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.lookup_by_cert(kh)
    }

    fn lookup_by_cert_fpr(&self, fingerprint: &Fingerprint)
        -> Result<Arc<LazyCert<'a>>>
    {
        let inner = self.inner.read().unwrap();
        inner.lookup_by_cert_fpr(fingerprint)
    }

    fn lookup_by_cert_or_subkey(&self, kh: &KeyHandle) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.lookup_by_cert_or_subkey(kh)
    }

    fn select_userid(&self, query: &UserIDQueryParams, pattern: &str)
        -> Result<Vec<Arc<LazyCert<'a>>>>
    {
        let inner = self.inner.read().unwrap();
        inner.select_userid(query, pattern)
    }

    fn lookup_by_userid(&self, userid: &UserID) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.lookup_by_userid(userid)
    }

    fn grep_userid(&self, pattern: &str) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.grep_userid(pattern)
    }

    fn lookup_by_email(&self, email: &str) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.lookup_by_email(email)
    }

    fn grep_email(&self, pattern: &str) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.grep_email(pattern)
    }

    fn lookup_by_email_domain(&self, domain: &str) -> Result<Vec<Arc<LazyCert<'a>>>> {
        let inner = self.inner.read().unwrap();
        inner.lookup_by_email_domain(domain)
    }

    fn fingerprints<'b>(&'b self) -> Box<dyn Iterator<Item=Fingerprint> + 'b> {
        // We aren't lazy to avoid holding the lock.
        let inner = self.inner.read().unwrap();
        let fprs: Vec<_> = inner.fingerprints().collect();
        Box::new(fprs.into_iter())
    }

    fn certs<'b>(&'b self)
        -> Box<dyn Iterator<Item=Arc<LazyCert<'a>>> + 'b>
        where 'a: 'b
    {
        // We aren't lazy to avoid holding the lock.
        let inner = self.inner.read().unwrap();
        let certs: Vec<_> = inner.certs().collect();
        Box::new(certs.into_iter())
    }

    fn prefetch_all(&self) {
        let mut inner = self.inner.write().unwrap();
        inner.prefetch(true, &[])
    }

    fn prefetch_some(&self, certs: &[KeyHandle]) {
        let mut inner = self.inner.write().unwrap();
        inner.prefetch(false, certs)
    }
}

impl CertsInner<'_> {
    /// Prefetches the certs identified by `khs`, or `all`.
    fn prefetch(&mut self, all: bool, khs: &[KeyHandle]) {
        // LazyCert is Sync and Send, but keeping a reference to the
        // RawCerts in self.certs prevents us from later updating
        // self.certs.  This requires a bit of acrobatics to get
        // right.

        tracer!(TRACE, "Certs::prefetch_some");
        t!("Prefetch: {} certificates", khs.len());

        let mut certs: Vec<RawCert>
            = self.certs.iter().filter_map(|(fpr, cert)| {
                if ! cert.is_parsed() {
                    if all
                        || khs.iter()
                               .any(|kh| {
                                   kh.aliases(&KeyHandle::from(fpr.clone()))
                               })
                    {
                        // Unfortunately we have to clone the bytes,
                        // because we cannot keep self.certs borrowed.
                        t!("Queuing {} to be prefetched", fpr);
                        use std::ops::Deref;
                        cert.deref().clone().into_raw_cert().ok()
                    } else {
                        None
                    }
                } else {
                    None
                }
            }).collect();
        let cert_count = certs.len();

        // Use not more threads than we have cores.  No point in
        // spawning a single thread, just do it inline.  Also, some
        // platforms don't support spawning threads.
        let mut threads = num_cpus::get();
        if threads <= 1 {
            threads = 0;
        } else if cert_count < 16 {
            // The keyring is small, limit the number of threads.
            threads = 2;
        } else {
            // Sort the certificates so they are ordered from most
            // packets to least.  More packets implies more work, and
            // this will hopefully result in a more equal distribution
            // of load.
            certs.sort_unstable_by_key(|c| {
                usize::MAX - c.count()
            });
        }
        t!("Using {} threads", threads);

        let mut count = 0;

        if threads == 0 {
            for raw in certs {
                t!("Inline processing {}!", raw.keyid());
                // Silently ignore errors.  This will be caught later
                // when the caller looks this one up.
                match Cert::try_from(&raw) {
                    Ok(cert) => {
                        let fpr = cert.fingerprint();
                        t!("Caching {}", fpr);
                        self.certs.insert(fpr, Arc::new(cert.into()));
                        count += 1;
                    }
                    Err(err) => {
                        t!("Parsing raw cert {}: {}",
                           raw.keyid(), err);
                    }
                }
            }
        } else {
            use crossbeam::thread;
            use crossbeam::channel::unbounded as channel;

            // Avoid an extra level of indentation.
            let result = thread::scope(|thread_scope| {
            // A communication channel for sending work to the workers.
            let (work_tx, work_rx) = channel();
            // A communication channel for returning results to the
            // main thread.
            let (results_tx, results_rx) = channel();

            let mut threads_extant = Vec::new();

            for cert in certs.into_iter() {
                if threads_extant.len() < threads {
                    let tid = threads_extant.len();
                    t!("Starting thread {} of {}",
                       tid, threads);

                    let mut work = Some(Ok(cert));

                    // The thread's state.
                    let work_rx = work_rx.clone();
                    let results_tx = results_tx.clone();

                    threads_extant.push(thread_scope.spawn(move |_| {
                        loop {
                            match work.take().unwrap_or_else(|| work_rx.recv()) {
                                Err(_) => break,
                                Ok(raw) => {
                                    t!("Thread {} dequeuing {}!",
                                       tid, raw.keyid());
                                    // Silently ignore errors.  This will be
                                    // caught later when the caller looks
                                    // this one up.
                                    match Cert::try_from(&raw) {
                                        Ok(cert) => {
                                            let _ = results_tx.send(cert);
                                        }
                                        Err(err) => {
                                            t!("Parsing raw cert {}: {}",
                                               raw.keyid(), err);
                                        }
                                    }
                                }
                            }
                        }

                        t!("Thread {} exiting", tid);
                    }));
                } else {
                    work_tx.send(cert).unwrap();
                }
            }

            // When the threads see this drop, they will exit.
            drop(work_tx);

            // Drop our reference to results_tx.  When the last thread
            // exits, the last reference will be dropped and the loop
            // below will exit.
            drop(results_tx);

            while let Ok(cert) = results_rx.recv() {
                let fpr = cert.fingerprint();
                t!("Caching {}", fpr);
                self.certs.insert(fpr, Arc::new(cert.into()));
                count += 1;
            }
            }); // thread scope.

            // We're just caching results so we can ignore errors.
            if let Err(err) = result {
                t!("{:?}", err);
            }
        }

        t!("Prefetched {} certificates, ({} RawCerts had errors)",
           count, cert_count - count);
    }
}

impl<'a> StoreUpdate<'a> for Certs<'a> {
    fn update_by(&self, cert: Arc<LazyCert<'a>>,
                 merge_strategy: &dyn MergeCerts<'a>)
        -> Result<Arc<LazyCert<'a>>>
    {
        tracer!(TRACE, "Certs::update_by");

        let mut inner = self.inner.write().unwrap();

        let fpr = cert.fingerprint();

        // Add the cert fingerprint -> cert entry.
        let merged: Arc<LazyCert>;
        match inner.certs.entry(fpr.clone()) {
            Entry::Occupied(mut oe) => {
                t!("Updating {}", fpr);

                let old = oe.get().clone();

                merged = merge_strategy.merge_public(cert, Some(old))
                    .with_context(|| {
                        format!("Merging two version of {}", fpr)
                    })?;

                *oe.get_mut() = merged.clone();
            }
            Entry::Vacant(ve) => {
                t!("Inserting {}", fpr);

                merged = merge_strategy.merge_public(cert, None)?;
                ve.insert(merged.clone());
            }
        }

        // Populate the key map.  This is a merge so we are not
        // removing anything.
        for k in merged.keys() {
            match inner.keys.entry(k.keyid()) {
                Entry::Occupied(mut oe) => {
                    let fprs = oe.get_mut();
                    if ! fprs.contains(&fpr) {
                        fprs.push(fpr.clone());
                    }
                }
                Entry::Vacant(ve) => {
                    ve.insert(smallvec![ fpr.clone() ]);
                }
            }
        }

        inner.userid_index.insert(&fpr, merged.userids());

        Ok(merged)
    }
}