Correctly format SS58-prefixed addresses in the CLI (#845)

* Fix SS58 formatting of addresses.

* cargo fmt --all

* Use only lifetime hint.

* Update relays/bin-substrate/src/cli.rs

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

* Remove unnecessary optimisation.

* Add re-formatting test.

* cargo fmt --all

Co-authored-by: Hernando Castano <HCastano@users.noreply.github.com>

bridges/relays/bin-substrate/src/cli.rs

@@ -16,6 +16,8 @@
 
 //! Deal with CLI args of substrate-to-substrate relay.
 
+use std::convert::TryInto;
+
 use bp_messages::LaneId;
 use codec::{Decode, Encode};
 use sp_runtime::app_crypto::Ss58Codec;
@@ -239,48 +241,76 @@ arg_enum! {
 }
 
 /// Generic account id with custom parser.
-#[derive(Debug)]
+#[derive(Debug, Clone)]
 pub struct AccountId {
 	account: sp_runtime::AccountId32,
-	version: sp_core::crypto::Ss58AddressFormat,
+	ss58_format: sp_core::crypto::Ss58AddressFormat,
+}
+
+impl std::fmt::Display for AccountId {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		write!(fmt, "{}", self.account.to_ss58check_with_version(self.ss58_format))
+	}
 }
 
 impl std::str::FromStr for AccountId {
 	type Err = String;
 
 	fn from_str(s: &str) -> Result<Self, Self::Err> {
-		let (account, version) = sp_runtime::AccountId32::from_ss58check_with_version(s)
+		let (account, ss58_format) = sp_runtime::AccountId32::from_ss58check_with_version(s)
 			.map_err(|err| format!("Unable to decode SS58 address: {:?}", err))?;
-		Ok(Self { account, version })
+		Ok(Self { account, ss58_format })
 	}
 }
 
+const SS58_FORMAT_PROOF: &str = "u16 -> Ss58Format is infallible; qed";
+
 impl AccountId {
-	/// Perform runtime checks of SS58 version and get Rialto's AccountId.
-	pub fn into_rialto(self) -> bp_rialto::AccountId {
-		self.check_and_get("Rialto", rialto_runtime::SS58Prefix::get())
+	/// Create new SS58-formatted address from raw account id.
+	pub fn from_raw<T: AccountChain>(account: sp_runtime::AccountId32) -> Self {
+		Self {
+			account,
+			ss58_format: T::ss58_format().try_into().expect(SS58_FORMAT_PROOF),
+		}
 	}
 
-	/// Perform runtime checks of SS58 version and get Millau's AccountId.
-	pub fn into_millau(self) -> bp_millau::AccountId {
-		self.check_and_get("Millau", millau_runtime::SS58Prefix::get())
-	}
-
-	/// Check SS58Prefix and return the account id.
-	fn check_and_get(self, net: &str, expected_prefix: u8) -> sp_runtime::AccountId32 {
-		let version: u16 = self.version.into();
-		println!("Version: {} vs {}", version, expected_prefix);
-		if version != expected_prefix as u16 {
+	/// Enforces formatting account to be for given `AccountChain` type.
+	///
+	/// This will change the `ss58format` of the account to match the requested one.
+	/// Note that a warning will be produced in case the current format does not match
+	/// the requested one, but the conversion always succeeds.
+	pub fn enforce_chain<T: AccountChain>(&mut self) {
+		let original = self.clone();
+		self.ss58_format = T::ss58_format().try_into().expect(SS58_FORMAT_PROOF);
+		log::debug!("{} SS58 format: {} (RAW: {})", self, self.ss58_format, self.account);
+		if original.ss58_format != self.ss58_format {
 			log::warn!(
 				target: "bridge",
-				"Following address: {} does not seem to match {}'s format, got: {}",
-				self.account,
-				net,
-				self.version,
+				"Address {} does not seem to match {}'s SS58 format (got: {}, expected: {}).\nConverted to: {}",
+				original,
+				T::NAME,
+				original.ss58_format,
+				self.ss58_format,
+				self,
 			)
 		}
-		self.account
 	}
+
+	/// Returns the raw (no SS58-prefixed) account id.
+	pub fn raw_id(&self) -> sp_runtime::AccountId32 {
+		self.account.clone()
+	}
+}
+
+/// A trait representing an account address bound to a specific chain.
+///
+/// Can be used to convert [`AccountId`] formatting to be chain-specific.
+pub trait AccountChain {
+	/// Network name associated with the SS58 format.
+	const NAME: &'static str;
+
+	/// Numeric value of SS58 format.
+	fn ss58_format() -> u16;
 }
 
 /// Lane id.
@@ -414,3 +444,31 @@ macro_rules! declare_chain_options {
 		}
 	};
 }
+
+#[cfg(test)]
+mod tests {
+	use std::str::FromStr;
+
+	use super::*;
+
+	#[test]
+	fn should_format_addresses_with_ss58_format() {
+		// given
+		let rialto1 = "5sauUXUfPjmwxSgmb3tZ5d6yx24eZX4wWJ2JtVUBaQqFbvEU";
+		let rialto2 = "5rERgaT1Z8nM3et2epA5i1VtEBfp5wkhwHtVE8HK7BRbjAH2";
+		let millau1 = "752paRyW1EGfq9YLTSSqcSJ5hqnBDidBmaftGhBo8fy6ypW9";
+		let millau2 = "74GNQjmkcfstRftSQPJgMREchqHM56EvAUXRc266cZ1NYVW5";
+
+		let expected = vec![rialto1, rialto2, millau1, millau2];
+
+		// when
+		let parsed = expected
+			.iter()
+			.map(|s| AccountId::from_str(s).unwrap())
+			.collect::<Vec<_>>();
+
+		let actual = parsed.iter().map(|a| format!("{}", a)).collect::<Vec<_>>();
+
+		assert_eq!(actual, expected)
+	}
+}

bridges/relays/bin-substrate/src/main.rs

@@ -18,7 +18,7 @@
 
 #![warn(missing_docs)]
 
-use relay_utils::initialize::initialize_relay;
+use relay_utils::initialize::initialize_logger;
 
 mod cli;
 mod finality_pipeline;
@@ -31,7 +31,7 @@ mod messages_target;
 mod rialto_millau;
 
 fn main() {
-	initialize_relay();
+	initialize_logger(false);
 	let command = cli::parse_args();
 	let run = command.run();
 	let result = async_std::task::block_on(run);

bridges/relays/bin-substrate/src/rialto_millau/mod.rs

@@ -30,7 +30,7 @@ pub type RialtoClient = relay_substrate_client::Client<Rialto>;
 /// Westend node client.
 pub type WestendClient = relay_substrate_client::Client<Westend>;
 
-use crate::cli::{ExplicitOrMaximal, HexBytes, Origins};
+use crate::cli::{AccountChain, AccountId, ExplicitOrMaximal, HexBytes, Origins};
 use codec::{Decode, Encode};
 use frame_support::weights::{GetDispatchInfo, Weight};
 use pallet_bridge_dispatch::{CallOrigin, MessagePayload};
@@ -426,23 +426,21 @@ async fn run_estimate_fee(cmd: cli::EstimateFee) -> Result<(), String> {
 
 async fn run_derive_account(cmd: cli::DeriveAccount) -> Result<(), String> {
 	match cmd {
-		cli::DeriveAccount::RialtoToMillau { account } => {
-			let account = account.into_rialto();
-			let acc = bp_runtime::SourceAccount::Account(account.clone());
+		cli::DeriveAccount::RialtoToMillau { mut account } => {
+			account.enforce_chain::<Rialto>();
+			let acc = bp_runtime::SourceAccount::Account(account.raw_id());
 			let id = bp_millau::derive_account_from_rialto_id(acc);
-			println!(
-				"{} (Rialto)\n\nCorresponding (derived) account id:\n-> {} (Millau)",
-				account, id
-			)
+			let derived_account = AccountId::from_raw::<Millau>(id);
+			println!("Source address:\n{} (Rialto)", account);
+			println!("->Corresponding (derived) address:\n{} (Millau)", derived_account);
 		}
-		cli::DeriveAccount::MillauToRialto { account } => {
-			let account = account.into_millau();
-			let acc = bp_runtime::SourceAccount::Account(account.clone());
+		cli::DeriveAccount::MillauToRialto { mut account } => {
+			account.enforce_chain::<Millau>();
+			let acc = bp_runtime::SourceAccount::Account(account.raw_id());
 			let id = bp_rialto::derive_account_from_millau_id(acc);
-			println!(
-				"{} (Millau)\n\nCorresponding (derived) account id:\n-> {} (Rialto)",
-				account, id
-			)
+			let derived_account = AccountId::from_raw::<Rialto>(id);
+			println!("Source address:\n{} (Millau)", account);
+			println!("->Corresponding (derived) address:\n{} (Rialto)", derived_account);
 		}
 	}
 
@@ -650,9 +648,10 @@ impl cli::MillauToRialtoMessagePayload {
 		match self {
 			Self::Raw { data } => MessagePayload::decode(&mut &*data.0)
 				.map_err(|e| format!("Failed to decode Millau's MessagePayload: {:?}", e)),
-			Self::Message { message, sender } => {
+			Self::Message { message, mut sender } => {
 				let spec_version = rialto_runtime::VERSION.spec_version;
-				let origin = CallOrigin::SourceAccount(sender.into_millau());
+				sender.enforce_chain::<Millau>();
+				let origin = CallOrigin::SourceAccount(sender.raw_id());
 				let call = message.into_call()?;
 				let weight = call.get_dispatch_info().weight;
 
@@ -670,9 +669,10 @@ impl cli::RialtoToMillauMessagePayload {
 		match self {
 			Self::Raw { data } => MessagePayload::decode(&mut &*data.0)
 				.map_err(|e| format!("Failed to decode Rialto's MessagePayload: {:?}", e)),
-			Self::Message { message, sender } => {
+			Self::Message { message, mut sender } => {
 				let spec_version = millau_runtime::VERSION.spec_version;
-				let origin = CallOrigin::SourceAccount(sender.into_rialto());
+				sender.enforce_chain::<Rialto>();
+				let origin = CallOrigin::SourceAccount(sender.raw_id());
 				let call = message.into_call()?;
 				let weight = call.get_dispatch_info().weight;
 
@@ -750,9 +750,9 @@ impl cli::ToRialtoMessage {
 					),
 				)))
 			}
-			cli::ToRialtoMessage::Transfer { recipient, amount } => {
-				let recipient = recipient.into_rialto();
-				rialto_runtime::Call::Balances(rialto_runtime::BalancesCall::transfer(recipient, amount))
+			cli::ToRialtoMessage::Transfer { mut recipient, amount } => {
+				recipient.enforce_chain::<Rialto>();
+				rialto_runtime::Call::Balances(rialto_runtime::BalancesCall::transfer(recipient.raw_id(), amount))
 			}
 			cli::ToRialtoMessage::MillauSendMessage { lane, payload, fee } => {
 				let payload = cli::RialtoToMillauMessagePayload::Raw { data: payload }.into_payload()?;
@@ -787,9 +787,9 @@ impl cli::ToMillauMessage {
 					),
 				)))
 			}
-			cli::ToMillauMessage::Transfer { recipient, amount } => {
-				let recipient = recipient.into_millau();
-				millau_runtime::Call::Balances(millau_runtime::BalancesCall::transfer(recipient, amount))
+			cli::ToMillauMessage::Transfer { mut recipient, amount } => {
+				recipient.enforce_chain::<Millau>();
+				millau_runtime::Call::Balances(millau_runtime::BalancesCall::transfer(recipient.raw_id(), amount))
 			}
 			cli::ToMillauMessage::RialtoSendMessage { lane, payload, fee } => {
 				let payload = cli::MillauToRialtoMessagePayload::Raw { data: payload }.into_payload()?;
@@ -808,6 +808,22 @@ impl cli::ToMillauMessage {
 	}
 }
 
+impl AccountChain for Rialto {
+	const NAME: &'static str = "Rialto";
+
+	fn ss58_format() -> u16 {
+		rialto_runtime::SS58Prefix::get() as u16
+	}
+}
+
+impl AccountChain for Millau {
+	const NAME: &'static str = "Millau";
+
+	fn ss58_format() -> u16 {
+		millau_runtime::SS58Prefix::get() as u16
+	}
+}
+
 #[cfg(test)]
 mod tests {
 	use super::*;
@@ -977,4 +993,25 @@ mod tests {
 			extra_bytes_in_transaction,
 		);
 	}
+
+	#[test]
+	fn should_reformat_addresses() {
+		// given
+		let mut rialto1: AccountId = "5sauUXUfPjmwxSgmb3tZ5d6yx24eZX4wWJ2JtVUBaQqFbvEU".parse().unwrap();
+		let mut millau1: AccountId = "752paRyW1EGfq9YLTSSqcSJ5hqnBDidBmaftGhBo8fy6ypW9".parse().unwrap();
+
+		// when
+		rialto1.enforce_chain::<Millau>();
+		millau1.enforce_chain::<Rialto>();
+
+		// then
+		assert_eq!(
+			&format!("{}", rialto1),
+			"752paRyW1EGfq9YLTSSqcSJ5hqnBDidBmaftGhBo8fy6ypW9"
+		);
+		assert_eq!(
+			&format!("{}", millau1),
+			"5sauUXUfPjmwxSgmb3tZ5d6yx24eZX4wWJ2JtVUBaQqFbvEU"
+		);
+	}
 }

bridges/relays/utils/src/initialize.rs

@@ -16,40 +16,80 @@
 
 //! Relayer initialization functions.
 
-use std::io::Write;
+use std::{fmt::Display, io::Write};
 
 /// Initialize relay environment.
 pub fn initialize_relay() {
+	initialize_logger(true);
+}
+
+/// Initialize Relay logger instance.
+pub fn initialize_logger(with_timestamp: bool) {
 	let mut builder = env_logger::Builder::new();
 	builder.filter_level(log::LevelFilter::Warn);
 	builder.filter_module("bridge", log::LevelFilter::Info);
 	builder.parse_default_env();
-	builder.format(move |buf, record| {
-		writeln!(buf, "{}", {
+	if with_timestamp {
+		builder.format(move |buf, record| {
 			let timestamp = time::OffsetDateTime::try_now_local()
 				.unwrap_or_else(|_| time::OffsetDateTime::now_utc())
 				.format("%Y-%m-%d %H:%M:%S %z");
-			if cfg!(windows) {
-				format!("{} {} {} {}", timestamp, record.level(), record.target(), record.args())
+
+			let log_level = color_level(record.level());
+			let log_target = color_target(record.target());
+			let timestamp = if cfg!(windows) {
+				Either::Left(timestamp)
 			} else {
-				use ansi_term::Colour as Color;
-				let log_level = match record.level() {
-					log::Level::Error => Color::Fixed(9).bold().paint(record.level().to_string()),
-					log::Level::Warn => Color::Fixed(11).bold().paint(record.level().to_string()),
-					log::Level::Info => Color::Fixed(10).paint(record.level().to_string()),
-					log::Level::Debug => Color::Fixed(14).paint(record.level().to_string()),
-					log::Level::Trace => Color::Fixed(12).paint(record.level().to_string()),
-				};
-				format!(
-					"{} {} {} {}",
-					Color::Fixed(8).bold().paint(timestamp),
-					log_level,
-					Color::Fixed(8).paint(record.target()),
-					record.args()
-				)
-			}
-		})
-	});
+				Either::Right(ansi_term::Colour::Fixed(8).bold().paint(timestamp))
+			};
+
+			writeln!(buf, "{} {} {} {}", timestamp, log_level, log_target, record.args(),)
+		});
+	} else {
+		builder.format(move |buf, record| {
+			let log_level = color_level(record.level());
+			let log_target = color_target(record.target());
+
+			writeln!(buf, "{} {} {}", log_level, log_target, record.args(),)
+		});
+	}
 
 	builder.init();
 }
+
+enum Either<A, B> {
+	Left(A),
+	Right(B),
+}
+impl<A: Display, B: Display> Display for Either<A, B> {
+	fn fmt(&self, fmt: &mut std::fmt::Formatter) -> std::fmt::Result {
+		match self {
+			Self::Left(a) => write!(fmt, "{}", a),
+			Self::Right(b) => write!(fmt, "{}", b),
+		}
+	}
+}
+
+fn color_target(target: &str) -> impl Display + '_ {
+	if cfg!(windows) {
+		Either::Left(target)
+	} else {
+		Either::Right(ansi_term::Colour::Fixed(8).paint(target))
+	}
+}
+
+fn color_level(level: log::Level) -> impl Display {
+	if cfg!(windows) {
+		Either::Left(level)
+	} else {
+		let s = level.to_string();
+		use ansi_term::Colour as Color;
+		Either::Right(match level {
+			log::Level::Error => Color::Fixed(9).bold().paint(s),
+			log::Level::Warn => Color::Fixed(11).bold().paint(s),
+			log::Level::Info => Color::Fixed(10).paint(s),
+			log::Level::Debug => Color::Fixed(14).paint(s),
+			log::Level::Trace => Color::Fixed(12).paint(s),
+		})
+	}
+}