diff --git a/plotting-service/plotting_service/plotting_api.py b/plotting-service/plotting_service/plotting_api.py index 555136a..84fec70 100644 --- a/plotting-service/plotting_service/plotting_api.py +++ b/plotting-service/plotting_service/plotting_api.py @@ -10,6 +10,7 @@ from fastapi import FastAPI, HTTPException from h5grove.fastapi_utils import router, settings # type: ignore from starlette.middleware.cors import CORSMiddleware +from starlette.middleware.gzip import GZipMiddleware from starlette.requests import Request from plotting_service.auth import get_experiments_for_user, get_user_from_token @@ -48,6 +49,7 @@ allow_methods=["*"], allow_headers=["*"], ) +app.add_middleware(GZipMiddleware, minimum_size=1000) CEPH_DIR = os.environ.get("CEPH_DIR", "/ceph") logger.info("Setting ceph directory to %s", CEPH_DIR) diff --git a/plotting-service/plotting_service/routers/imat.py b/plotting-service/plotting_service/routers/imat.py index 4210561..af0ba7d 100644 --- a/plotting-service/plotting_service/routers/imat.py +++ b/plotting-service/plotting_service/routers/imat.py @@ -7,16 +7,20 @@ from pathlib import Path from fastapi import APIRouter, HTTPException, Query -from starlette.responses import JSONResponse +from PIL import Image +from starlette.responses import JSONResponse, Response from plotting_service.services.image_service import ( + IMAGE_SUFFIXES, convert_image_to_rgb_array, find_latest_image_in_directory, ) +from plotting_service.utils import safe_check_filepath ImatRouter = APIRouter() IMAT_DIR: Path = Path(os.getenv("IMAT_DIR", "/imat")).resolve() +CEPH_DIR = os.environ.get("CEPH_DIR", "/ceph") stdout_handler = logging.StreamHandler(stream=sys.stdout) logging.basicConfig( @@ -87,3 +91,83 @@ async def get_latest_imat_image( "downsampleFactor": effective_downsample, } return JSONResponse(payload) + + +@ImatRouter.get("/imat/list-images", summary="List images in a directory") +async def list_imat_images( + path: typing.Annotated[ + str, Query(..., description="Path to the directory containing images, relative to CEPH_DIR") + ], +) -> list[str]: + """Return a sorted list of TIFF images in the given directory.""" + + dir_path = (Path(CEPH_DIR) / path).resolve() + # Security: Ensure path is within CEPH_DIR + try: + safe_check_filepath(dir_path, CEPH_DIR) + except FileNotFoundError as err: + raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Directory not found") from err + + if not dir_path.exists() or not dir_path.is_dir(): + raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Directory not found") + + images = [entry.name for entry in dir_path.iterdir() if entry.is_file() and entry.suffix.lower() in IMAGE_SUFFIXES] + + return sorted(images) + + +@ImatRouter.get("/imat/image", summary="Fetch a specific TIFF image as raw data") +async def get_imat_image( + path: typing.Annotated[str, Query(..., description="Path to the TIFF image file, relative to CEPH_DIR")], + downsample_factor: typing.Annotated[ + int, + Query( + ge=1, + le=64, + description="Integer factor to reduce each dimension by (1 keeps original resolution).", + ), + ] = 1, +) -> Response: + """Return the raw data of a specific TIFF image as binary.""" + + image_path = (Path(CEPH_DIR) / path).resolve() + # Security: Ensure path is within CEPH_DIR + try: + safe_check_filepath(image_path, CEPH_DIR) + except FileNotFoundError as err: + raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Directory not found") from err + + if not image_path.exists() or not image_path.is_file(): + raise HTTPException(status_code=HTTPStatus.NOT_FOUND, detail="Image not found") + + try: + with Image.open(image_path) as img: + original_width, original_height = img.size + + if downsample_factor > 1: + target_width = max(1, round(original_width / downsample_factor)) + target_height = max(1, round(original_height / downsample_factor)) + display_img = img.resize((target_width, target_height), Image.Resampling.NEAREST) + else: + display_img = img + + sampled_width, sampled_height = display_img.size + # For 16-bit TIFFs, tobytes() returns raw 16-bit bytes + data_bytes = display_img.tobytes() + + headers = { + "X-Image-Width": str(sampled_width), + "X-Image-Height": str(sampled_height), + "X-Original-Width": str(original_width), + "X-Original-Height": str(original_height), + "X-Downsample-Factor": str(downsample_factor), + "Access-Control-Expose-Headers": ( + "X-Image-Width, X-Image-Height, X-Original-Width, X-Original-Height, X-Downsample-Factor" + ), + } + + return Response(content=data_bytes, media_type="application/octet-stream", headers=headers) + + except Exception as exc: + logger.error(f"Failed to process image {image_path}: {exc}") + raise HTTPException(HTTPStatus.INTERNAL_SERVER_ERROR, "Unable to process image") from exc diff --git a/plotting-service/pyproject.toml b/plotting-service/pyproject.toml index 89428d3..4c27ee6 100644 --- a/plotting-service/pyproject.toml +++ b/plotting-service/pyproject.toml @@ -18,7 +18,7 @@ dependencies = [ [project.optional-dependencies] formatting = [ - "ruff==0.12.3", + "ruff==0.15.0", "mypy==1.17.0", "plotting-service[test]", "types-requests==2.32.4.20250611", diff --git a/plotting-service/test/test_plotting_api.py b/plotting-service/test/test_plotting_api.py index 6b918fd..2e50477 100644 --- a/plotting-service/test/test_plotting_api.py +++ b/plotting-service/test/test_plotting_api.py @@ -204,3 +204,150 @@ def test_get_latest_imat_image_with_mock_rb_folder(tmp_path, monkeypatch): assert payload["shape"] == [2, 4, 3] assert payload["downsampleFactor"] == 2 # noqa: PLR2004 assert payload["data"] == expected_bytes + + +def test_list_imat_images(tmp_path, monkeypatch): + """Verify that /imat/list-images correctly filters and sorts image files.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + + # Create test directory + data_dir = tmp_path / "test_data" + data_dir.mkdir() + (data_dir / "image2.tiff").touch() + (data_dir / "image1.tif").touch() + (data_dir / "not_an_image.txt").touch() + + client = TestClient(plotting_api.app) + response = client.get("/imat/list-images", params={"path": "test_data"}, headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.OK + assert response.json() == ["image1.tif", "image2.tiff"] + + +def test_list_imat_images_not_found(tmp_path, monkeypatch): + """Ensure 404 is returned when the requested directory does not exist.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + + client = TestClient(plotting_api.app) + response = client.get("/imat/list-images", params={"path": "non_existent"}, headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.NOT_FOUND + + +def test_list_imat_images_forbidden(tmp_path, monkeypatch): + """Verify that path traversal attempts are blocked with 403 Forbidden.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + + client = TestClient(plotting_api.app) + # safe_check_filepath should block this + response = client.get("/imat/list-images", params={"path": "../.."}, headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.FORBIDDEN + + +def test_get_imat_image(tmp_path, monkeypatch): + """Ensure /imat/image returns raw binary data and correct metadata headers.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + + # Create a tiny 16-bit TIFF (4x4, all 1000) + image_path = tmp_path / "test.tif" + image = Image.new("I;16", (4, 4), color=1000) + image.save(image_path, format="TIFF") + image.close() + + client = TestClient(plotting_api.app) + response = client.get("/imat/image", params={"path": "test.tif"}, headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.OK + assert response.headers["X-Image-Width"] == "4" + assert response.headers["X-Image-Height"] == "4" + assert response.headers["X-Original-Width"] == "4" + assert response.headers["X-Original-Height"] == "4" + assert response.headers["X-Downsample-Factor"] == "1" + assert response.content == Image.open(image_path).tobytes() + assert response.headers["Content-Type"] == "application/octet-stream" + + +def test_get_imat_image_downsampled(tmp_path, monkeypatch): + """Verify that downsampling works and headers reflect the sampled dimensions.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + + image_path = tmp_path / "test.tif" + image = Image.new("I;16", (8, 4), color=1000) + image.save(image_path, format="TIFF") + image.close() + + client = TestClient(plotting_api.app) + response = client.get( + "/imat/image", params={"path": "test.tif", "downsample_factor": 2}, headers={"Authorization": "Bearer foo"} + ) + + assert response.status_code == HTTPStatus.OK + assert response.headers["X-Image-Width"] == "4" + assert response.headers["X-Image-Height"] == "2" + assert response.headers["X-Original-Width"] == "8" + assert response.headers["X-Original-Height"] == "4" + + +def test_get_latest_imat_image_no_rb_folders(tmp_path, monkeypatch): + """Ensure 404 is returned if no RB folders are present in the IMAT directory.""" + monkeypatch.setattr(imat, "IMAT_DIR", tmp_path) + + client = TestClient(plotting_api.app) + response = client.get("/imat/latest-image", headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.NOT_FOUND + assert "No RB folders" in response.json()["detail"] + + +def test_get_imat_image_not_found(tmp_path, monkeypatch): + """Ensure /imat/image returns 404 for a missing file.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + + client = TestClient(plotting_api.app) + response = client.get("/imat/image", params={"path": "not_there.tif"}, headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.NOT_FOUND + + +def test_get_latest_imat_image_no_images_in_rb(tmp_path, monkeypatch): + """Ensure 404 is returned if RB folders exist but contain no valid image files.""" + monkeypatch.setattr(imat, "IMAT_DIR", tmp_path) + (tmp_path / "RB1234").mkdir() + + client = TestClient(plotting_api.app) + response = client.get("/imat/latest-image", headers={"Authorization": "Bearer foo"}) + assert response.status_code == HTTPStatus.NOT_FOUND + assert "No images found" in response.json()["detail"] + + +def test_get_imat_image_internal_error(tmp_path, monkeypatch): + """Verify that an exception during image processing returns a 500 error.""" + monkeypatch.setattr(imat, "CEPH_DIR", str(tmp_path)) + (tmp_path / "corrupt.tif").touch() + + client = TestClient(plotting_api.app) + with mock.patch("PIL.Image.open", side_effect=Exception("Simulated failure")): + response = client.get("/imat/image", params={"path": "corrupt.tif"}, headers={"Authorization": "Bearer foo"}) + + assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert "Unable to process image" in response.json()["detail"] + + +def test_get_latest_imat_image_conversion_error(tmp_path, monkeypatch): + """Verify that an error during RB latest image conversion returns a 500 error.""" + monkeypatch.setattr(imat, "IMAT_DIR", tmp_path) + rb_dir = tmp_path / "RB1234" + rb_dir.mkdir() + (rb_dir / "test.tif").touch() + + client = TestClient(plotting_api.app) + with mock.patch( + "plotting_service.routers.imat.convert_image_to_rgb_array", side_effect=Exception("Conversion failed") + ): + response = client.get( + "/imat/latest-image", params={"downsample_factor": 1}, headers={"Authorization": "Bearer foo"} + ) + + assert response.status_code == HTTPStatus.INTERNAL_SERVER_ERROR + assert "Unable to convert IMAT image" in response.json()["detail"]