use anyhow::Result;
use async_trait::async_trait;
use axum::{extract::State, routing::post, Json, Router};
use std::net::SocketAddr;
use std::sync::atomic::{AtomicBool, Ordering};
use std::sync::Arc;
use tokio::sync::{mpsc, Mutex};
use tower::ServiceBuilder;
use tower_http::cors::CorsLayer;
use tracing::{debug, error, info};
use super::{
ConnectionInfo, ConnectionStats, Transport, TransportConnection, TransportMessage,
TransportStats, TransportType,
};
use super::Deserialize;
use super::Serialize;
#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct HttpConfig {
pub bind_addr: String,
pub tls: bool,
pub cert_path: Option<String>,
pub key_path: Option<String>,
pub max_body_size: usize,
pub request_timeout_ms: u64,
}
impl Default for HttpConfig {
fn default() -> Self {
Self {
bind_addr: "127.0.0.1:8080".to_string(),
tls: false,
cert_path: None,
key_path: None,
max_body_size: 10 * 1024 * 1024, request_timeout_ms: 30000,
}
}
}
pub struct HttpTransport {
config: HttpConfig,
running: AtomicBool,
stats: Arc<Mutex<TransportStats>>,
shutdown_tx: Option<mpsc::Sender<()>>,
}
impl HttpTransport {
pub fn new(config: serde_json::Value) -> Result<Self> {
let config: HttpConfig = serde_json::from_value(config)?;
Ok(Self {
config,
running: AtomicBool::new(false),
stats: Arc::new(Mutex::new(TransportStats::default())),
shutdown_tx: None,
})
}
async fn start_server(&mut self) -> Result<mpsc::Sender<()>> {
let (shutdown_tx, mut shutdown_rx) = mpsc::channel(1);
let addr: SocketAddr = self.config.bind_addr.parse()?;
let stats = self.stats.clone();
let app = Router::new()
.route("/rpc", post(handle_rpc))
.layer(
ServiceBuilder::new()
.layer(CorsLayer::permissive())
.into_inner(),
)
.with_state(stats);
let listener = tokio::net::TcpListener::bind(&addr).await?;
info!("HTTP server listening on {}", addr);
tokio::spawn(async move {
axum::serve(listener, app)
.with_graceful_shutdown(async move {
let _ = shutdown_rx.recv().await;
info!("HTTP server shutting down");
})
.await
.map_err(|e| error!("Server error: {}", e))
.ok();
});
Ok(shutdown_tx)
}
}
#[async_trait]
impl Transport for HttpTransport {
fn transport_type(&self) -> TransportType {
TransportType::Http
}
async fn start(&mut self) -> Result<()> {
if self.running.load(Ordering::Relaxed) {
return Err(anyhow::anyhow!("Transport already running"));
}
let shutdown_tx = self.start_server().await?;
self.shutdown_tx = Some(shutdown_tx);
self.running.store(true, Ordering::Relaxed);
info!("Started HTTP transport on {}", self.config.bind_addr);
Ok(())
}
async fn stop(&mut self) -> Result<()> {
if let Some(shutdown_tx) = self.shutdown_tx.take() {
let _ = shutdown_tx.send(()).await;
}
self.running.store(false, Ordering::Relaxed);
info!("Stopped HTTP transport");
Ok(())
}
async fn accept(&mut self) -> Result<Box<dyn TransportConnection>> {
if !self.running.load(Ordering::Relaxed) {
return Err(anyhow::anyhow!("Transport not running"));
}
let mut stats = self.stats.lock().await;
stats.connections_accepted += 1;
stats.active_connections += 1;
drop(stats);
Ok(Box::new(HttpConnection::new(
"127.0.0.1:12345".to_string(),
self.stats.clone(),
)))
}
async fn connect(&mut self, address: &str) -> Result<Box<dyn TransportConnection>> {
Ok(Box::new(HttpConnection::new(
address.to_string(),
self.stats.clone(),
)))
}
fn is_running(&self) -> bool {
self.running.load(Ordering::Relaxed)
}
fn get_stats(&self) -> TransportStats {
if let Ok(stats) = self.stats.try_lock() {
stats.clone()
} else {
TransportStats::default()
}
}
async fn set_option(&mut self, key: &str, value: serde_json::Value) -> Result<()> {
match key {
"max_body_size" => {
if let Some(size) = value.as_u64() {
self.config.max_body_size = size as usize;
}
},
"request_timeout_ms" => {
if let Some(timeout) = value.as_u64() {
self.config.request_timeout_ms = timeout;
}
},
_ => {},
}
Ok(())
}
}
pub struct HttpConnection {
info: ConnectionInfo,
remote_addr: String,
connected: AtomicBool,
stats: Arc<Mutex<ConnectionStats>>,
transport_stats: Arc<Mutex<TransportStats>>,
message_queue: Arc<Mutex<mpsc::UnboundedReceiver<TransportMessage>>>,
#[allow(dead_code)] message_tx: mpsc::UnboundedSender<TransportMessage>,
}
impl HttpConnection {
fn new(remote_addr: String, transport_stats: Arc<Mutex<TransportStats>>) -> Self {
let (message_tx, message_rx) = mpsc::unbounded_channel();
let info = ConnectionInfo {
id: uuid::Uuid::new_v4().to_string(),
transport_type: TransportType::Http,
client_id: None,
remote_addr: Some(remote_addr.clone()),
connected_at: chrono::Utc::now(),
security_info: None, };
Self {
info,
remote_addr,
connected: AtomicBool::new(true),
stats: Arc::new(Mutex::new(ConnectionStats::default())),
transport_stats,
message_queue: Arc::new(Mutex::new(message_rx)),
message_tx,
}
}
}
#[async_trait]
impl TransportConnection for HttpConnection {
fn connection_info(&self) -> &ConnectionInfo {
&self.info
}
async fn send(&mut self, message: TransportMessage) -> Result<()> {
if !self.connected.load(Ordering::Relaxed) {
return Err(anyhow::anyhow!("Connection closed"));
}
debug!("HTTP send to {}: {}", self.remote_addr, message.id);
let mut stats = self.stats.lock().await;
stats.messages_sent += 1;
stats.bytes_sent += serde_json::to_vec(&message.payload)?.len() as u64;
drop(stats);
let mut transport_stats = self.transport_stats.lock().await;
transport_stats.messages_sent += 1;
drop(transport_stats);
Ok(())
}
async fn receive(&mut self) -> Result<Option<TransportMessage>> {
if !self.connected.load(Ordering::Relaxed) {
return Ok(None);
}
let mut queue = self.message_queue.lock().await;
match queue.recv().await {
Some(message) => {
let mut stats = self.stats.lock().await;
stats.messages_received += 1;
drop(stats);
let mut transport_stats = self.transport_stats.lock().await;
transport_stats.messages_received += 1;
drop(transport_stats);
Ok(Some(message))
},
None => Ok(None),
}
}
async fn close(&mut self) -> Result<()> {
self.connected.store(false, Ordering::Relaxed);
let mut transport_stats = self.transport_stats.lock().await;
transport_stats.active_connections = transport_stats.active_connections.saturating_sub(1);
drop(transport_stats);
info!("Closed HTTP connection to {}", self.remote_addr);
Ok(())
}
fn is_connected(&self) -> bool {
self.connected.load(Ordering::Relaxed)
}
fn get_stats(&self) -> ConnectionStats {
if let Ok(stats) = self.stats.try_lock() {
stats.clone()
} else {
ConnectionStats::default()
}
}
async fn set_option(&mut self, key: &str, value: serde_json::Value) -> Result<()> {
if key == "client_id" {
if let Some(id) = value.as_str() {
self.info.client_id = Some(id.to_string());
}
}
Ok(())
}
}
async fn handle_rpc(
State(stats): State<Arc<Mutex<TransportStats>>>,
Json(payload): Json<serde_json::Value>,
) -> Result<Json<serde_json::Value>, (axum::http::StatusCode, String)> {
let mut s = stats.lock().await;
s.messages_received += 1;
Ok(Json(serde_json::json!({
"jsonrpc": "2.0",
"result": payload,
"id": payload.get("id").cloned().unwrap_or(serde_json::Value::Null)
})))
}