AmqpConnection.java

/**
 * Licensed to the Apache Software Foundation (ASF) under one or more
 * contributor license agreements.  See the NOTICE file distributed with
 * this work for additional information regarding copyright ownership.
 * The ASF licenses this file to You 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 org.apache.activemq.transport.amqp.client;

import static org.apache.activemq.transport.amqp.AmqpSupport.CONNECTION_OPEN_FAILED;

import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.util.Collections;
import java.util.List;
import java.util.Map;
import java.util.UUID;
import java.util.concurrent.ScheduledExecutorService;
import java.util.concurrent.ScheduledThreadPoolExecutor;
import java.util.concurrent.ThreadFactory;
import java.util.concurrent.TimeUnit;
import java.util.concurrent.atomic.AtomicBoolean;
import java.util.concurrent.atomic.AtomicLong;

import org.apache.activemq.transport.InactivityIOException;
import org.apache.activemq.transport.amqp.client.sasl.SaslAuthenticator;
import org.apache.activemq.transport.amqp.client.transport.NettyTransportListener;
import org.apache.activemq.transport.amqp.client.util.AsyncResult;
import org.apache.activemq.transport.amqp.client.util.ClientFuture;
import org.apache.activemq.transport.amqp.client.util.IdGenerator;
import org.apache.activemq.transport.amqp.client.util.NoOpAsyncResult;
import org.apache.activemq.transport.amqp.client.util.UnmodifiableProxy;
import org.apache.qpid.proton.amqp.Symbol;
import org.apache.qpid.proton.engine.Collector;
import org.apache.qpid.proton.engine.Connection;
import org.apache.qpid.proton.engine.Delivery;
import org.apache.qpid.proton.engine.EndpointState;
import org.apache.qpid.proton.engine.Event;
import org.apache.qpid.proton.engine.Event.Type;
import org.apache.qpid.proton.engine.Sasl;
import org.apache.qpid.proton.engine.Transport;
import org.apache.qpid.proton.engine.impl.CollectorImpl;
import org.apache.qpid.proton.engine.impl.TransportImpl;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;

import io.netty.buffer.ByteBuf;
import io.netty.buffer.Unpooled;
import io.netty.util.ReferenceCountUtil;

public class AmqpConnection extends AmqpAbstractResource<Connection> implements NettyTransportListener {

   private static final Logger LOG = LoggerFactory.getLogger(AmqpConnection.class);

   private static final NoOpAsyncResult NOOP_REQUEST = new NoOpAsyncResult();

   private static final int DEFAULT_MAX_FRAME_SIZE = 1024 * 1024 * 1;
   // NOTE: Limit default channel max to signed short range to deal with
   //       brokers that don't currently handle the unsigned range well.
   private static final int DEFAULT_CHANNEL_MAX = 32767;
   private static final IdGenerator CONNECTION_ID_GENERATOR = new IdGenerator();

   public static final long DEFAULT_CONNECT_TIMEOUT = 515000;
   public static final long DEFAULT_CLOSE_TIMEOUT = 30000;
   public static final long DEFAULT_DRAIN_TIMEOUT = 60000;

   private ScheduledThreadPoolExecutor serializer;
   private final AtomicBoolean closed = new AtomicBoolean();
   private final AtomicBoolean connected = new AtomicBoolean();
   private final AtomicLong sessionIdGenerator = new AtomicLong();
   private final AtomicLong txIdGenerator = new AtomicLong();
   private final Collector protonCollector = new CollectorImpl();
   private final org.apache.activemq.transport.amqp.client.transport.NettyTransport transport;
   private final Transport protonTransport = Transport.Factory.create();

   private final String username;
   private final String password;
   private final URI remoteURI;
   private final String connectionId;
   private List<Symbol> desiredCapabilities = Collections.emptyList();
   private List<Symbol> offeredCapabilities = Collections.emptyList();
   private Map<Symbol, Object> offeredProperties = Collections.emptyMap();

   private volatile AmqpFrameValidator sentFrameInspector;
   private volatile AmqpFrameValidator receivedFrameInspector;
   private AmqpConnectionListener listener;
   private SaslAuthenticator authenticator;
   private String mechanismRestriction;
   private String authzid;

   private int idleTimeout = 0;
   private boolean idleProcessingDisabled;
   private String containerId;
   private boolean authenticated;
   private int channelMax = DEFAULT_CHANNEL_MAX;
   private long connectTimeout = DEFAULT_CONNECT_TIMEOUT;
   private long closeTimeout = DEFAULT_CLOSE_TIMEOUT;
   private long drainTimeout = DEFAULT_DRAIN_TIMEOUT;
   private boolean trace;
   private boolean noContainerID = false;

   public AmqpConnection(org.apache.activemq.transport.amqp.client.transport.NettyTransport transport,
                         String username,
                         String password) {
      setEndpoint(Connection.Factory.create());
      getEndpoint().collect(protonCollector);

      this.transport = transport;
      this.username = username;
      this.password = password;
      this.connectionId = CONNECTION_ID_GENERATOR.generateId();
      this.remoteURI = transport.getRemoteLocation();

      this.serializer = new ScheduledThreadPoolExecutor(1, new ThreadFactory() {

         @Override
         public Thread newThread(Runnable runner) {
            Thread serial = new Thread(runner);
            serial.setDaemon(true);
            serial.setName(toString());
            return serial;
         }
      });

      // Ensure timely shutdown
      this.serializer.setExecuteExistingDelayedTasksAfterShutdownPolicy(false);
      this.serializer.setContinueExistingPeriodicTasksAfterShutdownPolicy(false);

      this.transport.setTransportListener(this);
   }

   public void connect() throws Exception {
      if (connected.compareAndSet(false, true)) {
         transport.connect();

         final ClientFuture future = new ClientFuture();
         serializer.execute(new Runnable() {
            @Override
            public void run() {
               if (!noContainerID) {
                  getEndpoint().setContainer(safeGetContainerId());
               }
               getEndpoint().setHostname(remoteURI.getHost());
               if (!getDesiredCapabilities().isEmpty()) {
                  getEndpoint().setDesiredCapabilities(getDesiredCapabilities().toArray(new Symbol[0]));
               }
               if (!getOfferedCapabilities().isEmpty()) {
                  getEndpoint().setOfferedCapabilities(getOfferedCapabilities().toArray(new Symbol[0]));
               }
               if (!getOfferedProperties().isEmpty()) {
                  getEndpoint().setProperties(getOfferedProperties());
               }

               if (getIdleTimeout() > 0) {
                  protonTransport.setIdleTimeout(getIdleTimeout());
               }
               protonTransport.setMaxFrameSize(getMaxFrameSize());
               protonTransport.setChannelMax(getChannelMax());
               protonTransport.bind(getEndpoint());
               Sasl sasl = protonTransport.sasl();
               if (sasl != null) {
                  sasl.client();
               }
               authenticator = new SaslAuthenticator(sasl, username, password, authzid, mechanismRestriction);
               ((TransportImpl) protonTransport).setProtocolTracer(new AmqpProtocolTracer(AmqpConnection.this));
               open(future);

               pumpToProtonTransport(future);
            }
         });

         try {
            if (connectTimeout <= 0) {
               future.sync();
            } else {
               future.sync(connectTimeout, TimeUnit.MILLISECONDS);
               if (getEndpoint().getRemoteState() != EndpointState.ACTIVE) {
                  throw new IOException("Failed to connect after configured timeout.");
               }
            }
         } catch (Throwable e) {
            try {
               close();
            } catch (Throwable ignore) {
            }

            throw e;
         }
      }
   }

   public boolean isConnected() {
      return transport.isConnected() && connected.get();
   }

   public void close() {
      if (closed.compareAndSet(false, true)) {
         final ClientFuture request = new ClientFuture();
         serializer.execute(new Runnable() {

            @Override
            public void run() {
               try {

                  // If we are not connected then there is nothing we can do now
                  // just signal success.
                  if (!transport.isConnected()) {
                     request.onSuccess();
                  }

                  if (getEndpoint() != null) {
                     close(request);
                  } else {
                     request.onSuccess();
                  }

                  pumpToProtonTransport(request);
               } catch (Exception e) {
                  LOG.debug("Caught exception while closing proton connection");
               }
            }
         });

         try {
            if (closeTimeout <= 0) {
               request.sync();
            } else {
               request.sync(closeTimeout, TimeUnit.MILLISECONDS);
            }
         } catch (IOException e) {
            LOG.warn("Error caught while closing Provider: ", e.getMessage());
         } finally {
            if (transport != null) {
               try {
                  transport.close();
               } catch (Exception e) {
                  LOG.debug("Cuaght exception while closing down Transport: {}", e.getMessage());
               }
            }

            serializer.shutdownNow();
            try {
               if (!serializer.awaitTermination(10, TimeUnit.SECONDS)) {
                  LOG.warn("Serializer didn't shutdown cleanly");
               }
            } catch (InterruptedException e) {
            }
         }
      }
   }

   /**
    * Creates a new Session instance used to create AMQP resources like
    * senders and receivers.
    *
    * @return a new AmqpSession that can be used to create links.
    * @throws Exception if an error occurs during creation.
    */
   public AmqpSession createSession() throws Exception {
      checkClosed();

      final AmqpSession session = new AmqpSession(AmqpConnection.this, getNextSessionId());
      final ClientFuture request = new ClientFuture();

      serializer.execute(new Runnable() {

         @Override
         public void run() {
            checkClosed();
            session.setEndpoint(getEndpoint().session());
            session.setStateInspector(getStateInspector());
            session.open(request);
            pumpToProtonTransport(request);
         }
      });

      request.sync();

      return session;
   }

   //----- Access to low level IO for specific test cases -------------------//

   public void sendRawBytes(final byte[] rawData) throws Exception {
      checkClosed();

      final ClientFuture request = new ClientFuture();

      serializer.execute(new Runnable() {

         @Override
         public void run() {
            checkClosed();
            try {
               transport.send(Unpooled.wrappedBuffer(rawData));
            } catch (IOException e) {
               fireClientException(e);
            } finally {
               request.onSuccess();
            }
         }
      });

      request.sync();
   }

   //----- Configuration accessors ------------------------------------------//

   /**
    * @return the user name that was used to authenticate this connection.
    */
   public String getUsername() {
      return username;
   }

   /**
    * @return the password that was used to authenticate this connection.
    */
   public String getPassword() {
      return password;
   }

   public void setAuthzid(String authzid) {
      this.authzid = authzid;
   }

   public String getAuthzid() {
      return authzid;
   }

   /**
    * @return the URI of the remote peer this connection attached to.
    */
   public URI getRemoteURI() {
      return remoteURI;
   }

   /**
    * @return the container ID that will be set as the container Id.
    */
   public String getContainerId() {
      return this.containerId;
   }

   /**
    * Sets the container Id that will be configured on the connection prior to
    * connecting to the remote peer.  Calling this after connect has no effect.
    *
    * @param containerId the container Id to use on the connection.
    */
   public void setContainerId(String containerId) {
      this.containerId = containerId;
   }

   /**
    * @return the currently set Max Frame Size value.
    */
   public int getMaxFrameSize() {
      return DEFAULT_MAX_FRAME_SIZE;
   }

   public int getChannelMax() {
      return channelMax;
   }

   public void setChannelMax(int channelMax) {
      this.channelMax = channelMax;
   }

   public long getConnectTimeout() {
      return connectTimeout;
   }

   public void setConnectTimeout(long connectTimeout) {
      this.connectTimeout = connectTimeout;
   }

   public long getCloseTimeout() {
      return closeTimeout;
   }

   public void setCloseTimeout(long closeTimeout) {
      this.closeTimeout = closeTimeout;
   }

   public long getDrainTimeout() {
      return drainTimeout;
   }

   public void setDrainTimeout(long drainTimeout) {
      this.drainTimeout = drainTimeout;
   }

   public List<Symbol> getDesiredCapabilities() {
      return desiredCapabilities;
   }

   public void setDesiredCapabilities(List<Symbol> desiredCapabilities) {
      if (desiredCapabilities == null) {
         desiredCapabilities = Collections.emptyList();
      }

      this.desiredCapabilities = desiredCapabilities;
   }

   public List<Symbol> getOfferedCapabilities() {
      return offeredCapabilities;
   }

   public void setOfferedCapabilities(List<Symbol> offeredCapabilities) {
      if (offeredCapabilities == null) {
         offeredCapabilities = Collections.emptyList();
      }

      this.offeredCapabilities = offeredCapabilities;
   }

   public Map<Symbol, Object> getOfferedProperties() {
      return offeredProperties;
   }

   public void setOfferedProperties(Map<Symbol, Object> offeredProperties) {
      if (offeredProperties == null) {
         offeredProperties = Collections.emptyMap();
      }

      this.offeredProperties = offeredProperties;
   }

   public Connection getConnection() {
      return UnmodifiableProxy.connectionProxy(getEndpoint());
   }

   public AmqpConnectionListener getListener() {
      return listener;
   }

   public void setListener(AmqpConnectionListener listener) {
      this.listener = listener;
   }

   public int getIdleTimeout() {
      return idleTimeout;
   }

   public void setIdleTimeout(int idleTimeout) {
      this.idleTimeout = idleTimeout;
   }

   public void setIdleProcessingDisabled(boolean value) {
      this.idleProcessingDisabled = value;
   }

   public boolean isIdleProcessingDisabled() {
      return idleProcessingDisabled;
   }

   /**
    * Sets a restriction on the SASL mechanism to use (if offered by the server).
    *
    * @param mechanismRestriction the mechanism to use
    */
   public void setMechanismRestriction(String mechanismRestriction) {
      this.mechanismRestriction = mechanismRestriction;
   }

   public String getMechanismRestriction() {
      return mechanismRestriction;
   }

   public boolean isTraceFrames() {
      return trace;
   }

   public void setTraceFrames(boolean trace) {
      this.trace = trace;
   }

   public AmqpFrameValidator getSentFrameInspector() {
      return sentFrameInspector;
   }

   public void setSentFrameInspector(AmqpFrameValidator amqpFrameInspector) {
      this.sentFrameInspector = amqpFrameInspector;
   }

   public AmqpFrameValidator getReceivedFrameInspector() {
      return receivedFrameInspector;
   }

   public void setReceivedFrameInspector(AmqpFrameValidator amqpFrameInspector) {
      this.receivedFrameInspector = amqpFrameInspector;
   }

   //----- Internal getters used from the child AmqpResource classes --------//

   ScheduledExecutorService getScheduler() {
      return this.serializer;
   }

   Connection getProtonConnection() {
      return getEndpoint();
   }

   String getConnectionId() {
      return this.connectionId;
   }

   AmqpTransactionId getNextTransactionId() {
      return new AmqpTransactionId(connectionId + ":" + txIdGenerator.incrementAndGet());
   }

   void pumpToProtonTransport() {
      pumpToProtonTransport(NOOP_REQUEST);
   }

   void pumpToProtonTransport(AsyncResult request) {
      try {
         boolean done = false;
         while (!done) {
            ByteBuffer toWrite = protonTransport.getOutputBuffer();
            if (toWrite != null && toWrite.hasRemaining()) {
               ByteBuf outbound = transport.allocateSendBuffer(toWrite.remaining());
               outbound.writeBytes(toWrite);
               transport.send(outbound);
               protonTransport.outputConsumed();
            } else {
               done = true;
            }
         }
      } catch (IOException e) {
         fireClientException(e);
         request.onFailure(e);
      }
   }

   //----- Transport listener event hooks -----------------------------------//

   @Override
   public void onData(final ByteBuf incoming) {

      // We need to retain until the serializer gets around to processing it.
      ReferenceCountUtil.retain(incoming);

      serializer.execute(new Runnable() {

         @Override
         public void run() {
            ByteBuffer source = incoming.nioBuffer();
            LOG.trace("Client Received from Broker {} bytes:", source.remaining());

            if (protonTransport.isClosed()) {
               LOG.debug("Ignoring incoming data because transport is closed");
               return;
            }

            do {
               ByteBuffer buffer = protonTransport.getInputBuffer();
               int limit = Math.min(buffer.remaining(), source.remaining());
               ByteBuffer duplicate = source.duplicate();
               duplicate.limit(source.position() + limit);
               buffer.put(duplicate);
               protonTransport.processInput();
               source.position(source.position() + limit);
            } while (source.hasRemaining());

            ReferenceCountUtil.release(incoming);

            // Process the state changes from the latest data and then answer back
            // any pending updates to the Broker.
            processUpdates();
            pumpToProtonTransport();
         }
      });
   }

   @Override
   public void onTransportClosed() {
      LOG.debug("The transport has unexpectedly closed");
      failed(getOpenAbortException());
   }

   @Override
   public void onTransportError(Throwable cause) {
      fireClientException(cause);
   }

   //----- Internal implementation ------------------------------------------//

   @Override
   protected void doOpenCompletion() {
      // If the remote indicates that a close is pending, don't open.
      if (getEndpoint().getRemoteProperties() == null || !getEndpoint().getRemoteProperties().containsKey(CONNECTION_OPEN_FAILED)) {

         if (!isIdleProcessingDisabled()) {
            // Using nano time since it is not related to the wall clock, which may change
            long initialNow = TimeUnit.NANOSECONDS.toMillis(System.nanoTime());
            long initialKeepAliveDeadline = protonTransport.tick(initialNow);
            if (initialKeepAliveDeadline > 0) {

               getScheduler().schedule(new Runnable() {

                  @Override
                  public void run() {
                     try {
                        if (getEndpoint().getLocalState() != EndpointState.CLOSED) {
                           LOG.debug("Client performing next idle check");
                           // Using nano time since it is not related to the wall clock, which may change
                           long now = TimeUnit.NANOSECONDS.toMillis(System.nanoTime());
                           long rescheduleAt = protonTransport.tick(now) - now;
                           pumpToProtonTransport();
                           if (protonTransport.isClosed()) {
                              LOG.debug("Transport closed after inactivity check.");
                              throw new InactivityIOException("Channel was inactive for to long");
                           }

                           if (rescheduleAt > 0) {
                              getScheduler().schedule(this, rescheduleAt, TimeUnit.MILLISECONDS);
                           }
                        }
                     } catch (Exception e) {
                        try {
                           transport.close();
                        } catch (IOException e1) {
                        }
                        fireClientException(e);
                     }
                  }
               }, initialKeepAliveDeadline - initialNow, TimeUnit.MILLISECONDS);
            }
         }
         super.doOpenCompletion();
      }
   }

   @Override
   protected void doOpenInspection() {
      try {
         getStateInspector().inspectOpenedResource(getConnection());
      } catch (Throwable error) {
         getStateInspector().markAsInvalid(error.getMessage());
      }
   }

   @Override
   protected void doClosedInspection() {
      try {
         getStateInspector().inspectClosedResource(getConnection());
      } catch (Throwable error) {
         getStateInspector().markAsInvalid(error.getMessage());
      }
   }

   protected void fireClientException(Throwable ex) {
      AmqpConnectionListener listener = this.listener;
      if (listener != null) {
         listener.onException(ex);
      }
   }

   protected void checkClosed() throws IllegalStateException {
      if (closed.get()) {
         throw new IllegalStateException("The Connection is already closed");
      }
   }

   private void processUpdates() {
      try {
         Event protonEvent = null;
         while ((protonEvent = protonCollector.peek()) != null) {
            if (!protonEvent.getType().equals(Type.TRANSPORT)) {
               LOG.trace("Client: New Proton Event: {}", protonEvent.getType());
            }

            AmqpEventSink amqpEventSink = null;
            switch (protonEvent.getType()) {
               case CONNECTION_REMOTE_CLOSE:
                  amqpEventSink = (AmqpEventSink) protonEvent.getConnection().getContext();
                  amqpEventSink.processRemoteClose(this);
                  break;
               case CONNECTION_REMOTE_OPEN:
                  amqpEventSink = (AmqpEventSink) protonEvent.getConnection().getContext();
                  amqpEventSink.processRemoteOpen(this);
                  break;
               case SESSION_REMOTE_CLOSE:
                  amqpEventSink = (AmqpEventSink) protonEvent.getSession().getContext();
                  amqpEventSink.processRemoteClose(this);
                  break;
               case SESSION_REMOTE_OPEN:
                  amqpEventSink = (AmqpEventSink) protonEvent.getSession().getContext();
                  amqpEventSink.processRemoteOpen(this);
                  break;
               case LINK_REMOTE_CLOSE:
                  amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                  amqpEventSink.processRemoteClose(this);
                  break;
               case LINK_REMOTE_DETACH:
                  amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                  amqpEventSink.processRemoteDetach(this);
                  break;
               case LINK_REMOTE_OPEN:
                  amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                  amqpEventSink.processRemoteOpen(this);
                  break;
               case LINK_FLOW:
                  amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                  amqpEventSink.processFlowUpdates(this);
                  break;
               case DELIVERY:
                  amqpEventSink = (AmqpEventSink) protonEvent.getLink().getContext();
                  amqpEventSink.processDeliveryUpdates(this, (Delivery) protonEvent.getContext());
                  break;
               default:
                  break;
            }

            protonCollector.pop();
         }

         // We have to do this to pump SASL bytes in as SASL is not event driven yet.
         if (!authenticated) {
            processSaslAuthentication();
         }
      } catch (Exception ex) {
         LOG.warn("Caught Exception during update processing: {}", ex.getMessage(), ex);
         fireClientException(ex);
      }
   }

   private void processSaslAuthentication() {
      if (authenticated || authenticator == null) {
         return;
      }

      try {
         if (authenticator.authenticate()) {
            authenticator = null;
            authenticated = true;
         }
      } catch (SecurityException ex) {
         failed(ex);
      }
   }

   private String getNextSessionId() {
      return connectionId + ":" + sessionIdGenerator.incrementAndGet();
   }

   private String safeGetContainerId() {
      String containerId = getContainerId();
      if (containerId == null || containerId.isEmpty()) {
         containerId = UUID.randomUUID().toString();
      }

      return containerId;
   }

   @Override
   public String toString() {
      return "AmqpConnection { " + connectionId + " }";
   }

   public void setNoContainerID() {
      noContainerID = true;
   }
}