diff --git a/Cargo.lock b/Cargo.lock index ca3e47878f1d84..ab8261e80c13c7 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -1723,6 +1723,7 @@ dependencies = [ "clap 4.5.21", "deadpool-redis", "enum_dispatch", + "firebase-token", "futures", "hex", "ipnet", @@ -8201,6 +8202,18 @@ version = "1.2.0" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "8fcfdc7a0362c9f4444381a9e697c79d435fe65b52a37466fc2c1184cee9edc6" +[[package]] +name = "firebase-token" +version = "0.3.0" +source = "git+https://github.com/geekflyer/firebase-token?rev=34ea512d3d1fad6c11df3e7d82ff72beccc05836#34ea512d3d1fad6c11df3e7d82ff72beccc05836" +dependencies = [ + "jsonwebtoken 8.3.0", + "reqwest 0.11.23", + "serde", + "tokio", + "tracing", +] + [[package]] name = "firestore" version = "0.43.0" diff --git a/Cargo.toml b/Cargo.toml index cc429d5c1435df..d5b901956d8563 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -593,6 +593,7 @@ fail = "0.5.0" ff = { version = "0.13", features = ["derive"] } field_count = "0.1.1" file_diff = "1.0.0" +firebase-token = { git = "https://github.com/geekflyer/firebase-token", rev = "34ea512d3d1fad6c11df3e7d82ff72beccc05836" } firestore = "0.43.0" fixed = "1.25.1" flate2 = "1.0.24" diff --git a/crates/aptos-faucet/core/Cargo.toml b/crates/aptos-faucet/core/Cargo.toml index 7d3ffa848d9fbf..a078e9c95dec00 100644 --- a/crates/aptos-faucet/core/Cargo.toml +++ b/crates/aptos-faucet/core/Cargo.toml @@ -24,6 +24,7 @@ captcha = { version = "0.0.9" } clap = { workspace = true } deadpool-redis = { version = "0.11.1", features = ["rt_tokio_1"], default-features = false } enum_dispatch = { workspace = true } +firebase-token = { workspace = true } futures = { workspace = true } hex = { workspace = true } ipnet = { workspace = true } diff --git a/crates/aptos-faucet/core/src/checkers/auth_token.rs b/crates/aptos-faucet/core/src/checkers/auth_token.rs index 6716343e01f4f6..2652edb9a94e90 100644 --- a/crates/aptos-faucet/core/src/checkers/auth_token.rs +++ b/crates/aptos-faucet/core/src/checkers/auth_token.rs @@ -5,6 +5,7 @@ use super::{CheckerData, CheckerTrait}; use crate::{ common::{ListManager, ListManagerConfig}, endpoints::{AptosTapError, RejectionReason, RejectionReasonCode}, + firebase_jwt::X_IS_JWT_HEADER, }; use anyhow::Result; use aptos_logger::info; @@ -33,6 +34,11 @@ impl CheckerTrait for AuthTokenChecker { data: CheckerData, _dry_run: bool, ) -> Result, AptosTapError> { + // Don't check if the request has X_IS_JWT_HEADER set. + if data.headers.contains_key(X_IS_JWT_HEADER) { + return Ok(vec![]); + } + let auth_token = match data .headers .get(AUTHORIZATION) diff --git a/crates/aptos-faucet/core/src/checkers/memory_ratelimit.rs b/crates/aptos-faucet/core/src/checkers/memory_ratelimit.rs index 3d2f4e969c9da3..26d8524aac366c 100644 --- a/crates/aptos-faucet/core/src/checkers/memory_ratelimit.rs +++ b/crates/aptos-faucet/core/src/checkers/memory_ratelimit.rs @@ -27,7 +27,7 @@ impl MemoryRatelimitCheckerConfig { } /// Simple in memory storage that rejects if we've ever seen a request from an -/// IP that has succeeded. +/// IP that has succeeded. This does not support JWT-based ratelimiting. pub struct MemoryRatelimitChecker { pub max_requests_per_day: u32, @@ -81,7 +81,7 @@ impl CheckerTrait for MemoryRatelimitChecker { "IP {} has exceeded the daily limit of {} requests", data.source_ip, self.max_requests_per_day ), - RejectionReasonCode::IpUsageLimitExhausted, + RejectionReasonCode::UsageLimitExhausted, )]); } else if !dry_run { *requests_today += 1; diff --git a/crates/aptos-faucet/core/src/checkers/redis_ratelimit.rs b/crates/aptos-faucet/core/src/checkers/redis_ratelimit.rs index 16f7fb7b68a98d..2b009560754855 100644 --- a/crates/aptos-faucet/core/src/checkers/redis_ratelimit.rs +++ b/crates/aptos-faucet/core/src/checkers/redis_ratelimit.rs @@ -4,16 +4,79 @@ use super::{CheckerData, CheckerTrait, CompleteData}; use crate::{ endpoints::{AptosTapError, AptosTapErrorCode, RejectionReason, RejectionReasonCode}, + firebase_jwt::{FirebaseJwtVerifier, FirebaseJwtVerifierConfig}, helpers::{days_since_tap_epoch, get_current_time_secs, seconds_until_next_day}, }; use anyhow::{Context, Result}; use async_trait::async_trait; use deadpool_redis::{ - redis::{AsyncCommands, ConnectionAddr, ConnectionInfo, RedisConnectionInfo}, + redis::{self, AsyncCommands, ConnectionAddr, ConnectionInfo, RedisConnectionInfo}, Config, Connection, Pool, Runtime, }; use serde::{Deserialize, Serialize}; -use std::net::IpAddr; + +#[derive(Clone, Debug, Default, Deserialize, Serialize)] +pub enum RatelimitKeyProviderConfig { + #[default] + Ip, + Jwt(FirebaseJwtVerifierConfig), +} + +impl RatelimitKeyProviderConfig { + pub fn ratelimit_key_prefix(&self) -> &'static str { + match self { + RatelimitKeyProviderConfig::Ip => "ip", + RatelimitKeyProviderConfig::Jwt(_) => "jwt", + } + } + + /* + /// If the faucet is configured to ratelimit by IP, this will be the client's IP + /// address. If the faucet is configured to ratelimit by JWT, this will be the + /// user's Firebase UID (taken from the JWT's `sub` field). We also return a prefix + /// for use in any ratelimiting key (e.g. for Redis). + pub async fn ratelimit_key_value(&self, data: &CheckerData) -> Result { + match self { + RatelimitingApproachConfig::Ip => Ok(data.source_ip.to_string()), + RatelimitingApproachConfig::Jwt(config) => { + let jwt_verifier = FirebaseJwtVerifier::new(config.clone()).await?; + jwt_verifier.validate_jwt(data.headers).await?.ok_or(AptosTapError::new( + "Failed to extract Firebase UID from JWT".to_string(), + AptosTapErrorCode::AuthTokenInvalid, + )) + } + } + } + */ +} + +/// This is what produces the key we use for ratelimiting in Redis. +pub enum RatelimitKeyProvider { + Ip, + Jwt(FirebaseJwtVerifier), +} + +impl RatelimitKeyProvider { + pub fn ratelimit_key_prefix(&self) -> &'static str { + match self { + RatelimitKeyProvider::Ip => "ip", + RatelimitKeyProvider::Jwt(_) => "jwt", + } + } + + /// If the faucet is configured to ratelimit by IP, this will be the client's IP + /// address. If the faucet is configured to ratelimit by JWT, we verify the JWT + /// first. If it is valid, this will be the user's Firebase UID (taken from the + /// JWT's `sub` field). + pub async fn ratelimit_key_value(&self, data: &CheckerData) -> Result { + match self { + RatelimitKeyProvider::Ip => Ok(data.source_ip.to_string()), + RatelimitKeyProvider::Jwt(jwt_verifier) => { + jwt_verifier.validate_jwt(data.headers.clone()).await + }, + } + } +} #[derive(Clone, Debug, Deserialize, Serialize)] pub struct RedisRatelimitCheckerConfig { @@ -35,9 +98,13 @@ pub struct RedisRatelimitCheckerConfig { /// The password of the given user, if necessary. pub database_password: Option, - /// Max number of requests per IP per day. 500s are not counted, because - /// they are not the user's fault, but everything else is. - pub max_requests_per_ip_per_day: u32, + /// Max number of requests per key per day. 500s are not counted, because they are + /// not the user's fault, but everything else is. + pub max_requests_per_day: u32, + + /// The way we ratelimit. + #[serde(default)] + pub ratelimit_key_provider_config: RatelimitKeyProviderConfig, } impl RedisRatelimitCheckerConfig { @@ -76,10 +143,12 @@ impl RedisRatelimitCheckerConfig { /// request. Instead, it uses counters to track limits. This is heavily inspired /// by https://redis.com/redis-best-practices/basic-rate-limiting/. /// +/// We use a generic key (e.g. IP address or Firebase UID). +/// /// If we're not careful, it is possible for people to exceed the intended limit -/// by sending many requests simulatenously. We avoid this problem with this +/// by sending many requests simultaneously. We avoid this problem with this /// order of operations: -/// 1. Read the current value of the limit for source IP. +/// 1. Read the current value of the limit for the given key (e.g. IP / Firebase UID). /// 2. If value is greater than limit, reject. /// 3. Otherwise, increment and set TTL if necessary. /// 4. Increment returns the new value. Check if this is greater than the limit also. @@ -92,18 +161,19 @@ impl RedisRatelimitCheckerConfig { /// they're under their limit vs more unnecessary writes when they're over their /// limit, but we'll happily take more reads over more writes. /// -/// Note: Previously I made an attempt (d4fbf6db675e9036a967b52bf8d13e1b2566787e) at +/// Note: Previously we made an attempt (d4fbf6db675e9036a967b52bf8d13e1b2566787e) at /// doing these steps atomically, but it became very unwieldy: /// 1. Start a transaction. -/// 2. Increment current value for limit for source IP, set TTL if necessary. +/// 2. Increment current value for limit for source key, set TTL if necessary. /// 3. If value is greater than limit, revert the transaction. /// /// This second way leaves a small window for someone to slip in multiple requests, -/// therein blowing past the configured limit, but it's a very small window, so -/// we'll worry about it as a followup: https://github.com/aptos-labs/aptos-tap/issues/15. +/// thereby blowing past the configured limit, but it's a very small window. We'll +/// worry about it as a follow-up: https://github.com/aptos-labs/aptos-tap/issues/15. pub struct RedisRatelimitChecker { args: RedisRatelimitCheckerConfig, db_pool: Pool, + ratelimit_key_provider: RatelimitKeyProvider, } impl RedisRatelimitChecker { @@ -116,7 +186,18 @@ impl RedisRatelimitChecker { .await .context("Failed to connect to redis on startup")?; - Ok(Self { args, db_pool }) + let ratelimit_key_provider = match args.ratelimit_key_provider_config.clone() { + RatelimitKeyProviderConfig::Ip => RatelimitKeyProvider::Ip, + RatelimitKeyProviderConfig::Jwt(config) => { + RatelimitKeyProvider::Jwt(FirebaseJwtVerifier::new(config).await?) + }, + }; + + Ok(Self { + args, + db_pool, + ratelimit_key_provider, + }) } pub async fn get_redis_connection(&self) -> Result { @@ -128,28 +209,37 @@ impl RedisRatelimitChecker { }) } - // Returns the key and the seconds until the next day. - fn get_key_and_secs_until_next_day(&self, source_ip: &IpAddr) -> (String, u64) { + // Returns the key and the seconds until the next day. This daily suffix ensures + // that we only count requests for the current day. + fn get_key_and_secs_until_next_day( + &self, + ratelimit_key_prefix: &str, + ratelimit_key_value: &str, + ) -> (String, u64) { let now_secs = get_current_time_secs(); let seconds_until_next_day = seconds_until_next_day(now_secs); - let key = format!("ip:{}:{}", source_ip, days_since_tap_epoch(now_secs)); + let key = format!( + "{}:{}:{}", + ratelimit_key_prefix, + ratelimit_key_value, + days_since_tap_epoch(now_secs) + ); (key, seconds_until_next_day) } fn check_limit_value( &self, - data: &CheckerData, limit_value: Option, seconds_until_next_day: u64, ) -> Option { - if limit_value.unwrap_or(0) > self.args.max_requests_per_ip_per_day as i64 { + if limit_value.unwrap_or(0) > self.args.max_requests_per_day as i64 { Some( RejectionReason::new( format!( - "IP {} has reached the maximum allowed number of requests per day: {}", - data.source_ip, self.args.max_requests_per_ip_per_day + "You have reached the maximum allowed number of requests per day: {}", + self.args.max_requests_per_day ), - RejectionReasonCode::IpUsageLimitExhausted, + RejectionReasonCode::UsageLimitExhausted, ) .retry_after(seconds_until_next_day), ) @@ -171,11 +261,17 @@ impl CheckerTrait for RedisRatelimitChecker { .await .map_err(|e| AptosTapError::new_with_error_code(e, AptosTapErrorCode::StorageError))?; - // Generate a key corresponding to this IP address and the current day. - let (key, seconds_until_next_day) = self.get_key_and_secs_until_next_day(&data.source_ip); + // Generate a key corresponding to this identifier and the current day. + let key_prefix = self.ratelimit_key_provider.ratelimit_key_prefix(); + let key_value = self + .ratelimit_key_provider + .ratelimit_key_value(&data) + .await?; + let (key, seconds_until_next_day) = + self.get_key_and_secs_until_next_day(key_prefix, &key_value); - // Get the value for the key, indicating how many non-500 requests we - // have serviced for this it today. + // Get the value for the key, indicating how many non-500 requests we have + // serviced for this key today. let limit_value: Option = conn.get(&key).await.map_err(|e| { AptosTapError::new_with_error_code( format!("Failed to get value for redis key {}: {}", key, e), @@ -183,18 +279,16 @@ impl CheckerTrait for RedisRatelimitChecker { ) })?; - // If the limit value is greater than what we allow per day, signal - // that we should reject this request. - if let Some(rejection_reason) = - self.check_limit_value(&data, limit_value, seconds_until_next_day) + // If the limit value is greater than what we allow per day, signal that we + // should reject this request. + if let Some(rejection_reason) = self.check_limit_value(limit_value, seconds_until_next_day) { return Ok(vec![rejection_reason]); } - // Atomically increment the counter for the given IP, creating it and - // setting the expiration time if it doesn't already exist. + // Atomically increment the counter for the given key, creating it and setting + // the expiration time if it doesn't already exist. if !dry_run { - // If the limit value already exists, just increment. let incremented_limit_value = match limit_value { Some(_) => conn.incr(&key, 1).await.map_err(|e| { AptosTapError::new_with_error_code( @@ -202,17 +296,12 @@ impl CheckerTrait for RedisRatelimitChecker { AptosTapErrorCode::StorageError, ) })?, - // If the limit value doesn't exist, create it and set the - // expiration time. None => { let (incremented_limit_value,): (i64,) = redis::pipe() .atomic() .incr(&key, 1) // Expire at the end of the day roughly. .expire(&key, seconds_until_next_day as usize) - // Only set the expiration if one isn't already set. - // Only works with Redis 7 sadly. - // .arg("NX") .ignore() .query_async(&mut *conn) .await @@ -226,9 +315,9 @@ impl CheckerTrait for RedisRatelimitChecker { }, }; - // Check limit again, to ensure there wasn't a get / set race. + // Check limit again, to ensure there wasn't a get/set race. if let Some(rejection_reason) = - self.check_limit_value(&data, Some(incremented_limit_value), seconds_until_next_day) + self.check_limit_value(Some(incremented_limit_value), seconds_until_next_day) { return Ok(vec![rejection_reason]); } @@ -237,8 +326,8 @@ impl CheckerTrait for RedisRatelimitChecker { Ok(vec![]) } - /// All we have to do here is decrement the counter if the request was a - /// failure due to something wrong on our end. + /// All we have to do here is decrement the counter if the request was a failure due + /// to something wrong on our end. async fn complete(&self, data: CompleteData) -> Result<(), AptosTapError> { if !data.response_is_500 { return Ok(()); @@ -249,8 +338,15 @@ impl CheckerTrait for RedisRatelimitChecker { .await .map_err(|e| AptosTapError::new_with_error_code(e, AptosTapErrorCode::StorageError))?; - // Generate a key corresponding to this IP address and the current day. - let (key, _) = self.get_key_and_secs_until_next_day(&data.checker_data.source_ip); + // Generate a key corresponding to this identifier and the current day. In the + // JWT case we re-verify the JWT. This is inefficient, but these failures are + // extremely rare so I don't refactor for now. + let key_prefix = self.ratelimit_key_provider.ratelimit_key_prefix(); + let key_value = self + .ratelimit_key_provider + .ratelimit_key_value(&data.checker_data) + .await?; + let (key, _) = self.get_key_and_secs_until_next_day(key_prefix, &key_value); conn.decr(&key, 1).await.map_err(|e| { AptosTapError::new_with_error_code( @@ -262,6 +358,6 @@ impl CheckerTrait for RedisRatelimitChecker { } fn cost(&self) -> u8 { - 50 + 100 } } diff --git a/crates/aptos-faucet/core/src/endpoints/errors.rs b/crates/aptos-faucet/core/src/endpoints/errors.rs index db18257a73b8c9..9caab60b0e1115 100644 --- a/crates/aptos-faucet/core/src/endpoints/errors.rs +++ b/crates/aptos-faucet/core/src/endpoints/errors.rs @@ -64,7 +64,7 @@ impl AptosTapError { pub fn status_and_retry_after(&self) -> (StatusCode, Option) { let (mut status_code, mut retry_after) = (self.error_code.status(), None); for rejection_reason in &self.rejection_reasons { - if rejection_reason.code == RejectionReasonCode::IpUsageLimitExhausted { + if rejection_reason.code == RejectionReasonCode::UsageLimitExhausted { status_code = StatusCode::TOO_MANY_REQUESTS; retry_after = rejection_reason.retry_after; break; @@ -134,6 +134,9 @@ pub enum AptosTapErrorCode { /// The user tried to call an endpoint that is not enabled. EndpointNotEnabled = 45, + /// The user provided an invalid auth token. + AuthTokenInvalid = 46, + /// Failed when making requests to the Aptos API. AptosApiError = 50, @@ -170,7 +173,8 @@ impl AptosTapErrorCode { | AptosTapErrorCode::EndpointNotEnabled => StatusCode::BAD_REQUEST, AptosTapErrorCode::Rejected | AptosTapErrorCode::SourceIpMissing - | AptosTapErrorCode::TransactionFailed => StatusCode::FORBIDDEN, + | AptosTapErrorCode::TransactionFailed + | AptosTapErrorCode::AuthTokenInvalid => StatusCode::FORBIDDEN, AptosTapErrorCode::AptosApiError | AptosTapErrorCode::TransactionTimedOut | AptosTapErrorCode::SerializationError @@ -233,8 +237,8 @@ pub enum RejectionReasonCode { /// Account already has funds. AccountAlreadyExists = 100, - /// IP has exhausted its usage limit. - IpUsageLimitExhausted = 101, + /// Key (IP / Firebase UID) has exhausted its usage limit. + UsageLimitExhausted = 101, /// IP is in the blocklist. IpInBlocklist = 102, diff --git a/crates/aptos-faucet/core/src/endpoints/fund.rs b/crates/aptos-faucet/core/src/endpoints/fund.rs index f43ded751a9711..c916e00f63390a 100644 --- a/crates/aptos-faucet/core/src/endpoints/fund.rs +++ b/crates/aptos-faucet/core/src/endpoints/fund.rs @@ -9,6 +9,7 @@ use crate::{ bypasser::{Bypasser, BypasserTrait}, checkers::{Checker, CheckerData, CheckerTrait, CompleteData}, endpoints::AptosTapErrorCode, + firebase_jwt::jwt_sub, funder::{Funder, FunderTrait}, helpers::{get_current_time_secs, transaction_hashes}, }; @@ -307,6 +308,7 @@ impl FundApiComponents { // Include some additional logging that the logging middleware doesn't do. info!( source_ip = checker_data.source_ip, + jwt_sub = jwt_sub(checker_data.headers.clone()).ok(), address = checker_data.receiver, requested_amount = fund_request.amount, txn_hashes = txn_hashes, diff --git a/crates/aptos-faucet/core/src/firebase_jwt.rs b/crates/aptos-faucet/core/src/firebase_jwt.rs new file mode 100644 index 00000000000000..5bcfd7d1bf7a5e --- /dev/null +++ b/crates/aptos-faucet/core/src/firebase_jwt.rs @@ -0,0 +1,103 @@ +// Copyright © Aptos Foundation +// SPDX-License-Identifier: Apache-2.0 + +use crate::endpoints::{AptosTapError, AptosTapErrorCode}; +use anyhow::Result; +use firebase_token::JwkAuth; +use poem::http::{header::AUTHORIZATION, HeaderMap}; +use serde::{Deserialize, Serialize}; +use std::sync::Arc; + +pub const X_IS_JWT_HEADER: &str = "x-is-jwt"; + +#[derive(Clone, Debug, Deserialize, Serialize)] +pub struct FirebaseJwtVerifierConfig { + pub identity_platform_gcp_project: String, +} + +/// This verifies that the value in the Authorization header is a valid Firebase JWT. +/// Since we already have achecker that looks for API keys using the Authorization +/// header, we mandate that a `x-is-jwt` header is present as well. +pub struct FirebaseJwtVerifier { + pub jwt_verifier: JwkAuth, +} + +impl FirebaseJwtVerifier { + pub async fn new(config: FirebaseJwtVerifierConfig) -> Result { + let jwt_verifier = JwkAuth::new(config.identity_platform_gcp_project).await; + Ok(Self { jwt_verifier }) + } + + /// First, we mandate that the caller indicated that they're including a JWT by + /// checking for the presence of X_IS_JWT_HEADER. If they didn't include this + /// header, we reject them immediately. + /// + /// If they did include X_IS_JWT_HEADER and the Authorization header was present + /// and well-formed, we extract the token from the Authorization header and verify + /// it with Firebase. If the token is invalid, we reject them. If it is valid, we + /// return the UID (from the sub field). + pub async fn validate_jwt(&self, headers: Arc) -> Result { + let auth_token = jwt_sub(headers)?; + + let verify = self.jwt_verifier.verify::(&auth_token); + let token_data = match verify.await { + Some(token_data) => token_data, + None => { + return Err(AptosTapError::new( + "Failed to verify JWT token".to_string(), + AptosTapErrorCode::AuthTokenInvalid, + )); + }, + }; + let claims = token_data.claims; + + if !claims.email_verified { + return Err(AptosTapError::new( + "The JWT token is not verified".to_string(), + AptosTapErrorCode::AuthTokenInvalid, + )); + } + + Ok(claims.sub) + } +} + +/// Returns the sub field from a JWT if it is present (the Firebase UID). +pub fn jwt_sub(headers: Arc) -> Result { + headers + .get(X_IS_JWT_HEADER) + .and_then(|v| v.to_str().ok()) + .map(|v| v.eq_ignore_ascii_case("true")) + .ok_or_else(|| { + AptosTapError::new( + format!( + "The {} header must be present and set to 'true'", + X_IS_JWT_HEADER + ), + AptosTapErrorCode::AuthTokenInvalid, + ) + })?; + + match headers + .get(AUTHORIZATION) + .and_then(|v| v.to_str().ok()) + .and_then(|v| v.split_whitespace().nth(1)) +{ + Some(auth_token) => Ok(auth_token.to_string()), + None => Err(AptosTapError::new( + "Either the Authorization header is missing or it is not in the form of 'Bearer '".to_string(), + AptosTapErrorCode::AuthTokenInvalid, + )), +} +} + +#[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] +pub struct JwtClaims { + pub aud: String, + pub exp: i64, + pub iss: String, + pub sub: String, + pub iat: i64, + pub email: String, + pub email_verified: bool, +} diff --git a/crates/aptos-faucet/core/src/lib.rs b/crates/aptos-faucet/core/src/lib.rs index 65e2b0b8b832e2..382d1a24b83c28 100644 --- a/crates/aptos-faucet/core/src/lib.rs +++ b/crates/aptos-faucet/core/src/lib.rs @@ -5,6 +5,7 @@ pub mod bypasser; pub mod checkers; pub mod common; pub mod endpoints; +pub mod firebase_jwt; pub mod funder; pub mod helpers; pub mod middleware;