feat: strict behavior (#474)

server/src/builder.rs

@@ -99,6 +99,10 @@ pub(crate) fn build_offline_mode(
 }
 
 fn build_offline(offline_args: OfflineArgs) -> EdgeResult<CacheContainer> {
+    if offline_args.tokens.is_empty() {
+        return Err(EdgeError::NoTokens("No tokens provided. Tokens must be specified when running in offline mode".into()));
+    }
+
     if let Some(bootstrap) = offline_args.bootstrap_file {
         let file = File::open(bootstrap.clone()).map_err(|_| EdgeError::NoFeaturesFile)?;
 
@@ -151,6 +155,10 @@ async fn get_data_source(args: &EdgeArgs) -> Option<Arc<dyn EdgePersistence>> {
 }
 
 async fn build_edge(args: &EdgeArgs) -> EdgeResult<EdgeInfo> {
+    if args.strict && args.tokens.is_empty() {
+        return Err(EdgeError::NoTokens("No tokens provided. Tokens must be specified when running with strict behavior".into()));
+    }
+
     let (token_cache, feature_cache, engine_cache) = build_caches();
 
     let persistence = get_data_source(args).await;
@@ -183,6 +191,7 @@ async fn build_edge(args: &EdgeArgs) -> EdgeResult<EdgeInfo> {
         engine_cache.clone(),
         Duration::seconds(args.features_refresh_interval_seconds.try_into().unwrap()),
         persistence.clone(),
+        args.strict,
     ));
     let _ = token_validator.register_tokens(args.tokens.clone()).await;
 
@@ -223,3 +232,52 @@ pub async fn build_caches_and_refreshers(args: CliArgs) -> EdgeResult<EdgeInfo>
         _ => unreachable!(),
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use crate::{builder::{build_edge, build_offline}, cli::{EdgeArgs, OfflineArgs, TokenHeader}};
+
+    #[test]
+    fn should_fail_with_empty_tokens_when_offline_mode() {
+        let args = OfflineArgs {
+            bootstrap_file: None,
+            tokens: vec![],
+            reload_interval: Default::default()
+        };
+
+        let result = build_offline(args);
+        assert!(result.is_err());
+        assert_eq!(result
+          .err()
+          .unwrap()
+          .to_string(), "No tokens provided. Tokens must be specified when running in offline mode");
+    }
+
+    #[tokio::test]
+    async fn should_fail_with_empty_tokens_when_strict() {
+        let args = EdgeArgs {
+            upstream_url: Default::default(),
+            backup_folder: None,
+            metrics_interval_seconds: Default::default(),
+            features_refresh_interval_seconds: Default::default(),
+            strict: true,
+            tokens: vec![],
+            redis: None,
+            client_identity: Default::default(),
+            skip_ssl_verification: false,
+            upstream_request_timeout: Default::default(),
+            upstream_socket_timeout: Default::default(),
+            custom_client_headers: Default::default(),
+            token_header: TokenHeader { token_header: "Authorization".into() },
+            upstream_certificate_file: Default::default(),
+            token_revalidation_interval_seconds: Default::default(),
+        };
+
+        let result = build_edge(&args).await;
+        assert!(result.is_err());
+        assert_eq!(result
+          .err()
+          .unwrap()
+          .to_string(), "No tokens provided. Tokens must be specified when running with strict behavior");
+    }
+  }

server/src/cli.rs

@@ -174,6 +174,10 @@ pub struct EdgeArgs {
     /// Token header to use for both edge authorization and communication with the upstream server.
     #[clap(long, env, global = true, default_value = "Authorization")]
     pub token_header: TokenHeader,
+
+    /// If set to true, Edge starts with strict behavior. Strict behavior means that Edge will refuse tokens outside of the scope of the startup tokens
+    #[clap(long, env, default_value_t = false)]
+    pub strict: bool,
 }
 
 pub fn string_to_header_tuple(s: &str) -> Result<(String, String), String> {

server/src/client_api.rs

@@ -918,7 +918,7 @@ mod tests {
     }
 
     #[tokio::test]
-    async fn calling_client_features_endpoint_with_new_token_hydrates_from_upstream() {
+    async fn calling_client_features_endpoint_with_new_token_hydrates_from_upstream_when_dynamic() {
         let upstream_features_cache: Arc<DashMap<String, ClientFeatures>> =
             Arc::new(DashMap::default());
         let upstream_token_cache: Arc<DashMap<String, EdgeToken>> = Arc::new(DashMap::default());
@@ -949,6 +949,7 @@ mod tests {
             engine_cache: engine_cache.clone(),
             refresh_interval: Duration::seconds(6000),
             persistence: None,
+            strict: false,
         });
         let token_validator = Arc::new(TokenValidator {
             unleash_client: unleash_client.clone(),
@@ -977,6 +978,65 @@ mod tests {
         assert_eq!(res.status(), StatusCode::OK);
     }
 
+    #[tokio::test]
+    async fn calling_client_features_endpoint_with_new_token_does_not_hydrate_when_strict() {
+        let upstream_features_cache: Arc<DashMap<String, ClientFeatures>> =
+            Arc::new(DashMap::default());
+        let upstream_token_cache: Arc<DashMap<String, EdgeToken>> = Arc::new(DashMap::default());
+        let upstream_engine_cache: Arc<DashMap<String, EngineState>> = Arc::new(DashMap::default());
+        let server = upstream_server(
+            upstream_token_cache.clone(),
+            upstream_features_cache.clone(),
+            upstream_engine_cache.clone(),
+        )
+        .await;
+        let upstream_features = features_from_disk("../examples/hostedexample.json");
+        let mut upstream_known_token = EdgeToken::from_str("dx:development.secret123").unwrap();
+        upstream_known_token.status = TokenValidationStatus::Validated;
+        upstream_known_token.token_type = Some(TokenType::Client);
+        upstream_token_cache.insert(
+            upstream_known_token.token.clone(),
+            upstream_known_token.clone(),
+        );
+        upstream_features_cache.insert(cache_key(&upstream_known_token), upstream_features.clone());
+        let unleash_client = Arc::new(UnleashClient::new(server.url("/").as_str(), None).unwrap());
+        let features_cache: Arc<DashMap<String, ClientFeatures>> = Arc::new(DashMap::default());
+        let token_cache: Arc<DashMap<String, EdgeToken>> = Arc::new(DashMap::default());
+        let engine_cache: Arc<DashMap<String, EngineState>> = Arc::new(DashMap::default());
+        let feature_refresher = Arc::new(FeatureRefresher {
+            unleash_client: unleash_client.clone(),
+            features_cache: features_cache.clone(),
+            engine_cache: engine_cache.clone(),
+            refresh_interval: Duration::seconds(6000),
+            ..Default::default()
+        });
+        let token_validator = Arc::new(TokenValidator {
+            unleash_client: unleash_client.clone(),
+            token_cache: token_cache.clone(),
+            persistence: None,
+        });
+        let local_app = test::init_service(
+            App::new()
+                .app_data(Data::from(token_validator.clone()))
+                .app_data(Data::from(features_cache.clone()))
+                .app_data(Data::from(engine_cache.clone()))
+                .app_data(Data::from(token_cache.clone()))
+                .app_data(Data::from(feature_refresher.clone()))
+                .wrap(middleware::as_async_middleware::as_async_middleware(
+                    middleware::validate_token::validate_token,
+                ))
+                .service(web::scope("/api").configure(configure_client_api)),
+        )
+        .await;
+        let req = test::TestRequest::get()
+            .uri("/api/client/features")
+            .insert_header(ContentType::json())
+            .insert_header(("Authorization", upstream_known_token.token.clone()))
+            .to_request();
+        let res = test::call_service(&local_app, req).await;
+        assert_eq!(res.status(), StatusCode::FORBIDDEN);
+    }
+
     #[tokio::test]
     pub async fn gets_feature_by_name() {
         let features_cache: Arc<DashMap<String, ClientFeatures>> = Arc::new(DashMap::default());
@@ -1041,7 +1101,7 @@ mod tests {
         assert_eq!(result.status(), StatusCode::NOT_FOUND);
     }
     #[tokio::test]
-    pub async fn still_subsumes_tokens_after_moving_registration_to_initial_hydration() {
+    pub async fn still_subsumes_tokens_after_moving_registration_to_initial_hydration_when_dynamic() {
         let upstream_features_cache: Arc<DashMap<String, ClientFeatures>> =
             Arc::new(DashMap::default());
         let upstream_token_cache: Arc<DashMap<String, EdgeToken>> = Arc::new(DashMap::default());
@@ -1066,13 +1126,14 @@ mod tests {
         let features_cache: Arc<DashMap<String, ClientFeatures>> = Arc::new(DashMap::default());
         let token_cache: Arc<DashMap<String, EdgeToken>> = Arc::new(DashMap::default());
         let engine_cache: Arc<DashMap<String, EngineState>> = Arc::new(DashMap::default());
-        let feature_refresher = Arc::new(FeatureRefresher::new(
-            unleash_client.clone(),
-            features_cache.clone(),
-            engine_cache.clone(),
-            Duration::seconds(6000),
-            None,
-        ));
+        let feature_refresher = Arc::new(FeatureRefresher {
+            unleash_client: unleash_client.clone(),
+            features_cache: features_cache.clone(),
+            engine_cache: engine_cache.clone(),
+            refresh_interval: Duration::seconds(6000),
+            strict: false,
+            ..Default::default()
+        });
         let token_validator = Arc::new(TokenValidator {
             unleash_client: unleash_client.clone(),
             token_cache: token_cache.clone(),

server/src/error.rs

@@ -105,10 +105,12 @@ pub enum EdgeError {
     EdgeTokenParseError,
     InvalidBackupFile(String, String),
     InvalidServerUrl(String),
+    InvalidTokenWithStrictBehavior,
     HealthCheckError(String),
     JsonParseError(String),
     NoFeaturesFile,
     NoTokenProvider,
+    NoTokens(String),
     NotReady,
     ReadyCheckError(String),
     TlsError,
@@ -129,6 +131,7 @@ impl Display for EdgeError {
             EdgeError::NoFeaturesFile => write!(f, "No features file located"),
             EdgeError::AuthorizationDenied => write!(f, "Not allowed to access"),
             EdgeError::NoTokenProvider => write!(f, "Could not get a TokenProvider"),
+            EdgeError::NoTokens(msg) => write!(f, "{msg}"),
             EdgeError::TokenParseError(token) => write!(f, "Could not parse edge token: {token}"),
             EdgeError::PersistenceError(msg) => write!(f, "{msg}"),
             EdgeError::JsonParseError(msg) => write!(f, "{msg}"),
@@ -204,6 +207,7 @@ impl Display for EdgeError {
             EdgeError::NotReady => {
                 write!(f, "Edge is not ready to serve requests")
             }
+            EdgeError::InvalidTokenWithStrictBehavior => write!(f, "Edge is running with strict behavior and the token is not subsumed by any registered tokens"),
         }
     }
 }
@@ -216,6 +220,7 @@ impl ResponseError for EdgeError {
             EdgeError::NoFeaturesFile => StatusCode::INTERNAL_SERVER_ERROR,
             EdgeError::AuthorizationDenied => StatusCode::FORBIDDEN,
             EdgeError::NoTokenProvider => StatusCode::INTERNAL_SERVER_ERROR,
+            EdgeError::NoTokens(_) => StatusCode::INTERNAL_SERVER_ERROR,
             EdgeError::TokenParseError(_) => StatusCode::FORBIDDEN,
             EdgeError::ClientBuildError(_) => StatusCode::INTERNAL_SERVER_ERROR,
             EdgeError::ClientFeaturesParseError(_) => StatusCode::INTERNAL_SERVER_ERROR,
@@ -240,6 +245,7 @@ impl ResponseError for EdgeError {
             EdgeError::ClientCacheError => StatusCode::INTERNAL_SERVER_ERROR,
             EdgeError::FrontendExpectedToBeHydrated(_) => StatusCode::INTERNAL_SERVER_ERROR,
             EdgeError::NotReady => StatusCode::SERVICE_UNAVAILABLE,
+            EdgeError::InvalidTokenWithStrictBehavior => StatusCode::FORBIDDEN,
         }
     }