/*
* VNCHook.java
* Copyright (C) 2011,2012 Wannes De Smet
*
* This program is free software: you can redistribute it and/or modify
* it under the terms of the GNU General Public License as published by
* the Free Software Foundation, either version 3 of the License, or
* (at your option) any later version.
*
* This program is distributed in the hope that it will be useful,
* but WITHOUT ANY WARRANTY; without even the implied warranty of
* MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
* GNU General Public License for more details.
*
* You should have received a copy of the GNU General Public License
* along with this program. If not, see <http://www.gnu.org/licenses/>.
*/
package org.xenmaster.web;
import com.google.gson.Gson;
import com.google.gson.JsonElement;
import java.io.IOException;
import java.net.InetSocketAddress;
import java.net.Socket;
import java.net.URI;
import java.net.URISyntaxException;
import java.nio.ByteBuffer;
import java.util.ArrayList;
import java.util.Iterator;
import java.util.Map.Entry;
import java.util.UUID;
import java.util.concurrent.ConcurrentHashMap;
import java.util.concurrent.TimeUnit;
import net.wgr.server.web.handling.WebCommandHandler;
import net.wgr.utility.GlobalExecutorService;
import net.wgr.wcp.Commander;
import net.wgr.wcp.Scope;
import net.wgr.wcp.command.Command;
import net.wgr.wcp.command.CommandException;
import org.apache.commons.codec.binary.Base64;
import org.apache.log4j.Logger;
import org.xenmaster.api.entity.Console;
import org.xenmaster.api.entity.VM;
import org.xenmaster.connectivity.ConnectionMultiplexer;
import org.xenmaster.controller.Controller;
/**
*
* @created Dec 15, 2011
* @author double-u
*/
public class VNCHook extends WebCommandHandler {
protected static ConnectionMultiplexer cm;
protected static ConcurrentHashMap<String, Connection> connections;
protected ConnectionMultiplexer.ActivityListener al;
protected static int connectionSlots, available = -1;
protected static Gson gson;
public VNCHook() {
super("vnc");
if (cm == null) {
gson = new Gson();
setupInfrastructure();
}
}
protected static String buildHttpConnect(URI uri) {
StringBuilder sb = new StringBuilder();
sb.append("CONNECT ");
sb.append(uri.getPath()).append('?').append(uri.getQuery());
sb.append(" HTTP/1.1").append("\r\n");
sb.append("Cookie: session_id=").append(Controller.getSession().getReference());
sb.append("\r\n\r\n");
return sb.toString();
}
@Override
public Object execute(Command cmd) {
try {
switch (cmd.getName()) {
case "openConnection":
if (!cmd.getData().isJsonObject() || !cmd.getData().getAsJsonObject().has("ref")) {
throw new IllegalArgumentException("No VM reference parameter given");
}
Connection conn = new Connection(cmd.getConnection().getId());
VM vm = new VM(cmd.getData().getAsJsonObject().get("ref").getAsString(), false);
for (Console c : vm.getConsoles()) {
if (c.getProtocol() == Console.Protocol.RFB) {
try {
conn.console = c;
URI uri = new URI(c.getLocation());
conn.uri = uri;
InetSocketAddress isa = new InetSocketAddress(uri.getHost(), 80);
conn.waitForAddress = isa;
conn.lastWriteTime = System.currentTimeMillis();
connections.put(conn.getReference(), conn);
cm.addConnection(isa);
break;
} catch (URISyntaxException ex) {
Logger.getLogger(getClass()).error("Failed to parse URI", ex);
} catch (IOException | InterruptedException ex) {
Logger.getLogger(getClass()).error("Failed to create connection", ex);
}
}
}
return conn.getReference();
case "write":
Arguments data = Arguments.fromJson(cmd.getData());
if (!connections.containsKey(data.ref)) {
return new CommandException("Tried to write to unexisting connection", data.ref);
}
Connection c = connections.get(data.ref);
c.lastWriteTime = System.currentTimeMillis();
byte[] bytes = Base64.decodeBase64(data.data);
cm.write(c.connection, ByteBuffer.wrap(bytes));
break;
case "closeConnection":
Arguments close = Arguments.fromJson(cmd.getData());
Connection ci = connections.get(close.ref);
if (ci != null) {
// Indicate that we've initiated the connection close
ci.lastWriteTime = -1;
cm.close(ci.connection);
}
break;
case "connectionHeartbeat":
Arguments heartbeat = Arguments.fromJson(cmd.getData());
if (!connections.containsKey(heartbeat.ref)) {
return new CommandException("Tried to write to unexisting connection", heartbeat.ref);
}
Connection ch = connections.get(heartbeat.ref);
ch.lastHeartbeat = System.currentTimeMillis();
break;
}
} catch (IOException | IllegalArgumentException ex) {
Logger.getLogger(getClass()).error("Command failed : " + cmd.getName(), ex);
}
return null;
}
protected static int allocateConnection() {
if (available != -1) {
int temp = available;
available = -1;
return temp;
}
if (connectionSlots != Integer.MAX_VALUE) {
return ++connectionSlots;
}
for (int i = 0; i < Integer.MAX_VALUE; i++) {
if (!connections.containsKey("ConnectionRef:" + i)) {
return i;
}
}
throw new Error("Maxed out at " + Integer.MAX_VALUE + " active VNC connections. Something's wrong");
}
protected static class Connection {
public UUID clientId;
public int connection;
protected String reference;
public InetSocketAddress waitForAddress;
public long lastWriteTime;
public long lastHeartbeat;
public URI uri;
public boolean dismissedHttpOK;
public Console console;
public Connection(UUID client) {
this.reference = "ConnectionRef:" + allocateConnection();
this.clientId = client;
this.lastHeartbeat = System.currentTimeMillis();
}
public String getReference() {
return reference;
}
}
protected static class Arguments {
public String data;
public String ref;
public Arguments() {
}
public static Arguments fromJson(JsonElement data) {
Arguments na = gson.fromJson(data, Arguments.class);
if (na == null || na.ref == null || !na.ref.startsWith("ConnectionRef:")) {
throw new IllegalArgumentException("Illegal reference, is not a ConnectionRef");
}
return na;
}
public Arguments(String data, String ref) {
this.data = data;
this.ref = ref;
}
}
protected static class AL implements ConnectionMultiplexer.ActivityListener {
@Override
public void dataReceived(ByteBuffer buffer, int connection, ConnectionMultiplexer cm) {
Connection conn = null;
for (Entry<String, Connection> entry : connections.entrySet()) {
if (entry.getValue().connection == connection) {
conn = entry.getValue();
break;
}
}
if (conn == null) {
Logger.getLogger(getClass()).warn("Received data on inactive connection " + connection + ". Closing ...");
try {
cm.close(connection);
} catch (IOException ex) {
Logger.getLogger(getClass()).error("Failed to close connection " + connection, ex);
}
return;
}
byte[] data = buffer.array();
if (!conn.dismissedHttpOK) {
String content = new String(data);
// RFB xxx.xxx denotes start of VNC handshake
if (content.contains("RFB")) {
conn.dismissedHttpOK = true;
data = content.substring(content.indexOf("RFB")).getBytes();
} else {
return;
}
}
if (data.length < 1) {
return;
}
Arguments vncData = new Arguments();
vncData.ref = conn.getReference();
vncData.data = Base64.encodeBase64String(data).replace("\r\n", "");
Command cmd = new Command("vnc", "updateScreen", vncData);
ArrayList<UUID> ids = new ArrayList<>();
ids.add(conn.clientId);
Scope scope = new Scope(ids);
Commander.get().commandeer(cmd, scope);
}
@Override
public void connectionClosed(int connection) {
for (Iterator<Entry<String, Connection>> it = connections.entrySet().iterator(); it.hasNext();) {
Entry<String, Connection> entry = it.next();
if (entry.getValue().connection == connection) {
Connection conn = entry.getValue();
// Check if this disconnect was initiated by a user
if (conn.lastWriteTime == -1) {
Command cmd = new Command("vnc", "connectionClosed", new Arguments("", entry.getKey()));
Commander.get().commandeer(cmd, new Scope(Scope.Target.ALL));
it.remove();
} else {
try {
// Try to reconnect
URI uri = new URI(conn.console.getLocation());
conn.uri = uri;
InetSocketAddress isa = new InetSocketAddress(uri.getHost(), 80);
conn.waitForAddress = isa;
conn.lastWriteTime = System.currentTimeMillis();
cm.addConnection(isa);
} catch (URISyntaxException | IOException | InterruptedException ex) {
Logger.getLogger(getClass()).error("Failed to reinitiate connection", ex);
}
}
break;
}
}
}
@Override
public void connectionEstablished(int connection, Socket socket) {
Connection conn = null;
InetSocketAddress isa = null;
for (Entry<String, Connection> entry : connections.entrySet()) {
isa = entry.getValue().waitForAddress;
if (isa != null && isa.equals(socket.getRemoteSocketAddress())) {
conn = entry.getValue();
conn.waitForAddress = null;
break;
}
}
if (conn == null) {
Logger.getLogger(getClass()).warn("Unknown connection established to " + ((InetSocketAddress) socket.getRemoteSocketAddress()).getHostString());
return;
}
conn.connection = connection;
cm.write(conn.connection, ByteBuffer.wrap(buildHttpConnect(conn.uri).getBytes()));
Command cmd = new Command("vnc", "connectionEstablished", new Arguments("", conn.getReference()));
ArrayList<UUID> ids = new ArrayList<>();
ids.add(conn.clientId);
Scope scope = new Scope(ids);
Commander.get().commandeer(cmd, scope);
}
}
protected static void setupInfrastructure() {
cm = new ConnectionMultiplexer();
cm.addActivityListener(new AL());
cm.start();
connections = new ConcurrentHashMap<>();
GlobalExecutorService.get().scheduleAtFixedRate(new Reaper(), 0, 10, TimeUnit.SECONDS);
}
protected static class Reaper implements Runnable {
@Override
public void run() {
// Send a heartbeat
Command cmd = new Command("vnc", "connectionHeartbeat", new Arguments());
Commander.get().commandeer(cmd, new Scope(Scope.Target.ALL));
for (Iterator<Entry<String, Connection>> it = connections.entrySet().iterator(); it.hasNext();) {
Entry<String, Connection> entry = it.next();
// Skipped 2 hearbeats, is probably dead.
if (System.currentTimeMillis() - entry.getValue().lastHeartbeat > 1000 * 20) {
try {
Logger.getLogger(getClass()).info("Reaper closing inactive connection " + entry.getValue().connection);
cm.close(entry.getValue().connection);
it.remove();
} catch (IOException ex) {
Logger.getLogger(getClass()).error("Failed to close connection", ex);
}
}
}
}
}
}