use dusa_collection_utils::{
core::errors::{ErrorArrayItem, Errors},
core::logger::LogLevel,
core::types::{pathtype::PathType, stringy::Stringy},
log,
};
use rand::RngExt;
use serde::{Deserialize, Serialize};
use std::{
io::{Read, Write},
time::Duration,
};
use tokio::time::sleep;
use crate::{encryption::simple_encrypt, timestamp::current_timestamp};
#[cfg(target_os = "linux")]
use dusa_collection_utils::platform::functions::{create_hash, truncate};
pub const IDENTITYPATHSTR: &str = "/opt/artisan/.identity";
pub const HASH_LENGTH: usize = 28;
pub const CUSTOM_EPOCH: u64 = 1_047_587_400;
pub struct SnowflakeIDGenerator {
custom_epoch: u64,
datacenter_id: u8,
machine_id: u8,
sequence: u16,
last_timestamp: u64,
}
#[cfg(target_os = "linux")]
impl SnowflakeIDGenerator {
pub fn new(datacenter_id: u8, machine_id: u8) -> Result<Self, ()> {
if datacenter_id > 31 {
log!(LogLevel::Error, "Datacenter ID must be between 0 and 31");
return Err(());
}
if machine_id > 31 {
log!(LogLevel::Error, "Machine ID must be between 0 and 31");
return Err(());
}
Ok(Self {
custom_epoch: CUSTOM_EPOCH,
datacenter_id,
machine_id,
sequence: 0,
last_timestamp: 0,
})
}
fn wait_for_next_millis(last_timestamp: u64) -> u64 {
let mut timestamp = current_timestamp();
while timestamp <= last_timestamp {
timestamp = current_timestamp();
}
timestamp
}
pub async fn generate_id(&mut self) -> u64 {
let mut timestamp = current_timestamp();
if timestamp < self.last_timestamp {
sleep(Duration::from_millis(10)).await;
if timestamp < self.last_timestamp {
log!(
LogLevel::Error,
"Clock moved backwards. Refusing to generate ID."
);
return 0;
}
}
if timestamp == self.last_timestamp {
self.sequence = (self.sequence + 1) & 0xFFF; if self.sequence == 0 {
timestamp = Self::wait_for_next_millis(self.last_timestamp);
}
} else {
self.sequence = 0;
}
self.last_timestamp = timestamp;
((timestamp - self.custom_epoch) << 22)
| ((self.datacenter_id as u64) << 17)
| ((self.machine_id as u64) << 12)
| (self.sequence as u64)
}
}
#[derive(Serialize, Deserialize, Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct Identifier {
pub id: u64,
_signature: Stringy,
}
#[cfg(target_os = "linux")]
impl Identifier {
fn generate_signature(id: u64) -> Stringy {
truncate(&*create_hash(format!("{}", id)), HASH_LENGTH)
}
pub async fn new() -> Result<Self, ErrorArrayItem> {
let datacenter_id = rand::rng().random_range(1..=5);
let machine_id = rand::rng().random_range(1..=5);
let mut big_id: SnowflakeIDGenerator = SnowflakeIDGenerator::new(datacenter_id, machine_id)
.map_err(|_| {
ErrorArrayItem::new(
Errors::GeneralError,
"Error generating system ID".to_owned(),
)
})?;
let id = big_id.generate_id().await;
Ok(Self {
id,
_signature: Self::generate_signature(id),
})
}
pub async fn verify(&self) -> bool {
let given_signature = self._signature.clone();
let new_signature = Self::generate_signature(self.id);
given_signature == new_signature
}
pub async fn load() -> Result<Option<Self>, ErrorArrayItem> {
let identifier_path: PathType = PathType::Str(IDENTITYPATHSTR.into());
if identifier_path.exists() {
match Self::load_from_file() {
Ok(data) => return Ok(Some(data)),
Err(err) => {
log!(LogLevel::Trace, "ERROR: Failed to load identity: {}", err);
return Ok(None);
}
}
} else {
Ok(None)
}
}
pub fn save_to_file(&self) -> Result<(), ErrorArrayItem> {
let serialized_id = serde_json::to_string_pretty(&self)?;
let mut file = std::fs::File::create(PathType::Str(IDENTITYPATHSTR.into()))?;
let mut flag = std::fs::File::create(PathType::Str("/opt/artisan/.system_ready".into()))?;
file.write_all(serialized_id.as_bytes())?;
flag.write_all(serialized_id.as_bytes())?;
Ok(())
}
pub fn load_from_file() -> Result<Self, ErrorArrayItem> {
let mut file = std::fs::File::open(PathType::Str(IDENTITYPATHSTR.into()))?;
let mut content = String::new();
file.read_to_string(&mut content)?;
let identifier: Identifier = serde_json::from_str(&content)?;
Ok(identifier)
}
pub fn to_json(&self) -> Result<String, ErrorArrayItem> {
let json_representation = serde_json::to_string_pretty(self)?;
Ok(json_representation)
}
pub async fn to_encrypted_json(&self) -> Result<Stringy, ErrorArrayItem> {
let json_representation = self.to_json().map_err(|e| {
ErrorArrayItem::new(
dusa_collection_utils::core::errors::Errors::JsonCreation,
e.to_string(),
)
})?;
let encrypted_data = simple_encrypt(json_representation.as_bytes())?;
Ok(encrypted_data)
}
pub fn display_id(&self) {
log!(LogLevel::Debug, "ID: {}", self.id);
}
pub fn display_sig(&self) {
log!(LogLevel::Debug, "SIG: {}", self._signature);
}
}
use std::fmt;
pub const AIS_PREFIX: &str = "ais_";
pub fn ais_name(component: &str) -> String {
format!("{AIS_PREFIX}{component}")
}
pub fn strip_ais_prefix(name: &str) -> Option<&str> {
name.strip_prefix(AIS_PREFIX)
}
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct IdParseError {
pub expected: &'static str,
pub got: String,
}
impl fmt::Display for IdParseError {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "expected {}, got {:?}", self.expected, self.got)
}
}
impl std::error::Error for IdParseError {}
fn is_lowercase_hex(s: &str, len: usize) -> bool {
s.len() == len && s.chars().all(|c| c.is_ascii_hexdigit() && !c.is_ascii_uppercase())
}
#[derive(Debug, Clone, PartialEq, Eq, PartialOrd, Ord, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct ProjectId(String);
impl ProjectId {
pub fn parse(s: &str) -> Result<Self, IdParseError> {
if is_lowercase_hex(s, 8) {
Ok(Self(s.to_owned()))
} else {
Err(IdParseError {
expected: "8 lowercase hex characters",
got: s.to_owned(),
})
}
}
pub fn from_parts(user: &str, repo: &str, branch: &str) -> Self {
let hash_input = format!("{branch}-{repo}-{user}");
let hash = dusa_collection_utils::platform::functions::create_hash(hash_input);
let truncated = dusa_collection_utils::platform::functions::truncate(&*hash, 8);
Self(truncated.to_string())
}
pub fn as_str(&self) -> &str {
&self.0
}
pub fn ais_name(&self) -> String {
ais_name(&self.0)
}
}
impl fmt::Display for ProjectId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl AsRef<str> for ProjectId {
fn as_ref(&self) -> &str {
&self.0
}
}
impl From<ProjectId> for String {
fn from(id: ProjectId) -> Self {
id.0
}
}
pub fn generate_project_id(user: &str, repo: &str, branch: &str) -> ProjectId {
ProjectId::from_parts(user, repo, branch)
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct UuidId(String);
impl UuidId {
pub fn new_v4() -> Self {
Self(uuid::Uuid::new_v4().to_string())
}
pub fn parse(s: &str) -> Result<Self, IdParseError> {
uuid::Uuid::parse_str(s)
.map(|u| Self(u.to_string()))
.map_err(|_| IdParseError {
expected: "a UUID",
got: s.to_owned(),
})
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for UuidId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
macro_rules! uuid_id_newtype {
($name:ident, $doc:literal) => {
#[doc = $doc]
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct $name(UuidId);
impl $name {
pub fn new_v4() -> Self {
Self(UuidId::new_v4())
}
pub fn parse(s: &str) -> Result<Self, IdParseError> {
UuidId::parse(s).map(Self)
}
pub fn as_str(&self) -> &str {
self.0.as_str()
}
}
impl fmt::Display for $name {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
};
}
uuid_id_newtype!(
OrganizationId,
"Identifies an Organization -- the tenancy boundary that owns every other resource."
);
uuid_id_newtype!(
InstanceId,
"Identifies one running copy of a Project on one Node."
);
uuid_id_newtype!(SessionId, "Identifies a Runpod GPU compute session.");
macro_rules! wire_string_u64_id {
($name:ident, $doc:literal) => {
#[doc = $doc]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct $name(pub u64);
impl $name {
pub fn get(&self) -> u64 {
self.0
}
}
impl fmt::Display for $name {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<u64> for $name {
fn from(v: u64) -> Self {
Self(v)
}
}
impl Serialize for $name {
fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
where
S: serde::Serializer,
{
serializer.serialize_str(&self.0.to_string())
}
}
impl<'de> Deserialize<'de> for $name {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: serde::Deserializer<'de>,
{
struct WireVisitor;
impl<'de> serde::de::Visitor<'de> for WireVisitor {
type Value = u64;
fn expecting(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "a u64 or a decimal string")
}
fn visit_u64<E>(self, v: u64) -> Result<u64, E> {
Ok(v)
}
fn visit_i64<E>(self, v: i64) -> Result<u64, E>
where
E: serde::de::Error,
{
u64::try_from(v).map_err(|_| E::custom("negative value for u64 id"))
}
fn visit_str<E>(self, v: &str) -> Result<u64, E>
where
E: serde::de::Error,
{
v.parse::<u64>().map_err(E::custom)
}
}
deserializer.deserialize_any(WireVisitor).map($name)
}
}
};
}
wire_string_u64_id!(NodeId, "Identifies a compute Node.");
wire_string_u64_id!(DomainId, "Identifies a Domain.");
wire_string_u64_id!(VmId, "Identifies a Proxmox-managed Vm.");
impl From<Identifier> for NodeId {
fn from(identifier: Identifier) -> Self {
NodeId(identifier.id)
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
#[serde(transparent)]
pub struct EnvironmentId(String);
impl EnvironmentId {
pub fn new(s: impl Into<String>) -> Self {
Self(s.into())
}
pub fn as_str(&self) -> &str {
&self.0
}
}
impl fmt::Display for EnvironmentId {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", self.0)
}
}
impl From<String> for EnvironmentId {
fn from(s: String) -> Self {
Self(s)
}
}
impl From<&str> for EnvironmentId {
fn from(s: &str) -> Self {
Self(s.to_owned())
}
}
#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
pub struct SecretRef {
pub project_id: ProjectId,
pub environment_id: EnvironmentId,
pub key: String,
}
impl fmt::Display for SecretRef {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}/{}/{}", self.project_id, self.environment_id, self.key)
}
}
pub use crate::urn::ResourceType;
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum Action {
Read,
Write,
Control,
Delete,
Grant,
Purchase,
}
impl Action {
pub fn as_str(&self) -> &'static str {
match self {
Action::Read => "read",
Action::Write => "write",
Action::Control => "control",
Action::Delete => "delete",
Action::Grant => "grant",
Action::Purchase => "purchase",
}
}
pub fn from_str(s: &str) -> Option<Self> {
Some(match s {
"read" => Action::Read,
"write" => Action::Write,
"control" => Action::Control,
"delete" => Action::Delete,
"grant" => Action::Grant,
"purchase" => Action::Purchase,
_ => return None,
})
}
}
#[cfg(test)]
mod action_tests {
use super::Action;
#[test]
fn every_action_round_trips_through_as_str_and_from_str() {
for action in [
Action::Read,
Action::Write,
Action::Control,
Action::Delete,
Action::Grant,
Action::Purchase,
] {
assert_eq!(Action::from_str(action.as_str()), Some(action));
}
}
#[test]
fn from_str_rejects_an_unrecognized_action() {
assert_eq!(Action::from_str("execute"), None);
assert_eq!(Action::from_str(""), None);
assert_eq!(Action::from_str("READ"), None); }
}