testnet/brain
3
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Enhanced kick contract

Browse Source 
calculating refund inside transaction query
db function to calculate price per minute for app
binded all kick contract query to prevent sql injection
proper logging for kick contract
fix table schema app_node typo and default time on deleted_app
removed wrapper contract struct
removed migrating timer.sql
fixed some clippy warnings
This commit is contained in:
Noor 2025-05-23 00:42:54 +05:30
parent 540e5bfa4c
commit db60a3f807
Signed by: noormohammedb
GPG Key ID: D83EFB8B3B967146
13 changed files with 126 additions and 158 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
saved_data.yaml
View File

@ -128,7 +128,7 @@ operators:
- 7Xw3RxbP5pvfjZ8U6yA3HHVSS9YXjKH5Vkas3JRbQYd9
app_nodes: []
7V3rEuh6j8VuwMVB5PyGqWKLmjJ4fYSv6WtrTL51NZTB:
escrow: 0
escrow: 888888888899999
email: ""
banned_users: []
vm_nodes: []

3
src/constants.rs
View File

@ -5,8 +5,7 @@ pub const CERT_PATH: &str = "/etc/detee/brain/brain-crt.pem";
pub const CERT_KEY_PATH: &str = "/etc/detee/brain/brain-key.pem";
pub const CONFIG_PATH: &str = "/etc/detee/brain/config.ini";
pub const DB_SCHEMA_FILES: [&str; 3] =
["surql/tables.sql", "surql/timer.sql", "surql/functions.sql"];
pub const DB_SCHEMA_FILES: [&str; 2] = ["surql/tables.sql", "surql/functions.sql"];
pub static ADMIN_ACCOUNTS: LazyLock<Vec<String>> = LazyLock::new(|| {
let default_admin_keys = vec![

10
src/db/app.rs
View File

@ -351,8 +351,12 @@ impl From<ActiveAppWithNode> for ActiveApp {
impl ActiveAppWithNode {
pub async fn get_by_uuid(db: &Surreal<Client>, uuid: &str) -> Result<Option<Self>, Error> {
let contract: Option<Self> =
db.query(format!("select * from {ACTIVE_APP}:{uuid} fetch out;")).await?.take(0)?;
let contract: Option<Self> = db
.query(format!("select * from {ACTIVE_APP} where id = $uuid_input fetch out;"))
.bind(("uuid_input", RecordId::from((ACTIVE_APP, uuid))))
.await?
.take(0)?;
Ok(contract)
}
@ -460,7 +464,7 @@ impl From<&old_brain::BrainData> for Vec<ActiveApp> {
.clone()
.unwrap_or_default()
.iter()
.map(|byte| format!("{:02X}", byte))
.map(|byte| format!("{byte:02X}"))
.collect();
contracts.push(ActiveApp {

138
src/db/general.rs
View File

@ -276,114 +276,60 @@ impl Operator {
}
}
pub enum WrapperContract {
Vm(ActiveVmWithNode),
App(ActiveAppWithNode),
}
impl WrapperContract {
pub async fn kick_contract(
pub async fn kick_contract(
db: &Surreal<Client>,
operator_wallet: &str,
contract_uuid: &str,
reason: &str,
) -> Result<u64, Error> {
let (
operator_id,
admin_id,
contract_id,
collected_at,
price_per_mint,
deleted_table,
platform_specific_query,
) = if let Some(active_vm) = ActiveVmWithNode::get_by_uuid(db, contract_uuid).await? {
let price_per_minute = active_vm.price_per_minute();
(
active_vm.vm_node.operator,
active_vm.admin,
active_vm.id,
active_vm.collected_at,
price_per_minute,
"deleted_vm",
"
hostname = $contract.hostname,
public_ipv4 = $contract.public_ipv4,
public_ipv6 = $contract.public_ipv6,
dtrfs_sha = $contract.dtrfs_sha,
kernel_sha = $contract.kernel_sha,
",
)
) -> Result<u64, Error> {
let (contract_id, operator_id, admin_id, app_or_vm) =
if let Some(active_vm) = ActiveVmWithNode::get_by_uuid(db, contract_uuid).await? {
(active_vm.id, active_vm.vm_node.operator, active_vm.admin, "vm")
} else if let Some(active_app) = ActiveAppWithNode::get_by_uuid(db, contract_uuid).await? {
let price_per_minute = Into::<ActiveApp>::into(active_app.clone()).price_per_minute();
(
active_app.app_node.operator,
active_app.admin,
active_app.id,
active_app.collected_at,
price_per_minute,
"deleted_app",
"
app_name = $contract.app_name,
host_ipv4 = $contract.host_ipv4,
mr_enclave = $contract.mr_enclave,
package_url= $contract.package_url,
hratls_pubkey = $contract.hratls_pubkey,
",
)
(active_app.id, active_app.app_node.operator, active_app.admin, "app")
} else {
return Err(Error::ContractNotFound);
};
let operator = operator_id.key().to_string();
let admin = admin_id.key().to_string();
if operator != operator_wallet {
if operator_id.key().to_string() != operator_wallet {
return Err(Error::AccessDenied);
}
let mut minutes_to_refund =
chrono::Utc::now().signed_duration_since(*collected_at).num_minutes().unsigned_abs();
let one_week_minute = 10080;
if minutes_to_refund > one_week_minute {
minutes_to_refund = one_week_minute;
}
let refund_amount = minutes_to_refund * price_per_mint;
log::debug!("Removing {refund_amount} escrow from {} and giving it to {}", operator, admin);
log::debug!("Kicking contract {contract_id} by operator {operator_id} for reason: '{reason}'",);
let transaction_query = format!(
"
BEGIN TRANSACTION;
LET $contract = {contract_id};
LET $operator_account = {operator_id};
LET $reason = '{reason}';
LET $refund_amount = {refund_amount};
LET $deleted_contract = {deleted_table}:{contract_uuid};
LET $id = record::id($contract.id);
LET $reason = $reason_str_input;
LET $contract_id = record::id($contract.id);
LET $admin = $contract.in;
LET $node = $contract.out;
-- move contract into deleted state
LET $active_contract = (select * from $contract)[0];
LET $deleted_contract = $active_contract.patch([{{
'op': 'replace',
'path': 'id',
'value': type::record('deleted_{app_or_vm}:' + $contract_id)
}}]);
LET $deleted_contract = (INSERT RELATION INTO deleted_{app_or_vm} ( $deleted_contract ) RETURN AFTER)[0];
RELATE $admin->{deleted_table}->$node
SET id = $id,
{platform_specific_query}
mapped_ports = $contract.mapped_ports,
disk_size_gb = $contract.disk_size_gb,
vcpus = $contract.vcpus,
memory_mb = $contract.memory_mb,
created_at = $contract.created_at,
deleted_at = time::now(),
price_per_unit = $contract.price_per_unit
;
-- calculating refund minutes
LET $one_week_minutes = duration::mins(1w);
LET $uncollected_minutes = (time::now() - $active_contract.collected_at).mins();
DELETE $contract;
LET $minutes_to_refund = if $uncollected_minutes > $one_week_minutes {{
$one_week_minutes;
}} ELSE {{
$uncollected_minutes;
}};
-- calculating refund amount
LET $prince_per_minute = fn::{app_or_vm}_price_per_minute($active_contract.id);
LET $refund_amount = $prince_per_minute * $minutes_to_refund;
LET $refund = IF SELECT * FROM {KICK} WHERE out = $admin.id AND created_at > time::now() - 24h {{
0
}} ELSE IF $operator_account.escrow <= $refund_amount {{
@ -392,14 +338,14 @@ impl WrapperContract {
$refund_amount;
}};
RELATE $operator_account->{KICK}->$admin
SET id = $id,
SET id = $contract_id,
reason = $reason,
contract = $deleted_contract,
node = $node,
contract = $deleted_contract.id,
node = $node.id,
created_at = time::now()
;
DELETE $active_contract.id;
-- update balances
UPDATE $operator_account SET escrow -= $refund;
@ -408,16 +354,28 @@ impl WrapperContract {
}};
UPDATE $admin SET balance += $refund;
SELECT * FROM $refund;
$refund;
COMMIT TRANSACTION;
",
);
log::trace!("kick_contract transaction_query: {}", &transaction_query);
let refunded: Option<u64> = db.query(transaction_query).await?.take(14)?;
let mut query_res =
db.query(transaction_query).bind(("reason_str_input", reason.to_string())).await?;
log::trace!("transaction_query response: {:?}", &query_res);
let query_error = query_res.take_errors();
if !query_error.is_empty() {
log::error!("kick_contract query error: {query_error:?}");
return Err(Error::FailedKickContract(contract_id.to_string()));
}
let refunded: Option<u64> = query_res.take(20)?;
let refunded_amount = refunded.ok_or(Error::FailedToCreateDBEntry("Refund".to_string()))?;
log::info!("Refunded: {refunded_amount} to {admin_id}");
Ok(refunded_amount)
}
}

9
src/db/mod.rs
View File

@ -43,6 +43,8 @@ pub enum Error {
AccessDenied,
#[error("Failed to delete contract {0}")]
FailedToDeleteContract(String),
#[error("Failed to kick contract {0}")]
FailedKickContract(String),
}
pub mod prelude {
@ -77,8 +79,11 @@ pub async fn migration0(
let active_vm: Vec<ActiveVm> = old_data.into();
let active_app: Vec<ActiveApp> = old_data.into();
for schema in DB_SCHEMA_FILES.map(std::fs::read_to_string) {
db.query(schema?).await?;
for schema_data in DB_SCHEMA_FILES.map(|path| (std::fs::read_to_string(path), path)) {
let schema_file = schema_data.1;
println!("Loading schema from {schema_file}");
let schema = schema_data.0?;
db.query(schema).await?;
}
println!("Inserting accounts...");

19
src/db/vm.rs
View File

@ -219,7 +219,7 @@ impl NewVmReq {
->vm_node:{vm_node}
CONTENT {{
created_at: time::now(), hostname: '{}', vcpus: {}, memory_mb: {}, disk_size_gb: {},
extra_ports: {}, public_ipv4: {:?}, public_ipv6: {:?},
extra_ports: {:?}, public_ipv4: {:?}, public_ipv6: {:?},
dtrfs_url: '{}', dtrfs_sha: '{}', kernel_url: '{}', kernel_sha: '{}',
price_per_unit: {}, locked_nano: {locked_nano}, error: ''
}};
@ -229,7 +229,7 @@ impl NewVmReq {
self.vcpus,
self.memory_mb,
self.disk_size_gb,
format!("{:?}", self.extra_ports,),
self.extra_ports,
self.public_ipv4,
self.public_ipv6,
self.dtrfs_url,
@ -242,14 +242,12 @@ impl NewVmReq {
let mut query_resp = db.query(query).await?;
let resp_err = query_resp.take_errors();
if let Some(insufficient_funds_error) = resp_err.get(&1) {
if let surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error)) =
insufficient_funds_error
if let Some(surrealdb::Error::Api(surrealdb::error::Api::Query(tx_query_error))) =
resp_err.get(&1)
{
log::error!("Transaction error: {}", tx_query_error);
log::error!("Transaction error: {tx_query_error}");
return Err(Error::InsufficientFunds);
}
}
Ok(())
}
@ -753,8 +751,11 @@ impl From<ActiveVmWithNode> for ActiveVm {
impl ActiveVmWithNode {
pub async fn get_by_uuid(db: &Surreal<Client>, uuid: &str) -> Result<Option<Self>, Error> {
let contract: Option<Self> =
db.query(format!("select * from {ACTIVE_VM}:{uuid} fetch out;")).await?.take(0)?;
let contract: Option<Self> = db
.query(format!("select * from {ACTIVE_VM} where id = $uuid_input fetch out;"))
.bind(("uuid_input", RecordId::from((ACTIVE_VM, uuid))))
.await?
.take(0)?;
Ok(contract)
}

2
src/grpc/app.rs
View File

@ -115,7 +115,7 @@ impl BrainAppDaemon for AppDaemonServer {
let mut req_stream = req.into_inner();
let pubkey: String;
if let Some(Ok(msg)) = req_stream.next().await {
log::debug!("App daemon_messages received auth message: {:?}", msg);
log::debug!("App daemon_messages received auth message: {msg:?}");
if let Some(daemon_message_app::Msg::Auth(auth)) = msg.msg {
pubkey = auth.pubkey.clone();
check_sig_from_parts(

11
src/grpc/general.rs
View File

@ -127,12 +127,8 @@ impl BrainGeneralCli for GeneralCliServer {
async fn kick_contract(&self, req: Request<KickReq>) -> Result<Response<KickResp>, Status> {
let req = check_sig_from_req(req)?;
match db::WrapperContract::kick_contract(
&self.db,
&req.operator_wallet,
&req.contract_uuid,
&req.reason,
)
log::info!("Kicking contract: {}, by: {}", req.contract_uuid, req.operator_wallet);
match db::kick_contract(&self.db, &req.operator_wallet, &req.contract_uuid, &req.reason)
.await
{
Ok(nano_lp) => Ok(Response::new(KickResp { nano_lp })),
@ -144,10 +140,11 @@ impl BrainGeneralCli for GeneralCliServer {
| db::Error::FailedToDeleteContract(_)
) =>
{
log::warn!("Failed to kick contract: {e:?}");
Err(Status::failed_precondition(e.to_string()))
}
Err(e) => {
log::info!("Failed to kick contract: {e:?}");
log::error!("Failed to kick contract: {e:?}");
Err(Status::unknown(
"Unknown error. Please try again or contact the DeTEE devs team.",
))

8
src/grpc/mod.rs
View File

@ -58,7 +58,7 @@ impl_pubkey_getter!(RegisterAppNodeReq);
impl_pubkey_getter!(AppNodeFilters);
pub fn check_sig_from_req<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) -> Result<T, Status> {
log::trace!("Checking signature from request: {:?}", req);
log::trace!("Checking signature from request: {req:?}");
let time = match req.metadata().get("timestamp") {
Some(t) => t.clone(),
None => return Err(Status::unauthenticated("Timestamp not found in metadata.")),
@ -73,8 +73,7 @@ pub fn check_sig_from_req<T: std::fmt::Debug + PubkeyGetter>(req: Request<T>) ->
let seconds_elapsed = now.signed_duration_since(parsed_time).num_seconds();
if !(-4..=4).contains(&seconds_elapsed) {
return Err(Status::unauthenticated(format!(
"Date is not within 4 sec of the time of the server: CLI {} vs Server {}",
parsed_time, now
"Date is not within 4 sec of the time of the server: CLI {parsed_time} vs Server {now}",
)));
}
@ -131,8 +130,7 @@ pub fn check_sig_from_parts(pubkey: &str, time: &str, msg: &str, sig: &str) -> R
let seconds_elapsed = now.signed_duration_since(parsed_time).num_seconds();
if !(-4..=4).contains(&seconds_elapsed) {
return Err(Status::unauthenticated(format!(
"Date is not within 4 sec of the time of the server: CLI {} vs Server {}",
parsed_time, now
"Date is not within 4 sec of the time of the server: CLI {parsed_time} vs Server {now}",
)));
}

14
src/grpc/types.rs
View File

@ -263,15 +263,15 @@ impl From<db::ActiveAppWithNode> for AppContract {
node_pubkey: value.app_node.id.key().to_string(),
public_ipv4: value.host_ipv4,
resource: Some(AppResource {
memory_mb: value.memory_mb as u32,
disk_mb: value.disk_size_gb as u32,
vcpu: value.vcpus as u32,
ports: value.mapped_ports.iter().map(|(_, g)| *g as u32).collect(),
memory_mb: value.memory_mb,
disk_mb: value.disk_size_gb,
vcpu: value.vcpus,
ports: value.mapped_ports.iter().map(|(_, g)| *g).collect(),
}),
mapped_ports: value
.mapped_ports
.iter()
.map(|(h, g)| MappedPort { host_port: *h as u32, guest_port: *g as u32 })
.map(|(h, g)| MappedPort { host_port: *h, guest_port: *g })
.collect(),
created_at: value.created_at.to_rfc3339(),
@ -294,7 +294,7 @@ impl From<NewAppReq> for db::NewAppReq {
.public_package_mr_enclave
.unwrap_or_default()
.iter()
.fold(String::new(), |acc, x| acc + &format!("{:02x?}", x));
.fold(String::new(), |acc, x| acc + &format!("{x:02x?}"));
Self {
id: RecordId::from((NEW_APP_REQ, nanoid!(40, &ID_ALPHABET))),
@ -377,7 +377,7 @@ impl From<db::ActiveApp> for NewAppRes {
let mapped_ports = val
.mapped_ports
.iter()
.map(|(h, g)| MappedPort { host_port: *h as u32, guest_port: *g as u32 })
.map(|(h, g)| MappedPort { host_port: *h, guest_port: *g })
.collect();
Self {
uuid: val.id.key().to_string(),

10
surql/functions.sql
View File

@ -25,3 +25,13 @@ DEFINE FUNCTION OVERWRITE fn::delete_vm(
DELETE $vm.id;
};
DEFINE FUNCTION OVERWRITE fn::app_price_per_minute(
$app_id: record
) {
LET $app = (select * from $app_id)[0];
RETURN
(($app.vcpus * 5) +
($app.memory_mb / 200) +
($app.disk_size_gb / 10))
* $app.price_per_unit;
};

4
surql/tables.sql
View File

@ -129,7 +129,7 @@ DEFINE FIELD vcpus ON TABLE deleted_app TYPE int;
DEFINE FIELD memory_mb ON TABLE deleted_app TYPE int;
DEFINE FIELD disk_size_gb ON TABLE deleted_app TYPE int;
DEFINE FIELD created_at ON TABLE deleted_app TYPE datetime;
DEFINE FIELD deleted_at ON TABLE deleted_app TYPE datetime;
DEFINE FIELD deleted_at ON TABLE deleted_app TYPE datetime DEFAULT time::now();;
DEFINE FIELD price_per_unit ON TABLE deleted_app TYPE int;
DEFINE FIELD mr_enclave ON TABLE deleted_app TYPE string;
DEFINE FIELD package_url ON TABLE deleted_app TYPE string;
@ -142,7 +142,7 @@ DEFINE TABLE kick TYPE RELATION FROM account TO account;
DEFINE FIELD created_at ON TABLE kick TYPE datetime;
DEFINE FIELD reason ON TABLE kick TYPE string;
DEFINE FIELD contract ON TABLE kick TYPE record<deleted_vm|deleted_app>;
DEFINE FIELD node ON TABLE kick TYPE record<vm_node|app_name>;
DEFINE FIELD node ON TABLE kick TYPE record<vm_node|app_node>;
DEFINE TABLE report TYPE RELATION FROM account TO vm_node|app_node;
DEFINE FIELD created_at ON TABLE report TYPE datetime;

16
tests/grpc_general_test.rs
View File

@ -222,22 +222,18 @@ async fn test_kick_contract() {
// 7. refund of multiple contract kick in a day for same user
env_logger::builder()
.filter_level(log::LevelFilter::Trace)
.filter_level(log::LevelFilter::Info)
.filter_module("tungstenite", log::LevelFilter::Debug)
.filter_module("tokio_tungstenite", log::LevelFilter::Debug)
.init();
let db_conn = prepare_test_db().await.unwrap();
let operator_wallet = "BFopWmwcZAMF1h2PFECZNdEucdZfnZZ32p6R9ZaBiVsS";
let contract_uuid = "26577f1c98674a1780a86cf0490f1270";
let reason = "test reason";
let contract_uuid = "e3d01f252b2a410b80e312f44e474334";
let operator_wallet = "7V3rEuh6j8VuwMVB5PyGqWKLmjJ4fYSv6WtrTL51NZTB";
let reason = "'; THROW 'Injected error'; --"; // sql injection query
let kick_response = surreal_brain::db::general::WrapperContract::kick_contract(
&db_conn,
operator_wallet,
contract_uuid,
reason,
)
let kick_response =
surreal_brain::db::general::kick_contract(&db_conn, operator_wallet, contract_uuid, reason)
.await;
match kick_response {
Ok(refund_amount) => {