Skip to content
This repository was archived by the owner on Nov 1, 2023. It is now read-only.

Commit 9a2ea7b

Browse files
chkeitachkeita
committed
unmanaged node prototype
1 parent ecc0992 commit 9a2ea7b

File tree

16 files changed

+214
-75
lines changed

16 files changed

+214
-75
lines changed

‎src/ApiService/ApiService/Functions/AgentRegistration.cs‎

Lines changed: 27 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
1-
using System.Web;
2-
using Azure.Storage.Sas;
1+
using Azure.Storage.Sas;
32
using Microsoft.Azure.Functions.Worker;
43
using Microsoft.Azure.Functions.Worker.Http;
54

@@ -31,17 +30,13 @@ public Async.Task<HttpResponseData> Run(
3130
});
3231

3332
private async Async.Task<HttpResponseData> Get(HttpRequestData req) {
34-
var uri = HttpUtility.ParseQueryString(req.Url.Query);
35-
var rawMachineId = uri["machine_id"];
36-
if (rawMachineId is null || !Guid.TryParse(rawMachineId, out var machineId)) {
37-
return await _context.RequestHandling.NotOk(
38-
req,
39-
new Error(
40-
ErrorCode.INVALID_REQUEST,
41-
new string[] { "'machine_id' query parameter must be provided" }),
42-
"agent registration");
33+
var request = await RequestHandling.ParseUri<AgentRegistrationGet>(req);
34+
if (!request.IsOk) {
35+
return await _context.RequestHandling.NotOk(req, request.ErrorV, "agent registration");
4336
}
4437

38+
var machineId = request.OkV.MachineId;
39+
4540
var agentNode = await _context.NodeOperations.GetByMachineId(machineId);
4641
if (agentNode is null) {
4742
return await _context.RequestHandling.NotOk(
@@ -82,32 +77,18 @@ private async Async.Task<AgentRegistrationResponse> CreateRegistrationResponse(S
8277
WorkQueue: workQueue);
8378
}
8479

80+
// todo: add agent registration post
8581
private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
86-
var uri = HttpUtility.ParseQueryString(req.Url.Query);
87-
var rawMachineId = uri["machine_id"];
88-
if (rawMachineId is null || !Guid.TryParse(rawMachineId, out var machineId)) {
89-
return await _context.RequestHandling.NotOk(
90-
req,
91-
new Error(
92-
ErrorCode.INVALID_REQUEST,
93-
new string[] { "'machine_id' query parameter must be provided" }),
94-
"agent registration");
95-
}
96-
97-
var rawPoolName = uri["pool_name"];
98-
if (rawPoolName is null || !PoolName.TryParse(rawPoolName, out var poolName)) {
99-
return await _context.RequestHandling.NotOk(
100-
req,
101-
new Error(
102-
ErrorCode.INVALID_REQUEST,
103-
new string[] { "'pool_name' query parameter must be provided" }),
104-
"agent registration");
82+
var request = await RequestHandling.ParseUri<AgentRegistrationPost>(req);
83+
if (!request.IsOk) {
84+
return await _context.RequestHandling.NotOk(req, request.ErrorV, "agent registration");
10585
}
10686

107-
var rawScalesetId = uri["scaleset_id"];
108-
var scalesetId = rawScalesetId is null ? (Guid?)null : Guid.Parse(rawScalesetId);
109-
110-
var version = uri["version"] ?? "1.0.0";
87+
var machineId = request.OkV.MachineId;
88+
var poolName = request.OkV.PoolName;
89+
var scalesetId = request.OkV.ScalesetId;
90+
var version = request.OkV.Version;
91+
var os = request.OkV.Os;
11192

11293
_log.Info($"registration request: machine_id: {machineId} pool_name: {poolName} scaleset_id: {scalesetId} version: {version}");
11394
var poolResult = await _context.PoolOperations.GetByName(poolName);
@@ -127,12 +108,23 @@ private async Async.Task<HttpResponseData> Post(HttpRequestData req) {
127108
await _context.NodeOperations.Delete(existingNode);
128109
}
129110

111+
if (os != null && pool.Os != os) {
112+
return await _context.RequestHandling.NotOk(
113+
req,
114+
new Error(
115+
Code: ErrorCode.INVALID_REQUEST,
116+
Errors: new[] { $"OS mismatch: pool '{poolName}' is configured for '{pool.Os}', but agent is running '{os}'" }),
117+
"agent registration");
118+
}
119+
130120
var node = new Service.Node(
131121
PoolName: poolName,
132122
PoolId: pool.PoolId,
133123
MachineId: machineId,
134124
ScalesetId: scalesetId,
135-
Version: version);
125+
Version: version,
126+
Os: os ?? pool.Os
127+
);
136128

137129
await _context.NodeOperations.Replace(node);
138130

‎src/ApiService/ApiService/Functions/GetPoolConfig.cs‎

Lines changed: 58 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,58 @@
1+
using System.Net;
2+
using System.Security.Claims;
3+
using System.Threading.Tasks;
4+
using Microsoft.Azure.Functions.Worker;
5+
using Microsoft.Azure.Functions.Worker.Http;
6+
7+
namespace Microsoft.OneFuzz.Service.Functions;
8+
9+
public class GetPoolConfig {
10+
private readonly ILogTracer _log;
11+
private readonly IEndpointAuthorization _auth;
12+
private readonly IOnefuzzContext _context;
13+
14+
public GetPoolConfig(ILogTracer log, IEndpointAuthorization auth, IOnefuzzContext context) {
15+
_log = log;
16+
_auth = auth;
17+
_context = context;
18+
}
19+
20+
[Function("GetPoolConfig")]
21+
public Async.Task<HttpResponseData> Run([HttpTrigger(AuthorizationLevel.Anonymous, "GET", Route = "pool/getconfig")] HttpRequestData req, ClaimsPrincipal principal)
22+
=> _auth.CallIfUser(req, r => r.Method switch {
23+
"GET" => Get(r),
24+
var m => throw new InvalidOperationException("Unsupported HTTP method {m}"),
25+
});
26+
27+
private async Task<HttpResponseData> Get(HttpRequestData req) {
28+
var request = await RequestHandling.ParseRequest<PoolSearch>(req);
29+
if (!request.IsOk) {
30+
return await _context.RequestHandling.NotOk(req, request.ErrorV, "pool get config");
31+
}
32+
33+
var search = request.OkV;
34+
OneFuzzResult<Service.Pool> poolResult;
35+
if (search.PoolId is Guid poolId) {
36+
poolResult = await _context.PoolOperations.GetById(poolId);
37+
} else if (search.Name is PoolName name) {
38+
poolResult = await _context.PoolOperations.GetByName(name);
39+
} else {
40+
return await _context.RequestHandling.NotOk(req, new Error(ErrorCode.INVALID_REQUEST, new[] { "missing pool name or id" }), "pool get config");
41+
}
42+
43+
if (!poolResult.IsOk) {
44+
return await _context.RequestHandling.NotOk(req, poolResult.ErrorV, context: search.ToString());
45+
}
46+
47+
if (!poolResult.OkV.Managed) {
48+
return await _context.RequestHandling.NotOk(req, new Error(ErrorCode.INVALID_REQUEST, new[] { "Pool is a managed pool" }), context: search.ToString());
49+
}
50+
var poolConfig = await _context.Extensions.CreatePoolConfig(poolResult.OkV);
51+
52+
var response = req.CreateResponse(HttpStatusCode.OK);
53+
await response.WriteAsJsonAsync(poolConfig);
54+
return response;
55+
}
56+
57+
58+
}

‎src/ApiService/ApiService/OneFuzzTypes/Model.cs‎

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -96,6 +96,7 @@ public record Node
9696
DateTimeOffset? Heartbeat = null,
9797
DateTimeOffset? InitializedAt = null,
9898
NodeState State = NodeState.Init,
99+
Os? Os = null,
99100

100101
Guid? ScalesetId = null,
101102
bool ReimageRequested = false,
@@ -162,7 +163,11 @@ public sealed override string ToString() {
162163
}
163164
};
164165

165-
public record UserInfo(Guid? ApplicationId, Guid? ObjectId, String? Upn);
166+
public record UserInfo(Guid? ApplicationId, Guid? ObjectId, String? Upn, List<string> Roles) {
167+
public static UserInfo Create() {
168+
return new UserInfo(null, null, null, new List<string>());
169+
}
170+
}
166171

167172
public record TaskDetails(
168173
TaskType Type,

‎src/ApiService/ApiService/OneFuzzTypes/Requests.cs‎

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -252,3 +252,17 @@ public record WebhookUpdate(
252252
string? SecretToken,
253253
WebhookMessageFormat? MessageFormat
254254
);
255+
256+
257+
public record AgentRegistrationGet(
258+
Guid MachineId
259+
);
260+
261+
262+
public record AgentRegistrationPost(
263+
PoolName PoolName,
264+
Guid? ScalesetId,
265+
Guid MachineId,
266+
string Version,
267+
Os? Os
268+
);

‎src/ApiService/ApiService/OneFuzzTypes/Responses.cs‎

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -103,6 +103,7 @@ public record PoolGetResult(
103103
AgentConfig? Config,
104104
List<WorkSetSummary>? WorkQueue,
105105
List<ScalesetSummary>? ScalesetSummary
106+
//List<Node>? UnmanagedNodes
106107
) : BaseResponse();
107108

108109
public record ScalesetResponse(

‎src/ApiService/ApiService/UserCredentials.cs‎

Lines changed: 23 additions & 16 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
1-
using System.Net.Http.Headers;
1+
using System.IdentityModel.Tokens.Jwt;
2+
using System.Net.Http.Headers;
23
using System.Threading.Tasks;
34
using Microsoft.Azure.Functions.Worker.Http;
45
using Microsoft.IdentityModel.Tokens;
@@ -15,10 +16,12 @@ public interface IUserCredentials {
1516
public class UserCredentials : IUserCredentials {
1617
ILogTracer _log;
1718
IConfigOperations _instanceConfig;
19+
private JwtSecurityTokenHandler _tokenHandler;
1820

1921
public UserCredentials(ILogTracer log, IConfigOperations instanceConfig) {
2022
_log = log;
2123
_instanceConfig = instanceConfig;
24+
_tokenHandler = new JwtSecurityTokenHandler();
2225
}
2326

2427
public string? GetBearerToken(HttpRequestData req) {
@@ -57,6 +60,8 @@ from t in r.AllowedAadTenants
5760
}
5861

5962
public virtual async Task<OneFuzzResult<UserInfo>> ParseJwtToken(HttpRequestData req) {
63+
64+
6065
var authToken = GetAuthToken(req);
6166
if (authToken is null) {
6267
return OneFuzzResult<UserInfo>.Error(ErrorCode.INVALID_REQUEST, new[] { "unable to find authorization token" });
@@ -65,22 +70,24 @@ public virtual async Task<OneFuzzResult<UserInfo>> ParseJwtToken(HttpRequestData
6570
var allowedTenants = await GetAllowedTenants();
6671
if (allowedTenants.IsOk) {
6772
if (allowedTenants.OkV is not null && allowedTenants.OkV.Contains(token.Issuer)) {
68-
Guid? applicationId = (
69-
from t in token.Claims
70-
where t.Type == "appId"
71-
select (Guid.Parse(t.Value))).FirstOrDefault();
72-
73-
Guid? objectId = (
74-
from t in token.Claims
75-
where t.Type == "oid"
76-
select (Guid.Parse(t.Value))).FirstOrDefault();
77-
78-
string? upn = (
79-
from t in token.Claims
80-
where t.Type == "upn"
81-
select t.Value).FirstOrDefault();
73+
var userInfo =
74+
token.Payload.Claims.Aggregate(UserInfo.Create(), (acc, claim) => {
75+
switch (claim.Type) {
76+
case "oid":
77+
return acc with { ObjectId = Guid.Parse(claim.Value) };
78+
case "appId":
79+
return acc with { ApplicationId = Guid.Parse(claim.Value) };
80+
case "upn":
81+
return acc with { Upn = claim.Value };
82+
case "roles":
83+
acc.Roles.Add(claim.Value);
84+
return acc;
85+
default:
86+
return acc;
87+
}
88+
});
8289

83-
return OneFuzzResult<UserInfo>.Ok(new(applicationId, objectId, upn));
90+
return OneFuzzResult<UserInfo>.Ok(userInfo);
8491
} else {
8592
_log.Error($"issuer not from allowed tenant: {token.Issuer} - {allowedTenants}");
8693
return OneFuzzResult<UserInfo>.Error(ErrorCode.INVALID_REQUEST, new[] { "unauthorized AAD issuer" });

‎src/ApiService/ApiService/onefuzzlib/EndpointAuthorization.cs‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,7 @@
55
namespace Microsoft.OneFuzz.Service;
66

77
public interface IEndpointAuthorization {
8+
89
Async.Task<HttpResponseData> CallIfAgent(
910
HttpRequestData req,
1011
Func<HttpRequestData, Async.Task<HttpResponseData>> method)
@@ -183,6 +184,9 @@ private GroupMembershipChecker CreateGroupMembershipChecker(InstanceConfig confi
183184
}
184185

185186
public async Async.Task<bool> IsAgent(UserInfo tokenData) {
187+
// todo: handle unmanaged node here
188+
// check if the request is comming from a geristered app with apprile agent
189+
186190
if (tokenData.ObjectId != null) {
187191
var scalesets = _context.ScalesetOperations.GetByObjectId(tokenData.ObjectId.Value);
188192
if (await scalesets.AnyAsync()) {
@@ -202,6 +206,10 @@ public async Async.Task<bool> IsAgent(UserInfo tokenData) {
202206
return true;
203207
}
204208

209+
// if (tokenData.Roles.Contains("unmanagedNode")) {
210+
// return true;
211+
// }
212+
205213
return false;
206214
}
207215
}

‎src/ApiService/ApiService/onefuzzlib/Extension.cs‎

Lines changed: 13 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@ public interface IExtensions {
1212
Async.Task<IList<VirtualMachineScaleSetExtensionData>> FuzzExtensions(Pool pool, Scaleset scaleset);
1313

1414
Async.Task<Dictionary<string, VirtualMachineExtensionData>> ReproExtensions(AzureLocation region, Os reproOs, Guid reproId, ReproConfig reproConfig, Container? setupContainer);
15+
16+
Async.Task<Uri?> BuildPoolConfig(Pool pool);
17+
Task<AgentConfig> CreatePoolConfig(Pool pool);
1518
Task<IList<VMExtensionWrapper>> ProxyManagerExtensions(string region, Guid proxyId);
1619
}
1720

@@ -211,6 +214,14 @@ public static VMExtensionWrapper GenevaExtension(AzureLocation region) {
211214

212215

213216
public async Async.Task<Uri?> BuildPoolConfig(Pool pool) {
217+
var config = await CreatePoolConfig(pool);
218+
219+
var fileName = $"{pool.Name}/config.json";
220+
await _context.Containers.SaveBlob(new Container("vm-scripts"), fileName, (JsonSerializer.Serialize(config, EntityConverter.GetJsonSerializerOptions())), StorageType.Config);
221+
return await ConfigUrl(new Container("vm-scripts"), fileName, false);
222+
}
223+
224+
public async Task<AgentConfig> CreatePoolConfig(Pool pool) {
214225
var instanceId = await _context.Containers.GetInstanceId();
215226

216227
var queueSas = await _context.Queue.GetQueueSas("node-heartbeat", StorageType.Config, QueueSasPermissions.Add);
@@ -223,11 +234,8 @@ public static VMExtensionWrapper GenevaExtension(AzureLocation region) {
223234
MicrosoftTelemetryKey: _context.ServiceConfiguration.OneFuzzTelemetry,
224235
MultiTenantDomain: _context.ServiceConfiguration.MultiTenantDomain,
225236
InstanceId: instanceId
226-
);
227-
228-
var fileName = $"{pool.Name}/config.json";
229-
await _context.Containers.SaveBlob(new Container("vm-scripts"), fileName, (JsonSerializer.Serialize(config, EntityConverter.GetJsonSerializerOptions())), StorageType.Config);
230-
return await ConfigUrl(new Container("vm-scripts"), fileName, false);
237+
);
238+
return config;
231239
}
232240

233241

‎src/ApiService/IntegrationTests/JobsTests.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,7 @@ public async Async.Task Post_CreatesJob_AndContainer() {
157157
var func = new Jobs(auth, Context);
158158

159159
// need user credentials to put into the job object
160-
var userInfo = new UserInfo(Guid.NewGuid(), Guid.NewGuid(), "upn");
160+
var userInfo = new UserInfo(Guid.NewGuid(), Guid.NewGuid(), "upn", new List<string>());
161161
Context.UserCredentials = new TestUserCredentials(Logger, Context.ConfigOperations, OneFuzzResult.Ok(userInfo));
162162

163163
var result = await func.Run(TestHttpRequestData.FromJson("POST", _config));

‎src/ApiService/IntegrationTests/NodeTests.cs‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -161,7 +161,7 @@ await Context.InsertAll(
161161
var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
162162

163163
// override the found user credentials
164-
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: Guid.NewGuid(), "upn");
164+
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: Guid.NewGuid(), "upn", new List<string>());
165165
Context.UserCredentials = new TestUserCredentials(Logger, Context.ConfigOperations, OneFuzzResult<UserInfo>.Ok(userInfo));
166166

167167
var req = new NodeGet(MachineId: _machineId);
@@ -189,7 +189,7 @@ await Context.InsertAll(
189189
var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
190190

191191
// override the found user credentials
192-
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: Guid.NewGuid(), "upn");
192+
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: Guid.NewGuid(), "upn", new List<string>());
193193
Context.UserCredentials = new TestUserCredentials(Logger, Context.ConfigOperations, OneFuzzResult<UserInfo>.Ok(userInfo));
194194

195195
var req = new NodeGet(MachineId: _machineId);
@@ -219,7 +219,7 @@ await Context.InsertAll(
219219
var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
220220

221221
// override the found user credentials
222-
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: userObjectId, "upn");
222+
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: userObjectId, "upn", new List<string>());
223223
Context.UserCredentials = new TestUserCredentials(Logger, Context.ConfigOperations, OneFuzzResult<UserInfo>.Ok(userInfo));
224224

225225
var req = new NodeGet(MachineId: _machineId);
@@ -250,7 +250,7 @@ await Context.InsertAll(
250250
var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
251251

252252
// override the found user credentials
253-
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: userObjectId, "upn");
253+
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: userObjectId, "upn", new List<string>());
254254
Context.UserCredentials = new TestUserCredentials(Logger, Context.ConfigOperations, OneFuzzResult<UserInfo>.Ok(userInfo));
255255

256256
var req = new NodeGet(MachineId: _machineId);
@@ -279,7 +279,7 @@ await Context.InsertAll(
279279
var auth = new TestEndpointAuthorization(RequestType.User, Logger, Context);
280280

281281
// override the found user credentials
282-
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: Guid.NewGuid(), "upn");
282+
var userInfo = new UserInfo(ApplicationId: Guid.NewGuid(), ObjectId: Guid.NewGuid(), "upn", new List<string>());
283283
Context.UserCredentials = new TestUserCredentials(Logger, Context.ConfigOperations, OneFuzzResult<UserInfo>.Ok(userInfo));
284284

285285
// all of these operations use NodeGet

0 commit comments

Comments
 (0)