Skip to content

Commit 74a4012

Browse files
hierynomushierynomus
committed
Merge pull request #208 from lguerin/bandwidth
SCP : limit the used bandwidth
2 parents c98ad22 + d18e9d9 commit 74a4012

File tree

7 files changed

+228
-25
lines changed

7 files changed

+228
-25
lines changed

‎build.gradle

+3
Original file line numberDiff line numberDiff line change
@@ -21,6 +21,9 @@ configurations {
2121
}
2222

2323
test {
24+
testLogging {
25+
exceptionFormat = 'full'
26+
}
2427
include "**/*Test.*"
2528
if (!project.hasProperty("allTests")) {
2629
useJUnit {

‎src/main/java/net/schmizz/sshj/xfer/scp/AbstractSCPClient.java

+16
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,16 @@
1+
package net.schmizz.sshj.xfer.scp;
2+
3+
abstract class AbstractSCPClient {
4+
5+
protected final SCPEngine engine;
6+
protected int bandwidthLimit;
7+
8+
AbstractSCPClient(SCPEngine engine) {
9+
this.engine = engine;
10+
}
11+
12+
AbstractSCPClient(SCPEngine engine, int bandwidthLimit) {
13+
this.engine = engine;
14+
this.bandwidthLimit = bandwidthLimit;
15+
}
16+
}

‎src/main/java/net/schmizz/sshj/xfer/scp/SCPDownloadClient.java

+15-11
Original file line numberDiff line numberDiff line change
@@ -24,18 +24,22 @@
2424
import java.io.OutputStream;
2525
import java.util.ArrayList;
2626
import java.util.Arrays;
27-
import java.util.LinkedList;
2827
import java.util.List;
2928

29+
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArgument;
30+
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArguments;
31+
3032
/** Support for uploading files over a connected link using SCP. */
31-
public final class SCPDownloadClient {
33+
public final class SCPDownloadClient extends AbstractSCPClient {
3234

3335
private boolean recursiveMode = true;
3436

35-
private final SCPEngine engine;
36-
3737
SCPDownloadClient(SCPEngine engine) {
38-
this.engine = engine;
38+
super(engine);
39+
}
40+
41+
SCPDownloadClient(SCPEngine engine, int bandwidthLimit) {
42+
super(engine, bandwidthLimit);
3943
}
4044

4145
/** Download a file from {@code sourcePath} on the connected host to {@code targetPath} locally. */
@@ -60,12 +64,12 @@ public void setRecursiveMode(boolean recursive) {
6064

6165
void startCopy(String sourcePath, LocalDestFile targetFile)
6266
throws IOException {
63-
List<Arg> args = new LinkedList<Arg>();
64-
args.add(Arg.SOURCE);
65-
args.add(Arg.QUIET);
66-
args.add(Arg.PRESERVE_TIMES);
67-
if (recursiveMode)
68-
args.add(Arg.RECURSIVE);
67+
List<SCPArgument> args = SCPArguments.with(Arg.SOURCE)
68+
.and(Arg.QUIET)
69+
.and(Arg.PRESERVE_TIMES)
70+
.and(Arg.RECURSIVE, recursiveMode)
71+
.and(Arg.LIMIT, String.valueOf(bandwidthLimit), (bandwidthLimit > 0))
72+
.arguments();
6973
engine.execSCPWith(args, sourcePath);
7074

7175
engine.signal("Start status OK");

‎src/main/java/net/schmizz/sshj/xfer/scp/SCPEngine.java

+85-3
Original file line numberDiff line numberDiff line change
@@ -28,6 +28,7 @@
2828
import java.io.IOException;
2929
import java.io.InputStream;
3030
import java.io.OutputStream;
31+
import java.util.LinkedList;
3132
import java.util.List;
3233

3334
/** @see <a href="http://blogs.sun.com/janp/entry/how_the_scp_protocol_works">SCP Protocol</a> */
@@ -39,7 +40,8 @@ enum Arg {
3940
RECURSIVE('r'),
4041
VERBOSE('v'),
4142
PRESERVE_TIMES('p'),
42-
QUIET('q');
43+
QUIET('q'),
44+
LIMIT('l');
4345

4446
private final char a;
4547

@@ -97,10 +99,10 @@ void cleanSlate() {
9799
exitStatus = -1;
98100
}
99101

100-
void execSCPWith(List<Arg> args, String path)
102+
void execSCPWith(List<SCPArgument> args, String path)
101103
throws SSHException {
102104
final StringBuilder cmd = new StringBuilder(SCP_COMMAND);
103-
for (Arg arg : args) {
105+
for (SCPArgument arg : args) {
104106
cmd.append(" ").append(arg);
105107
}
106108
cmd.append(" ");
@@ -186,4 +188,84 @@ TransferListener getTransferListener() {
186188
return listener;
187189
}
188190

191+
public static class SCPArgument {
192+
193+
private Arg name;
194+
private String value;
195+
196+
private SCPArgument(Arg name, String value) {
197+
this.name = name;
198+
this.value = value;
199+
}
200+
201+
public static SCPArgument addArgument(Arg name, String value) {
202+
return new SCPArgument(name, value);
203+
}
204+
205+
@Override
206+
public String toString() {
207+
String option = name.toString();
208+
if (value != null) {
209+
option = option + value;
210+
}
211+
return option;
212+
}
213+
}
214+
215+
public static class SCPArguments {
216+
217+
private static List<SCPArgument> args = null;
218+
219+
private SCPArguments() {
220+
this.args = new LinkedList<SCPArgument>();
221+
}
222+
223+
private static void addArgument(Arg name, String value, boolean accept) {
224+
if (accept) {
225+
args.add(SCPArgument.addArgument(name, value));
226+
}
227+
}
228+
229+
public static SCPArguments with(Arg name) {
230+
return with(name, null, true);
231+
}
232+
233+
public static SCPArguments with(Arg name, String value) {
234+
return with(name, value, true);
235+
}
236+
237+
public static SCPArguments with(Arg name, boolean accept) {
238+
return with(name, null, accept);
239+
}
240+
241+
public static SCPArguments with(Arg name, String value, boolean accept) {
242+
SCPArguments scpArguments = new SCPArguments();
243+
addArgument(name, value, accept);
244+
return scpArguments;
245+
}
246+
247+
public SCPArguments and(Arg name) {
248+
addArgument(name, null, true);
249+
return this;
250+
}
251+
252+
public SCPArguments and(Arg name, String value) {
253+
addArgument(name, value, true);
254+
return this;
255+
}
256+
257+
public SCPArguments and(Arg name, boolean accept) {
258+
addArgument(name, null, accept);
259+
return this;
260+
}
261+
262+
public SCPArguments and(Arg name, String value, boolean accept) {
263+
addArgument(name, value, accept);
264+
return this;
265+
}
266+
267+
public List<SCPArgument> arguments() {
268+
return args;
269+
}
270+
}
189271
}

‎src/main/java/net/schmizz/sshj/xfer/scp/SCPFileTransfer.java

+13-2
Original file line numberDiff line numberDiff line change
@@ -28,18 +28,23 @@ public class SCPFileTransfer
2828
extends AbstractFileTransfer
2929
implements FileTransfer {
3030

31+
/** Default bandwidth limit for SCP transfer in kilobit/s (-1 means unlimited) */
32+
private static final int DEFAULT_BANDWIDTH_LIMIT = -1;
33+
3134
private final SessionFactory sessionFactory;
35+
private int bandwidthLimit;
3236

3337
public SCPFileTransfer(SessionFactory sessionFactory) {
3438
this.sessionFactory = sessionFactory;
39+
this.bandwidthLimit = DEFAULT_BANDWIDTH_LIMIT;
3540
}
3641

3742
public SCPDownloadClient newSCPDownloadClient() {
38-
return new SCPDownloadClient(newSCPEngine());
43+
return new SCPDownloadClient(newSCPEngine(), bandwidthLimit);
3944
}
4045

4146
public SCPUploadClient newSCPUploadClient() {
42-
return new SCPUploadClient(newSCPEngine());
47+
return new SCPUploadClient(newSCPEngine(), bandwidthLimit);
4348
}
4449

4550
private SCPEngine newSCPEngine() {
@@ -70,4 +75,10 @@ public void upload(LocalSourceFile localFile, String remotePath)
7075
newSCPUploadClient().copy(localFile, remotePath);
7176
}
7277

78+
public SCPFileTransfer bandwidthLimit(int limit) {
79+
if (limit > 0) {
80+
this.bandwidthLimit = limit;
81+
}
82+
return this;
83+
}
7384
}

‎src/main/java/net/schmizz/sshj/xfer/scp/SCPUploadClient.java

+14-9
Original file line numberDiff line numberDiff line change
@@ -24,17 +24,22 @@
2424

2525
import java.io.IOException;
2626
import java.io.InputStream;
27-
import java.util.LinkedList;
2827
import java.util.List;
2928

29+
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArgument;
30+
import static net.schmizz.sshj.xfer.scp.SCPEngine.SCPArguments;
31+
3032
/** Support for uploading files over a connected link using SCP. */
31-
public final class SCPUploadClient {
33+
public final class SCPUploadClient extends AbstractSCPClient {
3234

33-
private final SCPEngine engine;
3435
private LocalFileFilter uploadFilter;
3536

3637
SCPUploadClient(SCPEngine engine) {
37-
this.engine = engine;
38+
super(engine);
39+
}
40+
41+
SCPUploadClient(SCPEngine engine, int bandwidthLimit) {
42+
super(engine, bandwidthLimit);
3843
}
3944

4045
/** Upload a local file from {@code localFile} to {@code targetPath} on the remote host. */
@@ -55,11 +60,11 @@ public void setUploadFilter(LocalFileFilter uploadFilter) {
5560

5661
private synchronized void startCopy(LocalSourceFile sourceFile, String targetPath)
5762
throws IOException {
58-
List<Arg> args = new LinkedList<Arg>();
59-
args.add(Arg.SINK);
60-
args.add(Arg.RECURSIVE);
61-
if (sourceFile.providesAtimeMtime())
62-
args.add(Arg.PRESERVE_TIMES);
63+
List<SCPArgument> args = SCPArguments.with(Arg.SINK)
64+
.and(Arg.RECURSIVE)
65+
.and(Arg.PRESERVE_TIMES, sourceFile.providesAtimeMtime())
66+
.and(Arg.LIMIT, String.valueOf(bandwidthLimit), (bandwidthLimit > 0))
67+
.arguments();
6368
engine.execSCPWith(args, targetPath);
6469
engine.check("Start status OK");
6570
process(engine.getTransferListener(), sourceFile);

‎src/test/java/net/schmizz/sshj/xfer/scp/SCPFileTransferTest.java

+82
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,82 @@
1+
package net.schmizz.sshj.xfer.scp;
2+
3+
import com.hierynomus.sshj.test.SshFixture;
4+
import com.hierynomus.sshj.test.util.FileUtil;
5+
import net.schmizz.sshj.SSHClient;
6+
import org.junit.After;
7+
import org.junit.Before;
8+
import org.junit.Rule;
9+
import org.junit.Test;
10+
import org.junit.rules.TemporaryFolder;
11+
12+
import java.io.File;
13+
import java.io.IOException;
14+
15+
import static junit.framework.Assert.assertFalse;
16+
import static junit.framework.Assert.assertTrue;
17+
18+
public class SCPFileTransferTest {
19+
20+
public static final String DEFAULT_FILE_NAME = "my_file.txt";
21+
File targetDir;
22+
File sourceFile;
23+
File targetFile;
24+
SSHClient sshClient;
25+
26+
@Rule
27+
public SshFixture fixture = new SshFixture();
28+
29+
@Rule
30+
public TemporaryFolder tempFolder = new TemporaryFolder();
31+
32+
@Before
33+
public void init() throws IOException {
34+
sourceFile = tempFolder.newFile(DEFAULT_FILE_NAME);
35+
FileUtil.writeToFile(sourceFile, "This is my file");
36+
targetDir = tempFolder.newFolder();
37+
targetFile = new File(targetDir + File.separator + DEFAULT_FILE_NAME);
38+
sshClient = fixture.setupConnectedDefaultClient();
39+
sshClient.authPassword("test", "test");
40+
}
41+
42+
@After
43+
public void cleanup() {
44+
if (targetFile.exists()) {
45+
targetFile.delete();
46+
}
47+
}
48+
49+
@Test
50+
public void shouldSCPUploadFile() throws IOException {
51+
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer();
52+
assertFalse(targetFile.exists());
53+
scpFileTransfer.upload(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
54+
assertTrue(targetFile.exists());
55+
}
56+
57+
@Test
58+
public void shouldSCPUploadFileWithBandwidthLimit() throws IOException {
59+
// Limit upload transfer at 2Mo/s
60+
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer().bandwidthLimit(16000);
61+
assertFalse(targetFile.exists());
62+
scpFileTransfer.upload(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
63+
assertTrue(targetFile.exists());
64+
}
65+
66+
@Test
67+
public void shouldSCPDownloadFile() throws IOException {
68+
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer();
69+
assertFalse(targetFile.exists());
70+
scpFileTransfer.download(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
71+
assertTrue(targetFile.exists());
72+
}
73+
74+
@Test
75+
public void shouldSCPDownloadFileWithBandwidthLimit() throws IOException {
76+
// Limit download transfer at 128Ko/s
77+
SCPFileTransfer scpFileTransfer = sshClient.newSCPFileTransfer().bandwidthLimit(1024);
78+
assertFalse(targetFile.exists());
79+
scpFileTransfer.download(sourceFile.getAbsolutePath(), targetDir.getAbsolutePath());
80+
assertTrue(targetFile.exists());
81+
}
82+
}

0 commit comments

Comments
 (0)