[RFC][PATCH 14/X][libsanitizer] Introduce HWASAN block-scope poisoning

[Matthew Malcomson <Matthew dot Malcomson at arm dot com>]

Here we use exactly the same mechanism as ASAN_MARK to poison/unpoison
variables on entry/exit of a block.

In order to simply use the exact same machinery we're using the same
internal functions until the SANOPT pass.  This means that all handling
of ASAN_MARK is the same.
This has the negative that the naming may be a little confusing, but a
positive that handling of the internal function doesn't have to be
duplicated for a function that behaves exactly the same but has a
different name.

gcc/ChangeLog:

2019-09-06  Matthew Malcomson  <matthew.malcomson@arm.com>

	* asan.c (asan_expand_mark_ifn): New.
	(asan_expand_poison_ifn): Add assertion.
	(hwasan_emit_prologue): Add assertion.
	(hwasan_emit_uncolour_frame): Clear hash map of block-scope
	variables.
	(hwasan_extract_tag): New.
	(hwasan_expand_mark_ifn): New.
	(hardware_memory_tagging_p): New.
	* asan.h (hwasan_extract_tag): New.
	(hwasan_expand_mark_ifn): New.
	(enum hwasan_mark_flags): New.
	(hardware_memory_tagging_p): New.
	* cfgexpand.c (expand_stack_vars): Neaten up an if clause.
	* gimple-pretty-print.c (dump_gimple_call_args): Handle
	HWASAN_MARK.
	* gimplify.c (asan_poison_variable): Set alignment for HWASAN.
	(gimplify_function_tree): Record marked variables for both ASAN
	and HWASAN.
	* internal-fn.c (expand_HWASAN_MARK): New.
	(expand_HWASAN_CHOOSE_COLOUR): Use new hwasan_extract_tag
	helper.
	* internal-fn.def (HWASAN_MARK): New.
	* sanopt.c (pass_sanopt::execute): Account for HWASAN.



###############     Attachment also inlined for ease of reply    ###############


diff --git a/gcc/asan.h b/gcc/asan.h
index e4e823080e4ca7489135ee2da9e0727de9bba8ae..6e5ba8be606e9a1eae2afe57f17ccca5562167fd 100644
--- a/gcc/asan.h
+++ b/gcc/asan.h
@@ -30,12 +30,15 @@ extern void hwasan_increment_tag ();
 extern rtx hwasan_with_tag (rtx, poly_int64);
 extern void hwasan_tag_init ();
 extern rtx hwasan_create_untagged_base (rtx);
+extern rtx hwasan_extract_tag (rtx tagged_pointer);
 extern rtx hwasan_base ();
 extern void hwasan_emit_prologue (rtx *, rtx *, poly_int64 *, uint8_t *, size_t);
 extern rtx_insn *hwasan_emit_uncolour_frame (rtx, rtx, rtx_insn *);
 extern bool hwasan_expand_check_ifn (gimple_stmt_iterator *, bool);
+extern bool hwasan_expand_mark_ifn (gimple_stmt_iterator *);
 extern bool memory_tagging_p (void);
 extern bool gate_hwasan (void);
+extern bool hardware_memory_tagging_p (void);
 extern rtx_insn *asan_emit_stack_protection (rtx, rtx, unsigned int,
 					     HOST_WIDE_INT *, tree *, int);
 extern rtx_insn *asan_emit_allocas_unpoison (rtx, rtx, rtx_insn *);
@@ -142,6 +145,13 @@ enum asan_mark_flags
 #undef DEF
 };
 
+enum hwasan_mark_flags
+{
+#define DEF(X) HWASAN_MARK_##X
+  IFN_ASAN_MARK_FLAGS
+#undef DEF
+};
+
 /* Return true if STMT is ASAN_MARK with FLAG as first argument.  */
 extern bool asan_mark_p (gimple *stmt, enum asan_mark_flags flag);
 
diff --git a/gcc/asan.c b/gcc/asan.c
index fefd28cbd136d74ad3389cf8efbf1949e3815dfd..ad3d5a6451d3ecd9ff79b768c1e9a3fb92272a7e 100644
--- a/gcc/asan.c
+++ b/gcc/asan.c
@@ -3375,6 +3375,22 @@ asan_expand_mark_ifn (gimple_stmt_iterator *iter)
   unsigned HOST_WIDE_INT size_in_bytes = tree_to_shwi (len);
   gcc_assert (size_in_bytes);
 
+  if (memory_tagging_p ())
+    {
+      /* Here we swap the ASAN_MARK for HWASAN_MARK.
+	 This is because we are using the (possibly temporary) approach to
+	 always emit ASAN_MARK in TREE until here.
+	 That approach means we don't yet have to duplicate all the special
+	 cases for ASAN_MARK and ASAN_POISON with the exact same handling but
+	 called HWASAN_MARK etc.  */
+      gimple *hw_poison_call
+	= gimple_build_call_internal (IFN_HWASAN_MARK, 3,
+				      gimple_call_arg (g, 0),
+				      base, len);
+      gsi_replace (iter, hw_poison_call, false);
+      return false;
+    }
+
   g = gimple_build_assign (make_ssa_name (pointer_sized_int_node),
 			   NOP_EXPR, base);
   gimple_set_location (g, loc);
@@ -3673,6 +3689,7 @@ asan_expand_poison_ifn (gimple_stmt_iterator *iter,
 			bool *need_commit_edge_insert,
 			hash_map<tree, tree> &shadow_vars_mapping)
 {
+  gcc_assert (! memory_tagging_p ());
   gimple *g = gsi_stmt (*iter);
   tree poisoned_var = gimple_call_lhs (g);
   if (!poisoned_var || has_zero_uses (poisoned_var))
@@ -3968,6 +3985,7 @@ hwasan_emit_prologue (rtx *bases,
 	}
       else
 	{
+	  gcc_assert (known_ge (end, start));
 	  top = end;
 	  bot = start;
 	}
@@ -4037,6 +4055,12 @@ hwasan_emit_uncolour_frame (rtx dynamic, rtx vars, rtx_insn *before)
   do_pending_stack_adjust ();
   rtx_insn *insns = get_insns ();
   end_sequence ();
+
+  /* Clear the hash_map recording which variables are handled by HWASAN_MARK.
+     The only use in HWASAN is to decide which variables need to be coloured in
+     the prologue and which don't.  */
+  delete asan_handled_variables;
+  asan_handled_variables = NULL;
   return insns;
 }
 
@@ -4052,6 +4076,19 @@ hwasan_create_untagged_base (rtx orig_base)
   return untagged_base;
 }
 
+rtx
+hwasan_extract_tag (rtx tagged_pointer)
+{
+  rtx tag = expand_simple_binop (Pmode,
+				 LSHIFTRT,
+				 tagged_pointer,
+				 HWASAN_SHIFT_RTX,
+				 NULL_RTX,
+				 /* unsignedp = */0,
+				 OPTAB_DIRECT);
+  return gen_lowpart (QImode, tag);
+}
+
 /* Needs to be GTY(()), because cgraph_build_static_cdtor may
    invoke ggc_collect.  */
 static GTY(()) tree hwasan_ctor_statements;
@@ -4166,10 +4203,25 @@ hwasan_expand_check_ifn (gimple_stmt_iterator *iter, bool)
 }
 
 bool
+hwasan_expand_mark_ifn (gimple_stmt_iterator *)
+{
+  /* HWASAN_MARK should only ever be available after the sanopt pass.
+     It might be nicer to have it everywhere in the future, so I"m leaving this
+     function and the declaration in asan.h around in case that's requested
+     upstream.  */
+  gcc_unreachable ();
+}
+
+bool
 gate_hwasan ()
 {
   return memory_tagging_p ();
 }
+bool
+hardware_memory_tagging_p ()
+{
+  return memory_tagging_p () && HARDWARE_MEMORY_TAGGING;
+}
 
 namespace {
 
diff --git a/gcc/gimple-pretty-print.c b/gcc/gimple-pretty-print.c
index 10f79ea619e3ebfdcbe9f5159e174cf2bb08b7d8..52d104857df92bc80db6deb5b18c99ee2caa4151 100644
--- a/gcc/gimple-pretty-print.c
+++ b/gcc/gimple-pretty-print.c
@@ -739,6 +739,7 @@ dump_gimple_call_args (pretty_printer *buffer, gcall *gs, dump_flags_t flags)
 	  limit = ARRAY_SIZE (reduction_args);
 	  break;
 
+	case IFN_HWASAN_MARK:
 	case IFN_ASAN_MARK:
 #define DEF(X) #X
 	  static const char *const asan_mark_args[] = {IFN_ASAN_MARK_FLAGS};
diff --git a/gcc/gimplify.c b/gcc/gimplify.c
index 5bacb255ba75901a3cd95e3b99b1f274216bf85d..d89d81ec112d918912269c90faae468d1d8aa321 100644
--- a/gcc/gimplify.c
+++ b/gcc/gimplify.c
@@ -1202,8 +1202,10 @@ asan_poison_variable (tree decl, bool poison, gimple_stmt_iterator *it,
 
   /* It's necessary to have all stack variables aligned to ASAN granularity
      bytes.  */
-  if (DECL_ALIGN_UNIT (decl) <= ASAN_SHADOW_GRANULARITY)
-    SET_DECL_ALIGN (decl, BITS_PER_UNIT * ASAN_SHADOW_GRANULARITY);
+  unsigned shadow_granularity =
+    memory_tagging_p () ? HWASAN_TAG_GRANULE_SIZE : ASAN_SHADOW_GRANULARITY;
+  if (DECL_ALIGN_UNIT (decl) <= shadow_granularity)
+    SET_DECL_ALIGN (decl, BITS_PER_UNIT * shadow_granularity);
 
   HOST_WIDE_INT flags = poison ? ASAN_MARK_POISON : ASAN_MARK_UNPOISON;
 
@@ -13816,7 +13818,7 @@ gimplify_function_tree (tree fndecl)
       && !needs_to_live_in_memory (ret))
     DECL_GIMPLE_REG_P (ret) = 1;
 
-  if (asan_sanitize_use_after_scope () && sanitize_flags_p (SANITIZE_ADDRESS))
+  if (asan_sanitize_use_after_scope ())
     asan_poisoned_variables = new hash_set<tree> ();
   bind = gimplify_body (fndecl, true);
   if (asan_poisoned_variables)
diff --git a/gcc/internal-fn.c b/gcc/internal-fn.c
index 4eec9919b520691ab3e73a2920ef8b544cf55dfe..c530fe8951c30987c874df83e74be6d058730134 100644
--- a/gcc/internal-fn.c
+++ b/gcc/internal-fn.c
@@ -471,11 +471,7 @@ expand_HWASAN_CHOOSE_COLOUR (internal_fn, gcall *gc)
   machine_mode mode = GET_MODE (target);
   gcc_assert (mode == QImode);
 
-  rtx base_tag = expand_simple_binop (Pmode, LSHIFTRT, hwasan_base (),
-				      HWASAN_SHIFT_RTX,
-				      NULL_RTX, /* unsignedp = */0,
-				      OPTAB_DIRECT);
-
+  rtx base_tag = hwasan_extract_tag (hwasan_base ());
   gcc_assert (base_tag);
   rtx tag_offset = const_int_rtx[MAX_SAVED_CONST_INT + hwasan_current_tag ()];
   rtx chosen_tag = expand_simple_binop (QImode, PLUS, base_tag, tag_offset,
@@ -498,6 +494,39 @@ expand_HWASAN_CHOOSE_COLOUR (internal_fn, gcall *gc)
 }
 
 static void
+expand_HWASAN_MARK (internal_fn, gcall *gc)
+{
+  HOST_WIDE_INT flag = tree_to_shwi (gimple_call_arg (gc, 0));
+  bool is_poison = ((asan_mark_flags)flag) == ASAN_MARK_POISON;
+
+  tree base = gimple_call_arg (gc, 1);
+  gcc_checking_assert (TREE_CODE (base) == ADDR_EXPR);
+  rtx base_rtx = expand_normal (base);
+
+  rtx tag = is_poison ? const0_rtx : hwasan_extract_tag (base_rtx);
+  rtx address = hwasan_create_untagged_base (base_rtx);
+
+  tree len = gimple_call_arg (gc, 2);
+  gcc_assert (tree_fits_shwi_p (len));
+  unsigned HOST_WIDE_INT size_in_bytes = tree_to_shwi (len);
+  uint8_t tg_mask = HWASAN_TAG_GRANULE_SIZE - 1;
+  gcc_assert (size_in_bytes);
+  size_in_bytes = (size_in_bytes + tg_mask) & ~tg_mask;
+  rtx size = gen_int_mode (size_in_bytes, Pmode);
+
+  /* TODO Other options (i.e. inline options)  */
+  rtx func = init_one_libfunc ("__hwasan_tag_memory");
+  emit_library_call (func,
+      LCT_NORMAL,
+      VOIDmode,
+      address, ptr_mode,
+      tag, QImode,
+      size, ptr_mode);
+}
+
+/* This should get expanded in the sanopt pass.  */
+
+static void
 expand_ASAN_CHECK (internal_fn, gcall *)
 {
   gcc_unreachable ();
diff --git a/gcc/internal-fn.def b/gcc/internal-fn.def
index ed0c5bc110f16b2cdbc139403dbdbd8ebe7e2823..100f6fbad0af4fefdd0408384374b8fdcde9bee4 100644
--- a/gcc/internal-fn.def
+++ b/gcc/internal-fn.def
@@ -290,6 +290,7 @@ DEF_INTERNAL_FN (ABNORMAL_DISPATCHER, ECF_NORETURN, NULL)
 DEF_INTERNAL_FN (BUILTIN_EXPECT, ECF_CONST | ECF_LEAF | ECF_NOTHROW, NULL)
 DEF_INTERNAL_FN (HWASAN_CHOOSE_COLOUR, ECF_LEAF | ECF_NOTHROW, ".")
 DEF_INTERNAL_FN (HWASAN_CHECK, ECF_TM_PURE | ECF_LEAF | ECF_NOTHROW, "..R..")
+DEF_INTERNAL_FN (HWASAN_MARK, ECF_LEAF | ECF_NOTHROW, NULL)
 DEF_INTERNAL_FN (ASAN_CHECK, ECF_TM_PURE | ECF_LEAF | ECF_NOTHROW, "..R..")
 DEF_INTERNAL_FN (ASAN_MARK, ECF_LEAF | ECF_NOTHROW, NULL)
 DEF_INTERNAL_FN (ASAN_POISON, ECF_LEAF | ECF_NOTHROW | ECF_NOVOPS, NULL)
diff --git a/gcc/sanopt.c b/gcc/sanopt.c
index 31270153f3cf56bfbad593830de1b9334e7f65d1..29b12a43029cdcce5756354ac12d6d9e963e9ac8 100644
--- a/gcc/sanopt.c
+++ b/gcc/sanopt.c
@@ -1258,6 +1258,12 @@ sanitize_rewrite_addressable_params (function *fun)
 unsigned int
 pass_sanopt::execute (function *fun)
 {
+  /*
+     n.b. ASAN_MARK is used for both HWASAN and ASAN.
+     asan_num_accesses is used to count either HWASAN_CHECK or ASAN_CHECK
+     stuff.  This is fine because you can only have one of these active at a
+     time.
+ */
   basic_block bb;
   int asan_num_accesses = 0;
   bool contains_asan_mark = false;
@@ -1345,6 +1351,9 @@ pass_sanopt::execute (function *fun)
 						    &need_commit_edge_insert,
 						    shadow_vars_mapping);
 		  break;
+		case IFN_HWASAN_MARK:
+		  no_next = hwasan_expand_mark_ifn (&gsi);
+		  break;
 		default:
 		  break;
 		}

Attachment: hwasan-implementation13.patch
Description: hwasan-implementation13.patch