From 9f1cf8da5358a10ceaada28f2ee353a4306f8cde Mon Sep 17 00:00:00 2001
From: Justin Hiemstra <jhiemstra@wisc.edu>
Date: Tue, 1 Oct 2024 00:20:35 +0000
Subject: [PATCH] Implement TTL cache for unresolvable client IPs

This caches the randomly-assigned lat/long for a given unresolvable client IP,
ensuring that repeated requests from the same client are delivered a mostly-consistent
ordering of caches.
---
 director/sort.go      | 21 +++++++++++++++++++--
 director/sort_test.go | 18 ++++++++++++++++++
 2 files changed, 37 insertions(+), 2 deletions(-)

diff --git a/director/sort.go b/director/sort.go
index ae3ac8822..33cb91138 100644
--- a/director/sort.go
+++ b/director/sort.go
@@ -26,7 +26,9 @@ import (
 	"slices"
 	"sort"
 	"strings"
+	"time"
 
+	"github.com/jellydator/ttlcache/v3"
 	"github.com/pkg/errors"
 	log "github.com/sirupsen/logrus"
 
@@ -58,6 +60,9 @@ type (
 var (
 	invalidOverrideLogOnce = map[string]bool{}
 	geoIPOverrides         []GeoIPOverride
+
+	// Stores a mapping of client IPs that have been randomly assigned a coordinate
+	clientIpCache = ttlcache.New(ttlcache.WithTTL[netip.Addr, Coordinate](20 * time.Minute))
 )
 
 // Constants for the director sorting algorithm
@@ -194,7 +199,13 @@ func getClientLatLong(addr netip.Addr) (coord Coordinate) {
 	if !addr.IsValid() {
 		log.Warningf("Unable to sort servers based on client-server distance. Invalid client IP address: %s", addr.String())
 		coord.Lat, coord.Long = assignRandBoundedCoord(usLatMin, usLatMax, usLongMin, usLongMax)
-		log.Warningf("Assigning random location in the contiguous US to lat/long %f, %f ", coord.Lat, coord.Long)
+		cached, exists := clientIpCache.GetOrSet(addr, coord)
+		if exists {
+			log.Warningf("Retrieving pre-assigned lat/long for unresolved client IP %s: %f, %f.This assignment will be cached for 20 minutes unless the client IP is used again in that period.", addr.String(), coord.Lat, coord.Long)
+		} else {
+			log.Warningf("Assigning random location in the contiguous US to lat/long %f, %f. This assignment will be cached for 20 minutes unless the client IP is used again in that period.", coord.Lat, coord.Long)
+		}
+		coord = cached.Value()
 		return
 	}
 
@@ -204,7 +215,13 @@ func getClientLatLong(addr netip.Addr) (coord Coordinate) {
 			log.Warningf("Error while getting the client IP address: %v", err)
 		}
 		coord.Lat, coord.Long = assignRandBoundedCoord(usLatMin, usLatMax, usLongMin, usLongMax)
-		log.Warningf("Client IP %s has been re-assigned a random location in the contiguous US to lat/long %f, %f ", addr.String(), coord.Lat, coord.Long)
+		cached, exists := clientIpCache.GetOrSet(addr, coord)
+		if exists {
+			log.Warningf("Retrieving pre-assigned lat/long for client IP %s: %f, %f. This assignment will be cached for 20 minutes unless the client IP is used again in that period.", addr.String(), coord.Lat, coord.Long)
+		} else {
+			log.Warningf("Client IP %s has been re-assigned a random location in the contiguous US to lat/long %f, %f. This assignment will be cached for 20 minutes unless the client IP is used again in that period.", addr.String(), coord.Lat, coord.Long)
+		}
+		coord = cached.Value()
 	}
 	return
 }
diff --git a/director/sort_test.go b/director/sort_test.go
index 14087cc89..6d3b22b9f 100644
--- a/director/sort_test.go
+++ b/director/sort_test.go
@@ -450,11 +450,20 @@ func TestGetClientLatLong(t *testing.T) {
 		log.SetLevel(log.DebugLevel)
 
 		clientIp := netip.Addr{}
+		assert.False(t, clientIpCache.Has(clientIp))
 		coord := getClientLatLong(clientIp)
 
 		assert.True(t, coord.Lat <= usLatMax && coord.Lat >= usLatMin)
 		assert.True(t, coord.Long <= usLongMax && coord.Long >= usLongMin)
 		assert.Contains(t, logOutput.String(), "Unable to sort servers based on client-server distance. Invalid client IP address")
+		assert.NotContains(t, logOutput.String(), "Retrieving pre-assigned lat/long")
+
+		// Get it again to make sure it's coming from the cache
+		coord = getClientLatLong(clientIp)
+		assert.True(t, coord.Lat <= usLatMax && coord.Lat >= usLatMin)
+		assert.True(t, coord.Long <= usLongMax && coord.Long >= usLongMin)
+		assert.Contains(t, logOutput.String(), "Retrieving pre-assigned lat/long for unresolved client IP")
+		assert.True(t, clientIpCache.Has(clientIp))
 	})
 
 	t.Run("valid-ip-no-geoip-match", func(t *testing.T) {
@@ -463,10 +472,19 @@ func TestGetClientLatLong(t *testing.T) {
 		log.SetLevel(log.DebugLevel)
 
 		clientIp := netip.MustParseAddr("192.168.0.1")
+		assert.False(t, clientIpCache.Has(clientIp))
 		coord := getClientLatLong(clientIp)
 
 		assert.True(t, coord.Lat <= usLatMax && coord.Lat >= usLatMin)
 		assert.True(t, coord.Long <= usLongMax && coord.Long >= usLongMin)
 		assert.Contains(t, logOutput.String(), "Client IP 192.168.0.1 has been re-assigned a random location in the contiguous US to lat/long")
+		assert.NotContains(t, logOutput.String(), "Retrieving pre-assigned lat/long")
+
+		// Get it again to make sure it's coming from the cache
+		coord = getClientLatLong(clientIp)
+		assert.True(t, coord.Lat <= usLatMax && coord.Lat >= usLatMin)
+		assert.True(t, coord.Long <= usLongMax && coord.Long >= usLongMin)
+		assert.Contains(t, logOutput.String(), "Retrieving pre-assigned lat/long for client IP")
+		assert.True(t, clientIpCache.Has(clientIp))
 	})
 }