From 61c3a4d8f4eeb1f6907a1e20a75f9203d2040d89 Mon Sep 17 00:00:00 2001
From: azuchi <azuchi@chaintope.com>
Date: Sun, 21 Jan 2024 11:54:21 +0900
Subject: [PATCH] Implement Bitcoin::BIP324::Cipher#decrypt

---
 lib/bitcoin/bip324/cipher.rb             | 42 ++++++++++++++++++++----
 lib/bitcoin/bip324/fs_chacha20.rb        |  2 +-
 lib/bitcoin/bip324/fs_chacha_poly1305.rb | 18 ++++++----
 spec/bitcoin/bip324_spec.rb              | 19 +++++++++--
 4 files changed, 65 insertions(+), 16 deletions(-)

diff --git a/lib/bitcoin/bip324/cipher.rb b/lib/bitcoin/bip324/cipher.rb
index 79d1aae..5fae65b 100644
--- a/lib/bitcoin/bip324/cipher.rb
+++ b/lib/bitcoin/bip324/cipher.rb
@@ -5,6 +5,9 @@ class Cipher
       include Bitcoin::Util
 
       HEADER = [1 << 7].pack('C')
+      HEADER_LEN = 1
+      LENGTH_LEN = 3
+      EXPANSION = LENGTH_LEN + HEADER_LEN + 16
 
       attr_reader :key
       attr_reader :our_pubkey
@@ -30,22 +33,28 @@ def initialize(key, our_pubkey = nil)
       # Setup when the other side's public key is received.
       # @param [Bitcoin::BIP324::EllSwiftPubkey] their_pubkey
       # @param [Boolean] initiator Set true if we are the initiator establishing the v2 P2P connection.
-      def setup(their_pubkey, initiator)
+      # @param [Boolean] self_decrypt only for testing, and swaps encryption/decryption keys, so that encryption
+      # and decryption can be tested without knowing the other side's private key.
+      def setup(their_pubkey, initiator, self_decrypt = false)
         salt = 'bitcoin_v2_shared_secret' + Bitcoin.chain_params.magic_head.htb
         ecdh_secret = BIP324.v2_ecdh(key.priv_key, their_pubkey, our_pubkey, initiator).htb
         terminator = hkdf_sha256(ecdh_secret, salt, 'garbage_terminators')
-        if initiator
+        side = initiator != self_decrypt
+        if side
           self.send_l_cipher = FSChaCha20.new(hkdf_sha256(ecdh_secret, salt, 'initiator_L'))
           self.send_p_cipher = FSChaCha20Poly1305.new(hkdf_sha256(ecdh_secret, salt, 'initiator_P'))
           self.recv_l_cipher = FSChaCha20.new(hkdf_sha256(ecdh_secret, salt, 'responder_L'))
           self.recv_p_cipher = FSChaCha20Poly1305.new(hkdf_sha256(ecdh_secret, salt, 'responder_P'))
-          self.send_garbage_terminator = terminator[0...16].bth
-          self.recv_garbage_terminator = terminator[16..-1].bth
         else
           self.recv_l_cipher = FSChaCha20.new(hkdf_sha256(ecdh_secret, salt, 'initiator_L'))
           self.recv_p_cipher = FSChaCha20Poly1305.new(hkdf_sha256(ecdh_secret, salt, 'initiator_P'))
           self.send_l_cipher = FSChaCha20.new(hkdf_sha256(ecdh_secret, salt, 'responder_L'))
           self.send_p_cipher = FSChaCha20Poly1305.new(hkdf_sha256(ecdh_secret, salt, 'responder_P'))
+        end
+        if initiator
+          self.send_garbage_terminator = terminator[0...16].bth
+          self.recv_garbage_terminator = terminator[16..-1].bth
+        else
           self.recv_garbage_terminator = terminator[0...16].bth
           self.send_garbage_terminator = terminator[16..-1].bth
         end
@@ -53,7 +62,9 @@ def setup(their_pubkey, initiator)
       end
 
       # Encrypt a packet. Only after setup.
-      #
+      # @param [String] contents Packet with binary format.
+      # @param [String] aad AAD
+      # @param [Boolean] ignore Whether contains ignore bit or not.
       def encrypt(contents, aad: '', ignore: false)
         raise RuntimeError, "contents size over." unless contents.bytesize <= (2**24 - 1)
 
@@ -71,8 +82,27 @@ def encrypt(contents, aad: '', ignore: false)
         enc_plaintext_len + aead_ciphertext
       end
 
-      def decrypt
+      # Decrypt a packet. Only after setup.
+      # @param [String] input Packet to be decrypt.
+      # @param [String] aad AAD
+      # @param [Boolean] ignore Whether contains ignore bit or not.
+      # @return [String] Plaintext
+      def decrypt(input, aad: '', ignore: false)
+        len = decrypt_length(input[0...Bitcoin::BIP324::Cipher::LENGTH_LEN])
+        raise RuntimeError, "Packet length invalid." unless input.bytesize == len + EXPANSION
+        recv_p_cipher.decrypt(aad, input[Bitcoin::BIP324::Cipher::LENGTH_LEN..-1])
+      end
+
+      private
 
+      # Decrypt the length of a packet. Only after setup.
+      # @param [String] input Length packet with binary format.
+      # @return [Integer] length
+      def decrypt_length(input)
+        raise ArgumentError, "input length must be #{LENGTH_LEN}" unless input.bytesize == LENGTH_LEN
+        ret = recv_l_cipher.decrypt(input)
+        b0, b1, b2 = ret.unpack('CCC')
+        b0 + (b1 << 8) + (b2 << 16)
       end
     end
   end
diff --git a/lib/bitcoin/bip324/fs_chacha20.rb b/lib/bitcoin/bip324/fs_chacha20.rb
index d3ba017..b1d204a 100644
--- a/lib/bitcoin/bip324/fs_chacha20.rb
+++ b/lib/bitcoin/bip324/fs_chacha20.rb
@@ -95,7 +95,7 @@ def encrypt(chunk)
       # Decrypt a chunk
       # @param [String] chunk Chunk data with binary format.
       # @return [String] Decrypted data with binary format.
-      def decript(chunk)
+      def decrypt(chunk)
         crypt(chunk)
       end
 
diff --git a/lib/bitcoin/bip324/fs_chacha_poly1305.rb b/lib/bitcoin/bip324/fs_chacha_poly1305.rb
index e99e8bb..753426e 100644
--- a/lib/bitcoin/bip324/fs_chacha_poly1305.rb
+++ b/lib/bitcoin/bip324/fs_chacha_poly1305.rb
@@ -4,6 +4,7 @@ module BIP324
     class Poly1305
 
       MODULUS = 2**130 - 5
+      TAG_LEN = 16
 
       attr_reader :r
       attr_reader :s
@@ -33,7 +34,7 @@ def add(msg, length: nil, padding: false)
       # Compute the poly1305 tag.
       # @return Poly1305 tag wit binary format.
       def tag
-        ECDSA::Format::IntegerOctetString.encode((acc + s) & 0xffffffffffffffffffffffffffffffff, 16).reverse
+        ECDSA::Format::IntegerOctetString.encode((acc + s) & 0xffffffffffffffffffffffffffffffff, TAG_LEN).reverse
       end
     end
 
@@ -59,8 +60,12 @@ def encrypt(aad, plaintext)
       end
 
       # Decrypt a *ciphertext* with a specified +aad+.
+      # @param [String] aad AAD
+      # @param [String] ciphertext Data to be decrypted with binary format.
+      # @return [Array] [header, plaintext]
       def decrypt(aad, ciphertext)
-        crypt(aad, ciphertext, true)
+        contents = crypt(aad, ciphertext, true)
+        [contents[0], contents[1..-1]]
       end
 
       private
@@ -110,15 +115,14 @@ def chacha20_poly1305_decrypt(key, nonce, aad, ciphertext)
         poly1305.add(ciphertext, length: msg_len, padding: true)
         poly1305.add([aad.bytesize, msg_len].pack("Q<Q<"))
         return nil unless ciphertext[-16..-1] == poly1305.tag
-        ret = ""
-        ((msg_len + 63) / 64).times do |i|
+        ret = ((msg_len + 63) / 64).times.map do |i|
           now = [64, msg_len - 64 * i].min
           keystream = ChaCha20.block(key, nonce, i + 1)
-          now.times do |j|
-            ret += [(ciphertext[j + 64 * i].unpack1('C') ^ keystream[j].unpack1('C'))].pack('C')
+          now.times.map do |j|
+            ciphertext[j + 64 * i].unpack1('C') ^ keystream[j].unpack1('C')
           end
         end
-        ret
+        ret.flatten.pack('C*')
       end
     end
   end
diff --git a/spec/bitcoin/bip324_spec.rb b/spec/bitcoin/bip324_spec.rb
index fa1b4bd..2cecd78 100644
--- a/spec/bitcoin/bip324_spec.rb
+++ b/spec/bitcoin/bip324_spec.rb
@@ -77,18 +77,33 @@ def test_ellswift_xdh
         expect(cipher.session_id).to eq(v['out_session_id'])
 
         in_index = v['in_idx'].to_i
-        in_index.times do
+        dummies = in_index.times.map do
           cipher.encrypt("")
         end
         aad = v['in_aad'] ? v['in_aad'].htb : ''
         contents = v['in_contents'].htb * v['in_multiply'].to_i
-        ciphertext = cipher.encrypt(contents, aad: aad, ignore: v['in_ignore'] == '1')
+        ignore = v['in_ignore'] == '1'
+        ciphertext = cipher.encrypt(contents, aad: aad, ignore: ignore)
         if v['out_ciphertext']
           expect(ciphertext.bth).to eq(v['out_ciphertext'])
         end
         if v['out_ciphertext_endswith']
           expect(ciphertext.bth).to end_with(v['out_ciphertext_endswith'])
         end
+
+        # Decrypt
+        dec_cipher = Bitcoin::BIP324::Cipher.new(our_priv, our_ell)
+        dec_cipher.setup(their_ell, initiating, true)
+        expect(dec_cipher.session_id).to eq(v['out_session_id'])
+        expect(dec_cipher.send_garbage_terminator).to eq(v['mid_send_garbage_terminator'])
+        expect(dec_cipher.recv_garbage_terminator).to eq(v['mid_recv_garbage_terminator'])
+
+        in_index.times do |i|
+          dec_cipher.decrypt(dummies[i])
+        end
+
+        _, plaintext = dec_cipher.decrypt(ciphertext, aad: aad, ignore: ignore)
+        expect(plaintext.bth).to eq(contents.bth)
       end
     end
   end