ssh/tailssh: only call CloseWrite when both stdout and stderr are done

Updates #5209

Signed-off-by: Maisem Ali <maisem@tailscale.com>
This commit is contained in:
Maisem Ali 2022-11-15 11:46:15 -08:00 committed by Maisem Ali
parent 6ea2d01626
commit 2d653230ef

View File

@ -28,6 +28,7 @@
"strconv" "strconv"
"strings" "strings"
"sync" "sync"
"sync/atomic"
"time" "time"
gossh "github.com/tailscale/golang-x-crypto/ssh" gossh "github.com/tailscale/golang-x-crypto/ssh"
@ -1069,13 +1070,20 @@ func (ss *sshSession) run() {
ss.ctx.CloseWithError(err) ss.ctx.CloseWithError(err)
} }
}() }()
var openOutputStreams atomic.Int32
if ss.stderr != nil {
openOutputStreams.Store(2)
} else {
openOutputStreams.Store(1)
}
go func() { go func() {
defer ss.stdout.Close() defer ss.stdout.Close()
_, err := io.Copy(rec.writer("o", ss), ss.stdout) _, err := io.Copy(rec.writer("o", ss), ss.stdout)
if err != nil && !errors.Is(err, io.EOF) { if err != nil && !errors.Is(err, io.EOF) {
logf("stdout copy: %v", err) logf("stdout copy: %v", err)
ss.ctx.CloseWithError(err) ss.ctx.CloseWithError(err)
} else { }
if openOutputStreams.Add(-1) == 0 {
ss.CloseWrite() ss.CloseWrite()
} }
}() }()
@ -1086,6 +1094,9 @@ func (ss *sshSession) run() {
if err != nil { if err != nil {
logf("stderr copy: %v", err) logf("stderr copy: %v", err)
} }
if openOutputStreams.Add(-1) == 0 {
ss.CloseWrite()
}
}() }()
} }