github.com/yandex/pandora@v0.5.32/components/providers/http/decoders/raw.go (about)

     1  package decoders
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"io"
     7  	"net/http"
     8  	"strings"
     9  	"sync"
    10  
    11  	"github.com/yandex/pandora/components/providers/http/config"
    12  	"github.com/yandex/pandora/components/providers/http/decoders/ammo"
    13  	"github.com/yandex/pandora/components/providers/http/decoders/raw"
    14  	"github.com/yandex/pandora/core"
    15  	"golang.org/x/xerrors"
    16  )
    17  
    18  func newRawDecoder(file io.ReadSeeker, cfg config.Config, decodedConfigHeaders http.Header) *rawDecoder {
    19  	return &rawDecoder{
    20  		protoDecoder: protoDecoder{
    21  			file:                 file,
    22  			config:               cfg,
    23  			decodedConfigHeaders: decodedConfigHeaders,
    24  		},
    25  		pool:   &sync.Pool{New: func() any { return &ammo.RawAmmo{} }},
    26  		reader: bufio.NewReader(file),
    27  	}
    28  }
    29  
    30  /*
    31  Parses size-prefixed HTTP ammo files. Each ammo is prefixed with a header line (delimited with \n), which consists of
    32  two fields delimited by a space: ammo size and tag. Ammo size is in bytes (integer, including special characters like CR, LF).
    33  Tag is a string. Example:
    34  
    35  77 bad
    36  GET /abra HTTP/1.0
    37  Host: xxx.tanks.example.com
    38  User-Agent: xxx (shell 1)
    39  
    40  904
    41  POST /upload/2 HTTP/1.0
    42  Content-Length: 801
    43  Host: xxxxxxxxx.dev.example.com
    44  User-Agent: xxx (shell 1)
    45  
    46  ^.^........W.j^1^.^.^.²..^^.i.^B.P..-!(.l/Y..V^.      ...L?...S'NR.^^vm...3Gg@s...d'.\^.5N.$NF^,.Z^.aTE^.
    47  ._.[..k#L^ƨ`\RE.J.<.!,.q5.F^՚iΔĬq..^6..P..тH.`..i2
    48  .".uuzs^^F2...Rh.&.U.^^..J.P@.A......x..lǝy^?.u.p{4..g...m.,..R^.^.^......].^^.^J...p.ifTF0<.s.9V.o5<..%!6ļS.ƐǢ..㱋....C^&.....^.^y...v]^YT.1.#K.ibc...^.26...   ..7.
    49  b.$...j6.٨f...W.R7.^1.3....K`%.&^..d..{{      l0..^\..^X.g.^.r.(!.^^...4.1.$\ .%.8$(.n&..^^q.,.Q..^.D^.].^.R9.kE.^.$^.I..<..B^..^.h^^C.^E.|....3o^.@..Z.^.s.$[v.
    50  527
    51  POST /upload/3 HTTP/1.0
    52  Content-Length: 424
    53  Host: xxxxxxxxx.dev.example.com
    54  User-Agent: xxx (shell 1)
    55  
    56  ^.^........QMO.0^.++^zJw.ر^$^.^Ѣ.^V.J....vM.8r&.T+...{@pk%~C.G../z顲^.7....l...-.^W"cR..... .&^?u.U^^.^.....{^.^..8.^.^.I.EĂ.p...'^.3.Tq..@R8....RAiBU..1.Bd*".7+.
    57  .Ol.j=^.3..n....wp..,Wg.y^.T..~^..
    58  */
    59  type rawDecoder struct {
    60  	protoDecoder
    61  	reader *bufio.Reader
    62  	pool   *sync.Pool
    63  }
    64  
    65  func (d *rawDecoder) LoadAmmo(ctx context.Context) ([]DecodedAmmo, error) {
    66  	return d.protoDecoder.LoadAmmo(ctx, d.Scan)
    67  }
    68  
    69  func (d *rawDecoder) Release(a core.Ammo) {
    70  	if am, ok := a.(*ammo.RawAmmo); ok {
    71  		am.Reset()
    72  		d.pool.Put(am)
    73  	}
    74  }
    75  
    76  func (d *rawDecoder) Scan(ctx context.Context) (DecodedAmmo, error) {
    77  	var data string
    78  	var err error
    79  
    80  	if d.config.Limit != 0 && d.ammoNum >= d.config.Limit {
    81  		return nil, ErrAmmoLimit
    82  	}
    83  	for {
    84  		if ctx.Err() != nil {
    85  			return nil, ctx.Err()
    86  		}
    87  
    88  		data, err = d.reader.ReadString('\n')
    89  		if err == io.EOF {
    90  			d.passNum++
    91  			if d.config.Passes != 0 && d.passNum >= d.config.Passes {
    92  				return nil, ErrPassLimit
    93  			}
    94  			if d.ammoNum == 0 {
    95  				return nil, ErrNoAmmo
    96  			}
    97  			_, err := d.file.Seek(0, io.SeekStart)
    98  			if err != nil {
    99  				return nil, err
   100  			}
   101  			d.reader.Reset(d.file)
   102  			continue
   103  		}
   104  		position := filePosition(d.file)
   105  		if err != nil {
   106  			return nil, xerrors.Errorf("reading ammo failed with err: %w, at position: %v", err, position)
   107  		}
   108  		data = strings.TrimSpace(data)
   109  		if len(data) == 0 {
   110  			continue // skip empty lines
   111  		}
   112  		d.ammoNum++
   113  		reqSize, tag, err := raw.DecodeHeader(data)
   114  		if err != nil {
   115  			return nil, xerrors.Errorf("header decoding error for ammoNum %d: %w", d.ammoNum, err)
   116  		}
   117  
   118  		a := d.pool.Get().(*ammo.RawAmmo)
   119  		if reqSize != 0 {
   120  			buff := make([]byte, reqSize)
   121  			if n, err := io.ReadFull(d.reader, buff); err != nil {
   122  				return nil, xerrors.Errorf("failed to read ammo with err: %w, at position: %v; tried to read: %v; have read: %v", err, position, reqSize, n)
   123  			}
   124  
   125  			a.Setup(buff, tag, position, d.decodedConfigHeaders)
   126  		} else {
   127  			a.Setup(nil, "", position, d.decodedConfigHeaders)
   128  		}
   129  		return a, nil
   130  	}
   131  }