portforwardtester.go

Documentation: k8s.io/kubernetes/test/images/agnhost/port-forward-tester

     1  /*
     2  Copyright 2015 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  // A tiny binary for testing port forwarding. The following environment variables
    18  // control the binary's logic:
    19  //
    20  // BIND_PORT - the TCP port to use for the listener
    21  // EXPECTED_CLIENT_DATA - data that we expect to receive from the client; may be "".
    22  // CHUNKS - how many chunks of data we should send to the client
    23  // CHUNK_SIZE - how large each chunk should be
    24  // CHUNK_INTERVAL - the delay in between sending each chunk
    25  //
    26  // Log messages are written to stdout at various stages of the binary's execution.
    27  // Test code can retrieve this container's log and validate that the expected
    28  // behavior is taking place.
    29  
    30  package portforwardtester
    31  
    32  import (
    33  	"fmt"
    34  	"net"
    35  	"os"
    36  	"strconv"
    37  	"strings"
    38  	"time"
    39  
    40  	"github.com/spf13/cobra"
    41  )
    42  
    43  func getEnvInt(name string) int {
    44  	s := os.Getenv(name)
    45  	value, err := strconv.Atoi(s)
    46  	if err != nil {
    47  		fmt.Printf("Error parsing %s %q: %v\n", name, s, err)
    48  		os.Exit(1)
    49  	}
    50  	return value
    51  }
    52  
    53  // taken from net/http/server.go:
    54  //
    55  // rstAvoidanceDelay is the amount of time we sleep after closing the
    56  // write side of a TCP connection before closing the entire socket.
    57  // By sleeping, we increase the chances that the client sees our FIN
    58  // and processes its final data before they process the subsequent RST
    59  // from closing a connection with known unread data.
    60  // This RST seems to occur mostly on BSD systems. (And Windows?)
    61  // This timeout is somewhat arbitrary (~latency around the planet).
    62  const rstAvoidanceDelay = 500 * time.Millisecond
    63  
    64  // CmdPortForwardTester is used by agnhost Cobra.
    65  var CmdPortForwardTester = &cobra.Command{
    66  	Use:   "port-forward-tester",
    67  	Short: "Creates a TCP server that sends chunks of data",
    68  	Long: `Listens for TCP connections on a given address and port, optionally checks the data received,
    69  and sends a configurable number of data chunks, with a configurable interval between chunks.
    70  
    71  The subcommand is using the following environment variables:
    72  
    73  - BIND_ADDRESS (optional): The address on which it will start listening for TCP connections (default value: localhost)
    74  - BIND_PORT: The port on which it will start listening for TCP connections.
    75  - EXPECTED_CLIENT_DATA (optional): If set, it will check that the request sends the same exact data.
    76  - CHUNKS: How many chunks of data to write in the response.
    77  - CHUNK_SIZE: The expected size of each written chunk of data. If it does not match the actual size of the written data, it will exit with the exit code 4.
    78  - CHUNK_INTERVAL: The amount of time to wait in between chunks.`,
    79  	Args: cobra.MaximumNArgs(0),
    80  	Run:  main,
    81  }
    82  
    83  func main(cmd *cobra.Command, args []string) {
    84  	bindAddress := os.Getenv("BIND_ADDRESS")
    85  	if bindAddress == "" {
    86  		bindAddress = "localhost"
    87  	}
    88  	bindPort := os.Getenv("BIND_PORT")
    89  	addr, err := net.ResolveTCPAddr("tcp", net.JoinHostPort(bindAddress, bindPort))
    90  	if err != nil {
    91  		fmt.Printf("Error resolving: %v\n", err)
    92  		os.Exit(1)
    93  	}
    94  	listener, err := net.ListenTCP("tcp", addr)
    95  	if err != nil {
    96  		fmt.Printf("Error listening: %v\n", err)
    97  		os.Exit(1)
    98  	}
    99  
   100  	conn, err := listener.AcceptTCP()
   101  	if err != nil {
   102  		fmt.Printf("Error accepting connection: %v\n", err)
   103  		os.Exit(1)
   104  	}
   105  
   106  	fmt.Println("Accepted client connection")
   107  
   108  	expectedClientData := os.Getenv("EXPECTED_CLIENT_DATA")
   109  	if len(expectedClientData) > 0 {
   110  		buf := make([]byte, len(expectedClientData))
   111  		read, err := conn.Read(buf)
   112  		if read != len(expectedClientData) {
   113  			fmt.Printf("Expected to read %d bytes from client, but got %d instead. err=%v\n", len(expectedClientData), read, err)
   114  			os.Exit(2)
   115  		}
   116  		if expectedClientData != string(buf) {
   117  			fmt.Printf("Expect to read %q, but got %q. err=%v\n", expectedClientData, string(buf), err)
   118  			os.Exit(3)
   119  		}
   120  		if err != nil {
   121  			fmt.Printf("Read err: %v\n", err)
   122  		}
   123  		fmt.Println("Received expected client data")
   124  	}
   125  
   126  	chunks := getEnvInt("CHUNKS")
   127  	chunkSize := getEnvInt("CHUNK_SIZE")
   128  	chunkInterval := getEnvInt("CHUNK_INTERVAL")
   129  
   130  	stringData := strings.Repeat("x", chunkSize)
   131  	data := []byte(stringData)
   132  
   133  	for i := 0; i < chunks; i++ {
   134  		written, err := conn.Write(data)
   135  		if written != chunkSize {
   136  			fmt.Printf("Expected to write %d bytes from client, but wrote %d instead. err=%v\n", chunkSize, written, err)
   137  			os.Exit(4)
   138  		}
   139  		if err != nil {
   140  			fmt.Printf("Write err: %v\n", err)
   141  		}
   142  		if i+1 < chunks {
   143  			time.Sleep(time.Duration(chunkInterval) * time.Millisecond)
   144  		}
   145  	}
   146  
   147  	fmt.Println("Shutting down connection")
   148  
   149  	// set linger timeout to flush buffers. This is the official way according to the go api docs. But
   150  	// there are controversial discussions whether this value has any impact on most platforms
   151  	// (compare https://codereview.appspot.com/95320043).
   152  	conn.SetLinger(-1)
   153  
   154  	// Flush the connection cleanly, following https://blog.netherlabs.nl/articles/2009/01/18/the-ultimate-so_linger-page-or-why-is-my-tcp-not-reliable:
   155  	// 1. close write half of connection which sends a FIN packet
   156  	// 2. give client some time to receive the FIN
   157  	// 3. close the complete connection
   158  	conn.CloseWrite()
   159  	time.Sleep(rstAvoidanceDelay)
   160  	conn.Close()
   161  
   162  	fmt.Println("Done")
   163  }
   164
View as plain text