Bugfixes and improvements to rate limit check (#701)

- use the correct IP adress of the client
 - use DateTime(0) for first time connections
 - don't warn when capping the rate limit allowance
 - fix rate_limit specification with non-1 denominator

Project.toml

@@ -1,7 +1,7 @@
 name = "HTTP"
 uuid = "cd3eb016-35fb-5094-929b-558a96fad6f3"
 authors = ["Jacob Quinn", "Sam O'Connor", "contributors: https://github.com/JuliaWeb/HTTP.jl/graphs/contributors"]
-version = "0.9.5"
+version = "0.9.6"
 
 [deps]
 Base64 = "2a0f44e3-6c83-55bd-87e4-b1978d98bd5f"

src/Servers.jl

@@ -36,6 +36,9 @@ function update!(rl::RateLimit, rate_limit)
     timepassed = float(Dates.value(current - rl.lastcheck)) / 1000.0
     rl.lastcheck = current
     rl.allowance += timepassed * rate_limit
+    if rl.allowance > rate_limit
+        rl.allowance = rate_limit
+    end
     return nothing
 end
 
@@ -51,14 +54,9 @@ soon, it is closed and discarded, otherwise, the timestamp for the
 ip address is updated in the global cache.
 """
 function check_rate_limit(tcp, rate_limit::Rational{Int})
-    ip = Sockets.getsockname(tcp)[1]
-    rate = Float64(rate_limit.num)
-    rl = get!(RATE_LIMITS[Threads.threadid()], ip, RateLimit(rate, Dates.now()))
+    ip = Sockets.getpeername(tcp)[1]
+    rl = get!(RATE_LIMITS[Threads.threadid()], ip, RateLimit(rate_limit, Dates.DateTime(0)))
     update!(rl, rate_limit)
-    if rl.allowance > rate
-        @warn "throttling $ip"
-        rl.allowance = rate
-    end
     if rl.allowance < 1.0
         @warn "discarding connection from $ip due to rate limiting"
         return false

test/server.jl

@@ -191,6 +191,62 @@ end
 
 end # @testset
 
+@testset "HTTP.listen: rate_limit" begin
+    io = IOBuffer()
+    logger = Base.CoreLogging.SimpleLogger(io)
+    server = listen(IPv4(0), 8080)
+    @async Base.CoreLogging.with_logger(logger) do
+        HTTP.listen("0.0.0.0", 8080; server=server, rate_limit=2//1) do http
+            HTTP.setstatus(http, 200)
+            HTTP.setheader(http, "Content-Length" => "0")
+            HTTP.startwrite(http)
+            HTTP.close(http.stream) # close to force a new connection everytime
+        end
+    end
+    # Test requests from the same IP within the limit
+    for _ in 1:5
+        sleep(0.6) # rate limit allows 2 per second
+        @test HTTP.get("http://127.0.0.1:8080").status == 200
+    end
+    # Test requests from the same IP over the limit
+    try
+        for _ in 1:5
+            sleep(0.2) # rate limit allows 2 per second
+            r = HTTP.get("http://127.0.0.1:8080"; retry=false)
+            @test r.status == 200
+        end
+    catch e
+        @test e isa HTTP.IOExtras.IOError
+    end
+
+    close(server)
+    @test occursin("discarding connection from 127.0.0.1 due to rate limiting", String(take!(io)))
+
+    # # Tests to make sure the correct client IP is used (https://github.com/JuliaWeb/HTTP.jl/pull/701)
+    # # This test requires a second machine and thus commented out
+    #
+    # Machine 1
+    # @async HTTP.listen("0.0.0.0", 8080; rate_limit=2//1) do http
+    #     HTTP.setstatus(http, 200)
+    #     HTTP.setheader(http, "Content-Length" => "0")
+    #     HTTP.startwrite(http)
+    #     HTTP.close(http.stream) # close to force a new connection everytime
+    # end
+    # while true
+    #     sleep(0.6)
+    #     print("#") # to show some progress
+    #     HTTP.get("http://$(MACHINE_1_IPV4):8080"; retry=false)
+    # end
+    #
+    # # Machine 2
+    # while true
+    #     sleep(0.6)
+    #     print("#") # to show some progress
+    #     HTTP.get("http://$(MACHINE_1_IPV4):8080"; retry=false)
+    # end
+
+end
+
 @testset "on_shutdown" begin
     @test HTTP.Servers.shutdown(nothing) === nothing