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;
018
019import java.io.IOException;
020import java.util.concurrent.BlockingQueue;
021import java.util.concurrent.CountDownLatch;
022import java.util.concurrent.LinkedBlockingDeque;
023import java.util.concurrent.TimeUnit;
024
025import org.apache.activemq.transport.stomp.StompFrame;
026import org.apache.activemq.transport.stomp.StompWireFormat;
027import org.eclipse.jetty.websocket.api.Session;
028import org.eclipse.jetty.websocket.api.WebSocketAdapter;
029import org.eclipse.jetty.websocket.api.WebSocketListener;
030import org.slf4j.Logger;
031import org.slf4j.LoggerFactory;
032
033/**
034 * STOMP over WS based Connection class
035 */
036public class StompWSConnection extends WebSocketAdapter implements WebSocketListener {
037
038    private static final Logger LOG = LoggerFactory.getLogger(StompWSConnection.class);
039
040    private Session connection;
041    private final CountDownLatch connectLatch = new CountDownLatch(1);
042
043    private final BlockingQueue<String> prefetch = new LinkedBlockingDeque<String>();
044    private final StompWireFormat wireFormat = new StompWireFormat();
045
046    private int closeCode = -1;
047    private String closeMessage;
048
049    @Override
050    public boolean isConnected() {
051        return connection != null ? connection.isOpen() : false;
052    }
053
054    public void close() {
055        if (connection != null) {
056            connection.close();
057        }
058    }
059
060    protected Session getConnection() {
061        return connection;
062    }
063
064    //---- Send methods ------------------------------------------------------//
065
066    public synchronized void sendRawFrame(String rawFrame) throws Exception {
067        checkConnected();
068        connection.getRemote().sendString(rawFrame);
069    }
070
071    public synchronized void sendFrame(StompFrame frame) throws Exception {
072        checkConnected();
073        connection.getRemote().sendString(wireFormat.marshalToString(frame));
074    }
075
076    public synchronized void keepAlive() throws Exception {
077        checkConnected();
078        connection.getRemote().sendString("\n");
079    }
080
081    //----- Receive methods --------------------------------------------------//
082
083    public String receive() throws Exception {
084        checkConnected();
085        return prefetch.take();
086    }
087
088    public String receive(long timeout, TimeUnit unit) throws Exception {
089        checkConnected();
090        return prefetch.poll(timeout, unit);
091    }
092
093    public String receiveNoWait() throws Exception {
094        checkConnected();
095        return prefetch.poll();
096    }
097
098    //---- Blocking state change calls ---------------------------------------//
099
100    public void awaitConnection() throws InterruptedException {
101        connectLatch.await();
102    }
103
104    public boolean awaitConnection(long time, TimeUnit unit) throws InterruptedException {
105        return connectLatch.await(time, unit);
106    }
107
108    //----- Property Accessors -----------------------------------------------//
109
110    public int getCloseCode() {
111        return closeCode;
112    }
113
114    public String getCloseMessage() {
115        return closeMessage;
116    }
117
118    //----- WebSocket callback handlers --------------------------------------//
119
120    @Override
121    public void onWebSocketText(String data) {
122        if (data == null) {
123            return;
124        }
125
126        if (data.equals("\n")) {
127            LOG.debug("New incoming heartbeat read");
128        } else {
129            LOG.trace("New incoming STOMP Frame read: \n{}", data);
130            prefetch.add(data);
131        }
132    }
133
134    @Override
135    public void onWebSocketClose(int statusCode, String reason) {
136        LOG.trace("STOMP WS Connection closed, code:{} message:{}", statusCode, reason);
137
138        this.connection = null;
139        this.closeCode = statusCode;
140        this.closeMessage = reason;
141    }
142
143    @Override
144    public void onWebSocketConnect(org.eclipse.jetty.websocket.api.Session session) {
145        this.connection = session;
146        this.connectLatch.countDown();
147    }
148
149    //----- Internal implementation ------------------------------------------//
150
151    private void checkConnected() throws IOException {
152        if (!isConnected()) {
153            throw new IOException("STOMP WS Connection is closed.");
154        }
155    }
156}