Compare commits
3
Commits
| Author | SHA1 | Date | |
|---|---|---|---|
 Dreaded_X
|
f03c06fdf2
|
||
|
Dreaded_X
|
27f6557197
|
||
|
Dreaded_X
|
344f345854
|
Generated
+262
-1
@@ -52,6 +52,57 @@ dependencies = [
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-channel"
|
||||
version = "2.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "924ed96dd52d1b75e9c1a3e6275715fd320f5f9439fb5a4a11fa51f4221158d2"
|
||||
dependencies = [
|
||||
"concurrent-queue",
|
||||
"event-listener-strategy",
|
||||
"futures-core",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-fs"
|
||||
version = "2.2.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8034a681df4aed8b8edbd7fbe472401ecf009251c8b40556b304567052e294c5"
|
||||
dependencies = [
|
||||
"async-lock",
|
||||
"blocking",
|
||||
"futures-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-lock"
|
||||
version = "3.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "290f7f2596bd5b78a9fec8088ccd89180d7f9f55b94b0576823bbbdc72ee8311"
|
||||
dependencies = [
|
||||
"event-listener",
|
||||
"event-listener-strategy",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "async-task"
|
||||
version = "4.7.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8b75356056920673b02621b35afd0f7dda9306d03c79a30f5c56c44cf256e3de"
|
||||
|
||||
[[package]]
|
||||
name = "async-walkdir"
|
||||
version = "2.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "37672978ae0febce7516ae0a85b53e6185159a9a28787391eb63fc44ec36037d"
|
||||
dependencies = [
|
||||
"async-fs",
|
||||
"futures-lite",
|
||||
"thiserror",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "atomic-waker"
|
||||
version = "1.1.2"
|
||||
@@ -85,6 +136,19 @@ dependencies = [
|
||||
"generic-array",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "blocking"
|
||||
version = "1.7.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a70e4329df6cb94385eed412ec92375c3cdd8a6e502493d1229b6414e4036dfa"
|
||||
dependencies = [
|
||||
"async-channel",
|
||||
"async-task",
|
||||
"futures-io",
|
||||
"futures-lite",
|
||||
"piper",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "bstr"
|
||||
version = "1.13.1"
|
||||
@@ -187,6 +251,15 @@ version = "1.0.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d07550c9036bf2ae0c684c4297d503f838287c83c53686d05370d0e139ae570"
|
||||
|
||||
[[package]]
|
||||
name = "concurrent-queue"
|
||||
version = "2.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "4ca0197aee26d1ae37445ee532fefce43251d24cc7c166799f4d46817f1d3973"
|
||||
dependencies = [
|
||||
"crossbeam-utils",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "core-foundation"
|
||||
version = "0.10.1"
|
||||
@@ -216,10 +289,15 @@ dependencies = [
|
||||
name = "crete"
|
||||
version = "0.1.3"
|
||||
dependencies = [
|
||||
"async-walkdir",
|
||||
"byteorder",
|
||||
"clap",
|
||||
"clap_complete",
|
||||
"futures",
|
||||
"gix-discover",
|
||||
"minijinja",
|
||||
"netascii",
|
||||
"num_enum",
|
||||
"optional_struct",
|
||||
"reqwest",
|
||||
"schemars",
|
||||
@@ -228,9 +306,15 @@ dependencies = [
|
||||
"serde_json",
|
||||
"serde_yaml",
|
||||
"thiserror",
|
||||
"walkdir",
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "crossbeam-utils"
|
||||
version = "0.8.23"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a31eee39dddec8330830986fcd7625edb5a24ec90ea038215273bbc3adb08ac6"
|
||||
|
||||
[[package]]
|
||||
name = "crypto-common"
|
||||
version = "0.1.7"
|
||||
@@ -321,6 +405,26 @@ dependencies = [
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "event-listener"
|
||||
version = "5.4.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5a23add41df1562121a9393cb065eab5146a1242410f23a644851e90cfd669d2"
|
||||
dependencies = [
|
||||
"parking",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "event-listener-strategy"
|
||||
version = "0.5.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "8be9f3dfaaffdae2972880079a491a1a8bb7cbed0b8dd7a347f668b4150a3b93"
|
||||
dependencies = [
|
||||
"event-listener",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "faster-hex"
|
||||
version = "0.10.0"
|
||||
@@ -367,6 +471,21 @@ dependencies = [
|
||||
"percent-encoding",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "9a31d2a3fbaaeb2af2368bbdd904aa8e812d3c04a1ee10d3171f52d556e5d0a3"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-executor",
|
||||
"futures-io",
|
||||
"futures-sink",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-channel"
|
||||
version = "0.3.34"
|
||||
@@ -383,12 +502,36 @@ version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "92d699e522242e69e3003b94ecc1f960f3a5e015aa7c5d7486e65ad01dd94f5e"
|
||||
|
||||
[[package]]
|
||||
name = "futures-executor"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "031b47cf1a3c6cc8bc2fc76cd437f521619387907d469316e7c0bc278f1f5432"
|
||||
dependencies = [
|
||||
"futures-core",
|
||||
"futures-task",
|
||||
"futures-util",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-io"
|
||||
version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "53c0fa8157de1303bfffdaa1cc2a673bfffb60102f76b0ef4441659124373fed"
|
||||
|
||||
[[package]]
|
||||
name = "futures-lite"
|
||||
version = "2.6.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f78e10609fe0e0b3f4157ffab1876319b5b0db102a2c60dc4626306dc46b44ad"
|
||||
dependencies = [
|
||||
"fastrand",
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"parking",
|
||||
"pin-project-lite",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "futures-macro"
|
||||
version = "0.3.34"
|
||||
@@ -418,6 +561,7 @@ version = "0.3.34"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "0d50a92467f8ba5dd6e3ee5d4bd04d73ab2e4e1c44474a0674821dfce14b79bc"
|
||||
dependencies = [
|
||||
"futures-channel",
|
||||
"futures-core",
|
||||
"futures-io",
|
||||
"futures-macro",
|
||||
@@ -1110,6 +1254,34 @@ dependencies = [
|
||||
"tempfile",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "netascii"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "205c45083b6725eb4671392c7c399510ad1498dd0275f1a2e68a70a34ebdf1b1"
|
||||
|
||||
[[package]]
|
||||
name = "num_enum"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5d0bca838442ec211fa11de3a8b0e0e8f3a4522575b5c4c06ed722e005036f26"
|
||||
dependencies = [
|
||||
"num_enum_derive",
|
||||
"rustversion",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "num_enum_derive"
|
||||
version = "0.7.6"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "680998035259dcfcafe653688bf2aa6d3e2dc05e98be6ab46afb089dc84f1df8"
|
||||
dependencies = [
|
||||
"proc-macro-crate",
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "once_cell"
|
||||
version = "1.21.4"
|
||||
@@ -1186,6 +1358,12 @@ dependencies = [
|
||||
"syn 2.0.119",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "parking"
|
||||
version = "2.2.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f38d5652c16fde515bb1ecef450ab0f6a219d619a7274976324d5e377f7dceba"
|
||||
|
||||
[[package]]
|
||||
name = "parking_lot"
|
||||
version = "0.12.5"
|
||||
@@ -1221,6 +1399,17 @@ version = "0.2.17"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a89322df9ebe1c1578d689c92318e070967d1042b512afbe49518723f4e6d5cd"
|
||||
|
||||
[[package]]
|
||||
name = "piper"
|
||||
version = "0.2.5"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c835479a4443ded371d6c535cbfd8d31ad92c5d23ae9770a61bc155e4992a3c1"
|
||||
dependencies = [
|
||||
"atomic-waker",
|
||||
"fastrand",
|
||||
"futures-io",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "pkg-config"
|
||||
version = "0.3.34"
|
||||
@@ -1251,6 +1440,15 @@ dependencies = [
|
||||
"zerovec",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro-crate"
|
||||
version = "3.5.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "e67ba7e9b2b56446f1d419b1d807906278ffa1a658a8a5d8a39dcb1f5a78614f"
|
||||
dependencies = [
|
||||
"toml_edit",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "proc-macro2"
|
||||
version = "1.0.107"
|
||||
@@ -1596,6 +1794,16 @@ version = "2.0.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba"
|
||||
|
||||
[[package]]
|
||||
name = "signal-hook-registry"
|
||||
version = "1.4.8"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "c4db69cba1110affc0e9f7bcd48bbf87b3f4fc7c61fc9155afd4c469eb3d6c1b"
|
||||
dependencies = [
|
||||
"errno",
|
||||
"libc",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "slab"
|
||||
version = "0.4.12"
|
||||
@@ -1739,11 +1947,25 @@ dependencies = [
|
||||
"bytes",
|
||||
"libc",
|
||||
"mio",
|
||||
"parking_lot",
|
||||
"pin-project-lite",
|
||||
"signal-hook-registry",
|
||||
"socket2",
|
||||
"tokio-macros",
|
||||
"windows-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-macros"
|
||||
version = "2.7.2"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "78773a2a397f451582ce068015985c33193cf6dea8b74d2a639fe457b2f07b0e"
|
||||
dependencies = [
|
||||
"proc-macro2",
|
||||
"quote",
|
||||
"syn 3.0.4",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tokio-native-tls"
|
||||
version = "0.3.1"
|
||||
@@ -1767,6 +1989,36 @@ dependencies = [
|
||||
"tokio",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_datetime"
|
||||
version = "1.1.1+spec-1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3165f65f62e28e0115a00b2ebdd37eb6f3b641855f9d636d3cd4103767159ad7"
|
||||
dependencies = [
|
||||
"serde_core",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_edit"
|
||||
version = "0.25.15+spec-1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1340ea94a5856333492c9064b02c778b191dd2c853778d9609debdcdfea3a614"
|
||||
dependencies = [
|
||||
"indexmap",
|
||||
"toml_datetime",
|
||||
"toml_parser",
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "toml_parser"
|
||||
version = "1.1.3+spec-1.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "1d38ac1cf9b95face32296c0a3ede1fdc270627c9d9c02a7274dd6d960dc4d56"
|
||||
dependencies = [
|
||||
"winnow",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "tower"
|
||||
version = "0.5.3"
|
||||
@@ -2027,6 +2279,15 @@ dependencies = [
|
||||
"windows-link",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winnow"
|
||||
version = "1.0.4"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "23b97319f7b8343df12cc98938e5c3eb436064524c8d2b4e30a1d3a36eecdf81"
|
||||
dependencies = [
|
||||
"memchr",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "writeable"
|
||||
version = "0.6.4"
|
||||
|
||||
+6
-1
@@ -14,10 +14,15 @@ schemars = { version = "1.2.1", features = ["semver1"] }
|
||||
serde_json = "1.0.149"
|
||||
|
||||
[dependencies]
|
||||
async-walkdir = "2.1.0"
|
||||
byteorder = "1.5.0"
|
||||
clap = { version = "4.6.0", features = ["derive"] }
|
||||
clap_complete = "4.6.0"
|
||||
futures = "0.3.34"
|
||||
gix-discover = { version = "0.55.0", features = ["sha1"] }
|
||||
minijinja = { version = "2.18.0", features = ["json", "loader"] }
|
||||
netascii = "0.1.0"
|
||||
num_enum = "0.7.6"
|
||||
optional_struct = "0.5.2"
|
||||
reqwest = { version = "0.13.2", default-features = false, features = [
|
||||
"blocking",
|
||||
@@ -30,7 +35,7 @@ serde = { version = "1.0.228", features = ["derive"] }
|
||||
serde_json = { workspace = true }
|
||||
serde_yaml = "0.9.34"
|
||||
thiserror = "2.0.18"
|
||||
walkdir = "2.5.0"
|
||||
tokio = { version = "1.52.1", features = ["full"] }
|
||||
|
||||
[lib]
|
||||
test = false
|
||||
|
||||
+36
-16
@@ -1,13 +1,16 @@
|
||||
use std::net::Ipv4Addr;
|
||||
use std::path::PathBuf;
|
||||
use std::process::Command;
|
||||
use std::str::FromStr;
|
||||
|
||||
use async_walkdir::WalkDir;
|
||||
use futures::StreamExt;
|
||||
use futures::future::join_all;
|
||||
use futures::stream::FuturesUnordered;
|
||||
use minijinja::Environment;
|
||||
use optional_struct::{Applicable, optional_struct};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use walkdir::WalkDir;
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::environment::PathEnvironment;
|
||||
use crate::node::{Node, OptionalNodeDeserialize};
|
||||
@@ -68,7 +71,7 @@ struct ClusterDeserialize {
|
||||
base: Base,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[derive(Debug, Serialize, Clone)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Cluster {
|
||||
pub(crate) name: String,
|
||||
@@ -80,19 +83,23 @@ pub struct Cluster {
|
||||
}
|
||||
|
||||
impl Cluster {
|
||||
pub(crate) fn get(cluster_name: &str, env: &Environment) -> Self {
|
||||
pub(crate) async fn get(cluster_name: &str, env: &Environment<'static>) -> Self {
|
||||
let path = get_talos_path().join("clusters").join("default.yaml");
|
||||
|
||||
let default: OptionalClusterDeserialize = if path.exists() {
|
||||
let content = std::fs::read_to_string(path).unwrap();
|
||||
serde_yaml::from_str(&content).unwrap()
|
||||
let content = tokio::fs::read_to_string(path).await.unwrap();
|
||||
|
||||
// The deserialization makes a blocking web request
|
||||
tokio::task::spawn_blocking(move || serde_yaml::from_str(&content).unwrap())
|
||||
.await
|
||||
.unwrap()
|
||||
} else {
|
||||
Default::default()
|
||||
};
|
||||
|
||||
let mut path = get_talos_path().join("clusters").join(cluster_name);
|
||||
path.add_extension("yaml");
|
||||
let content = std::fs::read_to_string(path).unwrap();
|
||||
let content = tokio::fs::read_to_string(path).await.unwrap();
|
||||
|
||||
let mut cluster: OptionalClusterDeserialize = serde_yaml::from_str(&content).unwrap();
|
||||
|
||||
@@ -122,20 +129,23 @@ impl Cluster {
|
||||
|
||||
let nodes = nodes
|
||||
.iter()
|
||||
.map(|name| Node::get(name, env, &cluster, &default, &base))
|
||||
.collect();
|
||||
.map(|name| async { Node::get(name, env, &cluster, &default, &base).await });
|
||||
|
||||
let nodes = FuturesUnordered::from_iter(nodes).collect().await;
|
||||
|
||||
Self { nodes, ..cluster }
|
||||
}
|
||||
|
||||
pub fn get_all() -> Vec<Cluster> {
|
||||
pub async fn get_all() -> Vec<Cluster> {
|
||||
let env = PathEnvironment::new_patches();
|
||||
|
||||
WalkDir::new(get_talos_path().join("clusters"))
|
||||
.into_iter()
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.metadata().unwrap().is_file())
|
||||
.filter_map(|e| {
|
||||
let names: Vec<_> = WalkDir::new(get_talos_path().join("clusters"))
|
||||
.filter_map(|e| async { e.ok() })
|
||||
.filter(|e| {
|
||||
let e = e.clone();
|
||||
async move { e.metadata().await.unwrap().is_file() }
|
||||
})
|
||||
.filter_map(|e| async move {
|
||||
let mut path = PathBuf::from_str(&e.file_name().to_string_lossy()).unwrap();
|
||||
path.set_extension("");
|
||||
|
||||
@@ -145,9 +155,19 @@ impl Cluster {
|
||||
return None;
|
||||
}
|
||||
|
||||
Some(Cluster::get(&name, &env))
|
||||
Some(name)
|
||||
})
|
||||
.collect()
|
||||
.await;
|
||||
|
||||
// Get the clusters concurrently
|
||||
join_all(names.iter().map(|name| {
|
||||
let env = env.clone();
|
||||
let name = name.clone();
|
||||
|
||||
async move { Cluster::get(&name, &env).await }
|
||||
}))
|
||||
.await
|
||||
}
|
||||
|
||||
pub fn nodes(&self) -> &[Node] {
|
||||
|
||||
+20
-23
@@ -2,15 +2,17 @@ use std::net::Ipv4Addr;
|
||||
use std::ops::Deref;
|
||||
use std::path::{Path, PathBuf};
|
||||
|
||||
use async_walkdir::WalkDir;
|
||||
use futures::{Stream, StreamExt};
|
||||
use minijinja::value::ViaDeserialize;
|
||||
use minijinja::{AutoEscape, Environment, path_loader};
|
||||
use walkdir::WalkDir;
|
||||
|
||||
use crate::{get_configs_path, get_talos_path};
|
||||
|
||||
/// Transparent wrapper around minijinja::Environment that loads templates from a path and
|
||||
/// configures better defaults. It also implements IntoIter, making it possible to iterate over all
|
||||
/// the templates.
|
||||
#[derive(Clone)]
|
||||
pub struct PathEnvironment<'a> {
|
||||
env: Environment<'a>,
|
||||
path: PathBuf,
|
||||
@@ -76,6 +78,23 @@ impl<'a> PathEnvironment<'a> {
|
||||
pub(crate) fn new_patches() -> Self {
|
||||
Self::new(&get_talos_path().join("patches"))
|
||||
}
|
||||
|
||||
pub fn files(&self) -> impl Stream<Item = String> {
|
||||
// Find all templates in path
|
||||
WalkDir::new(&self.path)
|
||||
.filter_map(|e| async { e.ok() })
|
||||
.filter(|e| {
|
||||
let e = e.clone();
|
||||
async move { e.metadata().await.unwrap().is_file() }
|
||||
})
|
||||
.map(move |e| {
|
||||
e.path()
|
||||
.strip_prefix(&self.path)
|
||||
.expect("All paths should have prefix")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
// Make PathEnvironment act like a normal Environment transparently
|
||||
@@ -86,25 +105,3 @@ impl<'a> Deref for PathEnvironment<'a> {
|
||||
&self.env
|
||||
}
|
||||
}
|
||||
|
||||
// Iterate over all the files in the path
|
||||
impl<'a, 'b> IntoIterator for &'b PathEnvironment<'a> {
|
||||
type Item = String;
|
||||
|
||||
type IntoIter = impl Iterator<Item = Self::Item>;
|
||||
|
||||
fn into_iter(self) -> Self::IntoIter {
|
||||
// Find all templates in path
|
||||
WalkDir::new(&self.path)
|
||||
.into_iter()
|
||||
.filter_map(|e| e.ok())
|
||||
.filter(|e| e.metadata().unwrap().is_file())
|
||||
.map(|e| {
|
||||
e.path()
|
||||
.strip_prefix(&self.path)
|
||||
.expect("All paths should have prefix")
|
||||
.to_string_lossy()
|
||||
.to_string()
|
||||
})
|
||||
}
|
||||
}
|
||||
|
||||
+2
-1
@@ -1,4 +1,4 @@
|
||||
#![feature(impl_trait_in_assoc_type, path_absolute_method)]
|
||||
#![feature(path_absolute_method)]
|
||||
use std::path::{Path, PathBuf};
|
||||
use std::sync::OnceLock;
|
||||
|
||||
@@ -8,6 +8,7 @@ pub mod node;
|
||||
pub mod patch;
|
||||
pub mod schematic;
|
||||
pub mod secret;
|
||||
pub mod tftp;
|
||||
|
||||
pub(crate) static REPO_PATH: OnceLock<PathBuf> = OnceLock::new();
|
||||
|
||||
|
||||
+40
-23
@@ -1,14 +1,17 @@
|
||||
mod cli;
|
||||
|
||||
use std::process::Command;
|
||||
use std::pin::pin;
|
||||
|
||||
use clap::{CommandFactory, Parser};
|
||||
use clap_complete::{Shell, generate as generate_complete};
|
||||
use crete::cluster::Cluster;
|
||||
use crete::environment::PathEnvironment;
|
||||
use crete::{get_configs_path, get_repo_path, set_repo_path};
|
||||
use futures::StreamExt;
|
||||
use futures::future::join_all;
|
||||
use minijinja::context;
|
||||
use thiserror::Error;
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::cli::{Cli, Commands, GlobalOpts};
|
||||
|
||||
@@ -18,8 +21,8 @@ enum Error {
|
||||
NoClustersFound,
|
||||
}
|
||||
|
||||
fn run_command(mut command: Command) {
|
||||
let output = command.output().unwrap();
|
||||
async fn run_command(mut command: Command) {
|
||||
let output = command.output().await.unwrap();
|
||||
if !output.stdout.is_empty() {
|
||||
println!("{}", String::from_utf8_lossy(&output.stdout).trim_end());
|
||||
}
|
||||
@@ -28,24 +31,25 @@ fn run_command(mut command: Command) {
|
||||
}
|
||||
}
|
||||
|
||||
fn generate(opts: &GlobalOpts) -> Result<(), Error> {
|
||||
async fn generate(opts: &GlobalOpts) -> Result<(), Error> {
|
||||
set_repo_path(&opts.repo);
|
||||
|
||||
let clusters = Cluster::get_all();
|
||||
let clusters = Cluster::get_all().await;
|
||||
|
||||
if clusters.is_empty() {
|
||||
return Err(Error::NoClustersFound);
|
||||
}
|
||||
|
||||
let path = opts.repo.join("rendered");
|
||||
if path.exists() {
|
||||
std::fs::remove_dir_all(&path).unwrap();
|
||||
}
|
||||
std::fs::create_dir(&path).unwrap();
|
||||
// We don't really care if this fails, as that likely mean that the directory did not exist
|
||||
tokio::fs::remove_dir_all(&path).await.ok();
|
||||
tokio::fs::create_dir_all(&path).await.unwrap();
|
||||
|
||||
// Render templates
|
||||
let template_env = PathEnvironment::new(&get_repo_path().join("templates"));
|
||||
for template_name in &template_env {
|
||||
|
||||
let mut stream = pin!(template_env.files());
|
||||
while let Some(template_name) = stream.next().await {
|
||||
let template = template_env.get_template(&template_name).unwrap();
|
||||
|
||||
let content = template
|
||||
@@ -55,35 +59,48 @@ fn generate(opts: &GlobalOpts) -> Result<(), Error> {
|
||||
})
|
||||
.unwrap();
|
||||
|
||||
std::fs::write(path.join(template_name), content).unwrap();
|
||||
tokio::fs::write(path.join(template_name), content)
|
||||
.await
|
||||
.unwrap();
|
||||
}
|
||||
|
||||
// Remove existing config files and create output directory
|
||||
let path = get_configs_path();
|
||||
if path.exists() {
|
||||
std::fs::remove_dir_all(&path).unwrap();
|
||||
}
|
||||
std::fs::create_dir(&path).unwrap();
|
||||
tokio::fs::remove_dir_all(&path).await.ok();
|
||||
tokio::fs::create_dir_all(&path).await.unwrap();
|
||||
|
||||
// Generate config files
|
||||
for cluster in clusters {
|
||||
for node in cluster.nodes() {
|
||||
run_command(node.talosctl_gen_config_command(&cluster));
|
||||
join_all(clusters.iter().map(|cluster| async move {
|
||||
join_all(cluster.nodes().iter().map(|node| {
|
||||
let cluster = cluster.clone();
|
||||
async move {
|
||||
run_command(node.talosctl_gen_config_command(&cluster)).await;
|
||||
}
|
||||
}))
|
||||
.await;
|
||||
|
||||
run_command(cluster.talosctl_gen_config_command());
|
||||
run_command(cluster.talosctl_add_endpoint_command());
|
||||
run_command(cluster.talosctl_merge_command());
|
||||
// Run these commands in parallel with the the node commands
|
||||
// NOTE: These do need to happen in the correct order
|
||||
run_command(cluster.talosctl_gen_config_command()).await;
|
||||
run_command(cluster.talosctl_add_endpoint_command()).await;
|
||||
}))
|
||||
.await;
|
||||
|
||||
// Merge all the cluster talosconfig files into a single file
|
||||
// Since this all modifies the same output file this has to happen sequentially
|
||||
for cluster in clusters {
|
||||
run_command(cluster.talosctl_merge_command()).await;
|
||||
}
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn main() -> Result<(), Error> {
|
||||
#[tokio::main]
|
||||
async fn main() -> Result<(), Error> {
|
||||
let cli = Cli::parse();
|
||||
|
||||
match cli.command {
|
||||
Commands::Generate => generate(&cli.global_opts)?,
|
||||
Commands::Generate => generate(&cli.global_opts).await?,
|
||||
Commands::ShellCompletions => generate_complete(
|
||||
Shell::from_env().unwrap_or(Shell::Bash),
|
||||
&mut Cli::command(),
|
||||
|
||||
+6
-6
@@ -1,10 +1,10 @@
|
||||
use std::net::Ipv4Addr;
|
||||
use std::process::Command;
|
||||
|
||||
use minijinja::Environment;
|
||||
use optional_struct::{Applicable, optional_struct};
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use tokio::process::Command;
|
||||
|
||||
use crate::cluster::{Base, Cluster};
|
||||
use crate::patch::{OptionalPatches, OptionalPatchesString, Patches};
|
||||
@@ -90,7 +90,7 @@ pub(crate) struct NodeDeserialize {
|
||||
sops: Secret,
|
||||
}
|
||||
|
||||
#[derive(Debug, Serialize)]
|
||||
#[derive(Debug, Serialize, Clone)]
|
||||
#[serde(rename_all = "camelCase")]
|
||||
pub struct Node {
|
||||
hostname: String,
|
||||
@@ -106,9 +106,9 @@ pub struct Node {
|
||||
}
|
||||
|
||||
impl Node {
|
||||
pub(crate) fn get(
|
||||
pub(crate) async fn get(
|
||||
node_name: &str,
|
||||
env: &Environment,
|
||||
env: &Environment<'static>,
|
||||
cluster: &Cluster,
|
||||
default: &OptionalNodeDeserialize,
|
||||
base: &Base,
|
||||
@@ -121,7 +121,7 @@ impl Node {
|
||||
.to_string();
|
||||
|
||||
path.add_extension("yaml");
|
||||
let content = std::fs::read_to_string(path).unwrap();
|
||||
let content = tokio::fs::read_to_string(path).await.unwrap();
|
||||
|
||||
let node: OptionalNodeDeserialize = serde_yaml::from_str(&content).unwrap();
|
||||
|
||||
@@ -172,7 +172,7 @@ impl Node {
|
||||
};
|
||||
|
||||
// Render patches
|
||||
let patches = patches.render(env, cluster, &node);
|
||||
let patches = patches.render(env, cluster, &node).await;
|
||||
|
||||
Node { patches, ..node }
|
||||
}
|
||||
|
||||
+16
-7
@@ -16,15 +16,21 @@ pub(crate) struct Patches<T> {
|
||||
|
||||
pub(crate) type OptionalPatchesString = OptionalPatches<String>;
|
||||
|
||||
fn render(
|
||||
async fn render(
|
||||
patches: Vec<String>,
|
||||
env: &Environment,
|
||||
env: &Environment<'static>,
|
||||
cluster: &Cluster,
|
||||
node: &Node,
|
||||
) -> Vec<serde_yaml::Value> {
|
||||
let env = env.clone();
|
||||
let cluster = cluster.clone();
|
||||
let node = node.clone();
|
||||
|
||||
// Minijinja will still have blocking I/O internally
|
||||
tokio::task::spawn_blocking(move || {
|
||||
patches
|
||||
.into_iter()
|
||||
.map(|name| {
|
||||
.map(move |name| {
|
||||
let content = env
|
||||
.get_template(&name)
|
||||
.unwrap()
|
||||
@@ -37,6 +43,9 @@ fn render(
|
||||
serde_yaml::from_str(&content).unwrap()
|
||||
})
|
||||
.collect()
|
||||
})
|
||||
.await
|
||||
.unwrap()
|
||||
}
|
||||
|
||||
impl Patches<String> {
|
||||
@@ -50,15 +59,15 @@ impl Patches<String> {
|
||||
}
|
||||
}
|
||||
|
||||
pub(crate) fn render(
|
||||
pub(crate) async fn render(
|
||||
self,
|
||||
env: &Environment,
|
||||
env: &Environment<'static>,
|
||||
cluster: &Cluster,
|
||||
node: &Node,
|
||||
) -> Patches<serde_yaml::Value> {
|
||||
Patches {
|
||||
all: render(self.all.clone(), env, cluster, node),
|
||||
control_plane: render(self.control_plane.clone(), env, cluster, node),
|
||||
all: render(self.all.clone(), env, cluster, node).await,
|
||||
control_plane: render(self.control_plane.clone(), env, cluster, node).await,
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
+19
-8
@@ -1,5 +1,6 @@
|
||||
use schemars::JsonSchema;
|
||||
use serde::{Deserialize, Deserializer, Serialize};
|
||||
use tokio::runtime::Handle;
|
||||
|
||||
use crate::get_talos_path;
|
||||
|
||||
@@ -11,7 +12,7 @@ where
|
||||
let path = get_talos_path().join("schematics").join(name);
|
||||
let content = std::fs::read_to_string(path).unwrap().trim().to_owned();
|
||||
|
||||
let client = reqwest::blocking::ClientBuilder::new()
|
||||
let client = reqwest::ClientBuilder::new()
|
||||
.user_agent(format!(
|
||||
"{}/{}",
|
||||
std::env!("CARGO_PKG_NAME"),
|
||||
@@ -20,20 +21,30 @@ where
|
||||
.build()
|
||||
.unwrap();
|
||||
|
||||
let res = client
|
||||
.post("https://factory.talos.dev/schematics")
|
||||
.body(content)
|
||||
.send()
|
||||
.map_err(serde::de::Error::custom)?;
|
||||
|
||||
#[derive(Debug, Deserialize)]
|
||||
struct Response {
|
||||
id: String,
|
||||
}
|
||||
|
||||
let response: Response = serde_json::from_str(&res.text().map_err(serde::de::Error::custom)?)
|
||||
// We can't use async directly inside a custom deserializer, and using the blocking client also
|
||||
// does not work as it will try to start tokio inside of tokio.
|
||||
// So we have do this to work around it.
|
||||
// NOTE: This requires that the deserialization happens inside spawn_blocking
|
||||
let response = Handle::current().block_on(async {
|
||||
let res = client
|
||||
.post("https://factory.talos.dev/schematics")
|
||||
.body(content)
|
||||
.send()
|
||||
.await
|
||||
.map_err(serde::de::Error::custom)?;
|
||||
|
||||
let response: Response =
|
||||
serde_json::from_str(&res.text().await.map_err(serde::de::Error::custom)?)
|
||||
.map_err(serde::de::Error::custom)?;
|
||||
|
||||
Ok(response)
|
||||
})?;
|
||||
|
||||
Ok(response.id)
|
||||
}
|
||||
|
||||
|
||||
@@ -0,0 +1,20 @@
|
||||
use std::io;
|
||||
use std::string::FromUtf8Error;
|
||||
|
||||
use thiserror::Error;
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum Error {
|
||||
#[error("Invalid op code '{0}'")]
|
||||
InvalidOpCode(u16),
|
||||
#[error("Invalid error code '{0}'")]
|
||||
InvalidErrorCode(u16),
|
||||
#[error("Invalid mode '{0}'")]
|
||||
InvalidMode(String),
|
||||
#[error("String contains invalid characters")]
|
||||
InvalidString(#[from] FromUtf8Error),
|
||||
#[error("Unterminated string")]
|
||||
UnterminatedString,
|
||||
#[error("Io error: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
}
|
||||
@@ -0,0 +1,7 @@
|
||||
mod error;
|
||||
mod packet;
|
||||
mod server;
|
||||
|
||||
pub use error::Error;
|
||||
pub use packet::Packet;
|
||||
pub use server::{Error as ServerError, serve};
|
||||
@@ -0,0 +1,353 @@
|
||||
use std::io::{BufRead, BufReader, BufWriter, Cursor, Read, Write};
|
||||
use std::net::SocketAddr;
|
||||
|
||||
use byteorder::{NetworkEndian, ReadBytesExt, WriteBytesExt};
|
||||
use num_enum::{IntoPrimitive, TryFromPrimitive};
|
||||
use tokio::net::UdpSocket;
|
||||
|
||||
use crate::tftp::Error;
|
||||
|
||||
// 2 bytes op code + 2 bytes block number + max 512 bytes data
|
||||
pub(crate) const PACKET_SIZE: usize = 2 + 2 + 512;
|
||||
|
||||
fn decode_null_terminated<R>(reader: &mut BufReader<R>) -> Result<String, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let mut text = Vec::new();
|
||||
reader.read_until(0, &mut text)?;
|
||||
|
||||
// Make sure we ended with a null terminator
|
||||
match text.pop() {
|
||||
Some(0) => {}
|
||||
None | Some(_) => return Err(Error::UnterminatedString),
|
||||
}
|
||||
|
||||
Ok(String::from_utf8(text)?)
|
||||
}
|
||||
|
||||
fn encode_null_terminated<W>(writer: &mut BufWriter<W>, text: impl AsRef<str>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
let text = text.as_ref().as_bytes();
|
||||
writer.write_all(text)?;
|
||||
|
||||
// Add null terminator
|
||||
writer.write_u8(0)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
|
||||
pub enum Mode {
|
||||
NetAscii,
|
||||
Octet,
|
||||
Mail,
|
||||
}
|
||||
|
||||
impl Mode {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let mut mode = decode_null_terminated(reader)?;
|
||||
mode.make_ascii_lowercase();
|
||||
let mode = mode;
|
||||
|
||||
Ok(match mode.as_str() {
|
||||
"netascii" => Self::NetAscii,
|
||||
"octet" => Self::Octet,
|
||||
"mail" => Self::Mail,
|
||||
_ => return Err(Error::InvalidMode(mode)),
|
||||
})
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
let mode = match &self {
|
||||
Mode::NetAscii => "netascii",
|
||||
Mode::Octet => "octet",
|
||||
Mode::Mail => "mail",
|
||||
};
|
||||
|
||||
encode_null_terminated(writer, mode)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, PartialEq, Eq)]
|
||||
pub struct Request {
|
||||
filename: String,
|
||||
mode: Mode,
|
||||
}
|
||||
|
||||
impl Request {
|
||||
pub fn filename(&self) -> &str {
|
||||
&self.filename
|
||||
}
|
||||
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let filename = decode_null_terminated(reader)?;
|
||||
let mode = Mode::decode(reader)?;
|
||||
|
||||
Ok(Self { filename, mode })
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
encode_null_terminated(writer, &self.filename)?;
|
||||
self.mode.encode(writer)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub fn mode(&self) -> Mode {
|
||||
self.mode
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Data {
|
||||
block: u16,
|
||||
data: Vec<u8>,
|
||||
}
|
||||
|
||||
impl Data {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let block = reader.read_u16::<NetworkEndian>()?;
|
||||
|
||||
let mut data = Vec::new();
|
||||
reader.read_to_end(&mut data)?;
|
||||
|
||||
Ok(Self { data, block })
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
writer.write_u16::<NetworkEndian>(self.block)?;
|
||||
writer.write_all(&self.data)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct Ack(u16);
|
||||
|
||||
impl Ack {
|
||||
pub fn block(&self) -> u16 {
|
||||
self.0
|
||||
}
|
||||
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let number = reader.read_u16::<NetworkEndian>()?;
|
||||
|
||||
Ok(Self(number))
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
writer.write_u16::<NetworkEndian>(self.0)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug, TryFromPrimitive, IntoPrimitive, Clone, Copy)]
|
||||
#[repr(u16)]
|
||||
pub enum ErrorCode {
|
||||
NotDefined,
|
||||
FileNotFound,
|
||||
AccessViolation,
|
||||
DiskFull,
|
||||
Illegal,
|
||||
UnknownTransferId,
|
||||
FileAlreadyExists,
|
||||
NoSuchUser,
|
||||
}
|
||||
|
||||
impl ErrorCode {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let code = reader.read_u16::<NetworkEndian>()?;
|
||||
Self::try_from_primitive(code).map_err(|err| Error::InvalidErrorCode(err.number))
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
let code: u16 = (*self).into();
|
||||
|
||||
writer.write_u16::<NetworkEndian>(code)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub struct ErrorMessage {
|
||||
code: ErrorCode,
|
||||
message: String,
|
||||
}
|
||||
|
||||
impl ErrorMessage {
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let code = ErrorCode::decode(reader)?;
|
||||
let message = decode_null_terminated(reader)?;
|
||||
|
||||
Ok(Self { code, message })
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
self.code.encode(writer)?;
|
||||
|
||||
encode_null_terminated(writer, &self.message)?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
#[derive(Debug)]
|
||||
pub enum Packet {
|
||||
Read(Request),
|
||||
Write(Request),
|
||||
Data(Data),
|
||||
Ack(Ack),
|
||||
Error(ErrorMessage),
|
||||
}
|
||||
|
||||
impl Packet {
|
||||
pub fn data(block: u16, data: &[u8]) -> Self {
|
||||
Self::Data(Data {
|
||||
block,
|
||||
data: data.into(),
|
||||
})
|
||||
}
|
||||
|
||||
pub fn ack(block: u16) -> Self {
|
||||
Self::Ack(Ack(block))
|
||||
}
|
||||
|
||||
pub fn error(code: ErrorCode, message: impl Into<String>) -> Self {
|
||||
Self::Error(ErrorMessage {
|
||||
code,
|
||||
message: message.into(),
|
||||
})
|
||||
}
|
||||
|
||||
fn decode<R>(reader: &mut BufReader<R>) -> Result<Self, Error>
|
||||
where
|
||||
R: Read + ?Sized,
|
||||
{
|
||||
let op_code = reader.read_u16::<NetworkEndian>()?;
|
||||
|
||||
Ok(match op_code {
|
||||
1 => Self::Read(Request::decode(reader)?),
|
||||
2 => Self::Write(Request::decode(reader)?),
|
||||
3 => Self::Data(Data::decode(reader)?),
|
||||
4 => Self::Ack(Ack::decode(reader)?),
|
||||
5 => Self::Error(ErrorMessage::decode(reader)?),
|
||||
_ => return Err(Error::InvalidOpCode(op_code)),
|
||||
})
|
||||
}
|
||||
|
||||
fn encode<W>(&self, writer: &mut BufWriter<W>) -> Result<(), Error>
|
||||
where
|
||||
W: Write + ?Sized,
|
||||
{
|
||||
match &self {
|
||||
Self::Read(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(1)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Write(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(2)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Data(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(3)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Ack(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(4)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
Self::Error(inner) => {
|
||||
writer.write_u16::<NetworkEndian>(5)?;
|
||||
inner.encode(writer)?;
|
||||
}
|
||||
};
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn recv_from(socket: &UdpSocket) -> Result<(Self, SocketAddr), Error> {
|
||||
let mut buf = Vec::with_capacity(PACKET_SIZE);
|
||||
let (_, src) = socket.recv_buf_from(&mut buf).await?;
|
||||
|
||||
let packet = buf.try_into()?;
|
||||
|
||||
Ok((packet, src))
|
||||
}
|
||||
|
||||
pub async fn send_to(&self, socket: &UdpSocket, src: &SocketAddr) -> Result<(), Error> {
|
||||
let buf: Vec<_> = self.try_into()?;
|
||||
socket.send_to(&buf, src).await?;
|
||||
|
||||
Ok(())
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<&Packet> for Vec<u8> {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(packet: &Packet) -> Result<Self, Self::Error> {
|
||||
let mut buf = Vec::with_capacity(PACKET_SIZE);
|
||||
let mut writer = BufWriter::new(Cursor::new(&mut buf));
|
||||
packet.encode(&mut writer)?;
|
||||
|
||||
// Normally the writer would get dropped after we return, but since it has a mutable
|
||||
// borrow in buf we have to drop it first.
|
||||
drop(writer);
|
||||
|
||||
Ok(buf)
|
||||
}
|
||||
}
|
||||
|
||||
impl TryFrom<Vec<u8>> for Packet {
|
||||
type Error = Error;
|
||||
|
||||
fn try_from(buf: Vec<u8>) -> Result<Self, <Packet as TryFrom<Vec<u8>>>::Error> {
|
||||
let mut reader = BufReader::new(Cursor::new(buf));
|
||||
|
||||
Packet::decode(&mut reader)
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,162 @@
|
||||
use std::io;
|
||||
use std::net::SocketAddr;
|
||||
use std::time::Duration;
|
||||
|
||||
use thiserror::Error;
|
||||
use tokio::net::UdpSocket;
|
||||
|
||||
use crate::tftp::packet::{ErrorCode, Mode};
|
||||
use crate::tftp::{self, Packet};
|
||||
|
||||
#[derive(Debug, Error)]
|
||||
pub enum Error {
|
||||
#[error("{0}")]
|
||||
Tftp(#[from] tftp::Error),
|
||||
#[error("Io error: {0}")]
|
||||
Io(#[from] io::Error),
|
||||
#[error("Unsupported operation: {0:#?}")]
|
||||
UnsupportedOperation(Packet),
|
||||
#[error("File not found: {0}")]
|
||||
FileNotFound(String),
|
||||
#[error("Unsupported mode: {0:?}")]
|
||||
UnsupportedMode(Mode),
|
||||
#[error("Source changed: {0}")]
|
||||
SourceChanged(SocketAddr),
|
||||
#[error("UnexpectedPacket: {0:#?}")]
|
||||
UnexpectedPacket(Packet),
|
||||
#[error("Timeout")]
|
||||
Timeout,
|
||||
}
|
||||
|
||||
pub async fn handle_connection(
|
||||
src: SocketAddr,
|
||||
packet: Packet,
|
||||
callback: impl (Fn(&str) -> Option<Vec<u8>>) + Send + Clone + 'static,
|
||||
) -> Result<(), Error> {
|
||||
let socket = UdpSocket::bind("127.0.0.1:0").await?;
|
||||
println!(
|
||||
"{src}: Connecting from {}",
|
||||
socket.local_addr().expect("Socket is bound")
|
||||
);
|
||||
|
||||
let Packet::Read(request) = packet else {
|
||||
Packet::error(ErrorCode::NotDefined, "Operation not supported")
|
||||
.send_to(&socket, &src)
|
||||
.await?;
|
||||
|
||||
return Err(Error::UnsupportedOperation(packet));
|
||||
};
|
||||
|
||||
let filename = request.filename();
|
||||
println!("{src}: {filename} [{:?}]", request.mode());
|
||||
|
||||
let Some(data) = callback(filename) else {
|
||||
Packet::error(ErrorCode::FileNotFound, "Unknown file")
|
||||
.send_to(&socket, &src)
|
||||
.await?;
|
||||
|
||||
return Err(Error::FileNotFound(filename.into()));
|
||||
};
|
||||
|
||||
let mode = request.mode();
|
||||
let data = match mode {
|
||||
// In case of netascii mode we need to re-encode to escape certain sequences
|
||||
Mode::NetAscii => netascii::Netascii::from_bytes(data).collect(),
|
||||
Mode::Octet => data,
|
||||
_ => {
|
||||
Packet::error(
|
||||
ErrorCode::NotDefined,
|
||||
format!("Mode {mode:?} is not supported"),
|
||||
)
|
||||
.send_to(&socket, &src)
|
||||
.await?;
|
||||
|
||||
return Err(Error::UnsupportedMode(mode));
|
||||
}
|
||||
};
|
||||
|
||||
let mut chunks = data.chunks(512).enumerate();
|
||||
let mut chunk = chunks.next();
|
||||
|
||||
let mut timeout = 0;
|
||||
|
||||
loop {
|
||||
// Send the first byte
|
||||
if let Some((index, chunk)) = chunk {
|
||||
Packet::data((index + 1) as u16, chunk)
|
||||
.send_to(&socket, &src)
|
||||
.await?;
|
||||
}
|
||||
|
||||
tokio::select! {
|
||||
recv = Packet::recv_from(&socket) => {
|
||||
timeout = 0;
|
||||
let (packet, new_src) = recv?;
|
||||
|
||||
// I think this is might technically be possible and fine, but I don't
|
||||
// really know how to test it properly so we don't support it.
|
||||
if new_src != src {
|
||||
Packet::error(ErrorCode::Illegal, "Source changed")
|
||||
.send_to(&socket, &src)
|
||||
.await?;
|
||||
|
||||
return Err(Error::SourceChanged(new_src));
|
||||
}
|
||||
|
||||
match packet {
|
||||
Packet::Ack(ack) => {
|
||||
// If we receive an ack for the current chunk go to the next one
|
||||
if let Some((index, _)) = chunk
|
||||
&& (index + 1) as u16 == ack.block()
|
||||
{
|
||||
chunk = chunks.next();
|
||||
}
|
||||
|
||||
if chunk.is_none() {
|
||||
break;
|
||||
}
|
||||
}
|
||||
_ => {
|
||||
Packet::error(ErrorCode::Illegal, "Unexpected packet")
|
||||
.send_to(&socket, &src)
|
||||
.await?;
|
||||
|
||||
|
||||
return Err(Error::UnexpectedPacket(packet));
|
||||
}
|
||||
}
|
||||
},
|
||||
_ = tokio::time::sleep(Duration::from_secs(1)) => {
|
||||
timeout += 1;
|
||||
|
||||
if timeout == 5 {
|
||||
return Err(Error::Timeout);
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
println!("{src}: Transmission complete");
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
pub async fn serve(
|
||||
callback: impl (Fn(&str) -> Option<Vec<u8>>) + Send + Clone + 'static,
|
||||
) -> Result<(), Error> {
|
||||
let socket = UdpSocket::bind("127.0.0.1:69").await?;
|
||||
|
||||
println!("Server started");
|
||||
loop {
|
||||
let (packet, src) = Packet::recv_from(&socket).await?;
|
||||
|
||||
let callback = callback.clone();
|
||||
tokio::spawn({
|
||||
async move {
|
||||
if let Err(err) = handle_connection(src, packet, callback).await {
|
||||
println!("{src}: {err}")
|
||||
}
|
||||
}
|
||||
});
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user