From 4606613673682b9dffcd6f782986f1716e177c2a Mon Sep 17 00:00:00 2001 From: grossmj Date: Thu, 5 Mar 2026 12:38:28 +0800 Subject: [PATCH] Fix image handling --- gns3server/api/server.py | 2 +- gns3server/controller/appliance_manager.py | 28 ++++++++---------- gns3server/db/repositories/images.py | 16 +++++++---- gns3server/utils/images.py | 33 ++++++++++------------ tests/api/routes/controller/test_images.py | 29 +++++++++++++------ 5 files changed, 59 insertions(+), 49 deletions(-) diff --git a/gns3server/api/server.py b/gns3server/api/server.py index acdda6409..145394681 100644 --- a/gns3server/api/server.py +++ b/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, diff --git a/gns3server/controller/appliance_manager.py b/gns3server/controller/appliance_manager.py index 75d576747..ad824e51c 100644 --- a/gns3server/controller/appliance_manager.py +++ b/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 . -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.asyncio import wait_run_in_executor +from ..utils.images import InvalidImageError, write_image, read_image_info 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 - # FIXME: still necessary? the image should have been discovered and saved in the db already + # check if the image is on disk but it not yet in the database image_path = os.path.join(image_dir, appliance_file) - if os.path.exists(image_path) and \ - await wait_run_in_executor( - md5sum, - image_path, - cache_to_md5file=False - ) == image_checksum: - async with aiofiles.open(image_path, "rb") as f: - await write_image(appliance_file, image_path, f, images_repo, allow_raw_image=True) + if os.path.exists(image_path): + image_info = await read_image_info(image_path) + if image_info.get("checksum") == image_checksum: + log.info(f"Adding image '{image_path}' to the database") + try: + await images_repo.add_image(**image_info) + except SQLAlchemyError as e: + 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: """ diff --git a/gns3server/db/repositories/images.py b/gns3server/db/repositories/images.py index 9d83e72ee..8f6c64dc9 100644 --- a/gns3server/db/repositories/images.py +++ b/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) - result = await self._db_session.execute(query) - return result.scalars().first() + if image_dir: + query = select(models.Image).\ + 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 diff --git a/gns3server/utils/images.py b/gns3server/utils/images.py index beba1a5c2..dfb3152b3 100644 --- a/gns3server/utils/images.py +++ b/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, diff --git a/tests/api/routes/controller/test_images.py b/tests/api/routes/controller/test_images.py index e60dfec55..8760c9075 100644 --- a/tests/api/routes/controller/test_images.py +++ b/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: - # - # image_name = os.path.basename(qcow2_image) - # with open(qcow2_image, "rb") as f: - # image_data = f.read() - # response = await client.post( - # app.url_path_for("upload_image", image_path=image_name), - # content=image_data) - # assert response.status_code == status.HTTP_201_CREATED + async def test_same_image_is_uploaded(self, app: FastAPI, client: AsyncClient, qcow2_image: str) -> None: + + with open(qcow2_image, "rb") as f: + image_data = f.read() + response = await client.post( + app.url_path_for("upload_image", image_path="image1.qcow2"), + content=image_data) + 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: