From 074d176f039940d408bde2831c442500433f3ac4 Mon Sep 17 00:00:00 2001 From: Miek Gieben Date: Wed, 21 Feb 2018 21:13:41 +0000 Subject: [PATCH] Scrub: use binary search (#1543) Use binary search to find the minimal message size, that contains whole RRs and fits the client's buffer. This is better then just setting entire sections to `nil`. Extend the tests to test for additional and answer section truncation. In the first case we *don't* set the TC bit. This function now also set Compression to true. --- request/request.go | 88 ++++++++++++++++++++++++++++++++--------- request/request_test.go | 40 ++++++++++++++----- 2 files changed, 101 insertions(+), 27 deletions(-) diff --git a/request/request.go b/request/request.go index e2ff2008c..1105f919a 100644 --- a/request/request.go +++ b/request/request.go @@ -179,34 +179,86 @@ type Result int const ( // ScrubIgnored is returned when Scrub did nothing to the message. ScrubIgnored Result = iota - // ScrubDone is returned when the reply has been scrubbed. - ScrubDone + // ScrubExtra is returned when the reply has been scrubbed by removing RRs from the additional section. + ScrubExtra + // ScrubAnswer is returned when the reply has been scrubbed by removing RRs from the answer section. + ScrubAnswer ) -// Scrub scrubs the reply message so that it will fit the client's buffer. If -// even after dropping the additional section it does not fit, the answer will -// be cleared and the TC bit will be set on the message. Note, the TC bit will -// be set regardless of protocol, even TCP message will get the bit, the client -// should then retry with pigeons. TODO(referral). +// Scrub scrubs the reply message so that it will fit the client's buffer. It sets +// reply.Compress to true. +// Scrub uses binary search to find a save cut off point in the additional section. +// If even *without* the additional section the reply still doesn't fit we +// repeat this process for the answer section. If we scrub the answer section +// we set the TC bit on the reply; indicating the client should retry over TCP. +// Note, the TC bit will be set regardless of protocol, even TCP message will +// get the bit, the client should then retry with pigeons. func (r *Request) Scrub(reply *dns.Msg) (*dns.Msg, Result) { + reply.Compress = true + size := r.Size() - l := reply.Len() - if size >= l { + rl := reply.Len() + + if size >= rl { return reply, ScrubIgnored } - // TODO(miek): check for delegation - // If not delegation, drop additional section. - reply.Extra = nil - r.SizeAndDo(reply) - l = reply.Len() - if size >= l { - return reply, ScrubDone + origExtra := reply.Extra + re := len(reply.Extra) + l, m := 0, 0 + for l < re { + m = (l + re) / 2 + reply.Extra = origExtra[:m] + rl = reply.Len() + if rl < size { + l = m + 1 + continue + } + if rl > size { + re = m - 1 + continue + } + } + // We may come out of this loop with one rotation too many as we don't break on rl == size. + // I.e. m makes it too large, but m-1 works. + if rl > size && m > 0 { + reply.Extra = origExtra[:m-1] + rl = reply.Len() } + if rl < size { + r.SizeAndDo(reply) + return reply, ScrubExtra + } + + origAnswer := reply.Answer + ra := len(reply.Answer) + l, m = 0, 0 + for l < ra { + m = (l + ra) / 2 + reply.Answer = origAnswer[:m] + rl = reply.Len() + if rl < size { + l = m + 1 + continue + } + if rl > size { + ra = m - 1 + continue + } + } + // We may come out of this loop with one rotation too many as we don't break on rl == size. + // I.e. m makes it too large, but m-1 works. + if rl > size && m > 0 { + reply.Answer = origAnswer[:m-1] + // No need to recalc length, as we don't use it. We set truncated anyway. Doing + // this extra m-1 step does make it fit in the client's buffer however. + } + + // It now fits, but Truncated. + r.SizeAndDo(reply) reply.Truncated = true - reply.Answer = nil - return reply, ScrubDone + return reply, ScrubAnswer } // Type returns the type of the question as a string. If the request is malformed diff --git a/request/request_test.go b/request/request_test.go index 49b825627..e6fcf6935 100644 --- a/request/request_test.go +++ b/request/request_test.go @@ -61,7 +61,7 @@ func TestRequestMalformed(t *testing.T) { } } -func TestRequestScrub(t *testing.T) { +func TestRequestScrubAnswer(t *testing.T) { m := new(dns.Msg) m.SetQuestion("large.example.com.", dns.TypeSRV) req := Request{W: &test.ResponseWriter{}, Req: m} @@ -69,24 +69,46 @@ func TestRequestScrub(t *testing.T) { reply := new(dns.Msg) reply.SetReply(m) for i := 1; i < 200; i++ { - reply.Answer = append(reply.Answer, test.SRV(fmt.Sprintf( - "large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", - i, - ))) + reply.Answer = append(reply.Answer, test.SRV( + fmt.Sprintf("large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", i))) } - msg, got := req.Scrub(reply) - if want := ScrubDone; want != got { + _, got := req.Scrub(reply) + if want := ScrubAnswer; want != got { t.Errorf("want scrub result %d, got %d", want, got) } - if want, got := req.Size(), msg.Len(); want < got { + if want, got := req.Size(), reply.Len(); want < got { t.Errorf("want scrub to reduce message length below %d bytes, got %d bytes", want, got) } - if !msg.Truncated { + if !reply.Truncated { t.Errorf("want scrub to set truncated bit") } } +func TestRequestScrubExtra(t *testing.T) { + m := new(dns.Msg) + m.SetQuestion("large.example.com.", dns.TypeSRV) + req := Request{W: &test.ResponseWriter{}, Req: m} + + reply := new(dns.Msg) + reply.SetReply(m) + for i := 1; i < 200; i++ { + reply.Extra = append(reply.Extra, test.SRV( + fmt.Sprintf("large.example.com. 10 IN SRV 0 0 80 10-0-0-%d.default.pod.k8s.example.com.", i))) + } + + _, got := req.Scrub(reply) + if want := ScrubExtra; want != got { + t.Errorf("want scrub result %d, got %d", want, got) + } + if want, got := req.Size(), reply.Len(); want < got { + t.Errorf("want scrub to reduce message length below %d bytes, got %d bytes", want, got) + } + if reply.Truncated { + t.Errorf("want scrub to not set truncated bit") + } +} + func BenchmarkRequestDo(b *testing.B) { st := testRequest()