diff --git a/Cargo.lock b/Cargo.lock index 0a97b0f8..61a8a135 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1777,6 +1777,7 @@ dependencies = [ "once_cell", "parking_lot", "pbkdf2", + "pin-project-lite", "priority-queue", "protobuf", "quick-xml", diff --git a/core/Cargo.toml b/core/Cargo.toml index 1c8728ec..c62e54de 100644 --- a/core/Cargo.toml +++ b/core/Cargo.toml @@ -44,6 +44,7 @@ num-traits = "0.2" once_cell = "1" parking_lot = { version = "0.12", features = ["deadlock_detection"] } pbkdf2 = { version = "0.12", default-features = false, features = ["hmac"] } +pin-project-lite = "0.2" priority-queue = "2.0" protobuf = "3.5" quick-xml = { version = "0.36.1", features = ["serialize"] } diff --git a/core/src/session.rs b/core/src/session.rs index f934ed7b..b2887503 100644 --- a/core/src/session.rs +++ b/core/src/session.rs @@ -12,14 +12,18 @@ use std::{ use byteorder::{BigEndian, ByteOrder}; use bytes::Bytes; use futures_core::TryStream; -use futures_util::{future, ready, StreamExt, TryStreamExt}; +use futures_util::StreamExt; use librespot_protocol::authentication::AuthenticationType; use num_traits::FromPrimitive; use once_cell::sync::OnceCell; use parking_lot::RwLock; +use pin_project_lite::pin_project; use quick_xml::events::Event; use thiserror::Error; -use tokio::{sync::mpsc, time::Instant}; +use tokio::{ + sync::mpsc, + time::{sleep, Duration as TokioDuration, Instant as TokioInstant, Sleep}, +}; use tokio_stream::wrappers::UnboundedReceiverStream; use crate::{ @@ -82,7 +86,6 @@ struct SessionData { time_delta: i64, invalid: bool, user_data: UserData, - last_ping: Option, } struct SessionInternal { @@ -240,6 +243,8 @@ impl Session { } } + // This channel serves as a buffer for packets and serializes access to the TcpStream, such + // that `self.send_packet` can return immediately and needs no additional synchronization. let (tx_connection, rx_connection) = mpsc::unbounded_channel(); self.0 .tx_connection @@ -250,17 +255,20 @@ impl Session { let sender_task = UnboundedReceiverStream::new(rx_connection) .map(Ok) .forward(sink); - let receiver_task = DispatchTask(stream, self.weak()); - let timeout_task = Session::session_timeout(self.weak()); - + let session_weak = self.weak(); tokio::spawn(async move { - let result = future::try_join3(sender_task, receiver_task, timeout_task).await; - - if let Err(e) = result { + if let Err(e) = sender_task.await { error!("{}", e); + if let Some(session) = session_weak.try_upgrade() { + if !session.is_invalid() { + session.shutdown(); + } + } } }); + tokio::spawn(DispatchTask::new(self.weak(), stream)); + Ok(()) } @@ -302,33 +310,6 @@ impl Session { .get_or_init(|| TokenProvider::new(self.weak())) } - /// Returns an error, when we haven't received a ping for too long (2 minutes), - /// which means that we silently lost connection to Spotify servers. - async fn session_timeout(session: SessionWeak) -> io::Result<()> { - // pings are sent every 2 minutes and a 5 second margin should be fine - const SESSION_TIMEOUT: Duration = Duration::from_secs(125); - - while let Some(session) = session.try_upgrade() { - if session.is_invalid() { - break; - } - let last_ping = session.0.data.read().last_ping.unwrap_or_else(Instant::now); - if last_ping.elapsed() >= SESSION_TIMEOUT { - session.shutdown(); - // TODO: Optionally reconnect (with cached/last credentials?) - return Err(io::Error::new( - io::ErrorKind::TimedOut, - "session lost connection to server", - )); - } - // drop the strong reference before sleeping - drop(session); - // a potential timeout cannot occur at least until SESSION_TIMEOUT after the last_ping - tokio::time::sleep_until(last_ping + SESSION_TIMEOUT).await; - } - Ok(()) - } - pub fn time_delta(&self) -> i64 { self.0.data.read().time_delta } @@ -361,96 +342,6 @@ impl Session { } } - fn dispatch(&self, cmd: u8, data: Bytes) -> Result<(), Error> { - use PacketType::*; - - let packet_type = FromPrimitive::from_u8(cmd); - let cmd = match packet_type { - Some(cmd) => cmd, - None => { - trace!("Ignoring unknown packet {:x}", cmd); - return Err(SessionError::Packet(cmd).into()); - } - }; - - match packet_type { - Some(Ping) => { - let server_timestamp = BigEndian::read_u32(data.as_ref()) as i64; - let timestamp = SystemTime::now() - .duration_since(UNIX_EPOCH) - .unwrap_or(Duration::ZERO) - .as_secs() as i64; - - { - let mut data = self.0.data.write(); - data.time_delta = server_timestamp.saturating_sub(timestamp); - data.last_ping = Some(Instant::now()); - } - - self.debug_info(); - self.send_packet(Pong, vec![0, 0, 0, 0]) - } - Some(CountryCode) => { - let country = String::from_utf8(data.as_ref().to_owned())?; - info!("Country: {:?}", country); - self.0.data.write().user_data.country = country; - Ok(()) - } - Some(StreamChunkRes) | Some(ChannelError) => self.channel().dispatch(cmd, data), - Some(AesKey) | Some(AesKeyError) => self.audio_key().dispatch(cmd, data), - Some(MercuryReq) | Some(MercurySub) | Some(MercuryUnsub) | Some(MercuryEvent) => { - self.mercury().dispatch(cmd, data) - } - Some(ProductInfo) => { - let data = std::str::from_utf8(&data)?; - let mut reader = quick_xml::Reader::from_str(data); - - let mut buf = Vec::new(); - let mut current_element = String::new(); - let mut user_attributes: UserAttributes = HashMap::new(); - - loop { - match reader.read_event_into(&mut buf) { - Ok(Event::Start(ref element)) => { - std::str::from_utf8(element)?.clone_into(&mut current_element) - } - Ok(Event::End(_)) => { - current_element = String::new(); - } - Ok(Event::Text(ref value)) => { - if !current_element.is_empty() { - let _ = user_attributes - .insert(current_element.clone(), value.unescape()?.to_string()); - } - } - Ok(Event::Eof) => break, - Ok(_) => (), - Err(e) => warn!( - "Error parsing XML at position {}: {:?}", - reader.buffer_position(), - e - ), - } - } - - trace!("Received product info: {:#?}", user_attributes); - Self::check_catalogue(&user_attributes); - - self.0.data.write().user_data.attributes = user_attributes; - Ok(()) - } - Some(PongAck) - | Some(SecretBlock) - | Some(LegacyWelcome) - | Some(UnknownDataAllZeros) - | Some(LicenseVersion) => Ok(()), - _ => { - trace!("Ignoring {:?} packet with data {:#?}", cmd, data); - Err(SessionError::Packet(cmd as u8).into()) - } - } - } - pub fn send_packet(&self, cmd: PacketType, data: Vec) -> Result<(), Error> { match self.0.tx_connection.get() { Some(tx) => Ok(tx.send((cmd as u8, data))?), @@ -614,50 +505,277 @@ impl Drop for SessionInternal { } } -struct DispatchTask(S, SessionWeak) -where - S: TryStream + Unpin; +#[derive(Clone, Copy, Default, Debug, PartialEq)] +enum KeepAliveState { + #[default] + // Expecting a Ping from the server, either after startup or after a PongAck. + ExpectingPing, -impl Future for DispatchTask + // We need to send a Pong at the given time. + PendingPong, + + // We just sent a Pong and wait for it be ACK'd. + ExpectingPongAck, +} + +const INITIAL_PING_TIMEOUT: TokioDuration = TokioDuration::from_secs(20); +const PING_TIMEOUT: TokioDuration = TokioDuration::from_secs(80); // 60s expected + 20s buffer +const PONG_DELAY: TokioDuration = TokioDuration::from_secs(60); +const PONG_ACK_TIMEOUT: TokioDuration = TokioDuration::from_secs(20); + +impl KeepAliveState { + fn debug(&self, sleep: &Sleep) { + let delay = sleep + .deadline() + .checked_duration_since(TokioInstant::now()) + .map(|t| t.as_secs_f64()) + .unwrap_or(f64::INFINITY); + + trace!("keep-alive state: {:?}, timeout in {:.1}", self, delay); + } +} + +pin_project! { + struct DispatchTask + where + S: TryStream + { + session: SessionWeak, + keep_alive_state: KeepAliveState, + #[pin] + stream: S, + #[pin] + timeout: Sleep, + } + + impl PinnedDrop for DispatchTask + where + S: TryStream + { + fn drop(_this: Pin<&mut Self>) { + debug!("drop Dispatch"); + } + } +} + +impl DispatchTask where - S: TryStream + Unpin, - ::Ok: std::fmt::Debug, + S: TryStream, { - type Output = Result<(), S::Error>; + fn new(session: SessionWeak, stream: S) -> Self { + Self { + session, + keep_alive_state: KeepAliveState::ExpectingPing, + stream, + timeout: sleep(INITIAL_PING_TIMEOUT), + } + } - fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { - let session = match self.1.try_upgrade() { - Some(session) => session, - None => return Poll::Ready(Ok(())), + fn dispatch( + mut self: Pin<&mut Self>, + session: &Session, + cmd: u8, + data: Bytes, + ) -> Result<(), Error> { + use KeepAliveState::*; + use PacketType::*; + + let packet_type = FromPrimitive::from_u8(cmd); + let cmd = match packet_type { + Some(cmd) => cmd, + None => { + trace!("Ignoring unknown packet {:x}", cmd); + return Err(SessionError::Packet(cmd).into()); + } }; - loop { - let (cmd, data) = match ready!(self.0.try_poll_next_unpin(cx)) { - Some(Ok(t)) => t, - None => { - warn!("Connection to server closed."); - session.shutdown(); - return Poll::Ready(Ok(())); + match packet_type { + Some(Ping) => { + trace!("Received Ping"); + if self.keep_alive_state != ExpectingPing { + warn!("Received unexpected Ping from server") } - Some(Err(e)) => { - error!("Connection to server closed."); - session.shutdown(); - return Poll::Ready(Err(e)); - } - }; + let mut this = self.as_mut().project(); + *this.keep_alive_state = PendingPong; + this.timeout + .as_mut() + .reset(TokioInstant::now() + PONG_DELAY); + this.keep_alive_state.debug(&this.timeout); - if let Err(e) = session.dispatch(cmd, data) { - debug!("could not dispatch command: {}", e); + let server_timestamp = BigEndian::read_u32(data.as_ref()) as i64; + let timestamp = SystemTime::now() + .duration_since(UNIX_EPOCH) + .unwrap_or(Duration::ZERO) + .as_secs() as i64; + { + let mut data = session.0.data.write(); + data.time_delta = server_timestamp.saturating_sub(timestamp); + } + + session.debug_info(); + + Ok(()) + } + Some(PongAck) => { + trace!("Received PongAck"); + if self.keep_alive_state != ExpectingPongAck { + warn!("Received unexpected PongAck from server") + } + let mut this = self.as_mut().project(); + *this.keep_alive_state = ExpectingPing; + this.timeout + .as_mut() + .reset(TokioInstant::now() + PING_TIMEOUT); + this.keep_alive_state.debug(&this.timeout); + + Ok(()) + } + Some(CountryCode) => { + let country = String::from_utf8(data.as_ref().to_owned())?; + info!("Country: {:?}", country); + session.0.data.write().user_data.country = country; + Ok(()) + } + Some(StreamChunkRes) | Some(ChannelError) => session.channel().dispatch(cmd, data), + Some(AesKey) | Some(AesKeyError) => session.audio_key().dispatch(cmd, data), + Some(MercuryReq) | Some(MercurySub) | Some(MercuryUnsub) | Some(MercuryEvent) => { + session.mercury().dispatch(cmd, data) + } + Some(ProductInfo) => { + let data = std::str::from_utf8(&data)?; + let mut reader = quick_xml::Reader::from_str(data); + + let mut buf = Vec::new(); + let mut current_element = String::new(); + let mut user_attributes: UserAttributes = HashMap::new(); + + loop { + match reader.read_event_into(&mut buf) { + Ok(Event::Start(ref element)) => { + std::str::from_utf8(element)?.clone_into(&mut current_element) + } + Ok(Event::End(_)) => { + current_element = String::new(); + } + Ok(Event::Text(ref value)) => { + if !current_element.is_empty() { + let _ = user_attributes + .insert(current_element.clone(), value.unescape()?.to_string()); + } + } + Ok(Event::Eof) => break, + Ok(_) => (), + Err(e) => warn!( + "Error parsing XML at position {}: {:?}", + reader.buffer_position(), + e + ), + } + } + + trace!("Received product info: {:#?}", user_attributes); + Session::check_catalogue(&user_attributes); + + session.0.data.write().user_data.attributes = user_attributes; + Ok(()) + } + Some(SecretBlock) + | Some(LegacyWelcome) + | Some(UnknownDataAllZeros) + | Some(LicenseVersion) => Ok(()), + _ => { + trace!("Ignoring {:?} packet with data {:#?}", cmd, data); + Err(SessionError::Packet(cmd as u8).into()) } } } } -impl Drop for DispatchTask +impl Future for DispatchTask where - S: TryStream + Unpin, + S: TryStream, + ::Ok: std::fmt::Debug, { - fn drop(&mut self) { - debug!("drop Dispatch"); + type Output = Result<(), S::Error>; + + fn poll(mut self: Pin<&mut Self>, cx: &mut Context<'_>) -> Poll { + use KeepAliveState::*; + + let session = match self.session.try_upgrade() { + Some(session) => session, + None => return Poll::Ready(Ok(())), + }; + + // Process all messages that are immediately ready + loop { + match self.as_mut().project().stream.try_poll_next(cx) { + Poll::Ready(Some(Ok((cmd, data)))) => { + let result = self.as_mut().dispatch(&session, cmd, data); + if let Err(e) = result { + debug!("could not dispatch command: {}", e); + } + } + Poll::Ready(None) => { + warn!("Connection to server closed."); + session.shutdown(); + return Poll::Ready(Ok(())); + } + Poll::Ready(Some(Err(e))) => { + error!("Connection to server closed."); + session.shutdown(); + return Poll::Ready(Err(e)); + } + Poll::Pending => break, + } + } + + // Handle the keep-alive sequence, returning an error when we haven't received a + // Ping/PongAck for too long. + // + // The expected keepalive sequence is + // - Server: Ping + // - wait 60s + // - Client: Pong + // - Server: PongAck + // - wait 60s + // - repeat + // + // This means that we silently lost connection to Spotify servers if + // - we don't receive Ping immediately after connecting, + // - we don't receive a Ping 60s after the last PongAck or + // - we don't receive a PongAck immediately after our Pong. + // + // Currently, we add a safety margin of 20s to these expected deadlines. + let mut this = self.as_mut().project(); + if let Poll::Ready(()) = this.timeout.as_mut().poll(cx) { + match this.keep_alive_state { + ExpectingPing | ExpectingPongAck => { + if !session.is_invalid() { + session.shutdown(); + } + // TODO: Optionally reconnect (with cached/last credentials?) + return Poll::Ready(Err(io::Error::new( + io::ErrorKind::TimedOut, + format!( + "session lost connection to server ({:?})", + this.keep_alive_state + ), + ))); + } + PendingPong => { + trace!("Sending Pong"); + // TODO: Ideally, this should flush the `Framed as Sink` + // before starting the timeout. + let _ = session.send_packet(PacketType::Pong, vec![0, 0, 0, 0]); + *this.keep_alive_state = ExpectingPongAck; + this.timeout + .as_mut() + .reset(TokioInstant::now() + PONG_ACK_TIMEOUT); + this.keep_alive_state.debug(&this.timeout); + } + } + } + + Poll::Pending } }