bug fixes related to the unmanaged nodes

- fix the token request in the client - adding a is_unmnaged field to the agentConfig - fix typo in the appId claim type - fix typo in the UnmanagedNode claim value - fix query PoolOperation.GetByClientId
microsoft · Nov 16, 2022 · f36ddc7 · f36ddc7
1 parent 40c85b8
commit f36ddc7
Show file tree

Hide file tree

Showing 10 changed files with 31 additions and 27 deletions.
diff --git a/src/ApiService/ApiService/Functions/Pool.cs b/src/ApiService/ApiService/Functions/Pool.cs
@@ -133,7 +133,8 @@ private async Task<PoolGetResult> Populate(PoolGetResult p, bool skipSummaries =
                     HeartbeatQueue: queueSas,
                     InstanceId: instanceId,
                     ClientCredentials: null,
-                    MultiTenantDomain: _context.ServiceConfiguration.MultiTenantDomain)
+                    MultiTenantDomain: _context.ServiceConfiguration.MultiTenantDomain,
+                    IsUnmanaged: !p.Managed)
         };
     }
 }
diff --git a/src/ApiService/ApiService/OneFuzzTypes/Model.cs b/src/ApiService/ApiService/OneFuzzTypes/Model.cs
@@ -678,12 +678,11 @@ public record AgentConfig(
     string? InstanceTelemetryKey,
     string? MicrosoftTelemetryKey,
     string? MultiTenantDomain,
-    Guid InstanceId
+    Guid InstanceId,
+    bool? IsUnmanaged
 );
 
 
-
-
 public record Vm(
     string Name,
     Region Region,

diff --git a/src/ApiService/ApiService/UserCredentials.cs b/src/ApiService/ApiService/UserCredentials.cs
@@ -77,7 +77,7 @@ public virtual async Task<OneFuzzResult<UserAuthInfo>> ParseJwtToken(HttpRequest
                             switch (claim.Type) {
                                 case "oid":
                                     return acc with { UserInfo = acc.UserInfo with { ObjectId = Guid.Parse(claim.Value) } };
-                                case "appId":
+                                case "appid":
                                     return acc with { UserInfo = acc.UserInfo with { ApplicationId = Guid.Parse(claim.Value) } };
                                 case "upn":
                                     return acc with { UserInfo = acc.UserInfo with { Upn = claim.Value } };
@@ -88,7 +88,6 @@ public virtual async Task<OneFuzzResult<UserAuthInfo>> ParseJwtToken(HttpRequest
                                     return acc;
                             }
                         });
-
                     return OneFuzzResult<UserAuthInfo>.Ok(userInfo);
                 } else {
                     var tenantsStr = allowedTenants.OkV is null ? "null" : String.Join(';', allowedTenants.OkV!);

diff --git a/src/ApiService/ApiService/onefuzzlib/EndpointAuthorization.cs b/src/ApiService/ApiService/onefuzzlib/EndpointAuthorization.cs
@@ -30,8 +30,7 @@ public class EndpointAuthorization : IEndpointAuthorization {
     private readonly IOnefuzzContext _context;
     private readonly ILogTracer _log;
     private readonly GraphServiceClient _graphClient;
-
-    private static readonly HashSet<string> AgentRoles = new HashSet<string> { "UnmamagedNode", "ManagedNode" };
+    private static readonly HashSet<string> AgentRoles = new HashSet<string> { "UnmanagedNode", "ManagedNode" };
 
     public EndpointAuthorization(IOnefuzzContext context, ILogTracer log, GraphServiceClient graphClient) {
         _context = context;
@@ -46,8 +45,8 @@ public virtual async Async.Task<HttpResponseData> CallIf(HttpRequestData req, Fu
             return await _context.RequestHandling.NotOk(req, tokenResult.ErrorV, "token verification", HttpStatusCode.Unauthorized);
         }
 
-        var token = tokenResult.OkV;
-        if (await IsUser(token)) {
+        var token = tokenResult.OkV.UserInfo;
+        if (await IsUser(tokenResult.OkV)) {
             if (!allowUser) {
                 return await Reject(req, tokenResult.OkV.UserInfo);
             }
@@ -58,7 +57,7 @@ public virtual async Async.Task<HttpResponseData> CallIf(HttpRequestData req, Fu
             }
         }
 
-        if (await IsAgent(token) && !allowAgent) {
+        if (await IsAgent(tokenResult.OkV) && !allowAgent) {
             return await Reject(req, tokenResult.OkV.UserInfo);
         }
 
@@ -201,7 +200,9 @@ public async Async.Task<bool> IsAgent(UserAuthInfo authInfo) {
             }
 
             var principalId = await _context.Creds.GetScalesetPrincipalId();
-            return principalId == tokenData.ObjectId;
+            if (principalId == tokenData.ObjectId) {
+                return true;
+            }
         }
 
         if (!tokenData.ApplicationId.HasValue) {

diff --git a/src/ApiService/ApiService/onefuzzlib/PoolOperations.cs b/src/ApiService/ApiService/onefuzzlib/PoolOperations.cs
@@ -89,7 +89,7 @@ public async Task<bool> ScheduleWorkset(Pool pool, WorkSet workSet) {
     }
 
     public IAsyncEnumerable<Pool> GetByClientId(Guid clientId) {
-        return QueryAsync(filter: TableClient.CreateQueryFilter($"client_id eq {clientId}"));
+        return QueryAsync(filter: $"client_id eq '{clientId}'");
     }
 
     public string GetPoolQueue(Guid poolId)

diff --git a/src/agent/onefuzz-agent/src/config.rs b/src/agent/onefuzz-agent/src/config.rs
@@ -34,12 +34,15 @@ pub struct StaticConfig {
     pub heartbeat_queue: Option<Url>,
 
     pub instance_id: Uuid,
+
+    #[serde(default)]
+    pub is_unmanaged: bool,
 }
 
 // Temporary shim type to bridge the current service-provided config.
 #[derive(Clone, Debug, Deserialize, Eq, PartialEq)]
 struct RawStaticConfig {
-    pub credentials: Option<ClientCredentials>,
+    pub client_credentials: Option<ClientCredentials>,
 
     pub pool_name: String,
 
@@ -54,13 +57,16 @@ struct RawStaticConfig {
     pub heartbeat_queue: Option<Url>,
 
     pub instance_id: Uuid,
+
+    #[serde(default)]
+    pub is_unmanaged: bool,
 }
 
 impl StaticConfig {
     pub fn new(data: &[u8]) -> Result<Self> {
         let config: RawStaticConfig = serde_json::from_slice(data)?;
 
-        let credentials = match config.credentials {
+        let credentials = match config.client_credentials {
             Some(client) => client.into(),
             None => {
                 // Remove trailing `/`, which is treated as a distinct resource.
@@ -83,6 +89,7 @@ impl StaticConfig {
             instance_telemetry_key: config.instance_telemetry_key,
             heartbeat_queue: config.heartbeat_queue,
             instance_id: config.instance_id,
+            is_unmanaged: config.is_unmanaged,
         };
 
         Ok(config)
@@ -103,6 +110,7 @@ impl StaticConfig {
         let multi_tenant_domain = std::env::var("ONEFUZZ_MULTI_TENANT_DOMAIN").ok();
         let onefuzz_url = Url::parse(&std::env::var("ONEFUZZ_URL")?)?;
         let pool_name = std::env::var("ONEFUZZ_POOL")?;
+        let is_unmanaged = std::env::var("ONEFUZZ_IS_UNMANAGED").is_ok();
 
         let heartbeat_queue = if let Ok(key) = std::env::var("ONEFUZZ_HEARTBEAT") {
             Some(Url::parse(&key)?)
@@ -142,6 +150,7 @@ impl StaticConfig {
             microsoft_telemetry_key,
             heartbeat_queue,
             instance_id,
+            is_unmanaged,
         })
     }
 
@@ -213,7 +222,8 @@ impl Registration {
             .append_pair("machine_id", &machine_id.to_string())
             .append_pair("machine_name", &machine_name)
             .append_pair("pool_name", &config.pool_name)
-            .append_pair("version", env!("ONEFUZZ_VERSION"));
+            .append_pair("version", env!("ONEFUZZ_VERSION"))
+            .append_pair("os", std::env::consts::OS);
 
         if managed {
             let scaleset = onefuzz::machine_id::get_scaleset_name().await?;

diff --git a/src/agent/onefuzz-agent/src/main.rs b/src/agent/onefuzz-agent/src/main.rs
@@ -277,10 +277,10 @@ async fn run_agent(config: StaticConfig) -> Result<()> {
     let registration = match config::Registration::load_existing(config.clone()).await {
         Ok(registration) => registration,
         Err(_) => {
-            if scaleset.is_some() {
-                config::Registration::create_managed(config.clone()).await?
-            } else {
+            if config.is_unmanaged {
                 config::Registration::create_unmanaged(config.clone()).await?
+            } else {
+                config::Registration::create_managed(config.clone()).await?
             }
         }
     };

diff --git a/src/agent/onefuzz/src/auth.rs b/src/agent/onefuzz/src/auth.rs
@@ -134,7 +134,6 @@ impl ClientCredentials {
 
         let response = reqwest::Client::new()
             .post(url)
-            .header("Content-Length", "0")
             .form(&[
                 ("client_id", self.client_id.to_hyphenated().to_string()),
                 ("client_secret", self.client_secret.expose_ref().to_string()),

diff --git a/src/cli/onefuzz/api.py b/src/cli/onefuzz/api.py
@@ -1268,13 +1268,7 @@ def get_config(self, pool_name: primitives.PoolName) -> models.AgentConfig:
         if pool.config is None:
             raise Exception("Missing AgentConfig in response")
 
-        config = pool.config
-        config.client_credentials = models.ClientCredentials(  # nosec - bandit consider this a hard coded password
-            client_id=pool.client_id,
-            client_secret="<client secret>",
-        )
-
-        return config
+        return pool.config
 
     def shutdown(self, name: str, *, now: bool = False) -> responses.BoolResult:
         expanded_name = self._disambiguate(

diff --git a/src/pytypes/onefuzztypes/models.py b/src/pytypes/onefuzztypes/models.py
@@ -339,6 +339,7 @@ class AgentConfig(BaseModel):
     microsoft_telemetry_key: Optional[str]
     multi_tenant_domain: Optional[str]
     instance_id: UUID
+    is_unmanaged: Optional[bool]
 
 
 class TaskUnitConfig(BaseModel):