<commit_before>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package checkmgr

import (
	"reflect"
	"testing"

	"github.com/circonus-labs/circonus-gometrics/api"
)

func TestIsMetricActive(t *testing.T) {

	cm := &CheckManager{}

	cm.availableMetrics = map[string]bool{
		"foo": true,
	}

	t.Log("'foo' in active metric list")
	{
		if !cm.IsMetricActive("foo") {
			t.Error("Expected true")
		}
	}

	t.Log("'bar' not in active metric list")
	{
		if cm.IsMetricActive("bar") {
			t.Error("Expected false")
		}
	}
}

func TestActivateMetric(t *testing.T) {
	cm := &CheckManager{}
	cm.checkBundle = &api.CheckBundle{}
	cm.checkBundle.Metrics = []api.CheckBundleMetric{
		api.CheckBundleMetric{
			Name:   "foo",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "bar",
			Type:   "numeric",
			Status: "available",
		},
	}

	cm.availableMetrics = make(map[string]bool)
	cm.forceMetricActivation = false

	cm.inventoryMetrics()

	t.Log("'foo' already active")
	{
		if cm.ActivateMetric("foo") {
			t.Error("Expected false")
		}
	}

	t.Log("'bar' in list but not active [force=false]")
	{
		if cm.ActivateMetric("bar") {
			t.Error("Expected false")
		}
	}

	t.Log("'baz' not in list")
	{
		if !cm.ActivateMetric("baz") {
			t.Error("Expected true")
		}
	}

	cm.forceMetricActivation = true

	t.Log("'bar' in list but not active [force=true]")
	{
		if !cm.ActivateMetric("bar") {
			t.Error("Expected true")
		}
	}
}

func TestInventoryMetrics(t *testing.T) {
	cm := &CheckManager{}
	cm.checkBundle = &api.CheckBundle{}
	cm.checkBundle.Metrics = []api.CheckBundleMetric{
		api.CheckBundleMetric{
			Name:   "foo",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "bar",
			Type:   "numeric",
			Status: "available",
		},
	}

	cm.availableMetrics = make(map[string]bool)
	cm.inventoryMetrics()

	expectedMetrics := make(map[string]bool)
	expectedMetrics["foo"] = true
	expectedMetrics["bar"] = false

	t.Log("'foo', in inventory and active")
	{
		active, exists := cm.availableMetrics["foo"]
		if !active {
			t.Fatalf("Expected active")
		}
		if !exists {
			t.Fatalf("Expected exists")
		}
	}

	t.Log("'bar', in inventory and not active")
	{
		active, exists := cm.availableMetrics["bar"]
		if active {
			t.Fatalf("Expected not active")
		}
		if !exists {
			t.Fatalf("Expected exists")
		}
	}

	t.Log("'baz', not in inventory and not active")
	{
		active, exists := cm.availableMetrics["baz"]
		if active {
			t.Fatalf("Expected not active")
		}
		if exists {
			t.Fatalf("Expected not exists")
		}
	}
}

func TestAddMetricTags(t *testing.T) {
	cm := &CheckManager{}
	cm.checkBundle = &api.CheckBundle{}
	cm.metricTags = make(map[string][]string)

	t.Log("no tags")
	{
		if cm.AddMetricTags("foo", []string{}, false) {
			t.Fatalf("Expected false")
		}
	}

	t.Log("no metric named 'foo'")
	{
		if !cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected true")
		}
	}

	cm.checkBundle.Metrics = []api.CheckBundleMetric{
		api.CheckBundleMetric{
			Name:   "bar",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "foo",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "baz",
			Type:   "numeric",
			Status: "active",
			Tags:   []string{"cat1:tag1"},
		},
	}

	t.Log("metric named 'bar', add tag")
	{

		cm.metricTags = make(map[string][]string)
		// append, zero current
		if !cm.AddMetricTags("bar", []string{"cat:tag"}, true) {
			t.Fatalf("Expected true")
		}
		expected := make(map[string][]string)
		expected["bar"] = []string{"cat:tag"}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// tag already exists, no need to add
		if cm.AddMetricTags("bar", []string{"cat:tag"}, true) {
			t.Fatalf("Expected false")
		}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// append, zero tags
		if cm.AddMetricTags("bar", []string{}, true) {
			t.Fatalf("Expected false")
		}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}
	}

	t.Log("metric named 'baz', add tag")
	{
		cm.metricTags = make(map[string][]string)

		// append, current tag, should be noupdate
		if cm.AddMetricTags("baz", []string{"cat1:tag1"}, true) {
			t.Fatalf("Expected false")
		}

		if _, found := cm.metricTags["baz"]; found {
			t.Fatalf("expected not found")
		}

		// append, one current
		if !cm.AddMetricTags("baz", []string{"cat2:tag2"}, true) {
			t.Fatalf("Expected true")
		}
		expected := make(map[string][]string)
		expected["baz"] = []string{"cat1:tag1", "cat2:tag2"}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// append, tag already exists (should be noupdate)
		if cm.AddMetricTags("baz", []string{"cat2:tag2"}, true) {
			t.Fatalf("Expected false")
		}
	}

	t.Log("metric named 'foo', set tag")
	{
		expected := make(map[string][]string)
		cm.metricTags = make(map[string][]string)

		// set tag
		if !cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected true")
		}
		expected["foo"] = []string{"cat:tag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// set, reset (pass 0 tags)
		if !cm.AddMetricTags("foo", []string{}, false) {
			t.Fatalf("Expected true")
		}

		if _, found := cm.metricTags["foo"]; !found {
			t.Fatal("expected found")
		}

		expected["foo"] = []string{}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// set tag
		if !cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected true")
		}

		expected["foo"] = []string{"cat:tag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// set, no update (same tag)
		if cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected false")
		}

		expected["foo"] = []string{"cat:tag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// replace any existing
		if !cm.AddMetricTags("foo", []string{"cat:newtag"}, false) {
			t.Fatalf("Expected true")
		}

		expected["foo"] = []string{"cat:newtag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

	}
}

func TestAddNewMetrics(t *testing.T) {
	cm := &CheckManager{}

	newMetrics := make(map[string]*api.CheckBundleMetric)

	newMetrics["foo"] = &api.CheckBundleMetric{
		Name:   "foo",
		Type:   "numeric",
		Status: "active",
	}

	t.Log("no check bundle")
	{
		if cm.addNewMetrics(newMetrics) {
			t.Fatalf("Expected false")
		}
	}

	cm.checkBundle = &api.CheckBundle{}
	t.Log("no check bundle metrics")
	{
		if !cm.addNewMetrics(newMetrics) {
			t.Fatalf("Expected true")
		}
		if !cm.forceCheckUpdate {
			t.Fatal("Expected forceCheckUpdate to be true")
		}
	}
}
<commit_msg>add: additional isMetricActive test<commit_after>// Copyright 2016 Circonus, Inc. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package checkmgr

import (
	"reflect"
	"testing"

	"github.com/circonus-labs/circonus-gometrics/api"
)

func TestIsMetricActive(t *testing.T) {

	cm := &CheckManager{}

	cm.availableMetrics = map[string]bool{
		"foo": true,
		"baz": false,
	}

	t.Log("'foo' in active metric list")
	{
		if !cm.IsMetricActive("foo") {
			t.Error("Expected true")
		}
	}

	t.Log("'bar' not in active metric list")
	{
		if cm.IsMetricActive("bar") {
			t.Error("Expected false")
		}
	}

	t.Log("'baz' in active metric list, not active")
	{
		if cm.IsMetricActive("baz") {
			t.Error("Expected false")
		}
	}
}

func TestActivateMetric(t *testing.T) {
	cm := &CheckManager{}
	cm.checkBundle = &api.CheckBundle{}
	cm.checkBundle.Metrics = []api.CheckBundleMetric{
		api.CheckBundleMetric{
			Name:   "foo",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "bar",
			Type:   "numeric",
			Status: "available",
		},
	}

	cm.availableMetrics = make(map[string]bool)
	cm.forceMetricActivation = false

	cm.inventoryMetrics()

	t.Log("'foo' already active")
	{
		if cm.ActivateMetric("foo") {
			t.Error("Expected false")
		}
	}

	t.Log("'bar' in list but not active [force=false]")
	{
		if cm.ActivateMetric("bar") {
			t.Error("Expected false")
		}
	}

	t.Log("'baz' not in list")
	{
		if !cm.ActivateMetric("baz") {
			t.Error("Expected true")
		}
	}

	cm.forceMetricActivation = true

	t.Log("'bar' in list but not active [force=true]")
	{
		if !cm.ActivateMetric("bar") {
			t.Error("Expected true")
		}
	}
}

func TestInventoryMetrics(t *testing.T) {
	cm := &CheckManager{}
	cm.checkBundle = &api.CheckBundle{}
	cm.checkBundle.Metrics = []api.CheckBundleMetric{
		api.CheckBundleMetric{
			Name:   "foo",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "bar",
			Type:   "numeric",
			Status: "available",
		},
	}

	cm.availableMetrics = make(map[string]bool)
	cm.inventoryMetrics()

	expectedMetrics := make(map[string]bool)
	expectedMetrics["foo"] = true
	expectedMetrics["bar"] = false

	t.Log("'foo', in inventory and active")
	{
		active, exists := cm.availableMetrics["foo"]
		if !active {
			t.Fatalf("Expected active")
		}
		if !exists {
			t.Fatalf("Expected exists")
		}
	}

	t.Log("'bar', in inventory and not active")
	{
		active, exists := cm.availableMetrics["bar"]
		if active {
			t.Fatalf("Expected not active")
		}
		if !exists {
			t.Fatalf("Expected exists")
		}
	}

	t.Log("'baz', not in inventory and not active")
	{
		active, exists := cm.availableMetrics["baz"]
		if active {
			t.Fatalf("Expected not active")
		}
		if exists {
			t.Fatalf("Expected not exists")
		}
	}
}

func TestAddMetricTags(t *testing.T) {
	cm := &CheckManager{}
	cm.checkBundle = &api.CheckBundle{}
	cm.metricTags = make(map[string][]string)

	t.Log("no tags")
	{
		if cm.AddMetricTags("foo", []string{}, false) {
			t.Fatalf("Expected false")
		}
	}

	t.Log("no metric named 'foo'")
	{
		if !cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected true")
		}
	}

	cm.checkBundle.Metrics = []api.CheckBundleMetric{
		api.CheckBundleMetric{
			Name:   "bar",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "foo",
			Type:   "numeric",
			Status: "active",
		},
		api.CheckBundleMetric{
			Name:   "baz",
			Type:   "numeric",
			Status: "active",
			Tags:   []string{"cat1:tag1"},
		},
	}

	t.Log("metric named 'bar', add tag")
	{

		cm.metricTags = make(map[string][]string)
		// append, zero current
		if !cm.AddMetricTags("bar", []string{"cat:tag"}, true) {
			t.Fatalf("Expected true")
		}
		expected := make(map[string][]string)
		expected["bar"] = []string{"cat:tag"}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// tag already exists, no need to add
		if cm.AddMetricTags("bar", []string{"cat:tag"}, true) {
			t.Fatalf("Expected false")
		}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// append, zero tags
		if cm.AddMetricTags("bar", []string{}, true) {
			t.Fatalf("Expected false")
		}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}
	}

	t.Log("metric named 'baz', add tag")
	{
		cm.metricTags = make(map[string][]string)

		// append, current tag, should be noupdate
		if cm.AddMetricTags("baz", []string{"cat1:tag1"}, true) {
			t.Fatalf("Expected false")
		}

		if _, found := cm.metricTags["baz"]; found {
			t.Fatalf("expected not found")
		}

		// append, one current
		if !cm.AddMetricTags("baz", []string{"cat2:tag2"}, true) {
			t.Fatalf("Expected true")
		}
		expected := make(map[string][]string)
		expected["baz"] = []string{"cat1:tag1", "cat2:tag2"}

		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// append, tag already exists (should be noupdate)
		if cm.AddMetricTags("baz", []string{"cat2:tag2"}, true) {
			t.Fatalf("Expected false")
		}
	}

	t.Log("metric named 'foo', set tag")
	{
		expected := make(map[string][]string)
		cm.metricTags = make(map[string][]string)

		// set tag
		if !cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected true")
		}
		expected["foo"] = []string{"cat:tag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// set, reset (pass 0 tags)
		if !cm.AddMetricTags("foo", []string{}, false) {
			t.Fatalf("Expected true")
		}

		if _, found := cm.metricTags["foo"]; !found {
			t.Fatal("expected found")
		}

		expected["foo"] = []string{}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// set tag
		if !cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected true")
		}

		expected["foo"] = []string{"cat:tag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// set, no update (same tag)
		if cm.AddMetricTags("foo", []string{"cat:tag"}, false) {
			t.Fatalf("Expected false")
		}

		expected["foo"] = []string{"cat:tag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

		// replace any existing
		if !cm.AddMetricTags("foo", []string{"cat:newtag"}, false) {
			t.Fatalf("Expected true")
		}

		expected["foo"] = []string{"cat:newtag"}
		if !reflect.DeepEqual(cm.metricTags, expected) {
			t.Fatalf("expected %v got %+v", expected, cm.metricTags)
		}

	}
}

func TestAddNewMetrics(t *testing.T) {
	cm := &CheckManager{}

	newMetrics := make(map[string]*api.CheckBundleMetric)

	newMetrics["foo"] = &api.CheckBundleMetric{
		Name:   "foo",
		Type:   "numeric",
		Status: "active",
	}

	t.Log("no check bundle")
	{
		if cm.addNewMetrics(newMetrics) {
			t.Fatalf("Expected false")
		}
	}

	cm.checkBundle = &api.CheckBundle{}
	t.Log("no check bundle metrics")
	{
		if !cm.addNewMetrics(newMetrics) {
			t.Fatalf("Expected true")
		}
		if !cm.forceCheckUpdate {
			t.Fatal("Expected forceCheckUpdate to be true")
		}
	}
}
<|endoftext|>


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

package regexp_test

import (
	"fmt"
	"regexp"
)

func Example() {
	// Compile the expression once, usually at init time.
	// Use raw strings to avoid having to quote the backslashes.
	var validID = regexp.MustCompile(`^[a-z]+\[[0-9]+\]$`)

	fmt.Println(validID.MatchString("adam[23]"))
	fmt.Println(validID.MatchString("eve[7]"))
	fmt.Println(validID.MatchString("Job[48]"))
	fmt.Println(validID.MatchString("snakey"))
	// Output:
	// true
	// true
	// false
	// false
}

func ExampleMatchString() {
	matched, err := regexp.MatchString("foo.*", "seafood")
	fmt.Println(matched, err)
	matched, err = regexp.MatchString("bar.*", "seafood")
	fmt.Println(matched, err)
	matched, err = regexp.MatchString("a(b", "seafood")
	fmt.Println(matched, err)
	// Output:
	// true <nil>
	// false <nil>
	// false error parsing regexp: missing closing ): `a(b`
}

func ExampleQuoteMeta() {
	fmt.Println(regexp.QuoteMeta("Escaping symbols like: .+*?()|[]{}^$"))
	// Output:
	// Escaping symbols like: \.\+\*\?\(\)\|\[\]\{\}\^\$
}

func ExampleRegexp_FindString() {
	re := regexp.MustCompile("foo.?")
	fmt.Printf("%q\n", re.FindString("seafood fool"))
	fmt.Printf("%q\n", re.FindString("meat"))
	// Output:
	// "food"
	// ""
}

func ExampleRegexp_FindStringIndex() {
	re := regexp.MustCompile("ab?")
	fmt.Println(re.FindStringIndex("tablett"))
	fmt.Println(re.FindStringIndex("foo") == nil)
	// Output:
	// [1 3]
	// true
}

func ExampleRegexp_FindStringSubmatch() {
	re := regexp.MustCompile("a(x*)b(y|z)c")
	fmt.Printf("%q\n", re.FindStringSubmatch("-axxxbyc-"))
	fmt.Printf("%q\n", re.FindStringSubmatch("-abzc-"))
	// Output:
	// ["axxxbyc" "xxx" "y"]
	// ["abzc" "" "z"]
}

func ExampleRegexp_FindAllString() {
	re := regexp.MustCompile("a.")
	fmt.Println(re.FindAllString("paranormal", -1))
	fmt.Println(re.FindAllString("paranormal", 2))
	fmt.Println(re.FindAllString("graal", -1))
	fmt.Println(re.FindAllString("none", -1))
	// Output:
	// [ar an al]
	// [ar an]
	// [aa]
	// []
}

func ExampleRegexp_FindAllStringSubmatch() {
	re := regexp.MustCompile("a(x*)b")
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-ab-", -1))
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-axxb-", -1))
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-ab-axb-", -1))
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-axxb-ab-", -1))
	// Output:
	// [["ab" ""]]
	// [["axxb" "xx"]]
	// [["ab" ""] ["axb" "x"]]
	// [["axxb" "xx"] ["ab" ""]]
}

func ExampleRegexp_FindAllStringSubmatchIndex() {
	re := regexp.MustCompile("a(x*)b")
	// Indices:
	//    01234567   012345678
	//    -ab-axb-   -axxb-ab-
	fmt.Println(re.FindAllStringSubmatchIndex("-ab-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-axxb-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-ab-axb-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-axxb-ab-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-foo-", -1))
	// Output:
	// [[1 3 2 2]]
	// [[1 5 2 4]]
	// [[1 3 2 2] [4 7 5 6]]
	// [[1 5 2 4] [6 8 7 7]]
	// []
}

func ExampleRegexp_MatchString() {
	re := regexp.MustCompile("(gopher){2}")
	fmt.Println(re.MatchString("gopher"))
	fmt.Println(re.MatchString("gophergopher"))
	fmt.Println(re.MatchString("gophergophergopher"))
	// Output:
	// false
	// true
	// true
}

func ExampleRegexp_ReplaceAllLiteralString() {
	re := regexp.MustCompile("a(x*)b")
	fmt.Println(re.ReplaceAllLiteralString("-ab-axxb-", "T"))
	fmt.Println(re.ReplaceAllLiteralString("-ab-axxb-", "$1"))
	fmt.Println(re.ReplaceAllLiteralString("-ab-axxb-", "${1}"))
	// Output:
	// -T-T-
	// -$1-$1-
	// -${1}-${1}-
}

func ExampleRegexp_ReplaceAllString() {
	re := regexp.MustCompile("a(x*)b")
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "T"))
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "$1"))
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "$1W"))
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "${1}W"))
	// Output:
	// -T-T-
	// --xx-
	// ---
	// -W-xxW-
}

func ExampleRegexp_SubexpNames() {
	re := regexp.MustCompile("(?P<first>[a-zA-Z]+) (?P<last>[a-zA-Z]+)")
	fmt.Println(re.MatchString("Alan Turing"))
	fmt.Printf("%q\n", re.SubexpNames())
	reversed := fmt.Sprintf("${%s} ${%s}", re.SubexpNames()[2], re.SubexpNames()[1])
	fmt.Println(reversed)
	fmt.Println(re.ReplaceAllString("Alan Turing", reversed))
	// Output:
	// true
	// ["" "first" "last"]
	// ${last} ${first}
	// Turing Alan
}

func ExampleRegexp_Split() {
	a := regexp.MustCompile("a")
	fmt.Println(a.Split("banana", -1))
	fmt.Println(a.Split("banana", 0))
	fmt.Println(a.Split("banana", 1))
	fmt.Println(a.Split("banana", 2))
	zp := regexp.MustCompile("z+")
	fmt.Println(zp.Split("pizza", -1))
	fmt.Println(zp.Split("pizza", 0))
	fmt.Println(zp.Split("pizza", 1))
	fmt.Println(zp.Split("pizza", 2))
	// Output:
	// [b n n ]
	// []
	// [banana]
	// [b nana]
	// [pi a]
	// []
	// [pizza]
	// [pi a]
}

func ExampleRegexp_Expand() {
	content := []byte(`
	# comment line
	option1: value1
	option2: value2

	# another comment line
	option3: value3
`)

	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(`(?m)(?P<key>\w+):\s+(?P<value>\w+)$`)

	// Template to convert "key: value" to "key=value" by
	// referencing the values captured by the regex pattern.
	template := []byte("$key=$value\n")

	result := []byte{}

	// For each match of the regex in the content.
	for _, submatches := range pattern.FindAllSubmatchIndex(content, -1) {
		// Apply the captured submatches to the template and append the output
		// to the result.
		result = pattern.Expand(result, template, content, submatches)
	}
	fmt.Println(string(result))
	// Output:
	// option1=value1
	// option2=value2
	// option3=value3
}

func ExampleRegexp_ExpandString() {
	content := `
	# comment line
	option1: value1
	option2: value2

	# another comment line
	option3: value3
`

	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(`(?m)(?P<key>\w+):\s+(?P<value>\w+)$`)

	// Template to convert "key: value" to "key=value" by
	// referencing the values captured by the regex pattern.
	template := "$key=$value\n"

	result := []byte{}

	// For each match of the regex in the content.
	for _, submatches := range pattern.FindAllStringSubmatchIndex(content, -1) {
		// Apply the captured submatches to the template and append the output
		// to the result.
		result = pattern.ExpandString(result, template, content, submatches)
	}
	fmt.Println(string(result))
	// Output:
	// option1=value1
	// option2=value2
	// option3=value3
}
<commit_msg>regexp: examples for Regexp.FindIndex and Regexp.FindAllSubmatchIndex methods<commit_after>// Copyright 2013 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package regexp_test

import (
	"fmt"
	"regexp"
)

func Example() {
	// Compile the expression once, usually at init time.
	// Use raw strings to avoid having to quote the backslashes.
	var validID = regexp.MustCompile(`^[a-z]+\[[0-9]+\]$`)

	fmt.Println(validID.MatchString("adam[23]"))
	fmt.Println(validID.MatchString("eve[7]"))
	fmt.Println(validID.MatchString("Job[48]"))
	fmt.Println(validID.MatchString("snakey"))
	// Output:
	// true
	// true
	// false
	// false
}

func ExampleMatchString() {
	matched, err := regexp.MatchString("foo.*", "seafood")
	fmt.Println(matched, err)
	matched, err = regexp.MatchString("bar.*", "seafood")
	fmt.Println(matched, err)
	matched, err = regexp.MatchString("a(b", "seafood")
	fmt.Println(matched, err)
	// Output:
	// true <nil>
	// false <nil>
	// false error parsing regexp: missing closing ): `a(b`
}

func ExampleQuoteMeta() {
	fmt.Println(regexp.QuoteMeta("Escaping symbols like: .+*?()|[]{}^$"))
	// Output:
	// Escaping symbols like: \.\+\*\?\(\)\|\[\]\{\}\^\$
}

func ExampleRegexp_FindString() {
	re := regexp.MustCompile("foo.?")
	fmt.Printf("%q\n", re.FindString("seafood fool"))
	fmt.Printf("%q\n", re.FindString("meat"))
	// Output:
	// "food"
	// ""
}

func ExampleRegexp_FindStringIndex() {
	re := regexp.MustCompile("ab?")
	fmt.Println(re.FindStringIndex("tablett"))
	fmt.Println(re.FindStringIndex("foo") == nil)
	// Output:
	// [1 3]
	// true
}

func ExampleRegexp_FindStringSubmatch() {
	re := regexp.MustCompile("a(x*)b(y|z)c")
	fmt.Printf("%q\n", re.FindStringSubmatch("-axxxbyc-"))
	fmt.Printf("%q\n", re.FindStringSubmatch("-abzc-"))
	// Output:
	// ["axxxbyc" "xxx" "y"]
	// ["abzc" "" "z"]
}

func ExampleRegexp_FindAllString() {
	re := regexp.MustCompile("a.")
	fmt.Println(re.FindAllString("paranormal", -1))
	fmt.Println(re.FindAllString("paranormal", 2))
	fmt.Println(re.FindAllString("graal", -1))
	fmt.Println(re.FindAllString("none", -1))
	// Output:
	// [ar an al]
	// [ar an]
	// [aa]
	// []
}

func ExampleRegexp_FindAllStringSubmatch() {
	re := regexp.MustCompile("a(x*)b")
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-ab-", -1))
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-axxb-", -1))
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-ab-axb-", -1))
	fmt.Printf("%q\n", re.FindAllStringSubmatch("-axxb-ab-", -1))
	// Output:
	// [["ab" ""]]
	// [["axxb" "xx"]]
	// [["ab" ""] ["axb" "x"]]
	// [["axxb" "xx"] ["ab" ""]]
}

func ExampleRegexp_FindAllStringSubmatchIndex() {
	re := regexp.MustCompile("a(x*)b")
	// Indices:
	//    01234567   012345678
	//    -ab-axb-   -axxb-ab-
	fmt.Println(re.FindAllStringSubmatchIndex("-ab-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-axxb-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-ab-axb-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-axxb-ab-", -1))
	fmt.Println(re.FindAllStringSubmatchIndex("-foo-", -1))
	// Output:
	// [[1 3 2 2]]
	// [[1 5 2 4]]
	// [[1 3 2 2] [4 7 5 6]]
	// [[1 5 2 4] [6 8 7 7]]
	// []
}

func ExampleRegexp_MatchString() {
	re := regexp.MustCompile("(gopher){2}")
	fmt.Println(re.MatchString("gopher"))
	fmt.Println(re.MatchString("gophergopher"))
	fmt.Println(re.MatchString("gophergophergopher"))
	// Output:
	// false
	// true
	// true
}

func ExampleRegexp_ReplaceAllLiteralString() {
	re := regexp.MustCompile("a(x*)b")
	fmt.Println(re.ReplaceAllLiteralString("-ab-axxb-", "T"))
	fmt.Println(re.ReplaceAllLiteralString("-ab-axxb-", "$1"))
	fmt.Println(re.ReplaceAllLiteralString("-ab-axxb-", "${1}"))
	// Output:
	// -T-T-
	// -$1-$1-
	// -${1}-${1}-
}

func ExampleRegexp_ReplaceAllString() {
	re := regexp.MustCompile("a(x*)b")
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "T"))
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "$1"))
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "$1W"))
	fmt.Println(re.ReplaceAllString("-ab-axxb-", "${1}W"))
	// Output:
	// -T-T-
	// --xx-
	// ---
	// -W-xxW-
}

func ExampleRegexp_SubexpNames() {
	re := regexp.MustCompile("(?P<first>[a-zA-Z]+) (?P<last>[a-zA-Z]+)")
	fmt.Println(re.MatchString("Alan Turing"))
	fmt.Printf("%q\n", re.SubexpNames())
	reversed := fmt.Sprintf("${%s} ${%s}", re.SubexpNames()[2], re.SubexpNames()[1])
	fmt.Println(reversed)
	fmt.Println(re.ReplaceAllString("Alan Turing", reversed))
	// Output:
	// true
	// ["" "first" "last"]
	// ${last} ${first}
	// Turing Alan
}

func ExampleRegexp_Split() {
	a := regexp.MustCompile("a")
	fmt.Println(a.Split("banana", -1))
	fmt.Println(a.Split("banana", 0))
	fmt.Println(a.Split("banana", 1))
	fmt.Println(a.Split("banana", 2))
	zp := regexp.MustCompile("z+")
	fmt.Println(zp.Split("pizza", -1))
	fmt.Println(zp.Split("pizza", 0))
	fmt.Println(zp.Split("pizza", 1))
	fmt.Println(zp.Split("pizza", 2))
	// Output:
	// [b n n ]
	// []
	// [banana]
	// [b nana]
	// [pi a]
	// []
	// [pizza]
	// [pi a]
}

func ExampleRegexp_Expand() {
	content := []byte(`
	# comment line
	option1: value1
	option2: value2

	# another comment line
	option3: value3
`)

	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(`(?m)(?P<key>\w+):\s+(?P<value>\w+)$`)

	// Template to convert "key: value" to "key=value" by
	// referencing the values captured by the regex pattern.
	template := []byte("$key=$value\n")

	result := []byte{}

	// For each match of the regex in the content.
	for _, submatches := range pattern.FindAllSubmatchIndex(content, -1) {
		// Apply the captured submatches to the template and append the output
		// to the result.
		result = pattern.Expand(result, template, content, submatches)
	}
	fmt.Println(string(result))
	// Output:
	// option1=value1
	// option2=value2
	// option3=value3
}

func ExampleRegexp_ExpandString() {
	content := `
	# comment line
	option1: value1
	option2: value2

	# another comment line
	option3: value3
`

	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(`(?m)(?P<key>\w+):\s+(?P<value>\w+)$`)

	// Template to convert "key: value" to "key=value" by
	// referencing the values captured by the regex pattern.
	template := "$key=$value\n"

	result := []byte{}

	// For each match of the regex in the content.
	for _, submatches := range pattern.FindAllStringSubmatchIndex(content, -1) {
		// Apply the captured submatches to the template and append the output
		// to the result.
		result = pattern.ExpandString(result, template, content, submatches)
	}
	fmt.Println(string(result))
	// Output:
	// option1=value1
	// option2=value2
	// option3=value3
}

func ExampleRegexp_FindIndex() {
	content := []byte(`
	# comment line
	option1: value1
	option2: value2
`)
	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(`(?m)(?P<key>\w+):\s+(?P<value>\w+)$`)

	loc := pattern.FindIndex(content)
	fmt.Println(loc)
	fmt.Println(string(content[loc[0]:loc[1]]))
	// Output:
	// [18 33]
	// option1: value1
}
func ExampleRegexp_FindAllSubmatchIndex() {
	content := []byte(`
	# comment line
	option1: value1
	option2: value2
`)
	// Regex pattern captures "key: value" pair from the content.
	pattern := regexp.MustCompile(`(?m)(?P<key>\w+):\s+(?P<value>\w+)$`)
	allIndexes := pattern.FindAllSubmatchIndex(content, -1)
	for _, loc := range allIndexes {
		fmt.Println(loc)
		fmt.Println(string(content[loc[0]:loc[1]]))
		fmt.Println(string(content[loc[2]:loc[3]]))
		fmt.Println(string(content[loc[4]:loc[5]]))
	}
	// Output:
	// [18 33 18 25 27 33]
	// option1: value1
	// option1
	// value1
	// [35 50 35 42 44 50]
	// option2: value2
	// option2
	// value2
}
<|endoftext|>


<commit_before>package main

// commandline tool to grab, scrape and output a news article
//
// can grab article via http or from a file (raw html or the
// first response in a .warc)
//

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"errors"
	"flag"
	"fmt"
	"github.com/bcampbell/arts/arts"
	"github.com/bcampbell/warc/warc"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
)

func main() {
	var debug string
	var parseOnly bool
	flag.StringVar(&debug, "d", "", "log debug info to stderr (h=headline, c=content, a=authors d=dates u=urls s=cruft all=hcadus)")
	flag.BoolVar(&parseOnly, "parse", false, "just dump the parsed html and exit")
	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	flag.Parse()

	if len(flag.Args()) != 1 {
		fmt.Println("Usage: ", os.Args[0], "<article url>")
		os.Exit(1)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// set up the debug logging
	debug = strings.ToLower(debug)
	if debug == "name" {
		debug = ""
	}
	if debug == "all" {
		debug = "hcadus"
	}
	for _, flag := range debug {
		switch flag {
		case 'h':
			arts.Debug.HeadlineLogger = log.New(os.Stderr, "", 0)
		case 'c':
			arts.Debug.ContentLogger = log.New(os.Stderr, "", 0)
		case 'a':
			arts.Debug.AuthorsLogger = log.New(os.Stderr, "", 0)
		case 'd':
			arts.Debug.DatesLogger = log.New(os.Stderr, "", 0)
		case 'u':
			arts.Debug.URLLogger = log.New(os.Stderr, "", 0)
		case 's':
			arts.Debug.CruftLogger = log.New(os.Stderr, "", 0)
		}
	}

	var rawHTML []byte
	var artURL string

	srcName := flag.Arg(0)
	u, err := url.Parse(srcName)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s is not url: %s", srcName, err)
		os.Exit(1)
	}

	var in io.ReadCloser
	switch strings.ToLower(u.Scheme) {
	case "http", "https":
		artURL = srcName
		in, err = openHttp(srcName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: http fetch failed: %s", err)
			os.Exit(1)
		}
		rawHTML, err = ioutil.ReadAll(in)
		in.Close()
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: read failed: %s", err)
			os.Exit(1)
		}
	case "file", "":

		foo := strings.ToLower(u.Path)
		if strings.HasSuffix(foo, ".warc") || strings.HasSuffix(foo, ".warc.gz") {
			// it's a warc file
			rawHTML, artURL, err = fromWARC(u.Path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "ERROR: warc read failed: %s", err)
				os.Exit(1)
			}
		} else {
			// treat as plain html file (url will suck)
			in, err = os.Open(u.Path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "ERROR: open failed: %s", err)
				os.Exit(1)
			}
			rawHTML, err = ioutil.ReadAll(in)
			in.Close()
			if err != nil {
				fmt.Fprintf(os.Stderr, "ERROR: read failed: %s", err)
				os.Exit(1)
			}
		}
	}

	root, err := arts.ParseHTML(rawHTML)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: html parsing failed: %s", err)
		os.Exit(1)
	}

	if parseOnly {
		err = html.Render(os.Stdout, root)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: html render failed: %s", err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	art, err := arts.ExtractFromTree(root, artURL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: extraction failed: %s", err)
		os.Exit(1)
	}

	err = dumpArt(os.Stdout, art)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: dumping article to stdout: %s", err)
		os.Exit(1)
	}
}

// fetch html from a WARC file
// returns: html, url, err
func fromWARC(filename string) ([]byte, string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, "", err
	}
	defer f.Close()

	var in io.Reader
	if filepath.Ext(filename) == ".gz" {
		gin, err := gzip.NewReader(f)
		if err != nil {
			return nil, "", err
		}
		defer gin.Close()
		in = gin
	} else {
		in = f
	}

	warcReader := warc.NewReader(in)
	for {
		//	fmt.Printf("WARC\n")
		rec, err := warcReader.ReadRecord()
		if err != nil {
			return nil, "", fmt.Errorf("Error reading %s: %s", filename, err)
		}
		if rec.Header.Get("Warc-Type") != "response" {
			continue
		}
		reqURL := rec.Header.Get("Warc-Target-Uri")
		// parse response, grab raw html
		rdr := bufio.NewReader(bytes.NewReader(rec.Block))
		response, err := http.ReadResponse(rdr, nil)
		if err != nil {
			return nil, "", fmt.Errorf("Error parsing response: %s", err)
		}
		defer response.Body.Close()
		if response.StatusCode != 200 {
			return nil, "", fmt.Errorf("HTTP error: %d", response.StatusCode)
		}
		rawHTML, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return nil, "", err
		}
		return rawHTML, reqURL, err
	}
}

func openHttp(artURL string) (io.ReadCloser, error) {

	client := &http.Client{}

	req, err := http.NewRequest("GET", artURL, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Accept", "*/*")
	req.Header.Set("User-Agent", `Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:28.0) Gecko/20100101 Firefox/28.0`)
	//	req.Header.Set("User-Agent", "curl/7.47.0")

	response, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != 200 {
		return nil, errors.New(fmt.Sprintf("Request failed: %s", response.Status))
	}
	return response.Body, nil
}
<commit_msg>fix up github.com/bcampbell/warc import<commit_after>package main

// commandline tool to grab, scrape and output a news article
//
// can grab article via http or from a file (raw html or the
// first response in a .warc)
//

import (
	"bufio"
	"bytes"
	"compress/gzip"
	"errors"
	"flag"
	"fmt"
	"github.com/bcampbell/arts/arts"
	"github.com/bcampbell/warc"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"runtime/pprof"
	"strings"
)

func main() {
	var debug string
	var parseOnly bool
	flag.StringVar(&debug, "d", "", "log debug info to stderr (h=headline, c=content, a=authors d=dates u=urls s=cruft all=hcadus)")
	flag.BoolVar(&parseOnly, "parse", false, "just dump the parsed html and exit")
	var cpuprofile = flag.String("cpuprofile", "", "write cpu profile to file")
	flag.Parse()

	if len(flag.Args()) != 1 {
		fmt.Println("Usage: ", os.Args[0], "<article url>")
		os.Exit(1)
	}

	if *cpuprofile != "" {
		f, err := os.Create(*cpuprofile)
		if err != nil {
			log.Fatal(err)
		}
		pprof.StartCPUProfile(f)
		defer pprof.StopCPUProfile()
	}

	// set up the debug logging
	debug = strings.ToLower(debug)
	if debug == "name" {
		debug = ""
	}
	if debug == "all" {
		debug = "hcadus"
	}
	for _, flag := range debug {
		switch flag {
		case 'h':
			arts.Debug.HeadlineLogger = log.New(os.Stderr, "", 0)
		case 'c':
			arts.Debug.ContentLogger = log.New(os.Stderr, "", 0)
		case 'a':
			arts.Debug.AuthorsLogger = log.New(os.Stderr, "", 0)
		case 'd':
			arts.Debug.DatesLogger = log.New(os.Stderr, "", 0)
		case 'u':
			arts.Debug.URLLogger = log.New(os.Stderr, "", 0)
		case 's':
			arts.Debug.CruftLogger = log.New(os.Stderr, "", 0)
		}
	}

	var rawHTML []byte
	var artURL string

	srcName := flag.Arg(0)
	u, err := url.Parse(srcName)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: %s is not url: %s", srcName, err)
		os.Exit(1)
	}

	var in io.ReadCloser
	switch strings.ToLower(u.Scheme) {
	case "http", "https":
		artURL = srcName
		in, err = openHttp(srcName)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: http fetch failed: %s", err)
			os.Exit(1)
		}
		rawHTML, err = ioutil.ReadAll(in)
		in.Close()
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: read failed: %s", err)
			os.Exit(1)
		}
	case "file", "":

		foo := strings.ToLower(u.Path)
		if strings.HasSuffix(foo, ".warc") || strings.HasSuffix(foo, ".warc.gz") {
			// it's a warc file
			rawHTML, artURL, err = fromWARC(u.Path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "ERROR: warc read failed: %s", err)
				os.Exit(1)
			}
		} else {
			// treat as plain html file (url will suck)
			in, err = os.Open(u.Path)
			if err != nil {
				fmt.Fprintf(os.Stderr, "ERROR: open failed: %s", err)
				os.Exit(1)
			}
			rawHTML, err = ioutil.ReadAll(in)
			in.Close()
			if err != nil {
				fmt.Fprintf(os.Stderr, "ERROR: read failed: %s", err)
				os.Exit(1)
			}
		}
	}

	root, err := arts.ParseHTML(rawHTML)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: html parsing failed: %s", err)
		os.Exit(1)
	}

	if parseOnly {
		err = html.Render(os.Stdout, root)
		if err != nil {
			fmt.Fprintf(os.Stderr, "ERROR: html render failed: %s", err)
			os.Exit(1)
		}
		os.Exit(0)
	}

	art, err := arts.ExtractFromTree(root, artURL)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: extraction failed: %s", err)
		os.Exit(1)
	}

	err = dumpArt(os.Stdout, art)
	if err != nil {
		fmt.Fprintf(os.Stderr, "ERROR: dumping article to stdout: %s", err)
		os.Exit(1)
	}
}

// fetch html from a WARC file
// returns: html, url, err
func fromWARC(filename string) ([]byte, string, error) {
	f, err := os.Open(filename)
	if err != nil {
		return nil, "", err
	}
	defer f.Close()

	var in io.Reader
	if filepath.Ext(filename) == ".gz" {
		gin, err := gzip.NewReader(f)
		if err != nil {
			return nil, "", err
		}
		defer gin.Close()
		in = gin
	} else {
		in = f
	}

	warcReader := warc.NewReader(in)
	for {
		//	fmt.Printf("WARC\n")
		rec, err := warcReader.ReadRecord()
		if err != nil {
			return nil, "", fmt.Errorf("Error reading %s: %s", filename, err)
		}
		if rec.Header.Get("Warc-Type") != "response" {
			continue
		}
		reqURL := rec.Header.Get("Warc-Target-Uri")
		// parse response, grab raw html
		rdr := bufio.NewReader(bytes.NewReader(rec.Block))
		response, err := http.ReadResponse(rdr, nil)
		if err != nil {
			return nil, "", fmt.Errorf("Error parsing response: %s", err)
		}
		defer response.Body.Close()
		if response.StatusCode != 200 {
			return nil, "", fmt.Errorf("HTTP error: %d", response.StatusCode)
		}
		rawHTML, err := ioutil.ReadAll(response.Body)
		if err != nil {
			return nil, "", err
		}
		return rawHTML, reqURL, err
	}
}

func openHttp(artURL string) (io.ReadCloser, error) {

	client := &http.Client{}

	req, err := http.NewRequest("GET", artURL, nil)
	if err != nil {
		return nil, err
	}

	req.Header.Set("Accept", "*/*")
	req.Header.Set("User-Agent", `Mozilla/5.0 (X11; Ubuntu; Linux x86_64; rv:28.0) Gecko/20100101 Firefox/28.0`)
	//	req.Header.Set("User-Agent", "curl/7.47.0")

	response, err := client.Do(req)
	if err != nil {
		return nil, err
	}

	if response.StatusCode != 200 {
		return nil, errors.New(fmt.Sprintf("Request failed: %s", response.Status))
	}
	return response.Body, nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2012 VMware, Inc.

package gonatsd_test

import (
	"bufio"
	"code.google.com/p/gomock/gomock"
	. "gonatsd/gonatsd"
	. "gonatsd/gonatsd/mocks"
	"io"
	. "launchpad.net/gocheck"
	"time"
)

type ConnSuite struct {
	delegate        *TestReporterDelegate
	ctrl            *gomock.Controller
	conn            Conn
	tcpConn         *DummyTCPConn
	started         bool
	server          *MockServer
	serverCmds      chan ServerCmd
	heartbeatHelper *MockHeartbeatHelper
	pingTicker      chan time.Time
}

var _ = Suite(&ConnSuite{})

func (s *ConnSuite) SetUpTest(c *C) {
	s.delegate = &TestReporterDelegate{}
	s.ctrl = gomock.NewController(s.delegate)

	s.pingTicker = make(chan time.Time, 1)
	s.tcpConn = NewDummyTCPConn()
	s.serverCmds = make(chan ServerCmd, 1)

	s.heartbeatHelper = NewMockHeartbeatHelper(s.ctrl)
	s.heartbeatHelper.EXPECT().Ticker().Return(s.pingTicker).AnyTimes()
	s.heartbeatHelper.EXPECT().Stop().AnyTimes()
	NewHeartbeatHelper = func(Conn, time.Duration, int) HeartbeatHelper {
		return s.heartbeatHelper
	}

	config := &Config{}
	config.Limits.ControlLine = 20
	config.Limits.Pending = 1024

	s.server = NewMockServer(s.ctrl)
	s.server.EXPECT().Config().Return(config).AnyTimes()
	s.server.EXPECT().Info().Return(&dummyInfo).AnyTimes()
	s.server.EXPECT().Commands().Return(s.serverCmds).AnyTimes()
	s.server.EXPECT().Stats().Return(NewStats()).AnyTimes()
}

func (s *ConnSuite) TearDownTest(c *C) {
	NewHeartbeatHelper = NewRealHeartbeatHelper
	BUF_IO_SIZE = 262144

	if s.conn != nil {
		if !s.conn.Closed() {
			// Dummy server unregister mock
			go func() {
				cmd := <-s.serverCmds
				switch cmd.(type) {
				case *UnregisterConnCmd:
					cmd := cmd.(*UnregisterConnCmd)
					cmd.Done <- true
				}
			}()

			// Consume any unread output so write can finish
			go func() {
				buf := make([]byte, 1024)
				for {
					_, err := s.tcpConn.client.Read(buf)
					if err != nil {
						break
					}
				}
			}()

			s.conn.Close()
		}
	}
}

func (s *ConnSuite) TestNewConn(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.server.Config().Ping.IntervalDuration = 5 * time.Second
	s.server.Config().Ping.MaxOutstanding = 7

	NewHeartbeatHelper = func(conn Conn, interval time.Duration, maxOutstanding int) HeartbeatHelper {
		c.Check(interval, Equals, 5*time.Second)
		c.Check(maxOutstanding, Equals, 7)
		return s.heartbeatHelper
	}
	s.conn = NewConn(s.server, s.tcpConn)
}

func (s *ConnSuite) TestReadControlLine(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)

	go func() {
		io.WriteString(s.tcpConn.client, "Hello\r\n\r\n")
		s.tcpConn.client.Close()
	}()

	line, err := s.conn.ReadControlLine()
	c.Check(err, IsNil)
	c.Check(line, Equals, "Hello")
	line, err = s.conn.ReadControlLine()
	c.Check(err, IsNil)
	c.Check(line, Equals, "")
	line, err = s.conn.ReadControlLine()
	c.Check(err, Equals, io.EOF)
	c.Check(line, Equals, "")
}

func (s *ConnSuite) TestReadControlLineTooLong(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)

	s.server.Config().Limits.ControlLine = 4

	go func() {
		io.WriteString(s.tcpConn.client, "1234\r\n12345\r\n")
		s.tcpConn.client.Close()
	}()

	line, err := s.conn.ReadControlLine()
	c.Check(err, IsNil)
	c.Check(line, Equals, "1234")
	line, err = s.conn.ReadControlLine()
	c.Check(err, Equals, ErrProtocolOpTooBig)
}

func (s *ConnSuite) TestReadControlLineTooLongPartial(c *C) {
	BUF_IO_SIZE = 16

	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)

	go func() {
		io.WriteString(s.tcpConn.client, "12345678901234567890123456789012\r\n")
		s.tcpConn.client.Close()
	}()

	_, err := s.conn.ReadControlLine()
	c.Check(err, Equals, ErrProtocolOpTooBig)
}

func (s *ConnSuite) TestMessageDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	sm := &SubscribedMessage{&Subscription{}, &Message{"foo", "X", []byte("msg")}, false}
	s.conn.ServeMessage(sm)

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "MSG foo 0 X 3")
	checkReadLine(c, reader, "msg")
}

type TestRequest struct{}

func (t *TestRequest) Serve(c Conn) *Response {
	return NewStringResponse("TestReq")
}

func (t *TestRequest) Dispatch(c Conn) {
	c.ServeRequest(t)
}

func (s *ConnSuite) TestRequestDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.conn.ServeRequest(&TestRequest{})

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "TestReq")
}

type TestClientCmd struct{}

func (t *TestClientCmd) Process(c Conn) {
	c.Write(NewStringResponse("TestCmd"))
}

func (s *ConnSuite) TestCommandDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.conn.ServeCommand(&TestClientCmd{})

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "TestCmd")
}

func (s *ConnSuite) TestHeartbeatDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	var barrier = make(chan bool, 1)
	s.heartbeatHelper.EXPECT().Ping().Do(func() {
		barrier <- true
	})

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.pingTicker <- time.Now()

	// Give it a chance to be dispatched
	select {
	case <-barrier:
	case <-time.After(time.Second):
	}

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
}

func (s *ConnSuite) TestReadLoopTooBig(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.server.Config().Limits.ControlLine = 4

	go func() {
		cmd := <-s.serverCmds
		switch cmd.(type) {
		case *UnregisterConnCmd:
			cmd := cmd.(*UnregisterConnCmd)
			cmd.Done <- true
		}
	}()

	io.WriteString(s.tcpConn.client, "TESTCOMMAND\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "-ERR 'Protocol Operation size exceeded'")
}

func (s *ConnSuite) TestReadRequest(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	io.WriteString(s.tcpConn.client, "PING\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "PONG")
}

func (s *ConnSuite) TestReadBadRequest(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	io.WriteString(s.tcpConn.client, "PING foo\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "-ERR 'Unknown Protocol Operation'")
}

func (s *ConnSuite) TestReadUnknownRequest(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	io.WriteString(s.tcpConn.client, "HELLO\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "-ERR 'Unknown Protocol Operation'")
}

func (s *ConnSuite) TestClientClosed(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.tcpConn.client.Close()
	serverCmd := <-s.serverCmds
	(serverCmd.(*UnregisterConnCmd)).Done <- true
}

type TestServerCmd struct{}

func (c *TestServerCmd) Process(Server) {

}

func (s *ConnSuite) TestSendServerCmd(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.conn.SendServerCmd(&TestServerCmd{})

	select {
	case <-s.serverCmds:
	case <-time.After(time.Second):
		c.Errorf("Should have received server command")
	}
}

type CallbackClientCmd struct {
	cb func()
}

func (c *CallbackClientCmd) Process(conn Conn) {
	if c.cb != nil {
		c.cb()
	}
}

func (s *ConnSuite) TestSendServerCmdFull(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	barrier := make(chan bool, 1)

	s.conn.ServeCommand(&CallbackClientCmd{cb: func() {
		// Fill up server channel
		s.serverCmds <- &TestServerCmd{}

		go func() {
			// Fill Up connection channel
			sm := &SubscribedMessage{}
			sm.Message = &Message{}
			sm.Subscription = &Subscription{}
			for i := 0; i < MAX_CONN_CHAN_BACKLOG+1; i++ {
				s.conn.ServeMessage(sm)
			}
			// Process a server command to unblock connection from sending
			<-s.serverCmds
		}()

		s.conn.SendServerCmd(&TestServerCmd{})

		select {
		case <-s.serverCmds:
		case <-time.After(time.Second):
			c.Errorf("Should have received server command")
		}

		barrier <- true
	}})

	<-barrier
}

func (s *ConnSuite) TestWriteFull(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.server.Config().Limits.Pending = 16
	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	reader := bufio.NewReader(s.tcpConn.client)

	value := "12345678901234567890"

	reader.ReadLine()

	go func() {
		cmd := <-s.serverCmds
		switch cmd.(type) {
		case *UnregisterConnCmd:
			cmd := cmd.(*UnregisterConnCmd)
			cmd.Done <- true
		}
	}()

	go s.conn.Write(NewStringResponse(value))
	checkReadLine(c, reader, "-ERR 'Slow consumer detected, connection dropped'")
}

func checkReadLine(c *C, reader *bufio.Reader, expected string) {
	line, prefix, err := reader.ReadLine()
	c.Check(err, IsNil)
	c.Check(prefix, Equals, false)
	c.Check(string(line), Equals, expected)
}
<commit_msg>Fix race condition on connection close in tests<commit_after>// Copyright (c) 2012 VMware, Inc.

package gonatsd_test

import (
	"bufio"
	"code.google.com/p/gomock/gomock"
	. "gonatsd/gonatsd"
	. "gonatsd/gonatsd/mocks"
	"io"
	. "launchpad.net/gocheck"
	"time"
)

type ConnSuite struct {
	delegate        *TestReporterDelegate
	ctrl            *gomock.Controller
	conn            Conn
	tcpConn         *DummyTCPConn
	started         bool
	server          *MockServer
	serverCmds      chan ServerCmd
	heartbeatHelper *MockHeartbeatHelper
	pingTicker      chan time.Time
}

var _ = Suite(&ConnSuite{})

func (s *ConnSuite) SetUpTest(c *C) {
	s.delegate = &TestReporterDelegate{}
	s.ctrl = gomock.NewController(s.delegate)

	s.pingTicker = make(chan time.Time, 1)
	s.tcpConn = NewDummyTCPConn()
	s.serverCmds = make(chan ServerCmd, 1)

	s.heartbeatHelper = NewMockHeartbeatHelper(s.ctrl)
	s.heartbeatHelper.EXPECT().Ticker().Return(s.pingTicker).AnyTimes()
	s.heartbeatHelper.EXPECT().Stop().AnyTimes()
	NewHeartbeatHelper = func(Conn, time.Duration, int) HeartbeatHelper {
		return s.heartbeatHelper
	}

	config := &Config{}
	config.Limits.ControlLine = 20
	config.Limits.Pending = 1024

	s.server = NewMockServer(s.ctrl)
	s.server.EXPECT().Config().Return(config).AnyTimes()
	s.server.EXPECT().Info().Return(&dummyInfo).AnyTimes()
	s.server.EXPECT().Commands().Return(s.serverCmds).AnyTimes()
	s.server.EXPECT().Stats().Return(NewStats()).AnyTimes()
}

func (s *ConnSuite) TearDownTest(c *C) {
	NewHeartbeatHelper = NewRealHeartbeatHelper
	BUF_IO_SIZE = 262144

	if s.conn != nil {
		if !s.conn.Closed() {
			// Dummy server unregister mock
			go func() {
				cmd := <-s.serverCmds
				switch cmd.(type) {
				case *UnregisterConnCmd:
					cmd := cmd.(*UnregisterConnCmd)
					cmd.Done <- true
				}
			}()

			// Consume any unread output so write can finish
			go func() {
				buf := make([]byte, 1024)
				for {
					_, err := s.tcpConn.client.Read(buf)
					if err != nil {
						break
					}
				}
			}()

			// Instead of just closing connection we schedule a CLOSE_CMD
			// in order to let dispatch loop read it, otherwise there is
			// a possible race condition: dispatch loop might kick in after
			// teardown closes the connection.
			s.conn.ServeCommand(CLOSE_CMD)
		}
	}
}

func (s *ConnSuite) TestNewConn(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.server.Config().Ping.IntervalDuration = 5 * time.Second
	s.server.Config().Ping.MaxOutstanding = 7

	NewHeartbeatHelper = func(conn Conn, interval time.Duration, maxOutstanding int) HeartbeatHelper {
		c.Check(interval, Equals, 5*time.Second)
		c.Check(maxOutstanding, Equals, 7)
		return s.heartbeatHelper
	}
	s.conn = NewConn(s.server, s.tcpConn)
}

func (s *ConnSuite) TestReadControlLine(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)

	go func() {
		io.WriteString(s.tcpConn.client, "Hello\r\n\r\n")
		s.tcpConn.client.Close()
	}()

	line, err := s.conn.ReadControlLine()
	c.Check(err, IsNil)
	c.Check(line, Equals, "Hello")
	line, err = s.conn.ReadControlLine()
	c.Check(err, IsNil)
	c.Check(line, Equals, "")
	line, err = s.conn.ReadControlLine()
	c.Check(err, Equals, io.EOF)
	c.Check(line, Equals, "")
}

func (s *ConnSuite) TestReadControlLineTooLong(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)

	s.server.Config().Limits.ControlLine = 4

	go func() {
		io.WriteString(s.tcpConn.client, "1234\r\n12345\r\n")
		s.tcpConn.client.Close()
	}()

	line, err := s.conn.ReadControlLine()
	c.Check(err, IsNil)
	c.Check(line, Equals, "1234")
	line, err = s.conn.ReadControlLine()
	c.Check(err, Equals, ErrProtocolOpTooBig)
}

func (s *ConnSuite) TestReadControlLineTooLongPartial(c *C) {
	BUF_IO_SIZE = 16

	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)

	go func() {
		io.WriteString(s.tcpConn.client, "12345678901234567890123456789012\r\n")
		s.tcpConn.client.Close()
	}()

	_, err := s.conn.ReadControlLine()
	c.Check(err, Equals, ErrProtocolOpTooBig)
}

func (s *ConnSuite) TestMessageDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	sm := &SubscribedMessage{&Subscription{}, &Message{"foo", "X", []byte("msg")}, false}
	s.conn.ServeMessage(sm)

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "MSG foo 0 X 3")
	checkReadLine(c, reader, "msg")
}

type TestRequest struct{}

func (t *TestRequest) Serve(c Conn) *Response {
	return NewStringResponse("TestReq")
}

func (t *TestRequest) Dispatch(c Conn) {
	c.ServeRequest(t)
}

func (s *ConnSuite) TestRequestDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.conn.ServeRequest(&TestRequest{})

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "TestReq")
}

type TestClientCmd struct{}

func (t *TestClientCmd) Process(c Conn) {
	c.Write(NewStringResponse("TestCmd"))
}

func (s *ConnSuite) TestCommandDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.conn.ServeCommand(&TestClientCmd{})

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "TestCmd")
}

func (s *ConnSuite) TestHeartbeatDispatch(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	var barrier = make(chan bool, 1)
	s.heartbeatHelper.EXPECT().Ping().Do(func() {
		barrier <- true
	})

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.pingTicker <- time.Now()

	// Give it a chance to be dispatched
	select {
	case <-barrier:
	case <-time.After(time.Second):
	}

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
}

func (s *ConnSuite) TestReadLoopTooBig(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.server.Config().Limits.ControlLine = 4

	go func() {
		cmd := <-s.serverCmds
		switch cmd.(type) {
		case *UnregisterConnCmd:
			cmd := cmd.(*UnregisterConnCmd)
			cmd.Done <- true
		}
	}()

	io.WriteString(s.tcpConn.client, "TESTCOMMAND\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "-ERR 'Protocol Operation size exceeded'")
}

func (s *ConnSuite) TestReadRequest(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	io.WriteString(s.tcpConn.client, "PING\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "PONG")
}

func (s *ConnSuite) TestReadBadRequest(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	io.WriteString(s.tcpConn.client, "PING foo\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "-ERR 'Unknown Protocol Operation'")
}

func (s *ConnSuite) TestReadUnknownRequest(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	io.WriteString(s.tcpConn.client, "HELLO\r\n")

	reader := bufio.NewReader(s.tcpConn.client)
	reader.ReadLine()
	checkReadLine(c, reader, "-ERR 'Unknown Protocol Operation'")
}

func (s *ConnSuite) TestClientClosed(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.tcpConn.client.Close()
	serverCmd := <-s.serverCmds
	(serverCmd.(*UnregisterConnCmd)).Done <- true
}

type TestServerCmd struct{}

func (c *TestServerCmd) Process(Server) {

}

func (s *ConnSuite) TestSendServerCmd(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	s.conn.SendServerCmd(&TestServerCmd{})

	select {
	case <-s.serverCmds:
	case <-time.After(time.Second):
		c.Errorf("Should have received server command")
	}
}

type CallbackClientCmd struct {
	cb func()
}

func (c *CallbackClientCmd) Process(conn Conn) {
	if c.cb != nil {
		c.cb()
	}
}

func (s *ConnSuite) TestSendServerCmdFull(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	barrier := make(chan bool, 1)

	s.conn.ServeCommand(&CallbackClientCmd{cb: func() {
		// Fill up server channel
		s.serverCmds <- &TestServerCmd{}

		go func() {
			// Fill Up connection channel
			sm := &SubscribedMessage{}
			sm.Message = &Message{}
			sm.Subscription = &Subscription{}
			for i := 0; i < MAX_CONN_CHAN_BACKLOG+1; i++ {
				s.conn.ServeMessage(sm)
			}
			// Process a server command to unblock connection from sending
			<-s.serverCmds
		}()

		s.conn.SendServerCmd(&TestServerCmd{})

		select {
		case <-s.serverCmds:
		case <-time.After(time.Second):
			c.Errorf("Should have received server command")
		}

		barrier <- true
	}})

	<-barrier
}

func (s *ConnSuite) TestWriteFull(c *C) {
	s.delegate.Set(c)
	defer s.ctrl.Finish()

	s.server.Config().Limits.Pending = 16
	s.conn = NewConn(s.server, s.tcpConn)
	go s.conn.Start()

	reader := bufio.NewReader(s.tcpConn.client)

	value := "12345678901234567890"

	reader.ReadLine()

	go func() {
		cmd := <-s.serverCmds
		switch cmd.(type) {
		case *UnregisterConnCmd:
			cmd := cmd.(*UnregisterConnCmd)
			cmd.Done <- true
		}
	}()

	go s.conn.Write(NewStringResponse(value))
	checkReadLine(c, reader, "-ERR 'Slow consumer detected, connection dropped'")
}

func checkReadLine(c *C, reader *bufio.Reader, expected string) {
	line, prefix, err := reader.ReadLine()
	c.Check(err, IsNil)
	c.Check(prefix, Equals, false)
	c.Check(string(line), Equals, expected)
}
<|endoftext|>


<commit_before>package google

import (
	"fmt"
	"regexp"
	"strings"

	"google.golang.org/api/googleapi"
)

const (
	resolveImageFamilyRegex = "[-_a-zA-Z0-9]*"
	resolveImageImageRegex  = "[-_a-zA-Z0-9]*"
)

var (
	resolveImageProjectImage           = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamily          = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/family/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageGlobalImage            = regexp.MustCompile(fmt.Sprintf("^global/images/(%s)$", resolveImageImageRegex))
	resolveImageGlobalFamily           = regexp.MustCompile(fmt.Sprintf("^global/images/family/(%s)$", resolveImageFamilyRegex))
	resolveImageFamilyFamily           = regexp.MustCompile(fmt.Sprintf("^family/(%s)$", resolveImageFamilyRegex))
	resolveImageProjectImageShorthand  = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamilyShorthand = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageFamily                 = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageFamilyRegex))
	resolveImageImage                  = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageImageRegex))
	resolveImageLink                   = regexp.MustCompile(fmt.Sprintf("^https://www.googleapis.com/compute/[a-z0-9]+/projects/(%s)/global/images/(%s)", ProjectRegex, resolveImageImageRegex))

	windowsSqlImage         = regexp.MustCompile("^sql-([0-9]{4})-([a-z]+)-windows-([0-9]{4})(?:-r([0-9]+))?-dc-v[0-9]+$")
	canonicalUbuntuLtsImage = regexp.MustCompile("^ubuntu-(minimal-)?([0-9]+)-")
)

// built-in projects to look for images/families containing the string
// on the left in
var imageMap = map[string]string{
	"centos":      "centos-cloud",
	"coreos":      "coreos-cloud",
	"debian":      "debian-cloud",
	"opensuse":    "opensuse-cloud",
	"rhel":        "rhel-cloud",
	"sles":        "suse-cloud",
	"ubuntu":      "ubuntu-os-cloud",
	"windows":     "windows-cloud",
	"windows-sql": "windows-sql-cloud",
}

func resolveImageImageExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.Get(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if image %s exists: %s", name, err)
	}
}

func resolveImageFamilyExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.GetFromFamily(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if family %s exists: %s", name, err)
	}
}

func sanityTestRegexMatches(expected int, got []string, regexType, name string) error {
	if len(got)-1 != expected { // subtract one, index zero is the entire matched expression
		return fmt.Errorf("Expected %d %s regex matches, got %d for %s", expected, regexType, len(got)-1, name)
	}
	return nil
}

// If the given name is a URL, return it.
// If it's in the form projects/{project}/global/images/{image}, return it
// If it's in the form projects/{project}/global/images/family/{family}, return it
// If it's in the form global/images/{image}, return it
// If it's in the form global/images/family/{family}, return it
// If it's in the form family/{family}, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {project}/{family-or-image}, check if it's an image in the named project. If it is, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the named project. If it is, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {family-or-image}, check if it's an image in the current project. If it is, return it as global/images/{image}.
//    If not, check if it could be a GCP-provided image, and if it exists. If it does, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}
func resolveImage(c *Config, project, name string) (string, error) {
	var builtInProject string
	for k, v := range imageMap {
		if strings.Contains(name, k) {
			builtInProject = v
			break
		}
	}
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/%s", res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/family/%s", res[1]), nil
	case resolveImageFamilyFamily.MatchString(name): // family/xyz
		res := resolveImageFamilyFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	case resolveImageProjectImageShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectImageShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
		}
		fallthrough // check if it's a family
	case resolveImageProjectFamilyShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectFamilyShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
		}
	case resolveImageImage.MatchString(name): // xyz
		res := resolveImageImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "image", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the images GCP provides
			if ok, err := resolveImageImageExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/%s", builtInProject, res[1]), nil
			}
		}
		fallthrough // check if the name is a family, instead of an image
	case resolveImageFamily.MatchString(name): // xyz
		res := resolveImageFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the families GCP provides
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	}
	return "", fmt.Errorf("Could not find image or family %s", name)
}

// resolvedImageSelfLink takes the output of resolveImage and coerces it into a self_link.
// In the event that a global/images/IMAGE or global/images/family/FAMILY reference is
// returned from resolveImage, providerProject will be used as the project for the self_link.
func resolvedImageSelfLink(providerProject, name string) (string, error) {
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/%s", providerProject, res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/family/%s", providerProject, res[1]), nil
	}
	return "", fmt.Errorf("Could not expand image or family %q into a self_link", name)

}
<commit_msg>add server options to windows sql image (#1053)<commit_after>package google

import (
	"fmt"
	"regexp"
	"strings"

	"google.golang.org/api/googleapi"
)

const (
	resolveImageFamilyRegex = "[-_a-zA-Z0-9]*"
	resolveImageImageRegex  = "[-_a-zA-Z0-9]*"
)

var (
	resolveImageProjectImage           = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamily          = regexp.MustCompile(fmt.Sprintf("projects/(%s)/global/images/family/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageGlobalImage            = regexp.MustCompile(fmt.Sprintf("^global/images/(%s)$", resolveImageImageRegex))
	resolveImageGlobalFamily           = regexp.MustCompile(fmt.Sprintf("^global/images/family/(%s)$", resolveImageFamilyRegex))
	resolveImageFamilyFamily           = regexp.MustCompile(fmt.Sprintf("^family/(%s)$", resolveImageFamilyRegex))
	resolveImageProjectImageShorthand  = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageImageRegex))
	resolveImageProjectFamilyShorthand = regexp.MustCompile(fmt.Sprintf("^(%s)/(%s)$", ProjectRegex, resolveImageFamilyRegex))
	resolveImageFamily                 = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageFamilyRegex))
	resolveImageImage                  = regexp.MustCompile(fmt.Sprintf("^(%s)$", resolveImageImageRegex))
	resolveImageLink                   = regexp.MustCompile(fmt.Sprintf("^https://www.googleapis.com/compute/[a-z0-9]+/projects/(%s)/global/images/(%s)", ProjectRegex, resolveImageImageRegex))

	windowsSqlImage         = regexp.MustCompile("^sql-(?:server-)?([0-9]{4})-([a-z]+)-windows-(?:server-)?([0-9]{4})(?:-r([0-9]+))?-dc-v[0-9]+$")
	canonicalUbuntuLtsImage = regexp.MustCompile("^ubuntu-(minimal-)?([0-9]+)-")
)

// built-in projects to look for images/families containing the string
// on the left in
var imageMap = map[string]string{
	"centos":      "centos-cloud",
	"coreos":      "coreos-cloud",
	"debian":      "debian-cloud",
	"opensuse":    "opensuse-cloud",
	"rhel":        "rhel-cloud",
	"sles":        "suse-cloud",
	"ubuntu":      "ubuntu-os-cloud",
	"windows":     "windows-cloud",
	"windows-sql": "windows-sql-cloud",
}

func resolveImageImageExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.Get(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if image %s exists: %s", name, err)
	}
}

func resolveImageFamilyExists(c *Config, project, name string) (bool, error) {
	if _, err := c.clientCompute.Images.GetFromFamily(project, name).Do(); err == nil {
		return true, nil
	} else if gerr, ok := err.(*googleapi.Error); ok && gerr.Code == 404 {
		return false, nil
	} else {
		return false, fmt.Errorf("Error checking if family %s exists: %s", name, err)
	}
}

func sanityTestRegexMatches(expected int, got []string, regexType, name string) error {
	if len(got)-1 != expected { // subtract one, index zero is the entire matched expression
		return fmt.Errorf("Expected %d %s regex matches, got %d for %s", expected, regexType, len(got)-1, name)
	}
	return nil
}

// If the given name is a URL, return it.
// If it's in the form projects/{project}/global/images/{image}, return it
// If it's in the form projects/{project}/global/images/family/{family}, return it
// If it's in the form global/images/{image}, return it
// If it's in the form global/images/family/{family}, return it
// If it's in the form family/{family}, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {project}/{family-or-image}, check if it's an image in the named project. If it is, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the named project. If it is, return it as projects/{project}/global/images/family/{family}.
// If it's in the form {family-or-image}, check if it's an image in the current project. If it is, return it as global/images/{image}.
//    If not, check if it could be a GCP-provided image, and if it exists. If it does, return it as projects/{project}/global/images/{image}.
//    If not, check if it's a family in the current project. If it is, return it as global/images/family/{family}.
//    If not, check if it could be a GCP-provided family, and if it exists. If it does, return it as projects/{project}/global/images/family/{family}
func resolveImage(c *Config, project, name string) (string, error) {
	var builtInProject string
	for k, v := range imageMap {
		if strings.Contains(name, k) {
			builtInProject = v
			break
		}
	}
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/%s", res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("global/images/family/%s", res[1]), nil
	case resolveImageFamilyFamily.MatchString(name): // family/xyz
		res := resolveImageFamilyFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	case resolveImageProjectImageShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectImageShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/%s", res[1], res[2]), nil
		}
		fallthrough // check if it's a family
	case resolveImageProjectFamilyShorthand.MatchString(name): // xyz/xyz
		res := resolveImageProjectFamilyShorthand.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family shorthand", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, res[1], res[2]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("projects/%s/global/images/family/%s", res[1], res[2]), nil
		}
	case resolveImageImage.MatchString(name): // xyz
		res := resolveImageImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "image", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageImageExists(c, project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the images GCP provides
			if ok, err := resolveImageImageExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/%s", builtInProject, res[1]), nil
			}
		}
		fallthrough // check if the name is a family, instead of an image
	case resolveImageFamily.MatchString(name): // xyz
		res := resolveImageFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "family", name); err != nil {
			return "", err
		}
		if ok, err := resolveImageFamilyExists(c, c.Project, res[1]); err != nil {
			return "", err
		} else if ok {
			return fmt.Sprintf("global/images/family/%s", res[1]), nil
		}
		if builtInProject != "" {
			// check the families GCP provides
			if ok, err := resolveImageFamilyExists(c, builtInProject, res[1]); err != nil {
				return "", err
			} else if ok {
				return fmt.Sprintf("projects/%s/global/images/family/%s", builtInProject, res[1]), nil
			}
		}
	}
	return "", fmt.Errorf("Could not find image or family %s", name)
}

// resolvedImageSelfLink takes the output of resolveImage and coerces it into a self_link.
// In the event that a global/images/IMAGE or global/images/family/FAMILY reference is
// returned from resolveImage, providerProject will be used as the project for the self_link.
func resolvedImageSelfLink(providerProject, name string) (string, error) {
	switch {
	case resolveImageLink.MatchString(name): // https://www.googleapis.com/compute/v1/projects/xyz/global/images/xyz
		return name, nil
	case resolveImageProjectImage.MatchString(name): // projects/xyz/global/images/xyz
		res := resolveImageProjectImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/%s", res[1], res[2]), nil
	case resolveImageProjectFamily.MatchString(name): // projects/xyz/global/images/family/xyz
		res := resolveImageProjectFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(2, res, "project family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/family/%s", res[1], res[2]), nil
	case resolveImageGlobalImage.MatchString(name): // global/images/xyz
		res := resolveImageGlobalImage.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global image", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/%s", providerProject, res[1]), nil
	case resolveImageGlobalFamily.MatchString(name): // global/images/family/xyz
		res := resolveImageGlobalFamily.FindStringSubmatch(name)
		if err := sanityTestRegexMatches(1, res, "global family", name); err != nil {
			return "", err
		}
		return fmt.Sprintf("https://www.googleapis.com/compute/v1/projects/%s/global/images/family/%s", providerProject, res[1]), nil
	}
	return "", fmt.Errorf("Could not expand image or family %q into a self_link", name)

}
<|endoftext|>


<commit_before>package shadowsocks

import (
	"crypto/md5"
	"io"
	"encoding/binary"
	"bytes"
)

func GetTable(key string) (encryptTable []byte, decryptTable []byte) {
	encryptTable = make([]byte, 256)
	decryptTable = make([]byte, 256)
	table := make([]uint64, 256)

	h := md5.New()
	io.WriteString(h, key)

	s := h.Sum(nil)

	var a uint64
	buf := bytes.NewBuffer(s)
	binary.Read(buf, binary.LittleEndian, &a)
	var i uint64
	for i = 0; i < 256; i++ {
		table[i] = i
	}
	for i = 1; i < 1024; i++ {
		table = Sort(table, func(x, y uint64) int64 {
				return int64(a%uint64(x + i) - a%uint64(y + i))
			})
	}
	for i = 0; i < 256; i++ {
		encryptTable[i] = byte(table[i])
	}
	for i = 0; i < 256; i++ {
		decryptTable[encryptTable[i]] = byte(i)
	}

	return
}

func Encrypt(table []byte, buf []byte) []byte {
	var result = make([]byte, len(buf))
	for i := 0; i < len(buf); i++ {
		result[i] = table[buf[i]]
	}
	return result
}
<commit_msg>Add Encrypt2 function with given src and result buffer.<commit_after>package shadowsocks

import (
	"bytes"
	"crypto/md5"
	"encoding/binary"
	"io"
)

func GetTable(key string) (encryptTable []byte, decryptTable []byte) {
	encryptTable = make([]byte, 256)
	decryptTable = make([]byte, 256)
	table := make([]uint64, 256)

	h := md5.New()
	io.WriteString(h, key)

	s := h.Sum(nil)

	var a uint64
	buf := bytes.NewBuffer(s)
	binary.Read(buf, binary.LittleEndian, &a)
	var i uint64
	for i = 0; i < 256; i++ {
		table[i] = i
	}
	for i = 1; i < 1024; i++ {
		table = Sort(table, func(x, y uint64) int64 {
			return int64(a%uint64(x+i) - a%uint64(y+i))
		})
	}
	for i = 0; i < 256; i++ {
		encryptTable[i] = byte(table[i])
	}
	for i = 0; i < 256; i++ {
		decryptTable[encryptTable[i]] = byte(i)
	}

	return
}

func Encrypt2(table []byte, buf, result []byte) {
	for i := 0; i < len(buf); i++ {
		result[i] = table[buf[i]]
	}
}

func Encrypt(table []byte, buf []byte) []byte {
	var result = make([]byte, len(buf), len(buf))
	Encrypt2(table, buf, result)
	return result
}
<|endoftext|>