From 6a7912e37ad587ad71500d06f3db3c1e44ca5b35 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@tailscale.com>
Date: Thu, 8 Apr 2021 14:54:25 -0700
Subject: [PATCH] cmd/tailscale/cli: add debug mode to push slowly for testing

Also set Content-Length when known, and fail explicitly on sending
directories for now.

Signed-off-by: Brad Fitzpatrick <bradfitz@tailscale.com>
---
 cmd/tailscale/cli/push.go | 37 ++++++++++++++++++++++++++++++++++++-
 1 file changed, 36 insertions(+), 1 deletion(-)

diff --git a/cmd/tailscale/cli/push.go b/cmd/tailscale/cli/push.go
index 98fefe846..1eac0fe96 100644
--- a/cmd/tailscale/cli/push.go
+++ b/cmd/tailscale/cli/push.go
@@ -17,10 +17,12 @@ import (
 	"net/http"
 	"net/url"
 	"os"
+	"strconv"
 	"time"
 	"unicode/utf8"
 
 	"github.com/peterbourgon/ff/v2/ffcli"
+	"golang.org/x/time/rate"
 	"tailscale.com/ipn"
 	"tailscale.com/ipn/ipnstate"
 )
@@ -62,6 +64,7 @@ func runPush(ctx context.Context, args []string) error {
 
 	var fileContents io.Reader
 	var name = pushArgs.name
+	var contentLength int64 = -1
 	if fileArg == "-" {
 		fileContents = os.Stdin
 		if name == "" {
@@ -76,10 +79,22 @@ func runPush(ctx context.Context, args []string) error {
 			return err
 		}
 		defer f.Close()
-		fileContents = f
+		fi, err := f.Stat()
+		if err != nil {
+			return err
+		}
+		if fi.IsDir() {
+			return errors.New("directories not supported")
+		}
+		contentLength = fi.Size()
+		fileContents = io.LimitReader(f, contentLength)
 		if name == "" {
 			name = fileArg
 		}
+
+		if slow, _ := strconv.ParseBool(os.Getenv("TS_DEBUG_SLOW_PUSH")); slow {
+			fileContents = &slowReader{r: fileContents}
+		}
 	}
 
 	dstURL := "http://" + net.JoinHostPort(ip, fmt.Sprint(peerAPIPort)) + "/v0/put/" + url.PathEscape(name)
@@ -87,6 +102,7 @@ func runPush(ctx context.Context, args []string) error {
 	if err != nil {
 		return err
 	}
+	req.ContentLength = contentLength
 	if pushArgs.verbose {
 		log.Printf("sending to %v ...", dstURL)
 	}
@@ -171,3 +187,22 @@ func pickStdinFilename() (name string, r io.Reader, err error) {
 	}
 	return "stdin" + ext(sniff), io.MultiReader(bytes.NewReader(sniff), os.Stdin), nil
 }
+
+type slowReader struct {
+	r  io.Reader
+	rl *rate.Limiter
+}
+
+func (r *slowReader) Read(p []byte) (n int, err error) {
+	const burst = 4 << 10
+	plen := len(p)
+	if plen > burst {
+		plen = burst
+	}
+	if r.rl == nil {
+		r.rl = rate.NewLimiter(rate.Limit(1<<10), burst)
+	}
+	n, err = r.r.Read(p[:plen])
+	r.rl.WaitN(context.Background(), n)
+	return
+}