From 58edb3abe6605d86becdcc3f343beb6feee47a50 Mon Sep 17 00:00:00 2001
From: Oleh Tryfonov <oleg.himik.trifon@gmail.com>
Date: Mon, 26 Aug 2024 16:47:17 +0300
Subject: [PATCH] Added range for Twisted ElGamal decryption

---
 examples/typescript/twisted_el_gamal.ts | 18 +++++++++++++-----
 src/core/crypto/twistedElGamal.ts       | 21 ++++++++++++++-------
 2 files changed, 27 insertions(+), 12 deletions(-)

diff --git a/examples/typescript/twisted_el_gamal.ts b/examples/typescript/twisted_el_gamal.ts
index f57424f4..c78ceffc 100644
--- a/examples/typescript/twisted_el_gamal.ts
+++ b/examples/typescript/twisted_el_gamal.ts
@@ -6,7 +6,7 @@
 
 import { TwistedElGamal, TwistedEd25519PrivateKey } from "@aptos-labs/ts-sdk";
 
-const AMOUNT = BigInt(12345)
+const AMOUNT = BigInt(12_345)
 
 
 const example = async () => {
@@ -15,7 +15,7 @@ const example = async () => {
   const privateKey = TwistedEd25519PrivateKey.generate()
   console.log("=== Key pair ===");
   console.log(`Private key: ${privateKey.toString()}`);
-  console.log(`Pablic key: ${privateKey.publicKey().toString()}\n\n`);
+  console.log(`Public key: ${privateKey.publicKey().toString()}\n\n`);
 
   const twistedElGamalInstance = new TwistedElGamal(privateKey)
 
@@ -29,9 +29,17 @@ const example = async () => {
   const decAmount1 = twistedElGamalInstance.decrypt(ciphertext)
   console.log(`Decrypted amount: ${decAmount1}\n`);
 
-  console.log("Decrypting the ciphertext starting at a specified amount");
-  const decAmount2 = twistedElGamalInstance.decrypt(ciphertext, 1000n)
-  console.log(`Decrypted amount: ${decAmount2}`);
+  console.log("Decrypting a ciphertext using a valid range");
+  const decAmount2 = twistedElGamalInstance.decrypt(ciphertext, { start: 1_000n, end: 20_000n})
+  console.log(`Decrypted amount: ${decAmount2}\n`);
+
+  console.log("Decrypting a ciphertext using a invalid range");
+  try {
+    twistedElGamalInstance.decrypt(ciphertext, { start: 1_000n, end: 2_000n})
+  } catch (e) {
+    console.error(e);
+    console.log("Failed to decrypt the amount");
+  }
 };
 
 
diff --git a/src/core/crypto/twistedElGamal.ts b/src/core/crypto/twistedElGamal.ts
index ece9780b..c4bf45d9 100644
--- a/src/core/crypto/twistedElGamal.ts
+++ b/src/core/crypto/twistedElGamal.ts
@@ -7,6 +7,11 @@ import { TwistedEd25519PrivateKey, TwistedEd25519PublicKey } from "./twistedEd25
 
 export type RistPoint = InstanceType<typeof RistrettoPoint>
 
+export interface DecryptionRange {
+  start?: bigint;
+  end?: bigint;
+}
+
 /**
  * Twisted ElGamal encryption/decryption
  * @see {@link https://drive.google.com/file/d/1wGo-pIOPOcCQA0gjngE5kmWUQ-TxktAF/view | Veiled coins with twisted ElGamal}
@@ -47,10 +52,10 @@ export class TwistedElGamal {
    * Decrypts the amount with Twisted ElGamal
    * 
    * @param ciphertext сiphertext points encrypted by Twisted ElGamal
-   * @param startAmount Start amount from which the decryption will begin
+   * @param decryptionRange The range of amounts to be used in decryption
    */
-  public decrypt(ciphertext: TwistedElGamalCiphertext, startAmount?: bigint): bigint {
-    return TwistedElGamal.decryptWithSK(ciphertext, this.privateKey, startAmount)
+  public decrypt(ciphertext: TwistedElGamalCiphertext, decryptionRange?: DecryptionRange): bigint {
+    return TwistedElGamal.decryptWithSK(ciphertext, this.privateKey, decryptionRange)
   }
 
   /**
@@ -85,12 +90,12 @@ export class TwistedElGamal {
    * Decrypts the amount with Twisted ElGamal
    * @param ciphertext сiphertext points encrypted by Twisted ElGamal
    * @param privateKey Twisted ElGamal Ed25519 private key.
-   * @param startAmount Start amount from which the decryption will begin
+   * @param decryptionRange The range of amounts to be used in decryption
    */
   static decryptWithSK(
     ciphertext: TwistedElGamalCiphertext,
     privateKey: TwistedEd25519PrivateKey,
-    startAmount?: bigint
+    decryptionRange?: DecryptionRange
   ): bigint {
     const { C, D } = ciphertext
     const H = RistrettoPoint.fromHex(TwistedElGamal.HASH_BASE_POINT)
@@ -98,7 +103,7 @@ export class TwistedElGamal {
     const sD = RistrettoPoint.fromHex(D.toRawBytes()).multiply(modS)
     const mH = RistrettoPoint.fromHex(C.toRawBytes()).subtract(sD)
 
-    let amount = startAmount ?? BigInt(0)
+    let amount = decryptionRange?.start ?? BigInt(0)
     if (amount === BigInt(0)){
       if (mH.equals(RistrettoPoint.ZERO)) return BigInt(0)
 
@@ -106,9 +111,11 @@ export class TwistedElGamal {
     } 
 
     let searchablePoint = H.multiply(amount)
+    const endAmount = decryptionRange?.end ?? ed25519.CURVE.n
 
     while (!mH.equals(searchablePoint)) {
-      if (amount >= ed25519.CURVE.n) throw new Error("Error when decrypting the amount")
+      if (amount >= endAmount)
+        throw new Error("Error when decrypting the amount: it is not possible to decrypt the amount in the specified range")
 
       amount += BigInt(1)
       searchablePoint = searchablePoint.add(H)