Queued DM overwrite, storage bucket ownership

config/deny.toml

@@ -141,6 +141,14 @@ ignore = [
   # for one is not a plan. All four are npm-ecosystem ids that cargo-deny reports unknown; osv-scanner honours
   # them.
   { id = "GHSA-rgj7-g3m4-5g8c", reason = "expires 2026-10-15" },
+  # unsafe-libyaml 0.2.11 (unmaintained; repository archived 2024-03-25).
+  # The c2rust-translated libyaml that serde_yaml (itself deprecated) parses with, reached from this workspace
+  # in two places: the agent registry's save/load in ws_server, and operator-supplied cluster YAML in the CLI.
+  # Both read local files from a trusted machine, never attacker-controlled network input. INFO-unmaintained,
+  # not a vulnerability, and no patched version exists because there is nothing to patch -- it is already
+  # listed in the workspace unmaintained metadata; this is its cargo-deny counterpart. Cleared only by
+  # dropping serde_yaml workspace-wide for a safe-Rust parser (saphyr / serde-saphyr).
+  { id = "RUSTSEC-2024-0449", reason = "expires 2026-10-15" },
 ]
 
 [licenses]

libs/edge-toolkit/src/ws_server.rs

@@ -243,6 +243,9 @@ impl<S: Clone + Send + 'static> AgentRegistry<S> {
     ///
     /// Returns `None` when `to_agent_id` is not in the registry. The inner `Option<S>` is the recipient's live session
     /// -- `Some` when connected, `None` when the message was queued for a disconnected agent.
+    ///
+    /// Entries are keyed by message id, so a second message queued from the same sender waits alongside the
+    /// first instead of replacing it. Message ids are UUIDv7, which keeps the map in arrival order.
     #[must_use]
     pub fn queue_direct(
         &self,
@@ -262,9 +265,11 @@ impl<S: Clone + Send + 'static> AgentRegistry<S> {
             message,
         };
         let session = recipient.session.clone();
+        // Keyed by message id: keying by sender let a second queued message from
+        // one sender silently replace the first, losing it forever.
         let _previous: Option<PendingDirectMessage> = recipient
             .pending_direct_messages
-            .insert(from_agent_id.to_string(), pending.clone());
+            .insert(pending.message_id.clone(), pending.clone());
         drop(agents);
 
         Some((pending, session))
@@ -294,18 +299,20 @@ impl<S: Clone + Send + 'static> AgentRegistry<S> {
             .get_mut(recipient_agent_id)
             .ok_or_else(|| AcknowledgeError::UnknownAgent(recipient_agent_id.to_string()))?;
 
-        let sender_agent_id = recipient
+        // Looked up by the id carried on the message rather than by map key: the key has been the message id
+        // since queued messages stopped being keyed by sender, and entries persisted before that change are
+        // still keyed by sender id. Matching the value acknowledges either shape.
+        let key = recipient
             .pending_direct_messages
             .iter()
-            .find_map(|(id, pending)| (pending.message_id == message_id).then(|| id.clone()))
+            .find(|(_, pending)| pending.message_id == message_id)
+            .map(|(key, _)| key.clone())
             .ok_or(AcknowledgeError::NoPendingMessage)?;
-        // The `find_map` above drops its iterator before we re-borrow `pending_direct_messages` mutably for the
-        // removal. The double lookup costs O(log n) but lets us share the `NoPendingMessage` error with the find-side
-        // case instead of asserting an invariant.
         let pending = recipient
             .pending_direct_messages
-            .remove(&sender_agent_id)
+            .remove(&key)
             .ok_or(AcknowledgeError::NoPendingMessage)?;
+        let sender_agent_id = pending.from_agent_id.clone();
         let sender_session = agents.get(&sender_agent_id).and_then(|record| record.session.clone());
         drop(agents);
 

libs/edge-toolkit/tests/registry.rs

@@ -169,3 +169,85 @@ fn agent_record_with_pending_builder_replaces_the_map() {
         .with_pending_direct_messages(BTreeMap::new());
     assert!(record.pending_direct_messages.is_empty());
 }
+
+#[test]
+fn queued_direct_messages_from_one_sender_do_not_replace_each_other() {
+    let registry = AgentRegistry::<String>::default();
+    let (_recipient, _status) = registry.connect_agent(None, "recipient".to_string(), "10.0.0.1", "sess-r".to_string());
+    let (_sender, _status) = registry.connect_agent(None, "sender".to_string(), "10.0.0.2", "sess-s".to_string());
+
+    let (first, recipient_session) = registry
+        .queue_direct(
+            "msg-1".to_string(),
+            "sender",
+            "recipient",
+            "t1".to_string(),
+            serde_json::json!({"n": 1}),
+        )
+        .expect("first message queues");
+    assert_eq!(first.message_id, "msg-1");
+    assert!(recipient_session.is_some(), "recipient is connected");
+
+    let (_second, _session) = registry
+        .queue_direct(
+            "msg-2".to_string(),
+            "sender",
+            "recipient",
+            "t2".to_string(),
+            serde_json::json!({"n": 2}),
+        )
+        .expect("second message queues");
+
+    // Both messages wait: the second must not have replaced the first.
+    let pending = registry.pending_messages_for("recipient");
+    let ids = pending.iter().map(|p| p.message_id.as_str()).collect::<Vec<_>>();
+    assert_eq!(ids, vec!["msg-1", "msg-2"], "both queued, in arrival order");
+
+    // Acknowledging the first leaves the second queued for the recipient.
+    let (acked, _session, sender) = registry
+        .acknowledge_message("recipient", "msg-1")
+        .expect("acknowledged");
+    assert_eq!(acked, "msg-1");
+    assert_eq!(sender, "sender");
+
+    let pending = registry.pending_messages_for("recipient");
+    assert_eq!(pending.len(), 1, "only the acknowledged message was removed");
+    assert_eq!(pending[0].message_id, "msg-2");
+    assert_eq!(pending[0].from_agent_id, "sender");
+}
+
+#[test]
+fn entries_persisted_with_the_legacy_sender_key_still_acknowledge() {
+    use edge_toolkit::ws_server::PendingDirectMessage;
+
+    let registry = AgentRegistry::<String>::default();
+    let (_recipient, _status) = registry.connect_agent(None, "recipient".to_string(), "10.0.0.1", "sess-r".to_string());
+    let (_sender, _status) = registry.connect_agent(None, "sender".to_string(), "10.0.0.2", "sess-s".to_string());
+
+    // A registry saved before queued messages were keyed by message id holds
+    // sender-keyed entries; acknowledging one must still work.
+    // Deserialised rather than built: the type is #[non_exhaustive] outside
+    // the crate, and the persisted shape is exactly what load() hands back.
+    let legacy_entry: PendingDirectMessage = serde_json::from_value(serde_json::json!({
+        "message_id": "legacy-msg",
+        "from_agent_id": "sender",
+        "server_received_at": "t0",
+        "message": {"n": 0},
+    }))
+    .expect("the persisted shape deserialises");
+    let legacy = BTreeMap::from([("sender".to_string(), legacy_entry)]);
+    let record = edge_toolkit::ws_server::AgentRecord::<String>::new(
+        AgentConnectionState::Connected,
+        Some("10.0.0.1".to_string()),
+        None,
+    )
+    .with_pending_direct_messages(legacy);
+    let registry = AgentRegistry::from_agents(BTreeMap::from([("recipient".to_string(), record)]));
+
+    let (acked, _session, sender) = registry
+        .acknowledge_message("recipient", "legacy-msg")
+        .expect("legacy entry acknowledges");
+    assert_eq!(acked, "legacy-msg");
+    assert_eq!(sender, "sender");
+    assert!(registry.pending_messages_for("recipient").is_empty());
+}

services/storage/src/lib.rs

@@ -92,6 +92,10 @@ pub enum StorageError {
     #[response(status = 404, reason = "NOT_FOUND")]
     AgentNotFound,
 
+    #[error("requester does not own this bucket")]
+    #[response(status = 403, reason = "FORBIDDEN")]
+    RequesterNotBucketOwner,
+
     #[error("no such object")]
     #[response(status = 404, reason = "NOT_FOUND")]
     ObjectNotFound,

services/storage/src/routes.rs

@@ -123,9 +123,18 @@ where
 {
     let (agent_id, filename) = agent_and_filename(&req)?;
 
-    if !registry.agents.lock()?.contains_key(&agent_id) {
+    let agents = registry.agents.lock()?;
+    let Some(record) = agents.get(&agent_id) else {
         return Err(StorageError::AgentNotFound);
+    };
+    // An id's existence is not ownership: the listener is reachable from every
+    // interface, so a PUT is accepted only from the address the agent
+    // registered from, or from loopback where the server's own runners and
+    // modules write on an agent's behalf.
+    if !requester_owns_bucket(record.last_known_ip.as_deref(), req.peer_addr()) {
+        return Err(StorageError::RequesterNotBucketOwner);
     }
+    drop(agents);
 
     // The body is buffered rather than streamed to the store.
     // `ObjectStore::put` takes a whole payload, and the alternative (`put_multipart`) buys streaming at the
@@ -149,6 +158,22 @@ where
     Ok(response.finish())
 }
 
+/// Whether the request peer may write to the bucket of the agent whose registered address is `registered_ip`.
+///
+/// The registered address is the identity: `connect_agent` records the IP the agent's WebSocket arrived from,
+/// and a module or runner writing on an agent's behalf does so from this machine. Loopback is therefore
+/// accepted as well. Anything else -- including an agent that never registered a usable address (recorded as
+/// "unknown") or a request with no peer at all -- is a stranger and is refused.
+fn requester_owns_bucket(registered_ip: Option<&str>, peer: Option<std::net::SocketAddr>) -> bool {
+    let Some(peer) = peer else {
+        return false;
+    };
+    if peer.ip().is_loopback() {
+        return true;
+    }
+    registered_ip.is_some_and(|registered| registered.parse::<std::net::IpAddr>().is_ok_and(|ip| ip == peer.ip()))
+}
+
 /// Log a stored image, and show it on the terminal when the `tty-image` feature is on.
 ///
 /// `put_file` is the only path a file reaches storage through, so it is where every stored file can be watched

services/storage/tests/put.rs

@@ -20,12 +20,12 @@ use edge_toolkit::ws_server::{AgentRecord, AgentRegistry};
 use et_storage_service::{StorageConfig, StorageError, build_store, configure, put_file};
 use tempfile::TempDir;
 
-/// Build a registry with a single connected agent.
+/// Build a registry with a single connected agent registered from `10.0.0.7`.
 fn registry_with_agent(agent_id: &str) -> AgentRegistry<()> {
     let mut agents = BTreeMap::new();
     let _previous: Option<AgentRecord<()>> = agents.insert(
         agent_id.to_string(),
-        AgentRecord::new(AgentConnectionState::Connected, None, Some(())),
+        AgentRecord::new(AgentConnectionState::Connected, Some("10.0.0.7".to_string()), Some(())),
     );
     AgentRegistry::from_agents(agents)
 }
@@ -99,6 +99,7 @@ async fn writes_file_for_registered_agent() {
     let body = b"hello from agent-1".as_ref();
     let req = test::TestRequest::put()
         .uri("/storage/agent-1/payload.txt")
+        .peer_addr("10.0.0.7:40123".parse().unwrap())
         .set_payload(body)
         .to_request();
     let resp = test::call_service(&app, req).await;
@@ -129,6 +130,7 @@ async fn stores_an_image_and_returns_200_even_though_tty_rendering_cannot_succee
     let body = b"\x89PNG\r\n\x1a\nnot a real png, just image-extension-shaped bytes".as_ref();
     let req = test::TestRequest::put()
         .uri("/storage/agent-1/capture.png")
+        .peer_addr("10.0.0.7:40123".parse().unwrap())
         .set_payload(body)
         .to_request();
     let resp = test::call_service(&app, req).await;
@@ -160,6 +162,7 @@ async fn get_and_head_expose_etag_and_head_reports_size() {
     let body = b"etag round-trip payload".as_ref();
     let put = test::TestRequest::put()
         .uri("/storage/agent-1/payload.txt")
+        .peer_addr("10.0.0.7:40123".parse().unwrap())
         .set_payload(body)
         .to_request();
     let put_resp = test::call_service(&app, put).await;
@@ -243,9 +246,122 @@ async fn surfaces_io_failure_as_500() {
 
     let req = test::TestRequest::put()
         .uri("/storage/agent-1/payload.txt")
+        .peer_addr("10.0.0.7:40123".parse().unwrap())
         .set_payload(b"hello".as_ref())
         .to_request();
     let resp = test::call_service(&app, req).await;
 
     assert_eq!(resp.status(), StatusCode::INTERNAL_SERVER_ERROR);
 }
+
+/// A PUT from another host is refused: the agent id's existence is not proof of ownership.
+#[actix_rt::test]
+async fn rejects_foreign_peer_with_403() {
+    let tmp = tempfile::tempdir().unwrap();
+    let config = storage_config(&tmp);
+    let registry = registry_with_agent("agent-1");
+    let app = test::init_service(
+        App::new()
+            .app_data(web::Data::new(registry))
+            .app_data(web::Data::new(config.clone()))
+            .configure(|cfg| configure::<()>(cfg, &config)),
+    )
+    .await;
+
+    let req = test::TestRequest::put()
+        .uri("/storage/agent-1/payload.txt")
+        .peer_addr("198.51.100.9:40123".parse().unwrap())
+        .set_payload(b"not mine".as_ref())
+        .to_request();
+    let resp = test::call_service(&app, req).await;
+
+    assert_eq!(resp.status(), StatusCode::FORBIDDEN);
+    assert!(
+        !tmp.path().join("agent-1").join("payload.txt").exists(),
+        "the refused write must not reach the store"
+    );
+}
+
+/// The address the agent registered from may write: that pair is the ownership check.
+#[actix_rt::test]
+async fn accepts_peer_matching_the_registered_address() {
+    let tmp = tempfile::tempdir().unwrap();
+    let config = storage_config(&tmp);
+    let registry = registry_with_agent("agent-1");
+    let app = test::init_service(
+        App::new()
+            .app_data(web::Data::new(registry))
+            .app_data(web::Data::new(config.clone()))
+            .configure(|cfg| configure::<()>(cfg, &config)),
+    )
+    .await;
+
+    let body = b"hello from the agent's own address";
+    let req = test::TestRequest::put()
+        .uri("/storage/agent-1/payload.txt")
+        .peer_addr("10.0.0.7:40123".parse().unwrap())
+        .set_payload(body.as_ref())
+        .to_request();
+    let resp = test::call_service(&app, req).await;
+
+    assert_eq!(resp.status(), StatusCode::OK);
+    let written = fs_err::read(tmp.path().join("agent-1").join("payload.txt")).unwrap();
+    assert_eq!(written, body);
+}
+
+/// An agent that never registered a usable address (persisted as "unknown") has no identity to match,
+/// so a non-loopback writer is refused rather than trusted.
+#[actix_rt::test]
+async fn rejects_unverifiable_agent_address_with_403() {
+    let tmp = tempfile::tempdir().unwrap();
+    let config = storage_config(&tmp);
+    let mut agents = BTreeMap::new();
+    let _previous: Option<AgentRecord<()>> = agents.insert(
+        "agent-1".to_string(),
+        AgentRecord::new(AgentConnectionState::Connected, None, Some(())),
+    );
+    let registry = AgentRegistry::from_agents(agents);
+    let app = test::init_service(
+        App::new()
+            .app_data(web::Data::new(registry))
+            .app_data(web::Data::new(config.clone()))
+            .configure(|cfg| configure::<()>(cfg, &config)),
+    )
+    .await;
+
+    let req = test::TestRequest::put()
+        .uri("/storage/agent-1/payload.txt")
+        .peer_addr("10.0.0.9:40123".parse().unwrap())
+        .set_payload(b"unverifiable".as_ref())
+        .to_request();
+    let resp = test::call_service(&app, req).await;
+
+    assert_eq!(resp.status(), StatusCode::FORBIDDEN);
+}
+
+/// The runners and modules write from this machine on an agent's behalf: loopback stays accepted.
+#[actix_rt::test]
+async fn accepts_loopback_peer_for_runner_writes() {
+    let tmp = tempfile::tempdir().unwrap();
+    let config = storage_config(&tmp);
+    let registry = registry_with_agent("agent-1");
+    let app = test::init_service(
+        App::new()
+            .app_data(web::Data::new(registry))
+            .app_data(web::Data::new(config.clone()))
+            .configure(|cfg| configure::<()>(cfg, &config)),
+    )
+    .await;
+
+    let body = b"module output";
+    let req = test::TestRequest::put()
+        .uri("/storage/agent-1/module.txt")
+        .peer_addr("127.0.0.1:40123".parse().unwrap())
+        .set_payload(body.as_ref())
+        .to_request();
+    let resp = test::call_service(&app, req).await;
+
+    assert_eq!(resp.status(), StatusCode::OK);
+    let written = fs_err::read(tmp.path().join("agent-1").join("module.txt")).unwrap();
+    assert_eq!(written, body);
+}

services/storage/tests/s3_backend.rs

@@ -53,7 +53,7 @@ fn registry_with_agent(agent_id: &str) -> AgentRegistry<()> {
     let mut agents = BTreeMap::new();
     let _inserted = agents.insert(
         agent_id.to_string(),
-        AgentRecord::new(AgentConnectionState::Connected, None, Some(())),
+        AgentRecord::new(AgentConnectionState::Connected, Some("10.0.0.7".to_string()), Some(())),
     );
     AgentRegistry::from_agents(agents)
 }
@@ -161,6 +161,7 @@ async fn round_trips_put_and_get_through_the_s3_backend() {
 
         let put = test::TestRequest::put()
             .uri(&format!("/storage/{AGENT}/{FILENAME}"))
+            .peer_addr("10.0.0.7:40123".parse().unwrap())
             .set_payload(BODY)
             .to_request();
         let put_status = test::call_service(&app, put).await.status();