Golang snowflake.Snowflake类代码示例

func (log *memLog) Latest(ctx scope.Context, n int, before snowflake.Snowflake) ([]proto.Message, error) {
	log.Lock()
	defer log.Unlock()

	end := len(log.msgs)
	if !before.IsZero() {
		for end > 0 && !log.msgs[end-1].ID.Before(before) {
			end--
		}
	}

	start := end - n
	if start < 0 {
		start = 0
	}

	slice := make([]*proto.Message, 0, n)
	for _, msg := range log.msgs[start:] {
		if time.Time(msg.Deleted).IsZero() {
			slice = append(slice, maybeTruncate(msg))
			if len(slice) >= n {
				break
			}
		}
	}
	if len(slice) == 0 {
		return []proto.Message{}, nil
	}

	messages := make([]proto.Message, len(slice))
	for i, msg := range slice {
		messages[i] = *msg
	}
	return messages, nil
}

func (c *Client) AuthenticateWithAgent(ctx scope.Context, backend Backend, agent *Agent, agentKey *security.ManagedKey) error {
	if agent.AccountID == "" {
		return nil
	}

	var accountID snowflake.Snowflake
	if err := accountID.FromString(agent.AccountID); err != nil {
		return err
	}

	account, err := backend.AccountManager().Get(ctx, accountID)
	if err != nil {
		if err == ErrAccountNotFound {
			return nil
		}
		return err
	}

	clientKey, err := agent.Unlock(agentKey)
	if err != nil {
		return fmt.Errorf("agent key error: %s", err)
	}

	c.Account = account
	c.Authorization.ClientKey = clientKey
	return nil
}

func (b *AccountManagerBinding) RevokeStaff(ctx scope.Context, accountID snowflake.Snowflake) error {
	_, err := b.DbMap.Exec(
		"DELETE FROM capability USING account"+
			" WHERE account.id = $1 AND capability.id = account.staff_capability_id",
		accountID.String())
	return err
}

func (s *Server) handleRoom(w http.ResponseWriter, r *http.Request) {
	ctx := s.rootCtx.Fork()

	// Resolve the room.
	// TODO: support room creation?
	roomName := mux.Vars(r)["room"]
	room, err := s.b.GetRoom(ctx, roomName)
	if err != nil {
		if err == proto.ErrRoomNotFound {
			http.Error(w, "404 page not found", http.StatusNotFound)
			return
		}
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	// Tag the agent. We use an authenticated but un-encrypted cookie.
	agent, cookie, agentKey, err := getAgent(ctx, s, r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	client := &proto.Client{Agent: agent}
	client.FromRequest(ctx, r)

	// Look up account associated with agent.
	var accountID snowflake.Snowflake
	if err := accountID.FromString(agent.AccountID); agent.AccountID != "" && err == nil {
		if err := client.AuthenticateWithAgent(ctx, s.b, room, agent, agentKey); err != nil {
			fmt.Printf("agent auth failed: %s\n", err)
			switch err {
			case proto.ErrAccessDenied:
				http.Error(w, err.Error(), http.StatusForbidden)
			default:
				http.Error(w, err.Error(), http.StatusInternalServerError)
			}
			return
		}
	}

	// Upgrade to a websocket and set cookie.
	headers := http.Header{}
	if cookie != nil {
		headers.Add("Set-Cookie", cookie.String())
	}
	conn, err := upgrader.Upgrade(w, r, headers)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}
	defer conn.Close()

	// Serve the session.
	session := newSession(ctx, s, conn, roomName, room, client, agentKey)
	if err = session.serve(); err != nil {
		// TODO: error handling
		return
	}
}

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

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

	account, ok := m.b.accounts[accountID.String()]
	if !ok {
		return proto.ErrAccountNotFound
	}
	memAcc := account.(*memAccount)

	kms := kmsCred.KMS()
	key := memAcc.sec.SystemKey.Clone()
	if err := kms.DecryptKey(&key); err != nil {
		return err
	}

	nonce, err := kms.GenerateNonce(key.KeyType.BlockSize())
	if err != nil {
		return err
	}

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

	memAcc.staffCapability = capability
	return nil
}

func (b *AccountManagerBinding) SetUserKey(
	ctx scope.Context, accountID snowflake.Snowflake, key *security.ManagedKey) error {

	if !key.Encrypted() {
		return security.ErrKeyMustBeEncrypted
	}

	res, err := b.DbMap.Exec(
		"UPDATE account SET encrypted_user_key = $2 WHERE id = $1", accountID.String(), key.Ciphertext)
	if err != nil {
		if err == sql.ErrNoRows {
			return proto.ErrAccountNotFound
		}
		return err
	}

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

	return nil
}

func getClient(ctx scope.Context, s *Server, r *http.Request) (*proto.Client, *http.Cookie, *security.ManagedKey, error) {
	agent, cookie, agentKey, err := getAgent(ctx, s, r)
	if err != nil {
		return nil, nil, nil, err
	}

	client := &proto.Client{
		Agent: agent,
		Authorization: proto.Authorization{
			ClientKey: agentKey,
		},
	}
	client.FromRequest(ctx, r)

	// Look up account associated with agent.
	var accountID snowflake.Snowflake
	if err := accountID.FromString(agent.AccountID); agent.AccountID != "" && err == nil {
		if err := client.AuthenticateWithAgent(ctx, s.b, agent, agentKey); err != nil {
			fmt.Printf("agent auth failed: %s\n", err)
			switch err {
			case proto.ErrAccessDenied:
				// allow session to proceed, but agent will not be logged into account
				agent.AccountID = ""
			default:
				return nil, nil, nil, err
			}
		}
	}

	return client, cookie, agentKey, nil
}

func (rb *RoomBinding) GetMessage(ctx scope.Context, id snowflake.Snowflake) (*proto.Message, error) {
	var msg Message

	nDays, err := rb.DbMap.SelectInt("SELECT retention_days FROM room WHERE name = $1", rb.Name)
	if err != nil {
		return nil, err
	}

	err = rb.DbMap.SelectOne(
		&msg,
		"SELECT room, id, previous_edit_id, parent, posted, edited, deleted,"+
			" session_id, sender_id, sender_name, server_id, server_era, content, encryption_key_id"+
			" FROM message WHERE room = $1 AND id = $2",
		rb.Name, id.String())
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, proto.ErrMessageNotFound
		}
		return nil, err
	}
	if nDays > 0 {
		threshold := time.Now().Add(time.Duration(-nDays) * 24 * time.Hour)
		if msg.Posted.Before(threshold) {
			return nil, proto.ErrMessageNotFound
		}
	}
	m := msg.ToBackend()
	return &m, nil
}

func (rb *RoomBinding) MessageKey(ctx scope.Context) (proto.RoomMessageKey, error) {
	var row struct {
		MessageKey
		RoomMessageKey
	}
	err := rb.DbMap.SelectOne(
		&row,
		"SELECT mk.id, mk.encrypted_key, mk.iv, mk.nonce,"+
			" r.room, r.key_id, r.activated, r.expired, r.comment"+
			" FROM master_key mk, room_master_key r"+
			" WHERE r.room = $1 AND mk.id = r.key_id AND r.expired < r.activated"+
			" ORDER BY r.activated DESC LIMIT 1",
		rb.Name)
	if err != nil {
		if err == sql.ErrNoRows {
			return nil, nil
		}
		return nil, err
	}

	msgKey := &security.ManagedKey{
		KeyType:      proto.RoomMessageKeyType,
		IV:           row.MessageKey.IV,
		Ciphertext:   row.MessageKey.EncryptedKey,
		ContextKey:   "room",
		ContextValue: rb.Room.Name,
	}
	var keyID snowflake.Snowflake
	if err := keyID.FromString(row.KeyID); err != nil {
		return nil, err
	}
	return NewRoomMessageKeyBinding(rb, keyID, msgKey, row.Nonce), nil
}

func (t *PMTracker) Room(ctx scope.Context, kms security.KMS, pmID snowflake.Snowflake, client *proto.Client) (proto.Room, *security.ManagedKey, error) {
	row, err := t.Backend.Get(PM{}, pmID.String())
	if row == nil || err != nil {
		if row == nil || err == sql.ErrNoRows {
			return nil, nil, proto.ErrPMNotFound
		}
	}

	pm := row.(*PM).ToBackend()
	pmKey, modified, otherName, err := pm.Access(ctx, t.Backend, kms, client)
	if err != nil {
		return nil, nil, err
	}

	if modified {
		_, err := t.Backend.DbMap.Exec(
			"UPDATE pm SET receiver = $2, receiver_mac = $3, encrypted_receiver_key = $4 WHERE id = $1",
			pm.ID.String(), string(pm.Receiver), pm.ReceiverMAC, pm.EncryptedReceiverKey.Ciphertext)
		if err != nil {
			return nil, nil, err
		}
	}

	room := &PMRoomBinding{
		RoomBinding: RoomBinding{
			RoomName:  fmt.Sprintf("pm:%s", pm.ID),
			RoomTitle: fmt.Sprintf("%s (private chat)", otherName),
			Backend:   t.Backend,
		},
		pm: pm,
	}

	return room, pmKey, nil
}

func (b *AccountManagerBinding) getOTP(db gorp.SqlExecutor, kms security.KMS, accountID snowflake.Snowflake) (*proto.OTP, error) {
	encryptedOTP, err := b.getRawOTP(db, accountID)
	if err != nil {
		return nil, err
	}

	key := security.ManagedKey{
		KeyType:      OTPKeyType,
		IV:           encryptedOTP.IV,
		Ciphertext:   encryptedOTP.EncryptedKey,
		ContextKey:   "account",
		ContextValue: accountID.String(),
	}
	if err := kms.DecryptKey(&key); err != nil {
		return nil, err
	}

	uriBytes, err := security.DecryptGCM(&key, encryptedOTP.IV, encryptedOTP.Digest, encryptedOTP.EncryptedURI, nil)
	if err != nil {
		return nil, err
	}

	otp := &proto.OTP{
		URI:       string(uriBytes),
		Validated: encryptedOTP.Validated,
	}
	return otp, 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 (m *accountManager) Get(ctx scope.Context, id snowflake.Snowflake) (proto.Account, error) {
	m.b.Lock()
	defer m.b.Unlock()

	account, ok := m.b.accounts[id.String()]
	if !ok {
		return nil, proto.ErrAccountNotFound
	}
	return account, nil
}

func (b *AccountManagerBinding) getRawOTP(db gorp.SqlExecutor, accountID snowflake.Snowflake) (*OTP, error) {
	row, err := db.Get(OTP{}, accountID.String())
	if row == nil || err != nil {
		if row == nil || err == sql.ErrNoRows {
			return nil, proto.ErrOTPNotEnrolled
		}
		return nil, err
	}
	return row.(*OTP), nil
}

func (c *console) resolveAccount(ctx scope.Context, ref string) (proto.Account, error) {
	idx := strings.IndexRune(ref, ':')
	if idx < 0 {
		var accountID snowflake.Snowflake
		if err := accountID.FromString(ref); err != nil {
			return nil, err
		}
		return c.backend.AccountManager().Get(ctx, accountID)
	}
	return c.backend.AccountManager().Resolve(ctx, ref[:idx], ref[idx+1:])
}