From 4dfd23e0b68fd5b9afbea3037de84ce733eda5b2 Mon Sep 17 00:00:00 2001
From: hanako mumei <81144685+2501babe@users.noreply.github.com>
Date: Wed, 7 Sep 2022 13:51:55 -0700
Subject: [PATCH] token-client: memo now explicitly declares signers

---
 token/cli/src/main.rs                         |  6 +--
 token/client/src/token.rs                     | 40 +++++++++++++++++--
 .../program-2022-test/tests/memo_transfer.rs  |  2 +-
 3 files changed, 40 insertions(+), 8 deletions(-)

diff --git a/token/cli/src/main.rs b/token/cli/src/main.rs
index aefff138..e9b723de 100644
--- a/token/cli/src/main.rs
+++ b/token/cli/src/main.rs
@@ -393,7 +393,7 @@ async fn command_create_token(
     }
 
     if let Some(text) = memo {
-        token.with_memo(text);
+        token.with_memo(text, vec![config.default_signer.pubkey()]);
     }
 
     let res = token
@@ -1073,7 +1073,7 @@ async fn command_burn(
 
     let token = token_client_from_config(config, &mint_info.program_id, &mint_info.address);
     if let Some(text) = memo {
-        token.with_memo(text);
+        token.with_memo(text, vec![config.default_signer.pubkey()]);
     }
 
     let res = token
@@ -1120,7 +1120,7 @@ async fn command_mint(
 
     let token = token_client_from_config(config, &mint_info.program_id, &mint_info.address);
     if let Some(text) = memo {
-        token.with_memo(text);
+        token.with_memo(text, vec![config.default_signer.pubkey()]);
     }
 
     let res = token
diff --git a/token/client/src/token.rs b/token/client/src/token.rs
index 28951b31..1c4c235e 100644
--- a/token/client/src/token.rs
+++ b/token/client/src/token.rs
@@ -60,6 +60,8 @@ pub enum TokenError {
     AccountDecryption,
     #[error("not enough funds in account")]
     NotEnoughFunds,
+    #[error("missing memo signer")]
+    MissingMemoSigner,
 }
 impl PartialEq for TokenError {
     fn eq(&self, other: &Self) -> bool {
@@ -77,6 +79,7 @@ impl PartialEq for TokenError {
             ) => true,
             (Self::AccountDecryption, Self::AccountDecryption) => true,
             (Self::NotEnoughFunds, Self::NotEnoughFunds) => true,
+            (Self::MissingMemoSigner, Self::MissingMemoSigner) => true,
             _ => false,
         }
     }
@@ -172,6 +175,20 @@ impl ExtensionInitializationParams {
 
 pub type TokenResult<T> = Result<T, TokenError>;
 
+#[derive(Debug)]
+struct TokenMemo {
+    text: String,
+    signers: Vec<Pubkey>,
+}
+impl TokenMemo {
+    pub fn to_instruction(&self) -> Instruction {
+        spl_memo::build_memo(
+            self.text.as_bytes(),
+            &self.signers.iter().collect::<Vec<_>>(),
+        )
+    }
+}
+
 pub struct Token<T> {
     client: Arc<dyn ProgramClient<T>>,
     pubkey: Pubkey, /*token mint*/
@@ -179,7 +196,7 @@ pub struct Token<T> {
     program_id: Pubkey,
     nonce_account: Option<Pubkey>,
     nonce_authority: Option<Pubkey>,
-    memo: Arc<RwLock<Option<String>>>,
+    memo: Arc<RwLock<Option<TokenMemo>>>,
 }
 
 impl<T> fmt::Debug for Token<T> {
@@ -187,6 +204,9 @@ impl<T> fmt::Debug for Token<T> {
         f.debug_struct("Token")
             .field("pubkey", &self.pubkey)
             .field("payer", &self.payer.pubkey())
+            .field("program_id", &self.program_id)
+            .field("nonce_account", &self.nonce_account)
+            .field("nonce_authority", &self.nonce_authority)
             .field("memo", &self.memo.read().unwrap())
             .finish()
     }
@@ -242,9 +262,12 @@ where
         }
     }
 
-    pub fn with_memo<M: AsRef<str>>(&self, memo: M) -> &Self {
+    pub fn with_memo<M: AsRef<str>>(&self, memo: M, signers: Vec<Pubkey>) -> &Self {
         let mut w_memo = self.memo.write().unwrap();
-        *w_memo = Some(memo.as_ref().to_string());
+        *w_memo = Some(TokenMemo {
+            text: memo.as_ref().to_string(),
+            signers,
+        });
         self
     }
 
@@ -305,7 +328,16 @@ where
         {
             let mut w_memo = self.memo.write().unwrap();
             if let Some(memo) = w_memo.take() {
-                instructions.push(spl_memo::build_memo(memo.as_bytes(), &[&payer_key]));
+                let signing_pubkeys = signing_keypairs.pubkeys();
+                if !memo
+                    .signers
+                    .iter()
+                    .all(|signer| signing_pubkeys.contains(signer))
+                {
+                    return Err(TokenError::MissingMemoSigner);
+                }
+
+                instructions.push(memo.to_instruction());
             }
         }
 
diff --git a/token/program-2022-test/tests/memo_transfer.rs b/token/program-2022-test/tests/memo_transfer.rs
index ead1c0e3..05aeb954 100644
--- a/token/program-2022-test/tests/memo_transfer.rs
+++ b/token/program-2022-test/tests/memo_transfer.rs
@@ -131,7 +131,7 @@ async fn test_memo_transfers(
 
     // transfer with memo
     token
-        .with_memo("🦖")
+        .with_memo("🦖", vec![alice.pubkey()])
         .transfer(
             &alice_account,
             &bob_account,