From 5824a13b57ef6277f42c27607406a571b546c140 Mon Sep 17 00:00:00 2001
From: Filippo Valsorda <hi@filippo.io>
Date: Tue, 24 May 2022 15:19:50 +0200
Subject: [PATCH] cmd/age,internal/plugin: print a message if waiting on a
 plugin for too long

---
 cmd/age/tui.go            |  3 +++
 internal/plugin/client.go | 20 ++++++++++++++++++--
 2 files changed, 21 insertions(+), 2 deletions(-)

diff --git a/cmd/age/tui.go b/cmd/age/tui.go
index 7aeb0a7..cc2fa95 100644
--- a/cmd/age/tui.go
+++ b/cmd/age/tui.go
@@ -181,4 +181,7 @@ var pluginTerminalUI = &plugin.ClientUI{
 			}
 		}
 	},
+	WaitTimer: func(name string) {
+		printf("waiting on %s plugin...", name)
+	},
 }
diff --git a/internal/plugin/client.go b/internal/plugin/client.go
index 5c52c9e..6f97167 100644
--- a/internal/plugin/client.go
+++ b/internal/plugin/client.go
@@ -15,6 +15,7 @@ import (
 	"os"
 	"strconv"
 	"strings"
+	"time"
 
 	exec "golang.org/x/sys/execabs"
 
@@ -87,7 +88,7 @@ func (r *Recipient) Wrap(fileKey []byte) (stanzas []*age.Stanza, err error) {
 	sr := format.NewStanzaReader(bufio.NewReader(conn))
 ReadLoop:
 	for {
-		s, err := sr.ReadStanza()
+		s, err := r.ui.readStanza(r.name, sr)
 		if err != nil {
 			return nil, err
 		}
@@ -228,7 +229,7 @@ func (i *Identity) Unwrap(stanzas []*age.Stanza) (fileKey []byte, err error) {
 	sr := format.NewStanzaReader(bufio.NewReader(conn))
 ReadLoop:
 	for {
-		s, err := sr.ReadStanza()
+		s, err := i.ui.readStanza(i.name, sr)
 		if err != nil {
 			return nil, err
 		}
@@ -297,6 +298,12 @@ type ClientUI struct {
 	// value are the choices provided to the user. no may be empty. The return
 	// value indicates whether the user selected the yes or no option.
 	Confirm func(name, prompt, yes, no string) (choseYes bool, err error)
+
+	// WaitTimer is invoked once (Un)Wrap has been waiting for 5 seconds on the
+	// plugin, for example because the plugin is waiting for an external event
+	// (e.g. a hardware token touch). Unlike the other callbacks, WaitTimer runs
+	// in a separate goroutine, and if missing it's simply ignored.
+	WaitTimer func(name string)
 }
 
 func (c *ClientUI) handle(name string, conn *clientConnection, s *format.Stanza) (ok bool, err error) {
@@ -350,6 +357,15 @@ func (c *ClientUI) handle(name string, conn *clientConnection, s *format.Stanza)
 	}
 }
 
+// readStanza calls r.ReadStanza and, if set, invokes WaitTimer in a separate
+// goroutine if the call takes longer than 5 seconds.
+func (c *ClientUI) readStanza(name string, r *format.StanzaReader) (*format.Stanza, error) {
+	if c.WaitTimer != nil {
+		defer time.AfterFunc(5*time.Second, func() { c.WaitTimer(name) }).Stop()
+	}
+	return r.ReadStanza()
+}
+
 type clientConnection struct {
 	cmd       *exec.Cmd
 	io.Reader // stdout