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/ssh/ssh.go | 13 ++++++++----- 1 file changed, 8 insertions(+), 5 deletions(-) (limited to 'forged/internal/incoming/ssh') diff --git a/forged/internal/incoming/ssh/ssh.go b/forged/internal/incoming/ssh/ssh.go index 9f9bdff..0c722c0 100644 --- a/forged/internal/incoming/ssh/ssh.go +++ b/forged/internal/incoming/ssh/ssh.go @@ -61,20 +61,23 @@ func New(config Config) (server *Server, err error) { func (server *Server) Run(ctx context.Context) (err error) { listener, err := misc.Listen(server.net, server.addr) + if err != nil { + return fmt.Errorf("listen for SSH: %w", err) + } defer func() { _ = listener.Close() }() - go func() { - <-ctx.Done() - shCtx, cancel := context.WithTimeout(context.Background(), time.Duration(server.shutdownTimeout)*time.Second) + stop := context.AfterFunc(ctx, func() { + shCtx, cancel := context.WithTimeout(context.WithoutCancel(ctx), time.Duration(server.shutdownTimeout)*time.Second) defer cancel() _ = server.gliderServer.Shutdown(shCtx) _ = listener.Close() - }() + }) + defer stop() if err = server.gliderServer.Serve(listener); err != nil { - if errors.Is(err, gliderssh.ErrServerClosed) { + if errors.Is(err, gliderssh.ErrServerClosed) || ctx.Err() != nil { return nil } return fmt.Errorf("serve SSH: %w", err) -- cgit v1.2.3