Contractless/src/records/memory/connections.rs

774 lines
27 KiB
Rust

use crate::common::binary_conversions::{binary_to_ip, ip_to_binary};
use crate::lazy_static;
use crate::log::{info, warn};
use crate::records::memory::enums::{ClientType, ConnectionType};
use crate::records::memory::network_mapping::monitor::{MONITOR_ACTION_ADD, MONITOR_ACTION_REMOVE};
use crate::records::memory::network_mapping::structs::{MonitorAddressParams, SignedMonitorEdit};
use crate::records::memory::network_mapping::NodeInfo;
use crate::records::memory::response_channels::{delete_entry, reserve_entry, Command};
use crate::records::memory::structs::{Connection, StoreConnectionParams};
use crate::rpc::client::handshake::connect_and_handshake;
use crate::rpc::client::handshake_processing::{bootstrap_peer_discovery, BootstrapParams};
use crate::rpc::client::structs::Connect;
use crate::rpc::command_maps::RPC_BLOCK_HEIGHT;
use crate::rpc::responses::RpcResponse;
use crate::sled::Db;
use crate::sleep;
use crate::thread_rng;
use crate::timeout;
use crate::wallets::structures::Wallet;
use crate::Arc;
use crate::AsyncWriteExt;
use crate::AtomicBool;
use crate::AtomicOrdering;
use crate::Duration;
use crate::IteratorRandom;
use crate::Mutex;
use crate::RwLock;
use crate::TcpStream;
fn split_ip_port_key(value: &str) -> Option<(String, u16)> {
// Connection keys are stored as ip:port strings; IPv6 addresses may arrive
// bracketed, so strip brackets before parsing the port.
let (ip_part, port_part) = value.rsplit_once(':')?;
let ip = ip_part
.strip_prefix('[')
.and_then(|inner| inner.strip_suffix(']'))
.unwrap_or(ip_part)
.to_string();
let port = port_part.parse::<u16>().ok()?;
Some((ip, port))
}
use crate::records::memory::structs::{ConnectionInfo, ConnectionKey};
#[derive(Clone)]
struct ReconnectContext {
db: Db,
wallet: Arc<Wallet>,
map: Arc<Mutex<Command>>,
}
lazy_static! {
static ref RECONNECT_CONTEXT: Mutex<Option<ReconnectContext>> = Mutex::new(None);
static ref RECONNECT_IN_PROGRESS: AtomicBool = AtomicBool::new(false);
}
fn try_start_reconnect() -> bool {
// Only one reconnect path should run at a time, whether it came from
// liveness failure or bootstrap recovery.
RECONNECT_IN_PROGRESS
.compare_exchange(false, true, AtomicOrdering::SeqCst, AtomicOrdering::SeqCst)
.is_ok()
}
fn finish_reconnect() {
// Release the reconnect gate after the async reconnect attempt finishes.
RECONNECT_IN_PROGRESS.store(false, AtomicOrdering::SeqCst);
}
pub async fn set_reconnect_context(db: Db, wallet: Arc<Wallet>, map: Arc<Mutex<Command>>) {
let mut context = RECONNECT_CONTEXT.lock().await;
// Store enough state for later liveness checks to reconnect without
// needing the original startup stack.
*context = Some(ReconnectContext { db, wallet, map });
}
async fn reconnect_replacement_inner(excluded_ip: &str) {
// When an outgoing peer disappears, try to replace it with another
// active node that is not already connected and is not the failed IP.
let context = {
let guard = RECONNECT_CONTEXT.lock().await;
guard.clone()
};
let Some(context) = context else {
warn!("[reconnect] no reconnect context configured");
return;
};
let excluded_ip_bytes = ip_to_binary(excluded_ip);
let live_connection = {
let guard = CONNECTIONS.read().await;
guard.as_ref().and_then(|conn| {
conn.connection_map.iter().find_map(|(key, info)| {
if key.ip == excluded_ip_bytes {
return None;
}
if ClientType::from_bytes(&info.client_type) != Some(ClientType::Miner) {
return None;
}
let ip = binary_to_ip(key.ip.clone());
let connections_key = format!("{}:{}", ip, key.port);
Some((connections_key, Arc::clone(&info.stream)))
})
})
};
let Some((connections_key, stream)) = live_connection else {
warn!("[reconnect] no live stream available for bootstrap recovery");
return;
};
let bootstrap_params = BootstrapParams {
stream,
connections_key,
wallet: context.wallet,
db: context.db,
map: context.map,
first: false,
};
if let Err(err) = bootstrap_peer_discovery(bootstrap_params).await {
warn!("[reconnect] bootstrap recovery failed: {err}");
}
}
async fn retry_dropped_outgoing(ip: String, port: u16) {
if !try_start_reconnect() {
warn!("[reconnect] reconnect attempt already in progress, skipping duplicate request");
return;
}
async {
let context = {
let guard = RECONNECT_CONTEXT.lock().await;
guard.clone()
};
let Some(context) = context else {
warn!("[reconnect] no reconnect context configured");
return;
};
let addr_string = format!("{ip}:{port}");
for attempt in 1..=3 {
sleep(Duration::from_secs(30)).await;
let socket_addr = match addr_string.parse() {
Ok(addr) => addr,
Err(err) => {
warn!("[reconnect] invalid dropped peer address {addr_string}: {err}");
break;
}
};
let connect = Connect {
addr: socket_addr,
node_ip: addr_string.clone(),
wallet: context.wallet.clone(),
db: context.db.clone(),
map: context.map.clone(),
first: false,
};
match connect_and_handshake(connect).await {
Ok(()) => {
info!("[reconnect] reconnected dropped peer {addr_string} on attempt {attempt}");
return;
}
Err(err) => {
warn!(
"[reconnect] failed to reconnect dropped peer {addr_string} on attempt {attempt}/3: {err}"
);
}
}
}
reconnect_replacement_inner(&ip).await;
}
.await;
finish_reconnect();
}
pub fn spawn_retry_dropped_outgoing(ip: String, port: u16) {
tokio::spawn(async move {
retry_dropped_outgoing(ip, port).await;
});
}
pub fn spawn_reconnect_bootstrap(params: BootstrapParams) {
if !try_start_reconnect() {
warn!("[reconnect] bootstrap recovery already in progress, skipping duplicate request");
return;
}
// Bootstrap discovery can perform network requests, so it runs detached
// from the caller that noticed the connection problem.
tokio::spawn(async move {
if let Err(err) = bootstrap_peer_discovery(params).await {
warn!("[reconnect] bootstrap recovery failed: {err}");
}
finish_reconnect();
});
}
impl Connection {
// Initialize the in-memory connection manager state.
pub fn new() -> Self {
Self::default()
}
// Store a live socket in memory along with its role, peer identity,
// and session metadata used by the RPC and peer-management paths.
pub fn store_connection(&mut self, params: StoreConnectionParams) -> bool {
let StoreConnectionParams {
connection_type,
ip,
port,
stream,
client_type,
wallet_public_key,
command_map,
} = params;
let ip_bytes = ip_to_binary(&ip);
let connection_key = ConnectionKey {
connection_type: connection_type.as_bytes(),
ip: ip_bytes.clone(),
port,
};
let connection_key2 = ConnectionKey {
connection_type: connection_type.opposite().as_bytes(),
ip: ip_bytes.clone(),
port,
};
// Miner nodes are identified by IP, not by port. A second node
// announcing the same IP is rejected even if it uses another
// socket port.
if client_type == ClientType::Miner
&& self.connection_map.iter().any(|(key, info)| {
key.ip == ip_bytes
&& ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner)
})
{
return false;
}
// Non-miner RPC clients still use the full socket key so short
// request/response connections do not collide unnecessarily.
if self.connection_map.contains_key(&connection_key)
|| self.connection_map.contains_key(&connection_key2)
{
return false;
}
if wallet_public_key.len() != Wallet::PUBLIC_KEY_LENGTH {
return false;
}
let connection_info = ConnectionInfo::new(
connection_type.as_bytes(),
ip_bytes,
port,
stream.clone(),
client_type.as_bytes(),
wallet_public_key.clone(),
);
self.connection_map.insert(connection_key, connection_info);
if client_type == ClientType::Miner {
spawn_monitor_update(
ip.clone(),
MONITOR_ACTION_ADD,
wallet_public_key.clone(),
port,
);
Connection::client_checkup(stream, connection_type, ip, port, command_map);
}
true
}
// Remove a specific connection entry by direction, IP, and port.
pub fn drop_connection(
&mut self,
connection_type: ConnectionType,
ip: String,
port: u16,
) -> Option<ConnectionInfo> {
let ip_bytes = ip_to_binary(&ip);
let connection_key = ConnectionKey {
connection_type: connection_type.as_bytes(),
ip: ip_bytes,
port,
};
let removed = self.connection_map.remove(&connection_key);
if let Some(connection_info) = removed.as_ref() {
if ClientType::from_bytes(&connection_info.client_type) == Some(ClientType::Miner) {
spawn_monitor_update(
ip.clone(),
MONITOR_ACTION_REMOVE,
connection_info.wallet_public_key.clone(),
port,
);
}
let stream = Arc::clone(&connection_info.stream);
tokio::spawn(async move {
let mut stream_guard = stream.lock().await;
let _ = stream_guard.shutdown().await;
});
}
if let Some(connection_info) = removed.as_ref() {
let client_role = ClientType::from_bytes(&connection_info.client_type)
.map(|client_type| client_type.as_str())
.unwrap_or("unknown");
info!(
"[connection_manager] connection dropped: role={} direction={} peer={}:{}",
client_role,
connection_type.as_str(),
ip,
port
);
}
removed
}
pub fn client_checkup(
stream: Arc<Mutex<TcpStream>>,
connection_type: ConnectionType,
ip: String,
port: u16,
command_map: Arc<Mutex<Command>>,
) {
tokio::spawn(async move {
let connection_key = ConnectionKey {
connection_type: connection_type.as_bytes(),
ip: ip_to_binary(&ip),
port,
};
let mut consecutive_failures = 0_u8;
loop {
sleep(Duration::from_secs(30)).await;
let still_monitoring_same_stream = {
let guard = CONNECTIONS.read().await;
guard
.as_ref()
.and_then(|conn| conn.connection_map.get(&connection_key))
.map(|connection_info| Arc::ptr_eq(&connection_info.stream, &stream))
.unwrap_or(false)
};
if !still_monitoring_same_stream {
break;
}
let message_type = RPC_BLOCK_HEIGHT; // Block-height request used as a lightweight checkup ping.
let (checkup_key, _checkup_tx, checkup_rx_mutex) =
reserve_entry(command_map.clone()).await;
// Send a lightweight ping message and wait for the reply
// routed back through the shared response hashmap.
let mut message: Vec<u8> = Vec::with_capacity(4);
message.push(message_type);
message.extend_from_slice(&checkup_key);
RpcResponse::send_raw(&stream, None, &message).await;
let response_result = {
let mut checkup_rx = checkup_rx_mutex.lock().await;
timeout(Duration::from_secs(30), checkup_rx.recv()).await
};
match response_result {
Ok(Some(_reply)) => {
consecutive_failures = 0;
info!(
"[connection_manager] liveness check ok: type={} peer={}:{}",
connection_type.as_str(),
ip,
port
);
}
_ => {
let still_monitoring_same_stream = {
let guard = CONNECTIONS.read().await;
guard
.as_ref()
.and_then(|conn| conn.connection_map.get(&connection_key))
.map(|connection_info| {
Arc::ptr_eq(&connection_info.stream, &stream)
})
.unwrap_or(false)
};
if !still_monitoring_same_stream {
delete_entry(command_map.clone(), checkup_key).await;
break;
}
consecutive_failures = consecutive_failures.saturating_add(1);
warn!(
"[connection_manager] liveness check failed: type={} peer={}:{} attempt={}/3",
connection_type.as_str(),
ip,
port,
consecutive_failures
);
delete_entry(command_map.clone(), checkup_key).await;
if consecutive_failures < 3 {
continue;
}
// Three consecutive timed-out or missing replies drop
// the connection, and outgoing peers trigger
// replacement discovery.
let mut guard = CONNECTIONS.write().await;
if let Some(conn) = guard.as_mut() {
let should_drop = conn
.connection_map
.get(&connection_key)
.map(|connection_info| {
Arc::ptr_eq(&connection_info.stream, &stream)
})
.unwrap_or(false);
if should_drop {
conn.drop_connection(connection_type, ip.clone(), port);
}
}
drop(guard);
if connection_type == ConnectionType::Outgoing {
spawn_retry_dropped_outgoing(ip.clone(), port);
}
break;
}
}
}
});
}
// Count active incoming peer connections.
pub fn count_incoming_connections(&self) -> usize {
self.connection_map
.values()
.filter(|info| {
ConnectionType::from_bytes(&info.connection_type) == Some(ConnectionType::Incoming)
})
.count()
}
// Count active outgoing peer connections.
pub fn count_outgoing_connections(&self) -> usize {
self.connection_map
.values()
.filter(|info| {
ConnectionType::from_bytes(&info.connection_type) == Some(ConnectionType::Outgoing)
})
.count()
}
// Return all live peer streams so broadcast-style paths can fan out
// messages without caring whether a peer was incoming or outgoing.
pub fn get_all_streams(&self) -> Vec<Arc<Mutex<TcpStream>>> {
self.connection_map
.values()
.filter(|connection_info| {
ClientType::from_bytes(&connection_info.client_type) == Some(ClientType::Miner)
})
.map(|connection_info| Arc::clone(&connection_info.stream))
.collect()
}
// Return all non-client peer streams so network-wide broadcasts can
// reach every reachable chain peer.
pub fn get_all_peer_streams(&self) -> Vec<Arc<Mutex<TcpStream>>> {
self.connection_map
.values()
.filter(|connection_info| {
ClientType::from_bytes(&connection_info.client_type) == Some(ClientType::Miner)
})
.map(|connection_info| Arc::clone(&connection_info.stream))
.collect()
}
// Resolve a stored outgoing node connection back to its live stream.
pub fn get_stream_for_outgoing(&self, ip: &str, port: u16) -> Option<Arc<Mutex<TcpStream>>> {
let ip_bytes = ip_to_binary(ip);
let connection_key = ConnectionKey {
connection_type: ConnectionType::Outgoing.as_bytes(),
ip: ip_bytes,
port,
};
self.connection_map
.get(&connection_key)
.filter(|info| ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner))
.map(|info| Arc::clone(&info.stream))
}
// Look up a live miner stream by the exact ip:port connection key.
// Network-map records only store bare IPs, so they must not be used
// to select an arbitrary live socket.
pub async fn get_stream_from_memory(key: &str) -> Option<Arc<Mutex<TcpStream>>> {
let (ip, port) = split_ip_port_key(key)?;
let lock = CONNECTIONS.read().await;
let conn = lock.as_ref()?;
let ip_bytes = ip_to_binary(&ip);
conn.connection_map
.iter()
.find_map(|(connection_key, info)| {
if connection_key.ip == ip_bytes
&& connection_key.port == port
&& ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner)
{
Some(Arc::clone(&info.stream))
} else {
None
}
})
}
pub async fn get_wallet_for_connection_key(key: &str) -> Option<Vec<u8>> {
let (ip, port) = split_ip_port_key(key)?;
let lock = CONNECTIONS.read().await;
let conn = lock.as_ref()?;
let ip_bytes = ip_to_binary(&ip);
conn.connection_map
.iter()
.find_map(|(connection_key, info)| {
if connection_key.ip == ip_bytes
&& connection_key.port == port
&& ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner)
{
Some(info.wallet_public_key.clone())
} else {
None
}
})
}
// Build the serialized connection key for a live stream when only
// the stream handle is known.
pub fn connection_key_for_stream(&self, stream: &Arc<Mutex<TcpStream>>) -> Option<String> {
self.connection_map
.iter()
.find_map(|(connection_key, connection_info)| {
if Arc::ptr_eq(&connection_info.stream, stream) {
let ip = binary_to_ip(connection_key.ip.clone());
Some(format!("{}:{}", ip, connection_key.port))
} else {
None
}
})
}
// Find the first stored connection record for the requested IP.
pub fn find_connection_info(&self, ip: &str) -> Option<(ConnectionType, u16)> {
let ip_bytes = ip_to_binary(ip);
for (key, _info) in self.connection_map.iter() {
if key.ip == ip_bytes {
let connection_type = ConnectionType::from_bytes(&key.connection_type)?;
return Some((connection_type, key.port));
}
}
None
}
// Find a stored connection by IP, constrained to a specific client role.
pub fn find_connection_info_by_client_type(
&self,
ip: &str,
client_type: ClientType,
) -> Option<(ConnectionType, u16)> {
let ip_bytes = ip_to_binary(ip);
let client_type_bytes = client_type.as_bytes();
for (key, info) in self.connection_map.iter() {
if key.ip == ip_bytes && info.client_type == client_type_bytes {
let connection_type = ConnectionType::from_bytes(&key.connection_type)?;
return Some((connection_type, key.port));
}
}
None
}
// Find the stored outgoing port for a peer IP so reconnect and
// cleanup logic can target the correct connection entry.
pub fn find_outgoing_port(&self, ip: &str) -> Option<u16> {
let ip_bytes = ip_to_binary(ip);
self.connection_map
.iter()
.find(|(key, _)| {
key.connection_type == ConnectionType::Outgoing.as_bytes() && key.ip == ip_bytes
})
.map(|(key, _)| key.port)
}
// Prefer a random incoming node connection, falling back to an
// outgoing node connection when no incoming peer is available.
pub fn get_random_connection(&self, excluded_key: Option<&str>) -> Option<(Vec<u8>, u16)> {
let mut rng = thread_rng();
let excluded = excluded_key.and_then(split_ip_port_key);
if let Some((key, _info)) = self
.connection_map
.iter()
.filter(|(key, info)| {
ConnectionType::from_bytes(&key.connection_type) == Some(ConnectionType::Incoming)
&& ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner)
&& excluded
.as_ref()
.map(|(ip, _)| key.ip != ip_to_binary(ip))
.unwrap_or(true)
})
.choose(&mut rng)
{
return Some((key.ip.clone(), key.port));
}
if let Some((key, _info)) = self
.connection_map
.iter()
.filter(|(key, info)| {
ConnectionType::from_bytes(&key.connection_type) == Some(ConnectionType::Outgoing)
&& ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner)
&& excluded
.as_ref()
.map(|(ip, _)| key.ip != ip_to_binary(ip))
.unwrap_or(true)
})
.choose(&mut rng)
{
return Some((key.ip.clone(), key.port));
}
None
}
}
fn spawn_monitor_update(ip: String, action: u8, peer_public_key: Vec<u8>, port: u16) {
tokio::spawn(async move {
let context = {
let guard = RECONNECT_CONTEXT.lock().await;
guard.clone()
};
let Some(context) = context else {
return;
};
let Some(monitored_address) = Wallet::public_key_bytes_to_short_address(&peer_public_key)
else {
return;
};
let monitoring_address = context.wallet.saved.short_address.clone();
if monitored_address == monitoring_address {
return;
}
let timestamp = crate::Utc::now().timestamp_millis() as u64;
let signature = NodeInfo::monitor_signature(
action,
&monitored_address,
&monitoring_address,
&ip,
timestamp,
&context.wallet,
)
.await;
let edit = SignedMonitorEdit {
action,
monitored_address,
monitoring_address,
target_ip: ip.clone(),
modified_timestamp: timestamp,
modified_signature: signature,
};
let params = MonitorAddressParams {
map: context.map.clone(),
edit,
remote_ip: String::new(),
db: context.db.clone(),
wallet: context.wallet.clone(),
connections_key: format!("{ip}:{port}"),
};
let _ = if action == MONITOR_ACTION_ADD {
NodeInfo::add_monitor(params).await
} else {
NodeInfo::remove_monitor(params).await
};
});
}
lazy_static! {
pub static ref CONNECTIONS: Arc<RwLock<Option<Connection>>> = Arc::new(RwLock::new(None));
}
pub async fn initialize_connection() {
// Lazily create the singleton connection manager the first time the
// node starts accepting or opening peer connections.
let mut connection_instance = CONNECTIONS.write().await;
if connection_instance.is_none() {
*connection_instance = Some(Connection::new());
}
}
pub async fn outgoing_connection_count() -> usize {
// Read the singleton connection manager and count live outgoing peers.
CONNECTIONS
.read()
.await
.as_ref()
.map(|connection| connection.count_outgoing_connections())
.unwrap_or(0)
}
pub async fn peer_connection_count() -> usize {
// Mining only needs proof that this node is connected to the live
// network; incoming and outgoing miner peers both satisfy that.
CONNECTIONS
.read()
.await
.as_ref()
.map(|connection| {
connection
.connection_map
.values()
.filter(|info| ClientType::from_bytes(&info.client_type) == Some(ClientType::Miner))
.count()
})
.unwrap_or(0)
}
pub async fn live_miner_peer_streams() -> Vec<(String, Arc<Mutex<TcpStream>>)> {
// Snapshot consensus and recovery checks vote only across currently
// connected miner peers, regardless of incoming/outgoing direction.
CONNECTIONS
.read()
.await
.as_ref()
.map(|connection| {
connection
.connection_map
.iter()
.filter_map(|(key, info)| {
if ClientType::from_bytes(&info.client_type) != Some(ClientType::Miner) {
return None;
}
let ip = binary_to_ip(key.ip.clone());
let connections_key = format!("{}:{}", ip, key.port);
Some((connections_key, Arc::clone(&info.stream)))
})
.collect()
})
.unwrap_or_default()
}
pub async fn get_client_type_from_memory(key: &str) -> Option<ClientType> {
// Recover the stored client role from the serialized connection key
// used throughout the RPC layer.
let (ip, port) = split_ip_port_key(key)?;
let ip_bytes = ip_to_binary(&ip);
let guard = CONNECTIONS.read().await;
let conn = guard.as_ref()?;
for (connection_key, info) in conn.connection_map.iter() {
if connection_key.ip == ip_bytes && connection_key.port == port {
return ClientType::from_bytes(&info.client_type);
}
}
None
}