001/**
002 * Licensed to the Apache Software Foundation (ASF) under one or more
003 * contributor license agreements.  See the NOTICE file distributed with
004 * this work for additional information regarding copyright ownership.
005 * The ASF licenses this file to You under the Apache License, Version 2.0
006 * (the "License"); you may not use this file except in compliance with
007 * the License.  You may obtain a copy of the License at
008 *
009 *      http://www.apache.org/licenses/LICENSE-2.0
010 *
011 * Unless required by applicable law or agreed to in writing, software
012 * distributed under the License is distributed on an "AS IS" BASIS,
013 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
014 * See the License for the specific language governing permissions and
015 * limitations under the License.
016 */
017
018package org.apache.activemq.transport.nio;
019
020import java.io.DataInputStream;
021import java.io.DataOutputStream;
022import java.io.EOFException;
023import java.io.IOException;
024import java.net.Socket;
025import java.net.SocketTimeoutException;
026import java.net.URI;
027import java.net.UnknownHostException;
028import java.nio.ByteBuffer;
029import java.nio.channels.SelectionKey;
030import java.nio.channels.Selector;
031import java.security.cert.X509Certificate;
032import java.util.concurrent.CountDownLatch;
033
034import javax.net.SocketFactory;
035import javax.net.ssl.SSLContext;
036import javax.net.ssl.SSLEngine;
037import javax.net.ssl.SSLEngineResult;
038import javax.net.ssl.SSLParameters;
039import javax.net.ssl.SSLEngineResult.HandshakeStatus;
040import javax.net.ssl.SSLPeerUnverifiedException;
041import javax.net.ssl.SSLSession;
042
043import org.apache.activemq.command.ConnectionInfo;
044import org.apache.activemq.openwire.OpenWireFormat;
045import org.apache.activemq.thread.TaskRunnerFactory;
046import org.apache.activemq.util.IOExceptionSupport;
047import org.apache.activemq.util.ServiceStopper;
048import org.apache.activemq.wireformat.WireFormat;
049import org.slf4j.Logger;
050import org.slf4j.LoggerFactory;
051
052public class NIOSSLTransport extends NIOTransport {
053
054    private static final Logger LOG = LoggerFactory.getLogger(NIOSSLTransport.class);
055
056    protected boolean needClientAuth;
057    protected boolean wantClientAuth;
058    protected String[] enabledCipherSuites;
059    protected String[] enabledProtocols;
060    protected boolean verifyHostName = false;
061
062    protected SSLContext sslContext;
063    protected SSLEngine sslEngine;
064    protected SSLSession sslSession;
065
066    protected volatile boolean handshakeInProgress = false;
067    protected SSLEngineResult.Status status = null;
068    protected SSLEngineResult.HandshakeStatus handshakeStatus = null;
069    protected TaskRunnerFactory taskRunnerFactory;
070
071    public NIOSSLTransport(WireFormat wireFormat, SocketFactory socketFactory, URI remoteLocation, URI localLocation) throws UnknownHostException, IOException {
072        super(wireFormat, socketFactory, remoteLocation, localLocation);
073    }
074
075    public NIOSSLTransport(WireFormat wireFormat, Socket socket, SSLEngine engine, InitBuffer initBuffer,
076            ByteBuffer inputBuffer) throws IOException {
077        super(wireFormat, socket, initBuffer);
078        this.sslEngine = engine;
079        if (engine != null) {
080            this.sslSession = engine.getSession();
081        }
082        this.inputBuffer = inputBuffer;
083    }
084
085    public void setSslContext(SSLContext sslContext) {
086        this.sslContext = sslContext;
087    }
088
089    volatile boolean hasSslEngine = false;
090
091    @Override
092    protected void initializeStreams() throws IOException {
093        if (sslEngine != null) {
094            hasSslEngine = true;
095        }
096        NIOOutputStream outputStream = null;
097        try {
098            channel = socket.getChannel();
099            channel.configureBlocking(false);
100
101            if (sslContext == null) {
102                sslContext = SSLContext.getDefault();
103            }
104
105            String remoteHost = null;
106            int remotePort = -1;
107
108            try {
109                URI remoteAddress = new URI(this.getRemoteAddress());
110                remoteHost = remoteAddress.getHost();
111                remotePort = remoteAddress.getPort();
112            } catch (Exception e) {
113            }
114
115            // initialize engine, the initial sslSession we get will need to be
116            // updated once the ssl handshake process is completed.
117            if (!hasSslEngine) {
118                if (remoteHost != null && remotePort != -1) {
119                    sslEngine = sslContext.createSSLEngine(remoteHost, remotePort);
120                } else {
121                    sslEngine = sslContext.createSSLEngine();
122                }
123
124                if (verifyHostName) {
125                    SSLParameters sslParams = new SSLParameters();
126                    sslParams.setEndpointIdentificationAlgorithm("HTTPS");
127                    sslEngine.setSSLParameters(sslParams);
128                }
129
130                sslEngine.setUseClientMode(false);
131                if (enabledCipherSuites != null) {
132                    sslEngine.setEnabledCipherSuites(enabledCipherSuites);
133                }
134
135                if (enabledProtocols != null) {
136                    sslEngine.setEnabledProtocols(enabledProtocols);
137                }
138
139                if (wantClientAuth) {
140                    sslEngine.setWantClientAuth(wantClientAuth);
141                }
142
143                if (needClientAuth) {
144                    sslEngine.setNeedClientAuth(needClientAuth);
145                }
146
147                sslSession = sslEngine.getSession();
148
149                inputBuffer = ByteBuffer.allocate(sslSession.getPacketBufferSize());
150                inputBuffer.clear();
151            }
152
153            outputStream = new NIOOutputStream(channel);
154            outputStream.setEngine(sslEngine);
155            this.dataOut = new DataOutputStream(outputStream);
156            this.buffOut = outputStream;
157
158            //If the sslEngine was not passed in, then handshake
159            if (!hasSslEngine) {
160                sslEngine.beginHandshake();
161            }
162            handshakeStatus = sslEngine.getHandshakeStatus();
163            if (!hasSslEngine) {
164                doHandshake();
165            }
166
167            selection = SelectorManager.getInstance().register(channel, new SelectorManager.Listener() {
168                @Override
169                public void onSelect(SelectorSelection selection) {
170                    try {
171                        initialized.await();
172                    } catch (InterruptedException error) {
173                        onException(IOExceptionSupport.create(error));
174                    }
175                    serviceRead();
176                }
177
178                @Override
179                public void onError(SelectorSelection selection, Throwable error) {
180                    if (error instanceof IOException) {
181                        onException((IOException) error);
182                    } else {
183                        onException(IOExceptionSupport.create(error));
184                    }
185                }
186            });
187            doInit();
188
189        } catch (Exception e) {
190            try {
191                if(outputStream != null) {
192                    outputStream.close();
193                }
194                super.closeStreams();
195            } catch (Exception ex) {}
196            throw new IOException(e);
197        }
198    }
199
200    final protected CountDownLatch initialized = new CountDownLatch(1);
201
202    protected void doInit() throws Exception {
203        taskRunnerFactory.execute(new Runnable() {
204
205            @Override
206            public void run() {
207                //Need to start in new thread to let startup finish first
208                //We can trigger a read because we know the channel is ready since the SSL handshake
209                //already happened
210                serviceRead();
211                initialized.countDown();
212            }
213        });
214    }
215
216    //Only used for the auto transport to abort the openwire init method early if already initialized
217    boolean openWireInititialized = false;
218
219    protected void doOpenWireInit() throws Exception {
220        //Do this later to let wire format negotiation happen
221        if (initBuffer != null && !openWireInititialized && this.wireFormat instanceof OpenWireFormat) {
222            initBuffer.buffer.flip();
223            if (initBuffer.buffer.hasRemaining()) {
224                nextFrameSize = -1;
225                receiveCounter += initBuffer.readSize;
226                processCommand(initBuffer.buffer);
227                processCommand(initBuffer.buffer);
228                initBuffer.buffer.clear();
229                openWireInititialized = true;
230            }
231        }
232    }
233
234    protected void finishHandshake() throws Exception {
235        if (handshakeInProgress) {
236            handshakeInProgress = false;
237            nextFrameSize = -1;
238
239            // Once handshake completes we need to ask for the now real sslSession
240            // otherwise the session would return 'SSL_NULL_WITH_NULL_NULL' for the
241            // cipher suite.
242            sslSession = sslEngine.getSession();
243        }
244    }
245
246    @Override
247    public void serviceRead() {
248        try {
249            if (handshakeInProgress) {
250                doHandshake();
251            }
252
253            doOpenWireInit();
254
255            ByteBuffer plain = ByteBuffer.allocate(sslSession.getApplicationBufferSize());
256            plain.position(plain.limit());
257
258            while (true) {
259                if (!plain.hasRemaining()) {
260
261                    int readCount = secureRead(plain);
262
263                    if (readCount == 0) {
264                        break;
265                    }
266
267                    // channel is closed, cleanup
268                    if (readCount == -1) {
269                        onException(new EOFException());
270                        selection.close();
271                        break;
272                    }
273
274                    receiveCounter += readCount;
275                }
276
277                if (status == SSLEngineResult.Status.OK && handshakeStatus != SSLEngineResult.HandshakeStatus.NEED_UNWRAP) {
278                    processCommand(plain);
279                }
280            }
281        } catch (IOException e) {
282            onException(e);
283        } catch (Throwable e) {
284            onException(IOExceptionSupport.create(e));
285        }
286    }
287
288    protected void processCommand(ByteBuffer plain) throws Exception {
289
290        // Are we waiting for the next Command or are we building on the current one
291        if (nextFrameSize == -1) {
292
293            // We can get small packets that don't give us enough for the frame size
294            // so allocate enough for the initial size value and
295            if (plain.remaining() < Integer.SIZE) {
296                if (currentBuffer == null) {
297                    currentBuffer = ByteBuffer.allocate(4);
298                }
299
300                // Go until we fill the integer sized current buffer.
301                while (currentBuffer.hasRemaining() && plain.hasRemaining()) {
302                    currentBuffer.put(plain.get());
303                }
304
305                // Didn't we get enough yet to figure out next frame size.
306                if (currentBuffer.hasRemaining()) {
307                    return;
308                } else {
309                    currentBuffer.flip();
310                    nextFrameSize = currentBuffer.getInt();
311                }
312
313            } else {
314
315                // Either we are completing a previous read of the next frame size or its
316                // fully contained in plain already.
317                if (currentBuffer != null) {
318
319                    // Finish the frame size integer read and get from the current buffer.
320                    while (currentBuffer.hasRemaining()) {
321                        currentBuffer.put(plain.get());
322                    }
323
324                    currentBuffer.flip();
325                    nextFrameSize = currentBuffer.getInt();
326
327                } else {
328                    nextFrameSize = plain.getInt();
329                }
330            }
331
332            if (wireFormat instanceof OpenWireFormat) {
333                long maxFrameSize = ((OpenWireFormat) wireFormat).getMaxFrameSize();
334                if (nextFrameSize > maxFrameSize) {
335                    throw new IOException("Frame size of " + (nextFrameSize / (1024 * 1024)) +
336                                          " MB larger than max allowed " + (maxFrameSize / (1024 * 1024)) + " MB");
337                }
338            }
339
340            // now we got the data, lets reallocate and store the size for the marshaler.
341            // if there's more data in plain, then the next call will start processing it.
342            currentBuffer = ByteBuffer.allocate(nextFrameSize + 4);
343            currentBuffer.putInt(nextFrameSize);
344
345        } else {
346            // If its all in one read then we can just take it all, otherwise take only
347            // the current frame size and the next iteration starts a new command.
348            if (currentBuffer != null) {
349                if (currentBuffer.remaining() >= plain.remaining()) {
350                    currentBuffer.put(plain);
351                } else {
352                    byte[] fill = new byte[currentBuffer.remaining()];
353                    plain.get(fill);
354                    currentBuffer.put(fill);
355                }
356
357                // Either we have enough data for a new command or we have to wait for some more.
358                if (currentBuffer.hasRemaining()) {
359                    return;
360                } else {
361                    currentBuffer.flip();
362                    Object command = wireFormat.unmarshal(new DataInputStream(new NIOInputStream(currentBuffer)));
363                    doConsume(command);
364                    nextFrameSize = -1;
365                    currentBuffer = null;
366               }
367            }
368        }
369    }
370
371    protected int secureRead(ByteBuffer plain) throws Exception {
372
373        if (!(inputBuffer.position() != 0 && inputBuffer.hasRemaining()) || status == SSLEngineResult.Status.BUFFER_UNDERFLOW) {
374            int bytesRead = channel.read(inputBuffer);
375
376            if (bytesRead == 0 && !(sslEngine.getHandshakeStatus().equals(SSLEngineResult.HandshakeStatus.NEED_UNWRAP))) {
377                return 0;
378            }
379
380            if (bytesRead == -1) {
381                sslEngine.closeInbound();
382                if (inputBuffer.position() == 0 || status == SSLEngineResult.Status.BUFFER_UNDERFLOW) {
383                    return -1;
384                }
385            }
386        }
387
388        plain.clear();
389
390        inputBuffer.flip();
391        SSLEngineResult res;
392        do {
393            res = sslEngine.unwrap(inputBuffer, plain);
394        } while (res.getStatus() == SSLEngineResult.Status.OK && res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.NEED_UNWRAP
395                && res.bytesProduced() == 0);
396
397        if (res.getHandshakeStatus() == SSLEngineResult.HandshakeStatus.FINISHED) {
398            finishHandshake();
399        }
400
401        status = res.getStatus();
402        handshakeStatus = res.getHandshakeStatus();
403
404        // TODO deal with BUFFER_OVERFLOW
405
406        if (status == SSLEngineResult.Status.CLOSED) {
407            sslEngine.closeInbound();
408            return -1;
409        }
410
411        inputBuffer.compact();
412        plain.flip();
413
414        return plain.remaining();
415    }
416
417    protected void doHandshake() throws Exception {
418        handshakeInProgress = true;
419        Selector selector = null;
420        SelectionKey key = null;
421        boolean readable = true;
422        try {
423            while (true) {
424                HandshakeStatus handshakeStatus = sslEngine.getHandshakeStatus();
425                switch (handshakeStatus) {
426                    case NEED_UNWRAP:
427                        if (readable) {
428                            secureRead(ByteBuffer.allocate(sslSession.getApplicationBufferSize()));
429                        }
430                        if (this.status == SSLEngineResult.Status.BUFFER_UNDERFLOW) {
431                            long now = System.currentTimeMillis();
432                            if (selector == null) {
433                                selector = Selector.open();
434                                key = channel.register(selector, SelectionKey.OP_READ);
435                            } else {
436                                key.interestOps(SelectionKey.OP_READ);
437                            }
438                            int keyCount = selector.select(this.getSoTimeout());
439                            if (keyCount == 0 && this.getSoTimeout() > 0 && ((System.currentTimeMillis() - now) >= this.getSoTimeout())) {
440                                throw new SocketTimeoutException("Timeout during handshake");
441                            }
442                            readable = key.isReadable();
443                        }
444                        break;
445                    case NEED_TASK:
446                        Runnable task;
447                        while ((task = sslEngine.getDelegatedTask()) != null) {
448                            task.run();
449                        }
450                        break;
451                    case NEED_WRAP:
452                        ((NIOOutputStream) buffOut).write(ByteBuffer.allocate(0));
453                        break;
454                    case FINISHED:
455                    case NOT_HANDSHAKING:
456                        finishHandshake();
457                        return;
458                }
459            }
460        } finally {
461            if (key!=null) try {key.cancel();} catch (Exception ignore) {}
462            if (selector!=null) try {selector.close();} catch (Exception ignore) {}
463        }
464    }
465
466    @Override
467    protected void doStart() throws Exception {
468        taskRunnerFactory = new TaskRunnerFactory("ActiveMQ NIOSSLTransport Task");
469        // no need to init as we can delay that until demand (eg in doHandshake)
470        super.doStart();
471    }
472
473    @Override
474    protected void doStop(ServiceStopper stopper) throws Exception {
475        initialized.countDown();
476
477        if (taskRunnerFactory != null) {
478            taskRunnerFactory.shutdownNow();
479            taskRunnerFactory = null;
480        }
481        if (channel != null) {
482            channel.close();
483            channel = null;
484        }
485        super.doStop(stopper);
486    }
487
488    /**
489     * Overriding in order to add the client's certificates to ConnectionInfo Commands.
490     *
491     * @param command
492     *            The Command coming in.
493     */
494    @Override
495    public void doConsume(Object command) {
496        if (command instanceof ConnectionInfo) {
497            ConnectionInfo connectionInfo = (ConnectionInfo) command;
498            connectionInfo.setTransportContext(getPeerCertificates());
499        }
500        super.doConsume(command);
501    }
502
503    /**
504     * @return peer certificate chain associated with the ssl socket
505     */
506    @Override
507    public X509Certificate[] getPeerCertificates() {
508
509        X509Certificate[] clientCertChain = null;
510        try {
511            if (sslEngine.getSession() != null) {
512                clientCertChain = (X509Certificate[]) sslEngine.getSession().getPeerCertificates();
513            }
514        } catch (SSLPeerUnverifiedException e) {
515            if (LOG.isTraceEnabled()) {
516                LOG.trace("Failed to get peer certificates.", e);
517            }
518        }
519
520        return clientCertChain;
521    }
522
523    public boolean isNeedClientAuth() {
524        return needClientAuth;
525    }
526
527    public void setNeedClientAuth(boolean needClientAuth) {
528        this.needClientAuth = needClientAuth;
529    }
530
531    public boolean isWantClientAuth() {
532        return wantClientAuth;
533    }
534
535    public void setWantClientAuth(boolean wantClientAuth) {
536        this.wantClientAuth = wantClientAuth;
537    }
538
539    public String[] getEnabledCipherSuites() {
540        return enabledCipherSuites;
541    }
542
543    public void setEnabledCipherSuites(String[] enabledCipherSuites) {
544        this.enabledCipherSuites = enabledCipherSuites;
545    }
546
547    public String[] getEnabledProtocols() {
548        return enabledProtocols;
549    }
550
551    public void setEnabledProtocols(String[] enabledProtocols) {
552        this.enabledProtocols = enabledProtocols;
553    }
554
555        public boolean isVerifyHostName() {
556        return verifyHostName;
557    }
558
559    public void setVerifyHostName(boolean verifyHostName) {
560        this.verifyHostName = verifyHostName;
561    }
562}