Skip to content

Commit

Permalink
Specify User Agent in web socket
Browse files Browse the repository at this point in the history
  • Loading branch information
@h3poteto
h3poteto committed Jan 15, 2023
1 parent 64972bc commit 0019427
Show file tree
Hide file tree
Showing 4 changed files with 60 additions and 8 deletions.
10 changes: 9 additions & 1 deletion src/mastodon/mastodon.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -25,6 +25,7 @@ pub struct Mastodon {
client: APIClient,
base_url: String,
access_token: Option<String>,
user_agent: Option<String>,
}

impl Mastodon {
Expand All @@ -34,11 +35,12 @@ impl Mastodon {
access_token: Option<String>,
user_agent: Option<String>,
) -> Mastodon {
let client = APIClient::new(base_url.clone(), access_token.clone(), user_agent);
let client = APIClient::new(base_url.clone(), access_token.clone(), user_agent.clone());
Mastodon {
client,
base_url,
access_token,
user_agent,
}
}

Expand Down Expand Up @@ -2805,6 +2807,7 @@ impl megalodon::Megalodon for Mastodon {
String::from("user"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2817,6 +2820,7 @@ impl megalodon::Megalodon for Mastodon {
String::from("public"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2829,6 +2833,7 @@ impl megalodon::Megalodon for Mastodon {
String::from("public:local"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2841,6 +2846,7 @@ impl megalodon::Megalodon for Mastodon {
String::from("direct"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2857,6 +2863,7 @@ impl megalodon::Megalodon for Mastodon {
String::from("hashtag"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2873,6 +2880,7 @@ impl megalodon::Megalodon for Mastodon {
String::from("list"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand Down
24 changes: 21 additions & 3 deletions src/mastodon/web_socket.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,13 @@ use std::thread;
use std::time::Duration;

use super::entities;
use crate::default::DEFAULT_UA;
use crate::error::{Error, Kind};
use crate::streaming::{Message, Streaming};
use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::{
connect_async, tungstenite::protocol::frame::coding::CloseCode,
tungstenite::protocol::Message as WebSocketMessage,
Expand All @@ -23,6 +25,7 @@ pub struct WebSocket {
stream: String,
params: Option<Vec<String>>,
access_token: Option<String>,
user_agent: String,
}

#[derive(Deserialize)]
Expand All @@ -37,12 +40,19 @@ impl WebSocket {
stream: String,
params: Option<Vec<String>>,
access_token: Option<String>,
user_agent: Option<String>,
) -> Self {
let ua: String;
match user_agent {
Some(agent) => ua = agent,
None => ua = DEFAULT_UA.to_string(),
}
Self {
url,
stream,
params,
access_token,
user_agent: ua,
}
}

Expand Down Expand Up @@ -133,11 +143,19 @@ impl WebSocket {
url: &str,
callback: &Box<dyn Fn(Message) + Send + Sync>,
) -> Result<(), InnerError> {
let (mut socket, response) =
connect_async(Url::parse(url).unwrap()).await.map_err(|e| {
log::error!("Failed to connect: {}", e);
let mut req = Url::parse(url)
.unwrap()
.into_client_request()
.map_err(|e| {
log::error!("Failed to parse url: {}", e);
InnerError::new(InnerKind::ConnectionError)
})?;
req.headers_mut()
.insert("User-Agent", self.user_agent.parse().unwrap());
let (mut socket, response) = connect_async(req).await.map_err(|e| {
log::error!("Failed to connect: {}", e);
InnerError::new(InnerKind::ConnectionError)
})?;

log::debug!("Connected to {}", url);
log::debug!("Response HTTP code: {}", response.status());
Expand Down
10 changes: 9 additions & 1 deletion src/pleroma/pleroma.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,16 +26,18 @@ pub struct Pleroma {
client: APIClient,
base_url: String,
access_token: Option<String>,
user_agent: Option<String>,
}

impl Pleroma {
/// Create a new [`Pleroma`].
pub fn new(base_url: String, access_token: Option<String>, user_agent: Option<String>) -> Self {
let client = APIClient::new(base_url.clone(), access_token.clone(), user_agent);
let client = APIClient::new(base_url.clone(), access_token.clone(), user_agent.clone());
Self {
client,
base_url,
access_token,
user_agent,
}
}

Expand Down Expand Up @@ -2806,6 +2808,7 @@ impl megalodon::Megalodon for Pleroma {
String::from("user"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2818,6 +2821,7 @@ impl megalodon::Megalodon for Pleroma {
String::from("public"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2830,6 +2834,7 @@ impl megalodon::Megalodon for Pleroma {
String::from("public:local"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2842,6 +2847,7 @@ impl megalodon::Megalodon for Pleroma {
String::from("direct"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2858,6 +2864,7 @@ impl megalodon::Megalodon for Pleroma {
String::from("hashtag"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand All @@ -2874,6 +2881,7 @@ impl megalodon::Megalodon for Pleroma {
String::from("list"),
Some(params),
self.access_token.clone(),
self.user_agent.clone(),
);

Box::new(c)
Expand Down
24 changes: 21 additions & 3 deletions src/pleroma/web_socket.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3,11 +3,13 @@ use std::thread;
use std::time::Duration;

use super::entities;
use crate::default::DEFAULT_UA;
use crate::error::{Error, Kind};
use crate::streaming::{Message, Streaming};
use async_trait::async_trait;
use futures_util::{SinkExt, StreamExt};
use serde::Deserialize;
use tokio_tungstenite::tungstenite::client::IntoClientRequest;
use tokio_tungstenite::{
connect_async, tungstenite::protocol::frame::coding::CloseCode,
tungstenite::protocol::Message as WebSocketMessage,
Expand All @@ -23,6 +25,7 @@ pub struct WebSocket {
stream: String,
params: Option<Vec<String>>,
access_token: Option<String>,
user_agent: String,
}

#[derive(Deserialize)]
Expand All @@ -37,12 +40,19 @@ impl WebSocket {
stream: String,
params: Option<Vec<String>>,
access_token: Option<String>,
user_agent: Option<String>,
) -> Self {
let ua: String;
match user_agent {
Some(agent) => ua = agent,
None => ua = DEFAULT_UA.to_string(),
}
Self {
url,
stream,
params,
access_token,
user_agent: ua,
}
}

Expand Down Expand Up @@ -133,11 +143,19 @@ impl WebSocket {
url: &str,
callback: &Box<dyn Fn(Message) + Send + Sync>,
) -> Result<(), InnerError> {
let (mut socket, response) =
connect_async(Url::parse(url).unwrap()).await.map_err(|e| {
log::error!("Failed to connect: {}", e);
let mut req = Url::parse(url)
.unwrap()
.into_client_request()
.map_err(|e| {
log::error!("Failed to parse url: {}", e);
InnerError::new(InnerKind::ConnectionError)
})?;
req.headers_mut()
.insert("User-Agent", self.user_agent.parse().unwrap());
let (mut socket, response) = connect_async(req).await.map_err(|e| {
log::error!("Failed to connect: {}", e);
InnerError::new(InnerKind::ConnectionError)
})?;

log::debug!("Connected to {}", url);
log::debug!("Response HTTP code: {}", response.status());
Expand Down

0 comments on commit 0019427

Please sign in to comment.