From 79a030960b21ce7792205724081eaab8127639b5 Mon Sep 17 00:00:00 2001
From: Mariano Cano <mariano@smallstep.com>
Date: Wed, 13 Feb 2019 14:31:08 -0800
Subject: [PATCH] Enable client certificate rotation with GetClientCertificate.

---
 autocert/examples/hello-mtls/README.md        |  2 +-
 .../hello-mtls/go-grpc/client/client.go       |  6 ++
 .../examples/hello-mtls/go/client/client.go   | 69 +++++++++++++++++--
 .../examples/hello-mtls/go/server/server.go   |  7 +-
 4 files changed, 72 insertions(+), 12 deletions(-)

diff --git a/autocert/examples/hello-mtls/README.md b/autocert/examples/hello-mtls/README.md
index d228bf13..7fb95093 100644
--- a/autocert/examples/hello-mtls/README.md
+++ b/autocert/examples/hello-mtls/README.md
@@ -59,7 +59,7 @@ languages are appreciated!
   - [ ] Root certificate rotation
 - [X] Client using autocert root certificate
   - [X] mTLS (send client certificate if server asks for it)
-  - [ ] Automatic certificate rotation
+  - [X] Automatic certificate rotation
   - [X] Restrict to safe ciphersuites and TLS versions
   - [ ] TLS stack configuration loaded from `step-ca`
   - [ ] Root certificate rotation
diff --git a/autocert/examples/hello-mtls/go-grpc/client/client.go b/autocert/examples/hello-mtls/go-grpc/client/client.go
index d3cc79c7..33fc6273 100644
--- a/autocert/examples/hello-mtls/go-grpc/client/client.go
+++ b/autocert/examples/hello-mtls/go-grpc/client/client.go
@@ -112,6 +112,12 @@ func main() {
 			tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
 			tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
 		},
+		// GetClientCertificate is called when a server requests a
+		// certificate from a client.
+		//
+		// In this example keep alives will cause the certificate to
+		// only be called once, but if we disable them,
+		// GetClientCertificate will be called on every request.
 		GetClientCertificate: r.getClientCertificate,
 	}
 
diff --git a/autocert/examples/hello-mtls/go/client/client.go b/autocert/examples/hello-mtls/go/client/client.go
index 7ef4de6c..0d2da881 100644
--- a/autocert/examples/hello-mtls/go/client/client.go
+++ b/autocert/examples/hello-mtls/go/client/client.go
@@ -10,6 +10,7 @@ import (
 	"net/http"
 	"os"
 	"strings"
+	"sync"
 	"time"
 )
 
@@ -18,8 +19,34 @@ const (
 	autocertKey      = "/var/run/autocert.step.sm/site.key"
 	autocertRoot     = "/var/run/autocert.step.sm/root.crt"
 	requestFrequency = 5 * time.Second
+	tickFrequency    = 15 * time.Second
 )
 
+type rotator struct {
+	sync.RWMutex
+	certificate *tls.Certificate
+}
+
+func (r *rotator) getClientCertificate(cri *tls.CertificateRequestInfo) (*tls.Certificate, error) {
+	r.RLock()
+	defer r.RUnlock()
+	return r.certificate, nil
+}
+
+func (r *rotator) loadCertificate(certFile, keyFile string) error {
+	r.Lock()
+	defer r.Unlock()
+
+	c, err := tls.LoadX509KeyPair(certFile, keyFile)
+	if err != nil {
+		return err
+	}
+
+	r.certificate = &c
+
+	return nil
+}
+
 func loadRootCertPool() (*x509.CertPool, error) {
 	root, err := ioutil.ReadFile(autocertRoot)
 	if err != nil {
@@ -37,34 +64,62 @@ func loadRootCertPool() (*x509.CertPool, error) {
 func main() {
 	url := os.Getenv("HELLO_MTLS_URL")
 
-	// Read our leaf certificate and key from disk
-	cert, err := tls.LoadX509KeyPair(autocertFile, autocertKey)
-	if err != nil {
-		log.Fatal(err)
-	}
-
 	// Read the root certificate for our CA from disk
 	roots, err := loadRootCertPool()
 	if err != nil {
 		log.Fatal(err)
 	}
 
+	// Load certificate
+	r := &rotator{}
+	if err := r.loadCertificate(autocertFile, autocertKey); err != nil {
+		log.Fatal("error loading certificate and key", err)
+	}
+
 	// Create an HTTPS client using our cert, key & pool
 	client := &http.Client{
 		Transport: &http.Transport{
 			TLSClientConfig: &tls.Config{
 				RootCAs:          roots,
-				Certificates:     []tls.Certificate{cert},
 				MinVersion:       tls.VersionTLS12,
 				CurvePreferences: []tls.CurveID{tls.CurveP521, tls.CurveP384, tls.CurveP256},
 				CipherSuites: []uint16{
 					tls.TLS_ECDHE_ECDSA_WITH_CHACHA20_POLY1305,
 					tls.TLS_ECDHE_ECDSA_WITH_AES_128_GCM_SHA256,
 				},
+				// GetClientCertificate is called when a server requests a
+				// certificate from a client.
+				//
+				// In this example keep alives will cause the certificate to
+				// only be called once, but if we disable them,
+				// GetClientCertificate will be called on every request.
+				GetClientCertificate: r.getClientCertificate,
 			},
+			// Add this line to get the certificate on every request.
+			// DisableKeepAlives: true,
 		},
 	}
 
+	// Schedule periodic re-load of certificate
+	done := make(chan struct{})
+	go func() {
+		ticker := time.NewTicker(tickFrequency)
+		defer ticker.Stop()
+		for {
+			select {
+			case <-ticker.C:
+				fmt.Println("Checking for new certificate...")
+				err := r.loadCertificate(autocertFile, autocertKey)
+				if err != nil {
+					log.Println("Error loading certificate and key", err)
+				}
+			case <-done:
+				return
+			}
+		}
+	}()
+	defer close(done)
+
 	for {
 		// Make request
 		r, err := client.Get(url)
diff --git a/autocert/examples/hello-mtls/go/server/server.go b/autocert/examples/hello-mtls/go/server/server.go
index 9449b5f1..6413b796 100644
--- a/autocert/examples/hello-mtls/go/server/server.go
+++ b/autocert/examples/hello-mtls/go/server/server.go
@@ -23,14 +23,13 @@ const (
 // to automatically rotate certificates when they're renewed.
 
 type rotator struct {
-	sync.Mutex
+	sync.RWMutex
 	certificate *tls.Certificate
 }
 
 func (r *rotator) getCertificate(hello *tls.ClientHelloInfo) (*tls.Certificate, error) {
-	r.Lock()
-	defer r.Unlock()
-
+	r.RLock()
+	defer r.RUnlock()
 	return r.certificate, nil
 }