github.com/argoproj/argo-cd/v3@v3.2.1/util/grpc/sanitizer.go (about)

     1  package grpc
     2  
     3  import (
     4  	"context"
     5  	"errors"
     6  	"regexp"
     7  	"strings"
     8  
     9  	"google.golang.org/grpc"
    10  	"google.golang.org/grpc/status"
    11  )
    12  
    13  type sanitizerKey string
    14  
    15  const (
    16  	contextKey sanitizerKey = "sanitizer"
    17  )
    18  
    19  // ErrorSanitizerUnaryServerInterceptor returns a new unary server interceptor that sanitizes error messages
    20  // and provides Sanitizer to define replacements.
    21  func ErrorSanitizerUnaryServerInterceptor() grpc.UnaryServerInterceptor {
    22  	return func(ctx context.Context, req any, _ *grpc.UnaryServerInfo, handler grpc.UnaryHandler) (resp any, err error) {
    23  		sanitizer := NewSanitizer()
    24  		resp, err = handler(ContextWithSanitizer(ctx, sanitizer), req)
    25  		if err == nil {
    26  			return resp, nil
    27  		}
    28  
    29  		if se, ok := err.(interface{ GRPCStatus() *status.Status }); ok {
    30  			return resp, status.Error(se.GRPCStatus().Code(), sanitizer.Replace(se.GRPCStatus().Message()))
    31  		}
    32  
    33  		return resp, errors.New(sanitizer.Replace(err.Error()))
    34  	}
    35  }
    36  
    37  // ContextWithSanitizer returns a new context with sanitizer set.
    38  func ContextWithSanitizer(ctx context.Context, sanitizer Sanitizer) context.Context {
    39  	return context.WithValue(ctx, contextKey, sanitizer)
    40  }
    41  
    42  // SanitizerFromContext returns sanitizer from context.
    43  func SanitizerFromContext(ctx context.Context) (Sanitizer, bool) {
    44  	res, ok := ctx.Value(contextKey).(Sanitizer)
    45  	return res, ok
    46  }
    47  
    48  // Sanitizer provides methods to define list of strings and replacements
    49  type Sanitizer interface {
    50  	Replace(s string) string
    51  	AddReplacement(val string, replacement string)
    52  	AddRegexReplacement(regex *regexp.Regexp, replacement string)
    53  }
    54  
    55  type sanitizer struct {
    56  	replacers []func(in string) string
    57  }
    58  
    59  // NewSanitizer returns a new Sanitizer instance
    60  func NewSanitizer() *sanitizer {
    61  	return &sanitizer{}
    62  }
    63  
    64  // AddReplacement adds a replacement to the Sanitizer
    65  func (s *sanitizer) AddReplacement(val string, replacement string) {
    66  	s.replacers = append(s.replacers, func(in string) string {
    67  		return strings.ReplaceAll(in, val, replacement)
    68  	})
    69  }
    70  
    71  // AddRegexReplacement adds a replacement to the sanitizer using regexp
    72  func (s *sanitizer) AddRegexReplacement(regex *regexp.Regexp, replacement string) {
    73  	s.replacers = append(s.replacers, func(in string) string {
    74  		return regex.ReplaceAllString(in, replacement)
    75  	})
    76  }
    77  
    78  // Replace replaces all occurrences of the configured values in the sanitizer with the replacements
    79  func (s *sanitizer) Replace(val string) string {
    80  	for _, replacer := range s.replacers {
    81  		val = replacer(val)
    82  	}
    83  	return val
    84  }