From 35dfeb665ae8c0cb1f82156a4f0f9d4104c65a09 Mon Sep 17 00:00:00 2001
From: Julius Liu <juliusl@microsoft.com>
Date: Thu, 22 Jun 2023 14:00:51 -0700
Subject: [PATCH] fix. filter streamable descs by platform

---
 src/content/descriptor.rs | 50 +++++++++++++++++++++++++++++++++++----
 src/content/platform.rs   |  2 +-
 src/plugins/teleport.rs   | 17 ++++++++++++-
 3 files changed, 63 insertions(+), 6 deletions(-)

diff --git a/src/content/descriptor.rs b/src/content/descriptor.rs
index 0008202..1b32c34 100644
--- a/src/content/descriptor.rs
+++ b/src/content/descriptor.rs
@@ -2,7 +2,7 @@ use std::collections::BTreeMap;
 
 use lifec::prelude::{AttributeIndex, Component, DefaultVecStorage, ThunkContext};
 use serde::{Deserialize, Serialize};
-use tracing::{trace, debug};
+use tracing::{debug, trace};
 
 use super::Platform;
 
@@ -88,10 +88,23 @@ impl Descriptor {
     /// ```
     ///
     pub fn try_parse_streamable_descriptor(&self) -> Option<Self> {
-        if !self.annotations.as_ref().map(|a| a.contains_key("streaming.mediaType")).unwrap_or_default() {
+        if !self
+            .annotations
+            .as_ref()
+            .map(|a| a.contains_key("streaming.mediaType"))
+            .unwrap_or_default()
+        {
+            let mut desc = self.clone();
             debug!("Updated artifact manifest detected, skipping annotation parsing");
+            if let Some(platform) = self.try_parse_streamable_platform() {
+                desc.platform = Some(Platform {
+                    architecture: platform.platform_arch.unwrap_or_default(),
+                    os: platform.platform_os.unwrap_or_default(),
+                    variant: None,
+                });
+            }
 
-            return Some(self.clone());
+            return Some(desc);
         }
 
         if let Some(annotations) = self
@@ -108,7 +121,11 @@ impl Descriptor {
                     annotations: None,
                     urls: None,
                     data: None,
-                    platform: None,
+                    platform: Some(Platform {
+                        architecture: streaming_desc.platform_arch.unwrap_or_default(),
+                        os: streaming_desc.platform_os.unwrap_or_default(),
+                        variant: None,
+                    }),
                 }),
                 Err(err) => {
                     trace!(
@@ -122,6 +139,18 @@ impl Descriptor {
             None
         }
     }
+
+    fn try_parse_streamable_platform(&self) -> Option<StreamingPlatform> {
+        if let Some(annotations) = self
+            .annotations
+            .as_ref()
+            .and_then(|a| serde_json::to_string(a).ok())
+        {
+            serde_json::from_str::<StreamingPlatform>(annotations.as_str()).ok()
+        } else {
+            None
+        }
+    }
 }
 
 #[allow(dead_code)]
@@ -143,6 +172,19 @@ struct StreamingDescriptor {
     platform_arch: Option<String>,
 }
 
+/// Struct to parse stremaing platform,
+///
+#[allow(dead_code)]
+#[derive(Deserialize)]
+struct StreamingPlatform {
+    #[serde(rename = "streaming.format")]
+    format: String,
+    #[serde(rename = "streaming.platform.os")]
+    platform_os: Option<String>,
+    #[serde(rename = "streaming.platform.arch")]
+    platform_arch: Option<String>,
+}
+
 #[allow(unused_imports)]
 mod tests {
     use serde_json::json;
diff --git a/src/content/platform.rs b/src/content/platform.rs
index c01f631..5774d63 100644
--- a/src/content/platform.rs
+++ b/src/content/platform.rs
@@ -12,5 +12,5 @@ pub struct Platform {
     pub os: String,
     /// Operating system variant
     #[serde(skip_serializing_if = "Option::is_none")]
-    variant: Option<String>,
+    pub variant: Option<String>,
 }
\ No newline at end of file
diff --git a/src/plugins/teleport.rs b/src/plugins/teleport.rs
index 1ac2ca6..3a75fd0 100644
--- a/src/plugins/teleport.rs
+++ b/src/plugins/teleport.rs
@@ -58,7 +58,22 @@ impl Plugin for Teleport {
 
                         let streamable = list.find_streamable_descriptors();
 
-                        let digest = if let Some(streamable_desc) = streamable.first() {
+                        let arch = if std::env::consts::ARCH == "x86_64" {
+                            "amd64"
+                        } else {
+                            std::env::consts::ARCH
+                        };
+                        let os = std::env::consts::OS;
+
+                        info!("Filtering streamable descriptors w/ os - {os}, arch - {arch}");
+
+                        let digest = if let Some(streamable_desc) = streamable.iter().find(|s| {
+                            if let Some(platform) = s.platform.as_ref() {
+                                platform.architecture == arch && platform.os == os
+                            } else {
+                                false
+                            }
+                        }) {
                             info!("Streamable descriptor was found");
                             streamable_desc.digest.clone()
                         } else {