Added wip load balancer

main.go

@@ -69,11 +69,11 @@ var (
 	/*
 		Handler Modules
 	*/
-	sysdb              *database.Database        //System database
-	authAgent          *auth.AuthAgent           //Authentication agent
-	tlsCertManager     *tlscert.Manager          //TLS / SSL management
-	redirectTable      *redirection.RuleTable    //Handle special redirection rule sets
-	loadbalancer       *loadbalance.RouteManager //Load balancer manager to get routing targets from proxy rules
+	sysdb          *database.Database     //System database
+	authAgent      *auth.AuthAgent        //Authentication agent
+	tlsCertManager *tlscert.Manager       //TLS / SSL management
+	redirectTable  *redirection.RuleTable //Handle special redirection rule sets
+
 	pathRuleHandler    *pathrule.Handler         //Handle specific path blocking or custom headers
 	geodbStore         *geodb.Store              //GeoIP database, for resolving IP into country code
 	accessController   *access.Controller        //Access controller, handle black list and white list
@@ -88,6 +88,7 @@ var (
 	acmeAutoRenewer    *acme.AutoRenewer         //Handler for ACME auto renew ticking
 	staticWebServer    *webserv.WebServer        //Static web server for hosting simple stuffs
 	forwardProxy       *forwardproxy.Handler     //HTTP Forward proxy, basically VPN for web browser
+	loadBalancer       *loadbalance.RouteManager //Global scope loadbalancer, store the state of the lb routing
 
 	//Helper modules
 	EmailSender       *email.Sender         //Email sender that handle email sending

mod/dynamicproxy/Server.go

@@ -20,6 +20,7 @@ import (
 		- Access Router
 			- Blacklist
 			- Whitelist
+		- Rate Limitor
 		- Basic Auth
 		- Vitrual Directory Proxy
 		- Subdomain Proxy
@@ -30,7 +31,7 @@ func (h *ProxyHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
 	/*
 		Special Routing Rules, bypass most of the limitations
 	*/
-	//Check if there are external routing rule matches.
+	//Check if there are external routing rule (rr) matches.
 	//If yes, route them via external rr
 	matchedRoutingRule := h.Parent.GetMatchingRoutingRule(r)
 	if matchedRoutingRule != nil {

mod/dynamicproxy/dynamicproxy.go

@@ -28,6 +28,7 @@ func NewDynamicProxy(option RouterOption) (*Router, error) {
 		Running:          false,
 		server:           nil,
 		routingRules:     []*RoutingRule{},
+		loadBalancer:     option.LoadBalancer,
 		rateLimitCounter: RequestCountPerIpTable{},
 	}
 

mod/dynamicproxy/loadbalance/loadbalance.go

@@ -1,9 +1,11 @@
 package loadbalance
 
 import (
+	"sync"
+	"time"
+
 	"imuslab.com/zoraxy/mod/geodb"
 	"imuslab.com/zoraxy/mod/info/logger"
-	"imuslab.com/zoraxy/mod/uptime"
 )
 
 /*
@@ -14,47 +16,81 @@ import (
 
 type BalancePolicy int
 
-const (
-	BalancePolicy_RoundRobin BalancePolicy = 0 //Round robin, will ignore upstream if down
-	BalancePolicy_Fallback   BalancePolicy = 1 //Fallback only. Will only switch to next node if the first one failed
-	BalancePolicy_Random     BalancePolicy = 2 //Random, randomly pick one from the list that is online
-	BalancePolicy_GeoRegion  BalancePolicy = 3 //Use the one defined for this geo-location, when down, pick the next avaible node
-)
+// Data structure to hold the upstream server information
+type Upstream struct {
+	ProxyTargetIpOrDomain string   //Target IP address or domain name with port
+	RequireTLS            bool     //Require TLS connection
+	PreferedCountryCode   []string //Prefered country codes, default to empty string slice (not assigned)
+	Priority              int      //Prirotiy of fallback, set all to 0 for round robin
+}
 
+// LoadBalanceRule stores the load balance option of this host. If there are no usable upstreams for lb, the primary one will be used
 type LoadBalanceRule struct {
-	Upstreams         []string      //Reverse proxy upstream servers
-	LoadBalancePolicy BalancePolicy //Policy in deciding which target IP to proxy
-	UseRegionLock     bool          //If this is enabled with BalancePolicy_Geo, when the main site failed, it will not pick another node
-	UseStickySession  bool          //Use sticky session, if you are serving EU countries, make sure to add the "Do you want cookie" warning
+	Upstreams        []Upstream //Reverse proxy upstream servers
+	UseStickySession bool       //Enable stick session
+}
 
-	parent *RouteManager
+// PrimaryRoutingRule stores the primary routing target of this host
+type PrimaryRoutingRule struct {
+	MatchingDomainOrIp string //Primary proxy upstream origin server
+	RequireTLS         bool   //Target require TLS
 }
 
 type Options struct {
-	Geodb         *geodb.Store    //GeoIP resolver for checking incoming request origin country
-	UptimeMonitor *uptime.Monitor //For checking if the target is online, this might be nil when the module starts
+	Geodb  *geodb.Store //GeoIP resolver for checking incoming request origin country
+	Logger *logger.Logger
 }
 
 type RouteManager struct {
-	Options Options
-	Logger  *logger.Logger
+	LoadBalanceMap  sync.Map //Sync map to store the last load balance state of a given node
+	OnlineStatusMap sync.Map //Sync map to store the online status of a given ip address or domain name
+
+	onlineStatusTickerStop chan bool //Stopping channel for the online status pinger
+	Options                Options   //Options for the load balancer
 }
 
-// Create a new load balance route manager
-func NewRouteManager(options *Options, logger *logger.Logger) *RouteManager {
-	newManager := RouteManager{
-		Options: *options,
-		Logger:  logger,
+// Create a new load balancer
+func NewLoadBalancer(options *Options) *RouteManager {
+	onlineStatusCheckerStopChan := make(chan bool)
+
+	return &RouteManager{
+		LoadBalanceMap:         sync.Map{},
+		OnlineStatusMap:        sync.Map{},
+		onlineStatusTickerStop: onlineStatusCheckerStopChan,
+		Options:                *options,
 	}
-	logger.PrintAndLog("INFO", "Load Balance Route Manager started", nil)
-	return &newManager
 }
 
-func (b *LoadBalanceRule) GetProxyTargetIP() {
-//TODO: Implement get proxy target IP logic here
+func (m *RouteManager) UpdateKeepAliveTargets(pingTargets []*PrimaryRoutingRule) {
+	ticker := time.NewTicker(1 * time.Minute)
+	defer ticker.Stop()
+
+	go func() {
+		for {
+			select {
+			case <-m.onlineStatusTickerStop:
+				ticker.Stop()
+				return
+			case <-ticker.C:
+				for _, target := range pingTargets {
+					isOnline := PingTarget(target)
+					m.LoadBalanceMap.Store(target.MatchingDomainOrIp, isOnline)
+				}
+			}
+		}
+	}()
+}
+
+// GetProxyTargetIP Get the proxy target given the primary routing rule and load balance options
+func (m *RouteManager) GetProxyTargetIP(pr *PrimaryRoutingRule, lbr *LoadBalanceRule) {
+
+}
+
+func (m *RouteManager) Close() {
+	m.onlineStatusTickerStop <- true
 }
 
 // Print debug message
 func (m *RouteManager) debugPrint(message string, err error) {
-	m.Logger.PrintAndLog("LB", message, err)
+	m.Options.Logger.PrintAndLog("LoadBalancer", message, err)
 }

mod/dynamicproxy/loadbalance/onlineStatus.go

@@ -0,0 +1,39 @@
+package loadbalance
+
+import (
+	"net/http"
+	"time"
+)
+
+// Return the last ping status to see if the target is online
+func (m *RouteManager) IsTargetOnline(matchingDomainOrIp string) bool {
+	value, ok := m.LoadBalanceMap.Load(matchingDomainOrIp)
+	if !ok {
+		return false
+	}
+
+	isOnline, ok := value.(bool)
+	return ok && isOnline
+}
+
+// Ping a target to see if it is online
+func PingTarget(target *PrimaryRoutingRule) bool {
+	client := &http.Client{
+		Timeout: 10 * time.Second,
+	}
+
+	url := target.MatchingDomainOrIp
+	if target.RequireTLS {
+		url = "https://" + url
+	} else {
+		url = "http://" + url
+	}
+
+	resp, err := client.Get(url)
+	if err != nil {
+		return false
+	}
+	defer resp.Body.Close()
+
+	return resp.StatusCode == http.StatusOK
+}

mod/dynamicproxy/proxyRequestHandler.go

@@ -16,6 +16,7 @@ import (
 	"imuslab.com/zoraxy/mod/websocketproxy"
 )
 
+// Check if the request URI matches any of the proxy endpoint
 func (router *Router) getTargetProxyEndpointFromRequestURI(requestURI string) *ProxyEndpoint {
 	var targetProxyEndpoint *ProxyEndpoint = nil
 	router.ProxyEndpoints.Range(func(key, value interface{}) bool {
@@ -30,6 +31,7 @@ func (router *Router) getTargetProxyEndpointFromRequestURI(requestURI string) *P
 	return targetProxyEndpoint
 }
 
+// Get the proxy endpoint from hostname, which might includes checking of wildcard certificates
 func (router *Router) getProxyEndpointFromHostname(hostname string) *ProxyEndpoint {
 	var targetSubdomainEndpoint *ProxyEndpoint = nil
 	ep, ok := router.ProxyEndpoints.Load(hostname)

mod/dynamicproxy/typedef.go

@@ -8,6 +8,7 @@ import (
 
 	"imuslab.com/zoraxy/mod/access"
 	"imuslab.com/zoraxy/mod/dynamicproxy/dpcore"
+	"imuslab.com/zoraxy/mod/dynamicproxy/loadbalance"
 	"imuslab.com/zoraxy/mod/dynamicproxy/permissionpolicy"
 	"imuslab.com/zoraxy/mod/dynamicproxy/redirection"
 	"imuslab.com/zoraxy/mod/geodb"
@@ -25,23 +26,26 @@ type ProxyHandler struct {
 	Parent *Router
 }
 
+/* Router Object Options */
 type RouterOption struct {
-	HostUUID           string //The UUID of Zoraxy, use for heading mod
-	HostVersion        string //The version of Zoraxy, use for heading mod
-	Port               int    //Incoming port
-	UseTls             bool   //Use TLS to serve incoming requsts
-	ForceTLSLatest     bool   //Force TLS1.2 or above
-	NoCache            bool   //Force set Cache-Control: no-store
-	ListenOnPort80     bool   //Enable port 80 http listener
-	ForceHttpsRedirect bool   //Force redirection of http to https endpoint
-	TlsManager         *tlscert.Manager
-	RedirectRuleTable  *redirection.RuleTable
-	GeodbStore         *geodb.Store       //GeoIP resolver
-	AccessController   *access.Controller //Blacklist / whitelist controller
-	StatisticCollector *statistic.Collector
-	WebDirectory       string //The static web server directory containing the templates folder
+	HostUUID           string                    //The UUID of Zoraxy, use for heading mod
+	HostVersion        string                    //The version of Zoraxy, use for heading mod
+	Port               int                       //Incoming port
+	UseTls             bool                      //Use TLS to serve incoming requsts
+	ForceTLSLatest     bool                      //Force TLS1.2 or above
+	NoCache            bool                      //Force set Cache-Control: no-store
+	ListenOnPort80     bool                      //Enable port 80 http listener
+	ForceHttpsRedirect bool                      //Force redirection of http to https endpoint
+	TlsManager         *tlscert.Manager          //TLS manager for serving SAN certificates
+	RedirectRuleTable  *redirection.RuleTable    //Redirection rules handler and table
+	GeodbStore         *geodb.Store              //GeoIP resolver
+	AccessController   *access.Controller        //Blacklist / whitelist controller
+	StatisticCollector *statistic.Collector      //Statistic collector for storing stats on incoming visitors
+	WebDirectory       string                    //The static web server directory containing the templates folder
+	LoadBalancer       *loadbalance.RouteManager //Load balancer that handle load balancing of proxy target
 }
 
+/* Router Object */
 type Router struct {
 	Option         *RouterOption
 	ProxyEndpoints *sync.Map
@@ -50,6 +54,7 @@ type Router struct {
 	mux            http.Handler
 	server         *http.Server
 	tlsListener    net.Listener
+	loadBalancer   *loadbalance.RouteManager //Load balancer routing manager
 	routingRules   []*RoutingRule
 
 	tlsRedirectStop  chan bool              //Stop channel for tls redirection server
@@ -57,6 +62,7 @@ type Router struct {
 	rateLimitCounter RequestCountPerIpTable //Request counter for rate limter
 }
 
+/* Basic Auth Related Data structure*/
 // Auth credential for basic auth on certain endpoints
 type BasicAuthCredentials struct {
 	Username     string
@@ -74,6 +80,7 @@ type BasicAuthExceptionRule struct {
 	PathPrefix string
 }
 
+/* Custom Header Related Data structure */
 // Header injection direction type
 type HeaderDirection int
 
@@ -90,6 +97,8 @@ type UserDefinedHeader struct {
 	IsRemove  bool //Instead of set, remove this key instead
 }
 
+/* Routing Rule Data Structures */
+
 // A Virtual Directory endpoint, provide a subset of ProxyEndpoint for better
 // program structure than directly using ProxyEndpoint
 type VirtualDirectoryEndpoint struct {

reverseproxy.go

@@ -96,6 +96,7 @@ func ReverseProxtInit() {
 		StatisticCollector: statisticCollector,
 		WebDirectory:       *staticWebServerRoot,
 		AccessController:   accessController,
+		LoadBalancer:       loadBalancer,
 	})
 	if err != nil {
 		SystemWideLogger.PrintAndLog("Proxy", "Unable to create dynamic proxy router", err)
@@ -145,9 +146,6 @@ func ReverseProxtInit() {
 			MaxRecordsStore: 288, //1 day
 		})
 
-		//Pass the pointer of this uptime monitor into the load balancer
-		loadbalancer.Options.UptimeMonitor = uptimeMonitor
-
 		SystemWideLogger.Println("Uptime Monitor background service started")
 	}()
 }

start.go

@@ -102,10 +102,11 @@ func startupSequence() {
 		panic(err)
 	}
 
-	//Create a load balance route manager
-	loadbalancer = loadbalance.NewRouteManager(&loadbalance.Options{
-		Geodb: geodbStore,
-	}, SystemWideLogger)
+	//Create a load balancer
+	loadBalancer = loadbalance.NewLoadBalancer(&loadbalance.Options{
+		Geodb:  geodbStore,
+		Logger: SystemWideLogger,
+	})
 
 	//Create the access controller
 	accessController, err = access.NewAccessController(&access.Options{