ccn/Packages/Infrastructure/Tests/Cases/Public/Middlewares/RateLimitMiddlewareTests.swift

import Hummingbird
import HummingbirdTesting
import Testing

@testable import Infrastructure

@Suite(
    "RateLimitMiddleware middleware",
    .tags(.middleware)
)
struct RateLimitMiddlewareTests {

    // MARK: Functional tests

    @Test
    func `admits requests within the limit`() async throws {
        try await app(
            configuration: .init(limit: 3)
        ).test(.router) { client in
            for _ in 1 ... 3 {
                try await client.execute(
                    uri: "/hello",
                    method: .get
                ) { response in
                    #expect(response.status == .ok)
                }
            }
        }
    }

    @Test
    func `rejects a request over the limit with a retry-after header`() async throws {
        try await app(
            configuration: .init(limit: 2)
        ).test(.router) { client in
            for _ in 1 ... 2 {
                try await client.execute(
                    uri: "/hello",
                    method: .get
                ) { response in
                    #expect(response.status == .ok)
                }
            }

            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.status == .tooManyRequests)

                let retryAfter = try #require(response.headers[.retryAfter])

                #expect(try #require(Int(retryAfter)) >= 1)
            }
        }
    }

    @Test
    func `admits requests again once the window resets`() async throws {
        try await app(
            configuration: .init(
                limit: 1,
                window: .milliseconds(50)
            )
        ).test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.status == .ok)
            }
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.status == .tooManyRequests)
            }

            try await Task.sleep(for: .milliseconds(100))

            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.status == .ok)
            }
        }
    }

    @Test
    func `separates clients by their forwarded address when trusted`() async throws {
        try await app(
            configuration: .init(
                limit: 1,
                trustForwardedFor: true
            )
        ).test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get,
                headers: [.xForwardedFor: "203.0.113.7"]
            ) { response in
                #expect(response.status == .ok)
            }
            try await client.execute(
                uri: "/hello",
                method: .get,
                headers: [.xForwardedFor: "203.0.113.8"]
            ) { response in
                #expect(response.status == .ok)
            }
            // The first entry names the client; the appended proxy hop must not change its key.
            try await client.execute(
                uri: "/hello",
                method: .get,
                headers: [.xForwardedFor: "203.0.113.7, 10.0.0.1"]
            ) { response in
                #expect(response.status == .tooManyRequests)
            }
        }
    }

    @Test
    func `ignores the forwarded address when not trusted`() async throws {
        try await app(
            configuration: .init(limit: 1)
        ).test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get,
                headers: [.xForwardedFor: "203.0.113.7"]
            ) { response in
                #expect(response.status == .ok)
            }
            // Without trust (and without a connection address in router-only testing), every client
            // shares one bucket, so a rotated header must not mint a fresh budget.
            try await client.execute(
                uri: "/hello",
                method: .get,
                headers: [.xForwardedFor: "203.0.113.8"]
            ) { response in
                #expect(response.status == .tooManyRequests)
            }
        }
    }

}

// MARK: - Helpers

private extension RateLimitMiddlewareTests {

    // MARK: Methods

    /// Builds an application whose router applies the rate-limit middleware ahead of a single
    /// `/hello` route returning a plain body.
    func app(
        configuration: RateLimitMiddleware<BasicRequestContext>.Configuration
    ) -> some ApplicationProtocol {
        let router = Router()

        router.addMiddleware {
            RateLimitMiddleware(configuration: configuration)
        }

        router.get("hello") { _, _ in
            "Hello!"
        }

        return Application(router: router)
    }

}