From 8b5f5ccadceb3fd0e8d0117bc78ae410bfc463ad Mon Sep 17 00:00:00 2001
From: Minh Vu <vuhoangminh97@gmail.com>
Date: Sun, 21 Jun 2026 22:08:46 +0200
Subject: [PATCH] Validate edge device ID headers

---
 nvflare/edge/web/views/feg_views.py           |  8 ++
 .../edge/web/views/feg_views_test.py          | 82 +++++++++++++++++++
 2 files changed, 90 insertions(+)
 create mode 100644 tests/unit_test/edge/web/views/feg_views_test.py

diff --git a/nvflare/edge/web/views/feg_views.py b/nvflare/edge/web/views/feg_views.py
index 5f432290e8..98e63b1817 100644
--- a/nvflare/edge/web/views/feg_views.py
+++ b/nvflare/edge/web/views/feg_views.py
@@ -62,6 +62,14 @@ def _process_headers() -> dict:
     device_info_header = headers.get(HttpHeaderKey.DEVICE_INFO, None)
     if device_info_header:
         device_info.from_query_string(device_info_header)
+        header_device_id = device_info.get(EdgeProtoKey.DEVICE_ID)
+        if header_device_id and header_device_id != device_id:
+            raise ApiError(
+                400,
+                EdgeApiStatus.INVALID_REQUEST,
+                f"Device ID mismatch: {HttpHeaderKey.DEVICE_ID} does not match {HttpHeaderKey.DEVICE_INFO}",
+            )
+    device_info.device_id = device_id
     d[EdgeProtoKey.DEVICE_INFO] = device_info
 
     user_info_header = headers.get(HttpHeaderKey.USER_INFO, None)
diff --git a/tests/unit_test/edge/web/views/feg_views_test.py b/tests/unit_test/edge/web/views/feg_views_test.py
new file mode 100644
index 0000000000..95341877b9
--- /dev/null
+++ b/tests/unit_test/edge/web/views/feg_views_test.py
@@ -0,0 +1,82 @@
+# Copyright (c) 2026, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from urllib.parse import urlencode
+
+import pytest
+from flask import Flask
+
+from nvflare.edge.constants import EdgeApiStatus, HttpHeaderKey
+from nvflare.edge.web.models.api_error import ApiError
+from nvflare.edge.web.models.base_model import EdgeProtoKey
+from nvflare.edge.web.views.feg_views import _process_headers
+
+
+@pytest.fixture
+def app():
+    return Flask(__name__)
+
+
+def _process_headers_with(app, headers):
+    with app.test_request_context("/job", headers=headers):
+        return _process_headers()
+
+
+def test_process_headers_uses_device_id_header_if_device_info_is_absent(app):
+    headers = {HttpHeaderKey.DEVICE_ID: "device-1"}
+
+    parsed = _process_headers_with(app, headers)
+
+    device_info = parsed[EdgeProtoKey.DEVICE_INFO]
+    assert device_info.device_id == "device-1"
+
+
+def test_process_headers_uses_device_id_header_if_device_info_has_no_device_id(app):
+    headers = {
+        HttpHeaderKey.DEVICE_ID: "device-1",
+        HttpHeaderKey.DEVICE_INFO: urlencode({"platform": "ios"}),
+    }
+
+    parsed = _process_headers_with(app, headers)
+
+    device_info = parsed[EdgeProtoKey.DEVICE_INFO]
+    assert device_info.device_id == "device-1"
+    assert device_info.platform == "ios"
+
+
+def test_process_headers_accepts_matching_device_info_device_id(app):
+    headers = {
+        HttpHeaderKey.DEVICE_ID: "device-1",
+        HttpHeaderKey.DEVICE_INFO: urlencode({EdgeProtoKey.DEVICE_ID: "device-1", "platform": "ios"}),
+    }
+
+    parsed = _process_headers_with(app, headers)
+
+    device_info = parsed[EdgeProtoKey.DEVICE_INFO]
+    assert device_info.device_id == "device-1"
+    assert device_info.platform == "ios"
+
+
+def test_process_headers_rejects_mismatched_device_info_device_id(app):
+    headers = {
+        HttpHeaderKey.DEVICE_ID: "device-1",
+        HttpHeaderKey.DEVICE_INFO: urlencode({EdgeProtoKey.DEVICE_ID: "device-2", "platform": "ios"}),
+    }
+
+    with pytest.raises(ApiError) as exc_info:
+        _process_headers_with(app, headers)
+
+    assert exc_info.value.status_code == 400
+    assert exc_info.value.status == EdgeApiStatus.INVALID_REQUEST
+    assert "Device ID mismatch" in str(exc_info.value)