From 44d3ddf9287a61fa60501fb55911f314c13d882f Mon Sep 17 00:00:00 2001
From: Vitaliy Filippov <vitalif@yourcmc.ru>
Date: Sun, 15 May 2016 23:05:25 +0300
Subject: [PATCH] Implement upload/download speed limit (fix #60)

---
 README.md                          |  1 +
 grive/src/main.cc                  |  7 +++++++
 libgrive/src/http/Agent.cc         | 10 ++++++++++
 libgrive/src/http/Agent.hh         |  6 ++++++
 libgrive/src/http/CurlAgent.cc     |  4 ++++
 libgrive/src/protocol/AuthAgent.cc | 10 ++++++++++
 libgrive/src/protocol/AuthAgent.hh |  3 +++
 7 files changed, 41 insertions(+)

diff --git a/README.md b/README.md
index e5ec3e9..d6407c6 100644
--- a/README.md
+++ b/README.md
@@ -75,6 +75,7 @@ Enjoy!
 - no-remote-new and upload-only modes
 - ignore regexp does not persist anymore (note that Grive will still track it to not
   accidentally delete remote files when changing ignore regexp)
+- added options to limit upload and download speed
 
 ### Grive2 v0.5
 
diff --git a/grive/src/main.cc b/grive/src/main.cc
index 6038313..281245b 100644
--- a/grive/src/main.cc
+++ b/grive/src/main.cc
@@ -123,6 +123,8 @@ int Main( int argc, char **argv )
 		( "dry-run",	"Only detect which files need to be uploaded/downloaded, "
 						"without actually performing them." )
 		( "ignore",		po::value<std::string>(), "Perl RegExp to ignore files (matched against relative paths, remembered for next runs)." )
+		( "upload-speed,U", po::value<unsigned>(), "Limit upload speed in kbytes per second" )
+		( "download-speed,D", po::value<unsigned>(), "Limit download speed in kbytes per second" )
 	;
 	
 	po::variables_map vm;
@@ -195,6 +197,11 @@ int Main( int argc, char **argv )
 	AuthAgent agent( token, http.get() ) ;
 	v2::Syncer2 syncer( &agent );
 
+	if ( vm.count( "upload-speed" ) > 0 )
+		agent.SetUploadSpeed( vm["upload-speed"].as<unsigned>() * 1000 );
+	if ( vm.count( "download-speed" ) > 0 )
+		agent.SetDownloadSpeed( vm["download-speed"].as<unsigned>() * 1000 );
+
 	Drive drive( &syncer, config.GetAll() ) ;
 	drive.DetectChanges() ;
 
diff --git a/libgrive/src/http/Agent.cc b/libgrive/src/http/Agent.cc
index b9143c6..ba04846 100644
--- a/libgrive/src/http/Agent.cc
+++ b/libgrive/src/http/Agent.cc
@@ -64,4 +64,14 @@ long Agent::Post(
 	return Request( "POST", url, &s, dest, h );
 }
 
+void Agent::SetUploadSpeed( unsigned kbytes )
+{
+	mMaxUpload = kbytes;
+}
+
+void Agent::SetDownloadSpeed( unsigned kbytes )
+{
+	mMaxDownload = kbytes;
+}
+
 } } // end of namespace
diff --git a/libgrive/src/http/Agent.hh b/libgrive/src/http/Agent.hh
index cb8594a..ce1883b 100644
--- a/libgrive/src/http/Agent.hh
+++ b/libgrive/src/http/Agent.hh
@@ -34,6 +34,9 @@ class Header ;
 
 class Agent
 {
+protected:
+	unsigned mMaxUpload, mMaxDownload ;
+
 public :
 	virtual ~Agent() {}
 
@@ -70,6 +73,9 @@ public :
 		DataStream			*dest,
 		const Header&		hdr ) = 0 ;
 	
+	virtual void SetUploadSpeed( unsigned kbytes ) ;
+	virtual void SetDownloadSpeed( unsigned kbytes ) ;
+	
 	virtual std::string LastError() const = 0 ;
 	virtual std::string LastErrorHeaders() const = 0 ;
 	
diff --git a/libgrive/src/http/CurlAgent.cc b/libgrive/src/http/CurlAgent.cc
index 4b23067..3a27955 100644
--- a/libgrive/src/http/CurlAgent.cc
+++ b/libgrive/src/http/CurlAgent.cc
@@ -86,6 +86,10 @@ void CurlAgent::Init()
 	::curl_easy_setopt( m_pimpl->curl, CURLOPT_HEADERFUNCTION,	&CurlAgent::HeaderCallback ) ;
 	::curl_easy_setopt( m_pimpl->curl, CURLOPT_HEADERDATA,		this ) ;
 	::curl_easy_setopt( m_pimpl->curl, CURLOPT_HEADER,			0L ) ;
+	if ( mMaxUpload > 0 )
+		::curl_easy_setopt( m_pimpl->curl, CURLOPT_MAX_SEND_SPEED_LARGE, mMaxUpload ) ;
+	if ( mMaxDownload > 0 )
+		::curl_easy_setopt( m_pimpl->curl, CURLOPT_MAX_RECV_SPEED_LARGE, mMaxDownload ) ;
 	m_pimpl->error = false;
 	m_pimpl->error_headers = "";
 	m_pimpl->error_data = "";
diff --git a/libgrive/src/protocol/AuthAgent.cc b/libgrive/src/protocol/AuthAgent.cc
index 77cb54f..65665f7 100644
--- a/libgrive/src/protocol/AuthAgent.cc
+++ b/libgrive/src/protocol/AuthAgent.cc
@@ -47,6 +47,16 @@ void AuthAgent::SetLog( http::ResponseLog *log )
 	return m_agent->SetLog( log );
 }
 
+void AuthAgent::SetUploadSpeed( unsigned kbytes )
+{
+	m_agent->SetUploadSpeed( kbytes );
+}
+
+void AuthAgent::SetDownloadSpeed( unsigned kbytes )
+{
+	m_agent->SetDownloadSpeed( kbytes );
+}
+
 http::Header AuthAgent::AppendHeader( const http::Header& hdr ) const
 {
 	http::Header h(hdr) ;
diff --git a/libgrive/src/protocol/AuthAgent.hh b/libgrive/src/protocol/AuthAgent.hh
index 3dd3402..7f16dc1 100644
--- a/libgrive/src/protocol/AuthAgent.hh
+++ b/libgrive/src/protocol/AuthAgent.hh
@@ -54,6 +54,9 @@ public :
 	std::string Escape( const std::string& str ) ;
 	std::string Unescape( const std::string& str ) ;
 
+	void SetUploadSpeed( unsigned kbytes ) ;
+	void SetDownloadSpeed( unsigned kbytes ) ;
+
 private :
 	http::Header AppendHeader( const http::Header& hdr ) const ;
 	bool CheckRetry( long response ) ;