From 3ab3c478c34bb952bf35b5e605829cbcbb952a2c Mon Sep 17 00:00:00 2001
From: Peter Cai <peter@typeblog.net>
Date: Sun, 4 Apr 2021 11:29:21 +0800
Subject: [PATCH] util: implement generic random_range

---
 src/client.rs |  4 ++--
 src/util.rs   | 18 ++++++++++++++++++
 2 files changed, 20 insertions(+), 2 deletions(-)

diff --git a/src/client.rs b/src/client.rs
index 6c049e6..442c6f7 100644
--- a/src/client.rs
+++ b/src/client.rs
@@ -50,7 +50,7 @@ impl Client {
 
     // Select an upstream randomly
     fn select_upstream(&self) -> String {
-        let idx = crate::util::random() * self.options.upstream_urls.len() as f64;
+        let idx = crate::util::random_range(0, self.options.upstream_urls.len() as u16);
         self.options.upstream_urls[idx as usize].clone()
     }
 
@@ -61,7 +61,7 @@ impl Client {
         let mut builder = MessageBuilder::new_udp();
         // Set up the header
         let header = builder.header_mut();
-        header.set_id((crate::util::random() * u16::MAX as f64) as u16);
+        header.set_id(crate::util::random_range(0, u16::MAX));
         header.set_qr(false); // For queries, QR = false
         header.set_opcode(Opcode::Query);
         header.set_rd(true); // Ask for recursive queries
diff --git a/src/util.rs b/src/util.rs
index ff432a6..cf1cd5e 100644
--- a/src/util.rs
+++ b/src/util.rs
@@ -1,6 +1,7 @@
 use bytes::Bytes;
 use domain_core::bits::message::Message;
 use js_sys::{Math, Promise};
+use std::ops::Add;
 use wasm_bindgen::prelude::*;
 use wasm_bindgen::JsValue;
 use wasm_bindgen_futures::JsFuture;
@@ -32,3 +33,20 @@ pub async fn fetch_rs(req: &Request) -> Result<JsValue, JsValue> {
 pub fn random() -> f64 {
     unsafe { Math::random() }
 }
+
+pub fn random_range<T>(min: T, max: T) -> T
+where
+    T: Ord + Into<f64> + FromFloat<f64> + Add<Output = T>,
+{
+    min + T::from_float(random() * max.into())
+}
+
+pub trait FromFloat<F> {
+    fn from_float(f: F) -> Self;
+}
+
+impl FromFloat<f64> for u16 {
+    fn from_float(f: f64) -> u16 {
+        f as u16
+    }
+}