From 3c7dc9b9adc9f8cce1ab6de71da2c1f47a57e673 Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Sun, 31 Jan 2021 01:27:13 +0000
Subject: [PATCH] Feature: Private Internet Access custom port

---
 internal/models/selection.go   |  3 +-
 internal/params/params.go      |  1 +
 internal/params/pia.go         |  7 ++++
 internal/provider/piav4.go     | 69 +++++++++++++++++++++++++---------
 internal/settings/providers.go |  4 ++
 5 files changed, 66 insertions(+), 18 deletions(-)

diff --git a/internal/models/selection.go b/internal/models/selection.go
index ff7eabae8..9b887f27b 100644
--- a/internal/models/selection.go
+++ b/internal/models/selection.go
@@ -33,7 +33,7 @@ type ServerSelection struct {
 	ISPs  []string `json:"isps"`
 	Owned bool     `json:"owned"`
 
-	// Mullvad, Windscribe
+	// Mullvad, Windscribe, PIA
 	CustomPort uint16 `json:"custom_port"`
 
 	// NordVPN
@@ -92,6 +92,7 @@ func (p *ProviderSettings) String() string {
 			"Regions: "+commaJoin(p.ServerSelection.Regions),
 			"Encryption preset: "+p.ExtraConfigOptions.EncryptionPreset,
 			"Port forwarding: "+p.PortForwarding.String(),
+			"Custom port: "+customPort,
 		)
 	case "mullvad":
 		settingsList = append(settingsList,
diff --git a/internal/params/params.go b/internal/params/params.go
index 475f20b56..8a2b93d9c 100644
--- a/internal/params/params.go
+++ b/internal/params/params.go
@@ -62,6 +62,7 @@ type Reader interface {
 	GetPortForwardingStatusFilepath() (filepath models.Filepath, err error)
 	GetPIAEncryptionPreset() (preset string, err error)
 	GetPIARegions() (regions []string, err error)
+	GetPIAPort() (port uint16, err error)
 
 	// Mullvad getters
 	GetMullvadCountries() (countries []string, err error)
diff --git a/internal/params/pia.go b/internal/params/pia.go
index 330e10128..239ad9da8 100644
--- a/internal/params/pia.go
+++ b/internal/params/pia.go
@@ -63,3 +63,10 @@ func (r *reader) GetPIAEncryptionPreset() (preset string, err error) {
 func (r *reader) GetPIARegions() (regions []string, err error) {
 	return r.env.CSVInside("REGION", constants.PIAGeoChoices())
 }
+
+// GetPIAPort obtains the port to reach the PIA server on from the
+// environment variable PORT.
+func (r *reader) GetPIAPort() (port uint16, err error) {
+	n, err := r.env.IntRange("PORT", 0, 65535, libparams.Default("0"))
+	return uint16(n), err
+}
diff --git a/internal/provider/piav4.go b/internal/provider/piav4.go
index 801f0b7d7..4c15e75b6 100644
--- a/internal/provider/piav4.go
+++ b/internal/provider/piav4.go
@@ -6,6 +6,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 	"encoding/json"
+	"errors"
 	"fmt"
 	"io/ioutil"
 	"math/rand"
@@ -40,29 +41,63 @@ func newPrivateInternetAccess(servers []models.PIAServer, timeNow timeNowFunc) *
 	}
 }
 
-func (p *pia) GetOpenVPNConnection(selection models.ServerSelection) (
-	connection models.OpenVPNConnection, err error) {
-	var port uint16
+var (
+	ErrInvalidPort = errors.New("invalid port number")
+)
+
+func (p *pia) getPort(selection models.ServerSelection) (port uint16, err error) {
+	if selection.CustomPort == 0 {
+		switch selection.Protocol {
+		case constants.TCP:
+			switch selection.EncryptionPreset {
+			case constants.PIAEncryptionPresetNormal:
+				port = 502
+			case constants.PIAEncryptionPresetStrong:
+				port = 501
+			}
+		case constants.UDP:
+			switch selection.EncryptionPreset {
+			case constants.PIAEncryptionPresetNormal:
+				port = 1198
+			case constants.PIAEncryptionPresetStrong:
+				port = 1197
+			}
+		}
+
+		if port == 0 {
+			return 0, fmt.Errorf(
+				"%w: combination of protocol %q and encryption %q does not yield any port number",
+				ErrInvalidPort, selection.Protocol, selection.EncryptionPreset)
+		}
+		return port, nil
+	}
+
+	port = selection.CustomPort
 	switch selection.Protocol {
 	case constants.TCP:
-		switch selection.EncryptionPreset {
-		case constants.PIAEncryptionPresetNormal:
-			port = 502
-		case constants.PIAEncryptionPresetStrong:
-			port = 501
+		switch port {
+		case 80, 110, 443: //nolint:gomnd
+		default:
+			return 0, fmt.Errorf("%w: %d for protocol %s",
+				ErrInvalidPort, port, selection.Protocol)
 		}
 	case constants.UDP:
-		switch selection.EncryptionPreset {
-		case constants.PIAEncryptionPresetNormal:
-			port = 1198
-		case constants.PIAEncryptionPresetStrong:
-			port = 1197
+		switch port {
+		case 53, 1194, 1197, 1198, 8080, 9201: //nolint:gomnd
+		default:
+			return 0, fmt.Errorf("%w: %d for protocol %s",
+				ErrInvalidPort, port, selection.Protocol)
 		}
 	}
-	if port == 0 {
-		return connection, fmt.Errorf(
-			"combination of protocol %q and encryption %q does not yield any port number",
-			selection.Protocol, selection.EncryptionPreset)
+
+	return port, nil
+}
+
+func (p *pia) GetOpenVPNConnection(selection models.ServerSelection) (
+	connection models.OpenVPNConnection, err error) {
+	port, err := p.getPort(selection)
+	if err != nil {
+		return connection, err
 	}
 
 	servers := p.servers
diff --git a/internal/settings/providers.go b/internal/settings/providers.go
index 27db51e03..11defb81a 100644
--- a/internal/settings/providers.go
+++ b/internal/settings/providers.go
@@ -29,6 +29,10 @@ func GetPIASettings(paramsReader params.Reader) (settings models.ProviderSetting
 	if err != nil {
 		return settings, err
 	}
+	settings.ServerSelection.CustomPort, err = paramsReader.GetPIAPort()
+	if err != nil {
+		return settings, err
+	}
 	settings.PortForwarding.Enabled, err = paramsReader.GetPortForwarding()
 	if err != nil {
 		return settings, err