From 707d55557eb3fd6e895207f601b5aefabbc51d47 Mon Sep 17 00:00:00 2001
From: Andrey Golovizin <ag@sologoc.com>
Date: Sat, 15 Apr 2023 12:45:30 +0200
Subject: [PATCH] Resolve endpoint addresses

---
 src/api.rs  | 17 ++++++++++++++++-
 src/main.rs |  2 +-
 2 files changed, 17 insertions(+), 2 deletions(-)

diff --git a/src/api.rs b/src/api.rs
index 032ae94..7948ad9 100644
--- a/src/api.rs
+++ b/src/api.rs
@@ -1,4 +1,5 @@
-use std::net::{IpAddr, Ipv4Addr, Ipv6Addr};
+use std::net::{IpAddr, Ipv4Addr, Ipv6Addr, SocketAddr, ToSocketAddrs};
+use anyhow::anyhow;
 
 use log::debug;
 use serde::{Deserialize, Serialize};
@@ -30,6 +31,20 @@ pub(crate) struct Location {
     pub pubkey: String,
 }
 
+impl Location {
+    pub(crate) fn get_endpoint(&self) -> anyhow::Result<SocketAddr> {
+        let endpoint_addr_with_port: String;
+        let endpoint_addr = if self.pool.contains(':') {
+            &self.pool
+        } else {
+            endpoint_addr_with_port = format!("{}:51820", self.pool);
+            &endpoint_addr_with_port
+        };
+        let mut endpoints = endpoint_addr.to_socket_addrs()?;
+        endpoints.next().ok_or_else(|| anyhow!("no valid endpoint address"))
+    }
+}
+
 #[derive(Serialize, Deserialize, Debug)]
 pub(crate) struct Addresses {
     pub status: String,
diff --git a/src/main.rs b/src/main.rs
index 4fe581e..bbf6913 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -146,7 +146,7 @@ fn write_config(
 
     writeln!(output, "[Peer]")?;
     writeln!(output, "PublicKey = {}", &location.pubkey)?;
-    writeln!(output, "Endpoint = {}:51820", &location.pool)?;
+    writeln!(output, "Endpoint = {}", &location.get_endpoint()?)?;
 
     let allowed_ips: &[&str] = if config_opts.no_ipv6 {
         &["0.0.0.0/0"]