fix(urn-validator): update urn validation logic (datahub-project#11952)

acryldata · Nov 26, 2024 · 301e5cd · 301e5cd
1 parent 9f9a8b1
commit 301e5cd
Show file tree

Hide file tree

Showing 2 changed files with 69 additions and 2 deletions.
diff --git a/...data-io-api/src/main/java/com/linkedin/metadata/entity/validation/ValidationApiUtils.java b/...data-io-api/src/main/java/com/linkedin/metadata/entity/validation/ValidationApiUtils.java
@@ -115,7 +115,8 @@ public static void validateUrn(
 
   /** Recursively process URN parts with URL decoding */
   private static Stream<String> processUrnPartRecursively(String urnPart) {
-    String decodedPart = URLDecoder.decode(urnPart, StandardCharsets.UTF_8);
+    String decodedPart =
+        URLDecoder.decode(URLEncodingFixer.fixURLEncoding(urnPart), StandardCharsets.UTF_8);
     if (decodedPart.startsWith("urn:li:")) {
       // Recursively process nested URN after decoding
       return UrnUtils.getUrn(decodedPart).getEntityKey().getParts().stream()
@@ -177,4 +178,53 @@ public static void validateRecordTemplate(
       RecordTemplateValidator.validate(aspect, resultFunction, validator);
     }
   }
+
+  /**
+   * Fixes malformed URL encoding by escaping unescaped % characters while preserving valid
+   * percent-encoded sequences.
+   */
+  private static class URLEncodingFixer {
+    /**
+     * @param input The potentially malformed URL-encoded string
+     * @return A string with proper URL encoding that can be safely decoded
+     */
+    public static String fixURLEncoding(String input) {
+      if (input == null) {
+        return null;
+      }
+
+      StringBuilder result = new StringBuilder(input.length() * 2);
+      int i = 0;
+
+      while (i < input.length()) {
+        char currentChar = input.charAt(i);
+
+        if (currentChar == '%') {
+          if (i + 2 < input.length()) {
+            // Check if the next two characters form a valid hex pair
+            String hexPair = input.substring(i + 1, i + 3);
+            if (isValidHexPair(hexPair)) {
+              // This is a valid percent-encoded sequence, keep it as is
+              result.append(currentChar);
+            } else {
+              // Invalid sequence, escape the % character
+              result.append("%25");
+            }
+          } else {
+            // % at the end of string, escape it
+            result.append("%25");
+          }
+        } else {
+          result.append(currentChar);
+        }
+        i++;
+      }
+
+      return result.toString();
+    }
+
+    private static boolean isValidHexPair(String pair) {
+      return pair.matches("[0-9A-Fa-f]{2}");
+    }
+  }
 }
diff --git a/...-io-api/src/test/java/com/linkedin/metadata/entity/validation/ValidationApiUtilsTest.java b/...-io-api/src/test/java/com/linkedin/metadata/entity/validation/ValidationApiUtilsTest.java
@@ -89,7 +89,24 @@ public void testValidComplexUrn() {
 
   @Test(expectedExceptions = NullPointerException.class)
   public void testUrnNull() {
-    // Act
     ValidationApiUtils.validateUrn(entityRegistry, null);
   }
+
+  @Test
+  public void testValidPartialUrlEncode() {
+    Urn validUrn = UrnUtils.getUrn("urn:li:assertion:123=-%28__% weekly__%29");
+
+    ValidationApiUtils.validateUrn(entityRegistry, validUrn);
+    // If no exception is thrown, test passes
+  }
+
+  @Test
+  public void testValidPartialUrlEncode2() {
+    Urn validUrn =
+        UrnUtils.getUrn(
+            "urn:li:dataset:(urn:li:dataPlatform:s3,urn:li:dataset:%28urn:li:dataPlatform:s3%2Ctest-datalake-concepts%prog_maintenance%2CPROD%29,PROD)");
+
+    ValidationApiUtils.validateUrn(entityRegistry, validUrn);
+    // If no exception is thrown, test passes
+  }
 }