From 19b1c31e602b39ec3f82705abcd426ad5932adfd Mon Sep 17 00:00:00 2001
From: David Anderson <danderson@tailscale.com>
Date: Thu, 3 Feb 2022 14:17:32 -0800
Subject: [PATCH] cmd/hello: refactor to use tsweb.Server.

Signed-off-by: David Anderson <danderson@tailscale.com>
---
 cmd/hello/hello.go | 65 +++++++++++++++++++++-------------------------
 1 file changed, 29 insertions(+), 36 deletions(-)

diff --git a/cmd/hello/hello.go b/cmd/hello/hello.go
index 97b98be32..4ae67388a 100644
--- a/cmd/hello/hello.go
+++ b/cmd/hello/hello.go
@@ -18,10 +18,10 @@
 	"net/http"
 	"os"
 	"strings"
-	"time"
 
 	"tailscale.com/client/tailscale"
 	"tailscale.com/client/tailscale/apitype"
+	"tailscale.com/tsweb"
 )
 
 var (
@@ -62,44 +62,37 @@ func main() {
 	http.HandleFunc("/", root)
 	log.Printf("Starting hello server.")
 
-	errc := make(chan error, 1)
-	if *httpAddr != "" {
-		log.Printf("running HTTP server on %s", *httpAddr)
-		go func() {
-			errc <- http.ListenAndServe(*httpAddr, nil)
-		}()
+	mainAddr := *httpsAddr
+	if mainAddr == "" {
+		mainAddr = *httpAddr
 	}
-	if *httpsAddr != "" {
-		log.Printf("running HTTPS server on %s", *httpsAddr)
-		go func() {
-			hs := &http.Server{
-				Addr: *httpsAddr,
-				TLSConfig: &tls.Config{
-					GetCertificate: func(hi *tls.ClientHelloInfo) (*tls.Certificate, error) {
-						switch hi.ServerName {
-						case "hello.ts.net":
-							return tailscale.GetCertificate(hi)
-						case "hello.ipn.dev":
-							c, err := tls.LoadX509KeyPair(
-								"/etc/hello/hello.ipn.dev.crt",
-								"/etc/hello/hello.ipn.dev.key",
-							)
-							if err != nil {
-								return nil, err
-							}
-							return &c, nil
-						}
-						return nil, errors.New("invalid SNI name")
-					},
-				},
-				IdleTimeout:       30 * time.Second,
-				ReadHeaderTimeout: 20 * time.Second,
-				MaxHeaderBytes:    10 << 10,
+	httpCfg := tsweb.ServerConfig{
+		Name:    "hello",
+		Addr:    mainAddr,
+		Handler: http.DefaultServeMux,
+	}
+	server := tsweb.NewServer(httpCfg)
+	if server.HTTPS != nil {
+		server.HTTPS.TLSConfig.GetCertificate = func(hi *tls.ClientHelloInfo) (*tls.Certificate, error) {
+			switch hi.ServerName {
+			case "hello.ts.net":
+				return tailscale.GetCertificate(hi)
+			case "hello.ipn.dev":
+				c, err := tls.LoadX509KeyPair(
+					"/etc/hello/hello.ipn.dev.crt",
+					"/etc/hello/hello.ipn.dev.key",
+				)
+				if err != nil {
+					return nil, err
+				}
+				return &c, nil
 			}
-			errc <- hs.ListenAndServeTLS("", "")
-		}()
+			return nil, errors.New("invalid SNI name")
+		}
+	}
+	if err := server.ListenAndServe(); err != nil {
+		log.Fatal(err)
 	}
-	log.Fatal(<-errc)
 }
 
 func devMode() bool { return *httpsAddr == "" && *httpAddr != "" }