main.go

Documentation: github.com/chaos-mesh/chaos-mesh/test/cmd/e2e_helper

     1  // Copyright 2020 Chaos Mesh Authors.
     2  //
     3  // Licensed under the Apache License, Version 2.0 (the "License");
     4  // you may not use this file except in compliance with the License.
     5  // You may obtain a copy of the License at
     6  //
     7  //     http://www.apache.org/licenses/LICENSE-2.0
     8  //
     9  // Unless required by applicable law or agreed to in writing, software
    10  // distributed under the License is distributed on an "AS IS" BASIS,
    11  // See the License for the specific language governing permissions and
    12  // limitations under the License.
    13  
    14  package main
    15  
    16  import (
    17  	"encoding/json"
    18  	"flag"
    19  	"fmt"
    20  	"io"
    21  	"io/ioutil"
    22  	"net"
    23  	"net/http"
    24  	"os"
    25  	"time"
    26  )
    27  
    28  func main() {
    29  	port := flag.Int("port", 8080, "listen port")
    30  	dataDir := flag.String("data-dir", "/var/run/data", "data dir is the dir to write temp file, only used in io test")
    31  
    32  	flag.Parse()
    33  
    34  	s := newServer(*dataDir)
    35  	err := s.setupUDPServer()
    36  	if err != nil {
    37  		fmt.Println("failed to serve udp server", err)
    38  		os.Exit(1)
    39  	}
    40  
    41  	addr := fmt.Sprintf("0.0.0.0:%d", *port)
    42  	if err := http.ListenAndServe(addr, s.mux); err != nil {
    43  		fmt.Println("failed to serve http server", err)
    44  		os.Exit(1)
    45  	}
    46  }
    47  
    48  type server struct {
    49  	mux     *http.ServeMux
    50  	dataDir string
    51  
    52  	// ONLY FOR TEST: a buf without lock
    53  	recvBuf []byte
    54  }
    55  
    56  func newServer(dataDir string) *server {
    57  	s := &server{
    58  		mux:     http.NewServeMux(),
    59  		dataDir: dataDir,
    60  		recvBuf: make([]byte, 5),
    61  	}
    62  	s.mux.HandleFunc("/ping", pong)
    63  	s.mux.HandleFunc("/time", s.timer)
    64  	s.mux.HandleFunc("/io", s.ioTest)
    65  	s.mux.HandleFunc("/network/send", s.networkSendTest)
    66  	s.mux.HandleFunc("/network/recv", s.networkRecvTest)
    67  	s.mux.HandleFunc("/network/ping", s.networkPingTest)
    68  	return s
    69  }
    70  
    71  func pong(w http.ResponseWriter, _ *http.Request) {
    72  	w.Write([]byte("pong"))
    73  }
    74  
    75  func (s *server) setupUDPServer() error {
    76  	pc, err := net.ListenPacket("udp", "0.0.0.0:1070")
    77  	if err != nil {
    78  		return err
    79  	}
    80  
    81  	go func() {
    82  		for {
    83  			_, _, err := pc.ReadFrom(s.recvBuf)
    84  			if err != nil {
    85  				return
    86  			}
    87  		}
    88  	}()
    89  
    90  	return nil
    91  }
    92  
    93  // a handler to print out the current time
    94  func (s *server) timer(w http.ResponseWriter, _ *http.Request) {
    95  	w.Write([]byte(time.Now().Format(time.RFC3339Nano)))
    96  }
    97  
    98  // a handler to test io chaos
    99  func (s *server) ioTest(w http.ResponseWriter, _ *http.Request) {
   100  	t1 := time.Now()
   101  	f, err := ioutil.TempFile(s.dataDir, "e2e-test")
   102  	if err != nil {
   103  		w.Write([]byte(fmt.Sprintf("failed to create temp file %v", err)))
   104  		return
   105  	}
   106  	if _, err := f.Write([]byte("hello world")); err != nil {
   107  		w.Write([]byte(fmt.Sprintf("failed to write file %v", err)))
   108  		return
   109  	}
   110  	t2 := time.Now()
   111  	w.Write([]byte(t2.Sub(t1).String()))
   112  }
   113  
   114  type networkSendTestBody struct {
   115  	TargetIP string `json:"targetIP"`
   116  }
   117  
   118  // a handler to test network chaos
   119  func (s *server) networkPingTest(w http.ResponseWriter, r *http.Request) {
   120  	var body networkSendTestBody
   121  
   122  	err := json.NewDecoder(r.Body).Decode(&body)
   123  	if err != nil {
   124  		http.Error(w, err.Error(), http.StatusBadRequest)
   125  		return
   126  	}
   127  
   128  	c := http.Client{
   129  		Timeout: 10 * time.Second,
   130  	}
   131  	startTime := time.Now()
   132  	resp, err := c.Get(fmt.Sprintf("http://%s:8080/ping", body.TargetIP))
   133  	endTime := time.Now()
   134  	out, err := ioutil.ReadAll(resp.Body)
   135  	defer resp.Body.Close()
   136  	if string(out) != "pong" {
   137  		http.Error(w, "response is not pong", http.StatusBadRequest)
   138  		return
   139  	}
   140  
   141  	w.Write([]byte(fmt.Sprintf("OK %d", endTime.UnixNano()-startTime.UnixNano())))
   142  }
   143  
   144  // a handler to test network chaos
   145  func (s *server) networkSendTest(w http.ResponseWriter, r *http.Request) {
   146  	var body networkSendTestBody
   147  
   148  	err := json.NewDecoder(r.Body).Decode(&body)
   149  	if err != nil {
   150  		http.Error(w, err.Error(), http.StatusBadRequest)
   151  		return
   152  	}
   153  
   154  	conn, err := net.DialUDP("udp", nil, &net.UDPAddr{
   155  		IP:   net.ParseIP(body.TargetIP),
   156  		Port: 1070,
   157  	})
   158  	if err != nil {
   159  		http.Error(w, err.Error(), http.StatusBadRequest)
   160  		return
   161  	}
   162  	defer conn.Close()
   163  
   164  	n, err := io.WriteString(conn, "ping\n")
   165  	if err != nil {
   166  		http.Error(w, err.Error(), http.StatusBadRequest)
   167  		return
   168  	}
   169  	if n != 5 {
   170  		http.Error(w, "udp send less than 5 bytes", http.StatusBadRequest)
   171  		return
   172  	}
   173  	w.Write([]byte("send successfully\n"))
   174  }
   175  
   176  // a handler to test network chaos
   177  func (s *server) networkRecvTest(w http.ResponseWriter, r *http.Request) {
   178  	w.Write(s.recvBuf)
   179  
   180  	for index := range s.recvBuf {
   181  		s.recvBuf[index] = 0
   182  	}
   183  }
   184