Fleet agent update improvements

internal/config/case.go

@@ -1227,6 +1227,30 @@ type FleetConfig struct {
 	// pushing the AFTER config. It proves the case really starts suppressed
 	// (otherwise "delivery after the update" is vacuous). Default 20.
 	BaselineSeconds int `yaml:"baseline_seconds"`
+
+	// DownloadGate (agent_download_gate scenario) lists the per-device probes the
+	// driver fires at the director's /dl/agent/<os>/<arch> endpoint: each GET
+	// carries X-VM-Device-Id and must return ExpectCode (200 = served, 403 =
+	// blocked). This exercises the director's authoritative update-download gate
+	// against devices whose update policy is set in the subject's config.
+	DownloadGate []DownloadGateProbe `yaml:"download_gate"`
+	// DownloadOSArch is the "<os>/<arch>" path segment probed (default "linux/amd64").
+	DownloadOSArch string `yaml:"download_os_arch"`
+	// ThrottleCount (agent_download_gate) fires this many CONCURRENT downloads for
+	// an allowed device; ThrottleMin429 asserts at least this many come back 429
+	// (the director's download throttle engaged). 0 skips the throttle check.
+	ThrottleCount  int `yaml:"throttle_count"`
+	ThrottleMin429 int `yaml:"throttle_min_429"`
+	// ThrottleDeviceID is the (allowed) device id used for the throttle burst
+	// (default "0" would be blocked; set it to an immediate/authorized device).
+	ThrottleDeviceID string `yaml:"throttle_device_id"`
+}
+
+// DownloadGateProbe is one agent_download_gate assertion: GET /dl as DeviceID,
+// expect HTTP ExpectCode.
+type DownloadGateProbe struct {
+	DeviceID   string `yaml:"device_id"`
+	ExpectCode int    `yaml:"expect_code"`
 }
 
 // BaselineSecondsOrDefault returns BaselineSeconds or 20 when unset.
@@ -1345,6 +1369,27 @@ func (tc *TestCase) validateFleet() error {
 		if !tc.UsesVerifier() {
 			return fmt.Errorf("case %q: fleet.scenario pipeline_verify_change requires a `verifier:` block (the object-store verdict)", tc.Name)
 		}
+	case "agent_download_gate":
+		// The director serves /dl for agent binary downloads; the driver probes it
+		// per-device and asserts the authorization gate (200/403) + throttle (429).
+		// The subject config carries the device update policies directly (plain
+		// YAML with an environments block — no deliver_config), so nothing extra is
+		// required here beyond the download_gate probe list.
+		if len(tc.Fleet.DownloadGate) == 0 && tc.Fleet.ThrottleCount == 0 {
+			return fmt.Errorf("case %q: fleet.scenario agent_download_gate requires fleet.download_gate probes and/or fleet.throttle_count", tc.Name)
+		}
+		// When the throttle burst is enabled it must actually be able to engage the
+		// throttle: require a positive 429 threshold (otherwise the 429 assertion is
+		// vacuous) and an explicit allowed device (the default "0" is blocked and
+		// would 403 before ever acquiring a download slot).
+		if tc.Fleet.ThrottleCount > 0 {
+			if tc.Fleet.ThrottleMin429 <= 0 {
+				return fmt.Errorf("case %q: fleet.scenario agent_download_gate with throttle_count > 0 requires throttle_min_429 > 0 (else the 429 assertion is vacuous)", tc.Name)
+			}
+			if tc.Fleet.ThrottleDeviceID == "" {
+				return fmt.Errorf("case %q: fleet.scenario agent_download_gate with throttle_count > 0 requires throttle_device_id (an allowed device; the default \"0\" is blocked and never reaches the throttle)", tc.Name)
+			}
+		}
 	case "connect", "config_update", "remote_check", "live_data", "console_log", "stats", "reconnect", "bad_token", "self_managed", "enrollment":
 	default:
 		return fmt.Errorf("case %q: unknown fleet.scenario %q", tc.Name, tc.Fleet.Scenario)

internal/runner/runner.go

@@ -3648,6 +3648,29 @@ func fleetSimSend(simContainer, dirID, command string, params map[string]any) (s
 	return string(out), nil
 }
 
+// fleetSimDLProbe asks the simulator to fire GET probes at the director's /dl
+// endpoint (params: url, device_id, token, count, concurrency) and returns the
+// observed HTTP status histogram (status-code string → count). Status "0" means
+// the request never got an HTTP response (director not up / connection error).
+func fleetSimDLProbe(simContainer string, params map[string]any) (map[string]int, error) {
+	body, _ := json.Marshal(params)
+	ctx, cancel := context.WithTimeout(context.Background(), 120*time.Second)
+	defer cancel()
+	out, err := exec.CommandContext(ctx, "docker", "exec", simContainer,
+		"wget", "-q", "-T", "100", "-O", "-", "--post-data="+string(body),
+		"http://127.0.0.1:8090/dlprobe").CombinedOutput()
+	if err != nil {
+		return nil, fmt.Errorf("dlprobe: %w (out: %.200s)", err, string(out))
+	}
+	var resp struct {
+		Histogram map[string]int `json:"histogram"`
+	}
+	if err := json.Unmarshal(out, &resp); err != nil {
+		return nil, fmt.Errorf("decode dlprobe: %w (raw: %.200s)", err, string(out))
+	}
+	return resp.Histogram, nil
+}
+
 // fleetWaitConnected polls until the director shows connected at the simulator.
 func fleetWaitConnected(simContainer, dirID string, deadline time.Time) bool {
 	for time.Now().Before(deadline) {
@@ -3899,17 +3922,23 @@ func (r *Runner) runFleetAutomationCorrectness(tc *config.TestCase, subject conf
 	}
 
 	// All other scenarios: wait for the director to connect first.
-	connDeadline := time.Now().Add(warmup + 90*time.Second)
-	if connDeadline.After(runDeadline) {
-		connDeadline = runDeadline
-	}
-	fmt.Printf("  waiting for the director to connect to the fleet simulator (up to %s)…\n", time.Until(connDeadline).Round(time.Second))
-	if !fleetWaitConnected(simContainer, dirID, connDeadline) {
-		errs = append(errs, "director never connected to the fleet simulator")
-		passed = false
-		return r.saveFleetResult(tc, subject, configName, startTime, finalCount, passed, errs, simContainer, subjectContainer)
+	// EXCEPTION: agent_download_gate does not use the fleet link at all — the
+	// director serves /dl straight from its own mounted config (which carries an
+	// environments block) and the simulator only probes /dl. Skip the connect
+	// gate (and the deliver_config step below is already a no-op without one).
+	if fc.Scenario != "agent_download_gate" {
+		connDeadline := time.Now().Add(warmup + 90*time.Second)
+		if connDeadline.After(runDeadline) {
+			connDeadline = runDeadline
+		}
+		fmt.Printf("  waiting for the director to connect to the fleet simulator (up to %s)…\n", time.Until(connDeadline).Round(time.Second))
+		if !fleetWaitConnected(simContainer, dirID, connDeadline) {
+			errs = append(errs, "director never connected to the fleet simulator")
+			passed = false
+			return r.saveFleetResult(tc, subject, configName, startTime, finalCount, passed, errs, simContainer, subjectContainer)
+		}
+		fmt.Println("  director connected to the fleet simulator ✓")
 	}
-	fmt.Println("  director connected to the fleet simulator ✓")
 
 	// Deliver the audited operational config (the platform's job). A fleet
 	// director loads its operational config (environments/devices/targets/routes/
@@ -4050,6 +4079,92 @@ func (r *Runner) runFleetAutomationCorrectness(tc *config.TestCase, subject conf
 			fmt.Println("  initial config request seen ✓")
 		}
 
+	case "agent_download_gate":
+		// The director serves /dl over plain HTTP on its discovery port (8080).
+		// It authorizes each agent binary download from the device's update policy
+		// in the subject config (immediate/scheduled/authorized → 200; user-action
+		// without authorization → 403), and throttles concurrent downloads (429).
+		// fleetsim probes /dl by the subject's network hostname.
+		osArch := fleetStr(fc.DownloadOSArch, "linux/amd64")
+		dlURL := "http://subject:8080/dl/agent/" + osArch
+		probe := func(devID string, count, conc int) (map[string]int, error) {
+			return fleetSimDLProbe(simContainer, map[string]any{
+				"url": dlURL, "device_id": devID, "count": count, "concurrency": conc,
+			})
+		}
+
+		// Readiness: the director needs its environment config + listeners up
+		// before /dl serves. An unmanaged probe returns a real HTTP status (403),
+		// while a connection error is bucketed as "0".
+		fmt.Println("  waiting for the director's /dl endpoint to serve…")
+		ready := false
+		// Snapshot the deadline once so the wait expires at a fixed settle-based
+		// point; recomputing scenarioDeadline() each iteration would slide it
+		// forward (to runDeadline) instead of failing fast like the other waits.
+		readyDeadline := scenarioDeadline()
+		for time.Now().Before(readyDeadline) {
+			if h, e := probe("0", 1, 1); e == nil && len(h) > 0 && h["0"] == 0 {
+				ready = true
+				break
+			}
+			time.Sleep(3 * time.Second)
+		}
+		if !ready {
+			errs = append(errs, "director /dl endpoint never returned an HTTP response")
+			break
+		}
+		fmt.Println("  /dl endpoint is serving ✓")
+
+		// Per-device gate assertions.
+		for _, p := range fc.DownloadGate {
+			h, e := probe(p.DeviceID, 1, 1)
+			if e != nil {
+				errs = append(errs, fmt.Sprintf("dl probe device=%s failed: %v", p.DeviceID, e))
+				continue
+			}
+			// probe(_, 1, 1) fires exactly one request, so the histogram must have
+			// exactly one nonzero bucket — the observed status. Require that
+			// explicitly rather than depending on map-iteration order.
+			got, nonzero := 0, 0
+			for codeStr, n := range h {
+				if n > 0 {
+					got, _ = strconv.Atoi(codeStr)
+					nonzero++
+				}
+			}
+			if nonzero != 1 {
+				errs = append(errs, fmt.Sprintf("device %s: /dl probe returned %d nonzero status buckets, want exactly 1 (hist %v)", p.DeviceID, nonzero, h))
+				continue
+			}
+			if got != p.ExpectCode {
+				errs = append(errs, fmt.Sprintf("device %s: /dl returned %d, want %d", p.DeviceID, got, p.ExpectCode))
+			} else {
+				fmt.Printf("  device %s -> %d ✓\n", p.DeviceID, got)
+			}
+		}
+
+		// Throttle burst: fire N concurrent downloads for an ALLOWED device (a
+		// blocked device 403s before acquiring a slot, so it can't exercise the
+		// throttle) and assert at least ThrottleMin429 come back 429.
+		if fc.ThrottleCount > 0 {
+			// validateFleet guarantees ThrottleDeviceID is set (an allowed device)
+			// whenever ThrottleCount > 0, so there is no blocked-"0" fallback here —
+			// that would 403 before acquiring a slot and never reach the 429 check.
+			tdev := fc.ThrottleDeviceID
+			h, e := probe(tdev, fc.ThrottleCount, fc.ThrottleCount)
+			if e != nil {
+				errs = append(errs, "throttle probe failed: "+e.Error())
+			} else {
+				got429 := h["429"]
+				finalCount = int64(got429)
+				if got429 < fc.ThrottleMin429 {
+					errs = append(errs, fmt.Sprintf("throttle: got %d×429 across %d concurrent, want >= %d (hist %v)", got429, fc.ThrottleCount, fc.ThrottleMin429, h))
+				} else {
+					fmt.Printf("  throttle: %d/%d concurrent downloads got 429 ✓\n", got429, fc.ThrottleCount)
+				}
+			}
+		}
+
 	case "remote_check":
 		expect := fc.ExpectRemoteResultOrDefault()
 		params := map[string]any{