MSan: handle callbr instructions

[glider]

Summary

Handling callbr is very similar to handling an inline assembly call:
MSan must checks the instruction's inputs.
callbr doesn't (yet) have outputs, so there's nothing to unpoison,
and conservative assembly handling doesn't apply either.

Fixes PR42479.

[eugenis]

LGTM w/ nit

llvm/test/Instrumentation/MemorySanitizer/msan_asm_conservative.ll

293

This feels too vague. Could you add something to check that the warning is for a value coming from a function argument?

[glider]

Fixed the test

llvm/lib/Transforms/Instrumentation/MemorySanitizer.cpp

default:
         visitInstruction(I);
       break;
     }
   }
 
   void visitCallSite(CallSite CS) {
     Instruction &I = *CS.getInstruction();
     assert(!I.getMetadata("nosanitize"));
-    assert((CS.isCall() || CS.isInvoke()) && "Unknown type of CallSite");
-    if (CS.isCall()) {
-      CallInst *Call = cast<CallInst>(&I);
-
-      // For inline asm, do the usual thing: check argument shadow and mark all
-      // outputs as clean. Note that any side effects of the inline asm that are
-      // not immediately visible in its constraints are not handled.
-      if (Call->isInlineAsm()) {
-        if (ClHandleAsmConservative && MS.CompileKernel)
-          visitAsmInstruction(I);
-        else
-          visitInstruction(I);
-        return;
-      }
-
+    assert((CS.isCall() || CS.isInvoke() || CS.isCallBr()) &&
+           "Unknown type of CallSite");
+    if (CS.isCallBr() || (CS.isCall() && cast<CallInst>(&I)->isInlineAsm())) {
+      // For inline asm (either a call to asm function, or callbr instruction),
+      // do the usual thing: check argument shadow and mark all outputs as
+      // clean. Note that any side effects of the inline asm that are not
+      // immediately visible in its constraints are not handled.
+    if (CS.isCall()) {
+      CallInst *Call = cast<CallInst>(&I);
       assert(!isa<IntrinsicInst>(&I) && "intrinsics are handled elsewhere");
 
       // We are going to insert code that relies on the fact that the callee
       // will become a non-readonly function after it is instrumented by us. To
       // prevent this code from being optimized out, mark that function
       // non-readonly in advance.
       if (Function *Func = Call->getCalledFunction()) {
         // Clear out readonly/readnone attributes.
if (!ElType->isSized())
       return;
     int Size = DL.getTypeStoreSize(ElType);
     Value *Ptr = IRB.CreatePointerCast(Operand, IRB.getInt8PtrTy());
     Value *SizeVal = ConstantInt::get(MS.IntptrTy, Size);
     IRB.CreateCall(MS.MsanInstrumentAsmStoreFn, {Ptr, SizeVal});
   }
 
   /// Get the number of output arguments returned by pointers.
-  int getNumOutputArgs(InlineAsm *IA, CallInst *CI) {
+  int getNumOutputArgs(InlineAsm *IA, CallBase *CB) {
     int NumRetOutputs = 0;
     int NumOutputs = 0;
-    Type *RetTy = dyn_cast<Value>(CI)->getType();
+    Type *RetTy = dyn_cast<Value>(CB)->getType();
     if (!RetTy->isVoidTy()) {
       // Register outputs are returned via the CallInst return value.
       StructType *ST = dyn_cast_or_null<StructType>(RetTy);
       if (ST)
         NumRetOutputs = ST->getNumElements();
       else
         NumRetOutputs = 1;
     }
void visitAsmInstruction(Instruction &I) {
     //  - nO other outputs ("=m" and others) are returned by pointer as first
     // nO operands of the CallInst;
     //  - nI inputs ("r", "m" and others) are passed to CallInst as the
     // remaining nI operands.
     // The total number of asm() arguments in the source is nR+nO+nI, and the
     // corresponding CallInst has nO+nI+1 operands (the last operand is the
     // function to be called).
     const DataLayout &DL = F.getParent()->getDataLayout();
-    CallInst *CI = dyn_cast<CallInst>(&I);
+    CallBase *CB = dyn_cast<CallBase>(&I);
     IRBuilder<> IRB(&I);
-    InlineAsm *IA = cast<InlineAsm>(CI->getCalledValue());
-    int OutputArgs = getNumOutputArgs(IA, CI);
+    InlineAsm *IA = cast<InlineAsm>(CB->getCalledValue());
+    int OutputArgs = getNumOutputArgs(IA, CB);
     // The last operand of a CallInst is the function itself.
-    int NumOperands = CI->getNumOperands() - 1;
+    int NumOperands = CB->getNumOperands() - 1;
 
     // Check input arguments. Doing so before unpoisoning output arguments, so
     // that we won't overwrite uninit values before checking them.
     for (int i = OutputArgs; i < NumOperands; i++) {
-      Value *Operand = CI->getOperand(i);
+      Value *Operand = CB->getOperand(i);
       instrumentAsmArgument(Operand, I, IRB, DL, /*isOutput*/ false);
     }
     // Unpoison output arguments. This must happen before the actual InlineAsm
     // call, so that the shadow for memory published in the asm() statement
     // remains valid.
     for (int i = 0; i < OutputArgs; i++) {
-      Value *Operand = CI->getOperand(i);
+      Value *Operand = CB->getOperand(i);
       instrumentAsmArgument(Operand, I, IRB, DL, /*isOutput*/ true);
     }
 
     setShadow(&I, getCleanShadow(&I));
     setOrigin(&I, getCleanOrigin());
   }
 
   void visitInstruction(Instruction &I) {

llvm/test/Instrumentation/MemorySanitizer/msan_asm_conservative.ll

entry:
   ret void
 }
 
 ; CHECK-LABEL: @f_3i_3o_complex_mem
 ; CHECK-CONS: call void @__msan_instrument_asm_store({{.*}}@pair2{{.*}}, i64 8)
 ; CHECK-CONS: call void @__msan_instrument_asm_store({{.*}}@c2{{.*}}, i64 1)
 ; CHECK-CONS: call void @__msan_instrument_asm_store({{.*}}@memcpy_d1{{.*}}, i64 8)
 ; CHECK: call void asm "", "=*m,=*m,=*m,*m,*m,*m,~{dirflag},~{fpsr},~{flags}"(%struct.pair* @pair2, i8* @c2, i8* (i8*, i8*, i32)** @memcpy_d1, %struct.pair* @pair1, i8* @c1, i8* (i8*, i8*, i32)** @memcpy_s1)
+
+
+; A simple asm goto construct to check that callbr is handled correctly:
+;  int asm_goto(int n) {
+;    int v = 1;
+;    asm goto("cmp %0, %1; jnz %l2;" :: "r"(n), "r"(v)::skip_label);
+;    return 0;
+;  skip_label:
+;    return 1;
+;  }
+; asm goto statements can't have outputs, and we do nothing for assembly inputs, so just
+; make sure the compiler doesn't crash.
+define dso_local i32 @asm_goto(i32 %n) sanitize_memory {
+entry:
+  callbr void asm sideeffect "cmp $0, $1; jnz ${2:l}", "r,r,X,~{dirflag},~{fpsr},~{flags}"(i32 %n, i32 1, i8* blockaddress(@asm_goto, %skip_label))
+          to label %cleanup [label %skip_label]
+
+skip_label:                                       ; preds = %entry
+  br label %cleanup
+
+cleanup:                                          ; preds = %entry, %skip_label
+  %retval.0 = phi i32 [ 2, %skip_label ], [ 1, %entry ]
+  ret i32 %retval.0
+}
+
+; CHECK-LABEL: @asm_goto
+; CHECK: call void @__msan_warning