浏览代码

[filebrowser] Improve file extension restrictions check for old upload operations (#4184)

- This commit improves file extension restrictions for uploads, improves code readability, and adds comprehensive unit tests for validation.
- The most significant changes include adding configuration options for allowed and restricted file extensions, implementing a utility function to enforce these restrictions, and creating extensive tests to ensure correct behavior.
- Additionally, minor code cleanup and refactoring were performed across multiple files.
Harsh Gupta 4 月之前
父节点
当前提交
e53562ff81

+ 11 - 3
apps/filebrowser/src/filebrowser/conf.py

@@ -18,7 +18,7 @@
 from django.utils.translation import gettext_lazy as _
 from django.utils.translation import gettext_lazy as _
 
 
 from desktop.conf import ENABLE_DOWNLOAD, is_oozie_enabled
 from desktop.conf import ENABLE_DOWNLOAD, is_oozie_enabled
-from desktop.lib.conf import Config, coerce_bool, coerce_csv
+from desktop.lib.conf import coerce_bool, coerce_csv, Config
 
 
 MAX_SNAPPY_DECOMPRESSION_SIZE = Config(
 MAX_SNAPPY_DECOMPRESSION_SIZE = Config(
   key="max_snappy_decompression_size", help=_("Max snappy decompression size in bytes."), private=True, default=1024 * 1024 * 25, type=int
   key="max_snappy_decompression_size", help=_("Max snappy decompression size in bytes."), private=True, default=1024 * 1024 * 25, type=int
@@ -104,10 +104,18 @@ FILE_DOWNLOAD_CACHE_CONTROL = Config(
 )
 )
 
 
 RESTRICT_FILE_EXTENSIONS = Config(
 RESTRICT_FILE_EXTENSIONS = Config(
-  key='restrict_file_extensions',
+  key="restrict_file_extensions",
+  default=None,
+  type=coerce_csv,
+  help=_("Specify file extensions that are not allowed, separated by commas. For example: .exe, .zip, .rar, .tar, .gz"),
+)
+
+ALLOW_FILE_EXTENSIONS = Config(
+  key="allow_file_extensions",
   default=None,
   default=None,
   type=coerce_csv,
   type=coerce_csv,
   help=_(
   help=_(
-    'Specify file extensions that are not allowed, separated by commas. For example: .exe, .zip, .rar, .tar, .gz'
+    "Specify file extensions that are allowed, separated by commas. "
+    "When set, only these extensions will be permitted. For example: .tsv, .csv, .xlsx"
   ),
   ),
 )
 )

+ 40 - 2
apps/filebrowser/src/filebrowser/utils.py

@@ -14,8 +14,8 @@
 # See the License for the specific language governing permissions and
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # limitations under the License.
 import io
 import io
-import os
 import logging
 import logging
+import os
 from datetime import datetime
 from datetime import datetime
 from urllib.parse import urlparse
 from urllib.parse import urlparse
 
 
@@ -23,7 +23,7 @@ import redis
 
 
 from desktop.conf import TASK_SERVER_V2
 from desktop.conf import TASK_SERVER_V2
 from desktop.lib.django_util import JsonResponse
 from desktop.lib.django_util import JsonResponse
-from filebrowser.conf import ARCHIVE_UPLOAD_TEMPDIR
+from filebrowser.conf import ALLOW_FILE_EXTENSIONS, ARCHIVE_UPLOAD_TEMPDIR, RESTRICT_FILE_EXTENSIONS
 
 
 LOG = logging.getLogger()
 LOG = logging.getLogger()
 
 
@@ -130,3 +130,41 @@ def release_reserved_space_for_file_uploads(uuid):
     LOG.exception("Failed to release reserved space: %s", str(e))
     LOG.exception("Failed to release reserved space: %s", str(e))
   finally:
   finally:
     redis_client.close()
     redis_client.close()
+
+
+def is_file_upload_allowed(file_name):
+  """
+  Check if a file upload is allowed based on file extension restrictions.
+
+  Args:
+    file_name: The name of the file being uploaded
+
+  Returns:
+    tuple: (is_allowed, error_message)
+      - is_allowed: Boolean indicating if the file upload is allowed
+      - error_message: String with error message if not allowed, None otherwise
+  """
+  if not file_name:
+    return True, None
+
+  _, file_type = os.path.splitext(file_name)
+  if file_type:
+    file_type = file_type.lower()
+
+  # Check allow list first - if set, only these extensions are allowed
+  allow_list = ALLOW_FILE_EXTENSIONS.get()
+  if allow_list:
+    # Normalize extensions to lowercase with dots
+    normalized_allow_list = [ext.lower() if ext.startswith(".") else f".{ext.lower()}" for ext in allow_list]
+    if file_type not in normalized_allow_list:
+      return False, f'File type "{file_type}" is not permitted. Modify file extension settings to allow this type.'
+
+  # Check restrict list - if set, these extensions are not allowed
+  restrict_list = RESTRICT_FILE_EXTENSIONS.get()
+  if restrict_list:
+    # Normalize extensions to lowercase with dots
+    normalized_restrict_list = [ext.lower() if ext.startswith(".") else f".{ext.lower()}" for ext in restrict_list]
+    if file_type in normalized_restrict_list:
+      return False, f'File type "{file_type}" is restricted. Update file extension restrictions to allow this type.'
+
+  return True, None

+ 291 - 0
apps/filebrowser/src/filebrowser/utils_test.py

@@ -0,0 +1,291 @@
+#!/usr/bin/env python
+# Licensed to Cloudera, Inc. under one
+# or more contributor license agreements.  See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership.  Cloudera, Inc. licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License.  You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from filebrowser.conf import ALLOW_FILE_EXTENSIONS, RESTRICT_FILE_EXTENSIONS
+from filebrowser.utils import is_file_upload_allowed
+
+
+class TestIsFileUploadAllowed:
+  def test_no_file_name(self):
+    # Test with None
+    is_allowed, error_msg = is_file_upload_allowed(None)
+    assert is_allowed is True
+    assert error_msg is None
+
+    # Test with empty string
+    is_allowed, error_msg = is_file_upload_allowed("")
+    assert is_allowed is True
+    assert error_msg is None
+
+  def test_no_restrictions_configured(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing(None)
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing(None)
+
+    try:
+      # All file types should be allowed
+      test_files = ["document.pdf", "script.exe", "archive.zip", "data.csv", "image.png", "video.mp4"]
+
+      for file_name in test_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is True, f"File '{file_name}' should be allowed when no restrictions are configured"
+        assert error_msg is None
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_allow_list_with_dots(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".csv", ".txt", ".json"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing(None)
+
+    try:
+      # Allowed files
+      allowed_files = ["data.csv", "notes.txt", "config.json", "DATA.CSV", "NOTES.TXT"]
+      for file_name in allowed_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is True, f"File '{file_name}' should be allowed"
+        assert error_msg is None
+
+      # Not allowed files
+      not_allowed_files = ["script.exe", "archive.zip", "image.png"]
+      for file_name in not_allowed_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is False, f"File '{file_name}' should not be allowed"
+        assert error_msg is not None
+        assert "is not permitted" in error_msg
+        assert "Modify file extension settings" in error_msg
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_allow_list_without_dots(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing(["csv", "txt", "json"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing(None)
+
+    try:
+      # Should still work - extensions are normalized
+      allowed_files = ["data.csv", "notes.txt", "config.json"]
+      for file_name in allowed_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is True, f"File '{file_name}' should be allowed"
+        assert error_msg is None
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_restrict_list_with_dots(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing(None)
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe", ".zip", ".rar"])
+
+    try:
+      # Restricted files
+      restricted_files = ["malware.exe", "archive.zip", "compressed.rar", "MALWARE.EXE"]
+      for file_name in restricted_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is False, f"File '{file_name}' should be restricted"
+        assert error_msg is not None
+        assert "is restricted" in error_msg
+        assert "Update file extension restrictions" in error_msg
+
+      # Allowed files
+      allowed_files = ["document.pdf", "data.csv", "image.png"]
+      for file_name in allowed_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is True, f"File '{file_name}' should be allowed"
+        assert error_msg is None
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_restrict_list_without_dots(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing(None)
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing(["exe", "zip", "rar"])
+
+    try:
+      # Should still work - extensions are normalized
+      restricted_files = ["malware.exe", "archive.zip", "compressed.rar"]
+      for file_name in restricted_files:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is False, f"File '{file_name}' should be restricted"
+        assert error_msg is not None
+        assert "is restricted" in error_msg
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_both_allow_and_restrict_lists(self):
+    # Allow list takes precedence - if file type is not in allow list, it's rejected
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".csv", ".txt", ".exe"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe", ".zip"])
+
+    try:
+      # File in allow list but also in restrict list - should check restrict list
+      is_allowed, error_msg = is_file_upload_allowed("script.exe")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is restricted" in error_msg
+
+      # File in allow list and not in restrict list - should be allowed
+      is_allowed, error_msg = is_file_upload_allowed("data.csv")
+      assert is_allowed is True
+      assert error_msg is None
+
+      # File not in allow list - should be rejected regardless of restrict list
+      is_allowed, error_msg = is_file_upload_allowed("image.png")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is not permitted" in error_msg
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_case_insensitive_extensions(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".CSV", ".TXT"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".EXE", ".ZIP"])
+
+    try:
+      # Test allow list with different cases
+      test_cases = [
+        ("data.csv", True),
+        ("data.CSV", True),
+        ("data.CsV", True),
+        ("notes.txt", True),
+        ("notes.TXT", True),
+        ("notes.TxT", True),
+      ]
+
+      for file_name, expected in test_cases:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is expected, f"File '{file_name}' case handling failed"
+
+      # Reset for restrict list test
+      reset_allow()
+      reset_restrict()
+
+      reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing(None)
+      reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".EXE", ".ZIP"])
+
+      # Test restrict list with different cases
+      restricted_cases = [
+        ("malware.exe", False),
+        ("malware.EXE", False),
+        ("malware.ExE", False),
+        ("archive.zip", False),
+        ("archive.ZIP", False),
+        ("archive.ZiP", False),
+      ]
+
+      for file_name, expected in restricted_cases:
+        is_allowed, error_msg = is_file_upload_allowed(file_name)
+        assert is_allowed is expected, f"File '{file_name}' case handling failed"
+        if not expected:
+          assert error_msg is not None
+          assert "is restricted" in error_msg
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_files_without_extensions(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".txt", ".csv"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe"])
+
+    try:
+      # File without extension with allow list - should not be in allow list
+      is_allowed, error_msg = is_file_upload_allowed("README")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is not permitted" in error_msg
+
+      # Reset for restrict list only test
+      reset_allow()
+      reset_restrict()
+
+      reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing(None)
+      reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe"])
+
+      # File without extension with only restrict list - should be allowed
+      is_allowed, error_msg = is_file_upload_allowed("README")
+      assert is_allowed is True
+      assert error_msg is None
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_files_with_multiple_dots(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".gz", ".txt"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe"])
+
+    try:
+      # Should use the last extension
+      is_allowed, error_msg = is_file_upload_allowed("archive.tar.gz")
+      assert is_allowed is True
+      assert error_msg is None
+
+      is_allowed, error_msg = is_file_upload_allowed("document.backup.txt")
+      assert is_allowed is True
+      assert error_msg is None
+
+      is_allowed, error_msg = is_file_upload_allowed("file.backup.exe")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is not permitted" in error_msg
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_hidden_files(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".txt", ".conf"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe"])
+
+    try:
+      # Hidden file with extension
+      is_allowed, error_msg = is_file_upload_allowed(".bashrc.txt")
+      assert is_allowed is True
+      assert error_msg is None
+
+      # Hidden file without what looks like an extension
+      is_allowed, error_msg = is_file_upload_allowed(".bashrc")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is not permitted" in error_msg
+    finally:
+      reset_allow()
+      reset_restrict()
+
+  def test_edge_cases(self):
+    reset_allow = ALLOW_FILE_EXTENSIONS.set_for_testing([".txt"])
+    reset_restrict = RESTRICT_FILE_EXTENSIONS.set_for_testing([".exe"])
+
+    try:
+      # File ending with dot
+      is_allowed, error_msg = is_file_upload_allowed("file.")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is not permitted" in error_msg
+
+      # Just a dot
+      is_allowed, error_msg = is_file_upload_allowed(".")
+      assert is_allowed is False
+      assert error_msg is not None
+      assert "is not permitted" in error_msg
+
+      # Multiple consecutive dots
+      is_allowed, error_msg = is_file_upload_allowed("file..txt")
+      assert is_allowed is True
+      assert error_msg is None
+    finally:
+      reset_allow()
+      reset_restrict()

+ 36 - 36
apps/filebrowser/src/filebrowser/views.py

@@ -15,20 +15,16 @@
 # See the License for the specific language governing permissions and
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # limitations under the License.
 
 
-import io
-import os
-import re
-import sys
-import json
-import stat as stat_module
 import errno
 import errno
+import json
 import logging
 import logging
-import operator
 import mimetypes
 import mimetypes
+import operator
+import os
 import posixpath
 import posixpath
-import urllib.error
-import urllib.request
-from builtins import object, str as new_str
+import re
+import stat as stat_module
+from builtins import str as new_str
 from bz2 import decompress
 from bz2 import decompress
 from datetime import datetime
 from datetime import datetime
 from functools import partial
 from functools import partial
@@ -37,29 +33,28 @@ from io import BytesIO, StringIO as string_io
 from urllib.parse import quote as urllib_quote, unquote as urllib_unquote, urlparse as lib_urlparse
 from urllib.parse import quote as urllib_quote, unquote as urllib_unquote, urlparse as lib_urlparse
 
 
 import pandas as pd
 import pandas as pd
-from avro import datafile, io
-from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload
-from django.core.paginator import EmptyPage, InvalidPage, Page, Paginator
-from django.http import Http404, HttpResponse, HttpResponseForbidden, HttpResponseNotModified, HttpResponseRedirect, StreamingHttpResponse
+from avro import datafile, io as avro_io
+from django.core.files.uploadhandler import StopUpload
+from django.core.paginator import EmptyPage, InvalidPage, Paginator
+from django.http import Http404, HttpResponse, HttpResponseNotModified, HttpResponseRedirect, StreamingHttpResponse
 from django.shortcuts import redirect
 from django.shortcuts import redirect
 from django.template.defaultfilters import filesizeformat, stringformat
 from django.template.defaultfilters import filesizeformat, stringformat
 from django.urls import reverse
 from django.urls import reverse
 from django.utils.html import escape
 from django.utils.html import escape
 from django.utils.http import http_date
 from django.utils.http import http_date
 from django.utils.translation import gettext as _
 from django.utils.translation import gettext as _
-from django.views.decorators.csrf import csrf_exempt
 from django.views.decorators.http import require_http_methods
 from django.views.decorators.http import require_http_methods
 from django.views.static import was_modified_since
 from django.views.static import was_modified_since
 
 
-from aws.s3.s3fs import S3FileSystemException, S3ListAllBucketsException, get_s3_home_directory
+from aws.s3.s3fs import get_s3_home_directory, S3FileSystemException, S3ListAllBucketsException
 from aws.s3.upload import S3FineUploaderChunkedUpload
 from aws.s3.upload import S3FineUploaderChunkedUpload
 from azure.abfs.upload import ABFSFineUploaderChunkedUpload
 from azure.abfs.upload import ABFSFineUploaderChunkedUpload
 from desktop import appmanager
 from desktop import appmanager
 from desktop.auth.backend import is_admin
 from desktop.auth.backend import is_admin
-from desktop.conf import ENABLE_NEW_STORAGE_BROWSER, RAZ, TASK_SERVER_V2
-from desktop.lib import fsmanager, i18n
+from desktop.conf import RAZ, TASK_SERVER_V2
+from desktop.lib import i18n
 from desktop.lib.conf import coerce_bool
 from desktop.lib.conf import coerce_bool
-from desktop.lib.django_util import JsonResponse, format_preserving_redirect, render
+from desktop.lib.django_util import format_preserving_redirect, JsonResponse, render
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.export_csvxls import file_reader
 from desktop.lib.export_csvxls import file_reader
 from desktop.lib.fs import splitpath
 from desktop.lib.fs import splitpath
@@ -94,17 +89,13 @@ from filebrowser.forms import (
   SetReplicationFactorForm,
   SetReplicationFactorForm,
   TouchForm,
   TouchForm,
   TrashPurgeForm,
   TrashPurgeForm,
-  UploadArchiveForm,
   UploadFileForm,
   UploadFileForm,
 )
 )
 from filebrowser.lib import xxd
 from filebrowser.lib import xxd
-from filebrowser.lib.archives import archive_factory
 from filebrowser.lib.rwx import filetype, rwx
 from filebrowser.lib.rwx import filetype, rwx
-from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.core_site import get_trash_interval
 from hadoop.core_site import get_trash_interval
 from hadoop.fs.exceptions import WebHdfsException
 from hadoop.fs.exceptions import WebHdfsException
 from hadoop.fs.fsutils import do_overwrite_save
 from hadoop.fs.fsutils import do_overwrite_save
-from hadoop.fs.hadoopfs import Hdfs
 from hadoop.fs.upload import HDFSFineUploaderChunkedUpload, LocalFineUploaderChunkedUpload
 from hadoop.fs.upload import HDFSFineUploaderChunkedUpload, LocalFineUploaderChunkedUpload
 from useradmin.models import Group, User
 from useradmin.models import Group, User
 
 
@@ -219,7 +210,6 @@ def download(request, path):
   content_type = mimetypes.guess_type(path)[0] or 'application/octet-stream'
   content_type = mimetypes.guess_type(path)[0] or 'application/octet-stream'
   stats = request.fs.stats(path)
   stats = request.fs.stats(path)
   mtime = stats['mtime']
   mtime = stats['mtime']
-  size = stats['size']
   if not was_modified_since(request.META.get('HTTP_IF_MODIFIED_SINCE'), mtime):
   if not was_modified_since(request.META.get('HTTP_IF_MODIFIED_SINCE'), mtime):
     return HttpResponseNotModified()
     return HttpResponseNotModified()
     # TODO(philip): Ideally a with statement would protect from leaks, but tricky to do here.
     # TODO(philip): Ideally a with statement would protect from leaks, but tricky to do here.
@@ -955,7 +945,7 @@ def _read_avro(fhandle, path, offset, length, stats):
   contents = ''
   contents = ''
   try:
   try:
     fhandle.seek(offset)
     fhandle.seek(offset)
-    data_file_reader = datafile.DataFileReader(fhandle, io.DatumReader())
+    data_file_reader = datafile.DataFileReader(fhandle, avro_io.DatumReader())
 
 
     try:
     try:
       contents_list = []
       contents_list = []
@@ -1072,11 +1062,11 @@ def detect_parquet(fhandle):
 def snappy_installed():
 def snappy_installed():
   '''Snappy is library that isn't supported by python2.4'''
   '''Snappy is library that isn't supported by python2.4'''
   try:
   try:
-    import snappy
+    import snappy  # noqa: F401
     return True
     return True
   except ImportError:
   except ImportError:
     return False
     return False
-  except Exception as e:
+  except Exception:
     logging.exception('failed to verify if snappy is installed')
     logging.exception('failed to verify if snappy is installed')
     return False
     return False
 
 
@@ -1122,12 +1112,12 @@ def formset_initial_value_extractor(request, parameter_names):
   The formsets should then handle construction on their own.
   The formsets should then handle construction on their own.
   """
   """
   def _intial_value_extractor(request):
   def _intial_value_extractor(request):
-    if not submitted:
+    if not submitted:  # noqa: F821
       return []
       return []
     # Build data with list of in order parameters receive in POST data
     # Build data with list of in order parameters receive in POST data
     # Size can be inferred from largest list returned in POST data
     # Size can be inferred from largest list returned in POST data
     data = []
     data = []
-    for param in submitted:
+    for param in submitted:  # noqa: F821
       i = 0
       i = 0
       for val in request.POST.getlist(param):
       for val in request.POST.getlist(param):
         if len(data) == i:
         if len(data) == i:
@@ -1136,7 +1126,7 @@ def formset_initial_value_extractor(request, parameter_names):
         i += 1
         i += 1
     # Extend every data object with recurring params
     # Extend every data object with recurring params
     for kwargs in data:
     for kwargs in data:
-      for recurrent in recurring:
+      for recurrent in recurring:  # noqa: F821
         kwargs[recurrent] = request.POST.get(recurrent)
         kwargs[recurrent] = request.POST.get(recurrent)
     initial_data = data
     initial_data = data
     return {'initial': initial_data}
     return {'initial': initial_data}
@@ -1254,7 +1244,7 @@ def generic_op(form_class, request, op, parameter_names, piggyback=None, templat
         if piggyback:
         if piggyback:
           piggy_path = form.cleaned_data.get(piggyback)
           piggy_path = form.cleaned_data.get(piggyback)
           ret["result"] = _massage_stats(request, stat_absolute_path(piggy_path, request.fs.stats(piggy_path)))
           ret["result"] = _massage_stats(request, stat_absolute_path(piggy_path, request.fs.stats(piggy_path)))
-      except Exception as e:
+      except Exception:
         # Hard to report these more naturally here.  These happen either
         # Hard to report these more naturally here.  These happen either
         # because of a bug in the piggy-back code or because of a
         # because of a bug in the piggy-back code or because of a
         # race condition.
         # race condition.
@@ -1538,11 +1528,11 @@ def upload_chunks(request):
     for _ in request.FILES.values():  # This processes the upload.
     for _ in request.FILES.values():  # This processes the upload.
       pass
       pass
   except StopUpload:
   except StopUpload:
-    return JsonResponse({'success': False, 'error': 'Error in upload'})
+    return JsonResponse({"success": False, "error": "Error in upload"})
 
 
   # case where file is larger than the single chunk size
   # case where file is larger than the single chunk size
   if int(request.GET.get("qqtotalparts", 0)) > 0:
   if int(request.GET.get("qqtotalparts", 0)) > 0:
-    return JsonResponse({'success': True, 'uuid': request.GET.get('qquuid')})
+    return JsonResponse({"success": True, "uuid": request.GET.get("qquuid")})
 
 
   # case where file is smaller than the chunk size
   # case where file is smaller than the chunk size
   if int(request.GET.get("qqtotalparts", 0)) == 0:
   if int(request.GET.get("qqtotalparts", 0)) == 0:
@@ -1550,9 +1540,14 @@ def upload_chunks(request):
     try:
     try:
       response = perform_upload_task(request, **chunks)
       response = perform_upload_task(request, **chunks)
       return JsonResponse(response)
       return JsonResponse(response)
+    except PopupException as e:
+      logger.error(f"Upload failed: {e}")
+      return JsonResponse({"success": False, "error": str(e)})
     except Exception as e:
     except Exception as e:
-      return JsonResponse({'success': False, 'error': 'Error in upload %s' % str(e)})
-  return JsonResponse({'success': False, 'error': 'Unsupported request method'})
+      # For unexpected exceptions, log the full error but return a generic message
+      logger.exception(f"Unexpected error during upload: {e}")
+      return JsonResponse({"success": False, "error": "Upload failed due to an unexpected error"})
+  return JsonResponse({"success": False, "error": "Unsupported request method"})
 
 
 
 
 @require_http_methods(["POST"])
 @require_http_methods(["POST"])
@@ -1568,8 +1563,13 @@ def upload_complete(request):
   try:
   try:
     response = perform_upload_task(request, **chunks)
     response = perform_upload_task(request, **chunks)
     return JsonResponse(response)
     return JsonResponse(response)
+  except PopupException as e:
+    logger.error(f"Upload failed: {e}")
+    return JsonResponse({"success": False, "error": str(e)})
   except Exception as e:
   except Exception as e:
-    return JsonResponse({'success': False, 'error': 'Error in upload'})
+    # For unexpected exceptions, log the full error but return a generic message
+    logger.exception(f"Unexpected error during upload: {e}")
+    return JsonResponse({"success": False, "error": "Upload failed due to an unexpected error"})
 
 
 
 
 @require_http_methods(["POST"])
 @require_http_methods(["POST"])

+ 3 - 0
desktop/conf.dist/hue.ini

@@ -1768,6 +1768,9 @@ submit_to=True
 # Specify file extensions that are not allowed, separated by commas.
 # Specify file extensions that are not allowed, separated by commas.
 ## restrict_file_extensions=.exe, .zip, .rar, .tar, .gz
 ## restrict_file_extensions=.exe, .zip, .rar, .tar, .gz
 
 
+# Specify file extensions that are allowed, separated by commas.
+## allow_file_extensions=.tsv, .csv, .xlsx
+
 ###########################################################################
 ###########################################################################
 # Settings to configure Pig
 # Settings to configure Pig
 ###########################################################################
 ###########################################################################

+ 2 - 0
desktop/conf/pseudo-distributed.ini.tmpl

@@ -1750,6 +1750,8 @@
   # Specify file extensions that are not allowed, separated by commas.
   # Specify file extensions that are not allowed, separated by commas.
   ## restrict_file_extensions=.exe, .zip, .rar, .tar, .gz
   ## restrict_file_extensions=.exe, .zip, .rar, .tar, .gz
 
 
+  # Specify file extensions that are allowed, separated by commas.
+  ## allow_file_extensions=.tsv, .csv, .xlsx
 
 
 ###########################################################################
 ###########################################################################
 # Settings to configure Pig
 # Settings to configure Pig

+ 12 - 6
desktop/core/src/desktop/lib/fs/gc/upload.py

@@ -21,7 +21,6 @@ Classes for a custom upload handler to stream into GS.
 See http://docs.djangoproject.com/en/1.9/topics/http/file-uploads/
 See http://docs.djangoproject.com/en/1.9/topics/http/file-uploads/
 """
 """
 
 
-import os
 import logging
 import logging
 from io import BytesIO as stream_io
 from io import BytesIO as stream_io
 
 
@@ -31,7 +30,7 @@ from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandler
 from desktop.lib.fs.gc import parse_uri
 from desktop.lib.fs.gc import parse_uri
 from desktop.lib.fs.gc.gs import GSFileSystemException
 from desktop.lib.fs.gc.gs import GSFileSystemException
 from desktop.lib.fsmanager import get_client
 from desktop.lib.fsmanager import get_client
-from filebrowser.conf import RESTRICT_FILE_EXTENSIONS
+from filebrowser.utils import is_file_upload_allowed
 
 
 LOG = logging.getLogger()
 LOG = logging.getLogger()
 
 
@@ -59,6 +58,7 @@ class GSFileUploadHandler(FileUploadHandler):
     self._request = request
     self._request = request
     self._mp = None
     self._mp = None
     self._part_num = 1
     self._part_num = 1
+    self._upload_rejected = False
 
 
     if self._is_gs_upload():
     if self._is_gs_upload():
       self._fs = get_client(fs='gs', user=request.user.username)
       self._fs = get_client(fs='gs', user=request.user.username)
@@ -76,11 +76,13 @@ class GSFileUploadHandler(FileUploadHandler):
     if self._is_gs_upload():
     if self._is_gs_upload():
       LOG.info('Using GSFileUploadHandler to handle file upload.')
       LOG.info('Using GSFileUploadHandler to handle file upload.')
 
 
-      _, file_type = os.path.splitext(file_name)
-      if RESTRICT_FILE_EXTENSIONS.get() and file_type.lower() in [ext.lower() for ext in RESTRICT_FILE_EXTENSIONS.get()]:
-        err_message = f'Uploading files with type "{file_type}" is not allowed. Hue is configured to restrict this type.'
+      # Check file extension restrictions
+      is_allowed, err_message = is_file_upload_allowed(file_name)
+      if not is_allowed:
         LOG.error(err_message)
         LOG.error(err_message)
-        raise Exception(err_message)
+        self._request.META['upload_failed'] = err_message
+        self._upload_rejected = True
+        return None
 
 
       super().new_file(field_name, file_name, *args, **kwargs)
       super().new_file(field_name, file_name, *args, **kwargs)
 
 
@@ -106,6 +108,8 @@ class GSFileUploadHandler(FileUploadHandler):
 
 
     This method is called for each data chunk received during the upload process.
     This method is called for each data chunk received during the upload process.
     """
     """
+    if self._upload_rejected:
+      return None
     if self._is_gs_upload():
     if self._is_gs_upload():
       try:
       try:
         LOG.debug("GSFileUploadHandler uploading file part: %d" % self._part_num)
         LOG.debug("GSFileUploadHandler uploading file part: %d" % self._part_num)
@@ -125,6 +129,8 @@ class GSFileUploadHandler(FileUploadHandler):
 
 
     This method is called when the entire file has been uploaded.
     This method is called when the entire file has been uploaded.
     """
     """
+    if self._upload_rejected:
+      return None
     if self._is_gs_upload():
     if self._is_gs_upload():
       LOG.info("GSFileUploadHandler has completed file upload to GS, total file size is: %d." % file_size)
       LOG.info("GSFileUploadHandler has completed file upload to GS, total file size is: %d." % file_size)
       self._mp.complete_upload()
       self._mp.complete_upload()

+ 23 - 11
desktop/core/src/desktop/lib/fs/ozone/upload.py

@@ -15,7 +15,6 @@
 # limitations under the License.
 # limitations under the License.
 
 
 import io
 import io
-import os
 import logging
 import logging
 import unicodedata
 import unicodedata
 
 
@@ -26,8 +25,7 @@ from django.utils.translation import gettext as _
 from desktop.conf import TASK_SERVER_V2
 from desktop.conf import TASK_SERVER_V2
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.fsmanager import get_client
 from desktop.lib.fsmanager import get_client
-from filebrowser.conf import RESTRICT_FILE_EXTENSIONS
-from filebrowser.utils import calculate_total_size, generate_chunks
+from filebrowser.utils import calculate_total_size, generate_chunks, is_file_upload_allowed
 from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.fs.exceptions import WebHdfsException
 from hadoop.fs.exceptions import WebHdfsException
 
 
@@ -52,13 +50,20 @@ class OFSFineUploaderChunkedUpload(object):
     self._part_size = UPLOAD_CHUNK_SIZE.get()
     self._part_size = UPLOAD_CHUNK_SIZE.get()
 
 
   def check_access(self):
   def check_access(self):
+    # Check file extension restrictions
+    is_allowed, err_message = is_file_upload_allowed(self.file_name)
+    if not is_allowed:
+      LOG.error(err_message)
+      self._request.META["upload_failed"] = err_message
+      raise PopupException(err_message)
+
     if self._is_ofs_upload():
     if self._is_ofs_upload():
       self._fs = self._get_ofs(self._request)
       self._fs = self._get_ofs(self._request)
 
 
       # Verify that the path exists
       # Verify that the path exists
       try:
       try:
         self._fs.stats(self.destination)
         self._fs.stats(self.destination)
-      except Exception as e:
+      except Exception:
         raise PopupException(_('Destination path does not exist: %s' % self.destination))
         raise PopupException(_('Destination path does not exist: %s' % self.destination))
 
 
       LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
       LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
@@ -82,7 +87,7 @@ class OFSFineUploaderChunkedUpload(object):
         # Verify that the path exists
         # Verify that the path exists
         try:
         try:
           self._fs.stats(self.destination)
           self._fs.stats(self.destination)
-        except Exception as e:
+        except Exception:
           raise PopupException(_('Destination path does not exist: %s' % self.destination))
           raise PopupException(_('Destination path does not exist: %s' % self.destination))
 
 
         LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
         LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
@@ -189,6 +194,7 @@ class OFSFileUploadHandler(FileUploadHandler):
     self.file = None
     self.file = None
     self._request = request
     self._request = request
     self._part_size = UPLOAD_CHUNK_SIZE.get()
     self._part_size = UPLOAD_CHUNK_SIZE.get()
+    self._upload_rejected = False
 
 
     if self._is_ofs_upload():
     if self._is_ofs_upload():
       self._fs = self._get_ofs(request)
       self._fs = self._get_ofs(request)
@@ -196,7 +202,7 @@ class OFSFileUploadHandler(FileUploadHandler):
       # Verify that the path exists
       # Verify that the path exists
       try:
       try:
         self._fs.stats(self.destination)
         self._fs.stats(self.destination)
-      except Exception as e:
+      except Exception:
         raise OFSFileUploadError(_('Destination path does not exist: %s' % self.destination))
         raise OFSFileUploadError(_('Destination path does not exist: %s' % self.destination))
 
 
     LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
     LOG.debug("Chunk size = %d" % UPLOAD_CHUNK_SIZE.get())
@@ -205,11 +211,13 @@ class OFSFileUploadHandler(FileUploadHandler):
     if self._is_ofs_upload():
     if self._is_ofs_upload():
       LOG.info('Using OFSFileUploadHandler to handle file upload.')
       LOG.info('Using OFSFileUploadHandler to handle file upload.')
 
 
-      _, file_type = os.path.splitext(file_name)
-      if RESTRICT_FILE_EXTENSIONS.get() and file_type.lower() in [ext.lower() for ext in RESTRICT_FILE_EXTENSIONS.get()]:
-        err_message = f'Uploading files with type "{file_type}" is not allowed. Hue is configured to restrict this type.'
+      # Check file extension restrictions
+      is_allowed, err_message = is_file_upload_allowed(file_name)
+      if not is_allowed:
         LOG.error(err_message)
         LOG.error(err_message)
-        raise Exception(err_message)
+        self._request.META['upload_failed'] = err_message
+        self._upload_rejected = True
+        return None
 
 
       super(OFSFileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
       super(OFSFileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
 
 
@@ -223,10 +231,12 @@ class OFSFileUploadHandler(FileUploadHandler):
         raise StopFutureHandlers()
         raise StopFutureHandlers()
       except (OFSFileUploadError, WebHdfsException) as e:
       except (OFSFileUploadError, WebHdfsException) as e:
         LOG.error("Encountered error in OFSUploadHandler check_access: %s" % e)
         LOG.error("Encountered error in OFSUploadHandler check_access: %s" % e)
-        self.request.META['upload_failed'] = e
+        self._request.META["upload_failed"] = e
         raise StopUpload()
         raise StopUpload()
 
 
   def receive_data_chunk(self, raw_data, start):
   def receive_data_chunk(self, raw_data, start):
+    if self._upload_rejected:
+      return None
     if self._is_ofs_upload():
     if self._is_ofs_upload():
       LOG.debug("OFSfileUploadHandler receive_data_chunk")
       LOG.debug("OFSfileUploadHandler receive_data_chunk")
       try:
       try:
@@ -240,6 +250,8 @@ class OFSFileUploadHandler(FileUploadHandler):
       return raw_data
       return raw_data
 
 
   def file_complete(self, file_size):
   def file_complete(self, file_size):
+    if self._upload_rejected:
+      return None
     if self._is_ofs_upload():
     if self._is_ofs_upload():
       # Finish the upload
       # Finish the upload
       LOG.info("OFSFileUploadHandler has completed file upload to OFS, total file size is: %d." % file_size)
       LOG.info("OFSFileUploadHandler has completed file upload to OFS, total file size is: %d." % file_size)

+ 22 - 10
desktop/libs/aws/src/aws/s3/upload.py

@@ -21,13 +21,12 @@ Classes for a custom upload handler to stream into S3.
 See http://docs.djangoproject.com/en/1.9/topics/http/file-uploads/
 See http://docs.djangoproject.com/en/1.9/topics/http/file-uploads/
 """
 """
 
 
-import os
 import logging
 import logging
 import unicodedata
 import unicodedata
 from io import BytesIO as stream_io
 from io import BytesIO as stream_io
 
 
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.core.files.uploadedfile import SimpleUploadedFile
-from django.core.files.uploadhandler import FileUploadHandler, SkipFile, StopFutureHandlers, StopUpload, UploadFileException
+from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload, UploadFileException
 from django.utils.translation import gettext as _
 from django.utils.translation import gettext as _
 
 
 from aws.s3 import parse_uri
 from aws.s3 import parse_uri
@@ -35,8 +34,7 @@ from aws.s3.s3fs import S3FileSystemException
 from desktop.conf import TASK_SERVER_V2
 from desktop.conf import TASK_SERVER_V2
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.fsmanager import get_client
 from desktop.lib.fsmanager import get_client
-from filebrowser.conf import RESTRICT_FILE_EXTENSIONS
-from filebrowser.utils import calculate_total_size, generate_chunks
+from filebrowser.utils import calculate_total_size, generate_chunks, is_file_upload_allowed
 
 
 DEFAULT_WRITE_SIZE = 1024 * 1024 * 128  # TODO: set in configuration (currently 128 MiB)
 DEFAULT_WRITE_SIZE = 1024 * 1024 * 128  # TODO: set in configuration (currently 128 MiB)
 
 
@@ -64,6 +62,13 @@ class S3FineUploaderChunkedUpload(object):
       self.chunk_size = kwargs.get('chunk_size')
       self.chunk_size = kwargs.get('chunk_size')
 
 
   def check_access(self):
   def check_access(self):
+    # Check file extension restrictions
+    is_allowed, err_message = is_file_upload_allowed(self.file_name)
+    if not is_allowed:
+      LOG.error(err_message)
+      self._request.META["upload_failed"] = err_message
+      raise PopupException(err_message)
+
     if self._is_s3_upload():
     if self._is_s3_upload():
       try:
       try:
         # Check access permissions before attempting upload
         # Check access permissions before attempting upload
@@ -73,7 +78,7 @@ class S3FineUploaderChunkedUpload(object):
         self._mp = self._bucket.initiate_multipart_upload(self.filepath)
         self._mp = self._bucket.initiate_multipart_upload(self.filepath)
       except (S3FileUploadError, S3FileSystemException) as e:
       except (S3FileUploadError, S3FileSystemException) as e:
         LOG.error("S3FineUploaderChunkedUpload: Encountered error in S3UploadHandler check_access: %s" % e)
         LOG.error("S3FineUploaderChunkedUpload: Encountered error in S3UploadHandler check_access: %s" % e)
-        self.request.META['upload_failed'] = e
+        self._request.META['upload_failed'] = e
         raise PopupException("S3FineUploaderChunkedUpload: Initiating S3 multipart upload to target path: %s failed" % self.filepath)
         raise PopupException("S3FineUploaderChunkedUpload: Initiating S3 multipart upload to target path: %s failed" % self.filepath)
 
 
     self.chunk_size = DEFAULT_WRITE_SIZE
     self.chunk_size = DEFAULT_WRITE_SIZE
@@ -89,7 +94,7 @@ class S3FineUploaderChunkedUpload(object):
         self._mp = self._bucket.initiate_multipart_upload(self.filepath)
         self._mp = self._bucket.initiate_multipart_upload(self.filepath)
       except (S3FileUploadError, S3FileSystemException) as e:
       except (S3FileUploadError, S3FileSystemException) as e:
         LOG.error("S3FineUploaderChunkedUpload: Encountered error in S3UploadHandler check_access: %s" % e)
         LOG.error("S3FineUploaderChunkedUpload: Encountered error in S3UploadHandler check_access: %s" % e)
-        self.request.META['upload_failed'] = e
+        self._request.META['upload_failed'] = e
         raise PopupException("S3FineUploaderChunkedUpload: Initiating S3 multipart upload to target path: %s failed" % self.filepath)
         raise PopupException("S3FineUploaderChunkedUpload: Initiating S3 multipart upload to target path: %s failed" % self.filepath)
 
 
     try:
     try:
@@ -148,6 +153,7 @@ class S3FileUploadHandler(FileUploadHandler):
     self._request = request
     self._request = request
     self._mp = None
     self._mp = None
     self._part_num = 1
     self._part_num = 1
+    self._upload_rejected = False
 
 
     if self._is_s3_upload():
     if self._is_s3_upload():
       self._fs = get_client(fs='s3a', user=request.user.username)
       self._fs = get_client(fs='s3a', user=request.user.username)
@@ -160,11 +166,13 @@ class S3FileUploadHandler(FileUploadHandler):
     if self._is_s3_upload():
     if self._is_s3_upload():
       LOG.info('Using S3FileUploadHandler to handle file upload.')
       LOG.info('Using S3FileUploadHandler to handle file upload.')
 
 
-      _, file_type = os.path.splitext(file_name)
-      if RESTRICT_FILE_EXTENSIONS.get() and file_type.lower() in [ext.lower() for ext in RESTRICT_FILE_EXTENSIONS.get()]:
-        err_message = f'Uploading files with type "{file_type}" is not allowed. Hue is configured to restrict this type.'
+      # Check file extension restrictions
+      is_allowed, err_message = is_file_upload_allowed(file_name)
+      if not is_allowed:
         LOG.error(err_message)
         LOG.error(err_message)
-        raise Exception(err_message)
+        self._request.META['upload_failed'] = err_message
+        self._upload_rejected = True
+        return None
 
 
       super(S3FileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
       super(S3FileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
 
 
@@ -184,6 +192,8 @@ class S3FileUploadHandler(FileUploadHandler):
         raise StopUpload()
         raise StopUpload()
 
 
   def receive_data_chunk(self, raw_data, start):
   def receive_data_chunk(self, raw_data, start):
+    if self._upload_rejected:
+      return None
     if self._is_s3_upload():
     if self._is_s3_upload():
       try:
       try:
         LOG.debug("S3FileUploadHandler uploading file part: %d" % self._part_num)
         LOG.debug("S3FileUploadHandler uploading file part: %d" % self._part_num)
@@ -199,6 +209,8 @@ class S3FileUploadHandler(FileUploadHandler):
       return raw_data
       return raw_data
 
 
   def file_complete(self, file_size):
   def file_complete(self, file_size):
+    if self._upload_rejected:
+      return None
     if self._is_s3_upload():
     if self._is_s3_upload():
       # Finish the upload
       # Finish the upload
       LOG.info("S3FileUploadHandler has completed file upload to S3, total file size is: %d." % file_size)
       LOG.info("S3FileUploadHandler has completed file upload to S3, total file size is: %d." % file_size)

+ 22 - 10
desktop/libs/azure/src/azure/abfs/upload.py

@@ -14,13 +14,12 @@
 # See the License for the specific language governing permissions and
 # See the License for the specific language governing permissions and
 # limitations under the License.
 # limitations under the License.
 
 
-import os
 import logging
 import logging
 import unicodedata
 import unicodedata
 from io import BytesIO
 from io import BytesIO
 
 
 from django.core.files.uploadedfile import SimpleUploadedFile
 from django.core.files.uploadedfile import SimpleUploadedFile
-from django.core.files.uploadhandler import FileUploadHandler, SkipFile, StopFutureHandlers, StopUpload, UploadFileException
+from django.core.files.uploadhandler import FileUploadHandler, StopFutureHandlers, StopUpload, UploadFileException
 from django.utils.translation import gettext as _
 from django.utils.translation import gettext as _
 
 
 from azure.abfs.__init__ import parse_uri
 from azure.abfs.__init__ import parse_uri
@@ -28,8 +27,7 @@ from azure.abfs.abfs import ABFSFileSystemException
 from desktop.conf import TASK_SERVER_V2
 from desktop.conf import TASK_SERVER_V2
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.fsmanager import get_client
 from desktop.lib.fsmanager import get_client
-from filebrowser.conf import RESTRICT_FILE_EXTENSIONS
-from filebrowser.utils import calculate_total_size, generate_chunks
+from filebrowser.utils import calculate_total_size, generate_chunks, is_file_upload_allowed
 
 
 DEFAULT_WRITE_SIZE = 100 * 1024 * 1024  # As per Azure doc, maximum blob size is 100MB
 DEFAULT_WRITE_SIZE = 100 * 1024 * 1024  # As per Azure doc, maximum blob size is 100MB
 
 
@@ -59,6 +57,13 @@ class ABFSFineUploaderChunkedUpload(object):
       self._fs.stats(self.destination)
       self._fs.stats(self.destination)
 
 
   def check_access(self):
   def check_access(self):
+    # Check file extension restrictions
+    is_allowed, err_message = is_file_upload_allowed(self.file_name)
+    if not is_allowed:
+      LOG.error(err_message)
+      self._request.META["upload_failed"] = err_message
+      raise PopupException(err_message)
+
     LOG.info('ABFSFineUploaderChunkedUpload: handle file upload wit temp file %s.' % self.file_name)
     LOG.info('ABFSFineUploaderChunkedUpload: handle file upload wit temp file %s.' % self.file_name)
     self.target_path = self._fs.join(self.destination, self.file_name)
     self.target_path = self._fs.join(self.destination, self.file_name)
     self.filepath = self.target_path
     self.filepath = self.target_path
@@ -73,7 +78,7 @@ class ABFSFineUploaderChunkedUpload(object):
         self._fs.create(self.target_path)
         self._fs.create(self.target_path)
     except (ABFSFileUploadError, ABFSFileSystemException) as e:
     except (ABFSFileUploadError, ABFSFileSystemException) as e:
       LOG.error("ABFSFineUploaderChunkedUpload: Encountered error in ABFSUploadHandler check_access: %s" % e)
       LOG.error("ABFSFineUploaderChunkedUpload: Encountered error in ABFSUploadHandler check_access: %s" % e)
-      self.request.META['upload_failed'] = e
+      self._request.META["upload_failed"] = e
       raise PopupException("ABFSFineUploaderChunkedUpload: Initiating ABFS upload to target path: %s failed %s" % (self.target_path, e))
       raise PopupException("ABFSFineUploaderChunkedUpload: Initiating ABFS upload to target path: %s failed %s" % (self.target_path, e))
 
 
     if self.totalfilesize != calculate_total_size(self.qquuid, self.qqtotalparts):
     if self.totalfilesize != calculate_total_size(self.qquuid, self.qqtotalparts):
@@ -90,7 +95,7 @@ class ABFSFineUploaderChunkedUpload(object):
         self._fs.create(self.target_path)
         self._fs.create(self.target_path)
       except (ABFSFileUploadError, ABFSFileSystemException) as e:
       except (ABFSFileUploadError, ABFSFileSystemException) as e:
         LOG.error("ABFSFineUploaderChunkedUpload: Encountered error in ABFSUploadHandler check_access: %s" % e)
         LOG.error("ABFSFineUploaderChunkedUpload: Encountered error in ABFSUploadHandler check_access: %s" % e)
-        self.request.META['upload_failed'] = e
+        self._request.META['upload_failed'] = e
         raise PopupException("ABFSFineUploaderChunkedUpload: Initiating ABFS upload to target path: %s failed %s" % (self.target_path, e))
         raise PopupException("ABFSFineUploaderChunkedUpload: Initiating ABFS upload to target path: %s failed %s" % (self.target_path, e))
 
 
     try:
     try:
@@ -162,6 +167,7 @@ class ABFSFileUploadHandler(FileUploadHandler):
     self.file = None
     self.file = None
     self._request = request
     self._request = request
     self._part_size = DEFAULT_WRITE_SIZE
     self._part_size = DEFAULT_WRITE_SIZE
+    self._upload_rejected = False
 
 
     if self._is_abfs_upload():
     if self._is_abfs_upload():
       self._fs = self._get_abfs(request)
       self._fs = self._get_abfs(request)
@@ -175,11 +181,13 @@ class ABFSFileUploadHandler(FileUploadHandler):
     if self._is_abfs_upload():
     if self._is_abfs_upload():
       LOG.info('Using ABFSFileUploadHandler to handle file upload wit temp file%s.' % file_name)
       LOG.info('Using ABFSFileUploadHandler to handle file upload wit temp file%s.' % file_name)
 
 
-      _, file_type = os.path.splitext(file_name)
-      if RESTRICT_FILE_EXTENSIONS.get() and file_type.lower() in [ext.lower() for ext in RESTRICT_FILE_EXTENSIONS.get()]:
-        err_message = f'Uploading files with type "{file_type}" is not allowed. Hue is configured to restrict this type.'
+      # Check file extension restrictions
+      is_allowed, err_message = is_file_upload_allowed(file_name)
+      if not is_allowed:
         LOG.error(err_message)
         LOG.error(err_message)
-        raise Exception(err_message)
+        self._request.META['upload_failed'] = err_message
+        self._upload_rejected = True
+        return None
 
 
       super(ABFSFileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
       super(ABFSFileUploadHandler, self).new_file(field_name, file_name, *args, **kwargs)
 
 
@@ -198,6 +206,8 @@ class ABFSFileUploadHandler(FileUploadHandler):
         raise StopUpload()
         raise StopUpload()
 
 
   def receive_data_chunk(self, raw_data, start):
   def receive_data_chunk(self, raw_data, start):
+    if self._upload_rejected:
+      return None
     if self._is_abfs_upload():
     if self._is_abfs_upload():
       try:
       try:
         LOG.debug("ABFSFileUploadHandler uploading file part with size: %s" % self._part_size)
         LOG.debug("ABFSFileUploadHandler uploading file part with size: %s" % self._part_size)
@@ -212,6 +222,8 @@ class ABFSFileUploadHandler(FileUploadHandler):
       return raw_data
       return raw_data
 
 
   def file_complete(self, file_size):
   def file_complete(self, file_size):
+    if self._upload_rejected:
+      return None
     if self._is_abfs_upload():
     if self._is_abfs_upload():
       # finish the upload
       # finish the upload
       self._fs.flush(self.target_path, {'position': int(file_size)})
       self._fs.flush(self.target_path, {'position': int(file_size)})

+ 33 - 11
desktop/libs/hadoop/src/hadoop/fs/upload.py

@@ -19,11 +19,11 @@
 Classes for a custom upload handler to stream into HDFS.
 Classes for a custom upload handler to stream into HDFS.
 """
 """
 
 
-import os
-import time
 import errno
 import errno
 import logging
 import logging
+import os
 import posixpath
 import posixpath
+import time
 import unicodedata
 import unicodedata
 from builtins import object
 from builtins import object
 
 
@@ -35,8 +35,8 @@ import hadoop.cluster
 from desktop.lib import fsmanager
 from desktop.lib import fsmanager
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.exceptions_renderable import PopupException
 from desktop.lib.fsmanager import get_client
 from desktop.lib.fsmanager import get_client
-from filebrowser.conf import ARCHIVE_UPLOAD_TEMPDIR, RESTRICT_FILE_EXTENSIONS
-from filebrowser.utils import calculate_total_size, generate_chunks
+from filebrowser.conf import ARCHIVE_UPLOAD_TEMPDIR
+from filebrowser.utils import calculate_total_size, generate_chunks, is_file_upload_allowed
 from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.conf import UPLOAD_CHUNK_SIZE
 from hadoop.fs.exceptions import WebHdfsException
 from hadoop.fs.exceptions import WebHdfsException
 
 
@@ -65,6 +65,12 @@ class LocalFineUploaderChunkedUpload(object):
     self.chunk_size = 0
     self.chunk_size = 0
 
 
   def check_access(self):
   def check_access(self):
+    # Check file extension restrictions
+    is_allowed, err_message = is_file_upload_allowed(self.file_name)
+    if not is_allowed:
+      LOG.error(err_message)
+      self._request.META["upload_failed"] = err_message
+      raise PopupException(err_message)
     pass
     pass
 
 
   def upload_chunks(self):
   def upload_chunks(self):
@@ -96,6 +102,13 @@ class HDFSFineUploaderChunkedUpload(object):
       self.chunk_size = kwargs.get('chunk_size')
       self.chunk_size = kwargs.get('chunk_size')
 
 
   def check_access(self):
   def check_access(self):
+    # Check file extension restrictions
+    is_allowed, err_message = is_file_upload_allowed(self.file_name)
+    if not is_allowed:
+      LOG.error(err_message)
+      self._request.META["upload_failed"] = err_message
+      raise PopupException(err_message)
+
     if self._request.fs.isdir(self.dest) and posixpath.sep in self.file_name:
     if self._request.fs.isdir(self.dest) and posixpath.sep in self.file_name:
       raise PopupException(_('HDFSFineUploaderChunkedUpload: Sorry, no "%(sep)s" in the filename %(name)s.' %
       raise PopupException(_('HDFSFineUploaderChunkedUpload: Sorry, no "%(sep)s" in the filename %(name)s.' %
                              {'sep': posixpath.sep, 'name': self.file_name}))
                              {'sep': posixpath.sep, 'name': self.file_name}))
@@ -200,7 +213,7 @@ class HDFStemporaryUploadedFile(object):
     try:
     try:
       self.size = size
       self.size = size
       self.close()
       self.close()
-    except Exception as ex:
+    except Exception:
       LOG.exception('Error uploading file to %s' % (self._path,))
       LOG.exception('Error uploading file to %s' % (self._path,))
       raise
       raise
 
 
@@ -328,6 +341,7 @@ class HDFSfileUploadHandler(FileUploadHandler):
     self._activated = False
     self._activated = False
     self._destination = request.GET.get('dest', None)  # GET param avoids infinite looping
     self._destination = request.GET.get('dest', None)  # GET param avoids infinite looping
     self.request = request
     self.request = request
+    self._upload_rejected = False
     fs = fsmanager.get_filesystem('default')
     fs = fsmanager.get_filesystem('default')
     if not fs:
     if not fs:
       LOG.warning('No HDFS set for HDFS upload')
       LOG.warning('No HDFS set for HDFS upload')
@@ -341,11 +355,13 @@ class HDFSfileUploadHandler(FileUploadHandler):
     if field_name.upper().startswith('HDFS'):
     if field_name.upper().startswith('HDFS'):
       LOG.info('Using HDFSfileUploadHandler to handle file upload.')
       LOG.info('Using HDFSfileUploadHandler to handle file upload.')
 
 
-      _, file_type = os.path.splitext(file_name)
-      if RESTRICT_FILE_EXTENSIONS.get() and file_type.lower() in [ext.lower() for ext in RESTRICT_FILE_EXTENSIONS.get()]:
-        err_message = f'Uploading files with type "{file_type}" is not allowed. Hue is configured to restrict this type.'
+      # Check file extension restrictions
+      is_allowed, err_message = is_file_upload_allowed(file_name)
+      if not is_allowed:
         LOG.error(err_message)
         LOG.error(err_message)
-        raise Exception(err_message)
+        self.request.META['upload_failed'] = err_message
+        self._upload_rejected = True
+        return None
 
 
       try:
       try:
         fs_ref = self.request.GET.get('fs', 'default')
         fs_ref = self.request.GET.get('fs', 'default')
@@ -362,6 +378,9 @@ class HDFSfileUploadHandler(FileUploadHandler):
       raise StopFutureHandlers()
       raise StopFutureHandlers()
 
 
   def receive_data_chunk(self, raw_data, start):
   def receive_data_chunk(self, raw_data, start):
+    if self._upload_rejected:
+      return None
+
     LOG.debug("HDFSfileUploadHandler receive_data_chunk")
     LOG.debug("HDFSfileUploadHandler receive_data_chunk")
 
 
     if not self._activated:
     if not self._activated:
@@ -379,6 +398,9 @@ class HDFSfileUploadHandler(FileUploadHandler):
       raise StopUpload()
       raise StopUpload()
 
 
   def file_complete(self, file_size):
   def file_complete(self, file_size):
+    if self._upload_rejected:
+      return None
+
     if not self._activated:
     if not self._activated:
       return None
       return None
 
 
@@ -481,7 +503,7 @@ class HDFSNewTemporaryUploadedFile(object):
     # Check access permissions before attempting upload
     # Check access permissions before attempting upload
     try:
     try:
       self._fs.check_access(destination, 'rw-')
       self._fs.check_access(destination, 'rw-')
-    except WebHdfsException as e:
+    except WebHdfsException:
       raise HDFSerror(_('User %s does not have permissions to write to path "%s".') % (username, destination))
       raise HDFSerror(_('User %s does not have permissions to write to path "%s".') % (username, destination))
 
 
     if self._fs.exists(self._path):
     if self._fs.exists(self._path):
@@ -504,7 +526,7 @@ class HDFSNewTemporaryUploadedFile(object):
     try:
     try:
       self.size = size
       self.size = size
       self.close()
       self.close()
-    except Exception as ex:
+    except Exception:
       LOG.exception('Error uploading file to %s' % (self._path))
       LOG.exception('Error uploading file to %s' % (self._path))
       raise
       raise