Skip to content

Commit c04fd13

Browse files
leekibleekib
committed
feat: Introduce CookieProcessor interface and refactor cookie handling
- Created a `CookieProcessor` interface along with its implementation `AwsCookieProcessor` to encapsulate cookie parsing and formatting logic. - Modified `AwsHttpServletResponse` and `AwsHttpServletRequest` to use the `CookieProcessor` for all cookie-related operations.
1 parent 3a72095 commit c04fd13

File tree

5 files changed

+317
-48
lines changed

5 files changed

+317
-48
lines changed

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsCookieProcessor.java‎

Lines changed: 275 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,275 @@
1+
package com.amazonaws.serverless.proxy.internal.servlet;
2+
3+
import com.amazonaws.serverless.proxy.internal.SecurityUtils;
4+
import jakarta.servlet.http.Cookie;
5+
import org.slf4j.Logger;
6+
import org.slf4j.LoggerFactory;
7+
8+
import java.text.DateFormat;
9+
import java.text.FieldPosition;
10+
import java.text.SimpleDateFormat;
11+
import java.util.*;
12+
13+
/**
14+
* Implementation of the CookieProcessor interface that provides cookie parsing and generation functionality.
15+
*/
16+
public class AwsCookieProcessor implements CookieProcessor {
17+
18+
// Cookie attribute constants
19+
static final String COOKIE_COMMENT_ATTR = "Comment";
20+
static final String COOKIE_DOMAIN_ATTR = "Domain";
21+
static final String COOKIE_MAX_AGE_ATTR = "Max-Age";
22+
static final String COOKIE_PATH_ATTR = "Path";
23+
static final String COOKIE_SECURE_ATTR = "Secure";
24+
static final String COOKIE_HTTP_ONLY_ATTR = "HttpOnly";
25+
static final String COOKIE_SAME_SITE_ATTR = "SameSite";
26+
static final String COOKIE_PARTITIONED_ATTR = "Partitioned";
27+
static final String EMPTY_STRING = "";
28+
29+
// BitSet to store valid token characters as defined in RFC 2616
30+
static final BitSet tokenValid = createTokenValidSet();
31+
32+
// BitSet to validate domain characters
33+
static final BitSet domainValid = createDomainValidSet();
34+
35+
static final String COOKIE_DATE_PATTERN = "EEE, dd MMM yyyy HH:mm:ss z";
36+
37+
// ThreadLocal to ensure thread-safe creation of DateFormat instances for each thread
38+
static final ThreadLocal<DateFormat> COOKIE_DATE_FORMAT = ThreadLocal.withInitial(() -> {
39+
DateFormat df = new SimpleDateFormat(COOKIE_DATE_PATTERN, Locale.US);
40+
df.setTimeZone(TimeZone.getTimeZone("GMT"));
41+
return df;
42+
});
43+
44+
static final String ANCIENT_DATE = COOKIE_DATE_FORMAT.get().format(new Date(10000));
45+
46+
static BitSet createTokenValidSet() {
47+
BitSet tokenSet = new BitSet(128);
48+
for (char c = '0'; c <= '9'; c++) tokenSet.set(c);
49+
for (char c = 'a'; c <= 'z'; c++) tokenSet.set(c);
50+
for (char c = 'A'; c <= 'Z'; c++) tokenSet.set(c);
51+
for (char c : "!#$%&'*+-.^_`|~".toCharArray()) tokenSet.set(c);
52+
return tokenSet;
53+
}
54+
55+
static BitSet createDomainValidSet() {
56+
BitSet domainValid = new BitSet(128);
57+
for (char c = '0'; c <= '9'; c++) domainValid.set(c);
58+
for (char c = 'a'; c <= 'z'; c++) domainValid.set(c);
59+
for (char c = 'A'; c <= 'Z'; c++) domainValid.set(c);
60+
domainValid.set('.');
61+
domainValid.set('-');
62+
return domainValid;
63+
}
64+
65+
private final Logger log = LoggerFactory.getLogger(AwsCookieProcessor.class);
66+
67+
@Override
68+
public Cookie[] parseCookieHeader(String cookieHeader) {
69+
// Return an empty array if the input is null or empty after trimming
70+
if (cookieHeader == null || cookieHeader.trim().isEmpty()) {
71+
return new Cookie[0];
72+
}
73+
74+
// Parse cookie header and convert to Cookie array
75+
return Arrays.stream(cookieHeader.split("\\s*;\\s*"))
76+
.map(this::parseCookiePair)
77+
.filter(Objects::nonNull) // Filter out invalid pairs
78+
.toArray(Cookie[]::new);
79+
}
80+
81+
/**
82+
* Parse a single cookie pair (name=value).
83+
*
84+
* @param cookiePair The cookie pair string.
85+
* @return A valid Cookie object or null if the pair is invalid.
86+
*/
87+
private Cookie parseCookiePair(String cookiePair) {
88+
String[] kv = cookiePair.split("=", 2);
89+
90+
if (kv.length != 2) {
91+
log.warn("Ignoring invalid cookie: {}", cookiePair);
92+
return null; // Skip malformed cookie pairs
93+
}
94+
95+
String cookieName = kv[0];
96+
String cookieValue = kv[1];
97+
98+
// Validate name and value
99+
if (!isToken(cookieName)){
100+
log.warn("Ignoring cookie with invalid name: {}={}", cookieName, cookieValue);
101+
return null; // Skip invalid cookie names
102+
}
103+
104+
if (!isValidCookieValue(cookieValue)) {
105+
log.warn("Ignoring cookie with invalid value: {}={}", cookieName, cookieValue);
106+
return null; // Skip invalid cookie values
107+
}
108+
109+
// Return a new Cookie object after security processing
110+
return new Cookie(SecurityUtils.crlf(cookieName), SecurityUtils.crlf(cookieValue));
111+
}
112+
113+
@Override
114+
public String generateHeader(Cookie cookie) {
115+
StringBuffer header = new StringBuffer();
116+
header.append(cookie.getName()).append('=');
117+
118+
String value = cookie.getValue();
119+
if (value != null && value.length() > 0) {
120+
validateCookieValue(value);
121+
header.append(value);
122+
}
123+
124+
int maxAge = cookie.getMaxAge();
125+
if (maxAge > -1) {
126+
header.append("; Expires=");
127+
if (maxAge == 0) {
128+
header.append(ANCIENT_DATE);
129+
} else {
130+
COOKIE_DATE_FORMAT.get().format(
131+
new Date(System.currentTimeMillis() + maxAge * 1000L), header, new FieldPosition(0));
132+
header.append("; Max-Age=").append(maxAge);
133+
}
134+
}
135+
136+
String domain = cookie.getDomain();
137+
if (domain != null && !domain.isEmpty()) {
138+
validateDomain(domain);
139+
header.append("; Domain=").append(domain);
140+
}
141+
142+
String path = cookie.getPath();
143+
if (path != null && !path.isEmpty()) {
144+
validatePath(path);
145+
header.append("; Path=").append(path);
146+
}
147+
148+
if (cookie.getSecure()) {
149+
header.append("; Secure");
150+
}
151+
152+
if (cookie.isHttpOnly()) {
153+
header.append("; HttpOnly");
154+
}
155+
156+
String sameSite = cookie.getAttribute(COOKIE_SAME_SITE_ATTR);
157+
if (sameSite != null) {
158+
header.append("; SameSite=").append(sameSite);
159+
}
160+
161+
String partitioned = cookie.getAttribute(COOKIE_PARTITIONED_ATTR);
162+
if (EMPTY_STRING.equals(partitioned)) {
163+
header.append("; Partitioned");
164+
}
165+
166+
addAdditionalAttributes(cookie, header);
167+
168+
return header.toString();
169+
}
170+
171+
private void addAdditionalAttributes(Cookie cookie, StringBuffer header) {
172+
for (Map.Entry<String, String> entry : cookie.getAttributes().entrySet()) {
173+
switch (entry.getKey()) {
174+
case COOKIE_COMMENT_ATTR:
175+
case COOKIE_DOMAIN_ATTR:
176+
case COOKIE_MAX_AGE_ATTR:
177+
case COOKIE_PATH_ATTR:
178+
case COOKIE_SECURE_ATTR:
179+
case COOKIE_HTTP_ONLY_ATTR:
180+
case COOKIE_SAME_SITE_ATTR:
181+
case COOKIE_PARTITIONED_ATTR:
182+
// Already handled attributes are ignored
183+
break;
184+
default:
185+
validateAttribute(entry.getKey(), entry.getValue());
186+
header.append("; ").append(entry.getKey());
187+
if (!EMPTY_STRING.equals(entry.getValue())) {
188+
header.append('=').append(entry.getValue());
189+
}
190+
break;
191+
}
192+
}
193+
}
194+
195+
private void validateCookieValue(String value) {
196+
if (!isValidCookieValue(value)) {
197+
throw new IllegalArgumentException("Invalid cookie value: " + value);
198+
}
199+
}
200+
201+
private void validateDomain(String domain) {
202+
if (!isValidDomain(domain)) {
203+
throw new IllegalArgumentException("Invalid cookie domain: " + domain);
204+
}
205+
}
206+
207+
private void validatePath(String path) {
208+
for (char ch : path.toCharArray()) {
209+
if (ch < 0x20 || ch > 0x7E || ch == ';') {
210+
throw new IllegalArgumentException("Invalid cookie path: " + path);
211+
}
212+
}
213+
}
214+
215+
private void validateAttribute(String name, String value) {
216+
if (!isToken(name)) {
217+
throw new IllegalArgumentException("Invalid cookie attribute name: " + name);
218+
}
219+
220+
for (char ch : value.toCharArray()) {
221+
if (ch < 0x20 || ch > 0x7E || ch == ';') {
222+
throw new IllegalArgumentException("Invalid cookie attribute value: " + ch);
223+
}
224+
}
225+
}
226+
227+
private boolean isValidCookieValue(String value) {
228+
int start = 0;
229+
int end = value.length();
230+
boolean quoted = end > 1 && value.charAt(0) == '"' && value.charAt(end - 1) == '"';
231+
232+
char[] chars = value.toCharArray();
233+
for (int i = start; i < end; i++) {
234+
if (quoted && (i == start || i == end - 1)) {
235+
continue;
236+
}
237+
char c = chars[i];
238+
if (!isValidCookieChar(c)) return false;
239+
}
240+
return true;
241+
}
242+
243+
private boolean isValidDomain(String domain) {
244+
if (domain.isEmpty()) {
245+
return false;
246+
}
247+
int prev = -1;
248+
for (char c : domain.toCharArray()) {
249+
if (!domainValid.get(c) || isInvalidLabelStartOrEnd(prev, c)) {
250+
return false;
251+
}
252+
prev = c;
253+
}
254+
return prev != '.' && prev != '-';
255+
}
256+
257+
private boolean isInvalidLabelStartOrEnd(int prev, char current) {
258+
return (prev == '.' || prev == -1) && (current == '.' || current == '-') ||
259+
(prev == '-' && current == '.');
260+
}
261+
262+
private boolean isToken(String s) {
263+
if (s.isEmpty()) return false;
264+
for (char c : s.toCharArray()) {
265+
if (!tokenValid.get(c)) {
266+
return false;
267+
}
268+
}
269+
return true;
270+
}
271+
272+
private boolean isValidCookieChar(char c) {
273+
return !(c < 0x21 || c > 0x7E || c == 0x22 || c == 0x2c || c == 0x3b || c == 0x5c || c == 0x7f);
274+
}
275+
}

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsHttpApiV2ProxyHttpServletRequest.java‎

Lines changed: 2 additions & 15 deletions
Original file line numberDiff line numberDiff line change
@@ -37,7 +37,6 @@
3737
import java.time.ZonedDateTime;
3838
import java.time.format.DateTimeParseException;
3939
import java.util.*;
40-
import java.util.stream.Collectors;
4140
import java.util.stream.Stream;
4241

4342
public class AwsHttpApiV2ProxyHttpServletRequest extends AwsHttpServletRequest {
@@ -81,26 +80,14 @@ public Cookie[] getCookies() {
8180
if (headers == null || !headers.containsKey(HttpHeaders.COOKIE)) {
8281
rhc = new Cookie[0];
8382
} else {
84-
rhc = parseCookieHeaderValue(headers.getFirst(HttpHeaders.COOKIE));
83+
rhc = getCookieProcessor().parseCookieHeader(headers.getFirst(HttpHeaders.COOKIE));
8584
}
8685

8786
Cookie[] rc;
8887
if (request.getCookies() == null) {
8988
rc = new Cookie[0];
9089
} else {
91-
rc = request.getCookies().stream()
92-
.map(c -> {
93-
int i = c.indexOf('=');
94-
if (i == -1) {
95-
return null;
96-
} else {
97-
String k = SecurityUtils.crlf(c.substring(0, i)).trim();
98-
String v = SecurityUtils.crlf(c.substring(i+1));
99-
return new Cookie(k, v);
100-
}
101-
})
102-
.filter(c -> c != null)
103-
.toArray(Cookie[]::new);
90+
rc = getCookieProcessor().parseCookieHeader(String.join("; ", request.getCookies()));
10491
}
10592

10693
return Stream.concat(Arrays.stream(rhc), Arrays.stream(rc)).toArray(Cookie[]::new);

‎aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/servlet/AwsHttpServletRequest.java‎

Lines changed: 9 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -90,6 +90,7 @@ public abstract class AwsHttpServletRequest implements HttpServletRequest {
9090
private String queryString;
9191
private Map<String, List<Part>> multipartFormParameters;
9292
private Map<String, List<String>> urlEncodedFormParameters;
93+
private CookieProcessor cookieProcessor;
9394

9495
protected AwsHttpServletResponse response;
9596
protected AwsLambdaServletContainerHandler containerHandler;
@@ -295,12 +296,7 @@ public void setServletContext(ServletContext context) {
295296
* @return An array of Cookie objects from the header
296297
*/
297298
protected Cookie[] parseCookieHeaderValue(String headerValue) {
298-
List<HeaderValue> parsedHeaders = this.parseHeaderValue(headerValue, ";", ",");
299-
300-
return parsedHeaders.stream()
301-
.filter(e -> e.getKey() != null)
302-
.map(e -> new Cookie(SecurityUtils.crlf(e.getKey()), SecurityUtils.crlf(e.getValue())))
303-
.toArray(Cookie[]::new);
299+
return getCookieProcessor().parseCookieHeader(headerValue);
304300
}
305301

306302

@@ -512,6 +508,13 @@ protected Map<String, List<String>> getFormUrlEncodedParametersMap() {
512508
return urlEncodedFormParameters;
513509
}
514510

511+
protected CookieProcessor getCookieProcessor(){
512+
if (cookieProcessor == null) {
513+
cookieProcessor = new AwsCookieProcessor();
514+
}
515+
return cookieProcessor;
516+
}
517+
515518
@Override
516519
public Collection<Part> getParts()
517520
throws IOException, ServletException {

0 commit comments

Comments
 (0)