fix(hosts): rate-limit heartbeats and close enrollment and audit gaps

apps/control_plane/lib/workharbor/config_trust.ex

@@ -25,46 +25,49 @@ defmodule WorkHarbor.ConfigTrust do
   def approve(%Principal{} = principal, descriptor) do
     with :ok <- Accounts.authorize(principal, :config_approve),
          {:ok, normalized, bytes, digest} <- normalize(descriptor) do
-      transaction(fn ->
-        authorize!(principal, :config_approve)
-
-        case Repo.get_by(Approval, [owner_id: principal.owner_id, descriptor_sha256: digest],
-               log: false
-             ) do
-          nil ->
-            approval =
+      result =
+        transaction(fn ->
+          authorize!(principal, :config_approve)
+
+          case Repo.get_by(Approval, [owner_id: principal.owner_id, descriptor_sha256: digest],
+                 log: false
+               ) do
+            nil ->
+              approval =
+                Repo.insert!(
+                  %Approval{
+                    owner_id: principal.owner_id,
+                    descriptor: normalized,
+                    descriptor_bytes: bytes,
+                    descriptor_sha256: digest,
+                    approved_at: Accounts.now()
+                  },
+                  log: false
+                )
+
               Repo.insert!(
-                %Approval{
+                %AuditEvent{
                   owner_id: principal.owner_id,
-                  descriptor: normalized,
-                  descriptor_bytes: bytes,
-                  descriptor_sha256: digest,
-                  approved_at: Accounts.now()
+                  event: "revision_approved",
+                  action: "config_approve",
+                  channel: channel(principal),
+                  outcome: "recorded",
+                  code: "ok",
+                  metadata: %{"approval_id" => approval.id}
                 },
                 log: false
               )
 
-            Repo.insert!(
-              %AuditEvent{
-                owner_id: principal.owner_id,
-                event: "revision_approved",
-                action: "config_approve",
-                channel: channel(principal),
-                outcome: "recorded",
-                code: "ok",
-                metadata: %{"approval_id" => approval.id}
-              },
-              log: false
-            )
-
-            {:ok, approval}
-
-          approval ->
-            if approval.descriptor_bytes == bytes and approval.descriptor == normalized,
-              do: {:ok, approval},
-              else: Repo.rollback(:trust_conflict)
-        end
-      end)
+              {:ok, approval}
+
+            approval ->
+              if approval.descriptor_bytes == bytes and approval.descriptor == normalized,
+                do: {:ok, approval},
+                else: Repo.rollback(:trust_conflict)
+          end
+        end)
+
+      with {:error, {:denied, action}} <- result, do: Accounts.deny(principal, action)
     end
   end
 
@@ -109,7 +112,7 @@ defmodule WorkHarbor.ConfigTrust do
   end
 
   defp authorize!(principal, action) do
-    if Accounts.authorize(principal, action) != :ok, do: Repo.rollback(:forbidden)
+    if Accounts.authorize(principal, action) != :ok, do: Repo.rollback({:denied, action})
   end
 
   defp channel(%Principal{kind: "owner_browser"}), do: "browser"

apps/control_plane/lib/workharbor/host_connections.ex

@@ -6,6 +6,11 @@ defmodule WorkHarbor.HostConnections do
   alias WorkHarbor.{Repo, Hosts, Accounts}
   alias WorkHarbor.HostConnections.Protocol
   alias WorkHarbor.Hosts.Peer
+  # Contract cadence is 10 s. Each heartbeat costs a transaction and an undeletable audit
+  # row, so arrivals are limited by a token bucket (capacity 3, one token per 5 s) measured
+  # on receive time, not on the previous commit; the first heartbeat after hello is free.
+  @heartbeat_refill_ms 5000
+  @heartbeat_capacity 3
   def start_link(opts), do: GenServer.start_link(__MODULE__, opts, name: __MODULE__)
   def open(peer, socket), do: call(__MODULE__, {:open, peer, socket})
   def receive_hello(ref, message), do: peer_call(ref, {:hello, message})
@@ -116,6 +121,8 @@ defmodule WorkHarbor.HostConnections do
             negotiated: false,
             opened: mono,
             seen: nil,
+            hb_arrived: nil,
+            hb_tokens: 0,
             lease: nil,
             expiry: expiry,
             expiry_timer: expiry_timer
@@ -213,44 +220,9 @@ defmodule WorkHarbor.HostConnections do
   end
 
   defp peer_message(state, ref, entry, {:heartbeat, message}) do
-    with true <- entry.negotiated,
-         {:ok, %{"kind" => "heartbeat"}} <-
-           Protocol.validate(message, entry.peer.owner_id, entry.peer.host_id) do
-      case Hosts.transaction(fn ->
-             current_slot!(entry)
-             if not Hosts.current_peer?(entry.peer), do: Repo.rollback(:unauthorized_host)
-             now = DateTime.from_unix!(Accounts.now()) |> DateTime.to_iso8601()
-             lease = DateTime.from_unix!(Accounts.now() + 30) |> DateTime.to_iso8601()
-
-             Hosts.sql("UPDATE hosts SET last_seen=?,lease_expires_at=? WHERE id=?", [
-               now,
-               lease,
-               entry.peer.host_id
-             ])
-
-             Hosts.sql(
-               "UPDATE host_connection_slots SET state='connected' WHERE host_id=? AND connection_id=?",
-               [entry.peer.host_id, entry.id]
-             )
-
-             Hosts.audit(
-               entry.peer.owner_id,
-               entry.peer.host_id,
-               "host_heartbeat",
-               "host_heartbeat",
-               "authenticated"
-             )
-           end) do
-        {:ok, _} ->
-          entry = %{entry | seen: monotonic(), lease: make_ref()}
-          Process.send_after(self(), {:stale, ref, entry.lease}, 30000)
-          {:reply, {:ok, ack(entry, "heartbeat_ack")}, put_entry(state, ref, entry)}
-
-        _ ->
-          {:reply, {:error, :persistence_failed}, close(state, ref, :persistence_failed)}
-      end
-    else
-      _ -> {:reply, {:error, :invalid_connection_message}, close(state, ref, :invalid_message)}
+    case take_heartbeat_token(entry, monotonic()) do
+      {:ok, entry} -> heartbeat(state, ref, entry, message)
+      :limited -> {:reply, {:error, :heartbeat_too_fast}, close(state, ref, :heartbeat_too_fast)}
     end
   end
 
@@ -429,6 +401,63 @@ defmodule WorkHarbor.HostConnections do
   defp fresh?(entry),
     do: entry.negotiated and not is_nil(entry.seen) and not stale?(entry) and valid_entry?(entry)
 
+  defp heartbeat(state, ref, entry, message) do
+    with true <- entry.negotiated,
+         {:ok, %{"kind" => "heartbeat"}} <-
+           Protocol.validate(message, entry.peer.owner_id, entry.peer.host_id) do
+      case Hosts.transaction(fn ->
+             current_slot!(entry)
+             if not Hosts.current_peer?(entry.peer), do: Repo.rollback(:unauthorized_host)
+             now = DateTime.from_unix!(Accounts.now()) |> DateTime.to_iso8601()
+             lease = DateTime.from_unix!(Accounts.now() + 30) |> DateTime.to_iso8601()
+
+             Hosts.sql("UPDATE hosts SET last_seen=?,lease_expires_at=? WHERE id=?", [
+               now,
+               lease,
+               entry.peer.host_id
+             ])
+
+             Hosts.sql(
+               "UPDATE host_connection_slots SET state='connected' WHERE host_id=? AND connection_id=?",
+               [entry.peer.host_id, entry.id]
+             )
+
+             Hosts.audit(
+               entry.peer.owner_id,
+               entry.peer.host_id,
+               "host_heartbeat",
+               "host_heartbeat",
+               "authenticated"
+             )
+           end) do
+        {:ok, _} ->
+          entry = %{entry | seen: monotonic(), lease: make_ref()}
+          Process.send_after(self(), {:stale, ref, entry.lease}, 30000)
+          {:reply, {:ok, ack(entry, "heartbeat_ack")}, put_entry(state, ref, entry)}
+
+        _ ->
+          {:reply, {:error, :persistence_failed}, close(state, ref, :persistence_failed)}
+      end
+    else
+      _ -> {:reply, {:error, :invalid_connection_message}, close(state, ref, :invalid_message)}
+    end
+  end
+
+  defp take_heartbeat_token(%{hb_arrived: nil} = entry, now),
+    do: {:ok, %{entry | hb_arrived: now, hb_tokens: @heartbeat_capacity * @heartbeat_refill_ms}}
+
+  defp take_heartbeat_token(entry, now) do
+    tokens =
+      min(
+        @heartbeat_capacity * @heartbeat_refill_ms,
+        entry.hb_tokens + max(0, now - entry.hb_arrived)
+      )
+
+    if tokens >= @heartbeat_refill_ms,
+      do: {:ok, %{entry | hb_arrived: now, hb_tokens: tokens - @heartbeat_refill_ms}},
+      else: :limited
+  end
+
   defp stale?(entry), do: monotonic() - (entry.seen || entry.opened) >= 30000
 
   defp monotonic,

apps/control_plane/lib/workharbor/hosts.ex

@@ -13,7 +13,7 @@ defmodule WorkHarbor.Hosts do
   alias WorkHarbor.Hosts.{Issuer, Peer}
 
   def issue_enrollment(principal, host_id) do
-    with :ok <- Accounts.authorize(principal, :host_enrollment_issue) do
+    with_authorization(principal, :host_enrollment_issue, fn ->
       transaction(fn ->
         authorize!(principal, :host_enrollment_issue)
 
@@ -39,6 +39,11 @@ defmodule WorkHarbor.Hosts do
         code = "whe1_" <> Base.url_encode64(:crypto.strong_rand_bytes(32), padding: false)
         now = Accounts.now()
 
+        sql(
+          "UPDATE host_enrollments SET revoked_at=? WHERE host_id=? AND consumed_at IS NULL AND revoked_at IS NULL",
+          [now, host.id]
+        )
+
         sql(
           "INSERT INTO host_enrollments(id,host_id,owner_id,code_hash,issued_at,expires_at) VALUES(?,?,?,?,?,?)",
           [Ecto.UUID.generate(), host.id, host.owner_id, digest(code), now, now + 300]
@@ -50,7 +55,7 @@ defmodule WorkHarbor.Hosts do
 
         %{host_id: host.id, owner_id: host.owner_id, code: code, expires_at: now + 300}
       end)
-    end
+    end)
   end
 
   def redeem(code, csr) do
@@ -236,7 +241,7 @@ defmodule WorkHarbor.Hosts do
   def activate(_), do: {:error, :unauthorized_host}
 
   def revoke(principal, host_id) do
-    with :ok <- Accounts.authorize(principal, :host_revoke) do
+    with_authorization(principal, :host_revoke, fn ->
       result =
         transaction(fn ->
           authorize!(principal, :host_revoke)
@@ -275,7 +280,7 @@ defmodule WorkHarbor.Hosts do
         do: Phoenix.PubSub.broadcast(WorkHarbor.PubSub, "host: " <> host_id, :identity_changed)
 
       normalize_ok(result)
-    end
+    end)
   end
 
   def status(principal, host_id) do
@@ -371,8 +376,19 @@ defmodule WorkHarbor.Hosts do
     Repo.get_by(Host, id: id, owner_id: owner) || Repo.rollback(:not_found)
   end
 
+  # The in-transaction check rolls back its own denial audit, so the denial is
+  # surfaced as a marker and audited here, after the rollback (authentication.md section 4).
   defp authorize!(p, action),
-    do: if(Accounts.authorize(p, action) != :ok, do: Repo.rollback(:forbidden))
+    do: if(Accounts.authorize(p, action) != :ok, do: Repo.rollback({:denied, action}))
+
+  defp with_authorization(principal, action, fun) do
+    with :ok <- Accounts.authorize(principal, action) do
+      case fun.() do
+        {:error, {:denied, ^action}} -> Accounts.deny(principal, action)
+        result -> result
+      end
+    end
+  end
 
   defp channel(%{kind: "cli_owner"}), do: "cli"
   defp channel(_), do: "browser"

apps/control_plane/test/workharbor/config_trust_test.exs

@@ -21,6 +21,18 @@ defmodule WorkHarbor.ConfigTrustTest do
     %{owner: owner, secret: secret, principal: principal}
   end
 
+  test "denial inside the transaction keeps its audit when authorization changes after the pre-check",
+       ctx do
+    expire_after_precheck(ctx.principal, :config_approve)
+    assert {:error, :forbidden} = ConfigTrust.approve(ctx.principal, DB.config())
+
+    assert WorkHarbor.Repo.query!(
+             "SELECT count(*) FROM audit_events WHERE event='authorization_denied' AND action='config_approve'",
+             [],
+             log: false
+           ).rows == [[1]]
+  end
+
   test "exact approval repeats without a second approval or audit", ctx do
     descriptor = DB.config()
     assert {:ok, first} = ConfigTrust.approve(ctx.principal, descriptor)
@@ -186,4 +198,31 @@ defmodule WorkHarbor.ConfigTrustTest do
     assert results |> Enum.map(fn {:ok, row} -> row.id end) |> Enum.uniq() |> length() == 1
     assert DB.count("revision_trust_approvals") == 1
   end
+
+  # Counts clock reads of one pre-check, then expires the credential right after it, so only
+  # the in-transaction check can deny.
+  defp expire_after_precheck(principal, action) do
+    {:ok, counter} = Agent.start_link(fn -> 0 end)
+    real = Application.get_env(:workharbor, :clock)
+    now = WorkHarbor.Accounts.now()
+
+    Application.put_env(:workharbor, :clock, fn ->
+      Agent.update(counter, &(&1 + 1))
+      now
+    end)
+
+    :ok = WorkHarbor.Accounts.authorize(principal, action)
+    reads = Agent.get(counter, & &1)
+    Agent.update(counter, fn _ -> 0 end)
+
+    Application.put_env(:workharbor, :clock, fn ->
+      if Agent.get_and_update(counter, &{&1 + 1, &1 + 1}) > reads, do: now + 10_000_000, else: now
+    end)
+
+    on_exit(fn ->
+      if real,
+        do: Application.put_env(:workharbor, :clock, real),
+        else: Application.delete_env(:workharbor, :clock)
+    end)
+  end
 end