From 900d0434b217b18519eb6a903b8bdb127118ec70 Mon Sep 17 00:00:00 2001
From: grisu48 <grisu42@gmail.com>
Date: Fri, 10 Feb 2023 10:26:38 +0100
Subject: [PATCH] Introduce ReadPasswordWithContext

Introduced the ReadPasswordWithContext function that allows to read a
password from the terminal with support for cancellation.

Signed-off-by: grisu48 <grisu42@gmail.com>
---
 term.go          |  9 +++++++
 term_unix.go     | 65 ++++++++++++++++++++++++++++++++++++++++++++++++
 term_windows.go  | 45 +++++++++++++++++++++++++++++++++
 terminal_test.go |  1 +
 4 files changed, 120 insertions(+)

diff --git a/term.go b/term.go
index 1a40d10..e8058dd 100644
--- a/term.go
+++ b/term.go
@@ -16,6 +16,8 @@
 // Note that on non-Unix systems os.Stdin.Fd() may not be 0.
 package term
 
+import "context"
+
 // State contains the state of a terminal.
 type State struct {
 	state
@@ -58,3 +60,10 @@ func GetSize(fd int) (width, height int, err error) {
 func ReadPassword(fd int) ([]byte, error) {
 	return readPassword(fd)
 }
+
+// ReadPasswordWithContext reads a line of input from a terminal without local
+// echo. This call is similar to ReadPassword but allows the input to be cancelled
+// via the provided Context.
+func ReadPasswordWithContext(fd int, ctx context.Context) ([]byte, error) {
+	return readPasswordWithContext(fd, ctx)
+}
diff --git a/term_unix.go b/term_unix.go
index a4e31ab..325f419 100644
--- a/term_unix.go
+++ b/term_unix.go
@@ -8,6 +8,10 @@
 package term
 
 import (
+	"context"
+	"syscall"
+	"time"
+
 	"golang.org/x/sys/unix"
 )
 
@@ -90,3 +94,64 @@ func readPassword(fd int) ([]byte, error) {
 
 	return readPasswordLine(passwordReader(fd))
 }
+
+func readPasswordWithContext(fd int, ctx context.Context) ([]byte, error) {
+	termios, err := unix.IoctlGetTermios(fd, ioctlReadTermios)
+	nonblocking := false
+	if err != nil {
+		return nil, err
+	}
+	newState := *termios
+	newState.Lflag &^= unix.ECHO
+	newState.Lflag |= unix.ICANON | unix.ISIG
+	newState.Iflag |= unix.ICRNL
+
+	if err := unix.IoctlSetTermios(fd, ioctlWriteTermios, &newState); err != nil {
+		return nil, err
+	}
+	defer func() {
+		if nonblocking {
+			unix.SetNonblock(fd, false)
+		}
+		unix.IoctlSetTermios(fd, ioctlWriteTermios, termios)
+	}()
+
+	// Set nonblocking IO
+	if err := unix.SetNonblock(fd, true); err != nil {
+		return nil, err
+	}
+	nonblocking = true
+
+	var ret []byte
+	var buf [1]byte
+	for {
+		if ctx.Err() != nil {
+			return ret, ctx.Err()
+		}
+		n, err := unix.Read(fd, buf[:])
+		if err != nil {
+			// Check for nonblocking error
+			if serr, ok := err.(syscall.Errno); ok {
+				if serr == 11 {
+					// Add (hopefully not noticable) latency to prevent CPU hogging
+					time.Sleep(50 * time.Millisecond)
+					continue
+				}
+			}
+			return ret, err
+		}
+		if n > 0 {
+			switch buf[0] {
+			case '\b':
+				if len(ret) > 0 {
+					ret = ret[:len(ret)-1]
+				}
+			case '\n':
+				return ret, nil
+			default:
+				ret = append(ret, buf[0])
+			}
+			continue
+		}
+	}
+}
diff --git a/term_windows.go b/term_windows.go
index 465f560..68851f6 100644
--- a/term_windows.go
+++ b/term_windows.go
@@ -5,6 +5,7 @@
 package term
 
 import (
+	"context"
 	"os"
 
 	"golang.org/x/sys/windows"
@@ -77,3 +78,47 @@ func readPassword(fd int) ([]byte, error) {
 	defer f.Close()
 	return readPasswordLine(f)
 }
+
+func readPasswordWithContext(fd int, ctx context.Context) ([]byte, error) {
+	var st uint32
+	if err := windows.GetConsoleMode(windows.Handle(fd), &st); err != nil {
+		return nil, err
+	}
+	old := st
+
+	st &^= (windows.ENABLE_ECHO_INPUT | windows.ENABLE_LINE_INPUT)
+	st |= (windows.ENABLE_PROCESSED_OUTPUT | windows.ENABLE_PROCESSED_INPUT)
+	if err := windows.SetConsoleMode(windows.Handle(fd), st); err != nil {
+		return nil, err
+	}
+
+	defer windows.SetConsoleMode(windows.Handle(fd), old)
+
+	var h windows.Handle
+	p, _ := windows.GetCurrentProcess()
+	if err := windows.DuplicateHandle(p, windows.Handle(fd), p, &h, 0, false, windows.DUPLICATE_SAME_ACCESS); err != nil {
+		return nil, err
+	}
+
+	f := os.NewFile(uintptr(h), "stdin")
+	defer f.Close()
+
+	// Buffer for reading in separate goroutine
+	type Buffer struct {
+		line []byte
+		err  error
+	}
+	doneChannel := make(chan Buffer, 1)
+	go func() {
+		// The following blocks and cannot be unblocked
+		ret, err := readPasswordLine(f)
+		doneChannel <- Buffer{line: ret, err: err}
+	}()
+	select {
+	case <-ctx.Done():
+		f.Close() // Blocks until terminal receives a return key :-(
+		return make([]byte, 0), ctx.Err()
+	case buf := <-doneChannel:
+		return buf.line, buf.err
+	}
+}
diff --git a/terminal_test.go b/terminal_test.go
index d5c1794..745e856 100644
--- a/terminal_test.go
+++ b/terminal_test.go
@@ -6,6 +6,7 @@ package term
 
 import (
 	"bytes"
+	"context"
 	"io"
 	"os"
 	"runtime"