Merge pull request #1729 from o1-labs/v2-encryption

Encryption v2

CHANGELOG.md

@@ -22,6 +22,8 @@ This project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.htm
 ### Added
 
 - `SmartContract.emitEventIf()` to conditionally emit an event https://github.com/o1-labs/o1js/pull/1746
+- Added `Encryption.encryptV2()` and `Encryption.decryptV2()` for an updated encryption algorithm that guarantees cipher text integrity.
+  - Also added `Encryption.encryptBytes()` and `Encryption.decryptBytes()` using the same algorithm.
 
 ### Changed
 

src/examples/encryptionv2.ts

@@ -0,0 +1,32 @@
+import assert from 'assert';
+import {
+  Bytes,
+  PrivateKey,
+  initializeBindings,
+  Encryption,
+  Encoding,
+} from 'o1js';
+
+await initializeBindings();
+
+class Bytes256 extends Bytes(256) {}
+const priv = PrivateKey.random();
+const pub = priv.toPublicKey();
+
+const plainMsg = 'The quick brown fox jumped over the angry dog.';
+
+console.log('en/decryption of field elements');
+const cipher2 = Encryption.encryptV2(Encoding.stringToFields(plainMsg), pub);
+const plainText2 = Encryption.decryptV2(cipher2, priv);
+
+assert(
+  Encoding.stringFromFields(plainText2) === plainMsg,
+  'Plain message and decrypted message are the same'
+);
+
+console.log('en/decryption of bytes');
+const message = Bytes256.fromString(plainMsg);
+console.log('plain message', plainMsg);
+const cipher = Encryption.encryptBytes(message, pub);
+const plainText = Encryption.decryptBytes(cipher, priv);
+console.log('decrypted message', Buffer.from(plainText.toBytes()).toString());

src/lib/provable/bytes.ts

@@ -1,7 +1,7 @@
 import { provableFromClass } from './types/provable-derivers.js';
 import type { ProvablePureExtended } from './types/struct.js';
 import { assert } from './gadgets/common.js';
-import { chunkString } from '../util/arrays.js';
+import { chunk, chunkString } from '../util/arrays.js';
 import { Provable } from './provable.js';
 import { UInt8 } from './int.js';
 import { randomBytes } from '../../bindings/crypto/random.js';
@@ -194,6 +194,15 @@ class Bytes {
     return Bytes.from(decodedB64Bytes);
   }
 
+  /**
+   * Returns an array of chunks, each of size `size`.
+   * @param size size of each chunk
+   * @returns an array of {@link UInt8} chunks
+   */
+  chunk(size: number) {
+    return chunk(this.bytes, size);
+  }
+
   // dynamic subclassing infra
   static _size?: number;
   static _provable?: ProvablePureExtended<

src/lib/provable/crypto/encryption.ts

@@ -2,15 +2,30 @@ import { Field, Scalar, Group } from '../wrapped.js';
 import { Poseidon } from './poseidon.js';
 import { Provable } from '../provable.js';
 import { PrivateKey, PublicKey } from './signature.js';
+import { bytesToWord, wordToBytes } from '../gadgets/bit-slices.js';
+import { Bytes } from '../bytes.js';
+import { UInt8 } from '../int.js';
+import { chunk } from '../../util/arrays.js';
 
-export { encrypt, decrypt };
+export {
+  encrypt,
+  decrypt,
+  encryptV2,
+  decryptV2,
+  encryptBytes,
+  decryptBytes,
+  CipherTextBytes,
+  CipherText,
+};
 
 type CipherText = {
   publicKey: Group;
   cipherText: Field[];
 };
+type CipherTextBytes = CipherText & { messageLength: number };
 
 /**
+ * @deprecated Use {@link encryptV2} instead.
  * Public Key Encryption, using a given array of {@link Field} elements and encrypts it using a {@link PublicKey}.
  */
 function encrypt(message: Field[], otherPublicKey: PublicKey) {
@@ -40,7 +55,8 @@ function encrypt(message: Field[], otherPublicKey: PublicKey) {
 }
 
 /**
- * Decrypts a {@link CipherText} using a {@link PrivateKey}.^
+ * @deprecated Use {@link decryptV2} instead.
+ * Decrypts a {@link CipherText} using a {@link PrivateKey}.
  */
 function decrypt(
   { publicKey, cipherText }: CipherText,
@@ -68,3 +84,119 @@ function decrypt(
 
   return message;
 }
+
+/**
+ * Decrypts a {@link CipherText} using a {@link PrivateKey}.
+ */
+function decryptV2(
+  { publicKey, cipherText }: CipherText,
+  privateKey: PrivateKey
+) {
+  // key exchange
+  const sharedSecret = publicKey.scale(privateKey.s);
+  const sponge = new Poseidon.Sponge();
+  sponge.absorb(sharedSecret.x);
+  const authenticationTag = cipherText.pop();
+
+  // decryption
+  const message = [];
+  for (let i = 0; i < cipherText.length; i++) {
+    // absorb frame tag
+    if (i === cipherText.length - 1) sponge.absorb(Field(1));
+    else sponge.absorb(Field(0));
+
+    const keyStream = sponge.squeeze();
+    const messageChunk = cipherText[i].sub(keyStream);
+
+    // push the message to our final messages
+    message.push(messageChunk);
+
+    // absorb the cipher text chunk
+    sponge.absorb(cipherText[i]);
+  }
+
+  // authentication tag
+  sponge.squeeze().assertEquals(authenticationTag!);
+
+  return message;
+}
+
+/**
+ * Public Key Encryption, encrypts Field elements using a {@link PublicKey}.
+ */
+function encryptV2(message: Field[], otherPublicKey: PublicKey): CipherText {
+  // key exchange
+  const privateKey = Provable.witness(Scalar, () => Scalar.random());
+  const publicKey = Group.generator.scale(privateKey);
+  const sharedSecret = otherPublicKey.toGroup().scale(privateKey);
+
+  const sponge = new Poseidon.Sponge();
+  sponge.absorb(sharedSecret.x);
+
+  // encryption
+  const cipherText = [];
+  for (let [n, chunk] of message.entries()) {
+    // absorb frame bit
+    if (n === message.length - 1) sponge.absorb(Field(1));
+    else sponge.absorb(Field(0));
+
+    const keyStream = sponge.squeeze();
+    const encryptedChunk = chunk.add(keyStream);
+    cipherText.push(encryptedChunk);
+
+    sponge.absorb(encryptedChunk);
+  }
+
+  // authentication tag
+  const authenticationTag = sponge.squeeze();
+  cipherText.push(authenticationTag);
+
+  return { publicKey, cipherText };
+}
+
+/**
+ * Public Key Encryption, encrypts Bytes using a {@link PublicKey}.
+ */
+function encryptBytes(
+  message: Bytes,
+  otherPublicKey: PublicKey
+): CipherTextBytes {
+  const bytes = message.bytes;
+  const messageLength = bytes.length;
+
+  // pad message to a multiple of 31 so they still fit into one field element
+  const multipleOf = 31;
+  const n = Math.ceil(messageLength / multipleOf) * multipleOf;
+
+  // create the padding
+  const padding = Array.from({ length: n - messageLength }, () =>
+    UInt8.from(0)
+  );
+
+  // convert message into chunks of 31 bytes
+  const chunks = chunk(bytes.concat(padding), 31);
+
+  // call into encryption() and convert chunk to field elements
+  return {
+    ...encryptV2(
+      chunks.map((chunk) => bytesToWord(chunk)),
+      otherPublicKey
+    ),
+    messageLength,
+  };
+}
+
+/**
+ * Decrypts a {@link CipherText} using a {@link PrivateKey}.
+ */
+function decryptBytes(cipherText: CipherTextBytes, privateKey: PrivateKey) {
+  // calculate padding
+  const messageLength = cipherText.messageLength;
+  const multipleOf = 31;
+  const n = Math.ceil(messageLength / multipleOf) * multipleOf;
+
+  // decrypt plain field elements and convert them into bytes
+  const message = decryptV2(cipherText, privateKey);
+  const bytes = message.map((m) => wordToBytes(m, 31));
+  return Bytes.from(bytes.flat().slice(0, messageLength - n));
+}