Repository: pion/dtls
Pull request #113 (2019-08-28 19:57:49)

Fix connect timeout doesn't stop outgoing handshake routine
### Description
Recently a ConnectTimeout option was added which waits for the handshake to complete in the given timeframe or otherwise moves on. However, the outbound handshake goroutine was never cancelled when a timeout occurs. This PR refactors the previous signaling mechanism and adds a call to `c.handshakeDoneSignal.Close()` at timeout, so that the outgoing handshake goroutine receives the signal and stops.

#### Reference issue
Fixes regression from #94

RyanGordon: Makes sense :+1: I will add a test tomorrow
RyanGordon: Okay, I added a regression test (~400ms after optimizing the flight interval and timeout) that fails without the regression fix

I also switched the `handshakeCompletedSignal` from a channel to a `context.WithCancel()` which tends to be safer than calling close on a channel to signal a goroutine shutdown
Sean-Der: Great work @RyanGordon 

We should move everything from `close` -> `Context`. Doing `close` is a very common pattern in the code bases, so I think it would be a great cleanup across the board. If you feel strongly about making this change I say we commit to doing it everywhere, just want to avoid the lava pattern where some code bases are more modern then others!

```
seaduboi@38f9d359441f:~/go/src/github.com/pion$ rg -l --fixed-strings 'close(' | grep -v dart | grep -v .js
stun/client.go
stun/client_test.go
mdns/conn.go
ice/gather.go
ice/agent.go
ice/agent_test.go
ice/candidate_base.go
ice/transport_test.go
ice/candidate_relay.go
ice/connectivity_vnet_test.go
ice/candidate.go
turn/server.go
turn/client_test.go
dtls/resume_test.go
dtls/conn.go
sctp/ack_timer_test.go
sctp/association_test.go
sctp/rtx_timer.go
sctp/rtx_timer_test.go
sctp/ack_timer.go
srtp/stream_srtp_test.go
srtp/session_srtp.go
srtp/session.go
sctp/association.go
srtp/session_srtcp.go
transport/test/bridge.go
transport/test/stress.go
quic/transport_test.go
transport/packetio/buffer_test.go
transport/packetio/buffer.go
webrtc/peerconnection_close_test.go
webrtc/sctptransport.go
webrtc/peerconnection_media_test.go
webrtc/quictransport_test.go
webrtc/datachannel.go
webrtc/icegathererstate.go
webrtc/datachannel_go_test.go
transport/vnet/router_test.go
transport/vnet/net_test.go
transport/vnet/stress_test.go
transport/vnet/conn_test.go
transport/vnet/conn.go
transport/vnet/router.go
webrtc/peerconnection.go
webrtc/rtpsender.go
webrtc/datachannel_ortc_test.go
webrtc/icetransport_test.go
webrtc/dtlstransportstate.go
webrtc/rtpreceiver.go
webrtc/peerconnection_go_test.go
dtls/e2e/e2e_lossy_test.go
stun/cmd/stun-traversal/main.go
turn/internal/client/transaction.go
turn/internal/client/trylock_test.go
turn/internal/client/periodic_timer.go
turn/internal/client/conn.go
turn/internal/allocation/allocation.go
dtls/internal/udp/conn.go
webrtc/internal/mux/mux.go
webrtc/internal/mux/endpoint.go
```
RyanGordon: Personally I find it a bit dangerous to use a chan for shutdown because its not idempotent and the failure case is a panic. I.e. `close()` is the only way to guarantee precedence if multiple channels are being listened to in a `select` and if `close()` ends up being called twice by the application for whatever reason that leads to the panic.

I do agree chan is a common pattern to use for signaling closing, I know there are a few opinions on the best approach out there but generally, even if it looks a little worse in implementation (handling `ctx` and `func()` from `context.WithCancel(...)` in practice its much safer than closing channels.

Happy to have more discussion on this, but if we have consensus on the approach, I can open up PRs for the additional places you've found. We could even create a wrapper struct to abstract it, like this:

```
type Closer struct {
   ctx context.Context
   shutdownFunc func()
}

NewCloser() *Closer {
   ctx, shutdownFunc := context.WithCancel(context.Background())
   return &Closer{
      ctx: ctx,
      shutdownFunc: shutdownFunc,
   }
}

func (c *Closer) Closed() <-chan struct{}  {
   return c.ctx.Done()
}

func (c *Closer) Close() {
   c.shutdownFunc()
}
```

Then we can use the abstraction like this in the various libraries:

```
...
a.handshakeCloser := NewCloser()
defer a.handshakeCloser.Close()
go a.someFunc()
...
``` 

```
func (a ...) someFunc() {
  for { 
     select {
     case <-a.handshakeCloser.Closed():
        return
     case ...
     }
  }
}
```
RyanGordon: Okay, I updated this, mind reviewing again when you get a chance?
Sean-Der: @RyanGordon oh yea that is great, merge away!

I think the old `close` method is a 'clever hack' I don't think if I saw that for the first time I would understand the intent. I will be copying this into `pion/webrtc` when I get the chance for sure :)

Patch:
diff --git a/client_handlers.go b/client_handlers.go
index 7311714..a70b160 100644
--- a/client_handlers.go
+++ b/client_handlers.go
@@ -235,7 +235,7 @@ func clientHandshakeHandler(c *Conn) (*alert, error) {
 
 		c.setLocalEpoch(1)
 		c.state.localSequenceNumber = 1
-		c.signalHandshakeComplete()
+		c.handshakeDoneSignal.Close()
 	default:
 		return &alert{alertLevelFatal, alertUnexpectedMessage}, fmt.Errorf("client asked to handle unknown flight (%d)", c.currFlight.get())
 	}
diff --git a/closer.go b/closer.go
new file mode 100644
index 0000000..8af9278
--- /dev/null
+++ b/closer.go
@@ -0,0 +1,39 @@
+package dtls
+
+import (
+	"context"
+)
+
+// Closer allows for each signaling a channel for shutdown
+type Closer struct {
+	ctx       context.Context
+	closeFunc func()
+}
+
+// NewCloser creates a new instance of Closer
+func NewCloser() *Closer {
+	ctx, closeFunc := context.WithCancel(context.Background())
+	return &Closer{
+		ctx:       ctx,
+		closeFunc: closeFunc,
+	}
+}
+
+// NewCloserWithParent creates a new instance of Closer with a parent context
+func NewCloserWithParent(ctx context.Context) *Closer {
+	ctx, closeFunc := context.WithCancel(ctx)
+	return &Closer{
+		ctx:       ctx,
+		closeFunc: closeFunc,
+	}
+}
+
+// Done returns a channel signaling when it is done
+func (c *Closer) Done() <-chan struct{} {
+	return c.ctx.Done()
+}
+
+// Close sends a signal to trigger the ctx done channel
+func (c *Closer) Close() {
+	c.closeFunc()
+}
diff --git a/conn.go b/conn.go
index 743f2eb..d6a2617 100644
--- a/conn.go
+++ b/conn.go
@@ -75,7 +75,7 @@ type Conn struct {
 
 	handshakeMessageHandler        handshakeMessageHandler
 	flightHandler                  flightHandler
-	handshakeCompletedSignal       chan bool
+	handshakeDoneSignal            *Closer
 	handshakeCompletedSuccessfully atomic.Value
 
 	connErr atomic.Value
@@ -126,6 +126,8 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 		connectTimeout = math.MaxInt64 * time.Nanosecond
 	}
 
+	handshakeDoneSignal := NewCloser()
+
 	c := &Conn{
 		nextConn:                    nextConn,
 		currFlight:                  newFlight(isClient, logger),
@@ -149,10 +151,10 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 		localPSKCallback:     config.PSK,
 		localPSKIdentityHint: config.PSKIdentityHint,
 
-		decrypted:                make(chan []byte),
-		workerTicker:             time.NewTicker(workerInterval),
-		handshakeCompletedSignal: make(chan bool),
-		log:                      logger,
+		decrypted:           make(chan []byte),
+		workerTicker:        time.NewTicker(workerInterval),
+		handshakeDoneSignal: handshakeDoneSignal,
+		log:                 logger,
 	}
 
 	// Use host from conn address when serverName is not provided
@@ -188,10 +190,11 @@ func createConn(nextConn net.Conn, flightHandler flightHandler, handshakeMessage
 	go c.inboundLoop()
 
 	select {
-	case <-c.handshakeCompletedSignal:
+	case <-c.handshakeDoneSignal.Done():
 		err = c.getConnErr()
 	case <-time.After(c.connectTimeout):
 		err = errConnectTimeout
+		c.handshakeDoneSignal.Close()
 	}
 
 	if err == nil {
@@ -491,14 +494,6 @@ func (c *Conn) notify(level alertLevel, desc alertDescription) {
 	c.state.localSequenceNumber++
 }
 
-func (c *Conn) signalHandshakeComplete() {
-	select {
-	case <-c.handshakeCompletedSignal:
-	default:
-		close(c.handshakeCompletedSignal)
-	}
-}
-
 func (c *Conn) setHandshakeCompletedSuccessfully() {
 	c.handshakeCompletedSuccessfully.Store(struct{ bool }{true})
 }
@@ -517,7 +512,7 @@ func (c *Conn) startHandshakeOutbound() {
 				err        error
 			)
 			select {
-			case <-c.handshakeCompletedSignal:
+			case <-c.handshakeDoneSignal.Done():
 				return
 			case <-c.workerTicker.C:
 				isFinished, alertPtr, err = c.flightHandler(c)
@@ -555,7 +550,7 @@ func (c *Conn) stopWithError(err error) {
 
 	c.workerTicker.Stop()
 
-	c.signalHandshakeComplete()
+	c.handshakeDoneSignal.Close()
 }
 
 func (c *Conn) getConnErr() error {
diff --git a/resume.go b/resume.go
index 07c6d46..7e55349 100644
--- a/resume.go
+++ b/resume.go
@@ -18,7 +18,7 @@ func Resume(state *State, conn net.Conn, config *Config) (*Conn, error) {
 	// Custom flight handler that sets imported data and signals as handshaked
 	flightHandler := func(c *Conn) (bool, *alert, error) {
 		c.state = *state
-		c.signalHandshakeComplete()
+		c.handshakeDoneSignal.Close()
 		return true, nil, nil
 	}
 
diff --git a/server_handlers.go b/server_handlers.go
index 3a53925..8d0617a 100644
--- a/server_handlers.go
+++ b/server_handlers.go
@@ -502,8 +502,7 @@ func serverFlightHandler(c *Conn) (bool, *alert, error) {
 				}},
 		}, true)
 
-		// TODO: Better way to end handshake
-		c.signalHandshakeComplete()
+		c.handshakeDoneSignal.Close()
 		return true, nil, nil
 	default:
 		return false, &alert{alertLevelFatal, alertUnexpectedMessage}, fmt.Errorf("unhandled flight %s", c.currFlight.get())


Test patch:
diff --git a/conn_test.go b/conn_test.go
index 53f05da..5f9795f 100644
--- a/conn_test.go
+++ b/conn_test.go
@@ -2,6 +2,7 @@ package dtls
 
 import (
 	"bytes"
+	"crypto/rand"
 	"crypto/x509"
 	"errors"
 	"fmt"
@@ -933,3 +934,123 @@ func TestPSKConfiguration(t *testing.T) {
 		}
 	}
 }
+
+func TestServerTimeout(t *testing.T) {
+	// Limit runtime in case of deadlocks
+	lim := test.TimeOut(time.Second * 20)
+	defer lim.Stop()
+
+	cookie := make([]byte, 20)
+	_, err := rand.Read(cookie)
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	var rand [28]byte
+	random := handshakeRandom{time.Unix(500, 0), rand}
+
+	cipherSuites := []cipherSuite{
+		&cipherSuiteTLSEcdheEcdsaWithAes128GcmSha256{},
+		&cipherSuiteTLSEcdheRsaWithAes128GcmSha256{},
+	}
+
+	extensions := []extension{
+		&extensionSupportedSignatureAlgorithms{
+			signatureHashAlgorithms: []signatureHashAlgorithm{
+				{HashAlgorithmSHA256, signatureAlgorithmECDSA},
+				{HashAlgorithmSHA384, signatureAlgorithmECDSA},
+				{HashAlgorithmSHA512, signatureAlgorithmECDSA},
+				{HashAlgorithmSHA256, signatureAlgorithmRSA},
+				{HashAlgorithmSHA384, signatureAlgorithmRSA},
+				{HashAlgorithmSHA512, signatureAlgorithmRSA},
+			},
+		},
+		&extensionSupportedEllipticCurves{
+			ellipticCurves: []namedCurve{namedCurveX25519, namedCurveP256, namedCurveP384},
+		},
+		&extensionSupportedPointFormats{
+			pointFormats: []ellipticCurvePointFormat{ellipticCurvePointFormatUncompressed},
+		},
+	}
+
+	record := &recordLayer{
+		recordLayerHeader: recordLayerHeader{
+			sequenceNumber:  0,
+			protocolVersion: protocolVersion1_2,
+		},
+		content: &handshake{
+			// sequenceNumber and messageSequence line up, may need to be re-evaluated
+			handshakeHeader: handshakeHeader{
+				messageSequence: 0,
+			},
+			handshakeMessage: &handshakeMessageClientHello{
+				version:            protocolVersion1_2,
+				cookie:             cookie,
+				random:             random,
+				cipherSuites:       cipherSuites,
+				compressionMethods: defaultCompressionMethods,
+				extensions:         extensions,
+			}},
+	}
+
+	packet, err := record.Marshal()
+	if err != nil {
+		t.Fatal(err)
+	}
+
+	ca, cb := net.Pipe()
+	defer func() {
+		err := ca.Close()
+		if err != nil {
+			t.Fatal(err)
+		}
+	}()
+
+	// Client reader
+	caReadChan := make(chan []byte, 1000)
+	go func() {
+		for {
+			data := make([]byte, 8192)
+			n, err := ca.Read(data)
+			if err != nil {
+				return
+			}
+
+			caReadChan <- data[:n]
+		}
+	}()
+
+	// Start sending ClientHello packets until server responds with first packet
+	go func() {
+		for {
+			select {
+			case <-time.After(10 * time.Millisecond):
+				_, err := ca.Write(packet)
+				if err != nil {
+					return
+				}
+			case <-caReadChan:
+				// Once we receive the first reply from the server, stop
+				return
+			}
+		}
+	}()
+
+	config := &Config{
+		CipherSuites:   []CipherSuiteID{TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256},
+		ConnectTimeout: ConnectTimeoutOption(50 * time.Millisecond),
+		FlightInterval: 100 * time.Millisecond,
+	}
+
+	if _, err := testServer(cb, config, true); err != errConnectTimeout {
+		t.Fatalf("TestServerTimeout: Client error exp(%v) failed(%v)", errConnectTimeout, err)
+	}
+
+	// Wait a little longer to ensure no additional messages have been sent by the server
+	time.Sleep(300 * time.Millisecond)
+	select {
+	case msg := <-caReadChan:
+		t.Fatalf("TestServerTimeout: Expected no additional messages from server, got: %+v", msg)
+	default:
+	}
+}