/*
 * JBoss, Home of Professional Open Source.
 * Copyright 2014 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.handlers;

import io.undertow.UndertowLogger;
import io.undertow.UndertowMessages;
import io.undertow.security.api.SecurityContext;
import io.undertow.server.ExchangeCompletionListener;
import io.undertow.server.HandlerWrapper;
import io.undertow.server.HttpHandler;
import io.undertow.server.HttpServerExchange;
import io.undertow.server.handlers.builder.HandlerBuilder;
import io.undertow.util.Headers;

import java.net.InetSocketAddress;
import java.sql.Connection;
import java.sql.PreparedStatement;
import java.sql.SQLException;
import java.sql.Timestamp;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Deque;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentLinkedDeque;
import java.util.concurrent.Executor;
import java.util.concurrent.atomic.AtomicIntegerFieldUpdater;
import javax.naming.InitialContext;
import javax.naming.NamingException;
import javax.sql.DataSource;

public class JDBCLogHandler implements HttpHandler, Runnable {

    private final HttpHandler next;
    private final String formatString;
    private final ExchangeCompletionListener exchangeCompletionListener = new JDBCLogCompletionListener();

    private final Deque<JDBCLogAttribute> pendingMessages;

    //0 = not running
    //1 = queued
    //2 = running
    @SuppressWarnings("unused")
    private volatile int state = 0;
    @SuppressWarnings("unused")
    private volatile Executor executor;

    private static final AtomicIntegerFieldUpdater<JDBCLogHandler> stateUpdater = AtomicIntegerFieldUpdater.newUpdater(JDBCLogHandler.class, "state");

    protected boolean useLongContentLength = false;

    private final DataSource dataSource;

    private String tableName;
    private String remoteHostField;
    private String userField;
    private String timestampField;
    private String virtualHostField;
    private String methodField;
    private String queryField;
    private String statusField;
    private String bytesField;
    private String refererField;
    private String userAgentField;

    @Deprecated
    public JDBCLogHandler(final HttpHandler next, final Executor logWriteExecutor, final String formatString, DataSource dataSource) {
        this(next, formatString, dataSource);
    }

    public JDBCLogHandler(final HttpHandler next, final String formatString, DataSource dataSource) {
        this.next = next;
        this.formatString = formatString;
        this.dataSource = dataSource;

        tableName = "access";
        remoteHostField = "remoteHost";
        userField = "userName";
        timestampField = "timestamp";
        virtualHostField = "virtualHost";
        methodField = "method";
        queryField = "query";
        statusField = "status";
        bytesField = "bytes";
        refererField = "referer";
        userAgentField = "userAgent";
        this.pendingMessages = new ConcurrentLinkedDeque<>();
    }

    @Override
    public void handleRequest(HttpServerExchange exchange) throws Exception {
        exchange.addExchangeCompleteListener(exchangeCompletionListener);
        next.handleRequest(exchange);
    }

    private class JDBCLogCompletionListener implements ExchangeCompletionListener {

        @Override
        public void exchangeEvent(final HttpServerExchange exchange, final NextListener nextListener) {
            try {
                logMessage(formatString, exchange);
            } finally {
                nextListener.proceed();
            }
        }
    }

    public void logMessage(String pattern, HttpServerExchange exchange) {
        JDBCLogAttribute jdbcLogAttribute = new JDBCLogAttribute();

        if (pattern.equals("combined")) {
            jdbcLogAttribute.pattern = pattern;
        }
        jdbcLogAttribute.remoteHost = ((InetSocketAddress) exchange.getConnection().getPeerAddress()).getAddress().getHostAddress();
        SecurityContext sc = exchange.getSecurityContext();
        if (sc == null || !sc.isAuthenticated()) {
            jdbcLogAttribute.user = null;
        } else {
            jdbcLogAttribute.user = sc.getAuthenticatedAccount().getPrincipal().getName();
        }
        jdbcLogAttribute.query = exchange.getQueryString();

        jdbcLogAttribute.bytes = exchange.getResponseContentLength();
        if (jdbcLogAttribute.bytes < 0) {
            jdbcLogAttribute.bytes = 0;
        }

        jdbcLogAttribute.status = exchange.getStatusCode();

        if (jdbcLogAttribute.pattern.equals("combined")) {
            jdbcLogAttribute.virtualHost = exchange.getRequestHeaders().getFirst(Headers.HOST);
            jdbcLogAttribute.method = exchange.getRequestMethod().toString();
            jdbcLogAttribute.referer = exchange.getRequestHeaders().getFirst(Headers.REFERER);
            jdbcLogAttribute.userAgent = exchange.getRequestHeaders().getFirst(Headers.USER_AGENT);
        }

        this.pendingMessages.add(jdbcLogAttribute);
        int state = stateUpdater.get(this);
        if (state == 0) {
            if (stateUpdater.compareAndSet(this, 0, 1)) {
                this.executor = exchange.getConnection().getWorker();
                this.executor.execute(this);
            }
        }
    }

    
insert the log record to database
/** * insert the log record to database */
@Override public void run() { if (!stateUpdater.compareAndSet(this, 1, 2)) { return; } List<JDBCLogAttribute> messages = new ArrayList<>(); JDBCLogAttribute msg = null; //only grab at most 1000 messages at a time for (int i = 0; i < 1000; ++i) { msg = pendingMessages.poll(); if (msg == null) { break; } messages.add(msg); } try { if (!messages.isEmpty()) { writeMessage(messages); } } finally { Executor executor = this.executor; stateUpdater.set(this, 0); //check to see if there is still more messages //if so then run this again if (!pendingMessages.isEmpty()) { if (stateUpdater.compareAndSet(this, 0, 1)) { executor.execute(this); } } } } private void writeMessage(List<JDBCLogAttribute> messages) { PreparedStatement ps = null; Connection conn = null; try { conn = dataSource.getConnection(); conn.setAutoCommit(true); ps = prepareStatement(conn); for (JDBCLogAttribute jdbcLogAttribute : messages) { int numberOfTries = 2; while (numberOfTries > 0) { try { ps.clearParameters(); ps.setString(1, jdbcLogAttribute.remoteHost); ps.setString(2, jdbcLogAttribute.user); ps.setTimestamp(3, jdbcLogAttribute.timestamp); ps.setString(4, jdbcLogAttribute.query); ps.setInt(5, jdbcLogAttribute.status); if (useLongContentLength) { ps.setLong(6, jdbcLogAttribute.bytes); } else { if (jdbcLogAttribute.bytes > Integer.MAX_VALUE) { jdbcLogAttribute.bytes = -1; } ps.setInt(6, (int) jdbcLogAttribute.bytes); } ps.setString(7, jdbcLogAttribute.virtualHost); ps.setString(8, jdbcLogAttribute.method); ps.setString(9, jdbcLogAttribute.referer); ps.setString(10, jdbcLogAttribute.userAgent); ps.executeUpdate(); numberOfTries = 0; } catch (SQLException e) { UndertowLogger.ROOT_LOGGER.failedToWriteJdbcAccessLog(e); } numberOfTries--; } } ps.close(); } catch (SQLException e) { UndertowLogger.ROOT_LOGGER.errorWritingJDBCLog(e); } finally { if (ps != null) { try { ps.close(); } catch (SQLException e) { UndertowLogger.ROOT_LOGGER.debug("Exception closing prepared statement", e); } } if (conn != null) { try { conn.close(); } catch (SQLException e) { UndertowLogger.ROOT_LOGGER.debug("Exception closing connection", e); } } } }
For tests only. Blocks the current thread until all messages are written Just does a busy wait.

DO NOT USE THIS OUTSIDE OF A TEST
/** * For tests only. Blocks the current thread until all messages are written Just does a busy wait. * <p/> * DO NOT USE THIS OUTSIDE OF A TEST */
void awaitWrittenForTest() throws InterruptedException { while (!pendingMessages.isEmpty()) { Thread.sleep(10); } while (state != 0) { Thread.sleep(10); } } private PreparedStatement prepareStatement(Connection conn) throws SQLException { return conn.prepareStatement("INSERT INTO " + tableName + " (" + remoteHostField + ", " + userField + ", " + timestampField + ", " + queryField + ", " + statusField + ", " + bytesField + ", " + virtualHostField + ", " + methodField + ", " + refererField + ", " + userAgentField + ") VALUES(?, ?, ?, ?, ?, ?, ?, ?, ?, ?)"); } private static class JDBCLogAttribute { protected String remoteHost = ""; protected String user = ""; protected String query = ""; protected long bytes = 0; protected int status = 0; protected String virtualHost = ""; protected String method = ""; protected String referer = ""; protected String userAgent = ""; protected String pattern = "common"; protected Timestamp timestamp = new Timestamp(System.currentTimeMillis()); } public boolean isUseLongContentLength() { return useLongContentLength; } public void setUseLongContentLength(boolean useLongContentLength) { this.useLongContentLength = useLongContentLength; } public String getTableName() { return tableName; } public void setTableName(String tableName) { this.tableName = tableName; } public String getRemoteHostField() { return remoteHostField; } public void setRemoteHostField(String remoteHostField) { this.remoteHostField = remoteHostField; } public String getUserField() { return userField; } public void setUserField(String userField) { this.userField = userField; } public String getTimestampField() { return timestampField; } public void setTimestampField(String timestampField) { this.timestampField = timestampField; } public String getVirtualHostField() { return virtualHostField; } public void setVirtualHostField(String virtualHostField) { this.virtualHostField = virtualHostField; } public String getMethodField() { return methodField; } public void setMethodField(String methodField) { this.methodField = methodField; } public String getQueryField() { return queryField; } public void setQueryField(String queryField) { this.queryField = queryField; } public String getStatusField() { return statusField; } public void setStatusField(String statusField) { this.statusField = statusField; } public String getBytesField() { return bytesField; } public void setBytesField(String bytesField) { this.bytesField = bytesField; } public String getRefererField() { return refererField; } public void setRefererField(String refererField) { this.refererField = refererField; } public String getUserAgentField() { return userAgentField; } public void setUserAgentField(String userAgentField) { this.userAgentField = userAgentField; } @Override public String toString() { return "JDBCLogHandler{" + "formatString='" + formatString + '\'' + '}'; } public static class Builder implements HandlerBuilder { @Override public String name() { return "jdbc-access-log"; } @Override public Map<String, Class<?>> parameters() { Map<String, Class<?>> params = new HashMap<>(); params.put("format", String.class); params.put("datasource", String.class); params.put("tableName", String.class); params.put("remoteHostField", String.class); params.put("userField", String.class); params.put("timestampField", String.class); params.put("virtualHostField", String.class); params.put("methodField", String.class); params.put("queryField", String.class); params.put("statusField", String.class); params.put("bytesField", String.class); params.put("refererField", String.class); params.put("userAgentField", String.class); return params; } @Override public Set<String> requiredParameters() { return Collections.singleton("datasource"); } @Override public String defaultParameter() { return "datasource"; } @Override public HandlerWrapper build(Map<String, Object> config) { String datasourceName = (String) config.get("datasource"); try { DataSource ds = (DataSource) new InitialContext().lookup((String) config.get("datasource")); String format = (String) config.get("format"); return new Wrapper(format, ds, (String)config.get("tableName"), (String)config.get("remoteHostField"), (String)config.get("userField"), (String)config.get("timestampField"), (String)config.get("virtualHostField"), (String)config.get("methodField"), (String)config.get("queryField"), (String)config.get("statusField"), (String)config.get("bytesField"), (String)config.get("refererField"), (String)config.get("userAgentField")); } catch (NamingException ex) { throw UndertowMessages.MESSAGES.datasourceNotFound(datasourceName); } } } private static class Wrapper implements HandlerWrapper { private final DataSource datasource; private final String format; private final String tableName; private final String remoteHostField; private final String userField; private final String timestampField; private final String virtualHostField; private final String methodField; private final String queryField; private final String statusField; private final String bytesField; private final String refererField; private final String userAgentField; private Wrapper(String format, DataSource datasource, String tableName, String remoteHostField, String userField, String timestampField, String virtualHostField, String methodField, String queryField, String statusField, String bytesField, String refererField, String userAgentField) { this.datasource = datasource; this.tableName = tableName; this.remoteHostField = remoteHostField; this.userField = userField; this.timestampField = timestampField; this.virtualHostField = virtualHostField; this.methodField = methodField; this.queryField = queryField; this.statusField = statusField; this.bytesField = bytesField; this.refererField = refererField; this.userAgentField = userAgentField; this.format = "combined".equals(format) ? "combined" : "common"; } @Override public HttpHandler wrap(HttpHandler handler) { JDBCLogHandler jdbc = new JDBCLogHandler(handler, format, datasource); if(tableName != null) { jdbc.setTableName(tableName); } if(remoteHostField != null) { jdbc.setRemoteHostField(remoteHostField); } if(userField != null) { jdbc.setUserField(userField); } if(timestampField != null) { jdbc.setTimestampField(timestampField); } if(virtualHostField != null) { jdbc.setVirtualHostField(virtualHostField); } if(methodField != null) { jdbc.setMethodField(methodField); } if(queryField != null) { jdbc.setQueryField(queryField); } if(statusField != null) { jdbc.setStatusField(statusField); } if(bytesField != null) { jdbc.setBytesField(bytesField); } if(refererField != null) { jdbc.setRefererField(refererField); } if(userAgentField != null) { jdbc.setUserAgentField(userAgentField); } return jdbc; } } }