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

     1  package decoders
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"errors"
     7  	"io"
     8  	"net/http"
     9  	"net/url"
    10  	"strings"
    11  	"sync"
    12  
    13  	"github.com/yandex/pandora/components/providers/http/config"
    14  	"github.com/yandex/pandora/components/providers/http/decoders/ammo"
    15  	"github.com/yandex/pandora/components/providers/http/decoders/uripost"
    16  	"github.com/yandex/pandora/components/providers/http/util"
    17  	"github.com/yandex/pandora/core"
    18  	"golang.org/x/xerrors"
    19  )
    20  
    21  func newURIPostDecoder(file io.ReadSeeker, cfg config.Config, decodedConfigHeaders http.Header) *uripostDecoder {
    22  	return &uripostDecoder{
    23  		protoDecoder: protoDecoder{
    24  			file:                 file,
    25  			config:               cfg,
    26  			decodedConfigHeaders: decodedConfigHeaders,
    27  		},
    28  		reader: bufio.NewReader(file),
    29  		header: http.Header{},
    30  		pool:   &sync.Pool{New: func() any { return &ammo.Ammo{} }},
    31  	}
    32  }
    33  
    34  type uripostDecoder struct {
    35  	protoDecoder
    36  	reader *bufio.Reader
    37  	header http.Header
    38  	line   uint
    39  	pool   *sync.Pool
    40  }
    41  
    42  func (d *uripostDecoder) Release(a core.Ammo) {
    43  	if am, ok := a.(*ammo.Ammo); ok {
    44  		am.Reset()
    45  		d.pool.Put(am)
    46  	}
    47  }
    48  
    49  func (d *uripostDecoder) LoadAmmo(ctx context.Context) ([]DecodedAmmo, error) {
    50  	return d.protoDecoder.LoadAmmo(ctx, d.Scan)
    51  }
    52  
    53  func (d *uripostDecoder) Scan(ctx context.Context) (DecodedAmmo, error) {
    54  	if d.config.Limit != 0 && d.ammoNum >= d.config.Limit {
    55  		return nil, ErrAmmoLimit
    56  	}
    57  	for i := 0; i < 2; i++ {
    58  		for {
    59  			if ctx.Err() != nil {
    60  				return nil, ctx.Err()
    61  			}
    62  
    63  			a, err := d.readBlock(d.reader, d.header)
    64  			if err == io.EOF {
    65  				break
    66  			}
    67  			if err != nil {
    68  				return nil, err
    69  			}
    70  			if a != nil {
    71  				d.ammoNum++
    72  				return a, nil
    73  			}
    74  			// here only if read header
    75  		}
    76  
    77  		// seek file
    78  		d.passNum++
    79  		if d.config.Passes != 0 && d.passNum >= d.config.Passes {
    80  			return nil, ErrPassLimit
    81  		}
    82  		if d.ammoNum == 0 {
    83  			return nil, ErrNoAmmo
    84  		}
    85  		d.header = make(http.Header)
    86  		_, err := d.file.Seek(0, io.SeekStart)
    87  		if err != nil {
    88  			return nil, err
    89  		}
    90  		d.reader.Reset(d.file)
    91  	}
    92  
    93  	return nil, errors.New("unexpected behavior")
    94  }
    95  
    96  // readBlock read one header at time and set to commonHeader or read full request
    97  func (d *uripostDecoder) readBlock(reader *bufio.Reader, commonHeader http.Header) (*ammo.Ammo, error) {
    98  	data, err := reader.ReadString('\n')
    99  	if err != nil {
   100  		return nil, err
   101  	}
   102  	data = strings.TrimSpace(data)
   103  	if len(data) == 0 {
   104  		return nil, nil // skip empty lines
   105  	}
   106  	if data[0] == '[' {
   107  		key, val, err := util.DecodeHeader(data)
   108  		if err != nil {
   109  			return nil, err
   110  		}
   111  		commonHeader.Set(key, val)
   112  		return nil, nil
   113  	}
   114  
   115  	bodySize, uri, tag, err := uripost.DecodeURI(data)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  	_, err = url.Parse(uri)
   120  	if err != nil {
   121  		return nil, err
   122  	}
   123  
   124  	buff := make([]byte, bodySize)
   125  	if bodySize != 0 {
   126  		if n, err := io.ReadFull(reader, buff); err != nil {
   127  			err = xerrors.Errorf("failed to read ammo with err: %w, at position: %v; tried to read: %v; have read: %v", err, filePosition(d.file), bodySize, n)
   128  			return nil, err
   129  		}
   130  	}
   131  
   132  	header := commonHeader.Clone()
   133  	for k, vv := range d.decodedConfigHeaders {
   134  		for _, v := range vv {
   135  			header.Set(k, v)
   136  		}
   137  	}
   138  	a := d.pool.Get().(*ammo.Ammo)
   139  	err = a.Setup("POST", uri, buff, header, tag)
   140  	return a, err
   141  }