atproto / test / identity / test_identity.ml

(** Identity tests for AT Protocol.

    Tests the DID resolver module with mock HTTP responses. *)

open Atproto_identity

(** {1 Mock HTTP Handler} *)

(** Global mock handler for HTTP GET *)
let mock_http_handler : (Uri.t -> Did_resolver.http_response) ref =
  ref (fun _ -> Did_resolver.{ status = 500; body = "No mock configured" })

(** Effect handler *)
let http_effect_handler : type a.
    a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option = function
  | Did_resolver.Http_get uri ->
      Some (fun k -> Effect.Deep.continue k (!mock_http_handler uri))
  | _ -> None

(** Run with mock HTTP *)
let run_with_mock_http ~handler f =
  mock_http_handler := handler;
  Effect.Deep.match_with f ()
    { retc = (fun x -> x); exnc = raise; effc = http_effect_handler }

(** {1 Sample DID Documents} *)

let sample_plc_doc =
  {|{
  "id": "did:plc:ewvi7nxzy7mbhbzdkr36ha",
  "alsoKnownAs": ["at://jay.bsky.social"],
  "verificationMethod": [
    {
      "id": "did:plc:ewvi7nxzy7mbhbzdkr36ha#atproto",
      "type": "Multikey",
      "controller": "did:plc:ewvi7nxzy7mbhbzdkr36ha",
      "publicKeyMultibase": "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9pribSF"
    }
  ],
  "service": [
    {
      "id": "#atproto_pds",
      "type": "AtprotoPersonalDataServer",
      "serviceEndpoint": "https://bsky.social"
    }
  ]
}|}

let sample_web_doc =
  {|{
  "id": "did:web:example.com",
  "alsoKnownAs": ["at://example.com"],
  "verificationMethod": [
    {
      "id": "did:web:example.com#atproto",
      "type": "Multikey",
      "controller": "did:web:example.com",
      "publicKeyMultibase": "zQ3shXjHeiBuRCKmM36cuYnm7YEMzhGnCmCyW92sRJ9pribSF"
    }
  ],
  "service": [
    {
      "id": "#atproto_pds",
      "type": "AtprotoPersonalDataServer",
      "serviceEndpoint": "https://pds.example.com"
    }
  ]
}|}

(** {1 Tests} *)

let test_resolve_plc () =
  let handler uri =
    let path = Uri.path uri in
    if path = "/did:plc:ewvi7nxzy7mbhbzdkr36ha" then
      Did_resolver.{ status = 200; body = sample_plc_doc }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:plc:ewvi7nxzy7mbhbzdkr36ha" with
      | Ok doc ->
          Alcotest.(check string) "id" "did:plc:ewvi7nxzy7mbhbzdkr36ha" doc.id;
          Alcotest.(check bool)
            "has alsoKnownAs" true
            (List.length doc.also_known_as > 0);
          Alcotest.(check bool)
            "has verification methods" true
            (List.length doc.verification_method > 0);
          Alcotest.(check bool) "has services" true (List.length doc.service > 0)
      | Error e -> Alcotest.fail (Did_resolver.error_to_string e))

let test_resolve_web () =
  let handler uri =
    let host = Uri.host uri |> Option.value ~default:"" in
    let path = Uri.path uri in
    if host = "example.com" && path = "/.well-known/did.json" then
      Did_resolver.{ status = 200; body = sample_web_doc }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:web:example.com" with
      | Ok doc ->
          Alcotest.(check string) "id" "did:web:example.com" doc.id;
          Alcotest.(check bool)
            "has alsoKnownAs" true
            (List.length doc.also_known_as > 0)
      | Error e -> Alcotest.fail (Did_resolver.error_to_string e))

let test_get_handle () =
  let handler _uri = Did_resolver.{ status = 200; body = sample_plc_doc } in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:plc:ewvi7nxzy7mbhbzdkr36ha" with
      | Ok doc -> (
          match Did_resolver.get_handle doc with
          | Some handle ->
              Alcotest.(check string)
                "handle" "jay.bsky.social"
                (Atproto_syntax.Handle.to_string handle)
          | None -> Alcotest.fail "expected handle")
      | Error e -> Alcotest.fail (Did_resolver.error_to_string e))

let test_get_pds_endpoint () =
  let handler _uri = Did_resolver.{ status = 200; body = sample_plc_doc } in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:plc:ewvi7nxzy7mbhbzdkr36ha" with
      | Ok doc -> (
          match Did_resolver.get_pds_endpoint doc with
          | Some pds ->
              Alcotest.(check string)
                "pds" "https://bsky.social" (Uri.to_string pds)
          | None -> Alcotest.fail "expected PDS endpoint")
      | Error e -> Alcotest.fail (Did_resolver.error_to_string e))

let test_get_signing_key () =
  let handler _uri = Did_resolver.{ status = 200; body = sample_plc_doc } in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:plc:ewvi7nxzy7mbhbzdkr36ha" with
      | Ok doc -> (
          match Did_resolver.get_signing_key doc with
          | Some key ->
              Alcotest.(check bool)
                "key starts with z" true
                (String.length key > 0 && key.[0] = 'z')
          | None -> Alcotest.fail "expected signing key")
      | Error e -> Alcotest.fail (Did_resolver.error_to_string e))

let test_not_found () =
  let handler _uri = Did_resolver.{ status = 404; body = "Not found" } in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:plc:notfound" with
      | Error Did_resolver.Not_found -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Not_found, got %s"
               (Did_resolver.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_http_error () =
  let handler _uri =
    Did_resolver.{ status = 500; body = "Internal Server Error" }
  in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:plc:test" with
      | Error (Did_resolver.Http_error (500, _)) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Http_error 500, got %s"
               (Did_resolver.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_invalid_did () =
  let handler _uri = Did_resolver.{ status = 200; body = sample_plc_doc } in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "invalid" with
      | Error (Did_resolver.Invalid_did _) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Invalid_did, got %s"
               (Did_resolver.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_unsupported_method () =
  let handler _uri = Did_resolver.{ status = 200; body = sample_plc_doc } in
  run_with_mock_http ~handler (fun () ->
      match Did_resolver.resolve "did:key:z123" with
      | Error (Did_resolver.Unsupported_method _) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Unsupported_method, got %s"
               (Did_resolver.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

(** {1 Handle Resolution Tests} *)

(** Mock DNS handler *)
let mock_dns_handler : (string -> Handle_resolver.dns_result) ref =
  ref (fun _ -> Handle_resolver.Dns_not_found)

(** Mock HTTP handler for handle resolution *)
let mock_handle_http_handler : (Uri.t -> Handle_resolver.http_response) ref =
  ref (fun _ -> Handle_resolver.{ status = 500; body = "No mock" })

(** Combined effect handler for handle resolution *)
let handle_effect_handler : type a.
    a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option = function
  | Handle_resolver.Dns_txt domain ->
      Some (fun k -> Effect.Deep.continue k (!mock_dns_handler domain))
  | Handle_resolver.Http_get uri ->
      Some (fun k -> Effect.Deep.continue k (!mock_handle_http_handler uri))
  | Did_resolver.Http_get uri ->
      Some (fun k -> Effect.Deep.continue k (!mock_http_handler uri))
  | _ -> None

(** Run with mock handlers for handle resolution *)
let run_with_handle_mocks ~dns_handler ~http_handler f =
  mock_dns_handler := dns_handler;
  mock_handle_http_handler := http_handler;
  Effect.Deep.match_with f ()
    { retc = (fun x -> x); exnc = raise; effc = handle_effect_handler }

let test_handle_resolve_via_dns () =
  let dns_handler domain =
    if domain = "_atproto.alice.bsky.social" then
      Handle_resolver.Dns_records [ "did=did:plc:alice123" ]
    else Handle_resolver.Dns_not_found
  in
  let http_handler _uri =
    Handle_resolver.{ status = 404; body = "Not found" }
  in
  run_with_handle_mocks ~dns_handler ~http_handler (fun () ->
      match Handle_resolver.resolve_string "alice.bsky.social" with
      | Ok did ->
          Alcotest.(check string)
            "did" "did:plc:alice123"
            (Atproto_syntax.Did.to_string did)
      | Error e -> Alcotest.fail (Handle_resolver.error_to_string e))

let test_handle_resolve_via_https () =
  let dns_handler _domain = Handle_resolver.Dns_not_found in
  let http_handler uri =
    let host = Uri.host uri |> Option.value ~default:"" in
    let path = Uri.path uri in
    if host = "bob.example.com" && path = "/.well-known/atproto-did" then
      Handle_resolver.{ status = 200; body = "did:web:bob.example.com" }
    else Handle_resolver.{ status = 404; body = "Not found" }
  in
  run_with_handle_mocks ~dns_handler ~http_handler (fun () ->
      match Handle_resolver.resolve_string "bob.example.com" with
      | Ok did ->
          Alcotest.(check string)
            "did" "did:web:bob.example.com"
            (Atproto_syntax.Did.to_string did)
      | Error e -> Alcotest.fail (Handle_resolver.error_to_string e))

let test_handle_dns_priority () =
  (* DNS should be tried first, even if HTTPS would work *)
  let dns_handler domain =
    if domain = "_atproto.test.example.com" then
      Handle_resolver.Dns_records [ "did=did:plc:from-dns" ]
    else Handle_resolver.Dns_not_found
  in
  let http_handler _uri =
    Handle_resolver.{ status = 200; body = "did:plc:from-https" }
  in
  run_with_handle_mocks ~dns_handler ~http_handler (fun () ->
      match Handle_resolver.resolve_string "test.example.com" with
      | Ok did ->
          Alcotest.(check string)
            "prefers DNS" "did:plc:from-dns"
            (Atproto_syntax.Did.to_string did)
      | Error e -> Alcotest.fail (Handle_resolver.error_to_string e))

let test_handle_not_found () =
  let dns_handler _domain = Handle_resolver.Dns_not_found in
  let http_handler _uri =
    Handle_resolver.{ status = 404; body = "Not found" }
  in
  run_with_handle_mocks ~dns_handler ~http_handler (fun () ->
      match Handle_resolver.resolve_string "notfound.example.com" with
      | Error Handle_resolver.No_did_record -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected No_did_record, got %s"
               (Handle_resolver.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_handle_invalid () =
  let dns_handler _domain = Handle_resolver.Dns_not_found in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_handle_mocks ~dns_handler ~http_handler (fun () ->
      match Handle_resolver.resolve_string "invalid" with
      | Error (Handle_resolver.Invalid_handle _) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Invalid_handle, got %s"
               (Handle_resolver.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

(** {1 Identity Verification Tests} *)

(** Combined effect handler for identity verification *)
let identity_effect_handler : type a.
    a Effect.t -> ((a, _) Effect.Deep.continuation -> _) option = function
  | Handle_resolver.Dns_txt domain ->
      Some (fun k -> Effect.Deep.continue k (!mock_dns_handler domain))
  | Handle_resolver.Http_get uri ->
      Some (fun k -> Effect.Deep.continue k (!mock_handle_http_handler uri))
  | Did_resolver.Http_get uri ->
      Some (fun k -> Effect.Deep.continue k (!mock_http_handler uri))
  | _ -> None

let run_with_identity_mocks ~did_handler ~dns_handler ~http_handler f =
  mock_http_handler := did_handler;
  mock_dns_handler := dns_handler;
  mock_handle_http_handler := http_handler;
  Effect.Deep.match_with f ()
    { retc = (fun x -> x); exnc = raise; effc = identity_effect_handler }

let test_verify_did_success () =
  (* Setup: DID doc has handle, handle resolves back to DID *)
  let did_handler uri =
    let path = Uri.path uri in
    if path = "/did:plc:test123" then
      Did_resolver.
        {
          status = 200;
          body =
            {|{
        "id": "did:plc:test123",
        "alsoKnownAs": ["at://alice.example.com"],
        "verificationMethod": [
          {"id": "#key", "type": "Multikey", "controller": "did:plc:test123", "publicKeyMultibase": "zTest123"}
        ],
        "service": [
          {"id": "#pds", "type": "AtprotoPersonalDataServer", "serviceEndpoint": "https://pds.example.com"}
        ]
      }|};
        }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  let dns_handler domain =
    if domain = "_atproto.alice.example.com" then
      Handle_resolver.Dns_records [ "did=did:plc:test123" ]
    else Handle_resolver.Dns_not_found
  in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let did = Atproto_syntax.Did.of_string_exn "did:plc:test123" in
      match Identity.verify_did did with
      | Ok identity ->
          Alcotest.(check string)
            "did" "did:plc:test123"
            (Atproto_syntax.Did.to_string identity.did);
          Alcotest.(check string)
            "handle" "alice.example.com"
            (Atproto_syntax.Handle.to_string identity.handle);
          Alcotest.(check bool)
            "has signing key" true
            (Option.is_some identity.signing_key);
          Alcotest.(check bool)
            "has pds" true
            (Option.is_some identity.pds_endpoint)
      | Error e -> Alcotest.fail (Identity.error_to_string e))

let test_verify_handle_success () =
  let did_handler uri =
    let path = Uri.path uri in
    if path = "/did:plc:bob456" then
      Did_resolver.
        {
          status = 200;
          body =
            {|{
        "id": "did:plc:bob456",
        "alsoKnownAs": ["at://bob.example.com"],
        "verificationMethod": [],
        "service": []
      }|};
        }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  let dns_handler domain =
    if domain = "_atproto.bob.example.com" then
      Handle_resolver.Dns_records [ "did=did:plc:bob456" ]
    else Handle_resolver.Dns_not_found
  in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let handle = Atproto_syntax.Handle.of_string_exn "bob.example.com" in
      match Identity.verify_handle handle with
      | Ok identity ->
          Alcotest.(check string)
            "did" "did:plc:bob456"
            (Atproto_syntax.Did.to_string identity.did);
          Alcotest.(check string)
            "handle" "bob.example.com"
            (Atproto_syntax.Handle.to_string identity.handle)
      | Error e -> Alcotest.fail (Identity.error_to_string e))

let test_verify_bidirectional_success () =
  let did_handler uri =
    let path = Uri.path uri in
    if path = "/did:plc:carol789" then
      Did_resolver.
        {
          status = 200;
          body =
            {|{
        "id": "did:plc:carol789",
        "alsoKnownAs": ["at://carol.example.com"],
        "verificationMethod": [],
        "service": []
      }|};
        }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  let dns_handler domain =
    if domain = "_atproto.carol.example.com" then
      Handle_resolver.Dns_records [ "did=did:plc:carol789" ]
    else Handle_resolver.Dns_not_found
  in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let did = Atproto_syntax.Did.of_string_exn "did:plc:carol789" in
      let handle = Atproto_syntax.Handle.of_string_exn "carol.example.com" in
      match Identity.verify_bidirectional did handle with
      | Ok identity ->
          Alcotest.(check string)
            "did" "did:plc:carol789"
            (Atproto_syntax.Did.to_string identity.did)
      | Error e -> Alcotest.fail (Identity.error_to_string e))

let test_verify_did_handle_mismatch () =
  (* Handle in doc doesn't match what we expect *)
  let did_handler uri =
    let path = Uri.path uri in
    if path = "/did:plc:mismatch" then
      Did_resolver.
        {
          status = 200;
          body =
            {|{
        "id": "did:plc:mismatch",
        "alsoKnownAs": ["at://wrong.example.com"],
        "verificationMethod": [],
        "service": []
      }|};
        }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  let dns_handler domain =
    if domain = "_atproto.wrong.example.com" then
      (* Handle resolves to different DID *)
      Handle_resolver.Dns_records [ "did=did:plc:different" ]
    else Handle_resolver.Dns_not_found
  in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let did = Atproto_syntax.Did.of_string_exn "did:plc:mismatch" in
      match Identity.verify_did did with
      | Error (Identity.Did_mismatch _) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Did_mismatch, got %s"
               (Identity.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_verify_no_handle_in_doc () =
  let did_handler uri =
    let path = Uri.path uri in
    if path = "/did:plc:nohandle" then
      Did_resolver.
        {
          status = 200;
          body =
            {|{
        "id": "did:plc:nohandle",
        "alsoKnownAs": [],
        "verificationMethod": [],
        "service": []
      }|};
        }
    else Did_resolver.{ status = 404; body = "Not found" }
  in
  let dns_handler _domain = Handle_resolver.Dns_not_found in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let did = Atproto_syntax.Did.of_string_exn "did:plc:nohandle" in
      match Identity.verify_did did with
      | Error Identity.No_handle_in_document -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected No_handle_in_document, got %s"
               (Identity.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_verify_did_resolution_failed () =
  let did_handler _uri = Did_resolver.{ status = 404; body = "Not found" } in
  let dns_handler _domain = Handle_resolver.Dns_not_found in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let did = Atproto_syntax.Did.of_string_exn "did:plc:notfound" in
      match Identity.verify_did did with
      | Error (Identity.Did_resolution_failed _) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Did_resolution_failed, got %s"
               (Identity.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

let test_verify_handle_resolution_failed () =
  let did_handler _uri =
    Did_resolver.
      {
        status = 200;
        body =
          {|{
    "id": "did:plc:test",
    "alsoKnownAs": ["at://test.example.com"],
    "verificationMethod": [],
    "service": []
  }|};
      }
  in
  let dns_handler _domain = Handle_resolver.Dns_not_found in
  let http_handler _uri = Handle_resolver.{ status = 404; body = "" } in
  run_with_identity_mocks ~did_handler ~dns_handler ~http_handler (fun () ->
      let handle = Atproto_syntax.Handle.of_string_exn "notfound.example.com" in
      match Identity.verify_handle handle with
      | Error (Identity.Handle_resolution_failed _) -> ()
      | Error e ->
          Alcotest.fail
            (Printf.sprintf "expected Handle_resolution_failed, got %s"
               (Identity.error_to_string e))
      | Ok _ -> Alcotest.fail "expected error")

(** {1 Test Suites} *)

let resolver_tests =
  [
    Alcotest.test_case "resolve did:plc" `Quick test_resolve_plc;
    Alcotest.test_case "resolve did:web" `Quick test_resolve_web;
    Alcotest.test_case "get handle" `Quick test_get_handle;
    Alcotest.test_case "get PDS endpoint" `Quick test_get_pds_endpoint;
    Alcotest.test_case "get signing key" `Quick test_get_signing_key;
    Alcotest.test_case "not found" `Quick test_not_found;
    Alcotest.test_case "http error" `Quick test_http_error;
    Alcotest.test_case "invalid did" `Quick test_invalid_did;
    Alcotest.test_case "unsupported method" `Quick test_unsupported_method;
  ]

let handle_resolver_tests =
  [
    Alcotest.test_case "resolve via DNS" `Quick test_handle_resolve_via_dns;
    Alcotest.test_case "resolve via HTTPS" `Quick test_handle_resolve_via_https;
    Alcotest.test_case "DNS priority" `Quick test_handle_dns_priority;
    Alcotest.test_case "not found" `Quick test_handle_not_found;
    Alcotest.test_case "invalid handle" `Quick test_handle_invalid;
  ]

let identity_tests =
  [
    Alcotest.test_case "verify DID success" `Quick test_verify_did_success;
    Alcotest.test_case "verify handle success" `Quick test_verify_handle_success;
    Alcotest.test_case "verify bidirectional" `Quick
      test_verify_bidirectional_success;
    Alcotest.test_case "DID mismatch" `Quick test_verify_did_handle_mismatch;
    Alcotest.test_case "no handle in doc" `Quick test_verify_no_handle_in_doc;
    Alcotest.test_case "DID resolution failed" `Quick
      test_verify_did_resolution_failed;
    Alcotest.test_case "handle resolution failed" `Quick
      test_verify_handle_resolution_failed;
  ]

let () =
  Alcotest.run "atproto-identity"
    [
      ("did_resolver", resolver_tests);
      ("handle_resolver", handle_resolver_tests);
      ("identity", identity_tests);
    ]