mirrors_public/schollz_croc
0
0
Fork 0
mirror of https://github.com/schollz/croc.git synced 2025-10-11 13:21:00 +02:00
Code Activity

added throttle upload feature

Browse source
This commit is contained in:
Marcel Battista 2021-10-19 02:46:54 +02:00
parent 696479024e
commit 8ff3fb9b89
4 changed files with 38 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

33
src/croc/croc.go
View file

@ -17,6 +17,7 @@ import (
"strings"
"sync"
"time"
"golang.org/x/time/rate"
"github.com/denisbrodbeck/machineid"
log "github.com/schollz/logger"
@ -73,6 +74,7 @@ type Options struct {
Overwrite bool
Curve string
HashAlgorithm string
ThrottleUpload string
}
// Client holds the state of the croc transfer
@ -105,6 +107,7 @@ type Client struct {
TotalSent int64
TotalChunksTransfered int
chunkMap map[uint64]struct{}
limiter *rate.Limiter
// tcp connections
conn []*comm.Comm
@ -175,6 +178,31 @@ func New(ops Options) (c *Client, err error) {
c.conn = make([]*comm.Comm, 16)
if len(c.Options.ThrottleUpload) > 0 && 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
}
// 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)
@ -1519,6 +1547,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