Merge branch 'poldi1405-master'

Author: melbahja

cmd/got/main.go

@@ -3,6 +3,7 @@ package main
 import (
 	"bufio"
 	"context"
+	"errors"
 	"fmt"
 	"log"
 	"net/url"
@@ -21,6 +22,8 @@ import (
 
 var version string
 
+var HeaderSlice []got.GotHeader
+
 func main() {
 
 	// New context.
@@ -67,6 +70,11 @@ func main() {
 				Usage:   "Chunks that will be downloaded concurrently.",
 				Aliases: []string{"c"},
 			},
+			&cli.StringSliceFlag{
+				Name:    "header",
+				Usage:   "Set these HTTP-Headers on the requests. The format has to be Key: Value",
+				Aliases: []string{"H"},
+			},
 		},
 		Version: version,
 		Authors: []*cli.Author{
@@ -100,20 +108,24 @@ func run(ctx context.Context, c *cli.Context) error {
 
 		// 55 is just an estimation of the text showed with the progress.
 		// it's working fine with $COLUMNS >= 47
-		// TODO: hide progress bar on terminal size of $COLUMNS <= 46
 		p.Width = getWidth() - 55
 
 		perc, err := progress.GetPercentage(float64(d.Size()), float64(d.TotalSize()))
 		if err != nil {
 			perc = 100
 		}
 
+		var bar string
+		if getWidth() <= 46 {
+			bar = ""
+		} else {
+			bar = r + color(p.GetBar(perc, 100)) + l
+		}
+
 		fmt.Printf(
-			" %6.2f%% %s%s%s %s/%s @ %s/s%s\r",
+			" %6.2f%% %s %s/%s @ %s/s%s\r",
 			perc,
-			r,
-			color(p.GetBar(perc, 100)),
-			l,
+			bar,
 			humanize.Bytes(d.Size()),
 			humanize.Bytes(d.TotalSize()),
 			humanize.Bytes(d.Speed()),
@@ -157,6 +169,18 @@ func run(ctx context.Context, c *cli.Context) error {
 		}
 	}
 
+	if c.StringSlice("header") != nil {
+		header := c.StringSlice("header")
+
+		for _, h := range header {
+			split := strings.SplitN(h, ":", 2)
+			if len(split) == 1 {
+				return errors.New("malformatted header " + h)
+			}
+			HeaderSlice = append(HeaderSlice, got.GotHeader{Key: split[0], Value: strings.TrimSpace(split[1])})
+		}
+	}
+
 	// Download from args.
 	for _, url := range c.Args().Slice() {
 
@@ -211,6 +235,7 @@ func download(ctx context.Context, c *cli.Context, g *got.Got, url string) (err
 		URL:         url,
 		Dir:         c.String("dir"),
 		Dest:        c.String("output"),
+		Header:      HeaderSlice,
 		Interval:    150,
 		ChunkSize:   c.Uint64("size"),
 		Concurrency: c.Uint("concurrency"),

download.go

@@ -38,6 +38,8 @@ type (
 
 		Interval, ChunkSize, MinChunkSize, MaxChunkSize uint64
 
+		Header []GotHeader
+
 		StopProgress bool
 
 		ctx context.Context
@@ -52,12 +54,17 @@ type (
 
 		startedAt time.Time
 	}
+
+	GotHeader struct {
+		Key   string
+		Value string
+	}
 )
 
 // GetInfo returns URL info, and error if any.
 func (d Download) GetInfo() (*Info, error) {
 
-	req, err := NewRequest(d.ctx, "HEAD", d.URL)
+	req, err := NewRequest(d.ctx, "HEAD", d.URL, d.Header)
 
 	if err != nil {
 		return nil, err
@@ -86,7 +93,7 @@ func (d *Download) getInfoFromGetRequest() (*Info, error) {
 		res *http.Response
 	)
 
-	if req, err = NewRequest(d.ctx, "GET", d.URL); err != nil {
+	if req, err = NewRequest(d.ctx, "GET", d.URL, d.Header); err != nil {
 		return nil, err
 	}
 
@@ -473,7 +480,7 @@ func (d *Download) DownloadChunk(c Chunk, dest *os.File) error {
 		res *http.Response
 	)
 
-	if req, err = NewRequest(d.ctx, "GET", d.URL); err != nil {
+	if req, err = NewRequest(d.ctx, "GET", d.URL, d.Header); err != nil {
 		return err
 	}
 

download_test.go

@@ -31,6 +31,7 @@ func TestGetInfoAndInit(t *testing.T) {
 	t.Run("getInfoTest", getInfoTest)
 	t.Run("okInitTest", okInitTest)
 	t.Run("errInitTest", errInitTest)
+	t.Run("sendHeadersTest", sendHeadersTest)
 }
 
 func TestDownloading(t *testing.T) {
@@ -68,6 +69,35 @@ func getInfoTest(t *testing.T) {
 	}
 }
 
+func sendHeadersTest(t *testing.T) {
+
+	tmpFile := createTemp()
+	defer clean(tmpFile)
+
+	dl := got.NewDownload(context.Background(), httpt.URL+"/header_values", tmpFile)
+	dl.Header = []got.GotHeader{
+		{
+			Key:   "x-test-header",
+			Value: "foobar",
+		},
+	}
+
+	info, err := dl.GetInfo()
+
+	if err != nil {
+		t.Error(err)
+		return
+	}
+
+	if info.Rangeable == false {
+		t.Error("rangeable should be true")
+	}
+
+	if info.Size != uint64(okFileStat.Size()) {
+		t.Errorf("Invalid file size, wants %d but got %d", okFileStat.Size(), info.Size)
+	}
+}
+
 func getFilenameTest(t *testing.T) {
 
 	tmpFile := createTemp()

got.go

@@ -76,12 +76,17 @@ func NewWithContext(ctx context.Context) *Got {
 }
 
 // NewRequest returns a new http.Request and error if any.
-func NewRequest(ctx context.Context, method, URL string) (req *http.Request, err error) {
+func NewRequest(ctx context.Context, method, URL string, header []GotHeader) (req *http.Request, err error) {
 
 	if req, err = http.NewRequestWithContext(ctx, method, URL, nil); err != nil {
 		return
 	}
 
 	req.Header.Set("User-Agent", DefaulUserAgent)
+
+	for _, h := range header {
+		req.Header.Set(h.Key, h.Value)
+	}
+
 	return
 }

got_test.go

@@ -53,6 +53,15 @@ func NewHttptestServer() *httptest.Server {
 			http.ServeFile(w, r, "go.mod")
 			return
 
+		case "/header_values":
+			if r.Header.Get("x-test-header") == "foobar" {
+				http.ServeFile(w, r, "go.mod")
+				return
+			}
+
+			w.WriteHeader(403)
+			return
+
 		case "/not_found":
 		}