Repository: google/exposure-notifications-server
Pull request #1355 (2021-02-05 22:41:10)

Move maintenance mode to middleware
This will allow us to use it elsewhere, and it decouples from the implementation.

**Release Note**

<!-- Enter your extended release note in the below block. If the PR requires
additional action from users switching to the new release, include the string
"action required". If no release note is required, write "NONE". -->

```release-note
Move maintenance mode handling to middleware
```

/assign @mikehelmick 

google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/1355#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
sethvargo: https://github.com/google/exposure-notifications-server/pull/1355/files?w=1

google-oss-robot: [APPROVALNOTIFIER] This PR is **APPROVED**

This pull-request has been approved by: *<a href="https://github.com/google/exposure-notifications-server/pull/1355#pullrequestreview-584757701" title="Approved">mikehelmick</a>*, *<a href="https://github.com/google/exposure-notifications-server/pull/1355#" title="Author self-approved">sethvargo</a>*

The full list of commands accepted by this bot can be found [here](https://go.k8s.io/bot-commands?repo=google%2Fexposure-notifications-server).

The pull request process is described [here](https://git.k8s.io/community/contributors/guide/owners.md#the-code-review-process)

<details >
Needs approval from an approver in each of these files:

- ~~[OWNERS](https://github.com/google/exposure-notifications-server/blob/main/OWNERS)~~ [mikehelmick,sethvargo]

Approvers can indicate their approval by writing `/approve` in a comment
Approvers can cancel approval by writing `/approve cancel` in a comment
</details>
<!-- META={"approvers":[]} -->
google-oss-robot: @sethvargo: The following test **failed**, say `/retest` to rerun all failed tests:

Test name | Commit | Details | Rerun command
--- | --- | --- | ---
pull-en-server-release-unit | 468de25fdeb9fa7088e3f213ff58078947b9405d | [link](https://oss-prow.knative.dev/view/gs/oss-prow/pr-logs/pull/google_exposure-notifications-server/1355/pull-en-server-release-unit/1357821603454390272) | `/test pull-en-server-release-unit`

<details>

Instructions for interacting with me using PR comments are available [here](https://git.k8s.io/community/contributors/guide/pull-requests.md).  If you have questions or suggestions related to my behavior, please file an issue against the [kubernetes/test-infra](https://github.com/kubernetes/test-infra/issues/new?title=Prow%20issue:) repository. I understand the commands that are listed [here](https://go.k8s.io/bot-commands).
</details>
<!-- test report -->
google-oss-robot: New changes are detected. LGTM label has been removed.

Patch:
diff --git a/internal/maintenance/maintenance.go b/internal/maintenance/maintenance.go
deleted file mode 100644
index 4971628..0000000
--- a/internal/maintenance/maintenance.go
+++ /dev/null
@@ -1,54 +0,0 @@
-// Copyright 2020 Google LLC
-//
-// Licensed under the Apache License, Version 2.0 (the "License");
-// you may not use this file except in compliance with the License.
-// You may obtain a copy of the License at
-//
-//      http://www.apache.org/licenses/LICENSE-2.0
-//
-// Unless required by applicable law or agreed to in writing, software
-// distributed under the License is distributed on an "AS IS" BASIS,
-// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-// See the License for the specific language governing permissions and
-// limitations under the License.
-
-// Package maintenance provides utilities for maintenance mode handling
-package maintenance
-
-import (
-	"fmt"
-	"net/http"
-)
-
-// Config is an interface that determines if
-// the implementer can supply maintenance mode settings.
-type Config interface {
-	MaintenanceMode() bool
-}
-
-// Responder is a handle to a configured maintenance mode responder.
-type Responder struct {
-	inMaintenance bool
-}
-
-// New creates a new maintenance mode responder.
-func New(c Config) *Responder {
-	return &Responder{
-		inMaintenance: c.MaintenanceMode(),
-	}
-}
-
-// Handle will either return the maintenance mode responder (if enabled)
-// or pass through to the next handler.
-func (r *Responder) Handle(next http.Handler) http.Handler {
-	if r.inMaintenance {
-		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-			w.Header().Set("Content-Type", "application/json")
-			w.WriteHeader(http.StatusTooManyRequests)
-			fmt.Fprintf(w, "{\"error\": \"please try again later\"}")
-		})
-	}
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		next.ServeHTTP(w, r)
-	})
-}
diff --git a/internal/middleware/maintenance.go b/internal/middleware/maintenance.go
new file mode 100644
index 0000000..5308766
--- /dev/null
+++ b/internal/middleware/maintenance.go
@@ -0,0 +1,43 @@
+// Copyright 2020 Google LLC
+//
+// Licensed under the Apache License, Version 2.0 (the "License");
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//      http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+
+package middleware
+
+import (
+	"fmt"
+	"net/http"
+
+	"github.com/gorilla/mux"
+)
+
+// Maintainable is an interface that determines if the implementer can supply
+// maintenance mode settings.
+type Maintainable interface {
+	MaintenanceMode() bool
+}
+
+func ProcessMaintenance(cfg Maintainable) mux.MiddlewareFunc {
+	return func(next http.Handler) http.Handler {
+		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			if cfg.MaintenanceMode() {
+				w.Header().Set("Content-Type", "application/json")
+				w.WriteHeader(http.StatusTooManyRequests)
+				fmt.Fprint(w, `{"error": "please try again later"}`)
+				return
+			}
+
+			next.ServeHTTP(w, r)
+		})
+	}
+}
diff --git a/internal/publish/config.go b/internal/publish/config.go
index ac1aaf9..b0395cb 100644
--- a/internal/publish/config.go
+++ b/internal/publish/config.go
@@ -21,7 +21,7 @@ import (
 
 	"github.com/google/exposure-notifications-server/internal/authorizedapp"
 	"github.com/google/exposure-notifications-server/internal/database"
-	"github.com/google/exposure-notifications-server/internal/maintenance"
+	"github.com/google/exposure-notifications-server/internal/middleware"
 	"github.com/google/exposure-notifications-server/internal/publish/model"
 	"github.com/google/exposure-notifications-server/internal/revision"
 	"github.com/google/exposure-notifications-server/internal/setup"
@@ -39,7 +39,7 @@ var _ setup.SecretManagerConfigProvider = (*Config)(nil)
 var _ setup.ObservabilityExporterConfigProvider = (*Config)(nil)
 var _ model.TransformerConfig = (*Config)(nil)
 var _ setup.KeyManagerConfigProvider = (*Config)(nil)
-var _ maintenance.Config = (*Config)(nil)
+var _ middleware.Maintainable = (*Config)(nil)
 
 // Config represents the configuration and associated environment variables for
 // the publish components.
diff --git a/internal/publish/publish.go b/internal/publish/publish.go
index 36e8796..d5646f9 100644
--- a/internal/publish/publish.go
+++ b/internal/publish/publish.go
@@ -137,6 +137,7 @@ func (s *Server) Routes(ctx context.Context, cfg *Config) *mux.Router {
 	r.Use(middleware.PopulateRequestID())
 	r.Use(middleware.PopulateObservability())
 	r.Use(middleware.PopulateLogger(logger))
+	r.Use(middleware.ProcessMaintenance(cfg))
 
 	r.Handle("/health", server.HandleHealthz())
 
diff --git a/internal/publish/publish_v1.go b/internal/publish/publish_v1.go
index cc4cbca..20cf253 100644
--- a/internal/publish/publish_v1.go
+++ b/internal/publish/publish_v1.go
@@ -23,7 +23,6 @@ import (
 	"go.opencensus.io/trace"
 
 	"github.com/google/exposure-notifications-server/internal/jsonutil"
-	"github.com/google/exposure-notifications-server/internal/maintenance"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1"
 	"github.com/google/exposure-notifications-server/pkg/logging"
 	"github.com/mikehelmick/go-chaff"
@@ -62,27 +61,25 @@ func (s *Server) handleRequest(w http.ResponseWriter, r *http.Request) *response
 
 // handlePublishV1 returns an http.Handler that can process V1 publish requests.
 func (s *Server) handlePublishV1() http.Handler {
-	mResponder := maintenance.New(s.config)
 	return s.tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"),
-		mResponder.Handle(
-			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				ctx := r.Context()
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
 
-				logger := logging.FromContext(ctx).Named("handlePublishV1")
+			logger := logging.FromContext(ctx).Named("handlePublishV1")
 
-				response := s.handleRequest(w, r)
+			response := s.handleRequest(w, r)
 
-				if padding, err := generatePadding(s.config.ResponsePaddingMinBytes, s.config.ResponsePaddingRange); err != nil {
-					stats.Record(ctx, mPaddingFailed.M(1))
-					logger.Errorw("failed to pad response", "error", err)
-				} else {
-					response.pubResponse.Padding = padding
-				}
+			if padding, err := generatePadding(s.config.ResponsePaddingMinBytes, s.config.ResponsePaddingRange); err != nil {
+				stats.Record(ctx, mPaddingFailed.M(1))
+				logger.Errorw("failed to pad response", "error", err)
+			} else {
+				response.pubResponse.Padding = padding
+			}
 
-				if response.metrics != nil {
-					response.metrics()
-				}
+			if response.metrics != nil {
+				response.metrics()
+			}
 
-				jsonutil.MarshalResponse(w, response.status, response.pubResponse)
-			})))
+			jsonutil.MarshalResponse(w, response.status, response.pubResponse)
+		}))
 }
diff --git a/internal/publish/publish_v1alpha1.go b/internal/publish/publish_v1alpha1.go
index 37afd1e..a203adc 100644
--- a/internal/publish/publish_v1alpha1.go
+++ b/internal/publish/publish_v1alpha1.go
@@ -23,7 +23,6 @@ import (
 	"go.opencensus.io/trace"
 
 	"github.com/google/exposure-notifications-server/internal/jsonutil"
-	"github.com/google/exposure-notifications-server/internal/maintenance"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1"
 	"github.com/google/exposure-notifications-server/pkg/api/v1alpha1"
 	"github.com/google/exposure-notifications-server/pkg/logging"
@@ -79,36 +78,34 @@ func (s *Server) handleV1Apha1Request(w http.ResponseWriter, r *http.Request) *r
 }
 
 func (s *Server) handlePublishV1Alpha1() http.Handler {
-	mResponder := maintenance.New(s.config)
 	return s.tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"),
-		mResponder.Handle(
-			http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-				ctx := r.Context()
-
-				logger := logging.FromContext(ctx).Named("handlePublishV1Alpha1")
-
-				response := s.handleV1Apha1Request(w, r)
-
-				if padding, err := generatePadding(s.config.ResponsePaddingMinBytes, s.config.ResponsePaddingRange); err != nil {
-					stats.Record(ctx, mPaddingFailed.M(1))
-					logger.Errorw("failed to pad response", "error", err)
-				} else {
-					response.pubResponse.Padding = padding
-				}
-
-				// Downgrade the v1 response to a v1alpha1 response.
-				alpha1Response := &v1alpha1.PublishResponse{
-					RevisionToken:     response.pubResponse.RevisionToken,
-					InsertedExposures: response.pubResponse.InsertedExposures,
-					Error:             response.pubResponse.ErrorMessage,
-					Padding:           response.pubResponse.Padding,
-					Warnings:          response.pubResponse.Warnings,
-				}
-
-				if response.metrics != nil {
-					response.metrics()
-				}
-
-				jsonutil.MarshalResponse(w, response.status, alpha1Response)
-			})))
+		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			ctx := r.Context()
+
+			logger := logging.FromContext(ctx).Named("handlePublishV1Alpha1")
+
+			response := s.handleV1Apha1Request(w, r)
+
+			if padding, err := generatePadding(s.config.ResponsePaddingMinBytes, s.config.ResponsePaddingRange); err != nil {
+				stats.Record(ctx, mPaddingFailed.M(1))
+				logger.Errorw("failed to pad response", "error", err)
+			} else {
+				response.pubResponse.Padding = padding
+			}
+
+			// Downgrade the v1 response to a v1alpha1 response.
+			alpha1Response := &v1alpha1.PublishResponse{
+				RevisionToken:     response.pubResponse.RevisionToken,
+				InsertedExposures: response.pubResponse.InsertedExposures,
+				Error:             response.pubResponse.ErrorMessage,
+				Padding:           response.pubResponse.Padding,
+				Warnings:          response.pubResponse.Warnings,
+			}
+
+			if response.metrics != nil {
+				response.metrics()
+			}
+
+			jsonutil.MarshalResponse(w, response.status, alpha1Response)
+		}))
 }
diff --git a/internal/publish/stats.go b/internal/publish/stats.go
index b588281..ca06611 100644
--- a/internal/publish/stats.go
+++ b/internal/publish/stats.go
@@ -22,7 +22,6 @@ import (
 	"time"
 
 	"github.com/google/exposure-notifications-server/internal/jsonutil"
-	"github.com/google/exposure-notifications-server/internal/maintenance"
 	"github.com/google/exposure-notifications-server/internal/publish/model"
 	verifyapi "github.com/google/exposure-notifications-server/pkg/api/v1"
 	"github.com/google/exposure-notifications-server/pkg/logging"
@@ -32,8 +31,7 @@ import (
 )
 
 func (s *Server) handleStats() http.Handler {
-	mResponder := maintenance.New(s.config)
-	return s.tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"), mResponder.Handle(
+	return s.tracker.HandleTrack(chaff.HeaderDetector("X-Chaff"),
 		http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
 			ctx, span := trace.StartSpan(r.Context(), "(*publish.HandleStats)")
 			defer span.End()
@@ -60,7 +58,7 @@ func (s *Server) handleStats() http.Handler {
 			s.addMetricsPadding(ctx, response)
 
 			jsonutil.MarshalResponse(w, status, response)
-		})))
+		}))
 }
 
 func (s *Server) addMetricsPadding(ctx context.Context, response *verifyapi.StatsResponse) {


Test patch:
diff --git a/internal/maintenance/maintenance_test.go b/internal/middleware/maintenance_test.go
similarity index 60%
rename from internal/maintenance/maintenance_test.go
rename to internal/middleware/maintenance_test.go
index c5dfc0e..7270ec8 100644
--- a/internal/maintenance/maintenance_test.go
+++ b/internal/middleware/maintenance_test.go
@@ -12,8 +12,7 @@
 // See the License for the specific language governing permissions and
 // limitations under the License.
 
-// Package maintenance provides utilities for maintenance mode handling
-package maintenance
+package middleware
 
 import (
 	"net/http"
@@ -29,36 +28,36 @@ func (t *testConfig) MaintenanceMode() bool {
 	return t.enabled
 }
 
-func testHandler(invoked *bool) http.Handler {
-	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
-		*invoked = true
-	})
-}
-
 func TestHandle_Enabled(t *testing.T) {
-	responder := New(&testConfig{true})
+	responder := ProcessMaintenance(&testConfig{true})
 
 	r := &http.Request{}
 	w := httptest.NewRecorder()
 
-	invoked := false
-	responder.Handle(testHandler(&invoked)).ServeHTTP(w, r)
+	responder(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		t.Errorf("handler was invoked")
+	})).ServeHTTP(w, r)
 
-	if invoked {
-		t.Fatalf("handler was invoked while in maintenance mode")
+	w.Flush()
+
+	if got, want := w.Code, 429; got != want {
+		t.Errorf("expected %d to be %d", got, want)
 	}
 }
 
 func TestHandle_Disabled(t *testing.T) {
-	responder := New(&testConfig{false})
+	responder := ProcessMaintenance(&testConfig{false})
 
 	r := &http.Request{}
 	w := httptest.NewRecorder()
 
-	invoked := false
-	responder.Handle(testHandler(&invoked)).ServeHTTP(w, r)
+	responder(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+		w.WriteHeader(200)
+	})).ServeHTTP(w, r)
+
+	w.Flush()
 
-	if !invoked {
-		t.Fatalf("handler was not invoked while not in maintenance mode")
+	if got, want := w.Code, 200; got != want {
+		t.Errorf("expected %d to be %d", got, want)
 	}
 }
diff --git a/internal/publish/publish_test.go b/internal/publish/publish_test.go
index e3695ec..69347a4 100644
--- a/internal/publish/publish_test.go
+++ b/internal/publish/publish_test.go
@@ -125,7 +125,6 @@ func TestPublishWithBypass(t *testing.T) {
 		ErrorCode          string
 		SkipVersions       map[version]bool
 		SkipKeys           map[int]bool // which keys should be skipped (partial success)
-		MaintenanceMode    bool
 	}{
 		{
 			Name:       "successful_insert_bypass_ha_verification",
@@ -421,34 +420,6 @@ func TestPublishWithBypass(t *testing.T) {
 			Code:      http.StatusUnauthorized,
 			Error:     "unable to validate diagnosis verification: token is expired by 1m",
 		},
-		{
-			Name: "maintenance_mode",
-			HealthAuthority: &vermodel.HealthAuthority{
-				Issuer:   issNames.next(),
-				Audience: "unit.test.server",
-				Name:     "Unit Test Gov DOH",
-			},
-			HealthAuthorityKey: &vermodel.HealthAuthorityKey{
-				Version: "v1",
-				From:    time.Now().Add(-1 * time.Minute),
-			},
-			TestRegion: regions.next(),
-			AuthorizedApp: func() *aamodel.AuthorizedApp {
-				authApp := aamodel.NewAuthorizedApp()
-				authApp.AppPackageName = names.next()
-				authApp.AllowedRegions[regions.current()] = struct{}{}
-				return authApp
-			}(),
-			Publish: verifyapi.Publish{
-				Keys:                util.GenerateExposureKeys(2, 5, false),
-				HealthAuthorityID:   names.current(),
-				VerificationPayload: "totally not a JWT",
-			},
-			Regions:         []string{regions.current()},
-			Code:            http.StatusTooManyRequests,
-			Error:           "try again",
-			MaintenanceMode: true,
-		},
 	}
 
 	for _, ver := range versions {
@@ -503,7 +474,6 @@ func TestPublishWithBypass(t *testing.T) {
 				config.ResponsePaddingRange = 100
 				config.MaxMagnitudeSymptomOnsetDays = 14
 				config.MaxSymptomOnsetReportDays = 28
-				config.Maintenance = tc.MaintenanceMode
 				env := serverenv.New(ctx,
 					serverenv.WithDatabase(testDB),
 					serverenv.WithAuthorizedAppProvider(aaProvider),
@@ -612,9 +582,6 @@ func TestPublishWithBypass(t *testing.T) {
 				if resp.StatusCode != tc.Code {
 					t.Fatalf("http response code want: %v, got %v.", tc.Code, resp.StatusCode)
 				}
-				if !tc.MaintenanceMode && len(response.Padding) == 0 {
-					t.Errorf("padding is empty on response: %+v", response)
-				}
 
 				if ver == useV1 {
 					// The extended data validation only happens on verifyapi.