<commit_before>package main

import (
	"encoding/csv"
	"flag"
	"io"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const expectedCsvFieldCount = 52

var (
	totalScrapes     = prometheus.NewCounter()
	scrapeFailures   = prometheus.NewCounter()
	csvParseFailures = prometheus.NewCounter()
)

type registry struct {
	prometheus.Registry
	serviceMetrics map[int]prometheus.Gauge
	backendMetrics map[int]prometheus.Gauge
}

func newRegistry() *registry {
	r := &registry{prometheus.NewRegistry(), make(map[int]prometheus.Gauge), make(map[int]prometheus.Gauge)}

	r.Register("haproxy_exporter_total_scrapes", "Current total HAProxy scrapes.", prometheus.NilLabels, totalScrapes)
	r.Register("haproxy_exporter_scrape_failures", "Number of errors while scraping HAProxy.", prometheus.NilLabels, scrapeFailures)
	r.Register("haproxy_exporter_csv_parse_failures", "Number of errors while parsing CSV.", prometheus.NilLabels, csvParseFailures)

	r.serviceMetrics = map[int]prometheus.Gauge{
		2: r.newGauge("haproxy_current_queue", "Current server queue length."),
		3: r.newGauge("haproxy_max_queue", "Maximum server queue length."),
	}

	r.backendMetrics = map[int]prometheus.Gauge{
		4:  r.newGauge("haproxy_current_sessions", "Current number of active sessions."),
		5:  r.newGauge("haproxy_max_sessions", "Maximum number of active sessions."),
		8:  r.newGauge("haproxy_bytes_in", "Current total of incoming bytes."),
		9:  r.newGauge("haproxy_bytes_out", "Current total of outgoing bytes."),
		17: r.newGauge("haproxy_instance_up", "Current health status of the instance (1 = UP, 0 = DOWN)."),
		33: r.newGauge("haproxy_current_session_rate", "Current number of sessions per second."),
		35: r.newGauge("haproxy_max_session_rate", "Maximum number of sessions per second."),
	}

	return r
}

func (r *registry) newGauge(metricName string, docString string) prometheus.Gauge {
	gauge := prometheus.NewGauge()
	r.Register(metricName, docString, prometheus.NilLabels, gauge)
	return gauge
}

// Exporter collects HAProxy stats from the given URI and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI   string
	reg   *registry
	mutex sync.RWMutex
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string) *Exporter {
	return &Exporter{
		URI: uri,
		reg: newRegistry(),
	}
}

// Registry returns a prometheus.Registry type with the complete state of the
// last stats collection run.
func (e *Exporter) Registry() prometheus.Registry {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	return e.reg
}

// Handler returns a http.HandlerFunc of the last finished prometheus.Registry.
func (e *Exporter) Handler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		reg := e.Registry()
		f := reg.Handler()

		f(w, r)
	}
}

// Scrape fetches the stats from configured HAProxy location. It creates a new
// prometheus.Registry object every time to not leak stale data from previous
// collections.
func (e *Exporter) Scrape() {
	csvRows := make(chan []string)
	quitChan := make(chan bool)
	reg := newRegistry()

	go e.scrape(csvRows, quitChan)
	reg.exportMetrics(csvRows, quitChan)

	e.mutex.Lock()
	defer e.mutex.Unlock()

	e.reg = reg
}

// ScrapePeriodically runs the Scrape function in the specified interval.
func (e *Exporter) ScrapePeriodically(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for _ = range ticker.C {
		e.Scrape()
	}
}

func (e *Exporter) scrape(csvRows chan []string, quitChan chan bool) {
	defer close(quitChan)
	defer totalScrapes.Increment(prometheus.NilLabels)

	resp, err := http.Get(e.URI)
	if err != nil {
		log.Printf("Error while scraping HAProxy: %v", err)
		scrapeFailures.Increment(prometheus.NilLabels)
		return
	}
	defer resp.Body.Close()

	reader := csv.NewReader(resp.Body)
	reader.TrailingComma = true
	reader.Comment = '#'

	for {
		row, err := reader.Read()

		if err == io.EOF {
			break
		}
		if err != nil {
			log.Printf("Error while reading CSV: %v", err)
			csvParseFailures.Increment(prometheus.NilLabels)
			return
		}
		if len(row) == 0 {
			continue
		}

		csvRows <- row
	}

}

func (r *registry) exportMetrics(csvRows chan []string, quitChan chan bool) {
	for {
		select {
		case row := <-csvRows:
			r.exportCsvRow(row)
		case <-quitChan:
			return
		}
	}
}

func (r *registry) exportCsvRow(csvRow []string) {
	if len(csvRow) != expectedCsvFieldCount {
		log.Printf("Wrong CSV field count: %d vs. %d", len(csvRow), expectedCsvFieldCount)
		csvParseFailures.Increment(prometheus.NilLabels)
		return
	}

	service, instance := csvRow[0], csvRow[1]

	if instance == "FRONTEND" {
		return
	}

	if instance == "BACKEND" {
		labels := map[string]string{
			"service": service,
		}

		exportCsvFields(labels, r.serviceMetrics, csvRow)
	} else {
		labels := map[string]string{
			"service":  service,
			"instance": instance,
		}

		exportCsvFields(labels, r.backendMetrics, csvRow)
	}
}

func exportCsvFields(labels map[string]string, fields map[int]prometheus.Gauge, csvRow []string) {
	for fieldIdx, gauge := range fields {
		valueStr := csvRow[fieldIdx]
		if valueStr == "" {
			continue
		}

		var value int64
		var err error
		switch valueStr {
		// UP or UP going down
		case "UP", "UP 1/3", "UP 2/3":
			value = 1
		// DOWN or DOWN going up
		case "DOWN", "DOWN 1/2":
			value = 0
		case "OPEN":
			value = 0
		default:
			value, err = strconv.ParseInt(valueStr, 10, 64)
			if err != nil {
				log.Printf("Error while parsing CSV field value %s: %v", valueStr, err)
				csvParseFailures.Increment(prometheus.NilLabels)
				continue
			}
		}
		gauge.Set(labels, float64(value))
	}
}

func main() {
	var (
		listeningAddress      = flag.String("telemetry.address", ":8080", "Address on which to expose JSON metrics.")
		metricsEndpoint       = flag.String("telemetry.endpoint", prometheus.ExpositionResource, "Path under which to expose metrics.")
		haProxyScrapeUri      = flag.String("haproxy.scrape_uri", "http://localhost/;csv", "URI on which to scrape HAProxy.")
		haProxyScrapeInterval = flag.Duration("haproxy.scrape_interval", 15, "Interval in seconds between scrapes.")
	)
	flag.Parse()

	exporter := NewExporter(*haProxyScrapeUri)
	go exporter.ScrapePeriodically(*haProxyScrapeInterval * time.Second)

	log.Printf("Starting Server: %s", *listeningAddress)
	http.Handle(*metricsEndpoint, exporter.Handler())
	log.Fatal(http.ListenAndServe(*listeningAddress, nil))
}
<commit_msg>Fix default interval value<commit_after>package main

import (
	"encoding/csv"
	"flag"
	"io"
	"log"
	"net/http"
	"strconv"
	"sync"
	"time"

	"github.com/prometheus/client_golang/prometheus"
)

const expectedCsvFieldCount = 52

var (
	totalScrapes     = prometheus.NewCounter()
	scrapeFailures   = prometheus.NewCounter()
	csvParseFailures = prometheus.NewCounter()
)

type registry struct {
	prometheus.Registry
	serviceMetrics map[int]prometheus.Gauge
	backendMetrics map[int]prometheus.Gauge
}

func newRegistry() *registry {
	r := &registry{prometheus.NewRegistry(), make(map[int]prometheus.Gauge), make(map[int]prometheus.Gauge)}

	r.Register("haproxy_exporter_total_scrapes", "Current total HAProxy scrapes.", prometheus.NilLabels, totalScrapes)
	r.Register("haproxy_exporter_scrape_failures", "Number of errors while scraping HAProxy.", prometheus.NilLabels, scrapeFailures)
	r.Register("haproxy_exporter_csv_parse_failures", "Number of errors while parsing CSV.", prometheus.NilLabels, csvParseFailures)

	r.serviceMetrics = map[int]prometheus.Gauge{
		2: r.newGauge("haproxy_current_queue", "Current server queue length."),
		3: r.newGauge("haproxy_max_queue", "Maximum server queue length."),
	}

	r.backendMetrics = map[int]prometheus.Gauge{
		4:  r.newGauge("haproxy_current_sessions", "Current number of active sessions."),
		5:  r.newGauge("haproxy_max_sessions", "Maximum number of active sessions."),
		8:  r.newGauge("haproxy_bytes_in", "Current total of incoming bytes."),
		9:  r.newGauge("haproxy_bytes_out", "Current total of outgoing bytes."),
		17: r.newGauge("haproxy_instance_up", "Current health status of the instance (1 = UP, 0 = DOWN)."),
		33: r.newGauge("haproxy_current_session_rate", "Current number of sessions per second."),
		35: r.newGauge("haproxy_max_session_rate", "Maximum number of sessions per second."),
	}

	return r
}

func (r *registry) newGauge(metricName string, docString string) prometheus.Gauge {
	gauge := prometheus.NewGauge()
	r.Register(metricName, docString, prometheus.NilLabels, gauge)
	return gauge
}

// Exporter collects HAProxy stats from the given URI and exports them using
// the prometheus metrics package.
type Exporter struct {
	URI   string
	reg   *registry
	mutex sync.RWMutex
}

// NewExporter returns an initialized Exporter.
func NewExporter(uri string) *Exporter {
	return &Exporter{
		URI: uri,
		reg: newRegistry(),
	}
}

// Registry returns a prometheus.Registry type with the complete state of the
// last stats collection run.
func (e *Exporter) Registry() prometheus.Registry {
	e.mutex.RLock()
	defer e.mutex.RUnlock()

	return e.reg
}

// Handler returns a http.HandlerFunc of the last finished prometheus.Registry.
func (e *Exporter) Handler() http.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request) {
		reg := e.Registry()
		f := reg.Handler()

		f(w, r)
	}
}

// Scrape fetches the stats from configured HAProxy location. It creates a new
// prometheus.Registry object every time to not leak stale data from previous
// collections.
func (e *Exporter) Scrape() {
	csvRows := make(chan []string)
	quitChan := make(chan bool)
	reg := newRegistry()

	go e.scrape(csvRows, quitChan)
	reg.exportMetrics(csvRows, quitChan)

	e.mutex.Lock()
	defer e.mutex.Unlock()

	e.reg = reg
}

// ScrapePeriodically runs the Scrape function in the specified interval.
func (e *Exporter) ScrapePeriodically(interval time.Duration) {
	ticker := time.NewTicker(interval)
	defer ticker.Stop()

	for _ = range ticker.C {
		e.Scrape()
	}
}

func (e *Exporter) scrape(csvRows chan []string, quitChan chan bool) {
	defer close(quitChan)
	defer totalScrapes.Increment(prometheus.NilLabels)

	resp, err := http.Get(e.URI)
	if err != nil {
		log.Printf("Error while scraping HAProxy: %v", err)
		scrapeFailures.Increment(prometheus.NilLabels)
		return
	}
	defer resp.Body.Close()

	reader := csv.NewReader(resp.Body)
	reader.TrailingComma = true
	reader.Comment = '#'

	for {
		row, err := reader.Read()

		if err == io.EOF {
			break
		}
		if err != nil {
			log.Printf("Error while reading CSV: %v", err)
			csvParseFailures.Increment(prometheus.NilLabels)
			return
		}
		if len(row) == 0 {
			continue
		}

		csvRows <- row
	}

}

func (r *registry) exportMetrics(csvRows chan []string, quitChan chan bool) {
	for {
		select {
		case row := <-csvRows:
			r.exportCsvRow(row)
		case <-quitChan:
			return
		}
	}
}

func (r *registry) exportCsvRow(csvRow []string) {
	if len(csvRow) != expectedCsvFieldCount {
		log.Printf("Wrong CSV field count: %d vs. %d", len(csvRow), expectedCsvFieldCount)
		csvParseFailures.Increment(prometheus.NilLabels)
		return
	}

	service, instance := csvRow[0], csvRow[1]

	if instance == "FRONTEND" {
		return
	}

	if instance == "BACKEND" {
		labels := map[string]string{
			"service": service,
		}

		exportCsvFields(labels, r.serviceMetrics, csvRow)
	} else {
		labels := map[string]string{
			"service":  service,
			"instance": instance,
		}

		exportCsvFields(labels, r.backendMetrics, csvRow)
	}
}

func exportCsvFields(labels map[string]string, fields map[int]prometheus.Gauge, csvRow []string) {
	for fieldIdx, gauge := range fields {
		valueStr := csvRow[fieldIdx]
		if valueStr == "" {
			continue
		}

		var value int64
		var err error
		switch valueStr {
		// UP or UP going down
		case "UP", "UP 1/3", "UP 2/3":
			value = 1
		// DOWN or DOWN going up
		case "DOWN", "DOWN 1/2":
			value = 0
		case "OPEN":
			value = 0
		default:
			value, err = strconv.ParseInt(valueStr, 10, 64)
			if err != nil {
				log.Printf("Error while parsing CSV field value %s: %v", valueStr, err)
				csvParseFailures.Increment(prometheus.NilLabels)
				continue
			}
		}
		gauge.Set(labels, float64(value))
	}
}

func main() {
	var (
		listeningAddress      = flag.String("telemetry.address", ":8080", "Address on which to expose JSON metrics.")
		metricsEndpoint       = flag.String("telemetry.endpoint", prometheus.ExpositionResource, "Path under which to expose metrics.")
		haProxyScrapeUri      = flag.String("haproxy.scrape_uri", "http://localhost/;csv", "URI on which to scrape HAProxy.")
		haProxyScrapeInterval = flag.Duration("haproxy.scrape_interval", 15 * time.Second, "Interval in seconds between scrapes.")
	)
	flag.Parse()

	exporter := NewExporter(*haProxyScrapeUri)
	go exporter.ScrapePeriodically(*haProxyScrapeInterval)

	log.Printf("Starting Server: %s", *listeningAddress)
	http.Handle(*metricsEndpoint, exporter.Handler())
	log.Fatal(http.ListenAndServe(*listeningAddress, nil))
}
<|endoftext|>


<commit_before>package main

import "fmt"

import "encoding/json"
import "io/ioutil"
import "net/http"
import "net/url"

// `q` is the YQL query
func yql(jrsp interface{}, q string) (err error) {
	// form the YQL URL:
	u := `http://query.yahooapis.com/v1/public/yql?q=` + url.QueryEscape(q) + `&format=json&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys`
	resp, err := http.Get(u)
	if err != nil {
		return
	}

	// read body:
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	// print JSON to console:
	//fmt.Printf("%s\n\n", body)

	// decode JSON:
	err = json.Unmarshal(body, jrsp)
	return
}

// Head to http://developer.yahoo.com/yql/console/?q=select%20*%20from%20yahoo.finance.quote%20where%20symbol%20in%20(%22YHOO%22%2C%22AAPL%22%2C%22GOOG%22%2C%22MSFT%22)&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys
// to understand this JSON structure.

type Quote struct {
	Symbol             string "symbol"
	LastTradePriceOnly string
}

type QuoteResponse struct {
	Query *struct {
		Results *struct {
			Quote *Quote "quote"
		} "results"
	} "query"
}

type Historical struct {
	Date  string
	Close string
}

type HistoricalResponse struct {
	Query *struct {
		Results *struct {
			Quote []*Historical "quote"
		} "results"
	} "query"
}

// main:
func main() {
	// get current price of MSFT:
	quot := new(QuoteResponse)
	err := yql(quot, `select symbol, LastTradePriceOnly from yahoo.finance.quote where symbol in ("MSFT")`)
	if err != nil {
		fmt.Print(err)
		return
	}

	fmt.Printf("%s\n", *quot.Query.Results.Quote)

	// get historical data for MSFT:
	hist := new(HistoricalResponse)
	err = yql(hist, `select Date, Close from yahoo.finance.historicaldata where symbol = "MSFT" and startDate = "2013-12-04" and endDate = "2013-12-06"`)
	if err != nil {
		fmt.Print(err)
		return
	}

	fmt.Printf("%s %s %s\n", *hist.Query.Results.Quote[0], *hist.Query.Results.Quote[1], *hist.Query.Results.Quote[2])
	return
}
<commit_msg>Added sqlite3 support.<commit_after>package main

import "fmt"

import "encoding/json"
import "io/ioutil"
import "net/http"
import "net/url"
import "database/sql"
import _ "github.com/mattn/go-sqlite3"

// `q` is the YQL query
func yql(jrsp interface{}, q string) (err error) {
	// form the YQL URL:
	u := `http://query.yahooapis.com/v1/public/yql?q=` + url.QueryEscape(q) + `&format=json&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys`
	resp, err := http.Get(u)
	if err != nil {
		return
	}

	// read body:
	defer resp.Body.Close()
	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return
	}

	// print JSON to console:
	//fmt.Printf("%s\n\n", body)

	// decode JSON:
	err = json.Unmarshal(body, jrsp)
	return
}

// Head to http://developer.yahoo.com/yql/console/?q=select%20*%20from%20yahoo.finance.quote%20where%20symbol%20in%20(%22YHOO%22%2C%22AAPL%22%2C%22GOOG%22%2C%22MSFT%22)&env=store%3A%2F%2Fdatatables.org%2Falltableswithkeys
// to understand this JSON structure.

type Quote struct {
	Symbol             string "symbol"
	LastTradePriceOnly string
}

type QuoteResponse struct {
	Query *struct {
		Results *struct {
			Quote *Quote "quote"
		} "results"
	} "query"
}

type Historical struct {
	Date  string
	Close string
}

type HistoricalResponse struct {
	Query *struct {
		Results *struct {
			Quote []*Historical "quote"
		} "results"
	} "query"
}

func db_create_schema(db *sql.DB) (err error) {
	_, err = db.Exec(`create table if not exists stock (symbol TEXT UNIQUE, purchase_price NUMERIC, purchase_date TEXT, trailing_stop_percent NUMERIC, last_stop_price NUMERIC)`)
	if err != nil {
		return
	}

	_, err = db.Exec(`create table if not exists stock_history (symbol TEXT, date TEXT, closing_price NUMERIC)`)
	if err != nil {
		return
	}

	return
}

// main:
func main() {
	// using sqlite 3.8.0 release
	db, err := sql.Open("sqlite3", "stocks.db")
	if err != nil {
		fmt.Println(err)
		return
	}
	defer db.Close()

	// Create our schema:
	err = db_create_schema(db)
	if err != nil {
		fmt.Println(err)
		return
	}

	_, err = db.Query(`select * from stock`)
	if err != nil {
		fmt.Println(err)
		return
	}

	// get current price of MSFT:
	quot := new(QuoteResponse)
	err = yql(quot, `select symbol, LastTradePriceOnly from yahoo.finance.quote where symbol in ("MSFT")`)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("%s\n", *quot.Query.Results.Quote)

	// get historical data for MSFT:
	hist := new(HistoricalResponse)
	err = yql(hist, `select Date, Close from yahoo.finance.historicaldata where symbol = "MSFT" and startDate = "2013-12-04" and endDate = "2013-12-06"`)
	if err != nil {
		fmt.Println(err)
		return
	}

	fmt.Printf("%s %s %s\n", *hist.Query.Results.Quote[0], *hist.Query.Results.Quote[1], *hist.Query.Results.Quote[2])
	return
}
<|endoftext|>


<commit_before>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2014 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2016 Andriy Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy

import (
    "bufio"
    "math/rand"
    "net"
    "time"
    "strconv"
    "strings"

    "sippy/conf"
    "sippy/log"
    "sippy/types"
)

func NewRtpProxyClient(opts *rtpProxyClientOpts, config sippy_conf.Config, logger sippy_log.ErrorLogger) (sippy_types.RtpProxyClient, error) {
    rtpp := NewRtp_proxy_client_base(nil, config, opts, logger)
    err := rtpp.Init()
    return rtpp, err
}

type Rtp_proxy_client_base struct {
    heir            sippy_types.RtpProxyClient
    opts            *rtpProxyClientOpts
    transport       rtp_proxy_transport
    proxy_address   string
    online          bool
    sbind_supported bool
    tnot_supported  bool
    copy_supported  bool
    stat_supported  bool
    wdnt_supported  bool
    caps_done       bool
    shut_down       bool
    active_sessions int64
    sessions_created int64
    active_streams  int64
    preceived       int64
    ptransmitted    int64
    _CAPSTABLE      []struct{ vers string; attr *bool }
    logger          sippy_log.ErrorLogger
    global_config   sippy_conf.Config
}

type rtp_proxy_transport interface {
    is_local() bool
    send_command(string, func(string))
    shutdown()
}

func (self *Rtp_proxy_client_base) IsLocal() bool {
    return self.transport.is_local()
}

func (self *Rtp_proxy_client_base) IsOnline() bool {
    return self.online
}

func (self *Rtp_proxy_client_base) SBindSupported() bool {
    return self.sbind_supported
}

func (self *Rtp_proxy_client_base) TNotSupported() bool {
    return self.tnot_supported
}

func (self *Rtp_proxy_client_base) GetProxyAddress() string {
    return self.proxy_address
}

func (self *Rtp_proxy_client_base) me() sippy_types.RtpProxyClient {
    if self.heir != nil {
        return self.heir
    }
    return self
}

func NewRtp_proxy_client_base(heir sippy_types.RtpProxyClient, global_config sippy_conf.Config, opts *rtpProxyClientOpts, logger sippy_log.ErrorLogger) (*Rtp_proxy_client_base) {
    if opts == nil {
        opts = NewRtpProxyClientOpts() // default settings
    }
    self := &Rtp_proxy_client_base{
        heir            : heir,
        caps_done       : false,
        shut_down       : false,
        logger          : logger,
        opts            : opts,
        global_config   : global_config,
    }
    self._CAPSTABLE = []struct{ vers string; attr *bool }{
        { "20071218", &self.copy_supported },
        { "20080403", &self.stat_supported },
        { "20081224", &self.tnot_supported },
        { "20090810", &self.sbind_supported },
        { "20150617", &self.wdnt_supported },
    }
    return self
}

func (self *Rtp_proxy_client_base) Init() error {
    var err error
    var rtpp_class func(sippy_types.RtpProxyClient, sippy_conf.Config, net.Addr) (rtp_proxy_transport, error)
    var rtppa net.Addr

    a := self.opts.spath
    if strings.HasPrefix(a, "udp:") {
        tmp := strings.SplitN(a, ":", 3)
        if len(tmp) == 2 {
            rtppa, err = net.ResolveUDPAddr("udp", tmp[1] + ":22222")
        } else {
            rtppa, err = net.ResolveUDPAddr("udp", tmp[1] + ":" + tmp[2])
        }
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_udp
    } else if strings.HasPrefix(a, "udp6:") {
        tmp := strings.SplitN(a, ":", 2)
        a := tmp[1]
        rtp_proxy_host, rtp_proxy_port := a, "22222"
        if a[len(a)-1] != ']' {
            idx := strings.LastIndexByte(a, ':')
            if idx < 0 {
                rtp_proxy_host = a
            } else {
                rtp_proxy_host, rtp_proxy_port = a[:idx], a[idx+1:]
            }
        }
        if rtp_proxy_host[0] != '[' {
            rtp_proxy_host = "[" + rtp_proxy_host + "]"
        }
        rtppa, err = net.ResolveUDPAddr("udp", rtp_proxy_host + ":" + rtp_proxy_port)
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_udp
    } else if strings.HasPrefix(a, "tcp:") {
        tmp := strings.SplitN(a, ":", 3)
        if len(tmp) == 2 {
            rtppa, err = net.ResolveTCPAddr("tcp", tmp[1] + ":22222")
        } else {
            rtppa, err = net.ResolveTCPAddr("tcp", tmp[1] + ":" + tmp[2])
        }
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_stream
    } else if strings.HasPrefix(a, "tcp6:") {
        tmp := strings.SplitN(a, ":", 2)
        a := tmp[1]
        rtp_proxy_host, rtp_proxy_port := a, "22222"
        if a[len(a)-1] != ']' {
            idx := strings.LastIndexByte(a, ':')
            if idx < 0 {
                rtp_proxy_host = a
            } else {
                rtp_proxy_host, rtp_proxy_port = a[:idx], a[idx+1:]
            }
        }
        if rtp_proxy_host[0] != '[' {
            rtp_proxy_host = "[" + rtp_proxy_host + "]"
        }
        rtppa, err = net.ResolveTCPAddr("tcp", rtp_proxy_host + ":" + rtp_proxy_port)
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_stream
    } else {
        if strings.HasPrefix(a, "unix:") {
            rtppa, err = net.ResolveUnixAddr("unix", a[5:])
        } else if strings.HasPrefix(a, "cunix:") {
            rtppa, err = net.ResolveUnixAddr("unix", a[6:])
        } else {
            rtppa, err = net.ResolveUnixAddr("unix", a)
        }
        self.proxy_address = self.global_config.SipAddress().String()
        rtpp_class = newRtp_proxy_client_stream
    }
    self.transport, err = rtpp_class(self.me(), self.global_config, rtppa)
    if err != nil {
        return err
    }
    self.start()
    return nil
}

func (self *Rtp_proxy_client_base) InitWithAddress(address net.Addr) error {
    var err error
    if strings.HasPrefix(address.Network(), "udp") {
        self.transport, err = newRtp_proxy_client_udp(self.me(), self.global_config, address)
        if err != nil {
            return err
        }
        self.proxy_address, _, err = net.SplitHostPort(address.String())
        if err != nil {
            return err
        }
    } else {
        self.transport, err = newRtp_proxy_client_stream(self.me(), self.global_config, address)
        if err != nil {
            return err
        }
        self.proxy_address = self.global_config.SipAddress().String()
    }
    self.start()
    return nil
}

func (self *Rtp_proxy_client_base) start() {
    if ! self.opts.no_version_check {
        self.version_check()
    } else {
        self.caps_done = true
        self.online = true
    }
}

func (self *Rtp_proxy_client_base) SendCommand(cmd string, cb func(string)) {
    self.transport.send_command(cmd, cb)
}
/*
    def reconnect(self, *args, **kwargs):
        self.rtpp_class.reconnect(self, *args, **kwargs)
*/

func (self *Rtp_proxy_client_base) version_check() {
    if self.shut_down {
        return
    }
    self.transport.send_command("V", self.version_check_reply)
}

func (self *Rtp_proxy_client_base) version_check_reply(version string) {
    if self.shut_down {
        return
    }
    if version == "20040107" {
        self.me().GoOnline()
    } else if self.online {
        self.me().GoOffline()
    } else {
        StartTimeout(self.version_check, nil, randomize(self.opts.hrtb_retr_ival, 0.1), 1, self.logger)
    }
}

func (self *Rtp_proxy_client_base) heartbeat() {
    //print "heartbeat", self, self.address
    if self.shut_down {
        return
    }
    self.transport.send_command("Ib", self.heartbeat_reply)
}

func (self *Rtp_proxy_client_base) heartbeat_reply(stats string) {
    //print "heartbeat_reply", self.address, stats, self.online
    if self.shut_down || ! self.online {
        return
    }
    if stats == "" {
        self.active_sessions = 0
        self.me().GoOffline()
    } else {
        sessions_created := int64(0)
        active_sessions := int64(0)
        active_streams := int64(0)
        preceived := int64(0)
        ptransmitted := int64(0)
        scanner := bufio.NewScanner(strings.NewReader(stats))
        for scanner.Scan() {
            line_parts := strings.SplitN(scanner.Text(), ":", 2)
            if len(line_parts) != 2 { continue }
            switch line_parts[0] {
            case "sessions created":
                sessions_created, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "active sessions":
                active_sessions, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "active streams":
                active_streams, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "packets received":
                preceived, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "packets transmitted":
                ptransmitted, _ = strconv.ParseInt(line_parts[1], 10, 64)
            }
        }
        self.update_active(active_sessions, sessions_created, active_streams, preceived, ptransmitted)
    }
    StartTimeout(self.heartbeat, nil, randomize(self.opts.hrtb_ival, 0.1), 1, self.logger)
}

func (self *Rtp_proxy_client_base) GoOnline() {
    if self.shut_down {
        return
    }
    if ! self.online {
        if ! self.caps_done {
            NewRtpp_caps_checker(self)
            return
        }
        self.online = true
        self.heartbeat()
    }
}

func (self *Rtp_proxy_client_base) GoOffline() {
    if self.shut_down {
        return
    }
    //print "go_offline", self.address, self.online
    if self.online {
        self.online = false
        StartTimeout(self.version_check, nil, randomize(self.opts.hrtb_retr_ival, 0.1), 1, self.logger)
    }
}

func (self *Rtp_proxy_client_base) update_active(active_sessions, sessions_created, active_streams, preceived, ptransmitted int64) {
    self.sessions_created = sessions_created
    self.active_sessions = active_sessions
    self.active_streams = active_streams
    self.preceived = preceived
    self.ptransmitted = ptransmitted
}

func (self *Rtp_proxy_client_base) Shutdown() {
    if self.shut_down { // do not crash when shutdown() called twice
        return
    }
    self.shut_down = true
    self.transport.shutdown()
    self.transport = nil
}

func (self *Rtp_proxy_client_base) GetOpts() sippy_types.RtpProxyClientOpts {
    return self.opts
}
/*
    def get_rtpc_delay(self):
        self.transport.get_rtpc_delay(self)
*/

type Rtpp_caps_checker struct {
    caps_requested  int
    caps_received   int
    rtpc            *Rtp_proxy_client_base
}

func NewRtpp_caps_checker(rtpc *Rtp_proxy_client_base) *Rtpp_caps_checker {
    self := &Rtpp_caps_checker{
        rtpc    : rtpc,
    }
    rtpc.caps_done = false
    self.caps_requested = len(rtpc._CAPSTABLE)
    for _, it := range rtpc._CAPSTABLE {
        attr := it.attr // For some reason the it.attr cannot be passed into the following
                        // function directly - the resulting value is always that of the
                        // last 'it.attr' value.
        rtpc.transport.send_command("VF " + it.vers, func(res string) { self.caps_query_done(res, attr) })
    }
    return self
}

func (self *Rtpp_caps_checker) caps_query_done(result string, attr *bool) {
    self.caps_received += 1
    if result == "1" {
        *attr = true
    } else {
        *attr = false
    }
    if self.caps_received == self.caps_requested {
        self.rtpc.caps_done = true
        self.rtpc.GoOnline()
        self.rtpc = nil
    }
}

func randomize(x time.Duration, p float64) time.Duration {
    return time.Duration(x.Seconds() * (1.0 + p * (1.0 - 2.0 * rand.Float64())) * float64(time.Second))
}
<commit_msg>Fix the inheritance issues.<commit_after>// Copyright (c) 2003-2005 Maxim Sobolev. All rights reserved.
// Copyright (c) 2006-2014 Sippy Software, Inc. All rights reserved.
// Copyright (c) 2016 Andriy Pylypenko. All rights reserved.
//
// All rights reserved.
//
// Redistribution and use in source and binary forms, with or without modification,
// are permitted provided that the following conditions are met:
//
// 1. Redistributions of source code must retain the above copyright notice, this
// list of conditions and the following disclaimer.
//
// 2. Redistributions in binary form must reproduce the above copyright notice,
// this list of conditions and the following disclaimer in the documentation and/or
// other materials provided with the distribution.
//
// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
// ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
// WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR
// ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
// (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
// LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
// ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
// (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
// SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
package sippy

import (
    "bufio"
    "math/rand"
    "net"
    "time"
    "strconv"
    "strings"

    "sippy/conf"
    "sippy/log"
    "sippy/types"
)

func NewRtpProxyClient(opts *rtpProxyClientOpts, config sippy_conf.Config, logger sippy_log.ErrorLogger) (sippy_types.RtpProxyClient, error) {
    rtpp := NewRtp_proxy_client_base(nil, config, opts, logger)
    err := rtpp.Init()
    return rtpp, err
}

type Rtp_proxy_client_base struct {
    heir            sippy_types.RtpProxyClient
    opts            *rtpProxyClientOpts
    transport       rtp_proxy_transport
    proxy_address   string
    online          bool
    sbind_supported bool
    tnot_supported  bool
    copy_supported  bool
    stat_supported  bool
    wdnt_supported  bool
    caps_done       bool
    shut_down       bool
    active_sessions int64
    sessions_created int64
    active_streams  int64
    preceived       int64
    ptransmitted    int64
    logger          sippy_log.ErrorLogger
    global_config   sippy_conf.Config
}

type rtp_proxy_transport interface {
    is_local() bool
    send_command(string, func(string))
    shutdown()
}

func (self *Rtp_proxy_client_base) IsLocal() bool {
    return self.transport.is_local()
}

func (self *Rtp_proxy_client_base) IsOnline() bool {
    return self.online
}

func (self *Rtp_proxy_client_base) SBindSupported() bool {
    return self.sbind_supported
}

func (self *Rtp_proxy_client_base) TNotSupported() bool {
    return self.tnot_supported
}

func (self *Rtp_proxy_client_base) GetProxyAddress() string {
    return self.proxy_address
}

func (self *Rtp_proxy_client_base) me() sippy_types.RtpProxyClient {
    if self.heir != nil {
        return self.heir
    }
    return self
}

func NewRtp_proxy_client_base(heir sippy_types.RtpProxyClient, global_config sippy_conf.Config, opts *rtpProxyClientOpts, logger sippy_log.ErrorLogger) (*Rtp_proxy_client_base) {
    if opts == nil {
        opts = NewRtpProxyClientOpts() // default settings
    }
    self := &Rtp_proxy_client_base{
        heir            : heir,
        caps_done       : false,
        shut_down       : false,
        logger          : logger,
        opts            : opts,
        global_config   : global_config,
    }
    return self
}

func (self *Rtp_proxy_client_base) Init() error {
    var err error
    var rtpp_class func(sippy_types.RtpProxyClient, sippy_conf.Config, net.Addr) (rtp_proxy_transport, error)
    var rtppa net.Addr

    a := self.opts.spath
    if strings.HasPrefix(a, "udp:") {
        tmp := strings.SplitN(a, ":", 3)
        if len(tmp) == 2 {
            rtppa, err = net.ResolveUDPAddr("udp", tmp[1] + ":22222")
        } else {
            rtppa, err = net.ResolveUDPAddr("udp", tmp[1] + ":" + tmp[2])
        }
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_udp
    } else if strings.HasPrefix(a, "udp6:") {
        tmp := strings.SplitN(a, ":", 2)
        a := tmp[1]
        rtp_proxy_host, rtp_proxy_port := a, "22222"
        if a[len(a)-1] != ']' {
            idx := strings.LastIndexByte(a, ':')
            if idx < 0 {
                rtp_proxy_host = a
            } else {
                rtp_proxy_host, rtp_proxy_port = a[:idx], a[idx+1:]
            }
        }
        if rtp_proxy_host[0] != '[' {
            rtp_proxy_host = "[" + rtp_proxy_host + "]"
        }
        rtppa, err = net.ResolveUDPAddr("udp", rtp_proxy_host + ":" + rtp_proxy_port)
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_udp
    } else if strings.HasPrefix(a, "tcp:") {
        tmp := strings.SplitN(a, ":", 3)
        if len(tmp) == 2 {
            rtppa, err = net.ResolveTCPAddr("tcp", tmp[1] + ":22222")
        } else {
            rtppa, err = net.ResolveTCPAddr("tcp", tmp[1] + ":" + tmp[2])
        }
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_stream
    } else if strings.HasPrefix(a, "tcp6:") {
        tmp := strings.SplitN(a, ":", 2)
        a := tmp[1]
        rtp_proxy_host, rtp_proxy_port := a, "22222"
        if a[len(a)-1] != ']' {
            idx := strings.LastIndexByte(a, ':')
            if idx < 0 {
                rtp_proxy_host = a
            } else {
                rtp_proxy_host, rtp_proxy_port = a[:idx], a[idx+1:]
            }
        }
        if rtp_proxy_host[0] != '[' {
            rtp_proxy_host = "[" + rtp_proxy_host + "]"
        }
        rtppa, err = net.ResolveTCPAddr("tcp", rtp_proxy_host + ":" + rtp_proxy_port)
        if err != nil { return err }
        self.proxy_address, _, err = net.SplitHostPort(rtppa.String())
        if err != nil { return err }
        rtpp_class = newRtp_proxy_client_stream
    } else {
        if strings.HasPrefix(a, "unix:") {
            rtppa, err = net.ResolveUnixAddr("unix", a[5:])
        } else if strings.HasPrefix(a, "cunix:") {
            rtppa, err = net.ResolveUnixAddr("unix", a[6:])
        } else {
            rtppa, err = net.ResolveUnixAddr("unix", a)
        }
        self.proxy_address = self.global_config.SipAddress().String()
        rtpp_class = newRtp_proxy_client_stream
    }
    self.transport, err = rtpp_class(self.me(), self.global_config, rtppa)
    if err != nil {
        return err
    }
    self.start()
    return nil
}

func (self *Rtp_proxy_client_base) InitWithAddress(address net.Addr) error {
    var err error
    if strings.HasPrefix(address.Network(), "udp") {
        self.transport, err = newRtp_proxy_client_udp(self.me(), self.global_config, address)
        if err != nil {
            return err
        }
        self.proxy_address, _, err = net.SplitHostPort(address.String())
        if err != nil {
            return err
        }
    } else {
        self.transport, err = newRtp_proxy_client_stream(self.me(), self.global_config, address)
        if err != nil {
            return err
        }
        self.proxy_address = self.global_config.SipAddress().String()
    }
    self.start()
    return nil
}

func (self *Rtp_proxy_client_base) start() {
    if ! self.opts.no_version_check {
        self.version_check()
    } else {
        self.caps_done = true
        self.online = true
    }
}

func (self *Rtp_proxy_client_base) SendCommand(cmd string, cb func(string)) {
    self.transport.send_command(cmd, cb)
}
/*
    def reconnect(self, *args, **kwargs):
        self.rtpp_class.reconnect(self, *args, **kwargs)
*/

func (self *Rtp_proxy_client_base) version_check() {
    if self.shut_down {
        return
    }
    self.transport.send_command("V", self.version_check_reply)
}

func (self *Rtp_proxy_client_base) version_check_reply(version string) {
    if self.shut_down {
        return
    }
    if version == "20040107" {
        self.me().GoOnline()
    } else if self.online {
        self.me().GoOffline()
    } else {
        StartTimeout(self.version_check, nil, randomize(self.opts.hrtb_retr_ival, 0.1), 1, self.logger)
    }
}

func (self *Rtp_proxy_client_base) heartbeat() {
    //print "heartbeat", self, self.address
    if self.shut_down {
        return
    }
    self.transport.send_command("Ib", self.heartbeat_reply)
}

func (self *Rtp_proxy_client_base) heartbeat_reply(stats string) {
    //print "heartbeat_reply", self.address, stats, self.online
    if self.shut_down || ! self.online {
        return
    }
    if stats == "" {
        self.active_sessions = 0
        self.me().GoOffline()
    } else {
        sessions_created := int64(0)
        active_sessions := int64(0)
        active_streams := int64(0)
        preceived := int64(0)
        ptransmitted := int64(0)
        scanner := bufio.NewScanner(strings.NewReader(stats))
        for scanner.Scan() {
            line_parts := strings.SplitN(scanner.Text(), ":", 2)
            if len(line_parts) != 2 { continue }
            switch line_parts[0] {
            case "sessions created":
                sessions_created, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "active sessions":
                active_sessions, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "active streams":
                active_streams, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "packets received":
                preceived, _ = strconv.ParseInt(line_parts[1], 10, 64)
            case "packets transmitted":
                ptransmitted, _ = strconv.ParseInt(line_parts[1], 10, 64)
            }
        }
        self.update_active(active_sessions, sessions_created, active_streams, preceived, ptransmitted)
    }
    StartTimeout(self.heartbeat, nil, randomize(self.opts.hrtb_ival, 0.1), 1, self.logger)
}

func (self *Rtp_proxy_client_base) GoOnline() {
    if self.shut_down {
        return
    }
    if ! self.online {
        if ! self.caps_done {
            newRtppCapsChecker(self)
            return
        }
        self.online = true
        self.heartbeat()
    }
}

func (self *Rtp_proxy_client_base) GoOffline() {
    if self.shut_down {
        return
    }
    //print "go_offline", self.address, self.online
    if self.online {
        self.online = false
        StartTimeout(self.version_check, nil, randomize(self.opts.hrtb_retr_ival, 0.1), 1, self.logger)
    }
}

func (self *Rtp_proxy_client_base) update_active(active_sessions, sessions_created, active_streams, preceived, ptransmitted int64) {
    self.sessions_created = sessions_created
    self.active_sessions = active_sessions
    self.active_streams = active_streams
    self.preceived = preceived
    self.ptransmitted = ptransmitted
}

func (self *Rtp_proxy_client_base) Shutdown() {
    if self.shut_down { // do not crash when shutdown() called twice
        return
    }
    self.shut_down = true
    self.transport.shutdown()
    self.transport = nil
}

func (self *Rtp_proxy_client_base) GetOpts() sippy_types.RtpProxyClientOpts {
    return self.opts
}
/*
    def get_rtpc_delay(self):
        self.transport.get_rtpc_delay(self)
*/

type rtppCapsChecker struct {
    caps_requested  int
    caps_received   int
    rtpc            *Rtp_proxy_client_base
}

func newRtppCapsChecker(rtpc *Rtp_proxy_client_base) *rtppCapsChecker {
    self := &rtppCapsChecker{
        rtpc    : rtpc,
    }
    rtpc.caps_done = false
    CAPSTABLE := []struct{ vers string; attr *bool }{
        { "20071218", &self.rtpc.copy_supported },
        { "20080403", &self.rtpc.stat_supported },
        { "20081224", &self.rtpc.tnot_supported },
        { "20090810", &self.rtpc.sbind_supported },
        { "20150617", &self.rtpc.wdnt_supported },
    }
    self.caps_requested = len(CAPSTABLE)
    for _, it := range CAPSTABLE {
        attr := it.attr // For some reason the it.attr cannot be passed into the following
                        // function directly - the resulting value is always that of the
                        // last 'it.attr' value.
        rtpc.transport.send_command("VF " + it.vers, func(res string) { self.caps_query_done(res, attr) })
    }
    return self
}

func (self *rtppCapsChecker) caps_query_done(result string, attr *bool) {
    self.caps_received += 1
    if result == "1" {
        *attr = true
    } else {
        *attr = false
    }
    if self.caps_received == self.caps_requested {
        self.rtpc.caps_done = true
        self.rtpc.GoOnline()
        self.rtpc = nil
    }
}

func randomize(x time.Duration, p float64) time.Duration {
    return time.Duration(x.Seconds() * (1.0 + p * (1.0 - 2.0 * rand.Float64())) * float64(time.Second))
}
<|endoftext|>


<commit_before>package consulacls

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

	"github.com/hashicorp/nomad/e2e/framework/provisioning"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
)

// DefaultTFStateFile is the location of the TF state file, as created for the
// e2e test framework. This file is used to extract the TF serial number, which
// is used to determine whether the consul bootstrap process is necessary or has
// already taken place.
const DefaultTFStateFile = "terraform/terraform.tfstate"

// A Manager is used to manipulate whether Consul ACLs are enabled or disabled.
// Only works with TF provisioned clusters.
type Manager interface {
	// Enable Consul ACLs in the Consul cluster. The Consul ACL master token
	// associated with the Consul cluster is returned.
	//
	// A complete bootstrap process will take place if necessary.
	//
	// Once enabled, Consul ACLs can be disabled with Disable.
	Enable(t *testing.T) string

	// Disable Consul ACLs in the Consul Cluster.
	//
	// Once disabled, Consul ACLs can be re-enabled with Enable.
	Disable(t *testing.T)
}

type tfManager struct {
	serial int
	token  string
}

func New(tfStateFile string) (*tfManager, error) {
	serial, err := extractSerial(tfStateFile)
	if err != nil {
		return nil, err
	}
	return &tfManager{
		serial: serial,
	}, nil
}

func (m *tfManager) Enable(t *testing.T) string {
	// Create the local script runner that will be used to run the ACL management
	// script, this time with the "enable" sub-command.
	var runner provisioning.LinuxRunner
	err := runner.Open(t)
	require.NoError(t, err)

	// Run the consul ACL bootstrap script, which will store the master token
	// in the deterministic path based on the TF state serial number. If the
	// bootstrap process had already taken place, ACLs will be activated but
	// without going through the bootstrap process again, re-using the already
	// existing Consul ACL master token.
	err = runner.Run(strings.Join([]string{
		"consulacls/consul-acls-manage.sh", "enable",
	}, " "))
	require.NoError(t, err)

	// Read the Consul ACL master token that was generated (or if the token
	// already existed because the bootstrap process had already taken place,
	// that one).
	token, err := m.readToken()
	require.NoError(t, err)
	return token
}

type tfState struct {
	Serial int `json:"serial"`
}

// extractSerial will parse the TF state file looking for the serial number.
func extractSerial(filename string) (int, error) {
	if filename == "" {
		filename = DefaultTFStateFile
	}
	b, err := ioutil.ReadFile(filename)
	if err != nil {
		return 0, errors.Wrap(err, "failed to extract TF serial")
	}
	var state tfState
	if err := json.Unmarshal(b, &state); err != nil {
		return 0, errors.Wrap(err, "failed to extract TF serial")
	}
	return state.Serial, nil
}

// tokenPath returns the expected path for the Consul ACL master token generated
// by the consul-acls-manage.sh bootstrap script for the current TF serial version.
func (m *tfManager) tokenPath() string {
	return fmt.Sprintf("/tmp/e2e-consul-bootstrap-%d.token", m.serial)
}

func (m *tfManager) readToken() (string, error) {
	b, err := ioutil.ReadFile(m.tokenPath())
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(b)), nil
}

func (m *tfManager) Disable(t *testing.T) {
	// Create the local script runner that will be used to run the ACL management
	// script, this time with the "disable" sub-command.
	var runner provisioning.LinuxRunner
	err := runner.Open(t)
	require.NoError(t, err)

	// Run the consul ACL bootstrap script, which will modify the Consul Server
	// ACL policies to disable ACLs, and then restart those agents.
	err = runner.Run(strings.Join([]string{
		"consulacls/consul-acls-manage.sh", "disable",
	}, " "))
	require.NoError(t, err)
}
<commit_msg>e2e: remove forgotten unused field from new struct<commit_after>package consulacls

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

	"github.com/hashicorp/nomad/e2e/framework/provisioning"
	"github.com/pkg/errors"
	"github.com/stretchr/testify/require"
)

// DefaultTFStateFile is the location of the TF state file, as created for the
// e2e test framework. This file is used to extract the TF serial number, which
// is used to determine whether the consul bootstrap process is necessary or has
// already taken place.
const DefaultTFStateFile = "terraform/terraform.tfstate"

// A Manager is used to manipulate whether Consul ACLs are enabled or disabled.
// Only works with TF provisioned clusters.
type Manager interface {
	// Enable Consul ACLs in the Consul cluster. The Consul ACL master token
	// associated with the Consul cluster is returned.
	//
	// A complete bootstrap process will take place if necessary.
	//
	// Once enabled, Consul ACLs can be disabled with Disable.
	Enable(t *testing.T) string

	// Disable Consul ACLs in the Consul Cluster.
	//
	// Once disabled, Consul ACLs can be re-enabled with Enable.
	Disable(t *testing.T)
}

type tfManager struct {
	serial int
}

func New(tfStateFile string) (*tfManager, error) {
	serial, err := extractSerial(tfStateFile)
	if err != nil {
		return nil, err
	}
	return &tfManager{
		serial: serial,
	}, nil
}

func (m *tfManager) Enable(t *testing.T) string {
	// Create the local script runner that will be used to run the ACL management
	// script, this time with the "enable" sub-command.
	var runner provisioning.LinuxRunner
	err := runner.Open(t)
	require.NoError(t, err)

	// Run the consul ACL bootstrap script, which will store the master token
	// in the deterministic path based on the TF state serial number. If the
	// bootstrap process had already taken place, ACLs will be activated but
	// without going through the bootstrap process again, re-using the already
	// existing Consul ACL master token.
	err = runner.Run(strings.Join([]string{
		"consulacls/consul-acls-manage.sh", "enable",
	}, " "))
	require.NoError(t, err)

	// Read the Consul ACL master token that was generated (or if the token
	// already existed because the bootstrap process had already taken place,
	// that one).
	token, err := m.readToken()
	require.NoError(t, err)
	return token
}

type tfState struct {
	Serial int `json:"serial"`
}

// extractSerial will parse the TF state file looking for the serial number.
func extractSerial(filename string) (int, error) {
	if filename == "" {
		filename = DefaultTFStateFile
	}
	b, err := ioutil.ReadFile(filename)
	if err != nil {
		return 0, errors.Wrap(err, "failed to extract TF serial")
	}
	var state tfState
	if err := json.Unmarshal(b, &state); err != nil {
		return 0, errors.Wrap(err, "failed to extract TF serial")
	}
	return state.Serial, nil
}

// tokenPath returns the expected path for the Consul ACL master token generated
// by the consul-acls-manage.sh bootstrap script for the current TF serial version.
func (m *tfManager) tokenPath() string {
	return fmt.Sprintf("/tmp/e2e-consul-bootstrap-%d.token", m.serial)
}

func (m *tfManager) readToken() (string, error) {
	b, err := ioutil.ReadFile(m.tokenPath())
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(b)), nil
}

func (m *tfManager) Disable(t *testing.T) {
	// Create the local script runner that will be used to run the ACL management
	// script, this time with the "disable" sub-command.
	var runner provisioning.LinuxRunner
	err := runner.Open(t)
	require.NoError(t, err)

	// Run the consul ACL bootstrap script, which will modify the Consul Server
	// ACL policies to disable ACLs, and then restart those agents.
	err = runner.Run(strings.Join([]string{
		"consulacls/consul-acls-manage.sh", "disable",
	}, " "))
	require.NoError(t, err)
}
<|endoftext|>


<commit_before>package libproxy

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"
)

const (
	maxBufferSize = 65536
)

type windowState struct {
	current uint64
	allowed uint64
}

func (w *windowState) String() string {
	return fmt.Sprintf("current %d, allowed %d", w.current, w.allowed)
}

func newWindowState() *windowState {
	return &windowState{}
}
func (w *windowState) size() int {
	return int(w.allowed - w.current)
}
func (w *windowState) isAlmostClosed() bool {
	return w.size() < maxBufferSize/2
}
func (w *windowState) advance() {
	w.allowed = w.current + uint64(maxBufferSize)
}

type channel struct {
	m             *sync.Mutex
	c             *sync.Cond
	multiplexer   *Multiplexer
	destination   Destination
	ID            uint32
	read          *windowState
	write         *windowState
	readPipe      *bufferedPipe
	closeReceived bool
	closeSent     bool
	shutdownSent  bool
	writeDeadline time.Time
}

// newChannel registers a channel through the multiplexer
func newChannel(multiplexer *Multiplexer, ID uint32, d Destination) *channel {
	var m sync.Mutex
	c := sync.NewCond(&m)
	readPipe := newBufferedPipe()
	return &channel{
		m:           &m,
		c:           c,
		multiplexer: multiplexer,
		destination: d,
		ID:          ID,
		read:        &windowState{},
		write:       &windowState{},
		readPipe:    readPipe,
	}
}

func (c *channel) sendWindowUpdate() error {
	c.m.Lock()
	c.read.advance()
	seq := c.read.allowed
	c.m.Unlock()
	return c.multiplexer.send(NewWindow(c.ID, seq))
}

func (c *channel) recvWindowUpdate(seq uint64) {
	c.m.Lock()
	c.write.allowed = seq
	c.c.Signal()
	c.m.Unlock()
}

func (c *channel) Read(p []byte) (int, error) {
	n, err := c.readPipe.Read(p)
	c.m.Lock()
	c.read.current = c.read.current + uint64(n)
	needUpdate := c.read.isAlmostClosed()
	c.m.Unlock()
	if needUpdate {
		c.sendWindowUpdate()
	}
	return n, err
}

func (c *channel) Write(p []byte) (int, error) {
	c.m.Lock()
	defer c.m.Unlock()
	written := 0
	for {
		if len(p) == 0 {
			return written, nil
		}
		if c.closeReceived || c.closeSent || c.shutdownSent {
			return written, io.EOF
		}
		if c.write.size() > 0 {
			toWrite := c.write.size()
			if toWrite > len(p) {
				toWrite = len(p)
			}
			// need to write the header and the payload together
			c.multiplexer.writeMutex.Lock()
			f := NewData(c.ID, uint32(toWrite))
			err1 := f.Write(c.multiplexer.connW)
			_, err2 := c.multiplexer.connW.Write(p[0:toWrite])
			err3 := c.multiplexer.connW.Flush()
			c.multiplexer.writeMutex.Unlock()

			if err1 != nil {
				return written, err1
			}
			if err2 != nil {
				return written, err2
			}
			if err3 != nil {
				return written, err3
			}
			c.write.current = c.write.current + uint64(toWrite)
			p = p[toWrite:]
			written = written + toWrite
			continue
		}

		// Wait for the write window to be increased (or a timeout)
		done := make(chan struct{})
		timeout := make(chan time.Time)
		if !c.writeDeadline.IsZero() {
			go func() {
				time.Sleep(time.Until(c.writeDeadline))
				close(timeout)
			}()
		}
		go func() {
			c.c.Wait()
			close(done)
		}()
		select {
		case <-timeout:
			// clean up the goroutine
			c.c.Broadcast()
			<-done
			return written, &errTimeout{}
		case <-done:
			// The timeout will still fire in the background
			continue
		}
	}
}

func (c *channel) Close() error {
	// Avoid a Write() racing with us and sending after we Close()
	c.m.Lock()
	c.closeSent = true
	c.m.Unlock()

	if err := c.multiplexer.send(NewClose(c.ID)); err != nil {
		return err
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.c.Broadcast()
	if c.closeSent && c.closeReceived {
		c.multiplexer.freeChannel(c.ID)
	}
	return nil
}

func (c *channel) CloseRead() error {
	return c.readPipe.CloseWrite()
}

func (c *channel) CloseWrite() error {
	// Avoid a Write() racing with us and sending after we Close()
	c.m.Lock()
	c.shutdownSent = true
	c.m.Unlock()

	if err := c.multiplexer.send(NewShutdown(c.ID)); err != nil {
		return err
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.c.Broadcast()
	return nil
}

func (c *channel) recvClose() error {
	c.m.Lock()
	defer c.m.Unlock()
	c.closeReceived = true
	c.c.Broadcast()
	return nil
}

func (c *channel) isClosed() bool {
	c.m.Lock()
	defer c.m.Unlock()
	return c.closeReceived && c.closeSent
}

func (c *channel) SetReadDeadline(timeout time.Time) error {
	return c.readPipe.SetReadDeadline(timeout)
}

func (c *channel) SetWriteDeadline(timeout time.Time) error {
	c.m.Lock()
	defer c.m.Unlock()
	c.writeDeadline = timeout
	c.c.Broadcast()
	return nil
}

func (c *channel) SetDeadline(timeout time.Time) error {
	if err := c.SetReadDeadline(timeout); err != nil {
		return err
	}
	return c.SetWriteDeadline(timeout)
}

func (c *channel) RemoteAddr() net.Addr {
	return &channelAddr{
		d: c.destination,
	}
}

func (c *channel) LocalAddr() net.Addr {
	return c.RemoteAddr() // There is no local address
}

type channelAddr struct {
	d Destination
}

func (a *channelAddr) Network() string {
	return "channel"
}

func (a *channelAddr) String() string {
	return a.d.String()
}

// Multiplexer muxes and demuxes sub-connections over a single connection
type Multiplexer struct {
	label         string
	conn          io.Closer
	connR         io.Reader // with buffering
	connW         *bufio.Writer
	writeMutex    *sync.Mutex // hold when writing on the channel
	channels      map[uint32]*channel
	nextChannelID uint32
	metadataMutex *sync.Mutex // hold when reading/modifying this structure
	pendingAccept []*channel  // incoming connections
	acceptCond    *sync.Cond
}

// NewMultiplexer constructs a multiplexer from a channel
func NewMultiplexer(label string, conn io.ReadWriteCloser) *Multiplexer {
	var writeMutex, metadataMutex sync.Mutex
	acceptCond := sync.NewCond(&metadataMutex)
	channels := make(map[uint32]*channel)
	connR := bufio.NewReader(conn)
	connW := bufio.NewWriter(conn)
	return &Multiplexer{
		label:         label,
		conn:          conn,
		connR:         connR,
		connW:         connW,
		writeMutex:    &writeMutex,
		channels:      channels,
		metadataMutex: &metadataMutex,
		acceptCond:    acceptCond,
	}
}

func (m *Multiplexer) send(f *Frame) error {
	m.writeMutex.Lock()
	defer m.writeMutex.Unlock()
	if err := f.Write(m.connW); err != nil {
		return err
	}
	return m.connW.Flush()
}

func (m *Multiplexer) findFreeChannelID() uint32 {
	// the metadataMutex is already held
	id := m.nextChannelID
	for {
		if _, ok := m.channels[id]; !ok {
			m.nextChannelID = id + 1
			return id
		}
		id++
	}
}

func (m *Multiplexer) freeChannel(ID uint32) {
	m.metadataMutex.Lock()
	defer m.metadataMutex.Unlock()
	delete(m.channels, ID)
}

// Dial opens a connection to the given destination
func (m *Multiplexer) Dial(d Destination) (Conn, error) {
	m.metadataMutex.Lock()
	id := m.findFreeChannelID()
	channel := newChannel(m, id, d)
	m.channels[id] = channel
	m.metadataMutex.Unlock()

	if err := m.send(NewOpen(id, d)); err != nil {
		return nil, err
	}
	if err := channel.sendWindowUpdate(); err != nil {
		return nil, err
	}
	return channel, nil
}

// Accept returns the next client connection
func (m *Multiplexer) Accept() (Conn, *Destination, error) {
	m.metadataMutex.Lock()
	defer m.metadataMutex.Unlock()
	for {
		if len(m.pendingAccept) > 0 {
			first := m.pendingAccept[0]
			m.pendingAccept = m.pendingAccept[1:]
			if err := first.sendWindowUpdate(); err != nil {
				return nil, nil, err
			}
			return first, &first.destination, nil
		}
		m.acceptCond.Wait()
	}
}

// Run starts handling the requests from the other side
func (m *Multiplexer) Run() {
	go func() {
		if err := m.run(); err != nil {
			log.Printf("Multiplexer main loop failed with %v", err)
		}
	}()
}

func (m *Multiplexer) run() error {
	for {
		f, err := unmarshalFrame(m.connR)
		if err != nil {
			return fmt.Errorf("Failed to unmarshal command frame: %v", err)
		}
		switch f.Command {
		case Open:
			o, err := f.Open()
			if err != nil {
				return fmt.Errorf("Failed to unmarshal open command: %v", err)
			}
			switch o.Connection {
			case Dedicated:
				return fmt.Errorf("Dedicated connections are not implemented yet")
			case Multiplexed:
				m.metadataMutex.Lock()
				channel := newChannel(m, f.ID, o.Destination)
				m.channels[f.ID] = channel
				m.pendingAccept = append(m.pendingAccept, channel)
				m.acceptCond.Signal()
				m.metadataMutex.Unlock()
			}
		case Window:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			w, err := f.Window()
			if err != nil {
				return err
			}
			channel.recvWindowUpdate(w.seq)
		case Data:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			d, err := f.Data()
			if err != nil {
				return err
			}
			if _, err := io.CopyN(channel.readPipe, m.connR, int64(d.payloadlen)); err != nil {
				return err
			}
		case Shutdown:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			if err := channel.readPipe.CloseWrite(); err != nil {
				return err
			}
		case Close:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			// this will unblock waiting Read calls
			if err := channel.readPipe.CloseWrite(); err != nil {
				return err
			}
			// this will unblock waiting Write calls
			if err := channel.recvClose(); err != nil {
				return err
			}
			if channel.isClosed() {
				m.freeChannel(channel.ID)
			}
		default:
			return fmt.Errorf("Unknown command type: %v", f)
		}
	}
}

// Forward runs the TCP/UDP forwarder over a sub-connection
func Forward(conn Conn, destination Destination, quit chan struct{}) {
	defer conn.Close()

	switch destination.Proto {
	case TCP:
		backendAddr := net.TCPAddr{IP: destination.IP, Port: int(destination.Port), Zone: ""}
		if err := HandleTCPConnection(conn, &backendAddr, quit); err != nil {
			log.Printf("Error setting up TCP proxy subconnection: %v", err)
			return
		}
	case Unix:
		backendAddr, err := net.ResolveUnixAddr("unix", destination.Path)
		if err != nil {
			log.Printf("Error resolving Unix address %s", destination.Path)
			return
		}
		if err := HandleUnixConnection(conn, backendAddr, quit); err != nil {
			log.Printf("Error setting up Unix proxy subconnection: %v", err)
			return
		}
	case UDP:
		backendAddr := &net.UDPAddr{IP: destination.IP, Port: int(destination.Port), Zone: ""}

		proxy, err := NewUDPProxy(backendAddr, NewUDPConn(conn), backendAddr)
		if err != nil {
			log.Printf("Failed to setup UDP proxy for %s: %#v", backendAddr, err)
			return
		}
		proxy.Run()
		return
	default:
		log.Printf("Unknown protocol: %d", destination.Proto)
		return
	}
}
<commit_msg>pkg/libproxy: add Multiplexer.IsRunning<commit_after>package libproxy

import (
	"bufio"
	"fmt"
	"io"
	"log"
	"net"
	"sync"
	"time"
)

const (
	maxBufferSize = 65536
)

type windowState struct {
	current uint64
	allowed uint64
}

func (w *windowState) String() string {
	return fmt.Sprintf("current %d, allowed %d", w.current, w.allowed)
}

func newWindowState() *windowState {
	return &windowState{}
}
func (w *windowState) size() int {
	return int(w.allowed - w.current)
}
func (w *windowState) isAlmostClosed() bool {
	return w.size() < maxBufferSize/2
}
func (w *windowState) advance() {
	w.allowed = w.current + uint64(maxBufferSize)
}

type channel struct {
	m             *sync.Mutex
	c             *sync.Cond
	multiplexer   *Multiplexer
	destination   Destination
	ID            uint32
	read          *windowState
	write         *windowState
	readPipe      *bufferedPipe
	closeReceived bool
	closeSent     bool
	shutdownSent  bool
	writeDeadline time.Time
}

// newChannel registers a channel through the multiplexer
func newChannel(multiplexer *Multiplexer, ID uint32, d Destination) *channel {
	var m sync.Mutex
	c := sync.NewCond(&m)
	readPipe := newBufferedPipe()
	return &channel{
		m:           &m,
		c:           c,
		multiplexer: multiplexer,
		destination: d,
		ID:          ID,
		read:        &windowState{},
		write:       &windowState{},
		readPipe:    readPipe,
	}
}

func (c *channel) sendWindowUpdate() error {
	c.m.Lock()
	c.read.advance()
	seq := c.read.allowed
	c.m.Unlock()
	return c.multiplexer.send(NewWindow(c.ID, seq))
}

func (c *channel) recvWindowUpdate(seq uint64) {
	c.m.Lock()
	c.write.allowed = seq
	c.c.Signal()
	c.m.Unlock()
}

func (c *channel) Read(p []byte) (int, error) {
	n, err := c.readPipe.Read(p)
	c.m.Lock()
	c.read.current = c.read.current + uint64(n)
	needUpdate := c.read.isAlmostClosed()
	c.m.Unlock()
	if needUpdate {
		c.sendWindowUpdate()
	}
	return n, err
}

func (c *channel) Write(p []byte) (int, error) {
	c.m.Lock()
	defer c.m.Unlock()
	written := 0
	for {
		if len(p) == 0 {
			return written, nil
		}
		if c.closeReceived || c.closeSent || c.shutdownSent {
			return written, io.EOF
		}
		if c.write.size() > 0 {
			toWrite := c.write.size()
			if toWrite > len(p) {
				toWrite = len(p)
			}
			// need to write the header and the payload together
			c.multiplexer.writeMutex.Lock()
			f := NewData(c.ID, uint32(toWrite))
			err1 := f.Write(c.multiplexer.connW)
			_, err2 := c.multiplexer.connW.Write(p[0:toWrite])
			err3 := c.multiplexer.connW.Flush()
			c.multiplexer.writeMutex.Unlock()

			if err1 != nil {
				return written, err1
			}
			if err2 != nil {
				return written, err2
			}
			if err3 != nil {
				return written, err3
			}
			c.write.current = c.write.current + uint64(toWrite)
			p = p[toWrite:]
			written = written + toWrite
			continue
		}

		// Wait for the write window to be increased (or a timeout)
		done := make(chan struct{})
		timeout := make(chan time.Time)
		if !c.writeDeadline.IsZero() {
			go func() {
				time.Sleep(time.Until(c.writeDeadline))
				close(timeout)
			}()
		}
		go func() {
			c.c.Wait()
			close(done)
		}()
		select {
		case <-timeout:
			// clean up the goroutine
			c.c.Broadcast()
			<-done
			return written, &errTimeout{}
		case <-done:
			// The timeout will still fire in the background
			continue
		}
	}
}

func (c *channel) Close() error {
	// Avoid a Write() racing with us and sending after we Close()
	c.m.Lock()
	c.closeSent = true
	c.m.Unlock()

	if err := c.multiplexer.send(NewClose(c.ID)); err != nil {
		return err
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.c.Broadcast()
	if c.closeSent && c.closeReceived {
		c.multiplexer.freeChannel(c.ID)
	}
	return nil
}

func (c *channel) CloseRead() error {
	return c.readPipe.CloseWrite()
}

func (c *channel) CloseWrite() error {
	// Avoid a Write() racing with us and sending after we Close()
	c.m.Lock()
	c.shutdownSent = true
	c.m.Unlock()

	if err := c.multiplexer.send(NewShutdown(c.ID)); err != nil {
		return err
	}
	c.m.Lock()
	defer c.m.Unlock()
	c.c.Broadcast()
	return nil
}

func (c *channel) recvClose() error {
	c.m.Lock()
	defer c.m.Unlock()
	c.closeReceived = true
	c.c.Broadcast()
	return nil
}

func (c *channel) isClosed() bool {
	c.m.Lock()
	defer c.m.Unlock()
	return c.closeReceived && c.closeSent
}

func (c *channel) SetReadDeadline(timeout time.Time) error {
	return c.readPipe.SetReadDeadline(timeout)
}

func (c *channel) SetWriteDeadline(timeout time.Time) error {
	c.m.Lock()
	defer c.m.Unlock()
	c.writeDeadline = timeout
	c.c.Broadcast()
	return nil
}

func (c *channel) SetDeadline(timeout time.Time) error {
	if err := c.SetReadDeadline(timeout); err != nil {
		return err
	}
	return c.SetWriteDeadline(timeout)
}

func (c *channel) RemoteAddr() net.Addr {
	return &channelAddr{
		d: c.destination,
	}
}

func (c *channel) LocalAddr() net.Addr {
	return c.RemoteAddr() // There is no local address
}

type channelAddr struct {
	d Destination
}

func (a *channelAddr) Network() string {
	return "channel"
}

func (a *channelAddr) String() string {
	return a.d.String()
}

// Multiplexer muxes and demuxes sub-connections over a single connection
type Multiplexer struct {
	label         string
	conn          io.Closer
	connR         io.Reader // with buffering
	connW         *bufio.Writer
	writeMutex    *sync.Mutex // hold when writing on the channel
	channels      map[uint32]*channel
	nextChannelID uint32
	metadataMutex *sync.Mutex // hold when reading/modifying this structure
	pendingAccept []*channel  // incoming connections
	acceptCond    *sync.Cond
	isRunning     bool
}

// NewMultiplexer constructs a multiplexer from a channel
func NewMultiplexer(label string, conn io.ReadWriteCloser) *Multiplexer {
	var writeMutex, metadataMutex sync.Mutex
	acceptCond := sync.NewCond(&metadataMutex)
	channels := make(map[uint32]*channel)
	connR := bufio.NewReader(conn)
	connW := bufio.NewWriter(conn)
	return &Multiplexer{
		label:         label,
		conn:          conn,
		connR:         connR,
		connW:         connW,
		writeMutex:    &writeMutex,
		channels:      channels,
		metadataMutex: &metadataMutex,
		acceptCond:    acceptCond,
	}
}

func (m *Multiplexer) send(f *Frame) error {
	m.writeMutex.Lock()
	defer m.writeMutex.Unlock()
	if err := f.Write(m.connW); err != nil {
		return err
	}
	return m.connW.Flush()
}

func (m *Multiplexer) findFreeChannelID() uint32 {
	// the metadataMutex is already held
	id := m.nextChannelID
	for {
		if _, ok := m.channels[id]; !ok {
			m.nextChannelID = id + 1
			return id
		}
		id++
	}
}

func (m *Multiplexer) freeChannel(ID uint32) {
	m.metadataMutex.Lock()
	defer m.metadataMutex.Unlock()
	delete(m.channels, ID)
}

// Dial opens a connection to the given destination
func (m *Multiplexer) Dial(d Destination) (Conn, error) {
	m.metadataMutex.Lock()
	id := m.findFreeChannelID()
	channel := newChannel(m, id, d)
	m.channels[id] = channel
	m.metadataMutex.Unlock()

	if err := m.send(NewOpen(id, d)); err != nil {
		return nil, err
	}
	if err := channel.sendWindowUpdate(); err != nil {
		return nil, err
	}
	return channel, nil
}

// Accept returns the next client connection
func (m *Multiplexer) Accept() (Conn, *Destination, error) {
	m.metadataMutex.Lock()
	defer m.metadataMutex.Unlock()
	for {
		if len(m.pendingAccept) > 0 {
			first := m.pendingAccept[0]
			m.pendingAccept = m.pendingAccept[1:]
			if err := first.sendWindowUpdate(); err != nil {
				return nil, nil, err
			}
			return first, &first.destination, nil
		}
		m.acceptCond.Wait()
	}
}

// Run starts handling the requests from the other side
func (m *Multiplexer) Run() {
	m.metadataMutex.Lock()
	m.isRunning = true
	m.metadataMutex.Unlock()
	go func() {
		if err := m.run(); err != nil {
			log.Printf("Multiplexer main loop failed with %v", err)
		}
		m.metadataMutex.Lock()
		m.isRunning = false
		m.metadataMutex.Unlock()
	}()
}

// IsRunning returns whether the multiplexer is running or not
func (m *Multiplexer) IsRunning() bool {
	m.metadataMutex.Lock()
	defer m.metadataMutex.Unlock()
	return m.isRunning
}

func (m *Multiplexer) run() error {
	for {
		f, err := unmarshalFrame(m.connR)
		if err != nil {
			return fmt.Errorf("Failed to unmarshal command frame: %v", err)
		}
		switch f.Command {
		case Open:
			o, err := f.Open()
			if err != nil {
				return fmt.Errorf("Failed to unmarshal open command: %v", err)
			}
			switch o.Connection {
			case Dedicated:
				return fmt.Errorf("Dedicated connections are not implemented yet")
			case Multiplexed:
				m.metadataMutex.Lock()
				channel := newChannel(m, f.ID, o.Destination)
				m.channels[f.ID] = channel
				m.pendingAccept = append(m.pendingAccept, channel)
				m.acceptCond.Signal()
				m.metadataMutex.Unlock()
			}
		case Window:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			w, err := f.Window()
			if err != nil {
				return err
			}
			channel.recvWindowUpdate(w.seq)
		case Data:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			d, err := f.Data()
			if err != nil {
				return err
			}
			if _, err := io.CopyN(channel.readPipe, m.connR, int64(d.payloadlen)); err != nil {
				return err
			}
		case Shutdown:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			if err := channel.readPipe.CloseWrite(); err != nil {
				return err
			}
		case Close:
			m.metadataMutex.Lock()
			channel, ok := m.channels[f.ID]
			m.metadataMutex.Unlock()
			if !ok {
				return fmt.Errorf("Unknown channel id: %v", f.ID)
			}
			// this will unblock waiting Read calls
			if err := channel.readPipe.CloseWrite(); err != nil {
				return err
			}
			// this will unblock waiting Write calls
			if err := channel.recvClose(); err != nil {
				return err
			}
			if channel.isClosed() {
				m.freeChannel(channel.ID)
			}
		default:
			return fmt.Errorf("Unknown command type: %v", f)
		}
	}
}

// Forward runs the TCP/UDP forwarder over a sub-connection
func Forward(conn Conn, destination Destination, quit chan struct{}) {
	defer conn.Close()

	switch destination.Proto {
	case TCP:
		backendAddr := net.TCPAddr{IP: destination.IP, Port: int(destination.Port), Zone: ""}
		if err := HandleTCPConnection(conn, &backendAddr, quit); err != nil {
			log.Printf("Error setting up TCP proxy subconnection: %v", err)
			return
		}
	case Unix:
		backendAddr, err := net.ResolveUnixAddr("unix", destination.Path)
		if err != nil {
			log.Printf("Error resolving Unix address %s", destination.Path)
			return
		}
		if err := HandleUnixConnection(conn, backendAddr, quit); err != nil {
			log.Printf("Error setting up Unix proxy subconnection: %v", err)
			return
		}
	case UDP:
		backendAddr := &net.UDPAddr{IP: destination.IP, Port: int(destination.Port), Zone: ""}

		proxy, err := NewUDPProxy(backendAddr, NewUDPConn(conn), backendAddr)
		if err != nil {
			log.Printf("Failed to setup UDP proxy for %s: %#v", backendAddr, err)
			return
		}
		proxy.Run()
		return
	default:
		log.Printf("Unknown protocol: %d", destination.Proto)
		return
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/fatih/set"
	"github.com/streadway/amqp"
)

const BROKER_NAME = "broker"

var (
	conf *config.Config
	log  = logger.New(BROKER_NAME)

	STORAGE_BACKEND = "redis"
	routeMap        = make(map[string]*set.Set)
	sessionsMap     = make(map[string]*sockjs.Session)
	globalMapMutex  sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	flagProfile      = flag.String("c", "", "Configuration profile from file")
	flagBrokerDomain = flag.String("a", "", "Send kontrol a custom domain istead of os.Hostname")
	flagKontrolUUID  = flag.String("u", "", "Enable Kontrol mode")
	flagBrokerType   = flag.String("b", "broker", "Define broker type. Available: broker and brokerKite. B")
	flagDebug        = flag.Bool("d", false, "Debug mode")
)

// Broker is a router/multiplexer that routes messages coming from a SockJS
// server to an AMQP exchange and vice versa. Broker basically listens to
// client messages (Koding users) from the SockJS server. The message is
// either passed to the appropriate exchange or a response is sent back to the
// client. Each message has an "action" field that defines how to act for a
// received message.
type Broker struct {
	Config            *config.Broker
	Hostname          string
	ServiceUniqueName string
	AuthAllExchange   string
	PublishConn       *amqp.Connection
	ConsumeConn       *amqp.Connection

	// Accepts SockJS connections
	listener net.Listener

	// Closed when SockJS server is ready to acccept connections
	ready chan struct{}
}

// NewBroker returns a new Broker instance with ServiceUniqueName and Hostname
// prepopulated. After creating a Broker instance, one has to call
// broker.Run() or broker.Start() to start the broker instance and call
// broker.Close() for a graceful stop.
func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(*flagBrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := BROKER_NAME + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
		ready:             make(chan struct{}),
	}
}

func main() {
	flag.Parse()
	if *flagProfile == "" {
		log.Fatal("Please specify profile via -c. Aborting.")
	}

	conf = config.MustConfig(*flagProfile)
	broker := NewBroker()

	switch *flagBrokerType {
	case "brokerKite":
		broker.Config = &conf.BrokerKite
	case "premiumBrokerKite":
		broker.Config = &conf.PremiumBrokerKite
	default:
		broker.Config = &conf.Broker
	}

	// update broker name
	log = logger.New(broker.Config.Name)
	var logLevel logger.Level
	if *flagDebug {
		logLevel = logger.DEBUG
	} else {
		logLevel = logger.GetLoggingLevelFromConfig(BROKER_NAME, *flagProfile)
	}

	log.SetLevel(logLevel)
	broker.Run()
}

// Run starts the broker.
func (b *Broker) Run() {
	// sets the maximum number of CPUs that can be executing simultaneously
	runtime.GOMAXPROCS(runtime.NumCPU())

	lifecycle.Startup(BROKER_NAME, false)
	logger.RunGaugesLoop(log)

	if err := b.registerToKontrol(); err != nil {
		log.Critical("Couldnt register to kontrol, stopping... %v", err)
		return
	}

	if err := b.startAMQP(); err != nil {
		log.Critical("Couldnt create amqp bindings, stopping... %v", err)
		return
	}

	b.startSockJS() // blocking
}

// Start is like Run() but waits until the SockJS listener is ready to be
// used.
func (b *Broker) Start() {
	go b.Run()
	<-b.ready
}

// Close close all amqp connections and closes the SockJS server listener
func (b *Broker) Close() {
	b.PublishConn.Close()
	b.ConsumeConn.Close()
	b.listener.Close()
}

// registerToKontrol registers the broker to KontrolDaemon. This is needed to
// populate a list of brokers and show them to the client. The list is
// available at: https://koding.com/-/services/broker?all
func (b *Broker) registerToKontrol() error {
	if err := kontrolhelper.RegisterToKontrol(
		conf,
		b.Config.Name,
		b.Config.Name, // servicGenericName
		b.ServiceUniqueName,
		*flagKontrolUUID,
		b.Hostname,
		b.Config.Port,
	); err != nil {
		return err
	}
	return nil
}

// startAMQP setups the the neccesary publisher and consumer connections for
// the broker broker.
func (b *Broker) startAMQP() error {
	b.PublishConn = amqputil.CreateConnection(conf, b.Config.Name)
	b.ConsumeConn = amqputil.CreateConnection(conf, b.Config.Name)
	consumeChannel := amqputil.CreateChannel(b.ConsumeConn)
	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		b.Config.Name,       // serviceType
		b.Config.Name,       // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", b.Config.Name, "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		return fmt.Errorf("Couldnt create updateInstances exchange  %v", err)
	}

	if err := consumeChannel.ExchangeBind(BROKER_NAME, "", "updateInstances", false, nil); err != nil {
		return fmt.Errorf("Couldnt bind to updateInstances exchange  %v", err)
	}

	go func(stream <-chan amqp.Delivery) {
		// start to listen from "broker" topic exchange
		for amqpMessage := range stream {
			broadcastUpdateInstances(amqpMessage)
		}

		b.Close()

	}(stream)

	return nil
}

func broadcastUpdateInstances(amqpMessage amqp.Delivery) {
	routingKeys := amqpMessage.RoutingKey

	payload := json.RawMessage(utils.FilterInvalidUTF8(amqpMessage.Body))

	splittedRoutingKeys := strings.Split(routingKeys, " ")
	for _, routingKey := range splittedRoutingKeys {
		processMessage(routingKey, payload)
	}
}

func processMessage(routingKey string, payload json.RawMessage) {
	pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
	for pos != -1 && pos < len(routingKey) {
		index := strings.IndexRune(routingKey[pos+1:], '.')
		pos += index + 1
		if index == -1 {
			pos = len(routingKey)
		}
		routingKeyPrefix := routingKey[:pos]
		globalMapMutex.Lock()

		if routes, ok := routeMap[routingKeyPrefix]; ok {
			routes.Each(func(sessionId interface{}) bool {
				if routeSession, ok := sessionsMap[sessionId.(string)]; ok {
					sendToClient(routeSession, routingKey, &payload)
				}
				return true
			})
		}
		globalMapMutex.Unlock()
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		conf.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	mux := &sockjs.Mux{
		Handlers: map[string]http.Handler{
			"/subscribe": service,
			"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				w.Write([]byte(strconv.Itoa(conf.BuildNumber)))
			}),
		},
	}

	server := &http.Server{Handler: mux}

	var err error
	b.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(b.Config.IP), Port: b.Config.Port})
	if err != nil {
		log.Fatal(err)
	}

	if b.Config.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(b.Config.CertFile, b.Config.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		b.listener = tls.NewListener(b.listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	// signal that we are ready now
	close(b.ready)

	lastErrorTime := time.Now()
	for {
		err := server.Serve(b.listener)
		if err != nil {
			// comes when the broker is closed with Close() method. This error
			// is defined in net/net.go as "var errClosing", unfortunaly it's
			// not exported.
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}

			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sockjsSession is called for every client connection and handles all the
// message trafic for a single client connection.
func (b *Broker) sockjsSession(session *sockjs.Session) {
	client, err := NewClient(session, b)
	if err != nil {
		log.Critical("Couldnt create client %v", err)
		return
	}

	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

	err = client.ControlChannel.Publish(b.Config.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(client.SocketId)})
	if err != nil {
		log.Critical("Couldnt publish to control channel %v", err)
	}

	sendToClient(session, "broker.connected", client.SocketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		client.handleSessionMessage(data)
	}
}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}
<commit_msg>Broker: add multiple data sending support in one message<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"koding/kontrol/kontrolhelper"
	"koding/tools/amqputil"
	"koding/tools/config"
	"koding/tools/lifecycle"
	"koding/tools/logger"
	"koding/tools/sockjs"
	"koding/tools/utils"
	"net"
	"net/http"
	"os"
	"os/signal"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/fatih/set"
	"github.com/streadway/amqp"
)

const BROKER_NAME = "broker"

var (
	conf *config.Config
	log  = logger.New(BROKER_NAME)

	STORAGE_BACKEND = "redis"
	routeMap        = make(map[string]*set.Set)
	sessionsMap     = make(map[string]*sockjs.Session)
	globalMapMutex  sync.Mutex

	changeClientsGauge          = lifecycle.CreateClientsGauge()
	changeNewClientsGauge       = logger.CreateCounterGauge("newClients", logger.NoUnit, true)
	changeWebsocketClientsGauge = logger.CreateCounterGauge("websocketClients", logger.NoUnit, false)

	flagProfile      = flag.String("c", "", "Configuration profile from file")
	flagBrokerDomain = flag.String("a", "", "Send kontrol a custom domain istead of os.Hostname")
	flagKontrolUUID  = flag.String("u", "", "Enable Kontrol mode")
	flagBrokerType   = flag.String("b", "broker", "Define broker type. Available: broker and brokerKite. B")
	flagDebug        = flag.Bool("d", false, "Debug mode")
)

// Broker is a router/multiplexer that routes messages coming from a SockJS
// server to an AMQP exchange and vice versa. Broker basically listens to
// client messages (Koding users) from the SockJS server. The message is
// either passed to the appropriate exchange or a response is sent back to the
// client. Each message has an "action" field that defines how to act for a
// received message.
type Broker struct {
	Config            *config.Broker
	Hostname          string
	ServiceUniqueName string
	AuthAllExchange   string
	PublishConn       *amqp.Connection
	ConsumeConn       *amqp.Connection

	// Accepts SockJS connections
	listener net.Listener

	// Closed when SockJS server is ready to acccept connections
	ready chan struct{}
}

// NewBroker returns a new Broker instance with ServiceUniqueName and Hostname
// prepopulated. After creating a Broker instance, one has to call
// broker.Run() or broker.Start() to start the broker instance and call
// broker.Close() for a graceful stop.
func NewBroker() *Broker {
	// returns os.Hostname() if config.BrokerDomain is empty, otherwise it just
	// returns config.BrokerDomain back
	brokerHostname := kontrolhelper.CustomHostname(*flagBrokerDomain)
	sanitizedHostname := strings.Replace(brokerHostname, ".", "_", -1)
	serviceUniqueName := BROKER_NAME + "|" + sanitizedHostname

	return &Broker{
		Hostname:          brokerHostname,
		ServiceUniqueName: serviceUniqueName,
		ready:             make(chan struct{}),
	}
}

func main() {
	flag.Parse()
	if *flagProfile == "" {
		log.Fatal("Please specify profile via -c. Aborting.")
	}

	conf = config.MustConfig(*flagProfile)
	broker := NewBroker()

	switch *flagBrokerType {
	case "brokerKite":
		broker.Config = &conf.BrokerKite
	case "premiumBrokerKite":
		broker.Config = &conf.PremiumBrokerKite
	default:
		broker.Config = &conf.Broker
	}

	// update broker name
	log = logger.New(broker.Config.Name)
	var logLevel logger.Level
	if *flagDebug {
		logLevel = logger.DEBUG
	} else {
		logLevel = logger.GetLoggingLevelFromConfig(BROKER_NAME, *flagProfile)
	}

	log.SetLevel(logLevel)
	broker.Run()
}

// Run starts the broker.
func (b *Broker) Run() {
	// sets the maximum number of CPUs that can be executing simultaneously
	runtime.GOMAXPROCS(runtime.NumCPU())

	lifecycle.Startup(BROKER_NAME, false)
	logger.RunGaugesLoop(log)

	if err := b.registerToKontrol(); err != nil {
		log.Critical("Couldnt register to kontrol, stopping... %v", err)
		return
	}

	if err := b.startAMQP(); err != nil {
		log.Critical("Couldnt create amqp bindings, stopping... %v", err)
		return
	}

	b.startSockJS() // blocking
}

// Start is like Run() but waits until the SockJS listener is ready to be
// used.
func (b *Broker) Start() {
	go b.Run()
	<-b.ready
}

// Close close all amqp connections and closes the SockJS server listener
func (b *Broker) Close() {
	b.PublishConn.Close()
	b.ConsumeConn.Close()
	b.listener.Close()
}

// registerToKontrol registers the broker to KontrolDaemon. This is needed to
// populate a list of brokers and show them to the client. The list is
// available at: https://koding.com/-/services/broker?all
func (b *Broker) registerToKontrol() error {
	if err := kontrolhelper.RegisterToKontrol(
		conf,
		b.Config.Name,
		b.Config.Name, // servicGenericName
		b.ServiceUniqueName,
		*flagKontrolUUID,
		b.Hostname,
		b.Config.Port,
	); err != nil {
		return err
	}
	return nil
}

// startAMQP setups the the neccesary publisher and consumer connections for
// the broker broker.
func (b *Broker) startAMQP() error {
	b.PublishConn = amqputil.CreateConnection(conf, b.Config.Name)
	b.ConsumeConn = amqputil.CreateConnection(conf, b.Config.Name)
	consumeChannel := amqputil.CreateChannel(b.ConsumeConn)
	presenceQueue := amqputil.JoinPresenceExchange(
		consumeChannel,      // channel
		"services-presence", // exchange
		b.Config.Name,       // serviceType
		b.Config.Name,       // serviceGenericName
		b.ServiceUniqueName, // serviceUniqueName
		false,               // loadBalancing
	)

	go func() {
		sigusr1Channel := make(chan os.Signal)
		signal.Notify(sigusr1Channel, syscall.SIGUSR1)
		<-sigusr1Channel
		consumeChannel.QueueDelete(presenceQueue, false, false, false)
	}()

	stream := amqputil.DeclareBindConsumeQueue(consumeChannel, "topic", b.Config.Name, "#", false)

	if err := consumeChannel.ExchangeDeclare(
		"updateInstances", // name
		"fanout",          // kind
		false,             // durable
		false,             // autoDelete
		false,             // internal
		false,             // noWait
		nil,               // args
	); err != nil {
		return fmt.Errorf("Couldnt create updateInstances exchange  %v", err)
	}

	if err := consumeChannel.ExchangeBind(BROKER_NAME, "", "updateInstances", false, nil); err != nil {
		return fmt.Errorf("Couldnt bind to updateInstances exchange  %v", err)
	}

	go func(stream <-chan amqp.Delivery) {
		// start to listen from "broker" topic exchange
		for amqpMessage := range stream {
			broadcastUpdateInstances(amqpMessage)
		}

		b.Close()

	}(stream)

	return nil
}

func broadcastUpdateInstances(amqpMessage amqp.Delivery) {
	routingKey := amqpMessage.RoutingKey
	payloadsByte := utils.FilterInvalidUTF8(amqpMessage.Body)

	var payloads map[string]interface{}

	if err := json.Unmarshal(payloadsByte, &payloads); err != nil {
		log.Error("Error while unmarshalling %v", err)
		return
	}

	// payloadsArr := payloads.([]interface{})
	for _, payload := range payloads {
		payloadByte, err := json.Marshal(payload)
		if err != nil {
			log.Error("Error while marshalling %v", err)
			continue
		}
		processMessage(routingKey, payloadByte)
	}
}

func processMessage(routingKey string, payload interface{}) {
	pos := strings.IndexRune(routingKey, '.') // skip first dot, since we want at least two components to always include the secret
	for pos != -1 && pos < len(routingKey) {
		index := strings.IndexRune(routingKey[pos+1:], '.')
		pos += index + 1
		if index == -1 {
			pos = len(routingKey)
		}
		routingKeyPrefix := routingKey[:pos]
		globalMapMutex.Lock()

		if routes, ok := routeMap[routingKeyPrefix]; ok {
			routes.Each(func(sessionId interface{}) bool {
				if routeSession, ok := sessionsMap[sessionId.(string)]; ok {
					sendToClient(routeSession, routingKey, &payload)
				}
				return true
			})
		}
		globalMapMutex.Unlock()
	}
}

// startSockJS starts a new HTTPS listener that implies the SockJS protocol.
func (b *Broker) startSockJS() {
	service := sockjs.NewService(
		conf.Client.StaticFilesBaseUrl+"/js/sock.js",
		10*time.Minute,
		b.sockjsSession,
	)
	defer service.Close()

	service.MaxReceivedPerSecond = 50
	service.ErrorHandler = log.LogError

	// TODO use http.Mux instead of sockjs.Mux.
	mux := &sockjs.Mux{
		Handlers: map[string]http.Handler{
			"/subscribe": service,
			"/buildnumber": http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				w.Header().Set("Content-Type", "text/plain")
				w.Write([]byte(strconv.Itoa(conf.BuildNumber)))
			}),
		},
	}

	server := &http.Server{Handler: mux}

	var err error
	b.listener, err = net.ListenTCP("tcp", &net.TCPAddr{IP: net.ParseIP(b.Config.IP), Port: b.Config.Port})
	if err != nil {
		log.Fatal(err)
	}

	if b.Config.CertFile != "" {
		cert, err := tls.LoadX509KeyPair(b.Config.CertFile, b.Config.KeyFile)
		if err != nil {
			log.Fatal(err)
		}
		b.listener = tls.NewListener(b.listener, &tls.Config{
			NextProtos:   []string{"http/1.1"},
			Certificates: []tls.Certificate{cert},
		})
	}

	// signal that we are ready now
	close(b.ready)

	lastErrorTime := time.Now()
	for {
		err := server.Serve(b.listener)
		if err != nil {
			// comes when the broker is closed with Close() method. This error
			// is defined in net/net.go as "var errClosing", unfortunaly it's
			// not exported.
			if strings.Contains(err.Error(), "use of closed network connection") {
				return
			}

			log.Warning("Server error: %v", err)
			if time.Now().Sub(lastErrorTime) < time.Second {
				log.Fatal(nil)
			}
			lastErrorTime = time.Now()
		}
	}

}

// sockjsSession is called for every client connection and handles all the
// message trafic for a single client connection.
func (b *Broker) sockjsSession(session *sockjs.Session) {
	client, err := NewClient(session, b)
	if err != nil {
		log.Critical("Couldnt create client %v", err)
		return
	}

	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

	err = client.ControlChannel.Publish(b.Config.AuthAllExchange, "broker.clientConnected", false, false, amqp.Publishing{Body: []byte(client.SocketId)})
	if err != nil {
		log.Critical("Couldnt publish to control channel %v", err)
	}

	sendToClient(session, "broker.connected", client.SocketId)

	for data := range session.ReceiveChan {
		if data == nil || session.Closed {
			break
		}

		client.handleSessionMessage(data)
	}
}

// sendToClient sends the given payload back to the client. It attachs the
// routintKey along with the payload. It closes the session if sending fails.
func sendToClient(session *sockjs.Session, routingKey string, payload interface{}) {
	var message struct {
		RoutingKey string      `json:"routingKey"`
		Payload    interface{} `json:"payload"`
	}
	message.RoutingKey = routingKey
	message.Payload = payload
	if !session.Send(message) {
		session.Close()
		log.Warning("Dropped session because of broker to client buffer overflow. %v", session.Tag)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"context"
	"fmt"
	"time"
)

func doSomething(ctx context.Context) {
	ctx, cancelCtx := context.WithCancel(ctx)
	printCh := make(chan int)

	go doAnother(ctx, printCh)
	for num := 1; num <= 3; num++{
		printCh <- num
	}

	cancelCtx()

	time.Sleep(100 * time.Millisecond)

	fmt.Printf("doSomething ended\n")
}

func doAnother(ctx context.Context, printCh <-chan int) {
	for {
		select {
		case <-ctx.Done():
			if err := ctx.Err(); err != nil {
				fmt.Printf("doAnother err: %s\n", err)
			}
			fmt.Printf("doAnother ended\n")
			return
		case num := <-printCh:
			fmt.Printf("doAnother: %d\n", num)

		}
	}
}

func main() {
	ctx := context.TODO()

	ctx = context.WithValue(ctx, "myKey","myValue")

	doSomething(ctx)
}
<commit_msg>context with deadline<commit_after>package main

import (
	"context"
	"fmt"
	"time"
)

func doSomething(ctx context.Context) {
	deadline := time.Now().Add(1500 * time.Millisecond)
	ctx, cancelCtx := context.WithDeadline(ctx, deadline)
	defer cancelCtx()

	printCh := make(chan int)

	go doAnother(ctx, printCh)
	for num := 1; num <= 3; num++{
		select {
			case printCh <- num:
				time.Sleep(1 * time.Second)
			case <-ctx.Done():
				break
		}
	}

	cancelCtx()

	time.Sleep(100 * time.Millisecond)

	fmt.Printf("doSomething ended\n")
}

func doAnother(ctx context.Context, printCh <-chan int) {
	for {
		select {
		case <-ctx.Done():
			if err := ctx.Err(); err != nil {
				fmt.Printf("doAnother err: %s\n", err)
			}
			fmt.Printf("doAnother ended\n")
			return
		case num := <-printCh:
			fmt.Printf("doAnother: %d\n", num)

		}
	}
}

func main() {
	ctx := context.TODO()

	ctx = context.WithValue(ctx, "myKey","myValue")

	doSomething(ctx)
}
<|endoftext|>