From 5e659dc5b391075dfc374e6e16a31ea7d64e961e Mon Sep 17 00:00:00 2001
From: Quentin McGaw <quentin.mcgaw@gmail.com>
Date: Mon, 6 Jun 2022 03:04:58 +0000
Subject: [PATCH] feat(storage): add `keep` field for servers

---
 internal/models/server.go      | 15 ++++++++++
 internal/models/server_test.go |  2 ++
 internal/storage/merge.go      | 50 +++++++++++++++++++++++++++-------
 3 files changed, 57 insertions(+), 10 deletions(-)

diff --git a/internal/models/server.go b/internal/models/server.go
index de92ba37..695d44e4 100644
--- a/internal/models/server.go
+++ b/internal/models/server.go
@@ -1,8 +1,10 @@
 package models
 
 import (
+	"fmt"
 	"net"
 	"reflect"
+	"strings"
 )
 
 type Server struct {
@@ -25,6 +27,7 @@ type Server struct {
 	Free        bool     `json:"free,omitempty"`
 	Stream      bool     `json:"stream,omitempty"`
 	PortForward bool     `json:"port_forward,omitempty"`
+	Keep        bool     `json:"keep,omitempty"`
 	IPs         []net.IP `json:"ips,omitempty"`
 }
 
@@ -52,3 +55,15 @@ func ipsAreEqual(a, b []net.IP) (equal bool) {
 
 	return true
 }
+
+func (s *Server) Key() (key string) {
+	var protocols []string
+	if s.TCP {
+		protocols = append(protocols, "tcp")
+	}
+	if s.UDP {
+		protocols = append(protocols, "udp")
+	}
+
+	return fmt.Sprintf("%s-%s-%s", s.VPN, strings.Join(protocols, "-"), s.Hostname)
+}
diff --git a/internal/models/server_test.go b/internal/models/server_test.go
index 98daa516..4fa9e7c9 100644
--- a/internal/models/server_test.go
+++ b/internal/models/server_test.go
@@ -62,6 +62,7 @@ func Test_Server_Equal(t *testing.T) {
 				Stream:      true,
 				PortForward: true,
 				IPs:         []net.IP{net.IPv4(1, 2, 3, 4)},
+				Keep:        true,
 			},
 			b: Server{
 				VPN:         "vpn",
@@ -83,6 +84,7 @@ func Test_Server_Equal(t *testing.T) {
 				Stream:      true,
 				PortForward: true,
 				IPs:         []net.IP{net.IPv4(1, 2, 3, 4)},
+				Keep:        true,
 			},
 			equal: true,
 		},
diff --git a/internal/storage/merge.go b/internal/storage/merge.go
index 3b76517b..ad1d31e2 100644
--- a/internal/storage/merge.go
+++ b/internal/storage/merge.go
@@ -1,6 +1,7 @@
 package storage
 
 import (
+	"sort"
 	"time"
 
 	"github.com/qdm12/gluetun/internal/constants/providers"
@@ -26,18 +27,47 @@ func (s *Storage) mergeServers(hardcoded, persisted models.AllServers) models.Al
 
 func (s *Storage) mergeProviderServers(provider string,
 	hardcoded, persisted models.Servers) (merged models.Servers) {
-	if persisted.Timestamp <= hardcoded.Timestamp {
-		return hardcoded
+	if persisted.Timestamp > hardcoded.Timestamp {
+		diff := time.Unix(persisted.Timestamp, 0).Sub(time.Unix(hardcoded.Timestamp, 0))
+		if diff < 0 {
+			diff = -diff
+		}
+		diff = diff.Truncate(time.Second)
+		message := "Using " + provider + " servers from file which are " +
+			diff.String() + " more recent"
+		s.logger.Info(message)
+
+		return persisted
 	}
 
-	diff := time.Unix(persisted.Timestamp, 0).Sub(time.Unix(hardcoded.Timestamp, 0))
-	if diff < 0 {
-		diff = -diff
+	persistedServerKeyToServer := make(map[string]models.Server)
+	for _, persistedServer := range persisted.Servers {
+		if persistedServer.Keep {
+			persistedServerKeyToServer[persistedServer.Key()] = persistedServer
+		}
 	}
-	diff = diff.Truncate(time.Second)
-	message := "Using " + provider + " servers from file which are " +
-		diff.String() + " more recent"
-	s.logger.Info(message)
 
-	return persisted
+	merged = hardcoded // use all fields from hardcoded
+	merged.Servers = make([]models.Server, 0, len(hardcoded.Servers)+len(persistedServerKeyToServer))
+
+	for _, hardcodedServer := range hardcoded.Servers {
+		hardcodedServerKey := hardcodedServer.Key()
+		persistedServerToKeep, has := persistedServerKeyToServer[hardcodedServerKey]
+		if has {
+			// Drop hardcoded server and use persisted server matching the key.
+			merged.Servers = append(merged.Servers, persistedServerToKeep)
+			delete(persistedServerKeyToServer, hardcodedServerKey)
+		} else {
+			merged.Servers = append(merged.Servers, hardcodedServer)
+		}
+	}
+
+	// Add remaining persisted servers to keep
+	for _, persistedServer := range persistedServerKeyToServer {
+		merged.Servers = append(merged.Servers, persistedServer)
+	}
+
+	sort.Sort(models.SortableServers(merged.Servers))
+
+	return merged
 }