diff --git a/Cargo.lock b/Cargo.lock index d41134a6b3..fa1f2cb8e2 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -510,7 +510,7 @@ dependencies = [ "fastrand", "hex", "http 1.4.1", - "sha1 0.10.6", + "sha1", "time", "tokio", "tracing", @@ -561,7 +561,6 @@ dependencies = [ "aws-credential-types", "aws-sigv4", "aws-smithy-async", - "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-runtime", "aws-smithy-runtime-api", @@ -570,9 +569,7 @@ dependencies = [ "bytes", "bytes-utils", "fastrand", - "http 0.2.12", "http 1.4.1", - "http-body 0.4.6", "http-body 1.0.1", "percent-encoding", "pin-project-lite", @@ -632,42 +629,6 @@ dependencies = [ "url", ] -[[package]] -name = "aws-sdk-s3" -version = "1.135.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f97e3e7e7d86fd26fcdc18bc382da5ca9e8b2ff8d54030d187fd0dac8a236d96" -dependencies = [ - "arc-swap", - "aws-credential-types", - "aws-runtime", - "aws-sigv4", - "aws-smithy-async", - "aws-smithy-checksums", - "aws-smithy-eventstream", - "aws-smithy-http", - "aws-smithy-json", - "aws-smithy-observability", - "aws-smithy-runtime", - "aws-smithy-runtime-api", - "aws-smithy-types", - "aws-smithy-xml", - "aws-types", - "bytes", - "fastrand", - "hex", - "hmac 0.13.0", - "http 0.2.12", - "http 1.4.1", - "http-body 1.0.1", - "lru", - "percent-encoding", - "regex-lite", - "sha2 0.11.0", - "tracing", - "url", -] - [[package]] name = "aws-sdk-secretsmanager" version = "1.107.0" @@ -826,24 +787,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "bae38512beae0ffee7010fc24e7a8a123c53efdfef42a61e80fda4882418dc71" dependencies = [ "aws-credential-types", - "aws-smithy-eventstream", "aws-smithy-http", "aws-smithy-runtime-api", "aws-smithy-types", "bytes", - "crypto-bigint", "form_urlencoded", "hex", "hmac 0.13.0", "http 0.2.12", "http 1.4.1", - "p256", "percent-encoding", "sha2 0.11.0", - "subtle", "time", "tracing", - "zeroize", ] [[package]] @@ -857,45 +813,12 @@ dependencies = [ "tokio", ] -[[package]] -name = "aws-smithy-checksums" -version = "0.64.8" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e9e8e65f4f81fcccdeb6c3eca2af17ac21d421a1786a26a394aecf421d616d3a" -dependencies = [ - "aws-smithy-http", - "aws-smithy-types", - "bytes", - "crc-fast", - "hex", - "http 1.4.1", - "http-body 1.0.1", - "http-body-util", - "md-5 0.11.0", - "pin-project-lite", - "sha1 0.11.0", - "sha2 0.11.0", - "tracing", -] - -[[package]] -name = "aws-smithy-eventstream" -version = "0.60.20" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "faf09d74e5e32f76b8762da505a3cd59303e367a664ca67295387baa8c1d7548" -dependencies = [ - "aws-smithy-types", - "bytes", - "crc32fast", -] - [[package]] name = "aws-smithy-http" version = "0.63.6" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "ba1ab2dc1c2c3749ead27180d333c42f11be8b0e934058fb4b2258ee8dbe5231" dependencies = [ - "aws-smithy-eventstream", "aws-smithy-runtime-api", "aws-smithy-types", "bytes", @@ -1219,12 +1142,6 @@ dependencies = [ "windows-link", ] -[[package]] -name = "base16ct" -version = "0.2.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4c7f02d4ea65f2c1853089ffd8d2787bdbc63de2f0d29dedbcf8ccdfa0ccd4cf" - [[package]] name = "base64" version = "0.22.1" @@ -1432,6 +1349,17 @@ version = "0.2.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "613afe47fcd5fac7ccf1db93babcb082c5994d996f20b8b159f2ad1658eb5724" +[[package]] +name = "chacha20" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d524456ba66e72eb8b115ff89e01e497f8e6d11d78b70b1aa13c0fbd97540a81" +dependencies = [ + "cfg-if", + "cpufeatures 0.3.0", + "rand_core 0.10.1", +] + [[package]] name = "chrono" version = "0.4.44" @@ -1662,25 +1590,6 @@ version = "2.5.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "217698eaf96b4a3f0bc4f3662aaa55bdf913cd54d7204591faa790070c6d0853" -[[package]] -name = "crc-fast" -version = "1.10.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e75b2483e97a5a7da73ac68a05b629f9c53cff58d8ed1c77866079e18b00dba5" -dependencies = [ - "digest 0.10.7", - "spin 0.10.0", -] - -[[package]] -name = "crc32fast" -version = "1.5.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" -dependencies = [ - "cfg-if", -] - [[package]] name = "criterion" version = "0.5.1" @@ -1781,18 +1690,6 @@ 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" @@ -2083,20 +1980,6 @@ version = "1.0.20" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d0881ea181b1df73ff77ffaaf9c7544ecc11e82fba9b5f27b262a3c73a332555" -[[package]] -name = "ecdsa" -version = "0.16.9" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "ee27f32b5c5292967d2d4a9d7f1e0b0aed2c15daded5a60300e4abb9d8020bca" -dependencies = [ - "der", - "digest 0.10.7", - "elliptic-curve", - "rfc6979", - "signature 2.2.0", - "spki", -] - [[package]] name = "ed25519" version = "1.5.3" @@ -2127,26 +2010,6 @@ dependencies = [ "serde", ] -[[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 0.10.7", - "ff", - "generic-array", - "group", - "pem-rfc7468", - "pkcs8", - "rand_core 0.6.4", - "sec1", - "subtle", - "zeroize", -] - [[package]] name = "encoding_rs" version = "0.8.35" @@ -2266,16 +2129,6 @@ version = "2.4.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9f1f227452a390804cdb637b74a86990f2a7d7ba4b7d5693aac9b4dd6defd8d6" -[[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 = "find-msvc-tools" version = "0.1.9" @@ -2484,7 +2337,6 @@ checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" dependencies = [ "typenum", "version_check", - "zeroize", ] [[package]] @@ -2523,6 +2375,7 @@ dependencies = [ "cfg-if", "libc", "r-efi 6.0.0", + "rand_core 0.10.1", "wasip2", "wasip3", ] @@ -2539,17 +2392,6 @@ version = "0.3.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "0cc23270f6e1808e30a928bdc84dea0b9b4136a8bc82338574f23baf47bbd280" -[[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 = "h2" version = "0.3.27" @@ -2628,8 +2470,6 @@ version = "0.16.1" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "841d1cc9bed7f9236f321df977030373f4a4163ae1a7dbfe1a51a2c1a51d9100" dependencies = [ - "allocator-api2", - "equivalent", "foldhash 0.2.0", ] @@ -2769,6 +2609,12 @@ version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" +[[package]] +name = "humantime" +version = "2.4.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15cdd26707701c53297e2fa6afb323d55fbc1d0810c3aec078ae3ef0424c3c15" + [[package]] name = "hybrid-array" version = "0.4.12" @@ -3091,7 +2937,7 @@ dependencies = [ "log", "num-format", "once_cell", - "quick-xml", + "quick-xml 0.26.0", "rgb", "str_stack", ] @@ -3119,7 +2965,6 @@ dependencies = [ "ampc-anon-stats", "ampc-server-utils", "aws-config", - "aws-sdk-s3", "aws-sdk-secretsmanager", "aws-sdk-sns", "aws-sdk-sqs", @@ -3139,6 +2984,7 @@ dependencies = [ "iris-mpc-utils", "itertools 0.13.0", "metrics", + "object_store", "pprof", "rand 0.8.6", "reqwest", @@ -3166,11 +3012,9 @@ dependencies = [ "async_from", "aws-config", "aws-sdk-kms", - "aws-sdk-s3", "aws-sdk-secretsmanager", "aws-sdk-sns", "aws-sdk-sqs", - "aws-smithy-types", "axum 0.7.9", "base64", "blake3", @@ -3198,6 +3042,7 @@ dependencies = [ "metrics-exporter-statsd", "metrics-tracing-context", "metrics-util 0.20.4", + "object_store", "rand 0.8.6", "reqwest", "ring", @@ -3226,10 +3071,10 @@ dependencies = [ "ampc-anon-stats", "ampc-secret-sharing", "ampc-server-utils", + "async-trait", "aws-config", "aws-credential-types", "aws-sdk-kms", - "aws-sdk-s3", "aws-sdk-secretsmanager", "aws-sdk-sns", "aws-sdk-sqs", @@ -3250,6 +3095,7 @@ dependencies = [ "metrics-exporter-statsd", "nix 0.29.0", "num_cpus", + "object_store", "percent-encoding", "rand 0.8.6", "rayon", @@ -3284,7 +3130,6 @@ dependencies = [ "ampc-server-utils", "async-channel", "async-trait", - "aws-sdk-s3", "backon", "base64", "bincode", @@ -3309,6 +3154,7 @@ dependencies = [ "num-traits", "num_cpus", "num_enum", + "object_store", "rand 0.8.6", "rand_chacha 0.3.1", "rand_distr", @@ -3386,13 +3232,14 @@ dependencies = [ "ampc-server-utils", "async-trait", "aws-config", - "aws-sdk-s3", "bytemuck", + "bytes", "dotenvy", "eyre", "futures", "iris-mpc-common", "itertools 0.13.0", + "object_store", "rand 0.8.6", "sqlx", "tokio", @@ -3438,7 +3285,6 @@ dependencies = [ "ampc-server-utils", "aws-config", "aws-sdk-rds", - "aws-sdk-s3", "axum 0.7.9", "bincode", "blake3", @@ -3452,6 +3298,7 @@ dependencies = [ "iris-mpc-utils", "itertools 0.13.0", "metrics", + "object_store", "rand 0.8.6", "serde", "serde_json", @@ -3474,7 +3321,6 @@ dependencies = [ "async_from", "aws-config", "aws-sdk-rds", - "aws-sdk-s3", "aws-sdk-secretsmanager", "aws-sdk-sns", "aws-sdk-sqs", @@ -3490,6 +3336,7 @@ dependencies = [ "iris-mpc-cpu", "iris-mpc-store", "itertools 0.13.0", + "object_store", "rand 0.8.6", "rayon", "serde", @@ -3689,15 +3536,6 @@ version = "0.4.31" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "113b30b4cd05f7c06868fdb2854f66a7b9fece9a48425351cd532e810d74024f" -[[package]] -name = "lru" -version = "0.16.4" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "7f66e8d5d03f609abc3a39e6f08e4164ebf1447a732906d39eb9b99b7919ef39" -dependencies = [ - "hashbrown 0.16.1", -] - [[package]] name = "lru-slab" version = "0.1.2" @@ -3745,16 +3583,6 @@ dependencies = [ "digest 0.10.7", ] -[[package]] -name = "md-5" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69b6441f590336821bb897fb28fc622898ccceb1d6cea3fde5ea86b090c4de98" -dependencies = [ - "cfg-if", - "digest 0.11.3", -] - [[package]] name = "memchr" version = "2.8.1" @@ -4121,6 +3949,46 @@ dependencies = [ "memchr", ] +[[package]] +name = "object_store" +version = "0.13.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "622acbc9100d3c10e2ee15804b0caa40e55c933d5aa53814cd520805b7958a49" +dependencies = [ + "async-trait", + "base64", + "bytes", + "chrono", + "form_urlencoded", + "futures-channel", + "futures-core", + "futures-util", + "http 1.4.1", + "http-body-util", + "httparse", + "humantime", + "hyper 1.10.1", + "itertools 0.14.0", + "md-5", + "parking_lot", + "percent-encoding", + "quick-xml 0.39.4", + "rand 0.10.2", + "reqwest", + "ring", + "rustls-pki-types", + "serde", + "serde_json", + "serde_urlencoded", + "thiserror 2.0.18", + "tokio", + "tracing", + "url", + "walkdir", + "wasm-bindgen-futures", + "web-time", +] + [[package]] name = "once_cell" version = "1.21.4" @@ -4280,18 +4148,6 @@ version = "0.5.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "1a80800c0488c3a21695ea981a54918fbb37abf04f4d0720c453632255e2ff0e" -[[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 0.10.9", -] - [[package]] name = "parking" version = "2.2.1" @@ -4536,15 +4392,6 @@ dependencies = [ "syn", ] -[[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" @@ -4611,7 +4458,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "be769465445e8c1474e9c5dac2018218498557af32d9ed057325ec9a41ae81bf" dependencies = [ "heck", - "itertools 0.10.5", + "itertools 0.14.0", "log", "multimap", "once_cell", @@ -4644,7 +4491,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8a56d757972c98b346a9b766e3f02746cde6dd1cd1d1d563472929fdd74bec4d" dependencies = [ "anyhow", - "itertools 0.10.5", + "itertools 0.14.0", "proc-macro2", "quote", "syn", @@ -4749,6 +4596,16 @@ dependencies = [ "memchr", ] +[[package]] +name = "quick-xml" +version = "0.39.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "cdcc8dd4e2f670d309a5f0e83fe36dfdc05af317008fea29144da1a2ac858e5e" +dependencies = [ + "memchr", + "serde", +] + [[package]] name = "quinn" version = "0.11.9" @@ -4856,6 +4713,17 @@ dependencies = [ "rand_core 0.9.5", ] +[[package]] +name = "rand" +version = "0.10.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c7f5fa3a058cd35567ef9bfa5e75732bee0f9e4c55fa90477bef2dfcdbc4be80" +dependencies = [ + "chacha20", + "getrandom 0.4.2", + "rand_core 0.10.1", +] + [[package]] name = "rand_chacha" version = "0.3.1" @@ -4895,6 +4763,12 @@ dependencies = [ "getrandom 0.3.4", ] +[[package]] +name = "rand_core" +version = "0.10.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "63b8176103e19a2643978565ca18b50549f6101881c443590420e4dc998a3c69" + [[package]] name = "rand_distr" version = "0.4.3" @@ -5087,6 +4961,7 @@ dependencies = [ "pin-project-lite", "quinn", "rustls 0.23.40", + "rustls-native-certs", "rustls-pki-types", "serde", "serde_json", @@ -5095,26 +4970,18 @@ dependencies = [ "tokio", "tokio-native-tls", "tokio-rustls 0.26.4", + "tokio-util", "tower 0.5.3", "tower-http", "tower-service", "url", "wasm-bindgen", "wasm-bindgen-futures", + "wasm-streams", "web-sys", "webpki-roots 1.0.7", ] -[[package]] -name = "rfc6979" -version = "0.4.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f8dd2a808d456c4a54e300a23e9f5a67e122c3024119acbfd73e3bf664491cb2" -dependencies = [ - "hmac 0.12.1", - "subtle", -] - [[package]] name = "rgb" version = "0.8.53" @@ -5391,20 +5258,6 @@ dependencies = [ "untrusted", ] -[[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 = "security-framework" version = "3.7.0" @@ -5586,17 +5439,6 @@ dependencies = [ "digest 0.10.7", ] -[[package]] -name = "sha1" -version = "0.11.0" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "aacc4cc499359472b4abe1bf11d0b12e688af9a805fa5e3016f9a386dc2d0214" -dependencies = [ - "cfg-if", - "cpufeatures 0.3.0", - "digest 0.11.3", -] - [[package]] name = "sha2" version = "0.10.9" @@ -5873,14 +5715,14 @@ dependencies = [ "hmac 0.12.1", "itoa", "log", - "md-5 0.10.6", + "md-5", "memchr", "once_cell", "percent-encoding", "rand 0.8.6", "rsa", "serde", - "sha1 0.10.6", + "sha1", "sha2 0.10.9", "smallvec", "sqlx-core", @@ -5912,7 +5754,7 @@ dependencies = [ "home", "itoa", "log", - "md-5 0.10.6", + "md-5", "memchr", "once_cell", "rand 0.8.6", @@ -6910,6 +6752,19 @@ dependencies = [ "wasmparser", ] +[[package]] +name = "wasm-streams" +version = "0.4.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "15053d8d85c7eccdbefef60f06769760a563c7f0a9d6902a13d35c7800b0ad65" +dependencies = [ + "futures-util", + "js-sys", + "wasm-bindgen", + "wasm-bindgen-futures", + "web-sys", +] + [[package]] name = "wasmparser" version = "0.244.0" diff --git a/Cargo.toml b/Cargo.toml index 8c561698a4..de44bb13f2 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,12 +31,13 @@ repository = "https://github.com/worldcoin/iris-mpc" [workspace.dependencies] async-stream = "0.3" aws-config = { version = "1.5.10", features = ["behavior-version-latest"] } +aws-credential-types = "1.2.10" aws-sdk-kms = { version = "1.44.0" } aws-sdk-rds = { version = "1.44.0" } aws-sdk-sns = { version = "1.44.0" } aws-sdk-sqs = { version = "1.89.0" } -aws-sdk-s3 = { version = "1.65.0" } aws-sdk-secretsmanager = { version = "1.47.0" } +object_store = { version = "0.13.1", features = ["aws", "azure", "gcp", "http"] } async-trait = "0.1.83" axum = "0.7" blake3 = "1.8.2" diff --git a/deny.toml b/deny.toml index 395d4f8250..9c326646fb 100644 --- a/deny.toml +++ b/deny.toml @@ -32,7 +32,6 @@ ignore = [ { id = "RUSTSEC-2026-0098", reason = "rustls-webpki 0.101.7 pulled transitively through aws-config 1.5.x via the legacy rustls 0.21 chain; bug silently accepts URI name constraints instead of rejecting them — only exploitable via CA misissuance of URI-constrained certs, and our TLS usage targets AWS endpoints where this isn't used. The 0.101.x line has no fix; drop once aws-config is bumped to drop the legacy rustls 0.21 chain." }, { id = "RUSTSEC-2026-0099", reason = "same rustls-webpki 0.101.7 chain as RUSTSEC-2026-0098; sibling bug accepts DNS name constraints for wildcard-asserting certs. Same misissuance-required exploit model, same AWS-TLS-only attack surface, same fix path (bump aws-config off the legacy rustls 0.21 chain)." }, { id = "RUSTSEC-2026-0104", reason = "rustls-webpki 0.101.7 pulled transitively through aws-config 1.5.x via the legacy rustls 0.21 chain; reachable panic in certificate revocation list parsing — only triggered when parsing CRLs, which iris-mpc does not do (TLS to AWS endpoints only). The 0.101.x line has no fix; drop once aws-config is bumped to drop the legacy rustls 0.21 chain." }, - { id = "RUSTSEC-2026-0002", reason = "lru 0.12.5 pulled transitively through aws-sdk-s3 1.65.0; Stacked Borrows soundness issue in IterMut — iris-mpc does not use lru directly and aws-sdk-s3 uses it for internal endpoint/credential caching that does not iterate mutably. Patched in lru >=0.16.3; drop once aws-sdk-s3 is bumped to a release pulling the patched lru. Note: cargo-deny 0.19+ demotes unsound informational advisories to notes locally, so this is only surfaced by the CI action (currently pinned to cargo-deny 0.18.6)." }, { id = "RUSTSEC-2026-0176", reason = "pyo3 0.24.2 out-of-bounds read in PyList/PyTuple iterators; fixed in >=0.29.0. iris-mpc-py is a test-only tool with no production exposure; drop once iris-mpc-py updates to pyo3 >=0.29.0" }, { id = "RUSTSEC-2026-0177", reason = "pyo3 0.24.2 missing Sync bound on PyCFunction::new_closure; fixed in >=0.29.0. iris-mpc-py is a test-only tool with no production exposure; drop once iris-mpc-py updates to pyo3 >=0.29.0" }, { id = "RUSTSEC-2026-0194", reason = "quick-xml 0.26.0 quadratic-attribute-check DoS; reached only via pprof 0.15 -> inferno ^0.11 (flamegraph feature, transitively through ampc-server-utils), which is unbumpable (every pprof release pins inferno ^0.11 -> quick-xml ^0.26; patched only in quick-xml >=0.41.0). Not exploitable: inferno uses quick-xml solely to GENERATE flamegraph SVGs from our own profiling data, never to parse untrusted XML via the vulnerable Attributes/NsReader read paths. Matches the ampc-common ignore; drop once pprof widens its inferno bound." }, diff --git a/iris-mpc-bins/Cargo.toml b/iris-mpc-bins/Cargo.toml index 619344759d..10a536b5de 100644 --- a/iris-mpc-bins/Cargo.toml +++ b/iris-mpc-bins/Cargo.toml @@ -27,8 +27,8 @@ aws-config = { workspace = true } aws-sdk-kms = { workspace = true, optional = true } aws-sdk-sns = { workspace = true, optional = true } aws-sdk-sqs = { workspace = true, optional = true } -aws-sdk-s3 = { workspace = true } aws-sdk-secretsmanager = { workspace = true, optional = true } +object_store.workspace = true base64.workspace = true cudarc.workspace = true chrono.workspace = true @@ -83,7 +83,6 @@ iris-mpc-store = { path = "../iris-mpc-store" } iris-mpc-upgrade-hawk = { path = "../iris-mpc-upgrade-hawk" } iris-mpc-utils = { path = "../iris-mpc-utils" } blake3.workspace = true -aws-smithy-types = "1.2.9" clap_builder = "4.5.51" # --------------------- diff --git a/iris-mpc-bins/bin/iris-mpc-anon-stats-server/main.rs b/iris-mpc-bins/bin/iris-mpc-anon-stats-server/main.rs index 1b2c8d4e00..c1aec5c481 100644 --- a/iris-mpc-bins/bin/iris-mpc-anon-stats-server/main.rs +++ b/iris-mpc-bins/bin/iris-mpc-anon-stats-server/main.rs @@ -14,14 +14,13 @@ use ampc_server_utils::{ init_heartbeat_task, shutdown_handler::ShutdownHandler, wait_for_others_ready, wait_for_others_unready, TaskMonitor, }; -use aws_sdk_s3::{config::Builder as S3ConfigBuilder, Client as S3Client}; use aws_sdk_sns::{config::Region, types::MessageAttributeValue, Client as SNSClient}; -use aws_smithy_types::retry::RetryConfig; use chrono::{DateTime, Utc}; use clap_builder::Parser; use eyre::{bail, eyre, Context, Result}; use iris_mpc_common::config::{ENV_PROD, ENV_STAGE}; use iris_mpc_common::helpers::sqs_s3_helper::upload_file_to_s3; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_common::{ helpers::{ smpc_request::{ANONYMIZED_STATISTICS_2D_MESSAGE_TYPE, ANONYMIZED_STATISTICS_MESSAGE_TYPE}, @@ -187,7 +186,7 @@ struct AnonStatsProcessor { config: Arc, store: AnonStatsStore, sns_client: SNSClient, - s3_client: S3Client, + s3_client: ObjectStoreClient, publish: PublishTargets, sync_failures: HashMap<(AnonStatsOrigin, AnonStatsOperation), usize>, last_report_times: HashMap<(AnonStatsOrigin, AnonStatsOperation), DateTime>, @@ -198,7 +197,7 @@ impl AnonStatsProcessor { config: Arc, store: AnonStatsStore, sns_client: SNSClient, - s3_client: S3Client, + s3_client: ObjectStoreClient, ) -> Self { let publish = PublishTargets { topic_arn: config.results_topic_arn.clone(), @@ -1089,26 +1088,22 @@ async fn build_sns_client(config: &AnonStatsServerConfig) -> Result { Ok(SNSClient::from_conf(sns_config_builder.build())) } -async fn build_s3_client(config: &AnonStatsServerConfig) -> Result { +async fn build_s3_client(config: &AnonStatsServerConfig) -> Result { let force_path_style = config.environment != ENV_PROD && config.environment != ENV_STAGE; - let retry_config = RetryConfig::standard().with_max_attempts(5); - + let region = config.aws.as_ref().and_then(|aws| aws.region.clone()); let mut loader = aws_config::from_env(); - if let Some(aws) = &config.aws { - if let Some(region) = &aws.region { - loader = loader.region(Region::new(region.clone())); - } + if let Some(region) = ®ion { + loader = loader.region(Region::new(region.clone())); } - let shared_config = loader.load().await; - let mut s3_config = S3ConfigBuilder::from(&shared_config).retry_config(retry_config.clone()); + let mut client = + ObjectStoreClient::new(region, force_path_style).with_aws_sdk_config(&shared_config); if let Some(aws) = &config.aws { if let Some(endpoint) = &aws.endpoint { - s3_config = s3_config.endpoint_url(endpoint); - } - if force_path_style { - s3_config = s3_config.force_path_style(force_path_style); + client = client + .with_option("aws_endpoint", endpoint) + .with_option("aws_allow_http", endpoint.starts_with("http://")); } } - Ok(S3Client::from_conf(s3_config.build())) + Ok(client) } diff --git a/iris-mpc-bins/bin/iris-mpc-common/key_manager.rs b/iris-mpc-bins/bin/iris-mpc-common/key_manager.rs index 778f479b86..d3225c0f44 100644 --- a/iris-mpc-bins/bin/iris-mpc-common/key_manager.rs +++ b/iris-mpc-bins/bin/iris-mpc-common/key_manager.rs @@ -1,15 +1,14 @@ use aws_config::SdkConfig; -use aws_sdk_s3::{ - config::Region as S3Region, operation::put_object::PutObjectOutput, Client as S3Client, - Error as S3Error, -}; use aws_sdk_secretsmanager::{ + config::Region as S3Region, operation::{get_secret_value::GetSecretValueOutput, put_secret_value::PutSecretValueOutput}, Client as SecretsManagerClient, Error as SecretsManagerError, }; use base64::{engine::general_purpose::STANDARD, Engine}; use clap::{Parser, Subcommand}; use eyre::Result; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; +use object_store::PutResult; use rand::{thread_rng, Rng}; use reqwest::Client; use sodiumoxide::crypto::box_::{curve25519xsalsa20poly1305, PublicKey, SecretKey, Seed}; @@ -188,16 +187,21 @@ async fn rotate_keys( rng.fill(&mut seedbuf); let pk_seed = Seed(seedbuf); - let mut s3_config_builder = aws_sdk_s3::config::Builder::from(sdk_config); let mut sm_config_builder = aws_sdk_secretsmanager::config::Builder::from(sdk_config); + let region = sdk_config + .region() + .map(ToString::to_string) + .unwrap_or_else(|| "us-east-1".to_owned()); + let mut s3_client = ObjectStoreClient::new(Some(region), endpoint_url.is_some()) + .with_aws_sdk_config(sdk_config); if let Some(endpoint_url) = endpoint_url.as_ref() { - s3_config_builder = s3_config_builder.endpoint_url(endpoint_url); - s3_config_builder = s3_config_builder.force_path_style(true); + s3_client = s3_client + .with_option("aws_endpoint", endpoint_url) + .with_option("aws_allow_http", endpoint_url.starts_with("http://")); sm_config_builder = sm_config_builder.endpoint_url(endpoint_url); } - let s3_client = S3Client::from_conf(s3_config_builder.build()); let sm_client = SecretsManagerClient::from_conf(sm_config_builder.build()); let (public_key, private_key) = generate_key_pairs(pk_seed); @@ -298,18 +302,15 @@ async fn upload_private_key_to_asm( } async fn upload_public_key_to_s3( - client: &S3Client, + client: &ObjectStoreClient, bucket: &str, key: &str, content: &str, -) -> Result { - Ok(client - .put_object() - .bucket(bucket) - .key(key) - .body(content.to_string().into_bytes().into()) - .send() - .await?) +) -> object_store::Result { + client + .store(bucket)? + .put(&path(key)?, content.as_bytes().to_vec().into()) + .await } fn generate_key_pairs(seed: Seed) -> (PublicKey, SecretKey) { diff --git a/iris-mpc-bins/bin/iris-mpc-cpu/db_sanity_check.rs b/iris-mpc-bins/bin/iris-mpc-cpu/db_sanity_check.rs index a8f9b84bcf..70c1a3090d 100644 --- a/iris-mpc-bins/bin/iris-mpc-cpu/db_sanity_check.rs +++ b/iris-mpc-bins/bin/iris-mpc-cpu/db_sanity_check.rs @@ -2,6 +2,7 @@ use clap::Parser; use eyre::Result; +use iris_mpc_common::object_store::{path as object_path, ObjectStoreClient, ObjectStoreExt}; use iris_mpc_common::{ config::{ENV_PROD, ENV_STAGE}, helpers::smpc_request::{ @@ -1094,7 +1095,7 @@ async fn main() -> Result<()> { None => None, }; - // --- Load graph from S3 checkpoint, then replay any mutations recorded after it --- + // --- Load graph from an object-store checkpoint, then replay later mutations --- let bucket = config.graph_checkpoint_bucket_name.as_str(); rpt!(rpt, "--- Loading graph from S3 checkpoint ---"); let checkpoint_state = load_checkpoint_state( @@ -1105,8 +1106,8 @@ async fn main() -> Result<()> { ) .await?; - // Build the checkpoint S3 client (mirrors AwsClients::checkpoint_s3_client, - // which may target a different region from the general S3 client). + // Build the checkpoint object-store client, which may target a different + // region from the general object-store client. let checkpoint_s3_client = build_checkpoint_s3_client( &config.graph_checkpoint_bucket_region, config.force_path_style(), @@ -2370,22 +2371,23 @@ fn parse_s3_uri(uri: &str) -> Result<(String, String)> { Ok((bucket.to_string(), prefix.to_string())) } -fn build_s3_client(config: &aws_config::SdkConfig, force_path_style: bool) -> aws_sdk_s3::Client { - let retry_config = aws_config::retry::RetryConfig::standard().with_max_attempts(5); - let s3_config = aws_sdk_s3::config::Builder::from(config) - .force_path_style(force_path_style) - .retry_config(retry_config) - .build(); - aws_sdk_s3::Client::from_conf(s3_config) +fn build_s3_client( + sdk_config: &aws_config::SdkConfig, + force_path_style: bool, +) -> ObjectStoreClient { + ObjectStoreClient::new(None, force_path_style).with_aws_sdk_config(sdk_config) } /// Build an S3 client for the graph-checkpoint bucket, allowing the region to /// differ from the ambient default (genesis writes checkpoints into a bucket /// that may live in a different region from the iris/exclusions buckets). -async fn build_checkpoint_s3_client(region: &str, force_path_style: bool) -> aws_sdk_s3::Client { - let loader = aws_config::from_env().region(aws_sdk_s3::config::Region::new(region.to_owned())); - let config = loader.load().await; - build_s3_client(&config, force_path_style) +async fn build_checkpoint_s3_client(region: &str, force_path_style: bool) -> ObjectStoreClient { + let sdk_config = aws_config::from_env() + .region(aws_config::Region::new(region.to_owned())) + .load() + .await; + ObjectStoreClient::new(Some(region.to_owned()), force_path_style) + .with_aws_sdk_config(&sdk_config) } async fn download_exclusions_from_s3( @@ -2396,12 +2398,15 @@ async fn download_exclusions_from_s3( eyre::ensure!(!key.is_empty(), "S3 URI must include an object key"); println!("Downloading exclusions from s3://{bucket}/{key}"); - let config = aws_config::from_env().load().await; - let client = build_s3_client(&config, force_path_style); - - let response = client.get_object().bucket(&bucket).key(&key).send().await?; - let body = response.body.collect().await?; - let exclusions: ExclusionsFile = serde_json::from_slice(&body.into_bytes())?; + let sdk_config = aws_config::from_env().load().await; + let client = build_s3_client(&sdk_config, force_path_style); + let body = client + .store(&bucket)? + .get(&object_path(&key)?) + .await? + .bytes() + .await?; + let exclusions: ExclusionsFile = serde_json::from_slice(&body)?; println!( " Loaded {} excluded serial IDs", @@ -2414,8 +2419,9 @@ async fn upload_to_s3(s3_uri: &str, files: &[PathBuf], force_path_style: bool) - let (bucket, prefix) = parse_s3_uri(s3_uri)?; println!("--- Uploading to S3: s3://{bucket}/{prefix} ---"); - let config = aws_config::from_env().load().await; - let client = build_s3_client(&config, force_path_style); + let sdk_config = aws_config::from_env().load().await; + let client = build_s3_client(&sdk_config, force_path_style); + let store = client.store(&bucket)?; for path in files { let file_name = path @@ -2429,20 +2435,8 @@ async fn upload_to_s3(s3_uri: &str, files: &[PathBuf], force_path_style: bool) - format!("{trimmed}/{file_name}") }; - let content_type = if file_name.ends_with(".json") { - "application/json" - } else { - "text/plain" - }; - let body = aws_sdk_s3::primitives::ByteStream::from_path(path).await?; - client - .put_object() - .bucket(&bucket) - .key(&key) - .body(body) - .content_type(content_type) - .send() - .await?; + let body = tokio::fs::read(path).await?; + store.put(&object_path(&key)?, body.into()).await?; println!(" Uploaded s3://{bucket}/{key}"); } diff --git a/iris-mpc-bins/bin/iris-mpc-cpu/graph_mem_cli.rs b/iris-mpc-bins/bin/iris-mpc-cpu/graph_mem_cli.rs index 9aa19b7cf5..ee9d245c83 100644 --- a/iris-mpc-bins/bin/iris-mpc-cpu/graph_mem_cli.rs +++ b/iris-mpc-bins/bin/iris-mpc-cpu/graph_mem_cli.rs @@ -1,9 +1,8 @@ #![recursion_limit = "256"] -use aws_sdk_s3::config::Region as S3Region; -use aws_sdk_s3::Client as S3Client; use clap::{Parser, Subcommand}; use eyre::Result; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_cpu::hnsw::graph::test_utils::{DbContext, DiffMethod}; use iris_mpc_utils::misc::write_bin; use std::path::PathBuf; @@ -68,12 +67,12 @@ async fn main() -> Result<()> { command, } = Cli::parse(); - let mut builder = aws_config::from_env(); - if let Some(aws_region) = aws_region { - builder = builder.region(S3Region::new(aws_region)); + let mut sdk_config_loader = aws_config::from_env(); + if let Some(region) = &aws_region { + sdk_config_loader = sdk_config_loader.region(aws_config::Region::new(region.clone())); } - let shared_config = builder.load().await; - let s3_client = S3Client::new(&shared_config); + let sdk_config = sdk_config_loader.load().await; + let s3_client = ObjectStoreClient::new(aws_region, false).with_aws_sdk_config(&sdk_config); let db_context = DbContext::new(&db_url, &schema, s3_client, s3_bucket, party_id).await; diff --git a/iris-mpc-bins/bin/iris-mpc-cpu/init_test_dbs.rs b/iris-mpc-bins/bin/iris-mpc-cpu/init_test_dbs.rs index 195c09c980..fbce705148 100644 --- a/iris-mpc-bins/bin/iris-mpc-cpu/init_test_dbs.rs +++ b/iris-mpc-bins/bin/iris-mpc-cpu/init_test_dbs.rs @@ -1,9 +1,8 @@ #![recursion_limit = "256"] -use aws_sdk_s3::config::Region as S3Region; -use aws_sdk_s3::Client as S3Client; use clap::Parser; use eyre::Result; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_common::{iris_db::iris::IrisCode, SerialId, VectorId}; use iris_mpc_cpu::{ execution::hawk_main::STORE_IDS, @@ -450,12 +449,13 @@ async fn main() -> Result<()> { } async fn init_dbs(args: &Args) -> Vec { - let mut builder = aws_config::from_env(); - if let Some(aws_region) = args.aws_region.clone() { - builder = builder.region(S3Region::new(aws_region)); + let mut sdk_config_loader = aws_config::from_env(); + if let Some(region) = &args.aws_region { + sdk_config_loader = sdk_config_loader.region(aws_config::Region::new(region.clone())); } - let shared_config = builder.load().await; - let s3_client = S3Client::new(&shared_config); + let sdk_config = sdk_config_loader.load().await; + let s3_client = + ObjectStoreClient::new(args.aws_region.clone(), false).with_aws_sdk_config(&sdk_config); let mut dbs = Vec::new(); for (party_id, (url, schema)) in izip!(args.db_urls().iter(), args.db_schemas().iter()) diff --git a/iris-mpc-bins/bin/iris-mpc-cpu/sidecar.rs b/iris-mpc-bins/bin/iris-mpc-cpu/sidecar.rs index ddf8375405..66bf0aeee6 100644 --- a/iris-mpc-bins/bin/iris-mpc-cpu/sidecar.rs +++ b/iris-mpc-bins/bin/iris-mpc-cpu/sidecar.rs @@ -189,8 +189,12 @@ async fn main() -> Result<()> { let graph_store: GraphPg = GraphPg::new(&postgres).await?; - let aws_cfg = aws_config::from_env().load().await; - let s3_client = aws_sdk_s3::Client::new(&aws_cfg); + let aws_config = aws_config::from_env().load().await; + let s3_client = iris_mpc_common::object_store::ObjectStoreClient::new( + None, + std::env::var("AWS_ENDPOINT_URL").is_ok(), + ) + .with_aws_sdk_config(&aws_config); // TLS is supplied via SMPC__TLS__* env vars (same path as Hawk Main), not // CLI flags; load it from the SMPC-prefixed config and feed it through. diff --git a/iris-mpc-bins/bin/iris-mpc-upgrade/rerandomize_db.rs b/iris-mpc-bins/bin/iris-mpc-upgrade/rerandomize_db.rs index 58910d38b3..3cf68c991c 100644 --- a/iris-mpc-bins/bin/iris-mpc-upgrade/rerandomize_db.rs +++ b/iris-mpc-bins/bin/iris-mpc-upgrade/rerandomize_db.rs @@ -1,7 +1,6 @@ use std::ops::Range; use ampc_server_utils::TaskMonitor; -use aws_sdk_s3::{operation::put_object::PutObjectOutput, Client as S3Client, Error as S3Error}; use aws_sdk_secretsmanager::{ operation::put_secret_value::PutSecretValueOutput, Client as SecretsManagerClient, Error as SecretsManagerError, @@ -18,6 +17,7 @@ use iris_mpc_common::galois_engine::degree4::{ GaloisRingIrisCodeShare, GaloisRingTrimmedMaskCodeShare, }; use iris_mpc_common::id::PartyID; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use iris_mpc_common::postgres::{AccessMode, PostgresClient}; use iris_mpc_store::{DbStoredIris, Store, StoredIrisRef}; use iris_mpc_upgrade::config::{ @@ -30,6 +30,7 @@ use iris_mpc_upgrade::{ utils::{install_tracing, spawn_healthcheck_server}, }; use itertools::Itertools; +use object_store::PutResult; use tokio::task::JoinSet; use tracing::Level; @@ -61,18 +62,15 @@ async fn upload_private_key_to_asm( } async fn upload_public_key_to_s3( - client: &S3Client, + client: &ObjectStoreClient, bucket: &str, key: &str, content: &str, -) -> Result { - Ok(client - .put_object() - .bucket(bucket) - .key(key) - .body(content.to_string().into_bytes().into()) - .send() - .await?) +) -> object_store::Result { + client + .store(bucket)? + .put(&path(key)?, content.as_bytes().to_vec().into()) + .await } async fn keygen_main(config: KeyGenConfig) -> Result<()> { @@ -84,9 +82,9 @@ async fn keygen_main(config: KeyGenConfig) -> Result<()> { config.env, config.party_id ); - let s3_config_builder = aws_sdk_s3::config::Builder::from(&sdk_config); let sm_config_builder = aws_sdk_secretsmanager::config::Builder::from(&sdk_config); - let s3_client = S3Client::from_conf(s3_config_builder.build()); + let s3_client = ObjectStoreClient::new(sdk_config.region().map(ToString::to_string), false) + .with_aws_sdk_config(&sdk_config); let sm_client = SecretsManagerClient::from_conf(sm_config_builder.build()); // Generate keys only when the secret does not exist diff --git a/iris-mpc-bins/bin/iris-mpc/server.rs b/iris-mpc-bins/bin/iris-mpc/server.rs index e8ce1b4370..d8b2c71582 100644 --- a/iris-mpc-bins/bin/iris-mpc/server.rs +++ b/iris-mpc-bins/bin/iris-mpc/server.rs @@ -12,7 +12,6 @@ use ampc_server_utils::{ init_heartbeat_task, set_node_ready, shutdown_handler::ShutdownHandler, start_coordination_server, wait_for_others_ready, wait_for_others_unready, TaskMonitor, }; -use aws_sdk_s3::Client as S3Client; use aws_sdk_secretsmanager::Client as SecretsManagerClient; use aws_sdk_sns::{types::MessageAttributeValue, Client as SNSClient}; use aws_sdk_sqs::Client; @@ -543,7 +542,7 @@ async fn server_main(config: Config) -> Result<()> { // Start thread that will be responsible for communicating back the results let (tx, mut rx) = mpsc::channel::(32); // TODO: pick some buffer value let sns_client_bg = aws_clients.sns_client.clone(); - let s3_client_bg = aws_clients.s3_client.clone(); + let s3_client_bg = aws_clients.object_store_client.clone(); let config_bg = config.clone(); let store_bg = store.clone(); let shutdown_handler_bg = Arc::clone(&shutdown_handler); @@ -1138,7 +1137,7 @@ async fn server_main(config: Config) -> Result<()> { party_id, aws_clients.sqs_client.clone(), aws_clients.sns_client.clone(), - aws_clients.s3_client.clone(), + aws_clients.object_store_client.clone(), config.clone(), shares_encryption_key_pair.clone(), shutdown_handler.clone(), diff --git a/iris-mpc-common/Cargo.toml b/iris-mpc-common/Cargo.toml index a2d49bf933..2157867661 100644 --- a/iris-mpc-common/Cargo.toml +++ b/iris-mpc-common/Cargo.toml @@ -10,11 +10,12 @@ repository.workspace = true [features] default = ["helpers"] helpers = [ + "dep:async-trait", "dep:aws-config", + "dep:aws-credential-types", "dep:aws-sdk-kms", "dep:aws-sdk-sns", "dep:aws-sdk-sqs", - "dep:aws-sdk-s3", "dep:aws-sdk-secretsmanager", "dep:reqwest", ] @@ -22,11 +23,13 @@ explicit-sns-batching = ["dep:zstd"] [dependencies] aws-config = { workspace = true, optional = true } +aws-credential-types = { workspace = true, optional = true } aws-sdk-kms = { workspace = true, optional = true } aws-sdk-sns = { workspace = true, optional = true } aws-sdk-sqs = { workspace = true, optional = true } -aws-sdk-s3 = { workspace = true, optional = true } aws-sdk-secretsmanager = { workspace = true, optional = true } +object_store.workspace = true +async-trait = { workspace = true, optional = true } clap.workspace = true rand.workspace = true bytemuck.workspace = true @@ -71,4 +74,3 @@ zstd = { version = "0.13", optional = true } [dev-dependencies] float_eq = "1" -aws-credential-types = "1.2.1" diff --git a/iris-mpc-common/src/config/mod.rs b/iris-mpc-common/src/config/mod.rs index 0bd3d954e1..de5080b5e5 100644 --- a/iris-mpc-common/src/config/mod.rs +++ b/iris-mpc-common/src/config/mod.rs @@ -79,6 +79,8 @@ pub struct Config { pub public_key_base_url: String, #[serde(default = "default_shares_bucket_name")] + /// Object-store location for request shares. Plain values are treated as + /// legacy S3 bucket names; full `object_store` URLs are also accepted. pub shares_bucket_name: String, #[serde(default)] @@ -132,6 +134,8 @@ pub struct Config { pub enable_s3_importer: bool, #[serde(default)] + /// Object-store location for exported database chunks. Plain values remain + /// compatible with legacy S3 bucket names. pub db_chunks_bucket_name: String, #[serde(default = "default_db_chunks_bucket_region")] @@ -200,6 +204,12 @@ pub struct Config { #[serde(default = "default_enable_deletion")] pub enable_deletion: bool, + /// Object-store location for the Iris deletion snapshot. Plain values are + /// treated as legacy S3 bucket names; provider URLs select other backends. + /// If omitted, the historical environment-specific bucket is used. + #[serde(default)] + pub iris_deletions_store_location: Option, + #[serde(default)] pub hnsw_schema_name_suffix: String, @@ -331,6 +341,8 @@ pub struct Config { pub sns_retry_max_attempts: u32, #[serde(default = "default_graph_checkpoint_bucket_name")] + /// Object-store location for graph checkpoints. Plain values are treated + /// as legacy S3 bucket names; provider URLs select other backends. pub graph_checkpoint_bucket_name: String, #[serde(default = "default_graph_checkpoint_bucket_region")] @@ -613,6 +625,15 @@ impl Config { self.anon_stats_schema_name.clone() } + /// Returns the configured Iris deletion snapshot location, falling back to + /// the historical environment-specific S3 bucket. + pub fn iris_deletions_store_location(&self) -> String { + resolve_iris_deletions_store_location( + &self.environment, + self.iris_deletions_store_location.as_deref(), + ) + } + /// Returns the name of a database schema for connecting to a node's gpu dB. pub fn get_gpu_db_schema(&self) -> String { self.format_db_schema(&self.gpu_schema_name_suffix) @@ -642,6 +663,17 @@ impl Config { } } +/// Resolves an optional Iris deletion snapshot location while preserving the +/// bucket convention used before object-store URLs were supported. +pub fn resolve_iris_deletions_store_location( + environment: &str, + configured_location: Option<&str>, +) -> String { + configured_location + .map(ToOwned::to_owned) + .unwrap_or_else(|| format!("wf-smpcv2-{environment}-sync-protocol")) +} + /// Encapsulates database configuration settings. #[derive(Clone, Serialize, Deserialize, Default)] pub struct DbConfig { @@ -736,6 +768,7 @@ pub struct CommonConfig { enable_reset: bool, enable_recovery: bool, enable_deletion: bool, + iris_deletions_store_location: String, hawk_request_parallelism: usize, hawk_connection_parallelism: usize, hnsw_param_ef_constr: usize, @@ -824,6 +857,7 @@ impl From for CommonConfig { enable_reset, enable_recovery, enable_deletion, + iris_deletions_store_location, hawk_request_parallelism, hawk_connection_parallelism, hnsw_param_ef_constr, @@ -874,6 +908,11 @@ impl From for CommonConfig { "hnsw_disable_memory_persistence requires disable_persistence to also be true" ); + let iris_deletions_store_location = resolve_iris_deletions_store_location( + &environment, + iris_deletions_store_location.as_deref(), + ); + Self { environment, results_topic_arn, @@ -902,6 +941,7 @@ impl From for CommonConfig { enable_reset, enable_recovery, enable_deletion, + iris_deletions_store_location, hawk_request_parallelism, hawk_connection_parallelism, hnsw_param_ef_constr, @@ -935,6 +975,28 @@ impl From for CommonConfig { mod tests { use super::*; + #[test] + fn iris_deletions_store_location_preserves_legacy_default_and_allows_override() { + let legacy: Config = serde_json::from_value(serde_json::json!({ + "environment": "dev" + })) + .unwrap(); + assert_eq!( + legacy.iris_deletions_store_location(), + "wf-smpcv2-dev-sync-protocol" + ); + + let configured: Config = serde_json::from_value(serde_json::json!({ + "environment": "dev", + "iris_deletions_store_location": "gs://deletions/snapshots" + })) + .unwrap(); + assert_eq!( + configured.iris_deletions_store_location(), + "gs://deletions/snapshots" + ); + } + #[test] fn db_backed_ingest_is_part_of_common_config_equality() { // The equality check across parties is the gate that prevents mixed diff --git a/iris-mpc-common/src/helpers/key_pair.rs b/iris-mpc-common/src/helpers/key_pair.rs index 5c63bebb20..a443adea40 100644 --- a/iris-mpc-common/src/helpers/key_pair.rs +++ b/iris-mpc-common/src/helpers/key_pair.rs @@ -44,18 +44,12 @@ pub enum SharesDecodingError { url: String, message: String, }, - #[error("Received error message from S3 for key {}: {}", .key, .message)] - S3ResponseContent { key: String, message: String }, + #[error("Object store error for key {}: {}", .key, .message)] + ObjectStoreResponse { key: String, message: String }, #[error(transparent)] SerdeError(#[from] serde_json::error::Error), - #[error(transparent)] - PresigningConfigError(#[from] aws_sdk_s3::presigning::PresigningConfigError), - #[error(transparent)] - PresignedRequestError( - #[from] Box>, - ), - #[error("Upload share file error")] - UploadS3Error, + #[error("Object store upload error")] + ObjectStoreUploadError, } impl From> for SharesDecodingError { @@ -64,16 +58,6 @@ impl From> for SharesDecodingError { } } -impl From> - for SharesDecodingError -{ - fn from( - value: aws_sdk_s3::error::SdkError, - ) -> Self { - Self::PresignedRequestError(Box::new(value)) - } -} - #[derive(Clone, Debug)] pub struct SharesEncryptionKeyPairs { pub current_key_pair: SharesEncryptionKeyPair, diff --git a/iris-mpc-common/src/helpers/smpc_request.rs b/iris-mpc-common/src/helpers/smpc_request.rs index d98408a860..b988f30a9c 100644 --- a/iris-mpc-common/src/helpers/smpc_request.rs +++ b/iris-mpc-common/src/helpers/smpc_request.rs @@ -1,6 +1,6 @@ use super::{key_pair::SharesDecodingError, sha256::sha256_as_hex_string}; use crate::helpers::key_pair::SharesEncryptionKeyPairs; -use aws_sdk_s3::Client as S3Client; +use crate::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use aws_sdk_sns::types::MessageAttributeValue; use aws_sdk_sqs::{ error::SdkError, @@ -255,33 +255,36 @@ impl SharesS3Object { pub async fn get_iris_data_by_party_id( s3_key: &str, party_id: usize, - bucket_name: &String, - s3_client: &S3Client, + bucket_name: &str, + object_store_client: &ObjectStoreClient, ) -> Result<(String, String), SharesDecodingError> { - let response = s3_client - .get_object() - .bucket(bucket_name) - .key(s3_key) - .send() - .await - .map_err(|err| { - tracing::error!("Failed to download file: {}", err); - SharesDecodingError::S3ResponseContent { - key: s3_key.to_string(), - message: err.to_string(), - } - })?; + let store = object_store_client.store(bucket_name).map_err(|err| { + tracing::error!("Failed to create object store: {err}"); + SharesDecodingError::ObjectStoreResponse { + key: s3_key.to_string(), + message: err.to_string(), + } + })?; + let location = path(s3_key).map_err(|err| SharesDecodingError::ObjectStoreResponse { + key: s3_key.to_string(), + message: err.to_string(), + })?; + let response = store.get(&location).await.map_err(|err| { + tracing::error!("Failed to download file: {}", err); + SharesDecodingError::ObjectStoreResponse { + key: s3_key.to_string(), + message: err.to_string(), + } + })?; - let object_body = response.body.collect().await.map_err(|e| { + let bytes = response.bytes().await.map_err(|e| { tracing::error!("Failed to get object body: {}", e); - SharesDecodingError::S3ResponseContent { + SharesDecodingError::ObjectStoreResponse { key: s3_key.to_string(), message: e.to_string(), } })?; - let bytes = object_body.into_bytes(); - let shares_file: SharesS3Object = serde_json::from_slice(&bytes)?; let field_name = format!("iris_share_{}", party_id); diff --git a/iris-mpc-common/src/helpers/sqs_s3_helper.rs b/iris-mpc-common/src/helpers/sqs_s3_helper.rs index e1e9df6f3c..483e10deb3 100644 --- a/iris-mpc-common/src/helpers/sqs_s3_helper.rs +++ b/iris-mpc-common/src/helpers/sqs_s3_helper.rs @@ -1,32 +1,28 @@ use crate::helpers::key_pair::SharesDecodingError; -use aws_sdk_s3::{ - primitives::{ByteStream, SdkBody}, - Client, -}; +use crate::object_store::{path, ObjectStoreClient, ObjectStoreExt}; pub async fn upload_file_to_s3( bucket: &str, key: &str, - s3_client: Client, + object_store_client: ObjectStoreClient, contents: &[u8], ) -> Result { - let content_bytestream = ByteStream::new(SdkBody::from(contents)); + let store = object_store_client.store(bucket).map_err(|e| { + tracing::error!("Failed to create object store for {bucket}: {e}"); + SharesDecodingError::ObjectStoreUploadError + })?; + let location = path(key).map_err(|e| { + tracing::error!("Invalid object key {key}: {e}"); + SharesDecodingError::ObjectStoreUploadError + })?; - // Create a PutObject request - match s3_client - .put_object() - .bucket(bucket) - .key(key) - .body(content_bytestream) - .send() - .await - { + match store.put(&location, contents.to_vec().into()).await { Ok(_) => { - tracing::info!("File {} uploaded to s3 successfully", key); + tracing::info!("File {key} uploaded to object store successfully"); } Err(e) => { - tracing::error!("Failed to upload file {} to s3: {:?}", key, e); - return Err(SharesDecodingError::UploadS3Error); + tracing::error!("Failed to upload file {key}: {e:?}"); + return Err(SharesDecodingError::ObjectStoreUploadError); } } Ok(key.to_string()) diff --git a/iris-mpc-common/src/lib.rs b/iris-mpc-common/src/lib.rs index f336357579..bfc8b3d851 100644 --- a/iris-mpc-common/src/lib.rs +++ b/iris-mpc-common/src/lib.rs @@ -5,6 +5,7 @@ pub mod galois_engine; pub mod helpers; pub mod iris_db; pub mod job; +pub mod object_store; pub mod postgres; pub mod shamir; #[cfg(feature = "helpers")] diff --git a/iris-mpc-common/src/object_store.rs b/iris-mpc-common/src/object_store.rs new file mode 100644 index 0000000000..d7b19eb0a9 --- /dev/null +++ b/iris-mpc-common/src/object_store.rs @@ -0,0 +1,668 @@ +//! Backend-independent object storage construction. +//! +//! Existing configuration values historically contain an S3 bucket name. To +//! preserve compatibility, values without a URL scheme are interpreted as S3 +//! buckets. Full URLs can be used to select any backend supported by the +//! `object_store` crate, for example `gs://bucket`, `az://container`, or +//! `file:///var/lib/iris-mpc/objects`. + +use object_store::{ + aws::{AmazonS3Builder, AmazonS3ConfigKey, AwsCredentialProvider}, + path::Path, + prefix::PrefixStore, + ObjectStore, ObjectStoreScheme, +}; +use std::{ + collections::HashMap, + sync::{Arc, RwLock}, +}; +use url::Url; + +#[cfg(feature = "helpers")] +use aws_credential_types::provider::{ProvideCredentials, SharedCredentialsProvider}; +#[cfg(feature = "helpers")] +use object_store::{aws::AwsCredential, CredentialProvider}; +#[cfg(feature = "helpers")] +use std::time::{Duration, SystemTime}; + +pub use object_store::ObjectStoreExt; + +/// A shared object store rooted at the configured bucket/container and prefix. +pub type ObjectStoreRef = Arc; + +/// Lazily constructs and caches object stores for configured storage locations. +/// +/// This retains the old client-plus-bucket calling convention while removing +/// the storage implementation's dependency on `aws-sdk-s3`. +#[derive(Clone, Debug)] +pub struct ObjectStoreClient { + region: Option, + force_path_style: bool, + aws_credentials: Option, + options: Arc>, + stores: Arc>>, +} + +impl ObjectStoreClient { + pub fn new(region: Option, force_path_style: bool) -> Self { + Self { + region, + force_path_style, + aws_credentials: None, + options: Arc::default(), + stores: Arc::default(), + } + } + + /// Use an already-resolved AWS SDK credential chain for S3 stores. + /// + /// This preserves shared-profile, SSO, `credential_process`, and assume-role + /// support while still constructing the storage backend through + /// `object_store`. The provider remains dynamic, so expiring credentials can + /// be refreshed by the AWS SDK chain. + #[cfg(feature = "helpers")] + pub fn with_aws_sdk_config(mut self, sdk_config: &aws_config::SdkConfig) -> Self { + if self.region.is_none() { + self.region = sdk_config.region().map(ToString::to_string); + } + if let Some(endpoint) = sdk_config.endpoint_url() { + Arc::make_mut(&mut self.options).push(("aws_endpoint".to_owned(), endpoint.to_owned())); + } + if let Some(provider) = sdk_config.credentials_provider() { + self.aws_credentials = Some(Arc::new(AwsSdkCredentialAdapter::new(provider))); + } + self + } + + pub fn with_option(mut self, key: impl Into, value: impl ToString) -> Self { + Arc::make_mut(&mut self.options).push((key.into(), value.to_string())); + self + } + + /// Returns the store for `location`. + /// + /// A value without a URL scheme is treated as an S3 bucket for backwards + /// compatibility. A URL may also contain a path prefix; returned stores are + /// rooted at that prefix. + pub fn store(&self, location: &str) -> object_store::Result { + if let Some(store) = self.stores.read().unwrap().get(location).cloned() { + return Ok(store); + } + + let url = storage_url(location)?; + let (scheme, prefix) = ObjectStoreScheme::parse(&url)?; + let mut options: Vec<(String, String)> = std::env::vars().collect(); + if let Some(region) = &self.region { + options.push(("aws_region".to_owned(), region.clone())); + } + let mut addressing_options = options.clone(); + addressing_options.extend(self.options.iter().cloned()); + options.push(( + "aws_virtual_hosted_style_request".to_owned(), + automatic_virtual_hosted_style(&url, self.force_path_style, &addressing_options)? + .to_string(), + )); + options.extend(self.options.iter().cloned()); + + if scheme == ObjectStoreScheme::AmazonS3 { + normalize_s3_options(&url, &mut options)?; + } + + let store: Box = if scheme == ObjectStoreScheme::AmazonS3 { + let mut builder = AmazonS3Builder::new().with_url(url.as_str()); + for (key, value) in options { + if let Ok(key) = key.to_ascii_lowercase().parse::() { + builder = builder.with_config(key, value); + } + } + if let Some(credentials) = &self.aws_credentials { + builder = builder.with_credentials(credentials.clone()); + } + Box::new(builder.build()?) + } else { + object_store::parse_url_opts(&url, options)?.0 + }; + let store: ObjectStoreRef = if prefix.as_ref().is_empty() { + Arc::from(store) + } else { + Arc::new(PrefixStore::new(store, prefix)) + }; + + let mut stores = self.stores.write().unwrap(); + Ok(stores + .entry(location.to_owned()) + .or_insert_with(|| store.clone()) + .clone()) + } + + /// Registers a pre-built store, primarily for backend-independent tests. + pub fn insert(&self, location: impl Into, store: ObjectStoreRef) { + self.stores.write().unwrap().insert(location.into(), store); + } +} + +#[cfg(feature = "helpers")] +#[derive(Debug)] +struct AwsSdkCredentialAdapter { + provider: SharedCredentialsProvider, + cached: tokio::sync::Mutex>, +} + +#[cfg(feature = "helpers")] +#[derive(Debug)] +struct CachedAwsCredential { + credential: Arc, + expires_at: SystemTime, +} + +#[cfg(feature = "helpers")] +impl AwsSdkCredentialAdapter { + // Match the AWS SDK lazy identity cache defaults so switching storage + // implementations does not change credential loading behavior. + const LOAD_TIMEOUT: Duration = Duration::from_secs(5); + const DEFAULT_EXPIRATION: Duration = Duration::from_secs(15 * 60); + const REFRESH_BUFFER: Duration = Duration::from_secs(10); + + fn new(provider: SharedCredentialsProvider) -> Self { + Self { + provider, + cached: tokio::sync::Mutex::new(None), + } + } + + fn is_fresh(cached: &CachedAwsCredential) -> bool { + cached + .expires_at + .duration_since(SystemTime::now()) + .is_ok_and(|remaining| remaining > Self::REFRESH_BUFFER) + } + + async fn get_credential_with_timeout( + &self, + timeout: Duration, + ) -> object_store::Result> { + tokio::time::timeout(timeout, self.get_credential_inner()) + .await + .map_err(|source| object_store::Error::Generic { + store: "AWS SDK credential provider", + source: Box::new(source), + })? + } + + async fn get_credential_inner(&self) -> object_store::Result> { + let mut cached = self.cached.lock().await; + if let Some(cached) = cached.as_ref().filter(|cached| Self::is_fresh(cached)) { + return Ok(cached.credential.clone()); + } + + let credentials = self + .provider + .provide_credentials() + .await + .map_err(|source| object_store::Error::Generic { + store: "AWS SDK credential provider", + source: Box::new(source), + })?; + let expires_at = credentials + .expiry() + .unwrap_or_else(|| SystemTime::now() + Self::DEFAULT_EXPIRATION); + let credential = Arc::new(AwsCredential { + key_id: credentials.access_key_id().to_owned(), + secret_key: credentials.secret_access_key().to_owned(), + token: credentials.session_token().map(ToOwned::to_owned), + }); + *cached = Some(CachedAwsCredential { + credential: credential.clone(), + expires_at, + }); + Ok(credential) + } +} + +#[cfg(feature = "helpers")] +#[async_trait::async_trait] +impl CredentialProvider for AwsSdkCredentialAdapter { + type Credential = AwsCredential; + + async fn get_credential(&self) -> object_store::Result> { + self.get_credential_with_timeout(Self::LOAD_TIMEOUT).await + } +} + +/// Parse a configured object key using the portable `object_store` path rules. +pub fn path(key: &str) -> object_store::Result { + Path::parse(key).map_err(Into::into) +} + +fn storage_url(location: &str) -> object_store::Result { + let value = if location.contains("://") { + location.to_owned() + } else { + format!("s3://{location}") + }; + Url::parse(&value).map_err(|source| object_store::Error::Generic { + store: "configuration", + source: Box::new(source), + }) +} + +fn normalize_s3_options( + storage_url: &Url, + options: &mut Vec<(String, String)>, +) -> object_store::Result<()> { + let endpoint = effective_s3_endpoint(options).map(str::to_owned); + + if endpoint + .as_deref() + .is_some_and(|value| value.starts_with("http://")) + && !has_option(options, &["aws_allow_http", "allow_http"]) + { + options.push(("aws_allow_http".to_owned(), "true".to_owned())); + } + + let virtual_hosted = option_value( + options, + &[ + "aws_virtual_hosted_style_request", + "virtual_hosted_style_request", + ], + ) + .and_then(|value| value.parse::().ok()) + .unwrap_or(false); + if !virtual_hosted { + return Ok(()); + } + + let Some(bucket) = s3_bucket(storage_url) else { + return Ok(()); + }; + if !is_virtual_hostable_s3_bucket(&bucket, endpoint.as_deref())? { + options.push(( + "aws_virtual_hosted_style_request".to_owned(), + "false".to_owned(), + )); + return Ok(()); + } + let Some(endpoint) = endpoint else { + return Ok(()); + }; + match bucket_qualified_endpoint(&endpoint, &bucket)? { + Some(endpoint) => options.push(("aws_endpoint_url_s3".to_owned(), endpoint)), + None => options.push(( + "aws_virtual_hosted_style_request".to_owned(), + "false".to_owned(), + )), + } + Ok(()) +} + +fn automatic_virtual_hosted_style( + storage_url: &Url, + force_path_style: bool, + options: &[(String, String)], +) -> object_store::Result { + if force_path_style { + return Ok(false); + } + let Some(bucket) = s3_bucket(storage_url) else { + return Ok(false); + }; + is_virtual_hostable_s3_bucket(&bucket, effective_s3_endpoint(options)) +} + +/// Mirrors the AWS SDK's automatic choice between virtual-hosted and path-style +/// S3 addressing. HTTPS requires a single DNS-compatible label so wildcard TLS +/// certificates remain valid; custom HTTP endpoints may use dotted buckets. +fn is_virtual_hostable_s3_bucket( + bucket: &str, + endpoint: Option<&str>, +) -> object_store::Result { + let allow_subdomains = if let Some(endpoint) = endpoint { + let url = Url::parse(endpoint).map_err(|source| object_store::Error::Generic { + store: "S3 endpoint configuration", + source: Box::new(source), + })?; + match url.host() { + Some(url::Host::Domain(host)) if !host.eq_ignore_ascii_case("localhost") => { + url.scheme() == "http" + } + _ => return Ok(false), + } + } else { + false + }; + + if allow_subdomains { + Ok(bucket.split('.').all(is_virtual_hostable_s3_segment)) + } else { + Ok(!bucket.contains('.') && is_virtual_hostable_s3_segment(bucket)) + } +} + +fn is_virtual_hostable_s3_segment(segment: &str) -> bool { + let bytes = segment.as_bytes(); + (3..=63).contains(&bytes.len()) + && bytes + .first() + .is_some_and(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit()) + && bytes + .last() + .is_some_and(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit()) + && bytes + .iter() + .all(|byte| byte.is_ascii_lowercase() || byte.is_ascii_digit() || *byte == b'-') +} + +fn effective_s3_endpoint(options: &[(String, String)]) -> Option<&str> { + option_value(options, &["aws_endpoint_url_s3"]).or_else(|| { + option_value( + options, + &[ + "aws_endpoint_url", + "aws_endpoint", + "endpoint_url", + "endpoint", + ], + ) + }) +} + +fn option_value<'a>(options: &'a [(String, String)], keys: &[&str]) -> Option<&'a str> { + options.iter().rev().find_map(|(key, value)| { + keys.iter() + .any(|candidate| key.eq_ignore_ascii_case(candidate)) + .then_some(value.as_str()) + }) +} + +fn has_option(options: &[(String, String)], keys: &[&str]) -> bool { + option_value(options, keys).is_some() +} + +fn s3_bucket(url: &Url) -> Option { + match url.scheme() { + "s3" | "s3a" => url.host_str().map(ToOwned::to_owned), + "https" => { + let host = url.host_str()?; + if host.starts_with("s3") || host.ends_with("r2.cloudflarestorage.com") { + url.path_segments()?.next().map(ToOwned::to_owned) + } else { + host.split_once(".s3").map(|(bucket, _)| bucket.to_owned()) + } + } + _ => None, + } +} + +fn bucket_qualified_endpoint(endpoint: &str, bucket: &str) -> object_store::Result> { + let mut url = Url::parse(endpoint).map_err(|source| object_store::Error::Generic { + store: "S3 endpoint configuration", + source: Box::new(source), + })?; + let Some(host) = url.host() else { + return Ok(None); + }; + let url::Host::Domain(host) = host else { + return Ok(None); + }; + if host.eq_ignore_ascii_case("localhost") { + return Ok(None); + } + if host == bucket || host.starts_with(&format!("{bucket}.")) { + return Ok(Some(url.to_string().trim_end_matches('/').to_owned())); + } + url.set_host(Some(&format!("{bucket}.{host}"))) + .map_err(|_| object_store::Error::Generic { + store: "S3 endpoint configuration", + source: "unable to add bucket to S3 endpoint host".into(), + })?; + Ok(Some(url.to_string().trim_end_matches('/').to_owned())) +} + +#[cfg(test)] +mod tests { + use super::*; + use object_store::memory::InMemory; + use wiremock::{ + matchers::{method, path as request_path}, + Mock, MockServer, ResponseTemplate, + }; + + #[test] + fn caches_registered_stores() { + let client = ObjectStoreClient::new(None, false); + let expected: ObjectStoreRef = Arc::new(InMemory::new()); + client.insert("test", expected.clone()); + + let actual = client.store("test").unwrap(); + assert!(Arc::ptr_eq(&expected, &actual)); + } + + #[test] + fn plain_locations_remain_s3_buckets() { + assert_eq!(storage_url("bucket").unwrap().as_str(), "s3://bucket"); + } + + #[test] + fn service_specific_http_endpoint_is_allowed() { + let url = storage_url("bucket").unwrap(); + let mut options = vec![ + ( + "AWS_ENDPOINT_URL".to_owned(), + "https://global.example".to_owned(), + ), + ( + "AWS_ENDPOINT_URL_S3".to_owned(), + "http://s3.example".to_owned(), + ), + ( + "aws_virtual_hosted_style_request".to_owned(), + "false".to_owned(), + ), + ]; + + normalize_s3_options(&url, &mut options).unwrap(); + + assert_eq!(effective_s3_endpoint(&options), Some("http://s3.example")); + assert_eq!(option_value(&options, &["aws_allow_http"]), Some("true")); + } + + #[test] + fn virtual_hosted_custom_endpoint_is_bucket_qualified() { + let url = storage_url("bucket").unwrap(); + let mut options = vec![ + ( + "aws_endpoint_url_s3".to_owned(), + "https://storage.example:9443/base".to_owned(), + ), + ( + "aws_virtual_hosted_style_request".to_owned(), + "true".to_owned(), + ), + ]; + + normalize_s3_options(&url, &mut options).unwrap(); + + assert_eq!( + effective_s3_endpoint(&options), + Some("https://bucket.storage.example:9443/base") + ); + } + + #[test] + fn automatic_addressing_falls_back_for_non_virtual_hostable_buckets() { + let options = Vec::new(); + assert!(automatic_virtual_hosted_style( + &storage_url("valid-bucket").unwrap(), + false, + &options, + ) + .unwrap()); + + for bucket in ["bucket.name", "BucketName", "aa"] { + assert!( + !automatic_virtual_hosted_style(&storage_url(bucket).unwrap(), false, &options,) + .unwrap(), + "{bucket} must use path-style addressing" + ); + } + assert!(!automatic_virtual_hosted_style( + &storage_url("valid-bucket").unwrap(), + true, + &options, + ) + .unwrap()); + } + + #[test] + fn dotted_buckets_are_virtual_hosted_only_for_custom_http_endpoints() { + let url = storage_url("bucket.name").unwrap(); + let https_options = vec![( + "aws_endpoint_url_s3".to_owned(), + "https://storage.example".to_owned(), + )]; + let http_options = vec![( + "aws_endpoint_url_s3".to_owned(), + "http://storage.example".to_owned(), + )]; + + assert!(!automatic_virtual_hosted_style(&url, false, &https_options).unwrap()); + assert!(automatic_virtual_hosted_style(&url, false, &http_options).unwrap()); + } + + #[test] + fn already_qualified_custom_endpoint_is_unchanged() { + assert_eq!( + bucket_qualified_endpoint("https://bucket.storage.example", "bucket").unwrap(), + Some("https://bucket.storage.example".to_owned()) + ); + assert_eq!( + bucket_qualified_endpoint("http://localhost:4566", "bucket").unwrap(), + None + ); + } + + #[cfg(feature = "helpers")] + #[tokio::test] + async fn aws_sdk_credentials_are_bridged() { + use aws_credential_types::Credentials; + + let credentials = Credentials::new( + "profile-access-key", + "profile-secret-key", + Some("profile-session-token".to_owned()), + None, + "test-profile", + ); + let sdk_config = aws_config::SdkConfig::builder() + .credentials_provider(SharedCredentialsProvider::new(credentials)) + .endpoint_url("https://profile-endpoint.example") + .build(); + let client = ObjectStoreClient::new(None, false).with_aws_sdk_config(&sdk_config); + + assert_eq!( + option_value(&client.options, &["aws_endpoint"]), + Some("https://profile-endpoint.example") + ); + let provider = client.aws_credentials.as_ref().unwrap(); + let credentials = provider.get_credential().await.unwrap(); + assert_eq!(credentials.key_id, "profile-access-key"); + assert_eq!(credentials.secret_key, "profile-secret-key"); + assert_eq!(credentials.token.as_deref(), Some("profile-session-token")); + let cached_credentials = provider.get_credential().await.unwrap(); + assert!(Arc::ptr_eq(&credentials, &cached_credentials)); + } + + #[cfg(feature = "helpers")] + #[tokio::test] + async fn credentials_without_expiry_are_refreshed() { + use aws_credential_types::{credential_fn::provide_credentials_fn, Credentials}; + use std::sync::atomic::{AtomicUsize, Ordering}; + + let calls = Arc::new(AtomicUsize::new(0)); + let provider_calls = calls.clone(); + let provider = SharedCredentialsProvider::new(provide_credentials_fn(move || { + let call = provider_calls.fetch_add(1, Ordering::SeqCst) + 1; + async move { + Ok(Credentials::new( + format!("access-key-{call}"), + "secret-key", + None, + None, + "rotating-test-provider", + )) + } + })); + let adapter = AwsSdkCredentialAdapter::new(provider); + + let first = adapter.get_credential().await.unwrap(); + let cached = adapter.get_credential().await.unwrap(); + assert_eq!(first.key_id, "access-key-1"); + assert!(Arc::ptr_eq(&first, &cached)); + assert_eq!(calls.load(Ordering::SeqCst), 1); + + // Providers such as credential_process may omit an expiry. They still + // receive a finite cache lifetime, matching the AWS SDK identity cache. + adapter.cached.lock().await.as_mut().unwrap().expires_at = SystemTime::UNIX_EPOCH; + + let refreshed = adapter.get_credential().await.unwrap(); + assert_eq!(refreshed.key_id, "access-key-2"); + assert_eq!(calls.load(Ordering::SeqCst), 2); + } + + #[cfg(feature = "helpers")] + #[tokio::test] + async fn credential_loading_has_a_timeout() { + use aws_credential_types::credential_fn::provide_credentials_fn; + + let provider = SharedCredentialsProvider::new(provide_credentials_fn(|| async { + std::future::pending::().await + })); + let adapter = AwsSdkCredentialAdapter::new(provider); + + let error = adapter + .get_credential_with_timeout(Duration::from_millis(1)) + .await + .unwrap_err(); + + assert!(error.to_string().contains("AWS SDK credential provider")); + } + + #[tokio::test] + async fn http_ip_endpoint_falls_back_to_path_style() { + let server = MockServer::start().await; + Mock::given(method("PUT")) + .and(request_path("/bucket/checkpoint.bin")) + .respond_with(ResponseTemplate::new(200).insert_header("ETag", "\"test-etag\"")) + .mount(&server) + .await; + + let client = ObjectStoreClient::new(Some("us-east-1".to_owned()), false) + .with_option("aws_endpoint_url_s3", server.uri()) + .with_option("aws_access_key_id", "test") + .with_option("aws_secret_access_key", "test"); + let store = client.store("bucket").unwrap(); + store + .put( + &path("checkpoint.bin").unwrap(), + b"checkpoint".to_vec().into(), + ) + .await + .unwrap(); + } + + #[tokio::test] + async fn full_urls_select_non_s3_backends_and_prefixes() { + let client = ObjectStoreClient::new(None, false); + let store = client.store("memory:///iris-mpc").unwrap(); + let location = path("checkpoint.bin").unwrap(); + + store + .put(&location, b"checkpoint".to_vec().into()) + .await + .unwrap(); + let bytes = store.get(&location).await.unwrap().bytes().await.unwrap(); + + assert_eq!(bytes.as_ref(), b"checkpoint"); + } +} diff --git a/iris-mpc-common/tests/smpc_request.rs b/iris-mpc-common/tests/smpc_request.rs index 1c58043875..83d880693d 100644 --- a/iris-mpc-common/tests/smpc_request.rs +++ b/iris-mpc-common/tests/smpc_request.rs @@ -1,6 +1,4 @@ mod tests { - use aws_credential_types::{provider::SharedCredentialsProvider, Credentials}; - use aws_sdk_s3::Client as S3Client; use base64::{engine::general_purpose::STANDARD, Engine}; use iris_mpc_common::helpers::{ key_pair::{SharesDecodingError, SharesEncryptionKeyPairs}, @@ -10,10 +8,11 @@ mod tests { ReAuthRequest, UniquenessRequest, }, }; + use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; + use object_store::memory::InMemory; use serde_json::json; use sodiumoxide::crypto::{box_::PublicKey, sealedbox}; use std::sync::Arc; - use wiremock::{matchers::method, Mock, MockServer, ResponseTemplate}; const PREVIOUS_PUBLIC_KEY: &str = "1UY8lKlS7aVj5ZnorSfLIHlG3jg+L4ToVi4K+mLKqFQ="; const PREVIOUS_PRIVATE_KEY: &str = "X26wWfzP5fKMP7QMz0X3eZsEeF4NhJU92jT69wZg6x8="; @@ -59,7 +58,6 @@ mod tests { #[tokio::test] async fn test_retrieve_iris_shares_from_s3_success() { - let mock_server = MockServer::start().await; let bucket_name = "bobTheBucket"; let key = "kateTheKey"; let response_body = json!({ @@ -71,33 +69,16 @@ mod tests { "iris_hashes_2": "hash_2" }); - let data = response_body.to_string(); - - Mock::given(method("GET")) - .respond_with( - ResponseTemplate::new(200) - .insert_header("Content-Type", "application/octet-stream") - .set_body_raw(data, "application/octet-stream"), + let store = Arc::new(InMemory::new()); + store + .put( + &path(key).unwrap(), + response_body.to_string().into_bytes().into(), ) - .mount(&mock_server) - .await; - - let credentials = - Credentials::new("test-access-key", "test-secret-key", None, None, "test"); - let credentials_provider = SharedCredentialsProvider::new(credentials); - // Configure the S3Client to point to the mock server - let config = aws_config::from_env() - .region("us-west-2") - .endpoint_url(mock_server.uri()) - .credentials_provider(credentials_provider) - .load() - .await; - let s3_config = aws_sdk_s3::config::Builder::from(&config) - .endpoint_url(mock_server.uri()) - .force_path_style(true) - .build(); - - let s3_client = Arc::new(S3Client::from_conf(s3_config)); + .await + .unwrap(); + let s3_client = ObjectStoreClient::new(None, false); + s3_client.insert(bucket_name, store); let smpc_request = UniquenessRequest { signup_id: "test_signup_id".to_string(), @@ -108,13 +89,9 @@ mod tests { disable_anonymized_stats: None, }; - let result = get_iris_data_by_party_id( - smpc_request.s3_key.as_str(), - 0, - &bucket_name.to_string(), - &s3_client, - ) - .await; + let result = + get_iris_data_by_party_id(smpc_request.s3_key.as_str(), 0, bucket_name, &s3_client) + .await; assert!(result.is_ok()); let (share, hash) = result.unwrap(); diff --git a/iris-mpc-cpu/Cargo.toml b/iris-mpc-cpu/Cargo.toml index 949099df8f..11d2f86f0f 100644 --- a/iris-mpc-cpu/Cargo.toml +++ b/iris-mpc-cpu/Cargo.toml @@ -11,7 +11,7 @@ repository.workspace = true aes-prng = { git = "https://github.com/tf-encrypted/aes-prng.git", branch = "dragos/display" } async-channel = "2.3.1" async-trait = "~0.1" -aws-sdk-s3.workspace = true +object_store.workspace = true base64.workspace = true blake3.workspace = true bincode.workspace = true diff --git a/iris-mpc-cpu/src/checkpoint_protocol/materializer.rs b/iris-mpc-cpu/src/checkpoint_protocol/materializer.rs index 931378992d..35c4c5eab8 100644 --- a/iris-mpc-cpu/src/checkpoint_protocol/materializer.rs +++ b/iris-mpc-cpu/src/checkpoint_protocol/materializer.rs @@ -1,5 +1,5 @@ use async_trait::async_trait; -use aws_sdk_s3::Client as S3Client; +use iris_mpc_common::object_store::ObjectStoreClient; use crate::checkpoint_protocol::{ CheckpointMeta, CycleError, FreezeHeight, Graph, Materializer, MutationStore, @@ -24,7 +24,7 @@ pub enum CheckpointDownload { /// Rebuilds the graph from an S3 checkpoint plus WAL replay. pub struct RebuildFromCheckpoint<'a, V: VectorStore> { pub graph_store: &'a GraphPg, - pub s3_client: &'a S3Client, + pub s3_client: &'a ObjectStoreClient, pub bucket: String, pub download: CheckpointDownload, } @@ -32,7 +32,7 @@ pub struct RebuildFromCheckpoint<'a, V: VectorStore> { impl<'a, V: VectorStore + Send + Sync> RebuildFromCheckpoint<'a, V> { pub fn new( graph_store: &'a GraphPg, - s3_client: &'a S3Client, + s3_client: &'a ObjectStoreClient, bucket: String, download: CheckpointDownload, ) -> Self { diff --git a/iris-mpc-cpu/src/checkpoint_protocol/runner.rs b/iris-mpc-cpu/src/checkpoint_protocol/runner.rs index e007f1c35e..13ddc641db 100644 --- a/iris-mpc-cpu/src/checkpoint_protocol/runner.rs +++ b/iris-mpc-cpu/src/checkpoint_protocol/runner.rs @@ -9,8 +9,8 @@ use std::time::{Duration, Instant}; use ampc_actor_utils::network::mpc::NetworkHandle; use ampc_actor_utils::network::tcp::deserialize_yaml_json_string; -use aws_sdk_s3::Client as S3Client; use eyre::{eyre, Result}; +use iris_mpc_common::object_store::ObjectStoreClient; use serde::Deserialize; use serde_with::{serde_as, DurationSeconds}; use tokio::sync::RwLock; @@ -119,7 +119,7 @@ pub struct SidecarConfig { pub async fn sidecar_main( cfg: SidecarConfig, graph_store: &GraphPg, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, networking: &mut Box, shutdown_ct: CancellationToken, ) -> Result<()> { @@ -206,7 +206,7 @@ pub async fn sidecar_main( async fn sidecar_cycle( cfg: &SidecarConfig, graph_store: &GraphPg, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, networking: &mut Box, ) -> Result { let channel = tokio::time::timeout(cfg.make_connections_timeout, networking.control_channel()) @@ -286,7 +286,7 @@ const RESTART_RETRY_DELAY: Duration = Duration::from_secs(5); /// proceeds even when there are no new mutations beyond the base. pub async fn restart_from_checkpoint( graph_store: &GraphPg, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, bucket: String, networking: &mut Box, target: BothEyes>>, diff --git a/iris-mpc-cpu/src/checkpoint_protocol/terminal.rs b/iris-mpc-cpu/src/checkpoint_protocol/terminal.rs index d0a9a528fb..ac50fd1afe 100644 --- a/iris-mpc-cpu/src/checkpoint_protocol/terminal.rs +++ b/iris-mpc-cpu/src/checkpoint_protocol/terminal.rs @@ -1,5 +1,5 @@ use async_trait::async_trait; -use aws_sdk_s3::Client as S3Client; +use iris_mpc_common::object_store::ObjectStoreClient; use std::sync::Arc; use tokio::sync::{oneshot, RwLock}; @@ -24,7 +24,7 @@ use crate::utils::serialization::graph::GraphFormat; /// track per-mutation iris ids; the field is observability-only on this path). pub struct UploadAndRecord<'a, V: VectorStore> { pub graph_store: &'a GraphPg, - pub s3_client: &'a S3Client, + pub s3_client: &'a ObjectStoreClient, pub bucket: String, pub party_id: usize, pub is_archival: bool, @@ -34,7 +34,7 @@ pub struct UploadAndRecord<'a, V: VectorStore> { impl<'a, V: VectorStore + Send + Sync> UploadAndRecord<'a, V> { pub fn new( graph_store: &'a GraphPg, - s3_client: &'a S3Client, + s3_client: &'a ObjectStoreClient, bucket: String, party_id: usize, is_archival: bool, diff --git a/iris-mpc-cpu/src/genesis/state_accessor.rs b/iris-mpc-cpu/src/genesis/state_accessor.rs index b1b7de6c9d..4340a0567a 100644 --- a/iris-mpc-cpu/src/genesis/state_accessor.rs +++ b/iris-mpc-cpu/src/genesis/state_accessor.rs @@ -3,8 +3,8 @@ use crate::{ execution::hawk_main::HawkOps, hawkers::aby3::aby3_store::Aby3Store, hnsw::graph::graph_store::GraphPg, }; -use aws_sdk_s3::Client as S3_Client; use eyre::Result; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use iris_mpc_common::{config::Config, helpers::sync::Modification, SerialId}; use iris_mpc_store::Store; use serde::{de::DeserializeOwned, Deserialize, Serialize}; @@ -27,7 +27,7 @@ pub const STATE_KEY_LAST_INDEXED_MODIFICATION_ID: &str = "last_indexed_modificat /// # Arguments /// /// * `config` - Application configuration instance. -/// * `s3_client` - A configured AWS S3 client instance. +/// * `s3_client` - A configured object-store client. /// * `max_indexation_id` - Maximum Iris serial identifier to be indexed. /// /// # Returns @@ -36,7 +36,7 @@ pub const STATE_KEY_LAST_INDEXED_MODIFICATION_ID: &str = "last_indexed_modificat /// pub async fn get_iris_deletions( config: &Config, - s3_client: &S3_Client, + s3_client: &ObjectStoreClient, max_indexation_id: SerialId, ) -> Result, IndexationError> { // Struct for deserialization. @@ -45,35 +45,33 @@ pub async fn get_iris_deletions( deleted_serial_ids: Vec, } - // Set bucket and key based on environment - let s3_bucket = format!("wf-smpcv2-{}-sync-protocol", config.environment); - let s3_key = format!("{}_deleted_serial_ids.json", config.environment); + let store_location = config.iris_deletions_store_location(); + let object_key = format!("{}_deleted_serial_ids.json", config.environment); tracing::info!( - "Fetching deleted serial ids from S3 bucket: {}, key: {}", - s3_bucket, - s3_key + "Fetching deleted serial ids from object store: {}, key: {}", + store_location, + object_key ); - // Fetch from S3. - let s3_response = s3_client - .get_object() - .bucket(&s3_bucket) - .key(&s3_key) - .send() - .await - .map_err(|err| { - tracing::error!("Failed to download file from S3: {}", err); - IndexationError::AwsS3ObjectDownload - })?; + // Fetch from object storage. + let store = s3_client.store(&store_location).map_err(|err| { + tracing::error!("Failed to construct object store: {err}"); + IndexationError::AwsS3ObjectDownload + })?; + let location = path(&object_key).map_err(|err| { + tracing::error!("Invalid object key: {err}"); + IndexationError::AwsS3ObjectDownload + })?; + let response = store.get(&location).await.map_err(|err| { + tracing::error!("Failed to download file from object store: {}", err); + IndexationError::AwsS3ObjectDownload + })?; - // Consume S3 object stream. - let s3_object_body = s3_response.body.collect().await.map_err(|err| { + let s3_object_bytes = response.bytes().await.map_err(|err| { tracing::error!("Failed to get object body: {}", err); IndexationError::AwsS3ObjectDeserialize })?; - // Decode S3 object bytes. - let s3_object_bytes = s3_object_body.into_bytes(); let S3Object { deleted_serial_ids } = serde_json::from_slice(&s3_object_bytes).map_err(|err| { tracing::error!("Failed to deserialize S3 object: {}", err); @@ -282,6 +280,42 @@ async fn unset_state_element( Ok(()) } +#[cfg(test)] +mod object_store_tests { + use super::get_iris_deletions; + use iris_mpc_common::{ + config::Config, + object_store::{path, ObjectStoreClient, ObjectStoreExt}, + }; + + #[tokio::test] + async fn reads_deletions_from_configured_object_store_location() { + let store_location = "memory:///deletion-snapshots"; + let config: Config = serde_json::from_value(serde_json::json!({ + "environment": "test", + "iris_deletions_store_location": store_location + })) + .unwrap(); + let client = ObjectStoreClient::new(None, false); + let store = client.store(store_location).unwrap(); + store + .put( + &path("test_deleted_serial_ids.json").unwrap(), + serde_json::to_vec(&serde_json::json!({ + "deleted_serial_ids": [1, 3, 8] + })) + .unwrap() + .into(), + ) + .await + .unwrap(); + + let deletions = get_iris_deletions(&config, &client, 5).await.unwrap(); + + assert_eq!(deletions, vec![1, 3]); + } +} + #[cfg(test)] #[cfg(feature = "db_dependent")] mod tests { diff --git a/iris-mpc-cpu/src/graph_checkpoint/s3_client/mod.rs b/iris-mpc-cpu/src/graph_checkpoint/s3_client/mod.rs index c709cd2446..6160391d24 100644 --- a/iris-mpc-cpu/src/graph_checkpoint/s3_client/mod.rs +++ b/iris-mpc-cpu/src/graph_checkpoint/s3_client/mod.rs @@ -3,10 +3,9 @@ mod streaming; mod streaming_download; use std::{io::Cursor, time::Instant}; -use aws_sdk_s3::primitives::ByteStream; -use aws_sdk_s3::Client as S3Client; use bytes::Bytes; use eyre::{bail, eyre, Result}; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use crate::{ execution::hawk_main::{BothEyes, GraphRef, LEFT, RIGHT}, @@ -32,7 +31,7 @@ pub async fn upload_graph_checkpoint( bucket: &str, party_id: usize, graph_mem: &BothEyes, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, last_indexed_iris_id: SerialId, last_indexed_modification_id: i64, graph_mutation_id: Option, @@ -71,7 +70,7 @@ pub async fn upload_graph_checkpoint_plaintext( bucket: &str, party_id: usize, graph_mem: &BothEyes, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, last_indexed_iris_id: SerialId, last_indexed_modification_id: i64, graph_mutation_id: Option, @@ -106,7 +105,7 @@ pub async fn upload_graph_checkpoint_plaintext( async fn _upload_graph_checkpoint( bucket: &str, party_id: usize, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, last_indexed_iris_id: SerialId, last_indexed_modification_id: i64, graph_mutation_id: Option, @@ -169,7 +168,7 @@ async fn _upload_graph_checkpoint( } pub async fn download_graph_checkpoint( - s3_client: &S3Client, + s3_client: &ObjectStoreClient, bucket: &str, state: &GraphCheckpointState, ) -> Result> { @@ -200,7 +199,7 @@ pub async fn download_graph_checkpoint( // this is used for the genesis integration tests. // it does not convert between graph types pub async fn download_genesis_checkpoint_plaintext( - s3_client: &S3Client, + s3_client: &ObjectStoreClient, bucket: &str, state: &GraphCheckpointState, ) -> Result> { @@ -280,7 +279,7 @@ pub async fn save_checkpoint_state( /// recorded by all parties (the startup-agreed common checkpoint). pub async fn cleanup_checkpoints( bucket: &str, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, current_state: &GraphCheckpointState, retain_from_id: Option, graph_store: &GraphPg, @@ -320,54 +319,42 @@ pub async fn cleanup_checkpoints( Ok(()) } -/// Verifies that the S3 client has read, write, and delete access to the -/// checkpoint bucket. Uploads a small sentinel object, reads it back, and +/// Verifies read, write, and delete access to the checkpoint store. Uploads a +/// small sentinel object, reads it back, and /// deletes it. This catches misconfigured buckets/regions/IAM before any /// mutations occur. pub async fn verify_s3_checkpoint_access( - s3_client: &S3Client, + s3_client: &ObjectStoreClient, bucket: &str, party_id: usize, ) -> Result<()> { let key = format!("genesis/{party_id}/_access_check"); let body = b"access_check"; + let store = s3_client.store(bucket)?; + let location = path(&key)?; + // Write - s3_client - .put_object() - .bucket(bucket) - .key(&key) - .body(ByteStream::from_static(body)) - .send() + store + .put(&location, Bytes::from_static(body).into()) .await - .map_err(|e| eyre!("S3 checkpoint bucket write check failed: {e}"))?; + .map_err(|e| eyre!("Checkpoint object-store write check failed: {e}"))?; // Read - let resp = s3_client - .get_object() - .bucket(bucket) - .key(&key) - .send() + let data = store + .get(&location) .await - .map_err(|e| eyre!("S3 checkpoint bucket read check failed: {e}"))?; - let data = resp - .body - .collect() + .map_err(|e| eyre!("Checkpoint object-store read check failed: {e}"))? + .bytes() .await - .map_err(|e| eyre!("S3 checkpoint bucket read check failed to collect body: {e}"))?; - if data.into_bytes().as_ref() != body { - bail!("S3 checkpoint bucket read check returned unexpected content"); + .map_err(|e| eyre!("Checkpoint object-store body read failed: {e}"))?; + if data.as_ref() != body { + bail!("Checkpoint object-store read check returned unexpected content"); } // Delete - if let Err(e) = s3_client - .delete_object() - .bucket(bucket) - .key(&key) - .send() - .await - { - tracing::warn!("S3 checkpoint bucket delete check failed: {e}"); + if let Err(e) = store.delete(&location).await { + tracing::warn!("Checkpoint object-store delete check failed: {e}"); } Ok(()) @@ -379,29 +366,21 @@ pub async fn verify_s3_checkpoint_access( /// /// # Arguments /// -/// * `s3_client` - Authenticated S3 client. -/// * `bucket` - Name of the S3 bucket to query. +/// * `s3_client` - Configured object-store client. +/// * `bucket` - Object-store location to query. /// * `key` - Object key to check for existence. -pub async fn s3_key_exists(s3_client: &S3Client, bucket: &str, key: &str) -> Result { - match s3_client.head_object().bucket(bucket).key(key).send().await { +pub async fn s3_key_exists(s3_client: &ObjectStoreClient, bucket: &str, key: &str) -> Result { + let store = s3_client.store(bucket)?; + let location = path(key)?; + match store.head(&location).await { Ok(_) => Ok(true), - Err(e) => { - // `head_object` returns a 404 when the key does not exist. - // The SDK surfaces this as a `NotFound` service error. - if e.as_service_error() - .map(|se| se.is_not_found()) - .unwrap_or(false) - { - Ok(false) - } else { - Err(eyre!("S3 head_object failed for s3://{bucket}/{key}: {e}")) - } - } + Err(object_store::Error::NotFound { .. }) => Ok(false), + Err(e) => Err(eyre!("Object metadata failed for {bucket}/{key}: {e}")), } } async fn download_and_hash( - s3_client: &S3Client, + s3_client: &ObjectStoreClient, bucket: &str, state: &GraphCheckpointState, ) -> Result { diff --git a/iris-mpc-cpu/src/graph_checkpoint/s3_client/multipart.rs b/iris-mpc-cpu/src/graph_checkpoint/s3_client/multipart.rs index af281b3599..3caebc8025 100644 --- a/iris-mpc-cpu/src/graph_checkpoint/s3_client/multipart.rs +++ b/iris-mpc-cpu/src/graph_checkpoint/s3_client/multipart.rs @@ -1,354 +1,152 @@ -//! S3 Graph Checkpoint Module -//! -//! This module provides functionality for storing and loading graph checkpoints -//! needed by genesis and hawk. +//! Backend-independent graph checkpoint upload and download helpers. -use aws_sdk_s3::{ - types::{CompletedMultipartUpload, CompletedPart}, - Client as S3Client, -}; use bytes::{Bytes, BytesMut}; use eyre::{eyre, Result}; -use std::{sync::Arc, time::Duration}; -use tokio::{sync::Semaphore, task::JoinSet, time::sleep}; +use futures::{stream, StreamExt, TryStreamExt}; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; +use std::{future::Future, time::Duration}; +use tokio::time::sleep; -pub const DEFAULT_CHECKPOINT_CHUNK_SIZE: usize = 100 * 1024 * 1024; // 100 MB chunks +pub const DEFAULT_CHECKPOINT_CHUNK_SIZE: usize = 100 * 1024 * 1024; pub const DEFAULT_CHECKPOINT_PARALLELISM: usize = 32; -const MULTIPART_THRESHOLD: usize = 5 * 1024 * 1024; // 5MB - S3 multipart minimum part size +const MULTIPART_THRESHOLD: usize = 5 * 1024 * 1024; +const MAX_RETRIES: usize = 3; -/// Uploads checkpoint data to S3. -/// Uses simple PUT for files under 5MB, multipart upload for larger files. pub async fn upload_graph( - s3_client: &S3Client, - bucket: &str, + client: &ObjectStoreClient, + store_location: &str, key: &str, data: Bytes, ) -> Result<()> { tracing::info!( - "Uploading graph checkpoint to S3: bucket={}, key={}, size={}", - bucket, + "Uploading graph checkpoint: store={}, key={}, size={}", + store_location, key, data.len() ); - let chunk_size = DEFAULT_CHECKPOINT_CHUNK_SIZE; - let upload_parallelism = DEFAULT_CHECKPOINT_PARALLELISM; - + let store = client.store(store_location)?; + let location = path(key)?; if data.len() < MULTIPART_THRESHOLD { - return upload_graph_simple(s3_client, bucket, key, data).await; + return upload_graph_simple(&store, &location, data).await; } - let mut join_set = JoinSet::new(); - let semaphore = Arc::new(Semaphore::new(upload_parallelism)); - - // Initiate Multipart Upload - let multipart_res = s3_client - .create_multipart_upload() - .bucket(bucket) - .key(key) - .send() + let mut upload = store + .put_multipart(&location) .await - .map_err(|e| eyre!("Failed to initiate: {:?}", e))?; - - let upload_id = multipart_res - .upload_id() - .ok_or_else(|| eyre!("S3 did not return an upload ID"))?; - - // Build chunks, merging last chunk if it's under 5MB (S3 minimum part size) - let data_len = data.len(); - let chunk_size = std::cmp::max(chunk_size, MULTIPART_THRESHOLD); - let mut chunks: Vec<(usize, usize)> = (0..data_len) + .map_err(|e| eyre!("Failed to initiate multipart upload: {e}"))?; + let chunk_size = DEFAULT_CHECKPOINT_CHUNK_SIZE.max(MULTIPART_THRESHOLD); + let part_futures: Vec<_> = (0..data.len()) .step_by(chunk_size) .map(|start| { - let end = (start + chunk_size).min(data_len); - (start, end) + let end = (start + chunk_size).min(data.len()); + upload.put_part(data.slice(start..end).into()) }) .collect(); - if chunks.len() >= 2 { - if let Some((last_start, last_end)) = chunks.last().copied() { - if last_end - last_start < MULTIPART_THRESHOLD { - // Merge last two chunks by adjusting the second-to-last to include the remainder - let last_two_start = chunks[chunks.len() - 2].0; - chunks.pop(); - chunks.pop(); - chunks.push((last_two_start, data_len)); - } - } - } - - // Spawn Workers for Chunks - for (i, (start, end)) in chunks.into_iter().enumerate() { - let part_number = (i + 1) as i32; - let client = s3_client.clone(); - let bucket = bucket.to_string(); - let key = key.to_string(); - let upload_id = upload_id.to_string(); - let permit = semaphore - .clone() - .acquire_owned() - .await - .map_err(|e| eyre!("failed to acquire semaphore: {}", e))?; - let body = data.slice(start..end); - - join_set.spawn(async move { - let mut attempts = 0; - let max_retries = 3; - - loop { - match client - .upload_part() - .bucket(&bucket) - .key(&key) - .upload_id(&upload_id) - .part_number(part_number) - .body(body.clone().into()) - .send() - .await - { - Ok(res) => { - let etag = res.e_tag().map(|s| s.to_string()); - tracing::debug!("part {} uploaded: e_tag={:?}", part_number, etag); - - let etag = etag.ok_or_else(|| { - eyre!("s3 didn't return ETag for part {}", part_number) - })?; - - drop(permit); - return Ok(CompletedPart::builder() - .e_tag(etag) - .part_number(part_number) - .build()); - } - Err(_e) if attempts < max_retries => { - attempts += 1; - tracing::warn!("Retry {} for part {}", attempts, part_number); - sleep(Duration::from_secs(2)).await; - } - Err(e) => { - drop(permit); - return Err(eyre!("Part {} failed: {:?}", part_number, e)); - } - } - } - }); - } - // Collect & Sort ETags - let mut completed_parts = Vec::new(); - let mut error: Option = None; - while let Some(result) = join_set.join_next().await { - match result { - Ok(Ok(part)) => completed_parts.push(part), - Ok(Err(e)) => { - error.replace(e); - break; - } - Err(e) => { - error.replace(eyre!("Join error: {:?}", e)); - break; - } + let result: object_store::Result> = stream::iter(part_futures) + .buffer_unordered(DEFAULT_CHECKPOINT_PARALLELISM) + .try_collect() + .await; + if let Err(error) = result { + if let Err(abort_error) = upload.abort().await { + tracing::warn!("Failed to abort multipart upload: {abort_error}"); } + return Err(eyre!("Multipart upload failed: {error}")); } - if let Some(e) = error { - join_set.abort_all(); - let _ = s3_client - .abort_multipart_upload() - .bucket(bucket) - .key(key) - .upload_id(upload_id) - .send() - .await; - return Err(e); - } - - completed_parts.sort_by_key(|p| p.part_number); - - // Complete Upload - s3_client - .complete_multipart_upload() - .bucket(bucket) - .key(key) - .upload_id(upload_id) - .multipart_upload( - CompletedMultipartUpload::builder() - .set_parts(Some(completed_parts)) - .build(), - ) - .send() + upload + .complete() .await - .map_err(|e| eyre!("Failed to complete upload: {:?}", e))?; - - tracing::info!("Successfully uploaded graph checkpoint to S3: key={}", key); - + .map_err(|e| eyre!("Failed to complete multipart upload: {e}"))?; + tracing::info!("Successfully uploaded graph checkpoint: key={key}"); Ok(()) } -/// Downloads the graph from s3 -pub async fn download_graph(s3_client: &S3Client, bucket: &str, key: &str) -> Result { +pub async fn download_graph( + client: &ObjectStoreClient, + store_location: &str, + key: &str, +) -> Result { tracing::info!( - "Downloading graph checkpoint from S3: bucket={}, key={}", - bucket, + "Downloading graph checkpoint: store={}, key={}", + store_location, key ); - - let chunk_size = DEFAULT_CHECKPOINT_CHUNK_SIZE; - let download_parallelism = DEFAULT_CHECKPOINT_PARALLELISM; - - // Get object metadata to find total size - let head = s3_client - .head_object() - .bucket(bucket) - .key(key) - .send() - .await - .map_err(|e| { - eyre!( - "failed to get s3 checkpoint metadata for bucket {}:, key: {}, error: {}", - bucket, - key, - e - ) - })?; - let total_size = head - .content_length() - .ok_or_else(|| eyre!("Missing content length"))? - .try_into()?; - - let mut final_data = BytesMut::zeroed(total_size); - let semaphore = Arc::new(Semaphore::new(download_parallelism)); - let mut join_set = JoinSet::new(); - - tracing::info!("Starting parallel download: {} bytes", total_size); - - // Spawn range-request workers - for start in (0..total_size).step_by(chunk_size) { - let end = std::cmp::min(start + chunk_size - 1, total_size - 1); - let client = s3_client.clone(); - let (b, k) = (bucket.to_string(), key.to_string()); - let permit = semaphore - .clone() - .acquire_owned() - .await - .map_err(|e| eyre!("failed to acquire semaphore: {e}"))?; - - join_set.spawn(async move { - let mut attempts = 0; - let range = format!("bytes={}-{}", start, end); - - loop { - let res = client - .get_object() - .bucket(&b) - .key(&k) - .range(&range) - .send() - .await; - - match res { - Ok(output) => { - let data = output - .body - .collect() - .await - .map_err(|e| eyre!("Body collect error: {:?}", e))?; - - drop(permit); - return Ok((start, data.into_bytes())); - } - Err(_e) if attempts < 3 => { - attempts += 1; - tracing::warn!("Retry {} for range {}", attempts, range); - sleep(Duration::from_secs(2)).await; - } - Err(e) => { - drop(permit); - return Err(eyre!("Range {} failed: {:?}", range, e)); - } - } - } - }); + let store = client.store(store_location)?; + let location = path(key)?; + let total_size: usize = store.head(&location).await?.size.try_into()?; + if total_size == 0 { + return Ok(Bytes::new()); } - // Assemble the pieces - while let Some(result) = join_set.join_next().await { - match result { - Ok(Ok((offset, chunk_data))) => { - let len = chunk_data.len(); - final_data[offset..offset + len].copy_from_slice(&chunk_data); - } - Ok(Err(e)) => { - join_set.abort_all(); - return Err(e); - } - Err(e) => { - join_set.abort_all(); - return Err(eyre!("Join error: {:?}", e)); + let ranges: Vec<_> = (0..total_size) + .step_by(DEFAULT_CHECKPOINT_CHUNK_SIZE) + .map(|start| start..(start + DEFAULT_CHECKPOINT_CHUNK_SIZE).min(total_size)) + .collect(); + let chunks: Vec<(usize, Bytes)> = stream::iter(ranges) + .map(|range| { + let store = store.clone(); + let location = location.clone(); + async move { + let start = range.start; + let range = (range.start as u64)..(range.end as u64); + let bytes = retry("range download", || { + store.get_range(&location, range.clone()) + }) + .await?; + Ok::<_, eyre::Report>((start, bytes)) } - } - } - - tracing::info!( - "Successfully downloaded graph checkpoint from S3: key={}, size={}", - key, - final_data.len() - ); + }) + .buffer_unordered(DEFAULT_CHECKPOINT_PARALLELISM) + .try_collect() + .await?; - Ok(final_data.freeze()) + let mut result = BytesMut::zeroed(total_size); + for (offset, chunk) in chunks { + result[offset..offset + chunk.len()].copy_from_slice(&chunk); + } + Ok(result.freeze()) } -pub async fn delete_graph(s3_client: &S3Client, bucket: &str, key: &str) -> Result<()> { - tracing::info!("Deleting graph checkpoint: bucket={}, key={}", bucket, key,); - s3_client - .delete_object() - .bucket(bucket) - .key(key) - .send() - .await - .map_err(|e| { - tracing::error!("Failed to delete graph from S3: {:?}", e); - e - })?; +pub async fn delete_graph( + client: &ObjectStoreClient, + store_location: &str, + key: &str, +) -> Result<()> { + let store = client.store(store_location)?; + store.delete(&path(key)?).await?; Ok(()) } -/// Simple PUT upload for small files (under 5MB). async fn upload_graph_simple( - s3_client: &S3Client, - bucket: &str, - key: &str, + store: &iris_mpc_common::object_store::ObjectStoreRef, + location: &object_store::path::Path, data: Bytes, ) -> Result<()> { - tracing::info!( - "Using simple PUT upload for small file: bucket={}, key={}, size={}", - bucket, - key, - data.len() - ); + retry("object upload", || store.put(location, data.clone().into())) + .await + .map(|_| ()) +} +async fn retry(operation: &str, mut f: F) -> Result +where + F: FnMut() -> Fut, + Fut: Future>, +{ let mut attempts = 0; - let max_retries = 3; - loop { - match s3_client - .put_object() - .bucket(bucket) - .key(key) - .body(data.to_vec().into()) - .send() - .await - { - Ok(res) => { - tracing::info!("Simple PUT upload completed: e_tag={:?}", res.e_tag()); - return Ok(()); - } - Err(_e) if attempts < max_retries => { + match f().await { + Ok(value) => return Ok(value), + Err(_) if attempts < MAX_RETRIES => { attempts += 1; - tracing::warn!("Retry {} for simple PUT upload", attempts); + tracing::warn!("Retry {attempts} for {operation}"); sleep(Duration::from_secs(2)).await; } - Err(e) => { - return Err(eyre!("Simple PUT upload failed: {:?}", e)); + Err(error) => { + return Err(eyre!( + "{operation} failed after {attempts} retries: {error}" + )); } } } diff --git a/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming.rs b/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming.rs index ac8ea7e665..d12f4624dc 100644 --- a/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming.rs +++ b/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming.rs @@ -1,12 +1,12 @@ -//! Streaming serialize + S3 multipart upload. +//! Streaming serialize + multipart object-store upload. //! -//! [`stream_serialize_and_upload_with`] serializes a value directly into an -//! S3 multipart upload without materializing the full byte buffer in memory. +//! [`stream_serialize_and_upload_with`] serializes a value directly into a +//! multipart object-store upload without materializing the full byte buffer in memory. //! The serializer runs on a blocking thread and writes into a bounded async //! duplex pipe; the read side spawns up to `parallelism` concurrent //! `UploadPart` tasks as full chunks become available. Outside of the //! serializer's own working set, peak memory is roughly -//! `(parallelism + 1) * part_size` plus the AWS SDK's request buffers. +//! `(parallelism + 1) * part_size` plus the backend's request buffers. //! //! Compare to [`super::multipart::upload_graph`], which requires the caller //! to pass a fully buffered `Bytes` payload (doubling memory for large @@ -15,26 +15,19 @@ use std::{ io::{BufWriter, Write}, sync::Arc, - time::Duration, }; -use aws_sdk_s3::{ - types::{CompletedMultipartUpload, CompletedPart}, - Client as S3Client, -}; use bytes::Bytes; use eyre::{eyre, Result}; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; +use object_store::MultipartUpload; use tokio::{ io::{AsyncReadExt, DuplexStream}, sync::Semaphore, task::JoinSet, - time::sleep, }; use tokio_util::io::SyncIoBridge; -const UPLOAD_PART_MAX_RETRIES: u32 = 3; -const UPLOAD_PART_RETRY_DELAY: Duration = Duration::from_secs(2); - /// Suggested part size; also the size of the duplex pipe, which bounds /// back-pressure on the serializer. pub const DEFAULT_STREAMING_PART_SIZE: usize = 100 * 1024 * 1024; @@ -92,7 +85,7 @@ impl Write for BlakeTeeWriter { } } -/// Serialize directly into an S3 multipart upload without buffering the full +/// Serialize directly into a multipart object-store upload without buffering the full /// payload in memory. /// /// `serialize` runs on a `spawn_blocking` thread, so it must be `Send + @@ -102,15 +95,15 @@ impl Write for BlakeTeeWriter { /// a fully buffered `Vec` to `w` defeats the streaming intent and /// reintroduces a full-size second copy. /// -/// On any failure (serializer error, S3 error, task panic), the multipart +/// On any failure (serializer error, object-store error, task panic), the multipart /// upload is aborted before the function returns. On success, /// `complete_multipart_upload` is called and the object is durable. /// /// Callers without a reason to override should pass /// [`DEFAULT_STREAMING_PART_SIZE`] and [`DEFAULT_STREAMING_PARALLELISM`]. pub async fn stream_serialize_and_upload_with( - s3_client: &S3Client, - bucket: &str, + client: &ObjectStoreClient, + store_location: &str, key: &str, serialize: F, part_size: usize, @@ -120,60 +113,29 @@ where F: FnOnce(&mut dyn Write) -> Result<()> + Send + 'static, { tracing::info!( - "Streaming serialize + upload: bucket={bucket}, key={key}, \ + "Streaming serialize + upload: store={store_location}, key={key}, \ part_size={part_size}, parallelism={parallelism}" ); - let init = s3_client - .create_multipart_upload() - .bucket(bucket) - .key(key) - .send() + let store = client.store(store_location)?; + let location = path(key)?; + let mut upload = store + .put_multipart(&location) .await - .map_err(|e| eyre!("create_multipart_upload failed: {e:?}"))?; - let upload_id = init - .upload_id() - .ok_or_else(|| eyre!("create_multipart_upload returned no upload_id"))? - .to_string(); - - match drive_upload( - s3_client, - bucket, - key, - &upload_id, - serialize, - part_size, - parallelism, - ) - .await - { - Ok(parts) => { - s3_client - .complete_multipart_upload() - .bucket(bucket) - .key(key) - .upload_id(&upload_id) - .multipart_upload( - CompletedMultipartUpload::builder() - .set_parts(Some(parts)) - .build(), - ) - .send() + .map_err(|e| eyre!("multipart upload creation failed: {e:?}"))?; + + match drive_upload(&mut *upload, serialize, part_size, parallelism).await { + Ok(()) => { + upload + .complete() .await - .map_err(|e| eyre!("complete_multipart_upload failed: {e:?}"))?; + .map_err(|e| eyre!("multipart upload completion failed: {e:?}"))?; tracing::info!("Streaming upload complete: key={key}"); Ok(()) } Err(e) => { - if let Err(abort_err) = s3_client - .abort_multipart_upload() - .bucket(bucket) - .key(key) - .upload_id(&upload_id) - .send() - .await - { - tracing::warn!("abort_multipart_upload after failure also failed: {abort_err:?}"); + if let Err(abort_err) = upload.abort().await { + tracing::warn!("multipart upload abort after failure also failed: {abort_err:?}"); } Err(e) } @@ -181,14 +143,11 @@ where } async fn drive_upload( - s3_client: &S3Client, - bucket: &str, - key: &str, - upload_id: &str, + upload: &mut dyn MultipartUpload, serialize: F, part_size: usize, parallelism: usize, -) -> Result> +) -> Result<()> where F: FnOnce(&mut dyn Write) -> Result<()> + Send + 'static, { @@ -208,16 +167,7 @@ where Ok(()) }); - let upload_result = run_upload_loop( - s3_client.clone(), - bucket.to_string(), - key.to_string(), - upload_id.to_string(), - reader, - part_size, - parallelism, - ) - .await; + let upload_result = run_upload_loop(upload, reader, part_size, parallelism).await; let serialize_result = serialize_handle .await @@ -227,24 +177,20 @@ where // part failure it drops the reader, which surfaces as a broken-pipe error // in the serializer. The upload error is the root cause. match (upload_result, serialize_result) { - (Ok(parts), Ok(())) => Ok(parts), + (Ok(()), Ok(())) => Ok(()), (Err(upload_err), _) => Err(upload_err), (Ok(_), Err(serialize_err)) => Err(serialize_err), } } async fn run_upload_loop( - s3_client: S3Client, - bucket: String, - key: String, - upload_id: String, + upload: &mut dyn MultipartUpload, mut reader: DuplexStream, part_size: usize, parallelism: usize, -) -> Result> { +) -> Result<()> { let semaphore = Arc::new(Semaphore::new(parallelism)); - let mut join_set: JoinSet> = JoinSet::new(); - let mut parts: Vec = Vec::new(); + let mut join_set: JoinSet> = JoinSet::new(); let mut part_number: i32 = 1; loop { @@ -252,7 +198,7 @@ async fn run_upload_loop( // spawning uploads after the upload has already gone wrong. while let Some(res) = join_set.try_join_next() { match res { - Ok(Ok(part)) => parts.push(part), + Ok(Ok(part_number)) => tracing::debug!("uploaded part {part_number}"), Ok(Err(e)) => { join_set.abort_all(); return Err(e); @@ -290,47 +236,15 @@ async fn run_upload_loop( .map_err(|e| eyre!("semaphore acquire: {e}"))?; let pn = part_number; - let client = s3_client.clone(); - let bucket = bucket.clone(); - let key = key.clone(); - let upload_id = upload_id.clone(); - // `Bytes` so retries clone cheaply (Arc bump) instead of copying. let body = Bytes::from(buf); + let upload_part = upload.put_part(body.into()); join_set.spawn(async move { let _permit = permit; - let mut attempts: u32 = 0; - loop { - match client - .upload_part() - .bucket(&bucket) - .key(&key) - .upload_id(&upload_id) - .part_number(pn) - .body(body.clone().into()) - .send() - .await - { - Ok(res) => { - let etag = res - .e_tag() - .ok_or_else(|| eyre!("upload_part {pn} returned no e_tag"))? - .to_string(); - tracing::debug!("uploaded part {pn} (etag={etag})"); - return Ok(CompletedPart::builder().e_tag(etag).part_number(pn).build()); - } - Err(_) if attempts < UPLOAD_PART_MAX_RETRIES => { - attempts += 1; - tracing::warn!("Retry {attempts} for part {pn}"); - sleep(UPLOAD_PART_RETRY_DELAY).await; - } - Err(e) => { - return Err(eyre!( - "upload_part {pn} failed after {attempts} retries: {e:?}" - )); - } - } - } + upload_part + .await + .map_err(|e| eyre!("upload part {pn} failed: {e}"))?; + Ok(pn) }); part_number += 1; @@ -343,7 +257,7 @@ async fn run_upload_loop( let mut first_error: Option = None; while let Some(res) = join_set.join_next().await { match res { - Ok(Ok(part)) => parts.push(part), + Ok(Ok(part_number)) => tracing::debug!("uploaded part {part_number}"), Ok(Err(e)) => { first_error = Some(e); break; @@ -360,8 +274,7 @@ async fn run_upload_loop( return Err(e); } - parts.sort_by_key(|p| p.part_number); - Ok(parts) + Ok(()) } #[cfg(test)] diff --git a/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming_download.rs b/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming_download.rs index 6c899a1c37..3779ad86bd 100644 --- a/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming_download.rs +++ b/iris-mpc-cpu/src/graph_checkpoint/s3_client/streaming_download.rs @@ -1,4 +1,4 @@ -//! Streaming S3 download + bincode deserialize. +//! Streaming object-store download + bincode deserialize. //! //! `stream_download_and_deserialize` fetches the object via a sequence of //! HTTP range GETs, tees the concatenated bytes through a BLAKE3 hasher @@ -36,16 +36,17 @@ //! the canonical bincode encoding of `T` and then EOFs. Any trailing bytes //! cause the tee task to fail with `BrokenPipe` (the deserializer drops //! its end of the duplex once bincode is done), which surfaces as an -//! error from the function. For S3 objects produced by +//! error from the function. For objects produced by //! `stream_serialize_and_upload` this is automatic; callers //! wiring up other sources must respect the contract. use std::time::Duration; -use aws_sdk_s3::Client as S3Client; use bytes::Bytes; use eyre::{eyre, Result}; use futures::stream::{self, Stream, StreamExt}; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt, ObjectStoreRef}; +use object_store::path::Path; use serde::de::DeserializeOwned; use tokio::io::{AsyncRead, AsyncReadExt, AsyncWriteExt}; use tokio::time::sleep; @@ -84,24 +85,24 @@ pub const DEFAULT_DOWNLOAD_PARALLELISM: usize = 32; /// this each one would block across the async bridge individually. const GRAPH_DECODE_BUFFER: usize = 1024 * 1024; -/// Stream the object at `s3://{bucket}/{key}` through BLAKE3 and bincode +/// Stream an object through BLAKE3 and bincode /// in one pass. Returns the deserialized value and the BLAKE3 digest of /// the downloaded bytes. /// /// Callers verify the digest against the expected checkpoint hash before -/// trusting `value`. A mismatch indicates the S3 object diverges from the +/// trusting `value`. A mismatch indicates the stored object diverges from the /// hash recorded in the database. pub async fn stream_download_and_deserialize( - s3_client: &S3Client, - bucket: &str, + client: &ObjectStoreClient, + store_location: &str, key: &str, ) -> Result<(T, [u8; 32])> where T: DeserializeOwned + Send + 'static, { stream_download_and_deserialize_with( - s3_client, - bucket, + client, + store_location, key, DEFAULT_DOWNLOAD_PIPE_CAPACITY, DEFAULT_DOWNLOAD_RANGE_SIZE, @@ -111,8 +112,8 @@ where } pub async fn stream_download_and_deserialize_with( - s3_client: &S3Client, - bucket: &str, + client: &ObjectStoreClient, + store_location: &str, key: &str, pipe_capacity: usize, range_size: usize, @@ -122,7 +123,7 @@ where T: DeserializeOwned + Send + 'static, { tracing::info!( - "Streaming download + deserialize: bucket={bucket}, key={key}, \ + "Streaming download + deserialize: store={store_location}, key={key}, \ pipe_capacity={pipe_capacity}, range_size={range_size}, parallelism={parallelism}" ); @@ -130,12 +131,13 @@ where return Err(eyre!("range_size must be > 0")); } - let total_size = head_object_size_with_retry(s3_client, bucket, key).await?; + let store = client.store(store_location)?; + let location = path(key)?; + let total_size = head_object_size_with_retry(&store, &location).await?; let stream = Box::pin(range_stream( - s3_client.clone(), - bucket.to_string(), - key.to_string(), + store, + location, total_size, range_size as u64, parallelism, @@ -144,19 +146,19 @@ where deserialize_and_hash_from(reader, pipe_capacity).await } -/// Stream the object at `s3://{bucket}/{key}` through BLAKE3 and deserialize +/// Stream an object through BLAKE3 and deserialize /// into a `[GraphMem; 2]`. Bytes are fed to the decoder incrementally as ranges /// arrive; the decode itself ([`read_graph_pair_streaming`]) is the standard /// derived path. pub async fn stream_download_and_deserialize_graph_pair( - s3_client: &S3Client, - bucket: &str, + client: &ObjectStoreClient, + store_location: &str, key: &str, format: GraphFormat, ) -> Result<([GraphMem; 2], [u8; 32])> { stream_download_and_deserialize_graph_pair_with( - s3_client, - bucket, + client, + store_location, key, format, DEFAULT_DOWNLOAD_PIPE_CAPACITY, @@ -169,8 +171,8 @@ pub async fn stream_download_and_deserialize_graph_pair( /// Like [`stream_download_and_deserialize_graph_pair`] but with explicit /// `pipe_capacity`, `range_size`, and `parallelism` knobs. pub async fn stream_download_and_deserialize_graph_pair_with( - s3_client: &S3Client, - bucket: &str, + client: &ObjectStoreClient, + store_location: &str, key: &str, format: GraphFormat, pipe_capacity: usize, @@ -178,7 +180,7 @@ pub async fn stream_download_and_deserialize_graph_pair_with( parallelism: usize, ) -> Result<([GraphMem; 2], [u8; 32])> { tracing::info!( - "Streaming download + deserialize graph pair: bucket={bucket}, key={key}, \ + "Streaming download + deserialize graph pair: store={store_location}, key={key}, \ format={format}, pipe_capacity={pipe_capacity}, range_size={range_size}, \ parallelism={parallelism}" ); @@ -191,11 +193,12 @@ pub async fn stream_download_and_deserialize_graph_pair_with( return Err(eyre!("pipe_capacity must be > 0")); } - let total_size = head_object_size_with_retry(s3_client, bucket, key).await?; + let store = client.store(store_location)?; + let location = path(key)?; + let total_size = head_object_size_with_retry(&store, &location).await?; let stream = Box::pin(range_stream( - s3_client.clone(), - bucket.to_string(), - key.to_string(), + store, + location, total_size, range_size as u64, parallelism, @@ -211,28 +214,20 @@ pub async fn stream_download_and_deserialize_graph_pair_with( .await } -/// `HeadObject` for `content_length`, retried per [`RANGE_MAX_RETRIES`]. -async fn head_object_size_with_retry(s3_client: &S3Client, bucket: &str, key: &str) -> Result { +/// Object metadata size probe, retried per [`RANGE_MAX_RETRIES`]. +async fn head_object_size_with_retry(store: &ObjectStoreRef, key: &Path) -> Result { let mut attempts: u32 = 0; loop { - match s3_client.head_object().bucket(bucket).key(key).send().await { - Ok(out) => { - let len = out - .content_length() - .ok_or_else(|| eyre!("head_object {bucket}/{key}: missing content_length"))?; - if len < 0 { - return Err(eyre!("head_object {bucket}/{key}: negative content_length")); - } - return Ok(len as u64); - } + match store.head(key).await { + Ok(out) => return Ok(out.size), Err(e) if attempts < RANGE_MAX_RETRIES => { attempts += 1; - tracing::warn!("Retry {attempts} for head_object s3://{bucket}/{key}: {e:?}"); + tracing::warn!("Retry {attempts} for object metadata {key}: {e:?}"); sleep(RANGE_RETRY_DELAY).await; } Err(e) => { return Err(eyre!( - "head_object s3://{bucket}/{key} failed after {attempts} retries: {e:?}" + "object metadata {key} failed after {attempts} retries: {e:?}" )); } } @@ -250,9 +245,8 @@ async fn head_object_size_with_retry(s3_client: &S3Client, bucket: &str, key: &s /// `StreamReader` fuses on the first `Err` and surfaces it to its reader, /// at which point the dropped stream cancels any still-in-flight fetches. fn range_stream( - s3_client: S3Client, - bucket: String, - key: String, + store: ObjectStoreRef, + key: Path, total_size: u64, range_size: u64, parallelism: usize, @@ -267,54 +261,47 @@ fn range_stream( stream::iter(ranges) .map(move |(start, end_inclusive)| { - let s3_client = s3_client.clone(); - let bucket = bucket.clone(); + let store = store.clone(); let key = key.clone(); async move { - let range = format!("bytes={start}-{end_inclusive}"); - fetch_range(&s3_client, &bucket, &key, &range) + fetch_range(&store, &key, start..end_inclusive + 1) .await .map_err(|e| { - std::io::Error::other(format!("range {range} of s3://{bucket}/{key}: {e}")) + std::io::Error::other(format!( + "range {start}-{end_inclusive} of {key}: {e}" + )) }) } }) .buffered(parallelism.max(1)) } -/// Fetch one S3 range. Buffers the response body in memory so a mid-body +/// Fetch one object range. Buffers the response body in memory so a mid-body /// network failure cleanly maps to a retry of the same range — no partial /// bytes leak into the downstream tee. -async fn fetch_range(s3_client: &S3Client, bucket: &str, key: &str, range: &str) -> Result { +async fn fetch_range( + store: &ObjectStoreRef, + key: &Path, + range: std::ops::Range, +) -> Result { let mut attempts: u32 = 0; loop { - let attempt = async { - let out = s3_client - .get_object() - .bucket(bucket) - .key(key) - .range(range) - .send() - .await - .map_err(|e| eyre!("get_object: {e:?}"))?; - let agg = out - .body - .collect() - .await - .map_err(|e| eyre!("body collect: {e:?}"))?; - Ok::<_, eyre::Report>(agg.into_bytes()) - } - .await; + let attempt = store + .get_range(key, range.clone()) + .await + .map_err(|e| eyre!("get range: {e:?}")); match attempt { Ok(bytes) => return Ok(bytes), Err(e) if attempts < RANGE_MAX_RETRIES => { attempts += 1; - tracing::warn!("Retry {attempts} for range {range}: {e}"); + tracing::warn!("Retry {attempts} for range {range:?}: {e}"); sleep(RANGE_RETRY_DELAY).await; } Err(e) => { - return Err(eyre!("range {range} failed after {attempts} retries: {e}")); + return Err(eyre!( + "range {range:?} failed after {attempts} retries: {e}" + )); } } } diff --git a/iris-mpc-cpu/src/graph_checkpoint/synchronize.rs b/iris-mpc-cpu/src/graph_checkpoint/synchronize.rs index e22639e64b..eac903d6b4 100644 --- a/iris-mpc-cpu/src/graph_checkpoint/synchronize.rs +++ b/iris-mpc-cpu/src/graph_checkpoint/synchronize.rs @@ -1,8 +1,8 @@ use std::collections::HashMap; -use aws_sdk_s3::Client; use eyre::{bail, eyre, Result}; use iris_mpc_common::helpers::sync::{SyncResult, SyncState}; +use iris_mpc_common::object_store::ObjectStoreClient; use itertools::izip; use super::{download_graph_checkpoint, GraphCheckpointState}; @@ -132,7 +132,7 @@ fn build_mutation_bytes(all_states: &[SyncState]) -> Result, wal_rows: Vec, diff --git a/iris-mpc-cpu/src/hnsw/graph/test_utils.rs b/iris-mpc-cpu/src/hnsw/graph/test_utils.rs index ca665e2844..0448841e3a 100644 --- a/iris-mpc-cpu/src/hnsw/graph/test_utils.rs +++ b/iris-mpc-cpu/src/hnsw/graph/test_utils.rs @@ -31,10 +31,10 @@ use crate::{ protocol::shared_iris::GaloisRingSharedIris, }; use aes_prng::AesRng; -use aws_sdk_s3::Client as S3Client; use clap::ValueEnum; use eyre::Result; use iris_mpc_common::iris_db::db::IrisDB; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_common::postgres::{AccessMode, PostgresClient}; use iris_mpc_store::{Store, StoredIrisRef}; use itertools::Itertools; @@ -58,7 +58,7 @@ pub struct DbContext { /// Postgres store to persist data against pub store: Store, graph_pg: GraphPg, - s3_client: S3Client, + s3_client: ObjectStoreClient, bucket: String, party_id: usize, } @@ -67,7 +67,7 @@ impl DbContext { pub async fn new( url: &str, schema: &str, - s3_client: S3Client, + s3_client: ObjectStoreClient, bucket: String, party_id: usize, ) -> Self { @@ -89,7 +89,7 @@ impl DbContext { /// A dummy (unconfigured) S3 client is created; calls to S3 will fail, but /// as long as no checkpoint operations are performed this is safe. pub async fn new_without_s3(url: &str, schema: &str, party_id: usize) -> Self { - let s3_client = S3Client::from_conf(aws_sdk_s3::config::Builder::new().build()); + let s3_client = ObjectStoreClient::new(None, false); Self::new(url, schema, s3_client, String::new(), party_id).await } diff --git a/iris-mpc-cpu/src/utils/serialization/graph.rs b/iris-mpc-cpu/src/utils/serialization/graph.rs index 5387fd0adc..3d6e1a3837 100644 --- a/iris-mpc-cpu/src/utils/serialization/graph.rs +++ b/iris-mpc-cpu/src/utils/serialization/graph.rs @@ -656,8 +656,8 @@ impl From for graph_v4::GraphV4 { /* ----------- Streaming Deserialization ---------- */ /// Read a graph pair from a byte stream. Identical to [`read_graph_pair`]; the -/// streaming benefit is in feeding `reader` incrementally (see the `s3_client` -/// streaming download), not in the decode. +/// streaming benefit is in feeding `reader` incrementally from object storage, +/// not in the decode. pub fn read_graph_pair_streaming( reader: &mut R, format: GraphFormat, diff --git a/iris-mpc-cpu/tests/streaming_s3_integration.rs b/iris-mpc-cpu/tests/streaming_s3_integration.rs index 22e217329e..08fbb91c77 100644 --- a/iris-mpc-cpu/tests/streaming_s3_integration.rs +++ b/iris-mpc-cpu/tests/streaming_s3_integration.rs @@ -1,23 +1,10 @@ -//! Integration test for `stream_serialize_and_upload_with` against an -//! S3-compatible endpoint (localstack). -//! -//! Set `S3_TEST_ENDPOINT=http://localhost:4566` to enable; otherwise both -//! tests are skipped. To run locally: -//! -//! ```sh -//! docker run --rm -d --name localstack -p 4566:4566 -e SERVICES=s3 \ -//! public.ecr.aws/localstack/localstack:4.9 -//! S3_TEST_ENDPOINT=http://localhost:4566 \ -//! cargo test -p iris-mpc-cpu --test streaming_s3_integration -//! ``` - -use std::collections::HashMap; - -use aws_sdk_s3::{ - config::{BehaviorVersion, Credentials, Region}, - Client as S3Client, Config, -}; +//! Backend-independent integration tests for checkpoint streaming using the +//! `object_store` in-memory backend. + +use std::{collections::HashMap, sync::Arc}; + use eyre::{eyre, Result}; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use iris_mpc_cpu::{ graph_checkpoint::{ stream_download_and_deserialize_graph_pair, stream_download_and_deserialize_with, @@ -33,54 +20,19 @@ use iris_mpc_cpu::{ }, }, }; +use object_store::memory::InMemory; -fn s3_test_endpoint() -> Option { - std::env::var("S3_TEST_ENDPOINT").ok() -} - -fn make_test_client(endpoint: &str) -> S3Client { - let creds = Credentials::new("test", "test", None, None, "test"); - let cfg = Config::builder() - .behavior_version(BehaviorVersion::latest()) - .region(Region::new("us-east-1")) - .credentials_provider(creds) - .endpoint_url(endpoint) - .force_path_style(true) - .build(); - S3Client::from_conf(cfg) -} - -async fn create_bucket(client: &S3Client, bucket: &str) -> Result<()> { +fn make_test_client(bucket: &str) -> ObjectStoreClient { + let client = ObjectStoreClient::new(None, false); + client.insert(bucket, Arc::new(InMemory::new())); client - .create_bucket() - .bucket(bucket) - .send() - .await - .map_err(|e| eyre!("create_bucket failed: {e:?}"))?; - Ok(()) -} - -async fn cleanup_bucket(client: &S3Client, bucket: &str) { - if let Ok(list) = client.list_objects_v2().bucket(bucket).send().await { - for obj in list.contents() { - if let Some(k) = obj.key() { - let _ = client.delete_object().bucket(bucket).key(k).send().await; - } - } - } - let _ = client.delete_bucket().bucket(bucket).send().await; } #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn streaming_upload_round_trip_multipart() -> Result<()> { - let Some(endpoint) = s3_test_endpoint() else { - eprintln!("S3_TEST_ENDPOINT not set; skipping"); - return Ok(()); - }; - let client = make_test_client(&endpoint); let bucket = format!("streaming-test-mp-{}", uuid::Uuid::new_v4()); + let client = make_test_client(&bucket); let key = "round-trip.bin"; - create_bucket(&client, &bucket).await?; // ~12 MiB payload so the upload spans 3 parts at 5 MiB each (last ~2 MiB). // S3's 5 MiB minimum applies to all-but-last, so this exercises the @@ -101,26 +53,17 @@ async fn streaming_upload_round_trip_multipart() -> Result<()> { 4, ) .await; - if let Err(e) = upload { - cleanup_bucket(&client, &bucket).await; - return Err(e); - } + upload?; let downloaded = client - .get_object() - .bucket(&bucket) - .key(key) - .send() - .await - .map_err(|e| eyre!("get_object: {e:?}"))? - .body - .collect() + .store(&bucket)? + .get(&path(key)?) .await - .map_err(|e| eyre!("collect body: {e:?}"))? + .map_err(|e| eyre!("get object: {e:?}"))? + .bytes() + .await? .to_vec(); - cleanup_bucket(&client, &bucket).await; - assert_eq!(downloaded.len(), buffered.len(), "size mismatch"); assert_eq!(downloaded, buffered, "content mismatch"); Ok(()) @@ -128,14 +71,9 @@ async fn streaming_upload_round_trip_multipart() -> Result<()> { #[tokio::test(flavor = "multi_thread", worker_threads = 2)] async fn streaming_upload_aborts_on_serializer_error() -> Result<()> { - let Some(endpoint) = s3_test_endpoint() else { - eprintln!("S3_TEST_ENDPOINT not set; skipping"); - return Ok(()); - }; - let client = make_test_client(&endpoint); let bucket = format!("streaming-test-err-{}", uuid::Uuid::new_v4()); + let client = make_test_client(&bucket); let key = "should-not-exist.bin"; - create_bucket(&client, &bucket).await?; let result = stream_serialize_and_upload_with( &client, @@ -154,30 +92,24 @@ async fn streaming_upload_aborts_on_serializer_error() -> Result<()> { ); // Object must not exist after a failed/aborted multipart upload. - let head = client.head_object().bucket(&bucket).key(key).send().await; + let head = client.store(&bucket)?.head(&path(key)?).await; assert!( head.is_err(), "object should not exist after aborted upload" ); - cleanup_bucket(&client, &bucket).await; Ok(()) } /// End-to-end: upload via `stream_serialize_and_upload_with`, download via /// `stream_download_and_deserialize_with`, assert byte-identical value and -/// blake3 hash. Exercises the S3 → `ByteStream::into_async_read` → tee → +/// blake3 hash. Exercises ranged object reads → tee → /// SyncIoBridge → bincode path that the unit tests bypass. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn streaming_download_round_trip() -> Result<()> { - let Some(endpoint) = s3_test_endpoint() else { - eprintln!("S3_TEST_ENDPOINT not set; skipping"); - return Ok(()); - }; - let client = make_test_client(&endpoint); let bucket = format!("streaming-test-dl-{}", uuid::Uuid::new_v4()); + let client = make_test_client(&bucket); let key = "round-trip.bin"; - create_bucket(&client, &bucket).await?; // ~12 MiB payload so the upload spans 3 parts; 3 MiB range_size forces // the download to issue multiple sequential GetObject(Range) calls and @@ -199,10 +131,7 @@ async fn streaming_download_round_trip() -> Result<()> { 4, ) .await; - if let Err(e) = upload { - cleanup_bucket(&client, &bucket).await; - return Err(e); - } + upload?; let download: Result<(Vec, [u8; 32])> = stream_download_and_deserialize_with( &client, @@ -213,7 +142,6 @@ async fn streaming_download_round_trip() -> Result<()> { 4, ) .await; - cleanup_bucket(&client, &bucket).await; let (got, got_hash) = download?; assert_eq!(got, payload, "deserialized value mismatch"); @@ -259,7 +187,7 @@ fn make_v3_pair() -> [GraphV3; 2] { [make_graph(100), make_graph(200)] } -/// Upload a `[GraphV3; 2]` to localstack S3, stream-download it via +/// Upload a `[GraphV3; 2]` to object storage, stream-download it via /// `stream_download_and_deserialize_graph_pair(…, GraphFormat::V3)`, and /// assert the result matches the reference `.into()` conversion. /// @@ -269,17 +197,10 @@ fn make_v3_pair() -> [GraphV3; 2] { /// - `last_update_seq_no` is 0 for both graphs (V3 has no seq_no field). #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn v3_graph_pair_streams_to_graphmem() -> Result<()> { - let Some(endpoint) = s3_test_endpoint() else { - eprintln!("S3_TEST_ENDPOINT not set; skipping"); - return Ok(()); - }; - - let client = make_test_client(&endpoint); let bucket = format!("streaming-v3-{}", uuid::Uuid::new_v4()); + let client = make_test_client(&bucket); let key = "v3-pair.bin"; - create_bucket(&client, &bucket).await?; - let pair = make_v3_pair(); // Upload as `[GraphV3; 2]` bincode. @@ -296,10 +217,7 @@ async fn v3_graph_pair_streams_to_graphmem() -> Result<()> { 4, ) .await; - if let Err(e) = upload { - cleanup_bucket(&client, &bucket).await; - return Err(e); - } + upload?; // Compute the expected BLAKE3 hash from the raw serialized bytes — the // same bytes the upload wrote and the download read. @@ -311,7 +229,6 @@ async fn v3_graph_pair_streams_to_graphmem() -> Result<()> { // Stream-download V3 bytes → `[GraphMem; 2]`. let download = stream_download_and_deserialize_graph_pair(&client, &bucket, key, GraphFormat::V3).await; - cleanup_bucket(&client, &bucket).await; let (graphs, hash) = download?; // Hash must equal BLAKE3 of the on-wire bytes. @@ -373,7 +290,7 @@ fn make_v4_pair() -> [GraphV4; 2] { [make_graph(100, 42), make_graph(200, 99)] } -/// Upload a `[GraphV4; 2]` to localstack S3, stream-download it via +/// Upload a `[GraphV4; 2]` to object storage, stream-download it via /// `stream_download_and_deserialize_graph_pair(…, GraphFormat::V4)`, and /// assert the result matches the reference `.into()` conversion. /// @@ -384,17 +301,10 @@ fn make_v4_pair() -> [GraphV4; 2] { /// - The returned BLAKE3 hash equals `blake3::hash(bincode_bytes)`. #[tokio::test(flavor = "multi_thread", worker_threads = 4)] async fn v4_graph_pair_streams_to_graphmem_seq_no_preserved() -> Result<()> { - let Some(endpoint) = s3_test_endpoint() else { - eprintln!("S3_TEST_ENDPOINT not set; skipping"); - return Ok(()); - }; - - let client = make_test_client(&endpoint); let bucket = format!("streaming-v4-{}", uuid::Uuid::new_v4()); + let client = make_test_client(&bucket); let key = "v4-pair.bin"; - create_bucket(&client, &bucket).await?; - let pair = make_v4_pair(); // Compute the expected BLAKE3 hash before moving `pair` into the upload. @@ -417,15 +327,11 @@ async fn v4_graph_pair_streams_to_graphmem_seq_no_preserved() -> Result<()> { 4, ) .await; - if let Err(e) = upload { - cleanup_bucket(&client, &bucket).await; - return Err(e); - } + upload?; // Stream-download V4 bytes → `[GraphMem; 2]`. let download = stream_download_and_deserialize_graph_pair(&client, &bucket, key, GraphFormat::V4).await; - cleanup_bucket(&client, &bucket).await; let (graphs, hash) = download?; // Hash must equal BLAKE3 of the on-wire bytes. diff --git a/iris-mpc-store/Cargo.toml b/iris-mpc-store/Cargo.toml index ee9bdf3b10..a5cb8ef300 100644 --- a/iris-mpc-store/Cargo.toml +++ b/iris-mpc-store/Cargo.toml @@ -8,9 +8,9 @@ license.workspace = true repository.workspace = true [dependencies] -aws-config.workspace = true -aws-sdk-s3.workspace = true async-trait.workspace = true +aws-config.workspace = true +object_store.workspace = true iris-mpc-common = { path = "../iris-mpc-common" } bytemuck.workspace = true dotenvy.workspace = true @@ -20,6 +20,7 @@ eyre.workspace = true itertools.workspace = true tracing.workspace = true tokio.workspace = true +bytes = "1.7" rand.workspace = true ampc-server-utils.workspace = true diff --git a/iris-mpc-store/src/loader.rs b/iris-mpc-store/src/loader.rs index f2cbdc075f..63413fe13e 100644 --- a/iris-mpc-store/src/loader.rs +++ b/iris-mpc-store/src/loader.rs @@ -1,9 +1,8 @@ -use crate::s3_importer::create_db_chunks_s3_client; +use crate::s3_importer::create_db_chunks_object_store_client; use crate::{ fetch_and_parse_chunks, last_snapshot_timestamp, DbStoredIris, S3Store, S3StoredIris, Store, }; use ampc_server_utils::shutdown_handler::ShutdownHandler; -use aws_config::Region; use eyre::{bail, Result}; use futures::stream::BoxStream; use futures::StreamExt; @@ -122,10 +121,12 @@ async fn load_iris_db_internal( let s3_load_safety_overlap_seconds = config.db_load_safety_overlap_seconds; // Construct s3 client and store - let region_provider = Region::new(region); - let shared_config = aws_config::from_env().region(region_provider).load().await; - let s3_client = create_db_chunks_s3_client(&shared_config, true); - let s3_store = S3Store::new(s3_client, s3_chunks_bucket_name.clone()); + let sdk_config = aws_config::from_env() + .region(aws_config::Region::new(region)) + .load() + .await; + let object_store_client = create_db_chunks_object_store_client(&sdk_config, true); + let s3_store = S3Store::new(object_store_client, s3_chunks_bucket_name.clone())?; let s3_arc = Arc::new(s3_store); // First fetch last snapshot from S3 diff --git a/iris-mpc-store/src/s3_importer.rs b/iris-mpc-store/src/s3_importer.rs index 07b46279aa..52ec1bf7d4 100644 --- a/iris-mpc-store/src/s3_importer.rs +++ b/iris-mpc-store/src/s3_importer.rs @@ -1,14 +1,12 @@ use ampc_server_utils::ShutdownHandler; use async_trait::async_trait; -use aws_config::{retry::RetryConfig, timeout::TimeoutConfig}; -use aws_sdk_s3::config::StalledStreamProtectionConfig; -use aws_sdk_s3::{config::Builder as S3ConfigBuilder, Client as S3Client}; -use aws_sdk_s3::{primitives::ByteStream, Client}; +use bytes::Bytes; use eyre::{bail, eyre, Result}; -use futures::{stream, StreamExt}; +use futures::{stream, StreamExt, TryStreamExt}; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt, ObjectStoreRef}; use iris_mpc_common::{VectorId, IRIS_CODE_LENGTH, MASK_CODE_LENGTH}; use std::{mem, sync::Arc, time::Duration}; -use tokio::{io::AsyncReadExt, sync::mpsc::Sender}; +use tokio::sync::mpsc::Sender; const SINGLE_ELEMENT_SIZE: usize = IRIS_CODE_LENGTH * mem::size_of::() * 2 + MASK_CODE_LENGTH * mem::size_of::() * 2 @@ -135,93 +133,51 @@ impl S3StoredIris { } } -/// Creates an S3 client specifically for database chunks with additional -/// configuration -pub fn create_db_chunks_s3_client( - shared_config: &aws_config::SdkConfig, +pub fn create_db_chunks_object_store_client( + sdk_config: &aws_config::SdkConfig, force_path_style: bool, -) -> S3Client { - let retry_config = RetryConfig::standard().with_max_attempts(5); - - // Increase S3 connect timeouts to 10s - let timeout_config = TimeoutConfig::builder() - .connect_timeout(Duration::from_secs(10)) - .build(); - - let db_chunks_s3_config = S3ConfigBuilder::from(shared_config) - // disable stalled stream protection to avoid panics during s3 import - .stalled_stream_protection(StalledStreamProtectionConfig::disabled()) - .retry_config(retry_config) - .timeout_config(timeout_config) - .force_path_style(force_path_style) - .build(); - - S3Client::from_conf(db_chunks_s3_config) +) -> ObjectStoreClient { + ObjectStoreClient::new( + sdk_config.region().map(ToString::to_string), + force_path_style, + ) + .with_aws_sdk_config(sdk_config) } #[async_trait] pub trait ObjectStore: Send + Sync + 'static { - async fn get_object(&self, key: &str, range: (usize, usize)) -> Result; + async fn get_object(&self, key: &str, range: (usize, usize)) -> Result; async fn list_objects(&self, prefix: &str) -> Result>; } pub struct S3Store { - client: Client, - bucket: String, + store: ObjectStoreRef, } impl S3Store { - pub fn new(client: Client, bucket: String) -> Self { - Self { client, bucket } + pub fn new(client: ObjectStoreClient, location: String) -> Result { + Ok(Self { + store: client.store(&location)?, + }) } } #[async_trait] impl ObjectStore for S3Store { - async fn get_object(&self, key: &str, range: (usize, usize)) -> Result { - let res = self - .client - .get_object() - .bucket(&self.bucket) - .key(key) - .range(format!("bytes={}-{}", range.0, range.1 - 1)) - .send() - .await?; - - Ok(res.body) + async fn get_object(&self, key: &str, range: (usize, usize)) -> Result { + Ok(self + .store + .get_range(&path(key)?, (range.0 as u64)..(range.1 as u64)) + .await?) } async fn list_objects(&self, prefix: &str) -> Result> { - let mut objects = Vec::new(); - let mut continuation_token = None; - - loop { - let mut request = self - .client - .list_objects_v2() - .bucket(&self.bucket) - .prefix(prefix); - - if let Some(token) = continuation_token { - request = request.continuation_token(token); - } - - let response = request.send().await?; - - objects.extend( - response - .contents() - .iter() - .filter_map(|obj| obj.key().map(String::from)), - ); - - match response.next_continuation_token() { - Some(token) => continuation_token = Some(token.to_string()), - None => break, - } - } - - Ok(objects) + Ok(self + .store + .list(Some(&path(prefix)?)) + .map_ok(|meta| meta.location.to_string()) + .try_collect() + .await?) } } @@ -396,7 +352,7 @@ async fn read_range_in_chunk( range_size: usize, tx: Sender, ) -> Result<()> { - let mut stream = store + let bytes = store .get_object( key, ( @@ -404,20 +360,11 @@ async fn read_range_in_chunk( (offset_within_chunk + range_size) * SINGLE_ELEMENT_SIZE, ), ) - .await? - .into_async_read(); - - let mut slice = vec![0_u8; SINGLE_ELEMENT_SIZE]; + .await?; - loop { - match stream.read_exact(&mut slice).await { - Ok(_) => { - let iris = S3StoredIris::from_bytes(&slice)?; - tx.send(iris).await?; - } - Err(e) if e.kind() == std::io::ErrorKind::UnexpectedEof => break, - Err(e) => return Err(e.into()), - } + for slice in bytes.chunks_exact(SINGLE_ELEMENT_SIZE) { + let iris = S3StoredIris::from_bytes(slice)?; + tx.send(iris).await?; } Ok(()) @@ -427,7 +374,6 @@ async fn read_range_in_chunk( mod tests { use super::*; use crate::DbStoredIris; - use aws_sdk_s3::primitives::SdkBody; use rand::Rng; use std::{ cmp::min, @@ -466,7 +412,7 @@ mod tests { #[async_trait] impl ObjectStore for MockStore { - async fn get_object(&self, key: &str, range: (usize, usize)) -> Result { + async fn get_object(&self, key: &str, range: (usize, usize)) -> Result { let bytes = self .objects .get(key) @@ -478,7 +424,7 @@ mod tests { let end = range.1.min(bytes.len()); let sliced_bytes = bytes[start..end].to_vec(); - Ok(ByteStream::from(SdkBody::from(sliced_bytes))) + Ok(Bytes::from(sliced_bytes)) } async fn list_objects(&self, _: &str) -> Result> { @@ -505,7 +451,7 @@ mod tests { #[async_trait::async_trait] impl ObjectStore for IntentionalFailureStore { - async fn get_object(&self, key: &str, range: (usize, usize)) -> Result { + async fn get_object(&self, key: &str, range: (usize, usize)) -> Result { let range_hash = format!("{}_{},{}", key, range.0, range.1); let mut failures = self.remaining_failures.lock().await; let n_remaining = failures @@ -531,7 +477,7 @@ mod tests { #[async_trait] impl ObjectStore for HangingStore { - async fn get_object(&self, _key: &str, _range: (usize, usize)) -> Result { + async fn get_object(&self, _key: &str, _range: (usize, usize)) -> Result { tokio::time::sleep(Duration::from_secs(3600)).await; Err(eyre::eyre!( "HangingStore: should have been cancelled before this" diff --git a/iris-mpc-upgrade-hawk/Cargo.toml b/iris-mpc-upgrade-hawk/Cargo.toml index 050e771749..385b5448aa 100644 --- a/iris-mpc-upgrade-hawk/Cargo.toml +++ b/iris-mpc-upgrade-hawk/Cargo.toml @@ -10,7 +10,7 @@ repository.workspace = true [dependencies] aws-config.workspace = true aws-sdk-rds.workspace = true -aws-sdk-s3.workspace = true +object_store.workspace = true axum.workspace = true blake3.workspace = true tokio.workspace = true diff --git a/iris-mpc-upgrade-hawk/src/genesis/graph_checkpoint.rs b/iris-mpc-upgrade-hawk/src/genesis/graph_checkpoint.rs index ea38501f67..83f5cb1e31 100644 --- a/iris-mpc-upgrade-hawk/src/genesis/graph_checkpoint.rs +++ b/iris-mpc-upgrade-hawk/src/genesis/graph_checkpoint.rs @@ -1,5 +1,5 @@ -use aws_sdk_s3::Client as S3Client; use eyre::{bail, Result}; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_common::SerialId; use iris_mpc_cpu::execution::hawk_main::{BothEyes, GraphRef, HawkOps}; use iris_mpc_cpu::genesis::state_accessor::set_last_indexed_iris_id; @@ -19,7 +19,7 @@ pub async fn upload_and_sync_genesis_checkpoint( checkpoint_bucket: &str, party_id: usize, imem_graph_stores: &Arc>, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, last_indexed_id: u32, max_modification_indexed_id: i64, is_archival: bool, diff --git a/iris-mpc-upgrade-hawk/src/genesis/mod.rs b/iris-mpc-upgrade-hawk/src/genesis/mod.rs index d39ff85b0d..f3bd0e346f 100644 --- a/iris-mpc-upgrade-hawk/src/genesis/mod.rs +++ b/iris-mpc-upgrade-hawk/src/genesis/mod.rs @@ -5,16 +5,13 @@ use ampc_server_utils::{ start_coordination_server_with_extra_routes, wait_for_others_ready, wait_for_others_unready, BatchSyncSharedState, TaskMonitor, }; -use aws_config::retry::RetryConfig; +use aws_sdk_rds::config::Region; use aws_sdk_rds::Client as RDSClient; -use aws_sdk_s3::{ - config::{Builder as S3ConfigBuilder, Region}, - Client as S3Client, -}; use axum::{routing::get, Router}; use chrono::Utc; use eyre::{bail, eyre, Report, Result}; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_common::{ config::{CommonConfig, Config, ENV_PROD, ENV_STAGE}, helpers::{smpc_request, sync::Modification}, @@ -262,7 +259,7 @@ async fn exec_setup( ExecutionContextInfo, Arc, TaskMonitor, - S3Client, + ObjectStoreClient, RDSClient, BothEyes, BothEyes>, @@ -559,7 +556,7 @@ async fn exec_delta( config: &Config, ctx: &ExecutionContextInfo, graph_store: Arc>>, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, imem_graph_stores: &Arc>, mut hawk_handle: GenesisHawkHandle, tx_results: &Sender, @@ -702,7 +699,7 @@ async fn exec_delta( /// # Arguments /// /// * `ctx` - Execution context information. -/// * `s3_client` - AWS S3 client for checkpoint uploads. +/// * `s3_client` - Object-store client for checkpoint uploads. /// * `registries` - Per-eye VectorId registries used by the batch generator. /// * `worker_pools` - Per-eye worker pools that own iris data and cache queries. /// * `imem_graph_stores` - In-memory graph stores for checkpoints. @@ -714,7 +711,7 @@ async fn exec_delta( #[allow(clippy::too_many_arguments)] async fn exec_indexation( ctx: &ExecutionContextInfo, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, registries: &BothEyes, worker_pools: &BothEyes>, imem_graph_stores: &Arc>, @@ -1049,19 +1046,19 @@ async fn get_service_clients( config: &Config, ) -> Result< ( - (S3Client, S3Client, RDSClient), + (ObjectStoreClient, ObjectStoreClient, RDSClient), (IrisStore, (IrisStore, GraphPg>)), ), Report, > { - /// Returns S3 clients and an RDS client. + /// Returns object-store clients and an RDS client. /// - /// Two S3 clients are constructed so the graph-checkpoint bucket can - /// live in a different AWS region than the iris-snapshot bucket. - async fn get_aws_clients(config: &Config) -> Result<(S3Client, S3Client, RDSClient)> { + /// Two object-store clients are constructed so the graph-checkpoint store + /// can live in a different region than the iris-snapshot store. + async fn get_aws_clients( + config: &Config, + ) -> Result<(ObjectStoreClient, ObjectStoreClient, RDSClient)> { let force_path_style = config.environment != ENV_PROD && config.environment != ENV_STAGE; - let retry_config = RetryConfig::standard().with_max_attempts(5); - let config_region = config.aws.clone().and_then(|aws| aws.region); let region_name = config_region .clone() @@ -1085,12 +1082,9 @@ async fn get_service_clients( sdk_config.endpoint_url(), ); - // S3 client for general AWS operations (iris snapshots, deletions) - let s3_config = S3ConfigBuilder::from(&sdk_config) - .force_path_style(force_path_style) - .retry_config(retry_config.clone()) - .build(); - let aws_s3_client = S3Client::from_conf(s3_config); + // Object-store client for general operations (iris snapshots, deletions). + let aws_s3_client = ObjectStoreClient::new(Some(region_name.clone()), force_path_style) + .with_aws_sdk_config(&sdk_config); // RDS client using general AWS configuration tracing::info!( @@ -1102,20 +1096,15 @@ async fn get_service_clients( // S3 client for graph checkpoint operations (may be in a different region) let checkpoint_region_name = config.graph_checkpoint_bucket_region.clone(); - let checkpoint_region = Region::new(checkpoint_region_name.clone()); - tracing::info!( "Checkpoint S3 client: region={}, endpoint={:?}", checkpoint_region_name, sdk_config.endpoint_url(), ); - let checkpoint_s3_config = S3ConfigBuilder::from(&sdk_config) - .region(checkpoint_region) - .force_path_style(force_path_style) - .retry_config(retry_config.clone()) - .build(); - let checkpoint_s3_client = S3Client::from_conf(checkpoint_s3_config); + let checkpoint_s3_client = + ObjectStoreClient::new(Some(checkpoint_region_name), force_path_style) + .with_aws_sdk_config(&sdk_config); Ok((aws_s3_client, checkpoint_s3_client, rds_client)) } @@ -1383,7 +1372,7 @@ async fn init_graph_from_stores( iris_store: &IrisStore, hawk_args: HawkArgs, hawk_networking: Box, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, shutdown_handler: Arc, max_indexation_id: usize, checkpoint: Option, diff --git a/iris-mpc-upgrade-hawk/tests/utils/genesis_runner.rs b/iris-mpc-upgrade-hawk/tests/utils/genesis_runner.rs index d97b27cf25..b300b442e0 100644 --- a/iris-mpc-upgrade-hawk/tests/utils/genesis_runner.rs +++ b/iris-mpc-upgrade-hawk/tests/utils/genesis_runner.rs @@ -76,7 +76,7 @@ pub async fn base_genesis_e2e_init(configs: &HawkConfigs, s3_deletion_ids: Vec = download_graph_checkpoint( - &aws_clients.checkpoint_s3_client, + &aws_clients.checkpoint_object_store_client, &config.graph_checkpoint_bucket_name, &checkpoint_state, ) @@ -171,7 +171,7 @@ impl MpcNodes { checkpoint.s3_key ); if let Err(e) = - delete_graph(&aws_clients.s3_client, bucket, &checkpoint.s3_key).await + delete_graph(&aws_clients.object_store_client, bucket, &checkpoint.s3_key).await { tracing::warn!( "Failed to delete S3 checkpoint {}: {:?}", diff --git a/iris-mpc-upgrade-hawk/tests/utils/s3_deletions.rs b/iris-mpc-upgrade-hawk/tests/utils/s3_deletions.rs index 676a579f73..24dc3f5fec 100644 --- a/iris-mpc-upgrade-hawk/tests/utils/s3_deletions.rs +++ b/iris-mpc-upgrade-hawk/tests/utils/s3_deletions.rs @@ -1,13 +1,13 @@ -use aws_sdk_s3::{primitives::ByteStream as S3_ByteStream, Client as S3_Client}; use eyre::{eyre, Result}; use iris_mpc::services::aws::clients::AwsClients; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use iris_mpc_common::{config::Config, SerialId}; use serde::Serialize; /// Uploads to an AWS S3 bucket a set of serial identifiers marked as deleted. pub async fn upload_iris_deletions( data: &Vec, - s3: &S3_Client, + s3: &ObjectStoreClient, environment: &str, ) -> Result<()> { // Set bucket/key based on environment. @@ -20,20 +20,15 @@ pub async fn upload_iris_deletions( ); // Set body of payload to be persisted. - let body = S3_ByteStream::from( - serde_json::to_string(&IrisDeletionsForS3 { - deleted_serial_ids: data.to_owned(), - }) - .unwrap() - .into_bytes(), - ); + let body = serde_json::to_string(&IrisDeletionsForS3 { + deleted_serial_ids: data.to_owned(), + }) + .unwrap() + .into_bytes(); // Upload payload. - s3.put_object() - .bucket(&s3_bucket) - .key(&s3_key) - .body(body) - .send() + s3.store(&s3_bucket)? + .put(&path(&s3_key)?, body.into()) .await .map_err(|err| { tracing::error!("Failed to upload file to S3: {}", err); diff --git a/iris-mpc-utils/Cargo.toml b/iris-mpc-utils/Cargo.toml index d2846b48f7..d0abe2d10d 100644 --- a/iris-mpc-utils/Cargo.toml +++ b/iris-mpc-utils/Cargo.toml @@ -12,8 +12,8 @@ aws-config.workspace = true aws-sdk-rds.workspace = true aws-sdk-sns.workspace = true aws-sdk-sqs.workspace = true -aws-sdk-s3.workspace = true aws-sdk-secretsmanager.workspace = true +object_store.workspace = true base64.workspace = true bincode.workspace = true clap.workspace = true diff --git a/iris-mpc-utils/src/aws/client.rs b/iris-mpc-utils/src/aws/client.rs index 85d4307e5b..99c8a008c1 100644 --- a/iris-mpc-utils/src/aws/client.rs +++ b/iris-mpc-utils/src/aws/client.rs @@ -2,10 +2,6 @@ use std::collections::HashSet; use async_from::AsyncFrom; use async_stream::stream; -use aws_sdk_s3::{ - primitives::{ByteStream, SdkBody}, - Client as S3Client, -}; use aws_sdk_secretsmanager::Client as SecretsManagerClient; use aws_sdk_sns::Client as SNSClient; use aws_sdk_sqs::Client as SQSClient; @@ -13,6 +9,7 @@ use futures::stream::Stream; use serde_json; use iris_mpc_common::helpers::smpc_response::create_sns_message_attributes; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use super::{ config::AwsClientConfig, @@ -31,8 +28,8 @@ pub struct AwsClient { /// Encryption public key set ... one per MPC node. public_keyset: Option, - /// Client for Amazon Simple Storage Service. - s3: S3Client, + /// Client for bucket/object storage. + object_store: ObjectStoreClient, /// Client for AWS Secrets Manager. #[allow(dead_code)] @@ -71,7 +68,7 @@ impl AwsClient { Self { config: config.to_owned(), public_keyset: None, - s3: S3Client::from(&config), + object_store: config.object_store_client(), secrets_manager: SecretsManagerClient::from(&config), sqs: SQSClient::from(&config), sns: SNSClient::from(&config), @@ -80,21 +77,24 @@ impl AwsClient { } impl AwsClient { - /// Enqueues data to an S3 bucket. + /// Writes data to the configured object store. pub async fn s3_put_object(&self, s3_obj_info: &S3ObjectInfo) -> Result<(), AwsClientError> { - tracing::debug!("AWS-S3: putting object -> {}", s3_obj_info); - self.s3 - .put_object() - .bucket(s3_obj_info.bucket()) - .key(s3_obj_info.key()) - .body(ByteStream::new(SdkBody::from(s3_obj_info.body()))) - .send() - .await - .map(|_| ()) - .map_err(|e| { - tracing::error!("AWS-S3 upload error: {}", e); - AwsClientError::S3UploadError(s3_obj_info.key().to_string(), e.to_string()) - }) + tracing::debug!("Object store: putting object -> {}", s3_obj_info); + let result = async { + let store = self.object_store.store(s3_obj_info.bucket())?; + store + .put( + &path(s3_obj_info.key())?, + s3_obj_info.body().to_vec().into(), + ) + .await?; + Ok::<_, object_store::Error>(()) + } + .await; + result.map(|_| ()).map_err(|e| { + tracing::error!("Object-store upload error: {}", e); + AwsClientError::S3UploadError(s3_obj_info.key().to_string(), e.to_string()) + }) } /// Downloads & assigns encryption keys. @@ -499,7 +499,10 @@ mod tests { impl AwsClient { fn assert_instance(&self) { - assert!(self.s3.config().region().is_some()); + assert!(self + .object_store + .store(self.config.s3_request_bucket_name()) + .is_ok()); assert!(self.secrets_manager.config().region().is_some()); assert!(self.sns.config().region().is_some()); assert!(self.sqs.config().region().is_some()); diff --git a/iris-mpc-utils/src/aws/config.rs b/iris-mpc-utils/src/aws/config.rs index 83de9f8bee..91c353b252 100644 --- a/iris-mpc-utils/src/aws/config.rs +++ b/iris-mpc-utils/src/aws/config.rs @@ -2,12 +2,13 @@ use std::{env, path::Path, time::Duration}; use async_from::AsyncFrom; use aws_config::{retry::RetryConfig, timeout::TimeoutConfig, SdkConfig}; -use aws_sdk_s3::{config::Builder as S3ConfigBuilder, Client as S3Client}; use aws_sdk_secretsmanager::Client as SecretsManagerClient; use aws_sdk_sns::Client as SNSClient; use aws_sdk_sqs::{config::Builder, config::Region, Client as SQSClient}; +use iris_mpc_common::config::resolve_iris_deletions_store_location; use iris_mpc_common::config::{ENV_PROD, ENV_STAGE}; +use iris_mpc_common::object_store::ObjectStoreClient; use crate::client::AwsOptions; @@ -26,6 +27,9 @@ pub struct AwsClientConfig { /// S3: request ingress queue URL. s3_request_bucket_name: String, + /// Object-store location for the Iris deletion snapshot. + iris_deletions_store_location: String, + /// SDK: associated AWS SDK configuration. sdk: SdkConfig, @@ -42,10 +46,12 @@ pub struct AwsClientConfig { #[async_from::async_trait] impl AsyncFrom for AwsClientConfig { async fn async_from(opts: AwsOptions) -> Self { - AwsClientConfig::new( + let iris_deletions_store_location = opts.iris_deletions_store_location(); + AwsClientConfig::new_with_iris_deletions_store_location( opts.environment().to_owned(), opts.public_key_base_url().to_owned(), opts.s3_request_bucket_name().to_owned(), + iris_deletions_store_location, opts.sns_request_topic_arn().to_owned(), opts.sqs_long_poll_wait_time().to_owned(), opts.sqs_response_queue_urls().to_owned(), @@ -67,6 +73,10 @@ impl AwsClientConfig { &self.s3_request_bucket_name } + pub(crate) fn iris_deletions_store_location(&self) -> &str { + &self.iris_deletions_store_location + } + pub(super) fn sdk(&self) -> &SdkConfig { &self.sdk } @@ -90,11 +100,35 @@ impl AwsClientConfig { sns_request_topic_arn: String, sqs_long_poll_wait_time: usize, sqs_response_queue_urls: Vec, + ) -> Self { + let iris_deletions_store_location = + resolve_iris_deletions_store_location(&environment, None); + Self::new_with_iris_deletions_store_location( + environment, + public_key_base_url, + s3_request_bucket_name, + iris_deletions_store_location, + sns_request_topic_arn, + sqs_long_poll_wait_time, + sqs_response_queue_urls, + ) + .await + } + + pub async fn new_with_iris_deletions_store_location( + environment: String, + public_key_base_url: String, + s3_request_bucket_name: String, + iris_deletions_store_location: String, + sns_request_topic_arn: String, + sqs_long_poll_wait_time: usize, + sqs_response_queue_urls: Vec, ) -> Self { Self { environment, public_key_base_url, s3_request_bucket_name, + iris_deletions_store_location, sdk: get_sdk_config().await, sns_request_topic_arn, sqs_long_poll_wait_time, @@ -149,15 +183,15 @@ impl From<&AwsClientConfig> for SNSClient { } } -impl From<&AwsClientConfig> for S3Client { - fn from(config: &AwsClientConfig) -> Self { - let force_path_style = - config.environment() != ENV_PROD && config.environment() != ENV_STAGE; - let config_builder = S3ConfigBuilder::from(config.sdk()) - .retry_config(RetryConfig::standard().with_max_attempts(5)) - .force_path_style(force_path_style); - - S3Client::from_conf(config_builder.build()) +impl AwsClientConfig { + pub(crate) fn object_store_client(&self) -> ObjectStoreClient { + let force_path_style = self.environment() != ENV_PROD && self.environment() != ENV_STAGE; + let region = self + .sdk() + .region() + .map(ToString::to_string) + .unwrap_or_else(|| AWS_DEFAULT_REGION.to_owned()); + ObjectStoreClient::new(Some(region), force_path_style).with_aws_sdk_config(self.sdk()) } } diff --git a/iris-mpc-utils/src/aws/ops.rs b/iris-mpc-utils/src/aws/ops.rs index 41fe1bf413..bb4c112947 100644 --- a/iris-mpc-utils/src/aws/ops.rs +++ b/iris-mpc-utils/src/aws/ops.rs @@ -24,27 +24,30 @@ impl AwsClient { }; let environment = self.config().environment(); - let s3_bucket = format!("wf-smpcv2-{}-sync-protocol", environment); - let s3_key = format!("{}_deleted_serial_ids.json", environment); - let s3_obj = S3ObjectInfo::new(&s3_bucket, &s3_key, &data); - self.s3_put_object(&s3_obj) + let object_key = format!("{}_deleted_serial_ids.json", environment); + let object = S3ObjectInfo::new( + self.config().iris_deletions_store_location(), + &object_key, + &data, + ); + self.s3_put_object(&object) .await .map_err(|e| AwsClientError::IrisDeletionsUploadError(e.to_string())) } - // Uploads JSON encoded Iris shares to AWS S3 bucket. + // Uploads JSON-encoded Iris shares to the configured object store. pub async fn s3_upload_iris_shares( &self, signup_id: &uuid::Uuid, shares: &BothEyes<[GaloisRingSharedIrisForUpload; N_PARTIES]>, ) -> Result { - // Set AWS-S3 JSON compatible shares. + // Build JSON-compatible shares for the existing wire format. let shares = create_iris_code_shares_s3( &create_iris_code_shares(signup_id, shares), &self.public_keyset(), ); - // Upload to AWS-S3. + // Upload to object storage. let s3_obj_info = S3ObjectInfo::new( self.config().s3_request_bucket_name(), &signup_id.to_string(), diff --git a/iris-mpc-utils/src/aws/types.rs b/iris-mpc-utils/src/aws/types.rs index 18d7f66357..bc336f4c49 100644 --- a/iris-mpc-utils/src/aws/types.rs +++ b/iris-mpc-utils/src/aws/types.rs @@ -8,7 +8,7 @@ use crate::client::{Request, RequestPayload}; const ENROLLMENT_REQUEST_TYPE: &str = "enrollment"; -// Helper type encapsulating AWS-S3 object information. +// Compatibility type encapsulating object-store location and data. #[derive(Debug)] pub struct S3ObjectInfo { // S3 object data. diff --git a/iris-mpc-utils/src/client/options/types.rs b/iris-mpc-utils/src/client/options/types.rs index bdefe7472c..8f73985229 100644 --- a/iris-mpc-utils/src/client/options/types.rs +++ b/iris-mpc-utils/src/client/options/types.rs @@ -6,6 +6,7 @@ use serde::{ Deserialize, Serialize, }; +use iris_mpc_common::config::resolve_iris_deletions_store_location; use iris_mpc_common::SerialId; use iris_mpc_cpu::utils::serialization::iris_ndjson::IrisSelection; use uuid::Uuid; @@ -24,6 +25,11 @@ pub struct AwsOptions { /// S3: request ingress queue URL. s3_request_bucket_name: String, + /// Object-store location for the Iris deletion snapshot. If omitted, the + /// historical environment-specific S3 bucket is used. + #[serde(default)] + iris_deletions_store_location: Option, + /// SNS: system request ingress queue topic. sns_request_topic_arn: String, @@ -50,6 +56,13 @@ impl AwsOptions { &self.s3_request_bucket_name } + pub fn iris_deletions_store_location(&self) -> String { + resolve_iris_deletions_store_location( + &self.environment, + self.iris_deletions_store_location.as_deref(), + ) + } + pub fn sns_request_topic_arn(&self) -> &str { &self.sns_request_topic_arn } @@ -755,6 +768,37 @@ fn simple_into_iter( mod tests { use super::*; + fn aws_options(location: Option<&str>) -> AwsOptions { + let location = location + .map(|value| format!("iris_deletions_store_location = \"{value}\"")) + .unwrap_or_default(); + toml::from_str(&format!( + r#" + environment = "dev" + public_key_base_url = "http://localhost/public-keys" + s3_request_bucket_name = "request-shares" + {location} + sns_request_topic_arn = "request-topic" + sqs_long_poll_wait_time = 5 + sqs_response_queue_urls = [] + sqs_wait_time_seconds = 5 + "# + )) + .unwrap() + } + + #[test] + fn aws_options_deletion_store_preserves_legacy_default_and_allows_override() { + assert_eq!( + aws_options(None).iris_deletions_store_location(), + "wf-smpcv2-dev-sync-protocol" + ); + assert_eq!( + aws_options(Some("az://deletions/snapshots")).iris_deletions_store_location(), + "az://deletions/snapshots" + ); + } + #[test] fn test_request_batch_options_simple_roundtrip() { let toml_str = r#" diff --git a/iris-mpc/Cargo.toml b/iris-mpc/Cargo.toml index db8c2bc46a..679c47d71c 100644 --- a/iris-mpc/Cargo.toml +++ b/iris-mpc/Cargo.toml @@ -11,8 +11,8 @@ repository.workspace = true aws-config.workspace = true aws-sdk-sns.workspace = true aws-sdk-sqs.workspace = true -aws-sdk-s3.workspace = true aws-sdk-secretsmanager.workspace = true +object_store.workspace = true axum.workspace = true tokio.workspace = true tokio-util.workspace = true diff --git a/iris-mpc/src/client/e2e.rs b/iris-mpc/src/client/e2e.rs index b019ec9f12..1bf8c466ca 100644 --- a/iris-mpc/src/client/e2e.rs +++ b/iris-mpc/src/client/e2e.rs @@ -3,7 +3,6 @@ use crate::client::iris_data::{ generate_party_shares, read_iris_data_from_file, IrisCodePartyShares, }; use aws_config::retry::RetryConfig; -use aws_sdk_s3::Client as S3Client; use aws_sdk_sns::{config::Region, types::MessageAttributeValue, Client as SnsClient}; use aws_sdk_sqs::Client as SqsClient; use base64::{engine::general_purpose, Engine}; @@ -21,6 +20,7 @@ use iris_mpc_common::helpers::{ smpc_response::{create_message_type_attribute_map, UniquenessResult}, sqs_s3_helper::upload_file_to_s3, }; +use iris_mpc_common::object_store::ObjectStoreClient; use rand::{rngs::StdRng, Rng, SeedableRng}; use serde_json::to_string; use sodiumoxide::crypto::{box_::PublicKey, sealedbox}; @@ -110,7 +110,7 @@ pub struct E2EClient { populate_file_data_limit: Option, // AWS clients - s3_client: S3Client, + s3_client: ObjectStoreClient, sns_client: Arc, sqs_client: SqsClient, @@ -136,15 +136,17 @@ impl E2EClient { .load() .await; - let mut s3_config_builder = aws_sdk_s3::config::Builder::from(&shared_config); let mut sns_config_builder = aws_sdk_sns::config::Builder::from(&shared_config); + let mut s3_client = + ObjectStoreClient::new(Some(opts.region.clone()), opts.endpoint_url.is_some()) + .with_aws_sdk_config(&shared_config); if let Some(endpoint_url) = opts.endpoint_url.as_ref() { - s3_config_builder = s3_config_builder.endpoint_url(endpoint_url); - s3_config_builder = s3_config_builder.force_path_style(true); + s3_client = s3_client + .with_option("aws_endpoint", endpoint_url) + .with_option("aws_allow_http", endpoint_url.starts_with("http://")); sns_config_builder = sns_config_builder.endpoint_url(endpoint_url); } - let s3_client = S3Client::from_conf(s3_config_builder.build()); let sns_client = Arc::new(SnsClient::from_conf(sns_config_builder.build())); let sqs_client = SqsClient::new(&shared_config); diff --git a/iris-mpc/src/server/mod.rs b/iris-mpc/src/server/mod.rs index 56cb7e0397..a22f9c33af 100644 --- a/iris-mpc/src/server/mod.rs +++ b/iris-mpc/src/server/mod.rs @@ -1,8 +1,8 @@ use crate::services::aws::clients::AwsClients; use crate::services::processors::batch::{receive_batch_stream, spawn_db_backed_ingest_task}; use crate::services::processors::job::{process_job_result, BatchTimings}; -use aws_sdk_s3::Client; use aws_sdk_sns::types::MessageAttributeValue; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_cpu::checkpoint_protocol::{restart_from_checkpoint, sidecar_main, RestartOutcome}; use iris_mpc_cpu::graph_checkpoint::sync_graph_mutations; use iris_mpc_cpu::hnsw::GraphMem; @@ -250,7 +250,7 @@ pub async fn server_main(config: Config) -> Result<()> { let mut hawk_actor = init_hawk_actor( &config, - &aws_clients.checkpoint_s3_client, + &aws_clients.checkpoint_object_store_client, &config.graph_checkpoint_bucket_name, &iris_store, &graph_store, @@ -363,7 +363,7 @@ pub async fn server_main(config: Config) -> Result<()> { sidecar_main( sc_config, &graph_store, - &aws_clients.checkpoint_s3_client, + &aws_clients.checkpoint_object_store_client, &mut network_handle, sidecar_ct, ) @@ -684,7 +684,7 @@ fn build_hawk_args(config: &Config) -> Result<(HawkArgs, Vec, Vec>, @@ -911,7 +911,7 @@ async fn run_main_server_loop( party_id, aws_clients.sqs_client.clone(), aws_clients.sns_client.clone(), - aws_clients.s3_client.clone(), + aws_clients.object_store_client.clone(), config.clone(), shares_encryption_key_pair.clone(), shutdown_handler.clone(), @@ -1058,7 +1058,7 @@ async fn run_main_server_loop( let dur_secs = pprof_start.elapsed().as_secs(); let ts = Utc::now().format("%Y-%m-%dT%H-%M-%SZ"); let party = format!("party{}", config.party_id); - let s3 = aws_clients.s3_client.clone(); + let s3 = aws_clients.object_store_client.clone(); let bucket = config.pprof_s3_bucket.clone(); let prefix = config.pprof_prefix.clone(); let run_id = config diff --git a/iris-mpc/src/services/aws/clients.rs b/iris-mpc/src/services/aws/clients.rs index b7d863bb18..9ad99de842 100644 --- a/iris-mpc/src/services/aws/clients.rs +++ b/iris-mpc/src/services/aws/clients.rs @@ -1,22 +1,20 @@ -use crate::services::aws::s3::create_s3_client; use crate::services::aws::sns::create_sns_client; use crate::services::aws::sqs::create_sqs_client; -use aws_sdk_s3::config::Region as S3Region; -use aws_sdk_s3::Client as S3Client; use aws_sdk_secretsmanager::Client as SecretsManagerClient; use aws_sdk_sns::Client as SNSClient; use aws_sdk_sqs::{config::Region, Client as SQSClient}; use eyre::Result; use iris_mpc_common::config::{Config, ENV_PROD, ENV_STAGE}; +use iris_mpc_common::object_store::ObjectStoreClient; const DEFAULT_REGION: &str = "eu-north-1"; pub struct AwsClients { pub sqs_client: SQSClient, pub sns_client: SNSClient, - pub s3_client: S3Client, - // used to obtain graph checkpoints. this bucket could be in a different region - pub checkpoint_s3_client: S3Client, + pub object_store_client: ObjectStoreClient, + // Graph checkpoints may use a store in a different region. + pub checkpoint_object_store_client: ObjectStoreClient, pub secrets_manager_client: SecretsManagerClient, } @@ -29,25 +27,28 @@ impl AwsClients { .and_then(|aws| aws.region) .unwrap_or_else(|| DEFAULT_REGION.to_owned()); - let region_provider = Region::new(region); + let region_provider = Region::new(region.clone()); let shared_config = aws_config::from_env().region(region_provider).load().await; let force_path_style = config.environment != ENV_PROD && config.environment != ENV_STAGE; let sns_client = create_sns_client(&shared_config, config.sns_retry_max_attempts); let sqs_client = create_sqs_client(&shared_config, config.sqs_long_poll_wait_time); - let s3_client = create_s3_client(&shared_config, force_path_style, None); + let object_store_client = ObjectStoreClient::new(Some(region), force_path_style) + .with_aws_sdk_config(&shared_config); let secrets_manager_client = SecretsManagerClient::new(&shared_config); - let checkpoint_region = S3Region::new(config.graph_checkpoint_bucket_region.clone()); - let checkpoint_s3_client = - create_s3_client(&shared_config, force_path_style, Some(checkpoint_region)); + let checkpoint_object_store_client = ObjectStoreClient::new( + Some(config.graph_checkpoint_bucket_region.clone()), + force_path_style, + ) + .with_aws_sdk_config(&shared_config); Ok(Self { sqs_client, sns_client, - s3_client, - checkpoint_s3_client, + object_store_client, + checkpoint_object_store_client, secrets_manager_client, }) } @@ -59,8 +60,8 @@ impl Clone for AwsClients { Self { sqs_client: self.sqs_client.clone(), sns_client: self.sns_client.clone(), - s3_client: self.s3_client.clone(), - checkpoint_s3_client: self.checkpoint_s3_client.clone(), + object_store_client: self.object_store_client.clone(), + checkpoint_object_store_client: self.checkpoint_object_store_client.clone(), secrets_manager_client: self.secrets_manager_client.clone(), } } diff --git a/iris-mpc/src/services/aws/mod.rs b/iris-mpc/src/services/aws/mod.rs index 0c4f43d1b5..cf3fc73e7f 100644 --- a/iris-mpc/src/services/aws/mod.rs +++ b/iris-mpc/src/services/aws/mod.rs @@ -1,4 +1,3 @@ pub mod clients; -pub mod s3; pub mod sns; pub mod sqs; diff --git a/iris-mpc/src/services/aws/s3.rs b/iris-mpc/src/services/aws/s3.rs deleted file mode 100644 index 58c31e300e..0000000000 --- a/iris-mpc/src/services/aws/s3.rs +++ /dev/null @@ -1,28 +0,0 @@ -use aws_config::retry::RetryConfig; -use aws_sdk_s3::{config::Builder as S3ConfigBuilder, config::Region, Client as S3Client}; - -/// Creates an S3 client with retry configuration -/// -/// # Arguments -/// * `shared_config` - Base AWS SDK configuration to build from -/// * `force_path_style` - Whether to use path-style S3 URLs -/// * `region_override` - Optional region to override the shared config's region -pub fn create_s3_client( - shared_config: &aws_config::SdkConfig, - force_path_style: bool, - region_override: Option, -) -> S3Client { - let retry_config = RetryConfig::standard().with_max_attempts(5); - - let mut builder = S3ConfigBuilder::from(shared_config) - .force_path_style(force_path_style) - .retry_config(retry_config.clone()); - - if let Some(region) = region_override { - builder = builder.region(region); - } - - let s3_config = builder.build(); - - S3Client::from_conf(s3_config) -} diff --git a/iris-mpc/src/services/processors/batch.rs b/iris-mpc/src/services/processors/batch.rs index 2112de9252..1e6cef8c19 100644 --- a/iris-mpc/src/services/processors/batch.rs +++ b/iris-mpc/src/services/processors/batch.rs @@ -6,7 +6,6 @@ use ampc_server_utils::{ get_approximate_number_of_messages, get_batch_sync_states, BatchSyncResult, BatchSyncSharedState, BatchSyncState, TaskMonitor, }; -use aws_sdk_s3::Client as S3Client; use aws_sdk_sns::types::MessageAttributeValue; use aws_sdk_sns::Client as SNSClient; use aws_sdk_sqs::Client; @@ -38,6 +37,7 @@ use iris_mpc_common::helpers::smpc_response::{ use iris_mpc_common::helpers::sync::Modification; use iris_mpc_common::helpers::sync::ModificationKey::{RequestId, RequestSerialId}; use iris_mpc_common::job::{BatchMetadata, BatchQuery, GaloisSharesBothSides}; +use iris_mpc_common::object_store::ObjectStoreClient; use iris_mpc_store::{normalize_sns_sequence_number, IngestedRequest, Store}; use rand::Rng; use std::collections::HashMap; @@ -313,7 +313,7 @@ pub fn receive_batch_stream( party_id: usize, client: Client, sns_client: SNSClient, - s3_client: S3Client, + s3_client: ObjectStoreClient, config: Config, shares_encryption_key_pairs: SharesEncryptionKeyPairs, shutdown_handler: Arc, @@ -394,7 +394,7 @@ async fn receive_batch( party_id: usize, client: &Client, sns_client: &SNSClient, - s3_client: &S3Client, + s3_client: &ObjectStoreClient, config: &Config, shares_encryption_key_pairs: SharesEncryptionKeyPairs, shutdown_handler: &ShutdownHandler, @@ -436,7 +436,7 @@ pub struct BatchProcessor<'a> { party_id: usize, client: &'a Client, sns_client: &'a SNSClient, - s3_client: &'a S3Client, + s3_client: &'a ObjectStoreClient, config: &'a Config, shares_encryption_key_pairs: SharesEncryptionKeyPairs, shutdown_handler: &'a ShutdownHandler, @@ -462,7 +462,7 @@ impl<'a> BatchProcessor<'a> { party_id: usize, client: &'a Client, sns_client: &'a SNSClient, - s3_client: &'a S3Client, + s3_client: &'a ObjectStoreClient, config: &'a Config, shares_encryption_key_pairs: SharesEncryptionKeyPairs, shutdown_handler: &'a ShutdownHandler, diff --git a/iris-mpc/src/services/processors/mod.rs b/iris-mpc/src/services/processors/mod.rs index 852327ae01..fc4fea274c 100644 --- a/iris-mpc/src/services/processors/mod.rs +++ b/iris-mpc/src/services/processors/mod.rs @@ -3,7 +3,6 @@ pub mod job; pub mod modifications_sync; pub mod result_message; -use aws_sdk_s3::Client as S3Client; use eyre::Result; use eyre::{Context, Report}; use iris_mpc_common::galois_engine::degree4::{ @@ -14,6 +13,7 @@ use iris_mpc_common::helpers::key_pair::SharesEncryptionKeyPairs; use iris_mpc_common::helpers::smpc_request::{ decrypt_iris_share, get_iris_data_by_party_id, validate_iris_share, ReceiveRequestError, }; +use iris_mpc_common::object_store::ObjectStoreClient; use std::sync::Arc; use tokio::sync::Semaphore; use tokio::task::{spawn_blocking, JoinHandle}; @@ -49,7 +49,7 @@ pub fn get_iris_shares_parse_task( party_id: usize, shares_encryption_key_pairs: SharesEncryptionKeyPairs, semaphore: Arc, - s3_client_arc: S3Client, + s3_client_arc: ObjectStoreClient, bucket_name: String, s3_key: String, ) -> Result, ReceiveRequestError> { diff --git a/iris-mpc/src/services/processors/modifications_sync.rs b/iris-mpc/src/services/processors/modifications_sync.rs index ef0933ef08..3b869d2c87 100644 --- a/iris-mpc/src/services/processors/modifications_sync.rs +++ b/iris-mpc/src/services/processors/modifications_sync.rs @@ -92,7 +92,7 @@ pub async fn sync_modifications( config.party_id, shares_encryption_key_pair.clone(), Arc::clone(&semaphore), - aws_clients.s3_client.clone(), + aws_clients.object_store_client.clone(), config.shares_bucket_name.clone(), s3_url, )? diff --git a/iris-mpc/tests/utils/cpu_node.rs b/iris-mpc/tests/utils/cpu_node.rs index dfc1650170..86c66dd90e 100644 --- a/iris-mpc/tests/utils/cpu_node.rs +++ b/iris-mpc/tests/utils/cpu_node.rs @@ -3,6 +3,8 @@ use crate::utils::CpuNodeConfig; use super::CpuConfigs; use eyre::eyre; +use futures::TryStreamExt; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use iris_mpc_common::{ postgres::{AccessMode, PostgresClient}, VectorId, MASK_CODE_LENGTH, @@ -137,11 +139,11 @@ impl DbStores { pub struct CpuNode { pub store: DbStores, pub config: CpuNodeConfig, - pub s3: aws_sdk_s3::Client, + pub s3: ObjectStoreClient, } impl CpuNode { - pub async fn new(config: CpuNodeConfig, s3: aws_sdk_s3::Client) -> eyre::Result { + pub async fn new(config: CpuNodeConfig, s3: ObjectStoreClient) -> eyre::Result { let stores = DbStores::new(&config).await?; Ok(Self { store: stores, @@ -188,13 +190,9 @@ impl CpuNode { pub async fn delete_latest_checkpoint(&self) -> eyre::Result<()> { if let Some(s3_key) = self.store.delete_latest_checkpoint().await? { // Best-effort cleanup: ignore errors if the S3 object is already gone - let _ = self - .s3 - .delete_object() - .bucket(&self.config.checkpoint_bucket) - .key(&s3_key) - .send() - .await; + if let Ok(store) = self.s3.store(&self.config.checkpoint_bucket) { + let _ = store.delete(&path(&s3_key)?).await; + } } Ok(()) } @@ -207,10 +205,8 @@ impl CpuNode { let checkpoints = self.store.graph.get_genesis_graph_checkpoints().await?; for row in checkpoints { self.s3 - .head_object() - .bucket(bucket) - .key(&row.s3_key) - .send() + .store(bucket)? + .head(&path(&row.s3_key)?) .await .map_err(|e| eyre!("S3 object {} missing: {}", row.s3_key, e))?; } @@ -222,29 +218,13 @@ impl CpuNode { /// Returns an empty vec when the bucket is empty. Intended for exec_assert /// checks that verify the pruning pass reduced the number of stored checkpoints. pub async fn list_s3_keys(&self, bucket: &str) -> eyre::Result> { - let mut keys = Vec::new(); - let mut continuation_token: Option = None; - loop { - let mut req = self.s3.list_objects_v2().bucket(bucket); - if let Some(ref token) = continuation_token { - req = req.continuation_token(token); - } - let resp = req - .send() - .await - .map_err(|e| eyre!("failed to list S3 objects in bucket {bucket}: {e}"))?; - for obj in resp.contents() { - if let Some(key) = obj.key() { - keys.push(key.to_string()); - } - } - if resp.is_truncated().unwrap_or(false) { - continuation_token = resp.next_continuation_token().map(|s| s.to_string()); - } else { - break; - } - } - Ok(keys) + self.s3 + .store(bucket)? + .list(None) + .map_ok(|meta| meta.location.to_string()) + .try_collect() + .await + .map_err(|e| eyre!("failed to list objects in {bucket}: {e}")) } /// Seed a genesis (seed) checkpoint for this party. @@ -477,7 +457,7 @@ impl CpuNode { pub struct CpuNodes(pub [CpuNode; 3]); impl CpuNodes { - pub async fn new(configs: &CpuConfigs, s3: aws_sdk_s3::Client) -> eyre::Result { + pub async fn new(configs: &CpuConfigs, s3: ObjectStoreClient) -> eyre::Result { // Construct all 3 concurrently. let (n0, n1, n2) = tokio::try_join!( CpuNode::new(configs[0].clone(), s3.clone()), @@ -489,7 +469,7 @@ impl CpuNodes { /// Create nodes and immediately truncate checkpoint tables and clear S3 buckets — /// the standard starting state for every workflow test. - pub async fn new_clean(configs: &CpuConfigs, s3: aws_sdk_s3::Client) -> eyre::Result { + pub async fn new_clean(configs: &CpuConfigs, s3: ObjectStoreClient) -> eyre::Result { let nodes = Self::new(configs, s3).await?; nodes.clear_all_s3_buckets(configs).await?; nodes.truncate_checkpoint_tables().await?; @@ -497,14 +477,11 @@ impl CpuNodes { // may or may not affect this test, but it is useful to know about. if let Ok(resp) = nodes.0[0] .s3 - .get_object() - .bucket("wf-smpcv2-dev-sync-protocol") - .key("dev_deleted_serial_ids.json") - .send() + .store("wf-smpcv2-dev-sync-protocol")? + .get(&path("dev_deleted_serial_ids.json")?) .await { - if let Ok(body) = resp.body.collect().await { - let bytes = body.into_bytes(); + if let Ok(bytes) = resp.bytes().await { let text = String::from_utf8_lossy(&bytes); tracing::warn!( exclusions = %text, @@ -595,13 +572,9 @@ impl CpuNodes { let rows = node.store.graph.get_genesis_graph_checkpoints().await?; for row in &rows { // Ignore errors — the object may already be gone. - let _ = node - .s3 - .delete_object() - .bucket(&config.checkpoint_bucket) - .key(&row.s3_key) - .send() - .await; + if let Ok(store) = node.s3.store(&config.checkpoint_bucket) { + let _ = store.delete(&path(&row.s3_key)?).await; + } } node.store.truncate_checkpoint_tables().await?; } @@ -709,13 +682,9 @@ impl CpuNodes { let keys = node.list_s3_keys(bucket).await?; for key in keys { // Ignore errors — objects may already be gone or other transient issues. - let _ = node - .s3 - .delete_object() - .bucket(bucket) - .key(&key) - .send() - .await; + if let Ok(store) = node.s3.store(bucket) { + let _ = store.delete(&path(&key)?).await; + } } } Ok(()) diff --git a/iris-mpc/tests/utils/key_rotation.rs b/iris-mpc/tests/utils/key_rotation.rs index 85e9f666f4..17ee7b7d35 100644 --- a/iris-mpc/tests/utils/key_rotation.rs +++ b/iris-mpc/tests/utils/key_rotation.rs @@ -11,10 +11,10 @@ //! 2. Rotates ECDH keys twice for each of the three MPC parties. use aws_config::SdkConfig; -use aws_sdk_s3::Client as S3Client; use aws_sdk_secretsmanager::Client as SecretsManagerClient; use base64::{engine::general_purpose::STANDARD, Engine}; use eyre::Result; +use iris_mpc_common::object_store::{path, ObjectStoreClient, ObjectStoreExt}; use rand::{thread_rng, Rng}; use sodiumoxide::crypto::box_::{curve25519xsalsa20poly1305, Seed}; @@ -145,23 +145,26 @@ pub async fn rotate_keys( } // Build SDK clients, optionally overriding the endpoint for LocalStack. - let mut s3_builder = aws_sdk_s3::config::Builder::from(sdk_config); let mut sm_builder = aws_sdk_secretsmanager::config::Builder::from(sdk_config); + let mut object_store = ObjectStoreClient::new( + sdk_config.region().map(ToString::to_string), + endpoint_url.is_some(), + ) + .with_aws_sdk_config(sdk_config); if let Some(ref url) = endpoint_url { - s3_builder = s3_builder.endpoint_url(url).force_path_style(true); + object_store = object_store + .with_option("aws_endpoint", url) + .with_option("aws_allow_http", url.starts_with("http://")); sm_builder = sm_builder.endpoint_url(url); } - let s3 = S3Client::from_conf(s3_builder.build()); let sm = SecretsManagerClient::from_conf(sm_builder.build()); // Upload public key to S3. - s3.put_object() - .bucket(&bucket_name) - .key(bucket_key_name) - .body(pub_key_str.into_bytes().into()) - .send() + object_store + .store(&bucket_name)? + .put(&path(bucket_key_name)?, pub_key_str.into_bytes().into()) .await .map_err(|e| { eyre::eyre!( @@ -200,7 +203,7 @@ pub async fn rotate_keys( /// `DEFAULT_REGION` in `docker-compose.hawk-db.yaml`). async fn base_sdk_config() -> SdkConfig { aws_config::from_env() - .region(aws_sdk_s3::config::Region::new("us-east-1")) + .region(aws_sdk_secretsmanager::config::Region::new("us-east-1")) .load() .await } diff --git a/iris-mpc/tests/utils/runner.rs b/iris-mpc/tests/utils/runner.rs index a9432e2a0c..3c136ce36d 100644 --- a/iris-mpc/tests/utils/runner.rs +++ b/iris-mpc/tests/utils/runner.rs @@ -1,5 +1,5 @@ use super::CpuConfigs; -use aws_config; +use iris_mpc_common::object_store::ObjectStoreClient; use tokio_util::sync::CancellationToken; /// Lifecycle trait implemented by each `wal_NNN` test struct. @@ -109,7 +109,7 @@ pub struct CpuTestContext { pub env: TestEnvironment, /// Pre-built S3 client pointed at the LocalStack instance for this environment. /// Shared across all lifecycle phases to guarantee a consistent endpoint configuration. - pub s3_client: aws_sdk_s3::Client, + pub s3_client: ObjectStoreClient, /// Test number for log tagging and config selection. pub kind: usize, /// Test run index (usually 1) for multi-run scenarios. @@ -144,13 +144,11 @@ impl CpuTestContext { } else { TestEnvironment::Local }; - let aws_config = aws_config::load_from_env().await; - let s3_client = aws_sdk_s3::Client::from_conf( - aws_sdk_s3::config::Builder::from(&aws_config) - .endpoint_url(env.s3_endpoint()) - .force_path_style(true) - .build(), - ); + let sdk_config = aws_config::load_from_env().await; + let s3_client = ObjectStoreClient::new(Some("us-east-1".to_owned()), true) + .with_aws_sdk_config(&sdk_config) + .with_option("aws_endpoint", env.s3_endpoint()) + .with_option("aws_allow_http", true); Self { configs: Self::load_configs(&env), s3_client,