<commit_before>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains code related to state tracking.  If enabled, state
// tracking will capture the initial READY packet and many other websocket
// events and maintain an in-memory state of of guilds, channels, users, and
// so forth.  This information can be accessed through the Session.State struct.

package discordgo

import (
	"errors"
	"fmt"
)

// ErrNilState is returned when the state is nil.
var ErrNilState = errors.New("State not instantiated, please use discordgo.New() or assign Session.State.")

// NewState creates an empty state.
func NewState() *State {
	return &State{
		Ready: Ready{
			PrivateChannels: []*Channel{},
			Guilds:          []*Guild{},
		},
	}
}

// OnReady takes a Ready event and updates all internal state.
func (s *State) OnReady(r *Ready) error {
	if s == nil {
		return ErrNilState
	}

	s.Lock()
	defer s.Unlock()

	s.Ready = *r

	for _, g := range s.Guilds {
		for _, c := range g.Channels {
			c.GuildID = g.ID
		}
	}

	return nil
}

// GuildAdd adds a guild to the current world state, or
// updates it if it already exists.
func (s *State) GuildAdd(guild *Guild) error {
	if s == nil {
		return ErrNilState
	}

	s.Lock()
	defer s.Unlock()

	// If the guild exists, replace it.
	for i, g := range s.Guilds {
		if g.ID == guild.ID {
			// Don't stomp on properties that don't come in updates.
			guild.Members = g.Members
			guild.Presences = g.Presences
			guild.Channels = g.Channels
			guild.VoiceStates = g.VoiceStates
			s.Guilds[i] = guild
			return nil
		}
	}

	s.Guilds = append(s.Guilds, guild)
	return nil
}

// GuildRemove removes a guild from current world state.
func (s *State) GuildRemove(guild *Guild) error {
	if s == nil {
		return ErrNilState
	}

	s.Lock()
	defer s.Unlock()

	for i, g := range s.Guilds {
		if g.ID == guild.ID {
			s.Guilds = append(s.Guilds[:i], s.Guilds[i+1:]...)
			return nil
		}
	}

	return errors.New("Guild not found.")
}

// Guild gets a guild by ID.
// Useful for querying if @me is in a guild:
//     _, err := discordgo.Session.State.Guild(guildID)
//     isInGuild := err == nil
func (s *State) Guild(guildID string) (*Guild, error) {
	if s == nil {
		return nil, ErrNilState
	}

	s.RLock()
	defer s.RUnlock()

	for _, g := range s.Guilds {
		if g.ID == guildID {
			return g, nil
		}
	}

	return nil, errors.New("Guild not found.")
}

// TODO: Consider moving Guild state update methods onto *Guild.

// MemberAdd adds a member to the current world state, or
// updates it if it already exists.
func (s *State) MemberAdd(member *Member) error {
	if s == nil {
		return ErrNilState
	}

	guild, err := s.Guild(member.GuildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, m := range guild.Members {
		if m.User.ID == member.User.ID {
			guild.Members[i] = member
			return nil
		}
	}

	guild.Members = append(guild.Members, member)
	return nil
}

// MemberRemove removes a member from current world state.
func (s *State) MemberRemove(member *Member) error {
	if s == nil {
		return ErrNilState
	}

	guild, err := s.Guild(member.GuildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, m := range guild.Members {
		if m.User.ID == member.User.ID {
			guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
			return nil
		}
	}

	return errors.New("Member not found.")
}

// Member gets a member by ID from a guild.
func (s *State) Member(guildID, userID string) (*Member, error) {
	if s == nil {
		return nil, ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, m := range guild.Members {
		if m.User.ID == userID {
			return m, nil
		}
	}

	return nil, errors.New("Member not found.")
}

// ChannelAdd adds a guild to the current world state, or
// updates it if it already exists.
// Channels may exist either as PrivateChannels or inside
// a guild.
func (s *State) ChannelAdd(channel *Channel) error {
	if s == nil {
		return ErrNilState
	}

	if channel.IsPrivate {
		s.Lock()
		defer s.Unlock()

		// If the channel exists, replace it.
		for i, c := range s.PrivateChannels {
			if c.ID == channel.ID {
				// Don't stomp on messages.
				channel.Messages = c.Messages
				s.PrivateChannels[i] = channel
				return nil
			}
		}

		s.PrivateChannels = append(s.PrivateChannels, channel)
	} else {
		guild, err := s.Guild(channel.GuildID)
		if err != nil {
			return err
		}

		s.Lock()
		defer s.Unlock()

		// If the channel exists, replace it.
		for i, c := range guild.Channels {
			if c.ID == channel.ID {
				// Don't stomp on messages.
				channel.Messages = c.Messages
				guild.Channels[i] = channel
				return nil
			}
		}

		guild.Channels = append(guild.Channels, channel)
	}

	return nil
}

// ChannelRemove removes a channel from current world state.
func (s *State) ChannelRemove(channel *Channel) error {
	if s == nil {
		return ErrNilState
	}

	if channel.IsPrivate {
		s.Lock()
		defer s.Unlock()

		for i, c := range s.PrivateChannels {
			if c.ID == channel.ID {
				s.PrivateChannels = append(s.PrivateChannels[:i], s.PrivateChannels[i+1:]...)
				return nil
			}
		}
	} else {
		guild, err := s.Guild(channel.GuildID)
		if err != nil {
			return err
		}

		s.Lock()
		defer s.Unlock()

		for i, c := range guild.Channels {
			if c.ID == channel.ID {
				guild.Channels = append(guild.Channels[:i], guild.Channels[i+1:]...)
				return nil
			}
		}
	}

	return errors.New("Channel not found.")
}

// GuildChannel gets a channel by ID from a guild.
func (s *State) GuildChannel(guildID, channelID string) (*Channel, error) {
	if s == nil {
		return nil, ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, c := range guild.Channels {
		if c.ID == channelID {
			return c, nil
		}
	}

	return nil, errors.New("Channel not found.")
}

// PrivateChannel gets a private channel by ID.
func (s *State) PrivateChannel(channelID string) (*Channel, error) {
	if s == nil {
		return nil, ErrNilState
	}

	s.RLock()
	defer s.RUnlock()

	for _, c := range s.PrivateChannels {
		if c.ID == channelID {
			return c, nil
		}
	}

	return nil, errors.New("Channel not found.")
}

// Channel gets a channel by ID, it will look in all guilds an private channels.
func (s *State) Channel(channelID string) (*Channel, error) {
	if s == nil {
		return nil, ErrNilState
	}

	c, err := s.PrivateChannel(channelID)
	if err == nil {
		return c, nil
	}

	for _, g := range s.Guilds {
		c, err := s.GuildChannel(g.ID, channelID)
		if err == nil {
			return c, nil
		}
	}

	return nil, errors.New("Channel not found.")
}

// Emoji returns an emoji for a guild and emoji id.
func (s *State) Emoji(guildID, emojiID string) (*Emoji, error) {
	if s == nil {
		return nil, ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, e := range guild.Emojis {
		if e.ID == emojiID {
			return e, nil
		}
	}

	return nil, errors.New("Emoji not found.")
}

// EmojiAdd adds an emoji to the current world state.
func (s *State) EmojiAdd(guildID string, emoji *Emoji) error {
	if s == nil {
		return ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, e := range guild.Emojis {
		if e.ID == emoji.ID {
			guild.Emojis[i] = emoji
			return nil
		}
	}

	guild.Emojis = append(guild.Emojis, emoji)
	return nil
}

// EmojisAdd adds multiple emojis to the world state.
func (s *State) EmojisAdd(guildID string, emojis []*Emoji) error {
	for _, e := range emojis {
		if err := s.EmojiAdd(guildID, e); err != nil {
			return err
		}
	}
	return nil
}

// MessageAdd adds a message to the current world state, or updates it if it exists.
// If the channel cannot be found, the message is discarded.
// Messages are kept in state up to s.MaxMessageCount
func (s *State) MessageAdd(message *Message) error {
	if s == nil {
		return ErrNilState
	}

	c, err := s.Channel(message.ChannelID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	// If the message exists, replace it.
	for i, m := range c.Messages {
		if m.ID == message.ID {
			c.Messages[i] = message
			return nil
		}
	}

	c.Messages = append(c.Messages, message)

	if len(c.Messages) > s.MaxMessageCount {
		s.Unlock()
		for len(c.Messages) > s.MaxMessageCount {
			err := s.MessageRemove(c.Messages[0])
			if err != nil {
				fmt.Println("message remove error: ", err)
			}
		}
		s.Lock()
	}
	return nil
}

// MessageRemove removes a message from the world state.
func (s *State) MessageRemove(message *Message) error {
	if s == nil {
		return ErrNilState
	}

	c, err := s.Channel(message.ChannelID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, m := range c.Messages {
		if m.ID == message.ID {
			c.Messages = append(c.Messages[:i], c.Messages[i+1:]...)
			return nil
		}
	}

	return errors.New("Message not found.")
}

func (s *State) voiceStateUpdate(update *VoiceStateUpdate) error {
	guild, err := s.Guild(update.GuildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	// Handle Leaving Channel
	if update.ChannelID == "" {
		for i, state := range guild.VoiceStates {
			if state.UserID == update.UserID {
				guild.VoiceStates = append(guild.VoiceStates[:i], guild.VoiceStates[i+1:]...)
				return nil
			}
		}
	} else {
		for i, state := range guild.VoiceStates {
			if state.UserID == update.UserID {
				guild.VoiceStates[i] = update.VoiceState
				return nil
			}
		}

		guild.VoiceStates = append(guild.VoiceStates, update.VoiceState)
	}

	return nil
}

// Message gets a message by channel and message ID.
func (s *State) Message(channelID, messageID string) (*Message, error) {
	if s == nil {
		return nil, ErrNilState
	}

	c, err := s.Channel(channelID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, m := range c.Messages {
		if m.ID == messageID {
			return m, nil
		}
	}

	return nil, errors.New("Message not found.")
}

// onInterface handles all events related to states.
func (s *State) onInterface(se *Session, i interface{}) (err error) {
	if s == nil {
		return ErrNilState
	}
	if !se.StateEnabled {
		return nil
	}

	switch t := i.(type) {
	case *Ready:
		err = s.OnReady(t)
	case *GuildCreate:
		err = s.GuildAdd(t.Guild)
	case *GuildUpdate:
		err = s.GuildAdd(t.Guild)
	case *GuildDelete:
		err = s.GuildRemove(t.Guild)
	case *GuildMemberAdd:
		err = s.MemberAdd(t.Member)
	case *GuildMemberUpdate:
		err = s.MemberAdd(t.Member)
	case *GuildMemberRemove:
		err = s.MemberRemove(t.Member)
	case *GuildEmojisUpdate:
		err = s.EmojisAdd(t.GuildID, t.Emojis)
	case *ChannelCreate:
		err = s.ChannelAdd(t.Channel)
	case *ChannelUpdate:
		err = s.ChannelAdd(t.Channel)
	case *ChannelDelete:
		err = s.ChannelRemove(t.Channel)
	case *MessageCreate:
		err = s.MessageAdd(t.Message)
	case *MessageUpdate:
		err = s.MessageAdd(t.Message)
	case *MessageDelete:
		err = s.MessageRemove(t.Message)
	case *VoiceStateUpdate:
		err = s.voiceStateUpdate(t)
	}

	return
}
<commit_msg>Fix channels inside a guild from GUILD_CREATE having a null GuildID<commit_after>// Discordgo - Discord bindings for Go
// Available at https://github.com/bwmarrin/discordgo

// Copyright 2015-2016 Bruce Marriner <bruce@sqls.net>.  All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// This file contains code related to state tracking.  If enabled, state
// tracking will capture the initial READY packet and many other websocket
// events and maintain an in-memory state of of guilds, channels, users, and
// so forth.  This information can be accessed through the Session.State struct.

package discordgo

import (
	"errors"
	"fmt"
)

// ErrNilState is returned when the state is nil.
var ErrNilState = errors.New("State not instantiated, please use discordgo.New() or assign Session.State.")

// NewState creates an empty state.
func NewState() *State {
	return &State{
		Ready: Ready{
			PrivateChannels: []*Channel{},
			Guilds:          []*Guild{},
		},
	}
}

// OnReady takes a Ready event and updates all internal state.
func (s *State) OnReady(r *Ready) error {
	if s == nil {
		return ErrNilState
	}

	s.Lock()
	defer s.Unlock()

	s.Ready = *r

	for _, g := range s.Guilds {
		for _, c := range g.Channels {
			c.GuildID = g.ID
		}
	}

	return nil
}

// GuildAdd adds a guild to the current world state, or
// updates it if it already exists.
func (s *State) GuildAdd(guild *Guild) error {
	if s == nil {
		return ErrNilState
	}

	s.Lock()
	defer s.Unlock()

	// If the guild exists, replace it.
	for i, g := range s.Guilds {
		if g.ID == guild.ID {
			// Don't stomp on properties that don't come in updates.
			guild.Members = g.Members
			guild.Presences = g.Presences
			guild.Channels = g.Channels
			guild.VoiceStates = g.VoiceStates
			s.Guilds[i] = guild
			return nil
		}
	}

	// Otherwise, update the channels to point to the right guild
	for _, c := range guild.Channels {
		c.GuildID = guild.ID
	}

	s.Guilds = append(s.Guilds, guild)
	return nil
}

// GuildRemove removes a guild from current world state.
func (s *State) GuildRemove(guild *Guild) error {
	if s == nil {
		return ErrNilState
	}

	s.Lock()
	defer s.Unlock()

	for i, g := range s.Guilds {
		if g.ID == guild.ID {
			s.Guilds = append(s.Guilds[:i], s.Guilds[i+1:]...)
			return nil
		}
	}

	return errors.New("Guild not found.")
}

// Guild gets a guild by ID.
// Useful for querying if @me is in a guild:
//     _, err := discordgo.Session.State.Guild(guildID)
//     isInGuild := err == nil
func (s *State) Guild(guildID string) (*Guild, error) {
	if s == nil {
		return nil, ErrNilState
	}

	s.RLock()
	defer s.RUnlock()

	for _, g := range s.Guilds {
		if g.ID == guildID {
			return g, nil
		}
	}

	return nil, errors.New("Guild not found.")
}

// TODO: Consider moving Guild state update methods onto *Guild.

// MemberAdd adds a member to the current world state, or
// updates it if it already exists.
func (s *State) MemberAdd(member *Member) error {
	if s == nil {
		return ErrNilState
	}

	guild, err := s.Guild(member.GuildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, m := range guild.Members {
		if m.User.ID == member.User.ID {
			guild.Members[i] = member
			return nil
		}
	}

	guild.Members = append(guild.Members, member)
	return nil
}

// MemberRemove removes a member from current world state.
func (s *State) MemberRemove(member *Member) error {
	if s == nil {
		return ErrNilState
	}

	guild, err := s.Guild(member.GuildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, m := range guild.Members {
		if m.User.ID == member.User.ID {
			guild.Members = append(guild.Members[:i], guild.Members[i+1:]...)
			return nil
		}
	}

	return errors.New("Member not found.")
}

// Member gets a member by ID from a guild.
func (s *State) Member(guildID, userID string) (*Member, error) {
	if s == nil {
		return nil, ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, m := range guild.Members {
		if m.User.ID == userID {
			return m, nil
		}
	}

	return nil, errors.New("Member not found.")
}

// ChannelAdd adds a guild to the current world state, or
// updates it if it already exists.
// Channels may exist either as PrivateChannels or inside
// a guild.
func (s *State) ChannelAdd(channel *Channel) error {
	if s == nil {
		return ErrNilState
	}

	if channel.IsPrivate {
		s.Lock()
		defer s.Unlock()

		// If the channel exists, replace it.
		for i, c := range s.PrivateChannels {
			if c.ID == channel.ID {
				// Don't stomp on messages.
				channel.Messages = c.Messages
				s.PrivateChannels[i] = channel
				return nil
			}
		}

		s.PrivateChannels = append(s.PrivateChannels, channel)
	} else {
		guild, err := s.Guild(channel.GuildID)
		if err != nil {
			return err
		}

		s.Lock()
		defer s.Unlock()

		// If the channel exists, replace it.
		for i, c := range guild.Channels {
			if c.ID == channel.ID {
				// Don't stomp on messages.
				channel.Messages = c.Messages
				guild.Channels[i] = channel
				return nil
			}
		}

		guild.Channels = append(guild.Channels, channel)
	}

	return nil
}

// ChannelRemove removes a channel from current world state.
func (s *State) ChannelRemove(channel *Channel) error {
	if s == nil {
		return ErrNilState
	}

	if channel.IsPrivate {
		s.Lock()
		defer s.Unlock()

		for i, c := range s.PrivateChannels {
			if c.ID == channel.ID {
				s.PrivateChannels = append(s.PrivateChannels[:i], s.PrivateChannels[i+1:]...)
				return nil
			}
		}
	} else {
		guild, err := s.Guild(channel.GuildID)
		if err != nil {
			return err
		}

		s.Lock()
		defer s.Unlock()

		for i, c := range guild.Channels {
			if c.ID == channel.ID {
				guild.Channels = append(guild.Channels[:i], guild.Channels[i+1:]...)
				return nil
			}
		}
	}

	return errors.New("Channel not found.")
}

// GuildChannel gets a channel by ID from a guild.
func (s *State) GuildChannel(guildID, channelID string) (*Channel, error) {
	if s == nil {
		return nil, ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, c := range guild.Channels {
		if c.ID == channelID {
			return c, nil
		}
	}

	return nil, errors.New("Channel not found.")
}

// PrivateChannel gets a private channel by ID.
func (s *State) PrivateChannel(channelID string) (*Channel, error) {
	if s == nil {
		return nil, ErrNilState
	}

	s.RLock()
	defer s.RUnlock()

	for _, c := range s.PrivateChannels {
		if c.ID == channelID {
			return c, nil
		}
	}

	return nil, errors.New("Channel not found.")
}

// Channel gets a channel by ID, it will look in all guilds an private channels.
func (s *State) Channel(channelID string) (*Channel, error) {
	if s == nil {
		return nil, ErrNilState
	}

	c, err := s.PrivateChannel(channelID)
	if err == nil {
		return c, nil
	}

	for _, g := range s.Guilds {
		c, err := s.GuildChannel(g.ID, channelID)
		if err == nil {
			return c, nil
		}
	}

	return nil, errors.New("Channel not found.")
}

// Emoji returns an emoji for a guild and emoji id.
func (s *State) Emoji(guildID, emojiID string) (*Emoji, error) {
	if s == nil {
		return nil, ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, e := range guild.Emojis {
		if e.ID == emojiID {
			return e, nil
		}
	}

	return nil, errors.New("Emoji not found.")
}

// EmojiAdd adds an emoji to the current world state.
func (s *State) EmojiAdd(guildID string, emoji *Emoji) error {
	if s == nil {
		return ErrNilState
	}

	guild, err := s.Guild(guildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, e := range guild.Emojis {
		if e.ID == emoji.ID {
			guild.Emojis[i] = emoji
			return nil
		}
	}

	guild.Emojis = append(guild.Emojis, emoji)
	return nil
}

// EmojisAdd adds multiple emojis to the world state.
func (s *State) EmojisAdd(guildID string, emojis []*Emoji) error {
	for _, e := range emojis {
		if err := s.EmojiAdd(guildID, e); err != nil {
			return err
		}
	}
	return nil
}

// MessageAdd adds a message to the current world state, or updates it if it exists.
// If the channel cannot be found, the message is discarded.
// Messages are kept in state up to s.MaxMessageCount
func (s *State) MessageAdd(message *Message) error {
	if s == nil {
		return ErrNilState
	}

	c, err := s.Channel(message.ChannelID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	// If the message exists, replace it.
	for i, m := range c.Messages {
		if m.ID == message.ID {
			c.Messages[i] = message
			return nil
		}
	}

	c.Messages = append(c.Messages, message)

	if len(c.Messages) > s.MaxMessageCount {
		s.Unlock()
		for len(c.Messages) > s.MaxMessageCount {
			err := s.MessageRemove(c.Messages[0])
			if err != nil {
				fmt.Println("message remove error: ", err)
			}
		}
		s.Lock()
	}
	return nil
}

// MessageRemove removes a message from the world state.
func (s *State) MessageRemove(message *Message) error {
	if s == nil {
		return ErrNilState
	}

	c, err := s.Channel(message.ChannelID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	for i, m := range c.Messages {
		if m.ID == message.ID {
			c.Messages = append(c.Messages[:i], c.Messages[i+1:]...)
			return nil
		}
	}

	return errors.New("Message not found.")
}

func (s *State) voiceStateUpdate(update *VoiceStateUpdate) error {
	guild, err := s.Guild(update.GuildID)
	if err != nil {
		return err
	}

	s.Lock()
	defer s.Unlock()

	// Handle Leaving Channel
	if update.ChannelID == "" {
		for i, state := range guild.VoiceStates {
			if state.UserID == update.UserID {
				guild.VoiceStates = append(guild.VoiceStates[:i], guild.VoiceStates[i+1:]...)
				return nil
			}
		}
	} else {
		for i, state := range guild.VoiceStates {
			if state.UserID == update.UserID {
				guild.VoiceStates[i] = update.VoiceState
				return nil
			}
		}

		guild.VoiceStates = append(guild.VoiceStates, update.VoiceState)
	}

	return nil
}

// Message gets a message by channel and message ID.
func (s *State) Message(channelID, messageID string) (*Message, error) {
	if s == nil {
		return nil, ErrNilState
	}

	c, err := s.Channel(channelID)
	if err != nil {
		return nil, err
	}

	s.RLock()
	defer s.RUnlock()

	for _, m := range c.Messages {
		if m.ID == messageID {
			return m, nil
		}
	}

	return nil, errors.New("Message not found.")
}

// onInterface handles all events related to states.
func (s *State) onInterface(se *Session, i interface{}) (err error) {
	if s == nil {
		return ErrNilState
	}
	if !se.StateEnabled {
		return nil
	}

	switch t := i.(type) {
	case *Ready:
		err = s.OnReady(t)
	case *GuildCreate:
		err = s.GuildAdd(t.Guild)
	case *GuildUpdate:
		err = s.GuildAdd(t.Guild)
	case *GuildDelete:
		err = s.GuildRemove(t.Guild)
	case *GuildMemberAdd:
		err = s.MemberAdd(t.Member)
	case *GuildMemberUpdate:
		err = s.MemberAdd(t.Member)
	case *GuildMemberRemove:
		err = s.MemberRemove(t.Member)
	case *GuildEmojisUpdate:
		err = s.EmojisAdd(t.GuildID, t.Emojis)
	case *ChannelCreate:
		err = s.ChannelAdd(t.Channel)
	case *ChannelUpdate:
		err = s.ChannelAdd(t.Channel)
	case *ChannelDelete:
		err = s.ChannelRemove(t.Channel)
	case *MessageCreate:
		err = s.MessageAdd(t.Message)
	case *MessageUpdate:
		err = s.MessageAdd(t.Message)
	case *MessageDelete:
		err = s.MessageRemove(t.Message)
	case *VoiceStateUpdate:
		err = s.voiceStateUpdate(t)
	}

	return
}
<|endoftext|>


<commit_before>// Package noise implements the Noise Protocol Framework.
//
// Noise is a low-level framework for building crypto protocols. Noise protocols
// support mutual and optional authentication, identity hiding, forward secrecy,
// zero round-trip encryption, and other advanced features. For more details,
// visit http://noiseprotocol.org.
package noise

import (
	"crypto/rand"
	"errors"
	"io"
)

// A CipherState provides symmetric encryption and decryption after a successful
// handshake.
type CipherState struct {
	cs CipherSuite
	c  Cipher
	k  [32]byte
	n  uint64

	invalid bool
}

// Encrypt encrypts the plaintext and then appends the ciphertext and an
// authentication tag across the ciphertext and optional authenticated data to
// out. This method automatically increments the nonce after every call, so
// messages must be decrypted in the same order.
func (s *CipherState) Encrypt(out, ad, plaintext []byte) []byte {
	if s.invalid {
		panic("noise: CipherSuite has been copied, state is invalid")
	}
	out = s.c.Encrypt(out, s.n, ad, plaintext)
	s.n++
	return out
}

// Decrypt checks the authenticity of the ciphertext and authenticated data and
// then decrypts and appends the plaintext to out. This method automatically
// increments the nonce after every call, messages must be provided in the same
// order that they were encrypted with no missing messages.
func (s *CipherState) Decrypt(out, ad, ciphertext []byte) ([]byte, error) {
	if s.invalid {
		panic("noise: CipherSuite has been copied, state is invalid")
	}
	out, err := s.c.Decrypt(out, s.n, ad, ciphertext)
	s.n++
	return out, err
}

// Cipher returns the low-level symmetric encryption primitive. It should only
// be used if nonces need to be managed manually, for example with a network
// protocol that can deliver out-of-order messages. This is dangerous, users
// must ensure that they are incrementing a nonce after every encrypt operation.
// After calling this method, it is an error to call Encrypt/Decrypt on the
// CipherState.
func (s *CipherState) Cipher() Cipher {
	s.invalid = true
	return s.c
}

type symmetricState struct {
	CipherState
	hasK   bool
	hasPSK bool
	ck     []byte
	h      []byte
}

func (s *symmetricState) InitializeSymmetric(handshakeName []byte) {
	h := s.cs.Hash()
	if len(handshakeName) <= h.Size() {
		s.h = make([]byte, h.Size())
		copy(s.h, handshakeName)
	} else {
		h.Write(handshakeName)
		s.h = h.Sum(nil)
	}
	s.ck = make([]byte, len(s.h))
	copy(s.ck, s.h)
}

func (s *symmetricState) MixKey(dhOutput []byte) {
	s.n = 0
	s.hasK = true
	var hk []byte
	s.ck, hk = hkdf(s.cs.Hash, s.ck[:0], s.k[:0], s.ck, dhOutput)
	copy(s.k[:], hk)
	s.c = s.cs.Cipher(s.k)
}

func (s *symmetricState) MixHash(data []byte) {
	h := s.cs.Hash()
	h.Write(s.h)
	h.Write(data)
	s.h = h.Sum(s.h[:0])
}

func (s *symmetricState) MixPresharedKey(presharedKey []byte) {
	var temp []byte
	s.ck, temp = hkdf(s.cs.Hash, s.ck[:0], nil, s.ck, presharedKey)
	s.MixHash(temp)
	s.hasPSK = true
}

func (s *symmetricState) EncryptAndHash(out, plaintext []byte) []byte {
	if !s.hasK {
		s.MixHash(plaintext)
		return append(out, plaintext...)
	}
	ciphertext := s.Encrypt(out, s.h, plaintext)
	s.MixHash(ciphertext[len(out):])
	return ciphertext
}

func (s *symmetricState) DecryptAndHash(out, data []byte) ([]byte, error) {
	if !s.hasK {
		s.MixHash(data)
		return append(out, data...), nil
	}
	plaintext, err := s.Decrypt(out, s.h, data)
	if err != nil {
		return nil, err
	}
	s.MixHash(data)
	return plaintext, nil
}

func (s *symmetricState) Split() (*CipherState, *CipherState) {
	s1, s2 := &CipherState{cs: s.cs}, &CipherState{cs: s.cs}
	hk1, hk2 := hkdf(s.cs.Hash, s1.k[:0], s2.k[:0], s.ck, nil)
	copy(s1.k[:], hk1)
	copy(s2.k[:], hk2)
	s1.c = s.cs.Cipher(s1.k)
	s2.c = s.cs.Cipher(s2.k)
	return s1, s2
}

// A MessagePattern is a single message or operation used in a Noise handshake.
type MessagePattern int

// A HandshakePattern is a list of messages and operations that are used to
// perform a specific Noise handshake.
type HandshakePattern struct {
	Name                 string
	InitiatorPreMessages []MessagePattern
	ResponderPreMessages []MessagePattern
	Messages             [][]MessagePattern
}

const (
	MessagePatternS MessagePattern = iota
	MessagePatternE
	MessagePatternDHEE
	MessagePatternDHES
	MessagePatternDHSE
	MessagePatternDHSS
)

// MaxMsgLen is the maximum number of bytes that can be sent in a single Noise
// message.
const MaxMsgLen = 65535

// A HandshakeState tracks the state of a Noise handshake. It may be discarded
// after the handshake is complete.
type HandshakeState struct {
	ss              symmetricState
	s               DHKey  // local static keypair
	e               DHKey  // local ephemeral keypair
	rs              []byte // remote party's static public key
	re              []byte // remote party's ephemeral public key
	messagePatterns [][]MessagePattern
	shouldWrite     bool
	msgIdx          int
	rng             io.Reader
}

// A Config provides the details necessary to process a Noise handshake. It is
// never modified by this package, and can be reused.
type Config struct {
	// CipherSuite is the set of cryptographic primitives that will be used.
	CipherSuite CipherSuite

	// Random is the source for cryptographically appropriate random bytes. If
	// zero, it is automatically configured.
	Random io.Reader

	// Pattern is the pattern for the handshake.
	Pattern HandshakePattern

	// Initiator must be true if the first message in the handshake will be sent
	// by this peer.
	Initiator bool

	// Prologue is an optional message that has already be communicated and must
	// be identical on both sides for the handshake to succeed.
	Prologue []byte

	// PresharedKey is the optional pre-shared key for the handshake.
	PresharedKey []byte

	// StaticKeypair is this peer's static keypair, required if part of the
	// handshake.
	StaticKeypair DHKey

	// EphemeralKeypair is this peer's ephemeral keypair that was provided as
	// a pre-message in the handshake.
	EphemeralKeypair DHKey

	// PeerStatic is the static public key of the remote peer that was provided
	// as a pre-message in the handshake.
	PeerStatic []byte

	// PeerEphemeral is the ephemeral public key of the remote peer that was
	// provided as a pre-message in the handshake.
	PeerEphemeral []byte
}

// NewHandshakeState starts a new handshake using the provided configuration.
func NewHandshakeState(c Config) *HandshakeState {
	hs := &HandshakeState{
		s:               c.StaticKeypair,
		e:               c.EphemeralKeypair,
		rs:              c.PeerStatic,
		messagePatterns: c.Pattern.Messages,
		shouldWrite:     c.Initiator,
		rng:             c.Random,
	}
	if hs.rng == nil {
		hs.rng = rand.Reader
	}
	if len(c.PeerEphemeral) > 0 {
		hs.re = make([]byte, len(c.PeerEphemeral))
		copy(hs.re, c.PeerEphemeral)
	}
	hs.ss.cs = c.CipherSuite
	namePrefix := "Noise_"
	if len(c.PresharedKey) > 0 {
		namePrefix = "NoisePSK_"
	}
	hs.ss.InitializeSymmetric([]byte(namePrefix + c.Pattern.Name + "_" + string(hs.ss.cs.Name())))
	hs.ss.MixHash(c.Prologue)
	if len(c.PresharedKey) > 0 {
		hs.ss.MixPresharedKey(c.PresharedKey)
	}
	for _, m := range c.Pattern.InitiatorPreMessages {
		switch {
		case c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.s.Public)
		case c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.e.Public)
		case !c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.rs)
		case !c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.re)
		}
	}
	for _, m := range c.Pattern.ResponderPreMessages {
		switch {
		case !c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.s.Public)
		case !c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.e.Public)
		case c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.rs)
		case c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.re)
		}
	}
	return hs
}

// WriteMessage appends a handshake message to out. The message will include the
// optional payload if provided. If the handshake is completed by the call, two
// CipherStates will be returned, one is used for encryption of messages to the
// remote peer, the other is used for decryption of messages from the remote
// peer. It is an error to call this method out of sync with the handshake
// pattern.
func (s *HandshakeState) WriteMessage(out, payload []byte) ([]byte, *CipherState, *CipherState) {
	if !s.shouldWrite {
		panic("noise: unexpected call to WriteMessage should be ReadMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}
	if len(payload) > MaxMsgLen {
		panic("noise: message is too long")
	}

	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE:
			s.e = s.ss.cs.GenerateKeypair(s.rng)
			out = append(out, s.e.Public...)
			s.ss.MixHash(s.e.Public)
			if s.ss.hasPSK {
				s.ss.MixKey(s.e.Public)
			}
		case MessagePatternS:
			if len(s.s.Public) == 0 {
				panic("noise: invalid state, s.Public is nil")
			}
			out = s.ss.EncryptAndHash(out, s.s.Public)
		case MessagePatternDHEE:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSE:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSS:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = false
	s.msgIdx++
	out = s.ss.EncryptAndHash(out, payload)

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.ss.Split()
		return out, cs1, cs2
	}

	return out, nil, nil
}

// ErrShortMessage is returned by ReadMessage if a message is not as long as it should be.
var ErrShortMessage = errors.New("noise: message is too short")

// ReadMessage processes a received handshake message and appends the payload,
// if any to out. If the handshake is completed by the call, two CipherStates
// will be returned, one is used for encryption of messages to the remote peer,
// the other is used for decryption of messages from the remote peer. It is an
// error to call this method out of sync with the handshake pattern.
func (s *HandshakeState) ReadMessage(out, message []byte) ([]byte, *CipherState, *CipherState, error) {
	if s.shouldWrite {
		panic("noise: unexpected call to ReadMessage should be WriteMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}

	var err error
	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE, MessagePatternS:
			expected := s.ss.cs.DHLen()
			if msg == MessagePatternS && s.ss.hasK {
				expected += 16
			}
			if len(message) < expected {
				return nil, nil, nil, ErrShortMessage
			}
			switch msg {
			case MessagePatternE:
				if cap(s.re) < s.ss.cs.DHLen() {
					s.re = make([]byte, s.ss.cs.DHLen())
				}
				s.re = s.re[:s.ss.cs.DHLen()]
				copy(s.re, message)
				s.ss.MixHash(s.re)
				if s.ss.hasPSK {
					s.ss.MixKey(s.re)
				}
			case MessagePatternS:
				if len(s.rs) > 0 {
					panic("noise: invalid state, rs is not nil")
				}
				s.rs, err = s.ss.DecryptAndHash(s.rs[:0], message[:expected])
			}
			if err != nil {
				return nil, nil, nil, err
			}
			message = message[expected:]
		case MessagePatternDHEE:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSE:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSS:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = true
	s.msgIdx++
	out, err = s.ss.DecryptAndHash(out, message)
	if err != nil {
		return nil, nil, nil, err
	}

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.ss.Split()
		return out, cs1, cs2, nil
	}

	return out, nil, nil, nil
}

// ChannelBinding provides a value that uniquely identifies the session and can
// be used as a channel binding. It is an error to call this method before the
// handshake is complete.
func (s *HandshakeState) ChannelBinding() []byte {
	return s.ss.h
}
<commit_msg>Add PeerStatic method to HandshakeState<commit_after>// Package noise implements the Noise Protocol Framework.
//
// Noise is a low-level framework for building crypto protocols. Noise protocols
// support mutual and optional authentication, identity hiding, forward secrecy,
// zero round-trip encryption, and other advanced features. For more details,
// visit http://noiseprotocol.org.
package noise

import (
	"crypto/rand"
	"errors"
	"io"
)

// A CipherState provides symmetric encryption and decryption after a successful
// handshake.
type CipherState struct {
	cs CipherSuite
	c  Cipher
	k  [32]byte
	n  uint64

	invalid bool
}

// Encrypt encrypts the plaintext and then appends the ciphertext and an
// authentication tag across the ciphertext and optional authenticated data to
// out. This method automatically increments the nonce after every call, so
// messages must be decrypted in the same order.
func (s *CipherState) Encrypt(out, ad, plaintext []byte) []byte {
	if s.invalid {
		panic("noise: CipherSuite has been copied, state is invalid")
	}
	out = s.c.Encrypt(out, s.n, ad, plaintext)
	s.n++
	return out
}

// Decrypt checks the authenticity of the ciphertext and authenticated data and
// then decrypts and appends the plaintext to out. This method automatically
// increments the nonce after every call, messages must be provided in the same
// order that they were encrypted with no missing messages.
func (s *CipherState) Decrypt(out, ad, ciphertext []byte) ([]byte, error) {
	if s.invalid {
		panic("noise: CipherSuite has been copied, state is invalid")
	}
	out, err := s.c.Decrypt(out, s.n, ad, ciphertext)
	s.n++
	return out, err
}

// Cipher returns the low-level symmetric encryption primitive. It should only
// be used if nonces need to be managed manually, for example with a network
// protocol that can deliver out-of-order messages. This is dangerous, users
// must ensure that they are incrementing a nonce after every encrypt operation.
// After calling this method, it is an error to call Encrypt/Decrypt on the
// CipherState.
func (s *CipherState) Cipher() Cipher {
	s.invalid = true
	return s.c
}

type symmetricState struct {
	CipherState
	hasK   bool
	hasPSK bool
	ck     []byte
	h      []byte
}

func (s *symmetricState) InitializeSymmetric(handshakeName []byte) {
	h := s.cs.Hash()
	if len(handshakeName) <= h.Size() {
		s.h = make([]byte, h.Size())
		copy(s.h, handshakeName)
	} else {
		h.Write(handshakeName)
		s.h = h.Sum(nil)
	}
	s.ck = make([]byte, len(s.h))
	copy(s.ck, s.h)
}

func (s *symmetricState) MixKey(dhOutput []byte) {
	s.n = 0
	s.hasK = true
	var hk []byte
	s.ck, hk = hkdf(s.cs.Hash, s.ck[:0], s.k[:0], s.ck, dhOutput)
	copy(s.k[:], hk)
	s.c = s.cs.Cipher(s.k)
}

func (s *symmetricState) MixHash(data []byte) {
	h := s.cs.Hash()
	h.Write(s.h)
	h.Write(data)
	s.h = h.Sum(s.h[:0])
}

func (s *symmetricState) MixPresharedKey(presharedKey []byte) {
	var temp []byte
	s.ck, temp = hkdf(s.cs.Hash, s.ck[:0], nil, s.ck, presharedKey)
	s.MixHash(temp)
	s.hasPSK = true
}

func (s *symmetricState) EncryptAndHash(out, plaintext []byte) []byte {
	if !s.hasK {
		s.MixHash(plaintext)
		return append(out, plaintext...)
	}
	ciphertext := s.Encrypt(out, s.h, plaintext)
	s.MixHash(ciphertext[len(out):])
	return ciphertext
}

func (s *symmetricState) DecryptAndHash(out, data []byte) ([]byte, error) {
	if !s.hasK {
		s.MixHash(data)
		return append(out, data...), nil
	}
	plaintext, err := s.Decrypt(out, s.h, data)
	if err != nil {
		return nil, err
	}
	s.MixHash(data)
	return plaintext, nil
}

func (s *symmetricState) Split() (*CipherState, *CipherState) {
	s1, s2 := &CipherState{cs: s.cs}, &CipherState{cs: s.cs}
	hk1, hk2 := hkdf(s.cs.Hash, s1.k[:0], s2.k[:0], s.ck, nil)
	copy(s1.k[:], hk1)
	copy(s2.k[:], hk2)
	s1.c = s.cs.Cipher(s1.k)
	s2.c = s.cs.Cipher(s2.k)
	return s1, s2
}

// A MessagePattern is a single message or operation used in a Noise handshake.
type MessagePattern int

// A HandshakePattern is a list of messages and operations that are used to
// perform a specific Noise handshake.
type HandshakePattern struct {
	Name                 string
	InitiatorPreMessages []MessagePattern
	ResponderPreMessages []MessagePattern
	Messages             [][]MessagePattern
}

const (
	MessagePatternS MessagePattern = iota
	MessagePatternE
	MessagePatternDHEE
	MessagePatternDHES
	MessagePatternDHSE
	MessagePatternDHSS
)

// MaxMsgLen is the maximum number of bytes that can be sent in a single Noise
// message.
const MaxMsgLen = 65535

// A HandshakeState tracks the state of a Noise handshake. It may be discarded
// after the handshake is complete.
type HandshakeState struct {
	ss              symmetricState
	s               DHKey  // local static keypair
	e               DHKey  // local ephemeral keypair
	rs              []byte // remote party's static public key
	re              []byte // remote party's ephemeral public key
	messagePatterns [][]MessagePattern
	shouldWrite     bool
	msgIdx          int
	rng             io.Reader
}

// A Config provides the details necessary to process a Noise handshake. It is
// never modified by this package, and can be reused.
type Config struct {
	// CipherSuite is the set of cryptographic primitives that will be used.
	CipherSuite CipherSuite

	// Random is the source for cryptographically appropriate random bytes. If
	// zero, it is automatically configured.
	Random io.Reader

	// Pattern is the pattern for the handshake.
	Pattern HandshakePattern

	// Initiator must be true if the first message in the handshake will be sent
	// by this peer.
	Initiator bool

	// Prologue is an optional message that has already be communicated and must
	// be identical on both sides for the handshake to succeed.
	Prologue []byte

	// PresharedKey is the optional pre-shared key for the handshake.
	PresharedKey []byte

	// StaticKeypair is this peer's static keypair, required if part of the
	// handshake.
	StaticKeypair DHKey

	// EphemeralKeypair is this peer's ephemeral keypair that was provided as
	// a pre-message in the handshake.
	EphemeralKeypair DHKey

	// PeerStatic is the static public key of the remote peer that was provided
	// as a pre-message in the handshake.
	PeerStatic []byte

	// PeerEphemeral is the ephemeral public key of the remote peer that was
	// provided as a pre-message in the handshake.
	PeerEphemeral []byte
}

// NewHandshakeState starts a new handshake using the provided configuration.
func NewHandshakeState(c Config) *HandshakeState {
	hs := &HandshakeState{
		s:               c.StaticKeypair,
		e:               c.EphemeralKeypair,
		rs:              c.PeerStatic,
		messagePatterns: c.Pattern.Messages,
		shouldWrite:     c.Initiator,
		rng:             c.Random,
	}
	if hs.rng == nil {
		hs.rng = rand.Reader
	}
	if len(c.PeerEphemeral) > 0 {
		hs.re = make([]byte, len(c.PeerEphemeral))
		copy(hs.re, c.PeerEphemeral)
	}
	hs.ss.cs = c.CipherSuite
	namePrefix := "Noise_"
	if len(c.PresharedKey) > 0 {
		namePrefix = "NoisePSK_"
	}
	hs.ss.InitializeSymmetric([]byte(namePrefix + c.Pattern.Name + "_" + string(hs.ss.cs.Name())))
	hs.ss.MixHash(c.Prologue)
	if len(c.PresharedKey) > 0 {
		hs.ss.MixPresharedKey(c.PresharedKey)
	}
	for _, m := range c.Pattern.InitiatorPreMessages {
		switch {
		case c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.s.Public)
		case c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.e.Public)
		case !c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.rs)
		case !c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.re)
		}
	}
	for _, m := range c.Pattern.ResponderPreMessages {
		switch {
		case !c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.s.Public)
		case !c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.e.Public)
		case c.Initiator && m == MessagePatternS:
			hs.ss.MixHash(hs.rs)
		case c.Initiator && m == MessagePatternE:
			hs.ss.MixHash(hs.re)
		}
	}
	return hs
}

// WriteMessage appends a handshake message to out. The message will include the
// optional payload if provided. If the handshake is completed by the call, two
// CipherStates will be returned, one is used for encryption of messages to the
// remote peer, the other is used for decryption of messages from the remote
// peer. It is an error to call this method out of sync with the handshake
// pattern.
func (s *HandshakeState) WriteMessage(out, payload []byte) ([]byte, *CipherState, *CipherState) {
	if !s.shouldWrite {
		panic("noise: unexpected call to WriteMessage should be ReadMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}
	if len(payload) > MaxMsgLen {
		panic("noise: message is too long")
	}

	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE:
			s.e = s.ss.cs.GenerateKeypair(s.rng)
			out = append(out, s.e.Public...)
			s.ss.MixHash(s.e.Public)
			if s.ss.hasPSK {
				s.ss.MixKey(s.e.Public)
			}
		case MessagePatternS:
			if len(s.s.Public) == 0 {
				panic("noise: invalid state, s.Public is nil")
			}
			out = s.ss.EncryptAndHash(out, s.s.Public)
		case MessagePatternDHEE:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSE:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSS:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = false
	s.msgIdx++
	out = s.ss.EncryptAndHash(out, payload)

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.ss.Split()
		return out, cs1, cs2
	}

	return out, nil, nil
}

// ErrShortMessage is returned by ReadMessage if a message is not as long as it should be.
var ErrShortMessage = errors.New("noise: message is too short")

// ReadMessage processes a received handshake message and appends the payload,
// if any to out. If the handshake is completed by the call, two CipherStates
// will be returned, one is used for encryption of messages to the remote peer,
// the other is used for decryption of messages from the remote peer. It is an
// error to call this method out of sync with the handshake pattern.
func (s *HandshakeState) ReadMessage(out, message []byte) ([]byte, *CipherState, *CipherState, error) {
	if s.shouldWrite {
		panic("noise: unexpected call to ReadMessage should be WriteMessage")
	}
	if s.msgIdx > len(s.messagePatterns)-1 {
		panic("noise: no handshake messages left")
	}

	var err error
	for _, msg := range s.messagePatterns[s.msgIdx] {
		switch msg {
		case MessagePatternE, MessagePatternS:
			expected := s.ss.cs.DHLen()
			if msg == MessagePatternS && s.ss.hasK {
				expected += 16
			}
			if len(message) < expected {
				return nil, nil, nil, ErrShortMessage
			}
			switch msg {
			case MessagePatternE:
				if cap(s.re) < s.ss.cs.DHLen() {
					s.re = make([]byte, s.ss.cs.DHLen())
				}
				s.re = s.re[:s.ss.cs.DHLen()]
				copy(s.re, message)
				s.ss.MixHash(s.re)
				if s.ss.hasPSK {
					s.ss.MixKey(s.re)
				}
			case MessagePatternS:
				if len(s.rs) > 0 {
					panic("noise: invalid state, rs is not nil")
				}
				s.rs, err = s.ss.DecryptAndHash(s.rs[:0], message[:expected])
			}
			if err != nil {
				return nil, nil, nil, err
			}
			message = message[expected:]
		case MessagePatternDHEE:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.re))
		case MessagePatternDHES:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.re))
		case MessagePatternDHSE:
			s.ss.MixKey(s.ss.cs.DH(s.e.Private, s.rs))
		case MessagePatternDHSS:
			s.ss.MixKey(s.ss.cs.DH(s.s.Private, s.rs))
		}
	}
	s.shouldWrite = true
	s.msgIdx++
	out, err = s.ss.DecryptAndHash(out, message)
	if err != nil {
		return nil, nil, nil, err
	}

	if s.msgIdx >= len(s.messagePatterns) {
		cs1, cs2 := s.ss.Split()
		return out, cs1, cs2, nil
	}

	return out, nil, nil, nil
}

// ChannelBinding provides a value that uniquely identifies the session and can
// be used as a channel binding. It is an error to call this method before the
// handshake is complete.
func (s *HandshakeState) ChannelBinding() []byte {
	return s.ss.h
}

// PeerStatic returns the static key provided by the remote peer during
// a handshake. It is an error to call this method if a handshake message
// containing a static key has not been read.
func (s *HandshakeState) PeerStatic() []byte {
	return s.rs
}
<|endoftext|>


<commit_before>package beanstalk

import (
	"strings"
	"sync"

	"gopkg.in/yaml.v2"
)

// TubeStat describes the statistics of a beanstalk tube.
type TubeStat struct {
	Name           string `yaml:"name"                  json:"name"`
	UrgentJobs     int    `yaml:"current-jobs-urgent"   json:"urgentJobs"`
	ReadyJobs      int    `yaml:"current-jobs-ready"    json:"readyJobs"`
	ReservedJobs   int    `yaml:"current-jobs-reserved" json:"reservedJobs"`
	DelayedJobs    int    `yaml:"current-jobs-delayed"  json:"delayedJobs"`
	BuriedJobs     int    `yaml:"current-jobs-buried"   json:"buriedJobs"`
	TotalJobs      int64  `yaml:"total-jobs"             json:"totalJobs"`
	Using          int    `yaml:"current-using"         json:"using"`
	Watching       int    `yaml:"curent-watching"       json:"watching"`
	Waiting        int    `yaml:"current-waiting"       json:"waiting"`
	DeleteCommands int64  `yaml:"cmd-delete"            json:"deleteCommands"`
	PauseCommands  int64  `yaml:"cmd-pause-tube"        json:"pauseCommands"`
}

type stats struct {
	TubeStats []TubeStat
	err       error
}

// TubeStats returns a list of beanstalk tube statistics.
func TubeStats(urls []string, options *Options, tubePrefix string) ([]TubeStat, error) {
	if options == nil {
		options = DefaultOptions()
	}

	var wg sync.WaitGroup
	wg.Add(len(urls))

	// Request the beanstalk tube statistics for all the beanstalk servers.
	results := make([]stats, len(urls))
	for i, url := range urls {
		go func(i int, url string) {
			defer wg.Done()

			err := func() error {
				// Dial into the beanstalk server.
				conn, err := dial(url, options)
				if err != nil {
					return err
				}
				defer conn.Close()

				// Fetch the list of tubes on this server.
				client := NewClient(conn, options)
				_, data, err := client.requestResponse("list-tubes")
				if err != nil {
					return err
				}

				var tubes []string
				if err = yaml.Unmarshal(data, &tubes); err != nil {
					return err
				}

				// Cycle over the list of tube names and skip any tube that doesn't
				// match the specified prefix.
				var tubeStats []TubeStat
				for _, tube := range tubes {
					if !strings.HasPrefix(tube, tubePrefix) {
						continue
					}

					// Fetch the statistics for this tube.
					if _, data, err = client.requestResponse("stats-tube %s", tube); err != nil {
						return err
					}

					var tubeStat TubeStat
					if err = yaml.Unmarshal(data, &tubeStat); err != nil {
						return err
					}

					tubeStats = append(tubeStats, tubeStat)
				}

				results[i] = stats{TubeStats: tubeStats}
				return nil
			}()

			if err != nil {
				results[i] = stats{err: err}
			}
		}(i, url)
	}

	wg.Wait()

	// Return the first error it encounters.
	for _, result := range results {
		if result.err != nil {
			return []TubeStat{}, result.err
		}
	}

	// Collect and merge the results.
	var tubeStats []TubeStat
	for _, result := range results {
		for _, tubeStat := range result.TubeStats {
			found := false

			for i, stat := range tubeStats {
				if stat.Name == tubeStat.Name {
					tubeStats[i].UrgentJobs += tubeStat.UrgentJobs
					tubeStats[i].ReadyJobs += tubeStat.ReadyJobs
					tubeStats[i].ReservedJobs += tubeStat.ReservedJobs
					tubeStats[i].DelayedJobs += tubeStat.DelayedJobs
					tubeStats[i].BuriedJobs += tubeStat.BuriedJobs
					tubeStats[i].TotalJobs += tubeStat.TotalJobs
					tubeStats[i].Using += tubeStat.Using
					tubeStats[i].Watching += tubeStat.Watching
					tubeStats[i].DeleteCommands += tubeStat.DeleteCommands
					tubeStats[i].PauseCommands += tubeStat.PauseCommands
					found = true
				}
			}

			if !found {
				tubeStats = append(tubeStats, tubeStat)
			}
		}
	}

	return tubeStats, nil
}
<commit_msg>Add another typo<commit_after>package beanstalk

import (
	"strings"
	"sync"

	"gopkg.in/yaml.v2"
)

// TubeStat describes the statistics of a beanstalk tube.
type TubeStat struct {
	Name           string `yaml:"name"                  json:"name"`
	UrgentJobs     int    `yaml:"current-jobs-urgent"   json:"urgentJobs"`
	ReadyJobs      int    `yaml:"current-jobs-ready"    json:"readyJobs"`
	ReservedJobs   int    `yaml:"current-jobs-reserved" json:"reservedJobs"`
	DelayedJobs    int    `yaml:"current-jobs-delayed"  json:"delayedJobs"`
	BuriedJobs     int    `yaml:"current-jobs-buried"   json:"buriedJobs"`
	TotalJobs      int64  `yaml:"total-jobs"            json:"totalJobs"`
	Using          int    `yaml:"current-using"         json:"using"`
	Watching       int    `yaml:"current-watching"      json:"watching"`
	Waiting        int    `yaml:"current-waiting"       json:"waiting"`
	DeleteCommands int64  `yaml:"cmd-delete"            json:"deleteCommands"`
	PauseCommands  int64  `yaml:"cmd-pause-tube"        json:"pauseCommands"`
}

type stats struct {
	TubeStats []TubeStat
	err       error
}

// TubeStats returns a list of beanstalk tube statistics.
func TubeStats(urls []string, options *Options, tubePrefix string) ([]TubeStat, error) {
	if options == nil {
		options = DefaultOptions()
	}

	var wg sync.WaitGroup
	wg.Add(len(urls))

	// Request the beanstalk tube statistics for all the beanstalk servers.
	results := make([]stats, len(urls))
	for i, url := range urls {
		go func(i int, url string) {
			defer wg.Done()

			err := func() error {
				// Dial into the beanstalk server.
				conn, err := dial(url, options)
				if err != nil {
					return err
				}
				defer conn.Close()

				// Fetch the list of tubes on this server.
				client := NewClient(conn, options)
				_, data, err := client.requestResponse("list-tubes")
				if err != nil {
					return err
				}

				var tubes []string
				if err = yaml.Unmarshal(data, &tubes); err != nil {
					return err
				}

				// Cycle over the list of tube names and skip any tube that doesn't
				// match the specified prefix.
				var tubeStats []TubeStat
				for _, tube := range tubes {
					if !strings.HasPrefix(tube, tubePrefix) {
						continue
					}

					// Fetch the statistics for this tube.
					if _, data, err = client.requestResponse("stats-tube %s", tube); err != nil {
						return err
					}

					var tubeStat TubeStat
					if err = yaml.Unmarshal(data, &tubeStat); err != nil {
						return err
					}

					tubeStats = append(tubeStats, tubeStat)
				}

				results[i] = stats{TubeStats: tubeStats}
				return nil
			}()

			if err != nil {
				results[i] = stats{err: err}
			}
		}(i, url)
	}

	wg.Wait()

	// Return the first error it encounters.
	for _, result := range results {
		if result.err != nil {
			return []TubeStat{}, result.err
		}
	}

	// Collect and merge the results.
	var tubeStats []TubeStat
	for _, result := range results {
		for _, tubeStat := range result.TubeStats {
			found := false

			for i, stat := range tubeStats {
				if stat.Name == tubeStat.Name {
					tubeStats[i].UrgentJobs += tubeStat.UrgentJobs
					tubeStats[i].ReadyJobs += tubeStat.ReadyJobs
					tubeStats[i].ReservedJobs += tubeStat.ReservedJobs
					tubeStats[i].DelayedJobs += tubeStat.DelayedJobs
					tubeStats[i].BuriedJobs += tubeStat.BuriedJobs
					tubeStats[i].TotalJobs += tubeStat.TotalJobs
					tubeStats[i].Using += tubeStat.Using
					tubeStats[i].Watching += tubeStat.Watching
					tubeStats[i].DeleteCommands += tubeStat.DeleteCommands
					tubeStats[i].PauseCommands += tubeStat.PauseCommands
					found = true
				}
			}

			if !found {
				tubeStats = append(tubeStats, tubeStat)
			}
		}
	}

	return tubeStats, nil
}
<|endoftext|>


<commit_before>package appdash

import (
	"encoding/gob"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"
)

// A Store stores and retrieves spans.
type Store interface {
	Collector

	// Trace gets a trace (a tree of spans) given its trace ID. If no
	// such trace exists, ErrTraceNotFound is returned.
	Trace(ID) (*Trace, error)
}

var (
	// ErrTraceNotFound is returned by Store.GetTrace when no trace is
	// found with the given ID.
	ErrTraceNotFound = errors.New("trace not found")
)

// A Queryer indexes spans and makes them queryable.
type Queryer interface {
	// Traces returns an implementation-defined list of traces. It is
	// a placeholder method that will be removed when other, more
	// useful methods are added to Queryer.
	Traces() ([]*Trace, error)
}

// NewMemoryStore creates a new in-memory store
func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		trace: map[ID]*Trace{},
		span:  map[ID]map[ID]*Trace{},
	}
}

// A MemoryStore is an in-memory Store.
type MemoryStore struct {
	trace map[ID]*Trace        // trace ID -> trace tree
	span  map[ID]map[ID]*Trace // trace ID -> span ID -> trace (sub)tree

	sync.Mutex // protects trace

	log bool
}

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*MemoryStore)(nil)

// Collect implements the Collector interface by collecting the events that
// occured in the span in-memory.
func (ms *MemoryStore) Collect(id SpanID, as ...Annotation) error {
	ms.Lock()
	defer ms.Unlock()

	if ms.log {
		log.Printf("Collect %v", id)
	}

	// Initialize span map if needed.
	if _, present := ms.span[id.Trace]; !present {
		ms.span[id.Trace] = map[ID]*Trace{}
	}

	// Create or update span.
	s, present := ms.span[id.Trace][id.Span]
	if !present {
		s = &Trace{Span: Span{ID: id, Annotations: as}}
		ms.span[id.Trace][id.Span] = s
	} else {
		if ms.log {
			if len(as) > 0 {
				log.Printf("Add %d annotations to %v", len(as), id)
			}
		}
		s.Annotations = append(s.Annotations, as...)
		return nil
	}

	// Create trace tree if it doesn't already exist.
	root, present := ms.trace[id.Trace]
	if !present {
		// Root span hasn't been seen yet, so make this the temporary
		// root (until we collect the actual root).
		if ms.log {
			if id.IsRoot() {
				log.Printf("Create trace %v root %v", id.Trace, id)
			} else {
				log.Printf("Create temporary trace %v root %v", id.Trace, id)
			}
		}
		ms.trace[id.Trace] = s
		root = s
	}

	// If there's a temp root and we just collected the real
	// root, fix up the tree. Or if we're the temp root's
	// parents, set us up as the new temp root.
	if isRoot, isTempRootParent := id.IsRoot(), root.Span.ID.Parent == id.Span; s != root && (isRoot || isTempRootParent) {
		oldRoot := root
		root = s
		if ms.log {
			if isRoot {
				log.Printf("Set real root %v and move temp root %v", root.Span.ID, oldRoot.Span.ID)
			} else {
				log.Printf("Set new temp root %v and move previous temp root %v (child of new temp root)", root.Span.ID, oldRoot.Span.ID)
			}
		}
		ms.trace[id.Trace] = root // set new root
		ms.reattachChildren(root, oldRoot)
		ms.insert(root, oldRoot) // reinsert the old root

		// Move the old temp root's temp children to the new
		// (possibly temp) root.
		var sub2 []*Trace
		for _, c := range oldRoot.Sub {
			if c.Span.ID.Parent != oldRoot.Span.ID.Span {
				if ms.log {
					log.Printf("Move %v from old root %v to new (possibly temp) root %v", c.Span.ID, oldRoot.Span.ID, root.Span.ID)
				}
				root.Sub = append(root.Sub, c)
			} else {
				sub2 = append(sub2, c)
			}
		}
		oldRoot.Sub = sub2
	}

	// Insert into trace tree. (We inserted the trace root span
	// above.)
	if !id.IsRoot() && s != root {
		ms.insert(root, s)
	}

	// See if we're the parent of any of the root's temporary
	// children.
	if s != root {
		ms.reattachChildren(s, root)
	}

	return nil
}

// insert inserts t into the trace tree whose root (or temp root) is
// root.
func (ms *MemoryStore) insert(root, t *Trace) {
	p, present := ms.span[t.ID.Trace][t.ID.Parent]
	if present {
		if ms.log {
			log.Printf("Add %v as a child of parent %v", t.Span.ID, p.Span.ID)
		}
		p.Sub = append(p.Sub, t)
	} else {
		// Add as temporary child of the root for now. When the
		// real parent is added, we'll fix it up later.
		if ms.log {
			log.Printf("Add %v as a temporary child of root %v", t.Span.ID, root.Span.ID)
		}
		root.Sub = append(root.Sub, t)
	}
}

// reattachChildren moves temporary children of src to dst, if dst is
// the node's parent.
func (ms *MemoryStore) reattachChildren(dst, src *Trace) {
	if dst == src {
		panic("dst == src")
	}
	var sub2 []*Trace
	for _, c := range src.Sub {
		if c.Span.ID.Parent == dst.Span.ID.Span {
			if ms.log {
				log.Printf("Move %v from src %v to dst %v", c.Span.ID, src.Span.ID, dst.Span.ID)
			}
			dst.Sub = append(dst.Sub, c)
		} else {
			sub2 = append(sub2, c)
		}
	}
	src.Sub = sub2
}

// Trace implements the Store interface by returning the Trace (a tree of
// spans) for the given trace span ID or, if no such trace exists, by returning
// ErrTraceNotFound.
func (ms *MemoryStore) Trace(id ID) (*Trace, error) {
	ms.Lock()
	defer ms.Unlock()

	return ms.traceNoLock(id)
}

func (ms *MemoryStore) traceNoLock(id ID) (*Trace, error) {
	t, present := ms.trace[id]
	if !present {
		return nil, ErrTraceNotFound
	}
	return t, nil
}

// Traces implements the Queryer interface.
func (ms *MemoryStore) Traces() ([]*Trace, error) {
	ms.Lock()
	defer ms.Unlock()

	var ts []*Trace
	for id := range ms.trace {
		t, err := ms.traceNoLock(id)
		if err != nil {
			return nil, err
		}
		ts = append(ts, t)
	}
	return ts, nil
}

// Delete implements the DeleteStore interface by deleting the traces given by
// their span ID's from this in-memory store.
func (ms *MemoryStore) Delete(traces ...ID) error {
	ms.Lock()
	defer ms.Unlock()

	for _, id := range traces {
		delete(ms.trace, id)
		delete(ms.span, id)
	}
	return nil
}

type memoryStoreData struct {
	Trace map[ID]*Trace
	Span  map[ID]map[ID]*Trace
}

// Write writes ms's internal data structures.
func (ms *MemoryStore) Write(w io.Writer) error {
	ms.Lock()
	defer ms.Unlock()

	data := memoryStoreData{ms.trace, ms.span}
	return gob.NewEncoder(w).Encode(data)
}

// ReadFrom loads ms's internal data structures from a reader.
func (ms *MemoryStore) ReadFrom(r io.Reader) (int64, error) {
	ms.Lock()
	defer ms.Unlock()

	var data memoryStoreData
	if err := gob.NewDecoder(r).Decode(&data); err != nil {
		return 0, err
	}
	ms.trace = data.Trace
	ms.span = data.Span
	return int64(len(ms.trace)), nil
}

// PersistentStore is a Store that can persist its data and read it
// back in.
type PersistentStore interface {
	Write(io.Writer) error
	ReadFrom(io.Reader) (int64, error)
	Store
}

// PersistEvery persists s's data to a file periodically.
func PersistEvery(s PersistentStore, interval time.Duration, file string) error {
	for {
		time.Sleep(interval)

		f, err := ioutil.TempFile("", "appdash")
		if err != nil {
			return err
		}
		if err := s.Write(f); err != nil {
			f.Close()
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
		if err := os.Rename(f.Name(), file); err != nil {
			return err
		}
	}
}

// A DeleteStore is a Store that can delete traces.
type DeleteStore interface {
	Store

	// Delete deletes traces given their trace IDs.
	Delete(...ID) error
}

// A RecentStore wraps another store and deletes old traces after a
// specified amount of time.
type RecentStore struct {
	// MinEvictAge is the minimum age of a trace before it is evicted.
	MinEvictAge time.Duration

	// DeleteStore is the underlying store that spans are saved to and
	// deleted from.
	DeleteStore

	// Debug is whether to log debug messages.
	Debug bool

	// created maps trace ID to the UnixNano time it was first seen.
	created map[ID]int64

	// lastEvicted is the last time the eviction process was run.
	lastEvicted time.Time

	mu sync.Mutex // mu guards created and lastEvicted
}

// Collect calls the underlying store's Collect and records the time
// that this trace was first seen.
func (rs *RecentStore) Collect(id SpanID, anns ...Annotation) error {
	rs.mu.Lock()
	if rs.created == nil {
		rs.created = map[ID]int64{}
	}
	if _, present := rs.created[id.Trace]; !present {
		rs.created[id.Trace] = time.Now().UnixNano()
	}
	if time.Since(rs.lastEvicted) > rs.MinEvictAge {
		rs.evictBefore(time.Now().Add(-1 * rs.MinEvictAge))
	}
	rs.mu.Unlock()

	return rs.DeleteStore.Collect(id, anns...)
}

// evictBefore evicts traces that were created before t. The rs.mu lock
// must be held while calling evictBefore.
func (rs *RecentStore) evictBefore(t time.Time) {
	evictStart := time.Now()
	tnano := t.UnixNano()
	var toEvict []ID
	for id, ct := range rs.created {
		if ct < tnano {
			toEvict = append(toEvict, id)
			delete(rs.created, id)
		}
	}
	if len(toEvict) == 0 {
		return
	}

	if rs.Debug {
		log.Printf("RecentStore: deleting %d traces created before %s (age check took %s)", len(toEvict), t, time.Since(evictStart))
	}

	// Spawn separate goroutine so we don't hold the rs.mu lock.
	go func() {
		deleteStart := time.Now()
		if err := rs.DeleteStore.Delete(toEvict...); err != nil {
			log.Printf("RecentStore: failed to delete traces: %s", err)
		}
		if rs.Debug {
			log.Printf("RecentStore: finished deleting %d traces created before %s (took %s)", len(toEvict), t, time.Since(deleteStart))
		}
	}()
}
<commit_msg>Document that MemoryStore implements the PersistentStore interface.<commit_after>package appdash

import (
	"encoding/gob"
	"errors"
	"io"
	"io/ioutil"
	"log"
	"os"
	"sync"
	"time"
)

// A Store stores and retrieves spans.
type Store interface {
	Collector

	// Trace gets a trace (a tree of spans) given its trace ID. If no
	// such trace exists, ErrTraceNotFound is returned.
	Trace(ID) (*Trace, error)
}

var (
	// ErrTraceNotFound is returned by Store.GetTrace when no trace is
	// found with the given ID.
	ErrTraceNotFound = errors.New("trace not found")
)

// A Queryer indexes spans and makes them queryable.
type Queryer interface {
	// Traces returns an implementation-defined list of traces. It is
	// a placeholder method that will be removed when other, more
	// useful methods are added to Queryer.
	Traces() ([]*Trace, error)
}

// NewMemoryStore creates a new in-memory store
func NewMemoryStore() *MemoryStore {
	return &MemoryStore{
		trace: map[ID]*Trace{},
		span:  map[ID]map[ID]*Trace{},
	}
}

// A MemoryStore is an in-memory Store that also implements the PersistentStore
// interface.
type MemoryStore struct {
	trace map[ID]*Trace        // trace ID -> trace tree
	span  map[ID]map[ID]*Trace // trace ID -> span ID -> trace (sub)tree

	sync.Mutex // protects trace

	log bool
}

// Compile-time "implements" check.
var _ interface {
	Store
	Queryer
} = (*MemoryStore)(nil)

// Collect implements the Collector interface by collecting the events that
// occured in the span in-memory.
func (ms *MemoryStore) Collect(id SpanID, as ...Annotation) error {
	ms.Lock()
	defer ms.Unlock()

	if ms.log {
		log.Printf("Collect %v", id)
	}

	// Initialize span map if needed.
	if _, present := ms.span[id.Trace]; !present {
		ms.span[id.Trace] = map[ID]*Trace{}
	}

	// Create or update span.
	s, present := ms.span[id.Trace][id.Span]
	if !present {
		s = &Trace{Span: Span{ID: id, Annotations: as}}
		ms.span[id.Trace][id.Span] = s
	} else {
		if ms.log {
			if len(as) > 0 {
				log.Printf("Add %d annotations to %v", len(as), id)
			}
		}
		s.Annotations = append(s.Annotations, as...)
		return nil
	}

	// Create trace tree if it doesn't already exist.
	root, present := ms.trace[id.Trace]
	if !present {
		// Root span hasn't been seen yet, so make this the temporary
		// root (until we collect the actual root).
		if ms.log {
			if id.IsRoot() {
				log.Printf("Create trace %v root %v", id.Trace, id)
			} else {
				log.Printf("Create temporary trace %v root %v", id.Trace, id)
			}
		}
		ms.trace[id.Trace] = s
		root = s
	}

	// If there's a temp root and we just collected the real
	// root, fix up the tree. Or if we're the temp root's
	// parents, set us up as the new temp root.
	if isRoot, isTempRootParent := id.IsRoot(), root.Span.ID.Parent == id.Span; s != root && (isRoot || isTempRootParent) {
		oldRoot := root
		root = s
		if ms.log {
			if isRoot {
				log.Printf("Set real root %v and move temp root %v", root.Span.ID, oldRoot.Span.ID)
			} else {
				log.Printf("Set new temp root %v and move previous temp root %v (child of new temp root)", root.Span.ID, oldRoot.Span.ID)
			}
		}
		ms.trace[id.Trace] = root // set new root
		ms.reattachChildren(root, oldRoot)
		ms.insert(root, oldRoot) // reinsert the old root

		// Move the old temp root's temp children to the new
		// (possibly temp) root.
		var sub2 []*Trace
		for _, c := range oldRoot.Sub {
			if c.Span.ID.Parent != oldRoot.Span.ID.Span {
				if ms.log {
					log.Printf("Move %v from old root %v to new (possibly temp) root %v", c.Span.ID, oldRoot.Span.ID, root.Span.ID)
				}
				root.Sub = append(root.Sub, c)
			} else {
				sub2 = append(sub2, c)
			}
		}
		oldRoot.Sub = sub2
	}

	// Insert into trace tree. (We inserted the trace root span
	// above.)
	if !id.IsRoot() && s != root {
		ms.insert(root, s)
	}

	// See if we're the parent of any of the root's temporary
	// children.
	if s != root {
		ms.reattachChildren(s, root)
	}

	return nil
}

// insert inserts t into the trace tree whose root (or temp root) is
// root.
func (ms *MemoryStore) insert(root, t *Trace) {
	p, present := ms.span[t.ID.Trace][t.ID.Parent]
	if present {
		if ms.log {
			log.Printf("Add %v as a child of parent %v", t.Span.ID, p.Span.ID)
		}
		p.Sub = append(p.Sub, t)
	} else {
		// Add as temporary child of the root for now. When the
		// real parent is added, we'll fix it up later.
		if ms.log {
			log.Printf("Add %v as a temporary child of root %v", t.Span.ID, root.Span.ID)
		}
		root.Sub = append(root.Sub, t)
	}
}

// reattachChildren moves temporary children of src to dst, if dst is
// the node's parent.
func (ms *MemoryStore) reattachChildren(dst, src *Trace) {
	if dst == src {
		panic("dst == src")
	}
	var sub2 []*Trace
	for _, c := range src.Sub {
		if c.Span.ID.Parent == dst.Span.ID.Span {
			if ms.log {
				log.Printf("Move %v from src %v to dst %v", c.Span.ID, src.Span.ID, dst.Span.ID)
			}
			dst.Sub = append(dst.Sub, c)
		} else {
			sub2 = append(sub2, c)
		}
	}
	src.Sub = sub2
}

// Trace implements the Store interface by returning the Trace (a tree of
// spans) for the given trace span ID or, if no such trace exists, by returning
// ErrTraceNotFound.
func (ms *MemoryStore) Trace(id ID) (*Trace, error) {
	ms.Lock()
	defer ms.Unlock()

	return ms.traceNoLock(id)
}

func (ms *MemoryStore) traceNoLock(id ID) (*Trace, error) {
	t, present := ms.trace[id]
	if !present {
		return nil, ErrTraceNotFound
	}
	return t, nil
}

// Traces implements the Queryer interface.
func (ms *MemoryStore) Traces() ([]*Trace, error) {
	ms.Lock()
	defer ms.Unlock()

	var ts []*Trace
	for id := range ms.trace {
		t, err := ms.traceNoLock(id)
		if err != nil {
			return nil, err
		}
		ts = append(ts, t)
	}
	return ts, nil
}

// Delete implements the DeleteStore interface by deleting the traces given by
// their span ID's from this in-memory store.
func (ms *MemoryStore) Delete(traces ...ID) error {
	ms.Lock()
	defer ms.Unlock()

	for _, id := range traces {
		delete(ms.trace, id)
		delete(ms.span, id)
	}
	return nil
}

type memoryStoreData struct {
	Trace map[ID]*Trace
	Span  map[ID]map[ID]*Trace
}

// Write implements the PersistentStore interface by gob-encoding and writing
// ms's internal data structures out to w.
func (ms *MemoryStore) Write(w io.Writer) error {
	ms.Lock()
	defer ms.Unlock()

	data := memoryStoreData{ms.trace, ms.span}
	return gob.NewEncoder(w).Encode(data)
}

// ReadFrom implements the PersistentStore interface by using gob-decoding to
// load ms's internal data structures from the reader r.
func (ms *MemoryStore) ReadFrom(r io.Reader) (int64, error) {
	ms.Lock()
	defer ms.Unlock()

	var data memoryStoreData
	if err := gob.NewDecoder(r).Decode(&data); err != nil {
		return 0, err
	}
	ms.trace = data.Trace
	ms.span = data.Span
	return int64(len(ms.trace)), nil
}

// PersistentStore is a Store that can persist its data and read it
// back in.
type PersistentStore interface {
	Write(io.Writer) error
	ReadFrom(io.Reader) (int64, error)
	Store
}

// PersistEvery persists s's data to a file periodically.
func PersistEvery(s PersistentStore, interval time.Duration, file string) error {
	for {
		time.Sleep(interval)

		f, err := ioutil.TempFile("", "appdash")
		if err != nil {
			return err
		}
		if err := s.Write(f); err != nil {
			f.Close()
			return err
		}
		if err := f.Close(); err != nil {
			return err
		}
		if err := os.Rename(f.Name(), file); err != nil {
			return err
		}
	}
}

// A DeleteStore is a Store that can delete traces.
type DeleteStore interface {
	Store

	// Delete deletes traces given their trace IDs.
	Delete(...ID) error
}

// A RecentStore wraps another store and deletes old traces after a
// specified amount of time.
type RecentStore struct {
	// MinEvictAge is the minimum age of a trace before it is evicted.
	MinEvictAge time.Duration

	// DeleteStore is the underlying store that spans are saved to and
	// deleted from.
	DeleteStore

	// Debug is whether to log debug messages.
	Debug bool

	// created maps trace ID to the UnixNano time it was first seen.
	created map[ID]int64

	// lastEvicted is the last time the eviction process was run.
	lastEvicted time.Time

	mu sync.Mutex // mu guards created and lastEvicted
}

// Collect calls the underlying store's Collect and records the time
// that this trace was first seen.
func (rs *RecentStore) Collect(id SpanID, anns ...Annotation) error {
	rs.mu.Lock()
	if rs.created == nil {
		rs.created = map[ID]int64{}
	}
	if _, present := rs.created[id.Trace]; !present {
		rs.created[id.Trace] = time.Now().UnixNano()
	}
	if time.Since(rs.lastEvicted) > rs.MinEvictAge {
		rs.evictBefore(time.Now().Add(-1 * rs.MinEvictAge))
	}
	rs.mu.Unlock()

	return rs.DeleteStore.Collect(id, anns...)
}

// evictBefore evicts traces that were created before t. The rs.mu lock
// must be held while calling evictBefore.
func (rs *RecentStore) evictBefore(t time.Time) {
	evictStart := time.Now()
	tnano := t.UnixNano()
	var toEvict []ID
	for id, ct := range rs.created {
		if ct < tnano {
			toEvict = append(toEvict, id)
			delete(rs.created, id)
		}
	}
	if len(toEvict) == 0 {
		return
	}

	if rs.Debug {
		log.Printf("RecentStore: deleting %d traces created before %s (age check took %s)", len(toEvict), t, time.Since(evictStart))
	}

	// Spawn separate goroutine so we don't hold the rs.mu lock.
	go func() {
		deleteStart := time.Now()
		if err := rs.DeleteStore.Delete(toEvict...); err != nil {
			log.Printf("RecentStore: failed to delete traces: %s", err)
		}
		if rs.Debug {
			log.Printf("RecentStore: finished deleting %d traces created before %s (took %s)", len(toEvict), t, time.Since(deleteStart))
		}
	}()
}
<|endoftext|>


<commit_before>package rkv

import (
	"fmt"
	"math/rand"
	"time"

	"golang.org/x/net/context"

	"github.com/relab/raft"
	commonpb "github.com/relab/raft/raftpb"
)

// Store is a key-value store backed by a map.
type Store struct {
	store map[string]string

	nextID  uint64
	clients map[string]uint64

	waiting map[string]chan interface{}

	raft raft.Raft
}

// NewStore returns a newly initialized store.
func NewStore(raft raft.Raft) *Store {
	s := &Store{
		store:   make(map[string]string),
		clients: make(map[string]uint64),
		waiting: make(map[string]chan interface{}),
		raft:    raft,
	}

	go s.run()

	return s
}

func (s *Store) run() {
	for {
		select {
		case entries := <-s.raft.Committed():
			for _, entry := range entries {
				switch entry.EntryType {
				case commonpb.EntryNormal:
					cmd, err := ParseCmd(string(entry.Data))

					if err != nil {
						// TODO Ignore malformed requests. Log?
						continue
					}
					var res interface{}

					switch cmd.CmdType {
					case CmdRegister:
						id := fmt.Sprintf("%d", entry.Index)
						res = id
						if _, ok := s.clients[id]; ok {
							// Don't reset sequence
							// number if we receive
							// a duplicate.
							break
						}
						s.clients[id] = 0
					case CmdInsert:
						// TODO We need to return
						// ErrSessionExpired, when the
						// session is expired. TODO If
						// we see seq > oldseq + 1, we
						// need to deal with that, i.e,
						// we might have stale data?
						if oldseq := s.clients[cmd.ID]; oldseq >= cmd.Seq {
							// Already applied.
							break
						}

						s.clients[cmd.ID] = cmd.Seq
						key := cmd.Key
						value := cmd.Value

						s.store[key] = value
					}

					if ch, ok := s.waiting[cmd.Cmd]; ok {
						ch <- res
						delete(s.waiting, cmd.Cmd)
					}
				}
			}
		}
	}
}

func (s *Store) waitOnResponse(ctx context.Context, cmd string) (interface{}, error) {
	done := make(chan interface{}, 1)
	s.waiting[cmd] = done

	select {
	case res := <-done:
		return res, nil
	case <-ctx.Done():
		return nil, ctx.Err()
	}
}

func (s *Store) Register() (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	// Random int used to uniquely identify command. There is a chance of a
	// collision, but I'll take those odds any day.
	cmd, err := NewCmdRegister(rand.Int63())

	if err != nil {
		return "", err
	}

	if err := s.raft.ProposeCmd(ctx, []byte(cmd.String())); err != nil {
		return "", err
	}

	// There is a race condition on the command being committed before we
	// actually start waiting. If that happens, the wait will timeout, and
	// the client will retry.
	res, err := s.waitOnResponse(ctx, cmd.String())

	if err != nil {
		return "", err
	}

	return res.(string), nil
}

// Lookup gets a value from the map given a key. It returns a bool indicating if
// the value was found or not. TODO We need to query the state machine. It's not
// safe to just read the value.
func (s *Store) Lookup(key string) (string, error) {
	// TODO Buffer reads to amortize cost.
	err := s.raft.Read(context.TODO())

	if err != nil {
		return "", err
	}

	// Empty string if not set. This depends on the application I think. We
	// don't really care about the key being set or not.
	value, _ := s.store[key]

	return value, nil
}

// Insert inserts value in map given a key. It returns true if the value was
// successfully inserted. False indicates that the client should retry, as we
// cannot know if the value was inserted successfully or not.
func (s *Store) Insert(id, seq, key, value string) error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	// id and seq uniquely identifies this command.
	cmd, err := NewCmdInsert(id, seq, key, value)

	if err != nil {
		return err
	}

	if err := s.raft.ProposeCmd(ctx, []byte(cmd.String())); err != nil {
		return err
	}

	// There is a race condition on the command being committed before we
	// actually start waiting. If that happens, the wait will timeout, and
	// the client will retry.
	if _, err := s.waitOnResponse(ctx, cmd.String()); err != nil {
		return err
	}

	return nil
}
<commit_msg>store.go: Eliminate potential race (although it was safe)<commit_after>package rkv

import (
	"fmt"
	"math/rand"
	"time"

	"golang.org/x/net/context"

	"github.com/relab/raft"
	commonpb "github.com/relab/raft/raftpb"
)

// Store is a key-value store backed by a map.
type Store struct {
	store map[string]string

	nextID  uint64
	clients map[string]uint64

	waiting map[string]chan interface{}

	raft raft.Raft
}

// NewStore returns a newly initialized store.
func NewStore(raft raft.Raft) *Store {
	s := &Store{
		store:   make(map[string]string),
		clients: make(map[string]uint64),
		waiting: make(map[string]chan interface{}),
		raft:    raft,
	}

	go s.run()

	return s
}

func (s *Store) run() {
	for {
		select {
		case entries := <-s.raft.Committed():
			for _, entry := range entries {
				switch entry.EntryType {
				case commonpb.EntryNormal:
					cmd, err := ParseCmd(string(entry.Data))

					if err != nil {
						// TODO Ignore malformed requests. Log?
						continue
					}
					var res interface{}

					switch cmd.CmdType {
					case CmdRegister:
						id := fmt.Sprintf("%d", entry.Index)
						res = id
						if _, ok := s.clients[id]; ok {
							// Don't reset sequence
							// number if we receive
							// a duplicate.
							break
						}
						s.clients[id] = 0
					case CmdInsert:
						// TODO We need to return
						// ErrSessionExpired, when the
						// session is expired. TODO If
						// we see seq > oldseq + 1, we
						// need to deal with that, i.e,
						// we might have stale data?
						if oldseq := s.clients[cmd.ID]; oldseq >= cmd.Seq {
							// Already applied.
							break
						}

						s.clients[cmd.ID] = cmd.Seq
						key := cmd.Key
						value := cmd.Value

						s.store[key] = value
					}

					if ch, ok := s.waiting[cmd.Cmd]; ok {
						ch <- res
						delete(s.waiting, cmd.Cmd)
					}
				}
			}
		}
	}
}

func (s *Store) Register() (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	// Random int used to uniquely identify command. There is a chance of a
	// collision, but I'll take those odds any day.
	cmd, err := NewCmdRegister(rand.Int63())

	if err != nil {
		return "", err
	}

	done := make(chan interface{}, 1)
	s.waiting[cmd.String()] = done

	if err := s.raft.ProposeCmd(ctx, []byte(cmd.String())); err != nil {
		return "", err
	}

	select {
	case res := <-done:
		return res.(string), nil
	case <-ctx.Done():
		return "", ctx.Err()
	}
}

// Lookup gets a value from the map given a key. It returns a bool indicating if
// the value was found or not. TODO We need to query the state machine. It's not
// safe to just read the value.
func (s *Store) Lookup(key string) (string, error) {
	// TODO Buffer reads to amortize cost.
	err := s.raft.Read(context.TODO())

	if err != nil {
		return "", err
	}

	// Empty string if not set. This depends on the application I think. We
	// don't really care about the key being set or not.
	value, _ := s.store[key]

	return value, nil
}

// Insert inserts value in map given a key. It returns true if the value was
// successfully inserted. False indicates that the client should retry, as we
// cannot know if the value was inserted successfully or not.
func (s *Store) Insert(id, seq, key, value string) error {
	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	defer cancel()

	// id and seq uniquely identifies this command.
	cmd, err := NewCmdInsert(id, seq, key, value)

	if err != nil {
		return err
	}

	done := make(chan interface{}, 1)
	s.waiting[cmd.String()] = done

	if err := s.raft.ProposeCmd(ctx, []byte(cmd.String())); err != nil {
		return err
	}

	select {
	case <-done:
		return nil
	case <-ctx.Done():
		return ctx.Err()
	}
}
<|endoftext|>