fix a couple of crashes with premature cancels

The first would fail as we'd decrement the WaitGroup before adding to
it, and the escond would fail with a nil pointer dereference panic.

Add tests to reproduce both, and fix them.
This commit is contained in:
Daniel Martí 2019-04-18 15:07:33 +09:00
parent 1a54253acd
commit f606ca9e73
2 changed files with 36 additions and 2 deletions

View File

@ -80,7 +80,7 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
if pc := FromContext(parent); pc != nil { if pc := FromContext(parent); pc != nil {
c.Allocator = pc.Allocator c.Allocator = pc.Allocator
c.Browser = pc.Browser c.Browser = pc.Browser
// don't inherit SessionID, so that NewContext can be used to // don't inherit Target, so that NewContext can be used to
// create a new tab on the same browser. // create a new tab on the same browser.
c.first = c.Browser == nil c.first = c.Browser == nil
@ -94,6 +94,7 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
} }
ctx = context.WithValue(ctx, contextKey{}, c) ctx = context.WithValue(ctx, contextKey{}, c)
c.wg.Add(1)
go func() { go func() {
<-ctx.Done() <-ctx.Done()
if c.first { if c.first {
@ -103,6 +104,13 @@ func NewContext(parent context.Context, opts ...ContextOption) (context.Context,
return return
} }
if c.Target == nil {
// This is a new tab, but we didn't create it and attach
// to it yet. Nothing to do.
c.wg.Done()
return
}
// Not the original browser tab; simply detach and close it. // Not the original browser tab; simply detach and close it.
// We need a new context, as ctx is cancelled; use a 1s timeout. // We need a new context, as ctx is cancelled; use a 1s timeout.
ctx, cancel := context.WithTimeout(context.Background(), time.Second) ctx, cancel := context.WithTimeout(context.Background(), time.Second)
@ -215,7 +223,6 @@ func (c *Context) newSession(ctx context.Context) error {
if err != nil { if err != nil {
return err return err
} }
c.wg.Add(1)
c.Target = c.Browser.newExecutorForTarget(ctx, targetID, sessionID) c.Target = c.Browser.newExecutorForTarget(ctx, targetID, sessionID)

View File

@ -198,3 +198,30 @@ func TestCancelError(t *testing.T) {
t.Fatalf("expected a non-nil error, got %v", err) t.Fatalf("expected a non-nil error, got %v", err)
} }
} }
func TestPrematureCancel(t *testing.T) {
t.Parallel()
// Cancel before the browser is allocated.
ctx, cancel := NewContext(context.Background())
cancel()
if err := Run(ctx); err != context.Canceled {
t.Fatalf("wanted canceled context error, got %v", err)
}
}
func TestPrematureCancelTab(t *testing.T) {
t.Parallel()
ctx1, cancel := NewContext(context.Background())
defer cancel()
if err := Run(ctx1); err != nil {
t.Fatal(err)
}
// Cancel after the browser is allocated, but before we've created a new
// tab.
ctx2, cancel := NewContext(ctx1)
cancel()
Run(ctx2)
}