[UNDERTOW-2325] change SecureCookieCommitListener to also look throug…

core/src/main/java/io/undertow/server/SecureCookieCommitListener.java

@@ -1,6 +1,9 @@
 package io.undertow.server;
 
 import io.undertow.server.handlers.Cookie;
+import io.undertow.server.handlers.CookieImpl;
+import io.undertow.util.HeaderValues;
+import io.undertow.util.Headers;
 
 /**
  * Sets the <pre>secure</pre> attribute on all response cookies.
@@ -11,8 +14,28 @@ public enum SecureCookieCommitListener implements ResponseCommitListener {
 
     @Override
     public void beforeCommit(HttpServerExchange exchange) {
+        handleResponseCookies(exchange);
+        handleCookiesSetViaHeaders(exchange);
+    }
+
+    private void handleResponseCookies(HttpServerExchange exchange) {
         for (Cookie cookie : exchange.responseCookies()) {
             cookie.setSecure(true);
         }
     }
+
+    private void handleCookiesSetViaHeaders(HttpServerExchange exchange) {
+        HeaderValues cookieHeaders = exchange.getResponseHeaders().get(Headers.SET_COOKIE);
+        if (cookieHeaders != null) {
+            for (String cookieHeader : cookieHeaders) {
+                String[] parts = cookieHeader.split("=", 2);
+                String cookieName = parts[0];
+                String cookieValue = (parts.length > 1) ? parts[1] : null;
+                CookieImpl cookie = new CookieImpl(cookieName, cookieValue);
+                cookie.setSecure(true);
+                exchange.setResponseCookie(cookie);
+            }
+            exchange.getResponseHeaders().remove(Headers.SET_COOKIE);
+        }
+    }
 }

core/src/test/java/io/undertow/server/handlers/SecureCookieHandlerTestCase.java

@@ -21,6 +21,7 @@
 import java.io.IOException;
 import java.security.GeneralSecurityException;
 
+import io.undertow.util.Headers;
 import org.apache.http.Header;
 import org.apache.http.HttpResponse;
 import org.apache.http.client.methods.HttpGet;
@@ -74,4 +75,31 @@ public void handleRequest(final HttpServerExchange exchange) throws Exception {
         }
     }
 
+    @Test
+    public void testSecureCookieHandlerWithManuallySetCookie() throws IOException, GeneralSecurityException {
+
+        DefaultServer.setRootHandler(new SecureCookieHandler(new HttpHandler() {
+            @Override
+            public void handleRequest(final HttpServerExchange exchange) throws Exception {
+                exchange.getResponseHeaders().put(Headers.SET_COOKIE, "cookie=value");
+            }
+        }));
+
+        DefaultServer.startSSLServer();
+        TestHttpClient client = new TestHttpClient();
+        client.setSSLContext(DefaultServer.getClientSSLContext());
+        try {
+            HttpGet get = new HttpGet(DefaultServer.getDefaultServerSSLAddress());
+            HttpResponse result = client.execute(get);
+            Assert.assertEquals(StatusCodes.OK, result.getStatusLine().getStatusCode());
+
+            Header header = result.getFirstHeader("set-cookie");
+            Assert.assertEquals("cookie=value; secure", header.getValue());
+            FileUtils.readFile(result.getEntity().getContent());
+        } finally {
+            client.getConnectionManager().shutdown();
+            DefaultServer.stopSSLServer();
+        }
+    }
+
 }