mirror of https://github.com/coder/coder.git
48 lines
919 B
Go
48 lines
919 B
Go
package agentssh
|
|
|
|
import (
|
|
"context"
|
|
"io"
|
|
"sync"
|
|
)
|
|
|
|
// Bicopy copies all of the data between the two connections and will close them
|
|
// after one or both of them are done writing. If the context is canceled, both
|
|
// of the connections will be closed.
|
|
func Bicopy(ctx context.Context, c1, c2 io.ReadWriteCloser) {
|
|
ctx, cancel := context.WithCancel(ctx)
|
|
defer cancel()
|
|
|
|
defer func() {
|
|
_ = c1.Close()
|
|
_ = c2.Close()
|
|
}()
|
|
|
|
var wg sync.WaitGroup
|
|
copyFunc := func(dst io.WriteCloser, src io.Reader) {
|
|
defer func() {
|
|
wg.Done()
|
|
// If one side of the copy fails, ensure the other one exits as
|
|
// well.
|
|
cancel()
|
|
}()
|
|
_, _ = io.Copy(dst, src)
|
|
}
|
|
|
|
wg.Add(2)
|
|
go copyFunc(c1, c2)
|
|
go copyFunc(c2, c1)
|
|
|
|
// Convert waitgroup to a channel so we can also wait on the context.
|
|
done := make(chan struct{})
|
|
go func() {
|
|
defer close(done)
|
|
wg.Wait()
|
|
}()
|
|
|
|
select {
|
|
case <-ctx.Done():
|
|
case <-done:
|
|
}
|
|
}
|