25w20a

proxy/matchrepo.go

@@ -2,12 +2,18 @@ package proxy
 
 import (
 	"bufio"
+	"bytes"
 	"compress/gzip"
 	"fmt"
 	"ghproxy/config"
 	"io"
+	"net/http"
 	"regexp"
 	"strings"
+
+	"github.com/cloudwego/hertz/pkg/app"
+	hresp "github.com/cloudwego/hertz/pkg/protocol/http1/resp"
+	"github.com/valyala/bytebufferpool"
 )
 
 // 定义错误类型, error承载描述, 便于处理
@@ -205,15 +211,136 @@ func matchString(target string, stringsToMatch []string) bool {
 	return exists
 }
 
-// processLinks 处理链接并将结果写入输出流
-func processLinks(input io.Reader, output io.Writer, compress string, host string, cfg *config.Config) (written int64, err error) {
+// processLinks 处理链接并返回一个 io.ReadCloser
+func processLinks(input io.Reader, compress string, host string, cfg *config.Config) (io.ReadCloser, error) {
 	var reader *bufio.Reader
 
 	if compress == "gzip" {
-		// 解压gzip
+		// 解压 gzip
 		gzipReader, err := gzip.NewReader(input)
 		if err != nil {
-			return 0, fmt.Errorf("gzip解压错误: %v", err)
+			return nil, fmt.Errorf("gzip 解压错误: %w", err)
+		}
+		reader = bufio.NewReader(gzipReader)
+	} else {
+		reader = bufio.NewReader(input)
+	}
+
+	// 创建一个缓冲区用于存储输出
+	var outputBuffer io.Writer
+	var gzipWriter *gzip.Writer
+	var output io.ReadCloser
+	var buf bytes.Buffer
+
+	if compress == "gzip" {
+		// 创建一个管道来连接 gzipWriter 和 output
+		pipeReader, pipeWriter := io.Pipe() // 创建一个管道
+		output = pipeReader                 // 将管道的读取端作为输出
+		outputBuffer = pipeWriter           // 将管道的写入端作为 outputBuffer
+		gzipWriter = gzip.NewWriter(outputBuffer)
+		go func() {
+			defer pipeWriter.Close() // 确保在 goroutine 结束时关闭 pipeWriter
+			writer := bufio.NewWriter(gzipWriter)
+			defer func() {
+				if err := writer.Flush(); err != nil {
+					logError("gzip writer 刷新失败: %v", err)
+				}
+				if err := gzipWriter.Close(); err != nil {
+					logError("gzipWriter 关闭失败: %v", err)
+				}
+			}()
+
+			scanner := bufio.NewScanner(reader)
+			urlPattern := regexp.MustCompile(`https?://[^\s'"]+`)
+			for scanner.Scan() {
+				line := scanner.Text()
+				modifiedLine := urlPattern.ReplaceAllStringFunc(line, func(originalURL string) string {
+					return modifyURL(originalURL, host, cfg)
+				})
+				if _, err := writer.WriteString(modifiedLine + "\n"); err != nil {
+					logError("写入 gzipWriter 失败: %v", err)
+					return // 在发生错误时退出 goroutine
+				}
+			}
+			if err := scanner.Err(); err != nil {
+				logError("读取输入错误: %v", err)
+			}
+		}()
+	} else {
+		outputBuffer = &buf
+		writer := bufio.NewWriter(outputBuffer)
+		defer func() {
+			if err := writer.Flush(); err != nil {
+				logError("writer 刷新失败: %v", err)
+			}
+		}()
+
+		urlPattern := regexp.MustCompile(`https?://[^\s'"]+`)
+		scanner := bufio.NewScanner(reader)
+		for scanner.Scan() {
+			line := scanner.Text()
+			modifiedLine := urlPattern.ReplaceAllStringFunc(line, func(originalURL string) string {
+				return modifyURL(originalURL, host, cfg)
+			})
+			if _, err := writer.WriteString(modifiedLine + "\n"); err != nil {
+				return nil, fmt.Errorf("写入文件错误: %w", err) // 传递错误
+			}
+		}
+		if err := scanner.Err(); err != nil {
+			return nil, fmt.Errorf("读取行错误: %w", err) // 传递错误
+		}
+		output = io.NopCloser(&buf)
+	}
+
+	return output, nil
+}
+
+func WriteChunkedBody(resp io.ReadCloser, c *app.RequestContext) {
+	defer resp.Close()
+
+	c.Response.HijackWriter(hresp.NewChunkedBodyWriter(&c.Response, c.GetWriter()))
+
+	bufWrapper := bytebufferpool.Get()
+	buf := bufWrapper.B
+	size := 32768 // 32KB
+	buf = buf[:cap(buf)]
+	if len(buf) < size {
+		buf = append(buf, make([]byte, size-len(buf))...)
+	}
+	buf = buf[:size] // 将缓冲区限制为 'size'
+	defer bytebufferpool.Put(bufWrapper)
+
+	for {
+		n, err := resp.Read(buf)
+		if err != nil {
+			if err == io.EOF {
+				break // 读取到文件末尾
+			}
+			fmt.Println("读取错误:", err)
+			c.String(http.StatusInternalServerError, "读取错误")
+			return
+		}
+
+		_, err = c.Write(buf[:n]) // 写入 chunk
+		if err != nil {
+			fmt.Println("写入 chunk 错误:", err)
+			return
+		}
+
+		c.Flush() // 刷新 chunk 到客户端
+	}
+}
+
+// processLinksAndWriteChunked 处理链接并将结果以 chunked 方式写入响应
+func ProcessLinksAndWriteChunked(input io.Reader, compress string, host string, cfg *config.Config, c *app.RequestContext) {
+	var reader *bufio.Reader
+
+	if compress == "gzip" {
+		// 解压 gzip
+		gzipReader, err := gzip.NewReader(input)
+		if err != nil {
+			c.String(http.StatusInternalServerError, fmt.Sprintf("gzip 解压错误: %v", err))
+			return
 		}
 		defer gzipReader.Close()
 		reader = bufio.NewReader(gzipReader)
@@ -221,36 +348,108 @@ func processLinks(input io.Reader, output io.Writer, compress string, host strin
 		reader = bufio.NewReader(input)
 	}
 
-	var writer *bufio.Writer
+	// 获取 chunked body writer
+	chunkedWriter := hresp.NewChunkedBodyWriter(&c.Response, c.GetWriter())
+
+	var writer io.Writer = chunkedWriter
 	var gzipWriter *gzip.Writer
 
-	// 根据是否gzip确定 writer 的创建
 	if compress == "gzip" {
-		gzipWriter = gzip.NewWriter(output)
-		writer = bufio.NewWriterSize(gzipWriter, 4096) //设置缓冲区大小
-	} else {
-		writer = bufio.NewWriterSize(output, 4096)
+		gzipWriter = gzip.NewWriter(writer)
+		writer = gzipWriter
+		defer func() {
+			if err := gzipWriter.Close(); err != nil {
+				logError("gzipWriter close failed: %v", err)
+			}
+		}()
 	}
 
-	//确保writer关闭
-	defer func() {
-		var closeErr error // 局部变量，用于保存defer中可能发生的错误
+	bufWrapper := bytebufferpool.Get()
+	buf := bufWrapper.B
+	size := 32768 // 32KB
+	buf = buf[:cap(buf)]
+	if len(buf) < size {
+		buf = append(buf, make([]byte, size-len(buf))...)
+	}
+	buf = buf[:size] // 将缓冲区限制为 'size'
+	defer bytebufferpool.Put(bufWrapper)
 
-		if gzipWriter != nil {
-			if closeErr = gzipWriter.Close(); closeErr != nil {
-				logError("gzipWriter close failed %v", closeErr)
-				// 如果已经存在错误，则保留。否则，记录此错误。
-				if err == nil {
-					err = closeErr
-				}
+	urlPattern := regexp.MustCompile(`https?://[^\s'"]+`)
+	scanner := bufio.NewScanner(reader)
+	for scanner.Scan() {
+		line := scanner.Text()
+		modifiedLine := urlPattern.ReplaceAllStringFunc(line, func(originalURL string) string {
+			return modifyURL(originalURL, host, cfg)
+		})
+		modifiedLineWithNewline := modifiedLine + "\n"
+
+		_, err := writer.Write([]byte(modifiedLineWithNewline))
+		if err != nil {
+			logError("写入 chunk 错误: %v", err)
+			return // 发生错误时退出
+		}
+
+		if compress != "gzip" {
+			if fErr := chunkedWriter.Flush(); fErr != nil {
+				logError("chunkedWriter flush failed: %v", fErr)
+				return
 			}
 		}
+	}
+
+	if err := scanner.Err(); err != nil {
+		logError("读取输入错误: %v", err)
+		c.String(http.StatusInternalServerError, fmt.Sprintf("读取输入错误: %v", err))
+		return
+	}
+
+	// 对于 gzip，chunkedWriter 的关闭会触发最后的 chunk
+	if compress != "gzip" {
+		if fErr := chunkedWriter.Flush(); fErr != nil {
+			logError("final chunkedWriter flush failed: %v", fErr)
+		}
+	}
+}
+
+func ProcessAndWriteChunkedBody(input io.Reader, compress string, host string, cfg *config.Config, c *app.RequestContext) error {
+	var reader *bufio.Reader
+
+	if compress == "gzip" {
+		// 解压gzip
+		gzipReader, err := gzip.NewReader(input)
+		if err != nil {
+			return fmt.Errorf("gzip解压错误: %v", err)
+		}
+		defer gzipReader.Close()
+		reader = bufio.NewReader(gzipReader)
+	} else {
+		reader = bufio.NewReader(input)
+	}
+
+	// 创建一个缓冲区用于存储输出
+	var outputBuffer io.Writer
+	var gzipWriter *gzip.Writer
+	var buf bytes.Buffer
+
+	if compress == "gzip" {
+		// 创建一个缓冲区
+		outputBuffer = &buf
+		gzipWriter = gzip.NewWriter(outputBuffer)
+		defer func() {
+			if gzipWriter != nil {
+				if closeErr := gzipWriter.Close(); closeErr != nil {
+					logError("gzipWriter close failed %v", closeErr)
+				}
+			}
+		}()
+	} else {
+		outputBuffer = &buf
+	}
+
+	writer := bufio.NewWriter(outputBuffer)
+	defer func() {
 		if flushErr := writer.Flush(); flushErr != nil {
 			logError("writer flush failed %v", flushErr)
-			// 如果已经存在错误，则保留。否则，记录此错误。
-			if err == nil {
-				err = flushErr
-			}
 		}
 	}()
 
@@ -262,7 +461,7 @@ func processLinks(input io.Reader, output io.Writer, compress string, host strin
 			if err == io.EOF {
 				break // 文件结束
 			}
-			return written, fmt.Errorf("读取行错误: %v", err) // 传递错误
+			return fmt.Errorf("读取行错误: %v", err) // 传递错误
 		}
 
 		// 替换所有匹配的 URL
@@ -270,17 +469,56 @@ func processLinks(input io.Reader, output io.Writer, compress string, host strin
 			return modifyURL(originalURL, host, cfg)
 		})
 
-		n, werr := writer.WriteString(modifiedLine)
-		written += int64(n) // 更新写入的字节数
+		_, werr := writer.WriteString(modifiedLine)
 		if werr != nil {
-			return written, fmt.Errorf("写入文件错误: %v", werr) // 传递错误
+			return fmt.Errorf("写入文件错误: %v", werr) // 传递错误
 		}
 	}
 
 	// 在返回之前，再刷新一次
 	if fErr := writer.Flush(); fErr != nil {
-		return written, fErr
+		return fErr
 	}
 
-	return written, nil
+	if compress == "gzip" {
+		if err := gzipWriter.Close(); err != nil {
+			return fmt.Errorf("gzipWriter close failed: %v", err)
+		}
+	}
+
+	// 将处理后的内容以分块的方式写入响应
+	c.Response.HijackWriter(hresp.NewChunkedBodyWriter(&c.Response, c.GetWriter()))
+
+	bufWrapper := bytebufferpool.Get()
+	bbuf := bufWrapper.B
+	size := 32768 // 32KB
+	if cap(bbuf) < size {
+		bbuf = make([]byte, size)
+	} else {
+		bbuf = bbuf[:size]
+	}
+	defer bytebufferpool.Put(bufWrapper)
+
+	// 将缓冲区内容写入响应
+	for {
+		n, err := buf.Read(bbuf)
+		if err != nil {
+			if err != io.EOF {
+				fmt.Println("读取错误:", err)
+				c.String(http.StatusInternalServerError, "读取错误")
+				return err
+			}
+			break // 读取到文件末尾
+		}
+
+		_, err = c.Write(bbuf[:n]) // 写入 chunk
+		if err != nil {
+			fmt.Println("写入 chunk 错误:", err)
+			return err
+		}
+
+		c.Flush() // 刷新 chunk 到客户端
+	}
+
+	return nil
 }