Mirrors/sshj
1
0
Fork 0
mirror of https://github.com/hierynomus/sshj.git synced 2025-12-06 07:10:53 +03:00
Code Issues Actions Packages Projects Releases Wiki Activity

Merge pull request #208 from lguerin/bandwidth

Browse Source 
SCP : limit the used bandwidth
This commit is contained in:
Jeroen van Erp
2015-08-18 17:23:55 +02:00
parent c98ad22a7a d18e9d9961
commit 74a4012023
7 changed files with 229 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
build.gradle
View File

@@ -21,6 +21,9 @@ configurations {
}
test {
testLogging {
exceptionFormat = 'full'
}
include "**/*Test.*"
if (!project.hasProperty("allTests")) {
useJUnit {

16
src/main/java/net/schmizz/sshj/xfer/scp/AbstractSCPClient.java Normal file
View File

@@ -0,0 +1,16 @@
package net.schmizz.sshj.xfer.scp;
abstract class AbstractSCPClient {
protected final SCPEngine engine;
protected int bandwidthLimit;
AbstractSCPClient(SCPEngine engine) {
this.engine = engine;
}
AbstractSCPClient(SCPEngine engine, int bandwidthLimit) {
this.engine = engine;
this.bandwidthLimit = bandwidthLimit;
}
}

26
src/main/java/net/schmizz/sshj/xfer/scp/SCPDownloadClient.java
View File

@@ -24,18 +24,22 @@ import java.io.IOException;
import java.io.OutputStream;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.LinkedList;
import java.util.List;
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArgument;
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArguments;
/** Support for uploading files over a connected link using SCP. */
public final class SCPDownloadClient {
public final class SCPDownloadClient extends AbstractSCPClient {
private boolean recursiveMode = true;
private final SCPEngine engine;
SCPDownloadClient(SCPEngine engine) {
this.engine = engine;
super(engine);
}
SCPDownloadClient(SCPEngine engine, int bandwidthLimit) {
super(engine, bandwidthLimit);
}
/** Download a file from {@code sourcePath} on the connected host to {@code targetPath} locally. */
@@ -60,12 +64,12 @@ public final class SCPDownloadClient {
void startCopy(String sourcePath, LocalDestFile targetFile)
throws IOException {
List<Arg> args = new LinkedList<Arg>();
args.add(Arg.SOURCE);
args.add(Arg.QUIET);
args.add(Arg.PRESERVE_TIMES);
if (recursiveMode)
args.add(Arg.RECURSIVE);
List<SCPArgument> args = SCPArguments.with(Arg.SOURCE)
.and(Arg.QUIET)
.and(Arg.PRESERVE_TIMES)
.and(Arg.RECURSIVE, recursiveMode)
.and(Arg.LIMIT, String.valueOf(bandwidthLimit), (bandwidthLimit > 0))
.arguments();
engine.execSCPWith(args, sourcePath);
engine.signal("Start status OK");

88
src/main/java/net/schmizz/sshj/xfer/scp/SCPEngine.java
View File

@@ -28,6 +28,7 @@ import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.io.OutputStream;
import java.util.LinkedList;
import java.util.List;
/** @see <a href="http://blogs.sun.com/janp/entry/how_the_scp_protocol_works">SCP Protocol</a> */
@@ -39,7 +40,8 @@ class SCPEngine {
RECURSIVE('r'),
VERBOSE('v'),
PRESERVE_TIMES('p'),
QUIET('q');
QUIET('q'),
LIMIT('l');
private final char a;
@@ -97,10 +99,10 @@ class SCPEngine {
exitStatus = -1;
}
void execSCPWith(List<Arg> args, String path)
void execSCPWith(List<SCPArgument> args, String path)
throws SSHException {
final StringBuilder cmd = new StringBuilder(SCP_COMMAND);
for (Arg arg : args) {
for (SCPArgument arg : args) {
cmd.append(" ").append(arg);
}
cmd.append(" ");
@@ -186,4 +188,84 @@ class SCPEngine {
return listener;
}
public static class SCPArgument {
private Arg name;
private String value;
private SCPArgument(Arg name, String value) {
this.name = name;
this.value = value;
}
public static SCPArgument addArgument(Arg name, String value) {
return new SCPArgument(name, value);
}
@Override
public String toString() {
String option = name.toString();
if (value != null) {
option = option + value;
}
return option;
}
}
public static class SCPArguments {
private static List<SCPArgument> args = null;
private SCPArguments() {
this.args = new LinkedList<SCPArgument>();
}
private static void addArgument(Arg name, String value, boolean accept) {
if (accept) {
args.add(SCPArgument.addArgument(name, value));
}
}
public static SCPArguments with(Arg name) {
return with(name, null, true);
}
public static SCPArguments with(Arg name, String value) {
return with(name, value, true);
}
public static SCPArguments with(Arg name, boolean accept) {
return with(name, null, accept);
}
public static SCPArguments with(Arg name, String value, boolean accept) {
SCPArguments scpArguments = new SCPArguments();
addArgument(name, value, accept);
return scpArguments;
}
public SCPArguments and(Arg name) {
addArgument(name, null, true);
return this;
}
public SCPArguments and(Arg name, String value) {
addArgument(name, value, true);
return this;
}
public SCPArguments and(Arg name, boolean accept) {
addArgument(name, null, accept);
return this;
}
public SCPArguments and(Arg name, String value, boolean accept) {
addArgument(name, value, accept);
return this;
}
public List<SCPArgument> arguments() {
return args;
}
}
}

15
src/main/java/net/schmizz/sshj/xfer/scp/SCPFileTransfer.java
View File

@@ -28,18 +28,23 @@ public class SCPFileTransfer
extends AbstractFileTransfer
implements FileTransfer {
/** Default bandwidth limit for SCP transfer in kilobit/s (-1 means unlimited) */
private static final int DEFAULT_BANDWIDTH_LIMIT = -1;
private final SessionFactory sessionFactory;
private int bandwidthLimit;
public SCPFileTransfer(SessionFactory sessionFactory) {
this.sessionFactory = sessionFactory;
this.bandwidthLimit = DEFAULT_BANDWIDTH_LIMIT;
}
public SCPDownloadClient newSCPDownloadClient() {
return new SCPDownloadClient(newSCPEngine());
return new SCPDownloadClient(newSCPEngine(), bandwidthLimit);
}
public SCPUploadClient newSCPUploadClient() {
return new SCPUploadClient(newSCPEngine());
return new SCPUploadClient(newSCPEngine(), bandwidthLimit);
}
private SCPEngine newSCPEngine() {
@@ -70,4 +75,10 @@ public class SCPFileTransfer
newSCPUploadClient().copy(localFile, remotePath);
}
public SCPFileTransfer bandwidthLimit(int limit) {
if (limit > 0) {
this.bandwidthLimit = limit;
}
return this;
}
}

25
src/main/java/net/schmizz/sshj/xfer/scp/SCPUploadClient.java
View File

@@ -24,17 +24,22 @@ import net.schmizz.sshj.xfer.scp.SCPEngine.Arg;
import java.io.IOException;
import java.io.InputStream;
import java.util.LinkedList;
import java.util.List;
/** Support for uploading files over a connected link using SCP. */
public final class SCPUploadClient {
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArgument;
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArguments;
/** Support for uploading files over a connected link using SCP. */
public final class SCPUploadClient extends AbstractSCPClient {
private final SCPEngine engine;
private LocalFileFilter uploadFilter;
SCPUploadClient(SCPEngine engine) {
this.engine = engine;
super(engine);
}
SCPUploadClient(SCPEngine engine, int bandwidthLimit) {
super(engine, bandwidthLimit);
}
/** Upload a local file from {@code localFile} to {@code targetPath} on the remote host. */
@@ -55,11 +60,11 @@ public final class SCPUploadClient {
private synchronized void startCopy(LocalSourceFile sourceFile, String targetPath)
throws IOException {
List<Arg> args = new LinkedList<Arg>();
args.add(Arg.SINK);
args.add(Arg.RECURSIVE);
if (sourceFile.providesAtimeMtime())
args.add(Arg.PRESERVE_TIMES);
List<SCPArgument> args = SCPArguments.with(Arg.SINK)
.and(Arg.RECURSIVE)
.and(Arg.PRESERVE_TIMES, sourceFile.providesAtimeMtime())
.and(Arg.LIMIT, String.valueOf(bandwidthLimit), (bandwidthLimit > 0))
.arguments();
engine.execSCPWith(args, targetPath);
engine.check("Start status OK");
process(engine.getTransferListener(), sourceFile);

82
src/test/java/net/schmizz/sshj/xfer/scp/SCPFileTransferTest.java Normal file
View File

@@ -0,0 +1,82 @@
package net.schmizz.sshj.xfer.scp;
import com.hierynomus.sshj.test.SshFixture;
import com.hierynomus.sshj.test.util.FileUtil;
import net.schmizz.sshj.SSHClient;
import org.junit.After;
import org.junit.Before;
import org.junit.Rule;
import org.junit.Test;
import org.junit.rules.TemporaryFolder;
import java.io.File;
import java.io.IOException;
import static junit.framework.Assert.assertFalse;
import static junit.framework.Assert.assertTrue;
public class SCPFileTransferTest {
public static final String DEFAULT_FILE_NAME = "my_file.txt";
File targetDir;
File sourceFile;
File targetFile;
SSHClient sshClient;
@Rule
public SshFixture fixture = new SshFixture();
@Rule
public TemporaryFolder tempFolder = new TemporaryFolder();
@Before
public void init() throws IOException {
sourceFile = tempFolder.newFile(DEFAULT_FILE_NAME);
FileUtil.writeToFile(sourceFile, "This is my file");
targetDir = tempFolder.newFolder();
targetFile = new File(targetDir + File.separator + DEFAULT_FILE_NAME);
sshClient = fixture.setupConnectedDefaultClient();
sshClient.authPassword("test", "test");
}
@After
public void cleanup() {
if (targetFile.exists()) {
targetFile.delete();
}
}
@Test
public void shouldSCPUploadFile() throws IOException {
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer();
assertFalse(targetFile.exists());
scpFileTransfer.upload(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
assertTrue(targetFile.exists());
}
@Test
public void shouldSCPUploadFileWithBandwidthLimit() throws IOException {
// Limit upload transfer at 2Mo/s
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer().bandwidthLimit(16000);
assertFalse(targetFile.exists());
scpFileTransfer.upload(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
assertTrue(targetFile.exists());
}
@Test
public void shouldSCPDownloadFile() throws IOException {
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer();
assertFalse(targetFile.exists());
scpFileTransfer.download(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
assertTrue(targetFile.exists());
}
@Test
public void shouldSCPDownloadFileWithBandwidthLimit() throws IOException {
// Limit download transfer at 128Ko/s
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer().bandwidthLimit(1024);
assertFalse(targetFile.exists());
scpFileTransfer.download(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
assertTrue(targetFile.exists());
}
}