feat: Add legacy and pilot token refresh (choice between them with partial)

Author: Robin-Van-de-Merghel

Pilot/dirac-pilot.py

@@ -84,7 +84,8 @@
             pilotUUID=pilotParams.pilotUUID,
             debugFlag=pilotParams.debugFlag,
             jwt=pilotParams.jwt,
-            legacy_logging=pilotParams.isLegacyLogging
+            legacy_logging=pilotParams.isLegacyLogging,
+            clientID=pilotParams.clientID
         )
         log.info("Remote logger activated")
         log.buffer.write(log.format_to_json(

Pilot/pilotTools.py

@@ -69,9 +69,23 @@ def load_module_from_path(module_name, path_to_module):
     basestring = str
 
 try:
-    from Pilot.proxyTools import getVO, BaseRequest, TokenBasedRequest, extract_diracx_payload
+    from Pilot.proxyTools import (
+        getVO, 
+        BaseRequest, 
+        TokenBasedRequest, 
+        extract_diracx_payload,
+        refreshPilotToken,
+        refreshUserToken
+    )
 except ImportError:
-    from proxyTools import getVO, BaseRequest, TokenBasedRequest, extract_diracx_payload
+    from proxyTools import (
+        getVO, 
+        BaseRequest, 
+        TokenBasedRequest, 
+        extract_diracx_payload,
+        refreshPilotToken,
+        refreshUserToken
+    )
 
 try:
     FileNotFoundError  # pylint: disable=used-before-assignment
@@ -525,9 +539,10 @@ def __init__(
         isPilotLoggerOn=True,
         pilotUUID="unknown",
         flushInterval=10,
-        bufsize=1000,
+        bufsize=250,
         jwt = {},
-        legacy_logging = False
+        legacy_logging = False,
+        clientID = ""
     ):
         """
         c'tor
@@ -538,7 +553,7 @@ def __init__(
         self.url = url
         self.pilotUUID = pilotUUID
         self.isPilotLoggerOn = isPilotLoggerOn
-        sendToURL = partial(sendMessage, url, pilotUUID, legacy_logging)
+        sendToURL = partial(sendMessage, url, pilotUUID, legacy_logging, clientID)
         self.buffer = FixedSizeBuffer(sendToURL, bufsize=bufsize, autoflush=flushInterval, jwt=jwt)
 
     def format_to_json(self, level, message):
@@ -622,7 +637,7 @@ class FixedSizeBuffer(object):
     Once it's full, a message is sent to a remote server and the buffer is renewed.
     """
 
-    def __init__(self, senderFunc, bufsize=1000, autoflush=10, jwt={}):
+    def __init__(self, senderFunc, bufsize=250, autoflush=10, jwt={}):
         """
         Constructor.
 
@@ -645,6 +660,10 @@ def __init__(self, senderFunc, bufsize=1000, autoflush=10, jwt={}):
         self._nlines = 0
         self.senderFunc = senderFunc
         self.jwt = jwt
+        # A fixed buffer used by a remote buffer can be deactivated:
+        # If there's a 403/401 error, instead of crashing the pilot, 
+        # we will deactivate the log sending, and prefer just running the pilot.
+        self.activated = True
 
     @synchronized
     def write(self, content_json):
@@ -657,13 +676,11 @@ def write(self, content_json):
         :return: None
         :rtype: None
         """
+        if not self.activated:
+            pass
 
         self.output.extend(content_json)
-        
-        try:
-            self._nlines += max(1, len(content_json))
-        except Exception:
-            raise ValueError(content_json)
+        self._nlines += max(1, len(content_json))
         self.sendFullBuffer()
 
     @synchronized
@@ -674,7 +691,11 @@ def sendFullBuffer(self):
         """
 
         if self._nlines >= self.bufsize:
-            self.flush()
+            try:
+                self.flush()
+            except Exception as e:
+                print("Deactivating fixed size buffer due to", str(e))
+                self.activated = False
             self.output = [] 
 
     @synchronized
@@ -685,6 +706,9 @@ def flush(self, force=False):
         :return: None
         :rtype:  None
         """
+        if not self.activated:
+            pass
+
         if force or (self.output and self._nlines > 0):
             self.senderFunc(self.jwt, self.output)
             self._nlines = 0
@@ -700,7 +724,7 @@ def cancelTimer(self):
             self._timer.cancel()
 
 
-def sendMessage(diracx_URL, pilotUUID, legacy=False, jwt={}, rawMessage = []):
+def sendMessage(diracx_URL, pilotUUID, legacy=False, clientID="", jwt={}, rawMessage = []):
     """
     Invoke a remote method on a Tornado server and pass a JSON message to it.
 
@@ -724,8 +748,21 @@ def sendMessage(diracx_URL, pilotUUID, legacy=False, jwt={}, rawMessage = []):
 
     if legacy:
         endpoint_path = "api/pilots/legacy/message"
+        refresh_callback = partial(
+            refreshUserToken,
+            diracx_URL,
+            pilotUUID,
+            jwt,
+            clientID
+        )
     else:
         endpoint_path = "api/pilots/internal/message"
+        refresh_callback = partial(
+            refreshPilotToken,
+            diracx_URL,
+            pilotUUID,
+            jwt
+        )
 
     config = TokenBasedRequest(
         diracx_URL=diracx_URL,
@@ -738,7 +775,8 @@ def sendMessage(diracx_URL, pilotUUID, legacy=False, jwt={}, rawMessage = []):
     # Do the request
     _res = config.executeRequest(
         raw_data=raw_data,
-        json_output=False 
+        json_output=False,
+        refresh_callback=refresh_callback
     )
 
 
@@ -775,7 +813,8 @@ def __init__(self, pilotParams):
                 flushInterval=interval,
                 bufsize=bufsize,
                 jwt=pilotParams.jwt,
-                legacy_logging=pilotParams.isLegacyLogging
+                legacy_logging=pilotParams.isLegacyLogging,
+                clientID=pilotParams.clientID
             )
 
         self.log.isPilotLoggerOn = isPilotLoggerOn
@@ -976,7 +1015,7 @@ def __init__(self):
         self.loggerURL = None
         self.isLegacyLogging = False
         self.loggerTimerInterval = 0
-        self.loggerBufsize = 1000
+        self.loggerBufsize = 250
         self.pilotUUID = "unknown"
         self.modules = ""
         self.userEnvVariables = ""

Pilot/proxyTools.py

@@ -252,9 +252,9 @@ def addJwtToHeader(self):
         # Adds the JWT in the HTTP request (in the Bearer field)
         self.headers["Authorization"] = "Bearer %s" % self.jwtData["access_token"]
 
-    def executeRequest(self, raw_data, insecure=False, content_type="json", json_output=True, tries_left=1):
+    def executeRequest(self, raw_data, insecure=False, content_type="json", json_output=True, tries_left=1, refresh_callback=None):
 
-        while (tries_left > 0):
+        while (tries_left >= 0):
 
             try:
                 return super(TokenBasedRequest, self).executeRequest(
@@ -268,11 +268,8 @@ def executeRequest(self, raw_data, insecure=False, content_type="json", json_out
                     raise e
 
                 # If we have an unauthorized error, then refresh and retry
-                refreshPilotToken(
-                    self.diracx_URL,
-                    self.pilotUUID,
-                    self.jwtData
-                )
+                if refresh_callback:
+                    refresh_callback()
 
                 self.addJwtToHeader()
 
@@ -309,10 +306,47 @@ def executeRequest(self, raw_data, insecure=False, content_type="json", json_out
             json_output=json_output
         )
 
+def refreshUserToken(url, pilotUUID, jwt, clientID):
+    """
+    Refresh the JWT token (as a user).
 
-def refreshPilotToken(url, pilotUUID, jwt):
+    :param str url: Server URL
+    :param str pilotUUID: Pilot unique ID
+    :param dict jwt: Shared dict with current JWT; updated in-place
+    :return: None
     """
-    Refresh the JWT token in a separate thread.
+
+    # PRECONDITION: jwt must contain "refresh_token"
+    if not jwt or "refresh_token" not in jwt:
+        raise ValueError("To refresh a token, a pilot needs a JWT with refresh_token")
+
+    # Get CA path from environment
+    caPath = os.getenv("X509_CERT_DIR")
+
+    # Create request object with required configuration
+    config = BaseRequest(
+        url=url + "api/auth/token",
+        caPath=caPath,
+        pilotUUID=pilotUUID,
+    )
+
+    # Perform the request to refresh the token
+    response = config.executeRequest(
+        raw_data={
+            "refresh_token": jwt["refresh_token"],
+            "grant_type": "refresh_token",
+            "client_id": clientID
+        },
+        content_type="x-www-form-urlencoded",
+    )
+
+    # Do NOT assign directly, because jwt is a reference, not a copy
+    jwt["access_token"] = response["access_token"]
+    jwt["refresh_token"] = response["refresh_token"]
+
+def refreshPilotToken(url, pilotUUID, jwt, _=None):
+    """
+    Refresh the JWT token (as a pilot).
 
     :param str url: Server URL
     :param str pilotUUID: Pilot unique ID
@@ -329,7 +363,7 @@ def refreshPilotToken(url, pilotUUID, jwt):
 
     # Create request object with required configuration
     config = BaseRequest(
-        url=url + "/api/auth/pilot-token",
+        url=url + "api/auth/pilot-token",
         caPath=caPath,
         pilotUUID=pilotUUID,
     )
@@ -365,9 +399,12 @@ def revokePilotToken(url, pilotUUID, jwt, clientID):
     # Get CA path from environment
     caPath = os.getenv("X509_CERT_DIR")
 
+    if not url.endswith("/"):
+        url = url + "/"
+
     # Create request object with required configuration
     config = BaseRequest(
-        url="%s/api/auth/revoke" % url,
+        url="%sapi/auth/revoke" % url,
         caPath=caPath,
         pilotUUID=pilotUUID
     )