1
+ /*
2
+ * Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
3
+ *
4
+ * Licensed under the Apache License, Version 2.0 (the "License").
5
+ * You may not use this file except in compliance with the License.
6
+ * A copy of the License is located at
7
+ *
8
+ * http://aws.amazon.com/apache2.0
9
+ *
10
+ * or in the "license" file accompanying this file. This file is distributed
11
+ * on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either
12
+ * express or implied. See the License for the specific language governing
13
+ * permissions and limitations under the License.
14
+ */
15
+
16
+ package software .amazon .awssdk .services .dsql ;
17
+
18
+ import java .time .Clock ;
19
+ import java .time .Instant ;
20
+ import software .amazon .awssdk .annotations .Immutable ;
21
+ import software .amazon .awssdk .annotations .SdkInternalApi ;
22
+ import software .amazon .awssdk .annotations .SdkTestInternalApi ;
23
+ import software .amazon .awssdk .auth .credentials .AwsCredentials ;
24
+ import software .amazon .awssdk .auth .credentials .CredentialUtils ;
25
+ import software .amazon .awssdk .auth .signer .Aws4Signer ;
26
+ import software .amazon .awssdk .auth .signer .params .Aws4PresignerParams ;
27
+ import software .amazon .awssdk .awscore .client .config .AwsClientOption ;
28
+ import software .amazon .awssdk .core .client .config .SdkClientConfiguration ;
29
+ import software .amazon .awssdk .http .SdkHttpFullRequest ;
30
+ import software .amazon .awssdk .http .SdkHttpMethod ;
31
+ import software .amazon .awssdk .identity .spi .AwsCredentialsIdentity ;
32
+ import software .amazon .awssdk .identity .spi .IdentityProvider ;
33
+ import software .amazon .awssdk .regions .Region ;
34
+ import software .amazon .awssdk .services .dsql .model .GenerateAuthTokenRequest ;
35
+ import software .amazon .awssdk .utils .CompletableFutureUtils ;
36
+ import software .amazon .awssdk .utils .Logger ;
37
+ import software .amazon .awssdk .utils .StringUtils ;
38
+
39
+ @ Immutable
40
+ @ SdkInternalApi
41
+ public final class DefaultDsqlUtilities implements DsqlUtilities {
42
+ private static final Logger log = Logger .loggerFor (DsqlUtilities .class );
43
+ private final Aws4Signer signer = Aws4Signer .create ();
44
+ private final Region region ;
45
+ private final IdentityProvider <? extends AwsCredentialsIdentity > credentialsProvider ;
46
+ private final Clock clock ;
47
+
48
+ public DefaultDsqlUtilities (DefaultBuilder builder ) {
49
+ this (builder , Clock .systemUTC ());
50
+ }
51
+
52
+ /**
53
+ * For testing purposes only
54
+ */
55
+ @ SdkTestInternalApi
56
+ public DefaultDsqlUtilities (DefaultBuilder builder , Clock clock ) {
57
+ this .credentialsProvider = builder .credentialsProvider ;
58
+ this .region = builder .region ;
59
+ this .clock = clock ;
60
+ }
61
+
62
+ /**
63
+ * Used by DSQL low-level client's utilities() method
64
+ */
65
+ @ SdkInternalApi
66
+ static DsqlUtilities create (SdkClientConfiguration clientConfiguration ) {
67
+ return new DefaultBuilder ().clientConfiguration (clientConfiguration ).build ();
68
+ }
69
+
70
+ @ Override
71
+ public String generateDbConnectAuthToken (GenerateAuthTokenRequest request ) {
72
+ return generateAuthToken (request , false );
73
+ }
74
+
75
+ @ Override
76
+ public String generateDbConnectAdminAuthToken (GenerateAuthTokenRequest request ) {
77
+ return generateAuthToken (request , true );
78
+ }
79
+
80
+ private String generateAuthToken (GenerateAuthTokenRequest request , boolean isAdmin ) {
81
+ String action = isAdmin ? "DbConnectAdmin" : "DbConnect" ;
82
+
83
+ SdkHttpFullRequest httpRequest = SdkHttpFullRequest .builder ()
84
+ .method (SdkHttpMethod .GET )
85
+ .protocol ("https" )
86
+ .host (request .hostname ())
87
+ .encodedPath ("/" )
88
+ .putRawQueryParameter ("Action" , action )
89
+ .build ();
90
+
91
+ Instant expirationTime = Instant .now (clock ).plus (request .expiresIn ());
92
+
93
+ Aws4PresignerParams presignRequest = Aws4PresignerParams .builder ()
94
+ .signingClockOverride (clock )
95
+ .expirationTime (expirationTime )
96
+ .awsCredentials (resolveCredentials (request ))
97
+ .signingName ("dsql" )
98
+ .signingRegion (resolveRegion (request ))
99
+ .build ();
100
+
101
+ SdkHttpFullRequest fullRequest = signer .presign (httpRequest , presignRequest );
102
+ String signedUrl = fullRequest .getUri ().toString ();
103
+
104
+ log .debug (() -> "Generated DSQL authentication token with expiration of " + expirationTime );
105
+ return StringUtils .replacePrefixIgnoreCase (signedUrl , "https://" , "" );
106
+ }
107
+
108
+ private Region resolveRegion (GenerateAuthTokenRequest request ) {
109
+ if (request .region () != null ) {
110
+ return request .region ();
111
+ }
112
+
113
+ if (this .region != null ) {
114
+ return this .region ;
115
+ }
116
+
117
+ throw new IllegalArgumentException ("Region must be provided in GenerateAuthTokenRequest or DsqlUtilities" );
118
+ }
119
+
120
+ // TODO: update this to use AwsCredentialsIdentity when we migrate Signers to accept the new type.
121
+ private AwsCredentials resolveCredentials (GenerateAuthTokenRequest request ) {
122
+ if (request .credentialsIdentityProvider () != null ) {
123
+ return CredentialUtils .toCredentials (
124
+ CompletableFutureUtils .joinLikeSync (request .credentialsIdentityProvider ().resolveIdentity ()));
125
+ }
126
+
127
+ if (this .credentialsProvider != null ) {
128
+ return CredentialUtils .toCredentials (CompletableFutureUtils .joinLikeSync (this .credentialsProvider .resolveIdentity ()));
129
+ }
130
+
131
+ throw new IllegalArgumentException ("CredentialsProvider must be provided in GenerateAuthTokenRequest or DsqlUtilities" );
132
+ }
133
+
134
+ @ SdkInternalApi
135
+ public static final class DefaultBuilder implements DsqlUtilities .Builder {
136
+ private Region region ;
137
+ private IdentityProvider <? extends AwsCredentialsIdentity > credentialsProvider ;
138
+
139
+ Builder clientConfiguration (SdkClientConfiguration clientConfiguration ) {
140
+ this .credentialsProvider = clientConfiguration .option (AwsClientOption .CREDENTIALS_IDENTITY_PROVIDER );
141
+ this .region = clientConfiguration .option (AwsClientOption .AWS_REGION );
142
+
143
+ return this ;
144
+ }
145
+
146
+ @ Override
147
+ public Builder region (Region region ) {
148
+ this .region = region ;
149
+ return this ;
150
+ }
151
+
152
+ @ Override
153
+ public Builder credentialsProvider (IdentityProvider <? extends AwsCredentialsIdentity > credentialsProvider ) {
154
+ this .credentialsProvider = credentialsProvider ;
155
+ return this ;
156
+ }
157
+
158
+ /**
159
+ * Construct a {@link DsqlUtilities} object.
160
+ */
161
+ @ Override
162
+ public DsqlUtilities build () {
163
+ return new DefaultDsqlUtilities (this );
164
+ }
165
+ }
166
+ }
0 commit comments