Merge pull request #2629 from GNS3/image-handling

Fix image handling

gns3server/api/server.py

@@ -195,7 +195,7 @@ async def http_exception_handler(request: Request, exc: HTTPException):
 
 
 @app.exception_handler(SQLAlchemyError)
-async def sqlalchemry_error_handler(request: Request, exc: SQLAlchemyError):
+async def sqlalchemy_error_handler(request: Request, exc: SQLAlchemyError):
     log.error(f"Controller database error in {request.url.path} ({request.method}): {exc}")
     return JSONResponse(
         status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,

gns3server/controller/appliance_manager.py

@@ -15,12 +15,9 @@
 # You should have received a copy of the GNU General Public License
 # along with this program.  If not, see <http://www.gnu.org/licenses/>.
 
-import sys
 import os
 import json
 import asyncio
-import aiofiles
-import shutil
 import platformdirs
 
 
@@ -29,6 +26,7 @@ from aiohttp.client_exceptions import ClientError
 
 from uuid import UUID
 from pydantic import ValidationError
+from sqlalchemy.exc import SQLAlchemyError
 
 from .appliance import Appliance
 from ..config import Config
@@ -36,8 +34,7 @@ from ..utils.asyncio import locking
 from ..utils.http_client import HTTPClient
 from .controller_error import ControllerBadRequestError, ControllerNotFoundError, ControllerError
 from .appliance_to_template import ApplianceToTemplate
-from ..utils.images import InvalidImageError, write_image, md5sum
+from ..utils.images import InvalidImageError, write_image, read_image_info
-from ..utils.asyncio import wait_run_in_executor
 
 from gns3server import schemas
 from gns3server.utils.images import default_images_directory
@@ -183,17 +180,16 @@ class ApplianceManager:
                         if image_in_db:
                             version_images[appliance_key] = image_in_db.filename
                         else:
-                            # check if the image is on disk
+                            # check if the image is on disk but it not yet in the database
-                            # FIXME: still necessary? the image should have been discovered and saved in the db already
                             image_path = os.path.join(image_dir, appliance_file)
-                            if os.path.exists(image_path) and \
+                            if os.path.exists(image_path):
-                                    await wait_run_in_executor(
+                                image_info = await read_image_info(image_path)
-                                        md5sum,
+                                if image_info.get("checksum") == image_checksum:
-                                        image_path,
+                                    log.info(f"Adding image '{image_path}' to the database")
-                                        cache_to_md5file=False
+                                    try:
-                                    ) == image_checksum:
+                                        await images_repo.add_image(**image_info)
-                                async with aiofiles.open(image_path, "rb") as f:
+                                    except SQLAlchemyError as e:
-                                    await write_image(appliance_file, image_path, f, images_repo, allow_raw_image=True)
+                                        log.warning(f"Error while adding image '{image['path']}' to the database: {e}")
                             else:
                                 # download the image if there is a direct download URL
                                 direct_download_url = image.get("direct_download_url")
@@ -310,7 +306,7 @@ class ApplianceManager:
                                                         f"appliance '{appliance_id}'")
 
             template_data = await self._appliance_to_template(appliance)
-            await self._create_template(template_data, templates_repo, rbac_repo, current_user)
+            return await self._create_template(template_data, templates_repo, rbac_repo, current_user)
 
     def load_appliances(self, symbol_theme: str = None) -> None:
         """

gns3server/db/repositories/images.py

@@ -50,14 +50,20 @@ class ImagesRepository(BaseRepository):
         result = await self._db_session.execute(query)
         return result.scalars().one_or_none()
 
-    async def get_image_by_checksum(self, checksum: str) -> Optional[models.Image]:
+    async def get_image_by_checksum(self, checksum: str, image_dir: str = None) -> Optional[models.Image]:
         """
         Get an image by its checksum.
         """
 
-        query = select(models.Image).where(models.Image.checksum == checksum)
+        if image_dir:
-        result = await self._db_session.execute(query)
+            query = select(models.Image).\
-        return result.scalars().first()
+                where(models.Image.checksum == checksum, models.Image.path.startswith(image_dir))
+            result = await self._db_session.execute(query)
+            return result.scalars().one_or_none()
+        else:
+            query = select(models.Image).where(models.Image.checksum == checksum)
+            result = await self._db_session.execute(query)
+            return result.scalars().first()
 
     async def get_images(self, image_type=None) -> List[models.Image]:
         """
@@ -114,7 +120,7 @@ class ImagesRepository(BaseRepository):
 
         await self._db_session.execute(query)
         await self._db_session.commit()
-        image_db = await self.get_image_by_checksum(checksum)
+        image_db = await self.get_image(image_path)
         if image_db:
             await self._db_session.refresh(image_db)  # force refresh of updated_at value
         return image_db

gns3server/utils/images.py

@@ -34,7 +34,6 @@ import gns3server.db.models as models
 from gns3server.db.repositories.images import ImagesRepository
 from gns3server.utils.asyncio import wait_run_in_executor
 
-
 import logging
 
 log = logging.getLogger(__name__)
@@ -276,7 +275,7 @@ def md5sum(path, working_dir=None, stopped_event=None, cache_to_md5file=True):
             while True:
                 if stopped_event is not None and stopped_event.is_set():
                     log.error(f"MD5 sum calculation of `{path}` has stopped due to cancellation")
-                    return
+                    return None
                 buf = f.read(DEFAULT_BUFFER_SIZE)
                 if not buf:
                     break
@@ -343,20 +342,15 @@ async def write_image(
         allow_raw_image=False
 ) -> models.Image:
 
-    db_image = await images_repo.get_image(image_path)
-    if db_image and os.path.exists(image_path):
-        # the image already exists in the database and on disk
-        log.info(f"Image {image_path} already exists")
-        return db_image
-
     image_dir, image_name = os.path.split(image_filename)
-    log.info(f"Writing image file to '{image_path}'")
     # Store the file under its final name only when the upload is completed
     tmp_path = image_path + ".tmp"
+    log.info(f"Writing image file to '{tmp_path}'")
     os.makedirs(os.path.dirname(image_path), exist_ok=True)
     checksum = hashlib.md5()
     header_magic_len = 7
     image_type = None
+    image_size = 0
     try:
         async with aiofiles.open(tmp_path, "wb") as f:
             async for chunk in stream:
@@ -370,15 +364,22 @@ async def write_image(
         if not image_size or image_size < header_magic_len:
             raise InvalidImageError("The image content is empty or too small to be valid")
 
-        checksum = checksum.hexdigest()
-        duplicate_image = await images_repo.get_image_by_checksum(checksum)
-        if duplicate_image and os.path.dirname(duplicate_image.path) == os.path.dirname(image_path):
-            raise InvalidImageError(f"Image {duplicate_image.filename} with "
-                                    f"same checksum already exists in the same directory")
         if not image_dir:
             directory = default_images_directory(image_type)
             os.makedirs(directory, exist_ok=True)
             image_path = os.path.abspath(os.path.join(directory, image_filename))
+
+        if os.path.exists(image_path):
+            raise InvalidImageError(f"File '{image_path}' already exists, "
+                                    f"please choose a different name or remove the existing image")
+
+        checksum = checksum.hexdigest()
+        image_dir = os.path.dirname(image_path)
+        duplicate_image = await images_repo.get_image_by_checksum(checksum, image_dir)
+        if duplicate_image:
+            raise InvalidImageError(f"Image '{duplicate_image.filename}' with the "
+                                    f"same checksum already exists in '{image_dir}'")
+
         shutil.move(tmp_path, image_path)
         os.chmod(image_path, stat.S_IWRITE | stat.S_IREAD | stat.S_IEXEC)
     finally:
@@ -388,10 +389,6 @@ async def write_image(
         except OSError:
             log.warning(f"Could not remove '{tmp_path}'")
 
-    if db_image:
-        # the image already exists in the database, no need to add it again
-        return db_image
-
     return await images_repo.add_image(
         image_name,
         image_type,

tests/api/routes/controller/test_images.py

@@ -173,15 +173,26 @@ class TestImageRoutes:
         assert response.status_code == status.HTTP_200_OK
         assert response.json()["filename"] == image_name
 
-    # async def test_same_image_is_uploaded(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None:
+    async def test_same_image_is_uploaded(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None:
-    #
+
-    #     image_name = os.path.basename(qcow2_image)
+        with open(qcow2_image, "rb") as f:
-    #     with open(qcow2_image, "rb") as f:
+            image_data = f.read()
-    #         image_data = f.read()
+        response = await client.post(
-    #     response = await client.post(
+            app.url_path_for("upload_image", image_path="image1.qcow2"),
-    #         app.url_path_for("upload_image", image_path=image_name),
+            content=image_data)
-    #         content=image_data)
+        assert response.status_code == status.HTTP_201_CREATED
-    #     assert response.status_code == status.HTTP_201_CREATED
+
+        # same image with same name is uploaded again, it should return 409 Conflict
+        response = await client.post(
+            app.url_path_for("upload_image", image_path="image1.qcow2"),
+            content=image_data)
+        assert response.status_code == status.HTTP_409_CONFLICT
+
+        # same image with different name but same checksum is uploaded again, it should return 409 Conflict
+        response = await client.post(
+           app.url_path_for("upload_image", image_path="image2.qcow2"),
+           content=image_data)
+        assert response.status_code == status.HTTP_409_CONFLICT
 
     async def test_image_delete(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None: