Fixes a bug which prevents MATLAB from exiting cleanly when token authentication is enabled.

Author: diningPhilosopher64

matlab_proxy/app_state.py

@@ -8,12 +8,11 @@
 import time
 from collections import deque
 from datetime import datetime, timedelta, timezone
-from typing import Final
+from typing import Final, Optional
 
 from matlab_proxy import util
 from matlab_proxy.settings import (
     get_process_startup_timeout,
-    get_default_mwi_log_file_path,
 )
 from matlab_proxy.constants import (
     CONNECTOR_SECUREPORT_FILENAME,
@@ -293,7 +292,24 @@ def _are_required_processes_ready(
 
         return True
 
-    async def _get_matlab_connector_status(self):
+    def _get_token_auth_headers(self) -> Optional[dict]:
+        """Returns token info as headers if authentication is enabled.
+
+        Returns:
+            [Dict | None]: Returns token authentication headers if any.
+        """
+        return (
+            {self.settings["mwi_auth_token_name"]: self.settings["mwi_auth_token_hash"]}
+            if self.settings["mwi_is_token_auth_enabled"]
+            else None
+        )
+
+    async def _get_matlab_connector_status(self) -> str:
+        """Returns the status of MATLABs Embedded Connector.
+
+        Returns:
+            str: Returns any of "up", "down" or "starting" indicating the status of Embedded Connector.
+        """
         if not self.matlab_session_files["matlab_ready_file"].exists():
             return "starting"
 
@@ -303,11 +319,7 @@ async def _get_matlab_connector_status(self):
         # As the matlab_view is now a protected endpoint, we need to pass token information through headers.
 
         # Include token information into the headers if authentication is enabled.
-        headers = (
-            {self.settings["mwi_auth_token_name"]: self.settings["mwi_auth_token_hash"]}
-            if self.settings["mwi_is_token_auth_enabled"]
-            else None
-        )
+        headers = self._get_token_auth_headers()
 
         embedded_connector_status = await mwi.embedded_connector.request.get_state(
             mwi_server_url=self.settings["mwi_server_url"],
@@ -683,10 +695,13 @@ async def __setup_env_for_matlab(self) -> dict:
             else:
                 # On windows stdout is not supported yet.
                 # So, use the default log file for MATLAB logs
-                default_matlab_logs_file = get_default_mwi_log_file_path()
+                matlab_logs_file = self.mwi_logs_dir / MATLAB_LOGS_FILE_NAME
                 # Write MATLAB logs
-                matlab_env["MW_DIAGNOSTIC_DEST"] = f"file={default_matlab_logs_file}"
+                matlab_env["MW_DIAGNOSTIC_DEST"] = f"file={matlab_logs_file}"
 
+            logger.info(
+                f"Writing MATLAB process logs to: {matlab_env['MW_DIAGNOSTIC_DEST']}"
+            )
             matlab_env[
                 "MW_DIAGNOSTIC_SPEC"
             ] = "connector::http::server=all;connector::lifecycle=all"
@@ -1069,12 +1084,16 @@ async def __send_stop_request_to_matlab(self):
 
         try:
             data = mwi.embedded_connector.helpers.get_data_to_eval_mcode("exit")
+            headers = self._get_token_auth_headers()
             url = mwi.embedded_connector.helpers.get_mvm_endpoint(
                 self.settings["mwi_server_url"]
             )
 
             resp_json = await mwi.embedded_connector.send_request(
-                url=url, method="POST", data=data, headers=None
+                url=url,
+                method="POST",
+                data=data,
+                headers=headers,
             )
 
             if resp_json["messages"]["EvalResponse"][0]["isError"]:

matlab_proxy/settings.py

@@ -150,10 +150,6 @@ def get_mwi_config_folder(dev=False):
         return config_folder_path
 
 
-def get_default_mwi_log_file_path():
-    return get_mwi_config_folder() / constants.MATLAB_LOGS_FILE_NAME
-
-
 def get_mwi_logs_root_dir(dev=False):
     return get_mwi_config_folder(dev) / "ports"
 

matlab_proxy/util/mwi/embedded_connector/request.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2022 The MathWorks, Inc.
+# Copyright 2020-2023 The MathWorks, Inc.
 
 """ 
 This file contains the methods to communicate with the embedded connector.

matlab_proxy/util/mwi/token_auth.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2020-2023 The MathWorks, Inc.
+# Copyright 2020-2023 The MathWorks, Inc.
 
 # This file contains functions required to enable token based authentication in the server.
 
@@ -135,6 +135,8 @@ async def protect_endpoint(request):
 
 
 ## Module Private Methods:
+
+
 async def _get_token_name(request):
     """Gets the name of the token from settings.
 

tests/unit/test_app_state.py

@@ -10,6 +10,7 @@
 
 from matlab_proxy.app_state import AppState
 from matlab_proxy.util.mwi.exceptions import LicensingError, MatlabError
+from tests.unit.util import MockResponse
 
 
 @pytest.fixture
@@ -24,6 +25,36 @@ def app_state_fixture():
     return app_state
 
 
+@pytest.fixture
+def sample_token_headers_fixture():
+    return {"mwi_auth_token_name": "asdf"}
+
+
+@pytest.fixture
+def app_state_with_token_auth_fixture(
+    app_state_fixture, sample_token_headers_fixture, tmp_path
+):
+    """Pytest fixture which returns AppState instance with token authentication enabled.
+
+    Args:
+        app_state_fixture (AppState): Pytest fixture
+        tmp_path (str): Built-in pytest fixture
+
+    Returns:
+        (AppState, dict): Instance of the AppState class with token authentication enabled and token headers
+    """
+    tmp_matlab_ready_file = Path(tmp_path) / "tmp_file.txt"
+    tmp_matlab_ready_file.touch()
+    ((mwi_auth_token_name, mwi_auth_token_hash),) = sample_token_headers_fixture.items()
+    app_state_fixture.matlab_session_files["matlab_ready_file"] = tmp_matlab_ready_file
+    app_state_fixture.settings["mwi_is_token_auth_enabled"] = True
+    app_state_fixture.settings["mwi_auth_token_name"] = mwi_auth_token_name
+    app_state_fixture.settings["mwi_auth_token_hash"] = mwi_auth_token_hash
+    app_state_fixture.settings["mwi_server_url"] = "http://localhost:8888"
+
+    return app_state_fixture
+
+
 @pytest.fixture
 def mocker_os_patching_fixture(mocker, platform):
     """A pytest fixture which patches the is_* functions in system.py module
@@ -38,12 +69,20 @@ def mocker_os_patching_fixture(mocker, platform):
     mocker.patch("matlab_proxy.app_state.system.is_linux", return_value=False)
     mocker.patch("matlab_proxy.app_state.system.is_windows", return_value=False)
     mocker.patch("matlab_proxy.app_state.system.is_mac", return_value=False)
+    mocker.patch("matlab_proxy.app_state.system.is_posix", return_value=False)
+
     if platform == "linux":
         mocker.patch("matlab_proxy.app_state.system.is_linux", return_value=True)
+        mocker.patch("matlab_proxy.app_state.system.is_posix", return_value=True)
+
     elif platform == "windows":
         mocker.patch("matlab_proxy.app_state.system.is_windows", return_value=True)
+        mocker.patch("matlab_proxy.app_state.system.is_posix", return_value=False)
+
     else:
         mocker.patch("matlab_proxy.app_state.system.is_mac", return_value=True)
+        mocker.patch("matlab_proxy.app_state.system.is_posix", return_value=True)
+
     return mocker
 
 
@@ -397,9 +436,12 @@ def test_env_variables_filtration_for_xvfb_process(
     assert filtered_env_vars.get(env_var) == is_filtered
 
 
-@pytest.mark.parametrize("platform", [("linux"), ("windows"), ("mac")])
+@pytest.mark.parametrize(
+    "platform, expected_output",
+    [("linux", "stdout"), ("windows", "file"), ("mac", "stdout")],
+)
 async def test_setup_env_for_matlab(
-    mocker_os_patching_fixture, platform, app_state_fixture, tmp_path
+    mocker_os_patching_fixture, platform, expected_output, app_state_fixture, tmp_path
 ):
     """Test to check MW_DIAGNOSTIC_DEST is set appropriately for posix and non-posix systems
 
@@ -411,7 +453,6 @@ async def test_setup_env_for_matlab(
     """
 
     # Arrange
-    expected_log_file_path = tmp_path / "matlab_logs.txt"
     app_state_fixture.licensing = {"type": "existing_license"}
     app_state_fixture.settings = {"mwapikey": None, "matlab_display": ":1"}
     app_state_fixture.mwi_logs_dir = tmp_path
@@ -423,7 +464,47 @@ async def test_setup_env_for_matlab(
     matlab_env = await app_state_fixture._AppState__setup_env_for_matlab()
 
     # Assert
-    if "linux" or "mac" in platform:
-        assert matlab_env["MW_DIAGNOSTIC_DEST"] == "stdout"
-    else:
-        assert matlab_env["MW_DIAGNOSTIC_DEST"] == expected_log_file_path
+    assert expected_output in matlab_env["MW_DIAGNOSTIC_DEST"]
+
+
+@pytest.mark.parametrize(
+    "function_to_call ,mock_response",
+    [
+        ("_get_matlab_connector_status", MockResponse(ok=True)),
+        (
+            "_AppState__send_stop_request_to_matlab",
+            MockResponse(
+                ok=True, payload={"messages": {"EvalResponse": [{"isError": None}]}}
+            ),
+        ),
+    ],
+    ids=["request matlab connector status", "send request to stop matlab"],
+)
+async def test_requests_sent_by_matlab_proxy_have_headers(
+    app_state_with_token_auth_fixture,
+    function_to_call,
+    mock_response,
+    mocker,
+    sample_token_headers_fixture,
+):
+    """Test to check if token headers are included in requests sent by matlab-proxy when authentication is enabled
+
+    Args:
+        app_state_fixture_with_token_auth (AppState): Instance of AppState class with token authentication enabled
+        mocker : Built-in pytest fixture
+    """
+    # Arrange
+    mocked_request = mocker.patch(
+        "aiohttp.ClientSession.request", return_value=mock_response
+    )
+
+    # Act
+    # Call the function passed as a string
+    method = getattr(app_state_with_token_auth_fixture, function_to_call)
+    _ = await method()
+
+    # Assert
+    connector_status_request_headers = list(mocked_request.call_args_list)[0].kwargs[
+        "headers"
+    ]
+    assert sample_token_headers_fixture == connector_status_request_headers

tests/unit/util/mwi/embedded_connector/test_request.py

@@ -1,4 +1,4 @@
-# Copyright 2022 The MathWorks, Inc.
+# Copyright 2022-2023 The MathWorks, Inc.
 
 import pytest
 from matlab_proxy.util import mwi
@@ -7,42 +7,58 @@
 
 
 async def test_send_request_success(mocker):
+    """Test to check the happy path for send_request
+    Args:
+        mocker : Built in pytest fixture
+    """
+    # Arrange
     json_data = {"hello": "world"}
-
     payload = json_data
     mock_resp = MockResponse(payload=payload, ok=True)
+    _ = mocker.patch("aiohttp.ClientSession.request", return_value=mock_resp)
 
-    mocked = mocker.patch("aiohttp.ClientSession.request", return_value=mock_resp)
+    # Act
     res = await mwi.embedded_connector.send_request(
         url="https://localhost:3000", data=json_data, method="GET"
     )
 
+    # Assert
     assert json_data["hello"] == res["hello"]
 
 
 async def test_send_request_failure(mocker):
+    """Test to check if send_request fails when
+    1) EC does not respond
+    2) url or method is not supplied
+    Args:
+        mocker : Built in pytest fixture
+    """
+
+    # Arrange
     json_data = {"hello": "world"}
 
     payload = json_data
     mock_resp = MockResponse(payload=payload, ok=False)
+    _ = mocker.patch("aiohttp.ClientSession.request", return_value=mock_resp)
+
+    # Doesnt have url or data or method
+    mock_resp = MockResponse(payload=payload, ok=False)
 
-    mocked = mocker.patch("aiohttp.ClientSession.request", return_value=mock_resp)
+    # Act
 
     # Failed to communicate with EmbeddedConnector
     with pytest.raises(EmbeddedConnectorError):
-        res = await mwi.embedded_connector.send_request(
+        _ = await mwi.embedded_connector.send_request(
             url="https://localhost:3000", method="GET", data=json_data
         )
 
-    # Doesnt have url or data or method
-    mock_resp = MockResponse(payload=payload, ok=False)
-
     for key in ["url", "method"]:
         options = {
             "url": "https://localhost:3000",
             "data": json_data,
             "method": "GET",
         }
         options[key] = ""
+        # Assert
         with pytest.raises(EmbeddedConnectorError):
-            res = await mwi.embedded_connector.send_request(**options)
+            _ = await mwi.embedded_connector.send_request(**options)