Merge pull request #157 from TryCli-Studio/fix/keys

Fix/keys

Author: karthikeyjoshi

client/src/pages/view.rs

@@ -448,7 +448,9 @@ pub fn ViewPage() -> impl IntoView {
                                                             let vip = if key.is_empty() {
                                                                 String::new()
                                                             } else {
-                                                                format!("{}/{}/{}?key={}", origin_clone, username_clone, slug_clone, key)
+                                                                // URL encode the key to handle special characters in base64
+                                                                let encoded_key = js_sys::encode_uri_component(key);
+                                                                format!("{}/{}/{}?key={}", origin_clone, username_clone, slug_clone, encoded_key)
                                                             };
                                                             set_vip_link.set(vip);
                                                         }
@@ -501,7 +503,9 @@ pub fn ViewPage() -> impl IntoView {
                                                             let vip = if key.is_empty() {
                                                                 String::new()
                                                             } else {
-                                                                format!("{}/{}/{}?key={}", origin_clone, username_clone, slug_clone, key)
+                                                                // URL encode the key to handle special characters in base64
+                                                                let encoded_key = js_sys::encode_uri_component(key);
+                                                                format!("{}/{}/{}?key={}", origin_clone, username_clone, slug_clone, encoded_key)
                                                             };
                                                             set_vip_link.set(vip);
                                                         }

server/Cargo.toml

@@ -26,3 +26,4 @@ openssl-sys = { version = "0.9", features = ["vendored"] }
 url = "2.5.8"
 dashmap = "6.1"
 governor = "0.10"
+urlencoding = "2.1"

server/src/handlers/project.rs

@@ -53,6 +53,7 @@ fn normalize_url(url_str: &str) -> Option<String> {
 
     // Require a valid host for http(s) URLs
     let host = url.host_str()?;
+    let port = url.port();
     let path = url.path();
 
     // Normalize trailing slashes for consistency
@@ -62,7 +63,11 @@ fn normalize_url(url_str: &str) -> Option<String> {
         path.trim_end_matches('/')
     };
 
-    Some(format!("{}://{}{}", scheme, host, normalized_path))
+    if let Some(p) = port {
+        Some(format!("{}://{}:{}{}", scheme, host, p, normalized_path))
+    } else {
+        Some(format!("{}://{}{}", scheme, host, normalized_path))
+    }
 }
 
 /// Validates CSRF protection by checking Origin/Referer headers against expected frontend URL.
@@ -398,21 +403,47 @@ pub async fn get_project(
 
     // 4. Dual-layer security for embeds (VIP key + Guest List)
     if !is_owner {
-        // VIP Pass: compare URL ?key with stored embed_key using a match on Options
-        let vip_allowed = match (params.get("key"), embed_key.as_ref()) {
-            (Some(request_key), Some(stored_key)) if request_key == stored_key => true,
-            _ => false,
+        // VIP Pass: compare URL ?key with stored embed_key
+        // Note: Axum's Query extractor already URL-decodes parameters
+            let vip_allowed = match (params.get("key"), embed_key.as_ref()) {
+                (Some(request_key), Some(stored_key)) => {
+                    // Axum already decodes query params; treat spaces as '+' to handle
+                    // unencoded base64 keys in query strings.
+                    let trimmed_request = request_key.trim();
+                    let normalized_request = trimmed_request.replace(' ', "+");
+                    let trimmed_stored = stored_key.trim();
+
+                    let matches = trimmed_request == trimmed_stored
+                        || normalized_request == trimmed_stored;
+
+                    tracing::debug!(
+                        "VIP key comparison for project {}: request='{}' normalized='{}' stored='{}' match={}",
+                        project_id,
+                        trimmed_request,
+                        normalized_request,
+                        trimmed_stored,
+                        matches
+                    );
+
+                    matches
+                }
+            _ => {
+                tracing::debug!("VIP key check failed for project {}: request_key={:?} stored_key={:?}", 
+                    project_id, params.get("key"), embed_key.as_ref());
+                false
+            }
         };
 
         if !vip_allowed {
-            // Guest List: validate Referer header against project_whitelists
+            // Guest List: validate parent page URL against project_whitelists.
+            // Prefer X-Embed-Referer (forwarded by embed page) and fallback to Referer.
             let referer = headers
-                .get("X-Embed-Referer")
+                .get("x-embed-referer")
                 .and_then(|v| v.to_str().ok())
                 .map(|s| s.to_string())
                 .or_else(|| {
                     headers
-                        .get("Referer")
+                        .get("referer")
                         .and_then(|v| v.to_str().ok())
                         .map(|s| s.to_string())
                 });
@@ -421,7 +452,7 @@ pub async fn get_project(
             let whitelist_allowed: Option<bool> = if let Some(referer_url) = referer {
                 // Normalize the Referer URL to prevent bypasses via query params or fragments
                 let normalized_referer = normalize_url(&referer_url);
-                
+
                 if let Some(normalized) = normalized_referer {
                     let exists_row: Option<(bool,)> = sqlx::query_as(
                         "SELECT TRUE FROM project_whitelists \
@@ -870,16 +901,16 @@ pub async fn resolve_secret_embed(
     State(state): State<AppState>,
     Path(token): Path<String>,
 ) -> Result<Html<String>, (StatusCode, String)> {
-    // 1. Validate Token from DB
-    let row: Option<(String, String)> = sqlx::query_as(
-        "SELECT owner_username, slug FROM projects WHERE embed_token = $1"
+    // 1. Validate Token from DB and fetch embed_key for VIP access
+    let row: Option<(String, String, Option<String>)> = sqlx::query_as(
+        "SELECT owner_username, slug, embed_key FROM projects WHERE embed_token = $1"
     )
     .bind(&token)
     .fetch_optional(&state.db)
     .await
     .map_err(|e| (StatusCode::INTERNAL_SERVER_ERROR, e.to_string()))?;
 
-    let (owner, slug) = row.ok_or((StatusCode::NOT_FOUND, "Invalid embed token".to_string()))?;
+    let (owner, slug, embed_key) = row.ok_or((StatusCode::NOT_FOUND, "Invalid embed token".to_string()))?;
 
     // 2. Get Configurations
     let api_url = std::env::var("API_URL").unwrap_or("http://localhost:3000".to_string());
@@ -889,8 +920,14 @@ pub async fn resolve_secret_embed(
     // 3. Generate Discovery URL for OEmbed (Points to Backend)
     let oembed_url = format!("{}/api/oembed?url={}/e/{}", api_url, api_url, token);
 
-    // 4. Construct Target URL (Points to Frontend)
-    let target_url = format!("{}/embed/{}/{}", frontend_url, owner, slug);
+    // 4. Construct Target URL with VIP key for bypass (Points to Frontend)
+    let target_url = if let Some(key) = embed_key {
+        // URL encode the key to handle special characters in base64
+        let encoded_key = urlencoding::encode(&key);
+        format!("{}/embed/{}/{}?key={}", frontend_url, owner, slug, encoded_key)
+    } else {
+        format!("{}/embed/{}/{}", frontend_url, owner, slug)
+    };
 
     // 5. Return HTML with <link> tags + Meta Refresh
     let html = format!(r#"<!DOCTYPE html>

server/src/router.rs

@@ -4,7 +4,7 @@ use axum::{
     http::Method,
 };
 use tower_sessions::{Expiry, MemoryStore, SessionManagerLayer};
-use axum::http::header::{CONTENT_TYPE, AUTHORIZATION};
+use axum::http::header::{CONTENT_TYPE, AUTHORIZATION, REFERER, ORIGIN};
 use crate::state::AppState;
 use crate::handlers::{auth, project, spawn, analytics, admin, oembed};
 use crate::services::websocket;
@@ -20,6 +20,9 @@ pub fn create_router(state: AppState) -> Result<Router, Box<dyn std::error::Erro
     let frontend_url = std::env::var("FRONTEND_URL")
         .unwrap_or_else(|_| "http://localhost:8080".to_string());
 
+    // Custom header for embed referer
+    let x_embed_referer = axum::http::HeaderName::from_static("x-embed-referer");
+
     let app = Router::new()
         .merge(auth::routes())
         .merge(spawn::routes())
@@ -32,7 +35,7 @@ pub fn create_router(state: AppState) -> Result<Router, Box<dyn std::error::Erro
             .allow_origin(frontend_url.parse::<axum::http::HeaderValue>()?)
             // 2. ALLOW DELETE METHOD
             .allow_methods([Method::GET, Method::POST, Method::DELETE, Method::OPTIONS])
-            .allow_headers([CONTENT_TYPE, AUTHORIZATION])
+            .allow_headers([CONTENT_TYPE, AUTHORIZATION, REFERER, ORIGIN, x_embed_referer])
             .allow_credentials(true) 
         )
         .layer(session_layer)