use std::sync::Arc;
use http::{Method, Uri, header};
use parking_lot::RwLock;
use tracing::debug;
use crate::body::{self, Body as EncodedBody};
use crate::client::TorClient;
use crate::cookies::CookieJar;
use crate::error::{Error, Result};
use crate::redirect::{
RedirectAction, RedirectGuard, RedirectPolicy, redirect_method_for_status,
should_remove_auth_on_redirect,
};
use crate::response::Response;
pub struct Session {
client: TorClient,
cookies: Arc<RwLock<CookieJar>>,
redirect_policy: RedirectPolicy,
}
impl Session {
pub fn new(client: TorClient) -> Self {
Self {
client,
cookies: Arc::new(RwLock::new(CookieJar::new())),
redirect_policy: RedirectPolicy::default(),
}
}
#[must_use]
pub fn with_redirect_policy(mut self, policy: RedirectPolicy) -> Self {
self.redirect_policy = policy;
self
}
#[must_use]
pub fn with_standard_redirects(mut self) -> Self {
self.redirect_policy = RedirectPolicy::standard();
self
}
pub fn client(&self) -> &TorClient {
&self.client
}
pub fn cookies(&self) -> &Arc<RwLock<CookieJar>> {
&self.cookies
}
pub fn clear_cookies(&self) {
self.cookies.write().clear();
}
pub fn get(&self, url: &str) -> Result<SessionRequest<'_>> {
let uri: Uri = url
.parse()
.map_err(|_| Error::invalid_url(url, "invalid URL"))?;
Ok(SessionRequest::new(self, Method::GET, uri))
}
pub fn post(&self, url: &str) -> Result<SessionRequest<'_>> {
let uri: Uri = url
.parse()
.map_err(|_| Error::invalid_url(url, "invalid URL"))?;
Ok(SessionRequest::new(self, Method::POST, uri))
}
pub fn put(&self, url: &str) -> Result<SessionRequest<'_>> {
let uri: Uri = url
.parse()
.map_err(|_| Error::invalid_url(url, "invalid URL"))?;
Ok(SessionRequest::new(self, Method::PUT, uri))
}
pub fn delete(&self, url: &str) -> Result<SessionRequest<'_>> {
let uri: Uri = url
.parse()
.map_err(|_| Error::invalid_url(url, "invalid URL"))?;
Ok(SessionRequest::new(self, Method::DELETE, uri))
}
pub fn head(&self, url: &str) -> Result<SessionRequest<'_>> {
let uri: Uri = url
.parse()
.map_err(|_| Error::invalid_url(url, "invalid URL"))?;
Ok(SessionRequest::new(self, Method::HEAD, uri))
}
}
impl Clone for Session {
fn clone(&self) -> Self {
Self {
client: self.client.clone(),
cookies: Arc::clone(&self.cookies),
redirect_policy: self.redirect_policy,
}
}
}
pub struct SessionRequest<'a> {
session: &'a Session,
method: Method,
uri: Uri,
headers: http::HeaderMap,
body: Option<bytes::Bytes>,
timeout: Option<std::time::Duration>,
follow_redirects: bool,
}
impl<'a> SessionRequest<'a> {
fn new(session: &'a Session, method: Method, uri: Uri) -> Self {
Self {
session,
method,
uri,
headers: http::HeaderMap::new(),
body: None,
timeout: None,
follow_redirects: session.redirect_policy.should_follow(),
}
}
pub fn header<K, V>(mut self, key: K, value: V) -> Self
where
K: TryInto<http::header::HeaderName>,
V: TryInto<http::header::HeaderValue>,
{
if let (Ok(key), Ok(value)) = (key.try_into(), value.try_into()) {
self.headers.insert(key, value);
}
self
}
pub fn body(mut self, body: impl Into<bytes::Bytes>) -> Self {
self.body = Some(body.into());
self
}
pub fn json(mut self, json_str: impl Into<bytes::Bytes>) -> Self {
let encoded = EncodedBody::json(json_str);
self.headers
.insert(header::CONTENT_TYPE, encoded.content_type);
self.body = Some(encoded.data);
self
}
pub fn form<I, K, V>(mut self, params: I) -> Result<Self>
where
I: IntoIterator<Item = (K, V)>,
K: AsRef<str>,
V: AsRef<str>,
{
let encoded = EncodedBody::form(params)?;
self.headers
.insert(header::CONTENT_TYPE, encoded.content_type);
self.body = Some(encoded.data);
Ok(self)
}
pub fn basic_auth(mut self, username: &str, password: &str) -> Self {
let auth_header = body::basic_auth(username, password);
self.headers.insert(header::AUTHORIZATION, auth_header);
self
}
pub fn bearer_auth(self, token: &str) -> Self {
self.header(header::AUTHORIZATION, format!("Bearer {}", token))
}
pub fn timeout(mut self, timeout: std::time::Duration) -> Self {
self.timeout = Some(timeout);
self
}
pub fn no_redirect(mut self) -> Self {
self.follow_redirects = false;
self
}
pub async fn send(self) -> Result<Response> {
let mut current_uri = self.uri.clone();
let mut current_method = self.method.clone();
let mut redirect_guard = RedirectGuard::new(
self.session
.redirect_policy
.max_redirects()
.unwrap_or(u32::MAX),
);
let mut request_body = self.body.clone();
let mut custom_headers = self.headers.clone();
loop {
let mut builder = self
.session
.client
.request(current_method.clone(), ¤t_uri.to_string())?;
for (key, value) in &custom_headers {
builder = builder.header(key, value);
}
let cookie_header = {
let cookies = self.session.cookies.read();
cookies.get_cookie_header(¤t_uri)
};
if let Some(cookie_str) = cookie_header {
builder = builder.header(header::COOKIE, cookie_str);
}
if let Some(ref body) = request_body {
builder = builder.body(body.clone());
}
if let Some(timeout) = self.timeout {
builder = builder.timeout(timeout);
}
let response = builder.send().await?;
{
let cookies = self.session.cookies.write();
for cookie_str in response.headers().get_all(header::SET_COOKIE) {
if let Ok(cookie_str) = cookie_str.to_str() {
cookies.add_from_header(cookie_str, ¤t_uri);
}
}
}
if !self.follow_redirects {
return Ok(response);
}
match redirect_guard.check_redirect(&response, ¤t_uri)? {
RedirectAction::Follow(new_uri) => {
debug!("Following redirect: {} -> {}", current_uri, new_uri);
if should_remove_auth_on_redirect(¤t_uri, &new_uri) {
custom_headers.remove(header::AUTHORIZATION);
}
current_method = redirect_method_for_status(response.status(), ¤t_method);
if current_method == Method::GET {
request_body = None;
}
current_uri = new_uri;
}
RedirectAction::Stop => {
return Ok(response);
}
}
}
}
}
#[cfg(test)]
mod tests {
#![allow(clippy::unwrap_used, clippy::expect_used)]
use super::*;
#[test]
fn test_session_config() {
let policy = RedirectPolicy::Limited(5);
assert!(policy.should_follow());
assert_eq!(policy.max_redirects(), Some(5));
}
}