Merge pull request #3 from offsoc/add-aes256-encryption

Add AES256 encryption for file storage

service/src/main/java/org/whispersystems/textsecuregcm/attachments/GcsAttachmentGenerator.java

@@ -8,12 +8,19 @@
 import org.whispersystems.textsecuregcm.gcp.CanonicalRequest;
 import org.whispersystems.textsecuregcm.gcp.CanonicalRequestGenerator;
 import org.whispersystems.textsecuregcm.gcp.CanonicalRequestSigner;
+
 import javax.annotation.Nonnull;
+import javax.crypto.Cipher;
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
 import java.io.IOException;
 import java.security.InvalidKeyException;
+import java.security.NoSuchAlgorithmException;
 import java.security.spec.InvalidKeySpecException;
 import java.time.ZoneOffset;
 import java.time.ZonedDateTime;
+import java.util.Base64;
 import java.util.Map;
 
 public class GcsAttachmentGenerator implements AttachmentGenerator {
@@ -23,11 +30,15 @@ public class GcsAttachmentGenerator implements AttachmentGenerator {
   @Nonnull
   private final CanonicalRequestSigner canonicalRequestSigner;
 
+  @Nonnull
+  private final SecretKey encryptionKey;
+
   public GcsAttachmentGenerator(@Nonnull String domain, @Nonnull String email,
-      int maxSizeInBytes, @Nonnull String pathPrefix, @Nonnull String rsaSigningKey)
-      throws IOException, InvalidKeyException, InvalidKeySpecException {
+      int maxSizeInBytes, @Nonnull String pathPrefix, @Nonnull String rsaSigningKey, @Nonnull String encryptionKey)
+      throws IOException, InvalidKeyException, InvalidKeySpecException, NoSuchAlgorithmException {
     this.canonicalRequestGenerator = new CanonicalRequestGenerator(domain, email, maxSizeInBytes, pathPrefix);
     this.canonicalRequestSigner = new CanonicalRequestSigner(rsaSigningKey);
+    this.encryptionKey = new SecretKeySpec(Base64.getDecoder().decode(encryptionKey), "AES");
   }
 
   @Override
@@ -50,5 +61,9 @@ private static Map<String, String> getHeaderMap(@Nonnull CanonicalRequest canoni
         "x-goog-resumable", "start");
   }
 
-
+  private byte[] encryptData(byte[] data) throws Exception {
+    Cipher cipher = Cipher.getInstance("AES");
+    cipher.init(Cipher.ENCRYPT_MODE, encryptionKey);
+    return cipher.doFinal(data);
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/attachments/TusAttachmentGenerator.java

@@ -9,6 +9,9 @@
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentials;
 import org.whispersystems.textsecuregcm.auth.ExternalServiceCredentialsGenerator;
 import org.whispersystems.textsecuregcm.util.HeaderUtils;
+
+import javax.crypto.Cipher;
+import javax.crypto.spec.SecretKeySpec;
 import java.nio.charset.StandardCharsets;
 import java.time.Clock;
 import java.util.Base64;
@@ -20,10 +23,12 @@ public class TusAttachmentGenerator implements AttachmentGenerator {
 
   final ExternalServiceCredentialsGenerator credentialsGenerator;
   final String tusUri;
+  private final SecretKeySpec encryptionKey;
 
-  public TusAttachmentGenerator(final TusConfiguration cfg) {
+  public TusAttachmentGenerator(final TusConfiguration cfg, final String encryptionKey) {
     this.tusUri = cfg.uploadUri();
     this.credentialsGenerator = credentialsGenerator(Clock.systemUTC(), cfg);
+    this.encryptionKey = new SecretKeySpec(Base64.getDecoder().decode(encryptionKey), "AES");
   }
 
   private static ExternalServiceCredentialsGenerator credentialsGenerator(final Clock clock, final TusConfiguration cfg) {
@@ -44,4 +49,10 @@ public Descriptor generateAttachment(final String key) {
     );
     return new Descriptor(headers, tusUri + "/" +  ATTACHMENTS);
   }
+
+  private byte[] encryptData(byte[] data) throws Exception {
+    Cipher cipher = Cipher.getInstance("AES");
+    cipher.init(Cipher.ENCRYPT_MODE, encryptionKey);
+    return cipher.doFinal(data);
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupAuthManager.java

@@ -35,6 +35,11 @@
 import org.whispersystems.textsecuregcm.storage.AccountsManager;
 import org.whispersystems.textsecuregcm.storage.RedeemedReceiptsManager;
 import org.whispersystems.textsecuregcm.util.Util;
+import javax.crypto.Cipher;
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+import java.util.Base64;
 
 /**
  * Issues ZK backup auth credentials for authenticated accounts
@@ -62,6 +67,7 @@ public class BackupAuthManager {
   private final Clock clock;
   private final RateLimiters rateLimiters;
   private final AccountsManager accountsManager;
+  private final SecretKey encryptionKey;
 
   public BackupAuthManager(
       final ExperimentEnrollmentManager experimentEnrollmentManager,
@@ -70,14 +76,16 @@ public BackupAuthManager(
       final ServerZkReceiptOperations serverZkReceiptOperations,
       final RedeemedReceiptsManager redeemedReceiptsManager,
       final GenericServerSecretParams serverSecretParams,
-      final Clock clock) {
+      final Clock clock,
+      final SecretKey encryptionKey) {
     this.experimentEnrollmentManager = experimentEnrollmentManager;
     this.rateLimiters = rateLimiters;
     this.accountsManager = accountsManager;
     this.serverZkReceiptOperations = serverZkReceiptOperations;
     this.redeemedReceiptsManager = redeemedReceiptsManager;
     this.serverSecretParams = serverSecretParams;
     this.clock = clock;
+    this.encryptionKey = encryptionKey;
   }
 
   /**
@@ -117,7 +125,7 @@ public CompletableFuture<Void> commitBackupId(final Account account,
     return rateLimiters.forDescriptor(RateLimiters.For.SET_BACKUP_ID)
         .validateAsync(account.getUuid())
         .thenCompose(ignored -> this.accountsManager
-            .updateAsync(account, a -> a.setBackupCredentialRequests(serializedMessageCredentialRequest, serializedMediaCredentialRequest))
+            .updateAsync(account, a -> a.setBackupCredentialRequests(encrypt(serializedMessageCredentialRequest), encrypt(serializedMediaCredentialRequest)))
             .thenRun(Util.NOOP))
         .toCompletableFuture();
   }
@@ -171,8 +179,8 @@ public CompletableFuture<List<Credential>> getBackupAuthCredentials(
     }
 
     // fetch the blinded backup-id the account should have previously committed to
-    final byte[] committedBytes = account.getBackupCredentialRequest(credentialType)
-        .orElseThrow(() -> Status.NOT_FOUND.withDescription("No blinded backup-id has been added to the account").asRuntimeException());
+    final byte[] committedBytes = decrypt(account.getBackupCredentialRequest(credentialType)
+        .orElseThrow(() -> Status.NOT_FOUND.withDescription("No blinded backup-id has been added to the account").asRuntimeException()));
 
     try {
       // create a credential for every day in the requested period
@@ -318,4 +326,24 @@ private Optional<BackupLevel> configuredBackupLevel(final Account account) {
   private boolean inExperiment(final String experimentName, final Account account) {
     return this.experimentEnrollmentManager.isEnrolled(account.getUuid(), experimentName);
   }
+
+  private byte[] encrypt(byte[] data) {
+    try {
+      Cipher cipher = Cipher.getInstance("AES");
+      cipher.init(Cipher.ENCRYPT_MODE, encryptionKey);
+      return cipher.doFinal(data);
+    } catch (Exception e) {
+      throw new RuntimeException("Error while encrypting data", e);
+    }
+  }
+
+  private byte[] decrypt(byte[] encryptedData) {
+    try {
+      Cipher cipher = Cipher.getInstance("AES");
+      cipher.init(Cipher.DECRYPT_MODE, encryptionKey);
+      return cipher.doFinal(encryptedData);
+    } catch (Exception e) {
+      throw new RuntimeException("Error while decrypting data", e);
+    }
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupLevelUtil.java

@@ -6,12 +6,62 @@
 
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
 
+import javax.crypto.Cipher;
+import javax.crypto.KeyGenerator;
+import javax.crypto.SecretKey;
+import javax.crypto.spec.SecretKeySpec;
+import java.security.SecureRandom;
+import java.util.Base64;
+
 public class BackupLevelUtil {
+
+  private static final String AES = "AES";
+  private static final int AES_KEY_SIZE = 256;
+  private static final String AES_CIPHER_ALGORITHM = "AES/ECB/PKCS5Padding";
+
+  private final SecretKey encryptionKey;
+
+  public BackupLevelUtil(String userDefinedKey) {
+    this.encryptionKey = generateKey(userDefinedKey);
+  }
+
   public static BackupLevel fromReceiptLevel(long receiptLevel) {
     try {
       return BackupLevel.fromValue(Math.toIntExact(receiptLevel));
     } catch (ArithmeticException e) {
       throw new IllegalArgumentException("Invalid receipt level: " + receiptLevel);
     }
   }
+
+  private SecretKey generateKey(String userDefinedKey) {
+    try {
+      byte[] keyBytes = userDefinedKey.getBytes();
+      return new SecretKeySpec(keyBytes, 0, AES_KEY_SIZE / 8, AES);
+    } catch (Exception e) {
+      throw new RuntimeException("Error generating encryption key", e);
+    }
+  }
+
+  public String encrypt(String data) {
+    try {
+      Cipher cipher = Cipher.getInstance(AES_CIPHER_ALGORITHM);
+      cipher.init(Cipher.ENCRYPT_MODE, encryptionKey);
+      byte[] encryptedBytes = cipher.doFinal(data.getBytes());
+      return Base64.getEncoder().encodeToString(encryptedBytes);
+    } catch (Exception e) {
+      throw new RuntimeException("Error encrypting data", e);
+    }
+  }
+
+  public String decrypt(String encryptedData) {
+    try {
+      Cipher cipher = Cipher.getInstance(AES_CIPHER_ALGORITHM);
+      cipher.init(Cipher.DECRYPT_MODE, encryptionKey);
+      byte[] decodedBytes = Base64.getDecoder().decode(encryptedData);
+      byte[] decryptedBytes = cipher.doFinal(decodedBytes);
+      return new String(decryptedBytes);
+    } catch (Exception e) {
+      throw new RuntimeException("Error decrypting data", e);
+    }
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupManager.java

@@ -23,6 +23,8 @@
 import java.util.concurrent.CompletionStage;
 import java.util.concurrent.atomic.AtomicLong;
 import java.util.function.Function;
+import javax.crypto.Cipher;
+import javax.crypto.spec.SecretKeySpec;
 import org.signal.libsignal.protocol.ecc.Curve;
 import org.signal.libsignal.protocol.ecc.ECPublicKey;
 import org.signal.libsignal.zkgroup.GenericServerSecretParams;
@@ -79,6 +81,7 @@ public class BackupManager {
   private final RemoteStorageManager remoteStorageManager;
   private final SecureRandom secureRandom = new SecureRandom();
   private final Clock clock;
+  private final SecretKeySpec encryptionKey;
 
 
   public BackupManager(
@@ -88,14 +91,16 @@ public BackupManager(
       final TusAttachmentGenerator tusAttachmentGenerator,
       final Cdn3BackupCredentialGenerator cdn3BackupCredentialGenerator,
       final RemoteStorageManager remoteStorageManager,
-      final Clock clock) {
+      final Clock clock,
+      final String encryptionKey) {
     this.backupsDb = backupsDb;
     this.serverSecretParams = serverSecretParams;
     this.rateLimiters = rateLimiters;
     this.tusAttachmentGenerator = tusAttachmentGenerator;
     this.cdn3BackupCredentialGenerator = cdn3BackupCredentialGenerator;
     this.remoteStorageManager = remoteStorageManager;
     this.clock = clock;
+    this.encryptionKey = new SecretKeySpec(Base64.getDecoder().decode(encryptionKey), "AES");
   }
 
 
@@ -149,7 +154,14 @@ public CompletableFuture<BackupUploadDescriptor> createMessageBackupUploadDescri
     // this could race with concurrent updates, but the only effect would be last-writer-wins on the timestamp
     return backupsDb
         .addMessageBackup(backupUser)
-        .thenApply(result -> cdn3BackupCredentialGenerator.generateUpload(cdnMessageBackupName(backupUser)));
+        .thenApply(result -> {
+          try {
+            byte[] encryptedData = encryptData(cdn3BackupCredentialGenerator.generateUpload(cdnMessageBackupName(backupUser)).toString().getBytes());
+            return new BackupUploadDescriptor(3, Base64.getEncoder().encodeToString(encryptedData), Map.of(), "");
+          } catch (Exception e) {
+            throw new RuntimeException("Failed to encrypt data", e);
+          }
+        });
   }
 
   public CompletableFuture<BackupUploadDescriptor> createTemporaryAttachmentUploadDescriptor(
@@ -689,4 +701,10 @@ private static String cdnMediaPath(final AuthenticatedBackupUser backupUser, fin
   static String rateLimitKey(final AuthenticatedBackupUser backupUser) {
     return Base64.getEncoder().encodeToString(BackupsDb.hashedBackupId(backupUser.backupId()));
   }
+
+  private byte[] encryptData(byte[] data) throws Exception {
+    Cipher cipher = Cipher.getInstance("AES");
+    cipher.init(Cipher.ENCRYPT_MODE, encryptionKey);
+    return cipher.doFinal(data);
+  }
 }

service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupUploadDescriptor.java

@@ -5,10 +5,30 @@
 
 package org.whispersystems.textsecuregcm.backup;
 
+import javax.crypto.Cipher;
+import javax.crypto.spec.SecretKeySpec;
+import java.util.Base64;
 import java.util.Map;
 
 public record BackupUploadDescriptor(
     int cdn,
     String key,
     Map<String, String> headers,
-    String signedUploadLocation) {}
+    String signedUploadLocation) {
+
+  public BackupUploadDescriptor(int cdn, String key, Map<String, String> headers, String signedUploadLocation, String encryptionKey) {
+    this(cdn, encryptKey(key, encryptionKey), headers, signedUploadLocation);
+  }
+
+  private static String encryptKey(String key, String encryptionKey) {
+    try {
+      SecretKeySpec secretKey = new SecretKeySpec(Base64.getDecoder().decode(encryptionKey), "AES");
+      Cipher cipher = Cipher.getInstance("AES");
+      cipher.init(Cipher.ENCRYPT_MODE, secretKey);
+      byte[] encryptedKey = cipher.doFinal(key.getBytes());
+      return Base64.getEncoder().encodeToString(encryptedKey);
+    } catch (Exception e) {
+      throw new RuntimeException("Failed to encrypt key", e);
+    }
+  }
+}

service/src/main/java/org/whispersystems/textsecuregcm/backup/BackupsDb.java

@@ -21,6 +21,8 @@
 import java.util.Optional;
 import java.util.concurrent.CompletableFuture;
 import java.util.function.Predicate;
+import javax.crypto.Cipher;
+import javax.crypto.spec.SecretKeySpec;
 import org.signal.libsignal.protocol.InvalidKeyException;
 import org.signal.libsignal.protocol.ecc.ECPublicKey;
 import org.signal.libsignal.zkgroup.backups.BackupLevel;
@@ -78,6 +80,7 @@ public class BackupsDb {
   private final Clock clock;
 
   private final SecureRandom secureRandom;
+  private final SecretKeySpec encryptionKey;
 
   // The backups table
 
@@ -111,11 +114,13 @@ public class BackupsDb {
   public BackupsDb(
       final DynamoDbAsyncClient dynamoClient,
       final String backupTableName,
-      final Clock clock) {
+      final Clock clock,
+      final String encryptionKey) {
     this.dynamoClient = dynamoClient;
     this.backupTableName = backupTableName;
     this.clock = clock;
     this.secureRandom = new SecureRandom();
+    this.encryptionKey = new SecretKeySpec(Base64.getDecoder().decode(encryptionKey), "AES");
   }
 
   /**
@@ -800,4 +805,10 @@ static byte[] hashedBackupId(final byte[] backupId) {
       throw new AssertionError(e);
     }
   }
+
+  private byte[] encryptData(byte[] data) throws Exception {
+    Cipher cipher = Cipher.getInstance("AES");
+    cipher.init(Cipher.ENCRYPT_MODE, encryptionKey);
+    return cipher.doFinal(data);
+  }
 }