commit 894d4e10b62399a9916f7bf6ade947d06f185ddc · seiso.moe/alky

go.mod

···

       7
       7
        
       require (

     

       8
       8
        
       	github.com/BurntSushi/toml v1.5.0

     

       9
       9
        
       	github.com/ClickHouse/clickhouse-go/v2 v2.34.0

     

       10
       10
       +
       	github.com/stretchr/testify v1.10.0

     

       11
       11
       +
       	golang.org/x/time v0.11.0

     

       10
       12
        
       	tangled.sh/seiso.moe/magna v0.0.1

     

       11
       13
        
       )

     

       12
       14
        
       

     

       13
       15
        
       require (

     

       14
       16
        
       	github.com/ClickHouse/ch-go v0.65.1 // indirect

     

       15
       17
        
       	github.com/andybalholm/brotli v1.1.1 // indirect

     

       18
       18
       +
       	github.com/davecgh/go-spew v1.1.1 // indirect

     

       16
       19
        
       	github.com/go-faster/city v1.0.1 // indirect

     

       17
       20
        
       	github.com/go-faster/errors v0.7.1 // indirect

     

       18
       21
        
       	github.com/google/uuid v1.6.0 // indirect

     

       19
       22
        
       	github.com/klauspost/compress v1.18.0 // indirect

     

       20
       23
        
       	github.com/paulmach/orb v0.11.1 // indirect

     

       21
       24
        
       	github.com/pierrec/lz4/v4 v4.1.22 // indirect

     

       25
       25
       +
       	github.com/pmezard/go-difflib v1.0.0 // indirect

     

       22
       26
        
       	github.com/segmentio/asm v1.2.0 // indirect

     

       23
       27
        
       	github.com/shopspring/decimal v1.4.0 // indirect

     

       24
       28
        
       	go.opentelemetry.io/otel v1.35.0 // indirect

go.sum

···

       92
       92
        
       golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=

     

       93
       93
        
       golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=

     

       94
       94
        
       golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=

     

       95
       95
       +
       golang.org/x/time v0.11.0 h1:/bpjEDfN9tkoN/ryeYHnv5hcMlc8ncjMcM4XBk5NWV0=

     

       96
       96
       +
       golang.org/x/time v0.11.0/go.mod h1:CDIdPxbZBQxdj6cxyCIdrNogrJKMJ7pr37NYpMcMDSg=

     

       95
       97
        
       golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

     

       96
       98
        
       golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo=

     

       97
       99
        
       golang.org/x/tools v0.0.0-20200619180055-7c47624df98f/go.mod h1:EkVYQZoAsY45+roYkvgYkIh4xh/qjgUK9TdY2XT94GE=

+1 -1

main.go

···

       26
       26
        
       func main() {

     

       27
       27
        
       	cfg, err := config.LoadConfig(configFlag)

     

       28
       28
        
       	if err != nil {

     

       29
       29
       -
       		log.Fatalf("failed to load config: ", err)

     

       29
       29
       +
       		log.Fatalf("failed to load config: %v", err)

     

       30
       30
        
       	}

     

       31
       31
        
       

     

       32
       32
        
       	logger := setupLogger(&cfg)

+3 -8

pkg/dns/cache.go

···

       149
       149
        
       		Additional: make([]magna.ResourceRecord, len(original.Additional)),

     

       150
       150
        
       	}

     

       151
       151
        
       

     

       152
       152
       -
       	remainingDuration := original.ExpireAt.Sub(now)

     

       153
       153
       -
       	if remainingDuration < 0 {

     

       154
       154
       -
       		remainingDuration = 0

     

       155
       155
       -
       	}

     

       152
       152
       +
       	remainingDuration := max(original.ExpireAt.Sub(now), 0)

     

       156
       153
        
       	remainingTTL := uint32(remainingDuration.Seconds())

     

       157
       154
        
       

     

       158
       155
        
       	copyAndAdjust := func(dest *[]magna.ResourceRecord, src []magna.ResourceRecord) {

     

       159
       156
        
       		for i, rr := range src {

     

       160
       157
        
       			(*dest)[i] = rr

     

       161
       158
        
       			rrExpireAt := original.CacheTime.Add(time.Duration(rr.TTL) * time.Second)

     

       162
       162
       -
       			rrRemainingDuration := rrExpireAt.Sub(now)

     

       163
       163
       -
       			if rrRemainingDuration < 0 {

     

       164
       164
       -
       				rrRemainingDuration = 0

     

       165
       165
       -
       			}

     

       159
       159
       +
       			rrRemainingDuration := max(rrExpireAt.Sub(now), 0)

     

       166
       160
        
       			rrRemainingTTL := uint32(rrRemainingDuration.Seconds())

     

       167
       161
        
       

     

       168
       162
        
       			finalTTL := min(rrRemainingTTL, remainingTTL)

     
···

       178
       172
        
       }

     

       179
       173
        
       

     

       180
       174
        
       func (c *LRUCache) Set(key string, entry *CacheEntry) {

     

       175
       175
       +
       	c.logger.Info("setting key", "key", key)

     

       181
       176
        
       	if entry == nil {

     

       182
       177
        
       		c.logger.Warn("attempted to set nil entry in cache", "key", key)

     

       183
       178
        
       		return

+57 -38

pkg/dns/ratelimit.go

···

       5
       5
        
       	"sync"

     

       6
       6
        
       	"time"

     

       7
       7
        
       

     

       8
       8
       +
       	"golang.org/x/time/rate"

     

       8
       9
        
       	"tangled.sh/seiso.moe/magna"

     

       9
       10
        
       )

     

       10
       11
        
       

     
···

       17
       18
        
       

     

       18
       19
        
       type rateLimiter struct {

     

       19
       20
        
       	config RateLimitConfig

     

       20
       20
       -
       	ipData map[string]*ipRateData

     

       21
       21
       +
       	limiters map[string]*ipRateLimiterEntry

     

       21
       22
        
       	mu     sync.RWMutex

     

       23
       23
       +
       	stopCleanup chan struct{}

     

       22
       24
        
       }

     

       23
       25
        
       

     

       24
       24
       -
       type ipRateData struct {

     

       25
       25
       -
       	time time.Time

     

       26
       26
       +
       type ipRateLimiterEntry struct {

     

       27
       27
       +
       	limiter    *rate.Limiter

     

       28
       28
       +
       	lastAccess time.Time

     

       26
       29
        
       }

     

       27
       30
        
       

     

       28
       31
        
       func NewDefaultRateLimitConfig() *RateLimitConfig {

     

       29
       32
        
       	return &RateLimitConfig{

     

       30
       33
        
       		Rate:           1,

     

       31
       34
        
       		Burst:          1,

     

       32
       32
       -
       		WindowLength:   time.Hour,

     

       33
       33
       -
       		ExpirationTime: time.Hour,

     

       35
       35
       +
       		WindowLength:   time.Second,

     

       36
       36
       +
       		ExpirationTime: 5 * time.Minute,

     

       34
       37
        
       	}

     

       35
       38
        
       }

     

       36
       39
        
       

     

       37
       40
        
       func newRateLimiter(config RateLimitConfig) *rateLimiter {

     

       38
       38
       -
       	return &rateLimiter{

     

       41
       41
       +
       	rl := &rateLimiter{

     

       39
       42
        
       		config: config,

     

       40
       40
       -
       		ipData: make(map[string]*ipRateData),

     

       43
       43
       +
       		limiters: make(map[string]*ipRateLimiterEntry),

     

       44
       44
       +
       		stopCleanup: make(chan struct{}),

     

       41
       45
        
       	}

     

       46
       46
       +
       

     

       47
       47
       +
       	go rl.cleanupLoop()

     

       48
       48
       +
       	return rl

     

       42
       49
        
       }

     

       43
       50
        
       

     

       44
       51
        
       func (rl *rateLimiter) allow(ip string) bool {

     

       45
       45
       -
       	rl.mu.Lock()

     

       52
       52
       +
       	rl.mu.Lock() 

     

       46
       53
        
       	defer rl.mu.Unlock()

     

       47
       54
        
       

     

       55
       55
       +
       	entry, exists := rl.limiters[ip]

     

       48
       56
        
       	now := time.Now()

     

       49
       49
       -
       	cost := time.Duration(float64(time.Second) / rl.config.Rate)

     

       50
       57
        
       

     

       51
       51
       -
       	data, exists := rl.ipData[ip]

     

       52
       58
        
       	if !exists {

     

       53
       53
       -
       		data = &ipRateData{time: now.Add(-rl.config.WindowLength)}

     

       54
       54
       -
       		rl.ipData[ip] = data

     

       55
       55
       -
       	}

     

       59
       59
       +
       		limiter := rate.NewLimiter(rate.Limit(rl.config.Rate), rl.config.Burst)

     

       60
       60
       +
       		entry := &ipRateLimiterEntry{

     

       61
       61
       +
       			limiter: limiter,

     

       62
       62
       +
       			lastAccess: now,

     

       63
       63
       +
       		}

     

       56
       64
        
       

     

       57
       57
       -
       	if data.time.Before(now.Add(-rl.config.WindowLength)) {

     

       58
       58
       -
       		data.time = now.Add(-rl.config.WindowLength)

     

       65
       65
       +
       		rl.limiters[ip] = entry

     

       66
       66
       +
       		return entry.limiter.Allow()

     

       59
       67
        
       	}

     

       60
       68
        
       

     

       61
       61
       -
       	nextTime := data.time.Add(cost)

     

       62
       62
       -
       	if now.Before(nextTime) {

     

       63
       63
       -
       		return false

     

       64
       64
       -
       	}

     

       69
       69
       +
       	entry.lastAccess = now

     

       70
       70
       +
       	return entry.limiter.Allow()

     

       71
       71
       +
       }

     

       72
       72
       +
       

     

       73
       73
       +
       func (rl *rateLimiter) cleanupLoop() {

     

       74
       74
       +
       	ticker := time.NewTicker(rl.config.ExpirationTime)

     

       75
       75
       +
       	defer ticker.Stop()

     

       65
       76
        
       

     

       66
       66
       -
       	if nextTime.Sub(now.Add(-rl.config.WindowLength)) > time.Duration(rl.config.Burst)*cost {

     

       67
       67
       -
       		nextTime = now.Add(cost)

     

       77
       77
       +
       	for {

     

       78
       78
       +
       		select {

     

       79
       79
       +
       		case <-ticker.C:

     

       80
       80
       +
       			rl.performCleanup()

     

       81
       81
       +
       		case <-rl.stopCleanup:

     

       82
       82
       +
       			return

     

       83
       83
       +
       		}

     

       68
       84
        
       	}

     

       69
       69
       -
       

     

       70
       70
       -
       	data.time = nextTime

     

       71
       71
       -
       	return true

     

       72
       85
        
       }

     

       73
       86
        
       

     

       74
       74
       -
       func (rl *rateLimiter) cleanup() {

     

       87
       87
       +
       func (rl *rateLimiter) performCleanup() {

     

       75
       88
        
       	rl.mu.Lock()

     

       76
       89
        
       	defer rl.mu.Unlock()

     

       77
       90
        
       

     

       78
       91
        
       	now := time.Now()

     

       79
       79
       -
       	for ip, data := range rl.ipData {

     

       80
       80
       -
       		if data.time.Before(now.Add(-rl.config.WindowLength)) {

     

       81
       81
       -
       			delete(rl.ipData, ip)

     

       92
       92
       +
       	expirationCutoff := now.Add(-rl.config.ExpirationTime)

     

       93
       93
       +
       

     

       94
       94
       +
       	for ip, entry := range rl.limiters {

     

       95
       95
       +
       		if entry.lastAccess.Before(expirationCutoff) {

     

       96
       96
       +
       			delete(rl.limiters, ip)

     

       82
       97
        
       		}

     

       83
       98
        
       	}

     

       84
       99
        
       }

     

       85
       100
        
       

     

       101
       101
       +
       func (rl *rateLimiter) Stop() {

     

       102
       102
       +
       	close(rl.stopCleanup)

     

       103
       103
       +
       }

     

       104
       104
       +
       

     

       86
       105
        
       func extractIP(addr net.Addr) string {

     

       87
       106
        
       	switch v := addr.(type) {

     

       88
       107
        
       	case *net.UDPAddr:

     
···

       103
       122
        
       		config = NewDefaultRateLimitConfig()

     

       104
       123
        
       	}

     

       105
       124
        
       

     

       106
       106
       -
       	rl := newRateLimiter(*config)

     

       107
       107
       -
       

     

       108
       108
       -
       	go func() {

     

       109
       109
       -
       		ticker := time.NewTicker(config.ExpirationTime)

     

       110
       110
       -
       		for range ticker.C {

     

       111
       111
       -
       			rl.cleanup()

     

       125
       125
       +
       	if config.Rate <= 0 {

     

       126
       126
       +
       		config.Rate = 1

     

       127
       127
       +
       		if config.Burst <= 0 {

     

       128
       128
       +
       			config.Burst = 1

     

       112
       129
        
       		}

     

       113
       113
       -
       	}()

     

       130
       130
       +
       	}

     

       131
       131
       +
       

     

       132
       132
       +
       	rl := newRateLimiter(*config)

     

       114
       133
        
       

     

       115
       134
        
       	return func(next Handler) Handler {

     

       116
       135
        
       		return HandlerFunc(func(w ResponseWriter, r *Request) {

     

       117
       117
       -
       			if !rl.allow(extractIP(r.RemoteAddr)) {

     

       118
       118
       -
       				r.Message.Header.RA = true

     

       136
       136
       +
       			clientIP := extractIP(r.RemoteAddr)

     

       137
       137
       +
       			if !rl.allow(clientIP) {

     

       119
       138
        
       				msg := r.Message.CreateReply(r.Message)

     

       120
       139
        
       				msg = r.Message.SetRCode(magna.REFUSED)

     

       140
       140
       +
       				msg.Header.RA = true

     

       121
       141
        
       

     

       122
       122
       -
       				// XXX: dont support edns yet and these get copied over on responses

     

       123
       142
        
       				msg.Header.ANCount = 0

     

       124
       143
        
       				msg.Header.NSCount = 0

     

       125
       144
        
       				msg.Header.ARCount = 0

+74 -21

pkg/dns/resolve.go

···

       6
       6
        
       	"log/slog"

     

       7
       7
        
       	"math/rand"

     

       8
       8
        
       	"net"

     

       9
       9
       +
       	"strings"

     

       9
       10
        
       	"time"

     

       10
       11
        
       

     

       11
       12
        
       	"tangled.sh/seiso.moe/magna"

     
···

       97
       98
        
       }

     

       98
       99
        
       

     

       99
       100
        
       func (h *QueryHandler) resolveQuestion(ctx context.Context, question magna.Question, servers []string) ([]magna.ResourceRecord, error) {

     

       101
       101
       +
       	depth := 0

     

       102
       102
       +
       	if d, ok := ctx.Value("recursion_depth").(int); ok {

     

       103
       103
       +
       		depth = d

     

       104
       104
       +
       	}

     

       105
       105
       +
       

     

       106
       106
       +
       	// TODO move to configuration file

     

       107
       107
       +
       	if depth > 15 {

     

       108
       108
       +
       		h.Logger.Warn("Max recursion depth exceeded", "question", question.QName, "depth", depth)

     

       109
       109
       +
       		return nil, fmt.Errorf("max recursion depth exceeded")

     

       110
       110
       +
       	}

     

       111
       111
       +
       	ctx = context.WithValue(ctx, "recursion_depth", depth+1)

     

       112
       112
       +
       

     

       100
       113
        
       	if h.Cache != nil {

     

       101
       114
        
       		cacheKey := GenerateCacheKey(question)

     

       102
       115
        
       		if entry, found := h.Cache.Get(cacheKey); found {

     
···

       116
       129
        
       

     

       117
       130
        
       	resolveCtx, cancel := context.WithTimeout(ctx, h.Timeout)

     

       118
       131
        
       	defer cancel()

     

       132
       132
       +
       

     

       119
       133
        
       	ch := make(chan queryResponse, len(servers))

     

       120
       134
        
       

     

       121
       135
        
       	for _, s := range servers {

     

       122
       136
        
       		go queryServer(resolveCtx, question, s, ch)

     

       123
       137
        
       	}

     

       138
       138
       +
       

     

       139
       139
       +
       	var lastError error

     

       124
       140
        
       

     

       125
       141
        
       	for range servers {

     

       126
       142
        
       		select {

     

       127
       143
        
       		case res := <-ch:

     

       128
       144
        
       			if res.Error != nil {

     

       129
       145
        
       				h.Logger.Warn("error", "question", question.QName, "server", res.Server, "error", res.Error)

     

       146
       146
       +
       				lastError = res.Error

     

       130
       147
        
       				continue

     

       131
       148
        
       			}

     

       132
       149
        
       

     
···

       139
       156
        
       				Authority:  msg.Authority,

     

       140
       157
        
       				Additional: msg.Additional,

     

       141
       158
        
       			}

     

       159
       159
       +
       			fullMsg.Header.ANCount = uint16(len(msg.Answer))

     

       160
       160
       +
       			fullMsg.Header.NSCount = uint16(len(msg.Authority))

     

       161
       161
       +
       			fullMsg.Header.ARCount = uint16(len(msg.Additional))

     

       142
       162
        
       

     

       143
       163
        
       			if msg.Header.RCode == magna.NXDOMAIN {

     

       144
       164
        
       				if h.Cache != nil {

     
···

       201
       221
        
       			if msg.Header.NSCount > 0 {

     

       202
       222
        
       				var nextServers []string

     

       203
       223
        
       				var nsRecords []magna.ResourceRecord

     

       224
       224
       +
       				nsRecordsWithoutGlue := 0

     

       204
       225
        
       

     

       205
       226
        
       				glueMap := make(map[string]string)

     

       206
       227
        
       				for _, rr := range msg.Additional {

     

       207
       228
        
       					if rr.RType == magna.AType {

     

       208
       208
       -
       						for _, nsRR := range msg.Authority {

     

       209
       209
       -
       							if nsRR.RType == magna.NSType && nsRR.RData.String() == rr.Name {

     

       210
       210
       -
       								glueMap[rr.Name] = rr.RData.String()

     

       211
       211
       -
       								break

     

       212
       212
       -
       							}

     

       213
       213
       -
       						}

     

       229
       229
       +
       						glueMap[strings.ToLower(rr.Name)] = rr.RData.String()

     

       214
       230
        
       					}

     

       215
       231
        
       				}

     

       216
       232
        
       

     
···

       221
       237
        
       				}

     

       222
       238
        
       

     

       223
       239
        
       				for _, nsRR := range nsRecords {

     

       224
       224
       -
       					nsName := nsRR.RData.String()

     

       240
       240
       +
       					nsName := strings.ToLower(nsRR.RData.String())

     

       241
       241
       +
       

     

       225
       242
        
       					if ip, found := glueMap[nsName]; found {

     

       226
       243
        
       						nextServers = append(nextServers, ip)

     

       227
       244
        
       					} else {

     

       245
       245
       +
       						nsRecordsWithoutGlue++

     

       228
       246
        
       						nsQuestion := magna.Question{QName: nsName, QType: magna.AType, QClass: magna.IN}

     

       229
       247
        
       						nsAnswers, err := h.resolveQuestion(resolveCtx, nsQuestion, h.RootServers)

     

       230
       248
        
       						if err != nil {

     

       231
       249
        
       							h.Logger.Warn("error resolving NS A record", "ns", nsName, "error", err)

     

       232
       250
        
       							continue

     

       251
       251
       +
       						} else {

     

       252
       252
       +
       							foundAddr := false

     

       253
       253
       +
       								for _, ans := range nsAnswers {

     

       254
       254
       +
       									if ans.RType == magna.AType && strings.ToLower(ans.Name) == nsName {

     

       255
       255
       +
       										ipAddr := ans.RData.String()

     

       256
       256
       +
       										nextServers = append(nextServers, ipAddr)

     

       257
       257
       +
       										foundAddr = true

     

       258
       258
       +
       									}

     

       259
       259
       +
       								}

     

       260
       260
       +
       							if !foundAddr {

     

       261
       261
       +
       								h.Logger.Warn("A record lookup for NS succeeded but yielded no matching A records", "ns", nsName)

     

       262
       262
       +
       							}

     

       233
       263
        
       						}

     

       234
       264
        
       

     

       235
       235
       -
       						for _, ans := range nsAnswers {

     

       236
       236
       -
       							if ans.RType == magna.AType {

     

       237
       237
       -
       								nextServers = append(nextServers, ans.RData.String())

     

       238
       238
       -
       							}

     

       239
       239
       -
       						}

     

       240
       265
        
       					}

     

       241
       266
        
       				}

     

       242
       267
        
       

     

       243
       268
        
       				if len(nextServers) > 0 {

     

       244
       244
       -
       					return h.resolveQuestion(resolveCtx, question, nextServers)

     

       269
       269
       +
       					return h.resolveQuestion(ctx, question, nextServers)

     

       270
       270
       +
       				} else {

     

       271
       271
       +
       					h.Logger.Warn("Processed delegation but failed to find any next server IPs", "question", question.QName, "server", res.Server, "ns_records_count", len(nsRecords), "ns_without_glue", nsRecordsWithoutGlue, "depth", depth)

     

       272
       272
       +
       					lastError = fmt.Errorf("failed to resolve NS addresses for delegation from %s", res.Server)

     

       245
       273
        
       				}

     

       246
       274
        
       

     

       247
       275
        
       				h.Logger.Warn("could not resolve any NS records for delegation", "question", question.QName)

     

       248
       276
        
       				continue

     

       249
       277
        
       			}

     

       250
       278
        
       

     

       251
       251
       -
       			if msg.Header.RCode == magna.NOERROR && msg.Header.ANCount == 0 {

     

       279
       279
       +
       			if msg.Header.RCode == magna.NOERROR && msg.Header.ANCount == 0 && msg.Header.NSCount == 0{

     

       280
       280
       +
       				h.Logger.Debug("Received NODATA response (NOERROR, ANCount=0)", "question", question.QName, "server", res.Server, "depth", depth)

     

       281
       281
       +
       

     

       252
       282
        
       				if h.Cache != nil {

     

       253
       283
        
       					entry := CreateCacheEntry(question, &fullMsg)

     

       254
       284
        
       					if entry != nil {

     

       255
       285
        
       						h.Cache.Set(GenerateCacheKey(question), entry)

     

       256
       286
        
       					}

     

       257
       287
        
       				}

     

       288
       288
       +
       

     

       258
       289
        
       				return []magna.ResourceRecord{}, nil

     

       259
       290
        
       			}

     

       260
       291
        
       

     

       261
       261
       -
       			h.Logger.Warn("unexpected response state", "question", question.QName, "server", res.Server, "rcode", msg.Header.RCode)

     

       262
       262
       -
       			continue

     

       292
       292
       +
       			h.Logger.Warn("Unhandled response state", "question", question.QName, "server", res.Server, "rcode", msg.Header.RCode.String(), "ancount", msg.Header.ANCount, "nscount", msg.Header.NSCount, "depth", depth)

     

       293
       293
       +
       			lastError = fmt.Errorf("unhandled response code %s from %s", msg.Header.RCode.String(), res.Server)

     

       263
       294
        
       

     

       264
       295
        
       		case <-resolveCtx.Done():

     

       265
       265
       -
       			h.Logger.Debug("resolution cancelled or timed out", "question", question.QName)

     

       266
       266
       -
       			return []magna.ResourceRecord{}, fmt.Errorf("resolution timed out or cancelled")

     

       296
       296
       +
       			h.Logger.Warn("Resolution step timed out or cancelled", "question", question.QName, "elapsed", time.Since(time.Now().Add(-h.Timeout)), "depth", depth, "error", resolveCtx.Err()) // Approx elapsed

     

       297
       297
       +
       			if lastError != nil {

     

       298
       298
       +
       				return nil, fmt.Errorf("resolution timed out after error: %w", lastError)

     

       299
       299
       +
       			}

     

       300
       300
       +
       			return nil, fmt.Errorf("resolution step timed out or cancelled: %w", resolveCtx.Err())

     

       267
       301
        
       		}

     

       268
       302
        
       	}

     

       269
       303
        
       

     

       270
       304
        
       	h.Logger.Warn("all resolution paths failed", "question", question.QName)

     

       305
       305
       +
       	if lastError != nil {

     

       306
       306
       +
       		return nil, fmt.Errorf("all resolution paths failed: %w", lastError)

     

       307
       307
       +
       	}

     

       271
       308
        
       	return []magna.ResourceRecord{}, fmt.Errorf("all resolution paths failed")

     

       272
       309
        
       }

     

       273
       310
        
       

     
···

       275
       312
        
       	var response queryResponse

     

       276
       313
        
       	response.Server = server

     

       277
       314
        
       

     

       315
       315
       +
       	bufPtr := resolverUDPBufferPool.Get().(*[]byte)

     

       316
       316
       +
       	defer resolverUDPBufferPool.Put(bufPtr)

     

       317
       317
       +
       	p := *bufPtr

     

       318
       318
       +
       

     

       278
       319
        
       	defer func() {

     

       320
       320
       +
       		if ctx.Err() != nil {

     

       321
       321
       +
       			if response.Error == nil {

     

       322
       322
       +
       				response.Error = ctx.Err()

     

       323
       323
       +
       			}

     

       324
       324
       +
       		}

     

       325
       325
       +
       

     

       279
       326
        
       		select {

     

       280
       327
        
       		case ch <- response:

     

       328
       328
       +
       		case <-ctx.Done():

     

       329
       329
       +
       			slog.Debug("queryServer response channel send cancelled by context", "server", server, "question", question.QName)

     

       281
       330
        
       		default:

     

       282
       282
       -
       			slog.Debug("queryServer response channel blocked or closed", "server", server)

     

       331
       331
       +
       			slog.Debug("queryServer response channel blocked or closed on send", "server", server, "question", question.QName)

     

       283
       332
        
       		}

     

       284
       333
        
       	}()

     

       285
       334
        
       

     
···

       293
       342
        
       	dialer := net.Dialer{Timeout: 2 * time.Second}

     

       294
       343
        
       	conn, err := dialer.DialContext(ctx, "udp", fmt.Sprintf("%s:53", server))

     

       295
       344
        
       	if err != nil {

     

       296
       296
       -
       		response.Error = fmt.Errorf("dial error: %w", err)

     

       345
       345
       +
       		if ctx.Err() != nil {

     

       346
       346
       +
       			response.Error = ctx.Err()

     

       347
       347
       +
       		} else {

     

       348
       348
       +
       			response.Error = fmt.Errorf("dial error: %w", err)

     

       349
       349
       +
       		}

     

       350
       350
       +
       

     

       297
       351
        
       		return

     

       298
       352
        
       	}

     

       299
       353
        
       	defer conn.Close()

     
···

       339
       393
        
       		return

     

       340
       394
        
       	}

     

       341
       395
        
       

     

       342
       342
       -
       	p := make([]byte, 512)

     

       343
       396
        
       	n, err := conn.Read(p)

     

       344
       397
        
       	if err != nil {

     

       345
       398
        
       		if ctx.Err() != nil {

+36 -4

pkg/dns/server.go

···

       16
       16
        
       type contextKey string

     

       17
       17
        
       

     

       18
       18
        
       const (

     

       19
       19
       -
       	cacheHitKey contextKey = "cache_hit"

     

       19
       19
       +
       	cacheHitKey contextKey   = "cache_hit"

     

       20
       20
       +
       	maxUDPBufferSize         = 4096

     

       21
       21
       +
       	maxResolverUDPBufferSize = 4096

     

       22
       22
       +
       )

     

       23
       23
       +
       

     

       24
       24
       +
       var (

     

       25
       25
       +
       	serverUDPBufferPool = sync.Pool {

     

       26
       26
       +
       		New: func() any {

     

       27
       27
       +
       			b := make([]byte, maxUDPBufferSize)

     

       28
       28
       +
       			return &b

     

       29
       29
       +
       		},

     

       30
       30
       +
       	}

     

       31
       31
       +
       

     

       32
       32
       +
       	resolverUDPBufferPool = sync.Pool {

     

       33
       33
       +
       		New: func() any {

     

       34
       34
       +
       			b := make([]byte, maxResolverUDPBufferSize)

     

       35
       35
       +
       			return &b

     

       36
       36
       +
       		},

     

       37
       37
       +
       	}

     

       20
       38
        
       )

     

       21
       39
        
       

     

       22
       40
        
       func setCacheHit(ctx context.Context) context.Context {

     
···

       240
       258
        
       	srv.Logger.Info("UDP listener started", "address", conn.LocalAddr())

     

       241
       259
        
       

     

       242
       260
        
       	for {

     

       243
       243
       -
       		buf := make([]byte, srv.UDPSize)

     

       261
       261
       +
       		bufPtr := serverUDPBufferPool.Get().(*[]byte)

     

       262
       262
       +
       		buffer := *bufPtr

     

       244
       263
        
       

     

       264
       264
       +
       		readDeadlineSet := false

     

       245
       265
        
       		if srv.ReadTimeout > 0 {

     

       246
       266
        
       			err := conn.SetReadDeadline(time.Now().Add(srv.ReadTimeout))

     

       247
       267
        
       			if err != nil {

     

       268
       268
       +
       				serverUDPBufferPool.Put(bufPtr)

     

       248
       269
        
       				return fmt.Errorf("error setting UDP read deadline: %w", err)

     

       249
       270
        
       			}

     

       271
       271
       +
       			readDeadlineSet = true

     

       250
       272
        
       		}

     

       251
       273
        
       

     

       252
       252
       -
       		n, remoteAddr, err := conn.ReadFromUDP(buf)

     

       274
       274
       +
       		n, remoteAddr, err := conn.ReadFromUDP(buffer)

     

       253
       275
        
       		if err != nil {

     

       276
       276
       +
       			serverUDPBufferPool.Put(bufPtr)

     

       254
       277
        
       			if netErr, ok := err.(net.Error); ok && netErr.Timeout() {

     

       278
       278
       +
       				if readDeadlineSet {

     

       279
       279
       +
       					continue

     

       280
       280
       +
       				}

     

       281
       281
       +
       

     

       282
       282
       +
       				srv.Logger.Warn("UDP read timeout occurred without explicit deadline set", "error", err)

     

       255
       283
        
       				continue

     

       256
       284
        
       			}

     

       257
       285
        
       

     
···

       264
       292
        
       			continue

     

       265
       293
        
       		}

     

       266
       294
        
       

     

       267
       267
       -
       		go srv.handleUDPQuery(conn, buf[:n], remoteAddr)

     

       295
       295
       +
       		queryData := make([]byte, n)

     

       296
       296
       +
       		copy(queryData, buffer[:n])

     

       297
       297
       +
       		serverUDPBufferPool.Put(bufPtr)

     

       298
       298
       +
       

     

       299
       299
       +
       		go srv.handleUDPQuery(conn, queryData, remoteAddr)

     

       268
       300
        
       	}

     

       269
       301
        
       }

     

       270
       302