serve_hostname.go

Documentation: k8s.io/kubernetes/test/images/agnhost/serve-hostname

     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  // A small utility to just serve the hostname on TCP and/or UDP.
    18  
    19  package servehostname
    20  
    21  import (
    22  	"fmt"
    23  	"log"
    24  	"net"
    25  	"net/http"
    26  	"os"
    27  	"os/signal"
    28  	"syscall"
    29  	"time"
    30  
    31  	"github.com/spf13/cobra"
    32  )
    33  
    34  // CmdServeHostname is used by agnhost Cobra.
    35  var CmdServeHostname = &cobra.Command{
    36  	Use:   "serve-hostname",
    37  	Short: "Serves the hostname",
    38  	Long:  `Serves the hostname through HTTP / TCP / UDP on the given port.`,
    39  	Args:  cobra.MaximumNArgs(0),
    40  	Run:   main,
    41  }
    42  
    43  var (
    44  	doTCP   bool
    45  	doUDP   bool
    46  	doHTTP  bool
    47  	doClose bool
    48  	port    int
    49  )
    50  
    51  func init() {
    52  	CmdServeHostname.Flags().BoolVar(&doTCP, "tcp", false, "Serve raw over TCP.")
    53  	CmdServeHostname.Flags().BoolVar(&doUDP, "udp", false, "Serve raw over UDP.")
    54  	CmdServeHostname.Flags().BoolVar(&doHTTP, "http", true, "Serve HTTP.")
    55  	CmdServeHostname.Flags().BoolVar(&doClose, "close", false, "Close connection per each HTTP request.")
    56  	CmdServeHostname.Flags().IntVar(&port, "port", 9376, "Port number.")
    57  }
    58  
    59  func main(cmd *cobra.Command, args []string) {
    60  	if doHTTP && (doTCP || doUDP) {
    61  		log.Fatalf("Can't server TCP/UDP mode and HTTP mode at the same time")
    62  	}
    63  
    64  	hostname, err := os.Hostname()
    65  	if err != nil {
    66  		log.Fatalf("Error from os.Hostname(): %s", err)
    67  	}
    68  
    69  	if doTCP {
    70  		listener, err := net.Listen("tcp", fmt.Sprintf(":%d", port))
    71  		if err != nil {
    72  			log.Fatalf("Error from net.Listen(): %s", err)
    73  		}
    74  		go func() {
    75  			for {
    76  				conn, err := listener.Accept()
    77  				if err != nil {
    78  					log.Fatalf("Error from Accept(): %s", err)
    79  				}
    80  				log.Printf("TCP request from %s", conn.RemoteAddr().String())
    81  				conn.Write([]byte(hostname))
    82  				conn.Close()
    83  			}
    84  		}()
    85  	}
    86  	if doUDP {
    87  		addr, err := net.ResolveUDPAddr("udp", fmt.Sprintf(":%d", port))
    88  		if err != nil {
    89  			log.Fatalf("Error from net.ResolveUDPAddr(): %s", err)
    90  		}
    91  		sock, err := net.ListenUDP("udp", addr)
    92  		if err != nil {
    93  			log.Fatalf("Error from ListenUDP(): %s", err)
    94  		}
    95  		go func() {
    96  			var buffer [16]byte
    97  			for {
    98  				_, cliAddr, err := sock.ReadFrom(buffer[0:])
    99  				if err != nil {
   100  					log.Fatalf("Error from ReadFrom(): %s", err)
   101  				}
   102  				log.Printf("UDP request from %s", cliAddr.String())
   103  				sock.WriteTo([]byte(hostname), cliAddr)
   104  			}
   105  		}()
   106  	}
   107  	if doHTTP {
   108  		http.HandleFunc("/", func(w http.ResponseWriter, r *http.Request) {
   109  			log.Printf("HTTP request from %s", r.RemoteAddr)
   110  
   111  			if doClose {
   112  				// Add this header to force to close the connection after serving the request.
   113  				w.Header().Add("Connection", "close")
   114  			}
   115  
   116  			fmt.Fprintf(w, "%s", hostname)
   117  		})
   118  		go func() {
   119  			// Run in a closure so http.ListenAndServe doesn't block
   120  			log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", port), nil))
   121  		}()
   122  	}
   123  	log.Printf("Serving on port %d.\n", port)
   124  	signals := make(chan os.Signal, 1)
   125  	signal.Notify(signals, syscall.SIGTERM)
   126  	sig := <-signals
   127  	log.Printf("Shutting down after receiving signal: %s.\n", sig)
   128  	log.Printf("Awaiting pod deletion.\n")
   129  	time.Sleep(60 * time.Second)
   130  }
   131
View as plain text