From 850db432540e1b8a7a33787372c4e2b4a22fb3bd Mon Sep 17 00:00:00 2001
From: Roderick van Domburg <roderick@vandomburg.net>
Date: Sat, 19 Jun 2021 22:47:39 +0200
Subject: [PATCH] Add token provider

---
 core/src/dealer/api.rs |   2 +
 core/src/lib.rs        |   1 +
 core/src/session.rs    |  10 ++++
 core/src/token.rs      | 124 +++++++++++++++++++++++++++++++++++++++++
 4 files changed, 137 insertions(+)
 create mode 100644 core/src/dealer/api.rs
 create mode 100644 core/src/token.rs

diff --git a/core/src/dealer/api.rs b/core/src/dealer/api.rs
new file mode 100644
index 00000000..d9dd2b9b
--- /dev/null
+++ b/core/src/dealer/api.rs
@@ -0,0 +1,2 @@
+// https://github.com/librespot-org/librespot-java/blob/27783e06f456f95228c5ac37acf2bff8c1a8a0c4/lib/src/main/java/xyz/gianlu/librespot/dealer/ApiClient.java
+
diff --git a/core/src/lib.rs b/core/src/lib.rs
index f26caf3d..7a07d94d 100644
--- a/core/src/lib.rs
+++ b/core/src/lib.rs
@@ -25,6 +25,7 @@ mod proxytunnel;
 pub mod session;
 mod socket;
 pub mod spotify_id;
+mod token;
 #[doc(hidden)]
 pub mod util;
 pub mod version;
diff --git a/core/src/session.rs b/core/src/session.rs
index 17452b20..fe8e4d5f 100644
--- a/core/src/session.rs
+++ b/core/src/session.rs
@@ -24,6 +24,7 @@ use crate::channel::ChannelManager;
 use crate::config::SessionConfig;
 use crate::connection::{self, AuthenticationError};
 use crate::mercury::MercuryManager;
+use crate::token::TokenProvider;
 
 #[derive(Debug, Error)]
 pub enum SessionError {
@@ -49,6 +50,7 @@ struct SessionInternal {
     audio_key: OnceCell<AudioKeyManager>,
     channel: OnceCell<ChannelManager>,
     mercury: OnceCell<MercuryManager>,
+    token_provider: OnceCell<TokenProvider>,
     cache: Option<Arc<Cache>>,
 
     handle: tokio::runtime::Handle,
@@ -119,6 +121,7 @@ impl Session {
             audio_key: OnceCell::new(),
             channel: OnceCell::new(),
             mercury: OnceCell::new(),
+            token_provider: OnceCell::new(),
             handle,
             session_id,
         }));
@@ -157,6 +160,12 @@ impl Session {
             .get_or_init(|| MercuryManager::new(self.weak()))
     }
 
+    pub fn token_provider(&self) -> &TokenProvider {
+        self.0
+            .token_provider
+            .get_or_init(|| TokenProvider::new(self.weak()))
+    }
+
     pub fn time_delta(&self) -> i64 {
         self.0.data.read().unwrap().time_delta
     }
@@ -181,6 +190,7 @@ impl Session {
     #[allow(clippy::match_same_arms)]
     fn dispatch(&self, cmd: u8, data: Bytes) {
         match cmd {
+            // TODO: add command types
             0x4 => {
                 let server_timestamp = BigEndian::read_u32(data.as_ref()) as i64;
                 let timestamp = match SystemTime::now().duration_since(UNIX_EPOCH) {
diff --git a/core/src/token.rs b/core/src/token.rs
new file mode 100644
index 00000000..239b40af
--- /dev/null
+++ b/core/src/token.rs
@@ -0,0 +1,124 @@
+// Ported from librespot-java. Relicensed under MIT with permission.
+
+use crate::mercury::MercuryError;
+
+use serde::Deserialize;
+
+use std::error::Error;
+use std::time::{Duration, Instant};
+
+component! {
+    TokenProvider : TokenProviderInner {
+        tokens: Vec<Token> = vec![],
+    }
+}
+
+#[derive(Clone, Debug)]
+pub struct Token {
+    expires_in: Duration,
+    access_token: String,
+    scopes: Vec<String>,
+    timestamp: Instant,
+}
+
+#[derive(Deserialize)]
+#[serde(rename_all = "camelCase")]
+struct TokenData {
+    expires_in: u64,
+    access_token: String,
+    scope: Vec<String>,
+}
+
+impl TokenProvider {
+    const KEYMASTER_CLIENT_ID: &'static str = "65b708073fc0480ea92a077233ca87bd";
+
+    fn find_token(&self, scopes: Vec<String>) -> Option<usize> {
+        self.lock(|inner| {
+            for i in 0..inner.tokens.len() {
+                if inner.tokens[i].in_scopes(scopes.clone()) {
+                    return Some(i);
+                }
+            }
+            None
+        })
+    }
+
+    pub async fn get_token(&self, scopes: Vec<String>) -> Result<Token, MercuryError> {
+        if scopes.is_empty() {
+            return Err(MercuryError);
+        }
+
+        if let Some(index) = self.find_token(scopes.clone()) {
+            let cached_token = self.lock(|inner| inner.tokens[index].clone());
+            if cached_token.is_expired() {
+                self.lock(|inner| inner.tokens.remove(index));
+            } else {
+                return Ok(cached_token);
+            }
+        }
+
+        trace!(
+            "Requested token in scopes {:?} unavailable or expired, requesting new token.",
+            scopes
+        );
+
+        let query_uri = format!(
+            "hm://keymaster/token/authenticated?scope={}&client_id={}&device_id={}",
+            scopes.join(","),
+            Self::KEYMASTER_CLIENT_ID,
+            self.session().device_id()
+        );
+        let request = self.session().mercury().get(query_uri);
+        let response = request.await?;
+
+        if response.status_code == 200 {
+            let data = response
+                .payload
+                .first()
+                .expect("No tokens received")
+                .to_vec();
+            let token = Token::new(String::from_utf8(data).unwrap()).map_err(|_| MercuryError)?;
+            trace!("Got token: {:?}", token);
+            self.lock(|inner| inner.tokens.push(token.clone()));
+            Ok(token)
+        } else {
+            Err(MercuryError)
+        }
+    }
+}
+
+impl Token {
+    const EXPIRY_THRESHOLD: Duration = Duration::from_secs(10);
+
+    pub fn new(body: String) -> Result<Self, Box<dyn Error>> {
+        let data: TokenData = serde_json::from_slice(body.as_ref())?;
+        Ok(Self {
+            expires_in: Duration::from_secs(data.expires_in),
+            access_token: data.access_token,
+            scopes: data.scope,
+            timestamp: Instant::now(),
+        })
+    }
+
+    pub fn is_expired(&self) -> bool {
+        self.timestamp + (self.expires_in - Self::EXPIRY_THRESHOLD) < Instant::now()
+    }
+
+    pub fn in_scope(&self, scope: String) -> bool {
+        for s in &self.scopes {
+            if *s == scope {
+                return true;
+            }
+        }
+        false
+    }
+
+    pub fn in_scopes(&self, scopes: Vec<String>) -> bool {
+        for s in scopes {
+            if !self.in_scope(s) {
+                return false;
+            }
+        }
+        true
+    }
+}