From efd0477c2bb3e76c2b2e43037fd54da822dd1b25 Mon Sep 17 00:00:00 2001
From: Frank Denis <github@pureftpd.org>
Date: Wed, 10 Jan 2018 09:46:27 +0100
Subject: [PATCH] Implement an actual estimator for the response size

Scale back the minimum question size when relevant.

Did I mention that this is yet another thing that was never properly
implemented in dnscrypt-proxy 1.x?
---
 dnscrypt-proxy/estimators.go | 24 ++++++++++++++++++++++--
 dnscrypt-proxy/main.go       |  2 ++
 2 files changed, 24 insertions(+), 2 deletions(-)

diff --git a/dnscrypt-proxy/estimators.go b/dnscrypt-proxy/estimators.go
index 0c591168..ddbe347e 100644
--- a/dnscrypt-proxy/estimators.go
+++ b/dnscrypt-proxy/estimators.go
@@ -1,14 +1,23 @@
 package main
 
-import "sync"
+import (
+	"sync"
+
+	"github.com/VividCortex/ewma"
+)
+
+const (
+	EwmaDecay = 100.0
+)
 
 type QuestionSizeEstimator struct {
 	sync.RWMutex
 	minQuestionSize int
+	ewma            ewma.MovingAverage
 }
 
 func NewQuestionSizeEstimator() QuestionSizeEstimator {
-	return QuestionSizeEstimator{minQuestionSize: InitialMinQuestionSize}
+	return QuestionSizeEstimator{minQuestionSize: InitialMinQuestionSize, ewma: ewma.NewMovingAverage(EwmaDecay)}
 }
 
 func (questionSizeEstimator *QuestionSizeEstimator) MinQuestionSize() int {
@@ -25,5 +34,16 @@ func (questionSizeEstimator *QuestionSizeEstimator) blindAdjust() {
 	} else {
 		questionSizeEstimator.minQuestionSize *= 2
 	}
+	questionSizeEstimator.ewma.Set(float64(questionSizeEstimator.minQuestionSize))
+	questionSizeEstimator.Unlock()
+}
+
+func (questionSizeEstimator *QuestionSizeEstimator) adjust(packetSize int) {
+	questionSizeEstimator.Lock()
+	questionSizeEstimator.ewma.Add(float64(packetSize))
+	ma, minQuestionSize := int(questionSizeEstimator.ewma.Value()), questionSizeEstimator.minQuestionSize
+	if ma > InitialMinQuestionSize && ma < minQuestionSize/2 {
+		questionSizeEstimator.minQuestionSize = Max(InitialMinQuestionSize, minQuestionSize/2)
+	}
 	questionSizeEstimator.Unlock()
 }
diff --git a/dnscrypt-proxy/main.go b/dnscrypt-proxy/main.go
index a81f8f13..63bd2c41 100644
--- a/dnscrypt-proxy/main.go
+++ b/dnscrypt-proxy/main.go
@@ -163,6 +163,8 @@ func (proxy *Proxy) processIncomingQuery(serverInfo *ServerInfo, serverProto str
 		clientPc.(net.PacketConn).WriteTo(response, *clientAddr)
 		if HasTCFlag(response) {
 			proxy.questionSizeEstimator.blindAdjust()
+		} else {
+			proxy.questionSizeEstimator.adjust(len(response))
 		}
 	} else {
 		response, err = PrefixWithSize(response)