use {
crate::{
bpf_writer::BpfWriter,
error::{Error, ErrorCode},
solana_program::{
account_info::AccountInfo, instruction::AccountMeta, pubkey::Pubkey, system_program,
},
AccountDeserialize, AccountSerialize, Accounts, AccountsExit, Key, Owner, Result,
ToAccountInfos, ToAccountMetas,
},
std::{
collections::BTreeSet,
ops::{Deref, DerefMut},
},
};
#[derive(Debug)]
pub enum MigrationInner<From, To> {
From(From),
To(To),
}
#[derive(Debug)]
pub struct Migration<'info, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
info: &'info AccountInfo<'info>,
inner: MigrationInner<From, To>,
}
impl<'info, From, To> Migration<'info, From, To>
where
From: AccountDeserialize + Owner,
To: AccountSerialize + Owner,
{
fn new(info: &'info AccountInfo<'info>, account: From) -> Self {
Self {
info,
inner: MigrationInner::From(account),
}
}
#[inline(always)]
pub fn is_migrated(&self) -> bool {
matches!(self.inner, MigrationInner::To(_))
}
pub fn try_as_from(&self) -> Result<&From> {
match &self.inner {
MigrationInner::From(from) => Ok(from),
MigrationInner::To(_) => Err(ErrorCode::AccountAlreadyMigrated.into()),
}
}
pub fn try_as_from_mut(&mut self) -> Result<&mut From> {
match &mut self.inner {
MigrationInner::From(from) => Ok(from),
MigrationInner::To(_) => Err(ErrorCode::AccountAlreadyMigrated.into()),
}
}
pub fn migrate(&mut self, new_data: To) -> Result<()> {
if self.is_migrated() {
return Err(ErrorCode::AccountAlreadyMigrated.into());
}
self.inner = MigrationInner::To(new_data);
Ok(())
}
pub fn into_inner(&mut self, new_data: To) -> &To {
if !self.is_migrated() {
self.inner = MigrationInner::To(new_data);
}
match &self.inner {
MigrationInner::To(to) => to,
_ => unreachable!(),
}
}
pub fn into_inner_mut(&mut self, new_data: To) -> &mut To {
if !self.is_migrated() {
self.inner = MigrationInner::To(new_data);
}
match &mut self.inner {
MigrationInner::To(to) => to,
_ => unreachable!(),
}
}
#[inline(never)]
pub fn try_from(info: &'info AccountInfo<'info>) -> Result<Self> {
if info.owner == &system_program::ID && info.lamports() == 0 {
return Err(ErrorCode::AccountNotInitialized.into());
}
if info.owner != &From::owner() {
return Err(Error::from(ErrorCode::AccountOwnedByWrongProgram)
.with_pubkeys((*info.owner, From::owner())));
}
let mut data: &[u8] = &info.try_borrow_data()?;
Ok(Self::new(info, From::try_deserialize(&mut data)?))
}
#[inline(never)]
pub fn try_from_unchecked(info: &'info AccountInfo<'info>) -> Result<Self> {
if info.owner == &system_program::ID && info.lamports() == 0 {
return Err(ErrorCode::AccountNotInitialized.into());
}
if info.owner != &From::owner() {
return Err(Error::from(ErrorCode::AccountOwnedByWrongProgram)
.with_pubkeys((*info.owner, From::owner())));
}
let mut data: &[u8] = &info.try_borrow_data()?;
Ok(Self::new(info, From::try_deserialize_unchecked(&mut data)?))
}
}
impl<'info, B, From, To> Accounts<'info, B> for Migration<'info, From, To>
where
From: AccountDeserialize + Owner,
To: AccountSerialize + Owner,
{
#[inline(never)]
fn try_accounts(
_program_id: &Pubkey,
accounts: &mut &'info [AccountInfo<'info>],
_ix_data: &[u8],
_bumps: &mut B,
_reallocs: &mut BTreeSet<Pubkey>,
) -> Result<Self> {
if accounts.is_empty() {
return Err(ErrorCode::AccountNotEnoughKeys.into());
}
let account = &accounts[0];
*accounts = &accounts[1..];
Self::try_from(account)
}
}
impl<'info, From, To> AccountsExit<'info> for Migration<'info, From, To>
where
From: AccountDeserialize + Owner,
To: AccountSerialize + Owner,
{
fn exit(&self, program_id: &Pubkey) -> Result<()> {
if crate::common::is_closed(self.info) {
return Ok(());
}
match &self.inner {
MigrationInner::From(_) => {
return Err(ErrorCode::AccountNotMigrated.into());
}
MigrationInner::To(to) => {
let expected_owner = To::owner();
if &expected_owner != program_id {
return Ok(());
}
let mut data = self.info.try_borrow_mut_data()?;
let dst: &mut [u8] = &mut data;
let mut writer = BpfWriter::new(dst);
to.try_serialize(&mut writer)?;
}
}
Ok(())
}
}
impl<From, To> ToAccountMetas for Migration<'_, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
fn to_account_metas(&self, is_signer: Option<bool>) -> Vec<AccountMeta> {
let is_signer = is_signer.unwrap_or(self.info.is_signer);
let meta = match self.info.is_writable {
false => AccountMeta::new_readonly(*self.info.key, is_signer),
true => AccountMeta::new(*self.info.key, is_signer),
};
vec![meta]
}
}
impl<'info, From, To> ToAccountInfos<'info> for Migration<'info, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
fn to_account_infos(&self) -> Vec<AccountInfo<'info>> {
vec![self.info.clone()]
}
}
impl<'info, From, To> AsRef<AccountInfo<'info>> for Migration<'info, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
fn as_ref(&self) -> &AccountInfo<'info> {
self.info
}
}
impl<From, To> Key for Migration<'_, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
fn key(&self) -> Pubkey {
*self.info.key
}
}
impl<'info, From, To> Deref for Migration<'info, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
type Target = From;
fn deref(&self) -> &Self::Target {
match &self.inner {
MigrationInner::From(from) => from,
MigrationInner::To(_) => {
crate::solana_program::msg!("Cannot deref to From: account is already migrated.");
panic!();
}
}
}
}
impl<'info, From, To> DerefMut for Migration<'info, From, To>
where
From: AccountDeserialize,
To: AccountSerialize,
{
fn deref_mut(&mut self) -> &mut Self::Target {
match &mut self.inner {
MigrationInner::From(from) => from,
MigrationInner::To(_) => {
crate::solana_program::msg!(
"Cannot deref_mut to From: account is already migrated."
);
panic!();
}
}
}
}
#[cfg(test)]
mod tests {
use {
super::*,
crate::{AnchorDeserialize, AnchorSerialize, Discriminator},
};
const TEST_DISCRIMINATOR_V1: [u8; 8] = [1, 2, 3, 4, 5, 6, 7, 8];
const TEST_DISCRIMINATOR_V2: [u8; 8] = [8, 7, 6, 5, 4, 3, 2, 1];
const TEST_OWNER: Pubkey = Pubkey::new_from_array([1u8; 32]);
#[derive(Debug, Clone, AnchorSerialize, AnchorDeserialize, PartialEq)]
struct AccountV1 {
pub data: u64,
}
impl Discriminator for AccountV1 {
const DISCRIMINATOR: &'static [u8] = &TEST_DISCRIMINATOR_V1;
}
impl Owner for AccountV1 {
fn owner() -> Pubkey {
TEST_OWNER
}
}
impl AccountSerialize for AccountV1 {
fn try_serialize<W: std::io::Write>(&self, writer: &mut W) -> Result<()> {
writer.write_all(&TEST_DISCRIMINATOR_V1)?;
AnchorSerialize::serialize(self, writer)?;
Ok(())
}
}
impl AccountDeserialize for AccountV1 {
fn try_deserialize(buf: &mut &[u8]) -> Result<Self> {
if buf.len() < 8 {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}
let disc = &buf[..8];
if disc != TEST_DISCRIMINATOR_V1 {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}
Self::try_deserialize_unchecked(buf)
}
fn try_deserialize_unchecked(buf: &mut &[u8]) -> Result<Self> {
let mut data = &buf[8..];
AnchorDeserialize::deserialize(&mut data)
.map_err(|_| ErrorCode::AccountDidNotDeserialize.into())
}
}
#[derive(Debug, Clone, AnchorSerialize, AnchorDeserialize, PartialEq)]
struct AccountV2 {
pub data: u64,
pub new_field: u64,
}
impl Discriminator for AccountV2 {
const DISCRIMINATOR: &'static [u8] = &TEST_DISCRIMINATOR_V2;
}
impl Owner for AccountV2 {
fn owner() -> Pubkey {
TEST_OWNER
}
}
impl AccountSerialize for AccountV2 {
fn try_serialize<W: std::io::Write>(&self, writer: &mut W) -> Result<()> {
writer.write_all(&TEST_DISCRIMINATOR_V2)?;
AnchorSerialize::serialize(self, writer)?;
Ok(())
}
}
impl AccountDeserialize for AccountV2 {
fn try_deserialize(buf: &mut &[u8]) -> Result<Self> {
if buf.len() < 8 {
return Err(ErrorCode::AccountDiscriminatorNotFound.into());
}
let disc = &buf[..8];
if disc != TEST_DISCRIMINATOR_V2 {
return Err(ErrorCode::AccountDiscriminatorMismatch.into());
}
Self::try_deserialize_unchecked(buf)
}
fn try_deserialize_unchecked(buf: &mut &[u8]) -> Result<Self> {
let mut data = &buf[8..];
AnchorDeserialize::deserialize(&mut data)
.map_err(|_| ErrorCode::AccountDidNotDeserialize.into())
}
}
fn create_account_info<'a>(
key: &'a Pubkey,
owner: &'a Pubkey,
lamports: &'a mut u64,
data: &'a mut [u8],
) -> AccountInfo<'a> {
AccountInfo::new(key, false, true, lamports, data, owner, false)
}
#[test]
fn test_is_migrated_returns_false_initially() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
assert!(!migration.is_migrated());
}
#[test]
fn test_is_migrated_returns_true_after_migrate() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
migration
.migrate(AccountV2 {
data: 42,
new_field: 100,
})
.unwrap();
assert!(migration.is_migrated());
}
#[test]
fn test_try_as_from_returns_data_before_migration() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
let from = migration.try_as_from().unwrap();
assert_eq!(from.data, 42);
}
#[test]
fn test_try_as_from_returns_error_after_migration() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
migration
.migrate(AccountV2 {
data: 42,
new_field: 100,
})
.unwrap();
assert!(migration.try_as_from().is_err());
}
#[test]
fn test_try_as_from_mut_works_before_migration() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
let from = migration.try_as_from_mut().unwrap();
from.data = 100;
assert_eq!(migration.try_as_from().unwrap().data, 100);
}
#[test]
fn test_migrate_fails_if_already_migrated() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
migration
.migrate(AccountV2 {
data: 42,
new_field: 100,
})
.unwrap();
let result = migration.migrate(AccountV2 {
data: 42,
new_field: 200,
});
assert!(result.is_err());
}
#[test]
fn test_into_inner_migrates_and_returns_reference() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
let to = migration.into_inner(AccountV2 {
data: 42,
new_field: 100,
});
assert_eq!(to.data, 42);
assert_eq!(to.new_field, 100);
assert!(migration.is_migrated());
}
#[test]
fn test_into_inner_is_idempotent() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
let to1 = migration.into_inner(AccountV2 {
data: 42,
new_field: 100,
});
assert_eq!(to1.new_field, 100);
let to2 = migration.into_inner(AccountV2 {
data: 42,
new_field: 999,
});
assert_eq!(to2.new_field, 100); }
#[test]
fn test_into_inner_mut_allows_mutation() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
let to = migration.into_inner_mut(AccountV2 {
data: 42,
new_field: 100,
});
to.new_field = 200;
let to_ref = migration.into_inner(AccountV2 {
data: 0,
new_field: 0,
});
assert_eq!(to_ref.new_field, 200);
}
#[test]
fn test_deref_works_before_migration() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
assert_eq!(migration.data, 42);
}
#[test]
#[should_panic]
fn test_deref_panics_after_migration() {
let key = Pubkey::default();
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &TEST_OWNER, &mut lamports, &mut data);
let mut migration: Migration<AccountV1, AccountV2> = Migration::try_from(&info).unwrap();
migration
.migrate(AccountV2 {
data: 42,
new_field: 100,
})
.unwrap();
let _ = migration.data;
}
#[test]
fn test_try_from_fails_with_wrong_owner() {
let key = Pubkey::default();
let wrong_owner = Pubkey::new_from_array([99u8; 32]);
let mut lamports = 100;
let v1 = AccountV1 { data: 42 };
let mut data = vec![0u8; 100];
data[..8].copy_from_slice(&TEST_DISCRIMINATOR_V1);
v1.serialize(&mut &mut data[8..]).unwrap();
let info = create_account_info(&key, &wrong_owner, &mut lamports, &mut data);
let result: Result<Migration<AccountV1, AccountV2>> = Migration::try_from(&info);
assert!(result.is_err());
}
#[test]
fn test_try_from_fails_with_uninitialized_account() {
let key = Pubkey::default();
let mut lamports = 0;
let mut data = vec![0u8; 100];
let info = create_account_info(&key, &system_program::ID, &mut lamports, &mut data);
let result: Result<Migration<AccountV1, AccountV2>> = Migration::try_from(&info);
assert!(result.is_err());
}
}