Main: Fix support for authenticated media

This patch greately improves the authenticated media support
by checking for advertised support first, then trying
unauthenticated first and correctly interpreting M_NOT_FOUND
errors to then request a token.

src/main.rs

@@ -1,4 +1,4 @@
-use std::fs;
+use std::{collections::HashSet, fs};
 use std::process::exit;
 use std::io::Read;
 
@@ -9,28 +9,47 @@ use log::{info, debug, error};
 use serde::Deserialize;
 use rustyline_async::{Readline, ReadlineEvent};
 
-use matrix_sdk::{matrix_auth::{MatrixSession, MatrixSessionTokens}, ruma::{api::client::{authenticated_media, media}, device_id, user_id, OwnedMxcUri}, Client, SessionMeta};
+use matrix_sdk::{matrix_auth::{MatrixSession, MatrixSessionTokens}, Client, SessionMeta};
 use matrix_sdk::ruma::events::room::{MediaSource, message::{MessageType, RoomMessageEventContent}};
+use matrix_sdk::ruma::api::client::error::{ErrorBody, ErrorKind};
+use matrix_sdk::ruma::api::client::{authenticated_media, media};
+use matrix_sdk::ruma::api::{MatrixVersion, OutgoingRequest};
+use matrix_sdk::ruma::{device_id, user_id, OwnedMxcUri};
+use matrix_sdk::reqwest::StatusCode;
 
 // NOTE: This has been stolen from the matrix-rust-sdk to patch out checks for authenticated media
 async fn get_media_content(
     client: &Client,
     source: &MediaSource,
-    use_auth: bool
+    use_authenticated_media_endpoints: bool
 ) -> matrix_sdk::Result<Vec<u8>> {
+    #[allow(deprecated)]
+    type DeprecatedMediaRequest = media::get_content::v3::Request;
+    type AuthenticatedMediaRequest = authenticated_media::get_content::v1::Request;
+
+    let deprecated_media_request_path = DeprecatedMediaRequest::METADATA.history.stable_endpoint_for(&[MatrixVersion::V1_1]).unwrap();
+    let authenticated_media_request_path = AuthenticatedMediaRequest::METADATA.history.stable_endpoint_for(&[MatrixVersion::V1_11]).unwrap();
+
+    let homeserver_base_url = client.homeserver().to_string();
+    let homeserver_base_url = homeserver_base_url.trim_end_matches("/");
+
     let content: Vec<u8> = match &source {
         MediaSource::Encrypted(file) => {
-            let content = if use_auth {
+            let content = if use_authenticated_media_endpoints {
-                let request = authenticated_media::get_content::v1::Request::from_uri(&file.url)?;
+                debug!("Attempting to download {} from {}{}", file.url, homeserver_base_url, authenticated_media_request_path);
+
+                let request = AuthenticatedMediaRequest::from_uri(&file.url)?;
                 client.send(request, None).await?.file
             } else {
+                debug!("Attempting to download {} from {}{}", file.url, homeserver_base_url, deprecated_media_request_path);
+
                 #[allow(deprecated)]
-                let request = media::get_content::v3::Request::from_url(&file.url)?;
+                let request = DeprecatedMediaRequest::from_url(&file.url)?;
                 client.send(request, None).await?.file
             };
 
             let content = {
-                debug!("Decrypting file");
+                debug!("Decrypting file using {}", file.key.alg);
 
                 let content_len = content.len();
                 let mut cursor = std::io::Cursor::new(content);
@@ -57,14 +76,18 @@ async fn get_media_content(
             content
         }
         MediaSource::Plain(uri) => {
-            debug!("File doesn't appear to be encrypted");
+            debug!("No encryption details provided, will not attempt decryption after download");
 
-            if use_auth {
+            if use_authenticated_media_endpoints {
-                let request = authenticated_media::get_content::v1::Request::from_uri(uri)?;
+                debug!("Attempting to download {} from {}{}", uri, homeserver_base_url, authenticated_media_request_path);
+
+                let request = AuthenticatedMediaRequest::from_uri(uri)?;
                 client.send(request, None).await?.file
             } else {
+                debug!("Attempting to download {} from {}{}", uri, homeserver_base_url, deprecated_media_request_path);
+
                 #[allow(deprecated)]
-                let request = media::get_content::v3::Request::from_url(uri)?;
+                let request = DeprecatedMediaRequest::from_url(uri)?;
                 client.send(request, None).await?.file
             }
         }
@@ -73,6 +96,49 @@ async fn get_media_content(
     Ok(content)
 }
 
+async fn supports_authenticated_media(client: &Client) -> Result<bool> {
+    debug!("Discovering supported versions and features");
+
+    let supported_versions_response = client.send(
+        matrix_sdk::ruma::api::client::discovery::get_supported_versions::Request::new(),
+        None
+    ).await?;
+
+    let known_supported_versions = supported_versions_response.known_versions().collect::<Vec<_>>();
+    let supported_unstable_features = supported_versions_response.unstable_features
+        .into_iter()
+        .filter(|(_, supported)| *supported)
+        .map(|(feature, _)| feature)
+        .collect::<HashSet<_>>();
+
+    const AUTHENTICATED_MEDIA_STABLE_FEATURE: &str = "org.matrix.msc3916.stable";
+
+    Ok(match (
+        known_supported_versions.contains(&MatrixVersion::V1_11),
+        supported_unstable_features.contains(AUTHENTICATED_MEDIA_STABLE_FEATURE)
+    ) {
+        (false, false) => {
+            debug!("Authenticated Media *is not* supported (the server neither advertises Matrix 1.11 nor org.matrix.msc3916.stable)");
+            false
+        }
+
+        (true, false) => {
+            debug!("Authenticated Media *is* supported (the server advertises support for Matrix 1.11 but not org.matrix.msc3916.stable)");
+            true
+        }
+
+        (false, true) => {
+            debug!("Authenticated Media *is* supported (the server doesn't advertise support for Matrix 1.11 but does for org.matrix.msc3916.stable)");
+            true
+        }
+
+        (true, true) => {
+            debug!("Authenticated Media *is* supported (the server advertises support for both Matrix 1.11 and org.matrix.msc3916.stable)");
+            true
+        }
+    })
+}
+
 async fn process_media(source: MediaSource, body: Option<String>, server: Option<String>) -> Result<()> {
     let uri = match source.clone() {
         MediaSource::Plain(uri) => uri,
@@ -98,44 +164,53 @@ async fn process_media(source: MediaSource, body: Option<String>, server: Option
         },
     }.build().await?;
 
-    debug!("Downloading {} from {}", uri, client.homeserver());
+    let supports_authenticated_media = supports_authenticated_media(&client).await?;
 
-    let content = match get_media_content(&client, &source, false).await {
+    let content = match supports_authenticated_media {
-        Ok(content) => content,
+        false => get_media_content(&client, &source, false).await
-        Err(error) => {
+            .map_err(|error| anyhow::Error::msg(error.to_string()))?,
-            let error_is_needs_authentication = match &error {
-                matrix_sdk::Error::Http(http) => match http {
-                    matrix_sdk::HttpError::IntoHttp(into_http) => {
-                        match into_http {
-                            matrix_sdk::ruma::api::error::IntoHttpError::NeedsAuthentication => true,
-                            _ => false,
-                        }
-                    },
-                    _ => false
-                },
-                _ => false,
-            };
 
-            if error_is_needs_authentication {
+        true => match get_media_content(&client, &source, false).await {
-                info!("The server did not allow downloading without authentication, please enter an access token:");
+            Ok(content) => content,
+            Err(error) => {
+                let is_probably_new_media = matches!(
+                    error,
+                    matrix_sdk::Error::Http(
+                        matrix_sdk::HttpError::Api(
+                            matrix_sdk::ruma::api::error::FromHttpResponseError::Server(
+                                matrix_sdk::RumaApiError::ClientApi(
+                                    matrix_sdk::ruma::api::client::Error {
+                                        status_code: StatusCode::NOT_FOUND,
+                                        body: ErrorBody::Standard { kind: ErrorKind::NotFound, message: _ },
+                                        ..
+                                    }
+                                )
+                            )
+                        )
+                    )
+                );
 
-                let token = rpassword::prompt_password("> ")?;
+                if is_probably_new_media {
+                    info!("The server returned 404 M_NOT_FOUND. Assuming the media is new, please enter an access token:");
 
-                client.restore_session(MatrixSession {
+                    let token = rpassword::prompt_password("> ")?;
-                    meta: SessionMeta {
-                        user_id: user_id!("@unused:example.com").to_owned(),
-                        device_id: device_id!("UNUSED").to_owned(),
-                    },
-                    tokens: MatrixSessionTokens {
-                        access_token: token,
-                        refresh_token: None,
-                    },
-                }).await?;
 
-                get_media_content(&client, &source, false).await
+                    client.restore_session(MatrixSession {
-                    .map_err(|error| anyhow::Error::msg(error.to_string()))?
+                        meta: SessionMeta {
-            } else {
+                            user_id: user_id!("@unused:example.com").to_owned(),
-                return Err(anyhow::Error::msg(error.to_string()));
+                            device_id: device_id!("UNUSED").to_owned(),
+                        },
+                        tokens: MatrixSessionTokens {
+                            access_token: token,
+                            refresh_token: None,
+                        },
+                    }).await?;
+
+                    get_media_content(&client, &source, true).await
+                        .map_err(|error| anyhow::Error::msg(error.to_string()))?
+                } else {
+                    return Err(anyhow::Error::msg(error.to_string()));
+                }
             }
         }
     };