<commit_before>package main

import (
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"time"

	"github.com/gorilla/websocket"
	sm "github.com/ifn/go-statemachine"
)

type DeskMsg struct {
	Desk [][]string `json:"desk"`
}

type PlayerMsg struct {
	Cmd  sm.EventType `json:"command"`
	Card string       `json:"card"`
}

//

var Suits []string = []string{
	"S",
	"C",
	"H",
	"D",
}

var CardValues map[string]int = map[string]int{
	"6": 6, "7": 7, "8": 8, "9": 9, "10": 10,
	"J": 11,
	"Q": 12,
	"K": 13,
	"A": 14,
}

func higher(c0, c1, t string) int {
	// c0 and c1 have the same suit
	if c0[0] == c1[0] {
		if CardValues[c0[1:]] > CardValues[c1[1:]] {
			return 1
		}
		if CardValues[c0[1:]] < CardValues[c1[1:]] {
			return -1
		}
		return 0
	}
	// c0 is trump, c1 is not
	if c0[:1] == t {
		return 1
	}
	// c1 is trump, c0 is not
	if c1[:1] == t {
		return -1
	}
	// suits are different, both are not trump
	return -2
}

var CardRE *regexp.Regexp = regexp.MustCompile(`[SCHD]([6-9JQKA]|10)`)

func isValid(c string) bool {
	return CardRE.MatchString(c)
}

//

const (
	cmdStart sm.EventType = iota
	cmdMove

	cmdCount
)

const (
	stateCollection sm.State = iota

	stateAttack
	stateDefense

	stateCount
)

type roundResult int

const (
	None roundResult = iota
	Beat
	NotBeat
)

func stateToString(s sm.State) string {
	return [...]string{
		stateCollection: "COLLECTION",
		stateDefense:    "DEFENSE",
		stateAttack:     "ATTACK",
	}[s]
}

func cmdToString(t sm.EventType) string {
	return [...]string{
		cmdStart: "START",
		cmdMove:  "MOVE",
	}[t]
}

//

type cmdArgs struct {
	conn *playerConn
	card string
}

//

func logOutOfTurn(pconn *playerConn) {
	log.Printf("out of turn: %v", pconn.conn.RemoteAddr())
}

func logWontBeat(c1, c2, t string) {
	log.Printf("%v won't bit %v, trump is %v", c1, c2, t)
}

func logNoQuorum(n int) {
	log.Printf("no quorum: %v player(s)", n)
}

//

type gameState struct {
	// 1. fields that don't change during a game

	sm  *sm.StateMachine
	hub *hub

	// trump suit
	trump string

	// 2. fields that don't change during a round

	deck []string

	// attacker that started a round
	aconnStart *playerConn
	// defender
	dconn *playerConn

	// 3. fields that change during a round

	// attacker
	aconn *playerConn
	// card that should be beaten
	cardToBeat string
}

func (self *gameState) popCard() (card string) {
	if deck := self.deck; len(deck) > 0 {
		card = deck[0]
		self.deck = deck[1:]
	}
	return
}

func (self *gameState) initDeck() {
	numCards := len(Suits) * len(CardValues)
	self.deck = make([]string, 0, numCards+ /*for trump*/ 1)

	deck := make([]string, 0, numCards)
	for i := range Suits {
		for cv := range CardValues {
			deck = append(deck, Suits[i]+cv)
		}
	}

	order := rand.Perm(numCards)
	for i := range order {
		self.deck = append(self.deck, deck[order[i]])
	}
}

func (self *gameState) setTrump() {
	card := self.popCard()
	self.trump = card[:1]
	self.deck = append(self.deck, card)
}

func (self *gameState) nextPlayer(c *playerConn) *playerConn {
	return self.hub.conns.(*mapRing).Next(c).(*playerConn)
}

func (self *gameState) chooseStarting() *playerConn {
	conns := self.hub.conns.(*mapRing)

	return conns.Nth(rand.Intn(conns.Len())).(*playerConn)
}

func (self *gameState) setRoles(res roundResult) {
	switch res {
	case None:
		self.aconn = self.chooseStarting()
	case Beat:
		self.aconn = self.dconn
	case NotBeat:
		self.aconn = self.nextPlayer(self.dconn)
	}
	self.dconn = self.nextPlayer(self.aconn)
	self.aconnStart = self.aconn
}

func (self *gameState) dealCards() {
	for i := 0; i < 6; i++ {
		for pc := range self.hub.conns.Enumerate() {
			pc.(*playerConn).takeCard()
		}
	}
}

func (self *gameState) takeCards() {
	conns := self.hub.conns.(*mapRing)

	takeCards := func(pc *playerConn) {
		for len(self.deck) > 0 && len(pc.cards) < 6 {
			pc.takeCard()
		}
	}

	for pc := range conns.EnumerateFrom(self.aconnStart) {
		if pc := pc.(*playerConn); pc != self.dconn {
			takeCards(pc)
		}
	}
	takeCards(self.dconn)
}

func (self *gameState) newRound(res roundResult) {
	if res == None {
		self.initDeck()
		self.dealCards()
		self.setTrump()
	} else {
		self.takeCards()
	}
	self.setRoles(res)
}

// event handlers
// event handlers are actually transition functions.
// in case error event handler should neither change the gameState,
// nor return the state value different from passed to it as an argument.

func (self *gameState) handleStartInCollection(s sm.State, e *sm.Event) sm.State {
	if n := self.hub.conns.(*mapRing).Len(); n < 2 {
		logNoQuorum(n)
		return s
	}

	self.newRound(None)
	return stateAttack
}

func (self *gameState) handleMoveInAttack(s sm.State, e *sm.Event) sm.State {
	conn := e.Data.(cmdArgs).conn
	card := e.Data.(cmdArgs).card

	// check that it's conn's turn to move
	if conn != self.aconn {
		logOutOfTurn(conn)
		return s
	}

	// attacker sent the card
	if card != "" {
		self.cardToBeat = card
		return stateDefense
	}

	// attacker sent no card

	aconn := self.nextPlayer(self.aconn)
	if aconn == self.dconn {
		aconn = self.nextPlayer(aconn)
	}

	// check if all attackers have been polled
	if aconn == self.aconnStart {
		self.newRound(Beat)
		return stateAttack
	}

	self.aconn = aconn
	return stateAttack
}

func (self *gameState) handleMoveInDefense(s sm.State, e *sm.Event) sm.State {
	conn := e.Data.(cmdArgs).conn
	card := e.Data.(cmdArgs).card

	// check that it's conn's turn to move
	if conn != self.dconn {
		logOutOfTurn(conn)
		return s
	}

	// defender takes the cards
	if card == "" {
		self.newRound(NotBeat)
		return stateAttack
	}

	// check that the sent card is capable to beat
	if higher(card, self.cardToBeat, self.trump) != 1 {
		logWontBeat(card, self.cardToBeat, self.trump)
		return s
	}

	return stateAttack
}

func (self *gameState) showDesk(s sm.State, e *sm.Event) sm.State {
	desk, err := json.Marshal(DeskMsg{})
	if err != nil {
		log.Println(err)
		return s
	}

	self.hub.bcastChan <- desk

	return s
}

func (self *gameState) log(s sm.State, e *sm.Event) sm.State {
	log.Println(self)
	return s
}

//

func NewGameState() *gameState {
	gst := new(gameState)

	gst.sm = sm.New(stateCollection, uint(stateCount), uint(cmdCount))

	//TODO: refactor this
	gst.sm.On(cmdStart,
		[]sm.State{stateCollection},
		gst.handleStartInCollection,
	)
	gst.sm.On(cmdStart,
		[]sm.State{stateCollection},
		gst.showDesk,
	)
	gst.sm.On(cmdStart,
		[]sm.State{stateCollection},
		gst.log,
	)

	gst.sm.On(cmdMove,
		[]sm.State{stateAttack},
		gst.handleMoveInAttack,
	)
	gst.sm.On(cmdMove,
		[]sm.State{stateDefense},
		gst.handleMoveInDefense,
	)

	gst.sm.On(cmdMove,
		[]sm.State{stateAttack, stateDefense},
		gst.showDesk,
	)
	gst.sm.On(cmdMove,
		[]sm.State{stateAttack, stateDefense},
		gst.log,
	)

	gst.hub = NewHub()

	return gst
}

//

type playerConn struct {
	gst *gameState

	cards map[string]struct{}

	conn      *websocket.Conn
	hubToConn chan []byte
}

func (self *playerConn) takeCard() {
	if card := self.gst.popCard(); card != "" {
		self.cards[card] = struct{}{}
	}
}

func (self *playerConn) write() {
	defer func() {
		err := self.conn.Close()
		if err != nil {
			log.Println(err)
		}
	}()

	for {
		select {
		case m, ok := <-self.hubToConn:
			if !ok {
				return
			}
			//TODO: text or binary?
			err := self.conn.WriteMessage(websocket.TextMessage, m)
			if err != nil {
				log.Println(err)
				return
			}
		}
	}
}

func (self *playerConn) read() {
	defer func() {
		err := self.conn.Close()
		if err != nil {
			log.Println(err)
		}
	}()

	var m PlayerMsg

	for {
		err := self.conn.ReadJSON(&m)
		if err != nil {
			log.Println(err)
			return
		}

		var event *sm.Event
		switch m.Cmd {
		case cmdStart:
			event = &sm.Event{cmdStart, nil}
		case cmdMove:
			event = &sm.Event{cmdMove, cmdArgs{self, m.Card}}
		default:
			log.Printf("unknown command: %v", m.Cmd)
			continue
		}

		err = self.gst.sm.Emit(event)
		if err != nil {
			log.Println(err)
		}
	}
}

//

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

func playerHandler(gst *gameState) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}

		p := &playerConn{gst, make(map[string]struct{}), conn, make(chan []byte)}

		gst.hub.regChan <- p
		defer func() {
			gst.hub.unregChan <- p
		}()

		go p.write()
		p.read()
	}
}

//

func startDurakSrv() error {
	gst := NewGameState()

	http.HandleFunc("/", playerHandler(gst))

	return http.ListenAndServe(":"+os.Getenv("PORT"), nil)
}

func main() {
	err := startDurakSrv()
	if err != nil {
		log.Fatal(err)
	}
}

func init() {
	rand.Seed(time.Now().UnixNano())

	runtime.GOMAXPROCS(runtime.NumCPU())
}
<commit_msg>added toDesk<commit_after>package main

import (
	"encoding/json"
	"log"
	"math/rand"
	"net/http"
	"os"
	"regexp"
	"runtime"
	"time"

	"github.com/gorilla/websocket"
	sm "github.com/ifn/go-statemachine"
)

type DeskMsg struct {
	Desk [][]string `json:"desk"`
}

type PlayerMsg struct {
	Cmd  sm.EventType `json:"command"`
	Card string       `json:"card"`
}

//

var Suits []string = []string{
	"S",
	"C",
	"H",
	"D",
}

var CardValues map[string]int = map[string]int{
	"6": 6, "7": 7, "8": 8, "9": 9, "10": 10,
	"J": 11,
	"Q": 12,
	"K": 13,
	"A": 14,
}

func higher(c0, c1, t string) int {
	// c0 and c1 have the same suit
	if c0[0] == c1[0] {
		if CardValues[c0[1:]] > CardValues[c1[1:]] {
			return 1
		}
		if CardValues[c0[1:]] < CardValues[c1[1:]] {
			return -1
		}
		return 0
	}
	// c0 is trump, c1 is not
	if c0[:1] == t {
		return 1
	}
	// c1 is trump, c0 is not
	if c1[:1] == t {
		return -1
	}
	// suits are different, both are not trump
	return -2
}

var CardRE *regexp.Regexp = regexp.MustCompile(`[SCHD]([6-9JQKA]|10)`)

func isValid(c string) bool {
	return CardRE.MatchString(c)
}

//

const (
	cmdStart sm.EventType = iota
	cmdMove

	cmdCount
)

const (
	stateCollection sm.State = iota

	stateAttack
	stateDefense

	stateCount
)

type roundResult int

const (
	None roundResult = iota
	Beat
	NotBeat
)

func stateToString(s sm.State) string {
	return [...]string{
		stateCollection: "COLLECTION",
		stateDefense:    "DEFENSE",
		stateAttack:     "ATTACK",
	}[s]
}

func cmdToString(t sm.EventType) string {
	return [...]string{
		cmdStart: "START",
		cmdMove:  "MOVE",
	}[t]
}

//

type cmdArgs struct {
	conn *playerConn
	card string
}

//

func logOutOfTurn(pconn *playerConn) {
	log.Printf("out of turn: %v", pconn.conn.RemoteAddr())
}

func logWontBeat(c1, c2, t string) {
	log.Printf("%v won't bit %v, trump is %v", c1, c2, t)
}

func logNoQuorum(n int) {
	log.Printf("no quorum: %v player(s)", n)
}

func logStrangeCard(c string, pc *playerConn) {
	log.Printf("%v doesn't own %v", pc, c)
}

//

type gameState struct {
	// 1. fields that don't change during a game

	sm  *sm.StateMachine
	hub *hub

	// trump suit
	trump string

	// 2. fields that don't change during a round

	deck []string

	// attacker that started a round
	aconnStart *playerConn
	// defender
	dconn *playerConn

	// 3. fields that change during a round

	desk [][]string

	// attacker
	aconn *playerConn
	// card that should be beaten
	cardToBeat string
}

func (self *gameState) popCard() (card string) {
	if deck := self.deck; len(deck) > 0 {
		card = deck[0]
		self.deck = deck[1:]
	}
	return
}

func (self *gameState) initDeck() {
	numCards := len(Suits) * len(CardValues)
	self.deck = make([]string, 0, numCards+ /*for trump*/ 1)

	deck := make([]string, 0, numCards)
	for i := range Suits {
		for cv := range CardValues {
			deck = append(deck, Suits[i]+cv)
		}
	}

	order := rand.Perm(numCards)
	for i := range order {
		self.deck = append(self.deck, deck[order[i]])
	}
}

func (self *gameState) setTrump() {
	card := self.popCard()
	self.trump = card[:1]
	self.deck = append(self.deck, card)
}

func (self *gameState) nextPlayer(c *playerConn) *playerConn {
	return self.hub.conns.(*mapRing).Next(c).(*playerConn)
}

func (self *gameState) chooseStarting() *playerConn {
	conns := self.hub.conns.(*mapRing)

	return conns.Nth(rand.Intn(conns.Len())).(*playerConn)
}

func (self *gameState) setRoles(res roundResult) {
	switch res {
	case None:
		self.aconn = self.chooseStarting()
	case Beat:
		self.aconn = self.dconn
	case NotBeat:
		self.aconn = self.nextPlayer(self.dconn)
	}
	self.dconn = self.nextPlayer(self.aconn)
	self.aconnStart = self.aconn
}

func (self *gameState) dealCards() {
	for i := 0; i < 6; i++ {
		for pc := range self.hub.conns.Enumerate() {
			pc.(*playerConn).takeCard()
		}
	}
}

func (self *gameState) takeCards() {
	conns := self.hub.conns.(*mapRing)

	takeCards := func(pc *playerConn) {
		for len(self.deck) > 0 && len(pc.cards) < 6 {
			pc.takeCard()
		}
	}

	for pc := range conns.EnumerateFrom(self.aconnStart) {
		if pc := pc.(*playerConn); pc != self.dconn {
			takeCards(pc)
		}
	}
	takeCards(self.dconn)
}

func (self *gameState) newRound(res roundResult) {
	switch res {
	case None:
		self.initDeck()
		self.dealCards()
		self.setTrump()
	case NotBeat:
		// add desk to defender's stock
		self.takeCards()
	case Beat:
		self.takeCards()
	}
	self.setRoles(res)
}

// event handlers
// event handlers are actually transition functions.
// in case error event handler should neither change the gameState,
// nor return the state value different from passed to it as an argument.

func (self *gameState) handleStartInCollection(s sm.State, e *sm.Event) sm.State {
	if n := self.hub.conns.(*mapRing).Len(); n < 2 {
		logNoQuorum(n)
		return s
	}

	self.newRound(None)
	return stateAttack
}

func (self *gameState) handleMoveInAttack(s sm.State, e *sm.Event) sm.State {
	conn := e.Data.(cmdArgs).conn
	card := e.Data.(cmdArgs).card

	// check that it's conn's turn to move
	if conn != self.aconn {
		logOutOfTurn(conn)
		return s
	}

	// attacker sent the card
	if card != "" {
		// throw card to desk
		if _, ok := conn.cards[card]; !ok {
			logStrangeCard(card, conn)
			return s
		}
		conn.toDesk(card)

		self.cardToBeat = card
		return stateDefense
	}

	// attacker sent no card

	aconn := self.nextPlayer(self.aconn)
	if aconn == self.dconn {
		aconn = self.nextPlayer(aconn)
	}

	// check if all attackers have been polled
	if aconn == self.aconnStart {
		self.newRound(Beat)
		return stateAttack
	}

	self.aconn = aconn
	return stateAttack
}

func (self *gameState) handleMoveInDefense(s sm.State, e *sm.Event) sm.State {
	conn := e.Data.(cmdArgs).conn
	card := e.Data.(cmdArgs).card

	// check that it's conn's turn to move
	if conn != self.dconn {
		logOutOfTurn(conn)
		return s
	}

	// defender takes the cards
	if card == "" {
		self.newRound(NotBeat)
		return stateAttack
	}

	// check that the sent card is capable to beat
	if higher(card, self.cardToBeat, self.trump) != 1 {
		logWontBeat(card, self.cardToBeat, self.trump)
		return s
	}

	// throw card to desk
	if _, ok := conn.cards[card]; !ok {
		logStrangeCard(card, conn)
		return s
	}
	conn.toDesk(card)

	return stateAttack
}

func (self *gameState) showDesk(s sm.State, e *sm.Event) sm.State {
	desk, err := json.Marshal(DeskMsg{})
	if err != nil {
		log.Println(err)
		return s
	}

	self.hub.bcastChan <- desk

	return s
}

func (self *gameState) log(s sm.State, e *sm.Event) sm.State {
	log.Println(self)
	return s
}

//

func NewGameState() *gameState {
	gst := new(gameState)

	gst.sm = sm.New(stateCollection, uint(stateCount), uint(cmdCount))

	//TODO: refactor this
	gst.sm.On(cmdStart,
		[]sm.State{stateCollection},
		gst.handleStartInCollection,
	)
	gst.sm.On(cmdStart,
		[]sm.State{stateCollection},
		gst.showDesk,
	)
	gst.sm.On(cmdStart,
		[]sm.State{stateCollection},
		gst.log,
	)

	gst.sm.On(cmdMove,
		[]sm.State{stateAttack},
		gst.handleMoveInAttack,
	)
	gst.sm.On(cmdMove,
		[]sm.State{stateDefense},
		gst.handleMoveInDefense,
	)

	gst.sm.On(cmdMove,
		[]sm.State{stateAttack, stateDefense},
		gst.showDesk,
	)
	gst.sm.On(cmdMove,
		[]sm.State{stateAttack, stateDefense},
		gst.log,
	)

	gst.hub = NewHub()

	return gst
}

//

type playerConn struct {
	gst *gameState

	cards map[string]struct{}

	conn      *websocket.Conn
	hubToConn chan []byte
}

func (self *playerConn) takeCard() {
	if card := self.gst.popCard(); card != "" {
		self.cards[card] = struct{}{}
	}
}

func (self *playerConn) toDesk(card string) {
	delete(self.cards, card)

	// add card to desk
}

func (self *playerConn) write() {
	defer func() {
		err := self.conn.Close()
		if err != nil {
			log.Println(err)
		}
	}()

	for {
		select {
		case m, ok := <-self.hubToConn:
			if !ok {
				return
			}
			//TODO: text or binary?
			err := self.conn.WriteMessage(websocket.TextMessage, m)
			if err != nil {
				log.Println(err)
				return
			}
		}
	}
}

func (self *playerConn) read() {
	defer func() {
		err := self.conn.Close()
		if err != nil {
			log.Println(err)
		}
	}()

	var m PlayerMsg

	for {
		err := self.conn.ReadJSON(&m)
		if err != nil {
			log.Println(err)
			return
		}

		var event *sm.Event
		switch m.Cmd {
		case cmdStart:
			event = &sm.Event{cmdStart, nil}
		case cmdMove:
			event = &sm.Event{cmdMove, cmdArgs{self, m.Card}}
		default:
			log.Printf("unknown command: %v", m.Cmd)
			continue
		}

		err = self.gst.sm.Emit(event)
		if err != nil {
			log.Println(err)
		}
	}
}

//

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin:     func(r *http.Request) bool { return true },
}

func playerHandler(gst *gameState) http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		conn, err := upgrader.Upgrade(w, r, nil)
		if err != nil {
			log.Println(err)
			return
		}

		p := &playerConn{gst, make(map[string]struct{}), conn, make(chan []byte)}

		gst.hub.regChan <- p
		defer func() {
			gst.hub.unregChan <- p
		}()

		go p.write()
		p.read()
	}
}

//

func startDurakSrv() error {
	gst := NewGameState()

	http.HandleFunc("/", playerHandler(gst))

	return http.ListenAndServe(":"+os.Getenv("PORT"), nil)
}

func main() {
	err := startDurakSrv()
	if err != nil {
		log.Fatal(err)
	}
}

func init() {
	rand.Seed(time.Now().UnixNano())

	runtime.GOMAXPROCS(runtime.NumCPU())
}
<|endoftext|>


<commit_before>// @author Robin Verlangen
// Main function of Dispenso ( https://github.com/RobinUS2/dispenso )

package main

// Imports
import (
	"flag"
	"fmt"
	"log"
)

// Constants
const defaultPort int = 8011

// Configuration
var seedNodes string
var serverPort int
var storeState bool

// Signal channels
var shutdown chan bool = make(chan bool)

// Set configuration from flags
func init() {
	flag.StringVar(&seedNodes, "seeds", "", "Seed nodes, comma separated host:port tuples (e.g. 12.34.56.78,23.34.45.56:8080")
	flag.IntVar(&serverPort, "port", defaultPort, fmt.Sprintf("Port to bind on (defaults to %d)", defaultPort))
	flag.BoolVar(&storeState, "store-state", true, "Allow to store cluster state on this node (default=true)")
	flag.Parse()
}

// Main function of dispenso
func main() {
	log.Println("INFO: Starting dispenso")

	// Start server
	var server *Server = NewServer()
	server.Start()

	// Wait for shutdown
	<-shutdown
}
<commit_msg>Raft test<commit_after>// @author Robin Verlangen
// Main function of Dispenso ( https://github.com/RobinUS2/dispenso )

package main

// Imports
import (
	"flag"
	"fmt"
	"log"
	"github.com/goraft/raft"
)

// Constants
const defaultPort int = 8011

// Configuration
var seedNodes string
var serverPort int
var storeState bool

// Signal channels
var shutdown chan bool = make(chan bool)

// Set configuration from flags
func init() {
	flag.Stringvar(&seedNodes, "seeds", "", "Seed nodes, comma separated host:port tuples (e.g. 12.34.56.78,23.34.45.56:8080")
	flag.Intvar(&serverPort, "port", defaultPort, fmt.Sprintf("Port to bind on (defaults to %d)", defaultPort))
	flag.Boolvar(&storeState, "store-state", true, "Allow to store cluster state on this node (default=true)")
	flag.Parse()
}

// Main function of dispenso
func main() {
	log.Println("INFO: Starting dispenso")

	// Start server
	var server *Server = NewServer()
	server.Start()

	// Wait for shutdown
	<-shutdown
}<|endoftext|>


<commit_before>package db

import (
	"errors"
	"fmt"
	"github.com/tolexo/aero/conf"
	"gopkg.in/mgo.v2"
)

// create the mongo connection string
func getMongoConnStr(container string) (string, string) {
	username := conf.String(container+".username", "")
	password := conf.String(container+".password", "")
	host := conf.String(container+".host", "")
	port := conf.String(container+".port", "")
	db := conf.String(container+".db", "")
	replicas := conf.String(container+".replicas", "")
	options := conf.String(container+".options", "")

	if port != "" {
		port = ":" + port
	}
	if replicas != "" {
		replicas = "," + replicas
	}
	if options != "" {
		options = "?" + options
	}
	auth := ""
	if username != "" || password != "" {
		auth = username + ":" + password + "@"
	}

	return fmt.Sprintf("mongodb://%s%s%s%s/%s%s", auth, host, port, replicas, db, options), db
}

// validate the container string
func validateContainer(container ...string) (db string, err error) {
	cLen := len(container)
	if cLen == 0 {
		container = append(container, "database.mongo")
	} else if cLen > 1 {
		err = errors.New("More than one container not supported")
		return
	}
	if conf.Exists(container[0]) == false {
		err = errors.New("mongo configuration missing")
		return
	}
	db = container[0]
	return
}

// create mongo connection
// TODO introduce parameter for connection additional settings like socket timeout
func GetMongoConn(container ...string) (sess *mgo.Session, mdb string, err error) {
	var db string
	if db, err = validateContainer(container...); err == nil {
		conn, mdb := getMongoConnStr(db)
		if mdb == "" {
			err = errors.New("mongo database name missing")
		}
		sess, err = mgo.Dial(conn)
	}
	return
}
<commit_msg>CLA-137 local mdb removed<commit_after>package db

import (
	"errors"
	"fmt"
	"github.com/tolexo/aero/conf"
	"gopkg.in/mgo.v2"
)

// create the mongo connection string
func getMongoConnStr(container string) (string, string) {
	username := conf.String(container+".username", "")
	password := conf.String(container+".password", "")
	host := conf.String(container+".host", "")
	port := conf.String(container+".port", "")
	db := conf.String(container+".db", "")
	replicas := conf.String(container+".replicas", "")
	options := conf.String(container+".options", "")

	if port != "" {
		port = ":" + port
	}
	if replicas != "" {
		replicas = "," + replicas
	}
	if options != "" {
		options = "?" + options
	}
	auth := ""
	if username != "" || password != "" {
		auth = username + ":" + password + "@"
	}

	return fmt.Sprintf("mongodb://%s%s%s%s/%s%s", auth, host, port, replicas, db, options), db
}

// validate the container string
func validateContainer(container ...string) (db string, err error) {
	cLen := len(container)
	if cLen == 0 {
		container = append(container, "database.mongo")
	} else if cLen > 1 {
		err = errors.New("More than one container not supported")
		return
	}
	if conf.Exists(container[0]) == false {
		err = errors.New("mongo configuration missing")
		return
	}
	db = container[0]
	return
}

// create mongo connection
// TODO introduce parameter for connection additional settings like socket timeout
func GetMongoConn(container ...string) (sess *mgo.Session, mdb string, err error) {
	var db string
	if db, err = validateContainer(container...); err == nil {
		var conn string
		conn, mdb = getMongoConnStr(db)
		if mdb == "" {
			err = errors.New("mongo database name missing")
		}
		sess, err = mgo.Dial(conn)
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

// Load RSA private key from pem file
func privateKeyFromFile() (*rsa.PrivateKey, error) {
	var err error
	var pemData []byte
	var block *pem.Block
	var privateKey *rsa.PrivateKey
	if pemData, err = ioutil.ReadFile("./priv.pem"); err != nil {
		err = fmt.Errorf("Error reading pem file: %s", err)
		return nil, err
	}
	if block, _ = pem.Decode(pemData); block == nil || block.Type != "RSA PRIVATE KEY" {
		err = errors.New("No valid PEM data found")
		return nil, err
	}
	if privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes); err != nil {
		err := fmt.Errorf("Private key can't be decoded: %s", err)
		return nil, err
	}
	return privateKey, nil
}

// Generate a new RSA private key and save to pem file
func generatePrivateKey() error {
	var privateKey *rsa.PrivateKey
	var err error
	if privateKey, err = rsa.GenerateKey(rand.Reader, 1024); err != nil {
		return err
	}
	privateKeyFile, err := os.Create("./priv.pem")
	defer privateKeyFile.Close()
	if err != nil {
		return err
	}
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	if err := pem.Encode(privateKeyFile, privateKeyPEM); err != nil {
		return err
	}
	return nil
}

// Get RSA private key from file or generate new one if necessary
func getPrivateKey() (*rsa.PrivateKey, error) {
	if _, err := os.Stat("./priv.pem"); os.IsNotExist(err) {
		err = generatePrivateKey()
		if err != nil {
			return nil, err
		}
	}
	privateKey, err := privateKeyFromFile()
	return privateKey, err
}

// Encrypt data using RSA public key
func encrypt(public_key *rsa.PublicKey, plain_text []byte) ([]byte, error) {
	var label, encrypted []byte
	var err error
	if encrypted, err = rsa.EncryptOAEP(sha256.New(), rand.Reader, public_key, plain_text, label); err != nil {
		return nil, err
	}
	return encrypted, nil
}

// Decrypt data using RSA private key
func decrypt(private_key *rsa.PrivateKey, encrypted []byte) ([]byte, error) {
	var label, decrypted []byte
	var err error
	if decrypted, err = rsa.DecryptOAEP(sha256.New(), rand.Reader, private_key, encrypted, label); err != nil {
		return nil, err
	}
	return decrypted, nil
}

// Sign message using RSA private key
func sign(privateKey *rsa.PrivateKey, message []byte) ([]byte, error) {
	hasher := crypto.SHA256.New()
	hasher.Write(message)
	hashed := hasher.Sum(nil)
	signature, err := rsa.SignPSS(rand.Reader, privateKey, crypto.SHA256, hashed, nil)
	if err != nil {
		return nil, err
	}
	return signature, nil
}

// Encrypt data using AES key
func encryptAES(key, data []byte) ([]byte, error) {
	var err error
	var block cipher.Block
	if block, err = aes.NewCipher(key); err != nil {
		return nil, err
	}
	encryptedData := make([]byte, aes.BlockSize+len(data))
	initVector := encryptedData[:aes.BlockSize]
	if _, err = io.ReadFull(rand.Reader, initVector); err != nil {
		return nil, err
	}
	stream := cipher.NewCFBEncrypter(block, initVector)
	stream.XORKeyStream(encryptedData[aes.BlockSize:], data)
	return encryptedData, nil
}

// Decrypt data using AES key
func decryptAES(key, encryptedData []byte) ([]byte, error) {
	var err error
	var block cipher.Block
	if block, err = aes.NewCipher(key); err != nil {
		return nil, err
	}

	if len(encryptedData) < aes.BlockSize {
		err = errors.New("encryptedData too short")
		return nil, err
	}
	initVector := encryptedData[:aes.BlockSize]
	encryptedData = encryptedData[aes.BlockSize:]
	stream := cipher.NewCFBDecrypter(block, initVector)
	stream.XORKeyStream(encryptedData, encryptedData)
	return encryptedData, nil
}

// Generate new AES (256) key
func generateAESKey() ([]byte, error) {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		return nil, err
	}
	return key, nil
}
<commit_msg>Updated client crypto<commit_after>package main

import (
	"crypto"
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"encoding/pem"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"os"
)

// Load RSA private key from pem file
func privateKeyFromFile() (*rsa.PrivateKey, error) {
	var err error
	var pemData []byte
	var block *pem.Block
	var privateKey *rsa.PrivateKey
	if pemData, err = ioutil.ReadFile("./priv.pem"); err != nil {
		err = fmt.Errorf("Error reading pem file: %s", err)
		return nil, err
	}
	if block, _ = pem.Decode(pemData); block == nil || block.Type != "RSA PRIVATE KEY" {
		err = errors.New("No valid PEM data found")
		return nil, err
	}
	if privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes); err != nil {
		err := fmt.Errorf("Private key can't be decoded: %s", err)
		return nil, err
	}
	return privateKey, nil
}

// Generate a new RSA private key and save to pem file
func generatePrivateKey() error {
	var privateKey *rsa.PrivateKey
	var err error
	if privateKey, err = rsa.GenerateKey(rand.Reader, 1024); err != nil {
		return err
	}
	privateKeyFile, err := os.Create("./priv.pem")
	defer privateKeyFile.Close()
	if err != nil {
		return err
	}
	privateKeyPEM := &pem.Block{Type: "RSA PRIVATE KEY", Bytes: x509.MarshalPKCS1PrivateKey(privateKey)}
	if err := pem.Encode(privateKeyFile, privateKeyPEM); err != nil {
		return err
	}
	return nil
}

// Get RSA private key from file or generate new one if necessary
func getPrivateKey() (*rsa.PrivateKey, error) {
	if _, err := os.Stat("./priv.pem"); os.IsNotExist(err) {
		err = generatePrivateKey()
		if err != nil {
			return nil, err
		}
	}
	privateKey, err := privateKeyFromFile()
	return privateKey, err
}

// Encrypt data using RSA public key
func encrypt(public_key *rsa.PublicKey, plain_text []byte) ([]byte, error) {
	var label, encrypted []byte
	var err error
	if encrypted, err = rsa.EncryptOAEP(sha256.New(), rand.Reader, public_key, plain_text, label); err != nil {
		return nil, err
	}
	return encrypted, nil
}

// Decrypt data using RSA private key
func decrypt(private_key *rsa.PrivateKey, encrypted []byte) ([]byte, error) {
	var label, decrypted []byte
	var err error
	if decrypted, err = rsa.DecryptOAEP(sha256.New(), rand.Reader, private_key, encrypted, label); err != nil {
		return nil, err
	}
	return decrypted, nil
}

// Sign message using RSA private key
func sign(privateKey *rsa.PrivateKey, message []byte) ([]byte, error) {
	hasher := crypto.SHA256.New()
	hasher.Write(message)
	hashed := hasher.Sum(nil)
    var opts rsa.PSSOptions
	signature, err := rsa.SignPSS(rand.Reader, privateKey, crypto.SHA256, hashed, &opts)
	if err != nil {
		return nil, err
	}
	return signature, nil
}

// Encrypt data using AES key
func encryptAES(key, data []byte) ([]byte, error) {
	var err error
	var block cipher.Block
	if block, err = aes.NewCipher(key); err != nil {
		return nil, err
	}
	encryptedData := make([]byte, aes.BlockSize+len(data))
	initVector := encryptedData[:aes.BlockSize]
	if _, err = io.ReadFull(rand.Reader, initVector); err != nil {
		return nil, err
	}
	stream := cipher.NewCFBEncrypter(block, initVector)
	stream.XORKeyStream(encryptedData[aes.BlockSize:], data)
	return encryptedData, nil
}

// Decrypt data using AES key
func decryptAES(key, encryptedData []byte) ([]byte, error) {
	var err error
	var block cipher.Block
	if block, err = aes.NewCipher(key); err != nil {
		return nil, err
	}

	if len(encryptedData) < aes.BlockSize {
		err = errors.New("encryptedData too short")
		return nil, err
	}
	initVector := encryptedData[:aes.BlockSize]
	encryptedData = encryptedData[aes.BlockSize:]
	stream := cipher.NewCFBDecrypter(block, initVector)
	stream.XORKeyStream(encryptedData, encryptedData)
	return encryptedData, nil
}

// Generate new AES (256) key
func generateAESKey() ([]byte, error) {
	key := make([]byte, 32)
	if _, err := rand.Read(key); err != nil {
		return nil, err
	}
	return key, nil
}
<|endoftext|>


<commit_before>package mpawselasticache

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	mp "github.com/mackerelio/go-mackerel-plugin"
)

var metricsdefMemcached = []string{
	"CPUUtilization", "SwapUsage", "FreeableMemory", "NetworkBytesIn", "NetworkBytesOut",
	"BytesUsedForCacheItems", "BytesReadIntoMemcached", "BytesWrittenOutFromMemcached", "CasBadval", "CasHits",
	"CasMisses", "CmdFlush", "CmdGet", "CmdSet", "CurrConnections",
	"CurrItems", "DecrHits", "DecrMisses", "DeleteHits", "DeleteMisses",
	"Evictions", "GetHits", "GetMisses", "IncrHits", "IncrMisses",
	"Reclaimed", "BytesUsedForHash", "CmdConfigGet", "CmdConfigSet", "CmdTouch",
	"CurrConfig", "EvictedUnfetched", "ExpiredUnfetched", "SlabsMoved", "TouchHits",
	"TouchMisses", "NewConnections", "NewItems", "UnusedMemory",
}

var metricsdefRedis = []string{
	"CPUUtilization", "SwapUsage", "FreeableMemory", "NetworkBytesIn", "NetworkBytesOut",
	"CurrConnections", "Evictions", "Reclaimed", "NewConnections", "BytesUsedForCache",
	"CacheHits", "CacheMisses", "ReplicationLag", "GetTypeCmds", "SetTypeCmds",
	"KeyBasedCmds", "StringBasedCmds", "HashBasedCmds", "ListBasedCmds", "SetBasedCmds",
	"SortedSetBasedCmds", "CurrItems",
}

var graphdefMemcached = map[string]mp.Graphs{
	"ecache.CPUUtilization": {
		Label: "ECache CPU Utilization",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "CPUUtilization", Label: "CPUUtilization"},
		},
	},
	"ecache.SwapUsage": {
		Label: "ECache Swap Usage",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "SwapUsage", Label: "SwapUsage"},
		},
	},
	"ecache.FreeableMemory": {
		Label: "ECache Freeable Memory",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "FreeableMemory", Label: "FreeableMemory"},
		},
	},
	"ecache.NetworkTraffic": {
		Label: "ECache Network Traffic",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "NetworkBytesIn", Label: "NetworkBytesIn"},
			{Name: "NetworkBytesOut", Label: "NetworkBytesOut"},
		},
	},
	"ecache.Command": {
		Label: "ECache Command",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CmdGet", Label: "CmdGet"},
			{Name: "CmdSet", Label: "CmdSet"},
			{Name: "CmdFlush", Label: "CmdFlush"},
			{Name: "CmdTouch", Label: "CmdTouch"},
			{Name: "CmdConfigGet", Label: "CmdConfigGet"},
			{Name: "CmdConfigSet", Label: "CmdConfigSet"},
		},
	},
	"ecache.CacheHitAndMiss": {
		Label: "ECache Hits/Misses",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "GetHits", Label: "GetHits"},
			{Name: "GetMisses", Label: "GetMisses"},
			{Name: "DeleteHits", Label: "DeleteHits"},
			{Name: "DeleteMisses", Label: "DeleteMisses"},
			{Name: "IncrHits", Label: "IncrHits"},
			{Name: "IncrMisses", Label: "IncrMisses"},
			{Name: "DecrHits", Label: "DecrHits"},
			{Name: "DecrMisses", Label: "DecrMisses"},
			{Name: "CasBadval", Label: "CasBadval"},
			{Name: "CasHits", Label: "CasHits"},
			{Name: "CasMisses", Label: "CasMisses"},
			{Name: "TouchHits", Label: "TouchHits"},
			{Name: "TouchMisses", Label: "TouchMisses"},
		},
	},
	"ecache.Evictions": {
		Label: "ECache Evictions",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "Evictions", Label: "Evictions"},
		},
	},
	"ecache.Unfetched": {
		Label: "ECache Unfetched",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "EvictedUnfetched", Label: "EvictedUnfetched"},
			{Name: "ExpiredUnfetched", Label: "ExpiredUnfetched"},
		},
	},
	"ecache.Bytes": {
		Label: "ECache Traffics",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "BytesReadIntoMemcached", Label: "BytesReadIntoMemcached"},
			{Name: "BytesWrittenOutFromMemcached", Label: "BytesWrittenOutFromMemcached"},
		},
	},
	"ecache.Connections": {
		Label: "ECache Connections",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrConnections", Label: "CurrConnections"},
			{Name: "NewConnections", Label: "NewConnections"},
		},
	},
	"ecache.Items": {
		Label: "ECache Items",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrItems", Label: "CurrItems"},
			{Name: "NewItems", Label: "NewItems"},
			{Name: "Reclaimed", Label: "Reclaimed"},
			{Name: "CurrConfig", Label: "CurrConfig"},
			{Name: "SlabsMoved", Label: "SlabsMoved"},
		},
	},
	"ecache.MemoryUsage": {
		Label: "ECache Memory Usage",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "UnusedMemory", Label: "UnusedMemory"},
			{Name: "BytesUsedForHash", Label: "BytesUsedForHash"},
			{Name: "BytesUsedForCacheItems", Label: "BytesUsedForCacheItems"},
		},
	},
}

var graphdefRedis = map[string]mp.Graphs{
	"ecache.CPUUtilization": {
		Label: "ECache CPU Utilization",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "CPUUtilization", Label: "CPUUtilization"},
		},
	},
	"ecache.SwapUsage": {
		Label: "ECache Swap Usage",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "SwapUsage", Label: "SwapUsage"},
		},
	},
	"ecache.FreeableMemory": {
		Label: "ECache Freeable Memory",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "FreeableMemory", Label: "FreeableMemory"},
		},
	},
	"ecache.NetworkTraffic": {
		Label: "ECache Network Traffic",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "NetworkBytesIn", Label: "NetworkBytesIn"},
			{Name: "NetworkBytesOut", Label: "NetworkBytesOut"},
		},
	},
	"ecache.Command": {
		Label: "ECache Command",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "GetTypeCmds", Label: "GetTypeCmds"},
			{Name: "SetTypeCmds", Label: "SetTypeCmds"},
			{Name: "KeyBasedCmds", Label: "KeyBasedCmds"},
			{Name: "StringBasedCmds", Label: "StringBasedCmds"},
			{Name: "HashBasedCmds", Label: "HashBasedCmds"},
			{Name: "ListBasedCmds", Label: "ListBasedCmds"},
			{Name: "SetBasedCmds", Label: "SetBasedCmds"},
			{Name: "SortedSetBasedCmds", Label: "SortedSetBasedCmds"},
		},
	},
	"ecache.CacheHitAndMiss": {
		Label: "ECache Hits/Misses",
		Unit:  "float",
		Metrics: []mp.Metrics{
			{Name: "CacheHits", Label: "CacheHits"},
			{Name: "CacheMisses", Label: "CacheMisses"},
		},
	},
	"ecache.Evictions": {
		Label: "ECache Evictions",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "Evictions", Label: "Evictions"},
		},
	},
	"ecache.Bytes": {
		Label: "ECache Traffics",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "BytesUsedForCache", Label: "BytesUsedForCache"},
		},
	},
	"ecache.Connections": {
		Label: "ECache Connections",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrConnections", Label: "CurrConnections"},
			{Name: "NewConnections", Label: "NewConnections"},
		},
	},
	"ecache.Items": {
		Label: "ECache Items",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrItems", Label: "CurrItems"},
			{Name: "Reclaimed", Label: "Reclaimed"},
		},
	},
}

// ECachePlugin mackerel plugin for elasticache
type ECachePlugin struct {
	Region          string
	AccessKeyID     string
	SecretAccessKey string
	CacheClusterID  string
	CacheNodeID     string
	ElastiCacheType string
	CacheMetrics    []string
}

func getLastPoint(cloudWatch *cloudwatch.CloudWatch, dimensions []*cloudwatch.Dimension, metricName string) (float64, error) {
	now := time.Now()

	response, err := cloudWatch.GetMetricStatistics(&cloudwatch.GetMetricStatisticsInput{
		Dimensions: dimensions,
		StartTime:  aws.Time(now.Add(time.Duration(180) * time.Second * -1)), // 3 mins (to fetch at least 1 data-point)
		EndTime:    aws.Time(now),
		MetricName: aws.String(metricName),
		Period:     aws.Int64(60),
		Statistics: []*string{aws.String("Average")},
		Namespace:  aws.String("AWS/ElastiCache"),
	})
	if err != nil {
		return 0, err
	}

	datapoints := response.Datapoints
	if len(datapoints) == 0 {
		return 0, errors.New("fetched no datapoints")
	}

	latest := new(time.Time)
	var latestVal float64
	for _, dp := range datapoints {
		if dp.Timestamp.Before(*latest) {
			continue
		}

		latest = dp.Timestamp
		latestVal = *dp.Average
	}

	return latestVal, nil
}

// FetchMetrics fetch elasticache values
func (p ECachePlugin) FetchMetrics() (map[string]float64, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}

	config := aws.NewConfig()
	if p.AccessKeyID != "" && p.SecretAccessKey != "" {
		config = config.WithCredentials(credentials.NewStaticCredentials(p.AccessKeyID, p.SecretAccessKey, ""))
	}
	if p.Region != "" {
		config = config.WithRegion(p.Region)
	}

	cloudWatch := cloudwatch.New(sess, config)

	stat := make(map[string]float64)

	perInstances := []*cloudwatch.Dimension{
		{
			Name:  aws.String("CacheClusterId"),
			Value: aws.String(p.CacheClusterID),
		},
		{
			Name:  aws.String("CacheNodeId"),
			Value: aws.String(p.CacheNodeID),
		},
	}

	for _, met := range p.CacheMetrics {
		v, err := getLastPoint(cloudWatch, perInstances, met)
		if err == nil {
			stat[met] = v
		} else {
			log.Printf("%s: %s", met, err)
		}
	}

	return stat, nil
}

// GraphDefinition graph definition
func (p ECachePlugin) GraphDefinition() map[string]mp.Graphs {
	switch p.ElastiCacheType {
	case "memcached":
		return graphdefMemcached
	case "redis":
		return graphdefRedis
	default:
		log.Printf("elasticache-type is 'memcached' or 'redis'.")
		return nil
	}
}

// Do the plugin
func Do() {
	optRegion := flag.String("region", "", "AWS Region")
	optAccessKeyID := flag.String("access-key-id", "", "AWS Access Key ID")
	optSecretAccessKey := flag.String("secret-access-key", "", "AWS Secret Access Key")
	optCacheClusterID := flag.String("cache-cluster-id", "", "Cache Cluster Id")
	optCacheNodeID := flag.String("cache-node-id", "0001", "Cache Node Id")
	optElastiCacheType := flag.String("elasticache-type", "", "ElastiCache type")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var ecache ECachePlugin

	if *optRegion == "" {
		ec2metadata := ec2metadata.New(session.New())
		if ec2metadata.Available() {
			ecache.Region, _ = ec2metadata.Region()
		}
	} else {
		ecache.Region = *optRegion
	}

	ecache.Region = *optRegion
	ecache.AccessKeyID = *optAccessKeyID
	ecache.SecretAccessKey = *optSecretAccessKey
	ecache.CacheClusterID = *optCacheClusterID
	ecache.CacheNodeID = *optCacheNodeID
	ecache.ElastiCacheType = *optElastiCacheType
	switch ecache.ElastiCacheType {
	case "memcached":
		ecache.CacheMetrics = metricsdefMemcached
	case "redis":
		ecache.CacheMetrics = metricsdefRedis
	default:
		log.Printf("elasticache-type is 'memcached' or 'redis'.")
		os.Exit(1)
	}

	helper := mp.NewMackerelPlugin(ecache)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.Tempfile = fmt.Sprintf("/tmp/mackerel-plugin-aws-elasticache-%s-%s", *optCacheClusterID, *optCacheNodeID)
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<commit_msg>[aws-elasticache] support MACKEREL_PLUGIN_WORKDIR by SetTempfileByBasename<commit_after>package mpawselasticache

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/credentials"
	"github.com/aws/aws-sdk-go/aws/ec2metadata"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/cloudwatch"
	mp "github.com/mackerelio/go-mackerel-plugin"
)

var metricsdefMemcached = []string{
	"CPUUtilization", "SwapUsage", "FreeableMemory", "NetworkBytesIn", "NetworkBytesOut",
	"BytesUsedForCacheItems", "BytesReadIntoMemcached", "BytesWrittenOutFromMemcached", "CasBadval", "CasHits",
	"CasMisses", "CmdFlush", "CmdGet", "CmdSet", "CurrConnections",
	"CurrItems", "DecrHits", "DecrMisses", "DeleteHits", "DeleteMisses",
	"Evictions", "GetHits", "GetMisses", "IncrHits", "IncrMisses",
	"Reclaimed", "BytesUsedForHash", "CmdConfigGet", "CmdConfigSet", "CmdTouch",
	"CurrConfig", "EvictedUnfetched", "ExpiredUnfetched", "SlabsMoved", "TouchHits",
	"TouchMisses", "NewConnections", "NewItems", "UnusedMemory",
}

var metricsdefRedis = []string{
	"CPUUtilization", "SwapUsage", "FreeableMemory", "NetworkBytesIn", "NetworkBytesOut",
	"CurrConnections", "Evictions", "Reclaimed", "NewConnections", "BytesUsedForCache",
	"CacheHits", "CacheMisses", "ReplicationLag", "GetTypeCmds", "SetTypeCmds",
	"KeyBasedCmds", "StringBasedCmds", "HashBasedCmds", "ListBasedCmds", "SetBasedCmds",
	"SortedSetBasedCmds", "CurrItems",
}

var graphdefMemcached = map[string]mp.Graphs{
	"ecache.CPUUtilization": {
		Label: "ECache CPU Utilization",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "CPUUtilization", Label: "CPUUtilization"},
		},
	},
	"ecache.SwapUsage": {
		Label: "ECache Swap Usage",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "SwapUsage", Label: "SwapUsage"},
		},
	},
	"ecache.FreeableMemory": {
		Label: "ECache Freeable Memory",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "FreeableMemory", Label: "FreeableMemory"},
		},
	},
	"ecache.NetworkTraffic": {
		Label: "ECache Network Traffic",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "NetworkBytesIn", Label: "NetworkBytesIn"},
			{Name: "NetworkBytesOut", Label: "NetworkBytesOut"},
		},
	},
	"ecache.Command": {
		Label: "ECache Command",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CmdGet", Label: "CmdGet"},
			{Name: "CmdSet", Label: "CmdSet"},
			{Name: "CmdFlush", Label: "CmdFlush"},
			{Name: "CmdTouch", Label: "CmdTouch"},
			{Name: "CmdConfigGet", Label: "CmdConfigGet"},
			{Name: "CmdConfigSet", Label: "CmdConfigSet"},
		},
	},
	"ecache.CacheHitAndMiss": {
		Label: "ECache Hits/Misses",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "GetHits", Label: "GetHits"},
			{Name: "GetMisses", Label: "GetMisses"},
			{Name: "DeleteHits", Label: "DeleteHits"},
			{Name: "DeleteMisses", Label: "DeleteMisses"},
			{Name: "IncrHits", Label: "IncrHits"},
			{Name: "IncrMisses", Label: "IncrMisses"},
			{Name: "DecrHits", Label: "DecrHits"},
			{Name: "DecrMisses", Label: "DecrMisses"},
			{Name: "CasBadval", Label: "CasBadval"},
			{Name: "CasHits", Label: "CasHits"},
			{Name: "CasMisses", Label: "CasMisses"},
			{Name: "TouchHits", Label: "TouchHits"},
			{Name: "TouchMisses", Label: "TouchMisses"},
		},
	},
	"ecache.Evictions": {
		Label: "ECache Evictions",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "Evictions", Label: "Evictions"},
		},
	},
	"ecache.Unfetched": {
		Label: "ECache Unfetched",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "EvictedUnfetched", Label: "EvictedUnfetched"},
			{Name: "ExpiredUnfetched", Label: "ExpiredUnfetched"},
		},
	},
	"ecache.Bytes": {
		Label: "ECache Traffics",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "BytesReadIntoMemcached", Label: "BytesReadIntoMemcached"},
			{Name: "BytesWrittenOutFromMemcached", Label: "BytesWrittenOutFromMemcached"},
		},
	},
	"ecache.Connections": {
		Label: "ECache Connections",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrConnections", Label: "CurrConnections"},
			{Name: "NewConnections", Label: "NewConnections"},
		},
	},
	"ecache.Items": {
		Label: "ECache Items",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrItems", Label: "CurrItems"},
			{Name: "NewItems", Label: "NewItems"},
			{Name: "Reclaimed", Label: "Reclaimed"},
			{Name: "CurrConfig", Label: "CurrConfig"},
			{Name: "SlabsMoved", Label: "SlabsMoved"},
		},
	},
	"ecache.MemoryUsage": {
		Label: "ECache Memory Usage",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "UnusedMemory", Label: "UnusedMemory"},
			{Name: "BytesUsedForHash", Label: "BytesUsedForHash"},
			{Name: "BytesUsedForCacheItems", Label: "BytesUsedForCacheItems"},
		},
	},
}

var graphdefRedis = map[string]mp.Graphs{
	"ecache.CPUUtilization": {
		Label: "ECache CPU Utilization",
		Unit:  "percentage",
		Metrics: []mp.Metrics{
			{Name: "CPUUtilization", Label: "CPUUtilization"},
		},
	},
	"ecache.SwapUsage": {
		Label: "ECache Swap Usage",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "SwapUsage", Label: "SwapUsage"},
		},
	},
	"ecache.FreeableMemory": {
		Label: "ECache Freeable Memory",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "FreeableMemory", Label: "FreeableMemory"},
		},
	},
	"ecache.NetworkTraffic": {
		Label: "ECache Network Traffic",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "NetworkBytesIn", Label: "NetworkBytesIn"},
			{Name: "NetworkBytesOut", Label: "NetworkBytesOut"},
		},
	},
	"ecache.Command": {
		Label: "ECache Command",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "GetTypeCmds", Label: "GetTypeCmds"},
			{Name: "SetTypeCmds", Label: "SetTypeCmds"},
			{Name: "KeyBasedCmds", Label: "KeyBasedCmds"},
			{Name: "StringBasedCmds", Label: "StringBasedCmds"},
			{Name: "HashBasedCmds", Label: "HashBasedCmds"},
			{Name: "ListBasedCmds", Label: "ListBasedCmds"},
			{Name: "SetBasedCmds", Label: "SetBasedCmds"},
			{Name: "SortedSetBasedCmds", Label: "SortedSetBasedCmds"},
		},
	},
	"ecache.CacheHitAndMiss": {
		Label: "ECache Hits/Misses",
		Unit:  "float",
		Metrics: []mp.Metrics{
			{Name: "CacheHits", Label: "CacheHits"},
			{Name: "CacheMisses", Label: "CacheMisses"},
		},
	},
	"ecache.Evictions": {
		Label: "ECache Evictions",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "Evictions", Label: "Evictions"},
		},
	},
	"ecache.Bytes": {
		Label: "ECache Traffics",
		Unit:  "bytes",
		Metrics: []mp.Metrics{
			{Name: "BytesUsedForCache", Label: "BytesUsedForCache"},
		},
	},
	"ecache.Connections": {
		Label: "ECache Connections",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrConnections", Label: "CurrConnections"},
			{Name: "NewConnections", Label: "NewConnections"},
		},
	},
	"ecache.Items": {
		Label: "ECache Items",
		Unit:  "integer",
		Metrics: []mp.Metrics{
			{Name: "CurrItems", Label: "CurrItems"},
			{Name: "Reclaimed", Label: "Reclaimed"},
		},
	},
}

// ECachePlugin mackerel plugin for elasticache
type ECachePlugin struct {
	Region          string
	AccessKeyID     string
	SecretAccessKey string
	CacheClusterID  string
	CacheNodeID     string
	ElastiCacheType string
	CacheMetrics    []string
}

func getLastPoint(cloudWatch *cloudwatch.CloudWatch, dimensions []*cloudwatch.Dimension, metricName string) (float64, error) {
	now := time.Now()

	response, err := cloudWatch.GetMetricStatistics(&cloudwatch.GetMetricStatisticsInput{
		Dimensions: dimensions,
		StartTime:  aws.Time(now.Add(time.Duration(180) * time.Second * -1)), // 3 mins (to fetch at least 1 data-point)
		EndTime:    aws.Time(now),
		MetricName: aws.String(metricName),
		Period:     aws.Int64(60),
		Statistics: []*string{aws.String("Average")},
		Namespace:  aws.String("AWS/ElastiCache"),
	})
	if err != nil {
		return 0, err
	}

	datapoints := response.Datapoints
	if len(datapoints) == 0 {
		return 0, errors.New("fetched no datapoints")
	}

	latest := new(time.Time)
	var latestVal float64
	for _, dp := range datapoints {
		if dp.Timestamp.Before(*latest) {
			continue
		}

		latest = dp.Timestamp
		latestVal = *dp.Average
	}

	return latestVal, nil
}

// FetchMetrics fetch elasticache values
func (p ECachePlugin) FetchMetrics() (map[string]float64, error) {
	sess, err := session.NewSession()
	if err != nil {
		return nil, err
	}

	config := aws.NewConfig()
	if p.AccessKeyID != "" && p.SecretAccessKey != "" {
		config = config.WithCredentials(credentials.NewStaticCredentials(p.AccessKeyID, p.SecretAccessKey, ""))
	}
	if p.Region != "" {
		config = config.WithRegion(p.Region)
	}

	cloudWatch := cloudwatch.New(sess, config)

	stat := make(map[string]float64)

	perInstances := []*cloudwatch.Dimension{
		{
			Name:  aws.String("CacheClusterId"),
			Value: aws.String(p.CacheClusterID),
		},
		{
			Name:  aws.String("CacheNodeId"),
			Value: aws.String(p.CacheNodeID),
		},
	}

	for _, met := range p.CacheMetrics {
		v, err := getLastPoint(cloudWatch, perInstances, met)
		if err == nil {
			stat[met] = v
		} else {
			log.Printf("%s: %s", met, err)
		}
	}

	return stat, nil
}

// GraphDefinition graph definition
func (p ECachePlugin) GraphDefinition() map[string]mp.Graphs {
	switch p.ElastiCacheType {
	case "memcached":
		return graphdefMemcached
	case "redis":
		return graphdefRedis
	default:
		log.Printf("elasticache-type is 'memcached' or 'redis'.")
		return nil
	}
}

// Do the plugin
func Do() {
	optRegion := flag.String("region", "", "AWS Region")
	optAccessKeyID := flag.String("access-key-id", "", "AWS Access Key ID")
	optSecretAccessKey := flag.String("secret-access-key", "", "AWS Secret Access Key")
	optCacheClusterID := flag.String("cache-cluster-id", "", "Cache Cluster Id")
	optCacheNodeID := flag.String("cache-node-id", "0001", "Cache Node Id")
	optElastiCacheType := flag.String("elasticache-type", "", "ElastiCache type")
	optTempfile := flag.String("tempfile", "", "Temp file name")
	flag.Parse()

	var ecache ECachePlugin

	if *optRegion == "" {
		ec2metadata := ec2metadata.New(session.New())
		if ec2metadata.Available() {
			ecache.Region, _ = ec2metadata.Region()
		}
	} else {
		ecache.Region = *optRegion
	}

	ecache.Region = *optRegion
	ecache.AccessKeyID = *optAccessKeyID
	ecache.SecretAccessKey = *optSecretAccessKey
	ecache.CacheClusterID = *optCacheClusterID
	ecache.CacheNodeID = *optCacheNodeID
	ecache.ElastiCacheType = *optElastiCacheType
	switch ecache.ElastiCacheType {
	case "memcached":
		ecache.CacheMetrics = metricsdefMemcached
	case "redis":
		ecache.CacheMetrics = metricsdefRedis
	default:
		log.Printf("elasticache-type is 'memcached' or 'redis'.")
		os.Exit(1)
	}

	helper := mp.NewMackerelPlugin(ecache)
	if *optTempfile != "" {
		helper.Tempfile = *optTempfile
	} else {
		helper.SetTempfileByBasename(fmt.Sprintf("mackerel-plugin-aws-elasticache-%s-%s", *optCacheClusterID, *optCacheNodeID))
	}

	if os.Getenv("MACKEREL_AGENT_PLUGIN_META") != "" {
		helper.OutputDefinitions()
	} else {
		helper.OutputValues()
	}
}
<|endoftext|>


<commit_before>// Copyright 2016 Mender Software AS
//
//    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 controller

import (
	"errors"
	"github.com/mendersoftware/deployments/resources/deployments"
)

// Errors
var (
	ErrModelMissingInput       = errors.New("Missing input deplyoment data")
	ErrModelInvalidDeviceID    = errors.New("Invalid device ID")
	ErrModelDeploymentNotFound = errors.New("Deployment not found")
	ErrModelInternal           = errors.New("Internal error")
)

// Domain model for deployment
type DeploymentsModel interface {
	CreateDeployment(constructor *deployments.DeploymentConstructor) (string, error)
	GetDeployment(deploymentID string) (*deployments.Deployment, error)
	GetDeploymentStats(deploymentID string) (deployments.Stats, error)
	GetDeploymentForDevice(deviceID string) (*deployments.DeploymentInstructions, error)
	HasDeploymentForDevice(deploymentID string, deviceID string) (bool, error)
	UpdateDeviceDeploymentStatus(deploymentID string, deviceID string, status string) error
	GetDeviceStatusesForDeployment(deploymentID string) ([]deployments.DeviceDeployment, error)
	LookupDeployment(query deployments.Query) ([]*deployments.Deployment, error)
	SaveDeviceDeploymentLog(deviceID string, deploymentID string, log *deployments.DeploymentLog) error
}
<commit_msg>deployments/controller: define error for invalid deployment log<commit_after>// Copyright 2016 Mender Software AS
//
//    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 controller

import (
	"errors"
	"github.com/mendersoftware/deployments/resources/deployments"
)

// Errors
var (
	ErrModelMissingInput       = errors.New("Missing input deplyoment data")
	ErrModelInvalidDeviceID    = errors.New("Invalid device ID")
	ErrModelDeploymentNotFound = errors.New("Deployment not found")
	ErrModelInternal           = errors.New("Internal error")
	ErrStorageInvalidLog       = errors.New("Invalid deployment log")
)

// Domain model for deployment
type DeploymentsModel interface {
	CreateDeployment(constructor *deployments.DeploymentConstructor) (string, error)
	GetDeployment(deploymentID string) (*deployments.Deployment, error)
	GetDeploymentStats(deploymentID string) (deployments.Stats, error)
	GetDeploymentForDevice(deviceID string) (*deployments.DeploymentInstructions, error)
	HasDeploymentForDevice(deploymentID string, deviceID string) (bool, error)
	UpdateDeviceDeploymentStatus(deploymentID string, deviceID string, status string) error
	GetDeviceStatusesForDeployment(deploymentID string) ([]deployments.DeviceDeployment, error)
	LookupDeployment(query deployments.Query) ([]*deployments.Deployment, error)
	SaveDeviceDeploymentLog(deviceID string, deploymentID string, log *deployments.DeploymentLog) error
}
<|endoftext|>


<commit_before>// +build !v2

package engine

import (
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/funkygao/dbus"
	"github.com/funkygao/dbus/pkg/cluster"
	czk "github.com/funkygao/dbus/pkg/cluster/zk"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/telemetry"
	"github.com/funkygao/gafka/telemetry/influxdb"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/golib/observer"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
	"github.com/gorilla/mux"
)

var (
	_ PluginHelper = &Engine{}
)

// Engine is the pipeline engine of the data bus system which manages the core loop.
type Engine struct {
	sync.RWMutex
	*conf.Conf

	zkSvr       string
	participant cluster.Participant
	controller  cluster.Controller
	epoch       int // cache of latest cluster leader epoch

	// API Server
	apiListener net.Listener
	apiServer   *http.Server
	apiRouter   *mux.Router

	// RPC Server
	rpcListener net.Listener
	rpcServer   *http.Server
	rpcRouter   *mux.Router

	// input plugin resources map
	irm   map[string][]cluster.Resource
	irmMu sync.Mutex

	// dataflow router
	router *Router

	InputRunners  map[string]*iRunner
	inputWrappers map[string]*pluginWrapper

	FilterRunners  map[string]FilterRunner
	filterWrappers map[string]*pluginWrapper

	OutputRunners  map[string]OutputRunner
	outputWrappers map[string]*pluginWrapper

	inputRecycleChans map[string]chan *Packet
	filterRecycleChan chan *Packet

	hostname      string
	pid           int
	stopper       chan struct{}
	shutdown      chan struct{}
	pluginPanicCh chan error
}

// New creates an engine.
func New(globals *GlobalConfig) *Engine {
	if globals == nil {
		globals = DefaultGlobals()
	}
	Globals = func() *GlobalConfig {
		return globals
	}

	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}

	// make the participant
	localIP, err := ctx.LocalIP()
	if err != nil {
		panic(err)
	}

	return &Engine{
		pid:           os.Getpid(),
		hostname:      hostname,
		stopper:       make(chan struct{}),
		shutdown:      make(chan struct{}),
		pluginPanicCh: make(chan error),

		router: newRouter(),

		InputRunners:   make(map[string]*iRunner),
		inputWrappers:  make(map[string]*pluginWrapper),
		FilterRunners:  make(map[string]FilterRunner),
		filterWrappers: make(map[string]*pluginWrapper),
		OutputRunners:  make(map[string]OutputRunner),
		outputWrappers: make(map[string]*pluginWrapper),

		inputRecycleChans: make(map[string]chan *Packet),
		filterRecycleChan: make(chan *Packet, globals.FilterRecyclePoolSize),

		participant: cluster.Participant{
			Endpoint: fmt.Sprintf("%s:%d", localIP.String(), globals.RPCPort),
			Weight:   runtime.NumCPU() * 100,
			State:    cluster.StateOnline,
			Revision: dbus.Revision,
			APIPort:  globals.APIPort,
		},
	}
}

func (e *Engine) stopInputRunner(name string) {
	e.Lock()
	e.InputRunners[name] = nil
	e.Unlock()
}

// ClonePacket is used for plugin Filter to generate new Packet: copy on write.
// The generated Packet will use dedicated filter recycle chan.
func (e *Engine) ClonePacket(p *Packet) *Packet {
	pack := <-e.filterRecycleChan
	pack.Reset()
	p.copyTo(pack)
	return pack
}

// ClusterManager returns the cluster manager.
// If cluster is disabled, returns nil.
func (e *Engine) ClusterManager() cluster.Manager {
	if Globals().ClusterEnabled {
		return e.controller.(cluster.Manager)
	}

	return nil
}

// SubmitDAG submits a DAG configuration to the engine.
func (e *Engine) SubmitDAG(cf *conf.Conf) *Engine {
	return e.loadConfig(cf)
}

func (e *Engine) loadConfig(cf *conf.Conf) *Engine {
	e.Conf = cf
	Globals().Conf = cf

	// 'plugins' section
	var pluginNames = make(map[string]struct{})
	for i := 0; i < len(e.List("plugins", nil)); i++ {
		section, err := e.Section(fmt.Sprintf("plugins[%d]", i))
		if err != nil {
			panic(err)
		}

		name := e.loadPluginSection(section)
		if _, duplicated := pluginNames[name]; duplicated {
			// router.metrics will be bad with dup name
			panic("duplicated plugin name: " + name)
		}
		pluginNames[name] = struct{}{}
	}

	// influxdb related section
	if c, err := influxdb.NewConfig(cf.String("influx_addr", ""),
		cf.String("influx_db", "dbus"), "", "",
		cf.Duration("influx_tick", time.Minute)); err == nil {
		telemetry.Default = influxdb.New(metrics.DefaultRegistry, c)
	}

	return e
}

// LoadFrom load the configuration by location.
// The location can be empty: use default zk zone /dbus/conf.
// If config is stored on file, the loc arg is file path.
// If config is stored on zookeeper, the loc arg is like localhost:2181/foo/bar.
func (e *Engine) LoadFrom(loc string) *Engine {
	if len(loc) == 0 {
		// if no location provided, use the default zk
		loc = fmt.Sprintf("%s%s", ctx.ZoneZkAddrs(ctx.DefaultZone()), Globals().ZrootConf)
	}

	zkSvr, realPath := parseConfigPath(loc)
	var (
		cf  *conf.Conf
		err error
	)
	if len(zkSvr) == 0 {
		// from file system
		cf, err = conf.Load(realPath)
	} else {
		// from zookeeper
		e.zkSvr = zkSvr
		cf, err = conf.Load(realPath, conf.WithZkSvr(zkSvr))
		if err != nil {
			err = fmt.Errorf("%s %v", loc, err)
		}
	}
	if err != nil {
		panic(err)
	}

	return e.loadConfig(cf)
}

func (e *Engine) loadPluginSection(section *conf.Conf) string {
	pluginCommons := new(pluginCommons)
	pluginCommons.loadConfig(section)

	wrapper := &pluginWrapper{
		name:          pluginCommons.name,
		configCreator: func() *conf.Conf { return section },
	}
	var ok bool
	if wrapper.pluginCreator, ok = availablePlugins[pluginCommons.class]; !ok {
		panic("unknown plugin type: " + pluginCommons.class)
	}

	pluginType := pluginTypeRegex.FindStringSubmatch(pluginCommons.class)
	if len(pluginType) < 2 {
		panic("invalid plugin type: " + pluginCommons.class)
	}

	plugin := wrapper.Create()
	pluginCategory := pluginType[1]
	if pluginCategory == "Input" {
		e.inputRecycleChans[wrapper.name] = make(chan *Packet, Globals().InputRecyclePoolSize)
		e.InputRunners[wrapper.name] = newInputRunner(plugin.(Input), pluginCommons, e.pluginPanicCh)
		e.inputWrappers[wrapper.name] = wrapper
		e.router.metrics.m[wrapper.name] = metrics.NewRegisteredMeter(wrapper.name, metrics.DefaultRegistry)
		return pluginCommons.name
	}

	foRunner := newFORunner(plugin, pluginCommons, e.pluginPanicCh)
	matcher := newMatcher(section.StringList("match", nil), foRunner)
	foRunner.matcher = matcher

	switch pluginCategory {
	case "Filter":
		e.router.addFilterMatcher(matcher)
		e.FilterRunners[foRunner.Name()] = foRunner
		e.filterWrappers[foRunner.Name()] = wrapper
		e.router.metrics.m[wrapper.name] = metrics.NewRegisteredMeter(wrapper.name, metrics.DefaultRegistry)

	case "Output":
		e.router.addOutputMatcher(matcher)
		e.OutputRunners[foRunner.Name()] = foRunner
		e.outputWrappers[foRunner.Name()] = wrapper
		e.router.metrics.m[wrapper.name] = metrics.NewRegisteredMeter(wrapper.name, metrics.DefaultRegistry)

	default:
		panic("unknown plugin: " + pluginCategory)
	}

	return pluginCommons.name
}

func (e *Engine) Shutdown() {
	close(e.shutdown)
}

func (e *Engine) ServeForever() (ret error) {
	var (
		outputsWg = new(sync.WaitGroup)
		filtersWg = new(sync.WaitGroup)
		inputsWg  = new(sync.WaitGroup)
		routerWg  = new(sync.WaitGroup)

		globals = Globals()
		err     error
	)

	log.Info("engine starting...")

	if globals.ClusterEnabled {
		e.controller = czk.NewController(e.zkSvr, globals.ZrootCluster, e.participant, cluster.StrategyRoundRobin, e.onControllerRebalance)
	}

	// setup signal handler first to avoid race condition
	// if Input terminates very fast, global.Shutdown will not be able to trap it
	globals.sigChan = make(chan os.Signal)
	signal.Notify(globals.sigChan, syscall.SIGINT, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)

	e.launchAPIServer()

	if telemetry.Default != nil {
		go func() {
			if err := telemetry.Default.Start(); err != nil {
				log.Error("telemetry[%s]: %s", telemetry.Default.Name(), err)
			}
		}()
	}

	for _, outputRunner := range e.OutputRunners {
		log.Debug("launching Output[%s]...", outputRunner.Name())

		outputsWg.Add(1)
		if err = outputRunner.start(e, outputsWg); err != nil {
			panic(err)
		}
	}

	for _, filterRunner := range e.FilterRunners {
		log.Debug("launching Filter[%s]...", filterRunner.Name())

		filtersWg.Add(1)
		if err = filterRunner.start(e, filtersWg); err != nil {
			panic(err)
		}
	}

	for inputName := range e.inputRecycleChans {
		for i := 0; i < globals.InputRecyclePoolSize; i++ {
			inputPack := newPacket(e.inputRecycleChans[inputName])
			e.inputRecycleChans[inputName] <- inputPack
		}
	}

	for i := 0; i < globals.FilterRecyclePoolSize; i++ {
		filterPack := newPacket(e.filterRecycleChan)
		e.filterRecycleChan <- filterPack
	}

	go e.runWatchdog(globals.WatchdogTick)

	routerWg.Add(1)
	go e.router.Start(routerWg)

	for _, inputRunner := range e.InputRunners {
		log.Debug("launching Input[%s]...", inputRunner.Name())

		inputsWg.Add(1)
		if err = inputRunner.start(e, inputsWg); err != nil {
			panic(err)
		}
	}

	if globals.ClusterEnabled {
		e.launchRPCServer()

		log.Trace("[%s] participant starting...", e.participant)
		if err = e.controller.Start(); err != nil {
			panic(err)
		}
		go e.watchUpgrade(e.ClusterManager().Upgrade())

		log.Info("[%s] participant started", e.participant)
	} else {
		log.Info("cluster disabled")
	}

	configChanged := make(chan *conf.Conf)
	go e.Conf.Watch(time.Second*10, e.stopper, configChanged)

	for !globals.Stopping {
		select {
		case <-configChanged:
			log.Info("%s changed, closing...", e.Conf.ConfPath())
			globals.Stopping = true

		case <-e.shutdown:
			log.Info("shutdown...")
			globals.Stopping = true

		case sig := <-globals.sigChan:
			log.Info("Got signal %s", strings.ToUpper(sig.String()))

			switch sig {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP:
				log.Info("shutdown...")
				globals.Stopping = true
				ret = ErrQuitingSigal

			case syscall.SIGUSR1:
				observer.Publish(SIGUSR1, nil)

			case syscall.SIGUSR2:
				observer.Publish(SIGUSR2, nil)
			}

		case err := <-e.pluginPanicCh:
			log.Info("plugin panic, stopping...")
			ret = err
			globals.Stopping = true
		}
	}

	close(e.stopper)

	if telemetry.Default != nil {
		telemetry.Default.Stop()
	}

	e.Lock()
	for _, inputRunner := range e.InputRunners {
		if inputRunner == nil {
			// the Input plugin already exit
			continue
		}

		log.Debug("Stop message sent to %s", inputRunner.Name())
		inputRunner.Input().Stop(inputRunner)
	}
	e.Unlock()
	inputsWg.Wait() // wait for all inputs done

	// ok, now we are sure no more inputs, but in route.inChan there
	// still may be filter injected packs and output not consumed packs
	// we must wait for all the packs to be consumed before shutdown

	for _, filterRunner := range e.FilterRunners {
		log.Debug("Stop message sent to %s", filterRunner.Name())
		e.router.removeFilterMatcher <- filterRunner.getMatcher()
	}
	filtersWg.Wait()

	for _, outputRunner := range e.OutputRunners {
		log.Debug("Stop message sent to %s", outputRunner.Name())
		e.router.removeOutputMatcher <- outputRunner.getMatcher()
	}
	outputsWg.Wait()

	e.router.Stop()
	routerWg.Wait()
	log.Info("Router stopped")

	e.stopAPIServer()

	if globals.ClusterEnabled {
		e.stopRPCServer()

		if err = e.controller.Stop(); err != nil {
			log.Error("%v", err)
		}
	}

	if ret != nil {
		log.Info("shutdown complete: %s!", ret)
	} else {
		log.Info("shutdown complete!")
	}

	return
}
<commit_msg>tweak of log<commit_after>// +build !v2

package engine

import (
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/funkygao/dbus"
	"github.com/funkygao/dbus/pkg/cluster"
	czk "github.com/funkygao/dbus/pkg/cluster/zk"
	"github.com/funkygao/gafka/ctx"
	"github.com/funkygao/gafka/telemetry"
	"github.com/funkygao/gafka/telemetry/influxdb"
	"github.com/funkygao/go-metrics"
	"github.com/funkygao/golib/observer"
	conf "github.com/funkygao/jsconf"
	log "github.com/funkygao/log4go"
	"github.com/gorilla/mux"
)

var (
	_ PluginHelper = &Engine{}
)

// Engine is the pipeline engine of the data bus system which manages the core loop.
type Engine struct {
	sync.RWMutex
	*conf.Conf

	zkSvr       string
	participant cluster.Participant
	controller  cluster.Controller
	epoch       int // cache of latest cluster leader epoch

	// API Server
	apiListener net.Listener
	apiServer   *http.Server
	apiRouter   *mux.Router

	// RPC Server
	rpcListener net.Listener
	rpcServer   *http.Server
	rpcRouter   *mux.Router

	// input plugin resources map
	irm   map[string][]cluster.Resource
	irmMu sync.Mutex

	// dataflow router
	router *Router

	InputRunners  map[string]*iRunner
	inputWrappers map[string]*pluginWrapper

	FilterRunners  map[string]FilterRunner
	filterWrappers map[string]*pluginWrapper

	OutputRunners  map[string]OutputRunner
	outputWrappers map[string]*pluginWrapper

	inputRecycleChans map[string]chan *Packet
	filterRecycleChan chan *Packet

	hostname      string
	pid           int
	stopper       chan struct{}
	shutdown      chan struct{}
	pluginPanicCh chan error
}

// New creates an engine.
func New(globals *GlobalConfig) *Engine {
	if globals == nil {
		globals = DefaultGlobals()
	}
	Globals = func() *GlobalConfig {
		return globals
	}

	hostname, err := os.Hostname()
	if err != nil {
		panic(err)
	}

	// make the participant
	localIP, err := ctx.LocalIP()
	if err != nil {
		panic(err)
	}

	return &Engine{
		pid:           os.Getpid(),
		hostname:      hostname,
		stopper:       make(chan struct{}),
		shutdown:      make(chan struct{}),
		pluginPanicCh: make(chan error),

		router: newRouter(),

		InputRunners:   make(map[string]*iRunner),
		inputWrappers:  make(map[string]*pluginWrapper),
		FilterRunners:  make(map[string]FilterRunner),
		filterWrappers: make(map[string]*pluginWrapper),
		OutputRunners:  make(map[string]OutputRunner),
		outputWrappers: make(map[string]*pluginWrapper),

		inputRecycleChans: make(map[string]chan *Packet),
		filterRecycleChan: make(chan *Packet, globals.FilterRecyclePoolSize),

		participant: cluster.Participant{
			Endpoint: fmt.Sprintf("%s:%d", localIP.String(), globals.RPCPort),
			Weight:   runtime.NumCPU() * 100,
			State:    cluster.StateOnline,
			Revision: dbus.Revision,
			APIPort:  globals.APIPort,
		},
	}
}

func (e *Engine) stopInputRunner(name string) {
	e.Lock()
	e.InputRunners[name] = nil
	e.Unlock()
}

// ClonePacket is used for plugin Filter to generate new Packet: copy on write.
// The generated Packet will use dedicated filter recycle chan.
func (e *Engine) ClonePacket(p *Packet) *Packet {
	pack := <-e.filterRecycleChan
	pack.Reset()
	p.copyTo(pack)
	return pack
}

// ClusterManager returns the cluster manager.
// If cluster is disabled, returns nil.
func (e *Engine) ClusterManager() cluster.Manager {
	if Globals().ClusterEnabled {
		return e.controller.(cluster.Manager)
	}

	return nil
}

// SubmitDAG submits a DAG configuration to the engine.
func (e *Engine) SubmitDAG(cf *conf.Conf) *Engine {
	return e.loadConfig(cf)
}

func (e *Engine) loadConfig(cf *conf.Conf) *Engine {
	e.Conf = cf
	Globals().Conf = cf

	// 'plugins' section
	var pluginNames = make(map[string]struct{})
	for i := 0; i < len(e.List("plugins", nil)); i++ {
		section, err := e.Section(fmt.Sprintf("plugins[%d]", i))
		if err != nil {
			panic(err)
		}

		name := e.loadPluginSection(section)
		if _, duplicated := pluginNames[name]; duplicated {
			// router.metrics will be bad with dup name
			panic("duplicated plugin name: " + name)
		}
		pluginNames[name] = struct{}{}
	}

	// influxdb related section
	if c, err := influxdb.NewConfig(cf.String("influx_addr", ""),
		cf.String("influx_db", "dbus"), "", "",
		cf.Duration("influx_tick", time.Minute)); err == nil {
		telemetry.Default = influxdb.New(metrics.DefaultRegistry, c)
	}

	return e
}

// LoadFrom load the configuration by location.
// The location can be empty: use default zk zone /dbus/conf.
// If config is stored on file, the loc arg is file path.
// If config is stored on zookeeper, the loc arg is like localhost:2181/foo/bar.
func (e *Engine) LoadFrom(loc string) *Engine {
	if len(loc) == 0 {
		// if no location provided, use the default zk
		loc = fmt.Sprintf("%s%s", ctx.ZoneZkAddrs(ctx.DefaultZone()), Globals().ZrootConf)
	}

	zkSvr, realPath := parseConfigPath(loc)
	var (
		cf  *conf.Conf
		err error
	)
	if len(zkSvr) == 0 {
		// from file system
		cf, err = conf.Load(realPath)
	} else {
		// from zookeeper
		e.zkSvr = zkSvr
		cf, err = conf.Load(realPath, conf.WithZkSvr(zkSvr))
		if err != nil {
			err = fmt.Errorf("%s %v", loc, err)
		}
	}
	if err != nil {
		panic(err)
	}

	return e.loadConfig(cf)
}

func (e *Engine) loadPluginSection(section *conf.Conf) string {
	pluginCommons := new(pluginCommons)
	pluginCommons.loadConfig(section)

	wrapper := &pluginWrapper{
		name:          pluginCommons.name,
		configCreator: func() *conf.Conf { return section },
	}
	var ok bool
	if wrapper.pluginCreator, ok = availablePlugins[pluginCommons.class]; !ok {
		panic("unknown plugin type: " + pluginCommons.class)
	}

	pluginType := pluginTypeRegex.FindStringSubmatch(pluginCommons.class)
	if len(pluginType) < 2 {
		panic("invalid plugin type: " + pluginCommons.class)
	}

	plugin := wrapper.Create()
	pluginCategory := pluginType[1]
	if pluginCategory == "Input" {
		e.inputRecycleChans[wrapper.name] = make(chan *Packet, Globals().InputRecyclePoolSize)
		e.InputRunners[wrapper.name] = newInputRunner(plugin.(Input), pluginCommons, e.pluginPanicCh)
		e.inputWrappers[wrapper.name] = wrapper
		e.router.metrics.m[wrapper.name] = metrics.NewRegisteredMeter(wrapper.name, metrics.DefaultRegistry)
		return pluginCommons.name
	}

	foRunner := newFORunner(plugin, pluginCommons, e.pluginPanicCh)
	matcher := newMatcher(section.StringList("match", nil), foRunner)
	foRunner.matcher = matcher

	switch pluginCategory {
	case "Filter":
		e.router.addFilterMatcher(matcher)
		e.FilterRunners[foRunner.Name()] = foRunner
		e.filterWrappers[foRunner.Name()] = wrapper
		e.router.metrics.m[wrapper.name] = metrics.NewRegisteredMeter(wrapper.name, metrics.DefaultRegistry)

	case "Output":
		e.router.addOutputMatcher(matcher)
		e.OutputRunners[foRunner.Name()] = foRunner
		e.outputWrappers[foRunner.Name()] = wrapper
		e.router.metrics.m[wrapper.name] = metrics.NewRegisteredMeter(wrapper.name, metrics.DefaultRegistry)

	default:
		panic("unknown plugin: " + pluginCategory)
	}

	return pluginCommons.name
}

func (e *Engine) Shutdown() {
	close(e.shutdown)
}

func (e *Engine) ServeForever() (ret error) {
	var (
		outputsWg = new(sync.WaitGroup)
		filtersWg = new(sync.WaitGroup)
		inputsWg  = new(sync.WaitGroup)
		routerWg  = new(sync.WaitGroup)

		globals = Globals()
		err     error
	)

	log.Trace("engine starting...")

	if globals.ClusterEnabled {
		e.controller = czk.NewController(e.zkSvr, globals.ZrootCluster, e.participant, cluster.StrategyRoundRobin, e.onControllerRebalance)
	}

	// setup signal handler first to avoid race condition
	// if Input terminates very fast, global.Shutdown will not be able to trap it
	globals.sigChan = make(chan os.Signal)
	signal.Notify(globals.sigChan, syscall.SIGINT, syscall.SIGHUP, syscall.SIGTERM, syscall.SIGUSR1, syscall.SIGUSR2)

	e.launchAPIServer()

	if telemetry.Default != nil {
		go func() {
			if err := telemetry.Default.Start(); err != nil {
				log.Error("telemetry[%s]: %s", telemetry.Default.Name(), err)
			}
		}()
	}

	for _, outputRunner := range e.OutputRunners {
		log.Debug("launching Output[%s]...", outputRunner.Name())

		outputsWg.Add(1)
		if err = outputRunner.start(e, outputsWg); err != nil {
			panic(err)
		}
	}

	for _, filterRunner := range e.FilterRunners {
		log.Debug("launching Filter[%s]...", filterRunner.Name())

		filtersWg.Add(1)
		if err = filterRunner.start(e, filtersWg); err != nil {
			panic(err)
		}
	}

	for inputName := range e.inputRecycleChans {
		for i := 0; i < globals.InputRecyclePoolSize; i++ {
			inputPack := newPacket(e.inputRecycleChans[inputName])
			e.inputRecycleChans[inputName] <- inputPack
		}
	}

	for i := 0; i < globals.FilterRecyclePoolSize; i++ {
		filterPack := newPacket(e.filterRecycleChan)
		e.filterRecycleChan <- filterPack
	}

	go e.runWatchdog(globals.WatchdogTick)

	routerWg.Add(1)
	go e.router.Start(routerWg)

	for _, inputRunner := range e.InputRunners {
		log.Debug("launching Input[%s]...", inputRunner.Name())

		inputsWg.Add(1)
		if err = inputRunner.start(e, inputsWg); err != nil {
			panic(err)
		}
	}

	if globals.ClusterEnabled {
		e.launchRPCServer()

		log.Trace("[%s] participant starting...", e.participant)
		if err = e.controller.Start(); err != nil {
			panic(err)
		}
		go e.watchUpgrade(e.ClusterManager().Upgrade())

		log.Info("[%s] participant started", e.participant)
	} else {
		log.Info("cluster disabled")
	}

	configChanged := make(chan *conf.Conf)
	go e.Conf.Watch(time.Second*10, e.stopper, configChanged)

	log.Info("engine started")
	for !globals.Stopping {
		select {
		case <-configChanged:
			log.Info("%s changed, closing...", e.Conf.ConfPath())
			globals.Stopping = true

		case <-e.shutdown:
			log.Info("shutdown...")
			globals.Stopping = true

		case sig := <-globals.sigChan:
			log.Info("Got signal %s", strings.ToUpper(sig.String()))

			switch sig {
			case syscall.SIGINT, syscall.SIGTERM, syscall.SIGHUP:
				log.Info("shutdown...")
				globals.Stopping = true
				ret = ErrQuitingSigal

			case syscall.SIGUSR1:
				observer.Publish(SIGUSR1, nil)

			case syscall.SIGUSR2:
				observer.Publish(SIGUSR2, nil)
			}

		case err := <-e.pluginPanicCh:
			log.Info("plugin panic, stopping...")
			ret = err
			globals.Stopping = true
		}
	}

	close(e.stopper)

	if telemetry.Default != nil {
		telemetry.Default.Stop()
	}

	e.Lock()
	for _, inputRunner := range e.InputRunners {
		if inputRunner == nil {
			// the Input plugin already exit
			continue
		}

		log.Debug("Stop message sent to %s", inputRunner.Name())
		inputRunner.Input().Stop(inputRunner)
	}
	e.Unlock()
	inputsWg.Wait() // wait for all inputs done

	// ok, now we are sure no more inputs, but in route.inChan there
	// still may be filter injected packs and output not consumed packs
	// we must wait for all the packs to be consumed before shutdown

	for _, filterRunner := range e.FilterRunners {
		log.Debug("Stop message sent to %s", filterRunner.Name())
		e.router.removeFilterMatcher <- filterRunner.getMatcher()
	}
	filtersWg.Wait()

	for _, outputRunner := range e.OutputRunners {
		log.Debug("Stop message sent to %s", outputRunner.Name())
		e.router.removeOutputMatcher <- outputRunner.getMatcher()
	}
	outputsWg.Wait()

	e.router.Stop()
	routerWg.Wait()
	log.Info("Router stopped")

	e.stopAPIServer()

	if globals.ClusterEnabled {
		e.stopRPCServer()

		if err = e.controller.Stop(); err != nil {
			log.Error("%v", err)
		}
	}

	if ret != nil {
		log.Info("shutdown complete: %s!", ret)
	} else {
		log.Info("shutdown complete!")
	}

	return
}
<|endoftext|>


<commit_before>package helpers

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"math"

	"sync"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/bwmarrin/discordgo"
)

const (
	LEFT_ARROW_EMOJI  = "⬅"
	RIGHT_ARROW_EMOJI = "➡"
	X_EMOJI           = "🇽"
)

// map of messageID to pagedEmbedMessage
var pagedEmbededMessages map[string]*pagedEmbedMessage

type pagedEmbedMessage struct {
	files           []*discordgo.File
	fullEmbed       *discordgo.MessageEmbed
	totalNumOfPages int
	currentPage     int
	fieldsPerPage   int
	color           int
	messageID       string
	channelID       string
	userId          string //user who triggered the message
	msgType         string // "image" - will cause the embed to page through the files instead of fields
}

func init() {
	pagedEmbededMessages = make(map[string]*pagedEmbedMessage)
}

// will remove all reactions from all paged embed messages.
//  mainly used on bot uninit to clean embeds
func RemoveReactionsFromPagedEmbeds() {
	var syncGroup sync.WaitGroup
	for _, pagedEmbed := range pagedEmbededMessages {
		syncGroup.Add(1)
		go func() {
			cache.GetSession().MessageReactionsRemoveAll(pagedEmbed.channelID, pagedEmbed.messageID)
			syncGroup.Done()
		}()
	}
	// make sure all reactions get removed before bot exits
	syncGroup.Wait()
}

// GetPagedMessage will return the paged message if there is one, otherwill will return nil
func GetPagedMessage(messageID string) *pagedEmbedMessage {
	pagedMessaged, _ := pagedEmbededMessages[messageID]
	return pagedMessaged
}

// CreatePagedMessage creates the paged messages
func SendPagedMessage(msg *discordgo.Message, embed *discordgo.MessageEmbed, fieldsPerPage int) error {

	// if there aren't multiple fields to be paged through, or if the amount of fields is less than the requested fields per page
	//  just send a normal embed
	if len(embed.Fields) < 2 || len(embed.Fields) <= fieldsPerPage {
		SendEmbed(msg.ChannelID, embed)
		return nil
	}

	// fields per page can not be less than 1
	if fieldsPerPage < 1 {
		return errors.New("fieldsPerPage can not be less than 1")
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		fullEmbed:       embed,
		channelID:       msg.ChannelID,
		currentPage:     1,
		fieldsPerPage:   fieldsPerPage,
		totalNumOfPages: int(math.Ceil(float64(len(embed.Fields)) / float64(fieldsPerPage))),
		userId:          msg.Author.ID,
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// SendPagedImageMessage creates the paged image messages
func SendPagedImageMessage(msg *discordgo.Message, msgSend *discordgo.MessageSend) error {
	if msgSend.Embed == nil {
		return errors.New("parameter msgSend must contain an embed")
	}

	// make sure the image url is set to the name of the first file incease it wasn't set
	msgSend.Embed.Image.URL = fmt.Sprintf("attachment://%s", msgSend.Files[0].Name)

	// check if there are multiple files, not just send it normally
	if len(msgSend.Files) < 2 {
		SendComplex(msg.ChannelID, msgSend)
		return nil
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		fullEmbed:       msgSend.Embed,
		channelID:       msg.ChannelID,
		currentPage:     1,
		fieldsPerPage:   1,
		totalNumOfPages: len(msgSend.Files),
		files:           msgSend.Files,
		userId:          msg.Author.ID,
		msgType:         "image",
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// UpdateMessagePage will update the page based on the given direction and current page
//  reactions must be the left or right arrow
func (p *pagedEmbedMessage) UpdateMessagePage(reaction *discordgo.MessageReactionAdd) {
	// check for valid reaction
	if LEFT_ARROW_EMOJI != reaction.Emoji.Name && RIGHT_ARROW_EMOJI != reaction.Emoji.Name && X_EMOJI != reaction.Emoji.Name {
		return
	}

	// check if user who made the embed message is closing it
	if reaction.UserID == p.userId && X_EMOJI == reaction.Emoji.Name {
		delete(pagedEmbededMessages, reaction.MessageID)
		cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)
		return
	}

	// update current page based on direction
	if LEFT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage--
		if p.currentPage == 0 {
			p.currentPage = p.totalNumOfPages
		}
	}
	if RIGHT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage++
		if p.currentPage > p.totalNumOfPages {
			p.currentPage = 1
		}
	}

	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// updated stats
	if p.msgType == "image" {
		// image embeds can't be edited, need to delete and remate it
		cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)

		// we need to split and reset the reader since a reader can only be used once
		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		// change the url of the embed to point to the new image
		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)

		// update footer and send message
		tempEmbed.Footer = p.getEmbedFooter()
		sentMessage, _ := SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

		// update map with new message id since
		originalmsgID := p.messageID
		p.messageID = sentMessage[0].ID
		p.addReactionsToMessage()
		pagedEmbededMessages[sentMessage[0].ID] = p
		delete(pagedEmbededMessages, originalmsgID)
	} else {

		// get start and end fields based on current page and fields per page
		startField := (p.currentPage - 1) * p.fieldsPerPage
		endField := startField + p.fieldsPerPage
		if endField > len(p.fullEmbed.Fields) {
			endField = len(p.fullEmbed.Fields)
		}

		tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		tempEmbed.Footer = p.getEmbedFooter()
		EditEmbed(p.channelID, p.messageID, tempEmbed)

		cache.GetSession().MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
	}

}

// setupAndSendFirstMessage
func (p *pagedEmbedMessage) setupAndSendFirstMessage() {
	var sentMessage []*discordgo.Message
	var err error

	// copy the embeded message so changes can be made to it
	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// set footer which will hold information about the page it is on
	tempEmbed.Footer = p.getEmbedFooter()

	if p.msgType == "image" {

		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)
		sentMessage, _ = SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

	} else {
		// reduce fields to the fields per page
		tempEmbed.Fields = tempEmbed.Fields[:p.fieldsPerPage]

		sentMessage, err = SendEmbed(p.channelID, tempEmbed)
	}

	if err != nil {
		// should probably handle this at some point lul
		return
	}
	p.messageID = sentMessage[0].ID

	p.addReactionsToMessage()
}

// getEmbedFooter is a simlple helper function to return the footer for the embed message
func (p *pagedEmbedMessage) getEmbedFooter() *discordgo.MessageEmbedFooter {
	return &discordgo.MessageEmbedFooter{
		Text: fmt.Sprintf("Page: %d / %d", p.currentPage, p.totalNumOfPages),
	}
}

func (p *pagedEmbedMessage) addReactionsToMessage() {
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, LEFT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, RIGHT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, X_EMOJI)
}
<commit_msg>[pagedEmbed] fixes issue with paged image embeds reacting to other users hitting the X reaction<commit_after>package helpers

import (
	"bytes"
	"errors"
	"fmt"
	"io"
	"math"

	"sync"

	"github.com/Seklfreak/Robyul2/cache"
	"github.com/bwmarrin/discordgo"
)

const (
	LEFT_ARROW_EMOJI  = "⬅"
	RIGHT_ARROW_EMOJI = "➡"
	X_EMOJI           = "🇽"
)

// map of messageID to pagedEmbedMessage
var pagedEmbededMessages map[string]*pagedEmbedMessage

type pagedEmbedMessage struct {
	files           []*discordgo.File
	fullEmbed       *discordgo.MessageEmbed
	totalNumOfPages int
	currentPage     int
	fieldsPerPage   int
	color           int
	messageID       string
	channelID       string
	userId          string //user who triggered the message
	msgType         string // "image" - will cause the embed to page through the files instead of fields
}

func init() {
	pagedEmbededMessages = make(map[string]*pagedEmbedMessage)
}

// will remove all reactions from all paged embed messages.
//  mainly used on bot uninit to clean embeds
func RemoveReactionsFromPagedEmbeds() {
	var syncGroup sync.WaitGroup
	for _, pagedEmbed := range pagedEmbededMessages {
		syncGroup.Add(1)
		go func() {
			cache.GetSession().MessageReactionsRemoveAll(pagedEmbed.channelID, pagedEmbed.messageID)
			syncGroup.Done()
		}()
	}
	// make sure all reactions get removed before bot exits
	syncGroup.Wait()
}

// GetPagedMessage will return the paged message if there is one, otherwill will return nil
func GetPagedMessage(messageID string) *pagedEmbedMessage {
	pagedMessaged, _ := pagedEmbededMessages[messageID]
	return pagedMessaged
}

// CreatePagedMessage creates the paged messages
func SendPagedMessage(msg *discordgo.Message, embed *discordgo.MessageEmbed, fieldsPerPage int) error {

	// if there aren't multiple fields to be paged through, or if the amount of fields is less than the requested fields per page
	//  just send a normal embed
	if len(embed.Fields) < 2 || len(embed.Fields) <= fieldsPerPage {
		SendEmbed(msg.ChannelID, embed)
		return nil
	}

	// fields per page can not be less than 1
	if fieldsPerPage < 1 {
		return errors.New("fieldsPerPage can not be less than 1")
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		fullEmbed:       embed,
		channelID:       msg.ChannelID,
		currentPage:     1,
		fieldsPerPage:   fieldsPerPage,
		totalNumOfPages: int(math.Ceil(float64(len(embed.Fields)) / float64(fieldsPerPage))),
		userId:          msg.Author.ID,
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// SendPagedImageMessage creates the paged image messages
func SendPagedImageMessage(msg *discordgo.Message, msgSend *discordgo.MessageSend) error {
	if msgSend.Embed == nil {
		return errors.New("parameter msgSend must contain an embed")
	}

	// make sure the image url is set to the name of the first file incease it wasn't set
	msgSend.Embed.Image.URL = fmt.Sprintf("attachment://%s", msgSend.Files[0].Name)

	// check if there are multiple files, not just send it normally
	if len(msgSend.Files) < 2 {
		SendComplex(msg.ChannelID, msgSend)
		return nil
	}

	// create paged message
	pagedMessage := &pagedEmbedMessage{
		fullEmbed:       msgSend.Embed,
		channelID:       msg.ChannelID,
		currentPage:     1,
		fieldsPerPage:   1,
		totalNumOfPages: len(msgSend.Files),
		files:           msgSend.Files,
		userId:          msg.Author.ID,
		msgType:         "image",
	}

	pagedMessage.setupAndSendFirstMessage()

	pagedEmbededMessages[pagedMessage.messageID] = pagedMessage
	return nil
}

// UpdateMessagePage will update the page based on the given direction and current page
//  reactions must be the left or right arrow
func (p *pagedEmbedMessage) UpdateMessagePage(reaction *discordgo.MessageReactionAdd) {
	// check for valid reaction
	if LEFT_ARROW_EMOJI != reaction.Emoji.Name && RIGHT_ARROW_EMOJI != reaction.Emoji.Name && X_EMOJI != reaction.Emoji.Name {
		return
	}

	// check if user who made the embed message is closing it
	if X_EMOJI == reaction.Emoji.Name {
		if reaction.UserID == p.userId {
			delete(pagedEmbededMessages, reaction.MessageID)
			cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)
		}
		return
	}

	// update current page based on direction
	if LEFT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage--
		if p.currentPage == 0 {
			p.currentPage = p.totalNumOfPages
		}
	}
	if RIGHT_ARROW_EMOJI == reaction.Emoji.Name {
		p.currentPage++
		if p.currentPage > p.totalNumOfPages {
			p.currentPage = 1
		}
	}

	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// updated stats
	if p.msgType == "image" {
		// image embeds can't be edited, need to delete and remate it
		cache.GetSession().ChannelMessageDelete(p.channelID, p.messageID)

		// we need to split and reset the reader since a reader can only be used once
		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		// change the url of the embed to point to the new image
		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)

		// update footer and send message
		tempEmbed.Footer = p.getEmbedFooter()
		sentMessage, _ := SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

		// update map with new message id since
		originalmsgID := p.messageID
		p.messageID = sentMessage[0].ID
		p.addReactionsToMessage()
		pagedEmbededMessages[sentMessage[0].ID] = p
		delete(pagedEmbededMessages, originalmsgID)
	} else {

		// get start and end fields based on current page and fields per page
		startField := (p.currentPage - 1) * p.fieldsPerPage
		endField := startField + p.fieldsPerPage
		if endField > len(p.fullEmbed.Fields) {
			endField = len(p.fullEmbed.Fields)
		}

		tempEmbed.Fields = tempEmbed.Fields[startField:endField]
		tempEmbed.Footer = p.getEmbedFooter()
		EditEmbed(p.channelID, p.messageID, tempEmbed)

		cache.GetSession().MessageReactionRemove(reaction.ChannelID, reaction.MessageID, reaction.Emoji.Name, reaction.UserID)
	}

}

// setupAndSendFirstMessage
func (p *pagedEmbedMessage) setupAndSendFirstMessage() {
	var sentMessage []*discordgo.Message
	var err error

	// copy the embeded message so changes can be made to it
	tempEmbed := &discordgo.MessageEmbed{}
	*tempEmbed = *p.fullEmbed

	// set footer which will hold information about the page it is on
	tempEmbed.Footer = p.getEmbedFooter()

	if p.msgType == "image" {

		var buf bytes.Buffer
		newReader := io.TeeReader(p.files[p.currentPage-1].Reader, &buf)
		p.files[p.currentPage-1].Reader = &buf

		tempEmbed.Image.URL = fmt.Sprintf("attachment://%s", p.files[p.currentPage-1].Name)
		sentMessage, _ = SendComplex(p.channelID, &discordgo.MessageSend{
			Embed: tempEmbed,
			Files: []*discordgo.File{&discordgo.File{
				Name:   p.files[p.currentPage-1].Name,
				Reader: newReader,
			}},
		})

	} else {
		// reduce fields to the fields per page
		tempEmbed.Fields = tempEmbed.Fields[:p.fieldsPerPage]

		sentMessage, err = SendEmbed(p.channelID, tempEmbed)
	}

	if err != nil {
		// should probably handle this at some point lul
		return
	}
	p.messageID = sentMessage[0].ID

	p.addReactionsToMessage()
}

// getEmbedFooter is a simlple helper function to return the footer for the embed message
func (p *pagedEmbedMessage) getEmbedFooter() *discordgo.MessageEmbedFooter {
	return &discordgo.MessageEmbedFooter{
		Text: fmt.Sprintf("Page: %d / %d", p.currentPage, p.totalNumOfPages),
	}
}

func (p *pagedEmbedMessage) addReactionsToMessage() {
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, LEFT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, RIGHT_ARROW_EMOJI)
	cache.GetSession().MessageReactionAdd(p.channelID, p.messageID, X_EMOJI)
}
<|endoftext|>