From 4696ec4ec2e344d15a4e2e4444caabd963da82ac Mon Sep 17 00:00:00 2001
From: Henry Camacho <henry.f.camacho@gmail.com>
Date: Sun, 19 Mar 2017 09:56:33 -0500
Subject: [PATCH] Added support for default load balancer method.

---
 go-rvpn-server.yaml                      |  3 +++
 main.go                                  |  8 ++++++--
 rvpn/genericlistener/listener_generic.go | 17 +++++++++--------
 rvpn/genericlistener/manager.go          |  6 +++++-
 4 files changed, 23 insertions(+), 11 deletions(-)

diff --git a/go-rvpn-server.yaml b/go-rvpn-server.yaml
index d0f88e1..b06cfe4 100644
--- a/go-rvpn-server.yaml
+++ b/go-rvpn-server.yaml
@@ -13,4 +13,7 @@ rvpn:
       secret: abc123
     test3.daplie.me:
       secret: abc123
+  loadbalancing:
+    defaultmethod: 'round-robin'
+    
 
diff --git a/main.go b/main.go
index 3d1e4c5..1c5115f 100644
--- a/main.go
+++ b/main.go
@@ -32,6 +32,7 @@ var (
 	idle                     int
 	dwell                    int
 	cancelcheck              int
+	lbDefaultMethod          string
 )
 
 func init() {
@@ -59,6 +60,7 @@ func main() {
 	idle = deadtime.(map[string]interface{})["idle"].(int)
 	dwell = deadtime.(map[string]interface{})["dwell"].(int)
 	cancelcheck = deadtime.(map[string]interface{})["cancelcheck"].(int)
+	lbDefaultMethod = viper.Get("rvpn.loadbalancing.defaultmethod").(string)
 
 	loginfo.Println("startup")
 
@@ -91,8 +93,10 @@ func main() {
 	connectionTable = genericlistener.NewTable(dwell, idle)
 	go connectionTable.Run(ctx)
 
-	genericListeners := genericlistener.NewGenerListeners(ctx, connectionTable, connectionTracking, secretKey, certbundle, wssHostName, adminHostName, cancelcheck)
-	go genericListeners.Run(ctx, argGenericBinding)
+	genericListeners := genericlistener.NewGenerListeners(ctx, connectionTable, connectionTracking, secretKey, certbundle, wssHostName,
+		adminHostName, cancelcheck, lbDefaultMethod)
+
+	genericListeners.Run(ctx, argGenericBinding)
 
 	//Run for 10 minutes and then shutdown cleanly
 	time.Sleep(6000 * time.Second)
diff --git a/rvpn/genericlistener/listener_generic.go b/rvpn/genericlistener/listener_generic.go
index 5ad20c8..f221b5b 100644
--- a/rvpn/genericlistener/listener_generic.go
+++ b/rvpn/genericlistener/listener_generic.go
@@ -26,14 +26,15 @@ type contextKey string
 
 //CtxConnectionTrack
 const (
-	ctxSecretKey            contextKey = "secretKey"
-	ctxConnectionTable      contextKey = "connectionTable"
-	ctxConfig               contextKey = "config"
-	ctxListenerRegistration contextKey = "listenerRegistration"
-	ctxConnectionTrack      contextKey = "connectionTrack"
-	ctxWssHostName          contextKey = "wsshostname"
-	ctxAdminHostName        contextKey = "adminHostName"
-	ctxCancelCheck          contextKey = "cancelcheck"
+	ctxSecretKey                contextKey = "secretKey"
+	ctxConnectionTable          contextKey = "connectionTable"
+	ctxConfig                   contextKey = "config"
+	ctxListenerRegistration     contextKey = "listenerRegistration"
+	ctxConnectionTrack          contextKey = "connectionTrack"
+	ctxWssHostName              contextKey = "wsshostname"
+	ctxAdminHostName            contextKey = "adminHostName"
+	ctxCancelCheck              contextKey = "cancelcheck"
+	ctxLoadbalanceDefaultMethod contextKey = "lbdefaultmethod"
 )
 
 const (
diff --git a/rvpn/genericlistener/manager.go b/rvpn/genericlistener/manager.go
index c05c58f..dc6c098 100644
--- a/rvpn/genericlistener/manager.go
+++ b/rvpn/genericlistener/manager.go
@@ -56,10 +56,12 @@ type GenericListeners struct {
 	wssHostName        string
 	adminHostName      string
 	cancelCheck        int
+	lbDefaultMethod    string
 }
 
 //NewGenerListeners --
-func NewGenerListeners(ctx context.Context, connectionTable *Table, connectionTrack *Tracking, secretKey string, certbundle tls.Certificate, wssHostName string, adminHostName string, cancelCheck int) (p *GenericListeners) {
+func NewGenerListeners(ctx context.Context, connectionTable *Table, connectionTrack *Tracking, secretKey string, certbundle tls.Certificate,
+	wssHostName string, adminHostName string, cancelCheck int, lbDefaultMethod string) (p *GenericListeners) {
 	p = new(GenericListeners)
 	p.listeners = make(map[*net.Listener]int)
 	p.ctx = ctx
@@ -71,6 +73,7 @@ func NewGenerListeners(ctx context.Context, connectionTable *Table, connectionTr
 	p.wssHostName = wssHostName
 	p.adminHostName = adminHostName
 	p.cancelCheck = cancelCheck
+	p.lbDefaultMethod = lbDefaultMethod
 	return
 }
 
@@ -93,6 +96,7 @@ func (gl *GenericListeners) Run(ctx context.Context, initialPort int) {
 	ctx = context.WithValue(ctx, ctxWssHostName, gl.wssHostName)
 	ctx = context.WithValue(ctx, ctxAdminHostName, gl.adminHostName)
 	ctx = context.WithValue(ctx, ctxCancelCheck, gl.cancelCheck)
+	ctx = context.WithValue(ctx, ctxLoadbalanceDefaultMethod, gl.lbDefaultMethod)
 
 	go func(ctx context.Context) {
 		for {