<commit_before>/*
 * Copyright (c) 2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use capnp::any::{AnyPointer};
use capnp::capability;
use capnp::capability::{ClientHook, PipelineHook, PipelineOp, RemotePromise,
                        RequestHook, Request, Server};
use capnp::common;
use capnp::layout::{FromStructReader, FromStructBuilder, HasStructSize};
use capnp::message::{DEFAULT_READER_OPTIONS, MessageReader, MessageBuilder, MallocMessageBuilder};
use capnp::serialize;
use capnp::serialize::{OwnedSpaceMessageReader};
use std;
use std::any::AnyRefExt;
use std::hashmap::HashMap;
use rpc_capnp::{Message, Return, CapDescriptor, PromisedAnswer};

pub type QuestionId = u32;
pub type AnswerId = QuestionId;
pub type ExportId = u32;
pub type ImportId = ExportId;

pub struct Question {
    chan : std::comm::Chan<~OwnedSpaceMessageReader>,
    is_awaiting_return : bool,
}

pub struct Answer {
    result_exports : ~[ExportId]
}

pub struct Export {
    server_chan : std::comm::Chan<~OwnedSpaceMessageReader>,
}

impl Export {
    pub fn new(server : ~Server) -> Export {
        let (port, chan) = std::comm::Chan::<~OwnedSpaceMessageReader>::new();
        std::task::spawn(proc () {
                loop {
                    port.recv();
                }
            });

        Export { server_chan : chan }
    }
}

pub struct Import;

pub struct ImportTable<T> {
    slots : HashMap<u32, T>,
}

impl <T> ImportTable<T> {
    pub fn new() -> ImportTable<T> {
        ImportTable { slots : HashMap::new() }
    }
}

pub struct ExportTable<T> {
    slots : ~[T],
}

impl <T> ExportTable<T> {
    pub fn new() -> ExportTable<T> {
        ExportTable { slots : ~[] }
    }

    pub fn next(&mut self) -> u32 {
        fail!()
    }
}

pub struct RpcConnectionState {
    exports : ExportTable<Export>,
    questions : ExportTable<Question>,
    answers : ImportTable<Answer>,
    imports : ImportTable<Import>,
}

impl RpcConnectionState {
    pub fn new() -> RpcConnectionState {
        RpcConnectionState {
            exports : ExportTable::new(),
            questions : ExportTable::new(),
            answers : ImportTable::new(),
            imports : ImportTable::new(),
        }
    }

    pub fn run<T : std::io::Reader + Send, U : std::io::Writer + Send>(
        self, inpipe: T, outpipe: U)
         -> std::comm::SharedChan<RpcEvent> {

        let (port, chan) = std::comm::SharedChan::<RpcEvent>::new();

        let listener_chan = chan.clone();

        spawn(proc() {
                let mut r = inpipe;
                loop {
                    let message = box serialize::new_reader(
                        &mut r,
                        DEFAULT_READER_OPTIONS);
                    listener_chan.send(IncomingMessage(message));
                }
            });

        let loop_chan = chan.clone();

        let (writer_port, writer_chan) = std::comm::Chan::<~MallocMessageBuilder>::new();

        spawn(proc() {
                let mut w = outpipe;
                loop {
                    let message = writer_port.recv();
                    serialize::write_message(&mut w, message);
                }
            });

        spawn(proc() {
                let RpcConnectionState {mut questions, mut exports, answers, imports} = self;
                loop {
                    match port.recv() {
                        IncomingMessage(mut message) => {
                            let mut the_cap_table : ~[Option<~ClientHook>] = ~[];
                            let mut question = None::<u32>;
                            // populate the cap table
                            {
                                let root = message.get_root::<Message::Reader>();

                                match root.which() {
                                    Some(Message::Return(ret)) => {
                                        println!("got a return with answer id {}", ret.get_answer_id());

                                        match ret.which() {
                                            Some(Return::Results(payload)) => {
                                                println!("with a payload");
                                                let cap_table = payload.get_cap_table();
                                                for ii in range(0, cap_table.size()) {
                                                    match cap_table[ii].which() {
                                                        Some(CapDescriptor::None(())) => {
                                                            the_cap_table.push(None)
                                                        }
                                                        Some(CapDescriptor::SenderHosted(id)) => {
                                                            the_cap_table.push(Some(
                                                                    (box ImportClient {
                                                                            channel : loop_chan.clone(),
                                                                            import_id : id})
                                                                        as ~ClientHook));
                                                            println!("sender hosted: {}", id);
                                                        }
                                                        _ => {}
                                                    }
                                                }
                                            }
                                            Some(Return::Exception(e)) => {
                                                println!("exception: {}", e.get_reason());
                                            }
                                            _ => {}
                                        }

                                        question = Some(ret.get_answer_id());
                                    }
                                    Some(Message::Unimplemented(_)) => {
                                        println!("unimplemented");
                                    }
                                    Some(Message::Abort(exc)) => {
                                        println!("abort: {}", exc.get_reason());
                                    }
                                    None => {
                                        println!("Nothing there");
                                    }
                                    _ => {
                                        println!("something else");
                                    }
                                }
                            }
                            message.init_cap_table(the_cap_table);

                            match question {
                                Some(id) => {
                                    questions.slots[id].chan.try_send(message);
                                }
                                None => {}
                            }

                        }
                        Outgoing(OutgoingMessage { message : mut m,
                                                   answer_chan,
                                                   question_chan} ) => {
                            let root = m.get_root::<Message::Builder>();
                            // add a question to the question table
                            match root.which() {
                                Some(Message::Which::Return(_)) => {}
                                Some(Message::Which::Call(call)) => {
                                    call.set_question_id(questions.slots.len() as u32);
                                    questions.slots.push(Question {is_awaiting_return : true,
                                                                   chan : answer_chan} );
                                    question_chan.try_send(call.get_question_id());
                                }
                                Some(Message::Which::Restore(res)) => {
                                    res.set_question_id(questions.slots.len() as u32);
                                    questions.slots.push(Question {is_awaiting_return : true,
                                                                   chan : answer_chan} );
                                    question_chan.try_send(res.get_question_id());
                                }
                                _ => {
                                    error!("NONE OF THOSE");
                                }
                            }

                            // send
                            writer_chan.send(m);
                        }
                        NewLocalServer(server, export_chan) => {
                            let export_id = exports.slots.len() as u32;
                            export_chan.send(export_id);
                            exports.slots.push(Export::new(server));
                        }
                    }
                }});
        return chan;
    }
}

// HACK
pub enum OwnedCapDescriptor {
    NoDescriptor,
    SenderHosted(ExportId),
    SenderPromise(ExportId),
    ReceiverHosted(ImportId),
    ReceiverAnswer(QuestionId, ~[PipelineOp::Type]),
}

pub struct ImportClient {
    priv channel : std::comm::SharedChan<RpcEvent>,
    import_id : ImportId,
}

impl ClientHook for ImportClient {
    fn copy(&self) -> ~ClientHook {
        (box ImportClient {channel : self.channel.clone(),
                           import_id : self.import_id}) as ~ClientHook
    }

    fn new_call(&self, interface_id : u64, method_id : u16,
                _size_hint : Option<common::MessageSize>)
                -> capability::Request<AnyPointer::Builder, AnyPointer::Reader, AnyPointer::Pipeline> {
        let mut message = box MallocMessageBuilder::new_default();
        {
            let root : Message::Builder = message.get_root();
            let call = root.init_call();
            call.set_interface_id(interface_id);
            call.set_method_id(method_id);
            let target = call.init_target();
            target.set_imported_cap(self.import_id);
        }
        let hook = box RpcRequest { channel : self.channel.clone(),
                                    message : message };
        Request::new(hook as ~RequestHook)
    }

    fn get_descriptor(&self) -> ~std::any::Any {
        (box ReceiverHosted(self.import_id)) as ~std::any::Any
    }
}

pub struct PipelineClient {
    priv channel : std::comm::SharedChan<RpcEvent>,
    ops : ~[PipelineOp::Type],
    question_id : QuestionId,
}

impl ClientHook for PipelineClient {
    fn copy(&self) -> ~ClientHook {
        (~PipelineClient { channel : self.channel.clone(),
                           ops : self.ops.clone(),
                           question_id : self.question_id,
            }) as ~ClientHook
    }

    fn new_call(&self, interface_id : u64, method_id : u16,
                _size_hint : Option<common::MessageSize>)
                -> capability::Request<AnyPointer::Builder, AnyPointer::Reader, AnyPointer::Pipeline> {
        let mut message = box MallocMessageBuilder::new_default();
        {
            let root : Message::Builder = message.get_root();
            let call = root.init_call();
            call.set_interface_id(interface_id);
            call.set_method_id(method_id);
            let target = call.init_target();
            let promised_answer = target.init_promised_answer();
            promised_answer.set_question_id(self.question_id);
            let transform = promised_answer.init_transform(self.ops.len());
            for ii in range(0, self.ops.len()) {
                match self.ops[ii] {
                    PipelineOp::Noop => transform[ii].set_noop(()),
                    PipelineOp::GetPointerField(idx) => transform[ii].set_get_pointer_field(idx),
                }
            }
        }
        let hook = box RpcRequest { channel : self.channel.clone(),
                                    message : message };
        Request::new(hook as ~RequestHook)
    }

    fn get_descriptor(&self) -> ~std::any::Any {
        (box ReceiverAnswer(self.question_id, self.ops.clone())) as ~std::any::Any
    }
}

pub struct RpcRequest {
    priv channel : std::comm::SharedChan<RpcEvent>,
    priv message : ~MallocMessageBuilder,
}

impl RequestHook for RpcRequest {
    fn message<'a>(&'a mut self) -> &'a mut MallocMessageBuilder {
        &mut *self.message
    }
    fn send(~self) -> RemotePromise<AnyPointer::Reader, AnyPointer::Pipeline> {

        let ~RpcRequest { channel, mut message } = self;

        {
            let cap_table = {
                let mut caps = box [];
                for cap in message.get_cap_table().iter() {
                    match cap {
                        &Some(ref client_hook) => {
                            caps.push(client_hook.get_descriptor())
                        }
                        &None => {}
                    }
                }
                caps
            };
            let root : Message::Builder = message.get_root();
            match root.which() {
                Some(Message::Which::Call(call)) => {
                    let new_cap_table = call.get_params().init_cap_table(cap_table.len());
                    for ii in range(0, cap_table.len()) {
                        match cap_table[ii].as_ref::<OwnedCapDescriptor>() {
                            None => {fail!("noncompliant client hook")}
                            Some(&NoDescriptor) => {}
                            Some(&ReceiverHosted(import_id)) => {
                                new_cap_table[ii].set_receiver_hosted(import_id);
                            }
                            Some(&ReceiverAnswer(question_id,ref ops)) => {
                                let promised_answer = new_cap_table[ii].init_receiver_answer();
                                promised_answer.set_question_id(question_id);
                                let transform = promised_answer.init_transform(ops.len());
                                for ii in range(0, ops.len()) {
                                    match ops[ii] {
                                        PipelineOp::Noop => transform[ii].set_noop(()),
                                        PipelineOp::GetPointerField(idx) => transform[ii].set_get_pointer_field(idx),
                                    }
                                }
                            }
                            Some(&SenderHosted(export_id)) => {
                                new_cap_table[ii].set_sender_hosted(export_id);
                            }
                            _ => {}
                        }
                    }
                }
                _ => {}
            }
        }

        let (event, answer_port, question_port) = RpcEvent::new_outgoing(message);
        channel.send(event);

        let question_id = question_port.recv();

        let pipeline = ~RpcPipeline {channel : channel, question_id : question_id};
        let typeless = AnyPointer::Pipeline::new(pipeline as ~PipelineHook);

        RemotePromise {answer_port : answer_port, answer_result : None,
                       pipeline : typeless  }
    }
}

pub struct RpcPipeline {
    channel : std::comm::SharedChan<RpcEvent>,
    question_id : ExportId,
}

impl PipelineHook for RpcPipeline {
    fn copy(&self) -> ~PipelineHook {
        (~RpcPipeline { channel : self.channel.clone(),
                        question_id : self.question_id }) as ~PipelineHook
    }
    fn get_pipelined_cap(&self, ops : ~[PipelineOp::Type]) -> ~ClientHook {
        (~PipelineClient { channel : self.channel.clone(),
                           ops : ops,
                           question_id : self.question_id,
        }) as ~ClientHook
    }
}

pub trait InitRequest<'a, T> {
    fn init(&'a mut self) -> T;
}

impl <'a, Params : FromStructBuilder<'a> + HasStructSize, Results, Pipeline> InitRequest<'a, Params>
for Request<Params, Results, Pipeline> {
    fn init(&'a mut self) -> Params {
        let message : Message::Builder = self.hook.message().get_root();
        match message.which() {
            Some(Message::Which::Call(call)) => {
                let params = call.init_params();
                params.get_content().init_as_struct()
            }
            _ => fail!(),
        }
    }
}

pub trait WaitForContent<'a, T> {
    fn wait(&'a mut self) -> T;
}

impl <'a, Results : FromStructReader<'a>, Pipeline> WaitForContent<'a, Results>
for RemotePromise<Results, Pipeline> {
    fn wait(&'a mut self) -> Results {
        // XXX should check that it's not already been received.
        let message = self.answer_port.recv();
        self.answer_result = Some(message);
        match self.answer_result {
            None => unreachable!(),
            Some(ref message) => {
                let root : Message::Reader = message.get_root();
                match root.which() {
                    Some(Message::Return(ret)) => {
                        match ret.which() {
                            Some(Return::Results(res)) => {
                                res.get_content().get_as_struct()
                            }
                            _ => fail!(),
                        }
                    }
                    _ => {fail!()}
                }
            }
        }
    }
}

pub struct OutgoingMessage {
    message : ~MallocMessageBuilder,
    answer_chan : std::comm::Chan<~OwnedSpaceMessageReader>,
    question_chan : std::comm::Chan<QuestionId>,
}

pub enum RpcEvent {
    IncomingMessage(~serialize::OwnedSpaceMessageReader),
    Outgoing(OutgoingMessage),
    NewLocalServer(~Server, std::comm::Chan<ExportId>),
}

impl RpcEvent {
    pub fn new_outgoing(message : ~MallocMessageBuilder)
                        -> (RpcEvent, std::comm::Port<~OwnedSpaceMessageReader>,
                            std::comm::Port<ExportId>) {
        let (answer_port, answer_chan) = std::comm::Chan::<~OwnedSpaceMessageReader>::new();

        let (question_port, question_chan) = std::comm::Chan::<ExportId>::new();

        (Outgoing(OutgoingMessage{ message : message,
                                   answer_chan : answer_chan,
                                   question_chan : question_chan }),
         answer_port,
         question_port)
    }
}
<commit_msg>make populate_cap_table() its own function<commit_after>/*
 * Copyright (c) 2014, David Renshaw (dwrenshaw@gmail.com)
 *
 * See the LICENSE file in the capnproto-rust root directory.
 */

use capnp::any::{AnyPointer};
use capnp::capability;
use capnp::capability::{CallContextHook, ClientHook, PipelineHook, PipelineOp, RemotePromise,
                        RequestHook, Request, Server};
use capnp::common;
use capnp::layout::{FromStructReader, FromStructBuilder, HasStructSize};
use capnp::message::{DEFAULT_READER_OPTIONS, MessageReader, MessageBuilder, MallocMessageBuilder};
use capnp::serialize;
use capnp::serialize::{OwnedSpaceMessageReader};
use std;
use std::any::AnyRefExt;
use std::hashmap::HashMap;
use rpc_capnp::{Message, Return, CapDescriptor, PromisedAnswer, MessageTarget};

pub type QuestionId = u32;
pub type AnswerId = QuestionId;
pub type ExportId = u32;
pub type ImportId = ExportId;

pub struct Question {
    chan : std::comm::Chan<~OwnedSpaceMessageReader>,
    is_awaiting_return : bool,
}

pub struct Answer {
    result_exports : ~[ExportId]
}

pub struct Export {
    server_chan : std::comm::Chan<~OwnedSpaceMessageReader>,
}

impl Export {
    pub fn new(server : ~Server) -> Export {
        let (port, chan) = std::comm::Chan::<~OwnedSpaceMessageReader>::new();
        std::task::spawn(proc () {
                loop {
                    port.recv();
                }
            });

        Export { server_chan : chan }
    }
}

pub struct Import;

pub struct ImportTable<T> {
    slots : HashMap<u32, T>,
}

impl <T> ImportTable<T> {
    pub fn new() -> ImportTable<T> {
        ImportTable { slots : HashMap::new() }
    }
}

pub struct ExportTable<T> {
    slots : ~[T],
}

impl <T> ExportTable<T> {
    pub fn new() -> ExportTable<T> {
        ExportTable { slots : ~[] }
    }

    pub fn next(&mut self) -> u32 {
        fail!()
    }
}

pub struct RpcConnectionState {
    exports : ExportTable<Export>,
    questions : ExportTable<Question>,
    answers : ImportTable<Answer>,
    imports : ImportTable<Import>,
}

fn populate_cap_table(message : &mut OwnedSpaceMessageReader,
                      loop_chan : &std::comm::SharedChan<RpcEvent>) {
    let mut the_cap_table : ~[Option<~ClientHook>] = ~[];
    {
        let root = message.get_root::<Message::Reader>();

        match root.which() {
            Some(Message::Return(ret)) => {
                match ret.which() {
                    Some(Return::Results(payload)) => {
                        let cap_table = payload.get_cap_table();
                        for ii in range(0, cap_table.size()) {
                            match cap_table[ii].which() {
                                Some(CapDescriptor::None(())) => {
                                    the_cap_table.push(None)
                                }
                                Some(CapDescriptor::SenderHosted(id)) => {
                                    the_cap_table.push(Some(
                                            (box ImportClient {
                                                    channel : loop_chan.clone(),
                                                    import_id : id})
                                                as ~ClientHook));
                                }
                                _ => {}
                            }
                        }
                    }
                    Some(Return::Exception(e)) => {
                    }
                    _ => {}
                }

            }
            Some(Message::Call(call)) => {
                match call.get_target().which() {
                    Some(MessageTarget::ImportedCap(import_id)) => {
//                        exports.slots[import_id].server_chan.send();
                    }
                    _ => {
                    }

                }
            }
            Some(Message::Unimplemented(_)) => {
            }
            Some(Message::Abort(exc)) => {
            }
            None => {
            }
            _ => {
            }
        }
    }
    message.init_cap_table(the_cap_table);

}

impl RpcConnectionState {
    pub fn new() -> RpcConnectionState {
        RpcConnectionState {
            exports : ExportTable::new(),
            questions : ExportTable::new(),
            answers : ImportTable::new(),
            imports : ImportTable::new(),
        }
    }

    pub fn run<T : std::io::Reader + Send, U : std::io::Writer + Send>(
        self, inpipe: T, outpipe: U)
         -> std::comm::SharedChan<RpcEvent> {

        let (port, chan) = std::comm::SharedChan::<RpcEvent>::new();

        let listener_chan = chan.clone();

        spawn(proc() {
                let mut r = inpipe;
                loop {
                    let message = box serialize::new_reader(
                        &mut r,
                        DEFAULT_READER_OPTIONS);
                    listener_chan.send(IncomingMessage(message));
                }
            });

        let loop_chan = chan.clone();

        let (writer_port, writer_chan) = std::comm::Chan::<~MallocMessageBuilder>::new();

        spawn(proc() {
                let mut w = outpipe;
                loop {
                    let message = writer_port.recv();
                    serialize::write_message(&mut w, message);
                }
            });

        spawn(proc() {
                let RpcConnectionState {mut questions, mut exports, answers, imports} = self;
                loop {
                    match port.recv() {
                        IncomingMessage(mut message) => {
                            populate_cap_table(message, &loop_chan);
                            let root = message.get_root::<Message::Reader>();
                            match root.which() {
                                Some(Message::Return(ret)) => {
                                    println!("got a return with answer id {}", ret.get_answer_id());
                                    questions.slots[ret.get_answer_id()].chan.try_send(message);
                                }
                                Some(Message::Call(call)) => {
                                    println!("a call");
                                    match call.get_target().which() {
                                        Some(MessageTarget::ImportedCap(import_id)) => {
                                            println!("import id: {}", import_id);
                                            //exports.slots[import_id].server_chan.send();
                                        }
                                        _ => {
                                            fail!("call targets something else");
                                        }
                                    }
                                }
                                Some(Message::Unimplemented(_)) => {
                                    println!("unimplemented");
                                }
                                Some(Message::Abort(exc)) => {
                                    println!("abort: {}", exc.get_reason());
                                }
                                None => {
                                    println!("Nothing there");
                                }
                                _ => {
                                    println!("something else");
                                }
                            }
                        }
                        Outgoing(OutgoingMessage { message : mut m,
                                                   answer_chan,
                                                   question_chan} ) => {
                            let root = m.get_root::<Message::Builder>();
                            // add a question to the question table
                            match root.which() {
                                Some(Message::Which::Return(_)) => {}
                                Some(Message::Which::Call(call)) => {
                                    call.set_question_id(questions.slots.len() as u32);
                                    questions.slots.push(Question {is_awaiting_return : true,
                                                                   chan : answer_chan} );
                                    question_chan.try_send(call.get_question_id());
                                }
                                Some(Message::Which::Restore(res)) => {
                                    res.set_question_id(questions.slots.len() as u32);
                                    questions.slots.push(Question {is_awaiting_return : true,
                                                                   chan : answer_chan} );
                                    question_chan.try_send(res.get_question_id());
                                }
                                _ => {
                                    error!("NONE OF THOSE");
                                }
                            }

                            // send
                            writer_chan.send(m);
                        }
                        NewLocalServer(server, export_chan) => {
                            let export_id = exports.slots.len() as u32;
                            export_chan.send(export_id);
                            exports.slots.push(Export::new(server));
                        }
                    }
                }});
        return chan;
    }
}

// HACK
pub enum OwnedCapDescriptor {
    NoDescriptor,
    SenderHosted(ExportId),
    SenderPromise(ExportId),
    ReceiverHosted(ImportId),
    ReceiverAnswer(QuestionId, ~[PipelineOp::Type]),
}

pub struct ImportClient {
    priv channel : std::comm::SharedChan<RpcEvent>,
    import_id : ImportId,
}

impl ClientHook for ImportClient {
    fn copy(&self) -> ~ClientHook {
        (box ImportClient {channel : self.channel.clone(),
                           import_id : self.import_id}) as ~ClientHook
    }

    fn new_call(&self, interface_id : u64, method_id : u16,
                _size_hint : Option<common::MessageSize>)
                -> capability::Request<AnyPointer::Builder, AnyPointer::Reader, AnyPointer::Pipeline> {
        let mut message = box MallocMessageBuilder::new_default();
        {
            let root : Message::Builder = message.get_root();
            let call = root.init_call();
            call.set_interface_id(interface_id);
            call.set_method_id(method_id);
            let target = call.init_target();
            target.set_imported_cap(self.import_id);
        }
        let hook = box RpcRequest { channel : self.channel.clone(),
                                    message : message };
        Request::new(hook as ~RequestHook)
    }

    fn get_descriptor(&self) -> ~std::any::Any {
        (box ReceiverHosted(self.import_id)) as ~std::any::Any
    }
}

pub struct PipelineClient {
    priv channel : std::comm::SharedChan<RpcEvent>,
    ops : ~[PipelineOp::Type],
    question_id : QuestionId,
}

impl ClientHook for PipelineClient {
    fn copy(&self) -> ~ClientHook {
        (~PipelineClient { channel : self.channel.clone(),
                           ops : self.ops.clone(),
                           question_id : self.question_id,
            }) as ~ClientHook
    }

    fn new_call(&self, interface_id : u64, method_id : u16,
                _size_hint : Option<common::MessageSize>)
                -> capability::Request<AnyPointer::Builder, AnyPointer::Reader, AnyPointer::Pipeline> {
        let mut message = box MallocMessageBuilder::new_default();
        {
            let root : Message::Builder = message.get_root();
            let call = root.init_call();
            call.set_interface_id(interface_id);
            call.set_method_id(method_id);
            let target = call.init_target();
            let promised_answer = target.init_promised_answer();
            promised_answer.set_question_id(self.question_id);
            let transform = promised_answer.init_transform(self.ops.len());
            for ii in range(0, self.ops.len()) {
                match self.ops[ii] {
                    PipelineOp::Noop => transform[ii].set_noop(()),
                    PipelineOp::GetPointerField(idx) => transform[ii].set_get_pointer_field(idx),
                }
            }
        }
        let hook = box RpcRequest { channel : self.channel.clone(),
                                    message : message };
        Request::new(hook as ~RequestHook)
    }

    fn get_descriptor(&self) -> ~std::any::Any {
        (box ReceiverAnswer(self.question_id, self.ops.clone())) as ~std::any::Any
    }
}

pub struct RpcRequest {
    priv channel : std::comm::SharedChan<RpcEvent>,
    priv message : ~MallocMessageBuilder,
}

impl RequestHook for RpcRequest {
    fn message<'a>(&'a mut self) -> &'a mut MallocMessageBuilder {
        &mut *self.message
    }
    fn send(~self) -> RemotePromise<AnyPointer::Reader, AnyPointer::Pipeline> {

        let ~RpcRequest { channel, mut message } = self;

        {
            let cap_table = {
                let mut caps = box [];
                for cap in message.get_cap_table().iter() {
                    match cap {
                        &Some(ref client_hook) => {
                            caps.push(client_hook.get_descriptor())
                        }
                        &None => {}
                    }
                }
                caps
            };
            let root : Message::Builder = message.get_root();
            match root.which() {
                Some(Message::Which::Call(call)) => {
                    let new_cap_table = call.get_params().init_cap_table(cap_table.len());
                    for ii in range(0, cap_table.len()) {
                        match cap_table[ii].as_ref::<OwnedCapDescriptor>() {
                            None => {fail!("noncompliant client hook")}
                            Some(&NoDescriptor) => {}
                            Some(&ReceiverHosted(import_id)) => {
                                new_cap_table[ii].set_receiver_hosted(import_id);
                            }
                            Some(&ReceiverAnswer(question_id,ref ops)) => {
                                let promised_answer = new_cap_table[ii].init_receiver_answer();
                                promised_answer.set_question_id(question_id);
                                let transform = promised_answer.init_transform(ops.len());
                                for ii in range(0, ops.len()) {
                                    match ops[ii] {
                                        PipelineOp::Noop => transform[ii].set_noop(()),
                                        PipelineOp::GetPointerField(idx) => transform[ii].set_get_pointer_field(idx),
                                    }
                                }
                            }
                            Some(&SenderHosted(export_id)) => {
                                new_cap_table[ii].set_sender_hosted(export_id);
                            }
                            _ => {}
                        }
                    }
                }
                _ => {}
            }
        }

        let (event, answer_port, question_port) = RpcEvent::new_outgoing(message);
        channel.send(event);

        let question_id = question_port.recv();

        let pipeline = ~RpcPipeline {channel : channel, question_id : question_id};
        let typeless = AnyPointer::Pipeline::new(pipeline as ~PipelineHook);

        RemotePromise {answer_port : answer_port, answer_result : None,
                       pipeline : typeless  }
    }
}

pub struct RpcPipeline {
    channel : std::comm::SharedChan<RpcEvent>,
    question_id : ExportId,
}

impl PipelineHook for RpcPipeline {
    fn copy(&self) -> ~PipelineHook {
        (~RpcPipeline { channel : self.channel.clone(),
                        question_id : self.question_id }) as ~PipelineHook
    }
    fn get_pipelined_cap(&self, ops : ~[PipelineOp::Type]) -> ~ClientHook {
        (~PipelineClient { channel : self.channel.clone(),
                           ops : ops,
                           question_id : self.question_id,
        }) as ~ClientHook
    }
}

pub struct RpcCallContext {
    call_message : ~OwnedSpaceMessageReader,
}

impl CallContextHook for RpcCallContext {
    fn get_params<'a>(&'a self) -> AnyPointer::Reader<'a> {
        let root : Message::Reader = self.call_message.get_root();
        match root.which() {
            Some(Message::Call(call)) => {
                call.get_params().get_content()
            }
            _ => fail!(),
        }
    }
    fn release_params(&self) {
        fail!()
    }
    fn get_results<'a>(&'a self) -> AnyPointer::Builder<'a> {
        fail!()
    }
}

pub trait InitRequest<'a, T> {
    fn init(&'a mut self) -> T;
}

impl <'a, Params : FromStructBuilder<'a> + HasStructSize, Results, Pipeline> InitRequest<'a, Params>
for Request<Params, Results, Pipeline> {
    fn init(&'a mut self) -> Params {
        let message : Message::Builder = self.hook.message().get_root();
        match message.which() {
            Some(Message::Which::Call(call)) => {
                let params = call.init_params();
                params.get_content().init_as_struct()
            }
            _ => fail!(),
        }
    }
}

pub trait WaitForContent<'a, T> {
    fn wait(&'a mut self) -> T;
}

impl <'a, Results : FromStructReader<'a>, Pipeline> WaitForContent<'a, Results>
for RemotePromise<Results, Pipeline> {
    fn wait(&'a mut self) -> Results {
        // XXX should check that it's not already been received.
        let message = self.answer_port.recv();
        self.answer_result = Some(message);
        match self.answer_result {
            None => unreachable!(),
            Some(ref message) => {
                let root : Message::Reader = message.get_root();
                match root.which() {
                    Some(Message::Return(ret)) => {
                        match ret.which() {
                            Some(Return::Results(res)) => {
                                res.get_content().get_as_struct()
                            }
                            _ => fail!(),
                        }
                    }
                    _ => {fail!()}
                }
            }
        }
    }
}

pub struct OutgoingMessage {
    message : ~MallocMessageBuilder,
    answer_chan : std::comm::Chan<~OwnedSpaceMessageReader>,
    question_chan : std::comm::Chan<QuestionId>,
}

pub enum RpcEvent {
    IncomingMessage(~serialize::OwnedSpaceMessageReader),
    Outgoing(OutgoingMessage),
    NewLocalServer(~Server, std::comm::Chan<ExportId>),
}

impl RpcEvent {
    pub fn new_outgoing(message : ~MallocMessageBuilder)
                        -> (RpcEvent, std::comm::Port<~OwnedSpaceMessageReader>,
                            std::comm::Port<ExportId>) {
        let (answer_port, answer_chan) = std::comm::Chan::<~OwnedSpaceMessageReader>::new();

        let (question_port, question_chan) = std::comm::Chan::<ExportId>::new();

        (Outgoing(OutgoingMessage{ message : message,
                                   answer_chan : answer_chan,
                                   question_chan : question_chan }),
         answer_port,
         question_port)
    }
}
<|endoftext|>


<commit_before>use std::collections::HashMap;

use rlp::{DecoderError as RlpDecoderError, Prototype as RlpPrototype,
          RlpStream,
          UntrustedRlp,
};

use errors::prelude::*;

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub enum Node {
    Leaf(Leaf),
    Extension(Extension),
    Full(FullNode),
    Hash(Vec<u8>),
    Blank,
}

impl Node {
    const RADIX: usize = 16;
    const FULL_SIZE: usize = Node::RADIX + 1;
    const PAIR_SIZE: usize = 2;
    const HASH_SIZE: usize = 32;
    const EMPTY_SIZE: usize = 0;
    const IS_LEAF_MASK: u8 = 0x20;
    const IS_PATH_ODD_MASK: u8 = 0x10;
}

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct FullNode {
    nodes: [Option<Box<Node>>; Node::RADIX],
    value: Option<Vec<u8>>,
}

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct Leaf {
    path: Vec<u8>,
    value: Vec<u8>,
}

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct Extension {
    path: Vec<u8>,
    next: Box<Node>,
}

impl rlp::Encodable for Node {
    fn rlp_append(&self, s: &mut RlpStream) {
        match *self {
            Node::Hash(ref hash) => {
                s.append_internal(&hash.as_slice());
            }
            Node::Leaf(ref pair) => {
                s.begin_list(Node::PAIR_SIZE);
                s.append(&pair.path);
                s.append(&pair.value);
            }
            Node::Extension(ref ext) => {
                s.begin_list(Node::PAIR_SIZE);
                s.append(&ext.path);
                s.append(ext.next.as_ref());
            }
            Node::Full(ref node) => {
                s.begin_list(Node::FULL_SIZE);
                for node in &node.nodes {
                    if let Some(ref node) = *node {
                        s.append(node.as_ref());
                    } else {
                        s.append_empty_data();
                    }
                }
                if let Some(ref value) = node.value {
                    s.append(value);
                } else {
                    s.append_empty_data();
                }
            }
            Node::Blank => {
                s.append_empty_data();
            }
        }
    }
}

impl rlp::Decodable for Node {
    fn decode(rlp: &UntrustedRlp) -> Result<Self, RlpDecoderError> {
        match rlp.prototype()? {
            RlpPrototype::List(Node::PAIR_SIZE) => {
                let path: Vec<u8> = rlp.at(0)?.as_val()?;
                if path[0] & Node::IS_LEAF_MASK == Node::IS_LEAF_MASK {
                    Ok(Node::Leaf(Leaf {
                        path: rlp.at(0)?.as_val()?,
                        value: rlp.at(1)?.as_val()?,
                    }))
                } else if path[0] & Node::IS_LEAF_MASK == 0x00 {
                    Ok(Node::Extension(Extension {
                        path: rlp.at(0)?.as_val()?,
                        next: Box::new(rlp.at(1)?.as_val()?),
                    }))
                } else {
                    error!("RLP for path in Patricia Merkle Trie contains incorrect flags byte {}", path[0]);
                    Err(RlpDecoderError::Custom("Path contains incorrect flags byte"))
                }
            }
            RlpPrototype::List(Node::FULL_SIZE) => {
                let mut nodes: [Option<Box<Node>>; Node::RADIX] = [None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None];
                for i in 0..Node::RADIX {
                    let cur = rlp.at(i)?;
                    match cur.prototype()? {
                        RlpPrototype::Data(0) => {
                            continue
                        }
                        _ => {
                            nodes[i] = Some(Box::new(cur.as_val()?));
                        }
                    }
                }
                let mut value: Option<Vec<u8>> = None;
                if !rlp.at(Node::RADIX)?.is_empty() {
                    value = Some(rlp.at(Node::RADIX)?.as_val()?)
                }
                Ok(Node::Full(FullNode {
                    nodes: nodes,
                    value: value,
                }))
            }
            RlpPrototype::Data(Node::HASH_SIZE) => {
                Ok(Node::Hash(rlp.as_val()?))
            }
            RlpPrototype::Data(Node::EMPTY_SIZE) => {
                Ok(Node::Blank)
            }
            _ => {
                error!("Unexpected data while parsing Patricia Merkle Trie: {:?}: {:?}", rlp.prototype(), rlp);
                Err(RlpDecoderError::Custom("Unexpected data"))
            }
        }
    }
}

type NodeHash = sha3::digest::generic_array::GenericArray<u8, <sha3::Sha3_256 as sha3::digest::FixedOutput>::OutputSize>;
pub type TrieDB<'a> = HashMap<NodeHash, &'a Node>;

impl Node {
    pub fn get_hash<'a>(&'a self) -> NodeHash{
        use rlp::{
            encode as rlp_encode
        };
        let encoded = rlp_encode(self);
        sha3::Digest::digest(encoded.to_vec().as_slice()) as NodeHash
    }
    pub fn get_str_value<'a, 'b>(&'a self, db: &'a TrieDB, path: &'b [u8]) -> IndyResult<Option<String>> {
        let value = self.get_value(db, path)?;
        if let Some(vec) = value {
            let str = String::from_utf8(vec)
                .to_indy(IndyErrorKind::InvalidStructure, "Patricia Merkle Trie contains malformed utf8 string")?;

            trace!("Str value from Patricia Merkle Trie {}", str);
            Ok(Some(str))
        } else {
            Ok(None)
        }
    }
    fn get_value<'a, 'b>(&'a self, db: &'a TrieDB, path: &'b [u8]) -> IndyResult<Option<Vec<u8>>> {
        let nibble_path = Node::path_to_nibbles(path);
        match self._get_value(db, nibble_path.as_slice())? {
            Some(v) => {
                trace!("Raw value from Patricia Merkle Trie {:?}", v);
                let mut vec: Vec<Vec<u8>> = UntrustedRlp::new(v.as_slice()).as_list().unwrap_or_default(); //default will cause error below

                if let Some(val) = vec.pop() {
                    if vec.is_empty() {
                        return Ok(Some(val));
                    }
                }

                Err(err_msg(IndyErrorKind::InvalidStructure, "Unexpected data format of value in Patricia Merkle Trie"))
            }
            None => Ok(None)
        }
    }
    fn _get_value<'a, 'b>(&'a self, db: &'a TrieDB, path: &'b [u8]) -> IndyResult<Option<&'a Vec<u8>>> {
        trace!("Check proof, cur node: {:?}", self);
        match *self {
            Node::Full(ref node) => {
                if path.is_empty() {
                    return Ok(node.value.as_ref());
                }
                if let Some(ref next) = node.nodes[path[0] as usize] {
                    return next._get_value(db, &path[1..]);
                }
                Ok(None)
            }
            Node::Hash(ref hash) => {
                let hash = NodeHash::from_slice(hash.as_slice());
                if let Some(ref next) = db.get(hash) {
                    return next._get_value(db, path);
                } else {
                    return Err(err_msg(IndyErrorKind::InvalidStructure, "Incomplete key-value DB for Patricia Merkle Trie to get value by the key"));
                }
            }
            Node::Leaf(ref pair) => {
                let (is_leaf, pair_path) = Node::parse_path(pair.path.as_slice());

                if !is_leaf {
                    return Err(err_msg(IndyErrorKind::InvalidState, "Incorrect Patricia Merkle Trie: node marked as leaf but path contains extension flag"));
                }

                trace!("Node::_get_value in Leaf searched path {:?}, stored path {:?}", String::from_utf8(path.to_vec()), String::from_utf8(pair_path.clone()));

                if pair_path == path {
                    Ok(Some(&pair.value))
                } else {
                    Ok(None)
                }
            }
            Node::Extension(ref pair) => {
                let (is_leaf, pair_path) = Node::parse_path(pair.path.as_slice());

                if is_leaf {
                    return Err(err_msg(IndyErrorKind::InvalidState, "Incorrect Patricia Merkle Trie: node marked as extension but path contains leaf flag"));
                }

                if path.starts_with(&pair_path) {
                    pair.next._get_value(db, &path[pair_path.len()..])
                } else {
                    Ok(None)
                }
            }
            Node::Blank => {
                Ok(None)
            }
        }
    }

    fn path_to_nibbles(path: &[u8]) -> Vec<u8> {
        let mut nibble_path: Vec<u8> = Vec::new();

        for s in path {
            nibble_path.push(s >> 4);
            nibble_path.push(s & 0x0F);
        }

        nibble_path
    }

    fn parse_path(path: &[u8]) -> (bool, Vec<u8>) {
        let is_leaf: bool = path[0] & Node::IS_LEAF_MASK == Node::IS_LEAF_MASK;
        let is_odd: bool = path[0] & Node::IS_PATH_ODD_MASK == Node::IS_PATH_ODD_MASK;
        let mut nibbles: Vec<u8> = Node::path_to_nibbles(&path[1..]); //TODO avoid copy
        if is_odd {
            nibbles.insert(0, path[0] & 0x0F);
        }
        (is_leaf, nibbles)
    }
}

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

    #[test]
    fn node_deserialize_works_for_emtpy() {
        assert_eq!(UntrustedRlp::new(&base64::decode("wYA=").unwrap()).as_list::<Node>().unwrap(),
                   vec![Node::Blank]);
    }

    #[test]
    fn node_serialize_works_for_emtpy() {
        assert_eq!(base64::encode(&rlp::encode_list(&vec![Node::Blank])),
                   "wYA=");
    }
}<commit_msg>use sha3::{digest::FixedOutput}<commit_after>use std::collections::HashMap;

use rlp::{DecoderError as RlpDecoderError, Prototype as RlpPrototype,
          RlpStream,
          UntrustedRlp,
};

use errors::prelude::*;

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub enum Node {
    Leaf(Leaf),
    Extension(Extension),
    Full(FullNode),
    Hash(Vec<u8>),
    Blank,
}

impl Node {
    const RADIX: usize = 16;
    const FULL_SIZE: usize = Node::RADIX + 1;
    const PAIR_SIZE: usize = 2;
    const HASH_SIZE: usize = 32;
    const EMPTY_SIZE: usize = 0;
    const IS_LEAF_MASK: u8 = 0x20;
    const IS_PATH_ODD_MASK: u8 = 0x10;
}

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct FullNode {
    nodes: [Option<Box<Node>>; Node::RADIX],
    value: Option<Vec<u8>>,
}

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct Leaf {
    path: Vec<u8>,
    value: Vec<u8>,
}

#[derive(Debug, Serialize, Deserialize, Eq, PartialEq)]
pub struct Extension {
    path: Vec<u8>,
    next: Box<Node>,
}

impl rlp::Encodable for Node {
    fn rlp_append(&self, s: &mut RlpStream) {
        match *self {
            Node::Hash(ref hash) => {
                s.append_internal(&hash.as_slice());
            }
            Node::Leaf(ref pair) => {
                s.begin_list(Node::PAIR_SIZE);
                s.append(&pair.path);
                s.append(&pair.value);
            }
            Node::Extension(ref ext) => {
                s.begin_list(Node::PAIR_SIZE);
                s.append(&ext.path);
                s.append(ext.next.as_ref());
            }
            Node::Full(ref node) => {
                s.begin_list(Node::FULL_SIZE);
                for node in &node.nodes {
                    if let Some(ref node) = *node {
                        s.append(node.as_ref());
                    } else {
                        s.append_empty_data();
                    }
                }
                if let Some(ref value) = node.value {
                    s.append(value);
                } else {
                    s.append_empty_data();
                }
            }
            Node::Blank => {
                s.append_empty_data();
            }
        }
    }
}

impl rlp::Decodable for Node {
    fn decode(rlp: &UntrustedRlp) -> Result<Self, RlpDecoderError> {
        match rlp.prototype()? {
            RlpPrototype::List(Node::PAIR_SIZE) => {
                let path: Vec<u8> = rlp.at(0)?.as_val()?;
                if path[0] & Node::IS_LEAF_MASK == Node::IS_LEAF_MASK {
                    Ok(Node::Leaf(Leaf {
                        path: rlp.at(0)?.as_val()?,
                        value: rlp.at(1)?.as_val()?,
                    }))
                } else if path[0] & Node::IS_LEAF_MASK == 0x00 {
                    Ok(Node::Extension(Extension {
                        path: rlp.at(0)?.as_val()?,
                        next: Box::new(rlp.at(1)?.as_val()?),
                    }))
                } else {
                    error!("RLP for path in Patricia Merkle Trie contains incorrect flags byte {}", path[0]);
                    Err(RlpDecoderError::Custom("Path contains incorrect flags byte"))
                }
            }
            RlpPrototype::List(Node::FULL_SIZE) => {
                let mut nodes: [Option<Box<Node>>; Node::RADIX] = [None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None];
                for i in 0..Node::RADIX {
                    let cur = rlp.at(i)?;
                    match cur.prototype()? {
                        RlpPrototype::Data(0) => {
                            continue
                        }
                        _ => {
                            nodes[i] = Some(Box::new(cur.as_val()?));
                        }
                    }
                }
                let mut value: Option<Vec<u8>> = None;
                if !rlp.at(Node::RADIX)?.is_empty() {
                    value = Some(rlp.at(Node::RADIX)?.as_val()?)
                }
                Ok(Node::Full(FullNode {
                    nodes: nodes,
                    value: value,
                }))
            }
            RlpPrototype::Data(Node::HASH_SIZE) => {
                Ok(Node::Hash(rlp.as_val()?))
            }
            RlpPrototype::Data(Node::EMPTY_SIZE) => {
                Ok(Node::Blank)
            }
            _ => {
                error!("Unexpected data while parsing Patricia Merkle Trie: {:?}: {:?}", rlp.prototype(), rlp);
                Err(RlpDecoderError::Custom("Unexpected data"))
            }
        }
    }
}

type NodeHash = sha3::digest::generic_array::GenericArray<u8, <sha3::Sha3_256 as sha3::digest::FixedOutput>::OutputSize>;
pub type TrieDB<'a> = HashMap<NodeHash, &'a Node>;

impl Node {
    pub fn get_hash(&self) -> NodeHash{
        use rlp::{
            encode as rlp_encode
        };
        use sha3::{digest::FixedOutput};
        use sha3::Digest;
        let encoded = rlp_encode(self);
        let mut hasher = sha3::Sha3_256::default();
        hasher.input(encoded.to_vec().as_slice());
        hasher.fixed_result()
    }
    pub fn get_str_value<'a, 'b>(&'a self, db: &'a TrieDB, path: &'b [u8]) -> IndyResult<Option<String>> {
        let value = self.get_value(db, path)?;
        if let Some(vec) = value {
            let str = String::from_utf8(vec)
                .to_indy(IndyErrorKind::InvalidStructure, "Patricia Merkle Trie contains malformed utf8 string")?;

            trace!("Str value from Patricia Merkle Trie {}", str);
            Ok(Some(str))
        } else {
            Ok(None)
        }
    }
    fn get_value<'a, 'b>(&'a self, db: &'a TrieDB, path: &'b [u8]) -> IndyResult<Option<Vec<u8>>> {
        let nibble_path = Node::path_to_nibbles(path);
        match self._get_value(db, nibble_path.as_slice())? {
            Some(v) => {
                trace!("Raw value from Patricia Merkle Trie {:?}", v);
                let mut vec: Vec<Vec<u8>> = UntrustedRlp::new(v.as_slice()).as_list().unwrap_or_default(); //default will cause error below

                if let Some(val) = vec.pop() {
                    if vec.is_empty() {
                        return Ok(Some(val));
                    }
                }

                Err(err_msg(IndyErrorKind::InvalidStructure, "Unexpected data format of value in Patricia Merkle Trie"))
            }
            None => Ok(None)
        }
    }
    fn _get_value<'a, 'b>(&'a self, db: &'a TrieDB, path: &'b [u8]) -> IndyResult<Option<&'a Vec<u8>>> {
        trace!("Check proof, cur node: {:?}", self);
        match *self {
            Node::Full(ref node) => {
                if path.is_empty() {
                    return Ok(node.value.as_ref());
                }
                if let Some(ref next) = node.nodes[path[0] as usize] {
                    return next._get_value(db, &path[1..]);
                }
                Ok(None)
            }
            Node::Hash(ref hash) => {
                let hash = NodeHash::from_slice(hash.as_slice());
                if let Some(ref next) = db.get(hash) {
                    return next._get_value(db, path);
                } else {
                    return Err(err_msg(IndyErrorKind::InvalidStructure, "Incomplete key-value DB for Patricia Merkle Trie to get value by the key"));
                }
            }
            Node::Leaf(ref pair) => {
                let (is_leaf, pair_path) = Node::parse_path(pair.path.as_slice());

                if !is_leaf {
                    return Err(err_msg(IndyErrorKind::InvalidState, "Incorrect Patricia Merkle Trie: node marked as leaf but path contains extension flag"));
                }

                trace!("Node::_get_value in Leaf searched path {:?}, stored path {:?}", String::from_utf8(path.to_vec()), String::from_utf8(pair_path.clone()));

                if pair_path == path {
                    Ok(Some(&pair.value))
                } else {
                    Ok(None)
                }
            }
            Node::Extension(ref pair) => {
                let (is_leaf, pair_path) = Node::parse_path(pair.path.as_slice());

                if is_leaf {
                    return Err(err_msg(IndyErrorKind::InvalidState, "Incorrect Patricia Merkle Trie: node marked as extension but path contains leaf flag"));
                }

                if path.starts_with(&pair_path) {
                    pair.next._get_value(db, &path[pair_path.len()..])
                } else {
                    Ok(None)
                }
            }
            Node::Blank => {
                Ok(None)
            }
        }
    }

    fn path_to_nibbles(path: &[u8]) -> Vec<u8> {
        let mut nibble_path: Vec<u8> = Vec::new();

        for s in path {
            nibble_path.push(s >> 4);
            nibble_path.push(s & 0x0F);
        }

        nibble_path
    }

    fn parse_path(path: &[u8]) -> (bool, Vec<u8>) {
        let is_leaf: bool = path[0] & Node::IS_LEAF_MASK == Node::IS_LEAF_MASK;
        let is_odd: bool = path[0] & Node::IS_PATH_ODD_MASK == Node::IS_PATH_ODD_MASK;
        let mut nibbles: Vec<u8> = Node::path_to_nibbles(&path[1..]); //TODO avoid copy
        if is_odd {
            nibbles.insert(0, path[0] & 0x0F);
        }
        (is_leaf, nibbles)
    }
}

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

    #[test]
    fn node_deserialize_works_for_emtpy() {
        assert_eq!(UntrustedRlp::new(&base64::decode("wYA=").unwrap()).as_list::<Node>().unwrap(),
                   vec![Node::Blank]);
    }

    #[test]
    fn node_serialize_works_for_emtpy() {
        assert_eq!(base64::encode(&rlp::encode_list(&vec![Node::Blank])),
                   "wYA=");
    }
}<|endoftext|>


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

use futures::{future, Future};
use reqwest;
use serde_derive::Serialize;
use tokio;
use log::{error, info};

use crate::util;
use crate::worker;

#[derive(Serialize, Clone, PartialEq, Eq, Debug)]
pub struct SlackAttachment {
    pub text: String,
    pub title: Option<String>,
    pub title_link: Option<String>,
    pub color: Option<String>,
}

impl SlackAttachment {
    pub fn new(text: &str) -> SlackAttachment {
        SlackAttachment {
            text: text.to_string(),
            title: None,
            title_link: None,
            color: None,
        }
    }
}

pub struct SlackAttachmentBuilder {
    attachment: SlackAttachment,
}

impl SlackAttachmentBuilder {
    pub fn new(text: &str) -> SlackAttachmentBuilder {
        SlackAttachmentBuilder {
            attachment: SlackAttachment::new(text),
        }
    }

    pub fn text<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.text = value.into();
        self
    }

    pub fn title<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.title = Some(value.into());
        self
    }
    pub fn title_link<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.title_link = Some(value.into());
        self
    }

    pub fn color<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.color = Some(value.into());
        self
    }

    pub fn build(&self) -> SlackAttachment {
        self.attachment.clone()
    }
}

#[derive(Serialize, Clone, PartialEq)]
struct SlackMessage {
    text: String,
    attachments: Vec<SlackAttachment>,
    channel: String,
}

// the main object for sending messages to slack
struct Slack {
    client: reqwest::r#async::Client,
    webhook_url: String,
    recent_messages: Mutex<Vec<SlackMessage>>,
}

const TRIM_MESSAGES_AT: usize = 200;
const TRIM_MESSAGES_TO: usize = 20;

impl Slack {
    pub fn new(webhook_url: Option<String>) -> Slack {
        Slack {
            client: reqwest::r#async::Client::new(),
            webhook_url: webhook_url.unwrap_or(String::new()),
            recent_messages: Mutex::new(Vec::new()),
        }
    }

    fn send(&self, channel: &str, msg: &str, attachments: Vec<SlackAttachment>) {
        if self.webhook_url.is_empty() {
            return
        }

        let slack_msg = SlackMessage {
            text: msg.to_string(),
            attachments: attachments,
            channel: channel.to_string(),
        };

        if !self.is_unique(&slack_msg) {
            info!("Skipping duplicate message to {}", channel);
            return;
        }

        info!("Sending message to #{}", channel);
        tokio::spawn(self.client.post(&self.webhook_url).json(&slack_msg).send().then(|res| {
            match res {
                Ok(_) => info!("Successfully sent slack message"),
                Err(e) => error!("Error sending slack message: {}", e),
            };
            future::ok::<(), ()>(())
        }));
    }

    fn is_unique(&self, req: &SlackMessage) -> bool {
        let mut recent_messages = self.recent_messages.lock().unwrap();
        util::check_unique_event(req.clone(), &mut *recent_messages, TRIM_MESSAGES_AT, TRIM_MESSAGES_TO)
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct SlackRequest {
    pub channel: String,
    pub msg: String,
    pub attachments: Vec<SlackAttachment>,
}

struct Runner {
    slack: Arc<Slack>,
}

pub fn req(channel: &str, msg: &str, attachments: Vec<SlackAttachment>) -> SlackRequest {
    SlackRequest {
        channel: channel.into(),
        msg: msg.into(),
        attachments: attachments,
    }
}

pub fn new_runner(webhook_url: Option<String>) -> Arc<dyn worker::Runner<SlackRequest>> {
    Arc::new(Runner {
        slack: Arc::new(Slack::new(webhook_url)),
    })
}

impl worker::Runner<SlackRequest> for Runner {
    fn handle(&self, req: SlackRequest) {
        self.slack.send(&req.channel, &req.msg, req.attachments);
    }
}
<commit_msg>Don't log the slack webhook URL (#289)<commit_after>use std::sync::{Arc, Mutex};

use futures::{future, Future};
use reqwest;
use serde_derive::Serialize;
use tokio;
use log::{error, info};

use crate::util;
use crate::worker;

#[derive(Serialize, Clone, PartialEq, Eq, Debug)]
pub struct SlackAttachment {
    pub text: String,
    pub title: Option<String>,
    pub title_link: Option<String>,
    pub color: Option<String>,
}

impl SlackAttachment {
    pub fn new(text: &str) -> SlackAttachment {
        SlackAttachment {
            text: text.to_string(),
            title: None,
            title_link: None,
            color: None,
        }
    }
}

pub struct SlackAttachmentBuilder {
    attachment: SlackAttachment,
}

impl SlackAttachmentBuilder {
    pub fn new(text: &str) -> SlackAttachmentBuilder {
        SlackAttachmentBuilder {
            attachment: SlackAttachment::new(text),
        }
    }

    pub fn text<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.text = value.into();
        self
    }

    pub fn title<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.title = Some(value.into());
        self
    }
    pub fn title_link<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.title_link = Some(value.into());
        self
    }

    pub fn color<S: Into<String>>(&mut self, value: S) -> &mut SlackAttachmentBuilder {
        self.attachment.color = Some(value.into());
        self
    }

    pub fn build(&self) -> SlackAttachment {
        self.attachment.clone()
    }
}

#[derive(Serialize, Clone, PartialEq)]
struct SlackMessage {
    text: String,
    attachments: Vec<SlackAttachment>,
    channel: String,
}

// the main object for sending messages to slack
struct Slack {
    client: reqwest::r#async::Client,
    webhook_url: String,
    recent_messages: Mutex<Vec<SlackMessage>>,
}

const TRIM_MESSAGES_AT: usize = 200;
const TRIM_MESSAGES_TO: usize = 20;

impl Slack {
    pub fn new(webhook_url: Option<String>) -> Slack {
        Slack {
            client: reqwest::r#async::Client::new(),
            webhook_url: webhook_url.unwrap_or(String::new()),
            recent_messages: Mutex::new(Vec::new()),
        }
    }

    fn send(&self, channel: &str, msg: &str, attachments: Vec<SlackAttachment>) {
        if self.webhook_url.is_empty() {
            return
        }

        let slack_msg = SlackMessage {
            text: msg.to_string(),
            attachments: attachments,
            channel: channel.to_string(),
        };

        if !self.is_unique(&slack_msg) {
            info!("Skipping duplicate message to {}", channel);
            return;
        }

        info!("Sending message to #{}", channel);
        let webhook_url = self.webhook_url.clone();
        tokio::spawn(self.client.post(&webhook_url).json(&slack_msg).send().then(move |res| {
            match res {
                Ok(_) => info!("Successfully sent slack message"),
                Err(e) => {
                    let msg = format!("{}", e).replace(&webhook_url, "<webhook url>");
                    error!("Error sending slack message: {}", msg);
                }
            };
            future::ok::<(), ()>(())
        }));
    }

    fn is_unique(&self, req: &SlackMessage) -> bool {
        let mut recent_messages = self.recent_messages.lock().unwrap();
        util::check_unique_event(req.clone(), &mut *recent_messages, TRIM_MESSAGES_AT, TRIM_MESSAGES_TO)
    }
}

#[derive(Debug, PartialEq, Clone)]
pub struct SlackRequest {
    pub channel: String,
    pub msg: String,
    pub attachments: Vec<SlackAttachment>,
}

struct Runner {
    slack: Arc<Slack>,
}

pub fn req(channel: &str, msg: &str, attachments: Vec<SlackAttachment>) -> SlackRequest {
    SlackRequest {
        channel: channel.into(),
        msg: msg.into(),
        attachments: attachments,
    }
}

pub fn new_runner(webhook_url: Option<String>) -> Arc<dyn worker::Runner<SlackRequest>> {
    Arc::new(Runner {
        slack: Arc::new(Slack::new(webhook_url)),
    })
}

impl worker::Runner<SlackRequest> for Runner {
    fn handle(&self, req: SlackRequest) {
        self.slack.send(&req.channel, &req.msg, req.attachments);
    }
}
<|endoftext|>


<commit_before>/*
 * Copyright 2019 Miklos Vajna. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#![warn(missing_docs)]

//! Takes an OSM way ID and turns it into a string that is readable and
//! e.g. OsmAnd can parse it as well.

extern crate atty;
extern crate serde_json;
extern crate url;

use std::io::Write;

/// A Result which allows any error that implements Error.
pub type BoxResult<T> = Result<T, Box<dyn std::error::Error>>;

#[derive(Debug)]
struct OsmifyError {
    details: String,
}

impl std::fmt::Display for OsmifyError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.details)
    }
}

impl std::error::Error for OsmifyError {
    fn description(&self) -> &str {
        &self.details
    }
}

/// Allows HTTP GET and POST requests.
pub trait Urllib: Send + Sync {
    /// If `data` is empty, that means HTTP GET, otherwise a HTTP POST.
    fn urlopen(&self, url: &str, data: &str) -> BoxResult<String>;
}

fn query_turbo(urllib: &dyn Urllib, query: &str) -> BoxResult<String> {
    let url = "http://overpass-api.de/api/interpreter";

    let buf = urllib.urlopen(url, query)?;
    Ok(buf)
}

fn query_nominatim(urllib: &dyn Urllib, query: &str) -> BoxResult<String> {
    let prefix = "http://nominatim.openstreetmap.org/search.php?";
    let encoded: String = url::form_urlencoded::Serializer::new(String::new())
        .append_pair("q", query)
        .append_pair("format", "json")
        .finish();
    let url = format!("{}{}", prefix, encoded);

    let buf = urllib.urlopen(url.as_str(), "")?;

    Ok(buf)
}

fn osmify(query: &str, urllib: &dyn Urllib) -> BoxResult<String> {
    let nominatim = query_nominatim(urllib, query)?;
    let json: serde_json::Value = match serde_json::from_str(&nominatim) {
        Ok(value) => value,
        Err(error) => {
            return Err(Box::new(OsmifyError {
                details: format!("failed to parse JSON from nominatim: {}", error.to_string()),
            }));
        }
    };
    let mut elements = json.as_array().ok_or("option::NoneError")?.clone();
    if elements.is_empty() {
        return Err(Box::new(OsmifyError {
            details: "no results from nominatim".to_string(),
        }));
    }

    if elements.len() > 1 {
        // There are multiple elements, prefer buildings if possible.
        let buildings: Vec<serde_json::Value> = elements
            .iter()
            .filter(|i| {
                let i = match i.as_object() {
                    Some(value) => value,
                    None => return false,
                };
                let class = match i.get("class") {
                    Some(value) => value.as_str(),
                    None => return false,
                };
                match class {
                    Some(value) => value == "building",
                    None => false,
                }
            })
            .cloned()
            .collect();

        if !buildings.is_empty() {
            elements = buildings;
        }
    }

    let element = elements[0].as_object().ok_or("option::NoneError")?;
    let lat = element["lat"]
        .as_str()
        .ok_or("no lat in json from nominatim")?;
    let lon = element["lon"]
        .as_str()
        .ok_or("no lon in json from nominatim")?;
    let object_type = element["osm_type"]
        .as_str()
        .ok_or("no osm_type in json from nominatim")?;
    let object_id = element["osm_id"]
        .as_u64()
        .ok_or("no osm_id in json from nominatim")?;

    // Use overpass to get the properties of the object.
    let overpass_query = format!(
        r#"[out:json];
(
    {}({});
);
out body;"#,
        object_type, object_id
    );
    let turbo = query_turbo(urllib, &overpass_query)?;
    let json: serde_json::Value = match serde_json::from_str(&turbo) {
        Ok(value) => value,
        Err(error) => {
            return Err(Box::new(OsmifyError {
                details: format!("Failed to parse JSON from overpass: {:?}", error),
            }));
        }
    };
    let json = json.as_object().ok_or("option::NoneError")?;
    let elements = &json["elements"].as_array().ok_or("option::NoneError")?;
    if elements.is_empty() {
        return Err(Box::new(OsmifyError {
            details: "No results from overpass".to_string(),
        }));
    }

    let element = &elements[0];
    let tags = element["tags"].as_object().ok_or("option::NoneError")?;
    let city = tags["addr:city"].as_str().ok_or("option::NoneError")?;
    let housenumber = tags["addr:housenumber"]
        .as_str()
        .ok_or("option::NoneError")?;
    let postcode = tags["addr:postcode"].as_str().ok_or("option::NoneError")?;
    let street = tags["addr:street"].as_str().ok_or("option::NoneError")?;
    let addr = format!("{} {}, {} {}", postcode, city, street, housenumber);

    // Print the result.
    Ok(format!("geo:{},{} ({})", lat, lon, addr))
}

fn spinner(
    rx: &std::sync::mpsc::Receiver<Result<String, String>>,
    stream: &mut dyn Write,
) -> BoxResult<()> {
    let spin_characters = vec!['\\', '|', '/', '-'];
    let mut spin_index = 0;
    loop {
        match rx.try_recv() {
            Ok(result) => {
                if atty::is(atty::Stream::Stdout) {
                    print!("\r");
                }
                std::io::stdout().flush()?;
                let result = result?;
                writeln!(stream, "{}", result)?;
                return Ok(());
            }
            Err(_) => {
                if atty::is(atty::Stream::Stdout) {
                    print!("\r [{}] ", spin_characters[spin_index]);
                }
                std::io::stdout().flush()?;
                spin_index = (spin_index + 1) % spin_characters.len();
                std::thread::sleep(std::time::Duration::from_millis(100));
            }
        }
    }
}

/// Similar to plain main(), but with an interface that allows testing.
pub fn main(args: Vec<String>, stream: &mut dyn Write, urllib: Box<dyn Urllib>) -> BoxResult<()> {
    if args.len() > 1 {
        let (tx, rx) = std::sync::mpsc::channel();
        std::thread::spawn(move || {
            let result = osmify(&args[1], &*urllib);
            match result {
                Ok(value) => tx.send(Ok(value)),
                Err(error) => tx.send(Err(format!("failed to osmify: {}", error.to_string()))),
            }
        });
        spinner(&rx, stream)?;
    } else {
        writeln!(stream, "usage: addr-osmify <query>")?;
        writeln!(stream)?;
        writeln!(stream, "e.g. addr-osmify 'Mészáros utca 58/a, Budapest'")?;
    }

    Ok(())
}

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

    struct URLRoute {
        url: String,
        data_path: String,
        result_path: String,
    }

    struct MockUrllib {
        routes: Vec<URLRoute>,
    }

    impl Urllib for MockUrllib {
        fn urlopen(&self, url: &str, data: &str) -> BoxResult<String> {
            for route in self.routes.iter() {
                if route.url != url {
                    continue;
                }

                if !route.data_path.is_empty() {
                    // Can't use assert_eq!() here, failure would result in a hang.
                    let contents = match std::fs::read_to_string(&route.data_path) {
                        Ok(value) => value,
                        Err(error) => {
                            return Err(Box::new(OsmifyError {
                                details: format!(
                                    "failed read {}: {}",
                                    route.data_path,
                                    error.to_string()
                                ),
                            }))
                        }
                    };
                    if data != contents {
                        return Err(Box::new(OsmifyError {
                            details: format!("unexpected data: '{}' != '{}'", data, contents),
                        }));
                    }
                }

                let contents = match std::fs::read_to_string(&route.result_path) {
                    Ok(value) => value,
                    Err(error) => {
                        return Err(Box::new(OsmifyError {
                            details: format!(
                                "failed read {}: {}",
                                route.result_path,
                                error.to_string()
                            ),
                        }))
                    }
                };
                return Ok(contents);
            }

            return Err(Box::new(OsmifyError {
                details: format!("unexpected url: {}", url),
            }));
        }
    }

    #[test]
    fn test_happy() -> BoxResult<()> {
        let args: Vec<String> = vec!["".to_string(), "Mészáros utca 58/a, Budapest".to_string()];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=M%C3%A9sz%C3%A1ros+utca+58%2Fa%2C+Budapest&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-happy.json".to_string()
            },
            URLRoute {
                url: "http://overpass-api.de/api/interpreter".to_string(),
                data_path: "mock/overpass-happy.expected-data".to_string(),
                result_path: "mock/overpass-happy.json".to_string()
            }
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        main(args, &mut buf, urllib)?;

        let buf_vec = buf.into_inner();
        let buf_string = match std::str::from_utf8(&buf_vec) {
            Ok(v) => v,
            Err(e) => panic!("invalid UTF-8 sequence: {}", e),
        };
        assert_eq!(
            buf_string,
            "geo:47.490592,19.030662 (1016 Budapest, Mészáros utca 58/a)\n"
        );

        Ok(())
    }

    #[test]
    fn test_nominatim_json() -> BoxResult<()> {
        let args: Vec<String> = vec!["".to_string(), "Mészáros utca 58/a, Budapest".to_string()];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=M%C3%A9sz%C3%A1ros+utca+58%2Fa%2C+Budapest&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-bad.json".to_string()
            },
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        let error = match main(args, &mut buf, urllib) {
            Ok(_) => panic!("unexpected success"),
            Err(e) => e,
        };

        assert_eq!(
            error.to_string(),
            "failed to osmify: failed to parse JSON from nominatim: EOF while parsing an object at line 2 column 0",
        );

        Ok(())
    }

    #[test]
    fn test_nominatim_no_result() -> BoxResult<()> {
        let args: Vec<String> = vec!["".to_string(), "Mészáros utca 58/a, Budapestt".to_string()];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=M%C3%A9sz%C3%A1ros+utca+58%2Fa%2C+Budapestt&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-no-result.json".to_string()
            },
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        let error = match main(args, &mut buf, urllib) {
            Ok(_) => panic!("unexpected success"),
            Err(e) => e,
        };

        assert_eq!(
            error.to_string(),
            "failed to osmify: no results from nominatim",
        );

        Ok(())
    }

    #[test]
    fn test_prefer_buildings() -> BoxResult<()> {
        let args: Vec<String> = vec![
            "".to_string(),
            "Karinthy Frigyes út 18, Budapest".to_string(),
        ];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=Karinthy+Frigyes+%C3%BAt+18%2C+Budapest&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-prefer-buildings.json".to_string()
            },
            URLRoute {
                url: "http://overpass-api.de/api/interpreter".to_string(),
                data_path: "mock/overpass-prefer-buildings.expected-data".to_string(),
                result_path: "mock/overpass-prefer-buildings.json".to_string()
            }
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        main(args, &mut buf, urllib)?;

        let buf_vec = buf.into_inner();
        let buf_string = match std::str::from_utf8(&buf_vec) {
            Ok(v) => v,
            Err(e) => panic!("invalid UTF-8 sequence: {}", e),
        };
        assert_eq!(
            buf_string,
            "geo:47.47690895,19.0512550758533 (1111 Budapest, Karinthy Frigyes út 18)\n"
        );

        Ok(())
    }
}
<commit_msg>addr-osmify-rust: add building tests<commit_after>/*
 * Copyright 2019 Miklos Vajna. All rights reserved.
 * Use of this source code is governed by a BSD-style license that can be
 * found in the LICENSE file.
 */

#![warn(missing_docs)]

//! Takes an OSM way ID and turns it into a string that is readable and
//! e.g. OsmAnd can parse it as well.

extern crate atty;
extern crate serde_json;
extern crate url;

use std::io::Write;

/// A Result which allows any error that implements Error.
pub type BoxResult<T> = Result<T, Box<dyn std::error::Error>>;

#[derive(Debug)]
struct OsmifyError {
    details: String,
}

impl std::fmt::Display for OsmifyError {
    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
        write!(f, "{}", self.details)
    }
}

impl std::error::Error for OsmifyError {
    fn description(&self) -> &str {
        &self.details
    }
}

/// Allows HTTP GET and POST requests.
pub trait Urllib: Send + Sync {
    /// If `data` is empty, that means HTTP GET, otherwise a HTTP POST.
    fn urlopen(&self, url: &str, data: &str) -> BoxResult<String>;
}

fn query_turbo(urllib: &dyn Urllib, query: &str) -> BoxResult<String> {
    let url = "http://overpass-api.de/api/interpreter";

    let buf = urllib.urlopen(url, query)?;
    Ok(buf)
}

fn query_nominatim(urllib: &dyn Urllib, query: &str) -> BoxResult<String> {
    let prefix = "http://nominatim.openstreetmap.org/search.php?";
    let encoded: String = url::form_urlencoded::Serializer::new(String::new())
        .append_pair("q", query)
        .append_pair("format", "json")
        .finish();
    let url = format!("{}{}", prefix, encoded);

    let buf = urllib.urlopen(url.as_str(), "")?;

    Ok(buf)
}

fn is_building(element: &serde_json::Value) -> bool {
    if !element.is_object() {
        return false;
    }
    let class = element.get("class");
    if class.is_none() {
        return false;
    }
    let class = class.unwrap();
    if !class.is_string() {
        return false;
    }
    class.as_str().unwrap() == "building"
}

fn osmify(query: &str, urllib: &dyn Urllib) -> BoxResult<String> {
    let nominatim = query_nominatim(urllib, query)?;
    let json: serde_json::Value = match serde_json::from_str(&nominatim) {
        Ok(value) => value,
        Err(error) => {
            return Err(Box::new(OsmifyError {
                details: format!("failed to parse JSON from nominatim: {}", error.to_string()),
            }));
        }
    };
    let mut elements = json.as_array().ok_or("option::NoneError")?.clone();
    if elements.is_empty() {
        return Err(Box::new(OsmifyError {
            details: "no results from nominatim".to_string(),
        }));
    }

    if elements.len() > 1 {
        // There are multiple elements, prefer buildings if possible.
        let buildings: Vec<serde_json::Value> = elements
            .iter()
            .filter(|i| is_building(i))
            .cloned()
            .collect();

        if !buildings.is_empty() {
            elements = buildings;
        }
    }

    let element = elements[0].as_object().ok_or("option::NoneError")?;
    let lat = element["lat"]
        .as_str()
        .ok_or("no lat in json from nominatim")?;
    let lon = element["lon"]
        .as_str()
        .ok_or("no lon in json from nominatim")?;
    let object_type = element["osm_type"]
        .as_str()
        .ok_or("no osm_type in json from nominatim")?;
    let object_id = element["osm_id"]
        .as_u64()
        .ok_or("no osm_id in json from nominatim")?;

    // Use overpass to get the properties of the object.
    let overpass_query = format!(
        r#"[out:json];
(
    {}({});
);
out body;"#,
        object_type, object_id
    );
    let turbo = query_turbo(urllib, &overpass_query)?;
    let json: serde_json::Value = match serde_json::from_str(&turbo) {
        Ok(value) => value,
        Err(error) => {
            return Err(Box::new(OsmifyError {
                details: format!("Failed to parse JSON from overpass: {:?}", error),
            }));
        }
    };
    let json = json.as_object().ok_or("option::NoneError")?;
    let elements = &json["elements"].as_array().ok_or("option::NoneError")?;
    if elements.is_empty() {
        return Err(Box::new(OsmifyError {
            details: "No results from overpass".to_string(),
        }));
    }

    let element = &elements[0];
    let tags = element["tags"].as_object().ok_or("option::NoneError")?;
    let city = tags["addr:city"].as_str().ok_or("option::NoneError")?;
    let housenumber = tags["addr:housenumber"]
        .as_str()
        .ok_or("option::NoneError")?;
    let postcode = tags["addr:postcode"].as_str().ok_or("option::NoneError")?;
    let street = tags["addr:street"].as_str().ok_or("option::NoneError")?;
    let addr = format!("{} {}, {} {}", postcode, city, street, housenumber);

    // Print the result.
    Ok(format!("geo:{},{} ({})", lat, lon, addr))
}

fn spinner(
    rx: &std::sync::mpsc::Receiver<Result<String, String>>,
    stream: &mut dyn Write,
) -> BoxResult<()> {
    let spin_characters = vec!['\\', '|', '/', '-'];
    let mut spin_index = 0;
    loop {
        match rx.try_recv() {
            Ok(result) => {
                if atty::is(atty::Stream::Stdout) {
                    print!("\r");
                }
                std::io::stdout().flush()?;
                let result = result?;
                writeln!(stream, "{}", result)?;
                return Ok(());
            }
            Err(_) => {
                if atty::is(atty::Stream::Stdout) {
                    print!("\r [{}] ", spin_characters[spin_index]);
                }
                std::io::stdout().flush()?;
                spin_index = (spin_index + 1) % spin_characters.len();
                std::thread::sleep(std::time::Duration::from_millis(100));
            }
        }
    }
}

/// Similar to plain main(), but with an interface that allows testing.
pub fn main(args: Vec<String>, stream: &mut dyn Write, urllib: Box<dyn Urllib>) -> BoxResult<()> {
    if args.len() > 1 {
        let (tx, rx) = std::sync::mpsc::channel();
        std::thread::spawn(move || {
            let result = osmify(&args[1], &*urllib);
            match result {
                Ok(value) => tx.send(Ok(value)),
                Err(error) => tx.send(Err(format!("failed to osmify: {}", error.to_string()))),
            }
        });
        spinner(&rx, stream)?;
    } else {
        writeln!(stream, "usage: addr-osmify <query>")?;
        writeln!(stream)?;
        writeln!(stream, "e.g. addr-osmify 'Mészáros utca 58/a, Budapest'")?;
    }

    Ok(())
}

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

    struct URLRoute {
        url: String,
        data_path: String,
        result_path: String,
    }

    struct MockUrllib {
        routes: Vec<URLRoute>,
    }

    impl Urllib for MockUrllib {
        fn urlopen(&self, url: &str, data: &str) -> BoxResult<String> {
            for route in self.routes.iter() {
                if route.url != url {
                    continue;
                }

                if !route.data_path.is_empty() {
                    // Can't use assert_eq!() here, failure would result in a hang.
                    let contents = match std::fs::read_to_string(&route.data_path) {
                        Ok(value) => value,
                        Err(error) => {
                            return Err(Box::new(OsmifyError {
                                details: format!(
                                    "failed read {}: {}",
                                    route.data_path,
                                    error.to_string()
                                ),
                            }))
                        }
                    };
                    if data != contents {
                        return Err(Box::new(OsmifyError {
                            details: format!("unexpected data: '{}' != '{}'", data, contents),
                        }));
                    }
                }

                let contents = match std::fs::read_to_string(&route.result_path) {
                    Ok(value) => value,
                    Err(error) => {
                        return Err(Box::new(OsmifyError {
                            details: format!(
                                "failed read {}: {}",
                                route.result_path,
                                error.to_string()
                            ),
                        }))
                    }
                };
                return Ok(contents);
            }

            return Err(Box::new(OsmifyError {
                details: format!("unexpected url: {}", url),
            }));
        }
    }

    #[test]
    fn test_happy() -> BoxResult<()> {
        let args: Vec<String> = vec!["".to_string(), "Mészáros utca 58/a, Budapest".to_string()];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=M%C3%A9sz%C3%A1ros+utca+58%2Fa%2C+Budapest&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-happy.json".to_string()
            },
            URLRoute {
                url: "http://overpass-api.de/api/interpreter".to_string(),
                data_path: "mock/overpass-happy.expected-data".to_string(),
                result_path: "mock/overpass-happy.json".to_string()
            }
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        main(args, &mut buf, urllib)?;

        let buf_vec = buf.into_inner();
        let buf_string = match std::str::from_utf8(&buf_vec) {
            Ok(v) => v,
            Err(e) => panic!("invalid UTF-8 sequence: {}", e),
        };
        assert_eq!(
            buf_string,
            "geo:47.490592,19.030662 (1016 Budapest, Mészáros utca 58/a)\n"
        );

        Ok(())
    }

    #[test]
    fn test_nominatim_json() -> BoxResult<()> {
        let args: Vec<String> = vec!["".to_string(), "Mészáros utca 58/a, Budapest".to_string()];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=M%C3%A9sz%C3%A1ros+utca+58%2Fa%2C+Budapest&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-bad.json".to_string()
            },
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        let error = match main(args, &mut buf, urllib) {
            Ok(_) => panic!("unexpected success"),
            Err(e) => e,
        };

        assert_eq!(
            error.to_string(),
            "failed to osmify: failed to parse JSON from nominatim: EOF while parsing an object at line 2 column 0",
        );

        Ok(())
    }

    #[test]
    fn test_nominatim_no_result() -> BoxResult<()> {
        let args: Vec<String> = vec!["".to_string(), "Mészáros utca 58/a, Budapestt".to_string()];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=M%C3%A9sz%C3%A1ros+utca+58%2Fa%2C+Budapestt&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-no-result.json".to_string()
            },
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        let error = match main(args, &mut buf, urllib) {
            Ok(_) => panic!("unexpected success"),
            Err(e) => e,
        };

        assert_eq!(
            error.to_string(),
            "failed to osmify: no results from nominatim",
        );

        Ok(())
    }

    #[test]
    fn test_prefer_buildings() -> BoxResult<()> {
        let args: Vec<String> = vec![
            "".to_string(),
            "Karinthy Frigyes út 18, Budapest".to_string(),
        ];

        let mut buf: std::io::Cursor<Vec<u8>> = std::io::Cursor::new(Vec::new());

        let routes = vec![
            URLRoute {
                url: "http://nominatim.openstreetmap.org/search.php?q=Karinthy+Frigyes+%C3%BAt+18%2C+Budapest&format=json".to_string(),
                data_path: "".to_string(),
                result_path: "mock/nominatim-prefer-buildings.json".to_string()
            },
            URLRoute {
                url: "http://overpass-api.de/api/interpreter".to_string(),
                data_path: "mock/overpass-prefer-buildings.expected-data".to_string(),
                result_path: "mock/overpass-prefer-buildings.json".to_string()
            }
        ];
        let urllib: Box<dyn Urllib> = Box::new(MockUrllib { routes: routes });
        main(args, &mut buf, urllib)?;

        let buf_vec = buf.into_inner();
        let buf_string = match std::str::from_utf8(&buf_vec) {
            Ok(v) => v,
            Err(e) => panic!("invalid UTF-8 sequence: {}", e),
        };
        assert_eq!(
            buf_string,
            "geo:47.47690895,19.0512550758533 (1111 Budapest, Karinthy Frigyes út 18)\n"
        );

        Ok(())
    }

    #[test]
    fn test_is_building() {
        {
            let json_null = serde_json::Value::Null;
            assert!(!is_building(&json_null));
        }

        {
            let map = serde_json::Map::new();
            let json_object = serde_json::Value::Object(map);
            assert!(!is_building(&json_object));
        }

        {
            let mut map = serde_json::Map::new();
            let json_null = serde_json::Value::Null;
            map.insert("class".to_string(), json_null);
            let json_object = serde_json::Value::Object(map);
            assert!(!is_building(&json_object));
        }
    }
}
<|endoftext|>