From 2e57ca42581102fd548dae7c94f42e85b436a06b Mon Sep 17 00:00:00 2001
From: Yash Karandikar <nerdstep710@gmail.com>
Date: Fri, 24 Jun 2022 12:44:17 +0530
Subject: [PATCH] Read TOML into hashmap - this might need major refactoring
 later

---
 Cargo.lock  | 31 +++++++++++++++++++++++++++++++
 Cargo.toml  |  2 ++
 src/main.rs | 19 ++++++++++++++++---
 3 files changed, 49 insertions(+), 3 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 5113e79..7f0d465 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -13,7 +13,9 @@ name = "atp"
 version = "0.1.0"
 dependencies = [
  "anyhow",
+ "serde",
  "tokio",
+ "toml",
 ]
 
 [[package]]
@@ -170,6 +172,26 @@ version = "1.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "d29ab0c6d3fc0ee92fe66e2d99f700eab17a8d57d1c1d3b748380fb20baa78cd"
 
+[[package]]
+name = "serde"
+version = "1.0.137"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "61ea8d54c77f8315140a05f4c7237403bf38b72704d031543aa1d16abbf517d1"
+dependencies = [
+ "serde_derive",
+]
+
+[[package]]
+name = "serde_derive"
+version = "1.0.137"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "1f26faba0c3959972377d3b2d306ee9f71faee9714294e41bb777f83f88578be"
+dependencies = [
+ "proc-macro2",
+ "quote",
+ "syn",
+]
+
 [[package]]
 name = "signal-hook-registry"
 version = "1.4.0"
@@ -237,6 +259,15 @@ dependencies = [
  "syn",
 ]
 
+[[package]]
+name = "toml"
+version = "0.5.9"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "8d82e1a7758622a465f8cee077614c73484dac5b836c02ff6a40d5d1010324d7"
+dependencies = [
+ "serde",
+]
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.1"
diff --git a/Cargo.toml b/Cargo.toml
index 7a6101d..5c089c8 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -7,4 +7,6 @@ edition = "2021"
 
 [dependencies]
 anyhow = "1.0.58"
+serde = { version = "1.0.137", features = ["derive"] }
 tokio = { version = "1.19.2", features = ["full"] }
+toml = "0.5.9"
diff --git a/src/main.rs b/src/main.rs
index 547d20f..81438ee 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,14 +1,27 @@
+use std::collections::HashMap;
+use std::net::SocketAddr;
 use tokio::io::copy;
 use tokio::net::{TcpListener, TcpStream};
 
 #[tokio::main]
 async fn main() -> anyhow::Result<()> {
+    let filename = std::env::args()
+        .next()
+        .unwrap_or_else(|| "config.toml".to_string());
+
+    let contents = std::fs::read_to_string(filename)?;
+
+    let mut ports: HashMap<u16, SocketAddr> = toml::from_str(&contents)?;
+
     let listener = TcpListener::bind("0.0.0.0:8000").await?;
 
     loop {
-        let (stream, _addr) = listener.accept().await?;
-        // TODO: use a lookup to see where it should be proxied to
-        let to_stream = TcpStream::connect("192.168.1.28:6667").await?; // TODO: remove this placeholder
+        let (stream, addr) = listener.accept().await?;
+        let to_addr = match ports.get(&addr.port()) {
+            Some(v) => *v,
+            None => continue,
+        };
+        let to_stream = TcpStream::connect(to_addr).await?; // TODO: remove this placeholder
 
         tokio::spawn(async move {
             proxy(stream, to_stream).await?;