From ab2b48da0d832efb9b22f920de6612d4400a12c3 Mon Sep 17 00:00:00 2001 From: Garrison Davis Date: Mon, 24 Oct 2022 13:32:43 -0600 Subject: [PATCH] Stop using string keys in contexts This fixes the OriginalIP and RequestUserID in the main featurebase package, and the Access and Refresh tokens, the UserInfo, and the []string of Indexes passed with context.Context(s) in the authn package. An empty struct was used for all of these keys (and relevant helper functions we added) to avoid allocations where possible while still using the context functionality. Some of the logic in the server.GetIndexes function was fixed. (cherry picked from commit 0f5a56c9585fcec819b11a9be954f9fe1395f4de) --- api.go | 12 ++--- api_test.go | 18 ++----- authn/context.go | 54 ++++++++++++++++++++ context.go | 28 +++++++++++ ctl/backup.go | 6 +-- ctl/backup_tar.go | 6 +-- ctl/import.go | 6 +-- ctl/import_test.go | 6 +-- ctl/restore.go | 6 +-- ctl/restore_tar.go | 2 +- http_handler.go | 32 +++++------- internal/clustertests/cluster_test.go | 12 +---- internal/clustertests/pause_node_test.go | 6 +-- internal_client.go | 12 ++--- internal_client_test.go | 9 ++-- server/grpc.go | 64 +++++++++--------------- server/grpc_test.go | 56 +++++++++++---------- server/handler_test.go | 1 - sql/show.go | 4 +- 19 files changed, 178 insertions(+), 162 deletions(-) create mode 100644 authn/context.go create mode 100644 context.go diff --git a/api.go b/api.go index 36de49b2b..1d7be7448 100644 --- a/api.go +++ b/api.go @@ -218,10 +218,7 @@ func (api *API) CreateIndex(ctx context.Context, indexName string, options Index // get the requestUserID from the context -- assumes the http handler has populated this from // authN/Z info - requestUserID, ok := ctx.Value(ContextRequestUserIdKey).(string) - if !ok { - requestUserID = "" - } + requestUserID, _ := UserIDFromContext(ctx) // requestUserID is "" if not in ctx if err := api.validate(apiCreateIndex); err != nil { return nil, errors.Wrap(err, "validating api method") @@ -310,10 +307,7 @@ func (api *API) CreateField(ctx context.Context, indexName string, fieldName str // get the requestUserID from the context -- assumes the http handler has populated this from // authN/Z info - requestUserID, ok := ctx.Value(ContextRequestUserIdKey).(string) - if !ok { - requestUserID = "" - } + requestUserID, _ := UserIDFromContext(ctx) // requestUserID is "" if not in ctx // Apply and validate functional options. fo, err := newFieldOptions(opts...) @@ -369,7 +363,7 @@ func (api *API) UpdateField(ctx context.Context, indexName, fieldName string, up // get the requestUserID from the context -- assumes the http handler has populated this from // authN/Z info - requestUserID, _ := ctx.Value(ContextRequestUserIdKey).(string) + requestUserID, _ := UserIDFromContext(ctx) cfm, err := index.UpdateField(ctx, fieldName, requestUserID, update) if err != nil { diff --git a/api_test.go b/api_test.go index 19b2b287b..8e505351b 100644 --- a/api_test.go +++ b/api_test.go @@ -1403,23 +1403,11 @@ func TestAuth_MultiNode(t *testing.T) { "test": "write" admin: "ac97c9e2-346b-42a2-b6da-18bcb61a32fe"` adminUser := makeUser(t, []authn.Group{{GroupID: "ac97c9e2-346b-42a2-b6da-18bcb61a32fe", GroupName: "adminGroup"}}, "admin", "DEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEF") - adminCtx := context.WithValue( - context.Background(), - "userinfo", - adminUser, - ) + adminCtx := authn.WithUserInfo(context.Background(), adminUser) readUser := makeUser(t, []authn.Group{{GroupID: "dca35310-ecda-4f23-86cd-876aee55906b", GroupName: "readGroup"}}, "reader", "DEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEF") - readCtx := context.WithValue( - context.Background(), - "userinfo", - readUser, - ) + readCtx := authn.WithUserInfo(context.Background(), readUser) writeUser := makeUser(t, []authn.Group{{GroupID: "dca35310-ecda-4f23-86cd-876aee55906f", GroupName: "writeGroup"}}, "writer", "DEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEEFDEADBEED") - writeCtx := context.WithValue( - context.Background(), - "userinfo", - writeUser, - ) + writeCtx := authn.WithUserInfo(context.Background(), writeUser) srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) { token, ok := r.Header["Authorization"] if !ok || len(token) == 0 { diff --git a/authn/context.go b/authn/context.go new file mode 100644 index 000000000..e84310003 --- /dev/null +++ b/authn/context.go @@ -0,0 +1,54 @@ +// Copyright 2022 Molecula Corp (DBA FeatureBase). All rights reserved. +package authn + +import "context" + +// Empty struct to avoid allocations +type contextKeyAccessToken struct{} +type contextKeyRefreshToken struct{} +type contextKeyUserInfo struct{} +type contextKeyIndexes struct{} + +// GetAccessToken gets the access token from a context. +func GetAccessToken(ctx context.Context) (token string, ok bool) { + token, ok = ctx.Value(contextKeyAccessToken{}).(string) + return +} + +// WithAccessToken makes a new Context with an access token. +func WithAccessToken(ctx context.Context, token string) context.Context { + return context.WithValue(ctx, contextKeyAccessToken{}, token) +} + +// GetRefreshToken gets the refresh token from a context. +func GetRefreshToken(ctx context.Context) (token string, ok bool) { + token, ok = ctx.Value(contextKeyRefreshToken{}).(string) + return +} + +// WithRefreshToken makes a new Context with a refresh token. +func WithRefreshToken(ctx context.Context, token string) context.Context { + return context.WithValue(ctx, contextKeyRefreshToken{}, token) +} + +// GetUserInfo gets the UserInfo from a context. +func GetUserInfo(ctx context.Context) (userInfo *UserInfo, ok bool) { + userInfo, ok = ctx.Value(contextKeyUserInfo{}).(*UserInfo) + return +} + +// WithUserInfo makes a new Context with UserInfo. +func WithUserInfo(ctx context.Context, userInfo *UserInfo) context.Context { + return context.WithValue(ctx, contextKeyUserInfo{}, userInfo) +} + +// GetIndexes get the indices from a context. +func GetIndexes(ctx context.Context) (indexes []string, ok bool) { + indexes, ok = ctx.Value(contextKeyIndexes{}).([]string) + return +} + +// WithIndexes makes a new Context with a []string containing the indicies. +func WithIndexes(ctx context.Context, indexes []string) context.Context { + return context.WithValue(ctx, contextKeyUserInfo{}, indexes) +} diff --git a/context.go b/context.go new file mode 100644 index 000000000..e016366ba --- /dev/null +++ b/context.go @@ -0,0 +1,28 @@ +// Copyright 2022 Molecula Corp (DBA FeatureBase). All rights reserved. +package pilosa + +import "context" + +// Empty struct to avoid allocations +type contextKeyOriginalIP struct{} +type contextKeyRequestUserID struct{} + +// OriginalIPFromContext gets the original IP from the context. +func OriginalIPFromContext(ctx context.Context) (originalIP string, ok bool) { + originalIP, ok = ctx.Value(contextKeyOriginalIP{}).(string) + return +} + +// WithOriginalIP makes a new context with the originalIP in the context. +func WithOriginalIP(ctx context.Context, originalIP string) context.Context { + return context.WithValue(ctx, contextKeyOriginalIP{}, originalIP) +} + +func UserIDFromContext(ctx context.Context) (userID string, ok bool) { + userID, ok = ctx.Value(contextKeyRequestUserID{}).(string) + return +} + +func WithUserID(ctx context.Context, userID string) context.Context { + return context.WithValue(ctx, contextKeyRequestUserID{}, userID) +} diff --git a/ctl/backup.go b/ctl/backup.go index 97295249f..ad2f74e24 100644 --- a/ctl/backup.go +++ b/ctl/backup.go @@ -113,11 +113,7 @@ func (cmd *BackupCommand) Run(ctx context.Context) (err error) { cmd.client = client if cmd.AuthToken != "" { - ctx = context.WithValue( - ctx, - authn.ContextValueAccessToken, - "Bearer "+cmd.AuthToken, - ) + ctx = authn.WithAccessToken(ctx, "Bearer "+cmd.AuthToken) } // Determine the field type in order to correctly handle the input data. diff --git a/ctl/backup_tar.go b/ctl/backup_tar.go index 7420f4324..fad3c1a13 100644 --- a/ctl/backup_tar.go +++ b/ctl/backup_tar.go @@ -102,11 +102,7 @@ func (cmd *BackupTarCommand) Run(ctx context.Context) (err error) { cmd.client = client if cmd.AuthToken != "" { - ctx = context.WithValue( - ctx, - authn.ContextValueAccessToken, - "Bearer "+cmd.AuthToken, - ) + ctx = authn.WithAccessToken(ctx, "Bearer "+cmd.AuthToken) } // Determine the field type in order to correctly handle the input data. diff --git a/ctl/import.go b/ctl/import.go index 1bf4c9dce..a3b2923c8 100644 --- a/ctl/import.go +++ b/ctl/import.go @@ -94,11 +94,7 @@ func (cmd *ImportCommand) Run(ctx context.Context) error { cmd.client = client if cmd.AuthToken != "" { - ctx = context.WithValue( - ctx, - authn.ContextValueAccessToken, - "Bearer "+cmd.AuthToken, - ) + ctx = authn.WithAccessToken(ctx, "Bearer "+cmd.AuthToken) } if cmd.CreateSchema { diff --git a/ctl/import_test.go b/ctl/import_test.go index f1c660ed2..dc492f0a2 100644 --- a/ctl/import_test.go +++ b/ctl/import_test.go @@ -722,11 +722,7 @@ func TestImport_AuthOn(t *testing.T) { cm.Field = test.Field cm.CreateSchema = test.CreateSchema cm.Paths = []string{file.Name()} - ctx := context.WithValue( - context.Background(), - authn.ContextValueAccessToken, - test.Token, - ) + ctx := authn.WithAccessToken(context.Background(), test.Token) err = cm.Run(ctx) if test.Err != nil { if !strings.Contains(err.Error(), test.Err.Error()) { diff --git a/ctl/restore.go b/ctl/restore.go index 5af812730..3c118f7f6 100644 --- a/ctl/restore.go +++ b/ctl/restore.go @@ -96,7 +96,7 @@ func (cmd *RestoreCommand) Run(ctx context.Context) (err error) { cmd.client = client if cmd.AuthToken != "" { - ctx = context.WithValue(ctx, authn.ContextValueAccessToken, "Bearer "+cmd.AuthToken) + ctx = authn.WithAccessToken(ctx, "Bearer "+cmd.AuthToken) } nodes, err := cmd.client.Nodes(ctx) @@ -157,7 +157,7 @@ func (cmd *RestoreCommand) restoreSchema(ctx context.Context, primary *disco.Nod req = req.WithContext(ctx) req.Header.Add("Accept", "application/json") - token, ok := ctx.Value(authn.ContextValueAccessToken).(string) + token, ok := authn.GetAccessToken(ctx) if ok && token != "" { req.Header.Set("Authorization", token) } @@ -327,7 +327,7 @@ func (cmd *RestoreCommand) restoreShard(ctx context.Context, filename string) er req = req.WithContext(ctx) req.Header.Set("Content-Type", "application/octet-stream") - token, ok := ctx.Value(authn.ContextValueAccessToken).(string) + token, ok := authn.GetAccessToken(ctx) if ok && token != "" { req.Header.Set("Authorization", token) } diff --git a/ctl/restore_tar.go b/ctl/restore_tar.go index e9ca669ae..3f60a1154 100644 --- a/ctl/restore_tar.go +++ b/ctl/restore_tar.go @@ -98,7 +98,7 @@ func (cmd *RestoreTarCommand) Run(ctx context.Context) (err error) { cmd.client = client if cmd.AuthToken != "" { - ctx = context.WithValue(ctx, authn.ContextValueAccessToken, "Bearer "+cmd.AuthToken) + ctx = authn.WithAccessToken(ctx, "Bearer "+cmd.AuthToken) } var tarReader *tar.Reader diff --git a/http_handler.go b/http_handler.go index bab12da32..65005a8ca 100644 --- a/http_handler.go +++ b/http_handler.go @@ -49,13 +49,6 @@ import ( "github.com/zeebo/blake3" ) -type ContextRequestUserIdKeyType string - -const ( - // ContextRequestUserIdKey is request userid key for a request ctx - ContextRequestUserIdKey = ContextRequestUserIdKeyType("request-user-id") -) - const ( // HeaderRequestUserID is request userid header HeaderRequestUserID = "X-Request-Userid" @@ -692,7 +685,7 @@ func (h *Handler) chkAllowedNetworks(r *http.Request) (bool, context.Context) { // if client IP is in allowed networks // add it to the context for key X-Molecula-Original-IP if h.auth.CheckAllowedNetworks(reqIP) { - ctx := context.WithValue(r.Context(), OriginalIPHeader, reqIP) + ctx := WithOriginalIP(r.Context(), reqIP) return true, ctx } return false, r.Context() @@ -704,7 +697,7 @@ func (h *Handler) chkAuthN(handler http.HandlerFunc) http.HandlerFunc { //if the request is unauthenticated and we have the appropriate header get the userid from the header requestUserID := r.Header.Get(HeaderRequestUserID) - ctx = context.WithValue(ctx, ContextRequestUserIdKey, requestUserID) + ctx = WithUserID(ctx, requestUserID) if h.auth == nil { handler.ServeHTTP(w, r.WithContext(ctx)) @@ -724,12 +717,13 @@ func (h *Handler) chkAuthN(handler http.HandlerFunc) http.HandlerFunc { http.Error(w, errors.Wrap(err, "authenticating").Error(), http.StatusUnauthorized) return } + // prefer the user id from an authenticated request over one in a header - ctx = context.WithValue(ctx, ContextRequestUserIdKey, uinfo.UserID) + ctx = WithUserID(ctx, uinfo.UserID) // just in case it got refreshed - ctx = context.WithValue(ctx, authn.ContextValueAccessToken, "Bearer "+access) - ctx = context.WithValue(ctx, authn.ContextValueRefreshToken, refresh) + ctx = authn.WithAccessToken(ctx, "Bearer"+access) + ctx = authn.WithRefreshToken(ctx, refresh) h.auth.SetCookie(w, uinfo.Token, uinfo.RefreshToken, uinfo.Expiry) handler.ServeHTTP(w, r.WithContext(ctx)) @@ -742,7 +736,7 @@ func (h *Handler) chkAuthZ(handler http.HandlerFunc, perm authz.Permission) http //if the request is unauthenticated and we have the appropriate header get the userid from the header requestUserID := r.Header.Get(HeaderRequestUserID) - ctx = context.WithValue(ctx, ContextRequestUserIdKey, requestUserID) + ctx = WithUserID(ctx, requestUserID) // handle the case when auth is not turned on if h.auth == nil { @@ -771,18 +765,18 @@ func (h *Handler) chkAuthZ(handler http.HandlerFunc, perm authz.Permission) http } // prefer the user id from an authenticated request over one in a header - ctx = context.WithValue(ctx, ContextRequestUserIdKey, uinfo.UserID) + ctx = WithUserID(ctx, uinfo.UserID) - ctx = context.WithValue(ctx, authn.ContextValueAccessToken, "Bearer "+access) - ctx = context.WithValue(ctx, authn.ContextValueRefreshToken, refresh) + ctx = authn.WithAccessToken(ctx, "Bearer "+access) + ctx = authn.WithRefreshToken(ctx, refresh) // just in case it got refreshed h.auth.SetCookie(w, uinfo.Token, uinfo.RefreshToken, uinfo.Expiry) // put the user's authN/Z info in the context ctx = context.WithValue(ctx, contextKeyGroupMembership, uinfo.Groups) - ctx = context.WithValue(ctx, authn.ContextValueAccessToken, "Bearer "+uinfo.Token) - ctx = context.WithValue(ctx, authn.ContextValueRefreshToken, uinfo.RefreshToken) + ctx = authn.WithAccessToken(ctx, "Bearer "+uinfo.Token) + ctx = authn.WithRefreshToken(ctx, uinfo.RefreshToken) // unlikely h.permissions will be nil, but we'll check to be safe if h.permissions == nil { h.logger.Errorf("authentication is turned on without authorization permissions set") @@ -869,7 +863,7 @@ func GetIP(r *http.Request) string { } // check if original IP is in the context - if ogIP, ok := r.Context().Value(OriginalIPHeader).(string); ok && ogIP != "" { + if ogIP, ok := OriginalIPFromContext(r.Context()); ok && ogIP != "" { return ogIP } diff --git a/internal/clustertests/cluster_test.go b/internal/clustertests/cluster_test.go index dfbb9b226..56d5db3b8 100644 --- a/internal/clustertests/cluster_test.go +++ b/internal/clustertests/cluster_test.go @@ -103,11 +103,7 @@ func TestClusterStuff(t *testing.T) { // generate auth token and add to context if auth { token = GetAuthToken(t) - ctx = context.WithValue( - ctx, - authn.ContextValueAccessToken, - "Bearer "+token, - ) + ctx = authn.WithAccessToken(ctx, "Bearer"+token) } if err := cli[0].CreateIndex(ctx, "testidx", pilosa.IndexOptions{}); err != nil { @@ -330,11 +326,7 @@ func TestRetryLogic(t *testing.T) { } if auth { token := GetAuthToken(t) - ctx = context.WithValue( - ctx, - authn.ContextValueAccessToken, - "Bearer "+token, - ) + ctx = authn.WithAccessToken(ctx, "Bearer "+token) } var addrs = []string{"pilosa1:10101", "pilosa2:10101", "pilosa3:10101"} diff --git a/internal/clustertests/pause_node_test.go b/internal/clustertests/pause_node_test.go index 8aebe5a26..9f24525fd 100644 --- a/internal/clustertests/pause_node_test.go +++ b/internal/clustertests/pause_node_test.go @@ -302,11 +302,7 @@ func TestPauseReplica(t *testing.T) { ctx := context.Background() if auth { token := GetAuthToken(t) - ctx = context.WithValue( - ctx, - authn.ContextValueAccessToken, - "Bearer "+token, - ) + ctx = authn.WithAccessToken(ctx, "Bearer "+token) } ctx, cancel := context.WithCancel(ctx) diff --git a/internal_client.go b/internal_client.go index 6dfa599f8..c4d40131e 100644 --- a/internal_client.go +++ b/internal_client.go @@ -228,29 +228,29 @@ func NewInternalClientFromURI(defaultURI *pnet.URI, remoteClient *http.Client, o // same for refresh tokens as well. func AddAuthToken(ctx context.Context, header *http.Header) { var access, refresh string - if token, ok := ctx.Value(authn.ContextValueAccessToken).(string); ok { + if token, ok := authn.GetAccessToken(ctx); ok { // the AccessToken value should be prefixed with "Bearer" access = token } - if token, ok := ctx.Value(authn.ContextValueRefreshToken).(string); ok { + if token, ok := authn.GetRefreshToken(ctx); ok { refresh = token } // not combining these ifs so we don't call ctx.Value unless we have to if access == "" || refresh == "" { - if uinfo := ctx.Value("userinfo"); uinfo != nil { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { if access == "" { // UserInfo.Token is not prefixed with "Bearer" - access = "Bearer " + uinfo.(*authn.UserInfo).Token + access = "Bearer " + uinfo.Token } if refresh == "" { - refresh = uinfo.(*authn.UserInfo).RefreshToken + refresh = uinfo.RefreshToken } } } // set ogIP to request for remote calls - if ogIP, ok := ctx.Value(OriginalIPHeader).(string); ok && ogIP != "" { + if ogIP, ok := OriginalIPFromContext(ctx); ok && ogIP != "" { header.Set(OriginalIPHeader, ogIP) } diff --git a/internal_client_test.go b/internal_client_test.go index fd7bc3d8f..fb6ad7fc3 100644 --- a/internal_client_test.go +++ b/internal_client_test.go @@ -1608,7 +1608,7 @@ func TestAddAuthToken(t *testing.T) { t.Fatalf("unexpected error: %v", err) } uinfo := &authn.UserInfo{Token: "ayo"} - pilosa.AddAuthToken(context.WithValue(context.Background(), "userinfo", uinfo), &req.Header) + pilosa.AddAuthToken(authn.WithUserInfo(context.Background(), uinfo), &req.Header) if got := req.Header.Get("Authorization"); got != "Bearer "+uinfo.Token { t.Fatalf("got '%v', expected 'Bearer %v'", got, uinfo.Token) } @@ -1620,10 +1620,7 @@ func TestAddAuthToken(t *testing.T) { } tok := "Bearer thisisatoken" pilosa.AddAuthToken( - context.WithValue(context.Background(), - authn.ContextValueAccessToken, - tok, - ), + authn.WithAccessToken(context.Background(), tok), &req.Header, ) if got := req.Header.Get("Authorization"); got != tok { @@ -1636,7 +1633,7 @@ func TestAddAuthToken(t *testing.T) { t.Fatalf("unexpected error: %v", err) } ogIP := "10.0.0.1" - pilosa.AddAuthToken(context.WithValue(context.Background(), pilosa.OriginalIPHeader, ogIP), &req.Header) + pilosa.AddAuthToken(pilosa.WithOriginalIP(context.Background(), ogIP), &req.Header) if got := req.Header.Get(pilosa.OriginalIPHeader); got != ogIP { t.Fatalf("got '%v', expected '%v'", got, ogIP) } diff --git a/server/grpc.go b/server/grpc.go index 350ac8c73..07de3c126 100644 --- a/server/grpc.go +++ b/server/grpc.go @@ -166,8 +166,7 @@ func isAllowed(requested []string, allowed []string) bool { // QuerySQL handles the SQL request and sends RowResponses to the stream. func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQLServer) error { ctx := stream.Context() - uinfo, ok := ctx.Value("userinfo").(*authn.UserInfo) - if ok && uinfo != nil { + if uinfo, ok := authn.GetUserInfo(ctx); ok && uinfo != nil { // authz m := sql.NewMapper() parsed, err := m.MapSQL(req.Sql) @@ -186,7 +185,7 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ if !isAllowed(parsed.Tables, allowed) { return status.Error(codes.PermissionDenied, "insufficient permissions to access requested tables") } - ctx = context.WithValue(ctx, "indices", allowed) + ctx = authn.WithIndexes(ctx, allowed) } LogQuery(ctx, "QuerySQL", req, h.queryLogger) } @@ -230,8 +229,7 @@ func (h *GRPCHandler) QuerySQL(req *pb.QuerySQLRequest, stream pb.Pilosa_QuerySQ // https://github.com/molecula/pilosa/pull/644 func (h *GRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest) (*pb.TableResponse, error) { start := time.Now() - uinfo := ctx.Value("userinfo") - if uinfo != nil { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { // authz m := sql.NewMapper() parsed, err := m.MapSQL(req.Sql) @@ -245,12 +243,12 @@ func (h *GRPCHandler) QuerySQLUnary(ctx context.Context, req *pb.QuerySQLRequest perm = authz.Admin } - allowed := h.perms.GetAuthorizedIndexList(uinfo.(*authn.UserInfo).Groups, perm) - if !h.perms.IsAdmin(uinfo.(*authn.UserInfo).Groups) { + allowed := h.perms.GetAuthorizedIndexList(uinfo.Groups, perm) + if !h.perms.IsAdmin(uinfo.Groups) { if !isAllowed(parsed.Tables, allowed) { return nil, status.Error(codes.PermissionDenied, "insufficient permissions to access requested tables") } - ctx = context.WithValue(ctx, "indices", allowed) + ctx = authn.WithIndexes(ctx, allowed) } } @@ -289,8 +287,7 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ } ctx := stream.Context() - uinfo := ctx.Value("userinfo") - if uinfo != nil { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { lperm := authz.Read q, err := pql.ParseString(req.Pql) if err != nil { @@ -299,8 +296,8 @@ func (h *GRPCHandler) QueryPQL(req *pb.QueryPQLRequest, stream pb.Pilosa_QueryPQ if q.WriteCallN() > 0 { lperm = authz.Write } - if !h.perms.IsAdmin(uinfo.(*authn.UserInfo).Groups) { - if !isAllowed([]string{req.Index}, h.perms.GetAuthorizedIndexList(uinfo.(*authn.UserInfo).Groups, lperm)) { + if !h.perms.IsAdmin(uinfo.Groups) { + if !isAllowed([]string{req.Index}, h.perms.GetAuthorizedIndexList(uinfo.Groups, lperm)) { return status.Error(codes.PermissionDenied, "insufficient permissions to access requested indexes") } } @@ -358,8 +355,7 @@ func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest Index: req.Index, Query: req.Pql, } - uinfo := ctx.Value("userinfo") - if uinfo != nil { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { lperm := authz.Read q, err := pql.ParseString(req.Pql) if err != nil { @@ -368,8 +364,8 @@ func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest if q.WriteCallN() > 0 { lperm = authz.Write } - if !h.perms.IsAdmin(uinfo.(*authn.UserInfo).Groups) { - if !isAllowed([]string{req.Index}, h.perms.GetAuthorizedIndexList(uinfo.(*authn.UserInfo).Groups, lperm)) { + if !h.perms.IsAdmin(uinfo.Groups) { + if !isAllowed([]string{req.Index}, h.perms.GetAuthorizedIndexList(uinfo.Groups, lperm)) { return nil, status.Error(codes.PermissionDenied, fmt.Sprintf("insufficient permissions for %v", req.Index)) } } @@ -419,9 +415,8 @@ func (h *GRPCHandler) QueryPQLUnary(ctx context.Context, req *pb.QueryPQLRequest // CreateIndex creates a new Index func (h *GRPCHandler) CreateIndex(ctx context.Context, req *pb.CreateIndexRequest) (*pb.CreateIndexResponse, error) { - uinfo := ctx.Value("userinfo") - if uinfo != nil { - if !h.perms.IsAdmin(uinfo.(*authn.UserInfo).Groups) { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { + if !h.perms.IsAdmin(uinfo.Groups) { return nil, status.Error(codes.PermissionDenied, "must be admin to create index") } } @@ -439,13 +434,11 @@ func (h *GRPCHandler) CreateIndex(ctx context.Context, req *pb.CreateIndexReques // GetIndex returns a single Index given a name func (h *GRPCHandler) GetIndex(ctx context.Context, req *pb.GetIndexRequest) (*pb.GetIndexResponse, error) { - uinfo := ctx.Value("userinfo") - if uinfo != nil { - pp, ok := uinfo.(*authn.UserInfo) + if uinfo, ok := authn.GetUserInfo(ctx); uinfo != nil { if !ok { return nil, status.Error(codes.InvalidArgument, "malformed auth header") } - p, err := h.perms.GetPermissions(pp, req.Name) + p, err := h.perms.GetPermissions(uinfo, req.Name) if err != nil { return nil, err } @@ -471,14 +464,9 @@ func (h *GRPCHandler) GetIndex(ctx context.Context, req *pb.GetIndexRequest) (*p // GetIndexes returns a list of all Indexes func (h *GRPCHandler) GetIndexes(ctx context.Context, req *pb.GetIndexesRequest) (*pb.GetIndexesResponse, error) { - uinfo := ctx.Value("userinfo") - var userInfo *authn.UserInfo - if uinfo != nil { - var ok bool - userInfo, ok = uinfo.(*authn.UserInfo) - if !ok { - return nil, status.Error(codes.InvalidArgument, "malformed auth header") - } + userInfo, ok := authn.GetUserInfo(ctx) + if !ok { + return nil, status.Error(codes.InvalidArgument, "malformed auth header") } schema, err := h.api.Schema(ctx, false) if err != nil { @@ -505,9 +493,8 @@ func (h *GRPCHandler) GetIndexes(ctx context.Context, req *pb.GetIndexesRequest) // DeleteIndex deletes an Index func (h *GRPCHandler) DeleteIndex(ctx context.Context, req *pb.DeleteIndexRequest) (*pb.DeleteIndexResponse, error) { - uinfo := ctx.Value("userinfo") - if uinfo != nil { - if !h.perms.IsAdmin(uinfo.(*authn.UserInfo).Groups) { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { + if !h.perms.IsAdmin(uinfo.Groups) { return nil, status.Error(codes.PermissionDenied, "must be admin to delete index") } } @@ -742,8 +729,7 @@ func (h *GRPCHandler) Inspect(req *pb.InspectRequest, stream pb.Pilosa_InspectSe }) ctx := stream.Context() - uinfo := ctx.Value("userinfo") - if uinfo != nil { + if uinfo, _ := authn.GetUserInfo(ctx); uinfo != nil { LogQuery(stream.Context(), "Inspect", req, h.queryLogger) } @@ -1628,7 +1614,7 @@ func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) { // reset the molecula-chip cookie just in case the token was refreshed md, ok := metadata.FromIncomingContext(ctx) - if uinfo, yeah := ctx.Value("userinfo").(*authn.UserInfo); ok && yeah { + if uinfo, yeah := authn.GetUserInfo(ctx); ok && yeah { server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken) } return handler(ctx, req) @@ -1640,7 +1626,7 @@ func NewGRPCServer(opts ...grpcServerOption) (*grpcServer, error) { } // reset the molecula-chip cookie just in case the token was refreshed md, ok := metadata.FromIncomingContext(ctx) - if uinfo, yeah := ctx.Value("userinfo").(*authn.UserInfo); ok && yeah { + if uinfo, yeah := authn.GetUserInfo(ctx); ok && yeah { server.auth.SetGRPCMetadata(ctx, md, uinfo.Token, uinfo.RefreshToken) } return handler(srv, &wrappedStream{ss, ctx}) @@ -1729,7 +1715,7 @@ func Valid(ctx context.Context, auth *authn.Auth) (context.Context, error) { return ctx, status.Errorf(codes.Unauthenticated, err.Error()) } - return context.WithValue(ctx, "userinfo", uinfo), nil + return authn.WithUserInfo(ctx, uinfo), nil } func getTokensFromMetadata(md metadata.MD) (string, string) { diff --git a/server/grpc_test.go b/server/grpc_test.go index eab8f0998..bc6e2a0b0 100644 --- a/server/grpc_test.go +++ b/server/grpc_test.go @@ -1184,23 +1184,11 @@ admin: "ac97c9e2-346b-42a2-b6da-18bcb61a32fe"` } user := makeUser([]authn.Group{{GroupID: "ac97c9e2-346b-42a2-b6da-18bcb61a32fe", GroupName: "adminGroup"}}, "admin") - adminCtx := context.WithValue( - ctx, - "userinfo", - user, - ) + adminCtx := authn.WithUserInfo(ctx, user) readuser := makeUser([]authn.Group{{GroupID: "dca35310-ecda-4f23-86cd-876aee55906b", GroupName: "readers"}}, "reader") - readCtx := context.WithValue( - ctx, - "userinfo", - readuser, - ) + readCtx := authn.WithUserInfo(ctx, readuser) writeuser := makeUser([]authn.Group{{GroupID: "dca35310-ecda-4f23-86cd-876aee55906f", GroupName: "writers"}}, "admin") - writeCtx := context.WithValue( - ctx, - "userinfo", - writeuser, - ) + writeCtx := authn.WithUserInfo(ctx, writeuser) sql := "select * from grouper" t.Run("test-auth-with-admin-sqlUnary", func(t *testing.T) { @@ -1537,11 +1525,11 @@ func TestCRUDIndexes(t *testing.T) { func TestLogQuery(t *testing.T) { method := "test!" - uinfo := authn.UserInfo{ + uinfo := &authn.UserInfo{ UserID: "ID", UserName: "name", } - ctx := context.WithValue(context.Background(), "userinfo", &uinfo) + ctx := authn.WithUserInfo(context.Background(), uinfo) cases := []struct { name string @@ -1963,8 +1951,8 @@ type MockStream struct { context context.Context } -func (ms MockStream) SetHeader(md metadata.MD) error { - ms.context = context.WithValue(context.Background(), "metadata", md) +func (ms *MockStream) SetHeader(md metadata.MD) error { + ms.context = contextWithMetadata(context.Background(), md) return nil } @@ -1977,7 +1965,7 @@ func (ms MockStream) SetTrailer(metadata.MD) {} func (ms MockStream) Context() context.Context { if ms.context == nil { md := metadata.New(map[string]string{}) - ms.context = context.WithValue(context.Background(), "metadata", md) + ms.context = contextWithMetadata(context.Background(), md) } return ms.context } @@ -1990,21 +1978,34 @@ func (ms MockStream) RecvMsg(m interface{}) error { return nil } -func fromIncomingContext(ctx context.Context) metadata.MD { - return ctx.Value("metadata").(metadata.MD) +type contextKeyMetadata struct{} + +func contextWithMetadata(ctx context.Context, metadata metadata.MD) context.Context { + return context.WithValue(ctx, contextKeyMetadata{}, metadata) +} + +func metadataFromContext(ctx context.Context) (meta metadata.MD, ok bool) { + meta, ok = ctx.Value(contextKeyMetadata{}).(metadata.MD) + return } func Test_ChainStreamInterceptor(t *testing.T) { salt := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { - md := fromIncomingContext(ss.Context()) + md, ok := metadataFromContext(ss.Context()) + if !ok { + t.Fatal("metadata was not in context") + } md.Append("ingredient", "with salt") err := ss.SetHeader(md) assert.NoError(t, err) return handler(srv, ss) } pepper := func(srv interface{}, ss grpc.ServerStream, info *grpc.StreamServerInfo, handler grpc.StreamHandler) error { - md := fromIncomingContext(ss.Context()) + md, ok := metadataFromContext(ss.Context()) + if !ok { + t.Fatal("metadata was not in context") + } md.Append("ingredient", "and pepper") err := ss.SetHeader(md) assert.NoError(t, err) @@ -2028,10 +2029,13 @@ func Test_ChainStreamInterceptor(t *testing.T) { result := make([]string, 0) srv := "asdf" md := metadata.New(map[string]string{}) - ss := MockStream{context: context.WithValue(context.Background(), "metadata", md)} + ss := &MockStream{context: contextWithMetadata(context.Background(), md)} info := &grpc.StreamServerInfo{} handler := func(srv interface{}, stream grpc.ServerStream) error { - md := fromIncomingContext(stream.Context()) + md, ok := metadataFromContext(stream.Context()) + if !ok { + t.Fatal("metadata was not in context") + } vals := md.Get("ingredient") result = append(result, "Soup") result = append(result, vals...) diff --git a/server/handler_test.go b/server/handler_test.go index b9ac45700..0364a7c3c 100644 --- a/server/handler_test.go +++ b/server/handler_test.go @@ -177,7 +177,6 @@ func TestHandler_Endpoints(t *testing.T) { }) i0 := hldr.MustCreateIndexIfNotExists("i0", pilosa.IndexOptions{}) - const shard = 0 tx0 := holder.Txf().NewWritableQcx() defer tx0.Abort() if f, err := i0.CreateFieldIfNotExists("f1", "", pilosa.OptFieldTypeDefault()); err != nil { diff --git a/sql/show.go b/sql/show.go index 66d17d0c6..9a3ada249 100644 --- a/sql/show.go +++ b/sql/show.go @@ -49,7 +49,7 @@ func (s *ShowHandler) execShowTables(ctx context.Context, showStmt *sqlparser.Sh return nil, errors.Wrap(err, "getting schema") } - allowed, ok := ctx.Value("indices").([]string) + allowed, ok := authn.GetIndexes(ctx) result := make(pproto.ConstRowser, 0) for _, ii := range indexInfo { @@ -83,7 +83,7 @@ func (s *ShowHandler) execShowTables(ctx context.Context, showStmt *sqlparser.Sh func (s *ShowHandler) execShowFields(ctx context.Context, showStmt *sqlparser.Show) (pproto.ToRowser, error) { indexName := showStmt.OnTable.ToViewName().Name.String() - allowed, ok := ctx.Value("indices").([]string) + allowed, ok := authn.GetIndexes(ctx) if ok { found := false for _, idx := range allowed {