mirror of https://gitlab.com/famedly/conduit.git
Matthias Ahouansou
9 months ago
13 changed files with 1386 additions and 353 deletions
@ -0,0 +1,236 @@ |
|||||||
|
use std::{collections::HashMap, num::NonZeroU64}; |
||||||
|
|
||||||
|
use bytesize::ByteSize; |
||||||
|
use ruma::api::Metadata; |
||||||
|
use serde::Deserialize; |
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize)] |
||||||
|
pub struct Config { |
||||||
|
#[serde(flatten)] |
||||||
|
pub target: ConfigFragment, |
||||||
|
pub global: ConfigFragment, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize)] |
||||||
|
pub struct ConfigFragment { |
||||||
|
pub client: ConfigClientFragment, |
||||||
|
pub federation: ConfigFederationFragment, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize)] |
||||||
|
pub struct ConfigClientFragment { |
||||||
|
pub map: HashMap<ClientRestriction, RequestLimitation>, |
||||||
|
pub media: ClientMediaConfig, |
||||||
|
// TODO: Only have available on target, not global (same with most authenticated endpoints too maybe)?
|
||||||
|
pub authentication_failures: RequestLimitation, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Debug, Clone, Deserialize)] |
||||||
|
pub struct ConfigFederationFragment { |
||||||
|
pub map: HashMap<FederationRestriction, RequestLimitation>, |
||||||
|
pub media: FederationMediaConfig, |
||||||
|
} |
||||||
|
|
||||||
|
impl ConfigFragment { |
||||||
|
pub fn get(&self, restriction: &Restriction) -> &RequestLimitation { |
||||||
|
// Maybe look into https://github.com/moriyoshi-kasuga/enum-table
|
||||||
|
match restriction { |
||||||
|
Restriction::Client(client_restriction) => { |
||||||
|
self.client.map.get(client_restriction).unwrap() |
||||||
|
} |
||||||
|
Restriction::Federation(federation_restriction) => { |
||||||
|
self.federation.map.get(federation_restriction).unwrap() |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] |
||||||
|
pub enum Restriction { |
||||||
|
Client(ClientRestriction), |
||||||
|
Federation(FederationRestriction), |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd)] |
||||||
|
#[serde(rename_all = "snake_case")] |
||||||
|
pub enum ClientRestriction { |
||||||
|
Registration, |
||||||
|
Login, |
||||||
|
RegistrationTokenValidity, |
||||||
|
|
||||||
|
SendEvent, |
||||||
|
|
||||||
|
Join, |
||||||
|
Invite, |
||||||
|
Knock, |
||||||
|
|
||||||
|
SendReport, |
||||||
|
CreateAlias, |
||||||
|
|
||||||
|
MediaDownload, |
||||||
|
MediaCreate, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Deserialize, Eq, Hash, Ord, PartialEq, PartialOrd)] |
||||||
|
#[serde(rename_all = "snake_case")] |
||||||
|
pub enum FederationRestriction { |
||||||
|
Join, |
||||||
|
Knock, |
||||||
|
Invite, |
||||||
|
|
||||||
|
// Transactions should be handled by a completely dedicated rate-limiter
|
||||||
|
Transaction, |
||||||
|
|
||||||
|
MediaDownload, |
||||||
|
} |
||||||
|
|
||||||
|
impl TryFrom<Metadata> for Restriction { |
||||||
|
type Error = (); |
||||||
|
|
||||||
|
fn try_from(value: Metadata) -> Result<Self, Self::Error> { |
||||||
|
use Restriction::*; |
||||||
|
use ruma::api::{ |
||||||
|
IncomingRequest, |
||||||
|
client::{ |
||||||
|
account::{check_registration_token_validity, register}, |
||||||
|
alias::create_alias, |
||||||
|
authenticated_media::{ |
||||||
|
get_content, get_content_as_filename, get_content_thumbnail, get_media_preview, |
||||||
|
}, |
||||||
|
knock::knock_room, |
||||||
|
media::{self, create_content, create_mxc_uri}, |
||||||
|
membership::{invite_user, join_room_by_id, join_room_by_id_or_alias}, |
||||||
|
message::send_message_event, |
||||||
|
reporting::report_user, |
||||||
|
room::{report_content, report_room}, |
||||||
|
session::login, |
||||||
|
state::send_state_event, |
||||||
|
}, |
||||||
|
federation::{ |
||||||
|
authenticated_media::{ |
||||||
|
get_content as federation_get_content, |
||||||
|
get_content_thumbnail as federation_get_content_thumbnail, |
||||||
|
}, |
||||||
|
membership::{create_invite, create_join_event, create_knock_event}, |
||||||
|
}, |
||||||
|
}; |
||||||
|
|
||||||
|
Ok(match value { |
||||||
|
register::v3::Request::METADATA => Client(ClientRestriction::Registration), |
||||||
|
check_registration_token_validity::v1::Request::METADATA => { |
||||||
|
Client(ClientRestriction::RegistrationTokenValidity) |
||||||
|
} |
||||||
|
login::v3::Request::METADATA => Client(ClientRestriction::Login), |
||||||
|
send_message_event::v3::Request::METADATA | send_state_event::v3::Request::METADATA => { |
||||||
|
Client(ClientRestriction::SendEvent) |
||||||
|
} |
||||||
|
join_room_by_id::v3::Request::METADATA |
||||||
|
| join_room_by_id_or_alias::v3::Request::METADATA => Client(ClientRestriction::Join), |
||||||
|
invite_user::v3::Request::METADATA => Client(ClientRestriction::Invite), |
||||||
|
knock_room::v3::Request::METADATA => Client(ClientRestriction::Knock), |
||||||
|
report_user::v3::Request::METADATA |
||||||
|
| report_content::v3::Request::METADATA |
||||||
|
| report_room::v3::Request::METADATA => Client(ClientRestriction::SendReport), |
||||||
|
create_alias::v3::Request::METADATA => Client(ClientRestriction::CreateAlias), |
||||||
|
// NOTE: handle async media upload in a way that doesn't half the number of uploads you can do within a short timeframe, while not allowing pre-generation of MXC uris to allow uploading double the number of media at once
|
||||||
|
create_content::v3::Request::METADATA | create_mxc_uri::v1::Request::METADATA => { |
||||||
|
Client(ClientRestriction::MediaCreate) |
||||||
|
} |
||||||
|
// Unauthenticate media is deprecated
|
||||||
|
#[allow(deprecated)] |
||||||
|
media::get_content::v3::Request::METADATA |
||||||
|
| media::get_content_as_filename::v3::Request::METADATA |
||||||
|
| media::get_content_thumbnail::v3::Request::METADATA |
||||||
|
| media::get_media_preview::v3::Request::METADATA |
||||||
|
| get_content::v1::Request::METADATA |
||||||
|
| get_content_as_filename::v1::Request::METADATA |
||||||
|
| get_content_thumbnail::v1::Request::METADATA |
||||||
|
| get_media_preview::v1::Request::METADATA => Client(ClientRestriction::MediaDownload), |
||||||
|
federation_get_content::v1::Request::METADATA |
||||||
|
| federation_get_content_thumbnail::v1::Request::METADATA => { |
||||||
|
Federation(FederationRestriction::MediaDownload) |
||||||
|
} |
||||||
|
// v1 is deprecated
|
||||||
|
#[allow(deprecated)] |
||||||
|
create_join_event::v1::Request::METADATA | create_join_event::v2::Request::METADATA => { |
||||||
|
Federation(FederationRestriction::Join) |
||||||
|
} |
||||||
|
create_knock_event::v1::Request::METADATA => Federation(FederationRestriction::Knock), |
||||||
|
create_invite::v1::Request::METADATA | create_invite::v2::Request::METADATA => { |
||||||
|
Federation(FederationRestriction::Invite) |
||||||
|
} |
||||||
|
|
||||||
|
_ => return Err(()), |
||||||
|
}) |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Deserialize)] |
||||||
|
pub struct RequestLimitation { |
||||||
|
#[serde(flatten)] |
||||||
|
pub timeframe: Timeframe, |
||||||
|
pub burst_capacity: NonZeroU64, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Deserialize, Clone, Copy, Debug)] |
||||||
|
#[serde(rename_all = "snake_case")] |
||||||
|
// When deserializing, we want this prefix
|
||||||
|
#[allow(clippy::enum_variant_names)] |
||||||
|
pub enum Timeframe { |
||||||
|
PerSecond(NonZeroU64), |
||||||
|
PerMinute(NonZeroU64), |
||||||
|
PerHour(NonZeroU64), |
||||||
|
PerDay(NonZeroU64), |
||||||
|
} |
||||||
|
|
||||||
|
impl Timeframe { |
||||||
|
pub fn nano_gap(&self) -> u64 { |
||||||
|
match self { |
||||||
|
Timeframe::PerSecond(t) => 1000 * 1000 * 1000 / t.get(), |
||||||
|
Timeframe::PerMinute(t) => 1000 * 1000 * 1000 * 60 / t.get(), |
||||||
|
Timeframe::PerHour(t) => 1000 * 1000 * 1000 * 60 * 60 / t.get(), |
||||||
|
Timeframe::PerDay(t) => 1000 * 1000 * 1000 * 60 * 60 * 24 / t.get(), |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Deserialize)] |
||||||
|
pub struct ClientMediaConfig { |
||||||
|
pub download: MediaLimitation, |
||||||
|
pub upload: MediaLimitation, |
||||||
|
pub fetch: MediaLimitation, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Deserialize)] |
||||||
|
pub struct FederationMediaConfig { |
||||||
|
pub download: MediaLimitation, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Clone, Copy, Debug, Deserialize)] |
||||||
|
pub struct MediaLimitation { |
||||||
|
#[serde(flatten)] |
||||||
|
pub timeframe: MediaTimeframe, |
||||||
|
pub burst_capacity: ByteSize, |
||||||
|
} |
||||||
|
|
||||||
|
#[derive(Deserialize, Clone, Copy, Debug)] |
||||||
|
#[serde(rename_all = "snake_case")] |
||||||
|
// When deserializing, we want this prefix
|
||||||
|
#[allow(clippy::enum_variant_names)] |
||||||
|
pub enum MediaTimeframe { |
||||||
|
PerSecond(ByteSize), |
||||||
|
PerMinute(ByteSize), |
||||||
|
PerHour(ByteSize), |
||||||
|
PerDay(ByteSize), |
||||||
|
} |
||||||
|
|
||||||
|
impl MediaTimeframe { |
||||||
|
pub fn bytes_per_sec(&self) -> u64 { |
||||||
|
match self { |
||||||
|
MediaTimeframe::PerSecond(t) => t.as_u64(), |
||||||
|
MediaTimeframe::PerMinute(t) => t.as_u64() / 60, |
||||||
|
MediaTimeframe::PerHour(t) => t.as_u64() / (60 * 60), |
||||||
|
MediaTimeframe::PerDay(t) => t.as_u64() / (60 * 60 * 24), |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
@ -0,0 +1,581 @@ |
|||||||
|
use std::{ |
||||||
|
collections::{HashMap, hash_map::Entry}, |
||||||
|
net::IpAddr, |
||||||
|
sync::Arc, |
||||||
|
time::Duration, |
||||||
|
}; |
||||||
|
|
||||||
|
use conduit_config::{ |
||||||
|
Config, |
||||||
|
rate_limiting::{MediaLimitation, RequestLimitation, Restriction}, |
||||||
|
}; |
||||||
|
use ruma::{ |
||||||
|
OwnedServerName, OwnedUserId, UserId, |
||||||
|
api::{ |
||||||
|
Metadata, |
||||||
|
client::error::{ErrorKind, RetryAfter}, |
||||||
|
}, |
||||||
|
}; |
||||||
|
use tokio::{ |
||||||
|
sync::{Mutex, MutexGuard, RwLock}, |
||||||
|
time::Instant, |
||||||
|
}; |
||||||
|
|
||||||
|
use crate::{Error, Result, service::appservice::RegistrationInfo, services}; |
||||||
|
|
||||||
|
#[derive(Clone, Debug, Eq, Hash, Ord, PartialEq, PartialOrd)] |
||||||
|
pub enum Target { |
||||||
|
User(OwnedUserId), |
||||||
|
// Server endpoints should be rate-limited on a server and room basis
|
||||||
|
Server(OwnedServerName), |
||||||
|
Appservice { id: String, rate_limited: bool }, |
||||||
|
Ip(IpAddr), |
||||||
|
} |
||||||
|
|
||||||
|
impl Target { |
||||||
|
pub fn from_client_request( |
||||||
|
registration_info: Option<RegistrationInfo>, |
||||||
|
sender_user: &UserId, |
||||||
|
) -> Self { |
||||||
|
if let Some(info) = registration_info { |
||||||
|
// `rate_limited` only effects "masqueraded users", "The sender [user?] is excluded"
|
||||||
|
return Target::Appservice { |
||||||
|
id: info.registration.id, |
||||||
|
rate_limited: info.registration.rate_limited.unwrap_or(true) |
||||||
|
&& !(sender_user.server_name() == services().globals.server_name() |
||||||
|
&& info.registration.sender_localpart == sender_user.localpart()), |
||||||
|
}; |
||||||
|
} |
||||||
|
|
||||||
|
Target::User(sender_user.to_owned()) |
||||||
|
} |
||||||
|
|
||||||
|
pub fn from_client_request_optional_auth( |
||||||
|
registration_info: Option<RegistrationInfo>, |
||||||
|
sender_user: &Option<OwnedUserId>, |
||||||
|
ip_addr: Option<IpAddr>, |
||||||
|
) -> Option<Self> { |
||||||
|
if let Some(sender_user) = sender_user.as_ref() { |
||||||
|
Some(Self::from_client_request(registration_info, sender_user)) |
||||||
|
} else { |
||||||
|
ip_addr.map(Self::Ip) |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
fn rate_limited(&self) -> bool { |
||||||
|
match self { |
||||||
|
Target::User(user_id) => user_id != services().globals.server_user(), |
||||||
|
Target::Appservice { |
||||||
|
id: _, |
||||||
|
rate_limited, |
||||||
|
} => *rate_limited, |
||||||
|
_ => true, |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
pub fn is_authenticated(&self) -> bool { |
||||||
|
!matches!(self, Target::Ip(_)) |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
// NOTE: still chokes on the global mutex around the map, which in theory only needed for inserts,
|
||||||
|
// but due to rust ownership model I don't think it's possible to (easily) only require to lock
|
||||||
|
// said mutex for inserts (and have a rwlock for removals if memory usage ever becomes a problem).
|
||||||
|
type MediaBucket = Mutex<HashMap<Target, Arc<Mutex<Instant>>>>; |
||||||
|
type GlobalMediaBucket = Arc<Mutex<Instant>>; |
||||||
|
|
||||||
|
type RequestPermittedAfter = Arc<Mutex<Instant>>; |
||||||
|
|
||||||
|
pub struct Service { |
||||||
|
buckets: Mutex<HashMap<(Target, Restriction), RequestPermittedAfter>>, |
||||||
|
global_bucket: Mutex<HashMap<Restriction, RequestPermittedAfter>>, |
||||||
|
|
||||||
|
media_upload: MediaBucket, |
||||||
|
media_fetch: MediaBucket, |
||||||
|
media_download: MediaBucket, |
||||||
|
|
||||||
|
global_media_upload: GlobalMediaBucket, |
||||||
|
global_media_fetch: GlobalMediaBucket, |
||||||
|
global_media_download_client: GlobalMediaBucket, |
||||||
|
global_media_download_federation: GlobalMediaBucket, |
||||||
|
|
||||||
|
authentication_failures: RwLock<HashMap<IpAddr, Arc<RwLock<Instant>>>>, |
||||||
|
} |
||||||
|
|
||||||
|
impl Service { |
||||||
|
pub fn build(config: &Config) -> Arc<Self> { |
||||||
|
let now = Instant::now(); |
||||||
|
let global_media_config = &config.rate_limiting.global; |
||||||
|
|
||||||
|
Arc::new(Self { |
||||||
|
buckets: Mutex::new(HashMap::new()), |
||||||
|
global_bucket: Mutex::new(HashMap::new()), |
||||||
|
|
||||||
|
media_upload: Mutex::new(HashMap::new()), |
||||||
|
media_fetch: Mutex::new(HashMap::new()), |
||||||
|
media_download: Mutex::new(HashMap::new()), |
||||||
|
|
||||||
|
global_media_upload: default_media_entry(global_media_config.client.media.upload, now), |
||||||
|
global_media_fetch: default_media_entry(global_media_config.client.media.fetch, now), |
||||||
|
global_media_download_client: default_media_entry( |
||||||
|
global_media_config.client.media.download, |
||||||
|
now, |
||||||
|
), |
||||||
|
global_media_download_federation: default_media_entry( |
||||||
|
global_media_config.federation.media.download, |
||||||
|
now, |
||||||
|
), |
||||||
|
|
||||||
|
authentication_failures: RwLock::new(HashMap::new()), |
||||||
|
}) |
||||||
|
} |
||||||
|
|
||||||
|
//TODO: use checked and saturating arithmetic
|
||||||
|
|
||||||
|
/// Takes the target and request, and either accepts the request while adding to the
|
||||||
|
/// bucket, or rejects the request, returning the duration that should be waited until
|
||||||
|
/// the request should be retried.
|
||||||
|
pub async fn check(&self, target: Option<Target>, request: Metadata) -> Result<()> { |
||||||
|
let Ok(restriction) = request.try_into() else { |
||||||
|
// Endpoint has no associated restriction
|
||||||
|
return Ok(()); |
||||||
|
}; |
||||||
|
let arrival = Instant::now(); |
||||||
|
|
||||||
|
let config = services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.global |
||||||
|
.get(&restriction); |
||||||
|
|
||||||
|
let mut map = self.global_bucket.lock().await; |
||||||
|
|
||||||
|
let entry = map.entry(restriction); |
||||||
|
let proposed_entry = match &entry { |
||||||
|
Entry::Occupied(occupied_entry) => { |
||||||
|
let entry = Arc::clone(occupied_entry.get()); |
||||||
|
let entry = entry.lock().await; |
||||||
|
|
||||||
|
if arrival.checked_duration_since(*entry).is_none() { |
||||||
|
return instant_to_err(&entry); |
||||||
|
} |
||||||
|
|
||||||
|
let min_instant = arrival |
||||||
|
- Duration::from_nanos( |
||||||
|
config.timeframe.nano_gap() * config.burst_capacity.get(), |
||||||
|
); |
||||||
|
entry.max(min_instant) + Duration::from_nanos(config.timeframe.nano_gap()) |
||||||
|
} |
||||||
|
Entry::Vacant(_) => { |
||||||
|
arrival |
||||||
|
- Duration::from_nanos( |
||||||
|
config.timeframe.nano_gap() * (config.burst_capacity.get() - 1), |
||||||
|
) |
||||||
|
} |
||||||
|
}; |
||||||
|
|
||||||
|
if let Some(target) = target { |
||||||
|
let config = services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
.get(&restriction); |
||||||
|
|
||||||
|
let mut map = self.buckets.lock().await; |
||||||
|
let entry = map.entry((target, restriction)); |
||||||
|
match entry { |
||||||
|
Entry::Occupied(occupied_entry) => { |
||||||
|
let entry = Arc::clone(occupied_entry.get()); |
||||||
|
let mut entry = entry.lock().await; |
||||||
|
|
||||||
|
if arrival.checked_duration_since(*entry).is_none() { |
||||||
|
return instant_to_err(&entry); |
||||||
|
} |
||||||
|
|
||||||
|
let min_instant = arrival |
||||||
|
- Duration::from_nanos( |
||||||
|
config.timeframe.nano_gap() * config.burst_capacity.get(), |
||||||
|
); |
||||||
|
*entry = |
||||||
|
entry.max(min_instant) + Duration::from_nanos(config.timeframe.nano_gap()); |
||||||
|
} |
||||||
|
Entry::Vacant(vacant_entry) => { |
||||||
|
vacant_entry.insert(Arc::new(Mutex::new( |
||||||
|
arrival |
||||||
|
- Duration::from_nanos( |
||||||
|
config.timeframe.nano_gap() * (config.burst_capacity.get() - 1), |
||||||
|
), |
||||||
|
))); |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
entry.insert_entry(Arc::new(Mutex::new(proposed_entry))); |
||||||
|
|
||||||
|
Ok(()) |
||||||
|
} |
||||||
|
|
||||||
|
pub async fn check_media_download(&self, target: Option<Target>, size: u64) -> Result<()> { |
||||||
|
// All targets besides servers use the client-server API
|
||||||
|
let (target_limitation, global_limitation, global_bucket) = |
||||||
|
if let Some(Target::Server(_)) = &target { |
||||||
|
( |
||||||
|
services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
.federation |
||||||
|
.media |
||||||
|
.download, |
||||||
|
services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.global |
||||||
|
.federation |
||||||
|
.media |
||||||
|
.download, |
||||||
|
&self.global_media_download_federation, |
||||||
|
) |
||||||
|
} else { |
||||||
|
( |
||||||
|
services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
.client |
||||||
|
.media |
||||||
|
.download, |
||||||
|
services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.global |
||||||
|
.client |
||||||
|
.media |
||||||
|
.download, |
||||||
|
&self.global_media_download_client, |
||||||
|
) |
||||||
|
}; |
||||||
|
|
||||||
|
check_media( |
||||||
|
target, |
||||||
|
size, |
||||||
|
target_limitation, |
||||||
|
global_limitation, |
||||||
|
&self.media_download, |
||||||
|
global_bucket, |
||||||
|
) |
||||||
|
.await |
||||||
|
} |
||||||
|
|
||||||
|
pub async fn check_media_upload(&self, target: Target, size: u64) -> Result<()> { |
||||||
|
let target_limitation = services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
// Media can only be uploaded on the client-server API
|
||||||
|
.client |
||||||
|
.media |
||||||
|
.upload; |
||||||
|
|
||||||
|
let global_limitation = services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.global |
||||||
|
// Media can only be uploaded on the client-server API
|
||||||
|
.client |
||||||
|
.media |
||||||
|
.upload; |
||||||
|
|
||||||
|
check_media( |
||||||
|
Some(target), |
||||||
|
size, |
||||||
|
target_limitation, |
||||||
|
global_limitation, |
||||||
|
&self.media_upload, |
||||||
|
&self.global_media_upload, |
||||||
|
) |
||||||
|
.await |
||||||
|
} |
||||||
|
|
||||||
|
pub async fn check_media_pre_fetch(&self, target: &Target) -> Result<()> { |
||||||
|
if !target.rate_limited() { |
||||||
|
return Ok(()); |
||||||
|
} |
||||||
|
|
||||||
|
let arrival = Instant::now(); |
||||||
|
|
||||||
|
let check = async |map: &MediaBucket, global_bucket: &GlobalMediaBucket| { |
||||||
|
let map = map.lock().await; |
||||||
|
if let Some(mutex) = map.get(target) { |
||||||
|
let mutex = mutex.lock().await; |
||||||
|
|
||||||
|
if arrival.checked_duration_since(*mutex).is_none() { |
||||||
|
return instant_to_err(&mutex); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
let global_bucket = global_bucket.lock().await; |
||||||
|
|
||||||
|
if arrival.checked_duration_since(*global_bucket).is_none() { |
||||||
|
return instant_to_err(&global_bucket); |
||||||
|
} |
||||||
|
|
||||||
|
Ok(()) |
||||||
|
}; |
||||||
|
|
||||||
|
// checking fetch
|
||||||
|
check(&self.media_fetch, &self.global_media_fetch).await?; |
||||||
|
|
||||||
|
// checking download as well
|
||||||
|
check(&self.media_download, &self.global_media_download_client).await |
||||||
|
} |
||||||
|
|
||||||
|
/// Checks whether the ip address is has been rate limited due to too many bad access tokens being sent.
|
||||||
|
pub async fn pre_auth_check(&self, ip_addr: IpAddr) -> Result<()> { |
||||||
|
let arrival = Instant::now(); |
||||||
|
|
||||||
|
if let Some(instant) = self.authentication_failures.read().await.get(&ip_addr) { |
||||||
|
let instant = instant.read().await; |
||||||
|
|
||||||
|
if arrival.checked_duration_since(*instant).is_none() { |
||||||
|
return instant_to_err(&instant); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
Ok(()) |
||||||
|
} |
||||||
|
|
||||||
|
/// Updates the bad auth rate limiter when a bad access token is sent where access tokens auth is an option.
|
||||||
|
pub async fn update_post_auth_failure(&self, ip_addr: IpAddr) { |
||||||
|
let arrival = Instant::now(); |
||||||
|
|
||||||
|
let RequestLimitation { |
||||||
|
timeframe, |
||||||
|
burst_capacity, |
||||||
|
} = services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
.client |
||||||
|
.authentication_failures; |
||||||
|
|
||||||
|
let mut map = self.authentication_failures.write().await; |
||||||
|
let entry = map.entry(ip_addr); |
||||||
|
|
||||||
|
match entry { |
||||||
|
Entry::Occupied(occupied_entry) => { |
||||||
|
let entry = Arc::clone(occupied_entry.get()); |
||||||
|
let mut entry = entry.write().await; |
||||||
|
|
||||||
|
let min_instant = |
||||||
|
arrival - Duration::from_nanos(timeframe.nano_gap() * burst_capacity.get()); |
||||||
|
*entry = entry.max(min_instant) + Duration::from_nanos(timeframe.nano_gap()); |
||||||
|
} |
||||||
|
Entry::Vacant(vacant_entry) => { |
||||||
|
vacant_entry.insert(Arc::new(RwLock::new( |
||||||
|
arrival - Duration::from_nanos(burst_capacity.get() / timeframe.nano_gap()), |
||||||
|
))); |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
pub async fn update_media_post_fetch(&self, target: Target, size: u64) { |
||||||
|
if !target.rate_limited() { |
||||||
|
return; |
||||||
|
} |
||||||
|
|
||||||
|
let arrival = Instant::now(); |
||||||
|
|
||||||
|
let update = async |map: &MediaBucket, |
||||||
|
target_limitation: &MediaLimitation, |
||||||
|
global_bucket: &GlobalMediaBucket, |
||||||
|
global_limitation: &MediaLimitation| { |
||||||
|
let mut map = map.lock().await; |
||||||
|
let entry = map.entry(target.clone()); |
||||||
|
|
||||||
|
match entry { |
||||||
|
Entry::Occupied(occupied_entry) => { |
||||||
|
let entry = Arc::clone(occupied_entry.get()); |
||||||
|
|
||||||
|
let _ = |
||||||
|
update_media_entry(size, target_limitation, &arrival, entry, false).await; |
||||||
|
} |
||||||
|
Entry::Vacant(vacant_entry) => { |
||||||
|
vacant_entry.insert(Arc::new(Mutex::new( |
||||||
|
arrival |
||||||
|
- Duration::from_nanos( |
||||||
|
target_limitation.burst_capacity.as_u64() |
||||||
|
/ target_limitation.timeframe.bytes_per_sec(), |
||||||
|
), |
||||||
|
))); |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
let _ = update_media_entry( |
||||||
|
size, |
||||||
|
global_limitation, |
||||||
|
&arrival, |
||||||
|
Arc::clone(global_bucket), |
||||||
|
false, |
||||||
|
) |
||||||
|
.await; |
||||||
|
}; |
||||||
|
|
||||||
|
// updating fetch
|
||||||
|
update( |
||||||
|
&self.media_fetch, |
||||||
|
&services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
.client |
||||||
|
.media |
||||||
|
.fetch, |
||||||
|
&self.global_media_fetch, |
||||||
|
&services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.global |
||||||
|
.client |
||||||
|
.media |
||||||
|
.fetch, |
||||||
|
) |
||||||
|
.await; |
||||||
|
|
||||||
|
// updating download as well
|
||||||
|
update( |
||||||
|
&self.media_download, |
||||||
|
&services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.target |
||||||
|
.client |
||||||
|
.media |
||||||
|
.download, |
||||||
|
&self.global_media_download_client, |
||||||
|
&services() |
||||||
|
.globals |
||||||
|
.config |
||||||
|
.rate_limiting |
||||||
|
.global |
||||||
|
.client |
||||||
|
.media |
||||||
|
.download, |
||||||
|
) |
||||||
|
.await; |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
async fn update_media_entry( |
||||||
|
size: u64, |
||||||
|
limitation: &MediaLimitation, |
||||||
|
arrival: &Instant, |
||||||
|
entry: Arc<Mutex<Instant>>, |
||||||
|
and_check: bool, |
||||||
|
) -> Result<()> { |
||||||
|
let mut entry = entry.lock().await; |
||||||
|
|
||||||
|
//TODO: use more precise conversion than secs
|
||||||
|
let proposed_entry = get_proposed_entry(size, limitation, arrival, &entry, and_check)?; |
||||||
|
|
||||||
|
*entry = proposed_entry; |
||||||
|
|
||||||
|
Ok(()) |
||||||
|
} |
||||||
|
|
||||||
|
fn get_proposed_entry( |
||||||
|
size: u64, |
||||||
|
limitation: &MediaLimitation, |
||||||
|
arrival: &Instant, |
||||||
|
entry: &MutexGuard<'_, Instant>, |
||||||
|
and_check: bool, |
||||||
|
) -> Result<Instant> { |
||||||
|
let min_instant = *arrival |
||||||
|
- Duration::from_secs( |
||||||
|
limitation.burst_capacity.as_u64() / limitation.timeframe.bytes_per_sec(), |
||||||
|
); |
||||||
|
|
||||||
|
let proposed_entry = |
||||||
|
entry.max(min_instant) + Duration::from_secs(size / limitation.timeframe.bytes_per_sec()); |
||||||
|
|
||||||
|
if and_check && arrival.checked_duration_since(proposed_entry).is_none() { |
||||||
|
return instant_to_err(&proposed_entry).map(|_| proposed_entry); |
||||||
|
} |
||||||
|
|
||||||
|
Ok(proposed_entry) |
||||||
|
} |
||||||
|
|
||||||
|
async fn check_media( |
||||||
|
target: Option<Target>, |
||||||
|
size: u64, |
||||||
|
target_limitation: MediaLimitation, |
||||||
|
global_limitation: MediaLimitation, |
||||||
|
target_map: &MediaBucket, |
||||||
|
global_bucket: &GlobalMediaBucket, |
||||||
|
) -> Result<()> { |
||||||
|
if !target.as_ref().is_some_and(Target::rate_limited) { |
||||||
|
return Ok(()); |
||||||
|
} |
||||||
|
|
||||||
|
let arrival = Instant::now(); |
||||||
|
|
||||||
|
let mut global_bucket = global_bucket.lock().await; |
||||||
|
let proposed = get_proposed_entry(size, &global_limitation, &arrival, &global_bucket, true)?; |
||||||
|
|
||||||
|
if let Some(target) = target { |
||||||
|
let mut map = target_map.lock().await; |
||||||
|
let entry = map.entry(target); |
||||||
|
|
||||||
|
match entry { |
||||||
|
Entry::Occupied(occupied_entry) => { |
||||||
|
let entry = Arc::clone(occupied_entry.get()); |
||||||
|
|
||||||
|
update_media_entry(size, &target_limitation, &arrival, entry, true).await?; |
||||||
|
} |
||||||
|
Entry::Vacant(vacant_entry) => { |
||||||
|
vacant_entry.insert(default_media_entry(target_limitation, arrival)); |
||||||
|
} |
||||||
|
} |
||||||
|
} |
||||||
|
|
||||||
|
*global_bucket = proposed; |
||||||
|
|
||||||
|
Ok(()) |
||||||
|
} |
||||||
|
|
||||||
|
fn default_media_entry( |
||||||
|
target_limitation: MediaLimitation, |
||||||
|
arrival: Instant, |
||||||
|
) -> Arc<Mutex<Instant>> { |
||||||
|
Arc::new(Mutex::new( |
||||||
|
arrival |
||||||
|
- Duration::from_nanos( |
||||||
|
target_limitation.burst_capacity.as_u64() |
||||||
|
/ target_limitation.timeframe.bytes_per_sec(), |
||||||
|
), |
||||||
|
)) |
||||||
|
} |
||||||
|
|
||||||
|
fn instant_to_err(instant: &Instant) -> Result<()> { |
||||||
|
let now = Instant::now(); |
||||||
|
|
||||||
|
Err(Error::BadRequest( |
||||||
|
ErrorKind::LimitExceeded { |
||||||
|
// Not using ::DateTime because conversion from Instant to SystemTime is convoluted
|
||||||
|
retry_after: Some(RetryAfter::Delay(instant.duration_since(now))), |
||||||
|
}, |
||||||
|
"Rate limit exceeded", |
||||||
|
)) |
||||||
|
} |
||||||
Loading…
Reference in new issue