From 3b0a2adf296791bdbd0c1b0fba6399910c1fac82 Mon Sep 17 00:00:00 2001
From: Runxi Yu <me@runxiyu.org>
Date: Sun, 17 Aug 2025 04:08:24 +0800
Subject: Context fixes

---
 forged/internal/incoming/hooks/hooks.go | 13 ++++++++++++-
 forged/internal/incoming/lmtp/config.go | 12 +++++++++++-
 forged/internal/incoming/ssh/ssh.go     | 14 +++++++++++++-
 forged/internal/incoming/web/web.go     | 14 +++++++++++++-
 4 files changed, 49 insertions(+), 4 deletions(-)

(limited to 'forged/internal/incoming')

diff --git a/forged/internal/incoming/hooks/hooks.go b/forged/internal/incoming/hooks/hooks.go
index 18b9742..c021406 100644
--- a/forged/internal/incoming/hooks/hooks.go
+++ b/forged/internal/incoming/hooks/hooks.go
@@ -1,6 +1,7 @@
 package hooks
 
 import (
+	"context"
 	"fmt"
 	"net"
 
@@ -40,15 +41,25 @@ func New(config Config) (server *Server) {
 	}
 }
 
-func (server *Server) Run() error {
+func (server *Server) Run(ctx context.Context) error {
 	listener, _, err := misc.ListenUnixSocket(server.socketPath)
 	if err != nil {
 		return fmt.Errorf("listen unix socket for hooks: %w", err)
 	}
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = listener.Close()
+		// TODO: Log the error
+	}()
 
 	for {
 		conn, err := listener.Accept()
 		if err != nil {
+			// TODO: Handle errors caused by context cancel
 			return fmt.Errorf("accept conn: %w", err)
 		}
 
diff --git a/forged/internal/incoming/lmtp/config.go b/forged/internal/incoming/lmtp/config.go
index e6db5a6..c622648 100644
--- a/forged/internal/incoming/lmtp/config.go
+++ b/forged/internal/incoming/lmtp/config.go
@@ -1,6 +1,7 @@
 package lmtp
 
 import (
+	"context"
 	"fmt"
 	"net"
 
@@ -33,11 +34,20 @@ func New(config Config) (server *Server) {
 	}
 }
 
-func (server *Server) Run() error {
+func (server *Server) Run(ctx context.Context) error {
 	listener, _, err := misc.ListenUnixSocket(server.socket)
 	if err != nil {
 		return fmt.Errorf("listen unix socket for LMTP: %w", err)
 	}
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = listener.Close()
+		// TODO: Log the error
+	}()
 
 	for {
 		conn, err := listener.Accept()
diff --git a/forged/internal/incoming/ssh/ssh.go b/forged/internal/incoming/ssh/ssh.go
index b0269e5..efa24f1 100644
--- a/forged/internal/incoming/ssh/ssh.go
+++ b/forged/internal/incoming/ssh/ssh.go
@@ -1,6 +1,7 @@
 package ssh
 
 import (
+	"context"
 	"fmt"
 	"os"
 
@@ -53,8 +54,19 @@ func New(config Config) (server *Server, err error) {
 	return
 }
 
-func (server *Server) Run() (err error) {
+func (server *Server) Run(ctx context.Context) (err error) {
 	listener, err := misc.Listen(server.net, server.addr)
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = server.gliderServer.Close()
+		_ = listener.Close() // unnecessary?
+		// TODO: Log the error
+	}()
+
 	if err = server.gliderServer.Serve(listener); err != nil {
 		return fmt.Errorf("serve SSH: %w", err)
 	}
diff --git a/forged/internal/incoming/web/web.go b/forged/internal/incoming/web/web.go
index e4675f0..a95b650 100644
--- a/forged/internal/incoming/web/web.go
+++ b/forged/internal/incoming/web/web.go
@@ -1,6 +1,7 @@
 package web
 
 import (
+	"context"
 	"fmt"
 	"net/http"
 	"time"
@@ -48,8 +49,19 @@ func New(config Config) (server *Server) {
 	}
 }
 
-func (server *Server) Run() (err error) {
+func (server *Server) Run(ctx context.Context) (err error) {
 	listener, err := misc.Listen(server.net, server.addr)
+	defer func() {
+		_ = listener.Close()
+	}()
+
+	go func() {
+		<-ctx.Done()
+		_ = server.httpServer.Close()
+		_ = listener.Close() // unnecessary?
+		// TODO: Log the error
+	}()
+
 	if err = server.httpServer.Serve(listener); err != nil {
 		return fmt.Errorf("serve web: %w", err)
 	}
-- 
cgit v1.2.3