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