_ Git - dnsseed-rust/blob - src/datastore.rs

   1 use std::{cmp, mem};
   2 use std::collections::{HashSet, HashMap, hash_map};
   3 use std::sync::{Arc, RwLock};
   4 use std::net::SocketAddr;
   5 use std::time::{Duration, Instant};
   6 use std::io::{BufRead, BufReader};
   7
   8 use bitcoin::network::address::Address;
   9
  10 use rand::thread_rng;
  11 use rand::seq::{SliceRandom, IteratorRandom};
  12
  13 use tokio::prelude::*;
  14 use tokio::fs::File;
  15 use tokio::io::write_all;
  16
  17 use regex::Regex;
  18
  19 #[derive(Clone, Copy, Hash, PartialEq, Eq)]
  20 pub enum AddressState {
  21         Untested,
  22         LowBlockCount,
  23         HighBlockCount,
  24         LowVersion,
  25         BadVersion,
  26         NotFullNode,
  27         ProtocolViolation,
  28         Timeout,
  29         TimeoutDuringRequest,
  30         Good,
  31         WasGood,
  32 }
  33
  34 #[derive(Hash, PartialEq, Eq)]
  35 pub enum U64Setting {
  36         ConnsPerSec,
  37         RunTimeout,
  38         WasGoodTimeout,
  39         RescanInterval(AddressState),
  40         MinProtocolVersion,
  41 }
  42
  43 #[derive(Hash, PartialEq, Eq)]
  44 pub enum RegexSetting {
  45         SubverRegex,
  46 }
  47
  48 struct Node {
  49         last_update: Instant,
  50         last_good: Instant, // Ignored unless state is Good or WasGood
  51         last_services: u64,
  52         state: AddressState,
  53 }
  54
  55 struct Nodes {
  56         good_node_services: HashMap<u8, HashSet<SocketAddr>>,
  57         nodes_to_state: HashMap<SocketAddr, Node>,
  58         state_next_scan: HashMap<AddressState, Vec<(Instant, SocketAddr)>>,
  59 }
  60 struct NodesMutRef<'a> {
  61         good_node_services: &'a mut HashMap<u8, HashSet<SocketAddr>>,
  62         nodes_to_state: &'a mut HashMap<SocketAddr, Node>,
  63         state_next_scan: &'a mut HashMap<AddressState, Vec<(Instant, SocketAddr)>>,
  64
  65 }
  66 impl Nodes {
  67         fn borrow_mut<'a>(&'a mut self) -> NodesMutRef<'a> {
  68                 NodesMutRef {
  69                         good_node_services: &mut self.good_node_services,
  70                         nodes_to_state: &mut self.nodes_to_state,
  71                         state_next_scan: &mut self.state_next_scan,
  72                 }
  73         }
  74 }
  75
  76 pub struct Store {
  77         u64_settings: RwLock<HashMap<U64Setting, u64>>,
  78         subver_regex: RwLock<Arc<Regex>>,
  79         nodes: RwLock<Nodes>,
  80         store: String,
  81 }
  82
  83 impl Store {
  84         pub fn new(store: String) -> impl Future<Item=Store, Error=()> {
  85                 let settings_future = File::open(store.clone() + "/settings").and_then(|f| {
  86                         let mut l = BufReader::new(f).lines();
  87                         macro_rules! try_read {
  88                                 ($lines: expr, $ty: ty) => { {
  89                                         match $lines.next() {
  90                                                 Some(line) => match line {
  91                                                         Ok(line) => match line.parse::<$ty>() {
  92                                                                 Ok(res) => res,
  93                                                                 Err(e) => return future::err(std::io::Error::new(std::io::ErrorKind::InvalidData, e)),
  94                                                         },
  95                                                         Err(e) => return future::err(e),
  96                                                 },
  97                                                 None => return future::err(std::io::Error::new(std::io::ErrorKind::UnexpectedEof, "")),
  98                                         }
  99                                 } }
 100                         }
 101                         let mut u64s = HashMap::with_capacity(15);
 102                         u64s.insert(U64Setting::ConnsPerSec, try_read!(l, u64));
 103                         u64s.insert(U64Setting::RunTimeout, try_read!(l, u64));
 104                         u64s.insert(U64Setting::WasGoodTimeout, try_read!(l, u64));
 105                         u64s.insert(U64Setting::MinProtocolVersion, try_read!(l, u64));
 106                         u64s.insert(U64Setting::RescanInterval(AddressState::Untested), try_read!(l, u64));
 107                         u64s.insert(U64Setting::RescanInterval(AddressState::LowBlockCount), try_read!(l, u64));
 108                         u64s.insert(U64Setting::RescanInterval(AddressState::HighBlockCount), try_read!(l, u64));
 109                         u64s.insert(U64Setting::RescanInterval(AddressState::LowVersion), try_read!(l, u64));
 110                         u64s.insert(U64Setting::RescanInterval(AddressState::BadVersion), try_read!(l, u64));
 111                         u64s.insert(U64Setting::RescanInterval(AddressState::NotFullNode), try_read!(l, u64));
 112                         u64s.insert(U64Setting::RescanInterval(AddressState::ProtocolViolation), try_read!(l, u64));
 113                         u64s.insert(U64Setting::RescanInterval(AddressState::Timeout), try_read!(l, u64));
 114                         u64s.insert(U64Setting::RescanInterval(AddressState::TimeoutDuringRequest), try_read!(l, u64));
 115                         u64s.insert(U64Setting::RescanInterval(AddressState::Good), try_read!(l, u64));
 116                         u64s.insert(U64Setting::RescanInterval(AddressState::WasGood), try_read!(l, u64));
 117                         future::ok((u64s, try_read!(l, Regex)))
 118                 }).or_else(|_| -> future::FutureResult<(HashMap<U64Setting, u64>, Regex), ()> {
 119                         let mut u64s = HashMap::with_capacity(15);
 120                         u64s.insert(U64Setting::ConnsPerSec, 10);
 121                         u64s.insert(U64Setting::RunTimeout, 120);
 122                         u64s.insert(U64Setting::WasGoodTimeout, 21600);
 123                         u64s.insert(U64Setting::RescanInterval(AddressState::Untested), 0);
 124                         u64s.insert(U64Setting::RescanInterval(AddressState::LowBlockCount), 3600);
 125                         u64s.insert(U64Setting::RescanInterval(AddressState::HighBlockCount), 7200);
 126                         u64s.insert(U64Setting::RescanInterval(AddressState::LowVersion), 21600);
 127                         u64s.insert(U64Setting::RescanInterval(AddressState::BadVersion), 21600);
 128                         u64s.insert(U64Setting::RescanInterval(AddressState::NotFullNode), 86400);
 129                         u64s.insert(U64Setting::RescanInterval(AddressState::ProtocolViolation), 86400);
 130                         u64s.insert(U64Setting::RescanInterval(AddressState::Timeout), 86400);
 131                         u64s.insert(U64Setting::RescanInterval(AddressState::TimeoutDuringRequest), 21600);
 132                         u64s.insert(U64Setting::RescanInterval(AddressState::Good), 1800);
 133                         u64s.insert(U64Setting::RescanInterval(AddressState::WasGood), 1800);
 134                         u64s.insert(U64Setting::MinProtocolVersion, 10000); //XXX
 135                         future::ok((u64s, Regex::new(".*").unwrap()))
 136                 });
 137
 138                 macro_rules! nodes_uninitd {
 139                         () => { {
 140                                 let mut state_vecs = HashMap::with_capacity(11);
 141                                 state_vecs.insert(AddressState::Untested, Vec::new());
 142                                 state_vecs.insert(AddressState::LowBlockCount, Vec::new());
 143                                 state_vecs.insert(AddressState::HighBlockCount, Vec::new());
 144                                 state_vecs.insert(AddressState::LowVersion, Vec::new());
 145                                 state_vecs.insert(AddressState::BadVersion, Vec::new());
 146                                 state_vecs.insert(AddressState::NotFullNode, Vec::new());
 147                                 state_vecs.insert(AddressState::ProtocolViolation, Vec::new());
 148                                 state_vecs.insert(AddressState::Timeout, Vec::new());
 149                                 state_vecs.insert(AddressState::TimeoutDuringRequest, Vec::new());
 150                                 state_vecs.insert(AddressState::Good, Vec::new());
 151                                 state_vecs.insert(AddressState::WasGood, Vec::new());
 152                                 let mut good_node_services = HashMap::with_capacity(64);
 153                                 for i in 0..64 {
 154                                         good_node_services.insert(i, HashSet::new());
 155                                 }
 156                                 Nodes {
 157                                         good_node_services,
 158                                         nodes_to_state: HashMap::new(),
 159                                         state_next_scan: state_vecs,
 160                                 }
 161                         } }
 162                 }
 163
 164                 let nodes_future = File::open(store.clone() + "/nodes").and_then(|f| {
 165                         let mut res = nodes_uninitd!();
 166                         let l = BufReader::new(f).lines();
 167                         for line_res in l {
 168                                 let line = match line_res {
 169                                         Ok(l) => l,
 170                                         Err(_) => return future::ok(res),
 171                                 };
 172                                 let mut line_iter = line.split(',');
 173                                 macro_rules! try_read {
 174                                         ($lines: expr, $ty: ty) => { {
 175                                                 match $lines.next() {
 176                                                         Some(line) => match line.parse::<$ty>() {
 177                                                                 Ok(res) => res,
 178                                                                 Err(_) => return future::ok(res),
 179                                                         },
 180                                                         None => return future::ok(res),
 181                                                 }
 182                                         } }
 183                                 }
 184                                 let sockaddr = try_read!(line_iter, SocketAddr);
 185                                 let state = try_read!(line_iter, u8);
 186                                 let last_services = try_read!(line_iter, u64);
 187                                 let node = Node {
 188                                         state: match state {
 189                                                 0x0 => AddressState::Untested,
 190                                                 0x1 => AddressState::LowBlockCount,
 191                                                 0x2 => AddressState::HighBlockCount,
 192                                                 0x3 => AddressState::LowVersion,
 193                                                 0x4 => AddressState::BadVersion,
 194                                                 0x5 => AddressState::NotFullNode,
 195                                                 0x6 => AddressState::ProtocolViolation,
 196                                                 0x7 => AddressState::Timeout,
 197                                                 0x8 => AddressState::TimeoutDuringRequest,
 198                                                 0x9 => AddressState::Good,
 199                                                 0xa => AddressState::WasGood,
 200                                                 _   => return future::ok(res),
 201                                         },
 202                                         last_services,
 203                                         last_update: Instant::now(),
 204                                         last_good: Instant::now(),
 205                                 };
 206                                 if node.state == AddressState::Good {
 207                                         for i in 0..64 {
 208                                                 if node.last_services & (1 << i) != 0 {
 209                                                         res.good_node_services.get_mut(&i).unwrap().insert(sockaddr);
 210                                                 }
 211                                         }
 212                                 }
 213                                 res.state_next_scan.get_mut(&node.state).unwrap().push((Instant::now(), sockaddr));
 214                                 res.nodes_to_state.insert(sockaddr, node);
 215                         }
 216                         future::ok(res)
 217                 }).or_else(|_| -> future::FutureResult<Nodes, ()> {
 218                         future::ok(nodes_uninitd!())
 219                 });
 220                 settings_future.join(nodes_future).and_then(move |((u64_settings, regex), nodes)| {
 221                         future::ok(Store {
 222                                 u64_settings: RwLock::new(u64_settings),
 223                                 subver_regex: RwLock::new(Arc::new(regex)),
 224                                 nodes: RwLock::new(nodes),
 225                                 store,
 226                         })
 227                 })
 228         }
 229
 230         pub fn get_u64(&self, setting: U64Setting) -> u64 {
 231                 *self.u64_settings.read().unwrap().get(&setting).unwrap()
 232         }
 233
 234         pub fn set_u64(&self, setting: U64Setting, value: u64) {
 235                 *self.u64_settings.write().unwrap().get_mut(&setting).unwrap() = value;
 236         }
 237
 238         pub fn get_node_count(&self, state: AddressState) -> usize {
 239                 self.nodes.read().unwrap().state_next_scan.get(&state).unwrap().len()
 240         }
 241
 242         pub fn get_regex(&self, _setting: RegexSetting) -> Arc<Regex> {
 243                 Arc::clone(&*self.subver_regex.read().unwrap())
 244         }
 245
 246         pub fn set_regex(&self, _setting: RegexSetting, value: Regex) {
 247                 *self.subver_regex.write().unwrap() = Arc::new(value);
 248         }
 249
 250         pub fn add_fresh_nodes(&self, addresses: &Vec<(u32, Address)>) {
 251                 let mut nodes = self.nodes.write().unwrap();
 252                 let cur_time = Instant::now();
 253                 for &(_, ref addr) in addresses {
 254                         if let Ok(socketaddr) = addr.socket_addr() {
 255                                 match nodes.nodes_to_state.entry(socketaddr.clone()) {
 256                                         hash_map::Entry::Vacant(e) => {
 257                                                 e.insert(Node {
 258                                                         state: AddressState::Untested,
 259                                                         last_services: 0,
 260                                                         last_update: cur_time,
 261                                                         last_good: Instant::now(),
 262                                                 });
 263                                                 nodes.state_next_scan.get_mut(&AddressState::Untested).unwrap().push((cur_time, socketaddr));
 264                                         },
 265                                         hash_map::Entry::Occupied(_) => {},
 266                                 }
 267                         } else {
 268                                 //TODO: Handle onions
 269                         }
 270                 }
 271         }
 272
 273         pub fn set_node_state(&self, addr: SocketAddr, state: AddressState, services: u64) -> bool {
 274                 let mut nodes_lock = self.nodes.write().unwrap();
 275                 let nodes = nodes_lock.borrow_mut();
 276                 let state_ref = nodes.nodes_to_state.get_mut(&addr).unwrap();
 277                 let ret = state != state_ref.state;
 278                 state_ref.last_update = Instant::now();
 279                 if (state_ref.state == AddressState::Good || state_ref.state == AddressState::WasGood)
 280                                 && state != AddressState::Good
 281                                 && state_ref.last_good >= state_ref.last_update + Duration::from_secs(self.get_u64(U64Setting::WasGoodTimeout)) {
 282                         state_ref.state = AddressState::WasGood;
 283                         for i in 0..64 {
 284                                 if state_ref.last_services & (1 << i) != 0 {
 285                                         nodes.good_node_services.get_mut(&i).unwrap().remove(&addr);
 286                                 }
 287                         }
 288                         state_ref.last_services = 0;
 289                         nodes.state_next_scan.get_mut(&AddressState::WasGood).unwrap().push((state_ref.last_update, addr));
 290                 } else {
 291                         state_ref.state = state;
 292                         if state == AddressState::Good {
 293                                 for i in 0..64 {
 294                                         if services & (1 << i) != 0 && state_ref.last_services & (1 << i) == 0 {
 295                                                 nodes.good_node_services.get_mut(&i).unwrap().insert(addr);
 296                                         } else if services & (1 << i) == 0 && state_ref.last_services & (1 << i) != 0 {
 297                                                 nodes.good_node_services.get_mut(&i).unwrap().remove(&addr);
 298                                         }
 299                                 }
 300                                 state_ref.last_services = services;
 301                                 state_ref.last_good = state_ref.last_update;
 302                         }
 303                         nodes.state_next_scan.get_mut(&state).unwrap().push((state_ref.last_update, addr));
 304                 }
 305                 ret
 306         }
 307
 308         pub fn save_data(&'static self) -> impl Future<Item=(), Error=()> {
 309                 let settings_file = self.store.clone() + "/settings";
 310                 let settings_future = File::create(settings_file.clone() + ".tmp").and_then(move |f| {
 311                         let settings_string = format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}",
 312                                 self.get_u64(U64Setting::ConnsPerSec),
 313                                 self.get_u64(U64Setting::RunTimeout),
 314                                 self.get_u64(U64Setting::WasGoodTimeout),
 315                                 self.get_u64(U64Setting::MinProtocolVersion),
 316                                 self.get_u64(U64Setting::RescanInterval(AddressState::Untested)),
 317                                 self.get_u64(U64Setting::RescanInterval(AddressState::LowBlockCount)),
 318                                 self.get_u64(U64Setting::RescanInterval(AddressState::HighBlockCount)),
 319                                 self.get_u64(U64Setting::RescanInterval(AddressState::LowVersion)),
 320                                 self.get_u64(U64Setting::RescanInterval(AddressState::BadVersion)),
 321                                 self.get_u64(U64Setting::RescanInterval(AddressState::NotFullNode)),
 322                                 self.get_u64(U64Setting::RescanInterval(AddressState::ProtocolViolation)),
 323                                 self.get_u64(U64Setting::RescanInterval(AddressState::Timeout)),
 324                                 self.get_u64(U64Setting::RescanInterval(AddressState::TimeoutDuringRequest)),
 325                                 self.get_u64(U64Setting::RescanInterval(AddressState::Good)),
 326                                 self.get_u64(U64Setting::RescanInterval(AddressState::WasGood)),
 327                                 self.get_regex(RegexSetting::SubverRegex).as_str());
 328                         write_all(f, settings_string).and_then(|(mut f, _)| {
 329                                 f.poll_sync_all()
 330                         }).and_then(|_| {
 331                                 tokio::fs::rename(settings_file.clone() + ".tmp", settings_file)
 332                         })
 333                 });
 334
 335                 let nodes_file = self.store.clone() + "/nodes";
 336                 let nodes_future = File::create(nodes_file.clone() + ".tmp").and_then(move |f| {
 337                         let mut nodes_buff = String::new();
 338                         {
 339                                 let nodes = self.nodes.read().unwrap();
 340                                 nodes_buff.reserve(nodes.nodes_to_state.len() * 20);
 341                                 for (ref sockaddr, ref node) in nodes.nodes_to_state.iter() {
 342                                         nodes_buff += &sockaddr.to_string();
 343                                         nodes_buff += ",";
 344                                         nodes_buff += &match node.state {
 345                                                 AddressState::Untested => 0u8,
 346                                                 AddressState::LowBlockCount => 1u8,
 347                                                 AddressState::HighBlockCount => 2u8,
 348                                                 AddressState::LowVersion => 3u8,
 349                                                 AddressState::BadVersion => 4u8,
 350                                                 AddressState::NotFullNode => 5u8,
 351                                                 AddressState::ProtocolViolation => 6u8,
 352                                                 AddressState::Timeout => 7u8,
 353                                                 AddressState::TimeoutDuringRequest => 8u8,
 354                                                 AddressState::Good => 9u8,
 355                                                 AddressState::WasGood => 10u8,
 356                                         }.to_string();
 357                                         nodes_buff += ",";
 358                                         nodes_buff += &node.last_services.to_string();
 359                                         nodes_buff += "\n";
 360                                 }
 361                         }
 362                         write_all(f, nodes_buff)
 363                 }).and_then(|(mut f, _)| {
 364                         f.poll_sync_all()
 365                 }).and_then(|_| {
 366                         tokio::fs::rename(nodes_file.clone() + ".tmp", nodes_file)
 367                 });
 368
 369                 let dns_file = self.store.clone() + "/nodes.dump";
 370                 let dns_future = File::create(dns_file.clone() + ".tmp").and_then(move |f| {
 371                         let mut dns_buff = String::new();
 372                         {
 373                                 let nodes = self.nodes.read().unwrap();
 374                                 let mut rng = thread_rng();
 375                                 for i in &[1u64, 4, 5, 8, 9, 12, 13, 1024, 1025, 1028, 1029, 1032, 1033, 1036, 1037] {
 376                                         let mut v6_set = Vec::new();
 377                                         let mut v4_set = Vec::new();
 378                                         if i.count_ones() == 1 {
 379                                                 for j in 0..64 {
 380                                                         if i & (1 << j) != 0 {
 381                                                                 let set_ref = nodes.good_node_services.get(&j).unwrap();
 382                                                                 v4_set = set_ref.iter().filter(|e| e.is_ipv4() && e.port() == 8333)
 383                                                                         .choose_multiple(&mut rng, 21).iter().map(|e| e.ip()).collect();
 384                                                                 v6_set = set_ref.iter().filter(|e| e.is_ipv6() && e.port() == 8333)
 385                                                                         .choose_multiple(&mut rng, 12).iter().map(|e| e.ip()).collect();
 386                                                                 break;
 387                                                         }
 388                                                 }
 389                                         } else if i.count_ones() == 2 {
 390                                                 let mut first_set = None;
 391                                                 let mut second_set = None;
 392                                                 for j in 0..64 {
 393                                                         if i & (1 << j) != 0 {
 394                                                                 if first_set == None {
 395                                                                         first_set = Some(nodes.good_node_services.get(&j).unwrap());
 396                                                                 } else {
 397                                                                         second_set = Some(nodes.good_node_services.get(&j).unwrap());
 398                                                                         break;
 399                                                                 }
 400                                                         }
 401                                                 }
 402                                                 v4_set = first_set.unwrap().intersection(second_set.unwrap())
 403                                                         .filter(|e| e.is_ipv4() && e.port() == 8333)
 404                                                         .choose_multiple(&mut rng, 21).iter().map(|e| e.ip()).collect();
 405                                                 v6_set = first_set.unwrap().intersection(second_set.unwrap())
 406                                                         .filter(|e| e.is_ipv6() && e.port() == 8333)
 407                                                         .choose_multiple(&mut rng, 12).iter().map(|e| e.ip()).collect();
 408                                         } else {
 409                                                 //TODO: Could optimize this one a bit
 410                                                 let mut intersection;
 411                                                 let mut intersection_set_ref = None;
 412                                                 for j in 0..64 {
 413                                                         if i & (1 << j) != 0 {
 414                                                                 if intersection_set_ref == None {
 415                                                                         intersection_set_ref = Some(nodes.good_node_services.get(&j).unwrap());
 416                                                                 } else {
 417                                                                         let new_intersection = intersection_set_ref.unwrap()
 418                                                                                 .intersection(nodes.good_node_services.get(&j).unwrap()).map(|e| (*e).clone()).collect();
 419                                                                         intersection = Some(new_intersection);
 420                                                                         intersection_set_ref = Some(intersection.as_ref().unwrap());
 421                                                                 }
 422                                                         }
 423                                                 }
 424                                                 v4_set = intersection_set_ref.unwrap().iter()
 425                                                         .filter(|e| e.is_ipv4() && e.port() == 8333)
 426                                                         .choose_multiple(&mut rng, 21).iter().map(|e| e.ip()).collect();
 427                                                 v6_set = intersection_set_ref.unwrap().iter()
 428                                                         .filter(|e| e.is_ipv6() && e.port() == 8333)
 429                                                         .choose_multiple(&mut rng, 12).iter().map(|e| e.ip()).collect();
 430                                         }
 431                                         for a in v4_set {
 432                                                 dns_buff += &format!("x{:x}.dnsseed\tIN\tA\t{}\n", i, a);
 433                                         }
 434                                         for a in v6_set {
 435                                                 dns_buff += &format!("x{:x}.dnsseed\tIN\tAAAA\t{}\n", i, a);
 436                                         }
 437                                 }
 438                         }
 439                         write_all(f, dns_buff)
 440                 }).and_then(|(mut f, _)| {
 441                         f.poll_sync_all()
 442                 }).and_then(|_| {
 443                         tokio::fs::rename(dns_file.clone() + ".tmp", dns_file)
 444                 });
 445
 446                 settings_future.join3(nodes_future, dns_future).then(|_| { future::ok(()) })
 447         }
 448
 449         pub fn get_next_scan_nodes(&self) -> Vec<SocketAddr> {
 450                 let mut res = Vec::with_capacity(600);
 451                 let cur_time = Instant::now();
 452                 let mut nodes = self.nodes.write().unwrap();
 453                 for (state, state_nodes) in nodes.state_next_scan.iter_mut() {
 454                         let cmp_time = cur_time - Duration::from_secs(self.get_u64(U64Setting::RescanInterval(*state)));
 455                         let split_point = cmp::min(cmp::min(600 - res.len(), 60),
 456                                         state_nodes.binary_search_by(|a| a.0.cmp(&cmp_time)).unwrap_or_else(|idx| idx));
 457                         let mut new_nodes = state_nodes.split_off(split_point);
 458                         mem::swap(&mut new_nodes, state_nodes);
 459                         for (_, node) in new_nodes.drain(..) {
 460                                 res.push(node);
 461                         }
 462                 }
 463                 res.shuffle(&mut thread_rng());
 464                 res
 465         }
 466 }