NASP/rpki
84
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

新增slurm功能

Browse Source
This commit is contained in:
xiuting.xu
2026-04-01 16:24:01 +08:00
parent 8d6173f605
commit cd0330e8ae
49 changed files with 2738 additions and 1052 deletions
Download Patch File Download Diff File Expand all files Collapse all files

137
src/bin/rtr_debug_client/main.rs
View File

@@ -7,18 +7,16 @@ use rustls::{ClientConfig as RustlsClientConfig, RootCertStore};
use rustls_pki_types::{CertificateDer, PrivateKeyDer, ServerName};
use tokio::io::{self as tokio_io, AsyncBufReadExt, AsyncRead, AsyncWrite, BufReader, WriteHalf};
use tokio::net::TcpStream;
use tokio::time::{timeout, Duration, Instant};
use tokio::time::{Duration, Instant, timeout};
use tokio_rustls::TlsConnector;
mod wire;
mod pretty;
mod protocol;
mod wire;
use crate::wire::{read_pdu, send_reset_query, send_serial_query};
use crate::pretty::{
parse_end_of_data_info, parse_serial_notify_serial, print_pdu, print_raw_pdu,
};
use crate::pretty::{parse_end_of_data_info, parse_serial_notify_serial, print_pdu, print_raw_pdu};
use crate::protocol::{PduHeader, PduType, QueryMode};
use crate::wire::{read_pdu, send_reset_query, send_serial_query};
const DEFAULT_READ_TIMEOUT_SECS: u64 = 30;
const DEFAULT_POLL_INTERVAL_SECS: u64 = 600;
@@ -38,7 +36,10 @@ async fn main() -> io::Result<()> {
println!("transport: {}", config.transport.describe());
println!("version : {}", config.version);
println!("timeout : {}s", config.read_timeout_secs);
println!("poll : {}s (default before EndOfData refresh is known)", config.default_poll_secs);
println!(
"poll : {}s (default before EndOfData refresh is known)",
config.default_poll_secs
);
println!("keep-after-error: {}", config.keep_after_error);
match &config.mode {
QueryMode::Reset => {
@@ -72,11 +73,7 @@ async fn main() -> io::Result<()> {
}
Err(err) => {
let delay = state.reconnect_delay_secs();
eprintln!(
"connect failed: {}. retry after {}s",
err,
delay
);
eprintln!("connect failed: {}. retry after {}s", err, delay);
tokio::time::sleep(Duration::from_secs(delay)).await;
}
}
@@ -171,10 +168,7 @@ async fn main() -> io::Result<()> {
if reconnect {
let delay = state.reconnect_delay_secs();
state.current_session_id = None;
println!(
"[reconnect] transport disconnected, retry after {}s",
delay
);
println!("[reconnect] transport disconnected, retry after {}s", delay);
tokio::time::sleep(Duration::from_secs(delay)).await;
}
}
@@ -189,8 +183,7 @@ async fn send_resume_query(
(Some(session_id), Some(serial)) => {
println!(
"reconnected, send Serial Query with session_id={}, serial={}",
session_id,
serial
session_id, serial
);
send_serial_query(writer, state.version, session_id, serial).await?;
}
@@ -294,26 +287,20 @@ async fn handle_incoming_pdu(
println!();
println!(
"[notify] received Serial Notify: session_id={}, notify_serial={:?}",
notify_session_id,
notify_serial
notify_session_id, notify_serial
);
match (state.session_id, state.serial, notify_serial) {
(Some(current_session_id), Some(current_serial), Some(_new_serial))
if current_session_id == notify_session_id =>
{
println!(
"received Serial Notify for current session {}, send Serial Query with serial {}",
current_session_id, current_serial
);
send_serial_query(
writer,
state.version,
current_session_id,
current_serial,
)
.await?;
}
if current_session_id == notify_session_id =>
{
println!(
"received Serial Notify for current session {}, send Serial Query with serial {}",
current_session_id, current_serial
);
send_serial_query(writer, state.version, current_session_id, current_serial)
.await?;
}
_ => {
println!(
@@ -366,10 +353,7 @@ async fn handle_incoming_pdu(
Ok(())
}
async fn handle_poll_tick(
writer: &mut ClientWriter,
state: &mut ClientState,
) -> io::Result<()> {
async fn handle_poll_tick(writer: &mut ClientWriter, state: &mut ClientState) -> io::Result<()> {
println!();
println!(
"[auto-poll] timer fired (interval={}s)",
@@ -422,23 +406,21 @@ async fn handle_console_command(
state.schedule_next_poll();
}
["serial"] => {
match (state.session_id, state.serial) {
(Some(session_id), Some(serial)) => {
println!(
"manual command: send Serial Query with current state: session_id={}, serial={}",
session_id, serial
);
send_serial_query(writer, state.version, session_id, serial).await?;
state.schedule_next_poll();
}
_ => {
println!(
"manual command failed: current session_id/serial not available, use `reset` or `serial <session_id> <serial>`"
);
}
["serial"] => match (state.session_id, state.serial) {
(Some(session_id), Some(serial)) => {
println!(
"manual command: send Serial Query with current state: session_id={}, serial={}",
session_id, serial
);
send_serial_query(writer, state.version, session_id, serial).await?;
state.schedule_next_poll();
}
}
_ => {
println!(
"manual command failed: current session_id/serial not available, use `reset` or `serial <session_id> <serial>`"
);
}
},
["serial", session_id, serial] => {
let session_id = match session_id.parse::<u16>() {
@@ -493,7 +475,10 @@ async fn handle_console_command(
"current effective poll interval: {}s",
state.effective_poll_secs()
);
println!("poll interval source : {}", state.poll_interval_source());
println!(
"poll interval source : {}",
state.poll_interval_source()
);
println!("stored refresh hint : {:?}", state.refresh);
println!("default poll interval : {}s", state.default_poll_secs);
println!("last_error_code : {:?}", state.last_error_code);
@@ -626,17 +611,20 @@ impl ClientState {
fn effective_poll_secs(&self) -> u64 {
if self.should_prefer_retry_poll() {
self.retry
.map(|v| v as u64)
.unwrap_or_else(|| self.refresh.map(|v| v as u64).unwrap_or(self.default_poll_secs))
self.retry.map(|v| v as u64).unwrap_or_else(|| {
self.refresh
.map(|v| v as u64)
.unwrap_or(self.default_poll_secs)
})
} else {
self.refresh.map(|v| v as u64).unwrap_or(self.default_poll_secs)
self.refresh
.map(|v| v as u64)
.unwrap_or(self.default_poll_secs)
}
}
fn schedule_next_poll(&mut self) {
self.next_poll_deadline =
Instant::now() + Duration::from_secs(self.effective_poll_secs());
self.next_poll_deadline = Instant::now() + Duration::from_secs(self.effective_poll_secs());
}
fn pause_auto_poll(&mut self) {
@@ -728,7 +716,10 @@ impl Config {
}
"--server-name" => {
let name = args.next().ok_or_else(|| {
io::Error::new(io::ErrorKind::InvalidInput, "--server-name requires a value")
io::Error::new(
io::ErrorKind::InvalidInput,
"--server-name requires a value",
)
})?;
ensure_tls_config(&mut transport)?.server_name = Some(name);
}
@@ -805,10 +796,7 @@ impl Config {
let serial = positional
.next()
.ok_or_else(|| {
io::Error::new(
io::ErrorKind::InvalidInput,
"serial mode requires serial",
)
io::Error::new(io::ErrorKind::InvalidInput, "serial mode requires serial")
})?
.parse::<u32>()
.map_err(|e| {
@@ -949,9 +937,15 @@ async fn connect_tls_stream(addr: &str, tls: &TlsConfig) -> io::Result<DynStream
format!("invalid TLS server name '{}': {}", server_name_str, err),
)
})?;
let tls_stream = connector.connect(server_name, stream).await.map_err(|err| {
io::Error::new(io::ErrorKind::ConnectionAborted, format!("TLS handshake failed: {}", err))
})?;
let tls_stream = connector
.connect(server_name, stream)
.await
.map_err(|err| {
io::Error::new(
io::ErrorKind::ConnectionAborted,
format!("TLS handshake failed: {}", err),
)
})?;
Ok(Box::new(tls_stream))
}
@@ -966,7 +960,10 @@ fn build_tls_connector(tls: &TlsConfig) -> io::Result<TlsConnector> {
if added == 0 {
return Err(io::Error::new(
io::ErrorKind::InvalidInput,
format!("no valid CA certificates found in {}", ca_cert_path.display()),
format!(
"no valid CA certificates found in {}",
ca_cert_path.display()
),
));
}

8
src/bin/rtr_debug_client/pretty.rs
View File

@@ -1,9 +1,8 @@
use std::net::{Ipv4Addr, Ipv6Addr};
use crate::protocol::{
flag_meaning, hex_bytes, PduHeader, PduType, ASPA_FIXED_BODY_LEN,
END_OF_DATA_V0_BODY_LEN, END_OF_DATA_V1_BODY_LEN, IPV4_PREFIX_BODY_LEN,
IPV6_PREFIX_BODY_LEN, ROUTER_KEY_FIXED_BODY_LEN,
ASPA_FIXED_BODY_LEN, END_OF_DATA_V0_BODY_LEN, END_OF_DATA_V1_BODY_LEN, IPV4_PREFIX_BODY_LEN,
IPV6_PREFIX_BODY_LEN, PduHeader, PduType, ROUTER_KEY_FIXED_BODY_LEN, flag_meaning, hex_bytes,
};
pub fn print_pdu(header: &PduHeader, body: &[u8]) {
@@ -143,8 +142,7 @@ fn print_error_report(header: &PduHeader, body: &[u8]) {
return;
}
let encapsulated_len =
u32::from_be_bytes([body[0], body[1], body[2], body[3]]) as usize;
let encapsulated_len = u32::from_be_bytes([body[0], body[1], body[2], body[3]]) as usize;
if body.len() < 4 + encapsulated_len + 4 {
println!("invalid ErrorReport: truncated encapsulated PDU");

11
src/bin/rtr_debug_client/wire.rs
View File

@@ -2,9 +2,7 @@ use std::io;
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
use crate::protocol::{
PduHeader, PduType, RawPdu, HEADER_LEN, MAX_PDU_LEN, SERIAL_QUERY_LEN,
};
use crate::protocol::{HEADER_LEN, MAX_PDU_LEN, PduHeader, PduType, RawPdu, SERIAL_QUERY_LEN};
pub async fn send_reset_query<S>(stream: &mut S, version: u8) -> io::Result<()>
where
@@ -56,10 +54,7 @@ where
if header.length < HEADER_LEN as u32 {
return Err(io::Error::new(
io::ErrorKind::InvalidData,
format!(
"invalid PDU length {} < {}",
header.length, HEADER_LEN
),
format!("invalid PDU length {} < {}", header.length, HEADER_LEN),
));
}
@@ -78,4 +73,4 @@ where
stream.read_exact(&mut body).await?;
Ok(RawPdu { header, body })
}
}

168
src/bin/slurm_apply_client.rs Normal file
View File

@@ -0,0 +1,168 @@
use std::env;
use std::fs::File;
use std::path::PathBuf;
use anyhow::{Context, Result, anyhow};
use serde::Serialize;
use rpki::rtr::payload::Payload;
use rpki::slurm::file::SlurmFile;
use rpki::source::ccr::{
find_latest_ccr_file, load_ccr_snapshot_from_file, snapshot_to_payloads_with_options,
};
#[derive(Debug)]
struct Cli {
ccr_path: PathBuf,
slurm_path: PathBuf,
strict_ccr: bool,
dump_payloads: bool,
}
#[derive(Debug, Serialize)]
struct Output {
ccr_path: String,
slurm_path: String,
produced_at: Option<String>,
slurm_version: u32,
input_payload_count: usize,
input_vrp_count: usize,
input_vap_count: usize,
output_payload_count: usize,
output_vrp_count: usize,
output_vap_count: usize,
invalid_vrps: Vec<String>,
invalid_vaps: Vec<String>,
sample_output_aspa_customers: Vec<u32>,
payloads: Option<Vec<Payload>>,
}
fn main() -> Result<()> {
let cli = parse_args(env::args().skip(1))?;
let snapshot = load_ccr_snapshot_from_file(&cli.ccr_path)
.with_context(|| format!("failed to load CCR snapshot: {}", cli.ccr_path.display()))?;
let slurm = load_slurm(&cli.slurm_path)?;
let conversion = snapshot_to_payloads_with_options(&snapshot, cli.strict_ccr)?;
let payloads = slurm.apply(&conversion.payloads);
let (input_vrp_count, input_vap_count) = count_vrps_and_vaps(&conversion.payloads);
let (output_vrp_count, output_vap_count) = count_vrps_and_vaps(&payloads);
let output = Output {
ccr_path: cli.ccr_path.display().to_string(),
slurm_path: cli.slurm_path.display().to_string(),
produced_at: snapshot.produced_at.clone(),
slurm_version: slurm.version().as_u32(),
input_payload_count: conversion.payloads.len(),
input_vrp_count,
input_vap_count,
output_payload_count: payloads.len(),
output_vrp_count,
output_vap_count,
invalid_vrps: conversion.invalid_vrps,
invalid_vaps: conversion.invalid_vaps,
sample_output_aspa_customers: sample_aspa_customers(&payloads, 8),
payloads: cli.dump_payloads.then_some(payloads),
};
println!("{}", serde_json::to_string_pretty(&output)?);
Ok(())
}
fn load_slurm(path: &PathBuf) -> Result<SlurmFile> {
let file = File::open(path)
.with_context(|| format!("failed to open SLURM file: {}", path.display()))?;
SlurmFile::from_reader(file)
.with_context(|| format!("failed to parse SLURM file: {}", path.display()))
}
fn parse_args(args: impl Iterator<Item = String>) -> Result<Cli> {
let mut strict_ccr = false;
let mut dump_payloads = false;
let mut positionals = Vec::new();
for arg in args {
match arg.as_str() {
"--strict-ccr" => strict_ccr = true,
"--dump-payloads" => dump_payloads = true,
"-h" | "--help" => {
print_help();
std::process::exit(0);
}
_ if arg.starts_with('-') => {
return Err(anyhow!("unknown option: {}", arg));
}
_ => positionals.push(arg),
}
}
if positionals.is_empty() {
return Ok(Cli {
ccr_path: find_latest_ccr_file("data")
.context("failed to find latest .ccr in ./data for default run")?,
slurm_path: PathBuf::from("data/example.slurm"),
strict_ccr,
dump_payloads,
});
}
if positionals.len() != 2 {
print_help();
return Err(anyhow!(
"expected: slurm_apply_client <snapshot.ccr> <policy.slurm>"
));
}
Ok(Cli {
ccr_path: PathBuf::from(&positionals[0]),
slurm_path: PathBuf::from(&positionals[1]),
strict_ccr,
dump_payloads,
})
}
fn print_help() {
eprintln!(
"Usage: cargo run --bin slurm_apply_client -- [--strict-ccr] [--dump-payloads] <snapshot.ccr> <policy.slurm>"
);
eprintln!();
eprintln!("Reads a CCR snapshot, converts it into payloads, applies SLURM, and prints JSON.");
eprintln!(
"When no arguments are provided, it defaults to the latest .ccr under ./data and ./data/example.slurm."
);
eprintln!("Use --dump-payloads to include the full payload list in the JSON output.");
}
fn count_vrps_and_vaps(payloads: &[Payload]) -> (usize, usize) {
let mut vrps = 0;
let mut vaps = 0;
for payload in payloads {
match payload {
Payload::RouteOrigin(_) => vrps += 1,
Payload::Aspa(_) => vaps += 1,
Payload::RouterKey(_) => {}
}
}
(vrps, vaps)
}
fn sample_aspa_customers(payloads: &[Payload], limit: usize) -> Vec<u32> {
let mut customers = Vec::new();
for payload in payloads {
if let Payload::Aspa(aspa) = payload {
let customer = aspa.customer_asn().into_u32();
if !customers.contains(&customer) {
customers.push(customer);
if customers.len() == limit {
break;
}
}
}
}
customers
}