@@ -2,14 +2,78 @@ package proxy
22
33import (
44 "bufio"
5+ "bytes"
56 "fmt"
67 "ghproxy/config"
78 "io"
89 "strings"
10+ "sync"
911
1012 "github.com/infinite-iroha/touka"
1113)
1214
15+ var (
16+ prefixGithub = []byte ("https://github.com" )
17+ prefixRawUser = []byte ("https://raw.githubusercontent.com" )
18+ prefixRaw = []byte ("https://raw.github.com" )
19+ prefixGistUser = []byte ("https://gist.githubusercontent.com" )
20+ prefixGist = []byte ("https://gist.github.com" )
21+ prefixAPI = []byte ("https://api.github.com" )
22+ prefixHTTP = []byte ("http://" )
23+ prefixHTTPS = []byte ("https://" )
24+ )
25+
26+ func EditorMatcherBytes (rawPath []byte , cfg * config.Config ) (bool , error ) {
27+ if bytes .HasPrefix (rawPath , prefixGithub ) {
28+ return true , nil
29+ }
30+ if bytes .HasPrefix (rawPath , prefixRawUser ) {
31+ return true , nil
32+ }
33+ if bytes .HasPrefix (rawPath , prefixRaw ) {
34+ return true , nil
35+ }
36+ if bytes .HasPrefix (rawPath , prefixGistUser ) {
37+ return true , nil
38+ }
39+ if bytes .HasPrefix (rawPath , prefixGist ) {
40+ return true , nil
41+ }
42+ if cfg .Shell .RewriteAPI {
43+ if bytes .HasPrefix (rawPath , prefixAPI ) {
44+ return true , nil
45+ }
46+ }
47+ return false , nil
48+ }
49+
50+ func modifyURLBytes (url []byte , host []byte , cfg * config.Config ) []byte {
51+ matched , err := EditorMatcherBytes (url , cfg )
52+ if err != nil || ! matched {
53+ return url
54+ }
55+
56+ var u []byte
57+ if bytes .HasPrefix (url , prefixHTTPS ) {
58+ u = url [len (prefixHTTPS ):]
59+ } else if bytes .HasPrefix (url , prefixHTTP ) {
60+ u = url [len (prefixHTTP ):]
61+ } else {
62+ u = url
63+ }
64+
65+ newLen := len (prefixHTTPS ) + len (host ) + 1 + len (u )
66+ newURL := make ([]byte , newLen )
67+
68+ written := 0
69+ written += copy (newURL [written :], prefixHTTPS )
70+ written += copy (newURL [written :], host )
71+ written += copy (newURL [written :], []byte ("/" ))
72+ copy (newURL [written :], u )
73+
74+ return newURL
75+ }
76+
1377func EditorMatcher (rawPath string , cfg * config.Config ) (bool , error ) {
1478 // 匹配 "https://github.com"开头的链接
1579 if strings .HasPrefix (rawPath , "https://github.com" ) {
@@ -64,87 +128,126 @@ func modifyURL(url string, host string, cfg *config.Config) string {
64128 return url
65129}
66130
67- // processLinks 处理链接，返回包含处理后数据的 io.Reader
68- func processLinks (input io.ReadCloser , host string , cfg * config.Config , c * touka.Context ) (readerOut io.Reader , written int64 , err error ) {
69- pipeReader , pipeWriter := io .Pipe () // 创建 io.Pipe
131+ var bufferPool = sync.Pool {
132+ New : func () interface {} {
133+ return new (bytes.Buffer )
134+ },
135+ }
136+
137+ // processLinksStreamingInternal is a link processing function that reads the input line by line.
138+ // It is memory-safe for large files but less performant due to numerous small allocations.
139+ func processLinksStreamingInternal (input io.ReadCloser , host string , cfg * config.Config , c * touka.Context ) (readerOut io.Reader , written int64 , err error ) {
140+ pipeReader , pipeWriter := io .Pipe ()
70141 readerOut = pipeReader
71142
72- go func () { // 在 Goroutine 中执行写入操作
143+ go func () {
73144 defer func () {
74- if pipeWriter != nil { // 确保 pipeWriter 关闭，即使发生错误
75- if err != nil {
76- if closeErr := pipeWriter .CloseWithError (err ); closeErr != nil { // 如果有错误，传递错误给 reader
77- c .Errorf ("pipeWriter close with error failed: %v, original error: %v" , closeErr , err )
78- }
79- } else {
80- if closeErr := pipeWriter .Close (); closeErr != nil { // 没有错误，正常关闭
81- c .Errorf ("pipeWriter close failed: %v" , closeErr )
82- if err == nil { // 如果之前没有错误，记录关闭错误
83- err = closeErr
84- }
85- }
86- }
145+ if err != nil {
146+ pipeWriter .CloseWithError (err )
147+ } else {
148+ pipeWriter .Close ()
87149 }
88150 }()
151+ defer input .Close ()
89152
90- defer func () {
91- if err := input .Close (); err != nil {
92- c .Errorf ("input close failed: %v" , err )
153+ bufReader := bufio .NewReader (input )
154+ bufWriter := bufio .NewWriterSize (pipeWriter , 4096 )
155+ defer bufWriter .Flush ()
156+
157+ for {
158+ line , readErr := bufReader .ReadString ('\n' )
159+ if readErr != nil && readErr != io .EOF {
160+ err = fmt .Errorf ("read error: %w" , readErr )
161+ return
93162 }
94163
95- }()
164+ modifiedLine := urlPattern .ReplaceAllStringFunc (line , func (originalURL string ) string {
165+ return modifyURL (originalURL , host , cfg )
166+ })
96167
97- var bufReader * bufio.Reader
168+ var n int
169+ n , err = bufWriter .WriteString (modifiedLine )
170+ written += int64 (n )
171+ if err != nil {
172+ err = fmt .Errorf ("write error: %w" , err )
173+ return
174+ }
98175
99- bufReader = bufio .NewReader (input )
176+ if readErr == io .EOF {
177+ break
178+ }
179+ }
180+ }()
100181
101- var bufWriter * bufio.Writer
182+ return readerOut , written , nil
183+ }
102184
103- bufWriter = bufio .NewWriterSize (pipeWriter , 4096 ) // 使用 pipeWriter
185+ // processLinks acts as a dispatcher, choosing the best processing strategy based on file size.
186+ // It uses a memory-safe streaming approach for large or unknown-size files,
187+ // and a high-performance buffered approach for smaller files.
188+ func processLinks (input io.ReadCloser , host string , cfg * config.Config , c * touka.Context , bodySize int ) (readerOut io.Reader , written int64 , err error ) {
189+ const sizeThreshold = 256 * 1024 // 256KB
190+
191+ // Use streaming for large or unknown size files to prevent OOM
192+ if bodySize == - 1 || bodySize > sizeThreshold {
193+ c .Debugf ("Using streaming processor for large/unknown size file (%d bytes)" , bodySize )
194+ return processLinksStreamingInternal (input , host , cfg , c )
195+ } else {
196+ c .Debugf ("Using buffered processor for small file (%d bytes)" , bodySize )
197+ return processLinksBufferedInternal (input , host , cfg , c )
198+ }
199+ }
200+
201+ // processLinksBufferedInternal a link processing function that reads the entire content into a buffer.
202+ // It is optimized for performance on smaller files but carries an OOM risk for large files.
203+ func processLinksBufferedInternal (input io.ReadCloser , host string , cfg * config.Config , c * touka.Context ) (readerOut io.Reader , written int64 , err error ) {
204+ pipeReader , pipeWriter := io .Pipe ()
205+ readerOut = pipeReader
206+ hostBytes := []byte (host )
104207
105- //确保writer关闭
208+ go func () {
209+ // 在 goroutine 退出时, 根据 err 是否为 nil, 带错误或正常关闭 pipeWriter
106210 defer func () {
107- if flushErr := bufWriter .Flush (); flushErr != nil {
108- c .Errorf ("writer flush failed %v" , flushErr )
109- // 如果已经存在错误，则保留。否则，记录此错误。
110- if err == nil {
111- err = flushErr
112- }
211+ if closeErr := input .Close (); closeErr != nil {
212+ c .Errorf ("input close failed: %v" , closeErr )
113213 }
114214 }()
115-
116- // 使用正则表达式匹配 http 和 https 链接
117- for {
118- line , readErr := bufReader .ReadString ('\n' )
119- if readErr != nil {
120- if readErr == io .EOF {
121- break // 文件结束
215+ defer func () {
216+ if err != nil {
217+ if closeErr := pipeWriter .CloseWithError (err ); closeErr != nil {
218+ c .Errorf ("pipeWriter close with error failed: %v" , closeErr )
219+ }
220+ } else {
221+ if closeErr := pipeWriter .Close (); closeErr != nil {
222+ c .Errorf ("pipeWriter close failed: %v" , closeErr )
122223 }
123- err = fmt .Errorf ("读取行错误: %v" , readErr ) // 传递错误
124- return // Goroutine 中使用 return 返回错误
125224 }
225+ }()
126226
127- // 替换所有匹配的 URL
128- modifiedLine := urlPattern .ReplaceAllStringFunc (line , func (originalURL string ) string {
129- return modifyURL (originalURL , host , cfg ) // 假设 modifyURL 函数已定义
130- })
227+ buf := bufferPool .Get ().(* bytes.Buffer )
228+ buf .Reset ()
229+ defer bufferPool .Put (buf )
131230
132- n , writeErr := bufWriter .WriteString (modifiedLine )
133- written += int64 (n ) // 更新写入的字节数
134- if writeErr != nil {
135- err = fmt .Errorf ("写入文件错误: %v" , writeErr ) // 传递错误
136- return // Goroutine 中使用 return 返回错误
137- }
231+ // 将全部输入读入复用的缓冲区
232+ if _ , err = buf .ReadFrom (input ); err != nil {
233+ err = fmt .Errorf ("reading input failed: %w" , err )
234+ return
138235 }
139236
140- // 在返回之前，再刷新一次 (虽然 defer 中已经有 flush，但这里再加一次确保及时刷新)
141- if flushErr := bufWriter .Flush (); flushErr != nil {
142- if err == nil { // 避免覆盖之前的错误
143- err = flushErr
144- }
145- return // Goroutine 中使用 return 返回错误
237+ // 使用 ReplaceAllFunc 和字节版本辅助函数, 实现准零分配
238+ modifiedBytes := urlPattern .ReplaceAllFunc (buf .Bytes (), func (originalURL []byte ) []byte {
239+ return modifyURLBytes (originalURL , hostBytes , cfg )
240+ })
241+
242+ // 将处理后的字节写回管道
243+ var n int
244+ n , err = pipeWriter .Write (modifiedBytes )
245+ if err != nil {
246+ err = fmt .Errorf ("writing to pipe failed: %w" , err )
247+ return
146248 }
249+ written = int64 (n )
147250 }()
148251
149- return readerOut , written , nil // 返回 reader 和 written，error 由 Goroutine 通过 pipeWriter.CloseWithError 传递
252+ return readerOut , written , nil
150253}
0 commit comments