Skip to content

Instantly share code, notes, and snippets.

@djbe
Last active July 13, 2022 23:47
Show Gist options
  • Save djbe/bd96e8ce8743f16f30b22035c9021b68 to your computer and use it in GitHub Desktop.
Save djbe/bd96e8ce8743f16f30b22035c9021b68 to your computer and use it in GitHub Desktop.
MitM for Vapor
#if canImport(Atlantis)
import Atlantis
#endif
import Vapor
#if canImport(Atlantis)
extension Application.Clients.Provider {
static var atlantis: Self {
.init { application in
application.clients.use {
let normal = EventLoopHTTPClient(http: $0.http.client.shared, eventLoop: $0.eventLoopGroup.next(), logger: $0.logger)
return AtlantisHTTPClient(original: normal)
}
}
}
}
private struct AtlantisHTTPClient: Client {
fileprivate let original: Client
var eventLoop: EventLoop {
original.eventLoop
}
func send(_ client: ClientRequest) -> EventLoopFuture<ClientResponse> {
let atlantisRequest = Request(
url: client.url.string,
method: client.method.string,
headers: client.headers.atlantisHeaders,
body: client.body.flatMap { Data(buffer: $0) }
)
return original.send(client).map { result in
let atlantisResponse = Response(statusCode: Int(result.status.code), headers: result.headers.atlantisHeaders)
Atlantis.add(request: atlantisRequest, response: atlantisResponse, responseBody: result.body.flatMap { Data(buffer: $0) })
return result
}
}
func delegating(to eventLoop: EventLoop) -> Client {
AtlantisHTTPClient(original: original.delegating(to: eventLoop))
}
func logging(to logger: Logger) -> Client {
AtlantisHTTPClient(original: original.logging(to: logger))
}
}
// Needed because both `HTTPClient.delegating(to:logger:)` and `EventLoopHTTPClient` are private
// So we need our own copy. Last checked on 2022-02-18
private struct EventLoopHTTPClient: Client {
let http: HTTPClient
let eventLoop: EventLoop
var logger: Logger?
func send(_ client: ClientRequest) -> EventLoopFuture<ClientResponse> {
let urlString = client.url.string
guard let url = URL(string: urlString) else {
logger?.debug("\(urlString) is an invalid URL")
return eventLoop.makeFailedFuture(Abort(.internalServerError, reason: "\(urlString) is an invalid URL"))
}
do {
let request = try HTTPClient.Request(
url: url,
method: client.method,
headers: client.headers,
body: client.body.map { .byteBuffer($0) }
)
return http.execute(request: request, eventLoop: .delegate(on: eventLoop), logger: logger)
.map { response in
ClientResponse(
status: response.status,
headers: response.headers,
body: response.body
)
}
} catch {
return eventLoop.makeFailedFuture(error)
}
}
func delegating(to eventLoop: EventLoop) -> Client {
EventLoopHTTPClient(http: http, eventLoop: eventLoop, logger: logger)
}
func logging(to logger: Logger) -> Client {
EventLoopHTTPClient(http: http, eventLoop: eventLoop, logger: logger)
}
}
#endif
#if canImport(Atlantis)
import Atlantis
#endif
import Vapor
#if canImport(Atlantis)
struct AtlantisMiddleware: AsyncMiddleware {
func respond(to request: Vapor.Request, chainingTo next: AsyncResponder) async throws -> Vapor.Response {
let bodyData = try await request.body.collect().get()
let atlantisRequest = Request(
url: request.url.atlantisURL(with: request.headers),
method: request.method.string,
headers: request.headers.atlantisHeaders,
body: bodyData.flatMap { Data(buffer: $0) }
)
do {
let result = try await next.respond(to: request)
let response = Response(statusCode: Int(result.status.code), headers: result.headers.atlantisHeaders)
Atlantis.add(request: atlantisRequest, response: response, responseBody: result.body.data)
return result
} catch let error as Abort {
let atlantisResponse = Response(statusCode: Int(error.status.code), headers: error.headers.atlantisHeaders)
Atlantis.add(request: atlantisRequest, response: atlantisResponse, responseBody: error.description.data(using: .utf8))
throw error
}
}
}
#endif
#if canImport(Atlantis)
import Atlantis
#endif
import Vapor
#if canImport(Atlantis)
extension HTTPHeaders {
var atlantisHeaders: [Header] {
map { Header(key: $0, value: $1) }
}
}
extension URI {
func atlantisURL(with headers: HTTPHeaders) -> String {
URI(scheme: "http", host: headers.first(name: .host), path: path, query: query, fragment: fragment).string
}
}
#endif

A couple of notes:

  • The files above are all in the Run target.
  • Add a Info.plist with a CFBundleIdentifier and CFBundleName, otherwise Atlantis won't show a nice app name. The unsafeFlags property in the Run target will embed it correctly.
  • This will only be used in debug mode on macOS. Deployment builds will not embed the Atlantis MitM.

Requests made using client to other hosts will be shown as expected in Proxyman. Incoming requests (and corresponding responses) will be show as if made to the current host.

// swift-tools-version:5.6
import PackageDescription
let package = Package(
name: "Example",
platforms: [
.macOS(.v12)
],
dependencies: [
.package(url: "https://github.com/ProxymanApp/atlantis", from: "1.15.0"),
.package(url: "https://github.com/vapor/vapor.git", from: "4.0.0")
],
targets: [
.executableTarget(
name: "Run",
dependencies: [
.target(name: "App"),
.product(name: "Atlantis", package: "atlantis", condition: .when(platforms: [.macOS]))
],
linkerSettings: [
.unsafeFlags(["-Xlinker", "-sectcreate", "-Xlinker", "__TEXT", "-Xlinker", "__info_plist", "-Xlinker", "Resources/Info.plist"], .when(platforms: [.macOS]))
]
),
.target(
name: "App",
dependencies: [
.product(name: "Vapor", package: "vapor")
],
swiftSettings: [
// Enable better optimizations when building in Release configuration. Despite the use of
// the `.unsafeFlags` construct required by SwiftPM, this flag is recommended for Release
// builds. See <https://github.com/swift-server/guides/blob/main/docs/building.md#building-for-production> for details.
.unsafeFlags(["-cross-module-optimization"], .when(configuration: .release))
]
)
]
)
import App
#if canImport(Atlantis)
import Atlantis
#endif
import Vapor
extension Application {
private static let runQueue = DispatchQueue(label: "Run")
/// We need to avoid blocking the main thread, so spin this off to a separate queue
func run() async throws {
try await withCheckedThrowingContinuation { (continuation: CheckedContinuation<Void, _>) in
Self.runQueue.async { [self] in
do {
try run()
continuation.resume()
} catch {
continuation.resume(throwing: error)
}
}
}
}
}
@main
enum Run {
static func main() async throws {
var env = try Environment.detect()
try LoggingSystem.bootstrap(from: &env)
let app = Application(env)
defer { app.shutdown() }
#if canImport(Atlantis)
Atlantis.start()
app.middleware.use(AtlantisMiddleware(), at: .beginning)
app.clients.use(.atlantis)
#endif
try await app.configure()
try await app.run()
}
}
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment