func (rb *RoomBinding) GenerateMessageKey(ctx scope.Context, kms security.KMS) (
	proto.RoomMessageKey, error) {

	rmkb, err := rb.Room.generateMessageKey(rb.Backend, kms)
	if err != nil {
		return nil, err
	}

	// Insert key and room association into the DB.
	transaction, err := rb.DbMap.Begin()
	if err != nil {
		return nil, err
	}

	if err := transaction.Insert(&rmkb.MessageKey); err != nil {
		if rerr := transaction.Rollback(); rerr != nil {
			backend.Logger(ctx).Printf("rollback error: %s", rerr)
		}
		return nil, err
	}

	if err := transaction.Insert(&rmkb.RoomMessageKey); err != nil {
		if rerr := transaction.Rollback(); rerr != nil {
			backend.Logger(ctx).Printf("rollback error: %s", rerr)
		}
		return nil, err
	}

	if err := transaction.Commit(); err != nil {
		return nil, err
	}

	return rmkb, nil
}

func (m *accountManager) Register(
	ctx scope.Context, kms security.KMS, namespace, id, password string,
	agentID string, agentKey *security.ManagedKey) (
	proto.Account, *security.ManagedKey, error) {

	m.b.Lock()
	defer m.b.Unlock()

	key := fmt.Sprintf("%s:%s", namespace, id)
	if _, ok := m.b.accountIDs[key]; ok {
		return nil, nil, proto.ErrPersonalIdentityInUse
	}

	account, clientKey, err := NewAccount(kms, password)
	if err != nil {
		return nil, nil, err
	}

	if m.b.accounts == nil {
		m.b.accounts = map[snowflake.Snowflake]proto.Account{account.ID(): account}
	} else {
		m.b.accounts[account.ID()] = account
	}

	pi := &personalIdentity{
		accountID: account.ID(),
		namespace: namespace,
		id:        id,
	}
	account.(*memAccount).personalIdentities = []proto.PersonalIdentity{pi}
	if m.b.accountIDs == nil {
		m.b.accountIDs = map[string]*personalIdentity{key: pi}
	} else {
		m.b.accountIDs[key] = pi
	}

	agent, err := m.b.AgentTracker().Get(ctx, agentID)
	if err != nil {
		backend.Logger(ctx).Printf(
			"error locating agent %s for new account %s:%s: %s", agentID, namespace, id, err)
	} else {
		if err := agent.SetClientKey(agentKey, clientKey); err != nil {
			backend.Logger(ctx).Printf(
				"error associating agent %s with new account %s:%s: %s", agentID, namespace, id, err)
		}
		agent.AccountID = account.ID().String()
	}

	return account, clientKey, nil
}

func Serve(ctx scope.Context, addr string) {
	http.Handle("/metrics", prometheus.Handler())

	listener, err := net.Listen("tcp", addr)
	if err != nil {
		ctx.Terminate(err)
	}

	closed := false
	m := sync.Mutex{}
	closeListener := func() {
		m.Lock()
		if !closed {
			listener.Close()
			closed = true
		}
		m.Unlock()
	}

	// Spin off goroutine to watch ctx and close listener if shutdown requested.
	go func() {
		<-ctx.Done()
		closeListener()
	}()

	backend.Logger(ctx).Printf("serving /metrics on %s", addr)
	if err := http.Serve(listener, nil); err != nil {
		fmt.Printf("http[%s]: %s\n", addr, err)
		ctx.Terminate(err)
	}

	closeListener()
	ctx.WaitGroup().Done()
}

func (b *Backend) part(ctx scope.Context, room *Room, session proto.Session) error {
	t, err := b.DbMap.Begin()
	if err != nil {
		return err
	}

	_, err = t.Exec(
		"DELETE FROM presence WHERE room = $1 AND server_id = $2 AND server_era = $3 AND session_id = $4",
		room.Name, b.desc.ID, b.desc.Era, session.ID())
	if err != nil {
		rollback(ctx, t)
		backend.Logger(ctx).Printf("failed to persist departure: %s", err)
		return err
	}

	// Broadcast a presence event.
	// TODO: make this an explicit action via the Room protocol, to support encryption
	if err := room.broadcast(ctx, t, proto.PartEventType, proto.PresenceEvent(*session.View()), session); err != nil {
		rollback(ctx, t)
		return err
	}

	if err := t.Commit(); err != nil {
		return err
	}

	b.Lock()
	if lm, ok := b.listeners[room.Name]; ok {
		delete(lm, session.ID())
	}
	b.Unlock()

	return nil
}

func (b *AccountManagerBinding) ChangeClientKey(
	ctx scope.Context, accountID snowflake.Snowflake, oldKey, newKey *security.ManagedKey) error {

	t, err := b.DbMap.Begin()
	if err != nil {
		return err
	}

	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}

	var account Account
	err = t.SelectOne(
		&account,
		"SELECT nonce, mac, encrypted_user_key, encrypted_private_key FROM account WHERE id = $1",
		accountID.String())
	if err != nil {
		rollback()
		if err == sql.ErrNoRows {
			return proto.ErrAccountNotFound
		}
		return err
	}

	sec := account.Bind(b.Backend).accountSecurity()
	if err := sec.ChangeClientKey(oldKey, newKey); err != nil {
		rollback()
		return err
	}

	res, err := t.Exec(
		"UPDATE account SET mac = $2, encrypted_user_key = $3 WHERE id = $1",
		accountID.String(), sec.MAC, sec.UserKey.Ciphertext)
	if err != nil {
		rollback()
		return err
	}

	n, err := res.RowsAffected()
	if err != nil {
		rollback()
		return err
	}
	if n == 0 {
		rollback()
		return proto.ErrAccountNotFound
	}

	if err := t.Commit(); err != nil {
		return err
	}

	return nil
}

func (r *memRoom) RenameUser(
	ctx scope.Context, session proto.Session, formerName string) (*proto.NickEvent, error) {

	backend.Logger(ctx).Printf(
		"renaming %s from %s to %s\n", session.ID(), formerName, session.Identity().Name())
	payload := &proto.NickEvent{
		SessionID: session.ID(),
		ID:        session.Identity().ID(),
		From:      formerName,
		To:        session.Identity().Name(),
	}
	return payload, r.broadcast(ctx, proto.NickType, payload, session)
}

func (b *AccountManagerBinding) Resolve(ctx scope.Context, namespace, id string) (proto.Account, error) {
	t, err := b.DbMap.Begin()
	account, err := b.resolve(t, namespace, id)
	if err != nil {
		if rerr := t.Rollback(); rerr != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
		return nil, err
	}
	if err := t.Commit(); err != nil {
		return nil, err
	}
	return account, nil
}

func (b *AccountManagerBinding) Get(ctx scope.Context, id snowflake.Snowflake) (proto.Account, error) {
	t, err := b.DbMap.Begin()
	account, err := b.get(t, id)
	if err != nil {
		if rerr := t.Rollback(); rerr != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
		return nil, err
	}
	if err := t.Commit(); err != nil {
		return nil, err
	}
	return account, nil
}

func (rb *RoomBinding) banAgent(ctx scope.Context, agentID proto.UserID, until time.Time) error {
	ban := &BannedAgent{
		AgentID: agentID.String(),
		Room: sql.NullString{
			String: rb.Name,
			Valid:  true,
		},
		Created: time.Now(),
		Expires: gorp.NullTime{
			Time:  until,
			Valid: !until.IsZero(),
		},
	}

	// Loop within transaction in read committed mode to simulate UPSERT.
	t, err := rb.DbMap.Begin()
	if err != nil {
		return err
	}
	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}
	for {
		// Try to insert; if this fails due to duplicate key value, try to update.
		if err := rb.DbMap.Insert(ban); err != nil {
			if !strings.HasPrefix(err.Error(), "pq: duplicate key value") {
				rollback()
				return err
			}
		} else {
			break
		}
		n, err := rb.DbMap.Update(ban)
		if err != nil {
			rollback()
			return err
		}
		if n > 0 {
			break
		}
	}
	if err := t.Commit(); err != nil {
		return err
	}

	bounceEvent := &proto.BounceEvent{Reason: "banned", AgentID: agentID.String()}
	return rb.broadcast(ctx, rb.Room, proto.BounceEventType, bounceEvent)
}

func Run(args []string) {
	out = tabwriter.NewWriter(os.Stdout, 0, 8, 1, '\t', 0)

	if len(args) == 0 {
		generalHelp()
		return
	}

	exe := filepath.Base(os.Args[0])
	cmd, ok := subcommands[args[0]]
	if !ok {
		fmt.Fprintf(os.Stderr, "%s: invalid command: %s\n", exe, args[0])
		fmt.Fprintf(os.Stderr, "Run '%s help' for usage.\n", exe)
		os.Exit(2)
	}

	flags := cmd.flags()
	if err := flags.Parse(args[1:]); err != nil {
		fmt.Fprintf(os.Stderr, "%s %s: %s\n", exe, args[0], err)
		os.Exit(2)
	}

	ctx := backend.LoggingContext(scope.New(), fmt.Sprintf("[%s] ", args[0]))
	backend.Logger(ctx).Printf("starting up")
	if err := cmd.run(ctx, flags.Args()); err != nil {
		fmt.Fprintln(os.Stderr, err.Error())
		os.Exit(1)
	}

	timeout := time.After(10 * time.Second)
	completed := make(chan struct{})
	go func() {
		ctx.WaitGroup().Wait()
		close(completed)
	}()

	fmt.Println("waiting for graceful shutdown...")
	select {
	case <-timeout:
		fmt.Println("timed out")
		os.Exit(1)
	case <-completed:
		fmt.Println("ok")
		os.Exit(0)
	}
}

// invalidatePeer must be called with lock held
func (b *Backend) invalidatePeer(ctx scope.Context, id, era string) {
	logger := backend.Logger(ctx)
	packet, err := proto.MakeEvent(&proto.NetworkEvent{
		Type:      "partition",
		ServerID:  id,
		ServerEra: era,
	})
	if err != nil {
		logger.Printf("cluster: make network event error: %s", err)
		return
	}
	for _, lm := range b.listeners {
		if err := lm.Broadcast(ctx, packet); err != nil {
			logger.Printf("cluster: network event error: %s", err)
		}
	}
}

func (b *AccountManagerBinding) RequestPasswordReset(
	ctx scope.Context, kms security.KMS, namespace, id string) (
	proto.Account, *proto.PasswordResetRequest, error) {

	t, err := b.DbMap.Begin()
	if err != nil {
		return nil, nil, err
	}

	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}

	account, err := b.resolve(t, namespace, id)
	if err != nil {
		rollback()
		return nil, nil, err
	}

	req, err := proto.GeneratePasswordResetRequest(kms, account.ID())
	if err != nil {
		rollback()
		return nil, nil, err
	}

	stored := &PasswordResetRequest{
		ID:        req.ID.String(),
		AccountID: req.AccountID.String(),
		Key:       req.Key,
		Requested: req.Requested,
		Expires:   req.Expires,
	}
	if err := t.Insert(stored); err != nil {
		rollback()
		return nil, nil, err
	}

	if err := t.Commit(); err != nil {
		rollback()
		return nil, nil, err
	}

	return account, req, nil
}

func (atb *AgentTrackerBinding) Register(ctx scope.Context, agent *proto.Agent) error {
	row := &Agent{
		ID:      agent.IDString(),
		IV:      agent.IV,
		MAC:     agent.MAC,
		Created: agent.Created,
	}
	if agent.EncryptedClientKey != nil {
		row.EncryptedClientKey = agent.EncryptedClientKey.Ciphertext
	}

	if err := atb.Backend.DbMap.Insert(row); err != nil {
		if strings.HasPrefix(err.Error(), "pq: duplicate key value") {
			return proto.ErrAgentAlreadyExists
		}
		return err
	}

	backend.Logger(ctx).Printf("registered agent %s", agent.IDString())
	return nil
}

func (rb *RoomBinding) RemoveManager(
	ctx scope.Context, actor proto.Account, actorKey *security.ManagedKey,
	formerManager proto.Account) error {

	t, err := rb.Backend.DbMap.Begin()
	if err != nil {
		return err
	}

	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}

	rmkb := NewRoomManagerKeyBinding(rb)
	rmkb.SetExecutor(t)

	if _, _, _, err := rmkb.Authority(ctx, actor, actorKey); err != nil {
		rollback()
		if err == proto.ErrCapabilityNotFound {
			return proto.ErrAccessDenied
		}
		return err
	}

	if err := rmkb.RevokeFromAccount(ctx, formerManager); err != nil {
		rollback()
		if err == proto.ErrCapabilityNotFound || err == proto.ErrAccessDenied {
			return proto.ErrManagerNotFound
		}
		return err
	}

	if err := t.Commit(); err != nil {
		return err
	}

	return nil
}

func (atb *AgentTrackerBinding) SetClientKey(
	ctx scope.Context, agentID string, accessKey *security.ManagedKey,
	accountID snowflake.Snowflake, clientKey *security.ManagedKey) error {

	t, err := atb.Backend.DbMap.Begin()
	if err != nil {
		return err
	}

	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}

	agent, err := atb.getFromDB(agentID, atb.Backend.DbMap)
	if err != nil {
		rollback()
		return err
	}

	if err := agent.SetClientKey(accessKey, clientKey); err != nil {
		rollback()
		return err
	}

	err = atb.setClientKeyInDB(
		agentID, accountID.String(), agent.EncryptedClientKey.Ciphertext, t)
	if err != nil {
		rollback()
		return err
	}

	if err := t.Commit(); err != nil {
		return err
	}

	return nil
}

func ScanLoop(ctx scope.Context, listener *pq.Listener) {
	defer ctx.WaitGroup().Done()

	logger := backend.Logger(ctx)
	for {
		select {
		case <-ctx.Done():
			logger.Printf("received cancellation signal, shutting down")
			return
		case notice := <-listener.Notify:
			if notice == nil {
				logger.Printf("received nil from listener")
				continue
			}

			var msg psql.BroadcastMessage

			if err := json.Unmarshal([]byte(notice.Extra), &msg); err != nil {
				logger.Printf("error: pq listen: invalid broadcast: %s", err)
				logger.Printf("         payload: %#v", notice.Extra)
				continue
			}

			switch msg.Event.Type {
			case proto.BounceEventType:
				bounceActivity.WithLabelValues(msg.Room).Inc()
			case proto.JoinEventType:
				joinActivity.WithLabelValues(msg.Room).Inc()
			case proto.PartEventType:
				partActivity.WithLabelValues(msg.Room).Inc()
			case proto.SendEventType:
				messageActivity.WithLabelValues(msg.Room).Inc()
			}
		}
	}
}

func (b *AccountManagerBinding) ConfirmPasswordReset(
	ctx scope.Context, kms security.KMS, confirmation, password string) error {

	id, mac, err := proto.ParsePasswordResetConfirmation(confirmation)
	if err != nil {
		return err
	}

	t, err := b.DbMap.Begin()
	if err != nil {
		return err
	}

	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}

	req := &proto.PasswordResetRequest{
		ID: id,
	}

	var (
		stored  PasswordResetRequest
		account *AccountBinding
	)

	cols, err := allColumns(b.DbMap, PasswordResetRequest{}, "")
	if err != nil {
		return err
	}
	err = t.SelectOne(
		&stored,
		fmt.Sprintf(
			"SELECT %s FROM password_reset_request WHERE id = $1 AND expires > NOW() AND invalidated IS NULL AND consumed IS NULL",
			cols),
		id.String())
	if err != nil && err != sql.ErrNoRows {
		rollback()
		return err
	}

	if err == nil {
		req.Key = stored.Key
		if err := req.AccountID.FromString(stored.AccountID); err == nil {
			account, err = b.get(t, req.AccountID)
			if err != nil && err != proto.ErrAccountNotFound {
				rollback()
				return err
			}
		}
	}

	if !req.VerifyMAC(mac) || account == nil {
		rollback()
		fmt.Printf("invalid mac or no account (%#v)\n", account)
		return proto.ErrInvalidConfirmationCode
	}

	sec, err := account.accountSecurity().ResetPassword(kms, password)
	if err != nil {
		rollback()
		fmt.Printf("reset password failed: %s\n", err)
		return err
	}

	_, err = t.Exec(
		"UPDATE account SET mac = $2, encrypted_user_key = $3 WHERE id = $1",
		account.ID().String(), sec.MAC, sec.UserKey.Ciphertext)
	if err != nil {
		rollback()
		fmt.Printf("update 1 failed: %s\n", err)
		return err
	}

	_, err = t.Exec("UPDATE password_reset_request SET consumed = NOW() where id = $1", id.String())
	if err != nil {
		rollback()
		fmt.Printf("update 2 failed: %s\n", err)
		return err
	}

	_, err = t.Exec(
		"UPDATE password_reset_request SET invalidated = NOW() where account_id = $1 AND id != $2",
		account.ID().String(), id)
	if err != nil {
		rollback()
		fmt.Printf("update 3 failed: %s\n", err)
		return err
	}

	if err := t.Commit(); err != nil {
		fmt.Printf("commit failed: %s\n", err)
		return err
	}

	return nil
}

func (lm ListenerMap) Broadcast(ctx scope.Context, event *proto.Packet, exclude ...string) error {
	payload, err := event.Payload()
	if err != nil {
		return err
	}

	excludeSet := map[string]struct{}{}
	for _, exc := range exclude {
		excludeSet[exc] = struct{}{}
	}

	// Inspect packet to see if it's a bounce event. If so, we'll deliver it
	// only to the bounced parties.
	bounceAgentID := ""
	bounceIP := ""
	if event.Type == proto.BounceEventType {
		if bounceEvent, ok := payload.(*proto.BounceEvent); ok {
			bounceAgentID = bounceEvent.AgentID
			bounceIP = bounceEvent.IP
		} else {
			backend.Logger(ctx).Printf("wtf? expected *proto.BounceEvent, got %T", payload)
		}
	}

	// Inspect packet to see if it's a join event. If so, we'll enable the excluded
	// listener, and look for aliased sessions to kick into fast-keepalive mode.
	fastKeepaliveAgentID := ""
	if event.Type == proto.JoinEventType {
		if presence, ok := payload.(*proto.PresenceEvent); ok {
			if idx := strings.IndexRune(string(presence.ID), '-'); idx >= 1 {
				fastKeepaliveAgentID = string(presence.ID[:idx])
			}
		}
		for _, sessionID := range exclude {
			listener, ok := lm[sessionID]
			if ok && !listener.enabled {
				listener.enabled = true
				lm[sessionID] = listener
			}
		}
	}

	for sessionID, listener := range lm {
		if _, ok := excludeSet[sessionID]; !ok {
			if bounceAgentID != "" {
				if listener.Session.Identity().ID().String() == bounceAgentID {
					backend.Logger(ctx).Printf("sending disconnect to %s: %#v", listener.ID(), payload)
					discEvent := &proto.DisconnectEvent{Reason: payload.(*proto.BounceEvent).Reason}
					if err := listener.Send(ctx, proto.DisconnectEventType, discEvent); err != nil {
						backend.Logger(ctx).Printf("error sending disconnect event to %s: %s",
							listener.ID(), err)
					}
				}
				continue
			}
			if bounceIP != "" {
				if listener.Client.IP == bounceIP {
					backend.Logger(ctx).Printf("sending disconnect to %s: %#v", listener.ID(), payload)
					discEvent := &proto.DisconnectEvent{Reason: payload.(*proto.BounceEvent).Reason}
					if err := listener.Send(ctx, proto.DisconnectEventType, discEvent); err != nil {
						backend.Logger(ctx).Printf("error sending disconnect event to %s: %s",
							listener.ID(), err)
					}
				}
				continue
			}
			if fastKeepaliveAgentID != "" && strings.HasPrefix(sessionID, fastKeepaliveAgentID) {
				if err := listener.CheckAbandoned(); err != nil {
					fmt.Errorf("fast keepalive to %s: %s", listener.ID(), err)
				}
			}
			if !listener.enabled {
				// The event occurred before the listener joined, so don't deliver it.
				backend.Logger(ctx).Printf("not delivering event %s before %s joined",
					event.Type, listener.ID())
				continue
			}
			if err := listener.Send(ctx, event.Type, payload); err != nil {
				// TODO: accumulate errors
				return fmt.Errorf("send message to %s: %s", listener.ID(), err)
			}
		}
	}

	return nil
}

func (b *Backend) join(ctx scope.Context, room *Room, session proto.Session) error {
	client := &proto.Client{}
	if !client.FromContext(ctx) {
		return fmt.Errorf("client data not found in scope")
	}

	bannedAgentCols, err := allColumns(b.DbMap, BannedAgent{}, "")
	if err != nil {
		return err
	}

	bannedIPCols, err := allColumns(b.DbMap, BannedIP{}, "")
	if err != nil {
		return err
	}

	t, err := b.DbMap.Begin()
	if err != nil {
		return err
	}

	rb := func() { rollback(ctx, t) }

	// Check for agent ID bans.
	agentBans, err := t.Select(
		BannedAgent{},
		fmt.Sprintf(
			"SELECT %s FROM banned_agent WHERE agent_id = $1 AND (room IS NULL OR room = $2) AND (expires IS NULL OR expires > NOW())",
			bannedAgentCols),
		session.Identity().ID().String(), room.Name)
	if err != nil {
		rb()
		return err
	}
	if len(agentBans) > 0 {
		backend.Logger(ctx).Printf("access denied to %s: %#v", session.Identity().ID(), agentBans)
		if err := t.Rollback(); err != nil {
			return err
		}
		return proto.ErrAccessDenied
	}

	// Check for IP bans.
	ipBans, err := t.Select(
		BannedIP{},
		fmt.Sprintf(
			"SELECT %s FROM banned_ip WHERE ip = $1 AND (room IS NULL OR room = $2) AND (expires IS NULL OR expires > NOW())",
			bannedIPCols),
		client.IP, room.Name)
	if err != nil {
		rb()
		return err
	}
	if len(ipBans) > 0 {
		backend.Logger(ctx).Printf("access denied to %s: %#v", client.IP, ipBans)
		if err := t.Rollback(); err != nil {
			return err
		}
		return proto.ErrAccessDenied
	}

	// Write to session log.
	// TODO: do proper upsert simulation
	entry := &SessionLog{
		SessionID: session.ID(),
		IP:        client.IP,
		Room:      room.Name,
		UserAgent: client.UserAgent,
		Connected: client.Connected,
	}
	if _, err := t.Delete(entry); err != nil {
		if rerr := t.Rollback(); rerr != nil {
			backend.Logger(ctx).Printf("rollback error: %s", rerr)
		}
		return err
	}
	if err := t.Insert(entry); err != nil {
		if rerr := t.Rollback(); rerr != nil {
			backend.Logger(ctx).Printf("rollback error: %s", rerr)
		}
		return err
	}

	// Broadcast a presence event.
	// TODO: make this an explicit action via the Room protocol, to support encryption

	presence := &Presence{
		Room:      room.Name,
		ServerID:  b.desc.ID,
		ServerEra: b.desc.Era,
		SessionID: session.ID(),
		Updated:   time.Now(),
	}
	err = presence.SetFact(&proto.Presence{
		SessionView:    *session.View(),
		LastInteracted: presence.Updated,
	})
	if err != nil {
		rb()
		return fmt.Errorf("presence marshal error: %s", err)
//.........这里部分代码省略.........

func (b *AccountManagerBinding) GrantStaff(
	ctx scope.Context, accountID snowflake.Snowflake, kmsCred security.KMSCredential) error {

	// Look up the target account's (system) encrypted client key. This is
	// not part of the transaction, because we want to interact with KMS
	// before we proceed. That should be fine, since this is an infrequently
	// used action.
	var row struct {
		EncryptedClientKey []byte `db:"encrypted_system_key"`
		Nonce              []byte `db:"nonce"`
	}
	err := b.DbMap.SelectOne(
		&row, "SELECT encrypted_system_key, nonce FROM account WHERE id = $1", accountID.String())
	if err != nil {
		if err == sql.ErrNoRows {
			return proto.ErrAccountNotFound
		}
		return err
	}

	// Use kmsCred to obtain kms and decrypt the client's key.
	kms := kmsCred.KMS()
	clientKey := &security.ManagedKey{
		KeyType:      proto.ClientKeyType,
		Ciphertext:   row.EncryptedClientKey,
		ContextKey:   "nonce",
		ContextValue: base64.URLEncoding.EncodeToString(row.Nonce),
	}
	if err := kms.DecryptKey(clientKey); err != nil {
		return err
	}

	// Grant staff capability. This involves marshalling kmsCred to JSON and
	// encrypting it with the client key.
	nonce, err := kms.GenerateNonce(clientKey.KeyType.BlockSize())
	if err != nil {
		return err
	}

	capability, err := security.GrantSharedSecretCapability(clientKey, nonce, kmsCred.KMSType(), kmsCred)
	if err != nil {
		return err
	}

	// Store capability and update account table.
	t, err := b.DbMap.Begin()
	if err != nil {
		return err
	}

	rollback := func() {
		if err := t.Rollback(); err != nil {
			backend.Logger(ctx).Printf("rollback error: %s", err)
		}
	}

	dbCap := &Capability{
		ID:                   capability.CapabilityID(),
		NonceBytes:           capability.Nonce(),
		EncryptedPrivateData: capability.EncryptedPayload(),
		PublicData:           capability.PublicPayload(),
	}
	if err := t.Insert(dbCap); err != nil {
		rollback()
		return err
	}

	result, err := t.Exec(
		"UPDATE account SET staff_capability_id = $2 WHERE id = $1",
		accountID.String(), capability.CapabilityID())
	if err != nil {
		rollback()
		return err
	}
	n, err := result.RowsAffected()
	if err != nil {
		rollback()
		return err
	}
	if n != 1 {
		rollback()
		return proto.ErrAccountNotFound
	}

	if err := t.Commit(); err != nil {
		return err
	}

	return nil
}