traefik-forward-auth/main_test.go

package main

import (
	"fmt"
	"time"
	// "reflect"
	"io/ioutil"
	"net/http"
	"net/http/httptest"
	"net/url"
	"strings"
	"testing"

	"github.com/op/go-logging"
)

/**
 * Utilities
 */

type TokenServerHandler struct{}

func (t *TokenServerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, `{"access_token":"123456789"}`)
}

type UserServerHandler struct{}

func (t *UserServerHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	fmt.Fprint(w, `{
    "id":"1",
    "email":"example@example.com",
    "verified_email":true,
    "hd":"example.com"
  }`)
}

func init() {
	// Remove for debugging
	logging.SetLevel(logging.INFO, "traefik-forward-auth")
}

func httpRequest(r *http.Request, c *http.Cookie) (*http.Response, string) {
	w := httptest.NewRecorder()

	// Set cookies on recorder
	if c != nil {
		http.SetCookie(w, c)
	}

	// Copy into request
	for _, c := range w.HeaderMap["Set-Cookie"] {
		r.Header.Add("Cookie", c)
	}

	handler(w, r)

	res := w.Result()
	body, _ := ioutil.ReadAll(res.Body)

	return res, string(body)
}

func newHttpRequest(uri string) *http.Request {
	r := httptest.NewRequest("", "http://example.com", nil)
	r.Header.Add("X-Forwarded-Uri", uri)
	return r
}

func qsDiff(one, two url.Values) {
	for k, _ := range one {
		if two.Get(k) == "" {
			fmt.Printf("Key missing: %s\n", k)
		}
		if one.Get(k) != two.Get(k) {
			fmt.Printf("Value different for %s: expected: '%s' got: '%s'\n", k, one.Get(k), two.Get(k))
		}
	}
	for k, _ := range two {
		if one.Get(k) == "" {
			fmt.Printf("Extra key: %s\n", k)
		}
	}
}

/**
 * Tests
 */

func TestHandler(t *testing.T) {
	fw = &ForwardAuth{
		Path:         "_oauth",
		ClientId:     "idtest",
		ClientSecret: "sectest",
		Scope:        "scopetest",
		LoginURL: &url.URL{
			Scheme: "http",
			Host:   "test.com",
			Path:   "/auth",
		},
		CookieName: "cookie_test",
		Lifetime:   time.Second * time.Duration(10),
	}

	// Should redirect vanilla request to login url
	req := newHttpRequest("foo")
	res, _ := httpRequest(req, nil)
	if res.StatusCode != 307 {
		t.Error("Vanilla request should be redirected with 307, got:", res.StatusCode)
	}
	fwd, _ := res.Location()
	if fwd.Scheme != "http" || fwd.Host != "test.com" || fwd.Path != "/auth" {
		t.Error("Vanilla request should be redirected to login url, got:", fwd)
	}

	// Should catch invalid cookie
	req = newHttpRequest("foo")

	c := fw.MakeCookie(req, "test@example.com")
	parts := strings.Split(c.Value, "|")
	c.Value = fmt.Sprintf("bad|%s|%s", parts[1], parts[2])

	res, _ = httpRequest(req, c)
	if res.StatusCode != 401 {
		t.Error("Request with invalid cookie shound't be authorised", res.StatusCode)
	}

	// Should validate email
	req = newHttpRequest("foo")

	c = fw.MakeCookie(req, "test@example.com")
	fw.Domain = []string{"test.com"}

	res, _ = httpRequest(req, c)
	if res.StatusCode != 401 {
		t.Error("Request with invalid cookie shound't be authorised", res.StatusCode)
	}

	// Should allow valid request email
	req = newHttpRequest("foo")

	c = fw.MakeCookie(req, "test@example.com")
	fw.Domain = []string{}

	res, _ = httpRequest(req, c)
	if res.StatusCode != 200 {
		t.Error("Valid request should be allowed, got:", res.StatusCode)
	}

	// Should pass through user
	users := res.Header["X-Forwarded-User"]
	if len(users) != 1 {
		t.Error("Valid request missing X-Forwarded-User header")
	} else if users[0] != "test@example.com" {
		t.Error("X-Forwarded-User should match user, got: ", users)
	}
}

func TestCallback(t *testing.T) {
	fw = &ForwardAuth{
		Path:         "_oauth",
		ClientId:     "idtest",
		ClientSecret: "sectest",
		Scope:        "scopetest",
		LoginURL: &url.URL{
			Scheme: "http",
			Host:   "test.com",
			Path:   "/auth",
		},
		CSRFCookieName: "csrf_test",
	}

	// Setup token server
	tokenServerHandler := &TokenServerHandler{}
	tokenServer := httptest.NewServer(tokenServerHandler)
	defer tokenServer.Close()
	tokenUrl, _ := url.Parse(tokenServer.URL)
	fw.TokenURL = tokenUrl

	// Setup user server
	userServerHandler := &UserServerHandler{}
	userServer := httptest.NewServer(userServerHandler)
	defer userServer.Close()
	userUrl, _ := url.Parse(userServer.URL)
	fw.UserURL = userUrl

	// Should pass auth response request to callback
	req := newHttpRequest("_oauth")
	res, _ := httpRequest(req, nil)
	if res.StatusCode != 401 {
		t.Error("Auth callback without cookie shound't be authorised, got:", res.StatusCode)
	}

	// Should catch invalid csrf cookie
	req = newHttpRequest("_oauth?state=12345678901234567890123456789012:http://redirect")
	c := fw.MakeCSRFCookie(req, "nononononononononononononononono")
	res, _ = httpRequest(req, c)
	if res.StatusCode != 401 {
		t.Error("Auth callback with invalid cookie shound't be authorised, got:", res.StatusCode)
	}

	// Should redirect valid request
	req = newHttpRequest("_oauth?state=12345678901234567890123456789012:http://redirect")
	c = fw.MakeCSRFCookie(req, "12345678901234567890123456789012")
	res, _ = httpRequest(req, c)
	if res.StatusCode != 307 {
		t.Error("Valid callback should be allowed, got:", res.StatusCode)
	}
	fwd, _ := res.Location()
	if fwd.Scheme != "http" || fwd.Host != "redirect" || fwd.Path != "" {
		t.Error("Valid request should be redirected to return url, got:", fwd)
	}
}