From fed04d4b8f95e2a2c3d43d50750dbbe08b77cf44 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Gr=C3=A9goire=20Henry?= <gregoire.henry@tezos.com>
Date: Sat, 2 Jun 2018 18:23:08 +0200
Subject: [PATCH] P2P: allow to limit the size of protocols

---
 src/lib_base/protocol.ml                 |  5 +++++
 src/lib_base/protocol.mli                |  2 ++
 src/lib_shell/distributed_db_message.ml  | 12 +++++++++++-
 src/lib_shell/distributed_db_message.mli |  1 +
 4 files changed, 19 insertions(+), 1 deletion(-)

diff --git a/src/lib_base/protocol.ml b/src/lib_base/protocol.ml
index 602dc81bc..6864b6b31 100644
--- a/src/lib_base/protocol.ml
+++ b/src/lib_base/protocol.ml
@@ -53,6 +53,11 @@ let encoding =
        (req "expected_env_version" env_version_encoding)
        (req "components" (list component_encoding)))
 
+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/protocol.mli b/src/lib_base/protocol.mli
index 516a6adb9..a86b738d0 100644
--- a/src/lib_base/protocol.mli
+++ b/src/lib_base/protocol.mli
@@ -29,6 +29,8 @@ include S.HASHABLE with type t := t
                     and type hash := Protocol_hash.t
 val of_bytes_exn: MBytes.t -> t
 
+val bounded_encoding: ?max_size:int -> unit -> t Data_encoding.t
+
 module Meta: sig
 
   type t = {
diff --git a/src/lib_shell/distributed_db_message.ml b/src/lib_shell/distributed_db_message.ml
index 82fb4b050..d3506357f 100644
--- a/src/lib_shell/distributed_db_message.ml
+++ b/src/lib_shell/distributed_db_message.ml
@@ -83,6 +83,16 @@ module Bounded_encoding = struct
   let operation_list = delayed (fun () -> !operation_list_cache)
   let operation_hash_list = delayed (fun () -> !operation_hash_list_cache)
 
+  let protocol_max_size = ref None
+  let protocol_cache =
+    ref (Protocol.bounded_encoding ?max_size:!protocol_max_size ())
+  let update_protocol_encoding () =
+    protocol_cache :=
+      Protocol.bounded_encoding ?max_size:!protocol_max_size ()
+  let set_protocol_max_size max =
+    protocol_max_size := max
+  let protocol = delayed (fun () -> !protocol_cache)
+
 end
 
 type t =
@@ -207,7 +217,7 @@ let encoding =
 
     case ~tag:0x41
       ~title:"Protocol"
-      (obj1 (req "protocol" Protocol.encoding))
+      (obj1 (req "protocol" Bounded_encoding.protocol))
       (function Protocol proto -> Some proto  | _ -> None)
       (fun proto -> Protocol proto);
 
diff --git a/src/lib_shell/distributed_db_message.mli b/src/lib_shell/distributed_db_message.mli
index aa6ac631f..3ce24c929 100644
--- a/src/lib_shell/distributed_db_message.mli
+++ b/src/lib_shell/distributed_db_message.mli
@@ -47,4 +47,5 @@ module Bounded_encoding : sig
   val set_operation_list_max_size: int option -> unit
   val set_operation_list_max_length: int option -> unit
   val set_operation_max_pass: int option -> unit
+  val set_protocol_max_size: int option -> unit
 end