From 0c82c7e97a74888f8e96341615630a662bca647a Mon Sep 17 00:00:00 2001
From: lhchavez <lhchavez@lhchavez.com>
Date: Mon, 13 Dec 2021 07:35:30 -0800
Subject: [PATCH] Add a trailer to write the uncompressed archive size

This change allows the caller to know what the uncompressed archive size
is.
---
 browser.go      |  4 ++++
 browser_test.go | 19 ++++++++++++++++++-
 2 files changed, 22 insertions(+), 1 deletion(-)

diff --git a/browser.go b/browser.go
index 8ada24c..eefe84d 100644
--- a/browser.go
+++ b/browser.go
@@ -621,6 +621,7 @@ func handleArchive(
 	}
 
 	w.Header().Set("Content-Type", contentType)
+	w.Header().Set("Trailer", "Omegaup-Uncompressed-Size")
 	var z archive
 	if contentType == "application/gzip" {
 		gz := gzip.NewWriter(w)
@@ -632,6 +633,7 @@ func handleArchive(
 	}
 	defer z.Close()
 
+	var uncompressedSize int64
 	err = tree.Walk(func(parent string, entry *git.TreeEntry) error {
 		select {
 		case <-ctx.Done():
@@ -664,6 +666,7 @@ func handleArchive(
 		defer blob.Free()
 
 		// Object is a blob.
+		uncompressedSize += blob.Size()
 		w, err := z.Create(fullPath, blob.Size())
 		if err != nil {
 			return errors.Wrap(
@@ -687,6 +690,7 @@ func handleArchive(
 			"failed to walk the repository",
 		)
 	}
+	w.Header().Set("Omegaup-Uncompressed-Size", strconv.FormatInt(uncompressedSize, 10))
 	return nil
 }
 
diff --git a/browser_test.go b/browser_test.go
index c88b3b4..f5818e6 100644
--- a/browser_test.go
+++ b/browser_test.go
@@ -6,6 +6,7 @@ import (
 	"bytes"
 	"compress/gzip"
 	"context"
+	"io"
 	"net/http"
 	"net/http/httptest"
 	"reflect"
@@ -209,12 +210,15 @@ func TestHandleArchiveCommitTarball(t *testing.T) {
 	if "application/gzip" != response.Header().Get("Content-Type") {
 		t.Fatalf("Content-Type. Expected %s, got %s", "application/gzip", response.Header().Get("Content-Type"))
 	}
+	trailers := response.Result().Trailer
+	if _, ok := trailers["Omegaup-Uncompressed-Size"]; !ok {
+		t.Errorf("Omegaup-Uncompressed-Size was not present in the trailers: %v", trailers)
+	}
 
 	gz, err := gzip.NewReader(bytes.NewReader(response.Body.Bytes()))
 	if err != nil {
 		t.Fatalf("Error opening gzip from response: %v", err)
 	}
-	defer gz.Close()
 
 	a := tar.NewReader(gz)
 	hdr, err := a.Next()
@@ -225,6 +229,19 @@ func TestHandleArchiveCommitTarball(t *testing.T) {
 	if "empty" != hdr.Name {
 		t.Errorf("Expected %s, got %v", "empty", hdr.Name)
 	}
+	_, err = io.Copy(io.Discard, a)
+	if err != nil {
+		t.Fatalf("Error reading tar file: %v", err)
+	}
+	hdr, err = a.Next()
+	if err == nil {
+		t.Fatalf("Tarball has unexpected extra files: %v", hdr)
+	}
+	gz.Close()
+
+	if "0" != trailers.Get("Omegaup-Uncompressed-Size") {
+		t.Errorf("Omegaup-Uncompressed-Size trailer. Expected 0, got %v", trailers.Get("Omegaup-Uncompressed-Size"))
+	}
 }
 
 func TestHandleArchiveCommitTarballMismatchedTree(t *testing.T) {