fix(api): add user-agent header to appsync websocket handshake request (

#3586) * fix(api): add user-agent header to appsync websocket handshake request * fix broken integ test cases
aws-amplify · Mar 26, 2024 · 9640d5b · 9640d5b
1 parent 014d9b1
commit 9640d5b
Show file tree

Hide file tree

Showing 4 changed files with 18 additions and 8 deletions.
diff --git a/...fyPlugins/API/Sources/AWSAPIPlugin/SubscriptionFactory/AppSyncRealTimeClientFactory.swift b/...fyPlugins/API/Sources/AWSAPIPlugin/SubscriptionFactory/AppSyncRealTimeClientFactory.swift
@@ -61,7 +61,10 @@ actor AppSyncRealTimeClientFactory: AppSyncRealTimeClientFactoryProtocol {
                 requestInterceptor: authInterceptor,
                 webSocketClient: WebSocketClient(
                     url: Self.appSyncRealTimeEndpoint(endpoint),
-                    protocols: ["graphql-ws"],
+                    handshakeHttpHeaders: [
+                        URLRequestConstants.Header.webSocketSubprotocols: "graphql-ws",
+                        URLRequestConstants.Header.userAgent: AmplifyAWSServiceConfiguration.userAgentLib
+                    ],
                     interceptor: authInterceptor
                 )
             )

diff --git a/AmplifyPlugins/API/Sources/AWSAPIPlugin/Support/Constants/URLRequestConstants.swift b/AmplifyPlugins/API/Sources/AWSAPIPlugin/Support/Constants/URLRequestConstants.swift
@@ -18,6 +18,7 @@ struct URLRequestConstants {
         static let userAgent = "User-Agent"
         static let xApiKey = "x-api-key"
         static let host = "Host"
+        static let webSocketSubprotocols = "Sec-WebSocket-Protocol"
     }
 
     struct ContentType {

diff --git a/...Plugins/API/Tests/APIHostApp/AWSAPIPluginFunctionalTests/AppSyncRealTimeClientTests.swift b/...Plugins/API/Tests/APIHostApp/AWSAPIPluginFunctionalTests/AppSyncRealTimeClientTests.swift
@@ -47,7 +47,10 @@ class AppSyncRealTimeClientTests: XCTestCase {
 
             let webSocketClient = WebSocketClient(
                 url: AppSyncRealTimeClientFactory.appSyncRealTimeEndpoint(URL(string: endpoint)!),
-                protocols: ["graphql-ws"],
+                handshakeHttpHeaders: [
+                    URLRequestConstants.Header.webSocketSubprotocols: "graphql-ws",
+                    URLRequestConstants.Header.userAgent: AmplifyAWSServiceConfiguration.userAgentLib + " (intg-test)"
+                ],
                 interceptor: APIKeyAuthInterceptor(apiKey: apiKey)
             )
             appSyncRealTimeClient = AppSyncRealTimeClient(

diff --git a/AmplifyPlugins/Core/AWSPluginsCore/WebSocket/WebSocketClient.swift b/AmplifyPlugins/Core/AWSPluginsCore/WebSocket/WebSocketClient.swift
@@ -23,8 +23,8 @@ public final actor WebSocketClient: NSObject {
 
     /// WebSocket server endpoint
     private let url: URL
-    /// WebSocket subprotocols
-    private let protocols: [String]
+    /// Additional Header for WebSocket handshake http request
+    private let handshakeHttpHeaders: [String: String]
     /// Interceptor for appending additional info before makeing the connection
     private var interceptor: WebSocketInterceptor?
     /// Internal wriable WebSocketEvent data stream
@@ -68,12 +68,12 @@ public final actor WebSocketClient: NSObject {
      */
     public init(
         url: URL,
-        protocols: [String] = [],
+        handshakeHttpHeaders: [String: String] = [:],
         interceptor: WebSocketInterceptor? = nil,
         networkMonitor: WebSocketNetworkMonitorProtocol = AmplifyNetworkMonitor()
     ) {
         self.url = Self.useWebSocketProtocolScheme(url: url)
-        self.protocols = protocols
+        self.handshakeHttpHeaders = handshakeHttpHeaders
         self.interceptor = interceptor
         self.autoConnectOnNetworkStatusChange = false
         self.autoRetryOnConnectionFailure = false
@@ -156,9 +156,12 @@ public final actor WebSocketClient: NSObject {
     }
 
     private func createWebSocketConnection() async -> URLSessionWebSocketTask {
-        let urlSession = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
         let decoratedURL = (await self.interceptor?.interceptConnection(url: self.url)) ?? self.url
-        return urlSession.webSocketTask(with: decoratedURL, protocols: self.protocols)
+        var urlRequest = URLRequest(url: decoratedURL)
+        self.handshakeHttpHeaders.forEach { urlRequest.setValue($0.value, forHTTPHeaderField: $0.key) }
+
+        let urlSession = URLSession(configuration: .default, delegate: self, delegateQueue: nil)
+        return urlSession.webSocketTask(with: urlRequest)
     }
 
     private func createConnectionAndRead() async {