Safely decode non-ascii SQL queries for display

SQL queries containing non-ascii byte strings would cause errors, both with and
without Pygments highlighting.

This updates the non-Pygments case to handle a simple decoding to ensure the
value is ascii-safe. It also removes passing an explicit "utf-8" encoding to
Pygments, since this causes errors when the bytes are not utf-8. When the
encoding is omitted, Pygments will default to "guess" the encoding by trying
utf-8 and falling back to latin-1.

Fixes #55

flask_debugtoolbar/__init__.py

@@ -1,15 +1,14 @@
 import os
 
-from flask import current_app, request, g
+from flask import Blueprint, current_app, request, g, send_from_directory
 from flask.globals import _request_ctx_stack
-from flask import send_from_directory
 from jinja2 import Environment, PackageLoader
 from werkzeug.exceptions import HTTPException
 from werkzeug.urls import url_quote_plus
 
-from flask_debugtoolbar.toolbar import DebugToolbar
 from flask_debugtoolbar.compat import iteritems
-from flask import Blueprint
+from flask_debugtoolbar.toolbar import DebugToolbar
+from flask_debugtoolbar.utils import decode_text
 
 
 module = Blueprint('debugtoolbar', __name__)
@@ -30,10 +29,7 @@ def replace_insensitive(string, target, replacement):
 
 def _printable(value):
     try:
-        value = repr(value)
-        if isinstance(value, bytes):
-            value = value.decode('ascii', 'replace')
-        return value
+        return decode_text(repr(value))
     except Exception as e:
         return '<repr(%s) raised %s: %s>' % (
                object.__repr__(value), type(e).__name__, e)

flask_debugtoolbar/utils.py

@@ -26,7 +26,8 @@ def format_fname(value):
     # If the file is absolute, try normalizing it relative to the project root
     # to handle it as a project file
     if os.path.isabs(value):
-        value = _shortest_relative_path(value, [current_app.root_path], os.path)
+        value = _shortest_relative_path(
+            value, [current_app.root_path], os.path)
 
     # If the value is a relative path, it is a project file
     if not os.path.isabs(value):
@@ -53,11 +54,24 @@ def _relative_paths(value, paths, path_module):
             yield relval
 
 
+def decode_text(value):
+    """
+        Decode a text-like value for display.
+
+        Unicode values are returned unchanged. Byte strings will be decoded
+        with a text-safe replacement for unrecognized characters.
+    """
+    if isinstance(value, bytes):
+        return value.decode('ascii', 'replace')
+    else:
+        return value
+
+
 def format_sql(query, args):
     if not HAVE_PYGMENTS:
-        return query
+        return decode_text(query)
 
     return Markup(highlight(
         query,
-        SqlLexer(encoding='utf-8'),
-        HtmlFormatter(encoding='utf-8', noclasses=True, style=PYGMENT_STYLE)))
+        SqlLexer(),
+        HtmlFormatter(noclasses=True, style=PYGMENT_STYLE)))

test/test_utils.py

@@ -3,7 +3,10 @@
 import posixpath
 import ntpath
 
-from flask_debugtoolbar.utils import _relative_paths, _shortest_relative_path
+from flask import Markup
+
+from flask_debugtoolbar.utils import (_relative_paths, _shortest_relative_path,
+   format_sql, decode_text, HAVE_PYGMENTS)
 
 
 @pytest.mark.parametrize('value,paths,expected,path_module', [
@@ -48,4 +51,70 @@ def test_relative_paths(value, paths, expected, path_module):
     ('c:\\foo\\bar\\baz', ['c:\\foo', 'c:\\foo\\bar'], 'baz', ntpath),
 ])
 def test_shortest_relative_path(value, paths, expected, path_module):
-    assert _shortest_relative_path(value, paths, path_module) == expected
+    assert _shortest_relative_path(value, paths, path_module) == expected
+
+
+def test_decode_text_unicode():
+    value = u'\uffff'
+    decoded = decode_text(value)
+    assert decoded == value
+
+
+def test_decode_text_ascii():
+    value = 'abc'
+    assert decode_text(value.encode('ascii')) == value
+
+
+def test_decode_text_non_ascii():
+    value = b'abc \xff xyz'
+    assert isinstance(value, bytes)
+
+    decoded = decode_text(value)
+    assert not isinstance(decoded, bytes)
+
+    assert decoded.startswith('abc')
+    assert decoded.endswith('xyz')
+
+
+@pytest.fixture()
+def no_pygments(monkeypatch):
+    monkeypatch.setattr('flask_debugtoolbar.utils.HAVE_PYGMENTS', False)
+
+
+def test_format_sql_no_pygments(no_pygments):
+    sql = 'select 1'
+    assert format_sql(sql, {}) == sql
+
+
+def test_format_sql_no_pygments_non_ascii(no_pygments):
+    sql = b"select '\xff'"
+    formatted = format_sql(sql, {})
+    assert formatted.startswith(u"select '")
+
+
+def test_format_sql_no_pygments_escape_html(no_pygments):
+    sql = 'select x < 1'
+    formatted = format_sql(sql, {})
+    assert not isinstance(formatted, Markup)
+    assert Markup('%s') % formatted == 'select x &lt; 1'
+
+
+@pytest.mark.skipif(not HAVE_PYGMENTS, reason='test requires the "Pygments" library')
+def test_format_sql_pygments():
+    sql = 'select 1'
+    html = format_sql(sql, {})
+    assert isinstance(html, Markup)
+    assert html.startswith('<div')
+    assert 'select' in html
+    assert '1' in html
+
+
+@pytest.mark.skipif(not HAVE_PYGMENTS, reason='test requires the "Pygments" library')
+def test_format_sql_pygments_non_ascii():
+    sql = b"select 'abc \xff xyz'"
+    html = format_sql(sql, {})
+    assert isinstance(html, Markup)
+    assert html.startswith('<div')
+    assert 'select' in html
+    assert 'abc' in html
+    assert 'xyz' in html

tox.ini

@@ -5,5 +5,6 @@ envlist = py26,py27,py34
 deps =
     pytest
     Flask-SQLAlchemy
+    Pygments
 commands =
     py.test