diff --git a/ostp-server/src/api.rs b/ostp-server/src/api.rs index 4e060b0..52f2527 100644 --- a/ostp-server/src/api.rs +++ b/ostp-server/src/api.rs @@ -54,6 +54,7 @@ pub struct ApiState { pub dns_server: std::sync::Arc, pub audit_logs: Arc>>, pub router: std::sync::Arc, + pub is_licensed: bool, } #[derive(Debug, Clone, Serialize, Deserialize)] @@ -208,7 +209,8 @@ pub fn create_api_router(state: ApiState) -> Router { .delete(handle_clear_audit), ) .route("/users/bulk", post(handle_bulk_create_users)) - .route("/router/rules", get(handle_get_rules).put(handle_put_rules)); + .route("/router/rules", get(handle_get_rules).put(handle_put_rules)) + .layer(axum::middleware::from_fn_with_state(state.clone(), license_middleware)); let webpath = state.webpath.clone(); let webpath = webpath.trim_matches('/'); @@ -236,6 +238,25 @@ pub fn create_api_router(state: ApiState) -> Router { .layer(cors) .with_state(state) } +async fn license_middleware( + axum::extract::State(state): axum::extract::State, + req: axum::extract::Request, + next: axum::middleware::Next, +) -> axum::response::Response { + if state.is_licensed { + return next.run(req).await; + } + + let path = req.uri().path(); + // Allow read-only access to users for relay, and server status + if (path == "/server/status" && req.method() == axum::http::Method::GET) || + (path == "/users" && req.method() == axum::http::Method::GET) + { + return next.run(req).await; + } + + (axum::http::StatusCode::PAYMENT_REQUIRED, "This feature requires an active OSTP license. Get yours at https://ostp.ospab.lol").into_response() +} /// Start the Management API server on the configured bind address. pub async fn start_api_server( @@ -247,6 +268,7 @@ pub async fn start_api_server( config_path: Option, dns_server: std::sync::Arc, router: std::sync::Arc, + is_licensed: bool, ) { let state = ApiState { access_keys, @@ -263,6 +285,7 @@ pub async fn start_api_server( dns_server, audit_logs: Arc::new(RwLock::new(Vec::new())), router, + is_licensed, }; let app = create_api_router(state); diff --git a/ostp-server/src/dispatcher.rs b/ostp-server/src/dispatcher.rs index dbfe7e6..3ebb6d4 100644 --- a/ostp-server/src/dispatcher.rs +++ b/ostp-server/src/dispatcher.rs @@ -81,12 +81,11 @@ pub struct Dispatcher { replay_cache: std::collections::HashMap, u64>, roaming_tokens: f64, last_token_regen: std::time::Instant, - max_sessions: Option, } #[allow(dead_code)] impl Dispatcher { - pub fn new(machine_config: ProtocolConfig, access_keys: Arc>>, max_sessions: Option) -> Self { + pub fn new(machine_config: ProtocolConfig, access_keys: Arc>>) -> Self { let mut initial_stats = HashMap::new(); for (key, meta) in access_keys.read().unwrap_or_else(|e| e.into_inner()).iter() { initial_stats.insert(key.clone(), Arc::new(UserStats::new(meta.limit_bytes))); @@ -100,7 +99,6 @@ impl Dispatcher { replay_cache: std::collections::HashMap::new(), roaming_tokens: 50.0, last_token_regen: std::time::Instant::now(), - max_sessions, } } @@ -373,11 +371,6 @@ impl Dispatcher { tracing::warn!("Replay cache full (100000 entries), rejecting handshake from {}", peer); return Ok(DispatchOutcome::Unauthorized); } - let limit = self.max_sessions.unwrap_or(30); - if self.peer_machines.len() >= limit { - tracing::warn!("drop session by {}, for more active clients buy our license here: https://ostp.ospab.lol/license", peer.ip()); - return Ok(DispatchOutcome::Unauthorized); - } self.replay_cache.insert(payload.to_vec(), ts); diff --git a/ostp-server/src/lib.rs b/ostp-server/src/lib.rs index 42d350f..bf702d9 100644 --- a/ostp-server/src/lib.rs +++ b/ostp-server/src/lib.rs @@ -117,16 +117,15 @@ pub async fn run_server( mtu: 1350, }; - let mut max_sessions = Some(30); + let mut is_licensed = false; + if let Some(key) = license_key { let host = server_public_ip.as_deref().unwrap_or("0.0.0.0"); match crate::license::verify_license(&key, host) { Ok(payload) => { tracing::info!("License verified successfully! Features: {:?}", payload.features); - if payload.features.contains(&"unlimited_connections".to_string()) { - max_sessions = None; - tracing::info!("Unlimited connections enabled."); - } + is_licensed = true; + if payload.features.contains(&"control_panel".to_string()) { tracing::info!("Spawning control panel child process..."); @@ -152,11 +151,9 @@ pub async fn run_server( tracing::error!("Failed to verify license: {:?}", e); } } - } else { - tracing::info!("No license key provided. Free version limited to 30 sessions."); } - let dispatcher = Dispatcher::new(protocol_config, shared_keys.clone(), max_sessions); + let dispatcher = Dispatcher::new(protocol_config, shared_keys.clone()); // Background config hot-reloader for access keys let shared_keys_clone = shared_keys.clone(); @@ -307,7 +304,7 @@ pub async fn run_server( let dns_server_api = dns_server.clone(); let router_api = router.clone(); tokio::spawn(async move { - api::start_api_server(api_cfg, api_keys, api_stats, server_host, server_port, config_path_api, dns_server_api, router_api).await; + api::start_api_server(api_cfg, api_keys, api_stats, server_host, server_port, config_path_api, dns_server_api, router_api, is_licensed).await; }); } }