diff --git a/api_test.go b/api_test.go index 0ebf16b5f..73792bf8d 100644 --- a/api_test.go +++ b/api_test.go @@ -1475,7 +1475,6 @@ func makeUser(t *testing.T, groups []authn.Group, name, secret string) *authn.Us if err != nil { t.Fatalf("signing string %v", err) } - validToken = "Bearer " + validToken return &authn.UserInfo{ UserID: "fake" + name, @@ -1520,7 +1519,7 @@ admin: "ac97c9e2-346b-42a2-b6da-18bcb61a32fe"` return } g := []authn.Group{} - switch token[0] { + switch strings.TrimPrefix(token[0], "Bearer ") { case adminUser.Token: g = adminUser.Groups case readUser.Token: diff --git a/internal_client.go b/internal_client.go index 1bff1e90a..cbfeb355f 100644 --- a/internal_client.go +++ b/internal_client.go @@ -144,10 +144,15 @@ func NewInternalClientFromURI(defaultURI *pnet.URI, remoteClient *http.Client, o return ic } +// AddAuthToken checks in a couple spots for our authorization token and adds it to +// the Authorization Header in the request if it finds it. func AddAuthToken(ctx context.Context, req *http.Request) *http.Request { - token, ok := ctx.Value("token").(string) - if ok && token != "" { + if token, ok := ctx.Value("token").(string); ok && token != "" { + // the "token" value should be prefixed with "Bearer" req.Header.Set("Authorization", token) + } else if uinfo := ctx.Value("userinfo"); uinfo != nil { + // UserInfo.Token is not prefixed with "Bearer" + req.Header.Set("Authorization", "Bearer "+uinfo.(*authn.UserInfo).Token) } return req } @@ -606,12 +611,6 @@ func (c *InternalClient) QueryNode(ctx context.Context, uri *pnet.URI, index str return nil, errors.Wrap(err, "creating request") } - uinfo := ctx.Value("userinfo") - if uinfo != nil { - token := uinfo.(*authn.UserInfo).Token - req.Header.Set("Authorization", token) - } - req = AddAuthToken(ctx, req) req.Header.Set("Content-Length", strconv.Itoa(len(buf))) diff --git a/internal_client_test.go b/internal_client_test.go index 204a85cae..8acb74751 100644 --- a/internal_client_test.go +++ b/internal_client_test.go @@ -7,6 +7,7 @@ import ( "context" "encoding/hex" "fmt" + "net/http" gohttp "net/http" "reflect" "strings" @@ -15,6 +16,7 @@ import ( "github.com/davecgh/go-spew/spew" pilosa "github.com/molecula/featurebase/v3" + "github.com/molecula/featurebase/v3/authn" "github.com/molecula/featurebase/v3/encoding/proto" "github.com/molecula/featurebase/v3/pql" "github.com/molecula/featurebase/v3/server" @@ -1568,3 +1570,38 @@ func TestClient_ImportRoaringExists(t *testing.T) { } } + +func TestAddAuthToken(t *testing.T) { + t.Run("none", func(t *testing.T) { + req, err := http.NewRequest("GET", "dontmatternone", strings.NewReader("this doesn't matter")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + pilosa.AddAuthToken(context.Background(), req) + if req.Header.Get("Authorization") != "" { + t.Fatalf("Authorization header set when it should be empty") + } + }) + t.Run("userinfo", func(t *testing.T) { + req, err := http.NewRequest("GET", "dontmatternone", strings.NewReader("this doesn't matter")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + uinfo := &authn.UserInfo{Token: "ayo"} + pilosa.AddAuthToken(context.WithValue(context.Background(), "userinfo", uinfo), req) + if got := req.Header.Get("Authorization"); got != "Bearer "+uinfo.Token { + t.Fatalf("got '%v', expected 'Bearer %v'", got, uinfo.Token) + } + }) + t.Run("token", func(t *testing.T) { + req, err := http.NewRequest("GET", "dontmatternone", strings.NewReader("this doesn't matter")) + if err != nil { + t.Fatalf("unexpected error: %v", err) + } + tok := "Bearer thisisatoken" + pilosa.AddAuthToken(context.WithValue(context.Background(), "token", tok), req) + if got := req.Header.Get("Authorization"); got != tok { + t.Fatalf("got '%v', expected '%v'", got, tok) + } + }) +}