prevent Symbol attack in Debug scanner

2 files changed, 28 insertions, 23 deletions

diff --git a/lib/coderay/scanners/debug.rb b/lib/coderay/scanners/debug.rb
index 9d10864..ac12c16 100644
--- a/lib/coderay/scanners/debug.rb
+++ b/lib/coderay/scanners/debug.rb
@@ -1,3 +1,5 @@
+require 'set'
+
 module CodeRay
 module Scanners
   
@@ -11,6 +13,11 @@ module Scanners
     
   protected
     
+    def setup
+      super
+      @known_token_kinds = TokenKinds.keys.map(&:to_s).to_set
+    end
+    
     def scan_tokens encoder, options
       
       opened_tokens = []
@@ -21,26 +28,24 @@ module Scanners
           encoder.text_token match, :space
           
         elsif match = scan(/ (\w+) \( ( [^\)\\]* ( \\. [^\)\\]* )* ) \)? /x)
-          # FIXME: cache attack
-          kind = self[1].to_sym
-          match = self[2].gsub(/\\(.)/m, '\1')
-          unless TokenKinds.has_key? kind
-            kind = :error
-            match = matched
+          if @known_token_kinds.include? self[1]
+            encoder.text_token self[2].gsub(/\\(.)/m, '\1'), self[1].to_sym
+          else
+            encoder.text_token matched, :error
           end
-          encoder.text_token match, kind
           
         elsif match = scan(/ (\w+) ([<\[]) /x)
-          # FIXME: cache attack
-          kind = self[1].to_sym
-          opened_tokens << kind
-          case self[2]
-          when '<'
-            encoder.begin_group kind
-          when '['
-            encoder.begin_line kind
-          else
-            raise 'CodeRay bug: This case should not be reached.'
+          if @known_token_kinds.include? self[1]
+            kind = self[1].to_sym
+            opened_tokens << kind
+            case self[2]
+            when '<'
+              encoder.begin_group kind
+            when '['
+              encoder.begin_line kind
+            else
+              raise 'CodeRay bug: This case should not be reached.'
+            end
           end
           
         elsif !opened_tokens.empty? && match = scan(/ > /x)
diff --git a/test/unit/debug.rb b/test/unit/debug.rb
index f2b80bd..616cda5 100644
--- a/test/unit/debug.rb
+++ b/test/unit/debug.rb
@@ -18,15 +18,15 @@ class DebugEncoderTest < Test::Unit::TestCase
     [:begin_group, :string],
     ['test', :content],
     [:end_group, :string],
-    [:begin_line, :test],
+    [:begin_line, :head],
     ["\n", :space],
     ["\n  \t", :space],
     ["   \n", :space],
     ["[]", :method],
-    [:end_line, :test],
+    [:end_line, :head],
   ].flatten
   TEST_OUTPUT = <<-'DEBUG'.chomp
-integer(10)operator((\\\))string<content(test)>test[
+integer(10)operator((\\\))string<content(test)>head[
 
   	   
 method([])]
@@ -51,7 +51,7 @@ class DebugScannerTest < Test::Unit::TestCase
   end
   
   TEST_INPUT = <<-'DEBUG'.chomp
-integer(10)operator((\\\))string<content(test)>test[
+integer(10)operator((\\\))string<content(test)>head[
 
   	   
 method([])]
@@ -62,10 +62,10 @@ method([])]
     [:begin_group, :string],
     ['test', :content],
     [:end_group, :string],
-    [:begin_line, :test],
+    [:begin_line, :head],
     ["\n\n  \t   \n", :space],
     ["[]", :method],
-    [:end_line, :test],
+    [:end_line, :head],
   ].flatten
   
   def test_filtering_text_tokens