<commit_before>//! Regular expression virtual machine.

use std;
use std::mem::swap;
use std::trie::TrieSet;

/// A compiled regular expression, ready to execute.
pub type Program = (~[uint], ~[State]);

/// A single state in the machine.
pub type State = (Want, ~[uint]);

pub enum Want {
    /// Match any code point in the range, inclusive.
    Range(char, char)
}

struct Thread {
    pc: uint
}

impl Thread {
    fn new(pc: uint) -> Thread {
        Thread {
            pc: pc
        }
    }

    fn with_pc(&self, pc: uint) -> Thread {
        Thread {
            pc: pc
        }
    }
}

struct ThreadList {
    threads: ~[Thread],
    indices: TrieSet
}

impl ThreadList {
    /// Create a new, empty, `ThreadList`.
    fn new() -> ThreadList {
        ThreadList {
            threads: ~[],
            indices: TrieSet::new()
        }
    }

    /// Clear all threads from the list.
    fn clear(&mut self) {
        self.threads.clear();
        self.indices.clear();
    }

    /// Add a thread to the list, if one with the same `pc` is not
    /// already present.
    fn add(&mut self, t: Thread) {
        if self.indices.insert(t.pc) {
            self.threads.push(t);
        }
    }

    /// Iterate over the list of threads.
    fn iter<'a>(&'a self) -> std::vec::Items<'a, Thread> {
        self.threads.iter()
    }
}

/// A regular expression virtual machine, loosely based on the Pike VM.
pub struct VM<'a> {
    priv states: &'a [State],
    priv threads: ThreadList,
    priv next: ThreadList,
    priv matched: bool
}

impl<'a> VM<'a> {
    pub fn new(&(ref initial, ref states): &'a Program) -> VM<'a> {
        let mut vm = VM {
            states: *states,
            threads: ThreadList::new(),
            next: ThreadList::new(),
            matched: false
        };

        // Add the initial thread
        vm.matched = follow(Thread::new(0), *initial, vm.states, &mut vm.threads);

        vm
    }

    /// Feed a character into the automaton.
    pub fn feed(&mut self, c: char) {
        self.matched = false;

        // Run through all the threads
        for &t in self.threads.iter() {
            let (want, ref tails) = self.states[t.pc];
            match want {
                Range(lo, hi) => if lo <= c && c <= hi {
                    if follow(t, *tails, self.states, &mut self.next) {
                        self.matched = true;
                        // Cut off lower priority threads
                        break;
                    }
                }
            }
        }

        // Swap the thread buffers
        swap(&mut self.threads, &mut self.next);
        self.next.clear();
    }

    /// Determine if we have a match, given the existing input.
    pub fn is_match(&self) -> bool {
        self.matched
    }
}

/// Add the targets of the current instruction to the thread list.
/// Returns `true` if a matching state is reached; otherwise `false`.
fn follow(t: Thread, tails: &[uint], states: &[State], threads: &mut ThreadList) -> bool {
    let mut matched = false;
    for &tail in tails.iter() {
        if tail == states.len() {
            matched = true;
        } else {
            threads.add(t.with_pc(tail));
        }
    }
    matched
}
<commit_msg>Cosmetic fixes<commit_after>//! Regular expression virtual machine.

use std;
use std::mem::swap;
use std::trie::TrieSet;

/// A compiled regular expression, ready to execute.
pub type Program = (~[uint], ~[State]);

/// A single state in the machine.
pub type State = (Want, ~[uint]);

pub enum Want {
    /// Match any code point in the range, inclusive.
    Range(char, char)
}

struct Thread {
    pc: uint
}

impl Thread {
    fn new(pc: uint) -> Thread {
        Thread {
            pc: pc
        }
    }

    fn with_pc(&self, pc: uint) -> Thread {
        Thread {
            pc: pc
        }
    }
}

struct ThreadList {
    threads: ~[Thread],
    indices: TrieSet
}

impl ThreadList {
    /// Create a new, empty, `ThreadList`.
    fn new() -> ThreadList {
        ThreadList {
            threads: ~[],
            indices: TrieSet::new()
        }
    }

    /// Clear all threads from the list.
    fn clear(&mut self) {
        self.threads.clear();
        self.indices.clear();
    }

    /// Add a thread to the list, if one with the same `pc` is not
    /// already present.
    fn add(&mut self, t: Thread) {
        if self.indices.insert(t.pc) {
            self.threads.push(t);
        }
    }

    /// Iterate over the list of threads.
    fn iter<'a>(&'a self) -> std::vec::Items<'a, Thread> {
        self.threads.iter()
    }
}

/// A regular expression virtual machine, loosely based on the Pike VM.
pub struct VM<'a> {
    priv states: &'a [State],
    priv threads: ThreadList,
    priv next: ThreadList,
    priv matched: bool
}

impl<'a> VM<'a> {
    pub fn new(&(ref initial, ref states): &'a Program) -> VM<'a> {
        let mut vm = VM {
            states: *states,
            threads: ThreadList::new(),
            next: ThreadList::new(),
            matched: false
        };

        // Add the initial thread
        vm.matched = follow(Thread::new(0), *initial, vm.states, &mut vm.threads);

        vm
    }

    /// Feed a character into the automaton.
    pub fn feed(&mut self, c: char) {
        self.matched = false;

        // Run through all the threads
        for &t in self.threads.iter() {
            let (want, ref tails) = self.states[t.pc];
            match want {
                Range(lo, hi) => if lo <= c && c <= hi {
                    if follow(t, *tails, self.states, &mut self.next) {
                        self.matched = true;
                        // Cut off lower priority threads
                        break
                    }
                }
            }
        }

        // Swap the thread buffers
        swap(&mut self.threads, &mut self.next);
        self.next.clear();
    }

    /// Determine if we have a match, given the existing input.
    pub fn is_match(&self) -> bool {
        self.matched
    }
}

/// Add the targets of the current instruction to the thread list.
/// Returns `true` if a matching state is reached; otherwise `false`.
fn follow(t: Thread, tails: &[uint], states: &[State], threads: &mut ThreadList) -> bool {
    let mut matched = false;
    for &tail in tails.iter() {
        if tail == states.len() {
            matched = true;
        } else {
            threads.add(t.with_pc(tail));
        }
    }
    matched
}
<|endoftext|>


<commit_before>type Cell = u8;

pub struct VM {
    cells: [Cell; 30_000],
    data_pointer: usize,
}

impl VM {
    pub fn new() -> VM {
        VM { cells: [0; 30_000], data_pointer: 0 }
    }

    pub fn output(&self) -> Cell {
        self.cells[self.data_pointer]
    }
}
<commit_msg>Add VM.increment(&mut self)<commit_after>type Cell = u8;

pub struct VM {
    cells: [Cell; 30_000],
    data_pointer: usize,
}

impl VM {
    pub fn new() -> VM {
        VM { cells: [0; 30_000], data_pointer: 0 }
    }

    pub fn output(&self) -> Cell {
        self.cells[self.data_pointer]
    }

    pub fn increment(&mut self) {
        let value = self.cells[self.data_pointer];
        self.cells[self.data_pointer] = value + 1;
    }
}
<|endoftext|>


<commit_before>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.
use std;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use futures::{Async, Future, Poll, Stream};

use tokio_core::reactor::Core;
use tokio_reactor::Handle;
use tokio_tcp;
use tokio_udp;

use trust_dns::error::*;
use trust_dns::serialize::binary::{BinDecodable, BinDecoder};
use trust_dns::tcp::TcpStream;
use trust_dns::udp::UdpStream;
use trust_dns::BufStreamHandle;

#[cfg(feature = "tls")]
use trust_dns_openssl::{tls_server, TlsStream};

#[cfg(feature = "tls")]
use trust_dns_openssl::tls_server::*;

use authority::MessageRequest;
use server::{Request, RequestHandler, ResponseHandle, TimeoutStream};

// TODO, would be nice to have a Slab for buffers here...

/// A Futures based implementation of a DNS server
pub struct ServerFuture<T: RequestHandler + 'static> {
    io_loop: Core,
    handler: Arc<T>,
}

impl<T: RequestHandler> ServerFuture<T> {
    /// Creates a new ServerFuture with the specified Handler.
    pub fn new(handler: T) -> io::Result<ServerFuture<T>> {
        Ok(ServerFuture {
            io_loop: Core::new()?,
            handler: Arc::new(handler),
        })
    }

    /// Register a UDP socket. Should be bound before calling this function.
    pub fn register_socket(&self, socket: tokio_udp::UdpSocket) {
        debug!("registered udp: {:?}", socket);

        // create the new UdpStream
        let (buf_stream, stream_handle) = UdpStream::with_bound(socket);
        //let request_stream = RequestStream::new(buf_stream, stream_handle);
        let handler = self.handler.clone();

        // this spawns a ForEach future which handles all the requests into a Handler.
        self.io_loop.handle().spawn(
            buf_stream
                .for_each(move |(buffer, src_addr)| {
                    Self::handle_request(buffer, src_addr, stream_handle.clone(), handler.clone())
                        .map_err(move |e| {
                            debug!("error parsing UDP request src: {:?} error: {}", src_addr, e)
                        })
                        .ok();

                    // continue processing...
                    Ok(())
                })
                .map_err(|e| panic!("error in UDP request_stream handler: {}", e)),
        );
    }

    /// Register a UDP socket. Should be bound before calling this function.
    pub fn register_socket_std(&self, socket: std::net::UdpSocket) {
        self.register_socket(tokio_udp::UdpSocket::from_std(socket, &Handle::current()).expect("bad handle?"))
    }

    /// Register a TcpListener to the Server. This should already be bound to either an IPv6 or an
    ///  IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    pub fn register_listener(
        &self,
        listener: tokio_tcp::TcpListener,
        timeout: Duration,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        debug!("registered tcp: {:?}", listener);

        // for each incoming request...
        self.io_loop.handle().spawn(
            listener
                .incoming()
                .for_each(move |tcp_stream| {
                    let src_addr = tcp_stream.peer_addr().unwrap();
                    debug!("accepted request from: {}", src_addr);
                    // take the created stream...
                    let (buf_stream, stream_handle) = TcpStream::from_stream(tcp_stream, src_addr);
                    let timeout_stream = TimeoutStream::new(buf_stream, timeout);
                    //let request_stream = RequestStream::new(timeout_stream, stream_handle);
                    let handler = handler.clone();

                    // and spawn to the io_loop
                    handle.spawn(
                        timeout_stream
                            .for_each(move |(buffer, src_addr)| {
                                Self::handle_request(
                                    buffer,
                                    src_addr,
                                    stream_handle.clone(),
                                    handler.clone(),
                                )
                            })
                            .map_err(move |e| {
                                debug!(
                                    "error in TCP request_stream src: {:?} error: {}",
                                    src_addr, e
                                )
                            }),
                    );

                    Ok(())
                })
                .map_err(|e| panic!("error in inbound tcp_stream: {}", e)),
        );

        Ok(())
    }

    /// Register a TcpListener to the Server. This should already be bound to either an IPv6 or an
    ///  IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    pub fn register_listener_std(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
    ) -> io::Result<()> {
        self.register_listener(tokio_tcp::TcpListener::from_std(listener, &Handle::current())?, timeout)
    }

    /// Register a TlsListener to the Server. The TlsListener should already be bound to either an
    /// IPv6 or an IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP (needs to be on a different port from standard TCP connections) socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    /// * `pkcs12` - certificate used to announce to clients
    #[cfg(feature = "tls")]
    pub fn register_tls_listener(
        &self,
        listener: tokio_tcp::TcpListener,
        timeout: Duration,
        pkcs12: ParsedPkcs12,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        debug!("registered tcp: {:?}", listener);

        let tls_acceptor = tls_server::new_acceptor(&pkcs12)?;

        // for each incoming request...
        self.io_loop.handle().spawn(
            listener
                .incoming()
                .for_each(move |tcp_stream| {
                    let src_addr = tcp_stream.peer_addr().unwrap();
                    debug!("accepted request from: {}", src_addr);
                    let handle = handle.clone();
                    let handler = handler.clone();

                    // take the created stream...
                    tls_acceptor
                        .accept_async(tcp_stream)
                        .map_err(|e| {
                            io::Error::new(
                                io::ErrorKind::ConnectionRefused,
                                format!("tls error: {}", e),
                            )
                        })
                        .and_then(move |tls_stream| {
                            let (buf_stream, stream_handle) =
                                TlsStream::from_stream(tls_stream, src_addr);
                            let timeout_stream = TimeoutStream::new(buf_stream, timeout);
                            //let request_stream = RequestStream::new(timeout_stream, stream_handle);
                            let handler = handler.clone();

                            // and spawn to the io_loop
                            handle.spawn(
                                timeout_stream
                                    .for_each(move |(buffer, addr)| {
                                        Self::handle_request(
                                            buffer,
                                            addr,
                                            stream_handle.clone(),
                                            handler.clone(),
                                        )
                                    })
                                    .map_err(move |e| {
                                        debug!(
                                            "error in TCP request_stream src: {:?} error: {}",
                                            src_addr, e
                                        )
                                    }),
                            );

                            Ok(())
                        })
                    //.map_err(move |e| debug!("error TLS handshake: {:?} error: {:?}", src_addr, e))
                })
                .map_err(|e| panic!("error in inbound tls_stream: {}", e)),
        );

        Ok(())
    }

    /// Register a TlsListener to the Server. The TlsListener should already be bound to either an
    /// IPv6 or an IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP (needs to be on a different port from standard TCP connections) socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    /// * `pkcs12` - certificate used to announce to clients
    #[cfg(feature = "dns-over-openssl")]
    pub fn register_tls_listener_std(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
        pkcs12: ParsedPkcs12,
    ) -> io::Result<()> {
        self.register_tls_listener(tokio_tcp::TcpListener::from_std(listener, &Handle::current())?, timeout, pkcs12)
    }

    /// TODO: how to do threads? should we do a bunch of listener threads and then query threads?
    /// Ideally the processing would be n-threads for recieving, which hand off to m-threads for
    ///  request handling. It would generally be the case that n <= m.
    pub fn listen(&mut self) -> io::Result<()> {
        info!("Server starting up");
        self.io_loop.run(Forever)?;

        Err(io::Error::new(
            io::ErrorKind::Interrupted,
            "Server stopping due to interruption",
        ))
    }

    /// Returns a reference to the tokio core loop driving this Server instance
    pub fn tokio_core(&mut self) -> &mut Core {
        &mut self.io_loop
    }

    fn handle_request(
        buffer: Vec<u8>,
        src_addr: SocketAddr,
        stream_handle: BufStreamHandle<ClientError>,
        handler: Arc<T>,
    ) -> io::Result<()> {
        let response_handle = ResponseHandle::new(src_addr, stream_handle);

        // TODO: rather than decoding the message here, this RequestStream should instead
        //       forward the request to another sender such that we could pull serialization off
        //       the IO thread.
        // decode any messages that are ready
        let mut decoder = BinDecoder::new(&buffer);
        let message = MessageRequest::read(&mut decoder)?;

        let request = Request {
            message: message,
            src: src_addr,
        };

        info!(
            "request: {} type: {:?} op_code: {:?} dnssec: {} {}",
            request.message.id(),
            request.message.message_type(),
            request.message.op_code(),
            request
                .message
                .edns()
                .map_or(false, |edns| edns.dnssec_ok()),
            request
                .message
                .queries()
                .first()
                .map(|q| q.original().to_string())
                .unwrap_or_else(|| "empty_queries".to_string()),
        );

        handler.handle_request(&request, response_handle)
    }
}

struct Forever;

impl Future for Forever {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // run forever...
        Ok(Async::NotReady)
    }
}
<commit_msg>update cfg feature references<commit_after>// Copyright 2015-2017 Benjamin Fry <benjaminfry@me.com>
//
// Licensed under the Apache License, Version 2.0, <LICENSE-APACHE or
// http://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.
use std;
use std::io;
use std::net::SocketAddr;
use std::sync::Arc;
use std::time::Duration;

use futures::{Async, Future, Poll, Stream};

use tokio_core::reactor::Core;
use tokio_reactor::Handle;
use tokio_tcp;
use tokio_udp;

use trust_dns::error::*;
use trust_dns::serialize::binary::{BinDecodable, BinDecoder};
use trust_dns::tcp::TcpStream;
use trust_dns::udp::UdpStream;
use trust_dns::BufStreamHandle;

#[cfg(feature = "dns-over-openssl")]
use trust_dns_openssl::{tls_server, TlsStream};

#[cfg(feature = "dns-over-openssl")]
use trust_dns_openssl::tls_server::*;

use authority::MessageRequest;
use server::{Request, RequestHandler, ResponseHandle, TimeoutStream};

// TODO, would be nice to have a Slab for buffers here...

/// A Futures based implementation of a DNS server
pub struct ServerFuture<T: RequestHandler + 'static> {
    io_loop: Core,
    handler: Arc<T>,
}

impl<T: RequestHandler> ServerFuture<T> {
    /// Creates a new ServerFuture with the specified Handler.
    pub fn new(handler: T) -> io::Result<ServerFuture<T>> {
        Ok(ServerFuture {
            io_loop: Core::new()?,
            handler: Arc::new(handler),
        })
    }

    /// Register a UDP socket. Should be bound before calling this function.
    pub fn register_socket(&self, socket: tokio_udp::UdpSocket) {
        debug!("registered udp: {:?}", socket);

        // create the new UdpStream
        let (buf_stream, stream_handle) = UdpStream::with_bound(socket);
        //let request_stream = RequestStream::new(buf_stream, stream_handle);
        let handler = self.handler.clone();

        // this spawns a ForEach future which handles all the requests into a Handler.
        self.io_loop.handle().spawn(
            buf_stream
                .for_each(move |(buffer, src_addr)| {
                    Self::handle_request(buffer, src_addr, stream_handle.clone(), handler.clone())
                        .map_err(move |e| {
                            debug!("error parsing UDP request src: {:?} error: {}", src_addr, e)
                        })
                        .ok();

                    // continue processing...
                    Ok(())
                })
                .map_err(|e| panic!("error in UDP request_stream handler: {}", e)),
        );
    }

    /// Register a UDP socket. Should be bound before calling this function.
    pub fn register_socket_std(&self, socket: std::net::UdpSocket) {
        self.register_socket(tokio_udp::UdpSocket::from_std(socket, &Handle::current()).expect("bad handle?"))
    }

    /// Register a TcpListener to the Server. This should already be bound to either an IPv6 or an
    ///  IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    pub fn register_listener(
        &self,
        listener: tokio_tcp::TcpListener,
        timeout: Duration,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        debug!("registered tcp: {:?}", listener);

        // for each incoming request...
        self.io_loop.handle().spawn(
            listener
                .incoming()
                .for_each(move |tcp_stream| {
                    let src_addr = tcp_stream.peer_addr().unwrap();
                    debug!("accepted request from: {}", src_addr);
                    // take the created stream...
                    let (buf_stream, stream_handle) = TcpStream::from_stream(tcp_stream, src_addr);
                    let timeout_stream = TimeoutStream::new(buf_stream, timeout);
                    //let request_stream = RequestStream::new(timeout_stream, stream_handle);
                    let handler = handler.clone();

                    // and spawn to the io_loop
                    handle.spawn(
                        timeout_stream
                            .for_each(move |(buffer, src_addr)| {
                                Self::handle_request(
                                    buffer,
                                    src_addr,
                                    stream_handle.clone(),
                                    handler.clone(),
                                )
                            })
                            .map_err(move |e| {
                                debug!(
                                    "error in TCP request_stream src: {:?} error: {}",
                                    src_addr, e
                                )
                            }),
                    );

                    Ok(())
                })
                .map_err(|e| panic!("error in inbound tcp_stream: {}", e)),
        );

        Ok(())
    }

    /// Register a TcpListener to the Server. This should already be bound to either an IPv6 or an
    ///  IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    pub fn register_listener_std(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
    ) -> io::Result<()> {
        self.register_listener(tokio_tcp::TcpListener::from_std(listener, &Handle::current())?, timeout)
    }

    /// Register a TlsListener to the Server. The TlsListener should already be bound to either an
    /// IPv6 or an IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP (needs to be on a different port from standard TCP connections) socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    /// * `pkcs12` - certificate used to announce to clients
    #[cfg(feature = "dns-over-openssl")]
    pub fn register_tls_listener(
        &self,
        listener: tokio_tcp::TcpListener,
        timeout: Duration,
        pkcs12: ParsedPkcs12,
    ) -> io::Result<()> {
        let handle = self.io_loop.handle();
        let handler = self.handler.clone();
        debug!("registered tcp: {:?}", listener);

        let tls_acceptor = tls_server::new_acceptor(&pkcs12)?;

        // for each incoming request...
        self.io_loop.handle().spawn(
            listener
                .incoming()
                .for_each(move |tcp_stream| {
                    let src_addr = tcp_stream.peer_addr().unwrap();
                    debug!("accepted request from: {}", src_addr);
                    let handle = handle.clone();
                    let handler = handler.clone();

                    // take the created stream...
                    tls_acceptor
                        .accept_async(tcp_stream)
                        .map_err(|e| {
                            io::Error::new(
                                io::ErrorKind::ConnectionRefused,
                                format!("tls error: {}", e),
                            )
                        })
                        .and_then(move |tls_stream| {
                            let (buf_stream, stream_handle) =
                                TlsStream::from_stream(tls_stream, src_addr);
                            let timeout_stream = TimeoutStream::new(buf_stream, timeout);
                            //let request_stream = RequestStream::new(timeout_stream, stream_handle);
                            let handler = handler.clone();

                            // and spawn to the io_loop
                            handle.spawn(
                                timeout_stream
                                    .for_each(move |(buffer, addr)| {
                                        Self::handle_request(
                                            buffer,
                                            addr,
                                            stream_handle.clone(),
                                            handler.clone(),
                                        )
                                    })
                                    .map_err(move |e| {
                                        debug!(
                                            "error in TCP request_stream src: {:?} error: {}",
                                            src_addr, e
                                        )
                                    }),
                            );

                            Ok(())
                        })
                    //.map_err(move |e| debug!("error TLS handshake: {:?} error: {:?}", src_addr, e))
                })
                .map_err(|e| panic!("error in inbound tls_stream: {}", e)),
        );

        Ok(())
    }

    /// Register a TlsListener to the Server. The TlsListener should already be bound to either an
    /// IPv6 or an IPv4 address.
    ///
    /// To make the server more resilient to DOS issues, there is a timeout. Care should be taken
    ///  to not make this too low depending on use cases.
    ///
    /// # Arguments
    /// * `listener` - a bound TCP (needs to be on a different port from standard TCP connections) socket
    /// * `timeout` - timeout duration of incoming requests, any connection that does not send
    ///               requests within this time period will be closed. In the future it should be
    ///               possible to create long-lived queries, but these should be from trusted sources
    ///               only, this would require some type of whitelisting.
    /// * `pkcs12` - certificate used to announce to clients
    #[cfg(feature = "dns-over-openssl")]
    pub fn register_tls_listener_std(
        &self,
        listener: std::net::TcpListener,
        timeout: Duration,
        pkcs12: ParsedPkcs12,
    ) -> io::Result<()> {
        self.register_tls_listener(tokio_tcp::TcpListener::from_std(listener, &Handle::current())?, timeout, pkcs12)
    }

    /// TODO: how to do threads? should we do a bunch of listener threads and then query threads?
    /// Ideally the processing would be n-threads for recieving, which hand off to m-threads for
    ///  request handling. It would generally be the case that n <= m.
    pub fn listen(&mut self) -> io::Result<()> {
        info!("Server starting up");
        self.io_loop.run(Forever)?;

        Err(io::Error::new(
            io::ErrorKind::Interrupted,
            "Server stopping due to interruption",
        ))
    }

    /// Returns a reference to the tokio core loop driving this Server instance
    pub fn tokio_core(&mut self) -> &mut Core {
        &mut self.io_loop
    }

    fn handle_request(
        buffer: Vec<u8>,
        src_addr: SocketAddr,
        stream_handle: BufStreamHandle<ClientError>,
        handler: Arc<T>,
    ) -> io::Result<()> {
        let response_handle = ResponseHandle::new(src_addr, stream_handle);

        // TODO: rather than decoding the message here, this RequestStream should instead
        //       forward the request to another sender such that we could pull serialization off
        //       the IO thread.
        // decode any messages that are ready
        let mut decoder = BinDecoder::new(&buffer);
        let message = MessageRequest::read(&mut decoder)?;

        let request = Request {
            message: message,
            src: src_addr,
        };

        info!(
            "request: {} type: {:?} op_code: {:?} dnssec: {} {}",
            request.message.id(),
            request.message.message_type(),
            request.message.op_code(),
            request
                .message
                .edns()
                .map_or(false, |edns| edns.dnssec_ok()),
            request
                .message
                .queries()
                .first()
                .map(|q| q.original().to_string())
                .unwrap_or_else(|| "empty_queries".to_string()),
        );

        handler.handle_request(&request, response_handle)
    }
}

struct Forever;

impl Future for Forever {
    type Item = ();
    type Error = io::Error;

    fn poll(&mut self) -> Poll<Self::Item, Self::Error> {
        // run forever...
        Ok(Async::NotReady)
    }
}
<|endoftext|>


<commit_before>//! A JSON-RPC 2.0 messages
//!
//! The main entrypoint here is the [Message](enum.Message.html). The others are just building
//! blocks and you should generally work with `Message` instead.

use std::str::FromStr;

use serde::ser::{Serialize, Serializer, SerializeStruct};
use serde::de::{Deserialize, Deserializer, Unexpected, Error};
use serde_json::{Value, from_value};

#[derive(Debug, Clone, PartialEq, Eq)]
struct Version;

impl Serialize for Version {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str("2.0")
    }
}

impl Deserialize for Version {
    fn deserialize<D: Deserializer>(deserializer: D) -> Result<Self, D::Error> {
        // The version is actually a string
        let parsed: String = Deserialize::deserialize(deserializer)?;
        if parsed == "2.0" {
            Ok(Version)
        } else {
            Err(D::Error::invalid_value(Unexpected::Str(&parsed), &"value 2.0"))
        }
    }
}

/// An RPC request
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct Request {
    jsonrpc: Version,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
    // TODO: Make private?
    pub id: Value,
}

/// An error code
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct RPCError {
    pub code: i64,
    pub message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<Value>,
}

/// A response to RPC
#[derive(Debug, Deserialize, Clone, PartialEq)]
pub struct Response {
    pub result: Result<Value, RPCError>,
    pub id: Value,
}

impl Serialize for Response {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut sub = serializer.serialize_struct("Response", 2)?;
        sub.serialize_field("id", &self.id)?;
        match self.result {
            Ok(ref value) => sub.serialize_field("result", value),
            Err(ref err) => sub.serialize_field("error", err),
        }?;
        sub.end()
    }
}

/// A notification (doesn't expect an answer)
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct Notification {
    jsonrpc: Version,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

/// One message of the JSON RPC protocol
///
/// One message, directly mapped from the structures of the protocol. See the
/// [specification](http://www.jsonrpc.org/specification) for more details.
///
/// Since the protocol allows one endpoint to be both client and server at the same time, the
/// message can decode and encode both directions of the protocol.
///
/// The `Unmatched` variant is for cases when the message that arrived is valid JSON, but doesn't
/// match the protocol. It allows for handling these non-fatal errors on higher level than the
/// parser.
///
/// It can be serialized and deserialized, or converted to and from a string.
///
/// The `Batch` variant is supposed to be created directly, without a constructor. The `Unmatched`
/// is something you may get from parsing but it is not expected you'd need to create it (though it
/// can be created directly as well).
#[derive(Debug, Clone, PartialEq)]
pub enum Message {
    Request(Request),
    Response(Response),
    Notification(Notification),
    Batch(Vec<Message>),
    Unmatched(Value),
}

impl Message {
    /// A constructor for a request.
    ///
    /// The ID is auto-generated.
    pub fn request(method: String, params: Option<Value>) -> Self {
        Message::Request(Request {
            jsonrpc: Version,
            method: method,
            params: params,
            // TODO!
            id: Value::Null,
        })
    }
    /// Answer the request with a (positive) reply.
    ///
    /// The ID is taken from the request.
    pub fn reply(request: &Request, reply: Value) -> Self {
        Message::Response(Response {
            result: Ok(reply),
            id: request.id.clone(),
        })
    }
    /// Answer the request with an error.
    ///
    /// The ID is taken from the request and the error structure is constructed.
    pub fn error(request: &Request, code: i64, message: String, data: Option<Value>) -> Self {
        Message::Response(Response {
            result: Err(RPCError {
                code: code,
                message: message,
                data: data,
            }),
            id: request.id.clone(),
        })
    }
    /// Create an error without a request.
    ///
    /// Create a top-level/free-standing error (one without an ID). This is the required answer for
    /// less serious protocol errors.
    pub fn top_error(code: i64, message: String, data: Option<Value>) -> Self {
        Message::Response(Response {
            result: Err(RPCError {
                code: code,
                message: message,
                data: data,
            }),
            id: Value::Null,
        })
    }
    /// A constructor for a notification.
    pub fn notification(method: String, params: Option<Value>) -> Self {
        Message::Notification(Notification {
            jsonrpc: Version,
            method: method,
            params: params,
        })
    }
}

impl Serialize for Message {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        match *self {
            Message::Request(ref req) => req.serialize(serializer),
            Message::Response(ref resp) => resp.serialize(serializer),
            Message::Notification(ref notif) => notif.serialize(serializer),
            Message::Batch(ref batch) => batch.serialize(serializer),
            Message::Unmatched(ref val) => val.serialize(serializer),
        }
    }
}

macro_rules! deser_branch {
    ($src:expr, $branch:ident) => {
        match from_value($src.clone()) {
            Ok(parsed) => return Message::$branch(parsed),
            Err(_) => (),
        }
    };
}

// TODO: This must be possible to do in less wasteful way. The cloning is stupid :-(
// There seems to be a feature coming in a future release that'd make all this unnecessary.
// https://github.com/serde-rs/serde/pull/739.
impl From<Value> for Message {
    fn from(v: Value) -> Self {
        // Try decoding it by each branch in sequence, taking the first one that matches
        deser_branch!(v, Request);
        deser_branch!(v, Response);
        deser_branch!(v, Notification);
        deser_branch!(v, Batch);
        Message::Unmatched(v)
    }
}

impl Deserialize for Message {
    fn deserialize<D: Deserializer>(deserializer: D) -> Result<Self, D::Error> {
        // Read it as a JSON (delegate the deserialization)
        let preparsed: Value = Deserialize::deserialize(deserializer)?;
        // Convert it
        Ok(Self::from(preparsed))
    }
}

impl FromStr for Message {
    type Err = ::serde_json::error::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        ::serde_json::de::from_str(s)
    }
}

impl Into<String> for Message {
    fn into(self) -> String {
        ::serde_json::ser::to_string(&self).unwrap()
    }
}
<commit_msg>Use the untagged enum feature of new serde<commit_after>//! A JSON-RPC 2.0 messages
//!
//! The main entrypoint here is the [Message](enum.Message.html). The others are just building
//! blocks and you should generally work with `Message` instead.

use std::str::FromStr;

use serde::ser::{Serialize, Serializer, SerializeStruct};
use serde::de::{Deserialize, Deserializer, Unexpected, Error};
use serde_json::Value;

#[derive(Debug, Clone, PartialEq, Eq)]
struct Version;

impl Serialize for Version {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        serializer.serialize_str("2.0")
    }
}

impl Deserialize for Version {
    fn deserialize<D: Deserializer>(deserializer: D) -> Result<Self, D::Error> {
        // The version is actually a string
        let parsed: String = Deserialize::deserialize(deserializer)?;
        if parsed == "2.0" {
            Ok(Version)
        } else {
            Err(D::Error::invalid_value(Unexpected::Str(&parsed), &"value 2.0"))
        }
    }
}

/// An RPC request
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct Request {
    jsonrpc: Version,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
    // TODO: Make private?
    pub id: Value,
}

/// An error code
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct RPCError {
    pub code: i64,
    pub message: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub data: Option<Value>,
}

/// A response to RPC
#[derive(Debug, Deserialize, Clone, PartialEq)]
pub struct Response {
    pub result: Result<Value, RPCError>,
    pub id: Value,
}

impl Serialize for Response {
    fn serialize<S: Serializer>(&self, serializer: S) -> Result<S::Ok, S::Error> {
        let mut sub = serializer.serialize_struct("Response", 2)?;
        sub.serialize_field("id", &self.id)?;
        match self.result {
            Ok(ref value) => sub.serialize_field("result", value),
            Err(ref err) => sub.serialize_field("error", err),
        }?;
        sub.end()
    }
}

// TODO: We probably need a custom deserialize here as well

/// A notification (doesn't expect an answer)
#[derive(Debug, Serialize, Deserialize, Clone, PartialEq)]
pub struct Notification {
    jsonrpc: Version,
    pub method: String,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub params: Option<Value>,
}

/// One message of the JSON RPC protocol
///
/// One message, directly mapped from the structures of the protocol. See the
/// [specification](http://www.jsonrpc.org/specification) for more details.
///
/// Since the protocol allows one endpoint to be both client and server at the same time, the
/// message can decode and encode both directions of the protocol.
///
/// The `Unmatched` variant is for cases when the message that arrived is valid JSON, but doesn't
/// match the protocol. It allows for handling these non-fatal errors on higher level than the
/// parser.
///
/// It can be serialized and deserialized, or converted to and from a string.
///
/// The `Batch` variant is supposed to be created directly, without a constructor. The `Unmatched`
/// is something you may get from parsing but it is not expected you'd need to create it (though it
/// can be created directly as well).
#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
#[serde(untagged)]
pub enum Message {
    Request(Request),
    Response(Response),
    Notification(Notification),
    Batch(Vec<Message>),
    Unmatched(Value),
}

impl Message {
    /// A constructor for a request.
    ///
    /// The ID is auto-generated.
    pub fn request(method: String, params: Option<Value>) -> Self {
        Message::Request(Request {
            jsonrpc: Version,
            method: method,
            params: params,
            // TODO: Generate the ID.
            id: Value::Null,
        })
    }
    /// Answer the request with a (positive) reply.
    ///
    /// The ID is taken from the request.
    pub fn reply(request: &Request, reply: Value) -> Self {
        Message::Response(Response {
            result: Ok(reply),
            id: request.id.clone(),
        })
    }
    /// Answer the request with an error.
    ///
    /// The ID is taken from the request and the error structure is constructed.
    pub fn error(request: &Request, code: i64, message: String, data: Option<Value>) -> Self {
        Message::Response(Response {
            result: Err(RPCError {
                code: code,
                message: message,
                data: data,
            }),
            id: request.id.clone(),
        })
    }
    /// Create an error without a request.
    ///
    /// Create a top-level/free-standing error (one without an ID). This is the required answer for
    /// less serious protocol errors.
    pub fn top_error(code: i64, message: String, data: Option<Value>) -> Self {
        Message::Response(Response {
            result: Err(RPCError {
                code: code,
                message: message,
                data: data,
            }),
            id: Value::Null,
        })
    }
    /// A constructor for a notification.
    pub fn notification(method: String, params: Option<Value>) -> Self {
        Message::Notification(Notification {
            jsonrpc: Version,
            method: method,
            params: params,
        })
    }
}

impl FromStr for Message {
    type Err = ::serde_json::error::Error;
    fn from_str(s: &str) -> Result<Self, Self::Err> {
        ::serde_json::de::from_str(s)
    }
}

impl Into<String> for Message {
    fn into(self) -> String {
        ::serde_json::ser::to_string(&self).unwrap()
    }
}
<|endoftext|>