104 lines
3.4 KiB
Rust
104 lines
3.4 KiB
Rust
use crate::Bytes;
|
|
use anyhow::anyhow;
|
|
use log::{debug, trace, warn};
|
|
use once_cell::sync::OnceCell;
|
|
use reqwest::header::HeaderMap;
|
|
use reqwest::multipart::{Form, Part};
|
|
use reqwest::{Body, Client, IntoUrl, StatusCode};
|
|
use serde_json::Value;
|
|
use std::fmt::Display;
|
|
use std::time::Duration;
|
|
use tokio::time::sleep;
|
|
|
|
// note: only delete has rate-limit handling
|
|
|
|
static CLIENT: OnceCell<Client> = OnceCell::new();
|
|
|
|
pub async fn upload_webhook<T: Into<Body>>(
|
|
webhook: &str,
|
|
file: T,
|
|
filename: &str,
|
|
) -> anyhow::Result<(String, u64)> {
|
|
debug!("Uploading '{}' to Discord", filename);
|
|
let client = CLIENT.get_or_init(Client::new);
|
|
let form = Form::new().part("file", Part::stream(file).file_name(filename.to_string()));
|
|
let resp = client
|
|
.post(webhook)
|
|
.multipart(form)
|
|
.send()
|
|
.await?
|
|
.json::<Value>()
|
|
.await?;
|
|
trace!("Received JSON from Discord: {}", resp);
|
|
if let (Some(u), Some(i)) = (resp["attachments"][0]["url"].as_str(), resp["id"].as_str().and_then(|f| f.parse::<u64>().ok())) {
|
|
Ok((u.into(), i))
|
|
} else {
|
|
Err(anyhow!("Discord response didn't include the URL or message ID"))
|
|
}
|
|
}
|
|
|
|
fn extract_headers(h: &HeaderMap) -> Option<(u64, String)> {
|
|
let content_length = h
|
|
.get("content-length")?
|
|
.to_str()
|
|
.ok()?
|
|
.parse::<u64>()
|
|
.ok()?;
|
|
let mime = h.get("content-type")?.to_str().ok()?.to_string();
|
|
Some((content_length, mime))
|
|
}
|
|
|
|
pub async fn head<U: IntoUrl + Display>(url: U) -> anyhow::Result<(u64, String)> {
|
|
debug!("Downloading headers of '{}' from Discord", url);
|
|
let client = CLIENT.get_or_init(Client::new);
|
|
let resp = client.head(url).send().await?;
|
|
let headers = resp.headers();
|
|
if let Some(o) = extract_headers(headers) {
|
|
Ok(o)
|
|
} else {
|
|
Err(anyhow!("Discord response didn't include the URL"))
|
|
}
|
|
}
|
|
|
|
pub async fn get<U: IntoUrl + Display>(url: U) -> anyhow::Result<(u64, String, Bytes)> {
|
|
debug!("Downloading '{}' from Discord", url);
|
|
let client = CLIENT.get_or_init(Client::new);
|
|
let resp = client.get(url).send().await?;
|
|
let headers = extract_headers(resp.headers());
|
|
let bytes = resp.bytes().await?;
|
|
if let Some(o) = headers {
|
|
Ok((o.0, o.1, bytes))
|
|
} else {
|
|
Err(anyhow!("Discord response didn't include the URL"))
|
|
}
|
|
}
|
|
|
|
pub async fn delete(webhook: &str, mid: u64) -> anyhow::Result<()> {
|
|
debug!("Deleting message with ID {}", mid);
|
|
let client = CLIENT.get_or_init(Client::new);
|
|
let resp = client.delete(format!("{}/messages/{}", webhook, mid)).send().await?;
|
|
if resp.status() != StatusCode::NO_CONTENT {
|
|
Err(anyhow!(resp.text().await?))
|
|
} else {
|
|
let rt_header = resp.headers()
|
|
.get("X-RateLimit-Remaining")
|
|
.and_then(|v| v.to_str().ok()?.parse::<u64>().ok())
|
|
.and_then(|v| {
|
|
if v == 0 {
|
|
resp.headers()
|
|
.get("X-RateLimit-Reset-After")
|
|
.and_then(|v| v.to_str().ok()?.parse::<f64>().ok())
|
|
} else {
|
|
Some(0.0)
|
|
}
|
|
});
|
|
if let Some(rt) = rt_header {
|
|
if rt > 0.0 {
|
|
sleep(Duration::from_secs_f64(rt)).await;
|
|
}
|
|
} else {
|
|
warn!("Couldn't await the rate-limit, because there was a problem with the rate-limit header in Discord's response")
|
|
}
|
|
Ok(())
|
|
}
|
|
} |