Skip to content

Commit ccdb65c

Browse files
committed
Fixed addCookie method in AwsHttpServletResponse to address #51 and added unit tests
1 parent 88adfdf commit ccdb65c

File tree

6 files changed

+189
-13
lines changed

6 files changed

+189
-13
lines changed

aws-serverless-java-container-core/src/main/java/com/amazonaws/serverless/proxy/internal/LambdaContainerHandler.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -76,7 +76,7 @@ protected LambdaContainerHandler(RequestReader<RequestType, ContainerRequestType
7676
// Methods - Abstract
7777
//-------------------------------------------------------------
7878

79-
protected abstract ContainerResponseType getContainerResponse(CountDownLatch latch);
79+
protected abstract ContainerResponseType getContainerResponse(ContainerRequestType request, CountDownLatch latch);
8080

8181

8282
protected abstract void handleRequest(ContainerRequestType containerRequest, ContainerResponseType containerResponse, Context lambdaContext)
@@ -113,8 +113,8 @@ public ResponseType proxy(RequestType request, Context context) {
113113
try {
114114
SecurityContext securityContext = securityContextWriter.writeSecurityContext(request, context);
115115
CountDownLatch latch = new CountDownLatch(1);
116-
ContainerResponseType containerResponse = getContainerResponse(latch);
117116
ContainerRequestType containerRequest = requestReader.readRequest(request, securityContext, context, config);
117+
ContainerResponseType containerResponse = getContainerResponse(containerRequest, latch);
118118

119119
handleRequest(containerRequest, containerResponse, context);
120120

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

Lines changed: 26 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -21,7 +21,6 @@
2121
import java.io.ByteArrayOutputStream;
2222
import java.io.IOException;
2323
import java.io.PrintWriter;
24-
import java.net.URLEncoder;
2524
import java.text.SimpleDateFormat;
2625
import java.util.*;
2726
import java.util.concurrent.CountDownLatch;
@@ -38,8 +37,8 @@ public class AwsHttpServletResponse
3837
// Constants
3938
//-------------------------------------------------------------
4039

41-
private static final String HEADER_DATE_FORMAT = "EEE, d MMM yyyy HH:mm:ss z";
42-
40+
static final String HEADER_DATE_PATTERN = "EEE, d MMM yyyy HH:mm:ss z";
41+
static final String COOKIE_DEFAULT_TIME_ZONE = "GMT";
4342

4443
//-------------------------------------------------------------
4544
// Variables - Private
@@ -51,6 +50,7 @@ public class AwsHttpServletResponse
5150
private String responseBody;
5251
private ByteArrayOutputStream bodyOutputStream = new ByteArrayOutputStream();
5352
private CountDownLatch writersCountDownLatch;
53+
private AwsHttpServletRequest request;
5454
private boolean isCommitted = false;
5555

5656

@@ -63,8 +63,9 @@ public class AwsHttpServletResponse
6363
* function while the response is asynchronously written by the underlying container/application
6464
* @param latch A latch used to inform the <code>ContainerHandler</code> that we are done receiving the response data
6565
*/
66-
public AwsHttpServletResponse(CountDownLatch latch) {
66+
public AwsHttpServletResponse(AwsHttpServletRequest req, CountDownLatch latch) {
6767
writersCountDownLatch = latch;
68+
request = req;
6869
}
6970

7071

@@ -79,6 +80,25 @@ public void addCookie(Cookie cookie) {
7980
if (cookie.getPath() != null) {
8081
cookieData += "; Path=" + cookie.getPath();
8182
}
83+
if (cookie.getSecure()) {
84+
cookieData += "; Secure";
85+
}
86+
if (cookie.isHttpOnly()) {
87+
cookieData += "; HttpOnly";
88+
}
89+
if (cookie.getDomain() != null && !"".equals(cookie.getDomain().trim())) {
90+
cookieData += "; Domain=" + cookie.getDomain();
91+
}
92+
cookieData += "; Max-Age=" + cookie.getMaxAge();
93+
94+
// we always set the timezone to GMT
95+
TimeZone gmtTimeZone = TimeZone.getTimeZone(COOKIE_DEFAULT_TIME_ZONE);
96+
Calendar currentTimestamp = Calendar.getInstance(gmtTimeZone);
97+
currentTimestamp.add(Calendar.SECOND, cookie.getMaxAge());
98+
SimpleDateFormat cookieDateFormatter = new SimpleDateFormat(HEADER_DATE_PATTERN);
99+
cookieDateFormatter.setTimeZone(gmtTimeZone);
100+
cookieData += "; Expires=" + cookieDateFormatter.format(currentTimestamp.getTime());
101+
82102
setHeader(HttpHeaders.SET_COOKIE, cookieData, false);
83103
}
84104

@@ -141,7 +161,7 @@ public void sendRedirect(String s) throws IOException {
141161

142162
@Override
143163
public void setDateHeader(String s, long l) {
144-
SimpleDateFormat sdf = new SimpleDateFormat(HEADER_DATE_FORMAT);
164+
SimpleDateFormat sdf = new SimpleDateFormat(HEADER_DATE_PATTERN);
145165
Date responseDate = new Date();
146166
responseDate.setTime(l);
147167
setHeader(s, sdf.format(responseDate), true);
@@ -150,7 +170,7 @@ public void setDateHeader(String s, long l) {
150170

151171
@Override
152172
public void addDateHeader(String s, long l) {
153-
SimpleDateFormat sdf = new SimpleDateFormat(HEADER_DATE_FORMAT);
173+
SimpleDateFormat sdf = new SimpleDateFormat(HEADER_DATE_PATTERN);
154174
Date responseDate = new Date();
155175
responseDate.setTime(l);
156176
setHeader(s, sdf.format(responseDate), false);
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
package com.amazonaws.serverless.proxy.internal.servlet;
2+
3+
4+
import org.junit.Test;
5+
6+
import javax.servlet.http.Cookie;
7+
import javax.ws.rs.core.HttpHeaders;
8+
9+
import java.text.ParseException;
10+
import java.text.SimpleDateFormat;
11+
import java.util.Calendar;
12+
import java.util.TimeZone;
13+
import java.util.regex.Matcher;
14+
import java.util.regex.Pattern;
15+
16+
import static org.junit.Assert.*;
17+
18+
19+
public class AwsHttpServletResponseTest {
20+
private static final String COOKIE_NAME = "session_id";
21+
private static final String COOKIE_VALUE = "123";
22+
private static final String COOKIE_PATH = "/api";
23+
private static final String COOKIE_DOMAIN = "mydomain.com";
24+
private static final int MAX_AGE_VALUE = 300;
25+
26+
private static final Pattern MAX_AGE_PATTERN = Pattern.compile("Max-Age=(-?[0-9]+)");
27+
private static final Pattern EXPIRES_PATTERN = Pattern.compile("Expires=(.*)$");
28+
29+
@Test
30+
public void cookie_addCookie_verifyPath() {
31+
AwsHttpServletResponse resp = new AwsHttpServletResponse(null, null);
32+
Cookie pathCookie = new Cookie(COOKIE_NAME, COOKIE_VALUE);
33+
pathCookie.setPath(COOKIE_PATH);
34+
35+
resp.addCookie(pathCookie);
36+
String cookieHeader = resp.getHeader(HttpHeaders.SET_COOKIE);
37+
System.out.println("Cookie string: " + cookieHeader);
38+
assertNotNull(cookieHeader);
39+
assertTrue(cookieHeader.contains("Path=" + COOKIE_PATH));
40+
assertTrue(cookieHeader.contains(COOKIE_NAME + "=" + COOKIE_VALUE));
41+
}
42+
43+
@Test
44+
public void cookie_addCookie_verifySecure() {
45+
AwsHttpServletResponse resp = new AwsHttpServletResponse(null, null);
46+
Cookie secureCookie = new Cookie(COOKIE_NAME, COOKIE_VALUE);
47+
secureCookie.setSecure(true);
48+
49+
resp.addCookie(secureCookie);
50+
String cookieHeader = resp.getHeader(HttpHeaders.SET_COOKIE);
51+
System.out.println("Cookie string: " + cookieHeader);
52+
assertNotNull(cookieHeader);
53+
assertTrue(cookieHeader.contains("; Secure"));
54+
assertTrue(cookieHeader.contains(COOKIE_NAME + "=" + COOKIE_VALUE));
55+
}
56+
57+
@Test
58+
public void cookie_addCookie_verifyDomain() {
59+
AwsHttpServletResponse resp = new AwsHttpServletResponse(null, null);
60+
Cookie domainCookie = new Cookie(COOKIE_NAME, COOKIE_VALUE);
61+
domainCookie.setDomain(COOKIE_DOMAIN);
62+
63+
resp.addCookie(domainCookie);
64+
String cookieHeader = resp.getHeader(HttpHeaders.SET_COOKIE);
65+
System.out.println("Cookie string: " + cookieHeader);
66+
assertNotNull(cookieHeader);
67+
assertTrue(cookieHeader.contains("; Domain=" + COOKIE_DOMAIN));
68+
assertTrue(cookieHeader.contains(COOKIE_NAME + "=" + COOKIE_VALUE));
69+
}
70+
71+
@Test
72+
public void cookie_addCookie_defaultMaxAgeIsNegative() {
73+
AwsHttpServletResponse resp = new AwsHttpServletResponse(null, null);
74+
Cookie maxAgeCookie = new Cookie(COOKIE_NAME, COOKIE_VALUE);
75+
maxAgeCookie.setDomain(COOKIE_DOMAIN);
76+
77+
resp.addCookie(maxAgeCookie);
78+
String cookieHeader = resp.getHeader(HttpHeaders.SET_COOKIE);
79+
System.out.println("Cookie string: " + cookieHeader);
80+
assertNotNull(cookieHeader);
81+
assertTrue(cookieHeader.contains("; Max-Age="));
82+
assertTrue(cookieHeader.contains(COOKIE_NAME + "=" + COOKIE_VALUE));
83+
84+
int maxAge = getMaxAge(cookieHeader);
85+
assertEquals(-1, maxAge);
86+
}
87+
88+
@Test
89+
public void cookie_addCookie_positiveMaxAgeIsPresent() {
90+
AwsHttpServletResponse resp = new AwsHttpServletResponse(null, null);
91+
Cookie maxAgeCookie = new Cookie(COOKIE_NAME, COOKIE_VALUE);
92+
maxAgeCookie.setMaxAge(MAX_AGE_VALUE);
93+
94+
resp.addCookie(maxAgeCookie);
95+
String cookieHeader = resp.getHeader(HttpHeaders.SET_COOKIE);
96+
System.out.println("Cookie string: " + cookieHeader);
97+
assertNotNull(cookieHeader);
98+
assertTrue(cookieHeader.contains("; Max-Age="));
99+
assertTrue(cookieHeader.contains(COOKIE_NAME + "=" + COOKIE_VALUE));
100+
101+
int maxAge = getMaxAge(cookieHeader);
102+
assertEquals(MAX_AGE_VALUE, maxAge);
103+
}
104+
105+
@Test
106+
public void cookie_addCookie_positiveMaxAgeExpiresDate() {
107+
AwsHttpServletResponse resp = new AwsHttpServletResponse(null, null);
108+
Cookie maxAgeCookie = new Cookie(COOKIE_NAME, COOKIE_VALUE);
109+
maxAgeCookie.setMaxAge(MAX_AGE_VALUE);
110+
111+
resp.addCookie(maxAgeCookie);
112+
Calendar testExpiration = Calendar.getInstance();
113+
testExpiration.add(Calendar.SECOND, MAX_AGE_VALUE);
114+
testExpiration.setTimeZone(TimeZone.getTimeZone(AwsHttpServletResponse.COOKIE_DEFAULT_TIME_ZONE));
115+
116+
String cookieHeader = resp.getHeader(HttpHeaders.SET_COOKIE);
117+
System.out.println("Cookie string: " + cookieHeader);
118+
assertNotNull(cookieHeader);
119+
assertTrue(cookieHeader.contains("; Max-Age="));
120+
assertTrue(cookieHeader.contains(COOKIE_NAME + "=" + COOKIE_VALUE));
121+
122+
SimpleDateFormat dateFormat = new SimpleDateFormat(AwsHttpServletResponse.HEADER_DATE_PATTERN);
123+
124+
Calendar expiration = getExpires(cookieHeader);
125+
System.out.println("Cookie date: " + dateFormat.format(expiration.getTime()));
126+
System.out.println("Test date: " + dateFormat.format(testExpiration.getTime()));
127+
// we need to compare strings because the millis time will be off
128+
assertEquals(dateFormat.format(testExpiration.getTime()), dateFormat.format(expiration.getTime()));
129+
}
130+
131+
private int getMaxAge(String header) {
132+
Matcher ageMatcher = MAX_AGE_PATTERN.matcher(header);
133+
assertTrue(ageMatcher.find());
134+
assertTrue(ageMatcher.groupCount() >= 1);
135+
String ageString = ageMatcher.group(1);
136+
System.out.println("Age string: " + ageString);
137+
return Integer.parseInt(ageString);
138+
}
139+
140+
private Calendar getExpires(String header) {
141+
Matcher ageMatcher = EXPIRES_PATTERN.matcher(header);
142+
assertTrue(ageMatcher.find());
143+
assertTrue(ageMatcher.groupCount() >= 1);
144+
String expiresString = ageMatcher.group(1);
145+
SimpleDateFormat sdf = new SimpleDateFormat(AwsHttpServletResponse.HEADER_DATE_PATTERN);
146+
Calendar cal = Calendar.getInstance();
147+
try {
148+
cal.setTime(sdf.parse(expiresString));
149+
} catch (ParseException e) {
150+
e.printStackTrace();
151+
fail("Could not parse expire date");
152+
}
153+
154+
return cal;
155+
}
156+
}

aws-serverless-java-container-jersey/src/main/java/com/amazonaws/serverless/proxy/jersey/JerseyLambdaContainerHandler.java

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -180,7 +180,7 @@ public void reload(ResourceConfig resourceConfig) {
180180
//-------------------------------------------------------------
181181

182182
@Override
183-
protected JerseyResponseWriter getContainerResponse(CountDownLatch latch) {
183+
protected JerseyResponseWriter getContainerResponse(ContainerRequest request, CountDownLatch latch) {
184184
return new JerseyResponseWriter(latch);
185185
}
186186

aws-serverless-java-container-spark/src/main/java/com/amazonaws/serverless/proxy/spark/SparkLambdaContainerHandler.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -137,8 +137,8 @@ public SparkLambdaContainerHandler(RequestReader<RequestType, AwsProxyHttpServle
137137
//-------------------------------------------------------------
138138

139139
@Override
140-
protected AwsHttpServletResponse getContainerResponse(CountDownLatch latch) {
141-
return new AwsHttpServletResponse(latch);
140+
protected AwsHttpServletResponse getContainerResponse(AwsProxyHttpServletRequest request, CountDownLatch latch) {
141+
return new AwsHttpServletResponse(request, latch);
142142
}
143143

144144

aws-serverless-java-container-spring/src/main/java/com/amazonaws/serverless/proxy/spring/SpringLambdaContainerHandler.java

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -99,8 +99,8 @@ public void setRefreshContext(boolean refreshContext) {
9999
}
100100

101101
@Override
102-
protected AwsHttpServletResponse getContainerResponse(CountDownLatch latch) {
103-
return new AwsHttpServletResponse(latch);
102+
protected AwsHttpServletResponse getContainerResponse(AwsProxyHttpServletRequest request, CountDownLatch latch) {
103+
return new AwsHttpServletResponse(request, latch);
104104
}
105105

106106
public void activateSpringProfiles(String... profiles) throws ContainerInitializationException {

0 commit comments

Comments
 (0)