1use std::sync::Arc;
8use std::time::Duration;
9
10use alien_bindings::presigned::PresignedRequest;
11use alien_bindings::traits::{Kv, PutOptions, Storage};
12use alien_error::{AlienError, Context, ContextError, IntoAlienError};
13use chrono::{DateTime, Utc};
14use hex;
15use hmac::{Hmac, Mac};
16use serde::{Deserialize, Serialize};
17use sha2::Sha256;
18use tracing::{debug, info, warn};
19use uuid::Uuid;
20
21use base64::{engine::general_purpose, Engine as _};
22use bytes::Bytes;
23use object_store::path::Path as StoragePath;
24
25use crate::error::{ErrorData, Result};
26use crate::types::*;
27use crate::INLINE_MAX_BYTES;
28
29const KV_VALUE_THRESHOLD: usize = 20_000;
32
33fn is_definite_dispatch_rejection(error: &crate::error::Error) -> bool {
34 matches!(
35 error.error.as_ref(),
36 Some(ErrorData::TransportDispatchRejected { .. })
37 )
38}
39
40pub mod axum_handlers;
41pub mod command_registry;
42pub mod storage;
43
44pub use crate::dispatchers::{CommandDispatcher, NullCommandDispatcher};
45pub use axum_handlers::{
46 create_axum_router, CommandPayloadResponse, HasCommandServer, StorePayloadRequest,
47};
48pub use command_registry::{
49 delivery_mode_for, select_command_target, validate_command_name, validate_command_target_id,
50 CommandEnvelopeData, CommandMetadata, CommandRegistry, CommandStatus, InMemoryCommandRegistry,
51 ResolvedCommandTarget,
52};
53
54#[derive(Debug, Clone, Serialize, Deserialize)]
60struct CommandParamsData {
61 pub params: BodySpec,
62}
63
64#[derive(Debug, Clone, Serialize, Deserialize)]
66struct CommandResponseData {
67 pub response: CommandResponse,
68}
69
70#[derive(Debug, Clone, Serialize, Deserialize)]
72struct LeaseData {
73 pub lease_id: String,
74 pub acquired_at: DateTime<Utc>,
75 pub expires_at: DateTime<Utc>,
76 pub owner: String,
77}
78
79#[derive(Debug, Clone, Serialize, Deserialize)]
81struct DeadlineIndexData {
82 pub command_id: String,
83 pub deadline: DateTime<Utc>,
84}
85
86pub struct CommandServer {
95 kv: Arc<dyn Kv>,
96 storage: Arc<dyn Storage>,
97 command_dispatcher: Arc<dyn CommandDispatcher>,
98 command_registry: Arc<dyn CommandRegistry>,
99 inline_max_bytes: usize,
100 base_url: String,
101 response_signing_key: Vec<u8>,
102}
103
104impl CommandServer {
105 pub fn new(
107 kv: Arc<dyn Kv>,
108 storage: Arc<dyn Storage>,
109 command_dispatcher: Arc<dyn CommandDispatcher>,
110 command_registry: Arc<dyn CommandRegistry>,
111 base_url: String,
112 response_signing_key: Vec<u8>,
113 ) -> Self {
114 Self {
115 kv,
116 storage,
117 command_dispatcher,
118 command_registry,
119 inline_max_bytes: INLINE_MAX_BYTES,
120 base_url,
121 response_signing_key,
122 }
123 }
124
125 pub fn with_inline_limit(
127 kv: Arc<dyn Kv>,
128 storage: Arc<dyn Storage>,
129 command_dispatcher: Arc<dyn CommandDispatcher>,
130 command_registry: Arc<dyn CommandRegistry>,
131 base_url: String,
132 inline_max_bytes: usize,
133 response_signing_key: Vec<u8>,
134 ) -> Self {
135 Self {
136 kv,
137 storage,
138 command_dispatcher,
139 command_registry,
140 inline_max_bytes,
141 base_url,
142 response_signing_key,
143 }
144 }
145
146 const MAX_RESPONSE_TOKEN_LIFETIME_SECS: i64 = 7200;
148 const RESPONSE_CREDENTIAL_LIFETIME_SECS: u64 = 7200;
152
153 fn sign_response_url(&self, command_id: &str) -> (String, i64) {
158 let expires = Utc::now().timestamp() + Self::RESPONSE_CREDENTIAL_LIFETIME_SECS as i64;
159 let message = format!("commands.v1:{}:{}", command_id, expires);
160
161 type HmacSha256 = Hmac<Sha256>;
162 let mut mac =
163 HmacSha256::new_from_slice(&self.response_signing_key).expect("HMAC accepts any key");
164 mac.update(message.as_bytes());
165 let result = mac.finalize();
166 let token = hex::encode(result.into_bytes());
167
168 (token, expires)
169 }
170
171 pub fn verify_response_token(&self, command_id: &str, token: &str, expires: i64) -> bool {
176 if token.len() != 64 {
178 return false;
179 }
180
181 let message = format!("commands.v1:{}:{}", command_id, expires);
182
183 type HmacSha256 = Hmac<Sha256>;
184 let mut mac =
185 HmacSha256::new_from_slice(&self.response_signing_key).expect("HMAC accepts any key");
186 mac.update(message.as_bytes());
187
188 let Ok(token_bytes) = hex::decode(token) else {
190 return false;
191 };
192 let hmac_valid = mac.verify_slice(&token_bytes).is_ok();
193
194 let now = Utc::now().timestamp();
196 let not_expired = now <= expires;
197 let within_max_lifetime = expires <= now + Self::MAX_RESPONSE_TOKEN_LIFETIME_SECS;
198
199 hmac_valid && not_expired && within_max_lifetime
200 }
201
202 pub async fn create_command(
214 &self,
215 request: CreateCommandRequest,
216 ) -> Result<CreateCommandResponse> {
217 self.validate_create_command(&request).await?;
219
220 let resolved_target = self
223 .command_registry
224 .resolve_target(
225 &request.deployment_id,
226 request.target_resource_id.as_deref(),
227 )
228 .await?;
229
230 let composed_idempotency_key = request.idempotency_key.as_ref().map(|idem_key| {
237 Self::compose_idempotency_key(
238 &request.deployment_id,
239 &resolved_target.target.resource_id,
240 &request.command,
241 idem_key,
242 )
243 });
244
245 if let Some(ref composed_key) = composed_idempotency_key {
247 if let Some(existing_id) = self.check_idempotency(composed_key).await? {
248 let status = self
250 .command_registry
251 .get_command_status(&existing_id)
252 .await?;
253 if let Some(s) = status {
254 return Ok(CreateCommandResponse {
255 command_id: existing_id,
256 state: s.state,
257 storage_upload: None,
258 inline_allowed_up_to: self.inline_max_bytes as u64,
259 next: "poll".to_string(),
260 });
261 }
262 }
263 }
264
265 let (initial_state, request_size_bytes) = match &request.params {
267 BodySpec::Inline { inline_base64 } => {
268 let size = inline_base64.len() as u64;
269 (CommandState::Pending, Some(size))
270 }
271 BodySpec::Storage { size, .. } => {
272 if size.unwrap_or(0) > self.inline_max_bytes as u64 {
273 (CommandState::PendingUpload, *size)
274 } else {
275 (CommandState::Pending, *size)
276 }
277 }
278 };
279
280 let metadata = self
282 .command_registry
283 .create_command(
284 &request.deployment_id,
285 &request.command,
286 &resolved_target,
287 initial_state,
288 request.deadline,
289 request_size_bytes,
290 )
291 .await?;
292
293 let command_id = metadata.command_id;
294 let delivery_mode = metadata.delivery_mode;
295
296 if let Some(ref composed_key) = composed_idempotency_key {
303 if let Some(winner_id) = self.store_idempotency(composed_key, &command_id).await? {
304 let _ = self
305 .command_registry
306 .complete_command(
307 &command_id,
308 CommandState::Failed,
309 Utc::now(),
310 None,
311 Some(serde_json::json!({
312 "code": "IDEMPOTENT_DUPLICATE",
313 "message": format!(
314 "Superseded by concurrent create '{}' with the same idempotency key",
315 winner_id
316 ),
317 })),
318 )
319 .await?;
320 let status = self.command_registry.get_command_status(&winner_id).await?;
321 let state = status.map(|s| s.state).unwrap_or(CommandState::Pending);
322 return Ok(CreateCommandResponse {
323 command_id: winner_id,
324 state,
325 storage_upload: None,
326 inline_allowed_up_to: self.inline_max_bytes as u64,
327 next: "poll".to_string(),
328 });
329 }
330 }
331
332 self.store_params(&command_id, &request.params).await?;
334
335 let storage_upload = if initial_state == CommandState::PendingUpload {
337 Some(self.generate_params_upload(&command_id).await?)
338 } else {
339 None
340 };
341
342 let (final_state, next_action) = if initial_state == CommandState::Pending {
344 match delivery_mode {
345 CommandDeliveryMode::Push => {
346 let state = self
348 .dispatch_command_push(&command_id, &request.deployment_id)
349 .await?;
350 (state, "poll")
351 }
352 CommandDeliveryMode::Pull => {
353 self.create_pending_index(
356 &request.deployment_id,
357 &metadata.target.resource_id,
358 &command_id,
359 )
360 .await?;
361 debug!("Command {} ready for target-scoped lease", command_id);
362 (CommandState::Pending, "poll")
363 }
364 }
365 } else {
366 (initial_state, "upload")
368 };
369
370 if let Some(deadline) = request.deadline {
372 self.create_deadline_index(&command_id, deadline).await?;
373 }
374
375 Ok(CreateCommandResponse {
376 command_id,
377 state: final_state,
378 storage_upload,
379 inline_allowed_up_to: self.inline_max_bytes as u64,
380 next: next_action.to_string(),
381 })
382 }
383
384 pub async fn upload_complete(
386 &self,
387 command_id: &str,
388 upload_request: UploadCompleteRequest,
389 ) -> Result<UploadCompleteResponse> {
390 let status = self
392 .command_registry
393 .get_command_status(command_id)
394 .await?
395 .ok_or_else(|| {
396 AlienError::new(ErrorData::CommandNotFound {
397 command_id: command_id.to_string(),
398 })
399 })?;
400
401 if status.state != CommandState::PendingUpload {
403 return Err(AlienError::new(ErrorData::InvalidStateTransition {
404 from: status.state.as_ref().to_string(),
405 to: CommandState::Pending.as_ref().to_string(),
406 }));
407 }
408
409 let storage_get_request = self.generate_storage_get_request(command_id).await?;
411 let params = BodySpec::Storage {
412 size: Some(upload_request.size),
413 storage_get_request: Some(storage_get_request),
414 storage_put_used: None,
415 };
416 self.store_params(command_id, ¶ms).await?;
417
418 let transitioned = self
420 .command_registry
421 .update_command_state(command_id, CommandState::Pending, None, None, None, None)
422 .await?;
423 if !transitioned {
424 let current_state = self
425 .command_registry
426 .get_command_status(command_id)
427 .await?
428 .map(|current| current.state)
429 .unwrap_or(status.state);
430 return Err(AlienError::new(ErrorData::InvalidStateTransition {
431 from: current_state.as_ref().to_string(),
432 to: CommandState::Pending.as_ref().to_string(),
433 }));
434 }
435
436 let metadata = self
438 .command_registry
439 .get_command_metadata(command_id)
440 .await?
441 .ok_or_else(|| {
442 AlienError::new(ErrorData::CommandNotFound {
443 command_id: command_id.to_string(),
444 })
445 })?;
446
447 let final_state = match metadata.delivery_mode {
448 CommandDeliveryMode::Push => {
449 self.dispatch_command_push(command_id, &status.deployment_id)
450 .await?
451 }
452 CommandDeliveryMode::Pull => {
453 self.create_pending_index(
454 &status.deployment_id,
455 &metadata.target.resource_id,
456 command_id,
457 )
458 .await?;
459 debug!(
460 "Command {} ready for pull after upload (target will poll)",
461 command_id
462 );
463 CommandState::Pending
464 }
465 };
466
467 Ok(UploadCompleteResponse {
468 command_id: command_id.to_string(),
469 state: final_state,
470 })
471 }
472
473 pub async fn get_command_status(&self, command_id: &str) -> Result<CommandStatusResponse> {
477 let status = self
479 .command_registry
480 .get_command_status(command_id)
481 .await?
482 .ok_or_else(|| {
483 AlienError::new(ErrorData::CommandNotFound {
484 command_id: command_id.to_string(),
485 })
486 })?;
487
488 if let Some(deadline) = status.deadline {
490 if Utc::now() > deadline && !status.state.is_terminal() {
491 let won = self
496 .command_registry
497 .complete_command(command_id, CommandState::Expired, Utc::now(), None, None)
498 .await?;
499 if won {
500 self.delete_pending_index(
502 &status.deployment_id,
503 &status.target.resource_id,
504 command_id,
505 )
506 .await?;
507
508 return Ok(CommandStatusResponse {
510 command_id: command_id.to_string(),
511 state: CommandState::Expired,
512 attempt: status.attempt,
513 target: status.target,
514 response: None,
515 });
516 }
517 let status = self
520 .command_registry
521 .get_command_status(command_id)
522 .await?
523 .ok_or_else(|| {
524 AlienError::new(ErrorData::CommandNotFound {
525 command_id: command_id.to_string(),
526 })
527 })?;
528 let response = if status.state.is_terminal() {
529 self.get_response(command_id).await?
530 } else {
531 None
532 };
533 return Ok(CommandStatusResponse {
534 command_id: command_id.to_string(),
535 state: status.state,
536 attempt: status.attempt,
537 target: status.target,
538 response,
539 });
540 }
541 }
542
543 let response = if status.state.is_terminal() {
545 self.get_response(command_id).await?
546 } else {
547 None
548 };
549
550 Ok(CommandStatusResponse {
551 command_id: command_id.to_string(),
552 state: status.state,
553 attempt: status.attempt,
554 target: status.target,
555 response,
556 })
557 }
558
559 pub async fn submit_command_response(
563 &self,
564 command_id: &str,
565 mut response: CommandResponse,
566 ) -> Result<()> {
567 let status = self
569 .command_registry
570 .get_command_status(command_id)
571 .await?
572 .ok_or_else(|| {
573 AlienError::new(ErrorData::CommandNotFound {
574 command_id: command_id.to_string(),
575 })
576 })?;
577
578 if status.state.is_terminal() {
580 debug!(
581 "Ignoring duplicate response for terminal command {}",
582 command_id
583 );
584 return Ok(());
585 }
586
587 if status.state != CommandState::Dispatched {
589 return Err(AlienError::new(ErrorData::InvalidStateTransition {
590 from: status.state.as_ref().to_string(),
591 to: CommandState::Succeeded.as_ref().to_string(),
592 }));
593 }
594
595 if let CommandResponse::Success {
597 response: ref mut body,
598 } = response
599 {
600 if let BodySpec::Storage {
601 size,
602 storage_get_request,
603 storage_put_used,
604 } = body
605 {
606 if storage_get_request.is_none() && storage_put_used.unwrap_or(false) {
607 let get_request = self
608 .generate_response_storage_get_request(command_id)
609 .await?;
610 *body = BodySpec::Storage {
611 size: *size,
612 storage_get_request: Some(get_request),
613 storage_put_used: *storage_put_used,
614 };
615 }
616 }
617 }
618
619 self.store_response(command_id, &response).await?;
621
622 let (new_state, error) = if response.is_success() {
631 (CommandState::Succeeded, None)
632 } else if let CommandResponse::Error { code, message, .. } = &response {
633 (
634 CommandState::Failed,
635 Some(serde_json::json!({ "code": code, "message": message })),
636 )
637 } else {
638 (CommandState::Failed, None)
639 };
640
641 let response_size = match &response {
642 CommandResponse::Success {
643 response: BodySpec::Inline { inline_base64 },
644 } => Some(inline_base64.len() as u64),
645 CommandResponse::Success {
646 response: BodySpec::Storage { size, .. },
647 } => *size,
648 _ => None,
649 };
650
651 let won = self
661 .command_registry
662 .complete_command(command_id, new_state, Utc::now(), response_size, error)
663 .await?;
664 if !won {
665 debug!(
666 "Ignoring duplicate response for command {} (lost the terminal transition race)",
667 command_id
668 );
669 return Ok(());
670 }
671 self.store_response(command_id, &response).await?;
672
673 if let Err(e) = self.delete_lease(command_id).await {
678 warn!(
679 command_id,
680 error = %e,
681 "Failed to clean up lease after terminal response; will be reaped on next lease scan"
682 );
683 }
684
685 if let Err(e) = self
688 .delete_pending_index(
689 &status.deployment_id,
690 &status.target.resource_id,
691 command_id,
692 )
693 .await
694 {
695 warn!(
696 command_id,
697 error = %e,
698 "Failed to clean up pending index after terminal response; will be reaped on next lease scan"
699 );
700 }
701
702 info!(
703 "Command {} completed with state {:?}",
704 command_id, new_state
705 );
706 Ok(())
707 }
708
709 pub async fn acquire_lease(
713 &self,
714 deployment_id: &str,
715 lease_request: &LeaseRequest,
716 ) -> Result<LeaseResponse> {
717 let mut leases = Vec::new();
718
719 let target_prefix = format!(
722 "target:{}:{}:pending:",
723 deployment_id, lease_request.target.resource_id
724 );
725 let scan_result = self
726 .kv
727 .scan_prefix(&target_prefix, Some(lease_request.max_leases * 2), None)
728 .await
729 .into_alien_error()
730 .context(ErrorData::KvOperationFailed {
731 operation: "scan_prefix".to_string(),
732 key: target_prefix.clone(),
733 message: "Failed to scan for pending commands".to_string(),
734 })?;
735
736 for (index_key, _) in scan_result.items {
737 if leases.len() >= lease_request.max_leases {
738 break;
739 }
740
741 let command_id = self.extract_command_id_from_index_key(&index_key)?;
742
743 let lease_id = format!("lease_{}", Uuid::new_v4());
745 let lease_duration = Duration::from_secs(lease_request.lease_seconds);
746 let expires_at =
747 Utc::now() + chrono::Duration::seconds(lease_request.lease_seconds as i64);
748
749 let lease_data = LeaseData {
750 lease_id: lease_id.clone(),
751 acquired_at: Utc::now(),
752 expires_at,
753 owner: deployment_id.to_string(),
754 };
755
756 let lease_key = format!("cmd:{}:lease", command_id);
757 let lease_value = serde_json::to_vec(&lease_data).into_alien_error().context(
758 ErrorData::SerializationFailed {
759 message: "Failed to serialize lease data".to_string(),
760 data_type: Some("LeaseData".to_string()),
761 },
762 )?;
763
764 let options = Some(PutOptions {
765 ttl: Some(lease_duration),
766 if_not_exists: true,
767 });
768
769 let success = self
770 .kv
771 .put(&lease_key, lease_value, options)
772 .await
773 .context(ErrorData::KvOperationFailed {
774 operation: "put".to_string(),
775 key: lease_key.clone(),
776 message: "Failed to create lease".to_string(),
777 })?;
778
779 if !success {
780 continue;
782 }
783
784 let reverse_key = format!("lease:{}", lease_id);
789 self.kv
790 .put(
791 &reverse_key,
792 command_id.as_bytes().to_vec(),
793 Some(PutOptions {
794 ttl: Some(lease_duration),
795 if_not_exists: false,
796 }),
797 )
798 .await
799 .context(ErrorData::KvOperationFailed {
800 operation: "put".to_string(),
801 key: reverse_key,
802 message: "Failed to create lease reverse index".to_string(),
803 })?;
804
805 let mut metadata = match self
807 .command_registry
808 .get_command_metadata(&command_id)
809 .await?
810 {
811 Some(m) => m,
812 None => {
813 self.delete_lease(&command_id).await?;
815 let _ = self.kv.delete(&index_key).await;
816 continue;
817 }
818 };
819
820 if metadata.state == CommandState::Dispatched {
828 self.command_registry.increment_attempt(&command_id).await?;
829 metadata.attempt += 1;
830 }
831
832 if metadata.target != lease_request.target {
840 self.delete_lease(&command_id).await?;
841 return Err(AlienError::new(ErrorData::Other {
842 message: format!(
843 "Pending index corruption: command '{}' is indexed under target '{}' \
844 but the registry says it belongs to target '{}' — refusing to deliver",
845 command_id, lease_request.target.resource_id, metadata.target.resource_id,
846 ),
847 }));
848 }
849
850 if metadata.state.is_terminal() {
852 self.delete_lease(&command_id).await?;
854 let _ = self.kv.delete(&index_key).await;
855 continue;
856 }
857
858 if let Some(deadline) = metadata.deadline {
861 if Utc::now() > deadline {
862 let _ = self
863 .command_registry
864 .complete_command(
865 &command_id,
866 CommandState::Expired,
867 Utc::now(),
868 None,
869 None,
870 )
871 .await?;
872 self.delete_lease(&command_id).await?;
873 let _ = self.kv.delete(&index_key).await;
874 continue;
875 }
876 }
877
878 let params = match self.get_params(&command_id).await? {
880 Some(p) => p,
881 None => {
882 self.delete_lease(&command_id).await?;
884 continue;
885 }
886 };
887
888 let dispatched = self
893 .command_registry
894 .mark_dispatched_if_not_terminal(&command_id, Utc::now())
895 .await?;
896 if !dispatched {
897 debug!(
898 command_id = %command_id,
899 "Command turned terminal while leasing; releasing the lease"
900 );
901 self.delete_lease(&command_id).await?;
902 let _ = self.kv.delete(&index_key).await;
903 continue;
904 }
905
906 let mut envelope = self.build_envelope(&command_id, &metadata, params).await?;
915 Self::relativize_manager_urls(&mut envelope, &self.base_url);
916
917 leases.push(LeaseInfo {
918 lease_id,
919 lease_expires_at: expires_at,
920 command_id: command_id.clone(),
921 attempt: metadata.attempt,
922 envelope,
923 });
924 }
925
926 Ok(LeaseResponse { leases })
927 }
928
929 fn relativize_manager_urls(envelope: &mut Envelope, base_url: &str) {
934 let Ok(mut lease_endpoint) = reqwest::Url::parse(base_url) else {
935 return;
938 };
939 let Ok(mut segments) = lease_endpoint.path_segments_mut() else {
940 return;
941 };
942 segments.pop_if_empty().push("commands").push("leases");
943 drop(segments);
944 lease_endpoint.set_query(None);
945 lease_endpoint.set_fragment(None);
946
947 let relativize = |target: &mut String| {
948 let Ok(mut target_url) = reqwest::Url::parse(target.as_str()) else {
949 return;
950 };
951 let suffix = target
952 .find(['?', '#'])
953 .map(|index| target[index..].to_string())
954 .unwrap_or_default();
955 target_url.set_query(None);
956 target_url.set_fragment(None);
957 if let Some(relative) = lease_endpoint.make_relative(&target_url) {
958 *target = format!("{relative}{suffix}");
959 }
960 };
961
962 relativize(&mut envelope.response_handling.submit_response_url);
963 if let alien_core::presigned::PresignedRequestBackend::Http { url, .. } =
964 &mut envelope.response_handling.storage_upload_request.backend
965 {
966 relativize(url);
967 }
968 if let BodySpec::Storage {
969 storage_get_request: Some(request),
970 ..
971 } = &mut envelope.params
972 {
973 if let alien_core::presigned::PresignedRequestBackend::Http { url, .. } =
974 &mut request.backend
975 {
976 relativize(url);
977 }
978 }
979 }
980
981 pub async fn release_lease(&self, command_id: &str, lease_id: &str) -> Result<()> {
985 let lease_key = format!("cmd:{}:lease", command_id);
986
987 if let Ok(Some(lease_data)) = self.kv.get(&lease_key).await {
989 let lease: LeaseData = serde_json::from_slice(&lease_data)
990 .into_alien_error()
991 .context(ErrorData::SerializationFailed {
992 message: "Failed to deserialize lease data".to_string(),
993 data_type: Some("LeaseData".to_string()),
994 })?;
995
996 if lease.lease_id != lease_id {
997 return Err(AlienError::new(ErrorData::LeaseNotFound {
998 lease_id: lease_id.to_string(),
999 }));
1000 }
1001
1002 self.delete_lease(command_id).await?;
1004 let _ = self.kv.delete(&format!("lease:{}", lease_id)).await;
1005
1006 let released = self
1009 .command_registry
1010 .update_command_state(command_id, CommandState::Pending, None, None, None, None)
1011 .await?;
1012 if released {
1013 self.command_registry.increment_attempt(command_id).await?;
1016 }
1017
1018 debug!(
1020 released,
1021 "Lease {} released for command {}", lease_id, command_id
1022 );
1023 }
1024
1025 Ok(())
1026 }
1027
1028 pub async fn get_command_deployment_id(&self, command_id: &str) -> Result<Option<String>> {
1033 let status = self.command_registry.get_command_status(command_id).await?;
1034 Ok(status.map(|s| s.deployment_id))
1035 }
1036
1037 pub async fn get_lease_owner(&self, lease_id: &str) -> Result<Option<(String, String)>> {
1045 let reverse_key = format!("lease:{}", lease_id);
1046 let Some(command_id_bytes) =
1047 self.kv
1048 .get(&reverse_key)
1049 .await
1050 .context(ErrorData::KvOperationFailed {
1051 operation: "get".to_string(),
1052 key: reverse_key.clone(),
1053 message: "Failed to look up lease reverse index".to_string(),
1054 })?
1055 else {
1056 return Ok(None);
1057 };
1058 let command_id = String::from_utf8(command_id_bytes).map_err(|_| {
1059 AlienError::new(ErrorData::Other {
1060 message: format!("Lease reverse index '{}' is not valid UTF-8", reverse_key),
1061 })
1062 })?;
1063
1064 let lease_key = format!("cmd:{}:lease", command_id);
1065 let Some(lease_data) =
1066 self.kv
1067 .get(&lease_key)
1068 .await
1069 .context(ErrorData::KvOperationFailed {
1070 operation: "get".to_string(),
1071 key: lease_key,
1072 message: "Failed to look up lease".to_string(),
1073 })?
1074 else {
1075 return Ok(None);
1076 };
1077 let lease: LeaseData = serde_json::from_slice(&lease_data)
1078 .into_alien_error()
1079 .context(ErrorData::SerializationFailed {
1080 message: "Failed to deserialize lease data".to_string(),
1081 data_type: Some("LeaseData".to_string()),
1082 })?;
1083 if lease.lease_id != lease_id {
1084 return Ok(None);
1085 }
1086
1087 Ok(Some((command_id, lease.owner)))
1088 }
1089
1090 pub async fn release_lease_by_id(&self, lease_id: &str) -> Result<()> {
1092 match self.get_lease_owner(lease_id).await? {
1093 Some((command_id, _owner)) => self.release_lease(&command_id, lease_id).await,
1094 None => Err(AlienError::new(ErrorData::LeaseNotFound {
1095 lease_id: lease_id.to_string(),
1096 })),
1097 }
1098 }
1099
1100 async fn validate_create_command(&self, request: &CreateCommandRequest) -> Result<()> {
1105 if request.command.is_empty() {
1106 return Err(AlienError::new(ErrorData::InvalidCommand {
1107 message: "Command name cannot be empty".to_string(),
1108 }));
1109 }
1110
1111 validate_command_name(&request.command)?;
1119
1120 if request.deployment_id.is_empty() {
1121 return Err(AlienError::new(ErrorData::InvalidCommand {
1122 message: "Deployment ID cannot be empty".to_string(),
1123 }));
1124 }
1125
1126 if let Some(deadline) = request.deadline {
1127 if deadline <= Utc::now() {
1128 return Err(AlienError::new(ErrorData::InvalidCommand {
1129 message: "Deadline must be in the future".to_string(),
1130 }));
1131 }
1132 }
1133
1134 Ok(())
1135 }
1136
1137 fn compose_idempotency_key(
1145 deployment_id: &str,
1146 target_resource_id: &str,
1147 command_name: &str,
1148 idem_key: &str,
1149 ) -> String {
1150 debug_assert!(
1160 !target_resource_id.contains(':'),
1161 "target_resource_id must be ':'-free before key composition: {target_resource_id}"
1162 );
1163 format!(
1164 "{}:{}:{}:{}",
1165 deployment_id, target_resource_id, command_name, idem_key
1166 )
1167 }
1168
1169 async fn check_idempotency(&self, idem_key: &str) -> Result<Option<String>> {
1170 let key = format!("idem:{}", idem_key);
1171 if let Some(data) = self
1172 .kv
1173 .get(&key)
1174 .await
1175 .context(ErrorData::KvOperationFailed {
1176 operation: "get".to_string(),
1177 key: key.clone(),
1178 message: "Failed to check idempotency".to_string(),
1179 })?
1180 {
1181 let command_id = String::from_utf8(data).into_alien_error().context(
1182 ErrorData::SerializationFailed {
1183 message: "Invalid idempotency data".to_string(),
1184 data_type: Some("String".to_string()),
1185 },
1186 )?;
1187 return Ok(Some(command_id));
1188 }
1189 Ok(None)
1190 }
1191
1192 async fn store_idempotency(&self, idem_key: &str, command_id: &str) -> Result<Option<String>> {
1200 let key = format!("idem:{}", idem_key);
1201 let ttl = Duration::from_secs(24 * 60 * 60); let won = self
1203 .kv
1204 .put(
1205 &key,
1206 command_id.as_bytes().to_vec(),
1207 Some(PutOptions {
1208 ttl: Some(ttl),
1209 if_not_exists: true,
1210 }),
1211 )
1212 .await
1213 .context(ErrorData::KvOperationFailed {
1214 operation: "put".to_string(),
1215 key: key.clone(),
1216 message: "Failed to store idempotency".to_string(),
1217 })?;
1218 if won {
1219 return Ok(None);
1220 }
1221 match self.check_idempotency(idem_key).await? {
1225 Some(winner_id) => Ok(Some(winner_id)),
1226 None => Err(AlienError::new(ErrorData::Other {
1227 message: format!(
1228 "Idempotency key '{}' was concurrently claimed but has no winner entry",
1229 key
1230 ),
1231 })),
1232 }
1233 }
1234
1235 pub async fn store_params(&self, command_id: &str, params: &BodySpec) -> Result<()> {
1238 let key = format!("cmd:{}:params", command_id);
1239
1240 let data = CommandParamsData {
1242 params: params.clone(),
1243 };
1244 let value = serde_json::to_vec(&data).into_alien_error().context(
1245 ErrorData::SerializationFailed {
1246 message: "Failed to serialize params".to_string(),
1247 data_type: Some("CommandParamsData".to_string()),
1248 },
1249 )?;
1250
1251 if value.len() <= KV_VALUE_THRESHOLD {
1253 self.kv
1254 .put(&key, value, None)
1255 .await
1256 .context(ErrorData::KvOperationFailed {
1257 operation: "put".to_string(),
1258 key: key.clone(),
1259 message: "Failed to store params".to_string(),
1260 })?;
1261 return Ok(());
1262 }
1263
1264 if let BodySpec::Inline { inline_base64 } = params {
1266 let raw_bytes = general_purpose::STANDARD
1267 .decode(inline_base64)
1268 .into_alien_error()
1269 .context(ErrorData::SerializationFailed {
1270 message: "Failed to decode inline base64 params for auto-promotion".to_string(),
1271 data_type: Some("base64".to_string()),
1272 })?;
1273
1274 let raw_len = raw_bytes.len() as u64;
1275 let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1276
1277 self.storage
1278 .put(&blob_path, Bytes::from(raw_bytes).into())
1279 .await
1280 .into_alien_error()
1281 .context(ErrorData::StorageOperationFailed {
1282 message: "Failed to auto-promote params to blob storage".to_string(),
1283 operation: Some("put".to_string()),
1284 path: Some(blob_path.to_string()),
1285 })?;
1286
1287 debug!(
1288 "Auto-promoted params for command {} to blob ({} bytes raw)",
1289 command_id, raw_len
1290 );
1291
1292 let promoted = CommandParamsData {
1294 params: BodySpec::Storage {
1295 size: Some(raw_len),
1296 storage_get_request: None,
1297 storage_put_used: Some(true),
1298 },
1299 };
1300 let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1301 ErrorData::SerializationFailed {
1302 message: "Failed to serialize promoted params reference".to_string(),
1303 data_type: Some("CommandParamsData".to_string()),
1304 },
1305 )?;
1306 self.kv.put(&key, promoted_value, None).await.context(
1307 ErrorData::KvOperationFailed {
1308 operation: "put".to_string(),
1309 key: key.clone(),
1310 message: "Failed to store promoted params reference".to_string(),
1311 },
1312 )?;
1313 return Ok(());
1314 }
1315
1316 self.kv
1318 .put(&key, value, None)
1319 .await
1320 .context(ErrorData::KvOperationFailed {
1321 operation: "put".to_string(),
1322 key: key.clone(),
1323 message: "Failed to store params".to_string(),
1324 })?;
1325 Ok(())
1326 }
1327
1328 pub async fn get_params(&self, command_id: &str) -> Result<Option<BodySpec>> {
1329 let key = format!("cmd:{}:params", command_id);
1330 if let Some(value) = self
1331 .kv
1332 .get(&key)
1333 .await
1334 .context(ErrorData::KvOperationFailed {
1335 operation: "get".to_string(),
1336 key: key.clone(),
1337 message: "Failed to get params".to_string(),
1338 })?
1339 {
1340 let data: CommandParamsData = serde_json::from_slice(&value)
1341 .into_alien_error()
1342 .context(ErrorData::SerializationFailed {
1343 message: "Failed to deserialize params".to_string(),
1344 data_type: Some("CommandParamsData".to_string()),
1345 })?;
1346 return Ok(Some(data.params));
1347 }
1348 Ok(None)
1349 }
1350
1351 pub async fn store_response(&self, command_id: &str, response: &CommandResponse) -> Result<()> {
1354 let key = format!("cmd:{}:response", command_id);
1355 let data = CommandResponseData {
1356 response: response.clone(),
1357 };
1358 let value = serde_json::to_vec(&data).into_alien_error().context(
1359 ErrorData::SerializationFailed {
1360 message: "Failed to serialize response".to_string(),
1361 data_type: Some("CommandResponseData".to_string()),
1362 },
1363 )?;
1364
1365 if value.len() <= KV_VALUE_THRESHOLD {
1367 self.kv
1368 .put(&key, value, None)
1369 .await
1370 .context(ErrorData::KvOperationFailed {
1371 operation: "put".to_string(),
1372 key: key.clone(),
1373 message: "Failed to store response".to_string(),
1374 })?;
1375 return Ok(());
1376 }
1377
1378 if let CommandResponse::Success {
1380 response: BodySpec::Inline { inline_base64 },
1381 } = response
1382 {
1383 let raw_bytes = general_purpose::STANDARD
1384 .decode(inline_base64)
1385 .into_alien_error()
1386 .context(ErrorData::SerializationFailed {
1387 message: "Failed to decode inline base64 response for auto-promotion"
1388 .to_string(),
1389 data_type: Some("base64".to_string()),
1390 })?;
1391
1392 let raw_len = raw_bytes.len() as u64;
1393 let blob_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1394
1395 self.storage
1396 .put(&blob_path, Bytes::from(raw_bytes).into())
1397 .await
1398 .into_alien_error()
1399 .context(ErrorData::StorageOperationFailed {
1400 message: "Failed to auto-promote response to blob storage".to_string(),
1401 operation: Some("put".to_string()),
1402 path: Some(blob_path.to_string()),
1403 })?;
1404
1405 let get_request = self
1407 .generate_response_storage_get_request(command_id)
1408 .await?;
1409
1410 debug!(
1411 "Auto-promoted response for command {} to blob ({} bytes raw)",
1412 command_id, raw_len
1413 );
1414
1415 let promoted = CommandResponseData {
1417 response: CommandResponse::Success {
1418 response: BodySpec::Storage {
1419 size: Some(raw_len),
1420 storage_get_request: Some(get_request),
1421 storage_put_used: Some(true),
1422 },
1423 },
1424 };
1425 let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1426 ErrorData::SerializationFailed {
1427 message: "Failed to serialize promoted response reference".to_string(),
1428 data_type: Some("CommandResponseData".to_string()),
1429 },
1430 )?;
1431 self.kv.put(&key, promoted_value, None).await.context(
1432 ErrorData::KvOperationFailed {
1433 operation: "put".to_string(),
1434 key: key.clone(),
1435 message: "Failed to store promoted response reference".to_string(),
1436 },
1437 )?;
1438 return Ok(());
1439 }
1440
1441 self.kv
1443 .put(&key, value, None)
1444 .await
1445 .context(ErrorData::KvOperationFailed {
1446 operation: "put".to_string(),
1447 key: key.clone(),
1448 message: "Failed to store response".to_string(),
1449 })?;
1450 Ok(())
1451 }
1452
1453 pub async fn get_response(&self, command_id: &str) -> Result<Option<CommandResponse>> {
1454 let key = format!("cmd:{}:response", command_id);
1455 if let Some(value) = self
1456 .kv
1457 .get(&key)
1458 .await
1459 .context(ErrorData::KvOperationFailed {
1460 operation: "get".to_string(),
1461 key: key.clone(),
1462 message: "Failed to get response".to_string(),
1463 })?
1464 {
1465 let data: CommandResponseData = serde_json::from_slice(&value)
1466 .into_alien_error()
1467 .context(ErrorData::SerializationFailed {
1468 message: "Failed to deserialize response".to_string(),
1469 data_type: Some("CommandResponseData".to_string()),
1470 })?;
1471 return Ok(Some(data.response));
1472 }
1473 Ok(None)
1474 }
1475
1476 async fn create_pending_index(
1479 &self,
1480 deployment_id: &str,
1481 target_resource_id: &str,
1482 command_id: &str,
1483 ) -> Result<()> {
1484 debug_assert!(
1488 !target_resource_id.contains(':'),
1489 "target_resource_id must be ':'-free in the pending index: {target_resource_id}"
1490 );
1491 let timestamp = Utc::now().timestamp_nanos_opt().unwrap_or(0);
1492 let key = format!(
1493 "target:{}:{}:pending:{}:{}",
1494 deployment_id, target_resource_id, timestamp, command_id
1495 );
1496
1497 self.kv
1499 .put(&key, vec![], None)
1500 .await
1501 .context(ErrorData::KvOperationFailed {
1502 operation: "put".to_string(),
1503 key: key.clone(),
1504 message: "Failed to create pending index".to_string(),
1505 })?;
1506 Ok(())
1507 }
1508
1509 async fn delete_pending_index(
1510 &self,
1511 deployment_id: &str,
1512 target_resource_id: &str,
1513 command_id: &str,
1514 ) -> Result<()> {
1515 let prefix = format!("target:{}:{}:pending:", deployment_id, target_resource_id);
1517 let scan_result = self
1518 .kv
1519 .scan_prefix(&prefix, Some(100), None)
1520 .await
1521 .into_alien_error()
1522 .context(ErrorData::KvOperationFailed {
1523 operation: "scan_prefix".to_string(),
1524 key: prefix.clone(),
1525 message: "Failed to scan pending index".to_string(),
1526 })?;
1527
1528 for (key, _) in scan_result.items {
1529 if key.ends_with(&format!(":{}", command_id)) {
1530 let _ = self.kv.delete(&key).await;
1531 break;
1532 }
1533 }
1534 Ok(())
1535 }
1536
1537 async fn delete_lease(&self, command_id: &str) -> Result<()> {
1540 let key = format!("cmd:{}:lease", command_id);
1541 let _ = self.kv.delete(&key).await;
1542 Ok(())
1543 }
1544
1545 pub async fn reap_expired_commands(&self) -> Result<u32> {
1557 let now = Utc::now();
1558 let mut expired = 0u32;
1559 let mut cursor: Option<String> = None;
1560 for _ in 0..64 {
1565 let scan = self
1566 .kv
1567 .scan_prefix("deadline:", Some(256), cursor.clone())
1568 .await
1569 .into_alien_error()
1570 .context(ErrorData::KvOperationFailed {
1571 operation: "scan_prefix".to_string(),
1572 key: "deadline:".to_string(),
1573 message: "Failed to scan the deadline index".to_string(),
1574 })?;
1575 let next_cursor = scan.next_cursor.clone();
1576 for (key, value) in scan.items {
1577 let Ok(data) = serde_json::from_slice::<DeadlineIndexData>(&value) else {
1578 warn!(key = %key, "Unparseable deadline index entry; deleting");
1579 let _ = self.kv.delete(&key).await;
1580 continue;
1581 };
1582 if data.deadline > now {
1583 continue;
1586 }
1587
1588 let status = self
1589 .command_registry
1590 .get_command_status(&data.command_id)
1591 .await?;
1592 match status {
1593 None => {
1594 let _ = self.kv.delete(&key).await;
1595 }
1596 Some(status) if status.state.is_terminal() => {
1597 let _ = self.kv.delete(&key).await;
1598 }
1599 Some(status) => {
1600 let won = self
1601 .command_registry
1602 .complete_command(
1603 &data.command_id,
1604 CommandState::Expired,
1605 now,
1606 None,
1607 Some(serde_json::json!({
1608 "code": "COMMAND_EXPIRED",
1609 "message": format!("Deadline {} elapsed", data.deadline.to_rfc3339()),
1610 })),
1611 )
1612 .await?;
1613 if won {
1614 expired += 1;
1615 info!(command_id = %data.command_id, "Expired overdue command");
1616 let _ = self.delete_lease(&data.command_id).await;
1617 let _ = self
1618 .delete_pending_index(
1619 &status.deployment_id,
1620 &status.target.resource_id,
1621 &data.command_id,
1622 )
1623 .await;
1624 }
1625 let _ = self.kv.delete(&key).await;
1626 }
1627 }
1628 }
1629 match next_cursor {
1630 Some(next) => cursor = Some(next),
1631 None => break,
1632 }
1633 }
1634 Ok(expired)
1635 }
1636
1637 async fn create_deadline_index(&self, command_id: &str, deadline: DateTime<Utc>) -> Result<()> {
1640 let key = format!(
1641 "deadline:{}:{}",
1642 deadline.timestamp_nanos_opt().unwrap_or(0),
1643 command_id
1644 );
1645
1646 let data = DeadlineIndexData {
1647 command_id: command_id.to_string(),
1648 deadline,
1649 };
1650 let value = serde_json::to_vec(&data).into_alien_error().context(
1651 ErrorData::SerializationFailed {
1652 message: "Failed to serialize deadline index".to_string(),
1653 data_type: Some("DeadlineIndexData".to_string()),
1654 },
1655 )?;
1656
1657 const DEADLINE_INDEX_GRACE: chrono::Duration = chrono::Duration::days(7);
1664 let ttl = deadline
1665 .signed_duration_since(Utc::now())
1666 .checked_add(&DEADLINE_INDEX_GRACE)
1667 .unwrap_or(DEADLINE_INDEX_GRACE);
1668 let options = (ttl.num_seconds() > 0).then(|| PutOptions {
1669 ttl: Some(Duration::from_secs(ttl.num_seconds() as u64)),
1670 if_not_exists: false,
1671 });
1672
1673 self.kv
1674 .put(&key, value, options)
1675 .await
1676 .context(ErrorData::KvOperationFailed {
1677 operation: "put".to_string(),
1678 key: key.clone(),
1679 message: "Failed to create deadline index".to_string(),
1680 })?;
1681 Ok(())
1682 }
1683
1684 async fn dispatch_command_push(
1687 &self,
1688 command_id: &str,
1689 deployment_id: &str,
1690 ) -> Result<CommandState> {
1691 let metadata = self
1693 .command_registry
1694 .get_command_metadata(command_id)
1695 .await?
1696 .ok_or_else(|| {
1697 AlienError::new(ErrorData::CommandNotFound {
1698 command_id: command_id.to_string(),
1699 })
1700 })?;
1701
1702 let params = self.get_params(command_id).await?.ok_or_else(|| {
1704 AlienError::new(ErrorData::CommandNotFound {
1705 command_id: command_id.to_string(),
1706 })
1707 })?;
1708
1709 let envelope = self.build_envelope(command_id, &metadata, params).await?;
1711
1712 if !self
1718 .command_registry
1719 .mark_dispatched_if_not_terminal(command_id, Utc::now())
1720 .await?
1721 {
1722 return Ok(self
1723 .command_registry
1724 .get_command_status(command_id)
1725 .await?
1726 .map(|status| status.state)
1727 .unwrap_or(CommandState::Dispatched));
1728 }
1729
1730 if let Err(error) = self.command_dispatcher.dispatch(&envelope).await {
1738 if is_definite_dispatch_rejection(&error) {
1739 let delivery_failure = CommandResponse::error(
1740 "DELIVERY_FAILED",
1741 "Worker runtime did not accept command delivery",
1742 );
1743 self.submit_command_response(command_id, delivery_failure)
1744 .await?;
1745 warn!(
1746 command_id,
1747 deployment_id,
1748 error = %error,
1749 "Push dispatch was definitely rejected; command marked Failed"
1750 );
1751 return Ok(CommandState::Failed);
1752 }
1753
1754 let error = error.context(ErrorData::TransportDispatchFailed {
1755 message: "Failed to dispatch command".to_string(),
1756 transport_type: None,
1757 target: Some(deployment_id.to_string()),
1758 });
1759 warn!(
1760 command_id,
1761 deployment_id,
1762 error = %error,
1763 "Push dispatch acknowledgement failed; command remains Dispatched for a possible late response"
1764 );
1765 return Ok(CommandState::Dispatched);
1766 }
1767
1768 info!("Command {} dispatched via push", envelope.command_id);
1769 Ok(CommandState::Dispatched)
1770 }
1771
1772 async fn build_envelope(
1773 &self,
1774 command_id: &str,
1775 metadata: &CommandEnvelopeData,
1776 mut params: BodySpec,
1777 ) -> Result<Envelope> {
1778 let response_handling = self.create_response_handling(command_id).await?;
1779
1780 if let BodySpec::Storage { size, .. } = ¶ms {
1783 let raw_size = size.unwrap_or(0) as usize;
1784 if raw_size > 0 && raw_size <= self.inline_max_bytes {
1785 let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1786 match self.storage.get(&blob_path).await {
1787 Ok(get_result) => match get_result.bytes().await {
1788 Ok(raw_bytes) => {
1789 params = BodySpec::inline(&raw_bytes);
1790 debug!(
1791 "Re-inlined params for command {} ({} bytes) into envelope",
1792 command_id, raw_size
1793 );
1794 }
1795 Err(e) => {
1796 debug!(
1797 "Failed to read blob bytes for re-inline (command {}), falling back to presigned URL: {}",
1798 command_id, e
1799 );
1800 }
1801 },
1802 Err(e) => {
1803 debug!(
1804 "Failed to read blob for re-inline (command {}), falling back to presigned URL: {}",
1805 command_id, e
1806 );
1807 }
1808 }
1809 }
1810 }
1811
1812 if let BodySpec::Storage {
1817 size,
1818 storage_get_request: _,
1819 storage_put_used,
1820 } = ¶ms
1821 {
1822 let get_request = self.generate_storage_get_request(command_id).await?;
1823 params = BodySpec::Storage {
1824 size: *size,
1825 storage_get_request: Some(get_request),
1826 storage_put_used: *storage_put_used,
1827 };
1828 }
1829
1830 Ok(Envelope::new(
1831 metadata.deployment_id.clone(),
1832 metadata.target.clone(),
1833 command_id.to_string(),
1834 metadata.attempt,
1835 metadata.deadline,
1836 metadata.command.clone(),
1837 params,
1838 response_handling,
1839 ))
1840 }
1841
1842 async fn create_response_handling(&self, command_id: &str) -> Result<ResponseHandling> {
1843 let upload_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1844 let expires_in = Duration::from_secs(Self::RESPONSE_CREDENTIAL_LIFETIME_SECS);
1845 let presigned = self
1846 .storage
1847 .presigned_put(&upload_path, expires_in)
1848 .await
1849 .context(ErrorData::StorageOperationFailed {
1850 message: "Failed to create response upload URL".to_string(),
1851 operation: Some("presigned_put".to_string()),
1852 path: Some(upload_path.to_string()),
1853 })?;
1854
1855 let (response_token, expires) = self.sign_response_url(command_id);
1856
1857 Ok(ResponseHandling {
1858 max_inline_bytes: self.inline_max_bytes as u64,
1859 submit_response_url: format!(
1860 "{}/commands/{}/response?response_token={}&expires={}",
1861 self.base_url.trim_end_matches('/'),
1862 command_id,
1863 response_token,
1864 expires,
1865 ),
1866 storage_upload_request: presigned,
1867 })
1868 }
1869
1870 async fn generate_params_upload(&self, command_id: &str) -> Result<StorageUpload> {
1871 let upload_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1872 let expires_in = Duration::from_secs(3600);
1873 let presigned = self
1874 .storage
1875 .presigned_put(&upload_path, expires_in)
1876 .await
1877 .into_alien_error()
1878 .context(ErrorData::StorageOperationFailed {
1879 message: "Failed to create presigned URL".to_string(),
1880 operation: Some("presigned_put".to_string()),
1881 path: Some(upload_path.to_string()),
1882 })?;
1883
1884 Ok(StorageUpload {
1885 put_request: presigned.clone(),
1886 expires_at: presigned.expiration,
1887 })
1888 }
1889
1890 async fn generate_storage_get_request(&self, command_id: &str) -> Result<PresignedRequest> {
1891 let path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1892 let expires_in = Duration::from_secs(3600);
1893 self.storage.presigned_get(&path, expires_in).await.context(
1894 ErrorData::StorageOperationFailed {
1895 message: "Failed to create storage get request".to_string(),
1896 operation: Some("presigned_get".to_string()),
1897 path: Some(path.to_string()),
1898 },
1899 )
1900 }
1901
1902 async fn generate_response_storage_get_request(
1903 &self,
1904 command_id: &str,
1905 ) -> Result<PresignedRequest> {
1906 let path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1907 let expires_in = Duration::from_secs(3600);
1908 self.storage.presigned_get(&path, expires_in).await.context(
1909 ErrorData::StorageOperationFailed {
1910 message: "Failed to create response storage get request".to_string(),
1911 operation: Some("presigned_get".to_string()),
1912 path: Some(path.to_string()),
1913 },
1914 )
1915 }
1916
1917 fn extract_command_id_from_index_key(&self, index_key: &str) -> Result<String> {
1918 index_key
1919 .split(':')
1920 .last()
1921 .ok_or_else(|| {
1922 AlienError::new(ErrorData::Other {
1923 message: format!("Invalid index key format: {}", index_key),
1924 })
1925 })
1926 .map(|s| s.to_string())
1927 }
1928}
1929
1930#[cfg(test)]
1931mod relative_url_tests {
1932 use super::*;
1933 use alien_core::presigned::PresignedOperation;
1934 use std::collections::HashMap;
1935
1936 fn http_request(url: &str, operation: PresignedOperation) -> PresignedRequest {
1937 PresignedRequest::new_http(
1938 url.to_string(),
1939 match operation {
1940 PresignedOperation::Get => "GET",
1941 PresignedOperation::Put => "PUT",
1942 PresignedOperation::Delete => "DELETE",
1943 }
1944 .to_string(),
1945 HashMap::new(),
1946 operation,
1947 "commands/test".to_string(),
1948 Utc::now() + chrono::Duration::minutes(5),
1949 )
1950 }
1951
1952 #[test]
1953 fn leased_manager_urls_are_relative_to_lease_endpoint() {
1954 let mut envelope = Envelope::new(
1955 "deployment",
1956 CommandTarget::new("daemon", CommandTargetType::Daemon),
1957 "command",
1958 1,
1959 None,
1960 "run",
1961 BodySpec::Storage {
1962 size: Some(2048),
1963 storage_get_request: Some(http_request(
1964 "http://manager.internal/storage/params?signature=params",
1965 PresignedOperation::Get,
1966 )),
1967 storage_put_used: Some(true),
1968 },
1969 ResponseHandling {
1970 max_inline_bytes: 1024,
1971 submit_response_url:
1972 "http://manager.internal/v1/commands/command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1973 .to_string(),
1974 storage_upload_request: http_request(
1975 "http://manager.internal/v1/storage/response?signature=DoNotCanonicalize%2FUpload",
1976 PresignedOperation::Put,
1977 ),
1978 },
1979 );
1980
1981 CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
1982
1983 assert_eq!(
1984 envelope.response_handling.submit_response_url,
1985 "command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1986 );
1987 assert_eq!(
1988 envelope.response_handling.storage_upload_request.url(),
1989 "../storage/response?signature=DoNotCanonicalize%2FUpload"
1990 );
1991 let BodySpec::Storage {
1992 storage_get_request: Some(params),
1993 ..
1994 } = &envelope.params
1995 else {
1996 panic!("storage params request");
1997 };
1998 assert_eq!(params.url(), "../../storage/params?signature=params");
1999
2000 envelope.response_handling.storage_upload_request = http_request(
2001 "https://storage.example.com/result?signature=cloud",
2002 PresignedOperation::Put,
2003 );
2004 CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
2005 assert_eq!(
2006 envelope.response_handling.storage_upload_request.url(),
2007 "https://storage.example.com/result?signature=cloud",
2008 "cloud-presigned URLs must remain byte-for-byte absolute"
2009 );
2010 }
2011}
2012
2013#[cfg(test)]
2014mod idempotency_key_tests {
2015 use super::*;
2016 use crate::server::{validate_command_name, validate_command_target_id};
2017
2018 #[test]
2019 fn definite_dispatch_rejection_is_classified_by_typed_error_data() {
2020 let rejected = AlienError::new(ErrorData::TransportDispatchRejected {
2021 message: "not accepted".to_string(),
2022 transport_type: Some("http".to_string()),
2023 target: Some("command-id".to_string()),
2024 });
2025 assert!(is_definite_dispatch_rejection(&rejected));
2026
2027 let ambiguous = AlienError::new(ErrorData::TransportDispatchFailed {
2028 message: "acknowledgement lost".to_string(),
2029 transport_type: Some("http".to_string()),
2030 target: Some("command-id".to_string()),
2031 });
2032 assert!(!is_definite_dispatch_rejection(&ambiguous));
2033 }
2034
2035 #[test]
2044 fn target_id_colon_guard_prevents_idempotency_key_collision() {
2045 let colliding = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "k");
2046 assert_eq!(colliding, "dep:svc:a:b:k");
2047
2048 assert!(
2050 validate_command_target_id("svc:a").is_err(),
2051 "a ':'-bearing target id must be rejected so it cannot forge the rid segment"
2052 );
2053 assert!(validate_command_target_id("svc").is_ok());
2054 }
2055
2056 #[test]
2064 fn command_name_colon_guard_prevents_idempotency_key_collision() {
2065 let forged = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "c");
2067 let legitimate = CommandServer::compose_idempotency_key("dep", "svc", "a", "b:c");
2068 assert_eq!(
2069 forged, legitimate,
2070 "these inputs are exactly the colliding pair the guard must separate"
2071 );
2072 assert_eq!(legitimate, "dep:svc:a:b:c");
2073
2074 let err = validate_command_name("a:b").expect_err("':'-bearing command must be rejected");
2079 assert_eq!(err.code, "INVALID_COMMAND");
2080 assert!(validate_command_name("a").is_ok());
2081 }
2082}