fix processing unicode symbols in query_string by Python 2 (#5552)

nikulin · carltongibson · commit 7a278b3540fc · 2017-11-06T11:46:37.000+01:00
* fix processing unicode symbols in query_string by Python 2

* Add comments for encoded test strings.

* Add file encoding for Python 2.
diff --git a/rest_framework/utils/urls.py b/rest_framework/utils/urls.py
@@ -1,3 +1,4 @@
+from django.utils.encoding import force_str
 from django.utils.six.moves.urllib import parse as urlparse
 
 
@@ -6,9 +7,9 @@ def replace_query_param(url, key, val):
     Given a URL and a key/val pair, set or replace an item in the query
     parameters of the URL, and return the new URL.
     """
-    (scheme, netloc, path, query, fragment) = urlparse.urlsplit(url)
+    (scheme, netloc, path, query, fragment) = urlparse.urlsplit(force_str(url))
     query_dict = urlparse.parse_qs(query, keep_blank_values=True)
-    query_dict[key] = [val]
+    query_dict[force_str(key)] = [force_str(val)]
     query = urlparse.urlencode(sorted(list(query_dict.items())), doseq=True)
     return urlparse.urlunsplit((scheme, netloc, path, query, fragment))
 
@@ -18,7 +19,7 @@ def remove_query_param(url, key):
     Given a URL and a key/val pair, remove an item in the query
     parameters of the URL, and return the new URL.
     """
-    (scheme, netloc, path, query, fragment) = urlparse.urlsplit(url)
+    (scheme, netloc, path, query, fragment) = urlparse.urlsplit(force_str(url))
     query_dict = urlparse.parse_qs(query, keep_blank_values=True)
     query_dict.pop(key, None)
     query = urlparse.urlencode(sorted(list(query_dict.items())), doseq=True)
diff --git a/tests/test_utils.py b/tests/test_utils.py
@@ -1,3 +1,4 @@
+# -*- coding: utf-8 -*-
 from __future__ import unicode_literals
 
 from django.conf.urls import url
@@ -11,6 +12,7 @@
 from rest_framework.serializers import ModelSerializer
 from rest_framework.utils import json
 from rest_framework.utils.breadcrumbs import get_breadcrumbs
+from rest_framework.utils.urls import remove_query_param, replace_query_param
 from rest_framework.views import APIView
 from rest_framework.viewsets import ModelViewSet
 from tests.models import BasicModel
@@ -205,3 +207,66 @@ class NonStrictJsonFloatTests(JsonFloatTests):
     """
     'STRICT_JSON = False' should not somehow affect internal json behavior
     """
+
+
+class UrlsReplaceQueryParamTests(TestCase):
+    """
+    Tests the replace_query_param functionality.
+    """
+    def test_valid_unicode_preserved(self):
+        # Encoded string: '查询'
+        q = '/?q=%E6%9F%A5%E8%AF%A2'
+        new_key = 'page'
+        new_value = 2
+        value = '%E6%9F%A5%E8%AF%A2'
+
+        assert new_key in replace_query_param(q, new_key, new_value)
+        assert value in replace_query_param(q, new_key, new_value)
+
+    def test_valid_unicode_replaced(self):
+        q = '/?page=1'
+        value = '1'
+        new_key = 'q'
+        new_value = '%E6%9F%A5%E8%AF%A2'
+
+        assert new_key in replace_query_param(q, new_key, new_value)
+        assert value in replace_query_param(q, new_key, new_value)
+
+    def test_invalid_unicode(self):
+        # Encoded string: '��<script>alert(313)</script>=1'
+        q = '/e/?%FF%FE%3C%73%63%72%69%70%74%3E%61%6C%65%72%74%28%33%31%33%29%3C%2F%73%63%72%69%70%74%3E=1'
+        key = 'from'
+        value = 'login'
+
+        assert key in replace_query_param(q, key, value)
+
+
+class UrlsRemoveQueryParamTests(TestCase):
+    """
+    Tests the remove_query_param functionality.
+    """
+    def test_valid_unicode_preserved(self):
+        q = '/?q=%E6%9F%A5%E8%AF%A2'
+        new_key = 'page'
+        new_value = 2
+        value = '%E6%9F%A5%E8%AF%A2'
+
+        assert new_key in replace_query_param(q, new_key, new_value)
+        assert value in replace_query_param(q, new_key, new_value)
+
+    def test_valid_unicode_removed(self):
+        q = '/?page=2345&q=%E6%9F%A5%E8%AF%A2'
+        key = 'page'
+        value = '2345'
+        removed_key = 'q'
+
+        assert key in remove_query_param(q, removed_key)
+        assert value in remove_query_param(q, removed_key)
+        assert '%' not in remove_query_param(q, removed_key)
+
+    def test_invalid_unicode(self):
+        q = '/?from=login&page=2&%FF%FE%3C%73%63%72%69%70%74%3E%61%6C%65%72%74%28%33%31%33%29%3C%2F%73%63%72%69%70%74%3E=1'
+        key = 'from'
+        removed_key = 'page'
+
+        assert key in remove_query_param(q, removed_key)
