diff --git a/dlrouter.yaml b/dlrouter.yaml index d77f708..527014a 100644 --- a/dlrouter.yaml +++ b/dlrouter.yaml @@ -1,4 +1,5 @@ geodb: GeoIP2-City.mmdb +dl_map: userdata.csv servers: - mirrors.tuna.tsinghua.edu.cn/armbian/ diff --git a/go.mod b/go.mod index 582e1a0..7f8d378 100644 --- a/go.mod +++ b/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 diff --git a/http.go b/http.go index 3c7b610..eefd6bf 100644 --- a/http.go +++ b/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)) diff --git a/main.go b/main.go index 8c34f25..c8b2152 100644 --- a/main.go +++ b/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) diff --git a/map.go b/map.go new file mode 100644 index 0000000..efec3c4 --- /dev/null +++ b/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 +} diff --git a/servers.go b/servers.go index 32f7946..7b156be 100644 --- a/servers.go +++ b/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") - server.Available = false + 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) } diff --git a/userdata.csv b/userdata.csv new file mode 100644 index 0000000..2dd563e --- /dev/null +++ b/userdata.csv @@ -0,0 +1 @@ +bananapi/Bullseye_current,bananapi/archive/Armbian_21.08.1_Bananapi_bullseye_current_5.10.60.img.xz \ No newline at end of file