Initial project setup - Phases 1-3 complete
This commit is contained in:
Trey t
298
ProxyCore/Sources/ProxyEngine/ConnectHandler.swift
Normal file
298
ProxyCore/Sources/ProxyEngine/ConnectHandler.swift
Normal file
@@ -0,0 +1,298 @@
|
||||
import Foundation
|
||||
import NIOCore
|
||||
import NIOPosix
|
||||
import NIOHTTP1
|
||||
|
||||
/// Handles incoming proxy requests:
|
||||
/// - HTTP CONNECT → establishes TCP tunnel (GlueHandler passthrough, or MITM in Phase 3)
|
||||
/// - Plain HTTP → connects upstream, forwards request, captures request+response
|
||||
final class ConnectHandler: ChannelInboundHandler, RemovableChannelHandler {
|
||||
typealias InboundIn = HTTPServerRequestPart
|
||||
typealias OutboundOut = HTTPServerResponsePart
|
||||
|
||||
private let trafficRepo: TrafficRepository
|
||||
|
||||
// Buffer request parts until we've connected upstream
|
||||
private var pendingHead: HTTPRequestHead?
|
||||
private var pendingBody: [ByteBuffer] = []
|
||||
private var pendingEnd: HTTPHeaders?
|
||||
private var receivedEnd = false
|
||||
|
||||
init(trafficRepo: TrafficRepository) {
|
||||
self.trafficRepo = trafficRepo
|
||||
}
|
||||
|
||||
func channelRead(context: ChannelHandlerContext, data: NIOAny) {
|
||||
let part = unwrapInboundIn(data)
|
||||
|
||||
switch part {
|
||||
case .head(let head):
|
||||
if head.method == .CONNECT {
|
||||
handleConnect(context: context, head: head)
|
||||
} else {
|
||||
pendingHead = head
|
||||
}
|
||||
case .body(let buffer):
|
||||
pendingBody.append(buffer)
|
||||
case .end(let trailers):
|
||||
if pendingHead != nil {
|
||||
pendingEnd = trailers
|
||||
receivedEnd = true
|
||||
handleHTTPRequest(context: context)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - CONNECT (HTTPS tunnel)
|
||||
|
||||
private func handleConnect(context: ChannelHandlerContext, head: HTTPRequestHead) {
|
||||
let components = head.uri.split(separator: ":")
|
||||
let host = String(components[0])
|
||||
let port = components.count > 1 ? Int(components[1]) ?? 443 : 443
|
||||
|
||||
// Check if this domain should be MITM'd (SSL Proxying enabled + domain in include list)
|
||||
let shouldMITM = shouldInterceptSSL(domain: host)
|
||||
|
||||
// Send 200 Connection Established
|
||||
let responseHead = HTTPResponseHead(version: .http1_1, status: .ok)
|
||||
context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
|
||||
context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil)
|
||||
|
||||
if shouldMITM {
|
||||
// MITM mode: strip HTTP handlers, install MITMHandler
|
||||
setupMITM(context: context, host: host, port: port)
|
||||
} else {
|
||||
// Passthrough mode: record domain-level entry, tunnel raw bytes
|
||||
recordConnectTraffic(host: host, port: port)
|
||||
|
||||
// We don't need to connect upstream ourselves — GlueHandler does raw forwarding
|
||||
// But GlueHandler pairs two channels, so we need the remote channel first
|
||||
ClientBootstrap(group: context.eventLoop)
|
||||
.channelOption(.socketOption(.so_reuseaddr), value: 1)
|
||||
.connect(host: host, port: port)
|
||||
.whenComplete { result in
|
||||
switch result {
|
||||
case .success(let remoteChannel):
|
||||
self.setupGlue(context: context, remoteChannel: remoteChannel)
|
||||
case .failure(let error):
|
||||
print("[Proxy] CONNECT passthrough failed to \(host):\(port): \(error)")
|
||||
context.close(promise: nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func shouldInterceptSSL(domain: String) -> Bool {
|
||||
guard IPCManager.shared.isSSLProxyingEnabled else { return false }
|
||||
guard CertificateManager.shared.hasCA else { return false }
|
||||
|
||||
// Check SSL proxying list from database
|
||||
let rulesRepo = RulesRepository()
|
||||
do {
|
||||
let entries = try rulesRepo.fetchAllSSLEntries()
|
||||
|
||||
// Check exclude list first
|
||||
for entry in entries where !entry.isInclude {
|
||||
if WildcardMatcher.matches(domain, pattern: entry.domainPattern) {
|
||||
return false
|
||||
}
|
||||
}
|
||||
|
||||
// Check include list
|
||||
for entry in entries where entry.isInclude {
|
||||
if WildcardMatcher.matches(domain, pattern: entry.domainPattern) {
|
||||
return true
|
||||
}
|
||||
}
|
||||
} catch {
|
||||
print("[Proxy] Failed to check SSL proxying list: \(error)")
|
||||
}
|
||||
|
||||
return false
|
||||
}
|
||||
|
||||
private func setupMITM(context: ChannelHandlerContext, host: String, port: Int) {
|
||||
let mitmHandler = MITMHandler(host: host, port: port, trafficRepo: trafficRepo)
|
||||
|
||||
// Remove HTTP handlers, keep raw bytes for MITMHandler
|
||||
context.channel.pipeline.handler(type: ByteToMessageHandler<HTTPRequestDecoder>.self)
|
||||
.whenSuccess { handler in
|
||||
context.channel.pipeline.removeHandler(handler, promise: nil)
|
||||
}
|
||||
|
||||
context.pipeline.removeHandler(context: context).whenComplete { _ in
|
||||
context.channel.pipeline.addHandler(mitmHandler).whenFailure { error in
|
||||
print("[Proxy] Failed to install MITM handler: \(error)")
|
||||
context.close(promise: nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
private func setupGlue(context: ChannelHandlerContext, remoteChannel: Channel) {
|
||||
let localGlue = GlueHandler()
|
||||
let remoteGlue = GlueHandler()
|
||||
localGlue.partner = remoteGlue
|
||||
remoteGlue.partner = localGlue
|
||||
|
||||
// Remove all HTTP handlers from the client channel, leaving raw bytes
|
||||
context.channel.pipeline.handler(type: ByteToMessageHandler<HTTPRequestDecoder>.self)
|
||||
.whenSuccess { handler in
|
||||
context.channel.pipeline.removeHandler(handler, promise: nil)
|
||||
}
|
||||
|
||||
context.pipeline.removeHandler(context: context).whenComplete { _ in
|
||||
context.channel.pipeline.addHandler(localGlue).whenSuccess {
|
||||
remoteChannel.pipeline.addHandler(remoteGlue).whenFailure { _ in
|
||||
context.close(promise: nil)
|
||||
remoteChannel.close(promise: nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - Plain HTTP forwarding
|
||||
|
||||
private func handleHTTPRequest(context: ChannelHandlerContext) {
|
||||
guard let head = pendingHead else { return }
|
||||
|
||||
// Parse host and port from the absolute URI or Host header
|
||||
guard let (host, port, path) = parseHTTPTarget(head: head) else {
|
||||
let responseHead = HTTPResponseHead(version: .http1_1, status: .badRequest)
|
||||
context.write(wrapOutboundOut(.head(responseHead)), promise: nil)
|
||||
context.writeAndFlush(wrapOutboundOut(.end(nil)), promise: nil)
|
||||
return
|
||||
}
|
||||
|
||||
// Rewrite the request URI to relative path (upstream expects /path, not http://host/path)
|
||||
var upstreamHead = head
|
||||
upstreamHead.uri = path
|
||||
// Ensure Host header is set
|
||||
if !upstreamHead.headers.contains(name: "Host") {
|
||||
upstreamHead.headers.add(name: "Host", value: host)
|
||||
}
|
||||
|
||||
let captureHandler = HTTPCaptureHandler(trafficRepo: trafficRepo, domain: host, scheme: "http")
|
||||
|
||||
ClientBootstrap(group: context.eventLoop)
|
||||
.channelOption(.socketOption(.so_reuseaddr), value: 1)
|
||||
.channelInitializer { channel in
|
||||
// Remote channel: decode HTTP responses, encode HTTP requests
|
||||
channel.pipeline.addHandler(HTTPRequestEncoder()).flatMap {
|
||||
channel.pipeline.addHandler(ByteToMessageHandler(HTTPResponseDecoder(leftOverBytesStrategy: .forwardBytes)))
|
||||
}.flatMap {
|
||||
channel.pipeline.addHandler(captureHandler)
|
||||
}.flatMap {
|
||||
channel.pipeline.addHandler(
|
||||
HTTPRelayHandler(clientContext: context, wrapResponse: self.wrapOutboundOut)
|
||||
)
|
||||
}
|
||||
}
|
||||
.connect(host: host, port: port)
|
||||
.whenComplete { result in
|
||||
switch result {
|
||||
case .success(let remoteChannel):
|
||||
// Forward the buffered request to upstream
|
||||
remoteChannel.write(NIOAny(HTTPClientRequestPart.head(upstreamHead)), promise: nil)
|
||||
for bodyBuffer in self.pendingBody {
|
||||
remoteChannel.write(NIOAny(HTTPClientRequestPart.body(.byteBuffer(bodyBuffer))), promise: nil)
|
||||
}
|
||||
remoteChannel.writeAndFlush(NIOAny(HTTPClientRequestPart.end(self.pendingEnd)), promise: nil)
|
||||
|
||||
// Clear buffered data
|
||||
self.pendingHead = nil
|
||||
self.pendingBody.removeAll()
|
||||
self.pendingEnd = nil
|
||||
|
||||
case .failure(let error):
|
||||
print("[Proxy] HTTP forward failed to \(host):\(port): \(error)")
|
||||
let responseHead = HTTPResponseHead(version: .http1_1, status: .badGateway)
|
||||
context.write(self.wrapOutboundOut(.head(responseHead)), promise: nil)
|
||||
context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - URL Parsing
|
||||
|
||||
private func parseHTTPTarget(head: HTTPRequestHead) -> (host: String, port: Int, path: String)? {
|
||||
// Absolute URI: "http://example.com:8080/path?query"
|
||||
if head.uri.hasPrefix("http://") || head.uri.hasPrefix("https://") {
|
||||
guard let url = URLComponents(string: head.uri) else { return nil }
|
||||
let host = url.host ?? ""
|
||||
let port = url.port ?? (head.uri.hasPrefix("https") ? 443 : 80)
|
||||
var path = url.path.isEmpty ? "/" : url.path
|
||||
if let query = url.query {
|
||||
path += "?\(query)"
|
||||
}
|
||||
return (host, port, path)
|
||||
}
|
||||
|
||||
// Relative URI with Host header
|
||||
if let hostHeader = head.headers.first(name: "Host") {
|
||||
let parts = hostHeader.split(separator: ":")
|
||||
let host = String(parts[0])
|
||||
let port = parts.count > 1 ? Int(parts[1]) ?? 80 : 80
|
||||
return (host, port, head.uri)
|
||||
}
|
||||
|
||||
return nil
|
||||
}
|
||||
|
||||
// MARK: - CONNECT traffic recording
|
||||
|
||||
private func recordConnectTraffic(host: String, port: Int) {
|
||||
var traffic = CapturedTraffic(
|
||||
domain: host,
|
||||
url: "https://\(host):\(port)",
|
||||
method: "CONNECT",
|
||||
scheme: "https",
|
||||
statusCode: 200,
|
||||
statusText: "Connection Established",
|
||||
startedAt: Date().timeIntervalSince1970,
|
||||
completedAt: Date().timeIntervalSince1970,
|
||||
durationMs: 0,
|
||||
isSslDecrypted: false
|
||||
)
|
||||
try? trafficRepo.insert(&traffic)
|
||||
IPCManager.shared.post(.newTrafficCaptured)
|
||||
}
|
||||
}
|
||||
|
||||
// MARK: - HTTPRelayHandler
|
||||
|
||||
/// Relays HTTP responses from the upstream server back to the proxy client.
|
||||
final class HTTPRelayHandler: ChannelInboundHandler, RemovableChannelHandler {
|
||||
typealias InboundIn = HTTPClientResponsePart
|
||||
|
||||
private let clientContext: ChannelHandlerContext
|
||||
private let wrapResponse: (HTTPServerResponsePart) -> NIOAny
|
||||
|
||||
init(clientContext: ChannelHandlerContext, wrapResponse: @escaping (HTTPServerResponsePart) -> NIOAny) {
|
||||
self.clientContext = clientContext
|
||||
self.wrapResponse = wrapResponse
|
||||
}
|
||||
|
||||
func channelRead(context: ChannelHandlerContext, data: NIOAny) {
|
||||
let part = unwrapInboundIn(data)
|
||||
|
||||
switch part {
|
||||
case .head(let head):
|
||||
let serverHead = HTTPResponseHead(version: head.version, status: head.status, headers: head.headers)
|
||||
clientContext.write(wrapResponse(.head(serverHead)), promise: nil)
|
||||
case .body(let buffer):
|
||||
clientContext.write(wrapResponse(.body(.byteBuffer(buffer))), promise: nil)
|
||||
case .end(let trailers):
|
||||
clientContext.writeAndFlush(wrapResponse(.end(trailers)), promise: nil)
|
||||
}
|
||||
}
|
||||
|
||||
func channelInactive(context: ChannelHandlerContext) {
|
||||
clientContext.close(promise: nil)
|
||||
}
|
||||
|
||||
func errorCaught(context: ChannelHandlerContext, error: Error) {
|
||||
print("[Proxy] Relay error: \(error)")
|
||||
context.close(promise: nil)
|
||||
clientContext.close(promise: nil)
|
||||
}
|
||||
}
|
||||
Reference in New Issue
Block a user