Merge pull request #432 from mbattista/master

added throttle upload feature

go.mod

@@ -33,6 +33,7 @@ require (
 	golang.org/x/sys v0.0.0-20211002104244-808efd93c36d // indirect
 	golang.org/x/term v0.0.0-20210927222741-03fcf44c2211 // indirect
 	golang.org/x/text v0.3.7 // indirect
+	golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac // indirect
 	gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 // indirect
 	gopkg.in/yaml.v3 v3.0.0-20200313102051-9f266ea9e77c // indirect
 )

go.sum

@@ -80,6 +80,8 @@ golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.6/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7 h1:olpwvP2KacW1ZWvsR7uQhoyTYvKAupfQrRGBFM352Gk=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
+golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac h1:7zkz7BUtwNFFqcowJ+RIgu2MaV/MapERkDIy+mwPyjs=
+golang.org/x/time v0.0.0-20210723032227-1f47c861a9ac/go.mod h1:tRJNPiyCQ0inRvYxbN9jk5I+vvW/OXSQhTDSoE431IQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=
 gopkg.in/check.v1 v0.0.0-20161208181325-20d25e280405/go.mod h1:Co6ibVJAznAaIkqp8huTwlJQCZ016jof/cbN4VW5Yz0=
 gopkg.in/check.v1 v1.0.0-20190902080502-41f04d3bba15 h1:YR8cESwS4TdDjEe65xsg0ogRM/Nc3DYOhEAlW+xobZo=

src/cli/cli.go

@@ -101,6 +101,7 @@ func Run() (err error) {
 		&cli.StringFlag{Name: "out", Value: ".", Usage: "specify an output folder to receive the file"},
 		&cli.StringFlag{Name: "pass", Value: models.DEFAULT_PASSPHRASE, Usage: "password for the relay", EnvVars: []string{"CROC_PASS"}},
 		&cli.StringFlag{Name: "socks5", Value: "", Usage: "add a socks5 proxy", EnvVars: []string{"SOCKS5_PROXY"}},
+		&cli.StringFlag{Name: "throttleUpload", Value: "", Usage: "Throttle the upload speed e.g. 500k"},
 	}
 	app.EnableBashCompletion = true
 	app.HideHelp = false
@@ -206,6 +207,7 @@ func send(c *cli.Context) (err error) {
 		Overwrite:      c.Bool("overwrite"),
 		Curve:          c.String("curve"),
 		HashAlgorithm:  c.String("hash"),
+		ThrottleUpload: c.String("throttleUpload"),
 	}
 	if crocOptions.RelayAddress != models.DEFAULT_RELAY {
 		crocOptions.RelayAddress6 = ""

src/croc/croc.go

@@ -16,6 +16,7 @@ import (
 	"strings"
 	"sync"
 	"time"
+	"golang.org/x/time/rate"
 
 	"github.com/denisbrodbeck/machineid"
 	log "github.com/schollz/logger"
@@ -72,6 +73,7 @@ type Options struct {
 	Overwrite      bool
 	Curve          string
 	HashAlgorithm  string
+	ThrottleUpload string
 }
 
 // Client holds the state of the croc transfer
@@ -104,6 +106,7 @@ type Client struct {
 	TotalSent             int64
 	TotalChunksTransfered int
 	chunkMap              map[uint64]struct{}
+	limiter               *rate.Limiter
 
 	// tcp connections
 	conn []*comm.Comm
@@ -174,6 +177,37 @@ func New(ops Options) (c *Client, err error) {
 
 	c.conn = make([]*comm.Comm, 16)
 
+	// initialize throttler
+	if len(c.Options.ThrottleUpload) > 1 && c.Options.IsSender {
+		upload := c.Options.ThrottleUpload[:len(c.Options.ThrottleUpload)-1]
+		uploadLimit, err := strconv.ParseInt(upload, 10, 64)
+		if err != nil {
+			panic("Could not parse given Upload Limit")
+		}
+		minBurstSize := models.TCP_BUFFER_SIZE
+		var rt rate.Limit
+		switch unit := string(c.Options.ThrottleUpload[len(c.Options.ThrottleUpload)-1:]); unit {
+		case "g", "G":
+			uploadLimit = uploadLimit*1024*1024*1024
+		case "m", "M":
+			uploadLimit = uploadLimit*1024*1024
+		case "k", "K":
+			uploadLimit = uploadLimit*1024
+		default:
+			uploadLimit, err = strconv.ParseInt(c.Options.ThrottleUpload, 10, 64)
+			if err != nil {
+				panic("Could not parse given Upload Limit")
+			}
+		}
+		// Somehow 4* is neccessary
+		rt = rate.Every(time.Second / (4*time.Duration(uploadLimit)))
+		if (int(uploadLimit) > minBurstSize) {
+			minBurstSize = int(uploadLimit)
+		}
+		c.limiter = rate.NewLimiter(rt, minBurstSize)
+		log.Debugf("Throttling Upload to %#v", c.limiter.Limit())
+	}
+
 	// initialize pake for recipient
 	if !c.Options.IsSender {
 		c.Pake, err = pake.InitCurve([]byte(c.Options.SharedSecret[5:]), 0, c.Options.Curve)
@@ -1518,6 +1552,11 @@ func (c *Client) sendData(i int) {
 		n, errRead := c.fread.ReadAt(data, readingPos)
 		// log.Debugf("%d read %d bytes", i, n)
 		readingPos += int64(n)
+		if (c.limiter != nil) {
+			r := c.limiter.ReserveN(time.Now(), n)
+			log.Debugf("Limiting Upload for %d", r.Delay())
+			time.Sleep(r.Delay())
+		}
 
 		if math.Mod(curi, float64(len(c.Options.RelayPorts))) == float64(i) {
 			// check to see if this is a chunk that the recipient wants