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

package pgsql

import (
	"time"

	cerrors "github.com/coreos/clair/utils/errors"
)

// Lock tries to set a temporary lock in the database.
//
// Lock does not block, instead, it returns true and its expiration time
// is the lock has been successfully acquired or false otherwise
func (pgSQL *pgSQL) Lock(name string, owner string, duration time.Duration, renew bool) (bool, time.Time) {
	if name == "" || owner == "" || duration == 0 {
		log.Warning("could not create an invalid lock")
		return false, time.Time{}
	}

	defer observeQueryTime("Lock", "all", time.Now())

	// Prune locks.
	pgSQL.pruneLocks()

	// Compute expiration.
	until := time.Now().Add(duration)

	if renew {
		// Renew lock.
		r, err := pgSQL.Exec(getQuery("u_lock"), name, owner, until)
		if err != nil {
			handleError("u_lock", err)
			return false, until
		}
		if n, _ := r.RowsAffected(); n > 0 {
			// Updated successfully.
			return true, until
		}
	}

	// Lock.
	_, err := pgSQL.Exec(getQuery("i_lock"), name, owner, until)
	if err != nil {
		if !isErrUniqueViolation(err) {
			handleError("i_lock", err)
		}
		return false, until
	}

	return true, until
}

// Unlock unlocks a lock specified by its name if I own it
func (pgSQL *pgSQL) Unlock(name, owner string) {
	if name == "" || owner == "" {
		log.Warning("could not delete an invalid lock")
		return
	}

	defer observeQueryTime("Unlock", "all", time.Now())

	pgSQL.Exec(getQuery("r_lock"), name, owner)
}

// FindLock returns the owner of a lock specified by its name and its
// expiration time.
func (pgSQL *pgSQL) FindLock(name string) (string, time.Time, error) {
	if name == "" {
		log.Warning("could not find an invalid lock")
		return "", time.Time{}, cerrors.NewBadRequestError("could not find an invalid lock")
	}

	defer observeQueryTime("FindLock", "all", time.Now())

	var owner string
	var until time.Time
	err := pgSQL.QueryRow(getQuery("f_lock"), name).Scan(&owner, &until)
	if err != nil {
		return owner, until, handleError("f_lock", err)
	}

	return owner, until, nil
}

// pruneLocks removes every expired locks from the database
func (pgSQL *pgSQL) pruneLocks() {
	defer observeQueryTime("pruneLocks", "all", time.Now())

	if _, err := pgSQL.Exec(getQuery("r_lock_expired")); err != nil {
		handleError("r_lock_expired", err)
	}
}
<commit_msg>database: don't prune locks when we renew one<commit_after>// Copyright 2015 clair authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pgsql

import (
	"time"

	cerrors "github.com/coreos/clair/utils/errors"
)

// Lock tries to set a temporary lock in the database.
//
// Lock does not block, instead, it returns true and its expiration time
// is the lock has been successfully acquired or false otherwise
func (pgSQL *pgSQL) Lock(name string, owner string, duration time.Duration, renew bool) (bool, time.Time) {
	if name == "" || owner == "" || duration == 0 {
		log.Warning("could not create an invalid lock")
		return false, time.Time{}
	}

	defer observeQueryTime("Lock", "all", time.Now())

	// Compute expiration.
	until := time.Now().Add(duration)

	if renew {
		// Renew lock.
		r, err := pgSQL.Exec(getQuery("u_lock"), name, owner, until)
		if err != nil {
			handleError("u_lock", err)
			return false, until
		}
		if n, _ := r.RowsAffected(); n > 0 {
			// Updated successfully.
			return true, until
		}
	} else {
		// Prune locks.
		pgSQL.pruneLocks()
	}

	// Lock.
	_, err := pgSQL.Exec(getQuery("i_lock"), name, owner, until)
	if err != nil {
		if !isErrUniqueViolation(err) {
			handleError("i_lock", err)
		}
		return false, until
	}

	return true, until
}

// Unlock unlocks a lock specified by its name if I own it
func (pgSQL *pgSQL) Unlock(name, owner string) {
	if name == "" || owner == "" {
		log.Warning("could not delete an invalid lock")
		return
	}

	defer observeQueryTime("Unlock", "all", time.Now())

	pgSQL.Exec(getQuery("r_lock"), name, owner)
}

// FindLock returns the owner of a lock specified by its name and its
// expiration time.
func (pgSQL *pgSQL) FindLock(name string) (string, time.Time, error) {
	if name == "" {
		log.Warning("could not find an invalid lock")
		return "", time.Time{}, cerrors.NewBadRequestError("could not find an invalid lock")
	}

	defer observeQueryTime("FindLock", "all", time.Now())

	var owner string
	var until time.Time
	err := pgSQL.QueryRow(getQuery("f_lock"), name).Scan(&owner, &until)
	if err != nil {
		return owner, until, handleError("f_lock", err)
	}

	return owner, until, nil
}

// pruneLocks removes every expired locks from the database
func (pgSQL *pgSQL) pruneLocks() {
	defer observeQueryTime("pruneLocks", "all", time.Now())

	if _, err := pgSQL.Exec(getQuery("r_lock_expired")); err != nil {
		handleError("r_lock_expired", err)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	"github.com/lib/pq"
)

func connectToDatabase() {
	var err error
	driver := "postgres"

	connect := fmt.Sprintf("dbname='%s' user='%s' password='%s' host='%s' port='%s' sslmode='%s' connect_timeout='%s'",
		SomaCfg.Database.Name,
		SomaCfg.Database.User,
		SomaCfg.Database.Pass,
		SomaCfg.Database.Host,
		SomaCfg.Database.Port,
		SomaCfg.Database.TlsMode,
		SomaCfg.Database.Timeout,
	)

	// enable handling of infinity timestamps
	pq.EnableInfinityTs(NegTimeInf, PosTimeInf)

	conn, err = sql.Open(driver, connect)
	if err != nil {
		log.Fatal(err)
	}
	if err = conn.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Print("Connected to database")
	if _, err = conn.Exec(`SET TIME ZONE 'UTC';`); err != nil {
		log.Fatal(err)
	}
}

func pingDatabase() {
	ticker := time.NewTicker(time.Second).C

	for {
		<-ticker
		err := conn.Ping()
		if err != nil {
			log.Print(err)
		}
	}
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<commit_msg>Add DB schema version checks on startup<commit_after>package main

import (
	"database/sql"
	"fmt"
	"log"
	"time"

	"github.com/lib/pq"
)

func connectToDatabase() {
	var err error
	var rows *sql.Rows
	var schema string
	var schemaVer int64

	driver := "postgres"

	connect := fmt.Sprintf("dbname='%s' user='%s' password='%s' host='%s' port='%s' sslmode='%s' connect_timeout='%s'",
		SomaCfg.Database.Name,
		SomaCfg.Database.User,
		SomaCfg.Database.Pass,
		SomaCfg.Database.Host,
		SomaCfg.Database.Port,
		SomaCfg.Database.TlsMode,
		SomaCfg.Database.Timeout,
	)

	// enable handling of infinity timestamps
	pq.EnableInfinityTs(NegTimeInf, PosTimeInf)

	conn, err = sql.Open(driver, connect)
	if err != nil {
		log.Fatal(err)
	}
	if err = conn.Ping(); err != nil {
		log.Fatal(err)
	}
	log.Print("Connected to database")
	if _, err = conn.Exec(`SET TIME ZONE 'UTC';`); err != nil {
		log.Fatal(err)
	}

	// required schema versions
	required := map[string]int64{
		"inventory": 201605060001,
		"root":      201605150001,
		"auth":      201605150002,
		"soma":      201605060001,
	}

	if rows, err = conn.Query(`
SELECT schema,
       MAX(version) AS version
FROM   public.schema_versions
GROUP  BY schema;`); err != nil {
		log.Fatal("Query db schema versions: ", err)
	}

	for rows.Next() {
		if err = rows.Scan(
			&schema,
			&schemaVer,
		); err != nil {
			log.Fatal("Schema check: ", err)
		}
		if rsv, ok := required[schema]; ok {
			if rsv != required[schema] {
				log.Fatal("Incompatible schema: ", schema, schemaVer)
			} else {
				log.Printf("DB Schema %s, version: %d", schema, schemaVer)
				delete(required, schema)
			}
		} else {
			log.Fatal("Unknown schema: ", schema)
		}
	}
	if err = rows.Err(); err != nil {
		log.Fatal("Schema check: ", err)
	}
	if len(required) != 0 {
		for s, _ := range required {
			log.Printf("Missing schema: %s", s)
		}
		log.Fatal("FATAL - database incomplete")
	}
}

func pingDatabase() {
	ticker := time.NewTicker(time.Second).C

	for {
		<-ticker
		err := conn.Ping()
		if err != nil {
			log.Print(err)
		}
	}
}

// vim: ts=4 sw=4 sts=4 noet fenc=utf-8 ffs=unix
<|endoftext|>


<commit_before>package docker_loader

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	dClient "github.com/docker/docker/client"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/loaders"
)

const (
	loggingPrefix = "[docker_loader] "
	watchInterval = 30 * time.Second
	loaderName    = "docker"
)

func init() {
	loaders.Register(loaderName, func() loaders.TargetLoader {
		return &dockerLoader{
			cfg:         new(cfg),
			wg:          new(sync.WaitGroup),
			lastTargets: make(map[string]*collector.TargetConfig),
			logger:      log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
		}
	})
}

type dockerLoader struct {
	cfg         *cfg
	client      *dClient.Client
	wg          *sync.WaitGroup
	lastTargets map[string]*collector.TargetConfig
	logger      *log.Logger
	fl          []*targetFilterComp
}

type targetFilterComp struct {
	fl   []filters.Args
	nt   filters.Args
	port string
}

type cfg struct {
	Address  string          `json:"address,omitempty" mapstructure:"address,omitempty"`
	Interval time.Duration   `json:"interval,omitempty" mapstructure:"interval,omitempty"`
	Timeout  time.Duration   `json:"timeout,omitempty" mapstructure:"timeout,omitempty"`
	Filters  []*targetFilter `json:"filters,omitempty" mapstructure:"filters,omitempty"`
	Debug    bool            `json:"debug,omitempty" mapstructure:"debug,omitempty"`
}

type targetFilter struct {
	Containers []map[string]string `json:"containers,omitempty" mapstructure:"containers,omitempty"`
	Network    map[string]string   `json:"network,omitempty" mapstructure:"network,omitempty"`
	Port       string              `json:"port,omitempty" mapstructure:"port,omitempty"`
}

func (d *dockerLoader) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	err := loaders.DecodeConfig(cfg, d.cfg)
	if err != nil {
		return err
	}
	if d.cfg.Interval <= 0 {
		d.cfg.Interval = watchInterval
	}
	if d.cfg.Timeout <= 0 || d.cfg.Timeout >= d.cfg.Interval {
		d.cfg.Timeout = d.cfg.Interval / 2
	}
	if len(d.cfg.Filters) == 0 {
		d.cfg.Filters = []*targetFilter{
			{
				Containers: []map[string]string{
					{"status": "running"},
				},
			},
		}
	}
	d.fl = make([]*targetFilterComp, 0, len(d.cfg.Filters))
	for _, fm := range d.cfg.Filters {
		// network filter
		nflt := filters.NewArgs()
		for k, v := range fm.Network {
			nflt.Add(k, v)
		}
		// container filters
		cflt := make([]filters.Args, 0, len(fm.Containers))
		for _, sfm := range fm.Containers {
			flt := filters.NewArgs(filters.KeyValuePair{
				Key:   "status",
				Value: "running",
			})
			for k, v := range sfm {
				flt.Add(k, v)
			}
			cflt = append(cflt, flt)
		}
		// target filters
		d.fl = append(d.fl, &targetFilterComp{
			fl:   cflt,
			nt:   nflt,
			port: fm.Port,
		})
	}

	d.client, err = d.createDockerClient()
	if err != nil {
		return err
	}
	if logger != nil {
		d.logger.SetOutput(logger.Writer())
		d.logger.SetFlags(logger.Flags())
	}
	d.logger.Printf("initialized loader type %q: %s", loaderName, d)
	return nil
}

func (d *dockerLoader) createDockerClient() (*dClient.Client, error) {
	var opts []dClient.Opt
	if d.cfg.Address == "" {
		opts = []dClient.Opt{
			dClient.FromEnv,
			dClient.WithTimeout(d.cfg.Timeout),
		}
	} else {
		opts = []dClient.Opt{
			dClient.WithAPIVersionNegotiation(),
			dClient.WithHost(d.cfg.Address),
			dClient.WithTimeout(d.cfg.Timeout),
		}
	}
	return dClient.NewClientWithOpts(opts...)
}

func (d *dockerLoader) Start(ctx context.Context) chan *loaders.TargetOperation {
	fmt.Println("Starting docker loader")
	opChan := make(chan *loaders.TargetOperation)
	go func() {
		defer close(opChan)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				d.logger.Printf("querying %q targets", loaderName)
				readTargets, err := d.getTargets(ctx)
				if _, ok := err.(*os.PathError); ok {
					time.Sleep(d.cfg.Interval)
					continue
				}
				if err != nil {
					d.logger.Printf("failed to read targets from docker daemon: %v", err)
					time.Sleep(d.cfg.Interval)
					continue
				}
				select {
				case <-ctx.Done():
					return
				case opChan <- d.diff(readTargets):
					time.Sleep(d.cfg.Interval)
				}
			}
		}
	}()
	return opChan
}

func (d *dockerLoader) getTargets(ctx context.Context) (map[string]*collector.TargetConfig, error) {
	d.wg = new(sync.WaitGroup)
	d.wg.Add(len(d.fl))
	readTargets := make(map[string]*collector.TargetConfig)
	m := new(sync.Mutex)
	errChan := make(chan error, len(d.fl))
	for _, targetFilter := range d.fl {
		go func(fl *targetFilterComp) {
			defer d.wg.Done()
			// get networks
			nrs, err := d.client.NetworkList(ctx, types.NetworkListOptions{
				Filters: fl.nt,
			})
			if err != nil {
				errChan <- fmt.Errorf("failed getting networks list using filter %+v: %v", fl.nt, err)
				return
			}

			for _, cfl := range fl.fl {
				conts, err := d.client.ContainerList(ctx, types.ContainerListOptions{
					Filters: cfl,
				})
				if err != nil {
					errChan <- fmt.Errorf("failed getting containers list using filter %+v: %v", cfl, err)
					return
				}
				for _, cont := range conts {
					name := cont.ID
					if len(cont.Names) > 0 {
						name = strings.TrimLeft(cont.Names[0], "/")
					}
					if d.cfg.Debug {
						d.logger.Printf("filter %v returned container %v", cfl, name)
					}
					tc := new(collector.TargetConfig)
					tc.Name = name
					switch strings.ToLower(cont.HostConfig.NetworkMode) {
					case "host":
						if d.cfg.Address == "" || strings.HasPrefix(d.cfg.Address, "unix://") {
							tc.Address = "localhost"
						} else {
							tc.Address, _, err = net.SplitHostPort(d.cfg.Address)
							if err != nil {
								errChan <- err
								continue
							}
						}
						if fl.port != "" {
							if !strings.Contains(fl.port, "=") {
								tc.Address = fmt.Sprintf("%s:%s", tc.Address, fl.port)
							} else {
								portLabel := strings.Replace(fl.port, "label=", "", 1)
								if p, ok := cont.Labels[portLabel]; ok {
									tc.Address = fmt.Sprintf("%s:%s", tc.Address, p)
								}
							}
						}
					default:
						for _, nr := range nrs {
							if n, ok := cont.NetworkSettings.Networks[nr.Name]; ok {
								if n.IPAddress != "" {
									tc.Address = n.IPAddress
									break
								}
								tc.Address = n.GlobalIPv6Address
								break
							}
						}
						if tc.Address == "" {
							d.logger.Printf("%q no address found", tc.Name)
							continue
						}
						if fl.port != "" {
							if !strings.Contains(fl.port, "=") {
								tc.Address = fmt.Sprintf("%s:%s", tc.Address, fl.port)
							} else {
								portLabel := strings.Replace(fl.port, "label=", "", 1)
								if p, ok := cont.Labels[portLabel]; ok {
									tc.Address = fmt.Sprintf("%s:%s", tc.Address, p)
								}
							}
						}
					}
					//
					if d.cfg.Debug {
						d.logger.Printf("discoved target config %s with filter: %v", tc, cfl)
					}
					m.Lock()
					readTargets[tc.Name] = tc
					m.Unlock()
				}
			}
		}(targetFilter)
	}
	var errors = make([]error, 0)
	go func() {
		for err := range errChan {
			errors = append(errors, err)
		}
	}()

	d.wg.Wait()
	close(errChan)
	if len(errors) > 0 {
		for _, err := range errors {
			d.logger.Printf("%v", err)
		}
		return nil, fmt.Errorf("there was %d error(s)", len(errors))
	}
	return readTargets, nil
}

func (d *dockerLoader) diff(m map[string]*collector.TargetConfig) *loaders.TargetOperation {
	result := loaders.Diff(d.lastTargets, m)
	for _, t := range result.Add {
		if _, ok := d.lastTargets[t.Name]; !ok {
			d.lastTargets[t.Name] = t
		}
	}
	for _, n := range result.Del {
		delete(d.lastTargets, n)
	}
	if d.cfg.Debug {
		b, err := json.MarshalIndent(result, "", "  ")
		if err != nil {
			d.logger.Printf("discovery diff result: %v", result)
		} else {
			d.logger.Printf("discovery diff result:\n%s", string(b))
		}
	}
	return result
}

func (d *dockerLoader) String() string {
	b, err := json.Marshal(d.cfg)
	if err != nil {
		return fmt.Sprintf("%+v", d.cfg)
	}
	return string(b)
}
<commit_msg>add config per filter<commit_after>package docker_loader

import (
	"context"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/docker/docker/api/types"
	"github.com/docker/docker/api/types/filters"
	dClient "github.com/docker/docker/client"
	"github.com/karimra/gnmic/collector"
	"github.com/karimra/gnmic/loaders"
	"github.com/mitchellh/mapstructure"
)

const (
	loggingPrefix = "[docker_loader] "
	watchInterval = 30 * time.Second
	loaderName    = "docker"
)

func init() {
	loaders.Register(loaderName, func() loaders.TargetLoader {
		return &dockerLoader{
			cfg:         new(cfg),
			wg:          new(sync.WaitGroup),
			lastTargets: make(map[string]*collector.TargetConfig),
			logger:      log.New(ioutil.Discard, loggingPrefix, log.LstdFlags|log.Lmicroseconds),
		}
	})
}

type dockerLoader struct {
	cfg         *cfg
	client      *dClient.Client
	wg          *sync.WaitGroup
	lastTargets map[string]*collector.TargetConfig
	logger      *log.Logger
	fl          []*targetFilterComp
}

type targetFilterComp struct {
	fl   []filters.Args
	nt   filters.Args
	port string
	cfg  map[string]interface{}
}

type cfg struct {
	Address  string          `json:"address,omitempty" mapstructure:"address,omitempty"`
	Interval time.Duration   `json:"interval,omitempty" mapstructure:"interval,omitempty"`
	Timeout  time.Duration   `json:"timeout,omitempty" mapstructure:"timeout,omitempty"`
	Filters  []*targetFilter `json:"filters,omitempty" mapstructure:"filters,omitempty"`
	Debug    bool            `json:"debug,omitempty" mapstructure:"debug,omitempty"`
}

type targetFilter struct {
	Containers []map[string]string    `json:"containers,omitempty" mapstructure:"containers,omitempty"`
	Network    map[string]string      `json:"network,omitempty" mapstructure:"network,omitempty"`
	Port       string                 `json:"port,omitempty" mapstructure:"port,omitempty"`
	Config     map[string]interface{} `json:"config,omitempty" mapstructure:"config,omitempty"`
}

func (d *dockerLoader) Init(ctx context.Context, cfg map[string]interface{}, logger *log.Logger) error {
	err := loaders.DecodeConfig(cfg, d.cfg)
	if err != nil {
		return err
	}
	if d.cfg.Interval <= 0 {
		d.cfg.Interval = watchInterval
	}
	if d.cfg.Timeout <= 0 || d.cfg.Timeout >= d.cfg.Interval {
		d.cfg.Timeout = d.cfg.Interval / 2
	}
	if len(d.cfg.Filters) == 0 {
		d.cfg.Filters = []*targetFilter{
			{
				Containers: []map[string]string{
					{"status": "running"},
				},
			},
		}
	}
	d.fl = make([]*targetFilterComp, 0, len(d.cfg.Filters))
	for _, fm := range d.cfg.Filters {
		// network filter
		nflt := filters.NewArgs()
		for k, v := range fm.Network {
			nflt.Add(k, v)
		}
		// container filters
		cflt := make([]filters.Args, 0, len(fm.Containers))
		for _, sfm := range fm.Containers {
			flt := filters.NewArgs(filters.KeyValuePair{
				Key:   "status",
				Value: "running",
			})
			for k, v := range sfm {
				flt.Add(k, v)
			}
			cflt = append(cflt, flt)
		}
		// target filters
		d.fl = append(d.fl, &targetFilterComp{
			fl:   cflt,
			nt:   nflt,
			port: fm.Port,
			cfg:  fm.Config,
		})
	}

	d.client, err = d.createDockerClient()
	if err != nil {
		return err
	}
	if logger != nil {
		d.logger.SetOutput(logger.Writer())
		d.logger.SetFlags(logger.Flags())
	}
	d.logger.Printf("initialized loader type %q: %s", loaderName, d)
	return nil
}

func (d *dockerLoader) createDockerClient() (*dClient.Client, error) {
	var opts []dClient.Opt
	if d.cfg.Address == "" {
		opts = []dClient.Opt{
			dClient.FromEnv,
			dClient.WithTimeout(d.cfg.Timeout),
		}
	} else {
		opts = []dClient.Opt{
			dClient.WithAPIVersionNegotiation(),
			dClient.WithHost(d.cfg.Address),
			dClient.WithTimeout(d.cfg.Timeout),
		}
	}
	return dClient.NewClientWithOpts(opts...)
}

func (d *dockerLoader) Start(ctx context.Context) chan *loaders.TargetOperation {
	opChan := make(chan *loaders.TargetOperation)
	go func() {
		defer close(opChan)
		for {
			select {
			case <-ctx.Done():
				return
			default:
				d.logger.Printf("querying %q targets", loaderName)
				readTargets, err := d.getTargets(ctx)
				if _, ok := err.(*os.PathError); ok {
					time.Sleep(d.cfg.Interval)
					continue
				}
				if err != nil {
					d.logger.Printf("failed to read targets from docker daemon: %v", err)
					time.Sleep(d.cfg.Interval)
					continue
				}
				select {
				case <-ctx.Done():
					return
				case opChan <- d.diff(readTargets):
					time.Sleep(d.cfg.Interval)
				}
			}
		}
	}()
	return opChan
}

func (d *dockerLoader) getTargets(ctx context.Context) (map[string]*collector.TargetConfig, error) {
	d.wg = new(sync.WaitGroup)
	d.wg.Add(len(d.fl))
	readTargets := make(map[string]*collector.TargetConfig)
	m := new(sync.Mutex)
	errChan := make(chan error, len(d.fl))
	for _, targetFilter := range d.fl {
		go func(fl *targetFilterComp) {
			defer d.wg.Done()
			// get networks
			nrs, err := d.client.NetworkList(ctx, types.NetworkListOptions{
				Filters: fl.nt,
			})
			if err != nil {
				errChan <- fmt.Errorf("failed getting networks list using filter %+v: %v", fl.nt, err)
				return
			}

			for _, cfl := range fl.fl {
				conts, err := d.client.ContainerList(ctx, types.ContainerListOptions{
					Filters: cfl,
				})
				if err != nil {
					errChan <- fmt.Errorf("failed getting containers list using filter %+v: %v", cfl, err)
					return
				}
				for _, cont := range conts {
					name := cont.ID
					if len(cont.Names) > 0 {
						name = strings.TrimLeft(cont.Names[0], "/")
					}
					if d.cfg.Debug {
						d.logger.Printf("filter %v returned container %v", cfl, name)
					}
					tc := new(collector.TargetConfig)
					if fl.cfg != nil {
						err = mapstructure.Decode(fl.cfg, tc)
						if err != nil {
							d.logger.Printf("failed to decode config map: %v", err)
						}
						if d.cfg.Debug {
							d.logger.Printf("target config before adding name and address: %v", tc)
						}
					}
					tc.Name = name
					switch strings.ToLower(cont.HostConfig.NetworkMode) {
					case "host":
						if d.cfg.Address == "" || strings.HasPrefix(d.cfg.Address, "unix://") {
							tc.Address = "localhost"
						} else {
							tc.Address, _, err = net.SplitHostPort(d.cfg.Address)
							if err != nil {
								errChan <- err
								continue
							}
						}
						if fl.port != "" {
							if !strings.Contains(fl.port, "=") {
								tc.Address = fmt.Sprintf("%s:%s", tc.Address, fl.port)
							} else {
								portLabel := strings.Replace(fl.port, "label=", "", 1)
								if p, ok := cont.Labels[portLabel]; ok {
									tc.Address = fmt.Sprintf("%s:%s", tc.Address, p)
								}
							}
						}
					default:
						for _, nr := range nrs {
							if n, ok := cont.NetworkSettings.Networks[nr.Name]; ok {
								if n.IPAddress != "" {
									tc.Address = n.IPAddress
									break
								}
								tc.Address = n.GlobalIPv6Address
								break
							}
						}
						if tc.Address == "" {
							d.logger.Printf("%q no address found", tc.Name)
							continue
						}
						if fl.port != "" {
							if !strings.Contains(fl.port, "=") {
								tc.Address = fmt.Sprintf("%s:%s", tc.Address, fl.port)
							} else {
								portLabel := strings.Replace(fl.port, "label=", "", 1)
								if p, ok := cont.Labels[portLabel]; ok {
									tc.Address = fmt.Sprintf("%s:%s", tc.Address, p)
								}
							}
						}
					}
					//
					if d.cfg.Debug {
						d.logger.Printf("discoved target config %s with filter: %v", tc, cfl)
					}
					m.Lock()
					readTargets[tc.Name] = tc
					m.Unlock()
				}
			}
		}(targetFilter)
	}
	var errors = make([]error, 0)
	go func() {
		for err := range errChan {
			errors = append(errors, err)
		}
	}()

	d.wg.Wait()
	close(errChan)
	if len(errors) > 0 {
		for _, err := range errors {
			d.logger.Printf("%v", err)
		}
		return nil, fmt.Errorf("there was %d error(s)", len(errors))
	}
	return readTargets, nil
}

func (d *dockerLoader) diff(m map[string]*collector.TargetConfig) *loaders.TargetOperation {
	result := loaders.Diff(d.lastTargets, m)
	for _, t := range result.Add {
		if _, ok := d.lastTargets[t.Name]; !ok {
			d.lastTargets[t.Name] = t
		}
	}
	for _, n := range result.Del {
		delete(d.lastTargets, n)
	}
	if d.cfg.Debug {
		b, err := json.MarshalIndent(result, "", "  ")
		if err != nil {
			d.logger.Printf("discovery diff result: %v", result)
		} else {
			d.logger.Printf("discovery diff result:\n%s", string(b))
		}
	}
	return result
}

func (d *dockerLoader) String() string {
	b, err := json.Marshal(d.cfg)
	if err != nil {
		return fmt.Sprintf("%+v", d.cfg)
	}
	return string(b)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"koding/kites/kontrol/kontrol"

	"github.com/koding/multiconfig"
)

func main() {
	m := multiconfig.New()

	conf := new(kontrol.Config)

	// Load the config, it's reads from the file, environment variables and
	// lastly from flags in order
	if err := m.Load(conf); err != nil {
		panic(err)
	}

	fmt.Printf("Kontrol loaded with following variables: %+v\n", conf)

	k := kontrol.New(conf)
	k.Run()
}
<commit_msg>Kontrol: run artifact server inside of the kontrol<commit_after>package main

import (
	"fmt"
	"koding/artifact"
	"koding/kites/kontrol/kontrol"

	"github.com/koding/multiconfig"
)

var Name = "kontrol"

func main() {
	m := multiconfig.New()

	conf := new(kontrol.Config)

	// Load the config, it's reads from the file, environment variables and
	// lastly from flags in order
	if err := m.Load(conf); err != nil {
		panic(err)
	}

	fmt.Printf("Kontrol loaded with following variables: %+v\n", conf)

	k := kontrol.New(conf)

	// TODO use kite's http server instead of creating another one here
	// this is used for application lifecycle management
	go artifact.StartDefaultServer(Name, conf.ArtifactPort)

	k.Run()
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Vitess Authors.

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

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

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

package topotests

import (
	"fmt"
	"strings"
	"testing"

	"github.com/stretchr/testify/require"

	"context"

	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/memorytopo"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// This file tests the CellInfo part of the topo.Server API.

func TestCellInfo(t *testing.T) {
	cell := "cell1"
	ctx := context.Background()
	ts := memorytopo.NewServer(cell)

	// Check GetCellInfo returns what memorytopo created.
	ci, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.Root != "" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	var cells []string
	cells, err = ts.ExpandCells(ctx, cell)
	require.NoError(t, err)
	require.EqualValues(t, []string{"cell1"}, cells)

	// Update the Server Address.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "new address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test update with no change.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "bad address"
		return topo.NewError(topo.NoUpdateNeeded, cell)
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test failing update.
	updateErr := fmt.Errorf("inside error")
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		return updateErr
	}); err != updateErr {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}

	// Test update on non-existing object.
	newCell := "new_cell"
	if err := ts.UpdateCellInfoFields(ctx, newCell, func(ci *topodatapb.CellInfo) error {
		ci.Root = "/"
		ci.ServerAddress = "good address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, newCell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "good address" || ci.Root != "/" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Add a record that should block CellInfo deletion for safety reasons.
	if err := ts.UpdateSrvKeyspace(ctx, cell, "keyspace", &topodatapb.SrvKeyspace{}); err != nil {
		t.Fatalf("UpdateSrvKeyspace failed: %v", err)
	}
	srvKeyspaces, err := ts.GetSrvKeyspaceNames(ctx, cell)
	if err != nil {
		t.Fatalf("GetSrvKeyspaceNames failed: %v", err)
	}
	if len(srvKeyspaces) == 0 {
		t.Fatalf("UpdateSrvKeyspace did not add SrvKeyspace.")
	}

	// Try to delete without force; it should fail.
	if err := ts.DeleteCellInfo(ctx, cell, false); err == nil {
		t.Fatalf("DeleteCellInfo should have failed without -force")
	}

	// Use the force.
	if err := ts.DeleteCellInfo(ctx, cell, true); err != nil {
		t.Fatalf("DeleteCellInfo failed even with -force: %v", err)
	}
	if _, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/); !topo.IsErrType(err, topo.NoNode) {
		t.Fatalf("GetCellInfo(non-existing cell) failed: %v", err)
	}
}

func TestExpandCells(t *testing.T) {
	ctx := context.Background()
	var cells []string
	var err error
	var allCells = "cell1,cell2,cell3"
	type testCase struct {
		name      string
		cellsIn   string
		cellsOut  []string
		errString string
	}

	testCases := []testCase{
		{"single", "cell1", []string{"cell1"}, ""},
		{"multiple", "cell1,cell2,cell3", []string{"cell1", "cell2", "cell3"}, ""},
		{"empty", "", []string{"cell1", "cell2", "cell3"}, ""},
		{"bad", "unknown", nil, "node doesn't exist"},
	}

	for _, tCase := range testCases {
		t.Run(tCase.name, func(t *testing.T) {
			cellsIn := tCase.cellsIn
			if cellsIn == "" {
				cellsIn = allCells
			}
			topoCells := strings.Split(cellsIn, ",")
			var ts *topo.Server
			if tCase.name == "bad" {
				ts = memorytopo.NewServer()
			} else {
				ts = memorytopo.NewServer(topoCells...)
			}
			cells, err = ts.ExpandCells(ctx, cellsIn)
			if tCase.errString != "" {
				require.Error(t, err)
				require.Contains(t, err.Error(), tCase.errString)
			} else {
				require.NoError(t, err)
			}
			require.EqualValues(t, tCase.cellsOut, cells)
		})
	}

}
<commit_msg>Add failing tests<commit_after>/*
Copyright 2019 The Vitess Authors.

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

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

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

package topotests

import (
	"fmt"
	"strings"
	"testing"

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

	"context"

	"vitess.io/vitess/go/vt/topo"
	"vitess.io/vitess/go/vt/topo/memorytopo"

	topodatapb "vitess.io/vitess/go/vt/proto/topodata"
)

// This file tests the CellInfo part of the topo.Server API.

func TestCellInfo(t *testing.T) {
	cell := "cell1"
	ctx := context.Background()
	ts := memorytopo.NewServer(cell)

	// Check GetCellInfo returns what memorytopo created.
	ci, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.Root != "" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	var cells []string
	cells, err = ts.ExpandCells(ctx, cell)
	require.NoError(t, err)
	require.EqualValues(t, []string{"cell1"}, cells)

	// Update the Server Address.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "new address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test update with no change.
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		ci.ServerAddress = "bad address"
		return topo.NewError(topo.NoUpdateNeeded, cell)
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, cell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "new address" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Test failing update.
	updateErr := fmt.Errorf("inside error")
	if err := ts.UpdateCellInfoFields(ctx, cell, func(ci *topodatapb.CellInfo) error {
		return updateErr
	}); err != updateErr {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}

	// Test update on non-existing object.
	newCell := "new_cell"
	if err := ts.UpdateCellInfoFields(ctx, newCell, func(ci *topodatapb.CellInfo) error {
		ci.Root = "/"
		ci.ServerAddress = "good address"
		return nil
	}); err != nil {
		t.Fatalf("UpdateCellInfoFields failed: %v", err)
	}
	ci, err = ts.GetCellInfo(ctx, newCell, true /*strongRead*/)
	if err != nil {
		t.Fatalf("GetCellInfo failed: %v", err)
	}
	if ci.ServerAddress != "good address" || ci.Root != "/" {
		t.Fatalf("unexpected CellInfo: %v", ci)
	}

	// Add a record that should block CellInfo deletion for safety reasons.
	if err := ts.UpdateSrvKeyspace(ctx, cell, "keyspace", &topodatapb.SrvKeyspace{}); err != nil {
		t.Fatalf("UpdateSrvKeyspace failed: %v", err)
	}
	srvKeyspaces, err := ts.GetSrvKeyspaceNames(ctx, cell)
	if err != nil {
		t.Fatalf("GetSrvKeyspaceNames failed: %v", err)
	}
	if len(srvKeyspaces) == 0 {
		t.Fatalf("UpdateSrvKeyspace did not add SrvKeyspace.")
	}

	// Try to delete without force; it should fail.
	if err := ts.DeleteCellInfo(ctx, cell, false); err == nil {
		t.Fatalf("DeleteCellInfo should have failed without -force")
	}

	// Use the force.
	if err := ts.DeleteCellInfo(ctx, cell, true); err != nil {
		t.Fatalf("DeleteCellInfo failed even with -force: %v", err)
	}
	if _, err := ts.GetCellInfo(ctx, cell, true /*strongRead*/); !topo.IsErrType(err, topo.NoNode) {
		t.Fatalf("GetCellInfo(non-existing cell) failed: %v", err)
	}
}

func TestExpandCells(t *testing.T) {
	ctx := context.Background()
	var cells []string
	var err error
	var allCells = "cell1,cell2,cell3"
	type testCase struct {
		name      string
		cellsIn   string
		cellsOut  []string
		errString string
	}

	testCases := []testCase{
		{"single", "cell1", []string{"cell1"}, ""},
		{"multiple", "cell1,cell2,cell3", []string{"cell1", "cell2", "cell3"}, ""},
		{"empty", "", []string{"cell1", "cell2", "cell3"}, ""},
		{"bad", "unknown", nil, "node doesn't exist"},
	}

	for _, tCase := range testCases {
		t.Run(tCase.name, func(t *testing.T) {
			cellsIn := tCase.cellsIn
			if cellsIn == "" {
				cellsIn = allCells
			}
			topoCells := strings.Split(cellsIn, ",")
			var ts *topo.Server
			if tCase.name == "bad" {
				ts = memorytopo.NewServer()
			} else {
				ts = memorytopo.NewServer(topoCells...)
			}
			cells, err = ts.ExpandCells(ctx, cellsIn)
			if tCase.errString != "" {
				require.Error(t, err)
				require.Contains(t, err.Error(), tCase.errString)
			} else {
				require.NoError(t, err)
			}
			require.EqualValues(t, tCase.cellsOut, cells)
		})
	}

	t.Run("aliases", func(t *testing.T) {
		cells := []string{"cell1", "cell2", "cell3"}
		ts := memorytopo.NewServer(cells...)
		err := ts.CreateCellsAlias(ctx, "alias", &topodatapb.CellsAlias{Cells: cells})
		require.NoError(t, err)

		tests := []struct {
			name      string
			in        string
			out       []string
			shouldErr bool
		}{
			{
				name: "alias only",
				in:   "alias",
				out:  []string{"cell1", "cell2", "cell3"},
			},
			{
				name: "alias and cell in alias", // test deduping logic
				in:   "alias,cell1",
				out:  []string{"cell1", "cell2", "cell3"},
			},
			{
				name: "just cells",
				in:   "cell1",
				out:  []string{"cell1"},
			},
			{
				name:      "missing alias",
				in:        "not_an_alias",
				shouldErr: true,
			},
		}

		for _, tt := range tests {
			tt := tt
			t.Run(tt.name, func(t *testing.T) {
				expanded, err := ts.ExpandCells(ctx, tt.in)
				if tt.shouldErr {
					assert.Error(t, err)
					return
				}

				require.NoError(t, err)
				assert.ElementsMatch(t, expanded, tt.out)
			})
		}
	})
}
<|endoftext|>


<commit_before>package main

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"time"
)

const socksTimeout = 2
const bufSiz = 1500

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

func proxy(local *net.TCPConn, ws *websocket.Conn) {
	var localToWs chan bool
	var wsToLocal chan bool

	// Local-to-WebSocket read loop.
	localToWs = make(chan bool, 1)
	go func() {
		buf := make([]byte, bufSiz)
		var err error
		for {
			n, er := local.Read(buf[:])
			if n > 0 {
				ew := websocket.Message.Send(ws, buf[:n])
				if ew != nil {
					err = ew
					break
				}
			}
			if er != nil {
				err = er
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseRead()
		ws.Close()

		localToWs <- true
	}()

	// WebSocket-to-local read loop.
	wsToLocal = make(chan bool, 1)
	go func() {
		var buf []byte
		var err error
		for {
			er := websocket.Message.Receive(ws, &buf)
			if er != nil {
				err = er
				break
			}
			n, ew := local.Write(buf)
			if ew != nil {
				err = ew
				break
			}
			if n != len(buf) {
				err = io.ErrShortWrite
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseWrite()
		ws.Close()

		wsToLocal <- true
	}()

	// Select twice, once for each read loop.
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
	dest, err := readSocks4aConnect(conn)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	// Disable deadline.
	conn.SetDeadline(time.Time{})
	logDebug("SOCKS request for %s", dest)

	// We need the parsed IP and port for the SOCKS reply.
	destAddr, err := net.ResolveTCPAddr("tcp", dest)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}

	wsUrl := url.URL{Scheme: "ws", Host: dest}
	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	defer ws.Close()
	logDebug("WebSocket connection to %s", ws.Config().Location.String())

	sendSocks4aResponseGranted(conn, destAddr)

	proxy(conn, ws)

	return nil
}

func socksAcceptLoop(ln *net.TCPListener) error {
	for {
		socks, err := ln.AcceptTCP()
		if err != nil {
			return err
		}
		go func() {
			err := handleConnection(socks)
			if err != nil {
				logDebug("SOCKS from %s: %s", socks.RemoteAddr(), err)
			}
		}()
	}
	return nil
}

func startListener(addrStr string) (*net.TCPListener, error) {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		err := socksAcceptLoop(ln)
		if err != nil {
			logDebug("accept: %s", err)
		}
	}()
	return ln, nil
}

func main() {
	const ptMethodName = "websocket"
	var socksAddrStrs = [...]string{"127.0.0.1:0", "[::1]:0"}

	ptClientSetup([]string{ptMethodName})

	for _, socksAddrStr := range socksAddrStrs {
		ln, err := startListener(socksAddrStr)
		if err != nil {
			ptCmethodError(ptMethodName, err.Error())
		}
		ptCmethod(ptMethodName, "socks4", ln.Addr())
	}
	ptCmethodsDone()

	select {}
}
<commit_msg>Do an explicit wait for SIGINT.<commit_after>package main

import (
	"code.google.com/p/go.net/websocket"
	"fmt"
	"io"
	"net"
	"net/url"
	"os"
	"os/signal"
	"time"
)

const socksTimeout = 2
const bufSiz = 1500

func logDebug(format string, v ...interface{}) {
	fmt.Fprintf(os.Stderr, format+"\n", v...)
}

func proxy(local *net.TCPConn, ws *websocket.Conn) {
	var localToWs chan bool
	var wsToLocal chan bool

	// Local-to-WebSocket read loop.
	localToWs = make(chan bool, 1)
	go func() {
		buf := make([]byte, bufSiz)
		var err error
		for {
			n, er := local.Read(buf[:])
			if n > 0 {
				ew := websocket.Message.Send(ws, buf[:n])
				if ew != nil {
					err = ew
					break
				}
			}
			if er != nil {
				err = er
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseRead()
		ws.Close()

		localToWs <- true
	}()

	// WebSocket-to-local read loop.
	wsToLocal = make(chan bool, 1)
	go func() {
		var buf []byte
		var err error
		for {
			er := websocket.Message.Receive(ws, &buf)
			if er != nil {
				err = er
				break
			}
			n, ew := local.Write(buf)
			if ew != nil {
				err = ew
				break
			}
			if n != len(buf) {
				err = io.ErrShortWrite
				break
			}
		}
		if err != nil && err != io.EOF {
			logDebug("%s", err)
		}
		local.CloseWrite()
		ws.Close()

		wsToLocal <- true
	}()

	// Select twice, once for each read loop.
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
	select {
	case <-localToWs:
	case <-wsToLocal:
	}
}

func handleConnection(conn *net.TCPConn) error {
	defer conn.Close()

	conn.SetDeadline(time.Now().Add(socksTimeout * time.Second))
	dest, err := readSocks4aConnect(conn)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	// Disable deadline.
	conn.SetDeadline(time.Time{})
	logDebug("SOCKS request for %s", dest)

	// We need the parsed IP and port for the SOCKS reply.
	destAddr, err := net.ResolveTCPAddr("tcp", dest)
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}

	wsUrl := url.URL{Scheme: "ws", Host: dest}
	ws, err := websocket.Dial(wsUrl.String(), "", wsUrl.String())
	if err != nil {
		sendSocks4aResponseFailed(conn)
		return err
	}
	defer ws.Close()
	logDebug("WebSocket connection to %s", ws.Config().Location.String())

	sendSocks4aResponseGranted(conn, destAddr)

	proxy(conn, ws)

	return nil
}

func socksAcceptLoop(ln *net.TCPListener) error {
	for {
		socks, err := ln.AcceptTCP()
		if err != nil {
			return err
		}
		go func() {
			err := handleConnection(socks)
			if err != nil {
				logDebug("SOCKS from %s: %s", socks.RemoteAddr(), err)
			}
		}()
	}
	return nil
}

func startListener(addrStr string) (*net.TCPListener, error) {
	addr, err := net.ResolveTCPAddr("tcp", addrStr)
	if err != nil {
		return nil, err
	}
	ln, err := net.ListenTCP("tcp", addr)
	if err != nil {
		return nil, err
	}
	go func() {
		err := socksAcceptLoop(ln)
		if err != nil {
			logDebug("accept: %s", err)
		}
	}()
	return ln, nil
}

func main() {
	const ptMethodName = "websocket"
	var socksAddrStrs = [...]string{"127.0.0.1:0", "[::1]:0"}

	ptClientSetup([]string{ptMethodName})

	for _, socksAddrStr := range socksAddrStrs {
		ln, err := startListener(socksAddrStr)
		if err != nil {
			ptCmethodError(ptMethodName, err.Error())
		}
		ptCmethod(ptMethodName, "socks4", ln.Addr())
	}
	ptCmethodsDone()

	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, os.Interrupt)
	<-signalChan
}
<|endoftext|>


<commit_before>package device

import (
	"bytes"
	"fmt"
	"sync/atomic"
	"time"
)

// Statistics represents a set of device statistics.
type Statistics interface {
	// BytesReceived returns the total bytes received since this instance was created
	BytesReceived() uint32

	// AddBytesReceived adds a certain number of bytes to the BytesReceived count.
	// Implementations will always be safe for concurrent access.
	AddBytesReceived(uint32)

	// MessagesReceived returns the total messages received since this instance was created
	MessagesReceived() uint32

	// AddMessagesReceived adds a certain number of messages to the MessagesReceived count.
	// Implementations will always be safe for concurrent access.
	AddMessagesReceived(uint32)

	// BytesSent returns the total bytes sent since this instance was created
	BytesSent() uint32

	// AddBytesSent adds a certain number of bytes to the BytesSent count.
	// Implementations will always be safe for concurrent access.
	AddBytesSent(uint32)

	// MessagesSent returns the total messages sent since this instance was created
	MessagesSent() uint32

	// AddMessagesSent adds a certain number of messages to the MessagesSent count.
	// Implementations will always be safe for concurrent access.
	AddMessagesSent(uint32)

	// ConnectedAt returns the connection time at which this statistics began tracking
	ConnectedAt() time.Time
}

// NewStatistics creates a Statistics instance with the given connection time
func NewStatistics(connectedAt time.Time) Statistics {
	return &statistics{
		connectedAt: connectedAt,
	}
}

// statistics is the internal Statistics implementation
type statistics struct {
	bytesReceived    uint32
	bytesSent        uint32
	messagesReceived uint32
	messagesSent     uint32
	connectedAt      time.Time
}

func (s *statistics) BytesReceived() uint32 {
	return atomic.LoadUint32(&s.bytesReceived)
}

func (s *statistics) AddBytesReceived(delta uint32) {
	atomic.AddUint32(&s.bytesReceived, delta)
}

func (s *statistics) MessagesReceived() uint32 {
	return atomic.LoadUint32(&s.messagesReceived)
}

func (s *statistics) AddMessagesReceived(delta uint32) {
	atomic.AddUint32(&s.messagesReceived, delta)
}

func (s *statistics) BytesSent() uint32 {
	return atomic.LoadUint32(&s.bytesSent)
}

func (s *statistics) AddBytesSent(delta uint32) {
	atomic.AddUint32(&s.bytesSent, delta)
}

func (s *statistics) MessagesSent() uint32 {
	return atomic.LoadUint32(&s.messagesSent)
}

func (s *statistics) AddMessagesSent(delta uint32) {
	atomic.AddUint32(&s.messagesSent, delta)
}

func (s *statistics) ConnectedAt() time.Time {
	return s.connectedAt
}

func (s *statistics) MarshalJSON() ([]byte, error) {
	output := bytes.NewBuffer(make([]byte, 0, 150))
	fmt.Fprintf(
		output,
		`{"bytesSent": %d, "messagesSent": %d, "bytesReceived": %d, "messagesReceived": %d, "connectedAt": "%s"}`,
		s.BytesSent(),
		s.MessagesSent(),
		s.BytesReceived(),
		s.MessagesReceived(),
		s.ConnectedAt().Format(time.RFC3339),
	)

	return output.Bytes(), nil
}
<commit_msg>added String() for debug output<commit_after>package device

import (
	"bytes"
	"fmt"
	"sync/atomic"
	"time"
)

// Statistics represents a set of device statistics.
type Statistics interface {
	// BytesReceived returns the total bytes received since this instance was created
	BytesReceived() uint32

	// AddBytesReceived adds a certain number of bytes to the BytesReceived count.
	// Implementations will always be safe for concurrent access.
	AddBytesReceived(uint32)

	// MessagesReceived returns the total messages received since this instance was created
	MessagesReceived() uint32

	// AddMessagesReceived adds a certain number of messages to the MessagesReceived count.
	// Implementations will always be safe for concurrent access.
	AddMessagesReceived(uint32)

	// BytesSent returns the total bytes sent since this instance was created
	BytesSent() uint32

	// AddBytesSent adds a certain number of bytes to the BytesSent count.
	// Implementations will always be safe for concurrent access.
	AddBytesSent(uint32)

	// MessagesSent returns the total messages sent since this instance was created
	MessagesSent() uint32

	// AddMessagesSent adds a certain number of messages to the MessagesSent count.
	// Implementations will always be safe for concurrent access.
	AddMessagesSent(uint32)

	// ConnectedAt returns the connection time at which this statistics began tracking
	ConnectedAt() time.Time
}

// NewStatistics creates a Statistics instance with the given connection time
func NewStatistics(connectedAt time.Time) Statistics {
	return &statistics{
		connectedAt: connectedAt,
	}
}

// statistics is the internal Statistics implementation
type statistics struct {
	bytesReceived    uint32
	bytesSent        uint32
	messagesReceived uint32
	messagesSent     uint32
	connectedAt      time.Time
}

func (s *statistics) BytesReceived() uint32 {
	return atomic.LoadUint32(&s.bytesReceived)
}

func (s *statistics) AddBytesReceived(delta uint32) {
	atomic.AddUint32(&s.bytesReceived, delta)
}

func (s *statistics) MessagesReceived() uint32 {
	return atomic.LoadUint32(&s.messagesReceived)
}

func (s *statistics) AddMessagesReceived(delta uint32) {
	atomic.AddUint32(&s.messagesReceived, delta)
}

func (s *statistics) BytesSent() uint32 {
	return atomic.LoadUint32(&s.bytesSent)
}

func (s *statistics) AddBytesSent(delta uint32) {
	atomic.AddUint32(&s.bytesSent, delta)
}

func (s *statistics) MessagesSent() uint32 {
	return atomic.LoadUint32(&s.messagesSent)
}

func (s *statistics) AddMessagesSent(delta uint32) {
	atomic.AddUint32(&s.messagesSent, delta)
}

func (s *statistics) ConnectedAt() time.Time {
	return s.connectedAt
}

func (s *statistics) String() string {
	data, _ := s.MarshalJSON()
	return string(data)
}

func (s *statistics) MarshalJSON() ([]byte, error) {
	output := bytes.NewBuffer(make([]byte, 0, 150))
	fmt.Fprintf(
		output,
		`{"bytesSent": %d, "messagesSent": %d, "bytesReceived": %d, "messagesReceived": %d, "connectedAt": "%s"}`,
		s.BytesSent(),
		s.MessagesSent(),
		s.BytesReceived(),
		s.MessagesReceived(),
		s.ConnectedAt().Format(time.RFC3339),
	)

	return output.Bytes(), nil
}
<|endoftext|>


<commit_before>package types

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strings"
)

var reg = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_.]+$`)

type TimeSeriesData struct {
	Metric    string            `json:"metric"`    //sys.cpu.idle
	DataType  string            `json:"data_type"` //COUNTER,GAUGE,DERIVE
	Value     float64           `json:"value"`     //99.00
	Timestamp int64             `json:"timestamp"` //unix timestamp
	Cycle     int               `json:"cycle,omitempty"`
	Tags      map[string]string `json:"tags"` //{"product":"app01", "group":"dev02"}
}

func (m *TimeSeriesData) Validate() bool {
	if !reg.MatchString(m.Metric) || m.Metric == "" {
		return false
	}
	switch strings.ToLower(m.DataType) {
	case "gauge", "counter", "derive":
	default:
		return false
	}
	return true
}

func (tsd *TimeSeriesData) Encode() []byte {
	data, _ := json.Marshal(tsd)
	return data
}

func (tsd *TimeSeriesData) Decode(data []byte) error {
	return json.Unmarshal(data, &tsd)
}

func (tsd TimeSeriesData) String() string {
	return fmt.Sprintf("{metric:%s, data_type:%s, value:%.2f, time:%d, cycle:%d, tags:%s}",
		tsd.Metric,
		tsd.DataType,
		tsd.Value,
		tsd.Timestamp,
		tsd.Cycle,
		tsd.Tags2String(),
	)
}

func (tsd *TimeSeriesData) Tags2String() string {
	if len(tsd.Tags) == 0 {
		return ""
	}
	taglen := len(tsd.Tags)
	keys := make([]string, taglen)
	i := 0
	for k := range tsd.Tags {
		keys[i] = k
		i++
	}
	sort.Strings(keys)
	ret := ""
	for _, k := range keys {
		ret += fmt.Sprintf("%s=%s,", strings.TrimSpace(k), strings.TrimSpace(tsd.Tags[k]))
	}
	return strings.Trim(ret, ",")
}

func (tsd *TimeSeriesData) PK() string {
	return fmt.Sprintf("%s.%s", tsd.Metric, tsd.Tags2String())
}

func (tsd *TimeSeriesData) GetMetric() string {
	metric := tsd.Metric
	if len(tsd.Tags2String()) > 0 {
		metric = fmt.Sprintf("%s/%s", metric, tsd.Tags2String())
	}
	return metric
}

func (tsd *TimeSeriesData) AddTags(tags map[string]string) {
	if tsd.Tags == nil {
		tsd.Tags = tags
		return
	}
	for k, v := range tags {
		tsd.Tags[k] = v
	}
}

//tag1=v1,tag2=v2,tag3=v3
//{"tag1":v1,"tag2":v2,"tag3":v3}
func ParseTags(name string) map[string]string {
	res := make(map[string]string)
	kv := strings.Split(name, ",")
	for _, v := range kv {
		tmp := strings.Split(v, "=")
		if len(tmp) != 2 {
			continue
		}
		res[tmp[0]] = tmp[1]
	}
	return res
}
<commit_msg>修改正则规则 && 增加 V4 版本协议适配<commit_after>package types

import (
	"encoding/json"
	"fmt"
	"regexp"
	"sort"
	"strings"
)

var reg = regexp.MustCompile(`^[a-zA-Z][a-zA-Z0-9_.-]+$`)

type TimeSeriesDataV4 struct {
	PerformanceIndex map[string]float64 `json:"performance_index"`
	Host             string             `json:"host"`
	AppName          string             `json:"app_name"`
	Domain           string             `json:"domain"`
}

func (tsdv4 *TimeSeriesDataV4) ToV5() []*TimeSeriesData {
	ts := []*TimeSeriesData{}
	for k, v := range tsdv4.PerformanceIndex {
		ts = append(ts, &TimeSeriesData{
			Metric:   fmt.Sprintf("%s.%s", strings.ToLower(tsdv4.AppName), strings.ToLower(k)),
			DataType: "gauge",
			Value:    v,
			Tags:     map[string]string{},
		})
	}
	return ts
}

type TimeSeriesData struct {
	Metric    string            `json:"metric"`    //sys.cpu.idle
	DataType  string            `json:"data_type"` //COUNTER,GAUGE,DERIVE
	Value     float64           `json:"value"`     //99.00
	Timestamp int64             `json:"timestamp"` //unix timestamp
	Cycle     int               `json:"cycle,omitempty"`
	Tags      map[string]string `json:"tags"` //{"product":"app01", "group":"dev02"}
}

func (m *TimeSeriesData) Validate() bool {
	if !reg.MatchString(m.Metric) || m.Metric == "" {
		return false
	}
	switch strings.ToLower(m.DataType) {
	case "gauge", "counter", "derive":
	default:
		return false
	}
	return true
}

func (tsd *TimeSeriesData) Encode() []byte {
	data, _ := json.Marshal(tsd)
	return data
}

func (tsd *TimeSeriesData) Decode(data []byte) error {
	return json.Unmarshal(data, &tsd)
}

func (tsd TimeSeriesData) String() string {
	return fmt.Sprintf("{metric:%s, data_type:%s, value:%.2f, time:%d, cycle:%d, tags:%s}",
		tsd.Metric,
		tsd.DataType,
		tsd.Value,
		tsd.Timestamp,
		tsd.Cycle,
		tsd.Tags2String(),
	)
}

func (tsd *TimeSeriesData) Tags2String() string {
	if len(tsd.Tags) == 0 {
		return ""
	}
	taglen := len(tsd.Tags)
	keys := make([]string, taglen)
	i := 0
	for k := range tsd.Tags {
		keys[i] = k
		i++
	}
	sort.Strings(keys)
	ret := ""
	for _, k := range keys {
		ret += fmt.Sprintf("%s=%s,", strings.TrimSpace(k), strings.TrimSpace(tsd.Tags[k]))
	}
	return strings.Trim(ret, ",")
}

func (tsd *TimeSeriesData) PK() string {
	return fmt.Sprintf("%s.%s", tsd.Metric, tsd.Tags2String())
}

func (tsd *TimeSeriesData) GetMetric() string {
	metric := tsd.Metric
	if len(tsd.Tags2String()) > 0 {
		metric = fmt.Sprintf("%s/%s", metric, tsd.Tags2String())
	}
	return metric
}

func (tsd *TimeSeriesData) AddTags(tags map[string]string) {
	if tsd.Tags == nil {
		tsd.Tags = tags
		return
	}
	for k, v := range tags {
		tsd.Tags[k] = v
	}
}

func (tsd *TimeSeriesData) RemoveTag(tagk string) {
	if tsd.Tags == nil {
		return
	}
	delete(tsd.Tags, tagk)
}

//tag1=v1,tag2=v2,tag3=v3
//{"tag1":v1,"tag2":v2,"tag3":v3}
func ParseTags(name string) map[string]string {
	res := make(map[string]string)
	kv := strings.Split(name, ",")
	for _, v := range kv {
		tmp := strings.Split(v, "=")
		if len(tmp) != 2 {
			continue
		}
		res[tmp[0]] = tmp[1]
	}
	return res
}
<|endoftext|>


<commit_before>package ftp

import (
	"bytes"
	"io/ioutil"
	"net/textproto"
	"testing"
	"time"
)

const (
	testData = "Just some text"
	testDir  = "mydir"
)

func TestConnPASV(t *testing.T) {
	testConn(t, true)
}

func TestConnEPSV(t *testing.T) {
	testConn(t, false)
}

func testConn(t *testing.T, passive bool) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("localhost:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if passive {
		delete(c.features, "EPSV")
	}

	err = c.Login("anonymous", "anonymous")
	if err != nil {
		t.Fatal(err)
	}

	err = c.NoOp()
	if err != nil {
		t.Error(err)
	}

	err = c.ChangeDir("incoming")
	if err != nil {
		t.Error(err)
	}

	data := bytes.NewBufferString(testData)
	err = c.Stor("test", data)
	if err != nil {
		t.Error(err)
	}

	_, err = c.List(".")
	if err != nil {
		t.Error(err)
	}

	err = c.Rename("test", "tset")
	if err != nil {
		t.Error(err)
	}

	r, err := c.Retr("tset")
	if err != nil {
		t.Error(err)
	} else {
		buf, err := ioutil.ReadAll(r)
		if err != nil {
			t.Error(err)
		}
		if string(buf) != testData {
			t.Errorf("'%s'", buf)
		}
		r.Close()
	}

	r, err = c.RetrFrom("tset", 5)
	if err != nil {
		t.Error(err)
	} else {
		buf, err := ioutil.ReadAll(r)
		if err != nil {
			t.Error(err)
		}
		expected := testData[5:]
		if string(buf) != expected {
			t.Errorf("read %q, expected %q", buf, expected)
		}
		r.Close()
	}

	err = c.Delete("tset")
	if err != nil {
		t.Error(err)
	}

	err = c.MakeDir(testDir)
	if err != nil {
		t.Error(err)
	}

	err = c.ChangeDir(testDir)
	if err != nil {
		t.Error(err)
	}

	dir, err := c.CurrentDir()
	if err != nil {
		t.Error(err)
	} else {
		if dir != "/incoming/"+testDir {
			t.Error("Wrong dir: " + dir)
		}
	}

	err = c.ChangeDirToParent()
	if err != nil {
		t.Error(err)
	}

	entries, err := c.NameList("/")
	if err != nil {
		t.Error(err)
	}
	if len(entries) != 1 || entries[0] != "/incoming" {
		t.Errorf("Unexpected entries: %v", entries)
	}

	err = c.RemoveDir(testDir)
	if err != nil {
		t.Error(err)
	}

	err = c.Logout()
	if err != nil {
		if protoErr := err.(*textproto.Error); protoErr != nil {
			if protoErr.Code != StatusNotImplemented {
				t.Error(err)
			}
		} else {
			t.Error(err)
		}
	}

	c.Quit()

	err = c.NoOp()
	if err == nil {
		t.Error("Expected error")
	}
}

// ftp.mozilla.org uses multiline 220 response
func TestMultiline(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("ftp.mozilla.org:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	err = c.Login("anonymous", "anonymous")
	if err != nil {
		t.Fatal(err)
	}

	_, err = c.List(".")
	if err != nil {
		t.Error(err)
	}

	c.Quit()
}

// antioche.antioche.eu.org with IPv6
func TestConnIPv6(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("[::1]:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	err = c.Login("anonymous", "anonymous")
	if err != nil {
		t.Fatal(err)
	}

	_, err = c.List(".")
	if err != nil {
		t.Error(err)
	}

	c.Quit()
}
<commit_msg>Add tests for timeout and wrong login.<commit_after>package ftp

import (
	"bytes"
	"io/ioutil"
	"net/textproto"
	"testing"
	"time"
)

const (
	testData = "Just some text"
	testDir  = "mydir"
)

func TestConnPASV(t *testing.T) {
	testConn(t, true)
}

func TestConnEPSV(t *testing.T) {
	testConn(t, false)
}

func testConn(t *testing.T, passive bool) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("localhost:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	if passive {
		delete(c.features, "EPSV")
	}

	err = c.Login("anonymous", "anonymous")
	if err != nil {
		t.Fatal(err)
	}

	err = c.NoOp()
	if err != nil {
		t.Error(err)
	}

	err = c.ChangeDir("incoming")
	if err != nil {
		t.Error(err)
	}

	data := bytes.NewBufferString(testData)
	err = c.Stor("test", data)
	if err != nil {
		t.Error(err)
	}

	_, err = c.List(".")
	if err != nil {
		t.Error(err)
	}

	err = c.Rename("test", "tset")
	if err != nil {
		t.Error(err)
	}

	r, err := c.Retr("tset")
	if err != nil {
		t.Error(err)
	} else {
		buf, err := ioutil.ReadAll(r)
		if err != nil {
			t.Error(err)
		}
		if string(buf) != testData {
			t.Errorf("'%s'", buf)
		}
		r.Close()
	}

	r, err = c.RetrFrom("tset", 5)
	if err != nil {
		t.Error(err)
	} else {
		buf, err := ioutil.ReadAll(r)
		if err != nil {
			t.Error(err)
		}
		expected := testData[5:]
		if string(buf) != expected {
			t.Errorf("read %q, expected %q", buf, expected)
		}
		r.Close()
	}

	err = c.Delete("tset")
	if err != nil {
		t.Error(err)
	}

	err = c.MakeDir(testDir)
	if err != nil {
		t.Error(err)
	}

	err = c.ChangeDir(testDir)
	if err != nil {
		t.Error(err)
	}

	dir, err := c.CurrentDir()
	if err != nil {
		t.Error(err)
	} else {
		if dir != "/incoming/"+testDir {
			t.Error("Wrong dir: " + dir)
		}
	}

	err = c.ChangeDirToParent()
	if err != nil {
		t.Error(err)
	}

	entries, err := c.NameList("/")
	if err != nil {
		t.Error(err)
	}
	if len(entries) != 1 || entries[0] != "/incoming" {
		t.Errorf("Unexpected entries: %v", entries)
	}

	err = c.RemoveDir(testDir)
	if err != nil {
		t.Error(err)
	}

	err = c.Logout()
	if err != nil {
		if protoErr := err.(*textproto.Error); protoErr != nil {
			if protoErr.Code != StatusNotImplemented {
				t.Error(err)
			}
		} else {
			t.Error(err)
		}
	}

	c.Quit()

	err = c.NoOp()
	if err == nil {
		t.Error("Expected error")
	}
}

// ftp.mozilla.org uses multiline 220 response
func TestMultiline(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("ftp.mozilla.org:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	err = c.Login("anonymous", "anonymous")
	if err != nil {
		t.Fatal(err)
	}

	_, err = c.List(".")
	if err != nil {
		t.Error(err)
	}

	c.Quit()
}

func TestConnIPv6(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("[::1]:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}

	err = c.Login("anonymous", "anonymous")
	if err != nil {
		t.Fatal(err)
	}

	_, err = c.List(".")
	if err != nil {
		t.Error(err)
	}

	c.Quit()
}

func TestTimeout(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("localhost:2121", 1*time.Second)
	if err == nil {
		t.Fatal("expected timeout, got nil error")
		c.Quit()
	}
}

func TestWrongLogin(t *testing.T) {
	if testing.Short() {
		t.Skip("skipping test in short mode.")
	}

	c, err := DialTimeout("localhost:21", 5*time.Second)
	if err != nil {
		t.Fatal(err)
	}
	defer c.Quit()

	err = c.Login("zoo2Shia", "fei5Yix9")
	if err == nil {
		t.Fatal("expected error, got nil")
	}
}

<|endoftext|>


<commit_before>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	exec "os/exec"
	"strconv"
	dstrings "strings"

	"github.com/golang/glog"
	gcli "github.com/heketi/heketi/client/api/go-client"
	gapi "github.com/heketi/heketi/pkg/glusterfs/api"
	"github.com/kubernetes-incubator/external-storage/lib/controller"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	//storage "k8s.io/client-go/pkg/apis/storage/v1beta1"
)

const (
	provisionerName  = "gluster.org/glusterblock"
	defaultExecPath  = "./createiscsi"
	secretKeyName    = "key"
	shareIDAnn       = "glusterBlockShare"
	provisionerIDAnn = "glusterBlockProvisionerIdentity"
	creatorAnn       = "kubernetes.io/createdby"
	volumeTypeAnn    = "gluster.org/type"
	descriptionAnn   = "Description"
)

type glusterBlockProvisioner struct {
	// Kubernetes Client. Use to retrieve Gluster admin secret
	client kubernetes.Interface

	// Identity of this glusterBlockProvisioner, generated. Used to identify "this"
	// provisioner's PVs.
	identity string

	// Configuration of gluster block provisioner
	provConfig provisionerConfig

	volConfig glusterBlockVolume
}

type provisionerConfig struct {
	// Required:  this is the Rest Service Url ( ex: Heketi) for Gluster Block
	url string

	// Optional: Rest user who is capable of creating gluster block volumes.
	user string

	// Optional: Rest user key for above RestUser.
	userKey string

	// Optional:  secret name, namespace.
	secretNamespace string
	secretName      string
	secretValue     string

	// Optional:  clusterID from which the provisioner create the block volume
	clusterID string

	// Optional: high availability count in case of multipathing
	haCount int

	// Optional: Operation mode  (rest, executable)
	opMode string

	// Optional: Executable path if we are operating in executable mode.
	execPath string
}

type glusterBlockVolume struct {
	TargetPortal      string
	Portals           []string
	IQN               string
	Lun               int32
	FSType            string
	ISCSIInterface    string
	DiscoveryCHAPAuth bool
	SessionCHAPAuth   bool
	ReadOnly          bool
}

func NewglusterBlockProvisioner(client kubernetes.Interface, id string) controller.Provisioner {
	return &glusterBlockProvisioner{
		client:   client,
		identity: id,
	}
}

var _ controller.Provisioner = &glusterBlockProvisioner{}

// Provision creates a storage asset and returns a PV object representing it.
func (p *glusterBlockProvisioner) Provision(options controller.VolumeOptions) (*v1.PersistentVolume, error) {

	var err error
	if options.PVC.Spec.Selector != nil {
		return nil, fmt.Errorf("claim Selector is not supported")
	}

	glog.V(4).Infof("glusterblock: VolumeOptions %v", options)

	// If we want to retrieve storage class name.
	// scName := storageutil.GetClaimStorageClass(r.options.PVC)

	cfg, err := parseClassParameters(options.Parameters, p.client)
	if err != nil {
		return nil, err
	}
	p.provConfig = *cfg

	glog.V(4).Infof("glusterfs: creating volume with configuration %+v", p.provConfig)

	vol, err := p.createVolume(options.PVName)
	if err != nil {
		return nil, err
	}
	glog.V(1).Infof("Target Portal and IQN returned :%v %v", vol.TargetPortal, vol.IQN)

	// Create unique PVC identity.
	blockVolIdentity := fmt.Sprintf("kubernetes-dynamic-pvc-%s", uuid.NewUUID())

	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: options.PVName,
			Annotations: map[string]string{
				provisionerIDAnn: p.identity,
				shareIDAnn:       blockVolIdentity,
				creatorAnn:       "heketi-dynamic-provisioner",
				volumeTypeAnn:    "block",
				descriptionAnn:   "Gluster: Dynamically provisioned PV",
			},
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeReclaimPolicy: options.PersistentVolumeReclaimPolicy,
			AccessModes:                   options.PVC.Spec.AccessModes,
			Capacity: v1.ResourceList{
				v1.ResourceName(v1.ResourceStorage): options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)],
			},
			PersistentVolumeSource: v1.PersistentVolumeSource{
				ISCSI: &v1.ISCSIVolumeSource{
					TargetPortal: vol.TargetPortal,
					IQN:          vol.IQN,
					Lun:          0,
					FSType:       "ext4",
					ReadOnly:     false,
				},
			},
		},
	}
	glog.Infof("successfully created Gluster Block volume %+v", pv.Spec.PersistentVolumeSource.ISCSI)
	return pv, nil
}

// createVolume creates a gluster block volume i.e. the storage asset.

func (p *glusterBlockProvisioner) createVolume(PVName string) (*glusterBlockVolume, error) {

	switch p.provConfig.opMode {

	case "executable":
		cmd := exec.Command("sh", p.provConfig.execPath)
		err := cmd.Run()
		if err != nil {
			glog.Errorf("%v", err)
		}

		dtarget := os.Getenv("TARGET")
		diqn := os.Getenv("IQN")

		p.volConfig.TargetPortal = dtarget
		p.volConfig.IQN = diqn

	case "rest":
		cli := gcli.NewClient(p.provConfig.url, p.provConfig.user, p.provConfig.secretValue)
		if cli == nil {
			glog.Errorf("glusterfs: failed to create glusterfs rest client")
			return nil, fmt.Errorf("glusterfs: failed to create glusterfs rest client, REST server authentication failed")
		}

		//TODO:
		sz := 1
		volumeReq := &gapi.VolumeCreateRequest{Size: sz}
		_, err := cli.VolumeCreate(volumeReq)
		if err != nil {
			glog.Errorf("glusterfs: error creating volume %v ", err)
			return nil, fmt.Errorf("error creating volume %v", err)
		}

		p.volConfig.IQN = "iqn.2016-12.org.gluster-block:aafea465-9167-4880-b37c-2c36db8562ea"
		p.volConfig.TargetPortal = "192.168.1.11"

	default:
		return nil, fmt.Errorf("error parsing value for 'opmode' for volume plugin %s", provisionerName)
	}
	return &p.volConfig, nil

	/*
		if p.provConfig.opMode == "executable" {
			cmd := exec.Command("sh", p.provConfig.execPath)
			err := cmd.Run()
			if err != nil {
				glog.Errorf("%v", err)
			}
			dtarget = os.Getenv("TARGET")
			diqn = os.Getenv("IQN")
		} */

}

// Delete removes the storage asset that was created by Provision represented
// by the given PV.
func (p *glusterBlockProvisioner) Delete(volume *v1.PersistentVolume) error {
	ann, ok := volume.Annotations[provisionerIDAnn]
	if !ok {
		return errors.New("identity annotation not found on PV")
	}
	if ann != p.identity {
		return &controller.IgnoredError{"identity annotation on PV does not match this provisioners identity"}
	}

	blockVol, ok := volume.Annotations[shareIDAnn]
	if !ok {
		return errors.New("gluster block share annotation not found on PV")
	}

	// Delete this blockVol
	glog.V(1).Infof("blockVolume  %v", blockVol)

	// Unset the variables.
	os.Setenv("TARGET", "")
	os.Setenv("IQN", "")

	return nil
}

func parseClassParameters(params map[string]string, kubeclient kubernetes.Interface) (*provisionerConfig, error) {
	var cfg provisionerConfig
	var err error

	authEnabled := true

	// Default multipath count has been set to 3
	haCount := 3

	for k, v := range params {
		switch dstrings.ToLower(k) {
		case "resturl":
			cfg.url = v
		case "restuser":
			cfg.user = v
		case "restuserkey":
			cfg.userKey = v
		case "secretname":
			cfg.secretName = v
		case "secretnamespace":
			cfg.secretNamespace = v
		case "clusterID":
			if len(v) != 0 {
				cfg.clusterID = v
			}
		case "restauthenabled":
			authEnabled = dstrings.ToLower(v) == "true"
		case "hacount":
			haCount, err = strconv.Atoi(v)
			if err != nil {
				return nil, fmt.Errorf("glusterblock: failed to parse hacount %v ", k)
			}
		case "opmode":
			cfg.opMode = v
		case "execpath":
			cfg.execPath = v
		default:
			return nil, fmt.Errorf("glusterblock: invalid option %q for volume plugin %s", k, "glusterblock")
		}
	}

	if len(cfg.url) == 0 {
		return nil, fmt.Errorf("StorageClass for provisioner %s must contain 'resturl' parameter", "glusterblock")
	}

	if haCount == 0 {
		cfg.haCount = haCount
	}

	if len(cfg.opMode) == 0 {
		cfg.opMode = "executable"
	}

	if len(cfg.execPath) == 0 {
		cfg.execPath = defaultExecPath
	}
	if !authEnabled {
		cfg.user = ""
		cfg.secretName = ""
		cfg.secretNamespace = ""
		cfg.userKey = ""
		cfg.secretValue = ""
	}

	if len(cfg.secretName) != 0 || len(cfg.secretNamespace) != 0 {
		// secretName + Namespace has precedence over userKey
		if len(cfg.secretName) != 0 && len(cfg.secretNamespace) != 0 {
			cfg.secretValue, err = parseSecret(cfg.secretNamespace, cfg.secretName, kubeclient)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, fmt.Errorf("StorageClass for provisioner %q must have secretNamespace and secretName either both set or both empty", "glusterblock")
		}
	} else {
		cfg.secretValue = cfg.userKey
	}

	return &cfg, nil
}

// parseSecret finds a given Secret instance and reads user password from it.
func parseSecret(namespace, secretName string, kubeClient kubernetes.Interface) (string, error) {

	secretMap, err := GetSecretForPV(namespace, secretName, "glusterblock", kubeClient)
	if err != nil {
		glog.Errorf("failed to get secret %s/%s: %v", namespace, secretName, err)
		return "", fmt.Errorf("failed to get secret %s/%s: %v", namespace, secretName, err)
	}
	if len(secretMap) == 0 {
		return "", fmt.Errorf("empty secret map")
	}
	secret := ""
	for k, v := range secretMap {
		if k == secretKeyName {
			return v, nil
		}
		secret = v
	}
	// If not found, the last secret in the map wins as done before
	return secret, nil
}

// GetSecretForPV locates secret by name and namespace, verifies the secret type, and returns secret map

func GetSecretForPV(secretNamespace, secretName, volumePluginName string, kubeClient kubernetes.Interface) (map[string]string, error) {
	secret := make(map[string]string)
	if kubeClient == nil {
		return secret, fmt.Errorf("Cannot get kube client")
	}
	secrets, err := kubeClient.Core().Secrets(secretNamespace).Get(secretName, metav1.GetOptions{})
	if err != nil {
		return secret, err
	}
	if secrets.Type != v1.SecretType(volumePluginName) {
		return secret, fmt.Errorf("Cannot get secret of type %s", volumePluginName)
	}
	for name, data := range secrets.Data {
		secret[name] = string(data)
	}
	return secret, nil
}

var (
	master     = flag.String("master", "", "Master URL")
	kubeconfig = flag.String("kubeconfig", "", "Absolute path to the kubeconfig")
	id         = flag.String("id", "", "Unique provisioner identity")
)

func main() {
	flag.Parse()
	flag.Set("logtostderr", "true")

	// Create an InClusterConfig and use it to create a client for the controller
	// to use to communicate with Kubernetes

	var config *rest.Config
	var err error
	if *master != "" || *kubeconfig != "" {
		config, err = clientcmd.BuildConfigFromFlags(*master, *kubeconfig)
	} else {
		config, err = rest.InClusterConfig()
	}

	prId := string(uuid.NewUUID())

	if *id != "" {
		prId = *id
	}

	if err != nil {
		glog.Fatalf("Failed to create config: %v", err)
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		glog.Fatalf("Failed to create client: %v", err)
	}

	// The controller needs to know what the server version is because out-of-tree
	// provisioners aren't officially supported until 1.5
	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		glog.Fatalf("Error getting server version: %v", err)
	}

	// Create the provisioner: it implements the Provisioner interface expected by
	// the controller
	glusterBlockProvisioner := NewglusterBlockProvisioner(clientset, prId)

	// Start the provision controller which will dynamically provision glusterblock
	// PVs

	pc := controller.NewProvisionController(
		clientset,
		provisionerName,
		glusterBlockProvisioner,
		serverVersion.GitVersion,
	)

	pc.Run(wait.NeverStop)
}
<commit_msg>Remove unused code from the provisioner.<commit_after>/*
Copyright 2017 The Kubernetes Authors.

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

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

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

package main

import (
	"errors"
	"flag"
	"fmt"
	"os"
	exec "os/exec"
	"strconv"
	dstrings "strings"

	"github.com/golang/glog"
	gcli "github.com/heketi/heketi/client/api/go-client"
	gapi "github.com/heketi/heketi/pkg/glusterfs/api"
	"github.com/kubernetes-incubator/external-storage/lib/controller"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/uuid"
	"k8s.io/apimachinery/pkg/util/wait"
	"k8s.io/client-go/kubernetes"
	"k8s.io/client-go/pkg/api/v1"
	"k8s.io/client-go/rest"
	"k8s.io/client-go/tools/clientcmd"
	//storage "k8s.io/client-go/pkg/apis/storage/v1beta1"
)

const (
	provisionerName  = "gluster.org/glusterblock"
	defaultExecPath  = "./createiscsi"
	secretKeyName    = "key"
	shareIDAnn       = "glusterBlockShare"
	provisionerIDAnn = "glusterBlockProvisionerIdentity"
	creatorAnn       = "kubernetes.io/createdby"
	volumeTypeAnn    = "gluster.org/type"
	descriptionAnn   = "Description"
)

type glusterBlockProvisioner struct {
	// Kubernetes Client. Use to retrieve Gluster admin secret
	client kubernetes.Interface

	// Identity of this glusterBlockProvisioner, generated. Used to identify "this"
	// provisioner's PVs.
	identity string

	// Configuration of gluster block provisioner
	provConfig provisionerConfig

	volConfig glusterBlockVolume
}

type provisionerConfig struct {
	// Required:  this is the Rest Service Url ( ex: Heketi) for Gluster Block
	url string

	// Optional: Rest user who is capable of creating gluster block volumes.
	user string

	// Optional: Rest user key for above RestUser.
	userKey string

	// Optional:  secret name, namespace.
	secretNamespace string
	secretName      string
	secretValue     string

	// Optional:  clusterID from which the provisioner create the block volume
	clusterID string

	// Optional: high availability count in case of multipathing
	haCount int

	// Optional: Operation mode  (rest, executable)
	opMode string

	// Optional: Executable path if we are operating in executable mode.
	execPath string
}

type glusterBlockVolume struct {
	TargetPortal      string
	Portals           []string
	IQN               string
	Lun               int32
	FSType            string
	ISCSIInterface    string
	DiscoveryCHAPAuth bool
	SessionCHAPAuth   bool
	ReadOnly          bool
}

func NewglusterBlockProvisioner(client kubernetes.Interface, id string) controller.Provisioner {
	return &glusterBlockProvisioner{
		client:   client,
		identity: id,
	}
}

var _ controller.Provisioner = &glusterBlockProvisioner{}

// Provision creates a storage asset and returns a PV object representing it.
func (p *glusterBlockProvisioner) Provision(options controller.VolumeOptions) (*v1.PersistentVolume, error) {

	var err error
	if options.PVC.Spec.Selector != nil {
		return nil, fmt.Errorf("claim Selector is not supported")
	}

	glog.V(4).Infof("glusterblock: VolumeOptions %v", options)

	// If we want to retrieve storage class name.
	// scName := storageutil.GetClaimStorageClass(r.options.PVC)

	cfg, err := parseClassParameters(options.Parameters, p.client)
	if err != nil {
		return nil, err
	}
	p.provConfig = *cfg

	glog.V(4).Infof("glusterfs: creating volume with configuration %+v", p.provConfig)

	vol, err := p.createVolume(options.PVName)
	if err != nil {
		return nil, err
	}
	glog.V(1).Infof("Target Portal and IQN returned :%v %v", vol.TargetPortal, vol.IQN)

	// Create unique PVC identity.
	blockVolIdentity := fmt.Sprintf("kubernetes-dynamic-pvc-%s", uuid.NewUUID())

	pv := &v1.PersistentVolume{
		ObjectMeta: metav1.ObjectMeta{
			Name: options.PVName,
			Annotations: map[string]string{
				provisionerIDAnn: p.identity,
				shareIDAnn:       blockVolIdentity,
				creatorAnn:       "heketi-dynamic-provisioner",
				volumeTypeAnn:    "block",
				descriptionAnn:   "Gluster: Dynamically provisioned PV",
			},
		},
		Spec: v1.PersistentVolumeSpec{
			PersistentVolumeReclaimPolicy: options.PersistentVolumeReclaimPolicy,
			AccessModes:                   options.PVC.Spec.AccessModes,
			Capacity: v1.ResourceList{
				v1.ResourceName(v1.ResourceStorage): options.PVC.Spec.Resources.Requests[v1.ResourceName(v1.ResourceStorage)],
			},
			PersistentVolumeSource: v1.PersistentVolumeSource{
				ISCSI: &v1.ISCSIVolumeSource{
					TargetPortal: vol.TargetPortal,
					IQN:          vol.IQN,
					Lun:          0,
					FSType:       "ext4",
					ReadOnly:     false,
				},
			},
		},
	}
	glog.Infof("successfully created Gluster Block volume %+v", pv.Spec.PersistentVolumeSource.ISCSI)
	return pv, nil
}

// createVolume creates a gluster block volume i.e. the storage asset.

func (p *glusterBlockProvisioner) createVolume(PVName string) (*glusterBlockVolume, error) {

	switch p.provConfig.opMode {

	case "executable":
		cmd := exec.Command("sh", p.provConfig.execPath)
		err := cmd.Run()
		if err != nil {
			glog.Errorf("%v", err)
		}

		dtarget := os.Getenv("TARGET")
		diqn := os.Getenv("IQN")

		p.volConfig.TargetPortal = dtarget
		p.volConfig.IQN = diqn

	case "rest":
		cli := gcli.NewClient(p.provConfig.url, p.provConfig.user, p.provConfig.secretValue)
		if cli == nil {
			glog.Errorf("glusterfs: failed to create glusterfs rest client")
			return nil, fmt.Errorf("glusterfs: failed to create glusterfs rest client, REST server authentication failed")
		}

		//TODO:
		sz := 1
		volumeReq := &gapi.VolumeCreateRequest{Size: sz}
		_, err := cli.VolumeCreate(volumeReq)
		if err != nil {
			glog.Errorf("glusterfs: error creating volume %v ", err)
			return nil, fmt.Errorf("error creating volume %v", err)
		}

		p.volConfig.IQN = "iqn.2016-12.org.gluster-block:aafea465-9167-4880-b37c-2c36db8562ea"
		p.volConfig.TargetPortal = "192.168.1.11"

	default:
		return nil, fmt.Errorf("error parsing value for 'opmode' for volume plugin %s", provisionerName)
	}
	return &p.volConfig, nil
}

// Delete removes the storage asset that was created by Provision represented
// by the given PV.
func (p *glusterBlockProvisioner) Delete(volume *v1.PersistentVolume) error {
	ann, ok := volume.Annotations[provisionerIDAnn]
	if !ok {
		return errors.New("identity annotation not found on PV")
	}
	if ann != p.identity {
		return &controller.IgnoredError{"identity annotation on PV does not match this provisioners identity"}
	}

	blockVol, ok := volume.Annotations[shareIDAnn]
	if !ok {
		return errors.New("gluster block share annotation not found on PV")
	}

	// Delete this blockVol
	glog.V(1).Infof("blockVolume  %v", blockVol)

	// Unset the variables.
	os.Setenv("TARGET", "")
	os.Setenv("IQN", "")

	return nil
}

func parseClassParameters(params map[string]string, kubeclient kubernetes.Interface) (*provisionerConfig, error) {
	var cfg provisionerConfig
	var err error

	authEnabled := true

	// Default multipath count has been set to 3
	haCount := 3

	for k, v := range params {
		switch dstrings.ToLower(k) {
		case "resturl":
			cfg.url = v
		case "restuser":
			cfg.user = v
		case "restuserkey":
			cfg.userKey = v
		case "secretname":
			cfg.secretName = v
		case "secretnamespace":
			cfg.secretNamespace = v
		case "clusterID":
			if len(v) != 0 {
				cfg.clusterID = v
			}
		case "restauthenabled":
			authEnabled = dstrings.ToLower(v) == "true"
		case "hacount":
			haCount, err = strconv.Atoi(v)
			if err != nil {
				return nil, fmt.Errorf("glusterblock: failed to parse hacount %v ", k)
			}
		case "opmode":
			cfg.opMode = v
		case "execpath":
			cfg.execPath = v
		default:
			return nil, fmt.Errorf("glusterblock: invalid option %q for volume plugin %s", k, "glusterblock")
		}
	}

	if len(cfg.url) == 0 {
		return nil, fmt.Errorf("StorageClass for provisioner %s must contain 'resturl' parameter", "glusterblock")
	}

	if haCount == 0 {
		cfg.haCount = haCount
	}

	if len(cfg.opMode) == 0 {
		cfg.opMode = "executable"
	}

	if len(cfg.execPath) == 0 {
		cfg.execPath = defaultExecPath
	}
	if !authEnabled {
		cfg.user = ""
		cfg.secretName = ""
		cfg.secretNamespace = ""
		cfg.userKey = ""
		cfg.secretValue = ""
	}

	if len(cfg.secretName) != 0 || len(cfg.secretNamespace) != 0 {
		// secretName + Namespace has precedence over userKey
		if len(cfg.secretName) != 0 && len(cfg.secretNamespace) != 0 {
			cfg.secretValue, err = parseSecret(cfg.secretNamespace, cfg.secretName, kubeclient)
			if err != nil {
				return nil, err
			}
		} else {
			return nil, fmt.Errorf("StorageClass for provisioner %q must have secretNamespace and secretName either both set or both empty", "glusterblock")
		}
	} else {
		cfg.secretValue = cfg.userKey
	}

	return &cfg, nil
}

// parseSecret finds a given Secret instance and reads user password from it.
func parseSecret(namespace, secretName string, kubeClient kubernetes.Interface) (string, error) {

	secretMap, err := GetSecretForPV(namespace, secretName, "glusterblock", kubeClient)
	if err != nil {
		glog.Errorf("failed to get secret %s/%s: %v", namespace, secretName, err)
		return "", fmt.Errorf("failed to get secret %s/%s: %v", namespace, secretName, err)
	}
	if len(secretMap) == 0 {
		return "", fmt.Errorf("empty secret map")
	}
	secret := ""
	for k, v := range secretMap {
		if k == secretKeyName {
			return v, nil
		}
		secret = v
	}
	// If not found, the last secret in the map wins as done before
	return secret, nil
}

// GetSecretForPV locates secret by name and namespace, verifies the secret type, and returns secret map

func GetSecretForPV(secretNamespace, secretName, volumePluginName string, kubeClient kubernetes.Interface) (map[string]string, error) {
	secret := make(map[string]string)
	if kubeClient == nil {
		return secret, fmt.Errorf("Cannot get kube client")
	}
	secrets, err := kubeClient.Core().Secrets(secretNamespace).Get(secretName, metav1.GetOptions{})
	if err != nil {
		return secret, err
	}
	if secrets.Type != v1.SecretType(volumePluginName) {
		return secret, fmt.Errorf("Cannot get secret of type %s", volumePluginName)
	}
	for name, data := range secrets.Data {
		secret[name] = string(data)
	}
	return secret, nil
}

var (
	master     = flag.String("master", "", "Master URL")
	kubeconfig = flag.String("kubeconfig", "", "Absolute path to the kubeconfig")
	id         = flag.String("id", "", "Unique provisioner identity")
)

func main() {
	flag.Parse()
	flag.Set("logtostderr", "true")

	// Create an InClusterConfig and use it to create a client for the controller
	// to use to communicate with Kubernetes

	var config *rest.Config
	var err error
	if *master != "" || *kubeconfig != "" {
		config, err = clientcmd.BuildConfigFromFlags(*master, *kubeconfig)
	} else {
		config, err = rest.InClusterConfig()
	}

	prId := string(uuid.NewUUID())

	if *id != "" {
		prId = *id
	}

	if err != nil {
		glog.Fatalf("Failed to create config: %v", err)
	}
	clientset, err := kubernetes.NewForConfig(config)
	if err != nil {
		glog.Fatalf("Failed to create client: %v", err)
	}

	// The controller needs to know what the server version is because out-of-tree
	// provisioners aren't officially supported until 1.5
	serverVersion, err := clientset.Discovery().ServerVersion()
	if err != nil {
		glog.Fatalf("Error getting server version: %v", err)
	}

	// Create the provisioner: it implements the Provisioner interface expected by
	// the controller
	glusterBlockProvisioner := NewglusterBlockProvisioner(clientset, prId)

	// Start the provision controller which will dynamically provision glusterblock
	// PVs

	pc := controller.NewProvisionController(
		clientset,
		provisionerName,
		glusterBlockProvisioner,
		serverVersion.GitVersion,
	)

	pc.Run(wait.NeverStop)
}
<|endoftext|>