From 6cc29a137ed073e30c032ee112868f8d23e4aace Mon Sep 17 00:00:00 2001
From: Alois <aloisianer@proton.me>
Date: Sat, 16 May 2026 17:49:25 +0200
Subject: [PATCH] [Add] port options in nix flake

---
 flake.nix                     | 28 ++++++++++++++++++++++++++--
 src/main.rs                   |  7 ++++++-
 src/omega/omega_connection.rs | 17 ++++++++++++++++-
 3 files changed, 48 insertions(+), 4 deletions(-)

diff --git a/flake.nix b/flake.nix
index aeee942..03138e2 100644
--- a/flake.nix
+++ b/flake.nix
@@ -88,6 +88,24 @@
                 description = "Whether to open the firewall for ports used by the Omikron.";
               };
 
+              rhoPort = lib.mkOption {
+                type = lib.types.port;
+                default = 959;
+                description = "Port the Omikron server listens on for incoming QUIC connections.";
+              };
+
+              omegaHost = lib.mkOption {
+                type = lib.types.str;
+                default = "tensamin.net";
+                description = "Hostname of the upstream Omega server.";
+              };
+
+              omegaPort = lib.mkOption {
+                type = lib.types.port;
+                default = 9187;
+                description = "Port of the upstream Omega server.";
+              };
+
               package = lib.mkOption {
                 type = lib.types.package;
                 default = defaultPackage;
@@ -118,6 +136,12 @@
                 wantedBy = [ "multi-user.target" ];
                 after = [ "network.target" ];
 
+                environment = {
+                  RHO_PORT = toString cfg.rhoPort;
+                  OMEGA_HOST = cfg.omegaHost;
+                  OMEGA_PORT = toString cfg.omegaPort;
+                };
+
                 serviceConfig = {
                   Type = "simple";
                   User = "omikron";
@@ -164,8 +188,8 @@
               };
 
               networking.firewall = lib.mkIf cfg.openFirewall {
-                allowedTCPPorts = [ 959 ];
-                allowedUDPPorts = [ 959 ];
+                allowedTCPPorts = [ cfg.rhoPort ];
+                allowedUDPPorts = [ cfg.rhoPort ];
               };
             };
           };
diff --git a/src/main.rs b/src/main.rs
index fa511ea..f10b39f 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -39,9 +39,14 @@ async fn main() {
     dotenv().ok();
     startup();
 
+    let rho_port = env::var("RHO_PORT")
+        .ok()
+        .and_then(|s| s.parse().ok())
+        .unwrap_or(959);
+
     get_omega_connection();
     tokio::spawn(async move {
-        if let Err(e) = start(959).await {
+        if let Err(e) = start(rho_port).await {
             log_err!(0, util::logger::PrintType::General, "{}", e);
         }
     });
diff --git a/src/omega/omega_connection.rs b/src/omega/omega_connection.rs
index 211ed6d..a1fb418 100644
--- a/src/omega/omega_connection.rs
+++ b/src/omega/omega_connection.rs
@@ -26,6 +26,20 @@ use uuid::Uuid;
 
 const OMEGA_HOST_DEFAULT: &str = "tensamin.net";
 const OMEGA_PORT_DEFAULT: u16 = 9187;
+
+fn omega_host_and_port() -> (String, u16) {
+    let host = env::var("OMEGA_HOST")
+        .map(|s| s.trim().to_string())
+        .unwrap_or_else(|_| OMEGA_HOST_DEFAULT.to_string());
+
+    let port = env::var("OMEGA_PORT")
+        .ok()
+        .and_then(|s| s.trim().parse().ok())
+        .unwrap_or(OMEGA_PORT_DEFAULT);
+
+    (host, port)
+}
+
 const RECONNECT_DELAY: Duration = Duration::from_secs(5);
 const MAX_RECONNECT_DELAY: Duration = Duration::from_secs(300);
 const CONNECTION_TIMEOUT: Duration = Duration::from_secs(10);
@@ -104,7 +118,8 @@ pub struct OmegaConnection {
 
 impl OmegaConnection {
     pub fn new() -> Self {
-        Self::with_host(OMEGA_HOST_DEFAULT, OMEGA_PORT_DEFAULT)
+        let (host, port) = omega_host_and_port();
+        Self::with_host(&host, port)
     }
 
     pub fn with_host(host: &str, port: u16) -> Self {