diff --git a/crates/jmap/src/auth/oauth/mod.rs b/crates/jmap/src/auth/oauth/mod.rs index 6e30e90a..f0632c79 100644 --- a/crates/jmap/src/auth/oauth/mod.rs +++ b/crates/jmap/src/auth/oauth/mod.rs @@ -143,13 +143,13 @@ pub enum TokenResponse { #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] pub struct OAuthResponse { - access_token: String, - token_type: String, - expires_in: u64, + pub access_token: String, + pub token_type: String, + pub expires_in: u64, #[serde(skip_serializing_if = "Option::is_none")] - refresh_token: Option, + pub refresh_token: Option, #[serde(skip_serializing_if = "Option::is_none")] - scope: Option, + pub scope: Option, } #[derive(Debug, Serialize, Deserialize, PartialEq, Eq)] diff --git a/crates/smtp/src/core/eval.rs b/crates/smtp/src/core/eval.rs index c0479052..485f1158 100644 --- a/crates/smtp/src/core/eval.rs +++ b/crates/smtp/src/core/eval.rs @@ -240,7 +240,7 @@ impl SMTP { let value = params.next_as_integer(); self.get_lookup_store(store.as_ref()) - .counter_incr(key.into_owned().into_bytes(), value, None) + .counter_incr(key.into_owned().into_bytes(), value, None, true) .await .map(Variable::Integer) .unwrap_or_else(|err| { diff --git a/crates/smtp/src/scripts/plugins/bayes.rs b/crates/smtp/src/scripts/plugins/bayes.rs index 8b035ceb..9e8df73d 100644 --- a/crates/smtp/src/scripts/plugins/bayes.rs +++ b/crates/smtp/src/scripts/plugins/bayes.rs @@ -117,6 +117,7 @@ fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable { .finalize(), weights.into(), None, + false, ), ) .is_err() @@ -141,6 +142,7 @@ fn train(ctx: PluginContext<'_>, is_train: bool) -> Variable { .finalize(), weights.into(), None, + false, ), ) .is_err() diff --git a/crates/store/src/dispatch/lookup.rs b/crates/store/src/dispatch/lookup.rs index ab3eed28..b357a3ba 100644 --- a/crates/store/src/dispatch/lookup.rs +++ b/crates/store/src/dispatch/lookup.rs @@ -99,9 +99,21 @@ impl LookupStore { key: Vec, value: i64, expires: Option, + return_value: bool, ) -> crate::Result { match self { LookupStore::Store(store) => { + let result = if return_value { + store + .get_counter(ValueKey::from(ValueClass::Lookup(LookupClass::Counter( + key.clone(), + )))) + .await? + + 1 + } else { + 0 + }; + let mut batch = BatchBuilder::new(); if let Some(expires) = expires { @@ -122,7 +134,7 @@ impl LookupStore { store.write(batch.build()).await?; - Ok(0) + Ok(result) } #[cfg(feature = "redis")] LookupStore::Redis(store) => store.key_incr(key, value, expires).await, @@ -269,16 +281,8 @@ impl LookupStore { bucket.extend_from_slice(range_start.to_be_bytes().as_slice()); let requests = if !soft_check { - let requests = self.counter_incr(bucket, 1, expires_in.into()).await?; - if requests > 0 { - requests - } else { - // Increment and get not supported by store, fetch counter - let mut bucket = Vec::with_capacity(key.len() + U64_LEN); - bucket.extend_from_slice(key); - bucket.extend_from_slice(range_start.to_be_bytes().as_slice()); - self.counter_get(bucket).await? - } + self.counter_incr(bucket, 1, expires_in.into(), true) + .await? } else { self.counter_get(bucket).await? + 1 }; diff --git a/tests/src/jmap/auth_oauth.rs b/tests/src/jmap/auth_oauth.rs index c1038003..8a69ab51 100644 --- a/tests/src/jmap/auth_oauth.rs +++ b/tests/src/jmap/auth_oauth.rs @@ -434,15 +434,13 @@ fn parse_code_redirect(uri: String, state: &str) -> String { fn unwrap_token_response(response: TokenResponse) -> (String, Option, u64) { match response { - TokenResponse::Granted { - access_token, - token_type, - expires_in, - refresh_token, - .. - } => { - assert_eq!(token_type, "bearer"); - (access_token, refresh_token, expires_in) + TokenResponse::Granted(granted) => { + assert_eq!(granted.token_type, "bearer"); + ( + granted.access_token, + granted.refresh_token, + granted.expires_in, + ) } TokenResponse::Error { error } => panic!("Expected granted, got {:?}", error), } diff --git a/tests/src/smtp/lookup/sql.rs b/tests/src/smtp/lookup/sql.rs index 39ccdb37..905c4714 100644 --- a/tests/src/smtp/lookup/sql.rs +++ b/tests/src/smtp/lookup/sql.rs @@ -179,7 +179,7 @@ async fn lookup_sql() { "counter_get('sql', 'county') + '-' + counter_incr('sql', 'county', 1) + '-' ", "+ counter_incr('sql', 'county', 1) + '-' + counter_get('sql', 'county')" ), - "0-0-0-2", + "0-1-2-2", ), ] { let e = Expression::parse("test", expr, |name| { diff --git a/tests/src/store/lookup.rs b/tests/src/store/lookup.rs index 0d97ccf1..fb5ba695 100644 --- a/tests/src/store/lookup.rs +++ b/tests/src/store/lookup.rs @@ -82,16 +82,28 @@ pub async fn lookup_tests() { // Test counter let key = "abc".as_bytes().to_vec(); - store.counter_incr(key.clone(), 1, None).await.unwrap(); + store + .counter_incr(key.clone(), 1, None, false) + .await + .unwrap(); assert_eq!(1, store.counter_get(key.clone()).await.unwrap()); - store.counter_incr(key.clone(), 2, None).await.unwrap(); + store + .counter_incr(key.clone(), 2, None, false) + .await + .unwrap(); assert_eq!(3, store.counter_get(key.clone()).await.unwrap()); - store.counter_incr(key.clone(), -3, None).await.unwrap(); + store + .counter_incr(key.clone(), -3, None, false) + .await + .unwrap(); assert_eq!(0, store.counter_get(key.clone()).await.unwrap()); // Test counter expiry let key = "fgh".as_bytes().to_vec(); - store.counter_incr(key.clone(), 1, 1.into()).await.unwrap(); + store + .counter_incr(key.clone(), 1, 1.into(), false) + .await + .unwrap(); assert_eq!(1, store.counter_get(key.clone()).await.unwrap()); tokio::time::sleep(tokio::time::Duration::from_secs(1)).await; store.purge_expired().await.unwrap();