Added response validity check

src/main.rs

@@ -15,7 +15,7 @@ mod stun;
 #[tokio::main(flavor = "current_thread")]
 async fn main() -> io::Result<()> {
     let stream = candidates::get_candidates().await?.into_iter();
-    let futs = stream
+    let profiles = stream
         .map(|candidate| async move {
             let res = stun::check_candidate(candidate.clone()).await;
             match &res {
@@ -25,34 +25,39 @@ async fn main() -> io::Result<()> {
             res
         })
         .collect::<Vec<_>>();
-    let zz = Instant::now();
-    let results = futures::future::join_all(futs).await;
+    let timestamp = Instant::now();
+    let profiles = futures::future::join_all(profiles).await;
     let mut all_ok = 0;
     let mut dns_unresolved = 0;
     let mut partial_timeout = 0;
     let mut complete_timeout = 0;
-    results.iter().for_each(|res| {
+    let mut incorrect_mapping_returned = 0;
+    profiles.iter().for_each(|res| {
         match res {
             Ok(_) => { all_ok += 1; }
             Err(CheckError::DnsResolutionFailed) => { dns_unresolved += 1; }
             Err(CheckError::PartialTimeout) => { partial_timeout += 1; }
             Err(CheckError::Timeout) => { complete_timeout += 1; }
+            Err(CheckError::IncorrectMappingReturned) => { incorrect_mapping_returned += 1; }
         }
     });
-    println!("OK {} , DNS failure {} , p/Timeout {} , Timeout {}", all_ok, dns_unresolved, partial_timeout, complete_timeout);
+    println!(
+        "OK {} , DNS failure {} , p/Timeout {} , Timeout {} , Incorrect {}",
+        all_ok, dns_unresolved, partial_timeout, complete_timeout, incorrect_mapping_returned
+    );
 
-    let mut output = results.iter()
+    let mut output_hosts = profiles.iter()
         .filter_map(|res| res.as_ref().ok())
         .map(|profile| profile.candidate.clone())
         .collect::<Vec<_>>();
-    output.shuffle(&mut thread_rng());
-    let output_hosts = output.into_iter()
+    output_hosts.shuffle(&mut thread_rng());
+    let output_hosts = output_hosts.into_iter()
         .map(|candidate| String::from(candidate))
         .reduce(|a, b| format!("{}\n{}", a, b))
         .unwrap_or(String::from(""));
     tokio::fs::write("valid_hosts.txt", output_hosts).await?;
 
-    let output_ip4 = results.iter()
+    let output_ip4 = profiles.iter()
         .filter_map(|res| res.as_ref().ok())
         .flat_map(|profile| profile.addrs.clone().into_iter())
         .filter(|addr| addr.is_ipv4())
@@ -66,7 +71,7 @@ async fn main() -> io::Result<()> {
         .unwrap_or(String::from(""));
     tokio::fs::write("valid_ipv4s.txt", output_ip4).await?;
 
-    let output_ip6 = results.iter()
+    let output_ip6 = profiles.iter()
         .filter_map(|res| res.as_ref().ok())
         .flat_map(|profile| profile.addrs.clone().into_iter())
         .filter(|addr| addr.is_ipv6())
@@ -80,6 +85,6 @@ async fn main() -> io::Result<()> {
         .unwrap_or(String::from(""));
     tokio::fs::write("valid_ipv6s.txt", output_ip6).await?;
 
-    println!("Finished in {:?}", zz.elapsed());
+    println!("Finished in {:?}", timestamp.elapsed());
     Ok(())
 }

src/stun.rs

@@ -6,18 +6,19 @@ use tokio::io;
 use tokio::net::{lookup_host};
 use crate::candidates::StunCandidate;
 
-#[derive(Debug)]
+#[derive(Clone, Copy, Debug, PartialEq)]
 pub enum CheckError {
     DnsResolutionFailed,
+    IncorrectMappingReturned,
     PartialTimeout,
-    Timeout
+    Timeout,
 }
 
 #[derive(Debug)]
 pub struct CandidateProfile {
     pub candidate: StunCandidate,
     pub addrs: Vec<SocketAddr>,
-    pub rtt_ms: u32
+    pub rtt_ms: u32,
 }
 
 pub async fn check_candidate(candidate: StunCandidate) -> Result<CandidateProfile, CheckError> {
@@ -30,38 +31,58 @@ pub async fn check_candidate(candidate: StunCandidate) -> Result<CandidateProfil
             SocketAddr::V4(..) => { "0.0.0.0:0" }
             SocketAddr::V6(..) => { "[::]:0" }
         }.parse::<std::net::SocketAddr>().unwrap()).await.unwrap();
+        let local_address = u.local_addr().expect("Local address to be available");
         let mut client = stunclient::StunClient::new(*address);
         client.set_timeout(Duration::from_secs(1));
         let time = Instant::now();
         let res = client.query_external_address_async(&u).await;
-        (res, time.elapsed())
+        match res {
+            Ok(mapped_address) => if mapped_address.port() == local_address.port() {
+                Ok((mapped_address, time.elapsed()))
+            } else {
+                Err(CheckError::IncorrectMappingReturned)
+            },
+            Err(_) => { Err(CheckError::Timeout) }
+        }
     }).collect::<Vec<_>>();
 
     let responses = futures::future::join_all(responses).await;
 
     let ok_count = responses.iter()
-        .filter(|res_tuple| {
-            res_tuple.0.is_ok()
-        })
+        .filter(|response| { response.is_ok() })
         .count();
-    let are_all_ok = ok_count == responses.len();
-    let are_non_ok = ok_count == 0;
-    if !are_all_ok {
-        return if are_non_ok {
-            Err(CheckError::Timeout)
-        } else {
-            Err(CheckError::PartialTimeout)
-        }
+
+    if ok_count == responses.len() {
+        let rtt_ms = responses.iter()
+            .filter_map(|response| response.as_ref().ok())
+            .map(|response| response.1)
+            .map(|duration| duration.as_millis() as u32)
+            .sum::<u32>() / responses.len() as u32;
+
+        return Ok(CandidateProfile {
+            candidate,
+            addrs,
+            rtt_ms,
+        });
     }
 
-    let rtt_ms = responses.iter()
-        .map(|res_tuple| res_tuple.1)
-        .map(|duration| duration.as_millis() as u32)
-        .sum::<u32>() / responses.len() as u32;
+    let ip_fails = responses.iter()
+        .filter_map(|response| response.err())
+        .filter(|err| err == &CheckError::IncorrectMappingReturned)
+        .count();
 
-    Ok(CandidateProfile {
-        candidate,
-        addrs,
-        rtt_ms
-    })
-}
+    if ip_fails > 0 {
+        return Err(CheckError::IncorrectMappingReturned);
+    }
+
+    let timeouts = responses.iter()
+        .filter_map(|response| response.err())
+        .filter(|err| err == &CheckError::Timeout)
+        .count();
+
+    if timeouts < responses.len() {
+        return Err(CheckError::PartialTimeout);
+    }
+
+    return Err(CheckError::Timeout);
+}