@@ -4,6 +4,13 @@ import { HttpRequest } from '@smithy/protocol-http'
4
4
import { SignatureV4 } from '@smithy/signature-v4'
5
5
import { fromNodeProviderChain } from '@aws-sdk/credential-providers'
6
6
import { Sha256 } from '@aws-crypto/sha256-js'
7
+ import {
8
+ CognitoIdentityProviderClient ,
9
+ SignUpCommand ,
10
+ AdminConfirmSignUpCommand ,
11
+ AdminDeleteUserCommand ,
12
+ AdminInitiateAuthCommand ,
13
+ } from '@aws-sdk/client-cognito-identity-provider' ;
7
14
8
15
// The default headers to to sign the request
9
16
const DEFAULT_HEADERS = {
@@ -16,6 +23,56 @@ const AWS_APPSYNC_EVENTS_SUBPROTOCOL = 'aws-appsync-event-ws';
16
23
const realtimeUrl = process . env . EVENT_API_REALTIME_URL ;
17
24
const httpUrl = process . env . EVENT_API_HTTP_URL ;
18
25
const region = process . env . AWS_REGION ;
26
+ const API_KEY = process . env . API_KEY ;
27
+ const USER_POOL_ID = process . env . USER_POOL_ID ;
28
+ const CLIENT_ID = process . env . CLIENT_ID ;
29
+ const { username, password } = generateUsernamePassword ( 12 ) ;
30
+
31
+ const cognitoClient = new CognitoIdentityProviderClient ( ) ;
32
+
33
+ /**
34
+ * Utility function for generating a temporary password
35
+ * @param {int } length
36
+ * @returns
37
+ */
38
+ function generateUsernamePassword ( length ) {
39
+ const uppercaseChars = 'ABCDEFGHIJKLMNOPQRSTUVWXYZ' ;
40
+ const lowercaseChars = 'abcdefghijklmnopqrstuvwxyz' ;
41
+ const numberChars = '0123456789' ;
42
+ const specialChars = '!@#$%&' ;
43
+ const allChars = uppercaseChars + lowercaseChars + numberChars + specialChars ;
44
+
45
+ // Ensure length is at least 4 to accommodate required characters
46
+ const actualLength = Math . max ( length , 4 ) ;
47
+
48
+ // Start with one character from each required set
49
+ let password = [
50
+ uppercaseChars . charAt ( Math . floor ( Math . random ( ) * uppercaseChars . length ) ) ,
51
+ lowercaseChars . charAt ( Math . floor ( Math . random ( ) * lowercaseChars . length ) ) ,
52
+ numberChars . charAt ( Math . floor ( Math . random ( ) * numberChars . length ) ) ,
53
+ specialChars . charAt ( Math . floor ( Math . random ( ) * specialChars . length ) )
54
+ ] ;
55
+
56
+ // Fill the rest with random characters
57
+ for ( let i = 4 ; i < actualLength ; i ++ ) {
58
+ const randomIndex = Math . floor ( Math . random ( ) * allChars . length ) ;
59
+ password . push ( allChars . charAt ( randomIndex ) ) ;
60
+ }
61
+
62
+ // Shuffle the password array to randomize character positions
63
+ for ( let i = password . length - 1 ; i > 0 ; i -- ) {
64
+ const j = Math . floor ( Math . random ( ) * ( i + 1 ) ) ;
65
+ [ password [ i ] , password [ j ] ] = [ password [ j ] , password [ i ] ] ;
66
+ }
67
+
68
+ let username = '' ;
69
+ for ( let i = 0 ; i < 6 ; i ++ ) {
70
+ const randomIndex = Math . floor ( Math . random ( ) * lowercaseChars . length ) ;
71
+ username += lowercaseChars . charAt ( randomIndex ) ;
72
+ }
73
+
74
+ return { username, password : password . join ( '' ) } ;
75
+ }
19
76
20
77
/**
21
78
* Returns a signed authorization object
@@ -33,7 +90,7 @@ async function signWithAWSV4(httpDomain, region, body) {
33
90
sha256 : Sha256 ,
34
91
} )
35
92
36
- const url = new URL ( `https:// ${ httpDomain } /event ` )
93
+ const url = new URL ( `${ httpDomain } ` )
37
94
const request = new HttpRequest ( {
38
95
method : 'POST' ,
39
96
headers : {
@@ -55,13 +112,11 @@ async function signWithAWSV4(httpDomain, region, body) {
55
112
56
113
/**
57
114
* Returns a header value for the SubProtocol header
58
- * @param {string } httpDomain the AppSync Event API HTTP domain
59
- * @param {string } region the AWS region of your API
115
+ * @param {string } authHeaders the authorization headers
60
116
* @returns string a header string
61
117
*/
62
- async function getAuthProtocolForIAM ( httpDomain , region ) {
63
- const signed = await signWithAWSV4 ( httpDomain , region )
64
- const based64UrlHeader = btoa ( JSON . stringify ( signed ) )
118
+ function getAuthProtocolForIAM ( authHeaders ) {
119
+ const based64UrlHeader = btoa ( JSON . stringify ( authHeaders ) )
65
120
. replace ( / \+ / g, '-' ) // Convert '+' to '-'
66
121
. replace ( / \/ / g, '_' ) // Convert '/' to '_'
67
122
. replace ( / = + $ / , '' ) // Remove padding `=`
@@ -78,26 +133,114 @@ function sleep(ms) {
78
133
return new Promise ( resolve => setTimeout ( resolve , ms ) ) ;
79
134
}
80
135
136
+ /**
137
+ * Helper function for creating a Cognito user and confirming the user
138
+ * The function also deletes the user after the test is complete
139
+ * and it initiates and auth flow to get the ID token for testing the
140
+ * Event API auth flow with Cognito.
141
+ * @param {string } action - CREATE, DELETE, AUTH
142
+ * @returns
143
+ */
144
+ async function cognitoUserConfiguration ( action ) {
145
+ switch ( action ) {
146
+ case 'CREATE' :
147
+ const signUpUserInput = {
148
+ ClientId : CLIENT_ID ,
149
+ Username : username ,
150
+ Password : password ,
151
+ } ;
152
+ const signUpCommand = new SignUpCommand ( signUpUserInput ) ;
153
+ await cognitoClient . send ( signUpCommand ) ;
154
+ const confirmSignUpInput = {
155
+ UserPoolId : USER_POOL_ID ,
156
+ Username : username ,
157
+ } ;
158
+ const confirmSignUpCommand = new AdminConfirmSignUpCommand ( confirmSignUpInput ) ;
159
+ await cognitoClient . send ( confirmSignUpCommand ) ;
160
+ return { } ;
161
+ case 'DELETE' :
162
+ const deleteUserInput = {
163
+ UserPoolId : USER_POOL_ID ,
164
+ Username : username ,
165
+ } ;
166
+ const deleteUserCommand = new AdminDeleteUserCommand ( deleteUserInput ) ;
167
+ await cognitoClient . send ( deleteUserCommand ) ;
168
+ return ;
169
+ case 'AUTH' :
170
+ const authInput = {
171
+ UserPoolId : USER_POOL_ID ,
172
+ ClientId : CLIENT_ID ,
173
+ AuthFlow : 'ADMIN_USER_PASSWORD_AUTH' ,
174
+ AuthParameters : {
175
+ USERNAME : username ,
176
+ PASSWORD : password ,
177
+ } ,
178
+ } ;
179
+ const authCommand = new AdminInitiateAuthCommand ( authInput ) ;
180
+ const authRes = await cognitoClient . send ( authCommand ) ;
181
+ return authRes . AuthenticationResult . IdToken ;
182
+ }
183
+ }
184
+
185
+ /**
186
+ * Returns the appropriate headers depending on the auth mode selected
187
+ * @param {* } authMode - IAM, API_KEY, LAMBDA, USER_POOL, OIDC
188
+ * @param {* } event - the event payload for Publish operations, null by default
189
+ * @param {* } authToken - the token for LAMBDA auth modes
190
+ * @returns
191
+ */
192
+ async function getPublishAuthHeader ( authMode , event = { } , authToken = '' ) {
193
+ const url = new URL ( `${ httpUrl } ` )
194
+ const headers = {
195
+ host : url . hostname ,
196
+ } ;
197
+
198
+ switch ( authMode ) {
199
+ case 'IAM' :
200
+ return await signWithAWSV4 ( httpUrl , region , JSON . stringify ( event ) ) ;
201
+ case 'API_KEY' :
202
+ return {
203
+ 'x-api-key' : `${ API_KEY } ` ,
204
+ ...headers ,
205
+ }
206
+ case 'USER_POOL' :
207
+ return {
208
+ 'Authorization' : await cognitoUserConfiguration ( 'AUTH' ) ,
209
+ ...headers ,
210
+ }
211
+ case 'LAMBDA' :
212
+ return {
213
+ 'Authorization' : authToken ,
214
+ ...headers ,
215
+ }
216
+ default :
217
+ throw new Error ( `Unknown auth mode ${ authMode } ` )
218
+ }
219
+ }
220
+
81
221
/**
82
222
* Initiates a subscription to a channel and returns the response
83
223
*
84
224
* @param {string } channel the channel to subscribe to
225
+ * @param {string } authMode the authorization mode for the request
226
+ * @param {string } authToken the token used for Lambda auth mode
85
227
* @param {boolean } triggerPub whether to also publish in the method
86
228
* @returns {Object }
87
229
*/
88
- async function subscribe ( channel , triggerPub = false ) {
230
+ async function subscribe ( channel , authMode , authToken , triggerPub = false ) {
89
231
const response = { } ;
90
- const auth = await getAuthProtocolForIAM ( httpUrl , region )
232
+ const authHeader = await getPublishAuthHeader ( authMode , { } , authToken ) ;
233
+ const auth = getAuthProtocolForIAM ( authHeader ) ;
91
234
const socket = await new Promise ( ( resolve , reject ) => {
92
235
const socket = new WebSocket (
93
- `wss:// ${ realtimeUrl } /event/realtime ` ,
236
+ `${ realtimeUrl } ` ,
94
237
[ AWS_APPSYNC_EVENTS_SUBPROTOCOL , auth ] ,
95
238
{ headers : { ...DEFAULT_HEADERS } } ,
96
239
)
97
240
98
241
socket . onopen = ( ) => {
99
242
socket . send ( JSON . stringify ( { type : 'connection_init' } ) )
100
- console . log ( " Initialize connection" ) ;
243
+ console . log ( ' Initialize connection' ) ;
101
244
resolve ( socket )
102
245
}
103
246
@@ -113,18 +256,18 @@ async function subscribe(channel, triggerPub=false) {
113
256
console . log ( 'Data received' ) ;
114
257
response . pubStatusCode = 200 ;
115
258
response . pubMsg = JSON . parse ( payload . event ) . message ;
116
- } else if ( payload . type === " subscribe_error" ) {
259
+ } else if ( payload . type === ' subscribe_error' ) {
117
260
console . log ( payload ) ;
118
- if ( payload . errors . some ( ( error ) => error . errorType === " UnauthorizedException" ) ) {
119
- console . log ( " Error received" ) ;
261
+ if ( payload . errors . some ( ( error ) => error . errorType === ' UnauthorizedException' ) ) {
262
+ console . log ( ' Error received' ) ;
120
263
response . statusCode = 401 ;
121
- response . msg = " UnauthorizedException" ;
264
+ response . msg = ' UnauthorizedException' ;
122
265
} else if ( payload . errors . some ( error => error . errorType === 'AccessDeniedException' ) ) {
123
266
console . log ( 'Error received' ) ;
124
267
response . statusCode = 403 ;
125
268
response . msg = 'Forbidden' ;
126
269
} else {
127
- console . log ( " Error received" ) ;
270
+ console . log ( ' Error received' ) ;
128
271
response . statusCode = 400 ;
129
272
response . msg = payload . errors [ 0 ] . errorType ;
130
273
}
@@ -138,12 +281,12 @@ async function subscribe(channel, triggerPub=false) {
138
281
type : 'subscribe' ,
139
282
id : crypto . randomUUID ( ) ,
140
283
channel : subChannel ,
141
- authorization : await signWithAWSV4 ( httpUrl , region , JSON . stringify ( { channel : subChannel } ) ) ,
284
+ authorization : await getPublishAuthHeader ( authMode , { channel : subChannel } , authToken ) ,
142
285
} ) ) ;
143
286
144
287
if ( triggerPub ) {
145
288
await sleep ( 1000 ) ;
146
- await publish ( channel ) ;
289
+ await publish ( channel , authMode , authToken ) ;
147
290
}
148
291
await sleep ( 3000 ) ;
149
292
return response ;
@@ -153,19 +296,21 @@ async function subscribe(channel, triggerPub=false) {
153
296
* Publishes to a channel and returns the response
154
297
*
155
298
* @param {string } channel the channel to publish to
299
+ * @param {string } authMode the auth mode to use for publishing
300
+ * @param {string } authToken the auth token to use for Lambda auth mode
156
301
* @returns {Object }
157
302
*/
158
- async function publish ( channel ) {
303
+ async function publish ( channel , authMode , authToken ) {
159
304
const event = {
160
- " channel" : `/${ channel } /test` ,
161
- " events" : [
305
+ ' channel' : `/${ channel } /test` ,
306
+ ' events' : [
162
307
JSON . stringify ( { message :'Hello World!' } )
163
308
]
164
309
}
165
310
166
- const response = await fetch ( `https:// ${ httpUrl } /event ` , {
311
+ const response = await fetch ( `${ httpUrl } ` , {
167
312
method : 'POST' ,
168
- headers : await signWithAWSV4 ( httpUrl , region , JSON . stringify ( event ) ) ,
313
+ headers : await getPublishAuthHeader ( authMode , event , authToken ) ,
169
314
body : JSON . stringify ( event )
170
315
} ) ;
171
316
@@ -190,18 +335,34 @@ async function publish(channel) {
190
335
exports . handler = async function ( event ) {
191
336
const pubSubAction = event . action ;
192
337
const channel = event . channel ;
338
+ const authMode = event . authMode ;
339
+ const authToken = event . authToken ?? '' ;
340
+ const isCustomEndpoint = event . customEndpoint ?? false ;
193
341
342
+ // If custom endpoint, wait for 60 seconds for DNS to propagate
343
+ if ( isCustomEndpoint ) {
344
+ await sleep ( 60000 ) ;
345
+ }
346
+
347
+ if ( authMode === 'USER_POOL' ) {
348
+ await cognitoUserConfiguration ( 'CREATE' ) ;
349
+ }
350
+
351
+ let res ;
194
352
if ( pubSubAction === 'publish' ) {
195
- const res = await publish ( channel ) ;
353
+ res = await publish ( channel , authMode , authToken ) ;
196
354
console . log ( res ) ;
197
- return res ;
198
355
} else if ( pubSubAction === 'subscribe' ) {
199
- const res = await subscribe ( channel , false ) ;
356
+ res = await subscribe ( channel , authMode , authToken , false ) ;
200
357
console . log ( res ) ;
201
- return res ;
202
358
} else if ( pubSubAction === 'pubSub' ) {
203
- const res = await subscribe ( channel , true ) ;
359
+ res = await subscribe ( channel , authMode , authToken , true ) ;
204
360
console . log ( res ) ;
205
- return res ;
206
361
}
207
- } ;
362
+
363
+ if ( authMode === 'USER_POOL' ) {
364
+ await cognitoUserConfiguration ( 'DELETE' ) ;
365
+ }
366
+
367
+ return res ;
368
+ } ;
0 commit comments