okooooooo/NodeFlare
1
0 Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Fixes and improvements

Browse Source
This commit is contained in:
imengying committed 2026-09-12 02:34:48 +08:00
1 parent 7882bbbd37
commit 39fcfe50c6
49 files changed
+5585 -4450
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

agent/src/live.rs
Executable
+698
View File
@@ -0,0 +1,698 @@
use super::*;
pub(crate) struct LiveSender {
pub(crate) pending: Arc<(Mutex<VecDeque<Report>>, Condvar)>,
pub(crate) send_interval: Arc<Mutex<Duration>>,
pub(crate) persisted_through: Arc<AtomicI64>,
pub(crate) remote_config: Arc<Mutex<Option<RemoteConfig>>>,
pub(crate) stats: Arc<runtime_stats::RuntimeStats>,
}
pub(crate) struct LiveSenderWorker {
pub(crate) pending: Arc<(Mutex<VecDeque<Report>>, Condvar)>,
pub(crate) configured_interval: Arc<Mutex<Duration>>,
pub(crate) persisted_through: Arc<AtomicI64>,
pub(crate) remote_config: Arc<Mutex<Option<RemoteConfig>>>,
pub(crate) clock: SharedClock,
pub(crate) stats: Arc<runtime_stats::RuntimeStats>,
}
impl LiveSender {
pub(crate) fn start(
config: &RuntimeConfig,
clock: SharedClock,
stats: Arc<runtime_stats::RuntimeStats>,
) -> Result<Self> {
let pending = Arc::new((Mutex::new(VecDeque::new()), Condvar::new()));
let send_interval = Arc::new(Mutex::new(live_batch_interval(config.collect_interval)));
let persisted_through = Arc::new(AtomicI64::new(0));
let remote_config = Arc::new(Mutex::new(None));
let endpoint = live_endpoint(&config.endpoint)?;
let token = config.token.clone();
let sender_state = Arc::clone(&pending);
let sender_interval = Arc::clone(&send_interval);
let sender_persisted_through = Arc::clone(&persisted_through);
let sender_remote_config = Arc::clone(&remote_config);
let sender_stats = Arc::clone(&stats);
thread::Builder::new()
.name("nodeflare-live".to_string())
.spawn(move || {
live_sender_loop(
&endpoint,
&token,
LiveSenderWorker {
pending: sender_state,
configured_interval: sender_interval,
persisted_through: sender_persisted_through,
remote_config: sender_remote_config,
clock,
stats: sender_stats,
},
)
})?;
Ok(Self {
pending,
send_interval,
persisted_through,
remote_config,
stats,
})
}
pub(crate) fn send(&self, report: &Report) {
let (pending, ready) = &*self.pending;
if let Ok(mut pending) = pending.lock() {
if pending.len() >= LIVE_QUEUE_CAPACITY {
pending.pop_front();
self.stats.live_queue_dropped(1);
}
pending.push_back(report.clone());
ready.notify_one();
}
}
pub(crate) fn set_send_interval(&self, collect_interval: u64) {
if let Ok(mut interval) = self.send_interval.lock() {
*interval = live_batch_interval(collect_interval);
}
self.pending.1.notify_one();
}
pub(crate) fn persisted_through(&self) -> i64 {
self.persisted_through.load(Ordering::Acquire)
}
pub(crate) fn take_remote_config(&self) -> Option<RemoteConfig> {
self.remote_config
.lock()
.ok()
.and_then(|mut config| config.take())
}
}
pub(crate) fn live_batch_interval(collect_interval: u64) -> Duration {
Duration::from_secs(collect_interval.clamp(telemetry::MIN_UPLOAD_INTERVAL, 60))
}
pub(crate) fn live_endpoint(endpoint: &str) -> Result<String> {
let mut url = url::Url::parse(endpoint)?;
url.set_scheme(if url.scheme() == "https" { "wss" } else { "ws" })
.map_err(|_| "unsupported live endpoint scheme")?;
let path = format!("{}/api/agent/ws", url.path().trim_end_matches('/'));
url.set_path(&path);
Ok(url.to_string())
}
pub(crate) type LiveSocket = tungstenite::WebSocket<tungstenite::stream::MaybeTlsStream<TcpStream>>;
pub(crate) fn connect_live(endpoint: &str, token: &str) -> Result<(LiveSocket, Option<i64>)> {
let started_ms = unix_timestamp_millis();
let deadline = Instant::now() + LIVE_CONNECT_TIMEOUT;
let mut url = url::Url::parse(endpoint)?;
let mut redirects = 0;
let (socket, response) = loop {
let mut request = url.as_str().into_client_request()?;
request
.headers_mut()
.insert("Authorization", format!("Bearer {token}").parse()?);
request
.headers_mut()
.insert("User-Agent", format!("nodeflare-agent/{VERSION}").parse()?);
request
.headers_mut()
.insert(AGENT_PROTOCOL_HEADER, AGENT_PROTOCOL_VERSION.parse()?);
request.headers_mut().insert(
AGENT_CAPABILITIES_HEADER,
REQUIRED_AGENT_CAPABILITIES.join(",").parse()?,
);
let stream = connect_live_stream(&url, deadline)?;
match client_tls(request, stream) {
Ok(connected) => break connected,
Err(tungstenite::HandshakeError::Failure(WebSocketError::Http(response)))
if response.status().is_redirection() && redirects < 3 =>
{
let location = response
.headers()
.get("location")
.ok_or("WebSocket redirect has no location")?
.to_str()?;
let redirected = url.join(location)?;
// Never forward the Agent token to a different origin.
if redirected.origin() != url.origin() {
return Err("WebSocket redirect must stay on the configured origin".into());
}
url = redirected;
redirects += 1;
}
Err(error) => return Err(format!("WebSocket handshake failed: {error}").into()),
}
};
let ended_ms = unix_timestamp_millis();
let clock_offset_ms = response
.headers()
.get("date")
.and_then(|value| value.to_str().ok())
.and_then(|value| clock_offset_from_http_date(value, started_ms, ended_ms));
if !response
.headers()
.get(AGENT_PROTOCOL_HEADER)
.is_some_and(|value| value == AGENT_PROTOCOL_VERSION)
{
return Err(format!(
"backend protocol mismatch; update the backend to protocol {AGENT_PROTOCOL_VERSION}"
)
.into());
}
Ok((socket, clock_offset_ms))
}
pub(crate) fn connect_live_stream(url: &url::Url, deadline: Instant) -> Result<TcpStream> {
let host = url
.host_str()
.ok_or("WebSocket endpoint has no hostname")?
.trim_start_matches('[')
.trim_end_matches(']');
let port = url
.port_or_known_default()
.ok_or("WebSocket endpoint has no port")?;
let addresses = (host, port).to_socket_addrs()?;
let mut connected = None;
let mut last_error = io::Error::new(
ErrorKind::NotFound,
"WebSocket hostname resolved to no addresses",
);
for address in addresses {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(
io::Error::new(ErrorKind::TimedOut, "WebSocket connection timed out").into(),
);
}
match TcpStream::connect_timeout(&address, remaining.min(Duration::from_secs(3))) {
Ok(stream) => {
connected = Some(stream);
break;
}
Err(error) => last_error = error,
}
}
let stream = connected.ok_or(last_error)?;
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Err(io::Error::new(ErrorKind::TimedOut, "WebSocket connection timed out").into());
}
stream.set_nodelay(true)?;
stream.set_read_timeout(Some(remaining))?;
stream.set_write_timeout(Some(remaining))?;
Ok(stream)
}
pub(crate) fn set_live_read_timeout(socket: &mut LiveSocket, timeout: Option<Duration>) -> io::Result<()> {
match socket.get_mut() {
tungstenite::stream::MaybeTlsStream::Plain(stream) => stream.set_read_timeout(timeout),
tungstenite::stream::MaybeTlsStream::Rustls(stream) => {
stream.sock.set_read_timeout(timeout)
}
// An unmatched TLS backend keeps running without a read timeout rather
// than failing outright; the rustls-only feature set always matches an
// arm above, so this is defensive.
_ => Ok(()),
}
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct LiveAck {
#[serde(rename = "type")]
pub(crate) message_type: String,
pub(crate) ts: i64,
pub(crate) persisted: bool,
#[serde(rename = "persistenceError")]
pub(crate) persistence_error: bool,
#[serde(rename = "persistedThroughTs")]
pub(crate) persisted_through_ts: i64,
#[serde(rename = "nextPersistAfterMs")]
pub(crate) next_persist_after_ms: u64,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct LiveConfigMessage {
#[serde(rename = "type")]
pub(crate) message_type: String,
pub(crate) ts: i64,
pub(crate) config: RemoteConfig,
}
pub(crate) fn ack_persist_interval(ack: &LiveAck) -> Duration {
Duration::from_millis(ack.next_persist_after_ms)
.clamp(Duration::from_secs(1), Duration::from_secs(3600))
}
pub(crate) enum LiveRead {
Ack(LiveAck),
Config(RemoteConfig),
RemoteTask(RemoteTaskMessage),
TaskResultAck(String),
Closed,
Pending,
}
pub(crate) fn read_live_ack(socket: &mut LiveSocket) -> Result<LiveRead> {
match socket.read() {
Ok(Message::Text(text)) => {
let Ok(value) = serde_json::from_str::<serde_json::Value>(text.as_ref()) else {
return Ok(LiveRead::Pending);
};
match value.get("type").and_then(serde_json::Value::as_str) {
Some("ack") => {
let Ok(ack) = serde_json::from_value::<LiveAck>(value) else {
return Ok(LiveRead::Pending);
};
if ack.message_type != "ack" || ack.ts <= 0 {
return Ok(LiveRead::Pending);
}
Ok(LiveRead::Ack(ack))
}
Some("config") => {
let Ok(message) = serde_json::from_value::<LiveConfigMessage>(value) else {
return Ok(LiveRead::Pending);
};
if message.message_type != "config" || message.ts <= 0 {
return Ok(LiveRead::Pending);
}
Ok(LiveRead::Config(message.config))
}
Some("remote_task") => {
let Ok(task) = serde_json::from_value::<RemoteTaskMessage>(value) else {
return Ok(LiveRead::Pending);
};
if task.message_type != "remote_task"
|| !valid_remote_task_id(&task.task_id)
|| task.command.trim().is_empty()
|| task.command.len() > 16_384
{
return Ok(LiveRead::Pending);
}
Ok(LiveRead::RemoteTask(task))
}
Some("task_result_ack") => {
let Ok(ack) = serde_json::from_value::<TaskResultAckMessage>(value) else {
return Ok(LiveRead::Pending);
};
if ack.message_type != "task_result_ack" || !valid_remote_task_id(&ack.task_id)
{
return Ok(LiveRead::Pending);
}
Ok(LiveRead::TaskResultAck(ack.task_id))
}
_ => Ok(LiveRead::Pending),
}
}
Ok(Message::Ping(payload)) => {
socket.send(Message::Pong(payload))?;
Ok(LiveRead::Pending)
}
Ok(Message::Close(_)) => Ok(LiveRead::Closed),
Ok(_) => Ok(LiveRead::Pending),
Err(WebSocketError::Io(error))
if matches!(error.kind(), ErrorKind::TimedOut | ErrorKind::WouldBlock) =>
{
Ok(LiveRead::Pending)
}
Err(error) => Err(error.into()),
}
}
pub(crate) fn accept_remote_task(
pub(crate) socket: &mut LiveSocket,
pub(crate) executor: &mut RemoteExecutor,
pub(crate) task: RemoteTaskMessage,
) -> Result<()> {
let task_id = task.task_id.clone();
let rejected = executor.enqueue(task);
socket.send(Message::Text(
serde_json::json!({
"type": "task_received",
"task_id": task_id,
})
.to_string()
.into(),
))?;
if let Some(result) = rejected {
socket.send(Message::Text(serde_json::to_string(&result)?.into()))?;
}
flush_remote_results(socket, executor)?;
Ok(())
}
pub(crate) fn wait_for_live_ack(
pub(crate) socket: &mut LiveSocket,
pub(crate) remote_config: &Arc<Mutex<Option<RemoteConfig>>>,
pub(crate) remote_executor: &mut RemoteExecutor,
) -> Result<LiveRead> {
let deadline = Instant::now() + LIVE_ACK_READ_TIMEOUT;
loop {
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Ok(LiveRead::Pending);
}
set_live_read_timeout(socket, Some(remaining))?;
match read_live_ack(socket)? {
LiveRead::Ack(ack) => {
return Ok(LiveRead::Ack(ack));
}
LiveRead::Config(config) => {
if let Ok(mut target) = remote_config.lock() {
*target = Some(config);
}
}
LiveRead::RemoteTask(task) => {
accept_remote_task(socket, remote_executor, task)?;
}
LiveRead::TaskResultAck(task_id) => {
remote_executor.acknowledge(&task_id);
}
LiveRead::Closed => return Ok(LiveRead::Closed),
LiveRead::Pending => {}
}
}
}
pub(crate) fn live_update_payload(
pub(crate) reports: Vec<Report>,
pub(crate) persist: bool,
pub(crate) info: &mut Option<telemetry::Info>,
) -> Result<Vec<u8>> {
let mut next_info = info.clone();
let payload = telemetry::encode(&telemetry::Update::from_reports(
reports,
persist,
&mut next_info,
))?;
*info = next_info;
Ok(payload)
}
pub(crate) fn observe_persisted_through(target: &AtomicI64, ack: &LiveAck) {
if ack.persisted && !ack.persistence_error {
target.fetch_max(ack.persisted_through_ts.max(0), Ordering::AcqRel);
}
}
pub(crate) fn prune_live_queue(pending: &Arc<(Mutex<VecDeque<Report>>, Condvar)>, persisted_through: i64) {
if persisted_through <= 0 {
return;
}
if let Ok(mut queue) = pending.0.lock() {
queue.retain(|report| report.timestamp > persisted_through);
}
}
pub(crate) fn live_batch_after(queue: &VecDeque<Report>, timestamp: i64) -> Vec<Report> {
live_batch::batch_from(queue.iter().filter(|report| report.timestamp > timestamp))
}
/// Bounded exponential backoff for the live connection, with jitter derived
/// from the token so a fleet of agents that dropped together does not retry in
/// lockstep and hammer an unreachable server.
pub(crate) fn reconnect_delay(attempt: u32, token: &str) -> Duration {
let base_ms = LIVE_RECONNECT_DELAY.as_millis() as u64;
let backoff_ms = base_ms
.saturating_mul(1_u64 << attempt.min(5))
.min(LIVE_RECONNECT_MAX_DELAY.as_millis() as u64);
let digest = Sha256::digest(token.as_bytes());
let seed = u64::from(digest[0]) ^ (u64::from(attempt) << 8);
let jitter_ms = seed % (backoff_ms / 4 + 1);
Duration::from_millis(backoff_ms + jitter_ms)
}
pub(crate) struct ReconnectBackoff {
pub(crate) attempt: u32,
pub(crate) token: String,
}
impl ReconnectBackoff {
pub(crate) fn new(token: &str) -> Self {
Self {
attempt: 0,
token: token.to_string(),
}
}
/// Reset after a confirmed round trip, not merely after a TCP connect, so a
/// server that accepts then immediately drops connections cannot keep us at
/// the shortest delay.
pub(crate) fn reset(&mut self) {
self.attempt = 0;
}
pub(crate) fn next_delay(&mut self) -> Duration {
self.attempt = self.attempt.saturating_add(1);
reconnect_delay(self.attempt, &self.token)
}
}
pub(crate) fn live_sender_loop(endpoint: &str, token: &str, worker: LiveSenderWorker) {
let LiveSenderWorker {
pending,
configured_interval,
persisted_through,
remote_config,
clock,
stats,
} = worker;
let mut socket: Option<LiveSocket> = None;
let mut backoff = ReconnectBackoff::new(token);
let mut wss_interval = configured_interval.lock().map_or(
Duration::from_secs(telemetry::MIN_UPLOAD_INTERVAL),
|interval| *interval,
);
let mut info = None;
let mut accepted_through = persisted_through.load(Ordering::Acquire);
let mut next_send_at = Instant::now();
let mut next_probe_at = Instant::now();
let mut remote_executor = match RemoteExecutor::new() {
Ok(executor) => executor,
Err(error) => {
eprintln!("remote executor disabled: {error}");
RemoteExecutor::disabled()
}
};
'sender: loop {
if socket.is_none() {
match connect_live(endpoint, token) {
Ok((mut connected, offset_ms)) => {
observe_clock(&clock, offset_ms);
if set_live_read_timeout(&mut connected, Some(LIVE_HINT_READ_TIMEOUT)).is_err()
{
thread::sleep(backoff.next_delay());
continue;
}
socket = Some(connected);
info = None;
remote_executor.reset_delivery();
accepted_through = persisted_through.load(Ordering::Acquire);
next_send_at = Instant::now();
next_probe_at = Instant::now();
}
Err(error) => {
stats.live_connect_failed();
eprintln!("live connection failed: {error}");
thread::sleep(backoff.next_delay());
continue;
}
}
}
if let Some(connected) = socket.as_mut()
&& flush_remote_results(connected, &mut remote_executor).is_err()
{
socket = None;
thread::sleep(backoff.next_delay());
continue;
}
let (queue_lock, ready) = &*pending;
let Ok(mut queue) = queue_lock.lock() else {
return;
};
let (batch, persist, more_pending) = loop {
let now = Instant::now();
let persisted = persisted_through.load(Ordering::Acquire);
let has_unsent = queue.iter().any(|report| report.timestamp > accepted_through);
let persist =
now >= next_probe_at && queue.iter().any(|report| report.timestamp > persisted);
if persist || (has_unsent && now >= next_send_at) {
// Build (and serialize) the batch only once it is actually due;
// the wait below re-runs several times per second while idle.
let unsent = live_batch_after(&queue, accepted_through);
let last = unsent
.last()
.map_or(accepted_through, |report| report.timestamp);
let more_pending = queue.iter().any(|report| report.timestamp > last);
// A commit can be empty: samples already sent remain buffered by the server.
break (unsent, persist || more_pending, more_pending);
}
let wake_at = if has_unsent {
next_send_at.min(next_probe_at)
} else if queue.iter().any(|report| report.timestamp > persisted) {
next_probe_at
} else {
now + Duration::from_millis(250)
};
let wait = wake_at
.saturating_duration_since(now)
.min(Duration::from_millis(250));
queue = match ready.wait_timeout(queue, wait) {
Ok((queue, _)) => queue,
Err(_) => return,
};
drop(queue);
let mut drop_socket = false;
if let Some(connected) = socket.as_mut() {
if flush_remote_results(connected, &mut remote_executor).is_err()
|| set_live_read_timeout(connected, Some(LIVE_HINT_READ_TIMEOUT)).is_err()
{
drop_socket = true;
} else {
match read_live_ack(connected) {
Ok(LiveRead::Closed) => drop_socket = true,
Ok(LiveRead::Ack(ack)) => {
observe_persisted_through(&persisted_through, &ack);
if ack.persistence_error {
stats.live_persistence_failed();
drop_socket = true;
} else {
prune_live_queue(&pending, ack.persisted_through_ts);
accepted_through = accepted_through.max(ack.persisted_through_ts);
next_probe_at = Instant::now() + ack_persist_interval(&ack);
}
}
Ok(LiveRead::Config(config)) => {
if let Ok(mut target) = remote_config.lock() {
*target = Some(config);
}
}
Ok(LiveRead::RemoteTask(task)) => {
if accept_remote_task(connected, &mut remote_executor, task).is_err() {
drop_socket = true;
}
}
Ok(LiveRead::TaskResultAck(task_id)) => {
remote_executor.acknowledge(&task_id);
}
Ok(LiveRead::Pending) => {}
Err(error) => {
eprintln!("live hint read failed: {error}");
drop_socket = true;
}
}
}
}
if drop_socket {
socket = None;
accepted_through = persisted_through.load(Ordering::Acquire);
next_send_at = Instant::now();
next_probe_at = Instant::now();
thread::sleep(backoff.next_delay());
continue 'sender;
}
queue = match queue_lock.lock() {
Ok(queue) => queue,
Err(_) => return,
};
};
drop(queue);
if let Ok(interval) = configured_interval.lock() {
wss_interval = (*interval).clamp(
Duration::from_secs(telemetry::MIN_UPLOAD_INTERVAL),
Duration::from_secs(60),
);
}
let batch_last = batch.last().map(|report| report.timestamp);
let batch_len = batch.len();
let payload = match live_update_payload(batch, persist, &mut info) {
Ok(payload) => payload,
Err(error) => {
eprintln!("live payload encode failed: {error}");
next_send_at = Instant::now() + wss_interval;
continue;
}
};
let sent = socket
.as_mut()
.is_some_and(|socket| socket.send(Message::Binary(payload.into())).is_ok());
if !sent {
socket = None;
accepted_through = persisted_through.load(Ordering::Acquire);
next_send_at = Instant::now();
next_probe_at = Instant::now();
thread::sleep(backoff.next_delay());
continue;
}
stats.live_batch_sent(batch_len);
backoff.reset();
if let Some(timestamp) = batch_last {
accepted_through = accepted_through.max(timestamp);
}
next_send_at = Instant::now() + wss_interval;
if !persist {
continue;
}
let mut drop_socket = false;
if let Some(socket) = socket.as_mut() {
match wait_for_live_ack(socket, &remote_config, &mut remote_executor) {
Ok(LiveRead::Closed) => drop_socket = true,
Ok(LiveRead::Ack(ack)) => {
observe_persisted_through(&persisted_through, &ack);
if ack.persistence_error {
stats.live_persistence_failed();
drop_socket = true;
} else {
prune_live_queue(&pending, ack.persisted_through_ts);
accepted_through = accepted_through.max(ack.persisted_through_ts);
next_probe_at = Instant::now() + ack_persist_interval(&ack);
if more_pending
&& batch_last
.is_some_and(|timestamp| ack.persisted_through_ts >= timestamp)
{
next_probe_at = Instant::now();
}
}
}
// wait_for_live_ack consumes config and remote-task messages
// internally, so only Ack/Closed/Pending reach this match; these
// arms are kept defensively in case that contract changes.
Ok(LiveRead::Config(config)) => {
if let Ok(mut target) = remote_config.lock() {
*target = Some(config);
}
drop_socket = true;
}
Ok(LiveRead::RemoteTask(_) | LiveRead::TaskResultAck(_)) => {}
Ok(LiveRead::Pending) => {
eprintln!("live ACK timed out");
drop_socket = true;
}
Err(error) => {
eprintln!("live ACK read failed: {error}");
drop_socket = true;
}
}
}
if drop_socket {
socket = None;
accepted_through = persisted_through.load(Ordering::Acquire);
next_send_at = Instant::now();
next_probe_at = Instant::now();
thread::sleep(backoff.next_delay());
}
}
}
agent/src/live_batch.rs
Regular → Executable
+13 -9
View File
@@ -1,13 +1,17 @@
use std::collections::VecDeque;
use super::Report;
pub(crate) const MAX_LIVE_BATCH_BYTES: usize = nodeflare_telemetry::MAX_BATCH_BYTES;
pub(crate) fn batch_len(queue: &VecDeque<Report>) -> usize {
/// Collects the longest prefix of `queue` that fits one live batch, cloning only
/// the reports that are actually selected and serializing each candidate once
/// (the size check needs the encoded length).
pub(crate) fn batch_from<'a>(queue: impl Iterator<Item = &'a Report>) -> Vec<Report> {
let mut encoded_bytes = 192_usize;
let mut count = 0_usize;
let mut batch = Vec::new();
for report in queue {
if batch.len() >= super::LIVE_QUEUE_CAPACITY {
break;
}
let Ok(report_bytes) = serde_json::to_vec(report) else {
break;
};
@@ -18,19 +22,19 @@ pub(crate) fn batch_len(queue: &VecDeque<Report>) -> usize {
break;
}
encoded_bytes = next_bytes;
count += 1;
batch.push(report.clone());
if encoded_bytes >= MAX_LIVE_BATCH_BYTES {
break;
}
}
count.min(super::LIVE_QUEUE_CAPACITY)
batch
}
#[cfg(test)]
mod tests {
use std::collections::VecDeque;
use super::{MAX_LIVE_BATCH_BYTES, batch_len};
use super::{MAX_LIVE_BATCH_BYTES, batch_from};
use crate::Report;
#[test]
@@ -40,7 +44,7 @@ mod tests {
..Report::default()
};
let queue = std::iter::repeat_n(report, 10).collect::<VecDeque<_>>();
let count = batch_len(&queue);
let count = batch_from(queue.iter()).len();
assert!((1..10).contains(&count));
let payload = crate::live_update_payload(
queue.iter().take(count).cloned().collect::<Vec<_>>(),
@@ -92,7 +96,7 @@ mod tests {
..Report::default()
})
.collect::<VecDeque<_>>();
assert!(batch_len(&queue) <= crate::LIVE_QUEUE_CAPACITY);
assert!(batch_from(queue.iter()).len() <= crate::LIVE_QUEUE_CAPACITY);
let reports = crate::live_batch_after(&queue, 100);
assert_eq!(reports.first().unwrap().timestamp, 101);
for persist in [true, false] {
agent/src/main.rs
Regular → Executable
+64 -3654
View File
File diff suppressed because it is too large. Load diff
agent/src/metrics.rs
Executable
+1398
View File
File diff suppressed because it is too large. Load diff
agent/src/remote.rs
Executable
+587
View File
@@ -0,0 +1,587 @@
use super::*;
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct RemoteTaskMessage {
#[serde(rename = "type")]
pub(crate) message_type: String,
pub(crate) task_id: String,
pub(crate) command: String,
}
#[derive(Debug, Clone, Deserialize, Serialize, PartialEq, Eq)]
pub(crate) struct TaskResultMessage {
#[serde(rename = "type")]
pub(crate) message_type: String,
pub(crate) task_id: String,
pub(crate) status: String,
pub(crate) result: String,
pub(crate) exit_code: Option<i64>,
}
#[derive(Debug, Deserialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct TaskResultAckMessage {
#[serde(rename = "type")]
pub(crate) message_type: String,
pub(crate) task_id: String,
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(deny_unknown_fields)]
pub(crate) struct RemoteTaskJournal {
pub(crate) version: u32,
pub(crate) entries: Vec<RemoteTaskJournalEntry>,
}
#[derive(Debug, Deserialize, Serialize)]
#[serde(tag = "state", rename_all = "snake_case")]
pub(crate) enum RemoteTaskJournalEntry {
Active { task_id: String },
Completed { result: TaskResultMessage },
}
#[derive(Default)]
pub(crate) struct CapturedOutput {
pub(crate) bytes: Vec<u8>,
pub(crate) truncated: bool,
}
pub(crate) fn capture_remote_output<R>(mut reader: R) -> mpsc::Receiver<CapturedOutput>
where
R: Read + Send + 'static,
{
let (sender, receiver) = mpsc::sync_channel(1);
thread::spawn(move || {
let mut bytes = Vec::new();
let mut truncated = false;
let mut buffer = [0_u8; 8192];
// Keep draining after the limit so verbose commands do not receive SIGPIPE.
loop {
match reader.read(&mut buffer) {
Ok(0) => break,
Ok(read) => {
let retained = read.min(REMOTE_STREAM_OUTPUT_BYTES as usize - bytes.len());
bytes.extend_from_slice(&buffer[..retained]);
truncated |= retained < read;
}
Err(error) if error.kind() == ErrorKind::Interrupted => continue,
Err(_) => break,
}
}
let _ = sender.send(CapturedOutput { bytes, truncated });
});
receiver
}
pub(crate) fn terminate_remote_process_tree(child: &mut Child) {
#[cfg(unix)]
{
let process_group = format!("-{}", child.id());
let _ = Command::new("kill")
.args(["-KILL", "--", process_group.as_str()])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status();
}
#[cfg(target_os = "windows")]
{
let _ = Command::new("taskkill")
.args(["/PID", &child.id().to_string(), "/T", "/F"])
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.status();
}
let _ = child.kill();
let _ = child.wait();
}
pub(crate) fn receive_remote_output(
pub(crate) receiver: &mpsc::Receiver<CapturedOutput>,
pub(crate) child: &mut Child,
) -> CapturedOutput {
match receiver.recv_timeout(REMOTE_OUTPUT_DRAIN_TIMEOUT) {
Ok(output) => output,
Err(_) => {
// Descendants may keep the pipe open after the shell exits.
terminate_remote_process_tree(child);
receiver
.recv_timeout(REMOTE_OUTPUT_DRAIN_TIMEOUT)
.unwrap_or_default()
}
}
}
pub(crate) fn remote_result_text(stdout: &CapturedOutput, stderr: &CapturedOutput) -> String {
let truncated = stdout.truncated || stderr.truncated;
let stdout = String::from_utf8_lossy(&stdout.bytes);
let stderr = String::from_utf8_lossy(&stderr.bytes);
let mut combined = if stderr.is_empty() {
stdout.trim().to_string()
} else if stdout.is_empty() {
stderr.trim().to_string()
} else {
format!("{}\n{}", stdout.trim(), stderr.trim())
};
let result_truncated = combined.len() > REMOTE_RESULT_OUTPUT_BYTES;
if result_truncated {
let mut end = REMOTE_RESULT_OUTPUT_BYTES;
while !combined.is_char_boundary(end) {
end -= 1;
}
combined.truncate(end);
}
if truncated || result_truncated {
combined.push_str("\n[输出已截断]");
}
combined
}
pub(crate) fn wait_for_remote_child(
pub(crate) child: &mut Child,
pub(crate) timeout: Duration,
) -> io::Result<Option<std::process::ExitStatus>> {
let deadline = Instant::now() + timeout;
loop {
if let Some(status) = child.try_wait()? {
return Ok(Some(status));
}
let remaining = deadline.saturating_duration_since(Instant::now());
if remaining.is_zero() {
return Ok(None);
}
thread::sleep(remaining.min(Duration::from_millis(50)));
}
}
pub(crate) fn execute_remote_task_with_timeout(
pub(crate) task: &RemoteTaskMessage,
pub(crate) timeout: Duration,
) -> TaskResultMessage {
let task_id = task.task_id.clone();
if task.command.trim().is_empty() || task.command.len() > 16_384 {
return TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: "failed".to_string(),
result: "命令为空或长度超出限制".to_string(),
exit_code: Some(-1),
};
}
#[cfg(unix)]
let mut command = {
let mut command = Command::new("sh");
command.args(["-c", task.command.as_str()]);
command.process_group(0);
command
};
#[cfg(target_os = "windows")]
let mut command = {
let mut command = Command::new("cmd.exe");
command.args(["/C", task.command.as_str()]);
command
};
command.stdout(Stdio::piped()).stderr(Stdio::piped());
let mut child = match command.spawn() {
Ok(child) => child,
Err(error) => {
return TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: "failed".to_string(),
result: format!("无法启动命令：{error}"),
exit_code: Some(-1),
};
}
};
let stdout = capture_remote_output(child.stdout.take().expect("stdout is piped"));
let stderr = capture_remote_output(child.stderr.take().expect("stderr is piped"));
let status = match wait_for_remote_child(&mut child, timeout) {
Ok(Some(status)) => Some(status),
Ok(None) => {
terminate_remote_process_tree(&mut child);
None
}
Err(error) => {
terminate_remote_process_tree(&mut child);
let stdout = receive_remote_output(&stdout, &mut child);
let stderr = receive_remote_output(&stderr, &mut child);
let output = remote_result_text(&stdout, &stderr);
return TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: "failed".to_string(),
result: if output.is_empty() {
format!("等待命令结束失败：{error}")
} else {
format!("{output}\n等待命令结束失败：{error}")
},
exit_code: Some(-1),
};
}
};
let stdout = receive_remote_output(&stdout, &mut child);
let stderr = receive_remote_output(&stderr, &mut child);
let output = remote_result_text(&stdout, &stderr);
let Some(status) = status else {
let timeout_seconds = timeout.as_secs().max(1);
return TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: "failed".to_string(),
result: if output.is_empty() {
format!("命令执行超过 {timeout_seconds} 秒，已终止")
} else {
format!("{output}\n命令执行超过 {timeout_seconds} 秒，已终止")
},
exit_code: Some(-1),
};
};
TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: if status.success() {
"success"
} else {
"failed"
}
.to_string(),
result: output,
exit_code: status.code().map(i64::from),
}
}
pub(crate) fn execute_remote_task(task: &RemoteTaskMessage) -> TaskResultMessage {
execute_remote_task_with_timeout(task, REMOTE_TASK_TIMEOUT)
}
pub(crate) struct RemoteExecutor {
pub(crate) task_tx: mpsc::SyncSender<RemoteTaskMessage>,
pub(crate) result_rx: mpsc::Receiver<TaskResultMessage>,
pub(crate) active: HashSet<String>,
pub(crate) completed: HashMap<String, TaskResultMessage>,
pub(crate) task_order: VecDeque<String>,
pub(crate) last_sent: HashMap<String, Instant>,
pub(crate) journal_path: PathBuf,
pub(crate) disabled: bool,
}
impl RemoteExecutor {
pub(crate) fn new() -> Result<Self> {
Self::new_at(remote_task_journal_path()?)
}
pub(crate) fn new_at(journal_path: PathBuf) -> Result<Self> {
let journal = load_remote_task_journal(&journal_path)?;
let mut active = HashSet::new();
let mut completed = HashMap::new();
let mut task_order = VecDeque::new();
for entry in journal.entries {
match entry {
RemoteTaskJournalEntry::Active { task_id } => {
active.insert(task_id.clone());
task_order.push_back(task_id);
}
RemoteTaskJournalEntry::Completed { result } => {
task_order.push_back(result.task_id.clone());
completed.insert(result.task_id.clone(), result);
}
}
}
let (task_tx, task_rx) = mpsc::sync_channel::<RemoteTaskMessage>(16);
let (result_tx, result_rx) = mpsc::channel::<TaskResultMessage>();
thread::Builder::new()
.name("nodeflare-remote".to_string())
.spawn(move || {
while let Ok(task) = task_rx.recv() {
if result_tx.send(execute_remote_task(&task)).is_err() {
return;
}
}
})?;
let interrupted = active.drain().collect::<Vec<_>>();
for task_id in interrupted {
completed.insert(
task_id.clone(),
TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: "failed".to_string(),
result: "Agent 在命令执行期间重启，无法确认原命令状态；为避免重复操作，任务不会再次执行"
.to_string(),
exit_code: Some(-1),
},
);
}
let executor = Self {
task_tx,
result_rx,
active,
completed,
task_order,
last_sent: HashMap::new(),
journal_path,
disabled: false,
};
if !executor.completed.is_empty() {
executor.persist_snapshot(None)?;
}
Ok(executor)
}
pub(crate) fn disabled() -> Self {
let (task_tx, task_rx) = mpsc::sync_channel(1);
let (result_tx, result_rx) = mpsc::channel();
drop(task_rx);
drop(result_tx);
Self {
task_tx,
result_rx,
active: HashSet::new(),
completed: HashMap::new(),
task_order: VecDeque::new(),
last_sent: HashMap::new(),
journal_path: PathBuf::new(),
disabled: true,
}
}
pub(crate) fn enqueue(&mut self, task: RemoteTaskMessage) -> Option<TaskResultMessage> {
if self.disabled {
return Some(remote_task_rejected(
task.task_id,
"Agent 任务日志异常，远程执行已禁用",
));
}
if self.completed.contains_key(&task.task_id) {
self.last_sent.remove(&task.task_id);
return None;
}
if self.active.contains(&task.task_id) {
return None;
}
let task_id = task.task_id.clone();
if self.task_order.len() >= MAX_REMOTE_TASK_JOURNAL_ENTRIES {
return Some(remote_task_rejected(
task_id,
"待确认的远程任务过多，已拒绝执行",
));
}
self.active.insert(task_id.clone());
self.task_order.push_back(task_id.clone());
if let Err(error) = self.persist_snapshot(None) {
self.active.remove(&task_id);
self.task_order.retain(|id| id != &task_id);
eprintln!("remote task journal write failed: {error}");
return Some(remote_task_rejected(
task_id,
"Agent 无法持久化任务状态，已拒绝执行",
));
}
match self.task_tx.try_send(task) {
Ok(()) => None,
Err(error) => {
let message = match error {
mpsc::TrySendError::Full(_) => "远程命令队列已满，请稍后重试",
mpsc::TrySendError::Disconnected(_) => "远程命令执行线程不可用",
};
let result = remote_task_rejected(task_id, message);
self.remember_completed(result.clone());
None
}
}
}
pub(crate) fn drain_completed(&mut self) {
let results = self.result_rx.try_iter().collect::<Vec<_>>();
for result in results {
self.remember_completed(result);
}
}
pub(crate) fn acknowledge(&mut self, task_id: &str) {
if !self.completed.contains_key(task_id) {
return;
}
match self.persist_snapshot(Some(task_id)) {
Ok(()) => {
self.completed.remove(task_id);
self.task_order.retain(|id| id != task_id);
self.last_sent.remove(task_id);
}
Err(error) => {
eprintln!("remote task ACK persistence failed: {error}");
}
}
}
pub(crate) fn remember_completed(&mut self, result: TaskResultMessage) {
let task_id = result.task_id.clone();
let was_active = self.active.remove(&task_id);
if !was_active && !self.completed.contains_key(&task_id) {
self.task_order.push_back(task_id.clone());
}
self.completed.insert(task_id.clone(), result);
self.last_sent.remove(&task_id);
if let Err(error) = self.persist_snapshot(None) {
eprintln!("remote task result persistence failed: {error}");
}
}
pub(crate) fn due_results(&mut self) -> Vec<TaskResultMessage> {
self.drain_completed();
let now = Instant::now();
self.task_order
.iter()
.filter_map(|task_id| {
let result = self.completed.get(task_id)?;
let due = self.last_sent.get(task_id).is_none_or(|sent_at| {
now.duration_since(*sent_at) >= REMOTE_RESULT_RETRY_INTERVAL
});
due.then(|| result.clone())
})
.collect()
}
pub(crate) fn mark_sent(&mut self, task_id: &str) {
self.last_sent.insert(task_id.to_string(), Instant::now());
}
pub(crate) fn reset_delivery(&mut self) {
self.last_sent.clear();
}
pub(crate) fn persist_snapshot(&self, excluded_task_id: Option<&str>) -> Result<()> {
let entries = self
.task_order
.iter()
.filter(|task_id| excluded_task_id != Some(task_id.as_str()))
.filter_map(|task_id| {
if self.active.contains(task_id) {
Some(RemoteTaskJournalEntry::Active {
task_id: task_id.clone(),
})
} else {
self.completed
.get(task_id)
.cloned()
.map(|result| RemoteTaskJournalEntry::Completed { result })
}
})
.collect::<Vec<_>>();
write_remote_task_journal(&self.journal_path, entries)
}
}
pub(crate) fn remote_task_rejected(task_id: String, result: &str) -> TaskResultMessage {
TaskResultMessage {
message_type: "task_result".to_string(),
task_id,
status: "failed".to_string(),
result: result.to_string(),
exit_code: Some(-1),
}
}
pub(crate) fn valid_remote_task_id(value: &str) -> bool {
value.len() == 36
&& value.bytes().enumerate().all(|(index, byte)| {
if matches!(index, 8 | 13 | 18 | 23) {
byte == b'-'
} else {
byte.is_ascii_hexdigit()
}
})
}
pub(crate) fn load_remote_task_journal(path: &Path) -> Result<RemoteTaskJournal> {
let bytes = match fs::read(path) {
Ok(bytes) => bytes,
Err(error) if error.kind() == ErrorKind::NotFound => {
return Ok(RemoteTaskJournal {
version: REMOTE_TASK_JOURNAL_VERSION,
entries: Vec::new(),
});
}
Err(error) => return Err(error.into()),
};
if bytes.len() as u64 > MAX_REMOTE_TASK_JOURNAL_BYTES {
return Err("remote task journal is too large".into());
}
let journal = serde_json::from_slice::<RemoteTaskJournal>(&bytes)?;
if journal.version != REMOTE_TASK_JOURNAL_VERSION {
return Err("unsupported remote task journal version".into());
}
if journal.entries.len() > MAX_REMOTE_TASK_JOURNAL_ENTRIES {
return Err("remote task journal has too many entries".into());
}
let mut task_ids = HashSet::new();
for entry in &journal.entries {
let (task_id, valid_result) = match entry {
RemoteTaskJournalEntry::Active { task_id } => (task_id, true),
RemoteTaskJournalEntry::Completed { result } => (
&result.task_id,
result.message_type == "task_result"
&& matches!(result.status.as_str(), "success" | "failed")
&& result.result.len() <= REMOTE_RESULT_OUTPUT_BYTES + 64,
),
};
if !valid_remote_task_id(task_id) || !valid_result || !task_ids.insert(task_id) {
return Err("remote task journal contains an invalid entry".into());
}
}
Ok(journal)
}
pub(crate) fn write_remote_task_journal(path: &Path, entries: Vec<RemoteTaskJournalEntry>) -> Result<()> {
if entries.is_empty() {
return match fs::remove_file(path) {
Ok(()) => Ok(()),
Err(error) if error.kind() == ErrorKind::NotFound => Ok(()),
Err(error) => Err(error.into()),
};
}
let journal = RemoteTaskJournal {
version: REMOTE_TASK_JOURNAL_VERSION,
entries,
};
let mut encoded = serde_json::to_vec(&journal)?;
encoded.push(b'\n');
if encoded.len() as u64 > MAX_REMOTE_TASK_JOURNAL_BYTES {
return Err("remote task journal is too large".into());
}
let temporary = path.with_file_name(format!(
".{}.{}.tmp",
path.file_name()
.and_then(|value| value.to_str())
.unwrap_or("remote-tasks"),
std::process::id()
));
let mut file = OpenOptions::new()
.create(true)
.truncate(true)
.write(true)
.open(&temporary)?;
#[cfg(unix)]
fs::set_permissions(&temporary, fs::Permissions::from_mode(0o600))?;
file.write_all(&encoded)?;
file.sync_all()?;
drop(file);
// fs::rename maps to MoveFileEx(MOVEFILE_REPLACE_EXISTING) on Windows, so it
// replaces the destination atomically; deleting it first would open a window
// where a crash loses the whole journal.
if let Err(error) = fs::rename(&temporary, path) {
let _ = fs::remove_file(&temporary);
return Err(error.into());
}
Ok(())
}
agent/src/tests.rs
Executable
+811
View File
@@ -0,0 +1,811 @@
use super::*;
use clap::Parser;
use std::fs;
use std::path::PathBuf;
use std::time::{Duration, Instant, SystemTime, UNIX_EPOCH};
use super::{
CLOCK_CALIBRATION_MAX_AGE, CapturedOutput, CliOptions, ClockCalibration, DiskMetric,
GithubReleaseAsset, LatencyResult, LatencyTask, LiveAck, MAX_PENDING_LATENCY_RESULTS,
PUBLIC_IP_STALE_AFTER, PublicIpValue, REMOTE_RESULT_OUTPUT_BYTES,
REMOTE_STREAM_OUTPUT_BYTES, RemoteExecutor, RemoteTaskJournalEntry, RemoteTaskMessage,
Report, TaskResultMessage, UPDATE_CHECK_JITTER_MAX_SECONDS, ack_persist_interval,
advance_deadline, clock_offset_from_http_date, corrected_timestamp, dedupe_disks,
execute_remote_task_with_timeout, gpu_name_from_uevent, is_public_probe_ip, live_endpoint,
live_update_payload, monotonic_report_timestamp, normalized_version, parse_lspci_gpu_names,
parse_pciconf_gpu_names, parse_probe_target, parse_public_ip,
parse_system_profiler_gpu_names, ping_latencies, ping_latency, prune_report_samples,
release_asset_sha256, remote_result_text, sanitize_latency_tasks, selected_interface,
update_check_jitter, valid_endpoint, version_triplet, wildcard_match,
write_remote_task_journal,
};
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "freebsd"))]
use super::connection_counts_from_netstat;
#[cfg(target_os = "linux")]
use super::disk_device;
#[test]
fn parses_family_matched_public_ips() {
let v4: std::net::IpAddr = "203.0.113.7".parse().unwrap();
let v6: std::net::IpAddr = "2001:db8::1".parse().unwrap();
assert_eq!(parse_public_ip("203.0.113.7\n", false), Some(v4));
assert_eq!(parse_public_ip("2001:db8::1", true), Some(v6));
assert_eq!(parse_public_ip("2001:db8::1", false), None);
assert_eq!(parse_public_ip("203.0.113.7", true), None);
assert_eq!(parse_public_ip("not an ip", false), None);
assert_eq!(parse_public_ip("", false), None);
}
#[test]
fn expires_public_ip_after_sustained_probe_failures() {
let started_at = Instant::now();
let address = "203.0.113.7".parse().unwrap();
let mut value = PublicIpValue::default();
value.observe(Some(address), started_at);
value.observe(
None,
started_at + PUBLIC_IP_STALE_AFTER - Duration::from_secs(1),
);
assert_eq!(value.address, Some(address));
value.observe(None, started_at + PUBLIC_IP_STALE_AFTER);
assert_eq!(value.address, None);
}
#[test]
fn parses_platform_gpu_names() {
let lspci = "00:02.0 VGA compatible controller: Intel Corporation CometLake-S GT2 [UHD Graphics 630] (rev 05)\n\
01:00.0 Audio device: NVIDIA Corporation HDMI Audio\n\
02:00.0 3D controller: NVIDIA Corporation GA102 [GeForce RTX 3090] (rev a1)";
assert_eq!(
parse_lspci_gpu_names(lspci),
[
"Intel Corporation CometLake-S GT2 [UHD Graphics 630] (rev 05)",
"NVIDIA Corporation GA102 [GeForce RTX 3090] (rev a1)",
]
);
let profiler = "Graphics/Displays:\n\n Apple M3 Max:\n\n Chipset Model: Apple M3 Max\n Metal Support: Metal 3";
assert_eq!(parse_system_profiler_gpu_names(profiler), ["Apple M3 Max"]);
let pciconf = "vgapci0@pci0:0:2:0:\tclass=0x030000 rev=0x02 hdr=0x00 vendor=0x8086\n\
\tvendor = 'Intel Corporation'\n\
\tdevice = 'UHD Graphics 630'\n\
\tclass = display\n\
em0@pci0:0:25:0:\tclass=0x020000 rev=0x05 hdr=0x00 vendor=0x8086\n\
\tvendor = 'Intel Corporation'\n\
\tdevice = 'Ethernet Connection'\n\
\tclass = network";
assert_eq!(
parse_pciconf_gpu_names(pciconf),
["Intel Corporation UHD Graphics 630"]
);
assert_eq!(
gpu_name_from_uevent("DRIVER=i915\nPCI_CLASS=30000\nPCI_ID=8086:591B"),
Some("Intel Integrated Graphics".to_string())
);
assert_eq!(
gpu_name_from_uevent("DRIVER=virtio_gpu\nPCI_CLASS=30000"),
None
);
}
#[cfg(target_os = "linux")]
#[test]
fn selects_whole_disk_devices() {
assert!(disk_device("sda"));
assert!(disk_device("nvme0n1"));
assert!(!disk_device("sda1"));
assert!(!disk_device("nvme0n1p1"));
}
#[test]
fn filters_virtual_interfaces_and_supports_include_exclude_patterns() {
assert!(!selected_interface("lo", ""));
assert!(!selected_interface("docker0", ""));
assert!(!selected_interface("veth123", ""));
assert!(!selected_interface("br0", "!eth*"));
assert!(selected_interface("br0", "br0"));
assert!(!selected_interface("br0", "br*,!br0"));
assert!(!selected_interface("br0", "!br0,br*"));
assert!(!selected_interface("Software Loopback Interface 1", ""));
assert!(selected_interface("eth0", ""));
assert!(selected_interface("ens3", "eth*,ens*"));
assert!(!selected_interface("wlan0", "eth*,ens*"));
assert!(!selected_interface("eth0", "eth*,!eth0"));
assert!(selected_interface("ens3", "!eth*"));
assert!(wildcard_match("en?3", "ens3"));
assert!(!wildcard_match("en?3", "enp4s0"));
}
#[test]
fn deduplicates_mounts_for_the_same_device() {
let disks = dedupe_disks([
DiskMetric {
name: "/dev/vg/root".into(),
mount_point: "/var".into(),
total: 10,
used: 4,
..DiskMetric::default()
},
DiskMetric {
name: "/dev/vg/root".into(),
mount_point: "/".into(),
total: 20,
used: 8,
..DiskMetric::default()
},
DiskMetric {
name: "tank/data".into(),
mount_point: "/data".into(),
total: 30,
used: 9,
..DiskMetric::default()
},
DiskMetric {
name: "tank/archive".into(),
mount_point: "/archive".into(),
total: 25,
used: 7,
..DiskMetric::default()
},
DiskMetric {
name: "/var/lib/data".into(),
mount_point: "/var/lib/data".into(),
total: 40,
used: 12,
..DiskMetric::default()
},
]);
assert_eq!(disks.len(), 3);
assert_eq!(disks[0].mount_point, "/");
assert_eq!(disks[0].total, 20);
assert_eq!(disks[1].name, "tank/data");
assert_eq!(disks[2].name, "/var/lib/data");
}
#[cfg(target_os = "linux")]
#[test]
fn keeps_root_filesystem_but_filters_virtual_mounts() {
assert!(!super::excluded_filesystem("overlay", "/"));
assert!(super::excluded_filesystem("tmpfs", "/tmp"));
assert!(super::excluded_filesystem("proc", "/proc"));
assert!(!super::excluded_filesystem("btrfs", "/var"));
}
#[cfg(target_os = "linux")]
#[test]
fn memory_matches_komari_default_including_shared_memory_and_swap_cache() {
let memory = super::memory_sample(
"MemTotal: 1000000 kB\nMemFree: 100000 kB\nMemAvailable: 750000 kB\n\
Cached: 300000 kB\nSReclaimable: 50000 kB\nBuffers: 25000 kB\n\
Shmem: 25000 kB\nSwapTotal: 100000 kB\nSwapFree: 60000 kB\nSwapCached: 20000 kB\n",
);
assert_eq!(memory.total, 1_000_000 * 1024);
assert_eq!(memory.used, 550_000 * 1024);
assert_eq!(memory.swap_total, 100_000 * 1024);
assert_eq!(memory.swap_used, 20_000 * 1024);
}
#[cfg(target_os = "linux")]
#[test]
fn filling_file_cache_does_not_inflate_reported_memory() {
let before = super::memory_sample(
"MemTotal: 1000 kB\nMemFree: 400 kB\nCached: 300 kB\n\
Buffers: 50 kB\nSReclaimable: 50 kB\nShmem: 100 kB\n",
);
let after = super::memory_sample(
"MemTotal: 1000 kB\nMemFree: 200 kB\nCached: 500 kB\n\
Buffers: 50 kB\nSReclaimable: 50 kB\nShmem: 100 kB\n",
);
assert_eq!(before.used, 300 * 1024);
assert_eq!(after.used, before.used);
}
#[cfg(target_os = "linux")]
#[test]
fn memory_handles_missing_and_inconsistent_kernel_counters() {
let memory = super::memory_sample(
"MemTotal: 1000 kB\nMemFree: 200 kB\nCached: 1000 kB\n\
Shmem: 100 kB\nSwapTotal: 100 kB\nSwapFree: 20 kB\nSwapCached: 100 kB\n",
);
assert_eq!(memory.used, 900 * 1024);
assert_eq!(memory.swap_used, 80 * 1024);
for contents in ["", "MemTotal: invalid kB\n", "MemTotal: -1 kB\n"] {
let memory = super::memory_sample(contents);
assert_eq!(
(
memory.total,
memory.used,
memory.swap_total,
memory.swap_used
),
(0, 0, 0, 0)
);
}
let memory = super::memory_sample(
"MemTotal: 100 kB\nMemFree: 1000 kB\nSwapTotal: 100 kB\nSwapFree: 1000 kB\n",
);
assert_eq!((memory.used, memory.swap_used), (0, 0));
}
#[cfg(target_os = "linux")]
#[test]
fn reads_local_mounts_and_decodes_escaped_paths_once() {
let mounts = super::disk_mounts(
r"overlay / overlay rw 0 0
proc /proc proc rw 0 0
tmpfs /run tmpfs rw 0 0
/dev/vda1 /data\040volume ext4 rw 0 0
/dev/vda2 /literal\134040path ext4 rw 0 0
/dev/vda3 /tab\011dir ext4 rw 0 0
server:/export /remote nfs4 rw 0 0
//server/share /remote-share cifs rw 0 0
none /auto autofs rw 0 0
invalid",
)
.collect::<Vec<_>>();
assert_eq!(
mounts,
vec![
("overlay".into(), "/".into()),
("/dev/vda1".into(), "/data volume".into()),
("/dev/vda2".into(), "/literal\\040path".into()),
("/dev/vda3".into(), "/tab\tdir".into()),
]
);
}
#[cfg(target_os = "linux")]
#[test]
fn reads_root_disk_capacity_without_external_commands() {
let (total, used) = super::filesystem_space("/").unwrap();
assert!(total > 0);
assert!((0..=total).contains(&used));
assert!(
super::disk_usage()
.iter()
.any(|disk| disk.mount_point == "/")
);
assert!(super::filesystem_space("/proc/self/not-a-mount").is_none());
assert!(super::filesystem_space("/nul\0path").is_none());
}
#[test]
fn overdue_latency_tasks_wait_and_resume_without_duplicate_probes() {
let (task_tx, task_rx) = std::sync::mpsc::channel();
let (result_tx, result_rx) = std::sync::mpsc::channel();
let mut executor = super::LatencyExecutor {
task_tx,
result_rx,
in_flight: Default::default(),
stats: Default::default(),
};
let task = LatencyTask {
id: "queued-probe".into(),
name: "probe".into(),
task_type: "tcp".into(),
target: "example.com".into(),
port: Some(443),
interval_seconds: 30,
};
let mut deadlines = std::collections::HashMap::new();
let started = Instant::now();
executor.schedule(std::slice::from_ref(&task), &mut deadlines, started);
assert_eq!(task_rx.try_recv().unwrap().id, task.id);
assert_eq!(deadlines[&task.id], started + Duration::from_secs(30));
let overdue = started + Duration::from_secs(30);
executor.schedule(std::slice::from_ref(&task), &mut deadlines, overdue);
assert!(task_rx.try_recv().is_err());
assert_eq!(deadlines[&task.id], overdue + Duration::from_secs(1));
result_tx
.send(LatencyResult {
task_id: task.id.clone(),
..LatencyResult::default()
})
.unwrap();
assert_eq!(executor.drain().len(), 1);
let retry_at = deadlines[&task.id];
executor.schedule(std::slice::from_ref(&task), &mut deadlines, retry_at);
executor.schedule(std::slice::from_ref(&task), &mut deadlines, retry_at);
assert_eq!(task_rx.try_recv().unwrap().id, task.id);
assert!(task_rx.try_recv().is_err());
assert_eq!(deadlines[&task.id], retry_at + Duration::from_secs(30));
drop(task_rx);
executor.in_flight.clear();
let retry_at = deadlines[&task.id];
executor.schedule(std::slice::from_ref(&task), &mut deadlines, retry_at);
assert_eq!(deadlines[&task.id], retry_at + Duration::from_secs(1));
}
#[test]
fn parses_release_versions() {
assert_eq!(normalized_version("v1.2.3"), "1.2.3");
assert_eq!(version_triplet("v1.2.3"), Some((1, 2, 3)));
assert_eq!(version_triplet("1.2"), None);
assert_eq!(version_triplet("1.2.3.4"), None);
assert_eq!(version_triplet("rust-3"), None);
}
#[test]
fn validates_release_asset_digests() {
let hash = "0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF0123456789ABCDEF";
let mut asset = GithubReleaseAsset {
name: "agent-linux-x64-glibc".to_string(),
browser_download_url: "https://example.com/agent".to_string(),
digest: Some(format!("sha256:{hash}")),
};
let expected = hash.to_ascii_lowercase();
assert_eq!(
release_asset_sha256(&asset).as_deref(),
Some(expected.as_str())
);
asset.digest = Some("sha512:0123".to_string());
assert_eq!(release_asset_sha256(&asset), None);
asset.digest = None;
assert_eq!(release_asset_sha256(&asset), None);
}
#[test]
fn validates_endpoints() {
assert!(valid_endpoint("https://monitor.example.com/"));
assert!(valid_endpoint("http://127.0.0.1:8787"));
assert!(valid_endpoint("http://localhost:8787"));
assert!(valid_endpoint("http://[::1]:8787"));
assert!(!valid_endpoint("http://monitor.example.com"));
assert!(!valid_endpoint("http://127.0.0.2:8787"));
assert!(!valid_endpoint("monitor.example.com"));
assert!(!valid_endpoint("https://"));
assert!(!valid_endpoint("https://user@example.com"));
assert!(!valid_endpoint("https://monitor.example.com/?token=abc"));
assert!(!valid_endpoint("https://bad host.example"));
}
#[test]
fn builds_live_websocket_endpoints() {
assert_eq!(
live_endpoint("https://monitor.example.com").unwrap(),
"wss://monitor.example.com/api/agent/ws"
);
assert_eq!(
live_endpoint("https://monitor.example.com/base/").unwrap(),
"wss://monitor.example.com/base/api/agent/ws"
);
assert_eq!(
live_endpoint("http://127.0.0.1:8787").unwrap(),
"ws://127.0.0.1:8787/api/agent/ws"
);
}
#[test]
fn accepts_durable_ack_interval_and_limits_upload_frequency() {
let ack: LiveAck = serde_json::from_str(
r#"{"type":"ack","ts":100,"persisted":true,"persistenceError":false,"persistedThroughTs":100,"nextPersistAfterMs":60000}"#,
)
.unwrap();
assert_eq!(ack_persist_interval(&ack), Duration::from_secs(60));
let idle: LiveAck = serde_json::from_str(
r#"{"type":"ack","ts":100,"persisted":true,"persistenceError":false,"persistedThroughTs":100,"nextPersistAfterMs":120000}"#,
)
.unwrap();
assert_eq!(ack_persist_interval(&idle), Duration::from_secs(120));
assert_eq!(super::live_batch_interval(1), Duration::from_secs(3));
assert_eq!(super::live_batch_interval(3), Duration::from_secs(3));
assert_eq!(super::live_batch_interval(15), Duration::from_secs(15));
}
#[cfg(unix)]
#[test]
fn remote_commands_capture_output_and_exit_status() {
let result = execute_remote_task_with_timeout(
&RemoteTaskMessage {
message_type: "remote_task".to_string(),
task_id: "task-output".to_string(),
command: "printf stdout; printf stderr >&2; exit 7".to_string(),
},
Duration::from_secs(2),
);
assert_eq!(result.status, "failed");
assert_eq!(result.exit_code, Some(7));
assert_eq!(result.result, "stdout\nstderr");
}
#[cfg(unix)]
#[test]
fn remote_commands_are_terminated_after_the_deadline() {
let result = execute_remote_task_with_timeout(
&RemoteTaskMessage {
message_type: "remote_task".to_string(),
task_id: "task-timeout".to_string(),
command: "sleep 5".to_string(),
},
Duration::from_millis(50),
);
assert_eq!(result.status, "failed");
assert_eq!(result.exit_code, Some(-1));
assert!(result.result.contains("已终止"));
}
#[test]
fn combined_remote_output_is_bounded_at_utf8_boundaries() {
let result = remote_result_text(
&CapturedOutput {
bytes: vec![b'x'; REMOTE_STREAM_OUTPUT_BYTES as usize],
truncated: true,
},
&CapturedOutput {
bytes: "输出"
.repeat(REMOTE_STREAM_OUTPUT_BYTES as usize / 6)
.into_bytes(),
truncated: false,
},
);
assert!(result.len() < REMOTE_RESULT_OUTPUT_BYTES + 64);
assert!(result.ends_with("[输出已截断]"));
assert!(!result.contains('\u{fffd}'));
}
#[cfg(unix)]
#[test]
fn verbose_remote_commands_finish_after_output_is_truncated() {
for command in [
"set -e; head -c 1048576 /dev/zero; printf completed >&2",
"set -e; head -c 1048576 /dev/zero >&2; printf completed",
] {
let result = execute_remote_task_with_timeout(
&RemoteTaskMessage {
message_type: "remote_task".to_string(),
task_id: "verbose-output".to_string(),
command: command.to_string(),
},
Duration::from_secs(5),
);
assert_eq!(result.status, "success");
assert_eq!(result.exit_code, Some(0));
assert!(result.result.contains("completed"));
assert!(result.result.ends_with("[输出已截断]"));
assert!(result.result.len() < REMOTE_RESULT_OUTPUT_BYTES + 64);
}
}
fn temporary_remote_task_journal(label: &str) -> PathBuf {
let nonce = SystemTime::now()
.duration_since(UNIX_EPOCH)
.unwrap()
.as_nanos();
let directory = std::env::temp_dir().join(format!(
"nodeflare-agent-{label}-{}-{nonce}",
std::process::id()
));
fs::create_dir_all(&directory).unwrap();
directory.join("remote-tasks.json")
}
#[test]
fn interrupted_remote_tasks_are_failed_without_reexecution() {
let path = temporary_remote_task_journal("interrupted-task");
let task_id = "8dd70536-f721-4d47-af63-e2f17c83de25";
write_remote_task_journal(
&path,
vec![RemoteTaskJournalEntry::Active {
task_id: task_id.to_string(),
}],
)
.unwrap();
let mut executor = RemoteExecutor::new_at(path.clone()).unwrap();
assert!(executor.active.is_empty());
let due = executor.due_results();
assert_eq!(due.len(), 1);
assert_eq!(due[0].task_id, task_id);
assert_eq!(due[0].status, "failed");
assert!(due[0].result.contains("不会再次执行"));
executor.mark_sent(task_id);
assert!(executor.due_results().is_empty());
executor.reset_delivery();
assert_eq!(executor.due_results().len(), 1);
executor.acknowledge(task_id);
assert!(!path.exists());
let directory = path.parent().unwrap().to_path_buf();
drop(executor);
fs::remove_dir_all(directory).unwrap();
}
#[test]
fn completed_remote_tasks_return_the_cached_result() {
let path = temporary_remote_task_journal("completed-task");
let task_id = "6cbf4d36-f33a-40ef-b952-6af1b8491820";
let result = TaskResultMessage {
message_type: "task_result".to_string(),
task_id: task_id.to_string(),
status: "success".to_string(),
result: "already finished".to_string(),
exit_code: Some(0),
};
write_remote_task_journal(
&path,
vec![RemoteTaskJournalEntry::Completed {
result: result.clone(),
}],
)
.unwrap();
let mut executor = RemoteExecutor::new_at(path.clone()).unwrap();
assert!(
executor
.enqueue(RemoteTaskMessage {
message_type: "remote_task".to_string(),
task_id: task_id.to_string(),
command: "this command must not run".to_string(),
})
.is_none()
);
assert!(executor.active.is_empty());
assert_eq!(executor.due_results(), [result]);
executor.acknowledge(task_id);
let directory = path.parent().unwrap().to_path_buf();
drop(executor);
fs::remove_dir_all(directory).unwrap();
}
#[test]
fn encodes_realtime_samples_as_a_batch() {
let payload = live_update_payload(
vec![
Report {
timestamp: 10,
cpu: 20.0,
..Report::default()
},
Report {
timestamp: 15,
cpu: 30.0,
..Report::default()
},
],
false,
&mut None,
)
.unwrap();
let value: serde_json::Value = nodeflare_telemetry::decode(&payload).unwrap();
assert_eq!(value["persist"], false);
assert_eq!(value["samples"].as_array().unwrap().len(), 2);
assert_eq!(value["samples"][1]["metrics"]["cpu"], 30.0);
}
#[test]
fn calibrates_timestamps_from_http_date() {
let server_ms = 784_111_777_000_i64;
assert_eq!(
clock_offset_from_http_date(
"Sun, 06 Nov 1994 08:49:37 GMT",
server_ms - 1_000,
server_ms + 1_000,
),
Some(0),
);
assert_eq!(corrected_timestamp(100, 2_500), 102);
assert_eq!(corrected_timestamp(100, -2_500), 97);
assert_eq!(monotonic_report_timestamp(97, 100, 90), 101);
assert_eq!(monotonic_report_timestamp(110, 100, 120), 121);
assert_eq!(monotonic_report_timestamp(130, 120, 110), 130);
}
#[test]
fn ignores_small_clock_jitter_until_the_calibration_expires() {
let mut calibration = ClockCalibration::default();
assert!(calibration.observe(1_000));
assert!(!calibration.observe(2_000));
assert_eq!(calibration.offset_ms, 1_000);
assert!(calibration.observe(25_000));
calibration.calibrated_at = Some(Instant::now() - CLOCK_CALIBRATION_MAX_AGE);
assert!(calibration.observe(25_001));
}
#[test]
fn parses_cli_options() {
let parsed = CliOptions::try_parse_from([
"agent",
"-e",
"https://monitor.example.com",
"-t",
"agent-token",
"-i",
"60",
"--once",
])
.unwrap();
assert_eq!(parsed.endpoint, "https://monitor.example.com");
assert_eq!(parsed.token.as_deref(), Some("agent-token"));
assert_eq!(parsed.interval, 60);
assert!(parsed.once);
assert!(CliOptions::try_parse_from(["agent", "-t"]).is_err());
assert!(CliOptions::try_parse_from(["agent", "-t", "first", "-t", "second"]).is_err());
assert!(
CliOptions::try_parse_from([
"agent",
"-e",
"https://monitor.example.com",
"-t",
"agent-token",
"--token-file",
"/run/nodeflare/token",
])
.is_err()
);
assert!(CliOptions::try_parse_from(["agent", "--once", "--collect"]).is_err());
}
#[test]
fn filters_and_deduplicates_remote_latency_tasks_without_a_count_limit() {
let valid = LatencyTask {
id: "task-1".to_string(),
name: "Cloudflare".to_string(),
task_type: "tcp".to_string(),
target: "1.1.1.1".to_string(),
port: Some(443),
interval_seconds: 60,
};
let mut tasks = vec![valid.clone(), valid];
tasks.push(LatencyTask {
id: "bad".to_string(),
name: "Bad".to_string(),
task_type: "http".to_string(),
target: "https://example.com".to_string(),
port: Some(443),
interval_seconds: 1,
});
for index in 2..=300 {
tasks.push(LatencyTask {
id: format!("task-{index}"),
name: format!("Task {index}"),
task_type: "icmp".to_string(),
target: "1.1.1.1".to_string(),
port: None,
interval_seconds: 60,
});
}
let sanitized = sanitize_latency_tasks(&tasks);
assert_eq!(sanitized.len(), 300);
assert_eq!(sanitized[0].id, "task-1");
assert!(!sanitized.iter().any(|task| task.id == "bad"));
}
#[test]
fn advances_collection_deadlines_without_drift() {
let start = Instant::now();
let interval = Duration::from_secs(1);
assert_eq!(
advance_deadline(start, interval, start + Duration::from_millis(2500)),
start + Duration::from_secs(3)
);
}
#[test]
fn drops_samples_the_server_would_reject_as_expired() {
let mut samples = vec![
Report {
timestamp: 1_000,
..Report::default()
},
Report {
timestamp: 8_500,
..Report::default()
},
];
prune_report_samples(&mut samples, 9_000);
assert_eq!(samples.len(), 1);
assert_eq!(samples[0].timestamp, 8_500);
}
#[test]
fn bounds_pending_latency_results() {
let mut samples = (0..MAX_PENDING_LATENCY_RESULTS + 10)
.map(|index| Report {
timestamp: index as i64 + 1,
latency_results: vec![LatencyResult {
task_id: format!("task-{index}"),
timestamp: index as i64 + 1,
latency_ms: 10.0,
packet_loss: 0.0,
}],
..Report::default()
})
.collect::<Vec<_>>();
prune_report_samples(&mut samples, MAX_PENDING_LATENCY_RESULTS as i64 + 10);
assert_eq!(
samples
.iter()
.map(|sample| sample.latency_results.len())
.sum::<usize>(),
MAX_PENDING_LATENCY_RESULTS
);
assert!(samples.first().unwrap().latency_results.is_empty());
assert_eq!(samples.last().unwrap().latency_results.len(), 1);
}
#[cfg(any(target_os = "windows", target_os = "macos", target_os = "freebsd"))]
#[test]
fn parses_netstat_connection_counts() {
let output = "tcp4 0 0 host.443 peer.1 ESTABLISHED\nudp4 0 0 *.5353 *.*\nTCP host peer ESTABLISHED\n";
assert_eq!(connection_counts_from_netstat(output), (2, 1));
}
#[test]
fn parses_probe_targets() {
assert_eq!(
parse_probe_target("Example.COM", None),
Some(("example.com".to_string(), 443))
);
assert_eq!(
parse_probe_target("1.1.1.1", Some(8080)),
Some(("1.1.1.1".to_string(), 8080))
);
for target in [
"",
"https://example.com",
"example.com:0",
"example.com:65536",
"999.1.1.1",
"127.0.0.1:8080",
"169.254.169.254",
"192.168.1.1",
"router.local",
"localhost",
"[::1]:443",
"bad host",
] {
assert_eq!(parse_probe_target(target, None), None, "target: {target}");
}
}
#[test]
fn accepts_only_public_probe_addresses() {
for address in ["1.1.1.1", "8.8.8.8", "2606:4700:4700::1111"] {
assert!(is_public_probe_ip(address.parse().expect("IP address")));
}
for address in [
"0.0.0.0",
"10.0.0.1",
"100.64.0.1",
"127.0.0.1",
"169.254.169.254",
"172.16.0.1",
"192.168.1.1",
"198.18.0.1",
"224.0.0.1",
"::1",
"fe80::1",
"fd00:ec2::254",
] {
assert!(!is_public_probe_ip(address.parse().expect("IP address")));
}
}
#[test]
fn parses_ping_latency() {
assert_eq!(ping_latency("64 bytes time=12.34 ms"), Some(12.34));
assert_eq!(ping_latency("64 bytes time<1 ms"), Some(0.5));
assert_eq!(ping_latency("unreachable"), None);
assert_eq!(
ping_latencies("reply time=12.34 ms\nrequest timeout\nreply time=8.5 ms"),
vec![12.34, 8.5]
);
}
#[test]
fn update_check_jitter_is_stable_and_bounded() {
let first = update_check_jitter("agent-token");
assert_eq!(first, update_check_jitter("agent-token"));
assert!(first.as_secs() <= UPDATE_CHECK_JITTER_MAX_SECONDS);
}
agent/src/update.rs
Executable
+240
View File
@@ -0,0 +1,240 @@
use super::*;
#[derive(Debug, Deserialize)]
pub(crate) struct GithubRelease {
pub(crate) tag_name: String,
pub(crate) assets: Vec<GithubReleaseAsset>,
}
#[derive(Debug, Deserialize)]
pub(crate) struct GithubReleaseAsset {
pub(crate) name: String,
pub(crate) browser_download_url: String,
pub(crate) digest: Option<String>,
}
pub(crate) fn normalized_version(value: &str) -> &str {
let value = value.trim();
value.strip_prefix('v').unwrap_or(value)
}
pub(crate) fn version_triplet(value: &str) -> Option<(u64, u64, u64)> {
let mut parts = normalized_version(value).split('.');
let version = (
parts.next()?.parse().ok()?,
parts.next()?.parse().ok()?,
parts.next()?.parse().ok()?,
);
parts.next().is_none().then_some(version)
}
pub(crate) fn agent_artifact_name() -> Option<&'static str> {
match (env::consts::OS, env::consts::ARCH) {
("linux", "x86_64") => Some(if cfg!(target_env = "musl") {
"agent-linux-x64-musl"
} else {
"agent-linux-x64-glibc"
}),
("linux", "aarch64") => Some(if cfg!(target_env = "musl") {
"agent-linux-aarch64-musl"
} else {
"agent-linux-aarch64-glibc"
}),
("windows", "x86_64") => Some("agent-windows-x64.exe"),
("macos", "aarch64") => Some("agent-macos-aarch64"),
("freebsd", "x86_64") => Some("agent-freebsd-x64"),
("freebsd", "aarch64") => Some("agent-freebsd-aarch64"),
_ => None,
}
}
pub(crate) fn executable_format_valid(path: &Path) -> bool {
let mut magic = [0_u8; 4];
if fs::File::open(path)
.and_then(|mut file| file.read_exact(&mut magic))
.is_err()
{
return false;
}
match env::consts::OS {
"linux" | "freebsd" => magic == *b"\x7fELF",
"windows" => magic.starts_with(b"MZ"),
"macos" => matches!(
magic,
[0xcf, 0xfa, 0xed, 0xfe]
| [0xfe, 0xed, 0xfa, 0xcf]
| [0xca, 0xfe, 0xba, 0xbe]
| [0xbe, 0xba, 0xfe, 0xca]
),
_ => false,
}
}
pub(crate) fn sha256_file(path: &Path) -> Result<String> {
let mut file = fs::File::open(path)?;
let mut digest = Sha256::new();
let mut buffer = [0_u8; 64 * 1024];
loop {
let read = file.read(&mut buffer)?;
if read == 0 {
break;
}
digest.update(&buffer[..read]);
}
Ok(digest
.finalize()
.iter()
.map(|byte| format!("{byte:02x}"))
.collect())
}
pub(crate) fn release_asset_sha256(asset: &GithubReleaseAsset) -> Option<String> {
let digest = asset.digest.as_deref()?.strip_prefix("sha256:")?;
(digest.len() == 64
&& digest
.chars()
.all(|character| character.is_ascii_hexdigit()))
.then(|| digest.to_ascii_lowercase())
}
pub(crate) fn download_agent(
pub(crate) agent: &ureq::Agent,
pub(crate) url: &str,
pub(crate) destination: &Path,
pub(crate) expected_version: &str,
pub(crate) expected_sha256: &str,
) -> Result<bool> {
let Ok(response) = agent.get(url).call() else {
return Ok(false);
};
let response = response.into_body().into_reader();
let mut file = fs::File::create(destination)?;
let copied = io::copy(&mut response.take(MAX_AGENT_BINARY_BYTES + 1), &mut file)?;
file.flush()?;
drop(file);
if copied > MAX_AGENT_BINARY_BYTES
|| !executable_format_valid(destination)
|| sha256_file(destination)? != expected_sha256
{
let _ = fs::remove_file(destination);
return Ok(false);
}
#[cfg(unix)]
fs::set_permissions(destination, fs::Permissions::from_mode(0o755))?;
let downloaded_version = Command::new(destination)
.arg("--version")
.output()
.ok()
.filter(|output| output.status.success())
.and_then(|output| {
String::from_utf8_lossy(&output.stdout)
.split_whitespace()
.next_back()
.map(str::to_string)
});
let valid = downloaded_version.as_deref().map(normalized_version)
== Some(normalized_version(expected_version));
if !valid {
let _ = fs::remove_file(destination);
}
Ok(valid)
}
pub(crate) fn mirrored_download_url(mirror: &str, url: &str) -> String {
let mirror = mirror.trim().trim_end_matches('/');
if mirror.is_empty() {
url.to_string()
} else {
format!("{mirror}/{url}")
}
}
pub(crate) fn update(agent: &ureq::Agent, mirror: &str) -> Result<bool> {
let mut response = agent
.get(LATEST_RELEASE_API)
.header("Accept", "application/vnd.github+json")
.header("User-Agent", format!("nodeflare-agent/{VERSION}"))
.call()?;
let release = response.body_mut().read_json::<GithubRelease>()?;
let Some(remote_version) = version_triplet(&release.tag_name) else {
return Ok(false);
};
let Some(current_version) = version_triplet(VERSION) else {
return Ok(false);
};
if remote_version <= current_version {
return Ok(false);
}
let Some(artifact) = agent_artifact_name() else {
return Ok(false);
};
let current = env::current_exe()?;
let temporary = current.with_file_name(format!(
".nodeflare-agent.{}.download{}",
std::process::id(),
if cfg!(target_os = "windows") {
".exe"
} else {
""
}
));
let Some(release_asset) = release.assets.iter().find(|asset| asset.name == artifact) else {
return Err(format!("latest release does not contain {artifact}").into());
};
let Some(expected_sha256) = release_asset_sha256(release_asset) else {
return Err(
format!("latest release does not contain a SHA-256 digest for {artifact}").into(),
);
};
if !download_agent(
agent,
&mirrored_download_url(mirror, &release_asset.browser_download_url),
&temporary,
&release.tag_name,
&expected_sha256,
)? {
return Err("downloaded agent version does not match the configured version".into());
}
#[cfg(unix)]
{
fs::rename(&temporary, &current)?;
let error = Command::new(&current).args(env::args_os().skip(1)).exec();
Err(error.into())
}
#[cfg(target_os = "windows")]
{
let script = concat!(
"$ErrorActionPreference='Stop'; ",
"$targetPid=[int]$env:NODEFLARE_UPDATE_PID; ",
"Wait-Process -Id $targetPid; ",
"Move-Item -LiteralPath $env:NODEFLARE_UPDATE_NEW ",
"-Destination $env:NODEFLARE_UPDATE_CURRENT -Force; ",
"try { Start-ScheduledTask -TaskName 'nodeflare-agent' -ErrorAction Stop } ",
"catch { $restartArgs=@($env:NODEFLARE_UPDATE_ARGS | ConvertFrom-Json); ",
"Start-Process -FilePath $env:NODEFLARE_UPDATE_CURRENT -ArgumentList $restartArgs }"
);
Command::new("powershell.exe")
.args([
"-NoLogo",
"-NoProfile",
"-NonInteractive",
"-Command",
script,
])
.env("NODEFLARE_UPDATE_PID", std::process::id().to_string())
.env("NODEFLARE_UPDATE_NEW", &temporary)
.env("NODEFLARE_UPDATE_CURRENT", &current)
.env(
"NODEFLARE_UPDATE_ARGS",
serde_json::to_string(&env::args().skip(1).collect::<Vec<_>>())?,
)
.stdin(Stdio::null())
.stdout(Stdio::null())
.stderr(Stdio::null())
.spawn()?;
Ok(true)
}
}