Module Canary

Require Import Coqlib Wfsimpl Maps Errors Integers Values.
Require Import AST Linking Globalenvs.
Require Import Op Registers RTL CounterMeasures Memdata.
Require Compopts.

Open Scope Z.

Definition canary_chunk :=
  if Archi.canary64 then Mint64 else Mint32.

Definition canary_size :=
  if Archi.canary64 then 8 else 4.

Definition padding ofs al :=
  let md := Zmod ofs al in
  if zeq md 0
  then 0
  else al - md.

Definition canary_offset (f : function) :=
  f.(fn_stacksize) + (padding f.(fn_stacksize) (align_chunk canary_chunk)).

Definition canary_needed (f : function) : bool :=
  match getcanary with
  | Some _ =>
      zle 0 f.(fn_stacksize) &&
      zle (canary_offset f) Ptrofs.max_unsigned &&
      Compopts.stack_protector tt &&
      (Compopts.stack_protector_all tt ||
         zlt 0 f.(fn_stacksize))
  | None => false
  end.

Definition extra_canary_size (f : function) :=
  if canary_needed f
  then (padding f.(fn_stacksize) (align_chunk canary_chunk)) + canary_size
  else 0.

Definition entry_sequence offset tmpreg : seqinstr OneExit :=
  match getcanary with
  | None => Seinstr Inop
  | Some canary_op =>
        Ssinstr (Iop canary_op nil tmpreg);;
        Ssinstr (Istore canary_chunk (Ainstack (Ptrofs.repr offset)) nil tmpreg);;
        Seinstr (Iop clearcanary nil tmpreg)
  end.

Definition check_sequence offset tmpreg insn :=
  match getcanary with
  | None => Sereturn insn
  | Some canary_op =>
      let tmpreg2 := Pos.succ tmpreg in
      Ssinstr (Iload TRAP canary_chunk (Ainstack (Ptrofs.repr offset)) nil tmpreg);;
      Ssinstr (Iop canary_op nil tmpreg2);;
      Ssmerge1 (negate_condition canary_cmp) (tmpreg :: tmpreg2 :: nil) Secatch (Some false);;
      Sereturn insn
  end.

Definition transf_return offset tmpreg (_ : node) r :=
  check_sequence offset tmpreg (Ireturn r).
Definition transf_tailcall offset tmpreg (_ : node) sig ros rs :=
  check_sequence offset tmpreg (Itailcall sig ros rs).

Definition transf_function (f : function) : function :=
  let tmpreg := Pos.succ (max_reg_function f) in
  let next_pc := Pos.succ (max_pc_function f) in
  let offset := canary_offset f in
  let res :=
    if canary_needed f then
      transf_function
        (entry_sequence offset tmpreg)
        transf_nop_default transf_op_default
        transf_load_default transf_store_default
        transf_call_default (transf_tailcall offset tmpreg) transf_builtin_default
        transf_cond_default transf_jumptable_default
        (transf_return offset tmpreg)
        f
    else (fn_code f, fn_entrypoint f)
  in mkfunction (fn_sig f) (fn_params f) (fn_stacksize f + extra_canary_size f) (fst res) (snd res) (fn_attr f).

Definition transf_fundef (fd: fundef) : fundef :=
  transf_fundef transf_function fd.

Definition transf_program (p: program) : program :=
  transform_program transf_fundef p.