diff --git a/pipe.go b/pipe.go index d99eedb..eaa1523 100644 --- a/pipe.go +++ b/pipe.go @@ -3,6 +3,7 @@ package winio import ( + "context" "errors" "io" "net" @@ -137,9 +138,30 @@ func (s pipeAddress) String() string { return string(s) } +// tryDialPipe attempts to dial the pipe at `path` until `ctx` cancellation or timeout. +func tryDialPipe(ctx context.Context, path *string) (syscall.Handle, error) { + for { + select { + case <-ctx.Done(): + return syscall.Handle(0), ctx.Err() + default: + h, err := createFile(*path, syscall.GENERIC_READ|syscall.GENERIC_WRITE, 0, nil, syscall.OPEN_EXISTING, syscall.FILE_FLAG_OVERLAPPED|cSECURITY_SQOS_PRESENT|cSECURITY_ANONYMOUS, 0) + if err == nil { + return h, nil + } + if err != cERROR_PIPE_BUSY { + return h, &os.PathError{Err: err, Op: "open", Path: *path} + } + // Wait 10 msec and try again. This is a rather simplistic + // view, as we always try each 10 milliseconds. + time.Sleep(time.Millisecond * 10) + } + } +} + // DialPipe connects to a named pipe by path, timing out if the connection // takes longer than the specified duration. If timeout is nil, then we use -// a default timeout of 5 seconds. (We do not use WaitNamedPipe.) +// a default timeout of 2 seconds. (We do not use WaitNamedPipe.) func DialPipe(path string, timeout *time.Duration) (net.Conn, error) { var absTimeout time.Time if timeout != nil { @@ -147,23 +169,22 @@ func DialPipe(path string, timeout *time.Duration) (net.Conn, error) { } else { absTimeout = time.Now().Add(time.Second * 2) } + ctx, _ := context.WithDeadline(context.Background(), absTimeout) + conn, err := DialPipeContext(ctx, path) + if err == context.DeadlineExceeded { + return nil, ErrTimeout + } + return conn, err +} + +// DialPipeContext attempts to connect to a named pipe by `path` until `ctx` +// cancellation or timeout. +func DialPipeContext(ctx context.Context, path string) (net.Conn, error) { var err error var h syscall.Handle - for { - h, err = createFile(path, syscall.GENERIC_READ|syscall.GENERIC_WRITE, 0, nil, syscall.OPEN_EXISTING, syscall.FILE_FLAG_OVERLAPPED|cSECURITY_SQOS_PRESENT|cSECURITY_ANONYMOUS, 0) - if err != cERROR_PIPE_BUSY { - break - } - if time.Now().After(absTimeout) { - return nil, ErrTimeout - } - - // Wait 10 msec and try again. This is a rather simplistic - // view, as we always try each 10 milliseconds. - time.Sleep(time.Millisecond * 10) - } + h, err = tryDialPipe(ctx, &path) if err != nil { - return nil, &os.PathError{Op: "open", Path: path, Err: err} + return nil, err } var flags uint32 diff --git a/pipe_test.go b/pipe_test.go index 3869207..89e32d1 100644 --- a/pipe_test.go +++ b/pipe_test.go @@ -3,6 +3,7 @@ package winio import ( "bufio" "bytes" + "context" "io" "net" "os" @@ -37,6 +38,40 @@ func TestDialListenerTimesOut(t *testing.T) { } } +func TestDialContextListenerTimesOut(t *testing.T) { + l, err := ListenPipe(testPipeName, nil) + if err != nil { + t.Fatal(err) + } + defer l.Close() + var d = time.Duration(10 * time.Millisecond) + ctx, _ := context.WithTimeout(context.Background(), d) + _, err = DialPipeContext(ctx, testPipeName) + if err != context.DeadlineExceeded { + t.Fatalf("expected context.DeadlineExceeded, got %v", err) + } +} + +func TestDialListenerGetsCancelled(t *testing.T) { + ctx, cancel := context.WithCancel(context.Background()) + l, err := ListenPipe(testPipeName, nil) + if err != nil { + t.Fatal(err) + } + ch := make(chan error) + defer l.Close() + go func(ctx context.Context, ch chan error) { + _, err := DialPipeContext(ctx, testPipeName) + ch <- err + }(ctx, ch) + time.Sleep(time.Millisecond * 30) + cancel() + err = <-ch + if err != context.Canceled { + t.Fatalf("expected context.Canceled, got %v", err) + } +} + func TestDialAccessDeniedWithRestrictedSD(t *testing.T) { c := PipeConfig{ SecurityDescriptor: "D:P(A;;0x1200FF;;;WD)",