[ASTImporterTest] Fix potential use-after-free

[a.sidorin]

Summary

buildASTFromCodeWithArgs() accepts llvm::Twine as Code argument. However, if the argument is not a C string or std::string, the argument is being copied into a temporary buffer in order to get a null-terminated string. This lead to a potential UAF. Fixing this via calling .data() on StringRef since our Code is always null-terminated.

The issue was introduced by me in D44079 (sorry) but was not noticed.

[martong]

unittests/AST/ASTImporterTest.cpp

210–211

Can't we have the same problem with FileName?

Perhaps an other alternative would be to make the members real strings.

std::string Code;
std::string FileName;

[a.sidorin]

Hi Gabor,

Can't we have the same problem with FileName?

As I can see, no. FileName is copied into std::string while building compilation arguments.

Perhaps an other alternative would be to make the members real strings.

That was the initial solution and you was correct with it. However, it is still a workaround. The ultimate solution is to fix buildASTFromCodeWithArgs() so it can copy the buffer into a persistent storage if it is not null-terminated:

StringRef NullTerminated = Code.toNullTerminatedStringRef(CodeStorage);
auto CodeBuffer =
    CodeStorage.empty()
        ? llvm::MemoryBuffer::getMemBuffer(NullTerminated)
        : llvm::MemoryBuffer::getMemBufferCopy(NullTerminated);
InMemoryFileSystem->addFile(FileNameRef, 0, std::move(CodeBuffer));

Unfortunately, this patch has two problems. The first one is double copying of the input Twine, but it can be avoided with some additional code. The second one is copying of null-terminated StringRefs which, I guess, cannot be avoided at all because there is no legal way to check if StringRef is null-terminated or not. So, I'm summoning the initial authors of the code for the help.

@pcc, @klimek Could you propose a better idea?

[a.sidorin]

After a number of attempts of Twine'ifying all Code samples, I decided to restore the initial state of this code.

[a.sidorin]

Add forgotten context.

[martong]

LGTM!

[a.sidorin]

Closed with rC332256.

unittests/AST/ASTImporterTest.cpp

class ASTImporterTestBase : public ::testing::TestWithParam<ArgVector> {
   const char *const InputFileName = "input.cc";
   const char *const OutputFileName = "output.cc";
 
   // Buffer for the To context, must live in the test scope.
   std::string ToCode;
 
   struct TU {
     // Buffer for the context, must live in the test scope.
-    StringRef Code;
-    StringRef FileName;
+    std::string Code;
+    std::string FileName;
     std::unique_ptr<ASTUnit> Unit;
     TranslationUnitDecl *TUDecl = nullptr;
     TU(StringRef Code, StringRef FileName, ArgVector Args)
         : Code(Code), FileName(FileName),
           Unit(tooling::buildASTFromCodeWithArgs(this->Code, Args,
                                                  this->FileName)),
           TUDecl(Unit->getASTContext().getTranslationUnitDecl()) {}
   };