|
1 | 1 | package detectors |
2 | 2 |
|
3 | 3 | import ( |
| 4 | + "bytes" |
4 | 5 | "context" |
5 | 6 | "errors" |
| 7 | + "fmt" |
| 8 | + "io" |
6 | 9 | "net" |
7 | 10 | "net/http" |
8 | 11 | "slices" |
9 | 12 | "sync" |
10 | 13 | "time" |
11 | 14 |
|
| 15 | + "golang.org/x/sync/singleflight" |
| 16 | + |
12 | 17 | "github.com/trufflesecurity/trufflehog/v3/pkg/common" |
13 | 18 | "github.com/trufflesecurity/trufflehog/v3/pkg/feature" |
14 | 19 | ) |
@@ -175,3 +180,93 @@ func NewDetectorHttpClient(opts ...ClientOption) *http.Client { |
175 | 180 | client.Transport = common.NewInstrumentedTransport(client.Transport) |
176 | 181 | return client |
177 | 182 | } |
| 183 | + |
| 184 | +// bufferedResponse holds a fully-read HTTP response so it can be replayed to |
| 185 | +// every goroutine that was coalesced by singleflight. |
| 186 | +type bufferedResponse struct { |
| 187 | + statusCode int |
| 188 | + header http.Header |
| 189 | + body []byte |
| 190 | +} |
| 191 | + |
| 192 | +// singleflightTransport is an http.RoundTripper that coalesces concurrent requests |
| 193 | +// sharing the same deduplication key into a single network call. It is a no-op for |
| 194 | +// requests whose context does not carry a dedup key. |
| 195 | +type singleflightTransport struct { |
| 196 | + base http.RoundTripper |
| 197 | + group singleflight.Group |
| 198 | +} |
| 199 | + |
| 200 | +func (t *singleflightTransport) RoundTrip(req *http.Request) (*http.Response, error) { |
| 201 | + key, ok := req.Context().Value(dedupKeyContextKey{}).(string) |
| 202 | + if !ok || key == "" { |
| 203 | + return t.base.RoundTrip(req) |
| 204 | + } |
| 205 | + |
| 206 | + // DoChan is used instead of Do so each caller can independently respect its |
| 207 | + // own context cancellation without blocking on the shared in-flight call. |
| 208 | + ch := t.group.DoChan(key, func() (any, error) { |
| 209 | + // Detach the in-flight request from the first caller's cancellation so |
| 210 | + // that one goroutine timing out doesn't abort the shared network call |
| 211 | + // and propagate an error to all coalesced waiters. |
| 212 | + // |
| 213 | + // context.WithoutCancel also strips any deadline (e.g. from |
| 214 | + // http.Client.Timeout), so we re-attach the original deadline if |
| 215 | + // present. Without this the shared request has no timeout and a |
| 216 | + // hanging server would leak the goroutine and pin the singleflight |
| 217 | + // key indefinitely. |
| 218 | + sharedCtx := context.WithoutCancel(req.Context()) |
| 219 | + if deadline, ok := req.Context().Deadline(); ok { |
| 220 | + var cancel context.CancelFunc |
| 221 | + sharedCtx, cancel = context.WithDeadline(sharedCtx, deadline) |
| 222 | + defer cancel() |
| 223 | + } |
| 224 | + sharedReq := req.WithContext(sharedCtx) |
| 225 | + resp, err := t.base.RoundTrip(sharedReq) |
| 226 | + if err != nil { |
| 227 | + return nil, err |
| 228 | + } |
| 229 | + defer resp.Body.Close() |
| 230 | + |
| 231 | + body, err := io.ReadAll(resp.Body) |
| 232 | + if err != nil { |
| 233 | + return nil, err |
| 234 | + } |
| 235 | + |
| 236 | + return &bufferedResponse{ |
| 237 | + statusCode: resp.StatusCode, |
| 238 | + header: resp.Header.Clone(), |
| 239 | + body: body, |
| 240 | + }, nil |
| 241 | + }) |
| 242 | + |
| 243 | + select { |
| 244 | + case result := <-ch: |
| 245 | + if result.Err != nil { |
| 246 | + return nil, result.Err |
| 247 | + } |
| 248 | + br := result.Val.(*bufferedResponse) |
| 249 | + return &http.Response{ |
| 250 | + StatusCode: br.statusCode, |
| 251 | + Status: fmt.Sprintf("%d %s", br.statusCode, http.StatusText(br.statusCode)), |
| 252 | + Header: br.header.Clone(), |
| 253 | + Body: io.NopCloser(bytes.NewReader(br.body)), |
| 254 | + }, nil |
| 255 | + case <-req.Context().Done(): |
| 256 | + return nil, req.Context().Err() |
| 257 | + } |
| 258 | +} |
| 259 | + |
| 260 | +// NewClientWithDedup wraps base with a transport that deduplicates concurrent |
| 261 | +// verification requests sharing the same key. Detectors opt in per credential by |
| 262 | +// calling WithDedupKey on the request context before client.Do — no other changes |
| 263 | +// to request building or response reading are needed. |
| 264 | +func NewClientWithDedup(base *http.Client) *http.Client { |
| 265 | + clone := *base |
| 266 | + transport := base.Transport |
| 267 | + if transport == nil { |
| 268 | + transport = http.DefaultTransport |
| 269 | + } |
| 270 | + clone.Transport = &singleflightTransport{base: transport} |
| 271 | + return &clone |
| 272 | +} |
0 commit comments