Code change to support public 'SetPrincipal' method

server/auth/auth.go

@@ -8,11 +8,61 @@ import (
 	"io/ioutil"
 	"net/http"
 	"strings"
+	"sync"
 	"time"
 
+	"github.com/gorilla/context"
 	"github.com/pinterest/knox"
 )
 
+type contextKey int
+
+const (
+	principalCtxKey contextKey = iota
+)
+
+type PrincipalContext interface {
+	SetCurrentPrincipal(principal knox.Principal)
+	GetCurrentPrincipal() knox.Principal
+}
+
+type principalContext struct {
+	principalCtxKey contextKey
+	request         *http.Request
+	once            sync.Once
+	invocationCount int
+}
+
+func NewPrincipalContext(request *http.Request) PrincipalContext {
+	return &principalContext{
+		principalCtxKey,
+		request,
+		sync.Once{},
+		0,
+	}
+}
+
+func (ctx *principalContext) GetCurrentPrincipal() knox.Principal {
+	if rv := context.Get(ctx.request, principalCtxKey); rv != nil {
+		return rv.(knox.Principal)
+	}
+	return nil
+}
+
+func (ctx *principalContext) SetCurrentPrincipal(principal knox.Principal) {
+	if ctx.invocationCount > 0 {
+		panic("SetPrincipal was called more than once during the lifetime of the HTTP request")
+	}
+	ctx.setPrincipalInner(ctx.request, principal)
+}
+
+func (ctx *principalContext) setPrincipalInner(httpRequest *http.Request, principal knox.Principal) {
+	ctx.once.Do(func() {
+		context.Set(httpRequest, ctx.principalCtxKey, principal)
+		ctx.invocationCount += 1
+	})
+}
+
 // Provider is used for authenticating requests via the authentication decorator.
 type Provider interface {
 	Name() string

server/auth/auth_test.go

@@ -5,6 +5,7 @@ import (
 	"crypto/x509"
 	"encoding/base64"
 	"net/http"
+	"reflect"
 	"strings"
 	"time"
 
@@ -13,6 +14,53 @@ import (
 	"github.com/pinterest/knox"
 )
 
+func TestPrincipalContext(t *testing.T) {
+	originalPrincipal := NewUser("test", []string{"returntrue"}).(user)
+	newPrincipal := NewUser("hacker", []string{"returntrue"}).(user)
+
+	req, err := http.NewRequest("GET", "http://localhost/", nil)
+	if err != nil {
+		t.Fatal(err.Error())
+	}
+
+	ctx := NewPrincipalContext(req)
+
+	currentPrincipal := ctx.GetCurrentPrincipal()
+	if currentPrincipal != nil {
+		t.Error("Current principal was expected to be null")
+	}
+
+	ctx.SetCurrentPrincipal(originalPrincipal)
+	currentPrincipal = ctx.GetCurrentPrincipal().(user)
+
+	if !reflect.DeepEqual(currentPrincipal, originalPrincipal) {
+		t.Errorf(
+			"Current principal was expected to be user: '%s'. Instead got: '%s'",
+			originalPrincipal.GetID(),
+			currentPrincipal.GetID(),
+		)
+	}
+
+	panicRecovery := func() {
+		recoveryValue := recover()
+		if recoveryValue == nil {
+			t.Error("Expected 'SetCurrentPrincipal' to panic on second call")
+		}
+
+		currentPrincipal = ctx.GetCurrentPrincipal()
+		if !reflect.DeepEqual(currentPrincipal, originalPrincipal) {
+			t.Errorf(
+				"Current principal was expected to be user: '%s'. Instead got: '%s'",
+				originalPrincipal.GetID(),
+				currentPrincipal.GetID(),
+			)
+		}
+	}
+	defer panicRecovery()
+
+	ctx.SetCurrentPrincipal(newPrincipal)
+}
+
 func TestUserCanAccess(t *testing.T) {
 	u := NewUser("test", []string{"returntrue"})
 	a1 := knox.Access{ID: "test", AccessType: knox.Write, Type: knox.User}

server/decorators.go

@@ -35,14 +35,16 @@ func setAPIError(r *http.Request, val *HTTPError) {
 
 // GetPrincipal gets the principal authenticated through the authentication decorator
 func GetPrincipal(r *http.Request) knox.Principal {
-	if rv := context.Get(r, principalContext); rv != nil {
-		return rv.(knox.Principal)
-	}
-	return nil
+	ctx := getOrInitializePrincipalContext(r)
+	return ctx.GetCurrentPrincipal()
 }
 
-func setPrincipal(r *http.Request, val knox.Principal) {
-	context.Set(r, principalContext, val)
+// SetPrincipal sets the principal authenticated through the authentication decorator.
+// For security reasons, this method will only set the Principal in the context for
+// the first invocation. Subsequent invocations WILL cause a panic.
+func SetPrincipal(r *http.Request, val knox.Principal) {
+	ctx := getOrInitializePrincipalContext(r)
+	ctx.SetCurrentPrincipal(val)
 }
 
 // GetParams gets the parameters for the request through the parameters context.
@@ -68,6 +70,15 @@ func setDB(r *http.Request, val KeyManager) {
 	context.Set(r, dbContext, val)
 }
 
+func getOrInitializePrincipalContext(r *http.Request) auth.PrincipalContext {
+	if ctx := context.Get(r, principalContext); ctx != nil {
+		return ctx.(auth.PrincipalContext)
+	}
+	ctx := auth.NewPrincipalContext(r)
+	context.Set(r, principalContext, ctx)
+	return ctx
+}
+
 // GetRouteID gets the short form function name for the route being called. Used for logging/metrics.
 func GetRouteID(r *http.Request) string {
 	if rv := context.Get(r, idContext); rv != nil {
@@ -223,7 +234,7 @@ func Authentication(providers []auth.Provider) func(http.HandlerFunc) http.Handl
 				return
 			}
 
-			setPrincipal(r, knox.NewPrincipalMux(defaultPrincipal, allPrincipals))
+			SetPrincipal(r, knox.NewPrincipalMux(defaultPrincipal, allPrincipals))
 			f(w, r)
 			return
 		}