Refactor per-node-scan-start into a function
[dnsseed-rust] / src / peer.rs
1 use std::cmp;
2 use std::net::SocketAddr;
3 use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
4
5 use bitcoin::consensus::encode;
6 use bitcoin::consensus::encode::{Decodable, Encodable};
7 use bitcoin::network::address::Address;
8 use bitcoin::network::constants::Network;
9 use bitcoin::network::message::{RawNetworkMessage, NetworkMessage};
10 use bitcoin::network::message_network::VersionMessage;
11
12 use tokio::prelude::*;
13 use tokio::codec;
14 use tokio::codec::Framed;
15 use tokio::net::TcpStream;
16 use tokio::timer::Delay;
17
18 use futures::sync::mpsc;
19
20 use crate::printer::Printer;
21 use crate::timeout_stream::TimeoutStream;
22
23 struct BytesCoder<'a>(&'a mut bytes::BytesMut);
24 impl<'a> std::io::Write for BytesCoder<'a> {
25         fn write(&mut self, b: &[u8]) -> Result<usize, std::io::Error> {
26                 self.0.extend_from_slice(&b);
27                 Ok(b.len())
28         }
29         fn flush(&mut self) -> Result<(), std::io::Error> {
30                 Ok(())
31         }
32 }
33 struct BytesDecoder<'a> {
34         buf: &'a mut bytes::BytesMut,
35         pos: usize,
36 }
37 impl<'a> std::io::Read for BytesDecoder<'a> {
38         fn read(&mut self, b: &mut [u8]) -> Result<usize, std::io::Error> {
39                 let copy_len = cmp::min(b.len(), self.buf.len() - self.pos);
40                 b[..copy_len].copy_from_slice(&self.buf[self.pos..self.pos + copy_len]);
41                 self.pos += copy_len;
42                 Ok(copy_len)
43         }
44 }
45
46 struct MsgCoder<'a>(&'a Printer);
47 impl<'a> codec::Decoder for MsgCoder<'a> {
48         type Item = NetworkMessage;
49         type Error = std::io::Error;
50
51         fn decode(&mut self, bytes: &mut bytes::BytesMut) -> Result<Option<NetworkMessage>, std::io::Error> {
52                 let mut decoder = BytesDecoder {
53                         buf: bytes,
54                         pos: 0
55                 };
56                 match RawNetworkMessage::consensus_decode(&mut decoder) {
57                         Ok(res) => {
58                                 decoder.buf.advance(decoder.pos);
59                                 if res.magic == Network::Bitcoin.magic() {
60                                         Ok(Some(res.payload))
61                                 } else {
62                                         Err(std::io::Error::new(std::io::ErrorKind::InvalidData, "bad net magic"))
63                                 }
64                         },
65                         Err(e) => match e {
66                                 encode::Error::Io(_) => Ok(None),
67                                 encode::Error::UnrecognizedNetworkCommand(_msg) => {
68                                         decoder.buf.advance(decoder.pos);
69                                         //XXX(fixthese): self.0.add_line(format!("rust-bitcoin doesn't support {}!", msg), true);
70                                         Ok(None)
71                                 },
72                                 _ => {
73                                         self.0.add_line(format!("Error decoding message: {:?}", e), true);
74                                         Err(std::io::Error::new(std::io::ErrorKind::InvalidData, e))
75                                 },
76                         }
77                 }
78         }
79 }
80 impl<'a> codec::Encoder for MsgCoder<'a> {
81         type Item = NetworkMessage;
82         type Error = std::io::Error;
83
84         fn encode(&mut self, msg: NetworkMessage, res: &mut bytes::BytesMut) -> Result<(), std::io::Error> {
85                 if let Err(_) = (RawNetworkMessage {
86                         magic: Network::Bitcoin.magic(),
87                         payload: msg,
88                 }.consensus_encode(&mut BytesCoder(res))) {
89                         //XXX
90                 }
91                 Ok(())
92         }
93 }
94
95 pub struct Peer {}
96 impl Peer {
97         pub fn new(addr: SocketAddr, timeout: Duration, printer: &'static Printer) -> impl Future<Error=(), Item=(mpsc::Sender<NetworkMessage>, impl Stream<Item=NetworkMessage, Error=std::io::Error>)> {
98                 let connect_timeout = Delay::new(Instant::now() + timeout.clone()).then(|_| {
99                         future::err(std::io::Error::new(std::io::ErrorKind::TimedOut, "timeout reached"))
100                 });
101                 TcpStream::connect(&addr).select(connect_timeout)
102                         .or_else(move |_| {
103                                 Delay::new(Instant::now() + timeout / 10).then(|_| {
104                                         future::err(())
105                                 })
106                         }).and_then(move |stream| {
107                                 let (write, read) = Framed::new(stream.0, MsgCoder(printer)).split();
108                                 let (mut sender, receiver) = mpsc::channel(10); // We never really should send more than 10 messages unless they're dumb
109                                 tokio::spawn(write.sink_map_err(|_| { () }).send_all(receiver)
110                                         .then(|_| {
111                                                 future::err(())
112                                         }));
113                                 let _ = sender.try_send(NetworkMessage::Version(VersionMessage {
114                                         version: 70015,
115                                         services: (1 << 3), // NODE_WITNESS
116                                         timestamp: SystemTime::now().duration_since(UNIX_EPOCH).expect("time > 1970").as_secs() as i64,
117                                         receiver: Address::new(&addr, 0),
118                                         sender: Address::new(&"0.0.0.0:0".parse().unwrap(), 0),
119                                         nonce: 0xdeadbeef,
120                                         user_agent: "/rust-bitcoin:0.18/bluematt-tokio-client:0.1/".to_string(),
121                                         start_height: 0,
122                                         relay: true,
123                                 }));
124                                 future::ok((sender, TimeoutStream::new(read, timeout)))
125                         })
126         }
127 }