From 2c99e2461a98972fcbe292762d50839dabdccd4e Mon Sep 17 00:00:00 2001
From: Alexey Kondratov <kondratov.aleksey@gmail.com>
Date: Fri, 1 Oct 2021 15:14:42 +0300
Subject: [PATCH] Allow usage of the compute hostname in the proxy

---
 proxy/src/cplane_api.rs | 15 ++++++++++-----
 proxy/src/proxy.rs      |  2 +-
 2 files changed, 11 insertions(+), 6 deletions(-)

diff --git a/proxy/src/cplane_api.rs b/proxy/src/cplane_api.rs
index 51a8c96b71..2b3b8a64bd 100644
--- a/proxy/src/cplane_api.rs
+++ b/proxy/src/cplane_api.rs
@@ -1,6 +1,6 @@
-use anyhow::{bail, Result};
+use anyhow::{bail, Context, Result};
 use serde::{Deserialize, Serialize};
-use std::net::{IpAddr, SocketAddr};
+use std::net::{SocketAddr, ToSocketAddrs};
 
 pub struct CPlaneApi {
     auth_endpoint: &'static str,
@@ -8,7 +8,7 @@ pub struct CPlaneApi {
 
 #[derive(Serialize, Deserialize, Debug)]
 pub struct DatabaseInfo {
-    pub host: IpAddr, // TODO: allow host name here too
+    pub host: String,
     pub port: u16,
     pub dbname: String,
     pub user: String,
@@ -16,8 +16,13 @@ pub struct DatabaseInfo {
 }
 
 impl DatabaseInfo {
-    pub fn socket_addr(&self) -> SocketAddr {
-        SocketAddr::new(self.host, self.port)
+    pub fn socket_addr(&self) -> Result<SocketAddr> {
+        let host_port = format!("{}:{}", self.host, self.port);
+        host_port
+            .to_socket_addrs()
+            .with_context(|| format!("cannot resolve {} to SocketAddr", host_port))?
+            .next()
+            .ok_or_else(|| anyhow::Error::msg("cannot resolve at least one SocketAddr"))
     }
 
     pub fn conn_string(&self) -> String {
diff --git a/proxy/src/proxy.rs b/proxy/src/proxy.rs
index 5755a9cf9f..f246d4470a 100644
--- a/proxy/src/proxy.rs
+++ b/proxy/src/proxy.rs
@@ -248,7 +248,7 @@ databases without opening the browser.
 
 /// Create a TCP connection to a postgres database, authenticate with it, and receive the ReadyForQuery message
 async fn connect_to_db(db_info: DatabaseInfo) -> anyhow::Result<tokio::net::TcpStream> {
-    let mut socket = tokio::net::TcpStream::connect(db_info.socket_addr()).await?;
+    let mut socket = tokio::net::TcpStream::connect(db_info.socket_addr()?).await?;
     let config = db_info.conn_string().parse::<tokio_postgres::Config>()?;
     let _ = config.connect_raw(&mut socket, NoTls).await?;
     Ok(socket)