From e9953946a4d2b86ca152486f3b0fbb1d835c224a Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?=E9=87=8C=E8=A6=8B=20=E7=81=AF=E8=8A=B1?=
 <172008506+satomitouka@users.noreply.github.com>
Date: Tue, 28 Jan 2025 11:27:29 +0800
Subject: [PATCH] Update gitreq.go

---
 proxy/gitreq.go | 70 ++++++++++++++++++++++++++++++++++++++++++-------
 1 file changed, 61 insertions(+), 9 deletions(-)

diff --git a/proxy/gitreq.go b/proxy/gitreq.go
index 173fa6c..6dbd20d 100644
--- a/proxy/gitreq.go
+++ b/proxy/gitreq.go
@@ -6,6 +6,7 @@ import (
 	"ghproxy/config"
 	"io"
 	"net/http"
+	"strconv"
 	"time"
 
 	"github.com/gin-gonic/gin"
@@ -48,11 +49,24 @@ func GitReq(c *gin.Context, u string, cfg *config.Config, mode string, runMode s
 		HandleError(c, fmt.Sprintf("Failed to send request: %v", err))
 		return
 	}
-	defer headResp.Body.Close()
 
-	if err := HandleResponseSize(headResp, cfg, c); err != nil {
-		logWarning("%s %s %s %s %s Response-Size-Error: %v", c.ClientIP(), method, u, c.Request.Header.Get("User-Agent"), c.Request.Proto, err)
-		return
+	// defer headResp.Body.Close()
+	defer func(Body io.ReadCloser) {
+		if err := Body.Close(); err != nil {
+			logError("Failed to close response body: %v", err)
+		}
+	}(headResp.Body)
+
+	contentLength := headResp.Header.Get("Content-Length")
+	sizelimit := cfg.Server.SizeLimit * 1024 * 1024
+	if contentLength != "" {
+		size, err := strconv.Atoi(contentLength)
+		if err == nil && size > sizelimit {
+			finalURL := headResp.Request.URL.String()
+			c.Redirect(http.StatusMovedPermanently, finalURL)
+			logWarning("%s %s %s %s %s Final-URL: %s Size-Limit-Exceeded: %d", c.ClientIP(), c.Request.Method, c.Request.URL.String(), c.Request.Header.Get("User-Agent"), c.Request.Proto, finalURL, size)
+			return
+		}
 	}
 
 	body, err := readRequestBody(c)
@@ -77,14 +91,52 @@ func GitReq(c *gin.Context, u string, cfg *config.Config, mode string, runMode s
 		HandleError(c, fmt.Sprintf("Failed to send request: %v", err))
 		return
 	}
-	defer resp.Body.Close()
+	//defer resp.Body.Close()
+	defer func(Body io.ReadCloser) {
+		if err := Body.Close(); err != nil {
+			logError("Failed to close response body: %v", err)
+		}
+	}(resp.Body)
 
-	if err := HandleResponseSize(resp, cfg, c); err != nil {
-		logWarning("%s %s %s %s %s Response-Size-Error: %v", c.ClientIP(), method, u, c.Request.Header.Get("User-Agent"), c.Request.Proto, err)
-		return
+	/*
+		if err := HandleResponseSize(resp, cfg, c); err != nil {
+			logWarning("%s %s %s %s %s Response-Size-Error: %v", c.ClientIP(), method, u, c.Request.Header.Get("User-Agent"), c.Request.Proto, err)
+			return
+		}
+	*/
+	contentLength = resp.Header.Get("Content-Length")
+	if contentLength != "" {
+		size, err := strconv.Atoi(contentLength)
+		if err == nil && size > sizelimit {
+			finalURL := resp.Request.URL.String()
+			c.Redirect(http.StatusMovedPermanently, finalURL)
+			logWarning("%s %s %s %s %s Final-URL: %s Size-Limit-Exceeded: %d", c.ClientIP(), c.Request.Method, c.Request.URL.String(), c.Request.Header.Get("User-Agent"), c.Request.Proto, finalURL, size)
+			return
+		}
+	}
+
+	for key, values := range resp.Header {
+		for _, value := range values {
+			c.Header(key, value)
+		}
+	}
+
+	headersToRemove := map[string]struct{}{
+		"Content-Security-Policy":   {},
+		"Referrer-Policy":           {},
+		"Strict-Transport-Security": {},
+	}
+
+	for header := range headersToRemove {
+		resp.Header.Del(header)
+	}
+
+	if cfg.CORS.Enabled {
+		c.Header("Access-Control-Allow-Origin", "*")
+	} else {
+		c.Header("Access-Control-Allow-Origin", "")
 	}
 
-	CopyResponseHeaders(resp, c, cfg)
 	c.Status(resp.StatusCode)
 
 	if _, err := io.Copy(c.Writer, resp.Body); err != nil {