Fix ADLS Gen2 scripts and authentication checks (#1272)

* feedback from github issue

* fix authentication

* remove print, add test

---------

Co-authored-by: Matt Gotteiner <[email protected]>

Author: mattgotteiner

LoginAndAclSetup.md

@@ -200,8 +200,8 @@ The following environment variables are used to setup the optional login and doc
 * `AZURE_CLIENT_APP_ID`: Application ID of the Azure AD app for the client UI.
 * `AZURE_AUTH_TENANT_ID`: [Tenant ID](https://learn.microsoft.com/azure/active-directory/fundamentals/how-to-find-tenant) associated with the Azure AD used for login and document level access control. Defaults to `AZURE_TENANT_ID` if not defined.
 * `AZURE_ADLS_GEN2_STORAGE_ACCOUNT`: (Optional) Name of existing [Data Lake Storage Gen2 storage account](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-introduction) for storing sample data with [access control lists](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-access-control). Only used with the optional Data Lake Storage Gen2 [setup](#azure-data-lake-storage-gen2-setup) and [prep docs](#azure-data-lake-storage-gen2-prep-docs) scripts.
-* `AZURE_ADLS_GEN2_STORAGE_FILESYSTEM`: (Optional) Name of existing [Data Lake Storage Gen2 filesystem](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-introduction) for storing sample data with [access control lists](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-access-control). Only used with the optional Data Lake Storage Gen2 [setup](#azure-data-lake-storage-gen2-setup) and [prep docs](#azure-data-lake-storage-gen2-prep-docs) scripts.
-* `AZURE_ADLS_GEN2_STORAGE_FILESYSTEM_PATH`: (Optional) Name of existing path in a [Data Lake Storage Gen2 filesystem](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-introduction) for storing sample data with [access control lists](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-access-control). Only used with the optional Data Lake Storage Gen2 [prep docs](#azure-data-lake-storage-gen2-prep-docs) script.
+* `AZURE_ADLS_GEN2_FILESYSTEM`: (Optional) Name of existing [Data Lake Storage Gen2 filesystem](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-introduction) for storing sample data with [access control lists](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-access-control). Only used with the optional Data Lake Storage Gen2 [setup](#azure-data-lake-storage-gen2-setup) and [prep docs](#azure-data-lake-storage-gen2-prep-docs) scripts.
+* `AZURE_ADLS_GEN2_FILESYSTEM_PATH`: (Optional) Name of existing path in a [Data Lake Storage Gen2 filesystem](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-introduction) for storing sample data with [access control lists](https://learn.microsoft.com/azure/storage/blobs/data-lake-storage-access-control). Only used with the optional Data Lake Storage Gen2 [prep docs](#azure-data-lake-storage-gen2-prep-docs) script.
 
 ### Authentication behavior by environment
 

app/backend/core/authentication.py

@@ -241,12 +241,18 @@ async def get_auth_claims_if_enabled(self, headers: dict) -> dict[str, Any]:
     async def check_path_auth(self, path: str, auth_claims: dict[str, Any], search_client: SearchClient) -> bool:
         # Start with the standard security filter for all queries
         security_filter = self.build_security_filters(overrides={}, auth_claims=auth_claims)
-        # If there was no security filter, then the path is allowed
-        if not security_filter:
+        # If there was no security filter or no path, then the path is allowed
+        if not security_filter or len(path) == 0:
             return True
 
+        # Remove any fragment string from the path before checking
+        fragment_index = path.find("#")
+        if fragment_index != -1:
+            path = path[:fragment_index]
+
         # Filter down to only chunks that are from the specific source file
-        filter = f"{security_filter} and (sourcepage eq '{path}')"
+        # Sourcepage is used for GPT-4V
+        filter = f"{security_filter} and ((sourcefile eq '{path}') or (sourcepage eq '{path}'))"
 
         # If the filter returns any results, the user is allowed to access the document
         # Otherwise, access is denied

scripts/adlsgen2setup.py

@@ -99,6 +99,9 @@ async def run(self):
                                 await directory_client.update_access_control_recursive(
                                     acl=f"group:{groups[group_name]}:r-x"
                                 )
+                        if "oids" in access_control:
+                            for oid in access_control["oids"]:
+                                await directory_client.update_access_control_recursive(acl=f"user:{oid}:r-x")
                 finally:
                     for directory_client in directories.values():
                         await directory_client.close()

scripts/prepdocs.ps1

@@ -13,11 +13,11 @@ if ($env:AZURE_ADLS_GEN2_STORAGE_ACCOUNT) {
   $adlsGen2StorageAccountArg = "--datalakestorageaccount $env:AZURE_ADLS_GEN2_STORAGE_ACCOUNT"
   $adlsGen2FilesystemPathArg = ""
   if ($env:AZURE_ADLS_GEN2_FILESYSTEM_PATH) {
-    $adlsGen2FilesystemPathArg = "--datalakefilesystempath $env:ADLS_GEN2_FILESYSTEM_PATH"
+    $adlsGen2FilesystemPathArg = "--datalakefilesystempath $env:AZURE_ADLS_GEN2_FILESYSTEM_PATH"
   }
   $adlsGen2FilesystemArg = ""
   if ($env:AZURE_ADLS_GEN2_FILESYSTEM) {
-    $adlsGen2FilesystemArg = "--datalakefilesystem $env:ADLS_GEN2_FILESYSTEM"
+    $adlsGen2FilesystemArg = "--datalakefilesystem $env:AZURE_ADLS_GEN2_FILESYSTEM"
   }
 }
 if ($env:AZURE_USE_AUTHENTICATION) {

tests/test_authenticationhelper.py

@@ -220,12 +220,14 @@ async def mock_search(self, *args, **kwargs):
     )
     assert (
         filter
-        == "(oids/any(g:search.in(g, 'OID_X')) or groups/any(g:search.in(g, 'GROUP_Y, GROUP_Z'))) and (sourcepage eq 'Benefit_Options-2.pdf')"
+        == "(oids/any(g:search.in(g, 'OID_X')) or groups/any(g:search.in(g, 'GROUP_Y, GROUP_Z'))) and ((sourcefile eq 'Benefit_Options-2.pdf') or (sourcepage eq 'Benefit_Options-2.pdf'))"
     )
 
 
 @pytest.mark.asyncio
-async def test_check_path_auth_allowed(monkeypatch, mock_confidential_client_success, mock_validate_token_success):
+async def test_check_path_auth_allowed_sourcepage(
+    monkeypatch, mock_confidential_client_success, mock_validate_token_success
+):
     auth_helper_require_access_control = create_authentication_helper(require_access_control=True)
     filter = None
 
@@ -246,7 +248,88 @@ async def mock_search(self, *args, **kwargs):
     )
     assert (
         filter
-        == "(oids/any(g:search.in(g, 'OID_X')) or groups/any(g:search.in(g, 'GROUP_Y, GROUP_Z'))) and (sourcepage eq 'Benefit_Options-2.pdf')"
+        == "(oids/any(g:search.in(g, 'OID_X')) or groups/any(g:search.in(g, 'GROUP_Y, GROUP_Z'))) and ((sourcefile eq 'Benefit_Options-2.pdf') or (sourcepage eq 'Benefit_Options-2.pdf'))"
+    )
+
+
+@pytest.mark.asyncio
+async def test_check_path_auth_allowed_sourcefile(
+    monkeypatch, mock_confidential_client_success, mock_validate_token_success
+):
+    auth_helper_require_access_control = create_authentication_helper(require_access_control=True)
+    filter = None
+
+    async def mock_search(self, *args, **kwargs):
+        nonlocal filter
+        filter = kwargs.get("filter")
+        return MockAsyncPageIterator(data=[{"sourcefile": "Benefit_Options.pdf"}])
+
+    monkeypatch.setattr(SearchClient, "search", mock_search)
+
+    assert (
+        await auth_helper_require_access_control.check_path_auth(
+            path="Benefit_Options.pdf",
+            auth_claims={"oid": "OID_X", "groups": ["GROUP_Y", "GROUP_Z"]},
+            search_client=create_search_client(),
+        )
+        is True
+    )
+    assert (
+        filter
+        == "(oids/any(g:search.in(g, 'OID_X')) or groups/any(g:search.in(g, 'GROUP_Y, GROUP_Z'))) and ((sourcefile eq 'Benefit_Options.pdf') or (sourcepage eq 'Benefit_Options.pdf'))"
+    )
+
+
+@pytest.mark.asyncio
+async def test_check_path_auth_allowed_empty(
+    monkeypatch, mock_confidential_client_success, mock_validate_token_success
+):
+    auth_helper_require_access_control = create_authentication_helper(require_access_control=True)
+    filter = None
+
+    async def mock_search(self, *args, **kwargs):
+        nonlocal filter
+        filter = kwargs.get("filter")
+        return MockAsyncPageIterator(data=[{"sourcefile": "Benefit_Options.pdf"}])
+
+    monkeypatch.setattr(SearchClient, "search", mock_search)
+
+    assert (
+        await auth_helper_require_access_control.check_path_auth(
+            path="",
+            auth_claims={"oid": "OID_X", "groups": ["GROUP_Y", "GROUP_Z"]},
+            search_client=create_search_client(),
+        )
+        is True
+    )
+    assert filter is None
+
+
+@pytest.mark.asyncio
+async def test_check_path_auth_allowed_fragment(
+    monkeypatch, mock_confidential_client_success, mock_validate_token_success
+):
+    auth_helper_require_access_control = create_authentication_helper(require_access_control=True)
+    filter = None
+
+    async def mock_search(self, *args, **kwargs):
+        nonlocal filter
+        filter = kwargs.get("filter")
+        return MockAsyncPageIterator(data=[{"sourcefile": "Benefit_Options.pdf"}])
+
+    monkeypatch.setattr(SearchClient, "search", mock_search)
+
+    assert (
+        await auth_helper_require_access_control.check_path_auth(
+            path="Benefit_Options.pdf#textafterfragment",
+            auth_claims={"oid": "OID_X", "groups": ["GROUP_Y", "GROUP_Z"]},
+            search_client=create_search_client(),
+        )
+        is True
+    )
+    assert (
+        filter
+        == "(oids/any(g:search.in(g, 'OID_X')) or groups/any(g:search.in(g, 'GROUP_Y, GROUP_Z'))) and ((sourcefile eq 'Benefit_Options.pdf') or (sourcepage eq 'Benefit_Options.pdf'))"
     )