<commit_before>package fstab

import (
	"fmt"
	"strconv"
	"strings"
)

// Mount represetnts the filesystem info
type Mount struct {
	// The block special device or remote filesystem to be mounted
	Spec string

	// The mount point for the filesystem
	File string

	// The type of the filesystem
	VfsType string

	// Mount options associated with the filesystem
	MntOps []string

	// Used by dump to determine which filesystems need to be dumped.
	Freq int

	// Used by the fsck(8) program to determine the order in which filesystem checks are done at reboot time
	PassNo int
}

const (
	Path = 0
	Label
	UUID
	PartUUID
	PartLabel
)

type DeviceIdentifierType int

// parseOptions parses the options field into an array of strings
func parseOptions(options string) []string {
	return strings.Split(options, ",")
}

// String serializes the object into fstab format
func (mount *Mount) String() string {
	return fmt.Sprintf("%-21s %-21s %-21s %-21s %2d %2d", mount.Spec, mount.File, mount.VfsType, strings.Join(mount.MntOps, ","), mount.Freq, mount.PassNo)
}

func (mount *Mount) IsSwap() bool {
	return "swap" == mount.VfsType
}

func (mount *Mount) IsNFS() bool {
	return "nfs" == mount.VfsType
}

// SpecType returns the device identifier type
func (mount *Mount) SpecType() (spectype DeviceIdentifierType) {
	bits := strings.Split(mount.Spec, "=")
	switch strings.ToUpper(bits[0]) {
	case "UUID":
		spectype = UUID

	case "LABEL":
		spectype = Label

	case "PARTUUID":
		spectype = PartUUID

	case "PARTLABEL":
		spectype = PartLabel

	default:
		spectype = Path
	}
	return
}

// SpecType returns the device identifier value; that is if Spec is
// "UUID=vogons-ate-my-sandwich", it will return "vogons-ate-my-sandwich"
func (mount *Mount) SpecValue() string {
	bits := strings.Split(mount.Spec, "=")
	if 1 == len(bits) {
		return mount.Spec
	} else {
		return bits[1]
	}
}

// ParseLine parses a single line (of an fstab).
// It will most frequently return a Mount; however,
// If a parsing error occurs, `err` will be non-nil and provide an error message.
// If the line is either empy or a comment line, `mount` will also be nil.
func ParseLine(line string) (mount *Mount, err error) {
	line = strings.TrimSpace(line)

	// Lines starting with a pound sign (#) are comments, and are ignored. So are empty lines.
	if ("" == line) || (line[0] == '#') {
		return nil, nil
	}

	fields := strings.Fields(line)
	if len(fields) < 4 {
		return nil, fmt.Errorf("too few fields (%i), at least 4 are expected", len(fields))
	} else {
		mount = new(Mount)
		mount.Spec = fields[0]
		mount.File = fields[1]
		mount.VfsType = fields[2]
		mount.MntOps = parseOptions(fields[3])

		var convErr error

		if len(fields) > 4 {
			mount.Freq, convErr = strconv.Atoi(fields[4])
			if nil != convErr {
				return nil, fmt.Errorf("%s is not a number", fields[4])
			}
		}

		if len(fields) > 5 {
			mount.PassNo, convErr = strconv.Atoi(fields[5])
			if nil != convErr {
				return nil, fmt.Errorf("%s it not a number", fields[5])
			}
		}
	}

	return
}
<commit_msg>Added key-value parsing for options<commit_after>package fstab

import (
	"fmt"
	"strconv"
	"strings"
)

// Mount represetnts the filesystem info
type Mount struct {
	// The block special device or remote filesystem to be mounted
	Spec string

	// The mount point for the filesystem
	File string

	// The type of the filesystem
	VfsType string

	// Mount options associated with the filesystem
	MntOps map[string]string

	// Used by dump to determine which filesystems need to be dumped.
	Freq int

	// Used by the fsck(8) program to determine the order in which filesystem checks are done at reboot time
	PassNo int
}

const (
	Path = 0
	Label
	UUID
	PartUUID
	PartLabel
)

type DeviceIdentifierType int

// parseOptions parses the options field into an array of strings
func parseOptions(optionsString string) (options map[string]string) {
	options = make(map[string]string)
	for _, option := range strings.Split(optionsString, ",") {
		bits := strings.Split(strings.TrimSpace(option), "=")
		if len(bits) > 1 {
			options[bits[0]] = bits[1]
		} else {
			options[bits[0]] = ""
		}
	}
	return
}

func (mount Mount) MntOpsString() (opsstring string) {
	first := true
	for key, value := range mount.MntOps {
		if first {
			first = false
		} else {
			opsstring += ","
		}

		opsstring += key

		if "" != value {
			opsstring += "=" + value
		}
	}
	return
}

// String serializes the object into fstab format
func (mount Mount) String() string {
	return fmt.Sprintf("%-21s %-21s %-21s %-21s %2d %2d", mount.Spec, mount.File, mount.VfsType, mount.MntOpsString(), mount.Freq, mount.PassNo)
}

func (mount *Mount) IsSwap() bool {
	return "swap" == mount.VfsType
}

func (mount *Mount) IsNFS() bool {
	return "nfs" == mount.VfsType
}

// SpecType returns the device identifier type
func (mount *Mount) SpecType() (spectype DeviceIdentifierType) {
	bits := strings.Split(mount.Spec, "=")
	switch strings.ToUpper(bits[0]) {
	case "UUID":
		spectype = UUID

	case "LABEL":
		spectype = Label

	case "PARTUUID":
		spectype = PartUUID

	case "PARTLABEL":
		spectype = PartLabel

	default:
		spectype = Path
	}
	return
}

// SpecType returns the device identifier value; that is if Spec is
// "UUID=vogons-ate-my-sandwich", it will return "vogons-ate-my-sandwich"
func (mount *Mount) SpecValue() string {
	bits := strings.Split(mount.Spec, "=")
	if 1 == len(bits) {
		return mount.Spec
	} else {
		return bits[1]
	}
}

// ParseLine parses a single line (of an fstab).
// It will most frequently return a Mount; however,
// If a parsing error occurs, `err` will be non-nil and provide an error message.
// If the line is either empy or a comment line, `mount` will also be nil.
func ParseLine(line string) (mount *Mount, err error) {
	line = strings.TrimSpace(line)

	// Lines starting with a pound sign (#) are comments, and are ignored. So are empty lines.
	if ("" == line) || (line[0] == '#') {
		return nil, nil
	}

	fields := strings.Fields(line)
	if len(fields) < 4 {
		return nil, fmt.Errorf("too few fields (%i), at least 4 are expected", len(fields))
	} else {
		mount = new(Mount)
		mount.Spec = fields[0]
		mount.File = fields[1]
		mount.VfsType = fields[2]
		mount.MntOps = parseOptions(fields[3])

		var convErr error

		if len(fields) > 4 {
			mount.Freq, convErr = strconv.Atoi(fields[4])
			if nil != convErr {
				return nil, fmt.Errorf("%s is not a number", fields[4])
			}
		}

		if len(fields) > 5 {
			mount.PassNo, convErr = strconv.Atoi(fields[5])
			if nil != convErr {
				return nil, fmt.Errorf("%s it not a number", fields[5])
			}
		}
	}

	return
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	"time"
)

//DefaultConnMaxLifetimeInSecond default conn max lifetime
const DefaultConnMaxLifetimeInSecond = int64(30)

//PlainDBOption plain database init option interface.
type PlainDBOption interface {
	//Apply init plain database.
	Apply(*PlainDB) error
}

//Config database config
type Config struct {
	//Driver sql driver.
	Driver string
	//Type sql database type.
	Type string
	//Conn sql conn string.
	DataSource string
	//Prefix sql table prefix.
	Prefix string
	//MaxIdleConns max idle conns.
	MaxIdleConns int
	//ConnMaxLifetimeInSecond conn max Lifetime in second.
	ConnMaxLifetimeInSecond int64
	//MaxOpenConns max open conns.
	MaxOpenConns int
}

//Apply init plain database with config
func (c *Config) Apply(d *PlainDB) error {
	db, err := sql.Open(c.Driver, c.DataSource)
	if err != nil {
		return err
	}
	if c.MaxIdleConns > 0 {
		db.SetMaxIdleConns(c.MaxIdleConns)
	}
	if c.ConnMaxLifetimeInSecond > 0 {
		db.SetConnMaxLifetime(time.Duration(c.ConnMaxLifetimeInSecond) * time.Second)
	} else if c.ConnMaxLifetimeInSecond == 0 {
		db.SetConnMaxLifetime(time.Duration(DefaultConnMaxLifetimeInSecond) * time.Second)
	}
	if c.MaxOpenConns > 0 {
		db.SetMaxOpenConns(c.MaxOpenConns)
	}
	d.SetDB(db)
	d.SetPrefix(c.Prefix)
	return nil
}
<commit_msg>update<commit_after>package db

import (
	"database/sql"
	"time"
)

//DefaultConnMaxLifetimeInSecond default conn max lifetime
const DefaultConnMaxLifetimeInSecond = int64(30)

//PlainDBOption plain database init option interface.
type PlainDBOption interface {
	//Apply init plain database.
	Apply(*PlainDB) error
}

//Config database config
type Config struct {
	//Driver sql driver.
	Driver string
	//Type sql database type.
	Type string
	//Conn sql conn string.
	DataSource string
	//Prefix sql table prefix.
	Prefix string
	//MaxIdleConns max idle conns.
	MaxIdleConns int
	//ConnMaxLifetimeInSecond conn max Lifetime in second.
	ConnMaxLifetimeInSecond int64
	//MaxOpenConns max open conns.
	MaxOpenConns int
}

//Apply init plain database with config
func (c *Config) Apply(d *PlainDB) error {
	db, err := sql.Open(c.Driver, c.DataSource)
	if err != nil {
		return err
	}
	if c.MaxIdleConns > 0 {
		db.SetMaxIdleConns(c.MaxIdleConns)
	}
	if c.ConnMaxLifetimeInSecond > 0 {
		db.SetConnMaxLifetime(time.Duration(c.ConnMaxLifetimeInSecond) * time.Second)
	} else if c.ConnMaxLifetimeInSecond == 0 {
		db.SetConnMaxLifetime(time.Duration(DefaultConnMaxLifetimeInSecond) * time.Second)
	}
	if c.MaxOpenConns > 0 {
		db.SetMaxOpenConns(c.MaxOpenConns)
	}
	d.SetDB(db)
	d.SetPrefix(c.Prefix)
	if c.Type != "" {
		d.SetDriver(c.Type)
	} else {
		d.SetDriver(c.Driver)
	}
	return nil
}
<|endoftext|>


<commit_before>package models

import (
	"fmt"
	"math"
	"sort"
	"time"
	"unicode/utf8"
)

type StreamStatusGameHistory struct {
	Game  string    `bson:",omitempty"`
	Start time.Time `bson:",omitempty"`
}
type GamesHistory []StreamStatusGameHistory

func (history GamesHistory) Len() int {
	return len(history)
}
func (history GamesHistory) Less(i, j int) bool {
	return history[i].Start.Before(history[j].Start)
}

func (history GamesHistory) Swap(i, j int) {
	history[i], history[j] = history[j], history[i]
}

func (history GamesHistory) ReturnHistory() string {
	sort.Sort(sort.Reverse(history))
	stringHistory := ""
	if len(history) == 0 {
		return ""
	}
	for index := range history {
		newStringHistory := ""

		duration := time.Second * 1
		if index == 0 {
			duration = time.Now().Sub(history[0].Start)
		} else {
			duration = history[index-1].Start.Sub(history[index].Start)
		}
		minutes := float64(int(duration.Minutes() - math.Floor(duration.Minutes()/60)*60))
		hours := float64(int(duration.Hours()))
		stringDuration := fmt.Sprintf("%dh%dm]", int(hours), int(minutes))
		if minutes < 10 {
			stringDuration = fmt.Sprintf("%dh0%dm]", int(hours), int(minutes))
		}
		if stringHistory == "" {
			newStringHistory = history[index].Game + " [!NOW!" + stringDuration
		} else {
			newStringHistory = history[index].Game + " [" + stringDuration + " -> " + stringHistory
		}

		if utf8.RuneCountInString(newStringHistory) > 400 {
			break
		} else {
			stringHistory = newStringHistory
		}
	}
	return stringHistory
}

// StreamStatus Describes info about stream, when started, what game and title is, and if it is online
type StreamStatus struct {
	Online         bool
	Game           string       `bson:",omitempty"`
	Title          string       `bson:",omitempty"`
	Start          time.Time    `bson:",omitempty"`
	LastOnlineTime time.Time    `bson:",omitempty"`
	Viewers        int          `bson:",omitempty"`
	GamesHistory   GamesHistory `bson:",omitempty"`
}
<commit_msg>Fix: Reversed order of game history<commit_after>package models

import (
	"fmt"
	"math"
	"sort"
	"time"
	"unicode/utf8"
)

type StreamStatusGameHistory struct {
	Game  string    `bson:",omitempty"`
	Start time.Time `bson:",omitempty"`
}
type GamesHistory []StreamStatusGameHistory

func (history GamesHistory) Len() int {
	return len(history)
}
func (history GamesHistory) Less(i, j int) bool {
	return history[i].Start.Before(history[j].Start)
}

func (history GamesHistory) Swap(i, j int) {
	history[i], history[j] = history[j], history[i]
}

func (history GamesHistory) ReturnHistory() string {
	sort.Sort(sort.Reverse(history))
	stringHistory := ""
	if len(history) == 0 {
		return ""
	}
	for index := range history {
		newStringHistory := ""

		duration := time.Second * 1
		if index == 0 {
			duration = time.Now().Sub(history[0].Start)
		} else {
			duration = history[index-1].Start.Sub(history[index].Start)
		}
		minutes := float64(int(duration.Minutes() - math.Floor(duration.Minutes()/60)*60))
		hours := float64(int(duration.Hours()))
		stringDuration := fmt.Sprintf("%dh%dm]", int(hours), int(minutes))
		if minutes < 10 {
			stringDuration = fmt.Sprintf("%dh0%dm]", int(hours), int(minutes))
		}
		if stringHistory == "" {
			newStringHistory = "NOW" history[index].Game + " [" + stringDuration
		} else {
			newStringHistory = history[index].Game + " [" + stringDuration + " -> " + stringHistory
		}

		if utf8.RuneCountInString(newStringHistory) > 400 {
			break
		} else {
			stringHistory = newStringHistory
		}
	}
	return stringHistory
}

// StreamStatus Describes info about stream, when started, what game and title is, and if it is online
type StreamStatus struct {
	Online         bool
	Game           string       `bson:",omitempty"`
	Title          string       `bson:",omitempty"`
	Start          time.Time    `bson:",omitempty"`
	LastOnlineTime time.Time    `bson:",omitempty"`
	Viewers        int          `bson:",omitempty"`
	GamesHistory   GamesHistory `bson:",omitempty"`
}
<|endoftext|>


<commit_before>package proxy

import (
	"fmt"
	"io"
	"net"
	"net/http"

	"github.com/yryz/httpproxy/config"

	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
	log "github.com/sirupsen/logrus"
)

type ProxyServer struct {
	ssServer string
	ssCipher *ss.Cipher
}

func NewProxyServer() *http.Server {
	p := &ProxyServer{
		ssServer: config.Conf.SsServer,
	}

	var err error
	p.ssCipher, err = ss.NewCipher(config.Conf.SsCipher, config.Conf.SsPassword)
	if err != nil {
		panic("init cipher error: " + err.Error())
	}

	return &http.Server{
		Addr:    config.Conf.Listen,
		Handler: p,
	}
}

func (p *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer func() {
		if err := recover(); err != nil {
			http.Error(w, fmt.Sprint(err), http.StatusInternalServerError)
			log.Debugf("panic: %v\n", err)
		}
	}()

	if r.Method == "CONNECT" {
		p.HandleConnect(w, r)
	} else {
		p.HandleHttp(w, r)
	}
}

// 处理HTTPS、HTTP2代理请求
func (p *ProxyServer) HandleConnect(w http.ResponseWriter, r *http.Request) {
	log.Infof("%s %s", r.Method, r.Host)

	hj, _ := w.(http.Hijacker)
	conn, _, err := hj.Hijack()
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	ssConn, err := ss.Dial(r.URL.Host, p.ssServer, p.ssCipher.Copy())
	if err != nil {
		log.Error("ss dial: ", err)
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}

	conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))

	go ss.PipeThenClose(conn, ssConn)
	ss.PipeThenClose(ssConn, conn)
}

// 处理HTTP代理请求
func (p *ProxyServer) HandleHttp(w http.ResponseWriter, r *http.Request) {
	log.Infof("%s %s", r.Method, r.URL)

	// ss proxy
	tr := http.Transport{
		Dial: func(network, addr string) (net.Conn, error) {
			log.Infof("dial ss %v/%v", addr, network)
			return ss.Dial(addr, p.ssServer, p.ssCipher.Copy())
		},
	}

	// transport
	resp, err := tr.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		log.Error("request error: ", err)
		return
	}
	defer resp.Body.Close()

	// copy headers
	for k, values := range resp.Header {
		for _, v := range values {
			w.Header().Add(k, v)
		}
	}

	// copy body
	n, err := io.Copy(w, resp.Body)
	if err != nil && err != io.EOF {
		log.Errorf("copy response body error: %v", err)
	}

	log.Infof("copied %v bytes from %v.", n, r.Host)
}
<commit_msg>fix shadowsocks库<commit_after>package proxy

import (
	"fmt"
	"io"
	"net"
	"net/http"

	"github.com/yryz/httpproxy/config"

	ss "github.com/shadowsocks/shadowsocks-go/shadowsocks"
	log "github.com/sirupsen/logrus"
)

type ProxyServer struct {
	ssServer string
	ssCipher *ss.Cipher
}

func NewProxyServer() *http.Server {
	p := &ProxyServer{
		ssServer: config.Conf.SsServer,
	}

	var err error
	p.ssCipher, err = ss.NewCipher(config.Conf.SsCipher, config.Conf.SsPassword)
	if err != nil {
		panic("init cipher error: " + err.Error())
	}

	return &http.Server{
		Addr:    config.Conf.Listen,
		Handler: p,
	}
}

func (p *ProxyServer) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer func() {
		if err := recover(); err != nil {
			http.Error(w, fmt.Sprint(err), http.StatusInternalServerError)
			log.Debugf("panic: %v\n", err)
		}
	}()

	if r.Method == "CONNECT" {
		p.HandleConnect(w, r)
	} else {
		p.HandleHttp(w, r)
	}
}

// 处理HTTPS、HTTP2代理请求
func (p *ProxyServer) HandleConnect(w http.ResponseWriter, r *http.Request) {
	log.Infof("%s %s", r.Method, r.Host)

	hj, _ := w.(http.Hijacker)
	conn, _, err := hj.Hijack()
	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
		return
	}

	ssConn, err := ss.Dial(r.URL.Host, p.ssServer, p.ssCipher.Copy())
	if err != nil {
		log.Error("ss dial: ", err)
		http.Error(w, err.Error(), http.StatusBadGateway)
		return
	}

	conn.Write([]byte("HTTP/1.1 200 Connection Established\r\n\r\n"))

	go ss.PipeThenClose(conn, ssConn, nil)
	ss.PipeThenClose(ssConn, conn, nil)
}

// 处理HTTP代理请求
func (p *ProxyServer) HandleHttp(w http.ResponseWriter, r *http.Request) {
	log.Infof("%s %s", r.Method, r.URL)

	// ss proxy
	tr := http.Transport{
		Dial: func(network, addr string) (net.Conn, error) {
			log.Infof("dial ss %v/%v", addr, network)
			return ss.Dial(addr, p.ssServer, p.ssCipher.Copy())
		},
	}

	// transport
	resp, err := tr.RoundTrip(r)
	if err != nil {
		http.Error(w, err.Error(), http.StatusInternalServerError)
		log.Error("request error: ", err)
		return
	}
	defer resp.Body.Close()

	// copy headers
	for k, values := range resp.Header {
		for _, v := range values {
			w.Header().Add(k, v)
		}
	}

	// copy body
	n, err := io.Copy(w, resp.Body)
	if err != nil && err != io.EOF {
		log.Errorf("copy response body error: %v", err)
	}

	log.Infof("copied %v bytes from %v.", n, r.Host)
}
<|endoftext|>


<commit_before>package proxy

import (
	"encoding/json"
	"errors"
	"net"
	"strconv"
	"sync"
	"unicode/utf8"

	"github.com/cloudfoundry-incubator/diego-ssh/helpers"
	"github.com/cloudfoundry-incubator/runtime-schema/metric"
	"github.com/cloudfoundry/dropsonde/logs"
	"github.com/pivotal-golang/lager"
	"golang.org/x/crypto/ssh"
)

const (
	sshConnections = metric.Metric("ssh-connections")
)

type Waiter interface {
	Wait() error
}

type TargetConfig struct {
	Address         string `json:"address"`
	HostFingerprint string `json:"host_fingerprint"`
	User            string `json:"user,omitempty"`
	Password        string `json:"password,omitempty"`
	PrivateKey      string `json:"private_key,omitempty"`
}

type LogMessage struct {
	Guid    string `json:"guid"`
	Message string `json:"message"`
	Index   int    `json:"index"`
}

type Proxy struct {
	logger       lager.Logger
	serverConfig *ssh.ServerConfig

	connectionLock *sync.Mutex
	connections    int
}

func New(
	logger lager.Logger,
	serverConfig *ssh.ServerConfig,
) *Proxy {
	return &Proxy{
		logger:         logger,
		serverConfig:   serverConfig,
		connectionLock: &sync.Mutex{},
	}
}

func (p *Proxy) HandleConnection(netConn net.Conn) {
	logger := p.logger.Session("handle-connection")
	defer netConn.Close()

	serverConn, serverChannels, serverRequests, err := ssh.NewServerConn(netConn, p.serverConfig)
	if err != nil {
		return
	}
	defer serverConn.Close()

	clientConn, clientChannels, clientRequests, err := NewClientConn(logger, serverConn.Permissions)
	if err != nil {
		return
	}
	defer clientConn.Close()

	emitLogMessage(logger, serverConn.Permissions)

	fromClientLogger := logger.Session("from-client")
	fromDaemonLogger := logger.Session("from-daemon")

	go ProxyGlobalRequests(fromClientLogger, clientConn, serverRequests)
	go ProxyGlobalRequests(fromDaemonLogger, serverConn, clientRequests)

	go ProxyChannels(fromClientLogger, clientConn, serverChannels)
	go ProxyChannels(fromDaemonLogger, serverConn, clientChannels)

	p.connectionLock.Lock()
	p.connections++
	err = sshConnections.Send(p.connections)
	if err != nil {
		logger.Error("failed-to-send-ssh-connections-metric", err)
	}
	p.connectionLock.Unlock()

	defer func() {
		p.emitConnectionClosing(logger)
	}()

	Wait(logger, serverConn, clientConn)
}

func (p *Proxy) emitConnectionClosing(logger lager.Logger) {
	p.connectionLock.Lock()
	p.connections--
	err := sshConnections.Send(p.connections)
	p.connectionLock.Unlock()

	if err != nil {
		logger.Error("failed-to-send-ssh-connections-metric", err)
	}
}

func emitLogMessage(logger lager.Logger, perms *ssh.Permissions) {
	logMessageJson := perms.CriticalOptions["log-message"]
	if logMessageJson == "" {
		return
	}

	logMessage := &LogMessage{}
	err := json.Unmarshal([]byte(logMessageJson), logMessage)
	if err != nil {
		logger.Error("json-unmarshal-failed", err)
		return
	}

	logs.SendAppLog(logMessage.Guid, logMessage.Message, "SSH", strconv.Itoa(logMessage.Index))
}

func ProxyGlobalRequests(logger lager.Logger, conn ssh.Conn, reqs <-chan *ssh.Request) {
	logger = logger.Session("proxy-global-requests")

	logger.Info("started")
	defer logger.Info("completed")

	for req := range reqs {
		logger.Info("request", lager.Data{
			"type":      req.Type,
			"wantReply": req.WantReply,
			"payload":   req.Payload,
		})

		success, reply, err := conn.SendRequest(req.Type, req.WantReply, req.Payload)
		if err != nil {
			logger.Error("send-request-failed", err)
			continue
		}

		if req.WantReply {
			req.Reply(success, reply)
		}
	}
}

func ProxyChannels(logger lager.Logger, conn ssh.Conn, channels <-chan ssh.NewChannel) {
	logger = logger.Session("proxy-channels")

	logger.Info("started")
	defer func() {
		logger.Info("completed")
		conn.Close()
	}()

	for newChannel := range channels {
		handleNewChannel(logger, conn, newChannel)
	}
}

func handleNewChannel(logger lager.Logger, conn ssh.Conn, newChannel ssh.NewChannel) {
	logger.Info("new-channel", lager.Data{
		"channelType": newChannel.ChannelType(),
		"extraData":   newChannel.ExtraData(),
	})

	targetChan, targetReqs, err := conn.OpenChannel(newChannel.ChannelType(), newChannel.ExtraData())
	if err != nil {
		logger.Error("failed-to-open-channel", err)
		if openErr, ok := err.(*ssh.OpenChannelError); ok {
			newChannel.Reject(openErr.Reason, openErr.Message)
		} else {
			newChannel.Reject(ssh.ConnectionFailed, err.Error())
		}
		return
	}

	sourceChan, sourceReqs, err := newChannel.Accept()
	if err != nil {
		targetChan.Close()
		return
	}

	toTargetLogger := logger.Session("to-target")
	toSourceLogger := logger.Session("to-source")

	targetWg := &sync.WaitGroup{}
	sourceWg := &sync.WaitGroup{}

	targetWg.Add(1)
	sourceWg.Add(1)

	go func() {
		helpers.Copy(toTargetLogger, targetWg, targetChan, sourceChan)
		targetChan.CloseWrite()
	}()

	go func() {
		helpers.Copy(toSourceLogger, sourceWg, sourceChan, targetChan)
		sourceChan.CloseWrite()
	}()

	go ProxyRequests(toTargetLogger, newChannel.ChannelType(), sourceReqs, targetChan, targetWg)
	go ProxyRequests(toSourceLogger, newChannel.ChannelType(), targetReqs, sourceChan, sourceWg)
}

func ProxyRequests(logger lager.Logger, channelType string, reqs <-chan *ssh.Request, channel ssh.Channel, wg *sync.WaitGroup) {
	logger = logger.Session("proxy-requests", lager.Data{
		"channel-type": channelType,
	})

	logger.Info("started")
	defer func() {
		logger.Info("completed")
		wg.Wait()
		channel.Close()
	}()

	for req := range reqs {
		logger.Info("request", lager.Data{
			"type":      req.Type,
			"wantReply": req.WantReply,
			"payload":   req.Payload,
		})
		success, err := channel.SendRequest(req.Type, req.WantReply, req.Payload)
		if err != nil {
			logger.Error("send-request-failed", err)
			continue
		}

		if req.WantReply {
			req.Reply(success, nil)
		}

		if req.Type == "exit-status" {
			logger.Info("exiting-dawg")
			return
		}
	}
}

func Wait(logger lager.Logger, waiters ...Waiter) {
	wg := &sync.WaitGroup{}
	for _, waiter := range waiters {
		wg.Add(1)
		go func(waiter Waiter) {
			waiter.Wait()
			wg.Done()
		}(waiter)
	}
	wg.Wait()
}

func NewClientConn(logger lager.Logger, permissions *ssh.Permissions) (ssh.Conn, <-chan ssh.NewChannel, <-chan *ssh.Request, error) {
	if permissions == nil || permissions.CriticalOptions == nil {
		err := errors.New("Invalid permissions from authentication")
		logger.Error("permissions-and-critical-options-required", err)
		return nil, nil, nil, err
	}

	targetConfigJson := permissions.CriticalOptions["proxy-target-config"]
	logger = logger.Session("new-client-conn", lager.Data{
		"proxy-target-config": targetConfigJson,
	})

	var targetConfig TargetConfig
	err := json.Unmarshal([]byte(permissions.CriticalOptions["proxy-target-config"]), &targetConfig)
	if err != nil {
		logger.Error("unmarshal-failed", err)
		return nil, nil, nil, err
	}

	nConn, err := net.Dial("tcp", targetConfig.Address)
	if err != nil {
		logger.Error("dial-failed", err)
		return nil, nil, nil, err
	}

	clientConfig := &ssh.ClientConfig{}

	if targetConfig.User != "" {
		clientConfig.User = targetConfig.User
	}

	if targetConfig.PrivateKey != "" {
		key, err := ssh.ParsePrivateKey([]byte(targetConfig.PrivateKey))
		if err != nil {
			logger.Error("parsing-key-failed", err)
			return nil, nil, nil, err
		}
		clientConfig.Auth = append(clientConfig.Auth, ssh.PublicKeys(key))
	}

	if targetConfig.User != "" && targetConfig.Password != "" {
		clientConfig.Auth = append(clientConfig.Auth, ssh.Password(targetConfig.Password))
	}

	if targetConfig.HostFingerprint != "" {
		clientConfig.HostKeyCallback = func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			expectedFingerprint := targetConfig.HostFingerprint

			var actualFingerprint string
			switch utf8.RuneCountInString(expectedFingerprint) {
			case helpers.MD5_FINGERPRINT_LENGTH:
				actualFingerprint = helpers.MD5Fingerprint(key)
			case helpers.SHA1_FINGERPRINT_LENGTH:
				actualFingerprint = helpers.SHA1Fingerprint(key)
			}

			if expectedFingerprint != actualFingerprint {
				err := errors.New("Host fingerprint mismatch")
				logger.Error("host-key-fingerprint-mismatch", err)
				return err
			}

			return nil
		}
	}

	conn, ch, req, err := ssh.NewClientConn(nConn, targetConfig.Address, clientConfig)
	if err != nil {
		logger.Error("handshake-failed", err)
		return nil, nil, nil, err
	}

	return conn, ch, req, nil
}
<commit_msg>Remove test debugging log line<commit_after>package proxy

import (
	"encoding/json"
	"errors"
	"net"
	"strconv"
	"sync"
	"unicode/utf8"

	"github.com/cloudfoundry-incubator/diego-ssh/helpers"
	"github.com/cloudfoundry-incubator/runtime-schema/metric"
	"github.com/cloudfoundry/dropsonde/logs"
	"github.com/pivotal-golang/lager"
	"golang.org/x/crypto/ssh"
)

const (
	sshConnections = metric.Metric("ssh-connections")
)

type Waiter interface {
	Wait() error
}

type TargetConfig struct {
	Address         string `json:"address"`
	HostFingerprint string `json:"host_fingerprint"`
	User            string `json:"user,omitempty"`
	Password        string `json:"password,omitempty"`
	PrivateKey      string `json:"private_key,omitempty"`
}

type LogMessage struct {
	Guid    string `json:"guid"`
	Message string `json:"message"`
	Index   int    `json:"index"`
}

type Proxy struct {
	logger       lager.Logger
	serverConfig *ssh.ServerConfig

	connectionLock *sync.Mutex
	connections    int
}

func New(
	logger lager.Logger,
	serverConfig *ssh.ServerConfig,
) *Proxy {
	return &Proxy{
		logger:         logger,
		serverConfig:   serverConfig,
		connectionLock: &sync.Mutex{},
	}
}

func (p *Proxy) HandleConnection(netConn net.Conn) {
	logger := p.logger.Session("handle-connection")
	defer netConn.Close()

	serverConn, serverChannels, serverRequests, err := ssh.NewServerConn(netConn, p.serverConfig)
	if err != nil {
		return
	}
	defer serverConn.Close()

	clientConn, clientChannels, clientRequests, err := NewClientConn(logger, serverConn.Permissions)
	if err != nil {
		return
	}
	defer clientConn.Close()

	emitLogMessage(logger, serverConn.Permissions)

	fromClientLogger := logger.Session("from-client")
	fromDaemonLogger := logger.Session("from-daemon")

	go ProxyGlobalRequests(fromClientLogger, clientConn, serverRequests)
	go ProxyGlobalRequests(fromDaemonLogger, serverConn, clientRequests)

	go ProxyChannels(fromClientLogger, clientConn, serverChannels)
	go ProxyChannels(fromDaemonLogger, serverConn, clientChannels)

	p.connectionLock.Lock()
	p.connections++
	err = sshConnections.Send(p.connections)
	if err != nil {
		logger.Error("failed-to-send-ssh-connections-metric", err)
	}
	p.connectionLock.Unlock()

	defer func() {
		p.emitConnectionClosing(logger)
	}()

	Wait(logger, serverConn, clientConn)
}

func (p *Proxy) emitConnectionClosing(logger lager.Logger) {
	p.connectionLock.Lock()
	p.connections--
	err := sshConnections.Send(p.connections)
	p.connectionLock.Unlock()

	if err != nil {
		logger.Error("failed-to-send-ssh-connections-metric", err)
	}
}

func emitLogMessage(logger lager.Logger, perms *ssh.Permissions) {
	logMessageJson := perms.CriticalOptions["log-message"]
	if logMessageJson == "" {
		return
	}

	logMessage := &LogMessage{}
	err := json.Unmarshal([]byte(logMessageJson), logMessage)
	if err != nil {
		logger.Error("json-unmarshal-failed", err)
		return
	}

	logs.SendAppLog(logMessage.Guid, logMessage.Message, "SSH", strconv.Itoa(logMessage.Index))
}

func ProxyGlobalRequests(logger lager.Logger, conn ssh.Conn, reqs <-chan *ssh.Request) {
	logger = logger.Session("proxy-global-requests")

	logger.Info("started")
	defer logger.Info("completed")

	for req := range reqs {
		logger.Info("request", lager.Data{
			"type":      req.Type,
			"wantReply": req.WantReply,
			"payload":   req.Payload,
		})

		success, reply, err := conn.SendRequest(req.Type, req.WantReply, req.Payload)
		if err != nil {
			logger.Error("send-request-failed", err)
			continue
		}

		if req.WantReply {
			req.Reply(success, reply)
		}
	}
}

func ProxyChannels(logger lager.Logger, conn ssh.Conn, channels <-chan ssh.NewChannel) {
	logger = logger.Session("proxy-channels")

	logger.Info("started")
	defer func() {
		logger.Info("completed")
		conn.Close()
	}()

	for newChannel := range channels {
		handleNewChannel(logger, conn, newChannel)
	}
}

func handleNewChannel(logger lager.Logger, conn ssh.Conn, newChannel ssh.NewChannel) {
	logger.Info("new-channel", lager.Data{
		"channelType": newChannel.ChannelType(),
		"extraData":   newChannel.ExtraData(),
	})

	targetChan, targetReqs, err := conn.OpenChannel(newChannel.ChannelType(), newChannel.ExtraData())
	if err != nil {
		logger.Error("failed-to-open-channel", err)
		if openErr, ok := err.(*ssh.OpenChannelError); ok {
			newChannel.Reject(openErr.Reason, openErr.Message)
		} else {
			newChannel.Reject(ssh.ConnectionFailed, err.Error())
		}
		return
	}

	sourceChan, sourceReqs, err := newChannel.Accept()
	if err != nil {
		targetChan.Close()
		return
	}

	toTargetLogger := logger.Session("to-target")
	toSourceLogger := logger.Session("to-source")

	targetWg := &sync.WaitGroup{}
	sourceWg := &sync.WaitGroup{}

	targetWg.Add(1)
	sourceWg.Add(1)

	go func() {
		helpers.Copy(toTargetLogger, targetWg, targetChan, sourceChan)
		targetChan.CloseWrite()
	}()

	go func() {
		helpers.Copy(toSourceLogger, sourceWg, sourceChan, targetChan)
		sourceChan.CloseWrite()
	}()

	go ProxyRequests(toTargetLogger, newChannel.ChannelType(), sourceReqs, targetChan, targetWg)
	go ProxyRequests(toSourceLogger, newChannel.ChannelType(), targetReqs, sourceChan, sourceWg)
}

func ProxyRequests(logger lager.Logger, channelType string, reqs <-chan *ssh.Request, channel ssh.Channel, wg *sync.WaitGroup) {
	logger = logger.Session("proxy-requests", lager.Data{
		"channel-type": channelType,
	})

	logger.Info("started")
	defer func() {
		logger.Info("completed")
		wg.Wait()
		channel.Close()
	}()

	for req := range reqs {
		logger.Info("request", lager.Data{
			"type":      req.Type,
			"wantReply": req.WantReply,
			"payload":   req.Payload,
		})
		success, err := channel.SendRequest(req.Type, req.WantReply, req.Payload)
		if err != nil {
			logger.Error("send-request-failed", err)
			continue
		}

		if req.WantReply {
			req.Reply(success, nil)
		}

		if req.Type == "exit-status" {
			return
		}
	}
}

func Wait(logger lager.Logger, waiters ...Waiter) {
	wg := &sync.WaitGroup{}
	for _, waiter := range waiters {
		wg.Add(1)
		go func(waiter Waiter) {
			waiter.Wait()
			wg.Done()
		}(waiter)
	}
	wg.Wait()
}

func NewClientConn(logger lager.Logger, permissions *ssh.Permissions) (ssh.Conn, <-chan ssh.NewChannel, <-chan *ssh.Request, error) {
	if permissions == nil || permissions.CriticalOptions == nil {
		err := errors.New("Invalid permissions from authentication")
		logger.Error("permissions-and-critical-options-required", err)
		return nil, nil, nil, err
	}

	targetConfigJson := permissions.CriticalOptions["proxy-target-config"]
	logger = logger.Session("new-client-conn", lager.Data{
		"proxy-target-config": targetConfigJson,
	})

	var targetConfig TargetConfig
	err := json.Unmarshal([]byte(permissions.CriticalOptions["proxy-target-config"]), &targetConfig)
	if err != nil {
		logger.Error("unmarshal-failed", err)
		return nil, nil, nil, err
	}

	nConn, err := net.Dial("tcp", targetConfig.Address)
	if err != nil {
		logger.Error("dial-failed", err)
		return nil, nil, nil, err
	}

	clientConfig := &ssh.ClientConfig{}

	if targetConfig.User != "" {
		clientConfig.User = targetConfig.User
	}

	if targetConfig.PrivateKey != "" {
		key, err := ssh.ParsePrivateKey([]byte(targetConfig.PrivateKey))
		if err != nil {
			logger.Error("parsing-key-failed", err)
			return nil, nil, nil, err
		}
		clientConfig.Auth = append(clientConfig.Auth, ssh.PublicKeys(key))
	}

	if targetConfig.User != "" && targetConfig.Password != "" {
		clientConfig.Auth = append(clientConfig.Auth, ssh.Password(targetConfig.Password))
	}

	if targetConfig.HostFingerprint != "" {
		clientConfig.HostKeyCallback = func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			expectedFingerprint := targetConfig.HostFingerprint

			var actualFingerprint string
			switch utf8.RuneCountInString(expectedFingerprint) {
			case helpers.MD5_FINGERPRINT_LENGTH:
				actualFingerprint = helpers.MD5Fingerprint(key)
			case helpers.SHA1_FINGERPRINT_LENGTH:
				actualFingerprint = helpers.SHA1Fingerprint(key)
			}

			if expectedFingerprint != actualFingerprint {
				err := errors.New("Host fingerprint mismatch")
				logger.Error("host-key-fingerprint-mismatch", err)
				return err
			}

			return nil
		}
	}

	conn, ch, req, err := ssh.NewClientConn(nConn, targetConfig.Address, clientConfig)
	if err != nil {
		logger.Error("handshake-failed", err)
		return nil, nil, nil, err
	}

	return conn, ch, req, nil
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"log"
	"net/http"

	"github.com/gorilla/mux"
	_ "github.com/lib/pq"
)

var servport = ":6862"

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("[X] %s, %s", msg, err)
		panic(fmt.Sprintf("%s, %s", msg, err))
	}
}

func main() {
	log.Println("[i] Server started")

	// Connect to DB
	db, err := sql.Open("postgres", "user=appread dbname='quantifiedSelf' sslmode=disable")
	failOnError(err, "Error connecting to database")
	defer db.Close()

	http.HandleFunc("/data/all/", func(w http.ResponseWriter, r *http.Request) {
		// TODO: Allow filtering via URL
		// Get rows from DB
		var output string
		err := db.QueryRow(`SELECT json_agg(r) FROM (SELECT * FROM trello.cards) r;`).Scan(&output)
		if err != nil {
			log.Println("Error retriving from DB, ", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintln(w, "Error retriving from DB, ", err)
			return
		}

		// Print out returned
		fmt.Fprint(w, output)
	})

	// Restful handler
	r := mux.NewRouter()
	r.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "dla;jfkdlsajflkdsa;jfk;ldsajfklds;a")
	})

	r.HandleFunc("/api/totals/last/{num}", func(w http.ResponseWriter, r *http.Request) {
		// Grab vars
		vars := mux.Vars(r)

		var output string
		// This is bad... don't do this.... omg
		query := fmt.Sprintf(`SELECT json_agg(r) FROM (select to_char(day, 'YY Mon DD') as day, end_of_day_total from trello.dailytallies order by day DESC limit %s) r;`, vars["num"])
		err := db.QueryRow(query).Scan(&output)

		if err != nil {
			log.Println("Error retriving from DB, ", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintln(w, "Error retriving from DB, ", err)
			return
		}

		// Print out returned
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, output)
	})
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("../ui")))
	http.Handle("/", r)

	// Die gracefully
	// killchan := make(chan os.Signal)
	// signal.Notify(killchan, os.Interrupt, os.Kill)

	log.Println("[i] Serving on ", servport, "\n\tWaiting...")

	log.Fatal(http.ListenAndServe(servport, nil))
	// <-killchan

	log.Println("[i] Shutting down...")
}
<commit_msg>Changed date format a little to make the year more distinct from the day of month<commit_after>package main

import (
	"database/sql"
	"fmt"
	"log"
	"net/http"

	"github.com/gorilla/mux"
	_ "github.com/lib/pq"
)

var servport = ":6862"

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("[X] %s, %s", msg, err)
		panic(fmt.Sprintf("%s, %s", msg, err))
	}
}

func main() {
	log.Println("[i] Server started")

	// Connect to DB
	db, err := sql.Open("postgres", "user=appread dbname='quantifiedSelf' sslmode=disable")
	failOnError(err, "Error connecting to database")
	defer db.Close()

	http.HandleFunc("/data/all/", func(w http.ResponseWriter, r *http.Request) {
		// TODO: Allow filtering via URL
		// Get rows from DB
		var output string
		err := db.QueryRow(`SELECT json_agg(r) FROM (SELECT * FROM trello.cards) r;`).Scan(&output)
		if err != nil {
			log.Println("Error retriving from DB, ", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintln(w, "Error retriving from DB, ", err)
			return
		}

		// Print out returned
		fmt.Fprint(w, output)
	})

	// Restful handler
	r := mux.NewRouter()
	r.HandleFunc("/api", func(w http.ResponseWriter, r *http.Request) {
		fmt.Fprintln(w, "dla;jfkdlsajflkdsa;jfk;ldsajfklds;a")
	})

	r.HandleFunc("/api/totals/last/{num}", func(w http.ResponseWriter, r *http.Request) {
		// Grab vars
		vars := mux.Vars(r)

		var output string
		// This is bad... don't do this.... omg
		query := fmt.Sprintf(`SELECT json_agg(r) FROM (select to_char(day, 'YY-Mon DD') as day, end_of_day_total from trello.dailytallies order by day DESC limit %s) r;`, vars["num"])
		err := db.QueryRow(query).Scan(&output)

		if err != nil {
			log.Println("Error retriving from DB, ", err)
			w.WriteHeader(http.StatusInternalServerError)
			fmt.Fprintln(w, "Error retriving from DB, ", err)
			return
		}

		// Print out returned
		w.Header().Set("Content-Type", "application/json")
		fmt.Fprint(w, output)
	})
	r.PathPrefix("/").Handler(http.FileServer(http.Dir("../ui")))
	http.Handle("/", r)

	// Die gracefully
	// killchan := make(chan os.Signal)
	// signal.Notify(killchan, os.Interrupt, os.Kill)

	log.Println("[i] Serving on ", servport, "\n\tWaiting...")

	log.Fatal(http.ListenAndServe(servport, nil))
	// <-killchan

	log.Println("[i] Shutting down...")
}
<|endoftext|>


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

// Package gcslock is a scalable, distributed mutex that can be used
// to serialize computations anywhere on the global internet.
package gcslock

import (
	"bytes"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

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

const (
	defaultStorageLockURL   = "https://www.googleapis.com/upload/storage/v1"
	defaultStorageUnlockURL = "https://www.googleapis.com/storage/v1"
)

var (
	// These vars are used in the requests below. Having separate default
	// values makes it easy to reset the standard config during testing.
	storageLockURL   = defaultStorageLockURL
	storageUnlockURL = defaultStorageUnlockURL
)

// ContextLocker provides an extension of the sync.Locker interface.
type ContextLocker interface {
	sync.Locker
	ContextLock(context.Context) error
	ContextUnlock(context.Context) error
}

type mutex struct {
	bucket string
	object string
	client *http.Client
}

var _ ContextLocker = (*mutex)(nil)

// Lock waits indefinitely to acquire a mutex.
func (m *mutex) Lock() {
	m.ContextLock(context.Background())
}

// ContextLock waits indefinitely to acquire a mutex with timeout
// governed by passed context.
func (m *mutex) ContextLock(ctx context.Context) error {
	q := url.Values{
		"name":              {m.object},
		"uploadType":        {"media"},
		"ifGenerationMatch": {"0"},
	}
	url := fmt.Sprintf("%s/b/%s/o?%s", storageLockURL, m.bucket, q.Encode())
	// Starting w/ 10ms, backoff and retry exponentially.
	for i := 10; ; i *= 2 {
		done := make(chan bool)
		go func() {
			res, err := m.client.Post(url, "plain/text", bytes.NewReader([]byte("1")))
			if err == nil {
				res.Body.Close()
				if res.StatusCode == 200 {
					done <- true
					return
				}
			}
			done <- false
		}()
		select {
		case <-time.After(time.Duration(i) * time.Millisecond):
			continue
		case <-ctx.Done():
			return ctx.Err()
		case d := <-done:
			if d {
				return nil
			}
			continue
		}
	}
}

// Unlock waits indefinitely to release a mutex.
func (m *mutex) Unlock() {
	m.ContextUnlock(context.Background())
}

// ContextUnlock waits indefinitely to release a mutex with timeout
// governed by passed context.
func (m *mutex) ContextUnlock(ctx context.Context) error {
	url := fmt.Sprintf("%s/b/%s/o/%s?", storageUnlockURL, m.bucket, m.object)
	// Starting w/ 10ms, backoff and retry exponentially.
	for i := 10; ; i *= 2 {
		done := make(chan bool)
		go func() {
			req, err := http.NewRequest("DELETE", url, nil)
			if err == nil {
				res, err := m.client.Do(req)
				if err == nil {
					res.Body.Close()
					if res.StatusCode == 204 {
						done <- true
						return
					}
				}
			}
			done <- false
		}()
		select {
		case <-time.After(time.Duration(i) * time.Millisecond):
			continue
		case <-ctx.Done():
			return ctx.Err()
		case d := <-done:
			if d {
				return nil
			}
			continue
		}
	}
}

// httpClient is overwritten in tests
var httpClient = func(ctx context.Context) (*http.Client, error) {
	const scope = "https://www.googleapis.com/auth/devstorage.full_control"
	return google.DefaultClient(ctx, scope)
}

// New creates a GCS-based sync.Locker.
// It uses Application Default Credentials to make authenticated requests
// to Google Cloud Storage. See the DefaultClient function of the
// golang.org/x/oauth2/google package for App Default Credentials details.
//
// If ctx argument is nil, context.Background is used.
//
func New(ctx context.Context, bucket, object string) (ContextLocker, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	client, err := httpClient(ctx)
	if err != nil {
		return nil, err
	}
	m := &mutex{
		bucket: bucket,
		object: object,
		client: client,
	}
	return m, nil
}
<commit_msg>fixed ddos comments up through commit 14<commit_after>// Copyright 2016 Google Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package gcslock is a scalable, distributed mutex that can be used
// to serialize computations anywhere on the global internet.
package gcslock

import (
	"bytes"
	"fmt"
	"net/http"
	"net/url"
	"sync"
	"time"

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

const (
	defaultStorageLockURL   = "https://www.googleapis.com/upload/storage/v1"
	defaultStorageUnlockURL = "https://www.googleapis.com/storage/v1"
)

var (
	// These vars are used in the requests below. Having separate default
	// values makes it easy to reset the standard config during testing.
	storageLockURL   = defaultStorageLockURL
	storageUnlockURL = defaultStorageUnlockURL
)

// ContextLocker provides an extension of the sync.Locker interface.
type ContextLocker interface {
	sync.Locker
	ContextLock(context.Context) error
	ContextUnlock(context.Context) error
}

type mutex struct {
	bucket string
	object string
	client *http.Client
}

var _ ContextLocker = (*mutex)(nil)

// Lock waits indefinitely to acquire a mutex.
func (m *mutex) Lock() {
	m.ContextLock(context.Background())
}

// ContextLock waits indefinitely to acquire a mutex with timeout
// governed by passed context.
func (m *mutex) ContextLock(ctx context.Context) error {
	q := url.Values{
		"name":              {m.object},
		"uploadType":        {"media"},
		"ifGenerationMatch": {"0"},
	}
	url := fmt.Sprintf("%s/b/%s/o?%s", storageLockURL, m.bucket, q.Encode())
	// NOTE: ctx deadline/timeout and backoff are independent. The former is
	// an aggregate timeout and the latter is a per loop iteration delay.
	backoff := 10 * time.Millisecond
	for {
		req, err := http.NewRequest("POST", url, bytes.NewReader([]byte("1")))
		if err != nil {
			// Likely malformed URL - retry won't fix so return.
			return err
		}
		req.Header.Set("content-type", "text/plain")
		req = req.WithContext(ctx)
		res, err := m.client.Do(req)
		if err == nil {
			res.Body.Close()
			if res.StatusCode == 200 {
				return nil
			}
		}
		select {
		case <-time.After(backoff):
			backoff *= 2
			continue
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// Unlock waits indefinitely to release a mutex.
func (m *mutex) Unlock() {
	m.ContextUnlock(context.Background())
}

// ContextUnlock waits indefinitely to release a mutex with timeout
// governed by passed context.
func (m *mutex) ContextUnlock(ctx context.Context) error {
	url := fmt.Sprintf("%s/b/%s/o/%s?", storageUnlockURL, m.bucket, m.object)
	// NOTE: ctx deadline/timeout and backoff are independent. The former is
	// an aggregate timeout and the latter is a per loop iteration delay.
	backoff := 10 * time.Millisecond
	for {
		req, err := http.NewRequest("DELETE", url, nil)
		if err != nil {
			// Likely malformed URL - retry won't fix so return.
			return err
		}
		req.Header.Set("content-type", "text/plain")
		req = req.WithContext(ctx)
		res, err := m.client.Do(req)
		if err == nil {
			res.Body.Close()
			if res.StatusCode == 204 {
				return nil
			}
		}
		select {
		case <-time.After(backoff):
			backoff *= 2
			continue
		case <-ctx.Done():
			return ctx.Err()
		}
	}
}

// httpClient is overwritten in tests
var httpClient = func(ctx context.Context) (*http.Client, error) {
	const scope = "https://www.googleapis.com/auth/devstorage.full_control"
	return google.DefaultClient(ctx, scope)
}

// New creates a GCS-based sync.Locker.
// It uses Application Default Credentials to make authenticated requests
// to Google Cloud Storage. See the DefaultClient function of the
// golang.org/x/oauth2/google package for App Default Credentials details.
//
// If ctx argument is nil, context.Background is used.
//
func New(ctx context.Context, bucket, object string) (ContextLocker, error) {
	if ctx == nil {
		ctx = context.Background()
	}
	client, err := httpClient(ctx)
	if err != nil {
		return nil, err
	}
	m := &mutex{
		bucket: bucket,
		object: object,
		client: client,
	}
	return m, nil
}
<|endoftext|>


<commit_before>package detector

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	"github.com/mesos/mesos-go/upid"
)

var (
	pluginLock     sync.Mutex
	plugins        = map[string]PluginFactory{}
	EmptySpecError = errors.New("empty master specification")

	defaultFactory = PluginFactory(func(spec string) (Master, error) {
		if len(spec) == 0 {
			return nil, EmptySpecError
		}
		if strings.Index(spec, "@") < 0 {
			spec = "master@" + spec
		}
		if pid, err := upid.Parse(spec); err == nil {
			return NewStandalone(CreateMasterInfo(pid)), nil
		} else {
			return nil, err
		}
	})
)

type PluginFactory func(string) (Master, error)

// associates a plugin implementation with a Master specification prefix.
// packages that provide plugins are expected to invoke this func within
// their init() implementation. schedulers that wish to support plugins may
// anonymously import ("_") a package the auto-registers said plugins.
func Register(prefix string, f PluginFactory) error {
	if prefix == "" {
		return fmt.Errorf("illegal prefix: '%v'", prefix)
	}
	if f == nil {
		return fmt.Errorf("nil plugin factories are not allowed")
	}

	pluginLock.Lock()
	defer pluginLock.Unlock()

	if _, found := plugins[prefix]; found {
		return fmt.Errorf("detection plugin already registered for prefix '%s'", prefix)
	}
	plugins[prefix] = f
	return nil
}

// Create a new detector given the provided specification. Examples are:
//
//   - file://{path_to_local_file}
//   - {ipaddress}:{port}
//   - master@{ip_address}:{port}
//   - master({id})@{ip_address}:{port}
//
// Support for the file:// prefix is intentionally hardcoded so that it may
// not be inadvertently overridden by a custom plugin implementation. Custom
// plugins are supported via the Register and MatchingPlugin funcs.
//
// Furthermore it is expected that master detectors returned from this func
// are not yet running and will only begin to spawn requisite background
// processing upon, or some time after, the first invocation of their Detect.
//
func New(spec string) (m Master, err error) {
	if strings.HasPrefix(spec, "file://") {
		var body []byte
		path := spec[7:]
		body, err = ioutil.ReadFile(path)
		if err != nil {
			log.V(1).Infof("failed to read from file at '%s'", path)
		} else {
			m, err = New(string(body))
		}
	} else if f, ok := MatchingPlugin(spec); ok {
		m, err = f(spec)
	} else {
		m, err = defaultFactory(spec)
	}

	return
}

func MatchingPlugin(spec string) (PluginFactory, bool) {
	pluginLock.Lock()
	defer pluginLock.Unlock()

	for prefix, f := range plugins {
		if strings.HasPrefix(spec, prefix) {
			return f, true
		}
	}
	return nil, false
}

// Super-useful utility func that attempts to build a mesos.MasterInfo from a
// upid.UPID specification. An attempt is made to determine the IP address of
// the UPID's Host and any errors during such resolution will result in a nil
// returned result. A nil result is also returned upon errors parsing the Port
// specification of the UPID.
//
// TODO(jdef) make this a func of upid.UPID so that callers can invoke somePid.MasterInfo()?
//
func CreateMasterInfo(pid *upid.UPID) *mesos.MasterInfo {
	if pid == nil {
		return nil
	}
	port, err := strconv.Atoi(pid.Port)
	if err != nil {
		log.Errorf("failed to parse port: %v", err)
		return nil
	}
	//TODO(jdef) what about (future) ipv6 support?
	var ipv4 net.IP
	if addrs, err := net.LookupIP(pid.Host); err == nil {
		for _, ip := range addrs {
			if ip = ip.To4(); ip != nil {
				ipv4 = ip
				break
			}
		}
		if ipv4 == nil {
			log.Errorf("host does not resolve to an IPv4 address: %v", pid.Host)
			return nil
		}
	} else {
		log.Errorf("failed to lookup IPs for host '%v': %v", pid.Host, err)
		return nil
	}
	packedip := binary.BigEndian.Uint32(ipv4) // network byte order is big-endian
	mi := util.NewMasterInfo(pid.ID, packedip, uint32(port))
	mi.Pid = proto.String(pid.String())
	if pid.Host != "" {
		mi.Hostname = proto.String(pid.Host)
	}
	return mi
}
<commit_msg>Support PID hosts with IPs<commit_after>package detector

import (
	"encoding/binary"
	"errors"
	"fmt"
	"io/ioutil"
	"net"
	"strconv"
	"strings"
	"sync"

	"github.com/gogo/protobuf/proto"
	log "github.com/golang/glog"
	mesos "github.com/mesos/mesos-go/mesosproto"
	util "github.com/mesos/mesos-go/mesosutil"
	"github.com/mesos/mesos-go/upid"
)

var (
	pluginLock     sync.Mutex
	plugins        = map[string]PluginFactory{}
	EmptySpecError = errors.New("empty master specification")

	defaultFactory = PluginFactory(func(spec string) (Master, error) {
		if len(spec) == 0 {
			return nil, EmptySpecError
		}
		if strings.Index(spec, "@") < 0 {
			spec = "master@" + spec
		}
		if pid, err := upid.Parse(spec); err == nil {
			return NewStandalone(CreateMasterInfo(pid)), nil
		} else {
			return nil, err
		}
	})
)

type PluginFactory func(string) (Master, error)

// associates a plugin implementation with a Master specification prefix.
// packages that provide plugins are expected to invoke this func within
// their init() implementation. schedulers that wish to support plugins may
// anonymously import ("_") a package the auto-registers said plugins.
func Register(prefix string, f PluginFactory) error {
	if prefix == "" {
		return fmt.Errorf("illegal prefix: '%v'", prefix)
	}
	if f == nil {
		return fmt.Errorf("nil plugin factories are not allowed")
	}

	pluginLock.Lock()
	defer pluginLock.Unlock()

	if _, found := plugins[prefix]; found {
		return fmt.Errorf("detection plugin already registered for prefix '%s'", prefix)
	}
	plugins[prefix] = f
	return nil
}

// Create a new detector given the provided specification. Examples are:
//
//   - file://{path_to_local_file}
//   - {ipaddress}:{port}
//   - master@{ip_address}:{port}
//   - master({id})@{ip_address}:{port}
//
// Support for the file:// prefix is intentionally hardcoded so that it may
// not be inadvertently overridden by a custom plugin implementation. Custom
// plugins are supported via the Register and MatchingPlugin funcs.
//
// Furthermore it is expected that master detectors returned from this func
// are not yet running and will only begin to spawn requisite background
// processing upon, or some time after, the first invocation of their Detect.
//
func New(spec string) (m Master, err error) {
	if strings.HasPrefix(spec, "file://") {
		var body []byte
		path := spec[7:]
		body, err = ioutil.ReadFile(path)
		if err != nil {
			log.V(1).Infof("failed to read from file at '%s'", path)
		} else {
			m, err = New(string(body))
		}
	} else if f, ok := MatchingPlugin(spec); ok {
		m, err = f(spec)
	} else {
		m, err = defaultFactory(spec)
	}

	return
}

func MatchingPlugin(spec string) (PluginFactory, bool) {
	pluginLock.Lock()
	defer pluginLock.Unlock()

	for prefix, f := range plugins {
		if strings.HasPrefix(spec, prefix) {
			return f, true
		}
	}
	return nil, false
}

// Super-useful utility func that attempts to build a mesos.MasterInfo from a
// upid.UPID specification. An attempt is made to determine the IP address of
// the UPID's Host and any errors during such resolution will result in a nil
// returned result. A nil result is also returned upon errors parsing the Port
// specification of the UPID.
//
// TODO(jdef) make this a func of upid.UPID so that callers can invoke somePid.MasterInfo()?
//
func CreateMasterInfo(pid *upid.UPID) *mesos.MasterInfo {
	if pid == nil {
		return nil
	}
	port, err := strconv.Atoi(pid.Port)
	if err != nil {
		log.Errorf("failed to parse port: %v", err)
		return nil
	}
	//TODO(jdef) what about (future) ipv6 support?
	var ipv4 net.IP
	if ipv4 = net.ParseIP(pid.Host); ipv4 != nil {
		// Intentionally left empty. ipv4 is successfully parsed, so we can skip the other versions
        } else if addrs, err := net.LookupIP(pid.Host); err == nil {
		for _, ip := range addrs {
			if ip = ip.To4(); ip != nil {
				ipv4 = ip
				break
			}
		}
		if ipv4 == nil {
			log.Errorf("host does not resolve to an IPv4 address: %v", pid.Host)
			return nil
		}
	} else {
		log.Errorf("failed to lookup IPs for host '%v': %v", pid.Host, err)
		return nil
	}
	packedip := binary.BigEndian.Uint32(ipv4) // network byte order is big-endian
	mi := util.NewMasterInfo(pid.ID, packedip, uint32(port))
	mi.Pid = proto.String(pid.String())
	if pid.Host != "" {
		mi.Hostname = proto.String(pid.Host)
	}
	return mi
}
<|endoftext|>


<commit_before><commit_msg>Obliterate mappings/env variables if passed through, since we want to completely replace those arrays<commit_after><|endoftext|>


<commit_before><commit_msg>Fix differences in the usage information of the CLI<commit_after><|endoftext|>


<commit_before>package filepathfilter

import (
	"strings"

	"github.com/git-lfs/wildmatch"
	"github.com/rubyist/tracerx"
)

type Pattern interface {
	Match(filename string) bool
	// String returns a string representation (see: regular expressions) of
	// the underlying pattern used to match filenames against this Pattern.
	String() string
}

type Filter struct {
	include []Pattern
	exclude []Pattern
}

func NewFromPatterns(include, exclude []Pattern) *Filter {
	return &Filter{include: include, exclude: exclude}
}

func New(include, exclude []string) *Filter {
	return NewFromPatterns(
		convertToWildmatch(include),
		convertToWildmatch(exclude))
}

// Include returns the result of calling String() on each Pattern in the
// include set of this *Filter.
func (f *Filter) Include() []string { return wildmatchToString(f.include...) }

// Exclude returns the result of calling String() on each Pattern in the
// exclude set of this *Filter.
func (f *Filter) Exclude() []string { return wildmatchToString(f.exclude...) }

// wildmatchToString maps the given set of Pattern's to a string slice by
// calling String() on each pattern.
func wildmatchToString(ps ...Pattern) []string {
	s := make([]string, 0, len(ps))
	for _, p := range ps {
		s = append(s, p.String())
	}

	return s
}

func (f *Filter) Allows(filename string) bool {
	if f == nil {
		return true
	}

	var included bool
	for _, inc := range f.include {
		if included = inc.Match(filename); included {
			break
		}
	}

	tracerx.Printf("filepathfilter: rejecting %q via %v", filename, f.include)
	if !included && len(f.include) > 0 {
		return false
	}

	for _, ex := range f.exclude {
		if ex.Match(filename) {
			tracerx.Printf("filepathfilter: rejecting %q via %q", filename, ex.String())
			return false
		}
	}

	tracerx.Printf("filepathfilter: accepting %q", filename)
	return true
}

type wm struct {
	w    *wildmatch.Wildmatch
	p    string
	dirs bool
}

func (w *wm) Match(filename string) bool {
	return w.w.Match(w.chomp(filename))
}

func (w *wm) chomp(filename string) string {
	for _, suffix := range []string{`/`, `\`} {
		if strings.HasSuffix(filename, suffix) {
			return strings.TrimSuffix(filename, suffix)
		}
	}
	return filename
}

func (w *wm) String() string {
	return w.p
}

const (
	sep byte = '/'
)

func NewPattern(p string) Pattern {
	pp := p

	// Special case: the below patterns match anything according to existing
	// behavior.
	switch pp {
	case `*`, `*.*`, `.`, `./`, `.\`:
		pp = join("**", "*")
	}

	dirs := strings.Contains(pp, string(sep))
	rooted := strings.HasPrefix(pp, string(sep))
	wild := strings.Contains(pp, "*")

	if !dirs && !wild {
		// Special case: if pp is a literal string (optionally including
		// a character class), assume it is a substring match.
		pp = join("**", pp, "**")
	} else {
		if dirs && !rooted {
			// Special case: if there are any directory separators,
			// assume "pp" is rooted.
			if !wild {
				pp = join("**", pp, "**")
			}
		} else {
			if rooted {
				// Special case: if there are not any directory
				// separators, assume "pp" is a substring match.
				pp = join(pp, "**")
			} else {
				// Special case: if there are not any directory
				// separators, assume "pp" is a substring match.
				pp = join("**", pp)
			}
		}
	}
	tracerx.Printf("filepathfilter: rewrite %q as %q", p, pp)

	return &wm{
		p: p,
		w: wildmatch.NewWildmatch(
			pp,
			wildmatch.SystemCase,
		),
		dirs: dirs,
	}
}

// join joins path elements together via the separator "sep" and produces valid
// paths without multiple separators (unless multiple separators were included
// in the original paths []string).
func join(paths ...string) string {
	var joined string

	for i, path := range paths {
		joined = joined + path
		if i != len(paths)-1 && !strings.HasSuffix(path, string(sep)) {
			joined = joined + string(sep)
		}
	}

	return joined
}

func convertToWildmatch(rawpatterns []string) []Pattern {
	patterns := make([]Pattern, len(rawpatterns))
	for i, raw := range rawpatterns {
		patterns[i] = NewPattern(raw)
	}
	return patterns
}
<commit_msg>filepathfilter: remove filepath seaprator from pattern subjects<commit_after>package filepathfilter

import (
	"path/filepath"
	"strings"

	"github.com/git-lfs/wildmatch"
	"github.com/rubyist/tracerx"
)

type Pattern interface {
	Match(filename string) bool
	// String returns a string representation (see: regular expressions) of
	// the underlying pattern used to match filenames against this Pattern.
	String() string
}

type Filter struct {
	include []Pattern
	exclude []Pattern
}

func NewFromPatterns(include, exclude []Pattern) *Filter {
	return &Filter{include: include, exclude: exclude}
}

func New(include, exclude []string) *Filter {
	return NewFromPatterns(
		convertToWildmatch(include),
		convertToWildmatch(exclude))
}

// Include returns the result of calling String() on each Pattern in the
// include set of this *Filter.
func (f *Filter) Include() []string { return wildmatchToString(f.include...) }

// Exclude returns the result of calling String() on each Pattern in the
// exclude set of this *Filter.
func (f *Filter) Exclude() []string { return wildmatchToString(f.exclude...) }

// wildmatchToString maps the given set of Pattern's to a string slice by
// calling String() on each pattern.
func wildmatchToString(ps ...Pattern) []string {
	s := make([]string, 0, len(ps))
	for _, p := range ps {
		s = append(s, p.String())
	}

	return s
}

func (f *Filter) Allows(filename string) bool {
	if f == nil {
		return true
	}

	var included bool
	for _, inc := range f.include {
		if included = inc.Match(filename); included {
			break
		}
	}

	tracerx.Printf("filepathfilter: rejecting %q via %v", filename, f.include)
	if !included && len(f.include) > 0 {
		return false
	}

	for _, ex := range f.exclude {
		if ex.Match(filename) {
			tracerx.Printf("filepathfilter: rejecting %q via %q", filename, ex.String())
			return false
		}
	}

	tracerx.Printf("filepathfilter: accepting %q", filename)
	return true
}

type wm struct {
	w    *wildmatch.Wildmatch
	p    string
	dirs bool
}

func (w *wm) Match(filename string) bool {
	return w.w.Match(w.chomp(filename))
}

func (w *wm) chomp(filename string) string {
	return strings.TrimSuffix(filename, string(filepath.Separator))
}

func (w *wm) String() string {
	return w.p
}

const (
	sep byte = '/'
)

func NewPattern(p string) Pattern {
	pp := p

	// Special case: the below patterns match anything according to existing
	// behavior.
	switch pp {
	case `*`, `*.*`, `.`, `./`, `.\`:
		pp = join("**", "*")
	}

	dirs := strings.Contains(pp, string(sep))
	rooted := strings.HasPrefix(pp, string(sep))
	wild := strings.Contains(pp, "*")

	if !dirs && !wild {
		// Special case: if pp is a literal string (optionally including
		// a character class), assume it is a substring match.
		pp = join("**", pp, "**")
	} else {
		if dirs && !rooted {
			// Special case: if there are any directory separators,
			// assume "pp" is rooted.
			if !wild {
				pp = join("**", pp, "**")
			}
		} else {
			if rooted {
				// Special case: if there are not any directory
				// separators, assume "pp" is a substring match.
				pp = join(pp, "**")
			} else {
				// Special case: if there are not any directory
				// separators, assume "pp" is a substring match.
				pp = join("**", pp)
			}
		}
	}
	tracerx.Printf("filepathfilter: rewrite %q as %q", p, pp)

	return &wm{
		p: p,
		w: wildmatch.NewWildmatch(
			pp,
			wildmatch.SystemCase,
		),
		dirs: dirs,
	}
}

// join joins path elements together via the separator "sep" and produces valid
// paths without multiple separators (unless multiple separators were included
// in the original paths []string).
func join(paths ...string) string {
	var joined string

	for i, path := range paths {
		joined = joined + path
		if i != len(paths)-1 && !strings.HasSuffix(path, string(sep)) {
			joined = joined + string(sep)
		}
	}

	return joined
}

func convertToWildmatch(rawpatterns []string) []Pattern {
	patterns := make([]Pattern, len(rawpatterns))
	for i, raw := range rawpatterns {
		patterns[i] = NewPattern(raw)
	}
	return patterns
}
<|endoftext|>


<commit_before><commit_msg>More aggressive resets of state in xerial reader (#322)<commit_after><|endoftext|>


<commit_before>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package smart

import (
	"errors"
	"fmt"
	"strings"

	"github.com/intelsdi-x/pulse/control/plugin"
	"github.com/intelsdi-x/pulse/control/plugin/cpolicy"
)

const (
	name       = "pulse-collector-smart"
	version    = 3
	pluginType = plugin.CollectorPluginType
)

var sysUtilProvider SysutilProvider = NewSysutilProvider()

var namespace_prefix = []string{"intel", "disk"}
var namespace_suffix = []string{"smart"}

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(name, version, pluginType, []string{plugin.PulseGOBContentType}, []string{plugin.PulseGOBContentType})
}

func NewSmartCollector() *SmartCollector {
	return &SmartCollector{}
}

func makeName(device, metric string) []string {
	splited := strings.Split(metric, "/")

	name := []string{}
	name = append(name, namespace_prefix...)
	name = append(name, device)
	name = append(name, namespace_suffix...)
	name = append(name, splited...)

	return name
}

func parseName(namespace []string) (disk, attribute string) {
	disk = namespace[len(namespace_prefix)]
	smart_namespace := namespace[len(namespace_prefix)+len(namespace_suffix)+1:]
	attribute = strings.Join(smart_namespace, "/")
	return
}

func validateName(namespace []string) bool {
	for i, v := range namespace_prefix {
		if namespace[i] != v {
			return false
		}
	}

	offset := len(namespace_prefix) + 1
	for i, v := range namespace_suffix {
		if namespace[offset+i] != v {
			return false
		}
	}

	return true
}

type SmartCollector struct {
}

type smartResults map[string]interface{}

// CollectMetrics returns metrics from smart
func (sc *SmartCollector) CollectMetrics(mts []plugin.PluginMetricType) ([]plugin.PluginMetricType, error) {
	buffered_results := map[string]smartResults{}

	results := make([]plugin.PluginMetricType, len(mts))

	for i, mt := range mts {
		if !validateName(mt.Namespace()) {
			return nil, errors.New(fmt.Sprintf("%v is not valid metric", mt.Namespace()))
		}
		disk, attribute_path := parseName(mt.Namespace())
		buffered, ok := buffered_results[disk]
		if !ok {
			values, err := ReadSmartData(disk, sysUtilProvider)
			if err != nil {
				return nil, err
			}
			buffered = values.GetAttributes()
			buffered_results[disk] = buffered
		}

		attribute, ok := buffered[attribute_path]

		if !ok {
			return nil, errors.New("Unknown attribute " + attribute_path)
		}

		results[i] = plugin.PluginMetricType{
			Namespace_: mt.Namespace(),
			Data_:      attribute,
		}
	}

	return results, nil
}

// GetMetricTypes returns the metric types exposed by smart
func (sc *SmartCollector) GetMetricTypes(_ plugin.PluginConfigType) ([]plugin.PluginMetricType, error) {
	smart_metrics := ListAllKeys()
	devices, err := sysUtilProvider.ListDevices()
	if err != nil {
		return nil, err
	}
	mts := make([]plugin.PluginMetricType, 0, len(smart_metrics)*len(devices))

	for _, device := range devices {
		for _, metric := range smart_metrics {
			path := makeName(device, metric)
			mts = append(mts, plugin.PluginMetricType{Namespace_: path})
		}
	}

	return mts, nil
}

//GetConfigPolicy returns a ConfigPolicy
func (p *SmartCollector) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	c := cpolicy.New()
	return c, nil
}
<commit_msg>Update version to 4<commit_after>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package smart

import (
	"errors"
	"fmt"
	"strings"

	"github.com/intelsdi-x/pulse/control/plugin"
	"github.com/intelsdi-x/pulse/control/plugin/cpolicy"
)

const (
	name       = "pulse-collector-smart"
	version    = 4
	pluginType = plugin.CollectorPluginType
)

var sysUtilProvider SysutilProvider = NewSysutilProvider()

var namespace_prefix = []string{"intel", "disk"}
var namespace_suffix = []string{"smart"}

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(name, version, pluginType, []string{plugin.PulseGOBContentType}, []string{plugin.PulseGOBContentType})
}

func NewSmartCollector() *SmartCollector {
	return &SmartCollector{}
}

func makeName(device, metric string) []string {
	splited := strings.Split(metric, "/")

	name := []string{}
	name = append(name, namespace_prefix...)
	name = append(name, device)
	name = append(name, namespace_suffix...)
	name = append(name, splited...)

	return name
}

func parseName(namespace []string) (disk, attribute string) {
	disk = namespace[len(namespace_prefix)]
	smart_namespace := namespace[len(namespace_prefix)+len(namespace_suffix)+1:]
	attribute = strings.Join(smart_namespace, "/")
	return
}

func validateName(namespace []string) bool {
	for i, v := range namespace_prefix {
		if namespace[i] != v {
			return false
		}
	}

	offset := len(namespace_prefix) + 1
	for i, v := range namespace_suffix {
		if namespace[offset+i] != v {
			return false
		}
	}

	return true
}

type SmartCollector struct {
}

type smartResults map[string]interface{}

// CollectMetrics returns metrics from smart
func (sc *SmartCollector) CollectMetrics(mts []plugin.PluginMetricType) ([]plugin.PluginMetricType, error) {
	buffered_results := map[string]smartResults{}

	results := make([]plugin.PluginMetricType, len(mts))

	for i, mt := range mts {
		if !validateName(mt.Namespace()) {
			return nil, errors.New(fmt.Sprintf("%v is not valid metric", mt.Namespace()))
		}
		disk, attribute_path := parseName(mt.Namespace())
		buffered, ok := buffered_results[disk]
		if !ok {
			values, err := ReadSmartData(disk, sysUtilProvider)
			if err != nil {
				return nil, err
			}
			buffered = values.GetAttributes()
			buffered_results[disk] = buffered
		}

		attribute, ok := buffered[attribute_path]

		if !ok {
			return nil, errors.New("Unknown attribute " + attribute_path)
		}

		results[i] = plugin.PluginMetricType{
			Namespace_: mt.Namespace(),
			Data_:      attribute,
		}
	}

	return results, nil
}

// GetMetricTypes returns the metric types exposed by smart
func (sc *SmartCollector) GetMetricTypes(_ plugin.PluginConfigType) ([]plugin.PluginMetricType, error) {
	smart_metrics := ListAllKeys()
	devices, err := sysUtilProvider.ListDevices()
	if err != nil {
		return nil, err
	}
	mts := make([]plugin.PluginMetricType, 0, len(smart_metrics)*len(devices))

	for _, device := range devices {
		for _, metric := range smart_metrics {
			path := makeName(device, metric)
			mts = append(mts, plugin.PluginMetricType{Namespace_: path})
		}
	}

	return mts, nil
}

//GetConfigPolicy returns a ConfigPolicy
func (p *SmartCollector) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	c := cpolicy.New()
	return c, nil
}
<|endoftext|>


<commit_before>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package smart

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

	log "github.com/Sirupsen/logrus"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core"

	"github.com/intelsdi-x/snap-plugin-utilities/config"
)

const (
	PluginName = "smart-disk"
	version    = 8
	pluginType = plugin.CollectorPluginType

	nsVendor = "intel"
	nsClass  = "disk"
	nsType   = "smart"
	devname  = "device"
)

var (
	//procPath source of data for metrics
	procPath = "/proc"
	//devPath source of data for metrics
	devPath = "/dev"

	namespace_prefix = []string{nsVendor, nsClass, nsType}

	sysUtilProvider SysutilProvider
)

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		PluginName,
		version,
		pluginType,
		[]string{plugin.SnapGOBContentType},
		[]string{plugin.SnapGOBContentType},
		plugin.ConcurrencyCount(1),
	)
}

func NewSmartCollector() *SmartCollector {
	logger := log.New()
	imutex := new(sync.Mutex)
	return &SmartCollector{
		logger:           logger,
		initializedMutex: imutex,
		proc_path:        procPath,
		dev_path:         devPath,
	}
}

func parseName(namespace []string) (disk, attribute string) {
	disk = namespace[len(namespace_prefix)]
	smart_namespace := namespace[len(namespace_prefix)+1:]
	attribute = strings.Join(smart_namespace, "/")
	return
}

func validateName(namespace []string) bool {
	for i, v := range namespace_prefix {
		if namespace[i] != v {
			return false
		}
	}
	return true
}

// Function to check properness of configuration parameters
// and set plugin attribute accordingly
func (sc *SmartCollector) setProcDevPath(cfg interface{}) error {
	sc.initializedMutex.Lock()
	defer sc.initializedMutex.Unlock()
	if sc.initialized {
		return nil
	}
	procPath, err := config.GetConfigItem(cfg, "proc_path")
	if err == nil && len(procPath.(string)) > 0 {
		procPathStats, err := os.Stat(procPath.(string))
		if err != nil {
			return err
		}
		if !procPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", procPath.(string)))
		}
		sc.proc_path = procPath.(string)
	}
	devPath, err := config.GetConfigItem(cfg, "dev_path")
	if err == nil && len(devPath.(string)) > 0 {
		devPathStats, err := os.Stat(devPath.(string))
		if err != nil {
			return err
		}
		if !devPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", devPath.(string)))
		}
		sc.dev_path = devPath.(string)
	}
	if sysUtilProvider == nil {
		sysUtilProvider = NewSysutilProvider(sc.proc_path, sc.dev_path)
	}
	sc.initialized = true
	return nil
}

type SmartCollector struct {
	initialized      bool
	initializedMutex *sync.Mutex
	logger           *log.Logger
	proc_path        string
	dev_path         string
}

type smartResults map[string]interface{}

// DiskMetrics returns metrics from smart on given disk
func (sc *SmartCollector) DiskMetrics(ns []core.NamespaceElement,
	t time.Time, disk string, attribute_path string,
	buffered_results map[string]smartResults, errs []string) (bool, plugin.MetricType, error) {
	var result plugin.MetricType
	collected := false
	buffered, ok := buffered_results[disk]
	if !ok {
		values, err := ReadSmartData(disk, sysUtilProvider)
		if err != nil {
			return collected, result, err
		}
		buffered = values.GetAttributes()
		buffered_results[disk] = buffered
	}
	attribute, ok := buffered[attribute_path]
	if !ok {
		errs = append(errs, "Unknown attribute "+attribute_path)
	} else {
		ns1 := make([]core.NamespaceElement, len(ns))
		copy(ns1, ns)
		ns1[3].Value = disk
		result = plugin.MetricType{
			Namespace_: ns1,
			Timestamp_: t,
			Version_:   version,
			Data_:      attribute,
		}
		collected = true
	}
	return collected, result, nil
}

// CollectMetrics returns metrics from smart
func (sc *SmartCollector) CollectMetrics(mts []plugin.MetricType) ([]plugin.MetricType, error) {
	if err := sc.setProcDevPath(mts[0]); err != nil {
		return nil, err
	}

	buffered_results := map[string]smartResults{}
	results := []plugin.MetricType{}
	errs := make([]string, 0)
	something_collected := false
	t := time.Now()
	for _, mt := range mts {
		ns := mt.Namespace()
		if !validateName(ns.Strings()) {
			errs = append(errs, fmt.Sprintf("%s is not valid metric", ns.String()))
			continue
		}
		disk, attribute_path := parseName(ns.Strings())
		if disk == "*" {
			// All system disks requested
			devices, err := sysUtilProvider.ListDevices()
			if err != nil {
				return nil, err
			}
			for _, dev := range devices {
				collected, result, err := sc.DiskMetrics(ns, t, dev, attribute_path, buffered_results, errs)
				if err != nil {
					sc.logger.Error(fmt.Sprintf("Error collecting SMART %s data on %s disk: %#+v", attribute_path, dev, err))
				} else {
					if collected {
						results = append(results, result)
						something_collected = true
					}
				}
			}
		} else {
			// Single disk requested

			collected, result, err := sc.DiskMetrics(ns, t, disk, attribute_path, buffered_results, errs)
			if err != nil {
				sc.logger.Error(fmt.Sprintf("Error collecting SMART %s data on %s disk: %#+v", attribute_path, disk, err))
			} else {
				if collected {
					results = append(results, result)
					something_collected = true
				}
			}
		}
	}
	errsStr := strings.Join(errs, "; ")
	if something_collected {
		if len(errs) > 0 {
			sc.logger.Error(fmt.Sprintf("Data collected but error(s) occured: %v", errsStr))
		}
		return results, nil
	} else {
		return nil, errors.New(errsStr)
	}
}

// GetMetricTypes returns the metric types exposed by smart
func (sc *SmartCollector) GetMetricTypes(cfg plugin.ConfigType) ([]plugin.MetricType, error) {
	smart_metrics := ListAllKeys()
	mts := []plugin.MetricType{}
	for _, metric := range smart_metrics {
		ns := core.NewNamespace(namespace_prefix...).AddDynamicElement(devname, "SMART device")
		for _, elt := range strings.Split(metric, "/") {
			ns = ns.AddStaticElement(elt)
		}
		mts = append(mts, plugin.MetricType{
			Namespace_:   ns,
			Description_: "dynamic SMART metric: " + metric,
		})
	}
	return mts, nil
}

//GetConfigPolicy returns a ConfigPolicy
func (p *SmartCollector) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("proc_path", false, "/proc")
	node := cpolicy.NewPolicyNode()
	node.Add(rule)
	cp.Add([]string{nsVendor, nsClass, nsType}, node)
	rule, _ = cpolicy.NewStringRule("dev_path", false, "/dev")
	node.Add(rule)
	return cp, nil
}
<commit_msg>Fix for issue #27: smart plugin returns error if nothing is collected<commit_after>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package smart

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

	log "github.com/Sirupsen/logrus"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core"

	"github.com/intelsdi-x/snap-plugin-utilities/config"
)

const (
	PluginName = "smart-disk"
	version    = 8
	pluginType = plugin.CollectorPluginType

	nsVendor = "intel"
	nsClass  = "disk"
	nsType   = "smart"
	devname  = "device"
)

var (
	//procPath source of data for metrics
	procPath = "/proc"
	//devPath source of data for metrics
	devPath = "/dev"

	namespace_prefix = []string{nsVendor, nsClass, nsType}

	sysUtilProvider SysutilProvider
)

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(
		PluginName,
		version,
		pluginType,
		[]string{plugin.SnapGOBContentType},
		[]string{plugin.SnapGOBContentType},
		plugin.ConcurrencyCount(1),
	)
}

func NewSmartCollector() *SmartCollector {
	logger := log.New()
	imutex := new(sync.Mutex)
	return &SmartCollector{
		logger:           logger,
		initializedMutex: imutex,
		proc_path:        procPath,
		dev_path:         devPath,
	}
}

func parseName(namespace []string) (disk, attribute string) {
	disk = namespace[len(namespace_prefix)]
	smart_namespace := namespace[len(namespace_prefix)+1:]
	attribute = strings.Join(smart_namespace, "/")
	return
}

func validateName(namespace []string) bool {
	for i, v := range namespace_prefix {
		if namespace[i] != v {
			return false
		}
	}
	return true
}

// Function to check properness of configuration parameters
// and set plugin attribute accordingly
func (sc *SmartCollector) setProcDevPath(cfg interface{}) error {
	sc.initializedMutex.Lock()
	defer sc.initializedMutex.Unlock()
	if sc.initialized {
		return nil
	}
	procPath, err := config.GetConfigItem(cfg, "proc_path")
	if err == nil && len(procPath.(string)) > 0 {
		procPathStats, err := os.Stat(procPath.(string))
		if err != nil {
			return err
		}
		if !procPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", procPath.(string)))
		}
		sc.proc_path = procPath.(string)
	}
	devPath, err := config.GetConfigItem(cfg, "dev_path")
	if err == nil && len(devPath.(string)) > 0 {
		devPathStats, err := os.Stat(devPath.(string))
		if err != nil {
			return err
		}
		if !devPathStats.IsDir() {
			return errors.New(fmt.Sprintf("%s is not a directory", devPath.(string)))
		}
		sc.dev_path = devPath.(string)
	}
	if sysUtilProvider == nil {
		sysUtilProvider = NewSysutilProvider(sc.proc_path, sc.dev_path)
	}
	sc.initialized = true
	return nil
}

type SmartCollector struct {
	initialized      bool
	initializedMutex *sync.Mutex
	logger           *log.Logger
	proc_path        string
	dev_path         string
}

type smartResults map[string]interface{}

// DiskMetrics returns metrics from smart on given disk
func (sc *SmartCollector) DiskMetrics(ns []core.NamespaceElement,
	t time.Time, disk string, attribute_path string,
	buffered_results map[string]smartResults, errs []string) (bool, plugin.MetricType, error) {
	var result plugin.MetricType
	collected := false
	buffered, ok := buffered_results[disk]
	if !ok {
		values, err := ReadSmartData(disk, sysUtilProvider)
		if err != nil {
			return collected, result, err
		}
		buffered = values.GetAttributes()
		buffered_results[disk] = buffered
	}
	attribute, ok := buffered[attribute_path]
	if !ok {
		errs = append(errs, "Unknown attribute "+attribute_path)
	} else {
		ns1 := make([]core.NamespaceElement, len(ns))
		copy(ns1, ns)
		ns1[3].Value = disk
		result = plugin.MetricType{
			Namespace_: ns1,
			Timestamp_: t,
			Version_:   version,
			Data_:      attribute,
		}
		collected = true
	}
	return collected, result, nil
}

// CollectMetrics returns metrics from smart
func (sc *SmartCollector) CollectMetrics(mts []plugin.MetricType) ([]plugin.MetricType, error) {
	if err := sc.setProcDevPath(mts[0]); err != nil {
		return nil, err
	}

	buffered_results := map[string]smartResults{}
	results := []plugin.MetricType{}
	errs := make([]string, 0)
	something_collected := false
	t := time.Now()
	for _, mt := range mts {
		ns := mt.Namespace()
		if !validateName(ns.Strings()) {
			errs = append(errs, fmt.Sprintf("%s is not valid metric", ns.String()))
			continue
		}
		disk, attribute_path := parseName(ns.Strings())
		if disk == "*" {
			// All system disks requested
			devices, err := sysUtilProvider.ListDevices()
			if err != nil {
				return nil, err
			}
			for _, dev := range devices {
				collected, result, err := sc.DiskMetrics(ns, t, dev, attribute_path, buffered_results, errs)
				if err != nil {
					eStr := fmt.Sprintf("Error collecting SMART %s data on %s disk: %#+v", attribute_path, dev, err)
					sc.logger.Error(eStr)
					errs = append(errs, eStr)
				} else {
					if collected {
						results = append(results, result)
						something_collected = true
					}
				}
			}
		} else {
			// Single disk requested
			collected, result, err := sc.DiskMetrics(ns, t, disk, attribute_path, buffered_results, errs)
			if err != nil {
				eStr := fmt.Sprintf("Error collecting SMART %s data on %s disk: %#+v", attribute_path, disk, err)
				sc.logger.Error(eStr)
				errs = append(errs, eStr)
			} else {
				if collected {
					results = append(results, result)
					something_collected = true
				}
			}
		}
	}
	errsStr := strings.Join(errs, "; ")
	if len(errs) > 0 {
		if !something_collected {
			return nil, errors.New(errsStr)
		}
	}
	return results, nil
}

// GetMetricTypes returns the metric types exposed by smart
func (sc *SmartCollector) GetMetricTypes(cfg plugin.ConfigType) ([]plugin.MetricType, error) {
	smart_metrics := ListAllKeys()
	mts := []plugin.MetricType{}
	for _, metric := range smart_metrics {
		ns := core.NewNamespace(namespace_prefix...).AddDynamicElement(devname, "SMART device")
		for _, elt := range strings.Split(metric, "/") {
			ns = ns.AddStaticElement(elt)
		}
		mts = append(mts, plugin.MetricType{
			Namespace_:   ns,
			Description_: "dynamic SMART metric: " + metric,
		})
	}
	return mts, nil
}

//GetConfigPolicy returns a ConfigPolicy
func (p *SmartCollector) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	rule, _ := cpolicy.NewStringRule("proc_path", false, "/proc")
	node := cpolicy.NewPolicyNode()
	node.Add(rule)
	cp.Add([]string{nsVendor, nsClass, nsType}, node)
	rule, _ = cpolicy.NewStringRule("dev_path", false, "/dev")
	node.Add(rule)
	return cp, nil
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package versionupgrade

import (
	"flag"
	"fmt"
	"os"
	"path"
	"testing"

	"vitess.io/vitess/go/mysql"

	"vitess.io/vitess/go/test/endtoend/cluster"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	clusterInstance       *cluster.LocalProcessCluster
	vtParams              mysql.ConnParams
	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	totalTableCount       = 4
	createTable           = `
		CREATE TABLE %s (
			id bigint(20) NOT NULL,
			msg varchar(64),
			PRIMARY KEY (id)
		) ENGINE=InnoDB;
		`
	insertIntoTable = `
		INSERT INTO %s (id, msg) VALUES (17, 'abc');
		`
	selectFromTable = `
		SELECT id, msg FROM %s LIMIT 1;
		`
)

// TestMain is the main entry point
func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, true); err != nil {
			return 1, err
		}
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 1, false); err != nil {
			return 1, err
		}

		vtgateInstance := clusterInstance.NewVtgateInstance()
		// set the gateway we want to use
		vtgateInstance.GatewayImplementation = "tabletgateway"
		// Start vtgate
		if err := vtgateInstance.Setup(); err != nil {
			return 1, err
		}
		// ensure it is torn down during cluster TearDown
		clusterInstance.VtgateProcess = *vtgateInstance
		vtParams = mysql.ConnParams{
			Host: clusterInstance.Hostname,
			Port: clusterInstance.VtgateMySQLPort,
		}

		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestShards(t *testing.T) {
	defer cluster.PanicHandler(t)
	assert.Equal(t, 2, len(clusterInstance.Keyspaces[0].Shards))
}

func TestDeploySchema(t *testing.T) {
	defer cluster.PanicHandler(t)

	if clusterInstance.ReusingVTDATAROOT {
		// we assume data is already deployed
		return
	}
	// Create n tables, populate
	for i := 0; i < totalTableCount; i++ {
		tableName := fmt.Sprintf("vt_upgrade_test_%02d", i)

		{
			sqlQuery := fmt.Sprintf(createTable, tableName)
			err := clusterInstance.VtctlclientProcess.ApplySchema(keyspaceName, sqlQuery)
			require.Nil(t, err)
		}
		for i := range clusterInstance.Keyspaces[0].Shards {
			sqlQuery := fmt.Sprintf(insertIntoTable, tableName)
			tablet := clusterInstance.Keyspaces[0].Shards[i].Vttablets[0]
			_, err := tablet.VttabletProcess.QueryTablet(sqlQuery, keyspaceName, true)
			require.Nil(t, err)
		}
	}

	checkTables(t, "", totalTableCount)
}

func TestTablesExist(t *testing.T) {
	defer cluster.PanicHandler(t)

	checkTables(t, "", totalTableCount)
}

// checkTables checks the number of tables in the first two shards.
func checkTables(t *testing.T, showTableName string, expectCount int) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], showTableName, expectCount)
	}
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, showTableName string, expectCount int) {
	query := fmt.Sprintf(`show tables like '%%%s%%';`, showTableName)
	queryResult, err := tablet.VttabletProcess.QueryTablet(query, keyspaceName, true)
	require.Nil(t, err)
	assert.Equal(t, expectCount, len(queryResult.Rows))
}

// TestTablesData checks the data in tables
func TestTablesData(t *testing.T) {
	// Create n tables, populate
	for i := 0; i < totalTableCount; i++ {
		tableName := fmt.Sprintf("vt_upgrade_test_%02d", i)

		for i := range clusterInstance.Keyspaces[0].Shards {
			sqlQuery := fmt.Sprintf(selectFromTable, tableName)
			tablet := clusterInstance.Keyspaces[0].Shards[i].Vttablets[0]
			queryResult, err := tablet.VttabletProcess.QueryTablet(sqlQuery, keyspaceName, true)
			require.Nil(t, err)
			require.NotNil(t, queryResult)
			row := queryResult.Named().Row()
			require.NotNil(t, row)
			require.Equal(t, int64(17), row.AsInt64("id", 0))
			require.Equal(t, "abc", row.AsString("msg", ""))
		}
	}
}
<commit_msg>some comments explaining the purpose of upgrade_test<commit_after>/*
Copyright 2019 The Vitess Authors.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

/*

ABOUT THIS TEST
===============

This test plays part in testing an upgrade path from a previous version/tag. It takes a GitHub workflow file to complete the functionality.
What's in this file is the setting up of a cluster, sharded and unsharded keyspace, creating and populating some tables, then testing retrieval of data.
The twist here is that you can run this test over pre-existing vtdataroot, which means this test can reuse existing etcd, existing tables, existing mysql,
in which case it will not attempt to create keyspaces/schemas/tables, nor will it populate table data. Instead, it will only check for retrieval of data.

The game is to setup the cluster with a stable version (say `v8.0.0`), take it down (and preserve data), then setup a new cluster with a new version (namely the branch/PR head) and attempt to read the data.

Both executions must force some settings so that both reuse same directories, ports, etc. An invocation will look like:
go test ./go/test/endtoend/versionupgrade80/upgrade80_test.go --keep-data -force-vtdataroot /tmp/vtdataroot/vtroot_10901 --force-port-start 11900 --force-base-tablet-uid 1190

*/

package versionupgrade

import (
	"flag"
	"fmt"
	"os"
	"path"
	"testing"

	"vitess.io/vitess/go/mysql"

	"vitess.io/vitess/go/test/endtoend/cluster"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

var (
	clusterInstance       *cluster.LocalProcessCluster
	vtParams              mysql.ConnParams
	hostname              = "localhost"
	keyspaceName          = "ks"
	cell                  = "zone1"
	schemaChangeDirectory = ""
	totalTableCount       = 4
	createTable           = `
		CREATE TABLE %s (
			id bigint(20) NOT NULL,
			msg varchar(64),
			PRIMARY KEY (id)
		) ENGINE=InnoDB;
		`
	insertIntoTable = `
		INSERT INTO %s (id, msg) VALUES (17, 'abc');
		`
	selectFromTable = `
		SELECT id, msg FROM %s LIMIT 1;
		`
)

// TestMain is the main entry point
func TestMain(m *testing.M) {
	defer cluster.PanicHandler(nil)
	flag.Parse()

	exitcode, err := func() (int, error) {
		clusterInstance = cluster.NewCluster(cell, hostname)
		schemaChangeDirectory = path.Join("/tmp", fmt.Sprintf("schema_change_dir_%d", clusterInstance.GetAndReserveTabletUID()))
		defer os.RemoveAll(schemaChangeDirectory)
		defer clusterInstance.Teardown()

		if _, err := os.Stat(schemaChangeDirectory); os.IsNotExist(err) {
			_ = os.Mkdir(schemaChangeDirectory, 0700)
		}

		clusterInstance.VtctldExtraArgs = []string{
			"-schema_change_dir", schemaChangeDirectory,
			"-schema_change_controller", "local",
			"-schema_change_check_interval", "1"}

		if err := clusterInstance.StartTopo(); err != nil {
			return 1, err
		}

		// Start keyspace
		keyspace := &cluster.Keyspace{
			Name: keyspaceName,
		}

		if err := clusterInstance.StartUnshardedKeyspace(*keyspace, 2, true); err != nil {
			return 1, err
		}
		if err := clusterInstance.StartKeyspace(*keyspace, []string{"1"}, 1, false); err != nil {
			return 1, err
		}

		vtgateInstance := clusterInstance.NewVtgateInstance()
		// set the gateway we want to use
		vtgateInstance.GatewayImplementation = "tabletgateway"
		// Start vtgate
		if err := vtgateInstance.Setup(); err != nil {
			return 1, err
		}
		// ensure it is torn down during cluster TearDown
		clusterInstance.VtgateProcess = *vtgateInstance
		vtParams = mysql.ConnParams{
			Host: clusterInstance.Hostname,
			Port: clusterInstance.VtgateMySQLPort,
		}

		return m.Run(), nil
	}()
	if err != nil {
		fmt.Printf("%v\n", err)
		os.Exit(1)
	} else {
		os.Exit(exitcode)
	}

}

func TestShards(t *testing.T) {
	defer cluster.PanicHandler(t)
	assert.Equal(t, 2, len(clusterInstance.Keyspaces[0].Shards))
}

func TestDeploySchema(t *testing.T) {
	defer cluster.PanicHandler(t)

	if clusterInstance.ReusingVTDATAROOT {
		// we assume data is already deployed
		return
	}
	// Create n tables, populate
	for i := 0; i < totalTableCount; i++ {
		tableName := fmt.Sprintf("vt_upgrade_test_%02d", i)

		{
			sqlQuery := fmt.Sprintf(createTable, tableName)
			err := clusterInstance.VtctlclientProcess.ApplySchema(keyspaceName, sqlQuery)
			require.Nil(t, err)
		}
		for i := range clusterInstance.Keyspaces[0].Shards {
			sqlQuery := fmt.Sprintf(insertIntoTable, tableName)
			tablet := clusterInstance.Keyspaces[0].Shards[i].Vttablets[0]
			_, err := tablet.VttabletProcess.QueryTablet(sqlQuery, keyspaceName, true)
			require.Nil(t, err)
		}
	}

	checkTables(t, "", totalTableCount)
}

func TestTablesExist(t *testing.T) {
	defer cluster.PanicHandler(t)

	checkTables(t, "", totalTableCount)
}

// checkTables checks the number of tables in the first two shards.
func checkTables(t *testing.T, showTableName string, expectCount int) {
	for i := range clusterInstance.Keyspaces[0].Shards {
		checkTablesCount(t, clusterInstance.Keyspaces[0].Shards[i].Vttablets[0], showTableName, expectCount)
	}
}

// checkTablesCount checks the number of tables in the given tablet
func checkTablesCount(t *testing.T, tablet *cluster.Vttablet, showTableName string, expectCount int) {
	query := fmt.Sprintf(`show tables like '%%%s%%';`, showTableName)
	queryResult, err := tablet.VttabletProcess.QueryTablet(query, keyspaceName, true)
	require.Nil(t, err)
	assert.Equal(t, expectCount, len(queryResult.Rows))
}

// TestTablesData checks the data in tables
func TestTablesData(t *testing.T) {
	// Create n tables, populate
	for i := 0; i < totalTableCount; i++ {
		tableName := fmt.Sprintf("vt_upgrade_test_%02d", i)

		for i := range clusterInstance.Keyspaces[0].Shards {
			sqlQuery := fmt.Sprintf(selectFromTable, tableName)
			tablet := clusterInstance.Keyspaces[0].Shards[i].Vttablets[0]
			queryResult, err := tablet.VttabletProcess.QueryTablet(sqlQuery, keyspaceName, true)
			require.Nil(t, err)
			require.NotNil(t, queryResult)
			row := queryResult.Named().Row()
			require.NotNil(t, row)
			require.Equal(t, int64(17), row.AsInt64("id", 0))
			require.Equal(t, "abc", row.AsString("msg", ""))
		}
	}
}
<|endoftext|>