mirror of
https://github.com/dstotijn/hetty.git
synced 2025-07-01 18:47:29 -04:00
Add intercept module
This commit is contained in:
David Stotijn
452
pkg/proxy/intercept/intercept.go
Normal file
452
pkg/proxy/intercept/intercept.go
Normal file
@ -0,0 +1,452 @@
|
||||
package intercept
|
||||
|
||||
import (
|
||||
"context"
|
||||
"errors"
|
||||
"fmt"
|
||||
"net/http"
|
||||
"sort"
|
||||
"sync"
|
||||
|
||||
"github.com/oklog/ulid"
|
||||
|
||||
"github.com/dstotijn/hetty/pkg/log"
|
||||
"github.com/dstotijn/hetty/pkg/proxy"
|
||||
"github.com/dstotijn/hetty/pkg/search"
|
||||
)
|
||||
|
||||
var (
|
||||
ErrRequestAborted = errors.New("intercept: request was aborted")
|
||||
ErrRequestNotFound = errors.New("intercept: request not found")
|
||||
ErrRequestDone = errors.New("intercept: request is done")
|
||||
ErrResponseNotFound = errors.New("intercept: response not found")
|
||||
)
|
||||
|
||||
type contextKey int
|
||||
|
||||
const interceptResponseKey contextKey = 0
|
||||
|
||||
// Request represents a server received HTTP request, alongside a channel for sending a modified version of it to the
|
||||
// routine that's awaiting it. Also contains a channel for receiving a cancellation signal.
|
||||
type Request struct {
|
||||
req *http.Request
|
||||
ch chan<- *http.Request
|
||||
done <-chan struct{}
|
||||
}
|
||||
|
||||
// Response represents an HTTP response from a proxied request, alongside a channel for sending a modified version of it
|
||||
// to the routine that's awaiting it. Also contains a channel for receiving a cancellation signal.
|
||||
type Response struct {
|
||||
res *http.Response
|
||||
ch chan<- *http.Response
|
||||
done <-chan struct{}
|
||||
}
|
||||
|
||||
type Item struct {
|
||||
Request *http.Request
|
||||
Response *http.Response
|
||||
}
|
||||
|
||||
type Service struct {
|
||||
reqMu *sync.RWMutex
|
||||
resMu *sync.RWMutex
|
||||
requests map[ulid.ULID]Request
|
||||
responses map[ulid.ULID]Response
|
||||
logger log.Logger
|
||||
|
||||
requestsEnabled bool
|
||||
responsesEnabled bool
|
||||
reqFilter search.Expression
|
||||
resFilter search.Expression
|
||||
}
|
||||
|
||||
type Config struct {
|
||||
Logger log.Logger
|
||||
RequestsEnabled bool
|
||||
ResponsesEnabled bool
|
||||
RequestFilter search.Expression
|
||||
ResponseFilter search.Expression
|
||||
}
|
||||
|
||||
// RequestIDs implements sort.Interface.
|
||||
type RequestIDs []ulid.ULID
|
||||
|
||||
func NewService(cfg Config) *Service {
|
||||
s := &Service{
|
||||
reqMu: &sync.RWMutex{},
|
||||
resMu: &sync.RWMutex{},
|
||||
requests: make(map[ulid.ULID]Request),
|
||||
responses: make(map[ulid.ULID]Response),
|
||||
logger: cfg.Logger,
|
||||
requestsEnabled: cfg.RequestsEnabled,
|
||||
responsesEnabled: cfg.ResponsesEnabled,
|
||||
reqFilter: cfg.RequestFilter,
|
||||
resFilter: cfg.ResponseFilter,
|
||||
}
|
||||
|
||||
if s.logger == nil {
|
||||
s.logger = log.NewNopLogger()
|
||||
}
|
||||
|
||||
return s
|
||||
}
|
||||
|
||||
// RequestModifier is a proxy.RequestModifyMiddleware for intercepting HTTP requests.
|
||||
func (svc *Service) RequestModifier(next proxy.RequestModifyFunc) proxy.RequestModifyFunc {
|
||||
return func(req *http.Request) {
|
||||
// This is a blocking operation, that gets unblocked when either a modified request is returned or an error
|
||||
// (typically `context.Canceled`).
|
||||
modifiedReq, err := svc.InterceptRequest(req.Context(), req)
|
||||
|
||||
switch {
|
||||
case errors.Is(err, ErrRequestAborted):
|
||||
svc.logger.Debugw("Stopping intercept, request was aborted.")
|
||||
// Prevent further processing by replacing req.Context with a cancelled context value.
|
||||
// This will cause the http.Roundtripper in the `proxy` package to
|
||||
// handle this request as an error.
|
||||
ctx, cancel := context.WithCancel(context.Background())
|
||||
cancel()
|
||||
|
||||
*req = *req.WithContext(ctx)
|
||||
case errors.Is(err, context.Canceled):
|
||||
svc.logger.Debugw("Stopping intercept, context was cancelled.")
|
||||
case err != nil:
|
||||
svc.logger.Errorw("Failed to intercept request.",
|
||||
"error", err)
|
||||
default:
|
||||
*req = *modifiedReq
|
||||
next(req)
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// InterceptRequest adds an HTTP request to an array of pending intercepted requests, alongside channels used for
|
||||
// sending a cancellation signal and receiving a modified request. It's safe for concurrent use.
|
||||
func (svc *Service) InterceptRequest(ctx context.Context, req *http.Request) (*http.Request, error) {
|
||||
reqID, ok := proxy.RequestIDFromContext(ctx)
|
||||
if !ok {
|
||||
svc.logger.Errorw("Failed to intercept: context doesn't have an ID.")
|
||||
return req, nil
|
||||
}
|
||||
|
||||
if !svc.requestsEnabled {
|
||||
// If request intercept is disabled, return the incoming request as-is.
|
||||
svc.logger.Debugw("Bypassed request interception: feature disabled.")
|
||||
return req, nil
|
||||
}
|
||||
|
||||
if svc.reqFilter != nil {
|
||||
match, err := MatchRequestFilter(req, svc.reqFilter)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("intercept: failed to match request rules for request (id: %v): %w",
|
||||
reqID.String(), err,
|
||||
)
|
||||
}
|
||||
|
||||
if !match {
|
||||
svc.logger.Debugw("Bypassed request interception: request rules don't match.")
|
||||
return req, nil
|
||||
}
|
||||
}
|
||||
|
||||
ch := make(chan *http.Request)
|
||||
done := make(chan struct{})
|
||||
|
||||
svc.reqMu.Lock()
|
||||
svc.requests[reqID] = Request{
|
||||
req: req,
|
||||
ch: ch,
|
||||
done: done,
|
||||
}
|
||||
svc.reqMu.Unlock()
|
||||
|
||||
// Whatever happens next (modified request returned, or a context cancelled error), any blocked channel senders
|
||||
// should be unblocked, and the request should be removed from the requests queue.
|
||||
defer func() {
|
||||
close(done)
|
||||
svc.reqMu.Lock()
|
||||
defer svc.reqMu.Unlock()
|
||||
delete(svc.requests, reqID)
|
||||
}()
|
||||
|
||||
select {
|
||||
case modReq := <-ch:
|
||||
if modReq == nil {
|
||||
return nil, ErrRequestAborted
|
||||
}
|
||||
|
||||
return modReq, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// ModifyRequest sends a modified HTTP request to the related channel, or returns ErrRequestDone when the request was
|
||||
// cancelled. It's safe for concurrent use.
|
||||
func (svc *Service) ModifyRequest(reqID ulid.ULID, modReq *http.Request, modifyResponse *bool) error {
|
||||
svc.reqMu.RLock()
|
||||
req, ok := svc.requests[reqID]
|
||||
svc.reqMu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
return ErrRequestNotFound
|
||||
}
|
||||
|
||||
*modReq = *modReq.WithContext(req.req.Context())
|
||||
if modifyResponse != nil {
|
||||
*modReq = *modReq.WithContext(WithInterceptResponse(modReq.Context(), *modifyResponse))
|
||||
}
|
||||
|
||||
select {
|
||||
case <-req.done:
|
||||
return ErrRequestDone
|
||||
case req.ch <- modReq:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// CancelRequest ensures an intercepted request is dropped.
|
||||
func (svc *Service) CancelRequest(reqID ulid.ULID) error {
|
||||
return svc.ModifyRequest(reqID, nil, nil)
|
||||
}
|
||||
|
||||
func (svc *Service) ClearRequests() {
|
||||
svc.reqMu.Lock()
|
||||
defer svc.reqMu.Unlock()
|
||||
|
||||
for _, req := range svc.requests {
|
||||
select {
|
||||
case <-req.done:
|
||||
case req.ch <- nil:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
func (svc *Service) ClearResponses() {
|
||||
svc.resMu.Lock()
|
||||
defer svc.resMu.Unlock()
|
||||
|
||||
for _, res := range svc.responses {
|
||||
select {
|
||||
case <-res.done:
|
||||
case res.ch <- nil:
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
// Items returns a list of pending items (requests and responses). It's safe for concurrent use.
|
||||
func (svc *Service) Items() []Item {
|
||||
svc.reqMu.RLock()
|
||||
defer svc.reqMu.RUnlock()
|
||||
|
||||
svc.resMu.RLock()
|
||||
defer svc.resMu.RUnlock()
|
||||
|
||||
reqIDs := make([]ulid.ULID, 0, len(svc.requests)+len(svc.responses))
|
||||
|
||||
for id := range svc.requests {
|
||||
reqIDs = append(reqIDs, id)
|
||||
}
|
||||
|
||||
for id := range svc.responses {
|
||||
reqIDs = append(reqIDs, id)
|
||||
}
|
||||
|
||||
sort.Sort(RequestIDs(reqIDs))
|
||||
|
||||
items := make([]Item, len(reqIDs))
|
||||
|
||||
for i, id := range reqIDs {
|
||||
item := Item{}
|
||||
|
||||
if req, ok := svc.requests[id]; ok {
|
||||
item.Request = req.req
|
||||
}
|
||||
|
||||
if res, ok := svc.responses[id]; ok {
|
||||
item.Response = res.res
|
||||
}
|
||||
|
||||
items[i] = item
|
||||
}
|
||||
|
||||
return items
|
||||
}
|
||||
|
||||
func (svc *Service) UpdateSettings(settings Settings) {
|
||||
// When updating from requests `enabled` -> `disabled`, clear any pending reqs.
|
||||
if svc.requestsEnabled && !settings.RequestsEnabled {
|
||||
svc.ClearRequests()
|
||||
}
|
||||
|
||||
// When updating from responses `enabled` -> `disabled`, clear any pending responses.
|
||||
if svc.responsesEnabled && !settings.ResponsesEnabled {
|
||||
svc.ClearResponses()
|
||||
}
|
||||
|
||||
svc.requestsEnabled = settings.RequestsEnabled
|
||||
svc.responsesEnabled = settings.ResponsesEnabled
|
||||
svc.reqFilter = settings.RequestFilter
|
||||
svc.resFilter = settings.ResponseFilter
|
||||
}
|
||||
|
||||
// ItemByID returns an intercepted item (request and possible response) by ID. It's safe for concurrent use.
|
||||
func (svc *Service) ItemByID(id ulid.ULID) (Item, error) {
|
||||
svc.reqMu.RLock()
|
||||
defer svc.reqMu.RUnlock()
|
||||
|
||||
svc.resMu.RLock()
|
||||
defer svc.resMu.RUnlock()
|
||||
|
||||
item := Item{}
|
||||
found := false
|
||||
|
||||
if req, ok := svc.requests[id]; ok {
|
||||
item.Request = req.req
|
||||
found = true
|
||||
}
|
||||
|
||||
if res, ok := svc.responses[id]; ok {
|
||||
item.Response = res.res
|
||||
found = true
|
||||
}
|
||||
|
||||
if !found {
|
||||
return Item{}, ErrRequestNotFound
|
||||
}
|
||||
|
||||
return item, nil
|
||||
}
|
||||
|
||||
func (ids RequestIDs) Len() int {
|
||||
return len(ids)
|
||||
}
|
||||
|
||||
func (ids RequestIDs) Less(i, j int) bool {
|
||||
return ids[i].Compare(ids[j]) == -1
|
||||
}
|
||||
|
||||
func (ids RequestIDs) Swap(i, j int) {
|
||||
ids[i], ids[j] = ids[j], ids[i]
|
||||
}
|
||||
|
||||
func WithInterceptResponse(ctx context.Context, value bool) context.Context {
|
||||
return context.WithValue(ctx, interceptResponseKey, value)
|
||||
}
|
||||
|
||||
func ShouldInterceptResponseFromContext(ctx context.Context) (bool, bool) {
|
||||
shouldIntercept, ok := ctx.Value(interceptResponseKey).(bool)
|
||||
return shouldIntercept, ok
|
||||
}
|
||||
|
||||
// ResponseModifier is a proxy.ResponseModifyMiddleware for intercepting HTTP responses.
|
||||
func (svc *Service) ResponseModifier(next proxy.ResponseModifyFunc) proxy.ResponseModifyFunc {
|
||||
return func(res *http.Response) error {
|
||||
// This is a blocking operation, that gets unblocked when either a modified response is returned or an error.
|
||||
//nolint:bodyclose
|
||||
modifiedRes, err := svc.InterceptResponse(res.Request.Context(), res)
|
||||
if err != nil {
|
||||
return fmt.Errorf("failed to intercept response: %w", err)
|
||||
}
|
||||
|
||||
*res = *modifiedRes
|
||||
|
||||
return next(res)
|
||||
}
|
||||
}
|
||||
|
||||
// InterceptResponse adds an HTTP response to an array of pending intercepted responses, alongside channels used for
|
||||
// sending a cancellation signal and receiving a modified response. It's safe for concurrent use.
|
||||
func (svc *Service) InterceptResponse(ctx context.Context, res *http.Response) (*http.Response, error) {
|
||||
reqID, ok := proxy.RequestIDFromContext(ctx)
|
||||
if !ok {
|
||||
svc.logger.Errorw("Failed to intercept: context doesn't have an ID.")
|
||||
return res, nil
|
||||
}
|
||||
|
||||
shouldIntercept, ok := ShouldInterceptResponseFromContext(ctx)
|
||||
if ok && !shouldIntercept {
|
||||
// If the related request explicitly disabled response intercept, return the response as-is.
|
||||
svc.logger.Debugw("Bypassed response interception: related request explicitly disabled response intercept.")
|
||||
return res, nil
|
||||
}
|
||||
|
||||
// If global response intercept is disabled and interception is *not* explicitly enabled for this response: bypass.
|
||||
if !svc.responsesEnabled && !(ok && shouldIntercept) {
|
||||
svc.logger.Debugw("Bypassed response interception: feature disabled.")
|
||||
return res, nil
|
||||
}
|
||||
|
||||
if svc.resFilter != nil {
|
||||
match, err := MatchResponseFilter(res, svc.resFilter)
|
||||
if err != nil {
|
||||
return nil, fmt.Errorf("intercept: failed to match response rules for response (id: %v): %w",
|
||||
reqID.String(), err,
|
||||
)
|
||||
}
|
||||
|
||||
if !match {
|
||||
svc.logger.Debugw("Bypassed response interception: response rules don't match.")
|
||||
return res, nil
|
||||
}
|
||||
}
|
||||
|
||||
ch := make(chan *http.Response)
|
||||
done := make(chan struct{})
|
||||
|
||||
svc.resMu.Lock()
|
||||
svc.responses[reqID] = Response{
|
||||
res: res,
|
||||
ch: ch,
|
||||
done: done,
|
||||
}
|
||||
svc.resMu.Unlock()
|
||||
|
||||
// Whatever happens next (modified response returned, or a context cancelled error), any blocked channel senders
|
||||
// should be unblocked, and the response should be removed from the responses queue.
|
||||
defer func() {
|
||||
close(done)
|
||||
svc.resMu.Lock()
|
||||
defer svc.resMu.Unlock()
|
||||
delete(svc.responses, reqID)
|
||||
}()
|
||||
|
||||
select {
|
||||
case modRes := <-ch:
|
||||
if modRes == nil {
|
||||
return nil, ErrRequestAborted
|
||||
}
|
||||
|
||||
return modRes, nil
|
||||
case <-ctx.Done():
|
||||
return nil, ctx.Err()
|
||||
}
|
||||
}
|
||||
|
||||
// ModifyResponse sends a modified HTTP response to the related channel, or returns ErrRequestDone when the related
|
||||
// request was cancelled. It's safe for concurrent use.
|
||||
func (svc *Service) ModifyResponse(reqID ulid.ULID, modRes *http.Response) error {
|
||||
svc.resMu.RLock()
|
||||
res, ok := svc.responses[reqID]
|
||||
svc.resMu.RUnlock()
|
||||
|
||||
if !ok {
|
||||
return ErrRequestNotFound
|
||||
}
|
||||
|
||||
if modRes != nil {
|
||||
modRes.Request = res.res.Request
|
||||
}
|
||||
|
||||
select {
|
||||
case <-res.done:
|
||||
return ErrRequestDone
|
||||
case res.ch <- modRes:
|
||||
return nil
|
||||
}
|
||||
}
|
||||
|
||||
// CancelResponse ensures an intercepted response is dropped.
|
||||
func (svc *Service) CancelResponse(reqID ulid.ULID) error {
|
||||
return svc.ModifyResponse(reqID, nil)
|
||||
}
|
||||
Reference in New Issue
Block a user