(superset) 01/02: Verify image magic number

[di...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

diegopucci pushed a commit to branch geido/fix/dashboard-pdf-download
in repository https://gitbox.apache.org/repos/asf/superset.git

commit 8c646bc9326f70598c5462e2bee8622facd4f828
Author: Diego Pucci <di...@gmail.com>
AuthorDate: Mon Apr 29 15:39:50 2024 +0300

    Verify image magic number
---
 superset/initialization/__init__.py |  2 +-
 superset/views/img_proxy.py         | 95 ++++++++++++++++++++++++-------------
 2 files changed, 63 insertions(+), 34 deletions(-)

diff --git a/superset/initialization/__init__.py b/superset/initialization/__init__.py
index af53db2826..318278e5d1 100644
--- a/superset/initialization/__init__.py
+++ b/superset/initialization/__init__.py
@@ -315,7 +315,7 @@ class SupersetAppInitializer:  # pylint: disable=too-many-public-methods
         appbuilder.add_view_no_menu(TaggedObjectsModelView)
         appbuilder.add_view_no_menu(TagView)
         appbuilder.add_view_no_menu(ReportView)
-        appbuilder.add_view_no_menu(ImgProxyView())
+        appbuilder.add_view_no_menu(ImgProxyView)
 
         #
         # Add links
diff --git a/superset/views/img_proxy.py b/superset/views/img_proxy.py
index 58ade4246a..3d23f8d6d0 100644
--- a/superset/views/img_proxy.py
+++ b/superset/views/img_proxy.py
@@ -1,64 +1,93 @@
-from typing import Any, Dict
-import requests
+from io import BytesIO
+from typing import Any
 from urllib.parse import urlparse
-from flask import g, request, Response
+
+import requests
+from flask import abort, request, Response
 from flask_appbuilder.api import expose
+from flask_wtf.csrf import same_origin
+from PIL import Image
+
 from superset import event_logger
-from superset.utils.core import (
-    get_user_id,
-)
+from superset.superset_typing import FlaskResponse
+from superset.utils.core import get_user_id
+
 from .base import BaseSupersetView
 
+
 class ImgProxyView(BaseSupersetView):
     route_base = "/img_proxy"
 
     @expose("/")
     @event_logger.log_this
-    def img_proxy(self) -> Response:
+    def img_proxy(self) -> FlaskResponse:
         """
         Proxy to an external URL, to overcome CORS restrictions.
         Returns a HTTP response containing the resource fetched from the external URL.
         """
         if not get_user_id():
-            raise Exception("User context not found")
-        
-        url = request.args.get('url')
+            abort(403)
+
+        url = request.args.get("url")
+        allowed_content_types = ["image/"]
 
         if not url:
-            return Response("URL parameter 'url' is missing", status=400)
+            abort(400)
 
         parsed_url = urlparse(url)
-        if parsed_url.scheme not in ['http', 'https']:
-            return Response("Invalid URL scheme", status=400)
-        
+        if parsed_url.scheme not in ["http", "https"]:
+            abort(400)
+
+        if not request.referrer or (
+            request.referrer and not same_origin(request.referrer, request.url_root)
+        ):
+            abort(403)
+
         try:
             response = self.fetch_resource(url)
         except Exception:
-            return Response("Error fetching resource", status=500)
+            abort(500)
+
+        content_type = response.headers.get("content-type", "")
+
+        if not any(
+            content_type.startswith(content_type_prefix)
+            for content_type_prefix in allowed_content_types
+        ):
+            abort(400)
+
+        try:
+            is_valid_image = self.validate_image(response.content)
+            if not is_valid_image:
+                abort(400)
+        except Exception:
+            abort(500)
 
-        return self.build_response(response)
+        headers: dict[str, Any] = {
+            key: value for (key, value) in response.headers.items()
+        }
 
-    def fetch_resource(self, url: str) -> Any:
+        return Response(response.content, response.status_code, headers)
+
+    def validate_image(self, image_content) -> bool:
+        """
+        Load the image from bytes, Pillow will raise an IOError if the file is not an image
+        """
+        try:
+            with Image.open(BytesIO(image_content)) as img:
+                if img.format:
+                    return True
+            return False
+        except OSError as e:
+            raise e
+        except Exception as e:
+            raise e
+
+    def fetch_resource(self, url: str) -> requests.Response:
         """Fetch the resource from the external server and handle errors."""
         try:
             response = requests.get(url)
             response.raise_for_status()
-
             return response
         except requests.RequestException as e:
             raise e
-
-    def build_response(self, response: requests.Response) -> Response:
-        """Build the HTTP response to return based on the fetched resource."""
-        allowed_content_types = ['image/']
-        content_type = response.headers.get('content-type', '')
-        
-        if not any(content_type.startswith(content_type_prefix) for content_type_prefix in allowed_content_types):
-            return Response("Response is not an allowed resource type", status=400)
-
-        headers: Dict[str, Any] = {key: value for (key, value) in response.headers.items()}
-        
-        return Response(response.content, response.status_code, headers)
-        
-        
-