diff --git a/tof/src/main.rs b/tof/src/main.rs index 70a4339a..343ae8fa 100644 --- a/tof/src/main.rs +++ b/tof/src/main.rs @@ -646,10 +646,10 @@ async fn serve( Ok((stream, _)) => { let status = status.clone(); let imu_status = imu_status.clone(); - let frames = frames.subscribe(); - let imu_frames = imu_frames.subscribe(); + let frames = frames.clone(); + let imu_frames = imu_frames.clone(); tokio::spawn(async move { - if let Err(e) = subscriber(stream, &status, frames, &imu_status, imu_frames).await { + if let Err(e) = subscriber(stream, &status, &frames, &imu_status, &imu_frames).await { tracing::debug!(error = %e, "subscriber ended"); } }); @@ -669,12 +669,16 @@ async fn serve( } /// One subscriber: its request, then frames until it goes away. +/// Takes the two `Sender`s and subscribes only in the arm that matched, so a receiver exists +/// where the method is known and nowhere else. Subscribing on `accept` made `receiver_count()` +/// count connections rather than interest: a client asking for `head_imu.stream` held a depth +/// receiver, and one that connected and said nothing held both. async fn subscriber( stream: UnixStream, status: &Arc, - mut frames: tokio::sync::broadcast::Receiver, + frames: &tokio::sync::broadcast::Sender, imu_status: &Arc, - mut imu_frames: tokio::sync::broadcast::Receiver, + imu_frames: &tokio::sync::broadcast::Sender, ) -> Result<()> { let (read, mut write) = stream.into_split(); let mut reader = BufReader::new(read); @@ -704,11 +708,13 @@ async fn subscriber( Ok(proto::Call::TofStream) => { let response = proto::Response::ok(id, &status.result()); write_line(&mut write, &response).await?; + let mut frames = frames.subscribe(); return stream_tof(&mut write, &mut frames).await; } Ok(proto::Call::HeadImuStream) => { let response = proto::Response::ok(id, &imu_status.result()); write_line(&mut write, &response).await?; + let mut imu_frames = imu_frames.subscribe(); return stream_imu(&mut write, &mut imu_frames).await; } _ => { @@ -818,6 +824,77 @@ async fn write_line( mod tests { use super::*; + /// **A connection that has asked for nothing wants nothing**, and one that asked for depth + /// does not want the IMU. `accept` used to subscribe to both channels before reading a byte + /// of the request, so `receiver_count()` counted open connections rather than interest — + /// which is the number anything gating a sensor on demand would have to trust. + /// + /// Drives the real `subscriber` over a socket pair and counts receivers on the senders it + /// holds, since that count is the thing that was wrong. + #[tokio::test] + async fn a_connection_subscribes_only_to_what_it_asked_for() { + let dir = tempfile::tempdir().unwrap(); + let path = dir.path().join("tofd.sock"); + let listener = UnixListener::bind(&path).unwrap(); + + let (frames, _) = tokio::sync::broadcast::channel::(4); + let (imu_frames, _) = tokio::sync::broadcast::channel::(4); + let status = Arc::new(Status::new(15)); + let imu_status = Arc::new(ImuStatus::new(100)); + + let served = { + let (frames, imu_frames) = (frames.clone(), imu_frames.clone()); + tokio::spawn(async move { + let (stream, _) = listener.accept().await.unwrap(); + let _ = subscriber(stream, &status, &frames, &imu_status, &imu_frames).await; + }) + }; + + let client = tokio::net::UnixStream::connect(&path).await.unwrap(); + let (read, mut write) = client.into_split(); + let mut reader = BufReader::new(read); + + // Connected, and nothing asked for yet: neither channel has a receiver. + assert_eq!( + frames.receiver_count(), + 0, + "a silent connection wants no depth" + ); + assert_eq!(imu_frames.receiver_count(), 0, "nor the imu"); + + // Ask for depth only. + let request = format!( + "{}\n", + serde_json::to_string(&proto::Request::call( + proto::Id::Number(1), + &proto::Call::TofStream + )) + .unwrap() + ); + write.write_all(request.as_bytes()).await.unwrap(); + let mut answer = String::new(); + reader.read_line(&mut answer).await.unwrap(); + assert!(answer.contains("result"), "{answer}"); + + // The answer is written before the subscribe, so give the task its next poll. + for _ in 0..100 { + if frames.receiver_count() == 1 { + break; + } + tokio::time::sleep(std::time::Duration::from_millis(2)).await; + } + assert_eq!(frames.receiver_count(), 1, "depth was asked for"); + assert_eq!( + imu_frames.receiver_count(), + 0, + "a client that asked for depth must not hold an imu receiver" + ); + + // `stream_tof` is parked on the channel, not on the socket, so it would sit there until a + // frame arrived. Nothing here sends one, and the counts have already been taken. + served.abort(); + } + #[test] fn addresses_parse_in_both_bases() { assert_eq!(parse_address("0x29"), Ok(0x29));