Client: cache public keys to limit queries to remote/hardware signers

src/lib_client_base/client_keys.ml

@@ -50,15 +50,6 @@ end
 let uri_encoding =
   Data_encoding.(conv Uri.to_string Uri.of_string string)
 
-module Entity (Name : sig val name: string end) = struct
-  include Name
-  type t = Uri.t
-  let pp = Uri.pp_hum
-  let of_source s =  return (Uri.of_string s)
-  let to_source t = return (Uri.to_string t)
-  let encoding = uri_encoding
-end
-
 type pk_uri = Uri.t
 let make_pk_uri x = x
 
@@ -93,9 +84,38 @@ let sk_uri_param ?name ?desc params =
     params
 
 module Secret_key =
-  Client_aliases.Alias (Entity(struct let name = "secret_key" end))
+  Client_aliases.Alias (struct
+    let name = "secret_key"
+    type t = Uri.t
+    let pp = Uri.pp_hum
+    let of_source s = return (Uri.of_string s)
+    let to_source t = return (Uri.to_string t)
+    let encoding = uri_encoding
+  end)
+
 module Public_key =
-  Client_aliases.Alias (Entity(struct let name = "public_key" end))
+  Client_aliases.Alias (struct
+    let name = "public_key"
+    type t = Uri.t * Signature.Public_key.t option
+    let pp ppf (loc, _) = Uri.pp_hum ppf loc
+    let of_source s = return (Uri.of_string s, None)
+    let to_source (t, _) = return (Uri.to_string t)
+    let encoding =
+      let open Data_encoding in
+      union
+        [ case Json_only
+            ~title: "Locator_only"
+            uri_encoding
+            (function (uri, None) -> Some uri | (_, Some _) -> None)
+            (fun uri -> (uri, None)) ;
+          case Json_only
+            ~title: "Locator_and_full_key"
+            (obj2
+               (req "locator" uri_encoding)
+               (req "key" Signature.Public_key.encoding))
+            (function (uri, Some key) -> Some (uri, key) | (_, None) -> None)
+            (fun (uri, key) -> (uri, Some key)) ]
+  end)
 
 module type SIGNER = sig
   val scheme : string
@@ -103,7 +123,7 @@ module type SIGNER = sig
   val description : string
   val neuterize : sk_uri -> pk_uri tzresult Lwt.t
   val public_key : pk_uri -> Signature.Public_key.t tzresult Lwt.t
-  val public_key_hash : pk_uri -> Signature.Public_key_hash.t tzresult Lwt.t
+  val public_key_hash : pk_uri -> (Signature.Public_key_hash.t * Signature.Public_key.t option) tzresult Lwt.t
   val sign :
     ?watermark: Signature.watermark ->
     sk_uri -> MBytes.t -> Signature.t tzresult Lwt.t
@@ -154,7 +174,7 @@ let public_key pk_uri =
 
 let public_key_hash pk_uri =
   public_key pk_uri >>=? fun pk ->
-  return (Signature.Public_key.hash pk)
+  return (Signature.Public_key.hash pk, Some pk)
 
 let sign ?watermark sk_uri buf =
   let scheme = Option.unopt ~default:"" (Uri.scheme sk_uri) in
@@ -176,8 +196,8 @@ let check ?watermark pk_uri signature buf =
   public_key pk_uri >>=? fun pk ->
   return (Signature.check ?watermark pk signature buf)
 
-let register_key cctxt ?(force=false) (public_key_hash, pk_uri, sk_uri) name =
+let register_key cctxt ?(force=false) (public_key_hash, pk_uri, sk_uri) ?public_key name =
-  Public_key.add ~force cctxt name pk_uri >>=? fun () ->
+  Public_key.add ~force cctxt name (pk_uri, public_key) >>=? fun () ->
   Secret_key.add ~force cctxt name sk_uri >>=? fun () ->
   Public_key_hash.add ~force cctxt name public_key_hash >>=? fun () ->
   return ()
@@ -187,19 +207,15 @@ let raw_get_key (cctxt : #Client_context.wallet) pkh =
     Public_key_hash.rev_find cctxt pkh >>=? function
     | None -> failwith "no keys for the source contract manager"
     | Some n ->
-        Public_key.find_opt cctxt n >>=? fun pk_uri ->
         Secret_key.find_opt cctxt n >>=? fun sk_uri ->
-        begin
+        Public_key.find_opt cctxt n >>=? begin function
-          Option.unopt_map
+          | None -> return None
-            ~default:Lwt.return_none
+          | Some (_, Some pk) -> return (Some pk)
-            ~f:(fun pkh ->
+          | Some (pk_uri, None) ->
-                public_key pkh >>= function
+              public_key pk_uri >>=? fun pk ->
-                | Error e ->
+              Public_key.update cctxt n (pk_uri, Some pk) >>=? fun () ->
-                    Format.eprintf "PLOP: %a@." pp_print_error e ;
+              return (Some pk)
-                    Lwt.return_none
+        end >>=? fun pk ->
-                | Ok pk -> Lwt.return_some pk)
-            pk_uri
-        end >>= fun pk ->
         return (n, pk, sk_uri)
   end >>= function
   | (Ok (_, None, None) | Error _) as initial_result -> begin
@@ -232,9 +248,14 @@ let get_keys (cctxt : #Client_context.wallet) =
   Secret_key.load cctxt >>=? fun sks ->
   Lwt_list.filter_map_s begin fun (name, sk_uri) ->
     begin
-      Public_key.find cctxt name >>=? fun pk_uri ->
       Public_key_hash.find cctxt name >>=? fun pkh ->
+      Public_key.find cctxt name >>=? begin function
+        | _, Some pk -> return pk
+        | pk_uri, None ->
             public_key pk_uri >>=? fun pk ->
+            Public_key.update cctxt name (pk_uri, Some pk) >>=? fun () ->
+            return pk
+      end >>=? fun pk ->
       return (name, pkh, pk, sk_uri)
     end >>= function
     | Ok r -> Lwt.return (Some r)

src/lib_client_base/client_keys.mli

@@ -25,7 +25,7 @@ type error += Invalid_uri of Uri.t
 module Public_key_hash :
   Client_aliases.Alias with type t = Signature.Public_key_hash.t
 module Public_key :
-  Client_aliases.Alias with type t = pk_uri
+  Client_aliases.Alias with type t = pk_uri * Signature.Public_key.t option
 module Secret_key :
   Client_aliases.Alias with type t = sk_uri
 
@@ -50,8 +50,10 @@ module type SIGNER = sig
   val public_key : pk_uri -> Signature.Public_key.t tzresult Lwt.t
   (** [public_key pk] is the Ed25519 version of [pk]. *)
 
-  val public_key_hash : pk_uri -> Signature.Public_key_hash.t tzresult Lwt.t
+  val public_key_hash : pk_uri -> (Signature.Public_key_hash.t * Signature.Public_key.t option) tzresult Lwt.t
-  (** [public_key_hash pk] is the hash of [pk]. *)
+  (** [public_key_hash pk] is the hash of [pk].
+      As some signers will query the full public key to obtain the hash,
+      it can be optionally returned to reduce the amount of queries. *)
 
   val sign :
     ?watermark: Signature.watermark ->
@@ -69,7 +71,7 @@ val registered_signers : unit -> (string * (module SIGNER)) list
 
 val public_key : pk_uri -> Signature.Public_key.t tzresult Lwt.t
 
-val public_key_hash : pk_uri -> Signature.Public_key_hash.t tzresult Lwt.t
+val public_key_hash : pk_uri -> (Signature.Public_key_hash.t * Signature.Public_key.t option) tzresult Lwt.t
 
 val neuterize : sk_uri -> pk_uri tzresult Lwt.t
 
@@ -88,7 +90,9 @@ val check :
 val register_key :
   #Client_context.wallet ->
   ?force:bool ->
-  (Signature.Public_key_hash.t * pk_uri * sk_uri) -> string -> unit tzresult Lwt.t
+  (Signature.Public_key_hash.t * pk_uri * sk_uri) ->
+  ?public_key: Signature.Public_key.t ->
+  string -> unit tzresult Lwt.t
 
 val list_keys :
   #Client_context.wallet ->

src/lib_client_commands/client_keys_commands.ml

@@ -174,17 +174,17 @@ let commands () : Client_context.io_wallet Clic.command list =
          begin
            Public_key.find_opt cctxt name >>=? function
            | None -> return ()
-           | Some pk ->
+           | Some (pk_uri_found, _) ->
-               fail_unless (pk_uri = pk || force)
+               fail_unless (pk_uri = pk_uri_found || force)
                  (failure
                     "public and secret keys '%s' don't correspond, \
                      please don't use -force" name)
          end >>=? fun () ->
-         Client_keys.public_key_hash pk_uri >>=? fun pkh ->
+         Client_keys.public_key_hash pk_uri >>=? fun (pkh, public_key) ->
          cctxt#message
            "Tezos address added: %a"
            Signature.Public_key_hash.pp pkh >>= fun () ->
-         register_key cctxt ~force (pkh, pk_uri, sk_uri) name) ;
+         register_key cctxt ~force (pkh, pk_uri, sk_uri) ?public_key name) ;
 
     command ~group ~desc: "Add a public key to the wallet."
       (args1 (Public_key.force_switch ()))
@@ -195,12 +195,12 @@ let commands () : Client_context.io_wallet Clic.command list =
        @@ stop)
       (fun force name pk_uri (cctxt : Client_context.io_wallet) ->
          Public_key.of_fresh cctxt force name >>=? fun name ->
-         Client_keys.public_key_hash pk_uri >>=? fun pkh ->
+         Client_keys.public_key_hash pk_uri >>=? fun (pkh, public_key) ->
          Public_key_hash.add ~force cctxt name pkh >>=? fun () ->
          cctxt#message
            "Tezos address added: %a"
            Signature.Public_key_hash.pp pkh >>= fun () ->
-         Public_key.add ~force cctxt name pk_uri) ;
+         Public_key.add ~force cctxt name (pk_uri, public_key)) ;
 
     command ~group ~desc: "Add an identity to the wallet."
       (args1 (Public_key.force_switch ()))

src/lib_signer_backends/http_gen.ml

@@ -55,7 +55,7 @@ module Make(N : sig val scheme : string end) = struct
 
     let public_key_hash uri =
       public_key uri >>=? fun pk ->
-      return (Signature.Public_key.hash pk)
+      return (Signature.Public_key.hash pk, Some pk)
 
     let sign ?watermark uri msg =
       parse (uri : sk_uri :> Uri.t) >>=? fun (base, pkh) ->

src/lib_signer_backends/ledger.ml

@@ -188,10 +188,10 @@ let public_key (pk_uri : pk_uri) =
 
 let public_key_hash pk_uri =
   match Hashtbl.find_opt pkhs pk_uri with
-  | Some pkh -> return pkh
+  | Some pkh -> return (pkh, None)
   | None ->
-      public_key pk_uri >>=? fun _pk ->
+      public_key pk_uri >>=? fun pk ->
-      return (Hashtbl.find pkhs pk_uri)
+      return (Hashtbl.find pkhs pk_uri, Some pk)
 
 let sign ?watermark sk_uri msg =
   pkh_of_sk_uri sk_uri >>=? fun pkh ->
@@ -318,7 +318,7 @@ let commands =
                cctxt#message "Found a valid Tezos application running on %s %s at [%s]."
                  manufacturer product device_info.path >>= fun () ->
                public_key pk_uri >>=? fun pk ->
-               public_key_hash pk_uri >>=? fun pkh ->
+               public_key_hash pk_uri >>=? fun (pkh, _) ->
                let pkh_bytes = Signature.Public_key_hash.to_bytes pkh in
                sign ~watermark:Generic_operation
                  sk_uri pkh_bytes >>=? fun signature ->

src/lib_signer_backends/socket.ml

@@ -81,7 +81,7 @@ module Make(P : sig
 
     let public_key_hash uri =
       public_key uri >>=? fun pk ->
-      return (Signature.Public_key.hash pk)
+      return (Signature.Public_key.hash pk, Some pk)
 
     let sign ?watermark uri msg =
       parse (uri : sk_uri :> Uri.t) >>=? fun (path, pkh) ->
@@ -122,7 +122,7 @@ module Make(P : sig
 
     let public_key_hash uri =
       public_key uri >>=? fun pk ->
-      return (Signature.Public_key.hash pk)
+      return (Signature.Public_key.hash pk, Some pk)
 
     let sign ?watermark uri msg =
       parse (uri : sk_uri :> Uri.t) >>=? fun (path, pkh) ->

src/lib_signer_backends/unencrypted.ml

@@ -45,7 +45,7 @@ let neuterize sk_uri =
 
 let public_key_hash pk_uri =
   public_key pk_uri >>=? fun pk ->
-  return (Signature.Public_key.hash pk)
+  return (Signature.Public_key.hash pk, Some pk)
 
 let sign ?watermark sk_uri buf =
   secret_key sk_uri >>=? fun sk ->

src/proto_alpha/lib_client_commands/client_proto_programs_commands.ml

@@ -264,7 +264,7 @@ let commands () =
        @@ Clic.param ~name:"signature" ~desc:"the signature to check"
          signature_parameter
        @@ stop)
-      (fun quiet hashed (_, key_locator) signature (cctxt : #Proto_alpha.full) ->
+      (fun quiet hashed (_, (key_locator, _)) signature (cctxt : #Proto_alpha.full) ->
          Client_keys.check key_locator signature hashed >>=? fun res ->
          begin
            if quiet