From cbaef624cfb463921109b6c2d52270bd5b3f58bf Mon Sep 17 00:00:00 2001
From: Audrius Butkevicius <audrius.butkevicius@gmail.com>
Date: Mon, 22 Jun 2020 09:01:57 +0100
Subject: [PATCH]  lib/nat: Make service termination faster (#6777)

* lib/nat: Make service termination faster

* Newline
---
 lib/pmp/pmp.go   | 17 +++++++++++++----
 lib/upnp/upnp.go | 35 +++++++++++++++++++++++------------
 2 files changed, 36 insertions(+), 16 deletions(-)

diff --git a/lib/pmp/pmp.go b/lib/pmp/pmp.go
index 56ef99b2b..7af5021d4 100644
--- a/lib/pmp/pmp.go
+++ b/lib/pmp/pmp.go
@@ -15,6 +15,7 @@ import (
 
 	"github.com/jackpal/gateway"
 	"github.com/jackpal/go-nat-pmp"
+	"github.com/pkg/errors"
 
 	"github.com/syncthing/syncthing/lib/nat"
 	"github.com/syncthing/syncthing/lib/util"
@@ -44,10 +45,18 @@ func Discover(ctx context.Context, renewal, timeout time.Duration) []nat.Device
 	c := natpmp.NewClientWithTimeout(ip, timeout)
 	// Try contacting the gateway, if it does not respond, assume it does not
 	// speak NAT-PMP.
-	_, err = c.GetExternalAddress()
-	if err != nil && strings.Contains(err.Error(), "Timed out") {
-		l.Debugln("Timeout trying to get external address, assume no NAT-PMP available")
-		return nil
+	err = util.CallWithContext(ctx, func() error {
+		_, ierr := c.GetExternalAddress()
+		return ierr
+	})
+	if err != nil {
+		if errors.Cause(err) == context.Canceled {
+			return nil
+		}
+		if strings.Contains(err.Error(), "Timed out") {
+			l.Debugln("Timeout trying to get external address, assume no NAT-PMP available")
+			return nil
+		}
 	}
 
 	var localIP net.IP
diff --git a/lib/upnp/upnp.go b/lib/upnp/upnp.go
index a2a9b9d33..b18be59ac 100644
--- a/lib/upnp/upnp.go
+++ b/lib/upnp/upnp.go
@@ -119,19 +119,26 @@ func Discover(ctx context.Context, renewal, timeout time.Duration) []nat.Device
 	}()
 
 	seenResults := make(map[string]bool)
-	for result := range resultChan {
-		if seenResults[result.ID()] {
-			l.Debugf("Skipping duplicate result %s", result.ID())
-			continue
+
+	for {
+		select {
+		case result, ok := <-resultChan:
+			if !ok {
+				return results
+			}
+			if seenResults[result.ID()] {
+				l.Debugf("Skipping duplicate result %s", result.ID())
+				continue
+			}
+
+			results = append(results, result)
+			seenResults[result.ID()] = true
+
+			l.Debugf("UPnP discovery result %s", result.ID())
+		case <-ctx.Done():
+			return nil
 		}
-
-		results = append(results, result)
-		seenResults[result.ID()] = true
-
-		l.Debugf("UPnP discovery result %s", result.ID())
 	}
-
-	return results
 }
 
 // Search for UPnP InternetGatewayDevices for <timeout> seconds.
@@ -213,7 +220,11 @@ loop:
 		}
 		for _, igd := range igds {
 			igd := igd // Copy before sending pointer to the channel.
-			results <- &igd
+			select {
+			case results <- &igd:
+			case <-ctx.Done():
+				return
+			}
 		}
 	}
 	l.Debugln("Discovery for device type", deviceType, "on", intf.Name, "finished.")