Merge pull request #18439 from rintaro/sourcekit-editoffset-rdar34206143

[SourceKit] Defensive guard for invalid offset and length in edit request

Author: rintaro

test/SourceKit/Misc/Inputs/10bytes.swift

@@ -0,0 +1 @@
+let a = 1

test/SourceKit/Misc/edit_range.swift

@@ -0,0 +1,28 @@
+// RUN: %sourcekitd-test \
+// RUN:   -req=open %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift == \
+// RUN:   -req=edit -offset=10 -length=0 -replace="swift" %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift
+
+// RUN: %sourcekitd-test \
+// RUN:   -req=open %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift == \
+// RUN:   -req=edit -offset=5 -length=5 -replace="swift" %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift
+
+// 'offset' out of range.
+// RUN: not %sourcekitd-test \
+// RUN:   -req=open %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift == \
+// RUN:   -req=edit -offset=11 -length=0 -replace="swift" %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift \
+// RUN:   2>&1 | %FileCheck %s
+
+// 'offset' + 'length' out of range.
+// RUN: not %sourcekitd-test \
+// RUN:   -req=open %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift == \
+// RUN:   -req=edit -offset=5 -length=6 -replace="swift" %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift \
+// RUN:   2>&1 | %FileCheck %s
+
+// Out of range after edits.
+// RUN: not %sourcekitd-test \
+// RUN:   -req=open %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift == \
+// RUN:   -req=edit -offset=5 -length=5 -replace="" %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift == \
+// RUN:   -req=edit -offset=6 -length=0 -replace="swift" %S/Inputs/10bytes.swift -- %S/Inputs/10bytes.swift \
+// RUN:   2>&1 | %FileCheck %s
+
+// CHECK: 'offset' + 'length' is out of range

tools/SourceKit/include/SourceKit/Support/ImmutableTextBuffer.h

@@ -138,6 +138,8 @@ friend class EditableTextBuffer;
 
   ImmutableTextBufferRef getBuffer() const;
 
+  size_t getSize() const;
+
   bool isFromSameBuffer(ImmutableTextSnapshotRef Other) const {
     return Other->EditableBuf.get() == EditableBuf.get();
   }
@@ -169,6 +171,10 @@ class EditableTextBuffer : public ThreadSafeRefCountedBase<EditableTextBuffer> {
     return getSnapshot()->getBuffer();
   }
 
+  size_t getSize() const {
+    return getSnapshot()->getSize();
+  }
+
   ImmutableTextSnapshotRef insert(unsigned ByteOffset, StringRef Text);
   ImmutableTextSnapshotRef erase(unsigned ByteOffset, unsigned Length);
   ImmutableTextSnapshotRef replace(unsigned ByteOffset, unsigned Length,
@@ -178,6 +184,7 @@ class EditableTextBuffer : public ThreadSafeRefCountedBase<EditableTextBuffer> {
   ImmutableTextSnapshotRef addAtomicUpdate(ImmutableTextUpdateRef NewUpd);
   ImmutableTextBufferRef getBufferForSnapshot(
       const ImmutableTextSnapshot &Snap);
+  size_t getSizeForSnapshot(const ImmutableTextSnapshot &Snap) const;
   void refresh();
   friend class ImmutableTextSnapshot;
 };

tools/SourceKit/lib/Support/ImmutableTextBuffer.cpp

@@ -83,6 +83,10 @@ ImmutableTextBufferRef ImmutableTextSnapshot::getBuffer() const {
   return EditableBuf->getBufferForSnapshot(*this);
 }
 
+size_t ImmutableTextSnapshot::getSize() const {
+  return EditableBuf->getSizeForSnapshot(*this);
+}
+
 bool ImmutableTextSnapshot::precedesOrSame(ImmutableTextSnapshotRef Other) {
   assert(Other);
 
@@ -239,6 +243,36 @@ ImmutableTextBufferRef EditableTextBuffer::getBufferForSnapshot(
   return ImmBuf;
 }
 
+size_t EditableTextBuffer::getSizeForSnapshot(
+    const ImmutableTextSnapshot &Snap) const {
+  if (auto Buf = dyn_cast<ImmutableTextBuffer>(Snap.DiffEnd))
+    return Buf->getText().size();
+  ImmutableTextUpdateRef Next = Snap.DiffEnd->Next;
+  // FIXME: dyn_cast_null does not work with IntrusiveRefCntPtr.
+  if (Next)
+    if (auto Buf = dyn_cast<ImmutableTextBuffer>(Next))
+      return Buf->getText().size();
+
+  ImmutableTextBufferRef StartBuf = Snap.BufferStart;
+
+  // Find the last ImmutableTextBuffer.
+  ImmutableTextUpdateRef Upd = StartBuf;
+  while (Upd != Snap.DiffEnd) {
+    Upd = Upd->Next;
+    if (auto Buf = dyn_cast<ImmutableTextBuffer>(Upd))
+      StartBuf = Buf;
+  }
+
+  size_t Length = StartBuf->getText().size();
+  Upd = StartBuf;
+  while (Upd != Snap.DiffEnd) {
+    Upd = Upd->Next;
+    auto Edit = cast<ReplaceImmutableTextUpdate>(Upd);
+    Length = Length - Edit->getLength() + Edit->getText().size();
+  }
+  return Length;
+}
+
 // This should always be called under the mutex lock.
 void EditableTextBuffer::refresh() {
   while (CurrUpd->Next) {

tools/SourceKit/lib/SwiftLang/SwiftEditor.cpp

@@ -1709,10 +1709,16 @@ ImmutableTextSnapshotRef SwiftEditorDocument::initializeText(
 
 ImmutableTextSnapshotRef SwiftEditorDocument::replaceText(
     unsigned Offset, unsigned Length, llvm::MemoryBuffer *Buf,
-    bool ProvideSemanticInfo) {
+    bool ProvideSemanticInfo, std::string &error) {
 
   llvm::sys::ScopedLock L(Impl.AccessMtx);
 
+  // Validate offset and length.
+  if ((Offset + Length) > Impl.EditableBuffer->getSize()) {
+    error = "'offset' + 'length' is out of range";
+    return nullptr;
+  }
+
   Impl.Edited = true;
   llvm::StringRef Str = Buf->getBuffer();
 
@@ -2234,9 +2240,14 @@ void SwiftLangSupport::editorReplaceText(StringRef Name,
       StringRef PreEditTextRef(BufferStart + Offset, Length);
       PreEditText = PreEditTextRef.str();
     }
+    std::string error;
     Snapshot = EditorDoc->replaceText(Offset, Length, Buf,
-                                      Consumer.needsSemanticInfo());
-    assert(Snapshot);
+                                      Consumer.needsSemanticInfo(), error);
+    if (!Snapshot) {
+      assert(error.size());
+      Consumer.handleRequestError(error.c_str());
+      return;
+    }
 
     llvm::Optional<SyntaxParsingCache> SyntaxCache = llvm::None;
     if (EditorDoc->getSyntaxTree().hasValue()) {

tools/SourceKit/lib/SwiftLang/SwiftLangSupport.h

@@ -84,7 +84,8 @@ class SwiftEditorDocument :
                                           ArrayRef<const char *> Args);
   ImmutableTextSnapshotRef replaceText(unsigned Offset, unsigned Length,
                                        llvm::MemoryBuffer *Buf,
-                                       bool ProvideSemanticInfo);
+                                       bool ProvideSemanticInfo,
+                                       std::string &error);
 
   void updateSemaInfo();
 

unittests/SourceKit/Support/ImmutableTextBufferTest.cpp

@@ -18,26 +18,32 @@ using namespace llvm;
 
 TEST(EditableTextBuffer, Updates) {
   const char *Text = "hello world";
+  size_t Length = strlen(Text);
 
   EditableTextBufferManager BufMgr;
   EditableTextBufferRef EdBuf = BufMgr.getOrCreateBuffer("/a/test", Text);
   ImmutableTextBufferRef Buf = EdBuf->getBuffer();
 
   EXPECT_EQ(Buf->getText(), Text);
+  EXPECT_EQ(EdBuf->getSize(), Length);
 
   Buf = EdBuf->insert(6, "all ")->getBuffer();
   EXPECT_EQ(Buf->getText(), "hello all world");
+  EXPECT_EQ(EdBuf->getSize(), strlen("hello all world"));
 
   Buf = EdBuf->erase(9, 6)->getBuffer();
   EXPECT_EQ(Buf->getText(), "hello all");
+  EXPECT_EQ(EdBuf->getSize(), strlen("hello all"));
 
   Buf = EdBuf->replace(0, 5, "yo")->getBuffer();
   EXPECT_EQ(Buf->getText(), "yo all");
+  EXPECT_EQ(EdBuf->getSize(), strlen("yo all"));
 
   EdBuf = BufMgr.resetBuffer("/a/test", Text);
   EdBuf->insert(6, "all ");
   EdBuf->erase(9, 6);
   EdBuf->replace(0, 5, "yo");
+  EXPECT_EQ(EdBuf->getSize(), strlen("yo all"));
   Buf = EdBuf->getSnapshot()->getBuffer();
   EXPECT_EQ(Buf->getText(), "yo all");