From 75a50797cca1b8647418460edd8a1148e7caa42f Mon Sep 17 00:00:00 2001 From: Takanari Hayama Date: Thu, 6 Dec 2018 13:22:16 +0900 Subject: [PATCH] Wait for connect to finish sending all data Goroutines created in connect, transfer data from a handler to telnet client. However, as there's no synchronization between the goroutines and the telnet server, there's a case where the server sends a prompt before the goroutines finishsed all transmissions. This fix assure that the telnet server waits until all data are sent to the client, before sending a prompt. --- telsh/telnet_handler.go | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/telsh/telnet_handler.go b/telsh/telnet_handler.go index ca08c7b..755b816 100644 --- a/telsh/telnet_handler.go +++ b/telsh/telnet_handler.go @@ -229,27 +229,35 @@ func (telnetHandler *ShellHandler) ServeTELNET(ctx telnet.Context, writer telnet //@TODO: Wire up the stdin, stdout, stderr of the handler. + var stdoutDone chan bool if stdoutPipe, err := handler.StdoutPipe(); nil != err { //@TODO: } else if nil == stdoutPipe { //@TODO: } else { - connect(ctx, writer, stdoutPipe) + stdoutDone = connect(ctx, writer, stdoutPipe) } + var stderrDone chan bool if stderrPipe, err := handler.StderrPipe(); nil != err { //@TODO: } else if nil == stderrPipe { //@TODO: } else { - connect(ctx, writer, stderrPipe) + stderrDone = connect(ctx, writer, stderrPipe) } if err := handler.Run(); nil != err { //@TODO: } + if stdoutDone != nil { + <-stdoutDone + } + if stderrDone != nil { + <-stderrDone + } line.Reset() if _, err := oi.LongWrite(writer, promptBytes); nil != err { return @@ -270,10 +278,13 @@ func (telnetHandler *ShellHandler) ServeTELNET(ctx telnet.Context, writer telnet -func connect(ctx telnet.Context, writer io.Writer, reader io.Reader) { +func connect(ctx telnet.Context, writer io.Writer, reader io.Reader) chan bool { logger := ctx.Logger() + // Create a buffer to avoid the goroutine to be blocked. + // The goroutine may leave whenever its job is done. + done := make(chan bool, 1) go func(logger telnet.Logger){ var buffer [1]byte // Seems like the length of the buffer needs to be small, otherwise will have to wait for buffer to fill up. @@ -293,5 +304,7 @@ func connect(ctx telnet.Context, writer io.Writer, reader io.Reader) { oi.LongWrite(writer, p) //logger.Tracef("Sent: %q.", p) } + done <- true }(logger) + return done }