69 lines
1.6 KiB
Go
69 lines
1.6 KiB
Go
// Package udprelay preserves UDP datagram boundaries across a host socket flow.
|
|
package udprelay
|
|
|
|
import (
|
|
"context"
|
|
"errors"
|
|
"net"
|
|
"sync/atomic"
|
|
"time"
|
|
)
|
|
|
|
// Relay copies connected UDP datagrams until idle timeout or cancellation.
|
|
func Relay(ctx context.Context, left, right net.Conn, idleTimeout time.Duration) error {
|
|
parent := ctx
|
|
ctx, cancel := context.WithCancel(parent)
|
|
defer cancel()
|
|
if idleTimeout <= 0 {
|
|
idleTimeout = 60 * time.Second
|
|
}
|
|
var lastActivity atomic.Int64
|
|
lastActivity.Store(time.Now().UnixNano())
|
|
results := make(chan error, 2)
|
|
copyDatagrams := func(destination, source net.Conn) {
|
|
buffer := make([]byte, 65535)
|
|
for {
|
|
deadline := time.Now().Add(min(idleTimeout/2, time.Second))
|
|
_ = source.SetReadDeadline(deadline)
|
|
count, err := source.Read(buffer)
|
|
if err != nil {
|
|
if timeout, ok := err.(net.Error); ok && timeout.Timeout() {
|
|
last := time.Unix(0, lastActivity.Load())
|
|
if time.Since(last) < idleTimeout {
|
|
continue
|
|
}
|
|
}
|
|
results <- err
|
|
return
|
|
}
|
|
if _, err := destination.Write(buffer[:count]); err != nil {
|
|
results <- err
|
|
return
|
|
}
|
|
lastActivity.Store(time.Now().UnixNano())
|
|
}
|
|
}
|
|
go copyDatagrams(left, right)
|
|
go copyDatagrams(right, left)
|
|
go func() {
|
|
<-ctx.Done()
|
|
_ = left.Close()
|
|
_ = right.Close()
|
|
}()
|
|
err := <-results
|
|
cancel()
|
|
_ = left.Close()
|
|
_ = right.Close()
|
|
<-results
|
|
if parent.Err() != nil {
|
|
return parent.Err()
|
|
}
|
|
if timeout, ok := err.(net.Error); ok && timeout.Timeout() {
|
|
return nil
|
|
}
|
|
if errors.Is(err, net.ErrClosed) {
|
|
return nil
|
|
}
|
|
return err
|
|
}
|