From fb2b275781a1cd0c152244c36cc425abf0411de8 Mon Sep 17 00:00:00 2001
From: Waqar Ahmed <waqarahmedjoyia@live.com>
Date: Thu, 19 Dec 2024 21:20:50 +0500
Subject: [PATCH] Minor fixes

---
 .../plugins/acme_protocol_/authenticators/truenas_connect.py  | 2 +-
 .../plugins/truenas_connect/finalize_registration.py          | 4 ++++
 src/middlewared/middlewared/plugins/truenas_connect/mixin.py  | 2 +-
 3 files changed, 6 insertions(+), 2 deletions(-)

diff --git a/src/middlewared/middlewared/plugins/acme_protocol_/authenticators/truenas_connect.py b/src/middlewared/middlewared/plugins/acme_protocol_/authenticators/truenas_connect.py
index e704aece2a1f2..be41119e78144 100644
--- a/src/middlewared/middlewared/plugins/acme_protocol_/authenticators/truenas_connect.py
+++ b/src/middlewared/middlewared/plugins/acme_protocol_/authenticators/truenas_connect.py
@@ -40,7 +40,7 @@ def _perform_internal(self, domain, validation_name, validation_content):
         response = requests.post(LECA_DNS_URL, data=json.dumps({
             'token': validation_content,
             'hostnames': [domain],  # We should be using validation name here
-        }), headers=auth_headers(self.attributes))
+        }), headers=auth_headers(self.attributes), timeout=30)
         if response.status_code != 201:
             raise CallError(
                 f'Failed to perform {self.NAME} challenge for {domain!r} domain with '
diff --git a/src/middlewared/middlewared/plugins/truenas_connect/finalize_registration.py b/src/middlewared/middlewared/plugins/truenas_connect/finalize_registration.py
index dd7b67dff3d47..d60ccbe29831f 100644
--- a/src/middlewared/middlewared/plugins/truenas_connect/finalize_registration.py
+++ b/src/middlewared/middlewared/plugins/truenas_connect/finalize_registration.py
@@ -107,6 +107,10 @@ async def poll_once(self, claim_token, system_id):
 async def setup(middleware):
     tn_config = await middleware.call('tn_connect.config')
     if tn_config['status'] == Status.REGISTRATION_FINALIZATION_WAITING.name:
+        logger.debug(
+            'Registration finalization failed as middleware was restarted while waiting '
+            'for TNC registration finalization'
+        )
         # This means middleware got restarted or the system was rebooted while we were waiting for
         # registration to finalize, so in this case we set the state to registration failed
         await middleware.call('tn_connect.finalize.status_update', Status.REGISTRATION_FINALIZATION_FAILED)
diff --git a/src/middlewared/middlewared/plugins/truenas_connect/mixin.py b/src/middlewared/middlewared/plugins/truenas_connect/mixin.py
index f92725d268422..962c501fcb269 100644
--- a/src/middlewared/middlewared/plugins/truenas_connect/mixin.py
+++ b/src/middlewared/middlewared/plugins/truenas_connect/mixin.py
@@ -28,7 +28,7 @@ async def _call(
             headers = headers or {}
             headers['Content-Type'] = 'application/json'
         try:
-            async with aiohttp.timeout(timeout):
+            async with asyncio.timeout(timeout):
                 async with aiohttp.ClientSession(raise_for_status=True, trust_env=True) as session:
                     req = await getattr(session, mode)(
                         endpoint,