From 0ac070986c7bb9e1468704228fa7835e3ca858d6 Mon Sep 17 00:00:00 2001
From: Fabian Engelniederhammer <fabian.engelniederhammer@tngtech.com>
Date: Fri, 2 Feb 2024 10:39:06 +0100
Subject: [PATCH] feat: allow that `fields` only contains the primary key in
 protected mode #623

The result will not be of much use, since you only get all primary keys, but no additional information. This is required for the UShER integration of cov-spectrum.
---
 .idea/runConfigurations/LapisV2Protected.xml  |  2 +-
 .../auth/DataOpennessAuthorizationFilter.kt   | 36 +++++++++---
 .../auth/ProtectedDataAuthorizationTest.kt    | 57 +++++++++++++++----
 .../testData/protectedTestDatabaseConfig.yaml | 32 +++++++++++
 4 files changed, 106 insertions(+), 21 deletions(-)
 create mode 100644 siloLapisTests/testData/protectedTestDatabaseConfig.yaml

diff --git a/.idea/runConfigurations/LapisV2Protected.xml b/.idea/runConfigurations/LapisV2Protected.xml
index 4510b1bf..3b3c7b06 100644
--- a/.idea/runConfigurations/LapisV2Protected.xml
+++ b/.idea/runConfigurations/LapisV2Protected.xml
@@ -1,7 +1,7 @@
 <component name="ProjectRunConfigurationManager">
   <configuration default="false" name="LapisV2Protected" type="SpringBootApplicationConfigurationType" factoryName="Spring Boot">
     <module name="lapis-v2.main" />
-    <option name="PROGRAM_PARAMETERS" value="--silo.url=http://localhost:8081 --lapis.databaseConfig.path=siloLapisTests/testData/testDatabaseConfig.yaml --lapis.accessKeys.path=lapis2/src/test/resources/config/testAccessKeys.yaml --referenceGenomeFilename=siloLapisTests/testData/reference_genomes.json" />
+    <option name="PROGRAM_PARAMETERS" value="--silo.url=http://localhost:8081 --lapis.databaseConfig.path=siloLapisTests/testData/protectedTestDatabaseConfig.yaml --lapis.accessKeys.path=lapis2/src/test/resources/config/testAccessKeys.yaml --referenceGenomeFilename=siloLapisTests/testData/reference_genomes.json" />
     <option name="SPRING_BOOT_MAIN_CLASS" value="org.genspectrum.lapis.LapisApplicationKt" />
     <option name="WORKING_DIRECTORY" value="file://$PROJECT_DIR$" />
     <method v="2">
diff --git a/lapis2/src/main/kotlin/org/genspectrum/lapis/auth/DataOpennessAuthorizationFilter.kt b/lapis2/src/main/kotlin/org/genspectrum/lapis/auth/DataOpennessAuthorizationFilter.kt
index 90fde806..f5f3ef43 100644
--- a/lapis2/src/main/kotlin/org/genspectrum/lapis/auth/DataOpennessAuthorizationFilter.kt
+++ b/lapis2/src/main/kotlin/org/genspectrum/lapis/auth/DataOpennessAuthorizationFilter.kt
@@ -38,7 +38,7 @@ class DataOpennessAuthorizationFilterFactory(
             OpennessLevel.PROTECTED -> ProtectedDataAuthorizationFilter(
                 objectMapper,
                 accessKeysReader.read(),
-                databaseConfig.schema.metadata.filter { it.valuesAreUnique }.map { it.name },
+                databaseConfig,
             )
         }
 }
@@ -96,9 +96,14 @@ private class AlwaysAuthorizedAuthorizationFilter(objectMapper: ObjectMapper) :
 private class ProtectedDataAuthorizationFilter(
     objectMapper: ObjectMapper,
     private val accessKeys: AccessKeys,
-    private val fieldsThatServeNonAggregatedData: List<String>,
+    private val databaseConfig: DatabaseConfig,
 ) :
     DataOpennessAuthorizationFilter(objectMapper) {
+    private val fieldsThatServeNonAggregatedData = databaseConfig.schema
+        .metadata
+        .filter { it.valuesAreUnique }
+        .map { it.name }
+
     companion object {
         private val WHITELISTED_PATH_PREFIXES = listOf(
             "/swagger-ui",
@@ -130,15 +135,30 @@ private class ProtectedDataAuthorizationFilter(
             return AuthorizationResult.success()
         }
 
-        val endpointServesAggregatedData = ENDPOINTS_THAT_SERVE_AGGREGATED_DATA.contains(path) &&
-            fieldsThatServeNonAggregatedData.intersect(request.getRequestFieldNames()).isEmpty() &&
-            request.getStringArrayField(FIELDS_PROPERTY).intersect(fieldsThatServeNonAggregatedData.toSet())
-                .isEmpty()
-
-        if (endpointServesAggregatedData && accessKeys.aggregatedDataAccessKey == accessKey) {
+        if (accessKeys.aggregatedDataAccessKey == accessKey && endpointServesAggregatedData(request)) {
             return AuthorizationResult.success()
         }
 
         return AuthorizationResult.failure("You are not authorized to access $path.")
     }
+
+    private fun endpointServesAggregatedData(request: CachedBodyHttpServletRequest): Boolean {
+        val fields = request.getStringArrayField(FIELDS_PROPERTY)
+        if (containsOnlyPrimaryKey(fields)) {
+            return true
+        }
+
+        if (!ENDPOINTS_THAT_SERVE_AGGREGATED_DATA.contains(request.getProxyAwarePath())) {
+            return false
+        }
+
+        if (fieldsThatServeNonAggregatedData.intersect(request.getRequestFieldNames()).isNotEmpty()) {
+            return false
+        }
+
+        return fields.intersect(fieldsThatServeNonAggregatedData.toSet()).isEmpty()
+    }
+
+    private fun containsOnlyPrimaryKey(fields: List<String>) =
+        fields.size == 1 && fields.first() == databaseConfig.schema.primaryKey
 }
diff --git a/lapis2/src/test/kotlin/org/genspectrum/lapis/auth/ProtectedDataAuthorizationTest.kt b/lapis2/src/test/kotlin/org/genspectrum/lapis/auth/ProtectedDataAuthorizationTest.kt
index 8039f1b7..1be593d3 100644
--- a/lapis2/src/test/kotlin/org/genspectrum/lapis/auth/ProtectedDataAuthorizationTest.kt
+++ b/lapis2/src/test/kotlin/org/genspectrum/lapis/auth/ProtectedDataAuthorizationTest.kt
@@ -7,6 +7,7 @@ import io.mockk.verify
 import org.genspectrum.lapis.PRIMARY_KEY_FIELD
 import org.genspectrum.lapis.controller.AGGREGATED_ROUTE
 import org.genspectrum.lapis.controller.DATABASE_CONFIG_ROUTE
+import org.genspectrum.lapis.controller.DETAILS_ROUTE
 import org.genspectrum.lapis.controller.REFERENCE_GENOME_ROUTE
 import org.genspectrum.lapis.controller.getSample
 import org.genspectrum.lapis.controller.postSample
@@ -155,7 +156,6 @@ class ProtectedDataAuthorizationTest(
     @ParameterizedTest
     @ValueSource(
         strings = [
-            "fields=$PRIMARY_KEY_FIELD",
             "fields=$PRIMARY_KEY_FIELD,country",
             "fields=$PRIMARY_KEY_FIELD&fields=country",
         ],
@@ -180,21 +180,13 @@ class ProtectedDataAuthorizationTest(
             .andExpect(content().contentType(MediaType.APPLICATION_JSON))
     }
 
-    @ParameterizedTest
-    @ValueSource(
-        strings = [
-            """["$PRIMARY_KEY_FIELD"]""",
-            """["$PRIMARY_KEY_FIELD", "country"]""",
-        ],
-    )
-    fun `GIVEN aggregated access key in POST request but request stratifies too fine-grained THEN access is denied`(
-        fieldsJson: String,
-    ) {
+    @Test
+    fun `GIVEN aggregated access key in POST request but request stratifies too fine-grained THEN access is denied`() {
         mockMvc.perform(
             postRequestWithBody(
                 """ {
                     "accessKey": "testAggregatedDataAccessKey",
-                    "fields": $fieldsJson
+                    "fields": ["$PRIMARY_KEY_FIELD", "country"]
                 }""",
             ),
         )
@@ -203,6 +195,47 @@ class ProtectedDataAuthorizationTest(
             .andExpect(content().json(NOT_AUTHORIZED_TO_ACCESS_ENDPOINT_ERROR))
     }
 
+    @Test
+    fun `GIVEN aggregated access key in GET request where fields only contains primary key THEN access is granted`() {
+        mockMvc.perform(
+            getSample("$validRoute?accessKey=testAggregatedDataAccessKey&fields=$PRIMARY_KEY_FIELD"),
+        )
+            .andExpect(status().isOk)
+            .andExpect(content().contentType(MediaType.APPLICATION_JSON))
+    }
+
+    @Test
+    fun `GIVEN aggregated access key in POST request where fields only contains primary key THEN access is granted`() {
+        mockMvc.perform(
+            postRequestWithBody(
+                """ {
+                    "accessKey": "testAggregatedDataAccessKey",
+                    "fields": ["$PRIMARY_KEY_FIELD"]
+                }""",
+            ),
+        )
+            .andExpect(status().isOk)
+            .andExpect(content().contentType(MediaType.APPLICATION_JSON))
+    }
+
+    @Test
+    fun `GIVEN aggregated accessKey in details request where fields only contains primaryKey THEN access is granted`() {
+        every { siloQueryModelMock.getDetails(any()) } returns emptyList()
+
+        mockMvc.perform(
+            postSample(DETAILS_ROUTE)
+                .contentType(MediaType.APPLICATION_JSON)
+                .content(
+                    """{
+                        "accessKey": "testAggregatedDataAccessKey",
+                        "fields": ["$PRIMARY_KEY_FIELD"]
+                    }""",
+                ),
+        )
+            .andExpect(status().isOk)
+            .andExpect(content().contentType(MediaType.APPLICATION_JSON))
+    }
+
     @Test
     fun `given valid access key for full access in GET request to protected instance, then access is granted`() {
         mockMvc.perform(
diff --git a/siloLapisTests/testData/protectedTestDatabaseConfig.yaml b/siloLapisTests/testData/protectedTestDatabaseConfig.yaml
new file mode 100644
index 00000000..d77aecd3
--- /dev/null
+++ b/siloLapisTests/testData/protectedTestDatabaseConfig.yaml
@@ -0,0 +1,32 @@
+schema:
+  instanceName: sars_cov-2_minimal_test_config
+  opennessLevel: PROTECTED
+  metadata:
+    - name: primaryKey
+      type: string
+    - name: date
+      type: date
+    - name: region
+      type: string
+      generateIndex: true
+    - name: country
+      type: string
+      generateIndex: true
+    - name: pangoLineage
+      type: pango_lineage
+    - name: division
+      type: string
+      generateIndex: true
+    - name: age
+      type: int
+    - name: qc_value
+      type: float
+    - name: insertions
+      type: insertion
+    - name: aaInsertions
+      type: aaInsertion
+  features:
+    - name: sarsCoV2VariantQuery
+  primaryKey: primaryKey
+  dateToSortBy: date
+  partitionBy: pangoLineage