[C API] Prevent nullptr dereferences in C API for counting attributes.

See https://reviews.llvm.org/D26392

Patch by @maleadt

llvm-svn: 287044

4 files changed, 84 insertions, 0 deletions

diff --git a/llvm/tools/llvm-c-test/CMakeLists.txt b/llvm/tools/llvm-c-test/CMakeLists.txt
index 858f2b07214..a2bde0d9714 100644
--- a/llvm/tools/llvm-c-test/CMakeLists.txt
+++ b/llvm/tools/llvm-c-test/CMakeLists.txt
@@ -36,6 +36,7 @@ if (LLVM_COMPILER_IS_GCC_COMPATIBLE)
 endif ()
 
 add_llvm_tool(llvm-c-test
+  attributes.c
   calc.c
   diagnostic.c
   disassemble.c
diff --git a/llvm/tools/llvm-c-test/attributes.c b/llvm/tools/llvm-c-test/attributes.c
new file mode 100644
index 00000000000..7f07a220ad3
--- /dev/null
+++ b/llvm/tools/llvm-c-test/attributes.c
@@ -0,0 +1,75 @@
+/*===-- attributes.c - tool for testing libLLVM and llvm-c API ------------===*\
+|*                                                                            *|
+|*                     The LLVM Compiler Infrastructure                       *|
+|*                                                                            *|
+|* This file is distributed under the University of Illinois Open Source      *|
+|* License. See LICENSE.TXT for details.                                      *|
+|*                                                                            *|
+|*===----------------------------------------------------------------------===*|
+|*                                                                            *|
+|* This file implements the --test-attributes and --test-callsite-attributes  *|
+|* commands in llvm-c-test.                                                   *|
+|*                                                                            *|
+\*===----------------------------------------------------------------------===*/
+
+#include "llvm-c-test.h"
+
+#include <stdlib.h>
+
+int llvm_test_function_attributes(void) {
+  LLVMEnablePrettyStackTrace();
+
+  LLVMModuleRef M = llvm_load_module(false, true);
+
+  LLVMValueRef F = LLVMGetFirstFunction(M);
+  while (F) {
+    // Read attributes
+    for (int Idx = LLVMAttributeFunctionIndex, ParamCount = LLVMCountParams(F);
+         Idx <= ParamCount; ++Idx) {
+      int AttrCount = LLVMGetAttributeCountAtIndex(F, Idx);
+      LLVMAttributeRef *Attrs =
+          (LLVMAttributeRef *)malloc(AttrCount * sizeof(LLVMAttributeRef));
+      LLVMGetAttributesAtIndex(F, Idx, Attrs);
+      free(Attrs);
+    }
+    F = LLVMGetNextFunction(F);
+  }
+
+  LLVMDisposeModule(M);
+
+  return 0;
+}
+
+int llvm_test_callsite_attributes(void) {
+  LLVMEnablePrettyStackTrace();
+
+  LLVMModuleRef M = llvm_load_module(false, true);
+
+  LLVMValueRef F = LLVMGetFirstFunction(M);
+  while (F) {
+    LLVMBasicBlockRef BB;
+    for (BB = LLVMGetFirstBasicBlock(F); BB; BB = LLVMGetNextBasicBlock(BB)) {
+      LLVMValueRef I;
+      for (I = LLVMGetFirstInstruction(BB); I; I = LLVMGetNextInstruction(I)) {
+        if (LLVMIsACallInst(I)) {
+          // Read attributes
+          for (int Idx = LLVMAttributeFunctionIndex,
+                   ParamCount = LLVMCountParams(F);
+               Idx <= ParamCount; ++Idx) {
+            int AttrCount = LLVMGetCallSiteAttributeCount(I, Idx);
+            LLVMAttributeRef *Attrs = (LLVMAttributeRef *)malloc(
+                AttrCount * sizeof(LLVMAttributeRef));
+            LLVMGetCallSiteAttributes(I, Idx, Attrs);
+            free(Attrs);
+          }
+        }
+      }
+    }
+
+    F = LLVMGetNextFunction(F);
+  }
+
+  LLVMDisposeModule(M);
+
+  return 0;
+}
diff --git a/llvm/tools/llvm-c-test/llvm-c-test.h b/llvm/tools/llvm-c-test/llvm-c-test.h
index 0d1ade093bc..0401f1dbae3 100644
--- a/llvm/tools/llvm-c-test/llvm-c-test.h
+++ b/llvm/tools/llvm-c-test/llvm-c-test.h
@@ -52,6 +52,10 @@ int llvm_echo(void);
 // diagnostic.c
 int llvm_test_diagnostic_handler(void);
 
+// attributes.c
+int llvm_test_function_attributes();
+int llvm_test_callsite_attributes();
+
 #ifdef __cplusplus
 }
 #endif /* !defined(__cplusplus) */
diff --git a/llvm/tools/llvm-c-test/main.c b/llvm/tools/llvm-c-test/main.c
index 90d35017785..9bc0c96c3d6 100644
--- a/llvm/tools/llvm-c-test/main.c
+++ b/llvm/tools/llvm-c-test/main.c
@@ -88,6 +88,10 @@ int main(int argc, char **argv) {
     return llvm_add_named_metadata_operand();
   } else if (argc == 2 && !strcmp(argv[1], "--set-metadata")) {
     return llvm_set_metadata();
+  } else if (argc == 2 && !strcmp(argv[1], "--test-function-attributes")) {
+    return llvm_test_function_attributes();
+  } else if (argc == 2 && !strcmp(argv[1], "--test-callsite-attributes")) {
+    return llvm_test_callsite_attributes();
   } else if (argc == 2 && !strcmp(argv[1], "--echo")) {
     return llvm_echo();
   } else if (argc == 2 && !strcmp(argv[1], "--test-diagnostic-handler")) {