1
// Copyright 2019-2025 PureStake Inc.
2
// This file is part of Moonbeam.
3

            
4
// Moonbeam is free software: you can redistribute it and/or modify
5
// it under the terms of the GNU General Public License as published by
6
// the Free Software Foundation, either version 3 of the License, or
7
// (at your option) any later version.
8

            
9
// Moonbeam is distributed in the hope that it will be useful,
10
// but WITHOUT ANY WARRANTY; without even the implied warranty of
11
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
12
// GNU General Public License for more details.
13

            
14
// You should have received a copy of the GNU General Public License
15
// along with Moonbeam.  If not, see <http://www.gnu.org/licenses/>.
16

            
17
//! Module that provides types to extend xcm holding.
18

            
19
use core::marker::PhantomData;
20
use sp_core::{H160, U256};
21
use sp_std::collections::btree_map::BTreeMap;
22
use sp_std::vec::Vec;
23
use xcm::prelude::*;
24
use xcm_executor::traits::{FeeManager, FeeReason, XcmAssetTransfers};
25

            
26
environmental::environmental!(XCM_HOLDING_ERC20_ORIGINS: XcmHoldingErc20sOrigins);
27

            
28
#[cfg_attr(test, derive(PartialEq, Debug))]
29
pub(crate) enum DrainError {
30
	AssetNotFound,
31
	NotEnoughFounds,
32
	SplitError,
33
}
34

            
35
/// Xcm holding erc20 origins extension.
36
/// This extension track down the origin of alls erc20 tokens in the xcm holding.
37
#[derive(Default)]
38
pub(crate) struct XcmHoldingErc20sOrigins {
39
	map: BTreeMap<H160, Vec<(H160, U256)>>,
40
}
41
impl XcmHoldingErc20sOrigins {
42
	/// Take and remove a given amounts of erc20 tokens from the XCM holding.
43
	/// These tokens can come from one or more holders that we had tracked earlier in the XCM
44
	/// execution, so we return an array of (holder, balance).
45
6
	pub(crate) fn drain(
46
6
		&mut self,
47
6
		contract_address: H160,
48
6
		amount: U256,
49
6
	) -> Result<Vec<(H160, U256)>, DrainError> {
50
6
		let tokens_to_transfer = self.drain_inner(&contract_address, amount)?;
51

            
52
3
		self.map
53
3
			.entry(contract_address)
54
3
			.and_modify(|erc20_origins| {
55
3
				*erc20_origins = erc20_origins.split_off(tokens_to_transfer.len());
56
3
			});
57
3

            
58
3
		Ok(tokens_to_transfer)
59
6
	}
60
6
	fn drain_inner(
61
6
		&self,
62
6
		contract_address: &H160,
63
6
		mut amount: U256,
64
6
	) -> Result<Vec<(H160, U256)>, DrainError> {
65
6
		let mut tokens_to_transfer = Vec::new();
66
6
		if let Some(erc20_origins) = self.map.get(contract_address) {
67
6
			for (from, subamount) in erc20_origins {
68
5
				if &amount > subamount {
69
1
					tokens_to_transfer.push((*from, *subamount));
70
1
					amount -= *subamount;
71
4
				} else if &amount == subamount {
72
3
					tokens_to_transfer.push((*from, *subamount));
73
3
					return Ok(tokens_to_transfer);
74
				} else {
75
					// Each insertion of tokens must be drain at once
76
1
					return Err(DrainError::SplitError);
77
				}
78
			}
79
			// If there were enough tokens, we had to return in the for loop
80
1
			Err(DrainError::NotEnoughFounds)
81
		} else {
82
1
			Err(DrainError::AssetNotFound)
83
		}
84
6
	}
85
3
	pub(crate) fn insert(&mut self, contract_address: H160, who: H160, amount: U256) {
86
3
		self.map
87
3
			.entry(contract_address)
88
3
			.or_default()
89
3
			.push((who, amount));
90
3
	}
91
2
	pub(crate) fn with<R, F>(f: F) -> Option<R>
92
2
	where
93
2
		F: FnOnce(&mut Self) -> R,
94
2
	{
95
2
		XCM_HOLDING_ERC20_ORIGINS::with(|erc20s_origins| f(erc20s_origins))
96
2
	}
97
}
98

            
99
/// Xcm executor wrapper that inject xcm holding extension "XcmHoldingErc20sOrigins"
100
pub struct XcmExecutorWrapper<Config, InnerXcmExecutor>(PhantomData<(Config, InnerXcmExecutor)>);
101
impl<Config, InnerXcmExecutor> xcm::latest::ExecuteXcm<Config::RuntimeCall>
102
	for XcmExecutorWrapper<Config, InnerXcmExecutor>
103
where
104
	Config: xcm_executor::Config,
105
	InnerXcmExecutor: xcm::latest::ExecuteXcm<Config::RuntimeCall>,
106
{
107
	type Prepared = InnerXcmExecutor::Prepared;
108

            
109
43
	fn prepare(
110
43
		message: xcm::latest::Xcm<Config::RuntimeCall>,
111
43
	) -> Result<Self::Prepared, xcm::latest::Xcm<Config::RuntimeCall>> {
112
43
		InnerXcmExecutor::prepare(message)
113
43
	}
114

            
115
37
	fn execute(
116
37
		origin: impl Into<xcm::latest::Location>,
117
37
		pre: Self::Prepared,
118
37
		hash: &mut xcm::latest::XcmHash,
119
37
		weight_credit: xcm::latest::Weight,
120
37
	) -> xcm::latest::Outcome {
121
37
		let mut erc20s_origins = Default::default();
122
37
		XCM_HOLDING_ERC20_ORIGINS::using(&mut erc20s_origins, || {
123
37
			InnerXcmExecutor::execute(origin, pre, hash, weight_credit)
124
37
		})
125
37
	}
126

            
127
14
	fn charge_fees(
128
14
		location: impl Into<xcm::latest::Location>,
129
14
		fees: xcm::latest::Assets,
130
14
	) -> Result<(), xcm::latest::Error> {
131
14
		InnerXcmExecutor::charge_fees(location, fees)
132
14
	}
133
}
134

            
135
impl<Config, InnerXcmExecutor> XcmAssetTransfers for XcmExecutorWrapper<Config, InnerXcmExecutor>
136
where
137
	Config: xcm_executor::Config,
138
{
139
	type IsReserve = Config::IsReserve;
140
	type IsTeleporter = Config::IsTeleporter;
141
	type AssetTransactor = Config::AssetTransactor;
142
}
143

            
144
impl<Config, InnerXcmExecutor> FeeManager for XcmExecutorWrapper<Config, InnerXcmExecutor>
145
where
146
	Config: xcm_executor::Config,
147
	InnerXcmExecutor: FeeManager,
148
{
149
	fn is_waived(origin: Option<&Location>, r: FeeReason) -> bool {
150
		InnerXcmExecutor::is_waived(origin, r)
151
	}
152

            
153
	fn handle_fee(fee: Assets, context: Option<&XcmContext>, r: FeeReason) {
154
		InnerXcmExecutor::handle_fee(fee, context, r)
155
	}
156
}
157

            
158
#[cfg(test)]
159
mod tests {
160
	use super::*;
161

            
162
	#[test]
163
1
	fn test_xcm_holding_ext_erc20s_origins() {
164
		const TOKEN1: H160 = H160([1; 20]);
165
		const TOKEN2: H160 = H160([2; 20]);
166
		const USER1: H160 = H160([3; 20]);
167
		const USER2: H160 = H160([4; 20]);
168

            
169
		// Simple case
170
1
		let mut erc20s_origins_ = Default::default();
171
1
		XCM_HOLDING_ERC20_ORIGINS::using(&mut erc20s_origins_, || {
172
1
			XcmHoldingErc20sOrigins::with(|erc20s_origins| {
173
1
				erc20s_origins.insert(TOKEN1, USER1, U256::from(100));
174
1
				assert_eq!(
175
1
					erc20s_origins.drain(TOKEN2, U256::from(1)),
176
1
					Err(DrainError::AssetNotFound)
177
1
				);
178
1
				assert_eq!(
179
1
					erc20s_origins.drain(TOKEN1, U256::from(100)),
180
1
					Ok(vec![(USER1, U256::from(100))])
181
1
				);
182
1
			})
183
1
		});
184
1

            
185
1
		// Complex case
186
1
		let mut erc20s_origins_ = Default::default();
187
1
		XCM_HOLDING_ERC20_ORIGINS::using(&mut erc20s_origins_, || {
188
1
			XcmHoldingErc20sOrigins::with(|erc20s_origins| {
189
1
				erc20s_origins.insert(TOKEN1, USER1, U256::from(100));
190
1
				erc20s_origins.insert(TOKEN1, USER2, U256::from(200));
191
1
				assert_eq!(
192
1
					erc20s_origins.drain(TOKEN1, U256::from(100)),
193
1
					Ok(vec![(USER1, U256::from(100))])
194
1
				);
195
1
				assert_eq!(
196
1
					erc20s_origins.drain(TOKEN1, U256::from(201)),
197
1
					Err(DrainError::NotEnoughFounds)
198
1
				);
199
1
				assert_eq!(
200
1
					erc20s_origins.drain(TOKEN1, U256::from(199)),
201
1
					Err(DrainError::SplitError)
202
1
				);
203
1
				assert_eq!(
204
1
					erc20s_origins.drain(TOKEN1, U256::from(200)),
205
1
					Ok(vec![(USER2, U256::from(200))])
206
1
				);
207
1
			})
208
1
		});
209
1
	}
210
}