<commit_before>package modules

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/pajbot/pajbot2/pkg"
)

const garbageCollectionInterval = 1 * time.Minute
const maxMessageAge = 5 * time.Minute

type nukeModule struct {
	botChannel pkg.BotChannel

	server        *server
	messages      map[string][]nukeMessage
	messagesMutex sync.Mutex

	ticker *time.Ticker
}

type nukeMessage struct {
	user      pkg.User
	message   pkg.Message
	timestamp time.Time
}

func newNuke() pkg.Module {
	m := &nukeModule{
		server:   &_server,
		messages: make(map[string][]nukeMessage),
	}

	m.ticker = time.NewTicker(garbageCollectionInterval)

	go func() {
		for {
			select {
			case <-m.ticker.C:
				m.garbageCollect()
			}
		}
	}()

	return m
}

var nukeSpec = moduleSpec{
	id:    "nuke",
	name:  "Nuke",
	maker: newNuke,

	enabledByDefault: true,
}

func (m *nukeModule) Initialize(botChannel pkg.BotChannel, settings []byte) error {
	m.botChannel = botChannel

	return nil
}

func (m *nukeModule) Disable() error {
	return nil
}

func (m *nukeModule) Spec() pkg.ModuleSpec {
	return &nukeSpec
}

func (m *nukeModule) BotChannel() pkg.BotChannel {
	return m.botChannel
}

func (m *nukeModule) OnWhisper(bot pkg.BotChannel, user pkg.User, message pkg.Message) error {
	const usageString = `Usage: #channel !nuke phrase phrase phrase time`

	parts := strings.Split(message.GetText(), " ")
	// Minimum required parts: 4
	// !nuke PHRASE SCROLLBACK_LENGTH TIMEOUT_DURATION
	if len(parts) >= 4 {
		if parts[0] != "!nuke" {
			return nil
		}

		// TODO: Add another specific global/channel permission to check
		if !user.IsModerator() && !user.IsBroadcaster(bot.Channel()) && !user.HasChannelPermission(bot.Channel(), pkg.PermissionModeration) && !user.HasGlobalPermission(pkg.PermissionModeration) {
			return nil
		}

		phrase := strings.Join(parts[1:len(parts)-2], " ")
		scrollbackLength, err := time.ParseDuration(parts[len(parts)-2])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if scrollbackLength < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("scrollback length must be positive")
		}
		timeoutDuration, err := time.ParseDuration(parts[len(parts)-1])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if timeoutDuration < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("timeout duration must be positive")
		}

		m.nuke(user, bot, phrase, scrollbackLength, timeoutDuration)
	}

	return nil
}

func (m *nukeModule) OnMessage(bot pkg.BotChannel, user pkg.User, message pkg.Message, action pkg.Action) error {
	defer func() {
		m.addMessage(bot.Channel(), user, message)
	}()

	parts := strings.Split(message.GetText(), " ")
	// Minimum required parts: 4
	// !nuke PHRASE SCROLLBACK_LENGTH TIMEOUT_DURATION
	if len(parts) >= 4 {
		if parts[0] != "!nuke" {
			return nil
		}

		// TODO: Add another specific global/channel permission to check
		if !user.IsModerator() && !user.IsBroadcaster(bot.Channel()) && !user.HasChannelPermission(bot.Channel(), pkg.PermissionModeration) && !user.HasGlobalPermission(pkg.PermissionModeration) {
			return nil
		}

		phrase := strings.Join(parts[1:len(parts)-2], " ")
		scrollbackLength, err := time.ParseDuration(parts[len(parts)-2])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if scrollbackLength < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("scrollback length must be positive")
		}
		timeoutDuration, err := time.ParseDuration(parts[len(parts)-1])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if timeoutDuration < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("timeout duration must be positive")
		}

		m.nuke(user, bot, phrase, scrollbackLength, timeoutDuration)
	}

	return nil
}

func (m *nukeModule) garbageCollect() {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	now := time.Now()

	for channelID := range m.messages {
		for i := 0; i < len(m.messages[channelID]); i++ {
			diff := now.Sub(m.messages[channelID][i].timestamp)
			if diff < maxMessageAge {
				m.messages[channelID] = m.messages[channelID][i:]
				break
			}
		}
	}
}

func (m *nukeModule) nuke(source pkg.User, bot pkg.BotChannel, phrase string, scrollbackLength, timeoutDuration time.Duration) {
	if timeoutDuration > 24*time.Hour {
		timeoutDuration = 24 * time.Hour
	}

	lowercasePhrase := strings.ToLower(phrase)

	matcher := func(msg *nukeMessage) bool {
		return strings.Contains(strings.ToLower(msg.message.GetText()), lowercasePhrase)
	}

	reason := "Nuked '" + phrase + "'"

	if strings.HasPrefix(phrase, "/") && strings.HasSuffix(phrase, "/") {
		regex, err := regexp.Compile(phrase[1 : len(phrase)-1])
		if err == nil {
			reason = "Nuked r'" + phrase[1:len(phrase)-1] + "'"
			matcher = func(msg *nukeMessage) bool {
				return regex.MatchString(msg.message.GetText())
			}
		}
		// parse as regex
	}

	now := time.Now()
	timeoutDurationInSeconds := int(timeoutDuration.Seconds())

	if timeoutDurationInSeconds < 1 {
		// Timeout duration too short
		return
	}

	targets := make(map[string]pkg.User)

	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	messages := m.messages[bot.Channel().GetID()]

	for i := len(messages) - 1; i >= 0; i-- {
		diff := now.Sub(messages[i].timestamp)
		if diff > scrollbackLength {
			// We've gone far enough in the buffer, time to exit
			break
		}

		if matcher(&messages[i]) {
			targets[messages[i].user.GetID()] = messages[i].user
		}
	}

	for _, user := range targets {
		bot.SingleTimeout(user, timeoutDurationInSeconds, reason)
	}

	fmt.Printf("%s nuked %d users for the phrase %s in the last %s for %s\n", source.GetName(), len(targets), phrase, scrollbackLength, timeoutDuration)
}

func (m *nukeModule) addMessage(channel pkg.Channel, user pkg.User, message pkg.Message) {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	m.messages[channel.GetID()] = append(m.messages[channel.GetID()], nukeMessage{
		user:      user,
		message:   message,
		timestamp: time.Now(),
	})
}
<commit_msg>change max nuke duration from 1d to 3d<commit_after>package modules

import (
	"errors"
	"fmt"
	"regexp"
	"strings"
	"sync"
	"time"

	"github.com/pajbot/pajbot2/pkg"
)

const garbageCollectionInterval = 1 * time.Minute
const maxMessageAge = 5 * time.Minute

type nukeModule struct {
	botChannel pkg.BotChannel

	server        *server
	messages      map[string][]nukeMessage
	messagesMutex sync.Mutex

	ticker *time.Ticker
}

type nukeMessage struct {
	user      pkg.User
	message   pkg.Message
	timestamp time.Time
}

func newNuke() pkg.Module {
	m := &nukeModule{
		server:   &_server,
		messages: make(map[string][]nukeMessage),
	}

	m.ticker = time.NewTicker(garbageCollectionInterval)

	go func() {
		for {
			select {
			case <-m.ticker.C:
				m.garbageCollect()
			}
		}
	}()

	return m
}

var nukeSpec = moduleSpec{
	id:    "nuke",
	name:  "Nuke",
	maker: newNuke,

	enabledByDefault: true,
}

func (m *nukeModule) Initialize(botChannel pkg.BotChannel, settings []byte) error {
	m.botChannel = botChannel

	return nil
}

func (m *nukeModule) Disable() error {
	return nil
}

func (m *nukeModule) Spec() pkg.ModuleSpec {
	return &nukeSpec
}

func (m *nukeModule) BotChannel() pkg.BotChannel {
	return m.botChannel
}

func (m *nukeModule) OnWhisper(bot pkg.BotChannel, user pkg.User, message pkg.Message) error {
	const usageString = `Usage: #channel !nuke phrase phrase phrase time`

	parts := strings.Split(message.GetText(), " ")
	// Minimum required parts: 4
	// !nuke PHRASE SCROLLBACK_LENGTH TIMEOUT_DURATION
	if len(parts) >= 4 {
		if parts[0] != "!nuke" {
			return nil
		}

		// TODO: Add another specific global/channel permission to check
		if !user.IsModerator() && !user.IsBroadcaster(bot.Channel()) && !user.HasChannelPermission(bot.Channel(), pkg.PermissionModeration) && !user.HasGlobalPermission(pkg.PermissionModeration) {
			return nil
		}

		phrase := strings.Join(parts[1:len(parts)-2], " ")
		scrollbackLength, err := time.ParseDuration(parts[len(parts)-2])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if scrollbackLength < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("scrollback length must be positive")
		}
		timeoutDuration, err := time.ParseDuration(parts[len(parts)-1])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if timeoutDuration < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("timeout duration must be positive")
		}

		m.nuke(user, bot, phrase, scrollbackLength, timeoutDuration)
	}

	return nil
}

func (m *nukeModule) OnMessage(bot pkg.BotChannel, user pkg.User, message pkg.Message, action pkg.Action) error {
	defer func() {
		m.addMessage(bot.Channel(), user, message)
	}()

	parts := strings.Split(message.GetText(), " ")
	// Minimum required parts: 4
	// !nuke PHRASE SCROLLBACK_LENGTH TIMEOUT_DURATION
	if len(parts) >= 4 {
		if parts[0] != "!nuke" {
			return nil
		}

		// TODO: Add another specific global/channel permission to check
		if !user.IsModerator() && !user.IsBroadcaster(bot.Channel()) && !user.HasChannelPermission(bot.Channel(), pkg.PermissionModeration) && !user.HasGlobalPermission(pkg.PermissionModeration) {
			return nil
		}

		phrase := strings.Join(parts[1:len(parts)-2], " ")
		scrollbackLength, err := time.ParseDuration(parts[len(parts)-2])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if scrollbackLength < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("scrollback length must be positive")
		}
		timeoutDuration, err := time.ParseDuration(parts[len(parts)-1])
		if err != nil {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return err
		}
		if timeoutDuration < 0 {
			bot.Mention(user, "usage: !nuke bad phrase 1m 10m")
			return errors.New("timeout duration must be positive")
		}

		m.nuke(user, bot, phrase, scrollbackLength, timeoutDuration)
	}

	return nil
}

func (m *nukeModule) garbageCollect() {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	now := time.Now()

	for channelID := range m.messages {
		for i := 0; i < len(m.messages[channelID]); i++ {
			diff := now.Sub(m.messages[channelID][i].timestamp)
			if diff < maxMessageAge {
				m.messages[channelID] = m.messages[channelID][i:]
				break
			}
		}
	}
}

func (m *nukeModule) nuke(source pkg.User, bot pkg.BotChannel, phrase string, scrollbackLength, timeoutDuration time.Duration) {
	if timeoutDuration > 72*time.Hour {
		timeoutDuration = 72 * time.Hour
	}

	lowercasePhrase := strings.ToLower(phrase)

	matcher := func(msg *nukeMessage) bool {
		return strings.Contains(strings.ToLower(msg.message.GetText()), lowercasePhrase)
	}

	reason := "Nuked '" + phrase + "'"

	if strings.HasPrefix(phrase, "/") && strings.HasSuffix(phrase, "/") {
		regex, err := regexp.Compile(phrase[1 : len(phrase)-1])
		if err == nil {
			reason = "Nuked r'" + phrase[1:len(phrase)-1] + "'"
			matcher = func(msg *nukeMessage) bool {
				return regex.MatchString(msg.message.GetText())
			}
		}
		// parse as regex
	}

	now := time.Now()
	timeoutDurationInSeconds := int(timeoutDuration.Seconds())

	if timeoutDurationInSeconds < 1 {
		// Timeout duration too short
		return
	}

	targets := make(map[string]pkg.User)

	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	messages := m.messages[bot.Channel().GetID()]

	for i := len(messages) - 1; i >= 0; i-- {
		diff := now.Sub(messages[i].timestamp)
		if diff > scrollbackLength {
			// We've gone far enough in the buffer, time to exit
			break
		}

		if matcher(&messages[i]) {
			targets[messages[i].user.GetID()] = messages[i].user
		}
	}

	for _, user := range targets {
		bot.SingleTimeout(user, timeoutDurationInSeconds, reason)
	}

	fmt.Printf("%s nuked %d users for the phrase %s in the last %s for %s\n", source.GetName(), len(targets), phrase, scrollbackLength, timeoutDuration)
}

func (m *nukeModule) addMessage(channel pkg.Channel, user pkg.User, message pkg.Message) {
	m.messagesMutex.Lock()
	defer m.messagesMutex.Unlock()

	m.messages[channel.GetID()] = append(m.messages[channel.GetID()], nukeMessage{
		user:      user,
		message:   message,
		timestamp: time.Now(),
	})
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package present

import (
	"fmt"
	"net/url"
	"strings"
)

func init() {
	Register("link", parseLink)
}

type Link struct {
	URL   *url.URL
	Label string
}

func (l Link) TemplateName() string { return "link" }

func parseLink(ctx *Context, fileName string, lineno int, text string) (Elem, error) {
	args := strings.Fields(text)
	url, err := url.Parse(args[1])
	if err != nil {
		return nil, err
	}
	label := ""
	if len(args) > 2 {
		label = strings.Join(args[2:], " ")
	} else {
		scheme := url.Scheme + "://"
		if url.Scheme == "mailto" {
			scheme = "mailto:"
		}
		label = strings.Replace(url.String(), scheme, "", 1)
	}
	return Link{url, label}, nil
}

func renderLink(url, text string) string {
	text = font(text)
	if text == "" {
		text = url
	}
	return fmt.Sprintf(`<a href="%s" target="_blank">%s</a>`, url, text)
}

// parseInlineLink parses an inline link at the start of s, and returns
// a rendered HTML link and the total length of the raw inline link.
// If no inline link is present, it returns all zeroes.
func parseInlineLink(s string) (link string, length int) {
	if len(s) < 2 || s[:2] != "[[" {
		return
	}
	end := strings.Index(s, "]]")
	if end == -1 {
		return
	}
	urlEnd := strings.Index(s, "]")
	url := s[2:urlEnd]
	const badURLChars = `<>"{}|\^~[] ` + "`" // per RFC1738 section 2.2
	if strings.ContainsAny(url, badURLChars) {
		return
	}
	if urlEnd == end {
		return renderLink(url, ""), end + 2
	}
	if s[urlEnd:urlEnd+2] != "][" {
		return
	}
	text := s[urlEnd+2 : end]
	return renderLink(url, text), end + 2
}
<commit_msg>go.talks/pkg/present: remove tilde from bad url chars<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package present

import (
	"fmt"
	"net/url"
	"strings"
)

func init() {
	Register("link", parseLink)
}

type Link struct {
	URL   *url.URL
	Label string
}

func (l Link) TemplateName() string { return "link" }

func parseLink(ctx *Context, fileName string, lineno int, text string) (Elem, error) {
	args := strings.Fields(text)
	url, err := url.Parse(args[1])
	if err != nil {
		return nil, err
	}
	label := ""
	if len(args) > 2 {
		label = strings.Join(args[2:], " ")
	} else {
		scheme := url.Scheme + "://"
		if url.Scheme == "mailto" {
			scheme = "mailto:"
		}
		label = strings.Replace(url.String(), scheme, "", 1)
	}
	return Link{url, label}, nil
}

func renderLink(url, text string) string {
	text = font(text)
	if text == "" {
		text = url
	}
	return fmt.Sprintf(`<a href="%s" target="_blank">%s</a>`, url, text)
}

// parseInlineLink parses an inline link at the start of s, and returns
// a rendered HTML link and the total length of the raw inline link.
// If no inline link is present, it returns all zeroes.
func parseInlineLink(s string) (link string, length int) {
	if len(s) < 2 || s[:2] != "[[" {
		return
	}
	end := strings.Index(s, "]]")
	if end == -1 {
		return
	}
	urlEnd := strings.Index(s, "]")
	url := s[2:urlEnd]
	const badURLChars = `<>"{}|\^[] ` + "`" // per RFC2396 section 2.4.3
	if strings.ContainsAny(url, badURLChars) {
		return
	}
	if urlEnd == end {
		return renderLink(url, ""), end + 2
	}
	if s[urlEnd:urlEnd+2] != "][" {
		return
	}
	text := s[urlEnd+2 : end]
	return renderLink(url, text), end + 2
}
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"bytes"
	"hash/crc32"
	"strings"

	"github.com/CodisLabs/codis/pkg/proxy/redis"
	"github.com/CodisLabs/codis/pkg/utils/errors"
)

var charmap [256]byte

func init() {
	for i := range charmap {
		c := byte(i)
		switch {
		case c >= 'A' && c <= 'Z':
			charmap[i] = c
		case c >= 'a' && c <= 'z':
			charmap[i] = c - 'a' + 'A'
		}
	}
}

type OpFlag uint32

func (f OpFlag) IsNotAllowed() bool {
	return (f & FlagNotAllow) != 0
}

func (f OpFlag) IsReadOnly() bool {
	const mask = FlagWrite | FlagMayWrite
	return (f & mask) == 0
}

func (f OpFlag) IsMasterOnly() bool {
	const mask = FlagWrite | FlagMayWrite | FlagMasterOnly
	return (f & mask) == 0
}

type OpInfo struct {
	Name string
	Flag OpFlag
}

const (
	FlagWrite = 1 << iota
	FlagMasterOnly
	FlagMayWrite
	FlagNotAllow
)

var opTable = make(map[string]OpInfo, 256)

func init() {
	for _, i := range []OpInfo{
		{"APPEND", FlagWrite},
		{"ASKING", FlagNotAllow},
		{"AUTH", 0},
		{"BGREWRITEAOF", FlagNotAllow},
		{"BGSAVE", FlagNotAllow},
		{"BITCOUNT", 0},
		{"BITFIELD", FlagWrite},
		{"BITOP", FlagWrite | FlagNotAllow},
		{"BITPOS", 0},
		{"BLPOP", FlagWrite | FlagNotAllow},
		{"BRPOP", FlagWrite | FlagNotAllow},
		{"BRPOPLPUSH", FlagWrite | FlagNotAllow},
		{"CLIENT", FlagNotAllow},
		{"CLUSTER", FlagNotAllow},
		{"COMMAND", 0},
		{"CONFIG", FlagNotAllow},
		{"DBSIZE", FlagNotAllow},
		{"DEBUG", FlagNotAllow},
		{"DECR", FlagWrite},
		{"DECRBY", FlagWrite},
		{"DEL", FlagWrite},
		{"DISCARD", FlagNotAllow},
		{"DUMP", 0},
		{"ECHO", 0},
		{"EVAL", FlagWrite},
		{"EVALSHA", FlagWrite},
		{"EXEC", FlagNotAllow},
		{"EXISTS", 0},
		{"EXPIRE", FlagWrite},
		{"EXPIREAT", FlagWrite},
		{"FLUSHALL", FlagWrite | FlagNotAllow},
		{"FLUSHDB", FlagWrite | FlagNotAllow},
		{"GEOADD", FlagWrite},
		{"GEODIST", 0},
		{"GEOHASH", 0},
		{"GEOPOS", 0},
		{"GEORADIUS", FlagWrite},
		{"GEORADIUSBYMEMBER", FlagWrite},
		{"GET", 0},
		{"GETBIT", 0},
		{"GETRANGE", 0},
		{"GETSET", FlagWrite},
		{"HDEL", FlagWrite},
		{"HEXISTS", 0},
		{"HGET", 0},
		{"HGETALL", 0},
		{"HINCRBY", FlagWrite},
		{"HINCRBYFLOAT", FlagWrite},
		{"HKEYS", 0},
		{"HLEN", 0},
		{"HMGET", 0},
		{"HMSET", FlagWrite},
		{"HSCAN", FlagMasterOnly},
		{"HSET", FlagWrite},
		{"HSETNX", FlagWrite},
		{"HSTRLEN", 0},
		{"HVALS", 0},
		{"INCR", FlagWrite},
		{"INCRBY", FlagWrite},
		{"INCRBYFLOAT", FlagWrite},
		{"INFO", 0},
		{"KEYS", FlagNotAllow},
		{"LASTSAVE", FlagNotAllow},
		{"LATENCY", FlagNotAllow},
		{"LINDEX", 0},
		{"LINSERT", FlagWrite},
		{"LLEN", 0},
		{"LPOP", FlagWrite},
		{"LPUSH", FlagWrite},
		{"LPUSHX", FlagWrite},
		{"LRANGE", 0},
		{"LREM", FlagWrite},
		{"LSET", FlagWrite},
		{"LTRIM", FlagWrite},
		{"MGET", 0},
		{"MIGRATE", FlagWrite | FlagNotAllow},
		{"MONITOR", FlagNotAllow},
		{"MOVE", FlagWrite | FlagNotAllow},
		{"MSET", FlagWrite},
		{"MSETNX", FlagWrite | FlagNotAllow},
		{"MULTI", FlagNotAllow},
		{"OBJECT", FlagNotAllow},
		{"PERSIST", FlagWrite},
		{"PEXPIRE", FlagWrite},
		{"PEXPIREAT", FlagWrite},
		{"PFADD", FlagWrite},
		{"PFCOUNT", 0},
		{"PFDEBUG", FlagWrite},
		{"PFMERGE", FlagWrite},
		{"PFSELFTEST", 0},
		{"PING", 0},
		{"PSETEX", FlagWrite},
		{"PSUBSCRIBE", FlagNotAllow},
		{"PSYNC", FlagNotAllow},
		{"PTTL", 0},
		{"PUBLISH", FlagNotAllow},
		{"PUBSUB", 0},
		{"PUNSUBSCRIBE", FlagNotAllow},
		{"QUIT", 0},
		{"RANDOMKEY", FlagNotAllow},
		{"READONLY", FlagNotAllow},
		{"READWRITE", FlagNotAllow},
		{"RENAME", FlagWrite | FlagNotAllow},
		{"RENAMENX", FlagWrite | FlagNotAllow},
		{"REPLCONF", FlagNotAllow},
		{"RESTORE", FlagWrite | FlagNotAllow},
		{"RESTORE-ASKING", FlagWrite | FlagNotAllow},
		{"ROLE", 0},
		{"RPOP", FlagWrite},
		{"RPOPLPUSH", FlagWrite},
		{"RPUSH", FlagWrite},
		{"RPUSHX", FlagWrite},
		{"SADD", FlagWrite},
		{"SAVE", FlagNotAllow},
		{"SCAN", FlagMasterOnly | FlagNotAllow},
		{"SCARD", 0},
		{"SCRIPT", FlagNotAllow},
		{"SDIFF", 0},
		{"SDIFFSTORE", FlagWrite},
		{"SELECT", 0},
		{"SET", FlagWrite},
		{"SETBIT", FlagWrite},
		{"SETEX", FlagWrite},
		{"SETNX", FlagWrite},
		{"SETRANGE", FlagWrite},
		{"SHUTDOWN", FlagNotAllow},
		{"SINTER", 0},
		{"SINTERSTORE", FlagWrite},
		{"SISMEMBER", 0},
		{"SLAVEOF", FlagNotAllow},
		{"SLOTSCHECK", FlagNotAllow},
		{"SLOTSDEL", FlagWrite | FlagNotAllow},
		{"SLOTSHASHKEY", 0},
		{"SLOTSINFO", FlagMasterOnly},
		{"SLOTSMAPPING", 0},
		{"SLOTSMGRTONE", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTSLOT", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGONE", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGSLOT", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE", FlagWrite},
		{"SLOTSMGRTONE-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTSLOT-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGONE-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGSLOT-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRT-ASYNC-FENCE", FlagNotAllow},
		{"SLOTSMGRT-ASYNC-CANCEL", FlagNotAllow},
		{"SLOTSMGRT-ASYNC-STATUS", FlagNotAllow},
		{"SLOTSMGRT-EXEC-WRAPPER", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE-ASYNC-AUTH", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE-ASYNC-ACK", FlagWrite | FlagNotAllow},
		{"SLOTSSCAN", FlagMasterOnly},
		{"SLOWLOG", FlagNotAllow},
		{"SMEMBERS", 0},
		{"SMOVE", FlagWrite},
		{"SORT", FlagWrite},
		{"SPOP", FlagWrite},
		{"SRANDMEMBER", 0},
		{"SREM", FlagWrite},
		{"SSCAN", FlagMasterOnly},
		{"STRLEN", 0},
		{"SUBSCRIBE", FlagNotAllow},
		{"SUBSTR", 0},
		{"SUNION", 0},
		{"SUNIONSTORE", FlagWrite},
		{"SYNC", FlagNotAllow},
		{"TIME", FlagNotAllow},
		{"TOUCH", FlagWrite},
		{"TTL", 0},
		{"TYPE", 0},
		{"UNSUBSCRIBE", FlagNotAllow},
		{"UNWATCH", FlagNotAllow},
		{"WAIT", FlagNotAllow},
		{"WATCH", FlagNotAllow},
		{"ZADD", FlagWrite},
		{"ZCARD", 0},
		{"ZCOUNT", 0},
		{"ZINCRBY", FlagWrite},
		{"ZINTERSTORE", FlagWrite},
		{"ZLEXCOUNT", 0},
		{"ZRANGE", 0},
		{"ZRANGEBYLEX", 0},
		{"ZRANGEBYSCORE", 0},
		{"ZRANK", 0},
		{"ZREM", FlagWrite},
		{"ZREMRANGEBYLEX", FlagWrite},
		{"ZREMRANGEBYRANK", FlagWrite},
		{"ZREMRANGEBYSCORE", FlagWrite},
		{"ZREVRANGE", 0},
		{"ZREVRANGEBYLEX", 0},
		{"ZREVRANGEBYSCORE", 0},
		{"ZREVRANK", 0},
		{"ZSCAN", FlagMasterOnly},
		{"ZSCORE", 0},
		{"ZUNIONSTORE", FlagWrite},
	} {
		opTable[i.Name] = i
	}
}

var (
	ErrBadMultiBulk = errors.New("bad multi-bulk for command")
	ErrBadOpStrLen  = errors.New("bad command length, too short or too long")
)

const MaxOpStrLen = 64

func getOpInfo(multi []*redis.Resp) (string, OpFlag, error) {
	if len(multi) < 1 {
		return "", 0, ErrBadMultiBulk
	}

	var upper [MaxOpStrLen]byte

	var op = multi[0].Value
	if len(op) == 0 || len(op) > len(upper) {
		return "", 0, ErrBadOpStrLen
	}
	for i := range op {
		if c := charmap[op[i]]; c != 0 {
			upper[i] = c
		} else {
			return strings.ToUpper(string(op)), FlagMayWrite, nil
		}
	}
	op = upper[:len(op)]
	if r, ok := opTable[string(op)]; ok {
		return r.Name, r.Flag, nil
	}
	return string(op), FlagMayWrite, nil
}

func Hash(key []byte) uint32 {
	const (
		TagBeg = '{'
		TagEnd = '}'
	)
	if beg := bytes.IndexByte(key, TagBeg); beg >= 0 {
		if end := bytes.IndexByte(key[beg+1:], TagEnd); end >= 0 {
			key = key[beg+1 : beg+1+end]
		}
	}
	return crc32.ChecksumIEEE(key)
}

func getHashKey(multi []*redis.Resp, opstr string) []byte {
	var index = 1
	switch opstr {
	case "ZINTERSTORE", "ZUNIONSTORE", "EVAL", "EVALSHA":
		index = 3
	}
	if index < len(multi) {
		return multi[index].Value
	}
	return nil
}
<commit_msg>proxy: fix IsMasterOnly<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"bytes"
	"hash/crc32"
	"strings"

	"github.com/CodisLabs/codis/pkg/proxy/redis"
	"github.com/CodisLabs/codis/pkg/utils/errors"
)

var charmap [256]byte

func init() {
	for i := range charmap {
		c := byte(i)
		switch {
		case c >= 'A' && c <= 'Z':
			charmap[i] = c
		case c >= 'a' && c <= 'z':
			charmap[i] = c - 'a' + 'A'
		}
	}
}

type OpFlag uint32

func (f OpFlag) IsNotAllowed() bool {
	return (f & FlagNotAllow) != 0
}

func (f OpFlag) IsReadOnly() bool {
	const mask = FlagWrite | FlagMayWrite
	return (f & mask) == 0
}

func (f OpFlag) IsMasterOnly() bool {
	const mask = FlagWrite | FlagMayWrite | FlagMasterOnly
	return (f & mask) != 0
}

type OpInfo struct {
	Name string
	Flag OpFlag
}

const (
	FlagWrite = 1 << iota
	FlagMasterOnly
	FlagMayWrite
	FlagNotAllow
)

var opTable = make(map[string]OpInfo, 256)

func init() {
	for _, i := range []OpInfo{
		{"APPEND", FlagWrite},
		{"ASKING", FlagNotAllow},
		{"AUTH", 0},
		{"BGREWRITEAOF", FlagNotAllow},
		{"BGSAVE", FlagNotAllow},
		{"BITCOUNT", 0},
		{"BITFIELD", FlagWrite},
		{"BITOP", FlagWrite | FlagNotAllow},
		{"BITPOS", 0},
		{"BLPOP", FlagWrite | FlagNotAllow},
		{"BRPOP", FlagWrite | FlagNotAllow},
		{"BRPOPLPUSH", FlagWrite | FlagNotAllow},
		{"CLIENT", FlagNotAllow},
		{"CLUSTER", FlagNotAllow},
		{"COMMAND", 0},
		{"CONFIG", FlagNotAllow},
		{"DBSIZE", FlagNotAllow},
		{"DEBUG", FlagNotAllow},
		{"DECR", FlagWrite},
		{"DECRBY", FlagWrite},
		{"DEL", FlagWrite},
		{"DISCARD", FlagNotAllow},
		{"DUMP", 0},
		{"ECHO", 0},
		{"EVAL", FlagWrite},
		{"EVALSHA", FlagWrite},
		{"EXEC", FlagNotAllow},
		{"EXISTS", 0},
		{"EXPIRE", FlagWrite},
		{"EXPIREAT", FlagWrite},
		{"FLUSHALL", FlagWrite | FlagNotAllow},
		{"FLUSHDB", FlagWrite | FlagNotAllow},
		{"GEOADD", FlagWrite},
		{"GEODIST", 0},
		{"GEOHASH", 0},
		{"GEOPOS", 0},
		{"GEORADIUS", FlagWrite},
		{"GEORADIUSBYMEMBER", FlagWrite},
		{"GET", 0},
		{"GETBIT", 0},
		{"GETRANGE", 0},
		{"GETSET", FlagWrite},
		{"HDEL", FlagWrite},
		{"HEXISTS", 0},
		{"HGET", 0},
		{"HGETALL", 0},
		{"HINCRBY", FlagWrite},
		{"HINCRBYFLOAT", FlagWrite},
		{"HKEYS", 0},
		{"HLEN", 0},
		{"HMGET", 0},
		{"HMSET", FlagWrite},
		{"HSCAN", FlagMasterOnly},
		{"HSET", FlagWrite},
		{"HSETNX", FlagWrite},
		{"HSTRLEN", 0},
		{"HVALS", 0},
		{"INCR", FlagWrite},
		{"INCRBY", FlagWrite},
		{"INCRBYFLOAT", FlagWrite},
		{"INFO", 0},
		{"KEYS", FlagNotAllow},
		{"LASTSAVE", FlagNotAllow},
		{"LATENCY", FlagNotAllow},
		{"LINDEX", 0},
		{"LINSERT", FlagWrite},
		{"LLEN", 0},
		{"LPOP", FlagWrite},
		{"LPUSH", FlagWrite},
		{"LPUSHX", FlagWrite},
		{"LRANGE", 0},
		{"LREM", FlagWrite},
		{"LSET", FlagWrite},
		{"LTRIM", FlagWrite},
		{"MGET", 0},
		{"MIGRATE", FlagWrite | FlagNotAllow},
		{"MONITOR", FlagNotAllow},
		{"MOVE", FlagWrite | FlagNotAllow},
		{"MSET", FlagWrite},
		{"MSETNX", FlagWrite | FlagNotAllow},
		{"MULTI", FlagNotAllow},
		{"OBJECT", FlagNotAllow},
		{"PERSIST", FlagWrite},
		{"PEXPIRE", FlagWrite},
		{"PEXPIREAT", FlagWrite},
		{"PFADD", FlagWrite},
		{"PFCOUNT", 0},
		{"PFDEBUG", FlagWrite},
		{"PFMERGE", FlagWrite},
		{"PFSELFTEST", 0},
		{"PING", 0},
		{"PSETEX", FlagWrite},
		{"PSUBSCRIBE", FlagNotAllow},
		{"PSYNC", FlagNotAllow},
		{"PTTL", 0},
		{"PUBLISH", FlagNotAllow},
		{"PUBSUB", 0},
		{"PUNSUBSCRIBE", FlagNotAllow},
		{"QUIT", 0},
		{"RANDOMKEY", FlagNotAllow},
		{"READONLY", FlagNotAllow},
		{"READWRITE", FlagNotAllow},
		{"RENAME", FlagWrite | FlagNotAllow},
		{"RENAMENX", FlagWrite | FlagNotAllow},
		{"REPLCONF", FlagNotAllow},
		{"RESTORE", FlagWrite | FlagNotAllow},
		{"RESTORE-ASKING", FlagWrite | FlagNotAllow},
		{"ROLE", 0},
		{"RPOP", FlagWrite},
		{"RPOPLPUSH", FlagWrite},
		{"RPUSH", FlagWrite},
		{"RPUSHX", FlagWrite},
		{"SADD", FlagWrite},
		{"SAVE", FlagNotAllow},
		{"SCAN", FlagMasterOnly | FlagNotAllow},
		{"SCARD", 0},
		{"SCRIPT", FlagNotAllow},
		{"SDIFF", 0},
		{"SDIFFSTORE", FlagWrite},
		{"SELECT", 0},
		{"SET", FlagWrite},
		{"SETBIT", FlagWrite},
		{"SETEX", FlagWrite},
		{"SETNX", FlagWrite},
		{"SETRANGE", FlagWrite},
		{"SHUTDOWN", FlagNotAllow},
		{"SINTER", 0},
		{"SINTERSTORE", FlagWrite},
		{"SISMEMBER", 0},
		{"SLAVEOF", FlagNotAllow},
		{"SLOTSCHECK", FlagNotAllow},
		{"SLOTSDEL", FlagWrite | FlagNotAllow},
		{"SLOTSHASHKEY", 0},
		{"SLOTSINFO", FlagMasterOnly},
		{"SLOTSMAPPING", 0},
		{"SLOTSMGRTONE", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTSLOT", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGONE", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGSLOT", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE", FlagWrite},
		{"SLOTSMGRTONE-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTSLOT-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGONE-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRTTAGSLOT-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSMGRT-ASYNC-FENCE", FlagNotAllow},
		{"SLOTSMGRT-ASYNC-CANCEL", FlagNotAllow},
		{"SLOTSMGRT-ASYNC-STATUS", FlagNotAllow},
		{"SLOTSMGRT-EXEC-WRAPPER", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE-ASYNC", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE-ASYNC-AUTH", FlagWrite | FlagNotAllow},
		{"SLOTSRESTORE-ASYNC-ACK", FlagWrite | FlagNotAllow},
		{"SLOTSSCAN", FlagMasterOnly},
		{"SLOWLOG", FlagNotAllow},
		{"SMEMBERS", 0},
		{"SMOVE", FlagWrite},
		{"SORT", FlagWrite},
		{"SPOP", FlagWrite},
		{"SRANDMEMBER", 0},
		{"SREM", FlagWrite},
		{"SSCAN", FlagMasterOnly},
		{"STRLEN", 0},
		{"SUBSCRIBE", FlagNotAllow},
		{"SUBSTR", 0},
		{"SUNION", 0},
		{"SUNIONSTORE", FlagWrite},
		{"SYNC", FlagNotAllow},
		{"TIME", FlagNotAllow},
		{"TOUCH", FlagWrite},
		{"TTL", 0},
		{"TYPE", 0},
		{"UNSUBSCRIBE", FlagNotAllow},
		{"UNWATCH", FlagNotAllow},
		{"WAIT", FlagNotAllow},
		{"WATCH", FlagNotAllow},
		{"ZADD", FlagWrite},
		{"ZCARD", 0},
		{"ZCOUNT", 0},
		{"ZINCRBY", FlagWrite},
		{"ZINTERSTORE", FlagWrite},
		{"ZLEXCOUNT", 0},
		{"ZRANGE", 0},
		{"ZRANGEBYLEX", 0},
		{"ZRANGEBYSCORE", 0},
		{"ZRANK", 0},
		{"ZREM", FlagWrite},
		{"ZREMRANGEBYLEX", FlagWrite},
		{"ZREMRANGEBYRANK", FlagWrite},
		{"ZREMRANGEBYSCORE", FlagWrite},
		{"ZREVRANGE", 0},
		{"ZREVRANGEBYLEX", 0},
		{"ZREVRANGEBYSCORE", 0},
		{"ZREVRANK", 0},
		{"ZSCAN", FlagMasterOnly},
		{"ZSCORE", 0},
		{"ZUNIONSTORE", FlagWrite},
	} {
		opTable[i.Name] = i
	}
}

var (
	ErrBadMultiBulk = errors.New("bad multi-bulk for command")
	ErrBadOpStrLen  = errors.New("bad command length, too short or too long")
)

const MaxOpStrLen = 64

func getOpInfo(multi []*redis.Resp) (string, OpFlag, error) {
	if len(multi) < 1 {
		return "", 0, ErrBadMultiBulk
	}

	var upper [MaxOpStrLen]byte

	var op = multi[0].Value
	if len(op) == 0 || len(op) > len(upper) {
		return "", 0, ErrBadOpStrLen
	}
	for i := range op {
		if c := charmap[op[i]]; c != 0 {
			upper[i] = c
		} else {
			return strings.ToUpper(string(op)), FlagMayWrite, nil
		}
	}
	op = upper[:len(op)]
	if r, ok := opTable[string(op)]; ok {
		return r.Name, r.Flag, nil
	}
	return string(op), FlagMayWrite, nil
}

func Hash(key []byte) uint32 {
	const (
		TagBeg = '{'
		TagEnd = '}'
	)
	if beg := bytes.IndexByte(key, TagBeg); beg >= 0 {
		if end := bytes.IndexByte(key[beg+1:], TagEnd); end >= 0 {
			key = key[beg+1 : beg+1+end]
		}
	}
	return crc32.ChecksumIEEE(key)
}

func getHashKey(multi []*redis.Resp, opstr string) []byte {
	var index = 1
	switch opstr {
	case "ZINTERSTORE", "ZUNIONSTORE", "EVAL", "EVALSHA":
		index = 3
	}
	if index < len(multi) {
		return multi[index].Value
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"sync"

	"github.com/CodisLabs/codis/pkg/models"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
)

type Router struct {
	mu sync.RWMutex

	pool map[string]*SharedBackendConn

	slots [models.MaxSlotNum]Slot

	dispFunc

	config *Config
	online bool
	closed bool
}

func NewRouter(config *Config) *Router {
	s := &Router{config: config}
	s.pool = make(map[string]*SharedBackendConn)
	for i := range s.slots {
		s.slots[i].id = i
	}
	if config.BackendReadReplica {
		s.dispFunc = dispReadReplica
	}
	return s
}

func (s *Router) Start() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.online = true
}

func (s *Router) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.closed = true

	for i := range s.slots {
		s.fillSlot(&models.Slot{Id: i})
	}
}

func (s *Router) GetSlots() []*models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	slots := make([]*models.Slot, len(s.slots))
	for i := range s.slots {
		slots[i] = s.slots[i].model()
	}
	return slots
}

func (s *Router) GetSlot(id int) *models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if id < 0 || id >= len(s.slots) {
		return nil
	}
	slot := &s.slots[id]
	return slot.model()
}

var (
	ErrClosedRouter  = errors.New("use of closed router")
	ErrInvalidSlotId = errors.New("use of invalid slot id")
)

func (s *Router) FillSlot(m *models.Slot) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	return s.fillSlot(m)
}

func (s *Router) KeepAlive() error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return ErrClosedRouter
	}
	for _, bc := range s.pool {
		bc.KeepAlive()
	}
	return nil
}

func (s *Router) isOnline() bool {
	return s.online && !s.closed
}

func (s *Router) dispatch(r *Request) error {
	hkey := getHashKey(r.Multi, r.OpStr)
	var id = Hash(hkey) % uint32(len(s.slots))
	slot := &s.slots[id]
	return slot.forward(s.dispFunc, r, hkey)
}

func (s *Router) dispatchSlot(r *Request, id int) error {
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	return slot.forward(s.dispFunc, r, nil)
}

func (s *Router) dispatchAddr(r *Request, addr string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	bc := s.getBackendConn(addr, false)
	if bc == nil {
		return false
	} else {
		bc.PushBack(r)
		s.putBackendConn(bc)
		return true
	}
}

func (s *Router) getBackendConn(addr string, create bool) *SharedBackendConn {
	if bc := s.pool[addr]; bc != nil {
		return bc.Retain()
	} else if create {
		bc := NewSharedBackendConn(addr, s.config)
		s.pool[addr] = bc
		return bc
	} else {
		return nil
	}
}

func (s *Router) putBackendConn(bc *SharedBackendConn) {
	if bc != nil && bc.Release() {
		delete(s.pool, bc.Addr())
	}
}

func (s *Router) fillSlot(m *models.Slot) error {
	id := m.Id
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend)
	s.putBackendConn(slot.migrate)
	for i := range slot.replica {
		for _, bc := range slot.replica[i] {
			s.putBackendConn(bc)
		}
	}
	slot.backend = nil
	slot.migrate = nil
	slot.replica = nil

	if addr := m.BackendAddr; len(addr) != 0 {
		slot.backend = s.getBackendConn(addr, true)
	}
	if from := m.MigrateFrom; len(from) != 0 {
		slot.migrate = s.getBackendConn(from, true)
	}
	for i := range m.ReplicaGroup {
		var group []*SharedBackendConn
		for _, addr := range m.ReplicaGroup[i] {
			group = append(group, s.getBackendConn(addr, true))
		}
		slot.replica = append(slot.replica, group)
	}

	if !m.Locked {
		slot.unblock()
	}
	if !s.closed {
		if slot.migrate != nil {
			log.Warnf("fill slot %04d, backend.addr = %s, migrate.from = %s, locked = %t",
				id, slot.backend.Addr(), slot.migrate.Addr(), slot.lock.hold)
		} else {
			log.Warnf("fill slot %04d, backend.addr = %s, locked = %t",
				id, slot.backend.Addr(), slot.lock.hold)
		}
	}
	return nil
}
<commit_msg>router: don't fill replica slice if dispFunc is nil<commit_after>// Copyright 2016 CodisLabs. All Rights Reserved.
// Licensed under the MIT (MIT-LICENSE.txt) license.

package proxy

import (
	"sync"

	"github.com/CodisLabs/codis/pkg/models"
	"github.com/CodisLabs/codis/pkg/utils/errors"
	"github.com/CodisLabs/codis/pkg/utils/log"
)

type Router struct {
	mu sync.RWMutex

	pool map[string]*SharedBackendConn

	slots [models.MaxSlotNum]Slot

	dispFunc

	config *Config
	online bool
	closed bool
}

func NewRouter(config *Config) *Router {
	s := &Router{config: config}
	s.pool = make(map[string]*SharedBackendConn)
	for i := range s.slots {
		s.slots[i].id = i
	}
	if config.BackendReadReplica {
		s.dispFunc = dispReadReplica
	}
	return s
}

func (s *Router) Start() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.online = true
}

func (s *Router) Close() {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return
	}
	s.closed = true

	for i := range s.slots {
		s.fillSlot(&models.Slot{Id: i})
	}
}

func (s *Router) GetSlots() []*models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	slots := make([]*models.Slot, len(s.slots))
	for i := range s.slots {
		slots[i] = s.slots[i].model()
	}
	return slots
}

func (s *Router) GetSlot(id int) *models.Slot {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if id < 0 || id >= len(s.slots) {
		return nil
	}
	slot := &s.slots[id]
	return slot.model()
}

var (
	ErrClosedRouter  = errors.New("use of closed router")
	ErrInvalidSlotId = errors.New("use of invalid slot id")
)

func (s *Router) FillSlot(m *models.Slot) error {
	s.mu.Lock()
	defer s.mu.Unlock()
	if s.closed {
		return ErrClosedRouter
	}
	return s.fillSlot(m)
}

func (s *Router) KeepAlive() error {
	s.mu.RLock()
	defer s.mu.RUnlock()
	if s.closed {
		return ErrClosedRouter
	}
	for _, bc := range s.pool {
		bc.KeepAlive()
	}
	return nil
}

func (s *Router) isOnline() bool {
	return s.online && !s.closed
}

func (s *Router) dispatch(r *Request) error {
	hkey := getHashKey(r.Multi, r.OpStr)
	var id = Hash(hkey) % uint32(len(s.slots))
	slot := &s.slots[id]
	return slot.forward(s.dispFunc, r, hkey)
}

func (s *Router) dispatchSlot(r *Request, id int) error {
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	return slot.forward(s.dispFunc, r, nil)
}

func (s *Router) dispatchAddr(r *Request, addr string) bool {
	s.mu.RLock()
	defer s.mu.RUnlock()
	bc := s.getBackendConn(addr, false)
	if bc == nil {
		return false
	} else {
		bc.PushBack(r)
		s.putBackendConn(bc)
		return true
	}
}

func (s *Router) getBackendConn(addr string, create bool) *SharedBackendConn {
	if bc := s.pool[addr]; bc != nil {
		return bc.Retain()
	} else if create {
		bc := NewSharedBackendConn(addr, s.config)
		s.pool[addr] = bc
		return bc
	} else {
		return nil
	}
}

func (s *Router) putBackendConn(bc *SharedBackendConn) {
	if bc != nil && bc.Release() {
		delete(s.pool, bc.Addr())
	}
}

func (s *Router) fillSlot(m *models.Slot) error {
	id := m.Id
	if id < 0 || id >= len(s.slots) {
		return ErrInvalidSlotId
	}
	slot := &s.slots[id]
	slot.blockAndWait()

	s.putBackendConn(slot.backend)
	s.putBackendConn(slot.migrate)
	for i := range slot.replica {
		for _, bc := range slot.replica[i] {
			s.putBackendConn(bc)
		}
	}
	slot.backend = nil
	slot.migrate = nil
	slot.replica = nil

	if addr := m.BackendAddr; len(addr) != 0 {
		slot.backend = s.getBackendConn(addr, true)
	}
	if from := m.MigrateFrom; len(from) != 0 {
		slot.migrate = s.getBackendConn(from, true)
	}
	if s.dispFunc != nil {
		for i := range m.ReplicaGroup {
			var group []*SharedBackendConn
			for _, addr := range m.ReplicaGroup[i] {
				group = append(group, s.getBackendConn(addr, true))
			}
			slot.replica = append(slot.replica, group)
		}
	}

	if !m.Locked {
		slot.unblock()
	}
	if !s.closed {
		if slot.migrate != nil {
			log.Warnf("fill slot %04d, backend.addr = %s, migrate.from = %s, locked = %t",
				id, slot.backend.Addr(), slot.migrate.Addr(), slot.lock.hold)
		} else {
			log.Warnf("fill slot %04d, backend.addr = %s, locked = %t",
				id, slot.backend.Addr(), slot.lock.hold)
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package workflow

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"

	"github.com/go-gorp/gorp"

	"github.com/ovh/cds/engine/api/application"
	"github.com/ovh/cds/engine/api/cache"
	"github.com/ovh/cds/engine/api/environment"
	"github.com/ovh/cds/engine/api/pipeline"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/exportentities"
)

// Export a workflow
func Export(db gorp.SqlExecutor, cache cache.Store, key string, name string, f exportentities.Format, withPermissions bool, u *sdk.User, w io.Writer) (int, error) {
	wf, errload := Load(db, cache, key, name, u, LoadOptions{})
	if errload != nil {
		return 0, sdk.WrapError(errload, "workflow.Export> Cannot load workflow %s", name)
	}

	return exportWorkflow(*wf, f, withPermissions, w)
}

func exportWorkflow(wf sdk.Workflow, f exportentities.Format, withPermissions bool, w io.Writer) (int, error) {
	e, err := exportentities.NewWorkflow(wf, withPermissions)
	if err != nil {
		return 0, err
	}

	// Marshal to the desired format
	b, err := exportentities.Marshal(e, f)
	if err != nil {
		return 0, sdk.WrapError(err, "workflow.Export>")
	}

	return w.Write(b)
}

// Pull a workflow with all it dependencies; it writes a tar buffer in the writer
func Pull(db gorp.SqlExecutor, cache cache.Store, key string, name string, f exportentities.Format, withPermissions bool, encryptFunc sdk.EncryptFunc, u *sdk.User, w io.Writer) error {
	options := LoadOptions{
		DeepPipeline: true,
	}
	wf, errload := Load(db, cache, key, name, u, options)
	if errload != nil {
		return sdk.WrapError(errload, "workflow.Pull> Cannot load workflow %s", name)
	}

	apps := wf.GetApplications()
	envs := wf.GetEnvironments()
	pips := wf.GetPipelines()

	tw := tar.NewWriter(w)

	buffw := new(bytes.Buffer)
	size, errw := exportWorkflow(*wf, f, withPermissions, buffw)
	if errw != nil {
		tw.Close()
		return sdk.WrapError(errw, "workflow.Pull> Unable to export workflow")
	}

	hdr := &tar.Header{
		Name: fmt.Sprintf("%s.yml", wf.Name),
		Mode: 0644,
		Size: int64(size),
	}
	if err := tw.WriteHeader(hdr); err != nil {
		tw.Close()
		return sdk.WrapError(err, "workflow.Pull> Unable to write workflow header %+v", hdr)
	}
	if _, err := io.Copy(tw, buffw); err != nil {
		tw.Close()
		return sdk.WrapError(err, "workflow.Pull> Unable to copy workflow buffer")
	}

	for _, a := range apps {
		buff := new(bytes.Buffer)
		size, err := application.ExportApplication(db, a, f, withPermissions, encryptFunc, buff)
		if err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to export app %s", a.Name)
		}
		hdr := &tar.Header{
			Name: fmt.Sprintf("%s.app.yml", a.Name),
			Mode: 0644,
			Size: int64(size),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to write app header %+v", hdr)
		}
		if _, err := io.Copy(tw, buff); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to copy app buffer")
		}
	}

	for _, e := range envs {
		buff := new(bytes.Buffer)
		size, err := environment.ExportEnvironment(db, e, f, withPermissions, encryptFunc, buff)
		if err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to export env %s", e.Name)
		}

		hdr := &tar.Header{
			Name: fmt.Sprintf("%s.env.yml", e.Name),
			Mode: 0644,
			Size: int64(size),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to write env header %+v", hdr)
		}
		if _, err := io.Copy(tw, buff); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to copy env buffer")
		}
	}

	for _, p := range pips {
		buff := new(bytes.Buffer)
		size, err := pipeline.ExportPipeline(p, f, withPermissions, buff)
		if err != nil {
			return sdk.WrapError(err, "workflow.Pull> Unable to export pipeline %s", p.Name)
		}
		hdr := &tar.Header{
			Name: fmt.Sprintf("%s.pip.yml", p.Name),
			Mode: 0644,
			Size: int64(size),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to write pipeline header %+v", hdr)
		}
		if _, err := io.Copy(tw, buff); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to copy pip buffer")
		}
	}

	if err := tw.Close(); err != nil {
		return sdk.WrapError(err, "workflow.Pull> Unable to close tar writer")
	}
	return nil
}
<commit_msg>fix (api): workflow pull should not encrypt secret placeholder (#1867)<commit_after>package workflow

import (
	"archive/tar"
	"bytes"
	"fmt"
	"io"

	"github.com/go-gorp/gorp"

	"github.com/ovh/cds/engine/api/application"
	"github.com/ovh/cds/engine/api/cache"
	"github.com/ovh/cds/engine/api/environment"
	"github.com/ovh/cds/engine/api/pipeline"
	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/exportentities"
)

// Export a workflow
func Export(db gorp.SqlExecutor, cache cache.Store, key string, name string, f exportentities.Format, withPermissions bool, u *sdk.User, w io.Writer) (int, error) {
	wf, errload := Load(db, cache, key, name, u, LoadOptions{})
	if errload != nil {
		return 0, sdk.WrapError(errload, "workflow.Export> Cannot load workflow %s", name)
	}

	return exportWorkflow(*wf, f, withPermissions, w)
}

func exportWorkflow(wf sdk.Workflow, f exportentities.Format, withPermissions bool, w io.Writer) (int, error) {
	e, err := exportentities.NewWorkflow(wf, withPermissions)
	if err != nil {
		return 0, err
	}

	// Marshal to the desired format
	b, err := exportentities.Marshal(e, f)
	if err != nil {
		return 0, sdk.WrapError(err, "workflow.Export>")
	}

	return w.Write(b)
}

// Pull a workflow with all it dependencies; it writes a tar buffer in the writer
func Pull(db gorp.SqlExecutor, cache cache.Store, key string, name string, f exportentities.Format, withPermissions bool, encryptFunc sdk.EncryptFunc, u *sdk.User, w io.Writer) error {
	options := LoadOptions{
		DeepPipeline: true,
	}
	wf, errload := Load(db, cache, key, name, u, options)
	if errload != nil {
		return sdk.WrapError(errload, "workflow.Pull> Cannot load workflow %s", name)
	}

	apps := wf.GetApplications()
	envs := wf.GetEnvironments()
	pips := wf.GetPipelines()

	//Reload app to retrieve secrets
	for i := range apps {
		app := &apps[i]
		vars, errv := application.GetAllVariable(db, key, app.Name, application.WithClearPassword())
		if errv != nil {
			return sdk.WrapError(errv, "workflow.Pull> Cannot load application variables %s", app.Name)
		}
		app.Variable = vars

		if errk := application.LoadAllKeys(db, app); errk != nil {
			return sdk.WrapError(errk, "workflow.Pull> Cannot load application keys %s", app.Name)
		}
	}

	//Reload env to retrieve secrets
	for i := range envs {
		env := &envs[i]
		vars, errv := environment.GetAllVariable(db, key, env.Name, environment.WithClearPassword())
		if errv != nil {
			return sdk.WrapError(errv, "workflow.Pull> Cannot load environment variables %s", env.Name)
		}
		env.Variable = vars

		if errk := environment.LoadAllKeys(db, env); errk != nil {
			return sdk.WrapError(errk, "workflow.Pull> Cannot load environment keys %s", env.Name)
		}
	}

	tw := tar.NewWriter(w)

	buffw := new(bytes.Buffer)
	size, errw := exportWorkflow(*wf, f, withPermissions, buffw)
	if errw != nil {
		tw.Close()
		return sdk.WrapError(errw, "workflow.Pull> Unable to export workflow")
	}

	hdr := &tar.Header{
		Name: fmt.Sprintf("%s.yml", wf.Name),
		Mode: 0644,
		Size: int64(size),
	}
	if err := tw.WriteHeader(hdr); err != nil {
		tw.Close()
		return sdk.WrapError(err, "workflow.Pull> Unable to write workflow header %+v", hdr)
	}
	if _, err := io.Copy(tw, buffw); err != nil {
		tw.Close()
		return sdk.WrapError(err, "workflow.Pull> Unable to copy workflow buffer")
	}

	for _, a := range apps {
		buff := new(bytes.Buffer)
		size, err := application.ExportApplication(db, a, f, withPermissions, encryptFunc, buff)
		if err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to export app %s", a.Name)
		}
		hdr := &tar.Header{
			Name: fmt.Sprintf("%s.app.yml", a.Name),
			Mode: 0644,
			Size: int64(size),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to write app header %+v", hdr)
		}
		if _, err := io.Copy(tw, buff); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to copy app buffer")
		}
	}

	for _, e := range envs {
		buff := new(bytes.Buffer)
		size, err := environment.ExportEnvironment(db, e, f, withPermissions, encryptFunc, buff)
		if err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to export env %s", e.Name)
		}

		hdr := &tar.Header{
			Name: fmt.Sprintf("%s.env.yml", e.Name),
			Mode: 0644,
			Size: int64(size),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to write env header %+v", hdr)
		}
		if _, err := io.Copy(tw, buff); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to copy env buffer")
		}
	}

	for _, p := range pips {
		buff := new(bytes.Buffer)
		size, err := pipeline.ExportPipeline(p, f, withPermissions, buff)
		if err != nil {
			return sdk.WrapError(err, "workflow.Pull> Unable to export pipeline %s", p.Name)
		}
		hdr := &tar.Header{
			Name: fmt.Sprintf("%s.pip.yml", p.Name),
			Mode: 0644,
			Size: int64(size),
		}
		if err := tw.WriteHeader(hdr); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to write pipeline header %+v", hdr)
		}
		if _, err := io.Copy(tw, buff); err != nil {
			tw.Close()
			return sdk.WrapError(err, "workflow.Pull> Unable to copy pip buffer")
		}
	}

	if err := tw.Close(); err != nil {
		return sdk.WrapError(err, "workflow.Pull> Unable to close tar writer")
	}
	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

func runArtifactUpload(w *currentWorker) BuiltInAction {
	if w.currentJob.wJob == nil {
		return func(ctx context.Context, a *sdk.Action, buildID int64, params *[]sdk.Parameter, sendLog LoggerFunc) sdk.Result {
			res := sdk.Result{Status: sdk.StatusSuccess.String()}

			pipeline := sdk.ParameterValue(*params, "cds.pipeline")
			project := sdk.ParameterValue(*params, "cds.project")
			application := sdk.ParameterValue(*params, "cds.application")
			environment := sdk.ParameterValue(*params, "cds.environment")
			buildNumberString := sdk.ParameterValue(*params, "cds.buildNumber")

			path := sdk.ParameterValue(a.Parameters, "path")
			if path == "" {
				path = "."
			}

			tag := sdk.ParameterFind(a.Parameters, "tag")
			if tag == nil {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("tag variable is empty. aborting")
				sendLog(res.Reason)
				return res
			}
			tag.Value = strings.Replace(tag.Value, "/", "-", -1)
			tag.Value = url.QueryEscape(tag.Value)

			// Global all files matching filePath
			filesPath, err := filepath.Glob(path)
			if err != nil {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("cannot perform globbing of pattern '%s': %s", path, err)
				sendLog(res.Reason)
				return res
			}

			if len(filesPath) == 0 {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("Pattern '%s' matched no file", path)
				sendLog(res.Reason)
				return res
			}

			buildNumber, errBN := strconv.Atoi(buildNumberString)
			if errBN != nil {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("BuilNumber is not an integer %s", errBN)
				sendLog(res.Reason)
				return res
			}

			for _, filePath := range filesPath {
				filename := filepath.Base(filePath)
				throughTempURL, duration, err := sdk.UploadArtifact(project, pipeline, application, tag.Value, filePath, buildNumber, environment)
				if throughTempURL {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to object store", filename, duration.Seconds()))
				} else {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to CDS API", filename, duration.Seconds()))
				}
				if err != nil {
					res.Status = sdk.StatusFail.String()
					res.Reason = fmt.Sprintf("Error while uploading artifact '%s': %v", filename, err)
					sendLog(res.Reason)
					return res
				}
			}

			return res
		}
	}

	return func(ctx context.Context, a *sdk.Action, buildID int64, params *[]sdk.Parameter, sendLog LoggerFunc) sdk.Result {
		res := sdk.Result{Status: sdk.StatusSuccess.String()}

		path := sdk.ParameterValue(a.Parameters, "path")
		if path == "" {
			path = "."
		}

		tag := sdk.ParameterFind(a.Parameters, "tag")
		if tag == nil {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("tag variable is empty. aborting")
			sendLog(res.Reason)
			return res
		}
		tag.Value = strings.Replace(tag.Value, "/", "-", -1)
		tag.Value = url.QueryEscape(tag.Value)

		// Global all files matching filePath
		filesPath, err := filepath.Glob(path)
		if err != nil {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("cannot perform globbing of pattern '%s': %s", path, err)
			sendLog(res.Reason)
			return res
		}

		if len(filesPath) == 0 {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("Pattern '%s' matched no file", path)
			sendLog(res.Reason)
			return res
		}

		var globalError = &sdk.MultiError{}
		var chanError = make(chan error)
		var wg = new(sync.WaitGroup)
		var wgErrors = new(sync.WaitGroup)

		go func() {
			for err := range chanError {
				sendLog(err.Error())
				globalError.Append(err)
				wgErrors.Done()
			}
		}()

		wg.Add(len(filesPath))
		for _, p := range filesPath {
			filename := filepath.Base(p)
			go func(path string) {
				log.Debug("Uploading %s", path)
				defer wg.Done()
				throughTempURL, duration, err := w.client.QueueArtifactUpload(buildID, tag.Value, path)
				if err != nil {
					chanError <- sdk.WrapError(err, "Error while uploading artifact %s", path)
					wgErrors.Add(1)
					return
				}
				if throughTempURL {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to object store", filename, duration.Seconds()))
				} else {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to CDS API", filename, duration.Seconds()))
				}
			}(p)
			if len(filesPath) > 1 {
				//Wait 3 second to get the object storage to set up all the things
				time.Sleep(3 * time.Second)
			}
		}
		wg.Wait()
		close(chanError)
		<-chanError
		wgErrors.Wait()

		if !globalError.IsEmpty() {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("Error: %v", globalError.Error())
			return res
		}

		return res
	}
}
<commit_msg>fix: check tag variable (#1882)<commit_after>package main

import (
	"context"
	"fmt"
	"net/url"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/ovh/cds/sdk"
	"github.com/ovh/cds/sdk/log"
)

func runArtifactUpload(w *currentWorker) BuiltInAction {
	if w.currentJob.wJob == nil {
		return func(ctx context.Context, a *sdk.Action, buildID int64, params *[]sdk.Parameter, sendLog LoggerFunc) sdk.Result {
			res := sdk.Result{Status: sdk.StatusSuccess.String()}

			pipeline := sdk.ParameterValue(*params, "cds.pipeline")
			project := sdk.ParameterValue(*params, "cds.project")
			application := sdk.ParameterValue(*params, "cds.application")
			environment := sdk.ParameterValue(*params, "cds.environment")
			buildNumberString := sdk.ParameterValue(*params, "cds.buildNumber")

			path := sdk.ParameterValue(a.Parameters, "path")
			if path == "" {
				path = "."
			}

			tag := sdk.ParameterFind(a.Parameters, "tag")
			if tag == nil {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("tag variable is empty. aborting")
				sendLog(res.Reason)
				return res
			}
			if strings.Contains(tag.Value, "{") || strings.Contains(tag.Value, "}") || strings.Contains(tag.Value, " ") {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("tag variable invalid: %s", tag.Value)
				sendLog(res.Reason)
				return res
			}
			tag.Value = strings.Replace(tag.Value, "/", "-", -1)
			tag.Value = url.QueryEscape(tag.Value)

			// Global all files matching filePath
			filesPath, err := filepath.Glob(path)
			if err != nil {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("cannot perform globbing of pattern '%s': %s", path, err)
				sendLog(res.Reason)
				return res
			}

			if len(filesPath) == 0 {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("Pattern '%s' matched no file", path)
				sendLog(res.Reason)
				return res
			}

			buildNumber, errBN := strconv.Atoi(buildNumberString)
			if errBN != nil {
				res.Status = sdk.StatusFail.String()
				res.Reason = fmt.Sprintf("BuilNumber is not an integer %s", errBN)
				sendLog(res.Reason)
				return res
			}

			for _, filePath := range filesPath {
				filename := filepath.Base(filePath)
				throughTempURL, duration, err := sdk.UploadArtifact(project, pipeline, application, tag.Value, filePath, buildNumber, environment)
				if throughTempURL {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to object store", filename, duration.Seconds()))
				} else {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to CDS API", filename, duration.Seconds()))
				}
				if err != nil {
					res.Status = sdk.StatusFail.String()
					res.Reason = fmt.Sprintf("Error while uploading artifact '%s': %v", filename, err)
					sendLog(res.Reason)
					return res
				}
			}

			return res
		}
	}

	return func(ctx context.Context, a *sdk.Action, buildID int64, params *[]sdk.Parameter, sendLog LoggerFunc) sdk.Result {
		res := sdk.Result{Status: sdk.StatusSuccess.String()}

		path := sdk.ParameterValue(a.Parameters, "path")
		if path == "" {
			path = "."
		}

		tag := sdk.ParameterFind(a.Parameters, "tag")
		if tag == nil {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("tag variable is empty. aborting")
			sendLog(res.Reason)
			return res
		}
		tag.Value = strings.Replace(tag.Value, "/", "-", -1)
		tag.Value = url.QueryEscape(tag.Value)

		// Global all files matching filePath
		filesPath, err := filepath.Glob(path)
		if err != nil {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("cannot perform globbing of pattern '%s': %s", path, err)
			sendLog(res.Reason)
			return res
		}

		if len(filesPath) == 0 {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("Pattern '%s' matched no file", path)
			sendLog(res.Reason)
			return res
		}

		var globalError = &sdk.MultiError{}
		var chanError = make(chan error)
		var wg = new(sync.WaitGroup)
		var wgErrors = new(sync.WaitGroup)

		go func() {
			for err := range chanError {
				sendLog(err.Error())
				globalError.Append(err)
				wgErrors.Done()
			}
		}()

		wg.Add(len(filesPath))
		for _, p := range filesPath {
			filename := filepath.Base(p)
			go func(path string) {
				log.Debug("Uploading %s", path)
				defer wg.Done()
				throughTempURL, duration, err := w.client.QueueArtifactUpload(buildID, tag.Value, path)
				if err != nil {
					chanError <- sdk.WrapError(err, "Error while uploading artifact %s", path)
					wgErrors.Add(1)
					return
				}
				if throughTempURL {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to object store", filename, duration.Seconds()))
				} else {
					sendLog(fmt.Sprintf("File '%s' uploaded in %.2fs to CDS API", filename, duration.Seconds()))
				}
			}(p)
			if len(filesPath) > 1 {
				//Wait 3 second to get the object storage to set up all the things
				time.Sleep(3 * time.Second)
			}
		}
		wg.Wait()
		close(chanError)
		<-chanError
		wgErrors.Wait()

		if !globalError.IsEmpty() {
			res.Status = sdk.StatusFail.String()
			res.Reason = fmt.Sprintf("Error: %v", globalError.Error())
			return res
		}

		return res
	}
}
<|endoftext|>


<commit_before>package planner

import (
	"fmt"

	u "github.com/araddon/gou"
	"github.com/lytics/grid"

	"github.com/araddon/qlbridge/exec"
	"github.com/araddon/qlbridge/plan"
	"github.com/araddon/qlbridge/schema"
)

var (
	// ensure it meets interfaces
	_ exec.Executor = (*ExecutorGrid)(nil)

	_ = u.EMPTY
)

// Build a Sql Job which may be a Grid/Distributed job
func BuildSqlJob(ctx *plan.Context, gs *Server) (*ExecutorGrid, error) {
	sqlPlanner := plan.NewPlanner(ctx)
	baseJob := exec.NewExecutor(ctx, sqlPlanner)

	job := &ExecutorGrid{JobExecutor: baseJob}
	baseJob.Executor = job
	job.Executor = baseJob
	job.GridServer = gs
	job.Ctx = ctx
	//u.Debugf("buildsqljob: %T  %T", job, job.Executor)
	task, err := exec.BuildSqlJobPlanned(job.Planner, job, ctx)
	if err != nil {
		return nil, err
	}
	taskRunner, ok := task.(exec.TaskRunner)
	if !ok {
		return nil, fmt.Errorf("Expected TaskRunner but was %T", task)
	}
	job.RootTask = taskRunner
	return job, err
}

// Sql job that wraps the generic qlbridge job builder
// - contains ref to the shared GridServer which has info to
//   distribute tasks across servers
type ExecutorGrid struct {
	*exec.JobExecutor
	GridServer *Server
}

// Finalize is after the Dag of Relational-algebra tasks have been assembled
//  and just before we run them.
func (m *ExecutorGrid) Finalize(resultWriter exec.Task) error {
	//u.Debugf("planner.Finalize  %#v", m.JobExecutor.RootTask)

	m.JobExecutor.RootTask.Add(resultWriter)
	m.JobExecutor.Setup()
	//u.Debugf("finished finalize")
	return nil
}

// func (m *ExecutorGrid) WalkSource(p *plan.Source) (exec.Task, error) {
// 	u.Debugf("NewSource? %#v", p)
// 	if p.SourceExec {
// 		return m.WalkSourceExec(p)
// 	}
// 	return exec.NewSource(m.Ctx, p)
// }
func (m *ExecutorGrid) WalkSelect(p *plan.Select) (exec.Task, error) {
	if len(p.Stmt.From) > 0 {
		//u.Debugf("ExecutorGrid.WalkSelect ?  %s", p.Stmt.Raw)
	}

	if len(p.Stmt.With) > 0 && p.Stmt.With.Bool("distributed") {
		//u.Warnf("%p has distributed!!!!!: %#v", m, p.Stmt.With)

		// We are going to run tasks remotely, so need a local grid source for them
		//  remoteSink  -> nats ->  localSource
		localTask := exec.NewTaskSequential(m.Ctx)
		taskUint, err := NextId()
		if err != nil {
			u.Errorf("Could not create task id %v", err)
			return nil, err
		}
		flow := NewFlow(taskUint)

		rx, err := grid.NewReceiver(m.GridServer.Grid.Nats(), flow.Name(), 2, 0)
		if err != nil {
			u.Errorf("%v: error: %v", "ourtask", err)
			return nil, err
		}
		natsSource := NewSourceNats(m.Ctx, rx)
		localTask.Add(natsSource)

		// submit task in background node
		go func() {
			// task submission to worker actors
			if err := m.GridServer.SubmitTask(localTask, flow, p); err != nil {
				u.Errorf("Could not run task", err)
			}
			// need to send signal to quit
			ch := natsSource.MessageOut()
			close(ch)
		}()
		return localTask, nil
	}

	return m.JobExecutor.WalkSelect(p)
}
func (m *ExecutorGrid) WalkSelectPartition(p *plan.Select, part *schema.Partition) (exec.Task, error) {

	return m.JobExecutor.WalkSelect(p)
}
<commit_msg>fix executor override error<commit_after>package planner

import (
	"fmt"

	u "github.com/araddon/gou"
	"github.com/lytics/grid"

	"github.com/araddon/qlbridge/exec"
	"github.com/araddon/qlbridge/plan"
	"github.com/araddon/qlbridge/schema"
)

var (
	// ensure it meets interfaces
	_ exec.Executor = (*ExecutorGrid)(nil)

	_ = u.EMPTY
)

// Build a Sql Job which may be a Grid/Distributed job
func BuildSqlJob(ctx *plan.Context, gs *Server) (*ExecutorGrid, error) {
	sqlPlanner := plan.NewPlanner(ctx)
	baseJob := exec.NewExecutor(ctx, sqlPlanner)

	job := &ExecutorGrid{JobExecutor: baseJob}
	baseJob.Executor = job
	job.GridServer = gs
	job.Ctx = ctx
	//u.Debugf("buildsqljob: %T  %T", job, job.Executor)
	//u.Debugf("buildsqljob2: %T  %T", baseJob, baseJob.Executor)
	task, err := exec.BuildSqlJobPlanned(job.Planner, job, ctx)
	if err != nil {
		return nil, err
	}
	taskRunner, ok := task.(exec.TaskRunner)
	if !ok {
		return nil, fmt.Errorf("Expected TaskRunner but was %T", task)
	}
	job.RootTask = taskRunner
	return job, err
}

// Sql job that wraps the generic qlbridge job builder
// - contains ref to the shared GridServer which has info to
//   distribute tasks across servers
type ExecutorGrid struct {
	*exec.JobExecutor
	GridServer *Server
}

// Finalize is after the Dag of Relational-algebra tasks have been assembled
//  and just before we run them.
func (m *ExecutorGrid) Finalize(resultWriter exec.Task) error {
	//u.Debugf("planner.Finalize  %#v", m.JobExecutor.RootTask)

	m.JobExecutor.RootTask.Add(resultWriter)
	m.JobExecutor.Setup()
	//u.Debugf("finished finalize")
	return nil
}

func (m *ExecutorGrid) WalkSource(p *plan.Source) (exec.Task, error) {
	u.Debugf("%p %T  NewSource? ", m, m)
	if p.SourceExec {
		return m.WalkSourceExec(p)
	}
	return exec.NewSource(m.Ctx, p)
}
func (m *ExecutorGrid) WalkSelect(p *plan.Select) (exec.Task, error) {
	if len(p.Stmt.From) > 0 {
		//u.Debugf("ExecutorGrid.WalkSelect ?  %s", p.Stmt.Raw)
	}

	if len(p.Stmt.With) > 0 && p.Stmt.With.Bool("distributed") {
		//u.Warnf("%p has distributed!!!!!: %#v", m, p.Stmt.With)

		// We are going to run tasks remotely, so need a local grid source for them
		//  remoteSink  -> nats ->  localSource
		localTask := exec.NewTaskSequential(m.Ctx)
		taskUint, err := NextId()
		if err != nil {
			u.Errorf("Could not create task id %v", err)
			return nil, err
		}
		flow := NewFlow(taskUint)

		rx, err := grid.NewReceiver(m.GridServer.Grid.Nats(), flow.Name(), 2, 0)
		if err != nil {
			u.Errorf("%v: error: %v", "ourtask", err)
			return nil, err
		}
		natsSource := NewSourceNats(m.Ctx, rx)
		localTask.Add(natsSource)

		// submit task in background node
		go func() {
			// task submission to worker actors
			if err := m.GridServer.SubmitTask(localTask, flow, p); err != nil {
				u.Errorf("Could not run task", err)
			}
			// need to send signal to quit
			ch := natsSource.MessageOut()
			close(ch)
		}()
		return localTask, nil
	}

	return m.JobExecutor.WalkSelect(p)
}
func (m *ExecutorGrid) WalkSelectPartition(p *plan.Select, part *schema.Partition) (exec.Task, error) {

	return m.JobExecutor.WalkSelect(p)
}
<|endoftext|>


<commit_before>// The `s3` plugin for SHIELD is intended to be a back-end storage
// plugin, wrapping Amazon's Simple Storage Service (S3). It should
// be compatible with other services which emulate the S3 API, offering
// similar storage solutions for private cloud offerings (such as OpenStack
// Swift). However, this plugin has only been tested with Amazon S3.
//
// PLUGIN FEATURES
//
// This plugin implements functionality suitable for use with the following
// SHIELD Job components:
//
//  Target: no
//  Store:  yes
//
// PLUGIN CONFIGURATION
//
// The endpoint configuration passed to this plugin is used to determine
// how to connect to S3, and where to place/retrieve the data once connected.
// your endpoint JSON should look something like this:
//
//    {
//        "s3_host":             "s3.amazonaws.com", # default
//        "access_key_id":       "your-access-key-id",
//        "secret_access_key":   "your-secret-access-key",
//        "skip_ssl_validation":  false,
//        "bucket":              "bucket-name",
//        "prefix":              "/path/inside/bucket/to/place/backup/data",
//        "signature_version":   "4",  # should be 2 or 4. Defaults to 4
//        "socks5_proxy":        ""    # optionally defined SOCKS5 proxy to use for the s3 communications
//    }
//
// `prefix` will default to the empty string, and backups will be placed in the
// root of the bucket.
//
// STORE DETAILS
//
// When storing data, this plugin connects to the S3 service, and uploads the data
// into the specified bucket, using a path/filename with the following format:
//
//    <prefix>/<YYYY>/<MM>/<DD>/<HH-mm-SS>-<UUID>
//
// Upon successful storage, the plugin then returns this filename to SHIELD to use
// as the `store_key` when the data needs to be retrieved, or purged.
//
// RETRIEVE DETAILS
//
// When retrieving data, this plugin connects to the S3 service, and retrieves the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// PURGE DETAILS
//
// When purging data, this plugin connects to the S3 service, and deletes the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// DEPENDENCIES
//
// None.
//
package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	minio "github.com/minio/minio-go"
	"github.com/starkandwayne/goutils/ansi"
	"golang.org/x/net/proxy"

	"github.com/starkandwayne/shield/plugin"
)

const (
	DefaultS3Host            = "s3.amazonaws.com"
	DefaultPrefix            = ""
	DefaultSigVersion        = "4"
	DefaultSkipSSLValidation = false
)

func validSigVersion(v string) bool {
	return v == "2" || v == "4"
}

func main() {
	p := S3Plugin{
		Name:    "S3 Backup + Storage Plugin",
		Author:  "Stark & Wayne",
		Version: "0.0.1",
		Features: plugin.PluginFeatures{
			Target: "no",
			Store:  "yes",
		},
	}

	plugin.Run(p)
}

type S3Plugin plugin.PluginInfo

type S3ConnectionInfo struct {
	Host              string
	SkipSSLValidation bool
	AccessKey         string
	SecretKey         string
	Bucket            string
	PathPrefix        string
	SignatureVersion  string
	SOCKS5Proxy       string
}

func (p S3Plugin) Meta() plugin.PluginInfo {
	return plugin.PluginInfo(p)
}

func (p S3Plugin) Validate(endpoint plugin.ShieldEndpoint) error {
	var (
		s    string
		err  error
		fail bool
	)

	s, err = endpoint.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		ansi.Printf("@R{\u2717 s3_host              %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 s3_host}              @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("access_key_id")
	if err != nil {
		ansi.Printf("@R{\u2717 access_key_id        %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 access_key_id}        @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("secret_access_key")
	if err != nil {
		ansi.Printf("@R{\u2717 secret_access_key    %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 secret_access_key}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("bucket")
	if err != nil {
		ansi.Printf("@R{\u2717 bucket               %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 bucket}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		ansi.Printf("@R{\u2717 prefix               %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 prefix}               (none)\n")
	} else {
		ansi.Printf("@G{\u2713 prefix}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("signature_version", DefaultSigVersion)
	if err != nil {
		ansi.Printf("@R{\u2717 signature_version    %s}\n", err)
		fail = true
	} else if !validSigVersion(s) {
		ansi.Printf("@R{\u2717 signature_version    Unexpected signature version '%s' found (expecting '2' or '4')}\n", s)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 signature_version}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("socks5_proxy", "")
	if err != nil {
		ansi.Printf("@R{\u2717 socks5_proxy         %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 socks5_proxy}         (no proxy will be used)\n")
	} else {
		ansi.Printf("@G{\u2713 socks5_proxy}         @C{%s}\n", s)
	}

	tf, err := endpoint.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		ansi.Printf("@R{\u2717 skip_ssl_validation  %s}\n", err)
		fail = true
	} else if tf {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{yes}, SSL will @Y{NOT} be validated\n")
	} else {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{no}, SSL @Y{WILL} be validated\n")
	}

	if fail {
		return fmt.Errorf("s3: invalid configuration")
	}
	return nil
}

func (p S3Plugin) Backup(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Restore(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Store(endpoint plugin.ShieldEndpoint) (string, error) {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return "", err
	}
	client, err := s3.Connect()
	if err != nil {
		return "", err
	}

	path := s3.genBackupPath()
	plugin.DEBUG("Storing data in %s", path)

	// FIXME: should we do something with the size of the write performed?
	_, err = client.PutObject(s3.Bucket, path, os.Stdin, "application/x-gzip")
	if err != nil {
		return "", err
	}

	return path, nil
}

func (p S3Plugin) Retrieve(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	reader, err := client.GetObject(s3.Bucket, file)
	if err != nil {
		return err
	}
	if _, err = io.Copy(os.Stdout, reader); err != nil {
		return err
	}

	err = reader.Close()
	if err != nil {
		return err
	}

	return nil
}

func (p S3Plugin) Purge(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	err = client.RemoveObject(s3.Bucket, file)
	if err != nil {
		return err
	}

	return nil
}

func getS3ConnInfo(e plugin.ShieldEndpoint) (S3ConnectionInfo, error) {
	host, err := e.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	insecure_ssl, err := e.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	key, err := e.StringValue("access_key_id")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	secret, err := e.StringValue("secret_access_key")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	bucket, err := e.StringValue("bucket")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	prefix, err := e.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	sigVer, err := e.StringValueDefault("signature_version", DefaultSigVersion)
	if !validSigVersion(sigVer) {
		return S3ConnectionInfo{}, fmt.Errorf("Invalid `signature_version` specified (`%s`). Expected `2` or `4`", sigVer)
	}

	proxy, err := e.StringValueDefault("socks5_proxy", "")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	return S3ConnectionInfo{
		Host:              host,
		SkipSSLValidation: insecure_ssl,
		AccessKey:         key,
		SecretKey:         secret,
		Bucket:            bucket,
		PathPrefix:        prefix,
		SignatureVersion:  sigVer,
		SOCKS5Proxy:       proxy,
	}, nil
}

func (s3 S3ConnectionInfo) genBackupPath() string {
	t := time.Now()
	year, mon, day := t.Date()
	hour, min, sec := t.Clock()
	uuid := plugin.GenUUID()
	path := fmt.Sprintf("%s/%04d/%02d/%02d/%04d-%02d-%02d-%02d%02d%02d-%s", s3.PathPrefix, year, mon, day, year, mon, day, hour, min, sec, uuid)
	// Remove double slashes
	path = strings.Replace(path, "//", "/", -1)
	// Remove a leading slash
	if strings.HasPrefix(path, "/") {
		strings.Replace(path, "/", "", 1)
	}
	return path
}

func (s3 S3ConnectionInfo) Connect() (*minio.Client, error) {
	var s3Client *minio.Client
	var err error
	if s3.SignatureVersion == "2" {
		s3Client, err = minio.NewV2(s3.Host, s3.AccessKey, s3.SecretKey, false)
	} else {
		s3Client, err = minio.NewV4(s3.Host, s3.AccessKey, s3.SecretKey, false)
	}
	if err != nil {
		return nil, err
	}

	transport := http.DefaultTransport
	transport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: s3.SkipSSLValidation}
	if s3.SOCKS5Proxy != "" {
		dialer, err := proxy.SOCKS5("tcp", s3.SOCKS5Proxy, nil, proxy.Direct)
		if err != nil {
			fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
			os.Exit(1)
		}
		transport.(*http.Transport).Dial = dialer.Dial
	}

	s3Client.SetCustomTransport(transport)

	return s3Client, nil
}
<commit_msg>allow ssl to be configurable (#227)<commit_after>// The `s3` plugin for SHIELD is intended to be a back-end storage
// plugin, wrapping Amazon's Simple Storage Service (S3). It should
// be compatible with other services which emulate the S3 API, offering
// similar storage solutions for private cloud offerings (such as OpenStack
// Swift). However, this plugin has only been tested with Amazon S3.
//
// PLUGIN FEATURES
//
// This plugin implements functionality suitable for use with the following
// SHIELD Job components:
//
//  Target: no
//  Store:  yes
//
// PLUGIN CONFIGURATION
//
// The endpoint configuration passed to this plugin is used to determine
// how to connect to S3, and where to place/retrieve the data once connected.
// your endpoint JSON should look something like this:
//
//    {
//        "s3_host":             "s3.amazonaws.com", # default
//        "access_key_id":       "your-access-key-id",
//        "secret_access_key":   "your-secret-access-key",
//        "skip_ssl_validation":  false,
//        "bucket":              "bucket-name",
//        "prefix":              "/path/inside/bucket/to/place/backup/data",
//        "signature_version":   "4",  # should be 2 or 4. Defaults to 4
//        "socks5_proxy":        ""    # optionally defined SOCKS5 proxy to use for the s3 communications
//    }
//
// `prefix` will default to the empty string, and backups will be placed in the
// root of the bucket.
//
// STORE DETAILS
//
// When storing data, this plugin connects to the S3 service, and uploads the data
// into the specified bucket, using a path/filename with the following format:
//
//    <prefix>/<YYYY>/<MM>/<DD>/<HH-mm-SS>-<UUID>
//
// Upon successful storage, the plugin then returns this filename to SHIELD to use
// as the `store_key` when the data needs to be retrieved, or purged.
//
// RETRIEVE DETAILS
//
// When retrieving data, this plugin connects to the S3 service, and retrieves the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// PURGE DETAILS
//
// When purging data, this plugin connects to the S3 service, and deletes the data
// located in the specified bucket, identified by the `store_key` provided by SHIELD.
//
// DEPENDENCIES
//
// None.
//
package main

import (
	"crypto/tls"
	"fmt"
	"io"
	"net/http"
	"os"
	"strings"
	"time"

	minio "github.com/minio/minio-go"
	"github.com/starkandwayne/goutils/ansi"
	"golang.org/x/net/proxy"

	"github.com/starkandwayne/shield/plugin"
)

const (
	DefaultS3Host            = "s3.amazonaws.com"
	DefaultPrefix            = ""
	DefaultSigVersion        = "4"
	DefaultSkipSSLValidation = false
	DefaultSSL               = true
)

func validSigVersion(v string) bool {
	return v == "2" || v == "4"
}

func main() {
	p := S3Plugin{
		Name:    "S3 Backup + Storage Plugin",
		Author:  "Stark & Wayne",
		Version: "0.0.1",
		Features: plugin.PluginFeatures{
			Target: "no",
			Store:  "yes",
		},
	}

	plugin.Run(p)
}

type S3Plugin plugin.PluginInfo

type S3ConnectionInfo struct {
	Host              string
	UseSSL            bool
	SkipSSLValidation bool
	AccessKey         string
	SecretKey         string
	Bucket            string
	PathPrefix        string
	SignatureVersion  string
	SOCKS5Proxy       string
}

func (p S3Plugin) Meta() plugin.PluginInfo {
	return plugin.PluginInfo(p)
}

func (p S3Plugin) Validate(endpoint plugin.ShieldEndpoint) error {
	var (
		s    string
		err  error
		fail bool
	)

	s, err = endpoint.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		ansi.Printf("@R{\u2717 s3_host              %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 s3_host}              @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("access_key_id")
	if err != nil {
		ansi.Printf("@R{\u2717 access_key_id        %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 access_key_id}        @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("secret_access_key")
	if err != nil {
		ansi.Printf("@R{\u2717 secret_access_key    %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 secret_access_key}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValue("bucket")
	if err != nil {
		ansi.Printf("@R{\u2717 bucket               %s}\n", err)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 bucket}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		ansi.Printf("@R{\u2717 prefix               %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 prefix}               (none)\n")
	} else {
		ansi.Printf("@G{\u2713 prefix}               @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("signature_version", DefaultSigVersion)
	if err != nil {
		ansi.Printf("@R{\u2717 signature_version    %s}\n", err)
		fail = true
	} else if !validSigVersion(s) {
		ansi.Printf("@R{\u2717 signature_version    Unexpected signature version '%s' found (expecting '2' or '4')}\n", s)
		fail = true
	} else {
		ansi.Printf("@G{\u2713 signature_version}    @C{%s}\n", s)
	}

	s, err = endpoint.StringValueDefault("socks5_proxy", "")
	if err != nil {
		ansi.Printf("@R{\u2717 socks5_proxy         %s}\n", err)
		fail = true
	} else if s == "" {
		ansi.Printf("@G{\u2713 socks5_proxy}         (no proxy will be used)\n")
	} else {
		ansi.Printf("@G{\u2713 socks5_proxy}         @C{%s}\n", s)
	}

	tf, err := endpoint.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		ansi.Printf("@R{\u2717 skip_ssl_validation  %s}\n", err)
		fail = true
	} else if tf {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{yes}, SSL will @Y{NOT} be validated\n")
	} else {
		ansi.Printf("@G{\u2713 skip_ssl_validation}  @C{no}, SSL @Y{WILL} be validated\n")
	}

	tf, err = endpoint.BooleanValueDefault("use_ssl", DefaultSSL)
	if err != nil {
		ansi.Printf("@R{\u2717 use_ssl  %s}\n", err)
		fail = true
	} else if tf {
		ansi.Printf("@G{\u2713 use_ssl}  @C{yes}, SSL will @Y{NOT} be used\n")
	} else {
		ansi.Printf("@G{\u2713 use_ssl}  @C{no}, SSL @Y{WILL} be used\n")
	}

	if fail {
		return fmt.Errorf("s3: invalid configuration")
	}
	return nil
}

func (p S3Plugin) Backup(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Restore(endpoint plugin.ShieldEndpoint) error {
	return plugin.UNIMPLEMENTED
}

func (p S3Plugin) Store(endpoint plugin.ShieldEndpoint) (string, error) {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return "", err
	}
	client, err := s3.Connect()
	if err != nil {
		return "", err
	}

	path := s3.genBackupPath()
	plugin.DEBUG("Storing data in %s", path)

	// FIXME: should we do something with the size of the write performed?
	_, err = client.PutObject(s3.Bucket, path, os.Stdin, "application/x-gzip")
	if err != nil {
		return "", err
	}

	return path, nil
}

func (p S3Plugin) Retrieve(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	reader, err := client.GetObject(s3.Bucket, file)
	if err != nil {
		return err
	}
	if _, err = io.Copy(os.Stdout, reader); err != nil {
		return err
	}

	err = reader.Close()
	if err != nil {
		return err
	}

	return nil
}

func (p S3Plugin) Purge(endpoint plugin.ShieldEndpoint, file string) error {
	s3, err := getS3ConnInfo(endpoint)
	if err != nil {
		return err
	}
	client, err := s3.Connect()
	if err != nil {
		return err
	}

	err = client.RemoveObject(s3.Bucket, file)
	if err != nil {
		return err
	}

	return nil
}

func getS3ConnInfo(e plugin.ShieldEndpoint) (S3ConnectionInfo, error) {
	host, err := e.StringValueDefault("s3_host", DefaultS3Host)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	insecure_ssl, err := e.BooleanValueDefault("skip_ssl_validation", DefaultSkipSSLValidation)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	use_ssl, err := e.BooleanValueDefault("use_ssl", DefaultSSL)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	key, err := e.StringValue("access_key_id")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	secret, err := e.StringValue("secret_access_key")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	bucket, err := e.StringValue("bucket")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	prefix, err := e.StringValueDefault("prefix", DefaultPrefix)
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	sigVer, err := e.StringValueDefault("signature_version", DefaultSigVersion)
	if !validSigVersion(sigVer) {
		return S3ConnectionInfo{}, fmt.Errorf("Invalid `signature_version` specified (`%s`). Expected `2` or `4`", sigVer)
	}

	proxy, err := e.StringValueDefault("socks5_proxy", "")
	if err != nil {
		return S3ConnectionInfo{}, err
	}

	return S3ConnectionInfo{
		Host:              host,
		SkipSSLValidation: insecure_ssl,
		UseSSL:            use_ssl,
		AccessKey:         key,
		SecretKey:         secret,
		Bucket:            bucket,
		PathPrefix:        prefix,
		SignatureVersion:  sigVer,
		SOCKS5Proxy:       proxy,
	}, nil
}

func (s3 S3ConnectionInfo) genBackupPath() string {
	t := time.Now()
	year, mon, day := t.Date()
	hour, min, sec := t.Clock()
	uuid := plugin.GenUUID()
	path := fmt.Sprintf("%s/%04d/%02d/%02d/%04d-%02d-%02d-%02d%02d%02d-%s", s3.PathPrefix, year, mon, day, year, mon, day, hour, min, sec, uuid)
	// Remove double slashes
	path = strings.Replace(path, "//", "/", -1)
	// Remove a leading slash
	if strings.HasPrefix(path, "/") {
		strings.Replace(path, "/", "", 1)
	}
	return path
}

func (s3 S3ConnectionInfo) Connect() (*minio.Client, error) {
	var s3Client *minio.Client
	var err error
	if s3.SignatureVersion == "2" {
		s3Client, err = minio.NewV2(s3.Host, s3.AccessKey, s3.SecretKey, s3.UseSSL)
	} else {
		s3Client, err = minio.NewV4(s3.Host, s3.AccessKey, s3.SecretKey, s3.UseSSL)
	}
	if err != nil {
		return nil, err
	}

	transport := http.DefaultTransport
	transport.(*http.Transport).TLSClientConfig = &tls.Config{InsecureSkipVerify: s3.SkipSSLValidation}
	if s3.SOCKS5Proxy != "" {
		dialer, err := proxy.SOCKS5("tcp", s3.SOCKS5Proxy, nil, proxy.Direct)
		if err != nil {
			fmt.Fprintln(os.Stderr, "can't connect to the proxy:", err)
			os.Exit(1)
		}
		transport.(*http.Transport).Dial = dialer.Dial
	}

	s3Client.SetCustomTransport(transport)

	return s3Client, nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"

	_ "github.com/lib/pq"
)

// This is a type to hold our word definitions in
type item struct {
	ID         string `json:"id"`
	Word       string `json:"word"`
	Definition string `json:"definition"`
}

var db *sql.DB

func wordHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		rows, err := db.Query("SELECT id,word,definition FROM words")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer rows.Close()

		// Create an array of

		var items []*item
		items = make([]*item, 0)
		for rows.Next() {
			myitem := new(item)
			err = rows.Scan(&myitem.ID, &myitem.Word, &myitem.Definition)
			if err != nil {
				log.Fatal(err)
			}
			items = append(items, myitem)
		}

		jsonstr, err := json.Marshal(items)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(jsonstr)
		return
	case "PUT":
		r.ParseForm()
		_, err := db.Exec("INSERT INTO words (word,definition) VALUES ($1, $2)", r.Form.Get("word"), r.Form.Get("definition"))

		if err != nil {
			log.Println(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusAccepted)
		return
	}

	return
}

func main() {
	// Connect to database:
	// Connection string in $COMPOSE_POSTGRESQL_URL
	// Compose database certificate in $PATH_TO_POSTGRESQL_CERT

	myurl := os.Getenv("COMPOSE_POSTGRESQL_URL") + ("?sslmode=require&sslrootcert=" + os.Getenv("PATH_TO_POSTGRESQL_CERT"))
	fmt.Println(myurl)
	mydb, err := sql.Open("postgres", myurl)
	if err != nil {
		log.Fatal(err)
	}

	db = mydb // Copy to global
	defer db.Close()

	fmt.Println("Connected to PostgreSQL")

	_, err = db.Query(`CREATE TABLE IF NOT EXISTS words (
		id serial primary key,
		word varchar(256) NOT NULL,
		definition varchar(256) NOT NULL)`)

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

	fs := http.FileServer(http.Dir("public"))
	http.Handle("/", fs)
	http.HandleFunc("/words", wordHandler)
	fmt.Println("Listening on 8080")
	http.ListenAndServe(":8080", nil)
}
<commit_msg>Tidied startup code<commit_after>package main

import (
	"database/sql"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"os"

	_ "github.com/lib/pq"
)

// This is a type to hold our word definitions in
type item struct {
	ID         string `json:"id"`
	Word       string `json:"word"`
	Definition string `json:"definition"`
}

var db *sql.DB

func wordHandler(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		rows, err := db.Query("SELECT id,word,definition FROM words")
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		defer rows.Close()

		// Create an array of

		var items []*item
		items = make([]*item, 0)
		for rows.Next() {
			myitem := new(item)
			err = rows.Scan(&myitem.ID, &myitem.Word, &myitem.Definition)
			if err != nil {
				log.Fatal(err)
			}
			items = append(items, myitem)
		}

		jsonstr, err := json.Marshal(items)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}

		w.Header().Set("Content-Type", "application/json")
		w.Write(jsonstr)
		return
	case "PUT":
		r.ParseForm()
		_, err := db.Exec("INSERT INTO words (word,definition) VALUES ($1, $2)", r.Form.Get("word"), r.Form.Get("definition"))

		if err != nil {
			log.Println(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		w.WriteHeader(http.StatusAccepted)
		return
	}

	return
}

func main() {
	// Connect to database:
	// Connection string in $COMPOSE_POSTGRESQL_URL
	// Compose database certificate in $PATH_TO_POSTGRESQL_CERT

	myurl := os.Getenv("COMPOSE_POSTGRESQL_URL") +
		("?sslmode=require&sslrootcert=" + os.Getenv("PATH_TO_POSTGRESQL_CERT"))
	var err error
	db, err = sql.Open("postgres", myurl)
	if err != nil {
		log.Fatal(err)
	}

	defer db.Close()

	_, err = db.Query(`CREATE TABLE IF NOT EXISTS words (
		id serial primary key,
		word varchar(256) NOT NULL,
		definition varchar(256) NOT NULL)`)

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

	fs := http.FileServer(http.Dir("public"))
	http.Handle("/", fs)
	http.HandleFunc("/words", wordHandler)
	fmt.Println("Listening on 8080")
	http.ListenAndServe(":8080", nil)
}
<|endoftext|>


<commit_before>package bitfield

import "encoding/hex"

type BitField struct {
	b      []byte
	length uint32
}

// New creates a new BitField value of length bits.
func New(length uint32) BitField {
	return BitField{make([]byte, (length+7)/8), length}
}

// NewBytes returns a new BitField value from b.
// Bytes in b are not copied. Unused bits in last byte are cleared.
// Panics if b is not big enough to hold "length" bits.
func NewBytes(b []byte, length uint32) BitField {
	div, mod := divMod32(length, 8)
	lastByteIncomplete := mod != 0
	requiredBytes := div
	if lastByteIncomplete {
		requiredBytes++
	}
	if uint32(len(b)) < requiredBytes {
		panic("not enough bytes in slice for specified length")
	}
	if lastByteIncomplete {
		b[len(b)-1] &= ^(0xff >> mod)
	}
	return BitField{b[:requiredBytes], length}
}

// Bytes returns bytes in b. If you modify the returned slice the bits in b are modified too.
func (b BitField) Bytes() []byte { return b.b }

// Len returns the number of bits as given to New.
func (b BitField) Len() uint32 { return b.length }

// Hex returns bytes as string. If not all the bits in last byte are used, they encode as not set.
func (b BitField) Hex() string { return hex.EncodeToString(b.b) }

// Set bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b BitField) Set(i uint32) {
	b.checkIndex(i)
	div, mod := divMod32(i, 8)
	b.b[div] |= 1 << (7 - mod)
}

// SetTo sets bit i to value. Panics if i >= b.Len().
func (b BitField) SetTo(i uint32, value bool) {
	b.checkIndex(i)
	if value {
		b.Set(i)
	} else {
		b.Clear(i)
	}
}

// Clear bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b BitField) Clear(i uint32) {
	b.checkIndex(i)
	div, mod := divMod32(i, 8)
	b.b[div] &= ^(1 << (7 - mod))
}

// ClearAll clears all bits.
func (b BitField) ClearAll() {
	for i := range b.b {
		b.b[i] = 0
	}
}

// Test bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b BitField) Test(i uint32) bool {
	b.checkIndex(i)
	div, mod := divMod32(i, 8)
	return (b.b[div] & (1 << (7 - mod))) > 0
}

var countCache = [256]byte{
	0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8,
}

// Count returns the count of set bits.
func (b BitField) Count() uint32 {
	var total uint32
	for _, v := range b.b {
		total += uint32(countCache[v])
	}
	return total
}

// All returns true if all bits are set, false otherwise.
func (b BitField) All() bool {
	return b.Count() == b.length
}

func (b BitField) checkIndex(i uint32) {
	if i >= b.Len() {
		panic("index out of bound")
	}
}

func divMod32(a, b uint32) (uint32, uint32) { return a / b, a % b }
<commit_msg>pointer<commit_after>package bitfield

import "encoding/hex"

type BitField struct {
	b      []byte
	length uint32
}

// New creates a new BitField value of length bits.
func New(length uint32) BitField {
	return BitField{make([]byte, (length+7)/8), length}
}

// NewBytes returns a new BitField value from b.
// Bytes in b are not copied. Unused bits in last byte are cleared.
// Panics if b is not big enough to hold "length" bits.
func NewBytes(b []byte, length uint32) BitField {
	div, mod := divMod32(length, 8)
	lastByteIncomplete := mod != 0
	requiredBytes := div
	if lastByteIncomplete {
		requiredBytes++
	}
	if uint32(len(b)) < requiredBytes {
		panic("not enough bytes in slice for specified length")
	}
	if lastByteIncomplete {
		b[len(b)-1] &= ^(0xff >> mod)
	}
	return BitField{b[:requiredBytes], length}
}

// Bytes returns bytes in b. If you modify the returned slice the bits in b are modified too.
func (b *BitField) Bytes() []byte { return b.b }

// Len returns the number of bits as given to New.
func (b *BitField) Len() uint32 { return b.length }

// Hex returns bytes as string. If not all the bits in last byte are used, they encode as not set.
func (b *BitField) Hex() string { return hex.EncodeToString(b.b) }

// Set bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b *BitField) Set(i uint32) {
	b.checkIndex(i)
	div, mod := divMod32(i, 8)
	b.b[div] |= 1 << (7 - mod)
}

// SetTo sets bit i to value. Panics if i >= b.Len().
func (b *BitField) SetTo(i uint32, value bool) {
	b.checkIndex(i)
	if value {
		b.Set(i)
	} else {
		b.Clear(i)
	}
}

// Clear bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b *BitField) Clear(i uint32) {
	b.checkIndex(i)
	div, mod := divMod32(i, 8)
	b.b[div] &= ^(1 << (7 - mod))
}

// ClearAll clears all bits.
func (b *BitField) ClearAll() {
	for i := range b.b {
		b.b[i] = 0
	}
}

// Test bit i. 0 is the most significant bit. Panics if i >= b.Len().
func (b *BitField) Test(i uint32) bool {
	b.checkIndex(i)
	div, mod := divMod32(i, 8)
	return (b.b[div] & (1 << (7 - mod))) > 0
}

var countCache = [256]byte{
	0, 1, 1, 2, 1, 2, 2, 3, 1, 2, 2, 3, 2, 3, 3, 4,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	1, 2, 2, 3, 2, 3, 3, 4, 2, 3, 3, 4, 3, 4, 4, 5,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	2, 3, 3, 4, 3, 4, 4, 5, 3, 4, 4, 5, 4, 5, 5, 6,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	3, 4, 4, 5, 4, 5, 5, 6, 4, 5, 5, 6, 5, 6, 6, 7,
	4, 5, 5, 6, 5, 6, 6, 7, 5, 6, 6, 7, 6, 7, 7, 8,
}

// Count returns the count of set bits.
func (b *BitField) Count() uint32 {
	var total uint32
	for _, v := range b.b {
		total += uint32(countCache[v])
	}
	return total
}

// All returns true if all bits are set, false otherwise.
func (b *BitField) All() bool {
	return b.Count() == b.length
}

func (b *BitField) checkIndex(i uint32) {
	if i >= b.Len() {
		panic("index out of bound")
	}
}

func divMod32(a, b uint32) (uint32, uint32) { return a / b, a % b }
<|endoftext|>