[feat] add eTag encrypting/decrypting (#554)

* add etag encrypting

* fix

Author: ybz1013

dao-api/src/main/pegasus/pegasus/com/linkedin/metadata/events/IngestionAspectETag.pdl

@@ -8,10 +8,10 @@ record IngestionAspectETag {
   /**
    * aspect field name, e.g. "status"
    */
-  aspect_name: optional string = ""
+  aspect_alias: optional string = ""
 
   /**
-   * e.g. timestamp used for optimistic locking when writing new aspect value
+   * e.g. used for optimistic locking when writing new aspect value
    */
-  eTag: optional long = 0
+  etag: optional string = ""
 }

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/EbeanLocalDAO.java

@@ -22,6 +22,7 @@
 import com.linkedin.metadata.dao.urnpath.EmptyPathExtractor;
 import com.linkedin.metadata.dao.urnpath.UrnPathExtractor;
 import com.linkedin.metadata.dao.utils.EBeanDAOUtils;
+import com.linkedin.metadata.dao.utils.ETagUtils;
 import com.linkedin.metadata.dao.utils.ModelUtils;
 import com.linkedin.metadata.dao.utils.QueryUtils;
 import com.linkedin.metadata.dao.utils.RecordUtils;
@@ -631,16 +632,34 @@ public <ASPECT extends RecordTemplate> AuditStamp extractOptimisticLockForAspect
           continue;
         }
 
-        if (aspectAlias != null && aspectAlias.equalsIgnoreCase(ingestionAspectETag.getAspect_name()) && ingestionAspectETag.getETag() != null) {
-          optimisticLockAuditStamp = new AuditStamp();
-          optimisticLockAuditStamp.setTime(ingestionAspectETag.getETag());
-          break;
+        if (aspectAlias != null && aspectAlias.equalsIgnoreCase(ingestionAspectETag.getAspect_alias())) {
+          Long decryptedETag = getDecryptedETag(ingestionAspectETag);
+          if (decryptedETag != null) {
+            optimisticLockAuditStamp = new AuditStamp();
+            optimisticLockAuditStamp.setTime(decryptedETag);
+            break;
+          }
         }
       }
     }
     return optimisticLockAuditStamp;
   }
 
+  /**
+   * When eTag is null, it means this is a regular ingestion request, no read-modify-write consistency guarantee.
+   */
+  @Nullable
+  private Long getDecryptedETag(@Nonnull IngestionAspectETag ingestionAspectETag) {
+    try {
+      if (ingestionAspectETag.getEtag() == null) {
+        return null;
+      }
+      return ETagUtils.decrypt(ingestionAspectETag.getEtag());
+    } catch (Exception e) {
+      return null;
+    }
+  }
+
   @Override
   protected <ASPECT extends RecordTemplate> long saveLatest(@Nonnull URN urn, @Nonnull Class<ASPECT> aspectClass,
       @Nullable ASPECT oldValue, @Nullable AuditStamp optimisticLockAuditStamp, @Nullable ASPECT newValue,

dao-impl/ebean-dao/src/main/java/com/linkedin/metadata/dao/utils/ETagUtils.java

@@ -0,0 +1,72 @@
+package com.linkedin.metadata.dao.utils;
+
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import java.util.Base64;
+import javax.annotation.Nonnull;
+import javax.crypto.BadPaddingException;
+import javax.crypto.Cipher;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+
+
+/**
+ * Utility class for IngestionAspectETag. E.g., encrypting and decrypting timestamps using AES encryption.
+ */
+public final class ETagUtils {
+
+  public static final String AES = "AES";
+
+  private ETagUtils() {
+  }
+
+  // 16-char key for AES-128
+  private static final String SECRET_KEY = "9012312344567856";
+
+  /**
+   * Encrypts a timestamp using AES encryption.
+   * @param timestamp Timestamp to encrypt
+   * @return Encrypted timestamp as a Base64 encoded string
+   */
+  @Nonnull
+  public static String encrypt(long timestamp)
+      throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException, IllegalBlockSizeException,
+             BadPaddingException {
+    SecretKey secretKey = getSecretKey();
+
+    Cipher cipher = Cipher.getInstance(AES);
+    cipher.init(Cipher.ENCRYPT_MODE, secretKey);
+
+    byte[] inputBytes = Long.toString(timestamp).getBytes();
+    byte[] encrypted = cipher.doFinal(inputBytes);
+
+    return Base64.getEncoder().encodeToString(encrypted);
+  }
+
+  /**
+   * Decrypts a Base64 encoded encrypted timestamp string back to the original timestamp.
+   * @param encrypted Base64 encoded encrypted timestamp string
+   * @return Decrypted timestamp as a long
+   */
+  public static long decrypt(@Nonnull String encrypted)
+      throws NoSuchPaddingException, NoSuchAlgorithmException, InvalidKeyException, IllegalBlockSizeException,
+             BadPaddingException {
+    SecretKey secretKey = getSecretKey();
+
+    Cipher cipher = Cipher.getInstance(AES);
+    cipher.init(Cipher.DECRYPT_MODE, secretKey);
+
+    byte[] decoded = Base64.getDecoder().decode(encrypted);
+    byte[] decrypted = cipher.doFinal(decoded);
+
+    return Long.parseLong(new String(decrypted));
+  }
+
+  private static SecretKey getSecretKey() {
+    byte[] keyBytes = SECRET_KEY.getBytes();
+    return new SecretKeySpec(keyBytes, AES);
+  }
+
+}

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/EbeanLocalDAOTest.java

@@ -4064,22 +4064,24 @@ private void assertVersionMetadata(ListResultMetadata listResultMetadata, List<L
   }
 
   @Test
-  public void testExtractOptimisticLockForAspectFromIngestionParamsIfPossible() throws URISyntaxException {
+  public void testExtractOptimisticLockForAspectFromIngestionParamsIfPossible()
+      throws URISyntaxException {
     EbeanLocalDAO<EntityAspectUnion, FooUrn> dao = createDao(FooUrn.class);
 
     FooUrn urn = new FooUrn(1);
 
     IngestionAspectETag ingestionAspectETag = new IngestionAspectETag();
-    ingestionAspectETag.setAspect_name("aspectFoo");
-    ingestionAspectETag.setETag(1234L);
+    ingestionAspectETag.setAspect_alias("aspectFoo");
+    long timestamp = 1750796203701L;
+    ingestionAspectETag.setEtag("KsFkRXtjaBGQf37HjdEjDQ==");
 
     IngestionParams ingestionParams = new IngestionParams();
     ingestionParams.setIngestionETags(new IngestionAspectETagArray(ingestionAspectETag));
 
     AuditStamp result = dao.extractOptimisticLockForAspectFromIngestionParamsIfPossible(ingestionParams, AspectFoo.class,
         urn);
 
-    assertEquals(result.getTime(), Long.valueOf(1234L));
+    assertEquals(result.getTime(), Long.valueOf(timestamp));
   }
 
   @Test
@@ -4103,8 +4105,8 @@ public void testExtractOptimisticLockForAspectFromIngestionParamsIfPossibleAspec
     FooUrn urn = new FooUrn(1);
 
     IngestionAspectETag ingestionAspectETag = new IngestionAspectETag();
-    ingestionAspectETag.setAspect_name("aspectBar");
-    ingestionAspectETag.setETag(1234L);
+    ingestionAspectETag.setAspect_alias("aspectBar");
+    ingestionAspectETag.setEtag("KsFkRXtjaBGQf37HjdEjDQ==");
 
     IngestionParams ingestionParams = new IngestionParams();
     ingestionParams.setIngestionETags(new IngestionAspectETagArray(ingestionAspectETag));

dao-impl/ebean-dao/src/test/java/com/linkedin/metadata/dao/utils/ETagUtilsTest.java

@@ -0,0 +1,45 @@
+package com.linkedin.metadata.dao.utils;
+
+import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
+import javax.crypto.BadPaddingException;
+import javax.crypto.IllegalBlockSizeException;
+import javax.crypto.NoSuchPaddingException;
+import org.testng.annotations.Test;
+
+import static org.junit.Assert.*;
+
+
+public class ETagUtilsTest {
+
+  @Test
+  public void testEncryptAndDecrypt()
+      throws NoSuchPaddingException, IllegalBlockSizeException, NoSuchAlgorithmException, BadPaddingException,
+             InvalidKeyException {
+    long timestamp = System.currentTimeMillis();
+    String encrypted = ETagUtils.encrypt(timestamp);
+
+    long decrypted = ETagUtils.decrypt(encrypted);
+    assertEquals(decrypted, timestamp);
+  }
+
+  @Test
+  public void testEncrypt()
+      throws NoSuchPaddingException, IllegalBlockSizeException, NoSuchAlgorithmException, BadPaddingException,
+             InvalidKeyException {
+    long timestamp = 1750796203701L;
+    String encrypted = ETagUtils.encrypt(timestamp);
+
+    assertEquals("KsFkRXtjaBGQf37HjdEjDQ==", encrypted);
+  }
+
+  @Test
+  public void testDecrypt()
+      throws NoSuchPaddingException, IllegalBlockSizeException, NoSuchAlgorithmException, BadPaddingException,
+             InvalidKeyException {
+    String encrypted = "KsFkRXtjaBGQf37HjdEjDQ==";
+    long decrypted = ETagUtils.decrypt(encrypted);
+
+    assertEquals(1750796203701L, decrypted);
+  }
+}