use std::{collections::HashMap, io::Write}; use byteorder::{BigEndian, ByteOrder, WriteBytesExt}; use bytes::Bytes; use thiserror::Error; use tokio::sync::oneshot; use crate::{packet::PacketType, util::SeqGenerator, Error, FileId, SpotifyId}; #[derive(Debug, Hash, PartialEq, Eq, Copy, Clone)] pub struct AudioKey(pub [u8; 16]); #[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 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 = SeqGenerator::new(0), pending: HashMap>> = HashMap::new(), } } impl AudioKeyManager { 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)) .ok_or(AudioKeyError::Sequence(seq))?; 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 { let (tx, rx) = oneshot::channel(); let seq = self.lock(move |inner| { let seq = inner.sequence.get(); inner.pending.insert(seq, tx); seq }); self.send_key_request(seq, track, file)?; rx.await? } fn send_key_request(&self, seq: u32, track: SpotifyId, file: FileId) -> Result<(), Error> { let mut data: Vec = Vec::new(); data.write_all(&file.0)?; data.write_all(&track.to_raw())?; data.write_u32::(seq)?; data.write_u16::(0x0000)?; self.session().send_packet(PacketType::RequestKey, data) } }