eculocate

mdns.rs (14415B)

---

 // 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;
 		}
 	    }
 	}
     }
 }