Add custom user agent possibility

master
kirbylife 2021-08-30 14:29:46 -05:00
parent f6440ec7d5
commit 074ba2bc10
3 changed files with 60 additions and 1 deletions

13
build.rs 100644
View File

@ -0,0 +1,13 @@
use std::env;
fn main() {
println!("cargo:rerun-if-env-changed={}", "DEFAULT_USER_AGENT");
let user_agent = env::var("DEFAULT_USER_AGENT").unwrap_or(format!(
"Mozilla/5.0 ({}) Raspa/{}",
std::env::consts::OS,
env!("CARGO_PKG_VERSION")
));
println!("cargo:rustc-env=DEFAULT_USER_AGENT={}", user_agent);
}

View File

@ -1,9 +1,14 @@
use http::Method; use http::Method;
use std::collections::HashMap;
use url::ParseError; use url::ParseError;
use crate::into_url::IntoUrl; use crate::into_url::IntoUrl;
use crate::response; use crate::response;
// The DEFAULT_USER_AGENT is setted in the build.rs script
// and can be overwritten in compiled time
static DEFAULT_USER_AGENT: &str = env!("DEFAULT_USER_AGENT");
#[derive(Clone)] #[derive(Clone)]
pub enum ParameterType { pub enum ParameterType {
Number(isize), Number(isize),
@ -22,21 +27,30 @@ pub trait RequestBase {
#[derive(Debug)] #[derive(Debug)]
pub struct Request { pub struct Request {
pub url: url::Url, pub url: url::Url,
pub user_agent: String,
method: Method, method: Method,
params: Vec<(String, String)>, params: Vec<(String, String)>,
headers: HashMap<String, String>,
} }
impl RequestBase for Request { impl RequestBase for Request {
fn new<T: IntoUrl>(url: T) -> Result<Request, ParseError> { fn new<T: IntoUrl>(url: T) -> Result<Request, ParseError> {
url.into_url().map(|url_parsed| Request { url.into_url().map(|url_parsed| Request {
url: url_parsed, url: url_parsed,
user_agent: DEFAULT_USER_AGENT.to_string(),
method: Method::GET, method: Method::GET,
params: Vec::new(), params: Vec::new(),
headers: HashMap::new(),
}) })
} }
fn launch(mut self) -> response::Response { fn launch(mut self) -> response::Response {
let client = reqwest::blocking::Client::new(); let client = reqwest::blocking::Client::builder()
.user_agent(self.user_agent)
.build()
.unwrap();
// Set GET params
for (key, value) in &self.params { for (key, value) in &self.params {
self.url.query_pairs_mut().append_pair(key, value); self.url.query_pairs_mut().append_pair(key, value);
} }
@ -82,6 +96,11 @@ impl Request {
self self
} }
pub fn set_user_agent<S: AsRef<str>>(mut self, user_agent: S) -> Self {
self.user_agent = user_agent.as_ref().to_string();
self
}
pub fn add_params<S: AsRef<str>, T: ToString>(mut self, params: Vec<(S, T)>) -> Self { pub fn add_params<S: AsRef<str>, T: ToString>(mut self, params: Vec<(S, T)>) -> Self {
self.params = params self.params = params
.iter() .iter()

View File

@ -145,3 +145,30 @@ fn parameter_get_request() {
assert_eq!(resp["args"]["param_bool"].as_str(), Some("true")); assert_eq!(resp["args"]["param_bool"].as_str(), Some("true"));
assert_eq!(resp["args"]["param_float"].as_str(), Some("3.14159")); assert_eq!(resp["args"]["param_float"].as_str(), Some("3.14159"));
} }
#[test]
fn default_user_agent() {
let resp: Value = Request::new("https://httpbin.org/get")
.unwrap()
.launch()
.to_json()
.expect("Weird JSON parsed fail");
assert_eq!(
resp["headers"]["User-Agent"].as_str().unwrap(),
std::env::var("DEFAULT_USER_AGENT").unwrap()
);
}
#[test]
fn custom_user_agent() {
let resp: Value = Request::new("https://httpbin.org/get")
.unwrap()
.set_user_agent("test user agent")
.launch()
.to_json()
.expect("Weird JSON parsed fail");
assert_eq!(
resp["headers"]["User-Agent"].as_str(),
Some("test user agent")
);
}