如何在 ReverseProxy 中测试标头？

Question

我正在尝试对以下代码进行单元测试：

func(h *Handler)Forward(w http.ResponseWriter, r *http.Request) {

    url, err := url.Parse("http://test.com")
    if err != nil {
       return
    }

    reverseProxy := &httputil.ReverseProxy{
        Director: func(r *http.Request) {
            r.URL.Host = url.Host
            r.URL.Path = "/"
            r.URL.Scheme = url.Scheme
            r.Host = url.Host
            r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))
        },
    }


    reverseProxy.ServeHTTP(w, r)
}

我不知道如何测试标头是否被 Director 函数修改。我们如何在 Go 的反向代理中测试标头？

Answer 1

1。将外部依赖注入到您的测试单元中

我现在可以看到的最大问题是您转发到的 URL 是硬编码在您的函数中的。这使得单元测试变得非常困难。所以第一步是从函数中提取 URL。在不知道其余代码的情况下，Handler 似乎是一个不错的地方。简化：

type Handler struct {
    backend *url.URL
}

func NewHandler() (*Handler, error) {
    backend, err := url.Parse("http://test.com")
    if err != nil {
        return nil, err
    }
    return &Handler{backend}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    reverseProxy := &httputil.ReverseProxy{
        Director: func(r *http.Request) {
            r.URL.Host = h.backend.Host
            r.URL.Path = "/"
            r.URL.Scheme = h.backend.Scheme
            r.Host = url.Host
            r.Header.Set("X-Forwarded-Host", r.Header.Get("Host"))      
        },
    }
    reverseProxy.ServeHTTP(w, r)
}

请注意，我已将 Forward 重命名为 ServeHTTP 以简化此示例。

2。使用httptest 进行实时处理程序测试

下一步是进行基本测试：

func TestHandler(t *testing.T) {
    // 1. set-up a backend server
    // 2. set-up a reverse proxy with the handler we are testing
    // 3. call the reverse-proxy
    // 4. check that the backend server received the correct header

}

让我们从填写简单的部分开始：

// set-up a backend server 
backendServer := httptest.NewServer(http.DefaultServeMux)
defer backendServer.Close()

backendURL, err := url.Parse(backendServer.URL)
if err != nil {
    t.Fatal(err)
}

// set-up the reverse proxy
handler := &Handler{backend: backendURL} // <-- here we inject our own endpoint!
reverseProxy := httptest.NewServer(handler)
defer reverseProxy.Close()

reverseProxyURL, err := url.Parse(reverseProxy.URL)
if err != nil {
    t.Fatal(err)
}

// call the reverse proxy
res, err := http.Get(reverseProxy.URL)
if err != nil {
    t.Fatal(err)
}
// todo optional: assert properties of the response
_ = res


// check that the backend server received the correct header
// this comes next...

3。将结果从测试服务器传递给测试

现在我们需要一种将接收到的标头传达给主测试的方法。由于我们的测试服务器可以使用任意处理程序，让我们扩展后端服务器的设置。

var (
    mu     sync.Mutex
    header string
)

// create a backend server that checks the incoming headers
backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
    mu.Lock()
    defer mu.Unlock()
    header = r.Header.Get("X-Forwarded-Host")
    w.WriteHeader(http.StatusOK)
}))
defer backendServer.Close()

请注意我是如何使用互斥锁的，因为处理程序将在不同的 go-routine 中运行。您也可以使用频道。

此时，我们可以实现我们的断言了：

mu.Lock()
got := header
mu.Unlock()

// check that the header has been set
want := reverseProxyURL.Host
if got != want {
    t.Errorf("GET %s gives header %s, got %s", reverseProxy.URL, want, got)
}

请注意，这仍然会失败，但这一次是因为您的测试代码是错误的 :-) r.Header.Get("Host") 应替换为 r.Host。

附录：完整示例

package example

import (
    "net/http"
    "net/http/httptest"
    "net/http/httputil"
    "net/url"
    "sync"
    "testing"
)

type Handler struct {
    backend *url.URL
}

func NewHandler() (*Handler, error) {
    backend, err := url.Parse("http://test.com")
    if err != nil {
        return nil, err
    }
    return &Handler{backend}, nil
}

func (h *Handler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
    reverseProxy := &httputil.ReverseProxy{
        Director: func(r *http.Request) {
            r.URL.Host = h.backend.Host
            r.URL.Path = "/"
            r.URL.Scheme = h.backend.Scheme
            r.Header.Set("X-Forwarded-Host", r.Host)
            r.Host = h.backend.Host
        },
    }
    reverseProxy.ServeHTTP(w, r)
}

func TestHandler(t *testing.T) {
    var (
        mu     sync.Mutex
        header string
    )

    // create a backend server that checks the incoming headers
    backendServer := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
        mu.Lock()
        defer mu.Unlock()
        header = r.Header.Get("X-Forwarded-Host")
        w.WriteHeader(http.StatusOK)
    }))
    defer backendServer.Close()

    backendURL, err := url.Parse(backendServer.URL)
    if err != nil {
        t.Fatal(err)
    }

    // create a server for your reverse proxy
    handler := &Handler{backend: backendURL}
    reverseProxy := httptest.NewServer(handler)
    defer reverseProxy.Close()

    reverseProxyURL, err := url.Parse(reverseProxy.URL)
    if err != nil {
        t.Fatal(err)
    }

    // make a request to the reverse proxy
    res, err := http.Get(reverseProxy.URL)
    if err != nil {
        t.Fatal(err)
    }
    // todo optional: assert properties of the response
    _ = res

    mu.Lock()
    got := header
    mu.Unlock()

    // check that the header has been set
    want := reverseProxyURL.Host
    if got != want {
        t.Errorf("GET %s gives header %s, got %s", reverseProxy.URL, want, got)
    }
}