main.go

Documentation: k8s.io/kubernetes/test/images/agnhost/net

     1  /*
     2  Copyright 2014 The Kubernetes Authors.
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8      http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package net
    18  
    19  import (
    20  	"bytes"
    21  	"encoding/json"
    22  	"fmt"
    23  	"io"
    24  	"log"
    25  	"net/http"
    26  	"os"
    27  	"os/signal"
    28  	"strings"
    29  	"syscall"
    30  	"time"
    31  
    32  	"github.com/spf13/cobra"
    33  
    34  	"k8s.io/kubernetes/test/images/agnhost/net/common"
    35  	"k8s.io/kubernetes/test/images/agnhost/net/nat"
    36  )
    37  
    38  type runnerMap map[string]common.Runner
    39  
    40  var (
    41  	// flags for the command line. See usage args below for
    42  	// descriptions.
    43  	flags struct {
    44  		Serve         string
    45  		Runner        string
    46  		Options       string
    47  		DelayShutdown int
    48  	}
    49  	// runners is a map from runner name to runner instance.
    50  	runners = makeRunnerMap()
    51  )
    52  
    53  type logOutput struct {
    54  	b bytes.Buffer
    55  }
    56  
    57  // CmdNet is used by agnhost Cobra.
    58  var CmdNet = &cobra.Command{
    59  	Use:   "net",
    60  	Short: "Creates webserver or runner for various networking tests",
    61  	Long: `The subcommand will run the network tester in server mode if the "--serve" flag is given, and the runners are triggered through HTTP requests.
    62  
    63  Alternatively, if the "--runner" flag is given, it will execute the given runner directly. Note that "--runner" and "--serve" flags cannot be given at the same time.
    64  
    65  Examples:
    66  
    67      agnhost net --runner nat-closewait-client --options '{"RemoteAddr":"127.0.0.1:9999"}'
    68  
    69      agnhost net --serve :8889 && curl -v -X POST localhost:8889/run/nat-closewait-server -d '{"LocalAddr":"127.0.0.1:9999"}'
    70  `,
    71  	Args: cobra.MaximumNArgs(0),
    72  	Run:  main,
    73  }
    74  
    75  func init() {
    76  	legalRunners := ""
    77  	for k := range runners {
    78  		legalRunners += " " + k
    79  	}
    80  	CmdNet.Flags().StringVar(&flags.Serve, "serve", "",
    81  		"Address and port to bind to (e.g. 127.0.0.1:8080). Setting this will "+
    82  			"run the network tester in server mode runner are triggered through "+
    83  			"HTTP requests.")
    84  	CmdNet.Flags().StringVar(&flags.Runner, "runner", "", "Runner to execute (available:"+legalRunners+")")
    85  	CmdNet.Flags().StringVar(&flags.Options, "options", "", "JSON options to the Runner")
    86  	CmdNet.Flags().IntVar(&flags.DelayShutdown, "delay-shutdown", 0, "Number of seconds to delay shutdown when receiving SIGTERM.")
    87  }
    88  
    89  func main(cmd *cobra.Command, args []string) {
    90  	if flags.Runner == "" && flags.Serve == "" {
    91  		log.Fatalf("Must set either --runner or --serve, see --help")
    92  	}
    93  
    94  	if flags.DelayShutdown > 0 {
    95  		termCh := make(chan os.Signal, 1)
    96  		signal.Notify(termCh, syscall.SIGTERM)
    97  		go func() {
    98  			<-termCh
    99  			log.Printf("Sleeping %d seconds before terminating...", flags.DelayShutdown)
   100  			time.Sleep(time.Duration(flags.DelayShutdown) * time.Second)
   101  			os.Exit(0)
   102  		}()
   103  	}
   104  
   105  	log.SetFlags(log.Flags() | log.Lshortfile)
   106  
   107  	if flags.Serve == "" {
   108  		output, err := executeRunner(flags.Runner, flags.Options)
   109  		if err == nil {
   110  			fmt.Print("output:\n\n" + output.b.String())
   111  			os.Exit(0)
   112  		} else {
   113  			log.Printf("Error: %v", err)
   114  			fmt.Print("output:\n\n" + output.b.String())
   115  			os.Exit(1)
   116  		}
   117  	} else {
   118  		http.HandleFunc("/run/", handleRunRequest)
   119  		log.Printf("Running server on %v", flags.Serve)
   120  		log.Fatal(http.ListenAndServe(flags.Serve, nil))
   121  	}
   122  }
   123  
   124  func makeRunnerMap() runnerMap {
   125  	// runner name is <pkg>-<file>-<specific>.
   126  	return runnerMap{
   127  		"nat-closewait-client": nat.NewCloseWaitClient(),
   128  		"nat-closewait-server": nat.NewCloseWaitServer(),
   129  	}
   130  }
   131  
   132  func executeRunner(name string, rawOptions string) (logOutput, error) {
   133  	runner, ok := runners[name]
   134  	if ok {
   135  		options := runner.NewOptions()
   136  		if err := json.Unmarshal([]byte(rawOptions), options); err != nil {
   137  			return logOutput{}, fmt.Errorf("Invalid options JSON: %v", err)
   138  		}
   139  
   140  		log.Printf("Options: %+v", options)
   141  
   142  		output := logOutput{}
   143  		logger := log.New(&output.b, "# ", log.Lshortfile)
   144  
   145  		return output, runner.Run(logger, options)
   146  	}
   147  
   148  	return logOutput{}, fmt.Errorf("Invalid runner: '%v', see --help", runner)
   149  }
   150  
   151  // handleRunRequest handles a request JSON to the network tester.
   152  func handleRunRequest(w http.ResponseWriter, r *http.Request) {
   153  	log.Printf("handleRunRequest %v", *r)
   154  
   155  	urlParts := strings.Split(r.URL.Path, "/")
   156  	if len(urlParts) != 3 {
   157  		http.Error(w, fmt.Sprintf("invalid request to run: %v", urlParts), 400)
   158  		return
   159  	}
   160  
   161  	runner := urlParts[2]
   162  	if r.Body == nil || r.Body == http.NoBody {
   163  		http.Error(w, "Missing request body", 400)
   164  		return
   165  	}
   166  
   167  	body, err := io.ReadAll(r.Body)
   168  	if err != nil {
   169  		http.Error(w, fmt.Sprintf("error reading body: %v", err), 400)
   170  		return
   171  	}
   172  
   173  	var output logOutput
   174  	if output, err = executeRunner(runner, string(body)); err != nil {
   175  		contents := fmt.Sprintf("Error from runner: %v\noutput:\n\n%s",
   176  			err, output.b.String())
   177  		http.Error(w, contents, 500)
   178  		return
   179  	}
   180  
   181  	fmt.Fprintf(w, "ok\noutput:\n\n"+output.b.String())
   182  }
   183
View as plain text