<commit_before>//! Global `Arc`-based object interning infrastructure.
//!
//! Eventually this should probably be replaced with salsa-based interning.

use std::{
    collections::HashMap,
    fmt::{self, Debug},
    hash::{BuildHasherDefault, Hash},
    ops::Deref,
    sync::Arc,
};

use dashmap::{lock::RwLockWriteGuard, DashMap, SharedValue};
use once_cell::sync::OnceCell;
use rustc_hash::FxHasher;

use crate::generics::GenericParams;

type InternMap<T> = DashMap<Arc<T>, (), BuildHasherDefault<FxHasher>>;
type Guard<T> =
    RwLockWriteGuard<'static, HashMap<Arc<T>, SharedValue<()>, BuildHasherDefault<FxHasher>>>;

#[derive(Hash)]
pub struct Interned<T: Internable + ?Sized> {
    arc: Arc<T>,
}

impl<T: Internable> Interned<T> {
    pub fn new(obj: T) -> Self {
        match Interned::lookup(&obj) {
            Ok(this) => this,
            Err(shard) => {
                let arc = Arc::new(obj);
                Self::alloc(arc, shard)
            }
        }
    }
}

impl<T: Internable + ?Sized> Interned<T> {
    fn lookup(obj: &T) -> Result<Self, Guard<T>> {
        let storage = T::storage().get();
        let shard_idx = storage.determine_map(obj);
        let shard = &storage.shards()[shard_idx];
        let shard = shard.write();

        // Atomically,
        // - check if `obj` is already in the map
        //   - if so, clone its `Arc` and return it
        //   - if not, box it up, insert it, and return a clone
        // This needs to be atomic (locking the shard) to avoid races with other thread, which could
        // insert the same object between us looking it up and inserting it.

        // FIXME: avoid double lookup/hashing by using raw entry API (once stable, or when
        // hashbrown can be plugged into dashmap)
        match shard.get_key_value(obj) {
            Some((arc, _)) => Ok(Self { arc: arc.clone() }),
            None => Err(shard),
        }
    }

    fn alloc(arc: Arc<T>, mut shard: Guard<T>) -> Self {
        let arc2 = arc.clone();

        shard.insert(arc2, SharedValue::new(()));

        Self { arc }
    }
}

impl Interned<str> {
    pub fn new_str(s: &str) -> Self {
        match Interned::lookup(s) {
            Ok(this) => this,
            Err(shard) => {
                let arc = Arc::<str>::from(s);
                Self::alloc(arc, shard)
            }
        }
    }
}

impl<T: Internable + ?Sized> Drop for Interned<T> {
    #[inline]
    fn drop(&mut self) {
        // When the last `Ref` is dropped, remove the object from the global map.
        if Arc::strong_count(&self.arc) == 2 {
            // Only `self` and the global map point to the object.

            self.drop_slow();
        }
    }
}

impl<T: Internable + ?Sized> Interned<T> {
    #[cold]
    fn drop_slow(&mut self) {
        let storage = T::storage().get();
        let shard_idx = storage.determine_map(&self.arc);
        let shard = &storage.shards()[shard_idx];
        let mut shard = shard.write();

        // FIXME: avoid double lookup
        let (arc, _) = shard.get_key_value(&self.arc).expect("interned value removed prematurely");

        if Arc::strong_count(arc) != 2 {
            // Another thread has interned another copy
            return;
        }

        shard.remove(&self.arc);

        // Shrink the backing storage if the shard is less than 50% occupied.
        if shard.len() * 2 < shard.capacity() {
            shard.shrink_to_fit();
        }
    }
}

/// Compares interned `Ref`s using pointer equality.
impl<T: Internable> PartialEq for Interned<T> {
    // NOTE: No `?Sized` because `ptr_eq` doesn't work right with trait objects.

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.arc, &other.arc)
    }
}

impl<T: Internable> Eq for Interned<T> {}

impl PartialEq for Interned<str> {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.arc, &other.arc)
    }
}

impl Eq for Interned<str> {}

impl<T: Internable + ?Sized> AsRef<T> for Interned<T> {
    #[inline]
    fn as_ref(&self) -> &T {
        &self.arc
    }
}

impl<T: Internable + ?Sized> Deref for Interned<T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.arc
    }
}

impl<T: Internable + ?Sized> Clone for Interned<T> {
    fn clone(&self) -> Self {
        Self { arc: self.arc.clone() }
    }
}

impl<T: Debug + Internable + ?Sized> Debug for Interned<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        (*self.arc).fmt(f)
    }
}

pub struct InternStorage<T: ?Sized> {
    map: OnceCell<InternMap<T>>,
}

impl<T: ?Sized> InternStorage<T> {
    pub const fn new() -> Self {
        Self { map: OnceCell::new() }
    }
}

impl<T: Internable + ?Sized> InternStorage<T> {
    fn get(&self) -> &InternMap<T> {
        self.map.get_or_init(DashMap::default)
    }
}

pub trait Internable: Hash + Eq + 'static {
    fn storage() -> &'static InternStorage<Self>;
}

/// Implements `Internable` for a given list of types, making them usable with `Interned`.
#[macro_export]
#[doc(hidden)]
macro_rules! _impl_internable {
    ( $($t:path),+ $(,)? ) => { $(
        impl Internable for $t {
            fn storage() -> &'static InternStorage<Self> {
                static STORAGE: InternStorage<$t> = InternStorage::new();
                &STORAGE
            }
        }
    )+ };
}

pub use crate::_impl_internable as impl_internable;

impl_internable!(
    crate::type_ref::TypeRef,
    crate::type_ref::TraitRef,
    crate::path::ModPath,
    GenericParams,
    str,
);
<commit_msg>`Interned<T>`: Only hash the pointer<commit_after>//! Global `Arc`-based object interning infrastructure.
//!
//! Eventually this should probably be replaced with salsa-based interning.

use std::{
    collections::HashMap,
    fmt::{self, Debug},
    hash::{BuildHasherDefault, Hash, Hasher},
    ops::Deref,
    sync::Arc,
};

use dashmap::{lock::RwLockWriteGuard, DashMap, SharedValue};
use once_cell::sync::OnceCell;
use rustc_hash::FxHasher;

use crate::generics::GenericParams;

type InternMap<T> = DashMap<Arc<T>, (), BuildHasherDefault<FxHasher>>;
type Guard<T> =
    RwLockWriteGuard<'static, HashMap<Arc<T>, SharedValue<()>, BuildHasherDefault<FxHasher>>>;

pub struct Interned<T: Internable + ?Sized> {
    arc: Arc<T>,
}

impl<T: Internable> Interned<T> {
    pub fn new(obj: T) -> Self {
        match Interned::lookup(&obj) {
            Ok(this) => this,
            Err(shard) => {
                let arc = Arc::new(obj);
                Self::alloc(arc, shard)
            }
        }
    }
}

impl<T: Internable + ?Sized> Interned<T> {
    fn lookup(obj: &T) -> Result<Self, Guard<T>> {
        let storage = T::storage().get();
        let shard_idx = storage.determine_map(obj);
        let shard = &storage.shards()[shard_idx];
        let shard = shard.write();

        // Atomically,
        // - check if `obj` is already in the map
        //   - if so, clone its `Arc` and return it
        //   - if not, box it up, insert it, and return a clone
        // This needs to be atomic (locking the shard) to avoid races with other thread, which could
        // insert the same object between us looking it up and inserting it.

        // FIXME: avoid double lookup/hashing by using raw entry API (once stable, or when
        // hashbrown can be plugged into dashmap)
        match shard.get_key_value(obj) {
            Some((arc, _)) => Ok(Self { arc: arc.clone() }),
            None => Err(shard),
        }
    }

    fn alloc(arc: Arc<T>, mut shard: Guard<T>) -> Self {
        let arc2 = arc.clone();

        shard.insert(arc2, SharedValue::new(()));

        Self { arc }
    }
}

impl Interned<str> {
    pub fn new_str(s: &str) -> Self {
        match Interned::lookup(s) {
            Ok(this) => this,
            Err(shard) => {
                let arc = Arc::<str>::from(s);
                Self::alloc(arc, shard)
            }
        }
    }
}

impl<T: Internable + ?Sized> Drop for Interned<T> {
    #[inline]
    fn drop(&mut self) {
        // When the last `Ref` is dropped, remove the object from the global map.
        if Arc::strong_count(&self.arc) == 2 {
            // Only `self` and the global map point to the object.

            self.drop_slow();
        }
    }
}

impl<T: Internable + ?Sized> Interned<T> {
    #[cold]
    fn drop_slow(&mut self) {
        let storage = T::storage().get();
        let shard_idx = storage.determine_map(&self.arc);
        let shard = &storage.shards()[shard_idx];
        let mut shard = shard.write();

        // FIXME: avoid double lookup
        let (arc, _) = shard.get_key_value(&self.arc).expect("interned value removed prematurely");

        if Arc::strong_count(arc) != 2 {
            // Another thread has interned another copy
            return;
        }

        shard.remove(&self.arc);

        // Shrink the backing storage if the shard is less than 50% occupied.
        if shard.len() * 2 < shard.capacity() {
            shard.shrink_to_fit();
        }
    }
}

/// Compares interned `Ref`s using pointer equality.
impl<T: Internable> PartialEq for Interned<T> {
    // NOTE: No `?Sized` because `ptr_eq` doesn't work right with trait objects.

    #[inline]
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.arc, &other.arc)
    }
}

impl<T: Internable> Eq for Interned<T> {}

impl PartialEq for Interned<str> {
    fn eq(&self, other: &Self) -> bool {
        Arc::ptr_eq(&self.arc, &other.arc)
    }
}

impl Eq for Interned<str> {}

impl<T: Internable + ?Sized> Hash for Interned<T> {
    fn hash<H: Hasher>(&self, state: &mut H) {
        // NOTE: Cast disposes vtable pointer / slice/str length.
        state.write_usize(Arc::as_ptr(&self.arc) as *const () as usize)
    }
}

impl<T: Internable + ?Sized> AsRef<T> for Interned<T> {
    #[inline]
    fn as_ref(&self) -> &T {
        &self.arc
    }
}

impl<T: Internable + ?Sized> Deref for Interned<T> {
    type Target = T;

    #[inline]
    fn deref(&self) -> &Self::Target {
        &self.arc
    }
}

impl<T: Internable + ?Sized> Clone for Interned<T> {
    fn clone(&self) -> Self {
        Self { arc: self.arc.clone() }
    }
}

impl<T: Debug + Internable + ?Sized> Debug for Interned<T> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        (*self.arc).fmt(f)
    }
}

pub struct InternStorage<T: ?Sized> {
    map: OnceCell<InternMap<T>>,
}

impl<T: ?Sized> InternStorage<T> {
    pub const fn new() -> Self {
        Self { map: OnceCell::new() }
    }
}

impl<T: Internable + ?Sized> InternStorage<T> {
    fn get(&self) -> &InternMap<T> {
        self.map.get_or_init(DashMap::default)
    }
}

pub trait Internable: Hash + Eq + 'static {
    fn storage() -> &'static InternStorage<Self>;
}

/// Implements `Internable` for a given list of types, making them usable with `Interned`.
#[macro_export]
#[doc(hidden)]
macro_rules! _impl_internable {
    ( $($t:path),+ $(,)? ) => { $(
        impl Internable for $t {
            fn storage() -> &'static InternStorage<Self> {
                static STORAGE: InternStorage<$t> = InternStorage::new();
                &STORAGE
            }
        }
    )+ };
}

pub use crate::_impl_internable as impl_internable;

impl_internable!(
    crate::type_ref::TypeRef,
    crate::type_ref::TraitRef,
    crate::path::ModPath,
    GenericParams,
    str,
);
<|endoftext|>


<commit_before>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::epoch_ending::manifest::EpochEndingBackup,
    metrics::{
        restore::{EPOCH_ENDING_EPOCH, EPOCH_ENDING_VERSION},
        verify::{VERIFY_EPOCH_ENDING_EPOCH, VERIFY_EPOCH_ENDING_VERSION},
    },
    storage::{BackupStorage, FileHandle, FileHandleRef},
    utils::{
        read_record_bytes::ReadRecordBytes, storage_ext::BackupStorageExt, stream::StreamX,
        GlobalRestoreOptions, RestoreRunMode,
    },
};
use anyhow::{anyhow, ensure, Result};
use aptos_logger::prelude::*;
use aptos_types::{
    epoch_change::Verifier,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    transaction::Version,
    waypoint::Waypoint,
};
use clap::Parser;
use futures::StreamExt;
use std::{collections::HashMap, sync::Arc, time::Instant};

#[derive(Parser)]
pub struct EpochEndingRestoreOpt {
    #[clap(long = "epoch-ending-manifest")]
    pub manifest_handle: FileHandle,
}

pub struct EpochEndingRestoreController {
    storage: Arc<dyn BackupStorage>,
    run_mode: Arc<RestoreRunMode>,
    manifest_handle: FileHandle,
    target_version: Version,
    trusted_waypoints: Arc<HashMap<Version, Waypoint>>,
}

impl EpochEndingRestoreController {
    pub fn new(
        opt: EpochEndingRestoreOpt,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            run_mode: global_opt.run_mode,
            manifest_handle: opt.manifest_handle,
            target_version: global_opt.target_version,
            trusted_waypoints: global_opt.trusted_waypoints,
        }
    }

    pub async fn preheat(self) -> PreheatedEpochEndingRestore {
        PreheatedEpochEndingRestore {
            preheat_result: self.preheat_impl().await,
            controller: self,
        }
    }

    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        self.preheat()
            .await
            .run(previous_epoch_ending_ledger_info)
            .await
    }
}

impl EpochEndingRestoreController {
    fn name(&self) -> String {
        format!("epoch ending {}", self.run_mode.name())
    }

    async fn preheat_impl(&self) -> Result<EpochEndingRestorePreheatData> {
        let manifest: EpochEndingBackup =
            self.storage.load_json_file(&self.manifest_handle).await?;
        manifest.verify()?;

        let mut next_epoch = manifest.first_epoch;
        let mut waypoint_iter = manifest.waypoints.iter();

        let mut previous_li: Option<&LedgerInfoWithSignatures> = None;
        let mut ledger_infos = Vec::new();

        let mut past_target = false;
        for chunk in &manifest.chunks {
            if past_target {
                break;
            }

            let lis = self.read_chunk(&chunk.ledger_infos).await?;
            ensure!(
                chunk.first_epoch + lis.len() as u64 == chunk.last_epoch + 1,
                "Number of items in chunks doesn't match that in manifest. \
                first_epoch: {}, last_epoch: {}, items in chunk: {}",
                chunk.first_epoch,
                chunk.last_epoch,
                lis.len(),
            );

            for li in lis {
                if li.ledger_info().version() > self.target_version {
                    past_target = true;
                    break;
                }

                ensure!(
                    li.ledger_info().epoch() == next_epoch,
                    "LedgerInfo epoch not expected. Expected: {}, actual: {}.",
                    li.ledger_info().epoch(),
                    next_epoch,
                );
                let wp_manifest = waypoint_iter.next().ok_or_else(|| {
                    anyhow!("More LedgerInfo's found than waypoints in manifest.")
                })?;
                let wp_li = Waypoint::new_epoch_boundary(li.ledger_info())?;
                ensure!(
                    *wp_manifest == wp_li,
                    "Waypoints don't match. In manifest: {}, In chunk: {}",
                    wp_manifest,
                    wp_li,
                );
                if let Some(wp_trusted) = self.trusted_waypoints.get(&wp_li.version()) {
                    ensure!(
                        *wp_trusted == wp_li,
                        "Waypoints don't match. In backup: {}, trusted: {}",
                        wp_li,
                        wp_trusted,
                    );
                } else if let Some(pre_li) = previous_li {
                    pre_li
                        .ledger_info()
                        .next_epoch_state()
                        .ok_or_else(|| {
                            anyhow!(
                                "Next epoch state not found from LI at epoch {}.",
                                pre_li.ledger_info().epoch()
                            )
                        })?
                        .verify(&li)?;
                }
                ledger_infos.push(li);
                previous_li = ledger_infos.last();
                next_epoch += 1;
            }
        }

        Ok(EpochEndingRestorePreheatData {
            manifest,
            ledger_infos,
        })
    }

    async fn read_chunk(
        &self,
        file_handle: &FileHandleRef,
    ) -> Result<Vec<LedgerInfoWithSignatures>> {
        let mut file = self.storage.open_for_read(file_handle).await?;
        let mut chunk = vec![];

        while let Some(record_bytes) = file.read_record_bytes().await? {
            chunk.push(bcs::from_bytes(&record_bytes)?);
        }

        Ok(chunk)
    }
}

struct EpochEndingRestorePreheatData {
    manifest: EpochEndingBackup,
    ledger_infos: Vec<LedgerInfoWithSignatures>,
}

pub struct PreheatedEpochEndingRestore {
    controller: EpochEndingRestoreController,
    preheat_result: Result<EpochEndingRestorePreheatData>,
}

impl PreheatedEpochEndingRestore {
    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let name = self.controller.name();
        info!(
            "{} started. Manifest: {}",
            name, self.controller.manifest_handle
        );
        let res = self
            .run_impl(previous_epoch_ending_ledger_info)
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl PreheatedEpochEndingRestore {
    async fn run_impl(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let preheat_data = self
            .preheat_result
            .map_err(|e| anyhow!("Preheat failed: {}", e))?;

        let first_li = preheat_data
            .ledger_infos
            .first()
            .expect("Epoch ending backup can't be empty.");

        if let Some(li) = previous_epoch_ending_ledger_info {
            ensure!(
                li.next_block_epoch() == preheat_data.manifest.first_epoch,
                "Previous epoch ending LedgerInfo is not the one expected. \
                My first epoch: {}, previous LedgerInfo next_block_epoch: {}",
                preheat_data.manifest.first_epoch,
                li.next_block_epoch(),
            );
            // Waypoint has been verified in preheat if it's trusted, otherwise try to check
            // the signatures.
            if self
                .controller
                .trusted_waypoints
                .get(&first_li.ledger_info().version())
                .is_none()
            {
                li.next_epoch_state()
                    .ok_or_else(|| {
                        anyhow!("Previous epoch ending LedgerInfo doesn't end an epoch")
                    })?
                    .verify(first_li)?;
            }
        }

        let last_li = preheat_data
            .ledger_infos
            .last()
            .expect("Verified not empty.")
            .ledger_info();
        match self.controller.run_mode.as_ref() {
            RestoreRunMode::Restore { restore_handler } => {
                restore_handler.save_ledger_infos(&preheat_data.ledger_infos)?;

                EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
            RestoreRunMode::Verify => {
                VERIFY_EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                VERIFY_EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
        };

        Ok(preheat_data
            .ledger_infos
            .into_iter()
            .map(|x| x.ledger_info().clone())
            .collect())
    }
}

/// Represents a history of epoch changes since epoch 0.
#[derive(Clone)]
pub struct EpochHistory {
    pub epoch_endings: Vec<LedgerInfo>,
}

impl EpochHistory {
    pub fn verify_ledger_info(&self, li_with_sigs: &LedgerInfoWithSignatures) -> Result<()> {
        let epoch = li_with_sigs.ledger_info().epoch();
        ensure!(!self.epoch_endings.is_empty(), "Empty epoch history.",);
        if epoch > self.epoch_endings.len() as u64 {
            // TODO(aldenhu): fix this from upper level
            warn!(
                epoch = epoch,
                epoch_history_until = self.epoch_endings.len(),
                "Epoch too new, can't verify. Don't worry, node won't be able to start if data is not compatible.",
            );
            return Ok(());
        }
        if epoch == 0 {
            ensure!(
                li_with_sigs.ledger_info() == &self.epoch_endings[0],
                "Genesis epoch LedgerInfo info doesn't match.",
            );
            Ok(())
        } else {
            self.epoch_endings[epoch as usize - 1]
                .next_epoch_state()
                .ok_or_else(|| anyhow!("Shouldn't contain non- epoch bumping LIs."))?
                .verify(li_with_sigs)
        }
    }
}

pub struct EpochHistoryRestoreController {
    storage: Arc<dyn BackupStorage>,
    manifest_handles: Vec<FileHandle>,
    global_opt: GlobalRestoreOptions,
}

impl EpochHistoryRestoreController {
    pub fn new(
        manifest_handles: Vec<FileHandle>,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            manifest_handles,
            global_opt,
        }
    }

    pub async fn run(self) -> Result<EpochHistory> {
        let name = self.name();
        info!("{} started.", name,);
        let res = self
            .run_impl()
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl EpochHistoryRestoreController {
    fn name(&self) -> String {
        format!("epoch history {}", self.global_opt.run_mode.name())
    }

    async fn run_impl(self) -> Result<EpochHistory> {
        let timer = Instant::now();
        if self.manifest_handles.is_empty() {
            return Ok(EpochHistory {
                epoch_endings: Vec::new(),
            });
        }

        let futs_iter = self.manifest_handles.iter().map(|hdl| {
            EpochEndingRestoreController::new(
                EpochEndingRestoreOpt {
                    manifest_handle: hdl.clone(),
                },
                self.global_opt.clone(),
                self.storage.clone(),
            )
            .preheat()
        });
        let mut futs_stream = futures::stream::iter(futs_iter).buffered_x(
            self.global_opt.concurrent_downloads * 2, /* buffer size */
            self.global_opt.concurrent_downloads,     /* concurrency */
        );

        let mut next_epoch = 0u64;
        let mut previous_li = None;
        let mut epoch_endings = Vec::new();

        while let Some(preheated_restore) = futs_stream.next().await {
            let manifest_handle = preheated_restore.controller.manifest_handle.clone();
            let lis = preheated_restore.run(previous_li).await?;
            ensure!(
                !lis.is_empty(),
                "No epochs restored from {}",
                manifest_handle,
            );
            for li in &lis {
                ensure!(
                    li.epoch() == next_epoch,
                    "Restored LedgerInfo has epoch {}, expecting {}.",
                    li.epoch(),
                    next_epoch,
                );
                ensure!(
                    li.ends_epoch(),
                    "LedgerInfo is not one at an epoch ending. epoch: {}",
                    li.epoch(),
                );
                next_epoch += 1;
            }

            epoch_endings.extend(lis);
            previous_li = epoch_endings.last();
        }

        info!(
            "Epoch history recovered in {:.2} seconds",
            timer.elapsed().as_secs_f64()
        );
        Ok(EpochHistory { epoch_endings })
    }
}
<commit_msg>fixup! [backup] ignore "History until epoch X can't verify" error<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use crate::{
    backup_types::epoch_ending::manifest::EpochEndingBackup,
    metrics::{
        restore::{EPOCH_ENDING_EPOCH, EPOCH_ENDING_VERSION},
        verify::{VERIFY_EPOCH_ENDING_EPOCH, VERIFY_EPOCH_ENDING_VERSION},
    },
    storage::{BackupStorage, FileHandle, FileHandleRef},
    utils::{
        read_record_bytes::ReadRecordBytes, storage_ext::BackupStorageExt, stream::StreamX,
        GlobalRestoreOptions, RestoreRunMode,
    },
};
use anyhow::{anyhow, ensure, Result};
use aptos_logger::prelude::*;
use aptos_types::{
    epoch_change::Verifier,
    ledger_info::{LedgerInfo, LedgerInfoWithSignatures},
    transaction::Version,
    waypoint::Waypoint,
};
use clap::Parser;
use futures::StreamExt;
use std::{collections::HashMap, sync::Arc, time::Instant};

#[derive(Parser)]
pub struct EpochEndingRestoreOpt {
    #[clap(long = "epoch-ending-manifest")]
    pub manifest_handle: FileHandle,
}

pub struct EpochEndingRestoreController {
    storage: Arc<dyn BackupStorage>,
    run_mode: Arc<RestoreRunMode>,
    manifest_handle: FileHandle,
    target_version: Version,
    trusted_waypoints: Arc<HashMap<Version, Waypoint>>,
}

impl EpochEndingRestoreController {
    pub fn new(
        opt: EpochEndingRestoreOpt,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            run_mode: global_opt.run_mode,
            manifest_handle: opt.manifest_handle,
            target_version: global_opt.target_version,
            trusted_waypoints: global_opt.trusted_waypoints,
        }
    }

    pub async fn preheat(self) -> PreheatedEpochEndingRestore {
        PreheatedEpochEndingRestore {
            preheat_result: self.preheat_impl().await,
            controller: self,
        }
    }

    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        self.preheat()
            .await
            .run(previous_epoch_ending_ledger_info)
            .await
    }
}

impl EpochEndingRestoreController {
    fn name(&self) -> String {
        format!("epoch ending {}", self.run_mode.name())
    }

    async fn preheat_impl(&self) -> Result<EpochEndingRestorePreheatData> {
        let manifest: EpochEndingBackup =
            self.storage.load_json_file(&self.manifest_handle).await?;
        manifest.verify()?;

        let mut next_epoch = manifest.first_epoch;
        let mut waypoint_iter = manifest.waypoints.iter();

        let mut previous_li: Option<&LedgerInfoWithSignatures> = None;
        let mut ledger_infos = Vec::new();

        let mut past_target = false;
        for chunk in &manifest.chunks {
            if past_target {
                break;
            }

            let lis = self.read_chunk(&chunk.ledger_infos).await?;
            ensure!(
                chunk.first_epoch + lis.len() as u64 == chunk.last_epoch + 1,
                "Number of items in chunks doesn't match that in manifest. \
                first_epoch: {}, last_epoch: {}, items in chunk: {}",
                chunk.first_epoch,
                chunk.last_epoch,
                lis.len(),
            );

            for li in lis {
                if li.ledger_info().version() > self.target_version {
                    past_target = true;
                    break;
                }

                ensure!(
                    li.ledger_info().epoch() == next_epoch,
                    "LedgerInfo epoch not expected. Expected: {}, actual: {}.",
                    li.ledger_info().epoch(),
                    next_epoch,
                );
                let wp_manifest = waypoint_iter.next().ok_or_else(|| {
                    anyhow!("More LedgerInfo's found than waypoints in manifest.")
                })?;
                let wp_li = Waypoint::new_epoch_boundary(li.ledger_info())?;
                ensure!(
                    *wp_manifest == wp_li,
                    "Waypoints don't match. In manifest: {}, In chunk: {}",
                    wp_manifest,
                    wp_li,
                );
                if let Some(wp_trusted) = self.trusted_waypoints.get(&wp_li.version()) {
                    ensure!(
                        *wp_trusted == wp_li,
                        "Waypoints don't match. In backup: {}, trusted: {}",
                        wp_li,
                        wp_trusted,
                    );
                } else if let Some(pre_li) = previous_li {
                    pre_li
                        .ledger_info()
                        .next_epoch_state()
                        .ok_or_else(|| {
                            anyhow!(
                                "Next epoch state not found from LI at epoch {}.",
                                pre_li.ledger_info().epoch()
                            )
                        })?
                        .verify(&li)?;
                }
                ledger_infos.push(li);
                previous_li = ledger_infos.last();
                next_epoch += 1;
            }
        }

        Ok(EpochEndingRestorePreheatData {
            manifest,
            ledger_infos,
        })
    }

    async fn read_chunk(
        &self,
        file_handle: &FileHandleRef,
    ) -> Result<Vec<LedgerInfoWithSignatures>> {
        let mut file = self.storage.open_for_read(file_handle).await?;
        let mut chunk = vec![];

        while let Some(record_bytes) = file.read_record_bytes().await? {
            chunk.push(bcs::from_bytes(&record_bytes)?);
        }

        Ok(chunk)
    }
}

struct EpochEndingRestorePreheatData {
    manifest: EpochEndingBackup,
    ledger_infos: Vec<LedgerInfoWithSignatures>,
}

pub struct PreheatedEpochEndingRestore {
    controller: EpochEndingRestoreController,
    preheat_result: Result<EpochEndingRestorePreheatData>,
}

impl PreheatedEpochEndingRestore {
    pub async fn run(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let name = self.controller.name();
        info!(
            "{} started. Manifest: {}",
            name, self.controller.manifest_handle
        );
        let res = self
            .run_impl(previous_epoch_ending_ledger_info)
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl PreheatedEpochEndingRestore {
    async fn run_impl(
        self,
        previous_epoch_ending_ledger_info: Option<&LedgerInfo>,
    ) -> Result<Vec<LedgerInfo>> {
        let preheat_data = self
            .preheat_result
            .map_err(|e| anyhow!("Preheat failed: {}", e))?;

        let first_li = preheat_data
            .ledger_infos
            .first()
            .expect("Epoch ending backup can't be empty.");

        if let Some(li) = previous_epoch_ending_ledger_info {
            ensure!(
                li.next_block_epoch() == preheat_data.manifest.first_epoch,
                "Previous epoch ending LedgerInfo is not the one expected. \
                My first epoch: {}, previous LedgerInfo next_block_epoch: {}",
                preheat_data.manifest.first_epoch,
                li.next_block_epoch(),
            );
            // Waypoint has been verified in preheat if it's trusted, otherwise try to check
            // the signatures.
            if self
                .controller
                .trusted_waypoints
                .get(&first_li.ledger_info().version())
                .is_none()
            {
                li.next_epoch_state()
                    .ok_or_else(|| {
                        anyhow!("Previous epoch ending LedgerInfo doesn't end an epoch")
                    })?
                    .verify(first_li)?;
            }
        }

        let last_li = preheat_data
            .ledger_infos
            .last()
            .expect("Verified not empty.")
            .ledger_info();
        match self.controller.run_mode.as_ref() {
            RestoreRunMode::Restore { restore_handler } => {
                restore_handler.save_ledger_infos(&preheat_data.ledger_infos)?;

                EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
            RestoreRunMode::Verify => {
                VERIFY_EPOCH_ENDING_EPOCH.set(last_li.epoch() as i64);
                VERIFY_EPOCH_ENDING_VERSION.set(last_li.version() as i64);
            }
        };

        Ok(preheat_data
            .ledger_infos
            .into_iter()
            .map(|x| x.ledger_info().clone())
            .collect())
    }
}

/// Represents a history of epoch changes since epoch 0.
#[derive(Clone)]
pub struct EpochHistory {
    pub epoch_endings: Vec<LedgerInfo>,
}

impl EpochHistory {
    pub fn verify_ledger_info(&self, li_with_sigs: &LedgerInfoWithSignatures) -> Result<()> {
        let epoch = li_with_sigs.ledger_info().epoch();
        ensure!(!self.epoch_endings.is_empty(), "Empty epoch history.",);
        if epoch > self.epoch_endings.len() as u64 {
            // TODO(aldenhu): fix this from upper level
            warn!(
                epoch = epoch,
                epoch_history_until = self.epoch_endings.len(),
                "Epoch too new, can't verify, letting it pass. Don't worry, previous chunks are \
                verified, and node won't be able to start if this piece of data is indeed malicious.",
            );
            return Ok(());
        }
        if epoch == 0 {
            ensure!(
                li_with_sigs.ledger_info() == &self.epoch_endings[0],
                "Genesis epoch LedgerInfo info doesn't match.",
            );
            Ok(())
        } else {
            self.epoch_endings[epoch as usize - 1]
                .next_epoch_state()
                .ok_or_else(|| anyhow!("Shouldn't contain non- epoch bumping LIs."))?
                .verify(li_with_sigs)
        }
    }
}

pub struct EpochHistoryRestoreController {
    storage: Arc<dyn BackupStorage>,
    manifest_handles: Vec<FileHandle>,
    global_opt: GlobalRestoreOptions,
}

impl EpochHistoryRestoreController {
    pub fn new(
        manifest_handles: Vec<FileHandle>,
        global_opt: GlobalRestoreOptions,
        storage: Arc<dyn BackupStorage>,
    ) -> Self {
        Self {
            storage,
            manifest_handles,
            global_opt,
        }
    }

    pub async fn run(self) -> Result<EpochHistory> {
        let name = self.name();
        info!("{} started.", name,);
        let res = self
            .run_impl()
            .await
            .map_err(|e| anyhow!("{} failed: {}", name, e))?;
        info!("{} succeeded.", name);
        Ok(res)
    }
}

impl EpochHistoryRestoreController {
    fn name(&self) -> String {
        format!("epoch history {}", self.global_opt.run_mode.name())
    }

    async fn run_impl(self) -> Result<EpochHistory> {
        let timer = Instant::now();
        if self.manifest_handles.is_empty() {
            return Ok(EpochHistory {
                epoch_endings: Vec::new(),
            });
        }

        let futs_iter = self.manifest_handles.iter().map(|hdl| {
            EpochEndingRestoreController::new(
                EpochEndingRestoreOpt {
                    manifest_handle: hdl.clone(),
                },
                self.global_opt.clone(),
                self.storage.clone(),
            )
            .preheat()
        });
        let mut futs_stream = futures::stream::iter(futs_iter).buffered_x(
            self.global_opt.concurrent_downloads * 2, /* buffer size */
            self.global_opt.concurrent_downloads,     /* concurrency */
        );

        let mut next_epoch = 0u64;
        let mut previous_li = None;
        let mut epoch_endings = Vec::new();

        while let Some(preheated_restore) = futs_stream.next().await {
            let manifest_handle = preheated_restore.controller.manifest_handle.clone();
            let lis = preheated_restore.run(previous_li).await?;
            ensure!(
                !lis.is_empty(),
                "No epochs restored from {}",
                manifest_handle,
            );
            for li in &lis {
                ensure!(
                    li.epoch() == next_epoch,
                    "Restored LedgerInfo has epoch {}, expecting {}.",
                    li.epoch(),
                    next_epoch,
                );
                ensure!(
                    li.ends_epoch(),
                    "LedgerInfo is not one at an epoch ending. epoch: {}",
                    li.epoch(),
                );
                next_epoch += 1;
            }

            epoch_endings.extend(lis);
            previous_li = epoch_endings.last();
        }

        info!(
            "Epoch history recovered in {:.2} seconds",
            timer.elapsed().as_secs_f64()
        );
        Ok(EpochHistory { epoch_endings })
    }
}
<|endoftext|>


<commit_before>// 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.

extern crate byteorder;
extern crate chrono;
#[macro_use]
extern crate clap;
extern crate hex;
extern crate ring;
extern crate roughenough;
extern crate time;

use ring::rand;
use ring::rand::SecureRandom;

use byteorder::{LittleEndian, ReadBytesExt};

use chrono::offset::Utc;
use chrono::TimeZone;

use std::fs::File;
use std::io::Write;
use std::collections::HashMap;
use std::iter::Iterator;
use std::net::{ToSocketAddrs, UdpSocket};

use clap::{App, Arg};
use roughenough::merkle::root_from_paths;
use roughenough::sign::Verifier;
use roughenough::{RtMessage, Tag, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT, VERSION};

fn create_nonce() -> [u8; 64] {
    let rng = rand::SystemRandom::new();
    let mut nonce = [0u8; 64];
    rng.fill(&mut nonce).unwrap();

    nonce
}

fn make_request(nonce: &[u8]) -> Vec<u8> {
    let mut msg = RtMessage::new(1);
    msg.add_field(Tag::NONC, nonce).unwrap();
    msg.pad_to_kilobyte();

    msg.encode().unwrap()
}

fn receive_response(sock: &mut UdpSocket) -> RtMessage {
    let mut buf = [0; 744];
    let resp_len = sock.recv_from(&mut buf).unwrap().0;

    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
}

struct ResponseHandler {
    pub_key: Option<Vec<u8>>,
    msg: HashMap<Tag, Vec<u8>>,
    srep: HashMap<Tag, Vec<u8>>,
    cert: HashMap<Tag, Vec<u8>>,
    dele: HashMap<Tag, Vec<u8>>,
    nonce: [u8; 64],
}

struct ParsedResponse {
    verified: bool,
    midpoint: u64,
    radius: u32,
}

impl ResponseHandler {
    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
        let msg = response.into_hash_map();
        let srep = RtMessage::from_bytes(&msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let cert = RtMessage::from_bytes(&msg[&Tag::CERT])
            .unwrap()
            .into_hash_map();
        let dele = RtMessage::from_bytes(&cert[&Tag::DELE])
            .unwrap()
            .into_hash_map();

        ResponseHandler {
            pub_key,
            msg,
            srep,
            cert,
            dele,
            nonce,
        }
    }

    pub fn extract_time(&self) -> ParsedResponse {
        let midpoint = self.srep[&Tag::MIDP]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let radius = self.srep[&Tag::RADI]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let mut verified = false;

        if self.pub_key.is_some() {
            self.validate_dele();
            self.validate_srep();
            self.validate_merkle();
            self.validate_midpoint(midpoint);
            verified = true;
        }

        ParsedResponse {
            verified,
            midpoint,
            radius,
        }
    }

    fn validate_dele(&self) {
        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
        full_cert.extend(&self.cert[&Tag::DELE]);

        assert!(
            self.validate_sig(
                self.pub_key.as_ref().unwrap(),
                &self.cert[&Tag::SIG],
                &full_cert
            ),
            "Invalid signature on DELE tag!"
        );
    }

    fn validate_srep(&self) {
        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
        full_srep.extend(&self.msg[&Tag::SREP]);

        assert!(
            self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
            "Invalid signature on SREP tag!"
        );
    }

    fn validate_merkle(&self) {
        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let index = self.msg[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let paths = &self.msg[&Tag::PATH];

        let hash = root_from_paths(index as usize, &self.nonce, paths);

        assert_eq!(
            Vec::from(hash),
            srep[&Tag::ROOT],
            "Nonce not in merkle tree!"
        );
    }

    fn validate_midpoint(&self, midpoint: u64) {
        let mint = self.dele[&Tag::MINT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let maxt = self.dele[&Tag::MAXT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();

        assert!(
            midpoint >= mint,
            "Response midpoint {} lies before delegation span ({}, {})"
        );
        assert!(
            midpoint <= maxt,
            "Response midpoint {} lies after delegation span ({}, {})"
        );
    }

    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
        let mut verifier = Verifier::new(public_key);
        verifier.update(data);
        verifier.verify(sig)
    }
}

fn main() {
    let matches = App::new("roughenough client")
    .version(VERSION)
    .arg(Arg::with_name("host")
      .required(true)
      .help("The Roughtime server to connect to")
      .takes_value(true))
    .arg(Arg::with_name("port")
      .required(true)
      .help("The Roughtime server port to connect to")
      .takes_value(true))
    .arg(Arg::with_name("public-key")
      .short("p")
      .long("public-key")
      .takes_value(true)
      .help("The server public key used to validate responses. If unset, no validation will be performed"))
    .arg(Arg::with_name("time-format")
      .short("f")
      .long("time-format")
      .takes_value(true)
      .help("The strftime format string used to print the time recieved from the server")
      .default_value("%b %d %Y %H:%M:%S")
    )
    .arg(Arg::with_name("num-requests")
      .short("n")
      .long("num-requests")
      .takes_value(true)
      .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling")
      .default_value("1")
    )
    .arg(Arg::with_name("stress")
      .short("s")
      .long("stress")
      .help("Stress-tests the server by sending the same request as fast as possible. Please only use this on your own server")
    )
    .get_matches();

    let host = matches.value_of("host").unwrap();
    let port = value_t_or_exit!(matches.value_of("port"), u16);
    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
    let time_format = matches.value_of("time-format").unwrap();
    let stress = matches.is_present("stress");
    let pub_key = matches
        .value_of("public-key")
        .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));

    println!("Requesting time from: {:?}:{:?}", host, port);

    let addr = (host, port).to_socket_addrs().unwrap().next().unwrap();

    if stress {
        if !addr.ip().is_loopback() {
            println!(
                "ERROR: Cannot use non-loopback address {} for stress testing",
                addr.ip()
            );
            return;
        }

        println!("Stress-testing!");

        let nonce = create_nonce();
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);

        loop {
            socket.send_to(&request, addr).unwrap();
        }
    }

    let mut requests = Vec::with_capacity(num_requests);

    let mut file = File::create("requests.bin").unwrap();

    for _ in 0..num_requests {
        let nonce = create_nonce();
        let mut socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);
        file.write_all(&request);

        requests.push((nonce, request, socket));
    }

    for &mut (_, ref request, ref mut socket) in requests.iter_mut() {
        socket.send_to(request, addr).unwrap();
    }

    for (nonce, _, mut socket) in requests {
        let resp = receive_response(&mut socket);

        let ParsedResponse {
            verified,
            midpoint,
            radius,
        } = ResponseHandler::new(pub_key.clone(), resp.clone(), nonce).extract_time();

        let map = resp.into_hash_map();
        let index = map[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let seconds = midpoint / 10_u64.pow(6);
        let nsecs = (midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3);
        let spec = Utc.timestamp(seconds as i64, nsecs as u32);
        let out = spec.format(time_format).to_string();

        println!(
            "Received time from server: midpoint={:?}, radius={:?} (merkle_index={}, verified={})",
            out, radius, index, verified
        );
    }
}
<commit_msg>Add "output" argument to client<commit_after>// 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.

extern crate byteorder;
extern crate chrono;
#[macro_use]
extern crate clap;
extern crate hex;
extern crate ring;
extern crate roughenough;
extern crate time;

use ring::rand;
use ring::rand::SecureRandom;

use byteorder::{LittleEndian, ReadBytesExt};

use chrono::offset::Utc;
use chrono::TimeZone;

use std::fs::File;
use std::io::Write;
use std::collections::HashMap;
use std::iter::Iterator;
use std::net::{ToSocketAddrs, UdpSocket};

use clap::{App, Arg};
use roughenough::merkle::root_from_paths;
use roughenough::sign::Verifier;
use roughenough::{RtMessage, Tag, CERTIFICATE_CONTEXT, SIGNED_RESPONSE_CONTEXT, VERSION};

fn create_nonce() -> [u8; 64] {
    let rng = rand::SystemRandom::new();
    let mut nonce = [0u8; 64];
    rng.fill(&mut nonce).unwrap();

    nonce
}

fn make_request(nonce: &[u8]) -> Vec<u8> {
    let mut msg = RtMessage::new(1);
    msg.add_field(Tag::NONC, nonce).unwrap();
    msg.pad_to_kilobyte();

    msg.encode().unwrap()
}

fn receive_response(sock: &mut UdpSocket) -> RtMessage {
    let mut buf = [0; 744];
    let resp_len = sock.recv_from(&mut buf).unwrap().0;

    RtMessage::from_bytes(&buf[0..resp_len]).unwrap()
}

struct ResponseHandler {
    pub_key: Option<Vec<u8>>,
    msg: HashMap<Tag, Vec<u8>>,
    srep: HashMap<Tag, Vec<u8>>,
    cert: HashMap<Tag, Vec<u8>>,
    dele: HashMap<Tag, Vec<u8>>,
    nonce: [u8; 64],
}

struct ParsedResponse {
    verified: bool,
    midpoint: u64,
    radius: u32,
}

impl ResponseHandler {
    pub fn new(pub_key: Option<Vec<u8>>, response: RtMessage, nonce: [u8; 64]) -> ResponseHandler {
        let msg = response.into_hash_map();
        let srep = RtMessage::from_bytes(&msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let cert = RtMessage::from_bytes(&msg[&Tag::CERT])
            .unwrap()
            .into_hash_map();
        let dele = RtMessage::from_bytes(&cert[&Tag::DELE])
            .unwrap()
            .into_hash_map();

        ResponseHandler {
            pub_key,
            msg,
            srep,
            cert,
            dele,
            nonce,
        }
    }

    pub fn extract_time(&self) -> ParsedResponse {
        let midpoint = self.srep[&Tag::MIDP]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let radius = self.srep[&Tag::RADI]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let mut verified = false;

        if self.pub_key.is_some() {
            self.validate_dele();
            self.validate_srep();
            self.validate_merkle();
            self.validate_midpoint(midpoint);
            verified = true;
        }

        ParsedResponse {
            verified,
            midpoint,
            radius,
        }
    }

    fn validate_dele(&self) {
        let mut full_cert = Vec::from(CERTIFICATE_CONTEXT.as_bytes());
        full_cert.extend(&self.cert[&Tag::DELE]);

        assert!(
            self.validate_sig(
                self.pub_key.as_ref().unwrap(),
                &self.cert[&Tag::SIG],
                &full_cert
            ),
            "Invalid signature on DELE tag!"
        );
    }

    fn validate_srep(&self) {
        let mut full_srep = Vec::from(SIGNED_RESPONSE_CONTEXT.as_bytes());
        full_srep.extend(&self.msg[&Tag::SREP]);

        assert!(
            self.validate_sig(&self.dele[&Tag::PUBK], &self.msg[&Tag::SIG], &full_srep),
            "Invalid signature on SREP tag!"
        );
    }

    fn validate_merkle(&self) {
        let srep = RtMessage::from_bytes(&self.msg[&Tag::SREP])
            .unwrap()
            .into_hash_map();
        let index = self.msg[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();
        let paths = &self.msg[&Tag::PATH];

        let hash = root_from_paths(index as usize, &self.nonce, paths);

        assert_eq!(
            Vec::from(hash),
            srep[&Tag::ROOT],
            "Nonce not in merkle tree!"
        );
    }

    fn validate_midpoint(&self, midpoint: u64) {
        let mint = self.dele[&Tag::MINT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();
        let maxt = self.dele[&Tag::MAXT]
            .as_slice()
            .read_u64::<LittleEndian>()
            .unwrap();

        assert!(
            midpoint >= mint,
            "Response midpoint {} lies before delegation span ({}, {})"
        );
        assert!(
            midpoint <= maxt,
            "Response midpoint {} lies after delegation span ({}, {})"
        );
    }

    fn validate_sig(&self, public_key: &[u8], sig: &[u8], data: &[u8]) -> bool {
        let mut verifier = Verifier::new(public_key);
        verifier.update(data);
        verifier.verify(sig)
    }
}

fn main() {
    let matches = App::new("roughenough client")
    .version(VERSION)
    .arg(Arg::with_name("host")
      .required(true)
      .help("The Roughtime server to connect to")
      .takes_value(true))
    .arg(Arg::with_name("port")
      .required(true)
      .help("The Roughtime server port to connect to")
      .takes_value(true))
    .arg(Arg::with_name("public-key")
      .short("p")
      .long("public-key")
      .takes_value(true)
      .help("The server public key used to validate responses. If unset, no validation will be performed"))
    .arg(Arg::with_name("time-format")
      .short("f")
      .long("time-format")
      .takes_value(true)
      .help("The strftime format string used to print the time recieved from the server")
      .default_value("%b %d %Y %H:%M:%S")
    )
    .arg(Arg::with_name("num-requests")
      .short("n")
      .long("num-requests")
      .takes_value(true)
      .help("The number of requests to make to the server (each from a different source port). This is mainly useful for testing batch response handling")
      .default_value("1")
    )
    .arg(Arg::with_name("stress")
      .short("s")
      .long("stress")
      .help("Stress-tests the server by sending the same request as fast as possible. Please only use this on your own server")
    )
    .arg(Arg::with_name("output")
      .short("o")
      .long("output")
      .takes_value(true)
      .help("Writes all requsts to the specified file, in addition to sending them to the server. Useful for generating fuzer inputs")
    )
    .get_matches();

    let host = matches.value_of("host").unwrap();
    let port = value_t_or_exit!(matches.value_of("port"), u16);
    let num_requests = value_t_or_exit!(matches.value_of("num-requests"), u16) as usize;
    let time_format = matches.value_of("time-format").unwrap();
    let stress = matches.is_present("stress");
    let pub_key = matches
        .value_of("public-key")
        .map(|pkey| hex::decode(pkey).expect("Error parsing public key!"));
    let out = matches.value_of("output");

    println!("Requesting time from: {:?}:{:?}", host, port);

    let addr = (host, port).to_socket_addrs().unwrap().next().unwrap();

    if stress {
        if !addr.ip().is_loopback() {
            println!(
                "ERROR: Cannot use non-loopback address {} for stress testing",
                addr.ip()
            );
            return;
        }

        println!("Stress-testing!");

        let nonce = create_nonce();
        let socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);

        loop {
            socket.send_to(&request, addr).unwrap();
        }
    }

    let mut requests = Vec::with_capacity(num_requests);
    let mut file = out.map(|o | File::create(o).expect("Failed to create file!"));

    for _ in 0..num_requests {
        let nonce = create_nonce();
        let mut socket = UdpSocket::bind("0.0.0.0:0").expect("Couldn't open UDP socket");
        let request = make_request(&nonce);
        file.as_mut().map(|f| f.write_all(&request).expect("Failed to write to file!"));

        requests.push((nonce, request, socket));
    }

    for &mut (_, ref request, ref mut socket) in requests.iter_mut() {
        socket.send_to(request, addr).unwrap();
    }

    for (nonce, _, mut socket) in requests {
        let resp = receive_response(&mut socket);

        let ParsedResponse {
            verified,
            midpoint,
            radius,
        } = ResponseHandler::new(pub_key.clone(), resp.clone(), nonce).extract_time();

        let map = resp.into_hash_map();
        let index = map[&Tag::INDX]
            .as_slice()
            .read_u32::<LittleEndian>()
            .unwrap();

        let seconds = midpoint / 10_u64.pow(6);
        let nsecs = (midpoint - (seconds * 10_u64.pow(6))) * 10_u64.pow(3);
        let spec = Utc.timestamp(seconds as i64, nsecs as u32);
        let out = spec.format(time_format).to_string();

        println!(
            "Received time from server: midpoint={:?}, radius={:?} (merkle_index={}, verified={})",
            out, radius, index, verified
        );
    }
}
<|endoftext|>


<commit_before>#![cfg(target_os = "macos")]
use crate::{
    ContextError, CreationError, GlAttributes, PixelFormat,
    PixelFormatRequirements, Robustness, Rect,
};

use cgl::{
    kCGLCECrashOnRemovedFunctions, kCGLCPSurfaceOpacity, CGLEnable,
    CGLSetParameter,
};
use cocoa::appkit::{self, NSOpenGLContext, NSOpenGLPixelFormat};
use cocoa::base::{id, nil};
use cocoa::foundation::NSAutoreleasePool;
use core_foundation::base::TCFType;
use core_foundation::bundle::{
    CFBundleGetBundleWithIdentifier, CFBundleGetFunctionPointerForName,
};
use core_foundation::string::CFString;
use objc::runtime::{BOOL, NO};

use crate::platform::macos::WindowExtMacOS;
use winit;
use winit::dpi;
use winit::event_loop::EventLoopWindowTarget;
use winit::window::{Window, WindowBuilder};

use std::ops::Deref;
use std::os::raw;
use std::str::FromStr;

mod helpers;

#[derive(Debug)]
pub enum Context {
    WindowedContext(WindowedContext),
    HeadlessContext(HeadlessContext),
}

#[derive(Debug)]
pub struct WindowedContext {
    // NSOpenGLContext
    context: IdRef,
    pixel_format: PixelFormat,
}

#[derive(Debug)]
pub struct HeadlessContext {
    context: IdRef,
}

impl Context {
    #[inline]
    pub fn new_windowed<T>(
        wb: WindowBuilder,
        el: &EventLoopWindowTarget<T>,
        pf_reqs: &PixelFormatRequirements,
        gl_attr: &GlAttributes<&Context>,
    ) -> Result<(Window, Self), CreationError> {
        let transparent = wb.window.transparent;
        let win = wb.build(el)?;

        let share_ctx = gl_attr.sharing.map_or(nil, |c| *c.get_id());

        match gl_attr.robustness {
            Robustness::RobustNoResetNotification
            | Robustness::RobustLoseContextOnReset => {
                return Err(CreationError::RobustnessNotSupported);
            }
            _ => (),
        }

        let view = win.ns_view() as id;

        let gl_profile = helpers::get_gl_profile(gl_attr, pf_reqs)?;
        let attributes = helpers::build_nsattributes(pf_reqs, gl_profile)?;
        unsafe {
            let pixel_format = IdRef::new(
                NSOpenGLPixelFormat::alloc(nil)
                    .initWithAttributes_(&attributes),
            );
            let pixel_format = match pixel_format.non_nil() {
                None => return Err(CreationError::NoAvailablePixelFormat),
                Some(pf) => pf,
            };

            let gl_context = IdRef::new(
                NSOpenGLContext::alloc(nil)
                    .initWithFormat_shareContext_(*pixel_format, share_ctx),
            );
            let gl_context = match gl_context.non_nil() {
                Some(gl_context) => gl_context,
                None => {
                    return Err(CreationError::NotSupported(
                        "could not open gl context".to_string(),
                    ));
                }
            };

            let pixel_format = {
                let get_attr =
                    |attrib: appkit::NSOpenGLPixelFormatAttribute| -> i32 {
                        let mut value = 0;
                        NSOpenGLPixelFormat::getValues_forAttribute_forVirtualScreen_(
                        *pixel_format,
                        &mut value,
                        attrib,
                        NSOpenGLContext::currentVirtualScreen(*gl_context),
                    );
                        value
                    };

                PixelFormat {
                    hardware_accelerated: get_attr(
                        appkit::NSOpenGLPFAAccelerated,
                    ) != 0,
                    color_bits: (get_attr(appkit::NSOpenGLPFAColorSize)
                        - get_attr(appkit::NSOpenGLPFAAlphaSize))
                        as u8,
                    alpha_bits: get_attr(appkit::NSOpenGLPFAAlphaSize) as u8,
                    depth_bits: get_attr(appkit::NSOpenGLPFADepthSize) as u8,
                    stencil_bits: get_attr(appkit::NSOpenGLPFAStencilSize)
                        as u8,
                    stereoscopy: get_attr(appkit::NSOpenGLPFAStereo) != 0,
                    double_buffer: get_attr(appkit::NSOpenGLPFADoubleBuffer)
                        != 0,
                    multisampling: if get_attr(appkit::NSOpenGLPFAMultisample)
                        > 0
                    {
                        Some(get_attr(appkit::NSOpenGLPFASamples) as u16)
                    } else {
                        None
                    },
                    srgb: true,
                }
            };

            gl_context.setView_(view);
            let value = if gl_attr.vsync { 1 } else { 0 };
            gl_context.setValues_forParameter_(
                &value,
                appkit::NSOpenGLContextParameter::NSOpenGLCPSwapInterval,
            );

            if transparent {
                let mut opacity = 0;
                CGLSetParameter(
                    gl_context.CGLContextObj() as *mut _,
                    kCGLCPSurfaceOpacity,
                    &mut opacity,
                );
            }

            CGLEnable(
                gl_context.CGLContextObj() as *mut _,
                kCGLCECrashOnRemovedFunctions,
            );

            let context = WindowedContext {
                context: gl_context,
                pixel_format: pixel_format,
            };
            Ok((win, Context::WindowedContext(context)))
        }
    }

    #[inline]
    pub fn new_headless<T>(
        _el: &EventLoopWindowTarget<T>,
        pf_reqs: &PixelFormatRequirements,
        gl_attr: &GlAttributes<&Context>,
        _size: dpi::PhysicalSize,
    ) -> Result<Self, CreationError> {
        let gl_profile = helpers::get_gl_profile(gl_attr, pf_reqs)?;
        let attributes = helpers::build_nsattributes(pf_reqs, gl_profile)?;
        let context = unsafe {
            let pixelformat = NSOpenGLPixelFormat::alloc(nil)
                .initWithAttributes_(&attributes);
            if pixelformat == nil {
                return Err(CreationError::OsError(
                    "Could not create the pixel format".to_string(),
                ));
            }
            let context = NSOpenGLContext::alloc(nil)
                .initWithFormat_shareContext_(pixelformat, nil);
            if context == nil {
                return Err(CreationError::OsError(
                    "Could not create the rendering context".to_string(),
                ));
            }

            IdRef::new(context)
        };

        let headless = HeadlessContext { context };

        Ok(Context::HeadlessContext(headless))
    }

    pub fn resize(&self, _width: u32, _height: u32) {
        match *self {
            Context::WindowedContext(ref c) => unsafe { c.context.update() },
            _ => unreachable!(),
        }
    }

    #[inline]
    pub unsafe fn make_current(&self) -> Result<(), ContextError> {
        match *self {
            Context::WindowedContext(ref c) => {
                let _: () = msg_send![*c.context, update];
                c.context.makeCurrentContext();
            }
            Context::HeadlessContext(ref c) => {
                let _: () = msg_send![*c.context, update];
                c.context.makeCurrentContext();
            }
        }
        Ok(())
    }

    #[inline]
    pub unsafe fn make_not_current(&self) -> Result<(), ContextError> {
        if self.is_current() {
            match *self {
                Context::WindowedContext(ref c) => {
                    let _: () = msg_send![*c.context, update];
                    NSOpenGLContext::clearCurrentContext(nil);
                }
                Context::HeadlessContext(ref c) => {
                    let _: () = msg_send![*c.context, update];
                    NSOpenGLContext::clearCurrentContext(nil);
                }
            }
        }
        Ok(())
    }

    #[inline]
    pub fn is_current(&self) -> bool {
        unsafe {
            let context = match *self {
                Context::WindowedContext(ref c) => *c.context,
                Context::HeadlessContext(ref c) => *c.context,
            };

            let pool = NSAutoreleasePool::new(nil);
            let current = NSOpenGLContext::currentContext(nil);
            let res = if current != nil {
                let is_equal: BOOL = msg_send![current, isEqual: context];
                is_equal != NO
            } else {
                false
            };
            let _: () = msg_send![pool, release];
            res
        }
    }

    pub fn get_proc_address(&self, addr: &str) -> *const core::ffi::c_void {
        let symbol_name: CFString = FromStr::from_str(addr).unwrap();
        let framework_name: CFString =
            FromStr::from_str("com.apple.opengl").unwrap();
        let framework = unsafe {
            CFBundleGetBundleWithIdentifier(
                framework_name.as_concrete_TypeRef(),
            )
        };
        let symbol = unsafe {
            CFBundleGetFunctionPointerForName(
                framework,
                symbol_name.as_concrete_TypeRef(),
            )
        };
        symbol as *const _
    }

    #[inline]
    pub fn swap_buffers(&self) -> Result<(), ContextError> {
        unsafe {
            match *self {
                Context::WindowedContext(ref c) => {
                    let pool = NSAutoreleasePool::new(nil);
                    c.context.flushBuffer();
                    let _: () = msg_send![pool, release];
                }
                Context::HeadlessContext(_) => unreachable!(),
            }
        }
        Ok(())
    }

    #[inline]
    pub fn swap_buffers_with_damage(
        &self,
        rects: &[Rect],
    ) -> Result<(), ContextError> {
        Err(ContextError::OsError("buffer damage not suported".to_string()))
    }

    #[inline]
    pub fn swap_buffers_with_damage_supported(&self) -> bool {
        false
    }

    #[inline]
    pub fn get_api(&self) -> crate::Api {
        crate::Api::OpenGl
    }

    #[inline]
    pub fn get_pixel_format(&self) -> PixelFormat {
        match *self {
            Context::WindowedContext(ref c) => c.pixel_format.clone(),
            Context::HeadlessContext(_) => unreachable!(),
        }
    }

    #[inline]
    pub unsafe fn raw_handle(&self) -> *mut raw::c_void {
        match self {
            Context::WindowedContext(c) => {
                c.context.deref().CGLContextObj() as *mut _
            }
            Context::HeadlessContext(c) => {
                c.context.deref().CGLContextObj() as *mut _
            }
        }
    }

    #[inline]
    fn get_id(&self) -> IdRef {
        match self {
            Context::WindowedContext(w) => w.context.clone(),
            Context::HeadlessContext(h) => h.context.clone(),
        }
    }
}

#[derive(Debug)]
struct IdRef(id);

impl IdRef {
    fn new(i: id) -> IdRef {
        IdRef(i)
    }

    #[allow(dead_code)]
    fn retain(i: id) -> IdRef {
        if i != nil {
            let _: id = unsafe { msg_send![i, retain] };
        }
        IdRef(i)
    }

    fn non_nil(self) -> Option<IdRef> {
        if self.0 == nil {
            None
        } else {
            Some(self)
        }
    }
}

impl Drop for IdRef {
    fn drop(&mut self) {
        if self.0 != nil {
            let _: () = unsafe { msg_send![self.0, release] };
        }
    }
}

impl Deref for IdRef {
    type Target = id;
    fn deref<'a>(&'a self) -> &'a id {
        &self.0
    }
}

impl Clone for IdRef {
    fn clone(&self) -> IdRef {
        if self.0 != nil {
            let _: id = unsafe { msg_send![self.0, retain] };
        }
        IdRef(self.0)
    }
}

unsafe impl Send for Context {}
unsafe impl Sync for Context {}
<commit_msg>macos: Suppress warning about unused argument. (#1228)<commit_after>#![cfg(target_os = "macos")]
use crate::{
    ContextError, CreationError, GlAttributes, PixelFormat,
    PixelFormatRequirements, Robustness, Rect,
};

use cgl::{
    kCGLCECrashOnRemovedFunctions, kCGLCPSurfaceOpacity, CGLEnable,
    CGLSetParameter,
};
use cocoa::appkit::{self, NSOpenGLContext, NSOpenGLPixelFormat};
use cocoa::base::{id, nil};
use cocoa::foundation::NSAutoreleasePool;
use core_foundation::base::TCFType;
use core_foundation::bundle::{
    CFBundleGetBundleWithIdentifier, CFBundleGetFunctionPointerForName,
};
use core_foundation::string::CFString;
use objc::runtime::{BOOL, NO};

use crate::platform::macos::WindowExtMacOS;
use winit;
use winit::dpi;
use winit::event_loop::EventLoopWindowTarget;
use winit::window::{Window, WindowBuilder};

use std::ops::Deref;
use std::os::raw;
use std::str::FromStr;

mod helpers;

#[derive(Debug)]
pub enum Context {
    WindowedContext(WindowedContext),
    HeadlessContext(HeadlessContext),
}

#[derive(Debug)]
pub struct WindowedContext {
    // NSOpenGLContext
    context: IdRef,
    pixel_format: PixelFormat,
}

#[derive(Debug)]
pub struct HeadlessContext {
    context: IdRef,
}

impl Context {
    #[inline]
    pub fn new_windowed<T>(
        wb: WindowBuilder,
        el: &EventLoopWindowTarget<T>,
        pf_reqs: &PixelFormatRequirements,
        gl_attr: &GlAttributes<&Context>,
    ) -> Result<(Window, Self), CreationError> {
        let transparent = wb.window.transparent;
        let win = wb.build(el)?;

        let share_ctx = gl_attr.sharing.map_or(nil, |c| *c.get_id());

        match gl_attr.robustness {
            Robustness::RobustNoResetNotification
            | Robustness::RobustLoseContextOnReset => {
                return Err(CreationError::RobustnessNotSupported);
            }
            _ => (),
        }

        let view = win.ns_view() as id;

        let gl_profile = helpers::get_gl_profile(gl_attr, pf_reqs)?;
        let attributes = helpers::build_nsattributes(pf_reqs, gl_profile)?;
        unsafe {
            let pixel_format = IdRef::new(
                NSOpenGLPixelFormat::alloc(nil)
                    .initWithAttributes_(&attributes),
            );
            let pixel_format = match pixel_format.non_nil() {
                None => return Err(CreationError::NoAvailablePixelFormat),
                Some(pf) => pf,
            };

            let gl_context = IdRef::new(
                NSOpenGLContext::alloc(nil)
                    .initWithFormat_shareContext_(*pixel_format, share_ctx),
            );
            let gl_context = match gl_context.non_nil() {
                Some(gl_context) => gl_context,
                None => {
                    return Err(CreationError::NotSupported(
                        "could not open gl context".to_string(),
                    ));
                }
            };

            let pixel_format = {
                let get_attr =
                    |attrib: appkit::NSOpenGLPixelFormatAttribute| -> i32 {
                        let mut value = 0;
                        NSOpenGLPixelFormat::getValues_forAttribute_forVirtualScreen_(
                        *pixel_format,
                        &mut value,
                        attrib,
                        NSOpenGLContext::currentVirtualScreen(*gl_context),
                    );
                        value
                    };

                PixelFormat {
                    hardware_accelerated: get_attr(
                        appkit::NSOpenGLPFAAccelerated,
                    ) != 0,
                    color_bits: (get_attr(appkit::NSOpenGLPFAColorSize)
                        - get_attr(appkit::NSOpenGLPFAAlphaSize))
                        as u8,
                    alpha_bits: get_attr(appkit::NSOpenGLPFAAlphaSize) as u8,
                    depth_bits: get_attr(appkit::NSOpenGLPFADepthSize) as u8,
                    stencil_bits: get_attr(appkit::NSOpenGLPFAStencilSize)
                        as u8,
                    stereoscopy: get_attr(appkit::NSOpenGLPFAStereo) != 0,
                    double_buffer: get_attr(appkit::NSOpenGLPFADoubleBuffer)
                        != 0,
                    multisampling: if get_attr(appkit::NSOpenGLPFAMultisample)
                        > 0
                    {
                        Some(get_attr(appkit::NSOpenGLPFASamples) as u16)
                    } else {
                        None
                    },
                    srgb: true,
                }
            };

            gl_context.setView_(view);
            let value = if gl_attr.vsync { 1 } else { 0 };
            gl_context.setValues_forParameter_(
                &value,
                appkit::NSOpenGLContextParameter::NSOpenGLCPSwapInterval,
            );

            if transparent {
                let mut opacity = 0;
                CGLSetParameter(
                    gl_context.CGLContextObj() as *mut _,
                    kCGLCPSurfaceOpacity,
                    &mut opacity,
                );
            }

            CGLEnable(
                gl_context.CGLContextObj() as *mut _,
                kCGLCECrashOnRemovedFunctions,
            );

            let context = WindowedContext {
                context: gl_context,
                pixel_format: pixel_format,
            };
            Ok((win, Context::WindowedContext(context)))
        }
    }

    #[inline]
    pub fn new_headless<T>(
        _el: &EventLoopWindowTarget<T>,
        pf_reqs: &PixelFormatRequirements,
        gl_attr: &GlAttributes<&Context>,
        _size: dpi::PhysicalSize,
    ) -> Result<Self, CreationError> {
        let gl_profile = helpers::get_gl_profile(gl_attr, pf_reqs)?;
        let attributes = helpers::build_nsattributes(pf_reqs, gl_profile)?;
        let context = unsafe {
            let pixelformat = NSOpenGLPixelFormat::alloc(nil)
                .initWithAttributes_(&attributes);
            if pixelformat == nil {
                return Err(CreationError::OsError(
                    "Could not create the pixel format".to_string(),
                ));
            }
            let context = NSOpenGLContext::alloc(nil)
                .initWithFormat_shareContext_(pixelformat, nil);
            if context == nil {
                return Err(CreationError::OsError(
                    "Could not create the rendering context".to_string(),
                ));
            }

            IdRef::new(context)
        };

        let headless = HeadlessContext { context };

        Ok(Context::HeadlessContext(headless))
    }

    pub fn resize(&self, _width: u32, _height: u32) {
        match *self {
            Context::WindowedContext(ref c) => unsafe { c.context.update() },
            _ => unreachable!(),
        }
    }

    #[inline]
    pub unsafe fn make_current(&self) -> Result<(), ContextError> {
        match *self {
            Context::WindowedContext(ref c) => {
                let _: () = msg_send![*c.context, update];
                c.context.makeCurrentContext();
            }
            Context::HeadlessContext(ref c) => {
                let _: () = msg_send![*c.context, update];
                c.context.makeCurrentContext();
            }
        }
        Ok(())
    }

    #[inline]
    pub unsafe fn make_not_current(&self) -> Result<(), ContextError> {
        if self.is_current() {
            match *self {
                Context::WindowedContext(ref c) => {
                    let _: () = msg_send![*c.context, update];
                    NSOpenGLContext::clearCurrentContext(nil);
                }
                Context::HeadlessContext(ref c) => {
                    let _: () = msg_send![*c.context, update];
                    NSOpenGLContext::clearCurrentContext(nil);
                }
            }
        }
        Ok(())
    }

    #[inline]
    pub fn is_current(&self) -> bool {
        unsafe {
            let context = match *self {
                Context::WindowedContext(ref c) => *c.context,
                Context::HeadlessContext(ref c) => *c.context,
            };

            let pool = NSAutoreleasePool::new(nil);
            let current = NSOpenGLContext::currentContext(nil);
            let res = if current != nil {
                let is_equal: BOOL = msg_send![current, isEqual: context];
                is_equal != NO
            } else {
                false
            };
            let _: () = msg_send![pool, release];
            res
        }
    }

    pub fn get_proc_address(&self, addr: &str) -> *const core::ffi::c_void {
        let symbol_name: CFString = FromStr::from_str(addr).unwrap();
        let framework_name: CFString =
            FromStr::from_str("com.apple.opengl").unwrap();
        let framework = unsafe {
            CFBundleGetBundleWithIdentifier(
                framework_name.as_concrete_TypeRef(),
            )
        };
        let symbol = unsafe {
            CFBundleGetFunctionPointerForName(
                framework,
                symbol_name.as_concrete_TypeRef(),
            )
        };
        symbol as *const _
    }

    #[inline]
    pub fn swap_buffers(&self) -> Result<(), ContextError> {
        unsafe {
            match *self {
                Context::WindowedContext(ref c) => {
                    let pool = NSAutoreleasePool::new(nil);
                    c.context.flushBuffer();
                    let _: () = msg_send![pool, release];
                }
                Context::HeadlessContext(_) => unreachable!(),
            }
        }
        Ok(())
    }

    #[inline]
    pub fn swap_buffers_with_damage(
        &self,
        _rects: &[Rect],
    ) -> Result<(), ContextError> {
        Err(ContextError::OsError("buffer damage not suported".to_string()))
    }

    #[inline]
    pub fn swap_buffers_with_damage_supported(&self) -> bool {
        false
    }

    #[inline]
    pub fn get_api(&self) -> crate::Api {
        crate::Api::OpenGl
    }

    #[inline]
    pub fn get_pixel_format(&self) -> PixelFormat {
        match *self {
            Context::WindowedContext(ref c) => c.pixel_format.clone(),
            Context::HeadlessContext(_) => unreachable!(),
        }
    }

    #[inline]
    pub unsafe fn raw_handle(&self) -> *mut raw::c_void {
        match self {
            Context::WindowedContext(c) => {
                c.context.deref().CGLContextObj() as *mut _
            }
            Context::HeadlessContext(c) => {
                c.context.deref().CGLContextObj() as *mut _
            }
        }
    }

    #[inline]
    fn get_id(&self) -> IdRef {
        match self {
            Context::WindowedContext(w) => w.context.clone(),
            Context::HeadlessContext(h) => h.context.clone(),
        }
    }
}

#[derive(Debug)]
struct IdRef(id);

impl IdRef {
    fn new(i: id) -> IdRef {
        IdRef(i)
    }

    #[allow(dead_code)]
    fn retain(i: id) -> IdRef {
        if i != nil {
            let _: id = unsafe { msg_send![i, retain] };
        }
        IdRef(i)
    }

    fn non_nil(self) -> Option<IdRef> {
        if self.0 == nil {
            None
        } else {
            Some(self)
        }
    }
}

impl Drop for IdRef {
    fn drop(&mut self) {
        if self.0 != nil {
            let _: () = unsafe { msg_send![self.0, release] };
        }
    }
}

impl Deref for IdRef {
    type Target = id;
    fn deref<'a>(&'a self) -> &'a id {
        &self.0
    }
}

impl Clone for IdRef {
    fn clone(&self) -> IdRef {
        if self.0 != nil {
            let _: id = unsafe { msg_send![self.0, retain] };
        }
        IdRef(self.0)
    }
}

unsafe impl Send for Context {}
unsafe impl Sync for Context {}
<|endoftext|>


<commit_before>use futures_core::future::Future;
use futures_core::task::{LocalWaker, Poll};
use pin_utils::{unsafe_pinned, unsafe_unpinned};
use std::marker::PhantomPinned;
use std::pin::Pin;
use std::ptr;

/// Combinator for the
/// [`FutureTestExt::assert_unmoved`](super::FutureTestExt::assert_unmoved)
/// method.
#[derive(Debug, Clone)]
#[must_use = "futures do nothing unless polled"]
pub struct AssertUnmoved<Fut> {
    future: Fut,
    this_ptr: *const AssertUnmoved<Fut>,
    _pinned: PhantomPinned,
}

impl<Fut> AssertUnmoved<Fut> {
    unsafe_pinned!(future: Fut);
    unsafe_unpinned!(this_ptr: *const Self);

    pub(super) fn new(future: Fut) -> Self {
        Self {
            future,
            this_ptr: ptr::null(),
            _pinned: PhantomPinned,
        }
    }
}

impl<Fut: Future> Future for AssertUnmoved<Fut> {
    type Output = Fut::Output;

    fn poll(
        mut self: Pin<&mut Self>,
        lw: &LocalWaker,
    ) -> Poll<Self::Output> {
        let cur_this = &*self as *const Self;
        if self.this_ptr.is_null() {
            // First time being polled
            *self.as_mut().this_ptr() = cur_this;
        } else {
            assert_eq!(self.this_ptr, cur_this, "Future moved between poll calls");
        }
        self.as_mut().future().poll(lw)
    }
}

impl<Fut> Drop for AssertUnmoved<Fut> {
    fn drop(&mut self) {
        let cur_this = &*self as *const Self;
        assert_eq!(self.this_ptr, cur_this, "Future moved before drop");
    }
}
<commit_msg>Don't panic in AssertUnmoved while panicking<commit_after>use futures_core::future::Future;
use futures_core::task::{LocalWaker, Poll};
use pin_utils::{unsafe_pinned, unsafe_unpinned};
use std::marker::PhantomPinned;
use std::pin::Pin;
use std::ptr;
use std::thread::panicking;

/// Combinator for the
/// [`FutureTestExt::assert_unmoved`](super::FutureTestExt::assert_unmoved)
/// method.
#[derive(Debug, Clone)]
#[must_use = "futures do nothing unless polled"]
pub struct AssertUnmoved<Fut> {
    future: Fut,
    this_ptr: *const AssertUnmoved<Fut>,
    _pinned: PhantomPinned,
}

impl<Fut> AssertUnmoved<Fut> {
    unsafe_pinned!(future: Fut);
    unsafe_unpinned!(this_ptr: *const Self);

    pub(super) fn new(future: Fut) -> Self {
        Self {
            future,
            this_ptr: ptr::null(),
            _pinned: PhantomPinned,
        }
    }
}

impl<Fut: Future> Future for AssertUnmoved<Fut> {
    type Output = Fut::Output;

    fn poll(
        mut self: Pin<&mut Self>,
        lw: &LocalWaker,
    ) -> Poll<Self::Output> {
        let cur_this = &*self as *const Self;
        if self.this_ptr.is_null() {
            // First time being polled
            *self.as_mut().this_ptr() = cur_this;
        } else {
            assert_eq!(self.this_ptr, cur_this, "Future moved between poll calls");
        }
        self.as_mut().future().poll(lw)
    }
}

impl<Fut> Drop for AssertUnmoved<Fut> {
    fn drop(&mut self) {
        // If the thread is panicking then we can't panic again as that will
        // cause the process to be aborted.
        if !panicking() && !self.this_ptr.is_null() {
            let cur_this = &*self as *const Self;
            assert_eq!(self.this_ptr, cur_this, "Future moved before drop");
        }
    }
}

#[cfg(test)]
mod tests {
    use futures_core::future::Future;
    use futures_core::task::Poll;
    use futures_util::future::empty;
    use futures_util::task::noop_local_waker;
    use std::pin::Pin;

    use super::AssertUnmoved;

    #[test]
    fn dont_panic_when_not_polled() {
        // This shouldn't panic.
        let future = AssertUnmoved::new(empty::<()>());
        drop(future);
    }

    #[test]
    #[should_panic(expected = "Future moved between poll calls")]
    fn dont_double_panic() {
        // This test should only panic, not abort the process.
        let waker = noop_local_waker();

        // First we allocate the future on the stack and poll it.
        let mut future = AssertUnmoved::new(empty::<()>());
        let pinned_future = unsafe { Pin::new_unchecked(&mut future) };
        assert_eq!(pinned_future.poll(&waker), Poll::Pending);

        // Next we move it back to the heap and poll it again. This second call
        // should panic (as the future is moved), but we shouldn't panic again
        // whilst dropping `AssertUnmoved`.
        let mut future = Box::new(future);
        let pinned_boxed_future = unsafe { Pin::new_unchecked(&mut *future) };
        assert_eq!(pinned_boxed_future.poll(&waker), Poll::Pending);
    }
}
<|endoftext|>


<commit_before>use serde::Deserialize;
use serde_json::{Result, Value};

#[derive(Deserialize)]
struct Versions {
    stable: Option<String>,
    bottle: bool,
}
#[derive(Deserialize)]
struct Bottle {
    stable: Option<BottleStable>,
}
#[derive(Deserialize)]
struct BottleStable {
    rebuild: u8,
    files: Value,
}
#[derive(Deserialize)]
struct Formula {
    name: String,
    versions: Versions,
    bottle: Bottle,
}

#[derive(Deserialize)]
struct Formulae(Vec<Formula>);

#[derive(Deserialize)]
struct BottleInfo {
    url: String,
    sha256: String,
}

fn main() -> Result<()> {
    let f: Formulae = serde_json::from_reader(std::io::stdin())?;
    for f in f.0 {
        if f.versions.bottle {
            let name = f.name;
            let ver = f.versions.stable.unwrap();
            let rebuild = f.bottle.stable.as_ref().unwrap().rebuild;
            for (platform, v) in f
                .bottle
                .stable
                .as_ref()
                .unwrap()
                .files
                .clone()
                .as_object()
                .unwrap()
                .iter()
            {
                let v: BottleInfo = serde_json::from_value(v.clone()).unwrap();
                //a2ps-4.14.arm64_big_sur.bottle.4.tar.gz
                println!(
                    "{} {} bottles/{}-{}.{}.bottle.{}.tar.gz",
                    v.sha256, v.url, name, ver, platform, rebuild
                );
            }
        }
    }
    Ok(())
}
<commit_msg>homebrew-bottles: bottles-json refoactor<commit_after>use serde::Deserialize;
use serde_json::{Result, Value};

#[derive(Deserialize)]
struct Versions {
    stable: Option<String>,
    bottle: bool,
}
#[derive(Deserialize)]
struct Bottle {
    stable: Option<BottleStable>,
}
#[derive(Deserialize)]
struct BottleStable {
    rebuild: u8,
    files: Value,
}
#[derive(Deserialize)]
struct Formula {
    name: String,
    versions: Versions,
    bottle: Bottle,
}

#[derive(Deserialize)]
struct Formulae(Vec<Formula>);

#[derive(Deserialize)]
struct BottleInfo {
    url: String,
    sha256: String,
}

fn main() -> Result<()> {
    let f: Formulae = serde_json::from_reader(std::io::stdin())?;
    for f in f.0 {
        if f.versions.bottle {
            let name = f.name;
            let ver = f.versions.stable.unwrap();
            let bs = f.bottle.stable.as_ref().unwrap();
            let rebuild = bs.rebuild;
            for (platform, v) in bs
                .files
                .as_object()
                .unwrap()
                .iter()
            {
                let v: BottleInfo = serde_json::from_value(v.clone()).unwrap();
                //a2ps-4.14.arm64_big_sur.bottle.4.tar.gz
                println!(
                    "{} {} bottles/{}-{}.{}.bottle.{}.tar.gz",
                    v.sha256, v.url, name, ver, platform, rebuild
                );
            }
        }
    }
    Ok(())
}
<|endoftext|>


<commit_before>#![allow(non_snake_case)]
#![allow(non_camel_case_types)]
#![allow(dead_code)]

extern crate libc;
use libc::{c_char};

pub mod CPF
{
    pub mod GOM
    {
        pub const OK                               : u32 = 0x72D622E4;      // Sub-system: "Core" Code: "OK"
        pub const IN_USE                           : u32 = 0x72D65159;      // Sub-system: "Core" Code: "InUse"
        pub const ERROR                            : u32 = 0xF2D60ECC;      // Sub-system: "Core" Code: "Error"
        pub const UNKNOWN_INTERFACE                : u32 = 0xF2D623E7;      // Sub-system: "Core" Code: "UnknownInterface"
        pub const INVALID_PARAMETER                : u32 = 0xF2D60CEA;      // Sub-system: "Core" Code: "InvalidParameter"
        pub const OUT_OF_MEMORY                    : u32 = 0xF2D655B8;      // Sub-system: "Core" Code: "OutOfMemory"
        pub const UNKNOWN_CLASS                    : u32 = 0xF2D67F71;      // Sub-system: "Core" Code: "UnknownClass"
        pub const NOT_IMPLEMENTED                  : u32 = 0xF2D67EFE;      // Sub-system: "Core" Code: "NotImplemented"
        pub const INVALID                          : u32 = 0xF2D63EDB;      // Sub-system: "Core" Code: "Invalid"
        pub const NOT_ENOUGH_SPACE                 : u32 = 0xF2D6075E;      // Sub-system: "Core" Code: "NotEnoughSpace"
        pub const NOT_INITIALIZED                  : u32 = 0xF2D66D72;      // Sub-system: "Core" Code: "NotInitialized"
        pub const INITIALIZATION_FAILURE           : u32 = 0xF2D66346;      // Sub-system: "Core" Code: "InitializationFailure"
        pub const OUT_OF_RANGE                     : u32 = 0xF2D6436A;      // Sub-system: "Core" Code: "OutOfRange"
        pub const DUPLICATE                        : u32 = 0xF2D675B4;      // Sub-system: "Core" Code: "Duplicate"
        pub const REGISTRY_ERROR                   : u32 = 0xF2D62784;      // Sub-system: "Core" Code: "RegistryError"
        pub const NOT_RUNNING                      : u32 = 0xF2D67C4C;      // Sub-system: "Core" Code: "NotRunning"

        pub struct iBase_Vtbl
        {
            AddRef: extern fn() -> u32,
            Release: extern fn() -> u32,
            Cast: extern fn(id: u64, outIface: ::libc::c_void) -> u32,
        }
        pub struct iBase
        {
        	vtbl: *const iBase_Vtbl
        }
        pub struct IID_CID
        {}
    }
    pub mod Plugin
    {

        pub struct iRegistry_Vtbl
        {
            Load: extern fn(library: ::libc::c_char) -> u32,
            CanUnload: extern fn(library: ::libc::c_char) -> u32,
            Unload: extern fn(library: ::libc::c_char) -> u32,
            Install: extern fn(cid: u64, clsInst: iClassInstance) -> u32,
            Remove: extern fn(cid: u64) -> u32,
            GetClassInstance: extern fn(cid: u64, clsInst: iClassInstance) -> u32,
            Exists: extern fn(cid: u64) -> u32,
            Create: extern fn(outer: ::CPF::GOM::iBase, cid: u64, iid: u64, outIFace: ::libc::c_void) -> u32,
            ClassInstall: extern fn(count: i32, pairs: ::CPF::GOM::IID_CID) -> u32,
            ClassRemove: extern fn(count: i32, pairs: ::CPF::GOM::IID_CID) -> u32,
            GetClasses: extern fn(id: u64, count: i32, cid: u64) -> u32,
            InstanceInstall: extern fn(id: u64, instance: ::CPF::GOM::iBase) -> u32,
            InstanceRemove: extern fn(id: u64) -> u32,
            GetInstance: extern fn(id: u64, outIface: ::CPF::GOM::iBase) -> u32,
        }
        pub struct iRegistry
        {
        	vtbl: *const iRegistry_Vtbl
        }

        pub struct iClassInstance_Vtbl
        {
            CreateInstance: extern fn(registry: iRegistry, outer: ::CPF::GOM::iBase, outIface: ::CPF::GOM::iBase) -> u32,
        }
        pub struct iClassInstance
        {
        	vtbl: *const iClassInstance_Vtbl
        }
    }
}

#[no_mangle]
pub extern "stdcall" fn Install(registry: *const ::CPF::Plugin::iRegistry) -> u32
{
	CPF::GOM::OK
}

#[no_mangle]
pub extern "stdcall" fn CanUnload() -> u32
{
	CPF::GOM::OK
}

#[no_mangle]
pub extern "stdcall" fn Remove(registry: *const ::CPF::Plugin::iRegistry) -> u32
{
	CPF::GOM::OK
}
<commit_msg>Messy but getting close to full function.<commit_after>#![allow(non_snake_case)]
#![allow(non_camel_case_types)]
#![allow(dead_code)]

extern crate libc;

pub mod CPF
{
    pub mod GOM
    {
        pub const OK                               : u32 = 0x72D622E4;      // Sub-system: "Core" Code: "OK"
        pub const IN_USE                           : u32 = 0x72D65159;      // Sub-system: "Core" Code: "InUse"
        pub const ERROR                            : u32 = 0xF2D60ECC;      // Sub-system: "Core" Code: "Error"
        pub const UNKNOWN_INTERFACE                : u32 = 0xF2D623E7;      // Sub-system: "Core" Code: "UnknownInterface"
        pub const INVALID_PARAMETER                : u32 = 0xF2D60CEA;      // Sub-system: "Core" Code: "InvalidParameter"
        pub const OUT_OF_MEMORY                    : u32 = 0xF2D655B8;      // Sub-system: "Core" Code: "OutOfMemory"
        pub const UNKNOWN_CLASS                    : u32 = 0xF2D67F71;      // Sub-system: "Core" Code: "UnknownClass"
        pub const NOT_IMPLEMENTED                  : u32 = 0xF2D67EFE;      // Sub-system: "Core" Code: "NotImplemented"
        pub const INVALID                          : u32 = 0xF2D63EDB;      // Sub-system: "Core" Code: "Invalid"
        pub const NOT_ENOUGH_SPACE                 : u32 = 0xF2D6075E;      // Sub-system: "Core" Code: "NotEnoughSpace"
        pub const NOT_INITIALIZED                  : u32 = 0xF2D66D72;      // Sub-system: "Core" Code: "NotInitialized"
        pub const INITIALIZATION_FAILURE           : u32 = 0xF2D66346;      // Sub-system: "Core" Code: "InitializationFailure"
        pub const OUT_OF_RANGE                     : u32 = 0xF2D6436A;      // Sub-system: "Core" Code: "OutOfRange"
        pub const DUPLICATE                        : u32 = 0xF2D675B4;      // Sub-system: "Core" Code: "Duplicate"
        pub const REGISTRY_ERROR                   : u32 = 0xF2D62784;      // Sub-system: "Core" Code: "RegistryError"
        pub const NOT_RUNNING                      : u32 = 0xF2D67C4C;      // Sub-system: "Core" Code: "NotRunning"

        pub struct iBase_Vtbl
        {
            pub AddRef: extern "stdcall" fn() -> u32,
            pub Release: extern "stdcall" fn() -> u32,
            pub Cast: extern "stdcall" fn(id: u64, outIface: ::libc::c_void) -> u32,
        }
        pub struct iBase
        {
        	pub vtbl: *const iBase_Vtbl
        }
        pub struct IID_CID
        {}
    }
    pub mod Plugin
    {

        pub struct iRegistry_Vtbl
        {
            pub Load: extern "stdcall" fn(library: ::libc::c_char) -> u32,
            pub CanUnload: extern "stdcall" fn(library: ::libc::c_char) -> u32,
            pub Unload: extern "stdcall" fn(library: ::libc::c_char) -> u32,
            pub Install: extern "stdcall" fn(cid: u64, clsInst: iClassInstance) -> u32,
            pub Remove: extern "stdcall" fn(cid: u64) -> u32,
            pub GetClassInstance: extern "stdcall" fn(cid: u64, clsInst: iClassInstance) -> u32,
            pub Exists: extern "stdcall" fn(cid: u64) -> u32,
            pub Create: extern "stdcall" fn(outer: ::CPF::GOM::iBase, cid: u64, iid: u64, outIFace: ::libc::c_void) -> u32,
            pub ClassInstall: extern "stdcall" fn(count: i32, pairs: ::CPF::GOM::IID_CID) -> u32,
            pub ClassRemove: extern "stdcall" fn(count: i32, pairs: ::CPF::GOM::IID_CID) -> u32,
            pub GetClasses: extern "stdcall" fn(id: u64, count: i32, cid: u64) -> u32,
            pub InstanceInstall: extern "stdcall" fn(id: u64, instance: ::CPF::GOM::iBase) -> u32,
            pub InstanceRemove: extern "stdcall" fn(id: u64) -> u32,
            pub GetInstance: extern "stdcall" fn(id: u64, outIface: ::CPF::GOM::iBase) -> u32,
        }
        pub struct iRegistry
        {
        	pub vtbl: *const iRegistry_Vtbl
        }

        pub struct iClassInstance_Vtbl
        {
            pub AddRef: extern "stdcall" fn(this: *mut iClassInstance) -> u32,
            pub Release: extern "stdcall" fn(this: *mut iClassInstance) -> u32,
            pub Cast: extern "stdcall" fn(this: *mut iClassInstance, id: u64, outIface: ::libc::c_void) -> u32,
            pub CreateInstance: extern "stdcall" fn(this: *mut iClassInstance, registry: iRegistry, outer: ::CPF::GOM::iBase, outIface: ::CPF::GOM::iBase) -> u32,
        }
        pub struct iClassInstance
        {
        	pub vtbl: *const iClassInstance_Vtbl
        }
    }
}

// *****************
struct TestRustCreator
{
    ref_count: u32
}

impl TestRustCreator
{
    pub fn add_ref(&mut self) -> u32
    {
        self.ref_count = self.ref_count + 1;
        return self.ref_count;
    }
    pub fn release(&mut self) -> u32
    {
        self.ref_count -= 1;
        return self.ref_count;
    }
    pub fn cast(&self, id: u64, out_iface: ::libc::c_void) -> u32
    {
        ::CPF::GOM::UNKNOWN_CLASS
    }
    pub fn create_instance(&self, registry: ::CPF::Plugin::iRegistry, outer: ::CPF::GOM::iBase, out_iface: ::CPF::GOM::iBase) -> u32
    {
        ::CPF::GOM::UNKNOWN_CLASS
    }
}

extern "stdcall" fn add_ref(this: *mut ::CPF::Plugin::iClassInstance) -> u32
{
    unsafe
    {
        let ref mut This: TestRustCreator = *(this as *mut TestRustCreator);
        This.add_ref()
    }
}
extern "stdcall" fn release(this: *mut ::CPF::Plugin::iClassInstance) -> u32
{
    unsafe
    {
        let ref mut This: TestRustCreator = *(this as *mut TestRustCreator);
        This.release()
    }
}
extern "stdcall" fn cast(
    this: *mut ::CPF::Plugin::iClassInstance,
    iid: u64,
    out_iface: ::libc::c_void) -> u32
{
    ::CPF::GOM::UNKNOWN_CLASS
}

extern "stdcall" fn test_create(
    this: *mut ::CPF::Plugin::iClassInstance,
    registry: ::CPF::Plugin::iRegistry,
    outer: ::CPF::GOM::iBase,
    outIface: ::CPF::GOM::iBase) -> u32
{
    CPF::GOM::OK
}

static sTestClassInstance: ::CPF::Plugin::iClassInstance_Vtbl = ::CPF::Plugin::iClassInstance_Vtbl
{
    AddRef: add_ref,
    Release: release,
    Cast: cast,
    CreateInstance: test_create
};

// **************************
#[no_mangle]
pub extern "stdcall" fn Install(registry: *const ::CPF::Plugin::iRegistry) -> u32
{

//    ((*(*registry).vtbl).Install)(0, sTestClassInstance);
	CPF::GOM::OK
}

#[no_mangle]
pub extern "stdcall" fn CanUnload() -> u32
{
	CPF::GOM::OK
}

#[no_mangle]
pub extern "stdcall" fn Remove(registry: *const ::CPF::Plugin::iRegistry) -> u32
{
	CPF::GOM::OK
}
<|endoftext|>