eculocate/ecu-esp32/src/mdns.rs

// MDNS responder, mostly but not entirely standards-compliant

// ** Probe and Announce at startup

// at startup it does Probe and Announce for "all those resource
// records that a Multicast DNS responder desires to be unique on the
// local link".  "Probe" means send packets as QU (unicast response)
// containing ANY requests for each name: 3 packets with 250ms between
// them; "Announce" means send 2 <= n <= 8 packets one second apart,
// containing all the records we know

// ** additional data in responses

//  * for SRV, we send all the A and AAAA it names

//  * for PTR "Service Instance Enumeration or Selective Instance
//    Enumeration (subtype)": i.e. when the record data names an
//    instance of the service, we sent the SRV and TXT records for the
//    instance _and_ all the A/AAAA named by the SRV

// I haven't investigated what happens if the packet gets too
// large. There may be bugs

// ** don't answer queries if the answers in the query already include what we'd send

// This is unimplemented for now.  mdns is so fricking noisy already I
// don't know if this is worth the extra effort

// ** set the CACHE_FLUSH bit on unique records

// "all the records with that name, rrtype, and rrclass are conceptually
// under the control or ownership of a single responder"

// ** NSEC?

// not done

// ** unicast response requested (QU)

// - top bit in the class field of a DNS question is the unicast-response bit
// - send response unicast, unless the record has not been sent recently
//   ("within one quarter of its TTL")
// - also for queries received fom unicast
//   - but verify the source is local lan
//   - and the port is 5353 (otherwise see "legacy")

//  [ mostly done; we don't check "has been sent recently" and as we
//  didn't open a unicast socket then we can't receive queries from
//  unicast ]


// ** Legacy Unicast Responses

// if source port != 5353

// - respond via unicast to source
// - query id and question
// - don't set cache-flush bit
// - ttl <= 10s

// This is implemented but has not really been tested

// ** response compression

// Unimplemented


use alloc::vec::Vec;
use alloc::string::String;
use alloc::collections::BTreeSet;
use crate::alloc::string::ToString;

use embassy_net::{
    udp::{UdpSocket, UdpMetadata, PacketMetadata},
    IpAddress,
    IpEndpoint,
    Stack
};

use embassy_time::{Timer, Duration, Instant};
use esp_alloc as _;
use esp_backtrace as _;

use log::{info,error};

pub const DISCOVERY_NAME : &str = "_services._dns-sd._udp.local";

use dns_protocol::{Message, Question,
		   MessageType,
		   ResourceRecord, ResourceType, Flags};

const MDNS_IP: IpAddress = IpAddress::v4(224, 0, 0, 251);
const MDNS_PORT: u16 = 5353;
const MDNS_ENDPOINT: IpEndpoint = IpEndpoint::new(MDNS_IP, MDNS_PORT);

const CACHE_FLUSH : u16 = 1<<15;

pub enum ResponderState {
    Probing(Instant, usize),
    ProbeFailed,
    Announcing(Instant, usize),
    Running
}

pub struct Responder<'a> {
    sock : &'a UdpSocket<'a>,
    database : Database,
    pub state : ResponderState,
}

impl<'a> Responder<'a> {
    pub fn new(sock: &'a UdpSocket<'a>,
	       records: Vec<Resource>) -> Responder<'a> {

	Self {
	    sock: sock,
	    database: Database(records),
	    state: ResponderState::Probing(Instant::now(), 0)
	}
    }

    async fn advance_state(&mut self) {
	match self.state {
	    ResponderState::Probing(next_when, number_sent) => {
		let ts = Instant::now();
		if next_when < ts {
		    if number_sent >= 3 {
			self.state = ResponderState::Announcing(ts, 0);
		    } else {
			self.send_probe().await;
			self.state = ResponderState::Probing(ts + Duration::from_millis(250), number_sent + 1);
		    }
		}
	    },
	    ResponderState::Announcing(next_when, number_sent) => {
		let ts = Instant::now();
		if next_when < ts {
		    if number_sent > 6 {
			self.state = ResponderState::Running;
		    } else {
			self.advertise().await;
			self.state = ResponderState::Announcing(ts + Duration::from_millis(1000), number_sent + 1);
		    }
		}
	    },
	    ResponderState::Running => (),
	    ResponderState::ProbeFailed => (),
	}
    }

    async fn send_message(&self, message: Message<'_, '_>,  endpoint: IpEndpoint) {
	let mut data_buf = [0; 4096];
	// Serialize the message into the buffer
	assert!(message.space_needed() <= data_buf.len());

	match message.write(&mut data_buf) {
	    Result::Ok(len) => {
		if let Result::Err(x) = self.sock.send_to(&data_buf[..len], endpoint).await {
		    info!("sock.send failed: {:?}", x);
		}
	    },
	    Result::Err(err) => {
		error!("message.write {:?}", err);
	    }
	};
    }

    async fn send_probe(&self) {
	info!("sending probe at {}", Instant::now().as_millis());

	let db = &self.database;

	let mut questions : Vec<Question> = db.advertised().iter().filter_map( |record| {
	    // FIXME this sends duplicates because there's both an SRV
	    // and a TXT for the same name
	    if record.is_unique() {
		Some(
		    Question::new(
			&record.name[..],
			ResourceType::Wildcard,
			1 | (1<<15)	// QU
		    ))
	    } else {
		None
	    }
	}).collect();

	let message = Message::new(
	    0, Flags::standard_query(), &mut questions, &mut [], &mut [], &mut []
	);

	self.send_message(message, MDNS_ENDPOINT).await;
    }

    async fn advertise(&self) {
	info!("advertise at {}", Instant::now().as_millis());
	self.send_response(0, MDNS_ENDPOINT, None, self.database.advertised()).await;
    }

    pub async fn may_recv(&mut self) -> bool {
	self.advance_state().await;
	match self.state {
	    ResponderState::ProbeFailed => false,
	    _ => self.sock.may_recv()
	}
    }

    async fn send_response(&self,
			   id: u16,
			   endpoint: IpEndpoint,
			   question: Option<&Question<'_>>,
			   answers: Vec<&Resource>) {

	let legacy = question != None;

	let db = &self.database;
	let mut flags = Flags::new();

	flags.set_qr(MessageType::Reply);
	flags.set_authoritative(true);

	let mut additional : BTreeSet<ResourceRecord> = BTreeSet::new();
	let mut answers : Vec<ResourceRecord> = {
	    let mut rrs = Vec::new();
	    for a in answers {
		rrs.push(a.to_resource_record());
		for r in db.all_additional(a).into_iter() {
		    additional.insert(r.to_resource_record());
		};
	    }
	    rrs
	};

	// don't send anything in "additional" that's already in the answers
	for rr in &answers {
	    additional.remove(&rr);
	}

	let mut additional : Vec<ResourceRecord> = additional.into_iter().collect();

	let questions = match question {
	    Some(q) =>&mut [*q][..],
	    None => &mut [][..]
	};

	if legacy {
	    // force class = IN (no cache-flush) and ttl to 10 seconds
	    answers = answers.iter().map(|rr| {
		ResourceRecord::new(rr.name(), rr.ty(), 1, 10, rr.data())
	    }).collect();
	    additional = additional.iter().map(|rr| {
		ResourceRecord::new(rr.name(), rr.ty(), 1, 10, rr.data())
	    }).collect();
	}

	let message = Message::new(
	    id,
	    flags,
	    questions,
	    &mut answers,
	    &mut [],
	    &mut additional);

	self.send_message(message, endpoint).await;
    }

    async fn process_received_message(&mut self,
			     message: &Message<'_,'_>,
			     meta: &UdpMetadata) {

	match self.state {
	    ResponderState::Probing(_,_) => {
		for rr in message.answers() {
		    info!("received answer {:?}", rr.name());
		    if self.database.find(rr.name().to_string(), ResourceType::Wildcard)
			.iter().any(|record| record.is_unique()) {
			    info!("probe failed, received answer for {:?}", rr.name());
			    self.state = ResponderState::ProbeFailed;
			}
		}
	    },
	    ResponderState::ProbeFailed => (),
	    _ => {
		let q_id = message.id();

		for q in message.questions() {
		    let relevant = self.database.find(q.name().to_string(), q.ty());
		    if relevant.len() > 0 {
			log::info!("question {:?} {:?} {:?} flags {:?} from {:?}",
				   q.name(), q.ty(), q.class(),
				   message.flags().raw(),
				   meta.endpoint);

			let unicast = (q.class() & (1<<15)) != 0;
			let legacy = meta.endpoint.port != 5353;
			if legacy {
			    self.send_response(q_id, meta.endpoint, Some(q), relevant).await
			} else if unicast {
			    self.send_response(0, meta.endpoint, None, relevant).await
			} else {
			    self.send_response(0, MDNS_ENDPOINT, None, relevant).await
			}
		    }
		}
	    }
	}
    }

    pub async fn process_packet(&mut self) {
	let mut questions = [Question::default(); 64];
	let mut answers = [ResourceRecord::default(); 64];
	let mut authorities = [ResourceRecord::default(); 64];
	let mut addl = [ResourceRecord::default(); 64];

	let mut data_buf = [0; 4096];

        let res = self.sock.recv_from(&mut data_buf).await;
        if let Ok((n, meta)) = res {
	    let message = Message::read(
		&data_buf[..n],
		&mut questions,
		&mut answers,
		&mut authorities,
		&mut addl,
	    );

	    match message {
		Result::Ok(message) => {
		    self.process_received_message(&message, &meta).await;
		},
		Result::Err(err) => {
		    log::error!("parsing dns message from {:?} : {:?}",
				meta.endpoint,
				err);
		}
	    }
        }
    }
}


// dns_protocol's ResourceRecord doesn't make a good representation for
// our needs because it doesn't own its "data" attribute - it's just a
// reference.  So we need structures that can own the data

#[derive(Clone,Debug)]
pub enum RData {
    Srv { target: String, port: u16, priority: u16, weight: u16 },
    Ptr(String),
    A([u8; 4]),
    Txt(Vec<String>)
}

impl RData {
    fn copy_string<'a>(s: &str, out: &'a mut [u8], offset: usize) -> usize {
	let end = s.len() + offset + 1;
	out[offset] = (s.len() & 0xff) as u8;
	let _ = &out[(offset+1) .. end].clone_from_slice(s.as_bytes());
	end
    }

    fn encode_as_name<'a>(name: &str, out: &'a mut [u8], offset: usize) -> &'a [u8] {
	let mut offset = offset;

	for l in name.split('.') {
	    offset = Self::copy_string(l, out, offset);
	}
	out[offset] = 0;
	&out[..offset+1]
    }

    fn encode_as_strings<'a>(strings: &[String], out: &'a mut [u8], offset: usize) -> &'a [u8] {
	let mut offset = offset;

	for l in strings.into_iter() {
	    offset = Self::copy_string(l, out, offset);
	}
	out[offset] = 0;
	&out[..offset+1]
    }

    fn ty(&self) -> ResourceType {
	match self {
	    RData::Srv { .. } => ResourceType::Srv,
	    RData::Ptr(_name) => ResourceType::Ptr,
	    RData::A(_octets) => ResourceType::A,
	    RData::Txt(_strings) => ResourceType::Txt
	}
    }

    fn to_bytes<'a>(&self, out: &'a mut [u8]) -> &'a [u8]  {
	match self {
	    RData::Srv { target, port, priority, weight } => {
		out[0] = (priority >> 8) as u8;
		out[1] = (priority & 0xff) as u8;
		out[2] = (weight >> 8) as u8;
		out[3] = (weight & 0xff) as u8;
		out[4] = (port >> 8) as u8;
		out[5] = (port & 0xff) as u8;
		Self::encode_as_name(target, out, 6)
	    },
	    RData::Ptr(name) => {
		Self::encode_as_name(name, out, 0)
	    }
	    RData::A(octets) => {
		out[..4].clone_from_slice(octets);
		&out[..4]
	    },
	    RData::Txt(strings) => {
		Self::encode_as_strings(strings, out, 0)
	    }
	}
    }
}

#[derive(Debug)]
pub struct Resource {
    name : String,
    data : RData,
    data_raw : Vec<u8>
}

impl Resource {
    pub fn new(
	name: impl Into<String>,
	data: RData
    ) -> Self {
	let mut buf : [u8; 80] = [0; 80];
	let rdata_raw = Vec::from(data.to_bytes(&mut buf));
	Self {
	    name: name.into(),
	    data,
	    data_raw: rdata_raw
	}
    }

    // Should we send cache-flush with this record? If it's "unique",
    // which means "all the records with that name, rrtype, and
    // rrclass are conceptually under the control or ownership of a
    // single responder". This works out (perhaps by coincidence) to
    // be everything except the PTR records
    fn is_unique(&self) -> bool {
	self.data.ty() != ResourceType::Ptr
    }
    fn class(&self) -> u16 {
	if self.is_unique() {
	    1 | CACHE_FLUSH
	} else {
	    1
	}
    }

    fn to_resource_record(&self) -> ResourceRecord<'_> {
	let ttl = match self.data.ty() {
	    ResourceType::A => 120,
	    ResourceType::Srv => 120,
	    _ => 4500
	};

	ResourceRecord::new(
	    &self.name[..],
	    self.data.ty(),
	    self.class(),
	    ttl,
	    &self.data_raw[..]
	)
    }
}

struct Database(Vec<Resource>);

impl Database {
    fn find<'a>(&self, name: impl Into<String>, ty: ResourceType) -> Vec<&Resource> {
	let name = name.into();
	let results : Vec<&Resource> =
	    self.0.iter().filter(
		|r|
		((ty == ResourceType::Wildcard) || (r.data.ty() == ty)) &&
		    name == r.name
	    ).collect();
	results
    }

    // the records we announce unsolicited at startup
    fn advertised(&self) -> Vec<&Resource> {
	self.0.iter().filter( |r| r.name != DISCOVERY_NAME ).collect()
    }

    fn additional(&self, r: &Resource) -> Vec<&Resource> {
	match &r.data {
	    RData::Srv { target, port:_, priority:_, weight:_ } => {
		self.find(target, ResourceType::Wildcard)
	    },
	    RData::Ptr(name) => {
		if r.name.eq(DISCOVERY_NAME) {
		    Vec::from([])
		} else {
		    self.find(name, ResourceType::Wildcard)
		}
	    },
	    _ =>  Vec::from([])
	}
    }

    fn all_additional_<'a>(&'a self, r: &Resource, collect: &mut Vec<&'a Resource>) {
	for r in self.additional(r) {
	    collect.push(r);
	    self.all_additional_(r, collect);
	}
    }

    fn all_additional(&self, r: &Resource) -> Vec<&Resource> {
	let mut collect : Vec<&Resource>  = Vec::new();
	self.all_additional_(r, &mut collect);
	collect
    }
}

pub async fn responder(stack: Stack<'_>, records: Vec<Resource>)  {
    let mut rx_buffer = [0u8; 4096];
    let mut tx_buffer = [0u8; 4096];
    let mut rx_meta = [PacketMetadata::EMPTY; 512];
    let mut tx_meta = [PacketMetadata::EMPTY; 512];

    let mut sock = UdpSocket::new(
        stack,
        &mut rx_meta,
        &mut rx_buffer,
        &mut tx_meta,
        &mut tx_buffer,
    );
    sock.set_hop_limit(Some(255));
    log::info!("sock.bind(5353) res: {:?}", sock.bind(MDNS_PORT));
    log::info!(
        "multicast_res: {:?}",
        stack.join_multicast_group(MDNS_IP)
    );

    let mut mdns_responder = Responder::new(&mut sock, records);

    loop {
	match mdns_responder.state {
	    ResponderState::ProbeFailed =>  {
		sock.close();
		_ = stack.leave_multicast_group(MDNS_IP);
		return ;
	    },
	    _ => {
		if mdns_responder.may_recv().await {
		    mdns_responder.process_packet().await;
		} else {
		    Timer::after(Duration::from_millis(50)).await;
		}
	    }
	}
    }
}