Add bundle id check such as it should not crash when given a malformat bundle id

Author: shiltian

clang/include/clang/Driver/OffloadBundler.h

@@ -158,6 +158,11 @@ class CompressedOffloadBundle {
   static llvm::Expected<std::unique_ptr<llvm::MemoryBuffer>>
   decompress(const llvm::MemoryBuffer &Input, bool Verbose = false);
 };
+
+/// Check whether the bundle id is in the following format:
+/// <kind>-<triple>[-<target id>[:target features]]
+/// <triple> := <arch>-<vendor>-<os>-<env>
+bool checkOffloadBundleID(const llvm::StringRef Str);
 } // namespace clang
 
 #endif // LLVM_CLANG_DRIVER_OFFLOADBUNDLER_H

clang/lib/Driver/OffloadBundler.cpp

@@ -83,7 +83,6 @@ OffloadTargetInfo::OffloadTargetInfo(const StringRef Target,
                                      const OffloadBundlerConfig &BC)
     : BundlerConfig(BC) {
 
-  // TODO: Add error checking from ClangOffloadBundler.cpp
   // <kind>-<triple>[-<target id>[:target features]]
   // <triple> := <arch>-<vendor>-<os>-<env>
   SmallVector<StringRef, 6> Components;
@@ -1514,6 +1513,9 @@ Error OffloadBundler::UnbundleFiles() {
   StringMap<StringRef> Worklist;
   auto Output = BundlerConfig.OutputFileNames.begin();
   for (auto &Triple : BundlerConfig.TargetNames) {
+    if (!checkOffloadBundleID(Triple))
+      return createStringError(errc::invalid_argument,
+                               "invalid bundle id from bundle config");
     Worklist[Triple] = *Output;
     ++Output;
   }
@@ -1533,6 +1535,9 @@ Error OffloadBundler::UnbundleFiles() {
 
     StringRef CurTriple = **CurTripleOrErr;
     assert(!CurTriple.empty());
+    if (!checkOffloadBundleID(CurTriple))
+      return createStringError(errc::invalid_argument,
+                               "invalid bundle id read from the bundle");
 
     auto Output = Worklist.begin();
     for (auto E = Worklist.end(); Output != E; Output++) {
@@ -1591,6 +1596,8 @@ Error OffloadBundler::UnbundleFiles() {
         return createFileError(E.second, EC);
 
       // If this entry has a host kind, copy the input file to the output file.
+      // We don't need to check E.getKey() here through checkOffloadBundleID
+      // because the entire WorkList has been checked above.
       auto OffloadInfo = OffloadTargetInfo(E.getKey(), BundlerConfig);
       if (OffloadInfo.hasHostKind())
         OutputFile.write(Input.getBufferStart(), Input.getBufferSize());
@@ -1820,6 +1827,10 @@ Error OffloadBundler::UnbundleArchive() {
     // archive.
     while (!CodeObject.empty()) {
       SmallVector<StringRef> CompatibleTargets;
+      if (!checkOffloadBundleID(CodeObject)) {
+        return createStringError(errc::invalid_argument,
+                                 "Invalid bundle id read from code object");
+      }
       auto CodeObjectInfo = OffloadTargetInfo(CodeObject, BundlerConfig);
       if (getCompatibleOffloadTargets(CodeObjectInfo, CompatibleTargets,
                                       BundlerConfig)) {
@@ -1901,3 +1912,11 @@ Error OffloadBundler::UnbundleArchive() {
 
   return Error::success();
 }
+
+bool clang::checkOffloadBundleID(const llvm::StringRef Str) {
+  // <kind>-<triple>[-<target id>[:target features]]
+  // <triple> := <arch>-<vendor>-<os>-<env>
+  SmallVector<StringRef, 6> Components;
+  Str.split(Components, '-', /*MaxSplit=*/5);
+  return Components.size() == 5 || Components.size() == 6;
+}

clang/test/Driver/clang-offload-bundler.c

@@ -557,6 +557,12 @@
 // RUN: llvm-ar t %T/hip-openmp_906.a | FileCheck -check-prefix=OPENMPHIPCOMPAT %s
 // OPENMPHIPCOMPAT: hip_bundle1-hip-amdgcn-amd-amdhsa--gfx906
 
+// Check if a malformat bundle id can be detected and an error can be emitted.
+// RUN: not clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa -output=%T/hip-openmp_906.a -input=%T/hip_archive.a -hip-openmp-compatible 2>&1 | FileCheck %s -check-prefix=ERROR
+// ERROR: error: Targets need to follow the format '<offload kind>-<target triple>', where '<target triple>' follows the format '<kind>-<arch>-<vendor>-<os>-<env>[-<target id>[:target features]]'.
+// RUN: not clang-offload-bundler -unbundle -type=a -targets=openmp-amdgcn-amd-amdhsa-gfx906 -output=%T/hip-openmp_906.a -input=%T/hip_archive.a -hip-openmp-compatible 2>&1 | FileCheck %s -check-prefix=ERROR
+// ERROR: error: no compatible code object found for the target 'openmp-amdgcn-amd-amdhsa--'
+
 // Some code so that we can create a binary out of this file.
 int A = 0;
 void test_func(void) {

clang/tools/clang-offload-bundler/ClangOffloadBundler.cpp

@@ -355,6 +355,15 @@ int main(int argc, const char **argv) {
           errc::invalid_argument, "Duplicate targets are not allowed"));
     }
 
+    if (!checkOffloadBundleID(Target)) {
+      return reportError(createStringError(
+          errc::invalid_argument,
+          "Targets need to follow the format '<offload kind>-<target triple>', "
+          "where '<target triple>' follows the format "
+          "'<kind>-<arch>-<vendor>-<os>-<env>[-<target id>[:target "
+          "features]]'."));
+    }
+
     auto OffloadInfo = OffloadTargetInfo(Target, BundlerConfig);
     bool KindIsValid = OffloadInfo.isOffloadKindValid();
     bool TripleIsValid = OffloadInfo.isTripleValid();