Skip to main content

awsim_core/
lib.rs

1pub mod arn;
2pub mod auth;
3pub mod authz;
4pub mod bearer_token;
5pub mod body;
6pub mod body_store;
7pub mod error;
8pub mod events;
9pub mod gateway;
10pub mod idempotency;
11pub mod lifecycle;
12pub mod pagination;
13pub mod persistence;
14pub mod protocol;
15pub mod request_detail;
16pub mod request_event;
17pub mod router;
18pub mod sigv4_verify;
19pub mod state;
20pub mod tags;
21pub mod tick;
22pub mod totp;
23
24pub use authz::{
25    AuthzEngine, CloudMapRegistrar, GrantLookup, KmsKeyLookup, LambdaInvoker, NoopPrincipalLookup,
26    ParameterLookup, PrincipalLookup, ResolvedPrincipal, ResourcePolicyLookup, ScpLookup,
27    SecretLookup,
28};
29// `HandlerByteStream` and `HandlerResult` are defined further down in
30// this file; re-exported here for crate consumers.
31pub use arn::Arn;
32pub use body::Body;
33pub use body_store::{BlobInventory, BodyStore};
34pub use error::AwsError;
35pub use events::{EventBus, InternalEvent};
36pub use gateway::{AppState, BodyStoreHandle};
37pub use pagination::{
38    Page, cap_max_results, clamp_max_results_strict, decode_token, encode_token, paginate,
39};
40pub use persistence::PersistenceManager;
41pub use protocol::{Protocol, RouteDefinition};
42pub use request_detail::{
43    CapturedBody, CapturedHeader, DEFAULT_BODY_CAP, DEFAULT_RING_CAPACITY, RequestDetail,
44    RequestDetailStore, capture_body, capture_headers,
45};
46pub use request_event::{RequestEvent, RequestEventBus};
47pub use router::{DEFAULT_PARTITION, RequestContext};
48pub use state::{AccountRegionStore, Snapshottable};
49pub use tick::{TestDriver, WorkerPool};
50
51use bytes::Bytes;
52use futures::stream::BoxStream;
53use serde_json::Value;
54
55/// Boxed byte stream a handler may return when it wants to drive an
56/// HTTP response chunk-by-chunk (e.g. Bedrock's event-stream APIs)
57/// instead of buffering the whole response into a single `Value`.
58pub type HandlerByteStream = BoxStream<'static, Result<Bytes, AwsError>>;
59
60/// What `ServiceHandler::handle_streaming` returns. Most operations
61/// produce a single JSON `Value` (the existing path); a small set —
62/// notably Bedrock's `ConverseStream` and
63/// `InvokeModelWithResponseStream` — produce a continuous stream of
64/// already-encoded body bytes plus a content-type the gateway puts
65/// straight on the wire.
66pub enum HandlerResult {
67    /// Conventional single-shot response. The gateway runs it
68    /// through the normal protocol serializer.
69    Json(Value),
70    /// Streamed binary body. The gateway sends it via axum's
71    /// chunked-transfer body so the client sees bytes as they're
72    /// produced — no buffering on our side.
73    Streaming {
74        body: HandlerByteStream,
75        content_type: &'static str,
76    },
77}
78
79impl From<Value> for HandlerResult {
80    fn from(v: Value) -> Self {
81        HandlerResult::Json(v)
82    }
83}
84
85/// Trait that every AWS service crate must implement.
86///
87/// Each service (S3, SQS, DynamoDB, etc.) implements this trait in its own crate.
88/// The main `awsim` binary registers all service handlers with the gateway router.
89#[async_trait::async_trait]
90pub trait ServiceHandler: Send + Sync {
91    /// The AWS service name (e.g., "s3", "sqs", "dynamodb").
92    fn service_name(&self) -> &str;
93
94    /// The signing name used in SigV4 Authorization headers.
95    /// Usually the same as service_name, but not always.
96    fn signing_name(&self) -> &str {
97        self.service_name()
98    }
99
100    /// The primary protocol this service uses.
101    fn protocol(&self) -> Protocol;
102
103    /// Route definitions for REST-protocol services.
104    /// Not needed for RPC-style protocols (awsJson, awsQuery).
105    fn routes(&self) -> Vec<RouteDefinition> {
106        Vec::new()
107    }
108
109    /// Handle an AWS API operation.
110    async fn handle(
111        &self,
112        operation: &str,
113        input: Value,
114        ctx: &RequestContext,
115    ) -> Result<Value, AwsError>;
116
117    /// Streaming-aware variant. The default delegates to `handle` and
118    /// wraps the JSON result so existing services don't need to do
119    /// anything; services that genuinely stream (Bedrock data plane)
120    /// override this and return `HandlerResult::Streaming`.
121    async fn handle_streaming(
122        &self,
123        operation: &str,
124        input: Value,
125        ctx: &RequestContext,
126    ) -> Result<HandlerResult, AwsError> {
127        self.handle(operation, input, ctx)
128            .await
129            .map(HandlerResult::from)
130    }
131
132    /// Serialize the service's state to bytes for persistence.
133    ///
134    /// Return `None` if this service does not support snapshots.
135    fn snapshot(&self) -> Option<Vec<u8>> {
136        None
137    }
138
139    /// Restore the service's state from a previous snapshot.
140    ///
141    /// The default implementation is a no-op and always succeeds.
142    fn restore(&self, _data: &[u8]) -> Result<(), String> {
143        Ok(())
144    }
145
146    /// Called after every service has been restored from its
147    /// snapshot, before the gateway begins serving traffic.
148    ///
149    /// Use this to re-arm timers, restart event-source-mapping
150    /// pollers, and re-register tick-driven workers that don't
151    /// persist as data (because they are derivable from the
152    /// already-restored state). The default implementation is a
153    /// no-op so services that don't have background work pay
154    /// nothing.
155    ///
156    /// Two ordering guarantees the gateway makes:
157    /// - Every service's [`Self::restore`] completes before any
158    ///   service's `rehydrate` is invoked, so cross-service
159    ///   wiring (Lambda event source mappings reading SQS state,
160    ///   say) sees the fully restored peer.
161    /// - No request is dispatched until every service's
162    ///   `rehydrate` returns.
163    fn rehydrate(&self) -> Result<(), String> {
164        Ok(())
165    }
166
167    fn iam_action(&self, _operation: &str) -> Option<String> {
168        None
169    }
170
171    fn iam_resource(
172        &self,
173        _operation: &str,
174        _input: &serde_json::Value,
175        _ctx: &router::RequestContext,
176    ) -> Option<String> {
177        None
178    }
179
180    /// Periodic tick. The gateway spawns a single 1-second loop that
181    /// calls `tick` on every registered service after the server is up.
182    /// Use this hook for time-driven behavior that doesn't fit into the
183    /// request path: SQS visibility-timeout reclamation, DynamoDB TTL
184    /// expiry, Lambda event-source-mapping polling, S3 lifecycle
185    /// transitions, EventBridge schedule firing, SecretsManager
186    /// rotation, etc.
187    ///
188    /// **Contract:**
189    /// - `tick` must be idempotent — it may be called repeatedly, and
190    ///   missing a tick must not lose state. Use absolute deadlines
191    ///   (`Instant`/`SystemTime`) rather than per-call deltas.
192    /// - `tick` must return quickly (target <10 ms). Slow work
193    ///   (HTTP fan-out, subprocess invocation, large iterations)
194    ///   should be enqueued onto an internal worker the service spawns
195    ///   from elsewhere — `tick` enqueues, doesn't block.
196    /// - The default implementation is a no-op so existing services
197    ///   don't need to opt in.
198    async fn tick(&self) {}
199}