/
handler.go
134 lines (111 loc) · 2.89 KB
/
handler.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
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
package netx
import (
"bufio"
"context"
"io"
"net"
"time"
)
// A Handler manages a network connection.
//
// The ServeConn method is called by a Server when a new client connection is
// established, the method receives the connection and a context object that
// the server may use to indicate that it's shutting down.
//
// Servers recover from panics that escape the handlers and log the error and
// stack trace.
type Handler interface {
ServeConn(ctx context.Context, conn net.Conn)
}
// The HandlerFunc type allows simple functions to be used as connection
// handlers.
type HandlerFunc func(context.Context, net.Conn)
// ServeConn calls f.
func (f HandlerFunc) ServeConn(ctx context.Context, conn net.Conn) {
f(ctx, conn)
}
// CloseHandler wraps handler to ensure that the connections it receives are
// always closed after it returns.
func CloseHandler(handler Handler) Handler {
return HandlerFunc(func(ctx context.Context, conn net.Conn) {
defer conn.Close()
handler.ServeConn(ctx, conn)
})
}
var (
// Echo is the implementation of a connection handler that simply sends what
// it receives back to the client.
Echo Handler = HandlerFunc(echo)
// EchoLine is the implementation of a connection handler that reads lines
// and echos them back to the client, expecting the client not to send more
// than one line before getting it echoed back.
//
// The implementation supports cancellations and ensures that no partial
// lines are read from the connection.
//
// The maximum line length is limited to 8192 bytes.
EchoLine Handler = HandlerFunc(echoLine)
// Pass is the implementation of a connection that does nothing.
Pass Handler = HandlerFunc(pass)
)
func echo(ctx context.Context, conn net.Conn) {
ctx, cancel := context.WithCancel(ctx)
go func() {
defer cancel()
Copy(conn, conn)
}()
<-ctx.Done()
conn.Close()
}
func echoLine(ctx context.Context, conn net.Conn) {
r := bufio.NewReaderSize(conn, 8192)
for {
line, err := readLine(ctx, conn, r)
switch err {
case nil:
case io.EOF, context.Canceled:
return
default:
fatal(conn, err)
}
if _, err := conn.Write(line); err != nil {
fatal(conn, err)
}
}
}
func pass(ctx context.Context, conn net.Conn) {
// do nothing
}
func fatal(conn net.Conn, err error) {
conn.Close()
panic(err)
}
func readLine(ctx context.Context, conn net.Conn, r *bufio.Reader) ([]byte, error) {
for {
select {
default:
case <-ctx.Done():
return nil, ctx.Err()
}
conn.SetReadDeadline(time.Now().Add(1 * time.Second))
if _, err := r.Peek(1); err != nil {
if IsTimeout(err) {
continue
}
}
line, prefix, err := r.ReadLine()
switch {
case prefix:
line, err = nil, ErrLineTooLong
case err != nil:
line = nil
case r.Buffered() != 0:
line, err = nil, ErrNoPipeline
default:
if line = line[:len(line)+1]; line[len(line)-1] == '\r' {
line = line[:len(line)+1]
}
}
return line, err
}
}