Add timeout flag to dev client

code-asher · code-asher · commit 698c66a9d1cb · 2022-08-23T21:41:09.000-05:00
This makes it easier to test reconnects manually.
diff --git a/dev/client/main.go b/dev/client/main.go
@@ -9,6 +9,7 @@ import (
 	"os"
 	"os/signal"
 	"syscall"
+	"time"
 
 	"cdr.dev/wsep"
 	"github.com/spf13/pflag"
@@ -20,10 +21,11 @@ import (
 )
 
 type notty struct {
+	timeout time.Duration
 }
 
 func (c *notty) Run(fl *pflag.FlagSet) {
-	do(fl, false, "")
+	do(fl, false, "", c.timeout)
 }
 
 func (c *notty) Spec() cli.CommandSpec {
@@ -34,12 +36,17 @@ func (c *notty) Spec() cli.CommandSpec {
 	}
 }
 
+func (c *notty) RegisterFlags(fl *pflag.FlagSet) {
+	fl.DurationVar(&c.timeout, "timeout", 0, "disconnect after specified timeout")
+}
+
 type tty struct {
-	id string
+	id      string
+	timeout time.Duration
 }
 
 func (c *tty) Run(fl *pflag.FlagSet) {
-	do(fl, true, c.id)
+	do(fl, true, c.id, c.timeout)
 }
 
 func (c *tty) Spec() cli.CommandSpec {
@@ -52,9 +59,10 @@ func (c *tty) Spec() cli.CommandSpec {
 
 func (c *tty) RegisterFlags(fl *pflag.FlagSet) {
 	fl.StringVar(&c.id, "id", "", "sets id for reconnection")
+	fl.DurationVar(&c.timeout, "timeout", 0, "disconnect after the specified timeout")
 }
 
-func do(fl *pflag.FlagSet, tty bool, id string) {
+func do(fl *pflag.FlagSet, tty bool, id string, timeout time.Duration) {
 	ctx, cancel := context.WithCancel(context.Background())
 	defer cancel()
 
@@ -73,12 +81,18 @@ func do(fl *pflag.FlagSet, tty bool, id string) {
 	if len(fl.Args()) > 1 {
 		args = fl.Args()[1:]
 	}
+	width, height, err := term.GetSize(int(os.Stdin.Fd()))
+	if err != nil {
+		flog.Fatal("unable to get term size")
+	}
 	process, err := executor.Start(ctx, wsep.Command{
 		ID:      id,
 		Command: fl.Arg(0),
 		Args:    args,
 		TTY:     tty,
 		Stdin:   true,
+		Rows:    uint16(height),
+		Cols:    uint16(width),
 	})
 	if err != nil {
 		flog.Fatal("failed to start remote command: %v", err)
@@ -112,6 +126,15 @@ func do(fl *pflag.FlagSet, tty bool, id string) {
 		io.Copy(stdin, os.Stdin)
 	}()
 
+	if timeout != 0 {
+		timer := time.NewTimer(timeout)
+		defer timer.Stop()
+		go func() {
+			<-timer.C
+			conn.Close(websocket.StatusNormalClosure, "normal closure")
+		}()
+	}
+
 	err = process.Wait()
 	if err != nil {
 		flog.Error("process failed: %v", err)
