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 CommandAccessContext, CommandEnvelopeData, CommandMetadata, CommandRegistry, CommandStatus,
51 InMemoryCommandRegistry, 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_command_access_context(
1039 &self,
1040 command_id: &str,
1041 ) -> Result<Option<CommandAccessContext>> {
1042 self.command_registry
1043 .get_command_access_context(command_id)
1044 .await
1045 }
1046
1047 pub async fn get_lease_owner(&self, lease_id: &str) -> Result<Option<(String, String)>> {
1055 let reverse_key = format!("lease:{}", lease_id);
1056 let Some(command_id_bytes) =
1057 self.kv
1058 .get(&reverse_key)
1059 .await
1060 .context(ErrorData::KvOperationFailed {
1061 operation: "get".to_string(),
1062 key: reverse_key.clone(),
1063 message: "Failed to look up lease reverse index".to_string(),
1064 })?
1065 else {
1066 return Ok(None);
1067 };
1068 let command_id = String::from_utf8(command_id_bytes).map_err(|_| {
1069 AlienError::new(ErrorData::Other {
1070 message: format!("Lease reverse index '{}' is not valid UTF-8", reverse_key),
1071 })
1072 })?;
1073
1074 let lease_key = format!("cmd:{}:lease", command_id);
1075 let Some(lease_data) =
1076 self.kv
1077 .get(&lease_key)
1078 .await
1079 .context(ErrorData::KvOperationFailed {
1080 operation: "get".to_string(),
1081 key: lease_key,
1082 message: "Failed to look up lease".to_string(),
1083 })?
1084 else {
1085 return Ok(None);
1086 };
1087 let lease: LeaseData = serde_json::from_slice(&lease_data)
1088 .into_alien_error()
1089 .context(ErrorData::SerializationFailed {
1090 message: "Failed to deserialize lease data".to_string(),
1091 data_type: Some("LeaseData".to_string()),
1092 })?;
1093 if lease.lease_id != lease_id {
1094 return Ok(None);
1095 }
1096
1097 Ok(Some((command_id, lease.owner)))
1098 }
1099
1100 pub async fn release_lease_by_id(&self, lease_id: &str) -> Result<()> {
1102 match self.get_lease_owner(lease_id).await? {
1103 Some((command_id, _owner)) => self.release_lease(&command_id, lease_id).await,
1104 None => Err(AlienError::new(ErrorData::LeaseNotFound {
1105 lease_id: lease_id.to_string(),
1106 })),
1107 }
1108 }
1109
1110 async fn validate_create_command(&self, request: &CreateCommandRequest) -> Result<()> {
1115 if request.command.is_empty() {
1116 return Err(AlienError::new(ErrorData::InvalidCommand {
1117 message: "Command name cannot be empty".to_string(),
1118 }));
1119 }
1120
1121 validate_command_name(&request.command)?;
1129
1130 if request.deployment_id.is_empty() {
1131 return Err(AlienError::new(ErrorData::InvalidCommand {
1132 message: "Deployment ID cannot be empty".to_string(),
1133 }));
1134 }
1135
1136 if let Some(deadline) = request.deadline {
1137 if deadline <= Utc::now() {
1138 return Err(AlienError::new(ErrorData::InvalidCommand {
1139 message: "Deadline must be in the future".to_string(),
1140 }));
1141 }
1142 }
1143
1144 Ok(())
1145 }
1146
1147 fn compose_idempotency_key(
1155 deployment_id: &str,
1156 target_resource_id: &str,
1157 command_name: &str,
1158 idem_key: &str,
1159 ) -> String {
1160 debug_assert!(
1170 !target_resource_id.contains(':'),
1171 "target_resource_id must be ':'-free before key composition: {target_resource_id}"
1172 );
1173 format!(
1174 "{}:{}:{}:{}",
1175 deployment_id, target_resource_id, command_name, idem_key
1176 )
1177 }
1178
1179 async fn check_idempotency(&self, idem_key: &str) -> Result<Option<String>> {
1180 let key = format!("idem:{}", idem_key);
1181 if let Some(data) = self
1182 .kv
1183 .get(&key)
1184 .await
1185 .context(ErrorData::KvOperationFailed {
1186 operation: "get".to_string(),
1187 key: key.clone(),
1188 message: "Failed to check idempotency".to_string(),
1189 })?
1190 {
1191 let command_id = String::from_utf8(data).into_alien_error().context(
1192 ErrorData::SerializationFailed {
1193 message: "Invalid idempotency data".to_string(),
1194 data_type: Some("String".to_string()),
1195 },
1196 )?;
1197 return Ok(Some(command_id));
1198 }
1199 Ok(None)
1200 }
1201
1202 async fn store_idempotency(&self, idem_key: &str, command_id: &str) -> Result<Option<String>> {
1210 let key = format!("idem:{}", idem_key);
1211 let ttl = Duration::from_secs(24 * 60 * 60); let won = self
1213 .kv
1214 .put(
1215 &key,
1216 command_id.as_bytes().to_vec(),
1217 Some(PutOptions {
1218 ttl: Some(ttl),
1219 if_not_exists: true,
1220 }),
1221 )
1222 .await
1223 .context(ErrorData::KvOperationFailed {
1224 operation: "put".to_string(),
1225 key: key.clone(),
1226 message: "Failed to store idempotency".to_string(),
1227 })?;
1228 if won {
1229 return Ok(None);
1230 }
1231 match self.check_idempotency(idem_key).await? {
1235 Some(winner_id) => Ok(Some(winner_id)),
1236 None => Err(AlienError::new(ErrorData::Other {
1237 message: format!(
1238 "Idempotency key '{}' was concurrently claimed but has no winner entry",
1239 key
1240 ),
1241 })),
1242 }
1243 }
1244
1245 pub async fn store_params(&self, command_id: &str, params: &BodySpec) -> Result<()> {
1248 let key = format!("cmd:{}:params", command_id);
1249
1250 let data = CommandParamsData {
1252 params: params.clone(),
1253 };
1254 let value = serde_json::to_vec(&data).into_alien_error().context(
1255 ErrorData::SerializationFailed {
1256 message: "Failed to serialize params".to_string(),
1257 data_type: Some("CommandParamsData".to_string()),
1258 },
1259 )?;
1260
1261 if value.len() <= KV_VALUE_THRESHOLD {
1263 self.kv
1264 .put(&key, value, None)
1265 .await
1266 .context(ErrorData::KvOperationFailed {
1267 operation: "put".to_string(),
1268 key: key.clone(),
1269 message: "Failed to store params".to_string(),
1270 })?;
1271 return Ok(());
1272 }
1273
1274 if let BodySpec::Inline { inline_base64 } = params {
1276 let raw_bytes = general_purpose::STANDARD
1277 .decode(inline_base64)
1278 .into_alien_error()
1279 .context(ErrorData::SerializationFailed {
1280 message: "Failed to decode inline base64 params for auto-promotion".to_string(),
1281 data_type: Some("base64".to_string()),
1282 })?;
1283
1284 let raw_len = raw_bytes.len() as u64;
1285 let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1286
1287 self.storage
1288 .put(&blob_path, Bytes::from(raw_bytes).into())
1289 .await
1290 .into_alien_error()
1291 .context(ErrorData::StorageOperationFailed {
1292 message: "Failed to auto-promote params to blob storage".to_string(),
1293 operation: Some("put".to_string()),
1294 path: Some(blob_path.to_string()),
1295 })?;
1296
1297 debug!(
1298 "Auto-promoted params for command {} to blob ({} bytes raw)",
1299 command_id, raw_len
1300 );
1301
1302 let promoted = CommandParamsData {
1304 params: BodySpec::Storage {
1305 size: Some(raw_len),
1306 storage_get_request: None,
1307 storage_put_used: Some(true),
1308 },
1309 };
1310 let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1311 ErrorData::SerializationFailed {
1312 message: "Failed to serialize promoted params reference".to_string(),
1313 data_type: Some("CommandParamsData".to_string()),
1314 },
1315 )?;
1316 self.kv.put(&key, promoted_value, None).await.context(
1317 ErrorData::KvOperationFailed {
1318 operation: "put".to_string(),
1319 key: key.clone(),
1320 message: "Failed to store promoted params reference".to_string(),
1321 },
1322 )?;
1323 return Ok(());
1324 }
1325
1326 self.kv
1328 .put(&key, value, None)
1329 .await
1330 .context(ErrorData::KvOperationFailed {
1331 operation: "put".to_string(),
1332 key: key.clone(),
1333 message: "Failed to store params".to_string(),
1334 })?;
1335 Ok(())
1336 }
1337
1338 pub async fn get_params(&self, command_id: &str) -> Result<Option<BodySpec>> {
1339 let key = format!("cmd:{}:params", command_id);
1340 if let Some(value) = self
1341 .kv
1342 .get(&key)
1343 .await
1344 .context(ErrorData::KvOperationFailed {
1345 operation: "get".to_string(),
1346 key: key.clone(),
1347 message: "Failed to get params".to_string(),
1348 })?
1349 {
1350 let data: CommandParamsData = serde_json::from_slice(&value)
1351 .into_alien_error()
1352 .context(ErrorData::SerializationFailed {
1353 message: "Failed to deserialize params".to_string(),
1354 data_type: Some("CommandParamsData".to_string()),
1355 })?;
1356 return Ok(Some(data.params));
1357 }
1358 Ok(None)
1359 }
1360
1361 pub async fn store_response(&self, command_id: &str, response: &CommandResponse) -> Result<()> {
1364 let key = format!("cmd:{}:response", command_id);
1365 let data = CommandResponseData {
1366 response: response.clone(),
1367 };
1368 let value = serde_json::to_vec(&data).into_alien_error().context(
1369 ErrorData::SerializationFailed {
1370 message: "Failed to serialize response".to_string(),
1371 data_type: Some("CommandResponseData".to_string()),
1372 },
1373 )?;
1374
1375 if value.len() <= KV_VALUE_THRESHOLD {
1377 self.kv
1378 .put(&key, value, None)
1379 .await
1380 .context(ErrorData::KvOperationFailed {
1381 operation: "put".to_string(),
1382 key: key.clone(),
1383 message: "Failed to store response".to_string(),
1384 })?;
1385 return Ok(());
1386 }
1387
1388 if let CommandResponse::Success {
1390 response: BodySpec::Inline { inline_base64 },
1391 } = response
1392 {
1393 let raw_bytes = general_purpose::STANDARD
1394 .decode(inline_base64)
1395 .into_alien_error()
1396 .context(ErrorData::SerializationFailed {
1397 message: "Failed to decode inline base64 response for auto-promotion"
1398 .to_string(),
1399 data_type: Some("base64".to_string()),
1400 })?;
1401
1402 let raw_len = raw_bytes.len() as u64;
1403 let blob_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1404
1405 self.storage
1406 .put(&blob_path, Bytes::from(raw_bytes).into())
1407 .await
1408 .into_alien_error()
1409 .context(ErrorData::StorageOperationFailed {
1410 message: "Failed to auto-promote response to blob storage".to_string(),
1411 operation: Some("put".to_string()),
1412 path: Some(blob_path.to_string()),
1413 })?;
1414
1415 let get_request = self
1417 .generate_response_storage_get_request(command_id)
1418 .await?;
1419
1420 debug!(
1421 "Auto-promoted response for command {} to blob ({} bytes raw)",
1422 command_id, raw_len
1423 );
1424
1425 let promoted = CommandResponseData {
1427 response: CommandResponse::Success {
1428 response: BodySpec::Storage {
1429 size: Some(raw_len),
1430 storage_get_request: Some(get_request),
1431 storage_put_used: Some(true),
1432 },
1433 },
1434 };
1435 let promoted_value = serde_json::to_vec(&promoted).into_alien_error().context(
1436 ErrorData::SerializationFailed {
1437 message: "Failed to serialize promoted response reference".to_string(),
1438 data_type: Some("CommandResponseData".to_string()),
1439 },
1440 )?;
1441 self.kv.put(&key, promoted_value, None).await.context(
1442 ErrorData::KvOperationFailed {
1443 operation: "put".to_string(),
1444 key: key.clone(),
1445 message: "Failed to store promoted response reference".to_string(),
1446 },
1447 )?;
1448 return Ok(());
1449 }
1450
1451 self.kv
1453 .put(&key, value, None)
1454 .await
1455 .context(ErrorData::KvOperationFailed {
1456 operation: "put".to_string(),
1457 key: key.clone(),
1458 message: "Failed to store response".to_string(),
1459 })?;
1460 Ok(())
1461 }
1462
1463 pub async fn get_response(&self, command_id: &str) -> Result<Option<CommandResponse>> {
1464 let key = format!("cmd:{}:response", command_id);
1465 if let Some(value) = self
1466 .kv
1467 .get(&key)
1468 .await
1469 .context(ErrorData::KvOperationFailed {
1470 operation: "get".to_string(),
1471 key: key.clone(),
1472 message: "Failed to get response".to_string(),
1473 })?
1474 {
1475 let data: CommandResponseData = serde_json::from_slice(&value)
1476 .into_alien_error()
1477 .context(ErrorData::SerializationFailed {
1478 message: "Failed to deserialize response".to_string(),
1479 data_type: Some("CommandResponseData".to_string()),
1480 })?;
1481 return Ok(Some(data.response));
1482 }
1483 Ok(None)
1484 }
1485
1486 async fn create_pending_index(
1489 &self,
1490 deployment_id: &str,
1491 target_resource_id: &str,
1492 command_id: &str,
1493 ) -> Result<()> {
1494 debug_assert!(
1498 !target_resource_id.contains(':'),
1499 "target_resource_id must be ':'-free in the pending index: {target_resource_id}"
1500 );
1501 let timestamp = Utc::now().timestamp_nanos_opt().unwrap_or(0);
1502 let key = format!(
1503 "target:{}:{}:pending:{}:{}",
1504 deployment_id, target_resource_id, timestamp, command_id
1505 );
1506
1507 self.kv
1509 .put(&key, vec![], None)
1510 .await
1511 .context(ErrorData::KvOperationFailed {
1512 operation: "put".to_string(),
1513 key: key.clone(),
1514 message: "Failed to create pending index".to_string(),
1515 })?;
1516 Ok(())
1517 }
1518
1519 async fn delete_pending_index(
1520 &self,
1521 deployment_id: &str,
1522 target_resource_id: &str,
1523 command_id: &str,
1524 ) -> Result<()> {
1525 let prefix = format!("target:{}:{}:pending:", deployment_id, target_resource_id);
1527 let scan_result = self
1528 .kv
1529 .scan_prefix(&prefix, Some(100), None)
1530 .await
1531 .into_alien_error()
1532 .context(ErrorData::KvOperationFailed {
1533 operation: "scan_prefix".to_string(),
1534 key: prefix.clone(),
1535 message: "Failed to scan pending index".to_string(),
1536 })?;
1537
1538 for (key, _) in scan_result.items {
1539 if key.ends_with(&format!(":{}", command_id)) {
1540 let _ = self.kv.delete(&key).await;
1541 break;
1542 }
1543 }
1544 Ok(())
1545 }
1546
1547 async fn delete_lease(&self, command_id: &str) -> Result<()> {
1550 let key = format!("cmd:{}:lease", command_id);
1551 let _ = self.kv.delete(&key).await;
1552 Ok(())
1553 }
1554
1555 pub async fn reap_expired_commands(&self) -> Result<u32> {
1567 let now = Utc::now();
1568 let mut expired = 0u32;
1569 let mut cursor: Option<String> = None;
1570 for _ in 0..64 {
1575 let scan = self
1576 .kv
1577 .scan_prefix("deadline:", Some(256), cursor.clone())
1578 .await
1579 .into_alien_error()
1580 .context(ErrorData::KvOperationFailed {
1581 operation: "scan_prefix".to_string(),
1582 key: "deadline:".to_string(),
1583 message: "Failed to scan the deadline index".to_string(),
1584 })?;
1585 let next_cursor = scan.next_cursor.clone();
1586 for (key, value) in scan.items {
1587 let Ok(data) = serde_json::from_slice::<DeadlineIndexData>(&value) else {
1588 warn!(key = %key, "Unparseable deadline index entry; deleting");
1589 let _ = self.kv.delete(&key).await;
1590 continue;
1591 };
1592 if data.deadline > now {
1593 continue;
1596 }
1597
1598 let status = self
1599 .command_registry
1600 .get_command_status(&data.command_id)
1601 .await?;
1602 match status {
1603 None => {
1604 let _ = self.kv.delete(&key).await;
1605 }
1606 Some(status) if status.state.is_terminal() => {
1607 let _ = self.kv.delete(&key).await;
1608 }
1609 Some(status) => {
1610 let won = self
1611 .command_registry
1612 .complete_command(
1613 &data.command_id,
1614 CommandState::Expired,
1615 now,
1616 None,
1617 Some(serde_json::json!({
1618 "code": "COMMAND_EXPIRED",
1619 "message": format!("Deadline {} elapsed", data.deadline.to_rfc3339()),
1620 })),
1621 )
1622 .await?;
1623 if won {
1624 expired += 1;
1625 info!(command_id = %data.command_id, "Expired overdue command");
1626 let _ = self.delete_lease(&data.command_id).await;
1627 let _ = self
1628 .delete_pending_index(
1629 &status.deployment_id,
1630 &status.target.resource_id,
1631 &data.command_id,
1632 )
1633 .await;
1634 }
1635 let _ = self.kv.delete(&key).await;
1636 }
1637 }
1638 }
1639 match next_cursor {
1640 Some(next) => cursor = Some(next),
1641 None => break,
1642 }
1643 }
1644 Ok(expired)
1645 }
1646
1647 async fn create_deadline_index(&self, command_id: &str, deadline: DateTime<Utc>) -> Result<()> {
1650 let key = format!(
1651 "deadline:{}:{}",
1652 deadline.timestamp_nanos_opt().unwrap_or(0),
1653 command_id
1654 );
1655
1656 let data = DeadlineIndexData {
1657 command_id: command_id.to_string(),
1658 deadline,
1659 };
1660 let value = serde_json::to_vec(&data).into_alien_error().context(
1661 ErrorData::SerializationFailed {
1662 message: "Failed to serialize deadline index".to_string(),
1663 data_type: Some("DeadlineIndexData".to_string()),
1664 },
1665 )?;
1666
1667 const DEADLINE_INDEX_GRACE: chrono::Duration = chrono::Duration::days(7);
1674 let ttl = deadline
1675 .signed_duration_since(Utc::now())
1676 .checked_add(&DEADLINE_INDEX_GRACE)
1677 .unwrap_or(DEADLINE_INDEX_GRACE);
1678 let options = (ttl.num_seconds() > 0).then(|| PutOptions {
1679 ttl: Some(Duration::from_secs(ttl.num_seconds() as u64)),
1680 if_not_exists: false,
1681 });
1682
1683 self.kv
1684 .put(&key, value, options)
1685 .await
1686 .context(ErrorData::KvOperationFailed {
1687 operation: "put".to_string(),
1688 key: key.clone(),
1689 message: "Failed to create deadline index".to_string(),
1690 })?;
1691 Ok(())
1692 }
1693
1694 async fn dispatch_command_push(
1697 &self,
1698 command_id: &str,
1699 deployment_id: &str,
1700 ) -> Result<CommandState> {
1701 let metadata = self
1703 .command_registry
1704 .get_command_metadata(command_id)
1705 .await?
1706 .ok_or_else(|| {
1707 AlienError::new(ErrorData::CommandNotFound {
1708 command_id: command_id.to_string(),
1709 })
1710 })?;
1711
1712 let params = self.get_params(command_id).await?.ok_or_else(|| {
1714 AlienError::new(ErrorData::CommandNotFound {
1715 command_id: command_id.to_string(),
1716 })
1717 })?;
1718
1719 let envelope = self.build_envelope(command_id, &metadata, params).await?;
1721
1722 if !self
1728 .command_registry
1729 .mark_dispatched_if_not_terminal(command_id, Utc::now())
1730 .await?
1731 {
1732 return Ok(self
1733 .command_registry
1734 .get_command_status(command_id)
1735 .await?
1736 .map(|status| status.state)
1737 .unwrap_or(CommandState::Dispatched));
1738 }
1739
1740 if let Err(error) = self.command_dispatcher.dispatch(&envelope).await {
1748 if is_definite_dispatch_rejection(&error) {
1749 let delivery_failure = CommandResponse::error(
1750 "DELIVERY_FAILED",
1751 "Worker runtime did not accept command delivery",
1752 );
1753 self.submit_command_response(command_id, delivery_failure)
1754 .await?;
1755 warn!(
1756 command_id,
1757 deployment_id,
1758 error = %error,
1759 "Push dispatch was definitely rejected; command marked Failed"
1760 );
1761 return Ok(CommandState::Failed);
1762 }
1763
1764 let error = error.context(ErrorData::TransportDispatchFailed {
1765 message: "Failed to dispatch command".to_string(),
1766 transport_type: None,
1767 target: Some(deployment_id.to_string()),
1768 });
1769 warn!(
1770 command_id,
1771 deployment_id,
1772 error = %error,
1773 "Push dispatch acknowledgement failed; command remains Dispatched for a possible late response"
1774 );
1775 return Ok(CommandState::Dispatched);
1776 }
1777
1778 info!("Command {} dispatched via push", envelope.command_id);
1779 Ok(CommandState::Dispatched)
1780 }
1781
1782 async fn build_envelope(
1783 &self,
1784 command_id: &str,
1785 metadata: &CommandEnvelopeData,
1786 mut params: BodySpec,
1787 ) -> Result<Envelope> {
1788 let response_handling = self.create_response_handling(command_id).await?;
1789
1790 if let BodySpec::Storage { size, .. } = ¶ms {
1793 let raw_size = size.unwrap_or(0) as usize;
1794 if raw_size > 0 && raw_size <= self.inline_max_bytes {
1795 let blob_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1796 match self.storage.get(&blob_path).await {
1797 Ok(get_result) => match get_result.bytes().await {
1798 Ok(raw_bytes) => {
1799 params = BodySpec::inline(&raw_bytes);
1800 debug!(
1801 "Re-inlined params for command {} ({} bytes) into envelope",
1802 command_id, raw_size
1803 );
1804 }
1805 Err(e) => {
1806 debug!(
1807 "Failed to read blob bytes for re-inline (command {}), falling back to presigned URL: {}",
1808 command_id, e
1809 );
1810 }
1811 },
1812 Err(e) => {
1813 debug!(
1814 "Failed to read blob for re-inline (command {}), falling back to presigned URL: {}",
1815 command_id, e
1816 );
1817 }
1818 }
1819 }
1820 }
1821
1822 if let BodySpec::Storage {
1827 size,
1828 storage_get_request: _,
1829 storage_put_used,
1830 } = ¶ms
1831 {
1832 let get_request = self.generate_storage_get_request(command_id).await?;
1833 params = BodySpec::Storage {
1834 size: *size,
1835 storage_get_request: Some(get_request),
1836 storage_put_used: *storage_put_used,
1837 };
1838 }
1839
1840 Ok(Envelope::new(
1841 metadata.deployment_id.clone(),
1842 metadata.target.clone(),
1843 command_id.to_string(),
1844 metadata.attempt,
1845 metadata.deadline,
1846 metadata.command.clone(),
1847 params,
1848 response_handling,
1849 ))
1850 }
1851
1852 async fn create_response_handling(&self, command_id: &str) -> Result<ResponseHandling> {
1853 let upload_path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1854 let expires_in = Duration::from_secs(Self::RESPONSE_CREDENTIAL_LIFETIME_SECS);
1855 let presigned = self
1856 .storage
1857 .presigned_put(&upload_path, expires_in)
1858 .await
1859 .context(ErrorData::StorageOperationFailed {
1860 message: "Failed to create response upload URL".to_string(),
1861 operation: Some("presigned_put".to_string()),
1862 path: Some(upload_path.to_string()),
1863 })?;
1864
1865 let (response_token, expires) = self.sign_response_url(command_id);
1866
1867 Ok(ResponseHandling {
1868 max_inline_bytes: self.inline_max_bytes as u64,
1869 submit_response_url: format!(
1870 "{}/commands/{}/response?response_token={}&expires={}",
1871 self.base_url.trim_end_matches('/'),
1872 command_id,
1873 response_token,
1874 expires,
1875 ),
1876 storage_upload_request: presigned,
1877 })
1878 }
1879
1880 async fn generate_params_upload(&self, command_id: &str) -> Result<StorageUpload> {
1881 let upload_path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1882 let expires_in = Duration::from_secs(3600);
1883 let presigned = self
1884 .storage
1885 .presigned_put(&upload_path, expires_in)
1886 .await
1887 .into_alien_error()
1888 .context(ErrorData::StorageOperationFailed {
1889 message: "Failed to create presigned URL".to_string(),
1890 operation: Some("presigned_put".to_string()),
1891 path: Some(upload_path.to_string()),
1892 })?;
1893
1894 Ok(StorageUpload {
1895 put_request: presigned.clone(),
1896 expires_at: presigned.expiration,
1897 })
1898 }
1899
1900 async fn generate_storage_get_request(&self, command_id: &str) -> Result<PresignedRequest> {
1901 let path = StoragePath::from(format!("arc/commands/{}/params", command_id));
1902 let expires_in = Duration::from_secs(3600);
1903 self.storage.presigned_get(&path, expires_in).await.context(
1904 ErrorData::StorageOperationFailed {
1905 message: "Failed to create storage get request".to_string(),
1906 operation: Some("presigned_get".to_string()),
1907 path: Some(path.to_string()),
1908 },
1909 )
1910 }
1911
1912 async fn generate_response_storage_get_request(
1913 &self,
1914 command_id: &str,
1915 ) -> Result<PresignedRequest> {
1916 let path = StoragePath::from(format!("arc/commands/{}/response", command_id));
1917 let expires_in = Duration::from_secs(3600);
1918 self.storage.presigned_get(&path, expires_in).await.context(
1919 ErrorData::StorageOperationFailed {
1920 message: "Failed to create response storage get request".to_string(),
1921 operation: Some("presigned_get".to_string()),
1922 path: Some(path.to_string()),
1923 },
1924 )
1925 }
1926
1927 fn extract_command_id_from_index_key(&self, index_key: &str) -> Result<String> {
1928 index_key
1929 .split(':')
1930 .last()
1931 .ok_or_else(|| {
1932 AlienError::new(ErrorData::Other {
1933 message: format!("Invalid index key format: {}", index_key),
1934 })
1935 })
1936 .map(|s| s.to_string())
1937 }
1938}
1939
1940#[cfg(test)]
1941mod relative_url_tests {
1942 use super::*;
1943 use alien_core::presigned::PresignedOperation;
1944 use std::collections::HashMap;
1945
1946 fn http_request(url: &str, operation: PresignedOperation) -> PresignedRequest {
1947 PresignedRequest::new_http(
1948 url.to_string(),
1949 match operation {
1950 PresignedOperation::Get => "GET",
1951 PresignedOperation::Put => "PUT",
1952 PresignedOperation::Delete => "DELETE",
1953 }
1954 .to_string(),
1955 HashMap::new(),
1956 operation,
1957 "commands/test".to_string(),
1958 Utc::now() + chrono::Duration::minutes(5),
1959 )
1960 }
1961
1962 #[test]
1963 fn leased_manager_urls_are_relative_to_lease_endpoint() {
1964 let mut envelope = Envelope::new(
1965 "deployment",
1966 CommandTarget::new("daemon", CommandTargetType::Daemon),
1967 "command",
1968 1,
1969 None,
1970 "run",
1971 BodySpec::Storage {
1972 size: Some(2048),
1973 storage_get_request: Some(http_request(
1974 "http://manager.internal/storage/params?signature=params",
1975 PresignedOperation::Get,
1976 )),
1977 storage_put_used: Some(true),
1978 },
1979 ResponseHandling {
1980 max_inline_bytes: 1024,
1981 submit_response_url:
1982 "http://manager.internal/v1/commands/command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1983 .to_string(),
1984 storage_upload_request: http_request(
1985 "http://manager.internal/v1/storage/response?signature=DoNotCanonicalize%2FUpload",
1986 PresignedOperation::Put,
1987 ),
1988 },
1989 );
1990
1991 CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
1992
1993 assert_eq!(
1994 envelope.response_handling.submit_response_url,
1995 "command/response?response_token=DoNotCanonicalize%2FValue&expires=1"
1996 );
1997 assert_eq!(
1998 envelope.response_handling.storage_upload_request.url(),
1999 "../storage/response?signature=DoNotCanonicalize%2FUpload"
2000 );
2001 let BodySpec::Storage {
2002 storage_get_request: Some(params),
2003 ..
2004 } = &envelope.params
2005 else {
2006 panic!("storage params request");
2007 };
2008 assert_eq!(params.url(), "../../storage/params?signature=params");
2009
2010 envelope.response_handling.storage_upload_request = http_request(
2011 "https://storage.example.com/result?signature=cloud",
2012 PresignedOperation::Put,
2013 );
2014 CommandServer::relativize_manager_urls(&mut envelope, "http://manager.internal/v1");
2015 assert_eq!(
2016 envelope.response_handling.storage_upload_request.url(),
2017 "https://storage.example.com/result?signature=cloud",
2018 "cloud-presigned URLs must remain byte-for-byte absolute"
2019 );
2020 }
2021}
2022
2023#[cfg(test)]
2024mod idempotency_key_tests {
2025 use super::*;
2026 use crate::server::{validate_command_name, validate_command_target_id};
2027
2028 #[test]
2029 fn definite_dispatch_rejection_is_classified_by_typed_error_data() {
2030 let rejected = AlienError::new(ErrorData::TransportDispatchRejected {
2031 message: "not accepted".to_string(),
2032 transport_type: Some("http".to_string()),
2033 target: Some("command-id".to_string()),
2034 });
2035 assert!(is_definite_dispatch_rejection(&rejected));
2036
2037 let ambiguous = AlienError::new(ErrorData::TransportDispatchFailed {
2038 message: "acknowledgement lost".to_string(),
2039 transport_type: Some("http".to_string()),
2040 target: Some("command-id".to_string()),
2041 });
2042 assert!(!is_definite_dispatch_rejection(&ambiguous));
2043 }
2044
2045 #[test]
2054 fn target_id_colon_guard_prevents_idempotency_key_collision() {
2055 let colliding = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "k");
2056 assert_eq!(colliding, "dep:svc:a:b:k");
2057
2058 assert!(
2060 validate_command_target_id("svc:a").is_err(),
2061 "a ':'-bearing target id must be rejected so it cannot forge the rid segment"
2062 );
2063 assert!(validate_command_target_id("svc").is_ok());
2064 }
2065
2066 #[test]
2074 fn command_name_colon_guard_prevents_idempotency_key_collision() {
2075 let forged = CommandServer::compose_idempotency_key("dep", "svc", "a:b", "c");
2077 let legitimate = CommandServer::compose_idempotency_key("dep", "svc", "a", "b:c");
2078 assert_eq!(
2079 forged, legitimate,
2080 "these inputs are exactly the colliding pair the guard must separate"
2081 );
2082 assert_eq!(legitimate, "dep:svc:a:b:c");
2083
2084 let err = validate_command_name("a:b").expect_err("':'-bearing command must be rejected");
2089 assert_eq!(err.code, "INVALID_COMMAND");
2090 assert!(validate_command_name("a").is_ok());
2091 }
2092}