From 97ed6e0711df9ffebac83b3aabfae8444ec9f6f6 Mon Sep 17 00:00:00 2001
From: rht <rudyht@gmail.com>
Date: Thu, 21 May 2015 11:32:38 +0700
Subject: [PATCH] Fix progress bar for "get" command

License: MIT
Signed-off-by: rht <rhtbot@gmail.com>
---
 core/commands/get.go       | 24 ++++++++++++++--------
 test/sharness/t0090-get.sh |  4 ++--
 unixfs/tar/reader.go       | 42 ++++++++++++++++++++++++++++++++++++++
 3 files changed, 60 insertions(+), 10 deletions(-)

diff --git a/core/commands/get.go b/core/commands/get.go
index 4c573928a28f..c29dd5720ca8 100644
--- a/core/commands/get.go
+++ b/core/commands/get.go
@@ -63,11 +63,12 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 			return
 		}
 
-		reader, err := get(req.Context().Context, node, req.Arguments()[0], cmplvl)
+		reader, length, err := get(req.Context().Context, node, req.Arguments()[0], cmplvl)
 		if err != nil {
 			res.SetError(err, cmds.ErrNormal)
 			return
 		}
+		res.SetLength(length)
 		res.SetOutput(reader)
 	},
 	PostRun: func(req cmds.Request, res cmds.Response) {
@@ -105,6 +106,8 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 			}
 			defer file.Close()
 
+			// bar length is set to 0 because there is no deterministic way to
+			// pre-determine the size of the .tar.gz output
 			bar := pb.New(0).SetUnits(pb.U_BYTES)
 			bar.Output = os.Stderr
 			pbReader := bar.NewProxyReader(outReader)
@@ -122,8 +125,7 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 
 		fmt.Printf("Saving file(s) to %s\n", outPath)
 
-		// TODO: get total length of files
-		bar := pb.New(0).SetUnits(pb.U_BYTES)
+		bar := pb.New64(int64(res.Length())).SetUnits(pb.U_BYTES)
 		bar.Output = os.Stderr
 
 		// wrap the reader with the progress bar proxy reader
@@ -145,8 +147,7 @@ may also specify the level of compression by specifying '-l=<1-9>'.
 		defer bar.Finish()
 
 		extractor := &tar.Extractor{outPath}
-		err = extractor.Extract(reader)
-		if err != nil {
+		if err = extractor.Extract(reader); err != nil {
 			res.SetError(err, cmds.ErrNormal)
 		}
 	},
@@ -166,12 +167,19 @@ func getCompressOptions(req cmds.Request) (int, error) {
 	return gzip.NoCompression, nil
 }
 
-func get(ctx context.Context, node *core.IpfsNode, p string, compression int) (io.Reader, error) {
+func get(ctx context.Context, node *core.IpfsNode, p string, compression int) (io.Reader, uint64, error) {
 	pathToResolve := path.Path(p)
 	dagnode, err := core.Resolve(ctx, node, pathToResolve)
 	if err != nil {
-		return nil, err
+		return nil, 0, err
+	}
+
+	reader, err := utar.NewReader(pathToResolve, node.DAG, dagnode, compression)
+
+	length, err1 := utar.GetTarSize(node.DAG, dagnode)
+	if err1 != nil {
+		return nil, 0, err1
 	}
 
-	return utar.NewReader(pathToResolve, node.DAG, dagnode, compression)
+	return reader, length, err
 }
diff --git a/test/sharness/t0090-get.sh b/test/sharness/t0090-get.sh
index 393ba5b5b4d2..8802b6728be3 100755
--- a/test/sharness/t0090-get.sh
+++ b/test/sharness/t0090-get.sh
@@ -23,7 +23,7 @@ test_get_cmd() {
 	'
 	
 	test_expect_success "ipfs get succeeds" '
-	  echo "Hello Worlds!" >data &&
+	  echo "Hello Worlds#" >data &&
 	  HASH=`ipfs add -q data` &&
 	  ipfs get "$HASH" >actual
 	'
@@ -80,7 +80,7 @@ test_get_cmd() {
 	  mkdir -p dir &&
 	  touch dir/a &&
 	  mkdir -p dir/b &&
-	  echo "Hello, Worlds!" >dir/b/c &&
+	  echo "Hello, Worlds#" >dir/b/c &&
 	  HASH2=`ipfs add -r -q dir | tail -n 1` &&
 	  ipfs get "$HASH2" >actual
 	'
diff --git a/unixfs/tar/reader.go b/unixfs/tar/reader.go
index 20e18fe114d1..c1b11c9ebcbd 100644
--- a/unixfs/tar/reader.go
+++ b/unixfs/tar/reader.go
@@ -17,6 +17,8 @@ import (
 	upb "github.com/ipfs/go-ipfs/unixfs/pb"
 )
 
+const tarBlockSize = 512
+
 type Reader struct {
 	buf        bytes.Buffer
 	closed     bool
@@ -202,3 +204,43 @@ func (r *Reader) syncCopy(reader io.Reader) error {
 	}
 	return nil
 }
+
+func GetTarSize(dag mdag.DAGService, dagnode *mdag.Node) (uint64, error) {
+	return _GetTarSize(dag, dagnode, true)
+}
+
+func _GetTarSize(dag mdag.DAGService, dagnode *mdag.Node, isparent bool) (uint64, error) {
+	size := uint64(0)
+
+	pb := new(upb.Data)
+	if err := proto.Unmarshal(dagnode.Data, pb); err != nil {
+		return 0, err
+	}
+
+	if pb.GetType() == upb.Data_Directory {
+		//size += tarBlockSize
+
+		ctx, cancel := context.WithTimeout(context.TODO(), time.Second*60)
+		defer cancel()
+		for _, ng := range dag.GetDAG(ctx, dagnode) {
+			childNode, err := ng.Get(ctx)
+			if err != nil {
+				return 0, err
+			}
+			childSize, err := _GetTarSize(dag, childNode, false)
+			if err != nil {
+				return 0, err
+			}
+			size += childSize
+		}
+	} else {
+		unixSize := pb.GetFilesize()
+		// tar header + file size + round up to nearest 512 bytes
+		size = tarBlockSize + unixSize + (tarBlockSize - unixSize%tarBlockSize)
+	}
+
+	if isparent {
+		size += 2 * tarBlockSize // tar parent padding
+	}
+	return size, nil
+}