Stack zeroization on main with proofs

CHANGELOG.md

@@ -51,6 +51,16 @@
   `SMULTT`, `SMLABB`, `SMLABT`, `SMLATB`, `SMLATT`, `SMULWB`, and `SMULWT`
   ([PR #644](https://github.com/jasmin-lang/jasmin/pull/644)).
 
+- Support stack zeroization.
+  The compiler can introduce code that zeroizes the stack at the end of export
+  functions. The user can enable this feature using either annotations
+  (`stackzero` and `stackzerosize`), or compiler flags (`-stack-zero` and
+  `-stack-zero-size`). Three strategies are currently supported: `unrolled`
+  (the code is a sequence of writes as long as needed), `loop` (the code is
+  a loop) and `loopSCT` (same as `loop` but with a `LFENCE` at the end to defend
+  against Spectre).
+  ([PR #631](https://github.com/jasmin-lang/jasmin/pull/631)).
+
 ## Bug fixes
 
 - Type-checking rejects wrongly casted primitive operators

compiler/Makefile

@@ -5,7 +5,19 @@ JSJOBS    ?= 2
 CHECKPY   ?=
 CHECK     := $(CHECKPY) scripts/runtest --jobs="$(JSJOBS)"
 CHECK     += config/tests.config
-CHECKCATS ?= safety CCT x86-64-ATT x86-64-Intel x86-64-print x86-64-nolea arm-m4 arm-m4-print
+CHECKCATS ?= \
+	safety \
+	CCT \
+	x86-64-ATT \
+	x86-64-Intel \
+	x86-64-print \
+	x86-64-nolea \
+	arm-m4 \
+	arm-m4-print \
+	x86-64-stack-zero-loop \
+	x86-64-stack-zero-unrolled \
+	arm-m4-stack-zero-loop \
+	arm-m4-stack-zero-unrolled
 
 # --------------------------------------------------------------------
 DESTDIR  ?=

compiler/config/tests.config

@@ -53,3 +53,29 @@ bin = ./scripts/extract-and-check
 args = arm
 okdirs = examples/**/arm-m4 tests/success/**/arm-m4 tests/success/**/common
 exclude = !tests/success/noextract
+
+[test-x86-64-stack-zero-loop]
+bin     = ./scripts/check-x86-64
+args    = -stack-zero=loop
+okdirs  = examples/**/x86-64 tests/success/**/x86-64
+kodirs  = tests/fail/**/x86-64
+exclude = !tests/fail/warning
+
+[test-x86-64-stack-zero-unrolled]
+bin     = ./scripts/check-x86-64
+args    = -stack-zero=unrolled
+okdirs  = examples/**/x86-64 tests/success/**/x86-64
+kodirs  = tests/fail/**/x86-64
+exclude = !tests/fail/warning
+
+[test-arm-m4-stack-zero-loop]
+bin     = ./scripts/check-arm-m4
+args    = -stack-zero=loop
+okdirs  = examples/**/arm-m4 tests/success/**/arm-m4
+kodirs  = tests/fail/**/arm-m4
+
+[test-arm-m4-stack-zero-unrolled]
+bin     = ./scripts/check-arm-m4
+args    = -stack-zero=unrolled
+okdirs  = examples/**/arm-m4 tests/success/**/arm-m4
+kodirs  = tests/fail/**/arm-m4

compiler/src/annot.ml

@@ -88,13 +88,13 @@ let pos_int dfl ((id, _) as arg) =
 
 let string_of_ws ws = Prog.string_of_ws (Printer.ws_of_ws ws)
 
+let ws_strings =
+  List.map
+    (fun ws -> (string_of_ws ws, ws))
+    [ `W8; `W16; `W32; `W64; `W128; `W256 ]
+
 let ws_of_string =
-  let l =
-    List.map
-      (fun ws -> (string_of_ws ws, ws))
-      [ `W8; `W16; `W32; `W64; `W128; `W256 ]
-  in
-  fun s -> List.assoc s l
+  fun s -> List.assoc s ws_strings
 
 let wsize dfl arg =
   let error loc nid =

compiler/src/annot.mli

@@ -52,6 +52,7 @@ val pos_int :
   string Location.located * Annotations.simple_attribute Location.located option ->
   Z.t
 
+val ws_strings : (string * Annotations.wsize) list
 val ws_of_string : string -> Annotations.wsize
 
 val wsize :

compiler/src/annotations.ml

@@ -29,6 +29,7 @@ type f_annot = {
     stack_size            : Z.t option;
     stack_align           : wsize option;
     max_call_depth        : Z.t option;
+    stack_zero_strategy   : (Stack_zero_strategy.stack_zero_strategy * wsize option) option;
     f_user_annot          : annotations;
 }
 
@@ -38,6 +39,7 @@ let f_annot_empty = {
     stack_size            = None;
     stack_align           = None;
     max_call_depth        = None;
+    stack_zero_strategy   = None;
     f_user_annot          = [];
   }
 

compiler/src/compile.ml

@@ -210,6 +210,33 @@ let compile (type reg regx xreg rflag cond asm_op extra_op)
       try Hf.find ttbl fn with Not_found -> assert false
   in
 
+  let tbl_annot =
+    let tbl = Hf.create 17 in
+    let add (fn, cfd) =
+      let fd = fdef_of_cufdef fn cfd in
+      Hf.add tbl fn fd.f_annot
+    in
+    List.iter add cprog.Expr.p_funcs;
+    tbl
+  in
+
+  let get_annot fn =
+    try Hf.find tbl_annot fn
+    with Not_found ->
+           hierror
+             ~loc:Lnone
+             ~funname:fn.fn_name
+             ~kind:"compiler error"
+             ~internal:true
+             "invalid annotation table."
+  in
+
+  let szs_of_fn fn =
+    match (get_annot fn).stack_zero_strategy with
+    | Some (s, ows) -> Some (s, Option.map Pretyping.tt_ws ows)
+    | None -> None
+  in
+
   let cparams =
     {
       Compiler.rename_fd;
@@ -245,6 +272,7 @@ let compile (type reg regx xreg rflag cond asm_op extra_op)
       Compiler.fresh_id;
       Compiler.fresh_var_ident = Conv.fresh_var_ident;
       Compiler.slh_info;
+      Compiler.stack_zero_info = szs_of_fn;
     }
   in
 

compiler/src/fInfo.ml

@@ -5,5 +5,9 @@ type call_conv =
   | Subroutine of subroutine_info (** internal function that should not be inlined *)
   | Internal  (** internal function that should be inlined *)
 
+let is_subroutine = function
+  | Subroutine _ -> true
+  | _            -> false
+
 type t =
   Location.t * Annotations.f_annot * call_conv * Annotations.annotations list

compiler/src/glob_options.ml

@@ -36,6 +36,20 @@ let introduce_array_copy = ref true
 let print_dependencies = ref false 
 let lazy_regalloc = ref false
 
+let stack_zero_strategy = ref None
+let stack_zero_strategies =
+  let open Stack_zero_strategy in
+  let assoc = function
+    | SZSloop -> "loop"
+    | SZSloopSCT -> "loopSCT"
+    | SZSunrolled -> "unrolled"
+  in
+  List.map (fun s -> (assoc s, s)) stack_zero_strategy_list
+let set_stack_zero_strategy s =
+  stack_zero_strategy := Some (List.assoc s stack_zero_strategies)
+let stack_zero_size = ref None
+let set_stack_zero_size s = stack_zero_size := Some (Annot.ws_of_string s)
+
 type architecture =
   | X86_64
   | ARM_M4
@@ -156,6 +170,7 @@ let print_strings = function
   | Compiler.RegAllocation               -> "ralloc"   , "register allocation"
   | Compiler.DeadCode_RegAllocation      -> "rallocd"  , "dead code after register allocation"
   | Compiler.Linearization               -> "linear"   , "linearization"
+  | Compiler.StackZeroization            -> "stackzero", "stack zeroization"
   | Compiler.Tunneling                   -> "tunnel"   , "tunneling"
   | Compiler.Assembly                    -> "asm"      , "generation of assembly"
 
@@ -230,6 +245,12 @@ let options = [
     "-ATT", Arg.Unit (set_syntax `ATT), "use AT&T syntax (default is AT&T)"; 
     "-call-conv", Arg.Symbol (["windows"; "linux"], set_cc), ": select calling convention (default depend on host architecture)";
     "-arch", Arg.Symbol (["x86-64"; "arm-m4"], set_target_arch), ": select target arch (default is x86-64)";
+    "-stack-zero",
+      Arg.Symbol (List.map fst stack_zero_strategies, set_stack_zero_strategy),
+      ": select stack zeroization strategy for export functions";
+    "-stack-zero-size",
+      Arg.Symbol (List.map fst Annot.ws_strings, set_stack_zero_size),
+      ": select stack zeroization size for export functions";
   ] @  List.map print_option Compiler.compiler_step_list @ List.map stop_after_option Compiler.compiler_step_list
 
 let usage_msg = "Usage : jasminc [option] filename"

compiler/src/pretyping.ml

@@ -1831,16 +1831,63 @@ let tt_call_conv loc params returns cc =
 
 (* -------------------------------------------------------------------- *)
 
-let process_f_annot annot =
+let process_f_annot loc funname f_cc annot =
   let open A in
 
   let mk_ra = Annot.filter_string_list None ["stack", OnStack; "reg", OnReg] in
-
-  { retaddr_kind          = Annot.ensure_uniq1 "returnaddress"  mk_ra                annot;
+
+  let retaddr_kind =
+    let kind = Annot.ensure_uniq1 "returnaddress" mk_ra annot in
+    if kind <> None && not (FInfo.is_subroutine f_cc) then
+      hierror
+        ~loc:(Lone loc)
+        ~funname
+        ~kind:"unexpected annotation"
+        "returnaddress only applies to subroutines";
+    kind
+  in
+
+  let stack_zero_strategy =
+
+    let strategy =
+      let mk_szs = Annot.filter_string_list None Glob_options.stack_zero_strategies in
+      let strategy = Annot.ensure_uniq1 "stackzero" mk_szs annot in
+      if strategy <> None && f_cc <> Export then
+        hierror
+          ~loc:(Lone loc)
+          ~funname
+          ~kind:"unexpected annotation"
+          "stackzero only applies to export functions";
+      if Option.is_none strategy then
+        !Glob_options.stack_zero_strategy
+      else
+        strategy
+    in
+
+    let size =
+      let size = Annot.ensure_uniq1 "stackzerosize" (Annot.wsize None) annot in
+      if Option.is_none size then
+        !Glob_options.stack_zero_size
+      else
+        size
+    in
+
+    match strategy, size with
+    | None, None -> None
+    | None, Some _ ->
+        warning Always
+          (L.i_loc0 loc)
+          "\"stackzerosize\" is ignored, since you did not specify a strategy with attribute \"stackzero\"";
+        None
+    | Some szs, _ -> Some (szs, size)
+  in
+
+  { retaddr_kind;
     stack_allocation_size = Annot.ensure_uniq1 "stackallocsize" (Annot.pos_int None) annot;
     stack_size            = Annot.ensure_uniq1 "stacksize"      (Annot.pos_int None) annot;
     stack_align           = Annot.ensure_uniq1 "stackalign"     (Annot.wsize None)   annot;
     max_call_depth        = Annot.ensure_uniq1 "calldepth"      (Annot.pos_int None) annot;
+    stack_zero_strategy;
     f_user_annot          = annot;
   }
 
@@ -1911,11 +1958,12 @@ let tt_fundef arch_info (env0 : 'asm Env.env) loc (pf : S.pfundef) : 'asm Env.en
   let body, xret = tt_funbody arch_info envb pf.pdf_body in
   let f_cc = tt_call_conv loc args xret pf.pdf_cc in
   let args = List.map L.unloc args in
+  let name = L.unloc pf.pdf_name in
   let fdef =
     { P.f_loc   = loc;
-      P.f_annot = process_f_annot pf.pdf_annot;
+      P.f_annot = process_f_annot loc name f_cc pf.pdf_annot;
       P.f_cc    = f_cc;
-      P.f_name  = P.F.mk (L.unloc pf.pdf_name);
+      P.f_name  = P.F.mk name;
       P.f_tyin  = List.map (fun { P.v_ty } -> v_ty) args;
       P.f_args  = args;
       P.f_body  = body;

compiler/src/printLinear.ml

@@ -65,7 +65,7 @@ let pp_stackframe fmt (sz, ws) =
 
 let pp_meta fmt fd =
   F.fprintf fmt "(* %a *)"
-    pp_stackframe (fd.lfd_total_stack, fd.lfd_align)
+    pp_stackframe (fd.lfd_stk_max, fd.lfd_align)
 
 let pp_return is_export fmt =
   function

compiler/src/printer.ml

@@ -419,6 +419,7 @@ let pp_err ~debug fmt (pp_e : Compiler_util.pp_error) =
   let rec pp_err fmt pp_e =
     match pp_e with
     | Compiler_util.PPEstring s -> Format.fprintf fmt "%a" pp_string0 s
+    | Compiler_util.PPEz z -> Format.fprintf fmt "%a" Z.pp_print (Conv.z_of_cz z)
     | Compiler_util.PPEvar v -> Format.fprintf fmt "%a" pp_var v
     | Compiler_util.PPEvarinfo loc ->
       Format.fprintf fmt "%a" L.pp_loc loc

compiler/src/stackAlloc.ml

@@ -57,12 +57,21 @@ let pp_return fmt n =
 
 let pp_sao fmt sao =
   let open Stack_alloc in
-  Format.fprintf fmt "alignment = %s; size = %a; ioff = %a; extra size = %a; max size = %a@;max call depth = %a@;params =@;<2 2>@[<v>%a@]@;return = @[<hov>%a@]@;slots =@;<2 2>@[<v>%a@]@;alloc= @;<2 2>@[<v>%a@]@;saved register = @[<hov>%a@]@;saved stack = %a@;return address = %a"
+  let max_size = Conv.z_of_cz sao.sao_max_size in
+  let total_size =
+    (* if the function is export, we must take into account the alignment
+       of the stack *)
+    match sao.sao_return_address with
+    | RAnone -> Z.add max_size (Z.of_int (size_of_ws sao.sao_align - 1))
+    | _ -> max_size
+  in
+  Format.fprintf fmt "alignment = %s@;size = %a; ioff = %a; extra size = %a@;max size = %a; total size = %a@;max call depth = %a@;params =@;<2 2>@[<v>%a@]@;return = @[<hov>%a@]@;slots =@;<2 2>@[<v>%a@]@;alloc= @;<2 2>@[<v>%a@]@;saved register = @[<hov>%a@]@;saved stack = %a@;return address = %a"
     (string_of_ws sao.sao_align)
     Z.pp_print (Conv.z_of_cz sao.sao_size)
     Z.pp_print (Conv.z_of_cz sao.sao_ioff)
     Z.pp_print (Conv.z_of_cz sao.sao_extra_size)
-    Z.pp_print (Conv.z_of_cz sao.sao_max_size)
+    Z.pp_print max_size
+    Z.pp_print total_size
     Z.pp_print (Conv.z_of_cz sao.sao_max_call_depth)
     (pp_list "@;" pp_param_info) sao.sao_params
     (pp_list "@;" pp_return) sao.sao_return
@@ -247,17 +256,62 @@ let memory_analysis pp_err ~debug up =
           let max_call_depth = Z.max max_call_depth fn_max_call_depth in
           align, max_stk, max_call_depth
         ) sao.sao_calls (align, Z.zero, Z.zero) in
+    (* if we zeroize the stack, we may have to increase the alignment *)
+    let align =
+      match fd.f_cc, fd.f_annot.stack_zero_strategy with
+      | Export, Some (_, Some ws) ->
+          if Z.equal max_stk Z.zero
+            && Z.equal (Conv.z_of_cz csao.Stack_alloc.sao_size) Z.zero
+            && extra_size = 0
+          then
+            (* no stack to clear, we don't change the alignment *)
+            align
+          else
+            let ws = Pretyping.tt_ws ws in
+            if wsize_lt align ws then ws else align
+      | _, _ -> align
+    in
+    (* if we zeroize the stack, we ensure that the stack size of the export
+       function is a multiple of the alignment (this is what we systematically
+       do for subroutines). The difference is that, here, this is reflected
+       by increasing extra_size. *)
+    let extra_size =
+      let stk_size = 
+        Z.add (Conv.z_of_cz csao.Stack_alloc.sao_size)
+                   (Z.of_int extra_size) in
+      match fd.f_annot.stack_zero_strategy with
+      | Some _ ->
+          let round =
+              Conv.z_of_cz (Memory_model.round_ws align (Conv.cz_of_z stk_size))
+          in
+          Z.to_int (Z.sub round (Conv.z_of_cz csao.Stack_alloc.sao_size))
+      | None -> extra_size
+    in
+    (* if we zeroize the stack, we ensure that the max size is a multiple of the
+       size of the clear step. We use [fd.f_annot.stack_zero_strategy] and not
+       [align], this is on purpose! We know that the first one divides the
+       second one. *)
     let max_size = 
       let stk_size = 
         Z.add (Conv.z_of_cz csao.Stack_alloc.sao_size)
                    (Z.of_int extra_size) in
       let stk_size = 
         match fd.f_cc with
-        | Export       -> Z.add stk_size (Z.of_int (size_of_ws align - 1))
-        | Subroutine _ -> 
+        | Export -> stk_size
+        | Subroutine _ ->
           Conv.z_of_cz (Memory_model.round_ws align (Conv.cz_of_z stk_size))
         | Internal -> assert false in
-      Z.add max_stk stk_size in
+      let max_size = Z.add max_stk stk_size in
+      match fd.f_cc, fd.f_annot.stack_zero_strategy with
+      | Export, Some (_, ows) ->
+          let ws =
+            match ows with
+            | Some ws -> Pretyping.tt_ws ws
+            | None -> align (* the default clear step is the alignment *)
+          in
+          Conv.z_of_cz (Memory_model.round_ws ws (Conv.cz_of_z max_size))
+      | _, _ -> max_size
+    in
     let max_call_depth = Z.succ max_call_depth in
     let saved_stack = 
       if has_stack then