From 1634c2c11c76acb65e8df2e87e02f19537c9ef14 Mon Sep 17 00:00:00 2001
From: Adnan Maolood <me@adnano.co>
Date: Wed, 21 Oct 2020 17:47:34 -0400
Subject: [PATCH] Add Certificate helper function

---
 examples/auth.go | 146 +++++++++++++++++++++++++----------------------
 server.go        |  11 ++--
 2 files changed, 84 insertions(+), 73 deletions(-)

diff --git a/examples/auth.go b/examples/auth.go
index 7803fe7..058458b 100644
--- a/examples/auth.go
+++ b/examples/auth.go
@@ -33,19 +33,19 @@ var (
 )
 
 func main() {
-	handler := &gmi.ServeMux{}
-	handler.HandleFunc("/", welcome)
-	handler.HandleFunc("/login", login)
-	handler.HandleFunc("/login/password", loginPassword)
-	handler.HandleFunc("/profile", profile)
-	handler.HandleFunc("/admin", admin)
-	handler.HandleFunc("/logout", logout)
+	var mux gmi.ServeMux
+	mux.HandleFunc("/", welcome)
+	mux.HandleFunc("/login", login)
+	mux.HandleFunc("/login/password", loginPassword)
+	mux.HandleFunc("/profile", profile)
+	mux.HandleFunc("/admin", admin)
+	mux.HandleFunc("/logout", logout)
 
-	server := &gmi.Server{}
+	var server gmi.Server
 	if err := server.CertificateStore.Load("/var/lib/gemini/certs"); err != nil {
 		log.Fatal(err)
 	}
-	server.Register("localhost", handler)
+	server.Register("localhost", &mux)
 
 	if err := server.ListenAndServe(); err != nil {
 		log.Fatal(err)
@@ -59,78 +59,90 @@ func getSession(crt *x509.Certificate) (*session, bool) {
 }
 
 func welcome(w *gmi.ResponseWriter, r *gmi.Request) {
-	w.Write([]byte("Welcome to this example.\n=> /login Login\n"))
+	fmt.Fprintln(w, "Welcome to this example.")
+	fmt.Fprintln(w, "=> /login Login")
 }
 
 func login(w *gmi.ResponseWriter, r *gmi.Request) {
-	gmi.WithCertificate(w, r, func(cert *x509.Certificate) {
-		username, ok := gmi.Input(w, r, "Username")
-		if !ok {
-			return
-		}
-		fingerprint := gmi.Fingerprint(cert)
-		sessions[fingerprint] = &session{
-			username: username,
-		}
-		gmi.Redirect(w, r, "/login/password")
-	})
+	cert, ok := gmi.Certificate(w, r)
+	if !ok {
+		return
+	}
+	username, ok := gmi.Input(w, r, "Username")
+	if !ok {
+		return
+	}
+	fingerprint := gmi.Fingerprint(cert)
+	sessions[fingerprint] = &session{
+		username: username,
+	}
+	gmi.Redirect(w, r, "/login/password")
 }
 
 func loginPassword(w *gmi.ResponseWriter, r *gmi.Request) {
-	gmi.WithCertificate(w, r, func(cert *x509.Certificate) {
-		session, ok := getSession(cert)
-		if !ok {
-			gmi.CertificateNotAuthorized(w, r)
-			return
-		}
+	cert, ok := gmi.Certificate(w, r)
+	if !ok {
+		return
+	}
+	session, ok := getSession(cert)
+	if !ok {
+		gmi.CertificateNotAuthorized(w, r)
+		return
+	}
 
-		password, ok := gmi.SensitiveInput(w, r, "Password")
-		if !ok {
-			return
-		}
-		expected := logins[session.username].password
-		if password == expected {
-			session.authorized = true
-			gmi.Redirect(w, r, "/profile")
-		} else {
-			gmi.SensitiveInput(w, r, "Wrong password. Try again")
-		}
-	})
+	password, ok := gmi.SensitiveInput(w, r, "Password")
+	if !ok {
+		return
+	}
+	expected := logins[session.username].password
+	if password == expected {
+		session.authorized = true
+		gmi.Redirect(w, r, "/profile")
+	} else {
+		gmi.SensitiveInput(w, r, "Wrong password. Try again")
+	}
 }
 
 func logout(w *gmi.ResponseWriter, r *gmi.Request) {
-	gmi.WithCertificate(w, r, func(cert *x509.Certificate) {
-		fingerprint := gmi.Fingerprint(cert)
-		delete(sessions, fingerprint)
-	})
-	w.Write([]byte("Successfully logged out.\n"))
+	cert, ok := gmi.Certificate(w, r)
+	if !ok {
+		return
+	}
+	fingerprint := gmi.Fingerprint(cert)
+	delete(sessions, fingerprint)
+	fmt.Fprintln(w, "Successfully logged out.")
 }
 
 func profile(w *gmi.ResponseWriter, r *gmi.Request) {
-	gmi.WithCertificate(w, r, func(cert *x509.Certificate) {
-		session, ok := getSession(cert)
-		if !ok {
-			gmi.CertificateNotAuthorized(w, r)
-			return
-		}
-		user := logins[session.username]
-		profile := fmt.Sprintf("Username: %s\nAdmin: %t\n=> /logout Logout", session.username, user.admin)
-		w.Write([]byte(profile))
-	})
+	cert, ok := gmi.Certificate(w, r)
+	if !ok {
+		return
+	}
+	session, ok := getSession(cert)
+	if !ok {
+		gmi.CertificateNotAuthorized(w, r)
+		return
+	}
+	user := logins[session.username]
+	fmt.Fprintln(w, "Username:", session.username)
+	fmt.Fprintln(w, "Admin:", user.admin)
+	fmt.Fprintln(w, "=> /logout Logout")
 }
 
 func admin(w *gmi.ResponseWriter, r *gmi.Request) {
-	gmi.WithCertificate(w, r, func(cert *x509.Certificate) {
-		session, ok := getSession(cert)
-		if !ok {
-			gmi.CertificateNotAuthorized(w, r)
-			return
-		}
-		user := logins[session.username]
-		if !user.admin {
-			gmi.CertificateNotAuthorized(w, r)
-			return
-		}
-		w.Write([]byte("Welcome to the admin portal.\n"))
-	})
+	cert, ok := gmi.Certificate(w, r)
+	if !ok {
+		return
+	}
+	session, ok := getSession(cert)
+	if !ok {
+		gmi.CertificateNotAuthorized(w, r)
+		return
+	}
+	user := logins[session.username]
+	if !user.admin {
+		gmi.CertificateNotAuthorized(w, r)
+		return
+	}
+	fmt.Fprintln(w, "Welcome to the admin portal.")
 }
diff --git a/server.go b/server.go
index 4fde605..d33a234 100644
--- a/server.go
+++ b/server.go
@@ -311,15 +311,14 @@ func CertificateNotAuthorized(w *ResponseWriter, r *Request) {
 	w.WriteHeader(StatusCertificateNotAuthorized, "Certificate not authorized")
 }
 
-// WithCertificate either responds with CertificateRequired if the client did
-// not provide a certificate, or calls f with the first ceritificate provided.
-func WithCertificate(w *ResponseWriter, r *Request, f func(*x509.Certificate)) {
+// Certificate returns the request certificate. If one is not provided,
+// it returns nil and responds with StatusCertificateRequired.
+func Certificate(w *ResponseWriter, r *Request) (*x509.Certificate, bool) {
 	if len(r.TLS.PeerCertificates) == 0 {
 		CertificateRequired(w, r)
-		return
+		return nil, false
 	}
-	cert := r.TLS.PeerCertificates[0]
-	f(cert)
+	return r.TLS.PeerCertificates[0], true
 }
 
 // ResponderFunc is a wrapper around a bare function that implements Handler.