Validate server and token with a single request

ayon_api/server_api.py

@@ -74,6 +74,7 @@
     get_media_mime_type_for_stream,
     get_machine_name,
     fill_own_attribs,
+    UserInfo,
     get_user_info_by_token,
 )
 from ._api_helpers import (
@@ -421,6 +422,9 @@ def __init__(
         self._links_graphql_support_data: bool | None = None
 
         self._session = None
+        # Session without authorization used for availability check and
+        #   token validation, see '_get_validation_session'
+        self._validation_session: requests.Session | None = None
 
         self._base_functions_mapping = {
             RequestTypes.get: requests.get,
@@ -497,8 +501,9 @@ def set_ssl_verify(
         if self._ssl_verify == ssl_verify:
             return
         self._ssl_verify = ssl_verify
-        if self._session is not None:
-            self._session.verify = ssl_verify
+        for session in (self._session, self._validation_session):
+            if session is not None:
+                session.verify = ssl_verify
 
     def get_cert(self) -> str | None:
         """Current cert file used for connection to server.
@@ -519,8 +524,9 @@ def set_cert(self, cert: str | None) -> None:
         if cert == self._cert:
             return
         self._cert = cert
-        if self._session is not None:
-            self._session.cert = cert
+        for session in (self._session, self._validation_session):
+            if session is not None:
+                session.cert = cert
 
     ssl_verify = property(get_ssl_verify, set_ssl_verify)
     cert = property(get_cert, set_cert)
@@ -840,11 +846,20 @@ def as_username(
 
     @property
     def is_server_available(self) -> bool:
+        if (
+            self._server_available is None
+            and self._token_info.token
+            and self._token_info.is_valid is None
+        ):
+            # Token validation also finds out if server is available,
+            #   which saves a request
+            user_info, responded = self._request_token_user_info()
+            if responded:
+                self._store_token_validation(user_info)
+
         if self._server_available is None:
-            response = requests.get(
-                self._base_url,
-                cert=self._cert,
-                verify=self._ssl_verify,
+            response = self._get_validation_session().get(
+                f"{self._rest_url}/info",
                 timeout=self.timeout,
             )
             self._server_available = response.status_code == 200
@@ -872,29 +887,49 @@ def validate_token(self) -> bool:
             self.close_session()
             return False
 
-        # TODO add other possible validations
-        # - existence of 'user' key in info
-        # - validate that 'site_id' is in 'sites' in info
+        user_info, responded = self._request_token_user_info()
+        if not responded:
+            # Check server url, raises an error if the url does not lead
+            #   to AYON server
+            self._get_server_info()
+
+        self._store_token_validation(user_info)
+        return self._token_info.is_valid
 
-        # Check server url
-        self._get_server_info()
+    def _request_token_user_info(self) -> tuple[UserInfo, bool]:
+        """Request information about user of the token.
 
+        Server is marked as available if it did respond to the token.
+
+        Returns:
+            tuple[UserInfo, bool]: User information and if server did
+                respond with success or unauthorized status.
+
+        """
         user_info = get_user_info_by_token(
             self.base_url,
             self._token_info.token,
             verify=self._ssl_verify,
             cert=self._cert,
             timeout=self.timeout,
+            session=self._get_validation_session(),
+        )
+        responded = (
+            user_info.response is not None
+            and user_info.response.status_code in (200, 401)
         )
+        if responded:
+            self._server_available = True
+        return user_info, responded
+
+    def _store_token_validation(self, user_info: UserInfo) -> None:
         self._token_info.is_valid = user_info.is_valid
         self._token_info.unauthorized_response = user_info.response
         is_service = None
         if user_info.is_valid:
             is_service = user_info.is_service
         self._token_info.is_service = is_service
 
-        return self._token_info.is_valid
-
     def set_token(self, token: str | None) -> None:
         self._reset_token()
         self._token_info.token = token
@@ -936,10 +971,16 @@ def create_session(
             raise ValueError("Session is already created.")
 
         self._as_user_stack.clear()
-        # Validate token before session creation
-        self.validate_token()
+        # Validate token before session creation, unless it was already
+        #   validated (e.g. by server availability check)
+        if self._token_info.is_valid is None:
+            self.validate_token()
 
-        session = requests.Session()
+        # Reuse connection of availability check and token validation
+        session = self._validation_session
+        self._validation_session = None
+        if session is None:
+            session = requests.Session()
         session.cert = self._cert
         session.verify = self._ssl_verify
         session.headers.update(self.get_headers())
@@ -954,6 +995,10 @@ def create_session(
         self._session = session
 
     def close_session(self) -> None:
+        if self._validation_session is not None:
+            self._validation_session.close()
+            self._validation_session = None
+
         if self._session is None:
             return
 
@@ -962,6 +1007,25 @@ def close_session(self) -> None:
         self._session_functions_mapping = {}
         session.close()
 
+    def _get_validation_session(self) -> requests.Session:
+        """Session for requests made before authorized session is created.
+
+        Server availability check and token validation use the same
+            session, so they share one connection. The session is then used
+            as the authorized session by 'create_session', which reuses
+            the connection too.
+
+        Returns:
+            requests.Session: Session without authorization headers.
+
+        """
+        if self._validation_session is None:
+            session = requests.Session()
+            session.cert = self._cert
+            session.verify = self._ssl_verify
+            self._validation_session = session
+        return self._validation_session
+
     def create_websocket(
         self,
         endpoint: str,

ayon_api/utils.py

@@ -714,6 +714,7 @@ def get_user_info_by_token(
     verify: str | bool | None = None,
     cert: str | None = None,
     timeout: float | None = None,
+    session: requests.Session | None = None,
 ) -> UserInfo:
     """Get user information by url and token.
 
@@ -726,6 +727,9 @@ def get_user_info_by_token(
             'AYON_CERT_FILE' environment variable is used if not specified.
         timeout (float | None): Timeout for request. Value from
             'get_default_timeout' is used if not specified.
+        session (requests.Session | None): Session used for requests, so
+            its connection can be reused. New connection is used for each
+            request if not specified.
 
     Returns:
         UserInfo: User information if url and token are valid.
@@ -747,13 +751,20 @@ def get_user_info_by_token(
     base_headers = {
         "Content-Type": "application/json",
     }
-    for header_value, is_service in (
-        ({"Authorization": f"Bearer {token}"}, False),
-        ({"X-Api-Key": token}, True),
-    ):
+    # Server accepts only 64 characters long tokens in 'Authorization'
+    #   header (session tokens), a token of other length can be only
+    #   an API key. Skip the request that would fail.
+    auth_headers = [({"X-Api-Key": token}, True)]
+    if len(token) == 64:
+        auth_headers.insert(
+            0, ({"Authorization": f"Bearer {token}"}, False)
+        )
+
+    request_func = requests.get if session is None else session.get
+    for header_value, is_service in auth_headers:
         headers = base_headers.copy()
         headers.update(header_value)
-        response = requests.get(
+        response = request_func(
             f"{url}/api/users/me",
             headers=headers,
             timeout=timeout,

tests/test_token_validation.py

@@ -0,0 +1,169 @@
+"""Requests made to validate server url and token.
+
+Does not require running AYON server.
+"""
+import pytest
+import requests
+
+from ayon_api import GlobalServerAPI
+from ayon_api.constants import SERVER_API_ENV_KEY, SERVER_URL_ENV_KEY
+from ayon_api.exceptions import ServerNotReached
+from ayon_api.server_api import ServerAPI
+from ayon_api.utils import get_user_info_by_token
+
+from .fake_transfer import FakeResponse
+
+URL = "http://localhost:0"
+USER_TOKEN = "a" * 64
+API_KEY = "b" * 32
+LONG_API_KEY = "c" * 64
+
+
+class FakeServer:
+    """Answers 'users/me' and 'info' requests and records them."""
+    def __init__(self, user_tokens=(), api_keys=(), is_ayon=True):
+        self.user_tokens = set(user_tokens)
+        self.api_keys = set(api_keys)
+        self.is_ayon = is_ayon
+        self.requests = []
+        self.sessions = []
+
+    def get(self, url, headers=None, **kwargs):
+        headers = headers or {}
+        endpoint = url[len(URL):]
+        if "X-Api-Key" in headers:
+            auth = "x-api-key"
+            valid = headers["X-Api-Key"] in self.api_keys
+        elif "Authorization" in headers:
+            auth = "bearer"
+            token = headers["Authorization"].split(" ", 1)[-1]
+            valid = token in self.user_tokens
+        else:
+            auth = "none"
+            valid = False
+        self.requests.append((endpoint, auth))
+
+        if not self.is_ayon:
+            return FakeResponse(404)
+        if endpoint == "/api/info":
+            return FakeResponse(200, json_data={"version": "1.0.0"})
+        if endpoint == "/api/users/me":
+            if valid:
+                return FakeResponse(200, json_data={"name": "user"})
+            return FakeResponse(401)
+        return FakeResponse(404)
+
+    def create_session(self):
+        session = FakeSession(self)
+        self.sessions.append(session)
+        return session
+
+
+class FakeSession:
+    def __init__(self, server):
+        self.server = server
+        self.headers = {}
+        self.cert = None
+        self.verify = True
+        self.closed = False
+
+    def get(self, url, **kwargs):
+        return self.server.get(url, **kwargs)
+
+    def post(self, *args, **kwargs):
+        raise NotImplementedError
+
+    put = patch = delete = post
+
+    def close(self):
+        self.closed = True
+
+
+def _fail_request(*args, **kwargs):
+    pytest.fail("Request without a session")
+
+
+@pytest.fixture
+def server(monkeypatch):
+    server = FakeServer(user_tokens=[USER_TOKEN], api_keys=[API_KEY])
+    monkeypatch.setattr(requests, "Session", server.create_session)
+    monkeypatch.setattr(requests, "get", _fail_request)
+    return server
+
+
+@pytest.mark.parametrize(
+    "token, expected_requests, is_service",
+    [
+        (USER_TOKEN, [("/api/users/me", "bearer")], False),
+        # Only 64 characters long token can be used as 'Bearer'
+        (API_KEY, [("/api/users/me", "x-api-key")], True),
+        (
+            LONG_API_KEY,
+            [("/api/users/me", "bearer"), ("/api/users/me", "x-api-key")],
+            True,
+        ),
+    ],
+)
+def test_user_info_by_token(token, expected_requests, is_service):
+    server = FakeServer(user_tokens=[USER_TOKEN], api_keys=[LONG_API_KEY])
+    server.api_keys.add(API_KEY)
+    user_info = get_user_info_by_token(
+        URL, token, session=server.create_session()
+    )
+    assert user_info.is_valid
+    assert user_info.is_service is is_service
+    assert server.requests == expected_requests
+
+
+@pytest.mark.parametrize(
+    "token, auth", [(USER_TOKEN, "bearer"), (API_KEY, "x-api-key")]
+)
+def test_connection_validates_with_single_request(server, token, auth):
+    con = ServerAPI(URL, token)
+
+    assert con.has_valid_token
+    assert con.is_server_available
+    assert server.requests == [("/api/users/me", auth)]
+    # Authorized session reuses connection of the validation
+    assert len(server.sessions) == 1
+    assert con._session is server.sessions[0]
+
+
+def test_global_connection_validates_with_single_request(
+    server, monkeypatch
+):
+    monkeypatch.setenv(SERVER_URL_ENV_KEY, URL)
+    monkeypatch.setenv(SERVER_API_ENV_KEY, API_KEY)
+    con = GlobalServerAPI()
+
+    assert con.has_valid_token
+    assert server.requests == [("/api/users/me", "x-api-key")]
+    assert len(server.sessions) == 1
+
+
+def test_invalid_token_marks_server_available(server):
+    con = ServerAPI(URL, "d" * 32)
+
+    assert con.is_server_available
+    assert not con.has_valid_token
+    assert server.requests == [("/api/users/me", "x-api-key")]
+
+
+def test_availability_without_token(server):
+    con = ServerAPI(URL, create_session=False)
+
+    assert con.is_server_available
+    assert server.requests == [("/api/info", "none")]
+
+
+def test_url_of_other_server_is_not_available(monkeypatch):
+    server = FakeServer(api_keys=[API_KEY], is_ayon=False)
+    monkeypatch.setattr(requests, "Session", server.create_session)
+    monkeypatch.setattr(requests, "get", _fail_request)
+
+    with pytest.raises(ServerNotReached):
+        ServerAPI(URL, API_KEY)
+    assert server.requests == [
+        ("/api/users/me", "x-api-key"),
+        ("/api/info", "none"),
+    ]