GitVita
New repositoryImport repositoryNew issue
Your repositories Your dashboard Sign in

cloudflare/pint

Public

mirrored from https://github.com/cloudflare/pintAvailable

Watch0Fork0Star0
CodeCommitsIssuesPull requestsActionsInsightsSecurity
v0.83.0

Branches

Tags

  • No tags available.
0Branches0Tags
Go to file
Add file
Code

Clone

HTTPS

Download ZIP

pint/cmd/pint

cmd/pint/main_test.go

408lines · modecode

RawDownload
Latest commit unavailable.
1package main
2
3import (
4 "bytes"
5 "context"
6 "crypto/rand"
7 "crypto/rsa"
8 "crypto/x509"
9 "crypto/x509/pkix"
10 "encoding/pem"
11 "errors"
12 "fmt"
13 "io"
14 "math/big"
15 "net"
16 "net/http"
17 "os"
18 "path"
19 "regexp"
20 "slices"
21 "strconv"
22 "strings"
23 "sync"
24 "testing"
25 "time"
26
27 "github.com/itchyny/json2yaml"
28 "github.com/rogpeppe/go-internal/testscript"
29)
30
31func TestMain(m *testing.M) {
32 testscript.Main(m, map[string]func(){
33 "pint": main,
34 })
35}
36
37func TestScripts(t *testing.T) {
38 testscript.Run(t, testscript.Params{
39 Dir: "tests",
40 UpdateScripts: os.Getenv("UPDATE_SNAPSHOTS") == "1",
41 Cmds: map[string]func(ts *testscript.TestScript, neg bool, args []string){
42 "http": httpServer,
43 "cert": tlsCert,
44 },
45 Setup: func(env *testscript.Env) error {
46 env.Values["mocks"] = &httpMocks{resps: map[string][]httpMock{}}
47 return nil
48 },
49 })
50}
51
52func httpServer(ts *testscript.TestScript, _ bool, args []string) {
53 mocks := ts.Value("mocks").(*httpMocks)
54
55 if len(args) == 0 {
56 ts.Fatalf("! http command requires arguments")
57 }
58 cmd := args[0]
59
60 switch cmd {
61 case "response":
62 if len(args) < 5 {
63 ts.Fatalf("! http response command requires '$NAME $PATH $CODE $BODY' args, got [%s]", strings.Join(args, " "))
64 }
65 name := args[1]
66 path := regexp.MustCompile(args[2])
67 code, err := strconv.Atoi(args[3])
68 ts.Check(err)
69 body := strings.Join(args[4:], " ")
70 mocks.add(name, httpMock{pattern: path, handler: func(w http.ResponseWriter, r *http.Request) {
71 snapshotRequest(ts, r, name, path)
72 w.WriteHeader(code)
73 _, err = w.Write([]byte(body))
74 ts.Check(err)
75 }})
76 case "method":
77 if len(args) < 6 {
78 ts.Fatalf("! http response command requires '$NAME $METHOD $PATH $CODE $BODY' args, got [%s]", strings.Join(args, " "))
79 }
80 name := args[1]
81 meth := args[2]
82 path := regexp.MustCompile(args[3])
83 code, err := strconv.Atoi(args[4])
84 ts.Check(err)
85 body := strings.Join(args[5:], " ")
86 mocks.add(name, httpMock{pattern: path, method: meth, handler: func(w http.ResponseWriter, r *http.Request) {
87 snapshotRequest(ts, r, name, path)
88 w.WriteHeader(code)
89 _, err := w.Write([]byte(body))
90 ts.Check(err)
91 }})
92 // http auth-response name /200 user password 200 OK
93 case "auth-response":
94 if len(args) < 7 {
95 ts.Fatalf("! http response command requires '$NAME $PATH $USER $PASS $CODE $BODY' args, got [%s]", strings.Join(args, " "))
96 }
97 name := args[1]
98 path := regexp.MustCompile(args[2])
99 user := args[3]
100 pass := args[4]
101 code, err := strconv.Atoi(args[5])
102 ts.Check(err)
103 body := strings.Join(args[6:], " ")
104 mocks.add(name, httpMock{pattern: path, handler: func(w http.ResponseWriter, r *http.Request) {
105 username, password, ok := r.BasicAuth()
106 if ok && username == user && password == pass {
107 snapshotRequest(ts, r, name, path)
108 w.WriteHeader(code)
109 _, err := w.Write([]byte(body))
110 ts.Check(err)
111 return
112 }
113 w.WriteHeader(http.StatusUnauthorized)
114 }})
115 // http response name /200 200 OK
116 case "slow-response":
117 if len(args) < 6 {
118 ts.Fatalf("! http response command requires '$NAME $PATH $DELAY $CODE $BODY' args, got [%s]", strings.Join(args, " "))
119 }
120 name := args[1]
121 path := regexp.MustCompile(args[2])
122 delay, err := time.ParseDuration(args[3])
123 ts.Check(err)
124 code, err := strconv.Atoi(args[4])
125 ts.Check(err)
126 body := strings.Join(args[5:], " ")
127 mocks.add(name, httpMock{pattern: path, handler: func(w http.ResponseWriter, r *http.Request) {
128 snapshotRequest(ts, r, name, path)
129 time.Sleep(delay)
130 w.WriteHeader(code)
131 _, err := w.Write([]byte(body))
132 ts.Check(err)
133 }})
134 // http redirect name /foo/src /dst
135 case "redirect":
136 if len(args) != 4 {
137 ts.Fatalf("! http redirect command requires '$NAME $SRCPATH $DSTPATH' args, got [%s]", strings.Join(args, " "))
138 }
139 name := args[1]
140 srcpath := regexp.MustCompile(args[2])
141 dstpath := args[3]
142 mocks.add(name, httpMock{pattern: srcpath, handler: func(w http.ResponseWriter, r *http.Request) {
143 snapshotRequest(ts, r, name, srcpath)
144 w.Header().Set("Location", dstpath)
145 w.WriteHeader(http.StatusFound)
146 }})
147 // http start name 127.0.0.1:7088 [cert.pem cert.key]
148 case "start":
149 if len(args) < 3 {
150 ts.Fatalf("! http start command requires '$NAME $LISTEN [$TLS_CERT $TLS_KEY]' args, got [%s]", strings.Join(args, " "))
151 }
152 name := args[1]
153 listen := args[2]
154 var isTLS bool
155 var tlsCert, tlsKey string
156 if len(args) == 5 {
157 isTLS = true
158 tlsCert = args[3]
159 tlsKey = args[4]
160 }
161
162 var mtx sync.Mutex
163 mux := http.NewServeMux()
164 mux.Handle("/", http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
165 var done bool
166 for n, mockList := range mocks.responses() {
167 if n == name {
168 for _, mock := range mockList {
169 if mock.method != "" && mock.method != r.Method {
170 continue
171 }
172 if !mock.pattern.MatchString(r.URL.Path) {
173 continue
174 }
175 mtx.Lock()
176 mock.handler(w, r)
177 done = true
178 mtx.Unlock()
179 break
180 }
181 break
182 }
183 }
184 if !done {
185 w.WriteHeader(http.StatusNotFound)
186 }
187 }))
188
189 listener, err := net.Listen("tcp", listen)
190 ts.Check(err)
191 server := &http.Server{Addr: listen, Handler: mux}
192 go func() {
193 var serveErr error
194 if isTLS {
195 serveErr = server.ServeTLS(listener, tlsCert, tlsKey)
196 } else {
197 serveErr = server.Serve(listener)
198 }
199 if serveErr != nil && !errors.Is(serveErr, http.ErrServerClosed) {
200 fmt.Printf("http server failed to start: %s\n", serveErr)
201 ts.Fatalf("http server failed to start: %s", serveErr)
202 }
203 }()
204
205 ts.Defer(func() {
206 ts.Logf("http server %s shutting down", name)
207 ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
208 defer cancel()
209 _ = server.Shutdown(ctx)
210 })
211 default:
212 ts.Fatalf("! unknown http command: %v", args)
213 }
214}
215
216type httpMock struct {
217 pattern *regexp.Regexp
218 handler func(http.ResponseWriter, *http.Request)
219 method string
220}
221
222type httpMocks struct {
223 resps map[string][]httpMock
224 mtx sync.Mutex
225}
226
227func (m *httpMocks) add(name string, mock httpMock) {
228 m.mtx.Lock()
229 defer m.mtx.Unlock()
230 if _, ok := m.resps[name]; !ok {
231 m.resps[name] = []httpMock{}
232 }
233 m.resps[name] = append(m.resps[name], mock)
234}
235
236func (m *httpMocks) responses() map[string][]httpMock {
237 m.mtx.Lock()
238 defer m.mtx.Unlock()
239 return m.resps
240}
241
242func tlsCert(ts *testscript.TestScript, _ bool, args []string) {
243 if len(args) < 2 {
244 ts.Fatalf("! cert command requires '$DIRNAME $NAME' args, got [%s]", strings.Join(args, " "))
245 }
246 dirname := args[0]
247 name := args[1]
248
249 ts.Logf("test-script cert command: %s", strings.Join(args, " "))
250
251 ca := &x509.Certificate{
252 SerialNumber: big.NewInt(2019),
253 Subject: pkix.Name{
254 Organization: []string{"Company, INC."},
255 Country: []string{"US"},
256 Province: []string{""},
257 Locality: []string{"San Francisco"},
258 StreetAddress: []string{""},
259 PostalCode: []string{""},
260 },
261 NotBefore: time.Now(),
262 NotAfter: time.Now().AddDate(10, 0, 0),
263 IsCA: true,
264 ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
265 KeyUsage: x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
266 BasicConstraintsValid: true,
267 }
268
269 caPrivKey, err := rsa.GenerateKey(rand.Reader, 4096)
270 if err != nil {
271 ts.Fatalf("failed to generate CA private key: %s", err)
272 }
273
274 caBytes, err := x509.CreateCertificate(rand.Reader, ca, ca, &caPrivKey.PublicKey, caPrivKey)
275 if err != nil {
276 ts.Fatalf("failed to generate CA cert: %s", err)
277 }
278
279 writeCert(ts, dirname, name+"-ca.pem", &pem.Block{
280 Type: "CERTIFICATE",
281 Bytes: caBytes,
282 })
283 writeCert(ts, dirname, name+"-ca.key", &pem.Block{
284 Type: "RSA PRIVATE KEY",
285 Bytes: x509.MarshalPKCS1PrivateKey(caPrivKey),
286 })
287
288 cert := &x509.Certificate{
289 SerialNumber: big.NewInt(1658),
290 Subject: pkix.Name{
291 Organization: []string{""},
292 Country: []string{""},
293 Province: []string{""},
294 Locality: []string{""},
295 StreetAddress: []string{""},
296 PostalCode: []string{""},
297 },
298 DNSNames: []string{name},
299 IPAddresses: []net.IP{net.IPv4(127, 0, 0, 1), net.IPv6loopback},
300 NotBefore: time.Now(),
301 NotAfter: time.Now().AddDate(0, 0, 1),
302 SubjectKeyId: []byte{1, 2, 3, 4, 6},
303 ExtKeyUsage: []x509.ExtKeyUsage{x509.ExtKeyUsageClientAuth, x509.ExtKeyUsageServerAuth},
304 KeyUsage: x509.KeyUsageDigitalSignature,
305 }
306
307 certPrivKey, err := rsa.GenerateKey(rand.Reader, 4096)
308 if err != nil {
309 ts.Fatalf("failed to generate cert private key: %s", err)
310 }
311
312 certBytes, err := x509.CreateCertificate(rand.Reader, cert, ca, &certPrivKey.PublicKey, caPrivKey)
313 if err != nil {
314 ts.Fatalf("failed to generate cert: %s", err)
315 }
316
317 writeCert(ts, dirname, name+".pem", &pem.Block{
318 Type: "CERTIFICATE",
319 Bytes: certBytes,
320 })
321 writeCert(ts, dirname, name+".key", &pem.Block{
322 Type: "RSA PRIVATE KEY",
323 Bytes: x509.MarshalPKCS1PrivateKey(certPrivKey),
324 })
325}
326
327func writeCert(ts *testscript.TestScript, dirname, filename string, block *pem.Block) {
328 fullpath := path.Join(dirname, filename)
329
330 f, err := os.Create(fullpath)
331 if err != nil {
332 ts.Fatalf("failed to write %s: %s", fullpath, err)
333 }
334
335 if err = pem.Encode(f, block); err != nil {
336 ts.Fatalf("failed to encode %s: %s", fullpath, err)
337 }
338
339 if err = f.Close(); err != nil {
340 ts.Fatalf("failed to close %s: %s", fullpath, err)
341 }
342
343 ts.Logf("Wrote PEM file to %s", filename)
344}
345
346func snapshotRequest(ts *testscript.TestScript, r *http.Request, name string, path *regexp.Regexp) {
347 payload, err := io.ReadAll(r.Body)
348 ts.Check(err)
349 r.Body.Close()
350
351 var buf strings.Builder
352 buf.WriteString(r.Method)
353 buf.WriteRune(' ')
354 buf.WriteString(path.String())
355 buf.WriteRune('\n')
356
357 hKeys := make([]string, 0, len(r.Header))
358 for k := range r.Header {
359 if k == "Content-Length" || k == "User-Agent" {
360 continue
361 }
362 hKeys = append(hKeys, k)
363 }
364 slices.Sort(hKeys)
365 for _, k := range hKeys {
366 buf.WriteString(" ")
367 buf.WriteString(k)
368 buf.WriteRune(':')
369 buf.WriteRune(' ')
370 buf.WriteString(strings.Join(r.Header.Values(k), ";"))
371 buf.WriteRune('\n')
372 }
373
374 if len(payload) > 0 {
375 payload = sanitizePayload(payload)
376 buf.WriteString("--- BODY ---\n")
377 buf.Write(payload)
378 if !bytes.HasSuffix(payload, []byte("\n")) {
379 buf.WriteRune('\n')
380 }
381 buf.WriteString("--- END ---\n")
382 }
383 buf.WriteRune('\n')
384
385 f, err := os.OpenFile(ts.MkAbs(name+".got"), os.O_APPEND|os.O_WRONLY|os.O_CREATE, 0o644)
386 ts.Check(err)
387 _, err = f.WriteString(buf.String())
388 ts.Check(err)
389 ts.Check(f.Close())
390}
391
392func sanitizePayload(payload []byte) []byte {
393 bbDurationRe := regexp.MustCompile(`\{"value":([0-9]+),"title":"Checks duration","type":"DURATION"\}`)
394 payload = bbDurationRe.ReplaceAll(payload, []byte(`{"value":0,"title":"Checks duration","type":"DURATION"}`))
395
396 ghDurationRe := regexp.MustCompile(`\| Checks duration \| ([0-9]+[a-z]+) \|`)
397 payload = ghDurationRe.ReplaceAll(payload, []byte(`| Checks duration | 0 |`))
398
399 commitIDRe := regexp.MustCompile(`"commit_id":"([0-9a-zA-Z]{40})"`)
400 payload = commitIDRe.ReplaceAll(payload, []byte(`"commit_id":"<COMMIT ID>"`))
401
402 var output bytes.Buffer
403 if err := json2yaml.Convert(&output, bytes.NewReader(payload)); err != nil {
404 return payload
405 }
406
407 return output.Bytes()
408}
409