i/Aether
1
0
Fork 0
mirror of https://github.com/fawney19/Aether.git synced 2026-10-07 01:47:47 +08:00
Code Issues Packages Projects Releases Wiki Activity

feat(security): harden gateway boundaries and usage policies

Browse Source 
Consolidate subscription usage policy enforcement, privacy-safe persistence, and gateway security hardening into one reviewable change.

Includes bounded HTTP and execution envelopes, header and protocol guards, DNS and relay validation, authentication and secret projection hardening, secure backup/install paths, and regression coverage.
This commit is contained in:
elky
2026-09-04 03:45:52 +08:00
parent ddcbeb3ae9
commit 579f2c7cc1
1019 changed files with 190437 additions and 26080 deletions
Download Patch File Download Diff File Expand all files Collapse all files
apps/aether-gateway/src/data/decision_trace.rs
+1 -1
View File
@@ -172,7 +172,7 @@ mod tests {
candidates: vec![DecisionTraceCandidate {
candidate: sample_candidate("req-1"),
provider_name: Some("OpenAI".to_string()),
provider_website: Some("https://openai.com".to_string()),
provider_website: Some("https://openai.com/".to_string()),
provider_type: Some("custom".to_string()),
provider_priority: Some(0),
provider_keep_priority_on_conversion: Some(false),
apps/aether-gateway/src/data/state/auth.rs
+1334 -68
View File
File diff suppressed because it is too large Load Diff
apps/aether-gateway/src/data/state/catalog.rs
+275 -47
View File
@@ -1,23 +1,42 @@
use super::{
ApiKeyLastUsedDelta, DataLayerError, GatewayDataState, GeminiFileMappingListQuery,
GeminiFileMappingStats, ProviderCatalogKeyAdaptiveStateUpdate,
ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, PublicHealthStatusCount, PublicHealthTimelineBucket,
StoredGeminiFileMapping, StoredGeminiFileMappingListPage, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
StoredRequestCandidate, UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
ProviderCatalogKeyAdminCasUpdate, ProviderCatalogKeyCredentialsCasUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogProviderConfigCasUpdate, ProviderCatalogProxyCasUpdate, PublicHealthStatusCount,
PublicHealthTimelineBucket, StoredGeminiFileMapping, StoredGeminiFileMappingListPage,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider, StoredRequestCandidate,
UpsertGeminiFileMappingRecord, UpsertRequestCandidateRecord,
};
fn sanitize_request_candidate_rows(
mut candidates: Vec<StoredRequestCandidate>,
) -> Vec<StoredRequestCandidate> {
for candidate in &mut candidates {
candidate.sanitize_sensitive_diagnostics();
}
candidates
}
fn sanitize_request_candidate_row(mut candidate: StoredRequestCandidate) -> StoredRequestCandidate {
candidate.sanitize_sensitive_diagnostics();
candidate
}
impl GatewayDataState {
pub(crate) async fn list_request_candidates_by_request_id(
&self,
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_by_request_id(request_id).await,
Some(repository) => repository
.list_by_request_id(request_id)
.await
.map(sanitize_request_candidate_rows),
None => Ok(Vec::new()),
}
}
@@ -27,7 +46,10 @@ impl GatewayDataState {
request_id: &str,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_attempted_by_request_id(request_id).await,
Some(repository) => repository
.list_attempted_by_request_id(request_id)
.await
.map(sanitize_request_candidate_rows),
None => Ok(Vec::new()),
}
}
@@ -38,7 +60,10 @@ impl GatewayDataState {
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_by_provider_id(provider_id, limit).await,
Some(repository) => repository
.list_by_provider_id(provider_id, limit)
.await
.map(sanitize_request_candidate_rows),
None => Ok(Vec::new()),
}
}
@@ -48,7 +73,10 @@ impl GatewayDataState {
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => repository.list_recent(limit).await,
Some(repository) => repository
.list_recent(limit)
.await
.map(sanitize_request_candidate_rows),
None => Ok(Vec::new()),
}
}
@@ -60,11 +88,10 @@ impl GatewayDataState {
limit: usize,
) -> Result<Vec<StoredRequestCandidate>, DataLayerError> {
match &self.request_candidate_reader {
Some(repository) => {
repository
.list_finalized_by_endpoint_ids_since(endpoint_ids, since_unix_secs, limit)
.await
}
Some(repository) => repository
.list_finalized_by_endpoint_ids_since(endpoint_ids, since_unix_secs, limit)
.await
.map(sanitize_request_candidate_rows),
None => Ok(Vec::new()),
}
}
@@ -108,14 +135,19 @@ impl GatewayDataState {
pub(crate) async fn upsert_request_candidate(
&self,
candidate: UpsertRequestCandidateRecord,
mut candidate: UpsertRequestCandidateRecord,
) -> Result<Option<StoredRequestCandidate>, DataLayerError> {
candidate.sanitize_for_persistence();
crate::request_diagnostics::observe_db_operation(
"request_candidate_upsert",
self.database_pool_summary(),
async {
match &self.request_candidate_writer {
Some(repository) => repository.upsert(candidate).await.map(Some),
Some(repository) => repository
.upsert(candidate)
.await
.map(sanitize_request_candidate_row)
.map(Some),
None => Ok(None),
}
},
@@ -170,6 +202,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn upsert_gemini_file_mapping_if_owner_matches(
&self,
record: UpsertGeminiFileMappingRecord,
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => repository.upsert_if_owner_matches(record).await,
None => Ok(None),
}
}
pub(crate) async fn list_gemini_file_mappings(
&self,
query: &GeminiFileMappingListQuery,
@@ -183,6 +225,49 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_gemini_file_mapping_by_file_name(
&self,
file_name: &str,
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
match &self.gemini_file_mapping_reader {
Some(repository) => repository.find_by_file_name(file_name).await,
None => Ok(None),
}
}
pub(crate) async fn find_active_gemini_file_mapping_for_user(
&self,
file_name: &str,
user_id: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
match &self.gemini_file_mapping_reader {
Some(repository) => {
repository
.find_active_by_file_name_for_user(file_name, user_id, now_unix_secs)
.await
}
None => Ok(None),
}
}
pub(crate) async fn find_active_gemini_file_mapping_for_owner(
&self,
file_name: &str,
key_id: &str,
user_id: &str,
now_unix_secs: u64,
) -> Result<Option<StoredGeminiFileMapping>, DataLayerError> {
match &self.gemini_file_mapping_reader {
Some(repository) => {
repository
.find_active_by_file_name_for_owner(file_name, key_id, user_id, now_unix_secs)
.await
}
None => Ok(None),
}
}
pub(crate) async fn summarize_gemini_file_mappings(
&self,
now_unix_secs: u64,
@@ -208,6 +293,37 @@ impl GatewayDataState {
}
}
pub(crate) async fn delete_gemini_file_mapping_by_file_name_for_user(
&self,
file_name: &str,
user_id: &str,
) -> Result<bool, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => {
repository
.delete_by_file_name_for_user(file_name, user_id)
.await
}
None => Ok(false),
}
}
pub(crate) async fn delete_gemini_file_mapping_by_file_name_for_owner(
&self,
file_name: &str,
key_id: &str,
user_id: &str,
) -> Result<bool, DataLayerError> {
match &self.gemini_file_mapping_writer {
Some(repository) => {
repository
.delete_by_file_name_for_owner(file_name, key_id, user_id)
.await
}
None => Ok(false),
}
}
pub(crate) async fn delete_gemini_file_mapping_by_id(
&self,
mapping_id: &str,
@@ -357,38 +473,11 @@ impl GatewayDataState {
}
}
pub(crate) async fn update_provider_catalog_key_oauth_credentials(
&self,
key_id: &str,
encrypted_api_key: &str,
encrypted_auth_config: Option<&str>,
expires_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => {
repository
.update_key_oauth_credentials(
key_id,
encrypted_api_key,
encrypted_auth_config,
expires_at_unix_secs,
)
.await
}
None => Ok(false),
}?;
if updated {
self.clear_provider_catalog_cache();
}
Ok(updated)
}
pub(crate) async fn update_provider_catalog_key_oauth_runtime_state(
&self,
key_id: &str,
oauth_invalid_at_unix_secs: Option<u64>,
oauth_invalid_reason: Option<&str>,
encrypted_auth_config_update: Option<&str>,
updated_at_unix_secs: Option<u64>,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
@@ -398,7 +487,6 @@ impl GatewayDataState {
key_id,
oauth_invalid_at_unix_secs,
oauth_invalid_reason,
encrypted_auth_config_update,
updated_at_unix_secs,
)
.await
@@ -475,6 +563,30 @@ impl GatewayDataState {
Ok(updated)
}
pub(crate) async fn compare_and_swap_provider_catalog_provider_config(
&self,
update: &ProviderCatalogProviderConfigCasUpdate,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => repository.compare_and_swap_provider_config(update).await,
None => Ok(false),
}?;
self.clear_provider_catalog_cache();
Ok(updated)
}
pub(crate) async fn compare_and_swap_provider_catalog_provider_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => repository.compare_and_swap_provider_proxy(update).await,
None => Ok(false),
}?;
self.clear_provider_catalog_cache();
Ok(updated)
}
pub(crate) async fn delete_provider_catalog_provider(
&self,
provider_id: &str,
@@ -543,6 +655,18 @@ impl GatewayDataState {
Ok(updated)
}
pub(crate) async fn compare_and_swap_provider_catalog_endpoint_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => repository.compare_and_swap_endpoint_proxy(update).await,
None => Ok(false),
}?;
self.clear_provider_catalog_cache();
Ok(updated)
}
pub(crate) async fn delete_provider_catalog_endpoint(
&self,
endpoint_id: &str,
@@ -571,6 +695,32 @@ impl GatewayDataState {
Ok(updated)
}
pub(crate) async fn compare_and_swap_provider_catalog_key_proxy(
&self,
update: &ProviderCatalogProxyCasUpdate,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => repository.compare_and_swap_key_proxy(update).await,
None => Ok(false),
}?;
self.clear_provider_catalog_cache();
Ok(updated)
}
pub(crate) async fn compare_and_swap_provider_catalog_key_credentials(
&self,
update: &ProviderCatalogKeyCredentialsCasUpdate,
) -> Result<bool, DataLayerError> {
let updated = match &self.provider_catalog_writer {
Some(repository) => repository.compare_and_swap_key_credentials(update).await,
None => Ok(false),
}?;
// Clear on both outcomes: a CAS miss proves the cached credential
// generation was stale and the retry must observe the winning record.
self.clear_provider_catalog_cache();
Ok(updated)
}
pub(crate) async fn compare_and_update_provider_catalog_key_admin_state(
&self,
update: &ProviderCatalogKeyAdminCasUpdate,
@@ -843,3 +993,81 @@ impl GatewayDataState {
Ok(updated)
}
}
#[cfg(test)]
mod request_candidate_security_tests {
use serde_json::json;
use super::{
sanitize_request_candidate_row, sanitize_request_candidate_rows, StoredRequestCandidate,
};
use aether_data_contracts::repository::candidates::RequestCandidateStatus;
fn untrusted_candidate() -> StoredRequestCandidate {
let mut candidate = StoredRequestCandidate::new(
"candidate-untrusted".to_string(),
"request-1".to_string(),
None,
None,
None,
None,
0,
0,
Some("provider-1".to_string()),
Some("endpoint-1".to_string()),
Some("key-1".to_string()),
RequestCandidateStatus::Failed,
None,
false,
Some(500),
None,
None,
None,
None,
None,
None,
1,
None,
Some(2),
)
.expect("candidate should build");
candidate.skip_reason = Some("Bearer candidate-secret".to_string());
candidate.error_type = Some("candidate-secret".to_string());
candidate.error_message = Some("Bearer candidate-secret".to_string());
candidate.extra_data = Some(json!({
"gateway_execution_runtime": true,
"request_body": {"token": "candidate-secret"}
}));
candidate.required_capabilities = Some(json!({
"vision": 1,
"tenant_secret": "candidate-secret"
}));
candidate
}
fn assert_candidate_is_sanitized(candidate: &StoredRequestCandidate) {
assert_eq!(candidate.skip_reason.as_deref(), Some("unclassified_skip"));
assert_eq!(candidate.error_type.as_deref(), Some("unclassified_error"));
assert!(candidate.error_message.is_none());
assert_eq!(
candidate.extra_data,
Some(json!({"gateway_execution_runtime": true}))
);
assert_eq!(
candidate.required_capabilities,
Some(json!({"vision": true}))
);
assert!(!serde_json::to_string(candidate)
.expect("candidate should serialize")
.contains("candidate-secret"));
}
#[test]
fn gateway_candidate_boundary_sanitizes_repository_rows_and_write_results() {
let candidate = sanitize_request_candidate_row(untrusted_candidate());
assert_candidate_is_sanitized(&candidate);
let candidates = sanitize_request_candidate_rows(vec![untrusted_candidate()]);
assert_candidate_is_sanitized(&candidates[0]);
}
}
apps/aether-gateway/src/data/state/core.rs
+38
View File
@@ -895,6 +895,44 @@ impl GatewayDataState {
.await
}
pub(crate) async fn compare_and_set_system_config_string_value(
&self,
key: &str,
expected: &str,
replacement: &str,
) -> Result<bool, DataLayerError> {
if let Some(values) = &self.system_config_values {
let updated = {
let mut values = values.write().expect("system config values lock");
match values.get_mut(key) {
Some(entry) if entry.value.as_str() == Some(expected) => {
entry.value = serde_json::Value::String(replacement.to_string());
entry.updated_at_unix_secs =
Some(current_system_config_updated_at_unix_secs());
true
}
_ => false,
}
};
self.clear_cached_system_config_value(key);
return Ok(updated);
}
let result = match self.backends.as_ref() {
Some(backends) => {
crate::request_diagnostics::observe_db_operation(
"system_config_compare_and_set",
self.database_pool_summary(),
backends.compare_and_set_system_config_string_value(key, expected, replacement),
)
.await
}
None => Ok(false),
};
self.clear_cached_system_config_value(key);
result
}
pub(crate) async fn upsert_system_config_value(
&self,
key: &str,
apps/aether-gateway/src/data/state/integrations.rs
+89 -13
View File
@@ -11,7 +11,11 @@ use aether_data_contracts::repository::candidates::DecisionTrace;
use aether_data_contracts::repository::provider_catalog::{
StoredProviderCatalogEndpoint, StoredProviderCatalogKey, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::settlement::{StoredUsageSettlement, UsageSettlementInput};
use aether_data_contracts::repository::proxy_nodes::ProxyNodeTrafficMutation;
use aether_data_contracts::repository::settlement::{
ReconcileUsagePolicyCostInput, StoredUsagePolicyCostReservation, StoredUsageSettlement,
UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{
ProxyNodeCounterDelta, StoredRequestUsageAudit, UpsertUsageRecord, UsageWriteRepository,
};
@@ -34,7 +38,7 @@ const LEGACY_REQUEST_LOG_LEVEL_KEY: &str = "request_log_level";
fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestRecordLevel {
let Some(value) = value.and_then(Value::as_str).map(str::trim) else {
return UsageRequestRecordLevel::Full;
return UsageRequestRecordLevel::Basic;
};
if value.eq_ignore_ascii_case("basic")
@@ -45,7 +49,9 @@ fn usage_request_record_level_from_value(value: Option<&Value>) -> UsageRequestR
{
UsageRequestRecordLevel::Basic
} else {
UsageRequestRecordLevel::Full
// Raw HTTP payload capture is disabled at the runtime boundary. The setting remains
// accepted for compatibility, but no longer authorizes collecting request/response data.
UsageRequestRecordLevel::Basic
}
}
@@ -95,6 +101,14 @@ impl StoredVideoTaskReadSide for GatewayDataState {
) -> Result<Option<StoredVideoTask>, DataLayerError> {
GatewayDataState::find_video_task(self, key).await
}
async fn find_stored_video_task_for_user(
&self,
key: VideoTaskLookupKey<'_>,
user_id: &str,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
GatewayDataState::find_video_task_for_user(self, key, user_id).await
}
}
#[async_trait]
@@ -219,6 +233,13 @@ impl UsageSettlementWriter for GatewayDataState {
GatewayDataState::has_settlement_writer(self)
}
async fn reconcile_usage_policy_cost(
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
GatewayDataState::reconcile_usage_policy_cost(self, input).await
}
async fn settle_usage(
&self,
input: UsageSettlementInput,
@@ -285,10 +306,26 @@ impl aether_usage_runtime::ManualProxyNodeCounter for GatewayDataState {
failed_delta: i64,
latency_ms: Option<i64>,
) -> Result<(), DataLayerError> {
// This API predates incarnation fences and only receives a node id. Read
// the selected node first so every durable path can bind the delta to
// the observed generation. If the node disappeared, fail closed rather
// than allowing a bare id to target a replacement node.
let Some(node) = self.find_proxy_node(node_id).await? else {
return Ok(());
};
if !node.is_manual {
return Ok(());
}
let expected_tunnel_generation = node.tunnel_generation.trim().to_string();
if expected_tunnel_generation.is_empty() {
return Ok(());
}
if let Some(repository) = &self.usage_writer {
let enqueued = repository
.enqueue_proxy_node_counter_delta(ProxyNodeCounterDelta {
node_id: node_id.to_string(),
expected_tunnel_generation: Some(expected_tunnel_generation.clone()),
total_requests_delta: total_delta,
failed_requests_delta: failed_delta,
dns_failures_delta: 0,
@@ -300,14 +337,25 @@ impl aether_usage_runtime::ManualProxyNodeCounter for GatewayDataState {
}
}
match &self.proxy_node_writer {
Some(repository) => {
repository
.increment_manual_node_requests(node_id, total_delta, failed_delta, latency_ms)
.await
}
None => Ok(()),
// The legacy increment method has no generation argument and would
// re-read the current row, which is vulnerable to an id reuse between
// the read above and the write. Use the fenced traffic mutation as the
// only fallback. It intentionally omits the legacy latency-only field;
// preserving counters safely is more important than an unfenced write.
if let Some(repository) = &self.proxy_node_writer {
let _ = repository
.record_traffic(&ProxyNodeTrafficMutation {
node_id: node_id.to_string(),
expected_tunnel_generation: Some(expected_tunnel_generation),
total_requests_delta: total_delta,
failed_requests_delta: failed_delta,
dns_failures_delta: 0,
stream_errors_delta: 0,
})
.await?;
}
let _ = latency_ms;
Ok(())
}
}
@@ -452,6 +500,20 @@ mod tests {
assert_eq!(level, UsageRequestRecordLevel::Basic);
}
#[tokio::test]
async fn usage_runtime_access_disables_full_http_capture() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
"request_record_level".to_string(),
json!("full"),
)]);
let level = UsageRuntimeAccess::request_record_level(&state)
.await
.expect("request record level should read");
assert_eq!(level, UsageRequestRecordLevel::Basic);
}
#[tokio::test]
async fn usage_runtime_access_falls_back_to_legacy_request_log_level_alias() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
@@ -467,14 +529,14 @@ mod tests {
}
#[tokio::test]
async fn usage_runtime_access_defaults_missing_request_record_level_to_full() {
async fn usage_runtime_access_defaults_missing_request_record_level_to_basic() {
let state = GatewayDataState::disabled();
let level = UsageRuntimeAccess::request_record_level(&state)
.await
.expect("missing request record level should fall back");
assert_eq!(level, UsageRequestRecordLevel::Full);
assert_eq!(level, UsageRequestRecordLevel::Basic);
}
#[tokio::test]
@@ -488,6 +550,20 @@ mod tests {
.await
.expect("body capture policy should read");
assert_eq!(policy.record_level, UsageRequestRecordLevel::Full);
assert_eq!(policy.record_level, UsageRequestRecordLevel::Basic);
}
#[tokio::test]
async fn usage_runtime_access_fails_closed_for_unknown_record_level() {
let state = GatewayDataState::disabled().with_system_config_values_for_tests([(
"request_record_level".to_string(),
json!("everything"),
)]);
let level = UsageRuntimeAccess::request_record_level(&state)
.await
.expect("request record level should read");
assert_eq!(level, UsageRequestRecordLevel::Basic);
}
}
apps/aether-gateway/src/data/state/mod.rs
+35 -25
View File
@@ -27,8 +27,8 @@ use aether_data::repository::auth::{
StoredAuthApiKeyExportRecord, StoredAuthApiKeySnapshot,
};
use aether_data::repository::auth_modules::{
AuthModuleReadRepository, AuthModuleWriteRepository, StoredLdapModuleConfig,
StoredOAuthProviderModuleConfig,
AuthModuleReadRepository, AuthModuleWriteRepository, CompareAndSwapLdapConfigResult,
LdapBindPasswordUpdate, StoredLdapModuleConfig, StoredOAuthProviderModuleConfig,
};
use aether_data::repository::gemini_file_mappings::{
GeminiFileMappingListQuery, GeminiFileMappingReadRepository, GeminiFileMappingStats,
@@ -36,9 +36,10 @@ use aether_data::repository::gemini_file_mappings::{
UpsertGeminiFileMappingRecord,
};
use aether_data::repository::management_tokens::{
CreateManagementTokenRecord, ManagementTokenListQuery, ManagementTokenReadRepository,
ManagementTokenWriteRepository, RegenerateManagementTokenSecret, StoredManagementToken,
StoredManagementTokenListPage, StoredManagementTokenWithUser, UpdateManagementTokenRecord,
ActivateManagementTokenIfMatches, CreateManagementTokenRecord, ManagementTokenListQuery,
ManagementTokenReadRepository, ManagementTokenWriteRepository, RegenerateManagementTokenSecret,
StoredManagementToken, StoredManagementTokenListPage, StoredManagementTokenWithUser,
UpdateManagementTokenRecord,
};
use aether_data::repository::oauth_providers::{
OAuthProviderReadRepository, OAuthProviderWriteRepository, StoredOAuthProviderConfig,
@@ -63,21 +64,24 @@ pub(crate) use aether_data::repository::users::{
use aether_data::repository::wallet::{
AdjustWalletBalanceInput, AdminPaymentOrderListQuery, AdminRedeemCodeBatchListQuery,
AdminRedeemCodeListQuery, AdminWalletLedgerQuery, AdminWalletListQuery,
AdminWalletRefundRequestListQuery, CompleteAdminWalletRefundInput,
CreateAdminRedeemCodeBatchInput, CreateAdminRedeemCodeBatchResult,
CreateManualWalletRechargeInput, CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome,
CreateWalletRechargeOrderInput, CreateWalletRechargeOrderOutcome,
CreateWalletRefundRequestInput, CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
AdminWalletRefundRequestListQuery, CompareAndSwapPaymentOrderStripeClientSecretInput,
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
CreateAdminRedeemCodeBatchResult, CreateManualWalletRechargeInput,
CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput,
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput,
DeleteAdminRedeemCodeBatchInput, DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput,
FailAdminWalletRefundInput, ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput,
ProcessPaymentCallbackOutcome, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
FailAdminWalletRefundInput, FailWalletRechargeCheckoutInput, InitializeAuthWalletOutcome,
ProcessAdminWalletRefundInput, ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome,
ReclaimWalletRechargeCheckoutInput, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
StoredAdminPaymentCallback, StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder,
StoredAdminPaymentOrderPage, StoredAdminRedeemCode, StoredAdminRedeemCodeBatch,
StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage,
StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage,
StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
StoredAdminWalletTransactionPage, StoredWalletDailyUsageLedger,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, WalletLookupKey, WalletMutationOutcome,
StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot, UpdateAdminWalletRefundGatewayInput,
UpdateWalletRechargeCheckoutInput, WalletLookupKey, WalletMutationOutcome,
WalletReadRepository, WalletWriteRepository,
};
use aether_data::{
@@ -92,8 +96,9 @@ use aether_data_contracts::repository::background_tasks::{
use aether_data_contracts::repository::billing::{
AdminBillingCollectorRecord, AdminBillingCollectorWriteInput, AdminBillingMutationOutcome,
AdminBillingPresetApplyResult, AdminBillingRuleRecord, AdminBillingRuleWriteInput,
BillingPlanRecord, BillingPlanWriteInput, BillingReadRepository, PaymentGatewayConfigRecord,
PaymentGatewayConfigWriteInput, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
BillingPlanRecord, BillingPlanWriteInput, BillingReadRepository,
PaymentGatewayConfigCasWriteInput, PaymentGatewayConfigRecord, PaymentGatewayConfigWriteInput,
PaymentGatewaySecretCasUpdate, StoredBillingModelContext, UserDailyQuotaAvailabilityRecord,
UserPlanEntitlementRecord,
};
use aether_data_contracts::repository::candidate_selection::{
@@ -122,12 +127,14 @@ use aether_data_contracts::repository::pool_scores::{
};
use aether_data_contracts::repository::provider_catalog::{
ProviderCatalogKeyAdaptiveStateUpdate, ProviderCatalogKeyAdminCasUpdate,
ProviderCatalogKeyHealthStateUpdate, ProviderCatalogKeyListQuery,
ProviderCatalogKeyOAuthCredentialCasDelete, ProviderCatalogKeyOAuthRuntimeStateCasUpdate,
ProviderCatalogKeyRuntimeMetadataUpdate, ProviderCatalogKeyStatusSnapshotUpdate,
ProviderCatalogReadRepository, ProviderCatalogWriteRepository, StoredProviderCatalogEndpoint,
StoredProviderCatalogKey, StoredProviderCatalogKeyMaintenanceSummary,
StoredProviderCatalogKeyPage, StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
ProviderCatalogKeyCredentialsCasUpdate, ProviderCatalogKeyHealthStateUpdate,
ProviderCatalogKeyListQuery, ProviderCatalogKeyOAuthCredentialCasDelete,
ProviderCatalogKeyOAuthRuntimeStateCasUpdate, ProviderCatalogKeyRuntimeMetadataUpdate,
ProviderCatalogKeyStatusSnapshotUpdate, ProviderCatalogProviderConfigCasUpdate,
ProviderCatalogProxyCasUpdate, ProviderCatalogReadRepository, ProviderCatalogWriteRepository,
StoredProviderCatalogEndpoint, StoredProviderCatalogKey,
StoredProviderCatalogKeyMaintenanceSummary, StoredProviderCatalogKeyPage,
StoredProviderCatalogKeyStats, StoredProviderCatalogProvider,
};
use aether_data_contracts::repository::quota::{
ProviderQuotaReadRepository, ProviderQuotaWriteRepository, StoredProviderQuotaSnapshot,
@@ -136,7 +143,10 @@ use aether_data_contracts::repository::routing_profiles::{
RoutingGroupReadRepository, RoutingGroupWriteRepository,
};
use aether_data_contracts::repository::settlement::{
SettlementWriteRepository, StoredUsageSettlement, UsageSettlementInput,
ReconcileUsagePolicyCostInput, ReleaseUsagePolicyRequestAdmissionInput,
ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput,
ReserveUsagePolicyRequestOutcome, SettlementWriteRepository, StoredUsagePolicyCostReservation,
StoredUsagePolicyRequestAdmission, StoredUsageSettlement, UsageSettlementInput,
};
use aether_data_contracts::repository::usage::{
ApiKeyLastUsedDelta, ManagementTokenCounterDelta, PendingUsageCleanupSummary,
@@ -150,9 +160,9 @@ use aether_data_contracts::repository::video_tasks::{
use aether_runtime_state::RuntimeQueueStore;
pub(crate) use self::referrals::{
ReferralAdminStats, ReferralMutationStatus, ReferralRelationshipListQuery,
ReferralRelationshipRecord, ReferralRewardConfig, ReferralRewardListQuery,
ReferralRewardRecord, ReferralUserDashboard,
ReferralAdminStats, ReferralMutationStatus, ReferralReconciliationSummary,
ReferralRelationshipListQuery, ReferralRelationshipRecord, ReferralRewardConfig,
ReferralRewardListQuery, ReferralRewardRecord, ReferralUserDashboard,
};
#[derive(Clone, Default)]
apps/aether-gateway/src/data/state/referrals.rs
+12 -3
View File
@@ -4,9 +4,9 @@ use aether_data::DataLayerError;
use super::GatewayDataState;
pub(crate) use aether_data::backend::{
ReferralAdminStats, ReferralMutationStatus, ReferralRelationshipListQuery,
ReferralRelationshipRecord, ReferralRewardConfig, ReferralRewardListQuery,
ReferralRewardRecord, ReferralUserDashboard,
ReferralAdminStats, ReferralMutationStatus, ReferralReconciliationSummary,
ReferralRelationshipListQuery, ReferralRelationshipRecord, ReferralRewardConfig,
ReferralRewardListQuery, ReferralRewardRecord, ReferralUserDashboard,
};
impl GatewayDataState {
@@ -115,4 +115,13 @@ impl GatewayDataState {
.reverse_referral_rewards_for_order(order_id, amount_usd)
.await
}
pub(crate) async fn reconcile_referral_rewards_once(
&self,
reward_config: Option<ReferralRewardConfig>,
) -> Result<ReferralReconciliationSummary, DataLayerError> {
self.referrals()
.reconcile_referral_rewards_once(reward_config)
.await
}
}
apps/aether-gateway/src/data/state/runtime.rs
+257 -19
View File
@@ -7,33 +7,39 @@ use super::{
AdminWalletRefundRequestListQuery, AnnouncementListQuery, AuditLogListQuery,
BackgroundTaskListQuery, BackgroundTaskSummary, BillingModelContextCacheKey,
BillingModelContextCacheState, BillingModelContextInflightState, BillingPlanRecord,
BillingPlanWriteInput, CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
BillingPlanWriteInput, CompareAndSwapPaymentOrderStripeClientSecretInput,
CompleteAdminWalletRefundInput, CreateAdminRedeemCodeBatchInput,
CreateAdminRedeemCodeBatchResult, CreateAnnouncementRecord, CreateManualWalletRechargeInput,
CreatePlanPurchaseOrderInput, CreatePlanPurchaseOrderOutcome, CreateWalletRechargeOrderInput,
CreateWalletRechargeOrderOutcome, CreateWalletRefundRequestInput,
CreateWalletRefundRequestOutcome, CreditAdminPaymentOrderInput, DataLayerError,
DatabaseMaintenanceSummary, DecisionTrace, DeleteAdminRedeemCodeBatchInput,
DisableAdminRedeemCodeBatchInput, DisableAdminRedeemCodeInput, FailAdminWalletRefundInput,
GatewayDataState, GatewayProviderTransportSnapshot, LocalVideoTaskReadResponse,
PaymentGatewayConfigRecord, PaymentGatewayConfigWriteInput, ProcessAdminWalletRefundInput,
ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, RedeemWalletCodeInput,
RedeemWalletCodeOutcome, RequestAuditBundle, RequestCandidateTrace, StoredAdminAuditLogPage,
StoredAdminPaymentCallbackPage, StoredAdminPaymentOrder, StoredAdminPaymentOrderPage,
StoredAdminRedeemCodeBatch, StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage,
StoredAdminWalletLedgerPage, StoredAdminWalletListPage, StoredAdminWalletRefund,
StoredAdminWalletRefundPage, StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
FailWalletRechargeCheckoutInput, GatewayDataState, GatewayProviderTransportSnapshot,
LocalVideoTaskReadResponse, PaymentGatewayConfigCasWriteInput, PaymentGatewayConfigRecord,
PaymentGatewayConfigWriteInput, PaymentGatewaySecretCasUpdate, ProcessAdminWalletRefundInput,
ProcessPaymentCallbackInput, ProcessPaymentCallbackOutcome, ReclaimWalletRechargeCheckoutInput,
ReconcileUsagePolicyCostInput, RedeemWalletCodeInput, RedeemWalletCodeOutcome,
ReleaseUsagePolicyRequestAdmissionInput, RequestAuditBundle, RequestCandidateTrace,
ReserveUsagePolicyCostInput, ReserveUsagePolicyCostOutcome, ReserveUsagePolicyRequestInput,
ReserveUsagePolicyRequestOutcome, StoredAdminAuditLogPage, StoredAdminPaymentCallbackPage,
StoredAdminPaymentOrder, StoredAdminPaymentOrderPage, StoredAdminRedeemCodeBatch,
StoredAdminRedeemCodeBatchPage, StoredAdminRedeemCodePage, StoredAdminWalletLedgerPage,
StoredAdminWalletListPage, StoredAdminWalletRefund, StoredAdminWalletRefundPage,
StoredAdminWalletRefundRequestPage, StoredAdminWalletTransaction,
StoredAdminWalletTransactionPage, StoredAnnouncement, StoredAnnouncementPage,
StoredBackgroundTaskEvent, StoredBackgroundTaskRun, StoredBackgroundTaskRunPage,
StoredBillingModelContext, StoredProviderQuotaSnapshot, StoredProviderUsageSummary,
StoredRequestUsageAudit, StoredSuspiciousActivity, StoredUsageSettlement,
StoredUserAuditLogPage, StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary,
StoredVideoTask, StoredWalletDailyUsageLedger, StoredWalletDailyUsageLedgerPage,
StoredWalletSnapshot, UpdateAnnouncementRecord, UpsertBackgroundTaskEvent,
UpsertBackgroundTaskRun, UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput,
UserDailyQuotaAvailabilityRecord, UserPlanEntitlementRecord, VideoTaskLookupKey,
VideoTaskModelCount, VideoTaskQueryFilter, VideoTaskStatusCount,
WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult, WalletLookupKey,
WalletMutationOutcome,
StoredRequestUsageAudit, StoredSuspiciousActivity, StoredUsagePolicyCostReservation,
StoredUsagePolicyRequestAdmission, StoredUsageSettlement, StoredUserAuditLogPage,
StoredUserAuthRecord, StoredUserExportRow, StoredUserSummary, StoredVideoTask,
StoredWalletDailyUsageLedger, StoredWalletDailyUsageLedgerPage, StoredWalletSnapshot,
UpdateAdminWalletRefundGatewayInput, UpdateAnnouncementRecord,
UpdateWalletRechargeCheckoutInput, UpsertBackgroundTaskEvent, UpsertBackgroundTaskRun,
UpsertUsageRecord, UpsertVideoTask, UsageSettlementInput, UserDailyQuotaAvailabilityRecord,
UserPlanEntitlementRecord, VideoTaskLookupKey, VideoTaskModelCount, VideoTaskQueryFilter,
VideoTaskStatusCount, WalletDailyUsageAggregationInput, WalletDailyUsageAggregationResult,
WalletLookupKey, WalletMutationOutcome,
};
use aether_data_contracts::repository::usage::{
PendingUsageCleanupSummary, ProviderApiKeyWindowUsageRequest,
@@ -43,7 +49,9 @@ use aether_data_contracts::repository::usage::{
UsageDailyHeatmapQuery,
};
use aether_runtime_state::RuntimeQueueStore;
use aether_video_tasks_core::read_data_backed_video_task_response;
use aether_video_tasks_core::{
read_data_backed_video_task_response, read_data_backed_video_task_response_for_user,
};
use std::time::{Duration, Instant};
use tokio::time::timeout;
@@ -558,6 +566,17 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_video_task_for_user(
&self,
key: VideoTaskLookupKey<'_>,
user_id: &str,
) -> Result<Option<StoredVideoTask>, DataLayerError> {
match &self.video_task_reader {
Some(repository) => repository.find_for_user(key, user_id).await,
None => Ok(None),
}
}
pub(crate) async fn list_video_task_page(
&self,
filter: &VideoTaskQueryFilter,
@@ -894,6 +913,21 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_wallet_recharge_order_by_order_no(
&self,
user_id: &str,
order_no: &str,
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
match &self.wallet_reader {
Some(repository) => {
repository
.find_wallet_recharge_order_by_order_no(user_id, order_no)
.await
}
None => Ok(None),
}
}
pub(crate) async fn find_pending_plan_purchase_order_by_user_id(
&self,
user_id: &str,
@@ -909,6 +943,16 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_payment_order_by_order_no(
&self,
order_no: &str,
) -> Result<Option<StoredAdminPaymentOrder>, DataLayerError> {
match &self.wallet_reader {
Some(repository) => repository.find_payment_order_by_order_no(order_no).await,
None => Ok(None),
}
}
pub(crate) async fn find_wallet_refund(
&self,
wallet_id: &str,
@@ -934,6 +978,58 @@ impl GatewayDataState {
}
}
pub(crate) async fn update_wallet_recharge_checkout(
&self,
input: UpdateWalletRechargeCheckoutInput,
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
match &self.wallet_writer {
Some(repository) => repository
.update_wallet_recharge_checkout(input)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn compare_and_swap_payment_order_stripe_client_secret(
&self,
input: CompareAndSwapPaymentOrderStripeClientSecretInput,
) -> Result<Option<bool>, DataLayerError> {
match &self.wallet_writer {
Some(repository) => repository
.compare_and_swap_payment_order_stripe_client_secret(input)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn fail_wallet_recharge_checkout(
&self,
input: FailWalletRechargeCheckoutInput,
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
match &self.wallet_writer {
Some(repository) => repository
.fail_wallet_recharge_checkout(input)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn reclaim_wallet_recharge_checkout(
&self,
input: ReclaimWalletRechargeCheckoutInput,
) -> Result<Option<WalletMutationOutcome<StoredAdminPaymentOrder>>, DataLayerError> {
match &self.wallet_writer {
Some(repository) => repository
.reclaim_wallet_recharge_checkout(input)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn create_plan_purchase_order(
&self,
input: CreatePlanPurchaseOrderInput,
@@ -1009,6 +1105,19 @@ impl GatewayDataState {
}
}
pub(crate) async fn update_admin_wallet_refund_gateway(
&self,
input: UpdateAdminWalletRefundGatewayInput,
) -> Result<Option<WalletMutationOutcome<StoredAdminWalletRefund>>, DataLayerError> {
match &self.wallet_writer {
Some(repository) => repository
.update_admin_wallet_refund_gateway(input)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn complete_admin_wallet_refund(
&self,
input: CompleteAdminWalletRefundInput,
@@ -1151,6 +1260,83 @@ impl GatewayDataState {
}
}
pub(crate) async fn reserve_usage_policy_cost(
&self,
input: ReserveUsagePolicyCostInput,
) -> Result<Option<ReserveUsagePolicyCostOutcome>, DataLayerError> {
match &self.settlement_writer {
Some(repository) => repository.reserve_usage_policy_cost(input).await.map(Some),
None => Ok(None),
}
}
pub(crate) async fn reserve_usage_policy_request(
&self,
input: ReserveUsagePolicyRequestInput,
) -> Result<Option<ReserveUsagePolicyRequestOutcome>, DataLayerError> {
match &self.settlement_writer {
Some(repository) => repository
.reserve_usage_policy_request(input)
.await
.map(Some),
None => Ok(None),
}
}
pub(crate) async fn release_usage_policy_request_admission(
&self,
input: ReleaseUsagePolicyRequestAdmissionInput,
) -> Result<Option<StoredUsagePolicyRequestAdmission>, DataLayerError> {
match &self.settlement_writer {
Some(repository) => {
repository
.release_usage_policy_request_admission(input)
.await
}
None => Ok(None),
}
}
pub(crate) async fn cleanup_usage_policy_request_admissions(
&self,
now_unix_secs: u64,
batch_size: usize,
) -> Result<usize, DataLayerError> {
match &self.settlement_writer {
Some(repository) => {
repository
.cleanup_usage_policy_request_admissions(now_unix_secs, batch_size)
.await
}
None => Ok(0),
}
}
pub(crate) async fn reconcile_usage_policy_cost(
&self,
input: ReconcileUsagePolicyCostInput,
) -> Result<Option<StoredUsagePolicyCostReservation>, DataLayerError> {
match &self.settlement_writer {
Some(repository) => repository.reconcile_usage_policy_cost(input).await,
None => Ok(None),
}
}
pub(crate) async fn cleanup_usage_policy_cost_reservations(
&self,
now_unix_secs: u64,
batch_size: usize,
) -> Result<usize, DataLayerError> {
match &self.settlement_writer {
Some(repository) => {
repository
.cleanup_usage_policy_cost_reservations(now_unix_secs, batch_size)
.await
}
None => Ok(0),
}
}
pub(crate) async fn reset_due_provider_quotas(
&self,
now_unix_secs: u64,
@@ -2496,6 +2682,48 @@ impl GatewayDataState {
}
}
pub(crate) async fn find_payment_gateway_config_strong(
&self,
provider: &str,
) -> Result<Option<PaymentGatewayConfigRecord>, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.find_payment_gateway_config_strong(provider)
.await
}
None => Ok(None),
}
}
pub(crate) async fn compare_and_swap_payment_gateway_secret(
&self,
update: &PaymentGatewaySecretCasUpdate,
) -> Result<bool, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.compare_and_swap_payment_gateway_secret(update)
.await
}
None => Ok(false),
}
}
pub(crate) async fn compare_and_swap_payment_gateway_config(
&self,
input: &PaymentGatewayConfigCasWriteInput,
) -> Result<AdminBillingMutationOutcome<PaymentGatewayConfigRecord>, DataLayerError> {
match &self.billing_reader {
Some(repository) => {
repository
.compare_and_swap_payment_gateway_config(input)
.await
}
None => Ok(AdminBillingMutationOutcome::Unavailable),
}
}
pub(crate) async fn upsert_payment_gateway_config(
&self,
input: &PaymentGatewayConfigWriteInput,
@@ -2667,6 +2895,16 @@ impl GatewayDataState {
read_data_backed_video_task_response(self, route_family, request_path).await
}
pub(crate) async fn read_video_task_response_for_user(
&self,
route_family: Option<&str>,
request_path: &str,
user_id: &str,
) -> Result<Option<LocalVideoTaskReadResponse>, DataLayerError> {
read_data_backed_video_task_response_for_user(self, route_family, request_path, user_id)
.await
}
pub(crate) async fn find_background_task_run(
&self,
run_id: &str,
apps/aether-gateway/src/data/state/testing/mod.rs
+42
View File
@@ -1790,6 +1790,18 @@ impl GatewayDataState {
}
}
#[cfg(test)]
pub(crate) fn attach_auth_api_key_repository_for_tests<T>(mut self, repository: Arc<T>) -> Self
where
T: aether_data::repository::auth::AuthRepository + 'static,
{
let auth_api_key_reader: Arc<dyn AuthApiKeyReadRepository> = repository.clone();
let auth_api_key_writer: Arc<dyn AuthApiKeyWriteRepository> = repository;
self.auth_api_key_reader = Some(auth_api_key_reader);
self.auth_api_key_writer = Some(auth_api_key_writer);
self
}
#[cfg(test)]
pub(crate) fn with_decision_trace_readers_for_tests(
request_candidate_repository: Arc<dyn RequestCandidateReadRepository>,
@@ -2379,6 +2391,36 @@ impl GatewayDataState {
}
}
#[cfg(test)]
pub(crate) fn with_auth_candidate_selection_provider_catalog_request_candidate_and_gemini_file_mapping_repositories_for_tests<
T,
U,
V,
>(
auth_api_key_repository: Arc<dyn AuthApiKeyReadRepository>,
candidate_selection_repository: Arc<dyn MinimalCandidateSelectionReadRepository>,
provider_catalog_repository: Arc<U>,
request_candidate_repository: Arc<T>,
gemini_file_mapping_repository: Arc<V>,
encryption_key: impl Into<String>,
) -> Self
where
T: RequestCandidateRepository + 'static,
U: ProviderCatalogReadRepository + ProviderCatalogWriteRepository + 'static,
V: aether_data::repository::gemini_file_mappings::GeminiFileMappingRepository + 'static,
{
let mut state = Self::with_auth_candidate_selection_provider_catalog_and_request_candidate_repository_for_tests(
auth_api_key_repository,
candidate_selection_repository,
provider_catalog_repository,
request_candidate_repository,
encryption_key,
);
state.gemini_file_mapping_reader = Some(gemini_file_mapping_repository.clone());
state.gemini_file_mapping_writer = Some(gemini_file_mapping_repository);
state
}
#[cfg(test)]
pub(crate) fn with_auth_candidate_selection_provider_catalog_request_candidates_for_tests<
T,
apps/aether-gateway/src/data/state/testkit.rs
+103
View File
@@ -12,10 +12,113 @@ use aether_data_contracts::repository::usage::{
};
use aether_data::repository::auth::AuthApiKeyReadRepository;
use aether_data::repository::management_tokens::{
InMemoryManagementTokenRepository, ManagementTokenReadRepository,
ManagementTokenWriteRepository, StoredManagementToken, StoredManagementTokenUserSummary,
StoredManagementTokenWithUser,
};
use aether_data::repository::proxy_nodes::{InMemoryProxyNodeRepository, StoredProxyNode};
use aether_data::repository::proxy_nodes::{ProxyNodeReadRepository, ProxyNodeWriteRepository};
use aether_data::repository::users::{
InMemoryUserReadRepository, StoredUserAuthRecord, UserReadRepository,
};
use sha2::{Digest, Sha256};
use super::{GatewayDataConfig, GatewayDataState};
impl GatewayDataState {
pub(crate) fn with_tunnel_management_auth_for_testkit(
node_id: &str,
tunnel_generation: &str,
raw_token: &str,
encryption_key: impl Into<String>,
) -> Result<Self, aether_data::DataLayerError> {
const TOKEN_ID: &str = "token-tunnel-harness";
const USER_ID: &str = "user-tunnel-harness";
let node = StoredProxyNode::new(
node_id.to_string(),
"tunnel harness node".to_string(),
"127.0.0.1".to_string(),
0,
false,
"offline".to_string(),
30,
0,
0,
0,
0,
0,
true,
false,
0,
)?
.with_tunnel_generation(tunnel_generation.to_string());
let proxy_repository = Arc::new(InMemoryProxyNodeRepository::seed([node]));
let user_summary = StoredManagementTokenUserSummary::new(
USER_ID.to_string(),
Some("[email protected]".to_string()),
"tunnel_harness_admin".to_string(),
"admin".to_string(),
)?;
let token = StoredManagementToken::new(
TOKEN_ID.to_string(),
USER_ID.to_string(),
"tunnel harness token".to_string(),
)?
.with_permissions(Some(serde_json::json!(["admin:proxy_nodes:admin"])));
let token_hash = format!("{:x}", Sha256::digest(raw_token.as_bytes()));
let token_repository = Arc::new(InMemoryManagementTokenRepository::seed_with_hashes(
[StoredManagementTokenWithUser::new(token, user_summary)],
[(token_hash, TOKEN_ID.to_string())],
));
let token_reader: Arc<dyn ManagementTokenReadRepository> = token_repository.clone();
let token_writer: Arc<dyn ManagementTokenWriteRepository> = token_repository;
let user = StoredUserAuthRecord::new(
USER_ID.to_string(),
Some("[email protected]".to_string()),
true,
"tunnel_harness_admin".to_string(),
None,
"admin".to_string(),
"local".to_string(),
None,
None,
None,
true,
false,
None,
None,
)?;
let user_reader: Arc<dyn UserReadRepository> =
Arc::new(InMemoryUserReadRepository::seed_auth_users([user]));
let mut state =
Self::with_proxy_node_repository_for_testkit(proxy_repository, encryption_key);
state.management_token_reader = Some(token_reader);
state.management_token_writer = Some(token_writer);
state.user_reader = Some(user_reader);
Ok(state)
}
pub(crate) fn with_proxy_node_repository_for_testkit<T>(
repository: Arc<T>,
encryption_key: impl Into<String>,
) -> Self
where
T: ProxyNodeReadRepository + ProxyNodeWriteRepository + 'static,
{
let proxy_node_reader: Arc<dyn ProxyNodeReadRepository> = repository.clone();
let proxy_node_writer: Arc<dyn ProxyNodeWriteRepository> = repository;
let mut state = Self::disabled();
state.config = GatewayDataConfig::disabled().with_encryption_key(encryption_key);
state.proxy_node_reader = Some(proxy_node_reader);
state.proxy_node_writer = Some(proxy_node_writer);
state
}
pub(crate) fn with_openai_chat_pressure_repositories_for_testkit<T, U, V>(
auth_api_key_repository: Arc<dyn AuthApiKeyReadRepository>,
candidate_selection_repository: Arc<dyn MinimalCandidateSelectionReadRepository>,
apps/aether-gateway/src/data/tests.rs
+1 -1
View File
@@ -312,7 +312,7 @@ async fn data_state_checks_user_uniqueness_through_user_reader() {
Some("[email protected]".to_string()),
true,
"admin".to_string(),
Some(format!("$2b$12${}", "a".repeat(53))),
Some("$2b$12$4qL4tdcsFwVaDTw5Ck3xzu8GpNdre56DiNR6Dnw7t6gCXaEnqAe7G".to_string()),
"admin".to_string(),
"local".to_string(),
None,