Change panics into Result<_, librespot_core::Error>

core/src/audio_key.rs

@@ -1,54 +1,85 @@
+use std::{collections::HashMap, io::Write};
+
 use byteorder::{BigEndian, ByteOrder, WriteBytesExt};
 use bytes::Bytes;
-use std::collections::HashMap;
-use std::io::Write;
+use thiserror::Error;
 use tokio::sync::oneshot;
 
-use crate::file_id::FileId;
-use crate::packet::PacketType;
-use crate::spotify_id::SpotifyId;
-use crate::util::SeqGenerator;
+use crate::{packet::PacketType, util::SeqGenerator, Error, FileId, SpotifyId};
 
 #[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
 pub struct AudioKey(pub [u8; 16]);
 
-#[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)]
-pub struct AudioKeyError;
+#[derive(Debug, Error)]
+pub enum AudioKeyError {
+    #[error("audio key error")]
+    AesKey,
+    #[error("other end of channel disconnected")]
+    Channel,
+    #[error("unexpected packet type {0}")]
+    Packet(u8),
+    #[error("sequence {0} not pending")]
+    Sequence(u32),
+}
+
+impl From<AudioKeyError> for Error {
+    fn from(err: AudioKeyError) -> Self {
+        match err {
+            AudioKeyError::AesKey => Error::unavailable(err),
+            AudioKeyError::Channel => Error::aborted(err),
+            AudioKeyError::Sequence(_) => Error::aborted(err),
+            AudioKeyError::Packet(_) => Error::unimplemented(err),
+        }
+    }
+}
 
 component! {
     AudioKeyManager : AudioKeyManagerInner {
         sequence: SeqGenerator<u32> = SeqGenerator::new(0),
-        pending: HashMap<u32, oneshot::Sender<Result<AudioKey, AudioKeyError>>> = HashMap::new(),
+        pending: HashMap<u32, oneshot::Sender<Result<AudioKey, Error>>> = HashMap::new(),
     }
 }
 
 impl AudioKeyManager {
-    pub(crate) fn dispatch(&self, cmd: PacketType, mut data: Bytes) {
+    pub(crate) fn dispatch(&self, cmd: PacketType, mut data: Bytes) -> Result<(), Error> {
         let seq = BigEndian::read_u32(data.split_to(4).as_ref());
 
-        let sender = self.lock(|inner| inner.pending.remove(&seq));
+        let sender = self
+            .lock(|inner| inner.pending.remove(&seq))
+            .ok_or(AudioKeyError::Sequence(seq))?;
 
-        if let Some(sender) = sender {
-            match cmd {
-                PacketType::AesKey => {
-                    let mut key = [0u8; 16];
-                    key.copy_from_slice(data.as_ref());
-                    let _ = sender.send(Ok(AudioKey(key)));
-                }
-                PacketType::AesKeyError => {
-                    warn!(
-                        "error audio key {:x} {:x}",
-                        data.as_ref()[0],
-                        data.as_ref()[1]
-                    );
-                    let _ = sender.send(Err(AudioKeyError));
-                }
-                _ => (),
+        match cmd {
+            PacketType::AesKey => {
+                let mut key = [0u8; 16];
+                key.copy_from_slice(data.as_ref());
+                sender
+                    .send(Ok(AudioKey(key)))
+                    .map_err(|_| AudioKeyError::Channel)?
+            }
+            PacketType::AesKeyError => {
+                error!(
+                    "error audio key {:x} {:x}",
+                    data.as_ref()[0],
+                    data.as_ref()[1]
+                );
+                sender
+                    .send(Err(AudioKeyError::AesKey.into()))
+                    .map_err(|_| AudioKeyError::Channel)?
+            }
+            _ => {
+                trace!(
+                    "Did not expect {:?} AES key packet with data {:#?}",
+                    cmd,
+                    data
+                );
+                return Err(AudioKeyError::Packet(cmd as u8).into());
             }
         }
+
+        Ok(())
     }
 
-    pub async fn request(&self, track: SpotifyId, file: FileId) -> Result<AudioKey, AudioKeyError> {
+    pub async fn request(&self, track: SpotifyId, file: FileId) -> Result<AudioKey, Error> {
         let (tx, rx) = oneshot::channel();
 
         let seq = self.lock(move |inner| {
@@ -57,16 +88,16 @@ impl AudioKeyManager {
             seq
         });
 
-        self.send_key_request(seq, track, file);
-        rx.await.map_err(|_| AudioKeyError)?
+        self.send_key_request(seq, track, file)?;
+        rx.await?
     }
 
-    fn send_key_request(&self, seq: u32, track: SpotifyId, file: FileId) {
+    fn send_key_request(&self, seq: u32, track: SpotifyId, file: FileId) -> Result<(), Error> {
         let mut data: Vec<u8> = Vec::new();
-        data.write_all(&file.0).unwrap();
-        data.write_all(&track.to_raw()).unwrap();
-        data.write_u32::<BigEndian>(seq).unwrap();
-        data.write_u16::<BigEndian>(0x0000).unwrap();
+        data.write_all(&file.0)?;
+        data.write_all(&track.to_raw())?;
+        data.write_u32::<BigEndian>(seq)?;
+        data.write_u16::<BigEndian>(0x0000)?;
 
         self.session().send_packet(PacketType::RequestKey, data)
     }