forked from jmug/stoatchat
Fix rate limiter for bots.
Use local-request state to cache User object.
This commit is contained in:
@@ -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,
|
||||
))
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user