From e52aba264bfeb2d8ed1ce386f9ce4c229f17c206 Mon Sep 17 00:00:00 2001
From: Daisuke Maki <lestrrat+github@gmail.com>
Date: Sat, 22 Apr 2017 10:17:17 +0900
Subject: [PATCH] Checks for Host header when issuing selftest

1. Upon request to verifyDomain(), we record the domain that
we would like to test for reacheability
2. The selftest handler does a check on the Host header value
against this domain name. it returns 200 if the domain is present,
and 404 otherwise
3. Once the verifyDomain() test is done, this record is erased,
ensuring that future requests to the selftest endpoing returns 404
---
 pkg/acme/acme.go         |  7 +++++++
 pkg/acme/cert_request.go | 21 +++++++++++++++++++++
 pkg/acme/type.go         |  3 +++
 3 files changed, 31 insertions(+)

diff --git a/pkg/acme/acme.go b/pkg/acme/acme.go
index 9bdf8452..69c08f10 100644
--- a/pkg/acme/acme.go
+++ b/pkg/acme/acme.go
@@ -18,6 +18,7 @@ func New(kubeLego kubelego.KubeLego) *Acme {
 		kubelego:              kubeLego,
 		challengesHostToToken: map[string]string{},
 		challengesTokenToKey:  map[string]string{},
+		domainsToVerify:       map[string]struct{}{},
 		id:                    utils.RandomToken(16),
 	}
 	if kubeLego != nil {
@@ -57,6 +58,12 @@ func (a *Acme) Mux() *http.ServeMux {
 	mux.HandleFunc(kubelego.AcmeHttpChallengePath+"/", a.handleChallenge)
 
 	mux.HandleFunc(kubelego.AcmeHttpSelfTest, func(w http.ResponseWriter, r *http.Request) {
+		host := strings.Split(r.Host, ":")[0]
+		if !a.domainToVerifyExists(host) {
+			w.WriteHeader(http.StatusNotFound)
+			return
+		}
+
 		w.Header().Set("Content-Type", "text/plain")
 		w.WriteHeader(http.StatusOK)
 		fmt.Fprint(w, a.id)
diff --git a/pkg/acme/cert_request.go b/pkg/acme/cert_request.go
index 0553d93d..4dd5d553 100644
--- a/pkg/acme/cert_request.go
+++ b/pkg/acme/cert_request.go
@@ -73,7 +73,28 @@ func (a *Acme) testReachablilty(domain string) error {
 	return nil
 }
 
+func (a *Acme) enableDomainVerification(domain string) {
+	a.domainsMutex.Lock()
+	a.domainsToVerify[domain] = struct{}{}
+	a.domainsMutex.Unlock()
+}
+
+func (a *Acme) disableDomainVerification(domain string) {
+	a.domainsMutex.Lock()
+	delete(a.domainsToVerify, domain)
+	a.domainsMutex.Unlock()
+}
+
+func (a *Acme) domainToVerifyExists(domain string) bool {
+	a.domainsMutex.RLock()
+	_, ok := a.domainsToVerify[domain]
+	a.domainsMutex.RUnlock()
+	return ok
+}
+
 func (a *Acme) verifyDomain(domain string) (auth *acme.Authorization, err error) {
+	a.enableDomainVerification(domain)
+	defer a.disableDomainVerification(domain)
 	err = a.testReachablilty(domain)
 	if err != nil {
 		return nil, fmt.Errorf("reachability test failed: %s", err)
diff --git a/pkg/acme/type.go b/pkg/acme/type.go
index 6eeb61a3..f03b1367 100644
--- a/pkg/acme/type.go
+++ b/pkg/acme/type.go
@@ -23,6 +23,9 @@ type Acme struct {
 	challengesHostToToken map[string]string
 	challengesTokenToKey  map[string]string
 
+	domainsMutex    sync.RWMutex
+	domainsToVerify map[string]struct{}
+
 	notFound string // string displayed for 404 messages
 	id       string // identification (random string)