Add additional data support to timelock encryption for replay attack prevention

tlock.go

@@ -137,31 +137,60 @@
 // TimeLock encrypts the specified data for the given round number. The data
 // can't be decrypted until the specified round is reached by the network in use.
 func TimeLock(scheme crypto.Scheme, publicKey kyber.Point, roundNumber uint64, data []byte) (*ibe.Ciphertext, error) {
+	return TimeLockWithAdditionalData(scheme, publicKey, roundNumber, data, nil)
+}
+
+// TimeLockWithAdditionalData encrypts with additional data for replay attack prevention.
+// different AD produces different ciphertexts even for same message.
+func TimeLockWithAdditionalData(scheme crypto.Scheme, publicKey kyber.Point, roundNumber uint64, data []byte, additionalData []byte) (*ibe.Ciphertext, error) {
 	if publicKey.Equal(publicKey.Clone().Null()) {
 		return nil, ErrInvalidPublicKey
 	}
 
-	id := scheme.DigestBeacon(&chain.Beacon{
+	// Create beacon
+	beacon := &chain.Beacon{
 		Round: roundNumber,
-	})
+	}
+
+	// Get the ID from the beacon
+	id := scheme.DigestBeacon(beacon)
 
 	var cipherText *ibe.Ciphertext
 	var err error
-	switch scheme.Name {
-	case crypto.ShortSigSchemeID:
-		// the ShortSigSchemeID uses the wrong DST for G1, so we keep it for retro-compatibility
-		cipherText, err = ibe.EncryptCCAonG2(bls.NewBLS12381SuiteWithDST(bls.DefaultDomainG2(), bls.DefaultDomainG2()), publicKey, id, data)
-	case crypto.UnchainedSchemeID:
-		cipherText, err = ibe.EncryptCCAonG1(bls.NewBLS12381Suite(), publicKey, id, data)
-	case crypto.SigsOnG1ID:
-		cipherText, err = ibe.EncryptCCAonG2(bls.NewBLS12381Suite(), publicKey, id, data)
-	case crypto.BN254UnchainedOnG1SchemeID:
-		suite := bn.NewSuiteBn254()
-		suite.SetDomainG1([]byte("BLS_SIG_BN254G1_XMD:KECCAK-256_SVDW_RO_NUL_"))
-		suite.SetDomainG2([]byte("BLS_SIG_BN254G2_XMD:KECCAK-256_SVDW_RO_NUL_"))
-		cipherText, err = ibe.EncryptCCAonG2(suite, publicKey, id, data)
-	default:
-		return nil, fmt.Errorf("unsupported drand scheme '%s'", scheme.Name)
+
+	if len(additionalData) > 0 {
+		switch scheme.Name {
+		case crypto.ShortSigSchemeID:
+			// the ShortSigSchemeID uses the wrong DST for G1, so we keep it for retro-compatibility
+			cipherText, err = ibe.EncryptCCAonG2WithAD(bls.NewBLS12381SuiteWithDST(bls.DefaultDomainG2(), bls.DefaultDomainG2()), publicKey, id, data, additionalData)
+		case crypto.UnchainedSchemeID:
+			cipherText, err = ibe.EncryptCCAonG1WithAD(bls.NewBLS12381Suite(), publicKey, id, data, additionalData)
+		case crypto.SigsOnG1ID:
+			cipherText, err = ibe.EncryptCCAonG2WithAD(bls.NewBLS12381Suite(), publicKey, id, data, additionalData)
+		case crypto.BN254UnchainedOnG1SchemeID:
+			suite := bn.NewSuiteBn254()
+			suite.SetDomainG1([]byte("BLS_SIG_BN254G1_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			suite.SetDomainG2([]byte("BLS_SIG_BN254G2_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			cipherText, err = ibe.EncryptCCAonG2WithAD(suite, publicKey, id, data, additionalData)
+		default:
+			return nil, fmt.Errorf("unsupported drand scheme '%s'", scheme.Name)
+		}
+	} else {
+		switch scheme.Name {
+		case crypto.ShortSigSchemeID:
+			cipherText, err = ibe.EncryptCCAonG2(bls.NewBLS12381SuiteWithDST(bls.DefaultDomainG2(), bls.DefaultDomainG2()), publicKey, id, data)
+		case crypto.UnchainedSchemeID:
+			cipherText, err = ibe.EncryptCCAonG1(bls.NewBLS12381Suite(), publicKey, id, data)
+		case crypto.SigsOnG1ID:
+			cipherText, err = ibe.EncryptCCAonG2(bls.NewBLS12381Suite(), publicKey, id, data)
+		case crypto.BN254UnchainedOnG1SchemeID:
+			suite := bn.NewSuiteBn254()
+			suite.SetDomainG1([]byte("BLS_SIG_BN254G1_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			suite.SetDomainG2([]byte("BLS_SIG_BN254G2_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			cipherText, err = ibe.EncryptCCAonG2(suite, publicKey, id, data)
+		default:
+			return nil, fmt.Errorf("unsupported drand scheme '%s'", scheme.Name)
+		}
 	}
 
 	if err != nil {
@@ -174,50 +203,97 @@
 // TimeUnlock decrypts the specified ciphertext for the given beacon. The
 // ciphertext can't be decrypted until the specified round is reached by the network in use.
 func TimeUnlock(scheme crypto.Scheme, publicKey kyber.Point, beacon chain.Beacon, ciphertext *ibe.Ciphertext) ([]byte, error) {
+	return TimeUnlockWithAdditionalData(scheme, publicKey, beacon, ciphertext, nil)
+}
+
+// TimeUnlockWithAdditionalData decrypts with additional data verification.
+// wrong AD fails at cryptographic level during rP check.
+func TimeUnlockWithAdditionalData(scheme crypto.Scheme, publicKey kyber.Point, beacon chain.Beacon, ciphertext *ibe.Ciphertext, additionalData []byte) ([]byte, error) {
+	// Verify the beacon first
 	if err := scheme.VerifyBeacon(&beacon, publicKey); err != nil {
 		return nil, fmt.Errorf("verify beacon: %w", err)
 	}
 
-	var data []byte
+	var decryptedData []byte
 	var err error
-	switch scheme.Name {
-	case crypto.ShortSigSchemeID:
-		var signature bls.KyberG1
-		if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
-			return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
-		}
-		// the ShortSigSchemeID uses the wrong DST for G1, so we keep it for retro-compatibility
-		data, err = ibe.DecryptCCAonG2(bls.NewBLS12381SuiteWithDST(bls.DefaultDomainG2(), bls.DefaultDomainG2()), &signature, ciphertext)
-	case crypto.UnchainedSchemeID:
-		var signature bls.KyberG2
-		if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
-			return nil, fmt.Errorf("unmarshal kyber G2: %w", err)
-		}
-		data, err = ibe.DecryptCCAonG1(bls.NewBLS12381Suite(), &signature, ciphertext)
-	case crypto.SigsOnG1ID:
-		var signature bls.KyberG1
-		if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
-			return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+
+	if len(additionalData) > 0 {
+		switch scheme.Name {
+		case crypto.ShortSigSchemeID:
+			var signature bls.KyberG1
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+			}
+			suite := bls.NewBLS12381SuiteWithDST(bls.DefaultDomainG2(), bls.DefaultDomainG2())
+			decryptedData, err = ibe.DecryptCCAonG2WithAD(suite, &signature, ciphertext, additionalData)
+		case crypto.UnchainedSchemeID:
+			var signature bls.KyberG2
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G2: %w", err)
+			}
+			suite := bls.NewBLS12381Suite()
+			decryptedData, err = ibe.DecryptCCAonG1WithAD(suite, &signature, ciphertext, additionalData)
+		case crypto.SigsOnG1ID:
+			var signature bls.KyberG1
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+			}
+			suite := bls.NewBLS12381Suite()
+			decryptedData, err = ibe.DecryptCCAonG2WithAD(suite, &signature, ciphertext, additionalData)
+		case crypto.BN254UnchainedOnG1SchemeID:
+			suite := bn.NewSuiteBn254()
+			suite.SetDomainG1([]byte("BLS_SIG_BN254G1_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			suite.SetDomainG2([]byte("BLS_SIG_BN254G2_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			signature := suite.G1().Point()
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+			}
+			decryptedData, err = ibe.DecryptCCAonG2WithAD(suite, signature, ciphertext, additionalData)
+		default:
+			return nil, fmt.Errorf("unsupported drand scheme '%s'", scheme.Name)
 		}
-		data, err = ibe.DecryptCCAonG2(bls.NewBLS12381Suite(), &signature, ciphertext)
-	case crypto.BN254UnchainedOnG1SchemeID:
-		suite := bn.NewSuiteBn254()
-		suite.SetDomainG1([]byte("BLS_SIG_BN254G1_XMD:KECCAK-256_SVDW_RO_NUL_"))
-		suite.SetDomainG2([]byte("BLS_SIG_BN254G2_XMD:KECCAK-256_SVDW_RO_NUL_"))
-		signature := suite.G1().Point()
-		if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
-			return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+	} else {
+		switch scheme.Name {
+		case crypto.ShortSigSchemeID:
+			var signature bls.KyberG1
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+			}
+			suite := bls.NewBLS12381SuiteWithDST(bls.DefaultDomainG2(), bls.DefaultDomainG2())
+			decryptedData, err = ibe.DecryptCCAonG2(suite, &signature, ciphertext)
+		case crypto.UnchainedSchemeID:
+			var signature bls.KyberG2
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G2: %w", err)
+			}
+			suite := bls.NewBLS12381Suite()
+			decryptedData, err = ibe.DecryptCCAonG1(suite, &signature, ciphertext)
+		case crypto.SigsOnG1ID:
+			var signature bls.KyberG1
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+			}
+			suite := bls.NewBLS12381Suite()
+			decryptedData, err = ibe.DecryptCCAonG2(suite, &signature, ciphertext)
+		case crypto.BN254UnchainedOnG1SchemeID:
+			suite := bn.NewSuiteBn254()
+			suite.SetDomainG1([]byte("BLS_SIG_BN254G1_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			suite.SetDomainG2([]byte("BLS_SIG_BN254G2_XMD:KECCAK-256_SVDW_RO_NUL_"))
+			signature := suite.G1().Point()
+			if err := signature.UnmarshalBinary(beacon.Signature); err != nil {
+				return nil, fmt.Errorf("unmarshal kyber G1: %w", err)
+			}
+			decryptedData, err = ibe.DecryptCCAonG2(suite, signature, ciphertext)
+		default:
+			return nil, fmt.Errorf("unsupported drand scheme '%s'", scheme.Name)
 		}
-		data, err = ibe.DecryptCCAonG2(suite, signature, ciphertext)
-	default:
-		return nil, fmt.Errorf("unsupported drand scheme '%s'", scheme.Name)
 	}
 
 	if err != nil {
 		return nil, fmt.Errorf("decrypt dek: %w", err)
 	}
 
-	return data, nil
+	return decryptedData, nil
 }
 
 // =============================================================================

tlock_test.go

@@ -304,6 +304,165 @@ func TestTimeLockUnlock(t *testing.T) {
 	}
 }
 
+func TestTimeLockUnlockWithAdditionalData(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping live testing in short mode")
+	}
+	tests := []struct {
+		name      string
+		host      string
+		chainhash string
+	}{
+		{
+			"quicknetT",
+			testnetHost,
+			testnetQuicknetT,
+		},
+		{
+			"quicknet",
+			mainnetHost,
+			mainnetQuicknet,
+		},
+		{
+			"evmnet",
+			mainnetHost,
+			mainnetEvm,
+		},
+	}
+	for _, tt := range tests {
+		t.Run(tt.name, func(t *testing.T) {
+			network, err := http.NewNetwork(tt.host, tt.chainhash)
+			require.NoError(t, err)
+
+			futureRound := network.RoundNumber(time.Now())
+
+			id, err := network.Signature(futureRound)
+			require.NoError(t, err)
+
+			data := []byte(`sensitive data`)
+			additionalData := []byte("tx_hash_0x1234567890abcdef")
+
+			// Encrypt with additional data
+			cipherText, err := tlock.TimeLockWithAdditionalData(network.Scheme(), network.PublicKey(), futureRound, data, additionalData)
+			require.NoError(t, err)
+
+			beacon := chain.Beacon{
+				Round:     futureRound,
+				Signature: id,
+			}
+
+			// Decrypt with correct additional data should succeed
+			b, err := tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherText, additionalData)
+			require.NoError(t, err)
+			require.Equal(t, data, b, "decrypted data should match original")
+
+			// Decrypt with incorrect additional data should fail
+			incorrectAD := []byte("tx_hash_0xDIFFERENT")
+			_, err = tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherText, incorrectAD)
+			require.Error(t, err, "decryption with incorrect additional data should fail")
+			require.Contains(t, err.Error(), "invalid proof", "error should indicate cryptographic failure")
+
+			// Decrypt without additional data should fail
+			_, err = tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherText, nil)
+			require.Error(t, err, "decryption without additional data should fail")
+			require.Contains(t, err.Error(), "invalid proof", "error should indicate cryptographic failure")
+		})
+	}
+}
+
+func TestReplayAttackPreventionWithAdditionalData(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping live testing in short mode")
+	}
+
+	network, err := http.NewNetwork(testnetHost, testnetQuicknetT)
+	require.NoError(t, err)
+
+	futureRound := network.RoundNumber(time.Now())
+
+	id, err := network.Signature(futureRound)
+	require.NoError(t, err)
+
+	beacon := chain.Beacon{
+		Round:     futureRound,
+		Signature: id,
+	}
+
+	// Same message, different contexts
+	message := []byte("transfer 100 tokens")
+	contextAlice := []byte("recipient:alice|nonce:1")
+	contextBob := []byte("recipient:bob|nonce:1")
+
+	// Encrypt for Alice's context
+	cipherTextAlice, err := tlock.TimeLockWithAdditionalData(network.Scheme(), network.PublicKey(), futureRound, message, contextAlice)
+	require.NoError(t, err)
+
+	// Encrypt for Bob's context
+	cipherTextBob, err := tlock.TimeLockWithAdditionalData(network.Scheme(), network.PublicKey(), futureRound, message, contextBob)
+	require.NoError(t, err)
+
+	// Ciphertexts should be different even though message is the same
+	require.NotEqual(t, cipherTextAlice.U, cipherTextBob.U, "different contexts should produce different ciphertexts")
+
+	// Decrypt Alice's ciphertext with Alice's context - should succeed
+	decryptedAlice, err := tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherTextAlice, contextAlice)
+	require.NoError(t, err)
+	require.Equal(t, message, decryptedAlice)
+
+	// Try to "replay" Alice's ciphertext in Bob's context - should fail
+	_, err = tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherTextAlice, contextBob)
+	require.Error(t, err, "replay attack should be prevented")
+	require.Contains(t, err.Error(), "invalid proof")
+
+	// Decrypt Bob's ciphertext with Bob's context - should succeed
+	decryptedBob, err := tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherTextBob, contextBob)
+	require.NoError(t, err)
+	require.Equal(t, message, decryptedBob)
+}
+
+func TestBackwardCompatibilityWithoutAdditionalData(t *testing.T) {
+	if testing.Short() {
+		t.Skip("skipping live testing in short mode")
+	}
+
+	network, err := http.NewNetwork(testnetHost, testnetQuicknetT)
+	require.NoError(t, err)
+
+	futureRound := network.RoundNumber(time.Now())
+
+	id, err := network.Signature(futureRound)
+	require.NoError(t, err)
+
+	data := []byte(`test data`)
+
+	// Encrypt without additional data using TimeLock
+	cipherText1, err := tlock.TimeLock(network.Scheme(), network.PublicKey(), futureRound, data)
+	require.NoError(t, err)
+
+	// Encrypt without additional data using TimeLockWithAdditionalData with nil AD
+	cipherText2, err := tlock.TimeLockWithAdditionalData(network.Scheme(), network.PublicKey(), futureRound, data, nil)
+	require.NoError(t, err)
+
+	beacon := chain.Beacon{
+		Round:     futureRound,
+		Signature: id,
+	}
+
+	// Both should decrypt successfully
+	b1, err := tlock.TimeUnlock(network.Scheme(), network.PublicKey(), beacon, cipherText1)
+	require.NoError(t, err)
+	require.Equal(t, data, b1)
+
+	b2, err := tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherText2, nil)
+	require.NoError(t, err)
+	require.Equal(t, data, b2)
+
+	// Decrypt with TimeUnlock using TimeUnlockWithAdditionalData should also work
+	b3, err := tlock.TimeUnlockWithAdditionalData(network.Scheme(), network.PublicKey(), beacon, cipherText1, nil)
+	require.NoError(t, err)
+	require.Equal(t, data, b3)
+}
+
 func TestCannotEncryptWithPointAtInfinity(t *testing.T) {
 	suite := bls.NewBLS12381Suite()
 	t.Run("on G2", func(t *testing.T) {