Repository: google/go-safeweb (go)

Issue:
Expose methods for getting and setting the http.Request context from the safehttp.IncomingRequest
Currently, the `safehttp.IncomingRequest` API doesn't provide methods for retrieving and modifying the `context.Context` of the underlying `http.Request`. Plugins need this functionality to add XSRF tokens and CSP nonces.

Pull request:
Added functions for retrieving the context.Context and changing it for a safehttp.IncomingRequest
Fixes #70 

Extended the `safehttp.IncomingRequest` API to add functions for getting and setting the context of the underlying `http.Request` as well as tests.

Patch:
diff --git a/safehttp/incoming_request.go b/safehttp/incoming_request.go
index db56919..a0d48c4 100644
--- a/safehttp/incoming_request.go
+++ b/safehttp/incoming_request.go
@@ -15,6 +15,7 @@
 package safehttp
 
 import (
+	"context"
 	"crypto/tls"
 	"fmt"
 	"net/http"
@@ -145,3 +146,23 @@ func (r *IncomingRequest) Cookies() []*Cookie {
 	}
 	return res
 }
+
+// Context returns the context of a safehttp.IncomingRequest. This is always
+// non-nil and will default to the background context. The context of a
+// safehttp.IncomingRequest is the context of the underlying http.Request.
+//
+// The context is cancelled when the client's connection
+// closes, the request is canceled (with HTTP/2), or when the ServeHTTP method
+// returns.
+func (r *IncomingRequest) Context() context.Context {
+	return r.req.Context()
+}
+
+// SetContext sets the context of the safehttp.IncomingRequest to ctx. The
+// provided context must be non-nil, otherwise the method will panic.
+func (r *IncomingRequest) SetContext(ctx context.Context) {
+	if ctx == nil {
+		panic("nil context")
+	}
+	r.req = r.req.WithContext(ctx)
+}
diff --git a/safehttp/plugins/hsts/hsts.go b/safehttp/plugins/hsts/hsts.go
index e2bfc55..ef979bb 100644
--- a/safehttp/plugins/hsts/hsts.go
+++ b/safehttp/plugins/hsts/hsts.go
@@ -65,7 +65,7 @@ func NewPlugin() Plugin {
 // The function redirects HTTP requests to HTTPS. When HTTPS traffic
 // is received the Strict-Transport-Security header is applied to the
 // response.
-func (p *Plugin) Before(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+func (p Plugin) Before(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
 	if p.MaxAge < 0 {
 		return w.ServerError(safehttp.StatusInternalServerError)
 	}


Test patch:
diff --git a/safehttp/incoming_request_test.go b/safehttp/incoming_request_test.go
index babfe42..10a9145 100644
--- a/safehttp/incoming_request_test.go
+++ b/safehttp/incoming_request_test.go
@@ -15,11 +15,12 @@
 package safehttp_test
 
 import (
+	"context"
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-safeweb/safehttp"
 	"net/http"
 	"net/http/httptest"
 	"testing"
-
-	"github.com/google/go-safeweb/safehttp"
 )
 
 func TestIncomingRequestCookie(t *testing.T) {
@@ -134,5 +135,61 @@ func TestIncomingRequestCookies(t *testing.T) {
 				}
 			}
 		})
+
 	}
 }
+
+type pizza struct {
+	val string
+}
+
+type pizzaKey string
+
+func TestRequestSetValidContextWithValue(t *testing.T) {
+	tests := []struct {
+		name    string
+		key     pizzaKey
+		wantVal *pizza
+		wantOk  bool
+	}{
+		{
+			name:    "Value set for key",
+			key:     pizzaKey("1234"),
+			wantOk:  true,
+			wantVal: &pizza{val: "margeritta"},
+		},
+		{
+			name:    "Value not set for key",
+			key:     pizzaKey("5678"),
+			wantOk:  false,
+			wantVal: nil,
+		},
+	}
+	for _, test := range tests {
+		req := httptest.NewRequest("GET", "/", nil)
+		ir := safehttp.NewIncomingRequest(req)
+		ctx := context.WithValue(ir.Context(), pizzaKey("1234"), &pizza{val: "margeritta"})
+		ir.SetContext(ctx)
+
+		got, ok := ir.Context().Value(test.key).(*pizza)
+		if ok != test.wantOk {
+			t.Errorf("type match: got %v, want %v", ok, test.wantOk)
+		}
+		if diff := cmp.Diff(test.wantVal, got, cmp.AllowUnexported(pizza{})); diff != "" {
+			t.Errorf("ir.Context().Value(test.key): mismatch (-want +got): \n%s", diff)
+		}
+	}
+}
+
+func TestRequestSetNilContext(t *testing.T) {
+	req := httptest.NewRequest("GET", "/", nil)
+	ir := safehttp.NewIncomingRequest(req)
+
+	defer func() {
+		if r := recover(); r == nil {
+			t.Errorf(`ir.SetContext(nil): expected panic`)
+		}
+	}()
+
+	ir.SetContext(nil)
+}
diff --git a/tests/integration/hsts/hsts_test.go b/tests/integration/hsts/hsts_test.go
new file mode 100644
index 0000000..1fa51df
--- /dev/null
+++ b/tests/integration/hsts/hsts_test.go
@@ -0,0 +1,99 @@
+// 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
+//
+// 	https://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 hsts_test
+
+import (
+	"io"
+	"net/http"
+	"net/http/httptest"
+	"strings"
+	"testing"
+
+	"github.com/google/go-cmp/cmp"
+	"github.com/google/go-safeweb/safehttp"
+	"github.com/google/go-safeweb/safehttp/plugins/hsts"
+	"github.com/google/safehtml"
+)
+
+type dispatcher struct{}
+
+func (dispatcher) Write(rw http.ResponseWriter, resp safehttp.Response) error {
+	switch x := resp.(type) {
+	case safehtml.HTML:
+		_, err := rw.Write([]byte(x.String()))
+		return err
+	default:
+		panic("not a safe response type")
+	}
+}
+
+func (dispatcher) ExecuteTemplate(rw http.ResponseWriter, t safehttp.Template, data interface{}) error {
+	return nil
+}
+
+type responseRecorder struct {
+	header http.Header
+	writer io.Writer
+	status int
+}
+
+func newResponseRecorder(w io.Writer) *responseRecorder {
+	return &responseRecorder{header: http.Header{}, writer: w, status: http.StatusOK}
+}
+
+func (r *responseRecorder) Header() http.Header {
+	return r.header
+}
+
+func (r *responseRecorder) WriteHeader(statusCode int) {
+	r.status = statusCode
+}
+
+func (r *responseRecorder) Write(data []byte) (int, error) {
+	return r.writer.Write(data)
+}
+
+func TestHSTSServeMuxInstall(t *testing.T) {
+	mux := safehttp.NewServeMux(&dispatcher{}, "foo.com")
+
+	handler := safehttp.HandlerFunc(func(w safehttp.ResponseWriter, r *safehttp.IncomingRequest) safehttp.Result {
+		return w.Write(safehtml.HTMLEscaped("<h1>Hello World!</h1>"))
+	})
+	mux.Handle("/asdf", safehttp.MethodGet, handler)
+
+	mux.Install("hsts", hsts.NewPlugin())
+
+	b := strings.Builder{}
+	rr := newResponseRecorder(&b)
+
+	req := httptest.NewRequest(http.MethodGet, "https://foo.com/asdf", nil)
+
+	mux.ServeHTTP(rr, req)
+
+	if want := 200; rr.status != want {
+		t.Errorf("rr.status got: %v want: %v", rr.status, want)
+	}
+
+	wantHeaders := map[string][]string{
+		"Strict-Transport-Security": {"max-age=63072000; includeSubDomains"},
+	}
+	if diff := cmp.Diff(wantHeaders, map[string][]string(rr.header)); diff != "" {
+		t.Errorf("rr.header mismatch (-want +got):\n%s", diff)
+	}
+
+	if got, want := b.String(), "&lt;h1&gt;Hello World!&lt;/h1&gt;"; got != want {
+		t.Errorf("b.String() got: %v want: %v", got, want)
+	}
+}