1use {
4 crate::{
5 bpf_writer::BpfWriter,
6 error::{Error, ErrorCode},
7 solana_program::{
8 account_info::AccountInfo, instruction::AccountMeta, pubkey::Pubkey, system_program,
9 },
10 AccountDeserialize, AccountSerialize, Accounts, AccountsExit, Key, Owner, Result,
11 ToAccountInfos, ToAccountMetas,
12 },
13 std::{
14 collections::BTreeSet,
15 ops::{Deref, DerefMut},
16 },
17};
18
19#[derive(Debug)]
21pub enum MigrationInner<From, To> {
22 From(From),
24 To(To),
26}
27
28#[derive(Debug)]
154pub struct Migration<'info, From, To>
155where
156 From: AccountDeserialize,
157 To: AccountSerialize,
158{
159 info: &'info AccountInfo<'info>,
161 inner: MigrationInner<From, To>,
163}
164
165impl<'info, From, To> Migration<'info, From, To>
166where
167 From: AccountDeserialize + Owner,
168 To: AccountSerialize + Owner,
169{
170 fn new(info: &'info AccountInfo<'info>, account: From) -> Self {
172 Self {
173 info,
174 inner: MigrationInner::From(account),
175 }
176 }
177
178 #[inline(always)]
180 pub fn is_migrated(&self) -> bool {
181 matches!(self.inner, MigrationInner::To(_))
182 }
183
184 pub fn try_as_from(&self) -> Result<&From> {
189 match &self.inner {
190 MigrationInner::From(from) => Ok(from),
191 MigrationInner::To(_) => Err(ErrorCode::AccountAlreadyMigrated.into()),
192 }
193 }
194
195 pub fn try_as_from_mut(&mut self) -> Result<&mut From> {
200 match &mut self.inner {
201 MigrationInner::From(from) => Ok(from),
202 MigrationInner::To(_) => Err(ErrorCode::AccountAlreadyMigrated.into()),
203 }
204 }
205
206 pub fn migrate(&mut self, new_data: To) -> Result<()> {
214 if self.is_migrated() {
215 return Err(ErrorCode::AccountAlreadyMigrated.into());
216 }
217
218 self.inner = MigrationInner::To(new_data);
219 Ok(())
220 }
221
222 pub fn into_inner(&mut self, new_data: To) -> &To {
248 if !self.is_migrated() {
249 self.inner = MigrationInner::To(new_data);
250 }
251
252 match &self.inner {
253 MigrationInner::To(to) => to,
254 _ => unreachable!(),
255 }
256 }
257
258 pub fn into_inner_mut(&mut self, new_data: To) -> &mut To {
284 if !self.is_migrated() {
285 self.inner = MigrationInner::To(new_data);
286 }
287
288 match &mut self.inner {
289 MigrationInner::To(to) => to,
290 _ => unreachable!(),
291 }
292 }
293
294 #[inline(never)]
299 pub fn try_from(info: &'info AccountInfo<'info>) -> Result<Self> {
300 if info.owner == &system_program::ID && info.lamports() == 0 {
301 return Err(ErrorCode::AccountNotInitialized.into());
302 }
303
304 if info.owner != &From::owner() {
305 return Err(Error::from(ErrorCode::AccountOwnedByWrongProgram)
306 .with_pubkeys((*info.owner, From::owner())));
307 }
308
309 let mut data: &[u8] = &info.try_borrow_data()?;
310 Ok(Self::new(info, From::try_deserialize(&mut data)?))
311 }
312
313 #[inline(never)]
318 pub fn try_from_unchecked(info: &'info AccountInfo<'info>) -> Result<Self> {
319 if info.owner == &system_program::ID && info.lamports() == 0 {
320 return Err(ErrorCode::AccountNotInitialized.into());
321 }
322
323 if info.owner != &From::owner() {
324 return Err(Error::from(ErrorCode::AccountOwnedByWrongProgram)
325 .with_pubkeys((*info.owner, From::owner())));
326 }
327
328 let mut data: &[u8] = &info.try_borrow_data()?;
329 Ok(Self::new(info, From::try_deserialize_unchecked(&mut data)?))
330 }
331}
332
333impl<'info, B, From, To> Accounts<'info, B> for Migration<'info, From, To>
334where
335 From: AccountDeserialize + Owner,
336 To: AccountSerialize + Owner,
337{
338 #[inline(never)]
339 fn try_accounts(
340 _program_id: &Pubkey,
341 accounts: &mut &'info [AccountInfo<'info>],
342 _ix_data: &[u8],
343 _bumps: &mut B,
344 _reallocs: &mut BTreeSet<Pubkey>,
345 ) -> Result<Self> {
346 if accounts.is_empty() {
347 return Err(ErrorCode::AccountNotEnoughKeys.into());
348 }
349 let account = &accounts[0];
350 *accounts = &accounts[1..];
351 Self::try_from(account)
352 }
353}
354
355impl<'info, From, To> AccountsExit<'info> for Migration<'info, From, To>
356where
357 From: AccountDeserialize + Owner,
358 To: AccountSerialize + Owner,
359{
360 fn exit(&self, program_id: &Pubkey) -> Result<()> {
361 if crate::common::is_closed(self.info) {
363 return Ok(());
364 }
365
366 match &self.inner {
368 MigrationInner::From(_) => {
369 return Err(ErrorCode::AccountNotMigrated.into());
371 }
372 MigrationInner::To(to) => {
373 let expected_owner = To::owner();
375 if &expected_owner != program_id {
376 return Ok(());
377 }
378
379 let mut data = self.info.try_borrow_mut_data()?;
381 let dst: &mut [u8] = &mut data;
382 let mut writer = BpfWriter::new(dst);
383 to.try_serialize(&mut writer)?;
384 }
385 }
386
387 Ok(())
388 }
389}
390
391impl<From, To> ToAccountMetas for Migration<'_, From, To>
392where
393 From: AccountDeserialize,
394 To: AccountSerialize,
395{
396 fn to_account_metas(&self, is_signer: Option<bool>) -> Vec<AccountMeta> {
397 let is_signer = is_signer.unwrap_or(self.info.is_signer);
398 let meta = match self.info.is_writable {
399 false => AccountMeta::new_readonly(*self.info.key, is_signer),
400 true => AccountMeta::new(*self.info.key, is_signer),
401 };
402 vec![meta]
403 }
404}
405
406impl<'info, From, To> ToAccountInfos<'info> for Migration<'info, From, To>
407where
408 From: AccountDeserialize,
409 To: AccountSerialize,
410{
411 fn to_account_infos(&self) -> Vec<AccountInfo<'info>> {
412 vec![self.info.clone()]
413 }
414}
415
416impl<'info, From, To> AsRef<AccountInfo<'info>> for Migration<'info, From, To>
417where
418 From: AccountDeserialize,
419 To: AccountSerialize,
420{
421 fn as_ref(&self) -> &AccountInfo<'info> {
422 self.info
423 }
424}
425
426impl<From, To> Key for Migration<'_, From, To>
427where
428 From: AccountDeserialize,
429 To: AccountSerialize,
430{
431 fn key(&self) -> Pubkey {
432 *self.info.key
433 }
434}
435
436impl<'info, From, To> Deref for Migration<'info, From, To>
438where
439 From: AccountDeserialize,
440 To: AccountSerialize,
441{
442 type Target = From;
443
444 fn deref(&self) -> &Self::Target {
445 match &self.inner {
446 MigrationInner::From(from) => from,
447 MigrationInner::To(_) => {
448 crate::solana_program::msg!("Cannot deref to From: account is already migrated.");
449 panic!();
450 }
451 }
452 }
453}
454
455impl<'info, From, To> DerefMut for Migration<'info, From, To>
457where
458 From: AccountDeserialize,
459 To: AccountSerialize,
460{
461 fn deref_mut(&mut self) -> &mut Self::Target {
462 match &mut self.inner {
463 MigrationInner::From(from) => from,
464 MigrationInner::To(_) => {
465 crate::solana_program::msg!(
466 "Cannot deref_mut to From: account is already migrated."
467 );
468 panic!();
469 }
470 }
471 }
472}
473
474#[cfg(test)]
475mod tests {
476 use {
477 super::*,
478 crate::{AnchorDeserialize, AnchorSerialize, Discriminator},
479 };
480
481 const TEST_DISCRIMINATOR_V1: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8];
482 const TEST_DISCRIMINATOR_V2: [u8; 8] = [8, 7, 6, 5, 4, 3, 2, 1];
483 const TEST_OWNER: Pubkey = Pubkey::new_from_array([1u8; 32]);
484
485 #[derive(Debug, Clone, AnchorSerialize, AnchorDeserialize, PartialEq)]
486 struct AccountV1 {
487 pub data: u64,
488 }
489
490 impl Discriminator for AccountV1 {
491 const DISCRIMINATOR: &'static [u8] = &TEST_DISCRIMINATOR_V1;
492 }
493
494 impl Owner for AccountV1 {
495 fn owner() -> Pubkey {
496 TEST_OWNER
497 }
498 }
499
500 impl AccountSerialize for AccountV1 {
501 fn try_serialize<W: std::io::Write>(&self, writer: &mut W) -> Result<()> {
502 writer.write_all(&TEST_DISCRIMINATOR_V1)?;
503 AnchorSerialize::serialize(self, writer)?;
504 Ok(())
505 }
506 }
507
508 impl AccountDeserialize for AccountV1 {
509 fn try_deserialize(buf: &mut &[u8]) -> Result<Self> {
510 if buf.len() < 8 {
511 return Err(ErrorCode::AccountDiscriminatorNotFound.into());
512 }
513 let disc = &buf[..8];
514 if disc != TEST_DISCRIMINATOR_V1 {
515 return Err(ErrorCode::AccountDiscriminatorMismatch.into());
516 }
517 Self::try_deserialize_unchecked(buf)
518 }
519
520 fn try_deserialize_unchecked(buf: &mut &[u8]) -> Result<Self> {
521 let mut data = &buf[8..];
522 AnchorDeserialize::deserialize(&mut data)
523 .map_err(|_| ErrorCode::AccountDidNotDeserialize.into())
524 }
525 }
526
527 #[derive(Debug, Clone, AnchorSerialize, AnchorDeserialize, PartialEq)]
528 struct AccountV2 {
529 pub data: u64,
530 pub new_field: u64,
531 }
532
533 impl Discriminator for AccountV2 {
534 const DISCRIMINATOR: &'static [u8] = &TEST_DISCRIMINATOR_V2;
535 }
536
537 impl Owner for AccountV2 {
538 fn owner() -> Pubkey {
539 TEST_OWNER
540 }
541 }
542
543 impl AccountSerialize for AccountV2 {
544 fn try_serialize<W: std::io::Write>(&self, writer: &mut W) -> Result<()> {
545 writer.write_all(&TEST_DISCRIMINATOR_V2)?;
546 AnchorSerialize::serialize(self, writer)?;
547 Ok(())
548 }
549 }
550
551 impl AccountDeserialize for AccountV2 {
552 fn try_deserialize(buf: &mut &[u8]) -> Result<Self> {
553 if buf.len() < 8 {
554 return Err(ErrorCode::AccountDiscriminatorNotFound.into());
555 }
556 let disc = &buf[..8];
557 if disc != TEST_DISCRIMINATOR_V2 {
558 return Err(ErrorCode::AccountDiscriminatorMismatch.into());
559 }
560 Self::try_deserialize_unchecked(buf)
561 }
562
563 fn try_deserialize_unchecked(buf: &mut &[u8]) -> Result<Self> {
564 let mut data = &buf[8..];
565 AnchorDeserialize::deserialize(&mut data)
566 .map_err(|_| ErrorCode::AccountDidNotDeserialize.into())
567 }
568 }
569
570 fn create_account_info<'a>(
571 key: &'a Pubkey,
572 owner: &'a Pubkey,
573 lamports: &'a mut u64,
574 data: &'a mut [u8],
575 ) -> AccountInfo<'a> {
576 AccountInfo::new(key, false, true, lamports, data, owner, false)
577 }
578
579 #[test]
582 fn test_is_migrated_returns_false_initially() {
583 let key = Pubkey::default();
584 let mut lamports = 100;
585 let v1 = AccountV1 { data: 42 };
586 let mut data = vec![0u8; 100];
587 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
588 v1.serialize(&mut &mut data[8..]).unwrap();
589
590 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
591 let migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
592
593 assert!(!migration.is_migrated());
594 }
595
596 #[test]
599 fn test_is_migrated_returns_true_after_migrate() {
600 let key = Pubkey::default();
601 let mut lamports = 100;
602 let v1 = AccountV1 { data: 42 };
603 let mut data = vec![0u8; 100];
604 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
605 v1.serialize(&mut &mut data[8..]).unwrap();
606
607 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
608 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
609
610 migration
611 .migrate(AccountV2 {
612 data: 42,
613 new_field: 100,
614 })
615 .unwrap();
616
617 assert!(migration.is_migrated());
618 }
619
620 #[test]
623 fn test_try_as_from_returns_data_before_migration() {
624 let key = Pubkey::default();
625 let mut lamports = 100;
626 let v1 = AccountV1 { data: 42 };
627 let mut data = vec![0u8; 100];
628 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
629 v1.serialize(&mut &mut data[8..]).unwrap();
630
631 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
632 let migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
633
634 let from = migration.try_as_from().unwrap();
635 assert_eq!(from.data, 42);
636 }
637
638 #[test]
641 fn test_try_as_from_returns_error_after_migration() {
642 let key = Pubkey::default();
643 let mut lamports = 100;
644 let v1 = AccountV1 { data: 42 };
645 let mut data = vec![0u8; 100];
646 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
647 v1.serialize(&mut &mut data[8..]).unwrap();
648
649 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
650 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
651
652 migration
653 .migrate(AccountV2 {
654 data: 42,
655 new_field: 100,
656 })
657 .unwrap();
658
659 assert!(migration.try_as_from().is_err());
660 }
661
662 #[test]
665 fn test_try_as_from_mut_works_before_migration() {
666 let key = Pubkey::default();
667 let mut lamports = 100;
668 let v1 = AccountV1 { data: 42 };
669 let mut data = vec![0u8; 100];
670 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
671 v1.serialize(&mut &mut data[8..]).unwrap();
672
673 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
674 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
675
676 let from = migration.try_as_from_mut().unwrap();
677 from.data = 100;
678 assert_eq!(migration.try_as_from().unwrap().data, 100);
679 }
680
681 #[test]
684 fn test_migrate_fails_if_already_migrated() {
685 let key = Pubkey::default();
686 let mut lamports = 100;
687 let v1 = AccountV1 { data: 42 };
688 let mut data = vec![0u8; 100];
689 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
690 v1.serialize(&mut &mut data[8..]).unwrap();
691
692 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
693 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
694
695 migration
696 .migrate(AccountV2 {
697 data: 42,
698 new_field: 100,
699 })
700 .unwrap();
701 let result = migration.migrate(AccountV2 {
702 data: 42,
703 new_field: 200,
704 });
705
706 assert!(result.is_err());
707 }
708
709 #[test]
712 fn test_into_inner_migrates_and_returns_reference() {
713 let key = Pubkey::default();
714 let mut lamports = 100;
715 let v1 = AccountV1 { data: 42 };
716 let mut data = vec![0u8; 100];
717 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
718 v1.serialize(&mut &mut data[8..]).unwrap();
719
720 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
721 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
722
723 let to = migration.into_inner(AccountV2 {
724 data: 42,
725 new_field: 100,
726 });
727
728 assert_eq!(to.data, 42);
729 assert_eq!(to.new_field, 100);
730 assert!(migration.is_migrated());
731 }
732
733 #[test]
736 fn test_into_inner_is_idempotent() {
737 let key = Pubkey::default();
738 let mut lamports = 100;
739 let v1 = AccountV1 { data: 42 };
740 let mut data = vec![0u8; 100];
741 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
742 v1.serialize(&mut &mut data[8..]).unwrap();
743
744 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
745 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
746
747 let to1 = migration.into_inner(AccountV2 {
748 data: 42,
749 new_field: 100,
750 });
751 assert_eq!(to1.new_field, 100);
752
753 let to2 = migration.into_inner(AccountV2 {
755 data: 42,
756 new_field: 999,
757 });
758 assert_eq!(to2.new_field, 100); }
760
761 #[test]
764 fn test_into_inner_mut_allows_mutation() {
765 let key = Pubkey::default();
766 let mut lamports = 100;
767 let v1 = AccountV1 { data: 42 };
768 let mut data = vec![0u8; 100];
769 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
770 v1.serialize(&mut &mut data[8..]).unwrap();
771
772 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
773 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
774
775 let to = migration.into_inner_mut(AccountV2 {
776 data: 42,
777 new_field: 100,
778 });
779 to.new_field = 200;
780
781 let to_ref = migration.into_inner(AccountV2 {
782 data: 0,
783 new_field: 0,
784 });
785 assert_eq!(to_ref.new_field, 200);
786 }
787
788 #[test]
791 fn test_deref_works_before_migration() {
792 let key = Pubkey::default();
793 let mut lamports = 100;
794 let v1 = AccountV1 { data: 42 };
795 let mut data = vec![0u8; 100];
796 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
797 v1.serialize(&mut &mut data[8..]).unwrap();
798
799 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
800 let migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
801
802 assert_eq!(migration.data, 42);
803 }
804
805 #[test]
808 #[should_panic]
809 fn test_deref_panics_after_migration() {
810 let key = Pubkey::default();
811 let mut lamports = 100;
812 let v1 = AccountV1 { data: 42 };
813 let mut data = vec![0u8; 100];
814 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
815 v1.serialize(&mut &mut data[8..]).unwrap();
816
817 let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
818 let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
819
820 migration
821 .migrate(AccountV2 {
822 data: 42,
823 new_field: 100,
824 })
825 .unwrap();
826
827 let _ = migration.data;
829 }
830
831 #[test]
834 fn test_try_from_fails_with_wrong_owner() {
835 let key = Pubkey::default();
836 let wrong_owner = Pubkey::new_from_array([99u8; 32]);
837 let mut lamports = 100;
838 let v1 = AccountV1 { data: 42 };
839 let mut data = vec![0u8; 100];
840 data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
841 v1.serialize(&mut &mut data[8..]).unwrap();
842
843 let info = create_account_info(&key, &wrong_owner, &mut lamports, &mut data);
844 let result: Result<Migration<AccountV1, AccountV2>> = Migration::try_from(&info);
845
846 assert!(result.is_err());
847 }
848
849 #[test]
852 fn test_try_from_fails_with_uninitialized_account() {
853 let key = Pubkey::default();
854 let mut lamports = 0;
855 let mut data = vec![0u8; 100];
856
857 let info = create_account_info(&key, &system_program::ID, &mut lamports, &mut data);
858 let result: Result<Migration<AccountV1, AccountV2>> = Migration::try_from(&info);
859
860 assert!(result.is_err());
861 }
862}