Let copyLoop exit when either direction finishes.

Formerly we waiting until *both* directions finished. What this meant in
practice is that when the remote connection ended, copyLoop would become
useless but would continue blocking its caller until something else
finally closed the socks connection.
This commit is contained in:
David Fifield 2020-02-21 14:47:34 -07:00
parent ee2fb42d33
commit 904af9cb8a

View file

@ -5,7 +5,6 @@ import (
"io"
"log"
"net"
"sync"
"time"
)
@ -41,20 +40,19 @@ func Handler(socks net.Conn, snowflakes SnowflakeCollector) error {
// Exchanges bytes between two ReadWriters.
// (In this case, between a SOCKS and WebRTC connection.)
func copyLoop(socks, webRTC io.ReadWriter) {
var wg sync.WaitGroup
wg.Add(2)
done := make(chan struct{}, 2)
go func() {
if _, err := io.Copy(socks, webRTC); err != nil {
log.Printf("copying WebRTC to SOCKS resulted in error: %v", err)
}
wg.Done()
done <- struct{}{}
}()
go func() {
if _, err := io.Copy(webRTC, socks); err != nil {
log.Printf("copying SOCKS to WebRTC resulted in error: %v", err)
}
wg.Done()
done <- struct{}{}
}()
wg.Wait()
<-done
log.Println("copy loop ended")
}