<commit_before>package handlers

import (
	"log"
	"net/http"

	"hawx.me/code/riviera-admin/views"
	"hawx.me/code/riviera/subscriptions/opml"
)

func List(opmlPath, url, pathPrefix string) http.Handler {
	return &listHandler{opmlPath, url, pathPrefix}
}

type listHandler struct {
	opmlPath   string
	url        string
	pathPrefix string
}

func (h *listHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	outline, err := opml.Load(h.opmlPath)
	if err != nil {
		log.Println("list", err)
		w.WriteHeader(500)
		return
	}

	var list []feed
	for _, line := range outline.Body.Outline {
		if line.Type == "rss" {
			f := feed{
				FeedUrl:         line.XmlUrl,
				WebsiteUrl:      line.HtmlUrl,
				FeedTitle:       line.Title,
				FeedDescription: line.Description,
			}

			if f.FeedTitle == "" {
				f.FeedTitle = f.FeedUrl
			}

			list = append(list, f)
		}
	}

	w.Header().Add("Content-Type", "text/html")

	views.Index.Execute(w, struct {
		Url        string
		PathPrefix string
		Feeds      []feed
	}{h.url, h.pathPrefix, list})
}

type feed struct {
	FeedUrl         string
	WebsiteUrl      string
	FeedTitle       string
	FeedDescription string
}
<commit_msg>Sort by feed title<commit_after>package handlers

import (
	"log"
	"net/http"
	"sort"
	"strings"

	"hawx.me/code/riviera-admin/views"
	"hawx.me/code/riviera/subscriptions/opml"
)

func List(opmlPath, url, pathPrefix string) http.Handler {
	return &listHandler{opmlPath, url, pathPrefix}
}

type listHandler struct {
	opmlPath   string
	url        string
	pathPrefix string
}

func (h *listHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	outline, err := opml.Load(h.opmlPath)
	if err != nil {
		log.Println("list", err)
		w.WriteHeader(500)
		return
	}

	var list feeds
	for _, line := range outline.Body.Outline {
		if line.Type == "rss" {
			f := feed{
				FeedUrl:         line.XmlUrl,
				WebsiteUrl:      line.HtmlUrl,
				FeedTitle:       line.Title,
				FeedDescription: line.Description,
			}

			if f.FeedTitle == "" {
				f.FeedTitle = f.FeedUrl
			}

			list = append(list, f)
		}
	}

	sort.Sort(list)

	w.Header().Add("Content-Type", "text/html")

	views.Index.Execute(w, struct {
		Url        string
		PathPrefix string
		Feeds      []feed
	}{h.url, h.pathPrefix, list})
}

type feed struct {
	FeedUrl         string
	WebsiteUrl      string
	FeedTitle       string
	FeedDescription string
}

type feeds []feed

func (fs feeds) Len() int      { return len(fs) }
func (fs feeds) Swap(i, j int) { fs[i], fs[j] = fs[j], fs[i] }
func (fs feeds) Less(i, j int) bool {
	return strings.ToLower(fs[i].FeedTitle) < strings.ToLower(fs[j].FeedTitle)
}
<|endoftext|>


<commit_before>// +build go1.8

package handlers

import (
	"fmt"
	"net/http"
)

type loggingResponseWriter interface {
	commonLoggingResponseWriter
	http.Pusher
}

func (l *responseLogger) Push(target string, opts *http.PushOptions) error {
	p, ok := l.w.(http.Pusher)
	if !ok {
		return fmt.Errorf("responseLogger does not implement http.Pusher")
	}
	return p.Push(target, opts)
}
<commit_msg>Implement http.Pusher interface for compressResponseWriter (Go >= 1.8) (#138)<commit_after>// +build go1.8

package handlers

import (
	"fmt"
	"net/http"
)

type loggingResponseWriter interface {
	commonLoggingResponseWriter
	http.Pusher
}

func (l *responseLogger) Push(target string, opts *http.PushOptions) error {
	p, ok := l.w.(http.Pusher)
	if !ok {
		return fmt.Errorf("responseLogger does not implement http.Pusher")
	}
	return p.Push(target, opts)
}

func (c *compressResponseWriter) Push(target string, opts *http.PushOptions) error {
	p, ok := c.ResponseWriter.(http.Pusher)
	if !ok {
		return fmt.Errorf("compressResponseWriter does not implement http.Pusher")
	}
	return p.Push(target, opts)
}
<|endoftext|>


<commit_before>// Copyright 2019 ETH Zurich
// Copyright 2020 ETH Zurich, Anapaya Systems
//
// 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 json encodes AS topology information via JSON. All types exposed by this package are
// designed to be directly marshaled to / unmarshaled from JSON.
package json

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/scionproto/scion/go/lib/common"
	"github.com/scionproto/scion/go/lib/serrors"
	"github.com/scionproto/scion/go/lib/util"
)

// Attribute indicates the capability of a primary AS.
type Attribute string

const (
	// Authoritative indicates an authoritative AS.
	Authoritative Attribute = "authoritative"
	// AttrCore indicates a core AS.
	AttrCore Attribute = "core"
	// Issuing indicates an issuing AS.
	Issuing Attribute = "issuing"
	// Voting indicates a voting AS. A voting AS must also be a core AS.
	Voting Attribute = "voting"
)

// UnmarshalText checks that the attribute is valid. It can either be
// "authoritative", "core", "issuing", or "voting".
func (t *Attribute) UnmarshalText(b []byte) error {
	switch Attribute(b) {
	case Authoritative:
		*t = Authoritative
	case Issuing:
		*t = Issuing
	case Voting:
		*t = Voting
	case AttrCore:
		*t = AttrCore
	default:
		return serrors.New("invalid attribute", "input", string(b))
	}
	return nil
}

// Topology is the JSON type for the entire AS topology file.
type Topology struct {
	Timestamp      int64  `json:"timestamp,omitempty"`
	TimestampHuman string `json:"timestamp_human,omitempty"`
	IA             string `json:"isd_as"`
	MTU            int    `json:"mtu"`
	// Attributes are the primary AS attributes as described in
	// https://github.com/scionproto/scion/blob/master/doc/ControlPlanePKI.md#primary-ases
	Attributes          []Attribute             `json:"attributes"`
	BorderRouters       map[string]*BRInfo      `json:"border_routers,omitempty"`
	ControlService      map[string]*ServerInfo  `json:"control_service,omitempty"`
	DiscoveryService    map[string]*ServerInfo  `json:"discovery_service,omitempty"`
	HiddenSegmentLookup map[string]*ServerInfo  `json:"hidden_segment_lookup_service,omitempty"`
	HiddenSegmentReg    map[string]*ServerInfo  `json:"hidden_segment_registration_service,omitempty"`
	SIG                 map[string]*GatewayInfo `json:"sigs,omitempty"`
}

// ServerInfo contains the information for a SCION application running in the local AS.
type ServerInfo struct {
	Addr string `json:"addr"`
}

// BRInfo contains Border Router specific information.
type BRInfo struct {
	InternalAddr string                           `json:"internal_addr"`
	CtrlAddr     string                           `json:"ctrl_addr"`
	Interfaces   map[common.IFIDType]*BRInterface `json:"interfaces"`
}

// GatewayInfo contains SCION gateway information.
type GatewayInfo struct {
	CtrlAddr   string   `json:"ctrl_addr"`
	DataAddr   string   `json:"data_addr"`
	ProbeAddr  string   `json:"probe_addr,omitempty"`
	Interfaces []uint64 `json:"allow_interfaces,omitempty"`
}

// BRInterface contains the information for an data-plane BR socket that is external (i.e., facing
// the neighboring AS).
type BRInterface struct {
	Underlay  Underlay `json:"underlay,omitempty"`
	Bandwidth int      `json:"bandwidth"`
	IA        string   `json:"isd_as"`
	LinkTo    string   `json:"link_to"`
	MTU       int      `json:"mtu"`
	BFD       *BFD     `json:"bfd,omitempty"`
}

// Underlay is the underlay information for a BR interface.
type Underlay struct {
	Public string `json:"public"`
	Remote string `json:"remote"`
	Bind   string `json:"bind,omitempty"`
}

// BFD configuration.
type BFD struct {
	Disable               bool         `json:"disable,omitempty"`
	DetectMult            uint8        `json:"detect_mult,omitempty"`
	DesiredMinTxInterval  util.DurWrap `json:"desired_min_tx_interval,omitempty"`
	RequiredMinRxInterval util.DurWrap `json:"required_min_rx_interval,omitempty"`
}

func (i ServerInfo) String() string {
	return fmt.Sprintf("Addr: %s", i.Addr)
}

func (i BRInfo) String() string {
	var s []string
	s = append(s, fmt.Sprintf("Loc addrs:\n  %s\nControl addr:\n  %s\nInterfaces:",
		i.InternalAddr, i.CtrlAddr))
	for ifid, intf := range i.Interfaces {
		s = append(s, fmt.Sprintf("%d: %+v", ifid, intf))
	}
	return strings.Join(s, "\n")
}

// Load parses a topology from its raw byte representation.
func Load(b []byte) (*Topology, error) {
	rt := &Topology{}
	if err := json.Unmarshal(b, rt); err != nil {
		return nil, serrors.WrapStr("unable to parse topology from JSON", err)
	}
	return rt, nil
}

// LoadFromFile parses a topology from a file.
func LoadFromFile(path string) (*Topology, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, serrors.WrapStr("unable to open topology", err, "path", path)
	}
	return Load(b)
}
<commit_msg>topology: add more omitempty tags<commit_after>// Copyright 2019 ETH Zurich
// Copyright 2020 ETH Zurich, Anapaya Systems
//
// 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 json encodes AS topology information via JSON. All types exposed by this package are
// designed to be directly marshaled to / unmarshaled from JSON.
package json

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"strings"

	"github.com/scionproto/scion/go/lib/common"
	"github.com/scionproto/scion/go/lib/serrors"
	"github.com/scionproto/scion/go/lib/util"
)

// Attribute indicates the capability of a primary AS.
type Attribute string

const (
	// Authoritative indicates an authoritative AS.
	Authoritative Attribute = "authoritative"
	// AttrCore indicates a core AS.
	AttrCore Attribute = "core"
	// Issuing indicates an issuing AS.
	Issuing Attribute = "issuing"
	// Voting indicates a voting AS. A voting AS must also be a core AS.
	Voting Attribute = "voting"
)

// UnmarshalText checks that the attribute is valid. It can either be
// "authoritative", "core", "issuing", or "voting".
func (t *Attribute) UnmarshalText(b []byte) error {
	switch Attribute(b) {
	case Authoritative:
		*t = Authoritative
	case Issuing:
		*t = Issuing
	case Voting:
		*t = Voting
	case AttrCore:
		*t = AttrCore
	default:
		return serrors.New("invalid attribute", "input", string(b))
	}
	return nil
}

// Topology is the JSON type for the entire AS topology file.
type Topology struct {
	Timestamp      int64  `json:"timestamp,omitempty"`
	TimestampHuman string `json:"timestamp_human,omitempty"`
	IA             string `json:"isd_as"`
	MTU            int    `json:"mtu"`
	// Attributes are the primary AS attributes as described in
	// https://github.com/scionproto/scion/blob/master/doc/ControlPlanePKI.md#primary-ases
	Attributes          []Attribute             `json:"attributes"`
	BorderRouters       map[string]*BRInfo      `json:"border_routers,omitempty"`
	ControlService      map[string]*ServerInfo  `json:"control_service,omitempty"`
	DiscoveryService    map[string]*ServerInfo  `json:"discovery_service,omitempty"`
	HiddenSegmentLookup map[string]*ServerInfo  `json:"hidden_segment_lookup_service,omitempty"`
	HiddenSegmentReg    map[string]*ServerInfo  `json:"hidden_segment_registration_service,omitempty"`
	SIG                 map[string]*GatewayInfo `json:"sigs,omitempty"`
}

// ServerInfo contains the information for a SCION application running in the local AS.
type ServerInfo struct {
	Addr string `json:"addr"`
}

// BRInfo contains Border Router specific information.
type BRInfo struct {
	InternalAddr string                           `json:"internal_addr"`
	CtrlAddr     string                           `json:"ctrl_addr"`
	Interfaces   map[common.IFIDType]*BRInterface `json:"interfaces"`
}

// GatewayInfo contains SCION gateway information.
type GatewayInfo struct {
	CtrlAddr   string   `json:"ctrl_addr"`
	DataAddr   string   `json:"data_addr"`
	ProbeAddr  string   `json:"probe_addr,omitempty"`
	Interfaces []uint64 `json:"allow_interfaces,omitempty"`
}

// BRInterface contains the information for an data-plane BR socket that is external (i.e., facing
// the neighboring AS).
type BRInterface struct {
	Underlay  Underlay `json:"underlay,omitempty"`
	Bandwidth int      `json:"bandwidth,omitempty"`
	IA        string   `json:"isd_as"`
	LinkTo    string   `json:"link_to"`
	MTU       int      `json:"mtu"`
	BFD       *BFD     `json:"bfd,omitempty"`
}

// Underlay is the underlay information for a BR interface.
type Underlay struct {
	Public string `json:"public,omitempty"`
	Remote string `json:"remote,omitempty"`
	Bind   string `json:"bind,omitempty"`
}

// BFD configuration.
type BFD struct {
	Disable               bool         `json:"disable,omitempty"`
	DetectMult            uint8        `json:"detect_mult,omitempty"`
	DesiredMinTxInterval  util.DurWrap `json:"desired_min_tx_interval,omitempty"`
	RequiredMinRxInterval util.DurWrap `json:"required_min_rx_interval,omitempty"`
}

func (i ServerInfo) String() string {
	return fmt.Sprintf("Addr: %s", i.Addr)
}

func (i BRInfo) String() string {
	var s []string
	s = append(s, fmt.Sprintf("Loc addrs:\n  %s\nControl addr:\n  %s\nInterfaces:",
		i.InternalAddr, i.CtrlAddr))
	for ifid, intf := range i.Interfaces {
		s = append(s, fmt.Sprintf("%d: %+v", ifid, intf))
	}
	return strings.Join(s, "\n")
}

// Load parses a topology from its raw byte representation.
func Load(b []byte) (*Topology, error) {
	rt := &Topology{}
	if err := json.Unmarshal(b, rt); err != nil {
		return nil, serrors.WrapStr("unable to parse topology from JSON", err)
	}
	return rt, nil
}

// LoadFromFile parses a topology from a file.
func LoadFromFile(path string) (*Topology, error) {
	b, err := ioutil.ReadFile(path)
	if err != nil {
		return nil, serrors.WrapStr("unable to open topology", err, "path", path)
	}
	return Load(b)
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

// Package dbconfigs provides the registration for command line options
// to collect db connection parameters. Once registered and collected,
// it provides variables and functions to build connection parameters
// for connecting to the database.
package dbconfigs

import (
	"encoding/json"
	"flag"
	"fmt"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/log"
)

var (
	dbConfigs  = DBConfigs{userConfigs: make(map[string]*userConfig)}
	baseConfig = mysql.ConnParams{}
)

// DBConfigs stores all the data needed to build various connection
// parameters for the db. It stores credentials for app, appdebug,
// allprivs, dba, filtered and repl users.
// It contains other connection parameters like socket, charset, etc.
// It also stores the default db name, which it can combine with the
// rest of the data to build db-sepcific connection parameters.
// It also supplies the SidecarDBName. This is currently hardcoded
// to "_vt", but will soon become customizable.
// The life-cycle of this package is as follows:
// App must call RegisterFlags to request the types of connections
// it wants support for. This must be done before involing flags.Parse.
// After flag parsing, app invokes the Init function, which will return
// a DBConfigs object.
// The app must store the DBConfigs object internally, and use it to
// build connection parameters as needed.
// The DBName is initially empty and may later be set or changed by the app.
type DBConfigs struct {
	userConfigs   map[string]*userConfig
	DBName        sync2.AtomicString
	SidecarDBName sync2.AtomicString
}

type userConfig struct {
	useSSL bool
	param  mysql.ConnParams
}

// config flags
const (
	App      = "app"
	AppDebug = "appdebug"
	// AllPrivs user should have more privileges than App (should include possibility to do
	// schema changes and write to internal Vitess tables), but it shouldn't have SUPER
	// privilege like Dba has.
	AllPrivs = "allprivs"
	Dba      = "dba"
	Filtered = "filtered"
	Repl     = "repl"
)

// All can be used to register all flags: RegisterFlags(All...)
var All = []string{App, AppDebug, AllPrivs, Dba, Filtered, Repl}

// RegisterFlags registers the flags for the given DBConfigFlag.
// For instance, vttablet will register client, dba and repl.
// Returns all registered flags.
func RegisterFlags(userKeys ...string) {
	registerBaseFlags()
	for _, userKey := range userKeys {
		uc := &userConfig{}
		dbConfigs.userConfigs[userKey] = uc
		registerPerUserFlags(uc, userKey)
	}
}

func registerBaseFlags() {
	flag.StringVar(&baseConfig.UnixSocket, "db_socket", "", "The unix socket to connect on. If this is specifed, host and port will not be used.")
	flag.StringVar(&baseConfig.Host, "db_host", "", "The host name for the tcp connection.")
	flag.IntVar(&baseConfig.Port, "db_port", 0, "tcp port")
	flag.StringVar(&baseConfig.Charset, "db_charset", "utf8", "Character set. Only utf8 or latin1 based character sets are supported.")
	flag.Uint64Var(&baseConfig.Flags, "db_flags", 0, "Flag values as defined by MySQL.")
	flag.StringVar(&baseConfig.SslCa, "db_ssl_ca", "", "connection ssl ca")
	flag.StringVar(&baseConfig.SslCaPath, "db_ssl_ca_path", "", "connection ssl ca path")
	flag.StringVar(&baseConfig.SslCert, "db_ssl_cert", "", "connection ssl certificate")
	flag.StringVar(&baseConfig.SslKey, "db_ssl_key", "", "connection ssl key")
	flag.StringVar(&baseConfig.ServerName, "db_server_name", "", "server name of the DB we are connecting to.")

}

// The flags will change the global singleton
// TODO(sougou): deprecate the legacy flags.
func registerPerUserFlags(dbc *userConfig, userKey string) {
	newUserFlag := "db_" + userKey + "_user"
	flag.StringVar(&dbc.param.Uname, "db-config-"+userKey+"-uname", "vt_"+userKey, "deprecated: use "+newUserFlag)
	flag.StringVar(&dbc.param.Uname, newUserFlag, "vt_"+userKey, "db "+userKey+" user userKey")

	newPasswordFlag := "db_" + userKey + "_password"
	flag.StringVar(&dbc.param.Pass, "db-config-"+userKey+"-pass", "", "db "+userKey+" deprecated: use "+newPasswordFlag)
	flag.StringVar(&dbc.param.Pass, newPasswordFlag, "", "db "+userKey+" password")

	flag.BoolVar(&dbc.useSSL, "db_"+userKey+"_use_ssl", true, "Set this flag to false to make the "+userKey+" connection to not use ssl")

	flag.StringVar(&dbc.param.Host, "db-config-"+userKey+"-host", "", "deprecated: use db_host")
	flag.IntVar(&dbc.param.Port, "db-config-"+userKey+"-port", 0, "deprecated: use db_port")
	flag.StringVar(&dbc.param.UnixSocket, "db-config-"+userKey+"-unixsocket", "", "deprecated: use db_socket")
	flag.StringVar(&dbc.param.Charset, "db-config-"+userKey+"-charset", "utf8", "deprecated: use db_charset")
	flag.Uint64Var(&dbc.param.Flags, "db-config-"+userKey+"-flags", 0, "deprecated: use db_flags")
	flag.StringVar(&dbc.param.SslCa, "db-config-"+userKey+"-ssl-ca", "", "deprecated: use db_ssl_ca")
	flag.StringVar(&dbc.param.SslCaPath, "db-config-"+userKey+"-ssl-ca-path", "", "deprecated: use db_ssl_ca_path")
	flag.StringVar(&dbc.param.SslCert, "db-config-"+userKey+"-ssl-cert", "", "deprecated: use db_ssl_cert")
	flag.StringVar(&dbc.param.SslKey, "db-config-"+userKey+"-ssl-key", "", "deprecated: use db_ssl_key")
	flag.StringVar(&dbc.param.ServerName, "db-config-"+userKey+"-server_name", "", "deprecated: use db_server_name")

	flag.StringVar(&dbc.param.DeprecatedDBName, "db-config-"+userKey+"-dbname", "", "deprecated: dbname does not need to be explicitly configured")

}

// AppWithDB returns connection parameters for app with dbname set.
func (dbcfgs *DBConfigs) AppWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(App, true)
}

// AppDebugWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) AppDebugWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(AppDebug, true)
}

// AllPrivsWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) AllPrivsWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(AllPrivs, true)
}

// Dba returns connection parameters for dba with no dbname set.
func (dbcfgs *DBConfigs) Dba() *mysql.ConnParams {
	return dbcfgs.makeParams(Dba, false)
}

// DbaWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) DbaWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(Dba, true)
}

// FilteredWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) FilteredWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(Filtered, true)
}

// Repl returns connection parameters for appdebug with no dbname set.
func (dbcfgs *DBConfigs) Repl() *mysql.ConnParams {
	return dbcfgs.makeParams(Repl, false)
}

// AppWithDB returns connection parameters for app with dbname set.
func (dbcfgs *DBConfigs) makeParams(userKey string, withDB bool) *mysql.ConnParams {
	orig := dbcfgs.userConfigs[userKey]
	if orig == nil {
		return &mysql.ConnParams{}
	}
	result := orig.param
	if withDB {
		result.DbName = dbcfgs.DBName.Get()
	}
	return &result
}

// IsZero returns true if DBConfigs was uninitialized.
func (dbcfgs *DBConfigs) IsZero() bool {
	return len(dbcfgs.userConfigs) == 0
}

func (dbcfgs *DBConfigs) String() string {
	out := struct {
		Conn  mysql.ConnParams
		Users map[string]string
	}{
		Users: make(map[string]string),
	}
	if conn := dbcfgs.userConfigs[App]; conn != nil {
		out.Conn = conn.param
	} else if conn := dbcfgs.userConfigs[Dba]; conn != nil {
		out.Conn = conn.param
	}
	out.Conn.Pass = "****"
	for k, uc := range dbcfgs.userConfigs {
		out.Users[k] = uc.param.Uname
	}
	data, err := json.MarshalIndent(out, "", "  ")
	if err != nil {
		return err.Error()
	}
	return string(data)
}

// Copy returns a copy of the DBConfig.
func (dbcfgs *DBConfigs) Copy() *DBConfigs {
	result := &DBConfigs{userConfigs: make(map[string]*userConfig)}
	for k, u := range dbcfgs.userConfigs {
		newu := *u
		result.userConfigs[k] = &newu
	}
	result.DBName.Set(dbcfgs.DBName.Get())
	result.SidecarDBName.Set(dbcfgs.SidecarDBName.Get())
	return result
}

// HasConnectionParams returns true if connection parameters were
// specified in the command-line. This will allow the caller to
// search for alternate ways to connect, like looking in the my.cnf
// file.
func HasConnectionParams() bool {
	return baseConfig.Host != "" || baseConfig.UnixSocket != ""
}

// Init will initialize all the necessary connection parameters.
// Precedence is as follows: if baseConfig command line options are
// set, they supersede all other settings.
// If baseConfig is not set, the next priority is with per-user connection
// parameters. This is only for legacy support.
// If no per-user parameters are supplied, then the defaultSocketFile
// is used to initialize the per-user conn params.
func Init(defaultSocketFile string) (*DBConfigs, error) {
	// The new base configs, if set, supersede legacy settings.
	if HasConnectionParams() {
		for _, uc := range dbConfigs.userConfigs {
			uc.param.Host = baseConfig.Host
			uc.param.Port = baseConfig.Port
			uc.param.UnixSocket = baseConfig.UnixSocket
			uc.param.Charset = baseConfig.Charset
			uc.param.Flags = baseConfig.Flags
			if uc.useSSL {
				uc.param.SslCa = baseConfig.SslCa
				uc.param.SslCaPath = baseConfig.SslCaPath
				uc.param.SslCert = baseConfig.SslCert
				uc.param.SslKey = baseConfig.SslKey
				uc.param.ServerName = baseConfig.ServerName
			}
		}
	} else {
		// Use supplied socket value if conn parameters are not specified.
		for _, uc := range dbConfigs.userConfigs {
			if uc.param.UnixSocket == "" && uc.param.Host == "" {
				uc.param.UnixSocket = defaultSocketFile
			}
		}
	}

	// See if the CredentialsServer is working. We do not use the
	// result for anything, this is just a check.
	for _, uc := range dbConfigs.userConfigs {
		if _, err := WithCredentials(&uc.param); err != nil {
			return nil, fmt.Errorf("dbconfig cannot be initialized: %v", err)
		}
		// Check for only one.
		break
	}
	dbConfigs.SidecarDBName.Set("_vt")

	log.Infof("DBConfigs: %v\n", dbConfigs.String())
	return &dbConfigs, nil
}

// NewTestDBConfigs returns a DBConfigs meant for testing.
func NewTestDBConfigs(genParams, appDebugParams mysql.ConnParams, dbName string) *DBConfigs {
	dbcfgs := &DBConfigs{
		userConfigs: map[string]*userConfig{
			App:      {param: genParams},
			AppDebug: {param: appDebugParams},
			AllPrivs: {param: genParams},
			Dba:      {param: genParams},
			Filtered: {param: genParams},
			Repl:     {param: genParams},
		},
	}
	dbcfgs.DBName.Set(dbName)
	dbcfgs.SidecarDBName.Set("_vt")
	return dbcfgs
}
<commit_msg>Apply vttablet command line flags regardless of whether the tablet is running in managed or non-managed mode. Addresses https://github.com/vitessio/vitess/issues/4654<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

// Package dbconfigs provides the registration for command line options
// to collect db connection parameters. Once registered and collected,
// it provides variables and functions to build connection parameters
// for connecting to the database.
package dbconfigs

import (
	"encoding/json"
	"flag"
	"fmt"

	"vitess.io/vitess/go/mysql"
	"vitess.io/vitess/go/sync2"
	"vitess.io/vitess/go/vt/log"
)

var (
	dbConfigs  = DBConfigs{userConfigs: make(map[string]*userConfig)}
	baseConfig = mysql.ConnParams{}
)

// DBConfigs stores all the data needed to build various connection
// parameters for the db. It stores credentials for app, appdebug,
// allprivs, dba, filtered and repl users.
// It contains other connection parameters like socket, charset, etc.
// It also stores the default db name, which it can combine with the
// rest of the data to build db-sepcific connection parameters.
// It also supplies the SidecarDBName. This is currently hardcoded
// to "_vt", but will soon become customizable.
// The life-cycle of this package is as follows:
// App must call RegisterFlags to request the types of connections
// it wants support for. This must be done before involing flags.Parse.
// After flag parsing, app invokes the Init function, which will return
// a DBConfigs object.
// The app must store the DBConfigs object internally, and use it to
// build connection parameters as needed.
// The DBName is initially empty and may later be set or changed by the app.
type DBConfigs struct {
	userConfigs   map[string]*userConfig
	DBName        sync2.AtomicString
	SidecarDBName sync2.AtomicString
}

type userConfig struct {
	useSSL bool
	param  mysql.ConnParams
}

// config flags
const (
	App      = "app"
	AppDebug = "appdebug"
	// AllPrivs user should have more privileges than App (should include possibility to do
	// schema changes and write to internal Vitess tables), but it shouldn't have SUPER
	// privilege like Dba has.
	AllPrivs = "allprivs"
	Dba      = "dba"
	Filtered = "filtered"
	Repl     = "repl"
)

// All can be used to register all flags: RegisterFlags(All...)
var All = []string{App, AppDebug, AllPrivs, Dba, Filtered, Repl}

// RegisterFlags registers the flags for the given DBConfigFlag.
// For instance, vttablet will register client, dba and repl.
// Returns all registered flags.
func RegisterFlags(userKeys ...string) {
	registerBaseFlags()
	for _, userKey := range userKeys {
		uc := &userConfig{}
		dbConfigs.userConfigs[userKey] = uc
		registerPerUserFlags(uc, userKey)
	}
}

func registerBaseFlags() {
	flag.StringVar(&baseConfig.UnixSocket, "db_socket", "", "The unix socket to connect on. If this is specifed, host and port will not be used.")
	flag.StringVar(&baseConfig.Host, "db_host", "", "The host name for the tcp connection.")
	flag.IntVar(&baseConfig.Port, "db_port", 0, "tcp port")
	flag.StringVar(&baseConfig.Charset, "db_charset", "utf8", "Character set. Only utf8 or latin1 based character sets are supported.")
	flag.Uint64Var(&baseConfig.Flags, "db_flags", 0, "Flag values as defined by MySQL.")
	flag.StringVar(&baseConfig.SslCa, "db_ssl_ca", "", "connection ssl ca")
	flag.StringVar(&baseConfig.SslCaPath, "db_ssl_ca_path", "", "connection ssl ca path")
	flag.StringVar(&baseConfig.SslCert, "db_ssl_cert", "", "connection ssl certificate")
	flag.StringVar(&baseConfig.SslKey, "db_ssl_key", "", "connection ssl key")
	flag.StringVar(&baseConfig.ServerName, "db_server_name", "", "server name of the DB we are connecting to.")

}

// The flags will change the global singleton
// TODO(sougou): deprecate the legacy flags.
func registerPerUserFlags(dbc *userConfig, userKey string) {
	newUserFlag := "db_" + userKey + "_user"
	flag.StringVar(&dbc.param.Uname, "db-config-"+userKey+"-uname", "vt_"+userKey, "deprecated: use "+newUserFlag)
	flag.StringVar(&dbc.param.Uname, newUserFlag, "vt_"+userKey, "db "+userKey+" user userKey")

	newPasswordFlag := "db_" + userKey + "_password"
	flag.StringVar(&dbc.param.Pass, "db-config-"+userKey+"-pass", "", "db "+userKey+" deprecated: use "+newPasswordFlag)
	flag.StringVar(&dbc.param.Pass, newPasswordFlag, "", "db "+userKey+" password")

	flag.BoolVar(&dbc.useSSL, "db_"+userKey+"_use_ssl", true, "Set this flag to false to make the "+userKey+" connection to not use ssl")

	flag.StringVar(&dbc.param.Host, "db-config-"+userKey+"-host", "", "deprecated: use db_host")
	flag.IntVar(&dbc.param.Port, "db-config-"+userKey+"-port", 0, "deprecated: use db_port")
	flag.StringVar(&dbc.param.UnixSocket, "db-config-"+userKey+"-unixsocket", "", "deprecated: use db_socket")
	flag.StringVar(&dbc.param.Charset, "db-config-"+userKey+"-charset", "utf8", "deprecated: use db_charset")
	flag.Uint64Var(&dbc.param.Flags, "db-config-"+userKey+"-flags", 0, "deprecated: use db_flags")
	flag.StringVar(&dbc.param.SslCa, "db-config-"+userKey+"-ssl-ca", "", "deprecated: use db_ssl_ca")
	flag.StringVar(&dbc.param.SslCaPath, "db-config-"+userKey+"-ssl-ca-path", "", "deprecated: use db_ssl_ca_path")
	flag.StringVar(&dbc.param.SslCert, "db-config-"+userKey+"-ssl-cert", "", "deprecated: use db_ssl_cert")
	flag.StringVar(&dbc.param.SslKey, "db-config-"+userKey+"-ssl-key", "", "deprecated: use db_ssl_key")
	flag.StringVar(&dbc.param.ServerName, "db-config-"+userKey+"-server_name", "", "deprecated: use db_server_name")

	flag.StringVar(&dbc.param.DeprecatedDBName, "db-config-"+userKey+"-dbname", "", "deprecated: dbname does not need to be explicitly configured")

}

// AppWithDB returns connection parameters for app with dbname set.
func (dbcfgs *DBConfigs) AppWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(App, true)
}

// AppDebugWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) AppDebugWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(AppDebug, true)
}

// AllPrivsWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) AllPrivsWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(AllPrivs, true)
}

// Dba returns connection parameters for dba with no dbname set.
func (dbcfgs *DBConfigs) Dba() *mysql.ConnParams {
	return dbcfgs.makeParams(Dba, false)
}

// DbaWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) DbaWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(Dba, true)
}

// FilteredWithDB returns connection parameters for appdebug with dbname set.
func (dbcfgs *DBConfigs) FilteredWithDB() *mysql.ConnParams {
	return dbcfgs.makeParams(Filtered, true)
}

// Repl returns connection parameters for appdebug with no dbname set.
func (dbcfgs *DBConfigs) Repl() *mysql.ConnParams {
	return dbcfgs.makeParams(Repl, false)
}

// AppWithDB returns connection parameters for app with dbname set.
func (dbcfgs *DBConfigs) makeParams(userKey string, withDB bool) *mysql.ConnParams {
	orig := dbcfgs.userConfigs[userKey]
	if orig == nil {
		return &mysql.ConnParams{}
	}
	result := orig.param
	if withDB {
		result.DbName = dbcfgs.DBName.Get()
	}
	return &result
}

// IsZero returns true if DBConfigs was uninitialized.
func (dbcfgs *DBConfigs) IsZero() bool {
	return len(dbcfgs.userConfigs) == 0
}

func (dbcfgs *DBConfigs) String() string {
	out := struct {
		Conn  mysql.ConnParams
		Users map[string]string
	}{
		Users: make(map[string]string),
	}
	if conn := dbcfgs.userConfigs[App]; conn != nil {
		out.Conn = conn.param
	} else if conn := dbcfgs.userConfigs[Dba]; conn != nil {
		out.Conn = conn.param
	}
	out.Conn.Pass = "****"
	for k, uc := range dbcfgs.userConfigs {
		out.Users[k] = uc.param.Uname
	}
	data, err := json.MarshalIndent(out, "", "  ")
	if err != nil {
		return err.Error()
	}
	return string(data)
}

// Copy returns a copy of the DBConfig.
func (dbcfgs *DBConfigs) Copy() *DBConfigs {
	result := &DBConfigs{userConfigs: make(map[string]*userConfig)}
	for k, u := range dbcfgs.userConfigs {
		newu := *u
		result.userConfigs[k] = &newu
	}
	result.DBName.Set(dbcfgs.DBName.Get())
	result.SidecarDBName.Set(dbcfgs.SidecarDBName.Get())
	return result
}

// HasConnectionParams returns true if connection parameters were
// specified in the command-line. This will allow the caller to
// search for alternate ways to connect, like looking in the my.cnf
// file.
func HasConnectionParams() bool {
	return baseConfig.Host != "" || baseConfig.UnixSocket != ""
}

// Init will initialize all the necessary connection parameters.
// Precedence is as follows: if baseConfig command line options are
// set, they supersede all other settings.
// If baseConfig is not set, the next priority is with per-user connection
// parameters. This is only for legacy support.
// If no per-user parameters are supplied, then the defaultSocketFile
// is used to initialize the per-user conn params.
func Init(defaultSocketFile string) (*DBConfigs, error) {
	// The new base configs, if set, supersede legacy settings.
	for _, uc := range dbConfigs.userConfigs {
		if HasConnectionParams() {
			uc.param.Host = baseConfig.Host
			uc.param.Port = baseConfig.Port
			uc.param.UnixSocket = baseConfig.UnixSocket
		} else {
			uc.param.UnixSocket = defaultSocketFile
		}

		uc.param.Charset = baseConfig.Charset
		uc.param.Flags = baseConfig.Flags
		if uc.useSSL {
			uc.param.SslCa = baseConfig.SslCa
			uc.param.SslCaPath = baseConfig.SslCaPath
			uc.param.SslCert = baseConfig.SslCert
			uc.param.SslKey = baseConfig.SslKey
			uc.param.ServerName = baseConfig.ServerName
		}
	}

	// See if the CredentialsServer is working. We do not use the
	// result for anything, this is just a check.
	for _, uc := range dbConfigs.userConfigs {
		if _, err := WithCredentials(&uc.param); err != nil {
			return nil, fmt.Errorf("dbconfig cannot be initialized: %v", err)
		}
		// Check for only one.
		break
	}
	dbConfigs.SidecarDBName.Set("_vt")

	log.Infof("DBConfigs: %v\n", dbConfigs.String())
	return &dbConfigs, nil
}

// NewTestDBConfigs returns a DBConfigs meant for testing.
func NewTestDBConfigs(genParams, appDebugParams mysql.ConnParams, dbName string) *DBConfigs {
	dbcfgs := &DBConfigs{
		userConfigs: map[string]*userConfig{
			App:      {param: genParams},
			AppDebug: {param: appDebugParams},
			AllPrivs: {param: genParams},
			Dba:      {param: genParams},
			Filtered: {param: genParams},
			Repl:     {param: genParams},
		},
	}
	dbcfgs.DBName.Set(dbName)
	dbcfgs.SidecarDBName.Set("_vt")
	return dbcfgs
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Google Inc.

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

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

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

package servenv

import (
	"flag"
	"fmt"
	"net"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/servenv/grpcutils"
)

// This file handles gRPC server, on its own port.
// Clients register servers, based on service map:
//
// servenv.RegisterGRPCFlags()
// servenv.OnRun(func() {
//   if servenv.GRPCCheckServiceMap("XXX") {
//     pb.RegisterXXX(servenv.GRPCServer, XXX)
//   }
// }
//
// Note servenv.GRPCServer can only be used in servenv.OnRun,
// and not before, as it is initialized right before calling OnRun.
var (
	// GRPCPort is the port to listen on for gRPC. If not set or zero, don't listen.
	GRPCPort *int

	// GRPCCert is the cert to use if TLS is enabled
	GRPCCert *string

	// GRPCKey is the key to use if TLS is enabled
	GRPCKey *string

	// GRPCCA is the CA to use if TLS is enabled
	GRPCCA *string

	// GRPCMaxMessageSize is the maximum message size which the gRPC server will
	// accept. Larger messages will be rejected.
	GRPCMaxMessageSize *int

	// GRPCServer is the global server to serve gRPC.
	GRPCServer *grpc.Server
)

// isGRPCEnabled returns true if gRPC server is set
func isGRPCEnabled() bool {
	if GRPCPort != nil && *GRPCPort != 0 {
		return true
	}

	if SocketFile != nil && *SocketFile != "" {
		return true
	}

	return false
}

// createGRPCServer create the gRPC server we will be using.
// It has to be called after flags are parsed, but before
// services register themselves.
func createGRPCServer() {
	// skip if not registered
	if !isGRPCEnabled() {
		log.Infof("Skipping gRPC server creation")
		return
	}

	var opts []grpc.ServerOption
	if GRPCPort != nil && *GRPCCert != "" && *GRPCKey != "" {
		config, err := grpcutils.TLSServerConfig(*GRPCCert, *GRPCKey, *GRPCCA)
		if err != nil {
			log.Fatalf("Failed to log gRPC cert/key/ca: %v", err)
		}

		// create the creds server options
		creds := credentials.NewTLS(config)
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}
	// Override the default max message size for both send and receive
	// (which is 4 MiB in gRPC 1.0.0).
	// Large messages can occur when users try to insert or fetch very big
	// rows. If they hit the limit, they'll see the following error:
	// grpc: received message length XXXXXXX exceeding the max size 4194304
	// Note: For gRPC 1.0.0 it's sufficient to set the limit on the server only
	// because it's not enforced on the client side.
	if GRPCMaxMessageSize != nil {
		opts = append(opts, grpc.MaxRecvMsgSize(*GRPCMaxMessageSize))
		opts = append(opts, grpc.MaxSendMsgSize(*GRPCMaxMessageSize))
	}

	GRPCServer = grpc.NewServer(opts...)
}

func serveGRPC() {
	// skip if not registered
	if GRPCPort == nil || *GRPCPort == 0 {
		return
	}

	// listen on the port
	log.Infof("Listening for gRPC calls on port %v", *GRPCPort)
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *GRPCPort))
	if err != nil {
		log.Fatalf("Cannot listen on port %v for gRPC: %v", *GRPCPort, err)
	}

	// and serve on it
	go GRPCServer.Serve(listener)
}

// RegisterGRPCFlags registers the right command line flag to enable gRPC
func RegisterGRPCFlags() {
	GRPCPort = flag.Int("grpc_port", 0, "Port to listen on for gRPC calls")
	GRPCCert = flag.String("grpc_cert", "", "certificate to use, requires grpc_key, enables TLS")
	GRPCKey = flag.String("grpc_key", "", "key to use, requires grpc_cert, enables TLS")
	GRPCCA = flag.String("grpc_ca", "", "ca to use, requires TLS, and enforces client cert check")
	// Note: We're using 4 MiB as default value because that's the default in the
	// gRPC 1.0.0 Go server.
	GRPCMaxMessageSize = flag.Int("grpc_max_message_size", 4*1024*1024, "Maximum allowed RPC message size. Larger messages will be rejected by gRPC with the error 'exceeding the max size'.")
}

// GRPCCheckServiceMap returns if we should register a gRPC service
// (and also logs how to enable / disable it)
func GRPCCheckServiceMap(name string) bool {
	// Silently fail individual services if gRPC is not enabled in
	// the first place (either on a grpc port or on the socket file)
	if !isGRPCEnabled() {
		return false
	}

	// then check ServiceMap
	return CheckServiceMap("grpc", name)
}
<commit_msg>allow configuration of grpc max connection age<commit_after>/*
Copyright 2017 Google Inc.

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

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

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

package servenv

import (
	"flag"
	"fmt"
	"net"

	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"

	"math"
	"time"

	log "github.com/golang/glog"
	"github.com/youtube/vitess/go/vt/servenv/grpcutils"
	"google.golang.org/grpc/keepalive"
)

// This file handles gRPC server, on its own port.
// Clients register servers, based on service map:
//
// servenv.RegisterGRPCFlags()
// servenv.OnRun(func() {
//   if servenv.GRPCCheckServiceMap("XXX") {
//     pb.RegisterXXX(servenv.GRPCServer, XXX)
//   }
// }
//
// Note servenv.GRPCServer can only be used in servenv.OnRun,
// and not before, as it is initialized right before calling OnRun.
var (
	// GRPCPort is the port to listen on for gRPC. If not set or zero, don't listen.
	GRPCPort *int

	// GRPCCert is the cert to use if TLS is enabled
	GRPCCert *string

	// GRPCKey is the key to use if TLS is enabled
	GRPCKey *string

	// GRPCCA is the CA to use if TLS is enabled
	GRPCCA *string

	// GRPCMaxMessageSize is the maximum message size which the gRPC server will
	// accept. Larger messages will be rejected.
	GRPCMaxMessageSize *int

	// GRPCServer is the global server to serve gRPC.
	GRPCServer *grpc.Server

	// GRPCMaxConnectionAge is the maximum age of a client connection, before GoAway is sent.
	// This is useful for L4 loadbalancing to ensure rebalancing after scaling.
	GRPCMaxConnectionAge *time.Duration
)

// isGRPCEnabled returns true if gRPC server is set
func isGRPCEnabled() bool {
	if GRPCPort != nil && *GRPCPort != 0 {
		return true
	}

	if SocketFile != nil && *SocketFile != "" {
		return true
	}

	return false
}

// createGRPCServer create the gRPC server we will be using.
// It has to be called after flags are parsed, but before
// services register themselves.
func createGRPCServer() {
	// skip if not registered
	if !isGRPCEnabled() {
		log.Infof("Skipping gRPC server creation")
		return
	}

	var opts []grpc.ServerOption
	if GRPCPort != nil && *GRPCCert != "" && *GRPCKey != "" {
		config, err := grpcutils.TLSServerConfig(*GRPCCert, *GRPCKey, *GRPCCA)
		if err != nil {
			log.Fatalf("Failed to log gRPC cert/key/ca: %v", err)
		}

		// create the creds server options
		creds := credentials.NewTLS(config)
		opts = []grpc.ServerOption{grpc.Creds(creds)}
	}
	// Override the default max message size for both send and receive
	// (which is 4 MiB in gRPC 1.0.0).
	// Large messages can occur when users try to insert or fetch very big
	// rows. If they hit the limit, they'll see the following error:
	// grpc: received message length XXXXXXX exceeding the max size 4194304
	// Note: For gRPC 1.0.0 it's sufficient to set the limit on the server only
	// because it's not enforced on the client side.
	if GRPCMaxMessageSize != nil {
		opts = append(opts, grpc.MaxRecvMsgSize(*GRPCMaxMessageSize))
		opts = append(opts, grpc.MaxSendMsgSize(*GRPCMaxMessageSize))
	}

	if GRPCMaxConnectionAge != nil {
		opts = append(opts, grpc.KeepaliveParams(keepalive.ServerParameters{
			MaxConnectionAge: *GRPCMaxConnectionAge,
		}))
	}

	GRPCServer = grpc.NewServer(opts...)
}

func serveGRPC() {
	// skip if not registered
	if GRPCPort == nil || *GRPCPort == 0 {
		return
	}

	// listen on the port
	log.Infof("Listening for gRPC calls on port %v", *GRPCPort)
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *GRPCPort))
	if err != nil {
		log.Fatalf("Cannot listen on port %v for gRPC: %v", *GRPCPort, err)
	}

	// and serve on it
	go GRPCServer.Serve(listener)
}

// RegisterGRPCFlags registers the right command line flag to enable gRPC
func RegisterGRPCFlags() {
	GRPCPort = flag.Int("grpc_port", 0, "Port to listen on for gRPC calls")
	GRPCCert = flag.String("grpc_cert", "", "certificate to use, requires grpc_key, enables TLS")
	GRPCKey = flag.String("grpc_key", "", "key to use, requires grpc_cert, enables TLS")
	GRPCCA = flag.String("grpc_ca", "", "ca to use, requires TLS, and enforces client cert check")
	// Note: We're using 4 MiB as default value because that's the default in the
	// gRPC 1.0.0 Go server.
	GRPCMaxMessageSize = flag.Int("grpc_max_message_size", 4*1024*1024, "Maximum allowed RPC message size. Larger messages will be rejected by gRPC with the error 'exceeding the max size'.")
	// Default is effectively infinity, as defined in grpc.
	GRPCMaxConnectionAge = flag.Duration("grpc_max_connection_age", time.Duration(math.MaxInt64), "Maximum age of a client connection before GoAway is sent.")
}

// GRPCCheckServiceMap returns if we should register a gRPC service
// (and also logs how to enable / disable it)
func GRPCCheckServiceMap(name string) bool {
	// Silently fail individual services if gRPC is not enabled in
	// the first place (either on a grpc port or on the socket file)
	if !isGRPCEnabled() {
		return false
	}

	// then check ServiceMap
	return CheckServiceMap("grpc", name)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

import (
	"github.com/longbai/checkstyle"
)

var config = flag.String("config", "", "config json file")

var checker checkstyle.Checker

type Ignore struct {
	Files []string `json:"ignore"`
}

var ignore Ignore

func main() {
	flag.Parse()

	files := flag.Args()

	if config == nil {
		log.Println("No Config")
		return
	}
	conf, err := ioutil.ReadFile(*config)
	if err != nil {
		log.Printf("Open Config %v Fail %v\n", *config, err)
		return
	}

	err = json.Unmarshal(conf, &ignore)
	if err != nil {
		log.Printf("Parse Config %v Fail \n", *config, err)
		return
	}
	checker, err = checkstyle.New(conf)
	if err != nil {
		fmt.Errorf("New Checker Fail %s\n", err.Error())
		return
	}

	for _, v := range files {
		if isDir(v) {
			checkDir(v)
		} else {
			checkFile(v)
		}
	}
}

func isDir(filename string) bool {
	fi, err := os.Stat(filename)
	return err == nil && fi.IsDir()
}

func checkFile(fileName string) {
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		log.Printf("Read File Fail %v %v\n", fileName, err)
	}

	ps, err := checker.Check(fileName, file)
	if err != nil {
		log.Printf("Parse File Fail %v %v\n", fileName, err)
	}

	for _, p := range ps {
		fmt.Printf("%v: %s\n", p.Position, p.Description)
	}
}

func isIgnore(fileName string) bool {
	for _, v := range ignore.Files {
		if v == fileName {
			return true
		}
	}
	return false
}

func checkDir(dir string) {
	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() && strings.HasSuffix(path, ".go") && !isIgnore(path) {
			checkFile(path)
		}
		return err
	})
}
<commit_msg>log print<commit_after>package main

import (
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"os"
	"path/filepath"
	"strings"
)

import (
	"github.com/longbai/checkstyle"
)

var config = flag.String("config", "", "config json file")

var checker checkstyle.Checker

type Ignore struct {
	Files []string `json:"ignore"`
}

var ignore Ignore

func main() {
	flag.Parse()

	files := flag.Args()

	if config == nil {
		log.Println("No Config")
		return
	}
	conf, err := ioutil.ReadFile(*config)
	if err != nil {
		log.Printf("Open Config %v Fail %v\n", *config, err)
		return
	}

	err = json.Unmarshal(conf, &ignore)
	if err != nil {
		log.Printf("Parse Config %v Fail \n", *config, err)
		return
	}
	checker, err = checkstyle.New(conf)
	if err != nil {
		fmt.Errorf("New Checker Fail %s\n", err.Error())
		return
	}

	for _, v := range files {
		if isDir(v) {
			checkDir(v)
		} else {
			checkFile(v)
		}
	}
}

func isDir(filename string) bool {
	fi, err := os.Stat(filename)
	return err == nil && fi.IsDir()
}

func checkFile(fileName string) {
	file, err := ioutil.ReadFile(fileName)
	if err != nil {
		log.Printf("Read File Fail %v %v\n", fileName, err)
	}

	ps, err := checker.Check(fileName, file)
	if err != nil {
		log.Printf("Parse File Fail %v %v\n", fileName, err)
	}

	for _, p := range ps {
		log.Printf("%v: %s\n", p.Position, p.Description)
	}
}

func isIgnore(fileName string) bool {
	for _, v := range ignore.Files {
		if v == fileName {
			return true
		}
	}
	return false
}

func checkDir(dir string) {
	filepath.Walk(dir, func(path string, info os.FileInfo, err error) error {
		if err == nil && !info.IsDir() && strings.HasSuffix(path, ".go") && !isIgnore(path) {
			checkFile(path)
		}
		return err
	})
}
<|endoftext|>


<commit_before>package tchannel

// Copyright (c) 2015 Uber Technologies, Inc.

// Permission is hereby granted, free of charge, to any person obtaining a copy
// of this software and associated documentation files (the "Software"), to deal
// in the Software without restriction, including without limitation the rights
// to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
// copies of the Software, and to permit persons to whom the Software is
// furnished to do so, subject to the following conditions:
//
// The above copyright notice and this permission notice shall be included in
// all copies or substantial portions of the Software.
//
// THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
// IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
// FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
// AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
// LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
// OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
// THE SOFTWARE.

import (
	"bytes"
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel/golang/typed"
)

func TestNoFragmentation(t *testing.T) {
	in, out := buildChannels(ChecksumTypeCrc32)

	arg1 := []byte("Hello")
	w := newBodyWriter(out)
	require.Nil(t, w.WriteArgument(BytesOutput(arg1), true))

	// Should be a single frame
	// fragment flags(1), checksum type (1), checksum(5), chunk size(2), chunk(5)
	expectedFrames := combineBuffers([][]byte{
		[]byte{0x00, byte(ChecksumTypeCrc32)},
		ChecksumTypeCrc32.New().Add([]byte("Hello")),
		[]byte{0x00, 0x05},
		[]byte("Hello")})
	assertFramesEqual(t, expectedFrames, out.sentFragments, "no fragmentation")

	r1 := newBodyReader(in, true)
	rarg1 := make([]byte, len(arg1))
	if _, err := r1.Read(rarg1); err != nil {
		require.Nil(t, err)
	}

	assert.Equal(t, arg1, rarg1)
	require.Nil(t, r1.endArgument())
}

func TestFragmentationRoundTrip(t *testing.T) {
	in, out := buildChannels(ChecksumTypeCrc32)

	// Write three arguments, each of which should span fragments
	arg1 := make([]byte, MaxFramePayloadSize*2+756)
	for i := range arg1 {
		arg1[i] = byte(i % 0x0F)
	}
	w := newBodyWriter(out)
	if _, err := w.Write(arg1); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg2 := make([]byte, MaxFramePayloadSize+229)
	for i := range arg2 {
		arg2[i] = byte(i%0x0F) + 0x10
	}
	if _, err := w.Write(arg2); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg3 := make([]byte, MaxFramePayloadSize+72)
	for i := range arg3 {
		arg3[i] = byte(i%0x0F) + 0x20
	}
	if _, err := w.Write(arg3); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(true))

	// Read the three arguments
	r1 := newBodyReader(in, false)

	rarg1 := make([]byte, len(arg1))
	if _, err := r1.Read(rarg1); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg1, rarg1)
	require.Nil(t, r1.endArgument())

	r2 := newBodyReader(in, false)
	rarg2 := make([]byte, len(arg2))
	if _, err := r2.Read(rarg2); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg2, rarg2)
	require.Nil(t, r2.endArgument())

	r3 := newBodyReader(in, true)
	rarg3 := make([]byte, len(arg3))
	if _, err := r3.Read(rarg3); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg3, rarg3)
	require.Nil(t, r3.endArgument())
}

func TestArgEndOnFragmentBoundary(t *testing.T) {
	// Each argument should line up exactly at the end of each fragment
	in, out := buildChannels(ChecksumTypeCrc32)

	// Calculate the number of bytes available in the fragment content,
	// which is the size of the full frame minus the header content for the
	// fragment.  Header content consists of 1 byte flag, 1 byte checksum
	// type, 4 byte checksum value, for a total of 6 bytes
	fragmentContentSize := int(MaxFramePayloadSize) - 6
	arg1 := make([]byte, fragmentContentSize-2) // reserve 2 bytes for the arg chunk size
	for i := range arg1 {
		arg1[i] = byte(i % 0x0F)
	}
	w := newBodyWriter(out)
	if _, err := w.Write(arg1); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg2 := make([]byte, len(arg1)-2) // additional 2 byte trailing size for arg1
	for i := range arg2 {
		arg2[i] = byte(i % 0x1F)
	}
	if _, err := w.Write(arg2); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(false))

	arg3 := make([]byte, len(arg2)) // additional 2 byte trailing size for arg2
	for i := range arg3 {
		arg3[i] = byte(i % 0x2F)
	}
	if _, err := w.Write(arg3); err != nil {
		require.Nil(t, err)
	}
	require.Nil(t, w.endArgument(true))

	// We should have sent 4 fragments (one for arg1, one for zero arg1
	// size + arg2, one for zero arg2 size + arg3, one for zero arg3 size)
	sentFragments := out.sentFragments
	require.Equal(t, 4, len(sentFragments))
	lastFragment := sentFragments[len(sentFragments)-1]

	// 1 byte flags, 1 byte checksum type, 4 bytes checksum, 2 bytes size (0)
	require.Equal(t, 8, int(lastFragment.Header.PayloadSize()))
	r1 := newBodyReader(in, false)

	rarg1 := make([]byte, len(arg1))
	if _, err := r1.Read(rarg1); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg1, rarg1)
	require.Nil(t, r1.endArgument())

	r2 := newBodyReader(in, false)
	rarg2 := make([]byte, len(arg2))
	if _, err := r2.Read(rarg2); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg2, rarg2)
	require.Nil(t, r2.endArgument())

	r3 := newBodyReader(in, true)
	rarg3 := make([]byte, len(arg3))
	if _, err := r3.Read(rarg3); err != nil {
		require.Nil(t, err)
	}
	assert.Equal(t, arg3, rarg3)
	require.Nil(t, r3.endArgument())
}

func TestEmptyFragments(t *testing.T) {
	in, out := buildChannels(ChecksumTypeCrc32)

	w := newBodyWriter(out)
	require.Nil(t, w.WriteArgument(BytesOutput(nil), false))
	require.Nil(t, w.WriteArgument(BytesOutput(nil), true))

	r1 := newBodyReader(in, false)
	var arg1 BytesInput
	require.Nil(t, r1.ReadArgument(&arg1, false))
	assert.Equal(t, 0, len(arg1))

	r2 := newBodyReader(in, true)
	var arg2 BytesInput
	require.Nil(t, r2.ReadArgument(&arg2, true))
	assert.Equal(t, 0, len(arg2))
}

func buildChannels(checksumType ChecksumType) (*inFragments, *outFragments) {
	ch := make(chan *Frame, 512)

	in := &inFragments{ch: ch}
	out := &outFragments{ch: ch, checksum: checksumType.New()}
	return in, out
}

type inFragments struct {
	checksum Checksum
	ch       <-chan *Frame
	current  *inFragment
}

type sampleMessage struct{}

func (m *sampleMessage) ID() uint32                       { return 0xDEADBEEF }
func (m *sampleMessage) messageType() messageType         { return messageTypeCallReq }
func (m *sampleMessage) read(r *typed.ReadBuffer) error   { return nil }
func (m *sampleMessage) write(w *typed.WriteBuffer) error { return nil }

func (in *inFragments) waitForFragment() (*inFragment, error) {
	if in.current == nil || !in.current.hasMoreChunks() {
		var err error
		f := <-in.ch
		if in.current, err = newInboundFragment(f, &sampleMessage{}, in.checksum); err != nil {
			return nil, err
		}

		in.checksum = in.current.checksum
	}

	return in.current, nil
}

type outFragments struct {
	fragmentSize  int
	checksum      Checksum
	ch            chan<- *Frame
	sentFragments []*Frame
}

func (out *outFragments) beginFragment() (*outFragment, error) {
	return newOutboundFragment(NewFrame(MaxFramePayloadSize), &sampleMessage{}, out.checksum)
}

func (out *outFragments) flushFragment(toSend *outFragment, last bool) error {
	f := toSend.finish(last)
	out.ch <- f
	out.sentFragments = append(out.sentFragments, f)
	return nil
}

func assertFramesEqual(t *testing.T, expected [][]byte, frames []*Frame, msg string) {
	require.Equal(t, len(expected), len(frames), fmt.Sprintf("incorrect number of frames for %s", msg))

	for i := range expected {
		assert.Equal(t, len(expected[i]), int(frames[i].Header.PayloadSize()),
			fmt.Sprintf("incorrect size for frame %d of %s", i, msg))
		assert.Equal(t, expected[i], frames[i].Payload[:frames[i].Header.PayloadSize()])
	}
}

func combineBuffers(elements ...[][]byte) [][]byte {
	var buffers [][]byte
	for i := range elements {
		buffers = append(buffers, bytes.Join(elements[i], []byte{}))
	}

	return buffers
}
<commit_msg>Start on fragmentation test<commit_after>package tchannel

import (
	"bytes"
	"sync"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
	"github.com/uber/tchannel/golang/typed"
)

const (
	testFragmentHeaderSize = 1 /* flags */ + 1 /* checksum type */ + 4 /* CRC32 checksum */

	testFragmentPayloadSize = 10 // enough room for a small payload
	testFragmentSize        = testFragmentHeaderSize + testFragmentPayloadSize
)

func TestFragmentationEmptyArgs(t *testing.T) {
	sendCh := make(fragmentChannel, 10)
	recvCh := make(fragmentChannel, 10)

	w := newFragmentingWriter(sendCh, ChecksumTypeCrc32.New())
	r := newFragmentingReader(recvCh)

	var fragments [][]byte
	var args [][]byte
	var wg sync.WaitGroup
	wg.Add(1)
	go func() {
		defer wg.Done()
		for fragment := range sendCh {
			fragments = append(fragments, fragment)
			recvCh <- fragment
		}
	}()

	wg.Add(1)
	go func() {
		defer wg.Done()

		var arg BytesInput
		require.NoError(t, r.ReadArgument(&arg, false))
		args = append(args, arg)
		arg = nil
		require.NoError(t, r.ReadArgument(&arg, false))
		args = append(args, arg)
		arg = nil
		require.NoError(t, r.ReadArgument(&arg, true))
		args = append(args, arg)
		arg = nil
	}()

	require.NoError(t, w.WriteArgument(BytesOutput(nil), false))
	require.NoError(t, w.WriteArgument(BytesOutput(nil), false))
	require.NoError(t, w.WriteArgument(BytesOutput(nil), true))
	close(sendCh)

	wg.Wait()
	assert.Equal(t, [][]byte{[]byte{}, []byte{}, []byte{}}, args)

	// Make sure the fragments look as we expected
	expectedFragments := buffers([][]byte{{
		0x00, // flags
		byte(ChecksumTypeCrc32), 0x00, 0x00, 0x00, 0x00, // empty checksum
		0x00, 0x00, // arg 1 (length no body)
		0x00, 0x00, // arg 2 (length no body)
		0x00, 0x00}, // arg 3 (length no body)
	})
	assert.Equal(t, expectedFragments, fragments)
}

func TestSingleFragment(t *testing.T) {
}

func TestMultipleFragments(t *testing.T) {
}

func TestMiddleArgOnFragmentBoundary(t *testing.T) {
}

func TestLastArgOnFragmentBoundary(t *testing.T) {
}

type fragmentChannel chan []byte

func (ch fragmentChannel) newFragment(initial bool, checksum Checksum) (*writableFragment, error) {
	wbuf := typed.NewWriteBuffer(make([]byte, testFragmentSize))
	fragment := new(writableFragment)
	fragment.flagsRef = wbuf.DeferByte()
	wbuf.WriteByte(byte(checksum.TypeCode()))
	fragment.checksumRef = wbuf.DeferBytes(checksum.Size())
	fragment.checksum = checksum
	fragment.contents = wbuf
	return fragment, wbuf.Err()
}

func (ch fragmentChannel) flushFragment(fragment *writableFragment) error {
	var buf bytes.Buffer
	fragment.contents.FlushTo(&buf)
	ch <- buf.Bytes()
	return nil
}

func (ch fragmentChannel) recvNextFragment(initial bool) (*readableFragment, error) {
	rbuf := typed.NewReadBuffer(<-ch)
	fragment := new(readableFragment)
	fragment.flags = rbuf.ReadByte()
	fragment.checksumType = ChecksumType(rbuf.ReadByte())
	fragment.checksum = rbuf.ReadBytes(fragment.checksumType.ChecksumSize())
	fragment.contents = rbuf
	return fragment, rbuf.Err()
}

func buffers(elements ...[][]byte) [][]byte {
	var buffers [][]byte
	for i := range elements {
		buffers = append(buffers, bytes.Join(elements[i], []byte{}))
	}

	return buffers
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"strconv"
	"strings"
	"time"
)

type ServeConfig struct {
	Port string
	Dir  string
	Path string
}

type statusWriter struct {
	http.ResponseWriter
	status int
	length int
}

func (w *statusWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	w.length = len(b)
	return w.ResponseWriter.Write(b)
}

func logger(handle http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, request *http.Request) {
		start := time.Now()
		writer := statusWriter{w, 0, 0}
		handle.ServeHTTP(&writer, request)
		end := time.Now()
		latency := end.Sub(start)
		statusCode := writer.status
		log.Println(request.Method, request.URL.Path, statusCode, latency)
	}
}

func getServeConfig(args []string) ServeConfig {

	dir, _ := os.Getwd()
	config := ServeConfig{"8000", dir, "/"}

	// if no args are passed, use default config
	if len(args) > 0 {
		for _, element := range args {

			_, err := strconv.Atoi(element)

			if err == nil {
				config.Port = element
			} else if strings.Contains(element, ":") {
				s := strings.Split(element, ":")
				config.Dir, config.Port = s[0], s[1]
			} else {
				config.Dir = element
			}
		}
	}

	return config
}

func serve(config ServeConfig) {

	pathStat, err := os.Stat(config.Dir)
	if err != nil {
		log.Fatal(err)
	}

	if pathStat.IsDir() {
		http.Handle(config.Path, http.FileServer(http.Dir(config.Dir)))
	} else {
		http.HandleFunc(config.Path, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, config.Dir)
		})
	}

	fmt.Println("\nServe Config \n Directory : " + path.Base(config.Dir) + " \n Path      : http://localhost:" + config.Port + config.Path + "\n")
	log.Println("Starting server on port: " + config.Port)

	log.Fatal(http.ListenAndServe(":"+config.Port, logger(http.DefaultServeMux)))
}

func main() {
	flag.Parse()
	args := flag.Args()
	config := getServeConfig(args)
	serve(config)
}
<commit_msg>switch completely to flag arguements<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"net/http"
	"os"
	"path"
	"time"
)

type config struct {
	Port    string
	Dir     string
	Path    string
	listDir bool
}

type statusWriter struct {
	http.ResponseWriter
	status int
	length int
}

func (w *statusWriter) WriteHeader(status int) {
	w.status = status
	w.ResponseWriter.WriteHeader(status)
}

func (w *statusWriter) Write(b []byte) (int, error) {
	if w.status == 0 {
		w.status = http.StatusOK
	}
	w.length = len(b)
	return w.ResponseWriter.Write(b)
}

func logger(handle http.Handler) http.HandlerFunc {
	return func(w http.ResponseWriter, request *http.Request) {
		start := time.Now()
		writer := statusWriter{w, 0, 0}
		handle.ServeHTTP(&writer, request)
		latency := time.Since(start)
		statusCode := writer.status
		log.Println(request.Method, request.URL.Path, statusCode, latency)
	}
}

func (c config) serve() {
	file, err := os.Stat(c.Dir)
	if err != nil {
		log.Fatal(err)
	}

	if file.IsDir() {
		http.Handle(c.Path, http.FileServer(http.Dir(c.Dir)))
	} else {
		http.HandleFunc(c.Path, func(w http.ResponseWriter, r *http.Request) {
			http.ServeFile(w, r, c.Dir)
		})
	}

	fmt.Println("\nServe Config \n Directory : " + path.Base(c.Dir) + " \n Path      : http://localhost:" + c.Port + c.Path + "\n")
	log.Println("Starting server on port: " + c.Port)

	log.Fatal(http.ListenAndServe(":"+c.Port, logger(http.DefaultServeMux)))
}

func main() {
	c := config{}
	flag.StringVar(&c.Port, "p", "8000", "Port Number")
	flag.StringVar(&c.Dir, "d", ".", "Serve Directory")
	flag.StringVar(&c.Path, "path", "/", "Public Access Path")
	flag.BoolVar(&c.listDir, "disable-dir", false, "Disable Directory Listing")
	flag.Parse()
	c.serve()
}
<|endoftext|>


<commit_before>/*
   Copyright 2017 Shlomi Noach, GitHub Inc.

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

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

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

package kv

import (
	"fmt"
	"sync"
)

type KVPair struct {
	Key   string
	Value string
}

func NewKVPair(key string, value string) *KVPair {
	return &KVPair{Key: key, Value: value}
}

func (this *KVPair) String() string {
	return fmt.Sprintf("%s:%s", this.Key, this.Value)
}

type KVStore interface {
	PutKeyValue(key string, value string) (err error)
	GetKeyValue(key string) (value string, found bool, err error)
	DistributePairs(pairs [](*KVPair)) (failedDistributions []string, err error)
}

type KVStoreMap map[string](KVStore)

var kvMutex sync.Mutex
var kvInitOnce sync.Once
var kvStores = make(KVStoreMap)

// InitKVStores initializes the KV stores (duh), once in the lifetime of this app.
// Configuration reload does not affect a running instance.
func InitKVStores() {
	kvMutex.Lock()
	defer kvMutex.Unlock()

	kvInitOnce.Do(func() {
		kvStores = KVStoreMap{
			"internal": NewInternalKVStore(),
			"consul":   NewConsulStore(),
			"zk":       NewZkStore(),
		}
	})
}

func getKVStores() (stores KVStoreMap) {
	kvMutex.Lock()
	defer kvMutex.Unlock()

	stores = kvStores
	return stores
}

func GetValue(key string) (value string, found bool, err error) {
	for _, store := range getKVStores() {
		// It's really only the first (internal) that matters here
		return store.GetKeyValue(key)
	}
	return value, found, err
}

func PutValue(key string, value string) (err error) {
	for _, store := range getKVStores() {
		if err := store.PutKeyValue(key, value); err != nil {
			return err
		}
	}
	return nil
}

func PutKVPair(kvPair *KVPair) (err error) {
	if kvPair == nil {
		return nil
	}
	return PutValue(kvPair.Key, kvPair.Value)
}

func DistributePairs(canonicalPairs [](*KVPair), fullPairs [](*KVPair)) (failedDistributions []string, err error) {
	for _, store := range getKVStores() {
		if failedDistributions, err := store.DistributePairs(canonicalPairs); err != nil {
			return failedDistributions, err
		}
	}
	return failedDistributions, nil
}
<commit_msg>both map and array<commit_after>/*
   Copyright 2017 Shlomi Noach, GitHub Inc.

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

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

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

package kv

import (
	"fmt"
	"sync"
)

type KVPair struct {
	Key   string
	Value string
}

func NewKVPair(key string, value string) *KVPair {
	return &KVPair{Key: key, Value: value}
}

func (this *KVPair) String() string {
	return fmt.Sprintf("%s:%s", this.Key, this.Value)
}

type KVStore interface {
	PutKeyValue(key string, value string) (err error)
	GetKeyValue(key string) (value string, found bool, err error)
	DistributePairs(pairs [](*KVPair)) (failedDistributions []string, err error)
}

type KVStoreMap map[string](KVStore)

var kvMutex sync.Mutex
var kvInitOnce sync.Once
var kvStores = []KVStore{}
var kvStoresMap = make(KVStoreMap)

// InitKVStores initializes the KV stores (duh), once in the lifetime of this app.
// Configuration reload does not affect a running instance.
func InitKVStores() {
	kvMutex.Lock()
	defer kvMutex.Unlock()

	kvInitOnce.Do(func() {
		kvStoresMap = KVStoreMap{
			"internal": NewInternalKVStore(),
			"consul":   NewConsulStore(),
			"zk":       NewZkStore(),
		}
		kvStores = []KVStore{
			kvStoresMap["internal"],
			kvStoresMap["consul"],
			kvStoresMap["zk"],
		}
	})
}

func getKVStores() (stores []KVStore) {
	kvMutex.Lock()
	defer kvMutex.Unlock()

	stores = kvStores
	return stores
}

func GetValue(key string) (value string, found bool, err error) {
	for _, store := range getKVStores() {
		// It's really only the first (internal) that matters here
		return store.GetKeyValue(key)
	}
	return value, found, err
}

func PutValue(key string, value string) (err error) {
	for _, store := range getKVStores() {
		if err := store.PutKeyValue(key, value); err != nil {
			return err
		}
	}
	return nil
}

func PutKVPair(kvPair *KVPair) (err error) {
	if kvPair == nil {
		return nil
	}
	return PutValue(kvPair.Key, kvPair.Value)
}

func DistributePairs(canonicalPairs [](*KVPair), fullPairs [](*KVPair)) (failedDistributions []string, err error) {
	for _, store := range getKVStores() {
		if failedDistributions, err := store.DistributePairs(canonicalPairs); err != nil {
			return failedDistributions, err
		}
	}
	return failedDistributions, nil
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"image"
	"image/color"
	"image/png"
	"os"
	"os/signal"
	"strconv"
	"time"
)

const (
	SOLDER     = 0x01
	TERMINAL   = 0x02
	TIMESLIZES = 0x04
)

// Version number of this tool
const Version = "0.0.2"

// Data is a struct for each network packet
type Data struct {
	toa     int64  // Timestamp of arrival in microseconds
	payload []byte // Copied network packet
}

// configs represents all the configuration data
type configs struct {
	bpP   uint // Bits per Pixel
	ppI   uint // Number of packets per Image
	ts    uint // "Duration" for one Image
	limit uint // Number of network packets to process
	stil  uint // Type of illustration
}

func getBitsFromPacket(packet []byte, byteP, bitP *int, bpP uint) uint8 {
	var c uint8
	for i := 0; i < (int(bpP) / 3); i++ {
		if *byteP >= len(packet) {
			break
		}
		c |= (packet[*byteP] & (1 << uint8(7-*bitP)))
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	}
	return c
}

func createPixel(packet []byte, byteP, bitP *int, bpP uint) (c color.Color) {
	var r, g, b uint8

	if bpP == 1 {
		if (packet[*byteP] & (1 << uint8(7-*bitP))) == 0 {
			c = color.NRGBA{R: 0,
				G: 0,
				B: 0,
				A: 255}
		} else {
			c = color.NRGBA{R: 255,
				G: 255,
				B: 255,
				A: 255}
		}
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	} else {
		r = getBitsFromPacket(packet, byteP, bitP, bpP)
		g = getBitsFromPacket(packet, byteP, bitP, bpP)
		b = getBitsFromPacket(packet, byteP, bitP, bpP)

		c = color.NRGBA{R: r,
			G: g,
			B: b,
			A: 255}
	}
	return
}

func createTerminalVisualization(data Data, bitsPerPixel uint) {
	var bitPos int
	var bytePos int
	var packetLen int

	packetLen = len(data.payload)
	bitPos = 0
	bytePos = 0
	for {
		c := createPixel(data.payload, &bytePos, &bitPos, bitsPerPixel)
		r, g, b, _ := c.RGBA()
		fmt.Printf("\x1B[0m\x1B[38;2;%d;%d;%dm\u2588", uint8(r), uint8(g), uint8(b))
		if bytePos >= packetLen {
			break
		}
	}
	fmt.Printf("\x1B[m\n")

}
func createTimeVisualization(data []Data, xMax int, prefix string, ts uint, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int
	var firstPkg time.Time

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, int(ts)))

	for pkg := range data {
		if firstPkg.IsZero() {
			firstPkg = time.Unix(0, data[pkg].toa*int64(time.Microsecond))
		}
		packetLen = len(data[pkg].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[pkg].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, int(data[pkg].toa%int64(ts)), c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}
	}

	filename := prefix
	filename += "-"
	filename += firstPkg.Format(time.RFC3339Nano)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		fmt.Errorf("%s", err)
		return
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		fmt.Errorf("%s", err)
	}

	if err := f.Close(); err != nil {
		fmt.Errorf("%s", err)
	}

	return
}

func createFixedVisualization(data []Data, xMax int, prefix string, num int, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, len(data)))

	for yPos := range data {
		packetLen = len(data[yPos].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[yPos].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, yPos, c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}

	}

	filename := prefix
	filename += strconv.Itoa(num)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		fmt.Errorf("%s", err)
		return
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		fmt.Errorf("%s", err)
	}

	if err := f.Close(); err != nil {
		fmt.Errorf("%s", err)
	}

	return
}

func handlePackets(ps *gopacket.PacketSource, num uint, ch chan<- Data, done <-chan bool) {
	var count uint
	for packet := range ps.Packets() {
		var k Data
		select {
		case <-done:
			close(ch)
			return
		default:
		}
		count++
		if num != 0 && count > num {
			break
		}

		elements := packet.Data()
		if len(elements) == 0 {
			continue
		}
		k = Data{toa: (packet.Metadata().CaptureInfo.Timestamp.UnixNano() / int64(time.Microsecond)), payload: packet.Data()}
		ch <- k
	}
	close(ch)
	return
}

func availableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		fmt.Errorf("%s", err)
		os.Exit(1)
	}

	for _, device := range devices {
		if len(device.Addresses) == 0 {
			continue
		}
		fmt.Println("Interface: ", device.Name)
		for _, address := range device.Addresses {
			fmt.Println("   IP address:  ", address.IP)
			fmt.Println("   Subnet mask: ", address.Netmask)
		}
		fmt.Println("")
	}
}

func initSource(dev, file *string) (handle *pcap.Handle, err error) {
	if len(*dev) > 0 {
		handle, err = pcap.OpenLive(*dev, 4096, true, pcap.BlockForever)
		if err != nil {
			fmt.Errorf("%s", err)
			os.Exit(1)
		}
	} else if len(*file) > 0 {
		handle, err = pcap.OpenOffline(*file)
		if err != nil {
			fmt.Errorf("%s", err)
			os.Exit(1)
		}
	} else {
		return nil, fmt.Errorf("Source is missing\n")
	}
	return
}

func checkConfig(cfg configs) error {
	if cfg.bpP%3 != 0 && cfg.bpP != 1 {
		return fmt.Errorf("%d must be divisible by three or should be one", cfg.bpP)
	} else if cfg.bpP > 25 {
		return fmt.Errorf("%d must be smaller than 25", cfg.bpP)
	}

	if cfg.ts > 0 {
		cfg.stil |= TIMESLIZES
	}

	if cfg.stil == (TIMESLIZES | TERMINAL) {
		return fmt.Errorf("-timeslize and -terminal can't be combined")
	} else if cfg.stil == 0 {
		// If way of stil is provided, we will stick to the default one
		cfg.stil |= SOLDER
	}
	return nil
}

func init() {
	if len(os.Args) < 2 {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		os.Exit(1)
	}
}

func main() {
	var err error
	var handle *pcap.Handle
	var data []Data
	var xMax int
	var index int = 1
	osSig := make(chan os.Signal, 1)
	done := make(chan bool, 1)
	signal.Notify(osSig, os.Interrupt)
	var slicer int64
	var cfg configs
	ch := make(chan Data)

	go func() {
		<-osSig     // Blocking till interrupt signal is received
		osSig = nil // ignore further signals
		done <- true
	}()

	dev := flag.String("interface", "", "Choose an interface for online processing")
	file := flag.String("file", "", "Choose a file for offline processing")
	filter := flag.String("filter", "", "Set a specific filter")
	lst := flag.Bool("list_interfaces", false, "List available interfaces")
	vers := flag.Bool("version", false, "Show version")
	help := flag.Bool("help", false, "Show this help")
	terminalOut := flag.Bool("terminal", false, "Visualize on terminal")
	num := flag.Uint("count", 25, "Number of packets to process.\n\tIf argument is 0 the limit is removed")
	output := flag.String("prefix", "image", "Prefix of the resulting image")
	size := flag.Uint("size", 25, "Number of packets per image")
	bits := flag.Uint("bits", 24, "Number of bits per pixel.\n\tIt must be divisible by three and smaller than 25\n\tTo get black/white results, choose 1 as input.")
	ts := flag.Uint("timeslize", 0, "Number of microseconds per resulting image.\n\tSo each pixel of the height of the resulting image represents one microsecond")
	flag.Parse()

	if *lst {
		availableInterfaces()
		return
	}

	if *vers {
		fmt.Println("Version:", Version)
		return
	}

	if *help {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	cfg.bpP = *bits
	cfg.ppI = *size
	cfg.ts = *ts
	cfg.limit = *num
	cfg.stil = 0

	if *terminalOut == true {
		cfg.stil |= TERMINAL
	}
	if *ts != 0 {
		cfg.stil |= TIMESLIZES
	}

	if err = checkConfig(cfg); err != nil {
		fmt.Errorf("%s", err)
		os.Exit(1)
	}

	handle, err = initSource(dev, file)
	if err != nil {
		fmt.Errorf("%s", err)
		os.Exit(1)
	}
	defer handle.Close()

	if len(*filter) != 0 {
		err = handle.SetBPFFilter(*filter)
		if err != nil {
			fmt.Errorf("%s\nInvalid Filter: %s", err, *filter)
			os.Exit(1)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	packetSource.DecodeOptions = gopacket.Lazy

	go handlePackets(packetSource, cfg.limit, ch, done)

	switch cfg.stil {
	case SOLDER:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
			if len(data) >= int(*size) {
				xMax++
				createFixedVisualization(data, xMax, *output, index, cfg.bpP)
				xMax = 0
				index++
				data = data[:0]
			}
		}
	case TERMINAL:
		for i, ok := <-ch; ok; i, ok = <-ch {
			createTerminalVisualization(i, cfg.bpP)
		}
	case TIMESLIZES:
		for i, ok := <-ch; ok; i, ok = <-ch {
			if slicer == 0 {
				slicer = i.toa + int64(*ts)
			}
			if slicer < i.toa {
				xMax++
				createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
				xMax = 0
				data = data[:0]
				slicer = i.toa + int64(*ts)
			}
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
		}
	}

	if len(data) > 0 {
		xMax++
		switch cfg.stil {
		case SOLDER:
			createFixedVisualization(data, xMax, *output, index, cfg.bpP)
		case TIMESLIZES:
			createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
		}
	}

}
<commit_msg>Remove obsolete Code<commit_after>package main

import (
	"flag"
	"fmt"
	"github.com/google/gopacket"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"image"
	"image/color"
	"image/png"
	"os"
	"os/signal"
	"strconv"
	"time"
)

const (
	SOLDER     = 0x01
	TERMINAL   = 0x02
	TIMESLIZES = 0x04
)

// Version number of this tool
const Version = "0.0.2"

// Data is a struct for each network packet
type Data struct {
	toa     int64  // Timestamp of arrival in microseconds
	payload []byte // Copied network packet
}

// configs represents all the configuration data
type configs struct {
	bpP   uint // Bits per Pixel
	ppI   uint // Number of packets per Image
	ts    uint // "Duration" for one Image
	limit uint // Number of network packets to process
	stil  uint // Type of illustration
}

func getBitsFromPacket(packet []byte, byteP, bitP *int, bpP uint) uint8 {
	var c uint8
	for i := 0; i < (int(bpP) / 3); i++ {
		if *byteP >= len(packet) {
			break
		}
		c |= (packet[*byteP] & (1 << uint8(7-*bitP)))
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	}
	return c
}

func createPixel(packet []byte, byteP, bitP *int, bpP uint) (c color.Color) {
	var r, g, b uint8

	if bpP == 1 {
		if (packet[*byteP] & (1 << uint8(7-*bitP))) == 0 {
			c = color.NRGBA{R: 0,
				G: 0,
				B: 0,
				A: 255}
		} else {
			c = color.NRGBA{R: 255,
				G: 255,
				B: 255,
				A: 255}
		}
		*bitP += 1
		if *bitP%8 == 0 {
			*bitP = 0
			*byteP += 1
		}
	} else {
		r = getBitsFromPacket(packet, byteP, bitP, bpP)
		g = getBitsFromPacket(packet, byteP, bitP, bpP)
		b = getBitsFromPacket(packet, byteP, bitP, bpP)

		c = color.NRGBA{R: r,
			G: g,
			B: b,
			A: 255}
	}
	return
}

func createTerminalVisualization(data Data, bitsPerPixel uint) {
	var bitPos int
	var bytePos int
	var packetLen int

	packetLen = len(data.payload)
	bitPos = 0
	bytePos = 0
	for {
		c := createPixel(data.payload, &bytePos, &bitPos, bitsPerPixel)
		r, g, b, _ := c.RGBA()
		fmt.Printf("\x1B[0m\x1B[38;2;%d;%d;%dm\u2588", uint8(r), uint8(g), uint8(b))
		if bytePos >= packetLen {
			break
		}
	}
	fmt.Printf("\x1B[m\n")

}
func createTimeVisualization(data []Data, xMax int, prefix string, ts uint, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int
	var firstPkg time.Time

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, int(ts)))

	for pkg := range data {
		if firstPkg.IsZero() {
			firstPkg = time.Unix(0, data[pkg].toa*int64(time.Microsecond))
		}
		packetLen = len(data[pkg].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[pkg].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, int(data[pkg].toa%int64(ts)), c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}
	}

	filename := prefix
	filename += "-"
	filename += firstPkg.Format(time.RFC3339Nano)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		fmt.Errorf("%s", err)
		return
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		fmt.Errorf("%s", err)
	}

	if err := f.Close(); err != nil {
		fmt.Errorf("%s", err)
	}

	return
}

func createFixedVisualization(data []Data, xMax int, prefix string, num int, bitsPerPixel uint) {
	var xPos int
	var bitPos int
	var bytePos int
	var packetLen int

	img := image.NewNRGBA(image.Rect(0, 0, (xMax*8)/int(bitsPerPixel)+1, len(data)))

	for yPos := range data {
		packetLen = len(data[yPos].payload)
		xPos = 0
		bitPos = 0
		bytePos = 0
		for {
			c := createPixel(data[yPos].payload, &bytePos, &bitPos, bitsPerPixel)
			img.Set(xPos, yPos, c)
			xPos++
			if bytePos >= packetLen {
				break
			}
		}

	}

	filename := prefix
	filename += strconv.Itoa(num)
	filename += ".png"
	f, err := os.Create(filename)
	if err != nil {
		fmt.Errorf("%s", err)
		return
	}

	if err := png.Encode(f, img); err != nil {
		f.Close()
		fmt.Errorf("%s", err)
	}

	if err := f.Close(); err != nil {
		fmt.Errorf("%s", err)
	}

	return
}

func handlePackets(ps *gopacket.PacketSource, num uint, ch chan<- Data, done <-chan os.Signal) {
	var count uint
	for packet := range ps.Packets() {
		var k Data
		select {
		case <-done:
			close(ch)
			return
		default:
		}
		count++
		if num != 0 && count > num {
			break
		}

		elements := packet.Data()
		if len(elements) == 0 {
			continue
		}
		k = Data{toa: (packet.Metadata().CaptureInfo.Timestamp.UnixNano() / int64(time.Microsecond)), payload: packet.Data()}
		ch <- k
	}
	close(ch)
	return
}

func availableInterfaces() {
	devices, err := pcap.FindAllDevs()
	if err != nil {
		fmt.Errorf("%s", err)
		os.Exit(1)
	}

	for _, device := range devices {
		if len(device.Addresses) == 0 {
			continue
		}
		fmt.Println("Interface: ", device.Name)
		for _, address := range device.Addresses {
			fmt.Println("   IP address:  ", address.IP)
			fmt.Println("   Subnet mask: ", address.Netmask)
		}
		fmt.Println("")
	}
}

func initSource(dev, file *string) (handle *pcap.Handle, err error) {
	if len(*dev) > 0 {
		handle, err = pcap.OpenLive(*dev, 4096, true, pcap.BlockForever)
		if err != nil {
			fmt.Errorf("%s", err)
			os.Exit(1)
		}
	} else if len(*file) > 0 {
		handle, err = pcap.OpenOffline(*file)
		if err != nil {
			fmt.Errorf("%s", err)
			os.Exit(1)
		}
	} else {
		return nil, fmt.Errorf("Source is missing\n")
	}
	return
}

func checkConfig(cfg configs) error {
	if cfg.bpP%3 != 0 && cfg.bpP != 1 {
		return fmt.Errorf("%d must be divisible by three or should be one", cfg.bpP)
	} else if cfg.bpP > 25 {
		return fmt.Errorf("%d must be smaller than 25", cfg.bpP)
	}

	if cfg.ts > 0 {
		cfg.stil |= TIMESLIZES
	}

	if cfg.stil == (TIMESLIZES | TERMINAL) {
		return fmt.Errorf("-timeslize and -terminal can't be combined")
	} else if cfg.stil == 0 {
		// If way of stil is provided, we will stick to the default one
		cfg.stil |= SOLDER
	}
	return nil
}

func init() {
	if len(os.Args) < 2 {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		os.Exit(1)
	}
}

func main() {
	var err error
	var handle *pcap.Handle
	var data []Data
	var xMax int
	var index int = 1
	osSig := make(chan os.Signal, 1)
	signal.Notify(osSig, os.Interrupt)
	var slicer int64
	var cfg configs
	ch := make(chan Data)

	dev := flag.String("interface", "", "Choose an interface for online processing")
	file := flag.String("file", "", "Choose a file for offline processing")
	filter := flag.String("filter", "", "Set a specific filter")
	lst := flag.Bool("list_interfaces", false, "List available interfaces")
	vers := flag.Bool("version", false, "Show version")
	help := flag.Bool("help", false, "Show this help")
	terminalOut := flag.Bool("terminal", false, "Visualize on terminal")
	num := flag.Uint("count", 25, "Number of packets to process.\n\tIf argument is 0 the limit is removed")
	output := flag.String("prefix", "image", "Prefix of the resulting image")
	size := flag.Uint("size", 25, "Number of packets per image")
	bits := flag.Uint("bits", 24, "Number of bits per pixel.\n\tIt must be divisible by three and smaller than 25\n\tTo get black/white results, choose 1 as input.")
	ts := flag.Uint("timeslize", 0, "Number of microseconds per resulting image.\n\tSo each pixel of the height of the resulting image represents one microsecond")
	flag.Parse()

	if *lst {
		availableInterfaces()
		return
	}

	if *vers {
		fmt.Println("Version:", Version)
		return
	}

	if *help {
		fmt.Println(os.Args[0], "[-bits ...] [-count ...] [-file ... | -interface ...] [-filter ...] [-list_interfaces] [-help] [-prefix ...] [-size ... | -timeslize ... | -terminal] [-version]")
		flag.PrintDefaults()
		return
	}

	cfg.bpP = *bits
	cfg.ppI = *size
	cfg.ts = *ts
	cfg.limit = *num
	cfg.stil = 0

	if *terminalOut == true {
		cfg.stil |= TERMINAL
	}
	if *ts != 0 {
		cfg.stil |= TIMESLIZES
	}

	if err = checkConfig(cfg); err != nil {
		fmt.Errorf("%s", err)
		os.Exit(1)
	}

	handle, err = initSource(dev, file)
	if err != nil {
		fmt.Errorf("%s", err)
		os.Exit(1)
	}
	defer handle.Close()

	if len(*filter) != 0 {
		err = handle.SetBPFFilter(*filter)
		if err != nil {
			fmt.Errorf("%s\nInvalid Filter: %s", err, *filter)
			os.Exit(1)
		}
	}

	packetSource := gopacket.NewPacketSource(handle, layers.LayerTypeEthernet)
	packetSource.DecodeOptions = gopacket.Lazy

	go handlePackets(packetSource, cfg.limit, ch, osSig)

	switch cfg.stil {
	case SOLDER:
		for i, ok := <-ch; ok; i, ok = <-ch {
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
			if len(data) >= int(*size) {
				xMax++
				createFixedVisualization(data, xMax, *output, index, cfg.bpP)
				xMax = 0
				index++
				data = data[:0]
			}
		}
	case TERMINAL:
		for i, ok := <-ch; ok; i, ok = <-ch {
			createTerminalVisualization(i, cfg.bpP)
		}
	case TIMESLIZES:
		for i, ok := <-ch; ok; i, ok = <-ch {
			if slicer == 0 {
				slicer = i.toa + int64(*ts)
			}
			if slicer < i.toa {
				xMax++
				createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
				xMax = 0
				data = data[:0]
				slicer = i.toa + int64(*ts)
			}
			data = append(data, i)
			if xMax < len(i.payload) {
				xMax = len(i.payload)
			}
		}
	}

	if len(data) > 0 {
		xMax++
		switch cfg.stil {
		case SOLDER:
			createFixedVisualization(data, xMax, *output, index, cfg.bpP)
		case TIMESLIZES:
			createTimeVisualization(data, xMax, *output, *ts, cfg.bpP)
		}
	}

}
<|endoftext|>


<commit_before>// Copyright (c) 2012, mulander <netprobe@gmail.com>
// All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
package main

import (
	"code.google.com/p/plotinum/plot"
	"code.google.com/p/plotinum/plotter"
	"code.google.com/p/plotinum/plotutil"
	"fmt"
	"log"
	"os/exec"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Sampler interface {
	Sample(cmd *exec.Cmd, ticker *time.Ticker)
	GetData() []Data
	Stop()
}

type Data struct {
	label string
	data  []float64
}

type TopSampler struct {
	Samples []Data
	stop    chan bool
}

func (t *TopSampler) GetData() []Data {
	return t.Samples
}

func (t *TopSampler) Stop() {
	t.stop <- true
}

func (t *TopSampler) Sample(cmd *exec.Cmd, ticker *time.Ticker) {
	// %CPU(field=8) + %MEM(field=9)
	t.stop = make(chan bool)
	t.Samples = make([]Data, 2)
	t.Samples[0].label = "CPU"
	t.Samples[0].data = make([]float64, 1)
	t.Samples[1].label = "MEM"
	t.Samples[1].data = make([]float64, 1)
	raw := "(?m)%d.*$"
	r := regexp.MustCompile(fmt.Sprintf(raw, cmd.Process.Pid))
	for {
		select {
		case <-t.stop:
			return
		case <-ticker.C:
			top := exec.Command("top", "-b", "-n 1", fmt.Sprintf("-p %d", cmd.Process.Pid))
			log.Printf("Sampling the process")
			out, err := top.Output()
			if err != nil {
				log.Fatal(err)
			}
			fields := strings.Fields(r.FindString(fmt.Sprintf("%s", out)))
			if len(fields) != 0 {
				cpu, err := strconv.ParseFloat(fields[8], 64)
				if err != nil {
					log.Fatal(err)
				}
				mem, err := strconv.ParseFloat(fields[9], 64)
				if err != nil {
					log.Fatal(err)
				}
				t.Samples[0].data = append(t.Samples[0].data, cpu) // CPU
				t.Samples[1].data = append(t.Samples[1].data, mem) // MEM
				log.Printf("%+v", fields)
			}
		}
	}
}

type Project struct {
	Command    *exec.Cmd     // executable to run during the test
	Duration   time.Duration // duration of a single sample
	Sampling   time.Duration // trigger sampling based on this interval
	Executions int           // the total number of runs for a single test
	Sampler    Sampler
}

func (p *Project) Exec(samplers chan []Data) {
	err := p.Command.Start()
	if err != nil {
		log.Fatal(err)
	}

	ticker := time.NewTicker(p.Sampling)

	// Possibly more samplers in the future
	var wg sync.WaitGroup

	wg.Add(1)
	go func(n int) {
		defer wg.Done()
		p.Sampler.Sample(p.Command, ticker)
		samplers <- p.Sampler.GetData()
	}(1)
	go func() {
		wg.Wait()
		close(samplers)
	}()

	done := make(chan error)
	go func() {
		done <- p.Command.Wait()
	}()
	select {
	case <-time.After(p.Duration):
		if err := p.Command.Process.Kill(); err != nil {
			log.Fatal("Failed to kill: ", err)
		}
		<-done
		log.Println("Process killed")
	case err := <-done:
		log.Printf("Process done with error = %v", err)
	}
	log.Printf("Waiting for the ticker to stop")
	ticker.Stop()
	log.Printf("Stopping samplers")
	p.Sampler.Stop()
}

func main() {
	project := &Project{
		Command:    exec.Command("yes"),
		Duration:   10 * time.Second,
		Sampling:   1 * time.Second,
		Executions: 1,
		Sampler:    &TopSampler{}}

	samplers := make(chan []Data)

	project.Exec(samplers)

	log.Printf("Waiting for the samplers")
	for s := range samplers {
		log.Printf("%+v", s)

		for _, sample := range s {

			p, err := plot.New()
			if err != nil {
				log.Fatal(err)
			}

			p.Title.Text = fmt.Sprintf("%s graph", sample.label)
			p.X.Label.Text = fmt.Sprintf("tick per %s", project.Sampling)
			p.Y.Label.Text = sample.label

			pts := make(plotter.XYs, len(sample.data))
			for i := range pts {
				pts[i].X = float64(i)
				pts[i].Y = sample.data[i]
			}

			plotutil.AddLinePoints(p, sample.label, pts)
			plotFile := fmt.Sprintf("%s-%s.png", project.Command.Args[0], sample.label)

			if err := p.Save(4, 4, plotFile); err != nil {
				log.Fatal(err)
			}
		}
	}
}
<commit_msg>Switch current directory based on project configuration<commit_after>// Copyright (c) 2012, mulander <netprobe@gmail.com>
// All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.
package main

import (
	"code.google.com/p/plotinum/plot"
	"code.google.com/p/plotinum/plotter"
	"code.google.com/p/plotinum/plotutil"
	"fmt"
	"log"
	"os/exec"
	"path"
	"regexp"
	"strconv"
	"strings"
	"sync"
	"time"
)

type Sampler interface {
	Sample(cmd *exec.Cmd, ticker *time.Ticker)
	GetData() []Data
	Stop()
}

type Data struct {
	label string
	data  []float64
}

type TopSampler struct {
	Samples []Data
	stop    chan bool
}

func (t *TopSampler) GetData() []Data {
	return t.Samples
}

func (t *TopSampler) Stop() {
	t.stop <- true
}

func (t *TopSampler) Sample(cmd *exec.Cmd, ticker *time.Ticker) {
	// %CPU(field=8) + %MEM(field=9)
	t.stop = make(chan bool)
	t.Samples = make([]Data, 2)
	t.Samples[0].label = "CPU"
	t.Samples[0].data = make([]float64, 1)
	t.Samples[1].label = "MEM"
	t.Samples[1].data = make([]float64, 1)
	raw := "(?m)%d.*$"
	r := regexp.MustCompile(fmt.Sprintf(raw, cmd.Process.Pid))
	for {
		select {
		case <-t.stop:
			return
		case <-ticker.C:
			top := exec.Command("top", "-b", "-n 1", fmt.Sprintf("-p %d", cmd.Process.Pid))
			log.Printf("Sampling the process")
			out, err := top.Output()
			if err != nil {
				log.Fatal(err)
			}
			fields := strings.Fields(r.FindString(fmt.Sprintf("%s", out)))
			if len(fields) != 0 {
				cpu, err := strconv.ParseFloat(fields[8], 64)
				if err != nil {
					log.Fatal(err)
				}
				mem, err := strconv.ParseFloat(fields[9], 64)
				if err != nil {
					log.Fatal(err)
				}
				t.Samples[0].data = append(t.Samples[0].data, cpu) // CPU
				t.Samples[1].data = append(t.Samples[1].data, mem) // MEM
				log.Printf("%+v", fields)
			}
		}
	}
}

type Project struct {
	Command    *exec.Cmd     // executable to run during the test
	Directory  string        // working directory for running the project
	Duration   time.Duration // duration of a single sample
	Sampling   time.Duration // trigger sampling based on this interval
	Executions int           // the total number of runs for a single test
	Sampler    Sampler
}

func (p *Project) Exec(samplers chan []Data) {
	err := p.Command.Start()
	if err != nil {
		log.Fatal(err)
	}

	ticker := time.NewTicker(p.Sampling)

	// Possibly more samplers in the future
	var wg sync.WaitGroup

	wg.Add(1)
	go func(n int) {
		defer wg.Done()
		p.Sampler.Sample(p.Command, ticker)
		samplers <- p.Sampler.GetData()
	}(1)
	go func() {
		wg.Wait()
		close(samplers)
	}()

	done := make(chan error)
	go func() {
		done <- p.Command.Wait()
	}()
	select {
	case <-time.After(p.Duration):
		if err := p.Command.Process.Kill(); err != nil {
			log.Fatal("Failed to kill: ", err)
		}
		<-done
		log.Println("Process killed")
	case err := <-done:
		log.Printf("Process done with error = %v", err)
	}
	log.Printf("Waiting for the ticker to stop")
	ticker.Stop()
	log.Printf("Stopping samplers")
	p.Sampler.Stop()
}

func main() {
	project := &Project{
		Command:    exec.Command("yes"),
		Duration:   10 * time.Second,
		Sampling:   1 * time.Second,
		Executions: 1,
		Sampler:    &TopSampler{}}

	// Change the working directory if needed
	if project.Directory != "" {
		project.Command.Dir = project.Directory
		log.Printf("Set project directory to: %s", project.Directory)
	}

	samplers := make(chan []Data)

	project.Exec(samplers)

	log.Printf("Waiting for the samplers")
	for s := range samplers {
		log.Printf("%+v", s)

		for _, sample := range s {

			p, err := plot.New()
			if err != nil {
				log.Fatal(err)
			}

			p.Title.Text = fmt.Sprintf("%s graph", sample.label)
			p.X.Label.Text = fmt.Sprintf("tick per %s", project.Sampling)
			p.Y.Label.Text = sample.label

			pts := make(plotter.XYs, len(sample.data))
			for i := range pts {
				pts[i].X = float64(i)
				pts[i].Y = sample.data[i]
			}

			plotutil.AddLinePoints(p, sample.label, pts)
			plotFile := fmt.Sprintf("%s-%s.png", path.Base(project.Command.Args[0]), sample.label)

			if err := p.Save(4, 4, plotFile); err != nil {
				log.Fatal(err)
			}
		}
	}
}
<|endoftext|>