From f98400f753967917f0ddc125c28e6ab05ccf0f64 Mon Sep 17 00:00:00 2001
From: Filip Tibell <filip.tibell@gmail.com>
Date: Tue, 14 Feb 2023 19:05:35 +0100
Subject: [PATCH] Improve error for net serve failing to bind

---
 CHANGELOG.md                           |  4 ++++
 packages/lib/src/globals/net.rs        | 20 +++++++++++++++++---
 packages/lib/src/lua/task/scheduler.rs | 13 +++++++++++--
 3 files changed, 32 insertions(+), 5 deletions(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 5f284fb..dc1b795 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -10,6 +10,10 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## Unreleased
 
+### Changed
+
+- Improve error messages when `net.serve` fails
+
 ### Fixed
 
 - Fixed `task.delay` keeping the script running even if it was cancelled using `task.cancel`
diff --git a/packages/lib/src/globals/net.rs b/packages/lib/src/globals/net.rs
index 41f32f8..df544d9 100644
--- a/packages/lib/src/globals/net.rs
+++ b/packages/lib/src/globals/net.rs
@@ -6,6 +6,7 @@ use std::{
     task::{Context, Poll},
 };
 
+use console::style;
 use mlua::prelude::*;
 
 use hyper::{body::to_bytes, server::conn::AddrStream, service::Service};
@@ -175,11 +176,24 @@ async fn net_serve<'a>(
         lua.create_registry_value(handler)
             .expect("Failed to store websocket handler")
     }));
+    let sched = lua.app_data_mut::<&TaskScheduler>().unwrap();
+    // Bind first to make sure that we can bind to this address
+    let bound = match Server::try_bind(&([127, 0, 0, 1], port).into()) {
+        Err(e) => {
+            return Err(LuaError::external(format!(
+                "Failed to bind to localhost on port {port}\n{}",
+                format!("{e}").replace(
+                    "error creating server listener: ",
+                    &format!("{}", style("> ").dim())
+                )
+            )));
+        }
+        Ok(bound) => bound,
+    };
     // Register a background task to prevent
     // the task scheduler from exiting early
-    let sched = lua.app_data_mut::<&TaskScheduler>().unwrap();
     let task = sched.register_background_task();
-    let server = Server::bind(&([127, 0, 0, 1], port).into())
+    let server = bound
         .http1_only(true)
         .http1_keepalive(true)
         .executor(LocalExec)
@@ -189,12 +203,12 @@ async fn net_serve<'a>(
             server_websocket_callback,
         ))
         .with_graceful_shutdown(async move {
+            task.unregister(Ok(()));
             shutdown_rx
                 .recv()
                 .await
                 .expect("Server was stopped instantly");
             shutdown_rx.close();
-            task.unregister(Ok(()));
         });
     // Spawn a new tokio task so we don't block
     task::spawn_local(server);
diff --git a/packages/lib/src/lua/task/scheduler.rs b/packages/lib/src/lua/task/scheduler.rs
index 77528db..318759d 100644
--- a/packages/lib/src/lua/task/scheduler.rs
+++ b/packages/lib/src/lua/task/scheduler.rs
@@ -90,11 +90,20 @@ pub struct Task {
 #[must_use = "Background tasks must be unregistered"]
 #[derive(Debug)]
 pub struct TaskSchedulerBackgroundTaskHandle {
+    unregistered: bool,
     sender: mpsc::UnboundedSender<TaskSchedulerRegistrationMessage>,
 }
 
 impl TaskSchedulerBackgroundTaskHandle {
-    pub fn unregister(self, result: LuaResult<()>) {
+    pub fn new(sender: mpsc::UnboundedSender<TaskSchedulerRegistrationMessage>) -> Self {
+        Self {
+            unregistered: false,
+            sender,
+        }
+    }
+
+    pub fn unregister(mut self, result: LuaResult<()>) {
+        self.unregistered = true;
         self.sender
             .send(TaskSchedulerRegistrationMessage::Terminated(result))
             .unwrap_or_else(|_| {
@@ -770,7 +779,7 @@ impl<'fut> TaskScheduler<'fut> {
                     env!("CARGO_PKG_REPOSITORY")
                 )
             });
-        TaskSchedulerBackgroundTaskHandle { sender }
+        TaskSchedulerBackgroundTaskHandle::new(sender)
     }
 
     /**