From b52f43d1d50f156f262d61356d50d3010f4acaf7 Mon Sep 17 00:00:00 2001 From: Brandon Thomas Date: Wed, 7 Oct 2026 09:56:49 -0400 Subject: [PATCH] M12: lightcraft-segment: Apache-2.0, checked weights, pure-Rust model download MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - The crate ports Apache-2.0-only Hugging Face transformers code: it is now `license = "Apache-2.0"`, and every ported file carries the modification notice Apache-2.0 §4(b) asks for. - Weights::open checks every tensor's byte range (and the header) against the file's length: a truncated or damaged download is an error at load, never a short read later; regression test with hostile offsets. - fetch/: the model download, only when asked for. An ordered mirror list (env LIGHTCRAFT_SAM3_MIRRORS, a mirrors file, DEFAULT_MIRRORS — empty until LightCraft's CDN exists, TODO for the maintainer), fallback to the next mirror, resume via Range, redirects, connect and stall timeouts, cancel, size caps, exact size + SHA-256 (pinned from tools/install-sam3.sh) before an atomic rename; damaged files are deleted, HTTP error bodies never written. - HTTP/1.1 client over std sockets + rustls with the RustCrypto provider (rustls-rustcrypto) and webpki-roots: no ring / aws-lc-rs, nothing compiled from C (no new `cc` users in the app's build graph). - Tests against a local TcpListener server (no internet): fallback, redirects, chunked bodies, resume, 416, hash mismatch, too much data, stalls, cancel. - *.safetensors are gitignored. Co-Authored-By: Claude Opus 5.5 --- .gitignore | 3 + Cargo.lock | 755 +++++++++++++++++++++++++++++- crates/segment/Cargo.toml | 6 +- crates/segment/src/clip.rs | 4 + crates/segment/src/detector.rs | 4 + crates/segment/src/fetch/http.rs | 541 +++++++++++++++++++++ crates/segment/src/fetch/mod.rs | 424 +++++++++++++++++ crates/segment/src/fetch/tests.rs | 314 +++++++++++++ crates/segment/src/lib.rs | 13 +- crates/segment/src/mask.rs | 4 + crates/segment/src/neck.rs | 4 + crates/segment/src/nn.rs | 4 + crates/segment/src/tokenizer.rs | 4 + crates/segment/src/tracker.rs | 4 + crates/segment/src/vit.rs | 4 + crates/segment/src/weights.rs | 47 +- 16 files changed, 2119 insertions(+), 16 deletions(-) create mode 100644 crates/segment/src/fetch/http.rs create mode 100644 crates/segment/src/fetch/mod.rs create mode 100644 crates/segment/src/fetch/tests.rs diff --git a/.gitignore b/.gitignore index 823f92765..51531a34b 100644 --- a/.gitignore +++ b/.gitignore @@ -8,3 +8,6 @@ plan/ .DS_Store .claude/ /craft-fonts +# model weights are never committed (SAM 3: docs/ai-masks.md) +*.safetensors +*.safetensors.part diff --git a/Cargo.lock b/Cargo.lock index 6a07ead6b..3f0861610 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -146,6 +146,41 @@ version = "2.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" +[[package]] +name = "aead" +version = "0.5.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d122413f284cf2d62fb1b7db97e02edb8cda96d769b16e443a4f6195e35662b0" +dependencies = [ + "crypto-common", + "generic-array", +] + +[[package]] +name = "aes" +version = "0.8.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b169f7a6d4742236a0a00c541b845991d0ac43e546831af1249753ab4c3aa3a0" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "aes-gcm" +version = "0.10.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "831010a0f742e1209b3bcea8fab6a8e149051ba6099432c8cb2cc117dec3ead1" +dependencies = [ + "aead", + "aes", + "cipher", + "ctr", + "ghash", + "subtle", +] + [[package]] name = "ahash" version = "0.8.12" @@ -543,6 +578,18 @@ dependencies = [ "windows-link", ] +[[package]] +name = "base16ct" +version = "0.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" + +[[package]] +name = "base64ct" +version = "1.8.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" + [[package]] name = "bit-set" version = "0.8.0" @@ -600,6 +647,15 @@ version = "0.1.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0d8c1fef690941d3e7788d328517591fecc684c084084702d6ff1641e993699a" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "block2" version = "0.5.1" @@ -771,7 +827,7 @@ dependencies = [ "num_cpus", "objc2-foundation 0.3.2", "objc2-metal 0.3.2", - "rand", + "rand 0.9.5", "rand_distr", "rayon", "safetensors 0.7.0", @@ -856,6 +912,41 @@ dependencies = [ "libc", ] +[[package]] +name = "chacha20" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c3613f74bd2eac03dad61bd53dbe620703d4371614fe0bc3b9f04dd36fe4e818" +dependencies = [ + "cfg-if", + "cipher", + "cpufeatures", +] + +[[package]] +name = "chacha20poly1305" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "10cd79432192d1c0f4e1a0fef9527696cc039165d729fb41b3f4f4f354c2dc35" +dependencies = [ + "aead", + "chacha20", + "cipher", + "poly1305", + "zeroize", +] + +[[package]] +name = "cipher" +version = "0.4.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "773f3b9af64447d2ce9850330c473515014aa235e6a783b02db81ff39e4a3dad" +dependencies = [ + "crypto-common", + "inout", + "zeroize", +] + [[package]] name = "clipboard-win" version = "5.4.1" @@ -910,6 +1001,12 @@ dependencies = [ "crossbeam-utils", ] +[[package]] +name = "const-oid" +version = "0.9.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c2459377285ad874054d797f3ccebf984978aa39129f6eafde5cdc8315b612f8" + [[package]] name = "core-foundation" version = "0.9.4" @@ -950,6 +1047,15 @@ dependencies = [ "libc", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "crc32fast" version = "1.5.2" @@ -999,12 +1105,81 @@ version = "0.2.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "460fbee9c2c2f33933d720630a6a0bac33ba7053db5344fac858d4b8952d77d5" +[[package]] +name = "crypto-bigint" +version = "0.5.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0dc92fb57ca44df6db8059111ab3af99a63d5d0f8375d9972e319a379c6bab76" +dependencies = [ + "generic-array", + "rand_core 0.6.4", + "subtle", + "zeroize", +] + +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "ctr" +version = "0.9.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0369ee1ad671834580515889b80f2ea915f23b8be8d0daa4bbaf2ac5c7590835" +dependencies = [ + "cipher", +] + [[package]] name = "cursor-icon" version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f27ae1dd37df86211c42e150270f82743308803d90a6f6e6651cd730d5e1732f" +[[package]] +name = "curve25519-dalek" +version = "4.1.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "97fb8b7c4503de7d6ae7b42ab72a5a59857b4c937ec27a3d4539dba95b5ab2be" +dependencies = [ + "cfg-if", + "cpufeatures", + "curve25519-dalek-derive", + "digest", + "fiat-crypto", + "rustc_version", + "subtle", + "zeroize", +] + +[[package]] +name = "curve25519-dalek-derive" +version = "0.1.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f46882e17999c6cc590af592290432be3bce0428cb0d5f8b6715e4dc7b383eb3" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + +[[package]] +name = "der" +version = "0.7.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e7c1832837b905bbfb5101e07cc24c8deddf52f93225eee6ead5f4d63d53ddcb" +dependencies = [ + "const-oid", + "pem-rfc7468", + "zeroize", +] + [[package]] name = "dhat" version = "0.3.3" @@ -1021,6 +1196,18 @@ dependencies = [ "thousands", ] +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "const-oid", + "crypto-common", + "subtle", +] + [[package]] name = "dispatch" version = "0.2.0" @@ -1085,6 +1272,20 @@ version = "0.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e1d926b4d407d372f141f93bb444696142c29d32962ccbd3531117cf3aa0bfa9" +[[package]] +name = "ecdsa" +version = "0.16.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" +dependencies = [ + "der", + "digest", + "elliptic-curve", + "rfc6979", + "signature", + "spki", +] + [[package]] name = "ecolor" version = "0.36.2" @@ -1095,6 +1296,30 @@ dependencies = [ "emath", ] +[[package]] +name = "ed25519" +version = "2.2.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "115531babc129696a58c64a4fef0a8bf9e9698629fb97e9e40767d235cfbcd53" +dependencies = [ + "pkcs8", + "signature", +] + +[[package]] +name = "ed25519-dalek" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "70e796c081cee67dc755e1a36a0a172b897fab85fc3f6bc48307991f64e4eca9" +dependencies = [ + "curve25519-dalek", + "ed25519", + "serde", + "sha2", + "subtle", + "zeroize", +] + [[package]] name = "eframe" version = "0.36.2" @@ -1212,6 +1437,27 @@ version = "1.18.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "252afb9ae5eaa683babdc6a068b3f5726eb19e05070c731f9b2a23a7c3e8ed34" +[[package]] +name = "elliptic-curve" +version = "0.13.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b5e6043086bf7973472e0c7dff2142ea0b680d30e18d9cc40f267efbf222bd47" +dependencies = [ + "base16ct", + "crypto-bigint", + "digest", + "ff", + "generic-array", + "group", + "hkdf", + "pem-rfc7468", + "pkcs8", + "rand_core 0.6.4", + "sec1", + "subtle", + "zeroize", +] + [[package]] name = "emath" version = "0.36.2" @@ -1389,6 +1635,22 @@ version = "0.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b97b65636e5b9ef369943878ac74335ba1c55c1cb6adbf1e2c293c624248d693" +[[package]] +name = "ff" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0b50bfb653653f9ca9095b427bed08ab8d75a137839d9ad64eb11810d5b6393" +dependencies = [ + "rand_core 0.6.4", + "subtle", +] + +[[package]] +name = "fiat-crypto" +version = "0.2.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" + [[package]] name = "find-msvc-tools" version = "0.1.14" @@ -1414,7 +1676,7 @@ checksum = "719a903cc23e4a89e87962c2a80fdb45cdaad0983a89bd150bb57b4c8571a7d5" dependencies = [ "half", "num-traits", - "rand", + "rand 0.9.5", "rand_distr", ] @@ -1759,6 +2021,17 @@ dependencies = [ "seq-macro", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", + "zeroize", +] + [[package]] name = "gethostname" version = "1.1.0" @@ -1769,6 +2042,17 @@ dependencies = [ "windows-link", ] +[[package]] +name = "getrandom" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ff2abc00be7fca6ebc474524697ae276ad847ad0a6b3faa4bcb027e9a4614ad0" +dependencies = [ + "cfg-if", + "libc", + "wasi", +] + [[package]] name = "getrandom" version = "0.3.4" @@ -1792,6 +2076,16 @@ dependencies = [ "r-efi 6.0.0", ] +[[package]] +name = "ghash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0d8a4362ccb29cb0b265253fb0a2728f592895ee6854fd9bc13f2ffda266ff1" +dependencies = [ + "opaque-debug", + "polyval", +] + [[package]] name = "gif" version = "0.14.2" @@ -1927,6 +2221,17 @@ dependencies = [ "windows", ] +[[package]] +name = "group" +version = "0.13.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f0f9ef7462f7c099f518d754361858f86d8a07af53ba9af0fe635bbccb151a63" +dependencies = [ + "ff", + "rand_core 0.6.4", + "subtle", +] + [[package]] name = "guillotiere" version = "0.7.0" @@ -1946,7 +2251,7 @@ dependencies = [ "cfg-if", "crunchy", "num-traits", - "rand", + "rand 0.9.5", "rand_distr", "zerocopy", ] @@ -2003,6 +2308,24 @@ version = "0.4.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "7f24254aa9a54b5c858eaee2f5bccdb46aaf0e486a595ed5fd8f86ba55232a70" +[[package]] +name = "hkdf" +version = "0.12.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7b5f8eb2ad728638ea2c7d47a21db23b7b58a72ed6a38256b8a1849f15fbbdf7" +dependencies = [ + "hmac", +] + +[[package]] +name = "hmac" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "6c49c37c09c17a53d937dfbb742eb3a961d65a994e6bcdcf37e7399d0cc8ab5e" +dependencies = [ + "digest", +] + [[package]] name = "image" version = "0.25.10" @@ -2048,6 +2371,15 @@ dependencies = [ "hashbrown 0.17.1", ] +[[package]] +name = "inout" +version = "0.1.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879f10e63c20629ecabbb64a8010319738c66a5cd0c29b02d63d272b03751d01" +dependencies = [ + "generic-array", +] + [[package]] name = "interpolate_name" version = "0.2.4" @@ -2397,6 +2729,9 @@ name = "lazy_static" version = "1.5.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "20870f649af7073d53e38067b2a84312175d56ea15217e1b15bc83506ec50afb" +dependencies = [ + "spin", +] [[package]] name = "libc" @@ -2705,9 +3040,13 @@ dependencies = [ "candle-core", "candle-nn", "log", + "rustls", + "rustls-rustcrypto", "serde", "serde_json", + "sha2", "thiserror 2.0.21", + "webpki-roots", ] [[package]] @@ -3065,6 +3404,22 @@ dependencies = [ "num-traits", ] +[[package]] +name = "num-bigint-dig" +version = "0.8.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e661dda6640fad38e827a6d4a310ff4763082116fe217f279885c97f511bb0b7" +dependencies = [ + "lazy_static", + "libm", + "num-integer", + "num-iter", + "num-traits", + "rand 0.8.8", + "smallvec", + "zeroize", +] + [[package]] name = "num-complex" version = "0.4.6" @@ -3496,6 +3851,12 @@ version = "1.21.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f7c3e4beb33f85d45ae3e3a1792185706c8e16d043238c593331cc7cd313b50" +[[package]] +name = "opaque-debug" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c08d65885ee38876c4f86fa503fb49d7b507c2b62552df7c70b2fce627e06381" + [[package]] name = "orbclient" version = "0.3.55" @@ -3534,6 +3895,30 @@ dependencies = [ "ttf-parser", ] +[[package]] +name = "p256" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c9863ad85fa8f4460f9c48cb909d38a0d689dba1f6f6988a5e3e0d31071bcd4b" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + +[[package]] +name = "p384" +version = "0.13.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fe42f1670a52a47d448f14b6a5c61dd78fce51856e68edaa38f7ae3a46b8d6b6" +dependencies = [ + "ecdsa", + "elliptic-curve", + "primeorder", + "sha2", +] + [[package]] name = "parking" version = "2.2.1" @@ -3575,6 +3960,15 @@ version = "0.1.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "35fb2e5f958ec131621fdd531e9fc186ed768cbe395337403ae56c17a74c68ec" +[[package]] +name = "pem-rfc7468" +version = "0.7.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "88b39c9bfcfc231068454382784bb460aae594343fb030d46e9f50a645418412" +dependencies = [ + "base64ct", +] + [[package]] name = "peniko" version = "0.6.1" @@ -3674,6 +4068,38 @@ dependencies = [ "futures-io", ] +[[package]] +name = "pkcs1" +version = "0.7.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c8ffb9f10fa047879315e6625af03c164b16962a5368d724ed16323b68ace47f" +dependencies = [ + "der", + "pkcs8", + "spki", +] + +[[package]] +name = "pkcs5" +version = "0.7.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e847e2c91a18bfa887dd028ec33f2fe6f25db77db3619024764914affe8b69a6" +dependencies = [ + "der", + "spki", +] + +[[package]] +name = "pkcs8" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f950b2377845cebe5cf8b5165cb3cc1a5e0fa5cfa3e1f7f55707d8fd82e0a7b7" +dependencies = [ + "der", + "pkcs5", + "spki", +] + [[package]] name = "pkg-config" version = "0.3.34" @@ -3725,6 +4151,17 @@ version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bc6355899e1c9462875b6757c79f3caa011a1fdae12bbb1a2e72dd1f234f8336" +[[package]] +name = "poly1305" +version = "0.8.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8159bd90725d2df49889a078b54f4f79e87f1f8a8444194cdca81d38f5393abf" +dependencies = [ + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "polycool" version = "0.4.0" @@ -3734,6 +4171,18 @@ dependencies = [ "arrayvec", ] +[[package]] +name = "polyval" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9d1fe60d06143b2430aa532c94cfe9e29783047f06c0d7fd359a9a51b729fa25" +dependencies = [ + "cfg-if", + "cpufeatures", + "opaque-debug", + "universal-hash", +] + [[package]] name = "portable-atomic" version = "1.15.0" @@ -3764,6 +4213,15 @@ version = "0.3.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e8cf8e6a8aa66ce33f63993ffc4ea4271eb5b0530a9002db8455ea6050c77bfa" +[[package]] +name = "primeorder" +version = "0.13.6" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "353e1ca18966c16d9deb1c69278edbc5f194139612772bd9537af60ac231e1e6" +dependencies = [ + "elliptic-curve", +] + [[package]] name = "proc-macro-crate" version = "3.5.0" @@ -3811,8 +4269,8 @@ dependencies = [ "bit-vec 0.8.0", "bitflags 2.13.2", "num-traits", - "rand", - "rand_chacha", + "rand 0.9.5", + "rand_chacha 0.9.0", "rand_xorshift", "regex-syntax", "rusty-fork", @@ -3914,14 +4372,34 @@ version = "6.0.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "f8dcc9c7d52a811697d2151c701e0d08956f92b0e24136cf4cf27b57a6a0d9bf" +[[package]] +name = "rand" +version = "0.8.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e058c7de0b26af77780c769414d6257830bb240f3c38477dbc2c16e5f54d6d4c" +dependencies = [ + "rand_chacha 0.3.1", + "rand_core 0.6.4", +] + [[package]] name = "rand" version = "0.9.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b9ef1d0d795eb7d84685bca4f72f3649f064e6641543d3a8c415898726a57b41" dependencies = [ - "rand_chacha", - "rand_core", + "rand_chacha 0.9.0", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_chacha" +version = "0.3.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e6c10a63a0fa32252be49d21e7709d4d4baf8d231c2dbce1eaa8141b9b127d88" +dependencies = [ + "ppv-lite86", + "rand_core 0.6.4", ] [[package]] @@ -3931,7 +4409,16 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d3022b5f1df60f26e1ffddd6c66e8aa15de382ae63b3a0c1bfc0e4d3e3f325cb" dependencies = [ "ppv-lite86", - "rand_core", + "rand_core 0.9.5", +] + +[[package]] +name = "rand_core" +version = "0.6.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ec0be4795e2f6a28069bec0b5ff3e2ac9bafc99e6a9a7dc3547996c5c816922c" +dependencies = [ + "getrandom 0.2.17", ] [[package]] @@ -3950,7 +4437,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6a8615d50dcf34fa31f7ab52692afec947c4dd0ab803cc87cb3b0b4570ff7463" dependencies = [ "num-traits", - "rand", + "rand 0.9.5", ] [[package]] @@ -3959,7 +4446,7 @@ version = "0.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "513962919efc330f829edb2535844d1b912b0fbe2ca165d613e4e8788bb05a5a" dependencies = [ - "rand_core", + "rand_core 0.9.5", ] [[package]] @@ -3995,8 +4482,8 @@ dependencies = [ "num-traits", "paste", "profiling", - "rand", - "rand_chacha", + "rand 0.9.5", + "rand_chacha 0.9.0", "simd_helpers", "thiserror 2.0.21", "v_frame", @@ -4121,6 +4608,16 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19b30a45b0cd0bcca8037f3d0dc3421eaf95327a17cad11964fb8179b4fc4832" +[[package]] +name = "rfc6979" +version = "0.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" +dependencies = [ + "hmac", + "subtle", +] + [[package]] name = "rfd" version = "0.17.2" @@ -4154,6 +4651,41 @@ version = "0.8.53" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "47b34b781b31e5d73e9fbc8689c70551fd1ade9a19e3e28cfec8580a79290cc4" +[[package]] +name = "ring" +version = "0.17.14" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a4689e6c2294d81e88dc6261c768b63bc4fcdb852be6d1352498b114f61383b7" +dependencies = [ + "cc", + "cfg-if", + "getrandom 0.2.17", + "libc", + "untrusted", + "windows-sys 0.52.0", +] + +[[package]] +name = "rsa" +version = "0.9.10" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b8573f03f5883dcaebdfcf4725caa1ecb9c15b2ef50c43a07b816e06799bb12d" +dependencies = [ + "const-oid", + "digest", + "num-bigint-dig", + "num-integer", + "num-traits", + "pkcs1", + "pkcs8", + "rand_core 0.6.4", + "sha2", + "signature", + "spki", + "subtle", + "zeroize", +] + [[package]] name = "rustc-demangle" version = "0.1.28" @@ -4207,6 +4739,80 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "rustls" +version = "0.23.45" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0d41d731c7d2f962d1ccc364cec258de3c0e93b38c2fb3ba97ac74513048d634" +dependencies = [ + "once_cell", + "rustls-pki-types", + "rustls-webpki 0.103.15", + "subtle", + "zeroize", +] + +[[package]] +name = "rustls-pki-types" +version = "1.15.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "2f4925028c7eb5d1fcdaf196971378ed9d2c1c4efc7dc5d011256f76c99c0a96" +dependencies = [ + "zeroize", +] + +[[package]] +name = "rustls-rustcrypto" +version = "0.0.2-alpha" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f12052947763ab8515f753315357599e9b0b4dab3b8ba15f30f725fe6d025557" +dependencies = [ + "aead", + "aes-gcm", + "chacha20poly1305", + "crypto-common", + "der", + "digest", + "ecdsa", + "ed25519-dalek", + "hmac", + "p256", + "p384", + "paste", + "pkcs8", + "rand_core 0.6.4", + "rsa", + "rustls", + "rustls-pki-types", + "rustls-webpki 0.102.8", + "sec1", + "sha2", + "signature", + "x25519-dalek", +] + +[[package]] +name = "rustls-webpki" +version = "0.102.8" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "64ca1bc8749bd4cf37b5ce386cc146580777b4e8572c7b97baf22c83f444bee9" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + +[[package]] +name = "rustls-webpki" +version = "0.103.15" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f3c3cf1d8b1e7d4927e2d154c3fcb02979afb9939629c62cd9048d4f07b60ac2" +dependencies = [ + "ring", + "rustls-pki-types", + "untrusted", +] + [[package]] name = "rustversion" version = "1.0.23" @@ -4267,6 +4873,20 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "94143f37725109f92c262ed2cf5e59bce7498c01bcc1502d7b9afe439a4e9f49" +[[package]] +name = "sec1" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d3e97a565f76233a6003f9f5c54be1d9c5bdfa3eccfb189469f11ec4901c47dc" +dependencies = [ + "base16ct", + "der", + "generic-array", + "pkcs8", + "subtle", + "zeroize", +] + [[package]] name = "self_cell" version = "1.3.0" @@ -4348,6 +4968,17 @@ dependencies = [ "serde_core", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures", + "digest", +] + [[package]] name = "shlex" version = "2.0.1" @@ -4364,6 +4995,16 @@ dependencies = [ "libc", ] +[[package]] +name = "signature" +version = "2.2.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "77549399552de45a898a580c1b41d445bf730df867cc44e6c0233bbc4b8329de" +dependencies = [ + "digest", + "rand_core 0.6.4", +] + [[package]] name = "simd-adler32" version = "0.3.10" @@ -4504,6 +5145,12 @@ dependencies = [ "serde", ] +[[package]] +name = "spin" +version = "0.9.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3763264f6b73151db08c50ff20d7d8a0b8796e021cdea7ceedad07b80155fa0e" + [[package]] name = "spirv" version = "0.4.0+sdk-1.4.341.0" @@ -4513,6 +5160,16 @@ dependencies = [ "bitflags 2.13.2", ] +[[package]] +name = "spki" +version = "0.7.3" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d91ed6c858b01f942cd56b37a94b3e0a1798290327d1236e4d9cf4eaca44d29d" +dependencies = [ + "base64ct", + "der", +] + [[package]] name = "stable_deref_trait" version = "1.2.1" @@ -4525,6 +5182,12 @@ version = "1.1.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "a2eb9349b6444b326872e140eb1cf5e7c522154d69e7a0ffb0fb81c06b37543f" +[[package]] +name = "subtle" +version = "2.6.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "13c2bddecc57b384dee18652358fb23172facb8a2c51ccc10d74c157bdea3292" + [[package]] name = "syn" version = "2.0.119" @@ -4778,6 +5441,12 @@ version = "0.12.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8e28f89b80c87b8fb0cf04ab448d5dd0dd0ade2f8891bae878de66a75a28600e" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "uds_windows" version = "1.2.1" @@ -4854,6 +5523,22 @@ version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b4ac048d71ede7ee76d585517add45da530660ef4390e49b098733c6e897f254" +[[package]] +name = "universal-hash" +version = "0.5.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "fc1de2c688dc15305988b563c3854064043356019f97a4b46276fe734c4f07ea" +dependencies = [ + "crypto-common", + "subtle", +] + +[[package]] +name = "untrusted" +version = "0.9.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "8ecb6da28b8a351d773b68d5825ac39017e680750f980f3a1a85cd8dd28a47c1" + [[package]] name = "uuid" version = "1.26.1" @@ -4928,6 +5613,12 @@ dependencies = [ "winapi-util", ] +[[package]] +name = "wasi" +version = "0.11.1+wasi-snapshot-preview1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "ccf3ec651a847eb01de73ccad15eb7d99f80485de043efb2f370cd654f4ea44b" + [[package]] name = "wasip2" version = "1.0.4+wasi-0.2.12" @@ -5148,6 +5839,15 @@ dependencies = [ "wasm-bindgen", ] +[[package]] +name = "webpki-roots" +version = "1.0.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "7dcd9d09a39985f5344844e66b0c530a33843579125f23e21e9f0f220850f22a" +dependencies = [ + "rustls-pki-types", +] + [[package]] name = "weezl" version = "0.1.12" @@ -5729,6 +6429,17 @@ version = "0.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ea6fc2961e4ef194dcbfe56bb845534d0dc8098940c7e5c012a258bfec6701bd" +[[package]] +name = "x25519-dalek" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7e468321c81fb07fa7f4c636c3972b9100f0346e5b6a9f2bd0603a52f7ed277" +dependencies = [ + "curve25519-dalek", + "rand_core 0.6.4", + "zeroize", +] + [[package]] name = "xcursor" version = "0.3.11" @@ -5969,6 +6680,26 @@ dependencies = [ "synstructure 0.14.0", ] +[[package]] +name = "zeroize" +version = "1.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "e13084392c5e4bc371903e2935a5eaeed24905a7511356b883835e18a78f6879" +dependencies = [ + "zeroize_derive", +] + +[[package]] +name = "zeroize_derive" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3c50655cbb0fe3fc43170059e702f1ce5e19b84cec58dc87b037a09935c2f328" +dependencies = [ + "proc-macro2", + "quote", + "syn 2.0.119", +] + [[package]] name = "zip" version = "7.2.0" diff --git a/crates/segment/Cargo.toml b/crates/segment/Cargo.toml index 336b052fe..bd7b0d665 100644 --- a/crates/segment/Cargo.toml +++ b/crates/segment/Cargo.toml @@ -2,7 +2,7 @@ name = "lightcraft-segment" version.workspace = true edition.workspace = true -license.workspace = true +license = "Apache-2.0" description = "LightCraft segmentation: SAM 3 (Segment Anything with Concepts) inference in pure Rust (L3)" [lints] @@ -18,6 +18,10 @@ log = { workspace = true } [target.'cfg(not(target_arch = "wasm32"))'.dependencies] candle-core = { workspace = true } candle-nn = { workspace = true } +rustls = { version = "0.23", default-features = false, features = ["std", "tls12"] } +rustls-rustcrypto = "0.0.2-alpha" +webpki-roots = "1" +sha2 = { version = "0.10", default-features = false, features = ["std"] } [target.'cfg(target_os = "macos")'.dependencies] candle-core = { workspace = true, features = ["metal"] } diff --git a/crates/segment/src/clip.rs b/crates/segment/src/clip.rs index c392d5c5f..b69c148a4 100644 --- a/crates/segment/src/clip.rs +++ b/crates/segment/src/clip.rs @@ -1,5 +1,9 @@ //! The CLIP text encoder (24 layers, 1024 wide, causal) that turns a prompt into the token //! features SAM 3's detector attends to. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, CLIP text model), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{Device, Module, Tensor}; use candle_nn::{Embedding, LayerNorm, VarBuilder}; diff --git a/crates/segment/src/detector.rs b/crates/segment/src/detector.rs index b21b32d84..7c4eccbd3 100644 --- a/crates/segment/src/detector.rs +++ b/crates/segment/src/detector.rs @@ -1,6 +1,10 @@ //! Text prompts (SAM 3's concept detector): the prompt's CLIP features fused into the image //! features by a DETR encoder, 200 object queries refined by a DETR decoder (with box //! relative-position bias and a presence token), per-query masks from a pixel decoder. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{D, Device, IndexOp, Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, GroupNorm, LayerNorm, Linear, VarBuilder}; diff --git a/crates/segment/src/fetch/http.rs b/crates/segment/src/fetch/http.rs new file mode 100644 index 000000000..332f01477 --- /dev/null +++ b/crates/segment/src/fetch/http.rs @@ -0,0 +1,541 @@ +//! A small HTTP/1.1 GET client for model downloads, in pure Rust: `std::net` sockets and, for +//! `https://`, rustls with the RustCrypto provider (`rustls-rustcrypto`; no `ring` or +//! `aws-lc-rs`, so nothing is compiled from C or assembly) and the Mozilla root certificates +//! (`webpki-roots`). +//! +//! Only what a large-file download needs: one request per connection (`Connection: close`), +//! `Range` requests for resuming, redirects (followed by the caller), `Content-Length`, +//! chunked and read-to-close bodies. Every read waits at most `stall` for data and checks the +//! cancel flag at least twice a second, so a dead server or a cancelled download never hangs +//! the downloading thread. No proxies (see docs/ai-masks.md). + +use std::io::{ErrorKind, Read, Write}; +use std::net::{TcpStream, ToSocketAddrs}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::{Duration, Instant}; + +/// Why a request failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum HttpError { + /// The URL can't be used (scheme, host, characters). + BadUrl(String), + /// Name resolution or connecting failed (or took longer than the connect timeout). + Connect(String), + /// No data arrived for the stall timeout. + Stalled, + /// The cancel flag was raised. + Cancelled, + /// TLS failed (certificate, handshake). + Tls(String), + /// The server's response could not be understood. + Protocol(String), + /// Any other I/O error. + Io(String), +} + +impl std::fmt::Display for HttpError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + HttpError::BadUrl(e) => write!(f, "unusable URL: {e}"), + HttpError::Connect(e) => write!(f, "could not connect: {e}"), + HttpError::Stalled => write!(f, "the server stopped sending data"), + HttpError::Cancelled => write!(f, "cancelled"), + HttpError::Tls(e) => write!(f, "secure connection failed: {e}"), + HttpError::Protocol(e) => write!(f, "unexpected response: {e}"), + HttpError::Io(e) => write!(f, "{e}"), + } + } +} + +/// A parsed `http://` or `https://` URL. +#[derive(Debug, Clone, PartialEq, Eq)] +pub struct Url { + pub tls: bool, + pub host: String, + pub port: u16, + /// Path and query, starting with `/`. + pub path: String, +} + +impl Url { + pub fn parse(s: &str) -> Result { + let bad = |why: &str| HttpError::BadUrl(format!("{s}: {why}")); + let s = s.trim(); + if s.chars().any(|c| c.is_control() || c == ' ') { + return Err(bad("contains spaces or control characters")); + } + let (tls, rest) = if let Some(r) = s.strip_prefix("https://") { + (true, r) + } else if let Some(r) = s.strip_prefix("http://") { + (false, r) + } else { + return Err(bad("only http:// and https:// are supported")); + }; + let rest = rest.split('#').next().unwrap_or_default(); + let (authority, path) = match rest.find(['/', '?']) { + Some(i) => (rest.get(..i).unwrap_or_default(), rest.get(i..).unwrap_or_default()), + None => (rest, ""), + }; + let path = if path.starts_with('/') { path.to_string() } else { format!("/{path}") }; + if authority.contains('@') { + return Err(bad("user names and passwords in URLs are not supported")); + } + let (host, port) = if let Some(v6) = authority.strip_prefix('[') { + let end = v6.find(']').ok_or_else(|| bad("unterminated IPv6 address"))?; + let host = v6.get(..end).unwrap_or_default(); + let after = v6.get(end + 1..).unwrap_or_default(); + (host.to_string(), after.strip_prefix(':')) + } else { + match authority.rsplit_once(':') { + Some((h, p)) => (h.to_string(), Some(p)), + None => (authority.to_string(), None), + } + }; + let port = match port { + Some(p) => p.parse::().ok().filter(|p| *p != 0).ok_or_else(|| bad("bad port"))?, + None if tls => 443, + None => 80, + }; + if host.is_empty() || !host.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '-' | '.' | ':' | '_')) { + return Err(bad("bad host name")); + } + Ok(Url { tls, host: host.to_ascii_lowercase(), port, path }) + } + + /// `Location` of a redirect, resolved against this URL. + pub fn join(&self, location: &str) -> Result { + let l = location.trim(); + if l.starts_with("http://") || l.starts_with("https://") { + return Url::parse(l); + } + // another scheme (`ftp:`, `file:`, `data:`…) + if let Some((scheme, _)) = l.split_once(':') + && !scheme.is_empty() + && scheme.chars().all(|c| c.is_ascii_alphanumeric() || matches!(c, '+' | '-' | '.')) + { + return Err(HttpError::BadUrl(format!("redirect to an unsupported location `{scheme}:`"))); + } + let scheme = if self.tls { "https" } else { "http" }; + let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() }; + let base = format!("{scheme}://{host}:{}", self.port); + if let Some(rest) = l.strip_prefix("//") { + return Url::parse(&format!("{scheme}://{rest}")); + } + if l.starts_with('/') { + return Url::parse(&format!("{base}{l}")); + } + // relative to the current directory + let dir = self.path.split('?').next().unwrap_or("/"); + let dir = dir.rsplit_once('/').map_or("", |(d, _)| d); + Url::parse(&format!("{base}{dir}/{l}")) + } + + fn host_header(&self) -> String { + let host = if self.host.contains(':') { format!("[{}]", self.host) } else { self.host.clone() }; + let default = if self.tls { 443 } else { 80 }; + if self.port == default { host } else { format!("{host}:{}", self.port) } + } + + /// The URL without its query (for messages: a query may hold a signed token). + pub fn display(&self) -> String { + let scheme = if self.tls { "https" } else { "http" }; + format!("{scheme}://{}{}", self.host_header(), self.path.split('?').next().unwrap_or("/")) + } +} + +/// Timeouts and the cancel flag for one request. +pub struct Limits<'a> { + pub connect: Duration, + pub stall: Duration, + pub cancel: &'a AtomicBool, +} + +/// How often a blocked read wakes up to check the cancel flag. +const TICK: Duration = Duration::from_millis(250); +/// Longest status line + headers accepted. +const MAX_HEAD: usize = 64 * 1024; + +enum Stream { + Plain(TcpStream), + Tls(Box, TcpStream), +} + +fn is_timeout(e: &std::io::Error) -> bool { + matches!(e.kind(), ErrorKind::WouldBlock | ErrorKind::TimedOut | ErrorKind::Interrupted) +} + +fn tls_config() -> Result, HttpError> { + let roots = rustls::RootCertStore { roots: webpki_roots::TLS_SERVER_ROOTS.to_vec() }; + let config = rustls::ClientConfig::builder_with_provider(Arc::new(rustls_rustcrypto::provider())) + .with_safe_default_protocol_versions() + .map_err(|e| HttpError::Tls(e.to_string()))? + .with_root_certificates(roots) + .with_no_client_auth(); + Ok(Arc::new(config)) +} + +impl Stream { + fn connect(url: &Url, limits: &Limits) -> Result { + let addrs: Vec<_> = (url.host.as_str(), url.port).to_socket_addrs().map_err(|e| HttpError::Connect(format!("{}: {e}", url.host)))?.collect(); + if addrs.is_empty() { + return Err(HttpError::Connect(format!("{}: no address", url.host))); + } + let mut last = String::new(); + let mut sock = None; + for a in addrs { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match TcpStream::connect_timeout(&a, limits.connect) { + Ok(s) => { + sock = Some(s); + break; + } + Err(e) => last = format!("{}: {e}", url.host), + } + } + let sock = sock.ok_or(HttpError::Connect(last))?; + sock.set_read_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?; + sock.set_write_timeout(Some(TICK)).map_err(|e| HttpError::Io(e.to_string()))?; + let _ = sock.set_nodelay(true); + if !url.tls { + return Ok(Stream::Plain(sock)); + } + let name = rustls::pki_types::ServerName::try_from(url.host.clone()).map_err(|e| HttpError::Tls(format!("{}: {e}", url.host)))?; + let conn = rustls::ClientConnection::new(tls_config()?, name).map_err(|e| HttpError::Tls(e.to_string()))?; + let mut s = Stream::Tls(Box::new(conn), sock); + s.handshake(limits)?; + Ok(s) + } + + /// Wait for `f` to make progress, retrying timeouts until `stall` passes without any. + fn retry(limits: &Limits, mut f: impl FnMut() -> std::io::Result) -> Result { + let start = Instant::now(); + loop { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match f() { + Ok(v) => return Ok(v), + Err(e) if is_timeout(&e) => { + if start.elapsed() >= limits.stall { + return Err(HttpError::Stalled); + } + } + Err(e) if e.kind() == ErrorKind::InvalidData => return Err(HttpError::Tls(e.to_string())), + Err(e) => return Err(HttpError::Io(e.to_string())), + } + } + } + + fn handshake(&mut self, limits: &Limits) -> Result<(), HttpError> { + let Stream::Tls(conn, sock) = self else { return Ok(()) }; + while conn.is_handshaking() { + Self::retry(limits, || conn.complete_io(sock))?; + } + Ok(()) + } + + fn write_all(&mut self, data: &[u8], limits: &Limits) -> Result<(), HttpError> { + match self { + Stream::Plain(sock) => { + let mut rest = data; + while !rest.is_empty() { + let n = Self::retry(limits, || sock.write(rest))?; + if n == 0 { + return Err(HttpError::Io("connection closed while sending the request".into())); + } + rest = rest.get(n..).unwrap_or_default(); + } + Ok(()) + } + Stream::Tls(conn, sock) => { + // buffered by rustls, then sent as TLS records + conn.writer().write_all(data).map_err(|e| HttpError::Tls(e.to_string()))?; + while conn.wants_write() { + Self::retry(limits, || conn.write_tls(sock))?; + } + Ok(()) + } + } + } + + /// Read some bytes (0 at the end of the stream), waiting at most `stall` for them. + fn read(&mut self, buf: &mut [u8], limits: &Limits) -> Result { + match self { + Stream::Plain(sock) => Self::retry(limits, || sock.read(buf)), + Stream::Tls(conn, sock) => { + let start = Instant::now(); + loop { + if limits.cancel.load(Ordering::Relaxed) { + return Err(HttpError::Cancelled); + } + match conn.reader().read(buf) { + Ok(n) => return Ok(n), + // closed without close_notify: the caller checks the length it got + Err(e) if e.kind() == ErrorKind::UnexpectedEof => return Ok(0), + Err(e) if e.kind() == ErrorKind::WouldBlock => {} + Err(e) => return Err(HttpError::Tls(e.to_string())), + } + // no plaintext buffered: more TLS records from the socket + match conn.read_tls(sock) { + Ok(0) => return Ok(0), + Ok(_) => { + conn.process_new_packets().map_err(|e| HttpError::Tls(e.to_string()))?; + while conn.wants_write() { + Self::retry(limits, || conn.write_tls(sock))?; + } + } + Err(e) if is_timeout(&e) => { + if start.elapsed() >= limits.stall { + return Err(HttpError::Stalled); + } + } + Err(e) => return Err(HttpError::Io(e.to_string())), + } + } + } + } + } +} + +/// The connection with a read buffer (headers and chunk sizes are read line by line). +struct Conn { + stream: Stream, + buf: Vec, + pos: usize, +} + +impl Conn { + fn fill(&mut self, limits: &Limits) -> Result { + if self.pos >= self.buf.len() { + self.buf.clear(); + self.pos = 0; + } + let mut tmp = [0u8; 64 * 1024]; + let n = self.stream.read(&mut tmp, limits)?; + self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default()); + Ok(n) + } + + /// One line without its CRLF; at most `max` bytes. + fn line(&mut self, max: usize, limits: &Limits) -> Result { + loop { + let avail = self.buf.get(self.pos..).unwrap_or_default(); + if let Some(i) = avail.iter().position(|b| *b == b'\n') { + let line = avail.get(..i).unwrap_or_default(); + let line = line.strip_suffix(b"\r").unwrap_or(line); + let s = String::from_utf8_lossy(line).into_owned(); + self.pos = self.pos.saturating_add(i + 1); + return Ok(s); + } + if avail.len() > max { + return Err(HttpError::Protocol("header line too long".into())); + } + // keep the partial line, read more after it + if self.pos > 0 { + self.buf.drain(..self.pos.min(self.buf.len())); + self.pos = 0; + } + let mut tmp = [0u8; 16 * 1024]; + let n = self.stream.read(&mut tmp, limits)?; + if n == 0 { + return Err(HttpError::Protocol("connection closed in the middle of the response".into())); + } + self.buf.extend_from_slice(tmp.get(..n).unwrap_or_default()); + } + } + + /// Up to `out.len()` body bytes (0 at the end of the stream). + fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result { + if self.pos >= self.buf.len() && self.fill(limits)? == 0 { + return Ok(0); + } + let avail = self.buf.get(self.pos..).unwrap_or_default(); + let n = avail.len().min(out.len()); + if let (Some(dst), Some(src)) = (out.get_mut(..n), avail.get(..n)) { + dst.copy_from_slice(src); + } + self.pos = self.pos.saturating_add(n); + Ok(n) + } +} + +/// A response's status, headers and body. +pub struct Response { + pub status: u16, + headers: Vec<(String, String)>, + conn: Conn, + body: BodyMode, +} + +enum BodyMode { + Length(u64), + Chunked { left: u64, done: bool }, + Close, + Done, +} + +impl Response { + /// The first header named `name` (case-insensitive). + pub fn header(&self, name: &str) -> Option<&str> { + self.headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str()) + } + + /// The body's length when the server announced it. + pub fn content_length(&self) -> Option { + match self.body { + BodyMode::Length(n) => Some(n), + _ => None, + } + } + + /// Read body bytes into `out` (0 at the end of the body). + pub fn read(&mut self, out: &mut [u8], limits: &Limits) -> Result { + match &mut self.body { + BodyMode::Done => Ok(0), + BodyMode::Length(left) => { + if *left == 0 { + self.body = BodyMode::Done; + return Ok(0); + } + let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len()); + let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?; + if n == 0 { + return Err(HttpError::Protocol("the connection closed before the end of the file".into())); + } + *left = left.saturating_sub(n as u64); + Ok(n) + } + BodyMode::Close => { + let n = self.conn.read(out, limits)?; + if n == 0 { + self.body = BodyMode::Done; + } + Ok(n) + } + BodyMode::Chunked { left, done } => { + if *done { + return Ok(0); + } + if *left == 0 { + let line = self.conn.line(1024, limits)?; + let hex = line.split(';').next().unwrap_or_default().trim(); + let size = u64::from_str_radix(hex, 16).map_err(|_| HttpError::Protocol(format!("bad chunk size `{hex}`")))?; + if size == 0 { + // trailers, up to the empty line + let mut n = 0; + while !self.conn.line(MAX_HEAD, limits)?.is_empty() { + n += 1; + if n > 100 { + return Err(HttpError::Protocol("too many trailers".into())); + } + } + *done = true; + return Ok(0); + } + *left = size; + } + let want = usize::try_from(*left).unwrap_or(usize::MAX).min(out.len()); + let n = self.conn.read(out.get_mut(..want).unwrap_or_default(), limits)?; + if n == 0 { + return Err(HttpError::Protocol("the connection closed in the middle of a chunk".into())); + } + *left = left.saturating_sub(n as u64); + if *left == 0 && !self.conn.line(16, limits)?.is_empty() { + return Err(HttpError::Protocol("missing CRLF after a chunk".into())); + } + Ok(n) + } + } + } +} + +/// Send `GET url` with extra `headers` and read the status line and headers. +pub fn get(url: &Url, headers: &[(&str, String)], limits: &Limits) -> Result { + let mut req = format!( + "GET {} HTTP/1.1\r\nHost: {}\r\nUser-Agent: LightCraft/{}\r\nAccept: */*\r\nAccept-Encoding: identity\r\nConnection: close\r\n", + url.path, + url.host_header(), + env!("CARGO_PKG_VERSION") + ); + for (k, v) in headers { + if v.chars().any(|c| c == '\r' || c == '\n') || k.chars().any(|c| c == '\r' || c == '\n' || c == ':') { + return Err(HttpError::BadUrl("header with a line break".into())); + } + req.push_str(&format!("{k}: {v}\r\n")); + } + req.push_str("\r\n"); + let mut stream = Stream::connect(url, limits)?; + stream.write_all(req.as_bytes(), limits)?; + let mut conn = Conn { stream, buf: Vec::new(), pos: 0 }; + let status_line = conn.line(MAX_HEAD, limits)?; + let mut parts = status_line.split_whitespace(); + let version = parts.next().unwrap_or_default(); + if !version.starts_with("HTTP/1.") { + return Err(HttpError::Protocol(format!("not an HTTP/1.x response: {}", status_line.chars().take(60).collect::()))); + } + let status = parts.next().and_then(|s| s.parse::().ok()).ok_or_else(|| HttpError::Protocol("bad status line".into()))?; + let mut headers = Vec::new(); + let mut total = status_line.len(); + loop { + let line = conn.line(MAX_HEAD, limits)?; + if line.is_empty() { + break; + } + total = total.saturating_add(line.len()); + if total > MAX_HEAD || headers.len() > 200 { + return Err(HttpError::Protocol("response headers too large".into())); + } + if let Some((k, v)) = line.split_once(':') { + headers.push((k.trim().to_string(), v.trim().to_string())); + } + } + let find = |name: &str| headers.iter().find(|(k, _)| k.eq_ignore_ascii_case(name)).map(|(_, v)| v.as_str()); + let body = if status == 204 || status == 304 || (100..200).contains(&status) { + BodyMode::Done + } else if find("transfer-encoding").is_some_and(|v| v.to_ascii_lowercase().contains("chunked")) { + BodyMode::Chunked { left: 0, done: false } + } else if let Some(n) = find("content-length") { + BodyMode::Length(n.trim().parse::().map_err(|_| HttpError::Protocol(format!("bad Content-Length `{n}`")))?) + } else { + BodyMode::Close + }; + Ok(Response { status, headers, conn, body }) +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn parses_urls() { + let u = Url::parse("https://cdn.example.com/models/sam3/model.safetensors?sig=1#x").unwrap(); + assert_eq!((u.tls, u.host.as_str(), u.port, u.path.as_str()), (true, "cdn.example.com", 443, "/models/sam3/model.safetensors?sig=1")); + assert_eq!(u.display(), "https://cdn.example.com/models/sam3/model.safetensors"); + let u = Url::parse("http://127.0.0.1:8080").unwrap(); + assert_eq!((u.tls, u.port, u.path.as_str()), (false, 8080, "/")); + let u = Url::parse("http://[::1]:9/a").unwrap(); + assert_eq!((u.host.as_str(), u.port), ("::1", 9)); + for bad in ["ftp://x/y", "https://", "https://a b/", "https://u:p@h/", "http://h:0/", "http://h:99999/", "https://h/\r\nX: y", "file:///etc"] + { + assert!(Url::parse(bad).is_err(), "{bad}"); + } + } + + #[test] + fn resolves_redirects() { + let u = Url::parse("https://a.example/x/y/file?q").unwrap(); + assert_eq!(u.join("https://b.example/z").unwrap().host, "b.example"); + assert_eq!(u.join("/root").unwrap().path, "/root"); + assert_eq!(u.join("other").unwrap().path, "/x/y/other"); + assert_eq!(u.join("//c.example/p").unwrap().host, "c.example"); + assert!(u.join("ftp://nope").is_err()); + } + + #[test] + fn the_tls_config_builds_with_the_pure_rust_provider() { + assert!(tls_config().is_ok()); + } +} diff --git a/crates/segment/src/fetch/mod.rs b/crates/segment/src/fetch/mod.rs new file mode 100644 index 000000000..92c5b6ea0 --- /dev/null +++ b/crates/segment/src/fetch/mod.rs @@ -0,0 +1,424 @@ +//! Downloading the SAM 3 model, only when the user asked for it: from an ordered list of mirrors +//! (the next one when a mirror fails), resuming partial files, with connect and stall timeouts, +//! a size cap, progress and cancellation, and checked (exact size, SHA-256 where pinned) before +//! each file is moved into the model folder. Runs on whatever thread calls [`download`] (the +//! engine uses a background thread; the UI thread never waits for the network). +//! +//! Files are written to `.part` next to their final place and renamed when verified, so a +//! model folder never holds a half-written file under a real name, and an interrupted download +//! resumes from where it stopped. A file that fails its check is deleted, never kept. +//! +//! The weights are never part of LightCraft (SAM License, see docs/ai-masks.md); where they are +//! downloaded from is configured by [`DEFAULT_MIRRORS`] and the user's own list. + +pub mod http; + +use std::fs::{File, OpenOptions}; +use std::io::{Read, Write}; +use std::path::{Path, PathBuf}; +use std::sync::atomic::{AtomicBool, Ordering}; +use std::time::Duration; + +use sha2::{Digest, Sha256}; + +use http::{HttpError, Limits, Url}; + +/// One file of the model. +#[derive(Clone, Debug, PartialEq, Eq)] +pub struct FileSpec { + pub name: &'static str, + /// Exact size in bytes, when pinned (then the download resumes and is checked against it). + pub size: Option, + /// SHA-256 (lower-case hex), when pinned. + pub sha256: Option<&'static str>, + /// Largest size accepted when `size` is not pinned. + pub max: u64, +} + +/// The `facebook/sam3` checkpoint files [`crate::Sam3::load`] needs. +/// +/// `model.safetensors` is pinned to the official checkpoint (Hugging Face LFS SHA-256). The two +/// tokenizer files are small text files without a pinned hash yet: they are only parsed by the +/// tokenizer (never executed) and capped in size. +// TODO(maintainer): pin `vocab.json` and `merges.txt` (size + SHA-256) when the files are +// uploaded to LightCraft's CDN. +pub const SAM3_FILES: &[FileSpec] = &[ + FileSpec { name: "vocab.json", size: None, sha256: None, max: 16 << 20 }, + FileSpec { name: "merges.txt", size: None, sha256: None, max: 16 << 20 }, + FileSpec { + name: crate::WEIGHTS_FILE, + size: Some(SAM3_WEIGHTS_SIZE), + sha256: Some("6d06f0a5f84e435071fe6603e61d0b4cc7b40e0d39d487cfd4d67d8cc11cc14a"), + max: SAM3_WEIGHTS_SIZE, + }, +]; + +/// Size of the official `model.safetensors` (3.44 GB). +pub const SAM3_WEIGHTS_SIZE: u64 = 3_439_938_512; + +/// Where LightCraft downloads the model from, in order: each is a base URL, and file `f` is at +/// `/`. **Empty for now**: LightCraft has no download location of its own yet, so the +/// in-app download needs the user's list (`LIGHTCRAFT_SAM3_MIRRORS` or the mirrors file, see +/// [`mirrors`] and docs/ai-masks.md) until this is filled in. +/// +/// Hugging Face's `facebook/sam3` can't be a default: it is gated (each user must accept the +/// SAM License there and download with their own token). +// TODO(maintainer): add LightCraft's CDN locations (https, primary first), e.g. +// "https:///models/sam3/" — the files there must match `SAM3_FILES`. +pub const DEFAULT_MIRRORS: &[&str] = &[]; + +/// Environment variable with extra mirrors (base URLs separated by commas, spaces or newlines), +/// tried before the defaults. +pub const MIRRORS_ENV: &str = "LIGHTCRAFT_SAM3_MIRRORS"; + +/// Most mirrors used (a hostile list can't make a download loop forever). +const MAX_MIRRORS: usize = 16; +/// Redirects followed per request. +const MAX_REDIRECTS: usize = 8; + +/// The mirrors to try, in order: the environment variable's, then the mirrors file's (one base +/// URL per line, `#` comments), then [`DEFAULT_MIRRORS`]; duplicates and unusable URLs dropped. +pub fn mirrors(env: Option<&str>, file: Option<&Path>) -> Vec { + let mut out: Vec = Vec::new(); + let mut add = |s: &str| { + let s = s.trim().trim_end_matches('/'); + if !s.is_empty() && Url::parse(s).is_ok() && !out.iter().any(|m| m == s) && out.len() < MAX_MIRRORS { + out.push(s.to_string()); + } + }; + if let Some(env) = env { + env.split([',', ' ', '\n', '\t']).for_each(&mut add); + } + // a small text file; anything larger isn't a mirror list + if let Some(text) = file.and_then(|f| std::fs::metadata(f).ok().filter(|m| m.len() < 64 * 1024).and_then(|_| std::fs::read_to_string(f).ok())) { + text.lines().map(|l| l.split('#').next().unwrap_or_default()).for_each(&mut add); + } + DEFAULT_MIRRORS.iter().for_each(|m| add(m)); + out +} + +/// Download settings. +#[derive(Clone, Debug)] +pub struct Options { + /// Longest wait to connect to a mirror. + pub connect_timeout: Duration, + /// Longest wait for the next bytes once connected. + pub stall_timeout: Duration, + /// Tries per mirror and file (each resumes where the last one stopped). + pub attempts: u32, +} + +impl Default for Options { + fn default() -> Self { + Options { connect_timeout: Duration::from_secs(15), stall_timeout: Duration::from_secs(30), attempts: 3 } + } +} + +/// Where a download is. +#[derive(Clone, Debug, Default, PartialEq, Eq)] +pub struct Progress { + pub file: String, + /// Bytes of all files on disk so far, and the total expected (pinned sizes + announced ones). + pub done: u64, + pub total: u64, + /// The mirror in use (without query). + pub mirror: String, +} + +/// Why a download failed. +#[derive(Debug, Clone, PartialEq, Eq)] +pub enum DownloadError { + /// No mirror is configured. + NoMirrors, + Cancelled, + /// A local file problem (folder not writable, disk full…). + Disk(String), + /// Every mirror failed for `file`; the reason for each. + AllMirrorsFailed { + file: String, + errors: Vec, + }, +} + +impl std::fmt::Display for DownloadError { + fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result { + match self { + DownloadError::NoMirrors => write!( + f, + "no download location is configured for the SAM 3 model in this build (set {MIRRORS_ENV}, or put the files in the model folder yourself; see docs/ai-masks.md)" + ), + DownloadError::Cancelled => write!(f, "the download was cancelled"), + DownloadError::Disk(e) => write!(f, "{e}"), + DownloadError::AllMirrorsFailed { file, errors } => write!(f, "could not download {file}: {}", errors.join("; ")), + } + } +} + +/// One mirror's failure for one file; `Retry` resumes on the same mirror. +enum Fail { + Retry(String), + NextMirror(String), + Cancelled, + Disk(String), +} + +/// Download `files` into `dir` from `mirrors` (in order; see [`mirrors`]). Files already there +/// and verified are kept. `progress` is called as bytes arrive (often: throttle in it). +pub fn download( + files: &[FileSpec], + mirrors: &[String], + dir: &Path, + opts: &Options, + cancel: &AtomicBool, + progress: &mut dyn FnMut(&Progress), +) -> Result<(), DownloadError> { + std::fs::create_dir_all(dir).map_err(|e| DownloadError::Disk(format!("can't create {}: {e}", dir.display())))?; + // what's left to fetch + let mut todo = Vec::new(); + for f in files { + if !verified(f, &dir.join(f.name))? { + todo.push(f); + } + } + if todo.is_empty() { + return Ok(()); + } + if mirrors.is_empty() { + return Err(DownloadError::NoMirrors); + } + let total: u64 = files.iter().map(|f| f.size.unwrap_or(0)).sum(); + let mut before: u64 = files.iter().filter(|f| !todo.contains(f)).filter_map(|f| f.size).sum(); + for f in todo { + let mut errors = Vec::new(); + let mut ok = false; + 'mirrors: for m in mirrors { + let url = format!("{m}/{}", f.name); + for _ in 0..opts.attempts.max(1) { + if cancel.load(Ordering::Relaxed) { + return Err(DownloadError::Cancelled); + } + let mut report = |n: u64, announced: u64| { + let mirror = Url::parse(m).map(|u| u.display()).unwrap_or_default(); + progress(&Progress { + file: f.name.to_string(), + done: before.saturating_add(n), + total: total.max(before.saturating_add(announced)), + mirror, + }) + }; + match fetch_file(f, &url, dir, opts, cancel, &mut report) { + Ok(()) => { + ok = true; + break 'mirrors; + } + Err(Fail::Cancelled) => return Err(DownloadError::Cancelled), + Err(Fail::Disk(e)) => return Err(DownloadError::Disk(e)), + Err(Fail::Retry(e)) => { + log::warn!("SAM 3 download of {url}: {e} (retrying)"); + errors.push(format!("{}: {e}", short(m))); + } + Err(Fail::NextMirror(e)) => { + log::warn!("SAM 3 download of {url}: {e}"); + errors.push(format!("{}: {e}", short(m))); + continue 'mirrors; + } + } + } + } + if !ok { + return Err(DownloadError::AllMirrorsFailed { file: f.name.to_string(), errors }); + } + before = before.saturating_add(std::fs::metadata(dir.join(f.name)).map(|m| m.len()).unwrap_or(0)); + } + Ok(()) +} + +/// A mirror for messages (no query string: it may hold a token). +fn short(m: &str) -> String { + Url::parse(m).map(|u| u.display()).unwrap_or_else(|_| "mirror".into()) +} + +/// Whether `path` is already the right file (exact size and hash when pinned). +fn verified(f: &FileSpec, path: &Path) -> Result { + let Ok(meta) = std::fs::metadata(path) else { return Ok(false) }; + if !meta.is_file() || meta.len() == 0 { + return Ok(false); + } + if f.size.is_some_and(|s| s != meta.len()) || meta.len() > f.max { + return Ok(false); + } + match f.sha256 { + Some(want) => Ok(sha256_file(path).map_err(DownloadError::Disk)? == want), + None => Ok(true), + } +} + +fn sha256_file(path: &Path) -> Result { + let mut file = File::open(path).map_err(|e| format!("{}: {e}", path.display()))?; + let mut h = Sha256::new(); + let mut buf = vec![0u8; 1 << 20]; + loop { + let n = file.read(&mut buf).map_err(|e| format!("{}: {e}", path.display()))?; + if n == 0 { + break; + } + h.update(buf.get(..n).unwrap_or_default()); + } + Ok(h.finalize().iter().map(|b| format!("{b:02x}")).collect()) +} + +fn part_path(dir: &Path, name: &str) -> PathBuf { + dir.join(format!("{name}.part")) +} + +/// Fetch one file from one mirror into `.part` (resuming it when its size is pinned), +/// check it and move it into place. +fn fetch_file(f: &FileSpec, url: &str, dir: &Path, opts: &Options, cancel: &AtomicBool, report: &mut dyn FnMut(u64, u64)) -> Result<(), Fail> { + let part = part_path(dir, f.name); + let limit = f.size.unwrap_or(f.max); + let disk = |e: std::io::Error| Fail::Disk(format!("{}: {e}", part.display())); + // resume only a pinned file (an unpinned one can't be checked after resuming) + let mut have = match std::fs::metadata(&part) { + Ok(m) if f.size.is_some() && m.len() <= limit => m.len(), + Ok(_) => { + std::fs::remove_file(&part).map_err(disk)?; + 0 + } + Err(_) => 0, + }; + let fetched = f.size != Some(have) || have == 0; + if fetched { + let limits = Limits { connect: opts.connect_timeout, stall: opts.stall_timeout, cancel }; + let mut url = Url::parse(url).map_err(|e| Fail::NextMirror(e.to_string()))?; + let origin = url.host.clone(); + let mut redirects = 0; + let mut resp = loop { + let mut headers = vec![]; + if have > 0 { + headers.push(("Range", format!("bytes={have}-"))); + } + // a token for the configured host only (never sent on to a redirect's host) + if let Ok(t) = std::env::var("LIGHTCRAFT_SAM3_TOKEN") + && !t.trim().is_empty() + && url.tls + && url.host == origin + { + headers.push(("Authorization", format!("Bearer {}", t.trim()))); + } + let r = http::get(&url, &headers, &limits).map_err(http_fail)?; + match r.status { + 301 | 302 | 303 | 307 | 308 => { + redirects += 1; + if redirects > MAX_REDIRECTS { + return Err(Fail::NextMirror("too many redirects".into())); + } + let loc = r.header("location").ok_or_else(|| Fail::NextMirror("redirect without a location".into()))?; + url = url.join(loc).map_err(|e| Fail::NextMirror(e.to_string()))?; + } + _ => break r, + } + }; + let mut file = match resp.status { + 200 => { + // the whole file (from the start, also when a range was asked for) + if resp.content_length().is_some_and(|n| n > limit || f.size.is_some_and(|s| s != n)) { + return Err(Fail::NextMirror(format!("the server's file has the wrong size ({:?} bytes)", resp.content_length()))); + } + have = 0; + File::create(&part).map_err(disk)? + } + 206 if have > 0 => { + let range = resp.header("content-range").unwrap_or_default(); + let (start, total) = parse_content_range(range).ok_or_else(|| Fail::NextMirror(format!("bad Content-Range `{range}`")))?; + if start != have || total.is_some_and(|t| Some(t) != f.size) { + return Err(Fail::NextMirror(format!("the server resumed at the wrong place ({range})"))); + } + OpenOptions::new().append(true).open(&part).map_err(disk)? + } + 416 => { + // the partial file doesn't fit what the server has: start over + std::fs::remove_file(&part).map_err(disk)?; + return Err(Fail::Retry("the partial download didn't match; starting over".into())); + } + 401 | 403 => return Err(Fail::NextMirror(format!("access denied (HTTP {})", resp.status))), + 404 | 410 => return Err(Fail::NextMirror(format!("not found (HTTP {})", resp.status))), + s if s >= 500 => return Err(Fail::Retry(format!("server error (HTTP {s})"))), + s => return Err(Fail::NextMirror(format!("HTTP {s}"))), + }; + let announced = f.size.unwrap_or_else(|| resp.content_length().unwrap_or(0)); + report(have, announced); + let mut buf = vec![0u8; 256 * 1024]; + loop { + let n = match resp.read(&mut buf, &limits) { + Ok(n) => n, + Err(e) => { + let _ = file.flush(); + if f.size.is_none() { + drop(file); + let _ = std::fs::remove_file(&part); + } + return Err(http_fail(e)); + } + }; + if n == 0 { + break; + } + if have.saturating_add(n as u64) > limit { + drop(file); + let _ = std::fs::remove_file(&part); + return Err(Fail::NextMirror(format!("the server sent more than the expected {limit} bytes"))); + } + file.write_all(buf.get(..n).unwrap_or_default()).map_err(disk)?; + have = have.saturating_add(n as u64); + report(have, announced); + } + file.sync_all().map_err(disk)?; + } + // checks: exact size, hash + let len = std::fs::metadata(&part).map_err(disk)?.len(); + if f.size.is_some_and(|s| s != len) { + return Err(Fail::Retry(format!("incomplete ({len} of {} bytes)", f.size.unwrap_or(0)))); + } + if len == 0 { + let _ = std::fs::remove_file(&part); + return Err(Fail::NextMirror("the server sent an empty file".into())); + } + if let Some(want) = f.sha256 { + let got = sha256_file(&part).map_err(Fail::Disk)?; + if got != want { + std::fs::remove_file(&part).map_err(disk)?; + let why = format!("the file is damaged or not the official one (SHA-256 {got})"); + // a damaged partial file from an earlier run isn't this mirror's fault + return Err(if fetched { Fail::NextMirror(why) } else { Fail::Retry(why) }); + } + } + std::fs::rename(&part, dir.join(f.name)).map_err(disk)?; + Ok(()) +} + +fn http_fail(e: HttpError) -> Fail { + match e { + HttpError::Cancelled => Fail::Cancelled, + HttpError::Stalled | HttpError::Io(_) | HttpError::Protocol(_) => Fail::Retry(e.to_string()), + HttpError::BadUrl(_) | HttpError::Connect(_) | HttpError::Tls(_) => Fail::NextMirror(e.to_string()), + } +} + +/// `bytes a-b/total` → (a, total) (`*` total → None). +fn parse_content_range(v: &str) -> Option<(u64, Option)> { + let v = v.trim().strip_prefix("bytes")?.trim(); + let (range, total) = v.split_once('/')?; + let (a, b) = range.split_once('-')?; + let (a, b) = (a.trim().parse::().ok()?, b.trim().parse::().ok()?); + if b < a { + return None; + } + let total = match total.trim() { + "*" => None, + t => Some(t.parse::().ok()?), + }; + Some((a, total)) +} + +#[cfg(test)] +mod tests; diff --git a/crates/segment/src/fetch/tests.rs b/crates/segment/src/fetch/tests.rs new file mode 100644 index 000000000..81e38f018 --- /dev/null +++ b/crates/segment/src/fetch/tests.rs @@ -0,0 +1,314 @@ +//! The downloader against a local HTTP server (std `TcpListener`; never the internet): mirror +//! fallback, redirects, resuming, chunked bodies, stalls, hash mismatches, size caps, cancel. + +use std::io::{BufRead, BufReader, Write}; +use std::net::{TcpListener, TcpStream}; +use std::sync::Arc; +use std::sync::atomic::{AtomicBool, AtomicUsize, Ordering}; +use std::time::{Duration, Instant}; + +use super::*; + +/// The served file: 300 kB of a deterministic pattern. +fn payload() -> Vec { + (0..300_000u32).map(|i| (i.wrapping_mul(2_654_435_761) >> 13) as u8).collect() +} + +fn sha_hex(b: &[u8]) -> String { + Sha256::digest(b).iter().map(|b| format!("{b:02x}")).collect() +} + +fn leak(s: String) -> &'static str { + Box::leak(s.into_boxed_str()) +} + +fn pinned() -> FileSpec { + let p = payload(); + FileSpec { name: "model.bin", size: Some(p.len() as u64), sha256: Some(leak(sha_hex(&p))), max: p.len() as u64 } +} + +struct Server { + base: String, + hits: Arc, +} + +/// Serve on 127.0.0.1; the first path segment picks the behaviour. +fn server() -> Server { + let l = TcpListener::bind("127.0.0.1:0").unwrap(); + let base = format!("http://{}", l.local_addr().unwrap()); + let hits = Arc::new(AtomicUsize::new(0)); + let h = hits.clone(); + std::thread::spawn(move || { + for s in l.incoming().flatten() { + let h = h.clone(); + std::thread::spawn(move || handle(s, &h)); + } + }); + Server { base, hits } +} + +fn handle(mut s: TcpStream, hits: &AtomicUsize) { + let n = hits.fetch_add(1, Ordering::SeqCst); + let mut r = BufReader::new(s.try_clone().unwrap()); + let mut first = String::new(); + if r.read_line(&mut first).is_err() { + return; + } + let mut range: Option = None; + loop { + let mut l = String::new(); + if r.read_line(&mut l).unwrap_or(0) == 0 || l.trim().is_empty() { + break; + } + if let Some(v) = l.to_ascii_lowercase().strip_prefix("range: bytes=") { + range = v.trim().trim_end_matches('-').parse().ok(); + } + } + let path = first.split_whitespace().nth(1).unwrap_or("/").to_string(); + let mode = path.trim_start_matches('/').split('/').next().unwrap_or("").to_string(); + let body = payload(); + let send = |s: &mut TcpStream, status: &str, headers: &str, data: &[u8]| { + let _ = s.write_all(format!("HTTP/1.1 {status}\r\n{headers}Connection: close\r\n\r\n").as_bytes()); + let _ = s.write_all(data); + }; + match mode.as_str() { + "missing" => send(&mut s, "404 Not Found", "Content-Length: 9\r\n", b"not found"), + "error" => send(&mut s, "500 Oops", "Content-Length: 25\r\n", b"server error"), + "redirect" => send(&mut s, "302 Found", &format!("Location: {}\r\nContent-Length: 0\r\n", path.replacen("/redirect", "/good", 1)), b""), + "wrong" => { + let mut b = body.clone(); + b[1000] ^= 0xff; + send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", b.len()), &b) + } + "toolong" => { + // no length announced, more bytes than the file can have + let mut b = body.clone(); + b.extend_from_slice(&[7u8; 5000]); + send(&mut s, "200 OK", "", &b) + } + "chunked" => { + let _ = s.write_all(b"HTTP/1.1 200 OK\r\nTransfer-Encoding: chunked\r\nConnection: close\r\n\r\n"); + for c in body.chunks(70_001) { + let _ = s.write_all(format!("{:x};ext=1\r\n", c.len()).as_bytes()); + let _ = s.write_all(c); + let _ = s.write_all(b"\r\n"); + } + let _ = s.write_all(b"0\r\nX-Trailer: 1\r\n\r\n"); + } + "stall" => { + // half the file, then nothing + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body[..body.len() / 2]); + std::thread::sleep(Duration::from_secs(5)); + } + "slow" => { + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + for c in body.chunks(1000) { + if s.write_all(c).is_err() { + return; + } + std::thread::sleep(Duration::from_millis(20)); + } + } + "flaky" if n.is_multiple_of(2) && range.is_none() => { + // the first try breaks off after a third + let _ = s.write_all(format!("HTTP/1.1 200 OK\r\nContent-Length: {}\r\n\r\n", body.len()).as_bytes()); + let _ = s.write_all(&body[..body.len() / 3]); + } + "range416" if range.is_some() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""), + "norange" | "range416" => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body), + _ => match range { + Some(from) if from as usize >= body.len() => send(&mut s, "416 Range Not Satisfiable", "Content-Length: 0\r\n", b""), + Some(from) => send( + &mut s, + "206 Partial Content", + &format!("Content-Range: bytes {from}-{}/{}\r\nContent-Length: {}\r\n", body.len() - 1, body.len(), body.len() - from as usize), + &body[from as usize..], + ), + None => send(&mut s, "200 OK", &format!("Content-Length: {}\r\n", body.len()), &body), + }, + } +} + +fn tmp(name: &str) -> PathBuf { + let d = std::env::temp_dir().join(format!("lc-sam3-fetch-{name}-{}", std::process::id())); + let _ = std::fs::remove_dir_all(&d); + d +} + +fn quick() -> Options { + Options { connect_timeout: Duration::from_secs(2), stall_timeout: Duration::from_millis(600), attempts: 2 } +} + +fn run(files: &[FileSpec], mirrors: &[String], dir: &Path, cancel: &AtomicBool) -> (Result<(), DownloadError>, Vec) { + let mut seen = Vec::new(); + let r = download(files, mirrors, dir, &quick(), cancel, &mut |p| seen.push(p.clone())); + (r, seen) +} + +#[test] +fn falls_back_to_the_next_mirror_and_verifies() { + let srv = server(); + let dir = tmp("fallback"); + let mirrors = vec![ + format!("{}/missing", srv.base), + format!("{}/error", srv.base), + "http://127.0.0.1:1/unreachable".to_string(), + format!("{}/redirect", srv.base), + ]; + let (r, seen) = run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)); + r.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + assert!(!part_path(&dir, "model.bin").exists()); + let last = seen.last().unwrap(); + assert_eq!((last.done, last.total), (300_000, 300_000)); + // an error page never ends up in the file: nothing left over from the failed mirrors + assert_eq!(std::fs::read_dir(&dir).unwrap().count(), 1); + // installed: nothing is fetched again + let hits = srv.hits.load(Ordering::SeqCst); + run(&[pinned()], &mirrors, &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(srv.hits.load(Ordering::SeqCst), hits); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_hash_mismatch_is_deleted_and_the_next_mirror_used() { + let srv = server(); + let dir = tmp("hash"); + let (r, _) = run(&[pinned()], &[format!("{}/wrong", srv.base)], &dir, &AtomicBool::new(false)); + let e = r.unwrap_err().to_string(); + assert!(e.contains("SHA-256"), "{e}"); + assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists(), "a damaged file is never kept"); + run(&[pinned()], &[format!("{}/wrong", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn resumes_partial_files_and_restarts_bad_ones() { + let srv = server(); + let dir = tmp("resume"); + std::fs::create_dir_all(&dir).unwrap(); + // a third already there + std::fs::write(part_path(&dir, "model.bin"), &payload()[..100_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a partial file with the wrong bytes: the hash fails, it's deleted, the next try starts over + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), vec![0u8; 100_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a server ignoring Range: the whole file again, not appended + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), &payload()[..5]).unwrap(); + run(&[pinned()], &[format!("{}/norange", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // a connection breaking off is resumed on the same mirror + std::fs::remove_file(dir.join("model.bin")).unwrap(); + run(&[pinned()], &[format!("{}/flaky", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_416_or_a_bad_complete_part_starts_over() { + let srv = server(); + let dir = tmp("416"); + std::fs::create_dir_all(&dir).unwrap(); + // the server refuses the range: the part is dropped and the file fetched whole (no loop) + std::fs::write(part_path(&dir, "model.bin"), &payload()[..100]).unwrap(); + run(&[pinned()], &[format!("{}/range416", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + // as long as the file but wrong: hash fails → deleted → downloaded again from the same mirror + std::fs::remove_file(dir.join("model.bin")).unwrap(); + std::fs::write(part_path(&dir, "model.bin"), vec![1u8; 300_000]).unwrap(); + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn chunked_bodies_and_unpinned_files() { + let srv = server(); + let dir = tmp("chunked"); + let f = FileSpec { name: "vocab.json", size: None, sha256: None, max: 1 << 20 }; + run(std::slice::from_ref(&f), &[format!("{}/chunked", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("vocab.json")).unwrap(), payload()); + // more than the cap: refused, nothing kept + let small = FileSpec { max: 1000, ..f }; + let dir2 = tmp("cap"); + let e = run(&[small], &[format!("{}/good", srv.base)], &dir2, &AtomicBool::new(false)).0.unwrap_err().to_string(); + assert!(e.contains("wrong size") || e.contains("more than"), "{e}"); + assert!(!dir2.join("vocab.json").exists() && !part_path(&dir2, "vocab.json").exists()); + let _ = std::fs::remove_dir_all(&dir); + let _ = std::fs::remove_dir_all(&dir2); +} + +#[test] +fn a_server_sending_too_much_is_cut_off() { + let srv = server(); + let dir = tmp("toolong"); + let e = run(&[pinned()], &[format!("{}/toolong", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap_err().to_string(); + assert!(e.contains("more than"), "{e}"); + assert!(!dir.join("model.bin").exists() && !part_path(&dir, "model.bin").exists()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn a_stalled_server_times_out_and_the_next_mirror_finishes() { + let srv = server(); + let dir = tmp("stall"); + let t = Instant::now(); + let (r, _) = run(&[pinned()], &[format!("{}/stall", srv.base), format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)); + r.unwrap(); + assert!(t.elapsed() < Duration::from_secs(4), "{:?}", t.elapsed()); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn cancel_stops_quickly_and_keeps_the_part_for_later() { + let srv = server(); + let dir = tmp("cancel"); + let cancel = Arc::new(AtomicBool::new(false)); + let c = cancel.clone(); + std::thread::spawn(move || { + std::thread::sleep(Duration::from_millis(400)); + c.store(true, Ordering::SeqCst); + }); + let t = Instant::now(); + let (r, _) = run(&[pinned()], &[format!("{}/slow", srv.base)], &dir, &cancel); + assert_eq!(r, Err(DownloadError::Cancelled)); + assert!(t.elapsed() < Duration::from_secs(3), "{:?}", t.elapsed()); + assert!(!dir.join("model.bin").exists()); + let kept = std::fs::metadata(part_path(&dir, "model.bin")).map(|m| m.len()).unwrap_or(0); + assert!(kept > 0 && kept < 300_000, "{kept}"); + // and the next run resumes it + run(&[pinned()], &[format!("{}/good", srv.base)], &dir, &AtomicBool::new(false)).0.unwrap(); + assert_eq!(std::fs::read(dir.join("model.bin")).unwrap(), payload()); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn no_mirrors_is_a_clear_error_and_mirror_lists_are_parsed() { + let dir = tmp("nomirror"); + let e = run(&[pinned()], &[], &dir, &AtomicBool::new(false)).0.unwrap_err(); + assert_eq!(e, DownloadError::NoMirrors); + assert!(e.to_string().contains(MIRRORS_ENV)); + let file = std::env::temp_dir().join(format!("lc-sam3-mirrors-{}.txt", std::process::id())); + std::fs::write(&file, "# mine\nhttps://b.example/m/ # second\n\nnot a url\nhttps://a.example/x\n").unwrap(); + let m = mirrors(Some("https://a.example/x/, ftp://no"), Some(&file)); + assert_eq!(m[..2], ["https://a.example/x".to_string(), "https://b.example/m".to_string()]); + assert_eq!(m.len(), 2 + DEFAULT_MIRRORS.len()); + assert_eq!(mirrors(None, Some(Path::new("/nonexistent/mirrors.txt"))).len(), DEFAULT_MIRRORS.len()); + let _ = std::fs::remove_file(&file); + let _ = std::fs::remove_dir_all(&dir); +} + +#[test] +fn content_ranges() { + assert_eq!(parse_content_range("bytes 100-299/300"), Some((100, Some(300)))); + assert_eq!(parse_content_range("bytes 0-0/*"), Some((0, None))); + assert_eq!(parse_content_range("bytes 9-1/300"), None); + assert_eq!(parse_content_range("items 1-2/3"), None); +} diff --git a/crates/segment/src/lib.rs b/crates/segment/src/lib.rs index 83bcbbafe..42aed5380 100644 --- a/crates/segment/src/lib.rs +++ b/crates/segment/src/lib.rs @@ -13,7 +13,16 @@ //! read from a directory the user downloads them to ([`Sam3::load`]). //! //! Ported from the Hugging Face `transformers` implementation (Apache-2.0); see `NOTICE`. -//! No UI dependencies (L3). +//! This crate is therefore licensed under the Apache License 2.0 only (not the MIT option of +//! the rest of LightCraft). No UI dependencies (L3). +//! +//! [`fetch`] downloads the model, only when the user asks for it (from configurable mirrors, +//! verified before use). Nothing in LightCraft requires the model: without it, AI masks report +//! that it isn't installed and everything else works. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`src/transformers/models/sam3` and `sam3_tracker`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. #![forbid(unsafe_code)] #![deny(clippy::unwrap_used, clippy::expect_used, clippy::panic, clippy::unimplemented, clippy::todo, clippy::unreachable)] @@ -21,6 +30,8 @@ mod clip; #[cfg(not(target_arch = "wasm32"))] mod detector; +#[cfg(not(target_arch = "wasm32"))] +pub mod fetch; pub mod mask; #[cfg(not(target_arch = "wasm32"))] mod neck; diff --git a/crates/segment/src/mask.rs b/crates/segment/src/mask.rs index 1e61c1f7c..af349a262 100644 --- a/crates/segment/src/mask.rs +++ b/crates/segment/src/mask.rs @@ -1,4 +1,8 @@ //! Image preprocessing and mask post-processing (CPU, no model needed). +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/image_processing_sam3*.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use crate::{Error, Result}; diff --git a/crates/segment/src/neck.rs b/crates/segment/src/neck.rs index df4db88d1..c4f716dd6 100644 --- a/crates/segment/src/neck.rs +++ b/crates/segment/src/neck.rs @@ -1,5 +1,9 @@ //! The feature pyramid (`Sam3VisionNeck`): the backbone's 72 × 72 map at four scales //! (288, 144, 72 and 36 per side), each projected to 256 channels. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, VarBuilder}; diff --git a/crates/segment/src/nn.rs b/crates/segment/src/nn.rs index cec6a58b1..540196e5d 100644 --- a/crates/segment/src/nn.rs +++ b/crates/segment/src/nn.rs @@ -1,4 +1,8 @@ //! Shared building blocks: attention, MLPs, layer norms, sine position encodings. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`, `models/sam3_tracker`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{D, DType, Device, Module, Tensor}; use candle_nn::{LayerNorm, Linear, VarBuilder}; diff --git a/crates/segment/src/tokenizer.rs b/crates/segment/src/tokenizer.rs index c1a9f3c10..0819aa48f 100644 --- a/crates/segment/src/tokenizer.rs +++ b/crates/segment/src/tokenizer.rs @@ -1,5 +1,9 @@ //! The CLIP byte-pair-encoding tokenizer (lower-cased, `` word ends), padded to the //! 32-token context SAM 3's text encoder takes. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (the CLIP tokenizer used by `models/sam3`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use std::collections::HashMap; use std::path::Path; diff --git a/crates/segment/src/tracker.rs b/crates/segment/src/tracker.rs index 828d56b66..d51219c80 100644 --- a/crates/segment/src/tracker.rs +++ b/crates/segment/src/tracker.rs @@ -1,5 +1,9 @@ //! Point prompts (SAM 3's interactive "tracker" head, SAM 2 style): a prompt encoder for //! positive/negative clicks and a two-way-transformer mask decoder over the image features. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3_tracker/modeling_sam3_tracker.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{DType, Device, IndexOp, Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, ConvTranspose2d, ConvTranspose2dConfig, LayerNorm, VarBuilder}; diff --git a/crates/segment/src/vit.rs b/crates/segment/src/vit.rs index ad9018269..16b09a729 100644 --- a/crates/segment/src/vit.rs +++ b/crates/segment/src/vit.rs @@ -1,5 +1,9 @@ //! The SAM 3 vision backbone: a 32-layer ViT (Perception Encoder) with 2-D axial rotary //! position encoding, windowed attention (24 × 24) in most layers and global attention in four. +//! +//! Modified work (Apache License 2.0, §4(b)): ported by the LightCraft contributors in 2026 from +//! the Python/PyTorch SAM 3 code of Hugging Face Transformers (`models/sam3/modeling_sam3.py`), Copyright The HuggingFace +//! Team and Meta Platforms, Inc.; translated to Rust on candle and restructured. See NOTICE. use candle_core::{D, Device, Module, Tensor}; use candle_nn::{Conv2d, Conv2dConfig, LayerNorm, VarBuilder}; diff --git a/crates/segment/src/weights.rs b/crates/segment/src/weights.rs index 15de564ae..70ac92b22 100644 --- a/crates/segment/src/weights.rs +++ b/crates/segment/src/weights.rs @@ -31,10 +31,11 @@ pub struct Weights { impl Weights { pub fn open(path: &Path) -> Result> { let mut file = File::open(path).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; + let file_len = file.metadata().map_err(|e| Error::Model(format!("{}: {e}", path.display())))?.len(); let mut len = [0u8; 8]; file.read_exact(&mut len).map_err(|e| Error::Model(format!("{}: {e}", path.display())))?; let n = u64::from_le_bytes(len); - if n > 100 << 20 { + if n > 100 << 20 || n.saturating_add(8) > file_len { return Err(Error::Model(format!("{}: not a safetensors file (header of {n} bytes)", path.display()))); } let mut header = vec![0u8; n as usize]; @@ -62,6 +63,10 @@ impl Weights { if end.checked_sub(start) != want.map(|w| w as u64) { return Err(Error::Model(format!("{name}: size does not match its shape"))); } + // the bytes must be in the file (a damaged or truncated download fails here, not later) + if (8 + n).checked_add(end).is_none_or(|e| e > file_len) { + return Err(Error::Model(format!("{}: tensor {name} lies outside the file (incomplete or damaged download?)", path.display()))); + } entries.insert(name, Entry { shape, start, end }); } Ok(Arc::new(Self { file: Mutex::new(file), entries, data_start: 8 + n })) @@ -74,7 +79,9 @@ impl Weights { /// Tensor `name` as stored, on the CPU. pub fn read(&self, name: &str) -> Result { let e = self.entries.get(name).ok_or_else(|| Error::Model(format!("missing tensor `{name}`")))?; - let mut bytes = vec![0u8; e.end.saturating_sub(e.start) as usize]; + // checked against the file's length in `open` + let len = usize::try_from(e.end.saturating_sub(e.start)).map_err(|_| Error::Model(format!("{name}: too large")))?; + let mut bytes = vec![0u8; len]; { let mut f = self.file.lock().unwrap_or_else(std::sync::PoisonError::into_inner); f.seek(SeekFrom::Start(self.data_start + e.start)).map_err(|err| Error::Model(format!("{name}: {err}")))?; @@ -106,3 +113,39 @@ impl SimpleBackend for Backend { self.0.contains(name) } } + +#[cfg(test)] +mod tests { + use super::*; + + fn write(name: &str, header: &str, data_len: usize) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("lc-weights-{name}-{}.safetensors", std::process::id())); + let mut b = (header.len() as u64).to_le_bytes().to_vec(); + b.extend_from_slice(header.as_bytes()); + b.extend(std::iter::repeat_n(0u8, data_len)); + std::fs::write(&p, b).unwrap(); + p + } + + #[test] + fn tensors_must_lie_inside_the_file() { + let ok = write("ok", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 16); + let w = Weights::open(&ok).unwrap(); + assert_eq!(w.read("a").unwrap().dims(), &[2, 2]); + // a truncated download: the header promises bytes that aren't there + let short = write("short", r#"{"a":{"dtype":"F32","shape":[2,2],"data_offsets":[0,16]}}"#, 8); + assert!(Weights::open(&short).err().unwrap().to_string().contains("outside the file")); + // hostile offsets: huge, or beyond u64 + let huge = write("huge", r#"{"a":{"dtype":"F32","shape":[1073741824,1073741824],"data_offsets":[0,4611686018427387904]}}"#, 16); + assert!(Weights::open(&huge).is_err()); + let wrap = write("wrap", r#"{"a":{"dtype":"F32","shape":[1],"data_offsets":[18446744073709551612,18446744073709551616]}}"#, 16); + assert!(Weights::open(&wrap).is_err()); + // a header longer than the file + let p = std::env::temp_dir().join(format!("lc-weights-hdr-{}.safetensors", std::process::id())); + std::fs::write(&p, 1000u64.to_le_bytes()).unwrap(); + assert!(Weights::open(&p).is_err()); + for f in [ok, short, huge, wrap, p] { + let _ = std::fs::remove_file(f); + } + } +}