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

import Hummingbird
import HummingbirdTesting
import Testing

@testable import Infrastructure

@Suite(
    "SecurityHeadersMiddleware middleware",
    .tags(.middleware)
)
struct SecurityHeadersMiddlewareTests {

    // MARK: Functional tests

    @Test
    func `applies the default security headers to a response`() async throws {
        try await app().test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.status == .ok)
                #expect(response.headers[.contentSecurityPolicy] == .Security.contentSecurityPolicy)
                #expect(response.headers[.xContentTypeOptions] == .Security.contentTypeOptions)
                #expect(response.headers[.frameOptions] == .Security.frameOptions)
                #expect(response.headers[.referrerPolicy] == .Security.referrerPolicy)
                #expect(response.headers[.permissionsPolicy] == .Security.permissionsPolicy)
            }
        }
    }

    @Test
    func `omits strict-transport-security by default`() async throws {
        try await app().test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.headers[.strictTransportSecurity] == nil)
            }
        }
    }

    @Test
    func `applies strict-transport-security when configured`() async throws {
        let value = "max-age=31536000; includeSubDomains"

        try await app(
            configuration: .init(strictTransportSecurity: value)
        ).test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.headers[.strictTransportSecurity] == value)
            }
        }
    }

    @Test
    func `applies a custom header value`() async throws {
        let value = "default-src 'none'"

        try await app(
            configuration: .init(contentSecurityPolicy: value)
        ).test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.headers[.contentSecurityPolicy] == value)
            }
        }
    }

    @Test
    func `omits a header whose configured value is nil`() async throws {
        try await app(
            configuration: .init(contentTypeOptions: nil)
        ).test(.router) { client in
            try await client.execute(
                uri: "/hello",
                method: .get
            ) { response in
                #expect(response.headers[.xContentTypeOptions] == nil)
            }
        }
    }

    @Test
    func `replaces an existing header value set downstream`() async throws {
        try await app().test(.router) { client in
            try await client.execute(
                uri: "/weak",
                method: .get
            ) { response in
                #expect(response.headers[.xContentTypeOptions] == .Security.contentTypeOptions)
            }
        }
    }

    @Test
    func `applies the security headers to an error response`() async throws {
        try await app().test(.router) { client in
            try await client.execute(
                uri: "/throws",
                method: .get
            ) { response in
                #expect(response.status == .badRequest)
                #expect(response.headers[.contentSecurityPolicy] == .Security.contentSecurityPolicy)
                #expect(response.headers[.xContentTypeOptions] == .Security.contentTypeOptions)
                #expect(response.headers[.frameOptions] == .Security.frameOptions)
                #expect(response.headers[.referrerPolicy] == .Security.referrerPolicy)
                #expect(response.headers[.permissionsPolicy] == .Security.permissionsPolicy)
            }
        }
    }

}

// MARK: - Helpers

private extension SecurityHeadersMiddlewareTests {

    // MARK: Methods

    /// Builds an application whose router applies the security-headers middleware ahead of three
    /// routes: `/hello` returns a plain body, `/weak` returns a response that already carries a
    /// deliberately weak `X-Content-Type-Options` value for the middleware to override, and
    /// `/throws` fails with an `HTTPError` the way the controllers do on invalid input.
    func app(
        configuration: SecurityHeadersMiddleware<BasicRequestContext>.Configuration = .init()
    ) -> some ApplicationProtocol {
        let router = Router()

        router.addMiddleware {
            SecurityHeadersMiddleware(configuration: configuration)
        }

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

        router.get("weak") { _, _ -> Response in
            var response = Response(status: .ok)

            response.headers[.xContentTypeOptions] = "weak"

            return response
        }

        router.get("throws") { _, _ -> Response in
            throw HTTPError(.badRequest)
        }

        return Application(router: router)
    }

}