Using a custom http.ResponseWriter to write cookies based on the response from a proxied request?

Question

My original question here was flagged as a duplicate of this question. I had no luck implementing it and suspect my problem is misunderstood, so with my question closed, I'm starting fresh with a more specific question.

I'm trying to set a cookie based on a response header from within middleware in request that is reverse proxied.

Here's the workflow:

• User requests http://example.com/foo/bar
• Go app uses ReverseProxy to proxy that request to http://baz.com
• baz.com sets a response header X-FOO
• Go app modifies response by setting a MYAPPFOO cookie with the value of the X-FOO response header
• The cookie is written to the user's browser

It was suggested that a custom http.ResponseWriter will work, but after trying and searching for more information, it is not clear how to approach this.

Since I'm failing to grasp the concept of a custom ResponseWriter for my use case, I'll post code that demonstrates more precisely what I was trying to do at the point I got stuck:

package main

import (

    "github.com/gorilla/mux"
    "log"
    "net/http"
    "net/http/httputil"
    "net/url"
)

func setCookie(w http.ResponseWriter, name string, value string) {
    ...
    http.SetCookie(w, &cookie)
}

func handler(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {

        // setCookie() works here
        // but I cannot access w.Header().Get("X-FOO")

        next.ServeHTTP(w, r)

        // I can access w.Header().Get("X-FOO") here
        // but setCookie() does not cookie the user's browser

        // If I could do it all in one place, this is what I would do:
        if r.Method == "POST" && r.URL.String() == "/login" {
            foo := w.Header().Get("X-FOO")
            setCookie(w, "MYAPPFOO", foo)

        }
    })
}

func main() {

    r := mux.NewRouter()
    r.Use(handler)
    proxy := httputil.NewSingleHostReverseProxy("https://baz.example.com/")
    r.PathPrefix("/").Handler(proxy)
    log.Fatal(http.ListenAndServe(":9001", r))
}

As a side note, I was able to make this work with ReverseProxy.ModifyResponse as recommended in the comments of my last question but I'd really like to achieve this with middleware to keep the code that dynamically creates proxies from config clean. (not in the example code)

Answer 1

From the documentation on http.ResponseWriter methods: (emphasis added)

• Header() http.Header:

  Changing the header map after a call to WriteHeader (or Write) has no effect unless the modified headers are trailers.

• WriteHeader(statusCode int):

  WriteHeader sends an HTTP response header with the provided status code.

• Write([]byte) (int, error):

  If WriteHeader has not yet been called, Write calls WriteHeader(http.StatusOK) before writing the data.

This should highlight the reason why, you can't set a cookie after the next.ServeHTTP(w, r) call, which is that one of the handlers in the middleware chain executed by that call is calling either WriteHeader or Write directly or indirectly.

So to be able set the cookie after the next.ServeHTTP(w, r) call you need to make sure that none of the handlers in the middleware chain calls WriteHeader or Write on the original http.ResponseWriter instance. One way to do this is to wrap the original instance in a custom http.ResponseWriter implementation that will postpone the writing of the response until after you're done with setting the cookie.

---

For example something like this:

type responsewriter struct {
    w    http.ResponseWriter
    buf  bytes.Buffer
    code int
}

func (rw *responsewriter) Header() http.Header {
    return rw.w.Header()
}

func (rw *responsewriter) WriteHeader(statusCode int) {
    rw.code = statusCode
}

func (rw *responsewriter) Write(data []byte) (int, error) {
    return rw.buf.Write(data)
}

func (rw *responsewriter) Done() (int64, error) {
    if rw.code > 0 {
        rw.w.WriteHeader(rw.code)
    }
    return io.Copy(rw.w, &rw.buf)
}

And you would use it like this in your middleware:

func handler(next http.Handler) http.Handler {
    return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        rw := &responsewriter{w: w}
        next.ServeHTTP(rw, r)

        if r.Method == "POST" && r.URL.String() == "/login" {
            foo := rw.Header().Get("X-FOO")
            setCookie(rw, "MYAPPFOO", foo)
        }

        if _, err := rw.Done(); err != nil {
            log.Println(err)
        }
    })
}