winterbaume_ec2instanceconnect/
handlers.rs

1use std::future::Future;
2use std::pin::Pin;
3use std::sync::Arc;
4
5use serde_json::Value;
6use winterbaume_core::{
7    BackendState, MockRequest, MockResponse, MockService, StateChangeNotifier, json_error_response,
8};
9
10use crate::state::{Ec2InstanceConnectError, Ec2InstanceConnectState};
11use crate::views::Ec2InstanceConnectStateView;
12use crate::wire;
13
14pub struct Ec2InstanceConnectService {
15    pub(crate) state: Arc<BackendState<Ec2InstanceConnectState>>,
16    pub(crate) notifier: StateChangeNotifier<Ec2InstanceConnectStateView>,
17}
18
19impl Ec2InstanceConnectService {
20    pub fn new() -> Self {
21        Self {
22            state: Arc::new(BackendState::new()),
23            notifier: StateChangeNotifier::new(),
24        }
25    }
26}
27
28impl Default for Ec2InstanceConnectService {
29    fn default() -> Self {
30        Self::new()
31    }
32}
33
34impl MockService for Ec2InstanceConnectService {
35    fn service_name(&self) -> &str {
36        "ec2-instance-connect"
37    }
38
39    fn url_patterns(&self) -> Vec<&str> {
40        vec![
41            r"https?://ec2-instance-connect\..*\.amazonaws\.com",
42            r"https?://ec2-instance-connect\.amazonaws\.com",
43        ]
44    }
45
46    fn handle(
47        &self,
48        request: MockRequest,
49    ) -> Pin<Box<dyn Future<Output = MockResponse> + Send + '_>> {
50        Box::pin(async move { self.dispatch(request).await })
51    }
52}
53
54async fn shape_error(err: Ec2InstanceConnectError) -> MockResponse {
55    match err {
56        Ec2InstanceConnectError::MissingAction => {
57            json_error_response(400, "MissingAction", &err.to_string())
58        }
59        Ec2InstanceConnectError::SerializationException => {
60            json_error_response(400, "SerializationException", &err.to_string())
61        }
62        Ec2InstanceConnectError::InvalidAction { .. } => {
63            json_error_response(400, "InvalidAction", &err.to_string())
64        }
65        Ec2InstanceConnectError::InstanceIdRequired => {
66            json_error_response(400, "InvalidArgsException", &err.to_string())
67        }
68        Ec2InstanceConnectError::InstanceOSUserRequired => {
69            json_error_response(400, "InvalidArgsException", &err.to_string())
70        }
71        Ec2InstanceConnectError::SSHPublicKeyRequired => {
72            json_error_response(400, "InvalidArgsException", &err.to_string())
73        }
74    }
75}
76
77impl Ec2InstanceConnectService {
78    async fn dispatch(&self, request: MockRequest) -> MockResponse {
79        let region = winterbaume_core::auth::extract_region_from_uri(&request.uri);
80        let account_id = winterbaume_core::default_account_id();
81
82        // Extract action from X-Amz-Target header
83        // Format: "AWSEC2InstanceConnectService.SendSSHPublicKey"
84        let action = request
85            .headers
86            .get("x-amz-target")
87            .and_then(|v| v.to_str().ok())
88            .and_then(|v| v.split('.').next_back())
89            .map(|s| s.to_string());
90
91        let action = match action {
92            Some(a) => a,
93            None => {
94                return shape_error(Ec2InstanceConnectError::MissingAction).await;
95            }
96        };
97
98        let body: Value = match serde_json::from_slice(&request.body) {
99            Ok(v) => v,
100            Err(_) => {
101                return shape_error(Ec2InstanceConnectError::SerializationException).await;
102            }
103        };
104
105        // Ensure state is initialized for the region (even though near-stateless)
106        let _state = self.state.get(account_id, &region);
107
108        use winterbaume_core::StatefulService;
109        let response = match action.as_str() {
110            "SendSSHPublicKey" => self.handle_send_ssh_public_key(&body).await,
111            "SendSerialConsoleSSHPublicKey" => {
112                self.handle_send_serial_console_ssh_public_key(&body).await
113            }
114            _ => {
115                shape_error(Ec2InstanceConnectError::InvalidAction {
116                    action: action.clone(),
117                })
118                .await
119            }
120        };
121        if response.status / 100 == 2 {
122            self.notify_state_changed(account_id, &region).await;
123        }
124        response
125    }
126
127    async fn handle_send_ssh_public_key(&self, body: &Value) -> MockResponse {
128        // Validate required fields
129        if body.get("InstanceId").and_then(|v| v.as_str()).is_none() {
130            return shape_error(Ec2InstanceConnectError::InstanceIdRequired).await;
131        }
132        if body
133            .get("InstanceOSUser")
134            .and_then(|v| v.as_str())
135            .is_none()
136        {
137            return shape_error(Ec2InstanceConnectError::InstanceOSUserRequired).await;
138        }
139        if body.get("SSHPublicKey").and_then(|v| v.as_str()).is_none() {
140            return shape_error(Ec2InstanceConnectError::SSHPublicKeyRequired).await;
141        }
142
143        wire::serialize_send_s_s_h_public_key_response(&wire::SendSSHPublicKeyResponse {
144            request_id: Some(uuid::Uuid::new_v4().to_string()),
145            success: Some(true),
146        })
147    }
148
149    async fn handle_send_serial_console_ssh_public_key(&self, body: &Value) -> MockResponse {
150        // Validate required fields
151        if body.get("InstanceId").and_then(|v| v.as_str()).is_none() {
152            return shape_error(Ec2InstanceConnectError::InstanceIdRequired).await;
153        }
154        if body.get("SSHPublicKey").and_then(|v| v.as_str()).is_none() {
155            return shape_error(Ec2InstanceConnectError::SSHPublicKeyRequired).await;
156        }
157
158        wire::serialize_send_serial_console_s_s_h_public_key_response(
159            &wire::SendSerialConsoleSSHPublicKeyResponse {
160                request_id: Some(uuid::Uuid::new_v4().to_string()),
161                success: Some(true),
162            },
163        )
164    }
165}
winterbaume_ec2instanceconnect/handlers.rs

winterbaume_ec2instanceconnect/
handlers.rs
