106 lines
2.7 KiB
Go
106 lines
2.7 KiB
Go
package responsemodifers
|
|
|
|
import (
|
|
"bytes"
|
|
"io"
|
|
"ladder/proxychain"
|
|
"log"
|
|
"net/url"
|
|
"strings"
|
|
|
|
"golang.org/x/net/html"
|
|
)
|
|
|
|
type HTMLResourceURLRewriter struct {
|
|
proxyURL *url.URL // proxyURL is the URL of the proxy, not the upstream URL; TODO: implement
|
|
tokenizer *html.Tokenizer
|
|
currentToken html.Token
|
|
tokenBuffer *bytes.Buffer
|
|
currentTokenIndex int
|
|
currentTokenProcessed bool
|
|
}
|
|
|
|
func NewHTMLResourceURLRewriter(src io.ReadCloser, proxyURL *url.URL) *HTMLResourceURLRewriter {
|
|
log.Println("tokenize")
|
|
return &HTMLResourceURLRewriter{
|
|
tokenizer: html.NewTokenizer(src),
|
|
currentToken: html.Token{},
|
|
currentTokenIndex: 0,
|
|
tokenBuffer: new(bytes.Buffer),
|
|
proxyURL: proxyURL,
|
|
}
|
|
}
|
|
|
|
func (r *HTMLResourceURLRewriter) Close() error {
|
|
r.tokenBuffer.Reset()
|
|
r.currentToken = html.Token{}
|
|
r.currentTokenIndex = 0
|
|
r.currentTokenProcessed = false
|
|
return nil
|
|
}
|
|
|
|
func (r *HTMLResourceURLRewriter) Read(p []byte) (int, error) {
|
|
|
|
if r.currentToken.Data == "" || r.currentTokenProcessed {
|
|
tokenType := r.tokenizer.Next()
|
|
|
|
// done reading html, close out reader
|
|
if tokenType == html.ErrorToken {
|
|
if r.tokenizer.Err() == io.EOF {
|
|
return 0, io.EOF
|
|
}
|
|
return 0, r.tokenizer.Err()
|
|
}
|
|
|
|
// flush the current token into an internal buffer
|
|
// to handle fragmented tokens
|
|
r.currentToken = r.tokenizer.Token()
|
|
r.tokenBuffer.Reset()
|
|
r.tokenBuffer.WriteString(r.currentToken.String())
|
|
r.currentTokenProcessed = false
|
|
r.currentTokenIndex = 0
|
|
}
|
|
|
|
n, err := r.tokenBuffer.Read(p)
|
|
|
|
if err == io.EOF || r.tokenBuffer.Len() == 0 {
|
|
r.currentTokenProcessed = true
|
|
err = nil // Reset error to nil because EOF in this context is expected and not an actual error
|
|
}
|
|
return n, err
|
|
|
|
}
|
|
|
|
// RewriteHTMLResourceURLs updates src/href attributes in HTML content to route through the proxy.
|
|
func RewriteHTMLResourceURLs() proxychain.ResponseModification {
|
|
return func(chain *proxychain.ProxyChain) error {
|
|
log.Println("rhru")
|
|
ct := chain.Response.Header.Get("content-type")
|
|
log.Println(ct)
|
|
if !strings.HasPrefix(ct, "text/html") {
|
|
return nil
|
|
}
|
|
log.Println("rhru2")
|
|
// chain.Response.Body is an unread http.Response.Body
|
|
chain.Response.Body = NewHTMLResourceURLRewriter(chain.Response.Body, chain.Request.URL)
|
|
return nil
|
|
}
|
|
}
|
|
|
|
func rewriteToken(token *html.Token, baseURL *url.URL) {
|
|
log.Println(token.String())
|
|
attrsToRewrite := map[string]bool{"href": true, "src": true, "action": true, "srcset": true}
|
|
for i := range token.Attr {
|
|
attr := &token.Attr[i]
|
|
if attrsToRewrite[attr.Key] {
|
|
attr.Val = "/" + attr.Val
|
|
}
|
|
/*
|
|
if attrsToRewrite[attr.Key] && strings.HasPrefix(attr.Val, "/") {
|
|
// Make URL absolute
|
|
attr.Val = "/https://" + baseURL.Host + attr.Val
|
|
}
|
|
*/
|
|
}
|
|
}
|