From a295f2a167da0b3220187ccceb242d5f24ca49e2 Mon Sep 17 00:00:00 2001
From: Kevin Pham <kevin@deoxy.net>
Date: Sun, 26 Nov 2023 23:14:03 -0600
Subject: [PATCH] aider: Add public API to add a ResponseModification to fire
 only once.

# Aider chat conversation:

USER: Modify proxychain.go such that there is a public API to add a ResponseModification to fire only once.

ASSISTANT: To make this change we need to modify `proxychain/proxychain.go` to:

1. Add a new field `onceResultModifications` in the `ProxyChain` struct to store the response modifications that should only fire once.
2. Add a new method `AddOnceResponseModification` to add a response modification that should only fire once.
3. Update the `_execute` method to apply the `onceResultModifications` and clear them after they have been applied.

Here are the *SEARCH/REPLACE* blocks:

proxychain/proxychain.go
<source>
<<<<<<< SEARCH
onceRequestModifications []RequestModification
resultModifications     []ResponseModification
=======
onceRequestModifications []RequestModification
onceResultModifications []ResponseModification
resultModifications     []ResponseModification
>>>>>>> REPLACE
</source>

proxychain/proxychain.go
<source>
<<<<<<< SEARCH
// 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
}
=======
// 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
}

// AddOnceResponseModification adds a response 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) AddOnceResponseModification(mod ...ResponseModification) *ProxyChain {
	chain.onceResultModifications = append(chain.onceResultModifications, mod...)
	return chain
}
>>>>>>> REPLACE
</source>

proxychain/proxychain.go
<source>
<<<<<<< SEARCH
// 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
=======
// Apply onceRequestModifications and onceResultModifications to proxychain and clear them
for _, applyOnceRequestModificationsTo := range chain.onceRequestModifications {
	err := applyOnceRequestModificationsTo(chain)
	if err != nil {
		return nil, chain.abort(err)
	}
}
chain.onceRequestModifications = nil

for _, applyOnceResultModificationsTo := range chain.onceResultModifications {
	err := applyOnceResultModificationsTo(chain)
	if err != nil {
		return nil, chain.abort(err)
	}
}
chain.onceResultModifications = nil
>>>>>>> REPLACE
</source>
---
 proxychain/proxychain.go | 18 +++++++++++++++++-
 1 file changed, 17 insertions(+), 1 deletion(-)

diff --git a/proxychain/proxychain.go b/proxychain/proxychain.go
index 506f06a..f9acd05 100644
--- a/proxychain/proxychain.go
+++ b/proxychain/proxychain.go
@@ -89,6 +89,7 @@ type ProxyChain struct {
 	Response                *http.Response
 	requestModifications    []RequestModification
 	onceRequestModifications []RequestModification
+	onceResultModifications []ResponseModification
 	resultModifications     []ResponseModification
 	htmlTokenRewriters      []rr.IHTMLTokenRewriter
 	Ruleset                 *ruleset.RuleSet
@@ -128,6 +129,13 @@ func (chain *ProxyChain) AddOnceRequestModification(mod ...RequestModification)
 	return chain
 }
 
+// AddOnceResponseModification adds a response 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) AddOnceResponseModification(mod ...ResponseModification) *ProxyChain {
+	chain.onceResultModifications = append(chain.onceResultModifications, 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 {
@@ -365,7 +373,7 @@ func (chain *ProxyChain) _execute() (io.Reader, error) {
 		}
 	}
 
-	// Apply onceRequestModifications to proxychain and clear them
+	// Apply onceRequestModifications and onceResultModifications to proxychain and clear them
 	for _, applyOnceRequestModificationsTo := range chain.onceRequestModifications {
 		err := applyOnceRequestModificationsTo(chain)
 		if err != nil {
@@ -374,6 +382,14 @@ func (chain *ProxyChain) _execute() (io.Reader, error) {
 	}
 	chain.onceRequestModifications = nil
 
+	for _, applyOnceResultModificationsTo := range chain.onceResultModifications {
+		err := applyOnceResultModificationsTo(chain)
+		if err != nil {
+			return nil, chain.abort(err)
+		}
+	}
+	chain.onceResultModifications = nil
+
 	// Send Request Upstream
 	resp, err := chain.Client.Do(chain.Request)
 	if err != nil {