encapsulates AuthorizeNonceAccount implementation in the nonce module (#26202)

Follow up commit removes feature gate code separating durable nonce from
blockhash domain. This commit allows to encapsulate any logic
distinguishing legacy vs current nonce versions in the nonce module
after removing the feature gate.

runtime/src/nonce_keyed_account.rs

@@ -5,7 +5,7 @@ use {
         instruction::{checked_add, InstructionError},
         nonce::{
             self,
-            state::{DurableNonce, Versions},
+            state::{AuthorizeNonceError, DurableNonce, Versions},
             State,
         },
         pubkey::Pubkey,
@@ -268,30 +268,12 @@ pub fn authorize_nonce_account(
         );
         return Err(InstructionError::InvalidArgument);
     }
-
-    let state: Versions = account.get_state()?;
-    let separate_domains = state.separate_domains();
-    match state.state() {
-        State::Initialized(data) => {
-            if !signers.contains(&data.authority) {
-                ic_msg!(
-                    invoke_context,
-                    "Authorize nonce account: Account {} must sign",
-                    data.authority
-                );
-                return Err(InstructionError::MissingRequiredSignature);
-            }
-            let new_data = nonce::state::Data::new(
-                *nonce_authority,
-                data.durable_nonce,
-                data.get_lamports_per_signature(),
-            );
-            account.set_state(&Versions::new(
-                State::Initialized(new_data),
-                separate_domains,
-            ))
-        }
-        State::Uninitialized => {
+    match account
+        .get_state::<Versions>()?
+        .authorize(signers, *nonce_authority)
+    {
+        Ok(versions) => account.set_state(&versions),
+        Err(AuthorizeNonceError::Uninitialized) => {
             ic_msg!(
                 invoke_context,
                 "Authorize nonce account: Account {} state is invalid",
@@ -302,6 +284,14 @@ pub fn authorize_nonce_account(
                 merge_nonce_error_into_system_error,
             ))
         }
+        Err(AuthorizeNonceError::MissingRequiredSignature(account_authority)) => {
+            ic_msg!(
+                invoke_context,
+                "Authorize nonce account: Account {} must sign",
+                account_authority
+            );
+            Err(InstructionError::MissingRequiredSignature)
+        }
     }
 }
 

sdk/program/src/nonce/state/mod.rs

@@ -3,8 +3,9 @@
 mod current;
 pub use current::{Data, DurableNonce, State};
 use {
-    crate::hash::Hash,
+    crate::{hash::Hash, pubkey::Pubkey},
     serde_derive::{Deserialize, Serialize},
+    std::collections::HashSet,
 };
 
 #[derive(Debug, Serialize, Deserialize, PartialEq, Eq, Clone)]
@@ -14,6 +15,12 @@ pub enum Versions {
     Current(Box<State>),
 }
 
+#[derive(Debug, Eq, PartialEq)]
+pub enum AuthorizeNonceError {
+    MissingRequiredSignature(/*account authority:*/ Pubkey),
+    Uninitialized,
+}
+
 impl Versions {
     pub fn new(state: State, separate_domains: bool) -> Self {
         if separate_domains {
@@ -30,14 +37,6 @@ impl Versions {
         }
     }
 
-    /// Returns true if the durable nonce is not in the blockhash domain.
-    pub fn separate_domains(&self) -> bool {
-        match self {
-            Self::Legacy(_) => false,
-            Self::Current(_) => true,
-        }
-    }
-
     /// Checks if the recent_blockhash field in Transaction verifies, and
     /// returns nonce account data if so.
     pub fn verify_recent_blockhash(
@@ -63,7 +62,7 @@ impl Versions {
         }
     }
 
-    // Upgrades legacy nonces out of chain blockhash domains.
+    /// Upgrades legacy nonces out of chain blockhash domains.
     pub fn upgrade(self) -> Option<Self> {
         match self {
             Self::Legacy(mut state) => {
@@ -85,6 +84,35 @@ impl Versions {
             Self::Current(_) => None,
         }
     }
+
+    /// Updates the authority pubkey on the nonce account.
+    pub fn authorize(
+        self,
+        signers: &HashSet<Pubkey>,
+        authority: Pubkey,
+    ) -> Result<Self, AuthorizeNonceError> {
+        let data = match self.state() {
+            State::Uninitialized => return Err(AuthorizeNonceError::Uninitialized),
+            State::Initialized(data) => data,
+        };
+        if !signers.contains(&data.authority) {
+            return Err(AuthorizeNonceError::MissingRequiredSignature(
+                data.authority,
+            ));
+        }
+        let data = Data::new(
+            authority,
+            data.durable_nonce,
+            data.get_lamports_per_signature(),
+        );
+        let state = Box::new(State::Initialized(data));
+        // Preserve Version variant since cannot
+        // change durable_nonce field here.
+        Ok(match self {
+            Self::Legacy(_) => Self::Legacy,
+            Self::Current(_) => Self::Current,
+        }(state))
+    }
 }
 
 impl From<Versions> for State {
@@ -101,6 +129,7 @@ mod tests {
     use {
         super::*,
         crate::{fee_calculator::FeeCalculator, pubkey::Pubkey},
+        std::iter::repeat_with,
     };
 
     #[test]
@@ -231,4 +260,64 @@ mod tests {
         );
         assert_eq!(versions.upgrade(), None);
     }
+
+    #[test]
+    fn test_nonce_versions_authorize() {
+        // Uninitialized
+        let mut signers = repeat_with(Pubkey::new_unique).take(16).collect();
+        let versions = Versions::Legacy(Box::new(State::Uninitialized));
+        assert_eq!(
+            versions.authorize(&signers, Pubkey::new_unique()),
+            Err(AuthorizeNonceError::Uninitialized)
+        );
+        let versions = Versions::Current(Box::new(State::Uninitialized));
+        assert_eq!(
+            versions.authorize(&signers, Pubkey::new_unique()),
+            Err(AuthorizeNonceError::Uninitialized)
+        );
+        // Initialized, Legacy
+        let blockhash = Hash::from([171; 32]);
+        let durable_nonce =
+            DurableNonce::from_blockhash(&blockhash, /*separate_domains:*/ false);
+        let data = Data {
+            authority: Pubkey::new_unique(),
+            durable_nonce,
+            fee_calculator: FeeCalculator {
+                lamports_per_signature: 2718,
+            },
+        };
+        let account_authority = data.authority;
+        let versions = Versions::Legacy(Box::new(State::Initialized(data.clone())));
+        let authority = Pubkey::new_unique();
+        assert_ne!(authority, account_authority);
+        let data = Data { authority, ..data };
+        assert_eq!(
+            versions.clone().authorize(&signers, authority),
+            Err(AuthorizeNonceError::MissingRequiredSignature(
+                account_authority
+            )),
+        );
+        assert!(signers.insert(account_authority));
+        assert_eq!(
+            versions.authorize(&signers, authority),
+            Ok(Versions::Legacy(Box::new(State::Initialized(data.clone()))))
+        );
+        // Initialized, Current
+        let account_authority = data.authority;
+        let versions = Versions::Current(Box::new(State::Initialized(data.clone())));
+        let authority = Pubkey::new_unique();
+        assert_ne!(authority, account_authority);
+        let data = Data { authority, ..data };
+        assert_eq!(
+            versions.clone().authorize(&signers, authority),
+            Err(AuthorizeNonceError::MissingRequiredSignature(
+                account_authority
+            )),
+        );
+        assert!(signers.insert(account_authority));
+        assert_eq!(
+            versions.authorize(&signers, authority),
+            Ok(Versions::Current(Box::new(State::Initialized(data))))
+        );
+    }
 }