/
httptestfileservers.go
211 lines (184 loc) · 4.98 KB
/
httptestfileservers.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
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
package testutil
import (
"context"
"net"
"net/http"
"net/http/httptest"
"os"
"path"
"path/filepath"
"strconv"
"strings"
"sync"
"testing"
)
// HttpTestFileServer returns a http server that serves files from the given directory with some latency
func HttpTestFileServer(t *testing.T, dir string) (*httptest.Server, error) {
// start server with data to send
fileSystem := &SlowFileOpener{Dir: dir}
handler := http.FileServer(fileSystem)
svr := httptest.NewServer(handler)
return svr, nil
}
type HttpTestServer struct {
*httptest.Server
lk sync.Mutex
working bool
}
func (s *HttpTestServer) SetWorking(working bool) {
s.lk.Lock()
defer s.lk.Unlock()
s.working = working
}
// HttpTestUnstartedFileServer returns a http server that serves files from the given directory
func HttpTestUnstartedFileServer(t *testing.T, dir string) *HttpTestServer {
var tsrv *HttpTestServer
handler := http.FileServer(http.Dir(dir))
svr := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
tsrv.lk.Lock()
working := tsrv.working
tsrv.lk.Unlock()
if !working {
w.WriteHeader(401)
} else {
handler.ServeHTTP(w, r)
}
}))
tsrv = &HttpTestServer{Server: svr, working: true}
return tsrv
}
type unblockInfo struct {
ch chan struct{}
closeOnce sync.Once
}
// BlockingHttpTestServer returns an http server that blocks for a given file until the client unblocks the serving of the file.
type BlockingHttpTestServer struct {
URL string
svc *httptest.Server
mu sync.Mutex
unblock map[string]*unblockInfo
}
func NewBlockingHttpTestServer(t *testing.T, dir string) *BlockingHttpTestServer {
b := &BlockingHttpTestServer{
unblock: make(map[string]*unblockInfo),
}
svc := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// wait till serving the file is unblocked
name := path.Clean(strings.TrimPrefix(r.URL.Path, "/"))
b.mu.Lock()
ubi := b.unblock[name]
b.mu.Unlock()
<-ubi.ch
// serve the file
upath := r.URL.Path
if !strings.HasPrefix(upath, "/") {
upath = "/" + upath
r.URL.Path = upath
}
fp := path.Clean(r.URL.Path)
absPath := filepath.Join(dir, fp)
http.ServeFile(w, r, absPath)
}))
b.svc = svc
return b
}
func (b *BlockingHttpTestServer) AddFile(name string) {
b.mu.Lock()
b.unblock[name] = &unblockInfo{ch: make(chan struct{})}
b.mu.Unlock()
}
func (b *BlockingHttpTestServer) UnblockFile(name string) {
b.mu.Lock()
ubi := b.unblock[name]
b.mu.Unlock()
ubi.closeOnce.Do(func() {
close(ubi.ch)
})
}
func (b *BlockingHttpTestServer) Start() {
b.svc.Start()
b.URL = b.svc.URL
}
func (b *BlockingHttpTestServer) Close() {
b.mu.Lock()
defer b.mu.Unlock()
for _, ubi := range b.unblock {
ub := ubi
ub.closeOnce.Do(func() {
close(ubi.ch)
})
}
b.svc.CloseClientConnections()
b.svc.Close()
}
type contextKey struct {
key string
}
var ConnContextKey = &contextKey{"http-conn"}
func SaveConnInContext(ctx context.Context, c net.Conn) context.Context {
return context.WithValue(ctx, ConnContextKey, c)
}
func GetConn(r *http.Request) net.Conn {
return r.Context().Value(ConnContextKey).(net.Conn)
}
// HttpTestDisconnectingServer returns a test http server that serves files from the given directory but disconnects after sending `afterEvery` bytes
// starting at the start offset mentioned in the Range request.
func HttpTestDisconnectingServer(t *testing.T, dir string, afterEvery int64) *httptest.Server {
svr := httptest.NewUnstartedServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
// process the start offset
offset := r.Header.Get("Range")
finalOffset := strings.TrimSuffix(strings.TrimPrefix(offset, "bytes="), "-")
start, _ := strconv.ParseInt(finalOffset, 10, 64)
// only send `afterEvery` bytes and then disconnect
end := start + afterEvery
// open the file to serve
upath := r.URL.Path
if !strings.HasPrefix(upath, "/") {
upath = "/" + upath
r.URL.Path = upath
}
fp := path.Clean(r.URL.Path)
absPath := filepath.Join(dir, fp)
f, err := os.Open(absPath)
if err != nil {
t.Logf("failed to open file to serve: %s", err)
w.WriteHeader(500)
return
}
defer f.Close()
// prevent buffer overflow
fi, err := f.Stat()
if err != nil {
t.Logf("failed to stat file: %s", err)
w.WriteHeader(500)
return
}
if end > fi.Size() {
end = fi.Size()
}
// read (end-start) bytes from the file starting at the given offset and write them to the response
bz := make([]byte, end-start)
n, err := f.ReadAt(bz, start)
if err != nil {
t.Logf("failed to read file: %s", err)
w.WriteHeader(500)
return
}
if int64(n) != (end - start) {
w.WriteHeader(500)
return
}
w.WriteHeader(200)
_, err = w.Write(bz)
if err != nil {
t.Logf("failed to write file: %s", err)
w.WriteHeader(500)
return
}
// close the connection so client sees an error while reading the response
c := GetConn(r)
c.Close() //nolint:errcheck
}))
svr.Config.ConnContext = SaveConnInContext
return svr
}