From dc06d3e643411517e465276d75b66dbed2898850 Mon Sep 17 00:00:00 2001
From: Brad Fitzpatrick <bradfitz@golang.org>
Date: Wed, 11 Apr 2018 20:51:08 +0000
Subject: [PATCH] cmd/godoc: support http-01 ACME challenge in optional
 autocert support

Using same structure & naming as CL 91518.

Fixes golang/go#23627

Change-Id: Ifb73c77d2c39f9f669d425650f9c5bc31bace196
Reviewed-on: https://go-review.googlesource.com/106455
Run-TryBot: Brad Fitzpatrick <bradfitz@golang.org>
TryBot-Result: Gobot Gobot <gobot@golang.org>
Reviewed-by: Andrew Bonventre <andybons@golang.org>
---
 cmd/godoc/autocert.go | 21 ++++++++++++++++-----
 cmd/godoc/main.go     | 21 ++++++++++++++++-----
 2 files changed, 32 insertions(+), 10 deletions(-)

diff --git a/cmd/godoc/autocert.go b/cmd/godoc/autocert.go
index 9fc3a8fc..9c2da414 100644
--- a/cmd/godoc/autocert.go
+++ b/cmd/godoc/autocert.go
@@ -32,21 +32,28 @@ var (
 )
 
 func init() {
-	serveAutoCertHook = serveAutoCert
+	runHTTPS = runHTTPSAutocert
+	certInit = certInitAutocert
+	wrapHTTPMux = wrapHTTPMuxAutocert
 }
 
-func serveAutoCert(h http.Handler) error {
-	m := autocert.Manager{
+var autocertManager *autocert.Manager
+
+func certInitAutocert() {
+	autocertManager = &autocert.Manager{
 		Cache:  autocert.DirCache(*autoCertDirFlag),
 		Prompt: autocert.AcceptTOS,
 	}
 	if *autoCertHostFlag != "" {
-		m.HostPolicy = autocert.HostWhitelist(*autoCertHostFlag)
+		autocertManager.HostPolicy = autocert.HostWhitelist(*autoCertHostFlag)
 	}
+}
+
+func runHTTPSAutocert(h http.Handler) error {
 	srv := &http.Server{
 		Handler: h,
 		TLSConfig: &tls.Config{
-			GetCertificate: m.GetCertificate,
+			GetCertificate: autocertManager.GetCertificate,
 		},
 		IdleTimeout: 60 * time.Second,
 	}
@@ -58,6 +65,10 @@ func serveAutoCert(h http.Handler) error {
 	return srv.Serve(tls.NewListener(tcpKeepAliveListener{ln.(*net.TCPListener)}, srv.TLSConfig))
 }
 
+func wrapHTTPMuxAutocert(h http.Handler) http.Handler {
+	return autocertManager.HTTPHandler(h)
+}
+
 // tcpKeepAliveListener sets TCP keep-alive timeouts on accepted
 // connections. It's used by ListenAndServe and ListenAndServeTLS so
 // dead TCP connections (e.g. closing laptop mid-download) eventually
diff --git a/cmd/godoc/main.go b/cmd/godoc/main.go
index aa2996a4..5e5bc225 100644
--- a/cmd/godoc/main.go
+++ b/cmd/godoc/main.go
@@ -165,6 +165,10 @@ func main() {
 	flag.Usage = usage
 	flag.Parse()
 
+	if certInit != nil {
+		certInit()
+	}
+
 	playEnabled = *showPlayground
 
 	// Check usage: server and no args.
@@ -325,9 +329,9 @@ func main() {
 			go analysis.Run(pointerAnalysis, &corpus.Analysis)
 		}
 
-		if serveAutoCertHook != nil {
+		if runHTTPS != nil {
 			go func() {
-				if err := serveAutoCertHook(handler); err != nil {
+				if err := runHTTPS(handler); err != nil {
 					log.Fatalf("ListenAndServe TLS: %v", err)
 				}
 			}()
@@ -337,6 +341,9 @@ func main() {
 		if *verbose {
 			log.Println("starting HTTP server")
 		}
+		if wrapHTTPMux != nil {
+			handler = wrapHTTPMux(handler)
+		}
 		if err := http.ListenAndServe(*httpAddr, handler); err != nil {
 			log.Fatalf("ListenAndServe %s: %v", *httpAddr, err)
 		}
@@ -354,6 +361,10 @@ func main() {
 	}
 }
 
-// serveAutoCertHook if non-nil specifies a function to listen on port 443.
-// See autocert.go.
-var serveAutoCertHook func(http.Handler) error
+// Hooks that are set non-nil in autocert.go if the "autocert" build tag
+// is used.
+var (
+	certInit    func()
+	runHTTPS    func(http.Handler) error
+	wrapHTTPMux func(http.Handler) http.Handler
+)