Extract and host wheel METADATA on upload

tests/unit/forklift/test_legacy.py

@@ -72,19 +72,26 @@ def _get_tar_testdata(compression_type=""):
  return temp_f.getvalue()
 
 
+def _get_whl_testdata(name="fake_package", version="1.0"):
+ temp_f = io.BytesIO()
+ with zipfile.ZipFile(file=temp_f, mode="w") as zfp:
+ zfp.writestr(f"{name}-{version}.dist-info/METADATA", "Fake metadata")
+ return temp_f.getvalue()
+
+
+def _storage_hash(data):
+ return hashlib.blake2b(data, digest_size=256 // 8).hexdigest()
+
+
 _TAR_GZ_PKG_TESTDATA = _get_tar_testdata("gz")
 _TAR_GZ_PKG_MD5 = hashlib.md5(_TAR_GZ_PKG_TESTDATA).hexdigest()
 _TAR_GZ_PKG_SHA256 = hashlib.sha256(_TAR_GZ_PKG_TESTDATA).hexdigest()
-_TAR_GZ_PKG_STORAGE_HASH = hashlib.blake2b(
- _TAR_GZ_PKG_TESTDATA, digest_size=256 // 8
-).hexdigest()
+_TAR_GZ_PKG_STORAGE_HASH = _storage_hash(_TAR_GZ_PKG_TESTDATA)
 
 _TAR_BZ2_PKG_TESTDATA = _get_tar_testdata("bz2")
 _TAR_BZ2_PKG_MD5 = hashlib.md5(_TAR_BZ2_PKG_TESTDATA).hexdigest()
 _TAR_BZ2_PKG_SHA256 = hashlib.sha256(_TAR_BZ2_PKG_TESTDATA).hexdigest()
-_TAR_BZ2_PKG_STORAGE_HASH = hashlib.blake2b(
- _TAR_BZ2_PKG_TESTDATA, digest_size=256 // 8
-).hexdigest()
+_TAR_BZ2_PKG_STORAGE_HASH = _storage_hash(_TAR_BZ2_PKG_TESTDATA)
 
 
 class TestExcWithMessage:
@@ -2761,6 +2768,8 @@ def test_upload_succeeds_with_wheel(
  RoleFactory.create(user=user, project=project)
 
  filename = f"{project.name}-{release.version}-cp34-none-{plat}.whl"
+ filebody = _get_whl_testdata(project.name)
+ file_storage_hash = _storage_hash(filebody)
 
  pyramid_config.testing_securitypolicy(identity=user)
  db_request.user = user
@@ -2772,19 +2781,22 @@ def test_upload_succeeds_with_wheel(
  "version": release.version,
  "filetype": "bdist_wheel",
  "pyversion": "cp34",
- "md5_digest": _TAR_GZ_PKG_MD5,
+ "md5_digest": hashlib.md5(filebody).hexdigest(),
  "content": pretend.stub(
  filename=filename,
- file=io.BytesIO(_TAR_GZ_PKG_TESTDATA),
- type="application/tar",
+ file=io.BytesIO(filebody),
+ type="application/octet-stream",
  ),
  }
  )
 
  @pretend.call_recorder
  def storage_service_store(path, file_path, *, meta):
  with open(file_path, "rb") as fp:
- assert fp.read() == _TAR_GZ_PKG_TESTDATA
+ if file_path.endswith(".metadata"):
+ assert fp.read() == b"Fake metadata"
+ else:
+ assert fp.read() == filebody
 
  storage_service = pretend.stub(store=storage_service_store)
 
@@ -2808,9 +2820,9 @@ def storage_service_store(path, file_path, *, meta):
  pretend.call(
  "/".join(
  [
- _TAR_GZ_PKG_STORAGE_HASH[:2],
- _TAR_GZ_PKG_STORAGE_HASH[2:4],
- _TAR_GZ_PKG_STORAGE_HASH[4:],
+ file_storage_hash[:2],
+ file_storage_hash[2:4],
+ file_storage_hash[4:],
  filename,
  ]
  ),
@@ -2821,7 +2833,24 @@ def storage_service_store(path, file_path, *, meta):
  "package-type": "bdist_wheel",
  "python-version": "cp34",
  },
- )
+ ),
+ pretend.call(
+ "/".join(
+ [
+ file_storage_hash[:2],
+ file_storage_hash[2:4],
+ file_storage_hash[4:],
+ filename + ".metadata",
+ ]
+ ),
+ mock.ANY,
+ meta={
+ "project": project.normalized_name,
+ "version": release.version,
+ "package-type": "bdist_wheel",
+ "python-version": "cp34",
+ },
+ ),
  ]
 
  # Ensure that a File object has been created.
@@ -2874,6 +2903,8 @@ def test_upload_succeeds_with_wheel_after_sdist(
  RoleFactory.create(user=user, project=project)
 
  filename = f"{project.name}-{release.version}-cp34-none-any.whl"
+ filebody = _get_whl_testdata(project.name)
+ file_storage_hash = _storage_hash(filebody)
 
  pyramid_config.testing_securitypolicy(identity=user)
  db_request.user = user
@@ -2885,19 +2916,22 @@ def test_upload_succeeds_with_wheel_after_sdist(
  "version": release.version,
  "filetype": "bdist_wheel",
  "pyversion": "cp34",
- "md5_digest": "335c476dc930b959dda9ec82bd65ef19",
+ "md5_digest": hashlib.md5(filebody).hexdigest(),
  "content": pretend.stub(
  filename=filename,
- file=io.BytesIO(b"A fake file."),
- type="application/tar",
+ file=io.BytesIO(filebody),
+ type="application/zip",
  ),
  }
  )
 
  @pretend.call_recorder
  def storage_service_store(path, file_path, *, meta):
  with open(file_path, "rb") as fp:
- assert fp.read() == b"A fake file."
+ if file_path.endswith(".metadata"):
+ assert fp.read() == b"Fake metadata"
+ else:
+ assert fp.read() == filebody
 
  storage_service = pretend.stub(store=storage_service_store)
  db_request.find_service = pretend.call_recorder(
@@ -2920,9 +2954,9 @@ def storage_service_store(path, file_path, *, meta):
  pretend.call(
  "/".join(
  [
- "4e",
- "6e",
- "fa4c0ee2bbad071b4f5b5ea68f1aea89fa716e7754eb13e2314d45a5916e",
+ file_storage_hash[:2],
+ file_storage_hash[2:4],
+ file_storage_hash[4:],
  filename,
  ]
  ),
@@ -2933,7 +2967,24 @@ def storage_service_store(path, file_path, *, meta):
  "package-type": "bdist_wheel",
  "python-version": "cp34",
  },
- )
+ ),
+ pretend.call(
+ "/".join(
+ [
+ file_storage_hash[:2],
+ file_storage_hash[2:4],
+ file_storage_hash[4:],
+ filename + ".metadata",
+ ]
+ ),
+ mock.ANY,
+ meta={
+ "project": project.normalized_name,
+ "version": release.version,
+ "package-type": "bdist_wheel",
+ "python-version": "cp34",
+ },
+ ),
  ]
 
  # Ensure that a File object has been created.

warehouse/forklift/legacy.py

@@ -10,6 +10,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import base64
 import email
 import hashlib
 import hmac
@@ -790,6 +791,20 @@ def _is_duplicate_file(db_session, filename, hashes):
  return None
 
 
+def extract_wheel_metadata(path):
+ """
+ Extract METADATA file and return it as a content. The name of the
+ .whl file is used to find the corresponding .dist-info dir.
+
+ See https://www.python.org/dev/peps/pep-0658/#specification
+ """
+ filename = os.path.basename(path)
+ namever = _wheel_file_re.match(filename).group("namever")
+ metafile = namever + ".dist-info/METADATA"
+ with zipfile.ZipFile(path) as zfp:
+ return zfp.read(metafile)
+
+
 @view_config(
  route_name="forklift.legacy.file_upload",
  uses_session=True,
@@ -1317,11 +1332,19 @@ def file_upload(request):
  "Binary wheel '{filename}' has an unsupported "
  "platform tag '{plat}'.".format(filename=filename, plat=plat),
  )
+ wheel_metadata = extract_wheel_metadata(temporary_filename)
+ with open(temporary_filename + ".metadata", "wb") as fp:
+ fp.write(wheel_metadata)
+ metadata_hash = base64.b64encode(
+ hashlib.blake2s(wheel_metadata, digest_size=128 // 8).digest()
+ ).decode("utf-8")
+ else:
+ metadata_hash = None
 
  # Also buffer the entire signature file to disk.
  if "gpg_signature" in request.POST:
  has_signature = True
- with open(os.path.join(tmpdir, filename + ".asc"), "wb") as fp:
+ with open(temporary_filename + ".asc", "wb") as fp:
  signature_size = 0
  for chunk in iter(
  lambda: request.POST["gpg_signature"].file.read(8096), b""
@@ -1332,7 +1355,7 @@ def file_upload(request):
  fp.write(chunk)
 
  # Check whether signature is ASCII armored
- with open(os.path.join(tmpdir, filename + ".asc"), "rb") as fp:
+ with open(temporary_filename + ".asc", "rb") as fp:
  if not fp.read().startswith(b"-----BEGIN PGP SIGNATURE-----"):
  raise _exc_with_message(
  HTTPBadRequest, "PGP signature isn't ASCII armored."
@@ -1357,6 +1380,7 @@ def file_upload(request):
  md5_digest=file_hashes["md5"],
  sha256_digest=file_hashes["sha256"],
  blake2_256_digest=file_hashes["blake2_256"],
+ metadata_hash=metadata_hash,
  # Figure out what our filepath is going to be, we're going to use a
  # directory structure based on the hash of the file contents. This
  # will ensure that the contents of the file cannot change without
@@ -1412,18 +1436,29 @@ def file_upload(request):
  storage = request.find_service(IFileStorage, name="primary")
  storage.store(
  file_.path,
- os.path.join(tmpdir, filename),
+ temporary_filename,
  meta={
  "project": file_.release.project.normalized_name,
  "version": file_.release.version,
  "package-type": file_.packagetype,
  "python-version": file_.python_version,
  },
  )
+ if metadata_hash is not None:
+ storage.store(
+ file_.path + ".metadata",
+ temporary_filename + ".metadata",
+ meta={
+ "project": file_.release.project.normalized_name,
+ "version": file_.release.version,
+ "package-type": file_.packagetype,
+ "python-version": file_.python_version,
+ },
+ )
  if has_signature:
  storage.store(
  file_.pgp_path,
- os.path.join(tmpdir, filename + ".asc"),
+ temporary_filename + ".asc",
  meta={
  "project": file_.release.project.normalized_name,
  "version": file_.release.version,

warehouse/migrations/versions/9b9778779fe2_add_a_metadata_hash_column_to_file.py

@@ -0,0 +1,33 @@
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Add a metadata_hash column to File
+
+Revision ID: 9b9778779fe2
+Revises: d582fb87b94c
+Create Date: 2021-09-18 07:34:31.828437
+"""
+
+import sqlalchemy as sa
+
+from alembic import op
+
+revision = "9b9778779fe2"
+down_revision = "d582fb87b94c"
+
+
+def upgrade():
+ op.add_column("release_files", sa.Column("metadata_hash", sa.Text(), nullable=True))
+
+
+def downgrade():
+ op.drop_column("release_files", "metadata_hash")

warehouse/packaging/models.py

@@ -666,6 +666,7 @@ def __table_args__(cls): # noqa
  md5_digest = Column(Text, unique=True, nullable=False)
  sha256_digest = Column(CIText, unique=True, nullable=False)
  blake2_256_digest = Column(CIText, unique=True, nullable=False)
+ metadata_hash = Column(Text, unique=False, nullable=True)
  upload_time = Column(DateTime(timezone=False), server_default=func.now())
  uploaded_via = Column(Text)