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