diff --git a/Cargo.toml b/Cargo.toml index 056c3e4..e80c695 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -24,4 +24,6 @@ url = "2.5.0" itertools = "0.14.0" chrono = { version = "^0.4.38", features = ["serde"] } hex = "0.4.3" -m3u8-rs = "6.0.0" \ No newline at end of file +m3u8-rs = "6.0.0" +sha2 = "0.10.8" +data-encoding = "2.9.0" \ No newline at end of file diff --git a/crates/core/Cargo.toml b/crates/core/Cargo.toml index dce4db0..14af795 100644 --- a/crates/core/Cargo.toml +++ b/crates/core/Cargo.toml @@ -29,6 +29,8 @@ hex.workspace = true itertools.workspace = true futures-util = "0.3.30" m3u8-rs.workspace = true +sha2.workspace = true +data-encoding.workspace = true # srt srt-tokio = { version = "0.4.3", optional = true } diff --git a/crates/core/src/viewer.rs b/crates/core/src/viewer.rs index 52d28aa..f1d107f 100644 --- a/crates/core/src/viewer.rs +++ b/crates/core/src/viewer.rs @@ -4,6 +4,8 @@ use std::time::{Duration, Instant}; use uuid::Uuid; use tokio::task; use log::debug; +use sha2::{Digest, Sha256}; +use data_encoding::BASE32_NOPAD; #[derive(Debug, Clone)] pub struct ViewerInfo { @@ -35,8 +37,23 @@ impl ViewerTracker { tracker } - pub fn generate_viewer_token() -> String { - Uuid::new_v4().to_string() + pub fn generate_viewer_token(ip_address: &str, user_agent: Option<&str>) -> String { + // Create input string by combining IP address and user agent + let input = match user_agent { + Some(ua) => format!("{}{}", ip_address, ua), + None => ip_address.to_string(), + }; + + // Hash the input using SHA-256 + let mut hasher = Sha256::new(); + hasher.update(input.as_bytes()); + let hash = hasher.finalize(); + + // Take the first 8 bytes of the hash + let fingerprint = &hash[..8]; + + // Base32 encode the fingerprint (without padding) + BASE32_NOPAD.encode(fingerprint).to_lowercase() } pub fn track_viewer(&self, token: &str, stream_id: &str, ip_address: &str, user_agent: Option) { @@ -111,4 +128,72 @@ impl Default for ViewerTracker { fn default() -> Self { Self::new() } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_generate_viewer_token_consistency() { + // Test that the same IP and user agent always generate the same token + let ip = "192.168.1.1"; + let user_agent = Some("Mozilla/5.0 (Test Browser)"); + + let token1 = ViewerTracker::generate_viewer_token(ip, user_agent); + let token2 = ViewerTracker::generate_viewer_token(ip, user_agent); + + assert_eq!(token1, token2, "Same IP and user agent should generate identical tokens"); + } + + #[test] + fn test_generate_viewer_token_different_inputs() { + // Test that different inputs generate different tokens + let ip1 = "192.168.1.1"; + let ip2 = "192.168.1.2"; + let user_agent = Some("Mozilla/5.0 (Test Browser)"); + + let token1 = ViewerTracker::generate_viewer_token(ip1, user_agent); + let token2 = ViewerTracker::generate_viewer_token(ip2, user_agent); + + assert_ne!(token1, token2, "Different IPs should generate different tokens"); + } + + #[test] + fn test_generate_viewer_token_no_user_agent() { + // Test that tokens are generated even without user agent + let ip = "192.168.1.1"; + + let token1 = ViewerTracker::generate_viewer_token(ip, None); + let token2 = ViewerTracker::generate_viewer_token(ip, None); + + assert_eq!(token1, token2, "Same IP without user agent should generate identical tokens"); + } + + #[test] + fn test_generate_viewer_token_format() { + // Test that generated tokens have the expected base32 format + let ip = "192.168.1.1"; + let user_agent = Some("Mozilla/5.0 (Test Browser)"); + + let token = ViewerTracker::generate_viewer_token(ip, user_agent); + + // Should be base32 encoded (lowercase, no padding) + assert!(token.chars().all(|c| "abcdefghijklmnopqrstuvwxyz234567".contains(c)), + "Token should only contain base32 characters"); + assert!(token.len() > 10, "Token should be reasonably long"); + } + + #[test] + fn test_generate_viewer_token_different_user_agents() { + // Test that different user agents generate different tokens + let ip = "192.168.1.1"; + let user_agent1 = Some("Mozilla/5.0 (Chrome)"); + let user_agent2 = Some("Mozilla/5.0 (Firefox)"); + + let token1 = ViewerTracker::generate_viewer_token(ip, user_agent1); + let token2 = ViewerTracker::generate_viewer_token(ip, user_agent2); + + assert_ne!(token1, token2, "Different user agents should generate different tokens"); + } } \ No newline at end of file diff --git a/crates/zap-stream/Cargo.toml b/crates/zap-stream/Cargo.toml index 32b241b..217995c 100644 --- a/crates/zap-stream/Cargo.toml +++ b/crates/zap-stream/Cargo.toml @@ -38,7 +38,7 @@ nostr-sdk = { version = "0.38.0" } fedimint-tonic-lnd = { version = "0.2.0", default-features = false, features = ["invoicesrpc", "versionrpc"] } reqwest = { version = "0.12.9", features = ["stream", "json"] } base64 = { version = "0.22.1" } -sha2 = { version = "0.10.8" } +sha2.workspace = true pretty_env_logger = "0.5.0" clap = { version = "4.5.16", features = ["derive"] } futures-util = "0.3.31" diff --git a/crates/zap-stream/src/http.rs b/crates/zap-stream/src/http.rs index 9de9383..9a8c2c6 100644 --- a/crates/zap-stream/src/http.rs +++ b/crates/zap-stream/src/http.rs @@ -71,8 +71,8 @@ impl HttpServer { api.track_viewer(token, stream_id, &client_ip, user_agent.clone()); token.clone() } else { - // Generate new viewer token - let token = ViewerTracker::generate_viewer_token(); + // Generate new viewer token based on IP and user agent fingerprint + let token = ViewerTracker::generate_viewer_token(&client_ip, user_agent.as_deref()); api.track_viewer(&token, stream_id, &client_ip, user_agent); token };