mirror of
https://github.com/wweir/sower.git
synced 2024-04-21 12:42:15 +00:00
39 lines
712 B
Go
39 lines
712 B
Go
package proxy
|
|
|
|
import (
|
|
"io"
|
|
"net"
|
|
"sync"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
func withDefaultPort(addr string, port string) (address, host string) {
|
|
host, _, err := net.SplitHostPort(addr)
|
|
if err != nil {
|
|
return net.JoinHostPort(addr, port), addr
|
|
}
|
|
return addr, host
|
|
}
|
|
|
|
func relay(conn1, conn2 net.Conn) {
|
|
wg := &sync.WaitGroup{}
|
|
exitFlag := new(int32)
|
|
wg.Add(2)
|
|
go redirect(conn2, conn1, wg, exitFlag)
|
|
redirect(conn1, conn2, wg, exitFlag)
|
|
wg.Wait()
|
|
}
|
|
func redirect(dst, src net.Conn, wg *sync.WaitGroup, exitFlag *int32) {
|
|
io.Copy(dst, src)
|
|
|
|
if atomic.CompareAndSwapInt32(exitFlag, 0, 1) {
|
|
// wakeup blocked goroutine
|
|
now := time.Now()
|
|
src.SetDeadline(now)
|
|
dst.SetDeadline(now)
|
|
}
|
|
|
|
wg.Done()
|
|
}
|