[ruby/erb] escape.c: Fix UB in find_next_match_neon

(https://github.com/ruby/erb/pull/142)

Shifting by 64 or more bits is undefined behavior.

https://github.com/ruby/erb/commit/a9316104fc

ext/erb/escape/escape.c

@@ -100,10 +100,7 @@ find_next_match_sse2(search_state *search)
     int next_match_offset = trailing_zeros(search->matches_bitmap);
     search->matches_bitmap >>= (next_match_offset + 1);
     search->cstr += next_match_offset;
-    if (search->cstr > search->end) {
-        search->cstr = search->end;
-        return false;
-    }
+    RUBY_ASSERT(search->cstr <= search->end);
     return true;
 }
 
@@ -170,12 +167,16 @@ static inline uint32_t trailing_zeros64(uint64_t input)
 static inline bool
 find_next_match_neon(search_state *search)
 {
-    size_t next_match_offset = trailing_zeros64(search->matches_bitmap) / 4;
-    search->matches_bitmap >>= (next_match_offset + 1) * 4;
-    search->cstr += next_match_offset;
-    if (search->cstr > search->end) {
-        search->cstr = search->end;
-        return false;
+    uint32_t trailing_zeros = trailing_zeros64(search->matches_bitmap);
+
+    // uint64_t >>= 64 is undefined behaviour
+    if (trailing_zeros >= 63) {
+        search->matches_bitmap = 0;
+        search->cstr += 15;
+    }
+    else {
+        search->matches_bitmap >>= (trailing_zeros + 1);
+        search->cstr += trailing_zeros / 4;
     }
     return true;
 }
@@ -207,10 +208,10 @@ find_next_neon(search_state *search)
         const uint8x16_t matches = vorrq_u8(mask2, mask3);
 
         const uint8x8_t res = vshrn_n_u16(vreinterpretq_u16_u8(matches), 4);
-        const uint64_t bitmap = vget_lane_u64(vreinterpret_u64_u8(res), 0) & 0x8888888888888888ull;
+        const uint64_t bitmap = vget_lane_u64(vreinterpret_u64_u8(res), 0);
 
         if (bitmap) {
-            search->matches_bitmap = bitmap;
+            search->matches_bitmap = bitmap & 0x8888888888888888ull;
             return find_next_match_neon(search);
         }
         search->cstr += sizeof(uint8x16_t);

test/erb/test_erb.rb

@@ -43,34 +43,6 @@ class TestERB < Test::Unit::TestCase
     assert_match(/\Atest filename:201\b/, e.backtrace[0])
   end
 
-  def test_html_escape
-    assert_equal(" !&quot;\#$%&amp;&#39;()*+,-./0123456789:;&lt;=&gt;?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~",
-                 ERB::Util.html_escape(" !\"\#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~"))
-
-    assert_equal("", ERB::Util.html_escape(""))
-    assert_equal("abc", ERB::Util.html_escape("abc"))
-    assert_equal("&lt;&lt;", ERB::Util.html_escape("<\<"))
-    assert_equal("&#39;&amp;&quot;&gt;&lt;" * 10, ERB::Util.html_escape("'&\"><" * 10))
-
-    assert_equal("", ERB::Util.html_escape(nil))
-    assert_equal("123", ERB::Util.html_escape(123))
-
-    assert_equal(65536+5, ERB::Util.html_escape("x"*65536 + "&").size)
-    assert_equal(65536+5, ERB::Util.html_escape("&" + "x"*65536).size)
-  end
-
-  def test_html_escape_to_s
-    object = Object.new
-    def object.to_s
-      "object"
-    end
-    assert_equal("object", ERB::Util.html_escape(object))
-  end
-
-  def test_html_escape_extension
-    assert_nil(ERB::Util.method(:html_escape).source_location)
-  end if RUBY_ENGINE == 'ruby'
-
   def test_concurrent_default_binding
     # This test randomly fails with JRuby -- NameError: undefined local variable or method `template2'
     pend if RUBY_ENGINE == 'jruby'
@@ -770,13 +742,4 @@ class TestERBRactor < Test::Unit::TestCase
       assert_equal(["2", "2"], rs.map(&:value))
     RUBY
   end
-
-  def test_util_html_escape_in_ractor
-    assert_ractor(<<~RUBY, require: 'erb')
-      r = Ractor.new do
-        ERB::Util.html_escape("<script>")
-      end
-      assert_equal("&lt;script&gt;", r.value)
-    RUBY
-  end
 end

test/erb/test_erb_escape.rb

@@ -0,0 +1,57 @@
+require 'test/unit'
+require 'erb'
+
+class TestERBEscape < Test::Unit::TestCase
+  def test_html_escape
+    assert_equal(" !&quot;\#$%&amp;&#39;()*+,-./0123456789:;&lt;=&gt;?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~",
+                 ERB::Util.html_escape(" !\"\#$%&'()*+,-./0123456789:;<=>?@ABCDEFGHIJKLMNOPQRSTUVWXYZ[\\]^_`abcdefghijklmnopqrstuvwxyz{|}~"))
+
+    assert_equal("", ERB::Util.html_escape(""))
+    assert_equal("abc", ERB::Util.html_escape("abc"))
+    assert_equal("&lt;&lt;", ERB::Util.html_escape("<\<"))
+    assert_equal("&#39;&amp;&quot;&gt;&lt;" * 10, ERB::Util.html_escape("'&\"><" * 10))
+
+    assert_equal("", ERB::Util.html_escape(nil))
+    assert_equal("123", ERB::Util.html_escape(123))
+
+    assert_equal(65536+5, ERB::Util.html_escape("x"*65536 + "&").size)
+    assert_equal(65536+5, ERB::Util.html_escape("&" + "x"*65536).size)
+  end
+
+  def test_html_escape_to_s
+    object = Object.new
+    def object.to_s
+      "object"
+    end
+    assert_equal("object", ERB::Util.html_escape(object))
+  end
+
+  def test_html_escape_extension
+    assert_nil(ERB::Util.method(:html_escape).source_location)
+  end if RUBY_ENGINE == 'ruby'
+
+  def test_util_html_escape_in_ractor
+    assert_ractor(<<~RUBY, require: 'erb')
+      r = Ractor.new do
+        ERB::Util.html_escape("<script>")
+      end
+      assert_equal("&lt;script&gt;", r.value)
+    RUBY
+  end
+
+  def test_simd_coverage
+    16.times do |i|
+      str = "#{'.' * i}<#{'.' * (16-i-1)}#{'<' * 16}"
+      esc = "#{'.' * i}&lt;#{'.'* (16-i-1)}#{'&lt;' * 16}"
+      assert_equal esc, h(str)
+    end
+
+    assert_equal '&lt;' * 32, h('<' * 32)
+  end
+
+  private
+
+  def h(...)
+    ERB::Util.html_escape(...)
+  end
+end