-
Notifications
You must be signed in to change notification settings - Fork 1
/
handler_transport.go
81 lines (74 loc) · 1.85 KB
/
handler_transport.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
package apis
import (
"context"
"io"
"net/http"
"sync/atomic"
)
// HandlerTransport is a http.RoundTripper implemented using an http.Handler bypassing the
// the need to go through the http/tcp stack when your http client wants to call a handler.
type HandlerTransport struct {
http.Handler
}
func (h HandlerTransport) RoundTrip(request *http.Request) (*http.Response, error) {
reader, writer := io.Pipe()
resp := &response{
channel: make(chan *http.Response, 1),
header: http.Header{},
reader: reader,
writer: writer,
}
go func() {
h.ServeHTTP(resp, request)
resp.WriteHeader(200)
_ = writer.Close()
}()
select {
case <-request.Context().Done():
return nil, context.Canceled
case response := <-resp.channel:
response.Request = request
return response, nil
}
}
type response struct {
sendCounter int32
channel chan *http.Response
header http.Header
statusCode int
reader *io.PipeReader
writer *io.PipeWriter
}
func (t *response) Header() http.Header {
return t.header
}
func (t *response) WriteHeader(statusCode int) {
// The header can only be sent once...
if atomic.CompareAndSwapInt32(&t.sendCounter, 0, 1) {
// copy the header just in case the handler keeps modifying it async..
header := http.Header{}
for k, v := range t.header {
header[k] = v
}
t.channel <- &http.Response{
Status: http.StatusText(t.statusCode),
StatusCode: statusCode,
Proto: "http",
ProtoMajor: 1,
ProtoMinor: 1,
Header: header,
Body: t.reader,
ContentLength: -1,
TransferEncoding: nil,
Close: true,
Uncompressed: false,
Trailer: nil,
TLS: nil,
}
close(t.channel)
}
}
func (t *response) Write(bytes []byte) (int, error) {
t.WriteHeader(http.StatusOK)
return t.writer.Write(bytes)
}