solana_address_book/
registered_address.rs

1//! Registered address types and utilities for the address book.
2
3use crate::pda_seeds::find_pda_with_bump_and_strings;
4use anchor_lang::prelude::*;
5
6/// Role type for registered addresses, defining the purpose of each address
7#[derive(Debug, Clone, strum::Display, Hash, PartialEq, Eq)]
8pub enum AddressRole {
9    /// Standard user wallet
10    #[strum(serialize = "wallet")]
11    Wallet,
12
13    /// Token mint address
14    #[strum(serialize = "mint")]
15    Mint,
16
17    /// Associated Token Account with mint and owner references
18    #[strum(serialize = "ata")]
19    Ata { mint: Pubkey, owner: Pubkey },
20
21    /// Program Derived Address with seeds and program information
22    #[strum(serialize = "pda")]
23    Pda {
24        seeds: Vec<String>,
25        program_id: Pubkey,
26        bump: u8,
27    },
28
29    /// Smart contract program address
30    #[strum(serialize = "program")]
31    Program,
32
33    /// Custom user-defined role
34    #[strum(serialize = "custom")]
35    Custom(String),
36}
37
38/// Registered address with role information
39#[derive(Debug, Clone, Hash, PartialEq, Eq)]
40pub struct RegisteredAddress {
41    /// The address.
42    pub key: Pubkey,
43    /// The address's function within the program.
44    pub role: AddressRole,
45}
46
47impl RegisteredAddress {
48    /// Creates a new registered address with the specified role
49    ///
50    /// # Arguments
51    /// * `address` - The public key to register
52    /// * `role` - The role that defines the address's purpose
53    ///
54    /// # Example
55    /// ```
56    /// use anchor_lang::prelude::*;
57    /// use solana_address_book::{RegisteredAddress, AddressRole};
58    ///
59    /// let address = Pubkey::new_unique();
60    /// let registered = RegisteredAddress::new(address, AddressRole::Wallet);
61    /// ```
62    pub fn new(address: Pubkey, role: AddressRole) -> Self {
63        Self { key: address, role }
64    }
65
66    /// Creates a wallet-type registered address
67    ///
68    /// # Arguments
69    /// * `address` - The wallet's public key
70    ///
71    /// # Example
72    /// ```
73    /// use anchor_lang::prelude::*;
74    /// use solana_address_book::RegisteredAddress;
75    ///
76    /// let wallet = Pubkey::new_unique();
77    /// let registered = RegisteredAddress::wallet(wallet);
78    /// ```
79    pub fn wallet(address: Pubkey) -> Self {
80        Self::new(address, AddressRole::Wallet)
81    }
82
83    /// Creates a mint-type registered address
84    ///
85    /// # Arguments
86    /// * `address` - The mint's public key
87    ///
88    /// # Example
89    /// ```
90    /// use anchor_lang::prelude::*;
91    /// use solana_address_book::RegisteredAddress;
92    ///
93    /// let mint = Pubkey::new_unique();
94    /// let registered = RegisteredAddress::mint(mint);
95    /// ```
96    pub fn mint(address: Pubkey) -> Self {
97        Self::new(address, AddressRole::Mint)
98    }
99
100    /// Creates an Associated Token Account (ATA) registered address
101    ///
102    /// # Arguments
103    /// * `address` - The ATA's public key
104    /// * `mint` - The token mint public key
105    /// * `owner` - The owner's public key
106    ///
107    /// # Example
108    /// ```
109    /// use anchor_lang::prelude::*;
110    /// use solana_address_book::RegisteredAddress;
111    ///
112    /// let ata = Pubkey::new_unique();
113    /// let mint = Pubkey::new_unique();
114    /// let owner = Pubkey::new_unique();
115    /// let registered = RegisteredAddress::ata(ata, mint, owner);
116    /// ```
117    pub fn ata(address: Pubkey, mint: Pubkey, owner: Pubkey) -> Self {
118        Self::new(address, AddressRole::Ata { mint, owner })
119    }
120
121    /// Creates a custom-role registered address
122    ///
123    /// # Arguments
124    /// * `address` - The address's public key
125    /// * `custom_role` - A custom string describing the address's role
126    ///
127    /// # Example
128    /// ```
129    /// use anchor_lang::prelude::*;
130    /// use solana_address_book::RegisteredAddress;
131    ///
132    /// let address = Pubkey::new_unique();
133    /// let registered = RegisteredAddress::custom(address, "governance");
134    /// ```
135    pub fn custom(address: Pubkey, custom_role: &str) -> Self {
136        Self::new(address, AddressRole::Custom(custom_role.to_string()))
137    }
138
139    /// Creates a program-type registered address
140    ///
141    /// # Arguments
142    /// * `address` - The program's public key
143    ///
144    /// # Example
145    /// ```
146    /// use anchor_lang::prelude::*;
147    /// use solana_address_book::RegisteredAddress;
148    ///
149    /// let program = Pubkey::new_unique();
150    /// let registered = RegisteredAddress::program(program);
151    /// ```
152    pub fn program(address: Pubkey) -> Self {
153        Self::new(address, AddressRole::Program)
154    }
155
156    /// Creates a PDA registered address by finding the program address
157    ///
158    /// # Arguments
159    /// * `seeds` - The seeds used to derive the PDA
160    /// * `program_id` - The program that owns the PDA
161    ///
162    /// # Returns
163    /// A tuple containing:
164    /// * The derived PDA public key
165    /// * The bump seed
166    /// * The registered address
167    ///
168    /// # Example
169    /// ```
170    /// use anchor_lang::prelude::*;
171    /// use solana_address_book::RegisteredAddress;
172    ///
173    /// let program_id = Pubkey::new_unique();
174    /// let user = Pubkey::new_unique();
175    ///
176    /// let (pda_key, bump, registered) = RegisteredAddress::pda(&[b"vault", user.as_ref()], &program_id);
177    /// ```
178    pub fn pda(seeds: &[&[u8]], program_id: &Pubkey) -> (Pubkey, u8, Self) {
179        let derived_pda = find_pda_with_bump_and_strings(seeds, program_id);
180
181        (
182            derived_pda.key,
183            derived_pda.bump,
184            Self::new(
185                derived_pda.key,
186                AddressRole::Pda {
187                    seeds: derived_pda.seed_strings,
188                    program_id: *program_id,
189                    bump: derived_pda.bump,
190                },
191            ),
192        )
193    }
194
195    /// Creates a PDA registered address from existing PDA information
196    ///
197    /// # Arguments
198    /// * `pubkey` - The PDA's public key
199    /// * `seeds` - The string representations of seeds used
200    /// * `program_id` - The program that owns the PDA
201    /// * `bump` - The bump seed
202    ///
203    /// # Example
204    /// ```
205    /// use anchor_lang::prelude::*;
206    /// use solana_address_book::RegisteredAddress;
207    ///
208    /// let pda = Pubkey::new_unique();
209    /// let program_id = Pubkey::new_unique();
210    /// let seeds = vec!["vault".to_string(), "v1".to_string()];
211    ///
212    /// let registered = RegisteredAddress::pda_from_parts(pda, seeds, program_id, 255);
213    /// ```
214    pub fn pda_from_parts(
215        pubkey: Pubkey,
216        seeds: Vec<String>,
217        program_id: Pubkey,
218        bump: u8,
219    ) -> Self {
220        Self::new(
221            pubkey,
222            AddressRole::Pda {
223                seeds,
224                program_id,
225                bump,
226            },
227        )
228    }
229}
230
231impl std::fmt::Display for RegisteredAddress {
232    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
233        match &self.role {
234            AddressRole::Ata { mint, owner } => {
235                write!(f, "{} [ata mint:{} owner:{}]", self.key, mint, owner)
236            }
237            AddressRole::Pda { seeds, bump, .. } => {
238                write!(
239                    f,
240                    "{} [pda seeds:{} bump:{}]",
241                    self.key,
242                    seeds.join(","),
243                    bump
244                )
245            }
246            AddressRole::Custom(custom) => {
247                write!(f, "{} [{}]", self.key, custom)
248            }
249            _ => {
250                write!(f, "{} [{}]", self.key, self.role)
251            }
252        }
253    }
254}
255
256#[cfg(test)]
257mod tests {
258    use super::*;
259
260    #[test]
261    fn test_wallet_creation() {
262        let pubkey = Pubkey::new_unique();
263        let registered = RegisteredAddress::wallet(pubkey);
264
265        assert_eq!(registered.key, pubkey);
266        assert!(matches!(registered.role, AddressRole::Wallet));
267    }
268
269    #[test]
270    fn test_mint_creation() {
271        let pubkey = Pubkey::new_unique();
272        let registered = RegisteredAddress::mint(pubkey);
273
274        assert_eq!(registered.key, pubkey);
275        assert!(matches!(registered.role, AddressRole::Mint));
276    }
277
278    #[test]
279    fn test_ata_creation() {
280        let ata = Pubkey::new_unique();
281        let mint = Pubkey::new_unique();
282        let owner = Pubkey::new_unique();
283        let registered = RegisteredAddress::ata(ata, mint, owner);
284
285        assert_eq!(registered.key, ata);
286        if let AddressRole::Ata { mint: m, owner: o } = registered.role {
287            assert_eq!(m, mint);
288            assert_eq!(o, owner);
289        } else {
290            panic!("Expected ATA role");
291        }
292    }
293
294    #[test]
295    fn test_program_creation() {
296        let pubkey = Pubkey::new_unique();
297        let registered = RegisteredAddress::program(pubkey);
298
299        assert_eq!(registered.key, pubkey);
300        assert!(matches!(registered.role, AddressRole::Program));
301    }
302
303    #[test]
304    fn test_custom_creation() {
305        let pubkey = Pubkey::new_unique();
306        let registered = RegisteredAddress::custom(pubkey, "governance");
307
308        assert_eq!(registered.key, pubkey);
309        if let AddressRole::Custom(role) = registered.role {
310            assert_eq!(role, "governance");
311        } else {
312            panic!("Expected Custom role");
313        }
314    }
315
316    #[test]
317    fn test_pda_creation() {
318        let program_id = Pubkey::new_unique();
319        let (pubkey, bump, registered) = RegisteredAddress::pda(&[b"test", b"seed"], &program_id);
320
321        assert_eq!(registered.key, pubkey);
322        if let AddressRole::Pda {
323            seeds: pda_seeds,
324            program_id: pda_program_id,
325            bump: pda_bump,
326        } = &registered.role
327        {
328            assert_eq!(pda_seeds, &vec!["test".to_string(), "seed".to_string()]);
329            assert_eq!(*pda_program_id, program_id);
330            assert_eq!(*pda_bump, bump);
331        } else {
332            panic!("Expected PDA role");
333        }
334    }
335}