ostp/ostp-client/src/transport/dns.rs

104 lines
3.9 KiB
Rust

use std::sync::Arc;
use std::time::Duration;
use bytes::Bytes;
use tokio::net::UdpSocket;
use tokio::sync::{mpsc, Mutex};
use rand::Rng;
use ostp_core::dns::{
DnsPacket, DnsRecordType, encode_payload_to_domain, decode_domain_to_payload,
};
use crate::transport::Transport;
pub async fn start_dns_transport(domain: String, resolver: String, _pubkey: Option<String>) -> std::io::Result<Transport> {
let (app_tx, transport_rx) = mpsc::channel::<Bytes>(100);
let (transport_tx, app_rx) = mpsc::channel::<Bytes>(100);
let resolver_addr = if resolver.contains(':') {
resolver.clone()
} else {
format!("{}:53", resolver)
};
let socket = UdpSocket::bind("0.0.0.0:0").await?;
socket.connect(&resolver_addr).await?;
let socket = Arc::new(socket);
let sock_rx = socket.clone();
let sock_tx = socket;
let base_domain = domain.clone();
// Send task (reads from app, encodes into DNS TXT, sends to UDP socket)
tokio::spawn(async move {
let mut rx = transport_rx;
loop {
let data_opt = tokio::select! {
res = rx.recv() => res,
_ = tokio::time::sleep(Duration::from_secs(2)) => Some(Bytes::new()),
};
let data = match data_opt {
Some(d) => d,
None => break, // App closed
};
// Encode data to base32 domain
let fqdn = encode_payload_to_domain(&data, &base_domain);
let id: u16 = rand::thread_rng().gen();
// Randomly choose TXT or NULL for diversity (as requested)
let qtype = if rand::thread_rng().gen_bool(0.5) {
DnsRecordType::TXT
} else {
DnsRecordType::NULL
};
let packet = DnsPacket::new_query(id, &fqdn, qtype);
let encoded = packet.encode();
if let Err(e) = sock_tx.send(&encoded).await {
tracing::warn!("DNS transport send error: {}", e);
break;
}
}
});
// Receive task (reads from UDP socket, decodes DNS answer, sends to app)
let base_domain_rx = domain.clone();
tokio::spawn(async move {
let mut buf = vec![0u8; 65535];
loop {
match sock_rx.recv(&mut buf).await {
Ok(n) => {
if let Some(packet) = DnsPacket::decode(&buf[..n]) {
for answer in packet.answers {
if answer.rtype == DnsRecordType::TXT || answer.rtype == DnsRecordType::NULL {
// If it's a TXT record, the response might be base32 encoded payload?
// Actually, dnstt puts the payload in the TXT/NULL record data.
// We'll just assume the rdata is the raw payload, or base32 encoded if it was sent as such.
// Let's just pass the raw data (TXT strings are decoded in DnsPacket::decode)
// Wait, dnstt server responds with raw bytes in NULL, and base32/chunked strings in TXT.
// Our `DnsPacket::decode` already handles extracting TXT string bytes or NULL raw bytes into `rdata`.
// Let's just send `rdata` to the app.
if transport_tx.send(Bytes::from(answer.rdata)).await.is_err() {
return; // App closed
}
}
}
}
}
Err(e) => {
tracing::warn!("DNS transport recv error: {}", e);
break;
}
}
}
});
Ok(Transport::Dns {
tx: app_tx,
rx: Arc::new(Mutex::new(app_rx)),
})
}