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