Merge pull request #96 from sgodin/fixup-subscribe-alias-handling

englishm · web-flow · commit f488b5c4f48c · 2025-10-30T00:39:11.000-04:00
diff --git a/moq-transport/src/data/mod.rs b/moq-transport/src/data/mod.rs
@@ -1,11 +1,11 @@
 mod datagram;
 mod fetch;
 mod header;
-mod object;
+mod object_status;
 mod subgroup;
 
 pub use datagram::*;
 pub use fetch::*;
 pub use header::*;
-pub use object::*;
+pub use object_status::*;
 pub use subgroup::*;
diff --git a/moq-transport/src/data/object.rs b/moq-transport/src/data/object.rs
diff --git a/moq-transport/src/data/object_status.rs b/moq-transport/src/data/object_status.rs
@@ -0,0 +1,29 @@
+use crate::coding::{Decode, DecodeError, Encode, EncodeError};
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum ObjectStatus {
+    NormalObject = 0x0,
+    ObjectDoesNotExist = 0x1,
+    EndOfGroup = 0x3,
+    EndOfTrack = 0x4,
+}
+
+impl Decode for ObjectStatus {
+    fn decode<B: bytes::Buf>(r: &mut B) -> Result<Self, DecodeError> {
+        match u64::decode(r)? {
+            0x0 => Ok(Self::NormalObject),
+            0x1 => Ok(Self::ObjectDoesNotExist),
+            0x3 => Ok(Self::EndOfGroup),
+            0x4 => Ok(Self::EndOfTrack),
+            _ => Err(DecodeError::InvalidObjectStatus),
+        }
+    }
+}
+
+impl Encode for ObjectStatus {
+    fn encode<W: bytes::BufMut>(&self, w: &mut W) -> Result<(), EncodeError> {
+        let val = *self as u64;
+        val.encode(w)?;
+        Ok(())
+    }
+}
diff --git a/moq-transport/src/serve/track.rs b/moq-transport/src/serve/track.rs
@@ -173,7 +173,7 @@ impl TrackReader {
     }
 
     // Returns the largest group/sequence
-    pub fn latest(&self) -> Option<Location> {
+    pub fn largest(&self) -> Option<Location> {
         // We don't even know the mode yet.
         // TODO populate from SUBSCRIBE_OK
         None
diff --git a/moq-transport/src/session/publisher.rs b/moq-transport/src/session/publisher.rs
@@ -180,7 +180,7 @@ impl Publisher {
             .ok_or(ServeError::NotFound)?;
         let response;
 
-        if let Some(latest) = track.latest() {
+        if let Some(latest) = track.largest() {
             response = message::TrackStatusOk {
                 id: track_status_request.request_msg.id,
                 track_alias: 0, // TODO SLG - wire up track alias logic
@@ -363,20 +363,41 @@ impl Publisher {
         Ok(())
     }
 
-    pub(super) fn send_message<T: Into<message::Publisher> + Into<Message>>(&mut self, msg: T) {
+    /// Process a message before sending it, performing any necessary internal actions.
+    fn act_on_message_to_send<T: Into<message::Publisher>>(
+        &mut self,
+        msg: T,
+    ) -> message::Publisher {
         let msg = msg.into();
         match &msg {
-            message::Publisher::PublishDone(msg) => self.drop_subscribe(msg.id),
-            message::Publisher::SubscribeError(msg) => self.drop_subscribe(msg.id),
-            message::Publisher::PublishNamespaceDone(msg) => {
-                self.drop_publish_namespace(&msg.track_namespace)
+            message::Publisher::PublishDone(m) => self.drop_subscribe(m.id),
+            message::Publisher::SubscribeError(m) => self.drop_subscribe(m.id),
+            message::Publisher::PublishNamespaceDone(m) => {
+                self.drop_publish_namespace(&m.track_namespace);
             }
-            _ => (),
-        };
+            _ => {}
+        }
+        msg
+    }
 
+    /// Send a message without waiting for it to be sent.
+    pub(super) fn send_message<T: Into<message::Publisher> + Into<Message>>(&mut self, msg: T) {
+        let msg = self.act_on_message_to_send(msg);
         self.outgoing.push(msg.into()).ok();
     }
 
+    /// Send a message and wait until it is sent (or at least popped off the outgoing control message queue)
+    pub(super) async fn send_message_and_wait<T: Into<message::Publisher> + Into<Message>>(
+        &mut self,
+        msg: T,
+    ) {
+        let msg = self.act_on_message_to_send(msg);
+        self.outgoing
+            .push_and_wait_until_popped(msg.into())
+            .await
+            .ok();
+    }
+
     fn drop_subscribe(&mut self, id: u64) {
         self.subscribeds.lock().unwrap().remove(&id);
     }
diff --git a/moq-transport/src/session/subscribe.rs b/moq-transport/src/session/subscribe.rs
@@ -19,13 +19,15 @@ pub struct SubscribeInfo {
 
 struct SubscribeState {
     ok: bool,
+    track_alias: Option<u64>,
     closed: Result<(), ServeError>,
 }
 
 impl Default for SubscribeState {
     fn default() -> Self {
         Self {
             ok: Default::default(),
+            track_alias: None,
             closed: Ok(()),
         }
     }
@@ -121,19 +123,25 @@ pub(super) struct SubscribeRecv {
 }
 
 impl SubscribeRecv {
-    pub fn ok(&mut self) -> Result<(), ServeError> {
+    pub fn ok(&mut self, alias: u64) -> Result<(), ServeError> {
         let state = self.state.lock();
         if state.ok {
             return Err(ServeError::Duplicate);
         }
 
         if let Some(mut state) = state.into_mut() {
             state.ok = true;
+            state.track_alias = Some(alias);
         }
 
         Ok(())
     }
 
+    pub fn track_alias(&self) -> Option<u64> {
+        let state = self.state.lock();
+        state.track_alias
+    }
+
     pub fn error(mut self, err: ServeError) -> Result<(), ServeError> {
         if let Some(writer) = self.writer.take() {
             writer.close(err.clone())?;
@@ -163,7 +171,8 @@ impl SubscribeRecv {
 
         let writer = subgroups.create(serve::Subgroup {
             group_id: header.group_id,
-            subgroup_id: header.subgroup_id.unwrap(), // TODO SLG - subgroup_id may not be present
+            // When subgroup_id is not present in the header type, it implicitly means subgroup 0
+            subgroup_id: header.subgroup_id.unwrap_or(0),
             priority: header.publisher_priority,
         })?;
 
@@ -176,17 +185,20 @@ impl SubscribeRecv {
         let writer = self.writer.take().ok_or(ServeError::Done)?;
 
         let mut datagrams = match writer {
-            TrackWriterMode::Track(init) => init.datagrams()?, // TODO SLG - is this needed?
+            TrackWriterMode::Track(track) => track.datagrams()?, // TODO SLG - is this needed?
             TrackWriterMode::Datagrams(datagrams) => datagrams,
             _ => return Err(ServeError::Mode),
         };
 
         // TODO SLG - update with new datagram fields
         datagrams.write(serve::Datagram {
             group_id: datagram.group_id,
-            object_id: datagram.object_id.unwrap(), // TODO SLG - make safe
+            // When object_id is not present in the datagram type, it implicitly means object 0
+            object_id: datagram.object_id.unwrap_or(0),
             priority: datagram.publisher_priority,
-            payload: datagram.payload.unwrap(), // TODO SLG - datagram.payload is an Option
+            // TODO: Handle status datagrams separately - they don't have payload
+            // For now, use empty bytes as fallback (shouldn't happen in practice for payload datagrams)
+            payload: datagram.payload.unwrap_or_default(),
         })?;
 
         Ok(())
diff --git a/moq-transport/src/session/subscribed.rs b/moq-transport/src/session/subscribed.rs
@@ -102,24 +102,31 @@ impl Subscribed {
     }
 
     async fn serve_inner(&mut self, track: serve::TrackReader) -> Result<(), SessionError> {
-        let latest = track.latest();
+        // Update largest location before sending SubscribeOk
+        let largest_location = track.largest();
         self.state
             .lock_mut()
             .ok_or(ServeError::Cancel)?
-            .largest_location = latest;
-
-        self.publisher.send_message(message::SubscribeOk {
-            id: self.msg.id,
-            track_alias: self.msg.id, // TODO SLG - use subscription id for now, needs fixing
-            expires: 0,               // TODO SLG
-            group_order: message::GroupOrder::Descending, // TODO: resolve correct value from publisher / subscriber prefs
-            content_exists: latest.is_some(),
-            largest_location: latest,
-            params: Default::default(),
-        });
+            .largest_location = largest_location;
+
+        // Send SubscribeOk using send_message_and_wait to ensure it is sent at least to the QUIC stack before
+        // we start serving the track.  If a subscriber gets the stream before SubscribeOk
+        // then they won't recognize the track_alias in the stream header.
+        self.publisher
+            .send_message_and_wait(message::SubscribeOk {
+                id: self.msg.id,
+                track_alias: self.msg.id, // use subscription id as track alias
+                expires: 0,               // TODO SLG
+                group_order: message::GroupOrder::Descending, // TODO: resolve correct value from publisher / subscriber prefs
+                content_exists: largest_location.is_some(),
+                largest_location,
+                params: Default::default(),
+            })
+            .await;
 
         self.ok = true; // So we send SubscribeDone on drop
 
+        // Serve based on track mode
         match track.mode().await? {
             // TODO cancel track/datagrams on closed
             TrackReaderMode::Stream(_stream) => panic!("deprecated"),
@@ -204,7 +211,7 @@ impl Subscribed {
                     Ok(Some(subgroup)) => {
                         let header = data::SubgroupHeader {
                             header_type: data::StreamHeaderType::SubgroupIdExt,  // SubGroupId = Yes, Extensions = Yes, ContainsEndOfGroup = No
-                            track_alias: self.msg.id, // TODO SLG - use subscription id for now, needs fixing
+                            track_alias: self.msg.id, // use subscription id as track_alias
                             group_id: subgroup.group_id,
                             subgroup_id: Some(subgroup.subgroup_id),
                             publisher_priority: subgroup.priority,
@@ -367,7 +374,7 @@ impl Subscribed {
         while let Some(datagram) = datagrams.read().await? {
             let encoded_datagram = data::Datagram {
                 datagram_type: data::DatagramType::ObjectIdPayload, // TODO SLG
-                track_alias: self.msg.id, //  TODO SLG - use subscription id for now
+                track_alias: self.msg.id, // use subscription id as track_alias
                 group_id: datagram.group_id,
                 object_id: Some(datagram.object_id),
                 publisher_priority: datagram.priority,
@@ -403,7 +410,6 @@ impl Subscribed {
                     encoded_datagram.group_id,
                     encoded_datagram.object_id.unwrap(),
                 )?;
-            // TODO SLG - fix up safety of unwrap()
 
             datagram_count += 1;
         }
diff --git a/moq-transport/src/session/subscriber.rs b/moq-transport/src/session/subscriber.rs
diff --git a/moq-transport/src/watch/queue.rs b/moq-transport/src/watch/queue.rs

Original file line number	Diff line number	Diff line change
`@@ -173,7 +173,7 @@ impl TrackReader {`
`173`	`173`	`}`
`174`	`174`
`175`	`175`	`// Returns the largest group/sequence`
`176`		`- pub fn latest(&self) -> Option<Location> {`
	`176`	`+ pub fn largest(&self) -> Option<Location> {`
`177`	`177`	`// We don't even know the mode yet.`
`178`	`178`	`// TODO populate from SUBSCRIBE_OK`
`179`	`179`	`None`