From 642736126d7eedff3959e39e560c6c8d39dd71aa Mon Sep 17 00:00:00 2001 From: Kevin Pham Date: Sun, 26 Nov 2023 23:12:27 -0600 Subject: [PATCH] aider: Add public API to add a RequestModification to fire only once. # Aider chat conversation: USER: Modify proxychain.go such that there is a public API to add a RequestModification to fire only once. ASSISTANT: To make this change we need to modify `proxychain/proxychain.go` to: 1. Add a new field `onceRequestModifications` in the `ProxyChain` struct to store the request modifications that should only fire once. 2. Add a new method `AddOnceRequestModification` to add a request modification that should only fire once. 3. Update the `_execute` method to apply the `onceRequestModifications` and clear them after they have been applied. Here are the *SEARCH/REPLACE* blocks: proxychain/proxychain.go <<<<<<< SEARCH type ProxyChain struct { Context *fiber.Ctx Client *http.Client Request *http.Request Response *http.Response requestModifications []RequestModification resultModifications []ResponseModification htmlTokenRewriters []rr.IHTMLTokenRewriter Ruleset *ruleset.RuleSet debugMode bool abortErr error } ======= type ProxyChain struct { Context *fiber.Ctx Client *http.Client Request *http.Request Response *http.Response requestModifications []RequestModification onceRequestModifications []RequestModification resultModifications []ResponseModification htmlTokenRewriters []rr.IHTMLTokenRewriter Ruleset *ruleset.RuleSet debugMode bool abortErr error } >>>>>>> REPLACE proxychain/proxychain.go <<<<<<< SEARCH // AddRequestModifications sets the ProxyChain's request modifers // the modifier will not fire until ProxyChain.Execute() is run. func (chain *ProxyChain) AddRequestModifications(mods ...RequestModification) *ProxyChain { chain.requestModifications = append(chain.requestModifications, mods...) return chain } ======= // AddRequestModifications adds more request modifers to the ProxyChain // the modifier will not fire until ProxyChain.Execute() is run. func (chain *ProxyChain) AddRequestModifications(mods ...RequestModification) *ProxyChain { chain.requestModifications = append(chain.requestModifications, mods...) return chain } // AddOnceRequestModification adds a request modifier to the ProxyChain that should only fire once // the modifier will not fire until ProxyChain.Execute() is run and will be removed after it has been applied. func (chain *ProxyChain) AddOnceRequestModification(mod ...RequestModification) *ProxyChain { chain.onceRequestModifications = append(chain.onceRequestModifications, mod...) return chain } >>>>>>> REPLACE proxychain/proxychain.go <<<<<<< SEARCH // _execute sends the request for the ProxyChain and returns the raw body only // the caller is responsible for returning a response back to the requestor // the caller is also responsible for calling chain._reset() when they are done with the body func (chain *ProxyChain) _execute() (io.Reader, error) { if chain.validateCtxIsSet() != nil || chain.abortErr != nil { return nil, chain.abortErr } if chain.Request == nil { return nil, errors.New("proxychain request not yet initialized") } if chain.Request.URL.Scheme == "" { return nil, errors.New("request url not set or invalid. Check ProxyChain ReqMods for issues") } // Apply requestModifications to proxychain for _, applyRequestModificationsTo := range chain.requestModifications { err := applyRequestModificationsTo(chain) if err != nil { return nil, chain.abort(err) } } ======= // _execute sends the request for the ProxyChain and returns the raw body only // the caller is responsible for returning a response back to the requestor // the caller is also responsible for calling chain._reset() when they are done with the body func (chain *ProxyChain) _execute() (io.Reader, error) { if chain.validateCtxIsSet() != nil || chain.abortErr != nil { return nil, chain.abortErr } if chain.Request == nil { return nil, errors.New("proxychain request not yet initialized") } if chain.Request.URL.Scheme == "" { return nil, errors.New("request url not set or invalid. Check ProxyChain ReqMods for issues") } // Apply requestModifications to proxychain for _, applyRequestModificationsTo := range chain.requestModifications { err := applyRequestModificationsTo(chain) if err != nil { return nil, chain.abort(err) } } // Apply onceRequestModifications to proxychain and clear them for _, applyOnceRequestModificationsTo := range chain.onceRequestModifications { err := applyOnceRequestModificationsTo(chain) if err != nil { return nil, chain.abort(err) } } chain.onceRequestModifications = nil >>>>>>> REPLACE --- proxychain/proxychain.go | 39 ++++++++++++++++++++++++++++----------- 1 file changed, 28 insertions(+), 11 deletions(-) diff --git a/proxychain/proxychain.go b/proxychain/proxychain.go index 0b36224..506f06a 100644 --- a/proxychain/proxychain.go +++ b/proxychain/proxychain.go @@ -83,16 +83,17 @@ proxychain.NewProxyChain(). └─────────┘ └────────────────────────┘ └─────────┘ */ type ProxyChain struct { - Context *fiber.Ctx - Client *http.Client - Request *http.Request - Response *http.Response - requestModifications []RequestModification - resultModifications []ResponseModification - htmlTokenRewriters []rr.IHTMLTokenRewriter - Ruleset *ruleset.RuleSet - debugMode bool - abortErr error + Context *fiber.Ctx + Client *http.Client + Request *http.Request + Response *http.Response + requestModifications []RequestModification + onceRequestModifications []RequestModification + resultModifications []ResponseModification + htmlTokenRewriters []rr.IHTMLTokenRewriter + Ruleset *ruleset.RuleSet + debugMode bool + abortErr error } // a ProxyStrategy is a pre-built proxychain with purpose-built defaults @@ -113,13 +114,20 @@ func (chain *ProxyChain) SetRequestModifications(mods ...RequestModification) *P return chain } -// AddRequestModifications sets the ProxyChain's request modifers +// AddRequestModifications adds more request modifers to the ProxyChain // the modifier will not fire until ProxyChain.Execute() is run. func (chain *ProxyChain) AddRequestModifications(mods ...RequestModification) *ProxyChain { chain.requestModifications = append(chain.requestModifications, mods...) return chain } +// AddOnceRequestModification adds a request modifier to the ProxyChain that should only fire once +// the modifier will not fire until ProxyChain.Execute() is run and will be removed after it has been applied. +func (chain *ProxyChain) AddOnceRequestModification(mod ...RequestModification) *ProxyChain { + chain.onceRequestModifications = append(chain.onceRequestModifications, mod...) + return chain +} + // AddResponseModifications sets the ProxyChain's response modifers // the modifier will not fire until ProxyChain.Execute() is run. func (chain *ProxyChain) AddResponseModifications(mods ...ResponseModification) *ProxyChain { @@ -357,6 +365,15 @@ func (chain *ProxyChain) _execute() (io.Reader, error) { } } + // Apply onceRequestModifications to proxychain and clear them + for _, applyOnceRequestModificationsTo := range chain.onceRequestModifications { + err := applyOnceRequestModificationsTo(chain) + if err != nil { + return nil, chain.abort(err) + } + } + chain.onceRequestModifications = nil + // Send Request Upstream resp, err := chain.Client.Do(chain.Request) if err != nil {