From e17380706d974489a80e36e1d4d4b5902bc59e60 Mon Sep 17 00:00:00 2001
From: Oliver Chang <oliverchang@users.noreply.github.com>
Date: Thu, 17 Feb 2022 13:32:01 +1100
Subject: [PATCH] Allow querying by PURLs that have version numbers. (#303)

Part of #64.
---
 gcp/api/integration_tests.py |  9 +++++++++
 gcp/api/requirements.txt     |  1 +
 gcp/api/server.py            | 28 ++++++++++++++++++++++++++++
 3 files changed, 38 insertions(+)

diff --git a/gcp/api/integration_tests.py b/gcp/api/integration_tests.py
index d1c2319da76..24b2643933e 100644
--- a/gcp/api/integration_tests.py
+++ b/gcp/api/integration_tests.py
@@ -435,6 +435,15 @@ def test_query_purl(self):
     self.assert_results_equal({'vulns': [self._VULN_RUSTSEC_2020_0105]},
                               response.json())
 
+    response = requests.post(
+        _api() + '/v1/query',
+        data=json.dumps({'package': {
+            'purl': 'pkg:cargo/abi_stable@0.9.0',
+        }}))
+
+    self.assert_results_equal({'vulns': [self._VULN_RUSTSEC_2020_0105]},
+                              response.json())
+
 
 def print_logs(filename):
   """Print logs."""
diff --git a/gcp/api/requirements.txt b/gcp/api/requirements.txt
index 1d006519b45..d7d4df44ab3 100644
--- a/gcp/api/requirements.txt
+++ b/gcp/api/requirements.txt
@@ -9,6 +9,7 @@ google-cloud-ndb==1.7.2
 google-cloud-pubsub==2.2.0
 grpcio==1.34.0
 grpcio-tools==1.34.0
+packageurl-python==0.9.9
 packaging==20.9
 pygit2==1.4.0
 requests==2.25.1
diff --git a/gcp/api/server.py b/gcp/api/server.py
index d50139b44af..1f943e19767 100644
--- a/gcp/api/server.py
+++ b/gcp/api/server.py
@@ -23,6 +23,7 @@
 
 from google.cloud import ndb
 import grpc
+from packageurl import PackageURL
 
 import osv
 from osv import ecosystems
@@ -83,8 +84,25 @@ def QueryAffected(self, request, context):
       ecosystem = ''
       purl = ''
 
+    purl_version = None
+    if purl:
+      try:
+        parsed_purl = PackageURL.from_string(purl)
+        purl_version = parsed_purl.version
+        purl = _clean_purl(parsed_purl).to_string()
+      except ValueError:
+        context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Invalid Package URL.')
+        return None
+
     if request.query.WhichOneof('param') == 'commit':
       bugs = query_by_commit(request.query.commit, to_response=bug_to_response)
+    elif purl and purl_version:
+      bugs = query_by_version(
+          package_name,
+          ecosystem,
+          purl,
+          purl_version,
+          to_response=bug_to_response)
     elif request.query.WhichOneof('param') == 'version':
       bugs = query_by_version(
           package_name,
@@ -94,6 +112,7 @@ def QueryAffected(self, request, context):
           to_response=bug_to_response)
     else:
       context.abort(grpc.StatusCode.INVALID_ARGUMENT, 'Invalid query.')
+      return None
 
     return osv_service_v1_pb2.VulnerabilityList(vulns=bugs)
 
@@ -123,6 +142,15 @@ def _get_bugs(bug_ids, to_response=bug_to_response):
   ]
 
 
+def _clean_purl(purl):
+  """Clean a purl object."""
+  values = purl.to_dict()
+  values.pop('version', None)
+  values.pop('subpath', None)
+  values.pop('qualifiers', None)
+  return PackageURL(**values)
+
+
 def query_by_commit(commit, to_response=bug_to_response):
   """Query by commit."""
   query = osv.AffectedCommit.query(osv.AffectedCommit.commit == commit,