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

     1  package decoders
     2  
     3  import (
     4  	"bufio"
     5  	"context"
     6  	"fmt"
     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/util"
    16  	"github.com/yandex/pandora/core"
    17  )
    18  
    19  func newURIDecoder(file io.ReadSeeker, cfg config.Config, decodedConfigHeaders http.Header) *uriDecoder {
    20  	return &uriDecoder{
    21  		protoDecoder: protoDecoder{
    22  			file:                 file,
    23  			config:               cfg,
    24  			decodedConfigHeaders: decodedConfigHeaders,
    25  		},
    26  		scanner: bufio.NewScanner(file),
    27  		Header:  http.Header{},
    28  		pool:    &sync.Pool{New: func() any { return &ammo.Ammo{} }},
    29  	}
    30  }
    31  
    32  type uriDecoder struct {
    33  	protoDecoder
    34  	scanner *bufio.Scanner
    35  	Header  http.Header
    36  	line    uint
    37  	pool    *sync.Pool
    38  }
    39  
    40  func (d *uriDecoder) readLine(data string, commonHeader http.Header) (DecodedAmmo, error) {
    41  	data = strings.TrimSpace(data)
    42  	if len(data) == 0 {
    43  		return nil, nil // skip empty line
    44  	}
    45  	if data[0] == '[' {
    46  		key, val, err := util.DecodeHeader(data)
    47  		if err != nil {
    48  			err = fmt.Errorf("decoding header error: %w", err)
    49  			return nil, err
    50  		}
    51  		commonHeader.Set(key, val)
    52  		return nil, nil
    53  	}
    54  
    55  	var rawURL string
    56  	rawURL, tag, _ := strings.Cut(data, " ")
    57  	_, err := url.Parse(rawURL)
    58  	if err != nil {
    59  		return nil, err
    60  	}
    61  	header := commonHeader.Clone()
    62  	for k, vv := range d.decodedConfigHeaders {
    63  		for _, v := range vv {
    64  			header.Set(k, v)
    65  		}
    66  	}
    67  	a := d.pool.Get().(*ammo.Ammo)
    68  	if err := a.Setup("GET", rawURL, nil, header, tag); err != nil {
    69  		return nil, err
    70  	}
    71  	return a, nil
    72  }
    73  
    74  func (d *uriDecoder) Release(a core.Ammo) {
    75  	if am, ok := a.(*ammo.Ammo); ok {
    76  		am.Reset()
    77  		d.pool.Put(*am)
    78  	}
    79  }
    80  
    81  func (d *uriDecoder) LoadAmmo(ctx context.Context) ([]DecodedAmmo, error) {
    82  	return d.protoDecoder.LoadAmmo(ctx, d.Scan)
    83  }
    84  
    85  func (d *uriDecoder) Scan(ctx context.Context) (DecodedAmmo, error) {
    86  	if d.config.Limit != 0 && d.ammoNum >= d.config.Limit {
    87  		return nil, ErrAmmoLimit
    88  	}
    89  	for ; ; d.line++ {
    90  		if ctx.Err() != nil {
    91  			return nil, ctx.Err()
    92  		}
    93  		if !d.scanner.Scan() {
    94  			if d.scanner.Err() == nil { // assume as io.EOF; FIXME: check possible nil error with other reason
    95  				d.line = 0
    96  				d.passNum++
    97  				if d.config.Passes != 0 && d.passNum >= d.config.Passes {
    98  					return nil, ErrPassLimit
    99  				}
   100  				if d.ammoNum == 0 {
   101  					return nil, ErrNoAmmo
   102  				}
   103  				d.Header = http.Header{}
   104  				_, err := d.file.Seek(0, io.SeekStart)
   105  				if err != nil {
   106  					return nil, err
   107  				}
   108  				d.scanner = bufio.NewScanner(d.file)
   109  				continue
   110  			}
   111  			return nil, d.scanner.Err()
   112  		}
   113  		data := d.scanner.Text()
   114  		a, err := d.readLine(data, d.Header)
   115  		if err != nil {
   116  			return nil, fmt.Errorf("decode at line %d `%s` error: %w", d.line+1, data, err)
   117  		}
   118  		if a != nil {
   119  			d.ammoNum++
   120  			return a, nil
   121  		}
   122  	}
   123  }