<commit_before>package bslack

import (
	"errors"
	"fmt"
	"github.com/42wim/matterbridge/bridge/config"
	"github.com/42wim/matterbridge/matterhook"
	log "github.com/Sirupsen/logrus"
	"github.com/nlopes/slack"
	"html"
	"regexp"
	"strings"
	"time"
)

type MMMessage struct {
	Text     string
	Channel  string
	Username string
	UserID   string
	Raw      *slack.MessageEvent
}

type Bslack struct {
	mh       *matterhook.Client
	sc       *slack.Client
	Config   *config.Protocol
	rtm      *slack.RTM
	Plus     bool
	Remote   chan config.Message
	Users    []slack.User
	Account  string
	si       *slack.Info
	channels []slack.Channel
}

var flog *log.Entry
var protocol = "slack"

func init() {
	flog = log.WithFields(log.Fields{"module": protocol})
}

func New(cfg config.Protocol, account string, c chan config.Message) *Bslack {
	b := &Bslack{}
	b.Config = &cfg
	b.Remote = c
	b.Account = account
	return b
}

func (b *Bslack) Command(cmd string) string {
	return ""
}

func (b *Bslack) Connect() error {
	if b.Config.WebhookBindAddress != "" {
		if b.Config.WebhookURL != "" {
			flog.Info("Connecting using webhookurl (sending) and webhookbindaddress (receiving)")
			b.mh = matterhook.New(b.Config.WebhookURL,
				matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
					BindAddress: b.Config.WebhookBindAddress})
		} else if b.Config.Token != "" {
			flog.Info("Connecting using token (sending)")
			b.sc = slack.New(b.Config.Token)
			b.rtm = b.sc.NewRTM()
			go b.rtm.ManageConnection()
			flog.Info("Connecting using webhookbindaddress (receiving)")
			b.mh = matterhook.New(b.Config.WebhookURL,
				matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
					BindAddress: b.Config.WebhookBindAddress})
		} else {
			flog.Info("Connecting using webhookbindaddress (receiving)")
			b.mh = matterhook.New(b.Config.WebhookURL,
				matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
					BindAddress: b.Config.WebhookBindAddress})
		}
		go b.handleSlack()
		return nil
	}
	if b.Config.WebhookURL != "" {
		flog.Info("Connecting using webhookurl (sending)")
		b.mh = matterhook.New(b.Config.WebhookURL,
			matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
				DisableServer: true})
		if b.Config.Token != "" {
			flog.Info("Connecting using token (receiving)")
			b.sc = slack.New(b.Config.Token)
			b.rtm = b.sc.NewRTM()
			go b.rtm.ManageConnection()
			go b.handleSlack()
		}
	} else if b.Config.Token != "" {
		flog.Info("Connecting using token (sending and receiving)")
		b.sc = slack.New(b.Config.Token)
		b.rtm = b.sc.NewRTM()
		go b.rtm.ManageConnection()
		go b.handleSlack()
	}
	if b.Config.WebhookBindAddress == "" && b.Config.WebhookURL == "" && b.Config.Token == "" {
		return errors.New("No connection method found. See that you have WebhookBindAddress, WebhookURL or Token configured.")
	}
	return nil
}

func (b *Bslack) Disconnect() error {
	return nil

}

func (b *Bslack) JoinChannel(channel config.ChannelInfo) error {
	// we can only join channels using the API
	if b.Config.WebhookURL == "" && b.Config.WebhookBindAddress == "" {
		if strings.HasPrefix(b.Config.Token, "xoxb") {
			// TODO check if bot has already joined channel
			return nil
		}
		_, err := b.sc.JoinChannel(channel.Name)
		if err != nil {
			if err.Error() != "name_taken" {
				return err
			}
		}
	}
	return nil
}

func (b *Bslack) Send(msg config.Message) (string, error) {
	flog.Debugf("Receiving %#v", msg)
	if msg.Event == config.EVENT_USER_ACTION {
		msg.Text = "_" + msg.Text + "_"
	}
	nick := msg.Username
	message := msg.Text
	channel := msg.Channel
	if b.Config.PrefixMessagesWithNick {
		message = nick + " " + message
	}
	if b.Config.WebhookURL != "" {
		matterMessage := matterhook.OMessage{IconURL: b.Config.IconURL}
		matterMessage.Channel = channel
		matterMessage.UserName = nick
		matterMessage.Type = ""
		matterMessage.Text = message
		err := b.mh.Send(matterMessage)
		if err != nil {
			flog.Info(err)
			return "", err
		}
		return "", nil
	}
	schannel, err := b.getChannelByName(channel)
	if err != nil {
		return "", err
	}
	np := slack.NewPostMessageParameters()
	if b.Config.PrefixMessagesWithNick {
		np.AsUser = true
	}
	np.Username = nick
	np.IconURL = config.GetIconURL(&msg, b.Config)
	if msg.Avatar != "" {
		np.IconURL = msg.Avatar
	}
	np.Attachments = append(np.Attachments, slack.Attachment{CallbackID: "matterbridge"})
	np.Attachments = append(np.Attachments, b.createAttach(msg.Extra)...)

	// replace mentions
	np.LinkNames = 1

	if msg.Event == config.EVENT_MSG_DELETE {
		// some protocols echo deletes, but with empty ID
		if msg.ID == "" {
			return "", nil
		}
		// we get a "slack <ID>", split it
		ts := strings.Fields(msg.ID)
		b.sc.DeleteMessage(schannel.ID, ts[1])
		return "", nil
	}
	// if we have no ID it means we're creating a new message, not updating an existing one
	if msg.ID != "" {
		ts := strings.Fields(msg.ID)
		b.sc.UpdateMessage(schannel.ID, ts[1], message)
		return "", nil
	}
	_, id, err := b.sc.PostMessage(schannel.ID, message, np)
	if err != nil {
		return "", err
	}
	return "slack " + id, nil
}

func (b *Bslack) getAvatar(user string) string {
	var avatar string
	if b.Users != nil {
		for _, u := range b.Users {
			if user == u.Name {
				return u.Profile.Image48
			}
		}
	}
	return avatar
}

func (b *Bslack) getChannelByName(name string) (*slack.Channel, error) {
	if b.channels == nil {
		return nil, fmt.Errorf("%s: channel %s not found (no channels found)", b.Account, name)
	}
	for _, channel := range b.channels {
		if channel.Name == name {
			return &channel, nil
		}
	}
	return nil, fmt.Errorf("%s: channel %s not found", b.Account, name)
}

func (b *Bslack) getChannelByID(ID string) (*slack.Channel, error) {
	if b.channels == nil {
		return nil, fmt.Errorf("%s: channel %s not found (no channels found)", b.Account, ID)
	}
	for _, channel := range b.channels {
		if channel.ID == ID {
			return &channel, nil
		}
	}
	return nil, fmt.Errorf("%s: channel %s not found", b.Account, ID)
}

func (b *Bslack) handleSlack() {
	mchan := make(chan *MMMessage)
	if b.Config.WebhookBindAddress != "" {
		flog.Debugf("Choosing webhooks based receiving")
		go b.handleMatterHook(mchan)
	} else {
		flog.Debugf("Choosing token based receiving")
		go b.handleSlackClient(mchan)
	}
	time.Sleep(time.Second)
	flog.Debug("Start listening for Slack messages")
	for message := range mchan {
		// do not send messages from ourself
		if b.Config.WebhookURL == "" && b.Config.WebhookBindAddress == "" && message.Username == b.si.User.Name {
			continue
		}
		if (message.Text == "" || message.Username == "") && message.Raw.SubType != "message_deleted" {
			continue
		}
		text := message.Text
		text = b.replaceURL(text)
		text = html.UnescapeString(text)
		flog.Debugf("Sending message from %s on %s to gateway", message.Username, b.Account)
		msg := config.Message{Text: text, Username: message.Username, Channel: message.Channel, Account: b.Account, Avatar: b.getAvatar(message.Username), UserID: message.UserID, ID: "slack " + message.Raw.Timestamp}
		if message.Raw.SubType == "me_message" {
			msg.Event = config.EVENT_USER_ACTION
		}
		if message.Raw.SubType == "channel_leave" || message.Raw.SubType == "channel_join" {
			msg.Username = "system"
			msg.Event = config.EVENT_JOIN_LEAVE
		}
		// edited messages have a submessage, use this timestamp
		if message.Raw.SubMessage != nil {
			msg.ID = "slack " + message.Raw.SubMessage.Timestamp
		}
		if message.Raw.SubType == "message_deleted" {
			msg.Text = config.EVENT_MSG_DELETE
			msg.Event = config.EVENT_MSG_DELETE
			msg.ID = "slack " + message.Raw.DeletedTimestamp
		}
		flog.Debugf("Message is %#v", msg)
		b.Remote <- msg
	}
}

func (b *Bslack) handleSlackClient(mchan chan *MMMessage) {
	for msg := range b.rtm.IncomingEvents {
		switch ev := msg.Data.(type) {
		case *slack.MessageEvent:
			flog.Debugf("Receiving from slackclient %#v", ev)
			if len(ev.Attachments) > 0 {
				// skip messages we made ourselves
				if ev.Attachments[0].CallbackID == "matterbridge" {
					continue
				}
			}
			if !b.Config.EditDisable && ev.SubMessage != nil && ev.SubMessage.ThreadTimestamp != ev.SubMessage.Timestamp {
				flog.Debugf("SubMessage %#v", ev.SubMessage)
				ev.User = ev.SubMessage.User
				ev.Text = ev.SubMessage.Text + b.Config.EditSuffix
			}
			// use our own func because rtm.GetChannelInfo doesn't work for private channels
			channel, err := b.getChannelByID(ev.Channel)
			if err != nil {
				continue
			}
			m := &MMMessage{}
			if ev.BotID == "" && ev.SubType != "message_deleted" {
				user, err := b.rtm.GetUserInfo(ev.User)
				if err != nil {
					continue
				}
				m.UserID = user.ID
				m.Username = user.Name
			}
			m.Channel = channel.Name
			m.Text = ev.Text
			if m.Text == "" {
				for _, attach := range ev.Attachments {
					if attach.Text != "" {
						m.Text = attach.Text
					} else {
						m.Text = attach.Fallback
					}
				}
			}
			m.Raw = ev
			m.Text = b.replaceMention(m.Text)
			// when using webhookURL we can't check if it's our webhook or not for now
			if ev.BotID != "" && b.Config.WebhookURL == "" {
				bot, err := b.rtm.GetBotInfo(ev.BotID)
				if err != nil {
					continue
				}
				if bot.Name != "" {
					m.Username = bot.Name
					m.UserID = bot.ID
				}
			}
			mchan <- m
		case *slack.OutgoingErrorEvent:
			flog.Debugf("%#v", ev.Error())
		case *slack.ChannelJoinedEvent:
			b.Users, _ = b.sc.GetUsers()
		case *slack.ConnectedEvent:
			b.channels = ev.Info.Channels
			b.si = ev.Info
			b.Users, _ = b.sc.GetUsers()
			// add private channels
			groups, _ := b.sc.GetGroups(true)
			for _, g := range groups {
				channel := new(slack.Channel)
				channel.ID = g.ID
				channel.Name = g.Name
				b.channels = append(b.channels, *channel)
			}
		case *slack.InvalidAuthEvent:
			flog.Fatalf("Invalid Token %#v", ev)
		default:
		}
	}
}

func (b *Bslack) handleMatterHook(mchan chan *MMMessage) {
	for {
		message := b.mh.Receive()
		flog.Debugf("receiving from matterhook (slack) %#v", message)
		m := &MMMessage{}
		m.Username = message.UserName
		m.Text = message.Text
		m.Text = b.replaceMention(m.Text)
		m.Channel = message.ChannelName
		if m.Username == "slackbot" {
			continue
		}
		mchan <- m
	}
}

func (b *Bslack) userName(id string) string {
	for _, u := range b.Users {
		if u.ID == id {
			return u.Name
		}
	}
	return ""
}

func (b *Bslack) replaceMention(text string) string {
	results := regexp.MustCompile(`<@([a-zA-z0-9]+)>`).FindAllStringSubmatch(text, -1)
	for _, r := range results {
		text = strings.Replace(text, "<@"+r[1]+">", "@"+b.userName(r[1]), -1)

	}
	return text
}

func (b *Bslack) replaceURL(text string) string {
	results := regexp.MustCompile(`<(.*?)\|.*?>`).FindAllStringSubmatch(text, -1)
	for _, r := range results {
		text = strings.Replace(text, r[0], r[1], -1)
	}
	return text
}

func (b *Bslack) createAttach(extra []interface{}) []slack.Attachment {
	var attachs []slack.Attachment
	if extra != nil {
		for _, v := range extra {
			entry := v.(map[string]interface{})
			s := slack.Attachment{}
			s.Fallback = entry["fallback"].(string)
			s.Color = entry["color"].(string)
			s.Pretext = entry["pretext"].(string)
			s.AuthorName = entry["author_name"].(string)
			s.AuthorLink = entry["author_link"].(string)
			s.AuthorIcon = entry["author_icon"].(string)
			s.Title = entry["title"].(string)
			s.TitleLink = entry["title_link"].(string)
			s.Text = entry["text"].(string)
			s.ImageURL = entry["image_url"].(string)
			s.ThumbURL = entry["thumb_url"].(string)
			s.Footer = entry["footer"].(string)
			s.FooterIcon = entry["footer_icon"].(string)
			attachs = append(attachs, s)
		}
	}
	return attachs
}
<commit_msg>Remove unnecessary check, make megacheck happy<commit_after>package bslack

import (
	"errors"
	"fmt"
	"github.com/42wim/matterbridge/bridge/config"
	"github.com/42wim/matterbridge/matterhook"
	log "github.com/Sirupsen/logrus"
	"github.com/nlopes/slack"
	"html"
	"regexp"
	"strings"
	"time"
)

type MMMessage struct {
	Text     string
	Channel  string
	Username string
	UserID   string
	Raw      *slack.MessageEvent
}

type Bslack struct {
	mh       *matterhook.Client
	sc       *slack.Client
	Config   *config.Protocol
	rtm      *slack.RTM
	Plus     bool
	Remote   chan config.Message
	Users    []slack.User
	Account  string
	si       *slack.Info
	channels []slack.Channel
}

var flog *log.Entry
var protocol = "slack"

func init() {
	flog = log.WithFields(log.Fields{"module": protocol})
}

func New(cfg config.Protocol, account string, c chan config.Message) *Bslack {
	b := &Bslack{}
	b.Config = &cfg
	b.Remote = c
	b.Account = account
	return b
}

func (b *Bslack) Command(cmd string) string {
	return ""
}

func (b *Bslack) Connect() error {
	if b.Config.WebhookBindAddress != "" {
		if b.Config.WebhookURL != "" {
			flog.Info("Connecting using webhookurl (sending) and webhookbindaddress (receiving)")
			b.mh = matterhook.New(b.Config.WebhookURL,
				matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
					BindAddress: b.Config.WebhookBindAddress})
		} else if b.Config.Token != "" {
			flog.Info("Connecting using token (sending)")
			b.sc = slack.New(b.Config.Token)
			b.rtm = b.sc.NewRTM()
			go b.rtm.ManageConnection()
			flog.Info("Connecting using webhookbindaddress (receiving)")
			b.mh = matterhook.New(b.Config.WebhookURL,
				matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
					BindAddress: b.Config.WebhookBindAddress})
		} else {
			flog.Info("Connecting using webhookbindaddress (receiving)")
			b.mh = matterhook.New(b.Config.WebhookURL,
				matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
					BindAddress: b.Config.WebhookBindAddress})
		}
		go b.handleSlack()
		return nil
	}
	if b.Config.WebhookURL != "" {
		flog.Info("Connecting using webhookurl (sending)")
		b.mh = matterhook.New(b.Config.WebhookURL,
			matterhook.Config{InsecureSkipVerify: b.Config.SkipTLSVerify,
				DisableServer: true})
		if b.Config.Token != "" {
			flog.Info("Connecting using token (receiving)")
			b.sc = slack.New(b.Config.Token)
			b.rtm = b.sc.NewRTM()
			go b.rtm.ManageConnection()
			go b.handleSlack()
		}
	} else if b.Config.Token != "" {
		flog.Info("Connecting using token (sending and receiving)")
		b.sc = slack.New(b.Config.Token)
		b.rtm = b.sc.NewRTM()
		go b.rtm.ManageConnection()
		go b.handleSlack()
	}
	if b.Config.WebhookBindAddress == "" && b.Config.WebhookURL == "" && b.Config.Token == "" {
		return errors.New("No connection method found. See that you have WebhookBindAddress, WebhookURL or Token configured.")
	}
	return nil
}

func (b *Bslack) Disconnect() error {
	return nil

}

func (b *Bslack) JoinChannel(channel config.ChannelInfo) error {
	// we can only join channels using the API
	if b.Config.WebhookURL == "" && b.Config.WebhookBindAddress == "" {
		if strings.HasPrefix(b.Config.Token, "xoxb") {
			// TODO check if bot has already joined channel
			return nil
		}
		_, err := b.sc.JoinChannel(channel.Name)
		if err != nil {
			if err.Error() != "name_taken" {
				return err
			}
		}
	}
	return nil
}

func (b *Bslack) Send(msg config.Message) (string, error) {
	flog.Debugf("Receiving %#v", msg)
	if msg.Event == config.EVENT_USER_ACTION {
		msg.Text = "_" + msg.Text + "_"
	}
	nick := msg.Username
	message := msg.Text
	channel := msg.Channel
	if b.Config.PrefixMessagesWithNick {
		message = nick + " " + message
	}
	if b.Config.WebhookURL != "" {
		matterMessage := matterhook.OMessage{IconURL: b.Config.IconURL}
		matterMessage.Channel = channel
		matterMessage.UserName = nick
		matterMessage.Type = ""
		matterMessage.Text = message
		err := b.mh.Send(matterMessage)
		if err != nil {
			flog.Info(err)
			return "", err
		}
		return "", nil
	}
	schannel, err := b.getChannelByName(channel)
	if err != nil {
		return "", err
	}
	np := slack.NewPostMessageParameters()
	if b.Config.PrefixMessagesWithNick {
		np.AsUser = true
	}
	np.Username = nick
	np.IconURL = config.GetIconURL(&msg, b.Config)
	if msg.Avatar != "" {
		np.IconURL = msg.Avatar
	}
	np.Attachments = append(np.Attachments, slack.Attachment{CallbackID: "matterbridge"})
	np.Attachments = append(np.Attachments, b.createAttach(msg.Extra)...)

	// replace mentions
	np.LinkNames = 1

	if msg.Event == config.EVENT_MSG_DELETE {
		// some protocols echo deletes, but with empty ID
		if msg.ID == "" {
			return "", nil
		}
		// we get a "slack <ID>", split it
		ts := strings.Fields(msg.ID)
		b.sc.DeleteMessage(schannel.ID, ts[1])
		return "", nil
	}
	// if we have no ID it means we're creating a new message, not updating an existing one
	if msg.ID != "" {
		ts := strings.Fields(msg.ID)
		b.sc.UpdateMessage(schannel.ID, ts[1], message)
		return "", nil
	}
	_, id, err := b.sc.PostMessage(schannel.ID, message, np)
	if err != nil {
		return "", err
	}
	return "slack " + id, nil
}

func (b *Bslack) getAvatar(user string) string {
	var avatar string
	if b.Users != nil {
		for _, u := range b.Users {
			if user == u.Name {
				return u.Profile.Image48
			}
		}
	}
	return avatar
}

func (b *Bslack) getChannelByName(name string) (*slack.Channel, error) {
	if b.channels == nil {
		return nil, fmt.Errorf("%s: channel %s not found (no channels found)", b.Account, name)
	}
	for _, channel := range b.channels {
		if channel.Name == name {
			return &channel, nil
		}
	}
	return nil, fmt.Errorf("%s: channel %s not found", b.Account, name)
}

func (b *Bslack) getChannelByID(ID string) (*slack.Channel, error) {
	if b.channels == nil {
		return nil, fmt.Errorf("%s: channel %s not found (no channels found)", b.Account, ID)
	}
	for _, channel := range b.channels {
		if channel.ID == ID {
			return &channel, nil
		}
	}
	return nil, fmt.Errorf("%s: channel %s not found", b.Account, ID)
}

func (b *Bslack) handleSlack() {
	mchan := make(chan *MMMessage)
	if b.Config.WebhookBindAddress != "" {
		flog.Debugf("Choosing webhooks based receiving")
		go b.handleMatterHook(mchan)
	} else {
		flog.Debugf("Choosing token based receiving")
		go b.handleSlackClient(mchan)
	}
	time.Sleep(time.Second)
	flog.Debug("Start listening for Slack messages")
	for message := range mchan {
		// do not send messages from ourself
		if b.Config.WebhookURL == "" && b.Config.WebhookBindAddress == "" && message.Username == b.si.User.Name {
			continue
		}
		if (message.Text == "" || message.Username == "") && message.Raw.SubType != "message_deleted" {
			continue
		}
		text := message.Text
		text = b.replaceURL(text)
		text = html.UnescapeString(text)
		flog.Debugf("Sending message from %s on %s to gateway", message.Username, b.Account)
		msg := config.Message{Text: text, Username: message.Username, Channel: message.Channel, Account: b.Account, Avatar: b.getAvatar(message.Username), UserID: message.UserID, ID: "slack " + message.Raw.Timestamp}
		if message.Raw.SubType == "me_message" {
			msg.Event = config.EVENT_USER_ACTION
		}
		if message.Raw.SubType == "channel_leave" || message.Raw.SubType == "channel_join" {
			msg.Username = "system"
			msg.Event = config.EVENT_JOIN_LEAVE
		}
		// edited messages have a submessage, use this timestamp
		if message.Raw.SubMessage != nil {
			msg.ID = "slack " + message.Raw.SubMessage.Timestamp
		}
		if message.Raw.SubType == "message_deleted" {
			msg.Text = config.EVENT_MSG_DELETE
			msg.Event = config.EVENT_MSG_DELETE
			msg.ID = "slack " + message.Raw.DeletedTimestamp
		}
		flog.Debugf("Message is %#v", msg)
		b.Remote <- msg
	}
}

func (b *Bslack) handleSlackClient(mchan chan *MMMessage) {
	for msg := range b.rtm.IncomingEvents {
		switch ev := msg.Data.(type) {
		case *slack.MessageEvent:
			flog.Debugf("Receiving from slackclient %#v", ev)
			if len(ev.Attachments) > 0 {
				// skip messages we made ourselves
				if ev.Attachments[0].CallbackID == "matterbridge" {
					continue
				}
			}
			if !b.Config.EditDisable && ev.SubMessage != nil && ev.SubMessage.ThreadTimestamp != ev.SubMessage.Timestamp {
				flog.Debugf("SubMessage %#v", ev.SubMessage)
				ev.User = ev.SubMessage.User
				ev.Text = ev.SubMessage.Text + b.Config.EditSuffix
			}
			// use our own func because rtm.GetChannelInfo doesn't work for private channels
			channel, err := b.getChannelByID(ev.Channel)
			if err != nil {
				continue
			}
			m := &MMMessage{}
			if ev.BotID == "" && ev.SubType != "message_deleted" {
				user, err := b.rtm.GetUserInfo(ev.User)
				if err != nil {
					continue
				}
				m.UserID = user.ID
				m.Username = user.Name
			}
			m.Channel = channel.Name
			m.Text = ev.Text
			if m.Text == "" {
				for _, attach := range ev.Attachments {
					if attach.Text != "" {
						m.Text = attach.Text
					} else {
						m.Text = attach.Fallback
					}
				}
			}
			m.Raw = ev
			m.Text = b.replaceMention(m.Text)
			// when using webhookURL we can't check if it's our webhook or not for now
			if ev.BotID != "" && b.Config.WebhookURL == "" {
				bot, err := b.rtm.GetBotInfo(ev.BotID)
				if err != nil {
					continue
				}
				if bot.Name != "" {
					m.Username = bot.Name
					m.UserID = bot.ID
				}
			}
			mchan <- m
		case *slack.OutgoingErrorEvent:
			flog.Debugf("%#v", ev.Error())
		case *slack.ChannelJoinedEvent:
			b.Users, _ = b.sc.GetUsers()
		case *slack.ConnectedEvent:
			b.channels = ev.Info.Channels
			b.si = ev.Info
			b.Users, _ = b.sc.GetUsers()
			// add private channels
			groups, _ := b.sc.GetGroups(true)
			for _, g := range groups {
				channel := new(slack.Channel)
				channel.ID = g.ID
				channel.Name = g.Name
				b.channels = append(b.channels, *channel)
			}
		case *slack.InvalidAuthEvent:
			flog.Fatalf("Invalid Token %#v", ev)
		default:
		}
	}
}

func (b *Bslack) handleMatterHook(mchan chan *MMMessage) {
	for {
		message := b.mh.Receive()
		flog.Debugf("receiving from matterhook (slack) %#v", message)
		m := &MMMessage{}
		m.Username = message.UserName
		m.Text = message.Text
		m.Text = b.replaceMention(m.Text)
		m.Channel = message.ChannelName
		if m.Username == "slackbot" {
			continue
		}
		mchan <- m
	}
}

func (b *Bslack) userName(id string) string {
	for _, u := range b.Users {
		if u.ID == id {
			return u.Name
		}
	}
	return ""
}

func (b *Bslack) replaceMention(text string) string {
	results := regexp.MustCompile(`<@([a-zA-z0-9]+)>`).FindAllStringSubmatch(text, -1)
	for _, r := range results {
		text = strings.Replace(text, "<@"+r[1]+">", "@"+b.userName(r[1]), -1)

	}
	return text
}

func (b *Bslack) replaceURL(text string) string {
	results := regexp.MustCompile(`<(.*?)\|.*?>`).FindAllStringSubmatch(text, -1)
	for _, r := range results {
		text = strings.Replace(text, r[0], r[1], -1)
	}
	return text
}

func (b *Bslack) createAttach(extra []interface{}) []slack.Attachment {
	var attachs []slack.Attachment
	for _, v := range extra {
		entry := v.(map[string]interface{})
		s := slack.Attachment{}
		s.Fallback = entry["fallback"].(string)
		s.Color = entry["color"].(string)
		s.Pretext = entry["pretext"].(string)
		s.AuthorName = entry["author_name"].(string)
		s.AuthorLink = entry["author_link"].(string)
		s.AuthorIcon = entry["author_icon"].(string)
		s.Title = entry["title"].(string)
		s.TitleLink = entry["title_link"].(string)
		s.Text = entry["text"].(string)
		s.ImageURL = entry["image_url"].(string)
		s.ThumbURL = entry["thumb_url"].(string)
		s.Footer = entry["footer"].(string)
		s.FooterIcon = entry["footer_icon"].(string)
		attachs = append(attachs, s)
	}
	return attachs
}
<|endoftext|>


<commit_before>package broker_test

import (
	"net/http/httptest"

	"github.com/apoydence/talaria/broker"
	"github.com/apoydence/talaria/pb/messages"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Client", func() {

	var (
		mockServers []*mockServer
		servers     []*httptest.Server
		client      *broker.Client

		extraMockServer *mockServer
		extraServerURL  string
	)

	var createKnownMockServers = func() []string {
		var urls []string
		for i := 0; i < 3; i++ {
			mockServer, server := startMockServer()
			mockServers = append(mockServers, mockServer)
			servers = append(servers, server)
			urls = append(urls, convertToWs(server.URL))
		}

		return urls
	}

	var createExtraMockServer = func() {
		mockServer, server := startMockServer()
		extraMockServer = mockServer
		servers = append(servers, server)
		extraServerURL = convertToWs(server.URL)
	}

	BeforeEach(func() {
		mockServers = nil
		servers = nil
		URLs := createKnownMockServers()
		createExtraMockServer()

		var err error
		client, err = broker.NewClient(URLs...)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		client.Close()
		for _, server := range mockServers {
			close(server.serverCh)
		}

		for _, server := range servers {
			server.CloseClientConnections()
			server.Close()
		}
	})

	It("connects to each broker", func() {
		for _, server := range mockServers {
			Eventually(server.connEstablishedCh).Should(BeClosed())
		}
	})

	Describe("FetchWriter()", func() {

		var (
			expectedFileName string
			expectedData     []byte
		)

		BeforeEach(func() {
			expectedFileName = "some-file-1"
			expectedData = []byte("some-data")
		})

		Context("without errors", func() {

			BeforeEach(func() {

				mockServers[0].serverCh <- buildRemoteFileLocation(1, servers[1].URL)
				mockServers[1].serverCh <- buildFileLocation(1)

				for i := 0; i < 10; i++ {
					mockServers[1].serverCh <- buildFileOffset(uint64(i+2), 101)
				}
			})

			It("writes to the correct broker", func(done Done) {
				defer close(done)

				writer, err := client.FetchWriter(expectedFileName)
				Expect(err).ToNot(HaveOccurred())

				_, err = writer.WriteToFile(expectedData)
				Expect(err).ToNot(HaveOccurred())

				var msg *messages.Client
				Eventually(mockServers[1].clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_FetchFile))

				Eventually(mockServers[1].clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_WriteToFile))
				Expect(msg.GetWriteToFile().GetData()).To(Equal(expectedData))
			}, 5)

		})

		Measure("Writes to a file 1000 times in under a second", func(b Benchmarker) {
			runtime := b.Time("runtime", func() {
				mockServers[0].serverCh <- buildFileLocation(1)
				count := 1000

				go func() {
					for i := 0; i < count; i++ {
						mockServers[0].serverCh <- buildFileOffset(uint64(i+2), 101)
					}
				}()

				go func() {
					for _ = range mockServers[0].clientCh {
						//NOP
					}
				}()

				writer, _ := client.FetchWriter(expectedFileName)

				for i := 0; i < count; i++ {
					writer.WriteToFile(expectedData)
				}
			})

			Expect(runtime.Seconds()).To(BeNumerically("<", 1))

		}, 5)
	})

	Describe("FetchReader()", func() {

		var (
			expectedFileName string
			expectedData     []byte
		)

		BeforeEach(func() {
			expectedFileName = "some-file-1"
			expectedData = []byte("some-data")
		})

		Context("for known server", func() {

			BeforeEach(func() {
				mockServers[0].serverCh <- buildRemoteFileLocation(1, servers[1].URL)
				mockServers[1].serverCh <- buildFileLocation(1)
				mockServers[1].serverCh <- buildReadData(2, expectedData, 101)
			})

			It("reads from the correct broker", func(done Done) {
				defer close(done)

				reader, err := client.FetchReader(expectedFileName)
				Expect(err).ToNot(HaveOccurred())

				data, index, err := reader.ReadFromFile()
				Expect(err).ToNot(HaveOccurred())
				Expect(data).To(Equal(expectedData))
				Expect(index).To(BeEquivalentTo(101))

				var msg *messages.Client
				Eventually(mockServers[1].clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_FetchFile))
				Expect(msg.GetFetchFile().GetName()).To(Equal(expectedFileName))
			})
		})

		Context("for unknown server", func() {
			BeforeEach(func() {
				mockServers[0].serverCh <- buildRemoteFileLocation(1, extraServerURL)
				extraMockServer.serverCh <- buildFileLocation(1)
				extraMockServer.serverCh <- buildReadData(2, expectedData, 101)
			})

			It("reads from the new broker", func(done Done) {
				defer close(done)

				reader, err := client.FetchReader(expectedFileName)
				Expect(err).ToNot(HaveOccurred())

				data, index, err := reader.ReadFromFile()
				Expect(err).ToNot(HaveOccurred())
				Expect(data).To(Equal(expectedData))
				Expect(index).To(BeEquivalentTo(101))

				var msg *messages.Client
				Eventually(extraMockServer.clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_FetchFile))
				Expect(msg.GetFetchFile().GetName()).To(Equal(expectedFileName))
			})
		})

		Measure("Reads from a file 1000 times in under a second", func(b Benchmarker) {
			runtime := b.Time("runtime", func() {
				mockServers[0].serverCh <- buildFileLocation(1)
				count := 1000
				data := []byte("some-data")

				go func() {
					for i := 0; i < count; i++ {
						mockServers[0].serverCh <- buildReadData(uint64(i+2), data, 0)
					}
				}()

				go func() {
					for _ = range mockServers[0].clientCh {
						//NOP
					}
				}()

				reader, err := client.FetchReader("some-file")
				Expect(err).ToNot(HaveOccurred())
				for i := 0; i < count; i++ {
					reader.ReadFromFile()
				}
			})

			Expect(runtime.Seconds()).To(BeNumerically("<", 1))

		}, 5)
	})

	XDescribe("LeaderOf()", func() {
		var (
			id uint64
		)

		BeforeEach(func(done Done) {
			defer close(done)
			mockServers[0].serverCh <- buildFileLocation(1)
		})

		It("returns the broker URI that is the leader of the given file", func(done Done) {
			defer close(done)
			uri, err := client.LeaderOf(id)
			Expect(err).ToNot(HaveOccurred())
			Expect(uri).To(Equal(htmlToWs(servers[0].URL)))
		})

		It("returns an error for an unknown file ID", func(done Done) {
			defer close(done)
			_, err := client.LeaderOf(999)
			Expect(err).To(HaveOccurred())
		})
	})

})

func htmlToWs(URI string) string {
	return "ws" + URI[4:]
}

func startMockServer() (*mockServer, *httptest.Server) {
	mockServer := newMockServer()
	server := httptest.NewServer(mockServer)
	return mockServer, server
}
<commit_msg>Simplifies client tests to not cover multiple brokers. Tests only belong in ConnectionFetcher.<commit_after>package broker_test

import (
	"net/http/httptest"

	"github.com/apoydence/talaria/broker"
	"github.com/apoydence/talaria/pb/messages"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Client", func() {

	var (
		mockServer *mockServer
		server     *httptest.Server

		client *broker.Client
	)

	var createMockServer = func() string {
		mockServer, server = startMockServer()
		return convertToWs(server.URL)
	}

	BeforeEach(func() {
		URL := createMockServer()

		var err error
		client, err = broker.NewClient(URL)
		Expect(err).ToNot(HaveOccurred())
	})

	AfterEach(func() {
		client.Close()
		close(mockServer.serverCh)

		server.CloseClientConnections()
		server.Close()
	})

	Describe("FetchWriter()", func() {

		var (
			expectedFileName string
			expectedData     []byte
		)

		BeforeEach(func() {
			expectedFileName = "some-file-1"
			expectedData = []byte("some-data")
		})

		Context("without errors", func() {

			BeforeEach(func() {

				mockServer.serverCh <- buildFileLocation(1)

				for i := 0; i < 10; i++ {
					mockServer.serverCh <- buildFileOffset(uint64(i+2), 101)
				}
			})

			It("writes to the correct broker", func(done Done) {
				defer close(done)

				writer, err := client.FetchWriter(expectedFileName)
				Expect(err).ToNot(HaveOccurred())

				_, err = writer.WriteToFile(expectedData)
				Expect(err).ToNot(HaveOccurred())

				var msg *messages.Client
				Eventually(mockServer.clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_FetchFile))

				Eventually(mockServer.clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_WriteToFile))
				Expect(msg.GetWriteToFile().GetData()).To(Equal(expectedData))
			}, 5)

		})

		Measure("Writes to a file 1000 times in under a second", func(b Benchmarker) {
			runtime := b.Time("runtime", func() {
				mockServer.serverCh <- buildFileLocation(1)
				count := 1000

				go func() {
					for i := 0; i < count; i++ {
						mockServer.serverCh <- buildFileOffset(uint64(i+2), 101)
					}
				}()

				go func() {
					for _ = range mockServer.clientCh {
						//NOP
					}
				}()

				writer, _ := client.FetchWriter(expectedFileName)

				for i := 0; i < count; i++ {
					writer.WriteToFile(expectedData)
				}
			})

			Expect(runtime.Seconds()).To(BeNumerically("<", 1))

		}, 5)
	})

	Describe("FetchReader()", func() {

		var (
			expectedFileName string
			expectedData     []byte
		)

		BeforeEach(func() {
			expectedFileName = "some-file-1"
			expectedData = []byte("some-data")
		})

		Context("for known server", func() {

			BeforeEach(func() {
				mockServer.serverCh <- buildFileLocation(1)
				mockServer.serverCh <- buildReadData(2, expectedData, 101)
			})

			It("reads from the correct broker", func(done Done) {
				defer close(done)

				reader, err := client.FetchReader(expectedFileName)
				Expect(err).ToNot(HaveOccurred())

				data, index, err := reader.ReadFromFile()
				Expect(err).ToNot(HaveOccurred())
				Expect(data).To(Equal(expectedData))
				Expect(index).To(BeEquivalentTo(101))

				var msg *messages.Client
				Eventually(mockServer.clientCh).Should(Receive(&msg))
				Expect(msg.GetMessageType()).To(Equal(messages.Client_FetchFile))
				Expect(msg.GetFetchFile().GetName()).To(Equal(expectedFileName))
			})
		})

		Measure("Reads from a file 1000 times in under a second", func(b Benchmarker) {
			runtime := b.Time("runtime", func() {
				mockServer.serverCh <- buildFileLocation(1)
				count := 1000
				data := []byte("some-data")

				go func() {
					for i := 0; i < count; i++ {
						mockServer.serverCh <- buildReadData(uint64(i+2), data, 0)
					}
				}()

				go func() {
					for _ = range mockServer.clientCh {
						//NOP
					}
				}()

				reader, err := client.FetchReader("some-file")
				Expect(err).ToNot(HaveOccurred())
				for i := 0; i < count; i++ {
					reader.ReadFromFile()
				}
			})

			Expect(runtime.Seconds()).To(BeNumerically("<", 1))

		}, 5)
	})

	XDescribe("LeaderOf()", func() {
		var (
			id uint64
		)

		BeforeEach(func(done Done) {
			defer close(done)
			mockServer.serverCh <- buildFileLocation(1)
		})

		It("returns the broker URI that is the leader of the given file", func(done Done) {
			defer close(done)
			uri, err := client.LeaderOf(id)
			Expect(err).ToNot(HaveOccurred())
			Expect(uri).To(Equal(htmlToWs(server.URL)))
		})

		It("returns an error for an unknown file ID", func(done Done) {
			defer close(done)
			_, err := client.LeaderOf(999)
			Expect(err).To(HaveOccurred())
		})
	})

})

func htmlToWs(URI string) string {
	return "ws" + URI[4:]
}

func startMockServer() (*mockServer, *httptest.Server) {
	mockServer := newMockServer()
	server := httptest.NewServer(mockServer)
	return mockServer, server
}
<|endoftext|>


<commit_before>package canal

import (
	"bytes"
	"os"
	"sync"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/zjh1943/go-mysql/mysql"
	"github.com/siddontang/go/ioutil2"
)

type masterInfo struct {
	Addr     string `toml:"addr"`
	Name     string `toml:"bin_name"`
	Position uint32 `toml:"bin_pos"`

	name string

	l sync.RWMutex

	lastSaveTime time.Time
}

func loadMasterInfo(name string) (*masterInfo, error) {
	var m masterInfo

	m.name = name

	f, err := os.Open(name)
	if err != nil && !os.IsNotExist(errors.Cause(err)) {
		return nil, errors.Trace(err)
	} else if os.IsNotExist(errors.Cause(err)) {
		return &m, nil
	}
	defer f.Close()

	_, err = toml.DecodeReader(f, &m)

	return &m, err
}

func (m *masterInfo) Save(force bool) error {
	m.l.Lock()
	defer m.l.Unlock()

	n := time.Now()
	if !force && n.Sub(m.lastSaveTime) < time.Second {
		return nil
	}

	var buf bytes.Buffer
	e := toml.NewEncoder(&buf)

	e.Encode(m)

	var err error
	if err = ioutil2.WriteFileAtomic(m.name, buf.Bytes(), 0644); err != nil {
		log.Errorf("canal save master info to file %s err %v", m.name, err)
	}

	m.lastSaveTime = n

	return errors.Trace(err)
}

var lastLogTime int64 = 0
func (m *masterInfo) Update(name string, pos uint32) {
	m.l.Lock()
	m.Name = name
	m.Position = pos
	m.l.Unlock()
}

func (m *masterInfo) Pos() mysql.Position {
	var pos mysql.Position
	m.l.RLock()
	pos.Name = m.Name
	pos.Pos = m.Position
	m.l.RUnlock()

	return pos
}

func (m *masterInfo) Close() {
	m.Save(true)
}
<commit_msg>accidently deleted a log line<commit_after>package canal

import (
	"bytes"
	"os"
	"sync"
	"time"

	"github.com/BurntSushi/toml"
	"github.com/juju/errors"
	"github.com/ngaut/log"
	"github.com/zjh1943/go-mysql/mysql"
	"github.com/siddontang/go/ioutil2"
)

type masterInfo struct {
	Addr     string `toml:"addr"`
	Name     string `toml:"bin_name"`
	Position uint32 `toml:"bin_pos"`

	name string

	l sync.RWMutex

	lastSaveTime time.Time
}

func loadMasterInfo(name string) (*masterInfo, error) {
	var m masterInfo

	m.name = name

	f, err := os.Open(name)
	if err != nil && !os.IsNotExist(errors.Cause(err)) {
		return nil, errors.Trace(err)
	} else if os.IsNotExist(errors.Cause(err)) {
		return &m, nil
	}
	defer f.Close()

	_, err = toml.DecodeReader(f, &m)

	return &m, err
}

func (m *masterInfo) Save(force bool) error {
	m.l.Lock()
	defer m.l.Unlock()

	n := time.Now()
	if !force && n.Sub(m.lastSaveTime) < time.Second {
		return nil
	}

	var buf bytes.Buffer
	e := toml.NewEncoder(&buf)

	e.Encode(m)

	var err error
	if err = ioutil2.WriteFileAtomic(m.name, buf.Bytes(), 0644); err != nil {
		log.Errorf("canal save master info to file %s err %v", m.name, err)
	}

	m.lastSaveTime = n

	return errors.Trace(err)
}

var lastLogTime int64 = 0
func (m *masterInfo) Update(name string, pos uint32) {
	log.Debugf("update master position (%s, %d)", name, pos)
	m.l.Lock()
	m.Name = name
	m.Position = pos
	m.l.Unlock()
}

func (m *masterInfo) Pos() mysql.Position {
	var pos mysql.Position
	m.l.RLock()
	pos.Name = m.Name
	pos.Pos = m.Position
	m.l.RUnlock()

	return pos
}

func (m *masterInfo) Close() {
	m.Save(true)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/chrissnell/syslog"
)

type handler struct {
	// To simplify implementation of our handler we embed helper
	// syslog.BaseHandler struct.
	*syslog.BaseHandler
}

type LogentriesHostEntity struct {
	Response  string `json:"response"`
	Host      Host   `json:"host"`
	Host_key  string `json:"host_key"`
	Worker    string `json:"worker"`
	Agent_key string `json:"agent_key"`
}

type Host struct {
	C        float64 `json:"c"`
	Name     string  `json:"name"`
	Distver  string  `json:"distver"`
	Hostname string  `json:"hostname"`
	Object   string  `json:"object"`
	Distname string  `json:"distname"`
	Key      string  `json:"key"`
}

type LogentriesLogEntity struct {
	Response string `json:"response"`
	Log_key  string `json:"log_key"`
	Log      Log    `json:"log"`
}

type Log struct {
	Token     string  `json:"token"`
	Created   float64 `json:"created"`
	Name      string  `json:"name`
	Retention float64 `json:"retention"`
	Filename  string  `json:"filename"`
	Object    string  `json:"object"`
	Type      string  `json:"type"`
	Key       string  `json:"key"`
	Follow    string  `json:"folow"`
}

type LogLine struct {
	Line  syslog.Message
	Token string
}

var (
	logconsumerPtr        *string
	logentriesAPIKeyPtr   *string
	listenAddrPtr         *string
	logentities           = make(map[string]LogentriesLogEntity)
	hostentities          = make(map[string]LogentriesHostEntity)
	tokenchan             = make(chan string)
	logentities_filename  = "logentries-logentities.gob"
	hostentities_filename = "logentries-hostentities.gob"
)

func newHandler() *handler {
	msg := make(chan syslog.Message, 100)
	// Filter function name set to nil to disable filtering
	h := handler{syslog.NewBaseHandler(5, nil, false)}
	go h.mainLoop(msg)
	go ProcessLogMessage(msg)
	return &h
}

func (h *handler) mainLoop(msg chan syslog.Message) {
	for {
		m := h.Get()
		if m == nil {
			break
		}
		msg <- *m
	}
	fmt.Println("Exit handler")
	h.End()
}

func ProcessLogMessage(msg chan syslog.Message) {
	tokenfetchdone := make(chan bool, 1)
	logentrieschan := make(chan LogLine, 100)
	lh := make(chan struct{ host, log string }, 100)

	var logline LogLine

	for m := range msg {
		if m.Hostname == "" {
			m.Hostname = "NONE"
		}

		// for debugging
		// if m.Hostname == "testpiper001" {
		//	log.Printf("recv from %v: %v\n", m.Hostname, m.Content)
		// }

		go GetTokenForLog(tokenfetchdone, lh)
		lh <- struct{ host, log string }{m.Hostname, m.Tag}
		token := <-tokenchan
		<-tokenfetchdone

		logline.Token = token
		logline.Line = m

		go SendLogMessages(logentrieschan)
		logentrieschan <- logline

	}
}

func GetTokenForLog(tokenfetchdone chan bool, lh chan struct{ host, log string }) {
	select {
	case lht, msg_ok := <-lh:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {

			var hostentity LogentriesHostEntity
			var logentity LogentriesLogEntity

			l := strings.Join([]string{lht.host, lht.log}, "::")

			hostentity = hostentities[lht.host]
			if hostentity.Host.Key == "" {
				log.Printf("Registering host entity: %v\n", lht.host)
				hostentity = RegisterNewHost(lht.host)

				// Store our new host token in our map and sync it to disk
				hostentities[lht.host] = hostentity
				err := SyncHostEntitiesToDisk()
				if err != nil {
					log.Fatal(err)
				}
			}

			logentity = logentities[l]
			if logentity.Log.Token == "" {
				log.Printf("Registering new log entity [%v]: %v\n", lht.host, lht.log)
				logentity := RegisterNewLog(hostentity, lht.log)
				logentities[l] = logentity
				tokenchan <- logentity.Log.Token
				tokenfetchdone <- true
				err := SyncLogEntitiesToDisk()
				if err != nil {
					log.Fatal(err)
				}
			} else {
				tokenchan <- logentity.Log.Token
				tokenfetchdone <- true
			}
		}
	}
}

func DialLogEntries() (conn net.Conn, err error) {
	for {
		conn, err = net.Dial("tcp", *logconsumerPtr)
		if err == nil {
			return conn, err
		} else {
			fmt.Println("Could not connect to LogEntries log endpoint...retrying")
			// Wait for 5 seconds before redialing
			time.Sleep(5000 * time.Millisecond)
		}
	}
}

func SendLogMessages(msg chan LogLine) {
	conn, err := DialLogEntries()
	if err != nil {
		fmt.Println("Could not connect to LogEntries log endpoint ", err.Error())
	}

	select {
	case logline, msg_ok := <-msg:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {
			// For debugging...
			// if logline.Line.Hostname == "testpiper001" {
			//	log.Printf("send to LE from %v: %v\n", logline.Line.Hostname, logline.Line.Content)
			// }
			t := logline.Line.Time
			line := fmt.Sprintf("%v %v %v %v\n", logline.Token, t.Format(time.RFC3339), logline.Line.Hostname, logline.Line.Content)
			_, err = conn.Write([]byte(line))
			if err != nil {
				log.Print("Send to Logentries endpoint failed.")
				break
			}
			// fmt.Printf("Sending line: %v", line)
		}
	}
}

func SyncLogEntitiesToDisk() (err error) {
	m := new(bytes.Buffer)
	enc := gob.NewEncoder(m)
	enc.Encode(logentities)
	err = ioutil.WriteFile("logentries-logentities.gob", m.Bytes(), 0600)
	return (err)
}

func SyncHostEntitiesToDisk() (err error) {
	m := new(bytes.Buffer)
	enc := gob.NewEncoder(m)
	enc.Encode(hostentities)
	err = ioutil.WriteFile("logentries-hostentities.gob", m.Bytes(), 0600)
	return (err)
}

func LoadLogEntitiesFromDisk() (err error) {
	n, err := ioutil.ReadFile("logentries-logentities.gob")
	if err != nil {
		return (err)
	}
	p := bytes.NewBuffer(n)
	dec := gob.NewDecoder(p)
	err = dec.Decode(&logentities)
	return (err)
}

func LoadHostEntitiesFromDisk() (err error) {
	n, err := ioutil.ReadFile("logentries-hostentities.gob")
	if err != nil {
		return (err)
	}
	p := bytes.NewBuffer(n)
	dec := gob.NewDecoder(p)
	err = dec.Decode(&hostentities)
	return (err)
}

func RegisterNewHost(h string) (he LogentriesHostEntity) {
	v := url.Values{}
	v.Set("request", "register")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("name", h)
	v.Set("hostname", h)
	v.Set("distver", "")
	v.Set("system", "")
	v.Set("distname", "")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &he)
	return (he)
}

func RegisterNewLog(e LogentriesHostEntity, n string) (logentity LogentriesLogEntity) {
	v := url.Values{}
	v.Set("request", "new_log")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("host_key", e.Host.Key)
	v.Set("name", n)
	v.Set("filename", "")
	v.Set("retention", "-1")
	v.Set("source", "token")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &logentity)
	return (logentity)
}

func main() {

	logconsumerPtr = flag.String("consumer", "api.logentries.com:10000", "Logentries log consumer endpoint <host:port> (Default: api.logentries.com:10000)")
	logentriesAPIKeyPtr = flag.String("apikey", "", "Logentries API key")
	listenAddrPtr = flag.String("listen", "0.0.0.0:1987", "Host/port to listen for syslog messages <host:port> (Default: 0.0.0.0:1987)")

	flag.Parse()

	if *logentriesAPIKeyPtr == "" {
		log.Fatal("Must pass a Logentries API key. Use -h for help.")
	}

	if _, err := os.Stat(logentities_filename); err == nil {
		err = LoadLogEntitiesFromDisk()
		if err != nil {
			log.Fatal(err)
		}
	}

	if _, err := os.Stat(hostentities_filename); err == nil {
		err = LoadHostEntitiesFromDisk()
		if err != nil {
			log.Fatal(err)
		}
	}

	// Create a server with one handler and run one listen gorutine
	s := syslog.NewServer()
	s.AddAllowedRunes("-._")
	s.AddHandler(newHandler())
	s.Listen(*listenAddrPtr)

	// Wait for terminating signal
	sc := make(chan os.Signal, 2)
	signal.Notify(sc, syscall.SIGTERM, syscall.SIGINT)
	<-sc

	// Shutdown the server
	fmt.Println("Shutdown the server...")
	s.Shutdown()
	fmt.Println("Server is down")
}
<commit_msg>Experimental support for MySQL+groupcache token storage<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/chrissnell/syslog"
	_ "github.com/go-sql-driver/mysql"
	"github.com/golang/groupcache"
)

type handler struct {
	// To simplify implementation of our handler we embed helper
	// syslog.BaseHandler struct.
	*syslog.BaseHandler
}

type LogentriesHostEntity struct {
	Response  string `json:"response"`
	Host      Host   `json:"host"`
	Host_key  string `json:"host_key"`
	Worker    string `json:"worker"`
	Agent_key string `json:"agent_key"`
}

type Host struct {
	C        float64 `json:"c"`
	Name     string  `json:"name"`
	Distver  string  `json:"distver"`
	Hostname string  `json:"hostname"`
	Object   string  `json:"object"`
	Distname string  `json:"distname"`
	Key      string  `json:"key"`
}

type LogentriesLogEntity struct {
	Response string `json:"response"`
	Log_key  string `json:"log_key"`
	Log      Log    `json:"log"`
}

type Log struct {
	Token     string  `json:"token"`
	Created   float64 `json:"created"`
	Name      string  `json:"name`
	Retention float64 `json:"retention"`
	Filename  string  `json:"filename"`
	Object    string  `json:"object"`
	Type      string  `json:"type"`
	Key       string  `json:"key"`
	Follow    string  `json:"folow"`
}

type LogLine struct {
	Line  syslog.Message
	Token string
}

var (
	db                    *sql.DB
	hostTokenCache        *groupcache.Group
	logTokenCache         *groupcache.Group
	hostIDCache           *groupcache.Group
	ctx1, ctx2, ctx3      groupcache.Context
	logconsumerPtr        *string
	logentriesAPIKeyPtr   *string
	listenAddrPtr         *string
	logentities           = make(map[string]LogentriesLogEntity)
	hostentities          = make(map[string]LogentriesHostEntity)
	tokenchan             = make(chan string)
	dbConnectDone         = make(chan sql.DB)
	logentities_filename  = "logentries-logentities.gob"
	hostentities_filename = "logentries-hostentities.gob"
)

var hostTableSchema = `
CREATE TABLE IF NOT EXISTS host (
  id int NOT NULL AUTO_INCREMENT,
  hostname varchar(60) NOT NULL,
  host_key varchar(36) NOT NULL,
  PRIMARY KEY (id)
);
`

var logTableSchema = `
CREATE TABLE IF NOT EXISTS log (
  id int NOT NULL AUTO_INCREMENT,
  host_id int NOT NULL,
  logname varchar(80) NOT NULL,
  token varchar(36) NOT NULL,
  KEY host_ind (host_id),
  PRIMARY KEY (id),
  CONSTRAINT log_ibfk_1 FOREIGN KEY (host_id) REFERENCES host (id) ON DELETE CASCADE
);`

func newHandler() *handler {
	msg := make(chan syslog.Message, 100)
	// Filter function name set to nil to disable filtering
	h := handler{syslog.NewBaseHandler(5, nil, false)}
	go h.mainLoop(msg)
	go ProcessLogMessage(msg)
	return &h
}

func (h *handler) mainLoop(msg chan syslog.Message) {
	for {
		m := h.Get()
		if m == nil {
			break
		}
		msg <- *m
	}
	fmt.Println("Exit handler")
	h.End()
}

func ProcessLogMessage(msg chan syslog.Message) {
	tokenfetchdone := make(chan bool, 1)
	logentrieschan := make(chan LogLine, 100)
	lh := make(chan struct{ host, log string }, 100)

	var logline LogLine

	for m := range msg {
		if m.Hostname == "" {
			m.Hostname = "NONE"
		}

		// for debugging
		// if m.Hostname == "testpiper001" {
		//	log.Printf("recv from %v: %v\n", m.Hostname, m.Content)
		// }

		go GetTokenForLog(tokenfetchdone, lh)
		lh <- struct{ host, log string }{m.Hostname, m.Tag}
		token := <-tokenchan
		<-tokenfetchdone

		logline.Token = token
		logline.Line = m

		go SendLogMessages(logentrieschan)
		logentrieschan <- logline

	}
}

func GetTokenForLog(tokenfetchdone chan bool, lh chan struct{ host, log string }) {
	select {
	case lht, msg_ok := <-lh:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {

			var host_token, host_id, log_token string

			err := hostTokenCache.Get(ctx1, lht.host, groupcache.StringSink(&host_token))
			if err != nil {
				log.Fatal(err)
			}

			err = hostIDCache.Get(ctx3, lht.host, groupcache.StringSink(&host_id))
			if err != nil {
				log.Fatal(err)
			}

			if host_token == "" {
				log.Printf("Registering host entity: %v\n", lht.host)
				host_token = RegisterNewHost(lht.host)

				host_id = SaveHostTokenToDB(lht.host, host_token)
			}

			// log_token := GetLogTokenFromDB(lht.host, lht.log)
			hostandlog := lht.host + "::" + lht.log
			err = logTokenCache.Get(ctx3, hostandlog, groupcache.StringSink(&log_token))
			if err != nil {
				log.Fatal(err)
			}

			if log_token == "" {
				log.Printf("Registering new log entity [%v]: %v\n", lht.host, lht.log)
				log_token := RegisterNewLog(host_token, lht.log)

				log.Printf("Saving log token to db.  logname: %v  log_token: %v\n", lht.log, log_token)
				SaveLogTokenToDB(host_id, lht.log, log_token)

				tokenchan <- log_token
				tokenfetchdone <- true
			} else {
				tokenchan <- log_token
				tokenfetchdone <- true
			}
		}
	}
}

func GetHostTokenFromDB(h string) (token string) {

	log.Printf("Checking DB for host: %v\n", h)

	err := db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	err = db.QueryRow("SELECT token FROM host WHERE hostname = ?", h).Scan(&token)

	if err == sql.ErrNoRows {
		log.Print("Host does not exist in DB")
		return token
	}

	return token
}

func GetHostIDFromDB(h string) (host_id string) {

	log.Printf("Checking DB for host_id for host: %v\n", h)

	err := db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	err = db.QueryRow("SELECT id FROM host WHERE hostname = ?", h).Scan(&host_id)

	if err == sql.ErrNoRows {
		log.Print("Host does not exist in DB")
		return host_id
	}

	return host_id

}

func GetLogTokenFromDB(h, l string) (token string) {

	log.Printf("Checking DB for log: %v / %v\n", h, l)

	err := db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	err = db.QueryRow("SELECT log.token AS token FROM log INNER JOIN host ON log.host_id=host.id WHERE log.logname = ? AND host.hostname = ?", l, h).Scan(&token)

	if err == sql.ErrNoRows {
		return token
	}

	if err != nil {
		log.Fatal(err)
	}

	return (token)
}

func SaveHostTokenToDB(hostname string, token string) (host_id string) {

	err := db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	stmt, err := db.Prepare("INSERT INTO host (hostname, token) VALUES(?, ?)")
	if err != nil {
		log.Fatal(err)
	}

	res, err := stmt.Exec(hostname, token)
	if err != nil {
		log.Fatal(err)
	}

	host_id_int, _ := res.LastInsertId()
	host_id = string(host_id_int)
	rowCnt, _ := res.RowsAffected()
	log.Printf("INSERT INTO host sucessful.  Rows affected: %v", rowCnt)
	return host_id
}

func SaveLogTokenToDB(host_id, logname, token string) {

	err := db.Ping()
	if err != nil {
		log.Fatal(err)
	}

	stmt, err := db.Prepare("INSERT INTO log (host_id, logname, token) VALUES(?, ?, ?)")
	if err != nil {
		log.Fatalf("Error saving log token to DB: %v", err)
	}

	res, err := stmt.Exec(host_id, logname, token)
	if err != nil {
		log.Fatal(err)
	}

	rowCnt, _ := res.RowsAffected()
	log.Printf("INSERT INTO log sucessful.  Rows affected: %v", rowCnt)

}

func DialLogEntries() (conn net.Conn, err error) {
	for {
		conn, err = net.Dial("tcp", *logconsumerPtr)
		if err == nil {
			return conn, err
		} else {
			fmt.Println("Could not connect to LogEntries log endpoint...retrying")
			// Wait for 5 seconds before redialing
			time.Sleep(5000 * time.Millisecond)
		}
	}
}

func SendLogMessages(msg chan LogLine) {
	conn, err := DialLogEntries()
	if err != nil {
		fmt.Println("Could not connect to LogEntries log endpoint ", err.Error())
	}

	select {
	case logline, msg_ok := <-msg:
		if !msg_ok {
			fmt.Println("msg channel closed")
		} else {
			// For debugging...
			// if logline.Line.Hostname == "testpiper001" {
			//	log.Printf("send to LE from %v: %v\n", logline.Line.Hostname, logline.Line.Content)
			// }
			t := logline.Line.Time
			line := fmt.Sprintf("%v %v %v %v\n", logline.Token, t.Format(time.RFC3339), logline.Line.Hostname, logline.Line.Content)
			_, err = conn.Write([]byte(line))
			if err != nil {
				log.Print("Send to Logentries endpoint failed.")
				break
			}
			// fmt.Printf("Sending line: %v", line)
		}
	}
}

func RegisterNewHost(h string) (host_token string) {
	var he LogentriesHostEntity

	v := url.Values{}
	v.Set("request", "register")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("name", h)
	v.Set("hostname", h)
	v.Set("distver", "")
	v.Set("system", "")
	v.Set("distname", "")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &he)
	return (he.Host_key)
}

func RegisterNewLog(ht, n string) (log_token string) {
	var logentity LogentriesLogEntity

	v := url.Values{}
	v.Set("request", "new_log")
	v.Set("user_key", *logentriesAPIKeyPtr)
	v.Set("host_key", ht)
	v.Set("name", n)
	v.Set("filename", "")
	v.Set("retention", "-1")
	v.Set("source", "token")
	res, err := http.PostForm("http://api.logentries.com/", v)
	if err != nil {
		log.Fatal(err)
	}
	body, err := ioutil.ReadAll(res.Body)
	err = json.Unmarshal(body, &logentity)
	log_token = logentity.Log_key
	return (log_token)
}

func main() {
	var err error

	logconsumerPtr = flag.String("consumer", "api.logentries.com:10000", "Logentries log consumer endpoint <host:port> (Default: api.logentries.com:10000)")
	logentriesAPIKeyPtr = flag.String("apikey", "", "Logentries API key")
	listenAddrPtr = flag.String("listen", "0.0.0.0:1987", "Host/port to listen for syslog messages <host:port> (Default: 0.0.0.0:1987)")
	groupCachePeers := flag.String("peers", "", "groupcache peers <host:port> (Default: none)")
	groupCacheListenAddr := flag.String("cachelisten", "0.0.0.0:11000", "Interface to listen on for  <host:port> (Default: 0.0.0.0:11000)")

	flag.Parse()
	if *logentriesAPIKeyPtr == "" {
		log.Fatal("Must pass a Logentries API key. Use -h for help.")
	}

	// Set up groupcache peers
	peerSlice := make([]string, 1)

	if strings.Contains(*groupCachePeers, ",") {
		peerSlice = strings.Split(*groupCachePeers, ",")
	} else if *groupCachePeers != "" {
		peerSlice[0] = *groupCachePeers
	}

	if len(peerSlice) > 0 {
		for peer, value := range peerSlice {
			peerSlice[peer] = "http://" + value
		}
	}

	me := "http://" + *groupCacheListenAddr

	peerSlice = append(peerSlice, me)

	for peer, value := range peerSlice {
		log.Printf("Peer: %v  Value: %v\n", peer, value)
	}

	peers := groupcache.NewHTTPPool(me)
	peers.Set(peerSlice...)

	hostTokenCache = groupcache.NewGroup("HostTokenCache", 64<<10, groupcache.GetterFunc(
		func(ctx1 groupcache.Context, hostname string, dest groupcache.Sink) error {
			fmt.Printf("asking for %v from DB\n", hostname)
			result := GetHostTokenFromDB(hostname)
			dest.SetString(result)
			return nil
		}))

	logTokenCache = groupcache.NewGroup("LogTokenCache", 64<<10, groupcache.GetterFunc(
		func(ctx2 groupcache.Context, hostAndLog string, dest groupcache.Sink) error {
			s := strings.Split(hostAndLog, "::")
			hostname, logname := s[0], s[1]
			fmt.Printf("asking for %v, %v from DB\n", hostname, logname)
			result := GetLogTokenFromDB(hostname, logname)
			dest.SetString(result)
			return nil
		}))

	hostIDCache = groupcache.NewGroup("HostIDCache", 64<<10, groupcache.GetterFunc(
		func(ctx3 groupcache.Context, hostname string, dest groupcache.Sink) error {
			fmt.Printf("asking for %v from DB\n", hostname)
			result := GetHostIDFromDB(hostname)
			dest.SetString(string(result))
			return nil
		}))

	// Connect to the database
	db, err = sql.Open("mysql", "lerelay:UYsesKv(cjL7M4NUr}@tcp(f83ab18e40896fc40e90ce4e6b4a576f00544ea7.rackspaceclouddb.com:3306)/lerelay")
	if err != nil {
		log.Fatal(err)
	}

	// Create database schema if it doesn't exist already
	create, err := db.Prepare(hostTableSchema)
	if err != nil {
		log.Fatal(err)
	}

	_, err = create.Exec()
	if err != nil {
		log.Fatal(err)
	}

	create, err = db.Prepare(logTableSchema)
	if err != nil {
		log.Fatal(err)
	}

	_, err = create.Exec()
	if err != nil {
		log.Fatal(err)
	}

	// Create a server with one handler and run one listen gorutine
	s := syslog.NewServer()
	s.AddAllowedRunes("-._")
	s.AddHandler(newHandler())
	s.Listen(*listenAddrPtr)

	// Wait for terminating signal
	sc := make(chan os.Signal, 2)
	signal.Notify(sc, syscall.SIGTERM, syscall.SIGINT)
	<-sc

	// Shutdown the server
	fmt.Println("Shutdown the server...")
	s.Shutdown()
	fmt.Println("Server is down")

}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	//"errors"
	"container/list"
	"strconv"
	"sync"

	"github.com/reiver/go-porterstemmer"
)

var seperators []byte = []byte{10, 32, 44, 46}
var stopWords map[string]bool

//var dictionary map[string]*list.List
var groupLen int = 50

var dictionary = struct {
	sync.RWMutex
	m map[string]*list.List
}{m: make(map[string]*list.List)}

func main() {
	var corpusPath string
	flag.StringVar(&corpusPath, "corpus", "data/corpus", "File path of the corpus.")
	flag.Parse()
	dispatcher(corpusPath)
}

func dispatcher(corpusPath string) {
	file, err := os.Open(corpusPath)
	defer file.Close()
	if err != nil {
		panic(err)
	}
	collectStopWords()
	//dictionary = make(map[string]*list.List)
	scanner := bufio.NewScanner(io.Reader(file))

	scanner.Split(splitTokens)

	var (
		token string
		pos   uint32 = 0
		wg    sync.WaitGroup
	)
	for scanner.Scan() {
		token = scanner.Text()
		if len(token) > 0 {
			wg.Add(1)
			go addToken(token, pos, &wg)
			pos += 1
		}
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}
	wg.Wait()

	outFile, err := os.Create("res")
	defer outFile.Close()
	if err != nil {
		panic(err)
	}
	writer := bufio.NewWriter(io.Writer(outFile))
	writeIndex(writer)
}

func addToken(token string, pos uint32, wg *sync.WaitGroup) {
	defer wg.Done()

	// Check to see if token is a stop word
	_, ok := stopWords[token]
	if !ok {
		// Stem the token
		token = porterstemmer.StemString(token)
		dictionary.RLock()
		postingsList, ok := dictionary.m[token]
		dictionary.RUnlock()
		if ok {
			lastEl := postingsList.Back()
			lastGroup := lastEl.Value.([]uint32)
			if len(lastGroup) == groupLen {
				newGroup := []uint32{pos}
				postingsList.PushBack(newGroup)
			} else {
				lastGroup = append(lastGroup, pos)
				postingsList.Remove(lastEl)
				postingsList.PushBack(lastGroup)
			}
		} else {
			l := list.New()
			newGroup := []uint32{pos}
			l.PushBack(newGroup)
			dictionary.Lock()
			dictionary.m[token] = l
			dictionary.Unlock()
		}
	}
}

func writeIndex(writer *bufio.Writer) {
	for k, v := range dictionary.m {
		fmt.Println(k, v.Len(), len(v.Front().Value.([]uint32)))
		writer.WriteString("#" + k)
		var group []uint32
		for el := v.Front(); el != nil; el = el.Next() {
			group = el.Value.([]uint32)
			for _, posting := range group {
				val := strconv.FormatUint(uint64(posting), 10)
				writer.WriteString("," + string(val))
			}
		}
	}
	writer.Flush()
}

func splitTokens(data []byte, atEOF bool) (advance int, token []byte, err error) {
	var (
		found   = false
		started = false
		offset  = 0
	)
	for i := 0; i < len(data); i++ {
		found = false
		for s := 0; s < len(seperators); s++ {
			if data[i] == seperators[s] {
				advance = i + 1
				token = data[offset:i]
				if !started {
					offset += 1
				}
				found = true
				break
			}
		}
		if found {
			if started {
				break
			} else {
				continue
			}
		}
		if !started {
			started = true
		}
	}
	return
}

func collectStopWords() {
	file, err := os.Open("data/stopwords")
	defer file.Close()
	if err != nil {
		panic(err)
	}
	stopWords = make(map[string]bool)
	scanner := bufio.NewScanner(io.Reader(file))

	var word string
	for scanner.Scan() {
		word = scanner.Text()
		if len(word) > 0 {
			stopWords[word] = true
		}
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}
}
<commit_msg>Disable concurrency for addToken, add some error handling.<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"os"
	//"errors"
	"container/list"
	//"runtime"
	"sort"
	"strconv"
	"sync"
	"unsafe"

	"github.com/reiver/go-porterstemmer"
)

var seperators = [6]byte{10, 32, 34, 44, 46, 59}
var stopWords map[string]bool
var groupLen int = 50
var dictionary = struct {
	sync.RWMutex
	m map[string]*list.List
	keys []string
}{m: make(map[string]*list.List)}
var notStem = map[string]bool{
	"ION": true,
}
var indexMemoryLimit uint32 = 512 * 1024 * 1024
var indexMemoryConsumed = struct {
	sync.RWMutex
	v uint32
}{v: 0}

func main() {
	var corpusPath string
	flag.StringVar(&corpusPath, "corpus", "data/corpus", "File path of the corpus.")
	flag.Parse()

	fmt.Println("Creating index.")
	dispatcher(corpusPath)
}

func dispatcher(corpusPath string) {
	defer func() {
		err := recover()
		if err != nil {
			fmt.Println(err)
			os.Exit(1)
		}
	}()

	file, err := os.Open(corpusPath)
	defer file.Close()
	if err != nil {
		panic(err)
	}

	collectStopWords()
	scanner := bufio.NewScanner(io.Reader(file))
	scanner.Split(splitTokens)

	var (
		token string
		pos   uint32 = 0
	)
	for scanner.Scan() {
		token = scanner.Text()
		if len(token) > 0 {
			addToken(token, pos)
			pos += 1
		}
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}

	outFile, err := os.Create("res")
	defer outFile.Close()
	if err != nil {
		panic(err)
	}
	sort.Strings(dictionary.keys)
	writeIndex(io.Writer(outFile))
}

func addToken(token string, pos uint32) {
	defer func() {
		err := recover()
		if err != nil {
			fmt.Println(err)
			fmt.Println(token)
			os.Exit(1)
		}
	}()

	// Check to see if token is a stop word
	var memoryConsumed uint32 = 0
	if _, ok := stopWords[token]; !ok {
		// Stem the token
		if _, ok := notStem[token]; !ok {
			token = porterstemmer.StemString(token)
		}
		dictionary.RLock()
		postingsList, ok := dictionary.m[token]
		dictionary.RUnlock()
		if ok {
			lastEl := postingsList.Back()
			lastGroup := lastEl.Value.([]uint32)
			if len(lastGroup) == groupLen {
				newGroup := []uint32{pos}
				postingsList.PushBack(newGroup)
			} else {
				lastGroup = append(lastGroup, pos)
				postingsList.Remove(lastEl)
				postingsList.PushBack(lastGroup)
			}
		} else {
			l := list.New()
			newGroup := []uint32{pos}
			l.PushBack(newGroup)
			dictionary.Lock()
			dictionary.m[token] = l
			dictionary.Unlock()
			dictionary.keys = append(dictionary.keys, token)
			//fmt.Println(token, uint32(unsafe.Sizeof(token)))
			memoryConsumed += uint32(unsafe.Sizeof(token))
		}
	}
	indexMemoryConsumed.Lock()
	indexMemoryConsumed.v += memoryConsumed
	indexMemoryConsumed.Unlock()
}

func writeIndex(w io.Writer) {
	writer := bufio.NewWriter(w)
	var v *list.List
	for _, k := range dictionary.keys {
		v = dictionary.m[k]
		//fmt.Println(k, v.Len(), len(v.Front().Value.([]uint32)))
		writer.WriteString("#" + k)
		var group []uint32
		for el := v.Front(); el != nil; el = el.Next() {
			group = el.Value.([]uint32)
			for _, posting := range group {
				val := strconv.FormatUint(uint64(posting), 10)
				writer.WriteString("," + string(val))
			}
		}
	}
	writer.Flush()
}

func splitTokens(data []byte, atEOF bool) (advance int, token []byte, err error) {
	var (
		found   = false
		started = false
		offset  = 0
	)
	for i := 0; i < len(data); i++ {
		found = false
		for s := 0; s < len(seperators); s++ {
			if data[i] == seperators[s] {
				advance = i + 1
				token = data[offset:i]
				if !started {
					offset += 1
				}
				found = true
				break
			}
		}
		if found {
			if started {
				break
			} else {
				continue
			}
		}
		if !started {
			started = true
		}
	}
	return
}

func collectStopWords() {
	file, err := os.Open("data/stopwords")
	defer file.Close()
	if err != nil {
		panic(err)
	}
	stopWords = make(map[string]bool)
	scanner := bufio.NewScanner(io.Reader(file))

	var word string
	for scanner.Scan() {
		word = scanner.Text()
		if len(word) > 0 {
			stopWords[word] = true
		}
	}
	if err := scanner.Err(); err != nil {
		panic(err)
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package casbin

import (
	"testing"

	"github.com/casbin/casbin/util"
)

func testStringList(t *testing.T, title string, f func() []string, res []string) {
	t.Helper()
	myRes := f()
	t.Log(title+": ", myRes)

	if !util.ArrayEquals(res, myRes) {
		t.Error(title+": ", myRes, ", supposed to be ", res)
	}
}

func TestGetList(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testStringList(t, "Subjects", e.GetAllSubjects, []string{"alice", "bob", "data2_admin"})
	testStringList(t, "Objeccts", e.GetAllObjects, []string{"data1", "data2"})
	testStringList(t, "Actions", e.GetAllActions, []string{"read", "write"})
	testStringList(t, "Roles", e.GetAllRoles, []string{"data2_admin"})
}

func testGetPolicy(t *testing.T, e *Enforcer, res [][]string) {
	t.Helper()
	myRes := e.GetPolicy()
	t.Log("Policy: ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Policy: ", myRes, ", supposed to be ", res)
	}
}

func testGetFilteredPolicy(t *testing.T, e *Enforcer, fieldIndex int, res [][]string, fieldValues ...string) {
	t.Helper()
	myRes := e.GetFilteredPolicy(fieldIndex, fieldValues...)
	t.Log("Policy for ", util.ParamsToString(fieldValues...), ": ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Policy for ", util.ParamsToString(fieldValues...), ": ", myRes, ", supposed to be ", res)
	}
}

func testGetGroupingPolicy(t *testing.T, e *Enforcer, res [][]string) {
	t.Helper()
	myRes := e.GetGroupingPolicy()
	t.Log("Grouping policy: ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Grouping policy: ", myRes, ", supposed to be ", res)
	}
}

func testGetFilteredGroupingPolicy(t *testing.T, e *Enforcer, fieldIndex int, res [][]string, fieldValues ...string) {
	t.Helper()
	myRes := e.GetFilteredGroupingPolicy(fieldIndex, fieldValues...)
	t.Log("Grouping policy for ", util.ParamsToString(fieldValues...), ": ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Grouping policy for ", util.ParamsToString(fieldValues...), ": ", myRes, ", supposed to be ", res)
	}
}

func testHasPolicy(t *testing.T, e *Enforcer, policy []string, res bool) {
	t.Helper()
	myRes := e.HasPolicy(policy)
	t.Log("Has policy ", util.ArrayToString(policy), ": ", myRes)

	if res != myRes {
		t.Error("Has policy ", util.ArrayToString(policy), ": ", myRes, ", supposed to be ", res)
	}
}

func testHasGroupingPolicy(t *testing.T, e *Enforcer, policy []string, res bool) {
	t.Helper()
	myRes := e.HasGroupingPolicy(policy)
	t.Log("Has grouping policy ", util.ArrayToString(policy), ": ", myRes)

	if res != myRes {
		t.Error("Has grouping policy ", util.ArrayToString(policy), ": ", myRes, ", supposed to be ", res)
	}
}

func TestGetPolicyAPI(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testGetPolicy(t, e, [][]string{
		{"alice", "data1", "read"},
		{"bob", "data2", "write"},
		{"data2_admin", "data2", "read"},
		{"data2_admin", "data2", "write"}})

	testGetFilteredPolicy(t, e, 0, [][]string{{"alice", "data1", "read"}}, "alice")
	testGetFilteredPolicy(t, e, 0, [][]string{{"bob", "data2", "write"}}, "bob")
	testGetFilteredPolicy(t, e, 0, [][]string{{"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}}, "data2_admin")
	testGetFilteredPolicy(t, e, 1, [][]string{{"alice", "data1", "read"}}, "data1")
	testGetFilteredPolicy(t, e, 1, [][]string{{"bob", "data2", "write"}, {"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}}, "data2")
	testGetFilteredPolicy(t, e, 2, [][]string{{"alice", "data1", "read"}, {"data2_admin", "data2", "read"}}, "read")
	testGetFilteredPolicy(t, e, 2, [][]string{{"bob", "data2", "write"}, {"data2_admin", "data2", "write"}}, "write")

	testGetFilteredPolicy(t, e, 0, [][]string{{"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}}, "data2_admin", "data2")
	// Note: "" (empty string) in fieldValues means matching all values.
	testGetFilteredPolicy(t, e, 0, [][]string{{"data2_admin", "data2", "read"}}, "data2_admin", "", "read")
	testGetFilteredPolicy(t, e, 1, [][]string{{"bob", "data2", "write"}, {"data2_admin", "data2", "write"}}, "data2", "write")

	testHasPolicy(t, e, []string{"alice", "data1", "read"}, true)
	testHasPolicy(t, e, []string{"bob", "data2", "write"}, true)
	testHasPolicy(t, e, []string{"alice", "data2", "read"}, false)
	testHasPolicy(t, e, []string{"bob", "data3", "write"}, false)

	testGetGroupingPolicy(t, e, [][]string{{"alice", "data2_admin"}})

	testGetFilteredGroupingPolicy(t, e, 0, [][]string{{"alice", "data2_admin"}}, "alice")
	testGetFilteredGroupingPolicy(t, e, 0, [][]string{}, "bob")
	testGetFilteredGroupingPolicy(t, e, 1, [][]string{}, "data1_admin")
	testGetFilteredGroupingPolicy(t, e, 1, [][]string{{"alice", "data2_admin"}}, "data2_admin")
	// Note: "" (empty string) in fieldValues means matching all values.
	testGetFilteredGroupingPolicy(t, e, 0, [][]string{{"alice", "data2_admin"}}, "", "data2_admin")

	testHasGroupingPolicy(t, e, []string{"alice", "data2_admin"}, true)
	testHasGroupingPolicy(t, e, []string{"bob", "data2_admin"}, false)
}

func TestModifyPolicyAPI(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testGetPolicy(t, e, [][]string{
		{"alice", "data1", "read"},
		{"bob", "data2", "write"},
		{"data2_admin", "data2", "read"},
		{"data2_admin", "data2", "write"}})

	e.RemovePolicy("alice", "data1", "read")
	e.RemovePolicy("bob", "data2", "write")
	e.RemovePolicy("alice", "data1", "read")
	e.AddPolicy("eve", "data3", "read")
	e.AddPolicy("eve", "data3", "read")

	namedPolicy := []string{"eve", "data3", "read"}
	e.RemoveNamedPolicy("p", namedPolicy)
	e.AddNamedPolicy("p", namedPolicy)

	testGetPolicy(t, e, [][]string{
		{"data2_admin", "data2", "read"},
		{"data2_admin", "data2", "write"},
		{"eve", "data3", "read"}})

	e.RemoveFilteredPolicy(1, "data2")

	testGetPolicy(t, e, [][]string{{"eve", "data3", "read"}})
}

func TestModifyGroupingPolicyAPI(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testGetRoles(t, e, "alice", []string{"data2_admin"})
	testGetRoles(t, e, "bob", []string{})
	testGetRoles(t, e, "eve", []string{})
	testGetRoles(t, e, "non_exist", []string{})

	e.RemoveGroupingPolicy("alice", "data2_admin")
	e.AddGroupingPolicy("bob", "data1_admin")
	e.AddGroupingPolicy("eve", "data3_admin")

	namedGroupingPolicy := []string{"alice", "data2_admin"}
	testGetRoles(t, e, "alice", []string{})
	e.AddNamedGroupingPolicy("g", namedGroupingPolicy)
	testGetRoles(t, e, "alice", []string{"data2_admin"})
	e.RemoveNamedGroupingPolicy("g", namedGroupingPolicy)

	testGetRoles(t, e, "alice", []string{})
	testGetRoles(t, e, "bob", []string{"data1_admin"})
	testGetRoles(t, e, "eve", []string{"data3_admin"})
	testGetRoles(t, e, "non_exist", []string{})

	testGetUsers(t, e, "data1_admin", []string{"bob"})
	testGetUsers(t, e, "data2_admin", []string{})
	testGetUsers(t, e, "data3_admin", []string{"eve"})

	e.RemoveFilteredGroupingPolicy(0, "bob")

	testGetRoles(t, e, "alice", []string{})
	testGetRoles(t, e, "bob", []string{})
	testGetRoles(t, e, "eve", []string{"data3_admin"})
	testGetRoles(t, e, "non_exist", []string{})

	testGetUsers(t, e, "data1_admin", []string{})
	testGetUsers(t, e, "data2_admin", []string{})
	testGetUsers(t, e, "data3_admin", []string{"eve"})
}
<commit_msg>Fix test title typo from 'Objeccts' to 'Objects'<commit_after>// Copyright 2017 The casbin Authors. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package casbin

import (
	"testing"

	"github.com/casbin/casbin/util"
)

func testStringList(t *testing.T, title string, f func() []string, res []string) {
	t.Helper()
	myRes := f()
	t.Log(title+": ", myRes)

	if !util.ArrayEquals(res, myRes) {
		t.Error(title+": ", myRes, ", supposed to be ", res)
	}
}

func TestGetList(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testStringList(t, "Subjects", e.GetAllSubjects, []string{"alice", "bob", "data2_admin"})
	testStringList(t, "Objects", e.GetAllObjects, []string{"data1", "data2"})
	testStringList(t, "Actions", e.GetAllActions, []string{"read", "write"})
	testStringList(t, "Roles", e.GetAllRoles, []string{"data2_admin"})
}

func testGetPolicy(t *testing.T, e *Enforcer, res [][]string) {
	t.Helper()
	myRes := e.GetPolicy()
	t.Log("Policy: ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Policy: ", myRes, ", supposed to be ", res)
	}
}

func testGetFilteredPolicy(t *testing.T, e *Enforcer, fieldIndex int, res [][]string, fieldValues ...string) {
	t.Helper()
	myRes := e.GetFilteredPolicy(fieldIndex, fieldValues...)
	t.Log("Policy for ", util.ParamsToString(fieldValues...), ": ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Policy for ", util.ParamsToString(fieldValues...), ": ", myRes, ", supposed to be ", res)
	}
}

func testGetGroupingPolicy(t *testing.T, e *Enforcer, res [][]string) {
	t.Helper()
	myRes := e.GetGroupingPolicy()
	t.Log("Grouping policy: ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Grouping policy: ", myRes, ", supposed to be ", res)
	}
}

func testGetFilteredGroupingPolicy(t *testing.T, e *Enforcer, fieldIndex int, res [][]string, fieldValues ...string) {
	t.Helper()
	myRes := e.GetFilteredGroupingPolicy(fieldIndex, fieldValues...)
	t.Log("Grouping policy for ", util.ParamsToString(fieldValues...), ": ", myRes)

	if !util.Array2DEquals(res, myRes) {
		t.Error("Grouping policy for ", util.ParamsToString(fieldValues...), ": ", myRes, ", supposed to be ", res)
	}
}

func testHasPolicy(t *testing.T, e *Enforcer, policy []string, res bool) {
	t.Helper()
	myRes := e.HasPolicy(policy)
	t.Log("Has policy ", util.ArrayToString(policy), ": ", myRes)

	if res != myRes {
		t.Error("Has policy ", util.ArrayToString(policy), ": ", myRes, ", supposed to be ", res)
	}
}

func testHasGroupingPolicy(t *testing.T, e *Enforcer, policy []string, res bool) {
	t.Helper()
	myRes := e.HasGroupingPolicy(policy)
	t.Log("Has grouping policy ", util.ArrayToString(policy), ": ", myRes)

	if res != myRes {
		t.Error("Has grouping policy ", util.ArrayToString(policy), ": ", myRes, ", supposed to be ", res)
	}
}

func TestGetPolicyAPI(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testGetPolicy(t, e, [][]string{
		{"alice", "data1", "read"},
		{"bob", "data2", "write"},
		{"data2_admin", "data2", "read"},
		{"data2_admin", "data2", "write"}})

	testGetFilteredPolicy(t, e, 0, [][]string{{"alice", "data1", "read"}}, "alice")
	testGetFilteredPolicy(t, e, 0, [][]string{{"bob", "data2", "write"}}, "bob")
	testGetFilteredPolicy(t, e, 0, [][]string{{"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}}, "data2_admin")
	testGetFilteredPolicy(t, e, 1, [][]string{{"alice", "data1", "read"}}, "data1")
	testGetFilteredPolicy(t, e, 1, [][]string{{"bob", "data2", "write"}, {"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}}, "data2")
	testGetFilteredPolicy(t, e, 2, [][]string{{"alice", "data1", "read"}, {"data2_admin", "data2", "read"}}, "read")
	testGetFilteredPolicy(t, e, 2, [][]string{{"bob", "data2", "write"}, {"data2_admin", "data2", "write"}}, "write")

	testGetFilteredPolicy(t, e, 0, [][]string{{"data2_admin", "data2", "read"}, {"data2_admin", "data2", "write"}}, "data2_admin", "data2")
	// Note: "" (empty string) in fieldValues means matching all values.
	testGetFilteredPolicy(t, e, 0, [][]string{{"data2_admin", "data2", "read"}}, "data2_admin", "", "read")
	testGetFilteredPolicy(t, e, 1, [][]string{{"bob", "data2", "write"}, {"data2_admin", "data2", "write"}}, "data2", "write")

	testHasPolicy(t, e, []string{"alice", "data1", "read"}, true)
	testHasPolicy(t, e, []string{"bob", "data2", "write"}, true)
	testHasPolicy(t, e, []string{"alice", "data2", "read"}, false)
	testHasPolicy(t, e, []string{"bob", "data3", "write"}, false)

	testGetGroupingPolicy(t, e, [][]string{{"alice", "data2_admin"}})

	testGetFilteredGroupingPolicy(t, e, 0, [][]string{{"alice", "data2_admin"}}, "alice")
	testGetFilteredGroupingPolicy(t, e, 0, [][]string{}, "bob")
	testGetFilteredGroupingPolicy(t, e, 1, [][]string{}, "data1_admin")
	testGetFilteredGroupingPolicy(t, e, 1, [][]string{{"alice", "data2_admin"}}, "data2_admin")
	// Note: "" (empty string) in fieldValues means matching all values.
	testGetFilteredGroupingPolicy(t, e, 0, [][]string{{"alice", "data2_admin"}}, "", "data2_admin")

	testHasGroupingPolicy(t, e, []string{"alice", "data2_admin"}, true)
	testHasGroupingPolicy(t, e, []string{"bob", "data2_admin"}, false)
}

func TestModifyPolicyAPI(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testGetPolicy(t, e, [][]string{
		{"alice", "data1", "read"},
		{"bob", "data2", "write"},
		{"data2_admin", "data2", "read"},
		{"data2_admin", "data2", "write"}})

	e.RemovePolicy("alice", "data1", "read")
	e.RemovePolicy("bob", "data2", "write")
	e.RemovePolicy("alice", "data1", "read")
	e.AddPolicy("eve", "data3", "read")
	e.AddPolicy("eve", "data3", "read")

	namedPolicy := []string{"eve", "data3", "read"}
	e.RemoveNamedPolicy("p", namedPolicy)
	e.AddNamedPolicy("p", namedPolicy)

	testGetPolicy(t, e, [][]string{
		{"data2_admin", "data2", "read"},
		{"data2_admin", "data2", "write"},
		{"eve", "data3", "read"}})

	e.RemoveFilteredPolicy(1, "data2")

	testGetPolicy(t, e, [][]string{{"eve", "data3", "read"}})
}

func TestModifyGroupingPolicyAPI(t *testing.T) {
	e := NewEnforcer("examples/rbac_model.conf", "examples/rbac_policy.csv")

	testGetRoles(t, e, "alice", []string{"data2_admin"})
	testGetRoles(t, e, "bob", []string{})
	testGetRoles(t, e, "eve", []string{})
	testGetRoles(t, e, "non_exist", []string{})

	e.RemoveGroupingPolicy("alice", "data2_admin")
	e.AddGroupingPolicy("bob", "data1_admin")
	e.AddGroupingPolicy("eve", "data3_admin")

	namedGroupingPolicy := []string{"alice", "data2_admin"}
	testGetRoles(t, e, "alice", []string{})
	e.AddNamedGroupingPolicy("g", namedGroupingPolicy)
	testGetRoles(t, e, "alice", []string{"data2_admin"})
	e.RemoveNamedGroupingPolicy("g", namedGroupingPolicy)

	testGetRoles(t, e, "alice", []string{})
	testGetRoles(t, e, "bob", []string{"data1_admin"})
	testGetRoles(t, e, "eve", []string{"data3_admin"})
	testGetRoles(t, e, "non_exist", []string{})

	testGetUsers(t, e, "data1_admin", []string{"bob"})
	testGetUsers(t, e, "data2_admin", []string{})
	testGetUsers(t, e, "data3_admin", []string{"eve"})

	e.RemoveFilteredGroupingPolicy(0, "bob")

	testGetRoles(t, e, "alice", []string{})
	testGetRoles(t, e, "bob", []string{})
	testGetRoles(t, e, "eve", []string{"data3_admin"})
	testGetRoles(t, e, "non_exist", []string{})

	testGetUsers(t, e, "data1_admin", []string{})
	testGetUsers(t, e, "data2_admin", []string{})
	testGetUsers(t, e, "data3_admin", []string{"eve"})
}
<|endoftext|>


<commit_before>package controllers

import (
	"context"
	"fmt"
	"net/http"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/google/uuid"
	"github.com/microcosm-cc/bluemonday"
	"golang.org/x/crypto/bcrypt"
)

type UserCustomClaim struct {
	UUID uuid.UUID `json:"uuid"`
	jwt.StandardClaims
}

// Login, Logout, Sign-up

// getUser is a useful function for taking the user claim from the jwt
func getUser(tokenString string, hmacSecret []byte) (uuid.UUID, error) {
	u := *new(uuid.UUID)
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodECDSA); !ok {
			return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
		}
		return hmacSecret, nil
	})
	if err != nil {
		return u, err
	}
	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		u, err = uuid.Parse(claims["uuid"].(string))
		return u, err
	}
	return u, err
}

// UserCtx loads the user into the context if it exists
func (env *Env) UserCtx(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cookie, err := r.Cookie("authentication")
		if err != nil {
			next.ServeHTTP(w, r)
		}
		ID, err := uuid.Parse(cookie.Value)
		if err != nil {
			next.ServeHTTP(w, r)
		}
		user, err := env.DB.GetUserByID(ID)
		if err != nil {
			next.ServeHTTP(w, r)
		}
		ctx := context.WithValue(r.Context(), "user", user)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// Login takes a user and password from a login form and checks it against
// the hashed password in the database sets jwt if accepted
func (env *Env) Login(w http.ResponseWriter, r *http.Request) {
	// Clean everything but the password (password is hashed)
	s := bluemonday.UGCPolicy()
	uname := s.Sanitize(r.FormValue("user"))
	password := r.FormValue("password")
	t := time.Now().Add(time.Hour)

	// Database call and bcrypt compare hashed passwords
	u, err := env.DB.GetUserByUname(uname)
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
	}
	err = bcrypt.CompareHashAndPassword(u.Digest, []byte(password+env.salt))
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	// Generate new jwt with uuid claim and add cookie
	token := jwt.NewWithClaims(jwt.SigningMethodES512, UserCustomClaim{
		u.ID,
		jwt.StandardClaims{
			ExpiresAt: t.Unix(),
		},
	})
	tokenString, err := token.SignedString(env.hmac)
	cookie := http.Cookie{
		Name:     "authentication",
		Value:    tokenString,
		Expires:  t,
		HttpOnly: true,
	}
	http.SetCookie(w, &cookie)
	w.WriteHeader(http.StatusOK)
}

func (env *Env) Logout(w http.ResponseWriter, r *http.Request) {

}

// CreateAccount takes user, email, gpg, password, password2 from a form
// cleans any data that might show up on a page and returns 500 if the two passwords do not match
// 409 if there is an error in hashing and 200 if the user is added to the database successfully
func (env *Env) CreateAccount(w http.ResponseWriter, r *http.Request) {
	// Clean everything but the password (password is hashed)
	s := bluemonday.UGCPolicy()
	user := s.Sanitize(r.FormValue("user"))
	email := s.Sanitize(r.FormValue("email"))
	gpg := s.Sanitize(r.FormValue("gpg"))
	password := r.FormValue("password")
	password2 := r.FormValue("password2")

	// Are passwords the same?
	if password != password2 {
		w.WriteHeader(http.StatusUnauthorized)
	}

	// Generate the hash if there is an error in hashing we'll return 500
	digest, err := bcrypt.GenerateFromPassword([]byte(password+env.salt), 10)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	// Ignore the user struct returned and report back 409 if error or 200 if ok
	_, err = env.DB.InsertUser(user, digest, "MEMBER", email, gpg)
	if err != nil {
		w.WriteHeader(http.StatusConflict)
	}
	w.WriteHeader(http.StatusAccepted)
}
<commit_msg>logout will now delete cookie<commit_after>package controllers

import (
	"context"
	"fmt"
	"net/http"
	"time"

	jwt "github.com/dgrijalva/jwt-go"
	"github.com/google/uuid"
	"github.com/microcosm-cc/bluemonday"
	"golang.org/x/crypto/bcrypt"
)

type UserCustomClaim struct {
	UUID uuid.UUID `json:"uuid"`
	jwt.StandardClaims
}

// Login, Logout, Sign-up

// getUser is a useful function for taking the user claim from the jwt
func getUser(tokenString string, hmacSecret []byte) (uuid.UUID, error) {
	u := *new(uuid.UUID)
	token, err := jwt.Parse(tokenString, func(token *jwt.Token) (interface{}, error) {
		if _, ok := token.Method.(*jwt.SigningMethodECDSA); !ok {
			return nil, fmt.Errorf("Unexpected signing method: %v", token.Header["alg"])
		}
		return hmacSecret, nil
	})
	if err != nil {
		return u, err
	}
	if claims, ok := token.Claims.(jwt.MapClaims); ok && token.Valid {
		u, err = uuid.Parse(claims["uuid"].(string))
		return u, err
	}
	return u, err
}

// UserCtx loads the user into the context if it exists
func (env *Env) UserCtx(next http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		cookie, err := r.Cookie("authentication")
		if err != nil {
			next.ServeHTTP(w, r)
		}
		ID, err := uuid.Parse(cookie.Value)
		if err != nil {
			next.ServeHTTP(w, r)
		}
		user, err := env.DB.GetUserByID(ID)
		if err != nil {
			next.ServeHTTP(w, r)
		}
		ctx := context.WithValue(r.Context(), "user", user)
		next.ServeHTTP(w, r.WithContext(ctx))
	})
}

// Login takes a user and password from a login form and checks it against
// the hashed password in the database sets jwt if accepted
func (env *Env) Login(w http.ResponseWriter, r *http.Request) {
	// Clean everything but the password (password is hashed)
	s := bluemonday.UGCPolicy()
	uname := s.Sanitize(r.FormValue("user"))
	password := r.FormValue("password")
	t := time.Now().Add(time.Hour)

	// Database call and bcrypt compare hashed passwords
	u, err := env.DB.GetUserByUname(uname)
	if err != nil {
		w.WriteHeader(http.StatusNotFound)
	}
	err = bcrypt.CompareHashAndPassword(u.Digest, []byte(password+env.salt))
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	// Generate new jwt with uuid claim and add cookie
	token := jwt.NewWithClaims(jwt.SigningMethodES512, UserCustomClaim{
		u.ID,
		jwt.StandardClaims{
			ExpiresAt: t.Unix(),
		},
	})
	tokenString, err := token.SignedString(env.hmac)
	cookie := http.Cookie{
		Name:     "authentication",
		Value:    tokenString,
		Expires:  t,
		HttpOnly: true,
	}
	http.SetCookie(w, &cookie)
	w.WriteHeader(http.StatusOK)
}

func (env *Env) Logout(w http.ResponseWriter, r *http.Request) {
	// Setting MaxAge to < 0 will delete cookie now
	cookie := http.Cookie{
		Name:     "authentication",
		MaxAge:   -1,
		HttpOnly: true,
	}
	http.SetCookie(w, &cookie)
	w.WriteHeader(http.StatusOK)
}

// CreateAccount takes user, email, gpg, password, password2 from a form
// cleans any data that might show up on a page and returns 500 if the two passwords do not match
// 409 if there is an error in hashing and 200 if the user is added to the database successfully
func (env *Env) CreateAccount(w http.ResponseWriter, r *http.Request) {
	// Clean everything but the password (password is hashed)
	s := bluemonday.UGCPolicy()
	user := s.Sanitize(r.FormValue("user"))
	email := s.Sanitize(r.FormValue("email"))
	gpg := s.Sanitize(r.FormValue("gpg"))
	password := r.FormValue("password")
	password2 := r.FormValue("password2")

	// Are passwords the same?
	if password != password2 {
		w.WriteHeader(http.StatusUnauthorized)
	}

	// Generate the hash if there is an error in hashing we'll return 500
	digest, err := bcrypt.GenerateFromPassword([]byte(password+env.salt), 10)
	if err != nil {
		w.WriteHeader(http.StatusInternalServerError)
	}

	// Ignore the user struct returned and report back 409 if error or 202 if ok
	_, err = env.DB.InsertUser(user, digest, "MEMBER", email, gpg)
	if err != nil {
		w.WriteHeader(http.StatusConflict)
	}
	w.WriteHeader(http.StatusAccepted)
}
<|endoftext|>