[ruby/erb] Copy CGI.escapeHTML to ERB::Util.html_escape

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

5 files changed, 126 insertions, 2 deletions

diff --git a/ext/erb/erb.c b/ext/erb/erb.c
new file mode 100644
index 0000000000..92cfbd0769
--- /dev/null
+++ b/ext/erb/erb.c
@@ -0,0 +1,96 @@
+#include "ruby.h"
+#include "ruby/encoding.h"
+
+static VALUE rb_cERB, rb_mEscape;
+
+#define HTML_ESCAPE_MAX_LEN 6
+
+static const struct {
+    uint8_t len;
+    char str[HTML_ESCAPE_MAX_LEN+1];
+} html_escape_table[UCHAR_MAX+1] = {
+#define HTML_ESCAPE(c, str) [c] = {rb_strlen_lit(str), str}
+    HTML_ESCAPE('\'', "'"),
+    HTML_ESCAPE('&', "&"),
+    HTML_ESCAPE('"', """),
+    HTML_ESCAPE('<', "&lt;"),
+    HTML_ESCAPE('>', "&gt;"),
+#undef HTML_ESCAPE
+};
+
+static inline void
+preserve_original_state(VALUE orig, VALUE dest)
+{
+    rb_enc_associate(dest, rb_enc_get(orig));
+}
+
+static inline long
+escaped_length(VALUE str)
+{
+    const long len = RSTRING_LEN(str);
+    if (len >= LONG_MAX / HTML_ESCAPE_MAX_LEN) {
+        ruby_malloc_size_overflow(len, HTML_ESCAPE_MAX_LEN);
+    }
+    return len * HTML_ESCAPE_MAX_LEN;
+}
+
+static VALUE
+optimized_escape_html(VALUE str)
+{
+    VALUE vbuf;
+    char *buf = ALLOCV_N(char, vbuf, escaped_length(str));
+    const char *cstr = RSTRING_PTR(str);
+    const char *end = cstr + RSTRING_LEN(str);
+
+    char *dest = buf;
+    while (cstr < end) {
+        const unsigned char c = *cstr++;
+        uint8_t len = html_escape_table[c].len;
+        if (len) {
+            memcpy(dest, html_escape_table[c].str, len);
+            dest += len;
+        }
+        else {
+            *dest++ = c;
+        }
+    }
+
+    VALUE escaped;
+    if (RSTRING_LEN(str) < (dest - buf)) {
+        escaped = rb_str_new(buf, dest - buf);
+        preserve_original_state(str, escaped);
+    }
+    else {
+        escaped = rb_str_dup(str);
+    }
+    ALLOCV_END(vbuf);
+    return escaped;
+}
+
+static VALUE
+cgiesc_escape_html(VALUE self, VALUE str)
+{
+    StringValue(str);
+
+    if (rb_enc_str_asciicompat_p(str)) {
+        return optimized_escape_html(str);
+    }
+    else {
+        return rb_call_super(1, &str);
+    }
+}
+
+static VALUE
+erb_escape_html(VALUE self, VALUE str)
+{
+    str = rb_funcall(str, rb_intern("to_s"), 0);
+    return cgiesc_escape_html(self, str);
+}
+
+void
+Init_erb(void)
+{
+    rb_cERB = rb_define_class("ERB", rb_cObject);
+    rb_mEscape = rb_define_module_under(rb_cERB, "Escape");
+    rb_define_method(rb_mEscape, "html_escape", erb_escape_html, 1);
+}
diff --git a/ext/erb/extconf.rb b/ext/erb/extconf.rb
new file mode 100644
index 0000000000..00a7e92aea
--- /dev/null
+++ b/ext/erb/extconf.rb
@@ -0,0 +1,2 @@
+require 'mkmf'
+create_makefile 'erb'
diff --git a/lib/erb.gemspec b/lib/erb.gemspec
index 2e7e981ff1..419685c318 100644
--- a/lib/erb.gemspec
+++ b/lib/erb.gemspec
@@ -27,8 +27,11 @@ Gem::Specification.new do |spec|
   spec.executables   = ['erb']
   spec.require_paths = ['lib']
 
-  if RUBY_ENGINE != 'jruby'
+  if RUBY_ENGINE == 'jruby'
+    spec.platform = 'java'
+  else
     spec.required_ruby_version = '>= 2.7.0'
+    spec.extensions = ['ext/erb/extconf.rb']
   end
 
   spec.add_dependency 'cgi', '>= 0.3.3'
diff --git a/lib/erb.rb b/lib/erb.rb
index 962eeb6963..c588ae1a65 100644
--- a/lib/erb.rb
+++ b/lib/erb.rb
@@ -986,7 +986,6 @@ end
 class ERB
   # A utility module for conversion routines, often handy in HTML generation.
   module Util
-    public
     #
     # A utility method for escaping HTML tag characters in _s_.
     #
@@ -1002,6 +1001,17 @@ class ERB
     def html_escape(s)
       CGI.escapeHTML(s.to_s)
     end
+  end
+
+  begin
+    require 'erb.so'
+  rescue LoadError
+  else
+    private_constant :Escape
+    Util.prepend(Escape)
+  end
+
+  module Util
     alias h html_escape
     module_function :h
     module_function :html_escape
diff --git a/test/erb/test_erb.rb b/test/erb/test_erb.rb
index 424ddae87e..1db0e55f8a 100644
--- a/test/erb/test_erb.rb
+++ b/test/erb/test_erb.rb
@@ -73,11 +73,24 @@ class TestERB < Test::Unit::TestCase
     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;", ERB::Util.html_escape("'&\"><"))
 
     assert_equal("", ERB::Util.html_escape(nil))
     assert_equal("123", ERB::Util.html_escape(123))
   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'