trufflehog/pkg/roundtripper/roundtripper.go

package roundtripper

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"github.com/go-logr/logr"

	"github.com/trufflesecurity/trufflehog/v3/pkg/common"
	"github.com/trufflesecurity/trufflehog/v3/pkg/context"
)

type RoundTripper struct {
	logger   logr.Logger
	original http.RoundTripper
	*BasicAuthRoundTripper
	*RetryableRoundtripper
	*LoggingRoundtripper
}

type BasicAuthRoundTripper struct {
	username, password string
}

type RetryableRoundtripper struct {
	enabled                  bool
	maxRetries               uint
	shouldRetryError         bool
	shouldRetryErrorDuration time.Duration
	shouldRetry5XX           bool
	shouldRetry5XXDuration   time.Duration
	shouldRetry401           bool
	shouldRetry401Duration   time.Duration
	default429RetryDuration  time.Duration
}

type LoggingRoundtripper struct{}

// NewRoundTripper creates a new RoundTripper instance tailored for the application's specific needs.
// This custom RoundTripper provides a centralized place to manage outbound HTTP requests.
// By allowing configuration through functional options, it provides a clear, extensible, and maintainable
// way to adjust the behavior of HTTP calls. This ensures consistent logging, error handling,
// and other behaviors across all HTTP requests made in the application, reducing potential points of
// failure and simplifying debugging.
func NewRoundTripper(original http.RoundTripper, opts ...func(*RoundTripper)) *RoundTripper {
	r := &RoundTripper{
		logger:                context.Background().Logger().WithValues("component", "basic_auth_roundtripper"),
		original:              original,
		RetryableRoundtripper: &RetryableRoundtripper{enabled: false},
	}

	if original == nil {
		r.original = common.NewCustomTransport(nil)
	}

	for _, opt := range opts {
		opt(r)
	}

	return r
}

func (r *RoundTripper) RoundTrip(req *http.Request) (*http.Response, error) {
	retries := 0
	for {
		logger := r.logger.WithValues(
			"method", req.Method,
			"url", req.URL,
		)

		if r.BasicAuthRoundTripper != nil {
			req.SetBasicAuth(r.BasicAuthRoundTripper.username, r.BasicAuthRoundTripper.password)
		}

		authMethod := determineAuth(req)
		if r.LoggingRoundtripper != nil {
			logger.V(5).Info(fmt.Sprintf("sending request with %s", authMethod))
		}

		response, err := r.original.RoundTrip(req)
		if err != nil {
			logger.V(5).Info("got error while sending request",
				"error", err,
			)
		}

		if r.LoggingRoundtripper != nil && response != nil {
			logger.V(5).Info("got response",
				"status_code", response.StatusCode,
			)
		}

		reason, shouldRetry, duration := r.RetryableRoundtripper.shouldRetryRequest(response, err)
		if shouldRetry {
			if retries >= int(r.RetryableRoundtripper.maxRetries) {
				r.logger.V(2).Info("max retries reached",
					"host", req.Host,
					"reason", reason,
					"retries", retries,
				)
				return response, err
			}
			r.logger.V(2).Info("retrying request",
				"host", req.Host,
				"reason", reason,
				"retries", retries,
			)

			retries++
			time.Sleep(duration)
			continue
		}

		return response, err
	}
}

// WithInsecureTLS will disable TLS verification.
func WithInsecureTLS() func(*RoundTripper) {
	return func(r *RoundTripper) {
		r.original = common.NewCustomTransport(&http.Transport{
			TLSClientConfig: &tls.Config{InsecureSkipVerify: true},
		})
	}
}

func WithLogger(log logr.Logger) func(*RoundTripper) {
	return func(r *RoundTripper) {
		r.logger = log
	}
}

// Logging

func WithLogging() func(*RoundTripper) {
	return func(r *RoundTripper) {
		r.LoggingRoundtripper = &LoggingRoundtripper{}
	}
}

// Basic Auth

func WithBasicAuth(username, password string) func(*RoundTripper) {
	return func(r *RoundTripper) {
		r.BasicAuthRoundTripper = &BasicAuthRoundTripper{
			username: username,
			password: password,
		}
	}
}

// Retryable

func (r *RetryableRoundtripper) shouldRetryRequest(response *http.Response, err error) (reason string, shouldRetry bool, after time.Duration) {
	if !r.enabled {
		return "", false, 0
	}

	if err != nil {
		return err.Error(), r.shouldRetryError, r.shouldRetryErrorDuration
	}

	if response.StatusCode == http.StatusTooManyRequests || response.StatusCode == http.StatusServiceUnavailable {
		retryAfter := getRetryAfter(response)
		if retryAfter == 0 {
			retryAfter = r.default429RetryDuration
		}
		return strconv.Itoa(response.StatusCode), true, retryAfter
	}

	code := response.StatusCode
	switch {
	case code == 401:
		return strconv.Itoa(response.StatusCode), r.shouldRetry401, r.shouldRetry401Duration
	case code >= 500:
		return strconv.Itoa(response.StatusCode), r.shouldRetry5XX, r.shouldRetry5XXDuration
	default:
		return "", false, 0
	}
}

func getRetryAfter(response *http.Response) time.Duration {
	if s, ok := response.Header["Retry-After"]; ok {
		if sleep, err := strconv.ParseInt(s[0], 10, 64); err == nil {
			return time.Second * time.Duration(sleep)
		}
	}
	return 0
}

func WithRetryable(opts ...func(*RetryableRoundtripper)) func(*RoundTripper) {
	return func(r *RoundTripper) {
		rt := &RetryableRoundtripper{
			enabled:                  true,
			maxRetries:               3,
			shouldRetryError:         true,
			shouldRetryErrorDuration: time.Second * 5,
			shouldRetry5XX:           true,
			shouldRetry5XXDuration:   time.Second * 5,
			shouldRetry401:           false,
			shouldRetry401Duration:   time.Second * 5,
			default429RetryDuration:  time.Second * 30,
		}

		for _, opt := range opts {
			opt(rt)
		}

		r.RetryableRoundtripper = rt
	}
}

func WithShouldRetryError(should bool) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.shouldRetryError = should
	}
}

func WithShouldRetryErrorDuration(duration time.Duration) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.shouldRetryErrorDuration = duration
	}
}

func WithShouldRetry5XX(should bool) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.shouldRetry5XX = should
	}
}

func WithShouldRetry5XXDuration(duration time.Duration) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.shouldRetry5XXDuration = duration
	}
}

func WithShouldRetry401Duration(duration time.Duration) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.shouldRetry401 = true
		r.shouldRetry401Duration = duration
	}
}

func WithDefault429RetryDuration(duration time.Duration) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.default429RetryDuration = duration
	}
}

func WithMaxRetries(max uint) func(*RetryableRoundtripper) {
	return func(r *RetryableRoundtripper) {
		r.maxRetries = max
	}
}

// determineAuth will let us know if the Authorization header is set,
// and if it is, if it contains the prefix "Basic" or "Bearer"
func determineAuth(req *http.Request) string {
	switch authHeader := req.Header.Get("Authorization"); true {
	case authHeader == "":
		return "no auth"
	case strings.HasPrefix(authHeader, "Basic"):
		return "basic auth"
	case strings.HasPrefix(authHeader, "Bearer"):
		return "bearer token"
	default:
		return "unknown auth"
	}
}