From 22f2faca32a615f8e95ddf7b85e892f8812648e7 Mon Sep 17 00:00:00 2001 From: WofWca Date: Thu, 16 Jan 2025 17:27:49 +0400 Subject: [PATCH] feat: inverse rate limit option Add an option to limit _outgoing_ traffic from the upstream's perspective, instead of limiting _incoming_ traffic. This is a bit of a quick-and-dirty hack, and I haven't checked whether it really integrates well with the rest of the app (e.g. logging), but it worked for my use case. --- README.md | 2 ++ config.go | 1 + main.go | 9 +++++++-- 3 files changed, 10 insertions(+), 2 deletions(-) diff --git a/README.md b/README.md index 0a26ba1..7c80074 100644 --- a/README.md +++ b/README.md @@ -21,6 +21,7 @@ configuration file can be a JSON file with the following format: "127.0.0.1:12000" ], "upstream": "127.0.0.1:12001", + "inverse": false, "rate": "10M", "burst": "64KiB", "per_client": false, @@ -49,6 +50,7 @@ name: my-beloved-app addrs: - 127.0.0.1:12000 upstream: 127.0.0.1:12001 +inverse: false rate: 10M burst: 64KiB per_client: false diff --git a/config.go b/config.go index d8744c6..d2f9e7c 100644 --- a/config.go +++ b/config.go @@ -67,6 +67,7 @@ type configuration struct { Name string `yaml:"name" flag:"name,,instance name"` Addrs StringSlice `yaml:"addrs" flag:"addrs,127.0.0.1:12000,bind addresses"` Upstream string `yaml:"upstream" flag:"upstream,,upstream address"` + Inverse bool `yaml:"inverse" flag:"inverse,,instead of limiting incoming traffic from the upstream's perspective, limit outgoing traffic"` Rate HumanBytes `yaml:"rate" flag:"rate,,incoming traffic rate limit"` Burst HumanBytes `yaml:"burst" flag:"burst,,allowed traffic burst"` PerClient bool `yaml:"per_client" flag:"per-client,,apply rate limit per client"` diff --git a/main.go b/main.go index 61631c4..4b88df0 100644 --- a/main.go +++ b/main.go @@ -66,8 +66,13 @@ func handleClient(c *configuration, conn net.Conn, limiter *rate.Limiter) { } } wg.Add(2) - go forward(conn, uconn, true) - go forward(uconn, conn, false) + if !c.Inverse { + go forward(conn, uconn, true) + go forward(uconn, conn, false) + } else { + go forward(conn, uconn, false) + go forward(uconn, conn, true) + } wg.Wait() }