<commit_before>extern crate protobuf;
extern crate rand;
extern crate time;

use std::default::Default;
use std::fs::File;
use std::path::Path;

use rand::Rng;
use rand::StdRng;
use rand::SeedableRng;

use protobuf::Message;
use protobuf::MessageStatic;
use perftest_data::PerftestData;

mod perftest_data;

fn measure_ns<R, F: FnMut() -> R>(mut f: F) -> (u64, R) {
    let start = time::precise_time_ns();
    let r = f();
    (time::precise_time_ns() - start, r)
}

fn measure_and_print<R, F: FnMut() -> R>(title: &str, iter: u64, f: F) -> R {
    let (ns, r) = measure_ns(f);
    let ns_per_iter = ns / iter;
    println!("{}: {} ns per iter", title, ns_per_iter);
    r
}

fn run_test<M : Message + MessageStatic>(name: &str, data: &[M]) {
    let mut rng: StdRng = SeedableRng::from_seed(&[10, 20, 30, 40][..]);
    let mut random_data: Vec<M> = Vec::new();

    let mut total_size = 0;
    while total_size < 1000000 {
        let ref item = data[rng.gen_range(0, data.len())];
        random_data.push(item.clone());
        total_size += item.compute_size();
    }

    let mut buf = Vec::new();
    measure_and_print(&format!("{}: write", name), random_data.len() as u64, || {
        for m in random_data.iter() {
            m.write_length_delimited_to_writer(&mut buf).unwrap();
        }
    });

    let read_data = measure_and_print(&format!("{}: read", name), random_data.len() as u64, || {
        let mut r = Vec::new();
        let mut coded_input_stream = protobuf::CodedInputStream::from_bytes(&buf);
        while !coded_input_stream.eof().unwrap() {
            r.push(protobuf::parse_length_delimited_from::<M>(&mut coded_input_stream).unwrap());
        }
        r
    });

    assert_eq!(random_data, read_data);

    let merged = measure_and_print(&format!("{}: read reuse", name), random_data.len() as u64, || {
        let mut coded_input_stream = protobuf::CodedInputStream::from_bytes(&buf);
        let mut msg: M = Default::default();
        let mut count = 0;
        while !coded_input_stream.eof().unwrap() {
            msg.clear();
            coded_input_stream.merge_message(&mut msg).unwrap();
            count += 1;
        }
        count
    });

    assert_eq!(random_data.len(), merged);
}

struct TestRunner {
    selected: Option<String>,
    any_matched: bool,
}

impl TestRunner {
    fn test<M : Message + MessageStatic>(&mut self, name: &str, data: &[M]) {
        if self.selected.as_ref().map(|s| *s == name).unwrap_or(true) {
            run_test(name, data);
            self.any_matched = true;
        }
    }

    fn check(&self) {
        if !self.any_matched {
            let name = self.selected.as_ref().map(|s| &s[..]).unwrap_or("bug");
            panic!("no tests found with name {}", name);
        }
    }
}

fn main() {
    let args = std::env::args().collect::<Vec<_>>();
    if args.len() > 2 { panic!("usage: {} <test>", args[0]) }
    let selected = args.into_iter().nth(1);

    let mut runner = TestRunner { selected: selected, any_matched: false };

    let mut is = File::open(&Path::new("perftest_data.pbbin")).unwrap();
    let test_data = protobuf::parse_from_reader::<PerftestData>(&mut is).unwrap();

    runner.test("test1", test_data.get_test1());
    runner.test("test_repeated_bool", test_data.get_test_repeated_bool());
    runner.test("test_repeated_packed_int32", test_data.get_test_repeated_packed_int32());
    runner.test("test_repeated_messages", test_data.get_test_repeated_messages());
    runner.test("test_optional_messages", test_data.get_test_optional_messages());
    runner.test("test_strings", test_data.get_test_strings());
    runner.test("test_small_bytearrays", test_data.get_test_small_bytearrays());
    runner.test("test_large_bytearrays", test_data.get_test_large_bytearrays());
    runner.check();
}
<commit_msg>better assertion in perftest<commit_after>extern crate protobuf;
extern crate rand;
extern crate time;

use std::default::Default;
use std::fs::File;
use std::path::Path;

use rand::Rng;
use rand::StdRng;
use rand::SeedableRng;

use protobuf::Message;
use protobuf::MessageStatic;
use perftest_data::PerftestData;

mod perftest_data;

fn measure_ns<R, F: FnMut() -> R>(mut f: F) -> (u64, R) {
    let start = time::precise_time_ns();
    let r = f();
    (time::precise_time_ns() - start, r)
}

fn measure_and_print<R, F: FnMut() -> R>(title: &str, iter: u64, f: F) -> R {
    let (ns, r) = measure_ns(f);
    let ns_per_iter = ns / iter;
    println!("{}: {} ns per iter", title, ns_per_iter);
    r
}

fn run_test<M : Message + MessageStatic>(name: &str, data: &[M]) {
    assert!(data.len() > 0, "empty string for test: {}", name);

    let mut rng: StdRng = SeedableRng::from_seed(&[10, 20, 30, 40][..]);
    let mut random_data: Vec<M> = Vec::new();

    let mut total_size = 0;
    while total_size < 1000000 {
        let ref item = data[rng.gen_range(0, data.len())];
        random_data.push(item.clone());
        total_size += item.compute_size();
    }

    let mut buf = Vec::new();
    measure_and_print(&format!("{}: write", name), random_data.len() as u64, || {
        for m in random_data.iter() {
            m.write_length_delimited_to_writer(&mut buf).unwrap();
        }
    });

    let read_data = measure_and_print(&format!("{}: read", name), random_data.len() as u64, || {
        let mut r = Vec::new();
        let mut coded_input_stream = protobuf::CodedInputStream::from_bytes(&buf);
        while !coded_input_stream.eof().unwrap() {
            r.push(protobuf::parse_length_delimited_from::<M>(&mut coded_input_stream).unwrap());
        }
        r
    });

    assert_eq!(random_data, read_data);

    let merged = measure_and_print(&format!("{}: read reuse", name), random_data.len() as u64, || {
        let mut coded_input_stream = protobuf::CodedInputStream::from_bytes(&buf);
        let mut msg: M = Default::default();
        let mut count = 0;
        while !coded_input_stream.eof().unwrap() {
            msg.clear();
            coded_input_stream.merge_message(&mut msg).unwrap();
            count += 1;
        }
        count
    });

    assert_eq!(random_data.len(), merged);
}

struct TestRunner {
    selected: Option<String>,
    any_matched: bool,
}

impl TestRunner {
    fn test<M : Message + MessageStatic>(&mut self, name: &str, data: &[M]) {
        if self.selected.as_ref().map(|s| *s == name).unwrap_or(true) {
            run_test(name, data);
            self.any_matched = true;
        }
    }

    fn check(&self) {
        if !self.any_matched {
            let name = self.selected.as_ref().map(|s| &s[..]).unwrap_or("bug");
            panic!("no tests found with name {}", name);
        }
    }
}

fn main() {
    let args = std::env::args().collect::<Vec<_>>();
    if args.len() > 2 { panic!("usage: {} <test>", args[0]) }
    let selected = args.into_iter().nth(1);

    let mut runner = TestRunner { selected: selected, any_matched: false };

    let mut is = File::open(&Path::new("perftest_data.pbbin")).unwrap();
    let test_data = protobuf::parse_from_reader::<PerftestData>(&mut is).unwrap();

    runner.test("test1", test_data.get_test1());
    runner.test("test_repeated_bool", test_data.get_test_repeated_bool());
    runner.test("test_repeated_packed_int32", test_data.get_test_repeated_packed_int32());
    runner.test("test_repeated_messages", test_data.get_test_repeated_messages());
    runner.test("test_optional_messages", test_data.get_test_optional_messages());
    runner.test("test_strings", test_data.get_test_strings());
    runner.test("test_small_bytearrays", test_data.get_test_small_bytearrays());
    runner.test("test_large_bytearrays", test_data.get_test_large_bytearrays());
    runner.check();
}
<|endoftext|>


<commit_before>use std::{mem, sync::Arc};

use bifrost::{conshash::ConsistentHashing, raft::client::RaftClient, rpc::RPCError};
use dovahkiin::{expr::serde::Expr, types::OwnedValue};
use futures::stream::{FuturesUnordered, StreamExt};
use itertools::Itertools;

use crate::{
    client::{client_by_server_id, client_by_server_name},
    index::{
        ranged::{client::cursor::ClientCursor, lsm::btree::Ordering},
        EntryKey, IndexerClients, SCHEMA_SCAN_PATT_SIZE,
    },
    ram::cell::{OwnedCell, ReadError},
};

const SCAN_BUFFER_SIZE: u16 = 64;

pub struct IndexedDataClient {
    conshash: Arc<ConsistentHashing>,
    index_clients: Arc<IndexerClients>,
}

pub struct DataCursor<'a> {
    index_cursor: Option<ClientCursor>,
    buffer: Vec<OwnedCell>,
    projection: Vec<u64>,
    selection: Expr,
    proc: Expr,
    client: &'a IndexedDataClient,
    pos: usize,
}

impl IndexedDataClient {
    pub fn new(conshash: &Arc<ConsistentHashing>, raft_client: &Arc<RaftClient>) -> Self {
        Self {
            conshash: conshash.clone(),
            index_clients: Arc::new(IndexerClients::new(conshash, raft_client)),
        }
    }
    pub async fn range_index_scan<'a>(
        &self,
        schema: u32,
        field: u64,
        key: u64,
        selection: OwnedValue,
        projection: OwnedValue,
    ) -> DataCursor<'a> {
        unimplemented!()
    }
    pub async fn scan_all<'a>(
        &'a self,
        schema: u32,
        projection: Vec<u64>, // Column array
        selection: Expr,      // Checker expression
        proc: Expr,
        ordering: Ordering,
    ) -> Result<DataCursor<'a>, RPCError> {
        let key = EntryKey::for_schema(schema);
        let index_cursor = self
            .index_clients
            .range_seek(
                &key,
                ordering,
                SCAN_BUFFER_SIZE,
                Some(SCHEMA_SCAN_PATT_SIZE),
            )
            .await?;
        let mut cursor = DataCursor {
            index_cursor,
            projection,
            selection,
            proc,
            client: self,
            buffer: vec![],
            pos: 0,
        };
        cursor.refresh_batch().await;
        Ok(cursor)
    }
}

impl<'a> DataCursor<'a> {
    pub async fn next(&mut self) -> Result<Option<OwnedCell>, RPCError> {
        if self.buffer.len() <= self.pos {
            if self.next_block().await? {
                if !self.refresh_batch().await {
                    return Ok(None);
                }
            } else {
                return Ok(None);
            }
        }
        if self.buffer.is_empty() {
            return Ok(None);
        }
        let cell = mem::take(&mut self.buffer[self.pos]);
        self.pos += 1;
        return Ok(Some(cell));
    }

    pub async fn next_block(&mut self) -> Result<bool, RPCError> {
        if let Some(cursor) = &mut self.index_cursor {
            if cursor.next_block().await? {
                return Ok(true);
            }
        }
        // If cannot get next block, set the index cursor to none
        // We are done with this cursor
        self.index_cursor = None;
        self.buffer = vec![];
        self.pos = 0;
        return Ok(false);
    }

    pub async fn refresh_batch(&mut self) -> bool {
        if let Some(cursor) = &self.index_cursor {
            let all_ids = cursor.current_block();
            let mut tasks = all_ids
                .iter()
                .enumerate()
                .filter_map(|(i, id)| {
                    self.client
                        .conshash
                        .get_server_id_by(id)
                        .map(|sid| (i, sid, id))
                })
                .group_by(|(_i, sid, _id)| *sid)
                .into_iter()
                .map(|(sid, pairs)| {
                    let mut ids = vec![];
                    let mut idx = vec![];
                    for (i, _, id) in pairs {
                        idx.push(i);
                        ids.push(*id);
                    }
                    let projection = self.projection.clone();
                    let selection = self.selection.clone();
                    let proc = self.proc.clone();
                    let server_name = self.client.conshash.to_server_name(sid);
                    async move {
                        match client_by_server_name(sid, server_name).await {
                            Ok(client) => {
                                let read_res = client
                                    .read_all_cells_proced(ids, projection, selection, proc)
                                    .await
                                    .map(|v| v.into_iter().zip(idx).collect_vec());
                                match read_res {
                                    Ok(cells) => Ok(cells
                                        .into_iter()
                                        .filter_map(|(c, i)| c.ok().map(|c| (c, i)))
                                        .collect_vec()),
                                    Err(e) => return Err(e),
                                }
                            }
                            Err(e) => return Err(e),
                        }
                    }
                })
                .collect::<FuturesUnordered<_>>();
            let mut all_cells = vec![];
            while let Some(task_res) = tasks.next().await {
                if let Ok(mut cells) = task_res {
                    all_cells.append(&mut cells);
                }
            }
            all_cells.sort_by(|(_, i1), (_, i2)| i1.cmp(i2));
            self.buffer = all_cells.into_iter().map(|(c, _)| c).collect_vec();
            self.pos = 0;
            true
        } else {
            self.buffer = vec![];
            self.pos = 0;
            false
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        index::ranged::lsm::btree::Ordering,
        ram::{
            cell::OwnedCell,
            schema::{Field, IndexType, Schema},
        },
        server::*,
    };
    use bifrost_hasher::hash_str;
    use dovahkiin::{expr::serde::Expr, types::*, integrated::lisp::*};

    #[tokio::test(flavor = "multi_thread")]
    async fn scan_all() {
        const DATA_1: &'static str = "DATA_1";
        const DATA_2: &'static str = "DATA_2";
        let _ = env_logger::try_init();
        let server_addr = String::from("127.0.0.1:6701");
        let server_group = String::from("indexed_scan_all_test");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 8,
                memory_size: 512 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
                index_enabled: true,
                services: vec![Service::Cell, Service::Query],
            },
            &server_addr,
            &server_group,
        )
        .await;
        // Require schema to be scannable to insert special scan key to the range indexer
        let fields = Field::new(
            "*",
            Type::Map,
            false,
            false,
            Some(vec![
                Field::new(
                    DATA_1,
                    Type::U64,
                    false,
                    false,
                    None,
                    vec![IndexType::Ranged],
                ),
                Field::new(DATA_2, Type::U32, false, false, None, vec![]),
            ]),
            vec![],
        );
        let schema_id_1 = 123;
        let schema_id_2 = 234;
        let schema_1 = Schema::new_with_id(
            schema_id_1,
            &String::from("schema_1"),
            None,
            fields.clone(),
            false,
            true, // Scannable
        );
        let schema_2 = Schema::new_with_id(
            schema_id_2,
            &String::from("schema_2"),
            None,
            fields,
            false,
            true, // Scannable
        );
        let client = server.data_client(&vec![server_addr]).await.unwrap();
        client.new_schema_with_id(schema_1).await.unwrap().unwrap();
        client.new_schema_with_id(schema_2).await.unwrap().unwrap();
        let num = 1024;
        for i in 0..num {
            let id = Id::new(1, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 2) as u32);
            let cell = OwnedCell::new_with_id(schema_id_1, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let idx_data_client = server.indexed_data_client();
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell = cursor.next().await.unwrap().unwrap();
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        for i in 0..num {
            let id = Id::new(2, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 3) as u32);
            let cell = OwnedCell::new_with_id(schema_id_2, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_2,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(2, i);
            let cell = cursor.next().await.unwrap().unwrap();
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 3) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell = cursor.next().await.unwrap().unwrap();
            assert_eq!(id, cell.id());
        }
        // Testing selection
        let select_expr = parse_to_serde_expr("(and (>= DATA_1 10u64) (< DATA_1 100u64))").unwrap()[0].clone();
        let mut cursor = idx_data_client
        .scan_all(
            schema_id_1,
            vec![],
            select_expr,
            Expr::nothing(),
            Ordering::Forward,
        )
        .await
        .unwrap();
        // Start from 10 to 100 due to the selection expression
        for i in 10..100 {
            let id = Id::new(1, i);
            let cell = cursor.next().await.unwrap().unwrap();
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
    }
}
<commit_msg>formatting for presentation<commit_after>use std::{mem, sync::Arc};

use bifrost::{conshash::ConsistentHashing, raft::client::RaftClient, rpc::RPCError};
use dovahkiin::{expr::serde::Expr, types::OwnedValue};
use futures::stream::{FuturesUnordered, StreamExt};
use itertools::Itertools;

use crate::{
    client::{client_by_server_id, client_by_server_name},
    index::{
        ranged::{client::cursor::ClientCursor, lsm::btree::Ordering},
        EntryKey, IndexerClients, SCHEMA_SCAN_PATT_SIZE,
    },
    ram::cell::{OwnedCell, ReadError},
};

const SCAN_BUFFER_SIZE: u16 = 64;

pub struct IndexedDataClient {
    conshash: Arc<ConsistentHashing>,
    index_clients: Arc<IndexerClients>,
}

pub struct DataCursor<'a> {
    index_cursor: Option<ClientCursor>,
    buffer: Vec<OwnedCell>,
    projection: Vec<u64>,
    selection: Expr,
    proc: Expr,
    client: &'a IndexedDataClient,
    pos: usize,
}

impl IndexedDataClient {
    pub fn new(conshash: &Arc<ConsistentHashing>, raft_client: &Arc<RaftClient>) -> Self {
        Self {
            conshash: conshash.clone(),
            index_clients: Arc::new(IndexerClients::new(conshash, raft_client)),
        }
    }
    pub async fn range_index_scan<'a>(
        &self,
        schema: u32,
        field: u64,
        key: u64,
        selection: OwnedValue,
        projection: OwnedValue,
    ) -> DataCursor<'a> {
        unimplemented!()
    }
    pub async fn scan_all<'a>(
        &'a self,
        schema: u32,
        projection: Vec<u64>, // Column array
        selection: Expr,      // Checker expression
        proc: Expr,
        ordering: Ordering,
    ) -> Result<DataCursor<'a>, RPCError> {
        let key = EntryKey::for_schema(schema);
        let index_cursor = self
            .index_clients
            .range_seek(
                &key,
                ordering,
                SCAN_BUFFER_SIZE,
                Some(SCHEMA_SCAN_PATT_SIZE),
            )
            .await?;
        let mut cursor = DataCursor {
            index_cursor,
            projection,
            selection,
            proc,
            client: self,
            buffer: vec![],
            pos: 0,
        };
        cursor.refresh_batch().await;
        Ok(cursor)
    }
}

impl<'a> DataCursor<'a> {
    pub async fn next(&mut self) -> Result<Option<OwnedCell>, RPCError> {
        if self.buffer.len() <= self.pos {
            if self.next_block().await? {
                if !self.refresh_batch().await {
                    return Ok(None);
                }
            } else {
                return Ok(None);
            }
        }
        if self.buffer.is_empty() {
            return Ok(None);
        }
        let cell = mem::take(&mut self.buffer[self.pos]);
        self.pos += 1;
        return Ok(Some(cell));
    }

    pub async fn next_block(&mut self) -> Result<bool, RPCError> {
        if let Some(cursor) = &mut self.index_cursor {
            if cursor.next_block().await? {
                return Ok(true);
            }
        }
        // If cannot get next block, set the index cursor to none
        // We are done with this cursor
        self.index_cursor = None;
        self.buffer = vec![];
        self.pos = 0;
        return Ok(false);
    }

    pub async fn refresh_batch(&mut self) -> bool {
        if let Some(cursor) = &self.index_cursor {
            let all_ids = cursor.current_block();
            let mut tasks = all_ids
                .iter()
                .enumerate()
                .filter_map(|(i, id)| {
                    self.client
                        .conshash
                        .get_server_id_by(id)
                        .map(|sid| (i, sid, id))
                })
                .group_by(|(_i, sid, _id)| *sid)
                .into_iter()
                .map(|(sid, pairs)| {
                    let mut ids = vec![];
                    let mut idx = vec![];
                    for (i, _, id) in pairs {
                        idx.push(i);
                        ids.push(*id);
                    }
                    let projection = self.projection.clone();
                    let selection = self.selection.clone();
                    let proc = self.proc.clone();
                    let server_name = self.client.conshash.to_server_name(sid);
                    async move {
                        match client_by_server_name(sid, server_name).await {
                            Ok(client) => {
                                let read_res = client
                                    .read_all_cells_proced(ids, projection, selection, proc)
                                    .await
                                    .map(|v| v.into_iter().zip(idx).collect_vec());
                                match read_res {
                                    Ok(cells) => Ok(cells
                                        .into_iter()
                                        .filter_map(|(c, i)| c.ok().map(|c| (c, i)))
                                        .collect_vec()),
                                    Err(e) => return Err(e),
                                }
                            }
                            Err(e) => return Err(e),
                        }
                    }
                })
                .collect::<FuturesUnordered<_>>();
            let mut all_cells = vec![];
            while let Some(task_res) = tasks.next().await {
                if let Ok(mut cells) = task_res {
                    all_cells.append(&mut cells);
                }
            }
            all_cells.sort_by(|(_, i1), (_, i2)| i1.cmp(i2));
            self.buffer = all_cells.into_iter().map(|(c, _)| c).collect_vec();
            self.pos = 0;
            true
        } else {
            self.buffer = vec![];
            self.pos = 0;
            false
        }
    }
}

#[cfg(test)]
mod test {
    use crate::{
        index::ranged::lsm::btree::Ordering,
        ram::{
            cell::OwnedCell,
            schema::{Field, IndexType, Schema},
        },
        server::*,
    };
    use bifrost_hasher::hash_str;
    use dovahkiin::{expr::serde::Expr, integrated::lisp::*, types::*};

    #[tokio::test(flavor = "multi_thread")]
    async fn scan_all() {
        const DATA_1: &'static str = "DATA_1";
        const DATA_2: &'static str = "DATA_2";
        let _ = env_logger::try_init();
        let server_addr = String::from("127.0.0.1:6701");
        let server_group = String::from("indexed_scan_all_test");
        let server = NebServer::new_from_opts(
            &ServerOptions {
                chunk_count: 8,
                memory_size: 512 * 1024 * 1024,
                backup_storage: None,
                wal_storage: None,
                index_enabled: true,
                services: vec![Service::Cell, Service::Query],
            },
            &server_addr,
            &server_group,
        )
        .await;
        // Require schema to be scannable to insert special scan key to the range indexer
        let fields = Field::new(
            "*",
            Type::Map,
            false,
            false,
            Some(vec![
                Field::new(
                    DATA_1,
                    Type::U64,
                    false,
                    false,
                    None,
                    vec![IndexType::Ranged],
                ),
                Field::new(DATA_2, Type::U32, false, false, None, vec![]),
            ]),
            vec![],
        );
        let schema_id_1 = 123;
        let schema_id_2 = 234;
        let schema_1 = Schema::new_with_id(
            schema_id_1,
            &String::from("schema_1"),
            None,
            fields.clone(),
            false,
            true, // Scannable
        );
        let schema_2 = Schema::new_with_id(
            schema_id_2,
            &String::from("schema_2"),
            None,
            fields,
            false,
            true, // Scannable
        );
        let client = server.data_client(&vec![server_addr]).await.unwrap();
        client.new_schema_with_id(schema_1).await.unwrap().unwrap();
        client.new_schema_with_id(schema_2).await.unwrap().unwrap();
        let num = 1024;
        for i in 0..num {
            let id = Id::new(1, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 2) as u32);
            let cell = OwnedCell::new_with_id(schema_id_1, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        let idx_data_client = server.indexed_data_client();
        let mut cursor = idx_data_client
            .scan_all(
                schema_id_1,
                vec![],
                Expr::nothing(),
                Expr::nothing(),
                Ordering::Forward,
            )
            .await
            .unwrap();
        for i in 0..num {
            let id = Id::new(1, i);
            let cell = cursor.next().await.unwrap().unwrap();
            assert_eq!(id, cell.id());
            assert_eq!(*cell[DATA_1].u64().unwrap(), i);
            assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
            debug!("Checked cell id {:?} from index", id);
        }
        let out_of_range_item = cursor.next().await.unwrap();
        if let Some(cell) = out_of_range_item {
            panic!("Should not have any more cell. Got id {:?}", cell.id());
        }
        for i in 0..num {
            let id = Id::new(2, i);
            let mut value = OwnedValue::Map(OwnedMap::new());
            value[DATA_1] = OwnedValue::U64(i);
            value[DATA_2] = OwnedValue::U32((i * 3) as u32);
            let cell = OwnedCell::new_with_id(schema_id_2, &id, value);
            client.write_cell(cell).await.unwrap().unwrap();
        }
        {
            let mut cursor = idx_data_client
                .scan_all(
                    schema_id_2,
                    vec![],
                    Expr::nothing(),
                    Expr::nothing(),
                    Ordering::Forward,
                )
                .await
                .unwrap();
            for i in 0..num {
                let id = Id::new(2, i);
                let cell = cursor.next().await.unwrap().unwrap();
                assert_eq!(id, cell.id());
                assert_eq!(*cell[DATA_1].u64().unwrap(), i);
                assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 3) as u32);
                debug!("Checked cell id {:?} from index", id);
            }
            let out_of_range_item = cursor.next().await.unwrap();
            if let Some(cell) = out_of_range_item {
                panic!("Should not have any more cell. Got id {:?}", cell.id());
            }
        }
        {
            let mut cursor = idx_data_client
                .scan_all(
                    schema_id_1,
                    vec![],
                    Expr::nothing(),
                    Expr::nothing(),
                    Ordering::Forward,
                )
                .await
                .unwrap();
            for i in 0..num {
                let id = Id::new(1, i);
                let cell = cursor.next().await.unwrap().unwrap();
                assert_eq!(id, cell.id());
            }
        }
        {
            // Testing selection
            let select_expr = parse_to_serde_expr("(and (>= DATA_1 10u64) (< DATA_1 100u64))")
                .unwrap()[0]
                .clone();
            let mut cursor = idx_data_client
                .scan_all(
                    schema_id_1,
                    vec![],
                    select_expr,
                    Expr::nothing(),
                    Ordering::Forward,
                )
                .await
                .unwrap();
            // Start from 10 to 100 due to the selection expression
            for i in 10..100 {
                let id = Id::new(1, i);
                let cell = cursor.next().await.unwrap().unwrap();
                assert_eq!(id, cell.id());
                assert_eq!(*cell[DATA_1].u64().unwrap(), i);
                assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
                debug!("Checked cell id {:?} from index", id);
            }
            let out_of_range_item = cursor.next().await.unwrap();
            if let Some(cell) = out_of_range_item {
                panic!("Should not have any more cell. Got id {:?}", cell.id());
            }
        }
        {
            info!("Testing selection 2");
            let select_expr = parse_to_serde_expr("(= DATA_1 10u64)").unwrap()[0].clone();
            let mut cursor = idx_data_client
                .scan_all(
                    schema_id_1,
                    vec![],
                    select_expr,
                    Expr::nothing(),
                    Ordering::Forward,
                )
                .await
                .unwrap();
            // 10 and 500 due to the selection expression
            for i in vec![10] {
                let id = Id::new(1, i);
                let cell = cursor.next().await.unwrap().unwrap();
                assert_eq!(id, cell.id());
                assert_eq!(*cell[DATA_1].u64().unwrap(), i);
                assert_eq!(*cell[DATA_2].u32().unwrap(), (i * 2) as u32);
                debug!("-> Checked cell id {:?} from index", id);
            }
            let out_of_range_item = cursor.next().await.unwrap();
            if let Some(cell) = out_of_range_item {
                panic!("Should not have any more cell. Got id {:?}", cell.id());
            }
        }
    }
}
<|endoftext|>


<commit_before>//! Functions to read and write control registers.

pub use super::model_specific::{Efer, EferFlags};

use bitflags::bitflags;

/// Various control flags modifying the basic operation of the CPU.
#[derive(Debug)]
pub struct Cr0;

bitflags! {
    /// Configuration flags of the Cr0 register.
    pub struct Cr0Flags: u64 {
        /// Enables protected mode.
        const PROTECTED_MODE_ENABLE = 1;
        /// Enables monitoring of the coprocessor, typical for x87 instructions.
        ///
        /// Controls together with the `TASK_SWITCHED` flag whether a `wait` or `fwait`
        /// instruction should cause a device-not-available exception.
        const MONITOR_COPROCESSOR = 1 << 1;
        /// Force all x87 and MMX instructions to cause an exception.
        const EMULATE_COPROCESSOR = 1 << 2;
        /// Automatically set to 1 on _hardware_ task switch.
        ///
        /// This flags allows lazily saving x87/MMX/SSE instructions on hardware context switches.
        const TASK_SWITCHED = 1 << 3;
        /// Enables the native error reporting mechanism for x87 FPU errors.
        const NUMERIC_ERROR = 1 << 5;
        /// Controls whether supervisor-level writes to read-only pages are inhibited.
        ///
        /// When set, it is not possible to write to read-only pages from ring 0.
        const WRITE_PROTECT = 1 << 16;
        /// Enables automatic alignment checking.
        const ALIGNMENT_MASK = 1 << 18;
        /// Ignored. Used to control write-back/write-through cache strategy on older CPUs.
        const NOT_WRITE_THROUGH = 1 << 29;
        /// Disables internal caches (only for some cases).
        const CACHE_DISABLE = 1 << 30;
        /// Enables page translation.
        const PAGING = 1 << 31;
    }
}

/// Contains the Page Fault Linear Address (PFLA).
///
/// When page fault occurs, the CPU sets this register to the accessed address.
#[derive(Debug)]
pub struct Cr2;

/// Contains the physical address of the level 4 page table.
#[derive(Debug)]
pub struct Cr3;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr3Flags: u64 {
        /// Use a writethrough cache policy for the P4 table (else a writeback policy is used).
        const PAGE_LEVEL_WRITETHROUGH = 1 << 3;
        /// Disable caching for the P4 table.
        const PAGE_LEVEL_CACHE_DISABLE = 1 << 4;
    }
}

/// Various control flags modifying the basic operation of the CPU while in protected mode.
///
/// Note: The documention for the individual fields is taken from the AMD64 and Intel x86_64
/// manuals.
#[derive(Debug)]
pub struct Cr4;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr4Flags: u64 {
        /// Enables hardware-supported performance enhancements for software running in
        /// virtual-8086 mode.
        const VIRTUAL_8086_MODE_EXTENSIONS = 1;
        /// Enables support for protected-mode virtual interrupts.
        const PROTECTED_MODE_VIRTUAL_INTERRUPTS = 1 << 1;
        /// When set, only privilege-level 0 can execute the RDTSC or RDTSCP instructions.
        const TIMESTAMP_DISABLE = 1 << 2;
        /// Enables I/O breakpoint capability and enforces treatment of DR4 and DR5 registers
        /// as reserved.
        const DEBUGGING_EXTENSIONS = 1 << 3;
        /// Enables the use of 4MB physical frames; ignored in long mode.
        const PAGE_SIZE_EXTENSION = 1 << 4;
        /// Enables physical address extension and 2MB physical frames; required in long mode.
        const PHYSICAL_ADDRESS_EXTENSION = 1 << 5;
        /// Enables the machine-check exception mechanism.
        const MACHINE_CHECK_EXCEPTION = 1 << 6;
        /// Enables the global-page mechanism, which allows to make page translations global
        /// to all processes.
        const PAGE_GLOBAL = 1 << 7;
        /// Allows software running at any privilege level to use the RDPMC instruction.
        const PERFORMANCE_MONITOR_COUNTER = 1 << 8;
        /// Enable the use of legacy SSE instructions; allows using FXSAVE/FXRSTOR for saving
        /// processor state of 128-bit media instructions.
        const OSFXSR = 1 << 9;
        /// Enables the SIMD floating-point exception (#XF) for handling unmasked 256-bit and
        /// 128-bit media floating-point errors.
        const OSXMMEXCPT_ENABLE = 1 << 10;
        /// Prevents the execution of the SGDT, SIDT, SLDT, SMSW, and STR instructions by
        /// user-mode software.
        const USER_MODE_INSTRUCTION_PREVENTION = 1 << 11;
        /// Enables 5-level paging on supported CPUs.
        const L5_PAGING = 1 << 12;
        /// Enables VMX insturctions.
        const VIRTUAL_MACHINE_EXTENSIONS = 1 << 13;
        /// Enables SMX instructions.
        const SAFER_MODE_EXTENSIONS = 1 << 14;
        /// Enables software running in 64-bit mode at any privilege level to read and write
        /// the FS.base and GS.base hidden segment register state.
        const FSGSBASE = 1 << 16;
        /// Enables process-context identifiers (PCIDs).
        const PCID = 1 << 17;
        /// Enables extendet processor state management instructions, including XGETBV and XSAVE.
        const OSXSAVE = 1 << 18;
        /// Prevents the execution of instructions that reside in pages accessible by user-mode
        /// software when the processor is in supervisor-mode.
        const SUPERVISOR_MODE_EXECUTION_PROTECTION = 1 << 20;
        /// Enables restrictions for supervisor-mode software when reading data from user-mode
        /// pages.
        const SUPERVISOR_MODE_ACCESS_PREVENTION = 1 << 21;
        /// Enables 4-level paging to associate each linear address with a protection key.
        const PROTECTION_KEY = 1 << 22;
    }
}

#[cfg(target_arch = "x86_64")]
mod x86_64 {
    use super::*;
    use crate::structures::paging::PhysFrame;
    use crate::{PhysAddr, VirtAddr};

    impl Cr0 {
        /// Read the current set of CR0 flags.
        #[inline]
        pub fn read() -> Cr0Flags {
            Cr0Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR0 value.
        #[inline]
        pub fn read_raw() -> u64 {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr0, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr0();
            }

            value
        }

        /// Write CR0 flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by disabling paging.
        #[inline]
        pub unsafe fn write(flags: Cr0Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr0Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR0 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by disabling paging.
        #[inline]
        pub unsafe fn write_raw(value: u64) {
            #[cfg(feature = "inline_asm")]
            llvm_asm!("mov $0, %cr0" :: "r" (value) : "memory");

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_write_cr0(value);
        }

        /// Updates CR0 flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by disabling paging.
        #[inline]
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr0Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }

    impl Cr2 {
        /// Read the current page fault linear address from the CR3 register.
        #[inline]
        pub fn read() -> VirtAddr {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr2, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr2();
            }

            VirtAddr::new(value)
        }
    }

    impl Cr3 {
        /// Read the current P4 table address from the CR3 register.
        #[inline]
        pub fn read() -> (PhysFrame, Cr3Flags) {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr3, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr3();
            }

            let flags = Cr3Flags::from_bits_truncate(value);
            let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
            let frame = PhysFrame::containing_address(addr);
            (frame, flags)
        }

        /// Write a new P4 table address into the CR3 register.
        ///
        /// ## Safety
        /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
        /// changing the page mapping.
        #[inline]
        pub unsafe fn write(frame: PhysFrame, flags: Cr3Flags) {
            let addr = frame.start_address();
            let value = addr.as_u64() | flags.bits();

            #[cfg(feature = "inline_asm")]
            llvm_asm!("mov $0, %cr3" :: "r" (value) : "memory");

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_write_cr3(value)
        }
    }

    impl Cr4 {
        /// Read the current set of CR4 flags.
        #[inline]
        pub fn read() -> Cr4Flags {
            Cr4Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR4 value.
        #[inline]
        pub fn read_raw() -> u64 {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr4, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr4();
            }

            value
        }

        /// Write CR4 flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by overwriting the physical address extension
        /// flag.
        #[inline]
        pub unsafe fn write(flags: Cr4Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr4Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR4 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by overwriting the physical address extension
        /// flag.
        #[inline]
        pub unsafe fn write_raw(value: u64) {
            #[cfg(feature = "inline_asm")]
            llvm_asm!("mov $0, %cr4" :: "r" (value) : "memory");

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_write_cr4(value);
        }

        /// Updates CR4 flags.
        ///
        /// Preserves the value of reserved fields.
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by overwriting the physical address extension
        /// flag.
        #[inline]
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr4Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }
}
<commit_msg>Correct Cr2::read documentation (#161)<commit_after>//! Functions to read and write control registers.

pub use super::model_specific::{Efer, EferFlags};

use bitflags::bitflags;

/// Various control flags modifying the basic operation of the CPU.
#[derive(Debug)]
pub struct Cr0;

bitflags! {
    /// Configuration flags of the Cr0 register.
    pub struct Cr0Flags: u64 {
        /// Enables protected mode.
        const PROTECTED_MODE_ENABLE = 1;
        /// Enables monitoring of the coprocessor, typical for x87 instructions.
        ///
        /// Controls together with the `TASK_SWITCHED` flag whether a `wait` or `fwait`
        /// instruction should cause a device-not-available exception.
        const MONITOR_COPROCESSOR = 1 << 1;
        /// Force all x87 and MMX instructions to cause an exception.
        const EMULATE_COPROCESSOR = 1 << 2;
        /// Automatically set to 1 on _hardware_ task switch.
        ///
        /// This flags allows lazily saving x87/MMX/SSE instructions on hardware context switches.
        const TASK_SWITCHED = 1 << 3;
        /// Enables the native error reporting mechanism for x87 FPU errors.
        const NUMERIC_ERROR = 1 << 5;
        /// Controls whether supervisor-level writes to read-only pages are inhibited.
        ///
        /// When set, it is not possible to write to read-only pages from ring 0.
        const WRITE_PROTECT = 1 << 16;
        /// Enables automatic alignment checking.
        const ALIGNMENT_MASK = 1 << 18;
        /// Ignored. Used to control write-back/write-through cache strategy on older CPUs.
        const NOT_WRITE_THROUGH = 1 << 29;
        /// Disables internal caches (only for some cases).
        const CACHE_DISABLE = 1 << 30;
        /// Enables page translation.
        const PAGING = 1 << 31;
    }
}

/// Contains the Page Fault Linear Address (PFLA).
///
/// When page fault occurs, the CPU sets this register to the accessed address.
#[derive(Debug)]
pub struct Cr2;

/// Contains the physical address of the level 4 page table.
#[derive(Debug)]
pub struct Cr3;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr3Flags: u64 {
        /// Use a writethrough cache policy for the P4 table (else a writeback policy is used).
        const PAGE_LEVEL_WRITETHROUGH = 1 << 3;
        /// Disable caching for the P4 table.
        const PAGE_LEVEL_CACHE_DISABLE = 1 << 4;
    }
}

/// Various control flags modifying the basic operation of the CPU while in protected mode.
///
/// Note: The documention for the individual fields is taken from the AMD64 and Intel x86_64
/// manuals.
#[derive(Debug)]
pub struct Cr4;

bitflags! {
    /// Controls cache settings for the level 4 page table.
    pub struct Cr4Flags: u64 {
        /// Enables hardware-supported performance enhancements for software running in
        /// virtual-8086 mode.
        const VIRTUAL_8086_MODE_EXTENSIONS = 1;
        /// Enables support for protected-mode virtual interrupts.
        const PROTECTED_MODE_VIRTUAL_INTERRUPTS = 1 << 1;
        /// When set, only privilege-level 0 can execute the RDTSC or RDTSCP instructions.
        const TIMESTAMP_DISABLE = 1 << 2;
        /// Enables I/O breakpoint capability and enforces treatment of DR4 and DR5 registers
        /// as reserved.
        const DEBUGGING_EXTENSIONS = 1 << 3;
        /// Enables the use of 4MB physical frames; ignored in long mode.
        const PAGE_SIZE_EXTENSION = 1 << 4;
        /// Enables physical address extension and 2MB physical frames; required in long mode.
        const PHYSICAL_ADDRESS_EXTENSION = 1 << 5;
        /// Enables the machine-check exception mechanism.
        const MACHINE_CHECK_EXCEPTION = 1 << 6;
        /// Enables the global-page mechanism, which allows to make page translations global
        /// to all processes.
        const PAGE_GLOBAL = 1 << 7;
        /// Allows software running at any privilege level to use the RDPMC instruction.
        const PERFORMANCE_MONITOR_COUNTER = 1 << 8;
        /// Enable the use of legacy SSE instructions; allows using FXSAVE/FXRSTOR for saving
        /// processor state of 128-bit media instructions.
        const OSFXSR = 1 << 9;
        /// Enables the SIMD floating-point exception (#XF) for handling unmasked 256-bit and
        /// 128-bit media floating-point errors.
        const OSXMMEXCPT_ENABLE = 1 << 10;
        /// Prevents the execution of the SGDT, SIDT, SLDT, SMSW, and STR instructions by
        /// user-mode software.
        const USER_MODE_INSTRUCTION_PREVENTION = 1 << 11;
        /// Enables 5-level paging on supported CPUs.
        const L5_PAGING = 1 << 12;
        /// Enables VMX insturctions.
        const VIRTUAL_MACHINE_EXTENSIONS = 1 << 13;
        /// Enables SMX instructions.
        const SAFER_MODE_EXTENSIONS = 1 << 14;
        /// Enables software running in 64-bit mode at any privilege level to read and write
        /// the FS.base and GS.base hidden segment register state.
        const FSGSBASE = 1 << 16;
        /// Enables process-context identifiers (PCIDs).
        const PCID = 1 << 17;
        /// Enables extendet processor state management instructions, including XGETBV and XSAVE.
        const OSXSAVE = 1 << 18;
        /// Prevents the execution of instructions that reside in pages accessible by user-mode
        /// software when the processor is in supervisor-mode.
        const SUPERVISOR_MODE_EXECUTION_PROTECTION = 1 << 20;
        /// Enables restrictions for supervisor-mode software when reading data from user-mode
        /// pages.
        const SUPERVISOR_MODE_ACCESS_PREVENTION = 1 << 21;
        /// Enables 4-level paging to associate each linear address with a protection key.
        const PROTECTION_KEY = 1 << 22;
    }
}

#[cfg(target_arch = "x86_64")]
mod x86_64 {
    use super::*;
    use crate::structures::paging::PhysFrame;
    use crate::{PhysAddr, VirtAddr};

    impl Cr0 {
        /// Read the current set of CR0 flags.
        #[inline]
        pub fn read() -> Cr0Flags {
            Cr0Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR0 value.
        #[inline]
        pub fn read_raw() -> u64 {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr0, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr0();
            }

            value
        }

        /// Write CR0 flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by disabling paging.
        #[inline]
        pub unsafe fn write(flags: Cr0Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr0Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR0 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by disabling paging.
        #[inline]
        pub unsafe fn write_raw(value: u64) {
            #[cfg(feature = "inline_asm")]
            llvm_asm!("mov $0, %cr0" :: "r" (value) : "memory");

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_write_cr0(value);
        }

        /// Updates CR0 flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by disabling paging.
        #[inline]
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr0Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }

    impl Cr2 {
        /// Read the current page fault linear address from the CR2 register.
        #[inline]
        pub fn read() -> VirtAddr {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr2, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr2();
            }

            VirtAddr::new(value)
        }
    }

    impl Cr3 {
        /// Read the current P4 table address from the CR3 register.
        #[inline]
        pub fn read() -> (PhysFrame, Cr3Flags) {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr3, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr3();
            }

            let flags = Cr3Flags::from_bits_truncate(value);
            let addr = PhysAddr::new(value & 0x_000f_ffff_ffff_f000);
            let frame = PhysFrame::containing_address(addr);
            (frame, flags)
        }

        /// Write a new P4 table address into the CR3 register.
        ///
        /// ## Safety
        /// Changing the level 4 page table is unsafe, because it's possible to violate memory safety by
        /// changing the page mapping.
        #[inline]
        pub unsafe fn write(frame: PhysFrame, flags: Cr3Flags) {
            let addr = frame.start_address();
            let value = addr.as_u64() | flags.bits();

            #[cfg(feature = "inline_asm")]
            llvm_asm!("mov $0, %cr3" :: "r" (value) : "memory");

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_write_cr3(value)
        }
    }

    impl Cr4 {
        /// Read the current set of CR4 flags.
        #[inline]
        pub fn read() -> Cr4Flags {
            Cr4Flags::from_bits_truncate(Self::read_raw())
        }

        /// Read the current raw CR4 value.
        #[inline]
        pub fn read_raw() -> u64 {
            let value: u64;

            #[cfg(feature = "inline_asm")]
            unsafe {
                llvm_asm!("mov %cr4, $0" : "=r" (value));
            }

            #[cfg(not(feature = "inline_asm"))]
            unsafe {
                value = crate::asm::x86_64_asm_read_cr4();
            }

            value
        }

        /// Write CR4 flags.
        ///
        /// Preserves the value of reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by overwriting the physical address extension
        /// flag.
        #[inline]
        pub unsafe fn write(flags: Cr4Flags) {
            let old_value = Self::read_raw();
            let reserved = old_value & !(Cr4Flags::all().bits());
            let new_value = reserved | flags.bits();

            Self::write_raw(new_value);
        }

        /// Write raw CR4 flags.
        ///
        /// Does _not_ preserve any values, including reserved fields.
        ///
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by overwriting the physical address extension
        /// flag.
        #[inline]
        pub unsafe fn write_raw(value: u64) {
            #[cfg(feature = "inline_asm")]
            llvm_asm!("mov $0, %cr4" :: "r" (value) : "memory");

            #[cfg(not(feature = "inline_asm"))]
            crate::asm::x86_64_asm_write_cr4(value);
        }

        /// Updates CR4 flags.
        ///
        /// Preserves the value of reserved fields.
        /// ## Safety
        ///
        /// This function is unsafe because it's possible to violate memory
        /// safety through it, e.g. by overwriting the physical address extension
        /// flag.
        #[inline]
        pub unsafe fn update<F>(f: F)
        where
            F: FnOnce(&mut Cr4Flags),
        {
            let mut flags = Self::read();
            f(&mut flags);
            Self::write(flags);
        }
    }
}
<|endoftext|>


<commit_before>use std::default::Default;
use std::option::Option;

use time::Timespec;

// use resources::Club;
// use resources::Gear;
use resources::enums::ResourceState;

/// Athletes are Strava users, Strava users are athletes.
///
/// The object is returned in detailed, summary or meta representations.
///
/// See: http://strava.github.io/api/v3/athlete/
#[allow(dead_code)]
#[derive(Default)]
pub struct Athlete {
    id: Option<i32>,
    resource_state: ResourceState,
    firstname: Option<String>,
    lastname: Option<String>,
    profile_medium: Option<String>,
    profile: Option<String>,
    city: Option<String>,
    state: Option<String>,
    country: Option<String>,
    sex: Option<String>,
    friend: Option<String>,
    follower: Option<String>,
    premium: Option<bool>,
    created_at: Option<Timespec>,
    updated_at: Option<Timespec>,
    approve_followers: Option<bool>,
    follower_count: Option<i32>,
    friend_count: Option<i32>,
    mutual_friend_count: Option<i32>,
    date_preference: Option<String>,
    measurement_preference: Option<String>,
    email: Option<String>,
    ftp: Option<i32>,
    weight: Option<f32>
    // clubs: Vec<Club>,
    // shoes: Vec<Gear>,
    // bikes: Vec<Gear>
}

pub struct IOError;

impl Athlete {
    fn new() -> Athlete { Default::default() }

    pub fn get_current() -> Result<Athlete, IOError> {
        // let res = Http::new()
        //     .get("https://strava.com/api/v3/athlete")
        //     .exec().unwrap();

        // println!("{}", url);

        // TODO make http request
        Ok(Athlete::new())
    }

    pub fn get_by_id(id: i32) -> Result<Athlete, IOError> {
        let url = format!("https://strava.com/api/v3/athletes/{}", id);
        println!("{}", url);

        // TODO make http request
        Ok(Athlete::new())
    }
}
<commit_msg>Partial implementation of Athlete::get_current<commit_after>use std::default::Default;
use std::option::Option;

use time::Timespec;

// use resources::Club;
// use resources::Gear;
use resources::enums::ResourceState;

use http::Http;
use accesstoken::AccessToken;

/// Athletes are Strava users, Strava users are athletes.
///
/// The object is returned in detailed, summary or meta representations.
///
/// See: http://strava.github.io/api/v3/athlete/
#[allow(dead_code)]
#[derive(Default)]
pub struct Athlete {
    id: Option<i32>,
    resource_state: ResourceState,
    firstname: Option<String>,
    lastname: Option<String>,
    profile_medium: Option<String>,
    profile: Option<String>,
    city: Option<String>,
    state: Option<String>,
    country: Option<String>,
    sex: Option<String>,
    friend: Option<String>,
    follower: Option<String>,
    premium: Option<bool>,
    created_at: Option<Timespec>,
    updated_at: Option<Timespec>,
    approve_followers: Option<bool>,
    follower_count: Option<i32>,
    friend_count: Option<i32>,
    mutual_friend_count: Option<i32>,
    date_preference: Option<String>,
    measurement_preference: Option<String>,
    email: Option<String>,
    ftp: Option<i32>,
    weight: Option<f32>
    // clubs: Vec<Club>,
    // shoes: Vec<Gear>,
    // bikes: Vec<Gear>
}

#[derive(Debug, Copy, Clone)]
pub enum ApiError {
    InvalidAccessToken
}

impl Athlete {
    fn new() -> Athlete { Default::default() }

    pub fn get_current(token: &AccessToken) -> Result<Athlete, ApiError> {
        let url = format!("https://strava.com/api/v3/athlete?access_token={}", token.get());
        println!("access token: {}", url);
        let res = Http::new().get(url.as_ref()).unwrap();

        println!("{}", res);

        // TODO make http request
        Ok(Athlete::new())
    }

    pub fn get_by_id(id: i32) -> Result<Athlete, ApiError> {
        let url = format!("https://strava.com/api/v3/athletes/{}", id);
        println!("{}", url);

        // TODO make http request
        Ok(Athlete::new())
    }
}

#[cfg(test)]
mod tests {
    use accesstoken::AccessToken;
    use resources::athlete::Athlete;
    use std::result::Result;

    #[test]
    fn get_current_athlete() {
        let token = AccessToken::from("fake_token");

        let res = Athlete::get_current(&token);
        assert!(!res.is_err());
        assert!(res.is_err());
    }
}
<|endoftext|>


<commit_before>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::env;
use std::io::Error;
use std::ptr::null_mut;

use widestring::WideCString;
use winapi::winerror::*;
use winapi::{BOOL, LPCWSTR, LPDWORD, PSID, PSID_NAME_USE, SID_NAME_USE};

use super::sid::Sid;

extern "system" {
    fn LookupAccountNameW(
        lpSystemName: LPCWSTR,
        lpAccountName: LPCWSTR,
        Sid: PSID,
        cbSid: LPDWORD,
        ReferencedDomainName: LPCWSTR,
        cchReferencedDomainName: LPDWORD,
        peUse: PSID_NAME_USE,
    ) -> BOOL;
}

pub struct Account {
    pub name: String,
    pub system_name: Option<String>,
    pub domain: String,
    pub account_type: SID_NAME_USE,
    pub sid: Sid,
}

impl Account {
    pub fn from_name(name: &str) -> Option<Account> {
        lookup_account(name, None)
    }

    pub fn from_name_and_system(name: &str, system_name: &str) -> Option<Account> {
        lookup_account(name, Some(system_name.to_string()))
    }
}

fn lookup_account(name: &str, system_name: Option<String>) -> Option<Account> {
    // if this is a machine account, strip the terminating '$'
    // LookupAccountName will return the sid of the computer account
    // given the computer name. Windows forbids usernames to match the
    // computer name
    let stripped_name =
        if name.to_lowercase() == (env::var("COMPUTERNAME").unwrap().to_lowercase() + "$") {
            &name[..name.len() - 1]
        } else {
            name
        };
    let mut sid_size: u32 = 0;
    let mut domain_size: u32 = 0;
    let wide = WideCString::from_str(stripped_name).unwrap();
    unsafe {
        LookupAccountNameW(
            null_mut(),
            wide.as_ptr(),
            null_mut(),
            &mut sid_size as LPDWORD,
            null_mut(),
            &mut domain_size as LPDWORD,
            null_mut(),
        )
    };
    match Error::last_os_error().raw_os_error().unwrap() as u32 {
        ERROR_INSUFFICIENT_BUFFER => {}
        ERROR_NONE_MAPPED => return None,
        _ => {
            info!(
                "Error while looking up account for {}: {}",
                name,
                Error::last_os_error()
            );
            return None;
        }
    }

    let mut sid: Vec<u8> = Vec::with_capacity(sid_size as usize);
    let mut domain: Vec<u16> = Vec::with_capacity(domain_size as usize);
    let mut sid_type = SID_NAME_USE(0);

    let ret = unsafe {
        LookupAccountNameW(
            null_mut(),
            wide.as_ptr(),
            sid.as_mut_ptr() as PSID,
            &mut sid_size as LPDWORD,
            domain.as_mut_ptr(),
            &mut domain_size as LPDWORD,
            &mut sid_type as PSID_NAME_USE,
        )
    };
    if ret == 0 {
        info!(
            "Failed to retrieve SID for {}: {}",
            name,
            Error::last_os_error()
        );
        return None;
    }
    unsafe {
        domain.set_len(domain_size as usize);
        sid.set_len(sid_size as usize);
    }
    let domain_str = WideCString::new(domain).unwrap().to_string_lossy();
    Some(Account {
        name: name.to_string(),
        system_name: system_name,
        domain: domain_str,
        account_type: sid_type,
        sid: Sid { raw: sid },
    })
}

#[cfg(test)]
mod tests {
    use std::env;

    use winapi::winnt::{SidTypeUser, SidTypeWellKnownGroup};

    use super::*;

    #[test]
    fn real_account_returns_some() {
        assert_eq!(Account::from_name("Administrator").is_some(), true)
    }

    #[test]
    fn bogus_account_returns_none() {
        assert_eq!(Account::from_name("bogus").is_none(), true)
    }

    #[test]
    fn user_account_returns_user_type() {
        let acct_type = Account::from_name("Administrator").unwrap().account_type;
        assert_eq!(acct_type, SidTypeUser)
    }

    #[test]
    fn local_user_account_returns_local_machine_as_domain() {
        let acct_domain = Account::from_name("Administrator").unwrap().domain;
        assert_eq!(acct_domain, env::var("COMPUTERNAME").unwrap())
    }

    #[test]
    fn well_known_group_account_returns_correct_type() {
        let acct_type = Account::from_name("Everyone").unwrap().account_type;
        assert_eq!(acct_type, SidTypeWellKnownGroup)
    }

    #[test]
    fn well_known_group_account_has_well_known_sid() {
        let sid = Account::from_name("Everyone").unwrap().sid;
        assert_eq!(sid.to_string().unwrap(), "S-1-1-0")
    }

    #[test]
    fn mixing_case_returns_same_account() {
        let current_user = env::var("USERNAME").unwrap();
        let lower_sid = Account::from_name(current_user.to_lowercase().as_str())
            .unwrap()
            .sid;
        let upper_sid = Account::from_name(current_user.to_uppercase().as_str())
            .unwrap()
            .sid;
        assert_eq!(
            lower_sid.to_string().unwrap(),
            upper_sid.to_string().unwrap()
        )
    }

    #[test]
    fn machine_account_returns_some() {
        assert_eq!(
            Account::from_name((env::var("COMPUTERNAME").unwrap() + "$").as_str()).is_some(),
            true
        )
    }
}
<commit_msg>do not panic validating users and fail validation instead<commit_after>// Copyright (c) 2016-2017 Chef Software Inc. and/or applicable contributors
//
// 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.

use std::env;
use std::io::Error;
use std::ptr::null_mut;

use widestring::WideCString;
use winapi::winerror::*;
use winapi::{BOOL, LPCWSTR, LPDWORD, PSID, PSID_NAME_USE, SID_NAME_USE};

use super::sid::Sid;

extern "system" {
    fn LookupAccountNameW(
        lpSystemName: LPCWSTR,
        lpAccountName: LPCWSTR,
        Sid: PSID,
        cbSid: LPDWORD,
        ReferencedDomainName: LPCWSTR,
        cchReferencedDomainName: LPDWORD,
        peUse: PSID_NAME_USE,
    ) -> BOOL;
}

pub struct Account {
    pub name: String,
    pub system_name: Option<String>,
    pub domain: String,
    pub account_type: SID_NAME_USE,
    pub sid: Sid,
}

impl Account {
    pub fn from_name(name: &str) -> Option<Account> {
        lookup_account(name, None)
    }

    pub fn from_name_and_system(name: &str, system_name: &str) -> Option<Account> {
        lookup_account(name, Some(system_name.to_string()))
    }
}

fn lookup_account(name: &str, system_name: Option<String>) -> Option<Account> {
    // if this is a machine account, strip the terminating '$'
    // LookupAccountName will return the sid of the computer account
    // given the computer name. Windows forbids usernames to match the
    // computer name
    let stripped_name =
        if name.to_lowercase() == (env::var("COMPUTERNAME").unwrap().to_lowercase() + "$") {
            &name[..name.len() - 1]
        } else {
            name
        };
    let mut sid_size: u32 = 0;
    let mut domain_size: u32 = 0;
    let wide = WideCString::from_str(stripped_name).unwrap();
    unsafe {
        LookupAccountNameW(
            null_mut(),
            wide.as_ptr(),
            null_mut(),
            &mut sid_size as LPDWORD,
            null_mut(),
            &mut domain_size as LPDWORD,
            null_mut(),
        )
    };
    match Error::last_os_error().raw_os_error().unwrap() as u32 {
        ERROR_INSUFFICIENT_BUFFER => {}
        ERROR_NONE_MAPPED => return None,
        _ => {
            error!(
                "Error while looking up account for {}: {}",
                name,
                Error::last_os_error()
            );
            return None;
        }
    }

    let mut sid: Vec<u8> = Vec::with_capacity(sid_size as usize);
    let mut domain: Vec<u16> = Vec::with_capacity(domain_size as usize);
    let mut sid_type = SID_NAME_USE(0);

    let ret = unsafe {
        LookupAccountNameW(
            null_mut(),
            wide.as_ptr(),
            sid.as_mut_ptr() as PSID,
            &mut sid_size as LPDWORD,
            domain.as_mut_ptr(),
            &mut domain_size as LPDWORD,
            &mut sid_type as PSID_NAME_USE,
        )
    };
    if ret == 0 {
        error!(
            "Failed to retrieve SID for {}: {}",
            name,
            Error::last_os_error()
        );
        return None;
    }
    unsafe {
        domain.set_len(domain_size as usize);
        sid.set_len(sid_size as usize);
    }
    let domain_str = WideCString::new(domain).unwrap().to_string_lossy();
    Some(Account {
        name: name.to_string(),
        system_name: system_name,
        domain: domain_str,
        account_type: sid_type,
        sid: Sid { raw: sid },
    })
}

#[cfg(test)]
mod tests {
    use std::env;

    use winapi::winnt::{SidTypeUser, SidTypeWellKnownGroup};

    use super::*;

    #[test]
    fn real_account_returns_some() {
        assert_eq!(Account::from_name("Administrator").is_some(), true)
    }

    #[test]
    fn bogus_account_returns_none() {
        assert_eq!(Account::from_name("bogus").is_none(), true)
    }

    #[test]
    fn user_account_returns_user_type() {
        let acct_type = Account::from_name("Administrator").unwrap().account_type;
        assert_eq!(acct_type, SidTypeUser)
    }

    #[test]
    fn local_user_account_returns_local_machine_as_domain() {
        let acct_domain = Account::from_name("Administrator").unwrap().domain;
        assert_eq!(acct_domain, env::var("COMPUTERNAME").unwrap())
    }

    #[test]
    fn well_known_group_account_returns_correct_type() {
        let acct_type = Account::from_name("Everyone").unwrap().account_type;
        assert_eq!(acct_type, SidTypeWellKnownGroup)
    }

    #[test]
    fn well_known_group_account_has_well_known_sid() {
        let sid = Account::from_name("Everyone").unwrap().sid;
        assert_eq!(sid.to_string().unwrap(), "S-1-1-0")
    }

    #[test]
    fn mixing_case_returns_same_account() {
        let current_user = env::var("USERNAME").unwrap();
        let lower_sid = Account::from_name(current_user.to_lowercase().as_str())
            .unwrap()
            .sid;
        let upper_sid = Account::from_name(current_user.to_uppercase().as_str())
            .unwrap()
            .sid;
        assert_eq!(
            lower_sid.to_string().unwrap(),
            upper_sid.to_string().unwrap()
        )
    }

    #[test]
    fn machine_account_returns_some() {
        assert_eq!(
            Account::from_name((env::var("COMPUTERNAME").unwrap() + "$").as_str()).is_some(),
            true
        )
    }
}
<|endoftext|>