middleware/erratic: add delaying queries (#614)

* middleware/erratic: add delying queries

* Dont println

5 files changed, 131 insertions, 21 deletions

diff --git a/middleware/erratic/README.md b/middleware/erratic/README.md
index c255e3150..fe5aa5338 100644
--- a/middleware/erratic/README.md
+++ b/middleware/erratic/README.md
@@ -4,10 +4,6 @@
 queries, but the responses can be delayed by a random amount of time or dropped all together, i.e.
 no answer at all.
 
-~~~ txt
-._<transport>.qname. 0 IN SRV 0 0 <port> .
-~~~
-
 The *erratic* middleware will respond to every A or AAAA query. For any other type it will return
 a SERVFAIL response. The reply for A will return 192.0.2.53 (see RFC 5737), for AAAA it returns
 2001:DB8::53 (see RFC 3849).
@@ -16,11 +12,14 @@ a SERVFAIL response. The reply for A will return 192.0.2.53 (see RFC 5737), for
 
 ~~~ txt
 erratic {
-    drop AMOUNT
+    drop [AMOUNT]
+    delay [AMOUNT [DURATION]]
 }
 ~~~
 
-* **AMOUNT** drop 1 per **AMOUNT** of the queries, the default is 2.
+* `drop`: drop 1 per **AMOUNT** of the queries, the default is 2.
+* `delay`: delay 1 per **AMOUNT** of queries for **DURATION**, the default for **AMOUNT** is 2 and
+  the default for **DURATION** is 100ms.
 
 ## Examples
 
@@ -32,7 +31,7 @@ erratic {
 }
 ~~~
 
-Or even shorter if the defaults suits you:
+Or even shorter if the defaults suits you. Note this only drops queries, it does not delay them.
 
 ~~~ txt
 . {
@@ -40,6 +39,22 @@ Or even shorter if the defaults suits you:
 }
 ~~~
 
-## Bugs
+Delay 1 in 3 queries for 50ms, but also drop 1 in 2.
+
+~~~ txt
+. {
+    erratic {
+        delay 3 50ms
+    }
+}
+~~~
 
-Delaying answers is not implemented.
+To stop dropping you'll need to explicitally set that to 0:
+~~~ txt
+. {
+    erratic {
+        delay 3 50ms
+        drop 0
+    }
+}
+~~~
diff --git a/middleware/erratic/erratic.go b/middleware/erratic/erratic.go
index a3fbb2e8e..1f663a558 100644
--- a/middleware/erratic/erratic.go
+++ b/middleware/erratic/erratic.go
@@ -3,6 +3,7 @@ package erratic
 
 import (
 	"sync/atomic"
+	"time"
 
 	"github.com/coredns/coredns/request"
 
@@ -12,7 +13,10 @@ import (
 
 // Erratic is a middleware that returns erratic repsonses to each client.
 type Erratic struct {
-	amount uint64
+	drop uint64
+
+	delay    uint64
+	duration time.Duration
 
 	q uint64 // counter of queries
 }
@@ -20,16 +24,21 @@ type Erratic struct {
 // ServeDNS implements the middleware.Handler interface.
 func (e *Erratic) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg) (int, error) {
 	state := request.Request{W: w, Req: r}
-
 	drop := false
-	if e.amount > 0 {
-		queryNr := atomic.LoadUint64(&e.q)
+	delay := false
 
-		if queryNr%e.amount == 0 {
+	queryNr := atomic.LoadUint64(&e.q)
+	atomic.AddUint64(&e.q, 1)
+
+	if e.drop > 0 {
+		if queryNr%e.drop == 0 {
 			drop = true
 		}
-
-		atomic.AddUint64(&e.q, 1)
+	}
+	if e.delay > 0 {
+		if queryNr%e.delay == 0 {
+			delay = true
+		}
 	}
 
 	m := new(dns.Msg)
@@ -50,6 +59,9 @@ func (e *Erratic) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 		m.Answer = append(m.Answer, &rr)
 	default:
 		if !drop {
+			if delay {
+				time.Sleep(e.duration)
+			}
 			// coredns will return error.
 			return dns.RcodeServerFailure, nil
 		}
@@ -59,6 +71,10 @@ func (e *Erratic) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 		return 0, nil
 	}
 
+	if delay {
+		time.Sleep(e.duration)
+	}
+
 	state.SizeAndDo(m)
 	w.WriteMsg(m)
 
diff --git a/middleware/erratic/erratic_test.go b/middleware/erratic/erratic_test.go
index e63545b35..c99135bbd 100644
--- a/middleware/erratic/erratic_test.go
+++ b/middleware/erratic/erratic_test.go
@@ -11,7 +11,7 @@ import (
 )
 
 func TestErraticDrop(t *testing.T) {
-	e := &Erratic{amount: 2} // 50% drops
+	e := &Erratic{drop: 2} // 50% drops
 
 	tests := []struct {
 		expectedCode int
diff --git a/middleware/erratic/setup.go b/middleware/erratic/setup.go
index ac40d24fb..a642b6ff3 100644
--- a/middleware/erratic/setup.go
+++ b/middleware/erratic/setup.go
@@ -3,6 +3,7 @@ package erratic
 import (
 	"fmt"
 	"strconv"
+	"time"
 
 	"github.com/coredns/coredns/core/dnsserver"
 	"github.com/coredns/coredns/middleware"
@@ -31,7 +32,7 @@ func setupErratic(c *caddy.Controller) error {
 }
 
 func parseErratic(c *caddy.Controller) (*Erratic, error) {
-	e := &Erratic{amount: 2}
+	e := &Erratic{drop: 2}
 	for c.Next() { // 'erratic'
 		for c.NextBlock() {
 			switch c.Val() {
@@ -42,8 +43,30 @@ func parseErratic(c *caddy.Controller) (*Erratic, error) {
 				}
 
 				if len(args) == 0 {
-					return nil, nil
+					continue
 				}
+
+				amount, err := strconv.ParseInt(args[0], 10, 32)
+				if err != nil {
+					return nil, err
+				}
+				if amount < 0 {
+					return nil, fmt.Errorf("illegal amount value given %q", args[0])
+				}
+				e.drop = uint64(amount)
+			case "delay":
+				args := c.RemainingArgs()
+				if len(args) > 2 {
+					return nil, c.ArgErr()
+				}
+
+				// Defaults.
+				e.delay = 2
+				e.duration = time.Duration(100 * time.Millisecond)
+				if len(args) == 0 {
+					continue
+				}
+
 				amount, err := strconv.ParseInt(args[0], 10, 32)
 				if err != nil {
 					return nil, err
@@ -51,7 +74,15 @@ func parseErratic(c *caddy.Controller) (*Erratic, error) {
 				if amount < 0 {
 					return nil, fmt.Errorf("illegal amount value given %q", args[0])
 				}
-				e.amount = uint64(amount)
+				e.delay = uint64(amount)
+
+				if len(args) > 1 {
+					duration, err := time.ParseDuration(args[1])
+					if err != nil {
+						return nil, err
+					}
+					e.duration = duration
+				}
 			}
 		}
 	}
diff --git a/middleware/erratic/setup_test.go b/middleware/erratic/setup_test.go
index 99b545f5e..218973457 100644
--- a/middleware/erratic/setup_test.go
+++ b/middleware/erratic/setup_test.go
@@ -6,7 +6,7 @@ import (
 	"github.com/mholt/caddy"
 )
 
-func TestSetupWhoami(t *testing.T) {
+func TestSetupErratic(t *testing.T) {
 	c := caddy.NewTestController("dns", `erratic {
 		drop
 	}`)
@@ -26,3 +26,51 @@ func TestSetupWhoami(t *testing.T) {
 		t.Fatalf("Test 4, expected errors, but got: %q", err)
 	}
 }
+
+func TestParseErratic(t *testing.T) {
+	tests := []struct {
+		input     string
+		shouldErr bool
+		drop      uint64
+		delay     uint64
+	}{
+		// oks
+		{`erratic`, false, 2, 0},
+		{`erratic {
+			drop 2
+			delay 3 1ms
+
+		}`, false, 2, 3},
+		// fails
+		{`erratic {
+			drop -1
+		}`, true, 0, 0},
+		{`erraric {
+			drop 3
+			delay 3 bla
+		}`, true, 0, 0},
+	}
+	for i, test := range tests {
+		c := caddy.NewTestController("dns", test.input)
+		e, err := parseErratic(c)
+		if test.shouldErr && err == nil {
+			t.Errorf("Test %v: Expected error but found nil", i)
+			continue
+		} else if !test.shouldErr && err != nil {
+			t.Errorf("Test %v: Expected no error but found error: %v", i, err)
+			continue
+		}
+
+		if test.shouldErr {
+			continue
+		}
+
+		if test.delay != e.delay {
+			t.Errorf("Test %v: Expected delay %d but found: %d", i, test.delay, e.delay)
+		}
+
+		if test.drop != e.drop {
+			t.Errorf("Test %v: Expected drop %d but found: %d", i, test.drop, e.drop)
+		}
+	}
+}