From 6cbc078c57d230c58f91ce0274e4ac1a5f370c9d Mon Sep 17 00:00:00 2001
From: chzyer <0@0xdf.com>
Date: Thu, 1 Sep 2016 18:13:06 +0800
Subject: [PATCH] let stdin cancelable (#72)

---
 std.go      | 52 +++++++++++++++++++++++++++++++++++++++++++++++++++-
 terminal.go | 14 +++++++++++---
 2 files changed, 62 insertions(+), 4 deletions(-)

diff --git a/std.go b/std.go
index eb72882..3e68aed 100644
--- a/std.go
+++ b/std.go
@@ -7,7 +7,7 @@ import (
 )
 
 var (
-	Stdin  io.ReadCloser  = os.Stdin
+	Stdin  io.ReadCloser  = NewCancelableStdin()
 	Stdout io.WriteCloser = os.Stdout
 	Stderr io.WriteCloser = os.Stderr
 )
@@ -64,3 +64,53 @@ func Line(prompt string) (string, error) {
 	ins.SetPrompt(prompt)
 	return ins.Readline()
 }
+
+type CancelableStdin struct {
+	mutex  sync.Mutex
+	stop   chan struct{}
+	notify chan struct{}
+	data   []byte
+	read   int
+	err    error
+}
+
+func NewCancelableStdin() *CancelableStdin {
+	c := &CancelableStdin{
+		notify: make(chan struct{}),
+		stop:   make(chan struct{}),
+	}
+	go c.ioloop()
+	return c
+}
+
+func (c *CancelableStdin) ioloop() {
+loop:
+	for {
+		select {
+		case <-c.notify:
+			c.read, c.err = os.Stdin.Read(c.data)
+			<-c.notify
+		case <-c.stop:
+			break loop
+		}
+	}
+}
+
+func (c *CancelableStdin) Read(b []byte) (n int, err error) {
+	c.mutex.Lock()
+	defer c.mutex.Unlock()
+
+	c.data = b
+	c.notify <- struct{}{}
+	select {
+	case <-c.notify:
+		return c.read, c.err
+	case <-c.stop:
+		return 0, io.EOF
+	}
+}
+
+func (c *CancelableStdin) Close() error {
+	close(c.stop)
+	return nil
+}
diff --git a/terminal.go b/terminal.go
index 36e1bf1..b3d1caf 100644
--- a/terminal.go
+++ b/terminal.go
@@ -3,6 +3,7 @@ package readline
 import (
 	"bufio"
 	"fmt"
+	"io"
 	"strings"
 	"sync"
 	"sync/atomic"
@@ -109,7 +110,11 @@ func (t *Terminal) KickRead() {
 
 func (t *Terminal) ioloop() {
 	t.wg.Add(1)
-	defer t.wg.Done()
+	defer func() {
+		t.wg.Done()
+		close(t.outchan)
+	}()
+
 	var (
 		isEscape       bool
 		isEscapeEx     bool
@@ -182,7 +187,7 @@ func (t *Terminal) ioloop() {
 			t.outchan <- r
 		}
 	}
-	close(t.outchan)
+
 }
 
 func (t *Terminal) Bell() {
@@ -193,7 +198,10 @@ func (t *Terminal) Close() error {
 	if atomic.SwapInt32(&t.closed, 1) != 0 {
 		return nil
 	}
-	t.stopChan <- struct{}{}
+	if closer, ok := t.cfg.Stdin.(io.Closer); ok {
+		closer.Close()
+	}
+	close(t.stopChan)
 	t.wg.Wait()
 	return t.ExitRawMode()
 }