From f23e783600f958efd66b496e888558a4e91e46c8 Mon Sep 17 00:00:00 2001
From: Sean Bowe <ewillbefull@gmail.com>
Date: Mon, 16 Jul 2018 18:34:14 -0600
Subject: [PATCH] Decryption and tests of note/outgoing encryption.

---
 src/gtest/test_noteencryption.cpp | 94 +++++++++++++++++++++++++++++++
 src/zcash/Note.cpp                | 49 +++++++++++++++-
 src/zcash/Note.hpp                | 16 +++++-
 3 files changed, 157 insertions(+), 2 deletions(-)

diff --git a/src/gtest/test_noteencryption.cpp b/src/gtest/test_noteencryption.cpp
index 1aded18f2..8c7a1cb48 100644
--- a/src/gtest/test_noteencryption.cpp
+++ b/src/gtest/test_noteencryption.cpp
@@ -4,6 +4,7 @@
 #include <array>
 #include <stdexcept>
 
+#include "zcash/Note.hpp"
 #include "zcash/NoteEncryption.hpp"
 #include "zcash/prf.h"
 #include "zcash/Address.hpp"
@@ -19,6 +20,99 @@ public:
     }
 };
 
+TEST(noteencryption, NotePlaintext)
+{
+    using namespace libzcash;
+    auto sk = SaplingSpendingKey(uint256()).expanded_spending_key();
+    auto vk = sk.full_viewing_key();
+    auto ivk = vk.in_viewing_key();
+    SaplingPaymentAddress addr = *ivk.address({0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0});
+
+    std::array<unsigned char, ZC_MEMO_SIZE> memo;
+    for (size_t i = 0; i < ZC_MEMO_SIZE; i++) {
+        // Fill the message with dummy data
+        memo[i] = (unsigned char) i;
+    }
+
+    SaplingNote note(addr, 39393);
+    SaplingNotePlaintext pt(note, memo);
+
+    auto res = pt.encrypt(addr.pk_d);
+    if (!res) {
+        FAIL();
+    }
+
+    auto enc = res.get();
+
+    auto ct = enc.first;
+    auto encryptor = enc.second;
+    auto epk = encryptor.get_epk();
+
+    // Try to decrypt
+    auto foo = SaplingNotePlaintext::decrypt(
+        ct,
+        ivk,
+        epk
+    );
+
+    if (!foo) {
+        FAIL();
+    }
+
+    auto bar = foo.get();
+
+    ASSERT_TRUE(bar.value() == pt.value());
+    ASSERT_TRUE(bar.memo() == pt.memo());
+    ASSERT_TRUE(bar.d == pt.d);
+    ASSERT_TRUE(bar.rcm == pt.rcm);
+
+    auto foobar = bar.note(ivk);
+
+    if (!foobar) {
+        FAIL();
+    }
+
+    auto new_note = foobar.get();
+
+    ASSERT_TRUE(note.value() == new_note.value());
+    ASSERT_TRUE(note.d == new_note.d);
+    ASSERT_TRUE(note.pk_d == new_note.pk_d);
+    ASSERT_TRUE(note.r == new_note.r);
+    ASSERT_TRUE(note.cm() == new_note.cm());
+
+    SaplingOutgoingPlaintext out_pt;
+    out_pt.pk_d = note.pk_d;
+    out_pt.esk = encryptor.get_esk();
+
+    auto ovk = random_uint256();
+    auto cv = random_uint256();
+    auto cm = random_uint256();
+
+    auto out_ct = out_pt.encrypt(
+        ovk,
+        cv,
+        cm,
+        encryptor
+    );
+
+    auto decrypted_out_ct = out_pt.decrypt(
+        out_ct,
+        ovk,
+        cv,
+        cm,
+        encryptor.get_epk()
+    );
+
+    if (!decrypted_out_ct) {
+        FAIL();
+    }
+
+    auto decrypted_out_ct_unwrapped = decrypted_out_ct.get();
+
+    ASSERT_TRUE(decrypted_out_ct_unwrapped.pk_d == out_pt.pk_d);
+    ASSERT_TRUE(decrypted_out_ct_unwrapped.esk == out_pt.esk);
+}
+
 TEST(noteencryption, SaplingApi)
 {
     using namespace libzcash;
diff --git a/src/zcash/Note.cpp b/src/zcash/Note.cpp
index 7e776c756..42df3bf2e 100644
--- a/src/zcash/Note.cpp
+++ b/src/zcash/Note.cpp
@@ -154,12 +154,59 @@ boost::optional<SaplingNote> SaplingNotePlaintext::note(const SaplingIncomingVie
 {
     auto addr = ivk.address( d );
     if (addr) {
-        return SaplingNote(addr.get(), value_);
+        return SaplingNote(d, addr.get().pk_d, value_, rcm);
     } else {
         return boost::none;
     }
 }
 
+boost::optional<SaplingOutgoingPlaintext> SaplingOutgoingPlaintext::decrypt(
+    const SaplingOutCiphertext &ciphertext,
+    const uint256& ovk,
+    const uint256& cv,
+    const uint256& cm,
+    const uint256& epk
+)
+{
+    auto pt = AttemptSaplingOutDecryption(ciphertext, ovk, cv, cm, epk);
+    if (!pt) {
+        return boost::none;
+    }
+
+    // Deserialize from the plaintext
+    CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
+    ss << pt.get();
+
+    SaplingOutgoingPlaintext ret;
+    ss >> ret;
+
+    assert(ss.size() == 0);
+
+    return ret;
+}
+
+boost::optional<SaplingNotePlaintext> SaplingNotePlaintext::decrypt(
+    const SaplingEncCiphertext &ciphertext,
+    const uint256 &ivk,
+    const uint256 &epk
+)
+{
+    auto pt = AttemptSaplingEncDecryption(ciphertext, ivk, epk);
+    if (!pt) {
+        return boost::none;
+    }
+
+    // Deserialize from the plaintext
+    CDataStream ss(SER_NETWORK, PROTOCOL_VERSION);
+    ss << pt.get();
+
+    SaplingNotePlaintext ret;
+    ss >> ret;
+
+    assert(ss.size() == 0);
+
+    return ret;
+}
 
 boost::optional<SaplingNotePlaintextEncryptionResult> SaplingNotePlaintext::encrypt(const uint256& pk_d) const
 {
diff --git a/src/zcash/Note.hpp b/src/zcash/Note.hpp
index 8bdc2de3b..99d34f47b 100644
--- a/src/zcash/Note.hpp
+++ b/src/zcash/Note.hpp
@@ -116,7 +116,7 @@ public:
                                         ) const;
 };
 
-typedef std::pair<std::reference_wrapper<const SaplingEncCiphertext>, std::reference_wrapper<const SaplingNoteEncryption>> SaplingNotePlaintextEncryptionResult;
+typedef std::pair<SaplingEncCiphertext, SaplingNoteEncryption> SaplingNotePlaintextEncryptionResult;
 
 class SaplingNotePlaintext : public BaseNotePlaintext {
 public:
@@ -127,6 +127,12 @@ public:
 
     SaplingNotePlaintext(const SaplingNote& note, std::array<unsigned char, ZC_MEMO_SIZE> memo);
 
+    static boost::optional<SaplingNotePlaintext> decrypt(
+        const SaplingEncCiphertext &ciphertext,
+        const uint256 &ivk,
+        const uint256 &epk
+    );
+
     boost::optional<SaplingNote> note(const SaplingIncomingViewingKey& ivk) const;
 
     virtual ~SaplingNotePlaintext() {}
@@ -169,6 +175,14 @@ public:
         READWRITE(esk);         // 8 bytes
     }
 
+    static boost::optional<SaplingOutgoingPlaintext> decrypt(
+        const SaplingOutCiphertext &ciphertext,
+        const uint256& ovk,
+        const uint256& cv,
+        const uint256& cm,
+        const uint256& epk
+    );
+
     SaplingOutCiphertext encrypt(
         const uint256& ovk,
         const uint256& cv,