<commit_before>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn step(&mut self) {
        let counter = self.counter as usize;
        let opcode = self.memory[counter];
        let a = self.memory[counter + 1];
        let b = self.memory[counter + 2];
        let c = self.memory[counter + 3];
        self.exec(opcode, a, b, c);
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32, c: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // M[A] <- M[B] + M[C]
            7 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b + c;
                self.counter += 4;
            },

            // M[A] <- M[B] - M[C]
            8 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b - c;
                self.counter += 4;
            },

            // M[A] <- M[B] x M[C]
            9 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b * c;
                self.counter += 4;
            },

            // M[A] <- M[B] / M[C]
            10 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b / c;
                self.counter += 4;
            },

            // M[A] <- M[B] modulo M[C]
            11 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b % c;
                self.counter += 4;
            },

            // If M[B] < M[C], then M[A] <- 1, else M[A] <- 0
            12 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                if b < c {
                  self.memory[a as usize] = 1;
                } else {
                  self.memory[a as usize] = 0;
                }
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Machine;

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        // PC <- M[A]
        let mut m = Machine::with_data(vec![1, 4, 0, 0, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_then_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        // M[A] <- PC
        let mut m = Machine::with_data(vec![3, 4, 0, 0, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![3, 4, 0, 0, 0], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        // M[A] <- M[B]
        let mut m = Machine::with_data(vec![4, 4, 5, 0, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![4, 4, 5, 0, 7, 7], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        // M[A] <- M[M[B]]
        let mut m = Machine::with_data(vec![5, 4, 5, 0, 6, 7, 0, 8]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![5, 4, 5, 0, 8, 7, 0, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        // M[M[B]] <- M[A]
        let mut m = Machine::with_data(vec![6, 4, 5, 0, 8, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![6, 4, 5, 0, 8, 6, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_7() {
        // M[A] <- M[B] + M[C]
        let mut m = Machine::with_data(vec![7, 4, 2, 3, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![7, 4, 2, 3, 5], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_8() {
        // M[A] <- M[B] - M[C]
        let mut m = Machine::with_data(vec![8, 4, 5, 6, 0, 3, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![8, 4, 5, 6, 1, 3, 2], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_9() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(9, 0, 1, 2);
        assert_eq!(Some(&6), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_10() {
        let mut m = Machine::with_data(vec![1, 6, 2]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(10, 0, 1, 2);
        assert_eq!(Some(&3), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_11() {
        let mut m = Machine::with_data(vec![1, 11, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(11, 0, 1, 2);
        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_then_branch() {
        let mut m = Machine::with_data(vec![2, 3, 4]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_else_branch() {
        let mut m = Machine::with_data(vec![2, 4, 3]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<commit_msg>Use the same data for all the math opcodes.<commit_after>use std::vec::Vec;

pub struct Machine {
    memory: Vec<u32>,
    counter: u32,
}

impl Machine {
    pub fn new() -> Machine {
        Machine {
            memory: vec![0; 2097152],
            counter: 0,
        }
    }

    pub fn with_data(data: Vec<u32>) -> Machine {
        Machine {
            memory: data,
            counter: 0,
        }
    }

    pub fn loc(&self) -> u32 {
        self.counter
    }

    pub fn dump(&self) -> &Vec<u32> {
        &self.memory
    }

    pub fn load(&mut self, data: Vec<u32>) {
        self.memory = data;
    }

    pub fn step(&mut self) {
        let counter = self.counter as usize;
        let opcode = self.memory[counter];
        let a = self.memory[counter + 1];
        let b = self.memory[counter + 2];
        let c = self.memory[counter + 3];
        self.exec(opcode, a, b, c);
    }

    pub fn exec(&mut self, opcode: u32, a: u32, b: u32, c: u32) {
        match opcode {
            // PC <- M[A]
            1 => {
                self.counter = self.memory[a as usize];
            },

            // If M[B] = 0, then PC <- M[A]
            2 => {
                if 0 == self.memory[b as usize] {
                  self.counter = self.memory[a as usize];
                }
            },

            // M[A] <- PC
            3 => {
                self.memory[a as usize] = self.counter;
                self.counter += 4;
            },

            // M[A] <- M[B]
            4 => {
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[A] <- M[M[B]]
            5 => {
                let b = self.memory[b as usize];
                self.memory[a as usize] = self.memory[b as usize];
                self.counter += 4;
            },

            // M[M[B]] <- M[A]
            6 => {
                let b = self.memory[b as usize];
                self.memory[b as usize] = self.memory[a as usize];
                self.counter += 4;
            },

            // M[A] <- M[B] + M[C]
            7 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b + c;
                self.counter += 4;
            },

            // M[A] <- M[B] - M[C]
            8 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b - c;
                self.counter += 4;
            },

            // M[A] <- M[B] x M[C]
            9 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b * c;
                self.counter += 4;
            },

            // M[A] <- M[B] / M[C]
            10 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b / c;
                self.counter += 4;
            },

            // M[A] <- M[B] modulo M[C]
            11 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                self.memory[a as usize] = b % c;
                self.counter += 4;
            },

            // If M[B] < M[C], then M[A] <- 1, else M[A] <- 0
            12 => {
                let b = self.memory[b as usize];
                let c = self.memory[c as usize];
                if b < c {
                  self.memory[a as usize] = 1;
                } else {
                  self.memory[a as usize] = 0;
                }
                self.counter += 4;
            },

            // Unknown opcode
            _ => {
            },
        }
    }
}

#[cfg(test)]
mod tests {
    use super::Machine;

    #[test]
    fn it_loads_data() {
        let mut m = Machine::new();
        assert_eq!(Some(&0), m.dump().first());

        m.load(vec![1]);
        assert_eq!(Some(&1), m.dump().first());
    }

    #[test]
    fn it_runs_opcode_1() {
        // PC <- M[A]
        let mut m = Machine::with_data(vec![1, 4, 0, 0, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_then_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 0]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(2, m.loc());
    }

    #[test]
    fn it_runs_opcode_2_else_branch() {
        // If M[B] = 0, then PC <- M[A]
        let mut m = Machine::with_data(vec![2, 4, 5, 0, 2, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(0, m.loc());
    }

    #[test]
    fn it_runs_opcode_3() {
        // M[A] <- PC
        let mut m = Machine::with_data(vec![3, 4, 0, 0, 1]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![3, 4, 0, 0, 0], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_4() {
        // M[A] <- M[B]
        let mut m = Machine::with_data(vec![4, 4, 5, 0, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![4, 4, 5, 0, 7, 7], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_5() {
        // M[A] <- M[M[B]]
        let mut m = Machine::with_data(vec![5, 4, 5, 0, 6, 7, 0, 8]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![5, 4, 5, 0, 8, 7, 0, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_6() {
        // M[M[B]] <- M[A]
        let mut m = Machine::with_data(vec![6, 4, 5, 0, 8, 6, 7]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![6, 4, 5, 0, 8, 6, 8], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_7() {
        // M[A] <- M[B] + M[C]
        let mut m = Machine::with_data(vec![7, 4, 5, 6, 0, 11, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![7, 4, 5, 6, 13, 11, 2], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_8() {
        // M[A] <- M[B] - M[C]
        let mut m = Machine::with_data(vec![8, 4, 5, 6, 0, 11, 2]);

        assert_eq!(0, m.loc());
        m.step();
        assert_eq!(&vec![8, 4, 5, 6, 9, 11, 2], m.dump());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_9() {
        let mut m = Machine::with_data(vec![1, 2, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(9, 0, 1, 2);
        assert_eq!(Some(&6), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_10() {
        let mut m = Machine::with_data(vec![1, 6, 2]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(10, 0, 1, 2);
        assert_eq!(Some(&3), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_11() {
        let mut m = Machine::with_data(vec![1, 11, 3]);

        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(11, 0, 1, 2);
        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_then_branch() {
        let mut m = Machine::with_data(vec![2, 3, 4]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&1), m.dump().first());
        assert_eq!(4, m.loc());
    }

    #[test]
    fn it_runs_opcode_12_else_branch() {
        let mut m = Machine::with_data(vec![2, 4, 3]);

        assert_eq!(Some(&2), m.dump().first());
        assert_eq!(0, m.loc());
        m.exec(12, 0, 1, 2);
        assert_eq!(Some(&0), m.dump().first());
        assert_eq!(4, m.loc());
    }
}
<|endoftext|>


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

use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StateSyncConfig {
    // Size of chunk to request for state synchronization
    pub chunk_limit: u64,
    // The timeout of the state sync client to process a commit notification (in milliseconds)
    pub client_commit_timeout_ms: u64,
    // default timeout used for long polling to remote peer
    pub long_poll_timeout_ms: u64,
    // valid maximum chunk limit for sanity check
    pub max_chunk_limit: u64,
    // valid maximum timeout limit for sanity check
    pub max_timeout_ms: u64,
    // The timeout of the state sync coordinator to receive a commit ack from mempool (in milliseconds)
    pub mempool_commit_timeout_ms: u64,
    // default timeout to make state sync progress by sending chunk requests to a certain number of networks
    // if no progress is made by sending chunk requests to a number of networks,
    // the next sync request will be multicasted, i.e. sent to more networks
    pub multicast_timeout_ms: u64,
    // The timeout for ensuring sync requests are making progress (i.e., the maximum time between
    // commits when processing a sync request).
    pub sync_request_timeout_ms: u64,
    // interval used for checking state synchronization progress
    pub tick_interval_ms: u64,

    // Everything above belongs to state sync v1 and will be removed in the future.
    pub data_streaming_service: DataStreamingServiceConfig,
    pub aptos_data_client: AptosDataClientConfig,
    pub state_sync_driver: StateSyncDriverConfig,
    pub storage_service: StorageServiceConfig,
}

impl Default for StateSyncConfig {
    fn default() -> Self {
        Self {
            chunk_limit: 1000,
            client_commit_timeout_ms: 5_000,
            long_poll_timeout_ms: 10_000,
            max_chunk_limit: 1000,
            max_timeout_ms: 120_000,
            mempool_commit_timeout_ms: 5_000,
            multicast_timeout_ms: 30_000,
            sync_request_timeout_ms: 60_000,
            tick_interval_ms: 100,
            data_streaming_service: DataStreamingServiceConfig::default(),
            aptos_data_client: AptosDataClientConfig::default(),
            state_sync_driver: StateSyncDriverConfig::default(),
            storage_service: StorageServiceConfig::default(),
        }
    }
}

/// The bootstrapping mode determines how the node will bootstrap to the latest
/// blockchain state, e.g., directly download the latest account states.
#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
pub enum BootstrappingMode {
    ApplyTransactionOutputsFromGenesis, // Applies transaction outputs (starting at genesis)
    DownloadLatestAccountStates,        // Downloads the account states (at the latest version)
    ExecuteTransactionsFromGenesis,     // Executes transactions (starting at genesis)
}

/// The continuous syncing mode determines how the node will stay up-to-date
/// once it has bootstrapped and the blockchain continues to grow, e.g.,
/// continuously executing all transactions.
#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
pub enum ContinuousSyncingMode {
    ExecuteTransactions,     // Executes transactions to stay up-to-date
    ApplyTransactionOutputs, // Applies transaction outputs to stay up-to-date
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StateSyncDriverConfig {
    pub bootstrapping_mode: BootstrappingMode, // The mode by which to bootstrap
    pub enable_state_sync_v2: bool,            // If the node should sync with state sync v2
    pub continuous_syncing_mode: ContinuousSyncingMode, // The mode by which to sync after bootstrapping
    pub progress_check_interval_ms: u64, // The interval (ms) at which to check state sync progress
    pub max_connection_deadline_secs: u64, // The max time (secs) to wait for connections from peers
    pub max_consecutive_stream_notifications: u64, // The max number of notifications to process per driver loop
    pub max_pending_data_chunks: u64, // The max number of data chunks pending execution or commit
    pub max_stream_wait_time_ms: u64, // The max time (ms) to wait for a data stream notification
}

/// The default state sync driver config will be the one that gets (and keeps)
/// the node up-to-date as quickly and cheaply as possible.
impl Default for StateSyncDriverConfig {
    fn default() -> Self {
        Self {
            bootstrapping_mode: BootstrappingMode::ApplyTransactionOutputsFromGenesis,
            enable_state_sync_v2: true,
            continuous_syncing_mode: ContinuousSyncingMode::ApplyTransactionOutputs,
            progress_check_interval_ms: 100,
            max_connection_deadline_secs: 10,
            max_consecutive_stream_notifications: 10,
            max_pending_data_chunks: 100,
            max_stream_wait_time_ms: 10_000,
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StorageServiceConfig {
    pub max_account_states_chunk_sizes: u64, // Max num of accounts per chunk
    pub max_concurrent_requests: u64,        // Max num of concurrent storage server tasks
    pub max_epoch_chunk_size: u64,           // Max num of epoch ending ledger infos per chunk
    pub max_network_channel_size: u64,       // Max num of pending network messages
    pub max_transaction_chunk_size: u64,     // Max num of transactions per chunk
    pub max_transaction_output_chunk_size: u64, // Max num of transaction outputs per chunk
    pub storage_summary_refresh_interval_ms: u64, // The interval (ms) to refresh the storage summary
}

impl Default for StorageServiceConfig {
    fn default() -> Self {
        Self {
            max_account_states_chunk_sizes: 1000,
            max_concurrent_requests: 1000,
            max_epoch_chunk_size: 100,
            max_network_channel_size: 1000,
            max_transaction_chunk_size: 1000,
            max_transaction_output_chunk_size: 1000,
            storage_summary_refresh_interval_ms: 250,
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct DataStreamingServiceConfig {
    // The interval (milliseconds) at which to refresh the global data summary.
    pub global_summary_refresh_interval_ms: u64,

    // Maximum number of concurrent data client requests (per stream).
    pub max_concurrent_requests: u64,

    // Maximum channel sizes for each data stream listener. If messages are not
    // consumed, they will be dropped (oldest messages first). The remaining
    // messages will be retrieved using FIFO ordering.
    pub max_data_stream_channel_sizes: u64,

    // Maximum number of retries for a single client request before a data
    // stream will terminate.
    pub max_request_retry: u64,

    // Maximum number of notification ID to response context mappings held in
    // memory. Once the number grows beyond this value, garbage collection occurs.
    pub max_notification_id_mappings: u64,

    // The interval (milliseconds) at which to check the progress of each stream.
    pub progress_check_interval_ms: u64,
}

impl Default for DataStreamingServiceConfig {
    fn default() -> Self {
        Self {
            global_summary_refresh_interval_ms: 250,
            max_concurrent_requests: 1,
            max_data_stream_channel_sizes: 1000,
            max_request_retry: 3,
            max_notification_id_mappings: 2000,
            progress_check_interval_ms: 100,
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct AptosDataClientConfig {
    pub response_timeout_ms: u64, // Timeout (in milliseconds) when waiting for a response
    pub summary_poll_interval_ms: u64, // Interval (in milliseconds) between data summary polls
}

impl Default for AptosDataClientConfig {
    fn default() -> Self {
        Self {
            response_timeout_ms: 10000,
            summary_poll_interval_ms: 250,
        }
    }
}
<commit_msg>[Config] Disable SSv2.<commit_after>// Copyright (c) Aptos
// SPDX-License-Identifier: Apache-2.0

use serde::{Deserialize, Serialize};

#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StateSyncConfig {
    // Size of chunk to request for state synchronization
    pub chunk_limit: u64,
    // The timeout of the state sync client to process a commit notification (in milliseconds)
    pub client_commit_timeout_ms: u64,
    // default timeout used for long polling to remote peer
    pub long_poll_timeout_ms: u64,
    // valid maximum chunk limit for sanity check
    pub max_chunk_limit: u64,
    // valid maximum timeout limit for sanity check
    pub max_timeout_ms: u64,
    // The timeout of the state sync coordinator to receive a commit ack from mempool (in milliseconds)
    pub mempool_commit_timeout_ms: u64,
    // default timeout to make state sync progress by sending chunk requests to a certain number of networks
    // if no progress is made by sending chunk requests to a number of networks,
    // the next sync request will be multicasted, i.e. sent to more networks
    pub multicast_timeout_ms: u64,
    // The timeout for ensuring sync requests are making progress (i.e., the maximum time between
    // commits when processing a sync request).
    pub sync_request_timeout_ms: u64,
    // interval used for checking state synchronization progress
    pub tick_interval_ms: u64,

    // Everything above belongs to state sync v1 and will be removed in the future.
    pub data_streaming_service: DataStreamingServiceConfig,
    pub aptos_data_client: AptosDataClientConfig,
    pub state_sync_driver: StateSyncDriverConfig,
    pub storage_service: StorageServiceConfig,
}

impl Default for StateSyncConfig {
    fn default() -> Self {
        Self {
            chunk_limit: 1000,
            client_commit_timeout_ms: 5_000,
            long_poll_timeout_ms: 10_000,
            max_chunk_limit: 1000,
            max_timeout_ms: 120_000,
            mempool_commit_timeout_ms: 5_000,
            multicast_timeout_ms: 30_000,
            sync_request_timeout_ms: 60_000,
            tick_interval_ms: 100,
            data_streaming_service: DataStreamingServiceConfig::default(),
            aptos_data_client: AptosDataClientConfig::default(),
            state_sync_driver: StateSyncDriverConfig::default(),
            storage_service: StorageServiceConfig::default(),
        }
    }
}

/// The bootstrapping mode determines how the node will bootstrap to the latest
/// blockchain state, e.g., directly download the latest account states.
#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
pub enum BootstrappingMode {
    ApplyTransactionOutputsFromGenesis, // Applies transaction outputs (starting at genesis)
    DownloadLatestAccountStates,        // Downloads the account states (at the latest version)
    ExecuteTransactionsFromGenesis,     // Executes transactions (starting at genesis)
}

/// The continuous syncing mode determines how the node will stay up-to-date
/// once it has bootstrapped and the blockchain continues to grow, e.g.,
/// continuously executing all transactions.
#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
pub enum ContinuousSyncingMode {
    ExecuteTransactions,     // Executes transactions to stay up-to-date
    ApplyTransactionOutputs, // Applies transaction outputs to stay up-to-date
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StateSyncDriverConfig {
    pub bootstrapping_mode: BootstrappingMode, // The mode by which to bootstrap
    pub enable_state_sync_v2: bool,            // If the node should sync with state sync v2
    pub continuous_syncing_mode: ContinuousSyncingMode, // The mode by which to sync after bootstrapping
    pub progress_check_interval_ms: u64, // The interval (ms) at which to check state sync progress
    pub max_connection_deadline_secs: u64, // The max time (secs) to wait for connections from peers
    pub max_consecutive_stream_notifications: u64, // The max number of notifications to process per driver loop
    pub max_pending_data_chunks: u64, // The max number of data chunks pending execution or commit
    pub max_stream_wait_time_ms: u64, // The max time (ms) to wait for a data stream notification
}

/// The default state sync driver config will be the one that gets (and keeps)
/// the node up-to-date as quickly and cheaply as possible.
impl Default for StateSyncDriverConfig {
    fn default() -> Self {
        Self {
            bootstrapping_mode: BootstrappingMode::ApplyTransactionOutputsFromGenesis,
            enable_state_sync_v2: false,
            continuous_syncing_mode: ContinuousSyncingMode::ApplyTransactionOutputs,
            progress_check_interval_ms: 100,
            max_connection_deadline_secs: 10,
            max_consecutive_stream_notifications: 10,
            max_pending_data_chunks: 100,
            max_stream_wait_time_ms: 10_000,
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct StorageServiceConfig {
    pub max_account_states_chunk_sizes: u64, // Max num of accounts per chunk
    pub max_concurrent_requests: u64,        // Max num of concurrent storage server tasks
    pub max_epoch_chunk_size: u64,           // Max num of epoch ending ledger infos per chunk
    pub max_network_channel_size: u64,       // Max num of pending network messages
    pub max_transaction_chunk_size: u64,     // Max num of transactions per chunk
    pub max_transaction_output_chunk_size: u64, // Max num of transaction outputs per chunk
    pub storage_summary_refresh_interval_ms: u64, // The interval (ms) to refresh the storage summary
}

impl Default for StorageServiceConfig {
    fn default() -> Self {
        Self {
            max_account_states_chunk_sizes: 1000,
            max_concurrent_requests: 1000,
            max_epoch_chunk_size: 100,
            max_network_channel_size: 1000,
            max_transaction_chunk_size: 1000,
            max_transaction_output_chunk_size: 1000,
            storage_summary_refresh_interval_ms: 250,
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct DataStreamingServiceConfig {
    // The interval (milliseconds) at which to refresh the global data summary.
    pub global_summary_refresh_interval_ms: u64,

    // Maximum number of concurrent data client requests (per stream).
    pub max_concurrent_requests: u64,

    // Maximum channel sizes for each data stream listener. If messages are not
    // consumed, they will be dropped (oldest messages first). The remaining
    // messages will be retrieved using FIFO ordering.
    pub max_data_stream_channel_sizes: u64,

    // Maximum number of retries for a single client request before a data
    // stream will terminate.
    pub max_request_retry: u64,

    // Maximum number of notification ID to response context mappings held in
    // memory. Once the number grows beyond this value, garbage collection occurs.
    pub max_notification_id_mappings: u64,

    // The interval (milliseconds) at which to check the progress of each stream.
    pub progress_check_interval_ms: u64,
}

impl Default for DataStreamingServiceConfig {
    fn default() -> Self {
        Self {
            global_summary_refresh_interval_ms: 250,
            max_concurrent_requests: 1,
            max_data_stream_channel_sizes: 1000,
            max_request_retry: 3,
            max_notification_id_mappings: 2000,
            progress_check_interval_ms: 100,
        }
    }
}

#[derive(Copy, Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(default, deny_unknown_fields)]
pub struct AptosDataClientConfig {
    pub response_timeout_ms: u64, // Timeout (in milliseconds) when waiting for a response
    pub summary_poll_interval_ms: u64, // Interval (in milliseconds) between data summary polls
}

impl Default for AptosDataClientConfig {
    fn default() -> Self {
        Self {
            response_timeout_ms: 10000,
            summary_poll_interval_ms: 250,
        }
    }
}
<|endoftext|>


<commit_before>use sozu::channel::Channel;
use sozu::messages::Order;
use sozu_command::config::Config;
use sozu_command::data::{AnswerData,ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus,RunState};

use std::iter::FromIterator;
use std::collections::{HashMap,HashSet};
use rand::{thread_rng, Rng};

fn generate_id() -> String {
  let s: String = thread_rng().gen_ascii_chars().take(6).collect();
  format!("ID-{}", s)
}

pub fn save_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::SaveState(path.to_string()),
    None,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn load_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::LoadState(path.to_string()),
    None,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn dump_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::DumpState,
    None,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not dump proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state:\n{}", message.message);
        }
      }
    }
  }
}

pub fn soft_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let mut tags: HashMap<String,String> = HashMap::from_iter(config.proxies.keys().map(|tag| {
    println!("shutting down proxy \"{}\"", tag);
    let id = generate_id();
    channel.write_message(&ConfigMessage::new(
      id.clone(),
      ConfigCommand::ProxyConfiguration(Order::SoftStop),
      Some(tag.clone()),
      //FIXME: we should be able to soft stop one specific worker
      None,
    ));
    (id, tag.clone())
  }));

  loop {
    if tags.is_empty() {
      println!("all proxies shut down");
      break;
    }

    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if !tags.contains_key(&message.id) {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            tags.get(&message.id).map(|tag| println!("Proxy {} is processing: {}", tag, message.message));
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if let Some(tag) = tags.remove(&message.id) {
              println!("Proxy {} shut down: {}", tag, message.message);
            }
          }
        }
      }
    }
  }
}

pub fn hard_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let mut tags: HashMap<String,String> = HashMap::from_iter(config.proxies.keys().map(|tag| {
    println!("shutting down proxy \"{}\"", tag);
    let id = generate_id();
    channel.write_message(&ConfigMessage::new(
      id.clone(),
      ConfigCommand::ProxyConfiguration(Order::HardStop),
      Some(tag.clone()),
      //FIXME: we should be able to soft stop one specific worker
      None,
    ));
    (id, tag.clone())
  }));

  loop {
    if tags.is_empty() {
      println!("all proxies shut down");
      break;
    }

    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if !tags.contains_key(&message.id) {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            tags.get(&message.id).map(|tag| println!("Proxy {} is processing: {}", tag, message.message));
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if let Some(tag) = tags.remove(&message.id) {
              println!("Proxy {} shut down: {}", tag, message.message);
            }
          }
        }
      }
    }
  }
}

pub fn upgrade(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut launching: HashSet<String> = HashSet::new();
            let mut stopping:  HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::LaunchWorker(worker.tag.to_string()),
                None,
                None,
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              launching.insert(id);
            }

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::SoftStop),
                Some(worker.tag.to_string()),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              stopping.insert(id);
            }

            loop {
              println!("launching: {:?}\nstopping: {:?}", launching, stopping);
              if launching.is_empty() && stopping.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
            let id = generate_id();
            channel.write_message(&ConfigMessage::new(
              id.clone(),
              ConfigCommand::UpgradeMaster,
              None,
              None,
            ));

            loop {
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  if &id != &message.id {
                    println!("received message with invalid id: {:?}", message);
                    return;
                  }
                  match message.status {
                    ConfigMessageStatus::Processing => {
                      println!("master is processing: {}", message.message);
                    },
                    ConfigMessageStatus::Error => {
                      println!("could not upgrade the master: {}", message.message);
                      return;
                    },
                    ConfigMessageStatus::Ok => {
                      println!("successfully upgraded the master: {}", message.message);
                      return;
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

pub fn status(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut expecting: HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::Status),
                Some(worker.tag.to_string()),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              expecting.insert(id);
            }

            loop {
              println!("expecting: {:?}", expecting);
              if expecting.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
          }
        }
      }
    }
  }
}
<commit_msg>convert the ctl library to the new message format<commit_after>use sozu::channel::Channel;
use sozu::messages::Order;
use sozu_command::config::Config;
use sozu_command::data::{AnswerData,ConfigCommand,ConfigMessage,ConfigMessageAnswer,ConfigMessageStatus,RunState};

use std::iter::FromIterator;
use std::collections::{HashMap,HashSet};
use rand::{thread_rng, Rng};

fn generate_id() -> String {
  let s: String = thread_rng().gen_ascii_chars().take(6).collect();
  format!("ID-{}", s)
}

pub fn save_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::SaveState(path.to_string()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn load_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, path: &str) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::LoadState(path.to_string()),
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not save proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state saved to {}", path);
        }
      }
    }
  }
}

pub fn dump_state(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::DumpState,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          // do nothing here
          // for other messages, we would loop over read_message
          // until an error or ok message was sent
        },
        ConfigMessageStatus::Error => {
          println!("could not dump proxy state: {}", message.message);
        },
        ConfigMessageStatus::Ok => {
          println!("Proxy state:\n{}", message.message);
        }
      }
    }
  }
}

pub fn soft_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  println!("shutting down proxy");
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(Order::SoftStop),
    //FIXME: we should be able to soft stop one specific worker
    None,
  ));

  loop {
    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        if &id != &message.id {
          println!("received message with invalid id: {:?}", message);
          return;
        }
        match message.status {
          ConfigMessageStatus::Processing => {
            println!("Proxy is processing: {}", message.message);
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            println!("Proxy shut down: {}", message.message);
            break;
          }
        }
      }
    }
  }
}

pub fn hard_stop(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  println!("shutting down proxy");
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ProxyConfiguration(Order::HardStop),
    //FIXME: we should be able to soft stop one specific worker
    None,
  ));

  loop {
    match channel.read_message() {
      None          => println!("the proxy didn't answer"),
      Some(message) => {
        match message.status {
          ConfigMessageStatus::Processing => {
            println!("Proxy is processing: {}", message.message);
          },
          ConfigMessageStatus::Error => {
            println!("could not stop the proxy: {}", message.message);
          },
          ConfigMessageStatus::Ok => {
            if &id == &message.id {
              println!("Proxy shut down: {}", message.message);
              break;
            }
          }
        }
      }
    }
  }
}

pub fn upgrade(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut launching: HashSet<String> = HashSet::new();
            let mut stopping:  HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::LaunchWorker("BLAH".to_string()),
                None,
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              launching.insert(id);
            }

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::SoftStop),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              stopping.insert(id);
            }

            loop {
              println!("launching: {:?}\nstopping: {:?}", launching, stopping);
              if launching.is_empty() && stopping.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if launching.contains(&message.id) {
                        launching.remove(&message.id);
                        println!("launch message with ID {} done", message.id);
                      }
                      if stopping.contains(&message.id) {
                        stopping.remove(&message.id);
                        println!("stop message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
            let id = generate_id();
            channel.write_message(&ConfigMessage::new(
              id.clone(),
              ConfigCommand::UpgradeMaster,
              None,
            ));

            loop {
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  if &id != &message.id {
                    println!("received message with invalid id: {:?}", message);
                    return;
                  }
                  match message.status {
                    ConfigMessageStatus::Processing => {
                      println!("master is processing: {}", message.message);
                    },
                    ConfigMessageStatus::Error => {
                      println!("could not upgrade the master: {}", message.message);
                      return;
                    },
                    ConfigMessageStatus::Ok => {
                      println!("successfully upgraded the master: {}", message.message);
                      return;
                    }
                  }
                }
              }
            }
          }
        }
      }
    }
  }
}

pub fn status(channel: &mut Channel<ConfigMessage,ConfigMessageAnswer>, config: &Config) {
  let id = generate_id();
  channel.write_message(&ConfigMessage::new(
    id.clone(),
    ConfigCommand::ListWorkers,
    None,
  ));

  match channel.read_message() {
    None          => println!("the proxy didn't answer"),
    Some(message) => {
      if id != message.id {
        println!("received message with invalid id: {:?}", message);
        return;
      }
      match message.status {
        ConfigMessageStatus::Processing => {
          println!("should have obtained an answer immediately");
          return;
        },
        ConfigMessageStatus::Error => {
          println!("could not get the worker list: {}", message.message);
          return
        },
        ConfigMessageStatus::Ok => {
          println!("Worker list:\n{:?}", message.data);
          if let Some(AnswerData::Workers(ref workers)) = message.data {
            let mut expecting: HashSet<String> = HashSet::new();

            for ref worker in workers.iter().filter(|worker| worker.run_state == RunState::Running) {
              let id = generate_id();
              let msg = ConfigMessage::new(
                id.clone(),
                ConfigCommand::ProxyConfiguration(Order::Status),
                Some(worker.id),
              );
              println!("sending message: {:?}", msg);
              channel.write_message(&msg);
              expecting.insert(id);
            }

            loop {
              println!("expecting: {:?}", expecting);
              if expecting.is_empty() {
                break;
              }
              match channel.read_message() {
                None          => println!("the proxy didn't answer"),
                Some(message) => {
                  println!("received message: {:?}", message);
                  match message.status {
                    ConfigMessageStatus::Processing => {
                    },
                    ConfigMessageStatus::Error => {
                      println!("error for message[{}]: {}", message.id, message.message);
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    },
                    ConfigMessageStatus::Ok => {
                      if expecting.contains(&message.id) {
                        expecting.remove(&message.id);
                        println!("status message with ID {} done", message.id);
                      }
                    }
                  }
                }
              }
            }

            println!("worker upgrade done");
          }
        }
      }
    }
  }
}
<|endoftext|>


<commit_before>// See LICENSE file for copyright and license details.

use std::cmp;
use collections::hashmap::HashMap;
use stb_tt;
use cgmath::vector::{Vector3, Vector2};
use core::types::{Size2, Point2, MInt};
use core::misc::add_quad_to_vec;
use visualizer::texture::Texture;
use visualizer::types::MFloat;
use visualizer::mesh::Mesh;
use visualizer::shader::Shader;

struct Glyph {
    pos: Vector2<MInt>,
    size: Size2<MInt>,
    index: MInt,
    xoff: MInt,
    yoff: MInt,
}

pub struct FontStash {
    size: MFloat,
    font: stb_tt::Font,
    texture: Texture,
    texture_size: MInt,
    pos: Point2<MInt>,
    glyphs: HashMap<char, Glyph>,
    max_h: MInt,
}

impl FontStash {
    pub fn new(font_path: &str, size: MFloat) -> FontStash {
        // TODO: use updateble mesh
        let texture_size = 1024;
        let font = stb_tt::Font::new(font_path, size);
        let texture = Texture::new_empty(
            Size2{w: texture_size, h: texture_size});
        FontStash {
            size: size,
            font: font,
            texture: texture,
            texture_size: texture_size,
            pos: Vector2{x: 0, y: 0},
            glyphs: HashMap::new(),
            max_h: 0,
        }
    }

    pub fn get_glyph(&mut self, c: char) -> Glyph {
        match self.glyphs.find(&c) {
            Some(&r) => r,
            None => self.add_glyph(c),
        }
    }

    pub fn get_size(&mut self) -> MFloat {
        self.size
    }

    pub fn get_mesh(&mut self, text: &str, shader: &Shader) -> Mesh {
        // TODO: create mesh in c-tor, update mesh vertices data here
        let mut vertex_data = Vec::new();
        let mut tex_data = Vec::new();
        let s = self.texture_size as MFloat;
        let mut i = 0.0;
        for c in text.chars() {
            let glyph = self.get_glyph(c);
            let w = glyph.size.w as MFloat;
            let h = glyph.size.h as MFloat;
            let x1 = glyph.pos.x as MFloat / s;
            let y1 = glyph.pos.y as MFloat / s;
            let x2 = x1 + w / s;
            let y2 = y1 + h / s;
            add_quad_to_vec(
                &mut tex_data,
                Vector2{x: x1, y: y1},
                Vector2{x: x1, y: y2},
                Vector2{x: x2, y: y2},
                Vector2{x: x2, y: y1},
            );
            let yoff = -glyph.yoff as MFloat;
            add_quad_to_vec(
                &mut vertex_data,
                Vector3{x: i, y: yoff, z: 0.0},
                Vector3{x: i, y: yoff - h, z: 0.0},
                Vector3{x: w + i, y: yoff - h, z: 0.0},
                Vector3{x: w + i, y: yoff, z: 0.0},
            );
            i += w + glyph.xoff as MFloat;
        }
        let mut mesh = Mesh::new(vertex_data.as_slice());
        mesh.set_texture(self.texture, tex_data.as_slice());
        mesh.prepare(shader);
        mesh
    }

    fn add_glyph(&mut self, c: char) -> Glyph {
        assert!(self.glyphs.find(&c).is_none());
        let index = self.font.find_glyph_index(c);
        let (bitmap, w, h, xoff, yoff) = self.font.get_glyph(index);
        if self.pos.x + w > self.texture_size {
            self.pos.y += self.max_h;
            self.pos.x = 0;
            self.max_h = 0;
            assert!(self.pos.y < self.texture_size);
        }
        self.pos.y = cmp::max(h, self.pos.y);
        let pos = self.pos;
        let size = Size2{w: w, h: h};
        if w * h != 0 {
            let mut data = Vec::from_elem((w * h) as uint * 4, 0 as u8);
            for y in range(0, h) {
                for x in range(0, w) {
                    let n = (x + y * w) as uint * 4;
                    *data.get_mut(n + 0) = 0;
                    *data.get_mut(n + 1) = 0;
                    *data.get_mut(n + 2) = 0;
                    *data.get_mut(n + 3) = bitmap[(x + y * w) as uint];
                }
            }
            self.texture.bind();
            self.texture.set_sub_image(pos, size, &data);
        }
        let xoff = if c == ' ' {
            let space_width = (self.size / 3.0) as MInt; // TODO: get from ttf
            xoff + space_width
        } else {
            xoff
        };
        self.pos.x += w;
        let glyph = Glyph {
            pos: pos,
            size: size,
            index: index,
            xoff: xoff,
            yoff: yoff,
        };
        if self.max_h < h - yoff {
            self.max_h = h - yoff;
        }
        self.glyphs.insert(c, glyph);
        glyph
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<commit_msg>FontSatsh: Extracted few funcs from add_glyph(self, char)<commit_after>// See LICENSE file for copyright and license details.

use std::cmp;
use collections::hashmap::HashMap;
use stb_tt;
use cgmath::vector::{Vector3, Vector2};
use core::types::{Size2, Point2, MInt};
use core::misc::add_quad_to_vec;
use visualizer::texture::Texture;
use visualizer::types::MFloat;
use visualizer::mesh::Mesh;
use visualizer::shader::Shader;

struct Glyph {
    pos: Vector2<MInt>,
    size: Size2<MInt>,
    index: MInt,
    xoff: MInt,
    yoff: MInt,
}

pub struct FontStash {
    size: MFloat,
    font: stb_tt::Font,
    texture: Texture,
    texture_size: MInt,
    pos: Point2<MInt>,
    glyphs: HashMap<char, Glyph>,
    max_h: MInt,
}

impl FontStash {
    pub fn new(font_path: &str, size: MFloat) -> FontStash {
        // TODO: use updateble mesh
        let texture_size = 1024;
        let font = stb_tt::Font::new(font_path, size);
        let texture = Texture::new_empty(
            Size2{w: texture_size, h: texture_size});
        FontStash {
            size: size,
            font: font,
            texture: texture,
            texture_size: texture_size,
            pos: Vector2{x: 0, y: 0},
            glyphs: HashMap::new(),
            max_h: 0,
        }
    }

    pub fn get_glyph(&mut self, c: char) -> Glyph {
        match self.glyphs.find(&c) {
            Some(&r) => r,
            None => self.add_glyph(c),
        }
    }

    pub fn get_size(&mut self) -> MFloat {
        self.size
    }

    pub fn get_mesh(&mut self, text: &str, shader: &Shader) -> Mesh {
        // TODO: create mesh in c-tor, update mesh vertices data here
        let mut vertex_data = Vec::new();
        let mut tex_data = Vec::new();
        let s = self.texture_size as MFloat;
        let mut i = 0.0;
        for c in text.chars() {
            let glyph = self.get_glyph(c);
            let w = glyph.size.w as MFloat;
            let h = glyph.size.h as MFloat;
            let x1 = glyph.pos.x as MFloat / s;
            let y1 = glyph.pos.y as MFloat / s;
            let x2 = x1 + w / s;
            let y2 = y1 + h / s;
            add_quad_to_vec(
                &mut tex_data,
                Vector2{x: x1, y: y1},
                Vector2{x: x1, y: y2},
                Vector2{x: x2, y: y2},
                Vector2{x: x2, y: y1},
            );
            let yoff = -glyph.yoff as MFloat;
            add_quad_to_vec(
                &mut vertex_data,
                Vector3{x: i, y: yoff, z: 0.0},
                Vector3{x: i, y: yoff - h, z: 0.0},
                Vector3{x: w + i, y: yoff - h, z: 0.0},
                Vector3{x: w + i, y: yoff, z: 0.0},
            );
            i += w + glyph.xoff as MFloat;
        }
        let mut mesh = Mesh::new(vertex_data.as_slice());
        mesh.set_texture(self.texture, tex_data.as_slice());
        mesh.prepare(shader);
        mesh
    }

    fn insert_image_to_cache(
        &mut self,
        pos: Vector2<MInt>,
        size: Size2<MInt>,
        bitmap: ~[u8]
    ) {
        let mut data = Vec::from_elem((size.w * size.h) as uint * 4, 0 as u8);
        for y in range(0, size.h) {
            for x in range(0, size.w) {
                let n = (x + y * size.w) as uint * 4;
                *data.get_mut(n + 0) = 0;
                *data.get_mut(n + 1) = 0;
                *data.get_mut(n + 2) = 0;
                *data.get_mut(n + 3) = bitmap[(x + y * size.w) as uint];
            }
        }
        self.texture.bind();
        self.texture.set_sub_image(pos, size, &data);
    }

    fn start_new_row(&mut self) {
        self.pos.y += self.max_h;
        self.pos.x = 0;
        self.max_h = 0;
        assert!(self.pos.y < self.texture_size);
    }

    fn add_glyph(&mut self, c: char) -> Glyph {
        assert!(self.glyphs.find(&c).is_none());
        let index = self.font.find_glyph_index(c);
        let (bitmap, w, h, xoff, yoff) = self.font.get_glyph(index);
        if self.pos.x + w > self.texture_size {
            self.start_new_row();
        }
        self.pos.y = cmp::max(h, self.pos.y);
        let pos = self.pos;
        let size = Size2{w: w, h: h};
        if w * h != 0 {
            self.insert_image_to_cache(pos, size, bitmap);
        }
        let xoff = if c == ' ' {
            let space_width = (self.size / 3.0) as MInt; // TODO: get from ttf
            xoff + space_width
        } else {
            xoff
        };
        self.pos.x += w;
        let glyph = Glyph {
            pos: pos,
            size: size,
            index: index,
            xoff: xoff,
            yoff: yoff,
        };
        if self.max_h < h - yoff {
            self.max_h = h - yoff;
        }
        self.glyphs.insert(c, glyph);
        glyph
    }
}

// vim: set tabstop=4 shiftwidth=4 softtabstop=4 expandtab:
<|endoftext|>


<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A handler and associated types for the install date action.
//!
//! The install date action uses various heuristics to detect the operating
//! system install date.

use std::time::{SystemTime, Duration};

use log::error;
use crate::session::{self, Session};

#[cfg(target_family = "unix")]
use std::{fs, path::Path};

/// A response type for the install date action.
struct Response {
    /// Install date of the operating system, or `None` if the attemps to
    /// obtain install date failed.
    time: Option<SystemTime>,
}

impl super::Response for Response {

    // TODO: Change to RDFDatetime when the client will be capable to send
    // "raw" strings without wrapping them into protobuf.
    const RDF_NAME: Option<&'static str> = Some("DataBlob");

    type Proto = u64;

    fn into_proto(self) -> Self::Proto {
        let time = match self.time {
            Some(time) => time,
            None => {
                error!("cannot get install date, all methods failed");
                std::time::UNIX_EPOCH
            },
        };
        let since_epoch = match time.duration_since(std::time::UNIX_EPOCH) {
            Ok(duration) => duration,
            Err(err) => {
                error!(
                    "install date is {} seconds earlier than Unix epoch",
                    err.duration().as_secs()
                );
                Duration::from_secs(0)
            },
        };
        since_epoch.as_secs()
    }
}

/// Iterates over files given in `iter`, checking their modification times.
///
/// If the function cannot get modification time from all the files (e. g.
/// they all don't exist), the return value is `None`. Otherwise, the
/// modification time of the first file is returned.
#[cfg(target_family = "unix")]
fn get_modified_time<P, I>(iter: I) -> Option<SystemTime>
where
    P: AsRef<Path>,
    I: Iterator<Item = P>,
{
    iter.filter_map(|path| fs::metadata(path).ok()?.modified().ok()).next()
}

/// Contains utilities to work with ext2/3/4 filesystems. Currently, it is used
/// to run `dumpe2fs`, parse its output and get filesystem creation time.
#[cfg(target_os = "linux")]
mod e2fs_utils {
    use std::io::{Cursor, Read, BufRead, BufReader};
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::time::{self, SystemTime};

    use chrono::prelude::*;
    use proc_mounts::MountList;

    /// Returns the file name of the device that is mounted as root file
    /// system.
    ///
    /// This function returns `None` in case of errors or when the root
    /// filesystem doesn't have an associated device file (e.g. we created
    /// a tmpfs filesystem and chrooted into it).
    fn get_root_device() -> Option<PathBuf> {
        let mount_list = MountList::new().ok()?;
        let device = mount_list.get_mount_by_dest("/")?.source.clone();
        if !device.starts_with("/") {
            // Not a device file name.
            return None;
        }
        Some(device)
    }

    /// Parses `dumpe2fs` output to get the filesystem creation time.
    ///
    /// This function returns `None` in case of errors.
    fn parse_creation_time_from_dumpe2fs<R>(output: R) -> Option<SystemTime>
    where
        R: Read,
    {
        let reader = BufReader::new(output);
        for line in reader.lines() {
            const FIELD_NAME: &'static str = "Filesystem created:";
            let line = line.ok()?;
            if !line.starts_with(FIELD_NAME) {
                continue;
            }
            let line = line[FIELD_NAME.len() + 1 ..].trim();
            let local_time = Local.datetime_from_str(line, "%c").ok()?;
            let utc_time = local_time.with_timezone(&Utc);
            let time_since_epoch =
                (utc_time - Utc.timestamp(0, 0)).to_std().ok()?;
            return Some(time::UNIX_EPOCH + time_since_epoch);
        }
        None
    }

    /// Launches `dumpe2fs`, captures its output and obtains filesystem
    /// creation time.
    ///
    /// This function returns `None` in case of errors.
    pub fn creation_time_from_dumpe2fs() -> Option<SystemTime> {
        let output = Command::new("dumpe2fs")
            .arg(get_root_device()?)
            .stdout(Stdio::piped())
            .spawn().ok()?
            .wait_with_output().ok()?;
        if !output.status.success() {
            return None;
        }
        parse_creation_time_from_dumpe2fs(Cursor::new(output.stdout))
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use std::io::Cursor;

        #[test]
        fn test_parse_creation_time() {
            let correct = DateTime::parse_from_rfc3339("2020-05-31T19:02:03Z")
                .unwrap().with_timezone(&Utc);
            let creation_time = correct.with_timezone(&Local)
                .format("%c").to_string();
            let input = String::from(
                "Line 1\n\
                 Line 2\n\
                 Filesystem created:    "
            ) + &creation_time + "\n";
            let cursor = Cursor::new(input.as_bytes());
            let parsed = parse_creation_time_from_dumpe2fs(cursor).unwrap();
            let parsed = DateTime::<Utc>::from(parsed);
            assert_eq!(parsed, correct);
        }
    }
}

/// Obtains system install date (Linux version).
///
/// To obtain install date, it first checks the creation time of the root
/// filesystem. Then, it searches for the installation logs and gets their
/// modification date. It also checks for other files that were probably
/// unchanged from the installation.
///
/// This function returns `None` in case of errors.
#[cfg(target_os = "linux")]
fn get_install_time() -> Option<SystemTime> {
    // First, check the creation time of the root. Rust implementation of
    // `Metadata::created()` on Linux utilizes `statx` syscall, so this
    // method will work only on kernels >= 4.11.
    let time = fs::metadata("/").and_then(|metadata| metadata.created()).ok();
    if time.is_some() {
        return time;
    }

    // Then, try to detect filesystem creation time using dumpe2fs. This
    // works only on ext2/3/4 filesystems.
    let time = e2fs_utils::creation_time_from_dumpe2fs();
    if time.is_some() {
        return time;
    }

    // Then, search for various files that were potentially modified only
    // on installation:
    // * `/var/log/installer/`: The location of installation logs for
    //   Debian-based distributions.
    // * `/root/install.log`: The location of installation logs for RHEL-based
    //   distributions.
    // * `/etc/hostname`: In most cases, we may assume that the hostname didn't
    //   change after installation.
    // * `/lost+found`: This method was used by Python version on GRR client, so
    //   leaving it here.
    const CANDIDATES: [&str; 5] = [
        "/var/log/installer/syslog",
        "/var/log/installer/status",
        "/root/install.log",
        "/etc/hostname",
        "/lost+found",
    ];
    let time = get_modified_time(CANDIDATES.iter());
    if time.is_some() {
        return time;
    }

    // We tried our best and all the methods have failed, so just give up.
    None
}

/// Obtains system install date (OS X version).
///
/// This function searches for some files that were probably unchanged since
/// the installation.
///
/// This function returns `None` in case of errors.
#[cfg(target_os = "macos")]
fn get_install_time() -> Option<SystemTime> {
    // Here, we use the same way as Python version of GRR client does. We just
    // check the modification time for some of the paths.
    const CANDIDATES: [&str; 3] = [
        "/var/log/CDIS.custom",
        "/var",
        "/private",
    ];
    get_modified_time(CANDIDATES.iter())
}

/// Obtains system install date (Windows version).
///
/// To obtain install date, this function searches the corresponding value in
/// Windows registry.
///
/// This function returns `None` in case of errors.
#[cfg(target_os = "windows")]
fn get_install_time() -> Option<SystemTime> {
    use winreg::RegKey;

    // Don't use winreg::enums::KEY_WOW64_64KEY since it breaks on Windows 2000.
    let hklm = RegKey::predef(winreg::enums::HKEY_LOCAL_MACHINE);
    let install_time = hklm
        .open_subkey("Software\\Microsoft\\Windows NT\\CurrentVersion").ok()?
        .get_value::<u32, _>("InstallDate").ok()?;
    Some(SystemTime::UNIX_EPOCH + Duration::from_secs(install_time.into()))
}

/// Handles requests for the install date action.
pub fn handle<S: Session>(session: &mut S, _: ()) -> session::Result<()> {
    session.reply(Response {time: get_install_time()})?;
    Ok(())
}

#[cfg(test)]
mod tests {

    use super::*;
    use humantime::Timestamp;

    #[test]
    fn test_install_date() {
        let mut session = session::test::Fake::new();
        assert!(handle(&mut session, ()).is_ok());
        assert_eq!(session.reply_count(), 1);
        let response: &Response = session.reply(0);
        let time = response.time.unwrap();

        assert!(time <= SystemTime::now());
    }
}
<commit_msg>Remove unnecessary use directive.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

//! A handler and associated types for the install date action.
//!
//! The install date action uses various heuristics to detect the operating
//! system install date.

use std::time::{SystemTime, Duration};

use log::error;
use crate::session::{self, Session};

#[cfg(target_family = "unix")]
use std::{fs, path::Path};

/// A response type for the install date action.
struct Response {
    /// Install date of the operating system, or `None` if the attemps to
    /// obtain install date failed.
    time: Option<SystemTime>,
}

impl super::Response for Response {

    // TODO: Change to RDFDatetime when the client will be capable to send
    // "raw" strings without wrapping them into protobuf.
    const RDF_NAME: Option<&'static str> = Some("DataBlob");

    type Proto = u64;

    fn into_proto(self) -> Self::Proto {
        let time = match self.time {
            Some(time) => time,
            None => {
                error!("cannot get install date, all methods failed");
                std::time::UNIX_EPOCH
            },
        };
        let since_epoch = match time.duration_since(std::time::UNIX_EPOCH) {
            Ok(duration) => duration,
            Err(err) => {
                error!(
                    "install date is {} seconds earlier than Unix epoch",
                    err.duration().as_secs()
                );
                Duration::from_secs(0)
            },
        };
        since_epoch.as_secs()
    }
}

/// Iterates over files given in `iter`, checking their modification times.
///
/// If the function cannot get modification time from all the files (e. g.
/// they all don't exist), the return value is `None`. Otherwise, the
/// modification time of the first file is returned.
#[cfg(target_family = "unix")]
fn get_modified_time<P, I>(iter: I) -> Option<SystemTime>
where
    P: AsRef<Path>,
    I: Iterator<Item = P>,
{
    iter.filter_map(|path| fs::metadata(path).ok()?.modified().ok()).next()
}

/// Contains utilities to work with ext2/3/4 filesystems. Currently, it is used
/// to run `dumpe2fs`, parse its output and get filesystem creation time.
#[cfg(target_os = "linux")]
mod e2fs_utils {
    use std::io::{Cursor, Read, BufRead, BufReader};
    use std::path::PathBuf;
    use std::process::{Command, Stdio};
    use std::time::{self, SystemTime};

    use chrono::prelude::*;
    use proc_mounts::MountList;

    /// Returns the file name of the device that is mounted as root file
    /// system.
    ///
    /// This function returns `None` in case of errors or when the root
    /// filesystem doesn't have an associated device file (e.g. we created
    /// a tmpfs filesystem and chrooted into it).
    fn get_root_device() -> Option<PathBuf> {
        let mount_list = MountList::new().ok()?;
        let device = mount_list.get_mount_by_dest("/")?.source.clone();
        if !device.starts_with("/") {
            // Not a device file name.
            return None;
        }
        Some(device)
    }

    /// Parses `dumpe2fs` output to get the filesystem creation time.
    ///
    /// This function returns `None` in case of errors.
    fn parse_creation_time_from_dumpe2fs<R>(output: R) -> Option<SystemTime>
    where
        R: Read,
    {
        let reader = BufReader::new(output);
        for line in reader.lines() {
            const FIELD_NAME: &'static str = "Filesystem created:";
            let line = line.ok()?;
            if !line.starts_with(FIELD_NAME) {
                continue;
            }
            let line = line[FIELD_NAME.len() + 1 ..].trim();
            let local_time = Local.datetime_from_str(line, "%c").ok()?;
            let utc_time = local_time.with_timezone(&Utc);
            let time_since_epoch =
                (utc_time - Utc.timestamp(0, 0)).to_std().ok()?;
            return Some(time::UNIX_EPOCH + time_since_epoch);
        }
        None
    }

    /// Launches `dumpe2fs`, captures its output and obtains filesystem
    /// creation time.
    ///
    /// This function returns `None` in case of errors.
    pub fn creation_time_from_dumpe2fs() -> Option<SystemTime> {
        let output = Command::new("dumpe2fs")
            .arg(get_root_device()?)
            .stdout(Stdio::piped())
            .spawn().ok()?
            .wait_with_output().ok()?;
        if !output.status.success() {
            return None;
        }
        parse_creation_time_from_dumpe2fs(Cursor::new(output.stdout))
    }

    #[cfg(test)]
    mod tests {
        use super::*;
        use std::io::Cursor;

        #[test]
        fn test_parse_creation_time() {
            let correct = DateTime::parse_from_rfc3339("2020-05-31T19:02:03Z")
                .unwrap().with_timezone(&Utc);
            let creation_time = correct.with_timezone(&Local)
                .format("%c").to_string();
            let input = String::from(
                "Line 1\n\
                 Line 2\n\
                 Filesystem created:    "
            ) + &creation_time + "\n";
            let cursor = Cursor::new(input.as_bytes());
            let parsed = parse_creation_time_from_dumpe2fs(cursor).unwrap();
            let parsed = DateTime::<Utc>::from(parsed);
            assert_eq!(parsed, correct);
        }
    }
}

/// Obtains system install date (Linux version).
///
/// To obtain install date, it first checks the creation time of the root
/// filesystem. Then, it searches for the installation logs and gets their
/// modification date. It also checks for other files that were probably
/// unchanged from the installation.
///
/// This function returns `None` in case of errors.
#[cfg(target_os = "linux")]
fn get_install_time() -> Option<SystemTime> {
    // First, check the creation time of the root. Rust implementation of
    // `Metadata::created()` on Linux utilizes `statx` syscall, so this
    // method will work only on kernels >= 4.11.
    let time = fs::metadata("/").and_then(|metadata| metadata.created()).ok();
    if time.is_some() {
        return time;
    }

    // Then, try to detect filesystem creation time using dumpe2fs. This
    // works only on ext2/3/4 filesystems.
    let time = e2fs_utils::creation_time_from_dumpe2fs();
    if time.is_some() {
        return time;
    }

    // Then, search for various files that were potentially modified only
    // on installation:
    // * `/var/log/installer/`: The location of installation logs for
    //   Debian-based distributions.
    // * `/root/install.log`: The location of installation logs for RHEL-based
    //   distributions.
    // * `/etc/hostname`: In most cases, we may assume that the hostname didn't
    //   change after installation.
    // * `/lost+found`: This method was used by Python version on GRR client, so
    //   leaving it here.
    const CANDIDATES: [&str; 5] = [
        "/var/log/installer/syslog",
        "/var/log/installer/status",
        "/root/install.log",
        "/etc/hostname",
        "/lost+found",
    ];
    let time = get_modified_time(CANDIDATES.iter());
    if time.is_some() {
        return time;
    }

    // We tried our best and all the methods have failed, so just give up.
    None
}

/// Obtains system install date (OS X version).
///
/// This function searches for some files that were probably unchanged since
/// the installation.
///
/// This function returns `None` in case of errors.
#[cfg(target_os = "macos")]
fn get_install_time() -> Option<SystemTime> {
    // Here, we use the same way as Python version of GRR client does. We just
    // check the modification time for some of the paths.
    const CANDIDATES: [&str; 3] = [
        "/var/log/CDIS.custom",
        "/var",
        "/private",
    ];
    get_modified_time(CANDIDATES.iter())
}

/// Obtains system install date (Windows version).
///
/// To obtain install date, this function searches the corresponding value in
/// Windows registry.
///
/// This function returns `None` in case of errors.
#[cfg(target_os = "windows")]
fn get_install_time() -> Option<SystemTime> {
    use winreg::RegKey;

    // Don't use winreg::enums::KEY_WOW64_64KEY since it breaks on Windows 2000.
    let hklm = RegKey::predef(winreg::enums::HKEY_LOCAL_MACHINE);
    let install_time = hklm
        .open_subkey("Software\\Microsoft\\Windows NT\\CurrentVersion").ok()?
        .get_value::<u32, _>("InstallDate").ok()?;
    Some(SystemTime::UNIX_EPOCH + Duration::from_secs(install_time.into()))
}

/// Handles requests for the install date action.
pub fn handle<S: Session>(session: &mut S, _: ()) -> session::Result<()> {
    session.reply(Response {time: get_install_time()})?;
    Ok(())
}

#[cfg(test)]
mod tests {

    use super::*;

    #[test]
    fn test_install_date() {
        let mut session = session::test::Fake::new();
        assert!(handle(&mut session, ()).is_ok());
        assert_eq!(session.reply_count(), 1);
        let response: &Response = session.reply(0);
        let time = response.time.unwrap();

        assert!(time <= SystemTime::now());
    }
}
<|endoftext|>