[clang] [clang][ASTImporter] fix assert fail due to offset overflow (PR #79084)

[Qizhi Hu via cfe-commits]

https://github.com/jcsxky updated https://github.com/llvm/llvm-project/pull/79084

>From 0b127ff20bc9a6a1d4de7bcbfa2c5ad466cea14c Mon Sep 17 00:00:00 2001
From: huqizhi <huqizhi at feysh.com>
Date: Tue, 23 Jan 2024 10:09:44 +0800
Subject: [PATCH] [clang][ASTImporter] fix assert fail due to offset overflow

---
 clang/lib/AST/ASTImporter.cpp | 16 +++++++++++++++-
 1 file changed, 15 insertions(+), 1 deletion(-)

diff --git a/clang/lib/AST/ASTImporter.cpp b/clang/lib/AST/ASTImporter.cpp
index 12734d62ed9fb76..34f817a098b0fb3 100644
--- a/clang/lib/AST/ASTImporter.cpp
+++ b/clang/lib/AST/ASTImporter.cpp
@@ -9837,6 +9837,13 @@ Expected<SourceRange> ASTImporter::Import(SourceRange FromRange) {
   return SourceRange(ToBegin, ToEnd);
 }
 
+static bool isBufferSizeOverflow(SourceManager &SM, size_t BufferSize) {
+  unsigned Offset = SM.getNextLocalOffset();
+  unsigned FullSize = Offset + BufferSize + 1;
+  SourceLocation L = SourceLocation().getFromRawEncoding(FullSize);
+  return !L.isFileID() || FullSize <= Offset;
+}
+
 Expected<FileID> ASTImporter::Import(FileID FromID, bool IsBuiltin) {
   llvm::DenseMap<FileID, FileID>::iterator Pos = ImportedFileIDs.find(FromID);
   if (Pos != ImportedFileIDs.end())
@@ -9896,9 +9903,13 @@ Expected<FileID> ASTImporter::Import(FileID FromID, bool IsBuiltin) {
         // FIXME: The filename may be a virtual name that does probably not
         // point to a valid file and we get no Entry here. In this case try with
         // the memory buffer below.
-        if (Entry)
+        if (Entry) {
+          if (isBufferSizeOverflow(ToSM, Entry->getSize()))
+            return llvm::make_error<ASTImportError>(
+                ASTImportError::UnsupportedConstruct);
           ToID = ToSM.createFileID(*Entry, ToIncludeLocOrFakeLoc,
                                    FromSLoc.getFile().getFileCharacteristic());
+        }
       }
     }
 
@@ -9913,6 +9924,9 @@ Expected<FileID> ASTImporter::Import(FileID FromID, bool IsBuiltin) {
       std::unique_ptr<llvm::MemoryBuffer> ToBuf =
           llvm::MemoryBuffer::getMemBufferCopy(FromBuf->getBuffer(),
                                                FromBuf->getBufferIdentifier());
+      if (isBufferSizeOverflow(ToSM, ToBuf->getBufferSize()))
+        return llvm::make_error<ASTImportError>(
+            ASTImportError::UnsupportedConstruct);
       ToID = ToSM.createFileID(std::move(ToBuf),
                                FromSLoc.getFile().getFileCharacteristic());
     }