diff --git a/src/lib_base/block_header.ml b/src/lib_base/block_header.ml index 257553b45..15e4adf7d 100644 --- a/src/lib_base/block_header.ml +++ b/src/lib_base/block_header.ml @@ -77,6 +77,11 @@ let encoding = shell_header_encoding (obj1 (req "protocol_data" Variable.bytes))) +let bounded_encoding ?max_size () = + match max_size with + | None -> encoding + | Some max_size -> Data_encoding.check_size max_size encoding + let pp ppf op = Data_encoding.Json.pp ppf (Data_encoding.Json.construct encoding op) diff --git a/src/lib_base/block_header.mli b/src/lib_base/block_header.mli index 1e8f54b2b..3be5b4443 100644 --- a/src/lib_base/block_header.mli +++ b/src/lib_base/block_header.mli @@ -28,3 +28,5 @@ type t = { include S.HASHABLE with type t := t and type hash := Block_hash.t val of_bytes_exn: MBytes.t -> t + +val bounded_encoding: ?max_size:int -> unit -> t Data_encoding.t diff --git a/src/lib_base/block_locator.ml b/src/lib_base/block_locator.ml index 396922ec8..234ba685f 100644 --- a/src/lib_base/block_locator.ml +++ b/src/lib_base/block_locator.ml @@ -41,7 +41,17 @@ let encoding = (* TODO add a [description] *) (obj2 (req "current_head" (dynamic_size Block_header.encoding)) - (req "history" (dynamic_size (list Block_hash.encoding)))) + (req "history" (Variable.list Block_hash.encoding))) + +let bounded_encoding ?max_header_size ?max_length () = + let open Data_encoding in + (* TODO add a [description] *) + (obj2 + (req "current_head" + (dynamic_size + (Block_header.bounded_encoding ?max_size:max_header_size ()))) + (req "history" (Variable.list ?max_length Block_hash.encoding))) + type seed = { sender_id: P2p_peer.Id.t ; diff --git a/src/lib_base/block_locator.mli b/src/lib_base/block_locator.mli index 504d122cd..28f17bbfe 100644 --- a/src/lib_base/block_locator.mli +++ b/src/lib_base/block_locator.mli @@ -17,6 +17,10 @@ val raw: t -> raw val pp: Format.formatter -> t -> unit val pp_short: Format.formatter -> t -> unit val encoding: t Data_encoding.t +val bounded_encoding: + ?max_header_size:int -> + ?max_length:int -> + unit -> t Data_encoding.t type seed = { sender_id: P2p_peer.Id.t ; diff --git a/src/lib_shell/distributed_db_message.ml b/src/lib_shell/distributed_db_message.ml index e91bfec23..28a95bb63 100644 --- a/src/lib_shell/distributed_db_message.ml +++ b/src/lib_shell/distributed_db_message.ml @@ -7,6 +7,28 @@ (* *) (**************************************************************************) +module Bounded_encoding = struct + + open Data_encoding + + let block_header_max_size = ref None + let block_header_cache = ref Block_header.encoding + let block_locator_cache = ref Block_locator.encoding + + let update_block_header_encoding () = + block_header_cache := + Block_header.bounded_encoding ?max_size:!block_header_max_size () ; + block_locator_cache := + Block_locator.bounded_encoding ?max_header_size:!block_header_max_size () + + let set_block_header_max_size max = + block_header_max_size := max ; + update_block_header_encoding () + let block_header = delayed (fun () -> !block_header_cache) + let block_locator = delayed (fun () -> !block_locator_cache) + +end + type t = | Get_current_branch of Chain_id.t @@ -53,7 +75,7 @@ let encoding = ~title:"Current_branch" (obj2 (req "chain_id" Chain_id.encoding) - (req "current_branch" Block_locator.encoding)) + (req "current_branch" Bounded_encoding.block_locator)) (function | Current_branch (chain_id, locator) -> Some (chain_id, locator) | _ -> None) @@ -81,7 +103,7 @@ let encoding = ~title:"Current_head" (obj3 (req "chain_id" Chain_id.encoding) - (req "current_block_header" (dynamic_size Block_header.encoding)) + (req "current_block_header" (dynamic_size Bounded_encoding.block_header)) (req "current_mempool" Mempool.encoding)) (function | Current_head (chain_id, bh, mempool) -> Some (chain_id, bh, mempool) @@ -98,7 +120,7 @@ let encoding = case ~tag:0x21 ~title:"Block_header" - (obj1 (req "block_header" Block_header.encoding)) + (obj1 (req "block_header" Bounded_encoding.block_header)) (function | Block_header bh -> Some bh | _ -> None) diff --git a/src/lib_shell/distributed_db_message.mli b/src/lib_shell/distributed_db_message.mli index 29d43cdbe..d44213dea 100644 --- a/src/lib_shell/distributed_db_message.mli +++ b/src/lib_shell/distributed_db_message.mli @@ -40,3 +40,7 @@ type t = val cfg : t P2p.message_config val pp_json : Format.formatter -> t -> unit + +module Bounded_encoding : sig + val set_block_header_max_size: int option -> unit +end