From 4ef5969ea9497d1755e4657f09e735effda0288a Mon Sep 17 00:00:00 2001
From: Lonny Wong <lonnywong@qq.com>
Date: Sat, 27 Jul 2024 12:50:31 +0800
Subject: [PATCH] fix "early EOF" issue when using with git

---
 tssh/main.go  |  5 +++++
 tssh/trzsz.go | 26 +++++++++++++++++---------
 2 files changed, 22 insertions(+), 9 deletions(-)

diff --git a/tssh/main.go b/tssh/main.go
index a3c8fbd..115573b 100644
--- a/tssh/main.go
+++ b/tssh/main.go
@@ -272,5 +272,10 @@ func sshStart(args *sshArgs) error {
 	if args.Background {
 		_ = ss.client.Wait()
 	}
+
+	// wait for the output to be read by the parent process
+	if !isTerminal {
+		outputWaitGroup.Wait()
+	}
 	return nil
 }
diff --git a/tssh/trzsz.go b/tssh/trzsz.go
index 4a1a87c..ee704a9 100644
--- a/tssh/trzsz.go
+++ b/tssh/trzsz.go
@@ -32,11 +32,14 @@ import (
 	"os"
 	"runtime"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/trzsz/trzsz-go/trzsz"
 )
 
+var outputWaitGroup sync.WaitGroup
+
 func writeAll(dst io.Writer, data []byte) error {
 	m := 0
 	l := len(data)
@@ -53,6 +56,11 @@ func writeAll(dst io.Writer, data []byte) error {
 func wrapStdIO(serverIn io.WriteCloser, serverOut io.Reader, serverErr io.Reader, tty bool) {
 	win := runtime.GOOS == "windows"
 	forwardIO := func(reader io.Reader, writer io.WriteCloser, input bool) {
+		done := true
+		if !input {
+			done = false
+			outputWaitGroup.Add(1)
+		}
 		defer writer.Close()
 		buffer := make([]byte, 32*1024)
 		for {
@@ -72,20 +80,20 @@ func wrapStdIO(serverIn io.WriteCloser, serverOut io.Reader, serverErr io.Reader
 				}
 			}
 			if err == io.EOF {
-				if win && tty && input {
+				if win && isTerminal && tty && input {
 					_, _ = writer.Write([]byte{0x1A}) // ctrl + z
 					continue
 				}
-				if !input {
-					continue // ignore output EOF
+				if input {
+					return // input EOF
 				}
-				// delay and close
-				for {
-					time.Sleep(time.Second)
-					if lastTime := lastServerAliveTime.Load(); lastTime != nil && time.Since(*lastTime) > 2*time.Second {
-						return
-					}
+				// ignore output EOF
+				if !done {
+					outputWaitGroup.Done()
+					done = true
 				}
+				time.Sleep(100 * time.Millisecond)
+				continue
 			}
 			if err != nil {
 				return