Merge pull request #2629 from GNS3/image-handling

Fix image handling
This commit is contained in:
Jeremy Grossmann 2026-03-05 12:43:47 +08:00 committed by GitHub
commit ea68d8b959
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
5 changed files with 59 additions and 49 deletions

View File

@ -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,

View File

@ -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.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:
"""

View File

@ -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

View File

@ -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,

View File

@ -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: