From 6ac03c18ccaca3c60999f0d277aa75330ffffc4a Mon Sep 17 00:00:00 2001
From: Yifu Yu <root@jackyyf.com>
Date: Thu, 26 Dec 2019 01:10:26 +0800
Subject: [PATCH] Navigator now supports unweighted and weighted load
 balancing.

---
 mapping/elf/core.go | 53 ++++++++++++++++++++++++++++++++++++++++++++-
 mapping/map.go      | 15 ++++++++-----
 rules/map.starlark  | 19 ++++++++++++----
 3 files changed, 76 insertions(+), 11 deletions(-)

diff --git a/mapping/elf/core.go b/mapping/elf/core.go
index 09d3b42..1387d88 100644
--- a/mapping/elf/core.go
+++ b/mapping/elf/core.go
@@ -4,11 +4,14 @@ import (
 	"flag"
 	"fmt"
 	"log"
+	"math/rand"
 	"os"
 	"os/signal"
+	"sort"
 	"sync"
 	"sync/atomic"
 	"syscall"
+	"time"
 
 	"go.starlark.net/starlark"
 )
@@ -36,6 +39,7 @@ var (
 	reloadSignal = make(chan os.Signal)
 	pool         *threadPool
 	counter      uint64
+	rnd          = rand.New(rand.NewSource(time.Now().UnixNano()))
 )
 
 func newThreadPool(size int) *threadPool {
@@ -147,9 +151,56 @@ func GetMapping(ip string) string {
 		log.Println("Starlark execute error:", err.Error())
 		return ""
 	}
-	if r, ok := ret.(starlark.String); ok {
+	switch r := ret.(type) {
+	case starlark.String:
 		return string(r)
+	case *starlark.List:
+		val := r.Index(rnd.Intn(r.Len()))
+		if s, ok := val.(starlark.String); ok {
+			return string(s)
+		}
+		goto scriptError
+	case starlark.Tuple:
+		val := r.Index(rnd.Intn(r.Len()))
+		if s, ok := val.(starlark.String); ok {
+			return string(s)
+		}
+		goto scriptError
+	case *starlark.Dict:
+		values := r.Items()
+		options := make([]string, r.Len())
+		tot := int64(0)
+		for i := range values {
+			if v, ok := values[i].Index(0).(starlark.String); ok {
+				options[i] = string(v)
+			} else {
+				goto scriptError
+			}
+			if v, ok := values[i].Index(1).(starlark.Int); ok {
+				if weight, ok := v.Int64(); ok {
+					tot += weight
+				} else {
+					goto scriptError
+				}
+			} else {
+				goto scriptError
+			}
+		}
+		sort.Strings(options)
+		pos := rnd.Int63n(tot)
+		for _, option := range options {
+			w, _, _ := r.Get(starlark.String(option))
+			v, _ := w.(starlark.Int).Int64()
+			pos -= v
+			if pos < 0 {
+				return option
+			}
+		}
+		log.Println("Unexpected code execution,",
+			"check random logic for *starlark.Dict")
+		return options[len(options)-1]
 	}
+scriptError:
 	log.Println("Script returned unexpected result:", ret.String())
 	return ""
 }
diff --git a/mapping/map.go b/mapping/map.go
index 646cadc..92c88d1 100644
--- a/mapping/map.go
+++ b/mapping/map.go
@@ -7,8 +7,7 @@ import (
 
 var (
 	defaultServer = flag.String("fallback-node",
-		"xe-mci1-us.edge.eve.network",
-		"Default CDN node in case of any script error")
+		"xe-mci1-us", "Default CDN node in case of any script error")
 	defaultSuffix = flag.String("fallback-suffix",
 		".edge.eve.network",
 		"Default CDN suffix in case of any script error")
@@ -19,11 +18,15 @@ func Initialize() {
 }
 
 func Get(ip string) string {
-	ret := elf.GetMapping(ip)
-	if ret == "" {
-		return *defaultServer
+	node := elf.GetMapping(ip)
+	if node == "" {
+		node = *defaultServer
 	}
-	return ret
+	suffix := elf.GetSuffix(ip)
+	if suffix == "" {
+		suffix = *defaultSuffix
+	}
+	return node + suffix
 }
 
 func GetSuffix(ip string) string {
diff --git a/rules/map.starlark b/rules/map.starlark
index 2ed82eb..bd444c0 100644
--- a/rules/map.starlark
+++ b/rules/map.starlark
@@ -1,6 +1,9 @@
 WHOLESALE_INTERNET_10GE = "xe-mci1-us"
 HETZNER_FSN_1GE         = "ge-fsn1-de"
 HETZNER_HEL_1GE         = "ge-hel1-fi"
+CLOUDCONE_LAX1_1GE		= "ge-lax1-us"
+CLOUDCONE_LAX2_1GE		= "ge-lax2-us"
+CLOUDCONE_LAX_LB		= (CLOUDCONE_LAX1_1GE, CLOUDCONE_LAX2_1GE)
 default_server          = WHOLESALE_INTERNET_10GE
 
 CHINA_MAINLAND_SUFFIX = ".eveedge.link"
@@ -11,12 +14,20 @@ default_suffix        = GLOBAL_SUFFIX
 def getMapping(ip):
 	info = geoLookup(ip)
 	if not info:
-		return default_server + default_suffix
+		return default_server
 	if info.IspDomain == "ChinaMobile":
-		return HETZNER_FSN_1GE + CHINA_MAINLAND_SUFFIX
+		return HETZNER_FSN_1GE
+	if info.IspDomain == "ChinaTelecom":
+		return CLOUDCONE_LAX_LB
 	if info.CountryCode == "CN":
-		return default_server + CHINA_MAINLAND_SUFFIX
-	return default_server + GLOBAL_SUFFIX
+		return default_server
+	if ip == "0.11.45.14":
+		return {
+			default_server: 3,
+			CLOUDCONE_LAX1_1GE: 2,
+			CLOUDCONE_LAX2_1GE: 1,
+		}
+	return default_server
 
 def getNodes():
 	return ["xe-mci1-us", "ge-fsn1-de", "ge-lax1-us"]