1 files changed, 18 insertions, 7 deletions

diff --git a/plugin/loadbalance/loadbalance.go b/plugin/loadbalance/loadbalance.go
index 966121d6b..f2a1caed0 100644
--- a/plugin/loadbalance/loadbalance.go
+++ b/plugin/loadbalance/loadbalance.go
@@ -5,11 +5,19 @@ import (
 	"github.com/miekg/dns"
 )
 
-// RoundRobinResponseWriter is a response writer that shuffles A, AAAA and MX records.
-type RoundRobinResponseWriter struct{ dns.ResponseWriter }
+const (
+	ramdomShufflePolicy      = "round_robin"
+	weightedRoundRobinPolicy = "weighted"
+)
+
+// LoadBalanceResponseWriter is a response writer that shuffles A, AAAA and MX records.
+type LoadBalanceResponseWriter struct {
+	dns.ResponseWriter
+	shuffle func(*dns.Msg) *dns.Msg
+}
 
 // WriteMsg implements the dns.ResponseWriter interface.
-func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
+func (r *LoadBalanceResponseWriter) WriteMsg(res *dns.Msg) error {
 	if res.Rcode != dns.RcodeSuccess {
 		return r.ResponseWriter.WriteMsg(res)
 	}
@@ -18,11 +26,14 @@ func (r *RoundRobinResponseWriter) WriteMsg(res *dns.Msg) error {
 		return r.ResponseWriter.WriteMsg(res)
 	}
 
+	return r.ResponseWriter.WriteMsg(r.shuffle(res))
+}
+
+func randomShuffle(res *dns.Msg) *dns.Msg {
 	res.Answer = roundRobin(res.Answer)
 	res.Ns = roundRobin(res.Ns)
 	res.Extra = roundRobin(res.Extra)
-
-	return r.ResponseWriter.WriteMsg(res)
+	return res
 }
 
 func roundRobin(in []dns.RR) []dns.RR {
@@ -72,9 +83,9 @@ func roundRobinShuffle(records []dns.RR) {
 }
 
 // Write implements the dns.ResponseWriter interface.
-func (r *RoundRobinResponseWriter) Write(buf []byte) (int, error) {
+func (r *LoadBalanceResponseWriter) Write(buf []byte) (int, error) {
 	// Should we pack and unpack here to fiddle with the packet... Not likely.
-	log.Warning("RoundRobin called with Write: not shuffling records")
+	log.Warning("LoadBalance called with Write: not shuffling records")
 	n, err := r.ResponseWriter.Write(buf)
 	return n, err
 }