Allow CORS to be configured for gRPC Web (#1583)

glbrntt · Lukasa · web-flow · commit 4fd5a1054ff8 · 2023-04-05T12:26:39.000Z
Motivation:

The WebCORS handler unconditionally sets "Access-Control-Allow-Origin"
to "*" in response headers regardless of whether the request is a CORS
request or whether the client sends credentials. Moreover we don't
expose any knobs to control how CORS is configured.

Modifications:

- Add CORS configuration to the server and server builder
- Let the allowed origins be '.any' (i.e. '*") or '.only' (limited to
  the provided origins)
- Let the user configure what headers are permitted in responses.
- Let the user configure whether credentialed requests are accepted.

Result:

More control over CORS

Co-authored-by: Cory Benfield &lt;lukasa@apple.com&gt;
diff --git a/Sources/GRPC/GRPCServerPipelineConfigurator.swift b/Sources/GRPC/GRPCServerPipelineConfigurator.swift
@@ -191,7 +191,7 @@ final class GRPCServerPipelineConfigurator: ChannelInboundHandler, RemovableChan
       // we'll be on the right event loop and sync operations are fine.
       let sync = context.pipeline.syncOperations
       try sync.configureHTTPServerPipeline(withErrorHandling: true)
-      try sync.addHandler(WebCORSHandler())
+      try sync.addHandler(WebCORSHandler(configuration: self.configuration.webCORS))
       let scheme = self.configuration.tlsConfiguration == nil ? "http" : "https"
       try sync.addHandler(GRPCWebToHTTP2ServerCodec(scheme: scheme))
       // There's no need to normalize headers for HTTP/1.
diff --git a/Sources/GRPC/Server.swift b/Sources/GRPC/Server.swift
@@ -367,6 +367,9 @@ extension Server {
     /// the need to recalculate this dictionary each time we receive an rpc.
     internal var serviceProvidersByName: [Substring: CallHandlerProvider]
 
+    /// CORS configuration for gRPC-Web support.
+    public var webCORS = Configuration.CORS()
+
     #if canImport(NIOSSL)
     /// Create a `Configuration` with some pre-defined defaults.
     ///
@@ -401,11 +404,9 @@ extension Server {
     ) {
       self.target = target
       self.eventLoopGroup = eventLoopGroup
-      self
-        .serviceProvidersByName = Dictionary(
-          uniqueKeysWithValues: serviceProviders
-            .map { ($0.serviceName, $0) }
-        )
+      self.serviceProvidersByName = Dictionary(
+        uniqueKeysWithValues: serviceProviders.map { ($0.serviceName, $0) }
+      )
       self.errorDelegate = errorDelegate
       self.tlsConfiguration = tls.map { GRPCTLSConfiguration(transforming: $0) }
       self.connectionKeepalive = connectionKeepalive
@@ -451,6 +452,55 @@ extension Server {
   }
 }
 
+extension Server.Configuration {
+  public struct CORS: Hashable, GRPCSendable {
+    /// Determines which 'origin' header field values are permitted in a CORS request.
+    public var allowedOrigins: AllowedOrigins
+    /// Sets the headers which are permitted in a response to a CORS request.
+    public var allowedHeaders: [String]
+    /// Enabling this value allows sets the "access-control-allow-credentials" header field
+    /// to "true" in respones to CORS requests. This must be enabled if the client intends to send
+    /// credentials.
+    public var allowCredentialedRequests: Bool
+    /// The maximum age in seconds which pre-flight CORS requests may be cached for.
+    public var preflightCacheExpiration: Int
+
+    public init(
+      allowedOrigins: AllowedOrigins = .all,
+      allowedHeaders: [String] = ["content-type", "x-grpc-web", "x-user-agent"],
+      allowCredentialedRequests: Bool = false,
+      preflightCacheExpiration: Int = 86400
+    ) {
+      self.allowedOrigins = allowedOrigins
+      self.allowedHeaders = allowedHeaders
+      self.allowCredentialedRequests = allowCredentialedRequests
+      self.preflightCacheExpiration = preflightCacheExpiration
+    }
+  }
+}
+
+extension Server.Configuration.CORS {
+  public struct AllowedOrigins: Hashable, Sendable {
+    enum Wrapped: Hashable, Sendable {
+      case all
+      case only([String])
+    }
+
+    private(set) var wrapped: Wrapped
+    private init(_ wrapped: Wrapped) {
+      self.wrapped = wrapped
+    }
+
+    /// Allow all origin values.
+    public static let all = Self(.all)
+
+    /// Allow only the given origin values.
+    public static func only(_ allowed: [String]) -> Self {
+      return Self(.only(allowed))
+    }
+  }
+}
+
 extension ServerBootstrapProtocol {
   fileprivate func bind(to target: BindTarget) -> EventLoopFuture<Channel> {
     switch target.wrapped {
diff --git a/Sources/GRPC/ServerBuilder.swift b/Sources/GRPC/ServerBuilder.swift
@@ -165,6 +165,15 @@ extension Server.Builder {
   }
 }
 
+extension Server.Builder {
+  /// Set the CORS configuration for gRPC Web.
+  @discardableResult
+  public func withCORSConfiguration(_ configuration: Server.Configuration.CORS) -> Self {
+    self.configuration.webCORS = configuration
+    return self
+  }
+}
+
 extension Server.Builder {
   /// Sets the root server logger. Accepted connections will branch from this logger and RPCs on
   /// each connection will use a logger branched from the connections logger. This logger is made
diff --git a/Sources/GRPC/WebCORSHandler.swift b/Sources/GRPC/WebCORSHandler.swift
@@ -17,54 +17,110 @@ import NIOCore
 import NIOHTTP1
 
 /// Handler that manages the CORS protocol for requests incoming from the browser.
-internal class WebCORSHandler {
-  var requestMethod: HTTPMethod?
+internal final class WebCORSHandler {
+  let configuration: Server.Configuration.CORS
+
+  private var state: State = .idle
+  private enum State: Equatable {
+    /// Starting state.
+    case idle
+    /// CORS preflight request is in progress.
+    case processingPreflightRequest
+    /// "Real" request is in progress.
+    case processingRequest(origin: String?)
+  }
+
+  init(configuration: Server.Configuration.CORS) {
+    self.configuration = configuration
+  }
 }
 
 extension WebCORSHandler: ChannelInboundHandler {
   typealias InboundIn = HTTPServerRequestPart
+  typealias InboundOut = HTTPServerRequestPart
   typealias OutboundOut = HTTPServerResponsePart
 
   func channelRead(context: ChannelHandlerContext, data: NIOAny) {
-    // If the request is OPTIONS, the request is not propagated further.
     switch self.unwrapInboundIn(data) {
-    case let .head(requestHead):
-      self.requestMethod = requestHead.method
-      if self.requestMethod == .OPTIONS {
-        var headers = HTTPHeaders()
-        headers.add(name: "Access-Control-Allow-Origin", value: "*")
-        headers.add(name: "Access-Control-Allow-Methods", value: "POST")
-        headers.add(
-          name: "Access-Control-Allow-Headers",
-          value: "content-type,x-grpc-web,x-user-agent"
-        )
-        headers.add(name: "Access-Control-Max-Age", value: "86400")
-        context.write(
-          self.wrapOutboundOut(.head(HTTPResponseHead(
-            version: requestHead.version,
-            status: .ok,
-            headers: headers
-          ))),
-          promise: nil
-        )
-        return
+    case let .head(head):
+      self.receivedRequestHead(context: context, head)
+
+    case let .body(body):
+      self.receivedRequestBody(context: context, body)
+
+    case let .end(trailers):
+      self.receivedRequestEnd(context: context, trailers)
+    }
+  }
+
+  private func receivedRequestHead(context: ChannelHandlerContext, _ head: HTTPRequestHead) {
+    if head.method == .OPTIONS,
+       head.headers.contains(.accessControlRequestMethod),
+       let origin = head.headers.first(name: "origin") {
+      // If the request is OPTIONS with a access-control-request-method header it's a CORS
+      // preflight request and is not propagated further.
+      self.state = .processingPreflightRequest
+      self.handlePreflightRequest(context: context, head: head, origin: origin)
+    } else {
+      self.state = .processingRequest(origin: head.headers.first(name: "origin"))
+      context.fireChannelRead(self.wrapInboundOut(.head(head)))
+    }
+  }
+
+  private func receivedRequestBody(context: ChannelHandlerContext, _ body: ByteBuffer) {
+    // OPTIONS requests do not have a body, but still handle this case to be
+    // cautious.
+    if self.state == .processingPreflightRequest {
+      return
+    }
+
+    context.fireChannelRead(self.wrapInboundOut(.body(body)))
+  }
+
+  private func receivedRequestEnd(context: ChannelHandlerContext, _ trailers: HTTPHeaders?) {
+    if self.state == .processingPreflightRequest {
+      // End of OPTIONS request; reset state and finish the response.
+      self.state = .idle
+      context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
+    } else {
+      context.fireChannelRead(self.wrapInboundOut(.end(trailers)))
+    }
+  }
+
+  private func handlePreflightRequest(
+    context: ChannelHandlerContext,
+    head: HTTPRequestHead,
+    origin: String
+  ) {
+    let responseHead: HTTPResponseHead
+
+    if let allowedOrigin = self.configuration.allowedOrigins.header(origin) {
+      var headers = HTTPHeaders()
+      headers.reserveCapacity(4 + self.configuration.allowedHeaders.count)
+      headers.add(name: .accessControlAllowOrigin, value: allowedOrigin)
+      headers.add(name: .accessControlAllowMethods, value: "POST")
+
+      for value in self.configuration.allowedHeaders {
+        headers.add(name: .accessControlAllowHeaders, value: value)
       }
-    case .body:
-      if self.requestMethod == .OPTIONS {
-        // OPTIONS requests do not have a body, but still handle this case to be
-        // cautious.
-        return
+
+      if self.configuration.allowCredentialedRequests {
+        headers.add(name: .accessControlAllowCredentials, value: "true")
       }
 
-    case .end:
-      if self.requestMethod == .OPTIONS {
-        context.writeAndFlush(self.wrapOutboundOut(.end(nil)), promise: nil)
-        self.requestMethod = nil
-        return
+      if self.configuration.preflightCacheExpiration > 0 {
+        headers.add(
+          name: .accessControlMaxAge,
+          value: "\(self.configuration.preflightCacheExpiration)"
+        )
       }
+      responseHead = HTTPResponseHead(version: head.version, status: .ok, headers: headers)
+    } else {
+      // Not allowed; respond with 403. This is okay in a pre-flight request.
+      responseHead = HTTPResponseHead(version: head.version, status: .forbidden)
     }
-    // The OPTIONS request should be fully handled at this point.
-    context.fireChannelRead(data)
+
+    context.write(self.wrapOutboundOut(.head(responseHead)), promise: nil)
   }
 }
 
@@ -74,25 +130,76 @@ extension WebCORSHandler: ChannelOutboundHandler {
   func write(context: ChannelHandlerContext, data: NIOAny, promise: EventLoopPromise<Void>?) {
     let responsePart = self.unwrapOutboundIn(data)
     switch responsePart {
-    case let .head(responseHead):
-      var headers = responseHead.headers
-      // CORS requires all requests to have an Allow-Origin header.
-      headers.add(name: "Access-Control-Allow-Origin", value: "*")
-      //! FIXME: Check whether we can let browsers keep connections alive. It's not possible
-      // now as the channel has a state that can't be reused since the pipeline is modified to
-      // inject the gRPC call handler.
-      headers.add(name: "Connection", value: "close")
-
-      context.write(
-        self.wrapOutboundOut(.head(HTTPResponseHead(
-          version: responseHead.version,
-          status: responseHead.status,
-          headers: headers
-        ))),
-        promise: promise
-      )
-    default:
+    case var .head(responseHead):
+      switch self.state {
+      case let .processingRequest(origin):
+        self.prepareCORSResponseHead(&responseHead, origin: origin)
+        context.write(self.wrapOutboundOut(.head(responseHead)), promise: promise)
+
+      case .idle, .processingPreflightRequest:
+        assertionFailure("Writing response head when no request is in progress")
+        context.close(promise: nil)
+      }
+
+    case .body:
+      context.write(data, promise: promise)
+
+    case .end:
+      self.state = .idle
       context.write(data, promise: promise)
     }
   }
+
+  private func prepareCORSResponseHead(_ head: inout HTTPResponseHead, origin: String?) {
+    guard let header = origin.flatMap({ self.configuration.allowedOrigins.header($0) }) else {
+      // No origin or the origin is not allowed; don't treat it as a CORS request.
+      return
+    }
+
+    head.headers.replaceOrAdd(name: .accessControlAllowOrigin, value: header)
+
+    if self.configuration.allowCredentialedRequests {
+      head.headers.add(name: .accessControlAllowCredentials, value: "true")
+    }
+
+    //! FIXME: Check whether we can let browsers keep connections alive. It's not possible
+    // now as the channel has a state that can't be reused since the pipeline is modified to
+    // inject the gRPC call handler.
+    head.headers.replaceOrAdd(name: "Connection", value: "close")
+  }
+}
+
+extension HTTPHeaders {
+  fileprivate enum CORSHeader: String {
+    case accessControlRequestMethod = "access-control-request-method"
+    case accessControlRequestHeaders = "access-control-request-headers"
+    case accessControlAllowOrigin = "access-control-allow-origin"
+    case accessControlAllowMethods = "access-control-allow-methods"
+    case accessControlAllowHeaders = "access-control-allow-headers"
+    case accessControlAllowCredentials = "access-control-allow-credentials"
+    case accessControlMaxAge = "access-control-max-age"
+  }
+
+  fileprivate func contains(_ name: CORSHeader) -> Bool {
+    return self.contains(name: name.rawValue)
+  }
+
+  fileprivate mutating func add(name: CORSHeader, value: String) {
+    self.add(name: name.rawValue, value: value)
+  }
+
+  fileprivate mutating func replaceOrAdd(name: CORSHeader, value: String) {
+    self.replaceOrAdd(name: name.rawValue, value: value)
+  }
+}
+
+extension Server.Configuration.CORS.AllowedOrigins {
+  internal func header(_ origin: String) -> String? {
+    switch self.wrapped {
+    case .all:
+      return "*"
+    case let .only(allowed):
+      return allowed.contains(origin) ? origin : nil
+    }
+  }
 }
diff --git a/Tests/GRPCTests/WebCORSHandlerTests.swift b/Tests/GRPCTests/WebCORSHandlerTests.swift
