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 */
017package org.apache.activemq.transport.ws.jetty9;
018
019import java.io.IOException;
020import java.nio.ByteBuffer;
021import java.util.concurrent.TimeUnit;
022import java.util.concurrent.atomic.AtomicBoolean;
023
024import org.apache.activemq.transport.mqtt.MQTTCodec;
025import org.apache.activemq.transport.ws.AbstractMQTTSocket;
026import org.apache.activemq.util.ByteSequence;
027import org.apache.activemq.util.IOExceptionSupport;
028import org.eclipse.jetty.websocket.api.Session;
029import org.eclipse.jetty.websocket.api.WebSocketListener;
030import org.fusesource.hawtbuf.Buffer;
031import org.fusesource.hawtbuf.DataByteArrayInputStream;
032import org.fusesource.mqtt.codec.DISCONNECT;
033import org.fusesource.mqtt.codec.MQTTFrame;
034import org.slf4j.Logger;
035import org.slf4j.LoggerFactory;
036
037public class MQTTSocket extends AbstractMQTTSocket implements MQTTCodec.MQTTFrameSink, WebSocketListener {
038
039    private static final Logger LOG = LoggerFactory.getLogger(MQTTSocket.class);
040
041    private final int ORDERLY_CLOSE_TIMEOUT = 10;
042    private Session session;
043    private final AtomicBoolean receivedDisconnect = new AtomicBoolean();
044
045    private final MQTTCodec codec;
046
047    public MQTTSocket(String remoteAddress) {
048        super(remoteAddress);
049
050        this.codec = new MQTTCodec(this, getWireFormat());
051    }
052
053    @Override
054    public void sendToMQTT(MQTTFrame command) throws IOException {
055        ByteSequence bytes = wireFormat.marshal(command);
056        try {
057            //timeout after a period of time so we don't wait forever and hold the protocol lock
058            session.getRemote().sendBytesByFuture(
059                    ByteBuffer.wrap(bytes.getData(), 0, bytes.getLength())).get(getDefaultSendTimeOut(), TimeUnit.SECONDS);
060        } catch (Exception e) {
061            throw IOExceptionSupport.create(e);
062        }
063    }
064
065    @Override
066    public void handleStopped() throws IOException {
067        if (session != null && session.isOpen()) {
068            session.close();
069        }
070    }
071
072    //----- WebSocket.OnTextMessage callback handlers ------------------------//
073
074    @Override
075    public void onWebSocketBinary(byte[] bytes, int offset, int length) {
076        if (!transportStartedAtLeastOnce()) {
077            LOG.debug("Waiting for MQTTSocket to be properly started...");
078            try {
079                socketTransportStarted.await();
080            } catch (InterruptedException e) {
081                LOG.warn("While waiting for MQTTSocket to be properly started, we got interrupted!! Should be okay, but you could see race conditions...");
082            }
083        }
084
085        protocolLock.lock();
086        try {
087            receiveCounter += length;
088            codec.parse(new DataByteArrayInputStream(new Buffer(bytes, offset, length)), length);
089        } catch (Exception e) {
090            onException(IOExceptionSupport.create(e));
091        } finally {
092            protocolLock.unlock();
093        }
094    }
095
096    @Override
097    public void onWebSocketClose(int arg0, String arg1) {
098        try {
099            if (protocolLock.tryLock() || protocolLock.tryLock(ORDERLY_CLOSE_TIMEOUT, TimeUnit.SECONDS)) {
100                LOG.debug("MQTT WebSocket closed: code[{}] message[{}]", arg0, arg1);
101                //Check if we received a disconnect packet before closing
102                if (!receivedDisconnect.get()) {
103                    getProtocolConverter().onTransportError();
104                }
105                getProtocolConverter().onMQTTCommand(new DISCONNECT().encode());
106            }
107        } catch (Exception e) {
108            LOG.debug("Failed to close MQTT WebSocket cleanly", e);
109        } finally {
110            if (protocolLock.isHeldByCurrentThread()) {
111                protocolLock.unlock();
112            }
113        }
114    }
115
116    @Override
117    public void onWebSocketConnect(Session session) {
118        this.session = session;
119    }
120
121    @Override
122    public void onWebSocketError(Throwable arg0) {
123
124    }
125
126    @Override
127    public void onWebSocketText(String arg0) {
128    }
129
130    private static int getDefaultSendTimeOut() {
131        return Integer.getInteger("org.apache.activemq.transport.ws.MQTTSocket.sendTimeout", 30);
132    }
133
134    //----- MQTTCodec Frame Sink event point ---------------------------------//
135
136    @Override
137    public void onFrame(MQTTFrame mqttFrame) {
138        try {
139            if (mqttFrame.messageType() == DISCONNECT.TYPE) {
140                receivedDisconnect.set(true);
141            }
142            getProtocolConverter().onMQTTCommand(mqttFrame);
143        } catch (Exception e) {
144            onException(IOExceptionSupport.create(e));
145        }
146    }
147}