Resolve issue with returning a lower number of servers, add auth to reload

assets/status-down.svg

@@ -0,0 +1 @@
+<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="82" height="20" role="img" aria-label="status: down"><title>status: down</title><linearGradient id="s" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><clipPath id="r"><rect width="82" height="20" rx="3" fill="#fff"/></clipPath><g clip-path="url(#r)"><rect width="43" height="20" fill="#555"/><rect x="43" width="39" height="20" fill="#e05d44"/><rect width="82" height="20" fill="url(#s)"/></g><g fill="#fff" text-anchor="middle" font-family="Verdana,Geneva,DejaVu Sans,sans-serif" text-rendering="geometricPrecision" font-size="110"><text aria-hidden="true" x="225" y="150" fill="#010101" fill-opacity=".3" transform="scale(.1)" textLength="330">status</text><text x="225" y="140" transform="scale(.1)" fill="#fff" textLength="330">status</text><text aria-hidden="true" x="615" y="150" fill="#010101" fill-opacity=".3" transform="scale(.1)" textLength="290">down</text><text x="615" y="140" transform="scale(.1)" fill="#fff" textLength="290">down</text></g></svg>

assets/status-unknown.svg

@@ -0,0 +1 @@
+<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="104" height="20" role="img" aria-label="status: unknown"><title>status: unknown</title><linearGradient id="s" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><clipPath id="r"><rect width="104" height="20" rx="3" fill="#fff"/></clipPath><g clip-path="url(#r)"><rect width="43" height="20" fill="#555"/><rect x="43" width="61" height="20" fill="#9f9f9f"/><rect width="104" height="20" fill="url(#s)"/></g><g fill="#fff" text-anchor="middle" font-family="Verdana,Geneva,DejaVu Sans,sans-serif" text-rendering="geometricPrecision" font-size="110"><text aria-hidden="true" x="225" y="150" fill="#010101" fill-opacity=".3" transform="scale(.1)" textLength="330">status</text><text x="225" y="140" transform="scale(.1)" fill="#fff" textLength="330">status</text><text aria-hidden="true" x="725" y="150" fill="#010101" fill-opacity=".3" transform="scale(.1)" textLength="510">unknown</text><text x="725" y="140" transform="scale(.1)" fill="#fff" textLength="510">unknown</text></g></svg>

assets/status-up.svg

@@ -0,0 +1 @@
+<svg xmlns="http://www.w3.org/2000/svg" xmlns:xlink="http://www.w3.org/1999/xlink" width="66" height="20" role="img" aria-label="status: up"><title>status: up</title><linearGradient id="s" x2="0" y2="100%"><stop offset="0" stop-color="#bbb" stop-opacity=".1"/><stop offset="1" stop-opacity=".1"/></linearGradient><clipPath id="r"><rect width="66" height="20" rx="3" fill="#fff"/></clipPath><g clip-path="url(#r)"><rect width="43" height="20" fill="#555"/><rect x="43" width="23" height="20" fill="#4c1"/><rect width="66" height="20" fill="url(#s)"/></g><g fill="#fff" text-anchor="middle" font-family="Verdana,Geneva,DejaVu Sans,sans-serif" text-rendering="geometricPrecision" font-size="110"><text aria-hidden="true" x="225" y="150" fill="#010101" fill-opacity=".3" transform="scale(.1)" textLength="330">status</text><text x="225" y="140" transform="scale(.1)" fill="#fff" textLength="330">status</text><text aria-hidden="true" x="535" y="150" fill="#010101" fill-opacity=".3" transform="scale(.1)" textLength="130">up</text><text x="535" y="140" transform="scale(.1)" fill="#fff" textLength="130">up</text></g></svg>

config.go

@@ -14,6 +14,8 @@ import (
 )
 
 func reloadConfig() {
+	log.Info("Loading configuration...")
+
 	err := viper.ReadInConfig() // Find and read the config file
 
 	if err != nil { // Handle errors reading the config file
@@ -46,6 +48,25 @@ func reloadConfig() {
 	// Reload server list
 	reloadServers()
 
+	// Create mirror map
+	mirrors := make(map[string][]*Server)
+
+	for _, server := range servers {
+		mirrors[server.Continent] = append(mirrors[server.Continent], server)
+	}
+
+	mirrors["default"] = append(mirrors["NA"], mirrors["EU"]...)
+
+	regionMap = mirrors
+
+	hosts := make(map[string]*Server)
+
+	for _, server := range servers {
+		hosts[server.Host] = server
+	}
+
+	hostMap = hosts
+
 	// Check top choices size
 	if topChoices > len(servers) {
 		topChoices = len(servers)

http.go

@@ -3,6 +3,8 @@ package main
 import (
 	"encoding/json"
 	"fmt"
+	"github.com/jmcvetta/randutil"
+	"github.com/spf13/viper"
 	"net"
 	"net/http"
 	"net/url"
@@ -16,31 +18,6 @@ func statusHandler(w http.ResponseWriter, r *http.Request) {
 	w.Write([]byte("OK"))
 }
 
-func legacyMirrorsHandler(w http.ResponseWriter, r *http.Request) {
-	w.Header().Set("Content-Type", "application/json")
-
-	mirrors := make(map[string][]string)
-
-	for _, server := range servers {
-		u := &url.URL{
-			Scheme: r.URL.Scheme,
-			Host:   server.Host,
-			Path:   server.Path,
-		}
-
-		mirrors[server.Continent] = append(mirrors[server.Continent], u.String())
-	}
-
-	mirrors["default"] = append(mirrors["NA"], mirrors["EU"]...)
-
-	json.NewEncoder(w).Encode(mirrors)
-}
-
-func mirrorsHandler(w http.ResponseWriter, r *http.Request) {
-	w.Header().Set("Content-Type", "application/json")
-	json.NewEncoder(w).Encode(servers)
-}
-
 func redirectHandler(w http.ResponseWriter, r *http.Request) {
 	ipStr, _, err := net.SplitHostPort(r.RemoteAddr)
 
@@ -61,13 +38,49 @@ func redirectHandler(w http.ResponseWriter, r *http.Request) {
 		ip = net.ParseIP(overrideIP)
 	}
 
-	server, distance, err := servers.Closest(ip)
+	var server *Server
+	var distance float64
+
+	if strings.HasPrefix(r.URL.Path, "/region") {
+		parts := strings.Split(r.URL.Path, "/")
+
+		// region = parts[2]
+		if mirrors, ok := regionMap[parts[2]]; ok {
+			choices := make([]randutil.Choice, len(mirrors))
+
+			for i, item := range mirrors {
+				if !item.Available {
+					continue
+				}
+
+				choices[i] = randutil.Choice{
+					Weight: item.Weight,
+					Item:   item,
+				}
+			}
+
+			choice, err := randutil.WeightedChoice(choices)
 
 			if err != nil {
 				http.Error(w, err.Error(), http.StatusInternalServerError)
 				return
 			}
 
+			server = choice.Item.(*Server)
+
+			r.URL.Path = strings.Join(parts[3:], "/")
+		}
+	}
+
+	if server == nil {
+		server, distance, err = servers.Closest(ip)
+
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+	}
+
 	scheme := r.URL.Scheme
 
 	if scheme == "" {
@@ -77,13 +90,7 @@ func redirectHandler(w http.ResponseWriter, r *http.Request) {
 	redirectPath := path.Join(server.Path, r.URL.Path)
 
 	if dlMap != nil {
-		p := r.URL.Path
-
-		if p[0] != '/' {
-			p = "/" + p
-		}
-
-		if newPath, exists := dlMap[p]; exists {
+		if newPath, exists := dlMap[strings.TrimLeft(r.URL.Path, "/")]; exists {
 			downloadsMapped.Inc()
 			redirectPath = path.Join(server.Path, newPath)
 		}
@@ -102,12 +109,29 @@ func redirectHandler(w http.ResponseWriter, r *http.Request) {
 	server.Redirects.Inc()
 	redirectsServed.Inc()
 
+	if distance > 0 {
 		w.Header().Set("X-Geo-Distance", fmt.Sprintf("%f", distance))
+	}
+
 	w.Header().Set("Location", u.String())
 	w.WriteHeader(http.StatusFound)
 }
 
 func reloadHandler(w http.ResponseWriter, r *http.Request) {
+	token := r.Header.Get("Authorization")
+
+	if token == "" || !strings.HasPrefix(token, "Bearer") || !strings.Contains(token, " ") {
+		w.WriteHeader(http.StatusUnauthorized)
+		return
+	}
+
+	token = token[strings.Index(token, " ")+1:]
+
+	if token != viper.GetString("reloadToken") {
+		w.WriteHeader(http.StatusUnauthorized)
+		return
+	}
+
 	reloadConfig()
 
 	w.WriteHeader(http.StatusOK)

main.go

@@ -20,8 +20,10 @@ import (
 var (
 	db         *maxminddb.Reader
 	servers    ServerList
-
+	regionMap  map[string][]*Server
+	hostMap    map[string]*Server
 	dlMap      map[string]string
+	topChoices int
 
 	redirectsServed = promauto.NewCounter(prometheus.CounterOpts{
 		Name: "armbian_router_redirects",
@@ -34,8 +36,6 @@ var (
 	})
 
 	serverCache *lru.Cache
-
-	topChoices int
 )
 
 type LocationLookup struct {
@@ -79,14 +79,20 @@ type ServerConfig struct {
 
 var (
 	configFlag = flag.String("config", "", "configuration file path")
+	flagDebug  = flag.Bool("debug", false, "Enable debug logging")
 )
 
 func main() {
 	flag.Parse()
 
+	if *flagDebug {
+		log.SetLevel(log.DebugLevel)
+	}
+
 	viper.SetDefault("bind", ":8080")
 	viper.SetDefault("cacheSize", 1024)
 	viper.SetDefault("topChoices", 3)
+	viper.SetDefault("reloadKey", randSeq(32))
 
 	viper.SetConfigName("dlrouter")        // name of config file (without extension)
 	viper.SetConfigType("yaml")            // REQUIRED if the config file does not have the extension in the name
@@ -110,8 +116,10 @@ func main() {
 	r.Use(RealIPMiddleware)
 	r.Use(logger.Logger("router", log.StandardLogger()))
 
+	r.Head("/status", statusHandler)
 	r.Get("/status", statusHandler)
 	r.Get("/mirrors", legacyMirrorsHandler)
+	r.Get("/mirrors/{server}.svg", mirrorStatusHandler)
 	r.Get("/mirrors.json", mirrorsHandler)
 	r.Post("/reload", reloadHandler)
 	r.Get("/dl_map", dlMapHandler)

map.go

@@ -4,6 +4,7 @@ import (
 	"encoding/csv"
 	"io"
 	"os"
+	"strings"
 )
 
 func loadMap(file string) (map[string]string, error) {
@@ -32,7 +33,7 @@ func loadMap(file string) (map[string]string, error) {
 			return nil, err
 		}
 
-		m[row[0]] = row[1]
+		m[strings.TrimLeft(row[0], "/")] = strings.TrimLeft(row[1], "/")
 	}
 
 	return m, nil

mirrors.go

@@ -0,0 +1,69 @@
+package main
+
+import (
+	_ "embed"
+	"encoding/json"
+	"github.com/go-chi/chi/v5"
+	"net/http"
+	"strings"
+)
+
+func legacyMirrorsHandler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+
+	mirrorOutput := make(map[string][]string)
+
+	for region, mirrors := range regionMap {
+		list := make([]string, len(mirrors))
+
+		for i, mirror := range mirrors {
+			list[i] = r.URL.Scheme + "://" + mirror.Host + "/" + strings.TrimLeft(mirror.Path, "/")
+		}
+
+		mirrorOutput[region] = list
+	}
+
+	json.NewEncoder(w).Encode(mirrorOutput)
+}
+
+func mirrorsHandler(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+	json.NewEncoder(w).Encode(servers)
+}
+
+var (
+	//go:embed assets/status-up.svg
+	statusUp []byte
+
+	//go:embed assets/status-down.svg
+	statusDown []byte
+
+	//go:embed assets/status-unknown.svg
+	statusUnknown []byte
+)
+
+func mirrorStatusHandler(w http.ResponseWriter, r *http.Request) {
+	serverHost := chi.URLParam(r, "server")
+
+	w.Header().Set("Content-Type", "image/svg+xml;charset=utf-8")
+
+	if serverHost == "" {
+		w.Write(statusUnknown)
+		return
+	}
+
+	serverHost = strings.Replace(serverHost, "_", ".", -1)
+
+	server, ok := hostMap[serverHost]
+
+	if !ok {
+		w.Write(statusUnknown)
+		return
+	}
+
+	if server.Available {
+		w.Write(statusUp)
+	} else {
+		w.Write(statusDown)
+	}
+}

servers.go

@@ -7,6 +7,7 @@ import (
 	"math"
 	"net"
 	"net/http"
+	"net/url"
 	"runtime"
 	"sort"
 	"strings"
@@ -16,7 +17,7 @@ import (
 
 var (
 	checkClient = &http.Client{
-		Timeout: 10 * time.Second,
+		Timeout: 20 * time.Second,
 	}
 )
 
@@ -29,10 +30,11 @@ type Server struct {
 	Weight     int                `json:"weight"`
 	Continent  string             `json:"continent"`
 	Redirects  prometheus.Counter `json:"-"`
+	LastChange time.Time          `json:"lastChange"`
 }
 
 func (server *Server) checkStatus() {
-	req, err := http.NewRequest(http.MethodGet, "https://"+server.Host+"/"+strings.TrimLeft(server.Path, "/"), nil)
+	req, err := http.NewRequest(http.MethodGet, "http://"+server.Host+"/"+strings.TrimLeft(server.Path, "/"), nil)
 
 	req.Header.Set("User-Agent", "ArmbianRouter/1.0 (Go "+runtime.Version()+")")
 
@@ -55,6 +57,7 @@ func (server *Server) checkStatus() {
 			}).Info("Server went offline")
 
 			server.Available = false
+			server.LastChange = time.Now()
 		} else {
 			log.WithFields(log.Fields{
 				"server": server.Host,
@@ -70,8 +73,38 @@ func (server *Server) checkStatus() {
 	}
 
 	if res.StatusCode == http.StatusOK || res.StatusCode == http.StatusMovedPermanently || res.StatusCode == http.StatusFound || res.StatusCode == http.StatusNotFound {
+		if res.StatusCode == http.StatusMovedPermanently || res.StatusCode == http.StatusFound {
+			location := res.Header.Get("Location")
+
+			responseFields["url"] = location
+
+			log.WithFields(responseFields).Debug("Server responded with redirect")
+
+			newUrl, err := url.Parse(location)
+
+			if err != nil {
+				if server.Available {
+					log.WithFields(responseFields).Warning("Server returned invalid url")
+					server.Available = false
+					server.LastChange = time.Now()
+				}
+				return
+			}
+
+			if newUrl.Scheme == "https" {
+				if server.Available {
+					responseFields["url"] = location
+					log.WithFields(responseFields).Warning("Server returned https url for http request")
+					server.Available = false
+					server.LastChange = time.Now()
+				}
+				return
+			}
+		}
+
 		if !server.Available {
 			server.Available = true
+			server.LastChange = time.Now()
 			log.WithFields(responseFields).Info("Server is online")
 		}
 	} else {
@@ -80,6 +113,7 @@ func (server *Server) checkStatus() {
 		if server.Available {
 			log.WithFields(responseFields).Info("Server went offline")
 			server.Available = false
+			server.LastChange = time.Now()
 		}
 	}
 }
@@ -157,9 +191,19 @@ func (s ServerList) Closest(ip net.IP) (*Server, float64, error) {
 			return c[i].Distance < c[j].Distance
 		})
 
-		choices := make([]randutil.Choice, topChoices)
+		choiceCount := topChoices
+
+		if len(c) < topChoices {
+			choiceCount = len(c)
+		}
+
+		choices := make([]randutil.Choice, choiceCount)
+
+		for i, item := range c[0:choiceCount] {
+			if item.Server == nil {
+				continue
+			}
 
-		for i, item := range c[0:topChoices] {
 			choices[i] = randutil.Choice{
 				Weight: item.Server.Weight,
 				Item:   item,

util.go

@@ -0,0 +1,13 @@
+package main
+
+import "math/rand"
+
+var letters = []rune("abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ")
+
+func randSeq(n int) string {
+	b := make([]rune, n)
+	for i := range b {
+		b[i] = letters[rand.Intn(len(letters))]
+	}
+	return string(b)
+}