From dca47ebe61e1984450556019dc5a7ddcaf7d100a Mon Sep 17 00:00:00 2001
From: Kyle Ju <kyleju@google.com>
Date: Tue, 18 Oct 2022 15:22:38 -0700
Subject: [PATCH] Unify int arg parsing (#2357)

Co-authored-by: Kyle Ju <kyleju@chromium.org>
---
 api/features_api.py            | 12 ++----------
 framework/basehandlers.py      | 15 +++++++++++++++
 framework/basehandlers_test.py | 18 ++++++++++++++++++
 3 files changed, 35 insertions(+), 10 deletions(-)

diff --git a/api/features_api.py b/api/features_api.py
index 68470684..ba597fb2 100644
--- a/api/features_api.py
+++ b/api/features_api.py
@@ -55,16 +55,8 @@ class FeaturesAPI(basehandlers.APIHandler):
 
     user_query = self.request.args.get('q', '')
     sort_spec = self.request.args.get('sort')
-    try:
-      start = int(self.request.args.get('start', 0))
-      num = int(self.request.args.get(
-          'num', search.DEFAULT_RESULTS_PER_PAGE))
-    except ValueError:
-      self.abort(400, msg='Invalid start or num')
-    if start < 0:
-      self.abort(400, msg='Parameter out of range: start')
-    if num < 0:
-      self.abort(400, msg='Parameter out of range: num')
+    num = self.get_int_arg('num', search.DEFAULT_RESULTS_PER_PAGE)
+    start = self.get_int_arg('start', 0)
 
     features_on_page, total_count = search.process_query(
         user_query, sort_spec=sort_spec, show_unlisted=show_unlisted_features,
diff --git a/framework/basehandlers.py b/framework/basehandlers.py
index b6ca06d6..22fb1a60 100644
--- a/framework/basehandlers.py
+++ b/framework/basehandlers.py
@@ -138,6 +138,21 @@ class BaseHandler(flask.views.MethodView):
       self.abort(403, msg='Cannot view that feature')
     return feature
 
+  def get_int_arg(self, name, default=None):
+    """Get the specified integer from args."""
+    val = self.request.args.get(name, default)
+    if val is None:
+      return None
+
+    try:
+      num = int(val)
+    except ValueError:
+      self.abort(400, msg='Request parameter %r was not an int' % name)
+
+    if num < 0:
+      self.abort(400, msg='Request parameter %r out of range: %r' % (name, val))
+    return num
+
 
 class APIHandler(BaseHandler):
 
diff --git a/framework/basehandlers_test.py b/framework/basehandlers_test.py
index 31d78f33..815e9e3d 100644
--- a/framework/basehandlers_test.py
+++ b/framework/basehandlers_test.py
@@ -265,6 +265,24 @@ class BaseHandlerTests(testing_config.CustomTestCase):
       mock_abort.assert_called_once_with(
           400, msg="Parameter 'featureId' was not an int")
 
+  @mock.patch('framework.basehandlers.BaseHandler.abort')
+  def test_get_int_arg__bad(self, mock_abort):
+    mock_abort.side_effect = werkzeug.exceptions.BadRequest
+
+    with test_app.test_request_context('/test?num=abc'):
+      with self.assertRaises(werkzeug.exceptions.BadRequest):
+        self.handler.get_int_arg('num')
+      mock_abort.assert_called_once_with(
+          400, msg="Request parameter 'num' was not an int")
+
+  def test_get_int_arg(self):
+    with test_app.test_request_context('/test?num=1'):
+      actual = self.handler.get_int_arg('num')
+      self.assertEqual(1, actual)
+
+      actual = self.handler.get_int_arg('random')
+      self.assertEqual(None, actual)
+
 
 class RedirectorTests(testing_config.CustomTestCase):