Deduplicate artifact bytes and deliver files by authenticated hash

agent/ec_agent.py

@@ -33,7 +33,7 @@ import urllib.parse
 import urllib.request
 import uuid
 
-VERSION = "1.5.3"
+VERSION = "1.6.0"
 LABEL = "io.electric-cherry"
 MAX_FILE = 32 * 1024 * 1024
 MANIFEST_MAX_ENTRIES = 10000
@@ -1356,7 +1356,7 @@ class Agent:
                 if line.startswith("PRETTY_NAME="):
                     os_name = line.split("=", 1)[1].strip('"')
         update = self.update_inventory()
-        capabilities = ["library.remove", "instance.cancel-create", "instance.safety-override", "operations.read-lane", "operations.health-lane", "rcon.status-batch", "instance.heartbeat-state", "files.manifest"] + (["host.decommission"] if self.can_uninstall() else []) + (["agent.update"] if update["enabled"] else [])
+        capabilities = ["library.remove", "instance.cancel-create", "instance.safety-override", "operations.read-lane", "operations.health-lane", "rcon.status-batch", "instance.heartbeat-state", "files.manifest", "files.blobs"] + (["host.decommission"] if self.can_uninstall() else []) + (["agent.update"] if update["enabled"] else [])
         return {"hostname": socket.gethostname(), "agentVersion": VERSION, "protocolVersion": 1, "capabilities": capabilities, "agentUpdate": update, "agentLanes": self.lane_inventory(), "os": os_name, "osVersion": platform.release(), "architecture": platform.machine(), "cpu": platform.processor() or platform.machine(), "cpuCount": os.cpu_count() or 0, **resources, **({"cpuUsagePercent": cpu_usage} if cpu_usage is not None else {}), "memoryTotal": memory.get("MemTotal", 0), "memoryAvailable": memory.get("MemAvailable", 0), "storageTotal": usage.total, "storageAvailable": usage.free, "dockerVersion": docker_version, "dockerRunning": docker_running, "directoriesWritable": os.access(self.root, os.W_OK), "runtimeImages": sorted(set(images)), "libraries": libraries, "libraryCaches": caches, "currentOperation": self.current_operation, "portRange": self.config.get("portRange", {"min": 28960, "max": 29060})}
 
     def safe_to_disrupt(self, instance_id, generation):
@@ -1559,6 +1559,18 @@ class Agent:
             managed_write(root, item["path"], content, int(self.config.get("runtimeUid", 1000)), int(self.config.get("runtimeGid", 1000)))
 
     def decode_file(self, item):
+        if "contentSha256" in item:
+            sha = item["contentSha256"]
+            if not isinstance(sha, str) or not re.fullmatch(r"[a-f0-9]{64}", sha):
+                raise AgentError("Invalid file content hash")
+            url = self.config["controller"].rstrip("/") + "/api/agent/blobs/" + sha
+            with self.artifact_open("GET", url) as response:
+                content = response.read(MAX_FILE + 1)
+            if len(content) > MAX_FILE:
+                raise AgentError("File exceeds staging size limit")
+            if digest(content) != sha or item.get("sha256", sha) != sha:
+                raise AgentError("Staged file checksum mismatch")
+            return content
         try:
             content = base64.b64decode(item["contentBase64"], validate=True)
         except (ValueError, KeyError, TypeError):
@@ -1677,7 +1689,7 @@ class Agent:
             if item["path"] in paths:
                 raise AgentError("Duplicate staged file path")
             paths.add(item["path"])
-            deleting = item.get("delete") is True or item.get("contentBase64") is None
+            deleting = item.get("delete") is True or (item.get("contentBase64") is None and "contentSha256" not in item)
             content = None if deleting else self.decode_file(item)
             total_size += len(content) if content is not None else 0
             if total_size > 64 * 1024 * 1024:

agent/tests/test_agent.py

@@ -832,6 +832,26 @@ class AgentTests(unittest.TestCase):
         self.assertEqual(result["state"], "completed", result)
         return payload
 
+    def test_hash_file_downloads_are_verified_and_staged_as_writes(self):
+        self.create()
+        content = b"downloaded plugin"
+        sha = module.digest(content)
+        file = {"id": 1, "path": "configs/server.cfg", "contentSha256": sha, "sha256": sha}
+        with patch.object(self.agent, "artifact_open", return_value=Response(content)) as fetch:
+            result = self.agent.stage(INSTANCE, 1, {"stageId": "hash-stage", "files": [file]}, "hash-operation")
+        self.assertEqual(result["state"], "staged")
+        self.assertTrue(fetch.call_args.args[1].endswith("/api/agent/blobs/" + sha))
+        manifest = json.loads((self.agent.instance_dir(INSTANCE, 1) / "stages/hash-stage/manifest.json").read_text())
+        self.assertFalse(manifest["files"][0]["delete"])
+        self.assertEqual(base64.b64decode(manifest["files"][0]["contentBase64"]), content)
+        with patch.object(self.agent, "artifact_open", return_value=Response(b"wrong bytes")):
+            with self.assertRaisesRegex(AgentError, "checksum mismatch"):
+                self.agent.decode_file(file)
+        with patch.object(self.agent, "artifact_open") as fetch:
+            with self.assertRaisesRegex(AgentError, "Invalid file content hash"):
+                self.agent.decode_file({"contentSha256": "../../secret"})
+            fetch.assert_not_called()
+
     def test_same_id_reports_saved_result_without_executing_again(self):
         self.create()
         operation = self.operation("instance.restart")