Highlight relfected XSS vectors in view-source page.

Source/core/html/HTMLViewSourceDocument.cpp

Index: Source/core/html/HTMLViewSourceDocument.cpp
diff --git a/Source/core/html/HTMLViewSourceDocument.cpp b/Source/core/html/HTMLViewSourceDocument.cpp
index 389b7b9535c9781e51700fe0ea787fbe8a717703..1482f21e280f105683524541fd2d2646774f337b 100644
--- a/Source/core/html/HTMLViewSourceDocument.cpp
+++ b/Source/core/html/HTMLViewSourceDocument.cpp
@@ -47,6 +47,12 @@ namespace WebCore {
 
 using namespace HTMLNames;
 
+namespace {
+
+const char kXSSDetected[] = "Token contains a reflected XSS vector";
+
+} // namespace
+
 HTMLViewSourceDocument::HTMLViewSourceDocument(const DocumentInit& initializer, const String& mimeType)
     : HTMLDocument(initializer)
     , m_type(mimeType)
@@ -86,7 +92,7 @@ void HTMLViewSourceDocument::createContainingTable()
     m_lineNumber = 0;
 }
 
-void HTMLViewSourceDocument::addSource(const String& source, HTMLToken& token)
+void HTMLViewSourceDocument::addSource(const String& source, HTMLToken& token, SourceAnnotation annotation)
 {
     if (!m_current)
         createContainingTable();
@@ -103,13 +109,13 @@ void HTMLViewSourceDocument::addSource(const String& source, HTMLToken& token)
         break;
     case HTMLToken::StartTag:
     case HTMLToken::EndTag:
-        processTagToken(source, token);
+        processTagToken(source, token, annotation);
         break;
     case HTMLToken::Comment:
         processCommentToken(source, token);
         break;
     case HTMLToken::Character:
-        processCharacterToken(source, token);
+        processCharacterToken(source, token, annotation);
         break;
     }
 }
@@ -128,8 +134,9 @@ void HTMLViewSourceDocument::processEndOfFileToken(const String& source, HTMLTok
     m_current = m_td;
 }
 
-void HTMLViewSourceDocument::processTagToken(const String& source, HTMLToken& token)
+void HTMLViewSourceDocument::processTagToken(const String& source, HTMLToken& token, SourceAnnotation annotation)
 {
+    maybeAddSpanForAnnotation(annotation);
     m_current = addSpanWithClassName("webkit-html-tag");
 
     AtomicString tagName(token.name());
@@ -170,9 +177,9 @@ void HTMLViewSourceDocument::processCommentToken(const String& source, HTMLToken
     m_current = m_td;
 }
 
-void HTMLViewSourceDocument::processCharacterToken(const String& source, HTMLToken&)
+void HTMLViewSourceDocument::processCharacterToken(const String& source, HTMLToken&, SourceAnnotation annotation)
 {
-    addText(source, "");
+    addText(source, "", annotation);
 }
 
 PassRefPtrWillBeRawPtr<Element> HTMLViewSourceDocument::addSpanWithClassName(const AtomicString& className)
@@ -223,7 +230,7 @@ void HTMLViewSourceDocument::finishLine()
     m_current = m_tbody;
 }
 
-void HTMLViewSourceDocument::addText(const String& text, const AtomicString& className)
+void HTMLViewSourceDocument::addText(const String& text, const AtomicString& className, SourceAnnotation annotation)
 {
     if (text.isEmpty())
         return;
@@ -242,7 +249,10 @@ void HTMLViewSourceDocument::addText(const String& text, const AtomicString& cla
             finishLine();
             continue;
         }
+        RefPtrWillBeRawPtr<Element> oldElement = m_current;
+        maybeAddSpanForAnnotation(annotation);
         m_current->parserAppendChild(Text::create(*this, substring));
+        m_current = oldElement;
         if (i < size - 1)
             finishLine();
     }
@@ -294,6 +304,14 @@ PassRefPtrWillBeRawPtr<Element> HTMLViewSourceDocument::addLink(const AtomicStri
     return anchor.release();
 }
 
+void HTMLViewSourceDocument::maybeAddSpanForAnnotation(SourceAnnotation annotation)
+{
+    if (annotation == AnnotateSourceAsXSS) {
+        m_current = addSpanWithClassName("webkit-highlight");
+        m_current->setAttribute(titleAttr, kXSSDetected);
+    }
+}
+
 void HTMLViewSourceDocument::trace(Visitor* visitor)
 {
     visitor->trace(m_current);