Skip to content

Commit b920926

Browse files
AmyrAhmadyAmyrAhmady
committed
only rate limit if request is asking for the same thing
1 parent 24fc0eb commit b920926

File tree

1 file changed

+142
-111
lines changed

1 file changed

+142
-111
lines changed

‎src-tauri/src/query.rs‎

Lines changed: 142 additions & 111 deletions
Original file line numberDiff line numberDiff line change
@@ -14,10 +14,24 @@ use tokio::time::Instant;
1414

1515
use crate::{constants::*, errors::*, helpers};
1616

17+
#[derive(Clone, PartialEq)]
18+
struct QueryParams {
19+
info: bool,
20+
extra_info: bool,
21+
players: bool,
22+
rules: bool,
23+
ping: bool,
24+
}
25+
26+
struct RateLimitEntry {
27+
timestamp: u64,
28+
params: QueryParams,
29+
}
30+
1731
static OMP_EXTRA_INFO_LAST_UPDATE_LIST: Lazy<Mutex<HashMap<String, u64>>> =
1832
Lazy::new(|| Mutex::new(HashMap::new()));
1933

20-
static QUERY_RATE_LIMIT_LIST: Lazy<Mutex<HashMap<String, u64>>> =
34+
static QUERY_RATE_LIMIT_LIST: Lazy<Mutex<HashMap<String, RateLimitEntry>>> =
2135
Lazy::new(|| Mutex::new(HashMap::new()));
2236

2337
static CACHED_QUERY: Lazy<tokio::sync::Mutex<Option<(Query, String)>>> =
@@ -430,16 +444,34 @@ pub async fn query_server(
430444
};
431445

432446
let key = format!("{}:{}", ip, port);
447+
let current_params = QueryParams {
448+
info,
449+
extra_info,
450+
players,
451+
rules,
452+
ping,
453+
};
433454

434455
let should_allow = {
435456
let mut map = match QUERY_RATE_LIMIT_LIST.lock() {
436457
Ok(guard) => guard,
437458
Err(poisoned) => poisoned.into_inner(),
438459
};
439460
match map.get(&key) {
440-
Some(&last_time) if now - last_time < QUERY_RATE_LIMIT_MS => false,
461+
Some(entry)
462+
if now - entry.timestamp < QUERY_RATE_LIMIT_MS
463+
&& entry.params == current_params =>
464+
{
465+
false
466+
}
441467
_ => {
442-
map.insert(key.clone(), now);
468+
map.insert(
469+
key.clone(),
470+
RateLimitEntry {
471+
timestamp: now,
472+
params: current_params,
473+
},
474+
);
443475
true
444476
}
445477
}
@@ -470,52 +502,104 @@ pub async fn query_server(
470502
}
471503
};
472504

473-
let result = {
474-
let mut result = ServerQueryResponse {
475-
info: None,
476-
extra_info: None,
477-
players: None,
478-
rules: None,
479-
ping: None,
505+
let mut result = ServerQueryResponse {
506+
info: None,
507+
extra_info: None,
508+
players: None,
509+
rules: None,
510+
ping: None,
511+
};
512+
513+
if info {
514+
let _ = q.send('i').await;
515+
result.info = Some(match q.recv().await {
516+
Ok(p) => format!("{}", p),
517+
Err(e) => {
518+
let error_details = ErrorResponse {
519+
error: true,
520+
info: e.to_string(),
521+
};
522+
523+
serde_json::to_string(&error_details).unwrap_or_else(|_| {
524+
r#"{"error":true,"info":"Server information query failed"}""#.to_string()
525+
})
526+
}
527+
});
528+
}
529+
530+
if players {
531+
let _ = q.send('c').await;
532+
result.players = Some(match q.recv().await {
533+
Ok(p) => format!("{}", p),
534+
Err(e) => {
535+
let error_details = ErrorResponse {
536+
error: true,
537+
info: e.to_string(),
538+
};
539+
540+
serde_json::to_string(&error_details).unwrap_or_else(|_| {
541+
r#"{"error":true,"info":"Server players query failed"}""#.to_string()
542+
})
543+
}
544+
});
545+
}
546+
547+
if rules {
548+
let _ = q.send('r').await;
549+
result.rules = Some(match q.recv().await {
550+
Ok(p) => format!("{}", p),
551+
Err(e) => {
552+
let error_details = ErrorResponse {
553+
error: true,
554+
info: e.to_string(),
555+
};
556+
557+
serde_json::to_string(&error_details).unwrap_or_else(|_| {
558+
r#"{"error":true,"info":"Server rules query failed"}""#.to_string()
559+
})
560+
}
561+
});
562+
}
563+
564+
if extra_info {
565+
let now_secs = match SystemTime::now().duration_since(UNIX_EPOCH) {
566+
Ok(duration) => duration.as_secs(),
567+
Err(e) => {
568+
let error_details = ErrorResponse {
569+
error: true,
570+
info: format!("System time error: {}", e),
571+
};
572+
return Ok(serde_json::to_string(&error_details).unwrap_or_else(|_| {
573+
r#"{"error":true,"info":"System time error"}"#.to_string()
574+
}));
575+
}
480576
};
481577

482-
if info {
483-
let _ = q.send('i').await;
484-
result.info = Some(match q.recv().await {
485-
Ok(p) => format!("{}", p),
486-
Err(e) => {
487-
let error_details = ErrorResponse {
488-
error: true,
489-
info: e.to_string(),
490-
};
578+
let key = format!("{}:{}", ip, port);
491579

492-
serde_json::to_string(&error_details).unwrap_or_else(|_| {
493-
r#"{"error":true,"info":"Server information query failed"}""#.to_string()
494-
})
580+
let should_request = {
581+
let mut map = match OMP_EXTRA_INFO_LAST_UPDATE_LIST.lock() {
582+
Ok(guard) => guard,
583+
Err(poisoned) => {
584+
// Recover from poisoned mutex by getting the data anyway
585+
poisoned.into_inner()
495586
}
496-
});
497-
}
498-
499-
if players {
500-
let _ = q.send('c').await;
501-
result.players = Some(match q.recv().await {
502-
Ok(p) => format!("{}", p),
503-
Err(e) => {
504-
let error_details = ErrorResponse {
505-
error: true,
506-
info: e.to_string(),
507-
};
587+
};
508588

509-
serde_json::to_string(&error_details).unwrap_or_else(|_| {
510-
r#"{"error":true,"info":"Server players query failed"}""#.to_string()
511-
})
589+
match map.get(&key) {
590+
Some(&last_time) if now_secs - last_time < OMP_EXTRA_INFO_UPDATE_COOLDOWN_SECS => {
591+
false
512592
}
513-
});
514-
}
593+
_ => {
594+
map.insert(key.clone(), now_secs);
595+
true
596+
}
597+
}
598+
};
515599

516-
if rules {
517-
let _ = q.send('r').await;
518-
result.rules = Some(match q.recv().await {
600+
if should_request {
601+
let _ = q.send('o').await;
602+
result.extra_info = Some(match q.recv().await {
519603
Ok(p) => format!("{}", p),
520604
Err(e) => {
521605
let error_details = ErrorResponse {
@@ -524,89 +608,36 @@ pub async fn query_server(
524608
};
525609

526610
serde_json::to_string(&error_details).unwrap_or_else(|_| {
527-
r#"{"error":true,"info":"Server rules query failed"}""#.to_string()
611+
r#"{"error":true,"info":"Server open.mp information query failed"}""#
612+
.to_string()
528613
})
529614
}
530615
});
531616
}
617+
}
532618

533-
if extra_info {
534-
let now_secs = match SystemTime::now().duration_since(UNIX_EPOCH) {
535-
Ok(duration) => duration.as_secs(),
536-
Err(e) => {
537-
let error_details = ErrorResponse {
538-
error: true,
539-
info: format!("System time error: {}", e),
540-
};
541-
return Ok(serde_json::to_string(&error_details).unwrap_or_else(|_| {
542-
r#"{"error":true,"info":"System time error"}"#.to_string()
543-
}));
544-
}
545-
};
546-
547-
let key = format!("{}:{}", ip, port);
548-
549-
let should_request = {
550-
let mut map = match OMP_EXTRA_INFO_LAST_UPDATE_LIST.lock() {
551-
Ok(guard) => guard,
552-
Err(poisoned) => {
553-
// Recover from poisoned mutex by getting the data anyway
554-
poisoned.into_inner()
555-
}
556-
};
557-
558-
match map.get(&key) {
559-
Some(&last_time)
560-
if now_secs - last_time < OMP_EXTRA_INFO_UPDATE_COOLDOWN_SECS =>
561-
{
562-
false
563-
}
564-
_ => {
565-
map.insert(key.clone(), now_secs);
566-
true
567-
}
568-
}
569-
};
570-
571-
if should_request {
572-
let _ = q.send('o').await;
573-
result.extra_info = Some(match q.recv().await {
574-
Ok(p) => format!("{}", p),
575-
Err(e) => {
576-
let error_details = ErrorResponse {
577-
error: true,
578-
info: e.to_string(),
579-
};
580-
581-
serde_json::to_string(&error_details).unwrap_or_else(|_| {
582-
r#"{"error":true,"info":"Server open.mp information query failed"}""#.to_string()
583-
})
584-
}
585-
});
619+
if ping {
620+
let _ = q.send('p').await;
621+
let before = Instant::now();
622+
match q.recv().await {
623+
Ok(_p) => {
624+
result.ping = Some(before.elapsed().as_millis() as u32);
586625
}
587-
}
588-
589-
if ping {
590-
let _ = q.send('p').await;
591-
let before = Instant::now();
592-
match q.recv().await {
593-
Ok(_p) => {
594-
result.ping = Some(before.elapsed().as_millis() as u32);
595-
}
596-
Err(_) => {
597-
result.ping = Some(PING_TIMEOUT);
598-
}
626+
Err(_) => {
627+
result.ping = Some(PING_TIMEOUT);
599628
}
600629
}
601-
};
630+
}
602631

603632
// Store the query back in cache for next use
604633
{
605634
let mut cache = CACHED_QUERY.lock().await;
606635
*cache = Some((q, key));
607636
}
608637

609-
Ok(serde_json::to_string(&result).unwrap_or_else(|_| {
638+
let ret = serde_json::to_string(&result).unwrap_or_else(|_| {
610639
r#"{"error":true,"info":"Information serialization failed"}""#.to_string()
611-
}))
640+
});
641+
642+
Ok(ret)
612643
}

0 commit comments

Comments
 (0)