<commit_before>package main

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	nagios "source.datanerd.us/site-engineering/go_nagios"
)

const (
	API_VERSION = "v1.10"
)

type CliOpt struct {
	BaseUrl string
	CritDataSpace int
	WarnDataSpace int
	CritMetaSpace int
	WarnMetaSpace int
}

type DockerInfo struct {
	Containers     float64
	DriverStatus   [][]string
	DataSpaceUsed  float64
	DataSpaceTotal float64
	MetaSpaceUsed  float64
	MetaSpaceTotal float64
}

type HttpResponseFetcher interface {
	Fetch(url string) ([]byte, error)
}

type Fetcher struct{}

func float64String(num float64) string {
	return strconv.FormatFloat(num, 'f', 0, 64)
}

func megabytesFloat64(value string) (float64, error) {
	numberStr   := strings.Fields(value)[0]
	number, err := strconv.ParseFloat(numberStr, 64)

	if err != nil {
		return 0.00, err
	}

	return number, nil
}

// Look through a list of driveStatus slices and find the one that matches
func findDriverStatus(key string, driverStatus [][]string) string {
	for _, entry := range(driverStatus) {
		if entry[0] == key {
			return entry[1]
		}
	}

	return ""
}

// Connect to a Docker URL and return the contents as a []byte
func (Fetcher) Fetch(url string) ([]byte, error) {
	response, err := http.Get(url + "/" + API_VERSION + "/info")

	if err != nil {
		return nil, err
	}

	defer response.Body.Close()

	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return contents, nil
}

func populateInfo(contents []byte, info *DockerInfo) error {
	err := json.Unmarshal(contents, &info)
	if err != nil {
		return err
	}

	fields := map[string]*float64{
		"Data Space Used":      &info.DataSpaceUsed,
		"Data Space Total":     &info.DataSpaceTotal,
		"Metadata Space Used":  &info.MetaSpaceUsed,
		"Metadata Space Total": &info.MetaSpaceTotal,
	}

	for key, val := range(fields) {
	  *val, err = megabytesFloat64(findDriverStatus(key, info.DriverStatus))
	  if err != nil {
		return errors.New("Error parsing response from API! " + err.Error())
	  }
	}

	return nil
}

func fetchInfo(fetcher Fetcher, opts CliOpt, info *DockerInfo) error {
	contents, err := fetcher.Fetch(opts.BaseUrl)
	if err != nil {
		return err
	}

	err = populateInfo(contents, info)
	if err != nil {
		return err
	}

	return nil
}

func mapAlertStatuses(info *DockerInfo, opts *CliOpt) []*nagios.NagiosStatus {
	var statuses []*nagios.NagiosStatus

	type checkArgs struct {
		tag string
		value float64
		comparison float64
		statusVal nagios.NagiosStatusVal
	}

	var check = func(args checkArgs) *nagios.NagiosStatus {
		if args.value > args.comparison {
			return &nagios.NagiosStatus{args.tag + ": " + float64String(args.value), args.statusVal}
		}
		return nil
	}

	checks := []checkArgs{
		checkArgs{"Meta Space Used",
			info.MetaSpaceUsed / info.MetaSpaceTotal * 100,
			float64(opts.CritMetaSpace),
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Data Space Used",
			info.DataSpaceUsed / info.DataSpaceTotal * 100,
			float64(opts.CritDataSpace),
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Meta Space Used",
			info.MetaSpaceUsed / info.MetaSpaceTotal * 100,
			float64(opts.WarnMetaSpace),
			nagios.NAGIOS_WARNING,
		},
		checkArgs{"Data Space Used",
			info.DataSpaceUsed / info.DataSpaceTotal * 100,
			float64(opts.WarnDataSpace),
			nagios.NAGIOS_WARNING,
		},
	}

	for _, entry := range(checks) {
		result := check(entry)
		if result != nil {
			statuses = append(statuses, check(entry))
		}
	}

	return statuses
}

func main() {
	opts := CliOpt{
		BaseUrl:      "http://chi-staging-pool-1:4243/",
		WarnMetaSpace: 0,
		WarnDataSpace: 0,
		CritMetaSpace: 66,
		CritDataSpace: 66,
	}

	var fetcher Fetcher
	var info DockerInfo

	err := fetchInfo(fetcher, opts, &info)
	if err != nil {
		nagios.Critical(err)
	}

	statuses := mapAlertStatuses(&info, &opts)
	status   := nagios.NagiosStatus{float64String(info.Containers) + " containers", 0}

	status.Aggregate(statuses)
	nagios.ExitWithStatus(&status)
}
<commit_msg>Added better comments for godoc.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"io/ioutil"
	"net/http"
	"strconv"
	"strings"

	nagios "source.datanerd.us/site-engineering/go_nagios"
)

const (
	API_VERSION = "v1.10"
)

// A struct representing CLI opts that will be passed at runtime
type CliOpt struct {
	BaseUrl string
	CritDataSpace int
	WarnDataSpace int
	CritMetaSpace int
	WarnMetaSpace int
}

// Information describing the status of a Docker host
type DockerInfo struct {
	Containers     float64
	DriverStatus   [][]string
	DataSpaceUsed  float64
	DataSpaceTotal float64
	MetaSpaceUsed  float64
	MetaSpaceTotal float64
}

type HttpResponseFetcher interface {
	Fetch(url string) ([]byte, error)
}

type Fetcher struct{}

// Properly format a Float64 as a string
func float64String(num float64) string {
	return strconv.FormatFloat(num, 'f', 0, 64)
}

// Return a float from a Docker info string for megabytes
func megabytesFloat64(value string) (float64, error) {
	numberStr   := strings.Fields(value)[0]
	number, err := strconv.ParseFloat(numberStr, 64)

	if err != nil {
		return 0.00, err
	}

	return number, nil
}

// Look through a list of driveStatus slices and find the one that matches
func findDriverStatus(key string, driverStatus [][]string) string {
	for _, entry := range(driverStatus) {
		if entry[0] == key {
			return entry[1]
		}
	}

	return ""
}

// Connect to a Docker URL and return the contents as a []byte
func (Fetcher) Fetch(url string) ([]byte, error) {
	response, err := http.Get(url + "/" + API_VERSION + "/info")

	if err != nil {
		return nil, err
	}

	defer response.Body.Close()

	contents, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	return contents, nil
}

// Parses JSON and populates a DockerInfo
func populateInfo(contents []byte, info *DockerInfo) error {
	err := json.Unmarshal(contents, &info)
	if err != nil {
		return err
	}

	fields := map[string]*float64{
		"Data Space Used":      &info.DataSpaceUsed,
		"Data Space Total":     &info.DataSpaceTotal,
		"Metadata Space Used":  &info.MetaSpaceUsed,
		"Metadata Space Total": &info.MetaSpaceTotal,
	}

	for key, val := range(fields) {
	  *val, err = megabytesFloat64(findDriverStatus(key, info.DriverStatus))
	  if err != nil {
		return errors.New("Error parsing response from API! " + err.Error())
	  }
	}

	return nil
}

// Retrieves JSON from a Docker host and fills in a DockerInfo
func fetchInfo(fetcher Fetcher, opts CliOpt, info *DockerInfo) error {
	contents, err := fetcher.Fetch(opts.BaseUrl)
	if err != nil {
		return err
	}

	err = populateInfo(contents, info)
	if err != nil {
		return err
	}

	return nil
}

// Runs a set of checkes and returns an array of statuses
func mapAlertStatuses(info *DockerInfo, opts *CliOpt) []*nagios.NagiosStatus {
	var statuses []*nagios.NagiosStatus

	type checkArgs struct {
		tag string
		value float64
		comparison float64
		statusVal nagios.NagiosStatusVal
	}

	var check = func(args checkArgs) *nagios.NagiosStatus {
		if args.value > args.comparison {
			return &nagios.NagiosStatus{args.tag + ": " + float64String(args.value), args.statusVal}
		}
		return nil
	}

	checks := []checkArgs{
		checkArgs{"Meta Space Used",
			info.MetaSpaceUsed / info.MetaSpaceTotal * 100,
			float64(opts.CritMetaSpace),
			nagios.NAGIOS_CRITICAL,
		},
		checkArgs{"Data Space Used",
			info.DataSpaceUsed / info.DataSpaceTotal * 100,
			float64(opts.CritDataSpace),
			nagios.NAGIOS_CRITICAL,
		},
	}

	for _, entry := range(checks) {
		result := check(entry)
		if result != nil {
			statuses = append(statuses, check(entry))
		}
	}

	return statuses
}

func main() {
	opts := CliOpt{
		BaseUrl:      "http://chi-staging-pool-1:4243/",
		WarnMetaSpace: 0,
		WarnDataSpace: 0,
		CritMetaSpace: 66,
		CritDataSpace: 66,
	}

	var fetcher Fetcher
	var info DockerInfo

	err := fetchInfo(fetcher, opts, &info)
	if err != nil {
		nagios.Critical(err)
	}

	statuses := mapAlertStatuses(&info, &opts)
	status   := nagios.NagiosStatus{float64String(info.Containers) + " containers", 0}

	status.Aggregate(statuses)
	nagios.ExitWithStatus(&status)
}
<|endoftext|>


<commit_before>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package stats

import (
	"errors"
	"strconv"
	"time"

	"strings"

	"github.com/dustin/go-humanize"
	"gopkg.in/guregu/null.v3"
)

const (
	counterString = `"counter"`
	gaugeString   = `"gauge"`
	trendString   = `"trend"`
	rateString    = `"rate"`

	defaultString = `"default"`
	timeString    = `"time"`
	dataString    = `"data"`
)

// Possible values for MetricType.
const (
	Counter = MetricType(iota) // A counter that sums its data points
	Gauge                      // A gauge that displays the latest value
	Trend                      // A trend, min/max/avg/med are interesting
	Rate                       // A rate, displays % of values that aren't 0
)

// Possible values for ValueType.
const (
	Default = ValueType(iota) // Values are presented as-is
	Time                      // Values are timestamps (nanoseconds)
	Data                      // Values are data amounts (bytes)
)

// The serialized metric type is invalid.
var ErrInvalidMetricType = errors.New("Invalid metric type")

// The serialized value type is invalid.
var ErrInvalidValueType = errors.New("Invalid value type")

// A MetricType specifies the type of a metric.
type MetricType int

// MarshalJSON serializes a MetricType as a human readable string.
func (t MetricType) MarshalJSON() ([]byte, error) {
	switch t {
	case Counter:
		return []byte(counterString), nil
	case Gauge:
		return []byte(gaugeString), nil
	case Trend:
		return []byte(trendString), nil
	case Rate:
		return []byte(rateString), nil
	default:
		return nil, ErrInvalidMetricType
	}
}

// UnmarshalJSON deserializes a MetricType from a string representation.
func (t *MetricType) UnmarshalJSON(data []byte) error {
	switch string(data) {
	case counterString:
		*t = Counter
	case gaugeString:
		*t = Gauge
	case trendString:
		*t = Trend
	case rateString:
		*t = Rate
	default:
		return ErrInvalidMetricType
	}

	return nil
}

func (t MetricType) String() string {
	switch t {
	case Counter:
		return counterString
	case Gauge:
		return gaugeString
	case Trend:
		return trendString
	case Rate:
		return rateString
	default:
		return "[INVALID]"
	}
}

// The type of values a metric contains.
type ValueType int

// MarshalJSON serializes a ValueType as a human readable string.
func (t ValueType) MarshalJSON() ([]byte, error) {
	switch t {
	case Default:
		return []byte(defaultString), nil
	case Time:
		return []byte(timeString), nil
	case Data:
		return []byte(dataString), nil
	default:
		return nil, ErrInvalidValueType
	}
}

// UnmarshalJSON deserializes a ValueType from a string representation.
func (t *ValueType) UnmarshalJSON(data []byte) error {
	switch string(data) {
	case defaultString:
		*t = Default
	case timeString:
		*t = Time
	case dataString:
		*t = Data
	default:
		return ErrInvalidValueType
	}

	return nil
}

func (t ValueType) String() string {
	switch t {
	case Default:
		return defaultString
	case Time:
		return timeString
	case Data:
		return dataString
	default:
		return "[INVALID]"
	}
}

// A Sample is a single measurement.
type Sample struct {
	Metric *Metric
	Time   time.Time
	Tags   map[string]string
	Value  float64
}

// A Metric defines the shape of a set of data.
type Metric struct {
	Name       string      `json:"-"`
	Type       MetricType  `json:"type"`
	Contains   ValueType   `json:"contains"`
	Tainted    null.Bool   `json:"tainted"`
	Sink       Sink        `json:"sink"`
	Thresholds Thresholds  `json:"thresholds"`
	Submetrics []Submetric `json:"submetrics"`
}

func New(name string, typ MetricType, t ...ValueType) *Metric {
	vt := Default
	if len(t) > 0 {
		vt = t[0]
	}
	var sink Sink
	switch typ {
	case Counter:
		sink = &CounterSink{}
	case Gauge:
		sink = &GaugeSink{}
	case Trend:
		sink = &TrendSink{}
	case Rate:
		sink = &RateSink{}
	default:
		return nil
	}
	return &Metric{Name: name, Type: typ, Contains: vt, Sink: sink}
}

func (m Metric) HumanizeValue(v float64) string {
	switch m.Type {
	case Rate:
		return strconv.FormatFloat(100*v, 'f', 2, 64) + "%"
	default:
		switch m.Contains {
		case Time:
			d := ToD(v)
			switch {
			case d > time.Minute:
				d -= d % (1 * time.Second)
			case d > time.Second:
				d -= d % (10 * time.Millisecond)
			case d > time.Millisecond:
				d -= d % (10 * time.Microsecond)
			case d > time.Microsecond:
				d -= d % (10 * time.Nanosecond)
			}
			return d.String()
		case Data:
			return humanize.Bytes(uint64(v))
		default:
			return strconv.FormatFloat(v, 'f', -1, 64)
		}
	}
}

// A Submetric represents a filtered dataset based on a parent metric.
type Submetric struct {
	Name   string            `json:"name"`
	Tags   map[string]string `json:"tags"`
	Metric *Metric           `json:"metric"`
}

// Creates a submetric from a name.
func NewSubmetric(name string) (parentName string, sm Submetric) {
	parts := strings.SplitN(strings.TrimSuffix(name, "}"), "{", 2)
	if len(parts) == 1 {
		return parts[0], Submetric{Name: name}
	}

	kvs := strings.Split(parts[1], ",")
	tags := make(map[string]string, len(kvs))
	for _, kv := range kvs {
		if kv == "" {
			continue
		}
		parts := strings.SplitN(kv, ":", 2)

		key := strings.TrimSpace(strings.Trim(parts[0], `"'`))
		if len(parts) != 2 {
			tags[key] = ""
			continue
		}

		value := strings.TrimSpace(strings.Trim(parts[1], `"'`))
		tags[key] = value
	}
	return parts[0], Submetric{Name: name, Tags: tags}
}
<commit_msg>JSON: Include metric names, but not sinks<commit_after>/*
 *
 * k6 - a next-generation load testing tool
 * Copyright (C) 2016 Load Impact
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 *
 */

package stats

import (
	"errors"
	"strconv"
	"time"

	"strings"

	"github.com/dustin/go-humanize"
	"gopkg.in/guregu/null.v3"
)

const (
	counterString = `"counter"`
	gaugeString   = `"gauge"`
	trendString   = `"trend"`
	rateString    = `"rate"`

	defaultString = `"default"`
	timeString    = `"time"`
	dataString    = `"data"`
)

// Possible values for MetricType.
const (
	Counter = MetricType(iota) // A counter that sums its data points
	Gauge                      // A gauge that displays the latest value
	Trend                      // A trend, min/max/avg/med are interesting
	Rate                       // A rate, displays % of values that aren't 0
)

// Possible values for ValueType.
const (
	Default = ValueType(iota) // Values are presented as-is
	Time                      // Values are timestamps (nanoseconds)
	Data                      // Values are data amounts (bytes)
)

// The serialized metric type is invalid.
var ErrInvalidMetricType = errors.New("Invalid metric type")

// The serialized value type is invalid.
var ErrInvalidValueType = errors.New("Invalid value type")

// A MetricType specifies the type of a metric.
type MetricType int

// MarshalJSON serializes a MetricType as a human readable string.
func (t MetricType) MarshalJSON() ([]byte, error) {
	switch t {
	case Counter:
		return []byte(counterString), nil
	case Gauge:
		return []byte(gaugeString), nil
	case Trend:
		return []byte(trendString), nil
	case Rate:
		return []byte(rateString), nil
	default:
		return nil, ErrInvalidMetricType
	}
}

// UnmarshalJSON deserializes a MetricType from a string representation.
func (t *MetricType) UnmarshalJSON(data []byte) error {
	switch string(data) {
	case counterString:
		*t = Counter
	case gaugeString:
		*t = Gauge
	case trendString:
		*t = Trend
	case rateString:
		*t = Rate
	default:
		return ErrInvalidMetricType
	}

	return nil
}

func (t MetricType) String() string {
	switch t {
	case Counter:
		return counterString
	case Gauge:
		return gaugeString
	case Trend:
		return trendString
	case Rate:
		return rateString
	default:
		return "[INVALID]"
	}
}

// The type of values a metric contains.
type ValueType int

// MarshalJSON serializes a ValueType as a human readable string.
func (t ValueType) MarshalJSON() ([]byte, error) {
	switch t {
	case Default:
		return []byte(defaultString), nil
	case Time:
		return []byte(timeString), nil
	case Data:
		return []byte(dataString), nil
	default:
		return nil, ErrInvalidValueType
	}
}

// UnmarshalJSON deserializes a ValueType from a string representation.
func (t *ValueType) UnmarshalJSON(data []byte) error {
	switch string(data) {
	case defaultString:
		*t = Default
	case timeString:
		*t = Time
	case dataString:
		*t = Data
	default:
		return ErrInvalidValueType
	}

	return nil
}

func (t ValueType) String() string {
	switch t {
	case Default:
		return defaultString
	case Time:
		return timeString
	case Data:
		return dataString
	default:
		return "[INVALID]"
	}
}

// A Sample is a single measurement.
type Sample struct {
	Metric *Metric
	Time   time.Time
	Tags   map[string]string
	Value  float64
}

// A Metric defines the shape of a set of data.
type Metric struct {
	Name       string      `json:"name"`
	Type       MetricType  `json:"type"`
	Contains   ValueType   `json:"contains"`
	Tainted    null.Bool   `json:"tainted"`
	Thresholds Thresholds  `json:"thresholds"`
	Submetrics []Submetric `json:"submetrics"`
	Sink       Sink        `json:"-"`
}

func New(name string, typ MetricType, t ...ValueType) *Metric {
	vt := Default
	if len(t) > 0 {
		vt = t[0]
	}
	var sink Sink
	switch typ {
	case Counter:
		sink = &CounterSink{}
	case Gauge:
		sink = &GaugeSink{}
	case Trend:
		sink = &TrendSink{}
	case Rate:
		sink = &RateSink{}
	default:
		return nil
	}
	return &Metric{Name: name, Type: typ, Contains: vt, Sink: sink}
}

func (m Metric) HumanizeValue(v float64) string {
	switch m.Type {
	case Rate:
		return strconv.FormatFloat(100*v, 'f', 2, 64) + "%"
	default:
		switch m.Contains {
		case Time:
			d := ToD(v)
			switch {
			case d > time.Minute:
				d -= d % (1 * time.Second)
			case d > time.Second:
				d -= d % (10 * time.Millisecond)
			case d > time.Millisecond:
				d -= d % (10 * time.Microsecond)
			case d > time.Microsecond:
				d -= d % (10 * time.Nanosecond)
			}
			return d.String()
		case Data:
			return humanize.Bytes(uint64(v))
		default:
			return strconv.FormatFloat(v, 'f', -1, 64)
		}
	}
}

// A Submetric represents a filtered dataset based on a parent metric.
type Submetric struct {
	Name   string            `json:"name"`
	Tags   map[string]string `json:"tags"`
	Metric *Metric           `json:"metric"`
}

// Creates a submetric from a name.
func NewSubmetric(name string) (parentName string, sm Submetric) {
	parts := strings.SplitN(strings.TrimSuffix(name, "}"), "{", 2)
	if len(parts) == 1 {
		return parts[0], Submetric{Name: name}
	}

	kvs := strings.Split(parts[1], ",")
	tags := make(map[string]string, len(kvs))
	for _, kv := range kvs {
		if kv == "" {
			continue
		}
		parts := strings.SplitN(kv, ":", 2)

		key := strings.TrimSpace(strings.Trim(parts[0], `"'`))
		if len(parts) != 2 {
			tags[key] = ""
			continue
		}

		value := strings.TrimSpace(strings.Trim(parts[1], `"'`))
		tags[key] = value
	}
	return parts[0], Submetric{Name: name, Tags: tags}
}
<|endoftext|>


<commit_before>// Copyright 2017 fatedier, fatedier@gmail.com
//
// 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 client

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/fatedier/frp/models/config"
	"github.com/fatedier/frp/models/msg"
	"github.com/fatedier/frp/models/plugin"
	"github.com/fatedier/frp/models/proto/udp"
	"github.com/fatedier/frp/utils/errors"
	frpIo "github.com/fatedier/frp/utils/io"
	"github.com/fatedier/frp/utils/log"
	frpNet "github.com/fatedier/frp/utils/net"
	"github.com/fatedier/frp/utils/pool"
)

// Proxy defines how to deal with work connections for different proxy type.
type Proxy interface {
	Run() error

	// InWorkConn accept work connections registered to server.
	InWorkConn(conn frpNet.Conn)
	Close()
	log.Logger
}

func NewProxy(ctl *Control, pxyConf config.ProxyConf) (pxy Proxy) {
	baseProxy := BaseProxy{
		ctl:    ctl,
		Logger: log.NewPrefixLogger(pxyConf.GetName()),
	}
	switch cfg := pxyConf.(type) {
	case *config.TcpProxyConf:
		pxy = &TcpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.UdpProxyConf:
		pxy = &UdpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.HttpProxyConf:
		pxy = &HttpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.HttpsProxyConf:
		pxy = &HttpsProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.StcpProxyConf:
		pxy = &StcpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.XtcpProxyConf:
		pxy = &XtcpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	}
	return
}

type BaseProxy struct {
	ctl    *Control
	closed bool
	mu     sync.RWMutex
	log.Logger
}

// TCP
type TcpProxy struct {
	BaseProxy

	cfg         *config.TcpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *TcpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *TcpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *TcpProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// HTTP
type HttpProxy struct {
	BaseProxy

	cfg         *config.HttpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *HttpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *HttpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *HttpProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// HTTPS
type HttpsProxy struct {
	BaseProxy

	cfg         *config.HttpsProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *HttpsProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *HttpsProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *HttpsProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// STCP
type StcpProxy struct {
	BaseProxy

	cfg         *config.StcpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *StcpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *StcpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *StcpProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// XTCP
type XtcpProxy struct {
	BaseProxy

	cfg         *config.XtcpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *XtcpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *XtcpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *XtcpProxy) InWorkConn(conn frpNet.Conn) {
	defer conn.Close()
	var natHoleSidMsg msg.NatHoleSid
	err := msg.ReadMsgInto(conn, &natHoleSidMsg)
	if err != nil {
		pxy.Error("xtcp read from workConn error: %v", err)
		return
	}

	natHoleClientMsg := &msg.NatHoleClient{
		ProxyName: pxy.cfg.ProxyName,
		Sid:       natHoleSidMsg.Sid,
	}
	raddr, _ := net.ResolveUDPAddr("udp",
		fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerUdpPort))
	clientConn, err := net.DialUDP("udp", nil, raddr)
	defer clientConn.Close()

	err = msg.WriteMsg(clientConn, natHoleClientMsg)
	if err != nil {
		pxy.Error("send natHoleClientMsg to server error: %v", err)
		return
	}

	// Wait for client address at most 5 seconds.
	var natHoleRespMsg msg.NatHoleResp
	clientConn.SetReadDeadline(time.Now().Add(5 * time.Second))

	buf := pool.GetBuf(1024)
	n, err := clientConn.Read(buf)
	if err != nil {
		pxy.Error("get natHoleRespMsg error: %v", err)
		return
	}
	err = msg.ReadMsgInto(bytes.NewReader(buf[:n]), &natHoleRespMsg)
	if err != nil {
		pxy.Error("get natHoleRespMsg error: %v", err)
		return
	}
	clientConn.SetReadDeadline(time.Time{})
	clientConn.Close()
	pxy.Trace("get natHoleRespMsg, sid [%s], client address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr)

	// Send sid to visitor udp address.
	time.Sleep(time.Second)
	laddr, _ := net.ResolveUDPAddr("udp", clientConn.LocalAddr().String())
	daddr, err := net.ResolveUDPAddr("udp", natHoleRespMsg.VisitorAddr)
	if err != nil {
		pxy.Error("resolve visitor udp address error: %v", err)
		return
	}

	lConn, err := net.DialUDP("udp", laddr, daddr)
	if err != nil {
		pxy.Error("dial visitor udp address error: %v", err)
		return
	}
	lConn.Write([]byte(natHoleRespMsg.Sid))

	kcpConn, err := frpNet.NewKcpConnFromUdp(lConn, true, natHoleRespMsg.VisitorAddr)
	if err != nil {
		pxy.Error("create kcp connection from udp connection error: %v", err)
		return
	}

	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf,
		frpNet.WrapConn(kcpConn), []byte(pxy.cfg.Sk))
}

// UDP
type UdpProxy struct {
	BaseProxy

	cfg *config.UdpProxyConf

	localAddr *net.UDPAddr
	readCh    chan *msg.UdpPacket

	// include msg.UdpPacket and msg.Ping
	sendCh   chan msg.Message
	workConn frpNet.Conn
}

func (pxy *UdpProxy) Run() (err error) {
	pxy.localAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.cfg.LocalIp, pxy.cfg.LocalPort))
	if err != nil {
		return
	}
	return
}

func (pxy *UdpProxy) Close() {
	pxy.mu.Lock()
	defer pxy.mu.Unlock()

	if !pxy.closed {
		pxy.closed = true
		if pxy.workConn != nil {
			pxy.workConn.Close()
		}
		if pxy.readCh != nil {
			close(pxy.readCh)
		}
		if pxy.sendCh != nil {
			close(pxy.sendCh)
		}
	}
}

func (pxy *UdpProxy) InWorkConn(conn frpNet.Conn) {
	pxy.Info("incoming a new work connection for udp proxy, %s", conn.RemoteAddr().String())
	// close resources releated with old workConn
	pxy.Close()

	pxy.mu.Lock()
	pxy.workConn = conn
	pxy.readCh = make(chan *msg.UdpPacket, 1024)
	pxy.sendCh = make(chan msg.Message, 1024)
	pxy.closed = false
	pxy.mu.Unlock()

	workConnReaderFn := func(conn net.Conn, readCh chan *msg.UdpPacket) {
		for {
			var udpMsg msg.UdpPacket
			if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
				pxy.Warn("read from workConn for udp error: %v", errRet)
				return
			}
			if errRet := errors.PanicToError(func() {
				pxy.Trace("get udp package from workConn: %s", udpMsg.Content)
				readCh <- &udpMsg
			}); errRet != nil {
				pxy.Info("reader goroutine for udp work connection closed: %v", errRet)
				return
			}
		}
	}
	workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
		defer func() {
			pxy.Info("writer goroutine for udp work connection closed")
		}()
		var errRet error
		for rawMsg := range sendCh {
			switch m := rawMsg.(type) {
			case *msg.UdpPacket:
				pxy.Trace("send udp package to workConn: %s", m.Content)
			case *msg.Ping:
				pxy.Trace("send ping message to udp workConn")
			}
			if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
				pxy.Error("udp work write error: %v", errRet)
				return
			}
		}
	}
	heartbeatFn := func(conn net.Conn, sendCh chan msg.Message) {
		var errRet error
		for {
			time.Sleep(time.Duration(30) * time.Second)
			if errRet = errors.PanicToError(func() {
				sendCh <- &msg.Ping{}
			}); errRet != nil {
				pxy.Trace("heartbeat goroutine for udp work connection closed")
				break
			}
		}
	}

	go workConnSenderFn(pxy.workConn, pxy.sendCh)
	go workConnReaderFn(pxy.workConn, pxy.readCh)
	go heartbeatFn(pxy.workConn, pxy.sendCh)
	udp.Forwarder(pxy.localAddr, pxy.readCh, pxy.sendCh)
}

// Common handler for tcp work connections.
func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.Plugin,
	baseInfo *config.BaseProxyConf, workConn frpNet.Conn, encKey []byte) {

	var (
		remote io.ReadWriteCloser
		err    error
	)
	remote = workConn
	if baseInfo.UseEncryption {
		remote, err = frpIo.WithEncryption(remote, encKey)
		if err != nil {
			workConn.Error("create encryption stream error: %v", err)
			return
		}
	}
	if baseInfo.UseCompression {
		remote = frpIo.WithCompression(remote)
	}

	if proxyPlugin != nil {
		// if plugin is set, let plugin handle connections first
		workConn.Debug("handle by plugin: %s", proxyPlugin.Name())
		proxyPlugin.Handle(remote)
		workConn.Debug("handle by plugin finished")
		return
	} else {
		localConn, err := frpNet.ConnectServer("tcp", fmt.Sprintf("%s:%d", localInfo.LocalIp, localInfo.LocalPort))
		if err != nil {
			workConn.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIp, localInfo.LocalPort, err)
			return
		}

		workConn.Debug("join connections, localConn(l[%s] r[%s]) workConn(l[%s] r[%s])", localConn.LocalAddr().String(),
			localConn.RemoteAddr().String(), workConn.LocalAddr().String(), workConn.RemoteAddr().String())
		frpIo.Join(localConn, remote)
		workConn.Debug("join connections closed")
	}
}
<commit_msg>Close connection if frpc can't connection to local server<commit_after>// Copyright 2017 fatedier, fatedier@gmail.com
//
// 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 client

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	"github.com/fatedier/frp/models/config"
	"github.com/fatedier/frp/models/msg"
	"github.com/fatedier/frp/models/plugin"
	"github.com/fatedier/frp/models/proto/udp"
	"github.com/fatedier/frp/utils/errors"
	frpIo "github.com/fatedier/frp/utils/io"
	"github.com/fatedier/frp/utils/log"
	frpNet "github.com/fatedier/frp/utils/net"
	"github.com/fatedier/frp/utils/pool"
)

// Proxy defines how to deal with work connections for different proxy type.
type Proxy interface {
	Run() error

	// InWorkConn accept work connections registered to server.
	InWorkConn(conn frpNet.Conn)
	Close()
	log.Logger
}

func NewProxy(ctl *Control, pxyConf config.ProxyConf) (pxy Proxy) {
	baseProxy := BaseProxy{
		ctl:    ctl,
		Logger: log.NewPrefixLogger(pxyConf.GetName()),
	}
	switch cfg := pxyConf.(type) {
	case *config.TcpProxyConf:
		pxy = &TcpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.UdpProxyConf:
		pxy = &UdpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.HttpProxyConf:
		pxy = &HttpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.HttpsProxyConf:
		pxy = &HttpsProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.StcpProxyConf:
		pxy = &StcpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	case *config.XtcpProxyConf:
		pxy = &XtcpProxy{
			BaseProxy: baseProxy,
			cfg:       cfg,
		}
	}
	return
}

type BaseProxy struct {
	ctl    *Control
	closed bool
	mu     sync.RWMutex
	log.Logger
}

// TCP
type TcpProxy struct {
	BaseProxy

	cfg         *config.TcpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *TcpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *TcpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *TcpProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// HTTP
type HttpProxy struct {
	BaseProxy

	cfg         *config.HttpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *HttpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *HttpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *HttpProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// HTTPS
type HttpsProxy struct {
	BaseProxy

	cfg         *config.HttpsProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *HttpsProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *HttpsProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *HttpsProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// STCP
type StcpProxy struct {
	BaseProxy

	cfg         *config.StcpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *StcpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *StcpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *StcpProxy) InWorkConn(conn frpNet.Conn) {
	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf, conn,
		[]byte(config.ClientCommonCfg.PrivilegeToken))
}

// XTCP
type XtcpProxy struct {
	BaseProxy

	cfg         *config.XtcpProxyConf
	proxyPlugin plugin.Plugin
}

func (pxy *XtcpProxy) Run() (err error) {
	if pxy.cfg.Plugin != "" {
		pxy.proxyPlugin, err = plugin.Create(pxy.cfg.Plugin, pxy.cfg.PluginParams)
		if err != nil {
			return
		}
	}
	return
}

func (pxy *XtcpProxy) Close() {
	if pxy.proxyPlugin != nil {
		pxy.proxyPlugin.Close()
	}
}

func (pxy *XtcpProxy) InWorkConn(conn frpNet.Conn) {
	defer conn.Close()
	var natHoleSidMsg msg.NatHoleSid
	err := msg.ReadMsgInto(conn, &natHoleSidMsg)
	if err != nil {
		pxy.Error("xtcp read from workConn error: %v", err)
		return
	}

	natHoleClientMsg := &msg.NatHoleClient{
		ProxyName: pxy.cfg.ProxyName,
		Sid:       natHoleSidMsg.Sid,
	}
	raddr, _ := net.ResolveUDPAddr("udp",
		fmt.Sprintf("%s:%d", config.ClientCommonCfg.ServerAddr, config.ClientCommonCfg.ServerUdpPort))
	clientConn, err := net.DialUDP("udp", nil, raddr)
	defer clientConn.Close()

	err = msg.WriteMsg(clientConn, natHoleClientMsg)
	if err != nil {
		pxy.Error("send natHoleClientMsg to server error: %v", err)
		return
	}

	// Wait for client address at most 5 seconds.
	var natHoleRespMsg msg.NatHoleResp
	clientConn.SetReadDeadline(time.Now().Add(5 * time.Second))

	buf := pool.GetBuf(1024)
	n, err := clientConn.Read(buf)
	if err != nil {
		pxy.Error("get natHoleRespMsg error: %v", err)
		return
	}
	err = msg.ReadMsgInto(bytes.NewReader(buf[:n]), &natHoleRespMsg)
	if err != nil {
		pxy.Error("get natHoleRespMsg error: %v", err)
		return
	}
	clientConn.SetReadDeadline(time.Time{})
	clientConn.Close()
	pxy.Trace("get natHoleRespMsg, sid [%s], client address [%s]", natHoleRespMsg.Sid, natHoleRespMsg.ClientAddr)

	// Send sid to visitor udp address.
	time.Sleep(time.Second)
	laddr, _ := net.ResolveUDPAddr("udp", clientConn.LocalAddr().String())
	daddr, err := net.ResolveUDPAddr("udp", natHoleRespMsg.VisitorAddr)
	if err != nil {
		pxy.Error("resolve visitor udp address error: %v", err)
		return
	}

	lConn, err := net.DialUDP("udp", laddr, daddr)
	if err != nil {
		pxy.Error("dial visitor udp address error: %v", err)
		return
	}
	lConn.Write([]byte(natHoleRespMsg.Sid))

	kcpConn, err := frpNet.NewKcpConnFromUdp(lConn, true, natHoleRespMsg.VisitorAddr)
	if err != nil {
		pxy.Error("create kcp connection from udp connection error: %v", err)
		return
	}

	HandleTcpWorkConnection(&pxy.cfg.LocalSvrConf, pxy.proxyPlugin, &pxy.cfg.BaseProxyConf,
		frpNet.WrapConn(kcpConn), []byte(pxy.cfg.Sk))
}

// UDP
type UdpProxy struct {
	BaseProxy

	cfg *config.UdpProxyConf

	localAddr *net.UDPAddr
	readCh    chan *msg.UdpPacket

	// include msg.UdpPacket and msg.Ping
	sendCh   chan msg.Message
	workConn frpNet.Conn
}

func (pxy *UdpProxy) Run() (err error) {
	pxy.localAddr, err = net.ResolveUDPAddr("udp", fmt.Sprintf("%s:%d", pxy.cfg.LocalIp, pxy.cfg.LocalPort))
	if err != nil {
		return
	}
	return
}

func (pxy *UdpProxy) Close() {
	pxy.mu.Lock()
	defer pxy.mu.Unlock()

	if !pxy.closed {
		pxy.closed = true
		if pxy.workConn != nil {
			pxy.workConn.Close()
		}
		if pxy.readCh != nil {
			close(pxy.readCh)
		}
		if pxy.sendCh != nil {
			close(pxy.sendCh)
		}
	}
}

func (pxy *UdpProxy) InWorkConn(conn frpNet.Conn) {
	pxy.Info("incoming a new work connection for udp proxy, %s", conn.RemoteAddr().String())
	// close resources releated with old workConn
	pxy.Close()

	pxy.mu.Lock()
	pxy.workConn = conn
	pxy.readCh = make(chan *msg.UdpPacket, 1024)
	pxy.sendCh = make(chan msg.Message, 1024)
	pxy.closed = false
	pxy.mu.Unlock()

	workConnReaderFn := func(conn net.Conn, readCh chan *msg.UdpPacket) {
		for {
			var udpMsg msg.UdpPacket
			if errRet := msg.ReadMsgInto(conn, &udpMsg); errRet != nil {
				pxy.Warn("read from workConn for udp error: %v", errRet)
				return
			}
			if errRet := errors.PanicToError(func() {
				pxy.Trace("get udp package from workConn: %s", udpMsg.Content)
				readCh <- &udpMsg
			}); errRet != nil {
				pxy.Info("reader goroutine for udp work connection closed: %v", errRet)
				return
			}
		}
	}
	workConnSenderFn := func(conn net.Conn, sendCh chan msg.Message) {
		defer func() {
			pxy.Info("writer goroutine for udp work connection closed")
		}()
		var errRet error
		for rawMsg := range sendCh {
			switch m := rawMsg.(type) {
			case *msg.UdpPacket:
				pxy.Trace("send udp package to workConn: %s", m.Content)
			case *msg.Ping:
				pxy.Trace("send ping message to udp workConn")
			}
			if errRet = msg.WriteMsg(conn, rawMsg); errRet != nil {
				pxy.Error("udp work write error: %v", errRet)
				return
			}
		}
	}
	heartbeatFn := func(conn net.Conn, sendCh chan msg.Message) {
		var errRet error
		for {
			time.Sleep(time.Duration(30) * time.Second)
			if errRet = errors.PanicToError(func() {
				sendCh <- &msg.Ping{}
			}); errRet != nil {
				pxy.Trace("heartbeat goroutine for udp work connection closed")
				break
			}
		}
	}

	go workConnSenderFn(pxy.workConn, pxy.sendCh)
	go workConnReaderFn(pxy.workConn, pxy.readCh)
	go heartbeatFn(pxy.workConn, pxy.sendCh)
	udp.Forwarder(pxy.localAddr, pxy.readCh, pxy.sendCh)
}

// Common handler for tcp work connections.
func HandleTcpWorkConnection(localInfo *config.LocalSvrConf, proxyPlugin plugin.Plugin,
	baseInfo *config.BaseProxyConf, workConn frpNet.Conn, encKey []byte) {

	var (
		remote io.ReadWriteCloser
		err    error
	)
	remote = workConn
	if baseInfo.UseEncryption {
		remote, err = frpIo.WithEncryption(remote, encKey)
		if err != nil {
			workConn.Error("create encryption stream error: %v", err)
			return
		}
	}
	if baseInfo.UseCompression {
		remote = frpIo.WithCompression(remote)
	}

	if proxyPlugin != nil {
		// if plugin is set, let plugin handle connections first
		workConn.Debug("handle by plugin: %s", proxyPlugin.Name())
		proxyPlugin.Handle(remote)
		workConn.Debug("handle by plugin finished")
		return
	} else {
		localConn, err := frpNet.ConnectServer("tcp", fmt.Sprintf("%s:%d", localInfo.LocalIp, localInfo.LocalPort))
		if err != nil {
			workConn.Error("connect to local service [%s:%d] error: %v", localInfo.LocalIp, localInfo.LocalPort, err)
			remote.Close()
			return
		}

		workConn.Debug("join connections, localConn(l[%s] r[%s]) workConn(l[%s] r[%s])", localConn.LocalAddr().String(),
			localConn.RemoteAddr().String(), workConn.LocalAddr().String(), workConn.RemoteAddr().String())
		frpIo.Join(localConn, remote)
		workConn.Debug("join connections closed")
	}
}
<|endoftext|>


<commit_before>// Copyright © 2016 NAME HERE <EMAIL ADDRESS>
//
// 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 cmd

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/ezbuy/ezorm/db"
	"github.com/ezbuy/ezorm/parser"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
)

// genmsormCmd represents the genmsorm command
var genmsormCmd = &cobra.Command{
	Use:   "genmsorm",
	Short: "Generate sql server orm code",
	Run: func(cmd *cobra.Command, args []string) {
		sqlServer := db.GetSqlServer(dbConfig)
		if table != "all" {
			handler(table, sqlServer)
		} else {
			tables := getAllTables(sqlServer)
			for _, t := range tables {
				handler(t, sqlServer)
			}
		}

		fmt.Println("genmsorm called")
	},
}

var table string
var outputYaml string
var dbConfig string
var packageName string

type ColumnInfo struct {
	ColumnName    string        `db:"ColumnName"`
	DataType      string        `db:"DataType"`
	MaxLength     int           `db:"MaxLength"`
	Nullable      bool          `db:"Nullable"`
	IsPrimaryKey  bool          `db:"IsPrimaryKey"`
	Sort          int           `db:"Sort"`
	IndexId       sql.NullInt64 `db:"IndexId"`
	IndexColumnId sql.NullInt64 `db:"IndexColumnId"`
	IsUnique      sql.NullBool  `db:"IsUnique"`
}

func handler(table string, sqlServer *db.SqlServer) {
	columnsinfo := getColumnInfo(table, sqlServer)
	createYamlFile(table, columnsinfo)
	generate(table)
}

func getAllTables(sqlServer *db.SqlServer) (tables []string) {
	query := `SELECT name FROM sys.tables`
	err := sqlServer.Query(&tables, query)
	if err != nil {
		panic(err)
	}

	return tables
}

func createYamlFile(table string, columns []*ColumnInfo) {
	objs := mapper(table, columns)
	bs, err := yaml.Marshal(objs)
	if err != nil {
		panic(err)
	}
	fileName := outputYaml + "/" + strings.ToLower(table) + "_mssql.yaml"
	ioutil.WriteFile(fileName, bs, 0644)
}

func getIndexInfo(columns []*ColumnInfo) (multiColumnIndexes, multiColumnUniques [][]string,
	singleColumnIndexSet, singleColumnUniqueSet map[int64]struct{}) {
	indexIdToColumns := make(map[int64][]*ColumnInfo)
	for _, v := range columns {
		indexId := v.IndexId.Int64
		if indexId > 0 && !v.IsPrimaryKey {
			indexIdToColumns[indexId] = append(indexIdToColumns[indexId], v)
		}
	}

	singleColumnIndexSet = make(map[int64]struct{})
	singleColumnUniqueSet = make(map[int64]struct{})
	for indexId, indexColums := range indexIdToColumns {
		if len(indexColums) == 1 {
			if indexColums[0].IsUnique.Bool {
				singleColumnUniqueSet[indexId] = struct{}{}
			} else {
				singleColumnIndexSet[indexId] = struct{}{}
			}
		} else {
			columnNames := make([]string, 0, len(indexColums))
			// Note: columns are sorted by IndexColumdId
			for _, c := range indexColums {
				columnNames = append(columnNames, c.ColumnName)
			}

			if indexColums[0].IsUnique.Bool {
				multiColumnUniques = append(multiColumnUniques, columnNames)
			} else {
				multiColumnIndexes = append(multiColumnIndexes, columnNames)
			}
		}
	}

	return
}

type tbl struct {
	DB      string        `yaml:"db"`
	Fields  []interface{} `yaml:"fields"`
	Indexes [][]string    `yaml:"indexes,flow"`
	Uniques [][]string    `yaml:"uniques,flow"`
}

func mapper(table string, columns []*ColumnInfo) map[string]*tbl {
	multiColumnIndexes, multiColumnUniques, singleColumnIndexSet, singleColumnUniqueSet := getIndexInfo(columns)

	var t tbl
	t.DB = "mssql"
	t.Indexes = multiColumnIndexes
	t.Uniques = multiColumnUniques
	objs := make(map[string]*tbl)
	objs[table] = &t
	fields := make([]interface{}, len(columns))
	for i, v := range columns {
		dataitem := make(map[string]interface{}, len(columns))
		dataitem[v.ColumnName] = parser.DbToGoType(v.DataType)
		if dataitem[v.ColumnName] == "time.Time" {
			parser.HaveTime = true
		}

		if _, ok := singleColumnUniqueSet[v.IndexId.Int64]; ok {
			dataitem["flags"] = []string{"unique"}
		} else if _, ok := singleColumnIndexSet[v.IndexId.Int64]; ok {
			dataitem["flags"] = []string{"index"}
		}

		fields[i] = dataitem
	}
	t.Fields = fields
	return objs
}

func getColumnInfo(table string, sqlServer *db.SqlServer) []*ColumnInfo {
	// Note: sort columns by IndexId and IndexColumnId to simplify later process
	query := `SELECT DISTINCT c.name AS ColumnName, t.Name AS DataType, c.max_length AS MaxLength,
    c.is_nullable AS Nullable, ISNULL(i.is_primary_key, 0) AS IsPrimaryKey ,c.column_id AS Sort,
	i.index_id AS IndexId, ic.index_column_id AS IndexColumnId, i.is_unique AS IsUnique
	FROM
    sys.columns c
	INNER JOIN
    sys.types t ON c.user_type_id = t.user_type_id
	LEFT OUTER JOIN
    sys.index_columns ic ON ic.object_id = c.object_id AND ic.column_id = c.column_id
	LEFT OUTER JOIN
    sys.indexes i ON ic.object_id = i.object_id AND ic.index_id = i.index_id
	WHERE
    c.object_id = OBJECT_ID(?) ORDER BY IndexId, IndexColumnId`

	var columninfos []*ColumnInfo
	err := sqlServer.Query(&columninfos, query, table)
	if err != nil {
		panic(err)
	}

	return columninfos
}

func generate(table string) {
	var objs map[string]map[string]interface{}
	fileName := outputYaml + "/" + strings.ToLower(table) + "_mssql.yaml"
	data, _ := ioutil.ReadFile(fileName)
	_, err := os.Stat(fileName)
	if err != nil {
		panic(err)
	}
	err = yaml.Unmarshal([]byte(data), &objs)

	if packageName == "" {
		packageName = strings.ToLower(table)
	}
	for key, obj := range objs {
		metaObj := new(parser.Obj)
		metaObj.Package = packageName
		metaObj.Name = key
		metaObj.Db = obj["db"].(string)
		err := metaObj.Read(obj)
		if err != nil {
			panic(err)
		}

		for _, genType := range metaObj.GetGenTypes() {
			file, err := os.OpenFile(output+"/gen_"+metaObj.Name+"_"+genType+".go", os.O_WRONLY|os.O_CREATE|os.O_TRUNC, 0644)
			metaObj.TplWriter = file
			if err != nil {
				panic(err)
			}

			err = parser.Tpl.ExecuteTemplate(file, genType, metaObj)
			file.Close()
			if err != nil {
				panic(err)
			}
		}
	}

	cmd := exec.Command("gofmt", "-w", output)
	cmd.Run()
}

func init() {
	RootCmd.AddCommand(genmsormCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// genmsormCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// genmsormCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
	genmsormCmd.PersistentFlags().StringVarP(&table, "table", "t", "all", "table name, 'all' meaning all tables")
	genmsormCmd.PersistentFlags().StringVarP(&output, "output", "o", "", "output path")
	genmsormCmd.PersistentFlags().StringVarP(&outputYaml, "output yaml", "y", "", "output *.yaml path")
	genmsormCmd.PersistentFlags().StringVarP(&dbConfig, "db config", "d", "", "database configuration")
	genmsormCmd.PersistentFlags().StringVarP(&packageName, "package name", "p", "", "package name")
}
<commit_msg>Fix duplicate definitions for columns in multiple indexes<commit_after>// Copyright © 2016 NAME HERE <EMAIL ADDRESS>
//
// 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 cmd

import (
	"database/sql"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"strings"

	"github.com/ezbuy/ezorm/db"
	"github.com/ezbuy/ezorm/parser"
	"github.com/spf13/cobra"
	"gopkg.in/yaml.v2"
)

// genmsormCmd represents the genmsorm command
var genmsormCmd = &cobra.Command{
	Use:   "genmsorm",
	Short: "Generate sql server orm code",
	Run: func(cmd *cobra.Command, args []string) {
		sqlServer := db.GetSqlServer(dbConfig)
		if table != "all" {
			handler(table, sqlServer)
		} else {
			tables := getAllTables(sqlServer)
			for _, t := range tables {
				handler(t, sqlServer)
			}
		}

		fmt.Println("genmsorm called")
	},
}

var table string
var outputYaml string
var dbConfig string
var packageName string

type ColumnInfo struct {
	ColumnName    string        `db:"ColumnName"`
	DataType      string        `db:"DataType"`
	MaxLength     int           `db:"MaxLength"`
	Nullable      bool          `db:"Nullable"`
	IsPrimaryKey  bool          `db:"IsPrimaryKey"`
	Sort          int           `db:"Sort"`
	IndexId       sql.NullInt64 `db:"IndexId"`
	IndexColumnId sql.NullInt64 `db:"IndexColumnId"`
	IsUnique      sql.NullBool  `db:"IsUnique"`
}

func handler(table string, sqlServer *db.SqlServer) {
	columnsinfo := getColumnInfo(table, sqlServer)
	createYamlFile(table, columnsinfo)
	generate(table)
}

func getAllTables(sqlServer *db.SqlServer) (tables []string) {
	query := `SELECT name FROM sys.tables`
	err := sqlServer.Query(&tables, query)
	if err != nil {
		panic(err)
	}

	return tables
}

func createYamlFile(table string, columns []*ColumnInfo) {
	objs := mapper(table, columns)
	bs, err := yaml.Marshal(objs)
	if err != nil {
		panic(err)
	}
	fileName := getOutYamlFileName(table)
	ioutil.WriteFile(fileName, bs, 0644)
}

func getIndexInfo(columns []*ColumnInfo) (multiColumnIndexes, multiColumnUniques [][]string,
	singleColumnIndexSet, singleColumnUniqueSet map[int64]struct{}) {
	indexIdToColumns := make(map[int64][]*ColumnInfo)
	for _, v := range columns {
		indexId := v.IndexId.Int64
		if indexId > 0 && !v.IsPrimaryKey {
			indexIdToColumns[indexId] = append(indexIdToColumns[indexId], v)
		}
	}

	singleColumnIndexSet = make(map[int64]struct{})
	singleColumnUniqueSet = make(map[int64]struct{})
	for indexId, indexColums := range indexIdToColumns {
		if len(indexColums) == 1 {
			if indexColums[0].IsUnique.Bool {
				singleColumnUniqueSet[indexId] = struct{}{}
			} else {
				singleColumnIndexSet[indexId] = struct{}{}
			}
		} else {
			columnNames := make([]string, 0, len(indexColums))
			// Note: columns are sorted by IndexColumdId
			for _, c := range indexColums {
				columnNames = append(columnNames, c.ColumnName)
			}

			if indexColums[0].IsUnique.Bool {
				multiColumnUniques = append(multiColumnUniques, columnNames)
			} else {
				multiColumnIndexes = append(multiColumnIndexes, columnNames)
			}
		}
	}

	return
}

type tbl struct {
	DB      string        `yaml:"db"`
	Fields  []interface{} `yaml:"fields"`
	Indexes [][]string    `yaml:"indexes,flow"`
	Uniques [][]string    `yaml:"uniques,flow"`
}

func mapper(table string, columns []*ColumnInfo) map[string]*tbl {
	// 如果第一个字符不是大写，变为大写
	for _, v := range columns {
		colName := v.ColumnName
		firstChar := colName[:1]
		firstCharUpper := strings.ToUpper(firstChar)
		if firstChar != firstCharUpper {
			v.ColumnName = firstCharUpper + colName[1:]
		}
	}

	multiColumnIndexes, multiColumnUniques, singleColumnIndexSet, singleColumnUniqueSet := getIndexInfo(columns)

	var t tbl
	t.DB = "mssql"
	t.Indexes = multiColumnIndexes
	t.Uniques = multiColumnUniques
	objs := make(map[string]*tbl)
	objs[table] = &t
	lenColumns := len(columns)
	fields := make([]interface{}, 0, lenColumns)
	processedFields := make(map[string]struct{}, lenColumns)
	for _, v := range columns {
		// 有的字段可能出现在多个索引内，排除掉已经处理掉字段
		if _, ok := processedFields[v.ColumnName]; ok {
			continue
		}
		processedFields[v.ColumnName] = struct{}{}

		dataitem := make(map[string]interface{}, len(columns))
		dataitem[v.ColumnName] = parser.DbToGoType(v.DataType)
		isTypeTime := dataitem[v.ColumnName] == "time.Time"
		if isTypeTime {
			parser.HaveTime = true
		}

		var flags []string
		if _, ok := singleColumnUniqueSet[v.IndexId.Int64]; ok {
			flags = append(flags, "unique")
		} else if _, ok := singleColumnIndexSet[v.IndexId.Int64]; ok {
			flags = append(flags, "index")
		}

		if v.Nullable {
			flags = append(flags, "nullable")
		}

		if flags != nil {
			dataitem["flags"] = flags
		}
		fields = append(fields, dataitem)
	}
	t.Fields = fields
	return objs
}

func getColumnInfo(table string, sqlServer *db.SqlServer) []*ColumnInfo {
	// Note: sort columns by IndexId and IndexColumnId to simplify later process
	query := `SELECT DISTINCT c.name AS ColumnName, t.Name AS DataType, c.max_length AS MaxLength,
    c.is_nullable AS Nullable, ISNULL(i.is_primary_key, 0) AS IsPrimaryKey ,c.column_id AS Sort,
	i.index_id AS IndexId, ic.index_column_id AS IndexColumnId, i.is_unique AS IsUnique
	FROM
    sys.columns c
	INNER JOIN
    sys.types t ON c.user_type_id = t.user_type_id
	LEFT OUTER JOIN
    sys.index_columns ic ON ic.object_id = c.object_id AND ic.column_id = c.column_id
	LEFT OUTER JOIN
    sys.indexes i ON ic.object_id = i.object_id AND ic.index_id = i.index_id
	WHERE
    c.object_id = OBJECT_ID(?) ORDER BY IndexId, IndexColumnId`

	var columninfos []*ColumnInfo
	err := sqlServer.Query(&columninfos, query, table)
	if err != nil {
		panic(err)
	}

	return columninfos
}

func getOutYamlFileName(table string) string {
	return outputYaml + "/gen_" + strings.ToLower(table) + "_mssql.yaml"
}

func generate(table string) {
	var objs map[string]map[string]interface{}
	fileName := getOutYamlFileName(table)
	data, _ := ioutil.ReadFile(fileName)
	_, err := os.Stat(fileName)
	if err != nil {
		panic(err)
	}
	err = yaml.Unmarshal([]byte(data), &objs)

	if packageName == "" {
		packageName = strings.ToLower(table)
	}

	genConfigDone := false
	for key, obj := range objs {
		metaObj := new(parser.Obj)
		metaObj.Package = packageName
		metaObj.Name = key
		metaObj.Db = obj["db"].(string)
		err := metaObj.Read(obj)
		if err != nil {
			panic(err)
		}

		if !genConfigDone {
			if tpl, ok := metaObj.GetConfigTemplate(); ok {
				fileAbsPath := output + "/gen_" + metaObj.Db + "_config.go"
				executeTpl(fileAbsPath, tpl, metaObj)
				genConfigDone = true
			}
		}

		for _, genType := range metaObj.GetGenTypes() {
			fileAbsPath := output + "/gen_" + metaObj.Name + "_" + genType + ".go"
			executeTpl(fileAbsPath, genType, metaObj)
		}
	}

	cmd := exec.Command("gofmt", "-w", output)
	cmd.Run()
}

func init() {
	RootCmd.AddCommand(genmsormCmd)

	// Here you will define your flags and configuration settings.

	// Cobra supports Persistent Flags which will work for this command
	// and all subcommands, e.g.:
	// genmsormCmd.PersistentFlags().String("foo", "", "A help for foo")

	// Cobra supports local flags which will only run when this command
	// is called directly, e.g.:
	// genmsormCmd.Flags().BoolP("toggle", "t", false, "Help message for toggle")
	genmsormCmd.PersistentFlags().StringVarP(&table, "table", "t", "all", "table name, 'all' meaning all tables")
	genmsormCmd.PersistentFlags().StringVarP(&output, "output", "o", "", "output path")
	genmsormCmd.PersistentFlags().StringVarP(&outputYaml, "output yaml", "y", "", "output *.yaml path")
	genmsormCmd.PersistentFlags().StringVarP(&dbConfig, "db config", "d", "", "database configuration")
	genmsormCmd.PersistentFlags().StringVarP(&packageName, "package name", "p", "", "package name")
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"

	"github.com/aryann/difflib"
	"github.com/tucnak/climax"

	"bitbucket.org/homemade/scl"
)

func main() {
	app := climax.New("scl")
	app.Brief = "Scl is a tool for managing SCL soure code."
	app.Version = "1.1.1"

	app.AddCommand(runCommand(os.Stdout, os.Stderr))
	app.AddCommand(testCommand(os.Stdout, os.Stderr))

	os.Exit(app.Run())
}

func runCommand(stdout io.Writer, stderr io.Writer) climax.Command {

	return climax.Command{
		Name:  "run",
		Brief: "Transform one or more .scl files into HCL",
		Usage: `[options] <filename.scl> [filenameX.scl...]`,
		Help:  `Transform one or more .scl files into HCL. Output is written to stdout.`,

		Flags: standardParserParams(),

		Handle: func(ctx climax.Context) int {

			if len(ctx.Args) == 0 {
				fmt.Fprintf(stderr, "At least one filename is required. See `sep help run` for syntax")
				return 1
			}

			params, includePaths := parserParams(ctx)

			for _, fileName := range ctx.Args {

				parser, err := scl.NewParser(scl.NewDiskSystem())

				if err != nil {
					fmt.Fprintf(stderr, "Error: Unable to create new parser in CWD: %s\n", err.Error())
					return 1
				}

				for _, includeDir := range includePaths {
					parser.AddIncludePath(includeDir)
				}

				for _, p := range params {
					parser.SetParam(p.name, p.value)
				}

				if err := parser.Parse(fileName); err != nil {
					fmt.Fprintf(stderr, "Error: Unable to parse file: %s\n", err.Error())
					return 1
				}

				fmt.Fprintf(stdout, "/* %s */\n%s\n\n", fileName, parser)
			}

			return 0
		},
	}
}

func testCommand(stdout io.Writer, stderr io.Writer) climax.Command {

	return climax.Command{
		Name:  "test",
		Brief: "Parse each .scl file in a directory and compare the output to an .hcl file",
		Usage: `[options] [file-glob...]`,
		Help:  "Parse each .scl file in a directory and compare the output to an .hcl file",

		Flags: standardParserParams(),

		Handle: func(ctx climax.Context) int {

			errors := 0

			reportError := func(path string, err string, args ...interface{}) {
				fmt.Fprintf(stderr, "[%s] %s\n", path, fmt.Sprintf(err, args...))
				errors++
			}

			if len(ctx.Args) == 0 {
				fmt.Fprintf(stderr, "At least one file glob is required. See `sep help test` for syntax")
				return 1
			}

			newlineMatcher := regexp.MustCompile("\n\n")
			params, includePaths := parserParams(ctx)

			for _, fileName := range ctx.Args {

				fs := scl.NewDiskSystem()
				parser, err := scl.NewParser(fs)

				if err != nil {
					reportError("Unable to create new parser in CWD: %s", err.Error())
					continue
				}

				for _, includeDir := range includePaths {
					parser.AddIncludePath(includeDir)
				}

				for _, p := range params {
					parser.SetParam(p.name, p.value)
				}

				if err := parser.Parse(fileName); err != nil {
					reportError(fileName, "Unable to parse file: %s", err.Error())
					continue
				}

				hclFilePath := strings.TrimSuffix(fileName, ".scl") + ".hcl"
				hclFile, _, err := fs.ReadCloser(hclFilePath)

				if err != nil {
					fmt.Fprintf(stdout, "[%s] No HCL file; skipping gold standard test\n", fileName)
					continue
				}

				hcl, err := ioutil.ReadAll(hclFile)

				if err != nil {
					reportError(fileName, "Unable to read .hcl file: %s", err.Error())
					continue
				}

				hclLines := strings.Split(strings.TrimSuffix(newlineMatcher.ReplaceAllString(string(hcl), "\n"), "\n"), "\n")
				sclLines := strings.Split(parser.String(), "\n")

				diff := difflib.Diff(hclLines, sclLines)

				success := true

				for _, d := range diff {
					if d.Delta != difflib.Common {
						success = false
					}
				}

				if !success {
					reportError(fileName, "Diff failed:")

					fmt.Fprintln(stderr)

					for _, d := range diff {
						fmt.Fprintf(stderr, "\t%s\n", d.String())
					}

					fmt.Fprintln(stderr)

					continue
				}

				fmt.Fprintf(stdout, "[%s] OK\n", fileName)
			}

			if errors > 0 {
				fmt.Fprintf(stderr, "\n[FAIL] %d error(s)\n", errors)
				return 1
			}

			return 0
		},
	}
}

func standardParserParams() []climax.Flag {

	return []climax.Flag{
		{
			Name:     "include",
			Short:    "i",
			Usage:    `--include /path/to/lib1,/path/to/lib2`,
			Help:     `Comma-separated list of include paths`,
			Variable: true,
		},
		{
			Name:     "param",
			Short:    "p",
			Usage:    `--param param0=somthing,"param1='something else'"`,
			Help:     `Comma-separated list of include paths`,
			Variable: true,
		},
		{
			Name:  "no-env",
			Short: "ne",
			Usage: `--no-env`,
			Help:  `Don't import envionment variables when parsing the SCL`,
		},
	}

}

func parserParams(ctx climax.Context) (params paramSlice, includePaths []string) {

	if !ctx.Is("no-env") {
		for _, envVar := range os.Environ() {
			params.Set(envVar)
		}
	}

	if ps, set := ctx.Get("param"); set {
		for _, p := range strings.Split(ps, ",") {
			params.Set(p)
		}
	}

	if ps, set := ctx.Get("include"); set {
		for _, i := range strings.Split(ps, ",") {
			includePaths = append(includePaths, i)
		}
	}

	return
}
<commit_msg>make output consistent with go test<commit_after>package main

import (
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"regexp"
	"strings"
	"time"

	"github.com/aryann/difflib"
	"github.com/tucnak/climax"

	"bitbucket.org/homemade/scl"
)

func main() {
	app := climax.New("scl")
	app.Brief = "Scl is a tool for managing SCL soure code."
	app.Version = "1.1.1"

	app.AddCommand(runCommand(os.Stdout, os.Stderr))
	app.AddCommand(testCommand(os.Stdout, os.Stderr))

	os.Exit(app.Run())
}

func runCommand(stdout io.Writer, stderr io.Writer) climax.Command {

	return climax.Command{
		Name:  "run",
		Brief: "Transform one or more .scl files into HCL",
		Usage: `[options] <filename.scl> [filenameX.scl...]`,
		Help:  `Transform one or more .scl files into HCL. Output is written to stdout.`,

		Flags: standardParserParams(),

		Handle: func(ctx climax.Context) int {

			if len(ctx.Args) == 0 {
				fmt.Fprintf(stderr, "At least one filename is required. See `sep help run` for syntax")
				return 1
			}

			params, includePaths := parserParams(ctx)

			for _, fileName := range ctx.Args {

				parser, err := scl.NewParser(scl.NewDiskSystem())

				if err != nil {
					fmt.Fprintf(stderr, "Error: Unable to create new parser in CWD: %s\n", err.Error())
					return 1
				}

				for _, includeDir := range includePaths {
					parser.AddIncludePath(includeDir)
				}

				for _, p := range params {
					parser.SetParam(p.name, p.value)
				}

				if err := parser.Parse(fileName); err != nil {
					fmt.Fprintf(stderr, "Error: Unable to parse file: %s\n", err.Error())
					return 1
				}

				fmt.Fprintf(stdout, "/* %s */\n%s\n\n", fileName, parser)
			}

			return 0
		},
	}
}

func testCommand(stdout io.Writer, stderr io.Writer) climax.Command {

	return climax.Command{
		Name:  "test",
		Brief: "Parse each .scl file in a directory and compare the output to an .hcl file",
		Usage: `[options] [file-glob...]`,
		Help:  "Parse each .scl file in a directory and compare the output to an .hcl file",

		Flags: standardParserParams(),

		Handle: func(ctx climax.Context) int {

			errors := 0

			reportError := func(path string, err string, args ...interface{}) {
				fmt.Fprintf(stderr, "%-7s %s %s\n", "FAIL", path, fmt.Sprintf(err, args...))
				errors++
			}

			if len(ctx.Args) == 0 {
				fmt.Fprintf(stderr, "At least one file glob is required. See `sep help test` for syntax")
				return 1
			}

			newlineMatcher := regexp.MustCompile("\n\n")
			params, includePaths := parserParams(ctx)

			for _, fileName := range ctx.Args {

				fs := scl.NewDiskSystem()
				parser, err := scl.NewParser(fs)
				now := time.Now()

				if err != nil {
					reportError("Unable to create new parser in CWD: %s", err.Error())
					continue
				}

				for _, includeDir := range includePaths {
					parser.AddIncludePath(includeDir)
				}

				for _, p := range params {
					parser.SetParam(p.name, p.value)
				}

				if err := parser.Parse(fileName); err != nil {
					reportError(fileName, "Unable to parse file: %s", err.Error())
					continue
				}

				hclFilePath := strings.TrimSuffix(fileName, ".scl") + ".hcl"
				hclFile, _, err := fs.ReadCloser(hclFilePath)

				if err != nil {
					fmt.Fprintf(stdout, "%-7s %s [no .hcl file]\n", "?", fileName)
					continue
				}

				hcl, err := ioutil.ReadAll(hclFile)

				if err != nil {
					reportError(fileName, "Unable to read .hcl file: %s", err.Error())
					continue
				}

				hclLines := strings.Split(strings.TrimSuffix(newlineMatcher.ReplaceAllString(string(hcl), "\n"), "\n"), "\n")
				sclLines := strings.Split(parser.String(), "\n")

				diff := difflib.Diff(hclLines, sclLines)

				success := true

				for _, d := range diff {
					if d.Delta != difflib.Common {
						success = false
					}
				}

				if !success {
					reportError(fileName, "Diff failed:")

					fmt.Fprintln(stderr)

					for _, d := range diff {
						fmt.Fprintf(stderr, "\t%s\n", d.String())
					}

					fmt.Fprintln(stderr)

					continue
				}

				fmt.Fprintf(stdout, "%-7s %s\t%.3fs\n", "ok", fileName, time.Since(now).Seconds())
			}

			if errors > 0 {
				fmt.Fprintf(stderr, "\n[FAIL] %d error(s)\n", errors)
				return 1
			}

			return 0
		},
	}
}

func standardParserParams() []climax.Flag {

	return []climax.Flag{
		{
			Name:     "include",
			Short:    "i",
			Usage:    `--include /path/to/lib1,/path/to/lib2`,
			Help:     `Comma-separated list of include paths`,
			Variable: true,
		},
		{
			Name:     "param",
			Short:    "p",
			Usage:    `--param param0=somthing,"param1='something else'"`,
			Help:     `Comma-separated list of include paths`,
			Variable: true,
		},
		{
			Name:  "no-env",
			Short: "ne",
			Usage: `--no-env`,
			Help:  `Don't import envionment variables when parsing the SCL`,
		},
	}

}

func parserParams(ctx climax.Context) (params paramSlice, includePaths []string) {

	if !ctx.Is("no-env") {
		for _, envVar := range os.Environ() {
			params.Set(envVar)
		}
	}

	if ps, set := ctx.Get("param"); set {
		for _, p := range strings.Split(ps, ",") {
			params.Set(p)
		}
	}

	if ps, set := ctx.Get("include"); set {
		for _, i := range strings.Split(ps, ",") {
			includePaths = append(includePaths, i)
		}
	}

	return
}
<|endoftext|>


<commit_before>package stock

import (
	"bufio"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"regexp"
	"strings"

	"github.com/nzai/Tast/config"
)

const (
	stocksFileName = "stocks.txt"
	//	从智库百科-标准普尔500指数页面下载成份股
	sp100url = "http://wiki.mbalib.com/wiki/%E6%A0%87%E5%87%86%E6%99%AE%E5%B0%94100%E6%8C%87%E6%95%B0"
)

//	股票
type Stock struct {
	Code        string
	EnglishName string
	ChineseName string
}

//	更新股票列表
func UpdateAll() error {

	log.Println("开始更新股票列表")
	//	更新股票
	_, err := GetAll()

	log.Println("股票列表更新结束")

	return err
}

//	获取股票列表
func GetAll() ([]Stock, error) {

	//	数据保存目录
	dataDir, err := config.GetDataDir()
	if err != nil {
		return nil, err
	}

	//	股票列表文件路径
	filePath := filepath.Join(dataDir, stocksFileName)
	_, err = os.Stat(filePath)
	if os.IsNotExist(err) {
		//	如果股票列表文件不存在，则从智库百科下载
		stocks, err := downloadFromMbalib()
		if err != nil {
			return nil, err
		}
		//	保存下载的股票
		return stocks, save(stocks, filePath)
	}

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

	scanner := bufio.NewScanner(file)
	stocks := make([]Stock, 0)
	for scanner.Scan() {
		parts := strings.Split(scanner.Text(), "\t")
		if len(parts) != 3 {
			return nil, errors.New("股票列表文件格式不正确")
		}

		stocks = append(stocks, Stock{
			Code:        strings.ToUpper(parts[0]),
			EnglishName: parts[1],
			ChineseName: parts[2],
		})
	}

	return stocks, nil
}

//	从智库百科下载标普100成份股
func downloadFromMbalib() ([]Stock, error) {
	response, err := http.Get(sp100url)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	buffer, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	regex := regexp.MustCompile(`<td>(<a[^>]*?>)?(\w+)(</a>)?</td><td><a[^>]*?>([^<]+)</a>([^<]*)?</td><td><a[^>]*?>([^<]+?)</a>`)
	matches := regex.FindAllStringSubmatch(string(buffer), -1)

	stocks := make([]Stock, 0)
	for _, match := range matches {

		if len(match) != 7 {
			return nil, errors.New("股票列表格式不正确")
		}

		stocks = append(stocks, Stock{
			Code:        match[2],
			EnglishName: match[4],
			ChineseName: match[6],
		})
	}

	return stocks, nil
}

//	保存
func save(stocks []Stock, filePath string) error {

	//	打开文件
	file, err := os.OpenFile(filePath, os.O_CREATE, 0x777)
	if err != nil {
		return err
	}
	defer file.Close()

	for _, stock := range stocks {

		line := fmt.Sprintf("%s\t%s\t%s\n",
			stock.Code,
			stock.EnglishName,
			stock.ChineseName)

		//	将股票写入文件
		_, err = file.WriteString(line)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>因为智库百科中标准普尔100成份股很多都已经失效，所以改从纳斯达克100指数页面获取成份股<commit_after>package stock

import (
	"bufio"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"

	"github.com/nzai/Tast/config"
)

const (
	stocksFileName = "stocks.txt"
	//	从智库百科-标准普尔500指数页面下载成份股
	//sp100url     = "http://wiki.mbalib.com/wiki/%E6%A0%87%E5%87%86%E6%99%AE%E5%B0%94100%E6%8C%87%E6%95%B0"
	//	从纳斯达克100指数页面下载成份股
	nasdaq100Url = "http://www.nasdaq.com/quotes/nasdaq-100-stocks.aspx?render=download"
)

//	股票
type Stock struct {
	Code        string
	EnglishName string
}

//	更新股票列表
func UpdateAll() error {

	log.Println("开始更新股票列表")
	//	更新股票
	_, err := GetAll()

	log.Println("股票列表更新结束")

	return err
}

//	获取股票列表
func GetAll() ([]Stock, error) {

	//	数据保存目录
	dataDir, err := config.GetDataDir()
	if err != nil {
		return nil, err
	}

	//	股票列表文件路径
	filePath := filepath.Join(dataDir, stocksFileName)
	_, err = os.Stat(filePath)
	if os.IsNotExist(err) {
		//	如果股票列表文件不存在，则从纳斯达克下载
		stocks, err := downloadFromNasdaq100()
		if err != nil {
			return nil, err
		}

		//	保存下载的股票
		return stocks, save(stocks, filePath)
	}

	return load(filePath)
}

////	从智库百科下载标普100成份股
//func downloadFromMbalib() ([]Stock, error) {
//	response, err := http.Get(sp100url)
//	if err != nil {
//		return nil, err
//	}
//	defer response.Body.Close()

//	buffer, err := ioutil.ReadAll(response.Body)
//	if err != nil {
//		return nil, err
//	}

//	regex := regexp.MustCompile(`<td>(<a[^>]*?>)?(\w+)(</a>)?</td><td><a[^>]*?>([^<]+)</a>([^<]*)?</td><td><a[^>]*?>([^<]+?)</a>`)
//	matches := regex.FindAllStringSubmatch(string(buffer), -1)

//	stocks := make([]Stock, 0)
//	for _, match := range matches {

//		if len(match) != 7 {
//			return nil, errors.New("股票列表格式不正确")
//		}

//		stocks = append(stocks, Stock{
//			Code:        match[2],
//			EnglishName: match[4],
//			ChineseName: match[6],
//		})
//	}

//	return stocks, nil
//}

func downloadFromNasdaq100() ([]Stock, error) {

	response, err := http.Get(nasdaq100Url)
	if err != nil {
		return nil, err
	}
	defer response.Body.Close()

	buffer, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	lines := strings.Split(string(buffer), "\n")
	stocks := make([]Stock, 0)

	//	略过第一行的标题和最后一行空白
	for index := 1; index < len(lines)-1; index++ {

		parts := strings.Split(lines[index], ",")
		if len(parts) < 2 {
			return nil, errors.New("纳斯达克股票列表文件格式不正确")
		}

		stocks = append(stocks, Stock{
			Code:        strings.ToUpper(parts[0]),
			EnglishName: strings.Trim(parts[1], " "),
		})
	}

	return stocks, nil
}

//	保存
func save(stocks []Stock, filePath string) error {

	//	打开文件
	file, err := os.OpenFile(filePath, os.O_CREATE, 0x777)
	if err != nil {
		return err
	}
	defer file.Close()

	for _, stock := range stocks {

		line := fmt.Sprintf("%s\t%s\n", stock.Code, stock.EnglishName)

		//	将股票写入文件
		_, err = file.WriteString(line)
		if err != nil {
			return err
		}
	}
	return nil
}

//	读取
func load(filePath string) ([]Stock, error) {

	//	打开股票列表文件
	file, err := os.Open(filePath)
	if err != nil {
		return nil, err
	}
	defer file.Close()

	scanner := bufio.NewScanner(file)
	stocks := make([]Stock, 0)
	for scanner.Scan() {
		parts := strings.Split(scanner.Text(), "\t")
		if len(parts) != 2 {
			return nil, errors.New("股票列表文件格式不正确")
		}

		stocks = append(stocks, Stock{
			Code:        strings.ToUpper(parts[0]),
			EnglishName: parts[1],
		})
	}

	return stocks, nil
}
<|endoftext|>


<commit_before>package store

import (
	"encoding/json"
	"net"
	"net/url"
	"strings"
	"sync"

	validator "gopkg.in/go-playground/validator.v9"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/libkv"
	kv "github.com/docker/libkv/store"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/types"
	"github.com/pkg/errors"
)

func init() {
	registryStores()
}

func registryStores() {
	libkv.AddStore(kv.CONSUL, consul.New)
	libkv.AddStore(kv.ETCD, etcd.New)
}

type Store interface {
	GetServices() ([]types.Service, error)
	// GetService(serviceId string) (types.Service, error)
	AddService(svc *types.Service) error
	DeleteService(svc *types.Service) error
	SubscribeServices(ch chan []types.Service)
	WatchServices()

	GetDestinations() ([]types.Destination, error)
	AddDestination(svc *types.Service, dst *types.Destination) error
	DeleteDestination(svc *types.Service, dst *types.Destination) error
	SubscribeDestinations(ch chan []types.Destination)
	WatchDestinations()

	AddCheck(check types.CheckSpec) error
	DeleteCheck(check types.CheckSpec) error
	SubscribeChecks(ch chan []types.CheckSpec)
	WatchChecks()

	GetKV() kv.Store
}

var (
	ErrUnsupportedStore = errors.New("[store] Unsupported store")
)

type FusisStore struct {
	sync.Mutex
	kv kv.Store

	prefix string

	validate            *validator.Validate
	servicesChannels    []chan []types.Service
	destinationChannels []chan []types.Destination
	checksChannels      []chan []types.CheckSpec
}

func New(config *config.BalancerConfig) (Store, error) {
	u, err := url.Parse(config.StoreAddress)
	if err != nil {
		return nil, errors.Wrap(err, "error paring store address")
	}

	scheme := u.Scheme
	if scheme != "consul" && scheme != "etcd" {
		return nil, ErrUnsupportedStore
	}

	//Validating open connection
	_, err = net.Dial("tcp", u.Host)
	if err != nil {
		return nil, errors.Wrap(err, "Store connection failed. Make sure your store is up and running.")
	}

	kv, err := libkv.NewStore(
		kv.Backend(scheme),
		[]string{u.Host},
		nil,
	)
	if err != nil {
		kv.Close()
		return nil, errors.Wrap(err, "Cannot create store consul")
	}

	svcsChs := []chan []types.Service{}
	dstsChs := []chan []types.Destination{}
	checksChs := []chan []types.CheckSpec{}

	validate := validator.New()
	// Registering custom validations
	validate.RegisterValidation("protocols", validateValues(types.Protocols))
	validate.RegisterValidation("schedulers", validateValues(types.Schedulers))

	fusisStore := &FusisStore{
		kv:                  kv,
		prefix:              config.StorePrefix,
		validate:            validate,
		servicesChannels:    svcsChs,
		destinationChannels: dstsChs,
		checksChannels:      checksChs,
	}

	go fusisStore.WatchServices()
	go fusisStore.WatchDestinations()
	go fusisStore.WatchChecks()

	return fusisStore, nil
}

func (s *FusisStore) GetKV() kv.Store {
	return s.kv
}

func (s *FusisStore) GetServices() ([]types.Service, error) {
	svcs := []types.Service{}
	entries, err := s.kv.List(s.key("services"))
	if err != nil {
		if err == kv.ErrKeyNotFound {
			return svcs, nil
		}

		return nil, err
	}

	for _, pair := range entries {
		svc := types.Service{}
		if err := json.Unmarshal(pair.Value, &svc); err != nil {
			log.Error("[store] ", err)
		}

		svcs = append(svcs, svc)
	}

	return svcs, nil
}

func (s *FusisStore) GetDestinations() ([]types.Destination, error) {
	dsts := []types.Destination{}
	entries, err := s.kv.List("fusis/destinations")
	if err != nil {
		return nil, err
	}

	for _, pair := range entries {
		dst := types.Destination{}
		if err := json.Unmarshal(pair.Value, &dst); err != nil {
			log.Error("[store] ", err)
		}

		dsts = append(dsts, dst)
	}

	return dsts, nil
}

// AddService adds a new services to store. It validates the name
// uniqueness and the IPVS uniqueness by saving the IPVS key
// in the store, which consists in a combination of address, port
// and protocol.
func (s *FusisStore) AddService(svc *types.Service) error {
	svcKey := s.key("services", svc.GetId(), "config")
	ipvsKey := s.key("ipvs-ids", "services", svc.IpvsId())
	// Validating service
	if err := s.validateService(svc); err != nil {
		return err
	}
	if err := s.validateServiceNameUniqueness(svcKey); err != nil {
		return err
	}
	if err := s.validateServiceIpvsUniqueness(ipvsKey); err != nil {
		return err
	}

	// TODO: Make the persistence of service and ipvs id atomic.
	// Pesisting service
	value, err := json.Marshal(svc)
	if err != nil {
		return errors.Wrapf(err, "error marshaling service: %v", svc)
	}
	err = s.kv.Put(svcKey, value, nil)
	if err != nil {
		return errors.Wrapf(err, "error sending service to store: %v", svc)
	}

	// Persisting IPVS key. So it can be validated.
	err = s.kv.Put(ipvsKey, []byte("true"), nil)
	if err != nil {
		return errors.Wrapf(err, "error sending service ipvs id to store: %v", svc.IpvsId())
	}

	return nil
}

func (s *FusisStore) GetService(serviceId string) (*types.Service, error) {
	key := s.key("services", serviceId, "config")
	kvPair, err := s.kv.Get(key)
	if err != nil {
		if err == kv.ErrKeyNotFound {
			return nil, types.ErrServiceNotFound
		}
		return nil, errors.Wrap(err, "GetService from store failed")
	}

	svc := &types.Service{}
	if err := json.Unmarshal(kvPair.Value, svc); err != nil {
		return nil, errors.Wrap(err, "Unmarshal service from store failed")
	}

	return svc, nil
}

func (s *FusisStore) DeleteService(svc *types.Service) error {
	key := s.key("services", svc.GetId())
	err := s.kv.DeleteTree(key)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete service: %v", svc)
	}

	ipvsKey := s.key("ipvs-ids", "services", svc.IpvsId())
	err = s.kv.Delete(ipvsKey)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete service ipvs id: %s", ipvsKey)
	}

	// TODO: Delete all dependent destinations and checks

	return nil
}

func (s *FusisStore) SubscribeServices(updateCh chan []types.Service) {
	s.Lock()
	defer s.Unlock()
	s.servicesChannels = append(s.servicesChannels, updateCh)
}

func (s *FusisStore) WatchServices() {
	svcs := []types.Service{}

	stopCh := make(<-chan struct{})
	events, err := s.kv.WatchTree(s.key("services"), stopCh)
	if err != nil {
		log.Error("[store] ", err)
	}

	for {
		select {
		case entries := <-events:
			log.Debug("[store] Services received")

			for _, pair := range entries {
				svc := types.Service{}
				if err := json.Unmarshal(pair.Value, &svc); err != nil {
					log.Error("[store] ", err)
				}
				log.Debugf("[store] Got sevice: %#v ", svc)

				svcs = append(svcs, svc)
			}

			s.Lock()
			for _, ch := range s.servicesChannels {
				log.Debug("[store] Sending message to service ch")
				ch <- svcs
			}
			s.Unlock()

			//Cleaning up services slice
			svcs = []types.Service{}
		}
	}
}

func (s *FusisStore) AddDestination(svc *types.Service, dst *types.Destination) error {
	dstKey := s.key("destinations", svc.GetId(), dst.GetId())
	ipvsKey := s.key("ipvs-ids", "destinations", dst.IpvsId())

	// Validating destination
	if err := s.validateDestination(dst); err != nil {
		return err
	}
	if err := s.validateDestinationNameUniqueness(dstKey); err != nil {
		return err
	}
	if err := s.validateDestinationIpvsUniqueness(ipvsKey); err != nil {
		return err
	}

	// Persisting destination
	value, err := json.Marshal(dst)
	if err != nil {
		return errors.Wrapf(err, "error marshaling destination: %v", dst)
	}
	err = s.kv.Put(dstKey, value, nil)
	if err != nil {
		return errors.Wrapf(err, "error sending destination to store: %v", dst)
	}
	log.Debugf("[store] Added destination: %s with key: %s", value, dstKey)

	// Persisting IPVS key. So it can be validated.
	err = s.kv.Put(ipvsKey, []byte("true"), nil)
	if err != nil {
		return errors.Wrapf(err, "error sending destination ipvs id to store: %v", dst.IpvsId())
	}

	return nil
}

func (s *FusisStore) DeleteDestination(svc *types.Service, dst *types.Destination) error {
	key := s.key("destinations", svc.GetId(), dst.GetId())

	err := s.kv.DeleteTree(key)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete destination: %v", dst)
	}
	log.Debugf("[store] Deleted destination: %s", key)

	ipvsKey := s.key("ipvs-ids", "destinations", dst.IpvsId())
	err = s.kv.Delete(ipvsKey)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete destination ipvs id: %s", ipvsKey)
	}

	return nil
}

func (s *FusisStore) SubscribeDestinations(updateCh chan []types.Destination) {
	s.Lock()
	defer s.Unlock()
	s.destinationChannels = append(s.destinationChannels, updateCh)
}

func (s *FusisStore) WatchDestinations() {
	dsts := []types.Destination{}

	stopCh := make(<-chan struct{})
	events, err := s.kv.WatchTree(s.key("destinations"), stopCh)
	if err != nil {
		errors.Wrap(err, "failed watching fusis/destinations")
	}

	for {
		select {
		case entries := <-events:
			log.Debug("[store] Destinations received")

			for _, pair := range entries {
				dst := types.Destination{}
				if err := json.Unmarshal(pair.Value, &dst); err != nil {
					errors.Wrap(err, "failed unmarshall of destinations")
				}
				log.Debugf("[store] Got destination: %#v", dst)

				dsts = append(dsts, dst)
			}

			s.Lock()
			for _, ch := range s.destinationChannels {
				ch <- dsts
			}
			s.Unlock()

			//Cleaning up destinations slice
			dsts = []types.Destination{}
		}
	}
}

func (s *FusisStore) SubscribeChecks(updateCh chan []types.CheckSpec) {
	s.Lock()
	defer s.Unlock()
	s.checksChannels = append(s.checksChannels, updateCh)
}

func (s *FusisStore) AddCheck(spec types.CheckSpec) error {
	key := s.key("checks", spec.ServiceID)

	value, err := json.Marshal(spec)
	if err != nil {
		return errors.Wrapf(err, "error marshaling CheckSpec: %#v", spec)
	}

	err = s.kv.Put(key, value, nil)
	if err != nil {
		return errors.Wrapf(err, "error sending CheckSpec to store: %#v", spec)
	}

	return nil
}

func (s *FusisStore) DeleteCheck(spec types.CheckSpec) error {
	key := s.key("checks", spec.ServiceID)

	err := s.kv.DeleteTree(key)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete check: %#v", spec)
	}
	log.Debugf("[store] Deleted check: %s", key)

	return nil
}

func (s *FusisStore) WatchChecks() {
	specs := []types.CheckSpec{}

	stopCh := make(<-chan struct{})
	events, err := s.kv.WatchTree(s.key("checks"), stopCh)
	if err != nil {
		log.Error(errors.Wrap(err, "failed watching fusis/checks"))
	}

	for {
		select {
		case entries := <-events:
			log.Debug("[store] Checks received")

			for _, pair := range entries {
				spec := types.CheckSpec{}
				if err := json.Unmarshal(pair.Value, &spec); err != nil {
					log.Error(errors.Wrap(err, "failed unmarshall of checks"))
				}
				log.Debugf("[store] Got Check: %#v", spec)

				specs = append(specs, spec)
			}

			s.Lock()
			for _, ch := range s.checksChannels {
				ch <- specs
			}
			s.Unlock()

			//Cleaning up destinations slice
			specs = []types.CheckSpec{}
		}
	}
}

func (s *FusisStore) key(keys ...string) string {
	return strings.Join(append([]string{s.prefix}, keys...), "/")
}
<commit_msg>Deleting dependent destinations and checks for services<commit_after>package store

import (
	"encoding/json"
	"net"
	"net/url"
	"strings"
	"sync"

	validator "gopkg.in/go-playground/validator.v9"

	log "github.com/Sirupsen/logrus"
	"github.com/docker/libkv"
	kv "github.com/docker/libkv/store"
	"github.com/docker/libkv/store/consul"
	"github.com/docker/libkv/store/etcd"
	"github.com/luizbafilho/fusis/config"
	"github.com/luizbafilho/fusis/types"
	"github.com/pkg/errors"
)

func init() {
	registryStores()
}

func registryStores() {
	libkv.AddStore(kv.CONSUL, consul.New)
	libkv.AddStore(kv.ETCD, etcd.New)
}

type Store interface {
	GetServices() ([]types.Service, error)
	// GetService(serviceId string) (types.Service, error)
	AddService(svc *types.Service) error
	DeleteService(svc *types.Service) error
	SubscribeServices(ch chan []types.Service)
	WatchServices()

	GetDestinations() ([]types.Destination, error)
	AddDestination(svc *types.Service, dst *types.Destination) error
	DeleteDestination(svc *types.Service, dst *types.Destination) error
	SubscribeDestinations(ch chan []types.Destination)
	WatchDestinations()

	AddCheck(check types.CheckSpec) error
	DeleteCheck(check types.CheckSpec) error
	SubscribeChecks(ch chan []types.CheckSpec)
	WatchChecks()

	GetKV() kv.Store
}

var (
	ErrUnsupportedStore = errors.New("[store] Unsupported store")
)

type FusisStore struct {
	sync.Mutex
	kv kv.Store

	prefix string

	validate            *validator.Validate
	servicesChannels    []chan []types.Service
	destinationChannels []chan []types.Destination
	checksChannels      []chan []types.CheckSpec
}

func New(config *config.BalancerConfig) (Store, error) {
	u, err := url.Parse(config.StoreAddress)
	if err != nil {
		return nil, errors.Wrap(err, "error paring store address")
	}

	scheme := u.Scheme
	if scheme != "consul" && scheme != "etcd" {
		return nil, ErrUnsupportedStore
	}

	//Validating open connection
	_, err = net.Dial("tcp", u.Host)
	if err != nil {
		return nil, errors.Wrap(err, "Store connection failed. Make sure your store is up and running.")
	}

	kv, err := libkv.NewStore(
		kv.Backend(scheme),
		[]string{u.Host},
		nil,
	)
	if err != nil {
		kv.Close()
		return nil, errors.Wrap(err, "Cannot create store consul")
	}

	svcsChs := []chan []types.Service{}
	dstsChs := []chan []types.Destination{}
	checksChs := []chan []types.CheckSpec{}

	validate := validator.New()
	// Registering custom validations
	validate.RegisterValidation("protocols", validateValues(types.Protocols))
	validate.RegisterValidation("schedulers", validateValues(types.Schedulers))

	fusisStore := &FusisStore{
		kv:                  kv,
		prefix:              config.StorePrefix,
		validate:            validate,
		servicesChannels:    svcsChs,
		destinationChannels: dstsChs,
		checksChannels:      checksChs,
	}

	go fusisStore.WatchServices()
	go fusisStore.WatchDestinations()
	go fusisStore.WatchChecks()

	return fusisStore, nil
}

func (s *FusisStore) GetKV() kv.Store {
	return s.kv
}

func (s *FusisStore) GetServices() ([]types.Service, error) {
	svcs := []types.Service{}
	entries, err := s.kv.List(s.key("services"))
	if err != nil {
		if err == kv.ErrKeyNotFound {
			return svcs, nil
		}

		return nil, err
	}

	for _, pair := range entries {
		svc := types.Service{}
		if err := json.Unmarshal(pair.Value, &svc); err != nil {
			log.Error("[store] ", err)
		}

		svcs = append(svcs, svc)
	}

	return svcs, nil
}

func (s *FusisStore) GetDestinations() ([]types.Destination, error) {
	dsts := []types.Destination{}
	entries, err := s.kv.List("fusis/destinations")
	if err != nil {
		if err == kv.ErrKeyNotFound {
			return dsts, nil
		}
		return nil, err
	}

	for _, pair := range entries {
		dst := types.Destination{}
		if err := json.Unmarshal(pair.Value, &dst); err != nil {
			log.Error("[store] ", err)
		}

		dsts = append(dsts, dst)
	}

	return dsts, nil
}

// AddService adds a new services to store. It validates the name
// uniqueness and the IPVS uniqueness by saving the IPVS key
// in the store, which consists in a combination of address, port
// and protocol.
func (s *FusisStore) AddService(svc *types.Service) error {
	svcKey := s.key("services", svc.GetId(), "config")
	ipvsKey := s.key("ipvs-ids", "services", svc.IpvsId())
	// Validating service
	if err := s.validateService(svc); err != nil {
		return err
	}
	if err := s.validateServiceNameUniqueness(svcKey); err != nil {
		return err
	}
	if err := s.validateServiceIpvsUniqueness(ipvsKey); err != nil {
		return err
	}

	// TODO: Make the persistence of service and ipvs id atomic.
	// Pesisting service
	value, err := json.Marshal(svc)
	if err != nil {
		return errors.Wrapf(err, "error marshaling service: %v", svc)
	}
	err = s.kv.Put(svcKey, value, nil)
	if err != nil {
		return errors.Wrapf(err, "error sending service to store: %v", svc)
	}

	// Persisting IPVS key. So it can be validated.
	err = s.kv.Put(ipvsKey, []byte("true"), nil)
	if err != nil {
		return errors.Wrapf(err, "error sending service ipvs id to store: %v", svc.IpvsId())
	}

	return nil
}

func (s *FusisStore) GetService(serviceId string) (*types.Service, error) {
	key := s.key("services", serviceId, "config")
	kvPair, err := s.kv.Get(key)
	if err != nil {
		if err == kv.ErrKeyNotFound {
			return nil, types.ErrServiceNotFound
		}
		return nil, errors.Wrap(err, "GetService from store failed")
	}

	svc := &types.Service{}
	if err := json.Unmarshal(kvPair.Value, svc); err != nil {
		return nil, errors.Wrap(err, "Unmarshal service from store failed")
	}

	return svc, nil
}

func (s *FusisStore) DeleteService(svc *types.Service) error {
	// TODO: Make all these actions atomic

	key := s.key("services", svc.GetId())
	err := s.kv.DeleteTree(key)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete service: %v", svc)
	}

	ipvsKey := s.key("ipvs-ids", "services", svc.IpvsId())
	err = s.kv.Delete(ipvsKey)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete service ipvs id: %s", ipvsKey)
	}

	// Deleting dependent destionations
	dsts, err := s.GetDestinations()
	if err != nil {
		return errors.Wrapf(err, "get destinations from deleted service failed: %#v", svc)
	}

	for _, dst := range dsts {
		if dst.ServiceId == svc.GetId() {
			if err := s.DeleteDestination(svc, &dst); err != nil {
				return errors.Wrap(err, "deleting dependent destination failed")
			}
		}
	}

	// Deleting checks
	exists, err := s.kv.Exists(s.key("checks", svc.GetId()))
	if err != nil {
		return errors.Wrapf(err, "verifying health check to service: %s failed", svc.GetId())
	}

	if exists {
		if err := s.DeleteCheck(types.CheckSpec{ServiceID: svc.GetId()}); err != nil {
			return errors.Wrap(err, "deleting dependent check failed")
		}
	}

	return nil
}

func (s *FusisStore) SubscribeServices(updateCh chan []types.Service) {
	s.Lock()
	defer s.Unlock()
	s.servicesChannels = append(s.servicesChannels, updateCh)
}

func (s *FusisStore) WatchServices() {
	svcs := []types.Service{}

	stopCh := make(<-chan struct{})
	events, err := s.kv.WatchTree(s.key("services"), stopCh)
	if err != nil {
		log.Error("[store] ", err)
	}

	for {
		select {
		case entries := <-events:
			log.Debug("[store] Services received")

			for _, pair := range entries {
				svc := types.Service{}
				if err := json.Unmarshal(pair.Value, &svc); err != nil {
					log.Error("[store] ", err)
				}
				log.Debugf("[store] Got sevice: %#v ", svc)

				svcs = append(svcs, svc)
			}

			s.Lock()
			for _, ch := range s.servicesChannels {
				log.Debug("[store] Sending message to service ch")
				ch <- svcs
			}
			s.Unlock()

			//Cleaning up services slice
			svcs = []types.Service{}
		}
	}
}

func (s *FusisStore) AddDestination(svc *types.Service, dst *types.Destination) error {
	dstKey := s.key("destinations", svc.GetId(), dst.GetId())
	ipvsKey := s.key("ipvs-ids", "destinations", dst.IpvsId())

	// Validating destination
	if err := s.validateDestination(dst); err != nil {
		return err
	}
	if err := s.validateDestinationNameUniqueness(dstKey); err != nil {
		return err
	}
	if err := s.validateDestinationIpvsUniqueness(ipvsKey); err != nil {
		return err
	}

	// Persisting destination
	value, err := json.Marshal(dst)
	if err != nil {
		return errors.Wrapf(err, "error marshaling destination: %v", dst)
	}
	err = s.kv.Put(dstKey, value, nil)
	if err != nil {
		return errors.Wrapf(err, "error sending destination to store: %v", dst)
	}
	log.Debugf("[store] Added destination: %s with key: %s", value, dstKey)

	// Persisting IPVS key. So it can be validated.
	err = s.kv.Put(ipvsKey, []byte("true"), nil)
	if err != nil {
		return errors.Wrapf(err, "error sending destination ipvs id to store: %v", dst.IpvsId())
	}

	return nil
}

func (s *FusisStore) DeleteDestination(svc *types.Service, dst *types.Destination) error {
	key := s.key("destinations", svc.GetId(), dst.GetId())

	err := s.kv.DeleteTree(key)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete destination: %v", dst)
	}
	log.Debugf("[store] Deleted destination: %s", key)

	ipvsKey := s.key("ipvs-ids", "destinations", dst.IpvsId())
	err = s.kv.Delete(ipvsKey)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete destination ipvs id: %s", ipvsKey)
	}

	return nil
}

func (s *FusisStore) SubscribeDestinations(updateCh chan []types.Destination) {
	s.Lock()
	defer s.Unlock()
	s.destinationChannels = append(s.destinationChannels, updateCh)
}

func (s *FusisStore) WatchDestinations() {
	dsts := []types.Destination{}

	stopCh := make(<-chan struct{})
	events, err := s.kv.WatchTree(s.key("destinations"), stopCh)
	if err != nil {
		errors.Wrap(err, "failed watching fusis/destinations")
	}

	for {
		select {
		case entries := <-events:
			log.Debug("[store] Destinations received")

			for _, pair := range entries {
				dst := types.Destination{}
				if err := json.Unmarshal(pair.Value, &dst); err != nil {
					errors.Wrap(err, "failed unmarshall of destinations")
				}
				log.Debugf("[store] Got destination: %#v", dst)

				dsts = append(dsts, dst)
			}

			s.Lock()
			for _, ch := range s.destinationChannels {
				ch <- dsts
			}
			s.Unlock()

			//Cleaning up destinations slice
			dsts = []types.Destination{}
		}
	}
}

func (s *FusisStore) SubscribeChecks(updateCh chan []types.CheckSpec) {
	s.Lock()
	defer s.Unlock()
	s.checksChannels = append(s.checksChannels, updateCh)
}

func (s *FusisStore) AddCheck(spec types.CheckSpec) error {
	key := s.key("checks", spec.ServiceID)

	value, err := json.Marshal(spec)
	if err != nil {
		return errors.Wrapf(err, "error marshaling CheckSpec: %#v", spec)
	}

	err = s.kv.Put(key, value, nil)
	if err != nil {
		return errors.Wrapf(err, "error sending CheckSpec to store: %#v", spec)
	}

	return nil
}

func (s *FusisStore) DeleteCheck(spec types.CheckSpec) error {
	key := s.key("checks", spec.ServiceID)

	err := s.kv.DeleteTree(key)
	if err != nil {
		return errors.Wrapf(err, "error trying to delete check: %#v", spec)
	}
	log.Debugf("[store] Deleted check: %s", key)

	return nil
}

func (s *FusisStore) WatchChecks() {
	specs := []types.CheckSpec{}

	stopCh := make(<-chan struct{})
	events, err := s.kv.WatchTree(s.key("checks"), stopCh)
	if err != nil {
		log.Error(errors.Wrap(err, "failed watching fusis/checks"))
	}

	for {
		select {
		case entries := <-events:
			log.Debug("[store] Checks received")

			for _, pair := range entries {
				spec := types.CheckSpec{}
				if err := json.Unmarshal(pair.Value, &spec); err != nil {
					log.Error(errors.Wrap(err, "failed unmarshall of checks"))
				}
				log.Debugf("[store] Got Check: %#v", spec)

				specs = append(specs, spec)
			}

			s.Lock()
			for _, ch := range s.checksChannels {
				ch <- specs
			}
			s.Unlock()

			//Cleaning up destinations slice
			specs = []types.CheckSpec{}
		}
	}
}

func (s *FusisStore) key(keys ...string) string {
	return strings.Join(append([]string{s.prefix}, keys...), "/")
}
<|endoftext|>