Merge pull request #5 from abzcoding/bw-limit

added bandwidth limiting to hget

Author: abzcoding

Makefile

@@ -9,6 +9,8 @@ deps:
 	go get -d gopkg.in/cheggaaa/pb.v1
 	go get -d github.com/mattn/go-isatty
 	go get -d github.com/imkira/go-task
+	go get -d github.com/fujiwara/shapeio
+	go get -d github.com/alecthomas/units
 
 clean:
 	@echo "====> Remove installed binary"

README.md

@@ -17,12 +17,13 @@ Binary file will be built at ./bin/hget, you can copy to /usr/bin or /usr/local/
 ## Usage
 
 ```bash
-hget [-n parallel] [-skip-tls false] [-proxy proxy_server] [-file filename] [URL] # to download url, with n connections, and not skip tls certificate
+hget [-n parallel] [-skip-tls false] [-rate bwRate] [-proxy proxy_server] [-file filename] [URL] # to download url, with n connections, and not skip tls certificate
 hget tasks # get interrupted tasks
 hget resume [TaskName | URL] # to resume task
 hget -proxy "127.0.0.1:12345" URL # to download using socks5 proxy
 hget -proxy "http://sample-proxy.com:8080" URL # to download using http proxy
 hget -file sample.txt # to download a list of files
+hget -n 4 -rate 100KB URL # to download using 4 threads & limited to 100Kb per second
 ```
 
 ### Help
@@ -37,6 +38,10 @@ Usage of hget:
         proxy for downloading, ex
                 -proxy '127.0.0.1:12345' for socks5 proxy
                 -proxy 'http://proxy.com:8080' for http proxy
+  -rate string
+        bandwidth limit to use while downloading, ex
+                -rate 10kB
+                -rate 10MiB
   -skip-tls
         skip verify certificate for https (default true)
 ```

http.go

@@ -13,7 +13,9 @@ import (
 	"strings"
 	"sync"
 
+	"github.com/alecthomas/units"
 	"github.com/fatih/color"
+	"github.com/fujiwara/shapeio"
 	"golang.org/x/net/proxy"
 	pb "gopkg.in/cheggaaa/pb.v1"
 )
@@ -30,8 +32,10 @@ var (
 	contentLengthHeader = "Content-Length"
 )
 
+// HttpDownloader holds the required configurations
 type HttpDownloader struct {
 	proxy     string
+	rate      int64
 	url       string
 	file      string
 	par       int64
@@ -42,11 +46,11 @@ type HttpDownloader struct {
 	resumable bool
 }
 
-func NewHttpDownloader(url string, par int, skipTls bool, socks5_proxy string) *HttpDownloader {
+func NewHttpDownloader(url string, par int, skipTls bool, proxy_server string, bwLimit string) *HttpDownloader {
 	var resumable = true
-	
-	client := ProxyAwareHttpClient(socks5_proxy)
-	
+
+	client := ProxyAwareHttpClient(proxy_server)
+
 	parsed, err := stdurl.Parse(url)
 	FatalCheck(err)
 
@@ -94,6 +98,12 @@ func NewHttpDownloader(url string, par int, skipTls bool, socks5_proxy string) *
 
 	file := filepath.Base(url)
 	ret := new(HttpDownloader)
+	ret.rate = 0
+	bandwidthLimit, err := units.ParseStrictBytes(bwLimit)
+	if err == nil {
+		ret.rate = bandwidthLimit
+		Printf("Download with bandwidth limit set to %s[%d]\n", bwLimit, ret.rate)
+	}
 	ret.url = url
 	ret.file = file
 	ret.par = int64(par)
@@ -102,7 +112,7 @@ func NewHttpDownloader(url string, par int, skipTls bool, socks5_proxy string) *
 	ret.skipTls = skipTls
 	ret.parts = partCalculate(int64(par), len, url)
 	ret.resumable = resumable
-	ret.proxy = socks5_proxy
+	ret.proxy = proxy_server
 
 	return ret
 }
@@ -135,16 +145,16 @@ func partCalculate(par int64, len int64, url string) []Part {
 	return ret
 }
 
-func ProxyAwareHttpClient(socks5_proxy string) *http.Client {
+func ProxyAwareHttpClient(proxy_server string) *http.Client {
 	// setup a http client
 	httpTransport := &http.Transport{}
 	httpClient := &http.Client{Transport: httpTransport}
 	var dialer proxy.Dialer
 	dialer = proxy.Direct
-	
-	if len(socks5_proxy) > 0 {
-		if strings.HasPrefix(socks5_proxy, "http") {
-			proxyUrl, err := stdurl.Parse(socks5_proxy)
+
+	if len(proxy_server) > 0 {
+		if strings.HasPrefix(proxy_server, "http") {
+			proxyUrl, err := stdurl.Parse(proxy_server)
 			if err != nil {
 				fmt.Fprintln(os.Stderr, "invalid proxy: ", err)
 			}
@@ -155,12 +165,12 @@ func ProxyAwareHttpClient(socks5_proxy string) *http.Client {
 			}
 		} else {
 			// create a socks5 dialer
-			dialer, err := proxy.SOCKS5("tcp", socks5_proxy, nil, proxy.Direct)
+			dialer, err := proxy.SOCKS5("tcp", proxy_server, nil, proxy.Direct)
 			if err == nil {
 				httpTransport.Dial = dialer.Dial
 			}
 		}
-		
+
 	}
 	return httpClient
 }
@@ -247,7 +257,14 @@ func (d *HttpDownloader) Do(doneChan chan bool, fileChan chan string, errorChan
 			finishDownloadChan := make(chan bool)
 
 			go func() {
-				written, _ := io.Copy(writer, resp.Body)
+				var written int64
+				if d.rate != 0 {
+					reader := shapeio.NewReader(resp.Body)
+					reader.SetRateLimit(float64(d.rate))
+					written, _ = io.Copy(writer, reader)
+				} else {
+					written, _ = io.Copy(writer, resp.Body)
+				}
 				current += written
 				fileChan <- part.Path
 				finishDownloadChan <- true

main.go

@@ -17,12 +17,13 @@ var displayProgress = true
 
 func main() {
 	var err error
-	var proxy, filepath string
+	var proxy, filepath, bwLimit string
 
 	conn := flag.Int("n", runtime.NumCPU(), "connection")
 	skiptls := flag.Bool("skip-tls", true, "skip verify certificate for https")
 	flag.StringVar(&proxy, "proxy", "", "proxy for downloading, ex \n\t-proxy '127.0.0.1:12345' for socks5 proxy\n\t-proxy 'http://proxy.com:8080' for http proxy")
 	flag.StringVar(&filepath, "file", "", "filepath that contains links in each line")
+	flag.StringVar(&bwLimit, "rate", "", "bandwidth limit to use while downloading, ex\n\t -rate 10kB\n\t-rate 10MiB")
 
 	flag.Parse()
 	args := flag.Args()
@@ -46,7 +47,7 @@ func main() {
 					break
 				}
 
-				g1.AddChild(downloadTask(string(line), nil, *conn, *skiptls, proxy))
+				g1.AddChild(downloadTask(string(line), nil, *conn, *skiptls, proxy, bwLimit))
 			}
 			g1.Run(nil)
 			return
@@ -79,26 +80,26 @@ func main() {
 
 		state, err := Resume(task)
 		FatalCheck(err)
-		Execute(state.Url, state, *conn, *skiptls, proxy)
+		Execute(state.Url, state, *conn, *skiptls, proxy, bwLimit)
 		return
 	} else {
 		if ExistDir(FolderOf(command)) {
 			Warnf("Downloading task already exist, remove first \n")
 			err := os.RemoveAll(FolderOf(command))
 			FatalCheck(err)
 		}
-		Execute(command, nil, *conn, *skiptls, proxy)
+		Execute(command, nil, *conn, *skiptls, proxy, bwLimit)
 	}
 }
 
-func downloadTask(url string, state *State, conn int, skiptls bool, proxy string) task.Task {
+func downloadTask(url string, state *State, conn int, skiptls bool, proxy string, bwLimit string) task.Task {
 	run := func(t task.Task, ctx task.Context) {
-		Execute(url, state, conn, skiptls, proxy)
+		Execute(url, state, conn, skiptls, proxy, bwLimit)
 	}
 	return task.NewTaskWithFunc(run)
 }
 
-func Execute(url string, state *State, conn int, skiptls bool, proxy string) {
+func Execute(url string, state *State, conn int, skiptls bool, proxy string, bwLimit string) {
 	//otherwise is hget <URL> command
 
 	signal_chan := make(chan os.Signal, 1)
@@ -122,7 +123,7 @@ func Execute(url string, state *State, conn int, skiptls bool, proxy string) {
 
 	var downloader *HttpDownloader
 	if state == nil {
-		downloader = NewHttpDownloader(url, conn, skiptls, proxy)
+		downloader = NewHttpDownloader(url, conn, skiptls, proxy, bwLimit)
 	} else {
 		downloader = &HttpDownloader{url: state.Url, file: filepath.Base(state.Url), par: int64(len(state.Parts)), parts: state.Parts, resumable: true}
 	}