Add Total upload and download stats

go.mod

@@ -4,10 +4,12 @@ go 1.22.0
 
 toolchain go1.23.0
 
-require github.com/quic-go/quic-go v0.48.1
+require (
+	github.com/conduitio/bwlimit v0.1.0
+	github.com/quic-go/quic-go v0.48.1
+)
 
 require (
-	github.com/conduitio/bwlimit v0.1.0 // indirect
 	github.com/go-task/slim-sprig/v3 v3.0.0 // indirect
 	github.com/google/pprof v0.0.0-20241029010322-833c56d90c8e // indirect
 	github.com/onsi/ginkgo/v2 v2.20.2 // indirect

httppaths.go

@@ -10,7 +10,38 @@ import (
 	"sync/atomic"
 )
 
+type ResponseWriterWrapper struct {
+	http.ResponseWriter
+	size int
+}
+
+func (w *ResponseWriterWrapper) Size() int {
+	return w.size
+}
+
+// This overrides the http.ResponseWriter Write() function. Do not remove!
+func (w *ResponseWriterWrapper) Write(b []byte) (int, error) {
+	n, err := w.ResponseWriter.Write(b)
+	w.size += n
+	return n, err
+}
+
+func doRequest(request *http.Request) *http.Response {
+	resp, err := client.Do(request)
+	if err != nil {
+		log.Panic(err)
+	}
+
+	// Use `resp.ContentLength` for now, copying the whole body into memory uses a lot of memory lol!
+	// size, _ := httputil.DumpResponse(resp, true)
+	// atomic.AddInt64(&totalBandwidth.TotalDownload, int64(len(size)))
+	atomic.AddInt64(&totalBandwidth.TotalDownload, resp.ContentLength)
+
+	return resp
+}
+
 func videoplayback(w http.ResponseWriter, req *http.Request) {
+	w_ := &ResponseWriterWrapper{ResponseWriter: w}
 	q := req.URL.Query()
 	host := q.Get("host")
 	q.Del("host")
@@ -66,16 +97,14 @@ func videoplayback(w http.ResponseWriter, req *http.Request) {
 	body := []byte{0x78, 0} // protobuf body
 
 	request, err := http.NewRequest("POST", proxyURL.String(), bytes.NewReader(body))
-	copyHeaders(req.Header, request.Header, false)
-	request.Header.Set("User-Agent", ua)
 	if err != nil {
 		log.Panic(err)
 	}
 
-	resp, err := client.Do(request)
-	if err != nil {
-		log.Panic(err)
-	}
+	copyHeaders(req.Header, request.Header, false)
+	request.Header.Set("User-Agent", ua)
+
+	resp := doRequest(request)
 
 	if resp.StatusCode == 403 {
 		atomic.AddInt64(&stats_.RequestsForbidden.Videoplayback, 1)
@@ -84,8 +113,6 @@ func videoplayback(w http.ResponseWriter, req *http.Request) {
 		return
 	}
 
-	defer resp.Body.Close()
-
 	NoRewrite := strings.HasPrefix(resp.Header.Get("Content-Type"), "audio") || strings.HasPrefix(resp.Header.Get("Content-Type"), "video")
 	copyHeaders(resp.Header, w.Header(), NoRewrite)
 
@@ -118,7 +145,8 @@ func videoplayback(w http.ResponseWriter, req *http.Request) {
 
 		io.WriteString(w, strings.Join(lines, "\n"))
 	} else {
-		io.Copy(w, resp.Body)
+		io.Copy(w_, resp.Body)
+		atomic.AddInt64(&totalBandwidth.TotalUpload, int64(w_.Size()))
 	}
 }
 
@@ -145,17 +173,14 @@ func vi(w http.ResponseWriter, req *http.Request) {
 	proxyURL.RawQuery = q.Encode()
 
 	request, err := http.NewRequest(req.Method, proxyURL.String(), nil)
+	if err != nil {
+		log.Panic(err)
+	}
 
 	copyHeaders(req.Header, request.Header, false)
 	request.Header.Set("User-Agent", ua)
-	if err != nil {
-		log.Panic(err)
-	}
 
-	resp, err := client.Do(request)
-	if err != nil {
-		log.Panic(err)
-	}
+	resp := doRequest(request)
 
 	w.WriteHeader(resp.StatusCode)
 	if resp.StatusCode == 403 {
@@ -185,16 +210,14 @@ func ggpht(w http.ResponseWriter, req *http.Request) {
 	}
 
 	request, err := http.NewRequest(req.Method, proxyURL.String(), nil)
-	copyHeaders(req.Header, request.Header, false)
-	request.Header.Set("User-Agent", ua)
 	if err != nil {
 		log.Panic(err)
 	}
 
-	resp, err := client.Do(request)
-	if err != nil {
-		log.Panic(err)
-	}
+	copyHeaders(req.Header, request.Header, false)
+	request.Header.Set("User-Agent", ua)
+
+	resp := doRequest(request)
 
 	w.WriteHeader(resp.StatusCode)
 	if resp.StatusCode == 403 {

main.go

@@ -93,11 +93,25 @@ var ipv6_only = false
 
 var version string
 
+var programInit = time.Now()
+
+type TotalBandwidth struct {
+	TotalUpload   int64
+	TotalDownload int64
+}
+
+var totalBandwidth = TotalBandwidth{TotalUpload: 0, TotalDownload: 0}
+
 type statusJson struct {
 	Version          string        `json:"version"`
+	Uptime           time.Duration `json:"uptime"`
 	RequestCount     int64         `json:"requestCount"`
 	RequestPerSecond int64         `json:"requestPerSecond"`
 	RequestPerMinute int64         `json:"requestPerMinute"`
+	// CurrentUpload     any           `json:"currentUpload"`
+	// CurrentDownload   any           `json:"currentDownload"`
+	TotalUpload       string `json:"totalUpload"`
+	TotalDownload     string `json:"totalDownload"`
 	RequestsForbidden struct {
 		Videoplayback int64 `json:"videoplayback"`
 		Vi            int64 `json:"vi"`
@@ -107,9 +121,14 @@ type statusJson struct {
 
 var stats_ = statusJson{
 	Version:          version + "-" + runtime.GOARCH,
+	Uptime:           0,
 	RequestCount:     0,
 	RequestPerSecond: 0,
 	RequestPerMinute: 0,
+	// CurrentUpload:    nil,
+	// CurrentDownload:  nil,
+	TotalUpload:   "",
+	TotalDownload: "",
 	RequestsForbidden: struct {
 		Videoplayback int64 `json:"videoplayback"`
 		Vi            int64 `json:"vi"`
@@ -135,6 +154,12 @@ func root(w http.ResponseWriter, req *http.Request) {
 func stats(w http.ResponseWriter, req *http.Request) {
 	w.Header().Set("Content-Type", "application/json")
 
+	stats_.Uptime = time.Duration(time.Since(programInit).Seconds())
+	tu := float64(totalBandwidth.TotalUpload) / float64(1_048_576) // 1MiB = 1_048_576 (DON'T CONFUSE THEM WITH 1MB WHICH IS 1_000_000 BYTES)
+	td := float64(totalBandwidth.TotalDownload) / float64(1_048_576)
+	stats_.TotalUpload = fmt.Sprintf("%.2fMiB", tu)
+	stats_.TotalDownload = fmt.Sprintf("%.2fMiB", td)
+
 	if err := json.NewEncoder(w).Encode(stats_); err != nil {
 		http.Error(w, err.Error(), http.StatusInternalServerError)
 	}