Repository: pion/dtls
Pull request #257 (2020-05-27 10:30:31)

Start handshake only by handshake in Accept()
### Description
Avoid returning error from Accept() when receiving unrelated packet.
It also avoid allocating resources for unrelated packets.

### Reference issue
#255 


codecov[bot]: # [Codecov](https://codecov.io/gh/pion/dtls/pull/257?src=pr&el=h1) Report
> Merging [#257](https://codecov.io/gh/pion/dtls/pull/257?src=pr&el=desc) into [master](https://codecov.io/gh/pion/dtls/commit/52af2bfea38ea292b1f63b2067d3cb5aefd35430&el=desc) will **decrease** coverage by `0.04%`.
> The diff coverage is `84.00%`.

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

```diff
@@            Coverage Diff             @@
##           master     #257      +/-   ##
==========================================
- Coverage   77.77%   77.73%   -0.05%     
==========================================
  Files          76       76              
  Lines        4041     4055      +14     
==========================================
+ Hits         3143     3152       +9     
- Misses        620      622       +2     
- Partials      278      281       +3     
```

| Flag | Coverage Δ | |
|---|---|---|
| #go | `77.76% <84.00%> (+0.18%)` | :arrow_up: |
| #wasm | `72.52% <0.00%> (-0.26%)` | :arrow_down: |

| [Impacted Files](https://codecov.io/gh/pion/dtls/pull/257?src=pr&el=tree) | Coverage Δ | |
|---|---|---|
| [listener.go](https://codecov.io/gh/pion/dtls/pull/257/diff?src=pr&el=tree#diff-bGlzdGVuZXIuZ28=) | `45.71% <60.00%> (+3.40%)` | :arrow_up: |
| [internal/net/udp/conn.go](https://codecov.io/gh/pion/dtls/pull/257/diff?src=pr&el=tree#diff-aW50ZXJuYWwvbmV0L3VkcC9jb25uLmdv) | `91.72% <100.00%> (+2.66%)` | :arrow_up: |
| [crypto\_gcm.go](https://codecov.io/gh/pion/dtls/pull/257/diff?src=pr&el=tree#diff-Y3J5cHRvX2djbS5nbw==) | `65.38% <0.00%> (-3.85%)` | :arrow_down: |
| [conn.go](https://codecov.io/gh/pion/dtls/pull/257/diff?src=pr&el=tree#diff-Y29ubi5nbw==) | `83.99% <0.00%> (-0.72%)` | :arrow_down: |
| [internal/net/dpipe/dpipe.go](https://codecov.io/gh/pion/dtls/pull/257/diff?src=pr&el=tree#diff-aW50ZXJuYWwvbmV0L2RwaXBlL2RwaXBlLmdv) | `96.47% <0.00%> (+2.35%)` | :arrow_up: |

------

[Continue to review full report at Codecov](https://codecov.io/gh/pion/dtls/pull/257?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/257?src=pr&el=footer). Last update [52af2bf...acc5e5a](https://codecov.io/gh/pion/dtls/pull/257?src=pr&el=lastupdated). Read the [comment docs](https://docs.codecov.io/docs/pull-request-comments).


Patch:
diff --git a/internal/net/udp/conn.go b/internal/net/udp/conn.go
index 7e3ab40..9f680e8 100644
--- a/internal/net/udp/conn.go
+++ b/internal/net/udp/conn.go
@@ -23,10 +23,11 @@ var errListenQueueExceeded = errors.New("udp: listen queue exceeded")
 type listener struct {
 	pConn *net.UDPConn
 
-	accepting atomic.Value // bool
-	acceptCh  chan *Conn
-	doneCh    chan struct{}
-	doneOnce  sync.Once
+	accepting    atomic.Value // bool
+	acceptCh     chan *Conn
+	doneCh       chan struct{}
+	doneOnce     sync.Once
+	acceptFilter func([]byte) bool
 
 	connLock sync.Mutex
 	conns    map[string]*Conn
@@ -102,6 +103,10 @@ type ListenConfig struct {
 	// the request will be silently discarded, unlike TCP.
 	// Set zero to use default value 128 which is same as Linux default.
 	Backlog int
+
+	// AcceptFilter determines whether the new conn should be made for
+	// the incoming packet. If not set, any packet creates new conn.
+	AcceptFilter func([]byte) bool
 }
 
 // Listen creates a new listener based on the ListenConfig.
@@ -116,10 +121,11 @@ func (lc *ListenConfig) Listen(network string, laddr *net.UDPAddr) (net.Listener
 	}
 
 	l := &listener{
-		pConn:    conn,
-		acceptCh: make(chan *Conn, lc.Backlog),
-		conns:    make(map[string]*Conn),
-		doneCh:   make(chan struct{}),
+		pConn:        conn,
+		acceptCh:     make(chan *Conn, lc.Backlog),
+		conns:        make(map[string]*Conn),
+		doneCh:       make(chan struct{}),
+		acceptFilter: lc.AcceptFilter,
 	}
 	l.accepting.Store(true)
 	l.connWG.Add(1)
@@ -162,31 +168,38 @@ func (l *listener) readLoop() {
 		if err != nil {
 			return
 		}
-		conn, err := l.getConn(raddr)
+		conn, ok, err := l.getConn(raddr, buf[:n])
 		if err != nil {
 			continue
 		}
-		_, _ = conn.buffer.Write(buf[:n])
+		if ok {
+			_, _ = conn.buffer.Write(buf[:n])
+		}
 	}
 }
 
-func (l *listener) getConn(raddr net.Addr) (*Conn, error) {
+func (l *listener) getConn(raddr net.Addr, buf []byte) (*Conn, bool, error) {
 	l.connLock.Lock()
 	defer l.connLock.Unlock()
 	conn, ok := l.conns[raddr.String()]
 	if !ok {
 		if !l.accepting.Load().(bool) {
-			return nil, errClosedListener
+			return nil, false, errClosedListener
+		}
+		if l.acceptFilter != nil {
+			if !l.acceptFilter(buf) {
+				return nil, false, nil
+			}
 		}
 		conn = l.newConn(raddr)
 		select {
 		case l.acceptCh <- conn:
 			l.conns[raddr.String()] = conn
 		default:
-			return nil, errListenQueueExceeded
+			return nil, false, errListenQueueExceeded
 		}
 	}
-	return conn, nil
+	return conn, true, nil
 }
 
 // Conn augments a connection-oriented connection over a UDP PacketConn
diff --git a/listener.go b/listener.go
index 1fe3632..4c55de6 100644
--- a/listener.go
+++ b/listener.go
@@ -12,7 +12,20 @@ func Listen(network string, laddr *net.UDPAddr, config *Config) (net.Listener, e
 		return nil, err
 	}
 
-	parent, err := udp.Listen(network, laddr)
+	lc := udp.ListenConfig{
+		AcceptFilter: func(packet []byte) bool {
+			pkts, err := unpackDatagram(packet)
+			if err != nil || len(pkts) < 1 {
+				return false
+			}
+			h := &recordLayerHeader{}
+			if err := h.Unmarshal(pkts[0]); err != nil {
+				return false
+			}
+			return h.contentType == contentTypeHandshake
+		},
+	}
+	parent, err := lc.Listen(network, laddr)
 	if err != nil {
 		return nil, err
 	}


Test patch:
diff --git a/internal/net/udp/conn_test.go b/internal/net/udp/conn_test.go
index 5f9c8c8..718161e 100644
--- a/internal/net/udp/conn_test.go
+++ b/internal/net/udp/conn_test.go
@@ -145,6 +145,100 @@ func TestListenerCloseUnaccepted(t *testing.T) {
 	}
 }
 
+func TestListenerAcceptFilter(t *testing.T) {
+	// Limit runtime in case of deadlocks
+	lim := test.TimeOut(time.Second * 20)
+	defer lim.Stop()
+
+	// Check for leaking routines
+	report := test.CheckRoutines(t)
+	defer report()
+
+	testCases := map[string]struct {
+		packet []byte
+		accept bool
+	}{
+		"CreateConn": {
+			packet: []byte{0xAA},
+			accept: true,
+		},
+		"Discarded": {
+			packet: []byte{0x00},
+			accept: false,
+		},
+	}
+
+	for name, testCase := range testCases {
+		testCase := testCase
+		t.Run(name, func(t *testing.T) {
+			network, addr := getConfig()
+			listener, err := (&ListenConfig{
+				AcceptFilter: func(pkt []byte) bool {
+					return pkt[0] == 0xAA
+				},
+			}).Listen(network, addr)
+			if err != nil {
+				t.Fatal(err)
+			}
+
+			var wgAcceptLoop sync.WaitGroup
+			wgAcceptLoop.Add(1)
+			defer func() {
+				cerr := listener.Close()
+				if cerr != nil {
+					t.Fatal(cerr)
+				}
+				wgAcceptLoop.Wait()
+			}()
+
+			conn, derr := net.DialUDP(network, nil, listener.Addr().(*net.UDPAddr))
+			if derr != nil {
+				t.Fatal(derr)
+			}
+			if _, werr := conn.Write(testCase.packet); werr != nil {
+				t.Fatal(werr)
+			}
+			defer func() {
+				if cerr := conn.Close(); cerr != nil {
+					t.Error(cerr)
+				}
+			}()
+
+			chAccepted := make(chan struct{})
+			go func() {
+				defer wgAcceptLoop.Done()
+
+				conn, aerr := listener.Accept()
+				if aerr != nil {
+					if aerr != errClosedListener {
+						t.Error(aerr)
+					}
+					return
+				}
+				close(chAccepted)
+				if cerr := conn.Close(); cerr != nil {
+					t.Error(cerr)
+				}
+			}()
+
+			var accepted bool
+			select {
+			case <-chAccepted:
+				accepted = true
+			case <-time.After(10 * time.Millisecond):
+			}
+
+			if accepted != testCase.accept {
+				if testCase.accept {
+					t.Error("Packet should create new conn")
+				} else {
+					t.Error("Packet should not create new conn")
+				}
+			}
+		})
+	}
+}
+
 func TestListenerConcurrent(t *testing.T) {
 	// Limit runtime in case of deadlocks
 	lim := test.TimeOut(time.Second * 20)