From 70c5d8b9cea707b20a9f5aa1b0888776326d55be Mon Sep 17 00:00:00 2001
From: adnano <me@adnano.co>
Date: Mon, 28 Sep 2020 02:05:37 -0400
Subject: [PATCH] Add WithInput helper functions

---
 examples/auth/auth.go | 22 ++++++-------
 server.go             | 74 ++++++++++++++++++++++++++++---------------
 2 files changed, 57 insertions(+), 39 deletions(-)

diff --git a/examples/auth/auth.go b/examples/auth/auth.go
index c348648..c126de8 100644
--- a/examples/auth/auth.go
+++ b/examples/auth/auth.go
@@ -46,11 +46,11 @@ 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)
+	handler.HandleFunc("/login/", login)
+	handler.HandleFunc("/login/password/", loginPassword)
+	handler.HandleFunc("/profile/", profile)
+	handler.HandleFunc("/admin/", admin)
+	handler.HandleFunc("/logout/", logout)
 
 	server := &gmi.Server{
 		Certificate: cert,
@@ -74,15 +74,13 @@ func welcome(rw *gmi.ResponseWriter, req *gmi.Request) {
 
 func login(rw *gmi.ResponseWriter, req *gmi.Request) {
 	gmi.WithCertificate(rw, req, func(cert *x509.Certificate) {
-		if username := req.URL.RawQuery; username == "" {
-			gmi.Input(rw, req, "Username")
-		} else {
+		gmi.WithInput(rw, req, "Username", func(username string) {
 			fingerprint := gmi.Fingerprint(cert)
 			sessions[fingerprint] = &session{
 				username: username,
 			}
 			gmi.Redirect(rw, req, "/login/password")
-		}
+		})
 	})
 }
 
@@ -94,9 +92,7 @@ func loginPassword(rw *gmi.ResponseWriter, req *gmi.Request) {
 			return
 		}
 
-		if password := req.URL.RawQuery; password == "" {
-			gmi.SensitiveInput(rw, req, "Password")
-		} else {
+		gmi.WithSensitiveInput(rw, req, "Password", func(password string) {
 			expected := logins[session.username].password
 			if password == expected {
 				session.authorized = true
@@ -104,7 +100,7 @@ func loginPassword(rw *gmi.ResponseWriter, req *gmi.Request) {
 			} else {
 				gmi.SensitiveInput(rw, req, "Wrong password. Try again")
 			}
-		}
+		})
 	})
 }
 
diff --git a/server.go b/server.go
index 9fe4b40..0f5706d 100644
--- a/server.go
+++ b/server.go
@@ -213,6 +213,42 @@ func InputHandler(prompt string) Handler {
 	})
 }
 
+// WithInput either responds to the request with StatusInput if no input
+// is provided, or calls f with the input when provided.
+func WithInput(rw *ResponseWriter, req *Request, prompt string, f func(string)) {
+	input := req.URL.RawQuery
+	if input == "" {
+		Input(rw, req, prompt)
+		return
+	}
+	f(input)
+}
+
+// Sensitive responds to the request with a request for sensitive input
+// using the given prompt.
+func SensitiveInput(rw *ResponseWriter, req *Request, prompt string) {
+	rw.WriteHeader(StatusSensitiveInput, prompt)
+}
+
+// SensitiveInputHandler returns a simpler handler that responds to each request
+// with a request for sensitive input.
+func SensitiveInputHandler(prompt string) Handler {
+	return HandlerFunc(func(rw *ResponseWriter, req *Request) {
+		SensitiveInput(rw, req, prompt)
+	})
+}
+
+// WithSensitiveInput either responds to the request with StatusSensitiveInput
+// if no input is provided, or calls f with the input when provided.
+func WithSensitiveInput(rw *ResponseWriter, req *Request, prompt string, f func(string)) {
+	input := req.URL.RawQuery
+	if input == "" {
+		SensitiveInput(rw, req, prompt)
+		return
+	}
+	f(input)
+}
+
 // Redirect replies to the request with a redirect to the given URL.
 func Redirect(rw *ResponseWriter, req *Request, url string) {
 	rw.WriteHeader(StatusRedirect, url)
@@ -241,32 +277,6 @@ func PermanentRedirectHandler(url string) Handler {
 	})
 }
 
-// Sensitive responds to the request with a request for sensitive input
-// using the given prompt.
-func SensitiveInput(rw *ResponseWriter, req *Request, prompt string) {
-	rw.WriteHeader(StatusSensitiveInput, prompt)
-}
-
-// SensitiveInputHandler returns a simpler handler that responds to each request
-// with a request for sensitive input.
-func SensitiveInputHandler(prompt string) Handler {
-	return HandlerFunc(func(rw *ResponseWriter, req *Request) {
-		SensitiveInput(rw, req, prompt)
-	})
-}
-
-// CertificateRequired responds to the request with the CertificateRequired
-// status code.
-func CertificateRequired(rw *ResponseWriter, req *Request) {
-	rw.WriteHeader(StatusCertificateRequired, "Certificate required")
-}
-
-// CertificateNotAuthorized responds to the request with
-// the CertificateNotAuthorized status code.
-func CertificateNotAuthorized(rw *ResponseWriter, req *Request) {
-	rw.WriteHeader(StatusCertificateNotAuthorized, "Certificate not authorized")
-}
-
 // NotFound replies to the request with the NotFound status code.
 func NotFound(rw *ResponseWriter, req *Request) {
 	rw.WriteHeader(StatusNotFound, "Not found")
@@ -289,6 +299,18 @@ func GoneHandler() Handler {
 	return HandlerFunc(Gone)
 }
 
+// CertificateRequired responds to the request with the CertificateRequired
+// status code.
+func CertificateRequired(rw *ResponseWriter, req *Request) {
+	rw.WriteHeader(StatusCertificateRequired, "Certificate required")
+}
+
+// CertificateNotAuthorized responds to the request with
+// the CertificateNotAuthorized status code.
+func CertificateNotAuthorized(rw *ResponseWriter, req *Request) {
+	rw.WriteHeader(StatusCertificateNotAuthorized, "Certificate not authorized")
+}
+
 // WithCertificate responds with CertificateRequired if the client did not
 // provide a certificate, and calls f with the first ceritificate if they did.
 func WithCertificate(rw *ResponseWriter, req *Request, f func(*x509.Certificate)) {