Skip to content

Commit 48c7528

Browse files
mtdowlingmtdowling
committed
More sigv4 optimizations
1 parent 5307007 commit 48c7528

2 files changed

Lines changed: 114 additions & 25 deletions

File tree

‎aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigV4Signer.java‎

Lines changed: 91 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -7,13 +7,13 @@
77

88
import java.nio.ByteBuffer;
99
import java.nio.charset.StandardCharsets;
10+
import java.security.DigestException;
1011
import java.security.InvalidKeyException;
1112
import java.time.Clock;
1213
import java.time.Instant;
1314
import java.time.LocalDateTime;
1415
import java.time.ZoneOffset;
1516
import java.util.HexFormat;
16-
import java.util.List;
1717
import javax.crypto.spec.SecretKeySpec;
1818
import software.amazon.smithy.java.auth.api.SignResult;
1919
import software.amazon.smithy.java.auth.api.Signer;
@@ -35,18 +35,12 @@
3535
final class SigV4Signer implements Signer<HttpRequest, AwsCredentialsIdentity> {
3636

3737
private static final InternalLogger LOGGER = InternalLogger.getLogger(SigV4Signer.class);
38-
private static final List<String> HEADERS_TO_IGNORE_IN_LOWER_CASE = List.of(
39-
HeaderName.CONNECTION.name(),
40-
HeaderName.CONTENT_LENGTH.name(),
41-
HeaderName.X_AMZN_TRACE_ID.name(),
42-
HeaderName.USER_AGENT.name(),
43-
HeaderName.EXPECT.name());
44-
4538
private static final String ALGORITHM = "AWS4-HMAC-SHA256";
4639
private static final String TERMINATOR = "aws4_request";
4740
private static final String HMAC_SHA_256 = "HmacSHA256";
4841
private static final SigningCache SIGNER_CACHE = new SigningCache(300);
4942
private static final String EMPTY_BODY_HASH = "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855";
43+
private static final byte[] HEX_DIGITS = "0123456789abcdef".getBytes(StandardCharsets.US_ASCII);
5044

5145
private final SigningResources signingResources;
5246

@@ -189,8 +183,7 @@ private String signInPlace(
189183
canonicalLen,
190184
scope,
191185
requestTime,
192-
signingKey,
193-
sb);
186+
signingKey);
194187
var authorizationHeader = getAuthHeader(accessKeyId, scope, signedHeaders, signature, sb);
195188

196189
// Now mutate the actual request. setHeader is a single-key operation per header; no wholesale map replacement.
@@ -218,7 +211,7 @@ private static String buildSignedHeadersString(SigningResources r, StringBuilder
218211

219212
for (int i = 0; i < r.headerCount; i++) {
220213
String name = r.headers[i * 2];
221-
if (name.equals(previous) || HEADERS_TO_IGNORE_IN_LOWER_CASE.contains(name)) {
214+
if (name.equals(previous) || isIgnoredHeader(name)) {
222215
continue;
223216
}
224217
if (!sb.isEmpty()) {
@@ -276,7 +269,7 @@ private static void addCanonicalizedHeaderString(SigningResources r, StringBuild
276269
while (next < r.headerCount && name.equals(r.headers[next * 2])) {
277270
next++;
278271
}
279-
if (HEADERS_TO_IGNORE_IN_LOWER_CASE.contains(name)) {
272+
if (isIgnoredHeader(name)) {
280273
i = next;
281274
continue;
282275
}
@@ -396,6 +389,14 @@ private static void addCanonicalizedQueryString(SmithyUri uri, SigningResources
396389

397390
boolean canonical = isAlreadyCanonical(query);
398391
int len = query.length();
392+
if (canonical) {
393+
int amp = query.indexOf('&');
394+
if (amp == -1) {
395+
appendSingleCanonicalQueryPair(query, builder);
396+
return;
397+
}
398+
}
399+
399400
int start = 0;
400401
while (start <= len) {
401402
int amp = query.indexOf('&', start);
@@ -439,6 +440,18 @@ private static void addCanonicalizedQueryString(SmithyUri uri, SigningResources
439440
}
440441
}
441442

443+
private static void appendSingleCanonicalQueryPair(String query, StringBuilder builder) {
444+
int eq = query.indexOf('=');
445+
if (eq == 0) {
446+
return;
447+
}
448+
449+
builder.append(query);
450+
if (eq == -1) {
451+
builder.append('=');
452+
}
453+
}
454+
442455
// Returns true if every character in the query string is either an RFC 3986 unreserved
443456
// char or part of an already-uppercase {@code %XX} escape
444457
private static boolean isAlreadyCanonical(String s) {
@@ -514,6 +527,13 @@ private static boolean isWhiteSpace(char ch) {
514527
return ch == ' ' || (ch >= '\t' && ch <= '\f');
515528
}
516529

530+
private static boolean isIgnoredHeader(String name) {
531+
return switch (name) {
532+
case "connection", "content-length", "x-amzn-trace-id", "user-agent", "expect" -> true;
533+
default -> false;
534+
};
535+
}
536+
517537
/**
518538
* AWS4 uses a series of derived keys, formed by hashing different pieces of data
519539
*/
@@ -553,20 +573,24 @@ private String computeSignature(
553573
int canonicalRequestLength,
554574
String scope,
555575
String requestTime,
556-
byte[] signingKey,
557-
StringBuilder sb
576+
byte[] signingKey
558577
) {
559-
sb.setLength(0);
560-
sb.append(ALGORITHM)
561-
.append('\n')
562-
.append(requestTime)
563-
.append('\n')
564-
.append(scope)
565-
.append('\n')
566-
.append(HexFormat.of().formatHex(hash(canonicalRequest, 0, canonicalRequestLength)));
567-
var toSign = sb.toString();
568-
var signatureBytes = sign(toSign, signingKey);
569-
return HexFormat.of().formatHex(signatureBytes);
578+
byte[] canonicalRequestHash = hash(canonicalRequest, 0, canonicalRequestLength, signingResources.hashBytes);
579+
int stringToSignLength = ALGORITHM.length() + requestTime.length() + scope.length() + 67;
580+
byte[] stringToSign = signingResources.ensureStringToSignCapacity(stringToSignLength);
581+
582+
int pos = writeAscii(ALGORITHM, stringToSign, 0);
583+
stringToSign[pos++] = '\n';
584+
pos = writeAscii(requestTime, stringToSign, pos);
585+
stringToSign[pos++] = '\n';
586+
pos = writeAscii(scope, stringToSign, pos);
587+
stringToSign[pos++] = '\n';
588+
pos = writeHex(canonicalRequestHash, stringToSign, pos);
589+
590+
byte[] signatureBytes = sign(stringToSign, 0, pos, signingKey, signingResources.signatureBytes);
591+
byte[] signatureHex = signingResources.signatureHexBytes;
592+
writeHex(signatureBytes, signatureHex, 0);
593+
return new String(signatureHex, 0, signatureHex.length, StandardCharsets.US_ASCII);
570594
}
571595

572596
private byte[] sign(String data, byte[] key) {
@@ -580,6 +604,36 @@ private byte[] sign(String data, byte[] key) {
580604
}
581605
}
582606

607+
private byte[] sign(byte[] data, int offset, int length, byte[] key, byte[] output) {
608+
try {
609+
var sha256Mac = signingResources.sha256Mac;
610+
sha256Mac.reset();
611+
sha256Mac.init(new SecretKeySpec(key, HMAC_SHA_256));
612+
sha256Mac.update(data, offset, length);
613+
sha256Mac.doFinal(output, 0);
614+
return output;
615+
} catch (InvalidKeyException e) {
616+
throw new RuntimeException(e);
617+
} catch (javax.crypto.ShortBufferException e) {
618+
throw new RuntimeException(e);
619+
}
620+
}
621+
622+
private static int writeAscii(String value, byte[] dst, int offset) {
623+
for (int i = 0; i < value.length(); i++) {
624+
dst[offset++] = (byte) value.charAt(i);
625+
}
626+
return offset;
627+
}
628+
629+
private static int writeHex(byte[] bytes, byte[] dst, int offset) {
630+
for (byte b : bytes) {
631+
dst[offset++] = HEX_DIGITS[(b >>> 4) & 0x0F];
632+
dst[offset++] = HEX_DIGITS[b & 0x0F];
633+
}
634+
return offset;
635+
}
636+
583637
private byte[] hash(ByteBuffer data) {
584638
var sha256Digest = signingResources.sha256Digest;
585639
sha256Digest.reset();
@@ -593,4 +647,16 @@ private byte[] hash(byte[] data, int offset, int length) {
593647
sha256Digest.update(data, offset, length);
594648
return sha256Digest.digest();
595649
}
650+
651+
private byte[] hash(byte[] data, int offset, int length, byte[] output) {
652+
try {
653+
var sha256Digest = signingResources.sha256Digest;
654+
sha256Digest.reset();
655+
sha256Digest.update(data, offset, length);
656+
sha256Digest.digest(output, 0, output.length);
657+
return output;
658+
} catch (DigestException e) {
659+
throw new RuntimeException(e);
660+
}
661+
}
596662
}

‎aws/aws-sigv4/src/main/java/software/amazon/smithy/java/aws/client/auth/scheme/sigv4/SigningResources.java‎

Lines changed: 23 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -46,6 +46,10 @@ final class SigningResources {
4646
* StringBuilder chars directly as bytes without going through {@code String.getBytes}.
4747
*/
4848
byte[] canonicalRequestBytes;
49+
byte[] stringToSignBytes;
50+
final byte[] hashBytes;
51+
final byte[] signatureBytes;
52+
final byte[] signatureHexBytes;
4953

5054
/**
5155
* Ensure {@link #canonicalRequestBytes} has at least {@code minLength} bytes of capacity,
@@ -59,11 +63,27 @@ byte[] ensureCanonicalRequestCapacity(int minLength) {
5963
return canonicalRequestBytes;
6064
}
6165

66+
/**
67+
* Ensure {@link #stringToSignBytes} has at least {@code minLength} bytes of capacity,
68+
* growing to the next power of two if not.
69+
*/
70+
byte[] ensureStringToSignCapacity(int minLength) {
71+
if (stringToSignBytes.length < minLength) {
72+
int newLen = Integer.highestOneBit(minLength - 1) << 1;
73+
stringToSignBytes = new byte[newLen];
74+
}
75+
return stringToSignBytes;
76+
}
77+
6278
SigningResources() {
6379
this.sb = new StringBuilder(BUFFER_SIZE);
6480
this.headers = new String[INITIAL_HEADER_CAPACITY * 2];
6581
this.queryPairs = new String[INITIAL_HEADER_CAPACITY * 2];
6682
this.canonicalRequestBytes = new byte[BUFFER_SIZE];
83+
this.stringToSignBytes = new byte[BUFFER_SIZE];
84+
this.hashBytes = new byte[32];
85+
this.signatureBytes = new byte[32];
86+
this.signatureHexBytes = new byte[64];
6787

6888
try {
6989
this.sha256Digest = MessageDigest.getInstance("SHA-256");
@@ -93,6 +113,9 @@ void shrink() {
93113
// Reallocate in case the header array grew too large
94114
headers = new String[INITIAL_HEADER_CAPACITY * 2];
95115
}
116+
if (stringToSignBytes.length > BUFFER_SIZE) {
117+
stringToSignBytes = new byte[BUFFER_SIZE];
118+
}
96119
}
97120

98121
private void clearHeaderRefs() {

0 commit comments

Comments
 (0)