diff --git a/cmd/olliesrv/internal/session/session.go b/cmd/olliesrv/internal/session/session.go index 671283c..b55ea62 100644 --- a/cmd/olliesrv/internal/session/session.go +++ b/cmd/olliesrv/internal/session/session.go @@ -135,7 +135,7 @@ func (s *Session) runBypassLoop(notifyFn BypassNotifyFunc) { } } - req, err := conn.ReadBypassPending() + req, err := conn.ReadBypassPending(s.Ctx) conn.Close() if err != nil { select { diff --git a/cmd/olliesrv/server.go b/cmd/olliesrv/server.go index 14a99d8..d5183f9 100644 --- a/cmd/olliesrv/server.go +++ b/cmd/olliesrv/server.go @@ -198,8 +198,18 @@ func (s *Server) Start(conn net.Conn) { }() for { + // Set read deadline so we can check for shutdown periodically + conn.SetReadDeadline(time.Now().Add(1 * time.Second)) fc, err := plan9.ReadFcall(conn) if err != nil { + // Check if shutdown was requested + if connCtx.Err() != nil { + break + } + // Timeout - loop and check context again + if netErr, ok := err.(net.Error); ok && netErr.Timeout() { + continue + } if err != io.EOF { s.log.Error("read: %v", err) } @@ -704,11 +714,6 @@ func remove(root *fs.Tree, cs *connState, fc *plan9.Fcall) *plan9.Fcall { func (s *Server) Kill() { s.shutdownCancel() - s.connMu.Lock() - for conn := range s.activeConns { - conn.Close() - } - s.connMu.Unlock() session.Shutdown() done := make(chan struct{}) go func() { diff --git a/toolsrv/client9p.go b/toolsrv/client9p.go index 44cf00e..4f1f541 100644 --- a/toolsrv/client9p.go +++ b/toolsrv/client9p.go @@ -551,14 +551,25 @@ type BypassRequest struct { // ReadBypassPending blocks until a bypass request is available. // This should be called in a loop by the approval handler. -func (c *Conn) ReadBypassPending() (*BypassRequest, error) { +// Returns nil, ctx.Err() if context is cancelled. +func (c *Conn) ReadBypassPending(ctx context.Context) (*BypassRequest, error) { fid, err := c.fsys.Open("bypass/pending", plan9.OREAD) if err != nil { return nil, fmt.Errorf("open bypass/pending: %w", err) } - defer fid.Close() - // Blocking read - will return when a request is available + // Watch for context cancellation to unblock the read + done := make(chan struct{}) + defer close(done) + go func() { + select { + case <-ctx.Done(): + fid.Close() + case <-done: + } + }() + + // Blocking read - will return when a request is available or fid is closed var result []byte buf := make([]byte, 8192) for { @@ -570,9 +581,14 @@ func (c *Conn) ReadBypassPending() (*BypassRequest, error) { break } if err != nil { + fid.Close() + if ctx.Err() != nil { + return nil, ctx.Err() + } return nil, fmt.Errorf("read bypass/pending: %w", err) } } + fid.Close() var req BypassRequest if err := json.Unmarshal(result, &req); err != nil {