From f9c9ad34f79058d4d963331480e0a91a4a8ea099 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Fri, 22 Sep 2023 08:50:19 +0000
Subject: [PATCH] feat(protonvpn): check udp vs tcp port forwarded

---
 internal/provider/protonvpn/portforward.go | 49 +++++++++++++++-------
 1 file changed, 34 insertions(+), 15 deletions(-)

diff --git a/internal/provider/protonvpn/portforward.go b/internal/provider/protonvpn/portforward.go
index 6b89d3e5..c83fc24a 100644
--- a/internal/provider/protonvpn/portforward.go
+++ b/internal/provider/protonvpn/portforward.go
@@ -6,6 +6,7 @@ import (
 	"fmt"
 	"net/http"
 	"net/netip"
+	"strings"
 	"time"
 
 	"github.com/qdm12/gluetun/internal/natpmp"
@@ -32,29 +33,47 @@ func (p *Provider) PortForward(ctx context.Context, _ *http.Client,
 	}
 
 	logger.Info("gateway external IPv4 address is " + externalIPv4Address.String())
-	networkProtocols := []string{"udp", "tcp"}
 	const internalPort, externalPort = 0, 0
 	const lifetime = 60 * time.Second
-	for _, networkProtocol := range networkProtocols {
-		_, _, assignedExternalPort, assignedLiftetime, err :=
-			client.AddPortMapping(ctx, gateway, networkProtocol,
-				internalPort, externalPort, lifetime)
-		if err != nil {
-			return 0, fmt.Errorf("adding port mapping: %w", err)
-		}
 
-		if assignedLiftetime != lifetime {
-			logger.Warn(fmt.Sprintf("assigned lifetime %s differs"+
-				" from requested lifetime %s",
-				assignedLiftetime, lifetime))
-		}
-
-		port = assignedExternalPort
+	_, _, assignedUDPExternalPort, assignedLifetime, err :=
+		client.AddPortMapping(ctx, gateway, "udp",
+			internalPort, externalPort, lifetime)
+	if err != nil {
+		return 0, fmt.Errorf("adding UDP port mapping: %w", err)
 	}
+	checkLifetime(logger, "UDP", lifetime, assignedLifetime)
+
+	_, _, assignedTCPExternalPort, assignedLifetime, err :=
+		client.AddPortMapping(ctx, gateway, "tcp",
+			internalPort, externalPort, lifetime)
+	if err != nil {
+		return 0, fmt.Errorf("adding TCP port mapping: %w", err)
+	}
+	checkLifetime(logger, "TCP", lifetime, assignedLifetime)
+
+	checkExternalPorts(logger, assignedUDPExternalPort, assignedTCPExternalPort)
+	port = assignedTCPExternalPort
 
 	return port, nil
 }
 
+func checkLifetime(logger utils.Logger, protocol string,
+	requested, actual time.Duration) {
+	if requested != actual {
+		logger.Warn(fmt.Sprintf("assigned %s port lifetime %s differs"+
+			" from requested lifetime %s", strings.ToUpper(protocol),
+			actual, requested))
+	}
+}
+
+func checkExternalPorts(logger utils.Logger, udpPort, tcpPort uint16) {
+	if udpPort != tcpPort {
+		logger.Warn(fmt.Sprintf("UDP external port %d differs from TCP external port %d",
+			udpPort, tcpPort))
+	}
+}
+
 func (p *Provider) KeepPortForward(ctx context.Context, port uint16,
 	gateway netip.Addr, _ string, logger utils.Logger) (err error) {
 	client := natpmp.New()