multisigv2 more robust to spam

Set limits to the number of approvals and message sizes

src/test/contracts/multisig-v2.ligo

@@ -1,46 +1,90 @@
 // storage type
 type threshold_t is nat
+type max_proposal_t is nat
+type max_message_size_t is nat
 type addr_set_t is set(address)
-type message_store_t is big_map(bytes,addr_set_t)
+type message_store_t is map(bytes,addr_set_t)
+type counter_store_t is map(address,nat)
 
 type storage_t is record
   threshold : threshold_t ;
+  max_proposal : max_proposal_t ;
+  max_message_size : max_message_size_t ;
   auth : addr_set_t ;
   message_store : message_store_t ;
+  counter_store : counter_store_t ;
 end
 
 // I/O types
 type message_t is (unit -> list(operation))
 type send_pt is message_t
+type withdraw_pt is message_t
 
 type contract_return_t is (list(operation) * storage_t)
 
 type entry_point_t is
 | Send of send_pt
+| Withdraw of withdraw_pt
 
 function send (const param : send_pt; const s : storage_t) : contract_return_t is block {
 
   if not set_mem(sender,s.auth) then failwith("Unauthorized address") else skip ;
   
   var message : message_t := param ;
-  const packed_msg : bytes = bytes_pack(message) ;
+  var new_store : addr_set_t := set_empty ;
   var ret_ops : list(operation) := (nil : list(operation)) ;
 
-  var new_store : addr_set_t :=
+  const packed_msg : bytes = bytes_pack(message) ;
-    case map_get(packed_msg, s.message_store) of
+  if size(packed_msg) > s.max_message_size then failwith("Message size exceed maximum limit")
-    | Some(voters) -> set_add(sender,voters)
+  else skip ;
-    | None         -> set sender end end
+
-  ;
+  case map_get(packed_msg, s.message_store) of
+  | Some(voters) -> block {
+    if set_mem(sender,voters) then skip
+    else s.counter_store[sender] := get_force(sender,s.counter_store) + 1n ;
+    new_store := set_add(sender,voters)
+  }
+  | None -> block {
+    s.counter_store[sender] := get_force(sender,s.counter_store) + 1n ;
+    new_store := set [sender];
+  }
+  end ;
+
+  var sender_proposal_counter : nat := get_force(sender,s.counter_store) ;
+  if sender_proposal_counter > s.max_proposal then failwith("Maximum number of proposal reached")
+  else skip ;
 
   if size(new_store) >= s.threshold then block {
     remove packed_msg from map s.message_store ;
     ret_ops := message(unit) ;
+    s.counter_store[sender] := abs (sender_proposal_counter - 1n) ;
   } else
     s.message_store[packed_msg] := new_store
   
 } with ( ret_ops , s)
 
+function withdraw (const param : withdraw_pt; const s : storage_t) : contract_return_t is block {
+
+  var message : message_t := param ;
+  const packed_msg : bytes = bytes_pack(message) ;
+
+  case map_get(packed_msg, s.message_store) of
+  | Some(voters) -> block {
+    const new_set : addr_set_t = set_remove(sender,voters) ;
+
+    if size(voters) =/= size(new_set) then
+      s.counter_store[sender] := abs (get_force(sender,s.counter_store) - 1n)
+    else skip ;
+
+    if size(new_set) = 0n then remove packed_msg from map s.message_store
+    else s.message_store[packed_msg] := new_set
+  }
+  | None -> skip end
+
+} with ( (nil: list(operation)) , s)
+
 function main(const param : entry_point_t; const s : storage_t) : contract_return_t is 
   case param of
   | Send (p) -> send(p,s)
+  | Withdraw (p) -> withdraw(p,s)
 end

src/test/multisig_tests.ml

@@ -22,18 +22,6 @@ let compile_main () =
 
 open Ast_simplified
 
-let gen_keys = fun () ->
-  let open Tezos_crypto in
-  let (raw_pkh,raw_pk,raw_sk) = Signature.generate_key () in
-  (raw_pkh,raw_pk,raw_sk)
-
-let str_keys (raw_pkh, raw_pk, raw_sk) =
-  let open Tezos_crypto in
-  let sk_str = Signature.Secret_key.to_b58check raw_sk in
-  let pk_str = Signature.Public_key.to_b58check raw_pk in
-  let pkh_str = Signature.Public_key_hash.to_b58check raw_pkh in
-  (pkh_str,pk_str,sk_str)
-
 let init_storage threshold counter pkeys =
   let keys = List.map
     (fun el ->

src/test/multisig_v2_tests.ml

@@ -22,102 +22,223 @@ let compile_main () =
 
 open Ast_simplified
 
-let contract id = 
-  let open Proto_alpha_utils.Memory_proto_alpha in
-  let id = List.nth dummy_environment.identities id in
-  id.implicit_contract
-let addr id =
-  let open Proto_alpha_utils.Memory_proto_alpha in
-  Protocol.Alpha_context.Contract.to_b58check @@ contract id
-
 let empty_op_list = 
   (e_typed_list [] t_operation)
 let empty_message = e_lambda "arguments"
   (Some t_unit) (Some (t_list t_operation))
   empty_op_list
+let empty_message2 = e_lambda "arguments"
+  (Some t_unit) (Some (t_list t_operation))
+ ( e_let_in ("foo",Some t_unit) (e_unit ()) empty_op_list)
 
-let param = e_constructor "Send" empty_message
+let send_param msg = e_constructor "Send" msg
+let withdraw_param = e_constructor "Withdraw" empty_message
 
-let storage threshold id_list store_list = e_ez_record [
+type st_type = {
-  ("threshold", e_nat threshold) ;
+  threshold:int ;
-  ("auth"     , e_typed_set
+  max_proposal:int ; 
-    (List.fold_left (fun acc el -> (e_address @@ addr el)::acc) [] id_list) 
+  max_msg_size:int ;
-    t_address) ;
+  id_counter_list: (int * int) list ;
-  ("message_store" , e_typed_big_map store_list t_bytes (t_set t_address))
+  msg_store_list: (expression * expression) list ;
-]
+}
+let storage {threshold ; max_proposal ; max_msg_size ; id_counter_list ; msg_store_list} =
+  let auth_set,counter_store = List.fold_left
+    (fun (auth_set,counter_st) (id,ctr) ->
+      let addr_exp = e_address @@ addr id in
+      addr_exp::auth_set , (addr_exp, e_nat ctr)::counter_st)
+    ([],[])
+    id_counter_list in
+  e_ez_record [
+    ("threshold"       , e_nat threshold                                          ) ;
+    ("max_proposal"    , e_nat max_proposal                                       ) ;
+    ("max_message_size", e_nat max_msg_size                                       ) ;
+    ("auth"            , e_typed_set auth_set t_address                           ) ;
+    ("message_store"   , e_typed_map msg_store_list t_bytes (t_set t_address) ) ;
+    ("counter_store"   , e_typed_map counter_store t_address t_nat                ) ;
+  ]
 
 (* sender not stored in the authorized set *)
 let wrong_addr () =
   let%bind program,_ = get_program () in
-  let init_storage = storage 1 [1;2] [] in
+  let init_storage = storage {
-  let amount = Memory_proto_alpha.Protocol.Alpha_context.Tez.zero in
+    threshold = 1 ; max_proposal = 1 ; max_msg_size = 1 ;
+    id_counter_list = [1,0 ; 2,0] ;
+    msg_store_list = []
+  } in
   let source = contract 3 in
-  let options = Proto_alpha_utils.Memory_proto_alpha.make_options ~amount ~source () in
+  let options = Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
   let%bind () =
     let exp_failwith = "Unauthorized address" in
     expect_string_failwith ~options program "main"
-    (e_pair param init_storage) exp_failwith in
+    (e_pair (send_param empty_message) init_storage) exp_failwith in
+  ok ()
+
+(* send a message which exceed the size limit *)
+let message_size_exceeded () =
+  let%bind program,_ = get_program () in
+  let init_storage = storage {
+    threshold = 1 ; max_proposal = 1 ; max_msg_size = 1 ;
+    id_counter_list = [1,0] ;
+    msg_store_list = []
+  } in
+  let source = contract 1 in
+  let options = Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
+  let%bind () =
+    let exp_failwith = "Message size exceed maximum limit" in
+    expect_string_failwith ~options program "main"
+    (e_pair (send_param empty_message)  init_storage) exp_failwith in
+  ok ()
+
+(* sender has already has reached maximum number of proposal *)
+let maximum_number_of_proposal () =
+  let%bind program,_ = get_program () in
+  let%bind packed_payload1 = pack_payload program (send_param empty_message) in
+  let%bind bytes1 = e_bytes_ofbytes packed_payload1 in
+  let init_storage = storage {
+    threshold = 1 ; max_proposal = 1 ; max_msg_size = 15 ;
+    id_counter_list = [1,1] ;
+    msg_store_list = [(bytes1, e_set [e_address@@ addr 1])]
+  } in
+  let source = contract 1 in
+  let options = Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
+  let%bind () =
+    let exp_failwith = "Maximum number of proposal reached" in
+    expect_string_failwith ~options program "main"
+      (e_pair (send_param empty_message2) init_storage) exp_failwith in
   ok ()
 
 (* sender message is already stored in the message store *)
-let already_accounted () =
+let send_already_accounted () =
   let%bind program,_ = get_program () in
   let%bind packed_payload = pack_payload program empty_message in
   let%bind bytes = e_bytes_ofbytes packed_payload in
-  let init_storage = storage 2 [1;2]
+  let init_storage = storage {
-    [(bytes, e_set [e_address@@ addr 1])] in
+    threshold = 2 ;  max_proposal = 1 ;  max_msg_size = 15 ;
+    id_counter_list = [1,1 ; 2,0] ;
+    msg_store_list = [(bytes, e_set [e_address@@ addr 1])]
+  } in
   let options =
-    let amount = Memory_proto_alpha.Protocol.Alpha_context.Tez.zero in
     let source = contract 1 in
-    Proto_alpha_utils.Memory_proto_alpha.make_options ~amount ~source () in
+    Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
   expect_eq ~options program "main"
-    (e_pair param init_storage) (e_pair empty_op_list init_storage)
+    (e_pair (send_param empty_message) init_storage) (e_pair empty_op_list init_storage)
 
 (* sender message isn't stored in the message store *)
-let never_accounted () =
+let send_never_accounted () =
   let%bind program,_ = get_program () in
   let%bind packed_payload = pack_payload program empty_message in
   let%bind bytes = e_bytes_ofbytes packed_payload in
-  let init_storage = storage 2 [1;2]
+  let init_storage' = {
-    [] in
+    threshold = 2 ; max_proposal = 1 ;  max_msg_size = 15 ;
-  let final_storage = storage 2 [1;2]
+    id_counter_list = [1,0 ; 2,0] ;
-    [(bytes, e_set [e_address@@ addr 1])] in
+    msg_store_list = [] 
+  } in
+  let init_storage = storage init_storage' in
+  let final_storage = storage { init_storage' with
+    id_counter_list = [1,1 ; 2,0] ;
+    msg_store_list = [(bytes, e_set [e_address@@ addr 1])] ;
+  } in
   let options =
-    let amount = Memory_proto_alpha.Protocol.Alpha_context.Tez.zero in
     let source = contract 1 in
-    Proto_alpha_utils.Memory_proto_alpha.make_options ~amount ~source () in
+    Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
+  expect_eq ~options program "main"
+    (e_pair (send_param empty_message) init_storage) (e_pair empty_op_list final_storage)
+
+(* sender withdraw message is already binded to one address in the message store *)
+let withdraw_already_accounted_one () =
+  let%bind program,_ = get_program () in
+  let%bind packed_payload = pack_payload program empty_message in
+  let%bind bytes = e_bytes_ofbytes packed_payload in
+  let param = withdraw_param in
+  let init_storage' = {
+    threshold = 2 ; max_proposal = 1 ;  max_msg_size = 1 ;
+    id_counter_list = [1,1 ; 2,0] ;
+    msg_store_list = [(bytes, e_set [e_address@@ addr 1])] ;
+  } in
+  let init_storage = storage init_storage' in
+  let final_storage = storage { init_storage' with
+    id_counter_list = [1,0 ; 2,0] ;
+    msg_store_list = [] } in
+  let options =
+    let source = contract 1 in
+    Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
   expect_eq ~options program "main"
     (e_pair param init_storage) (e_pair empty_op_list final_storage)
 
+(* sender withdraw message is already binded to two addresses in the message store *)
+let withdraw_already_accounted_two () =
+  let%bind program,_ = get_program () in
+  let%bind packed_payload = pack_payload program empty_message in
+  let%bind bytes = e_bytes_ofbytes packed_payload in
+  let param = withdraw_param in
+  let init_storage' = {
+    threshold = 2 ; max_proposal = 2 ;  max_msg_size = 1 ;
+    id_counter_list = [1,1 ; 2,1] ;
+    msg_store_list = [(bytes, e_set [e_address@@ addr 1; e_address@@ addr 2])] ;
+  } in
+  let init_storage = storage init_storage' in
+  let final_storage = storage { init_storage' with
+    id_counter_list = [1,0 ; 2,1] ;
+    msg_store_list = [(bytes, e_set [e_address@@ addr 2])] } in
+  let options =
+    let source = contract 1 in
+    Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
+  expect_eq ~options program "main"
+    (e_pair param init_storage) (e_pair empty_op_list final_storage)
+
+(* sender withdraw message was never accounted *)
+let withdraw_never_accounted () =
+  let%bind program,_ = get_program () in
+  let param = withdraw_param in
+  let init_storage = storage {
+    threshold = 2 ; max_proposal = 1 ;  max_msg_size = 1 ;
+    id_counter_list = [1,0 ; 2,0] ;
+    msg_store_list = [] ;
+  } in
+  let options =
+    let source = contract 1 in
+    Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
+  expect_eq ~options program "main"
+    (e_pair param init_storage) (e_pair empty_op_list init_storage)
+
 (* successful storing in the message store *)
 let succeeded_storing () =
   let%bind program,_ = get_program () in
   let%bind packed_payload = pack_payload program empty_message in
   let%bind bytes = e_bytes_ofbytes packed_payload in
+  let init_storage th = {
+    threshold = th ; max_proposal = 1 ;  max_msg_size = 15 ;
+    id_counter_list = [1,0 ; 2,0 ; 3,0] ;
+    msg_store_list = [(bytes, e_typed_set [] t_address)] ;
+  } in
   let options =
-    let amount = Memory_proto_alpha.Protocol.Alpha_context.Tez.zero in
     let source = contract 1 in
-    Proto_alpha_utils.Memory_proto_alpha.make_options ~amount ~source () in
+    Proto_alpha_utils.Memory_proto_alpha.make_options ~source () in
   let%bind () = expect_eq_n_trace_aux ~options [1;2] program "main"
       (fun th ->
-        let init_storage = storage th [1;2;3]
+        let init_storage = storage (init_storage th) in
-        [(bytes, e_typed_set [] t_address)] in
+        ok @@ e_pair (send_param empty_message) init_storage
-        ok @@ e_pair param init_storage
       )
       (fun th ->
-        let final_msg_store, ret = match th with
+        let final_id_counter, final_msg_store, ret = match th with
-        | 1 -> [] , empty_op_list
+          | 1 -> [1,0 ; 2,0 ; 3,0] , []                                    , empty_op_list
-        | 2 -> [(bytes, e_set [e_address@@ addr 1])] , empty_op_list
+          | 2 -> [1,1 ; 2,0 ; 3,0] , [(bytes, e_set [e_address@@ addr 1])] , empty_op_list
-        | _ -> failwith "impossible" in
+          | _ -> assert false in
-        let final_storage = storage th [1;2;3] final_msg_store in
+        let final_storage = storage { (init_storage th) with
+          msg_store_list = final_msg_store ;
+          id_counter_list = final_id_counter } in
         ok @@ e_pair ret final_storage
       ) in
   ok ()
 
 let main = test_suite "Multisig v2" [
-    test "compile"           compile_main      ;
+    test "compile"                        compile_main                   ;
-    test "wrong_addr"        wrong_addr        ;
+    test "wrong_addr"                     wrong_addr                     ;
-    test "already_accounted" already_accounted ;
+    test "message_size_exceeded"          message_size_exceeded          ;
-    test "never_accounted"   never_accounted   ;
+    test "maximum_number_of_proposal"     maximum_number_of_proposal     ;
-    test "succeeded_storing" succeeded_storing ;
+    test "send_already_accounted"         send_already_accounted         ;
+    test "send_never_accounted"           send_never_accounted           ;
+    test "succeeded_storing"              succeeded_storing              ;
+    test "withdraw_already_accounted_one" withdraw_already_accounted_one ;
+    test "withdraw_already_accounted_two" withdraw_already_accounted_two ;
   ]

src/test/test_helpers.ml

@@ -49,6 +49,27 @@ let sign_message (program:Ast_typed.program) (payload : expression) sk : string
   let signature_str = Signature.to_b58check signed_data in
   ok signature_str
 
+let contract id = 
+  let open Proto_alpha_utils.Memory_proto_alpha in
+  let id = List.nth dummy_environment.identities id in
+  id.implicit_contract
+
+let addr id =
+  let open Proto_alpha_utils.Memory_proto_alpha in
+  Protocol.Alpha_context.Contract.to_b58check @@ contract id
+
+let gen_keys = fun () ->
+  let open Tezos_crypto in
+  let (raw_pkh,raw_pk,raw_sk) = Signature.generate_key () in
+  (raw_pkh,raw_pk,raw_sk)
+
+let str_keys (raw_pkh, raw_pk, raw_sk) =
+  let open Tezos_crypto in
+  let sk_str = Signature.Secret_key.to_b58check raw_sk in
+  let pk_str = Signature.Public_key.to_b58check raw_pk in
+  let pkh_str = Signature.Public_key_hash.to_b58check raw_pkh in
+  (pkh_str,pk_str,sk_str)
+
 open Ast_simplified.Combinators
 
 let expect ?options program entry_point input expecter =