<commit_before>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package exporter

import (
	"errors"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/metrics/datum"
	"github.com/google/mtail/internal/testutil"
)

const prefix = "prefix"

func TestCreateExporter(t *testing.T) {
	_, err := New(nil)
	if err == nil {
		t.Error("expecting error, got nil")
	}
	store := metrics.NewStore()
	_, err = New(store)
	if err != nil {
		t.Errorf("unexpected error:%s", err)
	}
	failopt := func(*Exporter) error {
		return errors.New("busted")
	}
	_, err = New(store, failopt)
	if err == nil {
		t.Errorf("unexpected success")
	}
}

func FakeSocketWrite(f formatter, m *metrics.Metric) []string {
	// TODO(jaq): urgh looking inside m to find preallocation size
	ret := make([]string, 0, len(m.LabelValues))
	lc := make(chan *metrics.LabelSet)
	go m.EmitLabelSets(lc)
	for l := range lc {
		ret = append(ret, f("gunstar", m, l))
	}
	sort.Strings(ret)
	return ret
}

func TestMetricToCollectd(t *testing.T) {
	ts, terr := time.Parse("2006/01/02 15:04:05", "2012/07/24 10:14:00")
	if terr != nil {
		t.Errorf("time parse error: %s", terr)
	}
	ms := metrics.NewStore()

	scalarMetric := metrics.NewMetric("foo", "prog", metrics.Counter, metrics.Int)
	d, _ := scalarMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	testutil.FatalIfErr(t, ms.Add(scalarMetric))

	r := FakeSocketWrite(metricToCollectd, scalarMetric)
	expected := []string{"PUTVAL \"gunstar/mtail-prog/counter-foo\" interval=60 1343124840:37\n"}
	testutil.ExpectNoDiff(t, expected, r)

	dimensionedMetric := metrics.NewMetric("bar", "prog", metrics.Gauge, metrics.Int, "label")
	d, _ = dimensionedMetric.GetDatum("quux")
	datum.SetInt(d, 37, ts)
	d, _ = dimensionedMetric.GetDatum("snuh")
	datum.SetInt(d, 37, ts)
	ms.ClearMetrics()
	testutil.FatalIfErr(t, ms.Add(dimensionedMetric))

	r = FakeSocketWrite(metricToCollectd, dimensionedMetric)
	expected = []string{
		"PUTVAL \"gunstar/mtail-prog/gauge-bar-label-quux\" interval=60 1343124840:37\n",
		"PUTVAL \"gunstar/mtail-prog/gauge-bar-label-snuh\" interval=60 1343124840:37\n"}
	testutil.ExpectNoDiff(t, expected, r)

	timingMetric := metrics.NewMetric("foo", "prog", metrics.Timer, metrics.Int)
	d, _ = timingMetric.GetDatum()
	datum.SetInt(d, 123, ts)
	testutil.FatalIfErr(t, ms.Add(timingMetric))

	r = FakeSocketWrite(metricToCollectd, timingMetric)
	expected = []string{"PUTVAL \"gunstar/mtail-prog/gauge-foo\" interval=60 1343124840:123\n"}
	testutil.ExpectNoDiff(t, expected, r)

	*collectdPrefix = prefix
	r = FakeSocketWrite(metricToCollectd, timingMetric)
	expected = []string{"PUTVAL \"gunstar/prefixmtail-prog/gauge-foo\" interval=60 1343124840:123\n"}
	testutil.ExpectNoDiff(t, expected, r)
}

func TestMetricToGraphite(t *testing.T) {
	ts, terr := time.Parse("2006/01/02 15:04:05", "2012/07/24 10:14:00")
	if terr != nil {
		t.Errorf("time parse error: %s", terr)
	}

	scalarMetric := metrics.NewMetric("foo", "prog", metrics.Counter, metrics.Int)
	d, _ := scalarMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	r := FakeSocketWrite(metricToGraphite, scalarMetric)
	expected := []string{"prog.foo 37 1343124840\n"}
	testutil.ExpectNoDiff(t, expected, r)

	dimensionedMetric := metrics.NewMetric("bar", "prog", metrics.Gauge, metrics.Int, "host")
	d, _ = dimensionedMetric.GetDatum("quux.com")
	datum.SetInt(d, 37, ts)
	d, _ = dimensionedMetric.GetDatum("snuh.teevee")
	datum.SetInt(d, 37, ts)
	r = FakeSocketWrite(metricToGraphite, dimensionedMetric)
	expected = []string{
		"prog.bar.host.quux_com 37 1343124840\n",
		"prog.bar.host.snuh_teevee 37 1343124840\n"}
	testutil.ExpectNoDiff(t, expected, r)

	*graphitePrefix = prefix
	r = FakeSocketWrite(metricToGraphite, dimensionedMetric)
	expected = []string{
		"prefixprog.bar.host.quux_com 37 1343124840\n",
		"prefixprog.bar.host.snuh_teevee 37 1343124840\n"}
	testutil.ExpectNoDiff(t, expected, r)
}

func TestMetricToStatsd(t *testing.T) {
	ts, terr := time.Parse("2006/01/02 15:04:05", "2012/07/24 10:14:00")
	if terr != nil {
		t.Errorf("time parse error: %s", terr)
	}

	scalarMetric := metrics.NewMetric("foo", "prog", metrics.Counter, metrics.Int)
	d, _ := scalarMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	r := FakeSocketWrite(metricToStatsd, scalarMetric)
	expected := []string{"prog.foo:37|c"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("String didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}

	dimensionedMetric := metrics.NewMetric("bar", "prog", metrics.Gauge, metrics.Int, "l")
	d, _ = dimensionedMetric.GetDatum("quux")
	datum.SetInt(d, 37, ts)
	d, _ = dimensionedMetric.GetDatum("snuh")
	datum.SetInt(d, 42, ts)
	r = FakeSocketWrite(metricToStatsd, dimensionedMetric)
	expected = []string{
		"prog.bar.l.quux:37|g",
		"prog.bar.l.snuh:42|g"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("String didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}

	timingMetric := metrics.NewMetric("foo", "prog", metrics.Timer, metrics.Int)
	d, _ = timingMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	r = FakeSocketWrite(metricToStatsd, timingMetric)
	expected = []string{"prog.foo:37|ms"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("String didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}

	*statsdPrefix = prefix
	r = FakeSocketWrite(metricToStatsd, timingMetric)
	expected = []string{"prefixprog.foo:37|ms"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("prefixed string didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}
}
<commit_msg>Initialise flag global in test to avoid side effects.<commit_after>// Copyright 2011 Google Inc. All Rights Reserved.
// This file is available under the Apache license.

package exporter

import (
	"errors"
	"reflect"
	"sort"
	"testing"
	"time"

	"github.com/google/mtail/internal/metrics"
	"github.com/google/mtail/internal/metrics/datum"
	"github.com/google/mtail/internal/testutil"
)

const prefix = "prefix"

func TestCreateExporter(t *testing.T) {
	_, err := New(nil)
	if err == nil {
		t.Error("expecting error, got nil")
	}
	store := metrics.NewStore()
	_, err = New(store)
	if err != nil {
		t.Errorf("unexpected error:%s", err)
	}
	failopt := func(*Exporter) error {
		return errors.New("busted")
	}
	_, err = New(store, failopt)
	if err == nil {
		t.Errorf("unexpected success")
	}
}

func FakeSocketWrite(f formatter, m *metrics.Metric) []string {
	ret := make([]string, 0)
	lc := make(chan *metrics.LabelSet)
	go m.EmitLabelSets(lc)
	for l := range lc {
		ret = append(ret, f("gunstar", m, l))
	}
	sort.Strings(ret)
	return ret
}

func TestMetricToCollectd(t *testing.T) {
	*collectdPrefix = ""
	ts, terr := time.Parse("2006/01/02 15:04:05", "2012/07/24 10:14:00")
	if terr != nil {
		t.Errorf("time parse error: %s", terr)
	}
	ms := metrics.NewStore()

	scalarMetric := metrics.NewMetric("foo", "prog", metrics.Counter, metrics.Int)
	d, _ := scalarMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	testutil.FatalIfErr(t, ms.Add(scalarMetric))

	r := FakeSocketWrite(metricToCollectd, scalarMetric)
	expected := []string{"PUTVAL \"gunstar/mtail-prog/counter-foo\" interval=60 1343124840:37\n"}
	testutil.ExpectNoDiff(t, expected, r)

	dimensionedMetric := metrics.NewMetric("bar", "prog", metrics.Gauge, metrics.Int, "label")
	d, _ = dimensionedMetric.GetDatum("quux")
	datum.SetInt(d, 37, ts)
	d, _ = dimensionedMetric.GetDatum("snuh")
	datum.SetInt(d, 37, ts)
	ms.ClearMetrics()
	testutil.FatalIfErr(t, ms.Add(dimensionedMetric))

	r = FakeSocketWrite(metricToCollectd, dimensionedMetric)
	expected = []string{
		"PUTVAL \"gunstar/mtail-prog/gauge-bar-label-quux\" interval=60 1343124840:37\n",
		"PUTVAL \"gunstar/mtail-prog/gauge-bar-label-snuh\" interval=60 1343124840:37\n"}
	testutil.ExpectNoDiff(t, expected, r)

	timingMetric := metrics.NewMetric("foo", "prog", metrics.Timer, metrics.Int)
	d, _ = timingMetric.GetDatum()
	datum.SetInt(d, 123, ts)
	testutil.FatalIfErr(t, ms.Add(timingMetric))

	r = FakeSocketWrite(metricToCollectd, timingMetric)
	expected = []string{"PUTVAL \"gunstar/mtail-prog/gauge-foo\" interval=60 1343124840:123\n"}
	testutil.ExpectNoDiff(t, expected, r)

	*collectdPrefix = prefix
	r = FakeSocketWrite(metricToCollectd, timingMetric)
	expected = []string{"PUTVAL \"gunstar/prefixmtail-prog/gauge-foo\" interval=60 1343124840:123\n"}
	testutil.ExpectNoDiff(t, expected, r)
}

func TestMetricToGraphite(t *testing.T) {
	*graphitePrefix = ""
	ts, terr := time.Parse("2006/01/02 15:04:05", "2012/07/24 10:14:00")
	if terr != nil {
		t.Errorf("time parse error: %s", terr)
	}

	scalarMetric := metrics.NewMetric("foo", "prog", metrics.Counter, metrics.Int)
	d, _ := scalarMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	r := FakeSocketWrite(metricToGraphite, scalarMetric)
	expected := []string{"prog.foo 37 1343124840\n"}
	testutil.ExpectNoDiff(t, expected, r)

	dimensionedMetric := metrics.NewMetric("bar", "prog", metrics.Gauge, metrics.Int, "host")
	d, _ = dimensionedMetric.GetDatum("quux.com")
	datum.SetInt(d, 37, ts)
	d, _ = dimensionedMetric.GetDatum("snuh.teevee")
	datum.SetInt(d, 37, ts)
	r = FakeSocketWrite(metricToGraphite, dimensionedMetric)
	expected = []string{
		"prog.bar.host.quux_com 37 1343124840\n",
		"prog.bar.host.snuh_teevee 37 1343124840\n"}
	testutil.ExpectNoDiff(t, expected, r)

	*graphitePrefix = prefix
	r = FakeSocketWrite(metricToGraphite, dimensionedMetric)
	expected = []string{
		"prefixprog.bar.host.quux_com 37 1343124840\n",
		"prefixprog.bar.host.snuh_teevee 37 1343124840\n"}
	testutil.ExpectNoDiff(t, expected, r)
}

func TestMetricToStatsd(t *testing.T) {
	*statsdPrefix = ""
	ts, terr := time.Parse("2006/01/02 15:04:05", "2012/07/24 10:14:00")
	if terr != nil {
		t.Errorf("time parse error: %s", terr)
	}

	scalarMetric := metrics.NewMetric("foo", "prog", metrics.Counter, metrics.Int)
	d, _ := scalarMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	r := FakeSocketWrite(metricToStatsd, scalarMetric)
	expected := []string{"prog.foo:37|c"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("String didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}

	dimensionedMetric := metrics.NewMetric("bar", "prog", metrics.Gauge, metrics.Int, "l")
	d, _ = dimensionedMetric.GetDatum("quux")
	datum.SetInt(d, 37, ts)
	d, _ = dimensionedMetric.GetDatum("snuh")
	datum.SetInt(d, 42, ts)
	r = FakeSocketWrite(metricToStatsd, dimensionedMetric)
	expected = []string{
		"prog.bar.l.quux:37|g",
		"prog.bar.l.snuh:42|g"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("String didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}

	timingMetric := metrics.NewMetric("foo", "prog", metrics.Timer, metrics.Int)
	d, _ = timingMetric.GetDatum()
	datum.SetInt(d, 37, ts)
	r = FakeSocketWrite(metricToStatsd, timingMetric)
	expected = []string{"prog.foo:37|ms"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("String didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}

	*statsdPrefix = prefix
	r = FakeSocketWrite(metricToStatsd, timingMetric)
	expected = []string{"prefixprog.foo:37|ms"}
	if !reflect.DeepEqual(expected, r) {
		t.Errorf("prefixed string didn't match:\n\texpected: %v\n\treceived: %v", expected, r)
	}
}
<|endoftext|>


<commit_before>// Copyright 2018 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package planner

import (
	"os"
	"testing"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/internal/ir"
)

func TestPlannerHelloWorld(t *testing.T) {

	// NOTE(tsandall): These tests are not meant to give comprehensive coverage
	// of the planner. Currently we have a suite of end-to-end tests in the
	// test/wasm/ directory that are specified in YAML, compiled into Wasm, and
	// executed inside of a node program. For the time being, the planner is
	// simple enough that exhaustive unit testing is not as valuable as
	// end-to-end testing. These tests provide a quick sanity check that the
	// planner is not failing on simple inputs.
	tests := []struct {
		note    string
		queries []string
		modules []string
		exp     ir.Policy
	}{
		{
			note:    "empty",
			queries: []string{},
		},
		{
			note:    "hello world",
			queries: []string{"input.a = 1"},
		},
		{
			note:    "conjunction",
			queries: []string{"1 = 1; 2 = 2"},
		},
		{
			note:    "disjunction",
			queries: []string{"input.a = 1", "input.b = 2"},
		},
		{
			note:    "iteration",
			queries: []string{"input.a[i] = 1; input.b = 2"},
		},
		{
			note:    "iteration: compare key",
			queries: []string{"input.a[i] = 1; input.b = i"},
		},
		{
			note:    "iteration: nested",
			queries: []string{"input.a[i] = 1; input.b[j] = 2"},
		},
		{
			note:    "iteration: chained",
			queries: []string{"input.a[i][j] = 1"},
		},
		{
			note:    "negation",
			queries: []string{"not input.x.y = 1"},
		},
		{
			note:    "array ref pattern match",
			queries: []string{"input.x = [1, [y]]"},
		},
		{
			note:    "arrays pattern match",
			queries: []string{"[x, 3, [2]] = [1, 3, [y]]"},
		},
		{
			note:    "sets",
			queries: []string{"x = {1,2,3}; x[y]"},
		},
		{
			note:    "vars",
			queries: []string{"x = 1"},
		},
		{
			note:    "complete rules",
			queries: []string{"true"},
			modules: []string{`
				package test
				p = x { x = 1 }
				p = y { y = 2 }
			`},
		},
		{
			note:    "complete rule reference",
			queries: []string{"data.test.p = 10"},
			modules: []string{`
				package test
				p = x { x = 10 }
			`},
		},
		{
			note:    "functions",
			queries: []string{"data.test.f([1,x])"},
			modules: []string{`
				package test
				f([a, b]) {
					a = b
				}
			`},
		},
		{
			note:    "else",
			queries: []string{"data.test.p = 1"},
			modules: []string{`
				package test
				p = 0 {
					false
				} else = 1 {
					true
				}
			`},
		},
		{
			note:    "partial set",
			queries: []string{"data.test.p = {1,2}"},
			modules: []string{`
				package test
				p[1]
				p[2]
			`},
		},
		{
			note:    "partial object",
			queries: []string{`data.test.p = {"a": 1, "b": 2}`},
			modules: []string{`
				package test
				p["a"] = 1
				p["b"] = 2
			`},
		},
		{
			note:    "virtual extent",
			queries: []string{`data`},
			modules: []string{`
				package test

				p = 1
			#	q = 2 { false }
			`},
		},
	}

	for _, tc := range tests {
		t.Run(tc.note, func(t *testing.T) {
			queries := make([]ast.Body, len(tc.queries))
			for i := range queries {
				queries[i] = ast.MustParseBody(tc.queries[i])
			}
			modules := make([]*ast.Module, len(tc.modules))
			for i := range modules {
				modules[i] = ast.MustParseModule(tc.modules[i])
			}
			planner := New().WithQueries(queries).WithModules(modules)
			policy, err := planner.Plan()
			if err != nil {
				t.Fatal(err)
			}
			_ = policy
			ir.Pretty(os.Stderr, policy)
		})
	}
}
<commit_msg>planner: Uncomment test rule<commit_after>// Copyright 2018 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package planner

import (
	"os"
	"testing"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/internal/ir"
)

func TestPlannerHelloWorld(t *testing.T) {

	// NOTE(tsandall): These tests are not meant to give comprehensive coverage
	// of the planner. Currently we have a suite of end-to-end tests in the
	// test/wasm/ directory that are specified in YAML, compiled into Wasm, and
	// executed inside of a node program. For the time being, the planner is
	// simple enough that exhaustive unit testing is not as valuable as
	// end-to-end testing. These tests provide a quick sanity check that the
	// planner is not failing on simple inputs.
	tests := []struct {
		note    string
		queries []string
		modules []string
		exp     ir.Policy
	}{
		{
			note:    "empty",
			queries: []string{},
		},
		{
			note:    "hello world",
			queries: []string{"input.a = 1"},
		},
		{
			note:    "conjunction",
			queries: []string{"1 = 1; 2 = 2"},
		},
		{
			note:    "disjunction",
			queries: []string{"input.a = 1", "input.b = 2"},
		},
		{
			note:    "iteration",
			queries: []string{"input.a[i] = 1; input.b = 2"},
		},
		{
			note:    "iteration: compare key",
			queries: []string{"input.a[i] = 1; input.b = i"},
		},
		{
			note:    "iteration: nested",
			queries: []string{"input.a[i] = 1; input.b[j] = 2"},
		},
		{
			note:    "iteration: chained",
			queries: []string{"input.a[i][j] = 1"},
		},
		{
			note:    "negation",
			queries: []string{"not input.x.y = 1"},
		},
		{
			note:    "array ref pattern match",
			queries: []string{"input.x = [1, [y]]"},
		},
		{
			note:    "arrays pattern match",
			queries: []string{"[x, 3, [2]] = [1, 3, [y]]"},
		},
		{
			note:    "sets",
			queries: []string{"x = {1,2,3}; x[y]"},
		},
		{
			note:    "vars",
			queries: []string{"x = 1"},
		},
		{
			note:    "complete rules",
			queries: []string{"true"},
			modules: []string{`
				package test
				p = x { x = 1 }
				p = y { y = 2 }
			`},
		},
		{
			note:    "complete rule reference",
			queries: []string{"data.test.p = 10"},
			modules: []string{`
				package test
				p = x { x = 10 }
			`},
		},
		{
			note:    "functions",
			queries: []string{"data.test.f([1,x])"},
			modules: []string{`
				package test
				f([a, b]) {
					a = b
				}
			`},
		},
		{
			note:    "else",
			queries: []string{"data.test.p = 1"},
			modules: []string{`
				package test
				p = 0 {
					false
				} else = 1 {
					true
				}
			`},
		},
		{
			note:    "partial set",
			queries: []string{"data.test.p = {1,2}"},
			modules: []string{`
				package test
				p[1]
				p[2]
			`},
		},
		{
			note:    "partial object",
			queries: []string{`data.test.p = {"a": 1, "b": 2}`},
			modules: []string{`
				package test
				p["a"] = 1
				p["b"] = 2
			`},
		},
		{
			note:    "virtual extent",
			queries: []string{`data`},
			modules: []string{`
				package test

				p = 1
				q = 2 { false }
			`},
		},
	}

	for _, tc := range tests {
		t.Run(tc.note, func(t *testing.T) {
			queries := make([]ast.Body, len(tc.queries))
			for i := range queries {
				queries[i] = ast.MustParseBody(tc.queries[i])
			}
			modules := make([]*ast.Module, len(tc.modules))
			for i := range modules {
				modules[i] = ast.MustParseModule(tc.modules[i])
			}
			planner := New().WithQueries(queries).WithModules(modules)
			policy, err := planner.Plan()
			if err != nil {
				t.Fatal(err)
			}
			_ = policy
			ir.Pretty(os.Stderr, policy)
		})
	}
}
<|endoftext|>


<commit_before>package provider

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"math/big"
	"strings"

	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
)

func resourceId() *schema.Resource {
	return &schema.Resource{
		Description: `
The resource ` + "`random_id`" + ` generates random numbers that are intended to be
used as unique identifiers for other resources.

This resource *does* use a cryptographic random number generator in order
to minimize the chance of collisions, making the results of this resource
when a 16-byte identifier is requested of equivalent uniqueness to a
type-4 UUID.

This resource can be used in conjunction with resources that have
the ` + "`create_before_destroy`" + ` lifecycle flag set to avoid conflicts with
unique names during the brief period where both the old and new resources
exist concurrently.
`,
		CreateContext: CreateID,
		ReadContext:   RepopulateEncodings,
		Delete:        schema.RemoveFromState,
		Importer: &schema.ResourceImporter{
			State: ImportID,
		},

		Schema: map[string]*schema.Schema{
			"keepers": {
				Description: "Arbitrary map of values that, when changed, will trigger recreation of " +
					"resource. See [the main provider documentation](../index.html) for more information.",
				Type:     schema.TypeMap,
				Optional: true,
				ForceNew: true,
			},

			"byte_length": {
				Description: "The number of random bytes to produce. The minimum value is 1, which produces " +
					"eight bits of randomness.",
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"prefix": {
				Description: "Arbitrary string to prefix the output value with. This string is supplied as-is, " +
					"meaning it is not guaranteed to be URL-safe or base64 encoded.",
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"b64_url": {
				Description: "The generated id presented in base64, using the URL-friendly character set: " +
					"case-sensitive letters, digits and the characters `_` and `-`.",
				Type:     schema.TypeString,
				Computed: true,
			},

			"b64_std": {
				Description: "The generated id presented in base64 without additional transformations.",
				Type:        schema.TypeString,
				Computed:    true,
			},

			"hex": {
				Description: "The generated id presented in padded hexadecimal digits. This result will " +
					"always be twice as long as the requested byte length.",
				Type:     schema.TypeString,
				Computed: true,
			},

			"dec": {
				Description: "The generated id presented in non-padded decimal digits.",
				Type:        schema.TypeString,
				Computed:    true,
			},

			"id": {
				Description: "The generated id presented in base64 without additional transformations or prefix.",
				Type:        schema.TypeString,
				Computed:    true,
			},
		},
	}
}

func CreateID(ctx context.Context, d *schema.ResourceData, meta interface{}) diag.Diagnostics {
	var diags diag.Diagnostics
	byteLength := d.Get("byte_length").(int)
	bytes := make([]byte, byteLength)

	n, err := rand.Reader.Read(bytes)
	if n != byteLength {
		return append(diags, diag.Errorf("generated insufficient random bytes: %s", err)...)
	}
	if err != nil {
		return append(diags, diag.Errorf("error generating random bytes: %s", err)...)
	}

	b64Str := base64.RawURLEncoding.EncodeToString(bytes)
	d.SetId(b64Str)

	repopEncsDiags := RepopulateEncodings(ctx, d, meta)
	if repopEncsDiags != nil {
		return append(diags, repopEncsDiags...)
	}

	return diags
}

func RepopulateEncodings(ctx context.Context, d *schema.ResourceData, _ interface{}) diag.Diagnostics {
	var diags diag.Diagnostics
	prefix := d.Get("prefix").(string)
	base64Str := d.Id()

	bytes, err := base64.RawURLEncoding.DecodeString(base64Str)
	if err != nil {
		return append(diags, diag.Errorf("error decoding ID: %s", err)...)
	}

	b64StdStr := base64.StdEncoding.EncodeToString(bytes)
	hexStr := hex.EncodeToString(bytes)

	bigInt := big.Int{}
	bigInt.SetBytes(bytes)
	decStr := bigInt.String()

	d.Set("b64_url", prefix+base64Str)
	d.Set("b64_std", prefix+b64StdStr)

	d.Set("hex", prefix+hexStr)
	d.Set("dec", prefix+decStr)

	return nil
}

func ImportID(d *schema.ResourceData, meta interface{}) ([]*schema.ResourceData, error) {
	id := d.Id()

	sep := strings.LastIndex(id, ",")
	if sep != -1 {
		d.Set("prefix", id[:sep])
		id = id[sep+1:]
	}

	bytes, err := base64.RawURLEncoding.DecodeString(id)
	if err != nil {
		return nil, errwrap.Wrapf("Error decoding ID: {{err}}", err)
	}

	d.Set("byte_length", len(bytes))
	d.SetId(id)

	return []*schema.ResourceData{d}, nil
}
<commit_msg>Using blank identifier for unused context passed to RepopulateEncodings func in resource_id (#230)<commit_after>package provider

import (
	"context"
	"crypto/rand"
	"encoding/base64"
	"encoding/hex"
	"math/big"
	"strings"

	"github.com/hashicorp/errwrap"
	"github.com/hashicorp/terraform-plugin-sdk/v2/diag"
	"github.com/hashicorp/terraform-plugin-sdk/v2/helper/schema"
)

func resourceId() *schema.Resource {
	return &schema.Resource{
		Description: `
The resource ` + "`random_id`" + ` generates random numbers that are intended to be
used as unique identifiers for other resources.

This resource *does* use a cryptographic random number generator in order
to minimize the chance of collisions, making the results of this resource
when a 16-byte identifier is requested of equivalent uniqueness to a
type-4 UUID.

This resource can be used in conjunction with resources that have
the ` + "`create_before_destroy`" + ` lifecycle flag set to avoid conflicts with
unique names during the brief period where both the old and new resources
exist concurrently.
`,
		CreateContext: CreateID,
		ReadContext:   RepopulateEncodings,
		Delete:        schema.RemoveFromState,
		Importer: &schema.ResourceImporter{
			State: ImportID,
		},

		Schema: map[string]*schema.Schema{
			"keepers": {
				Description: "Arbitrary map of values that, when changed, will trigger recreation of " +
					"resource. See [the main provider documentation](../index.html) for more information.",
				Type:     schema.TypeMap,
				Optional: true,
				ForceNew: true,
			},

			"byte_length": {
				Description: "The number of random bytes to produce. The minimum value is 1, which produces " +
					"eight bits of randomness.",
				Type:     schema.TypeInt,
				Required: true,
				ForceNew: true,
			},

			"prefix": {
				Description: "Arbitrary string to prefix the output value with. This string is supplied as-is, " +
					"meaning it is not guaranteed to be URL-safe or base64 encoded.",
				Type:     schema.TypeString,
				Optional: true,
				ForceNew: true,
			},

			"b64_url": {
				Description: "The generated id presented in base64, using the URL-friendly character set: " +
					"case-sensitive letters, digits and the characters `_` and `-`.",
				Type:     schema.TypeString,
				Computed: true,
			},

			"b64_std": {
				Description: "The generated id presented in base64 without additional transformations.",
				Type:        schema.TypeString,
				Computed:    true,
			},

			"hex": {
				Description: "The generated id presented in padded hexadecimal digits. This result will " +
					"always be twice as long as the requested byte length.",
				Type:     schema.TypeString,
				Computed: true,
			},

			"dec": {
				Description: "The generated id presented in non-padded decimal digits.",
				Type:        schema.TypeString,
				Computed:    true,
			},

			"id": {
				Description: "The generated id presented in base64 without additional transformations or prefix.",
				Type:        schema.TypeString,
				Computed:    true,
			},
		},
	}
}

func CreateID(ctx context.Context, d *schema.ResourceData, meta interface{}) diag.Diagnostics {
	var diags diag.Diagnostics
	byteLength := d.Get("byte_length").(int)
	bytes := make([]byte, byteLength)

	n, err := rand.Reader.Read(bytes)
	if n != byteLength {
		return append(diags, diag.Errorf("generated insufficient random bytes: %s", err)...)
	}
	if err != nil {
		return append(diags, diag.Errorf("error generating random bytes: %s", err)...)
	}

	b64Str := base64.RawURLEncoding.EncodeToString(bytes)
	d.SetId(b64Str)

	repopEncsDiags := RepopulateEncodings(ctx, d, meta)
	if repopEncsDiags != nil {
		return append(diags, repopEncsDiags...)
	}

	return diags
}

func RepopulateEncodings(_ context.Context, d *schema.ResourceData, _ interface{}) diag.Diagnostics {
	var diags diag.Diagnostics
	prefix := d.Get("prefix").(string)
	base64Str := d.Id()

	bytes, err := base64.RawURLEncoding.DecodeString(base64Str)
	if err != nil {
		return append(diags, diag.Errorf("error decoding ID: %s", err)...)
	}

	b64StdStr := base64.StdEncoding.EncodeToString(bytes)
	hexStr := hex.EncodeToString(bytes)

	bigInt := big.Int{}
	bigInt.SetBytes(bytes)
	decStr := bigInt.String()

	d.Set("b64_url", prefix+base64Str)
	d.Set("b64_std", prefix+b64StdStr)

	d.Set("hex", prefix+hexStr)
	d.Set("dec", prefix+decStr)

	return nil
}

func ImportID(d *schema.ResourceData, meta interface{}) ([]*schema.ResourceData, error) {
	id := d.Id()

	sep := strings.LastIndex(id, ",")
	if sep != -1 {
		d.Set("prefix", id[:sep])
		id = id[sep+1:]
	}

	bytes, err := base64.RawURLEncoding.DecodeString(id)
	if err != nil {
		return nil, errwrap.Wrapf("Error decoding ID: {{err}}", err)
	}

	d.Set("byte_length", len(bytes))
	d.SetId(id)

	return []*schema.ResourceData{d}, nil
}
<|endoftext|>


<commit_before>package util

type CompletionLatch interface {
	Success()
	Failure()
	Await() bool
}

type completionLatchImpl struct {
	count   int
	channel chan bool
}

func NewCompletionLatch(count int) CompletionLatch {
	return completionLatchImpl{count, make(chan bool, count)}
}

func (latch completionLatchImpl) Success() {
	latch.channel <- true
}

func (latch completionLatchImpl) Failure() {
	latch.channel <- false
}

func (latch completionLatchImpl) Await() bool {
	result := true
	for i := 0; i < latch.count; i++ {
		result = result && <-latch.channel
	}
	return result
}
<commit_msg>Fixed completion latch failure path (#20)<commit_after>package util

type CompletionLatch interface {
	Success()
	Failure()
	Await() bool
}

type completionLatchImpl struct {
	count   int
	channel chan bool
}

func NewCompletionLatch(count int) CompletionLatch {
	return completionLatchImpl{count, make(chan bool, count)}
}

func (latch completionLatchImpl) Success() {
	latch.channel <- true
}

func (latch completionLatchImpl) Failure() {
	latch.channel <- false
}

func (latch completionLatchImpl) Await() bool {
	finalResult := true
	for i := 0; i < latch.count; i++ {
		result := <-latch.channel
		finalResult = finalResult && result
	}
	return finalResult
}
<|endoftext|>


<commit_before>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

// go-to-protobuf generates a Protobuf IDL from a Go struct, respecting any
// existing IDL tags on the Go struct.
package protobuf

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"strings"

	"k8s.io/kubernetes/cmd/libs/go2idl/args"
	"k8s.io/kubernetes/cmd/libs/go2idl/generator"
	"k8s.io/kubernetes/cmd/libs/go2idl/namer"
	"k8s.io/kubernetes/cmd/libs/go2idl/parser"
	"k8s.io/kubernetes/cmd/libs/go2idl/types"

	flag "github.com/spf13/pflag"
)

type Generator struct {
	Common               args.GeneratorArgs
	Packages             string
	OutputBase           string
	ProtoImport          []string
	Conditional          string
	Clean                bool
	OnlyIDL              bool
	KeepGogoproto        bool
	SkipGeneratedRewrite bool
	DropEmbeddedFields   string
}

func New() *Generator {
	sourceTree := args.DefaultSourceTree()
	common := args.GeneratorArgs{
		OutputBase:       sourceTree,
		GoHeaderFilePath: filepath.Join(sourceTree, "k8s.io/kubernetes/hack/boilerplate/boilerplate.go.txt"),
	}
	defaultProtoImport := filepath.Join(sourceTree, "k8s.io", "kubernetes", "Godeps", "_workspace", "src", "github.com", "gogo", "protobuf", "protobuf")
	return &Generator{
		Common:      common,
		OutputBase:  sourceTree,
		ProtoImport: []string{defaultProtoImport},
		Packages: strings.Join([]string{
			`+k8s.io/kubernetes/pkg/util/intstr`,
			`+k8s.io/kubernetes/pkg/api/resource`,
			`+k8s.io/kubernetes/pkg/runtime`,
			`+k8s.io/kubernetes/pkg/watch/versioned`,
			`k8s.io/kubernetes/pkg/api/unversioned`,
			`k8s.io/kubernetes/pkg/api/v1`,
			`k8s.io/kubernetes/pkg/apis/extensions/v1beta1`,
			`k8s.io/kubernetes/pkg/apis/autoscaling/v1`,
			`k8s.io/kubernetes/pkg/apis/batch/v1`,
			`k8s.io/kubernetes/pkg/apis/apps/v1alpha1`,
		}, ","),
		DropEmbeddedFields: "k8s.io/kubernetes/pkg/api/unversioned.TypeMeta",
	}
}

func (g *Generator) BindFlags(flag *flag.FlagSet) {
	flag.StringVarP(&g.Common.GoHeaderFilePath, "go-header-file", "h", g.Common.GoHeaderFilePath, "File containing boilerplate header text. The string YEAR will be replaced with the current 4-digit year.")
	flag.BoolVar(&g.Common.VerifyOnly, "verify-only", g.Common.VerifyOnly, "If true, only verify existing output, do not write anything.")
	flag.StringVarP(&g.Packages, "packages", "p", g.Packages, "comma-separated list of directories to get input types from. Directories prefixed with '-' are not generated, directories prefixed with '+' only create types with explicit IDL instructions.")
	flag.StringVarP(&g.OutputBase, "output-base", "o", g.OutputBase, "Output base; defaults to $GOPATH/src/")
	flag.StringSliceVar(&g.ProtoImport, "proto-import", g.ProtoImport, "The search path for the core protobuf .protos, required, defaults to GODEPS on path.")
	flag.StringVar(&g.Conditional, "conditional", g.Conditional, "An optional Golang build tag condition to add to the generated Go code")
	flag.BoolVar(&g.Clean, "clean", g.Clean, "If true, remove all generated files for the specified Packages.")
	flag.BoolVar(&g.OnlyIDL, "only-idl", g.OnlyIDL, "If true, only generate the IDL for each package.")
	flag.BoolVar(&g.KeepGogoproto, "keep-gogoproto", g.KeepGogoproto, "If true, the generated IDL will contain gogoprotobuf extensions which are normally removed")
	flag.BoolVar(&g.SkipGeneratedRewrite, "skip-generated-rewrite", g.SkipGeneratedRewrite, "If true, skip fixing up the generated.pb.go file (debugging only).")
	flag.StringVar(&g.DropEmbeddedFields, "drop-embedded-fields", g.DropEmbeddedFields, "Comma-delimited list of embedded Go types to omit from generated protobufs")
}

func Run(g *Generator) {
	if g.Common.VerifyOnly {
		g.OnlyIDL = true
		g.Clean = false
	}

	b := parser.New()
	b.AddBuildTags("proto")

	omitTypes := map[types.Name]struct{}{}
	for _, t := range strings.Split(g.DropEmbeddedFields, ",") {
		name := types.Name{}
		if i := strings.LastIndex(t, "."); i != -1 {
			name.Package, name.Name = t[:i], t[i+1:]
		} else {
			name.Name = t
		}
		if len(name.Name) == 0 {
			log.Fatalf("--drop-embedded-types requires names in the form of [GOPACKAGE.]TYPENAME: %v", t)
		}
		omitTypes[name] = struct{}{}
	}

	boilerplate, err := g.Common.LoadGoBoilerplate()
	if err != nil {
		log.Fatalf("Failed loading boilerplate: %v", err)
	}

	protobufNames := NewProtobufNamer()
	outputPackages := generator.Packages{}
	for _, d := range strings.Split(g.Packages, ",") {
		generateAllTypes, outputPackage := true, true
		switch {
		case strings.HasPrefix(d, "+"):
			d = d[1:]
			generateAllTypes = false
		case strings.HasPrefix(d, "-"):
			d = d[1:]
			outputPackage = false
		}
		name := protoSafePackage(d)
		parts := strings.SplitN(d, "=", 2)
		if len(parts) > 1 {
			d = parts[0]
			name = parts[1]
		}
		p := newProtobufPackage(d, name, generateAllTypes, omitTypes)
		header := append([]byte{}, boilerplate...)
		header = append(header, p.HeaderText...)
		p.HeaderText = header
		protobufNames.Add(p)
		if outputPackage {
			outputPackages = append(outputPackages, p)
		}
	}

	if !g.Common.VerifyOnly {
		for _, p := range outputPackages {
			if err := p.(*protobufPackage).Clean(g.OutputBase); err != nil {
				log.Fatalf("Unable to clean package %s: %v", p.Name(), err)
			}
		}
	}

	if g.Clean {
		return
	}

	for _, p := range protobufNames.List() {
		if err := b.AddDir(p.Path()); err != nil {
			log.Fatalf("Unable to add directory %q: %v", p.Path(), err)
		}
	}

	c, err := generator.NewContext(
		b,
		namer.NameSystems{
			"public": namer.NewPublicNamer(3),
			"proto":  protobufNames,
		},
		"public",
	)
	c.Verify = g.Common.VerifyOnly
	c.FileTypes["protoidl"] = NewProtoFile()

	if err != nil {
		log.Fatalf("Failed making a context: %v", err)
	}

	if err := protobufNames.AssignTypesToPackages(c); err != nil {
		log.Fatalf("Failed to identify Common types: %v", err)
	}

	if err := c.ExecutePackages(g.OutputBase, outputPackages); err != nil {
		log.Fatalf("Failed executing generator: %v", err)
	}

	if g.OnlyIDL {
		return
	}

	if _, err := exec.LookPath("protoc"); err != nil {
		log.Fatalf("Unable to find 'protoc': %v", err)
	}

	searchArgs := []string{"-I", ".", "-I", g.OutputBase}
	if len(g.ProtoImport) != 0 {
		for _, s := range g.ProtoImport {
			searchArgs = append(searchArgs, "-I", s)
		}
	}
	args := append(searchArgs, fmt.Sprintf("--gogo_out=%s", g.OutputBase))

	buf := &bytes.Buffer{}
	if len(g.Conditional) > 0 {
		fmt.Fprintf(buf, "// +build %s\n\n", g.Conditional)
	}
	buf.Write(boilerplate)

	for _, outputPackage := range outputPackages {
		p := outputPackage.(*protobufPackage)

		path := filepath.Join(g.OutputBase, p.ImportPath())
		outputPath := filepath.Join(g.OutputBase, p.OutputPath())

		// generate the gogoprotobuf protoc
		cmd := exec.Command("protoc", append(args, path)...)
		out, err := cmd.CombinedOutput()
		if len(out) > 0 {
			log.Printf(string(out))
		}
		if err != nil {
			log.Println(strings.Join(cmd.Args, " "))
			log.Fatalf("Unable to generate protoc on %s: %v", p.PackageName, err)
		}

		if g.SkipGeneratedRewrite {
			continue
		}

		// alter the generated protobuf file to remove the generated types (but leave the serializers) and rewrite the
		// package statement to match the desired package name
		if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, buf.Bytes()); err != nil {
			log.Fatalf("Unable to rewrite generated %s: %v", outputPath, err)
		}

		// sort imports
		cmd = exec.Command("goimports", "-w", outputPath)
		out, err = cmd.CombinedOutput()
		if len(out) > 0 {
			log.Printf(string(out))
		}
		if err != nil {
			log.Println(strings.Join(cmd.Args, " "))
			log.Fatalf("Unable to rewrite imports for %s: %v", p.PackageName, err)
		}

		// format and simplify the generated file
		cmd = exec.Command("gofmt", "-s", "-w", outputPath)
		out, err = cmd.CombinedOutput()
		if len(out) > 0 {
			log.Printf(string(out))
		}
		if err != nil {
			log.Println(strings.Join(cmd.Args, " "))
			log.Fatalf("Unable to apply gofmt for %s: %v", p.PackageName, err)
		}
	}

	if g.SkipGeneratedRewrite {
		return
	}

	if !g.KeepGogoproto {
		// generate, but do so without gogoprotobuf extensions
		for _, outputPackage := range outputPackages {
			p := outputPackage.(*protobufPackage)
			p.OmitGogo = true
		}
		if err := c.ExecutePackages(g.OutputBase, outputPackages); err != nil {
			log.Fatalf("Failed executing generator: %v", err)
		}
	}

	for _, outputPackage := range outputPackages {
		p := outputPackage.(*protobufPackage)

		if len(p.StructTags) == 0 {
			continue
		}

		pattern := filepath.Join(g.OutputBase, p.PackagePath, "*.go")
		files, err := filepath.Glob(pattern)
		if err != nil {
			log.Fatalf("Can't glob pattern %q: %v", pattern, err)
		}

		for _, s := range files {
			if strings.HasSuffix(s, "_test.go") {
				continue
			}
			if err := RewriteTypesWithProtobufStructTags(s, p.StructTags); err != nil {
				log.Fatalf("Unable to rewrite with struct tags %s: %v", s, err)
			}
		}
	}
}
<commit_msg>add protobuf generation<commit_after>/*
Copyright 2015 The Kubernetes Authors All rights reserved.

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

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

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

// go-to-protobuf generates a Protobuf IDL from a Go struct, respecting any
// existing IDL tags on the Go struct.
package protobuf

import (
	"bytes"
	"fmt"
	"log"
	"os/exec"
	"path/filepath"
	"strings"

	"k8s.io/kubernetes/cmd/libs/go2idl/args"
	"k8s.io/kubernetes/cmd/libs/go2idl/generator"
	"k8s.io/kubernetes/cmd/libs/go2idl/namer"
	"k8s.io/kubernetes/cmd/libs/go2idl/parser"
	"k8s.io/kubernetes/cmd/libs/go2idl/types"

	flag "github.com/spf13/pflag"
)

type Generator struct {
	Common               args.GeneratorArgs
	Packages             string
	OutputBase           string
	ProtoImport          []string
	Conditional          string
	Clean                bool
	OnlyIDL              bool
	KeepGogoproto        bool
	SkipGeneratedRewrite bool
	DropEmbeddedFields   string
}

func New() *Generator {
	sourceTree := args.DefaultSourceTree()
	common := args.GeneratorArgs{
		OutputBase:       sourceTree,
		GoHeaderFilePath: filepath.Join(sourceTree, "k8s.io/kubernetes/hack/boilerplate/boilerplate.go.txt"),
	}
	defaultProtoImport := filepath.Join(sourceTree, "k8s.io", "kubernetes", "Godeps", "_workspace", "src", "github.com", "gogo", "protobuf", "protobuf")
	return &Generator{
		Common:      common,
		OutputBase:  sourceTree,
		ProtoImport: []string{defaultProtoImport},
		Packages: strings.Join([]string{
			`+k8s.io/kubernetes/pkg/util/intstr`,
			`+k8s.io/kubernetes/pkg/api/resource`,
			`+k8s.io/kubernetes/pkg/runtime`,
			`+k8s.io/kubernetes/pkg/watch/versioned`,
			`k8s.io/kubernetes/pkg/api/unversioned`,
			`k8s.io/kubernetes/pkg/api/v1`,
			`k8s.io/kubernetes/pkg/apis/extensions/v1beta1`,
			`k8s.io/kubernetes/pkg/apis/autoscaling/v1`,
			`k8s.io/kubernetes/pkg/apis/batch/v1`,
			`k8s.io/kubernetes/pkg/apis/apps/v1alpha1`,
			`k8s.io/kubernetes/federation/apis/federation/v1alpha1`,
		}, ","),
		DropEmbeddedFields: "k8s.io/kubernetes/pkg/api/unversioned.TypeMeta",
	}
}

func (g *Generator) BindFlags(flag *flag.FlagSet) {
	flag.StringVarP(&g.Common.GoHeaderFilePath, "go-header-file", "h", g.Common.GoHeaderFilePath, "File containing boilerplate header text. The string YEAR will be replaced with the current 4-digit year.")
	flag.BoolVar(&g.Common.VerifyOnly, "verify-only", g.Common.VerifyOnly, "If true, only verify existing output, do not write anything.")
	flag.StringVarP(&g.Packages, "packages", "p", g.Packages, "comma-separated list of directories to get input types from. Directories prefixed with '-' are not generated, directories prefixed with '+' only create types with explicit IDL instructions.")
	flag.StringVarP(&g.OutputBase, "output-base", "o", g.OutputBase, "Output base; defaults to $GOPATH/src/")
	flag.StringSliceVar(&g.ProtoImport, "proto-import", g.ProtoImport, "The search path for the core protobuf .protos, required, defaults to GODEPS on path.")
	flag.StringVar(&g.Conditional, "conditional", g.Conditional, "An optional Golang build tag condition to add to the generated Go code")
	flag.BoolVar(&g.Clean, "clean", g.Clean, "If true, remove all generated files for the specified Packages.")
	flag.BoolVar(&g.OnlyIDL, "only-idl", g.OnlyIDL, "If true, only generate the IDL for each package.")
	flag.BoolVar(&g.KeepGogoproto, "keep-gogoproto", g.KeepGogoproto, "If true, the generated IDL will contain gogoprotobuf extensions which are normally removed")
	flag.BoolVar(&g.SkipGeneratedRewrite, "skip-generated-rewrite", g.SkipGeneratedRewrite, "If true, skip fixing up the generated.pb.go file (debugging only).")
	flag.StringVar(&g.DropEmbeddedFields, "drop-embedded-fields", g.DropEmbeddedFields, "Comma-delimited list of embedded Go types to omit from generated protobufs")
}

func Run(g *Generator) {
	if g.Common.VerifyOnly {
		g.OnlyIDL = true
		g.Clean = false
	}

	b := parser.New()
	b.AddBuildTags("proto")

	omitTypes := map[types.Name]struct{}{}
	for _, t := range strings.Split(g.DropEmbeddedFields, ",") {
		name := types.Name{}
		if i := strings.LastIndex(t, "."); i != -1 {
			name.Package, name.Name = t[:i], t[i+1:]
		} else {
			name.Name = t
		}
		if len(name.Name) == 0 {
			log.Fatalf("--drop-embedded-types requires names in the form of [GOPACKAGE.]TYPENAME: %v", t)
		}
		omitTypes[name] = struct{}{}
	}

	boilerplate, err := g.Common.LoadGoBoilerplate()
	if err != nil {
		log.Fatalf("Failed loading boilerplate: %v", err)
	}

	protobufNames := NewProtobufNamer()
	outputPackages := generator.Packages{}
	for _, d := range strings.Split(g.Packages, ",") {
		generateAllTypes, outputPackage := true, true
		switch {
		case strings.HasPrefix(d, "+"):
			d = d[1:]
			generateAllTypes = false
		case strings.HasPrefix(d, "-"):
			d = d[1:]
			outputPackage = false
		}
		name := protoSafePackage(d)
		parts := strings.SplitN(d, "=", 2)
		if len(parts) > 1 {
			d = parts[0]
			name = parts[1]
		}
		p := newProtobufPackage(d, name, generateAllTypes, omitTypes)
		header := append([]byte{}, boilerplate...)
		header = append(header, p.HeaderText...)
		p.HeaderText = header
		protobufNames.Add(p)
		if outputPackage {
			outputPackages = append(outputPackages, p)
		}
	}

	if !g.Common.VerifyOnly {
		for _, p := range outputPackages {
			if err := p.(*protobufPackage).Clean(g.OutputBase); err != nil {
				log.Fatalf("Unable to clean package %s: %v", p.Name(), err)
			}
		}
	}

	if g.Clean {
		return
	}

	for _, p := range protobufNames.List() {
		if err := b.AddDir(p.Path()); err != nil {
			log.Fatalf("Unable to add directory %q: %v", p.Path(), err)
		}
	}

	c, err := generator.NewContext(
		b,
		namer.NameSystems{
			"public": namer.NewPublicNamer(3),
			"proto":  protobufNames,
		},
		"public",
	)
	c.Verify = g.Common.VerifyOnly
	c.FileTypes["protoidl"] = NewProtoFile()

	if err != nil {
		log.Fatalf("Failed making a context: %v", err)
	}

	if err := protobufNames.AssignTypesToPackages(c); err != nil {
		log.Fatalf("Failed to identify Common types: %v", err)
	}

	if err := c.ExecutePackages(g.OutputBase, outputPackages); err != nil {
		log.Fatalf("Failed executing generator: %v", err)
	}

	if g.OnlyIDL {
		return
	}

	if _, err := exec.LookPath("protoc"); err != nil {
		log.Fatalf("Unable to find 'protoc': %v", err)
	}

	searchArgs := []string{"-I", ".", "-I", g.OutputBase}
	if len(g.ProtoImport) != 0 {
		for _, s := range g.ProtoImport {
			searchArgs = append(searchArgs, "-I", s)
		}
	}
	args := append(searchArgs, fmt.Sprintf("--gogo_out=%s", g.OutputBase))

	buf := &bytes.Buffer{}
	if len(g.Conditional) > 0 {
		fmt.Fprintf(buf, "// +build %s\n\n", g.Conditional)
	}
	buf.Write(boilerplate)

	for _, outputPackage := range outputPackages {
		p := outputPackage.(*protobufPackage)

		path := filepath.Join(g.OutputBase, p.ImportPath())
		outputPath := filepath.Join(g.OutputBase, p.OutputPath())

		// generate the gogoprotobuf protoc
		cmd := exec.Command("protoc", append(args, path)...)
		out, err := cmd.CombinedOutput()
		if len(out) > 0 {
			log.Printf(string(out))
		}
		if err != nil {
			log.Println(strings.Join(cmd.Args, " "))
			log.Fatalf("Unable to generate protoc on %s: %v", p.PackageName, err)
		}

		if g.SkipGeneratedRewrite {
			continue
		}

		// alter the generated protobuf file to remove the generated types (but leave the serializers) and rewrite the
		// package statement to match the desired package name
		if err := RewriteGeneratedGogoProtobufFile(outputPath, p.ExtractGeneratedType, buf.Bytes()); err != nil {
			log.Fatalf("Unable to rewrite generated %s: %v", outputPath, err)
		}

		// sort imports
		cmd = exec.Command("goimports", "-w", outputPath)
		out, err = cmd.CombinedOutput()
		if len(out) > 0 {
			log.Printf(string(out))
		}
		if err != nil {
			log.Println(strings.Join(cmd.Args, " "))
			log.Fatalf("Unable to rewrite imports for %s: %v", p.PackageName, err)
		}

		// format and simplify the generated file
		cmd = exec.Command("gofmt", "-s", "-w", outputPath)
		out, err = cmd.CombinedOutput()
		if len(out) > 0 {
			log.Printf(string(out))
		}
		if err != nil {
			log.Println(strings.Join(cmd.Args, " "))
			log.Fatalf("Unable to apply gofmt for %s: %v", p.PackageName, err)
		}
	}

	if g.SkipGeneratedRewrite {
		return
	}

	if !g.KeepGogoproto {
		// generate, but do so without gogoprotobuf extensions
		for _, outputPackage := range outputPackages {
			p := outputPackage.(*protobufPackage)
			p.OmitGogo = true
		}
		if err := c.ExecutePackages(g.OutputBase, outputPackages); err != nil {
			log.Fatalf("Failed executing generator: %v", err)
		}
	}

	for _, outputPackage := range outputPackages {
		p := outputPackage.(*protobufPackage)

		if len(p.StructTags) == 0 {
			continue
		}

		pattern := filepath.Join(g.OutputBase, p.PackagePath, "*.go")
		files, err := filepath.Glob(pattern)
		if err != nil {
			log.Fatalf("Can't glob pattern %q: %v", pattern, err)
		}

		for _, s := range files {
			if strings.HasSuffix(s, "_test.go") {
				continue
			}
			if err := RewriteTypesWithProtobufStructTags(s, p.StructTags); err != nil {
				log.Fatalf("Unable to rewrite with struct tags %s: %v", s, err)
			}
		}
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"koding/databases/redis"
	"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)

	// routeMap holds the subscription list/set for any given routing key
	routeMap = make(map[string]*set.Set)

	// sessionsMap holds sessions with their socketIds
	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, premiumBroker and brokerKite, premiumBrokerKite. 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
	// we should open only one connection session to Redis for one broker
	RedisSingleton *redis.SingletonSession

	// 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(conf *config.Config) *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{}),
		RedisSingleton:    redis.Singleton(conf),
	}
}

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

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

	switch *flagBrokerType {
	case "premiumBroker":
		broker.Config = &conf.PremiumBroker
	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)

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

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

	// start listening/serving socket server
	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.ServiceGenericName, // 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.ServiceGenericName, // 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.ServiceGenericName, "#", 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 {
			sendMessageToClient(amqpMessage)
		}

		b.Close()

	}(stream)

	return nil
}

// sendMessageToClient takes an amqp messsage and delivers it to the related
// clients which are subscribed to the routing key
func sendMessageToClient(amqpMessage amqp.Delivery) {
	routingKey := amqpMessage.RoutingKey
	payloadsByte := utils.FilterInvalidUTF8(amqpMessage.Body)

	// We are sending multiple bodies for updateInstances exchange
	// so that there will be another operations, if exchange is not "updateInstances"
	// no need to add more overhead
	if amqpMessage.Exchange != "updateInstances" {
		payloadRaw := json.RawMessage(payloadsByte)
		processMessage(routingKey, &payloadRaw)
		return
	}

	// this part is only for updateInstances exchange
	var payloads []interface{}
	// unmarshal data to slice of interface
	if err := json.Unmarshal(payloadsByte, &payloads); err != nil {
		log.Error("Error while unmarshalling:%v data:%v routingKey:%v", err, string(payloadsByte), routingKey)
		return
	}

	// range over the slice and send all of them to the same routingkey
	for _, payload := range payloads {
		payloadByte, err := json.Marshal(payload)
		if err != nil {
			log.Error("Error while marshalling:%v data:%v routingKey:%v", err, string(payloadByte), routingKey)
			continue
		}
		payloadByteRaw := json.RawMessage(payloadByte)
		processMessage(routingKey, &payloadByteRaw)
	}
}

// processMessage gets routingKey and a payload for sending them to the client
// Gets subscription bindings from global routeMap
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)
	}
}
<commit_msg>Broker: close connection  if client is not created within 5 seconds<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"flag"
	"fmt"
	"koding/databases/redis"
	"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)

	// routeMap holds the subscription list/set for any given routing key
	routeMap = make(map[string]*set.Set)

	// sessionsMap holds sessions with their socketIds
	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, premiumBroker and brokerKite, premiumBrokerKite. 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
	// we should open only one connection session to Redis for one broker
	RedisSingleton *redis.SingletonSession

	// 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(conf *config.Config) *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{}),
		RedisSingleton:    redis.Singleton(conf),
	}
}

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

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

	switch *flagBrokerType {
	case "premiumBroker":
		broker.Config = &conf.PremiumBroker
	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)

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

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

	// start listening/serving socket server
	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.ServiceGenericName, // 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.ServiceGenericName, // 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.ServiceGenericName, "#", 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 {
			sendMessageToClient(amqpMessage)
		}

		b.Close()

	}(stream)

	return nil
}

// sendMessageToClient takes an amqp messsage and delivers it to the related
// clients which are subscribed to the routing key
func sendMessageToClient(amqpMessage amqp.Delivery) {
	routingKey := amqpMessage.RoutingKey
	payloadsByte := utils.FilterInvalidUTF8(amqpMessage.Body)

	// We are sending multiple bodies for updateInstances exchange
	// so that there will be another operations, if exchange is not "updateInstances"
	// no need to add more overhead
	if amqpMessage.Exchange != "updateInstances" {
		payloadRaw := json.RawMessage(payloadsByte)
		processMessage(routingKey, &payloadRaw)
		return
	}

	// this part is only for updateInstances exchange
	var payloads []interface{}
	// unmarshal data to slice of interface
	if err := json.Unmarshal(payloadsByte, &payloads); err != nil {
		log.Error("Error while unmarshalling:%v data:%v routingKey:%v", err, string(payloadsByte), routingKey)
		return
	}

	// range over the slice and send all of them to the same routingkey
	for _, payload := range payloads {
		payloadByte, err := json.Marshal(payload)
		if err != nil {
			log.Error("Error while marshalling:%v data:%v routingKey:%v", err, string(payloadByte), routingKey)
			continue
		}
		payloadByteRaw := json.RawMessage(payloadByte)
		processMessage(routingKey, &payloadByteRaw)
	}
}

// processMessage gets routingKey and a payload for sending them to the client
// Gets subscription bindings from global routeMap
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) {
	clientChan := make(chan *Client, 0)
	errChan := make(chan error, 0)

	// we don't use time.After because we are stopping the time if the given
	// functions are fast enough.
	timer := time.NewTicker(5 * time.Second)

	go createClient(b, session, clientChan, errChan)

	var client *Client
	var ok bool
	select {
	case client, ok = <-clientChan:
		if ok {
			timer.Stop()
		}
	case err, ok := <-errChan:
		if ok {
			timer.Stop()
			log.Critical("An error occured while creating client %v", err)
			return
		}
	case <-timer.C:
		timer.Stop()
		log.Critical("Client coulnt created in %s exiting ", "5s")
		return
	}

	sessionGaugeEnd := client.gaugeStart()

	defer sessionGaugeEnd()
	defer client.Close()

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

		client.handleSessionMessage(data)
	}
}

func createClient(b *Broker, session *sockjs.Session, clientChan chan *Client, errChan chan error) {
	// do not forget to close channels
	defer close(errChan)
	defer close(clientChan)

	client, err := NewClient(session, b)
	if err != nil {
		log.Critical("Couldnt create client %v", err)
		errChan <- err
		return
	}

	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)
		errChan <- err
		return
	}

	// if session is closed before the client creation no need to send
	// client object to listeners
	if !session.Closed {
		sendToClient(session, "broker.connected", client.SocketId)
		clientChan <- client
	}
}

// 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 tunnel

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Server satisfies the http.Handler interface. It is responsible of tracking
// tunnels and creating tunnels between remote and local connection.
type Server struct {
	tunnels    map[string]*Tunnels
	tunnelChan map[string]chan *Tunnels
	controls   *Controls
	hosts      *Hosts
	sync.Mutex
}

func NewServer() *Server {
	s := &Server{
		tunnels:    make(map[string]*Tunnels),
		tunnelChan: make(map[string]chan *Tunnels),
		controls:   NewControls(),
		hosts:      newHosts(),
	}

	http.HandleFunc(ControlPath, s.controlHandler)
	http.HandleFunc(TunnelPath, s.tunnelHandler)
	return s
}

func (s *Server) tunnelHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("tunnel Handler invoked", r.URL.String())
	if r.Method != "CONNECT" {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusMethodNotAllowed)
		io.WriteString(w, "405 must CONNECT\n")
		return
	}

	conn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Println("register hijacking ", r.RemoteAddr, ": ", err.Error())
		return
	}

	io.WriteString(conn, "HTTP/1.1 "+Connected+"\n\n")
	conn.SetDeadline(time.Time{})

	// set by control channel
	protocol := r.Header.Get("protocol")
	tunnelID := r.Header.Get("tunnelID")
	log.Println("protocol and tunnelID is", protocol, tunnelID)

	done, ok := s.tunnelChan[tunnelID]
	if !ok {
		log.Println("tunnelID channel does not exist")
		return
	}

	host := strings.ToLower(r.Host)

	tunnels, ok := s.GetTunnels(host)
	if !ok {
		// first time, create it
		tunnels = NewTunnels()
	}

	tunnels.addTunnel(protocol, NewTunnel(conn))
	s.AddTunnels(host, tunnels)

	// let any channel associated with this tunnel let know that we passed everything and
	// that our tunnel is now ready, also pass it back
	done <- tunnels
}

// controlHandler is used to register tunnel clients.
func (s *Server) controlHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("control Handler invoked", r.URL.String())
	if r.Method != "CONNECT" {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusMethodNotAllowed)
		io.WriteString(w, "405 must CONNECT\n")
		return
	}

	username := r.Header.Get("username")
	if username == "" {
		log.Printf("empty username is connected")
		http.Error(w, "username is not set", 405)
		return
	}

	_, ok := s.hosts.getHost(username)
	if !ok {
		log.Printf("no host is associated for username %s. please use server.AddHost()",
			username)
		http.Error(w, "there is no host defined for this username", 405)
		return
	}

	_, ok = s.GetControl(username)
	if ok {
		log.Printf("control conn for %s already exist\n", username)
		http.Error(w, "only one control connection is allowed", 405)
		return
	}

	conn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Println("register hijacking ", r.RemoteAddr, ": ", err.Error())
		return
	}

	// ok, everthing is ready
	io.WriteString(conn, "HTTP/1.1 "+Connected+"\n\n")
	conn.SetDeadline(time.Time{})

	// create a new control struct and add it
	control := NewControl(conn)
	s.AddControl(username, control)

	// delete and close the conn when the control connection is being closed
	defer func() {
		log.Println("closing control for", username)
		control.Close()
		s.DeleteControl(username)
	}()

	log.Println("control has started for", username)
	// blocking function
	control.run()
}

// ServeHTTP is a tunnel that creates an http/websocket tunnel between a
// public connection and the client connection.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if isWebsocket(r) {
		s.websocketHandleFunc(w, r)
		return
	}

	log.Println("http handler invoked", r.Host, r.URL.String())
	host := strings.ToLower(r.Host)

	tunnels, ok := s.GetTunnels(host)
	if !ok {
		var err error
		tunnels, err = s.requestTunnels("http", host)
		if err != nil {
			http.Error(w, err.Error(), 404)
			return
		}
	}

	// for http one single conn is enough
	tunnel, ok := tunnels.getTunnel("http")
	if !ok {
		err := fmt.Sprintf("no http tunnel: %s", host)
		http.Error(w, err, 404)
		return
	}

	s.Lock()
	defer s.Unlock()

	err := r.Write(tunnel.conn)
	if err != nil {
		err := fmt.Sprintf("write to tunnel %s", err)
		http.Error(w, err, 404)
		return
	}

	resp, err := http.ReadResponse(bufio.NewReader(tunnel.conn), r)
	if err != nil {
		errString := fmt.Sprintf("read from tunnel.con %s", err.Error())
		http.Error(w, errString, 404)
		return
	}

	defer resp.Body.Close()

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)

	io.Copy(w, resp.Body)
}

// requestTunnels makes a request to the control connection
func (s *Server) requestTunnels(protocol, host string) (*Tunnels, error) {
	fmt.Println("no tunnel getting control")
	// get the user associated with this user
	username, ok := s.GetUsername(host)
	if !ok {
		return nil, fmt.Errorf("no control available for %s", host)
	}

	// then grab the control connection that is associated with this username
	control, ok := s.GetControl(username)
	if !ok {
		return nil, fmt.Errorf("no control available for %s", host)
	}
	fmt.Println("got control preparing request")

	// create an unique id to used with that tunnel
	tunnelID := randomID(32)

	// request a new http tunnel
	control.SendMsg(protocol, tunnelID)

	// now wait until our tunnel is established if the tunnel has the
	// right ID it will send a message to this channel, which releases the
	// blocking channel. then we wait, until we got our done channel.
	fmt.Println("waiting for tunnel")
	s.tunnelChan[tunnelID] = make(chan *Tunnels)

	select {
	case tunnels := <-s.tunnelChan[tunnelID]:
		// remove and close it because we don't need it anymore
		close(s.tunnelChan[tunnelID])
		delete(s.tunnelChan, tunnelID)

		return tunnels, nil
	case <-time.After(time.Second * 10):
		return nil, errors.New("timeout")
	}
}

func (s *Server) websocketHandleFunc(w http.ResponseWriter, r *http.Request) {
	log.Println("websocket handler invoked", r.URL.String())

	host := strings.ToLower(r.Host)
	tunnels, ok := s.GetTunnels(host)
	if !ok {
		err := fmt.Sprintf("no such tunnel: %s", host)
		http.Error(w, err, 404)
		return
	}

	tunnel, ok := tunnels.getTunnel("websocket")
	if !ok {
		err := fmt.Sprintf("no such tunnel: %s", host)
		http.Error(w, err, 404)
		return
	}

	err := r.Write(tunnel.conn)
	if err != nil {
		err := fmt.Sprintf("write to tunnel %s", err)
		http.Error(w, err, 404)
		return
	}

	publicConn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Println("websocket hijacking ", err)
		return
	}
	defer publicConn.Close()

	err = <-join(tunnel.conn, publicConn)
	log.Println(err)
}

func (s *Server) GetTunnels(host string) (*Tunnels, bool) {
	s.Lock()
	defer s.Unlock()

	tunnels, ok := s.tunnels[host]
	return tunnels, ok
}

func (s *Server) AddTunnels(host string, tunnels *Tunnels) {
	s.Lock()
	defer s.Unlock()

	s.tunnels[host] = tunnels
}

func (s *Server) AddControl(username string, conn *Control) {
	s.controls.addControl(username, conn)
}

func (s *Server) GetControl(username string) (*Control, bool) {
	return s.controls.getControl(username)
}

func (s *Server) DeleteControl(username string) {
	s.controls.deleteControl(username)
}

func (s *Server) AddHost(host, username string) {
	s.hosts.addHost(host, username)
}

func (s *Server) DeleteHost(host, username string) {
	s.hosts.deleteHost(host)
}

func (s *Server) GetUsername(host string) (string, bool) {
	username, ok := s.hosts.getUsername(host)
	return username, ok
}
<commit_msg>tunnel: lock http read/request handling, just like in standart lib<commit_after>package tunnel

import (
	"bufio"
	"errors"
	"fmt"
	"io"
	"log"
	"net/http"
	"strings"
	"sync"
	"time"
)

// Server satisfies the http.Handler interface. It is responsible of tracking
// tunnels and creating tunnels between remote and local connection.
type Server struct {
	tunnels    map[string]*Tunnels
	tunnelChan map[string]chan *Tunnels
	controls   *Controls
	hosts      *Hosts
	sync.Mutex
}

func NewServer() *Server {
	s := &Server{
		tunnels:    make(map[string]*Tunnels),
		tunnelChan: make(map[string]chan *Tunnels),
		controls:   NewControls(),
		hosts:      newHosts(),
	}

	http.HandleFunc(ControlPath, s.controlHandler)
	http.HandleFunc(TunnelPath, s.tunnelHandler)
	return s
}

func (s *Server) tunnelHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("tunnel Handler invoked", r.URL.String())
	if r.Method != "CONNECT" {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusMethodNotAllowed)
		io.WriteString(w, "405 must CONNECT\n")
		return
	}

	conn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Println("register hijacking ", r.RemoteAddr, ": ", err.Error())
		return
	}

	io.WriteString(conn, "HTTP/1.1 "+Connected+"\n\n")
	conn.SetDeadline(time.Time{})

	// set by control channel
	protocol := r.Header.Get("protocol")
	tunnelID := r.Header.Get("tunnelID")
	log.Println("protocol and tunnelID is", protocol, tunnelID)

	done, ok := s.tunnelChan[tunnelID]
	if !ok {
		log.Println("tunnelID channel does not exist")
		return
	}

	host := strings.ToLower(r.Host)

	tunnels, ok := s.GetTunnels(host)
	if !ok {
		// first time, create it
		tunnels = NewTunnels()
	}

	tunnels.addTunnel(protocol, NewTunnel(conn))
	s.AddTunnels(host, tunnels)

	// let any channel associated with this tunnel let know that we passed everything and
	// that our tunnel is now ready, also pass it back
	done <- tunnels
}

// controlHandler is used to register tunnel clients.
func (s *Server) controlHandler(w http.ResponseWriter, r *http.Request) {
	fmt.Println("control Handler invoked", r.URL.String())
	if r.Method != "CONNECT" {
		w.Header().Set("Content-Type", "text/plain; charset=utf-8")
		w.WriteHeader(http.StatusMethodNotAllowed)
		io.WriteString(w, "405 must CONNECT\n")
		return
	}

	username := r.Header.Get("username")
	if username == "" {
		log.Printf("empty username is connected")
		http.Error(w, "username is not set", 405)
		return
	}

	_, ok := s.hosts.getHost(username)
	if !ok {
		log.Printf("no host is associated for username %s. please use server.AddHost()",
			username)
		http.Error(w, "there is no host defined for this username", 405)
		return
	}

	_, ok = s.GetControl(username)
	if ok {
		log.Printf("control conn for %s already exist\n", username)
		http.Error(w, "only one control connection is allowed", 405)
		return
	}

	conn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Println("register hijacking ", r.RemoteAddr, ": ", err.Error())
		return
	}

	// ok, everthing is ready
	io.WriteString(conn, "HTTP/1.1 "+Connected+"\n\n")
	conn.SetDeadline(time.Time{})

	// create a new control struct and add it
	control := NewControl(conn)
	s.AddControl(username, control)

	// delete and close the conn when the control connection is being closed
	defer func() {
		log.Println("closing control for", username)
		control.Close()
		s.DeleteControl(username)
	}()

	log.Println("control has started for", username)
	// blocking function
	control.run()
}

// ServeHTTP is a tunnel that creates an http/websocket tunnel between a
// public connection and the client connection.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if isWebsocket(r) {
		s.websocketHandleFunc(w, r)
		return
	}

	log.Println("http handler invoked", r.Host, r.URL.String())
	host := strings.ToLower(r.Host)

	tunnels, ok := s.GetTunnels(host)
	if !ok {
		var err error
		tunnels, err = s.requestTunnels("http", host)
		if err != nil {
			http.Error(w, err.Error(), 404)
			return
		}
	}

	// for http one single conn is enough
	tunnel, ok := tunnels.getTunnel("http")
	if !ok {
		err := fmt.Sprintf("no http tunnel: %s", host)
		http.Error(w, err, 404)
		return
	}

	err := r.Write(tunnel.conn)
	if err != nil {
		err := fmt.Sprintf("write to tunnel %s", err)
		http.Error(w, err, 404)
		return
	}

	resp, err := http.ReadResponse(bufio.NewReader(tunnel.conn), r)
	if err != nil {
		errString := fmt.Sprintf("read from tunnel.con %s", err.Error())
		http.Error(w, errString, 404)
		return
	}

	s.Lock()
	defer s.Unlock()

	defer resp.Body.Close()

	copyHeader(w.Header(), resp.Header)
	w.WriteHeader(resp.StatusCode)

	io.Copy(w, resp.Body)
}

// requestTunnels makes a request to the control connection to get a new
// tunnel from the client. It sends the request of a new tunnel directly to
// the client, which then opens a new tunnel to be used.
func (s *Server) requestTunnels(protocol, host string) (*Tunnels, error) {
	fmt.Println("no tunnel getting control")
	// get the user associated with this user
	username, ok := s.GetUsername(host)
	if !ok {
		return nil, fmt.Errorf("no control available for %s", host)
	}

	// then grab the control connection that is associated with this username
	control, ok := s.GetControl(username)
	if !ok {
		return nil, fmt.Errorf("no control available for %s", host)
	}
	fmt.Println("got control preparing request")

	// create an unique id to used with that tunnel
	tunnelID := randomID(32)

	// request a new http tunnel
	control.SendMsg(protocol, tunnelID)

	// now wait until our tunnel is established if the tunnel has the
	// right ID it will send a message to this channel, which releases the
	// blocking channel. then we wait, until we got our done channel.
	fmt.Println("waiting for tunnel")
	s.tunnelChan[tunnelID] = make(chan *Tunnels)

	select {
	case tunnels := <-s.tunnelChan[tunnelID]:
		// remove and close it because we don't need it anymore
		close(s.tunnelChan[tunnelID])
		delete(s.tunnelChan, tunnelID)

		return tunnels, nil
	case <-time.After(time.Second * 10):
		return nil, errors.New("timeout")
	}
}

func (s *Server) websocketHandleFunc(w http.ResponseWriter, r *http.Request) {
	log.Println("websocket handler invoked", r.URL.String())

	host := strings.ToLower(r.Host)
	tunnels, ok := s.GetTunnels(host)
	if !ok {
		err := fmt.Sprintf("no such tunnel: %s", host)
		http.Error(w, err, 404)
		return
	}

	tunnel, ok := tunnels.getTunnel("websocket")
	if !ok {
		err := fmt.Sprintf("no such tunnel: %s", host)
		http.Error(w, err, 404)
		return
	}

	err := r.Write(tunnel.conn)
	if err != nil {
		err := fmt.Sprintf("write to tunnel %s", err)
		http.Error(w, err, 404)
		return
	}

	publicConn, _, err := w.(http.Hijacker).Hijack()
	if err != nil {
		log.Println("websocket hijacking ", err)
		return
	}
	defer publicConn.Close()

	err = <-join(tunnel.conn, publicConn)
	log.Println(err)
}

func (s *Server) GetTunnels(host string) (*Tunnels, bool) {
	s.Lock()
	defer s.Unlock()

	tunnels, ok := s.tunnels[host]
	return tunnels, ok
}

func (s *Server) AddTunnels(host string, tunnels *Tunnels) {
	s.Lock()
	defer s.Unlock()

	s.tunnels[host] = tunnels
}

func (s *Server) AddControl(username string, conn *Control) {
	s.controls.addControl(username, conn)
}

func (s *Server) GetControl(username string) (*Control, bool) {
	return s.controls.getControl(username)
}

func (s *Server) DeleteControl(username string) {
	s.controls.deleteControl(username)
}

func (s *Server) AddHost(host, username string) {
	s.hosts.addHost(host, username)
}

func (s *Server) DeleteHost(host, username string) {
	s.hosts.deleteHost(host)
}

func (s *Server) GetUsername(host string) (string, bool) {
	username, ok := s.hosts.getUsername(host)
	return username, ok
}
<|endoftext|>


<commit_before>package models

import "time"

type Tag struct {
	Id        int64
	Slug      string
	CreatedAt time.Time
}
<commit_msg>Social: remove unused struct<commit_after><|endoftext|>


<commit_before>package rpc

import (
	"golang.org/x/net/context"
)

// CtxRpcKey is a type defining the context key for the RPC context
type CtxRpcKey int

const (
	// CtxRpcTagsKey defines a context key that can hold a slice of context keys
	CtxRpcTagsKey CtxRpcKey = iota
)

type CtxRpcTags map[string]interface{}

// AddRpcTagsToContext adds the given log tag mappings (logTagsToAdd) to the
// given context, creating a new one if necessary. Returns the resulting
// context with the new log tag mappings.
func AddRpcTagsToContext(ctx context.Context, logTagsToAdd CtxRpcTags) context.Context {
	currTags, ok := RpcTagsFromContext(ctx)
	if !ok {
		currTags = make(CtxRpcTags)
		ctx = context.WithValue(ctx, CtxRpcTagsKey, currTags)
	}
	for key, tag := range logTagsToAdd {
		currTags[key] = tag
	}
	return ctx
}

// RpcTagsFromContext returns the tags being passed along with the given context.
func RpcTagsFromContext(ctx context.Context) (CtxRpcTags, bool) {
	logTags, ok := ctx.Value(CtxRpcTagsKey).(CtxRpcTags)
	return logTags, ok
}
<commit_msg>copy the RPC tags before returning to avoid map concurrency errors<commit_after>package rpc

import "golang.org/x/net/context"

// CtxRpcKey is a type defining the context key for the RPC context
type CtxRpcKey int

const (
	// CtxRpcTagsKey defines a context key that can hold a slice of context keys
	CtxRpcTagsKey CtxRpcKey = iota
)

type CtxRpcTags map[string]interface{}

// AddRpcTagsToContext adds the given log tag mappings (logTagsToAdd) to the
// given context, creating a new one if necessary. Returns the resulting
// context with the new log tag mappings.
func AddRpcTagsToContext(ctx context.Context, logTagsToAdd CtxRpcTags) context.Context {
	currTags, ok := RpcTagsFromContext(ctx)
	if !ok {
		currTags = make(CtxRpcTags)
	}
	for key, tag := range logTagsToAdd {
		currTags[key] = tag
	}

	return context.WithValue(ctx, CtxRpcTagsKey, currTags)
}

// RpcTagsFromContext returns the tags being passed along with the given context.
func RpcTagsFromContext(ctx context.Context) (CtxRpcTags, bool) {
	logTags, ok := ctx.Value(CtxRpcTagsKey).(CtxRpcTags)
	if ok {
		ret := make(CtxRpcTags)
		for k, v := range logTags {
			ret[k] = v
		}
		return ret, true
	}
	return nil, false
}
<|endoftext|>