Add initial support for HTTP request intercepting

pkg/proxy/intercept/intercept.go

@@ -0,0 +1,189 @@
+package intercept
+
+import (
+	"context"
+	"errors"
+	"net/http"
+	"sort"
+	"sync"
+
+	"github.com/oklog/ulid"
+
+	"github.com/dstotijn/hetty/pkg/log"
+	"github.com/dstotijn/hetty/pkg/proxy"
+)
+
+var (
+	ErrRequestAborted  = errors.New("intercept: request was aborted")
+	ErrRequestNotFound = errors.New("intercept: request not found")
+	ErrRequestDone     = errors.New("intercept: request is done")
+)
+
+// 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{}
+}
+
+type Service struct {
+	mu       *sync.RWMutex
+	requests map[ulid.ULID]Request
+	logger   log.Logger
+}
+
+type Config struct {
+	Logger log.Logger
+}
+
+// RequestIDs implements sort.Interface.
+type RequestIDs []ulid.ULID
+
+func NewService(cfg Config) *Service {
+	s := &Service{
+		mu:       &sync.RWMutex{},
+		requests: make(map[ulid.ULID]Request),
+		logger:   cfg.Logger,
+	}
+
+	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.Intercept(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.WithContext(req.Context())
+			next(req)
+		}
+	}
+}
+
+// Intercept 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) Intercept(ctx context.Context, req *http.Request) (*http.Request, error) {
+	reqID, ok := proxy.RequestIDFromContext(ctx)
+	if !ok {
+		svc.logger.Errorw("Failed to intercept: request doesn't have an ID.")
+		return req, nil
+	}
+
+	ch := make(chan *http.Request)
+	done := make(chan struct{})
+
+	svc.mu.Lock()
+	svc.requests[reqID] = Request{
+		req:  req,
+		ch:   ch,
+		done: done,
+	}
+	svc.mu.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.mu.Lock()
+		defer svc.mu.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) error {
+	svc.mu.RLock()
+	req, ok := svc.requests[reqID]
+	svc.mu.RUnlock()
+
+	if !ok {
+		return ErrRequestNotFound
+	}
+
+	select {
+	case <-req.done:
+		return ErrRequestDone
+	case req.ch <- modReq:
+		return nil
+	}
+}
+
+func (svc *Service) ClearRequests() {
+	svc.mu.Lock()
+	defer svc.mu.Unlock()
+
+	for _, req := range svc.requests {
+		select {
+		case <-req.done:
+		case req.ch <- nil:
+		}
+	}
+}
+
+// Requests returns a list of pending intercepted requests. It's safe for concurrent use.
+func (svc *Service) Requests() []*http.Request {
+	svc.mu.RLock()
+	defer svc.mu.RUnlock()
+
+	ids := make([]ulid.ULID, 0, len(svc.requests))
+	for id := range svc.requests {
+		ids = append(ids, id)
+	}
+
+	sort.Sort(RequestIDs(ids))
+
+	reqs := make([]*http.Request, len(ids))
+	for i, id := range ids {
+		reqs[i] = svc.requests[id].req
+	}
+
+	return reqs
+}
+
+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]
+}