github.com/yandex/pandora@v0.5.32/components/providers/scenario/http/postprocessor/var_header.go (about)

     1  package postprocessor
     2  
     3  import (
     4  	"fmt"
     5  	"io"
     6  	"net/http"
     7  	"strconv"
     8  	"strings"
     9  
    10  	"github.com/yandex/pandora/lib/str"
    11  )
    12  
    13  type VarHeaderPostprocessor struct {
    14  	Mapping map[string]string
    15  }
    16  
    17  func (p *VarHeaderPostprocessor) ReturnedParams() []string {
    18  	result := make([]string, len(p.Mapping))
    19  	for k := range p.Mapping {
    20  		result = append(result, k)
    21  	}
    22  	return result
    23  }
    24  
    25  func (p *VarHeaderPostprocessor) Process(resp *http.Response, _ io.Reader) (map[string]any, error) {
    26  	if len(p.Mapping) == 0 {
    27  		return nil, nil
    28  	}
    29  	result := make(map[string]any, len(p.Mapping))
    30  	for k, v := range p.Mapping {
    31  		headerVal, modifier, err := p.parseValue(v)
    32  		if err != nil {
    33  			return nil, fmt.Errorf("failed to parse value %s: %w", v, err)
    34  		}
    35  		val := resp.Header.Get(headerVal)
    36  		if val == "" {
    37  			continue
    38  		}
    39  		result[k] = modifier(val)
    40  	}
    41  	return result, nil
    42  }
    43  
    44  func (p *VarHeaderPostprocessor) parseValue(v string) (value string, modifier func(in string) string, err error) {
    45  	vals := strings.Split(v, "|")
    46  	if len(vals) == 1 {
    47  		return vals[0], func(in string) string { return in }, nil
    48  	}
    49  
    50  	value = vals[0]
    51  	modifier = func(in string) string { return in }
    52  
    53  	for _, modStr := range vals[1:] {
    54  		mod, err := p.parseModifier(modStr)
    55  		if err != nil {
    56  			return "", nil, fmt.Errorf("failed to parse modifier %s: %w", modStr, err)
    57  		}
    58  		previousModifier := modifier
    59  		modifier = func(in string) string {
    60  			return mod(previousModifier(in))
    61  		}
    62  	}
    63  
    64  	return value, modifier, nil
    65  }
    66  
    67  func (p *VarHeaderPostprocessor) parseModifier(s string) (func(in string) string, error) {
    68  	name, args, err := str.ParseStringFunc(s)
    69  	if err != nil {
    70  		return nil, fmt.Errorf("failed to parse function name %s: %w", s, err)
    71  	}
    72  	switch name {
    73  	case "lower":
    74  		return func(in string) string { return strings.ToLower(in) }, nil
    75  	case "upper":
    76  		return func(in string) string { return strings.ToUpper(in) }, nil
    77  	case "substr":
    78  		return p.substr(args)
    79  	case "replace":
    80  		if len(args) != 2 {
    81  			return nil, fmt.Errorf("replace modifier requires 2 arguments")
    82  		}
    83  		return func(in string) string {
    84  			return strings.ReplaceAll(in, args[0], args[1])
    85  		}, nil
    86  	}
    87  	return nil, fmt.Errorf("unknown modifier %s", name)
    88  }
    89  
    90  func (p *VarHeaderPostprocessor) substr(args []string) (func(in string) string, error) {
    91  	if len(args) == 0 || len(args) > 2 {
    92  		return nil, fmt.Errorf("substr modifier requires one or two arguments")
    93  	}
    94  	start, err := strconv.Atoi(args[0])
    95  	if err != nil {
    96  		return nil, fmt.Errorf("substr modifier requires integer as first argument, got %s", args[0])
    97  	}
    98  	end := 0
    99  	if len(args) == 2 {
   100  		end, err = strconv.Atoi(args[1])
   101  		if err != nil {
   102  			return nil, fmt.Errorf("substr modifier requires integer as second argument, got %s", args[1])
   103  		}
   104  	}
   105  	return func(in string) string {
   106  		l := len(in)
   107  		if start < 0 {
   108  			start = l + start
   109  		}
   110  		if end <= 0 {
   111  			end = l + end
   112  		}
   113  		if end > l {
   114  			end = l
   115  		}
   116  		if start > end {
   117  			start, end = end, start
   118  		}
   119  		return in[start:end]
   120  	}, nil
   121  }