Skip to main content

manycastr/worker/outbound/
probe.rs

1use crate::custom_module::manycastr::{Address, ProtocolType};
2use crate::net::packet::{DnsProbeId, ProbePayload, create_dns, create_icmp, create_tcp};
3use crate::worker::outbound::{DISCOVERY_WORKER_ID_OFFSET, OutboundConfig, send_packet};
4use log::warn;
5use ratelimit_meter::{DirectRateLimiter, LeakyBucket, NonConformance};
6use socket2::Socket;
7use std::thread::sleep;
8use std::time::{Duration, Instant};
9
10/// Sends probes to the specified destination using the provided measurement configuration.
11/// This function constructs the appropriate packet based on the measurement type
12/// and sends it through the provided socket.
13///
14/// # Arguments
15/// * `config` - The outbound configuration containing worker details and settings.
16/// * `dst` - The destination address to which the probes will be sent.
17/// * `session_id` - 16-bit session ID encoded in the probe (ICMP/DNS-A only)
18/// * `socket` - Raw socket to send packets.
19/// * `limiter` - A rate limit bucket to control the sending rate of packets.
20/// * `is_discovery` - A boolean indicating whether the probes are for discovery purposes.
21///
22/// # Returns
23/// A tuple containing the number of successfully sent packets and the number of failed sends.
24pub fn send_probe(
25    config: &OutboundConfig,
26    dst: &Address,
27    session_id: u32,
28    socket: &Socket,
29    limiter: &mut DirectRateLimiter<LeakyBucket>,
30    is_discovery: bool,
31    packet_buffer: &mut Vec<u8>,
32) -> (u32, u32) {
33    let worker_id = if is_discovery {
34        config.worker_id as u32 + DISCOVERY_WORKER_ID_OFFSET // Use a different worker ID range for discovery probes
35    } else {
36        config.worker_id as u32
37    };
38
39    let mut sent = 0;
40    let mut failed = 0;
41
42    let probe_id = crate::probe_id(config.m_id, session_id);
43    let icmp_payload = ProbePayload {
44        worker_id,
45        probe_id,
46        trace_ttl: None,
47        info_url: config.info_url.as_deref(),
48    };
49
50    // Rate limit
51    if let Err(not_until) = limiter.check() {
52        let wait_time = not_until.wait_time_from(Instant::now());
53        if wait_time > Duration::ZERO {
54            sleep(wait_time);
55        }
56    }
57
58    packet_buffer.clear();
59
60    match config.p_type {
61        ProtocolType::Icmp => {
62            packet_buffer.extend_from_slice(&create_icmp(
63                &config.src,
64                dst,
65                config.dport, // ICMP identifier
66                2,            // ICMP seq
67                &icmp_payload,
68                255,
69            ));
70        }
71        ProtocolType::ADns | ProtocolType::ChaosDns => {
72            let dns_id = DnsProbeId {
73                worker_id,
74                probe_id,
75            };
76            packet_buffer.extend_from_slice(&create_dns(
77                &config.src,
78                dst,
79                config.sport,
80                &dns_id,
81                config.p_type == ProtocolType::ChaosDns,
82                config.qname.as_deref().expect("qname missing"),
83            ));
84        }
85        ProtocolType::Tcp => {
86            packet_buffer.extend_from_slice(&create_tcp(
87                &config.src,
88                dst,
89                config.sport,
90                config.dport,
91                worker_id,
92                is_discovery,
93                config.info_url.as_deref(),
94            ));
95        }
96    }
97
98    match send_packet(socket, packet_buffer, dst) {
99        Ok(()) => sent += 1,
100        Err(e) => {
101            warn!(
102                "[Worker outbound] Failed to send {} packet: {e}",
103                config.p_type
104            );
105            failed += 1;
106        }
107    }
108
109    (sent, failed)
110}