use std::{io, collections::HashSet};
use ciphersuite::{group::GroupEncoding, Ciphersuite};
use serai_client::primitives::{ExternalBalance, ExternalCoin, ExternalNetworkId};
use crate::{
Get, DbTxn, Db, Payment, Plan, create_db,
networks::{Output, Network},
multisigs::scheduler::{SchedulerAddendum, Scheduler as SchedulerTrait},
};
#[derive(Clone, PartialEq, Eq, Debug)]
pub struct Scheduler<N: Network> {
key: <N::Curve as Ciphersuite>::G,
coins: HashSet<ExternalCoin>,
rotated: bool,
}
#[derive(Clone, Copy, PartialEq, Eq, Debug)]
pub enum Addendum<N: Network> {
Nonce(u64),
RotateTo { nonce: u64, new_key: <N::Curve as Ciphersuite>::G },
}
impl<N: Network> SchedulerAddendum for Addendum<N> {
fn read<R: io::Read>(reader: &mut R) -> io::Result<Self> {
let mut kind = [0xff];
reader.read_exact(&mut kind)?;
match kind[0] {
0 => {
let mut nonce = [0; 8];
reader.read_exact(&mut nonce)?;
Ok(Addendum::Nonce(u64::from_le_bytes(nonce)))
}
1 => {
let mut nonce = [0; 8];
reader.read_exact(&mut nonce)?;
let nonce = u64::from_le_bytes(nonce);
let new_key = N::Curve::read_G(reader)?;
Ok(Addendum::RotateTo { nonce, new_key })
}
_ => Err(io::Error::other("reading unknown Addendum type"))?,
}
}
fn write<W: io::Write>(&self, writer: &mut W) -> io::Result<()> {
match self {
Addendum::Nonce(nonce) => {
writer.write_all(&[0])?;
writer.write_all(&nonce.to_le_bytes())
}
Addendum::RotateTo { nonce, new_key } => {
writer.write_all(&[1])?;
writer.write_all(&nonce.to_le_bytes())?;
writer.write_all(new_key.to_bytes().as_ref())
}
}
}
}
create_db! {
SchedulerDb {
LastNonce: () -> u64,
RotatedTo: (key: &[u8]) -> Vec<u8>,
}
}
impl<N: Network<Scheduler = Self>> SchedulerTrait<N> for Scheduler<N> {
type Addendum = Addendum<N>;
fn empty(&self) -> bool {
self.rotated
}
fn new<D: Db>(
_txn: &mut D::Transaction<'_>,
key: <N::Curve as Ciphersuite>::G,
network: ExternalNetworkId,
) -> Self {
assert!(N::branch_address(key).is_none());
assert!(N::change_address(key).is_none());
assert!(N::forward_address(key).is_none());
Scheduler { key, coins: network.coins().iter().copied().collect(), rotated: false }
}
fn from_db<D: Db>(
db: &D,
key: <N::Curve as Ciphersuite>::G,
network: ExternalNetworkId,
) -> io::Result<Self> {
Ok(Scheduler {
key,
coins: network.coins().iter().copied().collect(),
rotated: RotatedTo::get(db, key.to_bytes().as_ref()).is_some(),
})
}
fn can_use_branch(&self, _balance: ExternalBalance) -> bool {
false
}
fn schedule<D: Db>(
&mut self,
txn: &mut D::Transaction<'_>,
utxos: Vec<N::Output>,
payments: Vec<Payment<N>>,
key_for_any_change: <N::Curve as Ciphersuite>::G,
force_spend: bool,
) -> Vec<Plan<N>> {
for utxo in utxos {
assert!(self.coins.contains(&utxo.balance().coin));
}
let mut nonce = LastNonce::get(txn).unwrap_or(1);
let mut plans = vec![];
for chunk in payments.as_slice().chunks(N::MAX_OUTPUTS) {
assert!(!self.rotated);
plans.push(Plan {
key: self.key,
inputs: vec![],
payments: chunk.to_vec(),
change: None,
scheduler_addendum: Addendum::Nonce(nonce),
});
nonce += 1;
}
if force_spend && (!self.rotated) {
plans.push(Plan {
key: self.key,
inputs: vec![],
payments: vec![],
change: None,
scheduler_addendum: Addendum::RotateTo { nonce, new_key: key_for_any_change },
});
nonce += 1;
self.rotated = true;
RotatedTo::set(
txn,
self.key.to_bytes().as_ref(),
&key_for_any_change.to_bytes().as_ref().to_vec(),
);
}
LastNonce::set(txn, &nonce);
plans
}
fn consume_payments<D: Db>(&mut self, _txn: &mut D::Transaction<'_>) -> Vec<Payment<N>> {
vec![]
}
fn created_output<D: Db>(
&mut self,
_txn: &mut D::Transaction<'_>,
_expected: u64,
_actual: Option<u64>,
) {
panic!("Smart Contract Scheduler created a Branch output")
}
fn refund_plan<D: Db>(
&mut self,
txn: &mut D::Transaction<'_>,
output: N::Output,
refund_to: N::Address,
) -> Plan<N> {
let current_key = RotatedTo::get(txn, self.key.to_bytes().as_ref())
.and_then(|key_bytes| <N::Curve as Ciphersuite>::read_G(&mut key_bytes.as_slice()).ok())
.unwrap_or(self.key);
let nonce = LastNonce::get(txn).map_or(1, |nonce| nonce + 1);
LastNonce::set(txn, &(nonce + 1));
Plan {
key: current_key,
inputs: vec![],
payments: vec![Payment { address: refund_to, data: None, balance: output.balance() }],
change: None,
scheduler_addendum: Addendum::Nonce(nonce),
}
}
fn shim_forward_plan(_output: N::Output, _to: <N::Curve as Ciphersuite>::G) -> Option<Plan<N>> {
None
}
fn forward_plan<D: Db>(
&mut self,
_txn: &mut D::Transaction<'_>,
_output: N::Output,
_to: <N::Curve as Ciphersuite>::G,
) -> Option<Plan<N>> {
None
}
}