diff --git a/.mockery.yaml b/.mockery.yaml index 82c0e3e0..280d3fd6 100644 --- a/.mockery.yaml +++ b/.mockery.yaml @@ -27,3 +27,12 @@ packages: inpackage: true filename: "mocks_test.go" + # Resolver package - ChainedResolver interface + # Note: Using non-test filename (mocks.go) to allow cross-package access from server tests + github.com/0xERR0R/blocky/resolver: + interfaces: + ChainedResolver: + config: + inpackage: true + filename: "mocks.go" + diff --git a/config/config.go b/config/config.go index 29130ed3..f424421c 100644 --- a/config/config.go +++ b/config/config.go @@ -277,6 +277,7 @@ type Config struct { DNS64 DNS64 `yaml:"dns64"` DNSSEC DNSSEC `yaml:"dnssec"` HTTP3 HTTP3 `yaml:"http3"` + RateLimit RateLimit `yaml:"rateLimit"` // Deprecated options Deprecated struct { @@ -658,6 +659,10 @@ func (cfg *Config) validate(logger *logrus.Entry) { if err := cfg.DNS64.validate(logger, &cfg.Filtering, &cfg.Caching); err != nil { logger.Fatal(err) } + + if err := cfg.RateLimit.validate(); err != nil { + logger.Fatal(err) + } } // ConvertPort converts string representation into a valid port (0 - 65535) diff --git a/config/config_test.go b/config/config_test.go index 6d25dc1e..8ba8a3d7 100644 --- a/config/config_test.go +++ b/config/config_test.go @@ -1476,6 +1476,13 @@ var _ = Describe("TLSVersion", func() { }) }) +var _ = Describe("Config with RateLimit", func() { + It("exposes the RateLimit field", func() { + cfg := Config{} + Expect(cfg.RateLimit.IsEnabled()).Should(BeFalse()) + }) +}) + // Tiny helper to get a new pointer with a value. // // Avoids needing 2 lines: `x := new(T)` and `*x = val` diff --git a/config/rate_limit.go b/config/rate_limit.go new file mode 100644 index 00000000..17381b13 --- /dev/null +++ b/config/rate_limit.go @@ -0,0 +1,95 @@ +package config + +import ( + "errors" + "fmt" + "net" + + "github.com/sirupsen/logrus" +) + +const ( + ipv4MaxPrefix = 32 + ipv6MaxPrefix = 128 +) + +// RateLimit configures per-client rate limiting at the head of the resolver chain. +type RateLimit struct { + Enable bool `default:"false" yaml:"enable"` + Rate uint `default:"0" yaml:"rate"` + Burst uint `default:"0" yaml:"burst"` + IPv4Prefix uint8 `default:"32" yaml:"ipv4Prefix"` + IPv6Prefix uint8 `default:"64" yaml:"ipv6Prefix"` + Allowlist []string `yaml:"allowlist"` + + parsedAllowlist []*net.IPNet +} + +// IsEnabled implements `config.Configurable`. +func (c *RateLimit) IsEnabled() bool { return c.Enable } + +// LogConfig implements `config.Configurable`. +func (c *RateLimit) LogConfig(logger *logrus.Entry) { + logger.Infof("rate = %d qps", c.Rate) + logger.Infof("burst = %d", c.Burst) + logger.Infof("ipv4 prefix = /%d", c.IPv4Prefix) + logger.Infof("ipv6 prefix = /%d", c.IPv6Prefix) + logger.Infof("allowlist = %v", c.Allowlist) +} + +func (c *RateLimit) validate() error { + if !c.Enable { + return nil + } + if c.Rate == 0 { + return errors.New("rateLimit: rate must be > 0 when enabled") + } + if c.Burst == 0 { + // Sensible default: allow short spikes of up to 2× the sustained rate. + c.Burst = c.Rate * 2 + } + if c.Burst < c.Rate { + return fmt.Errorf("rateLimit: burst (%d) must be >= rate (%d)", c.Burst, c.Rate) + } + if c.IPv4Prefix > ipv4MaxPrefix { + return fmt.Errorf("rateLimit: ipv4Prefix (%d) must be in [0, %d]", c.IPv4Prefix, ipv4MaxPrefix) + } + if c.IPv6Prefix > ipv6MaxPrefix { + return fmt.Errorf("rateLimit: ipv6Prefix (%d) must be in [0, %d]", c.IPv6Prefix, ipv6MaxPrefix) + } + parsed := make([]*net.IPNet, 0, len(c.Allowlist)) + for _, s := range c.Allowlist { + ipNet, err := parseCIDRorIP(s) + if err != nil { + return fmt.Errorf("rateLimit: allowlist entry %q: %w", s, err) + } + parsed = append(parsed, ipNet) + } + c.parsedAllowlist = parsed + + return nil +} + +// ValidateForTest exposes validate for cross-package tests. +// Internal package callers should use the unexported validate. +func (c *RateLimit) ValidateForTest() error { return c.validate() } + +// ParsedAllowlist returns the parsed CIDR list populated by validate. +func (c *RateLimit) ParsedAllowlist() []*net.IPNet { return c.parsedAllowlist } + +func parseCIDRorIP(s string) (*net.IPNet, error) { + if _, ipNet, err := net.ParseCIDR(s); err == nil { + return ipNet, nil + } + if ip := net.ParseIP(s); ip != nil { + bits := ipv6MaxPrefix + if v4 := ip.To4(); v4 != nil { + ip = v4 + bits = ipv4MaxPrefix + } + + return &net.IPNet{IP: ip, Mask: net.CIDRMask(bits, bits)}, nil + } + + return nil, fmt.Errorf("not a valid CIDR or IP: %q", s) +} diff --git a/config/rate_limit_test.go b/config/rate_limit_test.go new file mode 100644 index 00000000..d3512c80 --- /dev/null +++ b/config/rate_limit_test.go @@ -0,0 +1,128 @@ +package config + +import ( + "net" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" +) + +var _ = Describe("RateLimit config", func() { + Describe("parseCIDRorIP", func() { + It("parses CIDR notation", func() { + n, err := parseCIDRorIP("10.0.0.0/8") + Expect(err).Should(Succeed()) + Expect(n.String()).Should(Equal("10.0.0.0/8")) + }) + + It("parses bare IPv4 as /32", func() { + n, err := parseCIDRorIP("192.168.1.10") + Expect(err).Should(Succeed()) + Expect(n.String()).Should(Equal("192.168.1.10/32")) + }) + + It("parses bare IPv6 as /128", func() { + n, err := parseCIDRorIP("::1") + Expect(err).Should(Succeed()) + Expect(n.IP.Equal(net.ParseIP("::1"))).Should(BeTrue()) + ones, bits := n.Mask.Size() + Expect(ones).Should(Equal(128)) + Expect(bits).Should(Equal(128)) + }) + + It("parses IPv6 CIDR", func() { + n, err := parseCIDRorIP("2001:db8::/32") + Expect(err).Should(Succeed()) + Expect(n.String()).Should(Equal("2001:db8::/32")) + }) + + It("rejects garbage", func() { + _, err := parseCIDRorIP("not-an-ip") + Expect(err).Should(HaveOccurred()) + }) + + It("rejects empty string", func() { + _, err := parseCIDRorIP("") + Expect(err).Should(HaveOccurred()) + }) + }) + + Describe("IsEnabled", func() { + It("is false when Enable is false", func() { + c := &RateLimit{Enable: false} + Expect(c.IsEnabled()).Should(BeFalse()) + }) + It("is true when Enable is true", func() { + c := &RateLimit{Enable: true} + Expect(c.IsEnabled()).Should(BeTrue()) + }) + }) + + Describe("LogConfig", func() { + suiteBeforeEach() + It("emits one line per relevant field", func() { + c := &RateLimit{ + Enable: true, Rate: 50, Burst: 100, + IPv4Prefix: 32, IPv6Prefix: 64, + Allowlist: []string{"10.0.0.0/8"}, + } + c.LogConfig(logger) + Expect(hook.Calls).ShouldNot(BeEmpty()) + }) + }) + + Describe("validate", func() { + It("accepts disabled config", func() { + c := &RateLimit{Enable: false} + Expect(c.validate()).Should(Succeed()) + }) + It("rejects enabled with Rate=0", func() { + c := &RateLimit{Enable: true, Rate: 0, Burst: 0, IPv4Prefix: 32, IPv6Prefix: 64} + Expect(c.validate()).Should(MatchError(ContainSubstring("rate must be > 0"))) + }) + It("rejects enabled with non-zero Burst < Rate", func() { + c := &RateLimit{Enable: true, Rate: 50, Burst: 10, IPv4Prefix: 32, IPv6Prefix: 64} + Expect(c.validate()).Should(MatchError(ContainSubstring("burst"))) + }) + It("defaults Burst to 2× Rate when left at zero", func() { + c := &RateLimit{Enable: true, Rate: 50, Burst: 0, IPv4Prefix: 32, IPv6Prefix: 64} + Expect(c.validate()).Should(Succeed()) + Expect(c.Burst).Should(BeNumerically("==", 100)) + }) + It("rejects IPv4Prefix > 32", func() { + c := &RateLimit{Enable: true, Rate: 1, Burst: 1, IPv4Prefix: 33, IPv6Prefix: 64} + Expect(c.validate()).Should(MatchError(ContainSubstring("ipv4Prefix"))) + }) + It("rejects IPv6Prefix > 128", func() { + c := &RateLimit{Enable: true, Rate: 1, Burst: 1, IPv4Prefix: 32, IPv6Prefix: 129} + Expect(c.validate()).Should(MatchError(ContainSubstring("ipv6Prefix"))) + }) + It("rejects unparseable allowlist entry", func() { + c := &RateLimit{ + Enable: true, Rate: 1, Burst: 1, + IPv4Prefix: 32, IPv6Prefix: 64, + Allowlist: []string{"garbage"}, + } + Expect(c.validate()).Should(MatchError(ContainSubstring("garbage"))) + }) + It("accepts bare IPv4 and IPv6 in allowlist", func() { + c := &RateLimit{ + Enable: true, Rate: 1, Burst: 1, + IPv4Prefix: 32, IPv6Prefix: 64, + Allowlist: []string{"127.0.0.1", "::1", "10.0.0.0/8"}, + } + Expect(c.validate()).Should(Succeed()) + Expect(c.parsedAllowlist).Should(HaveLen(3)) + }) + It("re-parses on repeated calls", func() { + c := &RateLimit{ + Enable: true, Rate: 1, Burst: 1, + IPv4Prefix: 32, IPv6Prefix: 64, + Allowlist: []string{"10.0.0.0/8"}, + } + Expect(c.validate()).Should(Succeed()) + Expect(c.validate()).Should(Succeed()) + Expect(c.parsedAllowlist).Should(HaveLen(1)) + }) + }) +}) diff --git a/docs/config.yml b/docs/config.yml index 8e497585..7d25cc76 100644 --- a/docs/config.yml +++ b/docs/config.yml @@ -300,6 +300,43 @@ filtering: queryTypes: - AAAA +# optional: per-client DNS query rate limiting at the head of the resolver chain. +# Intended for misbehaving clients, runaway scripts, and low-volume amplification +# abuse — NOT a DDoS defense (run fail2ban or crowdsec at the firewall for that). +# When a client exceeds its bucket, the query is dropped silently (no UDP response, +# no DNS error) so the amplification ratio against a spoofed-source attacker is zero. +# A WARN log line is emitted (throttled to <= 1 line per bucket per second) so +# fail2ban can match against it. See docs/configuration.md for the full sample filter. +rateLimit: + # enable rate limiting. default: false + enable: false + # sustained queries per second allowed per client bucket. Must be > 0 when enabled. + # default: 0 + rate: 50 + # token bucket capacity — i.e. the maximum number of queries a previously + # idle client may spend at once before being throttled to `rate` qps. + # Example: with rate=100 and burst=200, a client may fire 200 queries + # instantly, then is limited to 100/sec; after 2 seconds idle they + # regain a full 200-query burst. Must be >= rate when set explicitly. + # default: 0 — falls back to rate × 2 if left at zero + burst: 100 + # IPv4 aggregation prefix length. Clients sharing the same /ipv4Prefix share + # one bucket. Use 32 to limit per single address, 24 to limit per /24 subnet, etc. + # default: 32 + ipv4Prefix: 32 + # IPv6 aggregation prefix length. Defaults to /64 — the standard boundary for + # abuse mitigation, since a single subscriber typically gets a /64 from their ISP. + # default: 64 + ipv6Prefix: 64 + # optional: list of CIDRs or bare IPs that bypass the limiter entirely. + # Useful for trusted internal networks and monitoring probes. + # Accepts IPv4/IPv6 CIDR notation or single addresses. + # default: [] (no bypass) + allowlist: + - 10.0.0.0/8 + - 192.168.0.0/16 + - 127.0.0.1 + # optional: return NXDOMAIN for queries that are not FQDNs. fqdnOnly: # default: false diff --git a/docs/configuration.md b/docs/configuration.md index 69e0e2af..74da53ed 100644 --- a/docs/configuration.md +++ b/docs/configuration.md @@ -300,6 +300,66 @@ types, all queries with these types will be dropped (empty answer will be return This configuration will drop all 'AAAA' (IPv6) queries. +## Rate limiting per client IP + +Blocky can enforce a per-client query rate limit at the head of the resolver chain. **Disabled by default.** This is _not_ a DDoS defense — for that, run [fail2ban](https://www.fail2ban.org/) or [crowdsec](https://www.crowdsec.net/) at the firewall, which can drop traffic before it reaches Blocky. The limiter is intended for misbehaving clients, runaway scripts, and the kind of low-volume amplification that can enroll a public Blocky instance in someone else's DNS reflection attack ([issue #1135](https://github.com/0xERR0R/blocky/issues/1135)). + +When a client exceeds its bucket, the query is **dropped silently** (no UDP response, no DNS error) so that amplification ratio against a spoofed-source attacker is zero. A WARN log line is emitted (throttled to ≤ 1 line per bucket per second) so fail2ban can match against it. + +### Configuration + +```yaml +rateLimit: + enable: true + rate: 50 # avg queries per second per client + burst: 100 # optional, default rate × 2; must be >= rate + ipv4Prefix: 32 # default: aggregate by /32 (one IP = one client) + ipv6Prefix: 64 # default: aggregate by /64 (one subscriber prefix) + allowlist: + - 127.0.0.1/32 # localhost (recommended if /api/query is used) + - ::1/128 + - 10.0.0.0/8 + - 192.168.0.0/16 +``` + +The limiter is a **token bucket** per client. Tokens refill at `rate` per second, capped at `burst`. Each query consumes one token. When the bucket is empty, the query is dropped. + +- `rate` is the long-run sustained limit in queries/second. +- `burst` is the bucket capacity — how many queries an idle client may spend at once before being throttled to `rate`. Optional; defaults to `rate × 2` if omitted or set to 0. +- `ipv6Prefix` defaults to `/64` because residential IPv6 typically allocates a `/64` per subscriber; aggregating by `/64` keeps an attacker from cheaply rotating source addresses within their own subnet. + +**Example — `rate: 100, burst: 200`:** + +| Time | Bucket | Client sends | Result | +|------|--------|--------------|--------| +| 0.00 s | 200 (full) | 200 queries instantly | all 200 allowed, bucket drained | +| 0.00 s | 0 | 1 more query | **dropped** | +| 0.50 s | 50 (refilled at 100/s) | 50 queries | all 50 allowed | +| 0.50 s | 0 | 1 more query | **dropped** | +| 2.50 s | 200 (capped) | 200 queries instantly | all 200 allowed (burst recovered after 2 s idle) | + +Pick `burst == rate` for a strict steady-state limiter with no spike tolerance. Pick `burst` 2–10× `rate` to let typical browser page loads through (a single page can trigger 30–80 DNS lookups in well under a second) while still capping long-running abusers at `rate` qps. + +### Sample fail2ban filter + +`/etc/fail2ban/filter.d/blocky-ratelimit.conf`: + +```ini +[Definition] +failregex = ^.*rate-limiting: dropped query.*client_ip=.*$ +ignoreregex = +``` + +This filter assumes `logFormat: text` (the default); with `logFormat: json` the field becomes `"client_ip":""` and the regex must be adjusted accordingly. + +### Metrics + +- `blocky_rate_limit_drops_total{protocol="TCP|UDP"}` — counter, incremented on every drop (label cardinality is bounded). +- `blocky_rate_limit_cap_exhausted_total` — counter, incremented when a new client is dropped because the in-memory bucket store has reached its hard cap. A non-zero value means the store is under sustained pressure from many distinct clients. +- `blocky_rate_limit_active_buckets` — gauge, current number of in-memory token buckets. + +Per-client IP is intentionally not a Prometheus label; it lives in the log line. Per-IP analysis is better served by log aggregation (Loki, ELK) than by Prometheus. + ## FQDN only In domain environments, it may be useful to only respond to FQDN requests. If this option is enabled blocky will respond immediately diff --git a/e2e/rate_limit_test.go b/e2e/rate_limit_test.go new file mode 100644 index 00000000..2d35a7f3 --- /dev/null +++ b/e2e/rate_limit_test.go @@ -0,0 +1,196 @@ +package e2e + +import ( + "context" + "errors" + "io" + "net" + "net/http" + "strings" + "time" + + . "github.com/0xERR0R/blocky/helpertest" + "github.com/0xERR0R/blocky/util" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/testcontainers/testcontainers-go" +) + +var _ = Describe("Per-client rate limiting", func() { + var ( + e2eNet *testcontainers.DockerNetwork + blocky testcontainers.Container + err error + ) + + BeforeEach(func(ctx context.Context) { + e2eNet = getRandomNetwork(ctx) + _, err = createDNSMokkaContainer(ctx, "moka", e2eNet, + `A example.com/NOERROR("A 1.2.3.4 123")`, + ) + Expect(err).Should(Succeed()) + }) + + When("rate limiting is disabled", func() { + BeforeEach(func(ctx context.Context) { + blocky, err = createBlockyContainerFromString(ctx, e2eNet, dedent(` + upstreams: + groups: + default: + - moka + `)) + Expect(err).Should(Succeed()) + }) + + It("passes traffic through unhindered", func(ctx context.Context) { + for range 100 { + msg := util.NewMsgWithQuestion("example.com.", A) + Expect(doDNSRequest(ctx, blocky, msg)). + Should(BeDNSRecord("example.com.", A, "1.2.3.4")) + } + }) + }) + + When("rate=1 burst=1 (no allowlist)", func() { + BeforeEach(func(ctx context.Context) { + blocky, err = createBlockyContainerFromString(ctx, e2eNet, dedent(` + upstreams: + groups: + default: + - moka + rateLimit: + enable: true + rate: 1 + burst: 1 + `)) + Expect(err).Should(Succeed()) + }) + + It("drops the second back-to-back query", func(ctx context.Context) { + msg := util.NewMsgWithQuestion("example.com.", A) + + _, err := doDNSRequest(ctx, blocky, msg) + Expect(err).Should(Succeed()) + + _, err = doDNSRequest(ctx, blocky, msg) + Expect(err).Should(HaveOccurred()) + var netErr net.Error + Expect(errors.As(err, &netErr)).Should(BeTrue()) + Expect(netErr.Timeout()).Should(BeTrue()) + }) + + It("allows again after a refill window", func(ctx context.Context) { + msg := util.NewMsgWithQuestion("example.com.", A) + _, err := doDNSRequest(ctx, blocky, msg) + Expect(err).Should(Succeed()) + + Eventually(func() error { + _, err := doDNSRequest(ctx, blocky, msg) + + return err + }).WithTimeout(15 * time.Second).WithPolling(time.Second).Should(Succeed()) + }) + + It("emits a fail2ban-matchable WARN line", func(ctx context.Context) { + msg := util.NewMsgWithQuestion("example.com.", A) + _, _ = doDNSRequest(ctx, blocky, msg) + _, _ = doDNSRequest(ctx, blocky, msg) + + Eventually(func() bool { + lines, err := getContainerLogs(ctx, blocky) + if err != nil { + return false + } + for _, line := range lines { + if strings.Contains(line, "rate-limiting:") && + strings.Contains(line, "dropped query") && + strings.Contains(line, "client_ip=") { + return true + } + } + + return false + }, 5*time.Second, 200*time.Millisecond).Should(BeTrue()) + }) + }) + + When("rate=1 burst=1 with allowlist covering the test gateway", func() { + BeforeEach(func(ctx context.Context) { + blocky, err = createBlockyContainerFromString(ctx, e2eNet, dedent(` + upstreams: + groups: + default: + - moka + rateLimit: + enable: true + rate: 1 + burst: 1 + allowlist: + - 172.16.0.0/12 + - 10.0.0.0/8 + - 192.168.0.0/16 + `)) + Expect(err).Should(Succeed()) + }) + + It("never drops allowlisted clients", func(ctx context.Context) { + for range 5 { + msg := util.NewMsgWithQuestion("example.com.", A) + Expect(doDNSRequest(ctx, blocky, msg)). + Should(BeDNSRecord("example.com.", A, "1.2.3.4")) + } + }) + }) + + When("Prometheus is enabled and a drop occurs", func() { + BeforeEach(func(ctx context.Context) { + blocky, err = createBlockyContainerFromString(ctx, e2eNet, dedent(` + upstreams: + groups: + default: + - moka + rateLimit: + enable: true + rate: 1 + burst: 1 + ports: + http: 4000 + prometheus: + enable: true + path: /metrics + `)) + Expect(err).Should(Succeed()) + }) + + It("exposes blocky_rate_limit_drops_total > 0", func(ctx context.Context) { + msg := util.NewMsgWithQuestion("example.com.", A) + _, _ = doDNSRequest(ctx, blocky, msg) + _, _ = doDNSRequest(ctx, blocky, msg) + _, _ = doDNSRequest(ctx, blocky, msg) + + Eventually(func(g Gomega) { + host, port, err := getContainerHostPort(ctx, blocky, "4000/tcp") + g.Expect(err).Should(Succeed()) + + resp, err := http.Get("http://" + net.JoinHostPort(host, port) + "/metrics") + g.Expect(err).Should(Succeed()) + defer resp.Body.Close() + + body, err := io.ReadAll(resp.Body) + g.Expect(err).Should(Succeed()) + + // Look for the counter with a non-zero value + lines := strings.Split(string(body), "\n") + var foundNonZero bool + for _, line := range lines { + if strings.HasPrefix(line, "blocky_rate_limit_drops_total{") && !strings.HasSuffix(line, " 0") { + foundNonZero = true + + break + } + } + g.Expect(foundNonZero).Should(BeTrue(), "expected non-zero drops counter; metrics body=\n%s", string(body)) + }, 10*time.Second, 500*time.Millisecond).Should(Succeed()) + }) + }) +}) diff --git a/go.mod b/go.mod index c3f82826..977589cd 100644 --- a/go.mod +++ b/go.mod @@ -39,6 +39,7 @@ require ( github.com/x-cray/logrus-prefixed-formatter v0.5.2 golang.org/x/exp v0.0.0-20250718183923-645b1fa84792 golang.org/x/net v0.54.0 + golang.org/x/time v0.15.0 gopkg.in/yaml.v2 v2.4.0 gorm.io/driver/mysql v1.6.0 gorm.io/driver/postgres v1.6.0 diff --git a/go.sum b/go.sum index 6b709c08..90fdcc71 100644 --- a/go.sum +++ b/go.sum @@ -524,6 +524,8 @@ golang.org/x/text v0.7.0/go.mod h1:mrYo+phRRbMaCq/xk9113O4dZlRixOauAjOtrjsXDZ8= golang.org/x/text v0.9.0/go.mod h1:e1OnstbJyHTd6l/uOt8jFFHp6TRDWZR/bV3emEE/zU8= golang.org/x/text v0.37.0 h1:Cqjiwd9eSg8e0QAkyCaQTNHFIIzWtidPahFWR83rTrc= golang.org/x/text v0.37.0/go.mod h1:a5sjxXGs9hsn/AJVwuElvCAo9v8QYLzvavO5z2PiM38= +golang.org/x/time v0.15.0 h1:bbrp8t3bGUeFOx08pvsMYRTCVSMk89u4tKbNOZbp88U= +golang.org/x/time v0.15.0/go.mod h1:Y4YMaQmXwGQZoFaVFk4YpCt4FLQMYKZe9oeV/f4MSno= golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ= golang.org/x/tools v0.0.0-20191119224855-298f0cb1881e/go.mod h1:b+2E5dAYhXwXZwtnZ6UAqBI28+e2cm9otk0dWdXHAEo= golang.org/x/tools v0.0.0-20201224043029-2b0845dc783e/go.mod h1:emZCQorbCU4vsT4fOWvOPXz4eW1wZW4PmDk9uLelYpA= diff --git a/resolver/mocks.go b/resolver/mocks.go new file mode 100644 index 00000000..a607e1d3 --- /dev/null +++ b/resolver/mocks.go @@ -0,0 +1,366 @@ +// Code generated by mockery; DO NOT EDIT. +// github.com/vektra/mockery +// template: testify + +package resolver + +import ( + "context" + + "github.com/0xERR0R/blocky/model" + "github.com/sirupsen/logrus" + mock "github.com/stretchr/testify/mock" +) + +// NewMockChainedResolver creates a new instance of MockChainedResolver. It also registers a testing interface on the mock and a cleanup function to assert the mocks expectations. +// The first argument is typically a *testing.T value. +func NewMockChainedResolver(t interface { + mock.TestingT + Cleanup(func()) +}) *MockChainedResolver { + mock := &MockChainedResolver{} + mock.Mock.Test(t) + + t.Cleanup(func() { mock.AssertExpectations(t) }) + + return mock +} + +// MockChainedResolver is an autogenerated mock type for the ChainedResolver type +type MockChainedResolver struct { + mock.Mock +} + +type MockChainedResolver_Expecter struct { + mock *mock.Mock +} + +func (_m *MockChainedResolver) EXPECT() *MockChainedResolver_Expecter { + return &MockChainedResolver_Expecter{mock: &_m.Mock} +} + +// GetNext provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) GetNext() Resolver { + ret := _mock.Called() + + if len(ret) == 0 { + panic("no return value specified for GetNext") + } + + var r0 Resolver + if returnFunc, ok := ret.Get(0).(func() Resolver); ok { + r0 = returnFunc() + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(Resolver) + } + } + return r0 +} + +// MockChainedResolver_GetNext_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'GetNext' +type MockChainedResolver_GetNext_Call struct { + *mock.Call +} + +// GetNext is a helper method to define mock.On call +func (_e *MockChainedResolver_Expecter) GetNext() *MockChainedResolver_GetNext_Call { + return &MockChainedResolver_GetNext_Call{Call: _e.mock.On("GetNext")} +} + +func (_c *MockChainedResolver_GetNext_Call) Run(run func()) *MockChainedResolver_GetNext_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockChainedResolver_GetNext_Call) Return(resolver Resolver) *MockChainedResolver_GetNext_Call { + _c.Call.Return(resolver) + return _c +} + +func (_c *MockChainedResolver_GetNext_Call) RunAndReturn(run func() Resolver) *MockChainedResolver_GetNext_Call { + _c.Call.Return(run) + return _c +} + +// IsEnabled provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) IsEnabled() bool { + ret := _mock.Called() + + if len(ret) == 0 { + panic("no return value specified for IsEnabled") + } + + var r0 bool + if returnFunc, ok := ret.Get(0).(func() bool); ok { + r0 = returnFunc() + } else { + r0 = ret.Get(0).(bool) + } + return r0 +} + +// MockChainedResolver_IsEnabled_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'IsEnabled' +type MockChainedResolver_IsEnabled_Call struct { + *mock.Call +} + +// IsEnabled is a helper method to define mock.On call +func (_e *MockChainedResolver_Expecter) IsEnabled() *MockChainedResolver_IsEnabled_Call { + return &MockChainedResolver_IsEnabled_Call{Call: _e.mock.On("IsEnabled")} +} + +func (_c *MockChainedResolver_IsEnabled_Call) Run(run func()) *MockChainedResolver_IsEnabled_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockChainedResolver_IsEnabled_Call) Return(b bool) *MockChainedResolver_IsEnabled_Call { + _c.Call.Return(b) + return _c +} + +func (_c *MockChainedResolver_IsEnabled_Call) RunAndReturn(run func() bool) *MockChainedResolver_IsEnabled_Call { + _c.Call.Return(run) + return _c +} + +// LogConfig provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) LogConfig(logger *logrus.Entry) { + _mock.Called(logger) + return +} + +// MockChainedResolver_LogConfig_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'LogConfig' +type MockChainedResolver_LogConfig_Call struct { + *mock.Call +} + +// LogConfig is a helper method to define mock.On call +// - logger *logrus.Entry +func (_e *MockChainedResolver_Expecter) LogConfig(logger interface{}) *MockChainedResolver_LogConfig_Call { + return &MockChainedResolver_LogConfig_Call{Call: _e.mock.On("LogConfig", logger)} +} + +func (_c *MockChainedResolver_LogConfig_Call) Run(run func(logger *logrus.Entry)) *MockChainedResolver_LogConfig_Call { + _c.Call.Run(func(args mock.Arguments) { + var arg0 *logrus.Entry + if args[0] != nil { + arg0 = args[0].(*logrus.Entry) + } + run( + arg0, + ) + }) + return _c +} + +func (_c *MockChainedResolver_LogConfig_Call) Return() *MockChainedResolver_LogConfig_Call { + _c.Call.Return() + return _c +} + +func (_c *MockChainedResolver_LogConfig_Call) RunAndReturn(run func(logger *logrus.Entry)) *MockChainedResolver_LogConfig_Call { + _c.Run(run) + return _c +} + +// Next provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) Next(n Resolver) { + _mock.Called(n) + return +} + +// MockChainedResolver_Next_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Next' +type MockChainedResolver_Next_Call struct { + *mock.Call +} + +// Next is a helper method to define mock.On call +// - n Resolver +func (_e *MockChainedResolver_Expecter) Next(n interface{}) *MockChainedResolver_Next_Call { + return &MockChainedResolver_Next_Call{Call: _e.mock.On("Next", n)} +} + +func (_c *MockChainedResolver_Next_Call) Run(run func(n Resolver)) *MockChainedResolver_Next_Call { + _c.Call.Run(func(args mock.Arguments) { + var arg0 Resolver + if args[0] != nil { + arg0 = args[0].(Resolver) + } + run( + arg0, + ) + }) + return _c +} + +func (_c *MockChainedResolver_Next_Call) Return() *MockChainedResolver_Next_Call { + _c.Call.Return() + return _c +} + +func (_c *MockChainedResolver_Next_Call) RunAndReturn(run func(n Resolver)) *MockChainedResolver_Next_Call { + _c.Run(run) + return _c +} + +// Resolve provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) Resolve(ctx context.Context, req *model.Request) (*model.Response, error) { + ret := _mock.Called(ctx, req) + + if len(ret) == 0 { + panic("no return value specified for Resolve") + } + + var r0 *model.Response + var r1 error + if returnFunc, ok := ret.Get(0).(func(context.Context, *model.Request) (*model.Response, error)); ok { + return returnFunc(ctx, req) + } + if returnFunc, ok := ret.Get(0).(func(context.Context, *model.Request) *model.Response); ok { + r0 = returnFunc(ctx, req) + } else { + if ret.Get(0) != nil { + r0 = ret.Get(0).(*model.Response) + } + } + if returnFunc, ok := ret.Get(1).(func(context.Context, *model.Request) error); ok { + r1 = returnFunc(ctx, req) + } else { + r1 = ret.Error(1) + } + return r0, r1 +} + +// MockChainedResolver_Resolve_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Resolve' +type MockChainedResolver_Resolve_Call struct { + *mock.Call +} + +// Resolve is a helper method to define mock.On call +// - ctx context.Context +// - req *model.Request +func (_e *MockChainedResolver_Expecter) Resolve(ctx interface{}, req interface{}) *MockChainedResolver_Resolve_Call { + return &MockChainedResolver_Resolve_Call{Call: _e.mock.On("Resolve", ctx, req)} +} + +func (_c *MockChainedResolver_Resolve_Call) Run(run func(ctx context.Context, req *model.Request)) *MockChainedResolver_Resolve_Call { + _c.Call.Run(func(args mock.Arguments) { + var arg0 context.Context + if args[0] != nil { + arg0 = args[0].(context.Context) + } + var arg1 *model.Request + if args[1] != nil { + arg1 = args[1].(*model.Request) + } + run( + arg0, + arg1, + ) + }) + return _c +} + +func (_c *MockChainedResolver_Resolve_Call) Return(response *model.Response, err error) *MockChainedResolver_Resolve_Call { + _c.Call.Return(response, err) + return _c +} + +func (_c *MockChainedResolver_Resolve_Call) RunAndReturn(run func(ctx context.Context, req *model.Request) (*model.Response, error)) *MockChainedResolver_Resolve_Call { + _c.Call.Return(run) + return _c +} + +// String provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) String() string { + ret := _mock.Called() + + if len(ret) == 0 { + panic("no return value specified for String") + } + + var r0 string + if returnFunc, ok := ret.Get(0).(func() string); ok { + r0 = returnFunc() + } else { + r0 = ret.Get(0).(string) + } + return r0 +} + +// MockChainedResolver_String_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'String' +type MockChainedResolver_String_Call struct { + *mock.Call +} + +// String is a helper method to define mock.On call +func (_e *MockChainedResolver_Expecter) String() *MockChainedResolver_String_Call { + return &MockChainedResolver_String_Call{Call: _e.mock.On("String")} +} + +func (_c *MockChainedResolver_String_Call) Run(run func()) *MockChainedResolver_String_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockChainedResolver_String_Call) Return(s string) *MockChainedResolver_String_Call { + _c.Call.Return(s) + return _c +} + +func (_c *MockChainedResolver_String_Call) RunAndReturn(run func() string) *MockChainedResolver_String_Call { + _c.Call.Return(run) + return _c +} + +// Type provides a mock function for the type MockChainedResolver +func (_mock *MockChainedResolver) Type() string { + ret := _mock.Called() + + if len(ret) == 0 { + panic("no return value specified for Type") + } + + var r0 string + if returnFunc, ok := ret.Get(0).(func() string); ok { + r0 = returnFunc() + } else { + r0 = ret.Get(0).(string) + } + return r0 +} + +// MockChainedResolver_Type_Call is a *mock.Call that shadows Run/Return methods with type explicit version for method 'Type' +type MockChainedResolver_Type_Call struct { + *mock.Call +} + +// Type is a helper method to define mock.On call +func (_e *MockChainedResolver_Expecter) Type() *MockChainedResolver_Type_Call { + return &MockChainedResolver_Type_Call{Call: _e.mock.On("Type")} +} + +func (_c *MockChainedResolver_Type_Call) Run(run func()) *MockChainedResolver_Type_Call { + _c.Call.Run(func(args mock.Arguments) { + run() + }) + return _c +} + +func (_c *MockChainedResolver_Type_Call) Return(s string) *MockChainedResolver_Type_Call { + _c.Call.Return(s) + return _c +} + +func (_c *MockChainedResolver_Type_Call) RunAndReturn(run func() string) *MockChainedResolver_Type_Call { + _c.Call.Return(run) + return _c +} diff --git a/resolver/rate_limit_store.go b/resolver/rate_limit_store.go new file mode 100644 index 00000000..1d93297c --- /dev/null +++ b/resolver/rate_limit_store.go @@ -0,0 +1,107 @@ +package resolver + +import ( + "context" + "net" + "sync" + "sync/atomic" + "time" + + "golang.org/x/time/rate" +) + +const ( + ipv4Bits = 32 + ipv6Bits = 128 +) + +func bucketKey(ip net.IP, v4Prefix, v6Prefix uint8) string { + if v4 := ip.To4(); v4 != nil { + return v4.Mask(net.CIDRMask(int(v4Prefix), ipv4Bits)).String() + } + + return ip.To16().Mask(net.CIDRMask(int(v6Prefix), ipv6Bits)).String() +} + +type bucketEntry struct { + limiter *rate.Limiter + lastLogged atomic.Int64 +} + +type bucketStore struct { + limit rate.Limit + burst int + maxBuckets int + buckets sync.Map + size atomic.Int64 + janitorDone chan struct{} // closed when the janitor goroutine exits; nil until startJanitor is called +} + +func newBucketStore(limit rate.Limit, burst, maxBuckets int) *bucketStore { + return &bucketStore{limit: limit, burst: burst, maxBuckets: maxBuckets} +} + +// allowAt is the time-injectable Allow path used by tests and by the resolver +// (which passes r.clock()). Returns (nil, false) when the store is full — +// callers can distinguish cap exhaustion from a rate-limit drop by checking +// the entry for nil. +func (s *bucketStore) allowAt(key string, now time.Time) (*bucketEntry, bool) { + if v, ok := s.buckets.Load(key); ok { + e := v.(*bucketEntry) + + return e, e.limiter.AllowN(now, 1) + } + // Reserve a slot atomically before inserting so a concurrent burst of + // fresh keys cannot push size past maxBuckets. Roll back the reservation + // if another goroutine wins the LoadOrStore race for the same key. + for { + cur := s.size.Load() + if cur >= int64(s.maxBuckets) { + return nil, false + } + if s.size.CompareAndSwap(cur, cur+1) { + break + } + } + fresh := &bucketEntry{limiter: rate.NewLimiter(s.limit, s.burst)} + actual, loaded := s.buckets.LoadOrStore(key, fresh) + if loaded { + s.size.Add(-1) + } + e := actual.(*bucketEntry) + + return e, e.limiter.AllowN(now, 1) +} + +// sweep walks the map and evicts buckets whose limiter is fully refilled +// (idle = no state worth keeping; reconstruction yields identical behavior). +func (s *bucketStore) sweep() { + s.buckets.Range(func(k, v any) bool { + if v.(*bucketEntry).limiter.Tokens() >= float64(s.burst) { + s.buckets.Delete(k) + s.size.Add(-1) + } + + return true + }) +} + +// startJanitor launches a background sweep loop that exits when ctx is done. +// The returned-via-field janitorDone channel is closed when the goroutine exits, +// giving tests a deterministic signal without polling runtime.NumGoroutine. +func (s *bucketStore) startJanitor(ctx context.Context, interval time.Duration) { + s.janitorDone = make(chan struct{}) + go func() { + defer close(s.janitorDone) + t := time.NewTicker(interval) + defer t.Stop() + for { + select { + case <-ctx.Done(): + return + case <-t.C: + s.sweep() + } + } + }() +} diff --git a/resolver/rate_limit_store_test.go b/resolver/rate_limit_store_test.go new file mode 100644 index 00000000..83a7b0ec --- /dev/null +++ b/resolver/rate_limit_store_test.go @@ -0,0 +1,145 @@ +package resolver + +import ( + "context" + "net" + "strconv" + "sync" + "time" + + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "golang.org/x/time/rate" +) + +var _ = Describe("bucketKey", func() { + It("uses /32 mask for IPv4", func() { + Expect(bucketKey(net.ParseIP("192.0.2.5"), 32, 64)).Should(Equal("192.0.2.5")) + }) + + It("aggregates IPv4 by configured prefix", func() { + a := bucketKey(net.ParseIP("192.0.2.5"), 24, 64) + b := bucketKey(net.ParseIP("192.0.2.99"), 24, 64) + Expect(a).Should(Equal(b)) + }) + + It("aggregates IPv6 by /64 default", func() { + a := bucketKey(net.ParseIP("2001:db8::1"), 32, 64) + b := bucketKey(net.ParseIP("2001:db8::ffff"), 32, 64) + Expect(a).Should(Equal(b)) + }) + + It("separates different IPv6 /64s", func() { + a := bucketKey(net.ParseIP("2001:db8:1::1"), 32, 64) + b := bucketKey(net.ParseIP("2001:db8:2::1"), 32, 64) + Expect(a).ShouldNot(Equal(b)) + }) + + It("normalises IPv4-mapped IPv6 to v4", func() { + mapped := net.ParseIP("::ffff:192.0.2.5") + v4 := net.ParseIP("192.0.2.5") + Expect(bucketKey(mapped, 32, 64)).Should(Equal(bucketKey(v4, 32, 64))) + }) +}) + +var _ = Describe("bucketStore", func() { + now := time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + + It("allows the first burst requests for a fresh key", func() { + s := newBucketStore(rate.Limit(1), 3, 1024) + for range 3 { + _, ok := s.allowAt("k", now) + Expect(ok).Should(BeTrue()) + } + _, ok := s.allowAt("k", now) + Expect(ok).Should(BeFalse()) + }) + + It("refills tokens over time", func() { + s := newBucketStore(rate.Limit(1), 1, 1024) + _, ok := s.allowAt("k", now) + Expect(ok).Should(BeTrue()) + _, ok = s.allowAt("k", now) + Expect(ok).Should(BeFalse()) + _, ok = s.allowAt("k", now.Add(time.Second)) + Expect(ok).Should(BeTrue()) + }) + + It("keeps separate buckets per key", func() { + s := newBucketStore(rate.Limit(1), 1, 1024) + _, okA := s.allowAt("a", now) + _, okB := s.allowAt("b", now) + Expect(okA).Should(BeTrue()) + Expect(okB).Should(BeTrue()) + Expect(s.size.Load()).Should(BeNumerically("==", 2)) + }) + + It("drops new keys once cap is reached", func() { + s := newBucketStore(rate.Limit(1), 1, 2) + _, _ = s.allowAt("a", now) + _, _ = s.allowAt("b", now) + _, ok := s.allowAt("c", now) + Expect(ok).Should(BeFalse()) + Expect(s.size.Load()).Should(BeNumerically("==", 2)) + }) + + It("sweep removes fully-refilled buckets", func() { + s := newBucketStore(rate.Limit(1), 1, 1024) + _, _ = s.allowAt("idle", now) // bucket created, 0 tokens left + _, _ = s.allowAt("idle", now.Add(time.Hour)) // refills (capped at burst=1) + s.sweep() + Expect(s.size.Load()).Should(BeNumerically("==", 0)) + }) + + It("sweep keeps partially-drained buckets", func() { + s := newBucketStore(rate.Limit(1), 5, 1024) + recent := time.Now() // Use current time so bucket appears partially drained + _, _ = s.allowAt("active", recent) // 4 tokens left, not full + s.sweep() + Expect(s.size.Load()).Should(BeNumerically("==", 1)) + }) + + It("never exceeds cap under concurrent burst of distinct fresh keys", func() { + const capacity = 8 + s := newBucketStore(rate.Limit(1000), 1, capacity) + + var wg sync.WaitGroup + const N = 200 + for i := range N { + wg.Add(1) + go func(i int) { + defer wg.Done() + _, _ = s.allowAt("k-"+strconv.Itoa(i), time.Now()) + }(i) + } + wg.Wait() + + Expect(s.size.Load()).Should(BeNumerically("==", capacity)) + }) + + It("janitor exits when context is cancelled", func() { + s := newBucketStore(rate.Limit(1), 1, 1024) + ctx, cancel := context.WithCancel(context.Background()) + + s.startJanitor(ctx, time.Hour) // interval is irrelevant; we only test the cancel path + Consistently(s.janitorDone).WithTimeout(50 * time.Millisecond).ShouldNot(BeClosed()) + + cancel() + Eventually(s.janitorDone).WithTimeout(2 * time.Second).Should(BeClosed()) + }) + + It("exactly one limiter wins under concurrent LoadOrStore for same fresh key", func() { + s := newBucketStore(rate.Limit(1000), 1000, 1024) + var wg sync.WaitGroup + const N = 200 + for range N { + wg.Add(1) + go func() { + defer wg.Done() + _, _ = s.allowAt("shared", time.Now()) + }() + } + wg.Wait() + Expect(s.size.Load()).Should(BeNumerically("==", 1)) + }) +}) diff --git a/resolver/rate_limiting_resolver.go b/resolver/rate_limiting_resolver.go new file mode 100644 index 00000000..2ac917ff --- /dev/null +++ b/resolver/rate_limiting_resolver.go @@ -0,0 +1,136 @@ +package resolver + +import ( + "context" + "errors" + "net" + "time" + + "github.com/0xERR0R/blocky/config" + "github.com/0xERR0R/blocky/log" + "github.com/0xERR0R/blocky/metrics" + "github.com/0xERR0R/blocky/model" + "github.com/0xERR0R/blocky/util" + + "github.com/prometheus/client_golang/prometheus" + "github.com/sirupsen/logrus" + "golang.org/x/time/rate" +) + +// ErrRateLimited signals that handleReq must write no response. +var ErrRateLimited = errors.New("client rate-limited") + +const ( + // rateLimitBucketCap bounds the in-memory bucket store; hardcoded for v1. + rateLimitBucketCap = 16384 + + // rateLimitJanitorInterval is how often idle buckets are reclaimed. + rateLimitJanitorInterval = time.Minute +) + +type RateLimitingResolver struct { + configurable[*config.RateLimit] + NextResolver + typed + + store *bucketStore + clock func() time.Time + logger *logrus.Entry + + drops *prometheus.CounterVec + capExhausted prometheus.Counter + activeBuckets prometheus.GaugeFunc +} + +func NewRateLimitingResolver(ctx context.Context, cfg config.RateLimit) *RateLimitingResolver { + r := &RateLimitingResolver{ + configurable: withConfig(&cfg), + typed: withType("rate-limiting"), + clock: time.Now, + logger: log.PrefixedLog("rate-limiting"), + } + if !cfg.IsEnabled() { + return r + } + r.store = newBucketStore(rate.Limit(cfg.Rate), int(cfg.Burst), rateLimitBucketCap) + r.store.startJanitor(ctx, rateLimitJanitorInterval) + + r.drops = prometheus.NewCounterVec( + prometheus.CounterOpts{ + Name: "blocky_rate_limit_drops_total", + Help: "Total number of DNS queries dropped by the rate limiter, by protocol.", + }, + []string{logFieldProtocol}, + ) + r.capExhausted = prometheus.NewCounter( + prometheus.CounterOpts{ + Name: "blocky_rate_limit_cap_exhausted_total", + Help: "Total number of queries dropped because the in-memory bucket store was full.", + }, + ) + r.activeBuckets = prometheus.NewGaugeFunc( + prometheus.GaugeOpts{ + Name: "blocky_rate_limit_active_buckets", + Help: "Number of token buckets currently held in memory.", + }, + func() float64 { return float64(r.store.size.Load()) }, + ) + metrics.RegisterMetric(r.drops) + metrics.RegisterMetric(r.capExhausted) + metrics.RegisterMetric(r.activeBuckets) + + return r +} + +func (r *RateLimitingResolver) Resolve(ctx context.Context, req *model.Request) (*model.Response, error) { + if !r.IsEnabled() { + return r.next.Resolve(ctx, req) + } + ip := req.ClientIP + if ip == nil || r.isAllowlisted(ip) { + return r.next.Resolve(ctx, req) + } + key := bucketKey(ip, r.cfg.IPv4Prefix, r.cfg.IPv6Prefix) + entry, allowed := r.store.allowAt(key, r.clock()) + if allowed { + return r.next.Resolve(ctx, req) + } + if entry == nil { + r.capExhausted.Inc() + } + r.recordDrop(req, entry) + + return nil, ErrRateLimited +} + +func (r *RateLimitingResolver) isAllowlisted(ip net.IP) bool { + for _, n := range r.cfg.ParsedAllowlist() { + if n.Contains(ip) { + return true + } + } + + return false +} + +func (r *RateLimitingResolver) recordDrop(req *model.Request, e *bucketEntry) { + r.drops.WithLabelValues(req.Protocol.String()).Inc() + if e == nil { + return + } + now := r.clock().UnixNano() + prev := e.lastLogged.Load() + if now-prev < int64(time.Second) || !e.lastLogged.CompareAndSwap(prev, now) { + return + } + fields := logrus.Fields{ + "client_ip": req.ClientIP, + logFieldProtocol: req.Protocol, + "qname": util.QuestionToString(req.Req.Question), + "bucket_tokens": e.limiter.Tokens(), + } + if len(req.Req.Question) > 0 { + fields["qtype"] = req.Req.Question[0].Qtype + } + r.logger.WithFields(fields).Warn("dropped query") +} diff --git a/resolver/rate_limiting_resolver_test.go b/resolver/rate_limiting_resolver_test.go new file mode 100644 index 00000000..8f81399e --- /dev/null +++ b/resolver/rate_limiting_resolver_test.go @@ -0,0 +1,199 @@ +package resolver + +import ( + "context" + "errors" + "net" + "time" + + "github.com/0xERR0R/blocky/config" + . "github.com/0xERR0R/blocky/helpertest" + "github.com/0xERR0R/blocky/log" + . "github.com/0xERR0R/blocky/model" + + "github.com/miekg/dns" + . "github.com/onsi/ginkgo/v2" + . "github.com/onsi/gomega" + "github.com/prometheus/client_golang/prometheus/testutil" + "github.com/stretchr/testify/mock" +) + +var _ = Describe("RateLimitingResolver", func() { + var ( + sut *RateLimitingResolver + sutConfig config.RateLimit + m *mockResolver + ctx context.Context + cancelFn context.CancelFunc + ) + + BeforeEach(func() { + ctx, cancelFn = context.WithCancel(context.Background()) + DeferCleanup(cancelFn) + }) + + JustBeforeEach(func() { + sut = NewRateLimitingResolver(ctx, sutConfig) + m = &mockResolver{} + m.On("Resolve", mock.Anything).Return(&Response{Res: new(dns.Msg)}, nil) + sut.Next(m) + }) + + Describe("Type", func() { + It("follows conventions", func() { + expectValidResolverType(sut) + }) + }) + + Describe("LogConfig", func() { + BeforeEach(func() { + sutConfig = config.RateLimit{Enable: true, Rate: 1, Burst: 1, IPv4Prefix: 32, IPv6Prefix: 64} + }) + It("emits something", func() { + logger, hook := log.NewMockEntry() + sut.LogConfig(logger) + Expect(hook.Calls).ShouldNot(BeEmpty()) + }) + }) + + When("disabled", func() { + BeforeEach(func() { + sutConfig = config.RateLimit{Enable: false} + }) + It("delegates every request to next", func() { + req := newRequestWithClient("example.com.", A, "1.2.3.4") + _, err := sut.Resolve(ctx, req) + Expect(err).Should(Succeed()) + Expect(m.Calls).Should(HaveLen(1)) + }) + It("IsEnabled is false", func() { + Expect(sut.IsEnabled()).Should(BeFalse()) + }) + }) + + When("enabled with allowlist", func() { + BeforeEach(func() { + sutConfig = config.RateLimit{ + Enable: true, Rate: 1, Burst: 1, + IPv4Prefix: 32, IPv6Prefix: 64, + Allowlist: []string{"10.0.0.0/8"}, + } + Expect(sutConfig.ValidateForTest()).Should(Succeed()) + }) + It("bypasses allowlisted clients regardless of bucket", func() { + req := newRequestWithClient("example.com.", A, "10.1.2.3") + for range 5 { + _, err := sut.Resolve(ctx, req) + Expect(err).Should(Succeed()) + } + }) + It("bypasses requests with nil ClientIP", func() { + req := newRequest("example.com.", A) // no client IP set + for range 5 { + _, err := sut.Resolve(ctx, req) + Expect(err).Should(Succeed()) + } + }) + }) + + When("enabled with rate=1 burst=1", func() { + var fakeNow time.Time + BeforeEach(func() { + fakeNow = time.Date(2026, 1, 1, 0, 0, 0, 0, time.UTC) + sutConfig = config.RateLimit{ + Enable: true, Rate: 1, Burst: 1, + IPv4Prefix: 32, IPv6Prefix: 64, + } + Expect(sutConfig.ValidateForTest()).Should(Succeed()) + }) + JustBeforeEach(func() { + sut.clock = func() time.Time { return fakeNow } + }) + + It("allows first request, drops second", func() { + req := newRequestWithClient("example.com.", A, "1.2.3.4") + _, err := sut.Resolve(ctx, req) + Expect(err).Should(Succeed()) + + _, err = sut.Resolve(ctx, req) + Expect(errors.Is(err, ErrRateLimited)).Should(BeTrue()) + Expect(m.Calls).Should(HaveLen(1)) // second never reached next + }) + + It("allows again after enough time has elapsed", func() { + req := newRequestWithClient("example.com.", A, "1.2.3.4") + _, _ = sut.Resolve(ctx, req) + fakeNow = fakeNow.Add(2 * time.Second) + _, err := sut.Resolve(ctx, req) + Expect(err).Should(Succeed()) + }) + + It("uses separate buckets for different /32s", func() { + r1 := newRequestWithClient("example.com.", A, "1.2.3.4") + r2 := newRequestWithClient("example.com.", A, "5.6.7.8") + _, e1 := sut.Resolve(ctx, r1) + _, e2 := sut.Resolve(ctx, r2) + Expect(e1).Should(Succeed()) + Expect(e2).Should(Succeed()) + }) + + It("increments the drops counter on each drop", func() { + req := newRequestWithClient("example.com.", A, "1.2.3.4") + _, _ = sut.Resolve(ctx, req) // allowed + _, _ = sut.Resolve(ctx, req) // drop 1 + _, _ = sut.Resolve(ctx, req) // drop 2 + Expect(testutil.ToFloat64(sut.drops.WithLabelValues("UDP"))). + Should(BeNumerically("==", 2)) + }) + + It("logs at most once per second per bucket", func() { + req := newRequestWithClient("example.com.", A, "1.2.3.4") + logger, hook := log.NewMockEntry() + sut.logger = logger + + _, _ = sut.Resolve(ctx, req) // allowed + _, _ = sut.Resolve(ctx, req) // drop, log #1 + _, _ = sut.Resolve(ctx, req) // drop in same second, no log + Expect(hook.Calls).Should(HaveLen(1)) + + fakeNow = fakeNow.Add(2 * time.Second) + _, _ = sut.Resolve(ctx, req) // allowed (token refilled after 2 s) + _, _ = sut.Resolve(ctx, req) // drop, log #2 (window expired) + Expect(hook.Calls).Should(HaveLen(2)) + }) + + It("active-buckets gauge reflects store size", func() { + req := newRequestWithClient("example.com.", A, "1.2.3.4") + _, _ = sut.Resolve(ctx, req) + Expect(testutil.ToFloat64(sut.activeBuckets)).Should(BeNumerically("==", 1)) + }) + + It("aggregates IPv6 by /64 into one bucket", func() { + r1 := newRequestWithClient("example.com.", A, "2001:db8::1") + r2 := newRequestWithClient("example.com.", A, "2001:db8::ffff") + _, e1 := sut.Resolve(ctx, r1) + _, e2 := sut.Resolve(ctx, r2) + Expect(e1).Should(Succeed()) + Expect(errors.Is(e2, ErrRateLimited)).Should(BeTrue()) + }) + + It("normalises IPv4-mapped IPv6 to one bucket with IPv4", func() { + r1 := newRequestWithClient("example.com.", A, "192.0.2.5") + r2 := newRequestWithClient("example.com.", A, "::ffff:192.0.2.5") + _, e1 := sut.Resolve(ctx, r1) + _, e2 := sut.Resolve(ctx, r2) + Expect(e1).Should(Succeed()) + Expect(errors.Is(e2, ErrRateLimited)).Should(BeTrue()) + }) + + It("does not panic when a drop is logged for a malformed empty-question request", func() { + req := &Request{ + ClientIP: net.ParseIP("1.2.3.4"), + Req: new(dns.Msg), + Protocol: RequestProtocolUDP, + } + _, _ = sut.Resolve(ctx, req) // first allowed (or dropped — bucket=1/1 from prior tests; doesn't matter) + Expect(func() { _, _ = sut.Resolve(ctx, req) }).ShouldNot(Panic()) + }) + }) +}) diff --git a/resolver/resolver.go b/resolver/resolver.go index 4c1ca24c..e8dabf33 100644 --- a/resolver/resolver.go +++ b/resolver/resolver.go @@ -21,6 +21,7 @@ const ( logFieldDomain = "domain" logFieldAnswer = "answer" logFieldUpstream = "upstream" + logFieldProtocol = "protocol" localhostName = "localhost" loopbackIPv4Str = "127.0.0.1" ) diff --git a/resolver/upstream_resolver.go b/resolver/upstream_resolver.go index 3d52bdae..6cdefecc 100644 --- a/resolver/upstream_resolver.go +++ b/resolver/upstream_resolver.go @@ -467,7 +467,7 @@ func (r *UpstreamResolver) logResponse( "return_code": dns.RcodeToString[resp.Rcode], logFieldUpstream: r.cfg.String(), "upstream_ip": ip.String(), - "protocol": request.Protocol, + logFieldProtocol: request.Protocol, "net": r.cfg.Net, "response_time_ms": rtt.Milliseconds(), }).Debugf("received response from upstream") diff --git a/server/server.go b/server/server.go index ed4ba7e3..1d602c70 100644 --- a/server/server.go +++ b/server/server.go @@ -416,6 +416,7 @@ func createQueryResolver( } r := resolver.Chain( + resolver.NewRateLimitingResolver(ctx, cfg.RateLimit), resolver.NewFilteringResolver(cfg.Filtering), resolver.NewFQDNOnlyResolver(cfg.FQDNOnly), clientNames, @@ -654,14 +655,16 @@ type msgWriter interface { func (s *Server) handleReq(ctx context.Context, request *model.Request, w msgWriter) { response, err := s.resolve(ctx, request) - if err != nil { + switch { + case errors.Is(err, resolver.ErrRateLimited): + return + case err != nil: log.FromCtx(ctx).Error("error on processing request:", err) - m := new(dns.Msg) m.SetRcode(request.Req, dns.RcodeServerFailure) err := w.WriteMsg(m) util.LogOnError(ctx, "can't write message: ", err) - } else { + default: err := w.WriteMsg(response.Res) util.LogOnError(ctx, "can't write message: ", err) } diff --git a/server/server_test.go b/server/server_test.go index db7a32c7..225f7890 100644 --- a/server/server_test.go +++ b/server/server_test.go @@ -22,10 +22,10 @@ import ( "github.com/0xERR0R/blocky/resolver" "github.com/0xERR0R/blocky/util" "github.com/creasty/defaults" + "github.com/miekg/dns" . "github.com/onsi/ginkgo/v2" . "github.com/onsi/gomega" - - "github.com/miekg/dns" + "github.com/stretchr/testify/mock" ) const ( @@ -994,6 +994,29 @@ var _ = Describe("Running DNS server", func() { Expect(sut.Certificates).ShouldNot(BeEmpty()) }) }) + + Describe("Rate-limited request handling", func() { + It("writes no response when resolver returns ErrRateLimited", func() { + w := &countingMsgWriter{} + + m := resolver.NewMockChainedResolver(GinkgoT()) + m.EXPECT().Resolve(mock.Anything, mock.Anything).Return(nil, resolver.ErrRateLimited) + + s := &Server{ + queryResolver: m, + cfg: &config.Config{Upstreams: config.Upstreams{ + Timeout: config.Duration(time.Second), + }}, + } + req := &model.Request{ + Req: util.NewMsgWithQuestion("example.com.", A), + Protocol: model.RequestProtocolUDP, + ClientIP: net.ParseIP("1.2.3.4"), + } + s.handleReq(context.Background(), req, w) + Expect(w.writes).Should(BeZero()) + }) + }) }) func requestServer(ctx context.Context, request *dns.Msg) *dns.Msg { @@ -1032,6 +1055,16 @@ func requestServer(ctx context.Context, request *dns.Msg) *dns.Msg { return nil } +type countingMsgWriter struct { + writes int +} + +func (w *countingMsgWriter) WriteMsg(*dns.Msg) error { + w.writes++ + + return nil +} + func writeCertPem(tmpDir *TmpFolder) *TmpFile { return tmpDir.CreateStringFile("cert.pem", "-----BEGIN CERTIFICATE-----",