More cleanup and patches (specifically for servers), add support for download url mapping

dlrouter.yaml

@@ -1,4 +1,5 @@
 geodb: GeoIP2-City.mmdb
+dl_map: userdata.csv
 
 servers:
     - mirrors.tuna.tsinghua.edu.cn/armbian/

go.mod

@@ -5,10 +5,9 @@ go 1.17
 require (
 	github.com/oschwald/maxminddb-golang v1.8.0
 	github.com/sirupsen/logrus v1.8.1
+	github.com/spf13/viper v1.10.1
 )
 
-require github.com/spf13/viper v1.10.1
-
 require (
 	github.com/fsnotify/fsnotify v1.5.1 // indirect
 	github.com/hashicorp/hcl v1.0.0 // indirect

http.go

@@ -1,17 +1,25 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
 	"net"
 	"net/http"
 	"net/url"
 	"path"
+	"strings"
 )
 
 func statusRequest(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 }
 
+func serversRequest(w http.ResponseWriter, r *http.Request) {
+	w.Header().Set("Content-Type", "application/json")
+
+	json.NewEncoder(w).Encode(servers)
+}
+
 func redirectRequest(w http.ResponseWriter, r *http.Request) {
 	ipStr, _, err := net.SplitHostPort(r.RemoteAddr)
 
@@ -40,10 +48,18 @@ func redirectRequest(w http.ResponseWriter, r *http.Request) {
 		scheme = "https"
 	}
 
+	redirectPath := path.Join(server.Path, r.URL.Path)
+
+	if dlMap != nil {
+		if newPath, exists := dlMap[strings.TrimLeft(r.URL.Path, "/")]; exists {
+			redirectPath = path.Join(server.Path, newPath)
+		}
+	}
+
 	u := &url.URL{
 		Scheme: scheme,
 		Host:   server.Host,
-		Path:   path.Join(server.Path, r.URL.Path),
+		Path:   redirectPath,
 	}
 
 	w.Header().Set("X-Geo-Distance", fmt.Sprintf("%f", distance))

main.go

@@ -13,6 +13,8 @@ import (
 var (
 	db      *maxminddb.Reader
 	servers ServerList
+
+	dlMap map[string]string
 )
 
 // City represents a MaxmindDB city
@@ -41,6 +43,16 @@ func main() {
 		log.WithError(err).Fatalln("Unable to open database")
 	}
 
+	if mapFile := viper.GetString("dl_map"); mapFile != "" {
+		log.WithField("file", mapFile).Info("Loading download map")
+
+		dlMap, err = loadMap(mapFile)
+
+		if err != nil {
+			log.WithError(err).Fatalln("Unable to load download map")
+		}
+	}
+
 	serverList := viper.GetStringSlice("servers")
 
 	for _, server := range serverList {
@@ -108,6 +120,7 @@ func main() {
 	mux := http.NewServeMux()
 
 	mux.HandleFunc("/status", RealIPMiddleware(statusRequest))
+	mux.HandleFunc("/servers", RealIPMiddleware(serversRequest))
 	mux.HandleFunc("/", RealIPMiddleware(redirectRequest))
 
 	http.ListenAndServe(":8080", mux)

map.go

@@ -0,0 +1,38 @@
+package main
+
+import (
+	"encoding/csv"
+	"io"
+	"os"
+	"strings"
+)
+
+func loadMap(file string) (map[string]string, error) {
+	f, err := os.Open(file)
+
+	if err != nil {
+		return nil, err
+	}
+
+	defer f.Close()
+
+	m := make(map[string]string)
+
+	r := csv.NewReader(f)
+
+	for {
+		row, err := r.Read()
+
+		if err != nil {
+			if err == io.EOF {
+				break
+			}
+
+			return nil, err
+		}
+
+		m[strings.TrimLeft(row[0], "/")] = strings.TrimLeft(row[1], "/")
+	}
+
+	return m, nil
+}

servers.go

@@ -54,14 +54,18 @@ func (s ServerList) Check() {
 			if err != nil {
 				// This should never happen.
 				log.WithError(err).Warning("Invalid request! This should not happen, please check config.")
+				wg.Done()
 				return
 			}
 
 			res, err := checkClient.Do(req)
 
 			if err != nil {
+				if server.Available {
 					log.WithField("server", server.Host).Info("Server went offline")
 					server.Available = false
+				}
+				wg.Done()
 				return
 			}
 
@@ -70,6 +74,7 @@ func (s ServerList) Check() {
 				server.Available = true
 				log.WithField("server", server.Host).Info("Server is online")
 			}
+
 			wg.Done()
 		}(server)
 	}

userdata.csv

@@ -0,0 +1 @@
+bananapi/Bullseye_current,bananapi/archive/Armbian_21.08.1_Bananapi_bullseye_current_5.10.60.img.xz