[MSan] run materializeChecks() before materializeStores()

[glider]

Summary

When pointer checking is enabled, it's important that every pointer is checked before its value is used.
For stores MSan used to generate code that calculates shadow/origin addresses from a pointer before checking it.
For userspace this isn't a problem, because the shadow calculation code is quite simple and compiler is able to move it after the check on -O2.
But for KMSAN getShadowOriginPtr() creates a runtime call, so we want the check to be performed strictly before that call.

Swapping materializeChecks() and materializeStores() resolves the issue: both functions insert code before the given IR location, so the new insertion order guarantees that the code calculating shadow address is between the address check and the memory access.

[glider]

Uploaded the full context

[eugenis]

Could you simply pass the pointer shadow value as an argument to the runtime call in the kernel, and keep userspace instrumentation as it is?

lib/Transforms/Instrumentation/MemorySanitizer.cpp

1139

It's better to take a step back, add instrumentation, then take a step forward.
Must handle the start of the basic block separately.

[glider]

Reverted changes to getShadowOriginPtrUserspace()
Dima noticed that the problem we're trying to solve is essentially inserting two code snippets (one for shadow calculation and one for address check) right before the store to a given address so that the second snippet goes before the first one.
Therefore it's natural to insert the second snippet, and then the first one, i.e. to swap the order of calling materializeStores() and materializeChecks().

[glider]

Added a comment

[glider]

In D49473#1166913, @eugenis wrote:

Could you simply pass the pointer shadow value as an argument to the runtime call in the kernel, and keep userspace instrumentation as it is?

This'll cost us two registers (for the pointer shadow value and an origin) and is also conceptually incorrect, as calculation of the shadow for a given address should be considered a use of that address and shouldn't precede its checks for validity.

[glider]

I've also noticed that currently MSan instruments loads by placing the checks _after_ them, which is also incorrect.
E.g. for the following function:

int foo(int *v) {
  return *v;
}

the following IR is generated by the trunk compiler with -O2:

define dso_local i32 @foo(i32* nocapture readonly %v) local_unnamed_addr #0 {
entry:
  %0 = load i64, i64* getelementptr inbounds ([100 x i64], [100 x i64]* @__msan_param_tls, i64 0, i64 0), align 8
  %1 = load i32, i32* %v, align 4, !tbaa !2
  %2 = ptrtoint i32* %v to i64
  %3 = xor i64 %2, 87960930222080
  %4 = inttoptr i64 %3 to i32*
  %_msld = load i32, i32* %4, align 4
  store i32 %_msld, i32* bitcast ([100 x i64]* @__msan_retval_tls to i32*), align 8
  %_mscmp = icmp eq i64 %0, 0
  br i1 %_mscmp, label %6, label %5, !prof !6

; <label>:5:                                      ; preds = %entry
  call void @__msan_warning_noreturn() #1
  call void asm sideeffect "", ""() #1
  unreachable

; <label>:6:                                      ; preds = %entry
  ret i32 %1
}

[glider]

Sorry, please disregard the previous comment. It's my patch that breaks this order.

[eugenis]

lib/Transforms/Instrumentation/MemorySanitizer.cpp

1490

So all you really need is this chunk, and reordering of materializeStores and materializeChecks? Could you revert the rest of the change then?

I prefer the current, more explicit, way of inserting address checks.

[glider]

Yes, looks like this chunk is enough.
I'll finish comparing the results and update the patch tomorrow.

[glider]

Updated the patch.

[glider]

Fixed check_access_address.ll

[eugenis]

LGTM, thanks!

[glider]

Landed r337571, thanks!

lib/Transforms/Instrumentation/MemorySanitizer.cpp

for (StoreInst *SI : StoreList) {
       unsigned Alignment = SI->getAlignment();
       unsigned OriginAlignment = std::max(kMinOriginAlignment, Alignment);
       std::tie(ShadowPtr, OriginPtr) =
           getShadowOriginPtr(Addr, IRB, ShadowTy, Alignment, /*isStore*/ true);
 
       StoreInst *NewSI = IRB.CreateAlignedStore(Shadow, ShadowPtr, Alignment);
       LLVM_DEBUG(dbgs() << "  STORE: " << *NewSI << "\n");
 
-      if (ClCheckAccessAddress)
-        insertShadowCheck(Addr, NewSI);
-
       if (SI->isAtomic())
         SI->setOrdering(addReleaseOrdering(SI->getOrdering()));
 
       if (MS.TrackOrigins && !SI->isAtomic())
         storeOrigin(IRB, Addr, Shadow, getOrigin(Val), OriginPtr,
                     OriginAlignment, InstrumentWithCalls);
     }
   }
bool runOnFunction() {
     }
 
     VAHelper->finalizeInstrumentation();
 
     bool InstrumentWithCalls = ClInstrumentationWithCallThreshold >= 0 &&
                                InstrumentationList.size() + StoreList.size() >
                                    (unsigned)ClInstrumentationWithCallThreshold;
 
-    // Delayed instrumentation of StoreInst.
-    // This may add new checks to be inserted later.
-    materializeStores(InstrumentWithCalls);
-
     // Insert shadow value checks.
     materializeChecks(InstrumentWithCalls);
 
+    // Delayed instrumentation of StoreInst.
+    // This may not add new address checks.
+    materializeStores(InstrumentWithCalls);
+
     return true;
   }
 
   /// Compute the shadow type that corresponds to a given Value.
   Type *getShadowTy(Value *V) {
     return getShadowTy(V->getType());
   }
 
if (MS.TrackOrigins) {
         uint64_t Mask = kMinOriginAlignment - 1;
         OriginLong =
             IRB.CreateAnd(OriginLong, ConstantInt::get(MS.IntptrTy, ~Mask));
       }
       OriginPtr =
           IRB.CreateIntToPtr(OriginLong, PointerType::get(IRB.getInt32Ty(), 0));
     }
     return std::make_pair(ShadowPtr, OriginPtr);
   }
 
   std::pair<Value *, Value *> getShadowOriginPtr(Value *Addr, IRBuilder<> &IRB,
                                                  Type *ShadowTy,
                                                  unsigned Alignment,
                                                  bool isStore) {
     std::pair<Value *, Value *> ret =
         getShadowOriginPtrUserspace(Addr, IRB, ShadowTy, Alignment);
     return ret;
#endif
   }
 
   /// Instrument StoreInst
   ///
   /// Stores the corresponding shadow and (optionally) origin.
   /// Optionally, checks that the store address is fully defined.
   void visitStoreInst(StoreInst &I) {
     StoreList.push_back(&I);
+    if (ClCheckAccessAddress)
+      insertShadowCheck(I.getPointerOperand(), &I);
   }
 
   void handleCASOrRMW(Instruction &I) {
     assert(isa<AtomicRMWInst>(I) || isa<AtomicCmpXchgInst>(I));
 
     IRBuilder<> IRB(&I);
     Value *Addr = I.getOperand(0);
     Value *ShadowPtr = getShadowOriginPtr(Addr, IRB, I.getType(),

test/Instrumentation/MemorySanitizer/check_access_address.ll

 define void @Store(i32* nocapture %p, i32 %x) nounwind uwtable sanitize_memory {
 entry:
   store i32 %x, i32* %p, align 4
   ret void
 }
 
 ; CHECK-LABEL: @Store
 ; CHECK: load {{.*}} @__msan_param_tls
+; Shadow calculations must happen after the check.
+; CHECK-NOT: xor
 ; CHECK: icmp
 ; CHECK: br i1
 ; CHECK: <label>
 ; CHECK: call void @__msan_warning_noreturn
 ; CHECK: <label>
+; CHECK: xor
 ; CHECK: store
 ; CHECK: store i32 %x
 ; CHECK: ret void