package proxy import ( "bufio" "bytes" "fmt" "ghproxy/config" "io" "strings" "sync" "github.com/infinite-iroha/touka" ) var ( prefixGithub = []byte("https://github.com") prefixRawUser = []byte("https://raw.githubusercontent.com") prefixRaw = []byte("https://raw.github.com") prefixGistUser = []byte("https://gist.githubusercontent.com") prefixGist = []byte("https://gist.github.com") prefixAPI = []byte("https://api.github.com") prefixHTTP = []byte("http://") prefixHTTPS = []byte("https://") ) func EditorMatcherBytes(rawPath []byte, cfg *config.Config) (bool, error) { if bytes.HasPrefix(rawPath, prefixGithub) { return true, nil } if bytes.HasPrefix(rawPath, prefixRawUser) { return true, nil } if bytes.HasPrefix(rawPath, prefixRaw) { return true, nil } if bytes.HasPrefix(rawPath, prefixGistUser) { return true, nil } if bytes.HasPrefix(rawPath, prefixGist) { return true, nil } if cfg.Shell.RewriteAPI { if bytes.HasPrefix(rawPath, prefixAPI) { return true, nil } } return false, nil } func modifyURLBytes(url []byte, host []byte, cfg *config.Config) []byte { matched, err := EditorMatcherBytes(url, cfg) if err != nil || !matched { return url } var u []byte if bytes.HasPrefix(url, prefixHTTPS) { u = url[len(prefixHTTPS):] } else if bytes.HasPrefix(url, prefixHTTP) { u = url[len(prefixHTTP):] } else { u = url } newLen := len(prefixHTTPS) + len(host) + 1 + len(u) newURL := make([]byte, newLen) written := 0 written += copy(newURL[written:], prefixHTTPS) written += copy(newURL[written:], host) written += copy(newURL[written:], []byte("/")) copy(newURL[written:], u) return newURL } func EditorMatcher(rawPath string, cfg *config.Config) (bool, error) { // 匹配 "https://github.com"开头的链接 if strings.HasPrefix(rawPath, "https://github.com") { return true, nil } // 匹配 "https://raw.githubusercontent.com"开头的链接 if strings.HasPrefix(rawPath, "https://raw.githubusercontent.com") { return true, nil } // 匹配 "https://raw.github.com"开头的链接 if strings.HasPrefix(rawPath, "https://raw.github.com") { return true, nil } // 匹配 "https://gist.githubusercontent.com"开头的链接 if strings.HasPrefix(rawPath, "https://gist.githubusercontent.com") { return true, nil } // 匹配 "https://gist.github.com"开头的链接 if strings.HasPrefix(rawPath, "https://gist.github.com") { return true, nil } if cfg.Shell.RewriteAPI { // 匹配 "https://api.github.com/"开头的链接 if strings.HasPrefix(rawPath, "https://api.github.com") { return true, nil } } return false, nil } // 匹配文件扩展名是sh的rawPath func MatcherShell(rawPath string) bool { return strings.HasSuffix(rawPath, ".sh") } // LinkProcessor 是一个函数类型,用于处理提取到的链接。 type LinkProcessor func(string) string // 自定义 URL 修改函数 func modifyURL(url string, host string, cfg *config.Config) string { // 去除url内的https://或http:// matched, err := EditorMatcher(url, cfg) if err != nil { return url } if matched { var u = url u = strings.TrimPrefix(u, "https://") u = strings.TrimPrefix(u, "http://") return "https://" + host + "/" + u } return url } var bufferPool = sync.Pool{ New: func() interface{} { return new(bytes.Buffer) }, } // processLinksStreamingInternal is a link processing function that reads the input line by line. // It is memory-safe for large files but less performant due to numerous small allocations. func processLinksStreamingInternal(input io.ReadCloser, host string, cfg *config.Config, c *touka.Context) (readerOut io.Reader, written int64, err error) { pipeReader, pipeWriter := io.Pipe() readerOut = pipeReader go func() { defer func() { if err != nil { pipeWriter.CloseWithError(err) } else { pipeWriter.Close() } }() defer input.Close() bufReader := bufio.NewReader(input) bufWriter := bufio.NewWriterSize(pipeWriter, 4096) defer bufWriter.Flush() for { line, readErr := bufReader.ReadString('\n') if readErr != nil && readErr != io.EOF { err = fmt.Errorf("read error: %w", readErr) return } modifiedLine := urlPattern.ReplaceAllStringFunc(line, func(originalURL string) string { return modifyURL(originalURL, host, cfg) }) var n int n, err = bufWriter.WriteString(modifiedLine) written += int64(n) if err != nil { err = fmt.Errorf("write error: %w", err) return } if readErr == io.EOF { break } } }() return readerOut, written, nil } // processLinks acts as a dispatcher, choosing the best processing strategy based on file size. // It uses a memory-safe streaming approach for large or unknown-size files, // and a high-performance buffered approach for smaller files. func processLinks(input io.ReadCloser, host string, cfg *config.Config, c *touka.Context, bodySize int) (readerOut io.Reader, written int64, err error) { const sizeThreshold = 256 * 1024 // 256KB // Use streaming for large or unknown size files to prevent OOM if bodySize == -1 || bodySize > sizeThreshold { c.Debugf("Using streaming processor for large/unknown size file (%d bytes)", bodySize) return processLinksStreamingInternal(input, host, cfg, c) } else { c.Debugf("Using buffered processor for small file (%d bytes)", bodySize) return processLinksBufferedInternal(input, host, cfg, c) } } // processLinksBufferedInternal a link processing function that reads the entire content into a buffer. // It is optimized for performance on smaller files but carries an OOM risk for large files. func processLinksBufferedInternal(input io.ReadCloser, host string, cfg *config.Config, c *touka.Context) (readerOut io.Reader, written int64, err error) { pipeReader, pipeWriter := io.Pipe() readerOut = pipeReader hostBytes := []byte(host) go func() { // 在 goroutine 退出时, 根据 err 是否为 nil, 带错误或正常关闭 pipeWriter defer func() { if closeErr := input.Close(); closeErr != nil { c.Errorf("input close failed: %v", closeErr) } }() defer func() { if err != nil { if closeErr := pipeWriter.CloseWithError(err); closeErr != nil { c.Errorf("pipeWriter close with error failed: %v", closeErr) } } else { if closeErr := pipeWriter.Close(); closeErr != nil { c.Errorf("pipeWriter close failed: %v", closeErr) } } }() buf := bufferPool.Get().(*bytes.Buffer) buf.Reset() defer bufferPool.Put(buf) // 将全部输入读入复用的缓冲区 if _, err = buf.ReadFrom(input); err != nil { err = fmt.Errorf("reading input failed: %w", err) return } // 使用 ReplaceAllFunc 和字节版本辅助函数, 实现准零分配 modifiedBytes := urlPattern.ReplaceAllFunc(buf.Bytes(), func(originalURL []byte) []byte { return modifyURLBytes(originalURL, hostBytes, cfg) }) // 将处理后的字节写回管道 var n int n, err = pipeWriter.Write(modifiedBytes) if err != nil { err = fmt.Errorf("writing to pipe failed: %w", err) return } written = int64(n) }() return readerOut, written, nil }