cancelreader_windows.go

  1//go:build windows
  2// +build windows
  3
  4package cancelreader
  5
  6import (
  7	"fmt"
  8	"io"
  9	"os"
 10	"syscall"
 11	"time"
 12	"unicode/utf16"
 13
 14	"golang.org/x/sys/windows"
 15)
 16
 17var fileShareValidFlags uint32 = 0x00000007
 18
 19// NewReader returns a reader and a cancel function. If the input reader is a
 20// File with the same file descriptor as os.Stdin, the cancel function can
 21// be used to interrupt a blocking read call. In this case, the cancel function
 22// returns true if the call was canceled successfully. If the input reader is
 23// not a File with the same file descriptor as os.Stdin, the cancel
 24// function does nothing and always returns false. The Windows implementation
 25// is based on WaitForMultipleObject with overlapping reads from CONIN$.
 26func NewReader(reader io.Reader) (CancelReader, error) {
 27	if f, ok := reader.(File); !ok || f.Fd() != os.Stdin.Fd() {
 28		return newFallbackCancelReader(reader)
 29	}
 30
 31	// it is necessary to open CONIN$ (NOT windows.STD_INPUT_HANDLE) in
 32	// overlapped mode to be able to use it with WaitForMultipleObjects.
 33	conin, err := windows.CreateFile(
 34		&(utf16.Encode([]rune("CONIN$\x00"))[0]), windows.GENERIC_READ|windows.GENERIC_WRITE,
 35		fileShareValidFlags, nil, windows.OPEN_EXISTING, windows.FILE_FLAG_OVERLAPPED, 0)
 36	if err != nil {
 37		return nil, fmt.Errorf("open CONIN$ in overlapping mode: %w", err)
 38	}
 39
 40	resetConsole, err := prepareConsole(conin)
 41	if err != nil {
 42		return nil, fmt.Errorf("prepare console: %w", err)
 43	}
 44
 45	// flush input, otherwise it can contain events which trigger
 46	// WaitForMultipleObjects but which ReadFile cannot read, resulting in an
 47	// un-cancelable read
 48	err = flushConsoleInputBuffer(conin)
 49	if err != nil {
 50		return nil, fmt.Errorf("flush console input buffer: %w", err)
 51	}
 52
 53	cancelEvent, err := windows.CreateEvent(nil, 0, 0, nil)
 54	if err != nil {
 55		return nil, fmt.Errorf("create stop event: %w", err)
 56	}
 57
 58	return &winCancelReader{
 59		conin:              conin,
 60		cancelEvent:        cancelEvent,
 61		resetConsole:       resetConsole,
 62		blockingReadSignal: make(chan struct{}, 1),
 63	}, nil
 64}
 65
 66type winCancelReader struct {
 67	conin       windows.Handle
 68	cancelEvent windows.Handle
 69	cancelMixin
 70
 71	resetConsole       func() error
 72	blockingReadSignal chan struct{}
 73}
 74
 75func (r *winCancelReader) Read(data []byte) (int, error) {
 76	if r.isCanceled() {
 77		return 0, ErrCanceled
 78	}
 79
 80	err := r.wait()
 81	if err != nil {
 82		return 0, err
 83	}
 84
 85	if r.isCanceled() {
 86		return 0, ErrCanceled
 87	}
 88
 89	// windows.Read does not work on overlapping windows.Handles
 90	return r.readAsync(data)
 91}
 92
 93// Cancel cancels ongoing and future Read() calls and returns true if the
 94// cancelation of the ongoing Read() was successful. On Windows Terminal,
 95// WaitForMultipleObjects sometimes immediately returns without input being
 96// available. In this case, graceful cancelation is not possible and Cancel()
 97// returns false.
 98func (r *winCancelReader) Cancel() bool {
 99	r.setCanceled()
100
101	select {
102	case r.blockingReadSignal <- struct{}{}:
103		err := windows.SetEvent(r.cancelEvent)
104		if err != nil {
105			return false
106		}
107		<-r.blockingReadSignal
108	case <-time.After(100 * time.Millisecond):
109		// Read() hangs in a GetOverlappedResult which is likely due to
110		// WaitForMultipleObjects returning without input being available
111		// so we cannot cancel this ongoing read.
112		return false
113	}
114
115	return true
116}
117
118func (r *winCancelReader) Close() error {
119	err := windows.CloseHandle(r.cancelEvent)
120	if err != nil {
121		return fmt.Errorf("closing cancel event handle: %w", err)
122	}
123
124	err = r.resetConsole()
125	if err != nil {
126		return err
127	}
128
129	err = windows.Close(r.conin)
130	if err != nil {
131		return fmt.Errorf("closing CONIN$")
132	}
133
134	return nil
135}
136
137func (r *winCancelReader) wait() error {
138	event, err := windows.WaitForMultipleObjects([]windows.Handle{r.conin, r.cancelEvent}, false, windows.INFINITE)
139	switch {
140	case windows.WAIT_OBJECT_0 <= event && event < windows.WAIT_OBJECT_0+2:
141		if event == windows.WAIT_OBJECT_0+1 {
142			return ErrCanceled
143		}
144
145		if event == windows.WAIT_OBJECT_0 {
146			return nil
147		}
148
149		return fmt.Errorf("unexpected wait object is ready: %d", event-windows.WAIT_OBJECT_0)
150	case windows.WAIT_ABANDONED <= event && event < windows.WAIT_ABANDONED+2:
151		return fmt.Errorf("abandoned")
152	case event == uint32(windows.WAIT_TIMEOUT):
153		return fmt.Errorf("timeout")
154	case event == windows.WAIT_FAILED:
155		return fmt.Errorf("failed")
156	default:
157		return fmt.Errorf("unexpected error: %w", error(err))
158	}
159}
160
161// readAsync is necessary to read from a windows.Handle in overlapping mode.
162func (r *winCancelReader) readAsync(data []byte) (int, error) {
163	hevent, err := windows.CreateEvent(nil, 0, 0, nil)
164	if err != nil {
165		return 0, fmt.Errorf("create event: %w", err)
166	}
167
168	overlapped := windows.Overlapped{
169		HEvent: hevent,
170	}
171
172	var n uint32
173
174	err = windows.ReadFile(r.conin, data, &n, &overlapped)
175	if err != nil && err != windows.ERROR_IO_PENDING {
176		return int(n), err
177	}
178
179	r.blockingReadSignal <- struct{}{}
180	err = windows.GetOverlappedResult(r.conin, &overlapped, &n, true)
181	if err != nil {
182		return int(n), nil
183	}
184	<-r.blockingReadSignal
185
186	return int(n), nil
187}
188
189func prepareConsole(input windows.Handle) (reset func() error, err error) {
190	var originalMode uint32
191
192	err = windows.GetConsoleMode(input, &originalMode)
193	if err != nil {
194		return nil, fmt.Errorf("get console mode: %w", err)
195	}
196
197	var newMode uint32
198	newMode &^= windows.ENABLE_ECHO_INPUT
199	newMode &^= windows.ENABLE_LINE_INPUT
200	newMode &^= windows.ENABLE_MOUSE_INPUT
201	newMode &^= windows.ENABLE_WINDOW_INPUT
202	newMode &^= windows.ENABLE_PROCESSED_INPUT
203
204	newMode |= windows.ENABLE_EXTENDED_FLAGS
205	newMode |= windows.ENABLE_INSERT_MODE
206	newMode |= windows.ENABLE_QUICK_EDIT_MODE
207
208	// Enabling virtual terminal input is necessary for processing certain
209	// types of input like X10 mouse events and arrows keys with the current
210	// bytes-based input reader. It does, however, prevent cancelReader from
211	// being able to cancel input. The planned solution for this is to read
212	// Windows events in a more native fashion, rather than the current simple
213	// bytes-based input reader which works well on unix systems.
214	newMode |= windows.ENABLE_VIRTUAL_TERMINAL_INPUT
215
216	err = windows.SetConsoleMode(input, newMode)
217	if err != nil {
218		return nil, fmt.Errorf("set console mode: %w", err)
219	}
220
221	return func() error {
222		err := windows.SetConsoleMode(input, originalMode)
223		if err != nil {
224			return fmt.Errorf("reset console mode: %w", err)
225		}
226
227		return nil
228	}, nil
229}
230
231var (
232	modkernel32                 = windows.NewLazySystemDLL("kernel32.dll")
233	procFlushConsoleInputBuffer = modkernel32.NewProc("FlushConsoleInputBuffer")
234)
235
236func flushConsoleInputBuffer(consoleInput windows.Handle) error {
237	r, _, e := syscall.Syscall(procFlushConsoleInputBuffer.Addr(), 1,
238		uintptr(consoleInput), 0, 0)
239	if r == 0 {
240		return error(e)
241	}
242
243	return nil
244}