<commit_before>// +build ignore

package main

import (
	"errors"
	"flag"
	"fmt"
	"log"
	"os"
	"strings"
	"text/template"
)

type (
	Type struct {
		Notation  string
		Spec      string
		Prototype string
	}

	TypeMap map[string]Type

	Config struct {
		Package string
		Imports []string
		Args    string // arguments that we were invoked with; TODO(jdef) rename this to Flags?
		Types   TypeMap
	}
)

func (c *Config) String() string {
	if c == nil {
		return ""
	}
	return fmt.Sprintf("%#v", ([]string)(c.Imports))
}

func (c *Config) Set(s string) error {
	// TODO(jdef) this is gnarly, *Config shouldn't actually implement Set or String; define a new type for imports
	c.Imports = append(c.Imports, s)
	return nil
}

func (tm *TypeMap) Set(s string) error {
	tok := strings.SplitN(s, ":", 3)

	if len(tok) < 2 {
		return errors.New("expected {notation}:{type-spec} syntax, instead of " + s)
	}

	if *tm == nil {
		*tm = make(TypeMap)
	}

	t := (*tm)[tok[0]]
	t.Notation, t.Spec = tok[0], tok[1]

	if t.Notation == "" {
		return fmt.Errorf("type notation in %q may not be an empty string", s)
	}

	if t.Spec == "" {
		return fmt.Errorf("type specification in %q may not be an empty string", s)
	}

	if len(tok) == 3 {
		t.Prototype = tok[2]

		if t.Prototype == "" {
			return fmt.Errorf("prototype specification in %q may not be an empty string", s)
		}
	}

	(*tm)[tok[0]] = t
	return nil
}

func (tm *TypeMap) String() string {
	if tm == nil {
		return ""
	}
	return fmt.Sprintf("%#v", *tm)
}

func (c *Config) Var(notation string, names ...string) string {
	t := c.Type(notation)
	if t == "" || len(names) == 0 {
		return ""
	}
	return "var " + strings.Join(names, ",") + " " + t
}

func (c *Config) Arg(notation, name string) string {
	t := c.Type(notation)
	if t == "" {
		return ""
	}
	if name == "" {
		return t
	}
	if strings.HasSuffix(name, ",") {
		return strings.TrimSpace(name[:len(name)-1]+" "+t) + ", "
	}
	return name + " " + t
}

func (c *Config) Ref(notation, name string) (string, error) {
	t := c.Type(notation)
	if t == "" || name == "" {
		return "", nil
	}
	if strings.HasSuffix(name, ",") {
		if len(name) < 2 {
			return "", errors.New("expected ref name before comma")
		}
		return name[:len(name)-1] + ", ", nil
	}
	return name, nil
}

func (c *Config) RequireType(notation string) (string, error) {
	_, ok := c.Types[notation]
	if !ok {
		return "", fmt.Errorf("type %q is required but not specified", notation)
	}
	return "", nil
}

func (c *Config) RequirePrototype(notation string) (string, error) {
	t, ok := c.Types[notation]
	if !ok {
		// needed for optional types: don't require the prototype if the optional type is not defined
		return "", nil
	}
	if t.Prototype == "" {
		return "", fmt.Errorf("prototype for type %q is required but not specified", notation)
	}
	return "", nil
}

func (c *Config) Type(notation string) string {
	t, ok := c.Types[notation]
	if !ok {
		return ""
	}
	return t.Spec
}

func (c *Config) Prototype(notation string) string {
	t, ok := c.Types[notation]
	if !ok {
		return ""
	}
	return t.Prototype
}

func (c *Config) AddFlags(fs *flag.FlagSet) {
	fs.StringVar(&c.Package, "package", c.Package, "destination package")
	fs.Var(c, "import", "packages to import")
	fs.Var(&c.Types, "type", "auxilliary type mappings in {notation}:{type-spec}:{prototype-expr} format")
}

func NewConfig() *Config {
	var (
		c = Config{
			Package: os.Getenv("GOPACKAGE"),
		}
	)
	return &c
}

func Run(src, test *template.Template, args ...string) {
	if len(args) < 1 {
		panic(errors.New("expected at least one arg"))
	}
	var (
		c             = NewConfig()
		defaultOutput = "foo.go"
		output        string
	)
	if c.Package != "" {
		defaultOutput = c.Package + "_generated.go"
	}

	fs := flag.NewFlagSet(args[0], flag.PanicOnError)
	fs.StringVar(&output, "output", output, "path of the to-be-generated file")
	c.AddFlags(fs)

	if err := fs.Parse(args[1:]); err != nil {
		if err == flag.ErrHelp {
			fs.PrintDefaults()
		}
		panic(err)
	}

	c.Args = strings.Join(args[1:], " ")

	if c.Package == "" {
		c.Package = "foo"
	}

	if output == "" {
		output = defaultOutput
	}

	genmap := make(map[string]*template.Template)
	if src != nil {
		genmap[output] = src
	}
	if test != nil {
		testOutput := output + "_test"
		if strings.HasSuffix(output, ".go") {
			testOutput = output[:len(output)-3] + "_test.go"
		}
		genmap[testOutput] = test
	}
	if len(genmap) == 0 {
		panic(errors.New("neither src or test templates were provided"))
	}

	Generate(genmap, c, func(err error) { panic(err) })
}

func Generate(items map[string]*template.Template, data interface{}, eh func(error)) {
	for filename, t := range items {
		func() {
			f, err := os.Create(filename)
			if err != nil {
				eh(err)
				return
			}
			defer f.Close()
			log.Println("generating file", filename)
			err = t.Execute(f, data)
			if err != nil {
				eh(err)
			}
		}()
	}
}
<commit_msg>gen: safely close generated files<commit_after>// +build ignore

package main

import (
	"errors"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"strings"
	"text/template"
)

type (
	Type struct {
		Notation  string
		Spec      string
		Prototype string
	}

	TypeMap map[string]Type

	Config struct {
		Package string
		Imports []string
		Args    string // arguments that we were invoked with; TODO(jdef) rename this to Flags?
		Types   TypeMap
	}
)

func (c *Config) String() string {
	if c == nil {
		return ""
	}
	return fmt.Sprintf("%#v", ([]string)(c.Imports))
}

func (c *Config) Set(s string) error {
	// TODO(jdef) this is gnarly, *Config shouldn't actually implement Set or String; define a new type for imports
	c.Imports = append(c.Imports, s)
	return nil
}

func (tm *TypeMap) Set(s string) error {
	tok := strings.SplitN(s, ":", 3)

	if len(tok) < 2 {
		return errors.New("expected {notation}:{type-spec} syntax, instead of " + s)
	}

	if *tm == nil {
		*tm = make(TypeMap)
	}

	t := (*tm)[tok[0]]
	t.Notation, t.Spec = tok[0], tok[1]

	if t.Notation == "" {
		return fmt.Errorf("type notation in %q may not be an empty string", s)
	}

	if t.Spec == "" {
		return fmt.Errorf("type specification in %q may not be an empty string", s)
	}

	if len(tok) == 3 {
		t.Prototype = tok[2]

		if t.Prototype == "" {
			return fmt.Errorf("prototype specification in %q may not be an empty string", s)
		}
	}

	(*tm)[tok[0]] = t
	return nil
}

func (tm *TypeMap) String() string {
	if tm == nil {
		return ""
	}
	return fmt.Sprintf("%#v", *tm)
}

func (c *Config) Var(notation string, names ...string) string {
	t := c.Type(notation)
	if t == "" || len(names) == 0 {
		return ""
	}
	return "var " + strings.Join(names, ",") + " " + t
}

func (c *Config) Arg(notation, name string) string {
	t := c.Type(notation)
	if t == "" {
		return ""
	}
	if name == "" {
		return t
	}
	if strings.HasSuffix(name, ",") {
		return strings.TrimSpace(name[:len(name)-1]+" "+t) + ", "
	}
	return name + " " + t
}

func (c *Config) Ref(notation, name string) (string, error) {
	t := c.Type(notation)
	if t == "" || name == "" {
		return "", nil
	}
	if strings.HasSuffix(name, ",") {
		if len(name) < 2 {
			return "", errors.New("expected ref name before comma")
		}
		return name[:len(name)-1] + ", ", nil
	}
	return name, nil
}

func (c *Config) RequireType(notation string) (string, error) {
	_, ok := c.Types[notation]
	if !ok {
		return "", fmt.Errorf("type %q is required but not specified", notation)
	}
	return "", nil
}

func (c *Config) RequirePrototype(notation string) (string, error) {
	t, ok := c.Types[notation]
	if !ok {
		// needed for optional types: don't require the prototype if the optional type is not defined
		return "", nil
	}
	if t.Prototype == "" {
		return "", fmt.Errorf("prototype for type %q is required but not specified", notation)
	}
	return "", nil
}

func (c *Config) Type(notation string) string {
	t, ok := c.Types[notation]
	if !ok {
		return ""
	}
	return t.Spec
}

func (c *Config) Prototype(notation string) string {
	t, ok := c.Types[notation]
	if !ok {
		return ""
	}
	return t.Prototype
}

func (c *Config) AddFlags(fs *flag.FlagSet) {
	fs.StringVar(&c.Package, "package", c.Package, "destination package")
	fs.Var(c, "import", "packages to import")
	fs.Var(&c.Types, "type", "auxilliary type mappings in {notation}:{type-spec}:{prototype-expr} format")
}

func NewConfig() *Config {
	var (
		c = Config{
			Package: os.Getenv("GOPACKAGE"),
		}
	)
	return &c
}

func Run(src, test *template.Template, args ...string) {
	if len(args) < 1 {
		panic(errors.New("expected at least one arg"))
	}
	var (
		c             = NewConfig()
		defaultOutput = "foo.go"
		output        string
	)
	if c.Package != "" {
		defaultOutput = c.Package + "_generated.go"
	}

	fs := flag.NewFlagSet(args[0], flag.PanicOnError)
	fs.StringVar(&output, "output", output, "path of the to-be-generated file")
	c.AddFlags(fs)

	if err := fs.Parse(args[1:]); err != nil {
		if err == flag.ErrHelp {
			fs.PrintDefaults()
		}
		panic(err)
	}

	c.Args = strings.Join(args[1:], " ")

	if c.Package == "" {
		c.Package = "foo"
	}

	if output == "" {
		output = defaultOutput
	}

	genmap := make(map[string]*template.Template)
	if src != nil {
		genmap[output] = src
	}
	if test != nil {
		testOutput := output + "_test"
		if strings.HasSuffix(output, ".go") {
			testOutput = output[:len(output)-3] + "_test.go"
		}
		genmap[testOutput] = test
	}
	if len(genmap) == 0 {
		panic(errors.New("neither src or test templates were provided"))
	}

	Generate(genmap, c, func(err error) { panic(err) })
}

func Generate(items map[string]*template.Template, data interface{}, eh func(error)) {
	for filename, t := range items {
		func() {
			f, err := os.Create(filename)
			if err != nil {
				eh(err)
				return
			}
			closer := safeClose(f)
			defer closer()

			log.Println("generating file", filename)
			err = t.Execute(f, data)
			if err != nil {
				eh(err)
			}
			err = f.Sync()
			if err != nil {
				eh(err)
			}
			err = closer()
			if err != nil {
				eh(err)
			}
		}()
	}
}

func safeClose(c io.Closer) func() error {
	var s bool
	return func() error {
		if s {
			return nil
		}
		s = true
		return c.Close()
	}
}
<|endoftext|>


<commit_before>package devops

import (
	. "github.com/influxdata/influxdb-comparisons/bulk_data_gen/common"
	"time"
)

// A DevopsSimulator generates data similar to telemetry from Telegraf.
// It fulfills the Simulator interface.
type DevopsSimulator struct {
	madePoints int64
	madeValues int64
	maxPoints  int64

	simulatedMeasurementIndex int

	hostIndex int
	hosts     []Host

	timestampNow   time.Time
	timestampStart time.Time
	timestampEnd   time.Time
}

func (g *DevopsSimulator) SeenPoints() int64 {
	return g.madePoints
}

func (g *DevopsSimulator) SeenValues() int64 {
	return g.madeValues
}

func (g *DevopsSimulator) Total() int64 {
	return g.maxPoints
}

func (g *DevopsSimulator) Finished() bool {
	return g.madePoints >= g.maxPoints
}

// Type DevopsSimulatorConfig is used to create a DevopsSimulator.
type DevopsSimulatorConfig struct {
	Start time.Time
	End   time.Time

	HostCount int64
	HostOffset int64
}

func (d *DevopsSimulatorConfig) ToSimulator() *DevopsSimulator {
	hostInfos := make([]Host, d.HostCount)
	for i := 0; i < len(hostInfos); i++ {
		hostInfos[i] = NewHost(i, int(d.HostOffset), d.Start)
	}

	epochs := d.End.Sub(d.Start).Nanoseconds() / EpochDuration.Nanoseconds()
	maxPoints := epochs * (d.HostCount * NHostSims)
	dg := &DevopsSimulator{
		madePoints: 0,
		madeValues: 0,
		maxPoints:  maxPoints,

		simulatedMeasurementIndex: 0,

		hostIndex: 0,
		hosts:     hostInfos,

		timestampNow:   d.Start,
		timestampStart: d.Start,
		timestampEnd:   d.End,
	}

	return dg
}

// Next advances a Point to the next state in the generator.
func (d *DevopsSimulator) Next(p *Point) {
	// switch to the next metric if needed
	if d.hostIndex == len(d.hosts) {
		d.hostIndex = 0
		d.simulatedMeasurementIndex++
	}

	if d.simulatedMeasurementIndex == NHostSims {
		d.simulatedMeasurementIndex = 0

		for i := 0; i < len(d.hosts); i++ {
			d.hosts[i].TickAll(EpochDuration)
		}
	}

	host := &d.hosts[d.hostIndex]

	// Populate host-specific tags:
	p.AppendTag(MachineTagKeys[0], host.Name)
	p.AppendTag(MachineTagKeys[1], host.Region)
	p.AppendTag(MachineTagKeys[2], host.Datacenter)
	p.AppendTag(MachineTagKeys[3], host.Rack)
	p.AppendTag(MachineTagKeys[4], host.OS)
	p.AppendTag(MachineTagKeys[5], host.Arch)
	p.AppendTag(MachineTagKeys[6], host.Team)
	p.AppendTag(MachineTagKeys[7], host.Service)
	p.AppendTag(MachineTagKeys[8], host.ServiceVersion)
	p.AppendTag(MachineTagKeys[9], host.ServiceEnvironment)

	// Populate measurement-specific tags and fields:
	host.SimulatedMeasurements[d.simulatedMeasurementIndex].ToPoint(p)

	d.madePoints++
	d.hostIndex++
	d.madeValues += int64(len(p.FieldValues))

	return
}
<commit_msg>changed devops measurements interleaving<commit_after>package devops

import (
	. "github.com/influxdata/influxdb-comparisons/bulk_data_gen/common"
	"time"
)

// A DevopsSimulator generates data similar to telemetry from Telegraf.
// It fulfills the Simulator interface.
type DevopsSimulator struct {
	madePoints int64
	madeValues int64
	maxPoints  int64

	simulatedMeasurementIndex int

	hostIndex int
	hosts     []Host

	timestampNow   time.Time
	timestampStart time.Time
	timestampEnd   time.Time
}

func (g *DevopsSimulator) SeenPoints() int64 {
	return g.madePoints
}

func (g *DevopsSimulator) SeenValues() int64 {
	return g.madeValues
}

func (g *DevopsSimulator) Total() int64 {
	return g.maxPoints
}

func (g *DevopsSimulator) Finished() bool {
	return g.madePoints >= g.maxPoints
}

// Type DevopsSimulatorConfig is used to create a DevopsSimulator.
type DevopsSimulatorConfig struct {
	Start time.Time
	End   time.Time

	HostCount int64
	HostOffset int64
}

func (d *DevopsSimulatorConfig) ToSimulator() *DevopsSimulator {
	hostInfos := make([]Host, d.HostCount)
	for i := 0; i < len(hostInfos); i++ {
		hostInfos[i] = NewHost(i, int(d.HostOffset), d.Start)
	}

	epochs := d.End.Sub(d.Start).Nanoseconds() / EpochDuration.Nanoseconds()
	maxPoints := epochs * (d.HostCount * NHostSims)
	dg := &DevopsSimulator{
		madePoints: 0,
		madeValues: 0,
		maxPoints:  maxPoints,

		simulatedMeasurementIndex: 0,

		hostIndex: 0,
		hosts:     hostInfos,

		timestampNow:   d.Start,
		timestampStart: d.Start,
		timestampEnd:   d.End,
	}

	return dg
}

// Next advances a Point to the next state in the generator.
func (d *DevopsSimulator) Next(p *Point) {
	// switch to the next host if needed
	if d.simulatedMeasurementIndex == NHostSims {
		d.simulatedMeasurementIndex = 0
		d.hostIndex++
	}

	if d.hostIndex == len(d.hosts) {
		d.hostIndex = 0
		for i := 0; i < len(d.hosts); i++ {
			d.hosts[i].TickAll(EpochDuration)
		}
	}

	host := &d.hosts[d.hostIndex]

	// Populate host-specific tags:
	p.AppendTag(MachineTagKeys[0], host.Name)
	p.AppendTag(MachineTagKeys[1], host.Region)
	p.AppendTag(MachineTagKeys[2], host.Datacenter)
	p.AppendTag(MachineTagKeys[3], host.Rack)
	p.AppendTag(MachineTagKeys[4], host.OS)
	p.AppendTag(MachineTagKeys[5], host.Arch)
	p.AppendTag(MachineTagKeys[6], host.Team)
	p.AppendTag(MachineTagKeys[7], host.Service)
	p.AppendTag(MachineTagKeys[8], host.ServiceVersion)
	p.AppendTag(MachineTagKeys[9], host.ServiceEnvironment)

	// Populate measurement-specific tags and fields:
	host.SimulatedMeasurements[d.simulatedMeasurementIndex].ToPoint(p)

	d.madePoints++
	d.simulatedMeasurementIndex++
	d.madeValues += int64(len(p.FieldValues))

	return
}
<|endoftext|>


<commit_before>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"log"
	"net/http"
	"sync"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)
	//Called when no event handler for a specific event exists
	DefaultHandler func(*Server, *Client, []byte) []byte

	handlers     map[string]func(*Server, *Client, []byte) []byte
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       id,
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

var emitPool = &sync.Pool{New: func() interface{} { return emitJS{} }}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := emitPool.Get().(emitJS)
	defer emitPool.Put(js)

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, []byte) []byte)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	go s.listener()
	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.RLock()
	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	s.roomsLock.RUnlock()
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		return
	}

	s.roomsLock.Lock()
	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1
	s.joinedRoomsLock.RLock()
	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	s.joinedRoomsLock.RUnlock()
	if index == -1 {
		return
	}

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data
func (s *Server) Broadcast(room string, data string) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	rooms := make([]string, len(s.joinedRooms[id]))
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string `json:"id"`
	Data string `json:"data,string"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		js := reqPool.Get().(request)
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(string(js.Data))

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, js.Data)

			replyJs := replyPool.Get().(reply)
			replyJs.Id = js.Id
			replyJs.Data = string(rtrn)

			bytes, _ := json.Marshal(replyJs)
			c.Emit(string(bytes))

			reqPool.Put(js)
			replyPool.Put(replyJs)
		}()
	}
}

func (s *Server) listener() {
	for {
		c := <-s.newClient
		go c.listener(s)
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, []byte) []byte) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<commit_msg>EmitJSON: actually send out data<commit_after>//Copyright 2015 TF2Stadium. All rights reserved.
//Use of this source code is governed by the MIT
//that can be found in the LICENSE file.

//Package wsevent implements thread-safe event-driven communication similar to socket.IO,
//on the top of Gorilla's WebSocket implementation.
package wsevent

import (
	"crypto/rand"
	"encoding/base64"
	"encoding/json"
	"log"
	"net/http"
	"sync"

	ws "github.com/gorilla/websocket"
)

//Client
type Client struct {
	//Session ID
	id string

	conn     *ws.Conn
	connLock *sync.RWMutex
	request  *http.Request
}

//Server
type Server struct {
	//maps room string to a list of clients in it
	rooms     map[string]([]*Client)
	roomsLock *sync.RWMutex

	//maps client IDs to the list of rooms the corresponding client has joined
	joinedRooms     map[string][]string
	joinedRoomsLock *sync.RWMutex

	//The extractor function reads the byte array and the message type
	//and returns the event represented by the message.
	Extractor func(string) string
	//Called when the websocket connection closes. The disconnected client's
	//session ID is sent as an argument
	OnDisconnect func(string)
	//Called when no event handler for a specific event exists
	DefaultHandler func(*Server, *Client, []byte) []byte

	handlers     map[string]func(*Server, *Client, []byte) []byte
	handlersLock *sync.RWMutex

	newClient chan *Client
}

func genID() string {
	bytes := make([]byte, 32)
	rand.Read(bytes)

	return base64.URLEncoding.EncodeToString(bytes)
}

//Returns the client's unique session ID
func (c *Client) Id() string {
	return c.id
}

// Returns the first http request when established connection.
func (c *Client) Request() *http.Request {
	return c.request
}

func (s *Server) NewClientWithID(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request, id string) (*Client, error) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		return nil, err
	}

	client := &Client{
		id:       id,
		conn:     conn,
		connLock: new(sync.RWMutex),
		request:  r,
	}
	s.newClient <- client

	return client, nil
}

func (s *Server) NewClient(upgrader ws.Upgrader, w http.ResponseWriter, r *http.Request) (*Client, error) {
	return s.NewClientWithID(upgrader, w, r, genID())
}

func (c *Client) Close() error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.Close()
}

//A thread-safe variant of WriteMessage
func (c *Client) Emit(data string) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()
	return c.conn.WriteMessage(ws.TextMessage, []byte(data))
}

type emitJS struct {
	Id   int         `json:"id"`
	Data interface{} `json:"data"`
}

var emitPool = &sync.Pool{New: func() interface{} { return emitJS{} }}

//A thread-safe variant of EmitJSON
func (c *Client) EmitJSON(v interface{}) error {
	c.connLock.Lock()
	defer c.connLock.Unlock()

	js := emitPool.Get().(emitJS)
	defer emitPool.Put(js)

	js.Id = -1
	js.Data = v

	return c.conn.WriteJSON(js)
}

//Return a new server object
func NewServer() *Server {
	s := &Server{
		rooms:     make(map[string]([]*Client)),
		roomsLock: new(sync.RWMutex),

		//Maps socket ID -> list of rooms the client is in
		joinedRooms:     make(map[string][]string),
		joinedRoomsLock: new(sync.RWMutex),

		handlers:     make(map[string](func(*Server, *Client, []byte) []byte)),
		handlersLock: new(sync.RWMutex),

		newClient: make(chan *Client),
	}

	go s.listener()
	return s
}

//Add a client c to room r
func (s *Server) AddClient(c *Client, r string) {
	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		if r == room {
			//log.Printf("%s already in room %s", c.id, r)
			s.joinedRoomsLock.RUnlock()
			return
		}
	}
	s.joinedRoomsLock.RUnlock()

	s.roomsLock.Lock()
	defer s.roomsLock.Unlock()
	s.rooms[r] = append(s.rooms[r], c)

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()
	s.joinedRooms[c.id] = append(s.joinedRooms[c.id], r)
	//log.Printf("Added %s to room %s", c.id, r)
}

//Remove client c from room r
func (s *Server) RemoveClient(id, r string) {
	index := -1
	s.roomsLock.RLock()
	for i, client := range s.rooms[r] {
		if id == client.id {
			index = i
			break
		}
	}
	s.roomsLock.RUnlock()
	if index == -1 {
		//log.Printf("Client %s not found in room %s", id, r)
		return
	}

	s.roomsLock.Lock()
	s.rooms[r] = append(s.rooms[r][:index], s.rooms[r][index+1:]...)
	s.roomsLock.Unlock()

	index = -1
	s.joinedRoomsLock.RLock()
	for i, room := range s.joinedRooms[id] {
		if room == r {
			index = i
		}
	}
	s.joinedRoomsLock.RUnlock()
	if index == -1 {
		return
	}

	s.joinedRoomsLock.Lock()
	defer s.joinedRoomsLock.Unlock()

	s.joinedRooms[id] = append(s.joinedRooms[id][:index], s.joinedRooms[id][index+1:]...)
}

//Send all clients in room room data
func (s *Server) Broadcast(room string, data string) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s in room %s\n", client.id, room)
		go func(c *Client) {
			c.Emit(data)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (s *Server) BroadcastJSON(room string, v interface{}) {
	s.roomsLock.RLock()
	for _, client := range s.rooms[room] {
		//log.Printf("sending to %s %s\n", client.id, room)
		go func(c *Client) {
			c.EmitJSON(v)
		}(client)
	}
	s.roomsLock.RUnlock()
}

func (c *Client) cleanup(s *Server) {
	c.conn.Close()

	s.joinedRoomsLock.RLock()
	for _, room := range s.joinedRooms[c.id] {
		//log.Println(room)
		index := -1

		s.roomsLock.Lock()
		for i, client := range s.rooms[room] {
			if client.id == c.id {
				index = i
			}
		}

		s.rooms[room] = append(s.rooms[room][:index], s.rooms[room][index+1:]...)
		s.roomsLock.Unlock()
	}
	s.joinedRoomsLock.RUnlock()

	s.joinedRoomsLock.Lock()
	delete(s.joinedRooms, c.id)
	s.joinedRoomsLock.Unlock()

	if s.OnDisconnect != nil {
		s.OnDisconnect(c.id)
	}
}

//Returns an array of rooms the client c has been added to
func (s *Server) RoomsJoined(id string) []string {
	rooms := make([]string, len(s.joinedRooms[id]))
	s.joinedRoomsLock.RLock()
	defer s.joinedRoomsLock.RUnlock()

	copy(rooms, s.joinedRooms[id])

	return rooms
}

type request struct {
	Id   string
	Data json.RawMessage
}

type reply struct {
	Id   string `json:"id"`
	Data string `json:"data,string"`
}

var (
	reqPool   = &sync.Pool{New: func() interface{} { return request{} }}
	replyPool = &sync.Pool{New: func() interface{} { return reply{} }}
)

func (c *Client) listener(s *Server) {
	for {
		mtype, data, err := c.conn.ReadMessage()
		if err != nil {
			c.cleanup(s)
			return
		}

		js := reqPool.Get().(request)
		err = json.Unmarshal(data, &js)

		if err != nil || mtype != ws.TextMessage {
			log.Println(err)
			continue
		}

		callName := s.Extractor(string(js.Data))

		s.handlersLock.RLock()
		f, ok := s.handlers[callName]
		s.handlersLock.RUnlock()

		if !ok {
			if s.DefaultHandler != nil {
				f = s.DefaultHandler
				goto call
			}
			continue
		}
	call:
		go func() {
			rtrn := f(s, c, js.Data)

			replyJs := replyPool.Get().(reply)
			replyJs.Id = js.Id
			replyJs.Data = string(rtrn)

			bytes, _ := json.Marshal(replyJs)
			c.Emit(string(bytes))

			reqPool.Put(js)
			replyPool.Put(replyJs)
		}()
	}
}

func (s *Server) listener() {
	for {
		c := <-s.newClient
		go c.listener(s)
	}
}

//Registers a callback for the event string. The callback must take 2 arguments,
//The client from which the message was received and the string message itself.
func (s *Server) On(event string, f func(*Server, *Client, []byte) []byte) {
	s.handlersLock.Lock()
	s.handlers[event] = f
	s.handlersLock.Unlock()
}
<|endoftext|>


<commit_before>// Copyright 2016 Google Inc. All rights reserved.
//
// 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 version

var (
	// Major is the current major version of master branch.
	Major = 0
	// Minor is the current minor version of master branch.
	Minor = 3
	// Patch is the curernt patched version of the master branch.
	Patch = 1
	// Release is the current release level of the master branch. Valid values
	// are dev (developement unreleased), rcX (release candidate with current
	// iteration), stable (indicates a final released version).
	Release = "dev"
)
<commit_msg>Update release version<commit_after>// Copyright 2016 Google Inc. All rights reserved.
//
// 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 version

var (
	// Major is the current major version of master branch.
	Major = 0
	// Minor is the current minor version of master branch.
	Minor = 4
	// Patch is the curernt patched version of the master branch.
	Patch = 0
	// Release is the current release level of the master branch. Valid values
	// are dev (developement unreleased), rcX (release candidate with current
	// iteration), stable (indicates a final released version).
	Release = "stable"
)
<|endoftext|>


<commit_before>//  Copyright (c) 2018 Cisco and/or its affiliates.
//
//  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 agent

import (
	"os"
	"reflect"
	"strings"

	"github.com/ligato/cn-infra/infra"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logrus"
)

var infraLogger = logrus.NewLogger("infra")

func init() {
	if os.Getenv("DEBUG_INFRA") != "" {
		infraLogger.SetLevel(logging.DebugLevel)
		infraLogger.Debugf("infra debug logger enabled")
	}
}

var (
	// use DEBUG_INFRA=lookup to print plugin verbose lookup logs
	printPluginLookupDebugs = strings.Contains(strings.ToLower(os.Getenv("DEBUG_INFRA")), "lookup")
	// use DEBUG_INFRA=start to print plugin start durations
	printPluginStartDurations = strings.Contains(strings.ToLower(os.Getenv("DEBUG_INFRA")), "start")
)

func findPlugins(val reflect.Value, uniqueness map[infra.Plugin]struct{}, x ...int) (
	res []infra.Plugin, err error,
) {
	n := 0
	if len(x) > 0 {
		n = x[0]
	}
	var logf = func(f string, a ...interface{}) {
		for i := 0; i < n; i++ {
			f = "\t" + f
		}
		//infraLogger.Debugf(f, a...)

	}

	typ := val.Type()

	logf("=> %v (%v)", typ, typ.Kind())
	defer logf("== %v ", typ)

	if typ.Kind() == reflect.Interface {
		if val.IsNil() {
			logf(" - val is nil")
			return nil, nil
		}
		val = val.Elem()
		typ = val.Type()
		//logf(" - interface to elem: %v (%v)", typ, val.Kind())
	}

	if typ.Kind() == reflect.Ptr {
		typ = typ.Elem()
		//logrus.DefaultLogger().Debug(" - typ ptr kind: ", typ)
	}
	if val.Kind() == reflect.Ptr {
		val = val.Elem()
		//logrus.DefaultLogger().Debug(" - val ptr kind: ", val)
	}

	if !val.IsValid() {
		logf(" - val is invalid")
		return nil, nil
	}

	if typ.Kind() != reflect.Struct {
		logf(" - is not a struct: %v %v", typ.Kind(), val.Kind())
		return nil, nil
	}

	//logf(" -> checking %d fields", typ.NumField())

	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)

		// PkgPath is empty for exported fields
		if exported := field.PkgPath == ""; !exported {
			continue
		}

		fieldVal := val.Field(i)

		logf("-> field %d: %v - %v (%v)", i, field.Name, field.Type, fieldVal.Kind())

		var fieldPlug infra.Plugin

		plug, implementsPlugin := isFieldPlugin(field, fieldVal)
		if implementsPlugin {
			if plug == nil {
				logf(" - found nil plugin: %v", field.Name)
				continue
			}

			_, found := uniqueness[plug]
			if found {
				logf(" - found duplicate plugin: %v %v", field.Name, field.Type)
				continue
			}

			// TODO: perhaps add regexp for validation of plugin name

			uniqueness[plug] = struct{}{}
			fieldPlug = plug

			logf(" + FOUND PLUGIN: %v - %v (%v)", plug.String(), field.Name, field.Type)
		}

		// do recursive inspection only for plugins and fields Deps
		if fieldPlug != nil || (field.Anonymous && fieldVal.Kind() == reflect.Struct) {
			// try to inspect structure recursively
			l, err := findPlugins(fieldVal, uniqueness, n+1)
			if err != nil {
				logf(" - Bad field: %v %v", field.Name, err)
				continue
			}
			//logf(" - listed %v plugins from %v (%v)", len(l), field.Name, field.Type)
			res = append(res, l...)
		}

		if fieldPlug != nil {
			res = append(res, fieldPlug)
		}
	}

	logf("<- got %d plugins", len(res))

	return res, nil
}

var pluginType = reflect.TypeOf((*infra.Plugin)(nil)).Elem()

func isFieldPlugin(field reflect.StructField, fieldVal reflect.Value) (infra.Plugin, bool) {
	//logrus.DefaultLogger().Debugf(" - is field plugin: %v (%v) %v", field.Type, fieldVal.Kind(), fieldVal)

	switch fieldVal.Kind() {
	case reflect.Struct:
		ptrType := reflect.PtrTo(fieldVal.Type())
		if ptrType.Implements(pluginType) {
			if fieldVal.CanAddr() {
				if plug, ok := fieldVal.Addr().Interface().(infra.Plugin); ok {
					return plug, true
				}
			}
			return nil, true
		}
	case reflect.Ptr, reflect.Interface:
		if plug, ok := fieldVal.Interface().(infra.Plugin); ok {
			if fieldVal.IsNil() {
				return nil, true
			}
			return plug, true
		}
	}

	return nil, false
}
<commit_msg>Uncomment debug log for plugin lookup<commit_after>//  Copyright (c) 2018 Cisco and/or its affiliates.
//
//  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 agent

import (
	"os"
	"reflect"
	"strings"

	"github.com/ligato/cn-infra/infra"
	"github.com/ligato/cn-infra/logging"
	"github.com/ligato/cn-infra/logging/logrus"
)

var infraLogger = logrus.NewLogger("infra")

func init() {
	if os.Getenv("DEBUG_INFRA") != "" {
		infraLogger.SetLevel(logging.DebugLevel)
		infraLogger.Debugf("infra debug logger enabled")
	}
}

var (
	// use DEBUG_INFRA=lookup to print plugin verbose lookup logs
	printPluginLookupDebugs = strings.Contains(strings.ToLower(os.Getenv("DEBUG_INFRA")), "lookup")
	// use DEBUG_INFRA=start to print plugin start durations
	printPluginStartDurations = strings.Contains(strings.ToLower(os.Getenv("DEBUG_INFRA")), "start")
)

func findPlugins(val reflect.Value, uniqueness map[infra.Plugin]struct{}, x ...int) (
	res []infra.Plugin, err error,
) {
	n := 0
	if len(x) > 0 {
		n = x[0]
	}
	var logf = func(f string, a ...interface{}) {
		for i := 0; i < n; i++ {
			f = "\t" + f
		}
		if printPluginLookupDebugs {
			infraLogger.Debugf(f, a...)
		}
	}

	typ := val.Type()

	logf("=> %v (%v)", typ, typ.Kind())
	defer logf("== %v ", typ)

	if typ.Kind() == reflect.Interface {
		if val.IsNil() {
			logf(" - val is nil")
			return nil, nil
		}
		val = val.Elem()
		typ = val.Type()
		//logf(" - interface to elem: %v (%v)", typ, val.Kind())
	}

	if typ.Kind() == reflect.Ptr {
		typ = typ.Elem()
		//logrus.DefaultLogger().Debug(" - typ ptr kind: ", typ)
	}
	if val.Kind() == reflect.Ptr {
		val = val.Elem()
		//logrus.DefaultLogger().Debug(" - val ptr kind: ", val)
	}

	if !val.IsValid() {
		logf(" - val is invalid")
		return nil, nil
	}

	if typ.Kind() != reflect.Struct {
		logf(" - is not a struct: %v %v", typ.Kind(), val.Kind())
		return nil, nil
	}

	//logf(" -> checking %d fields", typ.NumField())

	for i := 0; i < typ.NumField(); i++ {
		field := typ.Field(i)

		// PkgPath is empty for exported fields
		if exported := field.PkgPath == ""; !exported {
			continue
		}

		fieldVal := val.Field(i)

		logf("-> field %d: %v - %v (%v)", i, field.Name, field.Type, fieldVal.Kind())

		var fieldPlug infra.Plugin

		plug, implementsPlugin := isFieldPlugin(field, fieldVal)
		if implementsPlugin {
			if plug == nil {
				logf(" - found nil plugin: %v", field.Name)
				continue
			}

			_, found := uniqueness[plug]
			if found {
				logf(" - found duplicate plugin: %v %v", field.Name, field.Type)
				continue
			}

			// TODO: perhaps add regexp for validation of plugin name

			uniqueness[plug] = struct{}{}
			fieldPlug = plug

			logf(" + FOUND PLUGIN: %v - %v (%v)", plug.String(), field.Name, field.Type)
		}

		// do recursive inspection only for plugins and fields Deps
		if fieldPlug != nil || (field.Anonymous && fieldVal.Kind() == reflect.Struct) {
			// try to inspect structure recursively
			l, err := findPlugins(fieldVal, uniqueness, n+1)
			if err != nil {
				logf(" - Bad field: %v %v", field.Name, err)
				continue
			}
			//logf(" - listed %v plugins from %v (%v)", len(l), field.Name, field.Type)
			res = append(res, l...)
		}

		if fieldPlug != nil {
			res = append(res, fieldPlug)
		}
	}

	logf("<- got %d plugins", len(res))

	return res, nil
}

var pluginType = reflect.TypeOf((*infra.Plugin)(nil)).Elem()

func isFieldPlugin(field reflect.StructField, fieldVal reflect.Value) (infra.Plugin, bool) {
	//logrus.DefaultLogger().Debugf(" - is field plugin: %v (%v) %v", field.Type, fieldVal.Kind(), fieldVal)

	switch fieldVal.Kind() {
	case reflect.Struct:
		ptrType := reflect.PtrTo(fieldVal.Type())
		if ptrType.Implements(pluginType) {
			if fieldVal.CanAddr() {
				if plug, ok := fieldVal.Addr().Interface().(infra.Plugin); ok {
					return plug, true
				}
			}
			return nil, true
		}
	case reflect.Ptr, reflect.Interface:
		if plug, ok := fieldVal.Interface().(infra.Plugin); ok {
			if fieldVal.IsNil() {
				return nil, true
			}
			return plug, true
		}
	}

	return nil, false
}
<|endoftext|>


<commit_before>package bridge

import (
	"fmt"
	"net"

	log "github.com/Sirupsen/logrus"
	"github.com/vishvananda/netlink"
)

var bridgeNetworks []*net.IPNet

func init() {
	// Here we don't follow the convention of using the 1st IP of the range for the gateway.
	// This is to use the same gateway IPs as the /24 ranges, which predate the /16 ranges.
	// In theory this shouldn't matter - in practice there's bound to be a few scripts relying
	// on the internal addressing or other stupid things like that.
	// They shouldn't, but hey, let's not break them unless we really have to.
	for _, addr := range []string{
		"172.17.42.1/16", // Don't use 172.16.0.0/16, it conflicts with EC2 DNS 172.16.0.23
		"10.0.42.1/16",   // Don't even try using the entire /8, that's too intrusive
		"10.1.42.1/16",
		"10.42.42.1/16",
		"172.16.42.1/24",
		"172.16.43.1/24",
		"172.16.44.1/24",
		"10.0.42.1/24",
		"10.0.43.1/24",
		"192.168.42.1/24",
		"192.168.43.1/24",
		"192.168.44.1/24",
	} {
		ip, net, err := net.ParseCIDR(addr)
		if err != nil {
			log.Errorf("Failed to parse address %s", addr)
			continue
		}
		net.IP = ip
		bridgeNetworks = append(bridgeNetworks, net)
	}
}

func SetupBridgeIPv4(i *Interface) error {
	bridgeIPv4, err := electBridgeIPv4(i.Config)
	if err != nil {
		return err
	}

	log.Debugf("Creating bridge interface %q with network %s", i.Config.BridgeName, bridgeIPv4)
	if err := netlink.AddrAdd(i.Link, &netlink.Addr{bridgeIPv4, ""}); err != nil {
		return fmt.Errorf("Failed to add IPv4 address %s to bridge: %v", bridgeIPv4, err)
	}

	return nil
}

func electBridgeIPv4(config *Configuration) (*net.IPNet, error) {
	// Use the requested IPv4 CIDR when available.
	if config.AddressIPv4 != nil {
		return config.AddressIPv4, nil
	}

	// We don't check for an error here, because we don't really care if we
	// can't read /etc/resolv.conf. So instead we skip the append if resolvConf
	// is nil. It either doesn't exist, or we can't read it for some reason.
	nameservers := []string{}
	if resolvConf, _ := readResolvConf(); resolvConf != nil {
		nameservers = append(nameservers, getNameserversAsCIDR(resolvConf)...)
	}

	// Try to automatically elect appropriate brige IPv4 settings.
	for _, n := range bridgeNetworks {
		if err := checkNameserverOverlaps(nameservers, n); err == nil {
			if err := checkRouteOverlaps(n); err == nil {
				return n, nil
			}
		}
	}

	return nil, fmt.Errorf("Couldn't find an address range for interface %q", config.BridgeName)
}

func checkNameserverOverlaps(nameservers []string, toCheck *net.IPNet) error {
	for _, ns := range nameservers {
		_, nsNetwork, err := net.ParseCIDR(ns)
		if err != nil {
			return err
		}
		if networkOverlaps(toCheck, nsNetwork) {
			return fmt.Errorf("Requested network %s overlaps with name server")
		}
	}
	return nil
}

func checkRouteOverlaps(toCheck *net.IPNet) error {
	networks, err := netlink.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		return err
	}

	for _, network := range networks {
		// TODO Is that right?
		if network.Dst != nil && networkOverlaps(toCheck, network.Dst) {
			return fmt.Errorf("Requested network %s overlaps with an existing network")
		}
	}
	return nil
}

func networkOverlaps(netX *net.IPNet, netY *net.IPNet) bool {
	if firstIP, _ := networkRange(netX); netY.Contains(firstIP) {
		return true
	}
	if firstIP, _ := networkRange(netY); netX.Contains(firstIP) {
		return true
	}
	return false
}

func networkRange(network *net.IPNet) (net.IP, net.IP) {
	var netIP net.IP
	if network.IP.To4() != nil {
		netIP = network.IP.To4()
	} else if network.IP.To16() != nil {
		netIP = network.IP.To16()
	} else {
		return nil, nil
	}

	lastIP := make([]byte, len(netIP), len(netIP))
	for i := 0; i < len(netIP); i++ {
		lastIP[i] = netIP[i] | ^network.Mask[i]
	}
	return netIP.Mask(network.Mask), net.IP(lastIP)
}
<commit_msg>Fix minor static analysis issue in setup_ipv4.go<commit_after>package bridge

import (
	"fmt"
	"net"

	log "github.com/Sirupsen/logrus"
	"github.com/vishvananda/netlink"
)

var bridgeNetworks []*net.IPNet

func init() {
	// Here we don't follow the convention of using the 1st IP of the range for the gateway.
	// This is to use the same gateway IPs as the /24 ranges, which predate the /16 ranges.
	// In theory this shouldn't matter - in practice there's bound to be a few scripts relying
	// on the internal addressing or other stupid things like that.
	// They shouldn't, but hey, let's not break them unless we really have to.
	for _, addr := range []string{
		"172.17.42.1/16", // Don't use 172.16.0.0/16, it conflicts with EC2 DNS 172.16.0.23
		"10.0.42.1/16",   // Don't even try using the entire /8, that's too intrusive
		"10.1.42.1/16",
		"10.42.42.1/16",
		"172.16.42.1/24",
		"172.16.43.1/24",
		"172.16.44.1/24",
		"10.0.42.1/24",
		"10.0.43.1/24",
		"192.168.42.1/24",
		"192.168.43.1/24",
		"192.168.44.1/24",
	} {
		ip, net, err := net.ParseCIDR(addr)
		if err != nil {
			log.Errorf("Failed to parse address %s", addr)
			continue
		}
		net.IP = ip
		bridgeNetworks = append(bridgeNetworks, net)
	}
}

func SetupBridgeIPv4(i *Interface) error {
	bridgeIPv4, err := electBridgeIPv4(i.Config)
	if err != nil {
		return err
	}

	log.Debugf("Creating bridge interface %q with network %s", i.Config.BridgeName, bridgeIPv4)
	if err := netlink.AddrAdd(i.Link, &netlink.Addr{bridgeIPv4, ""}); err != nil {
		return fmt.Errorf("Failed to add IPv4 address %s to bridge: %v", bridgeIPv4, err)
	}

	return nil
}

func electBridgeIPv4(config *Configuration) (*net.IPNet, error) {
	// Use the requested IPv4 CIDR when available.
	if config.AddressIPv4 != nil {
		return config.AddressIPv4, nil
	}

	// We don't check for an error here, because we don't really care if we
	// can't read /etc/resolv.conf. So instead we skip the append if resolvConf
	// is nil. It either doesn't exist, or we can't read it for some reason.
	nameservers := []string{}
	if resolvConf, _ := readResolvConf(); resolvConf != nil {
		nameservers = append(nameservers, getNameserversAsCIDR(resolvConf)...)
	}

	// Try to automatically elect appropriate brige IPv4 settings.
	for _, n := range bridgeNetworks {
		if err := checkNameserverOverlaps(nameservers, n); err == nil {
			if err := checkRouteOverlaps(n); err == nil {
				return n, nil
			}
		}
	}

	return nil, fmt.Errorf("Couldn't find an address range for interface %q", config.BridgeName)
}

func checkNameserverOverlaps(nameservers []string, toCheck *net.IPNet) error {
	for _, ns := range nameservers {
		_, nsNetwork, err := net.ParseCIDR(ns)
		if err != nil {
			return err
		}
		if networkOverlaps(toCheck, nsNetwork) {
			return fmt.Errorf("Requested network %s overlaps with name server", toCheck.String())
		}
	}
	return nil
}

func checkRouteOverlaps(toCheck *net.IPNet) error {
	networks, err := netlink.RouteList(nil, netlink.FAMILY_V4)
	if err != nil {
		return err
	}

	for _, network := range networks {
		// TODO Is that right?
		if network.Dst != nil && networkOverlaps(toCheck, network.Dst) {
			return fmt.Errorf("Requested network %s overlaps with an existing network", toCheck.String())
		}
	}
	return nil
}

func networkOverlaps(netX *net.IPNet, netY *net.IPNet) bool {
	if firstIP, _ := networkRange(netX); netY.Contains(firstIP) {
		return true
	}
	if firstIP, _ := networkRange(netY); netX.Contains(firstIP) {
		return true
	}
	return false
}

func networkRange(network *net.IPNet) (net.IP, net.IP) {
	var netIP net.IP
	if network.IP.To4() != nil {
		netIP = network.IP.To4()
	} else if network.IP.To16() != nil {
		netIP = network.IP.To16()
	} else {
		return nil, nil
	}

	lastIP := make([]byte, len(netIP), len(netIP))
	for i := 0; i < len(netIP); i++ {
		lastIP[i] = netIP[i] | ^network.Mask[i]
	}
	return netIP.Mask(network.Mask), net.IP(lastIP)
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"sync"
	"time"

	"github.com/miekg/dns"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet/udp"
)

const (
	CleanupInterval  = time.Second * 120
	CleanupThreshold = 512
)

var (
	pseudoDestination = net.UDPDestination(net.LocalHostIP, net.Port(53))

	multiQuestionDns = map[net.Address]bool{
		net.IPAddress([]byte{8, 8, 8, 8}): true,
		net.IPAddress([]byte{8, 8, 4, 4}): true,
		net.IPAddress([]byte{9, 9, 9, 9}): true,
	}
)

type ARecord struct {
	IPs    []net.IP
	Expire time.Time
}

type NameServer interface {
	QueryA(domain string) <-chan *ARecord
}

type PendingRequest struct {
	expire   time.Time
	response chan<- *ARecord
}

type UDPNameServer struct {
	sync.Mutex
	address     net.Destination
	requests    map[uint16]*PendingRequest
	udpServer   *udp.Dispatcher
	nextCleanup time.Time
}

func NewUDPNameServer(address net.Destination, dispatcher dispatcher.Interface) *UDPNameServer {
	s := &UDPNameServer{
		address:   address,
		requests:  make(map[uint16]*PendingRequest),
		udpServer: udp.NewDispatcher(dispatcher),
	}
	return s
}

func (v *UDPNameServer) Cleanup() {
	expiredRequests := make([]uint16, 0, 16)
	now := time.Now()
	v.Lock()
	for id, r := range v.requests {
		if r.expire.Before(now) {
			expiredRequests = append(expiredRequests, id)
			close(r.response)
		}
	}
	for _, id := range expiredRequests {
		delete(v.requests, id)
	}
	v.Unlock()
}

func (v *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
	var id uint16
	v.Lock()
	if len(v.requests) > CleanupThreshold && v.nextCleanup.Before(time.Now()) {
		v.nextCleanup = time.Now().Add(CleanupInterval)
		go v.Cleanup()
	}

	for {
		id = dice.RollUint16()
		if _, found := v.requests[id]; found {
			continue
		}
		log.Trace(newError("add pending request id ", id).AtDebug())
		v.requests[id] = &PendingRequest{
			expire:   time.Now().Add(time.Second * 8),
			response: response,
		}
		break
	}
	v.Unlock()
	return id
}

func (v *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err == dns.ErrTruncated {
		log.Trace(newError("truncated message received. DNS server should still work. If you see anything abnormal, please submit an issue to v2ray-core.").AtWarning())
	} else if err != nil {
		log.Trace(newError("failed to parse DNS response").Base(err).AtWarning())
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := uint32(3600) // an hour
	log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug())

	v.Lock()
	request, found := v.requests[id]
	if !found {
		v.Unlock()
		return
	}
	delete(v.requests, id)
	v.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (v *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {

	msg := new(dns.Msg)
	msg.Id = id
	msg.RecursionDesired = true
	msg.Question = []dns.Question{
		{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeA,
			Qclass: dns.ClassINET,
		}}
	if multiQuestionDns[v.address.Address] {
		msg.Question = append(msg.Question, dns.Question{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeAAAA,
			Qclass: dns.ClassINET,
		})
	}

	buffer := buf.New()
	common.Must(buffer.Reset(func(b []byte) (int, error) {
		writtenBuffer, err := msg.PackBuffer(b)
		return len(writtenBuffer), err
	}))

	return buffer
}

func (v *UDPNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)
	id := v.AssignUnusedID(response)

	ctx, cancel := context.WithCancel(context.Background())
	v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse)

	go func() {
		for i := 0; i < 2; i++ {
			time.Sleep(time.Second)
			v.Lock()
			_, found := v.requests[id]
			v.Unlock()
			if !found {
				break
			}
			v.udpServer.Dispatch(ctx, v.address, v.BuildQueryA(domain, id), v.HandleResponse)
		}
		cancel()
	}()

	return response
}

type LocalNameServer struct {
}

func (v *LocalNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)

	go func() {
		defer close(response)

		ips, err := net.LookupIP(domain)
		if err != nil {
			log.Trace(newError("failed to lookup IPs for domain ", domain).Base(err))
			return
		}

		response <- &ARecord{
			IPs:    ips,
			Expire: time.Now().Add(time.Hour),
		}
	}()

	return response
}
<commit_msg>refactor<commit_after>package server

import (
	"context"
	"sync"
	"time"

	"github.com/miekg/dns"
	"v2ray.com/core/app/dispatcher"
	"v2ray.com/core/app/log"
	"v2ray.com/core/common"
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
	"v2ray.com/core/transport/internet/udp"
)

const (
	CleanupInterval  = time.Second * 120
	CleanupThreshold = 512
)

var (
	multiQuestionDNS = map[net.Address]bool{
		net.IPAddress([]byte{8, 8, 8, 8}): true,
		net.IPAddress([]byte{8, 8, 4, 4}): true,
		net.IPAddress([]byte{9, 9, 9, 9}): true,
	}
)

type ARecord struct {
	IPs    []net.IP
	Expire time.Time
}

type NameServer interface {
	QueryA(domain string) <-chan *ARecord
}

type PendingRequest struct {
	expire   time.Time
	response chan<- *ARecord
}

type UDPNameServer struct {
	sync.Mutex
	address     net.Destination
	requests    map[uint16]*PendingRequest
	udpServer   *udp.Dispatcher
	nextCleanup time.Time
}

func NewUDPNameServer(address net.Destination, dispatcher dispatcher.Interface) *UDPNameServer {
	s := &UDPNameServer{
		address:   address,
		requests:  make(map[uint16]*PendingRequest),
		udpServer: udp.NewDispatcher(dispatcher),
	}
	return s
}

func (s *UDPNameServer) Cleanup() {
	expiredRequests := make([]uint16, 0, 16)
	now := time.Now()
	s.Lock()
	for id, r := range s.requests {
		if r.expire.Before(now) {
			expiredRequests = append(expiredRequests, id)
			close(r.response)
		}
	}
	for _, id := range expiredRequests {
		delete(s.requests, id)
	}
	s.Unlock()
}

func (s *UDPNameServer) AssignUnusedID(response chan<- *ARecord) uint16 {
	var id uint16
	s.Lock()
	if len(s.requests) > CleanupThreshold && s.nextCleanup.Before(time.Now()) {
		s.nextCleanup = time.Now().Add(CleanupInterval)
		go s.Cleanup()
	}

	for {
		id = dice.RollUint16()
		if _, found := s.requests[id]; found {
			continue
		}
		log.Trace(newError("add pending request id ", id).AtDebug())
		s.requests[id] = &PendingRequest{
			expire:   time.Now().Add(time.Second * 8),
			response: response,
		}
		break
	}
	s.Unlock()
	return id
}

func (s *UDPNameServer) HandleResponse(payload *buf.Buffer) {
	msg := new(dns.Msg)
	err := msg.Unpack(payload.Bytes())
	if err == dns.ErrTruncated {
		log.Trace(newError("truncated message received. DNS server should still work. If you see anything abnormal, please submit an issue to v2ray-core.").AtWarning())
	} else if err != nil {
		log.Trace(newError("failed to parse DNS response").Base(err).AtWarning())
		return
	}
	record := &ARecord{
		IPs: make([]net.IP, 0, 16),
	}
	id := msg.Id
	ttl := uint32(3600) // an hour
	log.Trace(newError("handling response for id ", id, " content: ", msg).AtDebug())

	s.Lock()
	request, found := s.requests[id]
	if !found {
		s.Unlock()
		return
	}
	delete(s.requests, id)
	s.Unlock()

	for _, rr := range msg.Answer {
		switch rr := rr.(type) {
		case *dns.A:
			record.IPs = append(record.IPs, rr.A)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		case *dns.AAAA:
			record.IPs = append(record.IPs, rr.AAAA)
			if rr.Hdr.Ttl < ttl {
				ttl = rr.Hdr.Ttl
			}
		}
	}
	record.Expire = time.Now().Add(time.Second * time.Duration(ttl))

	request.response <- record
	close(request.response)
}

func (s *UDPNameServer) BuildQueryA(domain string, id uint16) *buf.Buffer {
	msg := new(dns.Msg)
	msg.Id = id
	msg.RecursionDesired = true
	msg.Question = []dns.Question{
		{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeA,
			Qclass: dns.ClassINET,
		}}
	if multiQuestionDNS[s.address.Address] {
		msg.Question = append(msg.Question, dns.Question{
			Name:   dns.Fqdn(domain),
			Qtype:  dns.TypeAAAA,
			Qclass: dns.ClassINET,
		})
	}

	buffer := buf.New()
	common.Must(buffer.Reset(func(b []byte) (int, error) {
		writtenBuffer, err := msg.PackBuffer(b)
		return len(writtenBuffer), err
	}))

	return buffer
}

func (s *UDPNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)
	id := s.AssignUnusedID(response)

	ctx, cancel := context.WithCancel(context.Background())
	s.udpServer.Dispatch(ctx, s.address, s.BuildQueryA(domain, id), s.HandleResponse)

	go func() {
		for i := 0; i < 2; i++ {
			time.Sleep(time.Second)
			s.Lock()
			_, found := s.requests[id]
			s.Unlock()
			if !found {
				break
			}
			s.udpServer.Dispatch(ctx, s.address, s.BuildQueryA(domain, id), s.HandleResponse)
		}
		cancel()
	}()

	return response
}

type LocalNameServer struct {
}

func (*LocalNameServer) QueryA(domain string) <-chan *ARecord {
	response := make(chan *ARecord, 1)

	go func() {
		defer close(response)

		ips, err := net.LookupIP(domain)
		if err != nil {
			log.Trace(newError("failed to lookup IPs for domain ", domain).Base(err))
			return
		}

		response <- &ARecord{
			IPs:    ips,
			Expire: time.Now().Add(time.Hour),
		}
	}()

	return response
}
<|endoftext|>


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

package walk

import (
	"fmt"
	"sync"
	"syscall"
	"unsafe"
)

import . "github.com/lxn/go-winapi"

type CloseReason byte

const (
	CloseReasonUnknown CloseReason = iota
	CloseReasonUser
)

var syncFuncs struct {
	m     sync.Mutex
	funcs []func()
}

var syncMsgId uint32
var taskbarButtonCreatedMsgId uint32

func init() {
	syncMsgId = RegisterWindowMessage(syscall.StringToUTF16Ptr("WalkSync"))
	taskbarButtonCreatedMsgId = RegisterWindowMessage(syscall.StringToUTF16Ptr("TaskbarButtonCreated"))
}

func synchronize(f func()) {
	syncFuncs.m.Lock()
	defer syncFuncs.m.Unlock()
	syncFuncs.funcs = append(syncFuncs.funcs, f)
}

func runSynchronized() {
	// Clear the list of callbacks first to avoid deadlock
	// if a callback itself calls Synchronize()...
	syncFuncs.m.Lock()
	funcs := syncFuncs.funcs
	syncFuncs.funcs = nil
	syncFuncs.m.Unlock()
	for _, f := range funcs {
		f()
	}
}

type TopLevelWindow struct {
	ContainerBase
	owner                 RootWidget
	closingPublisher      CloseEventPublisher
	startingPublisher     EventPublisher
	titleChangedPublisher EventPublisher
	progressIndicator     *ProgressIndicator
	icon                  *Icon
	prevFocusHWnd         HWND
	isInRestoreState      bool
	closeReason           CloseReason
}

func (tlw *TopLevelWindow) init() {
	tlw.MustRegisterProperty("Title", NewProperty(
		func() interface{} {
			return tlw.Title()
		},
		func(v interface{}) error {
			return tlw.SetTitle(v.(string))
		},
		tlw.titleChangedPublisher.Event()))
}

func (tlw *TopLevelWindow) LayoutFlags() LayoutFlags {
	return ShrinkableHorz | ShrinkableVert | GrowableHorz | GrowableVert | GreedyHorz | GreedyVert
}

func (tlw *TopLevelWindow) SizeHint() Size {
	return tlw.dialogBaseUnitsToPixels(Size{252, 218})
}

func (tlw *TopLevelWindow) Enabled() bool {
	return tlw.enabled
}

func (tlw *TopLevelWindow) SetEnabled(enabled bool) {
	tlw.WidgetBase.SetEnabled(enabled)
}

func (tlw *TopLevelWindow) Font() *Font {
	if tlw.font != nil {
		return tlw.font
	}

	return defaultFont
}

func (tlw *TopLevelWindow) Title() string {
	return widgetText(tlw.hWnd)
}

func (tlw *TopLevelWindow) SetTitle(value string) error {
	return setWidgetText(tlw.hWnd, value)
}

func (tlw *TopLevelWindow) Run() int {
	tlw.startingPublisher.Publish()

	var msg MSG

	for tlw.hWnd != 0 {
		switch GetMessage(&msg, 0, 0, 0) {
		case 0:
			return int(msg.WParam)

		case -1:
			return -1
		}

		if !IsDialogMessage(tlw.hWnd, &msg) {
			TranslateMessage(&msg)
			DispatchMessage(&msg)
		}

		runSynchronized()
	}

	return 0
}

func (tlw *TopLevelWindow) Starting() *Event {
	return tlw.startingPublisher.Event()
}

func (tlw *TopLevelWindow) Owner() RootWidget {
	return tlw.owner
}

func (tlw *TopLevelWindow) SetOwner(value RootWidget) error {
	tlw.owner = value

	var ownerHWnd HWND
	if value != nil {
		ownerHWnd = value.Handle()
	}

	SetLastError(0)
	if 0 == SetWindowLong(
		tlw.hWnd,
		GWL_HWNDPARENT,
		int32(ownerHWnd)) && GetLastError() != 0 {

		return lastError("SetWindowLong")
	}

	return nil
}

func (tlw *TopLevelWindow) Icon() *Icon {
	return tlw.icon
}

func (tlw *TopLevelWindow) SetIcon(icon *Icon) {
	tlw.icon = icon

	var hIcon uintptr
	if icon != nil {
		hIcon = uintptr(icon.hIcon)
	}

	tlw.SendMessage(WM_SETICON, 0, hIcon)
	tlw.SendMessage(WM_SETICON, 1, hIcon)
}

func (tlw *TopLevelWindow) Hide() {
	tlw.widget.SetVisible(false)
}

func (tlw *TopLevelWindow) Show() {
	if appSingleton.settings != nil {
		tlw.RestoreState()
	}

	tlw.widget.SetVisible(true)
}

func (tlw *TopLevelWindow) close() error {
	if appSingleton.settings != nil {
		tlw.SaveState()
	}

	tlw.widget.Dispose()

	return nil
}

func (tlw *TopLevelWindow) Close() error {
	tlw.SendMessage(WM_CLOSE, 0, 0)

	return nil
}

func (tlw *TopLevelWindow) SaveState() error {
	var wp WINDOWPLACEMENT

	wp.Length = uint32(unsafe.Sizeof(wp))

	if !GetWindowPlacement(tlw.hWnd, &wp) {
		return lastError("GetWindowPlacement")
	}

	state := fmt.Sprint(
		wp.Flags, wp.ShowCmd,
		wp.PtMinPosition.X, wp.PtMinPosition.Y,
		wp.PtMaxPosition.X, wp.PtMaxPosition.Y,
		wp.RcNormalPosition.Left, wp.RcNormalPosition.Top,
		wp.RcNormalPosition.Right, wp.RcNormalPosition.Bottom)

	if err := tlw.putState(state); err != nil {
		return err
	}

	return tlw.ContainerBase.SaveState()
}

func (tlw *TopLevelWindow) RestoreState() error {
	if tlw.isInRestoreState {
		return nil
	}
	tlw.isInRestoreState = true
	defer func() {
		tlw.isInRestoreState = false
	}()

	state, err := tlw.getState()
	if err != nil {
		return err
	}
	if state == "" {
		return nil
	}

	var wp WINDOWPLACEMENT

	if _, err := fmt.Sscan(state,
		&wp.Flags, &wp.ShowCmd,
		&wp.PtMinPosition.X, &wp.PtMinPosition.Y,
		&wp.PtMaxPosition.X, &wp.PtMaxPosition.Y,
		&wp.RcNormalPosition.Left, &wp.RcNormalPosition.Top,
		&wp.RcNormalPosition.Right, &wp.RcNormalPosition.Bottom); err != nil {
		return err
	}

	wp.Length = uint32(unsafe.Sizeof(wp))

	if !SetWindowPlacement(tlw.hWnd, &wp) {
		return lastError("SetWindowPlacement")
	}

	if err := tlw.ContainerBase.RestoreState(); err != nil {
		return err
	}

	return nil
}

func (tlw *TopLevelWindow) Closing() *CloseEvent {
	return tlw.closingPublisher.Event()
}

func (tlw *TopLevelWindow) ProgressIndicator() *ProgressIndicator {
	return tlw.progressIndicator
}

func (tlw *TopLevelWindow) WndProc(hwnd HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case WM_ACTIVATE:
		switch LOWORD(uint32(wParam)) {
		case WA_ACTIVE, WA_CLICKACTIVE:
			if tlw.prevFocusHWnd != 0 {
				SetFocus(tlw.prevFocusHWnd)
			}

		case WA_INACTIVE:
			tlw.prevFocusHWnd = GetFocus()
		}
		return 0

	case WM_CLOSE:
		tlw.closeReason = CloseReasonUnknown
		var canceled bool
		tlw.closingPublisher.Publish(&canceled, tlw.closeReason)
		if !canceled {
			if tlw.owner != nil {
				tlw.owner.SetEnabled(true)
				if !SetWindowPos(tlw.owner.Handle(), HWND_NOTOPMOST, 0, 0, 0, 0, SWP_NOMOVE|SWP_NOSIZE|SWP_SHOWWINDOW) {
					lastError("SetWindowPos")
				}
			}

			tlw.close()
		}
		return 0

	case WM_GETMINMAXINFO:
		mmi := (*MINMAXINFO)(unsafe.Pointer(lParam))

		var layout Layout
		if container, ok := tlw.widget.(Container); ok {
			layout = container.Layout()
		}

		var min Size
		if layout != nil {
			min = tlw.sizeFromClientSize(layout.MinSize())
		}

		mmi.PtMinTrackSize = POINT{
			int32(maxi(min.Width, tlw.minSize.Width)),
			int32(maxi(min.Height, tlw.minSize.Height)),
		}
		return 0

	case WM_SETTEXT:
		tlw.titleChangedPublisher.Publish()

	case WM_SYSCOMMAND:
		if wParam == SC_CLOSE {
			tlw.closeReason = CloseReasonUser
		}

	case taskbarButtonCreatedMsgId:
		version := GetVersion()
		major := version & 0xFF
		minor := version & 0xFF00 >> 8
		// Check that the OS is Win 7 or later (Win 7 is v6.1).
		if major > 6 || (major == 6 && minor > 0) {
			tlw.progressIndicator, _ = newTaskbarList3(tlw.hWnd)
		}
	}

	return tlw.ContainerBase.WndProc(hwnd, msg, wParam, lParam)
}
<commit_msg>TopLevelWindow: Only restore + save state if persistent<commit_after>// Copyright 2012 The Walk Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package walk

import (
	"fmt"
	"sync"
	"syscall"
	"unsafe"
)

import . "github.com/lxn/go-winapi"

type CloseReason byte

const (
	CloseReasonUnknown CloseReason = iota
	CloseReasonUser
)

var syncFuncs struct {
	m     sync.Mutex
	funcs []func()
}

var syncMsgId uint32
var taskbarButtonCreatedMsgId uint32

func init() {
	syncMsgId = RegisterWindowMessage(syscall.StringToUTF16Ptr("WalkSync"))
	taskbarButtonCreatedMsgId = RegisterWindowMessage(syscall.StringToUTF16Ptr("TaskbarButtonCreated"))
}

func synchronize(f func()) {
	syncFuncs.m.Lock()
	defer syncFuncs.m.Unlock()
	syncFuncs.funcs = append(syncFuncs.funcs, f)
}

func runSynchronized() {
	// Clear the list of callbacks first to avoid deadlock
	// if a callback itself calls Synchronize()...
	syncFuncs.m.Lock()
	funcs := syncFuncs.funcs
	syncFuncs.funcs = nil
	syncFuncs.m.Unlock()
	for _, f := range funcs {
		f()
	}
}

type TopLevelWindow struct {
	ContainerBase
	owner                 RootWidget
	closingPublisher      CloseEventPublisher
	startingPublisher     EventPublisher
	titleChangedPublisher EventPublisher
	progressIndicator     *ProgressIndicator
	icon                  *Icon
	prevFocusHWnd         HWND
	isInRestoreState      bool
	closeReason           CloseReason
}

func (tlw *TopLevelWindow) init() {
	tlw.MustRegisterProperty("Title", NewProperty(
		func() interface{} {
			return tlw.Title()
		},
		func(v interface{}) error {
			return tlw.SetTitle(v.(string))
		},
		tlw.titleChangedPublisher.Event()))
}

func (tlw *TopLevelWindow) LayoutFlags() LayoutFlags {
	return ShrinkableHorz | ShrinkableVert | GrowableHorz | GrowableVert | GreedyHorz | GreedyVert
}

func (tlw *TopLevelWindow) SizeHint() Size {
	return tlw.dialogBaseUnitsToPixels(Size{252, 218})
}

func (tlw *TopLevelWindow) Enabled() bool {
	return tlw.enabled
}

func (tlw *TopLevelWindow) SetEnabled(enabled bool) {
	tlw.WidgetBase.SetEnabled(enabled)
}

func (tlw *TopLevelWindow) Font() *Font {
	if tlw.font != nil {
		return tlw.font
	}

	return defaultFont
}

func (tlw *TopLevelWindow) Title() string {
	return widgetText(tlw.hWnd)
}

func (tlw *TopLevelWindow) SetTitle(value string) error {
	return setWidgetText(tlw.hWnd, value)
}

func (tlw *TopLevelWindow) Run() int {
	tlw.startingPublisher.Publish()

	var msg MSG

	for tlw.hWnd != 0 {
		switch GetMessage(&msg, 0, 0, 0) {
		case 0:
			return int(msg.WParam)

		case -1:
			return -1
		}

		if !IsDialogMessage(tlw.hWnd, &msg) {
			TranslateMessage(&msg)
			DispatchMessage(&msg)
		}

		runSynchronized()
	}

	return 0
}

func (tlw *TopLevelWindow) Starting() *Event {
	return tlw.startingPublisher.Event()
}

func (tlw *TopLevelWindow) Owner() RootWidget {
	return tlw.owner
}

func (tlw *TopLevelWindow) SetOwner(value RootWidget) error {
	tlw.owner = value

	var ownerHWnd HWND
	if value != nil {
		ownerHWnd = value.Handle()
	}

	SetLastError(0)
	if 0 == SetWindowLong(
		tlw.hWnd,
		GWL_HWNDPARENT,
		int32(ownerHWnd)) && GetLastError() != 0 {

		return lastError("SetWindowLong")
	}

	return nil
}

func (tlw *TopLevelWindow) Icon() *Icon {
	return tlw.icon
}

func (tlw *TopLevelWindow) SetIcon(icon *Icon) {
	tlw.icon = icon

	var hIcon uintptr
	if icon != nil {
		hIcon = uintptr(icon.hIcon)
	}

	tlw.SendMessage(WM_SETICON, 0, hIcon)
	tlw.SendMessage(WM_SETICON, 1, hIcon)
}

func (tlw *TopLevelWindow) Hide() {
	tlw.widget.SetVisible(false)
}

func (tlw *TopLevelWindow) Show() {
	if p, ok := tlw.widget.(Persistable); ok && p.Persistent() && appSingleton.settings != nil {
		p.RestoreState()
	}

	tlw.widget.SetVisible(true)
}

func (tlw *TopLevelWindow) close() error {
	if p, ok := tlw.widget.(Persistable); ok && p.Persistent() && appSingleton.settings != nil {
		p.SaveState()
	}

	tlw.widget.Dispose()

	return nil
}

func (tlw *TopLevelWindow) Close() error {
	tlw.SendMessage(WM_CLOSE, 0, 0)

	return nil
}

func (tlw *TopLevelWindow) SaveState() error {
	var wp WINDOWPLACEMENT

	wp.Length = uint32(unsafe.Sizeof(wp))

	if !GetWindowPlacement(tlw.hWnd, &wp) {
		return lastError("GetWindowPlacement")
	}

	state := fmt.Sprint(
		wp.Flags, wp.ShowCmd,
		wp.PtMinPosition.X, wp.PtMinPosition.Y,
		wp.PtMaxPosition.X, wp.PtMaxPosition.Y,
		wp.RcNormalPosition.Left, wp.RcNormalPosition.Top,
		wp.RcNormalPosition.Right, wp.RcNormalPosition.Bottom)

	if err := tlw.putState(state); err != nil {
		return err
	}

	return tlw.ContainerBase.SaveState()
}

func (tlw *TopLevelWindow) RestoreState() error {
	if tlw.isInRestoreState {
		return nil
	}
	tlw.isInRestoreState = true
	defer func() {
		tlw.isInRestoreState = false
	}()

	state, err := tlw.getState()
	if err != nil {
		return err
	}
	if state == "" {
		return nil
	}

	var wp WINDOWPLACEMENT

	if _, err := fmt.Sscan(state,
		&wp.Flags, &wp.ShowCmd,
		&wp.PtMinPosition.X, &wp.PtMinPosition.Y,
		&wp.PtMaxPosition.X, &wp.PtMaxPosition.Y,
		&wp.RcNormalPosition.Left, &wp.RcNormalPosition.Top,
		&wp.RcNormalPosition.Right, &wp.RcNormalPosition.Bottom); err != nil {
		return err
	}

	wp.Length = uint32(unsafe.Sizeof(wp))

	if !SetWindowPlacement(tlw.hWnd, &wp) {
		return lastError("SetWindowPlacement")
	}

	if err := tlw.ContainerBase.RestoreState(); err != nil {
		return err
	}

	return nil
}

func (tlw *TopLevelWindow) Closing() *CloseEvent {
	return tlw.closingPublisher.Event()
}

func (tlw *TopLevelWindow) ProgressIndicator() *ProgressIndicator {
	return tlw.progressIndicator
}

func (tlw *TopLevelWindow) WndProc(hwnd HWND, msg uint32, wParam, lParam uintptr) uintptr {
	switch msg {
	case WM_ACTIVATE:
		switch LOWORD(uint32(wParam)) {
		case WA_ACTIVE, WA_CLICKACTIVE:
			if tlw.prevFocusHWnd != 0 {
				SetFocus(tlw.prevFocusHWnd)
			}

		case WA_INACTIVE:
			tlw.prevFocusHWnd = GetFocus()
		}
		return 0

	case WM_CLOSE:
		tlw.closeReason = CloseReasonUnknown
		var canceled bool
		tlw.closingPublisher.Publish(&canceled, tlw.closeReason)
		if !canceled {
			if tlw.owner != nil {
				tlw.owner.SetEnabled(true)
				if !SetWindowPos(tlw.owner.Handle(), HWND_NOTOPMOST, 0, 0, 0, 0, SWP_NOMOVE|SWP_NOSIZE|SWP_SHOWWINDOW) {
					lastError("SetWindowPos")
				}
			}

			tlw.close()
		}
		return 0

	case WM_GETMINMAXINFO:
		mmi := (*MINMAXINFO)(unsafe.Pointer(lParam))

		var layout Layout
		if container, ok := tlw.widget.(Container); ok {
			layout = container.Layout()
		}

		var min Size
		if layout != nil {
			min = tlw.sizeFromClientSize(layout.MinSize())
		}

		mmi.PtMinTrackSize = POINT{
			int32(maxi(min.Width, tlw.minSize.Width)),
			int32(maxi(min.Height, tlw.minSize.Height)),
		}
		return 0

	case WM_SETTEXT:
		tlw.titleChangedPublisher.Publish()

	case WM_SYSCOMMAND:
		if wParam == SC_CLOSE {
			tlw.closeReason = CloseReasonUser
		}

	case taskbarButtonCreatedMsgId:
		version := GetVersion()
		major := version & 0xFF
		minor := version & 0xFF00 >> 8
		// Check that the OS is Win 7 or later (Win 7 is v6.1).
		if major > 6 || (major == 6 && minor > 0) {
			tlw.progressIndicator, _ = newTaskbarList3(tlw.hWnd)
		}
	}

	return tlw.ContainerBase.WndProc(hwnd, msg, wParam, lParam)
}
<|endoftext|>


<commit_before>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

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

	"github.com/OWASP/Amass/amass/core"
	"github.com/OWASP/Amass/amass/utils"
)

// Crtsh is data source object type that implements the DataSource interface.
type Crtsh struct {
	BaseDataSource
}

type crtData struct {
	IssuerID          int    `json:"issuer_ca_id"`
	IssuerName        string `json:"issuer_name"`
	Name              string `json:"name_value"`
	MinCertID         int    `json:"min_cert_id"`
	MinEntryTimestamp string `json:"min_entry_timestamp"`
	NotBefore         string `json:"not_before"`
	NotAfter          string `json:"not_after"`
}

// NewCrtsh returns an initialized Crtsh as a DataSource.
func NewCrtsh(srv core.AmassService) DataSource {
	c := new(Crtsh)

	c.BaseDataSource = *NewBaseDataSource(srv, core.CERT, "crt.sh")
	return c
}

// Query returns the subdomain names discovered when querying this data source.
func (c *Crtsh) Query(domain, sub string) []string {
	var unique []string

	if domain != sub {
		return unique
	}
	// Pull the page that lists all certs for this domain
	url := c.getURL(domain)
	page, err := utils.GetWebPage(url, nil)
	if err != nil {
		c.Service.Config().Log.Printf("%s: %v", url, err)
		return unique
	}
	c.Service.SetActive()

	lines := json.NewDecoder(strings.NewReader(page))
	for {
		var line crtData
		if err := lines.Decode(&line); err == io.EOF {
			break
		} else if err != nil {
			c.Service.Config().Log.Printf("%s: %v", url, err)
		}
		unique = append(unique, line.Name)
	}
	return unique
}

func (c *Crtsh) getURL(domain string) string {
	return "https://crt.sh/?q=%25." + domain + "&output=json"
}
<commit_msg>Update crtsh.go<commit_after>// Copyright 2017 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package sources

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

	"github.com/OWASP/Amass/amass/core"
	"github.com/OWASP/Amass/amass/utils"
)

// Crtsh is data source object type that implements the DataSource interface.
type Crtsh struct {
	BaseDataSource
}

type crtData struct {
	IssuerID          int    `json:"issuer_ca_id"`
	IssuerName        string `json:"issuer_name"`
	Name              string `json:"name_value"`
	MinCertID         int    `json:"min_cert_id"`
	MinEntryTimestamp string `json:"min_entry_timestamp"`
	NotBefore         string `json:"not_before"`
	NotAfter          string `json:"not_after"`
}

// NewCrtsh returns an initialized Crtsh as a DataSource.
func NewCrtsh(srv core.AmassService) DataSource {
	c := new(Crtsh)

	c.BaseDataSource = *NewBaseDataSource(srv, core.CERT, "crt.sh")
	return c
}

// Query returns the subdomain names discovered when querying this data source.
func (c *Crtsh) Query(domain, sub string) []string {
	var unique []string

	if domain != sub {
		return unique
	}
	// Pull the page that lists all certs for this domain
	url := c.getURL(domain)
	page, err := utils.GetWebPage(url, nil)
	if err != nil {
		c.Service.Config().Log.Printf("%s: %v", url, err)
		return unique
	}
	c.Service.SetActive()

	lines := json.NewDecoder(strings.NewReader(page))
	for {
		var line crtData
		if err := lines.Decode(&line); err == io.EOF {
			break
		} else if err != nil {
			c.Service.Config().Log.Printf("%s: %v", url, err)
		}
		unique = utils.UniqueAppend(unique, removeAsteriskLabel(line.Name))
	}
	return unique
}

func (c *Crtsh) getURL(domain string) string {
	return "https://crt.sh/?q=%25." + domain + "&output=json"
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"net"
	"time"

	"golang.org/x/crypto/ssh"

	"log"

	"github.com/tg123/sshpiper/sshpiperd/auditor"
	"github.com/tg123/sshpiper/sshpiperd/challenger"
	"github.com/tg123/sshpiper/sshpiperd/registry"
	"github.com/tg123/sshpiper/sshpiperd/upstream"
)

type piperdConfig struct {
	ListenAddr     string        `short:"l" long:"listen" description:"Listening Address" default:"0.0.0.0" env:"SSHPIPERD_LISTENADDR" ini-name:"listen-address"`
	Port           uint          `short:"p" long:"port" description:"Listening Port" default:"2222" env:"SSHPIPERD_PORT" ini-name:"listen-port"`
	PiperKeyFile   string        `short:"i" long:"server-key" description:"Server key file for SSH Piper" default:"/etc/ssh/ssh_host_rsa_key" env:"SSHPIPERD_SERVER_KEY" ini-name:"server-key"`
	LoginGraceTime time.Duration `long:"login-grace-time" description:"Piper disconnects after this time if the pipe has not successfully established" default:"30s" env:"SSHPIPERD_LOGIN_GRACETIME" ini-name:"login-grace-time"`

	UpstreamDriver   string `short:"u" long:"upstream-driver" description:"Upstream provider driver" default:"workingdir" env:"SSHPIPERD_UPSTREAM_DRIVER" ini-name:"upstream-driver"`
	ChallengerDriver string `short:"c" long:"challenger-driver" description:"Additional challenger name, e.g. pam, empty for no additional challenge" env:"SSHPIPERD_CHALLENGER" ini-name:"challenger-driver"`
	AuditorDriver    string `long:"auditor-driver" description:"Auditor for ssh connections piped by SSH Piper" env:"SSHPIPERD_AUDITOR" ini-name:"auditor-driver"`

	BannerText string `long:"banner-text" description:"Display a banner before authentication, would be ignored if banner file was set" env:"SSHPIPERD_BANNERTEXT" ini-name:"banner-text" `
	BannerFile string `long:"banner-file" description:"Display a banner from file before authentication" env:"SSHPIPERD_BANNERFILE" ini-name:"banner-file" `
}

func getAndInstall(reg, name string, get func(n string) registry.Plugin, install func(plugin registry.Plugin) error, logger *log.Logger) error {
	if name == "" {
		return nil
	}

	p := get(name)

	if p == nil {
		return fmt.Errorf("%v driver %v not found", reg, name)
	}

	err := p.Init(logger)
	if err != nil {
		return err
	}
	return install(p)
}

func installDrivers(piper *ssh.PiperConfig, config *piperdConfig, logger *log.Logger) (auditor.Provider, error) {

	// install upstreamProvider driver
	if config.UpstreamDriver == "" {
		return nil, fmt.Errorf("must provider upstream driver")
	}

	var bigbro auditor.Provider

	for _, d := range []struct {
		reg     string
		name    string
		get     func(n string) registry.Plugin
		install func(plugin registry.Plugin) error
	}{
		// upstream driver
		{
			"Upstream",
			config.UpstreamDriver,
			func(n string) registry.Plugin {
				return upstream.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(upstream.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("upstream driver return nil handler")
				}

				piper.FindUpstream = handler
				return nil
			},
		},
		// challenger driver
		{
			"Challenger",
			config.ChallengerDriver,
			func(n string) registry.Plugin {
				return challenger.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(challenger.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("challenger driver return nil handler")
				}

				piper.AdditionalChallenge = handler
				return nil
			},
		},
		// auditor driver
		{
			"Auditor",
			config.AuditorDriver,
			func(n string) registry.Plugin {
				return auditor.Get(n)
			},
			func(plugin registry.Plugin) error {
				bigbro = plugin.(auditor.Provider)
				return nil
			},
		},
	} {
		err := getAndInstall(d.reg, d.name, d.get, d.install, logger)
		if err != nil {
			return nil, err
		}
	}

	return bigbro, nil
}

func startPiper(config *piperdConfig, logger *log.Logger) error {

	logger.Println("sshpiper is about to start")

	piper := &ssh.PiperConfig{}

	// drivers
	bigbro, err := installDrivers(piper, config, logger)
	if err != nil {
		return err
	}

	// listeners
	privateBytes, err := ioutil.ReadFile(config.PiperKeyFile)
	if err != nil {
		return err
	}

	private, err := ssh.ParsePrivateKey(privateBytes)
	if err != nil {
		return err
	}

	piper.AddHostKey(private)

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.ListenAddr, config.Port))
	if err != nil {
		return fmt.Errorf("failed to listen for connection: %v", err)
	}
	defer listener.Close()

	// banner
	if config.BannerFile != "" {

		piper.BannerCallback = func(conn ssh.ConnMetadata) string {

			msg, err := ioutil.ReadFile(config.BannerFile)

			if err != nil {
				logger.Printf("failed to read banner file: %v", err)
				return ""
			}

			return string(msg)
		}
	} else if config.BannerText != "" {
		piper.BannerCallback = func(conn ssh.ConnMetadata) string {
			return config.BannerText + "\n"
		}
	}

	logger.Printf("sshpiperd started")

	for {
		conn, err := listener.Accept()
		if err != nil {
			logger.Printf("failed to accept connection: %v", err)
			continue
		}

		logger.Printf("connection accepted: %v", conn.RemoteAddr())

		go func(c net.Conn) {
			defer c.Close()

			pipec := make(chan *ssh.PiperConn, 0)
			errorc := make(chan error, 0)

			go func() {
				p, err := ssh.NewSSHPiperConn(c, piper)

				if err != nil {
					errorc <- err
					return
				}

				pipec <- p
			}()

			var p *ssh.PiperConn

			select {
			case p = <-pipec:
			case err := <-errorc:
				logger.Printf("connection from %v establishing failed reason: %v", c.RemoteAddr(), err)
				return
			case <-time.After(config.LoginGraceTime):
				logger.Printf("pipe establishing timeout, disconnected connection from %v", c.RemoteAddr())
				return
			}

			defer p.Close()

			if bigbro != nil {
				a, err := bigbro.Create(p.DownstreamConnMeta())
				if err != nil {
					logger.Printf("connection from %v failed to create auditor reason: %v", c.RemoteAddr(), err)
					return
				}
				defer a.Close()

				p.HookUpstreamMsg = a.GetUpstreamHook()
				p.HookDownstreamMsg = a.GetDownstreamHook()
			}

			err = p.Wait()
			logger.Printf("connection from %v closed reason: %v", c.RemoteAddr(), err)
		}(conn)
	}
}
<commit_msg>Allow multiple host keys to be loaded by specifying a globbable path (#61)<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"net"
	"path/filepath"
	"time"

	"golang.org/x/crypto/ssh"

	"log"

	"github.com/tg123/sshpiper/sshpiperd/auditor"
	"github.com/tg123/sshpiper/sshpiperd/challenger"
	"github.com/tg123/sshpiper/sshpiperd/registry"
	"github.com/tg123/sshpiper/sshpiperd/upstream"
)

type piperdConfig struct {
	ListenAddr     string        `short:"l" long:"listen" description:"Listening Address" default:"0.0.0.0" env:"SSHPIPERD_LISTENADDR" ini-name:"listen-address"`
	Port           uint          `short:"p" long:"port" description:"Listening Port" default:"2222" env:"SSHPIPERD_PORT" ini-name:"listen-port"`
	PiperKeyFile   string        `short:"i" long:"server-key" description:"Server key file for SSH Piper" default:"/etc/ssh/ssh_host_rsa_key" env:"SSHPIPERD_SERVER_KEY" ini-name:"server-key"`
	LoginGraceTime time.Duration `long:"login-grace-time" description:"Piper disconnects after this time if the pipe has not successfully established" default:"30s" env:"SSHPIPERD_LOGIN_GRACETIME" ini-name:"login-grace-time"`

	UpstreamDriver   string `short:"u" long:"upstream-driver" description:"Upstream provider driver" default:"workingdir" env:"SSHPIPERD_UPSTREAM_DRIVER" ini-name:"upstream-driver"`
	ChallengerDriver string `short:"c" long:"challenger-driver" description:"Additional challenger name, e.g. pam, empty for no additional challenge" env:"SSHPIPERD_CHALLENGER" ini-name:"challenger-driver"`
	AuditorDriver    string `long:"auditor-driver" description:"Auditor for ssh connections piped by SSH Piper" env:"SSHPIPERD_AUDITOR" ini-name:"auditor-driver"`

	BannerText string `long:"banner-text" description:"Display a banner before authentication, would be ignored if banner file was set" env:"SSHPIPERD_BANNERTEXT" ini-name:"banner-text" `
	BannerFile string `long:"banner-file" description:"Display a banner from file before authentication" env:"SSHPIPERD_BANNERFILE" ini-name:"banner-file" `
}

func getAndInstall(reg, name string, get func(n string) registry.Plugin, install func(plugin registry.Plugin) error, logger *log.Logger) error {
	if name == "" {
		return nil
	}

	p := get(name)

	if p == nil {
		return fmt.Errorf("%v driver %v not found", reg, name)
	}

	err := p.Init(logger)
	if err != nil {
		return err
	}
	return install(p)
}

func installDrivers(piper *ssh.PiperConfig, config *piperdConfig, logger *log.Logger) (auditor.Provider, error) {

	// install upstreamProvider driver
	if config.UpstreamDriver == "" {
		return nil, fmt.Errorf("must provider upstream driver")
	}

	var bigbro auditor.Provider

	for _, d := range []struct {
		reg     string
		name    string
		get     func(n string) registry.Plugin
		install func(plugin registry.Plugin) error
	}{
		// upstream driver
		{
			"Upstream",
			config.UpstreamDriver,
			func(n string) registry.Plugin {
				return upstream.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(upstream.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("upstream driver return nil handler")
				}

				piper.FindUpstream = handler
				return nil
			},
		},
		// challenger driver
		{
			"Challenger",
			config.ChallengerDriver,
			func(n string) registry.Plugin {
				return challenger.Get(n)
			},
			func(plugin registry.Plugin) error {
				handler := plugin.(challenger.Provider).GetHandler()

				if handler == nil {
					return fmt.Errorf("challenger driver return nil handler")
				}

				piper.AdditionalChallenge = handler
				return nil
			},
		},
		// auditor driver
		{
			"Auditor",
			config.AuditorDriver,
			func(n string) registry.Plugin {
				return auditor.Get(n)
			},
			func(plugin registry.Plugin) error {
				bigbro = plugin.(auditor.Provider)
				return nil
			},
		},
	} {
		err := getAndInstall(d.reg, d.name, d.get, d.install, logger)
		if err != nil {
			return nil, err
		}
	}

	return bigbro, nil
}

func startPiper(config *piperdConfig, logger *log.Logger) error {

	logger.Println("sshpiper is about to start")

	piper := &ssh.PiperConfig{}

	// drivers
	bigbro, err := installDrivers(piper, config, logger)
	if err != nil {
		return err
	}

	// listeners
	privateKeys, err := filepath.Glob(config.PiperKeyFile)
	if err != nil {
		return err
	}

	logger.Println("Found host keys", privateKeys)
	for _, privateKey := range privateKeys {
		logger.Println("Loading host key", privateKey)
		privateBytes, err := ioutil.ReadFile(privateKey)
		if err != nil {
			return err
		}

		private, err := ssh.ParsePrivateKey(privateBytes)
		if err != nil {
			return err
		}

		piper.AddHostKey(private)
	}

	listener, err := net.Listen("tcp", fmt.Sprintf("%s:%d", config.ListenAddr, config.Port))
	if err != nil {
		return fmt.Errorf("failed to listen for connection: %v", err)
	}
	defer listener.Close()

	// banner
	if config.BannerFile != "" {

		piper.BannerCallback = func(conn ssh.ConnMetadata) string {

			msg, err := ioutil.ReadFile(config.BannerFile)

			if err != nil {
				logger.Printf("failed to read banner file: %v", err)
				return ""
			}

			return string(msg)
		}
	} else if config.BannerText != "" {
		piper.BannerCallback = func(conn ssh.ConnMetadata) string {
			return config.BannerText + "\n"
		}
	}

	logger.Printf("sshpiperd started")

	for {
		conn, err := listener.Accept()
		if err != nil {
			logger.Printf("failed to accept connection: %v", err)
			continue
		}

		logger.Printf("connection accepted: %v", conn.RemoteAddr())

		go func(c net.Conn) {
			defer c.Close()

			pipec := make(chan *ssh.PiperConn, 0)
			errorc := make(chan error, 0)

			go func() {
				p, err := ssh.NewSSHPiperConn(c, piper)

				if err != nil {
					errorc <- err
					return
				}

				pipec <- p
			}()

			var p *ssh.PiperConn

			select {
			case p = <-pipec:
			case err := <-errorc:
				logger.Printf("connection from %v establishing failed reason: %v", c.RemoteAddr(), err)
				return
			case <-time.After(config.LoginGraceTime):
				logger.Printf("pipe establishing timeout, disconnected connection from %v", c.RemoteAddr())
				return
			}

			defer p.Close()

			if bigbro != nil {
				a, err := bigbro.Create(p.DownstreamConnMeta())
				if err != nil {
					logger.Printf("connection from %v failed to create auditor reason: %v", c.RemoteAddr(), err)
					return
				}
				defer a.Close()

				p.HookUpstreamMsg = a.GetUpstreamHook()
				p.HookDownstreamMsg = a.GetDownstreamHook()
			}

			err = p.Wait()
			logger.Printf("connection from %v closed reason: %v", c.RemoteAddr(), err)
		}(conn)
	}
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

// CmdDatacenter subcommand
import (
	"errors"
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/fatih/color"
	"github.com/urfave/cli"
)

// ListDatacenters ...
var ListDatacenters = cli.Command{
	Name:      "list",
	Usage:     "List available datacenters.",
	ArgsUsage: " ",
	Description: `List available datacenters.

   Example:
    $ ernest datacenter list
	`,
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		datacenters, err := m.ListDatacenters(cfg.Token)
		if err != nil {
			color.Red(err.Error())
			return err
		}

		w := new(tabwriter.Writer)
		w.Init(os.Stdout, 0, 8, 0, '\t', 0)

		fmt.Fprintln(w, "NAME\tID")
		for _, datacenter := range datacenters {
			str := fmt.Sprintf("%s\t%d", datacenter.Name, datacenter.ID)
			fmt.Fprintln(w, str)
		}
		w.Flush()
		return nil
	},
}

// CreateDatacenter ...
var CreateDatacenter = cli.Command{
	Name:  "create",
	Usage: "Create a new datacenter.",
	Description: `Create a new datacenter on the targeted instance of Ernest.

   Example:
    $ ernest datacenter create --datacenter-user username --datacenter-password xxxx --datacenter-org MY-ORG-NAME --vse-url http://vse.url mydatacenter https://myernest.com MY-PUBLIC-NETWORK
	`,
	ArgsUsage: "<datacenter-name> <vcloud-url> <network-name>",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "datacenter-user",
			Value: "",
			Usage: "User to be configured with the datacenter",
		},
		cli.StringFlag{
			Name:  "datacenter-password",
			Value: "",
			Usage: "Password related with user",
		},
		cli.StringFlag{
			Name:  "datacenter-org",
			Value: "",
			Usage: "vCloud Organization name",
		},
		cli.StringFlag{
			Name:  "vse-url",
			Value: "",
			Usage: "VSE URL",
		},
	},
	Action: func(c *cli.Context) error {
		if len(c.Args()) < 3 {
			msg := "You should specify the datacenter name, vcloud URL and network name"
			color.Red(msg)
			return errors.New(msg)
		}
		name := c.Args()[0]
		user := c.String("datacenter-user") + "@" + c.String("datacenter-org")
		if user == "" {
			msg := "User not specified"
			color.Red(msg)
			return errors.New(msg)
		}
		password := c.String("datacenter-password")
		if password == "" {
			msg := "Password not specified"
			color.Red(msg)
			return errors.New("Password not specified")
		}
		m, cfg := setup(c)
		_, err := m.CreateDatacenter(cfg.Token, name, user, password, c.Args()[1], c.Args()[2], c.String("vse-url"))
		if err != nil {
			color.Red(err.Error())
		}
		return nil
	},
}

// CmdDatacenter ...
var CmdDatacenter = cli.Command{
	Name:  "datacenter",
	Usage: "Datacenter related subcommands",
	Subcommands: []cli.Command{
		ListDatacenters,
		CreateDatacenter,
	},
}
<commit_msg>Show a message if there is no datacenters created.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

package main

// CmdDatacenter subcommand
import (
	"errors"
	"fmt"
	"os"
	"text/tabwriter"

	"github.com/fatih/color"
	"github.com/urfave/cli"
)

// ListDatacenters ...
var ListDatacenters = cli.Command{
	Name:      "list",
	Usage:     "List available datacenters.",
	ArgsUsage: " ",
	Description: `List available datacenters.

   Example:
    $ ernest datacenter list
	`,
	Action: func(c *cli.Context) error {
		m, cfg := setup(c)
		datacenters, err := m.ListDatacenters(cfg.Token)
		if err != nil {
			color.Red(err.Error())
			return err
		}

		w := new(tabwriter.Writer)
		w.Init(os.Stdout, 0, 8, 0, '\t', 0)

		if len(datacenters) == 0 {
			fmt.Println("There is no datacenters created yet.")
			return nil
		}

		fmt.Fprintln(w, "NAME\tID")
		for _, datacenter := range datacenters {
			str := fmt.Sprintf("%s\t%d", datacenter.Name, datacenter.ID)
			fmt.Fprintln(w, str)
		}
		w.Flush()
		return nil
	},
}

// CreateDatacenter ...
var CreateDatacenter = cli.Command{
	Name:  "create",
	Usage: "Create a new datacenter.",
	Description: `Create a new datacenter on the targeted instance of Ernest.

   Example:
    $ ernest datacenter create --datacenter-user username --datacenter-password xxxx --datacenter-org MY-ORG-NAME --vse-url http://vse.url mydatacenter https://myernest.com MY-PUBLIC-NETWORK
	`,
	ArgsUsage: "<datacenter-name> <vcloud-url> <network-name>",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "datacenter-user",
			Value: "",
			Usage: "User to be configured with the datacenter",
		},
		cli.StringFlag{
			Name:  "datacenter-password",
			Value: "",
			Usage: "Password related with user",
		},
		cli.StringFlag{
			Name:  "datacenter-org",
			Value: "",
			Usage: "vCloud Organization name",
		},
		cli.StringFlag{
			Name:  "vse-url",
			Value: "",
			Usage: "VSE URL",
		},
	},
	Action: func(c *cli.Context) error {
		if len(c.Args()) < 3 {
			msg := "You should specify the datacenter name, vcloud URL and network name"
			color.Red(msg)
			return errors.New(msg)
		}
		name := c.Args()[0]
		user := c.String("datacenter-user") + "@" + c.String("datacenter-org")
		if user == "" {
			msg := "User not specified"
			color.Red(msg)
			return errors.New(msg)
		}
		password := c.String("datacenter-password")
		if password == "" {
			msg := "Password not specified"
			color.Red(msg)
			return errors.New("Password not specified")
		}
		m, cfg := setup(c)
		_, err := m.CreateDatacenter(cfg.Token, name, user, password, c.Args()[1], c.Args()[2], c.String("vse-url"))
		if err != nil {
			color.Red(err.Error())
		}
		return nil
	},
}

// CmdDatacenter ...
var CmdDatacenter = cli.Command{
	Name:  "datacenter",
	Usage: "Datacenter related subcommands",
	Subcommands: []cli.Command{
		ListDatacenters,
		CreateDatacenter,
	},
}
<|endoftext|>