tikhomirov@685: /*
tikhomirov@685:  * Copyright (c) 2013 TMate Software Ltd
tikhomirov@685:  *  
tikhomirov@685:  * This program is free software; you can redistribute it and/or modify
tikhomirov@685:  * it under the terms of the GNU General Public License as published by
tikhomirov@685:  * the Free Software Foundation; version 2 of the License.
tikhomirov@685:  *
tikhomirov@685:  * This program is distributed in the hope that it will be useful,
tikhomirov@685:  * but WITHOUT ANY WARRANTY; without even the implied warranty of
tikhomirov@685:  * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
tikhomirov@685:  * GNU General Public License for more details.
tikhomirov@685:  *
tikhomirov@685:  * For information on how to redistribute this software under
tikhomirov@685:  * the terms of a license other than GNU General Public License
tikhomirov@685:  * contact TMate Software at support@hg4j.com
tikhomirov@685:  */
tikhomirov@685: package org.tmatesoft.hg.internal.remote;
tikhomirov@685: 
tikhomirov@685: import java.io.BufferedReader;
tikhomirov@697: import java.io.ByteArrayInputStream;
tikhomirov@685: import java.io.Closeable;
tikhomirov@685: import java.io.EOFException;
tikhomirov@685: import java.io.File;
tikhomirov@685: import java.io.FilterInputStream;
tikhomirov@687: import java.io.FilterOutputStream;
tikhomirov@685: import java.io.IOException;
tikhomirov@685: import java.io.InputStream;
tikhomirov@685: import java.io.InputStreamReader;
tikhomirov@685: import java.io.OutputStream;
tikhomirov@697: import java.io.SequenceInputStream;
tikhomirov@685: import java.net.URL;
tikhomirov@685: import java.util.ArrayList;
tikhomirov@685: import java.util.Collection;
tikhomirov@685: import java.util.Collections;
tikhomirov@685: import java.util.List;
tikhomirov@685: 
tikhomirov@685: import org.tmatesoft.hg.core.HgRemoteConnectionException;
tikhomirov@685: import org.tmatesoft.hg.core.Nodeid;
tikhomirov@685: import org.tmatesoft.hg.core.SessionContext;
tikhomirov@685: import org.tmatesoft.hg.repo.HgRemoteRepository.Range;
tikhomirov@685: import org.tmatesoft.hg.repo.HgRuntimeException;
tikhomirov@685: import org.tmatesoft.hg.util.LogFacility.Severity;
tikhomirov@685: 
tikhomirov@696: import com.trilead.ssh2.Connection;
tikhomirov@696: import com.trilead.ssh2.ConnectionInfo;
tikhomirov@696: import com.trilead.ssh2.Session;
tikhomirov@696: import com.trilead.ssh2.StreamGobbler;
tikhomirov@685: 
tikhomirov@685: /**
tikhomirov@685:  * Remote repository via SSH
tikhomirov@685:  * 
tikhomirov@685:  * @author Artem Tikhomirov
tikhomirov@685:  * @author TMate Software Ltd.
tikhomirov@685:  */
tikhomirov@687: public class SshConnector implements Connector {
tikhomirov@685: 	private SessionContext sessionCtx;
tikhomirov@685: 	private URL url;
tikhomirov@685: 	private Connection conn;
tikhomirov@685: 	private Session session;
tikhomirov@685: 	private int sessionUse;
tikhomirov@685: 
tikhomirov@685: 	private StreamGobbler remoteErr, remoteOut;
tikhomirov@685: 	private OutputStream remoteIn;
tikhomirov@685: 	
tikhomirov@687: 	public void init(URL url, SessionContext sessionContext, Object globalConfig) throws HgRuntimeException {
tikhomirov@685: 		sessionCtx = sessionContext;
tikhomirov@685: 		this.url = url;
tikhomirov@687: 	}
tikhomirov@687: 	
tikhomirov@687: 	public void connect() throws HgRemoteConnectionException, HgRuntimeException {
tikhomirov@685: 		try {
tikhomirov@685: 			conn = new Connection(url.getHost(), url.getPort() == -1 ? 22 : url.getPort());
tikhomirov@685: 			conn.connect();
tikhomirov@685: 		} catch (IOException ex) {
tikhomirov@685: 			throw new HgRemoteConnectionException("Failed to establish connection");
tikhomirov@685: 		}
tikhomirov@685: 		try {
tikhomirov@685: 			conn.authenticateWithPublicKey(System.getProperty("user.name"), new File(System.getProperty("user.home"), ".ssh/id_rsa"), null);
tikhomirov@685: 			ConnectionInfo ci = conn.getConnectionInfo();
tikhomirov@685: 			System.out.printf("%s %s %s %d %s %s %s\n", ci.clientToServerCryptoAlgorithm, ci.clientToServerMACAlgorithm, ci.keyExchangeAlgorithm, ci.keyExchangeCounter, ci.serverHostKeyAlgorithm, ci.serverToClientCryptoAlgorithm, ci.serverToClientMACAlgorithm);
tikhomirov@685: 		} catch (IOException ex) {
tikhomirov@687: 			throw new HgRemoteConnectionException("Failed to authenticate", ex).setServerInfo(getServerLocation());
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public void disconnect() throws HgRemoteConnectionException {
tikhomirov@685: 		if (session != null) {
tikhomirov@685: 			forceSessionClose();
tikhomirov@685: 		}
tikhomirov@685: 		if (conn != null) {
tikhomirov@685: 			conn.close();
tikhomirov@685: 			conn = null;
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public void sessionBegin() throws HgRemoteConnectionException {
tikhomirov@685: 		if (sessionUse > 0) {
tikhomirov@685: 			assert session != null;
tikhomirov@685: 			sessionUse++;
tikhomirov@685: 			return;
tikhomirov@685: 		}
tikhomirov@685: 		try {
tikhomirov@685: 			session = conn.openSession();
tikhomirov@697: 			final String path = url.getPath();
tikhomirov@697: 			session.execCommand(String.format("hg -R %s serve --stdio", path.charAt(0) == '/' ? path.substring(1) : path));
tikhomirov@685: 			remoteErr = new StreamGobbler(session.getStderr());
tikhomirov@685: 			remoteOut = new StreamGobbler(session.getStdout());
tikhomirov@697: 			remoteIn = session.getStdin();
tikhomirov@685: 			sessionUse = 1;
tikhomirov@685: 		} catch (IOException ex) {
tikhomirov@685: 			throw new HgRemoteConnectionException("Failed to create ssh session", ex);
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public void sessionEnd() throws HgRemoteConnectionException {
tikhomirov@685: 		assert sessionUse > 0;
tikhomirov@685: 		assert session != null;
tikhomirov@685: 		if (sessionUse > 1) {
tikhomirov@685: 			sessionUse--;
tikhomirov@685: 			return;
tikhomirov@685: 		}
tikhomirov@685: 		forceSessionClose();
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@687: 	public String getServerLocation() {
tikhomirov@697: 		return url.toString(); // FIXME
tikhomirov@687: 	}
tikhomirov@687: 	
tikhomirov@687: 	public String getCapabilities() throws HgRemoteConnectionException {
tikhomirov@687: 		try {
tikhomirov@687: 			consume(remoteOut);
tikhomirov@687: 			consume(remoteErr);
tikhomirov@687: 			remoteIn.write(CMD_HELLO.getBytes());
tikhomirov@687: 			remoteIn.write('\n');
tikhomirov@687: 			remoteIn.write(CMD_CAPABILITIES.getBytes()); // see http connector for details
tikhomirov@687: 			remoteIn.write('\n');
tikhomirov@687: 			remoteIn.write(CMD_HEADS.getBytes());
tikhomirov@687: 			remoteIn.write('\n');
tikhomirov@687: 			checkError();
tikhomirov@687: 			int responseLen = readResponseLength();
tikhomirov@687: 			checkError();
tikhomirov@687: 			FilterStream s = new FilterStream(remoteOut, responseLen);
tikhomirov@687: 			BufferedReader r = new BufferedReader(new InputStreamReader(s));
tikhomirov@687: 			String line;
tikhomirov@687: 			while ((line = r.readLine()) != null) {
tikhomirov@687: 				if (line.startsWith(CMD_CAPABILITIES) && line.length() > (CMD_CAPABILITIES.length()+1)) {
tikhomirov@687: 					line = line.substring(CMD_CAPABILITIES.length());
tikhomirov@687: 					if (line.charAt(0) == ':') {
tikhomirov@687: 						return line.substring(CMD_CAPABILITIES.length() + 1);
tikhomirov@687: 					}
tikhomirov@687: 				}
tikhomirov@687: 			}
tikhomirov@687: 			r.close();
tikhomirov@687: 			consume(remoteOut);
tikhomirov@687: 			checkError();
tikhomirov@687: 			return new String();
tikhomirov@687: 		} catch (IOException ex) {
tikhomirov@687: 			throw new HgRemoteConnectionException("Failed to initiate dialog with server", ex).setRemoteCommand(CMD_HELLO).setServerInfo(getServerLocation());
tikhomirov@687: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	public InputStream heads() throws HgRemoteConnectionException {
tikhomirov@697: 		return executeCommand("heads", Collections.<Parameter>emptyList(), true);
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public InputStream between(Collection<Range> ranges) throws HgRemoteConnectionException {
tikhomirov@685: 		StringBuilder sb = new StringBuilder(ranges.size() * 82);
tikhomirov@685: 		for (Range r : ranges) {
tikhomirov@685: 			r.append(sb).append(' ');
tikhomirov@685: 		}
tikhomirov@685: 		if (!ranges.isEmpty()) {
tikhomirov@685: 			sb.setLength(sb.length() - 1);
tikhomirov@685: 		}
tikhomirov@697: 		return executeCommand("between", Collections.singletonList(new Parameter("pairs", sb.toString())), true);
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public InputStream branches(List<Nodeid> nodes) throws HgRemoteConnectionException {
tikhomirov@685: 		String l = join(nodes, ' ');
tikhomirov@697: 		return executeCommand("branches", Collections.singletonList(new Parameter("nodes", l)), true);
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public InputStream changegroup(List<Nodeid> roots) throws HgRemoteConnectionException, HgRuntimeException {
tikhomirov@685: 		String l = join(roots, ' ');
tikhomirov@697: 		InputStream cg = executeCommand("changegroup", Collections.singletonList(new Parameter("roots", l)), false);
tikhomirov@697: 		InputStream prefix = new ByteArrayInputStream("HG10UN".getBytes());
tikhomirov@697: 		return new SequenceInputStream(prefix, cg);
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@687: 	public OutputStream unbundle(long outputLen, List<Nodeid> remoteHeads) throws HgRemoteConnectionException, HgRuntimeException {
tikhomirov@685: 		String l = join(remoteHeads, ' ');
tikhomirov@687: 		try {
tikhomirov@687: 			consume(remoteOut);
tikhomirov@687: 			consume(remoteErr);
tikhomirov@687: 			remoteIn.write(CMD_UNBUNDLE.getBytes());
tikhomirov@687: 			remoteIn.write('\n');
tikhomirov@687: 			writeParameters(Collections.singletonList(new Parameter("heads", l)));
tikhomirov@687: 			checkError();
tikhomirov@687: 			return new FilterOutputStream(remoteIn) {
tikhomirov@687: 				@Override
tikhomirov@687: 				public void close() throws IOException {
tikhomirov@687: 					out.flush();
tikhomirov@687: 					@SuppressWarnings("unused")
tikhomirov@687: 					int responseLen = readResponseLength();
tikhomirov@687: 					checkError();
tikhomirov@687: 					// XXX perhaps, need to return responseLen to caller? 
tikhomirov@687: 				}
tikhomirov@687: 			};
tikhomirov@687: 		} catch (IOException ex) {
tikhomirov@687: 			throw new HgRemoteConnectionException("Communication failure", ex).setRemoteCommand(CMD_UNBUNDLE).setServerInfo(getServerLocation());
tikhomirov@687: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	public InputStream pushkey(String opName, String namespace, String key, String oldValue, String newValue) throws HgRemoteConnectionException, HgRuntimeException {
tikhomirov@685: 		ArrayList<Parameter> p = new ArrayList<Parameter>();
tikhomirov@685: 		p.add(new Parameter("namespace", namespace));
tikhomirov@685: 		p.add(new Parameter("key", key));
tikhomirov@685: 		p.add(new Parameter("old", oldValue));
tikhomirov@685: 		p.add(new Parameter("new", newValue));
tikhomirov@697: 		return executeCommand("pushkey", p, true);
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	public InputStream listkeys(String namespace, String actionName) throws HgRemoteConnectionException, HgRuntimeException {
tikhomirov@697: 		return executeCommand("listkeys", Collections.singletonList(new Parameter("namespace", namespace)), true);
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@697: 	private InputStream executeCommand(String cmd, List<Parameter> parameters, boolean expectResponseLength) throws HgRemoteConnectionException {
tikhomirov@685: 		try {
tikhomirov@685: 			consume(remoteOut);
tikhomirov@685: 			consume(remoteErr);
tikhomirov@685: 			remoteIn.write(cmd.getBytes());
tikhomirov@685: 			remoteIn.write('\n');
tikhomirov@687: 			writeParameters(parameters);
tikhomirov@685: 			checkError();
tikhomirov@697: 			if (expectResponseLength) {
tikhomirov@697: 				int responseLen = readResponseLength();
tikhomirov@697: 				checkError();
tikhomirov@697: 				return new FilterStream(remoteOut, responseLen);
tikhomirov@697: 			} else {
tikhomirov@697: 				return new FilterStream(remoteOut, Integer.MAX_VALUE);
tikhomirov@697: 			}
tikhomirov@685: 		} catch (IOException ex) {
tikhomirov@687: 			throw new HgRemoteConnectionException("Communication failure", ex).setRemoteCommand(cmd).setServerInfo(getServerLocation());
tikhomirov@687: 		}
tikhomirov@687: 	}
tikhomirov@687: 	
tikhomirov@687: 	private void writeParameters(List<Parameter> parameters) throws IOException {
tikhomirov@687: 		for (Parameter p : parameters) {
tikhomirov@687: 			remoteIn.write(p.name().getBytes());
tikhomirov@687: 			remoteIn.write(' ');
tikhomirov@687: 			remoteIn.write(String.valueOf(p.size()).getBytes());
tikhomirov@687: 			remoteIn.write('\n');
tikhomirov@687: 			remoteIn.write(p.data());
tikhomirov@687: 			remoteIn.write('\n');
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	private void consume(InputStream is) throws IOException {
tikhomirov@685: 		while (is.available() > 0) {
tikhomirov@685: 			is.read();
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@687: 	private void checkError() throws IOException {
tikhomirov@685: 		if (remoteErr.available() > 0) {
tikhomirov@685: 			StringBuilder sb = new StringBuilder();
tikhomirov@685: 			int c;
tikhomirov@685: 			while ((c = remoteErr.read()) != -1) {
tikhomirov@685: 				sb.append((char)c);
tikhomirov@685: 			}
tikhomirov@687: 			throw new IOException(sb.toString());
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 	
tikhomirov@685: 	private int readResponseLength() throws IOException {
tikhomirov@685: 		int c;
tikhomirov@685: 		StringBuilder sb = new StringBuilder();
tikhomirov@685: 		while ((c = remoteOut.read()) != -1) {
tikhomirov@685: 			if (c == '\n') {
tikhomirov@685: 				break;
tikhomirov@685: 			}
tikhomirov@685: 			sb.append((char) c);
tikhomirov@685: 		}
tikhomirov@685: 		if (c == -1) {
tikhomirov@685: 			throw new EOFException();
tikhomirov@685: 		}
tikhomirov@685: 		try {
tikhomirov@685: 			return Integer.parseInt(sb.toString());
tikhomirov@685: 		} catch (NumberFormatException ex) {
tikhomirov@685: 			throw new IOException(String.format("Expected response length instead of %s", sb));
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 
tikhomirov@685: 	private void forceSessionClose() {
tikhomirov@685: 		if (session != null) {
tikhomirov@685: 			closeQuietly(remoteErr);
tikhomirov@685: 			closeQuietly(remoteOut);
tikhomirov@685: 			remoteErr = remoteOut = null;
tikhomirov@685: 			closeQuietly(remoteIn);
tikhomirov@685: 			remoteIn = null;
tikhomirov@685: 			session.close();
tikhomirov@685: 			session = null;
tikhomirov@685: 		}
tikhomirov@685: 		sessionUse = 0;
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	private void closeQuietly(Closeable c) {
tikhomirov@685: 		try {
tikhomirov@685: 			if (c != null) {
tikhomirov@685: 				c.close();
tikhomirov@685: 			}
tikhomirov@685: 		} catch (IOException ex) {
tikhomirov@685: 			sessionCtx.getLog().dump(getClass(), Severity.Warn, ex, null);
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	private static String join(List<Nodeid> values, char sep) {
tikhomirov@685: 		StringBuilder sb = new StringBuilder(values.size() * 41);
tikhomirov@685: 		for (Nodeid n : values) {
tikhomirov@685: 			sb.append(n.toString());
tikhomirov@685: 			sb.append(sep);
tikhomirov@685: 		}
tikhomirov@685: 		if (!values.isEmpty()) {
tikhomirov@685: 			// strip last space 
tikhomirov@685: 			sb.setLength(sb.length() - 1);
tikhomirov@685: 		}
tikhomirov@685: 		return sb.toString();
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	private static final class Parameter {
tikhomirov@685: 		private final String name;
tikhomirov@685: 		private final byte[] data;
tikhomirov@685: 
tikhomirov@685: 		public Parameter(String paramName, String paramValue) {
tikhomirov@685: 			assert paramName != null;
tikhomirov@685: 			assert paramValue != null;
tikhomirov@685: 			name = paramName;
tikhomirov@685: 			data = paramValue.getBytes();
tikhomirov@685: 		}
tikhomirov@685: 
tikhomirov@685: 		public String name() {
tikhomirov@685: 			return name;
tikhomirov@685: 		}
tikhomirov@685: 		public int size() {
tikhomirov@685: 			return data.length;
tikhomirov@685: 		}
tikhomirov@685: 		public byte[] data() {
tikhomirov@685: 			return data;
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: 
tikhomirov@685: 	private static final class FilterStream extends FilterInputStream {
tikhomirov@685: 		private int length;
tikhomirov@685: 
tikhomirov@685: 		public FilterStream(InputStream is, int initialLength) {
tikhomirov@685: 			super(is);
tikhomirov@685: 			length = initialLength;
tikhomirov@685: 		}
tikhomirov@685: 		
tikhomirov@685: 		@Override
tikhomirov@685: 		public int available() throws IOException {
tikhomirov@685: 			return Math.min(super.available(), length);
tikhomirov@685: 		}
tikhomirov@685: 		@Override
tikhomirov@685: 		public int read() throws IOException {
tikhomirov@685: 			if (length == 0) {
tikhomirov@685: 				return -1;
tikhomirov@685: 			}
tikhomirov@685: 			int r = super.read();
tikhomirov@685: 			if (r >= 0) {
tikhomirov@685: 				length--;
tikhomirov@685: 			}
tikhomirov@685: 			return r;
tikhomirov@685: 		}
tikhomirov@685: 		@Override
tikhomirov@685: 		public int read(byte[] b, int off, int len) throws IOException {
tikhomirov@685: 			if (length == 0) {
tikhomirov@685: 				return -1;
tikhomirov@685: 			}
tikhomirov@685: 			int r = super.read(b, off, Math.min(len, length));
tikhomirov@685: 			if (r >= 0) {
tikhomirov@685: 				assert r <= length;
tikhomirov@685: 				length -= r;
tikhomirov@685: 			}
tikhomirov@685: 			return r;
tikhomirov@685: 		}
tikhomirov@685: 		@Override
tikhomirov@685: 		public void close() throws IOException {
tikhomirov@697: 			length = 0;
tikhomirov@685: 			// INTENTIONALLY DOES NOT CLOSE THE STREAM
tikhomirov@685: 		}
tikhomirov@685: 	}
tikhomirov@685: }