seiso.moe
1package main
2
3import (
4 "flag"
5 "log"
6 "log/slog"
7 "os"
8 "time"
9
10 "code.kiri.systems/kiri/alky/pkg/config"
11 "code.kiri.systems/kiri/alky/pkg/dns"
12 "code.kiri.systems/kiri/alky/pkg/metrics"
13 "code.kiri.systems/kiri/alky/pkg/rootservers"
14)
15
16var configFlag string
17
18func init() {
19 flag.StringVar(&configFlag, "config", "/etc/alky/alky.toml", "config file path for alky")
20
21 flag.Parse()
22}
23
24func main() {
25 cfg, err := config.LoadConfig(configFlag)
26 if err != nil {
27 log.Fatal(err)
28 }
29
30 logger := setupLogger(&cfg)
31
32 metricsClient, err := metrics.NewClickHouseMetrics(&cfg.Metrics, logger)
33 if err != nil {
34 log.Fatal(err)
35 }
36 defer metricsClient.Close()
37
38 rootServers, err := rootservers.DecodeRootHints(cfg.Server.RootHintsFile)
39 if err != nil {
40 log.Fatal(err)
41 }
42
43 cache := dns.NewMemoryCache(5000, 5*time.Minute)
44
45 handler := &dns.QueryHandler{
46 RootServers: rootServers,
47 Timeout: time.Duration(cfg.Advanced.QueryTimeout) * time.Second,
48 Cache: cache,
49 Logger: logger,
50 }
51
52 go monitorCacheMetrics(cache, metricsClient, logger)
53
54 rateLimitHandler := dns.RateLimitMiddleware(&dns.RateLimitConfig{
55 Rate: float64(cfg.Ratelimit.Rate),
56 Burst: cfg.Ratelimit.Burst,
57 WindowLength: time.Duration(cfg.Ratelimit.Window) * time.Second,
58 ExpirationTime: time.Duration(cfg.Ratelimit.ExpirationTime) * time.Second,
59 })(handler)
60
61 metricsHandler := metrics.MetricsMiddleware(metricsClient)(rateLimitHandler)
62
63 loggingHandler := dns.LoggingMiddleware(&dns.LogConfig{
64 Logger: logger,
65 Level: slog.LevelInfo,
66 })(metricsHandler)
67
68 s := dns.Server{
69 Address: cfg.Server.Address,
70 Port: cfg.Server.Port,
71 Handler: loggingHandler,
72 UDPSize: 512,
73 ReadTimeout: 2 * time.Second,
74 WriteTimeout: 2 * time.Second,
75 Logger: logger,
76 }
77
78 if err := s.ListenAndServe(); err != nil {
79 slog.Error("Failed to start server", "error", err)
80 }
81}
82
83func monitorCacheMetrics(cache *dns.MemoryCache, metricsClient *metrics.ClickHouseMetrics, logger *slog.Logger) {
84 ticker := time.NewTicker(1 * time.Minute)
85 defer ticker.Stop()
86
87 for range ticker.C {
88 stats := cache.GetStats()
89 metricsClient.RecordCacheStats(stats)
90 logger.Info("Cache metrics recorded to ClickHouse")
91 }
92}
93
94func setupLogger(cfg *config.Config) *slog.Logger {
95 var logger *slog.Logger
96
97 handlerOpts := &slog.HandlerOptions{
98 Level: slog.LevelDebug,
99 }
100
101 switch cfg.Logging.Output {
102 case "file":
103 f, err := os.OpenFile(cfg.Logging.FilePath, os.O_WRONLY|os.O_APPEND|os.O_CREATE, 0o644)
104 if err != nil {
105 log.Fatal(err)
106 }
107
108 logger = slog.New(slog.NewJSONHandler(f, handlerOpts))
109 default:
110 logger = slog.New(slog.NewJSONHandler(os.Stdout, handlerOpts))
111 }
112
113 return logger
114}