Give entity providers the connection's security context

lib/espex/connection.ex

@@ -740,8 +740,22 @@ defmodule Espex.Connection do
   end
 
   defp interpret_action(_socket, state, {:entity_command, command}) do
-    state.adapters.entity_provider.handle_command(command)
-    |> log_adapter_error(state.peer, "entity command")
+    provider = state.adapters.entity_provider
+
+    # Prefer the context-aware callback so the provider can refuse
+    # privileged commands (reboot / factory reset / firmware install) on an
+    # unauthenticated connection. Espex can't make that call itself —
+    # which entity is dangerous is the provider's knowledge — so it just
+    # supplies the security context. Providers that don't export /2 keep
+    # the original behaviour.
+    result =
+      if function_exported?(provider, :handle_command, 2) do
+        provider.handle_command(command, command_context(state))
+      else
+        provider.handle_command(command)
+      end
+
+    log_adapter_error(result, state.peer, "entity command")
 
     {:cont, state}
   end
@@ -1044,6 +1058,11 @@ defmodule Espex.Connection do
     :ok
   end
 
+  # A Noise session is the only authentication the ESPHome protocol
+  # actually provides here: AuthenticationRequest is answered
+  # unconditionally, so a keyless connection is anonymous by construction.
+  defp command_context(state), do: %{encrypted?: match?({:active, _, _}, state.encryption)}
+
   defp log_adapter_error(:ok, _peer, _what), do: :ok
 
   defp log_adapter_error({:error, reason}, peer, what) do

lib/espex/entity_provider.ex

@@ -16,6 +16,7 @@ defmodule Espex.EntityProvider do
   | `c:list_entities/0` | Once per connection, when the client sends `ListEntitiesRequest` |
   | `c:initial_states/0` | Once per connection, when the client sends `SubscribeStatesRequest` |
   | `c:handle_command/1` | Each time the client issues a command struct for one of your entities |
+  | `c:handle_command/2` | Optional. Same, but also given the connection's security context — preferred over `c:handle_command/1` when exported |
 
   Return values:
 
@@ -26,6 +27,10 @@ defmodule Espex.EntityProvider do
     * `c:handle_command/1` returns `:ok` on success or `{:error, term}`
       on failure. Espex currently logs errors and continues — it does
       not send an error back to the client.
+    * `c:handle_command/2` has the same return contract. Implement it to
+      refuse dangerous commands (reboot, factory reset, firmware install)
+      on an unauthenticated connection — see its docs for why that
+      judgement has to live in the provider rather than in Espex.
 
   ## Frozen-at-accept-time snapshot
 
@@ -223,4 +228,47 @@ defmodule Espex.EntityProvider do
   `%Espex.Proto.SwitchCommandRequest{}`, `%Espex.Proto.LightCommandRequest{}`.
   """
   @callback handle_command(command :: struct()) :: :ok | {:error, term()}
+
+  @doc """
+  Same as `c:handle_command/1`, but also given the originating
+  connection's security context. Preferred when exported — Espex calls
+  this in place of `c:handle_command/1`.
+
+  The context is currently:
+
+      %{encrypted?: boolean()}
+
+  `encrypted?` is true only once a Noise session is established. It is
+  false whenever the server is running keyless, because the ESPHome
+  protocol offers no other authentication: `AuthenticationRequest` carries
+  a password field, but Espex answers every such request with
+  `invalid_password: false`. On a keyless server, *any* host that can open
+  a TCP connection can therefore issue entity commands.
+
+  That is usually acceptable — the keyless window exists so Home Assistant
+  can adopt the device and provision a PSK, and most entity commands are
+  innocuous. It is not acceptable for commands that reboot, wipe or
+  reflash the device. Espex cannot tell those apart: which entity is
+  dangerous is the provider's knowledge, not the protocol's. Implement
+  this callback to make that judgement:
+
+      @impl Espex.EntityProvider
+      def handle_command(command, %{encrypted?: false}) do
+        if privileged?(command) do
+          Logger.warning("refusing privileged command on an unencrypted connection")
+          :ok
+        else
+          handle_command(command)
+        end
+      end
+
+      def handle_command(command, _context), do: handle_command(command)
+
+  Providers that don't export this keep receiving `c:handle_command/1`
+  unchanged.
+  """
+  @callback handle_command(command :: struct(), context :: %{encrypted?: boolean()}) ::
+              :ok | {:error, term()}
+
+  @optional_callbacks handle_command: 2
 end

test/espex/encrypted_integration_test.exs

@@ -13,17 +13,19 @@ defmodule Espex.EncryptedIntegrationTest do
 
     {:ok, sup_pid} =
       Espex.start_link(
-        name: sup_name,
-        server_name: server_name,
-        port: 0,
-        device_config: [
-          name: "espex-encrypted",
-          friendly_name: "Espex Encrypted",
-          project_name: "espex.demo",
-          project_version: "0.0.1",
-          mac_address: "AA:BB:CC:DD:EE:FF",
-          psk: @psk
-        ]
+        [
+          name: sup_name,
+          server_name: server_name,
+          port: 0,
+          device_config: [
+            name: "espex-encrypted",
+            friendly_name: "Espex Encrypted",
+            project_name: "espex.demo",
+            project_version: "0.0.1",
+            mac_address: "AA:BB:CC:DD:EE:FF",
+            psk: @psk
+          ]
+        ] ++ adapter_opts(context)
       )
 
     {:ok, port} = Espex.Supervisor.bound_port(sup_pid)
@@ -41,6 +43,37 @@ defmodule Espex.EncryptedIntegrationTest do
     %{port: port}
   end
 
+  describe "entity commands (security context)" do
+    setup do
+      :persistent_term.put(:espex_entity_command_test_pid, self())
+      on_exit(fn -> :persistent_term.erase(:espex_entity_command_test_pid) end)
+      :ok
+    end
+
+    @tag adapters: %{entity_provider: Espex.Test.ContextAwareEntityProvider}
+    test "an established Noise session reports encrypted?: true", %{port: port} do
+      sock = connect(port)
+      {tx, _rx} = do_handshake(sock)
+
+      {:ok, type, payload} = MessageTypes.encode_parts(%Proto.ButtonCommandRequest{key: 1})
+      inner = NoiseFrame.encode_inner(type, payload)
+      {:ok, _tx, ct} = Noise.encrypt(tx, <<>>, inner)
+      :ok = :gen_tcp.send(sock, NoiseFrame.encode_outer(ct))
+
+      assert_receive {:entity_command_2, %Proto.ButtonCommandRequest{key: 1}, %{encrypted?: true}},
+                     1_000
+
+      :gen_tcp.close(sock)
+    end
+  end
+
+  defp adapter_opts(context) do
+    case Map.get(context, :adapters) do
+      nil -> []
+      adapters -> Map.to_list(adapters)
+    end
+  end
+
   defp connect(port) do
     {:ok, sock} =
       :gen_tcp.connect(~c"127.0.0.1", port, [:binary, active: false, nodelay: true, packet: :raw])

test/espex/integration_test.exs

@@ -119,6 +119,38 @@ defmodule Espex.IntegrationTest do
     end
   end
 
+  describe "entity commands (security context)" do
+    setup do
+      :persistent_term.put(:espex_entity_command_test_pid, self())
+      on_exit(fn -> :persistent_term.erase(:espex_entity_command_test_pid) end)
+      :ok
+    end
+
+    @tag adapters: %{entity_provider: Espex.Test.FakeEntityProvider}
+    test "a provider without handle_command/2 still gets handle_command/1", %{port: port} do
+      socket = connect(port)
+      send_struct(socket, %Proto.ButtonCommandRequest{key: 1})
+
+      assert_receive {:entity_command_1, %Proto.ButtonCommandRequest{key: 1}}, 1_000
+      :gen_tcp.close(socket)
+    end
+
+    @tag adapters: %{entity_provider: Espex.Test.ContextAwareEntityProvider}
+    test "handle_command/2 is preferred and reports an unencrypted connection", %{port: port} do
+      socket = connect(port)
+      send_struct(socket, %Proto.ButtonCommandRequest{key: 1})
+
+      # Plaintext connection: the provider must be told, so it can refuse
+      # privileged commands. AuthenticationRequest is answered
+      # unconditionally, so there is no other signal available.
+      assert_receive {:entity_command_2, %Proto.ButtonCommandRequest{key: 1}, %{encrypted?: false}},
+                     1_000
+
+      refute_receive {:entity_command_1, _}, 100
+      :gen_tcp.close(socket)
+    end
+  end
+
   describe "disconnect" do
     test "DisconnectRequest gets a response and the server closes the socket", %{port: port} do
       socket = connect(port)

test/support/fake_adapters.ex

@@ -252,7 +252,53 @@ defmodule Espex.Test.FakeEntityProvider do
   end
 
   @impl true
-  def handle_command(_message), do: :ok
+  def handle_command(message) do
+    Espex.Test.EntityCommandProbe.notify({:entity_command_1, message})
+    :ok
+  end
+end
+
+defmodule Espex.Test.ContextAwareEntityProvider do
+  @moduledoc """
+  Entity provider exporting the optional `handle_command/2`, so a test can
+  assert Espex prefers it and hands over the connection's security context.
+  """
+  @behaviour Espex.EntityProvider
+
+  @impl true
+  def list_entities do
+    [%Espex.Proto.ListEntitiesBinarySensorResponse{object_id: "fake", key: 1, name: "Fake"}]
+  end
+
+  @impl true
+  def initial_states do
+    [%Espex.Proto.BinarySensorStateResponse{key: 1, state: true, missing_state: false}]
+  end
+
+  @impl true
+  def handle_command(message) do
+    Espex.Test.EntityCommandProbe.notify({:entity_command_1, message})
+    :ok
+  end
+
+  @impl true
+  def handle_command(message, context) do
+    Espex.Test.EntityCommandProbe.notify({:entity_command_2, message, context})
+    :ok
+  end
+end
+
+defmodule Espex.Test.EntityCommandProbe do
+  @moduledoc """
+  Routes entity-command callbacks to the test process registered under
+  `:persistent_term` key `:espex_entity_command_test_pid`.
+  """
+  def notify(msg) do
+    case :persistent_term.get(:espex_entity_command_test_pid, nil) do
+      pid when is_pid(pid) -> send(pid, msg)
+      _ -> :ok
+    end
+  end
 end
 
 defmodule Espex.Test.PidPskStore do