Extend LUCI_CONTEXT["local_auth"] protocol to understand accounts.

common/auth/localauth/server.go

 // Copyright 2017 The LUCI Authors. All rights reserved.
 // Use of this source code is governed under the Apache License, Version 2.0
 // that can be found in the LICENSE file.
 
 package localauth
 
 import (
         "crypto/subtle"
         "encoding/json"
         "fmt"
@@ skipping 53 matching lines @@
 
         // Code returns a code to put into RPC response alongside the error message.
         Code() int
 }
 
 // Server runs a local RPC server that hands out access tokens.
 //
 // Processes that need a token can discover location of this server by looking
 // at "local_auth" section of LUCI_CONTEXT.
 type Server struct {
-        // TokenGenerator produces access tokens.
-        TokenGenerator TokenGenerator
+        // TokenGenerators produce access tokens for given account IDs.
+        TokenGenerators map[string]TokenGenerator
+
+        // DefaultAccountID is account ID subprocesses should pick by default.
+        //
+        // It is put into "local_auth" section of LUIC_CONTEXT. If empty string,
+        // subprocesses won't attempt to use any account by default (they still can
+        // pick some non-default account though).
+        DefaultAccountID string
 
         // Port is a local TCP port to bind to or 0 to allow the OS to pick one.
         Port int
 
         l        sync.Mutex
         secret   []byte             // the clients are expected to send this secret
         listener net.Listener       // to know what to stop in Close, nil after Close
         wg       sync.WaitGroup     // +1 for each request being processed now
         ctx      context.Context    // derived from ctx in Initialize
         cancel   context.CancelFunc // cancels 'ctx'
 
         testingServeHook func() // called right before serving
 }
 
 // Initialize binds the server to a local port and prepares it for serving.
 //
 // The provided context is used as base context for request handlers and for
 // logging.
 //
-// Returns lucictx.LocalAuth structure that specifies how to contact the server.
-// It should be put into "local_auth" section of LUCI_CONTEXT where clients can
-// discover it.
+// Returns a copy of lucictx.LocalAuth structure that specifies how to contact
+// the server. It should be put into "local_auth" section of LUCI_CONTEXT where
+// clients can discover it.
 func (s *Server) Initialize(ctx context.Context) (*lucictx.LocalAuth, error) {
         s.l.Lock()
         defer s.l.Unlock()
 
         if s.ctx != nil {
                 return nil, fmt.Errorf("already initialized")
         }
 
         secret := make([]byte, 48)
         if _, err := cryptorand.Read(ctx, secret); err != nil {
                 return nil, err
         }
 
         ln, err := net.Listen("tcp", fmt.Sprintf("127.0.0.1:%d", s.Port))
         if err != nil {
                 return nil, err
         }
 
         s.ctx, s.cancel = context.WithCancel(ctx)
         s.listener = ln
         s.secret = secret
 
+        // Build a sorted list of LocalAuthAccount to put into the context.
+        ids := make([]string, 0, len(s.TokenGenerators))
+        for id := range s.TokenGenerators {
+                ids = append(ids, id)
+        }
+        sort.Strings(ids)
+        accounts := make([]lucictx.LocalAuthAccount, len(ids))
+        for i, id := range ids {
+                accounts[i] = lucictx.LocalAuthAccount{ID: id}
+        }
+
         return &lucictx.LocalAuth{
-                RPCPort: uint32(ln.Addr().(*net.TCPAddr).Port),
-                Secret:  secret,
+                RPCPort:          uint32(ln.Addr().(*net.TCPAddr).Port),
+                Secret:           secret,
+                Accounts:         accounts,
+                DefaultAccountID: s.DefaultAccountID,
         }, nil
 }
 
 // Serve runs a serving loop.
 //
 // It unblocks once Close is called and all pending requests are served.
 //
 // Returns nil if serving was stopped by Close or non-nil if it failed for some
 // other reason.
 func (s *Server) Serve() (err error) {
         s.l.Lock()
         switch {
         case s.ctx == nil:
                 err = fmt.Errorf("not initialized")
         case s.listener == nil:
                 err = fmt.Errorf("already closed")
         }
         if err != nil {
                 s.l.Unlock()
                 return
         }
         listener := s.listener // accessed outside the lock
         srv := http.Server{
                 Handler: &protocolHandler{
                         ctx:    s.ctx,
                         wg:     &s.wg,
                         secret: s.secret,
-                        tokens: s.TokenGenerator,
+                        tokens: s.TokenGenerators,
                 },
         }
         s.l.Unlock()
 
         // Notify unit tests that we have initialized.
         if s.testingServeHook != nil {
                 s.testingServeHook()
         }
 
         err = srv.Serve(listener) // blocks until Close() is called
@@ skipping 53 matching lines @@
 //  * The server sets Content-Length header in the response.
 //  * Protocol-level errors have non-200 HTTP status code.
 //  * Logic errors have 200 HTTP status code and error is communicated in
 //    the response body.
 //
 // The only supported method currently is 'GetOAuthToken':
 //
 //    Request body:
 //    {
 //      "scopes": [<string scope1>, <string scope2>, ...],
-//      "secret": <string from LUCI_CONTEXT.local_auth.secret>
+//      "secret": <string from LUCI_CONTEXT.local_auth.secret>,
+//      "account_id": <ID of some account from LUCI_CONTEXT.local_auth.accounts>
 //    }
 //    Response body:
 //    {
 //      "error_code": <int, on success not set or 0>,
 //      "error_message": <string, on success not set>,
 //      "access_token": <string with actual token (on success)>,
 //      "expiry": <int with unix timestamp in seconds (on success)>
 //    }
 //
 // See also python counterpart of this code:
 // https://github.com/luci/luci-py/blob/master/client/utils/auth_server.py
 type protocolHandler struct {
-        ctx    context.Context // the parent context
-        wg     *sync.WaitGroup // used for graceful shutdown
-        secret []byte          // expected "secret" value
-        tokens TokenGenerator  // the actual producer of tokens
+        ctx    context.Context           // the parent context
+        wg     *sync.WaitGroup           // used for graceful shutdown
+        secret []byte                    // expected "secret" value
+        tokens map[string]TokenGenerator // the actual producer of tokens (per account)
 }
 
 // protocolError triggers an HTTP reply with some non-200 status code.
 type protocolError struct {
         Status  int    // HTTP status to set
         Message string // the message to put in the body
 }
 
 func (e *protocolError) Error() string {
         return fmt.Sprintf("%s (HTTP %d)", e.Message, e.Status)
@@ skipping 115 matching lines @@
                         Status:  http.StatusNotFound,
                         Message: fmt.Sprintf("Unknown RPC method %q", method),
                 }
         }
 }
 
 ////////////////////////////////////////////////////////////////////////////////
 // RPC implementations.
 
 func (h *protocolHandler) handleGetOAuthToken(req *rpcs.GetOAuthTokenRequest) (*rpcs.GetOAuthTokenResponse, error) {
-        // Validate the request, verify the correct secret is passed.
+        // Validate the request.
         if err := req.Validate(); err != nil {
                 return nil, &protocolError{
                         Status:  400,
-                        Message: err.Error(),
+                        Message: fmt.Sprintf("Bad request: %s.", err.Error()),
+                }
+        }
+        // TODO(vadimsh): Remove this check once it is moved into Validate().
+        if req.AccountID == "" {
+                return nil, &protocolError{
+                        Status:  400,
+                        Message: `Bad request: field "account_id" is required.`,
                 }
         }
         if subtle.ConstantTimeCompare(h.secret, req.Secret) != 1 {
                 return nil, &protocolError{
                         Status:  403,
-                        Message: `Invalid secret.`,
+                        Message: "Invalid secret.",
+                }
+        }
+        generator := h.tokens[req.AccountID]
+        if generator == nil {
+                return nil, &protocolError{
+                        Status:  404,
+                        Message: fmt.Sprintf("Unrecognized account ID %q.", req.AccountID),
                 }
         }
 
         // Dedup and sort scopes.
         scopes := stringset.New(len(req.Scopes))
         for _, s := range req.Scopes {
                 scopes.Add(s)
         }
         sortedScopes := scopes.ToSlice()
         sort.Strings(sortedScopes)
 
         // Ask the token provider for the token. This may produce ErrorWithCode.
-        tok, err := h.tokens(h.ctx, sortedScopes, minTokenLifetime)
+        tok, err := generator(h.ctx, sortedScopes, minTokenLifetime)
         if err != nil {
                 return nil, err
         }
         return &rpcs.GetOAuthTokenResponse{
                 AccessToken: tok.AccessToken,
                 Expiry:      tok.Expiry.Unix(),
         }, nil
 }