From bed4b7937573af36a1c65d725867d699e70faf93 Mon Sep 17 00:00:00 2001 From: Denis Avvakumov Date: Fri, 8 May 2026 17:20:45 +0300 Subject: [PATCH] Add X-Forwarded-For IP extraction and rate limit logging --- Cargo.lock | 219 +++++++++++------------- Cargo.toml | 8 +- dev-env/config/config-single.toml | 2 + dev-env/config/config1.toml | 2 + dev-env/config/config2.toml | 2 + dev-env/config/config3.toml | 2 + src/config.rs | 160 ++++++++++++++++- src/config_runtime.rs | 16 +- src/http3/client_ip.rs | 107 ++++++++++++ src/http3/handlers/admin.rs | 16 +- src/http3/handlers/common/peer.rs | 9 +- src/http3/handlers/core.rs | 12 ++ src/http3/handlers/result/add_result.rs | 2 +- src/http3/handlers/task/add_task.rs | 2 +- src/http3/mod.rs | 1 + src/http3/rate_limits.rs | 76 ++++++-- src/http3/response.rs | 23 ++- src/http3/whitelist.rs | 8 +- tests/client_http_api/add_task.rs | 103 +++++++++++ tests/client_http_api/misc.rs | 43 ++++- tests/common/real_harness.rs | 5 + 21 files changed, 648 insertions(+), 170 deletions(-) create mode 100644 src/http3/client_ip.rs diff --git a/Cargo.lock b/Cargo.lock index 1829e80..3cf1890 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -219,9 +219,9 @@ dependencies = [ [[package]] name = "asn1-rs" -version = "0.7.1" +version = "0.7.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "56624a96882bb8c26d61312ae18cb45868e5a9992ea73c58e45c3101e56a1e60" +checksum = "b7f43a50ac4fdca5df8e885c21b835997f0a1cdee65494a6847694a98652d9d8" dependencies = [ "asn1-rs-derive", "asn1-rs-impl", @@ -258,9 +258,9 @@ dependencies = [ [[package]] name = "astral-tokio-tar" -version = "0.6.1" +version = "0.6.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4ce73b17c62717c4b6a9af10b43e87c578b0cac27e00666d48304d3b7d2c0693" +checksum = "cb50a7aae84a03bf55b067832bc376f4961b790c97e64d3eacee97d389b90277" dependencies = [ "filetime", "futures-core", @@ -392,18 +392,18 @@ dependencies = [ [[package]] name = "avif-serialize" -version = "0.8.8" +version = "0.8.9" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "375082f007bd67184fb9c0374614b29f9aaa604ec301635f72338bb65386a53d" +checksum = "e7178fe5f7d460b13895ebb9dcb28a3a6216d2df2574a0806cb51b555d297f38" dependencies = [ "arrayvec", ] [[package]] name = "aws-lc-rs" -version = "1.16.3" +version = "1.17.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "0ec6fb3fe69024a75fa7e1bfb48aa6cf59706a101658ea01bfd33b2b248a038f" +checksum = "5ec2f1fc3ec205783a5da9a7e6c1509cc69dedf09a1949e412c1e18469326d00" dependencies = [ "aws-lc-sys", "untrusted 0.7.1", @@ -412,9 +412,9 @@ dependencies = [ [[package]] name = "aws-lc-sys" -version = "0.40.0" +version = "0.41.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f50037ee5e1e41e7b8f9d161680a725bd1626cb6f8c7e901f91f942850852fe7" +checksum = "1a2f9779ce85b93ab6170dd940ad0169b5766ff848247aff13bb788b832fe3f4" dependencies = [ "cc", "cmake", @@ -504,6 +504,15 @@ version = "1.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "2af50177e190e07a26ab74f8b1efbfe2ef87da2116221318cb1c2e82baf7de06" +[[package]] +name = "bit-vec" +version = "0.9.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b71798fca2c1fe1086445a7258a4bc81e6e49dcd24c8d0dd9a1e57395b603f51" +dependencies = [ + "serde", +] + [[package]] name = "bit_field" version = "0.10.3" @@ -811,9 +820,9 @@ dependencies = [ [[package]] name = "cc" -version = "1.2.61" +version = "1.2.62" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d16d90359e986641506914ba71350897565610e87ce0ad9e6f28569db3dd5c6d" +checksum = "a1dce859f0832a7d088c4f1119888ab94ef4b5d6795d1ce05afb7fe159d79f98" dependencies = [ "find-msvc-tools", "jobserver", @@ -1184,9 +1193,9 @@ dependencies = [ [[package]] name = "crypto-common" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "77727bb15fa921304124b128af125e7e3b968275d1b108b379190264f4423710" +checksum = "ce6e4c961d6cd6c9a86db418387425e8bdeaf05b3c8bc1411e6dca4c252f1453" dependencies = [ "hybrid-array", ] @@ -1431,7 +1440,7 @@ checksum = "f1dd6dbb5841937940781866fa1281a1ff7bd3bf827091440879f9994983d5c2" dependencies = [ "block-buffer 0.12.0", "const-oid 0.10.2", - "crypto-common 0.2.1", + "crypto-common 0.2.2", "ctutils", ] @@ -1448,9 +1457,9 @@ dependencies = [ [[package]] name = "docker_credential" -version = "1.3.3" +version = "1.4.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a4564c274ebf369f501de192b02a0b81a5c4bda375abfe526aa70fc702fa6fa0" +checksum = "29547a1dc60885a552306986316bc9701ba120c1a8db6769fa68691529ad373d" dependencies = [ "base64", "serde", @@ -1510,9 +1519,9 @@ dependencies = [ [[package]] name = "either" -version = "1.15.0" +version = "1.16.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "48c757948c5ede0e46177b7add2e67155f70e33c07fea8284df6576da70b3719" +checksum = "91622ff5e7162018101f2fea40d6ebf4a78bbe5a49736a2020649edf9693679e" [[package]] name = "elliptic-curve" @@ -1735,13 +1744,12 @@ checksum = "28dea519a9695b9977216879a3ebfddf92f1c08c05d984f8996aecd6ecdc811d" [[package]] name = "filetime" -version = "0.2.27" +version = "0.2.29" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f98844151eee8917efc50bd9e8318cb963ae8b297431495d3f758616ea5c57db" +checksum = "5c287a33c7f0a620c38e641e7f60827713987b3c0f26e8ddc9462cc69cf75759" dependencies = [ "cfg-if", "libc", - "libredox", ] [[package]] @@ -2181,9 +2189,9 @@ dependencies = [ [[package]] name = "hashbrown" -version = "0.17.0" +version = "0.17.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4f467dd6dccf739c208452f8014c75c18bb8301b050ad1cfb27153803edb0f51" +checksum = "ed5909b6e89a2db4456e54cd5f673791d7eca6732202bbf2a9cc504fe2f9b84a" [[package]] name = "headers" @@ -2341,9 +2349,9 @@ checksum = "df3b46402a9d5adb4c86a0cf463f42e19994e3ee891101b1841f30a545cb49a9" [[package]] name = "hybrid-array" -version = "0.4.11" +version = "0.4.12" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "08d46837a0ed51fe95bd3b05de33cd64a1ee88fc797477ca48446872504507c5" +checksum = "9155a582abd142abc056962c29e3ce5ff2ad5469f4246b537ed42c5deba857da" dependencies = [ "typenum", ] @@ -2651,7 +2659,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d466e9454f08e4a911e14806c24e16fba1b4c121d1ea474396f396069cf949d9" dependencies = [ "equivalent", - "hashbrown 0.17.0", + "hashbrown 0.17.1", "serde", "serde_core", ] @@ -2715,16 +2723,6 @@ version = "2.12.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "d98f6fed1fde3f8c21bc40a1abb88dd75e67924f9cffc3ef95607bad8017f8e2" -[[package]] -name = "iri-string" -version = "0.7.12" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "25e659a4bb38e810ebc252e53b5814ff908a8c58c2a9ce2fae1bbec24cbf4e20" -dependencies = [ - "memchr", - "serde", -] - [[package]] name = "is_terminal_polyfill" version = "1.70.2" @@ -2832,9 +2830,9 @@ dependencies = [ [[package]] name = "js-sys" -version = "0.3.97" +version = "0.3.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a1840c94c045fbcf8ba2812c95db44499f7c64910a912551aaaa541decebcacf" +checksum = "67df7112613f8bfd9150013a0314e196f4800d3201ae742489d999db2f979f08" dependencies = [ "cfg-if", "futures-util", @@ -2863,9 +2861,9 @@ dependencies = [ [[package]] name = "kqueue-sys" -version = "1.1.0" +version = "1.1.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a7b65860415f949f23fa882e669f2dbd4a0f0eeb1acdd56790b30494afd7da2f" +checksum = "07293a4e297ac234359b510362495713f75ea345d5307140414f20c69ffeb087" dependencies = [ "bitflags", "libc", @@ -2916,9 +2914,9 @@ checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" [[package]] name = "libmimalloc-sys" -version = "0.1.47" +version = "0.1.48" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2d1eacfa31c33ec25e873c136ba5669f00f9866d0688bea7be4d3f7e43067df6" +checksum = "2892ae4ea6fa2cb7acb0e236a6880d39523239cd9089de71d220910ccc806790" dependencies = [ "cc", ] @@ -2929,10 +2927,7 @@ version = "0.1.16" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e02f3bb43d335493c96bf3fd3a321600bf6bd07ed34bc64118e9293bdffea46c" dependencies = [ - "bitflags", "libc", - "plain", - "redox_syscall 0.7.5", ] [[package]] @@ -2949,9 +2944,9 @@ checksum = "92daf443525c4cce67b150400bc2316076100ce0b3686209eb8cf3c31612e6f0" [[package]] name = "local-ip-address" -version = "0.6.12" +version = "0.6.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d7b0187df4e614e42405b49511b82ff7a1774fbd9a816060ee465067847cac22" +checksum = "aa08fb2b1ec3ea84575e94b489d06d4ce0cbf052d12acd515838f50e3c3d63e3" dependencies = [ "libc", "neli", @@ -3040,9 +3035,9 @@ dependencies = [ [[package]] name = "mimalloc" -version = "0.1.50" +version = "0.1.51" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b3627c4272df786b9260cabaa46aec1d59c93ede723d4c3ef646c503816b0640" +checksum = "ebca48a43116bc25f18a61360f1be98412f50cc218f5e52c823086b999a4a21a" dependencies = [ "libmimalloc-sys", ] @@ -3159,9 +3154,9 @@ dependencies = [ [[package]] name = "multra" -version = "1.0.0" +version = "1.0.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "3bc3fc14248e7e7891048c0b23a76bc862bcbf41a1ca3b3bbf5032cd0aea06f1" +checksum = "9de3f58a19fd337ad34dfba18efda0470f8fb3a618d088d2e5b9f94843fd4619" dependencies = [ "bytes", "encoding_rs", @@ -3329,9 +3324,9 @@ dependencies = [ [[package]] name = "num-conv" -version = "0.2.1" +version = "0.2.2" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "c6673768db2d862beb9b39a78fdcb1a69439615d5794a1be50caa9bc92c81967" +checksum = "521739c6d2bac4aa25192232afe6841231376b2b26d4d9fae5ecf8ca5772e441" [[package]] name = "num-derive" @@ -3541,7 +3536,7 @@ checksum = "2621685985a2ebf1c516881c026032ac7deafcda1a2c9b7850dc81e3dfcb64c1" dependencies = [ "cfg-if", "libc", - "redox_syscall 0.5.18", + "redox_syscall", "smallvec", "windows-link", ] @@ -3629,18 +3624,18 @@ dependencies = [ [[package]] name = "pin-project" -version = "1.1.11" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f1749c7ed4bcaf4c3d0a3efc28538844fb29bcdd7d2b67b2be7e20ba861ff517" +checksum = "2466b2336ed02bcdca6b294417127b90ec92038d1d5c4fbeac971a922e0e0924" dependencies = [ "pin-project-internal", ] [[package]] name = "pin-project-internal" -version = "1.1.11" +version = "1.1.13" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d9b20ed30f105399776b9c883e68e536ef602a16ae6f596d2c473591d6ad64c6" +checksum = "c96395f0a926bc13b1c17622aaddda1ecb55d49c8f1bf9777e4d877800a43f8b" dependencies = [ "proc-macro2", "quote", @@ -3680,12 +3675,6 @@ version = "0.3.33" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "19f132c84eca552bf34cab8ec81f1c1dcc229b811638f9d283dceabe58c5569e" -[[package]] -name = "plain" -version = "0.2.3" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b4596b6d070b27117e987119b4dac604f3c58cfb0b191112e24771b2faeac1a6" - [[package]] name = "png" version = "0.18.1" @@ -4192,9 +4181,9 @@ dependencies = [ [[package]] name = "rcgen" -version = "0.14.7" +version = "0.14.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "10b99e0098aa4082912d4c649628623db6aba77335e4f4569ff5083a6448b32e" +checksum = "57f6d249aad744e274e682777a50283a225a32705394ee6d5fcc01efa25e4055" dependencies = [ "aws-lc-rs", "pem", @@ -4214,15 +4203,6 @@ dependencies = [ "bitflags", ] -[[package]] -name = "redox_syscall" -version = "0.7.5" -source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "4666a1a60d8412eab19d94f6d13dcc9cea0a5ef4fdf6a5db306537413c661b1b" -dependencies = [ - "bitflags", -] - [[package]] name = "ref-cast" version = "1.0.25" @@ -4433,9 +4413,9 @@ dependencies = [ [[package]] name = "rust_decimal" -version = "1.41.0" +version = "1.42.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ce901f9a19d251159075a4c37af514c3b8ef99c22e02dd8c19161cf397ee94a" +checksum = "0c5108e3d4d903e21aac27f12ba5377b6b34f9f44b325e4894c7924169d06995" dependencies = [ "arrayvec", "borsh", @@ -4597,9 +4577,9 @@ checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" [[package]] name = "saa" -version = "5.5.1" +version = "5.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "bd8d438861332c3b1ac396c77bd9cac620ea1ff347efb63c05a83d8f0a593899" +checksum = "68f5acb362a0e75c2a963532fa7fabf13dff81626dc494df16488d30befcbea0" [[package]] name = "salvo" @@ -4810,9 +4790,9 @@ dependencies = [ [[package]] name = "scc" -version = "3.7.0" +version = "3.7.1" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "16c154cf1d115a1e901d7f4e3f279eb6eb455f0d670c1cf3c1aa74d50ad37fa9" +checksum = "5bcd12b6caff5213cc3c03123cde8c3db5e413008a63b0c0ba35e6275825ea92" dependencies = [ "saa", "sdd", @@ -5037,11 +5017,12 @@ dependencies = [ [[package]] name = "serde_with" -version = "3.19.0" +version = "3.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "f05839ce67618e14a09b286535c0d9c94e85ef25469b0e13cb4f844e5593eb19" +checksum = "e72c1c2cb7b223fafb600a619537a871c2818583d619401b785e7c0b746ccde2" dependencies = [ "base64", + "bs58", "chrono", "hex", "indexmap 1.9.3", @@ -5056,9 +5037,9 @@ dependencies = [ [[package]] name = "serde_with_macros" -version = "3.19.0" +version = "3.20.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "cf2ebbe86054f9b45bc3881e865683ccfaccce97b9b4cb53f3039d67f355a334" +checksum = "b90c488738ecb4fb0262f41f43bc40efc5868d9fb744319ddf5f5317f417bfac" dependencies = [ "darling 0.23.0", "proc-macro2", @@ -5557,9 +5538,9 @@ dependencies = [ [[package]] name = "tokio" -version = "1.52.2" +version = "1.52.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "110a78583f19d5cdb2c5ccf321d1290344e71313c6c37d43520d386027d18386" +checksum = "8fc7f01b389ac15039e4dc9531aa973a135d7a4135281b12d7c1bc79fd57fffe" dependencies = [ "bytes", "libc", @@ -5611,13 +5592,12 @@ dependencies = [ [[package]] name = "tokio-postgres-rustls" -version = "0.13.0" +version = "0.14.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "27d684bad428a0f2481f42241f821db42c54e2dc81d8c00db8536c506b0a0144" +checksum = "4c2ad44aa0ae96db89c4742212ed41645b2f597311ff6e1945542a4d9fadc2fb" dependencies = [ - "const-oid 0.9.6", - "ring", "rustls", + "sha2 0.11.0", "tokio", "tokio-postgres", "tokio-rustls", @@ -5724,9 +5704,9 @@ checksum = "756daf9b1013ebe47a8776667b466417e2d4c5679d441c26230efd9ef78692db" [[package]] name = "tonic" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "fec7c61a0695dc1887c1b53952990f3ad2e3a31453e1f49f10e75424943a93ec" +checksum = "ac2a5518c70fa84342385732db33fb3f44bc4cc748936eb5833d2df34d6445ef" dependencies = [ "async-trait", "axum", @@ -5753,9 +5733,9 @@ dependencies = [ [[package]] name = "tonic-prost" -version = "0.14.5" +version = "0.14.6" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "a55376a0bbaa4975a3f10d009ad763d8f4108f067c7c2e74f3001fb49778d309" +checksum = "50849f68853be452acf590cde0b146665b8d507b3b8af17261df47e02c209ea0" dependencies = [ "bytes", "prost", @@ -5783,20 +5763,20 @@ dependencies = [ [[package]] name = "tower-http" -version = "0.6.8" +version = "0.6.11" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "d4e6559d53cc268e5031cd8429d05415bc4cb4aefc4aa5d6cc35fbf5b924a1f8" +checksum = "4cfcf7e2740e6fc6d4d688b4ef00650406bb94adf4731e43c096c3a19fe40840" dependencies = [ "bitflags", "bytes", "futures-util", "http", "http-body", - "iri-string", "pin-project-lite", "tower", "tower-layer", "tower-service", + "url", ] [[package]] @@ -6205,9 +6185,9 @@ dependencies = [ [[package]] name = "wasm-bindgen" -version = "0.2.120" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "df52b6d9b87e0c74c9edfa1eb2d9bf85e5d63515474513aa50fa181b3c4f5db1" +checksum = "49ace1d07c165b0864824eee619580c4689389afa9dc9ed3a4c75040d82e6790" dependencies = [ "cfg-if", "once_cell", @@ -6218,9 +6198,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-futures" -version = "0.4.70" +version = "0.4.71" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "af934872acec734c2d80e6617bbb5ff4f12b052dd8e6332b0817bce889516084" +checksum = "96492d0d3ffba25305a7dc88720d250b1401d7edca02cc3bcd50633b424673b8" dependencies = [ "js-sys", "wasm-bindgen", @@ -6228,9 +6208,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro" -version = "0.2.120" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "78b1041f495fb322e64aca85f5756b2172e35cd459376e67f2a6c9dffcedb103" +checksum = "8e68e6f4afd367a562002c05637acb8578ff2dea1943df76afb9e83d177c8578" dependencies = [ "quote", "wasm-bindgen-macro-support", @@ -6238,9 +6218,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-macro-support" -version = "0.2.120" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "9dcd0ff20416988a18ac686d4d4d0f6aae9ebf08a389ff5d29012b05af2a1b41" +checksum = "d95a9ec35c64b2a7cb35d3fead40c4238d0940c86d107136999567a4703259f2" dependencies = [ "bumpalo", "proc-macro2", @@ -6251,9 +6231,9 @@ dependencies = [ [[package]] name = "wasm-bindgen-shared" -version = "0.2.120" +version = "0.2.121" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "49757b3c82ebf16c57d69365a142940b384176c24df52a087fb748e2085359ea" +checksum = "c4e0100b01e9f0d03189a92b96772a1fb998639d981193d7dbab487302513441" dependencies = [ "unicode-ident", ] @@ -6307,9 +6287,9 @@ dependencies = [ [[package]] name = "web-sys" -version = "0.3.97" +version = "0.3.98" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2eadbac71025cd7b0834f20d1fe8472e8495821b4e9801eb0a60bd1f19827602" +checksum = "4b572dff8bcf38bad0fa19729c89bb5748b2b9b1d8be70cf90df697e3a8f32aa" dependencies = [ "js-sys", "wasm-bindgen", @@ -6693,9 +6673,9 @@ checksum = "d6bbff5f0aada427a1e5a6da5f1f98158182f26556f345ac9e04d36d0ebed650" [[package]] name = "winnow" -version = "1.0.2" +version = "1.0.3" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "2ee1708bef14716a11bae175f579062d4554d95be2c6829f518df847b7b3fdd0" +checksum = "0592e1c9d151f854e6fd382574c3a0855250e1d9b2f99d9281c6e6391af352f1" dependencies = [ "memchr", ] @@ -6852,9 +6832,9 @@ dependencies = [ [[package]] name = "xml" -version = "1.2.1" +version = "1.3.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "b8aa498d22c9bbaf482329839bc5620c46be275a19a812e9a22a2b07529a642a" +checksum = "636f85e5ca6488e96401b61eb7de54f4e44755c988af0f52cf90230c312a1a89" [[package]] name = "xxhash-rust" @@ -6870,10 +6850,11 @@ checksum = "7a5a4b21e1a62b67a2970e6831bc091d7b87e119e7f9791aef9702e3bef04448" [[package]] name = "yasna" -version = "0.5.2" +version = "0.6.0" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "e17bb3549cc1321ae1296b9cdc2698e2b6cb1992adfa19a8c72e5b7a738f44cd" +checksum = "b5f6765e852b9b4dc8e2a76843e4d64d1cea8e79bcde0b6901aea8e7c7f08282" dependencies = [ + "bit-vec", "time", ] @@ -6922,9 +6903,9 @@ dependencies = [ [[package]] name = "zerofrom" -version = "0.1.7" +version = "0.1.8" source = "registry+https://github.com/rust-lang/crates.io-index" -checksum = "69faa1f2a1ea75661980b013019ed6687ed0e83d069bc1114e2cc74c6c04c4df" +checksum = "0ec05a11813ea801ff6d75110ad09cd0824ddba17dfe17128ea0d5f68e6c5272" dependencies = [ "zerofrom-derive", ] diff --git a/Cargo.toml b/Cargo.toml index 273d4f9..03fed7b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -46,7 +46,7 @@ h3-quinn = { version = "0.0.10" } http = "1.4.0" backon = "1.6.0" openraft = { version = "0.9.24", features = ["storage-v2", "serde"] } -tokio-postgres-rustls = "0.13.0" +tokio-postgres-rustls = "0.14.0" rustls = { version = "0.23.40" } rustls-platform-verifier = "0.7.0" toml = "1.1.2" @@ -56,10 +56,10 @@ serde_json = "1.0.149" rmp-serde = "1.3.1" anyhow = "1.0.102" bytes = "1.11.1" -rcgen = { version = "0.14.7", features = ["pem"] } +rcgen = { version = "0.14.8", features = ["pem"] } foldhash = "0.2.0" portable-atomic = "1.13.1" -scc = "3.7.0" +scc = "3.7.1" sdd = "4.8.6" schnorrkel = "0.11.5" base64-simd = "0.8" @@ -68,7 +68,7 @@ bs58 = "0.5.1" prometheus = "0.14.0" blake3 = "1.8.5" moka = { version = "0.12.15", features = ["sync"] } -mimalloc = { version = "0.1.50", default-features = false } +mimalloc = { version = "0.1.51", default-features = false } regex = "1.12.3" image = "0.25.10" chrono = { version = "0.4", features = ["serde", "clock"] } diff --git a/dev-env/config/config-single.toml b/dev-env/config/config-single.toml index 94769c9..18523f1 100644 --- a/dev-env/config/config-single.toml +++ b/dev-env/config/config-single.toml @@ -59,6 +59,8 @@ get_timeout_sec = 2 max_idle_timeout_sec = 4 keep_alive_interval_sec = 1 rate_limit_whitelist = [] +# Trust X-Forwarded-For only from the immediate proxy ranges that can reach the gateway. +trusted_proxy_cidrs = [] worker_whitelist = [] distributed_rate_limiter_max_capacity = 4096 allowed_origins = ["discord", "unity", "blender", "gen.404.xyz", "api"] diff --git a/dev-env/config/config1.toml b/dev-env/config/config1.toml index 222bc5b..06bfe86 100644 --- a/dev-env/config/config1.toml +++ b/dev-env/config/config1.toml @@ -59,6 +59,8 @@ get_timeout_sec = 2 max_idle_timeout_sec = 4 keep_alive_interval_sec = 1 rate_limit_whitelist = [] +# Trust X-Forwarded-For only from the immediate proxy ranges that can reach the gateway. +trusted_proxy_cidrs = [] worker_whitelist = [] distributed_rate_limiter_max_capacity = 4096 allowed_origins = ["discord", "unity", "blender", "gen.404.xyz", "api"] diff --git a/dev-env/config/config2.toml b/dev-env/config/config2.toml index 492ac58..bd74a44 100644 --- a/dev-env/config/config2.toml +++ b/dev-env/config/config2.toml @@ -59,6 +59,8 @@ get_timeout_sec = 2 max_idle_timeout_sec = 4 keep_alive_interval_sec = 1 rate_limit_whitelist = [] +# Trust X-Forwarded-For only from the immediate proxy ranges that can reach the gateway. +trusted_proxy_cidrs = [] worker_whitelist = [] distributed_rate_limiter_max_capacity = 4096 allowed_origins = ["discord", "unity", "blender", "gen.404.xyz", "api"] diff --git a/dev-env/config/config3.toml b/dev-env/config/config3.toml index dffdced..e7425c7 100644 --- a/dev-env/config/config3.toml +++ b/dev-env/config/config3.toml @@ -59,6 +59,8 @@ get_timeout_sec = 2 max_idle_timeout_sec = 4 keep_alive_interval_sec = 1 rate_limit_whitelist = [] +# Trust X-Forwarded-For only from the immediate proxy ranges that can reach the gateway. +trusted_proxy_cidrs = [] worker_whitelist = [] distributed_rate_limiter_max_capacity = 4096 allowed_origins = ["discord", "unity", "blender", "gen.404.xyz", "api"] diff --git a/src/config.rs b/src/config.rs index a32e686..5872003 100644 --- a/src/config.rs +++ b/src/config.rs @@ -1,6 +1,7 @@ use anyhow::{Result, anyhow}; use foldhash::HashSet; use serde::{Deserialize, Deserializer, Serialize}; +use std::net::{IpAddr, Ipv4Addr, Ipv6Addr}; use std::str::FromStr; use std::{fmt, path::Path, path::PathBuf}; use tracing::Level; @@ -137,6 +138,10 @@ pub struct HTTPConfig { pub basic_rate_limit: usize, pub add_task_unauthorized_per_ip_daily_rate_limit: usize, pub rate_limit_whitelist: HashSet, + /// CIDR ranges or literal IPs whose X-Forwarded-For headers are trusted. + /// Keep empty unless the gateway is only reachable through those proxies. + #[serde(default)] + pub trusted_proxy_cidrs: Vec, #[serde(default = "default_distributed_rate_limiter_max_capacity")] pub distributed_rate_limiter_max_capacity: usize, pub worker_per_minute_rate_limit: usize, @@ -452,10 +457,122 @@ pub fn validate_node_config(config: &NodeConfig) -> Result<()> { } tracing::warn!("TLS server verification is disabled; only use this in local development"); } + validate_trusted_proxy_cidrs(&config.http.trusted_proxy_cidrs)?; validate_raft_dns_config(config)?; Ok(()) } +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum TrustedProxyRange { + V4 { network: u32, prefix: u8 }, + V6 { network: u128, prefix: u8 }, +} + +impl TrustedProxyRange { + pub fn contains(&self, ip: IpAddr) -> bool { + match (*self, ip) { + (TrustedProxyRange::V4 { network, prefix }, IpAddr::V4(ip)) => { + u32::from(ip) & ipv4_prefix_mask(prefix) == network + } + (TrustedProxyRange::V6 { network, prefix }, IpAddr::V6(ip)) => { + u128::from(ip) & ipv6_prefix_mask(prefix) == network + } + _ => false, + } + } + + fn from_ip(ip: IpAddr) -> Self { + match ip { + IpAddr::V4(ip) => Self::from_ipv4_cidr(ip, 32), + IpAddr::V6(ip) => Self::from_ipv6_cidr(ip, 128), + } + } + + fn from_ipv4_cidr(ip: Ipv4Addr, prefix: u8) -> Self { + let network = u32::from(ip) & ipv4_prefix_mask(prefix); + Self::V4 { network, prefix } + } + + fn from_ipv6_cidr(ip: Ipv6Addr, prefix: u8) -> Self { + let network = u128::from(ip) & ipv6_prefix_mask(prefix); + Self::V6 { network, prefix } + } +} + +fn ipv4_prefix_mask(prefix: u8) -> u32 { + if prefix == 0 { + 0 + } else { + u32::MAX << (32 - u32::from(prefix)) + } +} + +fn ipv6_prefix_mask(prefix: u8) -> u128 { + if prefix == 0 { + 0 + } else { + u128::MAX << (128 - u32::from(prefix)) + } +} + +pub fn parse_trusted_proxy_cidr(entry: &str) -> Result { + let trimmed = entry.trim(); + if trimmed.is_empty() { + return Err(anyhow!( + "http.trusted_proxy_cidrs must not contain empty entries" + )); + } + + if let Ok(ip) = trimmed.parse::() { + return Ok(TrustedProxyRange::from_ip(ip)); + } + + if let Some((addr, prefix)) = trimmed.split_once('/') { + if prefix.contains('/') { + return Err(anyhow!( + "http.trusted_proxy_cidrs entry '{}' must be an IP address or CIDR block", + entry + )); + } + let ip = addr.trim().parse::().map_err(|_| { + anyhow!( + "http.trusted_proxy_cidrs entry '{}' must be an IP address or CIDR block", + entry + ) + })?; + let prefix = prefix.trim().parse::().map_err(|_| { + anyhow!( + "http.trusted_proxy_cidrs entry '{}' has an invalid CIDR prefix length", + entry + ) + })?; + return match ip { + IpAddr::V4(ip) if prefix <= 32 => Ok(TrustedProxyRange::from_ipv4_cidr(ip, prefix)), + IpAddr::V6(ip) if prefix <= 128 => Ok(TrustedProxyRange::from_ipv6_cidr(ip, prefix)), + IpAddr::V4(_) => Err(anyhow!( + "http.trusted_proxy_cidrs entry '{}' has IPv4 prefix length greater than 32", + entry + )), + IpAddr::V6(_) => Err(anyhow!( + "http.trusted_proxy_cidrs entry '{}' has IPv6 prefix length greater than 128", + entry + )), + }; + } + + Err(anyhow!( + "http.trusted_proxy_cidrs entry '{}' must be an IP address or CIDR block", + entry + )) +} + +fn validate_trusted_proxy_cidrs(entries: &[String]) -> Result<()> { + for entry in entries { + parse_trusted_proxy_cidr(entry)?; + } + Ok(()) +} + pub fn validate_multi_node_raft_dns_config(config: &NodeConfig) -> Result<()> { if config.raft.dns_name.trim().is_empty() { return Err(anyhow!( @@ -555,7 +672,10 @@ pub fn resolve_config_path(path: Option<&String>) -> Result { #[cfg(test)] mod tests { - use super::{NodeConfig, validate_multi_node_raft_dns_config, validate_node_config}; + use super::{ + NodeConfig, parse_trusted_proxy_cidr, validate_multi_node_raft_dns_config, + validate_node_config, + }; fn read_config_single() -> String { let path = std::path::PathBuf::from(env!("CARGO_MANIFEST_DIR")) @@ -582,6 +702,44 @@ mod tests { assert_eq!(config.http.generic_key_concurrent_limit, 3); } + #[test] + fn http_config_defaults_trusted_proxy_cidrs_when_missing() { + let config_text = read_config_single().replace("trusted_proxy_cidrs = []\n", ""); + let config: NodeConfig = + toml::from_str(&config_text).expect("parse config without trusted proxies"); + assert!(config.http.trusted_proxy_cidrs.is_empty()); + validate_node_config(&config).expect("empty trusted proxy cidrs are valid"); + } + + #[test] + fn http_config_rejects_invalid_trusted_proxy_cidrs() { + let config_text = read_config_single().replace( + "trusted_proxy_cidrs = []", + "trusted_proxy_cidrs = [\"not-a-cidr\"]", + ); + let config: NodeConfig = toml::from_str(&config_text).expect("parse config"); + let err = validate_node_config(&config).expect_err("reject invalid proxy cidr"); + assert!(err.to_string().contains("trusted_proxy_cidrs")); + } + + #[test] + fn trusted_proxy_cidr_parser_matches_ipv4_and_ipv6_ranges() { + let ipv4 = parse_trusted_proxy_cidr("35.191.0.0/16").expect("ipv4 cidr"); + assert!(ipv4.contains("35.191.22.10".parse().unwrap())); + assert!(!ipv4.contains("35.192.0.1".parse().unwrap())); + + let ipv6 = parse_trusted_proxy_cidr("2600:2d00:1:b029::/64").expect("ipv6 cidr"); + assert!(ipv6.contains("2600:2d00:1:b029::1".parse().unwrap())); + assert!(!ipv6.contains("2600:2d00:1:b02a::1".parse().unwrap())); + } + + #[test] + fn trusted_proxy_cidr_parser_accepts_literal_ip_as_single_host() { + let proxy = parse_trusted_proxy_cidr("127.0.0.1").expect("literal ip"); + assert!(proxy.contains("127.0.0.1".parse().unwrap())); + assert!(!proxy.contains("127.0.0.2".parse().unwrap())); + } + #[test] fn network_config_defaults_cluster_peer_egress_ips_when_missing() { let config_text = read_config_single(); diff --git a/src/config_runtime.rs b/src/config_runtime.rs index 1015270..faaa873 100644 --- a/src/config_runtime.rs +++ b/src/config_runtime.rs @@ -14,8 +14,8 @@ use tokio::sync::Notify; use tracing::{info, warn}; use crate::config::{ - HTTPConfig, ImageConfig, ModelParamsConfig, NodeConfig, PromptConfig, read_config_from_path, - validate_node_config, + HTTPConfig, ImageConfig, ModelParamsConfig, NodeConfig, PromptConfig, TrustedProxyRange, + parse_trusted_proxy_cidr, read_config_from_path, validate_node_config, }; use crate::http3::rate_limits::RateLimiters; use crate::http3::upload_limiter::ImageUploadLimiter; @@ -29,6 +29,7 @@ pub struct RuntimeConfigSnapshot { pub prompt_regex: Regex, pub rate_limit_whitelist: RateLimitWhitelist, pub cluster_ips: HashSet, + pub trusted_proxy_cidrs: Vec, pub rate_limiters: RateLimiters, pub rate_limit_service: RateLimitService, pub image_upload_limiter: ImageUploadLimiter, @@ -72,6 +73,10 @@ impl RuntimeConfigView { &self.snapshot.cluster_ips } + pub fn trusted_proxy_cidrs(&self) -> &[TrustedProxyRange] { + &self.snapshot.trusted_proxy_cidrs + } + pub fn rate_limits(&self) -> &RateLimitService { &self.snapshot.rate_limit_service } @@ -347,6 +352,12 @@ async fn build_runtime_snapshot(config: NodeConfig) -> Result>>()?; let rate_limit_service = RateLimitService::new(&config.http); let rate_limiters = RateLimiters::new(&config.http); @@ -359,6 +370,7 @@ async fn build_runtime_snapshot(config: NodeConfig) -> Result Option { + match req.remote_addr() { + salvo::conn::SocketAddr::IPv4(addr) => Some(IpAddr::V4(*addr.ip())), + salvo::conn::SocketAddr::IPv6(addr) => Some(IpAddr::V6(*addr.ip())), + _ => None, + } +} + +fn is_trusted_proxy(ip: IpAddr, trusted_proxy_cidrs: &[TrustedProxyRange]) -> bool { + trusted_proxy_cidrs.iter().any(|net| net.contains(ip)) +} + +fn gke_client_ip_from_x_forwarded_for(header: &str) -> Option { + let mut parts = header.split(',').rev().map(str::trim); + let _load_balancer_ip = parts.next()?.parse::().ok()?; + let client_ip = parts.next()?.parse::().ok()?; + + Some(client_ip) +} + +fn client_ip_from_parts( + remote_ip: Option, + x_forwarded_for: Option<&str>, + trusted_proxy_cidrs: &[TrustedProxyRange], +) -> Option { + let remote_ip = remote_ip?; + if !is_trusted_proxy(remote_ip, trusted_proxy_cidrs) { + return Some(remote_ip); + } + + x_forwarded_for + .and_then(gke_client_ip_from_x_forwarded_for) + .or(Some(remote_ip)) +} + +pub(crate) fn client_ip( + req: &Request, + trusted_proxy_cidrs: &[TrustedProxyRange], +) -> Option { + client_ip_from_parts( + remote_ip(req), + req.headers() + .get("x-forwarded-for") + .and_then(|v| v.to_str().ok()), + trusted_proxy_cidrs, + ) +} + +#[cfg(test)] +mod tests { + use super::client_ip_from_parts; + use crate::config::TrustedProxyRange; + use crate::config::parse_trusted_proxy_cidr; + use std::net::IpAddr; + + fn ip(value: &str) -> IpAddr { + value.parse().expect("ip") + } + + fn cidr(value: &str) -> TrustedProxyRange { + parse_trusted_proxy_cidr(value).expect("cidr") + } + + #[test] + fn ignores_forwarded_header_from_untrusted_remote() { + let trusted = [cidr("10.0.0.0/8")]; + + let resolved = client_ip_from_parts( + Some(ip("192.0.2.10")), + Some("198.51.100.99, 203.0.113.10"), + &trusted, + ); + + assert_eq!(resolved, Some(ip("192.0.2.10"))); + } + + #[test] + fn trusted_gke_header_uses_second_ip_from_right() { + let trusted = [cidr("192.0.2.0/24")]; + + let resolved = client_ip_from_parts( + Some(ip("192.0.2.10")), + Some("198.51.100.99, 203.0.113.20, 34.117.10.1"), + &trusted, + ); + + assert_eq!(resolved, Some(ip("203.0.113.20"))); + } + + #[test] + fn malformed_trusted_gke_header_falls_back_to_remote() { + let trusted = [cidr("192.0.2.0/24")]; + + let resolved = client_ip_from_parts( + Some(ip("192.0.2.10")), + Some("198.51.100.99, not-an-ip"), + &trusted, + ); + + assert_eq!(resolved, Some(ip("192.0.2.10"))); + } +} diff --git a/src/http3/handlers/admin.rs b/src/http3/handlers/admin.rs index 15ad660..1aeebda 100644 --- a/src/http3/handlers/admin.rs +++ b/src/http3/handlers/admin.rs @@ -6,6 +6,8 @@ use std::sync::Arc; use uuid::Uuid; use crate::api::response::GenericKeyResponse; +use crate::config::TrustedProxyRange; +use crate::http3::client_ip::client_ip; use crate::http3::depot_ext::DepotExt; use crate::http3::error::ServerError; use crate::http3::state::HttpState; @@ -56,7 +58,7 @@ pub async fn admin_key_check(depot: &mut Depot, req: &mut Request) -> Result<(), return Ok(()); } - if let Some(source_key) = admin_key_source_from_req(req) { + if let Some(source_key) = admin_key_source_from_req(req, cfg.trusted_proxy_cidrs()) { let limiter = state.admin_key_failure_limiter(); if limiter.is_blocked(&source_key).await || limiter.record_miss(source_key).await { return Err(ServerError::Json( @@ -74,10 +76,12 @@ pub async fn admin_key_check(depot: &mut Depot, req: &mut Request) -> Result<(), )) } -fn admin_key_source_from_req(req: &Request) -> Option> { - match req.remote_addr() { - salvo::conn::SocketAddr::IPv4(addr) => Some(Arc::::from(addr.ip().to_string())), - salvo::conn::SocketAddr::IPv6(addr) => Some(Arc::::from(addr.ip().to_string())), - _ => Some(Arc::::from(req.remote_addr().to_string())), +fn admin_key_source_from_req( + req: &Request, + trusted_proxy_cidrs: &[TrustedProxyRange], +) -> Option> { + if let Some(ip) = client_ip(req, trusted_proxy_cidrs) { + return Some(Arc::::from(ip.to_string())); } + Some(Arc::::from(req.remote_addr().to_string())) } diff --git a/src/http3/handlers/common/peer.rs b/src/http3/handlers/common/peer.rs index 6795475..25c8486 100644 --- a/src/http3/handlers/common/peer.rs +++ b/src/http3/handlers/common/peer.rs @@ -1,6 +1,13 @@ use salvo::prelude::Request; -pub(crate) fn request_ip(req: &Request) -> String { +use crate::http3::client_ip::client_ip; +use crate::http3::state::HttpState; + +pub(crate) fn request_ip(req: &Request, state: &HttpState) -> String { + let cfg = state.config(); + if let Some(ip) = client_ip(req, cfg.trusted_proxy_cidrs()) { + return ip.to_string(); + } match req.remote_addr() { salvo::conn::SocketAddr::IPv4(addr) => addr.ip().to_string(), salvo::conn::SocketAddr::IPv6(addr) => addr.ip().to_string(), diff --git a/src/http3/handlers/core.rs b/src/http3/handlers/core.rs index dbd0066..d74ee93 100644 --- a/src/http3/handlers/core.rs +++ b/src/http3/handlers/core.rs @@ -158,6 +158,10 @@ pub async fn api_or_generic_key_check( } if !context.key_is_uuid { + tracing::warn!( + "Unauthorized login attempt: Invalid API key format from IP {}", + context.source_addr.as_deref().unwrap_or("unknown") + ); return Err(ServerError::BadRequest( "Invalid API key format".to_string(), )); @@ -166,6 +170,10 @@ pub async fn api_or_generic_key_check( if context.has_authorized_key() { Ok(()) } else if context.auth_lookup_blocked { + tracing::warn!( + "Unauthorized login attempt: API key lookup blocked for IP {}", + context.source_addr.as_deref().unwrap_or("unknown") + ); Err(ServerError::Json( StatusCode::TOO_MANY_REQUESTS, json!({ @@ -174,6 +182,10 @@ pub async fn api_or_generic_key_check( }), )) } else { + tracing::warn!( + "Unauthorized login attempt: Invalid API key from IP {}", + context.source_addr.as_deref().unwrap_or("unknown") + ); Err(ServerError::Unauthorized("Invalid API key".to_string())) } } diff --git a/src/http3/handlers/result/add_result.rs b/src/http3/handlers/result/add_result.rs index 1219de2..3c5765d 100644 --- a/src/http3/handlers/result/add_result.rs +++ b/src/http3/handlers/result/add_result.rs @@ -41,7 +41,7 @@ pub async fn add_result_handler( res: &mut Response, ) -> Result<(), ServerError> { let state = depot.require::()?.clone(); - let worker_ip = request_ip(req); + let worker_ip = request_ip(req, &state); ResultSubmissionLogContext::parsing_started(&worker_ip); diff --git a/src/http3/handlers/task/add_task.rs b/src/http3/handlers/task/add_task.rs index 170f131..beb8e5e 100644 --- a/src/http3/handlers/task/add_task.rs +++ b/src/http3/handlers/task/add_task.rs @@ -38,7 +38,7 @@ pub async fn add_task_handler( ) -> Result<(), ServerError> { let state = depot.require::()?.clone(); let gateway_state = state.gateway_state().clone(); - let submitter_ip = request_ip(req); + let submitter_ip = request_ip(req, &state); let mut rollback_guard = PendingRateLimitRollbackGuard::new(None); let mut task_log: Option = None; diff --git a/src/http3/mod.rs b/src/http3/mod.rs index 546ea6a..ebb2a33 100644 --- a/src/http3/mod.rs +++ b/src/http3/mod.rs @@ -1,4 +1,5 @@ pub(crate) mod client; +pub(crate) mod client_ip; pub(crate) mod depot_ext; pub(crate) mod error; pub(crate) mod handlers; diff --git a/src/http3/rate_limits.rs b/src/http3/rate_limits.rs index d589f0c..858eeef 100644 --- a/src/http3/rate_limits.rs +++ b/src/http3/rate_limits.rs @@ -9,8 +9,9 @@ use std::sync::atomic::{AtomicU64, Ordering}; use std::time::Duration; use uuid::Uuid; -use crate::config::HTTPConfig; +use crate::config::{HTTPConfig, TrustedProxyRange}; use crate::db::GenerationBillingOwner; +use crate::http3::client_ip::client_ip; use crate::http3::depot_ext::DepotExt; use crate::http3::error::ServerError; use crate::http3::state::HttpState; @@ -198,28 +199,40 @@ impl RateLimitReservation { } } -fn decimal_ip_from_req(req: &mut Request) -> Option> { - match req.remote_addr() { - salvo::conn::SocketAddr::IPv4(addr) => { - let bits = addr.ip().to_bits(); +fn ip_from_req( + req: &Request, + trusted_proxy_cidrs: &[TrustedProxyRange], +) -> Option { + client_ip(req, trusted_proxy_cidrs) +} + +fn decimal_ip_from_req( + req: &mut Request, + trusted_proxy_cidrs: &[TrustedProxyRange], +) -> Option> { + let ip = ip_from_req(req, trusted_proxy_cidrs)?; + match ip { + std::net::IpAddr::V4(addr) => { + let bits = addr.to_bits(); let mut buf = itoa::Buffer::new(); Some(Arc::::from(buf.format(bits))) } - salvo::conn::SocketAddr::IPv6(addr) => { - let bits = addr.ip().to_bits(); + std::net::IpAddr::V6(addr) => { + let bits = addr.to_bits(); let mut buf = itoa::Buffer::new(); Some(Arc::::from(buf.format(bits))) } - _ => None, } } -fn source_addr_from_req(req: &mut Request) -> Option> { - match req.remote_addr() { - salvo::conn::SocketAddr::IPv4(addr) => Some(Arc::::from(addr.ip().to_string())), - salvo::conn::SocketAddr::IPv6(addr) => Some(Arc::::from(addr.ip().to_string())), - _ => Some(Arc::::from(req.remote_addr().to_string())), +fn source_addr_from_req( + req: &mut Request, + trusted_proxy_cidrs: &[TrustedProxyRange], +) -> Option> { + if let Some(ip) = ip_from_req(req, trusted_proxy_cidrs) { + return Some(Arc::::from(ip.to_string())); } + Some(Arc::::from(req.remote_addr().to_string())) } fn cached_decimal_ip(req: &mut Request, depot: &Depot) -> Option> { @@ -231,7 +244,13 @@ fn cached_decimal_ip(req: &mut Request, depot: &Depot) -> Option> { return Some(Arc::clone(addr)); } } - decimal_ip_from_req(req).or_else(|| source_addr_from_req(req)) + if let Ok(state) = depot.obtain::() { + let cfg = state.config(); + return decimal_ip_from_req(req, cfg.trusted_proxy_cidrs()) + .or_else(|| source_addr_from_req(req, cfg.trusted_proxy_cidrs())); + } + + decimal_ip_from_req(req, &[]).or_else(|| source_addr_from_req(req, &[])) } const UNAUTHORIZED_DAILY_COUNTER_TTL: Duration = Duration::from_secs(60 * 60 * 48); @@ -348,14 +367,14 @@ impl RateLimitContext { } fn base_rate_limit_context(req: &mut Request, state: &HttpState) -> RateLimitContext { + let cfg = state.config(); + let trusted_proxy_cidrs = cfg.trusted_proxy_cidrs(); let mut context = RateLimitContext { is_whitelisted_ip: is_whitelisted_ip(req, state), ..RateLimitContext::default() }; - context.decimal_ip = decimal_ip_from_req(req); - if context.decimal_ip.is_none() { - context.source_addr = source_addr_from_req(req); - } + context.decimal_ip = decimal_ip_from_req(req, trusted_proxy_cidrs); + context.source_addr = source_addr_from_req(req, trusted_proxy_cidrs); context } @@ -591,6 +610,7 @@ struct SubjectParams<'a> { active_limit: u64, daily_limit: u64, scope_label: &'a str, + user_email: Option<&'a str>, require_day_match: bool, } @@ -628,6 +648,7 @@ async fn check_subject_limit( let active_limit = params.active_limit; let daily_limit = params.daily_limit; let scope_label = params.scope_label; + let user_email = params.user_email; let require_day_match = params.require_day_match; let has_limits = active_limit > 0 || daily_limit > 0; @@ -668,6 +689,20 @@ async fn check_subject_limit( format!("{} daily task limit exceeded.", scope_label) } }; + + if subject != Subject::GenericGlobal { + if let Some(email) = user_email { + tracing::warn!( + "Rate limit exceeded for {} (user: {}): {:?}", + scope_label, + email, + rejection + ); + } else { + tracing::warn!("Rate limit exceeded for {}: {:?}", scope_label, rejection); + } + } + return Err(match rejection { RateLimitRejection::Active => ServerError::Json( StatusCode::TOO_MANY_REQUESTS, @@ -712,6 +747,7 @@ pub async fn reserve_add_task_rate_limit( active_limit: company.concurrent_limit, daily_limit: company.daily_limit, scope_label: company.name.as_ref(), + user_email: ctx.user_email.as_deref(), require_day_match: company.daily_limit > 0, }); } else if let (Some(user_id), Some((concurrent_limit, daily_limit))) = @@ -724,6 +760,7 @@ pub async fn reserve_add_task_rate_limit( active_limit: concurrent_limit, daily_limit, scope_label: "User", + user_email: ctx.user_email.as_deref(), require_day_match: daily_limit > 0, }); } else if ctx.key_is_uuid && ctx.is_generic_key { @@ -733,6 +770,7 @@ pub async fn reserve_add_task_rate_limit( active_limit: cfg.http().generic_key_concurrent_limit as u64, daily_limit: 0, scope_label: "Generic key", + user_email: None, require_day_match: false, }); } @@ -886,6 +924,7 @@ mod tests { active_limit: 0, daily_limit: 10, scope_label: "Company", + user_email: None, require_day_match: true, }; @@ -905,6 +944,7 @@ mod tests { active_limit: 0, daily_limit: 0, scope_label: "Company", + user_email: None, require_day_match: false, }; diff --git a/src/http3/response.rs b/src/http3/response.rs index 2f5d3fc..e3e38e4 100644 --- a/src/http3/response.rs +++ b/src/http3/response.rs @@ -2,6 +2,9 @@ use salvo::http::{StatusCode, header::ACCEPT}; use salvo::prelude::*; use tracing::error; +use crate::http3::client_ip::client_ip; +use crate::http3::state::HttpState; + const TOO_MANY_REQUESTS_HTML: &str = include_str!("responses/429.html"); const METHOD_NOT_ALLOWED_HTML: &str = include_str!("responses/405.html"); const BAD_REQUEST_HTML: &str = include_str!("responses/400.html"); @@ -50,17 +53,25 @@ fn accepts_html(req: &Request) -> bool { #[handler] pub async fn custom_response( req: &Request, - _depot: &mut Depot, + depot: &mut Depot, res: &mut Response, ctrl: &mut FlowCtrl, ) { if let Some(status) = res.status_code { if status.is_client_error() || status.is_server_error() { - let client_ip = match req.remote_addr() { - salvo::conn::SocketAddr::IPv4(addr_v4) => addr_v4.ip().to_string(), - salvo::conn::SocketAddr::IPv6(addr_v6) => addr_v6.ip().to_string(), - _ => "unknown".to_string(), - }; + let client_ip = depot + .obtain::() + .ok() + .and_then(|state| { + let cfg = state.config(); + client_ip(req, cfg.trusted_proxy_cidrs()) + }) + .map(|ip| ip.to_string()) + .unwrap_or_else(|| match req.remote_addr() { + salvo::conn::SocketAddr::IPv4(addr_v4) => addr_v4.ip().to_string(), + salvo::conn::SocketAddr::IPv6(addr_v6) => addr_v6.ip().to_string(), + _ => "unknown".to_string(), + }); log_error(status, client_ip, req, &res.body); } diff --git a/src/http3/whitelist.rs b/src/http3/whitelist.rs index 365ba2f..0519765 100644 --- a/src/http3/whitelist.rs +++ b/src/http3/whitelist.rs @@ -7,6 +7,7 @@ use std::sync::Arc; use tracing::warn; use crate::common::resolve::lookup_all_host_ips; +use crate::http3::client_ip::client_ip; use crate::http3::state::HttpState; #[derive(Clone)] @@ -107,11 +108,8 @@ pub async fn resolve_egress_ips(entries: &[String]) -> HashSet { } pub fn is_whitelisted_ip(req: &Request, state: &HttpState) -> bool { - let remote_ip = match req.remote_addr() { - salvo::conn::SocketAddr::IPv4(addr) => Some(IpAddr::V4(*addr.ip())), - salvo::conn::SocketAddr::IPv6(addr) => Some(IpAddr::V6(*addr.ip())), - _ => None, - }; + let cfg = state.config(); + let remote_ip = client_ip(req, cfg.trusted_proxy_cidrs()); if let Some(ip) = remote_ip { return state.gateway_state().is_rate_limit_whitelisted_ip(&ip); diff --git a/tests/client_http_api/add_task.rs b/tests/client_http_api/add_task.rs index 3d00813..2f8a34e 100644 --- a/tests/client_http_api/add_task.rs +++ b/tests/client_http_api/add_task.rs @@ -952,6 +952,109 @@ async fn add_task_unauthorized_limit_hot_updates_without_resetting_counters() { assert_eq!(third_status, StatusCode::UNAUTHORIZED); } +async fn post_unknown_key_with_xff( + h: &crate::support::TestHarness, + api_key: &str, + x_forwarded_for: &str, +) -> StatusCode { + let res = TestClient::post("http://localhost/add_task") + .add_header("x-api-key", api_key, true) + .add_header("x-forwarded-for", x_forwarded_for, true) + .json(&serde_json::json!({"prompt": "robot"})) + .send(&h.service) + .await; + let (status, _headers, _body) = read_response(res).await; + status +} + +async fn set_unauthorized_limit(h: &crate::support::TestHarness, limit: i32, whitelist: &[String]) { + set_gateway_runtime_settings( + h, + RateLimitPolicies::from_config(&h.config.http), + limit, + whitelist, + h.config.http.max_task_queue_len as i32, + h.config.http.request_file_size_limit as i64, + ) + .await; +} + +#[tokio::test] +async fn add_task_x_forwarded_for_cannot_spoof_whitelist_without_trusted_proxy() { + let h = build_harness().await; + let spoofed_whitelist = vec!["203.0.113.44".to_string()]; + set_unauthorized_limit(&h, 1, &spoofed_whitelist).await; + + let first_key = Uuid::new_v4().to_string(); + let second_key = Uuid::new_v4().to_string(); + + let first_status = + post_unknown_key_with_xff(&h, first_key.as_str(), "203.0.113.44, 34.117.10.1").await; + assert_eq!(first_status, StatusCode::UNAUTHORIZED); + + let second_status = + post_unknown_key_with_xff(&h, second_key.as_str(), "203.0.113.44, 34.117.10.1").await; + assert_eq!(second_status, StatusCode::TOO_MANY_REQUESTS); +} + +#[tokio::test] +async fn add_task_trusted_proxy_ignores_spoofed_xff_prefix() { + let h = build_harness_with_options(GatewayHarnessOptions { + trusted_proxy_cidrs: vec!["127.0.0.1/32".to_string()], + ..GatewayHarnessOptions::default() + }) + .await; + set_unauthorized_limit(&h, 1, &[]).await; + + let first_key = Uuid::new_v4().to_string(); + let second_key = Uuid::new_v4().to_string(); + + let first_status = post_unknown_key_with_xff( + &h, + first_key.as_str(), + "198.51.100.1, 203.0.113.20, 34.117.10.1", + ) + .await; + assert_eq!(first_status, StatusCode::UNAUTHORIZED); + + let second_status = post_unknown_key_with_xff( + &h, + second_key.as_str(), + "198.51.100.2, 203.0.113.20, 34.117.10.1", + ) + .await; + assert_eq!(second_status, StatusCode::TOO_MANY_REQUESTS); +} + +#[tokio::test] +async fn add_task_trusted_proxy_uses_gke_client_ip_position() { + let h = build_harness_with_options(GatewayHarnessOptions { + trusted_proxy_cidrs: vec!["127.0.0.1/32".to_string()], + ..GatewayHarnessOptions::default() + }) + .await; + set_unauthorized_limit(&h, 1, &[]).await; + + let first_key = Uuid::new_v4().to_string(); + let second_key = Uuid::new_v4().to_string(); + + let first_status = post_unknown_key_with_xff( + &h, + first_key.as_str(), + "198.51.100.1, 203.0.113.20, 34.117.10.1", + ) + .await; + assert_eq!(first_status, StatusCode::UNAUTHORIZED); + + let second_status = post_unknown_key_with_xff( + &h, + second_key.as_str(), + "198.51.100.1, 203.0.113.21, 34.117.10.1", + ) + .await; + assert_eq!(second_status, StatusCode::UNAUTHORIZED); +} + #[tokio::test] async fn add_task_generic_key_ok() { let h = build_harness().await; diff --git a/tests/client_http_api/misc.rs b/tests/client_http_api/misc.rs index c067287..64a6bc6 100644 --- a/tests/client_http_api/misc.rs +++ b/tests/client_http_api/misc.rs @@ -53,13 +53,44 @@ async fn get_load_gateway_last_update_advances() { } #[tokio::test] -async fn get_load_empty() { +async fn get_load_missing_gateway_info_returns_500() { + const MAX_GET_LOAD_ATTEMPTS: usize = 50; + let h = build_harness_with_gateway_info(false).await; - let res = TestClient::get("http://localhost/get_load") - .send(&h.service) - .await; - let (status, _headers, body) = read_response(res).await; - assert_eq!(status, StatusCode::INTERNAL_SERVER_ERROR, "body: {body:?}"); + + for attempt in 1..=MAX_GET_LOAD_ATTEMPTS { + let res = TestClient::get("http://localhost/get_load") + .send(&h.service) + .await; + let (status, _headers, body) = read_response(res).await; + + if status == StatusCode::INTERNAL_SERVER_ERROR { + return; + } + + let is_startup_empty_membership = status == StatusCode::OK + && serde_json::from_slice::(&body) + .ok() + .and_then(|payload| { + payload + .get("gateways") + .and_then(|value| value.as_array()) + .map(|gateways| gateways.is_empty()) + }) + == Some(true); + + assert!( + is_startup_empty_membership, + "expected get_load to return 500 for missing gateway info; status: {status}, body: {}", + String::from_utf8_lossy(&body) + ); + + if attempt < MAX_GET_LOAD_ATTEMPTS { + sleep(Duration::from_millis(50)).await; + } + } + + panic!("get_load kept returning an empty gateway list before gateway info became observable"); } async fn read_gateway_last_update(h: &crate::support::TestHarness) -> u64 { diff --git a/tests/common/real_harness.rs b/tests/common/real_harness.rs index 9ac1fa0..6d25b63 100644 --- a/tests/common/real_harness.rs +++ b/tests/common/real_harness.rs @@ -195,6 +195,7 @@ pub(crate) struct GatewayHarnessOptions { pub(crate) taskmanager_result_lifetime_secs: Option, pub(crate) api_keys_update_interval_secs: Option, pub(crate) generic_key_concurrent_limit: Option, + pub(crate) trusted_proxy_cidrs: Vec, } impl Default for GatewayHarnessOptions { @@ -206,6 +207,7 @@ impl Default for GatewayHarnessOptions { taskmanager_result_lifetime_secs: None, api_keys_update_interval_secs: None, generic_key_concurrent_limit: None, + trusted_proxy_cidrs: Vec::new(), } } } @@ -542,6 +544,9 @@ impl GatewayRuntimeHarness { if let Some(generic_key_concurrent_limit) = options.generic_key_concurrent_limit { config.http.generic_key_concurrent_limit = generic_key_concurrent_limit; } + if !options.trusted_proxy_cidrs.is_empty() { + config.http.trusted_proxy_cidrs = options.trusted_proxy_cidrs; + } let generic_key = config.http.generic_key.expect("generic key"); let admin_key = config.http.admin_key;