diff --git a/go.mod b/go.mod index d3a170a..d0fc567 100644 --- a/go.mod +++ b/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 diff --git a/httppaths.go b/httppaths.go index 003ab34..84c74ab 100644 --- a/httppaths.go +++ b/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 { diff --git a/main.go b/main.go index 585c0b2..82c8775 100644 --- a/main.go +++ b/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"` - RequestCount int64 `json:"requestCount"` - RequestPerSecond int64 `json:"requestPerSecond"` - RequestPerMinute int64 `json:"requestPerMinute"` + 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) }