forked from yandex/pandora
-
Notifications
You must be signed in to change notification settings - Fork 0
/
decoder.go
105 lines (94 loc) · 2.22 KB
/
decoder.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
// Copyright (c) 2017 Yandex LLC. All rights reserved.
// Use of this source code is governed by a MPL 2.0
// license that can be found in the LICENSE file.
// Author: Vladimir Skipor <skipor@yandex-team.ru>
package uri
import (
"bytes"
"context"
"net/http"
"strings"
"sync"
"github.com/pkg/errors"
"github.com/yandex/pandora/components/phttp/ammo/simple"
)
type decoder struct {
ctx context.Context
sink chan<- *simple.Ammo
pool *sync.Pool
ammoNum int
header http.Header
}
func newDecoder(ctx context.Context, sink chan<- *simple.Ammo, pool *sync.Pool) *decoder {
return &decoder{
sink: sink,
header: http.Header{},
pool: pool,
ctx: ctx,
}
}
func (d *decoder) Decode(line []byte) error {
if len(line) == 0 {
return errors.New("empty line")
}
switch line[0] {
case '/':
return d.decodeURI(line)
case '[':
return d.decodeHeader(line)
}
return errors.New("every line should begin with '[' or '/'")
}
func (d *decoder) decodeURI(line []byte) error {
// OPTIMIZE: reuse *http.Request, http.Header. Benchmark both variants.
parts := strings.SplitN(string(line), " ", 2)
url := parts[0]
var tag string
if len(parts) > 1 {
tag = parts[1]
}
req, err := http.NewRequest("GET", string(url), nil)
if err != nil {
return errors.Wrap(err, "uri decode")
}
for k, v := range d.header {
// http.Request.Write sends Host header based on Host or URL.Host.
if k == "Host" {
req.URL.Host = v[0]
req.Host = v[0]
} else {
req.Header[k] = v
}
}
sh := d.pool.Get().(*simple.Ammo)
sh.Reset(req, tag)
select {
case d.sink <- sh:
d.ammoNum++
return nil
case <-d.ctx.Done():
return d.ctx.Err()
}
}
func (d *decoder) decodeHeader(line []byte) error {
if len(line) < 3 || line[0] != '[' || line[len(line)-1] != ']' {
return errors.New("header line should be like '[key: value]")
}
line = line[1 : len(line)-1]
colonIdx := bytes.IndexByte(line, ':')
if colonIdx < 0 {
return errors.New("missing colon")
}
key := string(bytes.TrimSpace(line[:colonIdx]))
val := string(bytes.TrimSpace(line[colonIdx+1:]))
if key == "" {
return errors.New("missing header key")
}
d.header.Set(key, val)
return nil
}
func (d *decoder) ResetHeader() {
for k := range d.header {
delete(d.header, k)
}
}