model_signing.manifest

  1# Copyright 2024 The Sigstore Authors
  2#
  3# Licensed under the Apache License, Version 2.0 (the "License");
  4# you may not use this file except in compliance with the License.
  5# You may obtain a copy of the License at
  6#
  7#      http://www.apache.org/licenses/LICENSE-2.0
  8#
  9# Unless required by applicable law or agreed to in writing, software
 10# distributed under the License is distributed on an "AS IS" BASIS,
 11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 12# See the License for the specific language governing permissions and
 13# limitations under the License.
 14
 15"""An in-memory serialized representation of an ML model.
 16
 17A manifest pairs objects from the model (e.g., files, shards of files), with
 18their hashes. When signing the model we first generate a manifest from
 19serializing the model using a configured serialization method (see
 20`model_signing.signing`).
 21
 22When verifying the integrity of the model, after checking the authenticity of
 23the signature, we extract a manifest from it. Then, we serialize the local model
 24(the model being tested) and compare the two manifests.
 25
 26The serialization method used during signing must match the one used during
 27verification. We can auto detect the method to use during verification from the
 28signature, but it is recommended to be explicit when possible.
 29
 30Comparing the manifests can be done by checking that every item matches, both in
 31name and in associated hash.  In the future we will support partial object
 32matching. This is useful, for example, for the cases where the original model
 33contained files for multiple ML frameworks, but the user only uses the model
 34with one framework. This way, the user can verify the integrity only for the
 35files that are actually used.
 36
 37This API should not be used directly, we don't guarantee that it is fully stable
 38at the moment.
 39"""
 40
 41import abc
 42from collections.abc import Iterable, Iterator
 43import dataclasses
 44import pathlib
 45from typing import Any, Final, Self
 46
 47from typing_extensions import override
 48
 49from model_signing._hashing import hashing
 50
 51
 52@dataclasses.dataclass(frozen=True, order=True)
 53class _ResourceDescriptor:
 54    """A description of any content from any `Manifest`.
 55
 56    We aim this to be similar to in-toto's `ResourceDescriptor`. To support
 57    cases where in-toto cannot be directly used, we make this a dataclass that
 58    can be mapped to in-toto when needed, and used as its own otherwise.
 59
 60    Not all fields from in-toto are specified at this moment. All fields here
 61    must be present, unlike in-toto, where all are optional.
 62
 63    See github.com/in-toto/attestation/blob/main/spec/v1/resource_descriptor.md
 64    for the in-toto specification.
 65
 66    Attributes:
 67        identifier: A string that uniquely identifies this object within the
 68        manifest. Depending on serialized format, users might require the
 69        identifier to be unique across all manifests stored in a system.
 70        Producers and consumers can agree on additional requirements (e.g.,
 71        several descriptors must have a common pattern for the identifier and
 72        the integrity of the model implies integrity of all these items,
 73        ignoring any other descriptor). Corresponds to `name`, `uri`, or
 74        `content` in in-toto specification.
 75        digest: One digest for the item. Note that unlike in-toto, we only have
 76          one digest for the item and it is always required.
 77    """
 78
 79    identifier: str
 80    digest: hashing.Digest
 81
 82
 83class _ManifestKey(metaclass=abc.ABCMeta):
 84    """An object that can be a key for an item in a manifest.
 85
 86    We need to be able to convert the key to string when we serialize the
 87    manifest and to rebuild the object from the serialized representation when
 88    we deserialize the manifest from the signature.
 89    """
 90
 91    @classmethod
 92    @abc.abstractmethod
 93    def from_str(cls, s: str) -> Self:
 94        """Builds a manifest key from the string representation.
 95
 96        It is guaranteed that for any `key` of a derived class `C` and a (valid)
 97        string `s`, the following two round-trip properties hold:
 98
 99        ```
100        str(C.from_str(s)) == s
101        C.from_str(str(key)) == key
102        ```
103
104        Raises:
105            ValueError: if the string argument cannot be decoded correctly.
106        """
107
108
109@dataclasses.dataclass(frozen=True, order=True)
110class _File(_ManifestKey):
111    """A dataclass to hold information about a file as a manifest key.
112
113    Attributes:
114        path: The path to the file, relative to the model root.
115    """
116
117    path: pathlib.PurePath
118
119    def __str__(self) -> str:
120        return str(self.path)
121
122    @classmethod
123    @override
124    def from_str(cls, s: str) -> Self:
125        # Note that we always decode the string to a pure POSIX path
126        return cls(pathlib.PurePosixPath(s))
127
128
129@dataclasses.dataclass(frozen=True, order=True)
130class _Shard(_ManifestKey):
131    """A dataclass to hold information about a file shard as a manifest key.
132
133    Attributes:
134        path: The path to the file, relative to the model root.
135        start: The start offset of the shard (included).
136        end: The end offset of the shard (not included).
137    """
138
139    path: pathlib.PurePath
140    start: int
141    end: int
142
143    def __str__(self) -> str:
144        return f"{str(self.path)}:{self.start}:{self.end}"
145
146    @classmethod
147    @override
148    def from_str(cls, s: str) -> Self:
149        parts = s.split(":")
150        if len(parts) != 3:
151            raise ValueError(f"Expected 3 components separated by `:`, got {s}")
152
153        path = pathlib.PurePosixPath(parts[0])
154        start = int(parts[1])
155        end = int(parts[2])
156
157        return cls(path, start, end)
158
159
160class ManifestItem(metaclass=abc.ABCMeta):
161    """An individual object of a model, stored as an item in a manifest.
162
163    For example, this could be a file, or a file shard. All file paths are
164    relative to the model root, to allow moving or renaming the model, without
165    invalidating the signature.
166
167    The integrity of each `ManifestItem` can be verified individually and in
168    parallel. If the item is backed by a file, we recompute the hash of the
169    portion of the file that represents this item.
170
171    Attributes:
172        digest: The digest of the item. Use the `key` property to obtain a
173          canonical unique representation for the item.
174    """
175
176    digest: hashing.Digest
177
178    @property
179    @abc.abstractmethod
180    def key(self) -> _ManifestKey:
181        """A unique representation for the manifest item.
182
183        Two items in the same manifest must not share the same `key`. The
184        information contained in `key` should be sufficient to determine how to
185        compute the item's digest.
186        """
187
188
189class FileManifestItem(ManifestItem):
190    """A manifest item that records a filename path together with its digest.
191
192    Note that the path component is a `pathlib.PurePath`, relative to the model.
193    To ensure that the manifest is consistent across operating systems, we
194    convert the path to a POSIX path.
195    """
196
197    def __init__(self, *, path: pathlib.PurePath, digest: hashing.Digest):
198        """Builds a manifest item pairing a file with its digest.
199
200        Args:
201            path: The path to the file, relative to the model root.
202            digest: The digest of the file.
203        """
204        # Note: we need to force a PurePosixPath to canonicalize the manifest.
205        self._path = pathlib.PurePosixPath(path)
206        self.digest = digest
207
208    @property
209    @override
210    def key(self) -> _ManifestKey:
211        return _File(self._path)
212
213
214class ShardedFileManifestItem(ManifestItem):
215    """A manifest item that records a file shard together with its digest.
216
217    Note that the path component is a `pathlib.PurePath`, relative to the model.
218    To ensure that the manifest is consistent across operating systems, we
219    convert the path to a POSIX path.
220    """
221
222    def __init__(
223        self,
224        *,
225        path: pathlib.PurePath,
226        start: int,
227        end: int,
228        digest: hashing.Digest,
229    ):
230        """Builds a manifest item pairing a file shard with its digest.
231
232        Args:
233            path: The path to the file, relative to the model root.
234            start: The start offset of the shard (included).
235            end: The end offset of the shard (not included).
236            digest: The digest of the file shard.
237        """
238        # Note: we need to force a PurePosixPath to canonicalize the manifest.
239        self._path = pathlib.PurePosixPath(path)
240        self._start = start
241        self._end = end
242        self.digest = digest
243
244    @property
245    @override
246    def key(self) -> _ManifestKey:
247        return _Shard(self._path, self._start, self._end)
248
249
250class SerializationType(metaclass=abc.ABCMeta):
251    """A description of the serialization process that generated the manifest.
252
253    These should record all the parameters needed to ensure a reproducible
254    serialization. These are used to build a manifest back from the signature in
255    a backwards compatible way. We use these to determine what serialization to
256    use when verifying a signature.
257    """
258
259    @property
260    @abc.abstractmethod
261    def serialization_parameters(self) -> dict[str, Any]:
262        """The arguments of the serialization method."""
263
264    @classmethod
265    def from_args(cls, args: dict[str, Any]) -> Self:
266        """Builds an instance of this class based on the dict representation.
267
268        This is the reverse of `serialization_parameters`.
269
270        Args:
271            args: The arguments as a dictionary (equivalent to `**kwargs`).
272        """
273        serialization_type = args["method"]
274        for subclass in [_FileSerialization, _ShardSerialization]:
275            if serialization_type == subclass.method:
276                return subclass._from_args(args)
277        raise ValueError(f"Unknown serialization type {serialization_type}")
278
279    @classmethod
280    @abc.abstractmethod
281    def _from_args(cls, args: dict[str, Any]) -> Self:
282        """Performs the actual build from `from_dict`."""
283
284    @abc.abstractmethod
285    def new_item(self, name: str, digest: hashing.Digest) -> ManifestItem:
286        """Builds a `ManifestItem` of the correct type.
287
288        Each serialization type results in different types for the items in the
289        manifest. This method parses the `name` of the item according to the
290        serialization type to create the proper manifest item.
291
292        Args:
293            name: The name of the item, as shown in the manifest.
294            digest: The digest of the item.
295        """
296
297
298class _FileSerialization(SerializationType):
299    method: Final[str] = "files"
300
301    def __init__(
302        self,
303        hash_type: str,
304        allow_symlinks: bool = False,
305        ignore_paths: Iterable[pathlib.Path] = frozenset(),
306    ):
307        """Records the manifest serialization type for serialization by files.
308
309        We only need to record the hashing engine used and whether symlinks are
310        hashed or ignored.
311
312        Args:
313            hash_type: A string representation of the hash algorithm.
314            allow_symlinks: Controls whether symbolic links are included.
315            ignore_paths: Paths of files to ignore.
316        """
317        self._hash_type = hash_type
318        self._allow_symlinks = allow_symlinks
319        self._ignore_paths = [str(p) for p in ignore_paths]
320
321    @property
322    @override
323    def serialization_parameters(self) -> dict[str, Any]:
324        res = {
325            "method": self.method,
326            "hash_type": self._hash_type,
327            "allow_symlinks": self._allow_symlinks,
328        }
329        if self._ignore_paths:
330            res["ignore_paths"] = self._ignore_paths
331        return res
332
333    @classmethod
334    @override
335    def _from_args(cls, args: dict[str, Any]) -> Self:
336        return cls(
337            args["hash_type"],
338            args["allow_symlinks"],
339            args.get("ignore_paths", frozenset()),
340        )
341
342    @override
343    def new_item(self, name: str, digest: hashing.Digest) -> ManifestItem:
344        path = pathlib.PurePosixPath(name)
345        return FileManifestItem(path=path, digest=digest)
346
347
348class _ShardSerialization(SerializationType):
349    method: Final[str] = "shards"
350
351    def __init__(
352        self,
353        hash_type: str,
354        shard_size: int,
355        allow_symlinks: bool = False,
356        ignore_paths: Iterable[pathlib.Path] = frozenset(),
357    ):
358        """Records the manifest serialization type for serialization by shards.
359
360        We need to record the hashing engine used and whether symlinks are
361        hashed or ignored, just like for file serialization. We also need to
362        record the shard size used to split the files, since different shard
363        sizes results in different resources.
364
365        Args:
366            hash_type: A string representation of the hash algorithm.
367            shard_size: The size of each file shard.
368            allow_symlinks: Controls whether symbolic links are included.
369            ignore_paths: Paths of files to ignore.
370        """
371        self._hash_type = hash_type
372        self._allow_symlinks = allow_symlinks
373        self._shard_size = shard_size
374        self._ignore_paths = [str(p) for p in ignore_paths]
375
376    @property
377    @override
378    def serialization_parameters(self) -> dict[str, Any]:
379        res = {
380            "method": self.method,
381            "hash_type": self._hash_type,
382            "shard_size": self._shard_size,
383            "allow_symlinks": self._allow_symlinks,
384        }
385        if self._ignore_paths:
386            res["ignore_paths"] = self._ignore_paths
387        return res
388
389    @classmethod
390    @override
391    def _from_args(cls, args: dict[str, Any]) -> Self:
392        return cls(
393            args["hash_type"],
394            args["shard_size"],
395            args["allow_symlinks"],
396            args.get("ignore_paths", frozenset()),
397        )
398
399    @override
400    def new_item(self, name: str, digest: hashing.Digest) -> ManifestItem:
401        parts = name.split(":")
402        if len(parts) != 3:
403            raise ValueError(
404                "Invalid resource name: expected 3 components separated by "
405                f"`:`, got {name}"
406            )
407
408        path = pathlib.PurePosixPath(parts[0])
409        start = int(parts[1])
410        end = int(parts[2])
411        return ShardedFileManifestItem(
412            path=path, start=start, end=end, digest=digest
413        )
414
415
416class Manifest:
417    """Generic manifest file to represent a model."""
418
419    def __init__(
420        self,
421        model_name: str,
422        items: Iterable[ManifestItem],
423        serialization_type: SerializationType,
424    ):
425        """Builds a manifest from a collection of already hashed objects.
426
427        Args:
428            model_name: A name for the model that generated the manifest. This
429              is the final component of the model path, and is only informative.
430              See `model_name` property.
431            items: An iterable sequence of objects and their hashes.
432            serialization_type: The serialization method used to create items.
433        """
434        self._name = model_name
435        self._item_to_digest = {item.key: item.digest for item in items}
436        self._serialization_type = serialization_type
437
438    def __eq__(self, other: Self):
439        """Check equality based on item digests, ignoring model name."""
440        return self._item_to_digest == other._item_to_digest
441
442    def resource_descriptors(self) -> Iterator[_ResourceDescriptor]:
443        """Yields each resource from the manifest, one by one."""
444        for item, digest in sorted(self._item_to_digest.items()):
445            yield _ResourceDescriptor(identifier=str(item), digest=digest)
446
447    @property
448    def model_name(self) -> str:
449        """The name of the model when serialized (final component of the path).
450
451        This is only informative. Changing the name of the model should still
452        result in the same digests after serialization, it must not invalidate
453        signatures. As a result, two manifests with different model names but
454        with the same resource descriptors will compare equal.
455        """
456        return self._name
457
458    @property
459    def serialization_type(self) -> dict[str, Any]:
460        """The serialization (and arguments) used to build the manifest.
461
462        This is needed to record the serialization method used to generate the
463        manifest so that signature verification can use the same method.
464        """
465        return self._serialization_type.serialization_parameters