From 99483761af3bb1ef2aeec553cadb12ef4a56fd8e Mon Sep 17 00:00:00 2001
From: Jason Lee <jason@steeplesoft.com>
Date: Mon, 25 Nov 2024 10:58:41 -0600
Subject: [PATCH] [UNDERTOW-1881] - Add a new exchange attribute for SSL/TLS
 protocol version

Add and register new ExchangeAttribute implementation
Add support for AJP and TLS
Add test case
---
 .../attribute/SecureProtocolAttribute.java    |  85 ++++++++++++
 .../accesslog/DefaultAccessLogReceiver.java   |   2 +-
 ...ndertow.attribute.ExchangeAttributeBuilder |   1 +
 .../ssl/SecureProtocolAttributeTestCase.java  | 121 ++++++++++++++++++
 4 files changed, 208 insertions(+), 1 deletion(-)
 create mode 100644 core/src/main/java/io/undertow/attribute/SecureProtocolAttribute.java
 create mode 100644 core/src/test/java/io/undertow/server/ssl/SecureProtocolAttributeTestCase.java

diff --git a/core/src/main/java/io/undertow/attribute/SecureProtocolAttribute.java b/core/src/main/java/io/undertow/attribute/SecureProtocolAttribute.java
new file mode 100644
index 0000000000..d002bc8d3f
--- /dev/null
+++ b/core/src/main/java/io/undertow/attribute/SecureProtocolAttribute.java
@@ -0,0 +1,85 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2024 Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+
+package io.undertow.attribute;
+
+import javax.net.ssl.SSLSession;
+
+import io.undertow.server.HttpServerExchange;
+import io.undertow.server.SSLSessionInfo;
+import io.undertow.util.HeaderValues;
+
+public class SecureProtocolAttribute implements ExchangeAttribute {
+
+    public static final SecureProtocolAttribute INSTANCE = new SecureProtocolAttribute();
+
+    @Override
+    public String readAttribute(HttpServerExchange exchange) {
+        String secureProtocol = null;
+        String transportProtocol = exchange.getConnection().getTransportProtocol();
+        if ("ajp".equals(transportProtocol)) {
+            // TODO: wrong
+            HeaderValues headerValues = exchange.getRequestHeaders().get("AJP_SSL_PROTOCOL");
+            if (headerValues != null && !headerValues.isEmpty()) {
+                secureProtocol = headerValues.getFirst();
+            }
+        } else {
+            SSLSessionInfo ssl = exchange.getConnection().getSslSessionInfo();
+            if (ssl == null) {
+                return null;
+            }
+            SSLSession session = ssl.getSSLSession();
+            if (session != null) {
+                secureProtocol = session.getProtocol();
+            }
+        }
+
+        return secureProtocol;
+    }
+
+    @Override
+    public void writeAttribute(HttpServerExchange exchange, String newValue) throws ReadOnlyAttributeException {
+        throw new ReadOnlyAttributeException("Secure Protocol", newValue);
+    }
+
+    @Override
+    public String toString() {
+        return "%{SECURE_PROTOCOL}";
+    }
+
+    public static final class Builder implements ExchangeAttributeBuilder {
+
+        @Override
+        public String name() {
+            return "Secure Protocol";
+        }
+
+        @Override
+        public ExchangeAttribute build(final String token) {
+            if (token.equals("%{SECURE_PROTOCOL}")) {
+                return INSTANCE;
+            }
+            return null;
+        }
+
+        @Override
+        public int priority() {
+            return 0;
+        }
+    }
+}
diff --git a/core/src/main/java/io/undertow/server/handlers/accesslog/DefaultAccessLogReceiver.java b/core/src/main/java/io/undertow/server/handlers/accesslog/DefaultAccessLogReceiver.java
index d6f093c645..d048ae26fb 100644
--- a/core/src/main/java/io/undertow/server/handlers/accesslog/DefaultAccessLogReceiver.java
+++ b/core/src/main/java/io/undertow/server/handlers/accesslog/DefaultAccessLogReceiver.java
@@ -234,7 +234,7 @@ public void run() {
      * <p>
      * DO NOT USE THIS OUTSIDE OF A TEST
      */
-    void awaitWrittenForTest() throws InterruptedException {
+    protected void awaitWrittenForTest() throws InterruptedException {
         while (!pendingMessages.isEmpty() || forceLogRotation) {
             Thread.sleep(10);
         }
diff --git a/core/src/main/resources/META-INF/services/io.undertow.attribute.ExchangeAttributeBuilder b/core/src/main/resources/META-INF/services/io.undertow.attribute.ExchangeAttributeBuilder
index eea0057862..18ec2a2aa7 100644
--- a/core/src/main/resources/META-INF/services/io.undertow.attribute.ExchangeAttributeBuilder
+++ b/core/src/main/resources/META-INF/services/io.undertow.attribute.ExchangeAttributeBuilder
@@ -23,6 +23,7 @@ io.undertow.attribute.PredicateContextAttribute$Builder
 io.undertow.attribute.QueryParameterAttribute$Builder
 io.undertow.attribute.SslClientCertAttribute$Builder
 io.undertow.attribute.SslCipherAttribute$Builder
+io.undertow.attribute.SecureProtocolAttribute$Builder
 io.undertow.attribute.SslSessionIdAttribute$Builder
 io.undertow.attribute.ResponseTimeAttribute$Builder
 io.undertow.attribute.PathParameterAttribute$Builder
diff --git a/core/src/test/java/io/undertow/server/ssl/SecureProtocolAttributeTestCase.java b/core/src/test/java/io/undertow/server/ssl/SecureProtocolAttributeTestCase.java
new file mode 100644
index 0000000000..98dec3f298
--- /dev/null
+++ b/core/src/test/java/io/undertow/server/ssl/SecureProtocolAttributeTestCase.java
@@ -0,0 +1,121 @@
+/*
+ * JBoss, Home of Professional Open Source.
+ * Copyright 2024 Red Hat, Inc., and individual contributors
+ * as indicated by the @author tags.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     http://www.apache.org/licenses/LICENSE-2.0
+ *
+ *  Unless required by applicable law or agreed to in writing, software
+ *  distributed under the License is distributed on an "AS IS" BASIS,
+ *  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ *  See the License for the specific language governing permissions and
+ *  limitations under the License.
+ */
+package io.undertow.server.ssl;
+
+import java.io.BufferedWriter;
+import java.io.IOException;
+import java.nio.charset.StandardCharsets;
+import java.nio.file.Files;
+import java.nio.file.Path;
+import java.nio.file.Paths;
+import java.nio.file.StandardOpenOption;
+import java.util.concurrent.Executor;
+import javax.net.ssl.SSLContext;
+
+import io.undertow.UndertowLogger;
+import io.undertow.attribute.ExchangeAttribute;
+import io.undertow.attribute.ExchangeAttributes;
+import io.undertow.attribute.SubstituteEmptyWrapper;
+import io.undertow.server.handlers.accesslog.AccessLogReceiver;
+import io.undertow.server.handlers.accesslog.DefaultAccessLogReceiver;
+import io.undertow.testutils.DefaultServer;
+import io.undertow.testutils.HttpClientUtils;
+import io.undertow.testutils.TestHttpClient;
+import io.undertow.util.CompletionLatchHandler;
+import io.undertow.util.StatusCodes;
+import org.apache.http.HttpResponse;
+import org.apache.http.client.methods.HttpGet;
+import org.junit.Assert;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+
+@RunWith(DefaultServer.class)
+public class SecureProtocolAttributeTestCase {
+    private static final Path logDirectory = Paths.get(System.getProperty("java.io.tmpdir"), "logs");
+
+    @Test
+    public void testTlsRequestViaLogging() throws IOException, InterruptedException {
+        LocalAccessLogReceiver logReceiver
+                = new LocalAccessLogReceiver(DefaultServer.getWorker(), logDirectory, "server", ".log");
+
+        final String formatString = "Secure Protocol is %{SECURE_PROTOCOL}.";
+        CompletionLatchHandler latchHandler = new CompletionLatchHandler(
+                exchange -> {
+                    ExchangeAttribute tokens = ExchangeAttributes.parser(SecureProtocolAttributeTestCase.class.getClassLoader(),
+                            new SubstituteEmptyWrapper("-")).parse(formatString);
+                    exchange.getResponseSender().send(tokens.readAttribute(exchange));
+                });
+
+        DefaultServer.setRootHandler(latchHandler);
+
+        try (TestHttpClient client = new TestHttpClient()) {
+            DefaultServer.startSSLServer();
+            SSLContext sslContext = DefaultServer.getClientSSLContext();
+            client.setSSLContext(sslContext);
+
+            HttpResponse result = client.execute(new HttpGet(DefaultServer.getDefaultServerSSLAddress() + "/path"));
+            Assert.assertEquals(StatusCodes.OK, result.getStatusLine().getStatusCode());
+
+            String response = HttpClientUtils.readResponse(result);
+            Assert.assertEquals(
+                    formatString.replaceAll("%\\{SECURE_PROTOCOL}",
+                            "false".equals(System.getProperty("test.ajp")) ? sslContext.getProtocol() : "-"),
+                    response);
+        } finally {
+            DefaultServer.stopSSLServer();
+        }
+    }
+
+    private static class SimpleAccessLogReceiver implements AccessLogReceiver {
+        private final BufferedWriter writer;
+
+        SimpleAccessLogReceiver(Path logFile) {
+            try {
+                writer = Files.newBufferedWriter(logFile, StandardCharsets.UTF_8, StandardOpenOption.APPEND,
+                        StandardOpenOption.CREATE);
+            } catch (IOException e) {
+                throw new RuntimeException(e);
+            }
+        }
+
+        @Override
+        public void logMessage(String message) {
+            try {
+                writer.write(message);
+                writer.newLine();
+                writer.flush();
+            } catch (IOException e) {
+                UndertowLogger.ROOT_LOGGER.errorWritingAccessLog(e);
+            }
+        }
+    }
+
+    private static class LocalAccessLogReceiver extends DefaultAccessLogReceiver {
+        LocalAccessLogReceiver(final Executor logWriteExecutor,
+                               final Path outputDirectory,
+                               final String logBaseName,
+                               final String logNameSuffix) {
+            super(logWriteExecutor, outputDirectory, logBaseName, logNameSuffix, true);
+        }
+
+        public void awaitWrittenForTest() throws InterruptedException {
+            super.awaitWrittenForTest();
+        }
+    }
+
+}