Improve signal handling for graceful shutdown on Ctrl+C

Address issues where DMSG client utilities would not respond to Ctrl+C:

- Use signal-aware context for dmsgC.Serve() instead of context.Background() in internal/cli, pkg/direct, and dmsgpty-host
- Add context cancellation check in dmsg-socks5 accept loop to properly exit on signal
- Implement graceful HTTP server shutdown with 5-second timeout in dmsghttp
- Ensure dmsgC.Close() is called via defer in dmsgpty-host for proper cleanup
- Add WaitGroup wait in direct client stop function for clean goroutine shutdown
- Document temporary kill.go workaround for verification testing

These changes ensure client utilities properly respond to interrupt signals without requiring multiple Ctrl+C presses.

💘 Generated with Crush
Co-Authored-By: Crush <[email protected]>

Author: 0pcom

cmd/dmsg-socks5/commands/dmsg-socks5.go

@@ -157,11 +157,25 @@ var serveCmd = &cobra.Command{
 				dlog.Printf("Error closing listener: %v", err)
 			}
 		}()
+
+		go func() {
+			<-ctx.Done()
+			if err := dmsgL.Close(); err != nil {
+				dlog.WithError(err).Debug("Error closing listener on context cancellation")
+			}
+		}()
+
 		for {
 			respConn, err := dmsgL.Accept()
 			if err != nil {
-				dlog.Errorf("Error accepting initiator: %v", err)
-				continue
+				select {
+				case <-ctx.Done():
+					dlog.Info("Shutting down SOCKS5 server...")
+					return
+				default:
+					dlog.Errorf("Error accepting initiator: %v", err)
+					continue
+				}
 			}
 			dlog.Infof("Accepted connection from: %s", respConn.RemoteAddr())
 

cmd/dmsg/commands/kill.go

@@ -1,4 +1,4 @@
-// Package commands cmd/.../commands/kill.go
+// Package commands cmd/dmsg/commands/kill.go
 package commands
 
 import (
@@ -8,7 +8,11 @@ import (
 )
 
 func init() {
-	//the application must stop on ctrl+c
+	// TEMPORARY WORKAROUND: Force exit on Ctrl+C after 3 attempts
+	// This can be removed once the proper signal handling fixes are verified:
+	// - dmsgC.Serve() now uses signal-aware context (not context.Background())
+	// - Accept loops now check for context cancellation
+	// - HTTP servers now shutdown gracefully
 	c := make(chan os.Signal, 1)
 	signal.Notify(c, os.Interrupt, syscall.SIGTERM)
 	go func() {

cmd/dmsghttp/commands/dmsghttp.go

@@ -165,6 +165,16 @@ func server() {
 		WriteTimeout:      10 * time.Second,
 	}
 
+	// Gracefully shutdown HTTP server on context cancellation
+	go func() {
+		<-ctx.Done()
+		shutdownCtx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
+		defer cancel()
+		if err := serve.Shutdown(shutdownCtx); err != nil {
+			dlog.WithError(err).Warn("Server shutdown error")
+		}
+	}()
+
 	// Start serving
 	go func() {
 		dlog.WithField("dmsg_addr", lis.Addr().String()).Debug("Serving...\n")

cmd/dmsgpty-host/commands/root.go

@@ -104,7 +104,13 @@ var RootCmd = &cobra.Command{
 		dmsgC := dmsg.NewClient(pk, sk, disc.NewHTTP(conf.DmsgDisc, &http.Client{}, log), &dmsg.Config{
 			MinSessions: conf.DmsgSessions,
 		})
-		go dmsgC.Serve(context.Background())
+		defer func() {
+			if err := dmsgC.Close(); err != nil {
+				log.WithError(err).Warn("Failed to close dmsg client")
+			}
+		}()
+
+		go dmsgC.Serve(ctx)
 		select {
 		case <-ctx.Done():
 			return fmt.Errorf("failed to wait dmsg client to be ready: %w", ctx.Err())

internal/cli/cli.go

@@ -123,8 +123,8 @@ func StartDmsg(ctx context.Context, dlog *logging.Logger, pk cipher.PubKey, sk c
 	dmsgC = dmsg.NewClient(pk, sk, disc.NewHTTP(dmsgDisc, httpClient, dlog), &dmsg.Config{MinSessions: dmsgSessions})
 	dlog.Debug("Created dmsg client.")
 
-	go dmsgC.Serve(context.Background())
-	dlog.Debug("dmsgclient.Serve(context.Background())")
+	go dmsgC.Serve(ctx)
+	dlog.Debug("dmsgclient.Serve(ctx)")
 
 	stop = func() {
 		err := dmsgC.Close()

pkg/direct/direct.go

@@ -23,12 +23,13 @@ func StartDmsg(ctx context.Context, log *logging.Logger, pk cipher.PubKey, sk ci
 	wg.Add(1)
 	go func() {
 		defer wg.Done()
-		dmsgDC.Serve(context.Background())
+		dmsgDC.Serve(ctx)
 	}()
 
 	stop = func() {
 		err := dmsgDC.Close()
 		log.WithError(err).Debug("Disconnected from dmsg network.\n")
+		wg.Wait()
 	}
 
 	log.WithField("public_key", pk.String()).Debug("Connecting to dmsg network...\n")