Update node locations when they come in, and get the real IP addr of nodes

This commit is contained in:
James Wilson
2021-06-24 18:34:32 +01:00
parent e383866322
commit fb80edba47
11 changed files with 201 additions and 17 deletions
@@ -3,7 +3,7 @@ use std::sync::atomic::AtomicU64;
use futures::channel::mpsc;
use futures::{ future, Sink, SinkExt };
use super::inner_loop;
use super::find_location::{ self, find_location };
use crate::find_location::find_location;
use crate::state::NodeId;
use std::net::Ipv4Addr;
@@ -1,206 +0,0 @@
use std::net::Ipv4Addr;
use std::sync::Arc;
use parking_lot::RwLock;
use rustc_hash::FxHashMap;
use serde::Deserialize;
use futures::{Sink, SinkExt, StreamExt};
use futures::channel::mpsc;
use common::types::NodeLocation;
use tokio::sync::Semaphore;
/// The returned location is optional; it may be None if not found.
pub type Location = Option<Arc<NodeLocation>>;
/// This is responsible for taking an IP address and attempting
/// to find a geographical location from this
pub fn find_location<Id, R>(response_chan: R) -> mpsc::UnboundedSender<(Id, Ipv4Addr)>
where
R: Sink<(Id, Option<Arc<NodeLocation>>)> + Unpin + Send + Clone + 'static,
Id: Clone + Send + 'static
{
let (tx, mut rx) = mpsc::unbounded();
// cache entries
let mut cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>> = FxHashMap::default();
// Default entry for localhost
cache.insert(
Ipv4Addr::new(127, 0, 0, 1),
Some(Arc::new(NodeLocation {
latitude: 52.516_6667,
longitude: 13.4,
city: "Berlin".into(),
})),
);
// Create a locator with our cache. This is used to obtain locations.
let locator = Locator::new(cache);
// Spawn a loop to handle location requests
tokio::spawn(async move {
// Allow 4 requests at a time. acquiring a token will block while the
// number of concurrent location requests is more than this.
let semaphore = Arc::new(Semaphore::new(4));
loop {
while let Some((id, ip_address)) = rx.next().await {
let permit = semaphore.clone().acquire_owned().await.unwrap();
let mut response_chan = response_chan.clone();
let locator = locator.clone();
// Once we have acquired our permit, spawn a task to avoid
// blocking this loop so that we can handle concurrent requests.
tokio::spawn(async move {
match locator.locate(ip_address).await {
Ok(loc) => {
let _ = response_chan.send((id,loc)).await;
},
Err(e) => {
log::debug!("GET error for ip location: {:?}", e);
}
};
// ensure permit is moved into task by dropping it explicitly:
drop(permit);
});
}
}
});
tx
}
/// This struct can be used to make location requests, given
/// an IPV4 address.
#[derive(Clone)]
struct Locator {
client: reqwest::Client,
cache: Arc<RwLock<FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>>>>,
}
impl Locator {
pub fn new(cache: FxHashMap<Ipv4Addr, Option<Arc<NodeLocation>>>) -> Self {
let client = reqwest::Client::new();
Locator {
client,
cache: Arc::new(RwLock::new(cache))
}
}
pub async fn locate(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
// Return location quickly if it's cached:
let cached_loc = {
let cache_reader = self.cache.read();
cache_reader.get(&ip).map(|o| o.clone())
};
if let Some(loc) = cached_loc {
return Ok(loc);
}
// Look it up via the location services if not cached:
let location = self.iplocate_ipapi_co(ip).await?;
let location = match location {
Some(location) => Ok(Some(location)),
None => self.iplocate_ipinfo_io(ip).await,
}?;
self.cache.write().insert(ip, location.clone());
Ok(location)
}
async fn iplocate_ipapi_co(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
let location = self
.query(&format!("https://ipapi.co/{}/json", ip))
.await?
.map(Arc::new);
Ok(location)
}
async fn iplocate_ipinfo_io(&self, ip: Ipv4Addr) -> Result<Option<Arc<NodeLocation>>, reqwest::Error> {
let location = self
.query(&format!("https://ipinfo.io/{}/json", ip))
.await?
.and_then(|loc: IPApiLocate| loc.into_node_location().map(Arc::new));
Ok(location)
}
async fn query<T>(&self, url: &str) -> Result<Option<T>, reqwest::Error>
where for<'de> T: Deserialize<'de>
{
match self.client.get(url).send().await?.json::<T>().await {
Ok(result) => Ok(Some(result)),
Err(err) => {
log::debug!("JSON error for ip location: {:?}", err);
Ok(None)
}
}
}
}
/// This is the format returned from ipinfo.co, so we do
/// a little conversion to get it into the shape we want.
#[derive(Deserialize)]
struct IPApiLocate {
city: Box<str>,
loc: Box<str>,
}
impl IPApiLocate {
fn into_node_location(self) -> Option<NodeLocation> {
let IPApiLocate { city, loc } = self;
let mut loc = loc.split(',').map(|n| n.parse());
let latitude = loc.next()?.ok()?;
let longitude = loc.next()?.ok()?;
// Guarantee that the iterator has been exhausted
if loc.next().is_some() {
return None;
}
Some(NodeLocation {
latitude,
longitude,
city,
})
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn ipapi_locate_to_node_location() {
let ipapi = IPApiLocate {
loc: "12.5,56.25".into(),
city: "Foobar".into(),
};
let location = ipapi.into_node_location().unwrap();
assert_eq!(location.latitude, 12.5);
assert_eq!(location.longitude, 56.25);
assert_eq!(&*location.city, "Foobar");
}
#[test]
fn ipapi_locate_to_node_location_too_many() {
let ipapi = IPApiLocate {
loc: "12.5,56.25,1.0".into(),
city: "Foobar".into(),
};
let location = ipapi.into_node_location();
assert!(location.is_none());
}
}
+33 -5
View File
@@ -8,13 +8,13 @@ use common::{
util::now
};
use bimap::BiMap;
use std::{iter::FromIterator, net::Ipv4Addr, str::FromStr};
use std::{net::{IpAddr, Ipv4Addr}, str::FromStr};
use futures::channel::{ mpsc };
use futures::{ SinkExt, StreamExt };
use std::collections::{ HashMap, HashSet };
use crate::state::{ self, State, NodeId };
use crate::feed_message::{ self, FeedMessageSerializer };
use super::find_location;
use crate::find_location;
/// A unique Id is assigned per websocket connection (or more accurately,
/// per feed socket and per shard socket). This can be combined with the
@@ -143,7 +143,7 @@ pub struct InnerLoop {
/// These feeds want finality info, too.
feed_conn_id_finality: HashSet<ConnId>,
/// Send messages here to make location requests, which are sent back into the loop.
/// Send messages here to make geographical location requests.
tx_to_locator: mpsc::UnboundedSender<(NodeId, Ipv4Addr)>
}
@@ -184,8 +184,32 @@ impl InnerLoop {
}
}
/// Handle messages that come from the node geographical locator.
async fn handle_from_find_location(&mut self, node_id: NodeId, location: find_location::Location) {
// TODO: Update node location here
self.node_state.update_node_location(node_id, location.clone());
if let Some(loc) = location {
let mut feed_message_serializer = FeedMessageSerializer::new();
feed_message_serializer.push(feed_message::LocatedNode(
node_id,
loc.latitude,
loc.longitude,
&loc.city
));
if let Some(bytes) = feed_message_serializer.into_finalized() {
let chain_label = self.node_state
.get_node_chain(node_id)
.map(|chain| chain.label());
if let Some(chain_label) = chain_label {
// Don't hold onto lifetime from self because we call a mut fn next:
let label = chain_label.to_owned();
// Update location for any feeds subscribed to the node's chain.
self.broadcast_to_chain_feeds(&label, ToFeedWebsocket::Bytes(bytes)).await;
}
}
}
}
/// Handle messages coming from shards.
@@ -228,7 +252,10 @@ impl InnerLoop {
).await
}
// TODO: The node has been added. use it's IP to find a location.
// Currently we only geographically locate IPV4 addresses so ignore IPV6;
if let Some(IpAddr::V4(ip_v4)) = ip {
let _ = self.tx_to_locator.send((node_id, ip_v4)).await;
}
},
}
},
@@ -409,6 +436,7 @@ impl InnerLoop {
if let Some(feeds) = self.chain_to_feed_conn_ids.get(chain) {
for &feed_id in feeds {
// How much faster would it be if we processed these in parallel?
// Is it practical to do so given lifetimes and such?
if let Some(chan) = self.feed_channels.get_mut(&feed_id) {
chan.send(message.clone()).await;
}
-1
View File
@@ -1,6 +1,5 @@
mod aggregator;
mod inner_loop;
mod find_location;
// Expose the various message types that can be worked with externally:
pub use inner_loop::{ FromFeedWebsocket, FromShardWebsocket, ToFeedWebsocket, ToShardWebsocket };