test drop and undrop in udproxy.

This commit is contained in:
Matt Jaffee 2018-10-26 09:17:26 -05:00
parent 28591b4a91
commit af7ece74fd
No known key found for this signature in database
GPG key ID: 08A3DFFF987B11BF
2 changed files with 28 additions and 2 deletions

View file

@ -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)")
}

View file

@ -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)