package io.fathom.cloud.blobs.sftp;
import io.fathom.cloud.blobs.BlobData;
import io.fathom.cloud.blobs.BlobStore;
import io.fathom.cloud.blobs.BlobStoreBase;
import io.fathom.cloud.blobs.BlobStoreFactory;
import io.fathom.cloud.blobs.LocalFilesystemBlobStore;
import io.fathom.cloud.blobs.TempFile;
import io.fathom.cloud.sftp.RemoteFile;
import io.fathom.cloud.sftp.Sftp;
import io.fathom.cloud.ssh.SftpChannel;
import io.fathom.cloud.ssh.SshConfig;
import java.io.File;
import java.io.FileNotFoundException;
import java.io.FileOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.Arrays;
import java.util.List;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
import com.fathomdb.io.IoUtils;
import com.fathomdb.utils.Hex;
import com.google.common.base.Preconditions;
import com.google.common.base.Strings;
import com.google.common.collect.Lists;
import com.google.common.hash.HashCode;
import com.google.common.hash.Hasher;
import com.google.common.hash.Hashing;
import com.google.protobuf.ByteString;
public class SftpBlobStore extends BlobStoreBase {
private static final Logger log = LoggerFactory.getLogger(SftpBlobStore.class);
public static class Factory implements BlobStoreFactory {
final SshConfig sshConfig;
final RemoteFile remoteBaseDir;
final File localCacheDir;
public Factory(SshConfig sshConfig, RemoteFile remoteBaseDir, File localCacheDir) {
this.sshConfig = sshConfig;
this.remoteBaseDir = remoteBaseDir;
this.localCacheDir = localCacheDir;
}
@Override
public BlobStore get(String key) throws IOException {
Preconditions.checkArgument(!Strings.isNullOrEmpty(key));
Preconditions.checkArgument(LocalFilesystemBlobStore.isSafeFileName(key));
return new SftpBlobStore(sshConfig, new RemoteFile(remoteBaseDir, key), new File(localCacheDir, key));
}
}
final RemoteFile remoteBaseDir;
final RemoteFile tmpDir;
final SshConfig sshConfig;
private final File localCacheDir;
private final File localCacheDirTmp;
// TODO: Use pool? Move to SshChannel??
// private final SimplePool<SftpChannel> POOL = new
// SimplePool<SftpChannel>();
//
// private SftpChannel borrowSftp() throws CloudException {
// SftpChannel sftp = POOL.poll();
// if (sftp == null) {
// sftp = buildSftp();
// }
// return sftp;
// }
public SftpBlobStore(SshConfig sshConfig, RemoteFile remoteBaseDir, File localCacheDir) throws IOException {
super();
this.sshConfig = sshConfig;
this.remoteBaseDir = remoteBaseDir;
this.localCacheDir = localCacheDir;
this.tmpDir = new RemoteFile(remoteBaseDir, "_tmp");
mkdirs(tmpDir);
localCacheDirTmp = new File(localCacheDir, "_tmp");
IoUtils.mkdirs(localCacheDirTmp);
}
void mkdirs(RemoteFile dir) throws IOException {
try (Sftp sftp = buildSftp()) {
sftp.mkdirs(dir.getSshPath());
}
}
protected Sftp buildSftp() throws IOException {
SftpChannel sftp = sshConfig.getSftpChannel();
return new Sftp(sftp, tmpDir);
}
@Override
public BlobData find(ByteString key) throws IOException {
RemoteFile remoteFile = buildRemoteFile(key);
File localCache;
try {
localCache = buildCacheFile(key, true);
} catch (IOException e) {
throw new IOException("Error creating local cache file", e);
}
if (localCache.exists()) {
// TODO: Verify?
log.debug("Using local cache file: {}", localCache);
} else {
try (TempFile localTemp = TempFile.in(this.localCacheDirTmp)) {
HashCode md5;
try (Sftp sftp = buildSftp(); InputStream is = sftp.open(remoteFile.getSshPath())) {
try (OutputStream os = new FileOutputStream(localTemp.getFile())) {
Hasher hasher = Hashing.md5().newHasher();
byte[] buffer = new byte[8192];
while (true) {
int n = is.read(buffer);
if (n == -1) {
break;
}
hasher.putBytes(buffer, 0, n);
os.write(buffer, 0, n);
}
md5 = hasher.hash();
}
} catch (FileNotFoundException e) {
// Unable to find sftp file
return null;
}
if (!matches(key, md5)) {
throw new IOException("Copied file, but did not match hash");
}
if (!localTemp.renameTo(localCache)) {
// TODO: Maybe someone else concurrently downloaded it??
throw new IOException("Could not rename file to cache file");
}
}
}
if (localCache.exists()) {
BlobData is = new BlobData(localCache, key);
return is;
} else {
return null;
}
}
@Override
public boolean has(ByteString key, boolean checkCache) throws IOException {
if (checkCache) {
File localCache = buildCacheFile(key);
if (localCache.exists()) {
return true;
}
}
RemoteFile remoteFile = buildRemoteFile(key);
try (Sftp sftp = buildSftp()) {
return sftp.exists(remoteFile.getSshPath());
}
}
private static boolean matches(ByteString a, HashCode b) {
return Arrays.equals(a.toByteArray(), b.asBytes());
}
private RemoteFile buildRemoteFile(Sftp sftp, ByteString key, boolean mkdirs) throws IOException {
RemoteFile file = buildRemoteFile(key);
if (mkdirs) {
// TODO: Cache created / not created state?
sftp.mkdirs(file.getParentFile().getSshPath());
}
return file;
}
private RemoteFile buildRemoteFile(ByteString key) {
File file = buildFile(this.remoteBaseDir.getSshPath(), key);
RemoteFile remoteFile = new RemoteFile(file);
return remoteFile;
}
private File buildCacheFile(ByteString key, boolean mkdirs) throws IOException {
File file = buildCacheFile(key);
if (mkdirs) {
// TODO: Cache created / not created state?
IoUtils.mkdirs(file.getParentFile());
}
return file;
}
private File buildCacheFile(ByteString key) {
File file = buildFile(this.localCacheDir, key);
return file;
}
@Override
public void put(BlobData data) throws IOException {
ByteString key = data.getHash();
File localCacheFile;
try {
localCacheFile = buildCacheFile(key, true);
} catch (IOException e) {
throw new IOException("Error creating local cache file", e);
}
// TODO: Can we be more efficient if BlobData is itself a temp file?
try (TempFile localTemp = TempFile.in(localCacheDirTmp)) {
localTemp.copyFrom(data);
if (!localTemp.renameTo(localCacheFile)) {
// TODO: Maybe someone else concurrently downloaded it??
throw new IOException("Could not rename file to cache file");
}
}
try (Sftp sftp = buildSftp()) {
RemoteFile remoteFile = buildRemoteFile(sftp, key, true);
if (sftp.exists(remoteFile.getSshPath())) {
// TODO: Verify that the file contents match??
log.warn("TODO: Validate that remote files match for sftp");
return;
}
sftp.writeAtomic(remoteFile, localCacheFile);
}
}
@Override
public Iterable<ByteString> listWithPrefix(String prefix) throws IOException {
List<ByteString> ret = Lists.newArrayList();
try (Sftp sftp = buildSftp()) {
for (String name : sftp.ls(remoteBaseDir.getSshPath())) {
if (name.charAt(0) == '.') {
continue;
}
if (name.charAt(0) == '_') {
continue;
}
if (prefix.length() > name.length()) {
if (!prefix.startsWith(name)) {
continue;
}
} else {
if (!name.startsWith(prefix)) {
continue;
}
}
listWithPrefix(sftp, ret, new RemoteFile(remoteBaseDir, name), prefix);
}
}
return ret;
}
private void listWithPrefix(SftpChannel sftp, List<ByteString> ret, RemoteFile dir, String prefix)
throws IOException {
for (String name : sftp.ls(dir.getSshPath())) {
if (name.charAt(0) == '.') {
continue;
}
if (!name.startsWith(prefix)) {
continue;
}
try {
byte[] bytes = Hex.fromHex(name);
ret.add(ByteString.copyFrom(bytes));
} catch (IllegalArgumentException e) {
log.warn("Unable to parse filename: " + name, e);
continue;
}
}
}
@Override
public String toString() {
return "SftpBlobStore [" + sshConfig + ":" + remoteBaseDir + "]";
}
}