echo-server/main.go

package main

import (
	"crypto/tls"
	"encoding/pem"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"time"

	"golang.org/x/crypto/pkcs12"
)

func envOrDefault(key, def string) string {
	if v := os.Getenv(key); v != "" {
		return v
	}
	return def
}

func main() {
	port := flag.String("p", envOrDefault("PORT", "41912"), "port to run on")
	crt := flag.String("crt", envOrDefault("CRT", ""), "path to crt")
	key := flag.String("key", envOrDefault("KEY", ""), "path to key")
	tcp := flag.Bool("tcp", false, "accept TCP")
	udp := flag.Bool("udp", false, "accept udp")
	fail := flag.Bool("fail", false, "fail connections forever")
	block := flag.Int("block", 0, "seconds to block connections")
	status := flag.Int("status", http.StatusOK, "status to always return")
	body := flag.String("body", envOrDefault("BODY", ""), "body to always return")
	flag.Parse()

	// start server
	if *udp {
		startUDP(*port, *block)
	} else if *tcp {
		startTCP(*port, *fail, *block)
	} else {
		startHTTP(*port, *fail, *block, *crt, *key, *status, *body)
	}

	// catch stop
	stop := make(chan os.Signal)
	signal.Notify(stop, os.Interrupt)
	<-stop
}

func startUDP(port string, block int) {
	pc, err := net.ListenPacket("udp", ":"+port)
	if err != nil {
		panic(err)
	}
	log.Println("echoing", port)
	go func() {
		defer pc.Close()
		buff := make([]byte, 1024)
		for {
			n, addr, err := pc.ReadFrom(buff)
			if err != nil {
				panic(err)
			}
			log.Printf("%s: %s", addr, buff[:n])
			_, err = pc.WriteTo(buff[:n], addr)
			if err != nil {
				panic(err)
			}
		}
	}()
}

func startTCP(port string, fail bool, block int) {
	l, err := net.Listen("tcp", ":"+port)
	if err != nil {
		panic(err)
	}
	go func() {
		for {
			c, err := l.Accept()
			if err != nil {
				panic(err)
			}

			go func(c net.Conn) {
				defer c.Close()
				if block > 0 {
					time.Sleep(time.Second * time.Duration(block))
				}
				fmt.Fprintln(c, "Hello")
			}(c)
		}
		defer l.Close()
	}()
}

func startHTTP(port string, fail bool, block int, crt, key string, status int, body string) {
	srv := &http.Server{
		Addr:    ":" + port,
		Handler: makeHTTPHandler(fail, block, status, body),
	}
	go func() {
		if crt == "" || key == "" {
			if err := srv.ListenAndServe(); err != nil {
				panic(err)
			}
		} else {
			if _, err := os.Stat(key); err == nil {
				if err := srv.ListenAndServeTLS(crt, key); err != nil {
					panic(err)
				}
			} else {
				b, err := ioutil.ReadFile(crt)
				if err != nil {
					panic(err)
				}
				blocks, err := pkcs12.ToPEM(b, key)
				if err != nil {
					panic(err)
				}
				pd := []byte{}
				for _, block := range blocks {
					pd = append(pd, pem.EncodeToMemory(block)...)
				}
				cert, err := tls.X509KeyPair(pd, pd)
				if err != nil {
					panic(err)
				}
				srv.TLSConfig = &tls.Config{
					Certificates: []tls.Certificate{cert},
				}
				if err := srv.ListenAndServeTLS("", ""); err != nil {
					panic(err)
				}
			}
		}
	}()
	log.Println("Listening on", srv.Addr)
}

func makeHTTPHandler(fail bool, block int, status int, body string) http.HandlerFunc {
	if block > 0 {
		return blockH(block, fail, status, body)
	} else if fail {
		return http.HandlerFunc(failH)
	}
	return echoH(status, body)
}

func failH(w http.ResponseWriter, r *http.Request) {
	w.WriteHeader(http.StatusInternalServerError)
}

func blockH(block int, fail bool, status int, body string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		select {
		case <-time.After(time.Second * time.Duration(block)):
		case <-r.Context().Done():
			log.Println(r.Context().Err())
			fmt.Fprintln(w, r.Context().Err())
			w.WriteHeader(http.StatusInternalServerError)
			return
		}
		if fail {
			w.WriteHeader(http.StatusInternalServerError)
		} else {
			echoH(status, body)(w, r)
		}
	}
}

func echoH(status int, body string) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		log.Println("METHOD:", r.Method)
		log.Println("URL:", r.URL.String())
		log.Println("HEADER:", r.Header)
		b, err := ioutil.ReadAll(r.Body)
		if err != nil {
			return
		}
		s := string(b[:])
		log.Println("BODY:", s)

		w.WriteHeader(status)
		if len(body) == 0 {
			fmt.Fprintf(w, "===\nMETHOD: %q\nURL: %q\nHEADER: %q\nBODY: %q\n===\n",
				r.Method,
				r.URL.String(),
				fmt.Sprintf("%v", r.Header),
				s,
			)
		} else {
			fmt.Fprintln(w, body)
		}
	}
}