From 20ae76ff061bfaa984d3b0f23594417f5d0e796b Mon Sep 17 00:00:00 2001
From: Tyler <tystuyfzand@gmail.com>
Date: Thu, 31 Mar 2022 22:43:18 -0400
Subject: [PATCH] Fix distance sorting, add env variable for local IP to test

---
 http.go    | 12 +++++++++---
 servers.go | 30 ++++++++++++++----------------
 2 files changed, 23 insertions(+), 19 deletions(-)

diff --git a/http.go b/http.go
index 6dcc8d4..89e01e1 100644
--- a/http.go
+++ b/http.go
@@ -6,6 +6,7 @@ import (
 	"net"
 	"net/http"
 	"net/url"
+	"os"
 	"path"
 	"strings"
 )
@@ -50,9 +51,14 @@ func redirectHandler(w http.ResponseWriter, r *http.Request) {
 
 	ip := net.ParseIP(ipStr)
 
-	// TODO: This is temporary to allow testing on private addresses.
-	if ip.IsPrivate() {
-		ip = net.ParseIP("1.1.1.1")
+	if ip.IsLoopback() || ip.IsPrivate() {
+		overrideIP := os.Getenv("OVERRIDE_IP")
+
+		if overrideIP == "" {
+			overrideIP = "1.1.1.1"
+		}
+
+		ip = net.ParseIP(overrideIP)
 	}
 
 	server, distance, err := servers.Closest(ip)
diff --git a/servers.go b/servers.go
index c93e35d..22ef708 100644
--- a/servers.go
+++ b/servers.go
@@ -123,19 +123,6 @@ type ComputedDistance struct {
 // DistanceList is a list of Computed Distances with an easy "Choices" func
 type DistanceList []ComputedDistance
 
-func (d DistanceList) Choices() []randutil.Choice {
-	c := make([]randutil.Choice, len(d))
-
-	for i, item := range d {
-		c[i] = randutil.Choice{
-			Weight: item.Server.Weight,
-			Item:   item,
-		}
-	}
-
-	return c
-}
-
 // Closest will use GeoIP on the IP provided and find the closest servers.
 // When we have a list of x servers closest, we can choose a random or weighted one.
 // Return values are the closest server, the distance, and if an error occurred.
@@ -157,18 +144,29 @@ func (s ServerList) Closest(ip net.IP) (*Server, float64, error) {
 				continue
 			}
 
+			distance := Distance(city.Location.Latitude, city.Location.Longitude, server.Latitude, server.Longitude)
+
 			c[i] = ComputedDistance{
 				Server:   server,
-				Distance: Distance(city.Location.Latitude, city.Location.Longitude, server.Latitude, server.Longitude),
+				Distance: distance,
 			}
 		}
 
 		// Sort by distance
-		sort.Slice(s, func(i int, j int) bool {
+		sort.Slice(c, func(i int, j int) bool {
 			return c[i].Distance < c[j].Distance
 		})
 
-		choiceInterface = c[0:topChoices].Choices()
+		choices := make([]randutil.Choice, topChoices)
+
+		for i, item := range c[0:topChoices] {
+			choices[i] = randutil.Choice{
+				Weight: item.Server.Weight,
+				Item:   item,
+			}
+		}
+
+		choiceInterface = choices
 
 		serverCache.Add(ip.String(), choiceInterface)
 	}