From 29f215a57fa5df6ca8d7523be255d7b51c986b2a Mon Sep 17 00:00:00 2001
From: Mara Mihali <maramihali@google.com>
Date: Wed, 14 Oct 2020 12:20:37 +0300
Subject: [PATCH] Reintroduce the ContentType method to the Dispatcher
 interface and set the CT in the safehttp.ResponseWriter. Refactor some
 documentation.

---
 safehttp/default_dispatcher.go | 35 ++++++++++++++++++++++------------
 safehttp/response_writer.go    | 28 ++++++++++++++++++++++-----
 2 files changed, 46 insertions(+), 17 deletions(-)

diff --git a/safehttp/default_dispatcher.go b/safehttp/default_dispatcher.go
index 431ca7e8..91fe9656 100644
--- a/safehttp/default_dispatcher.go
+++ b/safehttp/default_dispatcher.go
@@ -26,6 +26,25 @@ import (
 // DefaultDispatcher is responsible for writing safe responses.
 type DefaultDispatcher struct{}
 
+// ContentType returns the Content-Type of a response if it's deemed safe and
+// an error otherwise.
+func (DefaultDispatcher) ContentType(resp Response) (string, error) {
+	switch x := resp.(type) {
+	case safehtml.HTML:
+		return "text/html; charset=utf-8", nil
+	case TemplateResponse:
+		_, ok := (*x.Template).(*template.Template)
+		if !ok {
+			return "", fmt.Errorf("%T is not a safe response type, a Content-Type cannot be provided", resp)
+		}
+		return "text/html; charset=utf-8", nil
+	case JSONResponse:
+		return "application/json; charset=utf-8", nil
+	default:
+		return "", fmt.Errorf("%T is not a safe response type, a Content-Type cannot be provided", resp)
+	}
+}
+
 // Write writes the response to the http.ResponseWriter if it's deemed safe.
 // A safe response is either safe HTML, a JSON object or safe HTML template. It
 // will  return a non-nil error if the response is deemed unsafe or if writing
@@ -33,8 +52,6 @@ type DefaultDispatcher struct{}
 func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
 	switch x := resp.(type) {
 	case JSONResponse:
-		rw.Header().Set("Content-Type", "application/json; charset=utf-8")
-		rw.WriteHeader(int(StatusOK))
 		io.WriteString(rw, ")]}',\n") // Break parsing of JavaScript in order to prevent XSSI.
 		return json.NewEncoder(rw).Encode(x.Data)
 	case TemplateResponse:
@@ -42,8 +59,6 @@ func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
 		if !ok {
 			return fmt.Errorf("%T is not a safe template and it cannot be parsed and written", t)
 		}
-		rw.Header().Set("Content-Type", "text/html; charset=utf-8")
-		rw.WriteHeader(int(StatusOK))
 		if len(x.FuncMap) == 0 {
 			return t.Execute(rw, x.Data)
 		}
@@ -52,14 +67,10 @@ func (DefaultDispatcher) Write(rw http.ResponseWriter, resp Response) error {
 			return err
 		}
 		return cloned.Funcs(x.FuncMap).Execute(rw, x.Data)
-	default:
-		r, ok := x.(safehtml.HTML)
-		if !ok {
-			return fmt.Errorf("%T is not a safe response type and it cannot be written", resp)
-		}
-		rw.Header().Set("Content-Type", "text/html; charset=utf-8")
-		rw.WriteHeader(int(StatusOK))
-		_, err := io.WriteString(rw, r.String())
+	case safehtml.HTML:
+		_, err := io.WriteString(rw, x.String())
 		return err
+	default:
+		return fmt.Errorf("%T is not a safe response type and it cannot be written", resp)
 	}
 }
diff --git a/safehttp/response_writer.go b/safehttp/response_writer.go
index 1e44fa53..6ac4a62a 100644
--- a/safehttp/response_writer.go
+++ b/safehttp/response_writer.go
@@ -66,9 +66,8 @@ func NotWritten() Result {
 	return Result{}
 }
 
-// Write dispatches the response to the Dispatcher. If the Dispatcher decides
-// the provided response is safe, it will be written to the underlying
-// http.ResponseWriter.
+// Write dispatches the response to the Dispatcher. This will be written to the
+// underlying http.ResponseWriter if the Dispatcher decides it's safe to do so.
 //
 // TODO: replace panics with proper error handling when writing the response fails.
 func (w *ResponseWriter) Write(resp Response) Result {
@@ -79,7 +78,16 @@ func (w *ResponseWriter) Write(resp Response) Result {
 	if w.written {
 		return Result{}
 	}
+
 	w.markWritten()
+
+	ct, err := w.d.ContentType(resp)
+	if err != nil {
+		panic(err)
+	}
+	w.rw.Header().Set("Content-Type", ct)
+	w.rw.WriteHeader(int(StatusOK))
+
 	if err := w.d.Write(w.rw, resp); err != nil {
 		panic(err)
 	}
@@ -155,8 +163,18 @@ func (w *ResponseWriter) SetCookie(c *Cookie) error {
 // The implementation of a custom Dispatcher should be thoroughly reviewed by
 // the security team to avoid introducing vulnerabilities.
 type Dispatcher interface {
-	// Write writes a Response to the http.ResponseWriter after properly setting
-	// the Content-Type header and the status code.
+	// Content-Type returns the Content-Type of the provided response if it is
+	// of a safe type, supported by the Dispatcher, and should return an error
+	// otherwise.
+	//
+	// Sending a response to the http.ResponseWriter without properly setting
+	// CT is error-prone and could introduce vulnerabilities. Therefore, this
+	// method should be used to set the Content-Type header before calling
+	// Dispatcher.Write. Writing should not proceed if ContentType returns an
+	// error.
+	ContentType(resp Response) (string, error)
+
+	// Write writes a Response to the underlying http.ResponseWriter.
 	//
 	// It should return an error if the writing operation fails or if the
 	// provided Response should not be written to the http.ResponseWriter