Pass the task tags to the agent (#2895)

* Pass the task tags to the agent * build fix * build fix * format
microsoft · Apr 10, 2023 · e35b9fa · e35b9fa
1 parent 35e9eff
commit e35b9fa
Show file tree

Hide file tree

Showing 6 changed files with 11 additions and 2 deletions.
diff --git a/src/ApiService/ApiService/OneFuzzTypes/Model.cs b/src/ApiService/ApiService/OneFuzzTypes/Model.cs
@@ -972,7 +972,8 @@ public record TaskUnitConfig(
     TaskType TaskType,
     string? InstanceTelemetryKey,
     string? MicrosoftTelemetryKey,
-    Uri HeartbeatQueue
+    Uri HeartbeatQueue,
+    Dictionary<string, string> Tags
     ) {
     public Uri? inputQueue { get; set; }
     public String? SupervisorExe { get; set; }

diff --git a/src/ApiService/ApiService/onefuzzlib/Config.cs b/src/ApiService/ApiService/onefuzzlib/Config.cs
@@ -70,7 +70,8 @@ private static BlobContainerSasPermissions ConvertPermissions(ContainerPermissio
             TaskType: task.Config.Task.Type,
             InstanceTelemetryKey: _serviceConfig.ApplicationInsightsInstrumentationKey,
             MicrosoftTelemetryKey: _serviceConfig.OneFuzzTelemetry,
-            HeartbeatQueue: await _queue.GetQueueSas("task-heartbeat", StorageType.Config, QueueSasPermissions.Add) ?? throw new Exception("unable to get heartbeat queue sas")
+            HeartbeatQueue: await _queue.GetQueueSas("task-heartbeat", StorageType.Config, QueueSasPermissions.Add) ?? throw new Exception("unable to get heartbeat queue sas"),
+            Tags: task.Config.Tags ?? new Dictionary<string, string>()
         );
 
         if (definition.MonitorQueue != null) {

diff --git a/src/agent/onefuzz-task/src/local/common.rs b/src/agent/onefuzz-task/src/local/common.rs
@@ -266,6 +266,7 @@ pub async fn build_local_context(
         microsoft_telemetry_key: None,
         logs: None,
         min_available_memory_mb: 0,
+        tags: Default::default(),
         from_agent_to_task_endpoint: "/".to_string(),
         from_task_to_agent_endpoint: "/".to_string(),
     };

diff --git a/src/agent/onefuzz-task/src/tasks/config.rs b/src/agent/onefuzz-task/src/tasks/config.rs
@@ -18,6 +18,7 @@ use onefuzz_telemetry::{
 use reqwest::Url;
 use serde::{self, Deserialize};
 use std::{
+    collections::HashMap,
     path::{Path, PathBuf},
     sync::Arc,
     time::Duration,
@@ -68,6 +69,9 @@ pub struct CommonConfig {
 
     pub machine_identity: MachineIdentity,
 
+    #[serde(default)]
+    pub tags: HashMap<String, String>,
+
     pub from_agent_to_task_endpoint: String,
     pub from_task_to_agent_endpoint: String,
 }

diff --git a/src/agent/onefuzz-task/src/tasks/fuzz/generator.rs b/src/agent/onefuzz-task/src/tasks/fuzz/generator.rs
@@ -308,6 +308,7 @@ mod tests {
                     machine_name: "test".to_string(),
                     scaleset_name: None,
                 },
+                tags: Default::default(),
                 from_agent_to_task_endpoint: "/".to_string(),
                 from_task_to_agent_endpoint: "/".to_string(),
             },

diff --git a/src/agent/onefuzz-task/src/tasks/fuzz/supervisor.rs b/src/agent/onefuzz-task/src/tasks/fuzz/supervisor.rs
@@ -399,6 +399,7 @@ mod tests {
                     machine_name: "test".to_string(),
                     scaleset_name: None,
                 },
+                tags: Default::default(),
                 from_agent_to_task_endpoint: "/".to_string(),
                 from_task_to_agent_endpoint: "/".to_string(),
             },