Repository: pion/dtls
Pull request #181 (2020-02-02 04:08:35)

Make Listener compatible with net.Listener
Breaking changes:

1. Remove timeout argument from `listener.Close`
Already accepted connections should be active after `listener.Close()`.
Previously, the connections were broken after the timeout.
In the new implementation, `readLoop` routine is remained after
`listener.Close()` and automatically stopped after closing
all connections. The final `conn.Close()` returns after stopping
the routine.

2. ~Return "use of closed connection" error~
~Double `Close()` calls returned `nil`, and `Write()` to closed connection
returned `io.EOF`. Return `"use of closed connection"` error on them to
make the behaviour close to `tls.Conn`.~

3. Unexpose listener implementation
Return as `net.Listener` interface instead of `*dtls.Listener`.

### Reference issue
Part of #148


codecov-io: # [Codecov](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=h1) Report
> Merging [#181](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=desc) into [master](https://codecov.io/gh/pion/dtls/commit/c9c352a9649817a84dc6c6b1be893c760b30e321?src=pr&el=desc) will **decrease** coverage by `0.26%`.
> The diff coverage is `93.18%`.

[![Impacted file tree graph](https://codecov.io/gh/pion/dtls/pull/181/graphs/tree.svg?width=650&token=bcCB2jB88I&height=150&src=pr)](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=tree)

```diff
@@            Coverage Diff             @@
##           master     #181      +/-   ##
==========================================
- Coverage    72.6%   72.34%   -0.27%     
==========================================
  Files          60       60              
  Lines        3431     3413      -18     
==========================================
- Hits         2491     2469      -22     
- Misses        649      652       +3     
- Partials      291      292       +1
```

| Flag | Coverage Δ | |
|---|---|---|
| #go | `72.22% <93.18%> (-0.38%)` | :arrow_down: |
| #wasm | `69.08% <0%> (+0.36%)` | :arrow_up: |

| [Impacted Files](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [listener.go](https://codecov.io/gh/pion/dtls/pull/181/diff?src=pr&el=tree#diff-bGlzdGVuZXIuZ28=) | `57.89% <83.33%> (ø)` | :arrow_up: |
| [internal/udp/conn.go](https://codecov.io/gh/pion/dtls/pull/181/diff?src=pr&el=tree#diff-aW50ZXJuYWwvdWRwL2Nvbm4uZ28=) | `82.17% <94.73%> (-0.18%)` | :arrow_down: |
| [config.go](https://codecov.io/gh/pion/dtls/pull/181/diff?src=pr&el=tree#diff-Y29uZmlnLmdv) | `95.23% <0%> (-4.77%)` | :arrow_down: |
| [crypto.go](https://codecov.io/gh/pion/dtls/pull/181/diff?src=pr&el=tree#diff-Y3J5cHRvLmdv) | `53.84% <0%> (-3.08%)` | :arrow_down: |
| [util.go](https://codecov.io/gh/pion/dtls/pull/181/diff?src=pr&el=tree#diff-dXRpbC5nbw==) | `85.29% <0%> (-1.97%)` | :arrow_down: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=continue).
> **Legend** - [Click here to learn more](https://docs.codecov.io/docs/codecov-delta)
> `Δ = absolute <relative> (impact)`, `ø = not affected`, `? = missing data`
> Powered by [Codecov](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=footer). Last update [c9c352a...68db4fe](https://codecov.io/gh/pion/dtls/pull/181?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).

at-wat: `Accept()`ed connection looks be allowed to be double closed in `net.TCPConn`.
Relaxing this part.
at-wat: Adding a test to check Close error propagation.
at-wat: @Sean-Der please take another look when you have time,
Sean-Der: Fantastic work! We should have been more strict about implementing stdlib interfaces.

I really love in the other PR that we are using the in-tree tests now :)

Patch:
diff --git a/examples/listen-psk/main.go b/examples/listen-psk/main.go
index f5ebc99..f2aed48 100644
--- a/examples/listen-psk/main.go
+++ b/examples/listen-psk/main.go
@@ -33,7 +33,7 @@ func main() {
 	listener, err := dtls.Listen("udp", addr, config)
 	util.Check(err)
 	defer func() {
-		util.Check(listener.Close(5 * time.Second))
+		util.Check(listener.Close())
 	}()
 
 	fmt.Println("Listening")
diff --git a/examples/listen/main.go b/examples/listen/main.go
index 436b52b..ec6e291 100644
--- a/examples/listen/main.go
+++ b/examples/listen/main.go
@@ -34,7 +34,7 @@ func main() {
 	listener, err := dtls.Listen("udp", addr, config)
 	util.Check(err)
 	defer func() {
-		util.Check(listener.Close(5 * time.Second))
+		util.Check(listener.Close())
 	}()
 
 	fmt.Println("Listening")
diff --git a/internal/udp/conn.go b/internal/udp/conn.go
index 5aa3487..ab65300 100644
--- a/internal/udp/conn.go
+++ b/internal/udp/conn.go
@@ -6,11 +6,11 @@ import (
 	"io"
 	"net"
 	"sync"
+	"sync/atomic"
 	"time"
 )
 
 const receiveMTU = 8192
-const closeRetryInterval = 100 * time.Millisecond
 
 var errClosedListener = errors.New("udp: listener closed")
 
@@ -18,20 +18,24 @@ var errClosedListener = errors.New("udp: listener closed")
 type Listener struct {
 	pConn *net.UDPConn
 
-	lock      sync.RWMutex
-	accepting bool
+	accepting atomic.Value // bool
 	acceptCh  chan *Conn
 	doneCh    chan struct{}
 	doneOnce  sync.Once
 
-	conns map[string]*Conn
+	connLock sync.Mutex
+	conns    map[string]*Conn
+	connWG   sync.WaitGroup
+
+	readWG   sync.WaitGroup
+	errClose atomic.Value // error
 }
 
 // Accept waits for and returns the next connection to the listener.
-// You have to either close or read on all connection that are created.
 func (l *Listener) Accept() (*Conn, error) {
 	select {
 	case c := <-l.acceptCh:
+		l.connWG.Add(1)
 		return c, nil
 
 	case <-l.doneCh:
@@ -41,41 +45,29 @@ func (l *Listener) Accept() (*Conn, error) {
 
 // Close closes the listener.
 // Any blocked Accept operations will be unblocked and return errors.
-func (l *Listener) Close(shutdownTimeout time.Duration) error {
-	l.lock.Lock()
-	defer l.lock.Unlock()
-
+func (l *Listener) Close() error {
 	var err error
 	l.doneOnce.Do(func() {
-		l.accepting = false
+		l.connWG.Done()
+		l.accepting.Store(false)
 		close(l.doneCh)
-		err = l.cleanupWithTimeout(shutdownTimeout)
-	})
 
-	return err
-}
+		l.connLock.Lock()
+		nConns := len(l.conns)
+		l.connLock.Unlock()
 
-func (l *Listener) cleanupWithTimeout(shutdownTimeout time.Duration) error {
-	timeoutTimer := time.NewTimer(shutdownTimeout)
-	for {
-		select {
-		case <-time.After(closeRetryInterval):
-			if len(l.conns) == 0 {
-				return l.cleanup()
+		if nConns == 0 {
+			// Wait if this is the final connection
+			l.readWG.Wait()
+			if errClose, ok := l.errClose.Load().(error); ok {
+				err = errClose
 			}
-		case <-timeoutTimer.C:
-			return l.cleanup()
+		} else {
+			err = nil
 		}
-	}
-}
+	})
 
-// cleanup closes the packet conn if it is no longer used
-// The caller should hold the read lock.
-func (l *Listener) cleanup() error {
-	if !l.accepting && len(l.conns) == 0 {
-		return l.pConn.Close()
-	}
-	return nil
+	return err
 }
 
 // Addr returns the listener's network address.
@@ -91,14 +83,23 @@ func Listen(network string, laddr *net.UDPAddr) (*Listener, error) {
 	}
 
 	l := &Listener{
-		pConn:     conn,
-		acceptCh:  make(chan *Conn),
-		conns:     make(map[string]*Conn),
-		accepting: true,
-		doneCh:    make(chan struct{}),
+		pConn:    conn,
+		acceptCh: make(chan *Conn),
+		conns:    make(map[string]*Conn),
+		doneCh:   make(chan struct{}),
 	}
+	l.accepting.Store(true)
+	l.connWG.Add(1)
+	l.readWG.Add(2) // wait readLoop and Close execution routine
 
 	go l.readLoop()
+	go func() {
+		l.connWG.Wait()
+		if err := l.pConn.Close(); err != nil {
+			l.errClose.Store(err)
+		}
+		l.readWG.Done()
+	}()
 
 	return l, nil
 }
@@ -108,9 +109,9 @@ func Listen(network string, laddr *net.UDPAddr) (*Listener, error) {
 //    It can therefore not be ended until all Conns are closed.
 // 2. Creating a new Conn when receiving from a new remote.
 func (l *Listener) readLoop() {
+	defer l.readWG.Done()
 	buf := make([]byte, receiveMTU)
 
-readLoop:
 	for {
 		n, raddr, err := l.pConn.ReadFrom(buf)
 		if err != nil {
@@ -120,23 +121,18 @@ readLoop:
 		if err != nil {
 			continue
 		}
-		select {
-		case cBuf := <-conn.readCh:
-			n = copy(cBuf, buf[:n])
-			conn.sizeCh <- n
-		case <-conn.doneCh:
-			continue readLoop
-		}
+		cBuf := <-conn.readCh
+		n = copy(cBuf, buf[:n])
+		conn.sizeCh <- n
 	}
 }
 
 func (l *Listener) getConn(raddr net.Addr) (*Conn, error) {
-	l.lock.Lock()
-	defer l.lock.Unlock()
-
+	l.connLock.Lock()
+	defer l.connLock.Unlock()
 	conn, ok := l.conns[raddr.String()]
 	if !ok {
-		if !l.accepting {
+		if !l.accepting.Load().(bool) {
 			return nil, errClosedListener
 		}
 		conn = l.newConn(raddr)
@@ -155,7 +151,6 @@ type Conn struct {
 	readCh chan []byte
 	sizeCh chan int
 
-	lock     sync.RWMutex
 	doneCh   chan struct{}
 	doneOnce sync.Once
 }
@@ -183,30 +178,29 @@ func (c *Conn) Read(p []byte) (int, error) {
 
 // Write writes len(p) bytes from p to the DTLS connection
 func (c *Conn) Write(p []byte) (n int, err error) {
-	c.lock.Lock()
-	l := c.listener
-	c.lock.Unlock()
-
-	if l == nil {
-		return 0, io.EOF
-	}
-
-	return l.pConn.WriteTo(p, c.rAddr)
+	return c.listener.pConn.WriteTo(p, c.rAddr)
 }
 
 // Close closes the conn and releases any Read calls
 func (c *Conn) Close() error {
-	c.lock.Lock()
-	defer c.lock.Unlock()
-
 	var err error
 	c.doneOnce.Do(func() {
+		c.listener.connWG.Done()
 		close(c.doneCh)
-		c.listener.lock.Lock()
+		c.listener.connLock.Lock()
 		delete(c.listener.conns, c.rAddr.String())
-		err = c.listener.cleanup()
-		c.listener.lock.Unlock()
-		c.listener = nil
+		nConns := len(c.listener.conns)
+		c.listener.connLock.Unlock()
+
+		if nConns == 0 && !c.listener.accepting.Load().(bool) {
+			// Wait if this is the final connection
+			c.listener.readWG.Wait()
+			if errClose, ok := c.listener.errClose.Load().(error); ok {
+				err = errClose
+			}
+		} else {
+			err = nil
+		}
 	})
 
 	return err
@@ -214,15 +208,7 @@ func (c *Conn) Close() error {
 
 // LocalAddr is a stub
 func (c *Conn) LocalAddr() net.Addr {
-	c.lock.Lock()
-	l := c.listener
-	c.lock.Unlock()
-
-	if l == nil {
-		return nil
-	}
-
-	return l.pConn.LocalAddr()
+	return c.listener.pConn.LocalAddr()
 }
 
 // RemoteAddr is a stub
diff --git a/listener.go b/listener.go
index c64a66f..d5151c6 100644
--- a/listener.go
+++ b/listener.go
@@ -2,13 +2,12 @@ package dtls
 
 import (
 	"net"
-	"time"
 
 	"github.com/pion/dtls/v2/internal/udp"
 )
 
 // Listen creates a DTLS listener
-func Listen(network string, laddr *net.UDPAddr, config *Config) (*Listener, error) {
+func Listen(network string, laddr *net.UDPAddr, config *Config) (net.Listener, error) {
 	if err := validateConfig(config); err != nil {
 		return nil, err
 	}
@@ -17,21 +16,21 @@ func Listen(network string, laddr *net.UDPAddr, config *Config) (*Listener, erro
 	if err != nil {
 		return nil, err
 	}
-	return &Listener{
+	return &listener{
 		config: config,
 		parent: parent,
 	}, nil
 }
 
 // Listener represents a DTLS listener
-type Listener struct {
+type listener struct {
 	config *Config
 	parent *udp.Listener
 }
 
 // Accept waits for and returns the next connection to the listener.
 // You have to either close or read on all connection that are created.
-func (l *Listener) Accept() (net.Conn, error) {
+func (l *listener) Accept() (net.Conn, error) {
 	c, err := l.parent.Accept()
 	if err != nil {
 		return nil, err
@@ -42,11 +41,11 @@ func (l *Listener) Accept() (net.Conn, error) {
 // Close closes the listener.
 // Any blocked Accept operations will be unblocked and return errors.
 // Already Accepted connections are not closed.
-func (l *Listener) Close(shutdownTimeout time.Duration) error {
-	return l.parent.Close(shutdownTimeout)
+func (l *listener) Close() error {
+	return l.parent.Close()
 }
 
 // Addr returns the listener's network address.
-func (l *Listener) Addr() net.Addr {
+func (l *listener) Addr() net.Addr {
 	return l.parent.Addr()
 }


Test patch:
diff --git a/e2e/e2e_test.go b/e2e/e2e_test.go
index 073a1a4..c633af7 100644
--- a/e2e/e2e_test.go
+++ b/e2e/e2e_test.go
@@ -71,7 +71,7 @@ func assertE2ECommunication(clientConfig, serverConfig *dtls.Config, serverPort
 		clientConn       net.Conn
 		serverMutex      sync.Mutex
 		serverConn       net.Conn
-		serverListener   *dtls.Listener
+		serverListener   net.Listener
 		serverReady      = make(chan struct{})
 		errChan          = make(chan error)
 		clientChan       = make(chan string)
@@ -141,7 +141,7 @@ func assertE2ECommunication(clientConfig, serverConfig *dtls.Config, serverPort
 			t.Fatal(err)
 		}
 
-		if err := serverListener.Close(2 * time.Second); err != nil {
+		if err := serverListener.Close(); err != nil {
 			t.Fatal(err)
 		}
 	}()
diff --git a/internal/udp/conn_test.go b/internal/udp/conn_test.go
index 0059704..91178d7 100644
--- a/internal/udp/conn_test.go
+++ b/internal/udp/conn_test.go
@@ -58,7 +58,7 @@ func stressDuplex(t *testing.T) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		err = listener.Close(5 * time.Second)
+		err = listener.Close()
 		if err != nil {
 			t.Fatal(err)
 		}
@@ -89,7 +89,7 @@ func TestListenerCloseTimeout(t *testing.T) {
 		t.Fatal(err)
 	}
 
-	err = listener.Close(500 * time.Millisecond)
+	err = listener.Close()
 	if err != nil {
 		t.Fatal(err)
 	}
@@ -149,24 +149,75 @@ func getConfig() (string, *net.UDPAddr) {
 	return "udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0}
 }
 
-// func TestConnClose(t *testing.T) {
-// 	lim := test.TimeOut(time.Second * 5)
-// 	defer lim.Stop()
-//
-// 	listener, ca, cb, err := pipe()
-// 	if err != nil {
-// 		t.Fatal(err)
-// 	}
-// 	err = ca.Close()
-// 	if err != nil {
-// 		t.Fatalf("Failed to close A side: %v\n", err)
-// 	}
-// 	err = cb.Close()
-// 	if err != nil {
-// 		t.Fatalf("Failed to close B side: %v\n", err)
-// 	}
-//  err = listener.Close(1 * time.Second)
-//  if err != nil {
-//  	t.Fatalf("Failed to close listener: %v\n", err)
-//  }
-// }
+func TestConnClose(t *testing.T) {
+	lim := test.TimeOut(time.Second * 5)
+	defer lim.Stop()
+
+	t.Run("Close", func(t *testing.T) {
+		// Check for leaking routines
+		report := test.CheckRoutines(t)
+		defer report()
+
+		listener, ca, cb, errPipe := pipe()
+		if errPipe != nil {
+			t.Fatal(errPipe)
+		}
+		if err := ca.Close(); err != nil {
+			t.Errorf("Failed to close A side: %v", err)
+		}
+		if err := cb.Close(); err != nil {
+			t.Errorf("Failed to close B side: %v", err)
+		}
+		if err := listener.Close(); err != nil {
+			t.Errorf("Failed to close listener: %v", err)
+		}
+	})
+	t.Run("CloseError1", func(t *testing.T) {
+		// Check for leaking routines
+		report := test.CheckRoutines(t)
+		defer report()
+
+		listener, ca, cb, errPipe := pipe()
+		if errPipe != nil {
+			t.Fatal(errPipe)
+		}
+		// Close listener.pConn to inject error.
+		if err := listener.pConn.Close(); err != nil {
+			t.Error(err)
+		}
+
+		if err := cb.Close(); err != nil {
+			t.Errorf("Failed to close A side: %v", err)
+		}
+		if err := ca.Close(); err != nil {
+			t.Errorf("Failed to close B side: %v", err)
+		}
+		if err := listener.Close(); err == nil {
+			t.Errorf("Error is not propagated to Listener.Close")
+		}
+	})
+	t.Run("CloseError2", func(t *testing.T) {
+		// Check for leaking routines
+		report := test.CheckRoutines(t)
+		defer report()
+
+		listener, ca, cb, errPipe := pipe()
+		if errPipe != nil {
+			t.Fatal(errPipe)
+		}
+		// Close listener.pConn to inject error.
+		if err := listener.pConn.Close(); err != nil {
+			t.Error(err)
+		}
+
+		if err := cb.Close(); err != nil {
+			t.Errorf("Failed to close A side: %v", err)
+		}
+		if err := listener.Close(); err != nil {
+			t.Errorf("Failed to close listener: %v", err)
+		}
+		if err := ca.Close(); err == nil {
+			t.Errorf("Error is not propagated to Conn.Close")
+		}
+	})
+}