#[allow(unused_imports)]
use crate::connection::{Message, WebSocketConnection, WebSocketError, WebSocketResult};
use async_trait::async_trait;
use std::sync::Arc;
#[cfg(feature = "di")]
use reinhardt_di::{Depends, Injectable, InjectionContext};
pub struct ConsumerContext {
pub connection: Arc<WebSocketConnection>,
pub headers: std::collections::HashMap<String, String>,
pub metadata: std::collections::HashMap<String, String>,
#[cfg(feature = "di")]
di_context: Option<Arc<InjectionContext>>,
}
impl ConsumerContext {
pub fn new(connection: Arc<WebSocketConnection>) -> Self {
Self {
connection,
headers: std::collections::HashMap::new(),
metadata: std::collections::HashMap::new(),
#[cfg(feature = "di")]
di_context: None,
}
}
#[cfg(feature = "di")]
pub fn with_di_context(
connection: Arc<WebSocketConnection>,
di_context: Arc<InjectionContext>,
) -> Self {
Self {
connection,
headers: std::collections::HashMap::new(),
metadata: std::collections::HashMap::new(),
di_context: Some(di_context),
}
}
pub fn with_header(mut self, key: String, value: String) -> Self {
self.headers.insert(key, value);
self
}
pub fn get_header(&self, key: &str) -> Option<&String> {
self.headers.get(key)
}
pub fn cookie_header(&self) -> Option<&str> {
self.headers.get("cookie").map(|s| s.as_str())
}
pub fn with_metadata(mut self, key: String, value: String) -> Self {
self.metadata.insert(key, value);
self
}
pub fn get_metadata(&self, key: &str) -> Option<&String> {
self.metadata.get(key)
}
#[cfg(feature = "di")]
pub fn di_context(&self) -> Option<&Arc<InjectionContext>> {
self.di_context.as_ref()
}
#[cfg(feature = "di")]
pub fn set_di_context(&mut self, ctx: Arc<InjectionContext>) {
self.di_context = Some(ctx);
}
#[cfg(feature = "di")]
pub async fn resolve<T>(&self) -> WebSocketResult<T>
where
T: Injectable + Clone + Send + Sync + 'static,
{
let ctx = self
.di_context
.as_ref()
.ok_or_else(|| WebSocketError::Internal("DI context not available".to_string()))?;
Depends::<T>::resolve(ctx, true)
.await
.map(|injected| injected.into_inner())
.map_err(|_| WebSocketError::Internal("dependency resolution failed".to_string()))
}
#[cfg(feature = "di")]
pub async fn resolve_uncached<T>(&self) -> WebSocketResult<T>
where
T: Injectable + Clone + Send + Sync + 'static,
{
let ctx = self
.di_context
.as_ref()
.ok_or_else(|| WebSocketError::Internal("DI context not available".to_string()))?;
Depends::<T>::resolve(ctx, false)
.await
.map(|injected| injected.into_inner())
.map_err(|_| WebSocketError::Internal("dependency resolution failed".to_string()))
}
#[cfg(feature = "di")]
pub async fn try_resolve<T>(&self) -> Option<T>
where
T: Injectable + Clone + Send + Sync + 'static,
{
let ctx = self.di_context.as_ref()?;
Depends::<T>::resolve(ctx, true)
.await
.ok()
.map(|injected| injected.into_inner())
}
}
#[async_trait]
pub trait WebSocketConsumer: Send + Sync {
async fn on_connect(&self, context: &mut ConsumerContext) -> WebSocketResult<()>;
async fn on_message(
&self,
context: &mut ConsumerContext,
message: Message,
) -> WebSocketResult<()>;
async fn on_disconnect(&self, context: &mut ConsumerContext) -> WebSocketResult<()>;
}
pub struct EchoConsumer {
prefix: String,
}
impl EchoConsumer {
pub fn new() -> Self {
Self {
prefix: "Echo".to_string(),
}
}
pub fn with_prefix(prefix: String) -> Self {
Self { prefix }
}
}
impl Default for EchoConsumer {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl WebSocketConsumer for EchoConsumer {
async fn on_connect(&self, context: &mut ConsumerContext) -> WebSocketResult<()> {
context
.connection
.send_text(format!("{}: Connection established", self.prefix))
.await
}
async fn on_message(
&self,
context: &mut ConsumerContext,
message: Message,
) -> WebSocketResult<()> {
match message {
Message::Text { data } => {
context
.connection
.send_text(format!("{}: {}", self.prefix, data))
.await
}
Message::Binary { data } => {
match String::from_utf8(data.clone()) {
Ok(text) => {
context
.connection
.send_text(format!("{}: {}", self.prefix, text))
.await
}
Err(_) => {
context
.connection
.send_text(format!("{}: binary({} bytes)", self.prefix, data.len()))
.await
}
}
}
Message::Close { code, reason } => {
context
.connection
.close_with_reason(code, reason)
.await
.ok();
Ok(())
}
_ => Ok(()),
}
}
async fn on_disconnect(&self, _context: &mut ConsumerContext) -> WebSocketResult<()> {
Ok(())
}
}
pub struct BroadcastConsumer {
room: Arc<crate::room::Room>,
}
impl BroadcastConsumer {
pub fn new(room: Arc<crate::room::Room>) -> Self {
Self { room }
}
}
#[async_trait]
impl WebSocketConsumer for BroadcastConsumer {
async fn on_connect(&self, context: &mut ConsumerContext) -> WebSocketResult<()> {
let client_id = context.connection.id().to_string();
self.room
.join(client_id.clone(), context.connection.clone())
.await
.map_err(|e| crate::connection::WebSocketError::Connection(e.to_string()))?;
context
.connection
.send_text("Joined broadcast room".to_string())
.await
}
async fn on_message(
&self,
_context: &mut ConsumerContext,
message: Message,
) -> WebSocketResult<()> {
let result = self.room.broadcast(message).await;
if result.is_complete_failure() {
return Err(crate::connection::WebSocketError::Send(
"broadcast failed for all clients".to_string(),
));
}
Ok(())
}
async fn on_disconnect(&self, context: &mut ConsumerContext) -> WebSocketResult<()> {
let client_id = context.connection.id();
let _ = self.room.leave(client_id).await;
context.connection.force_close().await;
Ok(())
}
}
pub struct JsonConsumer;
impl JsonConsumer {
pub fn new() -> Self {
Self
}
}
impl Default for JsonConsumer {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl WebSocketConsumer for JsonConsumer {
async fn on_connect(&self, context: &mut ConsumerContext) -> WebSocketResult<()> {
context
.connection
.send_json(&serde_json::json!({
"type": "connection",
"status": "connected"
}))
.await
}
async fn on_message(
&self,
context: &mut ConsumerContext,
message: Message,
) -> WebSocketResult<()> {
match message {
Message::Text { data } => {
let json: serde_json::Value = serde_json::from_str(&data)
.map_err(|e| crate::connection::WebSocketError::Protocol(e.to_string()))?;
let response = serde_json::json!({
"type": "echo",
"data": json,
"timestamp": chrono::Utc::now().to_rfc3339()
});
context.connection.send_json(&response).await
}
Message::Binary { data } => {
let text = String::from_utf8(data).map_err(|e| {
crate::connection::WebSocketError::BinaryPayload(format!(
"binary payload is not valid UTF-8: {}",
e
))
})?;
let json: serde_json::Value = serde_json::from_str(&text).map_err(|e| {
crate::connection::WebSocketError::BinaryPayload(format!(
"binary payload is not valid JSON: {}",
e
))
})?;
let response = serde_json::json!({
"type": "echo",
"data": json,
"source": "binary",
"timestamp": chrono::Utc::now().to_rfc3339()
});
context.connection.send_json(&response).await
}
_ => Ok(()),
}
}
async fn on_disconnect(&self, _context: &mut ConsumerContext) -> WebSocketResult<()> {
Ok(())
}
}
pub struct ConsumerChain {
consumers: Vec<Box<dyn WebSocketConsumer>>,
}
impl ConsumerChain {
pub fn new() -> Self {
Self {
consumers: Vec::new(),
}
}
pub fn add_consumer(&mut self, consumer: Box<dyn WebSocketConsumer>) {
self.consumers.push(consumer);
}
}
impl Default for ConsumerChain {
fn default() -> Self {
Self::new()
}
}
#[async_trait]
impl WebSocketConsumer for ConsumerChain {
async fn on_connect(&self, context: &mut ConsumerContext) -> WebSocketResult<()> {
for consumer in &self.consumers {
consumer.on_connect(context).await?;
}
Ok(())
}
async fn on_message(
&self,
context: &mut ConsumerContext,
message: Message,
) -> WebSocketResult<()> {
for consumer in &self.consumers {
consumer.on_message(context, message.clone()).await?;
}
Ok(())
}
async fn on_disconnect(&self, context: &mut ConsumerContext) -> WebSocketResult<()> {
for consumer in &self.consumers {
consumer.on_disconnect(context).await?;
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use rstest::rstest;
use tokio::sync::mpsc;
#[rstest]
#[tokio::test]
async fn test_consumer_context_creation() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let context = ConsumerContext::new(conn);
assert_eq!(context.connection.id(), "test");
}
#[rstest]
#[tokio::test]
async fn test_consumer_context_metadata() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let context =
ConsumerContext::new(conn).with_metadata("user_id".to_string(), "123".to_string());
assert_eq!(context.get_metadata("user_id").unwrap(), "123");
}
#[rstest]
#[tokio::test]
async fn test_echo_consumer_connect() {
let consumer = EchoConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
consumer.on_connect(&mut context).await.unwrap();
let msg = rx.recv().await.unwrap();
match msg {
Message::Text { data } => assert!(data.contains("Connection established")),
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_echo_consumer_message() {
let consumer = EchoConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
let msg = Message::text("Hello".to_string());
consumer.on_message(&mut context, msg).await.unwrap();
let received = rx.recv().await.unwrap();
match received {
Message::Text { data } => assert_eq!(data, "Echo: Hello"),
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_echo_consumer_binary_utf8_message() {
let consumer = EchoConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
let msg = Message::binary(b"Hello binary".to_vec());
consumer.on_message(&mut context, msg).await.unwrap();
let received = rx.recv().await.unwrap();
match received {
Message::Text { data } => assert_eq!(data, "Echo: Hello binary"),
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_echo_consumer_binary_non_utf8_message() {
let consumer = EchoConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
let msg = Message::binary(vec![0xFF, 0xFE, 0xFD]);
consumer.on_message(&mut context, msg).await.unwrap();
let received = rx.recv().await.unwrap();
match received {
Message::Text { data } => assert_eq!(data, "Echo: binary(3 bytes)"),
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_echo_consumer_handles_close_message() {
let consumer = EchoConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn.clone());
let msg = Message::Close {
code: 1000,
reason: "Normal closure".to_string(),
};
consumer.on_message(&mut context, msg).await.unwrap();
assert!(conn.is_closed().await);
let received = rx.recv().await.unwrap();
assert!(matches!(received, Message::Close { code: 1000, .. }));
}
#[rstest]
#[tokio::test]
async fn test_json_consumer_connect() {
let consumer = JsonConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
consumer.on_connect(&mut context).await.unwrap();
let msg = rx.recv().await.unwrap();
match msg {
Message::Text { data } => {
let json: serde_json::Value = serde_json::from_str(&data).unwrap();
assert_eq!(json["status"], "connected");
}
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_json_consumer_binary_valid_json() {
let consumer = JsonConsumer::new();
let (tx, mut rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
let msg = Message::binary(br#"{"key":"value"}"#.to_vec());
consumer.on_message(&mut context, msg).await.unwrap();
let received = rx.recv().await.unwrap();
match received {
Message::Text { data } => {
let json: serde_json::Value = serde_json::from_str(&data).unwrap();
assert_eq!(json["source"], "binary");
assert_eq!(json["data"]["key"], "value");
}
_ => panic!("Expected text message"),
}
}
#[rstest]
#[tokio::test]
async fn test_json_consumer_binary_invalid_utf8_returns_error() {
let consumer = JsonConsumer::new();
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
let msg = Message::binary(vec![0xFF, 0xFE]);
let result = consumer.on_message(&mut context, msg).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(matches!(err, WebSocketError::BinaryPayload(_)));
assert!(err.to_string().contains("not valid UTF-8"));
}
#[rstest]
#[tokio::test]
async fn test_json_consumer_binary_invalid_json_returns_error() {
let consumer = JsonConsumer::new();
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
let msg = Message::binary(b"not json at all".to_vec());
let result = consumer.on_message(&mut context, msg).await;
assert!(result.is_err());
let err = result.unwrap_err();
assert!(matches!(err, WebSocketError::BinaryPayload(_)));
assert!(err.to_string().contains("not valid JSON"));
}
#[rstest]
#[tokio::test]
async fn test_broadcast_consumer_disconnect_cleanup() {
let room = Arc::new(crate::room::Room::new("cleanup".to_string()));
let consumer = BroadcastConsumer::new(room.clone());
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("user1".to_string(), tx));
room.join("user1".to_string(), conn.clone()).await.unwrap();
let mut context = ConsumerContext::new(conn.clone());
consumer.on_disconnect(&mut context).await.unwrap();
assert!(conn.is_closed().await);
assert!(!room.has_client("user1").await);
}
#[rstest]
#[tokio::test]
async fn test_broadcast_consumer_disconnect_tolerates_already_removed() {
let room = Arc::new(crate::room::Room::new("tolerant".to_string()));
let consumer = BroadcastConsumer::new(room.clone());
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("ghost".to_string(), tx));
let mut context = ConsumerContext::new(conn.clone());
let result = consumer.on_disconnect(&mut context).await;
assert!(result.is_ok());
assert!(conn.is_closed().await);
}
#[rstest]
#[tokio::test]
async fn test_consumer_context_headers() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let context = ConsumerContext::new(conn)
.with_header("cookie".to_string(), "sessionid=abc123".to_string())
.with_header("origin".to_string(), "https://example.com".to_string());
assert_eq!(context.get_header("cookie").unwrap(), "sessionid=abc123");
assert_eq!(context.get_header("origin").unwrap(), "https://example.com");
}
#[rstest]
#[tokio::test]
async fn test_consumer_context_cookie_header() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let context = ConsumerContext::new(conn).with_header(
"cookie".to_string(),
"sessionid=abc123; csrftoken=xyz".to_string(),
);
assert_eq!(
context.cookie_header(),
Some("sessionid=abc123; csrftoken=xyz")
);
}
#[rstest]
#[tokio::test]
async fn test_consumer_context_cookie_header_missing() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let context = ConsumerContext::new(conn);
assert_eq!(context.cookie_header(), None);
}
#[rstest]
#[tokio::test]
async fn test_consumer_context_headers_default_empty() {
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let context = ConsumerContext::new(conn);
assert!(context.headers.is_empty());
}
#[rstest]
#[tokio::test]
async fn test_consumer_chain() {
let mut chain = ConsumerChain::new();
chain.add_consumer(Box::new(EchoConsumer::with_prefix("Consumer1".to_string())));
let (tx, _rx) = mpsc::unbounded_channel();
let conn = Arc::new(WebSocketConnection::new("test".to_string(), tx));
let mut context = ConsumerContext::new(conn);
assert!(chain.on_connect(&mut context).await.is_ok());
}
}