From a99d7976997d18d277f5f02d6370ec5d9129337e Mon Sep 17 00:00:00 2001
From: Brett Holman <brett.holman@canonical.com>
Date: Mon, 5 Aug 2024 15:38:56 -0600
Subject: [PATCH] fix: Fix ftp failures (#5585)

- fix exception handling when retr fails
- test: Close connection on failure
- test: Ensure server is running before it is queried
---
 cloudinit/url_helper.py                       | 42 ++++++++++---------
 .../datasources/test_nocloud.py               |  7 +++-
 2 files changed, 29 insertions(+), 20 deletions(-)

diff --git a/cloudinit/url_helper.py b/cloudinit/url_helper.py
index d409e3228584..cf74a257d552 100644
--- a/cloudinit/url_helper.py
+++ b/cloudinit/url_helper.py
@@ -140,6 +140,15 @@ def read_ftps(url: str, timeout: float = 5.0, **kwargs: dict) -> "FtpResponse":
                     user=user,
                     passwd=url_parts.password or "",
                 )
+                LOG.debug("Creating a secure connection")
+                ftp_tls.prot_p()
+                LOG.debug("Reading file: %s", url_parts.path)
+                ftp_tls.retrbinary(
+                    f"RETR {url_parts.path}", callback=buffer.write
+                )
+
+                response = FtpResponse(buffer.getvalue(), url)
+                LOG.debug("Closing connection")
             except ftplib.error_perm as e:
                 LOG.warning(
                     "Attempted to connect to an insecure ftp server but used "
@@ -156,15 +165,9 @@ def read_ftps(url: str, timeout: float = 5.0, **kwargs: dict) -> "FtpResponse":
                     headers=None,
                     url=url,
                 ) from e
-            LOG.debug("Creating a secure connection")
-            ftp_tls.prot_p()
-            LOG.debug("Reading file: %s", url_parts.path)
-            ftp_tls.retrbinary(f"RETR {url_parts.path}", callback=buffer.write)
-
-            response = FtpResponse(buffer.getvalue(), url)
-            LOG.debug("Closing connection")
-            ftp_tls.close()
-            return response
+            finally:
+                ftp_tls.close()
+                return response
         else:
             try:
                 ftp = ftplib.FTP()
@@ -176,6 +179,14 @@ def read_ftps(url: str, timeout: float = 5.0, **kwargs: dict) -> "FtpResponse":
                     port=port,
                     timeout=timeout or 5.0,  # uses float internally
                 )
+                LOG.debug("Attempting to login with user [%s]", user)
+                ftp.login(
+                    user=user,
+                    passwd=url_parts.password or "",
+                )
+                LOG.debug("Reading file: %s", url_parts.path)
+                ftp.retrbinary(f"RETR {url_parts.path}", callback=buffer.write)
+                response = FtpResponse(buffer.getvalue(), url)
             except ftplib.all_errors as e:
                 code = ftp_get_return_code_from_exception(e)
                 raise UrlError(
@@ -187,16 +198,9 @@ def read_ftps(url: str, timeout: float = 5.0, **kwargs: dict) -> "FtpResponse":
                     headers=None,
                     url=url,
                 ) from e
-            LOG.debug("Attempting to login with user [%s]", user)
-            ftp.login(
-                user=user,
-                passwd=url_parts.password or "",
-            )
-            LOG.debug("Reading file: %s", url_parts.path)
-            ftp.retrbinary(f"RETR {url_parts.path}", callback=buffer.write)
-            response = FtpResponse(buffer.getvalue(), url)
-            LOG.debug("Closing connection")
-            ftp.close()
+            finally:
+                LOG.debug("Closing connection")
+                ftp.close()
             return response
 
 
diff --git a/tests/integration_tests/datasources/test_nocloud.py b/tests/integration_tests/datasources/test_nocloud.py
index c3462c433a34..f6659d45d171 100644
--- a/tests/integration_tests/datasources/test_nocloud.py
+++ b/tests/integration_tests/datasources/test_nocloud.py
@@ -267,6 +267,8 @@ def _boot_with_cmdline(
                 #!/usr/bin/python3
                 import logging
 
+                from systemd.daemon import notify
+
                 from pyftpdlib.authorizers import DummyAuthorizer
                 from pyftpdlib.handlers import FTPHandler, TLS_FTPHandler
                 from pyftpdlib.servers import FTPServer
@@ -298,6 +300,9 @@ def _boot_with_cmdline(
                 handler.abstracted_fs = UnixFilesystem
                 server = FTPServer(("localhost", 2121), handler)
 
+                # tell systemd to proceed
+                notify("READY=1")
+
                 # start the ftp server
                 server.serve_forever()
                 """
@@ -357,7 +362,7 @@ def _boot_with_cmdline(
                 Before=cloud-init-network.service
 
                 [Service]
-                Type=exec
+                Type=notify
                 ExecStart=/server.py
 
                 [Install]