porter.go

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

     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 ports.
    18  //
    19  // Reads env vars; for every var of the form SERVE_PORT_X, where X is a valid
    20  // port number, porter starts an HTTP server which serves the env var's value
    21  // in response to any query.
    22  
    23  package porter
    24  
    25  import (
    26  	"encoding/json"
    27  	"fmt"
    28  	"log"
    29  	"net/http"
    30  	"os"
    31  	"strings"
    32  
    33  	"github.com/ishidawataru/sctp"
    34  	"github.com/spf13/cobra"
    35  )
    36  
    37  const tcpPrefix = "SERVE_PORT_"
    38  const sctpPrefix = "SERVE_SCTP_PORT_"
    39  const tlsPrefix = "SERVE_TLS_PORT_"
    40  
    41  // CmdPorter is used by agnhost Cobra.
    42  var CmdPorter = &cobra.Command{
    43  	Use:   "porter",
    44  	Short: "Serves requested data on ports specified in ENV variables",
    45  	Long: `Serves requested data on ports specified in environment variables of the form SERVE_{PORT,TLS_PORT,SCTP_PORT}_[NNNN]. 
    46  	
    47  eg:
    48  * SERVE_PORT_9001 - serve TCP connections on port 9001
    49  * SERVE_TLS_PORT_9002 - serve TLS-encrypted TCP connections on port 9002
    50  * SERVE_SCTP_PORT_9003 - serve SCTP connections on port 9003
    51  
    52  The included "localhost.crt" is a PEM-encoded TLS cert with SAN IPs "127.0.0.1" and "[::1]", expiring in January 2084, generated from "src/crypto/tls".
    53  
    54  To use a different cert/key, mount them into the pod and set the "CERT_FILE" and "KEY_FILE" environment variables to the desired paths.`,
    55  	Args: cobra.MaximumNArgs(0),
    56  	Run:  main,
    57  }
    58  
    59  // JSONResponse enables --json-response flag
    60  var JSONResponse bool
    61  
    62  type jsonResponse struct {
    63  	Method string
    64  	Body   string
    65  }
    66  
    67  func init() {
    68  	CmdPorter.Flags().BoolVar(&JSONResponse, "json-response", false, "Responds to requests with a json response that includes the default value with the http.Request Method")
    69  }
    70  
    71  func main(cmd *cobra.Command, args []string) {
    72  	for _, vk := range os.Environ() {
    73  		// Put everything before the first = sign in parts[0], and
    74  		// everything else in parts[1] (even if there are multiple =
    75  		// characters).
    76  		parts := strings.SplitN(vk, "=", 2)
    77  		key := parts[0]
    78  		value := parts[1]
    79  
    80  		switch {
    81  		case strings.HasPrefix(key, tcpPrefix):
    82  			port := strings.TrimPrefix(key, tcpPrefix)
    83  			go servePort(port, value)
    84  		case strings.HasPrefix(key, sctpPrefix):
    85  			port := strings.TrimPrefix(key, sctpPrefix)
    86  			go serveSCTPPort(port, value)
    87  		case strings.HasPrefix(key, tlsPrefix):
    88  			port := strings.TrimPrefix(key, tlsPrefix)
    89  			go serveTLSPort(port, value)
    90  		}
    91  	}
    92  
    93  	select {}
    94  }
    95  
    96  func servePort(port, value string) {
    97  
    98  	s := &http.Server{
    99  		Addr: "0.0.0.0:" + port,
   100  		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   101  			body := value
   102  			if JSONResponse {
   103  				j, err := json.Marshal(&jsonResponse{
   104  					Method: r.Method,
   105  					Body:   value})
   106  				if err != nil {
   107  					http.Error(w, fmt.Sprintf("Internal Server Error: %v", err), 500)
   108  					return
   109  				}
   110  				w.Header().Set("Content-Type", "application/json")
   111  				body = string(j)
   112  			}
   113  			fmt.Fprint(w, body)
   114  		}),
   115  	}
   116  	log.Printf("server on port %q failed: %v", port, s.ListenAndServe())
   117  }
   118  
   119  func serveTLSPort(port, value string) {
   120  	s := &http.Server{
   121  		Addr: "0.0.0.0:" + port,
   122  		Handler: http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
   123  			body := value
   124  			if JSONResponse {
   125  				j, err := json.Marshal(&jsonResponse{
   126  					Method: r.Method,
   127  					Body:   value})
   128  				if err != nil {
   129  					http.Error(w, fmt.Sprintf("Internal Server Error: %v", err), 500)
   130  					return
   131  				}
   132  				w.Header().Set("Content-Type", "application/json")
   133  				body = string(j)
   134  			}
   135  			fmt.Fprint(w, body)
   136  		}),
   137  	}
   138  	certFile := os.Getenv("CERT_FILE")
   139  	if len(certFile) == 0 {
   140  		certFile = "localhost.crt"
   141  	}
   142  	keyFile := os.Getenv("KEY_FILE")
   143  	if len(keyFile) == 0 {
   144  		keyFile = "localhost.key"
   145  	}
   146  	log.Printf("tls server on port %q with certFile=%q, keyFile=%q failed: %v", port, certFile, keyFile, s.ListenAndServeTLS(certFile, keyFile))
   147  }
   148  
   149  func serveSCTPPort(port, value string) {
   150  	serverAddress, err := sctp.ResolveSCTPAddr("sctp", "0.0.0.0:"+port)
   151  	if err != nil {
   152  		log.Fatal("Sctp: failed to resolve address. error:", err)
   153  	}
   154  
   155  	listener, err := sctp.ListenSCTP("sctp", serverAddress)
   156  	if err != nil {
   157  		log.Fatal("Failed to listen SCTP. error:", err)
   158  	}
   159  	log.Printf("Started SCTP server")
   160  
   161  	defer listener.Close()
   162  	defer func() {
   163  		log.Printf("SCTP server exited")
   164  	}()
   165  
   166  	for {
   167  		conn, err := listener.AcceptSCTP()
   168  		if err != nil {
   169  			log.Fatal("Failed to accept SCTP. error:", err)
   170  		}
   171  		go func(conn *sctp.SCTPConn) {
   172  			defer conn.Close()
   173  			log.Println("Sending response")
   174  			_, err = conn.Write([]byte(value))
   175  			if err != nil {
   176  				log.Println("Failed to send response", err)
   177  				return
   178  			}
   179  			log.Println("Response sent")
   180  		}(conn)
   181  	}
   182  }
   183
View as plain text