Add dl_map endpoint

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,18 +1,28 @@
 package main
 
 import (
+	"encoding/json"
 	"fmt"
+	log "github.com/sirupsen/logrus"
+	"github.com/spf13/viper"
 	"net"
 	"net/http"
 	"net/url"
 	"path"
+	"strings"
 )
 
-func statusRequest(w http.ResponseWriter, r *http.Request) {
+func statusHandler(w http.ResponseWriter, r *http.Request) {
 	w.WriteHeader(http.StatusOK)
 }
 
-func redirectRequest(w http.ResponseWriter, r *http.Request) {
+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)
 
 	if err != nil {
@@ -40,13 +50,51 @@ 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))
 	w.Header().Set("Location", u.String())
 	w.WriteHeader(http.StatusFound)
 }
+
+func reloadHandler(w http.ResponseWriter, r *http.Request) {
+	if mapFile := viper.GetString("dl_map"); mapFile != "" {
+		log.WithField("file", mapFile).Info("Loading download map")
+
+		newMap, err := loadMap(mapFile)
+
+		if err != nil {
+			http.Error(w, err.Error(), http.StatusInternalServerError)
+			return
+		}
+
+		dlMap = newMap
+
+		return
+	}
+
+	w.WriteHeader(http.StatusNotFound)
+}
+
+func dlMapHandler(w http.ResponseWriter, r *http.Request) {
+	if dlMap == nil {
+		w.WriteHeader(http.StatusNotFound)
+		return
+	}
+
+	w.Header().Set("Content-Type", "application/json")
+
+	json.NewEncoder(w).Encode(dlMap)
+}

main.go

@@ -13,6 +13,8 @@ import (
 var (
 	db      *maxminddb.Reader
 	servers ServerList
+
+	dlMap map[string]string
 )
 
 // City represents a MaxmindDB city
@@ -24,6 +26,8 @@ type City struct {
 }
 
 func main() {
+	viper.SetDefault("bind", ":8080")
+
 	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
 	viper.AddConfigPath("/etc/dlrouter/")  // path to look for the config file in
@@ -41,6 +45,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 {
@@ -107,8 +121,11 @@ func main() {
 
 	mux := http.NewServeMux()
 
-	mux.HandleFunc("/status", RealIPMiddleware(statusRequest))
+	mux.HandleFunc("/status", statusHandler)
-	mux.HandleFunc("/", RealIPMiddleware(redirectRequest))
+	mux.HandleFunc("/mirrors", mirrorsHandler)
+	mux.HandleFunc("/reload", reloadHandler)
+	mux.HandleFunc("/dl_map", dlMapHandler)
+	mux.HandleFunc("/", RealIPMiddleware(redirectHandler))
 
-	http.ListenAndServe(":8080", mux)
+	http.ListenAndServe(viper.GetString("bind"), 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 {
-				log.WithField("server", server.Host).Info("Server went offline")
+				if server.Available {
-				server.Available = false
+					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