Proper async server and ctrl-c handling

1 files changed, 64 insertions, 15 deletions

diff --git a/src/main.rs b/src/main.rs
index 088f27f..a113a30 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,40 +1,89 @@
-use std::net::{IpAddr, Ipv4Addr, SocketAddr, TcpListener, TcpStream};
+use std::net::{IpAddr, Ipv4Addr, SocketAddr};
 use std::time::Duration;
 
 use anyhow::Result;
-use craftio_rs::{CraftIo, CraftSyncReader, CraftSyncWriter, CraftTcpConnection};
+use craftio_rs::CraftIo;
 use mcproto_rs::nbt::Tag;
-use mcproto_rs::protocol::{HasPacketKind, State};
+use mcproto_rs::protocol::State;
 use mcproto_rs::Serializer;
 use mcproto_rs::status::{StatusPlayersSpec, StatusSpec, StatusVersionSpec};
 use mcproto_rs::types::{BytesSerializer, Chat, IntPosition, ItemStack, NamedNbtTag, RemainingBytes, Vec3};
 use mcproto_rs::uuid::UUID4;
 use mcproto_rs::v1_19_3::{BitSet, BlobArray, BookSettings, ChunkDataAndUpdateLightSpec, ChunkDataSpec, ChunkSection, CommandNode, CommandNodeSpec, CommandsSpec, EntityEventSpec, InitializeWorldBorderSpec, KeepAliveClientBoundSpec, LightDataSpec, PalettedContainer, PluginMessageSpec, RecipeBookAction, RecipeBookInitSpec, SetCenterChunkSpec, SetContainerContentSpec, SetDefaultSpawnPositionSpec, SynchronizePlayerPositionSpec, TagType, TypedTagList, UpdateRecipeBookSpec};
-use mcproto_rs::v1_19_3::{GameMode, HandshakeNextState, LoginPlaySpec, LoginSuccessSpec, Packet761, Packet761Kind, PingRequestSpec, PingResponseSpec, PreviousGameMode, RawPacket761, SetHeldItemClientSpec, StatusResponseSpec, UpdateRecipesSpec, UpdateTagsSpec};
+use mcproto_rs::v1_19_3::{GameMode, HandshakeNextState, LoginPlaySpec, LoginSuccessSpec, Packet761, Packet761Kind, PingRequestSpec, PingResponseSpec, PreviousGameMode, SetHeldItemClientSpec, StatusResponseSpec, UpdateRecipesSpec, UpdateTagsSpec};
 use tokio;
-use crate::connect::MinecraftClient;
+use tokio::net::TcpStream;
+use tokio::signal::unix::{signal, SignalKind};
+use tokio::sync::{broadcast};
+use tokio::sync::broadcast::Receiver;
+use tokio::task::JoinHandle;
 
+use crate::connect::MinecraftClient;
 
 pub mod connect;
 pub mod nbtdsl;
 
+
 #[tokio::main]
 async fn main() -> Result<()> {
+    let (shutdown_send, _recv) = broadcast::channel::<()>(1);
+
+    let bind = tokio::net::TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 25565)).await?;
+    let server_rx = shutdown_send.subscribe();
+    let server_handle = tokio::spawn(async {
+        listener(bind, server_rx).await;
+        ()
+    });
+
+    println!("Awaiting shutdown request");
+    signal(SignalKind::interrupt())?.recv().await;
+    println!("Shutdown request detected");
+    shutdown_send.send(())?;
+    server_handle.await?;
+
+    Ok(())
+}
+
+async fn listener(listener: tokio::net::TcpListener, mut receiver: Receiver<()>) {
     println!("Starting server");
-    let bind = TcpListener::bind(SocketAddr::new(IpAddr::V4(Ipv4Addr::LOCALHOST), 25565))?;
+    let mut handles: Vec<JoinHandle<()>> = Vec::new();
     loop {
-        if let Ok((socket, address)) = bind.accept() {
-            println!("Connection accepted from {}", address);
-            let client = MinecraftClient::from_stream(socket)?;
-            tokio::spawn(async {
-                if let Err(x) = handle_conn(client).await {
-                    println!("Error: {:?}", x);
-                }
-            });
-        }
+        let x: Option<(TcpStream, SocketAddr)> = tokio::select! {
+            _ = receiver.recv() => {
+                break;
+            }
+            conn = listener.accept() => { conn.ok() }
+        };
+        let (stream, from) = match x {
+            None => { continue }
+            Some(y) => { y }
+        };
+        println!("Connection accepted from {}", from);
+        let client = MinecraftClient::from_stream(stream);
+        let handle = tokio::spawn(async {
+            if let Err(x) = handle_conn(client).await {
+                println!("Error: {:?}", x);
+            }
+        });
+        handles.push(handle);
     }
+    println!("Shutdown started");
+    for x in &handles {
+        x.abort();
+    }
+    /*
+    TODO: proper kicks so that we can just join
+    for x in handles {
+        match x.await {
+            Ok(_) => {}
+            Err(_) => {
+                println!("Failed to join server during exit!")
+            }
+        };
+    }*/
 }
 
+
 async fn handle_conn(mut client: MinecraftClient) -> Result<()> {
     let hs = assert_packet!(Handshake, client, "Missing packet");
     if hs.next_state == HandshakeNextState::Login {