@@ -0,0 +1,3711 @@
+/*%0A * JBoss, Home of Professional Open Source.%0A * Copyright 2014 Red Hat, Inc., and individual contributors%0A * as indicated by the @author tags.%0A *%0A * Licensed under the Apache License, Version 2.0 (the %22License%22);%0A * you may not use this file except in compliance with the License.%0A * You may obtain a copy of the License at%0A *%0A *     http://www.apache.org/licenses/LICENSE-2.0%0A *%0A * Unless required by applicable law or agreed to in writing, software%0A * distributed under the License is distributed on an %22AS IS%22 BASIS,%0A * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.%0A * See the License for the specific language governing permissions and%0A * limitations under the License.%0A */%0A%0Apackage org.wildfly.security.auth.callback;%0A%0Aimport java.io.IOException;%0Aimport java.net.SocketAddress;%0Aimport java.util.concurrent.atomic.AtomicBoolean;%0A%0Aimport javax.security.auth.callback.Callback;%0Aimport javax.security.auth.callback.CallbackHandler;%0Aimport javax.security.auth.callback.UnsupportedCallbackException;%0A%0A/**%0A * A callback handler which delegates to another callback handler, passing the local and/or peer socket address to that%0A * callback handler on its first invocation.%0A *%0A * @author %3Ca href=%22mailto:david.lloyd@redhat.com%22%3EDavid M. Lloyd%3C/a%3E%0A */%0Apublic final class SocketAddressQueryCallbackHandler implements CallbackHandler %7B%0A    private final CallbackHandler delegate;%0A    private final SocketAddress localAddress;%0A    private final SocketAddress peerAddress;%0A%0A    private final AtomicBoolean once = new AtomicBoolean();%0A%0A    /**%0A     * Construct a new instance.%0A     *%0A     * @param delegate the callback handler to delegate to%0A     * @param localAddress the local socket address%0A     * @param peerAddress the peer socket address%0A     */%0A    public SocketAddressQueryCallbackHandler(final CallbackHandler delegate, final SocketAddress localAddress, final SocketAddress peerAddress) %7B%0A        this.delegate = delegate;%0A        this.localAddress = localAddress;%0A        this.peerAddress = peerAddress;%0A    %7D%0A%0A    public void handle(final Callback%5B%5D callbacks) throws IOException, UnsupportedCallbackException %7B%0A        if (localAddress == null && peerAddress == null %7C%7C ! once.compareAndSet(false, true)) %7B%0A            delegate.handle(callbacks);%0A            return;%0A        %7D%0A        final SocketAddressCallback localCallback = localAddress == null ? null : new SocketAddressCallback(localAddress, SocketAddressCallback.Kind.LOCAL);%0A        final SocketAddressCallback peerCallback = peerAddress == null ? null : new SocketAddressCallback(peerAddress, SocketAddressCallback.Kind.PEER);%0A        final Callback%5B%5D newCallbacks;%0A        final int length = callbacks.length;%0A        if (localCallback != null && peerCallback != null) %7B%0A            newCallbacks = new Callback%5Blength + 2%5D;%0A            newCallbacks%5B0%5D = localCallback;%0A            newCallbacks%5B1%5D = peerCallback;%0A            System.arraycopy(callbacks, 0, newCallbacks, 2, length);%0A        %7D else if (localCallback != null) %7B%0A            newCallbacks = new Callback%5Blength + 1%5D;%0A            newCallbacks%5B0%5D = localCallback;%0A            System.arraycopy(callbacks, 0, newCallbacks, 1, length);%0A        %7D else %7B%0A            // peerCallback != null%0A            newCallbacks = new Callback%5Blength + 1%5D;%0A            newCallbacks%5B0%5D = peerCallback;%0A            System.arraycopy(callbacks, 0, newCallbacks, 1, length);%0A        %7D%0A        try %7B%0A            delegate.handle(newCallbacks);%0A        %7D catch (UnsupportedCallbackException e) %7B%0A            if (e.getCallback() instanceof SocketAddressCallback) %7B%0A                delegate.handle(callbacks);%0A                return;%0A            %7D%0A            throw e;%0A        %7D%0A    %7D%0A%7D%0A

Add a callback handler wrapper which passes the local and peer socket address to the callback