2017-06-03 02:55:59 +02:00
|
|
|
package middleware
|
|
|
|
|
|
|
|
import (
|
2020-08-28 02:51:27 +02:00
|
|
|
"bytes"
|
2017-06-03 02:55:59 +02:00
|
|
|
"fmt"
|
2020-08-28 02:51:27 +02:00
|
|
|
"io/ioutil"
|
2020-02-24 18:29:34 +02:00
|
|
|
"net"
|
2017-06-03 02:55:59 +02:00
|
|
|
"net/http"
|
|
|
|
"net/http/httptest"
|
2017-06-04 02:37:35 +02:00
|
|
|
"net/url"
|
2018-02-21 20:44:17 +02:00
|
|
|
"testing"
|
2017-06-04 02:37:35 +02:00
|
|
|
|
2019-01-30 12:56:56 +02:00
|
|
|
"github.com/labstack/echo/v4"
|
2017-06-03 02:55:59 +02:00
|
|
|
"github.com/stretchr/testify/assert"
|
|
|
|
)
|
|
|
|
|
2020-08-28 21:47:02 +02:00
|
|
|
//Assert expected with url.EscapedPath method to obtain the path.
|
2017-06-03 02:55:59 +02:00
|
|
|
func TestProxy(t *testing.T) {
|
|
|
|
// Setup
|
|
|
|
t1 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
fmt.Fprint(w, "target 1")
|
|
|
|
}))
|
|
|
|
defer t1.Close()
|
2017-06-04 02:37:35 +02:00
|
|
|
url1, _ := url.Parse(t1.URL)
|
2017-06-03 02:55:59 +02:00
|
|
|
t2 := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
|
|
|
|
fmt.Fprint(w, "target 2")
|
|
|
|
}))
|
|
|
|
defer t2.Close()
|
2017-06-04 02:37:35 +02:00
|
|
|
url2, _ := url.Parse(t2.URL)
|
|
|
|
|
|
|
|
targets := []*ProxyTarget{
|
2018-02-21 20:44:17 +02:00
|
|
|
{
|
|
|
|
Name: "target 1",
|
|
|
|
URL: url1,
|
2017-06-04 02:37:35 +02:00
|
|
|
},
|
2018-02-21 20:44:17 +02:00
|
|
|
{
|
|
|
|
Name: "target 2",
|
|
|
|
URL: url2,
|
2017-06-04 02:37:35 +02:00
|
|
|
},
|
|
|
|
}
|
2018-02-21 20:44:17 +02:00
|
|
|
rb := NewRandomBalancer(nil)
|
|
|
|
// must add targets:
|
|
|
|
for _, target := range targets {
|
|
|
|
assert.True(t, rb.AddTarget(target))
|
|
|
|
}
|
|
|
|
|
|
|
|
// must ignore duplicates:
|
|
|
|
for _, target := range targets {
|
|
|
|
assert.False(t, rb.AddTarget(target))
|
|
|
|
}
|
2017-06-03 02:55:59 +02:00
|
|
|
|
|
|
|
// Random
|
|
|
|
e := echo.New()
|
2017-06-15 17:19:18 +02:00
|
|
|
e.Use(Proxy(rb))
|
2018-10-14 17:16:58 +02:00
|
|
|
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
2019-01-30 12:56:56 +02:00
|
|
|
rec := httptest.NewRecorder()
|
2017-06-03 02:55:59 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
|
|
|
body := rec.Body.String()
|
2017-06-04 02:37:35 +02:00
|
|
|
expected := map[string]bool{
|
2017-06-03 02:55:59 +02:00
|
|
|
"target 1": true,
|
|
|
|
"target 2": true,
|
|
|
|
}
|
|
|
|
assert.Condition(t, func() bool {
|
2017-06-04 02:37:35 +02:00
|
|
|
return expected[body]
|
2017-06-03 02:55:59 +02:00
|
|
|
})
|
|
|
|
|
2018-02-21 20:44:17 +02:00
|
|
|
for _, target := range targets {
|
|
|
|
assert.True(t, rb.RemoveTarget(target.Name))
|
|
|
|
}
|
|
|
|
|
|
|
|
assert.False(t, rb.RemoveTarget("unknown target"))
|
|
|
|
|
2017-06-03 02:55:59 +02:00
|
|
|
// Round-robin
|
2017-12-26 00:39:37 +02:00
|
|
|
rrb := NewRoundRobinBalancer(targets)
|
2017-06-03 02:55:59 +02:00
|
|
|
e = echo.New()
|
2017-06-15 17:19:18 +02:00
|
|
|
e.Use(Proxy(rrb))
|
2019-01-30 12:56:56 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2017-06-03 02:55:59 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
|
|
|
body = rec.Body.String()
|
|
|
|
assert.Equal(t, "target 1", body)
|
2019-01-30 12:56:56 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2017-06-03 02:55:59 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
|
|
|
body = rec.Body.String()
|
|
|
|
assert.Equal(t, "target 2", body)
|
2017-12-28 01:14:44 +02:00
|
|
|
|
|
|
|
// Rewrite
|
|
|
|
e = echo.New()
|
2017-12-28 21:24:34 +02:00
|
|
|
e.Use(ProxyWithConfig(ProxyConfig{
|
2017-12-28 01:14:44 +02:00
|
|
|
Balancer: rrb,
|
|
|
|
Rewrite: map[string]string{
|
|
|
|
"/old": "/new",
|
|
|
|
"/api/*": "/$1",
|
|
|
|
"/js/*": "/public/javascripts/$1",
|
|
|
|
"/users/*/orders/*": "/user/$1/order/$2",
|
|
|
|
},
|
|
|
|
}))
|
|
|
|
req.URL.Path = "/api/users"
|
2020-08-28 21:47:02 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2017-12-28 01:14:44 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
2020-08-28 21:47:02 +02:00
|
|
|
assert.Equal(t, "/users", req.URL.EscapedPath())
|
|
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
2017-12-28 01:14:44 +02:00
|
|
|
req.URL.Path = "/js/main.js"
|
2020-08-28 21:47:02 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2017-12-28 01:14:44 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
2020-08-28 21:47:02 +02:00
|
|
|
assert.Equal(t, "/public/javascripts/main.js", req.URL.EscapedPath())
|
|
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
2017-12-28 01:14:44 +02:00
|
|
|
req.URL.Path = "/old"
|
2020-08-28 21:47:02 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2017-12-28 01:14:44 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
2020-08-28 21:47:02 +02:00
|
|
|
assert.Equal(t, "/new", req.URL.EscapedPath())
|
|
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
2017-12-28 01:14:44 +02:00
|
|
|
req.URL.Path = "/users/jack/orders/1"
|
2020-08-28 21:47:02 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2017-12-28 01:14:44 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
2020-08-28 21:47:02 +02:00
|
|
|
assert.Equal(t, "/user/jack/order/1", req.URL.EscapedPath())
|
2018-10-13 08:10:19 +02:00
|
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
2020-08-28 02:51:27 +02:00
|
|
|
req.URL.Path = "/users/jill/orders/T%2FcO4lW%2Ft%2FVp%2F"
|
2020-08-28 21:47:02 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2020-04-08 17:19:22 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
2020-08-28 21:47:02 +02:00
|
|
|
assert.Equal(t, "/user/jill/order/T%2FcO4lW%2Ft%2FVp%2F", req.URL.EscapedPath())
|
2020-04-08 17:19:22 +02:00
|
|
|
assert.Equal(t, http.StatusOK, rec.Code)
|
2020-08-28 21:47:02 +02:00
|
|
|
req.URL.Path = "/users/jill/orders/%%%%"
|
|
|
|
rec = httptest.NewRecorder()
|
|
|
|
e.ServeHTTP(rec, req)
|
|
|
|
assert.Equal(t, http.StatusBadRequest, rec.Code)
|
2018-10-09 05:43:39 +02:00
|
|
|
|
2020-08-28 02:51:27 +02:00
|
|
|
// ModifyResponse
|
|
|
|
e = echo.New()
|
|
|
|
e.Use(ProxyWithConfig(ProxyConfig{
|
|
|
|
Balancer: rrb,
|
|
|
|
ModifyResponse: func(res *http.Response) error {
|
|
|
|
res.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("modified")))
|
|
|
|
res.Header.Set("X-Modified", "1")
|
|
|
|
return nil
|
|
|
|
},
|
|
|
|
}))
|
|
|
|
rec = httptest.NewRecorder()
|
|
|
|
e.ServeHTTP(rec, req)
|
|
|
|
assert.Equal(t, "modified", rec.Body.String())
|
|
|
|
assert.Equal(t, "1", rec.Header().Get("X-Modified"))
|
|
|
|
|
2018-10-09 05:43:39 +02:00
|
|
|
// ProxyTarget is set in context
|
|
|
|
contextObserver := func(next echo.HandlerFunc) echo.HandlerFunc {
|
|
|
|
return func(c echo.Context) (err error) {
|
|
|
|
next(c)
|
|
|
|
assert.Contains(t, targets, c.Get("target"), "target is not set in context")
|
|
|
|
return nil
|
|
|
|
}
|
|
|
|
}
|
|
|
|
rrb1 := NewRoundRobinBalancer(targets)
|
|
|
|
e = echo.New()
|
|
|
|
e.Use(contextObserver)
|
|
|
|
e.Use(Proxy(rrb1))
|
2019-01-30 12:56:56 +02:00
|
|
|
rec = httptest.NewRecorder()
|
2018-10-09 05:43:39 +02:00
|
|
|
e.ServeHTTP(rec, req)
|
2017-06-03 02:55:59 +02:00
|
|
|
}
|
2020-02-24 18:29:34 +02:00
|
|
|
|
|
|
|
func TestProxyRealIPHeader(t *testing.T) {
|
|
|
|
// Setup
|
|
|
|
upstream := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {}))
|
|
|
|
defer upstream.Close()
|
|
|
|
url, _ := url.Parse(upstream.URL)
|
|
|
|
rrb := NewRoundRobinBalancer([]*ProxyTarget{{Name: "upstream", URL: url}})
|
|
|
|
e := echo.New()
|
|
|
|
e.Use(Proxy(rrb))
|
|
|
|
req := httptest.NewRequest(http.MethodGet, "/", nil)
|
|
|
|
rec := httptest.NewRecorder()
|
|
|
|
|
|
|
|
remoteAddrIP, _, _ := net.SplitHostPort(req.RemoteAddr)
|
|
|
|
realIPHeaderIP := "203.0.113.1"
|
|
|
|
extractedRealIP := "203.0.113.10"
|
|
|
|
tests := []*struct {
|
|
|
|
hasRealIPheader bool
|
|
|
|
hasIPExtractor bool
|
|
|
|
extectedXRealIP string
|
|
|
|
}{
|
|
|
|
{false, false, remoteAddrIP},
|
|
|
|
{false, true, extractedRealIP},
|
|
|
|
{true, false, realIPHeaderIP},
|
|
|
|
{true, true, extractedRealIP},
|
|
|
|
}
|
|
|
|
|
|
|
|
for _, tt := range tests {
|
|
|
|
if tt.hasRealIPheader {
|
|
|
|
req.Header.Set(echo.HeaderXRealIP, realIPHeaderIP)
|
|
|
|
} else {
|
|
|
|
req.Header.Del(echo.HeaderXRealIP)
|
|
|
|
}
|
|
|
|
if tt.hasIPExtractor {
|
|
|
|
e.IPExtractor = func(*http.Request) string {
|
|
|
|
return extractedRealIP
|
|
|
|
}
|
|
|
|
} else {
|
|
|
|
e.IPExtractor = nil
|
|
|
|
}
|
|
|
|
e.ServeHTTP(rec, req)
|
|
|
|
assert.Equal(t, tt.extectedXRealIP, req.Header.Get(echo.HeaderXRealIP), "hasRealIPheader: %t / hasIPExtractor: %t", tt.hasRealIPheader, tt.hasIPExtractor)
|
|
|
|
}
|
|
|
|
}
|