Add timeout functionality when matching connections

An issue in the cockroachdb repo (cockroachdb/cockroach#143785)
uncovered an undesired behaviour in the goroutines handled by cmux.

Particularly, cmux can hang forever on trying to match a connection.

This can be resolved by users of cmux as illustrated in
cockroachdb/cockroach#144170. However, it is an
unorthodox solution.

This commit adds functionality to cmux to optionally specify a timeout on
matching (reading) connections.

Author: alyshanjahani-crl

cmux.go

@@ -20,6 +20,7 @@ import (
 	"io"
 	"net"
 	"sync"
+	"time"
 )
 
 // Matcher matches a connection based on its content.
@@ -32,12 +33,18 @@ type ErrorHandler func(error) bool
 var _ net.Error = ErrNotMatched{}
 
 // ErrNotMatched is returned whenever a connection is not matched by any of
-// the matchers registered in the multiplexer.
+// the matchers registered in the multiplexer. If the connection was not
+// matched due to a timeout, then that is indicated accordingly.
 type ErrNotMatched struct {
-	c net.Conn
+	c       net.Conn
+	timeout bool
 }
 
 func (e ErrNotMatched) Error() string {
+	if e.timeout {
+		return fmt.Sprintf("mux: connection %v not matched by an matcher (timeout)",
+			e.c.RemoteAddr())
+	}
 	return fmt.Sprintf("mux: connection %v not matched by an matcher",
 		e.c.RemoteAddr())
 }
@@ -46,7 +53,7 @@ func (e ErrNotMatched) Error() string {
 func (e ErrNotMatched) Temporary() bool { return true }
 
 // Timeout implements the net.Error interface.
-func (e ErrNotMatched) Timeout() bool { return false }
+func (e ErrNotMatched) Timeout() bool { return e.timeout }
 
 type errListenerClosed string
 
@@ -68,6 +75,18 @@ func New(l net.Listener) CMux {
 	}
 }
 
+// NewWithTimeout instantiates a new connection multiplexer that sets a timeout on matching
+// a connection.
+func NewWithTimeout(l net.Listener, timeout time.Duration) CMux {
+	return &cMux{
+		root:         l,
+		bufLen:       1024,
+		errh:         func(_ error) bool { return true },
+		donec:        make(chan struct{}),
+		matchTimeout: timeout,
+	}
+}
+
 // CMux is a multiplexer for network connections.
 type CMux interface {
 	// Match returns a net.Listener that sees (i.e., accepts) only
@@ -88,11 +107,12 @@ type matchersListener struct {
 }
 
 type cMux struct {
-	root   net.Listener
-	bufLen int
-	errh   ErrorHandler
-	donec  chan struct{}
-	sls    []matchersListener
+	root         net.Listener
+	bufLen       int
+	errh         ErrorHandler
+	donec        chan struct{}
+	sls          []matchersListener
+	matchTimeout time.Duration
 }
 
 func (m *cMux) Match(matchers ...Matcher) net.Listener {
@@ -137,11 +157,20 @@ func (m *cMux) Serve() error {
 func (m *cMux) serve(c net.Conn, donec <-chan struct{}, wg *sync.WaitGroup) {
 	defer wg.Done()
 
+	var timer *time.Timer
+	if m.matchTimeout > 0 {
+		_ = c.SetReadDeadline(time.Now().Add(m.matchTimeout))
+		timer = time.NewTimer(m.matchTimeout)
+		defer timer.Stop()
+	}
+
 	muc := newMuxConn(c)
 	for _, sl := range m.sls {
 		for _, s := range sl.ss {
 			matched := s(muc.getSniffer())
 			if matched {
+				// Reset (clear) any deadline once we've matched.
+				_ = c.SetReadDeadline(time.Time{})
 				select {
 				case sl.l.connc <- muc:
 				case <-donec:
@@ -152,8 +181,24 @@ func (m *cMux) serve(c net.Conn, donec <-chan struct{}, wg *sync.WaitGroup) {
 		}
 	}
 
+	// Check if the connection timed out.
+	// An alternative to using a timer is to read from the connection and see
+	// if a net.Timeout error is returned, however this is undesirable as it
+	// could lead to blocking on the read for matchTimeout duration.
+	// Another alternative is to plumb through the errors from the matchers reads
+	// which is a much bigger lift.
+	isTimeout := false
+	if timer != nil {
+		select {
+		case <-timer.C:
+			isTimeout = true
+		default:
+			// No timeout.
+		}
+	}
+
 	_ = c.Close()
-	err := ErrNotMatched{c: c}
+	err := ErrNotMatched{c: c, timeout: isTimeout}
 	if !m.handleErr(err) {
 		_ = m.root.Close()
 	}

cmux_test.go

@@ -418,6 +418,102 @@ func TestClose(t *testing.T) {
 	}
 }
 
+func readConnChan(conn net.Conn, buf []byte) <-chan struct {
+	n   int
+	err error
+} {
+	ch := make(chan struct {
+		n   int
+		err error
+	}, 1)
+	go func() {
+		n, err := conn.Read(buf)
+		ch <- struct {
+			n   int
+			err error
+		}{n, err}
+	}()
+	return ch
+}
+
+func TestServeWithTimeout(t *testing.T) {
+	defer leakCheck(t)()
+	errCh := make(chan error)
+	defer func() {
+		select {
+		case err := <-errCh:
+			t.Fatal(err)
+		default:
+		}
+	}()
+
+	t.Run("timeout expires", func(t *testing.T) {
+		l, cleanup := testListener(t)
+		defer cleanup()
+		mux := NewWithTimeout(l, 100*time.Millisecond)
+		httpl := mux.Match(HTTP1())
+		go runTestHTTPServer(errCh, httpl)
+		go safeServe(errCh, mux)
+
+		var handledTimeout bool
+		mux.HandleError(func(err error) bool {
+			if strings.Contains(err.Error(), "timeout") {
+				handledTimeout = true
+			}
+			return true
+		})
+
+		// Open a connection but send no data so it does not get matched.
+		_, err := net.Dial("tcp", l.Addr().String())
+		if err != nil {
+			t.Fatal(err)
+		}
+		// Wait for the timeout to occur
+		time.Sleep(200 * time.Millisecond)
+		if !handledTimeout {
+			t.Error("expected timeout error")
+		}
+	})
+
+	t.Run("timeout does not expire", func(t *testing.T) {
+		l, cleanup := testListener(t)
+		defer cleanup()
+		mux := NewWithTimeout(l, 100*time.Millisecond)
+		mux.Match(Any())
+		go safeServe(errCh, mux)
+
+		// Open a connection, it should match immediately.
+		matchedConn, err := net.Dial("tcp", l.Addr().String())
+		if err != nil {
+			t.Fatal(err)
+		}
+		defer matchedConn.Close()
+
+		// Timeout should not get triggered.
+		time.Sleep(200 * time.Millisecond)
+		// Create a channel that will send after 1 second.
+		timeoutCh := make(chan struct{})
+		go func() {
+			time.Sleep(1 * time.Second)
+			timeoutCh <- struct{}{}
+		}()
+		// The channel we created above should complete while reading from the connection should
+		// simply block indefinitley.
+		readBuf := make([]byte, 1)
+		select {
+		case <-timeoutCh:
+			break
+		case result := <-readConnChan(matchedConn, readBuf):
+			if result.err != nil {
+				t.Errorf("unexpected error reading from connection: %v", result.err)
+			} else {
+				t.Errorf("unexpected read of %d bytes from connection", result.n)
+			}
+		}
+	})
+
+}
+
 // Cribbed from google.golang.org/grpc/test/end2end_test.go.
 
 // interestingGoroutines returns all goroutines we care about for the purpose