spotifyio/crates/spotifyio/src/connection/mod.rs

pub use self::{codec::ApCodec, handshake::handshake, packet::PacketType};

use std::{io, time::Duration};

use futures_util::{SinkExt, StreamExt};
use num_traits::FromPrimitive;
use protobuf::Message;
use thiserror::Error;
use tokio::net::TcpStream;
use tokio_util::codec::Framed;
use tracing::{debug, trace};
use url::Url;

use crate::{authentication::AuthCredentials, util, Error};

use spotifyio_protocol::keyexchange::{APLoginFailed, ErrorCode};

mod codec;
mod diffie_hellman;
mod handshake;
mod packet;
mod proxytunnel;
mod socket;

pub type Transport = Framed<TcpStream, ApCodec>;

fn login_error_message(code: &ErrorCode) -> &'static str {
    pub use ErrorCode::*;
    match code {
        ProtocolError => "Protocol error",
        TryAnotherAP => "Try another access point",
        BadConnectionId => "Bad connection ID",
        TravelRestriction => "Travel restriction",
        PremiumAccountRequired => "Premium account required",
        BadCredentials => "Bad credentials",
        CouldNotValidateCredentials => "Could not validate credentials",
        AccountExists => "Account exists",
        ExtraVerificationRequired => "Extra verification required",
        InvalidAppKey => "Invalid app key",
        ApplicationBanned => "Application banned",
    }
}

#[derive(Debug, Error)]
pub enum AuthenticationError {
    #[error("Login failed with reason: {}", login_error_message(.0))]
    LoginFailed(ErrorCode),
    #[error("invalid packet {0}")]
    Packet(u8),
    #[error("transport returned no data")]
    Transport,
}

impl From<AuthenticationError> for Error {
    fn from(err: AuthenticationError) -> Self {
        match err {
            AuthenticationError::LoginFailed(_) => Error::permission_denied(err),
            AuthenticationError::Packet(_) => Error::unimplemented(err),
            AuthenticationError::Transport => Error::unavailable(err),
        }
    }
}

impl From<APLoginFailed> for AuthenticationError {
    fn from(login_failure: APLoginFailed) -> Self {
        Self::LoginFailed(login_failure.error_code())
    }
}

pub async fn connect(host: &str, port: u16, proxy: Option<&Url>) -> io::Result<Transport> {
    const TIMEOUT: Duration = Duration::from_secs(3);
    let socket = tokio::time::timeout(TIMEOUT, socket::connect(host, port, proxy)).await??;

    handshake(socket).await
}

pub async fn connect_with_retry(
    host: &str,
    port: u16,
    proxy: Option<&Url>,
    max_retries: u8,
) -> io::Result<Transport> {
    let mut num_retries = 0;
    loop {
        match connect(host, port, proxy).await {
            Ok(f) => return Ok(f),
            Err(e) => {
                debug!("Connection failed: {e}");
                if num_retries < max_retries {
                    num_retries += 1;
                    debug!("Retry access point...");
                    continue;
                }
                return Err(e);
            }
        }
    }
}

pub async fn authenticate(
    transport: &mut Transport,
    credentials: AuthCredentials,
    device_id: &str,
) -> Result<AuthCredentials, Error> {
    use spotifyio_protocol::authentication::{APWelcome, ClientResponseEncrypted, CpuFamily, Os};

    let mut packet = ClientResponseEncrypted::new();
    if let Some(username) = credentials.user_id {
        packet
            .login_credentials
            .mut_or_insert_default()
            .set_username(username);
    }
    packet
        .login_credentials
        .mut_or_insert_default()
        .set_typ(credentials.auth_type);
    packet
        .login_credentials
        .mut_or_insert_default()
        .set_auth_data(credentials.auth_data);
    packet
        .system_info
        .mut_or_insert_default()
        .set_cpu_family(CpuFamily::CPU_X86_64);
    packet
        .system_info
        .mut_or_insert_default()
        .set_os(Os::OS_LINUX);
    packet
        .system_info
        .mut_or_insert_default()
        .set_device_id(device_id.to_string());
    packet.set_version_string(util::SPOTIFY_VERSION.to_string());

    let cmd = PacketType::Login;
    let data = packet.write_to_bytes()?;

    debug!("Authenticating with AP using {:?}", credentials.auth_type);
    transport.send((cmd as u8, data)).await?;
    let (cmd, data) = transport
        .next()
        .await
        .ok_or(AuthenticationError::Transport)??;
    let packet_type = FromPrimitive::from_u8(cmd);
    let result = match packet_type {
        Some(PacketType::APWelcome) => {
            let welcome_data = APWelcome::parse_from_bytes(data.as_ref())?;

            let reusable_credentials = AuthCredentials {
                user_id: Some(welcome_data.canonical_username().to_owned()),
                auth_type: welcome_data.reusable_auth_credentials_type(),
                auth_data: welcome_data.reusable_auth_credentials().to_owned(),
            };

            Ok(reusable_credentials)
        }
        Some(PacketType::AuthFailure) => {
            let error_data = APLoginFailed::parse_from_bytes(data.as_ref())?;
            Err(error_data.into())
        }
        _ => {
            trace!(
                "Did not expect {:?} AES key packet with data {:#?}",
                cmd,
                data
            );
            Err(AuthenticationError::Packet(cmd))
        }
    };
    Ok(result?)
}