mirror of
https://github.com/EFForg/rayhunter.git
synced 2026-07-26 01:08:10 -07:00
Add support for WebDAV (#982)
* Add support for WebDAV * Fix get_unuploaded_entries_with_age to use start_time * Use better pattern when matching join! result Co-authored-by: Markus Unterwaditzer <markus-github@unterwaditzer.net> * Refactor Webdav with better string ownership * Unformat Cargo.toml * Add timeout config * Use a single url config instead of host and path * QMDL store returns single unuploaded entry at a time * Reset Cargo.lock --------- Co-authored-by: Markus Unterwaditzer <markus-github@unterwaditzer.net>
This commit is contained in:
John
@@ -0,0 +1,446 @@
|
||||
use std::fmt::Display;
|
||||
use std::{sync::Arc, time::Duration};
|
||||
|
||||
use chrono::TimeDelta;
|
||||
use log::{info, warn};
|
||||
use reqwest::header::{CONTENT_LENGTH, CONTENT_TYPE};
|
||||
use reqwest::{Body, Client, Response};
|
||||
use tokio::fs::File;
|
||||
use tokio::join;
|
||||
use tokio::{select, sync::RwLock, time};
|
||||
use tokio_util::io::ReaderStream;
|
||||
use tokio_util::{sync::CancellationToken, task::TaskTracker};
|
||||
|
||||
use crate::config::WebdavConfig;
|
||||
use crate::qmdl_store::RecordingStore;
|
||||
|
||||
pub struct WebdavUploadWorkerConfig {
|
||||
poll_interval: Duration,
|
||||
min_age: TimeDelta,
|
||||
url: String,
|
||||
username: Option<String>,
|
||||
password: Option<String>,
|
||||
timeout: Duration,
|
||||
delete_on_upload: bool,
|
||||
}
|
||||
|
||||
impl From<WebdavConfig> for WebdavUploadWorkerConfig {
|
||||
fn from(value: WebdavConfig) -> Self {
|
||||
WebdavUploadWorkerConfig {
|
||||
poll_interval: Duration::from_secs(value.poll_interval_secs),
|
||||
min_age: TimeDelta::seconds(value.min_age_secs),
|
||||
url: value.url,
|
||||
username: value.username,
|
||||
password: value.password,
|
||||
timeout: Duration::from_secs(value.upload_timeout_secs),
|
||||
delete_on_upload: value.delete_on_upload,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
enum FileKind {
|
||||
Analysis,
|
||||
Qmdl,
|
||||
}
|
||||
|
||||
impl FileKind {
|
||||
fn as_extension(&self) -> &'static str {
|
||||
match self {
|
||||
FileKind::Analysis => ".ndjson",
|
||||
FileKind::Qmdl => ".qmdl",
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
impl Display for FileKind {
|
||||
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
|
||||
match self {
|
||||
FileKind::Analysis => write!(f, "analysis"),
|
||||
FileKind::Qmdl => write!(f, "QMDL"),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone)]
|
||||
struct WebDavClient {
|
||||
client: Client,
|
||||
url: String,
|
||||
username: Option<String>,
|
||||
password: Option<String>,
|
||||
}
|
||||
|
||||
impl WebDavClient {
|
||||
fn new(
|
||||
mut url: String,
|
||||
username: Option<String>,
|
||||
password: Option<String>,
|
||||
timeout: Duration,
|
||||
) -> Result<Self, reqwest::Error> {
|
||||
if !url.ends_with('/') {
|
||||
url.push('/');
|
||||
}
|
||||
Ok(Self {
|
||||
client: reqwest::Client::builder().timeout(timeout).build()?,
|
||||
url,
|
||||
username,
|
||||
password,
|
||||
})
|
||||
}
|
||||
|
||||
async fn try_upload_file(&self, file: File, name: &str) -> anyhow::Result<Response> {
|
||||
let file_size = file.metadata().await?.len();
|
||||
|
||||
let stream = ReaderStream::new(file);
|
||||
let body = Body::wrap_stream(stream);
|
||||
|
||||
let target = format!("{}{}", self.url, name);
|
||||
|
||||
let client = self
|
||||
.client
|
||||
.put(&target)
|
||||
.header(CONTENT_TYPE, "application/octet-stream")
|
||||
.header(CONTENT_LENGTH, file_size);
|
||||
|
||||
let client = match (&self.username, &self.password) {
|
||||
(Some(username), Some(password)) => client.basic_auth(username, Some(password)),
|
||||
(Some(username), None) => client.basic_auth(username, None::<&str>),
|
||||
(None, None) => client,
|
||||
(None, Some(_)) => {
|
||||
warn!(
|
||||
"Got WebDAV auth setting with no username but with a password, skipping authentication"
|
||||
);
|
||||
client
|
||||
}
|
||||
};
|
||||
|
||||
let resp = client.body(body).send().await?.error_for_status();
|
||||
Ok(resp?)
|
||||
}
|
||||
}
|
||||
|
||||
async fn try_upload_entry(
|
||||
client: WebDavClient,
|
||||
store: Arc<RwLock<RecordingStore>>,
|
||||
entry_name: String,
|
||||
file_kind: FileKind,
|
||||
shutdown_token: CancellationToken,
|
||||
) -> Option<()> {
|
||||
let read_lock = store.read().await;
|
||||
let entry_idx = read_lock.entry_for_name(&entry_name)?.0;
|
||||
let file = match file_kind {
|
||||
FileKind::Analysis => read_lock.open_entry_analysis(entry_idx).await,
|
||||
FileKind::Qmdl => read_lock.open_entry_qmdl(entry_idx).await,
|
||||
};
|
||||
drop(read_lock);
|
||||
|
||||
let Ok(file) = file.map_err(|err| {
|
||||
warn!(
|
||||
"Unable to open entry: {} {} file: {:?}",
|
||||
entry_name, file_kind, err
|
||||
)
|
||||
}) else {
|
||||
return None;
|
||||
};
|
||||
|
||||
let file_name = format!("{}{}", entry_name, file_kind.as_extension());
|
||||
|
||||
let res = select! {
|
||||
_ = shutdown_token.cancelled() => {
|
||||
warn!(
|
||||
"Cancelling upload for entry {} {} file: received shutdown signal",
|
||||
entry_name, file_kind
|
||||
);
|
||||
return None;
|
||||
},
|
||||
res = client.try_upload_file(file, &file_name) => res,
|
||||
};
|
||||
|
||||
match res {
|
||||
Ok(_) => {
|
||||
info!("Uploaded {} file for entry {}", file_kind, entry_name);
|
||||
Some(())
|
||||
}
|
||||
Err(err) => {
|
||||
warn!(
|
||||
"Failed to upload {} file for entry {}: {:?}",
|
||||
file_kind, entry_name, err
|
||||
);
|
||||
None
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
pub fn run_webdav_upload_worker(
|
||||
task_tracker: &TaskTracker,
|
||||
shutdown_token: CancellationToken,
|
||||
qmdl_store_lock: Arc<RwLock<RecordingStore>>,
|
||||
config: WebdavUploadWorkerConfig,
|
||||
) {
|
||||
task_tracker.spawn(async move {
|
||||
let mut interval = time::interval(config.poll_interval);
|
||||
interval.set_missed_tick_behavior(time::MissedTickBehavior::Skip);
|
||||
|
||||
let webdav_client = match WebDavClient::new(
|
||||
config.url,
|
||||
config.username,
|
||||
config.password,
|
||||
config.timeout,
|
||||
) {
|
||||
Ok(client) => client,
|
||||
Err(err) => {
|
||||
warn!("Unable to create WebDAV client: {:?}", err);
|
||||
return;
|
||||
}
|
||||
};
|
||||
|
||||
loop {
|
||||
select! {
|
||||
_ = shutdown_token.cancelled() => break,
|
||||
_ = interval.tick() => {
|
||||
loop {
|
||||
let Some(unuploaded_entry) = qmdl_store_lock
|
||||
.read()
|
||||
.await
|
||||
.get_next_unuploaded_entry(config.min_age) else {
|
||||
break;
|
||||
};
|
||||
|
||||
let (Some(()), Some(())) = join!(
|
||||
try_upload_entry(
|
||||
webdav_client.clone(),
|
||||
qmdl_store_lock.clone(),
|
||||
unuploaded_entry.clone(),
|
||||
FileKind::Qmdl,
|
||||
shutdown_token.clone(),
|
||||
),
|
||||
try_upload_entry(
|
||||
webdav_client.clone(),
|
||||
qmdl_store_lock.clone(),
|
||||
unuploaded_entry.clone(),
|
||||
FileKind::Analysis,
|
||||
shutdown_token.clone()
|
||||
),
|
||||
) else {
|
||||
break;
|
||||
};
|
||||
|
||||
if config.delete_on_upload {
|
||||
match qmdl_store_lock.write().await.delete_entry(&unuploaded_entry).await {
|
||||
Ok(_) => info!("Successfully deleted entry: {} after upload to WebDAV", unuploaded_entry),
|
||||
Err(err) => warn!("Unable to delete entry: {} after upload to WebDAV: {}", unuploaded_entry, err),
|
||||
}
|
||||
} else {
|
||||
match qmdl_store_lock.write().await.mark_entry_as_uploaded(&unuploaded_entry, rayhunter::clock::get_adjusted_now()).await {
|
||||
Ok(_) => info!("Successfully marked entry: {} as uploaded", unuploaded_entry),
|
||||
Err(err) => warn!("Unable to mark entry: {} as uploaded: {}", unuploaded_entry, err),
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
|
||||
#[cfg(test)]
|
||||
mod tests {
|
||||
use super::*;
|
||||
use axum::{
|
||||
Router,
|
||||
body::Bytes,
|
||||
extract::{Path as AxumPath, State},
|
||||
http::{HeaderMap, StatusCode},
|
||||
routing::put,
|
||||
};
|
||||
use tempfile::Builder;
|
||||
use tokio::io::AsyncWriteExt;
|
||||
use tokio::net::TcpListener;
|
||||
use tokio::sync::Mutex;
|
||||
|
||||
#[derive(Clone, Debug)]
|
||||
struct RecordedPut {
|
||||
path: String,
|
||||
auth: Option<String>,
|
||||
body: Vec<u8>,
|
||||
}
|
||||
|
||||
async fn capture_put(
|
||||
State(state): State<Arc<Mutex<Vec<RecordedPut>>>>,
|
||||
AxumPath(path): AxumPath<String>,
|
||||
headers: HeaderMap,
|
||||
body: Bytes,
|
||||
) -> StatusCode {
|
||||
let auth = headers
|
||||
.get("authorization")
|
||||
.and_then(|v| v.to_str().ok())
|
||||
.map(String::from);
|
||||
state.lock().await.push(RecordedPut {
|
||||
path,
|
||||
auth,
|
||||
body: body.to_vec(),
|
||||
});
|
||||
StatusCode::CREATED
|
||||
}
|
||||
|
||||
async fn setup_webdav_server() -> (Arc<Mutex<Vec<RecordedPut>>>, String) {
|
||||
crate::crypto_provider::install_default();
|
||||
|
||||
let state = Arc::new(Mutex::new(Vec::new()));
|
||||
let app = Router::new()
|
||||
.route("/{*path}", put(capture_put))
|
||||
.with_state(state.clone());
|
||||
|
||||
let listener = TcpListener::bind("127.0.0.1:0").await.unwrap();
|
||||
let addr = listener.local_addr().unwrap();
|
||||
let url = format!("http://{}/dav", addr);
|
||||
|
||||
tokio::spawn(async move {
|
||||
axum::serve(listener, app).await.unwrap();
|
||||
});
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(100)).await;
|
||||
|
||||
(state, url)
|
||||
}
|
||||
|
||||
async fn cleanup_worker(shutdown: CancellationToken, tracker: TaskTracker) {
|
||||
shutdown.cancel();
|
||||
tracker.close();
|
||||
tracker.wait().await;
|
||||
}
|
||||
|
||||
async fn make_store_with_closed_entry(
|
||||
dir: &std::path::Path,
|
||||
) -> (Arc<RwLock<RecordingStore>>, String) {
|
||||
let mut store = RecordingStore::create(dir).await.unwrap();
|
||||
let (mut qmdl_file, mut analysis_file) = store.new_entry().await.unwrap();
|
||||
qmdl_file.write_all(b"fake qmdl payload").await.unwrap();
|
||||
qmdl_file.flush().await.unwrap();
|
||||
analysis_file
|
||||
.write_all(b"fake ndjson payload")
|
||||
.await
|
||||
.unwrap();
|
||||
analysis_file.flush().await.unwrap();
|
||||
let entry_index = store.current_entry.unwrap();
|
||||
let name = store.manifest.entries[entry_index].name.clone();
|
||||
store.update_entry_qmdl_size(entry_index, 17).await.unwrap();
|
||||
store.close_current_entry().await.unwrap();
|
||||
(Arc::new(RwLock::new(store)), name)
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_webdav_upload_worker_uploads_entry() {
|
||||
let (captured, url) = setup_webdav_server().await;
|
||||
|
||||
let dir = Builder::new().prefix("webdav_test").tempdir().unwrap();
|
||||
let (store, entry_name) = make_store_with_closed_entry(dir.path()).await;
|
||||
|
||||
let shutdown = CancellationToken::new();
|
||||
let tracker = TaskTracker::new();
|
||||
let config = WebdavUploadWorkerConfig {
|
||||
poll_interval: Duration::from_millis(50),
|
||||
min_age: TimeDelta::seconds(-1),
|
||||
url,
|
||||
username: Some("user".to_string()),
|
||||
password: Some("password".to_string()),
|
||||
timeout: Duration::from_secs(1),
|
||||
delete_on_upload: false,
|
||||
};
|
||||
|
||||
run_webdav_upload_worker(&tracker, shutdown.clone(), store.clone(), config);
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
cleanup_worker(shutdown, tracker).await;
|
||||
|
||||
let recorded = captured.lock().await;
|
||||
assert_eq!(recorded.len(), 2);
|
||||
let paths: Vec<&str> = recorded.iter().map(|r| r.path.as_str()).collect();
|
||||
let qmdl_path = format!("dav/{}.qmdl", entry_name);
|
||||
let ndjson_path = format!("dav/{}.ndjson", entry_name);
|
||||
assert!(paths.contains(&qmdl_path.as_str()));
|
||||
assert!(paths.contains(&ndjson_path.as_str()));
|
||||
for put in recorded.iter() {
|
||||
assert_eq!(put.auth.as_deref(), Some("Basic dXNlcjpwYXNzd29yZA=="));
|
||||
}
|
||||
let qmdl_body = recorded
|
||||
.iter()
|
||||
.find(|r| r.path == qmdl_path)
|
||||
.unwrap()
|
||||
.body
|
||||
.clone();
|
||||
let ndjson_body = recorded
|
||||
.iter()
|
||||
.find(|r| r.path == ndjson_path)
|
||||
.unwrap()
|
||||
.body
|
||||
.clone();
|
||||
drop(recorded);
|
||||
assert_eq!(qmdl_body, b"fake qmdl payload");
|
||||
assert_eq!(ndjson_body, b"fake ndjson payload");
|
||||
|
||||
let store_read = store.read().await;
|
||||
let (_, entry) = store_read.entry_for_name(&entry_name).unwrap();
|
||||
assert!(entry.upload_time.is_some());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_webdav_upload_worker_deletes_when_configured() {
|
||||
let (captured, url) = setup_webdav_server().await;
|
||||
|
||||
let dir = Builder::new().prefix("webdav_test").tempdir().unwrap();
|
||||
let (store, entry_name) = make_store_with_closed_entry(dir.path()).await;
|
||||
|
||||
let shutdown = CancellationToken::new();
|
||||
let tracker = TaskTracker::new();
|
||||
let config = WebdavUploadWorkerConfig {
|
||||
poll_interval: Duration::from_millis(50),
|
||||
min_age: TimeDelta::seconds(-1),
|
||||
url,
|
||||
username: None,
|
||||
password: None,
|
||||
timeout: Duration::from_secs(1),
|
||||
delete_on_upload: true,
|
||||
};
|
||||
|
||||
run_webdav_upload_worker(&tracker, shutdown.clone(), store.clone(), config);
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
cleanup_worker(shutdown, tracker).await;
|
||||
|
||||
assert_eq!(captured.lock().await.len(), 2);
|
||||
|
||||
let store_read = store.read().await;
|
||||
assert!(store_read.entry_for_name(&entry_name).is_none());
|
||||
}
|
||||
|
||||
#[tokio::test]
|
||||
async fn test_webdav_upload_worker_respects_min_age() {
|
||||
let (captured, url) = setup_webdav_server().await;
|
||||
|
||||
let dir = Builder::new().prefix("webdav_test").tempdir().unwrap();
|
||||
let (store, entry_name) = make_store_with_closed_entry(dir.path()).await;
|
||||
|
||||
let shutdown = CancellationToken::new();
|
||||
let tracker = TaskTracker::new();
|
||||
let config = WebdavUploadWorkerConfig {
|
||||
poll_interval: Duration::from_millis(50),
|
||||
min_age: TimeDelta::seconds(3600),
|
||||
url,
|
||||
username: None,
|
||||
password: None,
|
||||
timeout: Duration::from_secs(1),
|
||||
delete_on_upload: false,
|
||||
};
|
||||
|
||||
run_webdav_upload_worker(&tracker, shutdown.clone(), store.clone(), config);
|
||||
|
||||
tokio::time::sleep(Duration::from_millis(500)).await;
|
||||
cleanup_worker(shutdown, tracker).await;
|
||||
|
||||
assert!(captured.lock().await.is_empty());
|
||||
|
||||
let store_read = store.read().await;
|
||||
let (_, entry) = store_read.entry_for_name(&entry_name).unwrap();
|
||||
assert!(entry.upload_time.is_none());
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user