Fix goroutine leak, deadlock, and DTLS deadline issue

Makefile

@@ -24,6 +24,10 @@ app:
 	[ -d $(EXE_DIR) ] || mkdir -p $(EXE_DIR)
 	go build -o ${EXE_DIR}/application ./cmd/application
 
+app-dbg:
+	[ -d $(EXE_DIR) ] || mkdir -p $(EXE_DIR)
+	go build -tags debug -o ${EXE_DIR}/application ./cmd/application
+
 libtd:
 	cd ./libtapdance/ && make libtapdance.a
 

cmd/application/main.go

@@ -36,6 +36,8 @@ func main() {
 	flag.StringVar(&zmqAddress, "zmq-address", "ipc://@zmq-proxy", "Address of ZMQ proxy")
 	flag.Parse()
 
+	startPProf()
+
 	// Init stats
 	cj.Stat()
 

cmd/application/pprof.go

@@ -0,0 +1,6 @@
+//go:build !debug
+
+package main
+
+func startPProf() {
+}

cmd/application/pprof_dbg.go

@@ -0,0 +1,15 @@
+//go:build debug
+
+package main
+
+import (
+	"log"
+	"net/http"
+	_ "net/http/pprof"
+)
+
+func startPProf() {
+	go func() {
+		log.Println(http.ListenAndServe("localhost:6060", nil))
+	}()
+}

pkg/dtls/dial.go

@@ -6,6 +6,7 @@ import (
 	"crypto/x509"
 	"fmt"
 	"net"
+	"time"
 
 	"github.com/pion/dtls/v2"
 	"github.com/pion/dtls/v2/pkg/protocol/handshake"
@@ -33,6 +34,40 @@ func Client(conn net.Conn, config *Config) (net.Conn, error) {
 
 // DialWithContext creates a DTLS connection to the given network address using the given shared secret
 func ClientWithContext(ctx context.Context, conn net.Conn, config *Config) (net.Conn, error) {
+
+	dtlsConn, err := dtlsCtx(ctx, conn, config)
+	if err != nil {
+		return nil, fmt.Errorf("error creating dtls connection: %w", err)
+	}
+
+	ddl, ok := ctx.Deadline()
+	if ok {
+		err := conn.SetDeadline(ddl)
+		if err != nil {
+			return nil, fmt.Errorf("error setting deadline: %v", err)
+		}
+	}
+
+	wrappedConn, err := wrapSCTP(dtlsConn, config)
+	if err != nil {
+		dtlsConn.Close()
+		return nil, err
+	}
+
+	err = conn.SetDeadline(time.Time{})
+	if err != nil {
+		return nil, fmt.Errorf("error setting deadline: %v", err)
+	}
+
+	err = wrappedConn.SetDeadline(time.Time{})
+	if err != nil {
+		return nil, fmt.Errorf("error setting deadline: %v", err)
+	}
+
+	return wrappedConn, nil
+}
+
+func dtlsCtx(ctx context.Context, conn net.Conn, config *Config) (net.Conn, error) {
 	clientCert, serverCert, err := certsFromSeed(config.PSK)
 
 	if err != nil {
@@ -68,16 +103,6 @@ func ClientWithContext(ctx context.Context, conn net.Conn, config *Config) (net.
 		VerifyPeerCertificate: verifyServerCertificate,
 	}
 
-	dtlsConn, err := dtls.ClientWithContext(ctx, conn, dtlsConf)
-
-	if err != nil {
-		return nil, fmt.Errorf("error creating dtls connection: %w", err)
-	}
-
-	wrappedConn, err := wrapSCTP(dtlsConn, config)
-	if err != nil {
-		return nil, err
-	}
+	return dtls.ClientWithContext(ctx, conn, dtlsConf)
 
-	return wrappedConn, nil
 }

pkg/dtls/goroutine_leak_test.go

@@ -19,7 +19,7 @@ func passGoroutineLeak(testFunc func(*testing.T), t *testing.T) bool {
 }
 
 func TestGoroutineLeak(t *testing.T) {
-	testFuncs := []func(*testing.T){TestSend, TestServerFail, TestClientFail, TestListenSuccess, TestListenFail}
+	testFuncs := []func(*testing.T){TestSend, TestServerFail, TestClientFail, TestListenSuccess, TestListenFail, TestFailSCTP}
 
 	for _, test := range testFuncs {
 		require.True(t, passGoroutineLeak(test, t))

pkg/dtls/listener.go

@@ -149,23 +149,36 @@ func (l *Listener) Addr() net.Addr {
 
 func (l *Listener) verifyConnection(state *dtls.State) error {
 
-	certs, ok := l.connToCert[state.RemoteRandomBytes()]
-	if !ok {
-		return fmt.Errorf("no matching certificate found with client hello random")
+	certs, err := l.getCert(state.RemoteRandomBytes())
+	if err != nil {
+		return err
 	}
 
 	if len(state.PeerCertificates) != 1 {
 		return fmt.Errorf("expected 1 peer certificate, got %v", len(state.PeerCertificates))
 	}
 
-	err := verifyCert(state.PeerCertificates[0], certs.clientCert.Certificate[0])
+	err = verifyCert(state.PeerCertificates[0], certs.clientCert.Certificate[0])
 	if err != nil {
 		return fmt.Errorf("error verifying peer certificate: %v", err)
 	}
 
 	return nil
 }
 
+func (l *Listener) getCert(id [handshake.RandomBytesLength]byte) (*certPair, error) {
+	l.connToCertMutex.Lock()
+	defer l.connToCertMutex.Unlock()
+
+	certs, ok := l.connToCert[id]
+	if !ok {
+		return nil, fmt.Errorf("no matching certificate found with client hello random")
+	}
+
+	return certs, nil
+
+}
+
 // Accept accepts a connection with shared secret
 func (l *Listener) Accept(config *Config) (net.Conn, error) {
 	// Call the new function with a background context
@@ -174,6 +187,37 @@ func (l *Listener) Accept(config *Config) (net.Conn, error) {
 
 // AcceptWithContext accepts a connection with shared secret, with a context
 func (l *Listener) AcceptWithContext(ctx context.Context, config *Config) (net.Conn, error) {
+
+	conn, err := l.acceptDTLSConn(ctx, config)
+	if err != nil {
+		return nil, err
+	}
+
+	ddl, ok := ctx.Deadline()
+	if ok {
+		err := conn.SetDeadline(ddl)
+		if err != nil {
+			return nil, fmt.Errorf("error setting deadline: %v", err)
+		}
+	}
+	wrappedConn, err := wrapSCTP(conn, config)
+	if err != nil {
+		conn.Close()
+		return nil, err
+	}
+	err = conn.SetDeadline(time.Time{})
+	if err != nil {
+		return nil, fmt.Errorf("error setting deadline: %v", err)
+	}
+
+	err = wrappedConn.SetDeadline(time.Time{})
+	if err != nil {
+		return nil, fmt.Errorf("error setting deadline: %v", err)
+	}
+	return wrappedConn, nil
+}
+
+func (l *Listener) acceptDTLSConn(ctx context.Context, config *Config) (net.Conn, error) {
 	clientCert, serverCert, err := certsFromSeed(config.PSK)
 	if err != nil {
 		return &dtls.Conn{}, fmt.Errorf("error generating certificatess from seed: %v", err)
@@ -198,11 +242,7 @@ func (l *Listener) AcceptWithContext(ctx context.Context, config *Config) (net.C
 
 	select {
 	case conn := <-connCh:
-		wrappedConn, err := wrapSCTP(conn, config)
-		if err != nil {
-			return nil, err
-		}
-		return wrappedConn, nil
+		return conn, nil
 	case <-ctx.Done():
 		return nil, ctx.Err()
 	}

pkg/dtls/sctpconn.go

@@ -86,11 +86,9 @@ func (s *SCTPConn) Close() error {
 
 	s.closeOnce.Do(func() { close(s.closed) })
 
-	err := s.stream.Close()
-	if err != nil {
-		return err
-	}
-	return s.conn.Close()
+	s.stream.Close()
+	s.conn.Close()
+	return nil
 }
 
 func (s *SCTPConn) Write(b []byte) (int, error) {

pkg/dtls/server.go

@@ -6,6 +6,7 @@ import (
 	"crypto/x509"
 	"fmt"
 	"net"
+	"time"
 
 	"github.com/pion/dtls/v2"
 )
@@ -51,11 +52,30 @@ func ServerWithContext(ctx context.Context, conn net.Conn, config *Config) (net.
 		return nil, err
 	}
 
+	ddl, ok := ctx.Deadline()
+	if ok {
+		err := conn.SetDeadline(ddl)
+		if err != nil {
+			return nil, fmt.Errorf("error setting deadline: %v", err)
+		}
+	}
+
 	wrappedConn, err := wrapSCTP(dtlsConn, config)
 	if err != nil {
+		dtlsConn.Close()
 		return nil, err
 	}
 
+	err = conn.SetDeadline(time.Time{})
+	if err != nil {
+		return nil, fmt.Errorf("error setting deadline: %v", err)
+	}
+
+	err = wrappedConn.SetDeadline(time.Time{})
+	if err != nil {
+		return nil, fmt.Errorf("error setting deadline: %v", err)
+	}
+
 	return wrappedConn, nil
 }
 

pkg/dtls/server_test.go

@@ -85,3 +85,27 @@ func TestClientFail(t *testing.T) {
 		t.Fatalf("Connect does not respect context")
 	}
 }
+
+func TestFailSCTP(t *testing.T) {
+
+	ctxTime := 3 * time.Second
+	ctx, cancel := context.WithTimeout(context.Background(), ctxTime)
+	defer cancel()
+
+	server, client := net.Pipe()
+
+	before := time.Now()
+
+	go func() {
+		_, _ = dtlsCtx(ctx, client, &Config{PSK: sharedSecret, SCTP: ClientOpen})
+	}()
+
+	_, err := ServerWithContext(ctx, server, &Config{PSK: sharedSecret, SCTP: ServerAccept})
+
+	require.NotNil(t, err)
+	dur := time.Since(before)
+	if dur > ctxTime*2 {
+		t.Fatalf("Connect does not respect context")
+	}
+
+}

pkg/station/lib/registration_config_test.go

@@ -27,7 +27,7 @@ func TestConjureLibConfigResolveBlocklisted(t *testing.T) {
 	goodTestCases := map[string][]string{
 		"128.0.2.1:25":     {"128.0.2.1:25"},
 		"[2001:db8::1]:80": {"[2001:db8::1]:80"},
-		"example.com:1234": {"93.184.216.34:1234", "[2606:2800:220:1:248:1893:25c8:1946]:1234"},
+		"example.com:1234": {"93.184.215.14:1234", "[2606:2800:21f:cb07:6820:80da:af6b:8b2c]:1234"},
 		"[::2]:443":        {"[::2]:443"},
 	}
 

pkg/station/lib/registration_ingest.go

@@ -540,6 +540,8 @@ func handleConnectingTpReg(regManager *RegistrationManager, reg *DecoyRegistrati
 					return
 				}
 
+				defer conn.Close()
+
 				// regManager.connectingStats.AddCreatedToSuccessfulConnecting(asn, cc, transport.Name())
 
 				Stat().AddConn()

pkg/transports/connecting/dtls/dtls.go

@@ -77,8 +77,11 @@ func (t *Transport) Connect(ctx context.Context, reg transports.Registration) (n
 		return nil, fmt.Errorf("transport params is not *pb.DTLSTransportParams")
 	}
 
-	connCh := make(chan net.Conn, 2)
-	errCh := make(chan error, 2)
+	connCh := make(chan net.Conn)
+	errCh := make(chan error)
+
+	ctxCancel, cancel := context.WithCancel(ctx)
+	defer cancel()
 
 	go func() {
 
@@ -94,7 +97,10 @@ func (t *Transport) Connect(ctx context.Context, reg transports.Registration) (n
 
 		err := t.DNAT.AddEntry(&clientAddr.IP, uint16(clientAddr.Port), reg.PhantomIP(), reg.GetDstPort())
 		if err != nil {
-			errCh <- fmt.Errorf("error adding DNAT entry: %v", err)
+			select {
+			case errCh <- fmt.Errorf("error adding DNAT entry: %v", err):
+			case <-ctxCancel.Done():
+			}
 			return
 		}
 
@@ -108,30 +114,47 @@ func (t *Transport) Connect(ctx context.Context, reg transports.Registration) (n
 
 		udpConn, err := reuseport.Dial("udp", laddr.String(), clientAddr.String())
 		if err != nil {
-			errCh <- fmt.Errorf("error connecting to dtls client: %v", err)
+			select {
+			case errCh <- fmt.Errorf("error connecting to dtls client: %v", err):
+			case <-ctxCancel.Done():
+			}
 			return
 		}
 
-		dtlsConn, err := dtls.ClientWithContext(ctx, udpConn, &dtls.Config{PSK: reg.SharedSecret(), SCTP: dtls.ServerAccept, Unordered: params.GetUnordered()})
+		dtlsConn, err := dtls.ClientWithContext(ctxCancel, udpConn, &dtls.Config{PSK: reg.SharedSecret(), SCTP: dtls.ServerAccept, Unordered: params.GetUnordered()})
 		if err != nil {
-			errCh <- fmt.Errorf("error connecting to dtls client: %v", err)
+			select {
+			case errCh <- fmt.Errorf("error connecting to dtls client: %v", err):
+			case <-ctxCancel.Done():
+			}
 			return
 		}
-		t.logDialSuccess(&clientAddr.IP)
 
-		connCh <- dtlsConn
+		select {
+		case connCh <- dtlsConn:
+			t.logDialSuccess(&clientAddr.IP)
+		case <-ctxCancel.Done():
+			dtlsConn.Close()
+		}
 	}()
 
 	go func() {
-		conn, err := t.dtlsListener.AcceptWithContext(ctx, &dtls.Config{PSK: reg.SharedSecret(), SCTP: dtls.ServerAccept, Unordered: params.GetUnordered()})
+		conn, err := t.dtlsListener.AcceptWithContext(ctxCancel, &dtls.Config{PSK: reg.SharedSecret(), SCTP: dtls.ServerAccept, Unordered: params.GetUnordered()})
 		if err != nil {
-			errCh <- fmt.Errorf("error accepting dtls connection from secret: %v", err)
+			select {
+			case errCh <- fmt.Errorf("error accepting dtls connection from secret: %v", err):
+			case <-ctxCancel.Done():
+			}
 			return
 		}
-		logip := net.ParseIP(reg.GetRegistrationAddress())
-		t.logListenSuccess(&logip)
 
-		connCh <- conn
+		select {
+		case connCh <- conn:
+			logip := net.ParseIP(reg.GetRegistrationAddress())
+			t.logListenSuccess(&logip)
+		case <-ctxCancel.Done():
+			conn.Close()
+		}
 	}()
 
 	var errs []error