Validate Context ABI ref schemes

Author: nshkrdotcom

core/context_abi/lib/outer_brain/context_abi/context_packet.ex

@@ -44,17 +44,31 @@ defmodule OuterBrain.ContextABI.ContextPacket do
 
   def new(attrs) when is_map(attrs) do
     with :ok <- Validator.reject_raw_payload(attrs),
-         {:ok, tenant_ref} <- Validator.required_string(attrs, :tenant_ref),
-         {:ok, user_request_ref} <- Validator.required_string(attrs, :user_request_ref),
+         {:ok, tenant_ref} <- Validator.required_ref(attrs, :tenant_ref, ["tenant://"]),
+         {:ok, user_request_ref} <-
+           Validator.required_ref(attrs, :user_request_ref, [
+             "artifact://",
+             "user-request://",
+             "request://"
+           ]),
          {:ok, system_instruction_ref} <-
-           Validator.required_string(attrs, :system_instruction_ref),
-         {:ok, memory_refs} <- Validator.string_list(attrs, :memory_refs),
+           Validator.required_ref(attrs, :system_instruction_ref, [
+             "artifact://",
+             "system-instruction://"
+           ]),
+         {:ok, memory_refs} <-
+           Validator.string_ref_list(attrs, :memory_refs, ["memory://", "memory-candidate://"]),
          :ok <- reject_unpromoted_memory_refs(memory_refs),
-         {:ok, budget_ref} <- Validator.required_string(attrs, :budget_ref),
-         {:ok, model_class_allowlist} <- Validator.string_list(attrs, :model_class_allowlist),
+         {:ok, budget_ref} <- Validator.required_ref(attrs, :budget_ref, ["budget://"]),
+         {:ok, model_class_allowlist} <-
+           Validator.string_ref_list(attrs, :model_class_allowlist, [
+             "model-class://",
+             "class://"
+           ]),
          :ok <- require_nonempty(model_class_allowlist, :model_class_allowlist),
-         {:ok, route_policy_ref} <- Validator.required_string(attrs, :route_policy_ref),
-         {:ok, trace_ref} <- Validator.required_string(attrs, :trace_ref),
+         {:ok, route_policy_ref} <-
+           Validator.required_ref(attrs, :route_policy_ref, ["route-policy://"]),
+         {:ok, trace_ref} <- Validator.required_ref(attrs, :trace_ref, ["trace://"]),
          {:ok, extension_refs} <- Validator.optional_map(attrs, :extension_refs) do
       hash_input = %{
         schema_ref: Validator.fetch(attrs, :schema_ref, @schema_ref),

core/context_abi/lib/outer_brain/context_abi/validator.ex

@@ -63,6 +63,14 @@ defmodule OuterBrain.ContextABI.Validator do
     end
   end
 
+  @spec required_ref(map(), atom(), [String.t()]) :: {:ok, String.t()} | {:error, Failure.t()}
+  def required_ref(attrs, field, allowed_prefixes) do
+    with {:ok, value} <- required_string(attrs, field),
+         :ok <- validate_ref_scheme(value, field, allowed_prefixes) do
+      {:ok, value}
+    end
+  end
+
   @spec string_list(map(), atom()) :: {:ok, [String.t()]} | {:error, Failure.t()}
   def string_list(attrs, field) do
     case fetch(attrs, field, []) do
@@ -78,6 +86,24 @@ defmodule OuterBrain.ContextABI.Validator do
     end
   end
 
+  @spec string_ref_list(map(), atom(), [String.t()]) ::
+          {:ok, [String.t()]} | {:error, Failure.t()}
+  def string_ref_list(attrs, field, allowed_prefixes) do
+    with {:ok, values} <- string_list(attrs, field) do
+      values
+      |> Enum.reduce_while(:ok, fn value, :ok ->
+        case validate_ref_scheme(value, field, allowed_prefixes) do
+          :ok -> {:cont, :ok}
+          error -> {:halt, error}
+        end
+      end)
+      |> case do
+        :ok -> {:ok, values}
+        error -> error
+      end
+    end
+  end
+
   @spec optional_map(map(), atom()) :: {:ok, map()} | {:error, Failure.t()}
   def optional_map(attrs, field) do
     case fetch(attrs, field, %{}) do
@@ -139,6 +165,17 @@ defmodule OuterBrain.ContextABI.Validator do
     failure(:outer_brain, reason_code, safe_message: "context vocabulary value is invalid")
   end
 
+  defp validate_ref_scheme(value, field, allowed_prefixes) do
+    if Enum.any?(allowed_prefixes, &String.starts_with?(value, &1)) do
+      :ok
+    else
+      failure(:outer_brain, "outer_brain.context.invalid_ref_scheme.v1",
+        safe_message: "context ref uses an invalid scheme",
+        evidence_refs: ["field://#{Atom.to_string(field)}", "ref://#{value}"]
+      )
+    end
+  end
+
   defp find_raw_key(%{__struct__: _} = value), do: value |> Map.from_struct() |> find_raw_key()
 
   defp find_raw_key(%{} = map) do

core/context_abi/test/outer_brain/context_abi_test.exs

@@ -61,6 +61,16 @@ defmodule OuterBrain.ContextABITest do
     assert nested_failure.reason_code == "outer_brain.context.raw_payload_rejected.v1"
   end
 
+  test "compiler rejects invalid context ref schemes" do
+    assert {:error, %Failure{} = failure} =
+             compile_request()
+             |> Map.put(:tenant_ref, "not-a-tenant-ref")
+             |> ContextABI.compile()
+
+    assert failure.reason_code == "outer_brain.context.invalid_ref_scheme.v1"
+    assert "field://tenant_ref" in failure.evidence_refs
+  end
+
   test "compiler rejects unpromoted memory candidate refs in production packets" do
     assert {:error, %Failure{} = failure} =
              compile_request()