proxy forward test hooks (#1924)

fixing some bugs

Co-authored-by: stas <statis@microsoft.com>

src/ApiService/ApiService/OneFuzzTypes/Model.cs

@@ -103,8 +103,8 @@ public record Node
 
 public record Forward
 (
-    int SrcPort,
+    long SrcPort,
-    int DstPort,
+    long DstPort,
     string DstIp
 );
 
@@ -112,13 +112,13 @@ public record Forward
 public record ProxyForward
 (
     [PartitionKey] Region Region,
-    int Port,
+    [RowKey] string Port,
     Guid ScalesetId,
     Guid MachineId,
     Guid? ProxyId,
-    [RowKey] int DstPort,
+    long DstPort,
     string DstIp,
-    DateTimeOffset EndTime
+    [property: JsonPropertyName("endtime")] DateTimeOffset EndTime
 ) : EntityBase();
 
 public record ProxyConfig

src/ApiService/ApiService/TestHooks/InstanceConfigTestHooks.cs

@@ -54,10 +54,10 @@ namespace ApiService.TestHooks {
             } else {
 
                 var query = UriExtension.GetQueryComponents(req.Url);
-                bool isNew = UriExtension.GetBoolValue("isNew", query, false);
+                bool isNew = UriExtension.GetBool("isNew", query, false);
                 //requireEtag wont' work since our current schema does not return etag to the client when getting data form the table, so
                 // there is no way to know which etag to use
-                bool requireEtag = UriExtension.GetBoolValue("requireEtag", query, false);
+                bool requireEtag = UriExtension.GetBool("requireEtag", query, false);
 
                 await _configOps.Save(newInstanceConfig, isNew, requireEtag);
 

src/ApiService/ApiService/TestHooks/IpOperationsTestHooks.cs

@@ -59,7 +59,7 @@ namespace ApiService.TestHooks {
 
             var rg = query["rg"];
             var name = query["name"];
-            var yes = UriExtension.GetBoolValue("force", query, false);
+            var yes = UriExtension.GetBool("force", query, false);
 
             if (yes) {
                 await _ipOps.DeleteNic(rg, name);
@@ -79,7 +79,7 @@ namespace ApiService.TestHooks {
 
             var rg = query["rg"];
             var name = query["name"];
-            var yes = UriExtension.GetBoolValue("force", query, false);
+            var yes = UriExtension.GetBool("force", query, false);
 
             if (yes) {
                 await _ipOps.DeleteIp(rg, name);

src/ApiService/ApiService/TestHooks/JobOperationsTestHooks.cs

@@ -53,10 +53,10 @@ namespace ApiService.TestHooks {
             _log.Info("Search jobs by state");
 
             var query = UriExtension.GetQueryComponents(req.Url);
-            var init = UriExtension.GetBoolValue("init", query, false);
+            var init = UriExtension.GetBool("init", query, false);
-            var enabled = UriExtension.GetBoolValue("enabled", query, false);
+            var enabled = UriExtension.GetBool("enabled", query, false);
-            var stopping = UriExtension.GetBoolValue("stopping", query, false);
+            var stopping = UriExtension.GetBool("stopping", query, false);
-            var stopped = UriExtension.GetBoolValue("stopped", query, false);
+            var stopped = UriExtension.GetBool("stopped", query, false);
 
             var states = new List<JobState>();
             if (init) {

src/ApiService/ApiService/TestHooks/NodeOperationsTestHooks.cs

@@ -128,7 +128,7 @@ namespace ApiService.TestHooks {
             _log.Info("to reimage");
 
             var query = UriExtension.GetQueryComponents(req.Url);
-            var done = UriExtension.GetBoolValue("done", query, false);
+            var done = UriExtension.GetBool("done", query, false);
 
             var s = await req.ReadAsStringAsync();
             var node = JsonSerializer.Deserialize<Node>(s!, EntityConverter.GetJsonSerializerOptions());
@@ -158,35 +158,18 @@ namespace ApiService.TestHooks {
             _log.Info("search states");
 
             var query = UriExtension.GetQueryComponents(req.Url);
-
+            Guid? poolId = UriExtension.GetGuid("poolId", query);
-            Guid? poolId = default;
+            Guid? scaleSetId = UriExtension.GetGuid("scaleSetId", query);
-            if (query.ContainsKey("poolId")) {
-                poolId = Guid.Parse(query["poolId"]);
-            }
-
-            Guid? scaleSetId = default;
-            if (query.ContainsKey("scaleSetId")) {
-                scaleSetId = Guid.Parse(query["scaleSetId"]);
-            }
 
             List<NodeState>? states = default;
             if (query.ContainsKey("states")) {
                 states = query["states"].Split('-').Select(s => Enum.Parse<NodeState>(s)).ToList();
             }
+            string? poolName = UriExtension.GetString("poolName", query);
 
-            string? poolName = default;
+            var excludeUpdateScheduled = UriExtension.GetBool("excludeUpdateScheduled", query, false);
-            if (query.ContainsKey("poolName")) {
+            int? numResults = UriExtension.GetInt("numResults", query);
-                poolName = query["poolName"];
-            }
-
-            var excludeUpdateScheduled = UriExtension.GetBoolValue("excludeUpdateScheduled", query, false);
-            int? numResults = default;
-            if (query.ContainsKey("numResults")) {
-                numResults = int.Parse(query["numResults"]);
-            }
             var r = _nodeOps.SearchStates(poolId, scaleSetId, states, poolName, excludeUpdateScheduled, numResults);
-
-
             var json = JsonSerializer.Serialize(await r.ToListAsync(), EntityConverter.GetJsonSerializerOptions());
             var resp = req.CreateResponse(HttpStatusCode.OK);
             await resp.WriteStringAsync(json);
@@ -236,7 +219,7 @@ namespace ApiService.TestHooks {
 
             string version = query["version"];
 
-            bool isNew = UriExtension.GetBoolValue("isNew", query, false);
+            bool isNew = UriExtension.GetBool("isNew", query, false);
 
             var node = await _nodeOps.Create(poolId, poolName, machineId, scaleSetId, version, isNew);
 

src/ApiService/ApiService/TestHooks/NotificationOperationsTestHooks.cs

@@ -28,7 +28,7 @@ namespace ApiService.TestHooks {
 
             var container = query["container"];
             var fileName = query["fileName"];
-            var failTaskOnTransientError = UriExtension.GetBoolValue("failTaskOnTransientError", query, true);
+            var failTaskOnTransientError = UriExtension.GetBool("failTaskOnTransientError", query, true);
 
             await _notificationOps.NewFiles(new Container(container), fileName, failTaskOnTransientError);
             var resp = req.CreateResponse(HttpStatusCode.OK);

src/ApiService/ApiService/TestHooks/NsgOperationsTestHooks.cs

@@ -15,7 +15,7 @@ namespace ApiService.TestHooks {
         private readonly INsgOperations _nsgOperations;
 
         public NsgOperationsTestHooks(ILogTracer log, IConfigOperations configOps, INsgOperations nsgOperations) {
-            _log = log.WithTag("TestHooks", nameof(NotificationOperationsTestHooks));
+            _log = log.WithTag("TestHooks", nameof(NsgOperationsTestHooks));
             _configOps = configOps; ;
             _nsgOperations = nsgOperations;
         }

src/ApiService/ApiService/TestHooks/PoolOperationsTestHooks.cs

@@ -21,7 +21,7 @@ namespace ApiService.TestHooks {
 
 
         [Function("GetPoolTestHook")]
-        public async Task<HttpResponseData> GetNsg([HttpTrigger(AuthorizationLevel.Anonymous, "get", Route = "testhooks/poolOperations/pool")] HttpRequestData req) {
+        public async Task<HttpResponseData> GetPool([HttpTrigger(AuthorizationLevel.Anonymous, "get", Route = "testhooks/poolOperations/pool")] HttpRequestData req) {
             _log.Info("get pool");
 
             var query = UriExtension.GetQueryComponents(req.Url);

src/ApiService/ApiService/TestHooks/ProxyForwardTestHooks.cs

@@ -0,0 +1,43 @@
+using System.Net;
+using System.Text.Json;
+using System.Threading.Tasks;
+using Microsoft.Azure.Functions.Worker;
+using Microsoft.Azure.Functions.Worker.Http;
+using Microsoft.OneFuzz.Service;
+using Microsoft.OneFuzz.Service.OneFuzzLib.Orm;
+
+
+#if DEBUG
+namespace ApiService.TestHooks {
+    public class ProxyForwardTestHooks {
+        private readonly ILogTracer _log;
+        private readonly IConfigOperations _configOps;
+        private readonly IProxyForwardOperations _proxyForward;
+
+        public ProxyForwardTestHooks(ILogTracer log, IConfigOperations configOps, IProxyForwardOperations proxyForward) {
+            _log = log.WithTag("TestHooks", nameof(ProxyForwardTestHooks));
+            _configOps = configOps; ;
+            _proxyForward = proxyForward; ;
+        }
+
+        [Function("SearchForwardTestHook")]
+        public async Task<HttpResponseData> SearchForward([HttpTrigger(AuthorizationLevel.Anonymous, "get", Route = "testhooks/proxyForwardOperations/search")] HttpRequestData req) {
+            _log.Info("search proxy forward");
+
+            var query = UriExtension.GetQueryComponents(req.Url);
+
+            var poolRes = _proxyForward.SearchForward(
+                UriExtension.GetGuid("scaleSetId", query),
+                UriExtension.GetString("region", query),
+                UriExtension.GetGuid("machineId", query),
+                UriExtension.GetGuid("proxyId", query),
+                UriExtension.GetInt("dstPort", query));
+
+            var json = JsonSerializer.Serialize(await poolRes.ToListAsync(), EntityConverter.GetJsonSerializerOptions());
+            var resp = req.CreateResponse(HttpStatusCode.OK);
+            await resp.WriteStringAsync(json);
+            return resp;
+        }
+    }
+}
+#endif

src/ApiService/ApiService/TestHooks/UriExtension.cs

@@ -13,12 +13,20 @@
             return new Dictionary<string, string>(q);
         }
 
-        public static bool GetBoolValue(string key, IDictionary<string, string> query, bool defaultValue = false) {
+        public static bool GetBool(string key, IDictionary<string, string> query, bool defaultValue = false) {
             bool v;
             if (query.ContainsKey(key)) {
-                if (!bool.TryParse(query[key], out v)) {
+                v = bool.Parse(query[key]);
-                    v = defaultValue;
+            } else {
-                }
+                v = defaultValue;
+            }
+            return v;
+        }
+
+        public static int? GetInt(string key, IDictionary<string, string> query, int? defaultValue = null) {
+            int? v;
+            if (query.ContainsKey(key)) {
+                v = int.Parse(query[key]);
             } else {
                 v = defaultValue;
             }
@@ -26,5 +34,22 @@
         }
 
 
+        public static string? GetString(string key, IDictionary<string, string> query, string? defaultValue = null) {
+            if (query.ContainsKey(key)) {
+                return query[key];
+            } else {
+                return defaultValue;
+            }
+        }
+
+        public static Guid? GetGuid(string key, IDictionary<string, string> query, Guid? defaultValue = null) {
+            if (query.ContainsKey(key)) {
+                return Guid.Parse(query[key]);
+            } else {
+                return defaultValue;
+            }
+        }
+
+
     }
 }

src/ApiService/ApiService/onefuzzlib/ProxyForwardOperations.cs

@@ -19,13 +19,13 @@ public class ProxyForwardOperations : Orm<ProxyForward>, IProxyForwardOperations
         var conditions =
             new[] {
                 scalesetId != null ? $"scaleset_id eq '{scalesetId}'" : null,
-                region != null ? $"region eq '{region}'" : null ,
+                region != null ? $"PartitionKey eq '{region}'" : null ,
                 machineId != null ? $"machine_id eq '{machineId}'" : null ,
                 proxyId != null ? $"proxy_id eq '{proxyId}'" : null ,
-                dstPort != null ? $"dsp_port eq {dstPort }" : null ,
+                dstPort != null ? $"dst_port eq {dstPort}" : null ,
             }.Where(x => x != null);
 
-        var filter = string.Join(" and ", conditions);
+        var filter = Query.And(conditions!);
 
         return QueryAsync(filter);
     }

src/ApiService/ApiService/onefuzzlib/ProxyOperations.cs

@@ -108,7 +108,6 @@ public class ProxyOperations : StatefulOrm<Proxy, VmState>, IProxyOperations {
     }
 
 
-
     public async Async.Task SetState(Proxy proxy, VmState state) {
         if (proxy.State == state) {
             return;
@@ -127,7 +126,7 @@ public class ProxyOperations : StatefulOrm<Proxy, VmState>, IProxyOperations {
             if (entry.EndTime < DateTimeOffset.UtcNow) {
                 await _context.ProxyForwardOperations.Delete(entry);
             } else {
-                forwards.Add(new Forward(entry.Port, entry.DstPort, entry.DstIp));
+                forwards.Add(new Forward(long.Parse(entry.Port), entry.DstPort, entry.DstIp));
             }
         }
         return forwards;

src/ApiService/Tests/OrmModelsTest.cs

@@ -91,11 +91,11 @@ namespace Tests {
         }
 
         public static Gen<ProxyForward> ProxyForward() {
-            return Arb.Generate<Tuple<Tuple<string, int, Guid, Guid, Guid?, int>, Tuple<IPv4Address, DateTimeOffset>>>().Select(
+            return Arb.Generate<Tuple<Tuple<string, long, Guid, Guid, Guid?, long>, Tuple<IPv4Address, DateTimeOffset>>>().Select(
                 arg =>
                     new ProxyForward(
                         Region: arg.Item1.Item1,
-                        Port: arg.Item1.Item2,
+                        Port: arg.Item1.Item2.ToString(),
                         ScalesetId: arg.Item1.Item3,
                         MachineId: arg.Item1.Item4,
                         ProxyId: arg.Item1.Item5,