commit d8aebc080fd95eac7a8ebaa4600cb7951125bcfb · seiso.moe/alky

+4 -2

docs/alky.toml

···

       5
       5
        
       port = 2053

     

       6
       6
        
       # Location of root hints file.

     

       7
       7
        
       root_hints_file = "/etc/dns/root.hints"

     

       8
       8
       +
       udp_payload_size = 512

     

       8
       9
        
       

     

       9
       10
        
       [logging]

     

       10
       11
        
       # Logging output: "stdout" or "file".

     

       11
       12
        
       output = "stdout"

     

       13
       13
       +
       level = "debug"

     

       12
       14
        
       # This is used only if logging.output is "file".

     

       13
       15
        
       file_path = "/var/log/alky.log"

     

       14
       16
        
       

     
···

       26
       28
        
       # window: The interval (in seconds) at which the rate limit is checked and potentially reset

     

       27
       29
        
       # Implements a sliding window rate limit mechanism

     

       28
       30
        
       # Type: Integer

     

       29
       29
       -
       window = 3

     

       31
       31
       +
       window_seconds = "1s"

     

       30
       32
        
       

     

       31
       33
        
       # expiration_time: Duration (in seconds) for keeping a client's rate limit data in memory

     

       32
       34
        
       # After this period of inactivity, a client's rate limit data is removed to free up memory

     

       33
       35
        
       # Type: Integer

     

       34
       34
       -
       expiration_time = 300

     

       36
       36
       +
       cleanup_seconds = "5m"

     

       35
       37
        
       

     

       36
       38
        
       [metrics]

     

       37
       39
        
       # ClickHouse connection string

+109 -43

main.go

···

       6
       6
        
       	"log/slog"

     

       7
       7
        
       	"os"

     

       8
       8
        
       	"os/signal"

     

       9
       9
       +
       	"strings"

     

       9
       10
        
       	"syscall"

     

       10
       11
        
       	"time"

     

       11
       12
        
       

     
···

       19
       20
        
       

     

       20
       21
        
       func init() {

     

       21
       22
        
       	flag.StringVar(&configFlag, "config", "/etc/alky/alky.toml", "config file path for alky")

     

       22
       22
       -
       

     

       23
       23
        
       	flag.Parse()

     

       24
       24
        
       }

     

       25
       25
        
       

     

       26
       26
        
       func main() {

     

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

     

       28
       28
        
       	if err != nil {

     

       29
       29
       -
       		log.Fatal(err)

     

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

     

       30
       30
        
       	}

     

       31
       31
        
       

     

       32
       32
        
       	logger := setupLogger(&cfg)

     

       33
       33
       +
       	slog.SetDefault(logger)

     

       33
       34
        
       

     

       34
       34
       -
       	metricsClient, err := metrics.NewClickHouseMetrics(&cfg.Metrics, logger)

     

       35
       35
       +
       	logger.Info("Configuration loaded", "path", configFlag)

     

       36
       36
       +
       

     

       37
       37
       +
       	metricsClient, err := metrics.NewClickHouseMetrics(&cfg.Metrics, logger.With("component", "metrics"))

     

       35
       38
        
       	if err != nil {

     

       36
       36
       -
       		log.Fatal(err)

     

       39
       39
       +
       		logger.Error("failed to initialize metrics client", "error", err)

     

       40
       40
       +
       		os.Exit(1)

     

       37
       41
        
       	}

     

       38
       42
        
       	defer metricsClient.Close()

     

       43
       43
       +
       	logger.Info("Metrics client initialized")

     

       39
       44
        
       

     

       40
       45
        
       	rootServers, err := rootservers.DecodeRootHints(cfg.Server.RootHintsFile)

     

       41
       41
       -
       	if err != nil {

     

       42
       42
       -
       		log.Fatal(err)

     

       46
       46
       +
       	if err != nil || len(rootServers) == 0 {

     

       47
       47
       +
       		logger.Error("Failed to load root hints or no root servers found", "path", cfg.Server.RootHintsFile, "error", err)

     

       48
       48
       +
       		os.Exit(1)

     

       43
       49
        
       	}

     

       50
       50
       +
       	logger.Info("Root hints loaded", "count", len(rootServers), "path", cfg.Server.RootHintsFile)

     

       44
       51
        
       

     

       45
       45
       -
       	cache := dns.NewLRUCache(cfg.Cache.MaxItems, cfg.Cache.CleanupInterval.Duration)

     

       52
       52
       +
       	cache := dns.NewLRUCache(cfg.Cache.MaxItems, cfg.Cache.CleanupInterval.Duration, logger.With("component", "cache"))

     

       46
       53
        
       	defer cache.Stop()

     

       54
       54
       +
       	logger.Info("DNS cache initialized")

     

       47
       55
        
       

     

       48
       48
       -
       	handler := &dns.QueryHandler{

     

       56
       56
       +
       	queryHandler := &dns.QueryHandler{

     

       49
       57
        
       		RootServers: rootServers,

     

       50
       50
       -
       		Timeout:     time.Duration(cfg.Advanced.QueryTimeout) * time.Second,

     

       58
       58
       +
       		Timeout:     cfg.Advanced.QueryTimeout.Duration,

     

       51
       59
        
       		Cache:       cache,

     

       52
       52
       -
       		Logger:      logger,

     

       60
       60
       +
       		Logger:      logger.With("component", "resolver"),

     

       53
       61
        
       	}

     

       54
       62
        
       

     

       55
       55
       -
       	go monitorCacheMetrics(cache, metricsClient, logger)

     

       63
       63
       +
       	go monitorCacheMetrics(cache, metricsClient, logger.With("component", "cache-monitor"))

     

       56
       64
        
       

     

       57
       57
       -
       	rateLimitHandler := dns.RateLimitMiddleware(&dns.RateLimitConfig{

     

       58
       58
       -
       		Rate:           float64(cfg.Ratelimit.Rate),

     

       59
       59
       -
       		Burst:          cfg.Ratelimit.Burst,

     

       60
       60
       -
       		WindowLength:   time.Duration(cfg.Ratelimit.Window) * time.Second,

     

       61
       61
       -
       		ExpirationTime: time.Duration(cfg.Ratelimit.ExpirationTime) * time.Second,

     

       62
       62
       -
       	})(handler)

     

       65
       65
       +
       	var currentHandler dns.Handler = queryHandler

     

       63
       66
        
       

     

       64
       64
       -
       	metricsHandler := metrics.MetricsMiddleware(metricsClient)(rateLimitHandler)

     

       67
       67
       +
       	if cfg.Ratelimit.Rate > 0 {

     

       68
       68
       +
       		rateLimitHandler := dns.RateLimitMiddleware(&dns.RateLimitConfig{

     

       69
       69
       +
       			Rate:           float64(cfg.Ratelimit.Rate),

     

       70
       70
       +
       			Burst:          cfg.Ratelimit.Burst,

     

       71
       71
       +
       			WindowLength:   cfg.Ratelimit.Window.Duration,

     

       72
       72
       +
       			ExpirationTime: cfg.Ratelimit.ExpirationTime.Duration,

     

       73
       73
       +
       		})(currentHandler)

     

       74
       74
       +
       		currentHandler = rateLimitHandler

     

       75
       75
       +
       		logger.Info("Rate limiting enabled", "rate", cfg.Ratelimit.Rate, "burst", cfg.Ratelimit.Burst)

     

       76
       76
       +
       	} else {

     

       77
       77
       +
       		logger.Info("Rate limiting disabled")

     

       78
       78
       +
       	}

     

       79
       79
       +
       

     

       80
       80
       +
       	metricsHandler := metrics.MetricsMiddleware(metricsClient)(currentHandler)

     

       81
       81
       +
       	currentHandler = metricsHandler

     

       82
       82
       +
       	logger.Info("Metrics middleware enabled")

     

       65
       83
        
       

     

       66
       84
        
       	loggingHandler := dns.LoggingMiddleware(&dns.LogConfig{

     

       67
       67
       -
       		Logger: logger,

     

       85
       85
       +
       		Logger: logger.With("component", "query-log"),

     

       68
       86
        
       		Level:  slog.LevelInfo,

     

       69
       69
       -
       	})(metricsHandler)

     

       87
       87
       +
       	})(currentHandler)

     

       88
       88
       +
       	currentHandler = loggingHandler

     

       89
       89
       +
       	logger.Info("Logging middleware enabled")

     

       70
       90
        
       

     

       91
       91
       +
       	sigChan := make(chan os.Signal, 1)

     

       92
       92
       +
       	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

     

       71
       93
        
       	go func() {

     

       72
       72
       -
       		sigChan := make(chan os.Signal, 1)

     

       73
       73
       -
       		signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

     

       74
       74
       -
       		<-sigChan

     

       75
       75
       -
       		slog.Info("Shutting down...")

     

       94
       94
       +
       		sig := <-sigChan

     

       95
       95
       +
       		logger.Info("Received signal, shutting down gracefully...", "signal", sig.String())

     

       96
       96
       +
       

     

       97
       97
       +
       		logger.Info("Stopping cache...")

     

       76
       98
        
       		cache.Stop()

     

       99
       99
       +
       

     

       100
       100
       +
       		logger.Info("Closing metrics client...")

     

       77
       101
        
       		metricsClient.Close()

     

       102
       102
       +
       

     

       103
       103
       +
       		logger.Info("Shutdown complete.")

     

       78
       104
        
       		os.Exit(0)

     

       79
       105
        
       	}()

     

       80
       106
        
       

     

       81
       107
        
       	s := dns.Server{

     

       82
       108
        
       		Address:      cfg.Server.Address,

     

       83
       109
        
       		Port:         cfg.Server.Port,

     

       84
       84
       -
       		Handler:      loggingHandler,

     

       85
       85
       -
       		UDPSize:      512,

     

       86
       86
       -
       		ReadTimeout:  2 * time.Second,

     

       87
       87
       -
       		WriteTimeout: 2 * time.Second,

     

       88
       88
       -
       		Logger:       logger,

     

       110
       110
       +
       		Handler:      currentHandler,

     

       111
       111
       +
       		UDPSize:      cfg.Server.UDPSize,

     

       112
       112
       +
       		ReadTimeout:  cfg.Advanced.ReadTimeout.Duration,

     

       113
       113
       +
       		WriteTimeout: cfg.Advanced.WriteTimeout.Duration,

     

       114
       114
       +
       		Logger:       logger.With("component", "server"),

     

       89
       115
        
       	}

     

       90
       116
        
       

     

       117
       117
       +
       	logger.Info("Starting DNS server listener...")

     

       91
       118
        
       	if err := s.ListenAndServe(); err != nil {

     

       92
       92
       -
       		slog.Error("Failed to start server", "error", err)

     

       93
       93
       -
       		cache.Stop()

     

       94
       94
       -
       		metricsClient.Close()

     

       119
       119
       +
       		logger.Error("Failed to start server", "error", err)

     

       120
       120
       +
       		os.Exit(1)

     

       95
       121
        
       	}

     

       96
       122
        
       }

     

       97
       123
        
       

     

       98
       124
        
       func monitorCacheMetrics(cache dns.Cache, metricsClient *metrics.ClickHouseMetrics, logger *slog.Logger) {

     

       99
       99
       -
       	ticker := time.NewTicker(1 * time.Minute)

     

       125
       125
       +
       	interval := 1 * time.Minute

     

       126
       126
       +
       	logger.Info("Starting cache metrics monitoring", "interval", interval)

     

       127
       127
       +
       	ticker := time.NewTicker(interval)

     

       100
       128
        
       	defer ticker.Stop()

     

       101
       129
        
       

     

       102
       102
       -
       	for range ticker.C {

     

       103
       103
       -
       		stats := cache.GetStats()

     

       104
       104
       -
       		metricsClient.RecordCacheStats(stats)

     

       105
       105
       -
       		logger.Info("Cache metrics recorded to ClickHouse")

     

       130
       130
       +
       	for {

     

       131
       131
       +
       		select {

     

       132
       132
       +
       		case <-ticker.C:

     

       133
       133
       +
       			stats := cache.GetStats()

     

       134
       134
       +
       			logger.Debug("Recording cache metrics",

     

       135
       135
       +
       				"hits", stats.Hits.Load(),

     

       136
       136
       +
       				"misses", stats.Misses.Load(),

     

       137
       137
       +
       				"size", stats.Size.Load(),

     

       138
       138
       +
       				"evictions", stats.Evictions.Load(),

     

       139
       139
       +
       				"expired", stats.Expired.Load(),

     

       140
       140
       +
       				"pos_hits", stats.PositiveHits.Load(),

     

       141
       141
       +
       				"neg_hits", stats.NegativeHits.Load(),

     

       142
       142
       +
       			)

     

       143
       143
       +
       			metricsClient.RecordCacheStats(metrics.CacheMetric{

     

       144
       144
       +
       				Timestamp:    time.Now(),

     

       145
       145
       +
       				CacheHits:    stats.Hits.Load(),

     

       146
       146
       +
       				CacheMisses:  stats.Misses.Load(),

     

       147
       147
       +
       				NegativeHits: stats.NegativeHits.Load(),

     

       148
       148
       +
       				PositiveHits: stats.PositiveHits.Load(),

     

       149
       149
       +
       				Evictions:    stats.Evictions.Load() + stats.Expired.Load(),

     

       150
       150
       +
       				Size:         stats.Size.Load(),

     

       151
       151
       +
       			})

     

       152
       152
       +
       		}

     

       106
       153
        
       	}

     

       107
       154
        
       }

     

       108
       155
        
       

     

       109
       156
        
       func setupLogger(cfg *config.Config) *slog.Logger {

     

       110
       110
       -
       	var logger *slog.Logger

     

       157
       157
       +
       	var logLevel slog.Level

     

       158
       158
       +
       	switch strings.ToLower(cfg.Logging.Level) {

     

       159
       159
       +
       	case "debug":

     

       160
       160
       +
       		logLevel = slog.LevelDebug

     

       161
       161
       +
       	case "info":

     

       162
       162
       +
       		logLevel = slog.LevelInfo

     

       163
       163
       +
       	case "warn":

     

       164
       164
       +
       		logLevel = slog.LevelWarn

     

       165
       165
       +
       	case "error":

     

       166
       166
       +
       		logLevel = slog.LevelError

     

       167
       167
       +
       	default:

     

       168
       168
       +
       		logLevel = slog.LevelInfo

     

       169
       169
       +
       	}

     

       111
       170
        
       

     

       112
       171
        
       	handlerOpts := &slog.HandlerOptions{

     

       113
       113
       -
       		Level: slog.LevelDebug,

     

       172
       172
       +
       		Level:     logLevel,

     

       173
       173
       +
       		AddSource: logLevel <= slog.LevelDebug,

     

       114
       174
        
       	}

     

       115
       175
        
       

     

       176
       176
       +
       	var handler slog.Handler

     

       116
       177
        
       	switch cfg.Logging.Output {

     

       117
       178
        
       	case "file":

     

       118
       179
        
       		f, err := os.OpenFile(cfg.Logging.FilePath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0o644)

     

       119
       180
        
       		if err != nil {

     

       120
       120
       -
       			log.Fatal(err)

     

       181
       181
       +
       			log.Fatalf("Failed to open log file %s: %v", cfg.Logging.FilePath, err)

     

       121
       182
        
       		}

     

       122
       122
       -
       

     

       123
       123
       -
       		logger = slog.New(slog.NewJSONHandler(f, handlerOpts))

     

       183
       183
       +
       		handler = slog.NewJSONHandler(f, handlerOpts)

     

       184
       184
       +
       		log.Printf("Logging to file: %s at level: %s", cfg.Logging.FilePath, logLevel.String())

     

       185
       185
       +
       	case "text":

     

       186
       186
       +
       		handler = slog.NewTextHandler(os.Stdout, handlerOpts)

     

       187
       187
       +
       		log.Printf("Logging to stdout (text) at level: %s", logLevel.String())

     

       124
       188
        
       	default:

     

       125
       125
       -
       		logger = slog.New(slog.NewJSONHandler(os.Stdout, handlerOpts))

     

       189
       189
       +
       		handler = slog.NewJSONHandler(os.Stdout, handlerOpts)

     

       190
       190
       +
       		log.Printf("Logging to stdout (json) at level: %s", logLevel.String())

     

       126
       191
        
       	}

     

       127
       192
        
       

     

       193
       193
       +
       	logger := slog.New(handler)

     

       128
       194
        
       	return logger

     

       129
       195
        
       }

+60

migrations/00002_modified_metrics.sql

···

       1
       1
       +
       -- +goose Up

     

       2
       2
       +
       ALTER TABLE alky_dns_queries

     

       3
       3
       +
           MODIFY COLUMN timestamp DateTime CODEC(Delta, ZSTD(1)),

     

       4
       4
       +
           MODIFY COLUMN instance_id String CODEC(ZSTD(1)),

     

       5
       5
       +
           MODIFY COLUMN query_name String CODEC(ZSTD(1)),

     

       6
       6
       +
           MODIFY COLUMN query_type LowCardinality(String) CODEC(ZSTD(1)),

     

       7
       7
       +
           MODIFY COLUMN query_class LowCardinality(String) CODEC(ZSTD(1)),

     

       8
       8
       +
           MODIFY COLUMN remote_addr String CODEC(ZSTD(1)),

     

       9
       9
       +
           MODIFY COLUMN response_code LowCardinality(String) CODEC(ZSTD(1)),

     

       10
       10
       +
           MODIFY COLUMN duration Int64 CODEC(T64, ZSTD(1)),

     

       11
       11
       +
           MODIFY COLUMN cache_hit Bool CODEC(ZSTD(1));

     

       12
       12
       +
       

     

       13
       13
       +
       ALTER TABLE alky_dns_queries MODIFY TTL timestamp + INTERVAL 30 DAY;

     

       14
       14
       +
       

     

       15
       15
       +
       ALTER TABLE alky_dns_cache_metrics 

     

       16
       16
       +
           DROP COLUMN IF EXISTS total_queries,

     

       17
       17
       +
           MODIFY COLUMN timestamp DateTime CODEC(Delta, ZSTD(1)),

     

       18
       18
       +
           MODIFY COLUMN instance_id String CODEC(ZSTD(1)),

     

       19
       19
       +
           MODIFY COLUMN cache_hits Int64 CODEC(T64, ZSTD(1)),

     

       20
       20
       +
           MODIFY COLUMN cache_misses Int64 CODEC(T64, ZSTD(1)),

     

       21
       21
       +
           MODIFY COLUMN negative_hits Int64 CODEC(T64, ZSTD(1)),

     

       22
       22
       +
           MODIFY COLUMN positive_hits Int64 CODEC(T64, ZSTD(1)),

     

       23
       23
       +
           MODIFY COLUMN evictions Int64 CODEC(T64, ZSTD(1)),

     

       24
       24
       +
           MODIFY COLUMN size Int64 CODEC(T64, ZSTD(1));

     

       25
       25
       +
       

     

       26
       26
       +
       ALTER TABLE alky_dns_cache_metrics 

     

       27
       27
       +
           ADD COLUMN IF NOT EXISTS expired_count Int64 CODEC(T64, ZSTD(1));

     

       28
       28
       +
       

     

       29
       29
       +
       ALTER TABLE alky_dns_cache_metrics MODIFY TTL timestamp + INTERVAL 30 DAY;

     

       30
       30
       +
       

     

       31
       31
       +
       -- +goose Down

     

       32
       32
       +
       ALTER TABLE alky_dns_queries

     

       33
       33
       +
           MODIFY COLUMN timestamp DateTime,

     

       34
       34
       +
           MODIFY COLUMN instance_id String,

     

       35
       35
       +
           MODIFY COLUMN query_name String,

     

       36
       36
       +
           MODIFY COLUMN query_type String,

     

       37
       37
       +
           MODIFY COLUMN query_class String,

     

       38
       38
       +
           MODIFY COLUMN remote_addr String,

     

       39
       39
       +
           MODIFY COLUMN response_code String,

     

       40
       40
       +
           MODIFY COLUMN duration Int64,

     

       41
       41
       +
           MODIFY COLUMN cache_hit Bool;

     

       42
       42
       +
       

     

       43
       43
       +
       ALTER TABLE alky_dns_queries MODIFY TTL timestamp + toIntervalDay(30);

     

       44
       44
       +
       

     

       45
       45
       +
       ALTER TABLE alky_dns_cache_metrics 

     

       46
       46
       +
           ADD COLUMN IF NOT EXISTS total_queries Int64 AFTER instance_id;

     

       47
       47
       +
       

     

       48
       48
       +
       ALTER TABLE alky_dns_cache_metrics 

     

       49
       49
       +
           MODIFY COLUMN timestamp DateTime,

     

       50
       50
       +
           MODIFY COLUMN instance_id String,

     

       51
       51
       +
           MODIFY COLUMN cache_hits Int64,

     

       52
       52
       +
           MODIFY COLUMN cache_misses Int64,

     

       53
       53
       +
           MODIFY COLUMN negative_hits Int64,

     

       54
       54
       +
           MODIFY COLUMN positive_hits Int64,

     

       55
       55
       +
           MODIFY COLUMN evictions Int64,

     

       56
       56
       +
           MODIFY COLUMN size Int;

     

       57
       57
       +
       

     

       58
       58
       +
       ALTER TABLE alky_dns_cache_metrics DROP COLUMN IF EXISTS expired_count;

     

       59
       59
       +
       

     

       60
       60
       +
       ALTER TABLE alky_dns_cache_metrics MODIFY TTL timestamp + toIntervalDay(30);

+91 -30

pkg/config/config.go

···

       2
       2
        
       

     

       3
       3
        
       import (

     

       4
       4
        
       	"fmt"

     

       5
       5
       +
       	"os"

     

       6
       6
       +
       	"strings"

     

       5
       7
        
       	"time"

     

       6
       8
        
       

     

       7
       9
        
       	"github.com/BurntSushi/toml"

     
···

       11
       13
        
       	Address       string `toml:"address"`

     

       12
       14
        
       	Port          int    `toml:"port"`

     

       13
       15
        
       	RootHintsFile string `toml:"root_hints_file"`

     

       16
       16
       +
       	UDPSize       int    `toml:"udp_payload_size"`

     

       14
       17
        
       }

     

       15
       18
        
       

     

       16
       19
        
       type LoggingConfig struct {

     

       17
       20
        
       	Output   string `toml:"output"`

     

       18
       21
        
       	FilePath string `toml:"file_path"`

     

       22
       22
       +
       	Level    string `toml:"level"`

     

       19
       23
        
       }

     

       20
       24
        
       

     

       21
       25
        
       type RatelimitConfig struct {

     

       22
       22
       -
       	Rate           int `toml:"rate"`

     

       23
       23
       -
       	Burst          int `toml:"burst"`

     

       24
       24
       -
       	Window         int `toml:"window"`

     

       25
       25
       -
       	ExpirationTime int `toml:"expiration_time"`

     

       26
       26
       +
       	Rate           int      `toml:"rate"`

     

       27
       27
       +
       	Burst          int      `toml:"burst"`

     

       28
       28
       +
       	Window         duration `toml:"window_seconds"`

     

       29
       29
       +
       	ExpirationTime duration `toml:"cleanup_seconds"`

     

       26
       30
        
       }

     

       27
       31
        
       

     

       28
       32
        
       type MetricsConfig struct {

     
···

       33
       37
        
       }

     

       34
       38
        
       

     

       35
       39
        
       type AdvancedConfig struct {

     

       36
       36
       -
       	QueryTimeout int `toml:"query_timeout"`

     

       40
       40
       +
       	QueryTimeout duration `toml:"query_timeout"`

     

       41
       41
       +
       	ReadTimeout  duration `toml:"read_timeout"`

     

       42
       42
       +
       	WriteTimeout duration `toml:"write_timeout"`

     

       37
       43
        
       }

     

       38
       44
        
       

     

       39
       45
        
       type CacheConfig struct {

     
···

       57
       63
        
       func (d *duration) UnmarshalText(text []byte) error {

     

       58
       64
        
       	var err error

     

       59
       65
        
       	d.Duration, err = time.ParseDuration(string(text))

     

       60
       60
       -
       	return err

     

       66
       66
       +
       	if err != nil {

     

       67
       67
       +
       		var seconds int64

     

       68
       68
       +
       		seconds, err = parseInt64(string(text))

     

       69
       69
       +
       		if err == nil {

     

       70
       70
       +
       			d.Duration = time.Duration(seconds) * time.Second

     

       71
       71
       +
       			return nil

     

       72
       72
       +
       		}

     

       73
       73
       +
       		return fmt.Errorf("invalid duration format: %w", err)

     

       74
       74
       +
       	}

     

       75
       75
       +
       	return nil

     

       76
       76
       +
       }

     

       77
       77
       +
       

     

       78
       78
       +
       func parseInt64(s string) (int64, error) {

     

       79
       79
       +
       	var i int64

     

       80
       80
       +
       	_, err := fmt.Sscan(s, &i)

     

       81
       81
       +
       	return i, err

     

       61
       82
        
       }

     

       62
       83
        
       

     

       63
       84
        
       func LoadConfig(path string) (Config, error) {

     

       64
       85
        
       	cfg := Config{}

     

       65
       65
       -
       	if _, err := toml.DecodeFile(path, &cfg); err != nil {

     

       66
       66
       -
       		return cfg, err

     

       86
       86
       +
       	cfg.Server.Address = "127.0.0.1"

     

       87
       87
       +
       	cfg.Server.Port = 53

     

       88
       88
       +
       	cfg.Server.RootHintsFile = "/etc/alky/root.hints"

     

       89
       89
       +
       	cfg.Server.UDPSize = 512

     

       90
       90
       +
       

     

       91
       91
       +
       	cfg.Logging.Output = "stdout"

     

       92
       92
       +
       	cfg.Logging.Level = "info"

     

       93
       93
       +
       

     

       94
       94
       +
       	cfg.Ratelimit.Rate = 100

     

       95
       95
       +
       	cfg.Ratelimit.Burst = 200

     

       96
       96
       +
       	cfg.Ratelimit.Window.Duration = 1 * time.Second

     

       97
       97
       +
       	cfg.Ratelimit.ExpirationTime.Duration = 1 * time.Minute

     

       98
       98
       +
       

     

       99
       99
       +
       	cfg.Metrics.DSN = "clickhouse://localhost:9000/default"

     

       100
       100
       +
       	cfg.Metrics.BatchSize = 1000

     

       101
       101
       +
       	cfg.Metrics.FlushInterval.Duration = 10 * time.Second

     

       102
       102
       +
       	cfg.Metrics.RetentionPeriod.Duration = 30 * 24 * time.Hour

     

       103
       103
       +
       

     

       104
       104
       +
       	cfg.Cache.MaxItems = 10000

     

       105
       105
       +
       	cfg.Cache.CleanupInterval.Duration = 5 * time.Minute

     

       106
       106
       +
       

     

       107
       107
       +
       	cfg.Advanced.QueryTimeout.Duration = 5 * time.Second

     

       108
       108
       +
       	cfg.Advanced.ReadTimeout.Duration = 2 * time.Second

     

       109
       109
       +
       	cfg.Advanced.WriteTimeout.Duration = 2 * time.Second

     

       110
       110
       +
       

     

       111
       111
       +
       	md, err := toml.DecodeFile(path, &cfg)

     

       112
       112
       +
       	if err != nil {

     

       113
       113
       +
       		if os.IsNotExist(err) {

     

       114
       114
       +
       			fmt.Printf("Warning: Config file '%s' not found, using default settings.\n", path)

     

       115
       115
       +
       		} else {

     

       116
       116
       +
       			return cfg, fmt.Errorf("error decoding config file '%s': %w", path, err)

     

       117
       117
       +
       		}

     

       67
       118
        
       	}

     

       68
       119
        
       

     

       69
       69
       -
       	if cfg.Server.Address == "" {

     

       70
       70
       -
       		cfg.Server.Address = "127.0.0.1"

     

       120
       120
       +
       	if len(md.Undecoded()) > 0 {

     

       121
       121
       +
       		return cfg, fmt.Errorf("unknown configuration keys found: %v", md.Undecoded())

     

       71
       122
        
       	}

     

       72
       123
        
       

     

       73
       73
       -
       	if cfg.Server.Port == 0 {

     

       74
       74
       -
       		cfg.Server.Port = 53

     

       124
       124
       +
       	if cfg.Logging.Output == "file" && cfg.Logging.FilePath == "" {

     

       125
       125
       +
       		return cfg, fmt.Errorf("logging output is 'file' but 'file_path' is not set")

     

       75
       126
        
       	}

     

       76
       127
        
       

     

       77
       77
       -
       	if cfg.Server.RootHintsFile == "" {

     

       78
       78
       -
       		cfg.Server.RootHintsFile = "/etc/dns/root.hints"

     

       128
       128
       +
       	validLevels := []string{"debug", "info", "warn", "error"}

     

       129
       129
       +
       	isValidLevel := false

     

       130
       130
       +
       	for _, level := range validLevels {

     

       131
       131
       +
       		if strings.ToLower(cfg.Logging.Level) == level {

     

       132
       132
       +
       			isValidLevel = true

     

       133
       133
       +
       			break

     

       134
       134
       +
       		}

     

       79
       135
        
       	}

     

       80
       136
        
       

     

       81
       81
       -
       	if cfg.Logging.Output == "file" && cfg.Logging.FilePath == "" {

     

       82
       82
       -
       		return cfg, fmt.Errorf("If `[logging.output]` is `file` then `file_path` must be set.")

     

       137
       137
       +
       	if !isValidLevel {

     

       138
       138
       +
       		return cfg, fmt.Errorf("invalid logging level '%s', must be one of: %v", cfg.Logging.Level, validLevels)

     

       83
       139
        
       	}

     

       84
       140
        
       

     

       85
       85
       -
       	if cfg.Metrics.DSN == "" {

     

       86
       86
       -
       		cfg.Metrics.DSN = "clickhouse://localhost:9000/default"

     

       141
       141
       +
       	if cfg.Cache.MaxItems <= 0 {

     

       142
       142
       +
       		return cfg, fmt.Errorf("cache max_items must be positive")

     

       87
       143
        
       	}

     

       88
       144
        
       

     

       89
       89
       -
       	if cfg.Metrics.BatchSize == 0 {

     

       90
       90
       -
       		cfg.Metrics.BatchSize = 1000

     

       145
       145
       +
       	if cfg.Cache.CleanupInterval.Duration <= 0 {

     

       146
       146
       +
       		cfg.Cache.CleanupInterval.Duration = 5 * time.Minute

     

       91
       147
        
       	}

     

       92
       148
        
       

     

       93
       93
       -
       	if cfg.Metrics.FlushInterval.Duration == 0 {

     

       94
       94
       -
       		cfg.Metrics.FlushInterval.Duration = 10 * time.Second

     

       149
       149
       +
       	if cfg.Advanced.QueryTimeout.Duration <= 0 {

     

       150
       150
       +
       		cfg.Advanced.QueryTimeout.Duration = 5 * time.Second

     

       95
       151
        
       	}

     

       96
       152
        
       

     

       97
       97
       -
       	if cfg.Metrics.RetentionPeriod.Duration == 0 {

     

       98
       98
       -
       		cfg.Metrics.RetentionPeriod.Duration = 30 * 24 * time.Hour

     

       153
       153
       +
       	if cfg.Advanced.ReadTimeout.Duration <= 0 {

     

       154
       154
       +
       		cfg.Advanced.ReadTimeout.Duration = 2 * time.Second

     

       99
       155
        
       	}

     

       100
       156
        
       

     

       101
       101
       -
       	if cfg.Cache.MaxItems == 0 {

     

       102
       102
       -
       		cfg.Cache.MaxItems = 5000

     

       157
       157
       +
       	if cfg.Advanced.WriteTimeout.Duration <= 0 {

     

       158
       158
       +
       		cfg.Advanced.WriteTimeout.Duration = 2 * time.Second

     

       103
       159
        
       	}

     

       104
       160
        
       

     

       105
       105
       -
       	if cfg.Cache.CleanupInterval.Duration == 0 {

     

       106
       106
       -
       		cfg.Cache.CleanupInterval.Duration = 5 * time.Minute

     

       161
       161
       +
       	if cfg.Server.UDPSize <= 0 || cfg.Server.UDPSize > 4096 {

     

       162
       162
       +
       		cfg.Server.UDPSize = 512

     

       107
       163
        
       	}

     

       108
       164
        
       

     

       109
       109
       -
       	if cfg.Advanced.QueryTimeout == 0 {

     

       110
       110
       -
       		cfg.Advanced.QueryTimeout = 100

     

       165
       165
       +
       	if cfg.Ratelimit.Rate > 0 {

     

       166
       166
       +
       		if cfg.Ratelimit.Window.Duration <= 0 {

     

       167
       167
       +
       			cfg.Ratelimit.Window.Duration = 1 * time.Second

     

       168
       168
       +
       		}

     

       169
       169
       +
       		if cfg.Ratelimit.ExpirationTime.Duration <= 0 {

     

       170
       170
       +
       			cfg.Ratelimit.ExpirationTime.Duration = 1 * time.Minute

     

       171
       171
       +
       		}

     

       111
       172
        
       	}

     

       112
       173
        
       

     

       113
       174
        
       	return cfg, nil

+188 -171

pkg/dns/cache.go

···

       2
       2
        
       

     

       3
       3
        
       import (

     

       4
       4
        
       	"container/list"

     

       5
       5
       +
       	"fmt"

     

       6
       6
       +
       	"log/slog"

     

       5
       7
        
       	"strings"

     

       6
       8
        
       	"sync"

     

       7
       9
        
       	"sync/atomic"

     
···

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

     

       11
       13
        
       )

     

       12
       14
        
       

     

       13
       13
       -
       type BailiwickRule int

     

       14
       14
       -
       

     

       15
       15
        
       const (

     

       16
       16
       -
       	BailiwickSame BailiwickRule = iota

     

       17
       17
       -
       	BailiwickChild

     

       18
       18
       -
       	BailiwickOutside

     

       16
       16
       +
       	defaultNegativeTTL = 5 * time.Minute

     

       17
       17
       +
       	maxNegativeTTL     = 3 * time.Hour

     

       19
       18
        
       )

     

       20
       19
        
       

     

       21
       20
        
       type CacheEntry struct {

     

       22
       22
       -
       	Answer      []CachedResourceRecord

     

       23
       23
       -
       	Authority   []CachedResourceRecord

     

       24
       24
       -
       	Additional  []CachedResourceRecord

     

       25
       25
       -
       	NegativeTTL time.Duration

     

       26
       26
       -
       	ExpireAt    time.Time

     

       27
       27
       -
       	IsNegative  bool

     

       28
       28
       -
       }

     

       21
       21
       +
       	Key        string

     

       22
       22
       +
       	RCode      magna.RCode

     

       23
       23
       +
       	IsNegative bool

     

       24
       24
       +
       

     

       25
       25
       +
       	Answer     []magna.ResourceRecord

     

       26
       26
       +
       	Authority  []magna.ResourceRecord

     

       27
       27
       +
       	Additional []magna.ResourceRecord

     

       29
       28
        
       

     

       30
       30
       -
       type CachedResourceRecord struct {

     

       31
       31
       -
       	Record        magna.ResourceRecord

     

       32
       32
       -
       	ExpireAt      time.Time

     

       33
       33
       -
       	BailiwickRule BailiwickRule

     

       29
       29
       +
       	CacheTime time.Time

     

       30
       30
       +
       	ExpireAt  time.Time

     

       34
       31
        
       }

     

       35
       32
        
       

     

       36
       33
        
       type CacheStats struct {

     

       37
       37
       -
       	TotalQueries atomic.Int64

     

       38
       38
       -
       	CacheHits    atomic.Int64

     

       39
       39
       -
       	CacheMisses  atomic.Int64

     

       40
       40
       -
       	NegativeHits atomic.Int64

     

       41
       41
       -
       	PositiveHits atomic.Int64

     

       34
       34
       +
       	Hits         atomic.Int64

     

       35
       35
       +
       	Misses       atomic.Int64

     

       42
       36
        
       	Evictions    atomic.Int64

     

       37
       37
       +
       	Expired      atomic.Int64

     

       43
       38
        
       	Size         atomic.Int64

     

       39
       39
       +
       	NegativeHits atomic.Int64

     

       40
       40
       +
       	PositiveHits atomic.Int64

     

       44
       41
        
       }

     

       45
       42
        
       

     

       46
       43
        
       type LRUCache struct {

     
···

       51
       48
        
       	cacheMap        map[string]*list.Element

     

       52
       49
        
       	stats           CacheStats

     

       53
       50
        
       	stopCleanup     chan struct{}

     

       51
       51
       +
       	logger          *slog.Logger

     

       54
       52
        
       }

     

       55
       53
        
       

     

       56
       54
        
       type lruItem struct {

     

       57
       55
        
       	key   string

     

       58
       58
       -
       	value *CacheEntry

     

       56
       56
       +
       	entry *CacheEntry

     

       59
       57
        
       }

     

       60
       58
        
       

     

       61
       59
        
       type Cache interface {

     

       62
       60
        
       	Get(key string) (*CacheEntry, bool)

     

       63
       61
        
       	Set(key string, entry *CacheEntry)

     

       64
       64
       -
       	GetStats() *CacheStats

     

       62
       62
       +
       	GetStats() CacheStats

     

       65
       63
        
       	Stop()

     

       66
       64
        
       }

     

       67
       65
        
       

     

       68
       68
       -
       func NewLRUCache(maxSize int, cleanupInterval time.Duration) *LRUCache {

     

       66
       66
       +
       func NewLRUCache(maxSize int, cleanupInterval time.Duration, logger *slog.Logger) *LRUCache {

     

       69
       67
        
       	if maxSize <= 0 {

     

       70
       68
        
       		maxSize = 5000

     

       71
       69
        
       	}

     
···

       78
       76
        
       		maxSize:         maxSize,

     

       79
       77
        
       		cleanupInterval: cleanupInterval,

     

       80
       78
        
       		lruList:         list.New(),

     

       81
       81
       -
       		cacheMap:        make(map[string]*list.Element),

     

       79
       79
       +
       		cacheMap:        make(map[string]*list.Element, maxSize),

     

       82
       80
        
       		stopCleanup:     make(chan struct{}),

     

       81
       81
       +
       		logger:          logger.With("component", "cache"),

     

       83
       82
        
       	}

     

       84
       83
        
       

     

       84
       84
       +
       	cache.logger.Info("starting LRU cache", "max_size", maxSize, "cleanup_interval", cleanupInterval)

     

       85
       85
        
       	go cache.periodicCleanup()

     

       86
       86
        
       	return cache

     

       87
       87
        
       }

     

       88
       88
        
       

     

       89
       89
       +
       func GenerateCacheKey(q magna.Question) string {

     

       90
       90
       +
       	name := strings.ToLower(q.QName)

     

       91
       91
       +
       	if !strings.HasSuffix(name, ".") {

     

       92
       92
       +
       		name += "."

     

       93
       93
       +
       	}

     

       94
       94
       +
       	return fmt.Sprintf("%s:%s:%s", name, q.QType.String(), q.QClass.String())

     

       95
       95
       +
       }

     

       96
       96
       +
       

     

       89
       97
        
       func (c *LRUCache) Get(key string) (*CacheEntry, bool) {

     

       90
       90
       -
       	c.stats.TotalQueries.Add(1)

     

       91
       91
       -
       

     

       92
       98
        
       	c.mu.RLock()

     

       93
       99
        
       	element, exists := c.cacheMap[key]

     

       94
       100
        
       	c.mu.RUnlock()

     

       95
       101
        
       

     

       96
       102
        
       	if !exists {

     

       97
       97
       -
       		c.stats.CacheMisses.Add(1)

     

       103
       103
       +
       		c.stats.Misses.Add(1)

     

       98
       104
        
       		return nil, false

     

       99
       105
        
       	}

     

       100
       106
        
       

     
···

       102
       108
        
       	defer c.mu.Unlock()

     

       103
       109
        
       

     

       104
       110
        
       	element, exists = c.cacheMap[key]

     

       105
       105
       -
       	// element might have been evicted or cleaned between RUnlock and Lock

     

       106
       111
        
       	if !exists {

     

       112
       112
       +
       		c.stats.Misses.Add(1)

     

       107
       113
        
       		return nil, false

     

       108
       114
        
       	}

     

       109
       115
        
       

     

       110
       116
        
       	item := element.Value.(*lruItem)

     

       111
       111
       -
       	entry := item.value

     

       117
       117
       +
       	entry := item.entry

     

       112
       118
        
       

     

       113
       113
       -
       	if time.Now().After(entry.ExpireAt) {

     

       119
       119
       +
       	now := time.Now()

     

       120
       120
       +
       	if now.After(entry.ExpireAt) {

     

       121
       121
       +
       		c.stats.Misses.Add(1)

     

       122
       122
       +
       		c.stats.Expired.Add(1)

     

       114
       123
        
       		c.removeItem(element)

     

       115
       115
       -
       		c.stats.CacheMisses.Add(1)

     

       116
       124
        
       		return nil, false

     

       117
       125
        
       	}

     

       118
       126
        
       

     

       127
       127
       +
       	c.stats.Hits.Add(1)

     

       119
       128
        
       	if entry.IsNegative {

     

       120
       129
        
       		c.stats.NegativeHits.Add(1)

     

       121
       130
        
       	} else {

     

       122
       131
        
       		c.stats.PositiveHits.Add(1)

     

       123
       132
        
       	}

     

       124
       124
       -
       	c.stats.CacheHits.Add(1)

     

       133
       133
       +
       

     

       134
       134
       +
       	c.lruList.MoveToFront(element)

     

       135
       135
       +
       	respEntry := c.adjustTTLs(entry, now)

     

       125
       136
        
       

     

       126
       126
       -
       	return entry, true

     

       137
       137
       +
       	return respEntry, true

     

       138
       138
       +
       }

     

       139
       139
       +
       

     

       140
       140
       +
       func (c *LRUCache) adjustTTLs(original *CacheEntry, now time.Time) *CacheEntry {

     

       141
       141
       +
       	adjustedEntry := &CacheEntry{

     

       142
       142
       +
       		Key:        original.Key,

     

       143
       143
       +
       		RCode:      original.RCode,

     

       144
       144
       +
       		IsNegative: original.IsNegative,

     

       145
       145
       +
       		CacheTime:  original.CacheTime,

     

       146
       146
       +
       		ExpireAt:   original.ExpireAt,

     

       147
       147
       +
       		Answer:     make([]magna.ResourceRecord, len(original.Answer)),

     

       148
       148
       +
       		Authority:  make([]magna.ResourceRecord, len(original.Authority)),

     

       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
       +
       	}

     

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

     

       157
       157
       +
       

     

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

     

       159
       159
       +
       		for i, rr := range src {

     

       160
       160
       +
       			(*dest)[i] = rr

     

       161
       161
       +
       			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
       +
       			}

     

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

     

       167
       167
       +
       

     

       168
       168
       +
       			finalTTL := min(rrRemainingTTL, remainingTTL)

     

       169
       169
       +
       			(*dest)[i].TTL = finalTTL

     

       170
       170
       +
       		}

     

       171
       171
       +
       	}

     

       172
       172
       +
       

     

       173
       173
       +
       	copyAndAdjust(&adjustedEntry.Answer, original.Answer)

     

       174
       174
       +
       	copyAndAdjust(&adjustedEntry.Authority, original.Authority)

     

       175
       175
       +
       	copyAndAdjust(&adjustedEntry.Additional, original.Additional)

     

       176
       176
       +
       

     

       177
       177
       +
       	return adjustedEntry

     

       127
       178
        
       }

     

       128
       179
        
       

     

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

     

       130
       130
       -
       	if time.Now().After(entry.ExpireAt) {

     

       181
       181
       +
       	if entry == nil {

     

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

     

       183
       183
       +
       		return

     

       184
       184
       +
       	}

     

       185
       185
       +
       

     

       186
       186
       +
       	if entry.ExpireAt.Before(time.Now().Add(1 * time.Second)) {

     

       131
       187
        
       		return

     

       132
       188
        
       	}

     

       133
       189
        
       

     
···

       135
       191
        
       	defer c.mu.Unlock()

     

       136
       192
        
       

     

       137
       193
        
       	if element, exists := c.cacheMap[key]; exists {

     

       138
       138
       -
       		element.Value.(*lruItem).value = entry

     

       139
       194
        
       		c.lruList.MoveToFront(element)

     

       195
       195
       +
       		element.Value.(*lruItem).entry = entry

     

       140
       196
        
       		return

     

       141
       197
        
       	}

     

       142
       198
        
       

     

       143
       199
        
       	newItem := &lruItem{

     

       144
       200
        
       		key:   key,

     

       145
       145
       -
       		value: entry,

     

       201
       201
       +
       		entry: entry,

     

       146
       202
        
       	}

     

       147
       203
        
       	element := c.lruList.PushFront(newItem)

     

       148
       204
        
       	c.cacheMap[key] = element

     

       149
       149
       -
       	currentSize := int64(c.lruList.Len())

     

       150
       150
       -
       	c.stats.Size.Store(currentSize)

     

       205
       205
       +
       	c.stats.Size.Store(int64(c.lruList.Len()))

     

       151
       206
        
       

     

       152
       207
        
       	for int64(c.lruList.Len()) > int64(c.maxSize) {

     

       153
       208
        
       		c.evictLRU()

     
···

       169
       224
        
       	c.stats.Size.Store(int64(c.lruList.Len()))

     

       170
       225
        
       }

     

       171
       226
        
       

     

       172
       172
       -
       func (c *LRUCache) GetStats() *CacheStats {

     

       173
       173
       -
       	statsSnapshot := &CacheStats{}

     

       174
       174
       -
       	statsSnapshot.TotalQueries.Store(c.stats.TotalQueries.Load())

     

       175
       175
       -
       	statsSnapshot.CacheHits.Store(c.stats.CacheHits.Load())

     

       176
       176
       -
       	statsSnapshot.CacheMisses.Store(c.stats.CacheMisses.Load())

     

       177
       177
       -
       	statsSnapshot.NegativeHits.Store(c.stats.NegativeHits.Load())

     

       178
       178
       -
       	statsSnapshot.PositiveHits.Store(c.stats.PositiveHits.Load())

     

       227
       227
       +
       func (c *LRUCache) GetStats() CacheStats {

     

       228
       228
       +
       	statsSnapshot := CacheStats{}

     

       229
       229
       +
       	statsSnapshot.Hits.Store(c.stats.Hits.Load())

     

       230
       230
       +
       	statsSnapshot.Misses.Store(c.stats.Misses.Load())

     

       179
       231
        
       	statsSnapshot.Evictions.Store(c.stats.Evictions.Load())

     

       232
       232
       +
       	statsSnapshot.Expired.Store(c.stats.Expired.Load())

     

       180
       233
        
       	statsSnapshot.Size.Store(c.stats.Size.Load())

     

       234
       234
       +
       	statsSnapshot.NegativeHits.Store(c.stats.NegativeHits.Load())

     

       235
       235
       +
       	statsSnapshot.PositiveHits.Store(c.stats.PositiveHits.Load())

     

       181
       236
        
       	return statsSnapshot

     

       182
       237
        
       }

     

       183
       238
        
       

     
···

       200
       255
        
       	defer c.mu.Unlock()

     

       201
       256
        
       

     

       202
       257
        
       	now := time.Now()

     

       258
       258
       +
       	cleanedCount := 0

     

       203
       259
        
       	element := c.lruList.Back()

     

       260
       260
       +
       

     

       204
       261
        
       	for element != nil {

     

       205
       262
        
       		item := element.Value.(*lruItem)

     

       206
       206
       -
       		nextElement := element.Prev()

     

       263
       263
       +
       		prevElement := element.Prev()

     

       207
       264
        
       

     

       208
       208
       -
       		if now.After(item.value.ExpireAt) {

     

       265
       265
       +
       		if now.After(item.entry.ExpireAt) {

     

       209
       266
        
       			c.removeItem(element)

     

       267
       267
       +
       			c.stats.Expired.Add(1)

     

       268
       268
       +
       			cleanedCount++

     

       210
       269
        
       		}

     

       211
       270
        
       

     

       212
       212
       -
       		element = nextElement

     

       271
       271
       +
       		element = prevElement

     

       272
       272
       +
       	}

     

       273
       273
       +
       	if cleanedCount > 0 {

     

       274
       274
       +
       		c.logger.Info("Cache cleanup finished", "items_removed", cleanedCount, "current_size", c.stats.Size.Load())

     

       213
       275
        
       	}

     

       214
       276
        
       }

     

       215
       277
        
       

     
···

       217
       279
        
       	close(c.stopCleanup)

     

       218
       280
        
       }

     

       219
       281
        
       

     

       220
       220
       -
       func getMinTTL(records []magna.ResourceRecord) uint32 {

     

       221
       221
       -
       	if len(records) == 0 {

     

       222
       222
       -
       		return 60

     

       223
       223
       -
       	}

     

       224
       224
       -
       	minTTL := records[0].TTL

     

       225
       225
       -
       	for _, rr := range records[1:] {

     

       226
       226
       -
       		if rr.TTL == 0 {

     

       227
       227
       -
       			return 60

     

       228
       228
       -
       		}

     

       229
       229
       -
       		if rr.TTL < minTTL {

     

       230
       230
       -
       			minTTL = rr.TTL

     

       231
       231
       -
       		}

     

       232
       232
       -
       	}

     

       233
       233
       -
       	if minTTL < 60 {

     

       234
       234
       -
       		return 60

     

       235
       235
       -
       	}

     

       236
       236
       -
       	return minTTL

     

       237
       237
       -
       }

     

       238
       238
       -
       

     

       239
       239
       -
       func CreateCacheEntry(msg *magna.Message, zone string) *CacheEntry {

     

       282
       282
       +
       func CreateCacheEntry(query magna.Question, response *magna.Message) *CacheEntry {

     

       240
       283
        
       	now := time.Now()

     

       241
       284
        
       	entry := &CacheEntry{

     

       242
       242
       -
       		IsNegative: msg.Header.RCode == magna.NXDOMAIN,

     

       285
       285
       +
       		Key:        GenerateCacheKey(query),

     

       286
       286
       +
       		RCode:      response.Header.RCode,

     

       287
       287
       +
       		Answer:     response.Answer,

     

       288
       288
       +
       		Authority:  response.Authority,

     

       289
       289
       +
       		Additional: response.Additional,

     

       290
       290
       +
       		CacheTime:  now,

     

       243
       291
        
       	}

     

       244
       292
        
       

     

       245
       245
       -
       	if entry.IsNegative {

     

       246
       246
       -
       		var soaTTL uint32

     

       247
       247
       -
       		for _, auth := range msg.Authority {

     

       248
       248
       -
       			if auth.RType == magna.SOAType {

     

       249
       249
       -
       				soa, ok := auth.RData.(*magna.SOA)

     

       250
       250
       -
       				if ok {

     

       251
       251
       -
       					soaTTL = soa.Minimum

     

       252
       252
       -
       					if auth.TTL < soaTTL {

     

       253
       253
       -
       						soaTTL = auth.TTL

     

       293
       293
       +
       	var minTTL time.Duration = -1

     

       294
       294
       +
       

     

       295
       295
       +
       	if response.Header.RCode == magna.NXDOMAIN || (response.Header.RCode == magna.NOERROR && len(response.Answer) == 0) {

     

       296
       296
       +
       		entry.IsNegative = true

     

       297
       297
       +
       		negativeTTL := defaultNegativeTTL

     

       298
       298
       +
       

     

       299
       299
       +
       		for _, rr := range response.Authority {

     

       300
       300
       +
       			if rr.RType == magna.SOAType {

     

       301
       301
       +
       				soa, ok := rr.RData.(*magna.SOA)

     

       302
       302
       +
       				if ok && soa != nil {

     

       303
       303
       +
       					ttl := time.Duration(rr.TTL) * time.Second

     

       304
       304
       +
       					minimum := time.Duration(soa.Minimum) * time.Second

     

       305
       305
       +
       

     

       306
       306
       +
       					if minimum > maxNegativeTTL {

     

       307
       307
       +
       						negativeTTL = ttl

     

       308
       308
       +
       					} else {

     

       309
       309
       +
       						negativeTTL = minDuration(ttl, minimum)

     

       254
       310
        
       					}

     

       311
       311
       +
       

     

       255
       312
        
       				} else {

     

       256
       256
       -
       					soaTTL = auth.TTL

     

       313
       313
       +
       					negativeTTL = time.Duration(rr.TTL) * time.Second

     

       257
       314
        
       				}

     

       258
       315
        
       				break

     

       259
       316
        
       			}

     

       260
       317
        
       		}

     

       261
       318
        
       

     

       262
       262
       -
       		if soaTTL == 0 {

     

       263
       263
       -
       			soaTTL = 900

     

       264
       264
       -
       		}

     

       265
       265
       -
       

     

       266
       266
       -
       		if soaTTL < 60 {

     

       267
       267
       -
       			soaTTL = 60

     

       319
       319
       +
       		if negativeTTL < 0 {

     

       320
       320
       +
       			negativeTTL = defaultNegativeTTL

     

       268
       321
        
       		}

     

       269
       269
       -
       

     

       270
       270
       -
       		entry.NegativeTTL = time.Duration(soaTTL) * time.Second

     

       271
       271
       -
       		entry.ExpireAt = now.Add(entry.NegativeTTL)

     

       272
       272
       -
       

     

       273
       273
       -
       		entry.Authority = make([]CachedResourceRecord, len(msg.Authority))

     

       274
       274
       -
       		for i, rr := range msg.Authority {

     

       275
       275
       -
       			recordTTL := rr.TTL

     

       276
       276
       -
       			if recordTTL < 60 {

     

       277
       277
       -
       				recordTTL = 60

     

       278
       278
       -
       			}

     

       279
       279
       -
       			rule := determineBailiwickRule(zone, rr.Name)

     

       280
       280
       -
       			entry.Authority[i] = CachedResourceRecord{

     

       281
       281
       -
       				Record:        rr,

     

       282
       282
       -
       				ExpireAt:      now.Add(time.Duration(recordTTL) * time.Second),

     

       283
       283
       -
       				BailiwickRule: rule,

     

       284
       284
       -
       			}

     

       322
       322
       +
       		if negativeTTL > maxNegativeTTL {

     

       323
       323
       +
       			negativeTTL = maxNegativeTTL

     

       285
       324
        
       		}

     

       286
       325
        
       

     

       287
       287
       -
       	} else {

     

       288
       288
       -
       		minAnswerTTL := getMinTTL(msg.Answer)

     

       326
       326
       +
       		minTTL = negativeTTL

     

       327
       327
       +
       	} else if response.Header.RCode == magna.NOERROR {

     

       328
       328
       +
       		entry.IsNegative = false

     

       289
       329
        
       

     

       290
       290
       -
       		if len(msg.Answer) == 0 && len(msg.Authority) > 0 && msg.Authority[0].RType == magna.SOAType {

     

       330
       330
       +
       		if len(response.Answer) > 0 {

     

       331
       331
       +
       			minTTL = time.Duration(response.Answer[0].TTL) * time.Second

     

       291
       332
        
       

     

       292
       292
       -
       			soa, ok := msg.Authority[0].RData.(*magna.SOA)

     

       293
       293
       -
       			soaTTL := uint32(900)

     

       294
       294
       -
       			if ok {

     

       295
       295
       -
       				soaTTL = soa.Minimum

     

       296
       296
       -
       				if msg.Authority[0].TTL < soaTTL {

     

       297
       297
       -
       					soaTTL = msg.Authority[0].TTL

     

       333
       333
       +
       			for _, rr := range response.Answer[1:] {

     

       334
       334
       +
       				currentTTL := time.Duration(rr.TTL) * time.Second

     

       335
       335
       +
       				if currentTTL < minTTL {

     

       336
       336
       +
       					minTTL = currentTTL

     

       298
       337
        
       				}

     

       299
       299
       -
       			} else {

     

       300
       300
       -
       				soaTTL = msg.Authority[0].TTL

     

       301
       338
        
       			}

     

       302
       302
       -
       			if soaTTL < 60 {

     

       303
       303
       -
       				soaTTL = 60

     

       304
       304
       -
       			}

     

       305
       305
       -
       

     

       339
       339
       +
       		} else {

     

       306
       340
        
       			entry.IsNegative = true

     

       307
       307
       -
       			entry.NegativeTTL = time.Duration(soaTTL) * time.Second

     

       308
       308
       -
       			entry.ExpireAt = now.Add(entry.NegativeTTL)

     

       341
       341
       +
       			negativeTTL := defaultNegativeTTL

     

       342
       342
       +
       			for _, rr := range response.Authority {

     

       343
       343
       +
       				if rr.RType == magna.SOAType {

     

       344
       344
       +
       					soa, ok := rr.RData.(*magna.SOA)

     

       345
       345
       +
       					if ok && soa != nil {

     

       346
       346
       +
       						ttl := time.Duration(rr.TTL) * time.Second

     

       347
       347
       +
       						minimum := time.Duration(soa.Minimum) * time.Second

     

       348
       348
       +
       						if minimum > maxNegativeTTL {

     

       349
       349
       +
       							negativeTTL = ttl

     

       350
       350
       +
       						} else {

     

       351
       351
       +
       							negativeTTL = minDuration(ttl, minimum)

     

       352
       352
       +
       						}

     

       309
       353
        
       

     

       310
       310
       -
       			entry.Authority = make([]CachedResourceRecord, len(msg.Authority))

     

       311
       311
       -
       			for i, rr := range msg.Authority {

     

       312
       312
       -
       				recordTTL := rr.TTL

     

       313
       313
       -
       				if recordTTL < 60 {

     

       314
       314
       -
       					recordTTL = 60

     

       354
       354
       +
       					} else {

     

       355
       355
       +
       						negativeTTL = time.Duration(rr.TTL) * time.Second

     

       356
       356
       +
       					}

     

       357
       357
       +
       					break

     

       315
       358
        
       				}

     

       316
       316
       -
       				rule := determineBailiwickRule(zone, rr.Name)

     

       317
       317
       -
       				entry.Authority[i] = CachedResourceRecord{Record: rr, ExpireAt: now.Add(time.Duration(recordTTL) * time.Second), BailiwickRule: rule}

     

       359
       359
       +
       			}

     

       360
       360
       +
       			if negativeTTL < 0 {

     

       361
       361
       +
       				negativeTTL = defaultNegativeTTL

     

       318
       362
        
       			}

     

       319
       319
       -
       

     

       320
       320
       -
       		} else {

     

       321
       321
       -
       			entry.ExpireAt = now.Add(time.Duration(minAnswerTTL) * time.Second)

     

       322
       322
       -
       

     

       323
       323
       -
       			entry.Answer = make([]CachedResourceRecord, len(msg.Answer))

     

       324
       324
       -
       			for i, rr := range msg.Answer {

     

       325
       325
       -
       				recordTTL := rr.TTL

     

       326
       326
       -
       				if recordTTL < 60 {

     

       327
       327
       -
       					recordTTL = 60

     

       328
       328
       -
       				}

     

       329
       329
       -
       				rule := determineBailiwickRule(zone, rr.Name)

     

       330
       330
       -
       				entry.Answer[i] = CachedResourceRecord{Record: rr, ExpireAt: now.Add(time.Duration(recordTTL) * time.Second), BailiwickRule: rule}

     

       363
       363
       +
       			if negativeTTL > maxNegativeTTL {

     

       364
       364
       +
       				negativeTTL = maxNegativeTTL

     

       331
       365
        
       			}

     

       366
       366
       +
       			minTTL = negativeTTL

     

       367
       367
       +
       		}

     

       332
       368
        
       

     

       333
       333
       -
       			entry.Authority = make([]CachedResourceRecord, len(msg.Authority))

     

       334
       334
       -
       			for i, rr := range msg.Authority {

     

       335
       335
       -
       				recordTTL := rr.TTL

     

       336
       336
       -
       				if recordTTL < 60 {

     

       337
       337
       -
       					recordTTL = 60

     

       338
       338
       -
       				}

     

       339
       339
       -
       				rule := determineBailiwickRule(zone, rr.Name)

     

       340
       340
       -
       				entry.Authority[i] = CachedResourceRecord{Record: rr, ExpireAt: now.Add(time.Duration(recordTTL) * time.Second), BailiwickRule: rule}

     

       341
       341
       -
       			}

     

       369
       369
       +
       	} else {

     

       370
       370
       +
       		return nil

     

       371
       371
       +
       	}

     

       342
       372
        
       

     

       343
       343
       -
       			entry.Additional = make([]CachedResourceRecord, len(msg.Additional))

     

       344
       344
       -
       			for i, rr := range msg.Additional {

     

       345
       345
       -
       				recordTTL := rr.TTL

     

       346
       346
       -
       				if recordTTL < 60 {

     

       347
       347
       -
       					recordTTL = 60

     

       348
       348
       -
       				}

     

       349
       349
       -
       				rule := determineBailiwickRule(zone, rr.Name)

     

       350
       350
       -
       				entry.Additional[i] = CachedResourceRecord{Record: rr, ExpireAt: now.Add(time.Duration(recordTTL) * time.Second), BailiwickRule: rule}

     

       351
       351
       -
       			}

     

       352
       352
       -
       		}

     

       373
       373
       +
       	if minTTL < 0 {

     

       374
       374
       +
       		return nil

     

       353
       375
        
       	}

     

       354
       376
        
       

     

       377
       377
       +
       	entry.ExpireAt = now.Add(minTTL)

     

       355
       378
        
       	return entry

     

       356
       379
        
       }

     

       357
       380
        
       

     

       358
       358
       -
       func isSubdomainOf(a, b string) bool {

     

       359
       359
       -
       	if a == b {

     

       360
       360
       -
       		return true

     

       381
       381
       +
       func minDuration(a, b time.Duration) time.Duration {

     

       382
       382
       +
       	if a < b {

     

       383
       383
       +
       		return a

     

       361
       384
        
       	}

     

       362
       362
       -
       

     

       363
       363
       -
       	return strings.HasSuffix(b, "."+a)

     

       385
       385
       +
       	return b

     

       364
       386
        
       }

     

       365
       387
        
       

     

       366
       366
       -
       func determineBailiwickRule(zone, name string) BailiwickRule {

     

       367
       367
       -
       	if zone == name {

     

       368
       368
       -
       		return BailiwickSame

     

       388
       388
       +
       func min(a, b uint32) uint32 {

     

       389
       389
       +
       	if a < b {

     

       390
       390
       +
       		return a

     

       369
       391
        
       	}

     

       370
       370
       -
       

     

       371
       371
       -
       	if isSubdomainOf(zone, name) {

     

       372
       372
       -
       		return BailiwickChild

     

       373
       373
       -
       	}

     

       374
       374
       -
       

     

       375
       375
       -
       	return BailiwickOutside

     

       392
       392
       +
       	return b

     

       376
       393
        
       }

+122 -14

pkg/dns/resolve.go

···

       36
       36
        
       	}

     

       37
       37
        
       

     

       38
       38
        
       	question := r.Message.Question[0]

     

       39
       39
       -
       	records, err := h.resolveQuestion(r.Context, question, h.RootServers)

     

       40
       39
        
       

     

       41
       40
        
       	msg := r.Message.CreateReply(r.Message)

     

       42
       41
        
       	msg.Header.RA = true

     

       43
       42
        
       

     

       43
       43
       +
       	if h.Cache != nil {

     

       44
       44
       +
       		cacheKey := GenerateCacheKey(question)

     

       45
       45
       +
       		if entry, found := h.Cache.Get(cacheKey); found {

     

       46
       46
       +
       			h.Logger.Debug("cache hit", "question", question.QName, "type", question.QType)

     

       47
       47
       +
       

     

       48
       48
       +
       			r.Context = setCacheHit(r.Context)

     

       49
       49
       +
       

     

       50
       50
       +
       			msg.Header.RCode = entry.RCode

     

       51
       51
       +
       			msg.Answer = entry.Answer

     

       52
       52
       +
       			msg.Authority = entry.Authority

     

       53
       53
       +
       			msg.Additional = entry.Additional

     

       54
       54
       +
       			msg.Header.ANCount = uint16(len(msg.Answer))

     

       55
       55
       +
       			msg.Header.NSCount = uint16(len(msg.Authority))

     

       56
       56
       +
       			msg.Header.ARCount = uint16(len(msg.Additional))

     

       57
       57
       +
       

     

       58
       58
       +
       			w.WriteMsg(msg)

     

       59
       59
       +
       			return

     

       60
       60
       +
       		} else {

     

       61
       61
       +
       			r.Context = setCacheMiss(r.Context)

     

       62
       62
       +
       		}

     

       63
       63
       +
       	}

     

       64
       64
       +
       

     

       65
       65
       +
       	records, err := h.resolveQuestion(r.Context, question, h.RootServers)

     

       66
       66
       +
       

     

       44
       67
        
       	if err == errNXDOMAIN {

     

       45
       68
        
       		msg = msg.SetRCode(magna.NXDOMAIN)

     

       46
       69
        
       		msg.Authority = records

     
···

       48
       71
        
       		msg.Header.ANCount = 0

     

       49
       72
        
       		msg.Answer = nil

     

       50
       73
        
       	} else if err != nil {

     

       51
       51
       -
       		msg.SetRCode(magna.SERVFAIL)

     

       74
       74
       +
       		msg = msg.SetRCode(magna.SERVFAIL)

     

       52
       75
        
       	} else {

     

       53
       76
        
       		msg.Answer = records

     

       54
       77
        
       		msg.Header.ANCount = uint16(len(records))

     

       55
       78
        
       		msg = msg.SetRCode(magna.NOERROR)

     

       56
       79
        
       	}

     

       57
       80
        
       

     

       81
       81
       +
       	if h.Cache != nil {

     

       82
       82
       +
       		if err == nil || err == errNXDOMAIN {

     

       83
       83
       +
       			entry := CreateCacheEntry(question, msg)

     

       84
       84
       +
       			if entry != nil {

     

       85
       85
       +
       				h.Cache.Set(GenerateCacheKey(question), entry)

     

       86
       86
       +
       				h.Logger.Debug("cached response",

     

       87
       87
       +
       					"question", question.QName,

     

       88
       88
       +
       					"type", question.QType.String(),

     

       89
       89
       +
       					"negative", entry.IsNegative,

     

       90
       90
       +
       					"rcode", entry.RCode.String(),

     

       91
       91
       +
       					"expires", entry.ExpireAt.Sub(time.Now()).String())

     

       92
       92
       +
       			}

     

       93
       93
       +
       		}

     

       94
       94
       +
       	}

     

       95
       95
       +
       

     

       58
       96
        
       	w.WriteMsg(msg)

     

       59
       97
        
       }

     

       60
       98
        
       

     

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

     

       100
       100
       +
       	if h.Cache != nil {

     

       101
       101
       +
       		cacheKey := GenerateCacheKey(question)

     

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

     

       103
       103
       +
       			h.Logger.Debug("cache hit during recursion", "question", question.QName, "type", question.QType)

     

       104
       104
       +
       

     

       105
       105
       +
       			if entry.RCode == magna.NXDOMAIN {

     

       106
       106
       +
       				return entry.Authority, errNXDOMAIN

     

       107
       107
       +
       			}

     

       108
       108
       +
       

     

       109
       109
       +
       			if entry.IsNegative {

     

       110
       110
       +
       				return []magna.ResourceRecord{}, nil

     

       111
       111
       +
       			}

     

       112
       112
       +
       

     

       113
       113
       +
       			return entry.Answer, nil

     

       114
       114
       +
       		}

     

       115
       115
       +
       	}

     

       116
       116
       +
       

     

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

     

       63
       118
        
       	defer cancel()

     

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

     
···

       71
       126
        
       		select {

     

       72
       127
        
       		case res := <-ch:

     

       73
       128
        
       			if res.Error != nil {

     

       74
       74
       -
       				slog.Warn("error", "question", question, "server", res.Server, "error", res.Error)

     

       75
       75
       -
       				break

     

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

     

       130
       130
       +
       				continue

     

       76
       131
        
       			}

     

       77
       132
        
       

     

       78
       133
        
       			msg := res.MSG

     

       134
       134
       +
       

     

       135
       135
       +
       			fullMsg := magna.Message{

     

       136
       136
       +
       				Header:     msg.Header,

     

       137
       137
       +
       				Question:   []magna.Question{question},

     

       138
       138
       +
       				Answer:     msg.Answer,

     

       139
       139
       +
       				Authority:  msg.Authority,

     

       140
       140
       +
       				Additional: msg.Additional,

     

       141
       141
       +
       			}

     

       142
       142
       +
       

     

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

     

       144
       144
       +
       				if h.Cache != nil {

     

       145
       145
       +
       					entry := CreateCacheEntry(question, &fullMsg)

     

       146
       146
       +
       					if entry != nil {

     

       147
       147
       +
       						h.Cache.Set(GenerateCacheKey(question), entry)

     

       148
       148
       +
       					}

     

       149
       149
       +
       				}

     

       80
       150
        
       				return msg.Authority, errNXDOMAIN

     

       81
       151
        
       			}

     

       82
       152
        
       

     

       83
       153
        
       			if msg.Header.ANCount > 0 {

     

       84
       154
        
       				if msg.Answer[0].RType == magna.CNAMEType {

     

       85
       85
       -
       					cname_answers, err := h.resolveQuestion(resolveCtx, magna.Question{QName: msg.Answer[0].RData.String(), QType: question.QType, QClass: question.QClass}, h.RootServers)

     

       155
       155
       +
       					if h.Cache != nil {

     

       156
       156
       +
       						partialEntry := CreateCacheEntry(question, &fullMsg)

     

       157
       157
       +
       						if partialEntry != nil {

     

       158
       158
       +
       							h.Cache.Set(GenerateCacheKey(question), partialEntry)

     

       159
       159
       +
       						}

     

       160
       160
       +
       					}

     

       161
       161
       +
       

     

       162
       162
       +
       					cnameTarget := msg.Answer[0].RData.String()

     

       163
       163
       +
       					h.Logger.Debug("following CNAME", "from", question.QName, "to", cnameTarget)

     

       164
       164
       +
       

     

       165
       165
       +
       					cname_answers, err := h.resolveQuestion(resolveCtx, magna.Question{

     

       166
       166
       +
       						QName:  cnameTarget,

     

       167
       167
       +
       						QType:  question.QType,

     

       168
       168
       +
       						QClass: question.QClass,

     

       169
       169
       +
       					}, h.RootServers)

     

       86
       170
        
       					if err != nil {

     

       87
       171
        
       						if err == errNXDOMAIN {

     

       88
       172
        
       							return cname_answers, errNXDOMAIN

     

       89
       173
        
       						}

     

       90
       90
       -
       						slog.Warn("error resolving CNAME", "cname", msg.Answer[0].RData.String(), "error", err)

     

       91
       91
       -
       

     

       174
       174
       +
       						h.Logger.Warn("error resolving CNAME", "cname", cnameTarget, "error", err)

     

       92
       175
        
       						continue

     

       93
       176
        
       					}

     

       94
       177
        
       

     

       95
       95
       -
       					msg.Answer = append(msg.Answer, cname_answers...)

     

       178
       178
       +
       					combinedAnswer := append([]magna.ResourceRecord{}, msg.Answer[0])

     

       179
       179
       +
       					combinedAnswer = append(combinedAnswer, cname_answers...)

     

       180
       180
       +
       

     

       181
       181
       +
       					fullMsg.Answer = combinedAnswer

     

       182
       182
       +
       					if h.Cache != nil {

     

       183
       183
       +
       						fullEntry := CreateCacheEntry(question, &fullMsg)

     

       184
       184
       +
       						if fullEntry != nil {

     

       185
       185
       +
       							h.Cache.Set(GenerateCacheKey(question), fullEntry)

     

       186
       186
       +
       						}

     

       187
       187
       +
       					}

     

       188
       188
       +
       

     

       189
       189
       +
       					return combinedAnswer, nil

     

       96
       190
        
       				}

     

       97
       191
        
       

     

       192
       192
       +
       				if h.Cache != nil {

     

       193
       193
       +
       					entry := CreateCacheEntry(question, &fullMsg)

     

       194
       194
       +
       					if entry != nil {

     

       195
       195
       +
       						h.Cache.Set(GenerateCacheKey(question), entry)

     

       196
       196
       +
       					}

     

       197
       197
       +
       				}

     

       98
       198
        
       				return msg.Answer, nil

     

       99
       199
        
       			}

     

       100
       200
        
       

     
···

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

     

       126
       226
        
       						nextServers = append(nextServers, ip)

     

       127
       227
        
       					} else {

     

       128
       128
       -
       						nsAnswers, err := h.resolveQuestion(resolveCtx, magna.Question{QName: nsName, QType: magna.AType, QClass: magna.IN}, h.RootServers)

     

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

     

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

     

       129
       230
        
       						if err != nil {

     

       130
       130
       -
       							slog.Warn("error resolving NS A record", "ns", nsName, "error", err)

     

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

     

       131
       232
        
       							continue

     

       132
       233
        
       						}

     

       133
       234
        
       

     
···

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

     

       144
       245
        
       				}

     

       145
       246
        
       

     

       146
       146
       -
       				slog.Warn("could not resolve any NS records for delegation", "question", question.QName)

     

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

     

       147
       248
        
       				continue

     

       148
       249
        
       			}

     

       149
       250
        
       

     

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

     

       252
       252
       +
       				if h.Cache != nil {

     

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

     

       254
       254
       +
       					if entry != nil {

     

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

     

       256
       256
       +
       					}

     

       257
       257
       +
       				}

     

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

     

       152
       259
        
       			}

     

       153
       260
        
       

     

       154
       154
       -
       			slog.Warn("unexpected response state", "question", question, "server", res.Server, "rcode", msg.Header.RCode)

     

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

     

       155
       262
        
       			continue

     

       263
       263
       +
       

     

       156
       264
        
       		case <-resolveCtx.Done():

     

       157
       157
       -
       			slog.Debug("resolution cancelled or timed out", "question", question)

     

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

     

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

     

       159
       267
        
       		}

     

       160
       268
        
       	}

     

       161
       269
        
       

     

       162
       162
       -
       	slog.Warn("all resolution paths failed", "question", question)

     

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

     

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

     

       164
       272
        
       }

     

       165
       273

+155 -38

pkg/dns/server.go

···

       28
       28
        
       }

     

       29
       29
        
       

     

       30
       30
        
       func GetCacheHit(ctx context.Context) bool {

     

       31
       31
       -
       	v := ctx.Value(cacheHitKey)

     

       32
       32
       -
       	if v == nil {

     

       31
       31
       +
       	if ctx == nil {

     

       33
       32
        
       		return false

     

       34
       33
        
       	}

     

       35
       35
       -
       

     

       36
       36
       -
       	return v.(bool)

     

       34
       34
       +
       	v := ctx.Value(cacheHitKey)

     

       35
       35
       +
       	hit, ok := v.(bool)

     

       36
       36
       +
       	return ok && hit

     

       37
       37
        
       }

     

       38
       38
        
       

     

       39
       39
        
       type Handler interface {

     
···

       61
       61
        
       	addr         *net.UDPAddr

     

       62
       62
        
       	logger       *slog.Logger

     

       63
       63
        
       	writeTimeout time.Duration

     

       64
       64
       +
       	udpSize      int

     

       64
       65
        
       }

     

       65
       66
        
       

     

       66
       67
        
       func (w *udpResponseWriter) WriteMsg(msg *magna.Message) {

     

       67
       68
        
       	ans, err := msg.Encode()

     

       68
       69
        
       	if err != nil {

     

       69
       70
        
       		w.logger.Warn("err encoding msg", "error", err)

     

       70
       70
       -
       		return

     

       71
       71
       +
       

     

       72
       72
       +
       		failMsg := msg.CreateReply(msg)

     

       73
       73
       +
       		failMsg = failMsg.SetRCode(magna.SERVFAIL)

     

       74
       74
       +
       		failMsg.Answer = nil

     

       75
       75
       +
       		failMsg.Authority = nil

     

       76
       76
       +
       		failMsg.Additional = nil

     

       77
       77
       +
       		failMsg.Header.ANCount = 0

     

       78
       78
       +
       		failMsg.Header.NSCount = 0

     

       79
       79
       +
       		failMsg.Header.ARCount = 0

     

       80
       80
       +
       		ans, _ = failMsg.Encode()

     

       81
       81
       +
       		if ans == nil {

     

       82
       82
       +
       			return

     

       83
       83
       +
       		}

     

       84
       84
       +
       	}

     

       85
       85
       +
       

     

       86
       86
       +
       	if len(ans) > w.udpSize {

     

       87
       87
       +
       		w.logger.Debug("Response exceeds UDP size, setting TC bit", "size", len(ans), "limit", w.udpSize, "client", w.addr)

     

       88
       88
       +
       

     

       89
       89
       +
       		tcMsg := msg.CreateReply(msg)

     

       90
       90
       +
       		tcMsg.Header.TC = true

     

       91
       91
       +
       		tcMsg.Answer = nil

     

       92
       92
       +
       		tcMsg.Authority = nil

     

       93
       93
       +
       		tcMsg.Additional = nil

     

       94
       94
       +
       		tcMsg.Header.ANCount = 0

     

       95
       95
       +
       		tcMsg.Header.NSCount = 0

     

       96
       96
       +
       		tcMsg.Header.ARCount = 0

     

       97
       97
       +
       		tcMsg.Header.RCode = magna.NOERROR

     

       98
       98
       +
       

     

       99
       99
       +
       		ans, err = tcMsg.Encode()

     

       100
       100
       +
       		if err != nil {

     

       101
       101
       +
       			w.logger.Error("Error encoding truncated UDP response", "error", err, "client", w.addr)

     

       102
       102
       +
       			return

     

       103
       103
       +
       		}

     

       104
       104
       +
       		if len(ans) > w.udpSize {

     

       105
       105
       +
       			w.logger.Warn("Truncated message still exceeds UDP size limit!", "size", len(ans), "limit", w.udpSize)

     

       106
       106
       +
       		}

     

       71
       107
        
       	}

     

       72
       108
        
       

     

       73
       73
       -
       	if len(ans) > 512 {

     

       74
       74
       -
       		ans[3] |= 1 << 6 // set the truncated bit

     

       109
       109
       +
       	if w.udpConn == nil || w.addr == nil {

     

       110
       110
       +
       		w.logger.Error("UDP response writer used incorrectly (nil conn or addr)")

     

       111
       111
       +
       		return

     

       75
       112
        
       	}

     

       76
       113
        
       

     

       77
       114
        
       	err = w.udpConn.SetWriteDeadline(time.Now().Add(w.writeTimeout))

     
···

       95
       132
        
       	ans, err := msg.Encode()

     

       96
       133
        
       	if err != nil {

     

       97
       134
        
       		w.logger.Warn("err encoding msg", "error", err)

     

       135
       135
       +
       

     

       136
       136
       +
       		failMsg := msg.CreateReply(msg)

     

       137
       137
       +
       		failMsg = failMsg.SetRCode(magna.SERVFAIL)

     

       138
       138
       +
       		failMsg.Answer = nil

     

       139
       139
       +
       		failMsg.Authority = nil

     

       140
       140
       +
       		failMsg.Additional = nil

     

       141
       141
       +
       		failMsg.Header.ANCount = 0

     

       142
       142
       +
       		failMsg.Header.NSCount = 0

     

       143
       143
       +
       		failMsg.Header.ARCount = 0

     

       144
       144
       +
       		ans, _ = failMsg.Encode()

     

       145
       145
       +
       		return

     

       146
       146
       +
       	}

     

       147
       147
       +
       

     

       148
       148
       +
       	if w.tcpConn == nil {

     

       149
       149
       +
       		w.logger.Error("TCP response writer used incorrectly (nil conn)")

     

       98
       150
        
       		return

     

       99
       151
        
       	}

     

       100
       152
        
       

     
···

       123
       175
        
       	ReadTimeout  time.Duration

     

       124
       176
        
       	WriteTimeout time.Duration

     

       125
       177
        
       	Logger       *slog.Logger

     

       126
       126
       -
       	Cache        Cache

     

       127
       178
        
       }

     

       128
       179
        
       

     

       129
       180
        
       func (srv *Server) ListenAndServe() error {

     

       181
       181
       +
       	if srv.Logger == nil {

     

       182
       182
       +
       		srv.Logger = slog.Default()

     

       183
       183
       +
       	}

     

       184
       184
       +
       

     

       185
       185
       +
       	if srv.UDPSize <= 0 {

     

       186
       186
       +
       		srv.UDPSize = 512

     

       187
       187
       +
       	}

     

       188
       188
       +
       	srv.Logger.Info("Starting DNS server", "address", srv.Address, "port", srv.Port, "udp_size", srv.UDPSize)

     

       189
       189
       +
       

     

       130
       190
        
       	var wg sync.WaitGroup

     

       131
       191
        
       	errChan := make(chan error, 2)

     

       132
       192
        
       

     
···

       135
       195
        
       	go func() {

     

       136
       196
        
       		defer wg.Done()

     

       137
       197
        
       		if err := srv.serveTCP(); err != nil {

     

       138
       138
       -
       			errChan <- fmt.Errorf("TCP server error: %w", err)

     

       198
       198
       +
       			err = fmt.Errorf("TCP server error: %w", err)

     

       199
       199
       +
       			srv.Logger.Error(err.Error())

     

       200
       200
       +
       			select {

     

       201
       201
       +
       			case errChan <- err:

     

       202
       202
       +
       			default:

     

       203
       203
       +
       				srv.Logger.Warn("Error channel full, discarding TCP error")

     

       204
       204
       +
       			}

     

       139
       205
        
       		}

     

       140
       206
        
       	}()

     

       141
       207
        
       

     

       142
       208
        
       	go func() {

     

       143
       209
        
       		defer wg.Done()

     

       144
       210
        
       		if err := srv.serveUDP(); err != nil {

     

       145
       145
       -
       			errChan <- fmt.Errorf("TCP server error: %w", err)

     

       211
       211
       +
       			err = fmt.Errorf("UDP server error: %w", err)

     

       212
       212
       +
       			srv.Logger.Error(err.Error())

     

       213
       213
       +
       			select {

     

       214
       214
       +
       			case errChan <- err:

     

       215
       215
       +
       			default:

     

       216
       216
       +
       				srv.Logger.Warn("Error channel full, discarding UDP error")

     

       217
       217
       +
       			}

     

       146
       218
        
       		}

     

       147
       219
        
       	}()

     

       148
       220
        
       

     
···

       151
       223
        
       		close(errChan)

     

       152
       224
        
       	}()

     

       153
       225
        
       

     

       154
       154
       -
       	for err := range errChan {

     

       155
       155
       -
       		return err

     

       156
       156
       -
       	}

     

       157
       157
       -
       

     

       158
       158
       -
       	return nil

     

       226
       226
       +
       	err := <-errChan

     

       227
       227
       +
       	return err

     

       159
       228
        
       }

     

       160
       229
        
       

     

       161
       230
        
       func (srv *Server) serveUDP() error {

     

       162
       162
       -
       	addr := net.UDPAddr{

     

       231
       231
       +
       	addr := &net.UDPAddr{

     

       163
       232
        
       		Port: srv.Port,

     

       164
       233
        
       		IP:   net.ParseIP(srv.Address),

     

       165
       234
        
       	}

     

       166
       166
       -
       	conn, err := net.ListenUDP("udp", &addr)

     

       235
       235
       +
       	conn, err := net.ListenUDP("udp", addr)

     

       167
       236
        
       	if err != nil {

     

       168
       168
       -
       		return err

     

       237
       237
       +
       		return fmt.Errorf("failed to listen on UDP %s:%d: %w", srv.Address, srv.Port, err)

     

       169
       238
        
       	}

     

       170
       239
        
       	defer conn.Close()

     

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

     

       171
       241
        
       

     

       172
       242
        
       	for {

     

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

     

       174
       244
        
       

     

       175
       175
       -
       		err := conn.SetReadDeadline(time.Now().Add(srv.ReadTimeout))

     

       176
       176
       -
       		if err != nil {

     

       177
       177
       -
       			return fmt.Errorf("error setting read deadline: %w", err)

     

       245
       245
       +
       		if srv.ReadTimeout > 0 {

     

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

     

       247
       247
       +
       			if err != nil {

     

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

     

       249
       249
       +
       			}

     

       178
       250
        
       		}

     

       179
       251
        
       

     

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

     

       181
       253
        
       		if err != nil {

     

       182
       182
       -
       			// skip logging timeout errors

     

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

     

       184
       255
        
       				continue

     

       185
       256
        
       			}

     

       186
       257
        
       

     

       187
       187
       -
       			srv.Logger.Warn(err.Error())

     

       258
       258
       +
       			if opErr, ok := err.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {

     

       259
       259
       +
       				srv.Logger.Info("UDP listener stopping: connection closed")

     

       260
       260
       +
       				return nil

     

       261
       261
       +
       			}

     

       262
       262
       +
       

     

       263
       263
       +
       			srv.Logger.Warn("UDP read error", "error", err)

     

       188
       264
        
       			continue

     

       189
       265
        
       		}

     

       190
       266
        
       

     
···

       198
       274
        
       		addr:         remoteAddr,

     

       199
       275
        
       		logger:       srv.Logger,

     

       200
       276
        
       		writeTimeout: srv.WriteTimeout,

     

       277
       277
       +
       		udpSize:      srv.UDPSize,

     

       201
       278
        
       	}

     

       202
       279
        
       

     

       203
       280
        
       	srv.handleQuery(query, w, remoteAddr)

     
···

       211
       288
        
       

     

       212
       289
        
       	listener, err := net.ListenTCP("tcp", &addr)

     

       213
       290
        
       	if err != nil {

     

       214
       214
       -
       		return err

     

       291
       291
       +
       		return fmt.Errorf("failed to listen on TCP %s:%d: %w", srv.Address, srv.Port, err)

     

       215
       292
        
       	}

     

       216
       293
        
       	defer listener.Close()

     

       294
       294
       +
       	srv.Logger.Info("TCP listener started", "address", listener.Addr())

     

       217
       295
        
       

     

       218
       296
        
       	for {

     

       219
       219
       -
       		conn, err := listener.Accept()

     

       297
       297
       +
       		conn, err := listener.AcceptTCP()

     

       220
       298
        
       		if err != nil {

     

       221
       221
       -
       			srv.Logger.Warn("tcp accept error:", "error", err)

     

       299
       299
       +
       			if opErr, ok := err.(*net.OpError); ok && opErr.Err.Error() == "use of closed network connection" {

     

       300
       300
       +
       				srv.Logger.Info("TCP listener stopping: listener closed")

     

       301
       301
       +
       				return nil

     

       302
       302
       +
       			}

     

       303
       303
       +
       			srv.Logger.Warn("TCP accept error", "error", err)

     

       222
       304
        
       			continue

     

       223
       305
        
       		}

     

       224
       306
        
       

     

       307
       307
       +
       		conn.SetKeepAlive(true)

     

       308
       308
       +
       		conn.SetKeepAlivePeriod(3 * time.Minute)

     

       225
       309
        
       		go srv.handleTCPQuery(conn)

     

       226
       310
        
       	}

     

       227
       311
        
       }

     
···

       229
       313
        
       func (srv *Server) handleTCPQuery(conn net.Conn) {

     

       230
       314
        
       	defer conn.Close()

     

       231
       315
        
       

     

       232
       232
       -
       	err := conn.SetReadDeadline(time.Now().Add(srv.ReadTimeout))

     

       233
       233
       -
       	if err != nil {

     

       234
       234
       -
       		srv.Logger.Error("error setting read deadline", "error", err)

     

       235
       235
       -
       		return

     

       316
       316
       +
       	if srv.ReadTimeout > 0 {

     

       317
       317
       +
       		err := conn.SetReadDeadline(time.Now().Add(srv.ReadTimeout))

     

       318
       318
       +
       		if err != nil {

     

       319
       319
       +
       			srv.Logger.Warn("Error setting TCP initial read deadline", "error", err, "client", conn.RemoteAddr())

     

       320
       320
       +
       			return

     

       321
       321
       +
       		}

     

       236
       322
        
       	}

     

       237
       323
        
       

     

       238
       324
        
       	sizeBuffer := make([]byte, 2)

     

       239
       239
       -
       	if _, err := io.ReadFull(conn, sizeBuffer); err != nil {

     

       240
       240
       -
       		srv.Logger.Warn("tcp error occurred", "error", err)

     

       325
       325
       +
       	_, err := io.ReadFull(conn, sizeBuffer)

     

       326
       326
       +
       	if err != nil {

     

       327
       327
       +
       		if err != io.EOF && err != io.ErrUnexpectedEOF {

     

       328
       328
       +
       			srv.Logger.Warn("TCP error reading message length", "error", err, "client", conn.RemoteAddr())

     

       329
       329
       +
       		}

     

       241
       330
        
       		return

     

       242
       331
        
       	}

     

       243
       332
        
       

     

       244
       333
        
       	size := binary.BigEndian.Uint16(sizeBuffer)

     

       334
       334
       +
       	if size == 0 {

     

       335
       335
       +
       		srv.Logger.Debug("TCP received zero-length message", "client", conn.RemoteAddr())

     

       336
       336
       +
       		return

     

       337
       337
       +
       	}

     

       338
       338
       +
       

     

       339
       339
       +
       	maxTCPMessageSize := 65535

     

       340
       340
       +
       	if size > uint16(maxTCPMessageSize) {

     

       341
       341
       +
       		srv.Logger.Warn("TCP message size exceeds limit", "size", size, "limit", maxTCPMessageSize, "client", conn.RemoteAddr())

     

       342
       342
       +
       		return

     

       343
       343
       +
       	}

     

       344
       344
       +
       

     

       245
       345
        
       	data := make([]byte, size)

     

       246
       246
       -
       	if _, err := io.ReadFull(conn, data); err != nil {

     

       247
       247
       -
       		srv.Logger.Warn("tcp error occurred", "error", err)

     

       346
       346
       +
       	if srv.ReadTimeout > 0 {

     

       347
       347
       +
       		err = conn.SetReadDeadline(time.Now().Add(srv.ReadTimeout))

     

       348
       348
       +
       		if err != nil {

     

       349
       349
       +
       			srv.Logger.Warn("Error setting TCP read deadline for body", "error", err, "client", conn.RemoteAddr())

     

       350
       350
       +
       			return

     

       351
       351
       +
       		}

     

       352
       352
       +
       	}

     

       353
       353
       +
       

     

       354
       354
       +
       	_, err = io.ReadFull(conn, data)

     

       355
       355
       +
       	if err != nil {

     

       356
       356
       +
       		srv.Logger.Warn("TCP error reading message body", "error", err, "client", conn.RemoteAddr())

     

       248
       357
        
       		return

     

       249
       358
        
       	}

     

       250
       359
        
       

     
···

       259
       368
        
       

     

       260
       369
        
       func (srv *Server) handleQuery(messageBuffer []byte, w ResponseWriter, remoteAddr net.Addr) {

     

       261
       370
        
       	var query magna.Message

     

       262
       262
       -
       	if err := query.Decode(messageBuffer); err != nil {

     

       263
       263
       -
       		srv.Logger.Warn("message decode error", "error", err)

     

       371
       371
       +
       	err := query.Decode(messageBuffer)

     

       372
       372
       +
       	if err != nil {

     

       373
       373
       +
       		srv.Logger.Warn("Message decode error", "error", err, "client", remoteAddr)

     

       374
       374
       +
       		// TODO: find better way to handle failed decode drop for now.

     

       264
       375
        
       		return

     

       265
       376
        
       	}

     

       266
       377
        
       

     

       267
       378
        
       	r := &Request{

     

       379
       379
       +
       		Context:    context.Background(),

     

       268
       380
        
       		Message:    &query,

     

       269
       381
        
       		RemoteAddr: remoteAddr,

     

       270
       382
        
       	}

     

       271
       383
        
       

     

       272
       272
       -
       	r.Context = context.WithValue(context.Background(), contextKey("request"), r)

     

       273
       273
       -
       

     

       384
       384
       +
       	if srv.Handler == nil {

     

       385
       385
       +
       		srv.Logger.Error("No DNS handler configured!")

     

       386
       386
       +
       		reply := query.CreateReply(&query)

     

       387
       387
       +
       		reply = reply.SetRCode(magna.SERVFAIL)

     

       388
       388
       +
       		w.WriteMsg(reply)

     

       389
       389
       +
       		return

     

       390
       390
       +
       	}

     

       274
       391
        
       	srv.Handler.ServeDNS(w, r)

     

       275
       392
        
       }

+137 -66

pkg/metrics/clickhouse.go

···

       2
       2
        
       

     

       3
       3
        
       import (

     

       4
       4
        
       	"database/sql"

     

       5
       5
       +
       	"errors"

     

       5
       6
        
       	"fmt"

     

       6
       7
        
       	"log/slog"

     

       7
       8
        
       	"sync"

     
···

       10
       11
        
       	_ "github.com/ClickHouse/clickhouse-go/v2"

     

       11
       12
        
       

     

       12
       13
        
       	"tangled.sh/seiso.moe/alky/pkg/config"

     

       13
       13
       -
       	"tangled.sh/seiso.moe/alky/pkg/dns"

     

       14
       14
        
       )

     

       15
       15
        
       

     

       16
       16
        
       type ClickHouseMetrics struct {

     
···

       20
       20
        
       	cacheBuffer []CacheMetric

     

       21
       21
        
       	mu          sync.Mutex

     

       22
       22
        
       	stopChan    chan struct{}

     

       23
       23
       +
       	wg          sync.WaitGroup

     

       23
       24
        
       	logger      *slog.Logger

     

       24
       25
        
       }

     

       25
       26
        
       

     
···

       37
       38
        
       

     

       38
       39
        
       type CacheMetric struct {

     

       39
       40
        
       	Timestamp    time.Time

     

       40
       40
       -
       	TotalQueries int64

     

       41
       41
        
       	CacheHits    int64

     

       42
       42
        
       	CacheMisses  int64

     

       43
       43
        
       	NegativeHits int64

     
···

       47
       47
        
       }

     

       48
       48
        
       

     

       49
       49
        
       func NewClickHouseMetrics(config *config.MetricsConfig, logger *slog.Logger) (*ClickHouseMetrics, error) {

     

       50
       50
       +
       	if logger == nil {

     

       51
       51
       +
       		logger = slog.Default()

     

       52
       52
       +
       	}

     

       53
       53
       +
       	logger.Info("Connecting to ClickHouse", "dsn", config.DSN)

     

       54
       54
       +
       

     

       50
       55
        
       	db, err := sql.Open("clickhouse", config.DSN)

     

       51
       56
        
       	if err != nil {

     

       52
       52
       -
       		return nil, fmt.Errorf("failed to connect to ClickHouse: %w", err)

     

       57
       57
       +
       		return nil, fmt.Errorf("failed to initialize ClickHouse driver: %w", err)

     

       53
       58
        
       	}

     

       54
       59
        
       

     

       60
       60
       +
       	if err := db.Ping(); err != nil {

     

       61
       61
       +
       		db.Close()

     

       62
       62
       +
       		return nil, fmt.Errorf("failed to connect to ClickHouse (%s): %w", config.DSN, err)

     

       63
       63
       +
       	}

     

       64
       64
       +
       	logger.Info("Successfully connected to ClickHouse")

     

       65
       65
       +
       

     

       55
       66
        
       	m := &ClickHouseMetrics{

     

       56
       67
        
       		db:          db,

     

       57
       68
        
       		config:      config,

     
···

       61
       72
        
       		logger:      logger,

     

       62
       73
        
       	}

     

       63
       74
        
       

     

       64
       64
       -
       	if err := m.changeTTL(); err != nil {

     

       65
       65
       -
       		db.Close()

     

       66
       66
       -
       		return nil, fmt.Errorf("failed to initialize tables: %w", err)

     

       75
       75
       +
       	if err := m.checkAndUpdateTTL(); err != nil {

     

       76
       76
       +
       		logger.Warn("Failed to check/update table TTLs, using existing TTL", "error", err)

     

       77
       77
       +
       	} else {

     

       78
       78
       +
       		logger.Info("Table TTLs verified/updated", "retention", config.RetentionPeriod.Duration)

     

       67
       79
        
       	}

     

       68
       80
        
       

     

       81
       81
       +
       	m.wg.Add(1)

     

       69
       82
        
       	go m.flushLoop()

     

       70
       83
        
       	return m, nil

     

       71
       84
        
       }

     
···

       76
       89
        
       

     

       77
       90
        
       	m.queryBuffer = append(m.queryBuffer, metric)

     

       78
       91
        
       	if len(m.queryBuffer) >= m.config.BatchSize {

     

       79
       79
       -
       		m.flush()

     

       92
       92
       +
       		m.flushQueriesLocked()

     

       80
       93
        
       	}

     

       81
       94
        
       }

     

       82
       95
        
       

     

       83
       83
       -
       func (m *ClickHouseMetrics) RecordCacheMetrics(metric CacheMetric) {

     

       96
       96
       +
       func (m *ClickHouseMetrics) RecordCacheStats(metric CacheMetric) {

     

       84
       97
        
       	m.mu.Lock()

     

       85
       98
        
       	defer m.mu.Unlock()

     

       86
       99
        
       

     

       87
       100
        
       	m.cacheBuffer = append(m.cacheBuffer, metric)

     

       88
       101
        
       	if len(m.cacheBuffer) >= m.config.BatchSize {

     

       89
       89
       -
       		m.flush()

     

       102
       102
       +
       		m.flushCacheMetricsLocked()

     

       90
       103
        
       	}

     

       91
       104
        
       }

     

       92
       105
        
       

     

       93
       93
       -
       func (m *ClickHouseMetrics) RecordCacheStats(stats *dns.CacheStats) {

     

       94
       94
       -
       	m.RecordCacheMetrics(CacheMetric{

     

       95
       95
       -
       		Timestamp:    time.Now(),

     

       96
       96
       -
       		TotalQueries: stats.TotalQueries.Load(),

     

       97
       97
       -
       		CacheHits:    stats.CacheHits.Load(),

     

       98
       98
       -
       		CacheMisses:  stats.CacheMisses.Load(),

     

       99
       99
       -
       		NegativeHits: stats.NegativeHits.Load(),

     

       100
       100
       -
       		PositiveHits: stats.PositiveHits.Load(),

     

       101
       101
       -
       		Evictions:    stats.Evictions.Load(),

     

       102
       102
       -
       		Size:         stats.Size.Load(),

     

       103
       103
       -
       	})

     

       104
       104
       -
       }

     

       105
       105
       -
       

     

       106
       106
        
       func (m *ClickHouseMetrics) flushLoop() {

     

       107
       107
       +
       	defer m.wg.Done()

     

       107
       108
        
       	ticker := time.NewTicker(m.config.FlushInterval.Duration)

     

       108
       109
        
       	defer ticker.Stop()

     

       109
       110
        
       

     

       111
       111
       +
       	m.logger.Info("Metrics flush loop started", "interval", m.config.FlushInterval.Duration)

     

       112
       112
       +
       

     

       110
       113
        
       	for {

     

       111
       114
        
       		select {

     

       112
       115
        
       		case <-ticker.C:

     

       113
       116
        
       			m.mu.Lock()

     

       114
       114
       -
       			m.flush()

     

       117
       117
       +
       			m.flushLocked()

     

       115
       118
        
       			m.mu.Unlock()

     

       116
       119
        
       		case <-m.stopChan:

     

       120
       120
       +
       			m.logger.Info("Metrics flush loop received stop signal. Flushing remaining data...")

     

       121
       121
       +
       			m.mu.Lock()

     

       122
       122
       +
       			m.flushLocked()

     

       123
       123
       +
       			m.mu.Unlock()

     

       124
       124
       +
       			m.logger.Info("Metrics flush loop stopped.")

     

       117
       125
        
       			return

     

       118
       126
        
       		}

     

       119
       127
        
       	}

     

       120
       128
        
       }

     

       121
       129
        
       

     

       122
       122
       -
       func (m *ClickHouseMetrics) flush() {

     

       130
       130
       +
       func (m *ClickHouseMetrics) flushLocked() {

     

       123
       131
        
       	if len(m.queryBuffer) > 0 {

     

       124
       124
       -
       		if err := m.flushQueries(); err != nil {

     

       125
       125
       -
       			m.logger.Error("Failed to flush query metrics", "error", err)

     

       126
       126
       -
       		}

     

       127
       127
       -
       		m.queryBuffer = m.queryBuffer[:0]

     

       132
       132
       +
       		m.flushQueriesLocked()

     

       128
       133
        
       	}

     

       129
       129
       -
       

     

       130
       134
        
       	if len(m.cacheBuffer) > 0 {

     

       131
       131
       -
       		if err := m.flushCacheMetrics(); err != nil {

     

       132
       132
       -
       			m.logger.Error("Failed to flush cache metrics", "error", err)

     

       133
       133
       -
       		}

     

       134
       134
       -
       		m.cacheBuffer = m.cacheBuffer[:0]

     

       135
       135
       +
       		m.flushCacheMetricsLocked()

     

       135
       136
        
       	}

     

       136
       137
        
       }

     

       137
       138
        
       

     

       138
       138
       -
       func (m *ClickHouseMetrics) changeTTL() error {

     

       139
       139
       -
       	if _, err := m.db.Exec(

     

       140
       140
       -
       		"ALTER TABLE alky_dns_queries MODIFY TTL timestamp + toIntervalSecond(?)",

     

       141
       141
       -
       		int(m.config.RetentionPeriod.Seconds()),

     

       142
       142
       -
       	); err != nil {

     

       143
       143
       -
       		return fmt.Errorf("failed to update alky_dns_queries TTL: %w", err)

     

       139
       139
       +
       func (m *ClickHouseMetrics) checkAndUpdateTTL() error {

     

       140
       140
       +
       	retentionSeconds := int(m.config.RetentionPeriod.Duration.Seconds())

     

       141
       141
       +
       	if retentionSeconds <= 0 {

     

       142
       142
       +
       		m.logger.Warn("Invalid retention period configured, skipping TTL update", "retention", m.config.RetentionPeriod.Duration)

     

       143
       143
       +
       		return nil

     

       144
       144
       +
       	}

     

       145
       145
       +
       

     

       146
       146
       +
       	m.logger.Info("Checking/Updating table TTLs", "retention_seconds", retentionSeconds)

     

       147
       147
       +
       

     

       148
       148
       +
       	updateTableTTL := func(tableName string) error {

     

       149
       149
       +
       		query := fmt.Sprintf("ALTER TABLE %s MODIFY TTL timestamp + toIntervalSecond(?)", tableName)

     

       150
       150
       +
       		m.logger.Debug("Executing TTL update query", "query", query)

     

       151
       151
       +
       		_, err := m.db.Exec(query, retentionSeconds)

     

       152
       152
       +
       		if err != nil && errors.Is(err, sql.ErrNoRows) {

     

       153
       153
       +
       			m.logger.Warn("Table not found while updating TTL, likely needs migration", "table", tableName)

     

       154
       154
       +
       			return nil

     

       155
       155
       +
       		} else if err != nil {

     

       156
       156
       +
       			return fmt.Errorf("failed to update %s TTL: %w", tableName, err)

     

       157
       157
       +
       		}

     

       158
       158
       +
       		m.logger.Debug("TTL updated successfully", "table", tableName)

     

       159
       159
       +
       		return nil

     

       144
       160
        
       	}

     

       145
       161
        
       

     

       146
       146
       -
       	if _, err := m.db.Exec(

     

       147
       147
       -
       		"ALTER TABLE alky_dns_cache_metrics MODIFY TTL timestamp + toIntervalSecond(?)",

     

       148
       148
       -
       		int(m.config.RetentionPeriod.Seconds()),

     

       149
       149
       -
       	); err != nil {

     

       150
       150
       -
       		return fmt.Errorf("failed to update alky_dns_cache_metrics TTL: %w", err)

     

       162
       162
       +
       	if err := updateTableTTL("alky_dns_queries"); err != nil {

     

       163
       163
       +
       		return err

     

       164
       164
       +
       	}

     

       165
       165
       +
       	if err := updateTableTTL("alky_dns_cache_metrics"); err != nil {

     

       166
       166
       +
       		return err

     

       151
       167
        
       	}

     

       152
       168
        
       

     

       153
       169
        
       	return nil

     

       154
       170
        
       }

     

       155
       171
        
       

     

       156
       156
       -
       func (m *ClickHouseMetrics) flushQueries() error {

     

       172
       172
       +
       func (m *ClickHouseMetrics) flushQueriesLocked() {

     

       173
       173
       +
       	if len(m.queryBuffer) == 0 {

     

       174
       174
       +
       		return

     

       175
       175
       +
       	}

     

       176
       176
       +
       

     

       177
       177
       +
       	m.logger.Debug("Flushing query metrics", "count", len(m.queryBuffer))

     

       178
       178
       +
       	start := time.Now()

     

       179
       179
       +
       

     

       157
       180
        
       	tx, err := m.db.Begin()

     

       158
       181
        
       	if err != nil {

     

       159
       159
       -
       		return err

     

       182
       182
       +
       		m.logger.Error("Failed to begin transaction for query metrics", "error", err)

     

       183
       183
       +
       		m.queryBuffer = m.queryBuffer[:0]

     

       184
       184
       +
       		return

     

       160
       185
        
       	}

     

       161
       161
       -
       	defer tx.Rollback()

     

       186
       186
       +
       	defer func() {

     

       187
       187
       +
       		if err != nil {

     

       188
       188
       +
       			if rbErr := tx.Rollback(); rbErr != nil {

     

       189
       189
       +
       				m.logger.Error("Failed to rollback transaction for query metrics", "error", rbErr)

     

       190
       190
       +
       			}

     

       191
       191
       +
       		}

     

       192
       192
       +
       	}()

     

       162
       193
        
       

     

       163
       194
        
       	stmt, err := tx.Prepare(`

     

       164
       195
        
       		INSERT INTO alky_dns_queries (

     
···

       167
       198
        
       		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)

     

       168
       199
        
       	`)

     

       169
       200
        
       	if err != nil {

     

       170
       170
       -
       		return err

     

       201
       201
       +
       		m.logger.Error("Failed to prepare statement for query metrics", "error", err)

     

       202
       202
       +
       		m.queryBuffer = m.queryBuffer[:0]

     

       203
       203
       +
       		return

     

       171
       204
        
       	}

     

       172
       205
        
       	defer stmt.Close()

     

       173
       206
        
       

     

       207
       207
       +
       	count := 0

     

       174
       208
        
       	for _, metric := range m.queryBuffer {

     

       175
       175
       -
       		_, err := stmt.Exec(

     

       209
       209
       +
       		_, err = stmt.Exec(

     

       176
       210
        
       			metric.Timestamp,

     

       177
       211
        
       			metric.InstanceID,

     

       178
       212
        
       			metric.QueryName,

     
···

       184
       218
        
       			metric.CacheHit,

     

       185
       219
        
       		)

     

       186
       220
        
       		if err != nil {

     

       187
       187
       -
       			return err

     

       221
       221
       +
       			m.logger.Error("Failed to execute statement for query metric", "error", err, "metric_index", count)

     

       222
       222
       +
       			return

     

       188
       223
        
       		}

     

       224
       224
       +
       		count++

     

       189
       225
        
       	}

     

       190
       226
        
       

     

       191
       191
       -
       	return tx.Commit()

     

       227
       227
       +
       	err = tx.Commit()

     

       228
       228
       +
       	if err != nil {

     

       229
       229
       +
       		m.logger.Error("Failed to commit transaction for query metrics", "error", err)

     

       230
       230
       +
       		return

     

       231
       231
       +
       	}

     

       232
       232
       +
       

     

       233
       233
       +
       	m.logger.Debug("Successfully flushed query metrics", "count", count, "duration", time.Since(start))

     

       234
       234
       +
       	m.queryBuffer = m.queryBuffer[:0]

     

       192
       235
        
       }

     

       193
       236
        
       

     

       194
       194
       -
       func (m *ClickHouseMetrics) flushCacheMetrics() error {

     

       237
       237
       +
       func (m *ClickHouseMetrics) flushCacheMetricsLocked() {

     

       238
       238
       +
       	if len(m.cacheBuffer) == 0 {

     

       239
       239
       +
       		return

     

       240
       240
       +
       	}

     

       241
       241
       +
       

     

       242
       242
       +
       	m.logger.Debug("Flushing cache metrics", "count", len(m.cacheBuffer))

     

       243
       243
       +
       	start := time.Now()

     

       244
       244
       +
       

     

       195
       245
        
       	tx, err := m.db.Begin()

     

       196
       246
        
       	if err != nil {

     

       197
       197
       -
       		return err

     

       247
       247
       +
       		m.logger.Error("Failed to begin transaction for cache metrics", "error", err)

     

       248
       248
       +
       		m.cacheBuffer = m.cacheBuffer[:0]

     

       249
       249
       +
       		return

     

       198
       250
        
       	}

     

       199
       199
       -
       	defer tx.Rollback()

     

       251
       251
       +
       	defer func() {

     

       252
       252
       +
       		if err != nil {

     

       253
       253
       +
       			if rbErr := tx.Rollback(); rbErr != nil {

     

       254
       254
       +
       				m.logger.Error("Failed to rollback transaction for cache metrics", "error", rbErr)

     

       255
       255
       +
       			}

     

       256
       256
       +
       		}

     

       257
       257
       +
       	}()

     

       200
       258
        
       

     

       201
       259
        
       	stmt, err := tx.Prepare(`

     

       202
       260
        
       		INSERT INTO alky_dns_cache_metrics (

     

       203
       203
       -
       			timestamp, instance_id, total_queries, cache_hits, cache_misses,

     

       204
       204
       -
       			negative_hits, positive_hits, evictions, size

     

       205
       205
       -
       		) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?)

     

       261
       261
       +
       			timestamp, instance_id, cache_hits, cache_misses,

     

       262
       262
       +
       			negative_hits, positive_hits, evictions,  size

     

       263
       263
       +
       		) VALUES (?, ?, ?, ?, ?, ?, ?, ?)

     

       206
       264
        
       	`)

     

       207
       265
        
       	if err != nil {

     

       208
       208
       -
       		return err

     

       266
       266
       +
       		m.logger.Error("Failed to prepare statement for cache metrics", "error", err)

     

       267
       267
       +
       		m.cacheBuffer = m.cacheBuffer[:0]

     

       268
       268
       +
       		return

     

       209
       269
        
       	}

     

       210
       270
        
       	defer stmt.Close()

     

       211
       271
        
       

     

       272
       272
       +
       	count := 0

     

       212
       273
        
       	for _, metric := range m.cacheBuffer {

     

       213
       213
       -
       		_, err := stmt.Exec(

     

       274
       274
       +
       		instanceID := GetInstanceID()

     

       275
       275
       +
       

     

       276
       276
       +
       		_, err = stmt.Exec(

     

       214
       277
        
       			metric.Timestamp,

     

       215
       215
       -
       			GetInstanceID(),

     

       216
       216
       -
       			metric.TotalQueries,

     

       278
       278
       +
       			instanceID,

     

       217
       279
        
       			metric.CacheHits,

     

       218
       280
        
       			metric.CacheMisses,

     

       219
       281
        
       			metric.NegativeHits,

     
···

       222
       284
        
       			metric.Size,

     

       223
       285
        
       		)

     

       224
       286
        
       		if err != nil {

     

       225
       225
       -
       			return err

     

       287
       287
       +
       			m.logger.Error("Failed to execute statement for cache metric", "error", err, "metric_index", count)

     

       288
       288
       +
       			return

     

       226
       289
        
       		}

     

       290
       290
       +
       		count++

     

       227
       291
        
       	}

     

       228
       292
        
       

     

       229
       229
       -
       	return tx.Commit()

     

       293
       293
       +
       	err = tx.Commit()

     

       294
       294
       +
       	if err != nil {

     

       295
       295
       +
       		m.logger.Error("Failed to commit transaction for cache metrics", "error", err)

     

       296
       296
       +
       		return

     

       297
       297
       +
       	}

     

       298
       298
       +
       

     

       299
       299
       +
       	m.logger.Debug("Successfully flushed cache metrics", "count", count, "duration", time.Since(start))

     

       300
       300
       +
       	m.cacheBuffer = m.cacheBuffer[:0]

     

       230
       301
        
       }

     

       231
       302
        
       

     

       232
       303
        
       func (m *ClickHouseMetrics) Close() error {

     

       304
       304
       +
       	m.logger.Info("Closing metrics client...")

     

       233
       305
        
       	close(m.stopChan)

     

       234
       234
       -
       	m.mu.Lock()

     

       235
       235
       -
       	defer m.mu.Unlock()

     

       236
       236
       -
       	m.flush()

     

       306
       306
       +
       	m.wg.Wait()

     

       307
       307
       +
       	m.logger.Info("Flush loop stopped. Closing database connection.")

     

       237
       308
        
       	return m.db.Close()

     

       238
       309
        
       }

+39 -17

pkg/metrics/middleware.go

···

       3
       3
        
       import (

     

       4
       4
        
       	"context"

     

       5
       5
        
       	"fmt"

     

       6
       6
       +
       	"log/slog"

     

       6
       7
        
       	"os"

     

       8
       8
       +
       	"sync"

     

       7
       9
        
       	"time"

     

       8
       10
        
       

     

       9
       11
        
       	"tangled.sh/seiso.moe/alky/pkg/dns"

     
···

       11
       13
        
       

     

       12
       14
        
       var (

     

       13
       15
        
       	instanceID string

     

       16
       16
       +
       	initOnce   sync.Once

     

       14
       17
        
       	version    string

     

       15
       18
        
       )

     

       16
       19
        
       

     

       17
       17
       -
       func init() {

     

       18
       18
       -
       	var err error

     

       20
       20
       +
       func initializeInstanceID() {

     

       19
       21
        
       	hostname, err := os.Hostname()

     

       20
       22
        
       	if err != nil {

     

       21
       21
       -
       		hostname = "unknown"

     

       23
       23
       +
       		hostname = "unknown-host"

     

       22
       24
        
       	}

     

       23
       25
        
       

     

       26
       26
       +
       	pid := os.Getpid()

     

       27
       27
       +
       

     

       24
       28
        
       	if version != "" {

     

       25
       25
       -
       		instanceID = fmt.Sprintf("%s-%s", hostname, version)

     

       29
       29
       +
       		instanceID = fmt.Sprintf("%s-%d-%s", hostname, pid, version)

     

       26
       30
        
       	} else {

     

       27
       27
       -
       		instanceID = hostname

     

       31
       31
       +
       		instanceID = fmt.Sprintf("%s-%d-dev", hostname, pid)

     

       28
       32
        
       	}

     

       29
       33
        
       }

     

       30
       34
        
       

     

       31
       35
        
       func MetricsMiddleware(metrics *ClickHouseMetrics) func(dns.Handler) dns.Handler {

     

       36
       36
       +
       	if metrics == nil {

     

       37
       37
       +
       		slog := slog.Default()

     

       38
       38
       +
       		slog.Warn("Metrics client is nil, metrics middleware will be a no-op")

     

       39
       39
       +
       		return func(next dns.Handler) dns.Handler {

     

       40
       40
       +
       			return next

     

       41
       41
       +
       		}

     

       42
       42
       +
       	}

     

       43
       43
       +
       

     

       44
       44
       +
       	instanceID := GetInstanceID()

     

       45
       45
       +
       

     

       32
       46
        
       	return func(next dns.Handler) dns.Handler {

     

       33
       47
        
       		return dns.HandlerFunc(func(w dns.ResponseWriter, r *dns.Request) {

     

       34
       48
        
       			if r.Context == nil {

     
···

       39
       53
        
       			next.ServeDNS(w, r)

     

       40
       54
        
       			duration := time.Since(start)

     

       41
       55
        
       

     

       42
       42
       -
       			question := r.Message.Question[0]

     

       43
       43
       -
       			metrics.RecordQuery(QueryMetric{

     

       44
       44
       -
       				Timestamp:    time.Now(),

     

       45
       45
       -
       				InstanceID:   instanceID,

     

       46
       46
       -
       				QueryName:    question.QName,

     

       47
       47
       -
       				QueryType:    question.QType.String(),

     

       48
       48
       -
       				QueryClass:   question.QClass.String(),

     

       49
       49
       -
       				RemoteAddr:   r.RemoteAddr.String(),

     

       50
       50
       -
       				ResponseCode: r.Message.Header.RCode.String(),

     

       51
       51
       -
       				Duration:     duration.Nanoseconds(),

     

       52
       52
       -
       				CacheHit:     dns.GetCacheHit(r.Context),

     

       53
       53
       -
       			})

     

       56
       56
       +
       			if r.Message != nil && len(r.Message.Question) > 0 {

     

       57
       57
       +
       				question := r.Message.Question[0]

     

       58
       58
       +
       				cacheHit := dns.GetCacheHit(r.Context)

     

       59
       59
       +
       

     

       60
       60
       +
       				metrics.RecordQuery(QueryMetric{

     

       61
       61
       +
       					Timestamp:    start,

     

       62
       62
       +
       					InstanceID:   instanceID,

     

       63
       63
       +
       					QueryName:    question.QName,

     

       64
       64
       +
       					QueryType:    question.QType.String(),

     

       65
       65
       +
       					QueryClass:   question.QClass.String(),

     

       66
       66
       +
       					RemoteAddr:   r.RemoteAddr.String(),

     

       67
       67
       +
       					ResponseCode: r.Message.Header.RCode.String(),

     

       68
       68
       +
       					Duration:     duration.Nanoseconds(),

     

       69
       69
       +
       					CacheHit:     cacheHit,

     

       70
       70
       +
       				})

     

       71
       71
       +
       			} else {

     

       72
       72
       +
       				slog := slog.Default()

     

       73
       73
       +
       				slog.Warn("Metrics middleware received request with missing message or question", "remote_addr", r.RemoteAddr)

     

       74
       74
       +
       			}

     

       54
       75
        
       		})

     

       55
       76
        
       	}

     

       56
       77
        
       }

     

       57
       78
        
       

     

       58
       79
        
       func GetInstanceID() string {

     

       80
       80
       +
       	initOnce.Do(initializeInstanceID)

     

       59
       81
        
       	return instanceID

     

       60
       82
        
       }