Tc bits (#617)

* middleware/erratic: allow TC bit to be set

Add `truncate` as an option.

Fixes #593

5 files changed, 110 insertions, 21 deletions

diff --git a/middleware/erratic/README.md b/middleware/erratic/README.md
index fe5aa5338..5407c9a7b 100644
--- a/middleware/erratic/README.md
+++ b/middleware/erratic/README.md
@@ -1,8 +1,11 @@
 # erratic
 
 *erratic* is a middleware useful for testing client behavior. It returns a static response to all
-queries, but the responses can be delayed by a random amount of time or dropped all together, i.e.
-no answer at all.
+queries, but the responses can be:
+
+* delayed by some duration
+* dropped all together
+* the truncated bit can be set
 
 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
@@ -13,11 +16,13 @@ a SERVFAIL response. The reply for A will return 192.0.2.53 (see RFC 5737), for
 ~~~ txt
 erratic {
     drop [AMOUNT]
+    truncate [AMOUNT]
     delay [AMOUNT [DURATION]]
 }
 ~~~
 
-* `drop`: drop 1 per **AMOUNT** of the queries, the default is 2.
+* `drop`: drop 1 per **AMOUNT** of queries, the default is 2.
+* `truncate`: truncate 1 per **AMOUNT** of 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.
 
@@ -39,7 +44,7 @@ Or even shorter if the defaults suits you. Note this only drops queries, it does
 }
 ~~~
 
-Delay 1 in 3 queries for 50ms, but also drop 1 in 2.
+Delay 1 in 3 queries for 50ms
 
 ~~~ txt
 . {
@@ -49,12 +54,24 @@ Delay 1 in 3 queries for 50ms, but also drop 1 in 2.
 }
 ~~~
 
-To stop dropping you'll need to explicitally set that to 0:
+Delay 1 in 3 and truncate 1 in 5.
+
 ~~~ txt
 . {
     erratic {
-        delay 3 50ms
-        drop 0
+        delay 3 5ms
+        truncate 5
+    }
+}
+~~~
+
+Drop every second query.
+
+~~~ txt
+. {
+    erratic {
+        drop 2
+        truncate 2
     }
 }
 ~~~
diff --git a/middleware/erratic/erratic.go b/middleware/erratic/erratic.go
index 1f663a558..b05e45f03 100644
--- a/middleware/erratic/erratic.go
+++ b/middleware/erratic/erratic.go
@@ -18,6 +18,8 @@ type Erratic struct {
 	delay    uint64
 	duration time.Duration
 
+	truncate uint64
+
 	q uint64 // counter of queries
 }
 
@@ -26,25 +28,28 @@ func (e *Erratic) ServeDNS(ctx context.Context, w dns.ResponseWriter, r *dns.Msg
 	state := request.Request{W: w, Req: r}
 	drop := false
 	delay := false
+	trunc := false
 
 	queryNr := atomic.LoadUint64(&e.q)
 	atomic.AddUint64(&e.q, 1)
 
-	if e.drop > 0 {
-		if queryNr%e.drop == 0 {
-			drop = true
-		}
+	if e.drop > 0 && queryNr%e.drop == 0 {
+		drop = true
 	}
-	if e.delay > 0 {
-		if queryNr%e.delay == 0 {
-			delay = true
-		}
+	if e.delay > 0 && queryNr%e.delay == 0 {
+		delay = true
+	}
+	if e.truncate > 0 && queryNr&e.truncate == 0 {
+		trunc = true
 	}
 
 	m := new(dns.Msg)
 	m.SetReply(r)
 	m.Compress = true
 	m.Authoritative = true
+	if trunc {
+		m.Truncated = true
+	}
 
 	// small dance to copy rrA or rrAAAA into a non-pointer var that allows us to overwrite the ownername
 	// in a non-racy way.
diff --git a/middleware/erratic/erratic_test.go b/middleware/erratic/erratic_test.go
index c99135bbd..4b54e0c12 100644
--- a/middleware/erratic/erratic_test.go
+++ b/middleware/erratic/erratic_test.go
@@ -39,7 +39,41 @@ func TestErraticDrop(t *testing.T) {
 		}
 
 		if tc.drop && rec.Msg != nil {
-			t.Errorf("Test %d: Expected dropped packet, but got %q", i, rec.Msg.Question[0].Name)
+			t.Errorf("Test %d: Expected dropped message, but got %q", i, rec.Msg.Question[0].Name)
+		}
+	}
+}
+
+func TestErraticTruncate(t *testing.T) {
+	e := &Erratic{truncate: 2} // 50% drops
+
+	tests := []struct {
+		expectedCode int
+		expectedErr  error
+		truncate     bool
+	}{
+		{expectedCode: dns.RcodeSuccess, expectedErr: nil, truncate: true},
+		{expectedCode: dns.RcodeSuccess, expectedErr: nil, truncate: false},
+	}
+
+	ctx := context.TODO()
+
+	for i, tc := range tests {
+		req := new(dns.Msg)
+		req.SetQuestion("example.org.", dns.TypeA)
+
+		rec := dnsrecorder.New(&test.ResponseWriter{})
+		code, err := e.ServeDNS(ctx, rec, req)
+
+		if err != tc.expectedErr {
+			t.Errorf("Test %d: Expected error %q, but got %q", i, tc.expectedErr, err)
+		}
+		if code != int(tc.expectedCode) {
+			t.Errorf("Test %d: Expected status code %d, but got %d", i, tc.expectedCode, code)
+		}
+
+		if tc.truncate && !rec.Msg.Truncated {
+			t.Errorf("Test %d: Expected truncated message, but got %q", i, rec.Msg.Question[0].Name)
 		}
 	}
 }
diff --git a/middleware/erratic/setup.go b/middleware/erratic/setup.go
index a642b6ff3..714fc309b 100644
--- a/middleware/erratic/setup.go
+++ b/middleware/erratic/setup.go
@@ -33,6 +33,8 @@ func setupErratic(c *caddy.Controller) error {
 
 func parseErratic(c *caddy.Controller) (*Erratic, error) {
 	e := &Erratic{drop: 2}
+	drop := false // true if we've seen the drop keyword
+
 	for c.Next() { // 'erratic'
 		for c.NextBlock() {
 			switch c.Val() {
@@ -54,6 +56,7 @@ func parseErratic(c *caddy.Controller) (*Erratic, error) {
 					return nil, fmt.Errorf("illegal amount value given %q", args[0])
 				}
 				e.drop = uint64(amount)
+				drop = true
 			case "delay":
 				args := c.RemainingArgs()
 				if len(args) > 2 {
@@ -83,8 +86,30 @@ func parseErratic(c *caddy.Controller) (*Erratic, error) {
 					}
 					e.duration = duration
 				}
+			case "truncate":
+				args := c.RemainingArgs()
+				if len(args) > 1 {
+					return nil, c.ArgErr()
+				}
+
+				if len(args) == 0 {
+					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.truncate = uint64(amount)
 			}
 		}
 	}
+	if (e.delay > 0 || e.truncate > 0) && !drop { // delay is set, but we've haven't seen a drop keyword, remove default drop stuff
+		e.drop = 0
+	}
+
 	return e, nil
 }
diff --git a/middleware/erratic/setup_test.go b/middleware/erratic/setup_test.go
index 218973457..cd4a641d1 100644
--- a/middleware/erratic/setup_test.go
+++ b/middleware/erratic/setup_test.go
@@ -33,22 +33,28 @@ func TestParseErratic(t *testing.T) {
 		shouldErr bool
 		drop      uint64
 		delay     uint64
+		truncate  uint64
 	}{
 		// oks
-		{`erratic`, false, 2, 0},
+		{`erratic`, false, 2, 0, 0},
 		{`erratic {
 			drop 2
 			delay 3 1ms
 
-		}`, false, 2, 3},
+		}`, false, 2, 3, 0},
+		{`erratic {
+			truncate 2
+			delay 3 1ms
+
+		}`, false, 0, 3, 2},
 		// fails
 		{`erratic {
 			drop -1
-		}`, true, 0, 0},
+		}`, true, 0, 0, 0},
 		{`erraric {
 			drop 3
 			delay 3 bla
-		}`, true, 0, 0},
+		}`, true, 0, 0, 0},
 	}
 	for i, test := range tests {
 		c := caddy.NewTestController("dns", test.input)
@@ -68,9 +74,11 @@ func TestParseErratic(t *testing.T) {
 		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)
 		}
+		if test.truncate != e.truncate {
+			t.Errorf("Test %v: Expected truncate %d but found: %d", i, test.truncate, e.truncate)
+		}
 	}
 }