use isbot::Bots;
use actix_web::{
dev::Service,
http::header::{HeaderMap, USER_AGENT},
web, App, HttpRequest, HttpResponse, HttpServer,
};
use futures::{future, future::Either, future::FutureExt};
struct AppState {
bots: Bots,
}
async fn index(_: web::Data<AppState>) -> HttpResponse {
HttpResponse::Ok().body("Home")
}
async fn login(req: HttpRequest, data: web::Data<AppState>) -> HttpResponse {
if let Some(user_agent) = get_user_agent(req.headers()) {
if data.bots.is_bot(user_agent) {
return HttpResponse::Forbidden().body("Bots not allowed");
}
}
HttpResponse::Ok().body("Login")
}
fn get_user_agent(header_map: &HeaderMap) -> Option<&str> {
header_map.get(USER_AGENT)?.to_str().ok()
}
#[actix_web::main]
async fn main() -> std::io::Result<()> {
HttpServer::new(move || {
App::new()
.app_data(AppState {
bots: Bots::default(),
})
.wrap_fn(|sreq, srv| {
if let Some(data) = sreq.app_data::<web::Data<AppState>>() {
if let Some(user_agent) = get_user_agent(sreq.headers()) {
if data.bots.is_bot(user_agent) {
return Either::Right(future::ready(Ok(sreq.into_response(
HttpResponse::Forbidden().body("Bots not allowed"),
))));
}
}
}
Either::Left(srv.call(sreq).map(|res| res))
})
.route("/", web::get().to(index))
.route("/login", web::get().to(login))
})
.bind(("127.0.0.1", 8080))?
.run()
.await
}
#[cfg(test)]
mod tests {
use super::*;
use actix_web::{http::StatusCode, test, web, App};
use bytes::Bytes;
const VALID_BROWSER_USER_AGENT: &str = "Mozilla/5.0 (Macintosh; Intel Mac OS X 10_15_7) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/96.0.4664.110 Safari/537.36";
const KNOWN_BOT_USER_AGENT: &str = "Mozilla/5.0 (iPhone; CPU iPhone OS 9_1 like Mac OS X) AppleWebKit/601.1.46 (KHTML, like Gecko) Version/9.0 Mobile/13B143 Safari/601.1 (compatible; AdsBot-Google-Mobile; +http://www.google.com/mobile/adsbot.html)";
const CHROME_LIGHTHOUSE_BOT_USER_AGENT: &str = "Mozilla/5.0 (Linux; Android 7.0; Moto G (4)) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/84.0.4143.7 Mobile Safari/537.36 Chrome-Lighthouse";
async fn account(req: HttpRequest, data: web::Data<AppState>) -> HttpResponse {
if let Some(user_agent) = get_user_agent(&req.headers()) {
if data.bots.is_bot(user_agent) {
return HttpResponse::Forbidden().body("Bots not allowed");
}
}
HttpResponse::Ok().body("Account")
}
#[actix_rt::test]
async fn test_valid_browser() {
let mut app = test::init_service(
App::new()
.app_data(web::Data::new(AppState {
bots: Bots::default(),
}))
.route("/account", web::get().to(account)),
)
.await;
let req = test::TestRequest::with_uri("/account")
.insert_header((USER_AGENT, VALID_BROWSER_USER_AGENT))
.to_request();
let res = test::call_service(&mut app, req).await;
assert!(res.status().is_success());
let result = test::read_body(res).await;
assert_eq!(result, Bytes::from_static(b"Account"))
}
#[actix_rt::test]
async fn test_known_bot() {
let mut app = test::init_service(
App::new()
.app_data(web::Data::new(AppState {
bots: Bots::default(),
}))
.route("/account", web::get().to(account)),
)
.await;
let req = test::TestRequest::with_uri("/account")
.insert_header((USER_AGENT, KNOWN_BOT_USER_AGENT))
.to_request();
let resp = test::call_service(&mut app, req).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[actix_rt::test]
async fn test_exclude_known_bot() {
let mut bots = Bots::default();
bots.remove(&["Chrome-Lighthouse"]);
let data = web::Data::new(AppState { bots });
let mut app = test::init_service(
App::new()
.app_data(data.clone())
.route("/account", web::get().to(account)),
)
.await;
let req = test::TestRequest::with_uri("/account")
.insert_header((USER_AGENT, CHROME_LIGHTHOUSE_BOT_USER_AGENT))
.to_request();
let resp = test::call_service(&mut app, req).await;
assert!(resp.status().is_success());
}
#[actix_rt::test]
async fn test_middleware_known_bot() {
let mut app = test::init_service(
App::new()
.app_data(web::Data::new(AppState {
bots: Bots::default(),
}))
.wrap_fn(|sreq, srv| {
if let Some(data) = sreq.app_data::<web::Data<AppState>>() {
if let Some(user_agent) = get_user_agent(sreq.headers()) {
if data.bots.is_bot(user_agent) {
return Either::Right(future::ready(Ok(sreq.into_response(
HttpResponse::Forbidden().body("Bots not allowed"),
))));
}
}
}
Either::Left(srv.call(sreq).map(|res| res))
})
.route("/account", web::get().to(account)),
)
.await;
let req = test::TestRequest::with_uri("/account")
.insert_header((USER_AGENT, KNOWN_BOT_USER_AGENT))
.to_request();
let resp = test::call_service(&mut app, req).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
#[actix_rt::test]
async fn test_post_known_bot() {
let mut app = test::init_service(
App::new()
.app_data(web::Data::new(AppState {
bots: Bots::default(),
}))
.route("/account", web::post().to(account)),
)
.await;
let req = test::TestRequest::post()
.uri("/account")
.insert_header((USER_AGENT, KNOWN_BOT_USER_AGENT))
.to_request();
let resp = test::call_service(&mut app, req).await;
assert_eq!(resp.status(), StatusCode::FORBIDDEN);
}
}