From e31c9168627c040317e5cc8566724f88910439ae Mon Sep 17 00:00:00 2001
From: Mitch Riedstra <mitch@riedstra.us>
Date: Sat, 9 Jan 2021 00:49:16 -0500
Subject: Add some download stats, and a status endpoint.

---
 cmd/web/download.go | 41 +++++++++++++++++++++++++++++++++++------
 cmd/web/install.go  | 31 +++++++++++++++++++++++++------
 cmd/web/main.go     |  1 +
 steam/package.go    | 11 +++++++++--
 4 files changed, 70 insertions(+), 14 deletions(-)

diff --git a/cmd/web/download.go b/cmd/web/download.go
index 93c4ff5..1c70717 100644
--- a/cmd/web/download.go
+++ b/cmd/web/download.go
@@ -1,7 +1,10 @@
 package main
 
 import (
+	"io"
 	"net/http"
+	"time"
+	"fmt"
 
 	"github.com/gorilla/mux"
 )
@@ -20,14 +23,40 @@ func gameDownloader(w http.ResponseWriter, r *http.Request) {
 	}
 
 	w.Header().Add("Content-type", "application/tar")
+	w.Header().Add("Estimated-size", fmt.Sprintf("%d", g.Size)) 
 
 	Logger.Printf("Client %s is downloading: %s", r.RemoteAddr, game)
-	err := g.Package(w)
-	if err != nil {
-		Logger.Printf("Client %s Error Sending game: %s", r.RemoteAddr, err)
-		// Headers already sent, don't bother sending an error
-		return
+
+	// Invert the writer so we can break up the copy and get progress
+	// information in here
+	rdr, pwrtr := io.Pipe()
+	go func() {
+		err := g.Package(pwrtr)
+		if err != nil {
+			Logger.Println("Error in package writing: ", err)
+		}
+	}()
+
+	var total int64
+
+	start := time.Now()
+	for {
+		n, err := io.CopyN(w, rdr, 256*1024*1024)
+		if err == io.EOF {
+			break
+		}
+		if err != nil {
+			Logger.Printf("Client %s Error Sending game: %s", r.RemoteAddr, err)
+			// Headers already sent, don't bother sending an error
+			return
+		}
+		total += n
+		mb := float64(total / 1024 / 1024)
+		rate := mb / time.Since(start).Seconds()
+
+		Logger.Printf("Client %s is downloading %s: %0.1f%% done %.2f mb/s",
+			r.RemoteAddr, game, float64(total)/float64(g.Size)*100, rate)
 	}
+
 	Logger.Printf("Client %s finished downloading: %s", r.RemoteAddr, game)
 }
-
diff --git a/cmd/web/install.go b/cmd/web/install.go
index 63faaf6..678037a 100644
--- a/cmd/web/install.go
+++ b/cmd/web/install.go
@@ -1,18 +1,23 @@
 package main
 
 import (
+	"encoding/json"
+	"fmt"
+	"net/http"
 	"net/url"
 	"os"
 	"strings"
-	"fmt"
-	"net/http"
 	"sync"
+	"time"
 )
 
 type statusInfo struct {
-	Running bool
-	Error   error
-	Url     string
+	Running     bool
+	Error       error
+	Url         string
+	Transferred int64
+	Size        int64
+	Start       *time.Time
 }
 
 var (
@@ -27,6 +32,19 @@ var (
 	getPath = make(chan string)
 )
 
+func statsHandler(w http.ResponseWriter, r *http.Request) {
+	status.m.RLock()
+	defer status.m.RUnlock()
+
+	enc := json.NewEncoder(w)
+
+	err := enc.Encode(status.s)
+	if err != nil {
+		Logger.Println("While encoding status: ", err)
+	}
+	return
+}
+
 func installHttp(u string) error {
 	Logger.Println("Installer: loading from url")
 	resp, err := http.Get(u)
@@ -57,6 +75,8 @@ func installPath(p string) error {
 	return nil
 }
 
+// installer handles installing games either from a local path or a
+// remote URL
 func installer(urls <-chan string) {
 	var err error
 	for u := range urls {
@@ -87,7 +107,6 @@ func gameInstaller(w http.ResponseWriter, r *http.Request) {
 		return
 	}
 
-
 	err := r.ParseForm()
 	if err != nil {
 		Logger.Printf("Installer: While parsing form: %s", err)
diff --git a/cmd/web/main.go b/cmd/web/main.go
index 0f76670..a16f15f 100644
--- a/cmd/web/main.go
+++ b/cmd/web/main.go
@@ -153,6 +153,7 @@ func main() {
 	r.HandleFunc("/install", gameInstaller)
 	r.HandleFunc("/steam-export-web.exe", serveSelf)
 	r.HandleFunc("/download/{game}", gameDownloader)
+	r.HandleFunc("/status", statsHandler)
 	r.HandleFunc("/style.css", cssHandler)
 	r.HandleFunc("/", index)
 
diff --git a/steam/package.go b/steam/package.go
index 76914dc..7287e10 100644
--- a/steam/package.go
+++ b/steam/package.go
@@ -9,7 +9,7 @@ import (
 )
 
 // Package writes the package, returning bytes written and an error if any
-func (g *Game) Package(wr io.Writer) error {
+func (g *Game) Package(wr io.WriteCloser) error {
 	if err := os.Chdir(g.LibraryPath); err != nil {
 		return err
 	}
@@ -26,7 +26,14 @@ func (g *Game) Package(wr io.Writer) error {
 		}
 	}
 
-	return twriter.Flush()
+	err = twriter.Flush()
+	if err != nil {
+		return err
+	}
+
+	err = twriter.Close()
+
+	return wr.Close()
 }
 
 func (l *Library) Extract(r io.Reader) error {
-- 
cgit v1.2.3