1pub const HEADER_LEN: usize = 72;
59
60pub const PLANE_RECORD_LEN: usize = 56;
65
66pub const SUPPORTED_REQUIRED_MASK: u64 = 0;
70
71pub const HEADER_PID_OFFSET: usize = 28;
76
77pub const D3D11_HANDLE_BYTES: usize = 24;
86
87#[cfg(unix)]
92type BlobFd = std::os::fd::RawFd;
93#[cfg(not(unix))]
94type BlobFd = i32;
95
96#[derive(Debug, Clone, PartialEq, Eq)]
103pub enum BlobError {
104 Truncated {
106 need: usize,
108 have: usize,
110 },
111 UnsupportedRequirement(u64),
115 ForgedLength {
118 field: &'static str,
120 value: u64,
122 },
123 InvalidUtf8(&'static str),
125 Malformed(&'static str),
127}
128
129impl std::fmt::Display for BlobError {
130 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
131 match self {
132 Self::Truncated { need, have } => {
133 write!(f, "blob truncated: needs {need} bytes, buffer has {have}")
134 }
135 Self::UnsupportedRequirement(m) => write!(
136 f,
137 "blob requires unsupported features (required_mask {m:#x}); \
138 refusing rather than misreading it"
139 ),
140 Self::ForgedLength { field, value } => {
141 write!(f, "blob field `{field}` has implausible value {value}")
142 }
143 Self::InvalidUtf8(which) => write!(f, "blob string `{which}` is not valid UTF-8"),
144 Self::Malformed(why) => write!(f, "malformed blob: {why}"),
145 }
146 }
147}
148
149impl std::error::Error for BlobError {}
150
151#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
157pub struct BlobHeader {
158 pub size: u64,
162 pub required_mask: u64,
165 pub planes_bytes: u64,
167 pub storage_kind: u32,
169 pub pid: u32,
182 pub fence_fd: i32,
184 pub dtype: u32,
186 pub quant_axis: i32,
188 pub ndim: u32,
190 pub plane_count: u32,
192 pub quant_scales_len: u32,
194 pub quant_zero_points_len: u32,
196 pub strings_bytes: u32,
198 pub strides_len: u32,
204}
205
206impl BlobHeader {
207 pub fn empty() -> Self {
209 Self::default()
210 }
211
212 pub fn write_to(&self, buf: &mut [u8]) {
217 assert!(buf.len() >= HEADER_LEN, "header buffer too small");
218 buf[0..8].copy_from_slice(&self.size.to_le_bytes());
219 buf[8..16].copy_from_slice(&self.required_mask.to_le_bytes());
220 buf[16..24].copy_from_slice(&self.planes_bytes.to_le_bytes());
221 buf[24..28].copy_from_slice(&self.storage_kind.to_le_bytes());
222 buf[28..32].copy_from_slice(&self.pid.to_le_bytes());
223 buf[32..36].copy_from_slice(&self.fence_fd.to_le_bytes());
224 buf[36..40].copy_from_slice(&self.dtype.to_le_bytes());
225 buf[40..44].copy_from_slice(&self.quant_axis.to_le_bytes());
226 buf[44..48].copy_from_slice(&self.ndim.to_le_bytes());
227 buf[48..52].copy_from_slice(&self.plane_count.to_le_bytes());
228 buf[52..56].copy_from_slice(&self.quant_scales_len.to_le_bytes());
229 buf[56..60].copy_from_slice(&self.quant_zero_points_len.to_le_bytes());
230 buf[60..64].copy_from_slice(&self.strings_bytes.to_le_bytes());
231 buf[64..68].copy_from_slice(&self.strides_len.to_le_bytes());
232 buf[68..72].copy_from_slice(&0u32.to_le_bytes()); }
234}
235
236fn rd_u64(buf: &[u8], off: usize) -> u64 {
238 let mut b = [0u8; 8];
239 b.copy_from_slice(&buf[off..off + 8]);
240 u64::from_le_bytes(b)
241}
242
243fn rd_u32(buf: &[u8], off: usize) -> u32 {
245 let mut b = [0u8; 4];
246 b.copy_from_slice(&buf[off..off + 4]);
247 u32::from_le_bytes(b)
248}
249
250pub fn parse_header(buf: &[u8]) -> Result<BlobHeader, BlobError> {
261 if buf.len() < HEADER_LEN {
262 return Err(BlobError::Truncated {
263 need: HEADER_LEN,
264 have: buf.len(),
265 });
266 }
267 let required_mask = rd_u64(buf, 8);
268 if required_mask & !SUPPORTED_REQUIRED_MASK != 0 {
272 return Err(BlobError::UnsupportedRequirement(required_mask));
273 }
274 let size = rd_u64(buf, 0);
275 if size < HEADER_LEN as u64 {
276 return Err(BlobError::ForgedLength {
277 field: "size",
278 value: size,
279 });
280 }
281 if size > buf.len() as u64 {
284 return Err(BlobError::Truncated {
285 need: size as usize,
286 have: buf.len(),
287 });
288 }
289 if rd_u32(buf, 68) != 0 {
292 return Err(BlobError::Malformed("reserved header word is not zero"));
293 }
294 Ok(BlobHeader {
295 size,
296 required_mask,
297 planes_bytes: rd_u64(buf, 16),
298 storage_kind: rd_u32(buf, 24),
299 pid: rd_u32(buf, 28),
300 fence_fd: rd_u32(buf, 32) as i32,
301 dtype: rd_u32(buf, 36),
302 quant_axis: rd_u32(buf, 40) as i32,
303 ndim: rd_u32(buf, 44),
304 plane_count: rd_u32(buf, 48),
305 quant_scales_len: rd_u32(buf, 52),
306 quant_zero_points_len: rd_u32(buf, 56),
307 strings_bytes: rd_u32(buf, 60),
308 strides_len: rd_u32(buf, 64),
309 })
310}
311
312#[derive(Debug, Clone, Copy, PartialEq, Eq)]
317pub struct Regions {
318 pub shape: usize,
320 pub strides: usize,
322 pub quant_scales: usize,
324 pub quant_zero_points: usize,
326 pub strings: usize,
328 pub planes: usize,
330 pub end: usize,
332}
333
334fn advance(
341 off: usize,
342 count: u32,
343 elem: usize,
344 limit: usize,
345 field: &'static str,
346) -> Result<usize, BlobError> {
347 let bytes = (count as usize)
348 .checked_mul(elem)
349 .ok_or(BlobError::ForgedLength {
350 field,
351 value: count as u64,
352 })?;
353 let next = off.checked_add(bytes).ok_or(BlobError::ForgedLength {
354 field,
355 value: count as u64,
356 })?;
357 if next > limit {
358 return Err(BlobError::ForgedLength {
359 field,
360 value: count as u64,
361 });
362 }
363 Ok(next)
364}
365
366pub fn region_offsets(h: &BlobHeader) -> Result<Regions, BlobError> {
374 let limit = h.size as usize;
375 if limit < HEADER_LEN {
376 return Err(BlobError::ForgedLength {
377 field: "size",
378 value: h.size,
379 });
380 }
381 if h.strides_len != 0 && h.strides_len != h.ndim {
384 return Err(BlobError::ForgedLength {
385 field: "strides_len",
386 value: h.strides_len as u64,
387 });
388 }
389 let shape = HEADER_LEN;
390 let strides = advance(shape, h.ndim, 8, limit, "ndim")?;
391 let quant_scales = advance(strides, h.strides_len, 8, limit, "strides_len")?;
392 let quant_zero_points = advance(
393 quant_scales,
394 h.quant_scales_len,
395 4,
396 limit,
397 "quant_scales_len",
398 )?;
399 let strings = advance(
400 quant_zero_points,
401 h.quant_zero_points_len,
402 4,
403 limit,
404 "quant_zero_points_len",
405 )?;
406 let planes = advance(strings, h.strings_bytes, 1, limit, "strings_bytes")?;
407 if !planes.is_multiple_of(8) {
411 return Err(BlobError::Malformed(
412 "planes region is not 8-aligned; strings_bytes must include padding",
413 ));
414 }
415 let planes_bytes = usize::try_from(h.planes_bytes).map_err(|_| BlobError::ForgedLength {
416 field: "planes_bytes",
417 value: h.planes_bytes,
418 })?;
419 let end = planes
420 .checked_add(planes_bytes)
421 .ok_or(BlobError::ForgedLength {
422 field: "planes_bytes",
423 value: h.planes_bytes,
424 })?;
425 if end > limit {
426 return Err(BlobError::ForgedLength {
427 field: "planes_bytes",
428 value: h.planes_bytes,
429 });
430 }
431 let max_planes = planes_bytes / PLANE_RECORD_LEN;
435 if h.plane_count as usize > max_planes {
436 return Err(BlobError::ForgedLength {
437 field: "plane_count",
438 value: h.plane_count as u64,
439 });
440 }
441 Ok(Regions {
442 shape,
443 strides,
444 quant_scales,
445 quant_zero_points,
446 strings,
447 planes,
448 end,
449 })
450}
451
452#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
463pub struct BlobStrings<'a> {
464 pub format: &'a str,
466 pub color_space: &'a str,
468 pub color_transfer: &'a str,
470 pub color_encoding: &'a str,
472 pub color_range: &'a str,
474}
475
476const STRING_COUNT: usize = 5;
478
479pub fn strings_encoded_len(s: &BlobStrings<'_>) -> usize {
482 let raw: usize = STRING_COUNT * 4
483 + s.format.len()
484 + s.color_space.len()
485 + s.color_transfer.len()
486 + s.color_encoding.len()
487 + s.color_range.len();
488 raw.next_multiple_of(8)
489}
490
491pub fn write_strings(s: &BlobStrings<'_>, out: &mut Vec<u8>) {
498 let start = out.len();
499 for part in [
500 s.format,
501 s.color_space,
502 s.color_transfer,
503 s.color_encoding,
504 s.color_range,
505 ] {
506 out.extend_from_slice(&(part.len() as u32).to_le_bytes());
507 out.extend_from_slice(part.as_bytes());
508 }
509 while !(out.len() - start).is_multiple_of(8) {
510 out.push(0);
511 }
512}
513
514pub fn parse_strings(buf: &[u8]) -> Result<BlobStrings<'_>, BlobError> {
518 const NAMES: [&str; STRING_COUNT] = [
519 "format",
520 "color_space",
521 "color_transfer",
522 "color_encoding",
523 "color_range",
524 ];
525 let mut out: [&str; STRING_COUNT] = [""; STRING_COUNT];
526 let mut off = 0usize;
527 for (i, name) in NAMES.iter().enumerate() {
528 let name = *name;
529 let end = off
530 .checked_add(4)
531 .ok_or(BlobError::Malformed("string count overflows"))?;
532 if end > buf.len() {
533 return Err(BlobError::Truncated {
534 need: end,
535 have: buf.len(),
536 });
537 }
538 let len = rd_u32(buf, off) as usize;
539 off = end;
540 let stop = off.checked_add(len).ok_or(BlobError::ForgedLength {
541 field: name,
542 value: len as u64,
543 })?;
544 if stop > buf.len() {
545 return Err(BlobError::ForgedLength {
546 field: name,
547 value: len as u64,
548 });
549 }
550 out[i] = std::str::from_utf8(&buf[off..stop]).map_err(|_| BlobError::InvalidUtf8(name))?;
553 off = stop;
554 }
555 Ok(BlobStrings {
556 format: out[0],
557 color_space: out[1],
558 color_transfer: out[2],
559 color_encoding: out[3],
560 color_range: out[4],
561 })
562}
563
564#[derive(Debug, Clone, Copy, PartialEq, Eq)]
579pub struct BlobPlane<'a> {
580 pub handle: i64,
600 pub offset: u64,
602 pub stride: u64,
604 pub size: u64,
606 pub used: u64,
608 pub modifier: u64,
611 pub handle_bytes: &'a [u8],
623 pub data: &'a [u8],
625}
626
627impl BlobPlane<'_> {
628 pub fn is_inline(&self) -> bool {
630 self.handle < 0
631 }
632
633 pub fn encoded_len(&self) -> usize {
636 (PLANE_RECORD_LEN + self.handle_bytes.len() + self.data.len()).next_multiple_of(8)
637 }
638}
639
640pub fn write_plane(p: &BlobPlane<'_>, out: &mut Vec<u8>) {
642 let start = out.len();
643 out.extend_from_slice(&p.handle.to_le_bytes());
644 out.extend_from_slice(&p.offset.to_le_bytes());
645 out.extend_from_slice(&p.stride.to_le_bytes());
646 out.extend_from_slice(&p.size.to_le_bytes());
647 out.extend_from_slice(&p.used.to_le_bytes());
648 out.extend_from_slice(&p.modifier.to_le_bytes());
649 out.extend_from_slice(&(p.handle_bytes.len() as u32).to_le_bytes());
650 out.extend_from_slice(&(p.data.len() as u32).to_le_bytes());
651 out.extend_from_slice(p.handle_bytes);
652 out.extend_from_slice(p.data);
653 while !(out.len() - start).is_multiple_of(8) {
654 out.push(0);
655 }
656}
657
658fn rd_i64(buf: &[u8], off: usize) -> i64 {
660 let mut b = [0u8; 8];
661 b.copy_from_slice(&buf[off..off + 8]);
662 i64::from_le_bytes(b)
663}
664
665pub fn parse_planes(buf: &[u8], count: u32) -> Result<Vec<BlobPlane<'_>>, BlobError> {
672 let mut planes = Vec::with_capacity(count as usize);
673 let mut off = 0usize;
674 for _ in 0..count {
675 let head = off
676 .checked_add(PLANE_RECORD_LEN)
677 .ok_or(BlobError::Malformed("plane record offset overflows"))?;
678 if head > buf.len() {
679 return Err(BlobError::Truncated {
680 need: head,
681 have: buf.len(),
682 });
683 }
684 let hb_len = rd_u32(buf, off + 48) as usize;
685 let data_len = rd_u32(buf, off + 52) as usize;
686 let hb_end = head.checked_add(hb_len).ok_or(BlobError::ForgedLength {
687 field: "handle_bytes_len",
688 value: hb_len as u64,
689 })?;
690 let data_end = hb_end
691 .checked_add(data_len)
692 .ok_or(BlobError::ForgedLength {
693 field: "data_len",
694 value: data_len as u64,
695 })?;
696 if hb_end > buf.len() {
697 return Err(BlobError::ForgedLength {
698 field: "handle_bytes_len",
699 value: hb_len as u64,
700 });
701 }
702 if data_end > buf.len() {
703 return Err(BlobError::ForgedLength {
704 field: "data_len",
705 value: data_len as u64,
706 });
707 }
708 planes.push(BlobPlane {
709 handle: rd_i64(buf, off),
710 offset: rd_u64(buf, off + 8),
711 stride: rd_u64(buf, off + 16),
712 size: rd_u64(buf, off + 24),
713 used: rd_u64(buf, off + 32),
714 modifier: rd_u64(buf, off + 40),
715 handle_bytes: &buf[head..hb_end],
716 data: &buf[hb_end..data_end],
717 });
718 off = data_end.next_multiple_of(8);
719 }
720 Ok(planes)
721}
722
723#[derive(Debug, Clone, Copy, PartialEq, Eq)]
729pub enum TransportMode {
730 Reference,
733 Inline,
737}
738
739#[derive(Debug, Clone, Copy)]
743pub struct BlobView<'a> {
744 buf: &'a [u8],
745 header: BlobHeader,
746 regions: Regions,
747}
748
749impl<'a> BlobView<'a> {
750 pub fn parse(buf: &'a [u8]) -> Result<Self, BlobError> {
756 let header = parse_header(buf)?;
757 let regions = region_offsets(&header)?;
758 Ok(Self {
759 buf,
760 header,
761 regions,
762 })
763 }
764
765 pub fn header(&self) -> BlobHeader {
767 self.header
768 }
769
770 pub fn shape(&self) -> Vec<u64> {
772 (0..self.header.ndim as usize)
773 .map(|i| rd_u64(self.buf, self.regions.shape + i * 8))
774 .collect()
775 }
776
777 pub fn strides(&self) -> Vec<i64> {
779 (0..self.header.strides_len as usize)
780 .map(|i| rd_i64(self.buf, self.regions.strides + i * 8))
781 .collect()
782 }
783
784 pub fn strings(&self) -> BlobStrings<'a> {
786 parse_strings(&self.buf[self.regions.strings..self.regions.planes]).unwrap_or_default()
787 }
788
789 pub fn planes(&self) -> Result<Vec<BlobPlane<'a>>, BlobError> {
791 parse_planes(
792 &self.buf[self.regions.planes..self.regions.end],
793 self.header.plane_count,
794 )
795 }
796}
797
798pub fn export(t: &crate::TensorDyn, mode: TransportMode) -> crate::Result<(Vec<u8>, Vec<BlobFd>)> {
815 let fmt = t.format();
816 let shape = export_addressing_shape(t, fmt)?;
819 let esz = t.dtype().size() as i64;
820 let strides = export_c_strides(t, fmt, &shape, esz);
821
822 let quant = t.quantization();
823 let (quant_axis, scales, zeros): (i32, &[f32], &[i32]) = match quant {
824 None => (-2, &[], &[]),
825 Some(q) => (
826 q.axis().map(|a| a as i32).unwrap_or(-1),
827 q.scale(),
828 q.zero_point().unwrap_or(&[]),
829 ),
830 };
831
832 let colorimetry = t.colorimetry().unwrap_or_default();
833 let strings = BlobStrings {
834 format: fmt.map(|f| f.as_str()).unwrap_or(""),
835 color_space: colorimetry.space.map(|v| v.as_str()).unwrap_or(""),
836 color_transfer: colorimetry.transfer.map(|v| v.as_str()).unwrap_or(""),
837 color_encoding: colorimetry.encoding.map(|v| v.as_str()).unwrap_or(""),
838 color_range: colorimetry.range.map(|v| v.as_str()).unwrap_or(""),
839 };
840
841 let geoms = export_plane_geoms(t, fmt, esz)?;
844
845 let pin = match mode {
848 TransportMode::Inline => Some(t.pin_host(crate::CpuAccess::Read)?),
849 TransportMode::Reference => None,
850 };
851 let all_bytes: &[u8] = match &pin {
852 Some(p) => unsafe { std::slice::from_raw_parts(p.as_mut_ptr(), p.len()) },
853 None => &[],
854 };
855
856 let (handle, handle_bytes, fds): (i64, Vec<u8>, Vec<BlobFd>) = match mode {
860 TransportMode::Inline => (-1, Vec::new(), Vec::new()),
861 TransportMode::Reference => {
862 let raw = reference_handle(t)?;
863 #[cfg(target_os = "windows")]
864 {
865 (0, reference_handle_bytes(t, raw), Vec::new())
870 }
871 #[cfg(not(target_os = "windows"))]
872 {
873 let raw = i32::try_from(raw).map_err(|_| {
874 crate::Error::InvalidOperation("handle is not representable as an fd".into())
875 })?;
876 (0, Vec::new(), vec![raw])
878 }
879 }
880 };
881
882 let mut planes_buf = Vec::new();
883 write_export_planes(
884 &geoms,
885 mode,
886 handle,
887 &handle_bytes,
888 all_bytes,
889 &mut planes_buf,
890 )?;
891
892 let strings_bytes = strings_encoded_len(&strings);
893 let header = BlobHeader {
894 size: 0, required_mask: 0,
896 planes_bytes: planes_buf.len() as u64,
897 storage_kind: t.memory().code(),
898 pid: match mode {
900 TransportMode::Inline => 0,
901 TransportMode::Reference => std::process::id(),
902 },
903 fence_fd: -1,
906 dtype: t.dtype().code(),
907 quant_axis,
908 ndim: shape.len() as u32,
909 plane_count: geoms.len() as u32,
910 quant_scales_len: scales.len() as u32,
911 quant_zero_points_len: zeros.len() as u32,
912 strings_bytes: strings_bytes as u32,
913 strides_len: strides.len() as u32,
914 };
915
916 let mut out = vec![0u8; HEADER_LEN];
917 for d in &shape {
918 out.extend_from_slice(&d.to_le_bytes());
919 }
920 for st in &strides {
921 out.extend_from_slice(&st.to_le_bytes());
922 }
923 for sc in scales {
924 out.extend_from_slice(&sc.to_le_bytes());
925 }
926 for z in zeros {
927 out.extend_from_slice(&z.to_le_bytes());
928 }
929 write_strings(&strings, &mut out);
930 out.extend_from_slice(&planes_buf);
931
932 let mut header = header;
933 header.size = out.len() as u64;
934 header.write_to(&mut out);
935 Ok((out, fds))
936}
937
938fn export_addressing_shape(
939 t: &crate::TensorDyn,
940 fmt: Option<crate::PixelFormat>,
941) -> crate::Result<Vec<u64>> {
942 match fmt {
943 Some(f) => {
944 let (w, h) = image_dims(t, f)?;
945 Ok(f.addressing_shape(w, h)
946 .ok_or_else(|| {
947 crate::Error::InvalidArgument(format!("no addressing shape for {f:?}"))
948 })?
949 .iter()
950 .map(|d| *d as u64)
951 .collect())
952 }
953 None => Ok(t.shape().iter().map(|d| *d as u64).collect()),
954 }
955}
956
957fn export_c_strides(
958 t: &crate::TensorDyn,
959 fmt: Option<crate::PixelFormat>,
960 shape: &[u64],
961 esz: i64,
962) -> Vec<i64> {
963 let mut acc = esz;
964 let mut v = vec![0i64; shape.len()];
965 for i in (0..shape.len()).rev() {
966 v[i] = acc;
967 acc *= shape[i] as i64;
968 }
969 if let (Some(rs), true) = (t.row_stride(), shape.len() >= 2) {
970 let row_dim = match fmt.map(|f| f.layout()) {
971 Some(crate::PixelLayout::Planar) if shape.len() >= 3 => 1,
972 _ => 0,
973 };
974 v[row_dim] = rs as i64;
975 }
976 v
977}
978
979fn export_plane_geoms(
980 t: &crate::TensorDyn,
981 fmt: Option<crate::PixelFormat>,
982 esz: i64,
983) -> crate::Result<Vec<crate::PlaneGeometry>> {
984 match fmt {
985 Some(f) => {
986 let (w, h) = image_dims(t, f)?;
987 let rs = t
988 .effective_row_stride()
989 .unwrap_or(w * f.channels() * esz as usize);
990 f.plane_table(w, h, rs)
991 .ok_or_else(|| crate::Error::InvalidArgument(format!("no plane table for {f:?}")))
992 }
993 None => Ok(vec![crate::PlaneGeometry {
994 offset: 0,
995 stride: t.capacity_bytes() as u64,
996 size: t.capacity_bytes() as u64,
997 }]),
998 }
999}
1000
1001fn write_export_planes(
1002 geoms: &[crate::PlaneGeometry],
1003 mode: TransportMode,
1004 handle: i64,
1005 handle_bytes: &[u8],
1006 all_bytes: &[u8],
1007 planes_buf: &mut Vec<u8>,
1008) -> crate::Result<()> {
1009 for g in geoms {
1010 let data: &[u8] = match mode {
1011 TransportMode::Inline => {
1012 let start = g.offset as usize;
1013 let end = start
1014 .checked_add(g.size as usize)
1015 .ok_or(crate::Error::InvalidSize(g.size as usize))?;
1016 all_bytes
1017 .get(start..end)
1018 .ok_or(crate::Error::InvalidSize(end))?
1019 }
1020 TransportMode::Reference => &[],
1021 };
1022 write_plane(
1023 &BlobPlane {
1024 handle,
1025 offset: if mode == TransportMode::Inline {
1028 0
1029 } else {
1030 g.offset
1031 },
1032 stride: g.stride,
1033 size: if mode == TransportMode::Inline {
1035 data.len() as u64
1036 } else {
1037 g.size
1038 },
1039 used: if mode == TransportMode::Inline {
1040 data.len() as u64
1041 } else {
1042 g.size
1043 },
1044 modifier: 0,
1045 handle_bytes,
1050 data,
1051 },
1052 planes_buf,
1053 );
1054 }
1055 Ok(())
1056}
1057
1058fn image_dims(t: &crate::TensorDyn, f: crate::PixelFormat) -> crate::Result<(usize, usize)> {
1063 let shape = t.shape();
1064 let bad = || crate::Error::InvalidShape(format!("shape {shape:?} is not an image for {f:?}"));
1065 match f.layout() {
1066 crate::PixelLayout::Packed => {
1067 if shape.len() < 2 {
1068 return Err(bad());
1069 }
1070 Ok((shape[1], shape[0]))
1071 }
1072 crate::PixelLayout::Planar => {
1073 if shape.len() < 3 {
1074 return Err(bad());
1075 }
1076 Ok((shape[2], shape[1]))
1077 }
1078 crate::PixelLayout::SemiPlanar => {
1079 if shape.len() < 2 {
1080 return Err(bad());
1081 }
1082 let total = shape[0];
1086 let h = (0..=total)
1087 .find(|&h| f.combined_plane_height(h) == Some(total))
1088 .ok_or_else(bad)?;
1089 Ok((shape[1], h))
1090 }
1091 }
1092}
1093
1094fn reference_handle(t: &crate::TensorDyn) -> crate::Result<i64> {
1097 match t.memory() {
1098 #[cfg(target_os = "linux")]
1099 crate::TensorMemory::DmaBuf => {
1100 use std::os::fd::AsRawFd;
1101 t.dmabuf().map(|fd| fd.as_raw_fd() as i64).map_err(|e| {
1104 crate::Error::InvalidOperation(format!("dma-buf tensor has no fd: {e}"))
1105 })
1106 }
1107 #[cfg(any(target_os = "macos", target_os = "ios"))]
1108 crate::TensorMemory::DmaBuf => t
1109 .iosurface_id()
1110 .map(|id| id as i64)
1111 .ok_or_else(|| crate::Error::InvalidOperation("IOSurface tensor has no id".into())),
1112 #[cfg(target_os = "windows")]
1117 crate::TensorMemory::DmaBuf => {
1118 t.d3d11_shared_handle_value()
1119 .map(|h| h as i64)
1120 .ok_or_else(|| {
1121 crate::Error::InvalidOperation(
1122 "D3D11 texture tensor has no shared handle".into(),
1123 )
1124 })
1125 }
1126 other => Err(crate::Error::InvalidOperation(format!(
1130 "{other:?} has no shareable handle; export it with TransportMode::Inline"
1131 ))),
1132 }
1133}
1134
1135#[cfg(target_os = "windows")]
1150fn reference_handle_bytes(t: &crate::TensorDyn, raw: i64) -> Vec<u8> {
1151 let fence = crate::d3d11::device()
1152 .map(|d| d.exported_fence_handle_value())
1153 .unwrap_or(0);
1154 let (fence, value) = reference_completion(fence, t.gpu_write_value());
1155 let mut out = Vec::with_capacity(D3D11_HANDLE_BYTES);
1156 out.extend_from_slice(&(raw as u64).to_le_bytes());
1157 out.extend_from_slice(&fence.to_le_bytes());
1158 out.extend_from_slice(&value.to_le_bytes());
1159 out
1160}
1161
1162#[cfg(target_os = "windows")]
1171fn reference_completion(fence: usize, recorded: u64) -> (u64, u64) {
1172 if fence == 0 {
1173 (0, 0)
1174 } else {
1175 (fence as u64, recorded)
1176 }
1177}
1178
1179pub fn import(blob: &[u8], fds: &[BlobFd]) -> crate::Result<crate::TensorDyn> {
1201 #[cfg(not(unix))]
1202 let _ = fds;
1203 let v = BlobView::parse(blob).map_err(|e| crate::Error::InvalidArgument(e.to_string()))?;
1204 let h = v.header();
1205 let planes = v
1206 .planes()
1207 .map_err(|e| crate::Error::InvalidArgument(e.to_string()))?;
1208
1209 let inline = planes.first().map(BlobPlane::is_inline).unwrap_or(true);
1213 validate_import_planes(&planes, inline)?;
1214
1215 let dtype = crate::protocol::dtype_to_dtype(h.dtype)
1216 .ok_or_else(|| crate::Error::NotImplemented(format!("blob dtype code {}", h.dtype)))?;
1217 let strings = v.strings();
1218 let format = if strings.format.is_empty() {
1219 None
1220 } else {
1221 Some(
1222 crate::PixelFormat::from_str_code(strings.format).ok_or_else(|| {
1223 crate::Error::NotImplemented(format!("blob format {:?}", strings.format))
1224 })?,
1225 )
1226 };
1227
1228 if !inline {
1229 #[cfg(target_os = "windows")]
1230 {
1231 return import_referenced_d3d11_blob(&v, &h, &planes, format, &strings, dtype);
1232 }
1233 #[cfg(not(any(unix, target_os = "windows")))]
1234 {
1235 return Err(crate::Error::NotImplemented(
1236 "reference-mode blob import is not available on this platform".into(),
1237 ));
1238 }
1239 #[cfg(unix)]
1240 {
1241 return import_referenced_blob(&v, &h, &planes, fds, format, &strings, dtype);
1242 }
1243 }
1244
1245 let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
1249 let alloc_shape = alloc_shape_from_grid(format, &grid)?;
1250
1251 let mut t = crate::TensorDyn::new(&alloc_shape, dtype, Some(crate::TensorMemory::Mem), None)?;
1252 if let Some(f) = format {
1253 t.set_format(f)?;
1254 }
1255
1256 blit_inline_planes(&mut t, &planes, format, &grid, dtype)?;
1261
1262 t.set_colorimetry(colorimetry_from(&strings));
1263 if h.quant_axis != -2 {
1264 let scales: Vec<f32> = (0..h.quant_scales_len as usize)
1265 .map(|i| {
1266 let mut b = [0u8; 4];
1267 let off = v.regions.quant_scales + i * 4;
1268 b.copy_from_slice(&blob[off..off + 4]);
1269 f32::from_le_bytes(b)
1270 })
1271 .collect();
1272 let zeros: Vec<i32> = (0..h.quant_zero_points_len as usize)
1273 .map(|i| {
1274 let off = v.regions.quant_zero_points + i * 4;
1275 rd_u32(blob, off) as i32
1276 })
1277 .collect();
1278 if let Some(q) = quantization_from(h.quant_axis, &scales, &zeros) {
1279 t.set_quantization(q)?;
1280 }
1281 }
1282 Ok(t)
1283}
1284
1285fn validate_import_planes(planes: &[BlobPlane<'_>], inline: bool) -> crate::Result<()> {
1286 if planes.iter().any(|p| p.is_inline() != inline) {
1287 return Err(crate::Error::InvalidArgument(
1288 "blob mixes inline and referenced planes; a tensor has one transport mode".into(),
1289 ));
1290 }
1291 for p in planes {
1292 if p.is_inline() {
1293 if p.size as usize != p.data.len() {
1295 return Err(crate::Error::InvalidArgument(format!(
1296 "inline plane declares size {} but carries {} bytes",
1297 p.size,
1298 p.data.len()
1299 )));
1300 }
1301 if p.modifier != 0 || !p.handle_bytes.is_empty() {
1302 return Err(crate::Error::InvalidArgument(
1303 "inline plane must have modifier 0 and no handle_bytes".into(),
1304 ));
1305 }
1306 } else if !p.data.is_empty() {
1307 return Err(crate::Error::InvalidArgument(
1308 "referenced plane must not also carry inline bytes".into(),
1309 ));
1310 }
1311 }
1312 Ok(())
1313}
1314
1315fn alloc_shape_from_grid(
1316 format: Option<crate::PixelFormat>,
1317 grid: &[usize],
1318) -> crate::Result<Vec<usize>> {
1319 match format {
1320 Some(f) => {
1321 let (w, h_px) = dims_from_grid(f, grid)?;
1322 f.allocation_shape(w, h_px)
1323 .ok_or_else(|| crate::Error::InvalidShape(format!("no allocation shape for {f:?}")))
1324 }
1325 None => Ok(grid.to_vec()),
1326 }
1327}
1328
1329#[cfg(unix)]
1330fn import_referenced_blob(
1331 v: &BlobView<'_>,
1332 h: &BlobHeader,
1333 planes: &[BlobPlane<'_>],
1334 fds: &[BlobFd],
1335 format: Option<crate::PixelFormat>,
1336 strings: &BlobStrings<'_>,
1337 dtype: crate::DType,
1338) -> crate::Result<crate::TensorDyn> {
1339 let first = planes
1344 .first()
1345 .ok_or_else(|| crate::Error::InvalidArgument("referenced blob carries no planes".into()))?;
1346 let idx = usize::try_from(first.handle).map_err(|_| {
1347 crate::Error::InvalidArgument(format!("plane handle {} is not an index", first.handle))
1348 })?;
1349 let raw = *fds.get(idx).ok_or_else(|| {
1350 crate::Error::InvalidArgument(format!(
1351 "plane handle indexes fd {idx} but only {} were provided",
1352 fds.len()
1353 ))
1354 })?;
1355 for p in planes {
1356 let i = usize::try_from(p.handle).unwrap_or(usize::MAX);
1357 if i >= fds.len() {
1358 return Err(crate::Error::InvalidArgument(format!(
1359 "plane handle indexes fd {i} but only {} were provided",
1360 fds.len()
1361 )));
1362 }
1363 }
1364
1365 let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
1366 let alloc_shape = alloc_shape_from_grid(format, &grid)?;
1367 let mut t = open_referenced_tensor(h.storage_kind, raw, &alloc_shape, dtype)?;
1368 if let Some(f) = format {
1369 t.set_format(f)?;
1370 }
1371 if first.stride > 0 {
1372 let _ = t.set_row_stride(first.stride as usize);
1374 }
1375 t.set_colorimetry(colorimetry_from(strings));
1376 Ok(t)
1377}
1378
1379#[cfg(unix)]
1380fn open_referenced_tensor(
1381 storage_kind: u32,
1382 raw: BlobFd,
1383 alloc_shape: &[usize],
1384 dtype: crate::DType,
1385) -> crate::Result<crate::TensorDyn> {
1386 match crate::TensorMemory::from_code(storage_kind) {
1392 #[cfg(any(target_os = "macos", target_os = "ios"))]
1393 Some(crate::TensorMemory::DmaBuf) | Some(crate::TensorMemory::IoSurface) => {
1394 let id = u32::try_from(raw).map_err(|_| {
1395 crate::Error::InvalidArgument(format!("{raw} is not an IOSurface id"))
1396 })?;
1397 crate::TensorDyn::from_iosurface_id(id, alloc_shape, dtype, None)
1398 }
1399 #[cfg(all(unix, not(any(target_os = "macos", target_os = "ios"))))]
1400 Some(crate::TensorMemory::DmaBuf) | Some(crate::TensorMemory::Shm) => {
1401 use std::os::fd::{BorrowedFd, OwnedFd};
1402 let borrowed = unsafe { BorrowedFd::borrow_raw(raw) };
1407 let owned: OwnedFd = borrowed
1408 .try_clone_to_owned()
1409 .map_err(crate::Error::IoError)?;
1410 crate::TensorDyn::from_fd(owned, alloc_shape, dtype, None)
1411 }
1412 other => Err(crate::Error::NotImplemented(format!(
1413 "reference-mode import for {other:?} on this platform"
1414 ))),
1415 }
1416}
1417
1418#[cfg(target_os = "windows")]
1427fn import_referenced_d3d11_blob(
1428 v: &BlobView<'_>,
1429 h: &BlobHeader,
1430 planes: &[BlobPlane<'_>],
1431 format: Option<crate::PixelFormat>,
1432 strings: &BlobStrings<'_>,
1433 dtype: crate::DType,
1434) -> crate::Result<crate::TensorDyn> {
1435 let first = planes
1436 .first()
1437 .ok_or_else(|| crate::Error::InvalidArgument("referenced blob carries no planes".into()))?;
1438 for p in planes {
1445 if p.handle != 0 || p.handle_bytes != first.handle_bytes {
1446 return Err(crate::Error::InvalidArgument(
1447 "D3D11 reference blob planes must all carry handle 0 and identical \
1448 handle bytes; they name one texture"
1449 .into(),
1450 ));
1451 }
1452 }
1453 let grid: Vec<usize> = v.shape().iter().map(|d| *d as usize).collect();
1458 let mut t = open_referenced_texture(h.storage_kind, first, h.pid, &grid, dtype, format)?;
1459 t.set_colorimetry(colorimetry_from(strings));
1466 Ok(t)
1467}
1468
1469#[cfg(target_os = "windows")]
1484fn open_referenced_texture(
1485 storage_kind: u32,
1486 first: &BlobPlane<'_>,
1487 pid: u32,
1488 grid: &[usize],
1489 dtype: crate::DType,
1490 format: Option<crate::PixelFormat>,
1491) -> crate::Result<crate::TensorDyn> {
1492 use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
1493 use windows::Win32::Foundation::HANDLE;
1494 use windows::Win32::System::Threading::{OpenProcess, PROCESS_DUP_HANDLE};
1495
1496 if storage_kind != crate::TensorMemory::DmaBuf.code() {
1497 return Err(crate::Error::InvalidArgument(format!(
1498 "storage kind {storage_kind} is not referenceable on Windows"
1499 )));
1500 }
1501 if first.handle_bytes.len() != D3D11_HANDLE_BYTES {
1502 return Err(crate::Error::InvalidArgument(format!(
1503 "D3D11 reference plane needs {D3D11_HANDLE_BYTES} handle bytes, not {}",
1504 first.handle_bytes.len()
1505 )));
1506 }
1507 let format = format.ok_or_else(|| {
1508 crate::Error::InvalidArgument("D3D11 reference blob without a pixel format".into())
1509 })?;
1510 let word = |i: usize| {
1511 let mut b = [0u8; 8];
1512 b.copy_from_slice(&first.handle_bytes[i * 8..(i + 1) * 8]);
1513 u64::from_le_bytes(b)
1514 };
1515 let handle_at = |i: usize| HANDLE(word(i) as usize as *mut std::ffi::c_void);
1516 let (mut tex, mut fence, value) = (handle_at(0), handle_at(1), word(2));
1517 if tex.is_invalid() {
1518 return Err(crate::Error::InvalidArgument(
1519 "D3D11 reference plane carries no texture handle".into(),
1520 ));
1521 }
1522 if fence.is_invalid() && value != 0 {
1528 return Err(crate::Error::InvalidArgument(format!(
1529 "D3D11 reference plane names fence value {value} with no fence handle; \
1530 the completion it names cannot be waited on"
1531 )));
1532 }
1533
1534 let mut duplicates: Vec<OwnedHandle> = Vec::new();
1538 if std::process::id() != pid {
1539 let src = unsafe { OpenProcess(PROCESS_DUP_HANDLE, false, pid) }
1543 .map_err(|e| open_process_error(pid, &e))?;
1544 let src = unsafe { OwnedHandle::from_raw_handle(src.0) };
1547 for (h, which) in [(&mut tex, "texture"), (&mut fence, "fence")] {
1548 if h.is_invalid() {
1549 continue;
1550 }
1551 let dup = duplicate_out_of(&src, *h, pid, which)?;
1552 *h = HANDLE(dup.as_raw_handle());
1553 duplicates.push(dup);
1554 }
1555 }
1556
1557 let geometry = unsafe { crate::protocol::d3d11_geometry_checked(format, tex.0, grid) };
1561 let completion = (value != 0).then_some((fence.0, value));
1564 let opened = match geometry {
1565 Ok((width, height)) => unsafe {
1568 crate::TensorDyn::from_d3d11_shared_handle(
1569 tex.0,
1570 width,
1571 height,
1572 format,
1573 dtype,
1574 crate::CpuAccess::ReadWrite,
1575 completion,
1576 None,
1577 )
1578 },
1579 Err(e) => Err(e),
1580 };
1581 drop(duplicates);
1584 opened
1585}
1586
1587#[cfg(target_os = "windows")]
1596fn win32_error_kind(e: &windows::core::Error) -> std::io::ErrorKind {
1597 use std::io::ErrorKind;
1598 use windows::core::HRESULT;
1599 use windows::Win32::Foundation::{ERROR_ACCESS_DENIED, ERROR_INVALID_PARAMETER};
1600
1601 let code = e.code();
1602 if code == HRESULT::from_win32(ERROR_ACCESS_DENIED.0) {
1603 ErrorKind::PermissionDenied
1604 } else if code == HRESULT::from_win32(ERROR_INVALID_PARAMETER.0) {
1605 ErrorKind::NotFound
1610 } else {
1611 ErrorKind::Other
1612 }
1613}
1614
1615#[cfg(target_os = "windows")]
1617fn open_process_error(pid: u32, e: &windows::core::Error) -> crate::Error {
1618 use std::io::ErrorKind;
1619
1620 let kind = win32_error_kind(e);
1621 let what = match kind {
1622 ErrorKind::PermissionDenied => {
1623 format!("cannot open exporting process {pid} to duplicate its handles: access denied")
1624 }
1625 ErrorKind::NotFound => format!("exporting process {pid} is gone"),
1626 _ => format!(
1627 "cannot open exporting process {pid} to duplicate its handles: {} \
1628 (HRESULT 0x{:08X})",
1629 e.message(),
1630 e.code().0 as u32
1631 ),
1632 };
1633 crate::Error::IoError(std::io::Error::new(kind, what))
1634}
1635
1636#[cfg(target_os = "windows")]
1640fn duplicate_out_of(
1641 src: &std::os::windows::io::OwnedHandle,
1642 handle: windows::Win32::Foundation::HANDLE,
1643 pid: u32,
1644 which: &str,
1645) -> crate::Result<std::os::windows::io::OwnedHandle> {
1646 use std::os::windows::io::{AsRawHandle, FromRawHandle, OwnedHandle};
1647 use windows::Win32::Foundation::{DuplicateHandle, DUPLICATE_SAME_ACCESS, HANDLE};
1648 use windows::Win32::System::Threading::GetCurrentProcess;
1649
1650 let mut out = HANDLE(std::ptr::null_mut());
1651 unsafe {
1654 DuplicateHandle(
1655 HANDLE(src.as_raw_handle()),
1656 handle,
1657 GetCurrentProcess(),
1658 &mut out,
1659 0,
1660 false,
1661 DUPLICATE_SAME_ACCESS,
1662 )
1663 }
1664 .map_err(|e| {
1665 crate::Error::IoError(std::io::Error::new(
1669 win32_error_kind(&e),
1670 format!(
1671 "cannot duplicate the {which} handle out of exporting process {pid}: {} \
1672 (HRESULT 0x{:08X})",
1673 e.message(),
1674 e.code().0 as u32
1675 ),
1676 ))
1677 })?;
1678 Ok(unsafe { OwnedHandle::from_raw_handle(out.0) })
1681}
1682
1683fn blit_inline_planes(
1684 t: &mut crate::TensorDyn,
1685 planes: &[BlobPlane<'_>],
1686 format: Option<crate::PixelFormat>,
1687 grid: &[usize],
1688 dtype: crate::DType,
1689) -> crate::Result<()> {
1690 let pin = t.pin_host(crate::CpuAccess::Write)?;
1691 let dst = unsafe { std::slice::from_raw_parts_mut(pin.as_mut_ptr(), pin.len()) };
1692 let geoms: Vec<crate::PlaneGeometry> = match format {
1693 Some(f) => {
1694 let (w, h_px) = dims_from_grid(f, grid)?;
1695 let rs = t
1696 .effective_row_stride()
1697 .unwrap_or(w * f.channels() * dtype.size());
1698 f.plane_table(w, h_px, rs)
1699 .ok_or_else(|| crate::Error::InvalidShape(format!("no plane table for {f:?}")))?
1700 }
1701 None => vec![crate::PlaneGeometry {
1702 offset: 0,
1703 stride: dst.len() as u64,
1704 size: dst.len() as u64,
1705 }],
1706 };
1707 if geoms.len() != planes.len() {
1708 return Err(crate::Error::InvalidArgument(format!(
1709 "blob carries {} planes but {:?} has {}",
1710 planes.len(),
1711 format,
1712 geoms.len()
1713 )));
1714 }
1715 for (g, p) in geoms.iter().zip(planes) {
1716 let start = g.offset as usize;
1717 let end = start
1718 .checked_add(p.data.len())
1719 .ok_or(crate::Error::InvalidSize(p.data.len()))?;
1720 let slot = dst
1721 .get_mut(start..end)
1722 .ok_or(crate::Error::InvalidSize(end))?;
1723 slot.copy_from_slice(p.data);
1724 }
1725 Ok(())
1726}
1727
1728fn colorimetry_from(s: &BlobStrings<'_>) -> Option<crate::Colorimetry> {
1732 let c = crate::Colorimetry {
1733 space: crate::ColorSpace::from_str_code(s.color_space),
1734 transfer: crate::ColorTransfer::from_str_code(s.color_transfer),
1735 encoding: crate::ColorEncoding::from_str_code(s.color_encoding),
1736 range: crate::ColorRange::from_str_code(s.color_range),
1737 };
1738 if c == crate::Colorimetry::default() {
1739 None
1740 } else {
1741 Some(c)
1742 }
1743}
1744
1745fn quantization_from(axis: i32, scales: &[f32], zeros: &[i32]) -> Option<crate::Quantization> {
1747 let first = *scales.first()?;
1748 match (axis, zeros.first()) {
1749 (-1, None) => Some(crate::Quantization::per_tensor_symmetric(first)),
1750 (-1, Some(z)) => Some(crate::Quantization::per_tensor(first, *z)),
1751 (a, _) if a >= 0 => {
1752 crate::Quantization::per_channel_symmetric(scales.to_vec(), a as usize).ok()
1753 }
1754 _ => None,
1755 }
1756}
1757
1758fn dims_from_grid(f: crate::PixelFormat, grid: &[usize]) -> crate::Result<(usize, usize)> {
1762 let bad = || crate::Error::InvalidShape(format!("grid {grid:?} is not an image for {f:?}"));
1763 match f.layout() {
1764 crate::PixelLayout::Planar => {
1765 if grid.len() < 3 {
1766 return Err(bad());
1767 }
1768 Ok((grid[2], grid[1]))
1769 }
1770 _ => {
1773 if grid.len() < 2 {
1774 return Err(bad());
1775 }
1776 Ok((grid[1], grid[0]))
1777 }
1778 }
1779}
1780
1781#[cfg(test)]
1782mod tests {
1783 use super::*;
1784
1785 #[cfg(target_os = "windows")]
1790 #[test]
1791 fn a_reference_plane_drops_the_value_with_the_fence() {
1792 assert_eq!(super::reference_completion(0, 42), (0, 0));
1793 assert_eq!(super::reference_completion(0, 0), (0, 0));
1794 assert_eq!(super::reference_completion(0x1234, 0), (0x1234, 0));
1797 assert_eq!(super::reference_completion(0x1234, 42), (0x1234, 42));
1798 }
1799
1800 #[test]
1801 fn every_region_offset_is_constant_time_arithmetic() {
1802 let h = BlobHeader {
1804 size: 4096,
1805 ndim: 3,
1806 strides_len: 3,
1807 quant_scales_len: 4,
1808 quant_zero_points_len: 4,
1809 strings_bytes: 24,
1810 plane_count: 2,
1811 planes_bytes: 128,
1812 ..BlobHeader::empty()
1813 };
1814 let r = region_offsets(&h).unwrap();
1815 assert_eq!(r.shape, HEADER_LEN);
1816 assert_eq!(r.strides, HEADER_LEN + 3 * 8);
1817 assert_eq!(r.quant_scales, HEADER_LEN + 3 * 8 + 3 * 8);
1818 assert_eq!(r.quant_zero_points, r.quant_scales + 4 * 4);
1819 assert_eq!(r.strings, r.quant_zero_points + 4 * 4);
1820 assert_eq!(r.planes, r.strings + 24);
1821 assert_eq!(r.end, r.planes + 128);
1822 }
1823
1824 #[test]
1825 fn empty_strides_leave_the_region_zero_width() {
1826 let h = BlobHeader {
1827 size: 4096,
1828 ndim: 2,
1829 strides_len: 0,
1830 ..BlobHeader::empty()
1831 };
1832 let r = region_offsets(&h).unwrap();
1833 assert_eq!(r.strides, HEADER_LEN + 2 * 8);
1834 assert_eq!(
1835 r.quant_scales, r.strides,
1836 "densely-packed strides occupy no bytes"
1837 );
1838 }
1839
1840 #[test]
1841 fn strides_must_be_ndim_entries_or_none() {
1842 let h = BlobHeader {
1844 size: 4096,
1845 ndim: 3,
1846 strides_len: 2,
1847 ..BlobHeader::empty()
1848 };
1849 assert!(matches!(
1850 region_offsets(&h),
1851 Err(BlobError::ForgedLength {
1852 field: "strides_len",
1853 ..
1854 })
1855 ));
1856 }
1857
1858 #[test]
1859 fn a_forged_ndim_cannot_derive_an_offset_past_the_buffer() {
1860 let h = BlobHeader {
1861 size: 72,
1862 ndim: u32::MAX,
1863 strides_len: u32::MAX,
1864 ..BlobHeader::empty()
1865 };
1866 assert!(
1867 region_offsets(&h).is_err(),
1868 "ndim * 8 must be checked, not wrapped"
1869 );
1870 }
1871
1872 #[test]
1873 fn regions_that_overrun_size_are_refused() {
1874 let h = BlobHeader {
1877 size: 80,
1878 ndim: 2,
1879 strides_len: 2,
1880 plane_count: 1,
1881 planes_bytes: 4096,
1882 ..BlobHeader::empty()
1883 };
1884 assert!(region_offsets(&h).is_err());
1885 }
1886
1887 #[test]
1888 fn plane_count_is_bounded_by_the_planes_region_not_by_a_constant() {
1889 let h = BlobHeader {
1892 size: 4096,
1893 plane_count: u32::MAX,
1894 planes_bytes: 56,
1895 ..BlobHeader::empty()
1896 };
1897 assert!(matches!(
1898 region_offsets(&h),
1899 Err(BlobError::ForgedLength {
1900 field: "plane_count",
1901 ..
1902 })
1903 ));
1904 }
1905
1906 fn fill_ramp(t: &TensorDyn) {
1914 let pin = t.pin_host(CpuAccess::ReadWrite).expect("pin");
1915 t.sync_for_cpu(CpuAccess::ReadWrite).expect("sync for cpu");
1916 let s = unsafe { std::slice::from_raw_parts_mut(pin.as_mut_ptr(), pin.len()) };
1920 for (i, b) in s.iter_mut().enumerate() {
1921 *b = (i % 251) as u8;
1922 }
1923 t.sync_for_device(CpuAccess::ReadWrite)
1924 .expect("sync for device");
1925 }
1926
1927 fn bytes_of(t: &TensorDyn) -> Vec<u8> {
1928 let pin = t.pin_host(CpuAccess::Read).expect("pin");
1929 t.sync_for_cpu(CpuAccess::Read).expect("sync for cpu");
1930 let bytes = unsafe { std::slice::from_raw_parts(pin.as_mut_ptr(), pin.len()) }.to_vec();
1933 t.sync_for_device(CpuAccess::Read).expect("sync for device");
1934 bytes
1935 }
1936
1937 fn reference_capable(w: usize, h: usize) -> Option<TensorDyn> {
1940 use std::io::Write;
1941 match Tensor::<u8>::image(
1942 w,
1943 h,
1944 PixelFormat::Nv12,
1945 Some(TensorMemory::DmaBuf),
1946 CpuAccess::ReadWrite,
1947 ) {
1948 Ok(t) => Some(TensorDyn::from(t)),
1949 Err(e) => {
1950 let _ = writeln!(
1951 std::io::stderr(),
1952 "SKIP: no shareable backing on this host ({e:?}); \
1953 reference transport not exercised"
1954 );
1955 None
1956 }
1957 }
1958 }
1959
1960 #[test]
1961 fn a_reference_export_puts_an_index_in_the_handle_not_a_raw_fd() {
1962 let Some(src) = reference_capable(64, 48) else {
1963 return;
1964 };
1965 let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
1966 let v = BlobView::parse(&blob).unwrap();
1967 let planes = v.planes().unwrap();
1968 for p in &planes {
1969 assert!(!p.is_inline());
1970 }
1971 if cfg!(target_os = "windows") {
1972 assert!(fds.is_empty(), "an NT handle does not travel out of band");
1977 assert!(planes
1978 .iter()
1979 .all(|p| p.handle_bytes.len() == D3D11_HANDLE_BYTES));
1980 } else {
1981 assert!(!fds.is_empty(), "reference mode must carry an fd table");
1982 for p in &planes {
1983 assert!(
1984 (p.handle as usize) < fds.len(),
1985 "handle {} must index the {}-entry fd table, not be a raw fd",
1986 p.handle,
1987 fds.len()
1988 );
1989 }
1990 assert_eq!(fds.len(), 1, "planes sharing a buffer share a table entry");
1994 }
1995 assert!(planes.iter().all(|p| p.handle == 0));
1996 }
1997
1998 #[test]
1999 fn a_reference_round_trip_preserves_content_and_the_source_may_die() {
2000 let Some(src) = reference_capable(64, 48) else {
2001 return;
2002 };
2003 fill_ramp(&src);
2004 let expected = bytes_of(&src);
2005 let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
2006 let src_id = src.buffer_identity().id();
2007 let got = import(&blob, &fds).unwrap();
2008
2009 if !cfg!(target_os = "windows") {
2023 assert_eq!(
2024 got.buffer_identity().id(),
2025 src_id,
2026 "reference mode must share the buffer, not copy it"
2027 );
2028 }
2029
2030 drop(src);
2032 assert_eq!(
2033 bytes_of(&got),
2034 expected,
2035 "the imported tensor sees the bytes"
2036 );
2037 assert_eq!(got.format(), Some(PixelFormat::Nv12));
2038 }
2039
2040 #[test]
2041 fn an_fd_index_past_the_table_is_refused() {
2042 let Some(src) = reference_capable(64, 48) else {
2045 return;
2046 };
2047 let (blob, fds) = export(&src, TransportMode::Reference).unwrap();
2048 if cfg!(target_os = "windows") {
2049 let err = import(&patch_handle_word(&blob, 0, 0), &[])
2054 .expect_err("a null texture handle must be refused, not opened");
2055 assert!(
2056 matches!(err, crate::Error::InvalidArgument(_)),
2057 "expected InvalidArgument, got {err:?}"
2058 );
2059
2060 let named = patch_handle_word(&patch_handle_word(&blob, 1, 0), 2, 7);
2066 let err = import(&named, &[])
2067 .expect_err("a fence value with no fence handle must be refused");
2068 let msg = format!("{err}");
2069 assert!(
2070 matches!(err, crate::Error::InvalidArgument(_))
2071 && msg.contains("cannot be waited on"),
2072 "expected the unwaitable-completion refusal, got {err:?}"
2073 );
2074
2075 let mut tampered = blob.clone();
2078 tampered[HEADER_PID_OFFSET..HEADER_PID_OFFSET + 4]
2079 .copy_from_slice(&u32::MAX.to_le_bytes());
2080 let err =
2081 import(&tampered, &[]).expect_err("a pid that does not exist must be refused");
2082 match &err {
2083 crate::Error::IoError(e) => assert_eq!(
2084 e.kind(),
2085 std::io::ErrorKind::NotFound,
2086 "a pid no process holds is reported as gone, got {e:?}"
2087 ),
2088 other => panic!("expected IoError, got {other:?}"),
2089 }
2090 return;
2091 }
2092 let r = region_offsets(&BlobView::parse(&blob).unwrap().header()).unwrap();
2093 assert!(
2094 import(&blob, &[]).is_err(),
2095 "an empty table cannot satisfy an index"
2096 );
2097 let mut tampered = blob.clone();
2098 tampered[r.planes..r.planes + 8].copy_from_slice(&99i64.to_le_bytes());
2099 assert!(
2100 import(&tampered, &fds).is_err(),
2101 "an out-of-range index must be refused, not read"
2102 );
2103 }
2104
2105 fn patch_handle_word(blob: &[u8], word: usize, value: u64) -> Vec<u8> {
2116 let v = BlobView::parse(blob).unwrap();
2117 let r = region_offsets(&v.header()).unwrap();
2118 let stride = (PLANE_RECORD_LEN + D3D11_HANDLE_BYTES).next_multiple_of(8);
2119 let mut out = blob.to_vec();
2120 for i in 0..v.header().plane_count as usize {
2121 let at = r.planes + i * stride + PLANE_RECORD_LEN + word * 8;
2122 out[at..at + 8].copy_from_slice(&value.to_le_bytes());
2123 }
2124 out
2125 }
2126
2127 #[test]
2132 #[cfg(target_os = "windows")]
2133 fn a_garbage_texture_handle_is_refused_by_the_constructor() {
2134 let Some(src) = reference_capable(64, 48) else {
2135 return;
2136 };
2137 let (blob, _) = export(&src, TransportMode::Reference).unwrap();
2138 let err = import(&patch_handle_word(&blob, 0, 0x1234), &[])
2139 .expect_err("a handle that is not a shared texture must be refused");
2140 let msg = format!("{err}");
2141 assert!(
2142 msg.contains("OpenSharedResource"),
2143 "the refusal must come from opening the handle, got: {msg}"
2144 );
2145 }
2146
2147 #[test]
2150 #[cfg(target_os = "windows")]
2151 fn planes_that_disagree_about_the_texture_are_refused() {
2152 let Some(src) = reference_capable(64, 48) else {
2153 return;
2154 };
2155 let (blob, _) = export(&src, TransportMode::Reference).unwrap();
2156 let v = BlobView::parse(&blob).unwrap();
2157 let r = region_offsets(&v.header()).unwrap();
2158 assert_eq!(v.planes().unwrap().len(), 2, "NV12 exports two planes");
2159 let second = r.planes + (PLANE_RECORD_LEN + D3D11_HANDLE_BYTES).next_multiple_of(8);
2161 let mut tampered = blob.clone();
2162 tampered[second + PLANE_RECORD_LEN..second + PLANE_RECORD_LEN + 8]
2163 .copy_from_slice(&0xBEEFu64.to_le_bytes());
2164 let err = import(&tampered, &[]).expect_err("planes naming two textures must be refused");
2165 let msg = format!("{err}");
2166 assert!(
2167 matches!(err, crate::Error::InvalidArgument(_)) && msg.contains("one texture"),
2168 "expected the one-texture refusal, got {err:?}"
2169 );
2170 }
2171
2172 #[test]
2174 fn pid_offset_matches_the_writer() {
2175 let h = BlobHeader {
2176 size: HEADER_LEN as u64,
2177 pid: 0x0BAD_F00D,
2178 ..BlobHeader::empty()
2179 };
2180 let mut buf = vec![0u8; HEADER_LEN];
2181 h.write_to(&mut buf);
2182 assert_eq!(rd_u32(&buf, HEADER_PID_OFFSET), h.pid);
2183 }
2184
2185 #[test]
2186 fn inline_round_trip_preserves_content_and_every_metadata_field() {
2187 let src = nv12_mem(64, 48).with_colorimetry(crate::Colorimetry::jfif());
2188 fill_ramp(&src);
2189 let blob = export(&src, TransportMode::Inline).unwrap().0;
2190 let got = import(&blob, &[]).unwrap();
2191
2192 assert_eq!(got.dtype(), src.dtype());
2193 assert_eq!(got.format(), src.format());
2194 assert_eq!(got.colorimetry(), src.colorimetry());
2195 assert_eq!(got.shape(), src.shape(), "allocation geometry is rebuilt");
2196 assert_eq!(bytes_of(&got), bytes_of(&src), "every byte survived");
2197 }
2198
2199 #[test]
2200 fn inline_import_is_a_new_allocation_and_the_source_may_die() {
2201 let src = nv12_mem(64, 48);
2202 fill_ramp(&src);
2203 let expected = bytes_of(&src);
2204 let blob = export(&src, TransportMode::Inline).unwrap().0;
2205 drop(src);
2206 let got = import(&blob, &[]).unwrap();
2207 assert_eq!(bytes_of(&got), expected);
2208 }
2209
2210 #[test]
2211 fn a_bare_tensor_round_trips_including_byte_strides() {
2212 let src =
2213 TensorDyn::from(Tensor::<f32>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
2214 fill_ramp(&src);
2215 let blob = export(&src, TransportMode::Inline).unwrap().0;
2216 let got = import(&blob, &[]).unwrap();
2217 assert_eq!(got.shape(), &[3, 8]);
2218 assert_eq!(got.dtype(), crate::DType::F32);
2219 assert_eq!(got.format(), None);
2220 assert_eq!(bytes_of(&got), bytes_of(&src));
2221 }
2222
2223 fn build_blob(planes: &[BlobPlane<'_>]) -> Vec<u8> {
2231 let strings = BlobStrings::default();
2232 let mut planes_buf = Vec::new();
2233 for p in planes {
2234 write_plane(p, &mut planes_buf);
2235 }
2236 let mut out = vec![0u8; HEADER_LEN];
2237 out.extend_from_slice(&1u64.to_le_bytes()); write_strings(&strings, &mut out);
2239 out.extend_from_slice(&planes_buf);
2240 let h = BlobHeader {
2241 size: out.len() as u64,
2242 planes_bytes: planes_buf.len() as u64,
2243 dtype: crate::DType::U8.code(),
2244 quant_axis: -2,
2245 ndim: 1,
2246 plane_count: planes.len() as u32,
2247 strings_bytes: strings_encoded_len(&strings) as u32,
2248 ..BlobHeader::empty()
2249 };
2250 h.write_to(&mut out);
2251 out
2252 }
2253
2254 #[test]
2255 fn no_corrupted_header_value_can_panic_the_parser() {
2256 let good = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2262 for word in 0..HEADER_LEN / 4 {
2263 for pattern in [u32::MAX, u32::MAX / 2, 1, 0x8000_0000] {
2264 let mut bad = good.clone();
2265 bad[word * 4..word * 4 + 4].copy_from_slice(&pattern.to_le_bytes());
2266 if let Ok(v) = BlobView::parse(&bad) {
2269 let _ = v.shape();
2270 let _ = v.strides();
2271 let _ = v.strings();
2272 let _ = v.planes();
2273 }
2274 let _ = import(&bad, &[]);
2275 }
2276 }
2277 }
2278
2279 #[test]
2280 fn no_corrupted_plane_record_can_panic_the_parser() {
2281 let good = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2282 let r = region_offsets(&BlobView::parse(&good).unwrap().header()).unwrap();
2283 for word in 0..PLANE_RECORD_LEN / 4 {
2284 for pattern in [u32::MAX, u32::MAX / 2, 0x8000_0000] {
2285 let mut bad = good.clone();
2286 let at = r.planes + word * 4;
2287 bad[at..at + 4].copy_from_slice(&pattern.to_le_bytes());
2288 if let Ok(v) = BlobView::parse(&bad) {
2289 let _ = v.planes();
2290 }
2291 let _ = import(&bad, &[]);
2292 }
2293 }
2294 }
2295
2296 #[test]
2297 fn a_mixed_inline_and_referenced_plane_set_is_refused() {
2298 let payload = [1u8, 2, 3, 4];
2306 let referenced = BlobPlane {
2307 handle: 7,
2308 offset: 0,
2309 stride: 4,
2310 size: 4,
2311 used: 4,
2312 modifier: 0,
2313 handle_bytes: &[],
2314 data: &[],
2315 };
2316 let inlined = BlobPlane {
2317 handle: -1,
2318 offset: 0,
2319 stride: 4,
2320 size: payload.len() as u64,
2321 used: payload.len() as u64,
2322 modifier: 0,
2323 handle_bytes: &[],
2324 data: &payload,
2325 };
2326 let blob = build_blob(&[referenced, inlined]);
2327 let err =
2328 import(&blob, &[]).expect_err("a plane set mixing transport modes must be refused");
2329 assert!(
2330 err.to_string().contains("one transport mode"),
2331 "must fail on the mode mix, not incidentally: {err}"
2332 );
2333 }
2334
2335 #[test]
2336 fn a_uniformly_referenced_plane_set_is_not_rejected_as_mixed() {
2337 let referenced = BlobPlane {
2342 handle: 7,
2343 offset: 0,
2344 stride: 4,
2345 size: 4,
2346 used: 4,
2347 modifier: 0,
2348 handle_bytes: &[],
2349 data: &[],
2350 };
2351 let blob = build_blob(&[referenced, referenced]);
2352 let err = import(&blob, &[]).expect_err("reference-mode import is not implemented yet");
2353 assert!(
2354 !err.to_string().contains("one transport mode"),
2355 "a uniform plane set must clear the mode check: {err}"
2356 );
2357 }
2358
2359 #[test]
2360 fn an_inline_plane_whose_size_disagrees_with_its_data_is_refused() {
2361 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2362 let mut tampered = blob.clone();
2363 let r = region_offsets(&BlobView::parse(&blob).unwrap().header()).unwrap();
2364 tampered[r.planes + 24..r.planes + 32].copy_from_slice(&9999u64.to_le_bytes());
2365 assert!(import(&tampered, &[]).is_err());
2366 }
2367
2368 #[test]
2369 fn quantization_survives_the_round_trip() {
2370 let mut src =
2371 TensorDyn::from(Tensor::<i8>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
2372 src.set_quantization(crate::Quantization::per_tensor(0.5, -7))
2373 .unwrap();
2374 let blob = export(&src, TransportMode::Inline).unwrap().0;
2375 let got = import(&blob, &[]).unwrap();
2376 let q = got.quantization().expect("quantization survived");
2377 assert_eq!(q.scale(), &[0.5]);
2378 assert_eq!(q.zero_point(), Some(&[-7i32][..]));
2379 }
2380
2381 use crate::{CpuAccess, PixelFormat, Tensor, TensorDyn, TensorMemory};
2382
2383 fn nv12_mem(w: usize, h: usize) -> TensorDyn {
2384 TensorDyn::from(
2385 Tensor::<u8>::image(
2386 w,
2387 h,
2388 PixelFormat::Nv12,
2389 Some(TensorMemory::Mem),
2390 CpuAccess::ReadWrite,
2391 )
2392 .expect("alloc"),
2393 )
2394 }
2395
2396 #[test]
2397 fn export_carries_the_addressing_grid_not_the_allocation() {
2398 let blob = export(&nv12_mem(640, 480), TransportMode::Inline)
2399 .unwrap()
2400 .0;
2401 let v = BlobView::parse(&blob).unwrap();
2402 assert_eq!(v.shape(), &[480, 640]);
2404 assert_eq!(v.header().dtype, crate::DType::U8.code());
2405 assert_eq!(v.strings().format, "NV12");
2406 }
2407
2408 #[test]
2409 fn export_emits_one_plane_per_plane_table_entry() {
2410 let blob = export(&nv12_mem(640, 480), TransportMode::Inline)
2411 .unwrap()
2412 .0;
2413 let v = BlobView::parse(&blob).unwrap();
2414 let planes = v.planes().unwrap();
2415 assert_eq!(planes.len(), 2, "NV12 is two planes");
2416 assert_eq!(planes[0].size as usize, planes[0].data.len());
2418 assert_eq!(planes[1].size as usize, planes[1].data.len());
2419 assert_eq!(planes[0].size, 307_200);
2420 assert_eq!(planes[1].size, 153_600);
2421 }
2422
2423 #[test]
2424 fn inline_export_clears_pid_and_fence_fd() {
2425 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2428 let h = BlobView::parse(&blob).unwrap().header();
2429 assert_eq!(h.pid, 0, "no handle to reopen");
2430 assert_eq!(h.fence_fd, -1, "inlining means the producer already waited");
2431 }
2432
2433 #[test]
2434 fn inline_planes_carry_no_modifier_and_no_handle_bytes() {
2435 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2437 for p in BlobView::parse(&blob).unwrap().planes().unwrap() {
2438 assert!(p.is_inline());
2439 assert_eq!(p.modifier, 0);
2440 assert!(p.handle_bytes.is_empty());
2441 assert_eq!(p.size as usize, p.data.len());
2442 }
2443 }
2444
2445 #[test]
2446 fn reference_export_of_a_mem_tensor_is_refused() {
2447 assert!(export(&nv12_mem(64, 48), TransportMode::Reference).is_err());
2451 }
2452
2453 #[test]
2454 fn a_blob_parses_correctly_from_a_non_zero_offset() {
2455 let blob = export(&nv12_mem(64, 48), TransportMode::Inline).unwrap().0;
2458 let mut framed = vec![0xAAu8; 16];
2459 framed.extend_from_slice(&blob);
2460 let direct = BlobView::parse(&blob).unwrap();
2461 let framed_view = BlobView::parse(&framed[16..]).unwrap();
2462 assert_eq!(framed_view.header(), direct.header());
2463 assert_eq!(framed_view.shape(), direct.shape());
2464 assert_eq!(framed_view.strings(), direct.strings());
2465 }
2466
2467 #[test]
2468 fn colorimetry_axes_are_carried_as_strings() {
2469 let t = nv12_mem(64, 48).with_colorimetry(crate::Colorimetry::jfif());
2470 let blob = export(&t, TransportMode::Inline).unwrap().0;
2471 let v = BlobView::parse(&blob).unwrap();
2472 let s = v.strings();
2473 assert_eq!(s.color_space, "srgb");
2474 assert_eq!(s.color_transfer, "srgb");
2475 assert_eq!(s.color_encoding, "bt601");
2476 assert_eq!(s.color_range, "full");
2477 }
2478
2479 #[test]
2480 fn a_bare_tensor_carries_no_format_and_its_shape_is_its_grid() {
2481 let t =
2482 TensorDyn::from(Tensor::<f32>::new(&[3, 8], Some(TensorMemory::Mem), None).unwrap());
2483 let blob = export(&t, TransportMode::Inline).unwrap().0;
2484 let v = BlobView::parse(&blob).unwrap();
2485 assert_eq!(v.strings().format, "", "not an image");
2486 assert_eq!(v.shape(), &[3, 8]);
2487 assert_eq!(v.strides(), &[32i64, 4], "byte strides for f32");
2488 }
2489
2490 fn ref_plane() -> BlobPlane<'static> {
2491 BlobPlane {
2492 handle: 7,
2493 offset: 0,
2494 stride: 640,
2495 size: 307_200,
2496 used: 307_200,
2497 modifier: 0,
2498 handle_bytes: &[],
2499 data: &[],
2500 }
2501 }
2502
2503 #[test]
2504 fn a_plane_record_round_trips() {
2505 let mut out = Vec::new();
2506 write_plane(&ref_plane(), &mut out);
2507 let got = parse_planes(&out, 1).unwrap();
2508 assert_eq!(got.len(), 1);
2509 assert_eq!(got[0], ref_plane());
2510 }
2511
2512 #[test]
2513 fn a_plane_scalar_block_is_56_bytes() {
2514 let mut out = Vec::new();
2517 write_plane(&ref_plane(), &mut out);
2518 assert_eq!(out.len(), PLANE_RECORD_LEN);
2519 assert_eq!(PLANE_RECORD_LEN, 56);
2520 }
2521
2522 #[test]
2523 fn every_plane_record_starts_8_aligned_whatever_the_payload() {
2524 for n in 0..17usize {
2527 let payload: Vec<u8> = (0..n as u8).collect();
2528 let mut out = Vec::new();
2529 write_plane(
2530 &BlobPlane {
2531 handle: -1,
2532 size: n as u64,
2533 data: &payload,
2534 ..ref_plane()
2535 },
2536 &mut out,
2537 );
2538 write_plane(&ref_plane(), &mut out);
2539 assert_eq!(
2540 out.len() % 8,
2541 0,
2542 "payload of {n} bytes left the region unaligned"
2543 );
2544 let got = parse_planes(&out, 2).unwrap();
2545 assert_eq!(got[0].data, &payload[..], "payload {n} survived padding");
2546 assert_eq!(
2547 got[1],
2548 ref_plane(),
2549 "the second record parsed from an aligned offset"
2550 );
2551 }
2552 }
2553
2554 #[test]
2555 fn plane_payloads_round_trip_independently() {
2556 let hb = [0xDEu8, 0xAD, 0xBE, 0xEF];
2559 let data = [1u8, 2, 3];
2560 let mut out = Vec::new();
2561 write_plane(
2562 &BlobPlane {
2563 handle_bytes: &hb,
2564 data: &data,
2565 ..ref_plane()
2566 },
2567 &mut out,
2568 );
2569 let got = parse_planes(&out, 1).unwrap();
2570 assert_eq!(got[0].handle_bytes, &hb[..]);
2571 assert_eq!(got[0].data, &data[..]);
2572 }
2573
2574 #[test]
2575 fn a_forged_plane_payload_length_is_refused() {
2576 let mut out = vec![0u8; PLANE_RECORD_LEN];
2577 out[48..52].copy_from_slice(&u32::MAX.to_le_bytes()); assert!(matches!(
2579 parse_planes(&out, 1),
2580 Err(BlobError::ForgedLength { .. })
2581 ));
2582 }
2583
2584 #[test]
2585 fn a_plane_count_larger_than_the_records_present_is_refused() {
2586 let mut out = Vec::new();
2587 write_plane(&ref_plane(), &mut out);
2588 assert!(
2589 parse_planes(&out, 2).is_err(),
2590 "claims two records, holds one"
2591 );
2592 }
2593
2594 fn sample_strings() -> BlobStrings<'static> {
2595 BlobStrings {
2596 format: "NV12",
2597 color_space: "bt709",
2598 color_transfer: "bt709",
2599 color_encoding: "bt709",
2600 color_range: "limited",
2601 }
2602 }
2603
2604 #[test]
2605 fn the_five_strings_round_trip_in_message_order() {
2606 let mut out = Vec::new();
2607 write_strings(&sample_strings(), &mut out);
2608 let got = parse_strings(&out).unwrap();
2609 assert_eq!(got, sample_strings());
2610 }
2611
2612 #[test]
2613 fn string_order_is_positional_and_matches_the_message() {
2614 let s = BlobStrings {
2618 format: "a",
2619 color_space: "bb",
2620 color_transfer: "ccc",
2621 color_encoding: "dddd",
2622 color_range: "eeeee",
2623 };
2624 let mut out = Vec::new();
2625 write_strings(&s, &mut out);
2626 assert_eq!(parse_strings(&out).unwrap(), s);
2627 }
2628
2629 #[test]
2630 fn empty_strings_are_legal_and_distinct_from_absent() {
2631 let s = BlobStrings::default();
2633 let mut out = Vec::new();
2634 write_strings(&s, &mut out);
2635 assert_eq!(parse_strings(&out).unwrap(), s);
2636 assert!(
2637 !out.is_empty(),
2638 "five empty strings still occupy their counts"
2639 );
2640 }
2641
2642 #[test]
2643 fn the_strings_region_is_padded_so_planes_start_8_aligned() {
2644 for f in ["", "a", "ab", "abc", "NV12", "rgb8_planar"] {
2647 let mut out = Vec::new();
2648 write_strings(
2649 &BlobStrings {
2650 format: f,
2651 ..Default::default()
2652 },
2653 &mut out,
2654 );
2655 assert_eq!(out.len() % 8, 0, "region for format {f:?} is not 8-aligned");
2656 }
2657 }
2658
2659 #[test]
2660 fn a_non_utf8_string_is_refused() {
2661 let mut out = Vec::new();
2664 out.extend_from_slice(&2u32.to_le_bytes());
2665 out.extend_from_slice(&[0xff, 0xfe]);
2666 for _ in 0..4 {
2667 out.extend_from_slice(&0u32.to_le_bytes());
2668 }
2669 while out.len() % 8 != 0 {
2670 out.push(0);
2671 }
2672 assert!(matches!(
2673 parse_strings(&out),
2674 Err(BlobError::InvalidUtf8(_))
2675 ));
2676 }
2677
2678 #[test]
2679 fn a_forged_string_length_cannot_read_past_the_region() {
2680 let mut out = Vec::new();
2681 out.extend_from_slice(&u32::MAX.to_le_bytes());
2682 out.extend_from_slice(b"NV12");
2683 while out.len() % 8 != 0 {
2684 out.push(0);
2685 }
2686 assert!(matches!(
2687 parse_strings(&out),
2688 Err(BlobError::ForgedLength { .. })
2689 ));
2690 }
2691
2692 #[test]
2693 fn a_truncated_strings_region_is_refused() {
2694 assert!(parse_strings(&[0u8; 4]).is_err(), "needs five counts");
2695 }
2696
2697 #[test]
2698 fn an_overrunning_region_is_attributed_to_its_own_field() {
2699 let h = BlobHeader {
2705 size: 4096,
2706 ndim: 2,
2707 strides_len: 2,
2708 strings_bytes: u32::MAX,
2709 ..BlobHeader::empty()
2710 };
2711 assert!(
2712 matches!(
2713 region_offsets(&h),
2714 Err(BlobError::ForgedLength {
2715 field: "strings_bytes",
2716 ..
2717 })
2718 ),
2719 "the overrun must name strings_bytes, not a later field"
2720 );
2721 }
2722
2723 #[test]
2724 fn header_is_72_bytes_and_every_scalar_is_naturally_aligned() {
2725 assert_eq!(HEADER_LEN, 72);
2726 assert_eq!(HEADER_LEN % 8, 0, "the variable tail starts 8-aligned");
2727 for off in [0usize, 8, 16] {
2728 assert_eq!(off % 8, 0, "u64 field at {off} is not 8-aligned");
2729 }
2730 for off in [24usize, 28, 32, 36, 40, 44, 48, 52, 56, 60, 64] {
2731 assert_eq!(off % 4, 0, "u32 field at {off} is not 4-aligned");
2732 }
2733 }
2734
2735 #[test]
2736 fn a_header_round_trips_through_bytes() {
2737 let h = BlobHeader {
2738 size: 72,
2739 required_mask: 0,
2740 planes_bytes: 0,
2741 storage_kind: 2,
2742 pid: 4242,
2743 fence_fd: -1,
2744 dtype: 1,
2745 quant_axis: -2,
2746 ndim: 2,
2747 plane_count: 0,
2748 quant_scales_len: 0,
2749 quant_zero_points_len: 0,
2750 strings_bytes: 0,
2751 strides_len: 2,
2752 };
2753 let mut buf = vec![0u8; HEADER_LEN];
2754 h.write_to(&mut buf);
2755 assert_eq!(parse_header(&buf).unwrap(), h);
2756 }
2757
2758 #[test]
2759 fn the_header_is_little_endian_on_every_host() {
2760 let mut buf = vec![0u8; HEADER_LEN];
2761 BlobHeader {
2762 size: 0x0102_0304_0506_0708,
2763 ..BlobHeader::empty()
2764 }
2765 .write_to(&mut buf);
2766 assert_eq!(
2767 &buf[0..8],
2768 &[0x08, 0x07, 0x06, 0x05, 0x04, 0x03, 0x02, 0x01],
2769 "size must be little-endian regardless of host byte order"
2770 );
2771 }
2772
2773 #[test]
2774 fn a_truncated_buffer_is_refused_not_read() {
2775 assert!(matches!(
2776 parse_header(&[0u8; 71]),
2777 Err(BlobError::Truncated { .. })
2778 ));
2779 }
2780
2781 #[test]
2782 fn an_unknown_required_bit_is_refused() {
2783 let mut buf = vec![0u8; HEADER_LEN];
2784 BlobHeader {
2785 size: 72,
2786 required_mask: 1 << 63,
2787 ..BlobHeader::empty()
2788 }
2789 .write_to(&mut buf);
2790 assert!(matches!(
2791 parse_header(&buf),
2792 Err(BlobError::UnsupportedRequirement(_))
2793 ));
2794 }
2795
2796 #[test]
2797 fn a_size_smaller_than_the_header_is_refused() {
2798 let mut buf = vec![0u8; HEADER_LEN];
2799 BlobHeader {
2800 size: 8,
2801 ..BlobHeader::empty()
2802 }
2803 .write_to(&mut buf);
2804 assert!(
2805 parse_header(&buf).is_err(),
2806 "size must cover at least the header"
2807 );
2808 }
2809}