go-git progress notification output

clone-progress.go

package main

import (
	"bytes"
	"encoding/binary"
	"fmt"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"time"

	"gopkg.in/src-d/go-billy.v4/osfs"
	"gopkg.in/src-d/go-git.v4"
	"gopkg.in/src-d/go-git.v4/plumbing"
	"gopkg.in/src-d/go-git.v4/plumbing/format/pktline"
	"gopkg.in/src-d/go-git.v4/plumbing/transport"
	"gopkg.in/src-d/go-git.v4/plumbing/transport/client"
	httpproto "gopkg.in/src-d/go-git.v4/plumbing/transport/http"
	// gitproto "gopkg.in/src-d/go-git.v4/plumbing/transport/git"
	// sshproto "gopkg.in/src-d/go-git.v4/plumbing/transport/ssh"
	"gopkg.in/src-d/go-git.v4/storage"
	"gopkg.in/src-d/go-git.v4/storage/filesystem"
)

func main() {
	if len(os.Args) < 3 {
		fmt.Fprintln(os.Stderr, "Usage: ./clone-progress <url> <dir>")
		return
	}
	url, dir := os.Args[1], os.Args[2]

	// Use your own favorite progress bar library
	// (like https://github.com/cheggaaa/pb)
	// I recommend this one :)
	//   https://github.com/tyru/pgr
	pb := &progress{total: 0, current: 0}

	// Override http(s) default protocol to use our custom clients
	// https://github.com/src-d/go-git/blob/master/_examples/custom_http/main.go
	httpClient := getProgressClient(pb)
	client.InstallProtocol("http", httpClient)
	client.InstallProtocol("https", httpClient)
	// TODO
	// client.InstallProtocol("git", gitClient)
	// client.InstallProtocol("ssh", sshClient)

	// Filesystem abstraction
	fs := osfs.New(dir)
	// Git objects storer
	dot, err := fs.Chroot(".git")
	if err != nil {
		panic(err)
	}
	s, err := filesystem.NewStorage(dot)
	if err != nil {
		panic(err)
	}

	go func() {
		for {
			time.Sleep(500 * time.Millisecond)
			pb.Report()
		}
	}()

	_, err = git.Clone(&progressStorer{Storer: s, progress: pb}, fs, &git.CloneOptions{
		URL: url,
	})
	if err != nil {
		panic(err)
	}
}

func getProgressClient(pb *progress) transport.Transport {
	return httpproto.NewClient(&http.Client{
		Transport: &progressTransport{
			RoundTripper: http.DefaultTransport,
			progress:     pb,
		},
	})
}

type progress struct {
	total   uint32
	current uint32
}

func (p *progress) SetTotal(total uint32) {
	p.total = total
	// p.Report()    // slow
}

func (p *progress) Inc() {
	p.current++
	// p.Report()    // slow
}

func (p *progress) Report() {
	fmt.Print("\x1b\x5b2K\r")
	if p.total == 0 {
		fmt.Printf("(???) %d/???", p.current)
	} else {
		percent := int(float64(p.current) / float64(p.total) * 100)
		fmt.Printf("(%d%%) %d/%d", percent, p.current, p.total)
	}
}

type progressStorer struct {
	storage.Storer
	progress *progress
}

func (s *progressStorer) SetEncodedObject(o plumbing.EncodedObject) (plumbing.Hash, error) {
	hash, err := s.Storer.SetEncodedObject(o)
	s.progress.Inc()
	return hash, err
}

type progressTransport struct {
	http.RoundTripper
	progress *progress
}

func (t *progressTransport) RoundTrip(req *http.Request) (*http.Response, error) {
	res, err := t.RoundTripper.RoundTrip(req)
	if req.Method == "POST" && strings.HasSuffix(req.URL.String(), "/git-upload-pack") {
		if err := t.extractTotalInHeader(res); err != nil {
			return nil, err
		}
	}
	return res, err
}

func (t *progressTransport) extractTotalInHeader(res *http.Response) error {
	content, err := ioutil.ReadAll(res.Body)
	if err != nil {
		return err
	}
	res.Body.Close()
	res.Body = ioutil.NopCloser(bytes.NewBuffer(content))

	sc := pktline.NewScanner(bytes.NewBuffer(content))
	hi := 0
	var header [12]byte
	for sc.Scan() {
		b := sc.Bytes()
		if len(b) > 0 && b[0] == '\x01' {
			for i := 1; i < len(b) && hi < 12; i++ {
				header[hi] = b[i]
				hi++
			}
			if hi >= 12 {
				total := binary.BigEndian.Uint32(header[8:12])
				t.progress.SetTotal(total)
			}
		}
	}
	if sc.Err() != nil {
		return sc.Err()
	}
	return nil
}