1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16 package main
17
18 import (
19 "bytes"
20 "context"
21 "encoding/json"
22 "flag"
23 "fmt"
24 "io"
25 "net"
26 "net/http"
27 "os"
28 "path/filepath"
29 "time"
30
31 "github.com/containerd/cgroups"
32 )
33
34 func main() {
35 port := flag.Int("port", 8080, "listen port")
36 dataDir := flag.String("data-dir", "/var/run/data", "data dir is the dir to write temp file, only used in io test")
37
38 flag.Parse()
39
40 s := newServer(*dataDir)
41
42 ctx, cancel := context.WithCancel(context.Background())
43 defer cancel()
44 go s.childProcessTimeServer.Start(ctx)
45 err := s.setupUDPServer()
46 if err != nil {
47 fmt.Println("failed to serve udp server", err)
48 os.Exit(1)
49 }
50
51 addr := fmt.Sprintf("0.0.0.0:%d", *port)
52 if err := http.ListenAndServe(addr, s.mux); err != nil {
53 fmt.Println("failed to serve http server", err)
54 os.Exit(1)
55 }
56 }
57
58 type server struct {
59 mux *http.ServeMux
60 dataDir string
61 childProcessTimeServer childProcessTimeServer
62
63
64 recvBuf []byte
65 }
66
67 func newServer(dataDir string) *server {
68 s := &server{
69 mux: http.NewServeMux(),
70 dataDir: dataDir,
71 recvBuf: make([]byte, 5),
72 childProcessTimeServer: &defaultChildProcessTimeServer{},
73 }
74 s.mux.HandleFunc("/ping", pong)
75 s.mux.HandleFunc("/time", s.time)
76 s.mux.HandleFunc("/child-process-time", s.childProcessTime)
77 s.mux.HandleFunc("/io", s.ioTest)
78 s.mux.HandleFunc("/mistake", s.mistakeTest)
79 s.mux.HandleFunc("/network/send", s.networkSendTest)
80 s.mux.HandleFunc("/network/recv", s.networkRecvTest)
81 s.mux.HandleFunc("/network/ping", s.networkPingTest)
82 s.mux.HandleFunc("/dns", s.dnsTest)
83 s.mux.HandleFunc("/stress", s.stressCondition)
84 s.mux.HandleFunc("/http", s.httpEcho)
85 s.mux.HandleFunc("/setup_https", s.SetupHTTPSServer)
86 return s
87 }
88
89 func pong(w http.ResponseWriter, _ *http.Request) {
90 w.Write([]byte("pong"))
91 }
92
93 func (s *server) setupUDPServer() error {
94 pc, err := net.ListenPacket("udp", "0.0.0.0:1070")
95 if err != nil {
96 return err
97 }
98
99 go func() {
100 for {
101 _, _, err := pc.ReadFrom(s.recvBuf)
102 fmt.Println("receive buf " + string(s.recvBuf))
103 if err != nil {
104 return
105 }
106 }
107 }()
108
109 return nil
110 }
111
112
113 func (s *server) time(w http.ResponseWriter, _ *http.Request) {
114 w.Write([]byte(time.Now().Format(time.RFC3339Nano)))
115 }
116
117
118 func (s *server) childProcessTime(w http.ResponseWriter, _ *http.Request) {
119 now, err := s.childProcessTimeServer.Time()
120 if err != nil {
121 panic(err)
122 }
123 w.Write([]byte(now.Format(time.RFC3339Nano)))
124 }
125
126
127 func (s *server) mistakeTest(w http.ResponseWriter, _ *http.Request) {
128 path := filepath.Join(s.dataDir, "e2e-test")
129 origData := []byte("hello world!!!!!!!!!!!!")
130
131 err := os.WriteFile(path, origData, 0644)
132 if err != nil {
133 w.Write([]byte(fmt.Sprintf("failed to write file %v", err)))
134 return
135 }
136 gotData, err := os.ReadFile(path)
137 if err != nil {
138 w.Write([]byte(err.Error()))
139 return
140 }
141 result := bytes.Equal(origData, gotData)
142 if result {
143 w.Write([]byte("false"))
144 return
145 }
146 for i := 0; i < 10; i++ {
147 tmp, err := os.ReadFile(path)
148 if err != nil {
149 w.Write([]byte(err.Error()))
150 }
151 if !bytes.Equal(tmp, gotData) {
152 w.Write([]byte("true"))
153 return
154 }
155 }
156 w.Write([]byte("err"))
157 }
158
159
160 func (s *server) ioTest(w http.ResponseWriter, _ *http.Request) {
161 t1 := time.Now()
162 f, err := os.CreateTemp(s.dataDir, "e2e-test")
163 if err != nil {
164 w.Write([]byte(fmt.Sprintf("failed to create temp file %v", err)))
165 return
166 }
167 if _, err := f.Write([]byte("hello world")); err != nil {
168 w.Write([]byte(fmt.Sprintf("failed to write file %v", err)))
169 return
170 }
171 t2 := time.Now()
172 w.Write([]byte(t2.Sub(t1).String()))
173 }
174
175
176 func (s *server) dnsTest(w http.ResponseWriter, r *http.Request) {
177
178 url, ok := r.URL.Query()["url"]
179
180 if !ok || len(url[0]) < 1 {
181 http.Error(w, "failed", http.StatusBadRequest)
182 return
183 }
184
185 ips, err := net.LookupIP(url[0])
186 if err != nil {
187 http.Error(w, "failed", http.StatusBadRequest)
188 return
189 }
190
191 if len(ips) == 0 {
192 http.Error(w, "failed", http.StatusBadRequest)
193 return
194 }
195
196 w.Write([]byte(ips[0].String()))
197 }
198
199 type networkSendTestBody struct {
200 TargetIP string `json:"targetIP"`
201 }
202
203
204 func (s *server) networkPingTest(w http.ResponseWriter, r *http.Request) {
205 var body networkSendTestBody
206
207 err := json.NewDecoder(r.Body).Decode(&body)
208 if err != nil {
209 http.Error(w, err.Error(), http.StatusBadRequest)
210 return
211 }
212
213 c := http.Client{
214 Timeout: 2 * time.Second,
215 }
216 startTime := time.Now()
217 resp, err := c.Get(fmt.Sprintf("http://%s:8080/ping", body.TargetIP))
218 if err != nil {
219 http.Error(w, err.Error(), http.StatusBadRequest)
220 return
221 }
222 defer resp.Body.Close()
223
224 endTime := time.Now()
225 out, err := io.ReadAll(resp.Body)
226 if err != nil {
227 http.Error(w, err.Error(), http.StatusBadRequest)
228 return
229 }
230
231 if string(out) != "pong" {
232 http.Error(w, "response is not pong", http.StatusBadRequest)
233 return
234 }
235
236 w.Write([]byte(fmt.Sprintf("OK %d", endTime.UnixNano()-startTime.UnixNano())))
237 }
238
239
240 func (s *server) networkSendTest(w http.ResponseWriter, r *http.Request) {
241 var body networkSendTestBody
242
243 err := json.NewDecoder(r.Body).Decode(&body)
244 if err != nil {
245 http.Error(w, err.Error(), http.StatusBadRequest)
246 return
247 }
248
249 conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
250 IP: net.ParseIP(body.TargetIP),
251 Port: 1070,
252 })
253 if err != nil {
254 http.Error(w, err.Error(), http.StatusBadRequest)
255 return
256 }
257 defer conn.Close()
258
259 n, err := io.WriteString(conn, "ping\n")
260 if err != nil {
261 http.Error(w, err.Error(), http.StatusBadRequest)
262 return
263 }
264 if n != 5 {
265 http.Error(w, "udp send less than 5 bytes", http.StatusBadRequest)
266 return
267 }
268 w.Write([]byte("send successfully\n"))
269 }
270
271
272 func (s *server) networkRecvTest(w http.ResponseWriter, r *http.Request) {
273 w.Write(s.recvBuf)
274
275 s.recvBuf = []byte{}
276 }
277
278 func (s *server) stressCondition(w http.ResponseWriter, r *http.Request) {
279 control, err := cgroups.Load(cgroups.V1, cgroups.PidPath(1))
280 if err != nil {
281 http.Error(w, err.Error(), http.StatusInternalServerError)
282 return
283 }
284
285 stats, err := control.Stat(cgroups.IgnoreNotExist)
286 if err != nil {
287 http.Error(w, err.Error(), http.StatusInternalServerError)
288 return
289 }
290
291 response, err := json.Marshal(map[string]uint64{
292 "cpuTime": stats.CPU.Usage.Total,
293 "memoryUsage": stats.Memory.Usage.Usage - stats.Memory.Kernel.Usage - stats.Memory.Cache,
294 })
295 if err != nil {
296 http.Error(w, "fail to marshal response", http.StatusInternalServerError)
297 return
298 }
299
300 w.Write(response)
301 }
302
303 func (s *server) httpEcho(w http.ResponseWriter, r *http.Request) {
304 secrets := r.Header["Secret"]
305 if len(secrets) == 0 {
306 http.Error(w, "Forbidden", http.StatusForbidden)
307 return
308 }
309
310 for _, secret := range secrets {
311 w.Header().Add("Secret", secret)
312 }
313 defer r.Body.Close()
314 _, err := io.Copy(w, r.Body)
315 if err != nil {
316 http.Error(w, "fail to copy body between request and response", http.StatusInternalServerError)
317 return
318 }
319 }
320
321 type TLSServerKeys struct {
322 Cert []byte `json:"cert"`
323 Key []byte `json:"key"`
324 }
325
326 func (s *server) SetupHTTPSServer(w http.ResponseWriter, r *http.Request) {
327 var body TLSServerKeys
328 err := json.NewDecoder(r.Body).Decode(&body)
329 if err != nil {
330 http.Error(w, err.Error(), http.StatusBadRequest)
331 return
332 }
333 err = os.WriteFile("/tmp/server.crt", body.Cert, 0644)
334 if err != nil {
335 http.Error(w, err.Error(), http.StatusInternalServerError)
336 return
337 }
338 err = os.WriteFile("/tmp/server.key", body.Key, 0644)
339 if err != nil {
340 http.Error(w, err.Error(), http.StatusInternalServerError)
341 return
342 }
343 srv := &server{
344 mux: http.NewServeMux(),
345 }
346 srv.mux.HandleFunc("/ping", pong)
347
348 go func() {
349 panic(http.ListenAndServeTLS("0.0.0.0:8081", "/tmp/server.crt", "/tmp/server.key", srv.mux))
350 }()
351 if err != nil {
352 return
353 }
354 }
355