<commit_before>// Copyright 2014 Dmitry "Divius" Tantsur <divius.inside@gmail.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//

//! KRPC DHT service as described in
//! [BEP 0005](http://www.bittorrent.org/beps/bep_0005.html).
//!
//! Create as:
//! ```
//! let service = dht::bt::KRpcService::new(this_node);
//! ```

use std::io::IoResult;
use std::sync;

use super::super::base;
use super::super::knodetable;

use super::udpwrapper;

/// Implementation of basic KRPC DHT on which BitTorrent DHT is based.
///
/// No peer retrival is supported: just finding and pinging nodes.
pub struct KRpcService<TNodeTable: base::GenericNodeTable,
                       TSocket: udpwrapper::GenericSocketWrapper> {
    this_node: base::Node,
    node_table: sync::Arc<sync::RWLock<TNodeTable>>,
    socket: TSocket,
}

/// Default kind of KRpc service.
pub type DefaultKRpcService = KRpcService<knodetable::KNodeTable, udpwrapper::UdpSocketWrapper>;

fn handle_incoming<TSocket: udpwrapper::GenericSocketWrapper>(socket: TSocket) {
    // TODO(divius): implement
    drop(socket);
}

impl KRpcService<knodetable::KNodeTable, udpwrapper::UdpSocketWrapper> {
    /// New service with default node table.
    pub fn new_default(this_node: base::Node) -> IoResult<DefaultKRpcService> {
        let node_table = knodetable::KNodeTable::new(this_node.id.clone());
        let socket = try!(udpwrapper::UdpSocketWrapper::new(&this_node));
        KRpcService::new(this_node, node_table, socket)
    }
}

impl<TNodeTable: base::GenericNodeTable, TSocket: udpwrapper::GenericSocketWrapper>
KRpcService<TNodeTable, TSocket> {
    /// New service with given node table and socket.
    pub fn new(this_node: base::Node, node_table: TNodeTable, socket: TSocket)
            -> IoResult<KRpcService<TNodeTable, TSocket>> {
        let listening_socket = socket.clone();
        spawn(proc() handle_incoming(listening_socket));

        Ok(KRpcService {
            this_node: this_node,
            node_table: sync::Arc::new(sync::RWLock::new(node_table)),
            socket: socket,
        })
    }

    /// Get lock guarding access to a node table.
    pub fn node_table_lock(&self) -> &sync::RWLock<TNodeTable> {
        self.node_table.deref()
    }

    /// Get reference to a socket wrapper.
    ///
    /// Clone it if you want to get mutable copy.
    pub fn socket_ref(&self) -> &TSocket {
        &self.socket
    }
}

#[cfg(test)]
mod test {
    use std::io::IoResult;
    use std::io::net::ip;

    use num;

    use super::super::super::base::{mod, GenericNodeTable};
    use super::super::super::utils::test;

    use super::super::protocol;
    use super::super::udpwrapper::GenericSocketWrapper;

    use super::KRpcService;

    struct DummyNodeTable {
        last_node: Option<base::Node>,
    }

    impl GenericNodeTable for DummyNodeTable {
        fn random_id(&self) -> num::BigUint {
            // This number is random, I promise :)
            test::uint_to_id(42)
        }

        fn update(&mut self, node: &base::Node) -> bool {
            self.last_node = Some(node.clone());
            true
        }
        #[allow(unused_variable)]
        fn find(&self, id: &num::BigUint, count: uint) -> Vec<base::Node> {
            match self.last_node {
                Some(ref n) if n.id == *id => vec![n.clone()],
                _ => vec![]
            }
        }
        fn pop_oldest(&mut self) -> Vec<base::Node> {
            vec![]
        }
    }

    #[deriving(Clone)]
    struct DummySocketWrapper {
        last_package: Option<protocol::Package>,
        last_node: Option<base::Node>
    }

    impl GenericSocketWrapper for DummySocketWrapper {
        fn send(&mut self, package: &protocol::Package, node: &base::Node)
            -> IoResult<()> {
            self.last_package = Some(package.clone());
            self.last_node = Some(node.clone());
            Ok(())
        }

        fn receive(&mut self) -> IoResult<(protocol::Package, ip::SocketAddr)> {
            Ok((self.last_package.as_ref().unwrap().clone(),
                self.last_node.as_ref().unwrap().address.clone()))
        }
    }

    #[test]
    fn test_new() {
        let n = test::new_node_with_port(42, 8007);
        let sock = DummySocketWrapper { last_package: None, last_node: None };
        let s = KRpcService::new(n.clone(),
                                 DummyNodeTable { last_node: None },
                                 sock).unwrap();
        assert_eq!(n.id, s.this_node.id);
        let mut nt = s.node_table_lock().write();
        nt.update(&test::new_node(1));
        let nt2 = nt.downgrade();
        assert_eq!(nt2.random_id().to_u8().unwrap(), 42);
    }

    #[test]
    fn test_new_default() {
        let n = test::new_node_with_port(42, 8009);
        let s = KRpcService::new_default(n.clone()).unwrap();
        assert_eq!(n.id, s.this_node.id);
        let nt = s.node_table_lock().read();
        assert_eq!(0, nt.find(&test::uint_to_id(1), 1).len());
    }
}
<commit_msg>Fixed krpc docstring<commit_after>// Copyright 2014 Dmitry "Divius" Tantsur <divius.inside@gmail.com>
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.
//

//! KRPC DHT service as described in
//! [BEP 0005](http://www.bittorrent.org/beps/bep_0005.html).
//!
//! Create as:
//! ```
//! let service = dht::bt::KRpcService::new_default(this_node);
//! ```

use std::io::IoResult;
use std::sync;

use super::super::base;
use super::super::knodetable;

use super::udpwrapper;

/// Implementation of basic KRPC DHT on which BitTorrent DHT is based.
///
/// No peer retrival is supported: just finding and pinging nodes.
pub struct KRpcService<TNodeTable: base::GenericNodeTable,
                       TSocket: udpwrapper::GenericSocketWrapper> {
    this_node: base::Node,
    node_table: sync::Arc<sync::RWLock<TNodeTable>>,
    socket: TSocket,
}

/// Default kind of KRpc service.
pub type DefaultKRpcService = KRpcService<knodetable::KNodeTable, udpwrapper::UdpSocketWrapper>;

fn handle_incoming<TSocket: udpwrapper::GenericSocketWrapper>(socket: TSocket) {
    // TODO(divius): implement
    drop(socket);
}

impl KRpcService<knodetable::KNodeTable, udpwrapper::UdpSocketWrapper> {
    /// New service with default node table.
    pub fn new_default(this_node: base::Node) -> IoResult<DefaultKRpcService> {
        let node_table = knodetable::KNodeTable::new(this_node.id.clone());
        let socket = try!(udpwrapper::UdpSocketWrapper::new(&this_node));
        KRpcService::new(this_node, node_table, socket)
    }
}

impl<TNodeTable: base::GenericNodeTable, TSocket: udpwrapper::GenericSocketWrapper>
KRpcService<TNodeTable, TSocket> {
    /// New service with given node table and socket.
    pub fn new(this_node: base::Node, node_table: TNodeTable, socket: TSocket)
            -> IoResult<KRpcService<TNodeTable, TSocket>> {
        let listening_socket = socket.clone();
        spawn(proc() handle_incoming(listening_socket));

        Ok(KRpcService {
            this_node: this_node,
            node_table: sync::Arc::new(sync::RWLock::new(node_table)),
            socket: socket,
        })
    }

    /// Get lock guarding access to a node table.
    pub fn node_table_lock(&self) -> &sync::RWLock<TNodeTable> {
        self.node_table.deref()
    }

    /// Get reference to a socket wrapper.
    ///
    /// Clone it if you want to get mutable copy.
    pub fn socket_ref(&self) -> &TSocket {
        &self.socket
    }
}

#[cfg(test)]
mod test {
    use std::io::IoResult;
    use std::io::net::ip;

    use num;

    use super::super::super::base::{mod, GenericNodeTable};
    use super::super::super::utils::test;

    use super::super::protocol;
    use super::super::udpwrapper::GenericSocketWrapper;

    use super::KRpcService;

    struct DummyNodeTable {
        last_node: Option<base::Node>,
    }

    impl GenericNodeTable for DummyNodeTable {
        fn random_id(&self) -> num::BigUint {
            // This number is random, I promise :)
            test::uint_to_id(42)
        }

        fn update(&mut self, node: &base::Node) -> bool {
            self.last_node = Some(node.clone());
            true
        }
        #[allow(unused_variable)]
        fn find(&self, id: &num::BigUint, count: uint) -> Vec<base::Node> {
            match self.last_node {
                Some(ref n) if n.id == *id => vec![n.clone()],
                _ => vec![]
            }
        }
        fn pop_oldest(&mut self) -> Vec<base::Node> {
            vec![]
        }
    }

    #[deriving(Clone)]
    struct DummySocketWrapper {
        last_package: Option<protocol::Package>,
        last_node: Option<base::Node>
    }

    impl GenericSocketWrapper for DummySocketWrapper {
        fn send(&mut self, package: &protocol::Package, node: &base::Node)
            -> IoResult<()> {
            self.last_package = Some(package.clone());
            self.last_node = Some(node.clone());
            Ok(())
        }

        fn receive(&mut self) -> IoResult<(protocol::Package, ip::SocketAddr)> {
            Ok((self.last_package.as_ref().unwrap().clone(),
                self.last_node.as_ref().unwrap().address.clone()))
        }
    }

    #[test]
    fn test_new() {
        let n = test::new_node_with_port(42, 8007);
        let sock = DummySocketWrapper { last_package: None, last_node: None };
        let s = KRpcService::new(n.clone(),
                                 DummyNodeTable { last_node: None },
                                 sock).unwrap();
        assert_eq!(n.id, s.this_node.id);
        let mut nt = s.node_table_lock().write();
        nt.update(&test::new_node(1));
        let nt2 = nt.downgrade();
        assert_eq!(nt2.random_id().to_u8().unwrap(), 42);
    }

    #[test]
    fn test_new_default() {
        let n = test::new_node_with_port(42, 8009);
        let s = KRpcService::new_default(n.clone()).unwrap();
        assert_eq!(n.id, s.this_node.id);
        let nt = s.node_table_lock().read();
        assert_eq!(0, nt.find(&test::uint_to_id(1), 1).len());
    }
}
<|endoftext|>


<commit_before>use std::borrow::Cow;

use uuid::Uuid;
use package::Package;
use {UsernamePassword, Message};
use messages::mod_EventStore::mod_Client::mod_Messages::{WriteEvents, NewEvent};

pub struct Builder;

impl Builder {
    pub fn ping() -> SimpleBuilder {
        SimpleBuilder(Message::Ping)
    }

    pub fn authenticate() -> SimpleBuilder {
        SimpleBuilder(Message::Authenticate)
    }

    pub fn write_events() -> WriteEventsBuilder {
        WriteEventsBuilder::new()
    }
}

pub struct SimpleBuilder(Message);

impl SimpleBuilder {
    pub fn build_package(self, authentication: Option<UsernamePassword>, correlation_id: Option<Uuid>) -> Package {
        Package {
            authentication: authentication,
            correlation_id: correlation_id.unwrap_or_else(|| Uuid::new_v4()),
            message: self.0
        }
    }
}

pub struct WriteEventsBuilder {
    event_stream_id: Option<Cow<'static, str>>,
    expected_version: Option<ExpectedVersion>,
    require_master: Option<bool>,
    events: Vec<NewEvent<'static>>
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum ExpectedVersion {
    /// No optimistic locking
    Any,
    /// Expect a stream not to exist
    NewStream,
    /// Expect exact number of events in the stream
    Exact(StreamVersion)
}

impl Copy for ExpectedVersion {}

impl Into<i32> for ExpectedVersion {
    fn into(self) -> i32 {
        use self::ExpectedVersion::*;
        match self {
            Any => -2,
            NewStream => -1,
            Exact(StreamVersion(x)) => x as i32
        }
    }
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct StreamVersion(u32);

impl Copy for StreamVersion {}

impl StreamVersion {
    pub fn from_opt(version: u32) -> Option<StreamVersion> {
        // TODO: MAX_VALUE might be some magical value, should be lower?
        if version < i32::max_value() as u32 {
            Some(StreamVersion(version))
        } else {
            None
        }
    }
}

impl WriteEventsBuilder {

    pub fn new() -> Self {
        WriteEventsBuilder {
            event_stream_id: None,
            expected_version: None,
            require_master: None,
            events: Vec::new()
        }
    }

    /// Panics if the id is an empty string
    pub fn stream_id<S: Into<Cow<'static, str>>>(&mut self, id: S) -> &mut Self {
        let id = id.into();
        assert!(id.len() > 0);
        self.event_stream_id = Some(id);
        self
    }

    pub fn expected_version(&mut self, ver: ExpectedVersion) -> &mut Self {
        self.expected_version = Some(ver);
        self
    }

    /// Should the server only handle the request if it is the cluster master. Note that while only
    /// the master server can write, other cluster members can forward request to the master.
    ///
    /// Defaults to `false`.
    pub fn require_master(&mut self, require: bool) -> &mut Self {
        self.require_master = Some(require);
        self
    }

    pub fn new_event<'b>(&'b mut self) -> NewEventBuilder<'b> {
        NewEventBuilder::new(self)
    }

    fn push_event(&mut self, event: NewEvent<'static>) -> &mut Self {
        self.events.push(event);
        self
    }

    pub fn build_command(&mut self) -> WriteEvents<'static> {
        use std::mem;

        WriteEvents {
            event_stream_id: self.event_stream_id.take().unwrap(),
            expected_version: self.expected_version.take().unwrap_or(ExpectedVersion::Any).into(),
            require_master: self.require_master.take().unwrap_or(false),
            events: mem::replace(&mut self.events, Vec::new())
        }
    }

    pub fn build_message(&mut self) -> Message {
        self.build_command().into()
    }

    pub fn build_package(&mut self, authentication: Option<UsernamePassword>, correlation_id: Option<Uuid>) -> Package {
        Package {
            authentication: authentication,
            correlation_id: correlation_id.unwrap_or_else(|| Uuid::new_v4()),
            message: self.build_command().into()
        }
    }
}

pub struct NewEventBuilder<'a> {
    parent: &'a mut WriteEventsBuilder,

    event_id: Option<Uuid>,
    event_type: Option<Cow<'static, str>>,
    data_content_type: Option<i32>,
    metadata_content_type: Option<i32>,
    data: Option<Cow<'static, [u8]>>,
    metadata: Option<Cow<'static, [u8]>>,
}

impl<'a> NewEventBuilder<'a> {
    fn new(parent: &'a mut WriteEventsBuilder) -> NewEventBuilder<'a> {
        NewEventBuilder {
            parent: parent,
            event_id: None,
            event_type: None,
            data_content_type: None,
            data: None,
            metadata_content_type: None,
            metadata: None
        }
    }

    pub fn event_id(mut self, event_id: Uuid) -> Self {
        self.event_id = Some(event_id);
        self
    }

    pub fn event_type<S: Into<Cow<'static, str>>>(mut self, ty: S) -> Self {
        // TODO: is there some invariant?
        self.event_type = Some(ty.into());
        self
    }

    pub fn data<D: Into<Vec<u8>>>(mut self, data: D) -> Self {
        self.data = Some(Cow::Owned(data.into()));
        self
    }

    pub fn data_content_type(mut self, is_json: bool) -> Self {
        self.data_content_type = if is_json { Some(1) } else { Some(0) };
        self
    }

    pub fn metadata<D: Into<Vec<u8>>>(mut self, metadata: D) -> Self {
        self.metadata = Some(Cow::Owned(metadata.into()));
        self
    }

    pub fn metadata_content_type(mut self, is_json: bool) -> Self {
        self.metadata_content_type = if is_json { Some(1) } else { Some(0) };
        self
    }

    pub fn done(self) -> &'a mut WriteEventsBuilder {
        let event = NewEvent {
            event_id: Cow::Owned(self.event_id.unwrap_or_else(|| Uuid::new_v4()).as_bytes().into_iter().cloned().collect::<Vec<_>>()),
            event_type: self.event_type.unwrap(),
            data_content_type: self.data_content_type.unwrap_or(0),
            metadata_content_type: self.metadata_content_type.unwrap_or(0),
            data: self.data.unwrap(),
            metadata: self.metadata,
        };

        self.parent.push_event(event)
    }

    pub fn cancel(self) -> &'a mut WriteEventsBuilder {
        self.parent
    }
}

#[test]
fn build_new_event_for_write_events() {
    let _ = Builder::write_events()
        .stream_id("foobar")
        .new_event()
            .event_type("foo")
            .data(vec![0u8])
            .data_content_type(false)
        .done()
        .build_package(None, None);
}
<commit_msg>builder: add ReadEvent builder<commit_after>use std::borrow::Cow;

use uuid::Uuid;
use package::Package;
use {UsernamePassword, Message};
use messages::mod_EventStore::mod_Client::mod_Messages::{WriteEvents, NewEvent, ReadEvent};

pub struct Builder;

impl Builder {
    pub fn ping() -> SimpleBuilder {
        SimpleBuilder(Message::Ping)
    }

    pub fn authenticate() -> SimpleBuilder {
        SimpleBuilder(Message::Authenticate)
    }

    pub fn write_events() -> WriteEventsBuilder {
        WriteEventsBuilder::new()
    }

    pub fn read_event() -> ReadEventBuilder {
        ReadEventBuilder::new()
    }
}

pub struct SimpleBuilder(Message);

impl SimpleBuilder {
    pub fn build_package(self, authentication: Option<UsernamePassword>, correlation_id: Option<Uuid>) -> Package {
        Package {
            authentication: authentication,
            correlation_id: correlation_id.unwrap_or_else(|| Uuid::new_v4()),
            message: self.0
        }
    }
}

pub struct WriteEventsBuilder {
    event_stream_id: Option<Cow<'static, str>>,
    expected_version: Option<ExpectedVersion>,
    require_master: Option<bool>,
    events: Vec<NewEvent<'static>>
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum ExpectedVersion {
    /// No optimistic locking
    Any,
    /// Expect a stream not to exist
    NewStream,
    /// Expect exact number of events in the stream
    Exact(StreamVersion)
}

impl Copy for ExpectedVersion {}

impl Into<i32> for ExpectedVersion {
    fn into(self) -> i32 {
        use self::ExpectedVersion::*;
        match self {
            Any => -2,
            NewStream => -1,
            Exact(ver) => ver.into()
        }
    }
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct StreamVersion(u32);

impl Copy for StreamVersion {}

impl StreamVersion {
    pub fn from_opt(version: u32) -> Option<StreamVersion> {
        // TODO: MAX_VALUE might be some magical value, should be lower?
        if version < i32::max_value() as u32 {
            Some(StreamVersion(version))
        } else {
            None
        }
    }
}

impl Into<i32> for StreamVersion {
    fn into(self) -> i32 {
        self.0 as i32
    }
}

impl WriteEventsBuilder {

    pub fn new() -> Self {
        WriteEventsBuilder {
            event_stream_id: None,
            expected_version: None,
            require_master: None,
            events: Vec::new()
        }
    }

    /// Panics if the id is an empty string
    pub fn stream_id<S: Into<Cow<'static, str>>>(&mut self, id: S) -> &mut Self {
        let id = id.into();
        assert!(id.len() > 0);
        self.event_stream_id = Some(id);
        self
    }

    pub fn expected_version(&mut self, ver: ExpectedVersion) -> &mut Self {
        self.expected_version = Some(ver);
        self
    }

    /// Should the server only handle the request if it is the cluster master. Note that while only
    /// the master server can write, other cluster members can forward request to the master.
    ///
    /// Defaults to `false`.
    pub fn require_master(&mut self, require: bool) -> &mut Self {
        self.require_master = Some(require);
        self
    }

    pub fn new_event<'b>(&'b mut self) -> NewEventBuilder<'b> {
        NewEventBuilder::new(self)
    }

    fn push_event(&mut self, event: NewEvent<'static>) -> &mut Self {
        self.events.push(event);
        self
    }

    pub fn build_command(&mut self) -> WriteEvents<'static> {
        use std::mem;

        WriteEvents {
            event_stream_id: self.event_stream_id.take().unwrap(),
            expected_version: self.expected_version.take().unwrap_or(ExpectedVersion::Any).into(),
            require_master: self.require_master.take().unwrap_or(false),
            events: mem::replace(&mut self.events, Vec::new())
        }
    }

    pub fn build_message(&mut self) -> Message {
        self.build_command().into()
    }

    pub fn build_package(&mut self, authentication: Option<UsernamePassword>, correlation_id: Option<Uuid>) -> Package {
        Package {
            authentication: authentication,
            correlation_id: correlation_id.unwrap_or_else(|| Uuid::new_v4()),
            message: self.build_command().into()
        }
    }
}

pub struct NewEventBuilder<'a> {
    parent: &'a mut WriteEventsBuilder,

    event_id: Option<Uuid>,
    event_type: Option<Cow<'static, str>>,
    data_content_type: Option<i32>,
    metadata_content_type: Option<i32>,
    data: Option<Cow<'static, [u8]>>,
    metadata: Option<Cow<'static, [u8]>>,
}

impl<'a> NewEventBuilder<'a> {
    fn new(parent: &'a mut WriteEventsBuilder) -> NewEventBuilder<'a> {
        NewEventBuilder {
            parent: parent,
            event_id: None,
            event_type: None,
            data_content_type: None,
            data: None,
            metadata_content_type: None,
            metadata: None
        }
    }

    pub fn event_id(mut self, event_id: Uuid) -> Self {
        self.event_id = Some(event_id);
        self
    }

    pub fn event_type<S: Into<Cow<'static, str>>>(mut self, ty: S) -> Self {
        // TODO: is there some invariant?
        self.event_type = Some(ty.into());
        self
    }

    pub fn data<D: Into<Vec<u8>>>(mut self, data: D) -> Self {
        self.data = Some(Cow::Owned(data.into()));
        self
    }

    pub fn data_content_type(mut self, is_json: bool) -> Self {
        self.data_content_type = if is_json { Some(1) } else { Some(0) };
        self
    }

    pub fn metadata<D: Into<Vec<u8>>>(mut self, metadata: D) -> Self {
        self.metadata = Some(Cow::Owned(metadata.into()));
        self
    }

    pub fn metadata_content_type(mut self, is_json: bool) -> Self {
        self.metadata_content_type = if is_json { Some(1) } else { Some(0) };
        self
    }

    pub fn done(self) -> &'a mut WriteEventsBuilder {
        let event = NewEvent {
            event_id: Cow::Owned(self.event_id.unwrap_or_else(|| Uuid::new_v4()).as_bytes().into_iter().cloned().collect::<Vec<_>>()),
            event_type: self.event_type.unwrap(),
            data_content_type: self.data_content_type.unwrap_or(0),
            metadata_content_type: self.metadata_content_type.unwrap_or(0),
            data: self.data.unwrap(),
            metadata: self.metadata,
        };

        self.parent.push_event(event)
    }

    pub fn cancel(self) -> &'a mut WriteEventsBuilder {
        self.parent
    }
}

pub struct ReadEventBuilder {
    event_stream_id: Option<Cow<'static, str>>,
    event_number: Option<StreamVersion>,
    resolve_link_tos: Option<bool>,
    require_master: Option<bool>,
}

impl ReadEventBuilder {
    pub fn new() -> Self {
        ReadEventBuilder {
            event_stream_id: None,
            event_number: None,
            resolve_link_tos: None,
            require_master: None,
        }
    }

    /// Panics if the id is an empty string
    pub fn stream_id<S: Into<Cow<'static, str>>>(&mut self, id: S) -> &mut Self {
        let id = id.into();
        assert!(id.len() > 0);
        self.event_stream_id = Some(id);
        self
    }

    pub fn event_number(&mut self, ver: StreamVersion) -> &mut Self {
        self.event_number = Some(ver);
        self
    }

    pub fn resolve_link_tos(&mut self, resolve: bool) -> &mut Self {
        self.resolve_link_tos = Some(resolve);
        self
    }

    /// Should the server only handle the request if it is the cluster master. Note that while only
    /// the master server can write, other cluster members can forward request to the master.
    ///
    /// Defaults to `false`.
    pub fn require_master(&mut self, require: bool) -> &mut Self {
        self.require_master = Some(require);
        self
    }

    pub fn build_command(&mut self) -> ReadEvent<'static> {
        ReadEvent {
            event_stream_id: self.event_stream_id.take().unwrap(),
            event_number: self.event_number.take().unwrap().into(),
            resolve_link_tos: self.resolve_link_tos.take().unwrap_or(true),
            require_master: self.require_master.take().unwrap_or(false)
        }
    }

    pub fn build_message(&mut self) -> Message {
        self.build_command().into()
    }

    pub fn build_package(&mut self, authentication: Option<UsernamePassword>, correlation_id: Option<Uuid>) -> Package {
        build_package(self.build_message(), authentication, correlation_id)
    }
}

fn build_package<M: Into<Message>>(msg: M, authentication: Option<UsernamePassword>, correlation_id: Option<Uuid>) -> Package {
    Package {
        authentication: authentication,
        correlation_id: correlation_id.unwrap_or_else(|| Uuid::new_v4()),
        message: msg.into()
    }
}

#[test]
fn build_new_event_for_write_events() {
    let _ = Builder::write_events()
        .stream_id("foobar")
        .new_event()
            .event_type("foo")
            .data(vec![0u8])
            .data_content_type(false)
        .done()
        .build_package(None, None);
}
<|endoftext|>


<commit_before>use std::{thread, time};

/// Chipate Module
/// Rust emulation of the Chip-8
/// http://www.multigesturearticles/how-to-write-an-emulator-chip-8-interpreter/
#[allow(dead_code)]
pub struct Chipate {
    // Opcodes
    opcode: u16,

    // 4K Memory
    memory: [u8; 4096],
    // The systems memory map:
    // 0x000-0x1FF - Chip 8 interpreter (contains font set in emu)
    // 0x050-0x0A0 - Used for the built in 4x5 pixel font set (0-F)
    // 0x200-0xFFF - Program ROM and work RAM

    // Registers
    v: [u8; 16],

    // Index
    i: u16,

    // Program Counter
    pc: u16,

    gfx: [u8; 64 * 32],

    // Timers
    delay_timer: u8,
    sound_timer: u8,

    // Stack
    stack: [u8; 16],
    sp: u8,

    // Keypad
    key: [u8; 16],

    program: &'static str,
}

impl Chipate {
    pub fn init(&mut self) {
        debug!("Initialize Chip");
        // pc     = 0x200;  // Program counter starts at 0x200
        self.pc = 0x200;
        // opcode = 0;      // Reset current opcode
        self.opcode = 0x0;
        // I      = 0;      // Reset index register
        self.i = 0x0;
        // sp     = 0;      // Reset stack pointer
        self.sp = 0x0;

        // // Clear display
        // // Clear stack
        // // Clear registers V0-VF
        // // Clear memory

        // // Load fontset
        // for(int i = 0; i < 80; ++i)
        //     memory[i] = chip8_fontset[i];

        // // Reset timers
    }

    pub fn load_program(&mut self, program: &str) {
        debug!("Loading program {}", program);
    }

    pub fn emulate_cycle(&mut self) {
        debug!("Cycle Begin");

        self.fetch_opcode();

        let one_second = time::Duration::from_secs(1);
        thread::sleep(one_second);

        debug!("Cycle End");
    }

    pub fn draw_screen(&mut self) {
        debug!("Drawing to Screen")
    }

    pub fn set_keys(&mut self) {
        debug!("Saving Key State")
    }

    pub fn setup_testing_memory(&mut self) {
        debug!("Setting test memory");

        self.memory[self.pc as usize] = 0xA2;
        self.memory[(self.pc + 1) as usize] = 0xF0;

        debug!("Location: 0x{:X} data: 0x{:X}",
               self.pc,
               self.memory[self.pc as usize]);
        debug!("Location: 0x{:X} data: 0x{:X}",
               self.pc + 1,
               self.memory[(self.pc + 1) as usize]);
    }

    pub fn fetch_opcode(&mut self) {
        let op_a = self.memory[self.pc as usize];
        self.opcode = op_a as u16;
        self.opcode = self.opcode << 8;
        // debug!("location: 0x{:X} data: 0x{:X}", self.pc, self.opcode);

        let op_b = self.memory[(self.pc + 1) as usize] as u16;
        // debug!("location: 0x{:X} data: 0x{:X}", self.pc, op_b);

        self.opcode = self.opcode | op_b;
        debug!("Opcode: 0x{:X}", self.opcode);
    }

    pub fn decode_opcode(&mut self) {
        let op = self.opcode & 0xF000;
        debug!("Decode: 0x{:X}", op);

        match op {
            0xa000 => self.annn_opcode(),
            _ => {
                debug!("Catch all: 0x{:X}", self.opcode);
            }
        }
    }

    pub fn increase_pc(&mut self) {
        self.pc += 2;
        debug!("Program Counter: 0x{:X}", self.pc);
    }

    pub fn annn_opcode(&mut self) {
        debug!("ANNN: 0x{:X}", self.opcode);
        self.i = self.opcode & 0x0FFF;
        self.increase_pc();
    }

    pub fn new() -> Chipate {
        debug!("Creating New Chip");

        let chip = Chipate {
            opcode: 0,
            memory: [0; 4096],
            v: [0; 16],
            i: 0,
            pc: 0,
            gfx: [0; 64 * 32],
            delay_timer: 0,
            sound_timer: 0,
            stack: [0; 16],
            sp: 0,
            key: [0; 16],
            program: "",
        };

        return chip;
    }
}
<commit_msg>Consistent capitalization for hex values<commit_after>use std::{thread, time};

/// Chipate Module
/// Rust emulation of the Chip-8
/// http://www.multigesturearticles/how-to-write-an-emulator-chip-8-interpreter/
#[allow(dead_code)]
pub struct Chipate {
    // Opcodes
    opcode: u16,

    // 4K Memory
    memory: [u8; 4096],
    // The systems memory map:
    // 0x000-0x1FF - Chip 8 interpreter (contains font set in emu)
    // 0x050-0x0A0 - Used for the built in 4x5 pixel font set (0-F)
    // 0x200-0xFFF - Program ROM and work RAM

    // Registers
    v: [u8; 16],

    // Index
    i: u16,

    // Program Counter
    pc: u16,

    gfx: [u8; 64 * 32],

    // Timers
    delay_timer: u8,
    sound_timer: u8,

    // Stack
    stack: [u8; 16],
    sp: u8,

    // Keypad
    key: [u8; 16],

    program: &'static str,
}

impl Chipate {
    pub fn init(&mut self) {
        debug!("Initialize Chip");
        // pc     = 0x200;  // Program counter starts at 0x200
        self.pc = 0x200;
        // opcode = 0;      // Reset current opcode
        self.opcode = 0x0;
        // I      = 0;      // Reset index register
        self.i = 0x0;
        // sp     = 0;      // Reset stack pointer
        self.sp = 0x0;

        // // Clear display
        // // Clear stack
        // // Clear registers V0-VF
        // // Clear memory

        // // Load fontset
        // for(int i = 0; i < 80; ++i)
        //     memory[i] = chip8_fontset[i];

        // // Reset timers
    }

    pub fn load_program(&mut self, program: &str) {
        debug!("Loading program {}", program);
    }

    pub fn emulate_cycle(&mut self) {
        debug!("Cycle Begin");

        self.fetch_opcode();

        let one_second = time::Duration::from_secs(1);
        thread::sleep(one_second);

        debug!("Cycle End");
    }

    pub fn draw_screen(&mut self) {
        debug!("Drawing to Screen")
    }

    pub fn set_keys(&mut self) {
        debug!("Saving Key State")
    }

    pub fn setup_testing_memory(&mut self) {
        debug!("Setting test memory");

        self.memory[self.pc as usize] = 0xA2;
        self.memory[(self.pc + 1) as usize] = 0xF0;

        debug!("Location: 0x{:X} data: 0x{:X}",
               self.pc,
               self.memory[self.pc as usize]);
        debug!("Location: 0x{:X} data: 0x{:X}",
               self.pc + 1,
               self.memory[(self.pc + 1) as usize]);
    }

    pub fn fetch_opcode(&mut self) {
        let op_a = self.memory[self.pc as usize];
        self.opcode = op_a as u16;
        self.opcode = self.opcode << 8;
        // debug!("location: 0x{:X} data: 0x{:X}", self.pc, self.opcode);

        let op_b = self.memory[(self.pc + 1) as usize] as u16;
        // debug!("location: 0x{:X} data: 0x{:X}", self.pc, op_b);

        self.opcode = self.opcode | op_b;
        debug!("Opcode: 0x{:X}", self.opcode);
    }

    pub fn decode_opcode(&mut self) {
        let op = self.opcode & 0xF000;
        debug!("Decode: 0x{:X}", op);

        match op {
            0xA000 => self.annn_opcode(),
            _ => {
                debug!("Catch all: 0x{:X}", self.opcode);
            }
        }
    }

    pub fn increase_pc(&mut self) {
        self.pc += 2;
        debug!("Program Counter: 0x{:X}", self.pc);
    }

    pub fn annn_opcode(&mut self) {
        debug!("ANNN: 0x{:X}", self.opcode);
        self.i = self.opcode & 0x0FFF;
        self.increase_pc();
    }

    pub fn new() -> Chipate {
        debug!("Creating New Chip");

        let chip = Chipate {
            opcode: 0,
            memory: [0; 4096],
            v: [0; 16],
            i: 0,
            pc: 0,
            gfx: [0; 64 * 32],
            delay_timer: 0,
            sound_timer: 0,
            stack: [0; 16],
            sp: 0,
            key: [0; 16],
            program: "",
        };

        return chip;
    }
}
<|endoftext|>


<commit_before>#![recursion_limit = "1000"]

#[macro_use]
extern crate error_chain;
extern crate lalrpop_util;
extern crate reproto_core;
extern crate num;
extern crate linked_hash_map;

pub mod ast;
pub mod errors;
pub mod lexer;
#[allow(unused)]
mod parser;
mod token;
mod utils;

use lalrpop_util::ParseError;
use self::errors::*;
use std::io::Read;
use reproto_core::object::Object;
use std::sync::{Arc, Mutex};

pub fn read_reader<'a, R>(mut reader: R) -> Result<String>
    where R: AsMut<Read + 'a>
{
    let mut content = String::new();
    reader.as_mut().read_to_string(&mut content)?;
    Ok(content)
}

pub fn parse_string<'input>(object: Arc<Mutex<Box<Object>>>, input: &'input str) -> Result<ast::File<'input>> {
    use self::ErrorKind::*;

    let lexer = lexer::lex(input);

    match parser::parse_File(&object, lexer) {
        Ok(file) => Ok(file),
        Err(e) => {
            match e {
                ParseError::InvalidToken { location } => {
                    let pos = (object.clone(), location, location);
                    Err(Syntax(Some(pos.into()), vec![]).into())
                }
                ParseError::UnrecognizedToken { token, expected } => {
                    let pos = token.map(|(start, _, end)| (object.clone(), start, end));
                    Err(Syntax(pos.map(Into::into), expected).into())
                }
                ParseError::User { error } => {
                    match error {
                        token::Error::UnterminatedString { start } => {
                            let pos = (object.clone(), start, start);
                            return Err(Parse("unterminated string", pos.into()).into());
                        }
                        token::Error::UnterminatedEscape { start } => {
                            let pos = (object.clone(), start, start);
                            return Err(Parse("unterminated escape sequence", pos.into()).into());
                        }
                        token::Error::InvalidEscape { pos, message } => {
                            let pos = (object.clone(), pos, pos);
                            return Err(Parse(message, pos.into()).into());
                        }
                        token::Error::UnterminatedCodeBlock { start } => {
                            let pos = (object.clone(), start, start);
                            return Err(Parse("unterminated code block", pos.into()).into());
                        }
                        token::Error::InvalidNumber { pos, message } => {
                            let pos = (object.clone(), pos, pos);
                            return Err(Parse(message, pos.into()).into());
                        }
                        token::Error::Unexpected { pos } => {
                            let pos = (object.clone(), pos, pos);
                            return Err(Parse("unexpected input", pos.into()).into());
                        }
                        token::Error::InvalidVersion { start, end } => {
                            let pos = (object.clone(), start, end);
                            return Err(Parse("invalid version", pos.into()).into());
                        }
                        token::Error::InvalidVersionReq { start, end } => {
                            let pos = (object.clone(), start, end);
                            return Err(Parse("invalid version requirement", pos.into()).into());
                        }
                    }
                }
                _ => Err("parse error".into()),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use std::path::PathBuf;
    use super::*;
    use super::ast::*;

    /// Check that a parsed value equals expected.
    macro_rules! assert_value_eq {
        ($expected:expr, $input:expr) => {{
            let path = Rc::new(PathBuf::from(""));
            let v = parser::parse_Value(&path, parse($input)).unwrap();
            assert_eq!($expected, v);
        }}
    }

    macro_rules! assert_type_spec_eq {
        ($expected:expr, $input:expr) => {{
            let path = Rc::new(PathBuf::from(""));
            let v = parser::parse_TypeSpec(&path, parse($input)).unwrap();
            assert_eq!($expected, v);
        }}
    }

    const FILE1: &[u8] = include_bytes!("tests/file1.reproto");
    const INTERFACE1: &[u8] = include_bytes!("tests/interface1.reproto");

    fn parse(input: &'static str)
             -> Box<Iterator<Item = token::Result<(usize, token::Token<'static>, usize)>>> {
        Box::new(lexer::lex(input))
    }

    fn parse_file(input: &'static str) -> File {
        let path = Rc::new(PathBuf::from(""));
        parser::parse_File(&path, parse(input)).unwrap()
    }

    fn parse_member(input: &'static str) -> Member {
        let path = Rc::new(PathBuf::from(""));
        parser::parse_Member(&path, parse(input)).unwrap()
    }

    fn parse_type_spec(input: &'static str) -> RpType {
        let path = Rc::new(PathBuf::from(""));
        parser::parse_TypeSpec(&path, parse(input)).unwrap()
    }

    #[test]
    fn test_file1() {
        let input = ::std::str::from_utf8(FILE1).unwrap();
        let _ = parse_file(input);
    }

    #[test]
    fn test_array() {
        let ty = parse_type_spec("[string]");

        if let RpType::Array { inner } = ty {
            if let RpType::String = *inner {
                return;
            }
        }

        panic!("Expected Type::Array(Type::String)");
    }

    #[test]
    fn test_map() {
        let ty = parse_type_spec("{string: unsigned/123}");

        // TODO: use #![feature(box_patterns)]:
        // if let Type::Map(box Type::String, box Type::Unsigned(size)) = ty {
        // }
        if let RpType::Map { key, value } = ty {
            if let RpType::String = *key {
                if let RpType::Unsigned { size } = *value {
                    assert_eq!(Some(123usize), size);
                    return;
                }
            }
        }

        panic!("Expected Type::Array(Type::String)");
    }

    #[test]
    fn test_block_comment() {
        parse("/* hello \n world */");
    }

    #[test]
    fn test_line_comment() {
        parse("// hello world\n");
    }

    #[test]
    fn test_code() {
        parse_member("java{{\na { b { c } d } e\n}}");
    }

    #[test]
    fn test_interface() {
        let input = ::std::str::from_utf8(INTERFACE1).unwrap();
        let file = parse_file(input);
        assert_eq!(1, file.decls.len());
    }

    #[test]
    fn test_instance() {
        let path = Rc::new(PathBuf::from(""));

        let c = RpName {
            prefix: None,
            parts: vec!["Foo".to_owned(), "Bar".to_owned()],
        };

        let field = FieldInit {
            name: Loc::new("hello", (path.clone(), 8, 13)),
            value: Loc::new(Value::Number(12.into()), (path.clone(), 15, 17)),
        };

        let field = Loc::new(field, (path.clone(), 8, 17));

        let instance = Instance {
            name: c,
            arguments: Loc::new(vec![field], (path.clone(), 8, 17)),
        };

        let instance = Loc::new(instance, (path.clone(), 0, 18));
        let object = Loc::new(Object::Instance(instance), (path.clone(), 0, 18));

        assert_value_eq!(Value::Object(object), "Foo.Bar(hello: 12)");
    }

    #[test]
    fn test_strings() {
        assert_value_eq!(Value::String("foo\nbar".to_owned()), "\"foo\\nbar\"");
    }

    #[test]
    fn test_numbers() {
        assert_value_eq!(Value::Number(1.into()), "1");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: 125.into(),
                             decimal: 2,
                         }),
                         "1.25");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: 12500.into(),
                             decimal: 0,
                         }),
                         "1.25e4");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: (-12500).into(),
                             decimal: 0,
                         }),
                         "-1.25e4");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: (1234).into(),
                             decimal: 8,
                         }),
                         "0.00001234");
    }

    #[test]
    fn test_type_spec() {
        let c = RpName {
            prefix: None,
            parts: vec!["Hello".to_owned(), "World".to_owned()],
        };

        assert_type_spec_eq!(RpType::String, "string");
        assert_type_spec_eq!(RpType::Name { name: c }, "Hello.World");
    }

    #[test]
    fn test_option_decl() {
        let member = parse_member("foo_bar_baz true, foo, \"bar\", 12;");

        if let Member::Option(option) = member {
            assert_eq!("foo_bar_baz", option.name);
            assert_eq!(4, option.values.len());

            assert_eq!(Value::Boolean(true), *option.values[0].as_ref());
            assert_eq!(Value::Identifier("foo"), *option.values[1].as_ref());
            assert_eq!(Value::String("bar".to_owned()), *option.values[2].as_ref());
            assert_eq!(Value::Number(12u32.into()), *option.values[3].as_ref());
            return;
        }

        panic!("option did not match");
    }
}
<commit_msg>parser: fix tests<commit_after>#![recursion_limit = "1000"]

#[macro_use]
extern crate error_chain;
extern crate lalrpop_util;
extern crate reproto_core;
extern crate num;
extern crate linked_hash_map;

pub mod ast;
pub mod errors;
pub mod lexer;
#[allow(unused)]
mod parser;
mod token;
mod utils;

use lalrpop_util::ParseError;
use self::errors::*;
use std::io::Read;
use reproto_core::object;
use std::sync::{Arc, Mutex};

pub fn read_reader<'a, R>(mut reader: R) -> Result<String>
    where R: AsMut<Read + 'a>
{
    let mut content = String::new();
    reader.as_mut().read_to_string(&mut content)?;
    Ok(content)
}

pub fn parse_string<'input>(object: Arc<Mutex<Box<object::Object>>>, input: &'input str) -> Result<ast::File<'input>> {
    use self::ErrorKind::*;

    let lexer = lexer::lex(input);

    match parser::parse_File(&object, lexer) {
        Ok(file) => Ok(file),
        Err(e) => {
            match e {
                ParseError::InvalidToken { location } => {
                    let pos = (object.clone(), location, location);
                    Err(Syntax(Some(pos.into()), vec![]).into())
                }
                ParseError::UnrecognizedToken { token, expected } => {
                    let pos = token.map(|(start, _, end)| (object.clone(), start, end));
                    Err(Syntax(pos.map(Into::into), expected).into())
                }
                ParseError::User { error } => {
                    match error {
                        token::Error::UnterminatedString { start } => {
                            let pos = (object.clone(), start, start);
                            return Err(Parse("unterminated string", pos.into()).into());
                        }
                        token::Error::UnterminatedEscape { start } => {
                            let pos = (object.clone(), start, start);
                            return Err(Parse("unterminated escape sequence", pos.into()).into());
                        }
                        token::Error::InvalidEscape { pos, message } => {
                            let pos = (object.clone(), pos, pos);
                            return Err(Parse(message, pos.into()).into());
                        }
                        token::Error::UnterminatedCodeBlock { start } => {
                            let pos = (object.clone(), start, start);
                            return Err(Parse("unterminated code block", pos.into()).into());
                        }
                        token::Error::InvalidNumber { pos, message } => {
                            let pos = (object.clone(), pos, pos);
                            return Err(Parse(message, pos.into()).into());
                        }
                        token::Error::Unexpected { pos } => {
                            let pos = (object.clone(), pos, pos);
                            return Err(Parse("unexpected input", pos.into()).into());
                        }
                        token::Error::InvalidVersion { start, end } => {
                            let pos = (object.clone(), start, end);
                            return Err(Parse("invalid version", pos.into()).into());
                        }
                        token::Error::InvalidVersionReq { start, end } => {
                            let pos = (object.clone(), start, end);
                            return Err(Parse("invalid version requirement", pos.into()).into());
                        }
                    }
                }
                _ => Err("parse error".into()),
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use super::ast::*;
    use std::sync::{Arc, Mutex};
    use reproto_core::object;

    fn new_context() -> Arc<Mutex<Box<object::Object>>> {
        Arc::new(Mutex::new(Box::new(object::BytesObject::new(String::from(""), vec![]))))
    }

    /// Check that a parsed value equals expected.
    macro_rules! assert_value_eq {
        ($expected:expr, $input:expr) => {{
            let v = parser::parse_Value(&new_context(), parse($input)).unwrap();
            assert_eq!($expected, v);
        }}
    }

    macro_rules! assert_type_spec_eq {
        ($expected:expr, $input:expr) => {{
            let v = parser::parse_TypeSpec(&new_context(), parse($input)).unwrap();
            assert_eq!($expected, v);
        }}
    }

    const FILE1: &[u8] = include_bytes!("tests/file1.reproto");
    const INTERFACE1: &[u8] = include_bytes!("tests/interface1.reproto");

    fn parse(input: &'static str)
             -> Box<Iterator<Item = token::Result<(usize, token::Token<'static>, usize)>>> {
        Box::new(lexer::lex(input))
    }

    fn parse_file(input: &'static str) -> File {
        parser::parse_File(&new_context(), parse(input)).unwrap()
    }

    fn parse_member(input: &'static str) -> Member {
        parser::parse_Member(&new_context(), parse(input)).unwrap()
    }

    fn parse_type_spec(input: &'static str) -> RpType {
        parser::parse_TypeSpec(&new_context(), parse(input)).unwrap()
    }

    #[test]
    fn test_file1() {
        let input = ::std::str::from_utf8(FILE1).unwrap();
        let _ = parse_file(input);
    }

    #[test]
    fn test_array() {
        let ty = parse_type_spec("[string]");

        if let RpType::Array { inner } = ty {
            if let RpType::String = *inner {
                return;
            }
        }

        panic!("Expected Type::Array(Type::String)");
    }

    #[test]
    fn test_map() {
        let ty = parse_type_spec("{string: unsigned/123}");

        // TODO: use #![feature(box_patterns)]:
        // if let Type::Map(box Type::String, box Type::Unsigned(size)) = ty {
        // }
        if let RpType::Map { key, value } = ty {
            if let RpType::String = *key {
                if let RpType::Unsigned { size } = *value {
                    assert_eq!(Some(123usize), size);
                    return;
                }
            }
        }

        panic!("Expected Type::Array(Type::String)");
    }

    #[test]
    fn test_block_comment() {
        parse("/* hello \n world */");
    }

    #[test]
    fn test_line_comment() {
        parse("// hello world\n");
    }

    #[test]
    fn test_code() {
        parse_member("java{{\na { b { c } d } e\n}}");
    }

    #[test]
    fn test_interface() {
        let input = ::std::str::from_utf8(INTERFACE1).unwrap();
        let file = parse_file(input);
        assert_eq!(1, file.decls.len());
    }

    #[test]
    fn test_instance() {
        let context = new_context();

        let c = RpName {
            prefix: None,
            parts: vec!["Foo".to_owned(), "Bar".to_owned()],
        };

        let field = FieldInit {
            name: Loc::new("hello", (context.clone(), 8, 13)),
            value: Loc::new(Value::Number(12.into()), (context.clone(), 15, 17)),
        };

        let field = Loc::new(field, (context.clone(), 8, 17));

        let instance = Instance {
            name: c,
            arguments: Loc::new(vec![field], (context.clone(), 8, 17)),
        };

        let instance = Loc::new(instance, (context.clone(), 0, 18));
        let object = Loc::new(Object::Instance(instance), (context.clone(), 0, 18));

        assert_value_eq!(Value::Object(object), "Foo.Bar(hello: 12)");
    }

    #[test]
    fn test_strings() {
        assert_value_eq!(Value::String("foo\nbar".to_owned()), "\"foo\\nbar\"");
    }

    #[test]
    fn test_numbers() {
        assert_value_eq!(Value::Number(1.into()), "1");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: 125.into(),
                             decimal: 2,
                         }),
                         "1.25");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: 12500.into(),
                             decimal: 0,
                         }),
                         "1.25e4");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: (-12500).into(),
                             decimal: 0,
                         }),
                         "-1.25e4");

        assert_value_eq!(Value::Number(RpNumber {
                             digits: (1234).into(),
                             decimal: 8,
                         }),
                         "0.00001234");
    }

    #[test]
    fn test_type_spec() {
        let c = RpName {
            prefix: None,
            parts: vec!["Hello".to_owned(), "World".to_owned()],
        };

        assert_type_spec_eq!(RpType::String, "string");
        assert_type_spec_eq!(RpType::Name { name: c }, "Hello.World");
    }

    #[test]
    fn test_option_decl() {
        let member = parse_member("foo_bar_baz true, foo, \"bar\", 12;");

        if let Member::Option(option) = member {
            assert_eq!("foo_bar_baz", option.name);
            assert_eq!(4, option.values.len());

            assert_eq!(Value::Boolean(true), *option.values[0].as_ref());
            assert_eq!(Value::Identifier("foo"), *option.values[1].as_ref());
            assert_eq!(Value::String("bar".to_owned()), *option.values[2].as_ref());
            assert_eq!(Value::Number(12u32.into()), *option.values[3].as_ref());
            return;
        }

        panic!("option did not match");
    }
}
<|endoftext|>


<commit_before>//! Asynchronous DNS resolver

use std::io::{self, ErrorKind};
use std::mem;
use std::net::SocketAddr;
use std::sync::Arc;

use futures::Future;
use spin::Mutex;
use tokio::runtime::current_thread::Runtime;
use trust_dns_resolver::config::ResolverConfig;
use trust_dns_resolver::AsyncResolver;

use config::Config;

// Taken from
// bluejekyll/trust-dns/resolver/examples/global_resolver.rs
lazy_static! {
    static ref GLOBAL_DNS_ADDRESS: Mutex<Option<ResolverConfig>> = Mutex::new(None);

    // First we need to setup the global Resolver
    static ref GLOBAL_DNS_RESOLVER: AsyncResolver = init_resolver();
}

/// Set address for global DNS resolver
/// Must be called before servers are actually run
pub fn set_dns_config(addr: ResolverConfig) {
    *(&mut *GLOBAL_DNS_ADDRESS.lock()) = Some(addr);
}

fn get_dns_address() -> Option<ResolverConfig> {
    GLOBAL_DNS_ADDRESS.lock().clone()
}

fn init_resolver() -> AsyncResolver {
    use std::sync::{Arc, Condvar, Mutex};
    use std::thread;

    // We'll be using this condvar to get the Resolver from the thread...
    let pair = Arc::new((Mutex::new(None::<AsyncResolver>), Condvar::new()));
    let pair2 = pair.clone();

    // Spawn the runtime to a new thread...
    //
    // This thread will manage the actual resolution runtime
    thread::spawn(move || {
                      // A runtime for this new thread
                      let mut runtime = Runtime::new().expect("failed to launch Runtime");

                      // our platform independent future, result, see next blocks
                      let (resolver, bg) = {
                          // To make this independent, if targeting macOS, BSD, Linux, or Windows, we can use the system's configuration:
                          #[cfg(any(unix, windows))]
                          {
                              if let Some(conf) = get_dns_address() {
                                  use trust_dns_resolver::config::ResolverOpts;
                                  AsyncResolver::new(conf, ResolverOpts::default())
                              } else {
                                  use trust_dns_resolver::system_conf::read_system_conf;
                                  // use the system resolver configuration
                                  let (config, opts) = read_system_conf().expect("Failed to read global dns sysconf");
                                  AsyncResolver::new(config, opts)
                              }
                          }

                          // For other operating systems, we can use one of the preconfigured definitions
                          #[cfg(not(any(unix, windows)))]
                          {
                              // Directly reference the config types
                              use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};

                              if let Some(conf) = get_dns_address() {
                                  AsyncResolver::new(conf, ResolverOpts::default())
                              } else {
                                  // Get a new resolver with the google nameservers as the upstream recursive resolvers
                                  AsyncResolver::new(ResolverConfig::google(), ResolverOpts::default())
                              }
                          }
                      };

                      let &(ref lock, ref cvar) = &*pair2;
                      let mut started = lock.lock().unwrap();
                      *started = Some(resolver);
                      cvar.notify_one();
                      drop(started);

                      runtime.block_on(bg).expect("Failed to create DNS resolver");
                  });

    // Wait for the thread to start up.
    let &(ref lock, ref cvar) = &*pair;
    let mut resolver = lock.lock().unwrap();
    while resolver.is_none() {
        resolver = cvar.wait(resolver).unwrap();
    }

    // take the started resolver
    let resolver = mem::replace(&mut *resolver, None);

    // set the global resolver
    resolver.expect("resolver should not be none")
}

fn inner_resolve(config: Arc<Config>,
                 addr: &str,
                 port: u16,
                 check_forbidden: bool)
                 -> impl Future<Item = Vec<SocketAddr>, Error = io::Error> + Send {
    let owned_addr = addr.to_owned();
    let owned_addr2 = owned_addr.clone();

    GLOBAL_DNS_RESOLVER.lookup_ip(addr)
                       .map_err(move |err| {
                                    error!("Failed to resolve {}, err: {}", owned_addr2, err);
                                    io::Error::new(io::ErrorKind::Other, "dns resolve error")
                                })
                       .and_then(move |lookup_result| {
                                     let mut vaddr = Vec::new();
                                     for ip in lookup_result.iter() {
                                         if check_forbidden {
                                             let forbidden_ip = &config.forbidden_ip;
                                             if forbidden_ip.contains(&ip) {
                                                 debug!("Resolved {} => {}, which is skipped by forbidden_ip",
                                                        owned_addr, ip);
                                                 continue;
                                             }
                                         }
                                         vaddr.push(SocketAddr::new(ip, port));
                                     }

                                     if vaddr.is_empty() {
                                         let err = io::Error::new(ErrorKind::Other,
                                                       format!("resolved {} to empty address, all IPs are filtered",
                                                               owned_addr));
                                         Err(err)
                                     } else {
                                         debug!("Resolved {} => {:?}", owned_addr, vaddr);
                                         Ok(vaddr)
                                     }
                                 })
}

/// Resolve address to IP
pub fn resolve(config: Arc<Config>,
               addr: &str,
               port: u16,
               check_forbidden: bool)
               -> impl Future<Item = Vec<SocketAddr>, Error = io::Error> + Send {
    inner_resolve(config, addr, port, check_forbidden)
}
<commit_msg>Run trust-dns background resolving Future in current reactor<commit_after>//! Asynchronous DNS resolver

use std::io::{self, ErrorKind};
use std::net::SocketAddr;
use std::sync::Arc;

use futures::Future;
use spin::Mutex;
use tokio;
use trust_dns_resolver::config::ResolverConfig;
use trust_dns_resolver::AsyncResolver;

use config::Config;

// Taken from
// bluejekyll/trust-dns/resolver/examples/global_resolver.rs
lazy_static! {
    static ref GLOBAL_DNS_ADDRESS: Mutex<Option<ResolverConfig>> = Mutex::new(None);

    // First we need to setup the global Resolver
    static ref GLOBAL_DNS_RESOLVER: AsyncResolver = init_resolver();
}

/// Set address for global DNS resolver
/// Must be called before servers are actually run
pub fn set_dns_config(addr: ResolverConfig) {
    *(&mut *GLOBAL_DNS_ADDRESS.lock()) = Some(addr);
}

fn get_dns_address() -> Option<ResolverConfig> {
    GLOBAL_DNS_ADDRESS.lock().clone()
}

fn init_resolver() -> AsyncResolver {
    let (resolver, bg) = {
        // To make this independent, if targeting macOS, BSD, Linux, or Windows, we can use the system's configuration:
        #[cfg(any(unix, windows))]
        {
            if let Some(conf) = get_dns_address() {
                use trust_dns_resolver::config::ResolverOpts;
                AsyncResolver::new(conf, ResolverOpts::default())
            } else {
                use trust_dns_resolver::system_conf::read_system_conf;
                // use the system resolver configuration
                let (config, opts) = read_system_conf().expect("Failed to read global dns sysconf");
                AsyncResolver::new(config, opts)
            }
        }

        // For other operating systems, we can use one of the preconfigured definitions
        #[cfg(not(any(unix, windows)))]
        {
            // Directly reference the config types
            use trust_dns_resolver::config::{ResolverConfig, ResolverOpts};

            if let Some(conf) = get_dns_address() {
                AsyncResolver::new(conf, ResolverOpts::default())
            } else {
                // Get a new resolver with the google nameservers as the upstream recursive resolvers
                AsyncResolver::new(ResolverConfig::google(), ResolverOpts::default())
            }
        }
    };

    // NOTE: resolving will always be called inside a future.
    tokio::spawn(bg);

    resolver
}

fn inner_resolve(config: Arc<Config>,
                 addr: &str,
                 port: u16,
                 check_forbidden: bool)
                 -> impl Future<Item = Vec<SocketAddr>, Error = io::Error> + Send {
    let owned_addr = addr.to_owned();
    let owned_addr2 = owned_addr.clone();

    GLOBAL_DNS_RESOLVER.lookup_ip(addr)
                       .map_err(move |err| {
                                    error!("Failed to resolve {}, err: {}", owned_addr2, err);
                                    io::Error::new(io::ErrorKind::Other, "dns resolve error")
                                })
                       .and_then(move |lookup_result| {
                                     let mut vaddr = Vec::new();
                                     for ip in lookup_result.iter() {
                                         if check_forbidden {
                                             let forbidden_ip = &config.forbidden_ip;
                                             if forbidden_ip.contains(&ip) {
                                                 debug!("Resolved {} => {}, which is skipped by forbidden_ip",
                                                        owned_addr, ip);
                                                 continue;
                                             }
                                         }
                                         vaddr.push(SocketAddr::new(ip, port));
                                     }

                                     if vaddr.is_empty() {
                                         let err = io::Error::new(ErrorKind::Other,
                                                       format!("resolved {} to empty address, all IPs are filtered",
                                                               owned_addr));
                                         Err(err)
                                     } else {
                                         debug!("Resolved {} => {:?}", owned_addr, vaddr);
                                         Ok(vaddr)
                                     }
                                 })
}

/// Resolve address to IP
pub fn resolve(config: Arc<Config>,
               addr: &str,
               port: u16,
               check_forbidden: bool)
               -> impl Future<Item = Vec<SocketAddr>, Error = io::Error> + Send {
    inner_resolve(config, addr, port, check_forbidden)
}
<|endoftext|>


<commit_before>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = acl.check_name_in_proxy_list(query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => panic!("PTR records should not reach here"),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match acl.check_name_in_proxy_list(name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::from(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::from(*ip)),
                    RData::PTR(_) => panic!("PTR records should not reach here"),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        panic!("should not reach here")
    }
}

async fn acl_lookup<Local, Remote>(
    acl: &Option<AccessControl>,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await;

    if should_forward_by_response(acl, &local_response, query) {
        let remote_response = remote_response_fut.await;
        debug!("pick remote response (response): {:?}", remote_response);
        (remote_response, true)
    } else {
        debug!("pick local response (response): {:?}", local_response);
        (local_response, false)
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
        let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
        let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if request.queries().is_empty() {
                message.set_response_code(ResponseCode::FormErr);
            } else {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<commit_msg>Check for DNS RD bit<commit_after>use std::{
    io,
    net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr},
    sync::Arc,
    time::Duration,
};

use tokio::sync::mpsc;

use log::{debug, error, info, warn};
use trust_dns_proto::{
    op::{header::MessageType, response_code::ResponseCode, Message, Query},
    rr::{DNSClass, Name, RData, RecordType},
};

use crate::{
    acl::AccessControl,
    config::ConfigType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        sys::create_udp_socket,
        utils::try_timeout,
    },
};

mod upstream;

fn should_forward_by_ptr_name(acl: &AccessControl, name: &Name) -> bool {
    let mut iter = name.iter().rev();
    let mut next = || std::str::from_utf8(iter.next().unwrap_or(&[48])).unwrap_or("*");
    if !"arpa".eq_ignore_ascii_case(next()) {
        return acl.is_default_in_proxy_list();
    }
    match &next().to_ascii_lowercase()[..] {
        "in-addr" => {
            let mut octets: [u8; 4] = [0; 4];
            for octet in octets.iter_mut() {
                match next().parse() {
                    Ok(result) => *octet = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V4(Ipv4Addr::new(octets[0], octets[1], octets[2], octets[3])))
        }
        "ip6" => {
            let mut segments: [u16; 8] = [0; 8];
            for segment in segments.iter_mut() {
                match u16::from_str_radix(&[next(), next(), next(), next()].concat(), 16) {
                    Ok(result) => *segment = result,
                    Err(_) => return acl.is_default_in_proxy_list(),
                }
            }
            acl.check_ip_in_proxy_list(&IpAddr::V6(Ipv6Addr::new(
                segments[0], segments[1], segments[2], segments[3], segments[4], segments[5], segments[6], segments[7]
            )))
        }
        _ => acl.is_default_in_proxy_list(),
    }
}

/// given the query, determine whether remote/local query should be used, or inconclusive
fn should_forward_by_query(acl: &Option<AccessControl>, query: &Query) -> Option<bool> {
    if let Some(acl) = acl {
        if query.query_class() != DNSClass::IN {
            // unconditionally use default for all non-IN queries
            Some(acl.is_default_in_proxy_list())
        } else if query.query_type() == RecordType::PTR {
            Some(should_forward_by_ptr_name(acl, query.name()))
        } else {
            let result = acl.check_name_in_proxy_list(query.name());
            if result == None && match query.query_type() {
                RecordType::A => acl.is_ipv4_empty(),
                RecordType::AAAA => acl.is_ipv6_empty(),
                RecordType::ANY => acl.is_ipv4_empty() && acl.is_ipv6_empty(),
                RecordType::PTR => panic!("PTR records should not reach here"),
                _ => true,
            } {
                Some(acl.is_default_in_proxy_list())
            } else {
                result
            }
        }
    } else {
        Some(true)
    }
}

/// given the local response, determine whether remote response should be used instead
fn should_forward_by_response(
    acl: &Option<AccessControl>,
    local_response: &io::Result<Message>,
    query: &Query,
) -> bool {
    if let Some(acl) = acl {
        macro_rules! examine_record {
            ($rec:ident, $is_answer:expr) => {
                if let RData::CNAME(ref name) = $rec.rdata() {
                    match acl.check_name_in_proxy_list(name) {
                        Some(value) => return value,
                        None => continue,
                    }
                } else if $is_answer && !query.query_type().is_any() && $rec.record_type() != query.query_type() {
                    warn!("local DNS response has inconsistent answer type {} for query {}", $rec.record_type(), query);
                    return true;
                }
                let forward = match $rec.rdata() {
                    RData::A(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::from(*ip)),
                    RData::AAAA(ref ip) => acl.check_ip_in_proxy_list(&IpAddr::from(*ip)),
                    RData::PTR(_) => panic!("PTR records should not reach here"),
                    _ => acl.is_default_in_proxy_list(),
                };
                if !forward {
                    return false;
                }
            };
        }
        if let Ok(ref local_response) = local_response {
            for rec in local_response.answers() {
                examine_record!(rec, true);
            }
            for rec in local_response.additionals() {
                examine_record!(rec, false);
            }
        }
        true
    } else {
        panic!("should not reach here")
    }
}

async fn acl_lookup<Local, Remote>(
    acl: &Option<AccessControl>,
    local: Arc<Local>,
    remote: Arc<Remote>,
    query: &Query
) -> (io::Result<Message>, bool)
    where
        Local: upstream::Upstream,
        Remote: upstream::Upstream,
{
    // Start querying name servers
    debug!(
        "attempting lookup of {:?} {} with ns {:?} and {:?}",
        query.query_type(), query.name(), local, remote
    );

    let remote_response_fut = try_timeout(remote.lookup(query), Some(Duration::new(3, 0)));
    let local_response_fut = try_timeout(local.lookup(query), Some(Duration::new(3, 0)));

    match should_forward_by_query(acl, query) {
        Some(true) => {
            let remote_response = remote_response_fut.await;
            debug!("pick remote response (query): {:?}", remote_response);
            return (remote_response, true);
        }
        Some(false) => {
            let local_response = local_response_fut.await;
            debug!("pick local response (query): {:?}", local_response);
            return (local_response, false);
        }
        None => (),
    }

    // FIXME: spawn(remote_response_fut)
    let local_response = local_response_fut.await;

    if should_forward_by_response(acl, &local_response, query) {
        let remote_response = remote_response_fut.await;
        debug!("pick remote response (response): {:?}", remote_response);
        (remote_response, true)
    } else {
        debug!("pick local response (response): {:?}", local_response);
        (local_response, false)
    }
}

/// Start a DNS relay local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = match context.config().config_type {
        ConfigType::DnsLocal => {
            // Standalone server
            context.config().local_addr.as_ref().expect("local config")
        }
        c if c.is_local() => {
            // Integrated mode
            context.config().dns_local_addr.as_ref().expect("dns relay addr")
        }
        _ => {
            panic!("ConfigType must be DnsLocal");
        }
    };

    let bind_addr = local_addr.bind_addr(&context).await?;

    let socket = create_udp_socket(&bind_addr).await?;

    let actual_local_addr = socket.local_addr()?;
    info!("shadowsocks DNS relay listening on {}", actual_local_addr);

    let (mut rx, mut tx) = socket.split();
    let (qtx, mut qrx) = mpsc::channel::<(SocketAddr, Vec<u8>)>(1024);

    tokio::spawn(async move {
        while let Some((src, pkt)) = qrx.recv().await {
            if let Err(err) = tx.send_to(&pkt, &src).await {
                error!("failed to send packet {} bytes to {}, error: {}", pkt.len(), src, err);
            }
        }
    });

    let config = context.config();
    #[cfg(target_os = "android")]
        let local_upstream = Arc::new(upstream::UnixSocketUpstream {
        path: config.local_dns_path.clone().expect("local query DNS path"),
    });
    #[cfg(not(target_os = "android"))]
        let local_upstream = Arc::new(upstream::UdpUpstream {
        server: config.local_dns_addr.clone().expect("local query DNS address"),
    });
    // FIXME: We use TCP to send remote queries by default, which should be configuable.
    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Tcp).await;
    let remote_upstream = Arc::new(upstream::ProxyTcpUpstream {
        context: context.clone(),
        svr_cfg: move || balancer.pick_server().server_config().clone(),
        ns: config.remote_dns_addr.clone().expect("remote query DNS address"),
    });

    loop {
        let mut req_buffer: [u8; 512] = [0; 512];
        let (_, src) = match rx.recv_from(&mut req_buffer).await {
            Ok(x) => x,
            Err(e) => {
                error!("DNS relay read from UDP socket error: {}", e);
                continue;
            }
        };

        let request = match Message::from_vec(&req_buffer) {
            Ok(x) => x,
            Err(e) => {
                error!("failed to parse UDP query message, error: {:?}", e);
                continue;
            }
        };

        debug!("received src: {}, query: {:?}", src, request);

        let context = context.clone();
        let local_upstream = Arc::clone(&local_upstream);
        let remote_upstream = Arc::clone(&remote_upstream);
        let mut qtx = qtx.clone();

        tokio::spawn(async move {
            let mut message = Message::new();
            message.set_id(request.id());
            message.set_recursion_desired(true);
            message.set_recursion_available(true);
            message.set_message_type(MessageType::Response);

            if !request.recursion_desired() {
                message.set_recursion_desired(false);
                message.set_response_code(ResponseCode::NotImp);
            } else if request.query_count() > 0 {
                let question = &request.queries()[0];
                let (r, forward) = acl_lookup(context.acl(), local_upstream, remote_upstream, question).await;

                if let Ok(result) = r {
                    for rec in result.answers() {
                        debug!("dns answer: {:?}", rec);

                        match rec.rdata() {
                            RData::A(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            RData::AAAA(ref ip) => context.add_to_reverse_lookup_cache(&IpAddr::from(*ip), forward),
                            _ => (),
                        }
                    }

                    message = result;
                    message.set_id(request.id());
                } else {
                    message.set_response_code(ResponseCode::ServFail);
                }
            }

            debug!("DNS src: {}, final response: {:?}", src, message);

            match message.to_vec() {
                Err(err) => {
                    error!("failed to serialize message, error: {}", err);
                }
                Ok(res_buffer) => {
                    if let Err(..) = qtx.send((src, res_buffer)).await {
                        error!("DNS send back queue is closed unexpectly");
                    }
                }
            }
        });
    }
}
<|endoftext|>