<commit_before>use std::marker::PhantomData;

use Construct;
use Data;
use Count;
use Subspace;
use ToIndex;
use ToPos;

/// Dimension is natural number, position is the same as index.
pub struct Dimension<T>(PhantomData<T>);

impl<T> Construct for Dimension<T> {
    fn new() -> Dimension<T> { Dimension(PhantomData) }
}

impl Count<usize> for Dimension<Data> {
    fn count(&self, dim: usize) -> usize { dim }
}

impl<T, U> Count<(usize, U)> for Dimension<Subspace<T>>
    where
        T: Construct + Count<U>
{
    fn count(&self, (a, b): (usize, U)) -> usize {
        let subspace: T = Construct::new();
        a * subspace.count(b)
    }
}

impl ToIndex<usize, usize> for Dimension<Data> {
    fn to_index(&self, _dim: usize, pos: usize) -> usize { pos }
}

impl<T, U: Copy, V> ToIndex<(usize, U), (usize, V)> for Dimension<Subspace<T>>
    where
        T: Construct + Count<U> + ToIndex<U, V>
{
    fn to_index(&self, (_a, b): (usize, U), (pa, pb): (usize, V)) -> usize {
        let subspace: T = Construct::new();
        let count = subspace.count(b);
        pa * count + subspace.to_index(b, pb)
    }
}

impl ToPos<usize, usize> for Dimension<Data> {
    fn to_pos(&self, _dim: usize, index: usize, pos: &mut usize) {
        *pos = index;
    }
}

impl<T, U: Copy, V>
ToPos<(usize, U), (usize, V)> for Dimension<Subspace<T>>
    where
        T: Construct + Count<U> + ToPos<U, V>
{
    fn to_pos(
        &self,
        (_a, b): (usize, U),
        index: usize,
        &mut (ref mut head, ref mut tail): &mut (usize, V)
    ) {
        let subspace: T = Construct::new();
        let count = subspace.count(b);
        let x = index / count;
        *head = x;
        subspace.to_pos(b, index - x * count, tail)
    }
}

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

    #[test]
    fn features() {
        is_complete::<Dimension<Data>, usize, usize, usize>();
        does_count::<Dimension<Subspace<Pair<Data>>>, (usize, usize)>();
        // does_count::<Dimension<Of<Pair<Data>>>, usize>();
    }

    #[test]
    fn subspace() {
        type D2 = Dimension<Subspace<Dimension<Data>>>;
        type D3 = Dimension<Subspace<D2>>;

        let x: D3 = Construct::new();
        let dim = (3, (3, 3));
        assert_eq!(x.count(dim), 27);
        let pos = (1, (0, 2));
        let index = x.to_index(dim, pos);
        assert_eq!(index, 11);
        let mut new_pos = pos;
        x.to_pos(dim, index, &mut new_pos);
        assert_eq!(pos, new_pos);
    }
}
<commit_msg>Added `Count` for `Dimension<Of<T>>`<commit_after>use std::marker::PhantomData;

use Construct;
use Data;
use Count;
use Of;
use Subspace;
use ToIndex;
use ToPos;

/// Dimension is natural number, position is the same as index.
pub struct Dimension<T>(PhantomData<T>);

impl<T> Construct for Dimension<T> {
    fn new() -> Dimension<T> { Dimension(PhantomData) }
}

impl Count<usize> for Dimension<Data> {
    fn count(&self, dim: usize) -> usize { dim }
}

impl<T, U> Count<(usize, U)> for Dimension<Subspace<T>>
    where
        T: Construct + Count<U>
{
    fn count(&self, (a, b): (usize, U)) -> usize {
        let subspace: T = Construct::new();
        a * subspace.count(b)
    }
}

impl<T, U> Count<U> for Dimension<Of<T>>
    where
        T: Construct + Count<U>
{
    fn count(&self, dim: U) -> usize {
        let of: T = Construct::new();
        of.count(dim)
    }
}

impl ToIndex<usize, usize> for Dimension<Data> {
    fn to_index(&self, _dim: usize, pos: usize) -> usize { pos }
}

impl<T, U: Copy, V> ToIndex<(usize, U), (usize, V)> for Dimension<Subspace<T>>
    where
        T: Construct + Count<U> + ToIndex<U, V>
{
    fn to_index(&self, (_a, b): (usize, U), (pa, pb): (usize, V)) -> usize {
        let subspace: T = Construct::new();
        let count = subspace.count(b);
        pa * count + subspace.to_index(b, pb)
    }
}

impl ToPos<usize, usize> for Dimension<Data> {
    fn to_pos(&self, _dim: usize, index: usize, pos: &mut usize) {
        *pos = index;
    }
}

impl<T, U: Copy, V>
ToPos<(usize, U), (usize, V)> for Dimension<Subspace<T>>
    where
        T: Construct + Count<U> + ToPos<U, V>
{
    fn to_pos(
        &self,
        (_a, b): (usize, U),
        index: usize,
        &mut (ref mut head, ref mut tail): &mut (usize, V)
    ) {
        let subspace: T = Construct::new();
        let count = subspace.count(b);
        let x = index / count;
        *head = x;
        subspace.to_pos(b, index - x * count, tail)
    }
}

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

    #[test]
    fn features() {
        is_complete::<Dimension<Data>, usize, usize, usize>();
        does_count::<Dimension<Subspace<Pair<Data>>>, (usize, usize)>();
        does_count::<Dimension<Of<Pair<Data>>>, usize>();
    }

    #[test]
    fn subspace() {
        type D2 = Dimension<Subspace<Dimension<Data>>>;
        type D3 = Dimension<Subspace<D2>>;

        let x: D3 = Construct::new();
        let dim = (3, (3, 3));
        assert_eq!(x.count(dim), 27);
        let pos = (1, (0, 2));
        let index = x.to_index(dim, pos);
        assert_eq!(index, 11);
        let mut new_pos = pos;
        x.to_pos(dim, index, &mut new_pos);
        assert_eq!(pos, new_pos);
    }
}
<|endoftext|>


<commit_before>use base64;
use crypto::digest::Digest;
use crypto::mac::Mac;
use crypto;
use futures::sync::mpsc;
use futures::{Future, Stream, Poll};
use hyper::server::{Service, Request, Response, Http};
use hyper::{self, Get, Post, StatusCode};

#[cfg(feature = "with-dns-sd")]
use dns_sd::DNSService;

#[cfg(not(feature = "with-dns-sd"))]
use mdns;

use num_bigint::BigUint;
use rand;
use std::collections::BTreeMap;
use std::io;
use std::sync::Arc;
use tokio_core::reactor::Handle;
use url;

use core::diffie_hellman::{DH_GENERATOR, DH_PRIME};
use core::authentication::Credentials;
use core::util;
use core::config::ConnectConfig;

#[derive(Clone)]
struct Discovery(Arc<DiscoveryInner>);
struct DiscoveryInner {
    config: ConnectConfig,
    device_id: String,
    private_key: BigUint,
    public_key: BigUint,
    tx: mpsc::UnboundedSender<Credentials>,
}

impl Discovery {
    fn new(config: ConnectConfig, device_id: String)
        -> (Discovery, mpsc::UnboundedReceiver<Credentials>)
    {
        let (tx, rx) = mpsc::unbounded();

        let key_data = util::rand_vec(&mut rand::thread_rng(), 95);
        let private_key = BigUint::from_bytes_be(&key_data);
        let public_key = util::powm(&DH_GENERATOR, &private_key, &DH_PRIME);

        let discovery = Discovery(Arc::new(DiscoveryInner {
            config: config,
            device_id: device_id,
            private_key: private_key,
            public_key: public_key,
            tx: tx,
        }));

        (discovery, rx)
    }
}

impl Discovery {
    fn handle_get_info(&self, _params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let public_key = self.0.public_key.to_bytes_be();
        let public_key = base64::encode(&public_key);

        let result = json!({
            "status": 101,
            "statusString": "ERROR-OK",
            "spotifyError": 0,
            "version": "2.1.0",
            "deviceID": (self.0.device_id),
            "remoteName": (self.0.config.name),
            "activeUser": "",
            "publicKey": (public_key),
            "deviceType": (self.0.config.device_type.to_string().to_uppercase()),
            "libraryVersion": "0.1.0",
            "accountReq": "PREMIUM",
            "brandDisplayName": "librespot",
            "modelDisplayName": "librespot",
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn handle_add_user(&self, params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let username = params.get("userName").unwrap();
        let encrypted_blob = params.get("blob").unwrap();
        let client_key = params.get("clientKey").unwrap();

        let encrypted_blob = base64::decode(encrypted_blob).unwrap();

        let client_key = base64::decode(client_key).unwrap();
        let client_key = BigUint::from_bytes_be(&client_key);

        let shared_key = util::powm(&client_key, &self.0.private_key, &DH_PRIME);

        let iv = &encrypted_blob[0..16];
        let encrypted = &encrypted_blob[16..encrypted_blob.len() - 20];
        let cksum = &encrypted_blob[encrypted_blob.len() - 20..encrypted_blob.len()];

        let base_key = {
            let mut data = [0u8; 20];
            let mut h = crypto::sha1::Sha1::new();
            h.input(&shared_key.to_bytes_be());
            h.result(&mut data);
            data[..16].to_owned()
        };

        let checksum_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"checksum");
            h.result().code().to_owned()
        };

        let encryption_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"encryption");
            h.result().code().to_owned()
        };

        let mac = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &checksum_key);
            h.input(encrypted);
            h.result().code().to_owned()
        };

        assert_eq!(&mac[..], cksum);

        let decrypted = {
            let mut data = vec![0u8; encrypted.len()];
            let mut cipher = crypto::aes::ctr(crypto::aes::KeySize::KeySize128,
                                              &encryption_key[0..16], iv);
            cipher.process(encrypted, &mut data);
            String::from_utf8(data).unwrap()
        };

        let credentials = Credentials::with_blob(username.to_owned(), &decrypted, &self.0.device_id);

        self.0.tx.unbounded_send(credentials).unwrap();

        let result = json!({
            "status": 101,
            "spotifyError": 0,
            "statusString": "ERROR-OK"
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn not_found(&self)
        -> ::futures::Finished<Response, hyper::Error>
    {
        ::futures::finished(Response::new().with_status(StatusCode::NotFound))
    }
}

impl Service for Discovery {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item = Response, Error = hyper::Error>>;

    fn call(&self, request: Request) -> Self::Future {
        let mut params = BTreeMap::new();

        let (method, uri, _, _, body) = request.deconstruct();
        if let Some(query) = uri.query() {
            params.extend(url::form_urlencoded::parse(query.as_bytes()).into_owned());
        }

        if method != Get {
            debug!("{:?} {:?} {:?}", method, uri.path(), params);
        }

        let this = self.clone();
        Box::new(body.fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(chunk.as_ref());
            Ok::<_, hyper::Error>(acc)
        }).map(move |body| {
            params.extend(url::form_urlencoded::parse(&body).into_owned());
            params
        }).and_then(move |params| {
            match (method, params.get("action").map(AsRef::as_ref)) {
                (Get, Some("getInfo")) => this.handle_get_info(&params),
                (Post, Some("addUser")) => this.handle_add_user(&params),
                _ => this.not_found(),
            }
        }))
    }
}

#[cfg(feature = "with-dns-sd")]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: DNSService,
}

#[cfg(not(feature = "with-dns-sd"))]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: mdns::Service,
}

pub fn discovery(handle: &Handle, config: ConnectConfig, device_id: String, port: u16)
    -> io::Result<DiscoveryStream>
{
    let (discovery, creds_rx) = Discovery::new(config.clone(), device_id);

    let serve = {
        let http = Http::new();
        debug!("Zeroconf server listening on 0.0.0.0:{}", port);
        http.serve_addr_handle(&format!("0.0.0.0:{}", port).parse().unwrap(), &handle, move || Ok(discovery.clone())).unwrap()
    };

    #[cfg(feature = "with-dns-sd")]
    let port = serve.incoming_ref().local_addr().port();

    #[cfg(not(feature = "with-dns-sd"))]
    let addr = serve.incoming_ref().local_addr();

    let server_future = {
        let handle = handle.clone();
        serve.for_each(move |connection| {
                handle.spawn(connection.then(|_| Ok(())));
                Ok(())
            })
            .then(|_| Ok(()))
    };
    handle.spawn(server_future);

    #[cfg(feature = "with-dns-sd")]
    let svc = DNSService::register(Some(&*config.name),
       "_spotify-connect._tcp",
       None,
       None,
       port,
       &["VERSION=1.0", "CPath=/"]).unwrap();

    #[cfg(not(feature = "with-dns-sd"))]
    let responder = mdns::Responder::spawn(&handle)?;

    #[cfg(not(feature = "with-dns-sd"))]
    let svc = responder.register(
        "_spotify-connect._tcp".to_owned(),
        config.name,
        addr.port(),
        &["VERSION=1.0", "CPath=/"]);

    Ok(DiscoveryStream {
        credentials: creds_rx,
        _svc: svc,
    })
}

impl Stream for DiscoveryStream {
    type Item = Credentials;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.credentials.poll()
    }
}
<commit_msg>simplify dns-sd/mdns code<commit_after>use base64;
use crypto::digest::Digest;
use crypto::mac::Mac;
use crypto;
use futures::sync::mpsc;
use futures::{Future, Stream, Poll};
use hyper::server::{Service, Request, Response, Http};
use hyper::{self, Get, Post, StatusCode};

#[cfg(feature = "with-dns-sd")]
use dns_sd::DNSService;

#[cfg(not(feature = "with-dns-sd"))]
use mdns;

use num_bigint::BigUint;
use rand;
use std::collections::BTreeMap;
use std::io;
use std::sync::Arc;
use tokio_core::reactor::Handle;
use url;

use core::diffie_hellman::{DH_GENERATOR, DH_PRIME};
use core::authentication::Credentials;
use core::util;
use core::config::ConnectConfig;

#[derive(Clone)]
struct Discovery(Arc<DiscoveryInner>);
struct DiscoveryInner {
    config: ConnectConfig,
    device_id: String,
    private_key: BigUint,
    public_key: BigUint,
    tx: mpsc::UnboundedSender<Credentials>,
}

impl Discovery {
    fn new(config: ConnectConfig, device_id: String)
        -> (Discovery, mpsc::UnboundedReceiver<Credentials>)
    {
        let (tx, rx) = mpsc::unbounded();

        let key_data = util::rand_vec(&mut rand::thread_rng(), 95);
        let private_key = BigUint::from_bytes_be(&key_data);
        let public_key = util::powm(&DH_GENERATOR, &private_key, &DH_PRIME);

        let discovery = Discovery(Arc::new(DiscoveryInner {
            config: config,
            device_id: device_id,
            private_key: private_key,
            public_key: public_key,
            tx: tx,
        }));

        (discovery, rx)
    }
}

impl Discovery {
    fn handle_get_info(&self, _params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let public_key = self.0.public_key.to_bytes_be();
        let public_key = base64::encode(&public_key);

        let result = json!({
            "status": 101,
            "statusString": "ERROR-OK",
            "spotifyError": 0,
            "version": "2.1.0",
            "deviceID": (self.0.device_id),
            "remoteName": (self.0.config.name),
            "activeUser": "",
            "publicKey": (public_key),
            "deviceType": (self.0.config.device_type.to_string().to_uppercase()),
            "libraryVersion": "0.1.0",
            "accountReq": "PREMIUM",
            "brandDisplayName": "librespot",
            "modelDisplayName": "librespot",
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn handle_add_user(&self, params: &BTreeMap<String, String>)
        -> ::futures::Finished<Response, hyper::Error>
    {
        let username = params.get("userName").unwrap();
        let encrypted_blob = params.get("blob").unwrap();
        let client_key = params.get("clientKey").unwrap();

        let encrypted_blob = base64::decode(encrypted_blob).unwrap();

        let client_key = base64::decode(client_key).unwrap();
        let client_key = BigUint::from_bytes_be(&client_key);

        let shared_key = util::powm(&client_key, &self.0.private_key, &DH_PRIME);

        let iv = &encrypted_blob[0..16];
        let encrypted = &encrypted_blob[16..encrypted_blob.len() - 20];
        let cksum = &encrypted_blob[encrypted_blob.len() - 20..encrypted_blob.len()];

        let base_key = {
            let mut data = [0u8; 20];
            let mut h = crypto::sha1::Sha1::new();
            h.input(&shared_key.to_bytes_be());
            h.result(&mut data);
            data[..16].to_owned()
        };

        let checksum_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"checksum");
            h.result().code().to_owned()
        };

        let encryption_key = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &base_key);
            h.input(b"encryption");
            h.result().code().to_owned()
        };

        let mac = {
            let mut h = crypto::hmac::Hmac::new(crypto::sha1::Sha1::new(), &checksum_key);
            h.input(encrypted);
            h.result().code().to_owned()
        };

        assert_eq!(&mac[..], cksum);

        let decrypted = {
            let mut data = vec![0u8; encrypted.len()];
            let mut cipher = crypto::aes::ctr(crypto::aes::KeySize::KeySize128,
                                              &encryption_key[0..16], iv);
            cipher.process(encrypted, &mut data);
            String::from_utf8(data).unwrap()
        };

        let credentials = Credentials::with_blob(username.to_owned(), &decrypted, &self.0.device_id);

        self.0.tx.unbounded_send(credentials).unwrap();

        let result = json!({
            "status": 101,
            "spotifyError": 0,
            "statusString": "ERROR-OK"
        });

        let body = result.to_string();
        ::futures::finished(Response::new().with_body(body))
    }

    fn not_found(&self)
        -> ::futures::Finished<Response, hyper::Error>
    {
        ::futures::finished(Response::new().with_status(StatusCode::NotFound))
    }
}

impl Service for Discovery {
    type Request = Request;
    type Response = Response;
    type Error = hyper::Error;
    type Future = Box<Future<Item = Response, Error = hyper::Error>>;

    fn call(&self, request: Request) -> Self::Future {
        let mut params = BTreeMap::new();

        let (method, uri, _, _, body) = request.deconstruct();
        if let Some(query) = uri.query() {
            params.extend(url::form_urlencoded::parse(query.as_bytes()).into_owned());
        }

        if method != Get {
            debug!("{:?} {:?} {:?}", method, uri.path(), params);
        }

        let this = self.clone();
        Box::new(body.fold(Vec::new(), |mut acc, chunk| {
            acc.extend_from_slice(chunk.as_ref());
            Ok::<_, hyper::Error>(acc)
        }).map(move |body| {
            params.extend(url::form_urlencoded::parse(&body).into_owned());
            params
        }).and_then(move |params| {
            match (method, params.get("action").map(AsRef::as_ref)) {
                (Get, Some("getInfo")) => this.handle_get_info(&params),
                (Post, Some("addUser")) => this.handle_add_user(&params),
                _ => this.not_found(),
            }
        }))
    }
}

#[cfg(feature = "with-dns-sd")]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: DNSService,
}

#[cfg(not(feature = "with-dns-sd"))]
pub struct DiscoveryStream {
    credentials: mpsc::UnboundedReceiver<Credentials>,
    _svc: mdns::Service,
}

pub fn discovery(handle: &Handle, config: ConnectConfig, device_id: String, port: u16)
    -> io::Result<DiscoveryStream>
{
    let (discovery, creds_rx) = Discovery::new(config.clone(), device_id);

    let serve = {
        let http = Http::new();
        debug!("Zeroconf server listening on 0.0.0.0:{}", port);
        http.serve_addr_handle(&format!("0.0.0.0:{}", port).parse().unwrap(), &handle, move || Ok(discovery.clone())).unwrap()
    };

    let s_port = serve.incoming_ref().local_addr().port();

    let server_future = {
        let handle = handle.clone();
        serve.for_each(move |connection| {
                handle.spawn(connection.then(|_| Ok(())));
                Ok(())
            })
            .then(|_| Ok(()))
    };
    handle.spawn(server_future);

    #[cfg(feature = "with-dns-sd")]
    let svc = DNSService::register(Some(&*config.name),
       "_spotify-connect._tcp",
       None,
       None,
       s_port,
       &["VERSION=1.0", "CPath=/"]).unwrap();

    #[cfg(not(feature = "with-dns-sd"))]
    let responder = mdns::Responder::spawn(&handle)?;

    #[cfg(not(feature = "with-dns-sd"))]
    let svc = responder.register(
        "_spotify-connect._tcp".to_owned(),
        config.name,
        s_port,
        &["VERSION=1.0", "CPath=/"]);

    Ok(DiscoveryStream {
        credentials: creds_rx,
        _svc: svc,
    })
}

impl Stream for DiscoveryStream {
    type Item = Credentials;
    type Error = ();

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        self.credentials.poll()
    }
}
<|endoftext|>


<commit_before>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

//! Handles routing to devices in an address space.

use std::cmp::{Ord, Ordering, PartialEq, PartialOrd};
use std::collections::btree_map::BTreeMap;
use std::fmt;
use std::io;
use std::result;
use std::sync::{Arc, Mutex};

use crate::virtio::AsAny;

/// Trait for devices that respond to reads or writes in an arbitrary address space.
///
/// The device does not care where it exists in address space as each method is only given an offset
/// into its allocated portion of address space.
#[allow(unused_variables)]
pub trait BusDevice: AsAny + Send {
    /// Reads at `offset` from this device
    fn read(&mut self, offset: u64, data: &mut [u8]) {}
    /// Writes at `offset` into this device
    fn write(&mut self, offset: u64, data: &[u8]) {}
    /// Triggers the `irq_mask` interrupt on this device
    fn interrupt(&self, irq_mask: u32) -> io::Result<()> {
        Ok(())
    }
}

#[derive(Debug)]
pub enum Error {
    /// The insertion failed because the new device overlapped with an old device.
    Overlap,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match *self {
            Overlap => write!(f, "New device overlaps with an old device."),
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

#[derive(Debug, Copy, Clone)]
struct BusRange(u64, u64);

impl Eq for BusRange {}

impl PartialEq for BusRange {
    fn eq(&self, other: &BusRange) -> bool {
        self.0 == other.0
    }
}

impl Ord for BusRange {
    fn cmp(&self, other: &BusRange) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl PartialOrd for BusRange {
    fn partial_cmp(&self, other: &BusRange) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

/// A device container for routing reads and writes over some address space.
///
/// This doesn't have any restrictions on what kind of device or address space this applies to. The
/// only restriction is that no two devices can overlap in this address space.
#[derive(Clone, Default)]
pub struct Bus {
    devices: BTreeMap<BusRange, Arc<Mutex<dyn BusDevice>>>,
}

impl Bus {
    /// Constructs an a bus with an empty address space.
    pub fn new() -> Bus {
        Bus {
            devices: BTreeMap::new(),
        }
    }

    fn first_before(&self, addr: u64) -> Option<(BusRange, &Mutex<dyn BusDevice>)> {
        // for when we switch to rustc 1.17: self.devices.range(..addr).iter().rev().next()
        for (range, dev) in self.devices.iter().rev() {
            if range.0 <= addr {
                return Some((*range, dev));
            }
        }
        None
    }

    pub fn get_device(&self, addr: u64) -> Option<(u64, &Mutex<dyn BusDevice>)> {
        if let Some((BusRange(start, len), dev)) = self.first_before(addr) {
            let offset = addr - start;
            if offset < len {
                return Some((offset, dev));
            }
        }
        None
    }

    /// Puts the given device at the given address space.
    pub fn insert(&mut self, device: Arc<Mutex<dyn BusDevice>>, base: u64, len: u64) -> Result<()> {
        if len == 0 {
            return Err(Error::Overlap);
        }

        // Reject all cases where the new device's base is within an old device's range.
        if self.get_device(base).is_some() {
            return Err(Error::Overlap);
        }

        // The above check will miss an overlap in which the new device's base address is before the
        // range of another device. To catch that case, we search for a device with a range before
        // the new device's range's end. If there is no existing device in that range that starts
        // after the new device, then there will be no overlap.
        if let Some((BusRange(start, _), _)) = self.first_before(base + len - 1) {
            // Such a device only conflicts with the new device if it also starts after the new
            // device because of our initial `get_device` check above.
            if start >= base {
                return Err(Error::Overlap);
            }
        }

        if self.devices.insert(BusRange(base, len), device).is_some() {
            return Err(Error::Overlap);
        }

        Ok(())
    }

    /// Reads data from the device that owns the range containing `addr` and puts it into `data`.
    ///
    /// Returns true on success, otherwise `data` is untouched.
    pub fn read(&self, addr: u64, data: &mut [u8]) -> bool {
        if let Some((offset, dev)) = self.get_device(addr) {
            // OK to unwrap as lock() failing is a serious error condition and should panic.
            dev.lock()
                .expect("Failed to acquire device lock")
                .read(offset, data);
            true
        } else {
            false
        }
    }

    /// Writes `data` to the device that owns the range containing `addr`.
    ///
    /// Returns true on success, otherwise `data` is untouched.
    pub fn write(&self, addr: u64, data: &[u8]) -> bool {
        if let Some((offset, dev)) = self.get_device(addr) {
            // OK to unwrap as lock() failing is a serious error condition and should panic.
            dev.lock()
                .expect("Failed to acquire device lock")
                .write(offset, data);
            true
        } else {
            false
        }
    }
}

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

    struct DummyDevice;
    impl BusDevice for DummyDevice {}

    struct ConstantDevice;
    impl BusDevice for ConstantDevice {
        fn read(&mut self, offset: u64, data: &mut [u8]) {
            for (i, v) in data.iter_mut().enumerate() {
                *v = (offset as u8) + (i as u8);
            }
        }

        fn write(&mut self, offset: u64, data: &[u8]) {
            for (i, v) in data.iter().enumerate() {
                assert_eq!(*v, (offset as u8) + (i as u8))
            }
        }
    }

    #[test]
    fn bus_insert() {
        let mut bus = Bus::new();
        let dummy = Arc::new(Mutex::new(DummyDevice));
        assert!(bus.insert(dummy.clone(), 0x10, 0).is_err());
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_ok());

        let result = bus.insert(dummy.clone(), 0x0f, 0x10);
        assert!(result.is_err());
        assert_eq!(format!("{:?}", result), "Err(Overlap)");

        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_err());
        assert!(bus.insert(dummy.clone(), 0x10, 0x15).is_err());
        assert!(bus.insert(dummy.clone(), 0x12, 0x15).is_err());
        assert!(bus.insert(dummy.clone(), 0x12, 0x01).is_err());
        assert!(bus.insert(dummy.clone(), 0x0, 0x20).is_err());
        assert!(bus.insert(dummy.clone(), 0x20, 0x05).is_ok());
        assert!(bus.insert(dummy.clone(), 0x25, 0x05).is_ok());
        assert!(bus.insert(dummy.clone(), 0x0, 0x10).is_ok());
    }

    #[test]
    fn bus_read_write() {
        let mut bus = Bus::new();
        let dummy = Arc::new(Mutex::new(DummyDevice));
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_ok());
        assert!(bus.read(0x10, &mut [0, 0, 0, 0]));
        assert!(bus.write(0x10, &[0, 0, 0, 0]));
        assert!(bus.read(0x11, &mut [0, 0, 0, 0]));
        assert!(bus.write(0x11, &[0, 0, 0, 0]));
        assert!(bus.read(0x16, &mut [0, 0, 0, 0]));
        assert!(bus.write(0x16, &[0, 0, 0, 0]));
        assert!(!bus.read(0x20, &mut [0, 0, 0, 0]));
        assert!(!bus.write(0x20, &[0, 0, 0, 0]));
        assert!(!bus.read(0x06, &mut [0, 0, 0, 0]));
        assert!(!bus.write(0x06, &[0, 0, 0, 0]));
    }

    #[test]
    fn bus_read_write_values() {
        let mut bus = Bus::new();
        let dummy = Arc::new(Mutex::new(ConstantDevice));
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_ok());

        let mut values = [0, 1, 2, 3];
        assert!(bus.read(0x10, &mut values));
        assert_eq!(values, [0, 1, 2, 3]);
        assert!(bus.write(0x10, &values));
        assert!(bus.read(0x15, &mut values));
        assert_eq!(values, [5, 6, 7, 8]);
        assert!(bus.write(0x15, &values));
    }

    #[test]
    fn busrange_cmp_and_clone() {
        assert_eq!(BusRange(0x10, 2), BusRange(0x10, 3));
        assert_eq!(BusRange(0x10, 2), BusRange(0x10, 2));

        assert!(BusRange(0x10, 2) < BusRange(0x12, 1));
        assert!(BusRange(0x10, 2) < BusRange(0x12, 3));

        let bus_range = BusRange(0x10, 2);
        assert_eq!(bus_range, bus_range);

        let mut bus = Bus::new();
        let mut data = [1, 2, 3, 4];
        assert!(bus
            .insert(Arc::new(Mutex::new(DummyDevice)), 0x10, 0x10)
            .is_ok());
        assert!(bus.write(0x10, &data));
        let bus_clone = bus.clone();
        assert!(bus.read(0x10, &mut data));
        assert_eq!(data, [1, 2, 3, 4]);
        assert!(bus_clone.read(0x10, &mut data));
        assert_eq!(data, [1, 2, 3, 4]);
    }

    #[test]
    fn test_display_error() {
        assert_eq!(
            format!("{}", Error::Overlap),
            "New device overlaps with an old device."
        );
    }
}
<commit_msg>devices: add comments to unit test<commit_after>// Copyright 2018 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0
//
// Portions Copyright 2017 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the THIRD-PARTY file.

//! Handles routing to devices in an address space.

use std::cmp::{Ord, Ordering, PartialEq, PartialOrd};
use std::collections::btree_map::BTreeMap;
use std::fmt;
use std::io;
use std::result;
use std::sync::{Arc, Mutex};

use crate::virtio::AsAny;

/// Trait for devices that respond to reads or writes in an arbitrary address space.
///
/// The device does not care where it exists in address space as each method is only given an offset
/// into its allocated portion of address space.
#[allow(unused_variables)]
pub trait BusDevice: AsAny + Send {
    /// Reads at `offset` from this device
    fn read(&mut self, offset: u64, data: &mut [u8]) {}
    /// Writes at `offset` into this device
    fn write(&mut self, offset: u64, data: &[u8]) {}
    /// Triggers the `irq_mask` interrupt on this device
    fn interrupt(&self, irq_mask: u32) -> io::Result<()> {
        Ok(())
    }
}

#[derive(Debug)]
pub enum Error {
    /// The insertion failed because the new device overlapped with an old device.
    Overlap,
}

impl fmt::Display for Error {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        use self::Error::*;

        match *self {
            Overlap => write!(f, "New device overlaps with an old device."),
        }
    }
}

pub type Result<T> = result::Result<T, Error>;

#[derive(Debug, Copy, Clone)]
struct BusRange(u64, u64);

impl Eq for BusRange {}

impl PartialEq for BusRange {
    fn eq(&self, other: &BusRange) -> bool {
        self.0 == other.0
    }
}

impl Ord for BusRange {
    fn cmp(&self, other: &BusRange) -> Ordering {
        self.0.cmp(&other.0)
    }
}

impl PartialOrd for BusRange {
    fn partial_cmp(&self, other: &BusRange) -> Option<Ordering> {
        self.0.partial_cmp(&other.0)
    }
}

/// A device container for routing reads and writes over some address space.
///
/// This doesn't have any restrictions on what kind of device or address space this applies to. The
/// only restriction is that no two devices can overlap in this address space.
#[derive(Clone, Default)]
pub struct Bus {
    devices: BTreeMap<BusRange, Arc<Mutex<dyn BusDevice>>>,
}

impl Bus {
    /// Constructs an a bus with an empty address space.
    pub fn new() -> Bus {
        Bus {
            devices: BTreeMap::new(),
        }
    }

    fn first_before(&self, addr: u64) -> Option<(BusRange, &Mutex<dyn BusDevice>)> {
        // for when we switch to rustc 1.17: self.devices.range(..addr).iter().rev().next()
        for (range, dev) in self.devices.iter().rev() {
            if range.0 <= addr {
                return Some((*range, dev));
            }
        }
        None
    }

    pub fn get_device(&self, addr: u64) -> Option<(u64, &Mutex<dyn BusDevice>)> {
        if let Some((BusRange(start, len), dev)) = self.first_before(addr) {
            let offset = addr - start;
            if offset < len {
                return Some((offset, dev));
            }
        }
        None
    }

    /// Puts the given device at the given address space.
    pub fn insert(&mut self, device: Arc<Mutex<dyn BusDevice>>, base: u64, len: u64) -> Result<()> {
        if len == 0 {
            return Err(Error::Overlap);
        }

        // Reject all cases where the new device's base is within an old device's range.
        if self.get_device(base).is_some() {
            return Err(Error::Overlap);
        }

        // The above check will miss an overlap in which the new device's base address is before the
        // range of another device. To catch that case, we search for a device with a range before
        // the new device's range's end. If there is no existing device in that range that starts
        // after the new device, then there will be no overlap.
        if let Some((BusRange(start, _), _)) = self.first_before(base + len - 1) {
            // Such a device only conflicts with the new device if it also starts after the new
            // device because of our initial `get_device` check above.
            if start >= base {
                return Err(Error::Overlap);
            }
        }

        if self.devices.insert(BusRange(base, len), device).is_some() {
            return Err(Error::Overlap);
        }

        Ok(())
    }

    /// Reads data from the device that owns the range containing `addr` and puts it into `data`.
    ///
    /// Returns true on success, otherwise `data` is untouched.
    pub fn read(&self, addr: u64, data: &mut [u8]) -> bool {
        if let Some((offset, dev)) = self.get_device(addr) {
            // OK to unwrap as lock() failing is a serious error condition and should panic.
            dev.lock()
                .expect("Failed to acquire device lock")
                .read(offset, data);
            true
        } else {
            false
        }
    }

    /// Writes `data` to the device that owns the range containing `addr`.
    ///
    /// Returns true on success, otherwise `data` is untouched.
    pub fn write(&self, addr: u64, data: &[u8]) -> bool {
        if let Some((offset, dev)) = self.get_device(addr) {
            // OK to unwrap as lock() failing is a serious error condition and should panic.
            dev.lock()
                .expect("Failed to acquire device lock")
                .write(offset, data);
            true
        } else {
            false
        }
    }
}

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

    struct DummyDevice;
    impl BusDevice for DummyDevice {}

    struct ConstantDevice;
    impl BusDevice for ConstantDevice {
        fn read(&mut self, offset: u64, data: &mut [u8]) {
            for (i, v) in data.iter_mut().enumerate() {
                *v = (offset as u8) + (i as u8);
            }
        }

        fn write(&mut self, offset: u64, data: &[u8]) {
            for (i, v) in data.iter().enumerate() {
                assert_eq!(*v, (offset as u8) + (i as u8))
            }
        }
    }

    #[test]
    fn bus_insert() {
        let mut bus = Bus::new();
        let dummy = Arc::new(Mutex::new(DummyDevice));
        // Insert len should not be 0.
        assert!(bus.insert(dummy.clone(), 0x10, 0).is_err());
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_ok());

        let result = bus.insert(dummy.clone(), 0x0f, 0x10);
        // This overlaps the address space of the existing bus device at 0x10.
        assert!(result.is_err());
        assert_eq!(format!("{:?}", result), "Err(Overlap)");

        // This overlaps the address space of the existing bus device at 0x10.
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_err());
        // This overlaps the address space of the existing bus device at 0x10.
        assert!(bus.insert(dummy.clone(), 0x10, 0x15).is_err());
        // This overlaps the address space of the existing bus device at 0x10.
        assert!(bus.insert(dummy.clone(), 0x12, 0x15).is_err());
        // This overlaps the address space of the existing bus device at 0x10.
        assert!(bus.insert(dummy.clone(), 0x12, 0x01).is_err());
        // This overlaps the address space of the existing bus device at 0x10.
        assert!(bus.insert(dummy.clone(), 0x0, 0x20).is_err());
        assert!(bus.insert(dummy.clone(), 0x20, 0x05).is_ok());
        assert!(bus.insert(dummy.clone(), 0x25, 0x05).is_ok());
        assert!(bus.insert(dummy.clone(), 0x0, 0x10).is_ok());
    }

    #[test]
    fn bus_read_write() {
        let mut bus = Bus::new();
        let dummy = Arc::new(Mutex::new(DummyDevice));
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_ok());
        assert!(bus.read(0x10, &mut [0, 0, 0, 0]));
        assert!(bus.write(0x10, &[0, 0, 0, 0]));
        assert!(bus.read(0x11, &mut [0, 0, 0, 0]));
        assert!(bus.write(0x11, &[0, 0, 0, 0]));
        assert!(bus.read(0x16, &mut [0, 0, 0, 0]));
        assert!(bus.write(0x16, &[0, 0, 0, 0]));
        assert!(!bus.read(0x20, &mut [0, 0, 0, 0]));
        assert!(!bus.write(0x20, &[0, 0, 0, 0]));
        assert!(!bus.read(0x06, &mut [0, 0, 0, 0]));
        assert!(!bus.write(0x06, &[0, 0, 0, 0]));
    }

    #[test]
    fn bus_read_write_values() {
        let mut bus = Bus::new();
        let dummy = Arc::new(Mutex::new(ConstantDevice));
        assert!(bus.insert(dummy.clone(), 0x10, 0x10).is_ok());

        let mut values = [0, 1, 2, 3];
        assert!(bus.read(0x10, &mut values));
        assert_eq!(values, [0, 1, 2, 3]);
        assert!(bus.write(0x10, &values));
        assert!(bus.read(0x15, &mut values));
        assert_eq!(values, [5, 6, 7, 8]);
        assert!(bus.write(0x15, &values));
    }

    #[test]
    fn busrange_cmp_and_clone() {
        assert_eq!(BusRange(0x10, 2), BusRange(0x10, 3));
        assert_eq!(BusRange(0x10, 2), BusRange(0x10, 2));

        assert!(BusRange(0x10, 2) < BusRange(0x12, 1));
        assert!(BusRange(0x10, 2) < BusRange(0x12, 3));

        let bus_range = BusRange(0x10, 2);
        assert_eq!(bus_range, bus_range);

        let mut bus = Bus::new();
        let mut data = [1, 2, 3, 4];
        assert!(bus
            .insert(Arc::new(Mutex::new(DummyDevice)), 0x10, 0x10)
            .is_ok());
        assert!(bus.write(0x10, &data));
        let bus_clone = bus.clone();
        assert!(bus.read(0x10, &mut data));
        assert_eq!(data, [1, 2, 3, 4]);
        assert!(bus_clone.read(0x10, &mut data));
        assert_eq!(data, [1, 2, 3, 4]);
    }

    #[test]
    fn test_display_error() {
        assert_eq!(
            format!("{}", Error::Overlap),
            "New device overlaps with an old device."
        );
    }
}
<|endoftext|>


<commit_before>use crate::FileTime;
use std::ffi::CString;
use std::fs;
use std::io;
use std::os::unix::prelude::*;
use std::path::Path;

#[allow(dead_code)]
pub fn set_file_times(p: &Path, atime: FileTime, mtime: FileTime) -> io::Result<()> {
    set_times(p, Some(atime), Some(mtime), false)
}

#[allow(dead_code)]
pub fn set_file_mtime(p: &Path, mtime: FileTime) -> io::Result<()> {
    set_times(p, None, Some(mtime), false)
}

#[allow(dead_code)]
pub fn set_file_atime(p: &Path, atime: FileTime) -> io::Result<()> {
    set_times(p, Some(atime), None, false)
}

#[allow(dead_code)]
pub fn set_file_handle_times(
    f: &fs::File,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
) -> io::Result<()> {
    let (atime, mtime) = match get_times(atime, mtime, || f.metadata())? {
        Some(pair) => pair,
        None => return Ok(()),
    };
    let times = [to_timeval(&atime), to_timeval(&mtime)];
    let rc = unsafe { libc::futimes(f.as_raw_fd(), times.as_ptr()) };
    return if rc == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    };
}

fn get_times(
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
    current: impl FnOnce() -> io::Result<fs::Metadata>,
) -> io::Result<Option<(FileTime, FileTime)>> {
    let pair = match (atime, mtime) {
        (Some(a), Some(b)) => (a, b),
        (None, None) => return Ok(None),
        (Some(a), None) => {
            let meta = current()?;
            (a, FileTime::from_last_modification_time(&meta))
        }
        (None, Some(b)) => {
            let meta = current()?;
            (FileTime::from_last_access_time(&meta), b)
        }
    };
    Ok(Some(pair))
}

#[allow(dead_code)]
pub fn set_symlink_file_times(p: &Path, atime: FileTime, mtime: FileTime) -> io::Result<()> {
    set_times(p, Some(atime), Some(mtime), true)
}

pub fn set_times(
    p: &Path,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
    symlink: bool,
) -> io::Result<()> {
    let (atime, mtime) = match get_times(atime, mtime, || p.metadata())? {
        Some(pair) => pair,
        None => return Ok(()),
    };
    let p = CString::new(p.as_os_str().as_bytes())?;
    let times = [to_timeval(&atime), to_timeval(&mtime)];
    let rc = unsafe {
        if symlink {
            libc::lutimes(p.as_ptr(), times.as_ptr())
        } else {
            libc::utimes(p.as_ptr(), times.as_ptr())
        }
    };
    return if rc == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    };
}

fn to_timeval(ft: &FileTime) -> libc::timeval {
    libc::timeval {
        tv_sec: ft.seconds() as libc::time_t,
        tv_usec: (ft.nanoseconds() / 1000) as libc::suseconds_t,
    }
}
<commit_msg>Adding support for uclibc's funtimens in set_file_handle_times (#49)<commit_after>use crate::FileTime;
use std::ffi::CString;
use std::fs;
use std::io;
use std::os::unix::prelude::*;
use std::path::Path;

#[allow(dead_code)]
pub fn set_file_times(p: &Path, atime: FileTime, mtime: FileTime) -> io::Result<()> {
    set_times(p, Some(atime), Some(mtime), false)
}

#[allow(dead_code)]
pub fn set_file_mtime(p: &Path, mtime: FileTime) -> io::Result<()> {
    set_times(p, None, Some(mtime), false)
}

#[allow(dead_code)]
pub fn set_file_atime(p: &Path, atime: FileTime) -> io::Result<()> {
    set_times(p, Some(atime), None, false)
}

#[cfg(not(target_env = "uclibc"))]
#[allow(dead_code)]
pub fn set_file_handle_times(
    f: &fs::File,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
) -> io::Result<()> {
    let (atime, mtime) = match get_times(atime, mtime, || f.metadata())? {
        Some(pair) => pair,
        None => return Ok(()),
    };
    let times = [to_timeval(&atime), to_timeval(&mtime)];
    let rc = unsafe { libc::futimes(f.as_raw_fd(), times.as_ptr()) };
    return if rc == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    };
}

#[cfg(target_env = "uclibc")]
#[allow(dead_code)]
pub fn set_file_handle_times(
    f: &fs::File,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
) -> io::Result<()> {
    let (atime, mtime) = match get_times(atime, mtime, || f.metadata())? {
        Some(pair) => pair,
        None => return Ok(()),
    };
    let times = [to_timespec(&atime), to_timespec(&mtime)];
    let rc = unsafe { libc::futimens(f.as_raw_fd(), times.as_ptr()) };
    return if rc == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    };
}

fn get_times(
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
    current: impl FnOnce() -> io::Result<fs::Metadata>,
) -> io::Result<Option<(FileTime, FileTime)>> {
    let pair = match (atime, mtime) {
        (Some(a), Some(b)) => (a, b),
        (None, None) => return Ok(None),
        (Some(a), None) => {
            let meta = current()?;
            (a, FileTime::from_last_modification_time(&meta))
        }
        (None, Some(b)) => {
            let meta = current()?;
            (FileTime::from_last_access_time(&meta), b)
        }
    };
    Ok(Some(pair))
}

#[allow(dead_code)]
pub fn set_symlink_file_times(p: &Path, atime: FileTime, mtime: FileTime) -> io::Result<()> {
    set_times(p, Some(atime), Some(mtime), true)
}

pub fn set_times(
    p: &Path,
    atime: Option<FileTime>,
    mtime: Option<FileTime>,
    symlink: bool,
) -> io::Result<()> {
    let (atime, mtime) = match get_times(atime, mtime, || p.metadata())? {
        Some(pair) => pair,
        None => return Ok(()),
    };
    let p = CString::new(p.as_os_str().as_bytes())?;
    let times = [to_timeval(&atime), to_timeval(&mtime)];
    let rc = unsafe {
        if symlink {
            libc::lutimes(p.as_ptr(), times.as_ptr())
        } else {
            libc::utimes(p.as_ptr(), times.as_ptr())
        }
    };
    return if rc == 0 {
        Ok(())
    } else {
        Err(io::Error::last_os_error())
    };
}

fn to_timeval(ft: &FileTime) -> libc::timeval {
    libc::timeval {
        tv_sec: ft.seconds() as libc::time_t,
        tv_usec: (ft.nanoseconds() / 1000) as libc::suseconds_t,
    }
}

fn to_timespec(ft: &FileTime) -> libc::timespec {
    libc::timespec {
        tv_sec: ft.seconds() as libc::time_t,
        #[cfg(all(target_arch = "x86_64", target_pointer_width = "32"))]
        tv_nsec: (ft.nanoseconds()) as i64,
        #[cfg(not(all(target_arch = "x86_64", target_pointer_width = "32")))]
        tv_nsec: (ft.nanoseconds()) as libc::c_long,
    }
}
<|endoftext|>


<commit_before>//! This crate should eventually represent the structure at this repo:
//!
//! https://github.com/eth2-clients/eth2-testnets/tree/master/nimbus/testnet1
//!
//! It is not accurate at the moment, we include extra files and we also don't support a few
//! others. We are unable to conform to the repo until we have the following PR merged:
//!
//! https://github.com/sigp/lighthouse/pull/605

use eth2_libp2p::Enr;
use ssz::{Decode, Encode};
use std::fs::{create_dir_all, File};
use std::io::{Read, Write};
use std::path::PathBuf;
use types::{Address, BeaconState, EthSpec, YamlConfig};

pub const ADDRESS_FILE: &str = "deposit_contract.txt";
pub const DEPLOY_BLOCK_FILE: &str = "deploy_block.txt";
pub const BOOT_ENR_FILE: &str = "boot_enr.yaml";
pub const GENESIS_STATE_FILE: &str = "genesis.ssz";
pub const YAML_CONFIG_FILE: &str = "config.yaml";

pub const HARDCODED_YAML_CONFIG: &[u8] = include_bytes!("../testnet5/config.yaml");
pub const HARDCODED_DEPLOY_BLOCK: &[u8] = include_bytes!("../testnet5/deploy_block.txt");
pub const HARDCODED_DEPOSIT_CONTRACT: &[u8] = include_bytes!("../testnet5/deposit_contract.txt");
pub const HARDCODED_GENESIS_STATE: &[u8] = include_bytes!("../testnet5/genesis.ssz");
pub const HARDCODED_BOOT_ENR: &[u8] = include_bytes!("../testnet5/boot_enr.yaml");

/// Specifies an Eth2 testnet.
///
/// See the crate-level documentation for more details.
#[derive(Clone, PartialEq, Debug)]
pub struct Eth2TestnetConfig<E: EthSpec> {
    pub deposit_contract_address: String,
    pub deposit_contract_deploy_block: u64,
    pub boot_enr: Option<Vec<Enr>>,
    pub genesis_state: Option<BeaconState<E>>,
    pub yaml_config: Option<YamlConfig>,
}

impl<E: EthSpec> Eth2TestnetConfig<E> {
    // Creates the `Eth2TestnetConfig` that was included in the binary at compile time. This can be
    // considered the default Lighthouse testnet.
    //
    // Returns an error if those included bytes are invalid (this is unlikely).
    pub fn hard_coded() -> Result<Self, String> {
        Ok(Self {
            deposit_contract_address: serde_yaml::from_reader(HARDCODED_DEPOSIT_CONTRACT)
                .map_err(|e| format!("Unable to parse contract address: {:?}", e))?,
            deposit_contract_deploy_block: serde_yaml::from_reader(HARDCODED_DEPLOY_BLOCK)
                .map_err(|e| format!("Unable to parse deploy block: {:?}", e))?,
            boot_enr: Some(
                serde_yaml::from_reader(HARDCODED_BOOT_ENR)
                    .map_err(|e| format!("Unable to parse boot enr: {:?}", e))?,
            ),
            genesis_state: Some(
                BeaconState::from_ssz_bytes(HARDCODED_GENESIS_STATE)
                    .map_err(|e| format!("Unable to parse genesis state: {:?}", e))?,
            ),
            yaml_config: Some(
                serde_yaml::from_reader(HARDCODED_YAML_CONFIG)
                    .map_err(|e| format!("Unable to parse genesis state: {:?}", e))?,
            ),
        })
    }

    // Write the files to the directory, only if the directory doesn't already exist.
    pub fn write_to_file(&self, base_dir: PathBuf) -> Result<(), String> {
        if base_dir.exists() {
            return Err("Testnet directory already exists".to_string());
        }

        self.force_write_to_file(base_dir)
    }

    // Write the files to the directory, even if the directory already exists.
    pub fn force_write_to_file(&self, base_dir: PathBuf) -> Result<(), String> {
        create_dir_all(&base_dir)
            .map_err(|e| format!("Unable to create testnet directory: {:?}", e))?;

        macro_rules! write_to_yaml_file {
            ($file: ident, $variable: expr) => {
                File::create(base_dir.join($file))
                    .map_err(|e| format!("Unable to create {}: {:?}", $file, e))
                    .and_then(|mut file| {
                        let yaml = serde_yaml::to_string(&$variable)
                            .map_err(|e| format!("Unable to YAML encode {}: {:?}", $file, e))?;

                        // Remove the doc header from the YAML file.
                        //
                        // This allows us to play nice with other clients that are expecting
                        // plain-text, not YAML.
                        let no_doc_header = if yaml.starts_with("---\n") {
                            &yaml[4..]
                        } else {
                            &yaml
                        };

                        file.write_all(no_doc_header.as_bytes())
                            .map_err(|e| format!("Unable to write {}: {:?}", $file, e))
                    })?;
            };
        }

        write_to_yaml_file!(ADDRESS_FILE, self.deposit_contract_address);
        write_to_yaml_file!(DEPLOY_BLOCK_FILE, self.deposit_contract_deploy_block);

        if let Some(boot_enr) = &self.boot_enr {
            write_to_yaml_file!(BOOT_ENR_FILE, boot_enr);
        }

        if let Some(yaml_config) = &self.yaml_config {
            write_to_yaml_file!(YAML_CONFIG_FILE, yaml_config);
        }

        // The genesis state is a special case because it uses SSZ, not YAML.
        if let Some(genesis_state) = &self.genesis_state {
            let file = base_dir.join(GENESIS_STATE_FILE);

            File::create(&file)
                .map_err(|e| format!("Unable to create {:?}: {:?}", file, e))
                .and_then(|mut file| {
                    file.write_all(&genesis_state.as_ssz_bytes())
                        .map_err(|e| format!("Unable to write {:?}: {:?}", file, e))
                })?;
        }

        Ok(())
    }

    pub fn load(base_dir: PathBuf) -> Result<Self, String> {
        macro_rules! load_from_file {
            ($file: ident) => {
                File::open(base_dir.join($file))
                    .map_err(|e| format!("Unable to open {}: {:?}", $file, e))
                    .and_then(|file| {
                        serde_yaml::from_reader(file)
                            .map_err(|e| format!("Unable to parse {}: {:?}", $file, e))
                    })?;
            };
        }

        macro_rules! optional_load_from_file {
            ($file: ident) => {
                if base_dir.join($file).exists() {
                    Some(load_from_file!($file))
                } else {
                    None
                }
            };
        }

        let deposit_contract_address = load_from_file!(ADDRESS_FILE);
        let deposit_contract_deploy_block = load_from_file!(DEPLOY_BLOCK_FILE);
        let boot_enr = optional_load_from_file!(BOOT_ENR_FILE);
        let yaml_config = optional_load_from_file!(YAML_CONFIG_FILE);

        // The genesis state is a special case because it uses SSZ, not YAML.
        let file = base_dir.join(GENESIS_STATE_FILE);
        let genesis_state = if base_dir.join(&file).exists() {
            Some(
                File::open(base_dir.join(&file))
                    .map_err(|e| format!("Unable to open {:?}: {:?}", file, e))
                    .and_then(|mut file| {
                        let mut bytes = vec![];
                        file.read_to_end(&mut bytes)
                            .map_err(|e| format!("Unable to read {:?}: {:?}", file, e))?;

                        BeaconState::from_ssz_bytes(&bytes)
                            .map_err(|e| format!("Unable to SSZ decode {:?}: {:?}", file, e))
                    })?,
            )
        } else {
            None
        };

        Ok(Self {
            deposit_contract_address,
            deposit_contract_deploy_block,
            boot_enr,
            genesis_state,
            yaml_config,
        })
    }

    pub fn deposit_contract_address(&self) -> Result<Address, String> {
        if self.deposit_contract_address.starts_with("0x") {
            self.deposit_contract_address[2..]
                .parse()
                .map_err(|e| format!("Corrupted address, unable to parse: {:?}", e))
        } else {
            Err("Corrupted address, must start with 0x".to_string())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempdir::TempDir;
    use types::{Eth1Data, Hash256, MainnetEthSpec, YamlConfig};

    type E = MainnetEthSpec;

    #[test]
    fn hard_coded_works() {
        let dir: Eth2TestnetConfig<E> =
            Eth2TestnetConfig::hard_coded().expect("should decode hard_coded params");

        assert!(dir.boot_enr.is_some());
        assert!(dir.genesis_state.is_some());
        assert!(dir.yaml_config.is_some());
    }

    #[test]
    fn round_trip() {
        let spec = &E::default_spec();

        let eth1_data = Eth1Data {
            deposit_root: Hash256::zero(),
            deposit_count: 0,
            block_hash: Hash256::zero(),
        };

        // TODO: figure out how to generate ENR and add some here.
        let boot_enr = None;
        let genesis_state = Some(BeaconState::new(42, eth1_data, spec));
        let yaml_config = Some(YamlConfig::from_spec::<E>(spec));

        do_test::<E>(boot_enr, genesis_state, yaml_config);
        do_test::<E>(None, None, None);
    }

    fn do_test<E: EthSpec>(
        boot_enr: Option<Vec<Enr>>,
        genesis_state: Option<BeaconState<E>>,
        yaml_config: Option<YamlConfig>,
    ) {
        let temp_dir = TempDir::new("eth2_testnet_test").expect("should create temp dir");
        let base_dir = temp_dir.path().join("my_testnet");
        let deposit_contract_address = "0xBB9bc244D798123fDe783fCc1C72d3Bb8C189413".to_string();
        let deposit_contract_deploy_block = 42;

        let testnet: Eth2TestnetConfig<E> = Eth2TestnetConfig {
            deposit_contract_address,
            deposit_contract_deploy_block,
            boot_enr,
            genesis_state,
            yaml_config,
        };

        testnet
            .write_to_file(base_dir.clone())
            .expect("should write to file");

        let decoded = Eth2TestnetConfig::load(base_dir).expect("should load struct");

        assert_eq!(testnet, decoded, "should decode as encoded");
    }
}
<commit_msg>Ensure correct path for genesis state file (#951)<commit_after>//! This crate should eventually represent the structure at this repo:
//!
//! https://github.com/eth2-clients/eth2-testnets/tree/master/nimbus/testnet1
//!
//! It is not accurate at the moment, we include extra files and we also don't support a few
//! others. We are unable to conform to the repo until we have the following PR merged:
//!
//! https://github.com/sigp/lighthouse/pull/605

use eth2_libp2p::Enr;
use ssz::{Decode, Encode};
use std::fs::{create_dir_all, File};
use std::io::{Read, Write};
use std::path::PathBuf;
use types::{Address, BeaconState, EthSpec, YamlConfig};

pub const ADDRESS_FILE: &str = "deposit_contract.txt";
pub const DEPLOY_BLOCK_FILE: &str = "deploy_block.txt";
pub const BOOT_ENR_FILE: &str = "boot_enr.yaml";
pub const GENESIS_STATE_FILE: &str = "genesis.ssz";
pub const YAML_CONFIG_FILE: &str = "config.yaml";

pub const HARDCODED_YAML_CONFIG: &[u8] = include_bytes!("../testnet5/config.yaml");
pub const HARDCODED_DEPLOY_BLOCK: &[u8] = include_bytes!("../testnet5/deploy_block.txt");
pub const HARDCODED_DEPOSIT_CONTRACT: &[u8] = include_bytes!("../testnet5/deposit_contract.txt");
pub const HARDCODED_GENESIS_STATE: &[u8] = include_bytes!("../testnet5/genesis.ssz");
pub const HARDCODED_BOOT_ENR: &[u8] = include_bytes!("../testnet5/boot_enr.yaml");

/// Specifies an Eth2 testnet.
///
/// See the crate-level documentation for more details.
#[derive(Clone, PartialEq, Debug)]
pub struct Eth2TestnetConfig<E: EthSpec> {
    pub deposit_contract_address: String,
    pub deposit_contract_deploy_block: u64,
    pub boot_enr: Option<Vec<Enr>>,
    pub genesis_state: Option<BeaconState<E>>,
    pub yaml_config: Option<YamlConfig>,
}

impl<E: EthSpec> Eth2TestnetConfig<E> {
    // Creates the `Eth2TestnetConfig` that was included in the binary at compile time. This can be
    // considered the default Lighthouse testnet.
    //
    // Returns an error if those included bytes are invalid (this is unlikely).
    pub fn hard_coded() -> Result<Self, String> {
        Ok(Self {
            deposit_contract_address: serde_yaml::from_reader(HARDCODED_DEPOSIT_CONTRACT)
                .map_err(|e| format!("Unable to parse contract address: {:?}", e))?,
            deposit_contract_deploy_block: serde_yaml::from_reader(HARDCODED_DEPLOY_BLOCK)
                .map_err(|e| format!("Unable to parse deploy block: {:?}", e))?,
            boot_enr: Some(
                serde_yaml::from_reader(HARDCODED_BOOT_ENR)
                    .map_err(|e| format!("Unable to parse boot enr: {:?}", e))?,
            ),
            genesis_state: Some(
                BeaconState::from_ssz_bytes(HARDCODED_GENESIS_STATE)
                    .map_err(|e| format!("Unable to parse genesis state: {:?}", e))?,
            ),
            yaml_config: Some(
                serde_yaml::from_reader(HARDCODED_YAML_CONFIG)
                    .map_err(|e| format!("Unable to parse genesis state: {:?}", e))?,
            ),
        })
    }

    // Write the files to the directory, only if the directory doesn't already exist.
    pub fn write_to_file(&self, base_dir: PathBuf) -> Result<(), String> {
        if base_dir.exists() {
            return Err("Testnet directory already exists".to_string());
        }

        self.force_write_to_file(base_dir)
    }

    // Write the files to the directory, even if the directory already exists.
    pub fn force_write_to_file(&self, base_dir: PathBuf) -> Result<(), String> {
        create_dir_all(&base_dir)
            .map_err(|e| format!("Unable to create testnet directory: {:?}", e))?;

        macro_rules! write_to_yaml_file {
            ($file: ident, $variable: expr) => {
                File::create(base_dir.join($file))
                    .map_err(|e| format!("Unable to create {}: {:?}", $file, e))
                    .and_then(|mut file| {
                        let yaml = serde_yaml::to_string(&$variable)
                            .map_err(|e| format!("Unable to YAML encode {}: {:?}", $file, e))?;

                        // Remove the doc header from the YAML file.
                        //
                        // This allows us to play nice with other clients that are expecting
                        // plain-text, not YAML.
                        let no_doc_header = if yaml.starts_with("---\n") {
                            &yaml[4..]
                        } else {
                            &yaml
                        };

                        file.write_all(no_doc_header.as_bytes())
                            .map_err(|e| format!("Unable to write {}: {:?}", $file, e))
                    })?;
            };
        }

        write_to_yaml_file!(ADDRESS_FILE, self.deposit_contract_address);
        write_to_yaml_file!(DEPLOY_BLOCK_FILE, self.deposit_contract_deploy_block);

        if let Some(boot_enr) = &self.boot_enr {
            write_to_yaml_file!(BOOT_ENR_FILE, boot_enr);
        }

        if let Some(yaml_config) = &self.yaml_config {
            write_to_yaml_file!(YAML_CONFIG_FILE, yaml_config);
        }

        // The genesis state is a special case because it uses SSZ, not YAML.
        if let Some(genesis_state) = &self.genesis_state {
            let file = base_dir.join(GENESIS_STATE_FILE);

            File::create(&file)
                .map_err(|e| format!("Unable to create {:?}: {:?}", file, e))
                .and_then(|mut file| {
                    file.write_all(&genesis_state.as_ssz_bytes())
                        .map_err(|e| format!("Unable to write {:?}: {:?}", file, e))
                })?;
        }

        Ok(())
    }

    pub fn load(base_dir: PathBuf) -> Result<Self, String> {
        macro_rules! load_from_file {
            ($file: ident) => {
                File::open(base_dir.join($file))
                    .map_err(|e| format!("Unable to open {}: {:?}", $file, e))
                    .and_then(|file| {
                        serde_yaml::from_reader(file)
                            .map_err(|e| format!("Unable to parse {}: {:?}", $file, e))
                    })?;
            };
        }

        macro_rules! optional_load_from_file {
            ($file: ident) => {
                if base_dir.join($file).exists() {
                    Some(load_from_file!($file))
                } else {
                    None
                }
            };
        }

        let deposit_contract_address = load_from_file!(ADDRESS_FILE);
        let deposit_contract_deploy_block = load_from_file!(DEPLOY_BLOCK_FILE);
        let boot_enr = optional_load_from_file!(BOOT_ENR_FILE);
        let yaml_config = optional_load_from_file!(YAML_CONFIG_FILE);

        // The genesis state is a special case because it uses SSZ, not YAML.
        let genesis_file_path = base_dir.join(GENESIS_STATE_FILE);
        let genesis_state = if genesis_file_path.exists() {
            Some(
                File::open(&genesis_file_path)
                    .map_err(|e| format!("Unable to open {:?}: {:?}", genesis_file_path, e))
                    .and_then(|mut file| {
                        let mut bytes = vec![];
                        file.read_to_end(&mut bytes)
                            .map_err(|e| format!("Unable to read {:?}: {:?}", file, e))?;

                        BeaconState::from_ssz_bytes(&bytes)
                            .map_err(|e| format!("Unable to SSZ decode {:?}: {:?}", file, e))
                    })?,
            )
        } else {
            None
        };

        Ok(Self {
            deposit_contract_address,
            deposit_contract_deploy_block,
            boot_enr,
            genesis_state,
            yaml_config,
        })
    }

    pub fn deposit_contract_address(&self) -> Result<Address, String> {
        if self.deposit_contract_address.starts_with("0x") {
            self.deposit_contract_address[2..]
                .parse()
                .map_err(|e| format!("Corrupted address, unable to parse: {:?}", e))
        } else {
            Err("Corrupted address, must start with 0x".to_string())
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tempdir::TempDir;
    use types::{Eth1Data, Hash256, MainnetEthSpec, YamlConfig};

    type E = MainnetEthSpec;

    #[test]
    fn hard_coded_works() {
        let dir: Eth2TestnetConfig<E> =
            Eth2TestnetConfig::hard_coded().expect("should decode hard_coded params");

        assert!(dir.boot_enr.is_some());
        assert!(dir.genesis_state.is_some());
        assert!(dir.yaml_config.is_some());
    }

    #[test]
    fn round_trip() {
        let spec = &E::default_spec();

        let eth1_data = Eth1Data {
            deposit_root: Hash256::zero(),
            deposit_count: 0,
            block_hash: Hash256::zero(),
        };

        // TODO: figure out how to generate ENR and add some here.
        let boot_enr = None;
        let genesis_state = Some(BeaconState::new(42, eth1_data, spec));
        let yaml_config = Some(YamlConfig::from_spec::<E>(spec));

        do_test::<E>(boot_enr, genesis_state, yaml_config);
        do_test::<E>(None, None, None);
    }

    fn do_test<E: EthSpec>(
        boot_enr: Option<Vec<Enr>>,
        genesis_state: Option<BeaconState<E>>,
        yaml_config: Option<YamlConfig>,
    ) {
        let temp_dir = TempDir::new("eth2_testnet_test").expect("should create temp dir");
        let base_dir = temp_dir.path().join("my_testnet");
        let deposit_contract_address = "0xBB9bc244D798123fDe783fCc1C72d3Bb8C189413".to_string();
        let deposit_contract_deploy_block = 42;

        let testnet: Eth2TestnetConfig<E> = Eth2TestnetConfig {
            deposit_contract_address,
            deposit_contract_deploy_block,
            boot_enr,
            genesis_state,
            yaml_config,
        };

        testnet
            .write_to_file(base_dir.clone())
            .expect("should write to file");

        let decoded = Eth2TestnetConfig::load(base_dir).expect("should load struct");

        assert_eq!(testnet, decoded, "should decode as encoded");
    }
}
<|endoftext|>


<commit_before>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

extern crate libc;

use std::mem;

use constants::LINE_SIZE;

#[repr(C)]
#[allow(missing_copy_implementations)]
pub struct GCHeader {
    reference_count: libc::size_t,
    spans_lines: bool,
    forwarded: bool,
    logged: bool,
    marked: bool,
    new: bool,
}

#[repr(C)]
#[allow(missing_copy_implementations)]
pub struct GCRTTI {
    object_size: libc::size_t,
    variables: libc::size_t,
}

#[repr(C)]
pub struct GCObject {
    header: GCHeader,
    rtti: *const GCRTTI
}

pub type GCObjectRef = *mut GCObject;

impl GCRTTI {
    pub fn new(object_size: uint, variables: uint) -> GCRTTI {
        return GCRTTI {
            object_size: object_size as libc::size_t,
            variables: variables as libc::size_t,
        };
    }

    pub fn object_size(&self) -> uint {
        return self.object_size as uint;
    }

    pub fn variables(&self) -> uint {
        return self.variables as uint;
    }
}

impl GCObject {
    pub fn new(rtti: *const GCRTTI, mark: bool) -> GCObject {
        debug!("GCobject::new(rtti={}, mark={}", rtti, mark);
        let size = unsafe{ (*rtti).object_size() };
        return GCObject {
            header: GCHeader {
                reference_count: 0,
                spans_lines: size > LINE_SIZE,
                forwarded: false,
                logged: false,
                marked: mark,
                new: true,
            },
            rtti: rtti,
        }
    }

    pub fn object_size(&self) -> uint {
        return unsafe{ (*self.rtti).object_size() };
    }

    pub fn children(&mut self) -> Vec<GCObjectRef> {
        let base: *const GCObjectRef = unsafe{ mem::transmute(&self.rtti) };
        let variables = unsafe{ (*self.rtti).variables() };
        return range(1, variables + 1).map(|i| unsafe{ *base.offset(i as int) })
                                      .collect();
    }

    pub fn decrement(&mut self) -> bool {
        if self.header.reference_count == 0 {
            return false;
        }
        self.header.reference_count -= 1;
        debug!("Decrement object {:p} to {}", self, self.header.reference_count);
        return self.header.reference_count == 0;
    }

    pub fn increment(&mut self) -> bool {
        self.header.reference_count += 1;
        debug!("Increment object {:p} to {} (new={})", self,
            self.header.reference_count, self.header.new);
        if self.header.new {
            self.header.new = false;
            return true;
        }
        return false;
    }

    pub fn set_logged(&mut self, new: bool) -> bool {
        debug!("Set object {:p} logged={}", self, new);
        let logged = self.header.logged;
        self.header.logged = new;
        return logged;
    }

    pub fn set_marked(&mut self, next: bool) -> bool {
        debug!("Set object {:p} marked={}", self, next);
        let marked = self.header.marked;
        self.header.marked = next;
        return marked == next;
    }

    pub fn is_marked(&self, next: bool) -> bool {
        return self.header.marked == next;
    }

    pub fn spans_lines(&self) -> bool {
        return self.header.spans_lines;
    }
}

<commit_msg>Reorganize methods for GCObject<commit_after>// Copyright (c) <2015> <lummax>
// Licensed under MIT (http://opensource.org/licenses/MIT)

extern crate libc;

use std::mem;

use constants::LINE_SIZE;

#[repr(C)]
#[allow(missing_copy_implementations)]
pub struct GCHeader {
    reference_count: libc::size_t,
    spans_lines: bool,
    forwarded: bool,
    logged: bool,
    marked: bool,
    new: bool,
}

#[repr(C)]
#[allow(missing_copy_implementations)]
pub struct GCRTTI {
    object_size: libc::size_t,
    variables: libc::size_t,
}

#[repr(C)]
pub struct GCObject {
    header: GCHeader,
    rtti: *const GCRTTI
}

pub type GCObjectRef = *mut GCObject;

impl GCRTTI {
    pub fn new(object_size: uint, variables: uint) -> GCRTTI {
        return GCRTTI {
            object_size: object_size as libc::size_t,
            variables: variables as libc::size_t,
        };
    }

    pub fn object_size(&self) -> uint {
        return self.object_size as uint;
    }

    pub fn variables(&self) -> uint {
        return self.variables as uint;
    }
}

impl GCObject {
    pub fn new(rtti: *const GCRTTI, mark: bool) -> GCObject {
        debug!("GCobject::new(rtti={}, mark={}", rtti, mark);
        let size = unsafe{ (*rtti).object_size() };
        return GCObject {
            header: GCHeader {
                reference_count: 0,
                spans_lines: size > LINE_SIZE,
                forwarded: false,
                logged: false,
                marked: mark,
                new: true,
            },
            rtti: rtti,
        }
    }

    pub fn set_logged(&mut self, new: bool) -> bool {
        debug!("Set object {:p} logged={}", self, new);
        let logged = self.header.logged;
        self.header.logged = new;
        return logged;
    }

    pub fn set_marked(&mut self, next: bool) -> bool {
        debug!("Set object {:p} marked={}", self, next);
        let marked = self.header.marked;
        self.header.marked = next;
        return marked == next;
    }

    pub fn is_marked(&self, next: bool) -> bool {
        return self.header.marked == next;
    }

    pub fn spans_lines(&self) -> bool {
        return self.header.spans_lines;
    }

    pub fn object_size(&self) -> uint {
        return unsafe{ (*self.rtti).object_size() };
    }

    pub fn decrement(&mut self) -> bool {
        if self.header.reference_count == 0 {
            return false;
        }
        self.header.reference_count -= 1;
        debug!("Decrement object {:p} to {}", self, self.header.reference_count);
        return self.header.reference_count == 0;
    }

    pub fn increment(&mut self) -> bool {
        self.header.reference_count += 1;
        debug!("Increment object {:p} to {} (new={})", self,
            self.header.reference_count, self.header.new);
        if self.header.new {
            self.header.new = false;
            return true;
        }
        return false;
    }

    pub fn children(&mut self) -> Vec<GCObjectRef> {
        let base: *const GCObjectRef = unsafe{ mem::transmute(&self.rtti) };
        let variables = unsafe{ (*self.rtti).variables() };
        return range(1, variables + 1).map(|i| unsafe{ *base.offset(i as int) })
                                      .collect();
    }

}

<|endoftext|>


<commit_before>
#![crate_name = "geoip"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_upper_case_globals,
        unused_qualifications)]

extern crate libc;
extern crate serialize;
extern crate "geoip-sys" as geoip_sys;

use libc::{c_char, c_int, c_ulong};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

enum Charset {
    UTF8 = 1
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

impl Copy for Options { }

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

impl Copy for DBType { }

pub struct GeoIp {
    db: geoip_sys::RawGeoIp
}

impl Copy for GeoIp { }

#[deriving(Decodable, Encodable)]
pub struct ASInfo {
    pub asn: uint,
    pub name: String,
    pub netmask: uint
}

#[deriving(Decodable, Encodable)]
pub struct CityInfo {
    pub country_code: Option<String>,
    pub country_code3: Option<String>,
    pub country_name: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
    pub postal_code: Option<String>,
    pub latitude: f32,
    pub longitude: f32,
    pub dma_code: uint,
    pub area_code: uint,
    pub charset: uint,
    pub continent_code: Option<String>,
    pub netmask: uint
}

unsafe fn maybe_string(c_str: *const c_char) -> Option<String> {
    c_str.as_ref().and_then(|opt| {
        CString::new(opt, false).as_str().map(|s| s.to_string())
    })
}

impl CityInfo {
    unsafe fn from_geoiprecord(res: &geoip_sys::GeoIpRecord) -> CityInfo {
        CityInfo {
            country_code: maybe_string(res.country_code),
            country_code3: maybe_string(res.country_code3),
            country_name: maybe_string(res.country_name),
            region: maybe_string(res.region),
            city: maybe_string(res.city),
            postal_code: maybe_string(res.postal_code),
            latitude: res.latitude as f32,
            longitude: res.longitude as f32,
            dma_code: res.dma_code as uint,
            area_code: res.area_code as uint,
            charset: res.charset as uint,
            continent_code: maybe_string(res.continent_code),
            netmask: res.netmask as uint
        }
    }
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}", self.asn, self.name)
    }
}

enum CNetworkIp {
    V4(c_ulong),
    V6(geoip_sys::In6Addr)
}

impl CNetworkIp {
    fn new(ip: IpAddr) -> CNetworkIp {
        match ip {
            Ipv4Addr(a, b, c, d) => {
                CNetworkIp::V4((a as c_ulong << 24) | (b as c_ulong << 16) |
                               (c as c_ulong << 8)  | (d as c_ulong))
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                CNetworkIp::V6([(a >> 8) as u8, a as u8,
                                (b >> 8) as u8, b as u8,
                                (c >> 8) as u8, c as u8,
                                (d >> 8) as u8, d as u8,
                                (e >> 8) as u8, e as u8,
                                (f >> 8) as u8, f as u8,
                                (g >> 8) as u8, g as u8,
                                (h >> 8) as u8, h as u8])
            }
        }
    }
}

impl GeoIp {
    pub fn open(path: &Path, options: Options) -> Result<GeoIp, String> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            geoip_sys::GeoIP_open(file.to_c_str().into_inner(),
                                  options as c_int)
        };
        if db.is_null() {
            return Err(format!("Can't open {}", file));
        }
        if unsafe { geoip_sys::GeoIP_set_charset(db, Charset::UTF8 as c_int)
        } != 0 {
            return Err("Can't set charset to UTF8".to_string());
        }
        Ok(GeoIp { db: db })
    }

    pub fn city_info_by_ip(&self, ip: IpAddr) -> Option<CityInfo> {
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe {
                geoip_sys::GeoIP_record_by_ipnum(self.db, ip) },
            CNetworkIp::V6(ip) => unsafe {
                geoip_sys::GeoIP_record_by_ipnum_v6(self.db, ip) }
        };

        if cres.is_null() { return None; }

        unsafe {
            let city_info = CityInfo::from_geoiprecord(&*cres);
            geoip_sys::GeoIPRecord_delete(cres);
            std::mem::forget(cres);
            Some(city_info)
        }
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let mut gl = geoip_sys::GeoIpLookup::new();
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe {
                geoip_sys::GeoIP_name_by_ipnum_gl(self.db, ip, &mut gl) },
            CNetworkIp::V6(ip) => unsafe {
                geoip_sys::GeoIP_name_by_ipnum_v6_gl(self.db, ip, &mut gl) }
        };

        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(1, ' ');
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    from_str::<uint>(asn.slice_from(2)).unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let as_info = ASInfo {
            asn: asn,
            name: name.to_string(),
            netmask: gl.netmask as uint
        };
        Some(as_info)
    }
}

impl Drop for GeoIp {
    fn drop(&mut self) {
        unsafe {
            geoip_sys::GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.as_slice().contains("Telefun"));
    assert!(res.netmask == 22);
}

#[test]
fn geoip_test_city() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoLiteCity.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("8.8.8.8").unwrap();
    let res = geoip.city_info_by_ip(ip).unwrap();
    assert!(res.city.unwrap().as_slice() == "Mountain View");
}
<commit_msg>from_str -> parse<commit_after>
#![crate_name = "geoip"]
#![crate_type = "rlib"]

#![warn(non_camel_case_types,
        non_upper_case_globals,
        unused_qualifications)]

extern crate libc;
extern crate serialize;
extern crate "geoip-sys" as geoip_sys;

use libc::{c_char, c_int, c_ulong};
use std::c_str::CString;
use std::fmt;
use std::io::net::ip::{IpAddr,Ipv4Addr,Ipv6Addr};

enum Charset {
    UTF8 = 1
}

pub enum Options {
    Standard = 0,
    MemoryCache = 1,
    CheckCache = 2,
    IndexCache = 4,
    MmapCache = 8
}

impl Copy for Options { }

pub enum DBType {
    CountryEdition = 1,
    RegionEditionRev0 = 7,
    CityEditionRev0 = 6,
    ORGEdition = 5,
    ISPEdition = 4,
    CityEditionRev1 = 2,
    RegionEditionRev1 = 3,
    ProxyEdition = 8,
    ASNUMEdition = 9,
    NetSpeedEdition = 10,
    DomainEdition = 11,
    CountryEditionV6 = 12,
    LocationAEdition = 13,
    AccuracyRadiusEdition = 14,
    LargeCountryEdition = 17,
    LargeCountryEditionV6 = 18,
    ASNumEditionV6 = 21,
    ISPEditionV6 = 22,
    ORGEditionV6 = 23,
    DomainEditionV6 = 24,
    LoctionAEditionV6 = 25,
    RegistrarEdition = 26,
    RegistrarEditionV6 = 27,
    UserTypeEdition = 28,
    UserTypeEditionV6 = 29,
    CityEditionRev1V6 = 30,
    CityEditionRev0V6 = 31,
    NetSpeedEditionRev1 = 32,
    NetSpeedEditionRev1V6 = 33,
    CountryConfEdition = 34,
    CityConfEdition = 35,
    RegionConfEdition = 36,
    PostalConfEdition = 37,
    AccuracyRadiusEditionV6 = 38
}

impl Copy for DBType { }

pub struct GeoIp {
    db: geoip_sys::RawGeoIp
}

impl Copy for GeoIp { }

#[deriving(Decodable, Encodable)]
pub struct ASInfo {
    pub asn: uint,
    pub name: String,
    pub netmask: uint
}

#[deriving(Decodable, Encodable)]
pub struct CityInfo {
    pub country_code: Option<String>,
    pub country_code3: Option<String>,
    pub country_name: Option<String>,
    pub region: Option<String>,
    pub city: Option<String>,
    pub postal_code: Option<String>,
    pub latitude: f32,
    pub longitude: f32,
    pub dma_code: uint,
    pub area_code: uint,
    pub charset: uint,
    pub continent_code: Option<String>,
    pub netmask: uint
}

unsafe fn maybe_string(c_str: *const c_char) -> Option<String> {
    c_str.as_ref().and_then(|opt| {
        CString::new(opt, false).as_str().map(|s| s.to_string())
    })
}

impl CityInfo {
    unsafe fn from_geoiprecord(res: &geoip_sys::GeoIpRecord) -> CityInfo {
        CityInfo {
            country_code: maybe_string(res.country_code),
            country_code3: maybe_string(res.country_code3),
            country_name: maybe_string(res.country_name),
            region: maybe_string(res.region),
            city: maybe_string(res.city),
            postal_code: maybe_string(res.postal_code),
            latitude: res.latitude as f32,
            longitude: res.longitude as f32,
            dma_code: res.dma_code as uint,
            area_code: res.area_code as uint,
            charset: res.charset as uint,
            continent_code: maybe_string(res.continent_code),
            netmask: res.netmask as uint
        }
    }
}

impl fmt::Show for ASInfo {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}\t{}", self.asn, self.name)
    }
}

enum CNetworkIp {
    V4(c_ulong),
    V6(geoip_sys::In6Addr)
}

impl CNetworkIp {
    fn new(ip: IpAddr) -> CNetworkIp {
        match ip {
            Ipv4Addr(a, b, c, d) => {
                CNetworkIp::V4((a as c_ulong << 24) | (b as c_ulong << 16) |
                               (c as c_ulong << 8)  | (d as c_ulong))
            },
            Ipv6Addr(a, b, c, d, e, f, g, h) => {
                CNetworkIp::V6([(a >> 8) as u8, a as u8,
                                (b >> 8) as u8, b as u8,
                                (c >> 8) as u8, c as u8,
                                (d >> 8) as u8, d as u8,
                                (e >> 8) as u8, e as u8,
                                (f >> 8) as u8, f as u8,
                                (g >> 8) as u8, g as u8,
                                (h >> 8) as u8, h as u8])
            }
        }
    }
}

impl GeoIp {
    pub fn open(path: &Path, options: Options) -> Result<GeoIp, String> {
        let file = match path.as_str() {
            None => return Err(format!("Invalid path {}", path.display())),
            Some(file) => file
        };
        let db = unsafe {
            geoip_sys::GeoIP_open(file.to_c_str().into_inner(),
                                  options as c_int)
        };
        if db.is_null() {
            return Err(format!("Can't open {}", file));
        }
        if unsafe { geoip_sys::GeoIP_set_charset(db, Charset::UTF8 as c_int)
        } != 0 {
            return Err("Can't set charset to UTF8".to_string());
        }
        Ok(GeoIp { db: db })
    }

    pub fn city_info_by_ip(&self, ip: IpAddr) -> Option<CityInfo> {
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe {
                geoip_sys::GeoIP_record_by_ipnum(self.db, ip) },
            CNetworkIp::V6(ip) => unsafe {
                geoip_sys::GeoIP_record_by_ipnum_v6(self.db, ip) }
        };

        if cres.is_null() { return None; }

        unsafe {
            let city_info = CityInfo::from_geoiprecord(&*cres);
            geoip_sys::GeoIPRecord_delete(cres);
            std::mem::forget(cres);
            Some(city_info)
        }
    }

    pub fn as_info_by_ip(&self, ip: IpAddr) -> Option<ASInfo> {
        let mut gl = geoip_sys::GeoIpLookup::new();
        let cres = match CNetworkIp::new(ip) {
            CNetworkIp::V4(ip) => unsafe {
                geoip_sys::GeoIP_name_by_ipnum_gl(self.db, ip, &mut gl) },
            CNetworkIp::V6(ip) => unsafe {
                geoip_sys::GeoIP_name_by_ipnum_v6_gl(self.db, ip, &mut gl) }
        };

        if cres.is_null() {
            return None;
        }
        let description_cstr = unsafe { CString::new(cres, true) };
        let description = match description_cstr.as_str() {
            None => return None,
            Some(description) => description
        };
        let mut di = description.splitn(1, ' ');
        let asn = match di.next() {
            None => return None,
            Some(asn) => {
                if ! asn.starts_with("AS") {
                    return None
                } else {
                    asn.slice_from(2).parse::<uint>().unwrap()
                }
            }
        };
        let name = di.next().unwrap_or("(none)");
        let as_info = ASInfo {
            asn: asn,
            name: name.to_string(),
            netmask: gl.netmask as uint
        };
        Some(as_info)
    }
}

impl Drop for GeoIp {
    fn drop(&mut self) {
        unsafe {
            geoip_sys::GeoIP_delete(self.db);
        }
    }
}

#[test]
fn geoip_test_basic() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoIPASNum.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("91.203.184.192").unwrap();
    let res = geoip.as_info_by_ip(ip).unwrap();
    assert!(res.asn == 41064);
    assert!(res.name.as_slice().contains("Telefun"));
    assert!(res.netmask == 22);
}

#[test]
fn geoip_test_city() {
    let geoip = match GeoIp::open(&from_str("/opt/geoip/GeoLiteCity.dat").unwrap(), Options::MemoryCache) {
        Err(err) => panic!(err),
        Ok(geoip) => geoip
    };
    let ip = from_str("8.8.8.8").unwrap();
    let res = geoip.city_info_by_ip(ip).unwrap();
    assert!(res.city.unwrap().as_slice() == "Mountain View");
}
<|endoftext|>