chore: migrate store to redb
This commit is contained in:
Riz Ashraf committed
1 parent
7f0286dc40
commit
9a8b6e52e9
12 files changed
+1954
-1365
No files matched your search
+65
-59
@@ -1,9 +1,9 @@
|
||||
use tokio_util::io::StreamReader;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use futures_util::StreamExt;
|
||||
use std::sync::Arc;
|
||||
use tokio::io::AsyncBufReadExt;
|
||||
use tokio::sync::RwLock;
|
||||
use tokio::sync::mpsc;
|
||||
use tokio_util::io::StreamReader;
|
||||
|
||||
pub fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
|
||||
let rt = tokio::runtime::Runtime::new().unwrap();
|
||||
@@ -16,7 +16,9 @@ pub fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error + S
|
||||
let mut handle = stdin.lock();
|
||||
let mut buffer = String::new();
|
||||
while let Ok(bytes) = std::io::BufRead::read_line(&mut handle, &mut buffer) {
|
||||
if bytes == 0 { break; }
|
||||
if bytes == 0 {
|
||||
break;
|
||||
}
|
||||
let _ = msg_tx.blocking_send(buffer.clone());
|
||||
buffer.clear();
|
||||
}
|
||||
@@ -26,7 +28,7 @@ pub fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error + S
|
||||
let target_url = target_url.to_string();
|
||||
let post_url = Arc::new(RwLock::new(String::new()));
|
||||
let post_url_clone = Arc::clone(&post_url);
|
||||
|
||||
|
||||
let client = reqwest::Client::builder().build().unwrap();
|
||||
|
||||
tokio::spawn(async move {
|
||||
@@ -34,14 +36,16 @@ pub fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error + S
|
||||
loop {
|
||||
let url = post_url_clone.read().await.clone();
|
||||
if !url.is_empty() {
|
||||
let res = client.post(&url)
|
||||
let res = client
|
||||
.post(&url)
|
||||
.header("Accept", "application/json, text/event-stream")
|
||||
.header("Content-Type", "application/json")
|
||||
.body(msg.clone())
|
||||
.send().await;
|
||||
|
||||
.send()
|
||||
.await;
|
||||
|
||||
if res.is_ok() {
|
||||
break;
|
||||
break;
|
||||
}
|
||||
}
|
||||
tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
|
||||
@@ -49,67 +53,69 @@ pub fn run_proxy(target_url: &str) -> Result<bool, Box<dyn std::error::Error + S
|
||||
}
|
||||
});
|
||||
|
||||
loop {
|
||||
if shutdown_rx.try_recv().is_ok() {
|
||||
return Ok(false);
|
||||
}
|
||||
if shutdown_rx.try_recv().is_ok() {
|
||||
return Ok(false);
|
||||
}
|
||||
|
||||
let sse_url = format!("{}/sse", target_url);
|
||||
let client = reqwest::Client::builder().build().unwrap();
|
||||
|
||||
match client.get(&sse_url).header("Accept", "text/event-stream").send().await {
|
||||
Ok(resp) => {
|
||||
if resp.status() == reqwest::StatusCode::GONE {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let stream = resp.bytes_stream().map(|res| res.map_err(|e| std::io::Error::new(std::io::ErrorKind::Other, e)));
|
||||
let mut reader = tokio::io::BufReader::new(StreamReader::new(stream));
|
||||
let mut line = String::new();
|
||||
let mut is_message = false;
|
||||
let mut is_endpoint = false;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
return Ok(false);
|
||||
}
|
||||
res = reader.read_line(&mut line) => {
|
||||
match res {
|
||||
Ok(bytes) => {
|
||||
if bytes == 0 { break; }
|
||||
let trimmed = line.trim();
|
||||
if trimmed.starts_with("event: message") {
|
||||
is_message = true;
|
||||
is_endpoint = false;
|
||||
} else if trimmed.starts_with("event: endpoint") {
|
||||
is_endpoint = true;
|
||||
let sse_url = format!("{}/sse", target_url);
|
||||
let client = reqwest::Client::builder().build().unwrap();
|
||||
|
||||
match client
|
||||
.get(&sse_url)
|
||||
.header("Accept", "text/event-stream")
|
||||
.send()
|
||||
.await
|
||||
{
|
||||
Ok(resp) => {
|
||||
if resp.status() == reqwest::StatusCode::GONE {
|
||||
return Ok(true);
|
||||
}
|
||||
|
||||
let stream = resp.bytes_stream().map(|res| {
|
||||
res.map_err(std::io::Error::other)
|
||||
});
|
||||
let mut reader = tokio::io::BufReader::new(StreamReader::new(stream));
|
||||
let mut line = String::new();
|
||||
let mut is_message = false;
|
||||
let mut is_endpoint = false;
|
||||
|
||||
loop {
|
||||
tokio::select! {
|
||||
_ = shutdown_rx.recv() => {
|
||||
return Ok(false);
|
||||
}
|
||||
res = reader.read_line(&mut line) => {
|
||||
match res {
|
||||
Ok(bytes) => {
|
||||
if bytes == 0 { break; }
|
||||
let trimmed = line.trim();
|
||||
if trimmed.starts_with("event: message") {
|
||||
is_message = true;
|
||||
is_endpoint = false;
|
||||
} else if trimmed.starts_with("event: endpoint") {
|
||||
is_endpoint = true;
|
||||
is_message = false;
|
||||
} else if let Some(stripped) = trimmed.strip_prefix("data: ") {
|
||||
if is_message {
|
||||
println!("{}", stripped);
|
||||
is_message = false;
|
||||
} else if trimmed.starts_with("data: ") {
|
||||
if is_message {
|
||||
println!("{}", &trimmed[6..]);
|
||||
is_message = false;
|
||||
} else if is_endpoint {
|
||||
let ep = &trimmed[6..];
|
||||
let mut p = post_url.write().await;
|
||||
*p = format!("{}{}", target_url, ep);
|
||||
is_endpoint = false;
|
||||
}
|
||||
} else if is_endpoint {
|
||||
let mut p = post_url.write().await;
|
||||
*p = format!("{}{}", target_url, stripped);
|
||||
is_endpoint = false;
|
||||
}
|
||||
line.clear();
|
||||
}
|
||||
Err(_) => break,
|
||||
line.clear();
|
||||
}
|
||||
Err(_) => break,
|
||||
}
|
||||
}
|
||||
}
|
||||
*post_url.write().await = String::new();
|
||||
return Ok(true);
|
||||
}
|
||||
Err(_) => {
|
||||
return Ok(true);
|
||||
}
|
||||
*post_url.write().await = String::new();
|
||||
Ok(true)
|
||||
}
|
||||
Err(_) => Ok(true),
|
||||
}
|
||||
})
|
||||
}
|
||||
Reference in new issue
Block a user