mz_frontegg_mock/
server.rs

1// Copyright Materialize, Inc. and contributors. All rights reserved.
2//
3// Use of this software is governed by the Business Source License
4// included in the LICENSE file.
5//
6// As of the Change Date specified in that file, in accordance with
7// the Business Source License, use of this software will be governed
8// by the Apache License, Version 2.0.
9
10use crate::utils::RefreshTokenTarget;
11use axum::routing::{delete, get, post, put};
12use axum::{Router, middleware};
13use jsonwebtoken::{DecodingKey, EncodingKey};
14use mz_ore::now::NowFn;
15use mz_ore::retry::Retry;
16use mz_ore::task::JoinHandle;
17use std::borrow::Cow;
18use std::collections::BTreeMap;
19use std::future::IntoFuture;
20use std::net::{IpAddr, Ipv4Addr, SocketAddr};
21use std::sync::atomic::AtomicBool;
22use std::sync::{Arc, Mutex};
23use std::time::Duration;
24use tokio::net::TcpListener;
25use tokio::sync::mpsc::{UnboundedReceiver, UnboundedSender, unbounded_channel};
26
27use crate::handlers::*;
28use crate::middleware::*;
29use crate::models::*;
30
31const AUTH_API_TOKEN_PATH: &str = "/identity/resources/auth/v1/api-token";
32const AUTH_USER_PATH: &str = "/identity/resources/auth/v1/user";
33const AUTH_API_TOKEN_REFRESH_PATH: &str = "/identity/resources/auth/v1/api-token/token/refresh";
34const AUTH_VENDOR_PATH: &str = "/auth/vendor";
35const GROUPS_PATH: &str = "/frontegg/identity/resources/groups/v1";
36const GROUP_PATH: &str = "/frontegg/identity/resources/groups/v1/{:id}";
37const GROUP_ROLES_PATH: &str = "/frontegg/identity/resources/groups/v1/{:id}/roles";
38const GROUP_USERS_PATH: &str = "/frontegg/identity/resources/groups/v1/{:id}/users";
39const GROUP_PATH_WITH_SLASH: &str = "/frontegg/identity/resources/groups/v1/{:id}/";
40const MEMBERS_PATH: &str = "/frontegg/team/resources/members/v1";
41const USERS_ME_PATH: &str = "/identity/resources/users/v2/me";
42const USERS_API_TOKENS_PATH: &str = "/identity/resources/users/api-tokens/v1";
43const USER_API_TOKENS_PATH: &str = "/identity/resources/users/api-tokens/v1/{:id}";
44const TENANT_API_TOKENS_PATH: &str = "/identity/resources/tenants/api-tokens/v1";
45const TENANT_API_TOKEN_PATH: &str = "/identity/resources/tenants/api-tokens/v1/{:id}";
46const USER_PATH: &str = "/identity/resources/users/v1/{:id}";
47const USER_CREATE_PATH: &str = "/identity/resources/users/v2";
48const USERS_V3_PATH: &str = "/identity/resources/users/v3";
49const ROLES_PATH: &str = "/identity/resources/roles/v2";
50const SCIM_CONFIGURATIONS_PATH: &str = "/frontegg/directory/resources/v1/configurations/scim2";
51const SCIM_CONFIGURATION_PATH: &str = "/frontegg/directory/resources/v1/configurations/scim2/{:id}";
52const TENANTS_PATH: &str = "/tenants/resources/tenants/v1";
53const TENANT_PATH: &str = "/tenants/resources/tenants/v1/{:id}";
54const TENANT_METADATA_PATH: &str = "/tenants/resources/tenants/v1/{:id}/metadata";
55const TENANT_METADATA_KEY_PATH: &str = "/tenants/resources/tenants/v1/{:id}/metadata/{:key}";
56const SSO_CONFIGS_PATH: &str = "/frontegg/team/resources/sso/v1/configurations";
57const SSO_CONFIG_PATH: &str = "/frontegg/team/resources/sso/v1/configurations/{:id}";
58const SSO_CONFIG_DOMAINS_PATH: &str =
59    "/frontegg/team/resources/sso/v1/configurations/{:id}/domains";
60const SSO_CONFIG_DOMAIN_PATH: &str =
61    "/frontegg/team/resources/sso/v1/configurations/{:id}/domains/{:domain_id}";
62const SSO_CONFIG_GROUPS_PATH: &str = "/frontegg/team/resources/sso/v1/configurations/{:id}/groups";
63const SSO_CONFIG_GROUP_PATH: &str =
64    "/frontegg/team/resources/sso/v1/configurations/{:id}/groups/{:group_id}";
65const SSO_CONFIG_ROLES_PATH: &str = "/frontegg/team/resources/sso/v1/configurations/{:id}/roles";
66
67// Internal endpoints for testing
68const INTERNAL_USER_PASSWORD_PATH: &str = "/api/internal-mock/user-password";
69
70pub struct FronteggMockServer {
71    pub base_url: String,
72    pub refreshes: Arc<Mutex<u64>>,
73    pub enable_auth: Arc<AtomicBool>,
74    pub auth_requests: Arc<Mutex<u64>>,
75    pub role_updates_tx: UnboundedSender<(String, Vec<String>)>,
76    pub handle: JoinHandle<Result<(), std::io::Error>>,
77}
78
79impl FronteggMockServer {
80    /// Starts a [`FronteggMockServer`], must be started from within a [`tokio::runtime::Runtime`].
81    pub async fn start(
82        addr: Option<&SocketAddr>,
83        issuer: String,
84        encoding_key: EncodingKey,
85        decoding_key: DecodingKey,
86        users: BTreeMap<String, UserConfig>,
87        tenant_api_tokens: BTreeMap<ApiToken, TenantApiTokenConfig>,
88        role_permissions: Option<BTreeMap<String, Vec<String>>>,
89        now: NowFn,
90        expires_in_secs: i64,
91        latency: Option<Duration>,
92        roles: Option<Vec<UserRole>>,
93    ) -> Result<FronteggMockServer, anyhow::Error> {
94        let (role_updates_tx, role_updates_rx) = unbounded_channel();
95
96        let enable_auth = Arc::new(AtomicBool::new(true));
97        let refreshes = Arc::new(Mutex::new(0u64));
98        let auth_requests = Arc::new(Mutex::new(0u64));
99
100        let user_api_tokens: BTreeMap<ApiToken, String> = users
101            .iter()
102            .map(|(email, user)| {
103                user.initial_api_tokens
104                    .iter()
105                    .map(|token| (token.clone(), email.clone()))
106            })
107            .flatten()
108            .collect();
109        let role_permissions = role_permissions.unwrap_or_else(|| {
110            BTreeMap::from([
111                (
112                    "MaterializePlatformAdmin".to_owned(),
113                    vec![
114                        "materialize.environment.write".to_owned(),
115                        "materialize.invoice.read".to_owned(),
116                    ],
117                ),
118                (
119                    "MaterializePlatform".to_owned(),
120                    vec!["materialize.environment.read".to_owned()],
121                ),
122            ])
123        });
124
125        // Provide default roles if None is provided
126        let roles = roles.unwrap_or_else(|| {
127            vec![
128                UserRole {
129                    id: uuid::Uuid::new_v4().to_string(),
130                    name: "Organization Admin".to_string(),
131                    key: "MaterializePlatformAdmin".to_string(),
132                },
133                UserRole {
134                    id: uuid::Uuid::new_v4().to_string(),
135                    name: "Organization Member".to_string(),
136                    key: "MaterializePlatform".to_string(),
137                },
138            ]
139        });
140
141        let context = Arc::new(Context {
142            issuer,
143            encoding_key,
144            decoding_key,
145            users: Mutex::new(users),
146            user_api_tokens: Mutex::new(user_api_tokens),
147            tenant_api_tokens: Mutex::new(tenant_api_tokens),
148            role_updates_rx: Mutex::new(role_updates_rx),
149            role_permissions,
150            now,
151            expires_in_secs,
152            latency,
153            refresh_tokens: Mutex::new(BTreeMap::new()),
154            refreshes: Arc::clone(&refreshes),
155            enable_auth: Arc::clone(&enable_auth),
156            auth_requests: Arc::clone(&auth_requests),
157            roles: Arc::new(roles),
158            sso_configs: Mutex::new(BTreeMap::new()),
159            groups: Mutex::new(BTreeMap::new()),
160            scim_configurations: Mutex::new(BTreeMap::new()),
161            tenants: Mutex::new(BTreeMap::new()),
162        });
163
164        let router = Router::new()
165            .route(AUTH_API_TOKEN_PATH, post(handle_post_auth_api_token))
166            .route(AUTH_USER_PATH, post(handle_post_auth_user))
167            .route(AUTH_API_TOKEN_REFRESH_PATH, post(handle_post_token_refresh))
168            .route(AUTH_VENDOR_PATH, post(handle_post_auth_vendor))
169            .route(USERS_ME_PATH, get(handle_get_user_profile))
170            .route(
171                USERS_API_TOKENS_PATH,
172                get(handle_list_user_api_tokens).post(handle_post_user_api_token),
173            )
174            .route(USER_API_TOKENS_PATH, delete(handle_delete_user_api_token))
175            .route(
176                TENANT_API_TOKENS_PATH,
177                get(handle_list_tenant_api_tokens).post(handle_create_tenant_api_token),
178            )
179            .route(
180                TENANT_API_TOKEN_PATH,
181                delete(handle_delete_tenant_api_token),
182            )
183            .route(USER_PATH, get(handle_get_user).delete(handle_delete_user))
184            .route(USER_CREATE_PATH, post(handle_create_user))
185            .route(USERS_V3_PATH, get(handle_get_users_v3))
186            .route(MEMBERS_PATH, put(handle_update_user_roles))
187            .route(ROLES_PATH, get(handle_roles_request))
188            .route(
189                SSO_CONFIGS_PATH,
190                get(handle_list_sso_configs).post(handle_create_sso_config),
191            )
192            .route(
193                SSO_CONFIG_PATH,
194                get(handle_get_sso_config)
195                    .patch(handle_update_sso_config)
196                    .delete(handle_delete_sso_config),
197            )
198            .route(
199                SSO_CONFIG_DOMAINS_PATH,
200                get(handle_list_domains).post(handle_create_domain),
201            )
202            .route(
203                SSO_CONFIG_DOMAIN_PATH,
204                get(handle_get_domain)
205                    .patch(handle_update_domain)
206                    .delete(handle_delete_domain),
207            )
208            .route(
209                SSO_CONFIG_GROUPS_PATH,
210                get(handle_list_group_mappings).post(handle_create_group_mapping),
211            )
212            .route(
213                SSO_CONFIG_GROUP_PATH,
214                get(handle_get_group_mapping)
215                    .patch(handle_update_group_mapping)
216                    .delete(handle_delete_group_mapping),
217            )
218            .route(
219                SSO_CONFIG_ROLES_PATH,
220                get(handle_get_default_roles).put(handle_set_default_roles),
221            )
222            .route(
223                GROUPS_PATH,
224                get(handle_list_groups).post(handle_create_group),
225            )
226            .route(
227                GROUP_PATH,
228                get(handle_get_group)
229                    .patch(handle_update_group)
230                    .delete(handle_delete_group),
231            )
232            .route(GROUP_PATH_WITH_SLASH, get(handle_get_group))
233            .route(
234                GROUP_ROLES_PATH,
235                post(handle_add_roles_to_group).delete(handle_remove_roles_from_group),
236            )
237            .route(
238                GROUP_USERS_PATH,
239                post(handle_add_users_to_group).delete(handle_remove_users_from_group),
240            )
241            .route(
242                SCIM_CONFIGURATIONS_PATH,
243                get(handle_list_scim_configurations).post(handle_create_scim_configuration),
244            )
245            .route(
246                SCIM_CONFIGURATION_PATH,
247                delete(handle_delete_scim_configuration),
248            )
249            .route(
250                TENANTS_PATH,
251                get(handle_list_tenants).post(handle_create_tenant),
252            )
253            .route(
254                TENANT_PATH,
255                get(handle_get_tenant).delete(handle_delete_tenant),
256            )
257            .route(TENANT_METADATA_PATH, post(handle_set_tenant_metadata))
258            .route(
259                TENANT_METADATA_KEY_PATH,
260                delete(handle_delete_tenant_metadata),
261            )
262            .route(
263                INTERNAL_USER_PASSWORD_PATH,
264                post(internal_handle_get_user_password),
265            )
266            .layer(middleware::from_fn(logging_middleware))
267            .layer(middleware::from_fn_with_state(
268                Arc::clone(&context),
269                latency_middleware,
270            ))
271            .layer(middleware::from_fn_with_state(
272                Arc::clone(&context),
273                role_update_middleware,
274            ))
275            .with_state(context);
276
277        let addr = match addr {
278            Some(addr) => Cow::Borrowed(addr),
279            None => Cow::Owned(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 0)),
280        };
281        let listener = TcpListener::bind(*addr).await.unwrap_or_else(|e| {
282            panic!("error binding to {}: {}", addr, e);
283        });
284        let base_url = format!("http://{}", listener.local_addr().unwrap());
285        let server = axum::serve(
286            listener,
287            router.into_make_service_with_connect_info::<SocketAddr>(),
288        );
289        let handle = mz_ore::task::spawn(|| "mzcloud-mock-server", server.into_future());
290
291        Ok(FronteggMockServer {
292            base_url,
293            refreshes,
294            enable_auth,
295            auth_requests,
296            role_updates_tx,
297            handle,
298        })
299    }
300
301    pub fn wait_for_auth(&self, expires_in_secs: u64) {
302        let expected = *self.auth_requests.lock().unwrap() + 1;
303        Retry::default()
304            .factor(1.0)
305            .max_duration(Duration::from_secs(expires_in_secs + 20))
306            .retry(|_| {
307                let refreshes = *self.auth_requests.lock().unwrap();
308                if refreshes >= expected {
309                    Ok(())
310                } else {
311                    Err(format!(
312                        "expected refresh count {}, got {}",
313                        expected, refreshes
314                    ))
315                }
316            })
317            .unwrap();
318    }
319
320    pub fn auth_api_token_url(&self) -> String {
321        format!("{}{}", self.base_url, AUTH_API_TOKEN_PATH)
322    }
323}
324
325pub struct Context {
326    pub issuer: String,
327    pub encoding_key: EncodingKey,
328    pub decoding_key: DecodingKey,
329    pub users: Mutex<BTreeMap<String, UserConfig>>,
330    pub user_api_tokens: Mutex<BTreeMap<ApiToken, String>>,
331    pub tenant_api_tokens: Mutex<BTreeMap<ApiToken, TenantApiTokenConfig>>,
332    pub role_updates_rx: Mutex<UnboundedReceiver<(String, Vec<String>)>>,
333    pub role_permissions: BTreeMap<String, Vec<String>>,
334    pub now: NowFn,
335    pub expires_in_secs: i64,
336    pub latency: Option<Duration>,
337    pub refresh_tokens: Mutex<BTreeMap<String, RefreshTokenTarget>>,
338    pub refreshes: Arc<Mutex<u64>>,
339    pub enable_auth: Arc<AtomicBool>,
340    pub auth_requests: Arc<Mutex<u64>>,
341    pub roles: Arc<Vec<UserRole>>,
342    pub sso_configs: Mutex<BTreeMap<String, SSOConfigStorage>>,
343    pub groups: Mutex<BTreeMap<String, Group>>,
344    pub scim_configurations: Mutex<BTreeMap<String, SCIM2ConfigurationStorage>>,
345    pub tenants: Mutex<BTreeMap<uuid::Uuid, TenantConfig>>,
346}
mz_frontegg_mock/server.rs

mz_frontegg_mock/
server.rs
