-
Notifications
You must be signed in to change notification settings - Fork 1
/
client.go
87 lines (73 loc) · 1.42 KB
/
client.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
package peregrine
import (
"context"
"github.com/gobwas/ws"
"io"
"net/http"
"sync/atomic"
)
type writer struct {
n int64
w io.Writer
}
func (w *writer) WriteString(s string) (int, error) {
n, err := io.WriteString(w.w, s)
w.n += int64(n)
return n, err
}
func (w *writer) Write(p []byte) (int, error) {
n, err := w.w.Write(p)
w.n += int64(n)
return n, err
}
type ClientHeader struct {
http.Header
}
func (h *ClientHeader) WriteTo(w io.Writer) (int64, error) {
wr := writer{w: w}
err := h.Write(&wr)
return wr.n, err
}
type Client struct {
dialer ws.Dialer
ctx context.Context
addr string
}
func (c *Client) Dial(header *ClientHeader) (*ClientConn, error) {
dialer := c.dialer
// preprocess request header
if header != nil {
dialerHeader, ok := dialer.Header.(*ClientHeader)
if !ok {
dialerHeader = header
} else {
for key, values := range (map[string][]string)(header.Header) {
dialerHeader.Set(key, values[0])
}
}
dialer.Header = dialerHeader
}
conn, _, _, err := dialer.Dial(c.ctx, c.addr)
if err != nil {
return nil, err
}
return &ClientConn{
state: func() *atomic.Bool {
b := atomic.Bool{}
b.Store(true)
return &b
}(),
conn: conn,
}, nil
}
func NewClient(addr string, opts ...ClientOptionFunc) *Client {
client := &Client{
dialer: ws.DefaultDialer,
ctx: context.Background(),
addr: addr,
}
for _, opt := range opts {
opt(client)
}
return client
}