Implement new CdnKey protocol

Includes numbered CDNs and attachment key
whisperfish · Sep 27, 2020 · f84fdeb · f84fdeb
1 parent 98f4ca8
commit f84fdeb
Show file tree

Hide file tree

Showing 3 changed files with 52 additions and 38 deletions.
diff --git a/libsignal-service-actix/src/push_service.rs b/libsignal-service-actix/src/push_service.rs
@@ -131,11 +131,12 @@ impl PushService for AwcPushService {
 
     async fn get_from_cdn(
         &mut self,
+        cdn_id: u64,
         path: &str,
     ) -> Result<Self::ByteStream, ServiceError> {
         use futures::stream::TryStreamExt;
 
-        let url = Url::parse(&self.cfg.cdn_urls[0])
+        let url = Url::parse(&self.cfg.cdn_urls[&cdn_id])
             .expect("valid cdn base url")
             .join(path)
             .expect("valid CDN path");

diff --git a/libsignal-service/src/configuration.rs b/libsignal-service/src/configuration.rs
@@ -1,9 +1,11 @@
+use std::collections::HashMap;
+
 use crate::envelope::{CIPHER_KEY_SIZE, MAC_KEY_SIZE};
 
 #[derive(Clone)]
 pub struct ServiceConfiguration {
     pub service_urls: Vec<String>,
-    pub cdn_urls: Vec<String>,
+    pub cdn_urls: HashMap<u64, String>,
     pub contact_discovery_url: Vec<String>,
     pub certificate_authority: String,
 }
@@ -51,38 +53,40 @@ impl Into<ServiceConfiguration> for SignalServers {
         match self {
             // configuration with the Signal API staging endpoints
             // see: https://github.com/signalapp/Signal-Desktop/blob/master/config/default.json
-            SignalServers::Staging => {
-                ServiceConfiguration {
-                    service_urls: vec![
-                        "https://textsecure-service-staging.whispersystems.org".into(),
-                    ],
-                    cdn_urls: vec![
-                        "https://cdn-staging.signal.org".into(),
-                        "https://cdn2-staging.signal.org".into(),
-                    ],
-                    contact_discovery_url: vec![
-                        "https://api-staging.directory.signal.org".into(),
-                    ],
-                    certificate_authority: SIGNAL_ROOT_CA.into(),
-                }
+            SignalServers::Staging => ServiceConfiguration {
+                service_urls: vec![
+                    "https://textsecure-service-staging.whispersystems.org"
+                        .into(),
+                ],
+                cdn_urls: {
+                    let mut map = HashMap::new();
+                    map.insert(0, "https://cdn-staging.signal.org".into());
+                    map.insert(2, "https://cdn2-staging.signal.org".into());
+                    map
+                },
+                contact_discovery_url: vec![
+                    "https://api-staging.directory.signal.org".into(),
+                ],
+                certificate_authority: SIGNAL_ROOT_CA.into(),
             },
             // configuration with the Signal API production endpoints
             // https://github.com/signalapp/Signal-Desktop/blob/master/config/production.json
-            SignalServers::Production => {
-                ServiceConfiguration {
-                    service_urls: vec![
-                        "https://textsecure-service.whispersystems.org".into()
-                    ],
-                    cdn_urls: vec![
-                        "https://cdn.signal.org".into(),
-                        "https://cdn2.signal.org".into(),
-                    ],
-                    contact_discovery_url: vec![
-                        "https://api.directory.signal.org".into()
-                    ],
-                    certificate_authority: SIGNAL_ROOT_CA.into(),
-                }
-            }
+            SignalServers::Production => ServiceConfiguration {
+                service_urls: vec![
+                    "https://textsecure-service.whispersystems.org".into(),
+                ],
+                cdn_urls: {
+                    let mut map = HashMap::new();
+                    map.insert(0, "https://cdn.signal.org".into());
+                    map.insert(2, "https://cdn2.signal.org".into());
+                    map
+                },
+
+                contact_discovery_url: vec![
+                    "https://api.directory.signal.org".into()
+                ],
+                certificate_authority: SIGNAL_ROOT_CA.into(),
+            },
         }
     }
-}
+}
diff --git a/libsignal-service/src/push_service.rs b/libsignal-service/src/push_service.rs
@@ -4,7 +4,7 @@ use crate::{
     configuration::{Credentials, ServiceConfiguration},
     envelope::*,
     messagepipe::WebSocketService,
-    proto::AttachmentPointer,
+    proto::{attachment_pointer::AttachmentIdentifier, AttachmentPointer},
     utils::serde_base64,
 };
 
@@ -128,11 +128,11 @@ impl ServiceError {
             StatusCode::NO_CONTENT => Ok(()),
             StatusCode::UNAUTHORIZED | StatusCode::FORBIDDEN => {
                 Err(ServiceError::Unauthorized)
-            }
+            },
             StatusCode::PAYLOAD_TOO_LARGE => {
                 // This is 413 and means rate limit exceeded for Signal.
                 Err(ServiceError::RateLimitExceeded)
-            }
+            },
             // XXX: fill in rest from PushServiceSocket
             _ => Err(ServiceError::UnhandledResponseCode {
                 http_code: code.as_u16(),
@@ -162,6 +162,7 @@ pub trait PushService {
     /// Downloads larger files in streaming fashion, e.g. attachments.
     async fn get_from_cdn(
         &mut self,
+        cdn_id: u64,
         path: &str,
     ) -> Result<Self::ByteStream, ServiceError>;
 
@@ -193,17 +194,24 @@ pub trait PushService {
 
     async fn get_attachment_by_id(
         &mut self,
-        id: u64,
+        id: &str,
     ) -> Result<Self::ByteStream, ServiceError> {
         let path = format!("{}{}", ATTACHMENT_UPLOAD_PATH, id);
-        self.get_from_cdn(&path).await
+        self.get_from_cdn(0, &path).await
     }
 
     async fn get_attachment(
         &mut self,
         ptr: &AttachmentPointer,
     ) -> Result<Self::ByteStream, ServiceError> {
-        self.get_attachment_by_id(ptr.id()).await
+        match ptr.attachment_identifier.as_ref().unwrap() {
+            AttachmentIdentifier::CdnId(id) => {
+                self.get_attachment_by_id(&format!("{}", id)).await
+            },
+            AttachmentIdentifier::CdnKey(key) => {
+                self.get_attachment_by_id(key).await
+            },
+        }
     }
 
     async fn get_messages(
@@ -267,6 +275,7 @@ impl PushService for PanicingPushService {
 
     async fn get_from_cdn(
         &mut self,
+        _cdn_id: u64,
         _path: &str,
     ) -> Result<Self::ByteStream, ServiceError> {
         unimplemented!()