internal/worker,middleware: validate IAP header

Validate the header that comes from the IAP. This ensures that
no one is accessing the server in a way that bypasses the IAP.

The worker will perform this validation, but only if
an env var is set, so it won't affect local testing and
we can turn it on or off easily.

Change-Id: I5932548d8a4ef0502f0c936b394f711ac310bcd3
Reviewed-on: https://go-review.googlesource.com/c/pkgsite/+/256759
Trust: Jonathan Amsterdam <jba@google.com>
Run-TryBot: Jonathan Amsterdam <jba@google.com>
TryBot-Result: kokoro <noreply+kokoro@google.com>
Reviewed-by: Julie Qiu <julie@golang.org>
diff --git a/cmd/worker/main.go b/cmd/worker/main.go
index 50ccbc9..f3873e5 100644
--- a/cmd/worker/main.go
+++ b/cmd/worker/main.go
@@ -143,9 +143,15 @@
 	if err != nil {
 		log.Fatalf(ctx, "strconv.Atoi(%q): %v", timeout, err)
 	}
+	iap := middleware.Identity()
+	if aud := os.Getenv("GO_DISCOVERY_IAP_AUDIENCE"); aud != "" {
+		iap = middleware.ValidateIAPHeader(aud)
+	}
+
 	mw := middleware.Chain(
 		middleware.RequestLog(cmdconfig.Logger(ctx, cfg, "worker-log")),
 		middleware.Timeout(time.Duration(handlerTimeout)*time.Minute),
+		iap,
 		middleware.Experiment(experimenter),
 	)
 	http.Handle("/", mw(router))
diff --git a/internal/middleware/iapheader.go b/internal/middleware/iapheader.go
new file mode 100644
index 0000000..18a36b6
--- /dev/null
+++ b/internal/middleware/iapheader.go
@@ -0,0 +1,44 @@
+// Copyright 2019 The Go Authors. All rights reserved.
+// Use of this source code is governed by a BSD-style
+// license that can be found in the LICENSE file.
+
+package middleware
+
+import (
+	"context"
+	"errors"
+	"fmt"
+	"net/http"
+
+	"google.golang.org/api/idtoken"
+)
+
+// ValidateIAPHeader checks that the request has a header that proves it arrived
+// via the IAP.
+// See https://cloud.google.com/iap/docs/signed-headers-howto#securing_iap_headers.
+func ValidateIAPHeader(audience string) Middleware {
+	return func(h http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			// Health checks don't come from the IAP; allow them.
+			if r.URL.Path != "/healthz" {
+				// Adapted from https://github.com/GoogleCloudPlatform/golang-samples/blob/master/iap/validate.go
+				token := r.Header.Get("X-Goog-IAP-JWT-Assertion")
+				if err := validateIAPToken(r.Context(), token, audience); err != nil {
+					http.Error(w, err.Error(), http.StatusUnauthorized)
+					return
+				}
+			}
+			h.ServeHTTP(w, r)
+		})
+	}
+}
+
+func validateIAPToken(ctx context.Context, iapJWT, audience string) error {
+	if iapJWT == "" {
+		return errors.New("missing IAP token")
+	}
+	if _, err := idtoken.Validate(ctx, iapJWT, audience); err != nil {
+		return fmt.Errorf("validating IPA token: %v", err)
+	}
+	return nil
+}