90 lines
3.1 KiB
Rust
90 lines
3.1 KiB
Rust
|
use crate::records::memory::connections::get_client_type_from_memory;
|
||
|
|
use crate::records::memory::enums::ClientType;
|
||
|
|
use crate::rpc::command_maps::RPC_REPLY;
|
||
|
|
use crate::rpc::read_bytes_from_stream::read_first_byte;
|
||
|
|
use crate::rpc::server::connection_memory_manager::remove_stream_from_memory;
|
||
|
|
use crate::rpc::server::flood_protection::check_request_frequency_with_client_type;
|
||
|
|
use crate::rpc::server::structs::IncomingCommand;
|
||
|
|
use crate::log::warn;
|
||
|
|
use crate::sled::Db;
|
||
|
|
use crate::Arc;
|
||
|
|
use crate::Duration;
|
||
|
|
use crate::Mutex;
|
||
|
|
use crate::sleep;
|
||
|
|
use crate::TcpStream;
|
||
|
|
use crate::timeout;
|
||
|
|
|
||
|
|
async fn wait_for_stream_data(stream_locked: &Arc<Mutex<TcpStream>>) -> Result<bool, String> {
|
||
|
|
// Poll the socket with peek so the command byte stays in the stream
|
||
|
|
// until the normal byte reader consumes it.
|
||
|
|
let timeout_duration = Duration::from_millis(100);
|
||
|
|
|
||
|
|
loop {
|
||
|
|
let stream = stream_locked.lock().await;
|
||
|
|
let mut buffer = [0; 1];
|
||
|
|
|
||
|
|
match timeout(timeout_duration, async { stream.peek(&mut buffer).await }).await {
|
||
|
|
Ok(Ok(n)) => {
|
||
|
|
if n > 0 {
|
||
|
|
return Ok(true);
|
||
|
|
}
|
||
|
|
if stream.peer_addr().is_err() {
|
||
|
|
warn!("Dropped stream: {:?}", stream.peer_addr().unwrap_err());
|
||
|
|
drop(stream);
|
||
|
|
remove_stream_from_memory(stream_locked).await;
|
||
|
|
return Ok(false);
|
||
|
|
}
|
||
|
|
}
|
||
|
|
Ok(Err(_)) | Err(_) => {
|
||
|
|
drop(stream);
|
||
|
|
sleep(timeout_duration).await;
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
async fn peer_ip(stream_locked: &Arc<Mutex<TcpStream>>) -> String {
|
||
|
|
// Resolve the peer address once per incoming command so scoring and
|
||
|
|
// command handlers can use the same caller identity.
|
||
|
|
let stream = stream_locked.lock().await;
|
||
|
|
stream
|
||
|
|
.peer_addr()
|
||
|
|
.map(|addr| addr.ip().to_string())
|
||
|
|
.unwrap_or_else(|_| "unknown".into())
|
||
|
|
}
|
||
|
|
|
||
|
|
pub async fn next_incoming_command(
|
||
|
|
stream_locked: Arc<Mutex<TcpStream>>,
|
||
|
|
db: &Db,
|
||
|
|
connections_key: &str,
|
||
|
|
wallet_key: &str,
|
||
|
|
) -> Result<Option<IncomingCommand>, String> {
|
||
|
|
// A disconnected socket returns None so the caller can end the RPC
|
||
|
|
// loop without treating a clean disconnect as a command failure.
|
||
|
|
if !wait_for_stream_data(&stream_locked).await? {
|
||
|
|
return Ok(None);
|
||
|
|
}
|
||
|
|
|
||
|
|
let command =
|
||
|
|
read_first_byte(connections_key, stream_locked.clone()).await?;
|
||
|
|
let ip = peer_ip(&stream_locked).await;
|
||
|
|
|
||
|
|
// Connection memory is the source of truth for whether this stream
|
||
|
|
// belongs to a node/miner or a wallet-backed RPC client.
|
||
|
|
let client_type = get_client_type_from_memory(connections_key)
|
||
|
|
.await
|
||
|
|
.unwrap_or(ClientType::Miner);
|
||
|
|
|
||
|
|
// Replies belong to an existing request path, so only new inbound
|
||
|
|
// commands are counted against flood protection.
|
||
|
|
if command != RPC_REPLY {
|
||
|
|
check_request_frequency_with_client_type(db, ip.clone(), client_type, wallet_key).await;
|
||
|
|
}
|
||
|
|
|
||
|
|
Ok(Some(IncomingCommand {
|
||
|
|
command,
|
||
|
|
ip,
|
||
|
|
client_type,
|
||
|
|
}))
|
||
|
|
}
|