Skip to content

Fix four out-of-bounds defects in the Query serializer #1343

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
adwsingh wants to merge 1 commit into
base: main
Choose a base branch
from
Open

Fix four out-of-bounds defects in the Query serializer #1343

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions aws/client/aws-client-awsquery/build.gradle.kts
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,5 +1,6 @@
plugins {
id("smithy-java.module-conventions")
id("smithy-java.jmh-conventions")
id("smithy-java.protocol-testing-conventions")
}

Expand Down
61 changes: 61 additions & 0 deletions ...mh/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializerBenchmark.java
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,61 @@
/*
* Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
* SPDX-License-Identifier: Apache-2.0
*/

package software.amazon.smithy.java.aws.client.awsquery;

import java.nio.ByteBuffer;
import org.openjdk.jmh.annotations.Benchmark;
import org.openjdk.jmh.annotations.Param;
import org.openjdk.jmh.annotations.Scope;
import org.openjdk.jmh.annotations.Setup;
import org.openjdk.jmh.annotations.State;
import software.amazon.smithy.java.core.schema.PreludeSchemas;
import software.amazon.smithy.java.core.schema.Schema;
import software.amazon.smithy.model.shapes.ShapeId;

@State(Scope.Thread)
public class QueryFormSerializerBenchmark {

@Param({
"ascii_128",
"unicode_first_128",
"unicode_last_128",
"cjk_128",
"ascii_8192",
"unicode_last_8192",
})
public String testCaseId;

private Schema member;
private String value;

@Setup
public void setup() {
int separator = testCaseId.lastIndexOf('_');
int length = Integer.parseInt(testCaseId.substring(separator + 1));
value = switch (testCaseId.substring(0, separator)) {
case "ascii" -> "a".repeat(length);
case "unicode_first" -> "日" + "a".repeat(length - 1);
case "unicode_last" -> "a".repeat(length - 1) + "日";
case "cjk" -> "日".repeat(length);
default -> throw new IllegalArgumentException("Unknown test case: " + testCaseId);
};

Schema struct = Schema.structureBuilder(ShapeId.from("smithy.benchmark#Input"))
.putMember("value", PreludeSchemas.STRING)
.build();
member = struct.member("value");
}

@Benchmark
public ByteBuffer serializeString() {
QueryFormSerializer serializer = QueryFormSerializer.acquire(
QueryFormSerializer.QueryVariant.AWS_QUERY,
"Benchmark",
"2020-01-01");
serializer.writeString(member, value);
return serializer.finish();
}
}
81 changes: 77 additions & 4 deletions ...ry/src/main/java/software/amazon/smithy/java/aws/client/awsquery/QueryFormSerializer.java
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -73,6 +73,8 @@ enum QueryVariant {
private static final int DEFAULT_BUF_SIZE = 1024;
private static final int MAX_CACHEABLE_BUF = DEFAULT_BUF_SIZE * 4;

private static final int MAX_BYTES_PER_CHAR = 9;

record AcquireContext(QueryVariant variant, String action, String version) {}

private static final StripedPool<QueryFormSerializer, AcquireContext> POOL =
Expand Down Expand Up @@ -139,7 +141,14 @@ static QueryFormSerializer acquire(QueryVariant variant, String action, String v
}

ByteBuffer finish() {
ByteBuffer result = ByteBuffer.wrap(buf, 0, pos);
ByteBuffer result;
if (buf.length > MAX_CACHEABLE_BUF) {
byte[] resultBuf = buf;
buf = new byte[DEFAULT_BUF_SIZE];
result = ByteBuffer.wrap(resultBuf, 0, pos).slice();
} else {
result = ByteBuffer.wrap(Arrays.copyOf(buf, pos));
}
POOL.release(this);
return result;
}
Expand Down Expand Up @@ -223,7 +232,32 @@ private void writeUrlEncodedAsciiBytes(byte[] data, int dataLen) {

private void writeUrlEncoded(String s) {
int len = s.length();
int next = pos;
for (int i = 0; i < len; i++) {
char c = s.charAt(i);
if (c >= 0x80) {
ensureCapacity(next - pos + (len - i) * MAX_BYTES_PER_CHAR);
// Read buf after ensureCapacity: it may have replaced the array.
pos = writeUrlEncodedRemainder(buf, next, s, i);
return;
}
if (UNRESERVED[c]) {
buf[next++] = (byte) c;
} else {
int off = c * 3;
buf[next] = PERCENT_ENCODED[off];
buf[next + 1] = PERCENT_ENCODED[off + 1];
buf[next + 2] = PERCENT_ENCODED[off + 2];
next += 3;
}
}
pos = next;
}

// Encodes the remaining arbitrary characters. The caller must reserve nine bytes per char.
private static int writeUrlEncodedRemainder(byte[] buf, int pos, String s, int start) {
int len = s.length();
for (int i = start; i < len; i++) {
char c = s.charAt(i);
if (c < 0x80) {
if (UNRESERVED[c]) {
Expand Down Expand Up @@ -269,6 +303,16 @@ private void writeUrlEncoded(String s) {
pos += 3;
}
}
return pos;
}

/**
* Upper bound on the base-10 ASCII bytes {@link NumberCodec#writeBigInteger} writes to the form body.
*/
static int maxBigIntegerLength(BigInteger value) {
// log10(2) is just over 0.301; the larger factor plus constants cover rounding and a sign.
int digits = (int) (value.bitLength() * 0.302) + 2;
return 1 + digits;
}

/**
Expand Down Expand Up @@ -431,9 +475,18 @@ private <T> void writeEc2List(Schema schema, T listState, int size, BiConsumer<T
return;
}

// Save/restore for the same reason writeAwsQueryList does it.
var savedMemberNameBytes = listSerializer.memberNameBytes;
var savedFlattened = listSerializer.flattened;
var savedIndex = listSerializer.index;

listSerializer.reset(null, true);
consumer.accept(listState, listSerializer);

listSerializer.memberNameBytes = savedMemberNameBytes;
listSerializer.flattened = savedFlattened;
listSerializer.index = savedIndex;

if (schema.isMember()) {
popPrefix();
}
Expand Down Expand Up @@ -571,7 +624,7 @@ public void writeDouble(Schema schema, double value) {

@Override
public void writeBigInteger(Schema schema, BigInteger value) {
writeIndexedKeyPrefix(64);
writeIndexedKeyPrefix(maxBigIntegerLength(value));
pos = NumberCodec.writeBigInteger(buf, pos, value);
index++;
}
Expand Down Expand Up @@ -744,8 +797,16 @@ public <T> void writeList(Schema schema, T listState, int size, BiConsumer<T, Sh
memberNameBytes = xmlName != null ? xmlName.getValue().getBytes(StandardCharsets.UTF_8) : MEMBER;
}

var savedMemberNameBytes = listSerializer.memberNameBytes;
var savedFlattened = listSerializer.flattened;
var savedIndex = listSerializer.index;

listSerializer.reset(memberNameBytes, flattened);
consumer.accept(listState, listSerializer);

listSerializer.memberNameBytes = savedMemberNameBytes;
listSerializer.flattened = savedFlattened;
listSerializer.index = savedIndex;
}

@Override
Expand All @@ -762,8 +823,20 @@ public <T> void writeMap(Schema schema, T mapState, int size, BiConsumer<T, MapS
valueXmlName != null ? valueXmlName.getValue().getBytes(StandardCharsets.UTF_8) : VALUE;
byte[] entryNameBytes = flattened ? null : ENTRY;

var savedEntry = mapSerializer.entryNameBytes;
var savedKey = mapSerializer.keyNameBytes;
var savedValue = mapSerializer.valueNameBytes;
var savedFlattened = mapSerializer.flattened;
var savedIndex = mapSerializer.index;

mapSerializer.reset(entryNameBytes, keyNameBytes, valueNameBytes, flattened);
consumer.accept(mapState, mapSerializer);

mapSerializer.entryNameBytes = savedEntry;
mapSerializer.keyNameBytes = savedKey;
mapSerializer.valueNameBytes = savedValue;
mapSerializer.flattened = savedFlattened;
mapSerializer.index = savedIndex;
}

@Override
Expand Down Expand Up @@ -817,7 +890,7 @@ public void writeDouble(Schema schema, double value) {

@Override
public void writeBigInteger(Schema schema, BigInteger value) {
writePrefixEquals(64);
writePrefixEquals(maxBigIntegerLength(value));
pos = NumberCodec.writeBigInteger(buf, pos, value);
}

Expand Down Expand Up @@ -924,7 +997,7 @@ public void writeDouble(Schema schema, double value) {
@Override
public void writeBigInteger(Schema schema, BigInteger value) {
byte[] key = getMemberNameBytes(schema);
writeKeyPrefix(key, 64);
writeKeyPrefix(key, maxBigIntegerLength(value));
pos = NumberCodec.writeBigInteger(buf, pos, value);
}

Expand Down
Loading
Loading