- djm@cvs.openbsd.org 2010/09/22 22:58:51

     [atomicio.c atomicio.h misc.c misc.h scp.c sftp-client.c]
     [sftp-client.h sftp.1 sftp.c]
     add an option per-read/write callback to atomicio

     factor out bandwidth limiting code from scp(1) into a generic bandwidth
     limiter that can be attached using the atomicio callback mechanism

     add a bandwidth limit option to sftp(1) using the above
     "very nice" markus@

10 files changed, 282 insertions, 213 deletions

diff --git a/ChangeLog b/ChangeLog
index 5f391400..b3338dcb 100644
--- a/ChangeLog
+++ b/ChangeLog
@@ -44,6 +44,16 @@
      ssh_config.5: format the kexalgorithms in a more consistent
      (prettier!) way
      ok djm
+   - djm@cvs.openbsd.org 2010/09/22 22:58:51
+     [atomicio.c atomicio.h misc.c misc.h scp.c sftp-client.c]
+     [sftp-client.h sftp.1 sftp.c]
+     add an option per-read/write callback to atomicio
+
+     factor out bandwidth limiting code from scp(1) into a generic bandwidth
+     limiter that can be attached using the atomicio callback mechanism
+
+     add a bandwidth limit option to sftp(1) using the above
+     "very nice" markus@
 
 20100910
  - (dtucker) [openbsd-compat/port-linux.c] Check is_selinux_enabled for exact
diff --git a/atomicio.c b/atomicio.c
index a6b2d127..601b3c37 100644
--- a/atomicio.c
+++ b/atomicio.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: atomicio.c,v 1.25 2007/06/25 12:02:27 dtucker Exp $ */
+/* $OpenBSD: atomicio.c,v 1.26 2010/09/22 22:58:51 djm Exp $ */
 /*
  * Copyright (c) 2006 Damien Miller. All rights reserved.
  * Copyright (c) 2005 Anil Madhavapeddy. All rights reserved.
@@ -48,7 +48,8 @@
  * ensure all of data on socket comes through. f==read || f==vwrite
  */
 size_t
-atomicio(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n)
+atomicio6(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n,
+    int (*cb)(void *, size_t), void *cb_arg)
 {
 	char *s = _s;
 	size_t pos = 0;
@@ -73,17 +74,28 @@ atomicio(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n)
 			return pos;
 		default:
 			pos += (size_t)res;
+			if (cb != NULL && cb(cb_arg, (size_t)res) == -1) {
+				errno = EINTR;
+				return pos;
+			}
 		}
 	}
-	return (pos);
+	return pos;
+}
+
+size_t
+atomicio(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n)
+{
+	return atomicio6(f, fd, _s, n, NULL, NULL);
 }
 
 /*
  * ensure all of data on socket comes through. f==readv || f==writev
  */
 size_t
-atomiciov(ssize_t (*f) (int, const struct iovec *, int), int fd,
-    const struct iovec *_iov, int iovcnt)
+atomiciov6(ssize_t (*f) (int, const struct iovec *, int), int fd,
+    const struct iovec *_iov, int iovcnt,
+    int (*cb)(void *, size_t), void *cb_arg)
 {
 	size_t pos = 0, rem;
 	ssize_t res;
@@ -137,6 +149,17 @@ atomiciov(ssize_t (*f) (int, const struct iovec *, int), int fd,
 			iov[0].iov_base = ((char *)iov[0].iov_base) + rem;
 			iov[0].iov_len -= rem;
 		}
+		if (cb != NULL && cb(cb_arg, (size_t)res) == -1) {
+			errno = EINTR;
+			return pos;
+		}
 	}
 	return pos;
 }
+
+size_t
+atomiciov(ssize_t (*f) (int, const struct iovec *, int), int fd,
+    const struct iovec *_iov, int iovcnt)
+{
+	return atomiciov6(f, fd, _iov, iovcnt, NULL, NULL);
+}
diff --git a/atomicio.h b/atomicio.h
index 2fcd25d4..0d728ac8 100644
--- a/atomicio.h
+++ b/atomicio.h
@@ -1,4 +1,4 @@
-/* $OpenBSD: atomicio.h,v 1.10 2006/08/03 03:34:41 deraadt Exp $ */
+/* $OpenBSD: atomicio.h,v 1.11 2010/09/22 22:58:51 djm Exp $ */
 
 /*
  * Copyright (c) 2006 Damien Miller.  All rights reserved.
@@ -32,6 +32,9 @@
 /*
  * Ensure all of data on socket comes through. f==read || f==vwrite
  */
+size_t
+atomicio6(ssize_t (*f) (int, void *, size_t), int fd, void *_s, size_t n,
+    int (*cb)(void *, size_t), void *);
 size_t	atomicio(ssize_t (*)(int, void *, size_t), int, void *, size_t);
 
 #define vwrite (ssize_t (*)(int, void *, size_t))write
@@ -39,6 +42,9 @@ size_t	atomicio(ssize_t (*)(int, void *, size_t), int, void *, size_t);
 /*
  * ensure all of data on socket comes through. f==readv || f==writev
  */
+size_t
+atomiciov6(ssize_t (*f) (int, const struct iovec *, int), int fd,
+    const struct iovec *_iov, int iovcnt, int (*cb)(void *, size_t), void *);
 size_t	atomiciov(ssize_t (*)(int, const struct iovec *, int),
     int, const struct iovec *, int);
 
diff --git a/misc.c b/misc.c
index a82e7936..41c92a82 100644
--- a/misc.c
+++ b/misc.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: misc.c,v 1.80 2010/07/21 02:10:58 djm Exp $ */
+/* $OpenBSD: misc.c,v 1.81 2010/09/22 22:58:51 djm Exp $ */
 /*
  * Copyright (c) 2000 Markus Friedl.  All rights reserved.
  * Copyright (c) 2005,2006 Damien Miller.  All rights reserved.
@@ -860,6 +860,70 @@ timingsafe_bcmp(const void *b1, const void *b2, size_t n)
 		ret |= *p1++ ^ *p2++;
 	return (ret != 0);
 }
+
+void
+bandwidth_limit_init(struct bwlimit *bw, u_int64_t kbps, size_t buflen)
+{
+	bw->buflen = buflen;
+	bw->rate = kbps;
+	bw->thresh = bw->rate;
+	bw->lamt = 0;
+	timerclear(&bw->bwstart);
+	timerclear(&bw->bwend);
+}	
+
+/* Callback from read/write loop to insert bandwidth-limiting delays */
+void
+bandwidth_limit(struct bwlimit *bw, size_t read_len)
+{
+	u_int64_t waitlen;
+	struct timespec ts, rm;
+
+	if (!timerisset(&bw->bwstart)) {
+		gettimeofday(&bw->bwstart, NULL);
+		return;
+	}
+
+	bw->lamt += read_len;
+	if (bw->lamt < bw->thresh)
+		return;
+
+	gettimeofday(&bw->bwend, NULL);
+	timersub(&bw->bwend, &bw->bwstart, &bw->bwend);
+	if (!timerisset(&bw->bwend))
+		return;
+
+	bw->lamt *= 8;
+	waitlen = (double)1000000L * bw->lamt / bw->rate;
+
+	bw->bwstart.tv_sec = waitlen / 1000000L;
+	bw->bwstart.tv_usec = waitlen % 1000000L;
+
+	if (timercmp(&bw->bwstart, &bw->bwend, >)) {
+		timersub(&bw->bwstart, &bw->bwend, &bw->bwend);
+
+		/* Adjust the wait time */
+		if (bw->bwend.tv_sec) {
+			bw->thresh /= 2;
+			if (bw->thresh < bw->buflen / 4)
+				bw->thresh = bw->buflen / 4;
+		} else if (bw->bwend.tv_usec < 10000) {
+			bw->thresh *= 2;
+			if (bw->thresh > bw->buflen * 8)
+				bw->thresh = bw->buflen * 8;
+		}
+
+		TIMEVAL_TO_TIMESPEC(&bw->bwend, &ts);
+		while (nanosleep(&ts, &rm) == -1) {
+			if (errno != EINTR)
+				break;
+			ts = rm;
+		}
+	}
+
+	bw->lamt = 0;
+	gettimeofday(&bw->bwstart, NULL);
+}
 void
 sock_set_v6only(int s)
 {
diff --git a/misc.h b/misc.h
index bb799f61..f5aab029 100644
--- a/misc.h
+++ b/misc.h
@@ -1,4 +1,4 @@
-/* $OpenBSD: misc.h,v 1.43 2010/07/13 23:13:16 djm Exp $ */
+/* $OpenBSD: misc.h,v 1.44 2010/09/22 22:58:51 djm Exp $ */
 
 /*
  * Author: Tatu Ylonen <ylo@cs.hut.fi>
@@ -80,6 +80,15 @@ void		put_u32(void *, u_int32_t)
 void		put_u16(void *, u_int16_t)
     __attribute__((__bounded__( __minbytes__, 1, 2)));
 
+struct bwlimit {
+	size_t buflen;
+	u_int64_t rate, thresh, lamt;
+	struct timeval bwstart, bwend;
+};
+
+void bandwidth_limit_init(struct bwlimit *, u_int64_t, size_t);
+void bandwidth_limit(struct bwlimit *, size_t);
+
 
 /* readpass.c */
 
diff --git a/scp.c b/scp.c
index e07de42f..a4066c66 100644
--- a/scp.c
+++ b/scp.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: scp.c,v 1.166 2010/07/01 13:06:59 millert Exp $ */
+/* $OpenBSD: scp.c,v 1.167 2010/09/22 22:58:51 djm Exp $ */
 /*
  * scp - secure remote copy.  This is basically patched BSD rcp which
  * uses ssh to do the data transfer (instead of using rcmd).
@@ -120,13 +120,12 @@ extern char *__progname;
 
 int do_cmd(char *host, char *remuser, char *cmd, int *fdin, int *fdout);
 
-void bwlimit(int);
-
 /* Struct for addargs */
 arglist args;
 
 /* Bandwidth limit */
-off_t limit_rate = 0;
+long long limit_kbps = 0;
+struct bwlimit bwlimit;
 
 /* Name of current file being transferred. */
 char *curfile;
@@ -312,15 +311,14 @@ void sink(int, char *[]);
 void source(int, char *[]);
 void tolocal(int, char *[]);
 void toremote(char *, int, char *[]);
-size_t scpio(ssize_t (*)(int, void *, size_t), int, void *, size_t, off_t *);
 void usage(void);
 
 int
 main(int argc, char **argv)
 {
 	int ch, fflag, tflag, status, n;
-	double speed;
-	char *targ, *endp, **newargv;
+	char *targ, **newargv;
+	const char *errstr;
 	extern char *optarg;
 	extern int optind;
 
@@ -369,10 +367,12 @@ main(int argc, char **argv)
 			addargs(&args, "-oBatchmode yes");
 			break;
 		case 'l':
-			speed = strtod(optarg, &endp);
-			if (speed <= 0 || *endp != '\0')
+			limit_kbps = strtonum(optarg, 1, 100 * 1024 * 1024,
+			    &errstr);
+			if (errstr != NULL)
 				usage();
-			limit_rate = speed * 1024;
+			limit_kbps *= 1024; /* kbps */
+			bandwidth_limit_init(&bwlimit, limit_kbps, COPY_BUFLEN);
 			break;
 		case 'p':
 			pflag = 1;
@@ -474,41 +474,16 @@ main(int argc, char **argv)
 	exit(errs != 0);
 }
 
-/*
- * atomicio-like wrapper that also applies bandwidth limits and updates
- * the progressmeter counter.
- */
-size_t
-scpio(ssize_t (*f)(int, void *, size_t), int fd, void *_p, size_t l, off_t *c)
+/* Callback from atomicio6 to update progress meter and limit bandwidth */
+static int
+scpio(void *_cnt, size_t s)
 {
-	u_char *p = (u_char *)_p;
-	size_t offset;
-	ssize_t r;
-	struct pollfd pfd;
-
-	pfd.fd = fd;
-	pfd.events = f == read ? POLLIN : POLLOUT;
-	for (offset = 0; offset < l;) {
-		r = f(fd, p + offset, l - offset);
-		if (r == 0) {
-			errno = EPIPE;
-			return offset;
-		}
-		if (r < 0) {
-			if (errno == EINTR)
-				continue;
-			if (errno == EAGAIN || errno == EWOULDBLOCK) {
-				(void)poll(&pfd, 1, -1); /* Ignore errors */
-				continue;
-			}
-			return offset;
-		}
-		offset += (size_t)r;
-		*c += (off_t)r;
-		if (limit_rate)
-			bwlimit(r);
-	}
-	return offset;
+	off_t *cnt = (off_t *)_cnt;
+
+	*cnt += s;
+	if (limit_kbps > 0)
+		bandwidth_limit(&bwlimit, s);
+	return 0;
 }
 
 void
@@ -750,7 +725,7 @@ next:			if (fd != -1) {
 				(void)atomicio(vwrite, remout, bp->buf, amt);
 				continue;
 			}
-			if (scpio(vwrite, remout, bp->buf, amt,
+			if (atomicio6(vwrite, remout, bp->buf, amt, scpio,
 			    &statbytes) != amt)
 				haderr = errno;
 		}
@@ -825,60 +800,6 @@ rsource(char *name, struct stat *statp)
 }
 
 void
-bwlimit(int amount)
-{
-	static struct timeval bwstart, bwend;
-	static int lamt, thresh = 16384;
-	u_int64_t waitlen;
-	struct timespec ts, rm;
-
-	if (!timerisset(&bwstart)) {
-		gettimeofday(&bwstart, NULL);
-		return;
-	}
-
-	lamt += amount;
-	if (lamt < thresh)
-		return;
-
-	gettimeofday(&bwend, NULL);
-	timersub(&bwend, &bwstart, &bwend);
-	if (!timerisset(&bwend))
-		return;
-
-	lamt *= 8;
-	waitlen = (double)1000000L * lamt / limit_rate;
-
-	bwstart.tv_sec = waitlen / 1000000L;
-	bwstart.tv_usec = waitlen % 1000000L;
-
-	if (timercmp(&bwstart, &bwend, >)) {
-		timersub(&bwstart, &bwend, &bwend);
-
-		/* Adjust the wait time */
-		if (bwend.tv_sec) {
-			thresh /= 2;
-			if (thresh < 2048)
-				thresh = 2048;
-		} else if (bwend.tv_usec < 10000) {
-			thresh *= 2;
-			if (thresh > COPY_BUFLEN * 4)
-				thresh = COPY_BUFLEN * 4;
-		}
-
-		TIMEVAL_TO_TIMESPEC(&bwend, &ts);
-		while (nanosleep(&ts, &rm) == -1) {
-			if (errno != EINTR)
-				break;
-			ts = rm;
-		}
-	}
-
-	lamt = 0;
-	gettimeofday(&bwstart, NULL);
-}
-
-void
 sink(int argc, char **argv)
 {
 	static BUF buffer;
@@ -1071,7 +992,8 @@ bad:			run_err("%s: %s", np, strerror(errno));
 				amt = size - i;
 			count += amt;
 			do {
-				j = scpio(read, remin, cp, amt, &statbytes);
+				j = atomicio6(read, remin, cp, amt,
+				    scpio, &statbytes);
 				if (j == 0) {
 					run_err("%s", j != EPIPE ?
 					    strerror(errno) :
diff --git a/sftp-client.c b/sftp-client.c
index 9dab4778..4e009ef2 100644
--- a/sftp-client.c
+++ b/sftp-client.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: sftp-client.c,v 1.92 2010/07/19 03:16:33 djm Exp $ */
+/* $OpenBSD: sftp-client.c,v 1.93 2010/09/22 22:58:51 djm Exp $ */
 /*
  * Copyright (c) 2001-2004 Damien Miller <djm@openbsd.org>
  *
@@ -76,14 +76,26 @@ struct sftp_conn {
 #define SFTP_EXT_STATVFS	0x00000002
 #define SFTP_EXT_FSTATVFS	0x00000004
 	u_int exts;
+	u_int64_t limit_kbps;
+	struct bwlimit bwlimit_in, bwlimit_out;
 };
 
 static char *
-get_handle(int fd, u_int expected_id, u_int *len, const char *errfmt, ...)
-    __attribute__((format(printf, 4, 5)));
+get_handle(struct sftp_conn *conn, u_int expected_id, u_int *len,
+    const char *errfmt, ...) __attribute__((format(printf, 4, 5)));
+
+/* ARGSUSED */
+static int
+sftpio(void *_bwlimit, size_t amount)
+{
+	struct bwlimit *bwlimit = (struct bwlimit *)_bwlimit;
+
+	bandwidth_limit(bwlimit, amount);
+	return 0;
+}
 
 static void
-send_msg(int fd, Buffer *m)
+send_msg(struct sftp_conn *conn, Buffer *m)
 {
 	u_char mlen[4];
 	struct iovec iov[2];
@@ -98,19 +110,22 @@ send_msg(int fd, Buffer *m)
 	iov[1].iov_base = buffer_ptr(m);
 	iov[1].iov_len = buffer_len(m);
 
-	if (atomiciov(writev, fd, iov, 2) != buffer_len(m) + sizeof(mlen))
+	if (atomiciov6(writev, conn->fd_out, iov, 2,
+	    conn->limit_kbps > 0 ? sftpio : NULL, &conn->bwlimit_out) != 
+	    buffer_len(m) + sizeof(mlen))
 		fatal("Couldn't send packet: %s", strerror(errno));
 
 	buffer_clear(m);
 }
 
 static void
-get_msg(int fd, Buffer *m)
+get_msg(struct sftp_conn *conn, Buffer *m)
 {
 	u_int msg_len;
 
 	buffer_append_space(m, 4);
-	if (atomicio(read, fd, buffer_ptr(m), 4) != 4) {
+	if (atomicio6(read, conn->fd_in, buffer_ptr(m), 4,
+	    conn->limit_kbps > 0 ? sftpio : NULL, &conn->bwlimit_in) != 4) {
 		if (errno == EPIPE)
 			fatal("Connection closed");
 		else
@@ -122,7 +137,9 @@ get_msg(int fd, Buffer *m)
 		fatal("Received message too long %u", msg_len);
 
 	buffer_append_space(m, msg_len);
-	if (atomicio(read, fd, buffer_ptr(m), msg_len) != msg_len) {
+	if (atomicio6(read, conn->fd_in, buffer_ptr(m), msg_len,
+	    conn->limit_kbps > 0 ? sftpio : NULL, &conn->bwlimit_in)
+	    != msg_len) {
 		if (errno == EPIPE)
 			fatal("Connection closed");
 		else
@@ -131,7 +148,7 @@ get_msg(int fd, Buffer *m)
 }
 
 static void
-send_string_request(int fd, u_int id, u_int code, char *s,
+send_string_request(struct sftp_conn *conn, u_int id, u_int code, char *s,
     u_int len)
 {
 	Buffer msg;
@@ -140,14 +157,14 @@ send_string_request(int fd, u_int id, u_int code, char *s,
 	buffer_put_char(&msg, code);
 	buffer_put_int(&msg, id);
 	buffer_put_string(&msg, s, len);
-	send_msg(fd, &msg);
-	debug3("Sent message fd %d T:%u I:%u", fd, code, id);
+	send_msg(conn, &msg);
+	debug3("Sent message fd %d T:%u I:%u", conn->fd_out, code, id);
 	buffer_free(&msg);
 }
 
 static void
-send_string_attrs_request(int fd, u_int id, u_int code, char *s,
-    u_int len, Attrib *a)
+send_string_attrs_request(struct sftp_conn *conn, u_int id, u_int code,
+    char *s, u_int len, Attrib *a)
 {
 	Buffer msg;
 
@@ -156,19 +173,19 @@ send_string_attrs_request(int fd, u_int id, u_int code, char *s,
 	buffer_put_int(&msg, id);
 	buffer_put_string(&msg, s, len);
 	encode_attrib(&msg, a);
-	send_msg(fd, &msg);
-	debug3("Sent message fd %d T:%u I:%u", fd, code, id);
+	send_msg(conn, &msg);
+	debug3("Sent message fd %d T:%u I:%u", conn->fd_out, code, id);
 	buffer_free(&msg);
 }
 
 static u_int
-get_status(int fd, u_int expected_id)
+get_status(struct sftp_conn *conn, u_int expected_id)
 {
 	Buffer msg;
 	u_int type, id, status;
 
 	buffer_init(&msg);
-	get_msg(fd, &msg);
+	get_msg(conn, &msg);
 	type = buffer_get_char(&msg);
 	id = buffer_get_int(&msg);
 
@@ -183,11 +200,12 @@ get_status(int fd, u_int expected_id)
 
 	debug3("SSH2_FXP_STATUS %u", status);
 
-	return(status);
+	return status;
 }
 
 static char *
-get_handle(int fd, u_int expected_id, u_int *len, const char *errfmt, ...)
+get_handle(struct sftp_conn *conn, u_int expected_id, u_int *len,
+    const char *errfmt, ...)
 {
 	Buffer msg;
 	u_int type, id;
@@ -201,7 +219,7 @@ get_handle(int fd, u_int expected_id, u_int *len, const char *errfmt, ...)
 	va_end(args);
 
 	buffer_init(&msg);
-	get_msg(fd, &msg);
+	get_msg(conn, &msg);
 	type = buffer_get_char(&msg);
 	id = buffer_get_int(&msg);
 
@@ -225,14 +243,14 @@ get_handle(int fd, u_int expected_id, u_int *len, const char *errfmt, ...)
 }
 
 static Attrib *
-get_decode_stat(int fd, u_int expected_id, int quiet)
+get_decode_stat(struct sftp_conn *conn, u_int expected_id, int quiet)
 {
 	Buffer msg;
 	u_int type, id;
 	Attrib *a;
 
 	buffer_init(&msg);
-	get_msg(fd, &msg);
+	get_msg(conn, &msg);
 
 	type = buffer_get_char(&msg);
 	id = buffer_get_int(&msg);
@@ -260,14 +278,14 @@ get_decode_stat(int fd, u_int expected_id, int quiet)
 }
 
 static int
-get_decode_statvfs(int fd, struct sftp_statvfs *st, u_int expected_id,
-    int quiet)
+get_decode_statvfs(struct sftp_conn *conn, struct sftp_statvfs *st,
+    u_int expected_id, int quiet)
 {
 	Buffer msg;
 	u_int type, id, flag;
 
 	buffer_init(&msg);
-	get_msg(fd, &msg);
+	get_msg(conn, &msg);
 
 	type = buffer_get_char(&msg);
 	id = buffer_get_int(&msg);
@@ -311,21 +329,29 @@ get_decode_statvfs(int fd, struct sftp_statvfs *st, u_int expected_id,
 }
 
 struct sftp_conn *
-do_init(int fd_in, int fd_out, u_int transfer_buflen, u_int num_requests)
+do_init(int fd_in, int fd_out, u_int transfer_buflen, u_int num_requests,
+    u_int64_t limit_kbps)
 {
-	u_int type, exts = 0;
-	int version;
+	u_int type;
 	Buffer msg;
 	struct sftp_conn *ret;
 
+	ret = xmalloc(sizeof(*ret));
+	ret->fd_in = fd_in;
+	ret->fd_out = fd_out;
+	ret->transfer_buflen = transfer_buflen;
+	ret->num_requests = num_requests;
+	ret->exts = 0;
+	ret->limit_kbps = 0;
+
 	buffer_init(&msg);
 	buffer_put_char(&msg, SSH2_FXP_INIT);
 	buffer_put_int(&msg, SSH2_FILEXFER_VERSION);
-	send_msg(fd_out, &msg);
+	send_msg(ret, &msg);
 
 	buffer_clear(&msg);
 
-	get_msg(fd_in, &msg);
+	get_msg(ret, &msg);
 
 	/* Expecting a VERSION reply */
 	if ((type = buffer_get_char(&msg)) != SSH2_FXP_VERSION) {
@@ -334,9 +360,9 @@ do_init(int fd_in, int fd_out, u_int transfer_buflen, u_int num_requests)
 		buffer_free(&msg);
 		return(NULL);
 	}
-	version = buffer_get_int(&msg);
+	ret->version = buffer_get_int(&msg);
 
-	debug2("Remote version: %d", version);
+	debug2("Remote version: %u", ret->version);
 
 	/* Check for extensions */
 	while (buffer_len(&msg) > 0) {
@@ -346,15 +372,15 @@ do_init(int fd_in, int fd_out, u_int transfer_buflen, u_int num_requests)
 
 		if (strcmp(name, "posix-rename@openssh.com") == 0 &&
 		    strcmp(value, "1") == 0) {
-			exts |= SFTP_EXT_POSIX_RENAME;
+			ret->exts |= SFTP_EXT_POSIX_RENAME;
 			known = 1;
 		} else if (strcmp(name, "statvfs@openssh.com") == 0 &&
 		    strcmp(value, "2") == 0) {
-			exts |= SFTP_EXT_STATVFS;
+			ret->exts |= SFTP_EXT_STATVFS;
 			known = 1;
 		} if (strcmp(name, "fstatvfs@openssh.com") == 0 &&
 		    strcmp(value, "2") == 0) {
-			exts |= SFTP_EXT_FSTATVFS;
+			ret->exts |= SFTP_EXT_FSTATVFS;
 			known = 1;
 		}
 		if (known) {
@@ -369,26 +395,25 @@ do_init(int fd_in, int fd_out, u_int transfer_buflen, u_int num_requests)
 
 	buffer_free(&msg);
 
-	ret = xmalloc(sizeof(*ret));
-	ret->fd_in = fd_in;
-	ret->fd_out = fd_out;
-	ret->transfer_buflen = transfer_buflen;
-	ret->num_requests = num_requests;
-	ret->version = version;
-	ret->msg_id = 1;
-	ret->exts = exts;
-
 	/* Some filexfer v.0 servers don't support large packets */
-	if (version == 0)
+	if (ret->version == 0)
 		ret->transfer_buflen = MIN(ret->transfer_buflen, 20480);
 
-	return(ret);
+	ret->limit_kbps = limit_kbps;
+	if (ret->limit_kbps > 0) {
+		bandwidth_limit_init(&ret->bwlimit_in, ret->limit_kbps,
+		    ret->transfer_buflen);
+		bandwidth_limit_init(&ret->bwlimit_out, ret->limit_kbps,
+		    ret->transfer_buflen);
+	}
+
+	return ret;
 }
 
 u_int
 sftp_proto_version(struct sftp_conn *conn)
 {
-	return(conn->version);
+	return conn->version;
 }
 
 int
@@ -403,16 +428,16 @@ do_close(struct sftp_conn *conn, char *handle, u_int handle_len)
 	buffer_put_char(&msg, SSH2_FXP_CLOSE);
 	buffer_put_int(&msg, id);
 	buffer_put_string(&msg, handle, handle_len);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	debug3("Sent message SSH2_FXP_CLOSE I:%u", id);
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't close file: %s", fx2txt(status));
 
 	buffer_free(&msg);
 
-	return(status);
+	return status;
 }
 
 
@@ -430,14 +455,14 @@ do_lsreaddir(struct sftp_conn *conn, char *path, int printflag,
 	buffer_put_char(&msg, SSH2_FXP_OPENDIR);
 	buffer_put_int(&msg, id);
 	buffer_put_cstring(&msg, path);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 
 	buffer_clear(&msg);
 
-	handle = get_handle(conn->fd_in, id, &handle_len,
+	handle = get_handle(conn, id, &handle_len,
 	    "remote readdir(\"%s\")", path);
 	if (handle == NULL)
-		return(-1);
+		return -1;
 
 	if (dir) {
 		ents = 0;
@@ -454,11 +479,11 @@ do_lsreaddir(struct sftp_conn *conn, char *path, int printflag,
 		buffer_put_char(&msg, SSH2_FXP_READDIR);
 		buffer_put_int(&msg, id);
 		buffer_put_string(&msg, handle, handle_len);
-		send_msg(conn->fd_out, &msg);
+		send_msg(conn, &msg);
 
 		buffer_clear(&msg);
 
-		get_msg(conn->fd_in, &msg);
+		get_msg(conn, &msg);
 
 		type = buffer_get_char(&msg);
 		id = buffer_get_int(&msg);
@@ -537,7 +562,7 @@ do_lsreaddir(struct sftp_conn *conn, char *path, int printflag,
 		**dir = NULL;
 	}
 
-	return(0);
+	return 0;
 }
 
 int
@@ -566,9 +591,8 @@ do_rm(struct sftp_conn *conn, char *path)
 	debug2("Sending SSH2_FXP_REMOVE \"%s\"", path);
 
 	id = conn->msg_id++;
-	send_string_request(conn->fd_out, id, SSH2_FXP_REMOVE, path,
-	    strlen(path));
-	status = get_status(conn->fd_in, id);
+	send_string_request(conn, id, SSH2_FXP_REMOVE, path, strlen(path));
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't delete file: %s", fx2txt(status));
 	return(status);
@@ -580,10 +604,10 @@ do_mkdir(struct sftp_conn *conn, char *path, Attrib *a, int printflag)
 	u_int status, id;
 
 	id = conn->msg_id++;
-	send_string_attrs_request(conn->fd_out, id, SSH2_FXP_MKDIR, path,
+	send_string_attrs_request(conn, id, SSH2_FXP_MKDIR, path,
 	    strlen(path), a);
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK && printflag)
 		error("Couldn't create directory: %s", fx2txt(status));
 
@@ -596,10 +620,10 @@ do_rmdir(struct sftp_conn *conn, char *path)
 	u_int status, id;
 
 	id = conn->msg_id++;
-	send_string_request(conn->fd_out, id, SSH2_FXP_RMDIR, path,
+	send_string_request(conn, id, SSH2_FXP_RMDIR, path,
 	    strlen(path));
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't remove directory: %s", fx2txt(status));
 
@@ -613,11 +637,11 @@ do_stat(struct sftp_conn *conn, char *path, int quiet)
 
 	id = conn->msg_id++;
 
-	send_string_request(conn->fd_out, id,
+	send_string_request(conn, id,
 	    conn->version == 0 ? SSH2_FXP_STAT_VERSION_0 : SSH2_FXP_STAT,
 	    path, strlen(path));
 
-	return(get_decode_stat(conn->fd_in, id, quiet));
+	return(get_decode_stat(conn, id, quiet));
 }
 
 Attrib *
@@ -634,10 +658,10 @@ do_lstat(struct sftp_conn *conn, char *path, int quiet)
 	}
 
 	id = conn->msg_id++;
-	send_string_request(conn->fd_out, id, SSH2_FXP_LSTAT, path,
+	send_string_request(conn, id, SSH2_FXP_LSTAT, path,
 	    strlen(path));
 
-	return(get_decode_stat(conn->fd_in, id, quiet));
+	return(get_decode_stat(conn, id, quiet));
 }
 
 #ifdef notyet
@@ -647,10 +671,10 @@ do_fstat(struct sftp_conn *conn, char *handle, u_int handle_len, int quiet)
 	u_int id;
 
 	id = conn->msg_id++;
-	send_string_request(conn->fd_out, id, SSH2_FXP_FSTAT, handle,
+	send_string_request(conn, id, SSH2_FXP_FSTAT, handle,
 	    handle_len);
 
-	return(get_decode_stat(conn->fd_in, id, quiet));
+	return(get_decode_stat(conn, id, quiet));
 }
 #endif
 
@@ -660,10 +684,10 @@ do_setstat(struct sftp_conn *conn, char *path, Attrib *a)
 	u_int status, id;
 
 	id = conn->msg_id++;
-	send_string_attrs_request(conn->fd_out, id, SSH2_FXP_SETSTAT, path,
+	send_string_attrs_request(conn, id, SSH2_FXP_SETSTAT, path,
 	    strlen(path), a);
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't setstat on \"%s\": %s", path,
 		    fx2txt(status));
@@ -678,10 +702,10 @@ do_fsetstat(struct sftp_conn *conn, char *handle, u_int handle_len,
 	u_int status, id;
 
 	id = conn->msg_id++;
-	send_string_attrs_request(conn->fd_out, id, SSH2_FXP_FSETSTAT, handle,
+	send_string_attrs_request(conn, id, SSH2_FXP_FSETSTAT, handle,
 	    handle_len, a);
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't fsetstat: %s", fx2txt(status));
 
@@ -697,12 +721,12 @@ do_realpath(struct sftp_conn *conn, char *path)
 	Attrib *a;
 
 	expected_id = id = conn->msg_id++;
-	send_string_request(conn->fd_out, id, SSH2_FXP_REALPATH, path,
+	send_string_request(conn, id, SSH2_FXP_REALPATH, path,
 	    strlen(path));
 
 	buffer_init(&msg);
 
-	get_msg(conn->fd_in, &msg);
+	get_msg(conn, &msg);
 	type = buffer_get_char(&msg);
 	id = buffer_get_int(&msg);
 
@@ -756,13 +780,13 @@ do_rename(struct sftp_conn *conn, char *oldpath, char *newpath)
 	}
 	buffer_put_cstring(&msg, oldpath);
 	buffer_put_cstring(&msg, newpath);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	debug3("Sent message %s \"%s\" -> \"%s\"",
 	    (conn->exts & SFTP_EXT_POSIX_RENAME) ? "posix-rename@openssh.com" :
 	    "SSH2_FXP_RENAME", oldpath, newpath);
 	buffer_free(&msg);
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't rename file \"%s\" to \"%s\": %s", oldpath,
 		    newpath, fx2txt(status));
@@ -789,12 +813,12 @@ do_symlink(struct sftp_conn *conn, char *oldpath, char *newpath)
 	buffer_put_int(&msg, id);
 	buffer_put_cstring(&msg, oldpath);
 	buffer_put_cstring(&msg, newpath);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	debug3("Sent message SSH2_FXP_SYMLINK \"%s\" -> \"%s\"", oldpath,
 	    newpath);
 	buffer_free(&msg);
 
-	status = get_status(conn->fd_in, id);
+	status = get_status(conn, id);
 	if (status != SSH2_FX_OK)
 		error("Couldn't symlink file \"%s\" to \"%s\": %s", oldpath,
 		    newpath, fx2txt(status));
@@ -812,12 +836,11 @@ do_readlink(struct sftp_conn *conn, char *path)
 	Attrib *a;
 
 	expected_id = id = conn->msg_id++;
-	send_string_request(conn->fd_out, id, SSH2_FXP_READLINK, path,
-	    strlen(path));
+	send_string_request(conn, id, SSH2_FXP_READLINK, path, strlen(path));
 
 	buffer_init(&msg);
 
-	get_msg(conn->fd_in, &msg);
+	get_msg(conn, &msg);
 	type = buffer_get_char(&msg);
 	id = buffer_get_int(&msg);
 
@@ -871,10 +894,10 @@ do_statvfs(struct sftp_conn *conn, const char *path, struct sftp_statvfs *st,
 	buffer_put_int(&msg, id);
 	buffer_put_cstring(&msg, "statvfs@openssh.com");
 	buffer_put_cstring(&msg, path);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	buffer_free(&msg);
 
-	return get_decode_statvfs(conn->fd_in, st, id, quiet);
+	return get_decode_statvfs(conn, st, id, quiet);
 }
 
 #ifdef notyet
@@ -898,16 +921,16 @@ do_fstatvfs(struct sftp_conn *conn, const char *handle, u_int handle_len,
 	buffer_put_int(&msg, id);
 	buffer_put_cstring(&msg, "fstatvfs@openssh.com");
 	buffer_put_string(&msg, handle, handle_len);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	buffer_free(&msg);
 
-	return get_decode_statvfs(conn->fd_in, st, id, quiet);
+	return get_decode_statvfs(conn, st, id, quiet);
 }
 #endif
 
 static void
-send_read_request(int fd_out, u_int id, u_int64_t offset, u_int len,
-    char *handle, u_int handle_len)
+send_read_request(struct sftp_conn *conn, u_int id, u_int64_t offset,
+    u_int len, char *handle, u_int handle_len)
 {
 	Buffer msg;
 
@@ -918,7 +941,7 @@ send_read_request(int fd_out, u_int id, u_int64_t offset, u_int len,
 	buffer_put_string(&msg, handle, handle_len);
 	buffer_put_int64(&msg, offset);
 	buffer_put_int(&msg, len);
-	send_msg(fd_out, &msg);
+	send_msg(conn, &msg);
 	buffer_free(&msg);
 }
 
@@ -976,10 +999,10 @@ do_download(struct sftp_conn *conn, char *remote_path, char *local_path,
 	buffer_put_int(&msg, SSH2_FXF_READ);
 	attrib_clear(&junk); /* Send empty attributes */
 	encode_attrib(&msg, &junk);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	debug3("Sent message SSH2_FXP_OPEN I:%u P:%s", id, remote_path);
 
-	handle = get_handle(conn->fd_in, id, &handle_len,
+	handle = get_handle(conn, id, &handle_len,
 	    "remote open(\"%s\")", remote_path);
 	if (handle == NULL) {
 		buffer_free(&msg);
@@ -1032,12 +1055,12 @@ do_download(struct sftp_conn *conn, char *remote_path, char *local_path,
 			offset += buflen;
 			num_req++;
 			TAILQ_INSERT_TAIL(&requests, req, tq);
-			send_read_request(conn->fd_out, req->id, req->offset,
+			send_read_request(conn, req->id, req->offset,
 			    req->len, handle, handle_len);
 		}
 
 		buffer_clear(&msg);
-		get_msg(conn->fd_in, &msg);
+		get_msg(conn, &msg);
 		type = buffer_get_char(&msg);
 		id = buffer_get_int(&msg);
 		debug3("Received reply T:%u I:%u R:%d", type, id, max_req);
@@ -1092,7 +1115,7 @@ do_download(struct sftp_conn *conn, char *remote_path, char *local_path,
 				req->id = conn->msg_id++;
 				req->len -= len;
 				req->offset += len;
-				send_read_request(conn->fd_out, req->id,
+				send_read_request(conn, req->id,
 				    req->offset, req->len, handle, handle_len);
 				/* Reduce the request size */
 				if (len < buflen)
@@ -1327,12 +1350,12 @@ do_upload(struct sftp_conn *conn, char *local_path, char *remote_path,
 	buffer_put_cstring(&msg, remote_path);
 	buffer_put_int(&msg, SSH2_FXF_WRITE|SSH2_FXF_CREAT|SSH2_FXF_TRUNC);
 	encode_attrib(&msg, &a);
-	send_msg(conn->fd_out, &msg);
+	send_msg(conn, &msg);
 	debug3("Sent message SSH2_FXP_OPEN I:%u P:%s", id, remote_path);
 
 	buffer_clear(&msg);
 
-	handle = get_handle(conn->fd_in, id, &handle_len,
+	handle = get_handle(conn, id, &handle_len,
 	    "remote open(\"%s\")", remote_path);
 	if (handle == NULL) {
 		close(local_fd);
@@ -1381,7 +1404,7 @@ do_upload(struct sftp_conn *conn, char *local_path, char *remote_path,
 			buffer_put_string(&msg, handle, handle_len);
 			buffer_put_int64(&msg, offset);
 			buffer_put_string(&msg, data, len);
-			send_msg(conn->fd_out, &msg);
+			send_msg(conn, &msg);
 			debug3("Sent message SSH2_FXP_WRITE I:%u O:%llu S:%u",
 			    id, (unsigned long long)offset, len);
 		} else if (TAILQ_FIRST(&acks) == NULL)
@@ -1395,7 +1418,7 @@ do_upload(struct sftp_conn *conn, char *local_path, char *remote_path,
 			u_int r_id;
 
 			buffer_clear(&msg);
-			get_msg(conn->fd_in, &msg);
+			get_msg(conn, &msg);
 			type = buffer_get_char(&msg);
 			r_id = buffer_get_int(&msg);
 
diff --git a/sftp-client.h b/sftp-client.h
index 1d08c404..145fc38e 100644
--- a/sftp-client.h
+++ b/sftp-client.h
@@ -1,4 +1,4 @@
-/* $OpenBSD: sftp-client.h,v 1.18 2009/08/18 18:36:20 djm Exp $ */
+/* $OpenBSD: sftp-client.h,v 1.19 2010/09/22 22:58:51 djm Exp $ */
 
 /*
  * Copyright (c) 2001-2004 Damien Miller <djm@openbsd.org>
@@ -51,7 +51,7 @@ struct sftp_statvfs {
  * Initialise a SSH filexfer connection. Returns NULL on error or
  * a pointer to a initialized sftp_conn struct on success.
  */
-struct sftp_conn *do_init(int, int, u_int, u_int);
+struct sftp_conn *do_init(int, int, u_int, u_int, u_int64_t);
 
 u_int sftp_proto_version(struct sftp_conn *);
 
diff --git a/sftp.1 b/sftp.1
index 49d88de0..b2a19b72 100644
--- a/sftp.1
+++ b/sftp.1
@@ -1,4 +1,4 @@
-.\" $OpenBSD: sftp.1,v 1.84 2010/09/19 21:30:05 jmc Exp $
+.\" $OpenBSD: sftp.1,v 1.85 2010/09/22 22:58:51 djm Exp $
 .\"
 .\" Copyright (c) 2001 Damien Miller.  All rights reserved.
 .\"
@@ -22,7 +22,7 @@
 .\" (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
 .\" THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 .\"
-.Dd $Mdocdate: September 19 2010 $
+.Dd $Mdocdate: September 22 2010 $
 .Dt SFTP 1
 .Os
 .Sh NAME
@@ -38,6 +38,7 @@
 .Op Fl D Ar sftp_server_path
 .Op Fl F Ar ssh_config
 .Op Fl i Ar identity_file
+.Op Fl l Ar limit
 .Op Fl o Ar ssh_option
 .Op Fl P Ar port
 .Op Fl R Ar num_requests
@@ -159,6 +160,8 @@ Selects the file from which the identity (private key) for public key
 authentication is read.
 This option is directly passed to
 .Xr ssh 1 .
+.It Fl l Ar limit
+Limits the used bandwidth, specified in Kbit/s.
 .It Fl o Ar ssh_option
 Can be used to pass options to
 .Nm ssh
diff --git a/sftp.c b/sftp.c
index 229f1298..8ce7d91f 100644
--- a/sftp.c
+++ b/sftp.c
@@ -1,4 +1,4 @@
-/* $OpenBSD: sftp.c,v 1.125 2010/06/18 00:58:39 djm Exp $ */
+/* $OpenBSD: sftp.c,v 1.126 2010/09/22 22:58:51 djm Exp $ */
 /*
  * Copyright (c) 2001-2004 Damien Miller <djm@openbsd.org>
  *
@@ -2073,6 +2073,7 @@ main(int argc, char **argv)
 	int debug_level = 0, sshver = 2;
 	char *file1 = NULL, *sftp_server = NULL;
 	char *ssh_program = _PATH_SSH_PROGRAM, *sftp_direct = NULL;
+	const char *errstr;
 	LogLevel ll = SYSLOG_LEVEL_INFO;
 	arglist args;
 	extern int optind;
@@ -2080,6 +2081,7 @@ main(int argc, char **argv)
 	struct sftp_conn *conn;
 	size_t copy_buffer_len = DEFAULT_COPY_BUFLEN;
 	size_t num_requests = DEFAULT_NUM_REQUESTS;
+	long long limit_kbps = 0;
 
 	/* Ensure that fds 0, 1 and 2 are open or directed to /dev/null */
 	sanitise_stdfd();
@@ -2097,7 +2099,7 @@ main(int argc, char **argv)
 	infile = stdin;
 
 	while ((ch = getopt(argc, argv,
-	    "1246hpqrvCc:D:i:o:s:S:b:B:F:P:R:")) != -1) {
+	    "1246hpqrvCc:D:i:l:o:s:S:b:B:F:P:R:")) != -1) {
 		switch (ch) {
 		/* Passed through to ssh(1) */
 		case '4':
@@ -2158,6 +2160,13 @@ main(int argc, char **argv)
 		case 'D':
 			sftp_direct = optarg;
 			break;
+		case 'l':
+			limit_kbps = strtonum(optarg, 1, 100 * 1024 * 1024,
+			    &errstr);
+			if (errstr != NULL)
+				usage();
+			limit_kbps *= 1024; /* kbps */
+			break;
 		case 'r':
 			global_rflag = 1;
 			break;
@@ -2235,7 +2244,7 @@ main(int argc, char **argv)
 	}
 	freeargs(&args);
 
-	conn = do_init(in, out, copy_buffer_len, num_requests);
+	conn = do_init(in, out, copy_buffer_len, num_requests, limit_kbps);
 	if (conn == NULL)
 		fatal("Couldn't initialise connection to server");