/*
* Copyright 2015 Julien Viet
*
* Licensed 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 io.termd.core.http.websocket.client;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import io.termd.core.http.websocket.server.TaskStatusUpdateEvent;
import io.termd.core.util.ObjectWrapper;
import io.termd.core.util.Wait;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import javax.websocket.ClientEndpointConfig;
import javax.websocket.CloseReason;
import javax.websocket.ContainerProvider;
import javax.websocket.Endpoint;
import javax.websocket.EndpointConfig;
import javax.websocket.MessageHandler;
import javax.websocket.RemoteEndpoint;
import javax.websocket.Session;
import java.io.IOException;
import java.net.URI;
import java.nio.ByteBuffer;
import java.time.temporal.ChronoUnit;
import java.util.Optional;
import java.util.concurrent.TimeoutException;
import java.util.function.Consumer;
/**
* @author <a href="mailto:matejonnet@gmail.com">Matej Lazar</a>
* @see "https://github.com/undertow-io/undertow/blob/5bdddf327209a4abf18792e78148863686c26e9b/websockets-jsr/src/test/java/io/undertow/websockets/jsr/test/BinaryEndpointTest.java"
*/
public class Client {
public static final String WEB_SOCKET_TERMINAL_PATH = "/socket/term";
public static final String WEB_SOCKET_LISTENER_PATH = "/socket/process-status-updates";
private static final Logger log = LoggerFactory.getLogger(Client.class);
ProgramaticClientEndpoint endpoint = new ProgramaticClientEndpoint();
private Consumer<Session> onOpenConsumer;
private Consumer<String> onStringMessageConsumer;
private Consumer<byte[]> onBinaryMessageConsumer;
private Consumer<CloseReason> onCloseConsumer;
private Consumer<Throwable> onErrorConsumer;
public Endpoint connect(String websocketUrl) throws Exception {
ClientEndpointConfig clientEndpointConfig = ClientEndpointConfig.Builder.create().build();
ContainerProvider.getWebSocketContainer().connectToServer(endpoint, clientEndpointConfig, new URI(websocketUrl));
return endpoint;
}
public void close() throws Exception {
log.debug("Client is closing connection.");
endpoint.session.close();
// endpoint.closeLatch.await(10, TimeUnit.SECONDS);
}
public void onOpen(Consumer<Session> onOpen) {
onOpenConsumer = onOpen;
}
public void onStringMessage(Consumer<String> onStringMessage) {
onStringMessageConsumer = onStringMessage;
}
public void onBinaryMessage(Consumer<byte[]> onBinaryMessage) {
onBinaryMessageConsumer = onBinaryMessage;
}
public void onClose(Consumer<CloseReason> onClose) {
onCloseConsumer = onClose;
}
public void onError(Consumer<Throwable> onError) {
onErrorConsumer = onError;
}
public RemoteEndpoint.Basic getRemoteEndpoint() {
return endpoint.session.getBasicRemote();
}
public class ProgramaticClientEndpoint extends Endpoint {
volatile Session session;
@Override
public void onOpen(Session session, EndpointConfig config) {
log.debug("Client received open.");
this.session = session;
session.addMessageHandler(new MessageHandler.Whole<String>() {
@Override
public void onMessage(String message) {
log.trace("Client received text MESSAGE: {}", message);
if (onStringMessageConsumer != null) {
onStringMessageConsumer.accept(message);
}
}
});
session.addMessageHandler(new MessageHandler.Whole<byte[]>() {
@Override
public void onMessage(byte[] bytes) {
log.trace("Client received binary MESSAGE: {}", new String(bytes));
if (onBinaryMessageConsumer != null) {
onBinaryMessageConsumer.accept(bytes);
}
}
});
if (onOpenConsumer != null) {
onOpenConsumer.accept(session);
}
}
@Override
public void onClose(Session session, CloseReason closeReason) {
log.debug("Client received close.");
onCloseConsumer.accept(closeReason);
}
@Override
public void onError(Session session, Throwable thr) {
if (onErrorConsumer != null) {
onErrorConsumer.accept(thr);
} else {
log.error("No error handler defined. Received error was: ", thr);
}
}
}
public static Client initializeDefault() {
Client client = new Client();
Consumer<Session> onOpen = (session) -> {
log.info("Client connection opened.");
};
Consumer<CloseReason> onClose = (closeReason) -> {
log.info("Client connection closed. " + closeReason);
};
client.onOpen(onOpen);
client.onClose(onClose);
return client;
}
public static Client connectStatusListenerClient(String webSocketUrl, Consumer<TaskStatusUpdateEvent> onStatusUpdate) {
Client client = Client.initializeDefault();
Consumer<String> responseConsumer = (text) -> {
log.trace("Decoding response: {}", text);
ObjectMapper mapper = new ObjectMapper();
JsonNode jsonObject = null;
try {
jsonObject = mapper.readTree(text);
} catch (IOException e) {
log.error("Cannot read JSON string: " + text, e);
}
try {
TaskStatusUpdateEvent taskStatusUpdateEvent = TaskStatusUpdateEvent.fromJson(jsonObject.get("event").toString());
onStatusUpdate.accept(taskStatusUpdateEvent);
} catch (IOException e) {
log.error("Cannot deserialize TaskStatusUpdateEvent.", e);
}
};
client.onStringMessage(responseConsumer);
client.onClose(closeReason -> {
});
try {
client.connect(webSocketUrl + "/");
} catch (Exception e) {
throw new AssertionError("Failed to connect to remote client.", e);
}
return client;
}
public static Client connectCommandExecutingClient(String webSocketUrl, Optional<Consumer<String>> responseDataConsumer) throws InterruptedException, TimeoutException {
ObjectWrapper<Boolean> connected = new ObjectWrapper<>(false);
Client client = Client.initializeDefault();
Consumer<byte[]> responseConsumer = (bytes) -> {
String responseData = new String(bytes);
if ("% ".equals(responseData)) { //TODO use events
connected.set(true);
} else {
responseDataConsumer.ifPresent((rdc) -> rdc.accept(responseData));
;
}
};
client.onBinaryMessage(responseConsumer);
client.onClose(closeReason -> {
});
try {
client.connect(webSocketUrl + "/");
} catch (Exception e) {
throw new AssertionError("Failed to connect to remote client.", e);
}
Wait.forCondition(() -> connected.get(), 5, ChronoUnit.SECONDS, "Client was not connected within given timeout.");
return client;
}
public static void executeRemoteCommand(Client client, String command) {
log.info("Executing remote command ...");
RemoteEndpoint.Basic remoteEndpoint = client.getRemoteEndpoint();
String data = "{\"action\":\"read\",\"data\":\"" + command + "\\r\\n\"}";
try {
remoteEndpoint.sendBinary(ByteBuffer.wrap(data.getBytes()));
} catch (IOException e) {
e.printStackTrace();
}
}
}