Fix rate limiter for bots.

Use local-request state to cache User object.
This commit is contained in:
Paul
2021-09-18 20:06:00 +01:00
parent 0dcf8cb281
commit e5e0031cad
5 changed files with 60 additions and 78 deletions

View File

@@ -10,90 +10,72 @@ impl<'r> FromRequest<'r> for User {
type Error = rauth::util::Error;
async fn from_request(request: &'r Request<'_>) -> request::Outcome<Self, Self::Error> {
let header_bot_token = request
.headers()
.get("x-bot-token")
.next()
.map(|x| x.to_string());
let user: &Option<User> = request.local_cache_async(async {
let header_bot_token = request
.headers()
.get("x-bot-token")
.next()
.map(|x| x.to_string());
if let Some(bot_token) = header_bot_token {
return if let Ok(result) = get_collection("bots")
.find_one(
doc! {
"token": bot_token
},
None,
)
.await
{
if let Some(doc) = result {
let id = doc.get_str("_id").unwrap();
if let Ok(result) = get_collection("users")
.find_one(
doc! {
"_id": &id
},
None,
)
.await
{
if let Some(doc) = result {
Outcome::Success(from_document(doc).unwrap())
} else {
Outcome::Failure((
Status::Forbidden,
rauth::util::Error::InvalidSession,
))
}
} else {
Outcome::Failure((
Status::InternalServerError,
rauth::util::Error::DatabaseError {
operation: "find_one",
with: "user",
},
))
}
} else {
Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
}
} else {
Outcome::Failure((
Status::InternalServerError,
rauth::util::Error::DatabaseError {
operation: "find_one",
with: "bot",
},
))
};
} else {
if let Outcome::Success(session) = request.guard::<Session>().await {
if let Ok(result) = get_collection("users")
if let Some(bot_token) = header_bot_token {
if let Ok(result) = get_collection("bots")
.find_one(
doc! {
"_id": &session.user_id
"token": bot_token
},
None,
)
.await
{
if let Some(doc) = result {
Outcome::Success(from_document(doc).unwrap())
} else {
Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
let id = doc.get_str("_id").unwrap();
if let Ok(result) = get_collection("users")
.find_one(
doc! {
"_id": &id
},
None,
)
.await
{
if let Some(doc) = result {
if let Ok(user) = from_document(doc) {
return Some(user)
}
}
}
}
} else {
Outcome::Failure((
Status::InternalServerError,
rauth::util::Error::DatabaseError {
operation: "find_one",
with: "user",
},
))
}
} else {
Outcome::Failure((Status::Forbidden, rauth::util::Error::InvalidSession))
if let Outcome::Success(session) = request.guard::<Session>().await {
if let Ok(result) = get_collection("users")
.find_one(
doc! {
"_id": &session.user_id
},
None,
)
.await
{
if let Some(doc) = result {
if let Ok(user) = from_document(doc) {
return Some(user)
}
}
}
}
}
None
}).await;
if let Some(user) = user {
Outcome::Success(user.clone())
} else {
Outcome::Failure((
Status::Forbidden,
rauth::util::Error::InvalidSession,
))
}
}
}