From b6dd3ffce416fa86c223fad5d2f6c3db5d5727e4 Mon Sep 17 00:00:00 2001 From: Runxi Yu Date: Sun, 17 Aug 2025 04:48:47 +0800 Subject: A few other context fixes --- forged/internal/incoming/hooks/hooks.go | 21 +++++++++++++-------- 1 file changed, 13 insertions(+), 8 deletions(-) (limited to 'forged/internal/incoming/hooks/hooks.go') diff --git a/forged/internal/incoming/hooks/hooks.go b/forged/internal/incoming/hooks/hooks.go index 3be0811..52ccb0f 100644 --- a/forged/internal/incoming/hooks/hooks.go +++ b/forged/internal/incoming/hooks/hooks.go @@ -5,6 +5,7 @@ import ( "errors" "fmt" "net" + "time" "github.com/gliderlabs/ssh" "go.lindenii.runxiyu.org/forge/forged/internal/common/cmap" @@ -51,25 +52,29 @@ func (server *Server) Run(ctx context.Context) error { _ = listener.Close() }() - go func() { - <-ctx.Done() + stop := context.AfterFunc(ctx, func() { _ = listener.Close() - // TODO: Log the error - }() + }) + defer stop() for { conn, err := listener.Accept() if err != nil { - if errors.Is(err, net.ErrClosed) { + if errors.Is(err, net.ErrClosed) || ctx.Err() != nil { return nil } return fmt.Errorf("accept conn: %w", err) } - go server.handleConn(conn) + go server.handleConn(ctx, conn) } } -func (server *Server) handleConn(conn net.Conn) { - panic("TODO: handle hook connection") +func (server *Server) handleConn(ctx context.Context, conn net.Conn) { + defer conn.Close() + unblock := context.AfterFunc(ctx, func() { + _ = conn.SetDeadline(time.Now()) + _ = conn.Close() + }) + defer unblock() } -- cgit v1.2.3