[objcmt] Fix a buffer overflow crash than can occur while modernizing enums.

Note that due to the nature of the crash it requires libgmalloc or asan for it to crash consistently.

rdar://19932927

git-svn-id: https://llvm.org/svn/llvm-project/cfe/trunk@275600 91177308-0d34-0410-b5e6-96231b3b80d8

1 files changed, 17 insertions, 27 deletions

diff --git a/lib/ARCMigrate/ObjCMT.cpp b/lib/ARCMigrate/ObjCMT.cpp
index 940a9c95f2..241a7246b6 100644
--- a/lib/ARCMigrate/ObjCMT.cpp
+++ b/lib/ARCMigrate/ObjCMT.cpp
@@ -771,23 +771,11 @@ static void rewriteToNSMacroDecl(ASTContext &Ctx,
   ClassString += ", ";
   
   ClassString += TypedefDcl->getIdentifier()->getName();
-  ClassString += ')';
-  SourceLocation EndLoc;
-  if (EnumDcl->getIntegerTypeSourceInfo()) {
-    TypeSourceInfo *TSourceInfo = EnumDcl->getIntegerTypeSourceInfo();
-    TypeLoc TLoc = TSourceInfo->getTypeLoc();
-    EndLoc = TLoc.getLocEnd();
-    const char *lbrace = Ctx.getSourceManager().getCharacterData(EndLoc);
-    unsigned count = 0;
-    if (lbrace)
-      while (lbrace[count] != '{')
-        ++count;
-    if (count > 0)
-      EndLoc = EndLoc.getLocWithOffset(count-1);
-  }
-  else
-    EndLoc = EnumDcl->getLocStart();
-  SourceRange R(EnumDcl->getLocStart(), EndLoc);
+  ClassString += ") ";
+  SourceLocation EndLoc = EnumDcl->getBraceRange().getBegin();
+  if (EndLoc.isInvalid())
+    return;
+  CharSourceRange R = CharSourceRange::getCharRange(EnumDcl->getLocStart(), EndLoc);
   commit.replace(R, ClassString);
   // This is to remove spaces between '}' and typedef name.
   SourceLocation StartTypedefLoc = EnumDcl->getLocEnd();
@@ -1900,18 +1888,20 @@ void ObjCMigrateASTConsumer::HandleTranslationUnit(ASTContext &Ctx) {
         if (++N == DEnd)
           continue;
         if (const EnumDecl *ED = dyn_cast<EnumDecl>(*N)) {
-          if (++N != DEnd)
-            if (const TypedefDecl *TDF = dyn_cast<TypedefDecl>(*N)) {
-              // prefer typedef-follows-enum to enum-follows-typedef pattern.
-              if (migrateNSEnumDecl(Ctx, ED, TDF)) {
-                ++D; ++D;
-                CacheObjCNSIntegerTypedefed(TD);
-                continue;
+          if (canModify(ED)) {
+            if (++N != DEnd)
+              if (const TypedefDecl *TDF = dyn_cast<TypedefDecl>(*N)) {
+                // prefer typedef-follows-enum to enum-follows-typedef pattern.
+                if (migrateNSEnumDecl(Ctx, ED, TDF)) {
+                  ++D; ++D;
+                  CacheObjCNSIntegerTypedefed(TD);
+                  continue;
+                }
               }
+            if (migrateNSEnumDecl(Ctx, ED, TD)) {
+              ++D;
+              continue;
             }
-          if (migrateNSEnumDecl(Ctx, ED, TD)) {
-            ++D;
-            continue;
           }
         }
         CacheObjCNSIntegerTypedefed(TD);