<commit_before>package models

import (
	"fmt"
	"time"

	"github.com/koding/bongo"
)

// NotificationActivity stores each user NotificationActivity related to notification content.
// When a user makes duplicate NotificationActivity for the same content
// old one is set as obsolete and new one is added to NotificationActivity table
type NotificationActivity struct {
	// unique identifier of NotificationActivity
	Id int64 `json:"id"`

	// notification content foreign key
	NotificationContentId int64 `json:"notificationContentId" sql:"NOT NULL"`

	// notifier account foreign key
	ActorId int64 `json:"actorId,string" sql:"NOT NULL"`

	// activity creation time
	CreatedAt time.Time `json:"createdAt" sql:"NOT NULL"`

	// activity obsolete information
	Obsolete bool `json:"obsolete" sql:"NOT NULL"`
}

func (a *NotificationActivity) BeforeCreate() {
	a.CreatedAt = time.Now()
}

func (a *NotificationActivity) BeforeUpdate() {
	a.Obsolete = true
}

func (a *NotificationActivity) GetId() int64 {
	return a.Id
}

func NewNotificationActivity() *NotificationActivity {
	return &NotificationActivity{}
}

func (a NotificationActivity) TableName() string {
	return "notification.notification_activity"
}

// Create method creates a new activity with obsolete field set as false
// If there already exists one activity with same ActorId and
// NotificationContentId pair, old one is set as obsolete, and
// new one is created
func (a *NotificationActivity) Create() error {
	s := map[string]interface{}{
		"notification_content_id": a.NotificationContentId,
		"actor_id":                a.ActorId,
		"obsolete":                false,
	}

	q := bongo.NewQS(s)
	found := true
	if err := a.One(q); err != nil {
		if err != bongo.RecordNotFound {
			return err
		}
		found = false
	}

	if found {
		if err := bongo.B.Update(a); err != nil {
			return err
		}
		a.Id = 0
		a.Obsolete = false
	}

	return bongo.B.Create(a)
}

func (a *NotificationActivity) FetchByContentIds(ids []int64) ([]NotificationActivity, error) {
	activities := make([]NotificationActivity, 0)
	err := bongo.B.DB.Table(a.TableName()).
		Where("notification_content_id IN (?)", ids).
		Order("id asc").
		Find(&activities).Error

	if err != nil {
		return nil, err
	}

	return activities, nil
}

func (a *NotificationActivity) FetchMapByContentIds(ids []int64) (map[int64][]NotificationActivity, error) {
	if len(ids) == 0 {
		return make(map[int64][]NotificationActivity), nil
	}
	aList, err := a.FetchByContentIds(ids)
	if err != nil {
		return nil, err
	}

	aMap := make(map[int64][]NotificationActivity)
	for _, activity := range aList {
		aMap[activity.NotificationContentId] = append(aMap[activity.NotificationContentId], activity)
	}

	return aMap, nil
}

func (a *NotificationActivity) Fetch() error {

	return bongo.B.Fetch(a)
}

func (a *NotificationActivity) One(q *bongo.Query) error {

	return bongo.B.One(a, a, q)
}

func (a *NotificationActivity) Some(data interface{}, q *bongo.Query) error {

	return bongo.B.Some(a, data, q)
}

func (a *NotificationActivity) LastActivity() error {
	s := map[string]interface{}{
		"notification_content_id": a.NotificationContentId,
		"obsolete":                false,
	}

	q := bongo.NewQS(s)
	q.Sort = map[string]string{
		"id": "DESC",
	}

	return a.One(q)
}

func (a *NotificationActivity) FetchContent() (*NotificationContent, error) {
	if a.NotificationContentId == 0 {
		return nil, fmt.Errorf("NotificationContentId is not set")
	}
	nc := NewNotificationContent()
	if err := nc.ById(a.NotificationContentId); err != nil {
		return nil, err
	}

	return nc, nil
}

func (a *NotificationActivity) ById(id int64) error {
	return bongo.B.ById(a, id)
}
<commit_msg>Social: Fetch -> ById<commit_after>package models

import (
	"fmt"
	"time"

	"github.com/koding/bongo"
)

// NotificationActivity stores each user NotificationActivity related to notification content.
// When a user makes duplicate NotificationActivity for the same content
// old one is set as obsolete and new one is added to NotificationActivity table
type NotificationActivity struct {
	// unique identifier of NotificationActivity
	Id int64 `json:"id"`

	// notification content foreign key
	NotificationContentId int64 `json:"notificationContentId" sql:"NOT NULL"`

	// notifier account foreign key
	ActorId int64 `json:"actorId,string" sql:"NOT NULL"`

	// activity creation time
	CreatedAt time.Time `json:"createdAt" sql:"NOT NULL"`

	// activity obsolete information
	Obsolete bool `json:"obsolete" sql:"NOT NULL"`
}

func (a *NotificationActivity) BeforeCreate() {
	a.CreatedAt = time.Now()
}

func (a *NotificationActivity) BeforeUpdate() {
	a.Obsolete = true
}

func (a *NotificationActivity) GetId() int64 {
	return a.Id
}

func NewNotificationActivity() *NotificationActivity {
	return &NotificationActivity{}
}

func (a NotificationActivity) TableName() string {
	return "notification.notification_activity"
}

// Create method creates a new activity with obsolete field set as false
// If there already exists one activity with same ActorId and
// NotificationContentId pair, old one is set as obsolete, and
// new one is created
func (a *NotificationActivity) Create() error {
	s := map[string]interface{}{
		"notification_content_id": a.NotificationContentId,
		"actor_id":                a.ActorId,
		"obsolete":                false,
	}

	q := bongo.NewQS(s)
	found := true
	if err := a.One(q); err != nil {
		if err != bongo.RecordNotFound {
			return err
		}
		found = false
	}

	if found {
		if err := bongo.B.Update(a); err != nil {
			return err
		}
		a.Id = 0
		a.Obsolete = false
	}

	return bongo.B.Create(a)
}

func (a *NotificationActivity) FetchByContentIds(ids []int64) ([]NotificationActivity, error) {
	activities := make([]NotificationActivity, 0)
	err := bongo.B.DB.Table(a.TableName()).
		Where("notification_content_id IN (?)", ids).
		Order("id asc").
		Find(&activities).Error

	if err != nil {
		return nil, err
	}

	return activities, nil
}

func (a *NotificationActivity) FetchMapByContentIds(ids []int64) (map[int64][]NotificationActivity, error) {
	if len(ids) == 0 {
		return make(map[int64][]NotificationActivity), nil
	}
	aList, err := a.FetchByContentIds(ids)
	if err != nil {
		return nil, err
	}

	aMap := make(map[int64][]NotificationActivity)
	for _, activity := range aList {
		aMap[activity.NotificationContentId] = append(aMap[activity.NotificationContentId], activity)
	}

	return aMap, nil
}

func (a *NotificationActivity) One(q *bongo.Query) error {
	return bongo.B.One(a, a, q)
}

func (a *NotificationActivity) Some(data interface{}, q *bongo.Query) error {

	return bongo.B.Some(a, data, q)
}

func (a *NotificationActivity) LastActivity() error {
	s := map[string]interface{}{
		"notification_content_id": a.NotificationContentId,
		"obsolete":                false,
	}

	q := bongo.NewQS(s)
	q.Sort = map[string]string{
		"id": "DESC",
	}

	return a.One(q)
}

func (a *NotificationActivity) FetchContent() (*NotificationContent, error) {
	if a.NotificationContentId == 0 {
		return nil, fmt.Errorf("NotificationContentId is not set")
	}
	nc := NewNotificationContent()
	if err := nc.ById(a.NotificationContentId); err != nil {
		return nil, err
	}

	return nc, nil
}

func (a *NotificationActivity) ById(id int64) error {
	return bongo.B.ById(a, id)
}
<|endoftext|>


<commit_before>package extensions

import (
	"fmt"
	"github.com/pawelszydlo/papa-bot"
	"github.com/pawelszydlo/papa-bot/events"
	"github.com/pawelszydlo/papa-bot/utils"
	"math"
	"strconv"
	"strings"
	"text/template"
	"time"
)

// ExtensionCounters - enables the creation of custom counters.
type ExtensionCounters struct {
	counters map[int]*extensionCountersCounter
	bot      *papaBot.Bot
}

type extensionCountersCounter struct {
	transport string
	channel   string
	creator   string
	text      string
	textTmp   *template.Template
	interval  time.Duration
	date      time.Time
	nextTick  time.Time
}

// message will produce an announcement message for the counter.
func (cs *extensionCountersCounter) message(ext *ExtensionCounters) string {
	diff := time.Since(cs.date)
	days := int(math.Abs(diff.Hours())) / 24
	hours := int(math.Abs(diff.Hours())) - days*24
	minutes := int(math.Abs(diff.Minutes())) - hours*60 - days*1440
	vars := map[string]string{
		"days":    fmt.Sprintf("%d", days),
		"hours":   fmt.Sprintf("%d", hours),
		"minutes": fmt.Sprintf("%d", minutes),
		"since":   ext.bot.Humanizer.TimeDiffNow(cs.date, false),
	}
	return utils.Format(cs.textTmp, vars)
}

// Init initializes the extension.
func (ext *ExtensionCounters) Init(bot *papaBot.Bot) error {
	ext.bot = bot
	// Create database table to hold the counters.
	query := `
		-- Main URLs table.
		CREATE TABLE IF NOT EXISTS "counters" (
			"id" INTEGER PRIMARY KEY  AUTOINCREMENT  NOT NULL,
			"transport" VARCHAR NOT NULL,
			"channel" VARCHAR NOT NULL,
			"creator" VARCHAR NOT NULL,
			"announce_text" VARCHAR NOT NULL,
			"interval" INTEGER NOT NULL,
			"target_date" VARCHAR NOT NULL,
			"created" DATETIME DEFAULT (datetime('now','localtime')),
			FOREIGN KEY(creator) REFERENCES users(nick)
		);`
	if _, err := bot.Db.Exec(query); err != nil {
		bot.Log.Panic(err)
	}

	// Add commands for handling the counters.
	bot.RegisterCommand(&papaBot.BotCommand{
		[]string{"c", "counter"},
		true, false, true,
		"help / list / announce <id> / del <id> / add <date> <time> <interval> <channel> <text>",
		"Controls custom counters.",
		ext.commandCounters})

	// Load counters from the db.
	ext.loadCounters()

	// Attach to events.
	bot.EventDispatcher.RegisterListener(events.EventTick, ext.TickListener)

	return nil
}

// TickListener will announce all the counters if needed.
func (ext *ExtensionCounters) TickListener(message events.EventMessage) {
	// Check if it's time to announce the counter.
	for id, c := range ext.counters {
		if time.Since(c.nextTick) > 0 {
			sourceEvent := &events.EventMessage{
				c.transport,
				events.FormatPlain,
				events.EventChannelOps,
				ext.bot.Config.Name,
				"",
				c.channel,
				"",
				message.Context,
				false,
			}
			ext.bot.SendNotice(sourceEvent, c.message(ext))
			c.nextTick = c.nextTick.Add(c.interval * time.Hour)
			ext.bot.Log.Debugf("Counter %d, next tick: %s", id, c.nextTick)
		}
	}
}

// loadCounters will load the counters from the database.
func (ext *ExtensionCounters) loadCounters() {
	ext.counters = map[int]*extensionCountersCounter{}

	result, err := ext.bot.Db.Query(
		`SELECT id, channel, transport, creator, announce_text, interval, target_date FROM counters`)
	if err != nil {
		ext.bot.Log.Warningf("Error while loading counters: %s", err)
		return
	}
	defer result.Close()

	// Get vars.
	for result.Next() {
		var c extensionCountersCounter
		var dateStr string
		var id int
		var interval int
		if err = result.Scan(&id, &c.channel, &c.transport, &c.creator, &c.text, &interval, &dateStr); err != nil {
			ext.bot.Log.Warningf("Can't load counter: %s", err)
			continue
		}
		c.interval = time.Duration(interval)
		// Parse the text template.
		c.textTmp, err = template.New(fmt.Sprintf("counter_%d", id)).Parse(c.text)
		if err != nil {
			ext.bot.Log.Warningf("Can't parse counter template '%s': %s", c.text, err)
		}
		// Handle the date.
		c.date, err = time.Parse("2006-01-02 15:04:05", dateStr)
		if err != nil {
			ext.bot.Log.Fatalf("Can't parse counter date %s: %s", dateStr, err)
		}
		c.date = utils.MustForceLocalTimezone(c.date)
		// Calculate next tick. Start from next daily tick and move backwards.
		nextTick := ext.bot.NextDailyTick()
		for {
			c.nextTick = nextTick
			nextTick = nextTick.Add(-time.Duration(c.interval) * time.Hour)
			if time.Since(nextTick) > 0 { // We moved too far back.
				break
			}
		}
		ext.bot.Log.Debugf("Counter %d, next tick: %s", id, c.nextTick)

		ext.counters[id] = &c
	}
}

// commandCounters is a command for handling the counters.
func (ext *ExtensionCounters) commandCounters(bot *papaBot.Bot, sourceEvent *events.EventMessage, params []string) {

	if len(params) < 1 {
		return
	}
	command := params[0]

	// List.
	if command == "list" {
		if len(ext.counters) > 0 {
			bot.SendMessage(sourceEvent, "Counters:")
			for id, c := range ext.counters {
				bot.SendMessage(sourceEvent, fmt.Sprintf(
					"%d: %s | %s | interval %dh | %s", id, c.channel, c.date, c.interval, c.text))
			}
		} else {
			bot.SendMessage(sourceEvent, "No counters yet.")
		}
		return
	}

	if command == "help" {
		bot.SendMessage(sourceEvent, "To add a new counter:")
		bot.SendMessage(sourceEvent, "add <date> <time> <interval> <channel> <text>")
		bot.SendMessage(
			sourceEvent, `Where: date in format 'YYYY-MM-DD', time in format 'HH:MM:SS', interval is annouce`+
				` interval in hours, channel is the name of the channel to announce on, text is the announcement text.`)
		bot.SendMessage(
			sourceEvent,
			"Announcement text may contain placeholders: {{ .days }}, {{ .hours }}, {{ .minutes }}, {{ .since }}")
		return
	}

	// Force announce.
	if len(params) == 2 && command == "announce" {
		id, err := strconv.Atoi(params[1])
		if err != nil || ext.counters[id] == nil {
			bot.SendMessage(sourceEvent, "Wrong id.")
			return
		}
		bot.SendMessage(sourceEvent,
			fmt.Sprintf("Announcing counter %d to %s...", id, ext.counters[id].channel))
		fakeEvent := &events.EventMessage{
			ext.counters[id].transport,
			events.FormatPlain,
			events.EventChannelOps,
			ext.bot.Config.Name,
			"",
			ext.counters[id].channel,
			"",
			sourceEvent.Context,
			false,
		}
		bot.SendMessage(fakeEvent, ext.counters[id].message(ext))
	}

	// Delete.
	if len(params) == 2 && command == "del" {
		id := params[1]
		bot.SendMessage(sourceEvent, fmt.Sprintf("Deleting counter number %s...", id))
		query := ""
		// Bot owner can delete all counters.
		if bot.UserIsOwner(sourceEvent.UserId) {
			query = `DELETE FROM counters WHERE id=?;`
		} else {
			// User must be an admin, he can delete only his own counters.
			nick := bot.GetAuthenticatedNick(sourceEvent.UserId)
			query = fmt.Sprintf(`DELETE FROM counters WHERE id=? AND creator="%s";`, nick)
		}
		if _, err := bot.Db.Exec(query, id); err != nil {
			bot.Log.Warningf("Error while deleting a counter: %s", err)
			bot.SendMessage(sourceEvent, fmt.Sprintf("Error: %s", err))
			return
		}
		// Reload  counters.
		ext.loadCounters()
		return
	}

	// Add.
	if len(params) > 5 && command == "add" {
		// Sanity check parameters.
		if _, err := time.Parse("2006-01-0215:04:05", params[1]+params[2]); err != nil {
			bot.SendMessage(sourceEvent, "Date and time must be in format: 2015-12-31 12:54:00")
			return
		}
		dateStr := params[1] + " " + params[2]
		interval, err := strconv.ParseInt(params[3], 10, 32)
		if err != nil {
			bot.SendMessage(sourceEvent, "interval parameter must be a number.")
			return
		}
		channel := params[4]

		text := strings.Join(params[5:], " ")
		nick := bot.GetAuthenticatedNick(sourceEvent.UserId)
		// Add counter to database.
		// TODO: what about the transport?
		query := `
			INSERT INTO counters (channel, creator, announce_text, interval, target_date)
			VALUES (?, ?, ?, ?, ?);
			`
		if _, err := bot.Db.Exec(query, channel, nick, text, interval, dateStr); err != nil {
			bot.Log.Warningf("Error while adding a counter: %s", err)
			bot.SendMessage(sourceEvent, fmt.Sprintf("Error: %s", err))
			return
		}
		bot.SendMessage(sourceEvent, "Counter created.")
		// Reload  counters.
		ext.loadCounters()
		return
	}
}
<commit_msg>Fix counter extension addition.<commit_after>package extensions

import (
	"fmt"
	"github.com/pawelszydlo/papa-bot"
	"github.com/pawelszydlo/papa-bot/events"
	"github.com/pawelszydlo/papa-bot/utils"
	"math"
	"strconv"
	"strings"
	"text/template"
	"time"
)

// ExtensionCounters - enables the creation of custom counters.
type ExtensionCounters struct {
	counters map[int]*extensionCountersCounter
	bot      *papaBot.Bot
}

type extensionCountersCounter struct {
	transport string
	channel   string
	creator   string
	text      string
	textTmp   *template.Template
	interval  time.Duration
	date      time.Time
	nextTick  time.Time
}

// message will produce an announcement message for the counter.
func (cs *extensionCountersCounter) message(ext *ExtensionCounters) string {
	diff := time.Since(cs.date)
	days := int(math.Abs(diff.Hours())) / 24
	hours := int(math.Abs(diff.Hours())) - days*24
	minutes := int(math.Abs(diff.Minutes())) - hours*60 - days*1440
	vars := map[string]string{
		"days":    fmt.Sprintf("%d", days),
		"hours":   fmt.Sprintf("%d", hours),
		"minutes": fmt.Sprintf("%d", minutes),
		"since":   ext.bot.Humanizer.TimeDiffNow(cs.date, false),
	}
	return utils.Format(cs.textTmp, vars)
}

// Init initializes the extension.
func (ext *ExtensionCounters) Init(bot *papaBot.Bot) error {
	ext.bot = bot
	// Create database table to hold the counters.
	query := `
		-- Main URLs table.
		CREATE TABLE IF NOT EXISTS "counters" (
			"id" INTEGER PRIMARY KEY  AUTOINCREMENT  NOT NULL,
			"transport" VARCHAR NOT NULL,
			"channel" VARCHAR NOT NULL,
			"creator" VARCHAR NOT NULL,
			"announce_text" VARCHAR NOT NULL,
			"interval" INTEGER NOT NULL,
			"target_date" VARCHAR NOT NULL,
			"created" DATETIME DEFAULT (datetime('now','localtime')),
			FOREIGN KEY(creator) REFERENCES users(nick)
		);`
	if _, err := bot.Db.Exec(query); err != nil {
		bot.Log.Panic(err)
	}

	// Add commands for handling the counters.
	bot.RegisterCommand(&papaBot.BotCommand{
		[]string{"c", "counter"},
		true, false, true,
		"help / list / announce <id> / del <id> / add <date> <time> <interval> <channel> <text>",
		"Controls custom counters.",
		ext.commandCounters})

	// Load counters from the db.
	ext.loadCounters()

	// Attach to events.
	bot.EventDispatcher.RegisterListener(events.EventTick, ext.TickListener)

	return nil
}

// TickListener will announce all the counters if needed.
func (ext *ExtensionCounters) TickListener(message events.EventMessage) {
	// Check if it's time to announce the counter.
	for id, c := range ext.counters {
		if time.Since(c.nextTick) > 0 {
			sourceEvent := &events.EventMessage{
				c.transport,
				events.FormatPlain,
				events.EventChannelOps,
				ext.bot.Config.Name,
				"",
				c.channel,
				"",
				message.Context,
				false,
			}
			ext.bot.SendNotice(sourceEvent, c.message(ext))
			c.nextTick = c.nextTick.Add(c.interval * time.Hour)
			ext.bot.Log.Debugf("Counter %d, next tick: %s", id, c.nextTick)
		}
	}
}

// loadCounters will load the counters from the database.
func (ext *ExtensionCounters) loadCounters() {
	ext.counters = map[int]*extensionCountersCounter{}

	result, err := ext.bot.Db.Query(
		`SELECT id, channel, transport, creator, announce_text, interval, target_date FROM counters`)
	if err != nil {
		ext.bot.Log.Warningf("Error while loading counters: %s", err)
		return
	}
	defer result.Close()

	// Get vars.
	for result.Next() {
		var c extensionCountersCounter
		var dateStr string
		var id int
		var interval int
		if err = result.Scan(&id, &c.channel, &c.transport, &c.creator, &c.text, &interval, &dateStr); err != nil {
			ext.bot.Log.Warningf("Can't load counter: %s", err)
			continue
		}
		c.interval = time.Duration(interval)
		// Parse the text template.
		c.textTmp, err = template.New(fmt.Sprintf("counter_%d", id)).Parse(c.text)
		if err != nil {
			ext.bot.Log.Warningf("Can't parse counter template '%s': %s", c.text, err)
		}
		// Handle the date.
		c.date, err = time.Parse("2006-01-02 15:04:05", dateStr)
		if err != nil {
			ext.bot.Log.Fatalf("Can't parse counter date %s: %s", dateStr, err)
		}
		c.date = utils.MustForceLocalTimezone(c.date)
		// Calculate next tick. Start from next daily tick and move backwards.
		nextTick := ext.bot.NextDailyTick()
		for {
			c.nextTick = nextTick
			nextTick = nextTick.Add(-time.Duration(c.interval) * time.Hour)
			if time.Since(nextTick) > 0 { // We moved too far back.
				break
			}
		}
		ext.bot.Log.Debugf("Counter %d, next tick: %s", id, c.nextTick)

		ext.counters[id] = &c
	}
}

// commandCounters is a command for handling the counters.
func (ext *ExtensionCounters) commandCounters(bot *papaBot.Bot, sourceEvent *events.EventMessage, params []string) {

	if len(params) < 1 {
		return
	}
	command := params[0]

	// List.
	if command == "list" {
		if len(ext.counters) > 0 {
			bot.SendMessage(sourceEvent, "Counters:")
			for id, c := range ext.counters {
				bot.SendMessage(sourceEvent, fmt.Sprintf(
					"%d: %s (%s) | %s | interval %dh | %s", id, c.channel, c.transport, c.date, c.interval, c.text))
			}
		} else {
			bot.SendMessage(sourceEvent, "No counters yet.")
		}
		return
	}

	if command == "help" {
		bot.SendMessage(sourceEvent, "To add a new counter:")
		bot.SendMessage(sourceEvent, "add <date> <time> <interval> <channel> <text>")
		bot.SendMessage(
			sourceEvent, `Where: date in format 'YYYY-MM-DD', time in format 'HH:MM:SS', interval is annouce`+
				` interval in hours, channel is the name of the channel to announce on (on this transport),`+
				`text is the announcement text.`)
		bot.SendMessage(
			sourceEvent,
			"Announcement text may contain placeholders: {{ .days }}, {{ .hours }}, {{ .minutes }}, {{ .since }}")
		return
	}

	// Force announce.
	if len(params) == 2 && command == "announce" {
		id, err := strconv.Atoi(params[1])
		if err != nil || ext.counters[id] == nil {
			bot.SendMessage(sourceEvent, "Wrong id.")
			return
		}
		bot.SendMessage(sourceEvent,
			fmt.Sprintf("Announcing counter %d to %s...", id, ext.counters[id].channel))
		fakeEvent := &events.EventMessage{
			ext.counters[id].transport,
			events.FormatPlain,
			events.EventChannelOps,
			ext.bot.Config.Name,
			"",
			ext.counters[id].channel,
			"",
			sourceEvent.Context,
			false,
		}
		bot.SendMessage(fakeEvent, ext.counters[id].message(ext))
	}

	// Delete.
	if len(params) == 2 && command == "del" {
		id := params[1]
		bot.SendMessage(sourceEvent, fmt.Sprintf("Deleting counter number %s...", id))
		query := ""
		// Bot owner can delete all counters.
		if bot.UserIsOwner(sourceEvent.UserId) {
			query = `DELETE FROM counters WHERE id=?;`
		} else {
			// User must be an admin, he can delete only his own counters.
			nick := bot.GetAuthenticatedNick(sourceEvent.UserId)
			query = fmt.Sprintf(`DELETE FROM counters WHERE id=? AND creator="%s";`, nick)
		}
		if _, err := bot.Db.Exec(query, id); err != nil {
			bot.Log.Warningf("Error while deleting a counter: %s", err)
			bot.SendMessage(sourceEvent, fmt.Sprintf("Error: %s", err))
			return
		}
		// Reload  counters.
		ext.loadCounters()
		return
	}

	// Add.
	if len(params) > 5 && command == "add" {
		// Sanity check parameters.
		if _, err := time.Parse("2006-01-0215:04:05", params[1]+params[2]); err != nil {
			bot.SendMessage(sourceEvent, "Date and time must be in format: 2015-12-31 12:54:00")
			return
		}
		dateStr := params[1] + " " + params[2]
		interval, err := strconv.ParseInt(params[3], 10, 32)
		if err != nil {
			bot.SendMessage(sourceEvent, "Interval parameter must be a number of hours.")
			return
		}
		channel := params[4]

		text := strings.Join(params[5:], " ")
		nick := bot.GetAuthenticatedNick(sourceEvent.UserId)
		// Add counter to database.
		query := `
			INSERT INTO counters (channel, transport, creator, announce_text, interval, target_date)
			VALUES (?, ?, ?, ?, ?, ?);
			`
		if _, err := bot.Db.Exec(query, channel, sourceEvent.TransportName, nick, text, interval, dateStr); err != nil {
			bot.Log.Warningf("Error while adding a counter: %s", err)
			bot.SendMessage(sourceEvent, fmt.Sprintf("Error: %s", err))
			return
		}
		bot.SendMessage(sourceEvent, "Counter created.")
		// Reload  counters.
		ext.loadCounters()
		return
	}
}
<|endoftext|>


<commit_before>// Copyright e-Xpert Solutions SA. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sys

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/e-XpertSolutions/f5-rest-client/f5"
)

// FileSSLCRLConfigList holds a list of FileSSLCRL configuration.
type FileSSLCRLConfigList struct {
	Items    []FileSSLCRLConfig `json:"items"`
	Kind     string             `json:"kind"`
	SelfLink string             `json:"selflink"`
}

// FileSSLCRLConfig holds the configuration of a single FileSSLCRL.
type FileSSLCRLConfig struct {
	Name           string    `json:"name,omitempty"`
	Partition      string    `json:"partition,omitempty"`
	FullPath       string    `json:"fullPath,omitempty"`
	Generator      int64     `json:"generator,omitempty"`
	Checksum       string    `json:"checksum,omitempty"`
	CreateTime     time.Time `json:"createTime,omitempty"`
	CreatedBy      string    `json:"createdBy,omitempty"`
	LastUpdateTime time.Time `json:"lastUpdateTime,omitempty"`
	Mode           int64     `json:"mode,omitempty"`
	Revision       int64     `json:"revision,omitempty"`
	Size           int64     `json:"size,omitempty"`
	UpdatedBy      string    `json:"updatedBy,omitempty"`
}

// FileSSLCRLEndpoint represents the REST resource for managing FileSSLCRL.
const FileSSLCRLEndpoint = "/file/ssl-crl"

// FileSSLCRLResource provides an API to manage FileSSLCRL configurations.
type FileSSLCRLResource struct {
	c *f5.Client
}

// ListAll  lists all the FileSSLCRL configurations.
func (r *FileSSLCRLResource) ListAll() (*FileSSLCRLConfigList, error) {
	var list FileSSLCRLConfigList
	if err := r.c.ReadQuery(BasePath+FileSSLCRLEndpoint, &list); err != nil {
		return nil, err
	}
	return &list, nil
}

// Get a single FileSSLCRL configuration identified by id.
func (r *FileSSLCRLResource) Get(id string) (*FileSSLCRLConfig, error) {
	var item FileSSLCRLConfig
	if err := r.c.ReadQuery(BasePath+FileSSLCRLEndpoint+"/"+id, &item); err != nil {
		return nil, err
	}
	return &item, nil
}

// CreateFromFile uploads a CRL file in PEM and create or update its value.
func (r *FileSSLCRLResource) CreateFromFile(name string, crlPEMFile io.Reader, filesize int64) error {
	uploadResp, err := r.c.UploadFile(crlPEMFile, name+".crl", filesize)
	if err != nil {
		return fmt.Errorf("failed to create upload request: %v", err)
	}
	data := map[string]string{
		"name":        name + ".crl",
		"source-path": "file:" + uploadResp.LocalFilePath,
	}
	if err := r.c.ModQuery("PUT", BasePath+FileSSLCRLEndpoint+"/"+name+".crl", data); err != nil {
		return fmt.Errorf("failed to import crl file: %v", err)
	}
	return nil
}

// Edit a FileSSLCRL configuration identified by id.
func (r *FileSSLCRLResource) Edit(id, path string) error {
	info, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("failed to gather information about '%s': %v", path, err)
	}
	f, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("failed to read file from path: %v", err)
	}
	defer f.Close()

	uploadResp, err := r.c.UploadFile(f, filepath.Base(path), info.Size())
	if err != nil {
		return fmt.Errorf("failed to upload file %q: %v", path, err)
	}

	data := map[string]string{
		"source-path": "file:" + uploadResp.LocalFilePath,
	}
	if err := r.c.ModQuery("PUT", BasePath+FileSSLCRLEndpoint+"/"+id, data); err != nil {
		return fmt.Errorf("failed to create FileSSLCRL configuration: %v", err)
	}

	return nil
}

// Delete a single FileSSLCRL configuration identified by id.
func (r *FileSSLCRLResource) Delete(id string) error {
	if err := r.c.ModQuery("DELETE", BasePath+FileSSLCRLEndpoint+"/"+id, nil); err != nil {
		return err
	}
	return nil
}
<commit_msg>f5/sys: create a method for updating a CRL file<commit_after>// Copyright e-Xpert Solutions SA. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package sys

import (
	"fmt"
	"io"
	"os"
	"path/filepath"
	"time"

	"github.com/e-XpertSolutions/f5-rest-client/f5"
)

// FileSSLCRLConfigList holds a list of FileSSLCRL configuration.
type FileSSLCRLConfigList struct {
	Items    []FileSSLCRLConfig `json:"items"`
	Kind     string             `json:"kind"`
	SelfLink string             `json:"selflink"`
}

// FileSSLCRLConfig holds the configuration of a single FileSSLCRL.
type FileSSLCRLConfig struct {
	Name           string    `json:"name,omitempty"`
	Partition      string    `json:"partition,omitempty"`
	FullPath       string    `json:"fullPath,omitempty"`
	Generator      int64     `json:"generator,omitempty"`
	Checksum       string    `json:"checksum,omitempty"`
	CreateTime     time.Time `json:"createTime,omitempty"`
	CreatedBy      string    `json:"createdBy,omitempty"`
	LastUpdateTime time.Time `json:"lastUpdateTime,omitempty"`
	Mode           int64     `json:"mode,omitempty"`
	Revision       int64     `json:"revision,omitempty"`
	Size           int64     `json:"size,omitempty"`
	UpdatedBy      string    `json:"updatedBy,omitempty"`
}

// FileSSLCRLEndpoint represents the REST resource for managing FileSSLCRL.
const FileSSLCRLEndpoint = "/file/ssl-crl"

// FileSSLCRLResource provides an API to manage FileSSLCRL configurations.
type FileSSLCRLResource struct {
	c *f5.Client
}

// ListAll  lists all the FileSSLCRL configurations.
func (r *FileSSLCRLResource) ListAll() (*FileSSLCRLConfigList, error) {
	var list FileSSLCRLConfigList
	if err := r.c.ReadQuery(BasePath+FileSSLCRLEndpoint, &list); err != nil {
		return nil, err
	}
	return &list, nil
}

// Get a single FileSSLCRL configuration identified by id.
func (r *FileSSLCRLResource) Get(id string) (*FileSSLCRLConfig, error) {
	var item FileSSLCRLConfig
	if err := r.c.ReadQuery(BasePath+FileSSLCRLEndpoint+"/"+id, &item); err != nil {
		return nil, err
	}
	return &item, nil
}

// CreateFromFile uploads a CRL file in PEM and create or update its value.
func (r *FileSSLCRLResource) CreateFromFile(name string, crlPEMFile io.Reader, filesize int64) error {
	uploadResp, err := r.c.UploadFile(crlPEMFile, name+".crl", filesize)
	if err != nil {
		return fmt.Errorf("failed to create upload request: %v", err)
	}
	data := map[string]string{
		"name":        name + ".crl",
		"source-path": "file:" + uploadResp.LocalFilePath,
	}
	if err := r.c.ModQuery("POST", BasePath+FileSSLCRLEndpoint, data); err != nil {
		return fmt.Errorf("failed to import crl file: %v", err)
	}
	return nil
}

func (r *FileSSLCRLResource) EditFromFile(name string, crlPEMFile io.Reader, filesize int64) error {
	uploadResp, err := r.c.UploadFile(crlPEMFile, name+".crl", filesize)
	if err != nil {
		return fmt.Errorf("failed to create upload request: %v", err)
	}
	data := map[string]string{
		"name":        name + ".crl",
		"source-path": "file:" + uploadResp.LocalFilePath,
	}
	if err := r.c.ModQuery("PUT", BasePath+FileSSLCRLEndpoint+"/"+name+".crl", data); err != nil {
		return fmt.Errorf("failed to update imported crl file: %v", err)
	}
	return nil
}

// Edit a FileSSLCRL configuration identified by id.
func (r *FileSSLCRLResource) Edit(id, path string) error {
	info, err := os.Stat(path)
	if err != nil {
		return fmt.Errorf("failed to gather information about '%s': %v", path, err)
	}
	f, err := os.Open(path)
	if err != nil {
		return fmt.Errorf("failed to read file from path: %v", err)
	}
	defer f.Close()

	uploadResp, err := r.c.UploadFile(f, filepath.Base(path), info.Size())
	if err != nil {
		return fmt.Errorf("failed to upload file %q: %v", path, err)
	}

	data := map[string]string{
		"source-path": "file:" + uploadResp.LocalFilePath,
	}
	if err := r.c.ModQuery("PUT", BasePath+FileSSLCRLEndpoint+"/"+id, data); err != nil {
		return fmt.Errorf("failed to create FileSSLCRL configuration: %v", err)
	}

	return nil
}

// Delete a single FileSSLCRL configuration identified by id.
func (r *FileSSLCRLResource) Delete(id string) error {
	if err := r.c.ModQuery("DELETE", BasePath+FileSSLCRLEndpoint+"/"+id, nil); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2015 Google Inc. 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.

// More information about Google Distance Matrix API is available on
// https://developers.google.com/maps/documentation/distancematrix/

package maps

import (
	"errors"
	"net/url"
	"strings"
	"time"

	"golang.org/x/net/context"
)

var distanceMatrixAPI = &apiConfig{
	host:            "https://maps.googleapis.com",
	path:            "/maps/api/distancematrix/json",
	acceptsClientID: true,
}

// DistanceMatrix makes a Distance Matrix API request
func (c *Client) DistanceMatrix(ctx context.Context, r *DistanceMatrixRequest) (*DistanceMatrixResponse, error) {

	if len(r.Origins) == 0 {
		return nil, errors.New("maps: origins empty")
	}
	if len(r.Destinations) == 0 {
		return nil, errors.New("maps: destinations empty")
	}
	if r.DepartureTime != "" && r.ArrivalTime != "" {
		return nil, errors.New("maps: DepartureTime and ArrivalTime both specified")
	}
	if len(r.TransitMode) != 0 && r.Mode != TravelModeTransit {
		return nil, errors.New("maps: TransitMode specified while Mode != TravelModeTransit")
	}
	if r.TransitRoutingPreference != "" && r.Mode != TravelModeTransit {
		return nil, errors.New("maps: mode of transit '" + string(r.Mode) + "' invalid for TransitRoutingPreference")
	}
	if r.Mode == TravelModeTransit && r.TrafficModel != "" {
		return nil, errors.New("maps: cannot specify transit mode and traffic model together")
	}

	var response struct {
		commonResponse
		DistanceMatrixResponse
	}

	if err := c.getJSON(ctx, distanceMatrixAPI, r, &response); err != nil {
		return nil, err
	}

	if err := response.StatusError(); err != nil {
		return nil, err
	}

	return &response.DistanceMatrixResponse, nil
}

func (r *DistanceMatrixRequest) params() url.Values {
	q := make(url.Values)
	q.Set("origins", strings.Join(r.Origins, "|"))
	q.Set("destinations", strings.Join(r.Destinations, "|"))
	if r.Mode != "" {
		q.Set("mode", string(r.Mode))
	}
	if r.Language != "" {
		q.Set("language", r.Language)
	}
	if r.Avoid != "" {
		q.Set("avoid", string(r.Avoid))
	}
	if r.Units != "" {
		q.Set("units", string(r.Units))
	}
	if r.DepartureTime != "" {
		q.Set("departure_time", r.DepartureTime)
	}
	if r.ArrivalTime != "" {
		q.Set("arrival_time", r.ArrivalTime)
	}
	if r.TrafficModel != "" {
		q.Set("traffic_model", string(r.TrafficModel))
	}
	if len(r.TransitMode) != 0 {
		var transitMode []string
		for _, t := range r.TransitMode {
			transitMode = append(transitMode, string(t))
		}
		q.Set("transit_mode", strings.Join(transitMode, "|"))
	}
	if r.TransitRoutingPreference != "" {
		q.Set("transit_routing_preference", string(r.TransitRoutingPreference))
	}
	return q
}

// DistanceMatrixRequest is the request struct for Distance Matrix APi
type DistanceMatrixRequest struct {
	// Origins is a list of addresses and/or textual latitude/longitude values from which to calculate distance and time. Required.
	Origins []string
	// Destinations is a list of addresses and/or textual latitude/longitude values to which to calculate distance and time. Required.
	Destinations []string
	// Mode specifies the mode of transport to use when calculating distance. Valid values are `ModeDriving`, `ModeWalking`, `ModeBicycling`
	// and `ModeTransit`. Optional.
	Mode Mode
	// Language in which to return results. Optional.
	Language string
	// Avoid introduces restrictions to the route. Valid values are `AvoidTolls`, `AvoidHighways` and `AvoidFerries`. Optional.
	Avoid Avoid
	// Units Specifies the unit system to use when expressing distance as text. Valid values are `UnitsMetric` and `UnitsImperial`. Optional.
	Units Units
	// DepartureTime is the desired time of departure. You can specify the time as an integer in seconds since midnight, January 1, 1970 UTC.
	// Alternatively, you can specify a value of `"now"``. Optional.
	DepartureTime string
	// ArrivalTime specifies the desired time of arrival for transit requests, in seconds since midnight, January 1, 1970 UTC. You cannot
	// specify both `DepartureTime` and `ArrivalTime`. Optional.
	ArrivalTime string
	// TrafficModel determines the type of model that will be used when determining travel time when using depature times in the future
	// options are TrafficModelBestGuess, TrafficModelOptimistic or TrafficModelPessimistic. Optional. Default is TrafficModelBestGuess
	TrafficModel TrafficModel
	// TransitMode specifies one or more preferred modes of transit. This parameter may only be specified for requests where the mode is
	// `transit`. Valid values are `TransitModeBus`, `TransitModeSubway`, `TransitModeTrain`, `TransitModeTram`, and `TransitModeRail`.
	// Optional.
	TransitMode []TransitMode
	// TransitRoutingPreference Specifies preferences for transit requests. Valid values are `TransitRoutingPreferenceLessWalking` and
	// `TransitRoutingPreferenceFewerTransfers`. Optional.
	TransitRoutingPreference TransitRoutingPreference
}

// DistanceMatrixResponse represents a Distance Matrix API response.
type DistanceMatrixResponse struct {

	// OriginAddresses contains an array of addresses as returned by the API from your original request.
	OriginAddresses []string `json:"origin_addresses"`
	// DestinationAddresses contains an array of addresses as returned by the API from your original request.
	DestinationAddresses []string `json:"destination_addresses"`
	// Rows contains an array of elements.
	Rows []DistanceMatrixElementsRow `json:"rows"`
}

// DistanceMatrixElementsRow is a row of distance elements.
type DistanceMatrixElementsRow struct {
	Elements []*DistanceMatrixElement `json:"elements"`
}

// DistanceMatrixElement is the travel distance and time for a pair of origin and destination.
type DistanceMatrixElement struct {
	Status string `json:"status"`
	// Duration is the length of time it takes to travel this route.
	Duration time.Duration `json:"duration"`
	// DurationInTraffic is the length of time it takes to travel this route considering traffic.
	DurationInTraffic time.Duration `json:"duration_in_traffic"`
	// Distance is the total distance of this route.
	Distance Distance `json:"distance"`
}
<commit_msg>Wrap distance matrix comments to make them better fit godoc.<commit_after>// Copyright 2015 Google Inc. 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.

// More information about Google Distance Matrix API is available on
// https://developers.google.com/maps/documentation/distancematrix/

package maps

import (
	"errors"
	"net/url"
	"strings"
	"time"

	"golang.org/x/net/context"
)

var distanceMatrixAPI = &apiConfig{
	host:            "https://maps.googleapis.com",
	path:            "/maps/api/distancematrix/json",
	acceptsClientID: true,
}

// DistanceMatrix makes a Distance Matrix API request
func (c *Client) DistanceMatrix(ctx context.Context, r *DistanceMatrixRequest) (*DistanceMatrixResponse, error) {

	if len(r.Origins) == 0 {
		return nil, errors.New("maps: origins empty")
	}
	if len(r.Destinations) == 0 {
		return nil, errors.New("maps: destinations empty")
	}
	if r.DepartureTime != "" && r.ArrivalTime != "" {
		return nil, errors.New("maps: DepartureTime and ArrivalTime both specified")
	}
	if len(r.TransitMode) != 0 && r.Mode != TravelModeTransit {
		return nil, errors.New("maps: TransitMode specified while Mode != TravelModeTransit")
	}
	if r.TransitRoutingPreference != "" && r.Mode != TravelModeTransit {
		return nil, errors.New("maps: mode of transit '" + string(r.Mode) + "' invalid for TransitRoutingPreference")
	}
	if r.Mode == TravelModeTransit && r.TrafficModel != "" {
		return nil, errors.New("maps: cannot specify transit mode and traffic model together")
	}

	var response struct {
		commonResponse
		DistanceMatrixResponse
	}

	if err := c.getJSON(ctx, distanceMatrixAPI, r, &response); err != nil {
		return nil, err
	}

	if err := response.StatusError(); err != nil {
		return nil, err
	}

	return &response.DistanceMatrixResponse, nil
}

func (r *DistanceMatrixRequest) params() url.Values {
	q := make(url.Values)
	q.Set("origins", strings.Join(r.Origins, "|"))
	q.Set("destinations", strings.Join(r.Destinations, "|"))
	if r.Mode != "" {
		q.Set("mode", string(r.Mode))
	}
	if r.Language != "" {
		q.Set("language", r.Language)
	}
	if r.Avoid != "" {
		q.Set("avoid", string(r.Avoid))
	}
	if r.Units != "" {
		q.Set("units", string(r.Units))
	}
	if r.DepartureTime != "" {
		q.Set("departure_time", r.DepartureTime)
	}
	if r.ArrivalTime != "" {
		q.Set("arrival_time", r.ArrivalTime)
	}
	if r.TrafficModel != "" {
		q.Set("traffic_model", string(r.TrafficModel))
	}
	if len(r.TransitMode) != 0 {
		var transitMode []string
		for _, t := range r.TransitMode {
			transitMode = append(transitMode, string(t))
		}
		q.Set("transit_mode", strings.Join(transitMode, "|"))
	}
	if r.TransitRoutingPreference != "" {
		q.Set("transit_routing_preference", string(r.TransitRoutingPreference))
	}
	return q
}

// DistanceMatrixRequest is the request struct for Distance Matrix APi
type DistanceMatrixRequest struct {
	// Origins is a list of addresses and/or textual latitude/longitude values
	// from which to calculate distance and time. Required.
	Origins []string
	// Destinations is a list of addresses and/or textual latitude/longitude values
	// to which to calculate distance and time. Required.
	Destinations []string
	// Mode specifies the mode of transport to use when calculating distance.
	// Valid values are `ModeDriving`, `ModeWalking`, `ModeBicycling`
	// and `ModeTransit`. Optional.
	Mode Mode
	// Language in which to return results. Optional.
	Language string
	// Avoid introduces restrictions to the route. Valid values are `AvoidTolls`,
	// `AvoidHighways` and `AvoidFerries`. Optional.
	Avoid Avoid
	// Units Specifies the unit system to use when expressing distance as text.
	// Valid values are `UnitsMetric` and `UnitsImperial`. Optional.
	Units Units
	// DepartureTime is the desired time of departure. You can specify the time as
	// an integer in seconds since midnight, January 1, 1970 UTC. Alternatively,
	// you can specify a value of `"now"``. Optional.
	DepartureTime string
	// ArrivalTime specifies the desired time of arrival for transit requests,
	// in seconds since midnight, January 1, 1970 UTC. You cannot specify
	// both `DepartureTime` and `ArrivalTime`. Optional.
	ArrivalTime string
	// TrafficModel determines the type of model that will be used when determining
	// travel time when using depature times in the future. Options are
	// `TrafficModelBestGuess`, `TrafficModelOptimistic`` or `TrafficModelPessimistic`.
	// Optional. Default is `TrafficModelBestGuess``
	TrafficModel TrafficModel
	// TransitMode specifies one or more preferred modes of transit. This parameter
	// may only be specified for requests where the mode is `transit`. Valid values
	// are `TransitModeBus`, `TransitModeSubway`, `TransitModeTrain`, `TransitModeTram`,
	// and `TransitModeRail`. Optional.
	TransitMode []TransitMode
	// TransitRoutingPreference Specifies preferences for transit requests. Valid
	// values are `TransitRoutingPreferenceLessWalking` and
	// `TransitRoutingPreferenceFewerTransfers`. Optional.
	TransitRoutingPreference TransitRoutingPreference
}

// DistanceMatrixResponse represents a Distance Matrix API response.
type DistanceMatrixResponse struct {

	// OriginAddresses contains an array of addresses as returned by the API from
	// your original request.
	OriginAddresses []string `json:"origin_addresses"`
	// DestinationAddresses contains an array of addresses as returned by the API
	// from your original request.
	DestinationAddresses []string `json:"destination_addresses"`
	// Rows contains an array of elements.
	Rows []DistanceMatrixElementsRow `json:"rows"`
}

// DistanceMatrixElementsRow is a row of distance elements.
type DistanceMatrixElementsRow struct {
	Elements []*DistanceMatrixElement `json:"elements"`
}

// DistanceMatrixElement is the travel distance and time for a pair of origin
// and destination.
type DistanceMatrixElement struct {
	Status string `json:"status"`
	// Duration is the length of time it takes to travel this route.
	Duration time.Duration `json:"duration"`
	// DurationInTraffic is the length of time it takes to travel this route
	// considering traffic.
	DurationInTraffic time.Duration `json:"duration_in_traffic"`
	// Distance is the total distance of this route.
	Distance Distance `json:"distance"`
}
<|endoftext|>


<commit_before>// Copyright (c) 2014, Kevin Walsh.  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 main

import (
	"bufio"
	"crypto/rand"
	"crypto/rsa"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"math/big"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"cloudproxy/tao"
)

var serverHost = flag.String("host", "localhost", "address for client/server")
var serverPort = flag.String("port", "8123", "port for client/server")
var serverAddr string // see main()
var localMode = flag.Bool("local", true, "Run host demo")
var clientMode = flag.Bool("client", true, "Run demo client")
var serverMode = flag.Bool("server", true, "Run demo server")
var pingCount = flag.Int("n", 5, "Number of client/server pings")
var demoAuth = flag.String("auth", "tao", "\"tcp\", \"tls\", or \"tao\"")

// TCP mode client/server

func setupTCPServer() (net.Listener, error) {
	return net.Listen("tcp", serverAddr)
}

func setupTCPClient() (net.Conn, error) {
	return net.Dial("tcp", serverAddr)
}

// TLS mode client/server

const (
	x509duration = 24 * time.Hour
	x509keySize  = 2048
)

func GenerateX509() (*tls.Certificate, error) {
	priv, err := rsa.GenerateKey(rand.Reader, x509keySize)
	if err != nil {
		return nil, err
	}

	notBefore := time.Now()
	notAfter := notBefore.Add(x509duration)

	serialNumberLimit := new(big.Int).Lsh(big.NewInt(1), 128)
	serialNumber, err := rand.Int(rand.Reader, serialNumberLimit)
	if err != nil {
		return nil, err
	}

	template := x509.Certificate{
		SerialNumber: serialNumber,
		Subject: pkix.Name{
			Organization: []string{"Google Tao Demo"},
		},
		NotBefore:             notBefore,
		NotAfter:              notAfter,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
	}

	if ip := net.ParseIP(*serverHost); ip != nil {
		template.IPAddresses = append(template.IPAddresses, ip)
	} else {
		template.DNSNames = append(template.DNSNames, *serverHost)
	}

	// template.IsCA = true
	// template.KeyUsage |= x509.KeyUsageCertSign

	derBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, &priv.PublicKey, priv)
	if err != nil {
		return nil, err
	}

	certPem := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	keyPem := pem.EncodeToMemory(&pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(priv)})

	cert, err := tls.X509KeyPair(certPem, keyPem)
	if err != nil {
		fmt.Printf("can't parse my cert\n")
		return nil, err
	}

	return &cert, nil
}

func setupTLSServer() (net.Listener, error) {
	cert, err := GenerateX509()
	if err != nil {
		fmt.Printf("server: can't create key and cert: %s\n", err.Error())
		return nil, err
	}
	return tls.Listen("tcp", serverAddr, &tls.Config{
		RootCAs:            x509.NewCertPool(),
		Certificates:       []tls.Certificate{*cert},
		InsecureSkipVerify: true,
	})
}

func setupTLSClient() (net.Conn, error) {
	cert, err := GenerateX509()
	if err != nil {
		fmt.Printf("client: can't create key and cert: %s\n", err.Error())
		return nil, err
	}
	return tls.Dial("tcp", serverAddr, &tls.Config{
		RootCAs:            x509.NewCertPool(),
		Certificates:       []tls.Certificate{*cert},
		InsecureSkipVerify: true,
	})
}

// client/server driver

func doRequest() bool {
	fmt.Printf("client: connecting to %s using %s authentication.\n", serverAddr, *demoAuth)
	var conn net.Conn
	var err error
	switch *demoAuth {
	case "tcp":
		conn, err = setupTCPClient()
	case "tls":
		conn, err = setupTLSClient()
		// TODO(kwalsh) Tao-level authentication: use TLS, then exchange names and
		// delegation attestations
		// case "tao":
		// conn, err = setupTaoClient()
	}
	if err != nil {
		fmt.Printf("client: error connecting to %s: %s\n", serverAddr, err.Error())
		return false
	}
	defer conn.Close()

	_, err = fmt.Fprintf(conn, "Hello\n")
	if err != nil {
		fmt.Printf("client: can't write: %s\n", err.Error())
		return false
	}
	msg, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		fmt.Printf("client can't read: %s\n", err.Error())
		return false
	}
	msg = strings.TrimSpace(msg)
	fmt.Printf("client: got reply: %s\n", msg)
	return true
}

func doClient() {
	pingGood := 0
	pingFail := 0
	for i := 0; i < *pingCount || *pingCount < 0; i++ { // negative means forever
		if doRequest() {
			pingGood++
		} else {
			pingFail++
		}
		fmt.Printf("client: made %d connections, finished %d ok, %d bad pings\n",
			i+1, pingGood, pingFail)
	}
}

func doResponse(conn net.Conn, responseOk chan<- bool) {
	defer conn.Close()

	switch *demoAuth {
	case "tcp", "tls":
		// authentication already done by lower layers
	case "tao":
		// TODO(kwalsh) Tao-level authorization: exchange names and delegation
		// attestations.
	}

	msg, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		fmt.Printf("server: can't read: %s\n", err.Error())
		conn.Close()
		responseOk <- false
		return
	}
	msg = strings.TrimSpace(msg)
	fmt.Printf("server: got message: %s\n", msg)
	responseOk <- true
	fmt.Fprintf(conn, "echo(%s)\n", msg)
	conn.Close()
}

func doServer(stop chan bool, ready, done chan<- bool) {
	var sock net.Listener
	var err error
	switch *demoAuth {
	case "tcp":
		sock, err = setupTCPServer()
	case "tls", "tao":
		sock, err = setupTLSServer()
	}
	if err != nil {
		fmt.Printf("server: can't listen at %s: %s\n", serverAddr, err.Error())
		ready <- false
		done <- true
		return
	}
	fmt.Printf("server: listening at %s using %s authentication.\n", serverAddr, *demoAuth)
	ready <- true

	pings := make(chan bool, 10)
	connCount := 0

	go func() {
		for connCount = 0; connCount < *pingCount || *pingCount < 0; connCount++ { // negative means forever
			conn, err := sock.Accept()
			if err != nil {
				fmt.Printf("server: can't accept connection: %s\n", err.Error())
				stop <- true
				return
			}
			go doResponse(conn, pings)
		}
	}()

	pingGood := 0
	pingFail := 0

loop:
	for {
		select {
		case <-stop:
			break loop
		case ok := <-pings:
			if ok {
				pingGood++
			} else {
				pingFail++
			}
		}
	}

	sock.Close()
	fmt.Printf("server: handled %d connections, finished %d ok, %d bad pings\n",
		connCount, pingGood, pingFail)

	done <- true
}

// Tao Host demo

func hostTaoDemo() error {
	name, err := tao.Host().GetTaoName()
	if err != nil {
		return err
	}
	fmt.Printf("My root name is %s\n", name)

	args := make([]string, len(os.Args))
	for index, arg := range os.Args {
		args[index] = strconv.Quote(arg)
	}
	subprin := "Args(" + strings.Join(args, ", ") + ")"
	err = tao.Host().ExtendTaoName(subprin)
	if err != nil {
		return err
	}

	name, err = tao.Host().GetTaoName()
	if err != nil {
		return err
	}
	fmt.Printf("My full name is %s\n", name)

	random, err := tao.Host().GetRandomBytes(10)
	if err != nil {
		return err
	}
	fmt.Printf("Random bytes  : % x\n", random)

	n, err := tao.Host().Rand().Read(random)
	if err != nil {
		return err
	}
	fmt.Printf("%d more bytes : % x\n", n, random)

	secret, err := tao.Host().GetSharedSecret(10, tao.SharedSecretPolicyDefault)
	if err != nil {
		return err
	}
	fmt.Printf("Shared secret : % x\n", secret)

	sealed, err := tao.Host().Seal(random, tao.SealPolicyDefault)
	if err != nil {
		return err
	}
	fmt.Printf("Sealed bytes  : % x\n", sealed)

	unsealed, policy, err := tao.Host().Unseal(sealed)
	if err != nil {
		return err
	}
	if policy != tao.SealPolicyDefault {
		return errors.New("unexpected policy on unseal")
	}
	fmt.Printf("Unsealed bytes: % x\n", unsealed)

	return nil
}

func main() {
	flag.Parse()
	serverAddr = *serverHost + ":" + *serverPort
	switch *demoAuth {
	case "tcp", "tls", "tao":
	default:
		fmt.Printf("unrecognized authentication mode: %s\n", *demoAuth)
		return
	}

	fmt.Printf("Go Tao Demo\n")

	if !tao.HostAvailable() {
		fmt.Printf("can't continue: No host Tao available\n")
		return
	}

	if *localMode {
		err := hostTaoDemo()
		if err != nil {
			fmt.Printf("error: %s\n", err.Error())
			return
		}
	}

	serverStop := make(chan bool, 1)
	serverReady := make(chan bool, 1)
	serverDone := make(chan bool, 1)

	if *serverMode {
		go doServer(serverStop, serverReady, serverDone)
	} else {
		serverReady <- true
		serverDone <- true
	}

	if *clientMode {
		ok := <-serverReady
		if ok {
			doClient()
		}
		serverStop <- true
	}

	<-serverDone
	fmt.Printf("Done\n")
}
<commit_msg>go demo uses go keys<commit_after>// Copyright (c) 2014, Kevin Walsh.  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 main

import (
	"bufio"
	"crypto/tls"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"errors"
	"flag"
	"fmt"
	"net"
	"os"
	"strconv"
	"strings"
	"time"

	"cloudproxy/tao"
)

var serverHost = flag.String("host", "localhost", "address for client/server")
var serverPort = flag.String("port", "8123", "port for client/server")
var serverAddr string // see main()
var localMode = flag.Bool("local", true, "Run host demo")
var clientMode = flag.Bool("client", true, "Run demo client")
var serverMode = flag.Bool("server", true, "Run demo server")
var pingCount = flag.Int("n", 5, "Number of client/server pings")
var demoAuth = flag.String("auth", "tao", "\"tcp\", \"tls\", or \"tao\"")

// TCP mode client/server

func setupTCPServer() (net.Listener, error) {
	return net.Listen("tcp", serverAddr)
}

func setupTCPClient() (net.Conn, error) {
	return net.Dial("tcp", serverAddr)
}

// TLS mode client/server

const (
	x509duration = 24 * time.Hour
	x509keySize  = 2048
)

func GenerateX509() (*tls.Certificate, error) {
	keys, err := tao.NewTemporaryTaoDelegatedKeys(tao.Signing, tao.Host())
	if err != nil {
		return nil, err
	}

	/*
		if ip := net.ParseIP(*serverHost); ip != nil {
			template.IPAddresses = append(template.IPAddresses, ip)
		} else {
			template.DNSNames = append(template.DNSNames, *serverHost)
		}
	*/

	derBytes, err := keys.SigningKey.CreateSelfSignedX509(&pkix.Name{
		Organization: []string{"Google Tao Demo"}})
	if err != nil {
		return nil, err
	}

	certPem := pem.EncodeToMemory(&pem.Block{Type: "CERTIFICATE", Bytes: derBytes})
	keyBytes, err := tao.MarshalSignerDER(keys.SigningKey)
	if err != nil {
		return nil, err
	}
	keyPem := pem.EncodeToMemory(&pem.Block{Type: "ECDSA PRIVATE KEY", Bytes: keyBytes})

	cert, err := tls.X509KeyPair(certPem, keyPem)
	if err != nil {
		fmt.Printf("can't parse my cert\n")
		return nil, err
	}

	return &cert, nil
}

func setupTLSServer() (net.Listener, error) {
	cert, err := GenerateX509()
	if err != nil {
		fmt.Printf("server: can't create key and cert: %s\n", err.Error())
		return nil, err
	}
	return tls.Listen("tcp", serverAddr, &tls.Config{
		RootCAs:            x509.NewCertPool(),
		Certificates:       []tls.Certificate{*cert},
		InsecureSkipVerify: true,
	})
}

func setupTLSClient() (net.Conn, *tls.Certificate, error) {
	cert, err := GenerateX509()
	if err != nil {
		fmt.Printf("client: can't create key and cert: %s\n", err.Error())
		return nil, nil, err
	}
	conn, err := tls.Dial("tcp", serverAddr, &tls.Config{
		RootCAs:            x509.NewCertPool(),
		Certificates:       []tls.Certificate{*cert},
		InsecureSkipVerify: true,
	})
	return conn, cert, err
}

// Tao mode client/server

func setupTaoClient() (net.Conn, error) {
	conn, cert, err := setupTLSClient()
	if err != nil {
		return nil, err
	}
	_ = cert
	return conn, errors.New("not yet implemented")
}

// client/server driver

func doRequest() bool {
	fmt.Printf("client: connecting to %s using %s authentication.\n", serverAddr, *demoAuth)
	var conn net.Conn
	var err error
	switch *demoAuth {
	case "tcp":
		conn, err = setupTCPClient()
	case "tls":
		conn, _, err = setupTLSClient()
	case "tao":
		conn, err = setupTaoClient()
	}
	if err != nil {
		fmt.Printf("client: error connecting to %s: %s\n", serverAddr, err.Error())
		return false
	}
	defer conn.Close()

	_, err = fmt.Fprintf(conn, "Hello\n")
	if err != nil {
		fmt.Printf("client: can't write: %s\n", err.Error())
		return false
	}
	msg, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		fmt.Printf("client can't read: %s\n", err.Error())
		return false
	}
	msg = strings.TrimSpace(msg)
	fmt.Printf("client: got reply: %s\n", msg)
	return true
}

func doClient() {
	pingGood := 0
	pingFail := 0
	for i := 0; i < *pingCount || *pingCount < 0; i++ { // negative means forever
		if doRequest() {
			pingGood++
		} else {
			pingFail++
		}
		fmt.Printf("client: made %d connections, finished %d ok, %d bad pings\n",
			i+1, pingGood, pingFail)
	}
}

func doResponse(conn net.Conn, responseOk chan<- bool) {
	defer conn.Close()

	switch *demoAuth {
	case "tcp", "tls":
		// authentication already done by lower layers
	case "tao":
		// TODO(kwalsh) Tao-level authorization: exchange names and delegation
		// attestations.
	}

	msg, err := bufio.NewReader(conn).ReadString('\n')
	if err != nil {
		fmt.Printf("server: can't read: %s\n", err.Error())
		conn.Close()
		responseOk <- false
		return
	}
	msg = strings.TrimSpace(msg)
	fmt.Printf("server: got message: %s\n", msg)
	responseOk <- true
	fmt.Fprintf(conn, "echo(%s)\n", msg)
	conn.Close()
}

func doServer(stop chan bool, ready, done chan<- bool) {
	var sock net.Listener
	var err error
	switch *demoAuth {
	case "tcp":
		sock, err = setupTCPServer()
	case "tls", "tao":
		sock, err = setupTLSServer()
	}
	if err != nil {
		fmt.Printf("server: can't listen at %s: %s\n", serverAddr, err.Error())
		ready <- false
		done <- true
		return
	}
	fmt.Printf("server: listening at %s using %s authentication.\n", serverAddr, *demoAuth)
	ready <- true

	pings := make(chan bool, 10)
	connCount := 0

	go func() {
		for connCount = 0; connCount < *pingCount || *pingCount < 0; connCount++ { // negative means forever
			conn, err := sock.Accept()
			if err != nil {
				fmt.Printf("server: can't accept connection: %s\n", err.Error())
				stop <- true
				return
			}
			go doResponse(conn, pings)
		}
	}()

	pingGood := 0
	pingFail := 0

loop:
	for {
		select {
		case <-stop:
			break loop
		case ok := <-pings:
			if ok {
				pingGood++
			} else {
				pingFail++
			}
		}
	}

	sock.Close()
	fmt.Printf("server: handled %d connections, finished %d ok, %d bad pings\n",
		connCount, pingGood, pingFail)

	done <- true
}

// Tao Host demo

func hostTaoDemo() error {
	name, err := tao.Host().GetTaoName()
	if err != nil {
		return err
	}
	fmt.Printf("My root name is %s\n", name)

	args := make([]string, len(os.Args))
	for index, arg := range os.Args {
		args[index] = strconv.Quote(arg)
	}
	subprin := "Args(" + strings.Join(args, ", ") + ")"
	err = tao.Host().ExtendTaoName(subprin)
	if err != nil {
		return err
	}

	name, err = tao.Host().GetTaoName()
	if err != nil {
		return err
	}
	fmt.Printf("My full name is %s\n", name)

	random, err := tao.Host().GetRandomBytes(10)
	if err != nil {
		return err
	}
	fmt.Printf("Random bytes  : % x\n", random)

	n, err := tao.Host().Rand().Read(random)
	if err != nil {
		return err
	}
	fmt.Printf("%d more bytes : % x\n", n, random)

	secret, err := tao.Host().GetSharedSecret(10, tao.SharedSecretPolicyDefault)
	if err != nil {
		return err
	}
	fmt.Printf("Shared secret : % x\n", secret)

	sealed, err := tao.Host().Seal(random, tao.SealPolicyDefault)
	if err != nil {
		return err
	}
	fmt.Printf("Sealed bytes  : % x\n", sealed)

	unsealed, policy, err := tao.Host().Unseal(sealed)
	if err != nil {
		return err
	}
	if policy != tao.SealPolicyDefault {
		return errors.New("unexpected policy on unseal")
	}
	fmt.Printf("Unsealed bytes: % x\n", unsealed)

	return nil
}

func main() {
	flag.Parse()
	serverAddr = *serverHost + ":" + *serverPort
	switch *demoAuth {
	case "tcp", "tls", "tao":
	default:
		fmt.Printf("unrecognized authentication mode: %s\n", *demoAuth)
		return
	}

	fmt.Printf("Go Tao Demo\n")

	if !tao.HostAvailable() {
		fmt.Printf("can't continue: No host Tao available\n")
		return
	}

	if *localMode {
		err := hostTaoDemo()
		if err != nil {
			fmt.Printf("error: %s\n", err.Error())
			return
		}
	}

	serverStop := make(chan bool, 1)
	serverReady := make(chan bool, 1)
	serverDone := make(chan bool, 1)

	if *serverMode {
		go doServer(serverStop, serverReady, serverDone)
	} else {
		serverReady <- true
		serverDone <- true
	}

	if *clientMode {
		ok := <-serverReady
		if ok {
			doClient()
		}
		serverStop <- true
	}

	<-serverDone
	fmt.Printf("Done\n")
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"koding/tools/log"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"time"
)

func main() {
	utils.Startup("broker", false)
	utils.RunStatusLogger()

	utils.AmqpAutoReconnect(func(consumeConn, publishConn *amqp.Connection) {

		service := sockjs.NewService("http://localhost/sockjs.js", true, false, 10*time.Minute, 0, func(receiveChan <-chan interface{}, sendChan chan<- interface{}) {
			defer log.RecoverAndLog()

			socketId := fmt.Sprintf("%x", rand.Int63())
			clientQueue := "broker-client-" + socketId
			exchanges := make([]string, 0)

			utils.ChangeNumClients <- 1
			log.Debug("Client connected: " + socketId)
			defer func() {
				utils.ChangeNumClients <- -1
				log.Debug("Client disconnected: " + socketId)
			}()

			controlChannel := utils.CreateAmqpChannel(publishConn)
			defer func() { controlChannel.Close() }() // controlChannel is replaced on error

			body, _ := json.Marshal(map[string]string{"socket_id": socketId})
			err := controlChannel.Publish("private-broker", "connected", false, false, amqp.Publishing{Body: body})
			if err != nil {
				panic(err)
			}

			defer func() {
				body, _ = json.Marshal(map[string]interface{}{"socket_id": socketId, "exchanges": exchanges})
				err = controlChannel.Publish("private-broker", "disconnected", false, false, amqp.Publishing{Body: body})
				if err != nil {
					panic(err)
				}
				body, _ = json.Marshal(map[string]interface{}{"socket_id": socketId})
				for _, exchange := range exchanges {
					err = controlChannel.Publish(exchange, "disconnected", false, false, amqp.Publishing{Body: body})
					if err != nil {
						panic(err)
					}
				}
			}()

			consumeChannel := utils.CreateAmqpChannel(consumeConn)
			defer consumeChannel.Close()
			consumerFinished := make(chan bool)
			defer close(consumerFinished)

			_, err = consumeChannel.QueueDeclare(clientQueue, false, true, false, false, nil)
			if err != nil {
				panic(err)
			}

			stream, err := consumeChannel.Consume(clientQueue, "", true, false, false, false, nil)
			if err != nil {
				panic(err)
			}

			go func() {
				defer log.RecoverAndLog()
				defer func() { consumerFinished <- true }()

				for message := range stream {
					body, _ = json.Marshal(map[string]string{"event": message.RoutingKey, "channel": message.Exchange, "payload": string(message.Body)})
					sendChan <- string(body)
				}
			}()

			for data := range receiveChan {
				func() {
					defer func() {
						err := recover()
						if err != nil {
							log.LogError(err)
							controlChannel.Close()
							controlChannel = utils.CreateAmqpChannel(publishConn)
						}
					}()

					var message map[string]string
					err := json.Unmarshal([]byte(data.(string)), &message)
					if err != nil {
						panic(err)
					}
					log.Debug(message)

					event := message["event"]
					exchange := message["channel"]

					switch event {
					case "client-subscribe":
						err = controlChannel.QueueBind(clientQueue, "#", exchange, false, nil)
						if err != nil {
							panic(err)
						}
						exchanges = append(exchanges, exchange)

						body, _ = json.Marshal(map[string]string{"event": "broker:subscription_succeeded", "channel": exchange, "payload": ""})
						sendChan <- string(body)

					case "client-unsubscribe":
						err = controlChannel.QueueUnbind(clientQueue, "#", exchange, nil)
						if err != nil {
							panic(err)
						}
						for i, e := range exchanges {
							if e == exchange {
								exchanges[i] = exchanges[len(exchanges)-1]
								exchanges = exchanges[:len(exchanges)-1]
								break
							}
						}

					case "client-bind-event":

					case "client-unbind-event":

					case "client-presence":

					default:
						if strings.HasPrefix(event, "client-") && strings.HasPrefix(exchange, "secret-") {
							err := controlChannel.Publish(exchange, event, false, false, amqp.Publishing{Body: []byte(message["payload"])})
							if err != nil {
								panic(err)
							}
						} else if message["vhost"] != "" {
							// ignored
						} else {
							log.Warn(fmt.Sprintf("Invalid message: %v", message))
						}

					}
				}()
			}

			consumeChannel.Close()
			<-consumerFinished
		})
		defer service.Close()

		server := &http.Server{Handler: &sockjs.Mux{
			Services: map[string]*sockjs.Service{
				"/subscribe": service,
			},
		}}
		listener, err := net.Listen("tcp", ":8008")
		if err != nil {
			panic(err)
		}

		go func() {
			for _ = range consumeConn.NotifyClose(make(chan *amqp.Error)) {
				listener.Close()
			}
		}()

		go func() {
			for _ = range publishConn.NotifyClose(make(chan *amqp.Error)) {
				listener.Close()
			}
		}()

		err = server.Serve(listener)
		if err != nil {
			log.Warn("Server error: " + err.Error())
		}
	})
}
<commit_msg>Hopefully fix of broker issue.<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"github.com/streadway/amqp"
	"koding/tools/log"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"time"
)

func main() {
	utils.Startup("broker", false)
	utils.RunStatusLogger()

	utils.AmqpAutoReconnect(func(consumeConn, publishConn *amqp.Connection) {

		service := sockjs.NewService("http://localhost/sockjs.js", true, false, 10*time.Minute, 0, func(receiveChan <-chan interface{}, sendChan chan<- interface{}) {
			defer log.RecoverAndLog()

			socketId := fmt.Sprintf("%x", rand.Int63())
			clientQueue := "broker-client-" + socketId
			exchanges := make([]string, 0)

			utils.ChangeNumClients <- 1
			log.Debug("Client connected: " + socketId)
			defer func() {
				utils.ChangeNumClients <- -1
				log.Debug("Client disconnected: " + socketId)
			}()

			controlChannel := utils.CreateAmqpChannel(publishConn)
			defer func() { controlChannel.Close() }() // controlChannel is replaced on error

			body, _ := json.Marshal(map[string]string{"socket_id": socketId})
			err := controlChannel.Publish("private-broker", "connected", false, false, amqp.Publishing{Body: body})
			if err != nil {
				panic(err)
			}

			defer func() {
				body, _ = json.Marshal(map[string]interface{}{"socket_id": socketId, "exchanges": exchanges})
				err = controlChannel.Publish("private-broker", "disconnected", false, false, amqp.Publishing{Body: body})
				if err != nil {
					panic(err)
				}
				body, _ = json.Marshal(map[string]interface{}{"socket_id": socketId})
				for _, exchange := range exchanges {
					err = controlChannel.Publish(exchange, "disconnected", false, false, amqp.Publishing{Body: body})
					if err != nil {
						panic(err)
					}
				}
			}()

			consumeChannel := utils.CreateAmqpChannel(consumeConn)
			defer consumeChannel.Close()
			consumerFinished := make(chan bool)
			defer close(consumerFinished)

			_, err = consumeChannel.QueueDeclare(clientQueue, false, true, false, false, nil)
			if err != nil {
				panic(err)
			}

			stream, err := consumeChannel.Consume(clientQueue, "", true, false, false, false, nil)
			if err != nil {
				panic(err)
			}

			go func() {
				defer log.RecoverAndLog()
				defer func() { consumerFinished <- true }()

				for message := range stream {
					go func() {
						defer log.RecoverAndLog()

						body, _ = json.Marshal(map[string]string{"event": message.RoutingKey, "channel": message.Exchange, "payload": string(message.Body)})
						sendChan <- string(body)
					}()
				}
			}()

			for data := range receiveChan {
				func() {
					defer func() {
						err := recover()
						if err != nil {
							log.LogError(err)
							controlChannel.Close()
							controlChannel = utils.CreateAmqpChannel(publishConn)
						}
					}()

					var message map[string]string
					err := json.Unmarshal([]byte(data.(string)), &message)
					if err != nil {
						panic(err)
					}
					log.Debug(message)

					event := message["event"]
					exchange := message["channel"]

					switch event {
					case "client-subscribe":
						err = controlChannel.QueueBind(clientQueue, "#", exchange, false, nil)
						if err != nil {
							panic(err)
						}
						exchanges = append(exchanges, exchange)

						body, _ = json.Marshal(map[string]string{"event": "broker:subscription_succeeded", "channel": exchange, "payload": ""})
						sendChan <- string(body)

					case "client-unsubscribe":
						err = controlChannel.QueueUnbind(clientQueue, "#", exchange, nil)
						if err != nil {
							panic(err)
						}
						for i, e := range exchanges {
							if e == exchange {
								exchanges[i] = exchanges[len(exchanges)-1]
								exchanges = exchanges[:len(exchanges)-1]
								break
							}
						}

					case "client-bind-event":

					case "client-unbind-event":

					case "client-presence":

					default:
						if strings.HasPrefix(event, "client-") && strings.HasPrefix(exchange, "secret-") {
							err := controlChannel.Publish(exchange, event, false, false, amqp.Publishing{Body: []byte(message["payload"])})
							if err != nil {
								panic(err)
							}
						} else if message["vhost"] != "" {
							// ignored
						} else {
							log.Warn(fmt.Sprintf("Invalid message: %v", message))
						}

					}
				}()
			}

			consumeChannel.Close()
			<-consumerFinished
		})
		defer service.Close()

		server := &http.Server{Handler: &sockjs.Mux{
			Services: map[string]*sockjs.Service{
				"/subscribe": service,
			},
		}}
		listener, err := net.Listen("tcp", ":8008")
		if err != nil {
			panic(err)
		}

		go func() {
			for _ = range consumeConn.NotifyClose(make(chan *amqp.Error)) {
				listener.Close()
			}
		}()

		go func() {
			for _ = range publishConn.NotifyClose(make(chan *amqp.Error)) {
				listener.Close()
			}
		}()

		err = server.Serve(listener)
		if err != nil {
			log.Warn("Server error: " + err.Error())
		}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"koding/databases/redis"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/fatih/set"
	"github.com/streadway/amqp"
)

const BROKER_NAME = "broker"

var (
	conf *config.Config
	log  = logger.New(BROKER_NAME)

	// routeMap holds the subscription list/set for any given routing key
	routeMap = make(map[string]*set.Set)

	// sessionsMap holds sessions with their socketIds
	sessionsMap = make(map[string]*sockjs.Session)

	globalMapMutex sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	flagProfile      = flag.String("c", "", "Configuration profile from file")
	flagBrokerDomain = flag.String("a", "", "Send kontrol a custom domain istead of os.Hostname")
	flagKontrolUUID  = flag.String("u", "", "Enable Kontrol mode")
	flagBrokerType   = flag.String("b", "broker", "Define broker type. Available: broker, premiumBroker and brokerKite, premiumBrokerKite. B")
	flagDebug        = flag.Bool("d", false, "Debug mode")
)

// Broker is a router/multiplexer that routes messages coming from a SockJS
// server to an AMQP exchange and vice versa. Broker basically listens to
// client messages (Koding users) from the SockJS server. The message is
// either passed to the appropriate exchange or a response is sent back to the
// client. Each message has an "action" field that defines how to act for a
// received message.
type Broker struct {
	Config            *config.Broker
	Hostname          string
	ServiceUniqueName string
	AuthAllExchange   string
	PublishConn       *amqp.Connection
	ConsumeConn       *amqp.Connection
	// we should open only one connection session to Redis for one broker
	RedisSingleton *redis.SingletonSession

	// Accepts SockJS connections
	listener net.Listener

	// Closed when SockJS server is ready to acccept connections
	ready chan struct{}
}

// NewBroker returns a new Broker instance with ServiceUniqueName and Hostname
// prepopulated. After creating a Broker instance, one has to call
// broker.Run() or broker.Start() to start the broker instance and call
// broker.Close() for a graceful stop.
func NewBroker(conf *config.Config) *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(*flagBrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := BROKER_NAME + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
		ready:             make(chan struct{}),
		RedisSingleton:    redis.Singleton(conf),
	}
}

func main() {
	flag.Parse()
	if *flagProfile == "" {
		log.Fatal("Please specify profile via -c. Aborting.")
	}

	conf = config.MustConfig(*flagProfile)
	broker := NewBroker(conf)

	switch *flagBrokerType {
	case "premiumBroker":
		broker.Config = &conf.PremiumBroker
	case "brokerKite":
		broker.Config = &conf.BrokerKite
	case "premiumBrokerKite":
		broker.Config = &conf.PremiumBrokerKite
	default:
		broker.Config = &conf.Broker
	}

	// update broker name
	log = logger.New(broker.Config.Name)
	var logLevel logger.Level
	if *flagDebug {
		logLevel = logger.DEBUG
	} else {
		logLevel = logger.GetLoggingLevelFromConfig(BROKER_NAME, *flagProfile)
	}

	log.SetLevel(logLevel)
	broker.Run()
}

// Run starts the broker.
func (b *Broker) Run() {
	// sets the maximum number of CPUs that can be executing simultaneously
	runtime.GOMAXPROCS(runtime.NumCPU())

	lifecycle.Startup(BROKER_NAME, false)
	logger.RunGaugesLoop(log)

	// Register broker to kontrol
	if err := b.registerToKontrol(); err != nil {
		log.Critical("Couldnt register to kontrol, stopping... %v", err)
		return
	}

	// Create AMQP exchanges/queues/bindings
	if err := b.startAMQP(); err != nil {
		log.Critical("Couldnt create amqp bindings, stopping... %v", err)
		return
	}

	// start listening/serving socket server
	b.startSockJS() // blocking
}

// Start is like Run() but waits until the SockJS listener is ready to be
// used.
func (b *Broker) Start() {
	go b.Run()
	<-b.ready
}

// Close close all amqp connections and closes the SockJS server listener
func (b *Broker) Close() {
	b.PublishConn.Close()
	b.ConsumeConn.Close()
	b.listener.Close()
}

// registerToKontrol registers the broker to KontrolDaemon. This is needed to
// populate a list of brokers and show them to the client. The list is
// available at: https://koding.com/-/services/broker?all
func (b *Broker) registerToKontrol() error {
	if err := kontrolhelper.RegisterToKontrol(
		conf,
		b.Config.Name,
		b.Config.ServiceGenericName, // servicGenericName
		b.ServiceUniqueName,
		*flagKontrolUUID,
		b.Hostname,
		b.Config.Port,
	); err != nil {
		return err
	}
	return nil
}

// startAMQP setups the the neccesary publisher and consumer connections for
// the broker broker.
func (b *Broker) startAMQP() error {
	b.PublishConn = amqputil.CreateConnection(conf, b.Config.Name)
	b.ConsumeConn = amqputil.CreateConnection(conf, b.Config.Name)
	consumeChannel := amqputil.CreateChannel(b.ConsumeConn)
	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,              // channel
		"services-presence",         // exchange
		b.Config.Name,               // serviceType
		b.Config.ServiceGenericName, // serviceGenericName
		b.ServiceUniqueName,         // serviceUniqueName
		false,                       // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", b.Config.ServiceGenericName, "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		return fmt.Errorf("Couldnt create updateInstances exchange  %v", err)
	}

	if err := consumeChannel.ExchangeBind(BROKER_NAME, "", "updateInstances", false, nil); err != nil {
		return fmt.Errorf("Couldnt bind to updateInstances exchange  %v", err)
	}

	go func(stream <-chan amqp.Delivery) {
		// start to listen from "broker" topic exchange
		for amqpMessage := range stream {
			sendMessageToClient(amqpMessage)
		}

		b.Close()

	}(stream)

	return nil
}

// sendMessageToClient takes an amqp messsage and delivers it to the related
// clients which are subscribed to the routing key
func sendMessageToClient(amqpMessage amqp.Delivery) {
	routingKey := amqpMessage.RoutingKey
	payloadsByte := utils.FilterInvalidUTF8(amqpMessage.Body)

	// We are sending multiple bodies for updateInstances exchange
	// so that there will be another operations, if exchange is not "updateInstances"
	// no need to add more overhead
	if amqpMessage.Exchange != "updateInstances" {
		payloadRaw := json.RawMessage(payloadsByte)
		processMessage(routingKey, &payloadRaw)
		return
	}

	// this part is only for updateInstances exchange
	var payloads []interface{}
	// unmarshal data to slice of interface
	if err := json.Unmarshal(payloadsByte, &payloads); err != nil {
		log.Error("Error while unmarshalling:%v data:%v routingKey:%v", err, string(payloadsByte), routingKey)
		return
	}

	// range over the slice and send all of them to the same routingkey
	for _, payload := range payloads {
		payloadByte, err := json.Marshal(payload)
		if err != nil {
			log.Error("Error while marshalling:%v data:%v routingKey:%v", err, string(payloadByte), routingKey)
			continue
		}
		payloadByteRaw := json.RawMessage(payloadByte)
		processMessage(routingKey, &payloadByteRaw)
	}
}

// processMessage gets routingKey and a payload for sending them to the client
// Gets subscription bindings from global routeMap
func processMessage(routingKey string, payload interface{}) {
	pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
	for pos != -1 && pos < len(routingKey) {
		index := strings.IndexRune(routingKey[pos+1:], '.')
		pos += index + 1
		if index == -1 {
			pos = len(routingKey)
		}
		routingKeyPrefix := routingKey[:pos]
		globalMapMutex.Lock()

		if routes, ok := routeMap[routingKeyPrefix]; ok {
			routes.Each(func(sessionId interface{}) bool {
				if routeSession, ok := sessionsMap[sessionId.(string)]; ok {
					sendToClient(routeSession, routingKey, &payload)
				}
				return true
			})
		}
		globalMapMutex.Unlock()
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		conf.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	mux := &sockjs.Mux{
		Handlers: map[string]http.Handler{
			"/subscribe": service,
			"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				w.Write([]byte(strconv.Itoa(conf.BuildNumber)))
			}),
		},
	}

	server := &http.Server{Handler: mux}

	var err error
	b.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(b.Config.IP), Port: b.Config.Port})
	if err != nil {
		log.Fatal(err)
	}

	if b.Config.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(b.Config.CertFile, b.Config.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		b.listener = tls.NewListener(b.listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	// signal that we are ready now
	close(b.ready)

	lastErrorTime := time.Now()
	for {
		err := server.Serve(b.listener)
		if err != nil {
			// comes when the broker is closed with Close() method. This error
			// is defined in net/net.go as "var errClosing", unfortunaly it's
			// not exported.
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}

			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sockjsSession is called for every client connection and handles all the
// message trafic for a single client connection.
func (b *Broker) sockjsSession(session *sockjs.Session) {
	clientChan := make(chan *Client, 0)
	errChan := make(chan error, 0)

	// we don't use time.After because we are stopping the time if the given
	// functions are fast enough.
	timer := time.NewTicker(5 * time.Second)

	go createClient(b, session, clientChan, errChan)

	var client *Client
	var ok bool
	select {
	case client, ok = <-clientChan:
		if ok {
			timer.Stop()
		}
	case err, ok := <-errChan:
		if ok {
			timer.Stop()
			log.Critical("An error occured while creating client %v", err)
			return
		}
	case <-timer.C:
		timer.Stop()
		log.Critical("Client coulnt created in %s exiting ", "5s")
		return
	}

	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		client.handleSessionMessage(data)
	}
}

func createClient(b *Broker, session *sockjs.Session, clientChan chan *Client, errChan chan error) {
	// do not forget to close channels
	defer close(errChan)
	defer close(clientChan)

	client, err := NewClient(session, b)
	if err != nil {
		log.Critical("Couldnt create client %v", err)
		errChan <- err
		return
	}

	err = client.ControlChannel.Publish(b.Config.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(client.SocketId)})
	if err != nil {
		log.Critical("Couldnt publish to control channel %v", err)
		errChan <- err
		return
	}

	// if session is closed before the client creation no need to send
	// client object to listeners
	if !session.Closed {
		sendToClient(session, "broker.connected", client.SocketId)
		clientChan <- client
	}
}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}
<commit_msg>Broker: change ticker to time.after and get duration from flag<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"koding/databases/redis"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/fatih/set"
	"github.com/streadway/amqp"
)

const BROKER_NAME = "broker"

var (
	conf *config.Config
	log  = logger.New(BROKER_NAME)

	// routeMap holds the subscription list/set for any given routing key
	routeMap = make(map[string]*set.Set)

	// sessionsMap holds sessions with their socketIds
	sessionsMap = make(map[string]*sockjs.Session)

	globalMapMutex sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	flagProfile      = flag.String("c", "", "Configuration profile from file")
	flagBrokerDomain = flag.String("a", "", "Send kontrol a custom domain istead of os.Hostname")
	flagDuration     = flag.Duration("t", time.Second*5, "Duration for timeout in seconds - Duration flag accept any input valid for time.ParseDuration.")
	flagKontrolUUID  = flag.String("u", "", "Enable Kontrol mode")
	flagBrokerType   = flag.String("b", "broker", "Define broker type. Available: broker, premiumBroker and brokerKite, premiumBrokerKite. B")
	flagDebug        = flag.Bool("d", false, "Debug mode")
)

// Broker is a router/multiplexer that routes messages coming from a SockJS
// server to an AMQP exchange and vice versa. Broker basically listens to
// client messages (Koding users) from the SockJS server. The message is
// either passed to the appropriate exchange or a response is sent back to the
// client. Each message has an "action" field that defines how to act for a
// received message.
type Broker struct {
	Config            *config.Broker
	Hostname          string
	ServiceUniqueName string
	AuthAllExchange   string
	PublishConn       *amqp.Connection
	ConsumeConn       *amqp.Connection
	// we should open only one connection session to Redis for one broker
	RedisSingleton *redis.SingletonSession

	// Accepts SockJS connections
	listener net.Listener

	// Closed when SockJS server is ready to acccept connections
	ready chan struct{}
}

// NewBroker returns a new Broker instance with ServiceUniqueName and Hostname
// prepopulated. After creating a Broker instance, one has to call
// broker.Run() or broker.Start() to start the broker instance and call
// broker.Close() for a graceful stop.
func NewBroker(conf *config.Config) *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(*flagBrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := BROKER_NAME + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
		ready:             make(chan struct{}),
		RedisSingleton:    redis.Singleton(conf),
	}
}

func main() {
	flag.Parse()
	if *flagProfile == "" {
		log.Fatal("Please specify profile via -c. Aborting.")
	}

	conf = config.MustConfig(*flagProfile)
	broker := NewBroker(conf)

	switch *flagBrokerType {
	case "premiumBroker":
		broker.Config = &conf.PremiumBroker
	case "brokerKite":
		broker.Config = &conf.BrokerKite
	case "premiumBrokerKite":
		broker.Config = &conf.PremiumBrokerKite
	default:
		broker.Config = &conf.Broker
	}

	// update broker name
	log = logger.New(broker.Config.Name)
	var logLevel logger.Level
	if *flagDebug {
		logLevel = logger.DEBUG
	} else {
		logLevel = logger.GetLoggingLevelFromConfig(BROKER_NAME, *flagProfile)
	}

	log.SetLevel(logLevel)
	broker.Run()
}

// Run starts the broker.
func (b *Broker) Run() {
	// sets the maximum number of CPUs that can be executing simultaneously
	runtime.GOMAXPROCS(runtime.NumCPU())

	lifecycle.Startup(BROKER_NAME, false)
	logger.RunGaugesLoop(log)

	// Register broker to kontrol
	if err := b.registerToKontrol(); err != nil {
		log.Critical("Couldnt register to kontrol, stopping... %v", err)
		return
	}

	// Create AMQP exchanges/queues/bindings
	if err := b.startAMQP(); err != nil {
		log.Critical("Couldnt create amqp bindings, stopping... %v", err)
		return
	}

	// start listening/serving socket server
	b.startSockJS() // blocking
}

// Start is like Run() but waits until the SockJS listener is ready to be
// used.
func (b *Broker) Start() {
	go b.Run()
	<-b.ready
}

// Close close all amqp connections and closes the SockJS server listener
func (b *Broker) Close() {
	b.PublishConn.Close()
	b.ConsumeConn.Close()
	b.listener.Close()
}

// registerToKontrol registers the broker to KontrolDaemon. This is needed to
// populate a list of brokers and show them to the client. The list is
// available at: https://koding.com/-/services/broker?all
func (b *Broker) registerToKontrol() error {
	if err := kontrolhelper.RegisterToKontrol(
		conf,
		b.Config.Name,
		b.Config.ServiceGenericName, // servicGenericName
		b.ServiceUniqueName,
		*flagKontrolUUID,
		b.Hostname,
		b.Config.Port,
	); err != nil {
		return err
	}
	return nil
}

// startAMQP setups the the neccesary publisher and consumer connections for
// the broker broker.
func (b *Broker) startAMQP() error {
	b.PublishConn = amqputil.CreateConnection(conf, b.Config.Name)
	b.ConsumeConn = amqputil.CreateConnection(conf, b.Config.Name)
	consumeChannel := amqputil.CreateChannel(b.ConsumeConn)
	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,              // channel
		"services-presence",         // exchange
		b.Config.Name,               // serviceType
		b.Config.ServiceGenericName, // serviceGenericName
		b.ServiceUniqueName,         // serviceUniqueName
		false,                       // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", b.Config.ServiceGenericName, "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		return fmt.Errorf("Couldnt create updateInstances exchange  %v", err)
	}

	if err := consumeChannel.ExchangeBind(BROKER_NAME, "", "updateInstances", false, nil); err != nil {
		return fmt.Errorf("Couldnt bind to updateInstances exchange  %v", err)
	}

	go func(stream <-chan amqp.Delivery) {
		// start to listen from "broker" topic exchange
		for amqpMessage := range stream {
			sendMessageToClient(amqpMessage)
		}

		b.Close()

	}(stream)

	return nil
}

// sendMessageToClient takes an amqp messsage and delivers it to the related
// clients which are subscribed to the routing key
func sendMessageToClient(amqpMessage amqp.Delivery) {
	routingKey := amqpMessage.RoutingKey
	payloadsByte := utils.FilterInvalidUTF8(amqpMessage.Body)

	// We are sending multiple bodies for updateInstances exchange
	// so that there will be another operations, if exchange is not "updateInstances"
	// no need to add more overhead
	if amqpMessage.Exchange != "updateInstances" {
		payloadRaw := json.RawMessage(payloadsByte)
		processMessage(routingKey, &payloadRaw)
		return
	}

	// this part is only for updateInstances exchange
	var payloads []interface{}
	// unmarshal data to slice of interface
	if err := json.Unmarshal(payloadsByte, &payloads); err != nil {
		log.Error("Error while unmarshalling:%v data:%v routingKey:%v", err, string(payloadsByte), routingKey)
		return
	}

	// range over the slice and send all of them to the same routingkey
	for _, payload := range payloads {
		payloadByte, err := json.Marshal(payload)
		if err != nil {
			log.Error("Error while marshalling:%v data:%v routingKey:%v", err, string(payloadByte), routingKey)
			continue
		}
		payloadByteRaw := json.RawMessage(payloadByte)
		processMessage(routingKey, &payloadByteRaw)
	}
}

// processMessage gets routingKey and a payload for sending them to the client
// Gets subscription bindings from global routeMap
func processMessage(routingKey string, payload interface{}) {
	pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
	for pos != -1 && pos < len(routingKey) {
		index := strings.IndexRune(routingKey[pos+1:], '.')
		pos += index + 1
		if index == -1 {
			pos = len(routingKey)
		}
		routingKeyPrefix := routingKey[:pos]
		globalMapMutex.Lock()

		if routes, ok := routeMap[routingKeyPrefix]; ok {
			routes.Each(func(sessionId interface{}) bool {
				if routeSession, ok := sessionsMap[sessionId.(string)]; ok {
					sendToClient(routeSession, routingKey, &payload)
				}
				return true
			})
		}
		globalMapMutex.Unlock()
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		conf.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	mux := &sockjs.Mux{
		Handlers: map[string]http.Handler{
			"/subscribe": service,
			"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				w.Write([]byte(strconv.Itoa(conf.BuildNumber)))
			}),
		},
	}

	server := &http.Server{Handler: mux}

	var err error
	b.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(b.Config.IP), Port: b.Config.Port})
	if err != nil {
		log.Fatal(err)
	}

	if b.Config.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(b.Config.CertFile, b.Config.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		b.listener = tls.NewListener(b.listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	// signal that we are ready now
	close(b.ready)

	lastErrorTime := time.Now()
	for {
		err := server.Serve(b.listener)
		if err != nil {
			// comes when the broker is closed with Close() method. This error
			// is defined in net/net.go as "var errClosing", unfortunaly it's
			// not exported.
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}

			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sockjsSession is called for every client connection and handles all the
// message trafic for a single client connection.
func (b *Broker) sockjsSession(session *sockjs.Session) {
	clientChan := make(chan *Client, 0)
	errChan := make(chan error, 0)

	go createClient(b, session, clientChan, errChan)

	// Return if there is any error or if we don't get the result in 5 seconds back
	var client *Client
	select {
	case client = <-clientChan:
	case err := <-errChan:
		log.Critical("An error occured while creating client %v", err)
		return
	case <-time.After(*flagDuration):
		log.Critical("Client coulnt created in %s exiting ", flagDuration.String())
		return
	}

	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		client.handleSessionMessage(data)
	}
}

func createClient(b *Broker, session *sockjs.Session, clientChan chan *Client, errChan chan error) {
	// do not forget to close channels
	defer close(errChan)
	defer close(clientChan)

	client, err := NewClient(session, b)
	if err != nil {
		log.Critical("Couldnt create client %v", err)
		errChan <- err
		return
	}

	err = client.ControlChannel.Publish(b.Config.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(client.SocketId)})
	if err != nil {
		log.Critical("Couldnt publish to control channel %v", err)
		errChan <- err
		return
	}

	// if session is closed before the client creation no need to send
	// client object to listeners
	if !session.Closed {
		sendToClient(session, "broker.connected", client.SocketId)
		clientChan <- client
	}
}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}
<|endoftext|>