diff --git a/internal/udproxy/udproxy.go b/internal/udproxy/udproxy.go index fa30b22a5..a79c0121f 100644 --- a/internal/udproxy/udproxy.go +++ b/internal/udproxy/udproxy.go @@ -47,12 +47,18 @@ func New(listenIP string, listenPort int, upstreamIP string, upstreamPort int) ( } func (p *Proxy) Drop() { + // sleep to attempt to ensure all traffic that was supposed to pass, did + // pass. + time.Sleep(time.Millisecond * 3) p.dropLock.Lock() p.drop = true p.dropLock.Unlock() } func (p *Proxy) Undrop() { + // this sleep is a cheap attempt to ensure everything that was supposed to + // be dropped was. + time.Sleep(time.Millisecond * 3) p.dropLock.Lock() p.drop = false p.dropLock.Unlock() @@ -73,7 +79,7 @@ func (p *Proxy) run() error { return nil default: } - err := p.conn.SetReadDeadline(time.Now().Add(time.Millisecond * 10)) + err := p.conn.SetReadDeadline(time.Now().Add(time.Millisecond)) if err != nil { return errors.Wrap(err, "setting read deadline (run)") } @@ -111,7 +117,7 @@ func (p *Proxy) proxyBack(to *net.UDPAddr, from *net.UDPConn) error { return nil default: } - err := from.SetReadDeadline(time.Now().Add(time.Millisecond * 10)) + err := from.SetReadDeadline(time.Now().Add(time.Millisecond)) if err != nil { return errors.Wrap(err, "setting read deadline (proxyBack)") } diff --git a/internal/udproxy/udproxy_test.go b/internal/udproxy/udproxy_test.go index edb915d1f..be746ea72 100644 --- a/internal/udproxy/udproxy_test.go +++ b/internal/udproxy/udproxy_test.go @@ -35,6 +35,16 @@ func TestUDProxy(t *testing.T) { if err != nil { return errors.Wrap(err, "reading from proxy") } + p.Drop() + _, err = conn.Write([]byte("hello2")) + if err != nil { + return errors.Wrap(err, "writing to dropping proxy") + } + p.Undrop() + _, err = conn.Write([]byte("hello3")) + if err != nil { + return errors.Wrap(err, "writing to undropping proxy") + } return nil }) @@ -50,10 +60,20 @@ func TestUDProxy(t *testing.T) { if err != nil { t.Fatalf("writing response: %v", err) } + + _, _, err = uc.ReadFrom(req) + if err != nil { + t.Fatalf("upstream reading from proxy: %v", err) + } + if string(req[:6]) != "hello3" { + t.Fatalf("got unexpected request %s", req) + } + eg.Wait() if string(resp[:7]) != "goodbye" { t.Fatalf("got unexpected response '%v", resp) } + err = p.Close() if err != nil { t.Fatalf("err closing proxy: '%v'", err)