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

package check

import (
	"bytes"
	"fmt"
	"math/big"
	"reflect"
	"sort"
	"strings"
	"testing"

	"github.com/google/wuffs/lang/builtin"
	"github.com/google/wuffs/lang/parse"
	"github.com/google/wuffs/lang/render"

	a "github.com/google/wuffs/lang/ast"
	t "github.com/google/wuffs/lang/token"
)

func compareToWuffsfmt(tm *t.Map, tokens []t.Token, comments []string, src string) error {
	buf := &bytes.Buffer{}
	if err := render.Render(buf, tm, tokens, comments); err != nil {
		return err
	}
	got := strings.Split(buf.String(), "\n")
	want := strings.Split(src, "\n")
	for line := 1; len(got) != 0 || len(want) != 0; line++ {
		if len(got) == 0 {
			w := strings.TrimSpace(want[0])
			return fmt.Errorf("difference at line %d:\n%s\n%s", line, "", w)
		}
		if len(want) == 0 {
			g := strings.TrimSpace(got[0])
			return fmt.Errorf("difference at line %d:\n%s\n%s", line, g, "")
		}
		if g, w := strings.TrimSpace(got[0]), strings.TrimSpace(want[0]); g != w {
			return fmt.Errorf("difference at line %d:\n%s\n%s", line, g, w)
		}
		got = got[1:]
		want = want[1:]
	}
	return nil
}

func TestCheck(tt *testing.T) {
	const filename = "test.wuffs"
	src := strings.TrimSpace(`
		packageid "test"

		pri struct foo(
			i i32,
		)

		pri func foo.bar()() {
			var x u8
			var y i32 = +2
			var z u64[..123]
			var a[4] u8
			var b bool

			x = 0
			x = 1 + (x * 0)
			y = -y - 1
			y = this.i
			b = not true

			y = x as i32

			var p i32
			var q i32[0..8]

			assert true

			while:label p == q,
				pre true,
				inv true,
				post p != q,
			{
				// Redundant, but shows the labeled jump syntax.
				continue:label
			}
		}
	`) + "\n"

	tm := &t.Map{}

	tokens, comments, err := t.Tokenize(tm, filename, []byte(src))
	if err != nil {
		tt.Fatalf("Tokenize: %v", err)
	}

	file, err := parse.Parse(tm, filename, tokens, nil)
	if err != nil {
		tt.Fatalf("Parse: %v", err)
	}

	if err := compareToWuffsfmt(tm, tokens, comments, src); err != nil {
		tt.Fatalf("compareToWuffsfmt: %v", err)
	}

	c, err := Check(tm, []*a.File{file}, nil)
	if err != nil {
		tt.Fatalf("Check: %v", err)
	}

	funcs := c.Funcs()
	if len(funcs) != 1 {
		tt.Fatalf("Funcs: got %d elements, want 1", len(funcs))
	}
	fooBar := Func{}
	for _, f := range funcs {
		fooBar = f
		break
	}

	if got, want := fooBar.QQID.Str(tm), "foo.bar"; got != want {
		tt.Fatalf("Funcs[0] name: got %q, want %q", got, want)
	}

	got := [][2]string(nil)
	for id, typ := range fooBar.LocalVars {
		got = append(got, [2]string{
			id.Str(tm),
			typ.Str(tm),
		})
	}
	sort.Slice(got, func(i, j int) bool {
		return got[i][0] < got[j][0]
	})

	want := [][2]string{
		{"a", "[4] u8"},
		{"b", "bool"},
		{"in", "in"},
		{"out", "out"},
		{"p", "i32"},
		{"q", "i32[0..8]"},
		{"this", "ptr foo"},
		{"x", "u8"},
		{"y", "i32"},
		{"z", "u64[..123]"},
	}
	if !reflect.DeepEqual(got, want) {
		tt.Fatalf("\ngot  %v\nwant %v", got, want)
	}
}

func TestConstValues(tt *testing.T) {
	const filename = "test.wuffs"
	testCases := map[string]int64{
		"var i i32 = 42": 42,

		"var i i32 = +7": +7,
		"var i i32 = -7": -7,

		"var b bool = false":         0,
		"var b bool = not false":     1,
		"var b bool = not not false": 0,

		"var i i32 = 10  + 3": 13,
		"var i i32 = 10  - 3": 7,
		"var i i32 = 10  * 3": 30,
		"var i i32 = 10  / 3": 3,
		"var i i32 = 10 << 3": 80,
		"var i i32 = 10 >> 3": 1,
		"var i i32 = 10  & 3": 2,
		"var i i32 = 10 &^ 3": 8,
		"var i i32 = 10  | 3": 11,
		"var i i32 = 10  ^ 3": 9,

		"var b bool = 10 != 3": 1,
		"var b bool = 10  < 3": 0,
		"var b bool = 10 <= 3": 0,
		"var b bool = 10 == 3": 0,
		"var b bool = 10 >= 3": 1,
		"var b bool = 10  > 3": 1,

		"var b bool = false and true": 0,
		"var b bool = false  or true": 1,
	}

	tm := &t.Map{}
	for s, wantInt64 := range testCases {
		src := "packageid \"test\"\npri func foo()() {\n\t" + s + "\n}\n"

		tokens, _, err := t.Tokenize(tm, filename, []byte(src))
		if err != nil {
			tt.Errorf("%q: Tokenize: %v", s, err)
			continue
		}

		file, err := parse.Parse(tm, filename, tokens, nil)
		if err != nil {
			tt.Errorf("%q: Parse: %v", s, err)
			continue
		}

		c, err := Check(tm, []*a.File{file}, nil)
		if err != nil {
			tt.Errorf("%q: Check: %v", s, err)
			continue
		}

		funcs := c.Funcs()
		if len(funcs) != 1 {
			tt.Errorf("%q: Funcs: got %d elements, want 1", s, len(funcs))
			continue
		}
		foo := Func{}
		for _, f := range funcs {
			foo = f
			break
		}
		body := foo.Func.Body()
		if len(body) != 1 {
			tt.Errorf("%q: Body: got %d elements, want 1", s, len(body))
			continue
		}
		if body[0].Kind() != a.KVar {
			tt.Errorf("%q: Body[0]: got %s, want %s", s, body[0].Kind(), a.KVar)
			continue
		}

		got := body[0].Var().Value().ConstValue()
		want := big.NewInt(wantInt64)
		if got == nil || want == nil || got.Cmp(want) != 0 {
			tt.Errorf("%q: got %v, want %v", s, got, want)
			continue
		}
	}
}

func TestBitMask(tt *testing.T) {
	testCases := [][2]uint64{
		{0, 0},
		{1, 1},
		{2, 3},
		{3, 3},
		{4, 7},
		{5, 7},
		{50, 63},
		{500, 511},
		{5000, 8191},
		{50000, 65535},
		{500000, 524287},
		{1<<7 - 2, 1<<7 - 1},
		{1<<7 - 1, 1<<7 - 1},
		{1<<7 + 0, 1<<8 - 1},
		{1<<7 + 1, 1<<8 - 1},
		{1<<8 - 2, 1<<8 - 1},
		{1<<8 - 1, 1<<8 - 1},
		{1<<8 + 0, 1<<9 - 1},
		{1<<8 + 1, 1<<9 - 1},
		{1<<32 - 2, 1<<32 - 1},
		{1<<32 - 1, 1<<32 - 1},
		{1<<32 + 0, 1<<33 - 1},
		{1<<32 + 1, 1<<33 - 1},
		{1<<64 - 2, 1<<64 - 1},
		{1<<64 - 1, 1<<64 - 1},
	}

	for _, tc := range testCases {
		got := bitMask(big.NewInt(0).SetUint64(tc[0]).BitLen())
		want := big.NewInt(0).SetUint64(tc[1])
		if got.Cmp(want) != 0 {
			tt.Errorf("roundUpToPowerOf2Minus1(%v): got %v, want %v", tc[0], got, tc[1])
		}
	}
}

func TestBuiltInTypeMap(tt *testing.T) {
	if got, want := len(builtInTypeMap), len(builtin.Types); got != want {
		tt.Fatalf("lengths: got %d, want %d", got, want)
	}
	tm := t.Map{}
	for _, s := range builtin.Types {
		id := tm.ByName(s)
		if id == 0 {
			tt.Fatalf("ID(%q): got 0, want non-0", s)
		}
		if _, ok := builtInTypeMap[id]; !ok {
			tt.Fatalf("no builtInTypeMap entry for %q", s)
		}
	}
}
<commit_msg>Fix lang/check tests<commit_after>// Copyright 2017 The Wuffs Authors.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package check

import (
	"bytes"
	"fmt"
	"math/big"
	"reflect"
	"sort"
	"strings"
	"testing"

	"github.com/google/wuffs/lang/builtin"
	"github.com/google/wuffs/lang/parse"
	"github.com/google/wuffs/lang/render"

	a "github.com/google/wuffs/lang/ast"
	t "github.com/google/wuffs/lang/token"
)

func compareToWuffsfmt(tm *t.Map, tokens []t.Token, comments []string, src string) error {
	buf := &bytes.Buffer{}
	if err := render.Render(buf, tm, tokens, comments); err != nil {
		return err
	}
	got := strings.Split(buf.String(), "\n")
	want := strings.Split(src, "\n")
	for line := 1; len(got) != 0 || len(want) != 0; line++ {
		if len(got) == 0 {
			w := strings.TrimSpace(want[0])
			return fmt.Errorf("difference at line %d:\n%s\n%s", line, "", w)
		}
		if len(want) == 0 {
			g := strings.TrimSpace(got[0])
			return fmt.Errorf("difference at line %d:\n%s\n%s", line, g, "")
		}
		if g, w := strings.TrimSpace(got[0]), strings.TrimSpace(want[0]); g != w {
			return fmt.Errorf("difference at line %d:\n%s\n%s", line, g, w)
		}
		got = got[1:]
		want = want[1:]
	}
	return nil
}

func TestCheck(tt *testing.T) {
	const filename = "test.wuffs"
	src := strings.TrimSpace(`
		packageid "test"

		pri struct foo(
			i i32,
		)

		pri func foo.bar()() {
			var x u8
			var y i32 = +2
			var z u64[..123]
			var a[4] u8
			var b bool

			x = 0
			x = 1 + (x * 0)
			y = -y - 1
			y = this.i
			b = not true

			y = x as i32

			var p i32
			var q i32[0..8]

			assert true

			while:label p == q,
				pre true,
				inv true,
				post p != q,
			{
				// Redundant, but shows the labeled jump syntax.
				continue:label
			}
		}
	`) + "\n"

	tm := &t.Map{}

	tokens, comments, err := t.Tokenize(tm, filename, []byte(src))
	if err != nil {
		tt.Fatalf("Tokenize: %v", err)
	}

	file, err := parse.Parse(tm, filename, tokens, nil)
	if err != nil {
		tt.Fatalf("Parse: %v", err)
	}

	if err := compareToWuffsfmt(tm, tokens, comments, src); err != nil {
		tt.Fatalf("compareToWuffsfmt: %v", err)
	}

	c, err := Check(tm, []*a.File{file}, nil)
	if err != nil {
		tt.Fatalf("Check: %v", err)
	}

	if len(c.funcs) != 1 {
		tt.Fatalf("Funcs: got %d elements, want 1", len(c.funcs))
	}
	fooBar := (*a.Func)(nil)
	fooBarLocalVars := typeMap(nil)
	for _, f := range c.funcs {
		fooBar = f
		break
	}
	for _, v := range c.localVars {
		fooBarLocalVars = v
		break
	}

	if got, want := fooBar.QQID().Str(tm), "foo.bar"; got != want {
		tt.Fatalf("Funcs[0] name: got %q, want %q", got, want)
	}

	got := [][2]string(nil)
	for id, typ := range fooBarLocalVars {
		got = append(got, [2]string{
			id.Str(tm),
			typ.Str(tm),
		})
	}
	sort.Slice(got, func(i, j int) bool {
		return got[i][0] < got[j][0]
	})

	want := [][2]string{
		{"a", "[4] u8"},
		{"b", "bool"},
		{"in", "in"},
		{"out", "out"},
		{"p", "i32"},
		{"q", "i32[0..8]"},
		{"this", "ptr foo"},
		{"x", "u8"},
		{"y", "i32"},
		{"z", "u64[..123]"},
	}
	if !reflect.DeepEqual(got, want) {
		tt.Fatalf("\ngot  %v\nwant %v", got, want)
	}
}

func TestConstValues(tt *testing.T) {
	const filename = "test.wuffs"
	testCases := map[string]int64{
		"var i i32 = 42": 42,

		"var i i32 = +7": +7,
		"var i i32 = -7": -7,

		"var b bool = false":         0,
		"var b bool = not false":     1,
		"var b bool = not not false": 0,

		"var i i32 = 10  + 3": 13,
		"var i i32 = 10  - 3": 7,
		"var i i32 = 10  * 3": 30,
		"var i i32 = 10  / 3": 3,
		"var i i32 = 10 << 3": 80,
		"var i i32 = 10 >> 3": 1,
		"var i i32 = 10  & 3": 2,
		"var i i32 = 10 &^ 3": 8,
		"var i i32 = 10  | 3": 11,
		"var i i32 = 10  ^ 3": 9,

		"var b bool = 10 != 3": 1,
		"var b bool = 10  < 3": 0,
		"var b bool = 10 <= 3": 0,
		"var b bool = 10 == 3": 0,
		"var b bool = 10 >= 3": 1,
		"var b bool = 10  > 3": 1,

		"var b bool = false and true": 0,
		"var b bool = false  or true": 1,
	}

	tm := &t.Map{}
	for s, wantInt64 := range testCases {
		src := "packageid \"test\"\npri func foo()() {\n\t" + s + "\n}\n"

		tokens, _, err := t.Tokenize(tm, filename, []byte(src))
		if err != nil {
			tt.Errorf("%q: Tokenize: %v", s, err)
			continue
		}

		file, err := parse.Parse(tm, filename, tokens, nil)
		if err != nil {
			tt.Errorf("%q: Parse: %v", s, err)
			continue
		}

		c, err := Check(tm, []*a.File{file}, nil)
		if err != nil {
			tt.Errorf("%q: Check: %v", s, err)
			continue
		}

		if len(c.funcs) != 1 {
			tt.Errorf("%q: Funcs: got %d elements, want 1", s, len(c.funcs))
			continue
		}
		foo := (*a.Func)(nil)
		for _, f := range c.funcs {
			foo = f
			break
		}
		body := foo.Body()
		if len(body) != 1 {
			tt.Errorf("%q: Body: got %d elements, want 1", s, len(body))
			continue
		}
		if body[0].Kind() != a.KVar {
			tt.Errorf("%q: Body[0]: got %s, want %s", s, body[0].Kind(), a.KVar)
			continue
		}

		got := body[0].Var().Value().ConstValue()
		want := big.NewInt(wantInt64)
		if got == nil || want == nil || got.Cmp(want) != 0 {
			tt.Errorf("%q: got %v, want %v", s, got, want)
			continue
		}
	}
}

func TestBitMask(tt *testing.T) {
	testCases := [][2]uint64{
		{0, 0},
		{1, 1},
		{2, 3},
		{3, 3},
		{4, 7},
		{5, 7},
		{50, 63},
		{500, 511},
		{5000, 8191},
		{50000, 65535},
		{500000, 524287},
		{1<<7 - 2, 1<<7 - 1},
		{1<<7 - 1, 1<<7 - 1},
		{1<<7 + 0, 1<<8 - 1},
		{1<<7 + 1, 1<<8 - 1},
		{1<<8 - 2, 1<<8 - 1},
		{1<<8 - 1, 1<<8 - 1},
		{1<<8 + 0, 1<<9 - 1},
		{1<<8 + 1, 1<<9 - 1},
		{1<<32 - 2, 1<<32 - 1},
		{1<<32 - 1, 1<<32 - 1},
		{1<<32 + 0, 1<<33 - 1},
		{1<<32 + 1, 1<<33 - 1},
		{1<<64 - 2, 1<<64 - 1},
		{1<<64 - 1, 1<<64 - 1},
	}

	for _, tc := range testCases {
		got := bitMask(big.NewInt(0).SetUint64(tc[0]).BitLen())
		want := big.NewInt(0).SetUint64(tc[1])
		if got.Cmp(want) != 0 {
			tt.Errorf("roundUpToPowerOf2Minus1(%v): got %v, want %v", tc[0], got, tc[1])
		}
	}
}

func TestBuiltInTypeMap(tt *testing.T) {
	if got, want := len(builtInTypeMap), len(builtin.Types); got != want {
		tt.Fatalf("lengths: got %d, want %d", got, want)
	}
	tm := t.Map{}
	for _, s := range builtin.Types {
		id := tm.ByName(s)
		if id == 0 {
			tt.Fatalf("ID(%q): got 0, want non-0", s)
		}
		if _, ok := builtInTypeMap[id]; !ok {
			tt.Fatalf("no builtInTypeMap entry for %q", s)
		}
	}
}
<|endoftext|>


<commit_before>package configcommands

import (
	"fmt"
	"strconv"

	"github.com/pivotalservices/cf-mgmt/config"
)

//BaseConfigCommand - commmand that specifies config-dir
type BaseConfigCommand struct {
	ConfigDirectory string `long:"config-dir" env:"CONFIG_DIR" default:"config" description:"Name of the config directory"`
}

type UserRole struct {
	UserRoleAdd
	LDAPUsersToRemove  []string `long:"ldap-user-to-remove" description:"Ldap User to remove, specify multiple times"`
	UsersToRemove      []string `long:"user-to-remove" description:"User to remove, specify multiple times"`
	SamlUsersToRemove  []string `long:"saml-user-to-remove" description:"SAML user to remove, specify multiple times"`
	LDAPGroupsToRemove []string `long:"ldap-group-to-remove" description:"Group to remove, specify multiple times"`
}

type UserRoleAdd struct {
	LDAPUsers  []string `long:"ldap-user" description:"Ldap User to add, specify multiple times"`
	Users      []string `long:"user" description:"User to add, specify multiple times"`
	SamlUsers  []string `long:"saml-user" description:"SAML user to add, specify multiple times"`
	LDAPGroups []string `long:"ldap-group" description:"Group to add, specify multiple times"`
}

type OrgQuota struct {
	EnableOrgQuota          string `long:"enable-org-quota" description:"Enable the Org Quota in the config" choice:"true" choice:"false"`
	MemoryLimit             string `long:"memory-limit" description:"An Org's memory limit in Megabytes"`
	InstanceMemoryLimit     string `long:"instance-memory-limit" description:"Global Org Application instance memory limit in Megabytes"`
	TotalRoutes             string `long:"total-routes" description:"Total Routes capacity for an Org"`
	TotalServices           string `long:"total-services" description:"Total Services capacity for an Org"`
	PaidServicesAllowed     string `long:"paid-service-plans-allowed" description:"Allow paid services to appear in an org" choice:"true" choice:"false"`
	TotalPrivateDomains     string `long:"total-private-domains" description:"Total Private Domain capacity for an Org"`
	TotalReservedRoutePorts string `long:"total-reserved-route-ports" description:"Total Reserved Route Ports capacity for an Org"`
	TotalServiceKeys        string `long:"total-service-keys" description:"Total Service Keys capacity for an Org"`
	AppInstanceLimit        string `long:"app-instance-limit" description:"Total Service Keys capacity for an Org"`
}

type SpaceQuota struct {
	EnableSpaceQuota        string `long:"enable-space-quota" description:"Enable the Space Quota in the config" choice:"true" choice:"false"`
	MemoryLimit             string `long:"memory-limit" description:"An Space's memory limit in Megabytes"`
	InstanceMemoryLimit     string `long:"instance-memory-limit" description:"Space Application instance memory limit in Megabytes"`
	TotalRoutes             string `long:"total-routes" description:"Total Routes capacity for an Space"`
	TotalServices           string `long:"total-services" description:"Total Services capacity for an Space"`
	PaidServicesAllowed     string `long:"paid-service-plans-allowed" description:"Allow paid services to appear in an Space" choice:"true" choice:"false"`
	TotalPrivateDomains     string `long:"total-private-domains" description:"Total Private Domain capacity for an Space"`
	TotalReservedRoutePorts string `long:"total-reserved-route-ports" description:"Total Reserved Route Ports capacity for an Space"`
	TotalServiceKeys        string `long:"total-service-keys" description:"Total Service Keys capacity for an Space"`
	AppInstanceLimit        string `long:"app-instance-limit" description:"Total Service Keys capacity for an Space"`
}

func updateUsersBasedOnRole(userMgmt *config.UserMgmt, currentLDAPGroups []string, userRole *UserRole, errorString *string) {
	userMgmt.LDAPGroups = removeFromSlice(addToSlice(currentLDAPGroups, userRole.LDAPGroups, errorString), userRole.LDAPGroupsToRemove)
	userMgmt.Users = removeFromSlice(addToSlice(userMgmt.Users, userRole.Users, errorString), userRole.UsersToRemove)
	userMgmt.SamlUsers = removeFromSlice(addToSlice(userMgmt.SamlUsers, userRole.SamlUsers, errorString), userRole.SamlUsersToRemove)
	userMgmt.LDAPUsers = removeFromSlice(addToSlice(userMgmt.LDAPUsers, userRole.LDAPUsers, errorString), userRole.LDAPUsersToRemove)
	userMgmt.LDAPGroup = ""
}

func addUsersBasedOnRole(userMgmt *config.UserMgmt, currentLDAPGroups []string, userRole *UserRoleAdd, errorString *string) {
	userMgmt.LDAPGroups = addToSlice(currentLDAPGroups, userRole.LDAPGroups, errorString)
	userMgmt.Users = addToSlice(userMgmt.Users, userRole.Users, errorString)
	userMgmt.SamlUsers = addToSlice(userMgmt.SamlUsers, userRole.SamlUsers, errorString)
	userMgmt.LDAPUsers = addToSlice(userMgmt.LDAPUsers, userRole.LDAPUsers, errorString)
	userMgmt.LDAPGroup = ""
}

func convertToInt(parameterName string, currentValue *int, proposedValue string, errorString *string) {
	if proposedValue == "" {
		return
	}
	i, err := strconv.Atoi(proposedValue)
	if err != nil {
		*errorString += fmt.Sprintf("\n--%s must be an integer instead of [%s]", parameterName, proposedValue)
		return
	}
	*currentValue = i
}

func convertToBool(parameterName string, currentValue *bool, proposedValue string, errorString *string) {
	if proposedValue == "" {
		return
	}
	b, err := strconv.ParseBool(proposedValue)
	if err != nil {
		*errorString += fmt.Sprintf("\n--%s must be an boolean instead of [%s]", parameterName, proposedValue)
		return
	}
	*currentValue = b
}

func addToSlice(theSlice, sliceToAdd []string, errorString *string) []string {
	checkForDuplicates(sliceToAdd, errorString)
	sliceToReturn := theSlice
	valuesThatExist := sliceToMap(theSlice)
	for _, val := range sliceToAdd {
		if _, ok := valuesThatExist[val]; !ok && val != "" {
			sliceToReturn = append(sliceToReturn, val)
		} else {
			*errorString += fmt.Sprintf("\n--value [%s] already exists in %v", val, theSlice)
		}
	}
	return sliceToReturn
}

func checkForDuplicates(slice []string, errorString *string) {
	sliceMap := make(map[string]string)
	for _, val := range slice {
		if _, ok := sliceMap[val]; ok && val != "" {
			*errorString += fmt.Sprintf("\n--value [%s] cannot be specified more than once %v", val, slice)
		} else {
			sliceMap[val] = val
		}
	}
}

func removeFromSlice(theSlice, sliceToRemove []string) []string {
	var sliceToReturn []string
	valuesToRemove := sliceToMap(sliceToRemove)
	for _, val := range theSlice {
		if _, ok := valuesToRemove[val]; !ok && val != "" {
			sliceToReturn = append(sliceToReturn, val)
		}
	}
	return sliceToReturn
}

func sliceToMap(theSlice []string) map[string]string {
	theMap := make(map[string]string)
	for _, val := range theSlice {
		theMap[val] = val
	}
	return theMap
}

func updateOrgQuotaConfig(orgConfig *config.OrgConfig, orgQuota OrgQuota, errorString *string) {
	convertToBool("enable-org-quota", &orgConfig.EnableOrgQuota, orgQuota.EnableOrgQuota, errorString)
	convertToInt("memory-limit", &orgConfig.MemoryLimit, orgQuota.MemoryLimit, errorString)
	convertToInt("instance-memory-limit", &orgConfig.InstanceMemoryLimit, orgQuota.InstanceMemoryLimit, errorString)
	convertToInt("total-routes", &orgConfig.TotalRoutes, orgQuota.TotalRoutes, errorString)
	convertToInt("total-services", &orgConfig.TotalServices, orgQuota.TotalServices, errorString)
	convertToBool("paid-service-plans-allowed", &orgConfig.PaidServicePlansAllowed, orgQuota.PaidServicesAllowed, errorString)
	convertToInt("total-private-domains", &orgConfig.TotalPrivateDomains, orgQuota.TotalPrivateDomains, errorString)
	convertToInt("total-reserved-route-ports", &orgConfig.TotalReservedRoutePorts, orgQuota.TotalReservedRoutePorts, errorString)
	convertToInt("total-service-keys", &orgConfig.TotalServiceKeys, orgQuota.TotalServiceKeys, errorString)
	convertToInt("app-instance-limit", &orgConfig.AppInstanceLimit, orgQuota.AppInstanceLimit, errorString)
}

func updateSpaceQuotaConfig(spaceConfig *config.SpaceConfig, spaceQuota SpaceQuota, errorString *string) {
	convertToBool("enable-space-quota", &spaceConfig.EnableSpaceQuota, spaceQuota.EnableSpaceQuota, errorString)
	convertToInt("memory-limit", &spaceConfig.MemoryLimit, spaceQuota.MemoryLimit, errorString)
	convertToInt("instance-memory-limit", &spaceConfig.InstanceMemoryLimit, spaceQuota.InstanceMemoryLimit, errorString)
	convertToInt("total-routes", &spaceConfig.TotalRoutes, spaceQuota.TotalRoutes, errorString)
	convertToInt("total-services", &spaceConfig.TotalServices, spaceQuota.TotalServices, errorString)
	convertToBool("paid-service-plans-allowed", &spaceConfig.PaidServicePlansAllowed, spaceQuota.PaidServicesAllowed, errorString)
	convertToInt("total-private-domains", &spaceConfig.TotalPrivateDomains, spaceQuota.TotalPrivateDomains, errorString)
	convertToInt("total-reserved-route-ports", &spaceConfig.TotalReservedRoutePorts, spaceQuota.TotalReservedRoutePorts, errorString)
	convertToInt("total-service-keys", &spaceConfig.TotalServiceKeys, spaceQuota.TotalServiceKeys, errorString)
	convertToInt("app-instance-limit", &spaceConfig.AppInstanceLimit, spaceQuota.AppInstanceLimit, errorString)
}

func validateASGsExist(configuredASGs []config.ASGConfig, asgs []string, errorString *string) {
	asgMap := make(map[string]string)
	for _, configuredASG := range configuredASGs {
		asgMap[configuredASG.Name] = configuredASG.Name
	}
	for _, asg := range asgs {
		if _, ok := asgMap[asg]; !ok {
			*errorString += fmt.Sprintf("\n--[%s.json] does not exist in asgs directory", asg)
		}
	}
}
<commit_msg>update to config command to include app task limit property<commit_after>package configcommands

import (
	"fmt"
	"strconv"

	"github.com/pivotalservices/cf-mgmt/config"
)

//BaseConfigCommand - commmand that specifies config-dir
type BaseConfigCommand struct {
	ConfigDirectory string `long:"config-dir" env:"CONFIG_DIR" default:"config" description:"Name of the config directory"`
}

type UserRole struct {
	UserRoleAdd
	LDAPUsersToRemove  []string `long:"ldap-user-to-remove" description:"Ldap User to remove, specify multiple times"`
	UsersToRemove      []string `long:"user-to-remove" description:"User to remove, specify multiple times"`
	SamlUsersToRemove  []string `long:"saml-user-to-remove" description:"SAML user to remove, specify multiple times"`
	LDAPGroupsToRemove []string `long:"ldap-group-to-remove" description:"Group to remove, specify multiple times"`
}

type UserRoleAdd struct {
	LDAPUsers  []string `long:"ldap-user" description:"Ldap User to add, specify multiple times"`
	Users      []string `long:"user" description:"User to add, specify multiple times"`
	SamlUsers  []string `long:"saml-user" description:"SAML user to add, specify multiple times"`
	LDAPGroups []string `long:"ldap-group" description:"Group to add, specify multiple times"`
}

type OrgQuota struct {
	EnableOrgQuota          string `long:"enable-org-quota" description:"Enable the Org Quota in the config" choice:"true" choice:"false"`
	MemoryLimit             string `long:"memory-limit" description:"An Org's memory limit in Megabytes"`
	InstanceMemoryLimit     string `long:"instance-memory-limit" description:"Global Org Application instance memory limit in Megabytes"`
	TotalRoutes             string `long:"total-routes" description:"Total Routes capacity for an Org"`
	TotalServices           string `long:"total-services" description:"Total Services capacity for an Org"`
	PaidServicesAllowed     string `long:"paid-service-plans-allowed" description:"Allow paid services to appear in an org" choice:"true" choice:"false"`
	TotalPrivateDomains     string `long:"total-private-domains" description:"Total Private Domain capacity for an Org"`
	TotalReservedRoutePorts string `long:"total-reserved-route-ports" description:"Total Reserved Route Ports capacity for an Org"`
	TotalServiceKeys        string `long:"total-service-keys" description:"Total Service Keys capacity for an Org"`
	AppInstanceLimit        string `long:"app-instance-limit" description:"App Instance Limit an Org"`
	AppTaskLimit            string `long:"app-task-limit" description:"App Task Limit an Org"`
}

type SpaceQuota struct {
	EnableSpaceQuota        string `long:"enable-space-quota" description:"Enable the Space Quota in the config" choice:"true" choice:"false"`
	MemoryLimit             string `long:"memory-limit" description:"An Space's memory limit in Megabytes"`
	InstanceMemoryLimit     string `long:"instance-memory-limit" description:"Space Application instance memory limit in Megabytes"`
	TotalRoutes             string `long:"total-routes" description:"Total Routes capacity for an Space"`
	TotalServices           string `long:"total-services" description:"Total Services capacity for an Space"`
	PaidServicesAllowed     string `long:"paid-service-plans-allowed" description:"Allow paid services to appear in an Space" choice:"true" choice:"false"`
	TotalPrivateDomains     string `long:"total-private-domains" description:"Total Private Domain capacity for an Space"`
	TotalReservedRoutePorts string `long:"total-reserved-route-ports" description:"Total Reserved Route Ports capacity for an Space"`
	TotalServiceKeys        string `long:"total-service-keys" description:"Total Service Keys capacity for an Space"`
	AppInstanceLimit        string `long:"app-instance-limit" description:"App Instance Limit for a space"`
	AppTaskLimit            string `long:"app-task-limit" description:"App Task Limit for a space"`
}

func updateUsersBasedOnRole(userMgmt *config.UserMgmt, currentLDAPGroups []string, userRole *UserRole, errorString *string) {
	userMgmt.LDAPGroups = removeFromSlice(addToSlice(currentLDAPGroups, userRole.LDAPGroups, errorString), userRole.LDAPGroupsToRemove)
	userMgmt.Users = removeFromSlice(addToSlice(userMgmt.Users, userRole.Users, errorString), userRole.UsersToRemove)
	userMgmt.SamlUsers = removeFromSlice(addToSlice(userMgmt.SamlUsers, userRole.SamlUsers, errorString), userRole.SamlUsersToRemove)
	userMgmt.LDAPUsers = removeFromSlice(addToSlice(userMgmt.LDAPUsers, userRole.LDAPUsers, errorString), userRole.LDAPUsersToRemove)
	userMgmt.LDAPGroup = ""
}

func addUsersBasedOnRole(userMgmt *config.UserMgmt, currentLDAPGroups []string, userRole *UserRoleAdd, errorString *string) {
	userMgmt.LDAPGroups = addToSlice(currentLDAPGroups, userRole.LDAPGroups, errorString)
	userMgmt.Users = addToSlice(userMgmt.Users, userRole.Users, errorString)
	userMgmt.SamlUsers = addToSlice(userMgmt.SamlUsers, userRole.SamlUsers, errorString)
	userMgmt.LDAPUsers = addToSlice(userMgmt.LDAPUsers, userRole.LDAPUsers, errorString)
	userMgmt.LDAPGroup = ""
}

func convertToInt(parameterName string, currentValue *int, proposedValue string, errorString *string) {
	if proposedValue == "" {
		return
	}
	i, err := strconv.Atoi(proposedValue)
	if err != nil {
		*errorString += fmt.Sprintf("\n--%s must be an integer instead of [%s]", parameterName, proposedValue)
		return
	}
	*currentValue = i
}

func convertToBool(parameterName string, currentValue *bool, proposedValue string, errorString *string) {
	if proposedValue == "" {
		return
	}
	b, err := strconv.ParseBool(proposedValue)
	if err != nil {
		*errorString += fmt.Sprintf("\n--%s must be an boolean instead of [%s]", parameterName, proposedValue)
		return
	}
	*currentValue = b
}

func addToSlice(theSlice, sliceToAdd []string, errorString *string) []string {
	checkForDuplicates(sliceToAdd, errorString)
	sliceToReturn := theSlice
	valuesThatExist := sliceToMap(theSlice)
	for _, val := range sliceToAdd {
		if _, ok := valuesThatExist[val]; !ok && val != "" {
			sliceToReturn = append(sliceToReturn, val)
		} else {
			*errorString += fmt.Sprintf("\n--value [%s] already exists in %v", val, theSlice)
		}
	}
	return sliceToReturn
}

func checkForDuplicates(slice []string, errorString *string) {
	sliceMap := make(map[string]string)
	for _, val := range slice {
		if _, ok := sliceMap[val]; ok && val != "" {
			*errorString += fmt.Sprintf("\n--value [%s] cannot be specified more than once %v", val, slice)
		} else {
			sliceMap[val] = val
		}
	}
}

func removeFromSlice(theSlice, sliceToRemove []string) []string {
	var sliceToReturn []string
	valuesToRemove := sliceToMap(sliceToRemove)
	for _, val := range theSlice {
		if _, ok := valuesToRemove[val]; !ok && val != "" {
			sliceToReturn = append(sliceToReturn, val)
		}
	}
	return sliceToReturn
}

func sliceToMap(theSlice []string) map[string]string {
	theMap := make(map[string]string)
	for _, val := range theSlice {
		theMap[val] = val
	}
	return theMap
}

func updateOrgQuotaConfig(orgConfig *config.OrgConfig, orgQuota OrgQuota, errorString *string) {
	convertToBool("enable-org-quota", &orgConfig.EnableOrgQuota, orgQuota.EnableOrgQuota, errorString)
	convertToInt("memory-limit", &orgConfig.MemoryLimit, orgQuota.MemoryLimit, errorString)
	convertToInt("instance-memory-limit", &orgConfig.InstanceMemoryLimit, orgQuota.InstanceMemoryLimit, errorString)
	convertToInt("total-routes", &orgConfig.TotalRoutes, orgQuota.TotalRoutes, errorString)
	convertToInt("total-services", &orgConfig.TotalServices, orgQuota.TotalServices, errorString)
	convertToBool("paid-service-plans-allowed", &orgConfig.PaidServicePlansAllowed, orgQuota.PaidServicesAllowed, errorString)
	convertToInt("total-private-domains", &orgConfig.TotalPrivateDomains, orgQuota.TotalPrivateDomains, errorString)
	convertToInt("total-reserved-route-ports", &orgConfig.TotalReservedRoutePorts, orgQuota.TotalReservedRoutePorts, errorString)
	convertToInt("total-service-keys", &orgConfig.TotalServiceKeys, orgQuota.TotalServiceKeys, errorString)
	convertToInt("app-instance-limit", &orgConfig.AppInstanceLimit, orgQuota.AppInstanceLimit, errorString)
	convertToInt("app-task-limit", &orgConfig.AppTaskLimit, orgQuota.AppTaskLimit, errorString)
}

func updateSpaceQuotaConfig(spaceConfig *config.SpaceConfig, spaceQuota SpaceQuota, errorString *string) {
	convertToBool("enable-space-quota", &spaceConfig.EnableSpaceQuota, spaceQuota.EnableSpaceQuota, errorString)
	convertToInt("memory-limit", &spaceConfig.MemoryLimit, spaceQuota.MemoryLimit, errorString)
	convertToInt("instance-memory-limit", &spaceConfig.InstanceMemoryLimit, spaceQuota.InstanceMemoryLimit, errorString)
	convertToInt("total-routes", &spaceConfig.TotalRoutes, spaceQuota.TotalRoutes, errorString)
	convertToInt("total-services", &spaceConfig.TotalServices, spaceQuota.TotalServices, errorString)
	convertToBool("paid-service-plans-allowed", &spaceConfig.PaidServicePlansAllowed, spaceQuota.PaidServicesAllowed, errorString)
	convertToInt("total-private-domains", &spaceConfig.TotalPrivateDomains, spaceQuota.TotalPrivateDomains, errorString)
	convertToInt("total-reserved-route-ports", &spaceConfig.TotalReservedRoutePorts, spaceQuota.TotalReservedRoutePorts, errorString)
	convertToInt("total-service-keys", &spaceConfig.TotalServiceKeys, spaceQuota.TotalServiceKeys, errorString)
	convertToInt("app-instance-limit", &spaceConfig.AppInstanceLimit, spaceQuota.AppInstanceLimit, errorString)
	convertToInt("app-task-limit", &spaceConfig.AppTaskLimit, spaceQuota.AppTaskLimit, errorString)
}

func validateASGsExist(configuredASGs []config.ASGConfig, asgs []string, errorString *string) {
	asgMap := make(map[string]string)
	for _, configuredASG := range configuredASGs {
		asgMap[configuredASG.Name] = configuredASG.Name
	}
	for _, asg := range asgs {
		if _, ok := asgMap[asg]; !ok {
			*errorString += fmt.Sprintf("\n--[%s.json] does not exist in asgs directory", asg)
		}
	}
}
<|endoftext|>


<commit_before>package repo_manager

import (
	"bytes"
	"context"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"regexp"
	"strings"

	"go.skia.org/infra/autoroll/go/codereview"
	"go.skia.org/infra/autoroll/go/strategy"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/github"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	GITHUB_COMMIT_MSG_TMPL = `Roll {{.ChildPath}} {{.From}}..{{.To}} ({{.NumCommits}} commits)

{{.ChildRepoCompareUrl}}

git log {{.From}}..{{.To}} --no-merges --oneline
{{.LogStr}}

The AutoRoll server is located here: {{.ServerURL}}

Documentation for the AutoRoller is here:
https://skia.googlesource.com/buildbot/+/master/autoroll/README.md

If the roll is causing failures, please contact the current sheriff ({{.SheriffEmails}}), and stop
the roller if necessary.

`
)

var (
	// Use this function to instantiate a NewGithubRepoManager. This is able to be
	// overridden for testing.
	NewGithubRepoManager func(context.Context, *GithubRepoManagerConfig, string, *github.GitHub, string, string, *http.Client, codereview.CodeReview, bool) (RepoManager, error) = newGithubRepoManager

	githubCommitMsgTmpl = template.Must(template.New("githubCommitMsg").Parse(GITHUB_COMMIT_MSG_TMPL))

	pullRequestInLogRE = regexp.MustCompile(`(?m) \((#[0-9]+)\)$`)
)

// GithubRepoManagerConfig provides configuration for the Github RepoManager.
type GithubRepoManagerConfig struct {
	CommonRepoManagerConfig
	// URL of the parent repo.
	ParentRepoURL string `json:"parentRepoURL"`
	// URL of the child repo.
	ChildRepoURL string `json:"childRepoURL"`
	// The roller will update this file with the child repo's revision.
	RevisionFile string `json:"revisionFile"`
	// The default strategy to use.
	DefaultStrategy string `json:"defaultStrategy"`
	// GS Bucket and template to use if strategy.ROLL_STRATEGY_STORAGE_FILE is used.
	StorageBucket        string   `json:"storageBucket"`
	StoragePathTemplates []string `json:"storagePathTemplates"`
}

// githubRepoManager is a struct used by the autoroller for managing checkouts.
type githubRepoManager struct {
	*commonRepoManager
	githubClient    *github.GitHub
	parentRepo      *git.Checkout
	parentRepoURL   string
	childRepoURL    string
	revisionFile    string
	defaultStrategy string
	gsBucket        string
	gsPathTemplates []string
}

// newGithubRepoManager returns a RepoManager instance which operates in the given
// working directory and updates at the given frequency.
func newGithubRepoManager(ctx context.Context, c *GithubRepoManagerConfig, workdir string, githubClient *github.GitHub, recipeCfgFile, serverURL string, client *http.Client, cr codereview.CodeReview, local bool) (RepoManager, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}
	wd := path.Join(workdir, "github_repos")
	if _, err := os.Stat(wd); err != nil {
		if err := os.MkdirAll(wd, 0755); err != nil {
			return nil, err
		}
	}

	// Create and populate the parent directory if needed.
	_, repo := GetUserAndRepo(c.ParentRepoURL)
	userFork := fmt.Sprintf("git@github.com:%s/%s.git", cr.UserName(), repo)
	parentRepo, err := git.NewCheckout(ctx, userFork, wd)
	if err != nil {
		return nil, err
	}

	crm, err := newCommonRepoManager(c.CommonRepoManagerConfig, wd, serverURL, nil, client, cr, local)
	if err != nil {
		return nil, err
	}

	// Create and populate the child directory if needed.
	if _, err := os.Stat(crm.childDir); err != nil {
		if err := os.MkdirAll(crm.childDir, 0755); err != nil {
			return nil, err
		}
		if _, err := git.GitDir(crm.childDir).Git(ctx, "clone", c.ChildRepoURL, "."); err != nil {
			return nil, err
		}
	}

	gr := &githubRepoManager{
		commonRepoManager: crm,
		githubClient:      githubClient,
		parentRepo:        parentRepo,
		parentRepoURL:     c.ParentRepoURL,
		childRepoURL:      c.ChildRepoURL,
		revisionFile:      c.RevisionFile,
		defaultStrategy:   c.DefaultStrategy,
		gsBucket:          c.StorageBucket,
		gsPathTemplates:   c.StoragePathTemplates,
	}

	return gr, nil
}

// See documentation for RepoManager interface.
func (rm *githubRepoManager) Update(ctx context.Context) error {
	// Sync the projects.
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	// Update the repositories.
	if err := rm.parentRepo.Update(ctx); err != nil {
		return err
	}
	if err := rm.childRepo.Update(ctx); err != nil {
		return err
	}

	// Check to see whether there is an upstream yet.
	remoteOutput, err := rm.parentRepo.Git(ctx, "remote", "show")
	if err != nil {
		return err
	}
	remoteFound := false
	remoteLines := strings.Split(remoteOutput, "\n")
	for _, remoteLine := range remoteLines {
		if remoteLine == GITHUB_UPSTREAM_REMOTE_NAME {
			remoteFound = true
			break
		}
	}
	if !remoteFound {
		if _, err := rm.parentRepo.Git(ctx, "remote", "add", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentRepoURL); err != nil {
			return err
		}
	}
	// Fetch upstream.
	if _, err := rm.parentRepo.Git(ctx, "fetch", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentBranch); err != nil {
		return err
	}

	// Read the contents of the revision file to determine the last roll rev.
	revisionFileContents, err := rm.githubClient.ReadRawFile(rm.parentBranch, rm.revisionFile)
	if err != nil {
		return err
	}
	lastRollRev := strings.TrimRight(revisionFileContents, "\n")

	// Find the number of not-rolled child repo commits.
	notRolled, err := rm.getCommitsNotRolled(ctx, lastRollRev)
	if err != nil {
		return err
	}

	// Get the next roll revision.
	nextRollRev, err := rm.getNextRollRev(ctx, notRolled, lastRollRev)
	if err != nil {
		return err
	}

	rm.infoMtx.Lock()
	defer rm.infoMtx.Unlock()
	rm.lastRollRev = lastRollRev
	rm.nextRollRev = nextRollRev
	rm.commitsNotRolled = len(notRolled)

	sklog.Infof("lastRollRev is: %s", rm.lastRollRev)
	sklog.Infof("nextRollRev is: %s", nextRollRev)
	sklog.Infof("commitsNotRolled: %d", rm.commitsNotRolled)
	return nil
}

func (rm *githubRepoManager) cleanParent(ctx context.Context) error {
	if _, err := rm.parentRepo.Git(ctx, "clean", "-d", "-f", "-f"); err != nil {
		return err
	}
	_, _ = rm.parentRepo.Git(ctx, "rebase", "--abort")
	if _, err := rm.parentRepo.Git(ctx, "checkout", fmt.Sprintf("%s/%s", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentBranch), "-f"); err != nil {
		return err
	}
	_, _ = rm.parentRepo.Git(ctx, "branch", "-D", ROLL_BRANCH)
	return nil
}

// See documentation for RepoManager interface.
func (rm *githubRepoManager) CreateNewRoll(ctx context.Context, from, to string, emails []string, cqExtraTrybots string, dryRun bool) (int64, error) {
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	sklog.Info("Creating a new Github Roll")

	// Clean the checkout, get onto a fresh branch.
	if err := rm.cleanParent(ctx); err != nil {
		return 0, err
	}
	if _, err := rm.parentRepo.Git(ctx, "checkout", fmt.Sprintf("%s/%s", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentBranch), "-b", ROLL_BRANCH); err != nil {
		return 0, err
	}
	// Defer cleanup.
	defer func() {
		util.LogErr(rm.cleanParent(ctx))
	}()

	// Make sure the forked repo is at the same hash as the target repo before
	// creating the pull request.
	if _, err := rm.parentRepo.Git(ctx, "push", "origin", ROLL_BRANCH, "-f"); err != nil {
		return 0, err
	}

	// Make sure the right name and email are set.
	if !rm.local {
		if _, err := rm.parentRepo.Git(ctx, "config", "user.name", rm.codereview.UserName()); err != nil {
			return 0, err
		}
		if _, err := rm.parentRepo.Git(ctx, "config", "user.email", rm.codereview.UserEmail()); err != nil {
			return 0, err
		}
	}

	// Build the commit message.
	user, repo := GetUserAndRepo(rm.childRepoURL)
	githubCompareUrl := fmt.Sprintf("https://github.com/%s/%s/compare/%s...%s", user, repo, from[:12], to[:12])
	logStr, err := rm.childRepo.Git(ctx, "log", fmt.Sprintf("%s..%s", from, to), "--no-merges", "--oneline")
	if err != nil {
		return 0, err
	}
	logStr = strings.TrimSpace(logStr)
	// Github autolinks PR numbers to be of the same repository in logStr. Fix this by
	// explicitly adding the child repo to the PR number.
	logStr = pullRequestInLogRE.ReplaceAllString(logStr, fmt.Sprintf(" (%s/%s$1)", user, repo))

	data := struct {
		ChildPath           string
		ChildRepoCompareUrl string
		From                string
		To                  string
		NumCommits          int
		LogURL              string
		LogStr              string
		ServerURL           string
		SheriffEmails       string
	}{
		ChildPath:           rm.childPath,
		ChildRepoCompareUrl: githubCompareUrl,
		From:                from[:12],
		To:                  to[:12],
		NumCommits:          len(strings.Split(logStr, "\n")),
		LogStr:              logStr,
		ServerURL:           rm.serverURL,
		SheriffEmails:       strings.Join(emails, ","),
	}
	var buf bytes.Buffer
	if err := githubCommitMsgTmpl.Execute(&buf, data); err != nil {
		return 0, err
	}
	commitMsg := buf.String()

	versions, err := rm.childRepo.RevList(ctx, fmt.Sprintf("%s..%s", from, to))
	if err != nil {
		return 0, err
	}
	logStrList := strings.Split(logStr, "\n")
	for i := len(versions) - 1; i >= 0; i-- {
		version := versions[i]
		// Write the file.
		if err := ioutil.WriteFile(path.Join(rm.parentRepo.Dir(), rm.revisionFile), []byte(version+"\n"), os.ModePerm); err != nil {
			return 0, err
		}

		// Commit.
		if _, err := rm.parentRepo.Git(ctx, "commit", "-a", "-m", logStrList[i]); err != nil {
			return 0, err
		}

	}

	// Run the pre-upload steps.
	for _, s := range rm.PreUploadSteps() {
		if err := s(ctx, rm.httpClient, rm.parentRepo.Dir()); err != nil {
			return 0, fmt.Errorf("Error when running pre-upload step: %s", err)
		}
	}

	// Push to the forked repository.
	if _, err := rm.parentRepo.Git(ctx, "push", "origin", ROLL_BRANCH, "-f"); err != nil {
		return 0, err
	}

	// Grab the first line of the commit msg to use as the title of the pull request.
	title := strings.Split(commitMsg, "\n")[0]
	// Use the remaining part of the commit message as the pull request description.
	descComment := strings.Split(commitMsg, "\n")[1:]
	// Create a pull request.
	headBranch := fmt.Sprintf("%s:%s", rm.codereview.UserName(), ROLL_BRANCH)
	pr, err := rm.githubClient.CreatePullRequest(title, rm.parentBranch, headBranch, strings.Join(descComment, "\n"))
	if err != nil {
		return 0, err
	}

	// Add appropriate label to the pull request.
	label := github.COMMIT_LABEL
	if dryRun {
		label = github.DRYRUN_LABEL
	}
	if err := rm.githubClient.AddLabel(pr.GetNumber(), label); err != nil {
		return 0, err
	}

	// Mention the sheriffs on the pull request so that they are automatically
	// subscribed to it.
	mentions := []string{}
	for _, e := range emails {
		m := fmt.Sprintf("@%s", strings.Split(e, "@")[0])
		mentions = append(mentions, m)
	}
	if err := rm.githubClient.AddComment(pr.GetNumber(), fmt.Sprintf("%s : New roll has been created by %s", strings.Join(mentions, " "), rm.serverURL)); err != nil {
		return 0, err
	}

	return int64(pr.GetNumber()), nil
}

func GetUserAndRepo(githubRepo string) (string, string) {
	repoTokens := strings.Split(githubRepo, ":")
	user := strings.Split(repoTokens[1], "/")[0]
	repo := strings.TrimRight(strings.Split(repoTokens[1], "/")[1], ".git")
	return user, repo
}

// See documentation for RepoManager interface.
func (r *githubRepoManager) CreateNextRollStrategy(ctx context.Context, s string) (strategy.NextRollStrategy, error) {
	return strategy.GetNextRollStrategy(ctx, s, r.childBranch, DEFAULT_REMOTE, r.gsBucket, r.gsPathTemplates, r.childRepo, nil)
}

// See documentation for RepoManager interface.
func (r *githubRepoManager) DefaultStrategy() string {
	return r.defaultStrategy
}

// See documentation for RepoManager interface.
func (r *githubRepoManager) ValidStrategies() []string {
	return []string{
		strategy.ROLL_STRATEGY_GCS_FILE,
		strategy.ROLL_STRATEGY_SINGLE,
		strategy.ROLL_STRATEGY_BATCH,
	}
}
<commit_msg>[Autoroller] Fixes to github roller to handle merges and many commits<commit_after>package repo_manager

import (
	"bytes"
	"context"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path"
	"regexp"
	"strings"

	"go.skia.org/infra/autoroll/go/codereview"
	"go.skia.org/infra/autoroll/go/strategy"
	"go.skia.org/infra/go/git"
	"go.skia.org/infra/go/github"
	"go.skia.org/infra/go/sklog"
	"go.skia.org/infra/go/util"
)

const (
	GITHUB_COMMIT_MSG_TMPL = `Roll {{.ChildPath}} {{.From}}..{{.To}} ({{.NumCommits}} commits)

{{.ChildRepoCompareUrl}}

git log {{.From}}..{{.To}} --no-merges --oneline
{{.LogStr}}

The AutoRoll server is located here: {{.ServerURL}}

Documentation for the AutoRoller is here:
https://skia.googlesource.com/buildbot/+/master/autoroll/README.md

If the roll is causing failures, please contact the current sheriff ({{.SheriffEmails}}), and stop
the roller if necessary.

`
)

var (
	// Use this function to instantiate a NewGithubRepoManager. This is able to be
	// overridden for testing.
	NewGithubRepoManager func(context.Context, *GithubRepoManagerConfig, string, *github.GitHub, string, string, *http.Client, codereview.CodeReview, bool) (RepoManager, error) = newGithubRepoManager

	githubCommitMsgTmpl = template.Must(template.New("githubCommitMsg").Parse(GITHUB_COMMIT_MSG_TMPL))

	pullRequestInLogRE = regexp.MustCompile(`(?m) \((#[0-9]+)\)$`)
)

// GithubRepoManagerConfig provides configuration for the Github RepoManager.
type GithubRepoManagerConfig struct {
	CommonRepoManagerConfig
	// URL of the parent repo.
	ParentRepoURL string `json:"parentRepoURL"`
	// URL of the child repo.
	ChildRepoURL string `json:"childRepoURL"`
	// The roller will update this file with the child repo's revision.
	RevisionFile string `json:"revisionFile"`
	// The default strategy to use.
	DefaultStrategy string `json:"defaultStrategy"`
	// GS Bucket and template to use if strategy.ROLL_STRATEGY_STORAGE_FILE is used.
	StorageBucket        string   `json:"storageBucket"`
	StoragePathTemplates []string `json:"storagePathTemplates"`
}

// githubRepoManager is a struct used by the autoroller for managing checkouts.
type githubRepoManager struct {
	*commonRepoManager
	githubClient    *github.GitHub
	parentRepo      *git.Checkout
	parentRepoURL   string
	childRepoURL    string
	revisionFile    string
	defaultStrategy string
	gsBucket        string
	gsPathTemplates []string
}

// newGithubRepoManager returns a RepoManager instance which operates in the given
// working directory and updates at the given frequency.
func newGithubRepoManager(ctx context.Context, c *GithubRepoManagerConfig, workdir string, githubClient *github.GitHub, recipeCfgFile, serverURL string, client *http.Client, cr codereview.CodeReview, local bool) (RepoManager, error) {
	if err := c.Validate(); err != nil {
		return nil, err
	}
	wd := path.Join(workdir, "github_repos")
	if _, err := os.Stat(wd); err != nil {
		if err := os.MkdirAll(wd, 0755); err != nil {
			return nil, err
		}
	}

	// Create and populate the parent directory if needed.
	_, repo := GetUserAndRepo(c.ParentRepoURL)
	userFork := fmt.Sprintf("git@github.com:%s/%s.git", cr.UserName(), repo)
	parentRepo, err := git.NewCheckout(ctx, userFork, wd)
	if err != nil {
		return nil, err
	}

	crm, err := newCommonRepoManager(c.CommonRepoManagerConfig, wd, serverURL, nil, client, cr, local)
	if err != nil {
		return nil, err
	}

	// Create and populate the child directory if needed.
	if _, err := os.Stat(crm.childDir); err != nil {
		if err := os.MkdirAll(crm.childDir, 0755); err != nil {
			return nil, err
		}
		if _, err := git.GitDir(crm.childDir).Git(ctx, "clone", c.ChildRepoURL, "."); err != nil {
			return nil, err
		}
	}

	gr := &githubRepoManager{
		commonRepoManager: crm,
		githubClient:      githubClient,
		parentRepo:        parentRepo,
		parentRepoURL:     c.ParentRepoURL,
		childRepoURL:      c.ChildRepoURL,
		revisionFile:      c.RevisionFile,
		defaultStrategy:   c.DefaultStrategy,
		gsBucket:          c.StorageBucket,
		gsPathTemplates:   c.StoragePathTemplates,
	}

	return gr, nil
}

// See documentation for RepoManager interface.
func (rm *githubRepoManager) Update(ctx context.Context) error {
	// Sync the projects.
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	// Update the repositories.
	if err := rm.parentRepo.Update(ctx); err != nil {
		return err
	}
	if err := rm.childRepo.Update(ctx); err != nil {
		return err
	}

	// Check to see whether there is an upstream yet.
	remoteOutput, err := rm.parentRepo.Git(ctx, "remote", "show")
	if err != nil {
		return err
	}
	remoteFound := false
	remoteLines := strings.Split(remoteOutput, "\n")
	for _, remoteLine := range remoteLines {
		if remoteLine == GITHUB_UPSTREAM_REMOTE_NAME {
			remoteFound = true
			break
		}
	}
	if !remoteFound {
		if _, err := rm.parentRepo.Git(ctx, "remote", "add", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentRepoURL); err != nil {
			return err
		}
	}
	// Fetch upstream.
	if _, err := rm.parentRepo.Git(ctx, "fetch", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentBranch); err != nil {
		return err
	}

	// Read the contents of the revision file to determine the last roll rev.
	revisionFileContents, err := rm.githubClient.ReadRawFile(rm.parentBranch, rm.revisionFile)
	if err != nil {
		return err
	}
	lastRollRev := strings.TrimRight(revisionFileContents, "\n")

	// Find the number of not-rolled child repo commits.
	notRolled, err := rm.getCommitsNotRolled(ctx, lastRollRev)
	if err != nil {
		return err
	}

	// Get the next roll revision.
	nextRollRev, err := rm.getNextRollRev(ctx, notRolled, lastRollRev)
	if err != nil {
		return err
	}

	rm.infoMtx.Lock()
	defer rm.infoMtx.Unlock()
	rm.lastRollRev = lastRollRev
	rm.nextRollRev = nextRollRev
	rm.commitsNotRolled = len(notRolled)

	sklog.Infof("lastRollRev is: %s", rm.lastRollRev)
	sklog.Infof("nextRollRev is: %s", nextRollRev)
	sklog.Infof("commitsNotRolled: %d", rm.commitsNotRolled)
	return nil
}

func (rm *githubRepoManager) cleanParent(ctx context.Context) error {
	if _, err := rm.parentRepo.Git(ctx, "clean", "-d", "-f", "-f"); err != nil {
		return err
	}
	_, _ = rm.parentRepo.Git(ctx, "rebase", "--abort")
	if _, err := rm.parentRepo.Git(ctx, "checkout", fmt.Sprintf("%s/%s", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentBranch), "-f"); err != nil {
		return err
	}
	_, _ = rm.parentRepo.Git(ctx, "branch", "-D", ROLL_BRANCH)
	return nil
}

// See documentation for RepoManager interface.
func (rm *githubRepoManager) CreateNewRoll(ctx context.Context, from, to string, emails []string, cqExtraTrybots string, dryRun bool) (int64, error) {
	rm.repoMtx.Lock()
	defer rm.repoMtx.Unlock()

	sklog.Info("Creating a new Github Roll")

	// Clean the checkout, get onto a fresh branch.
	if err := rm.cleanParent(ctx); err != nil {
		return 0, err
	}
	if _, err := rm.parentRepo.Git(ctx, "checkout", fmt.Sprintf("%s/%s", GITHUB_UPSTREAM_REMOTE_NAME, rm.parentBranch), "-b", ROLL_BRANCH); err != nil {
		return 0, err
	}
	// Defer cleanup.
	defer func() {
		util.LogErr(rm.cleanParent(ctx))
	}()

	// Make sure the forked repo is at the same hash as the target repo before
	// creating the pull request.
	if _, err := rm.parentRepo.Git(ctx, "push", "origin", ROLL_BRANCH, "-f"); err != nil {
		return 0, err
	}

	// Make sure the right name and email are set.
	if !rm.local {
		if _, err := rm.parentRepo.Git(ctx, "config", "user.name", rm.codereview.UserName()); err != nil {
			return 0, err
		}
		if _, err := rm.parentRepo.Git(ctx, "config", "user.email", rm.codereview.UserEmail()); err != nil {
			return 0, err
		}
	}

	// Build the commit message.
	user, repo := GetUserAndRepo(rm.childRepoURL)
	githubCompareUrl := fmt.Sprintf("https://github.com/%s/%s/compare/%s...%s", user, repo, from[:12], to[:12])
	logStr, err := rm.childRepo.Git(ctx, "log", fmt.Sprintf("%s..%s", from, to), "--no-merges", "--oneline")
	if err != nil {
		return 0, err
	}
	logStr = strings.TrimSpace(logStr)
	// Github autolinks PR numbers to be of the same repository in logStr. Fix this by
	// explicitly adding the child repo to the PR number.
	logStr = pullRequestInLogRE.ReplaceAllString(logStr, fmt.Sprintf(" (%s/%s$1)", user, repo))

	data := struct {
		ChildPath           string
		ChildRepoCompareUrl string
		From                string
		To                  string
		NumCommits          int
		LogURL              string
		LogStr              string
		ServerURL           string
		SheriffEmails       string
	}{
		ChildPath:           rm.childPath,
		ChildRepoCompareUrl: githubCompareUrl,
		From:                from[:12],
		To:                  to[:12],
		NumCommits:          len(strings.Split(logStr, "\n")),
		LogStr:              logStr,
		ServerURL:           rm.serverURL,
		SheriffEmails:       strings.Join(emails, ","),
	}
	var buf bytes.Buffer
	if err := githubCommitMsgTmpl.Execute(&buf, data); err != nil {
		return 0, err
	}
	commitMsg := buf.String()

	versions, err := rm.childRepo.RevList(ctx, "--no-merges", fmt.Sprintf("%s..%s", from, to))
	if err != nil {
		return 0, err
	}
	logStrList := strings.Split(logStr, "\n")
	for i := len(versions) - 1; i >= 0; i-- {
		version := versions[i]
		// Write the file.
		if err := ioutil.WriteFile(path.Join(rm.parentRepo.Dir(), rm.revisionFile), []byte(version+"\n"), os.ModePerm); err != nil {
			return 0, err
		}

		// Commit.
		if _, err := rm.parentRepo.Git(ctx, "commit", "-a", "-m", logStrList[i]); err != nil {
			return 0, err
		}

	}

	// Run the pre-upload steps.
	for _, s := range rm.PreUploadSteps() {
		if err := s(ctx, rm.httpClient, rm.parentRepo.Dir()); err != nil {
			return 0, fmt.Errorf("Error when running pre-upload step: %s", err)
		}
	}

	// Push to the forked repository.
	if _, err := rm.parentRepo.Git(ctx, "push", "origin", ROLL_BRANCH, "-f"); err != nil {
		return 0, err
	}

	// Grab the first line of the commit msg to use as the title of the pull request.
	title := strings.Split(commitMsg, "\n")[0]
	// Use the remaining part of the commit message as the pull request description.
	commitMsgLines := strings.Split(commitMsg, "\n")
	var descComment []string
	if len(commitMsgLines) > 50 {
		// Truncate too large description comment because Github API cannot handle large comments.
		descComment = commitMsgLines[1:50]
		descComment = append(descComment, "...")
	} else {
		descComment = commitMsgLines[1:]
	}
	// Create a pull request.
	headBranch := fmt.Sprintf("%s:%s", rm.codereview.UserName(), ROLL_BRANCH)
	pr, err := rm.githubClient.CreatePullRequest(title, rm.parentBranch, headBranch, strings.Join(descComment, "\n"))
	if err != nil {
		return 0, err
	}

	// Add appropriate label to the pull request.
	label := github.COMMIT_LABEL
	if dryRun {
		label = github.DRYRUN_LABEL
	}
	if err := rm.githubClient.AddLabel(pr.GetNumber(), label); err != nil {
		return 0, err
	}

	// Mention the sheriffs on the pull request so that they are automatically
	// subscribed to it.
	mentions := []string{}
	for _, e := range emails {
		m := fmt.Sprintf("@%s", strings.Split(e, "@")[0])
		mentions = append(mentions, m)
	}
	if err := rm.githubClient.AddComment(pr.GetNumber(), fmt.Sprintf("%s : New roll has been created by %s", strings.Join(mentions, " "), rm.serverURL)); err != nil {
		return 0, err
	}

	return int64(pr.GetNumber()), nil
}

func GetUserAndRepo(githubRepo string) (string, string) {
	repoTokens := strings.Split(githubRepo, ":")
	user := strings.Split(repoTokens[1], "/")[0]
	repo := strings.TrimRight(strings.Split(repoTokens[1], "/")[1], ".git")
	return user, repo
}

// See documentation for RepoManager interface.
func (r *githubRepoManager) CreateNextRollStrategy(ctx context.Context, s string) (strategy.NextRollStrategy, error) {
	return strategy.GetNextRollStrategy(ctx, s, r.childBranch, DEFAULT_REMOTE, r.gsBucket, r.gsPathTemplates, r.childRepo, nil)
}

// See documentation for RepoManager interface.
func (r *githubRepoManager) DefaultStrategy() string {
	return r.defaultStrategy
}

// See documentation for RepoManager interface.
func (r *githubRepoManager) ValidStrategies() []string {
	return []string{
		strategy.ROLL_STRATEGY_GCS_FILE,
		strategy.ROLL_STRATEGY_SINGLE,
		strategy.ROLL_STRATEGY_BATCH,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
)

var server = flag.String("server", "punter.inf.ed.ac.uk", "server ip")
var port = flag.Int("port", 9001, "server port")
var name = flag.String("name", "iris punter", "bot name")

type HandshakeRequest struct {
	Me string `json:"me"`
}

type HandshakeResponse struct {
	You string `json:"you"`
}

type PunterID int
type SiteID int

type Site struct {
	ID SiteID `json:"id"`
}

type River struct {
	Source  SiteID   `json:"source"`
	Target  SiteID   `json:"target"`
	Claimed bool     `json:"claimed"`
	Owner   PunterID `json:"owner"`
}

type Map struct {
	Sites  []Site   `json:"sites"`
	Rivers []River  `json:"rivers"`
	Mines  []SiteID `json:"mines"`
}

type SetupRequest struct {
	Punter  PunterID `json:"punter"`
	Punters int      `json:"punters"`
	Map     Map      `json:"map"`
}

type SetupResponse struct {
	Ready PunterID `json:"ready"`
}

type Claim struct {
	Punter PunterID `json:"punter"`
	Source SiteID   `json:"source"`
	Target SiteID   `json:"target"`
}

type Pass struct {
	Punter PunterID `json:"punter"`
}

// Poor man's union type. Only one of Claim or Pass is non-null
type Move struct {
	Claim *Claim `json:"claim",omitempty`
	Pass  *Pass  `json:"pass",omitempty`
}

func (m Move) String() string {
	if m.Claim != nil {
		return fmt.Sprintf("claim:%+v", m.Claim)
	} else if m.Pass != nil {
		return fmt.Sprintf("pass:%+v", m.Pass)
	} else {
		return "empty"
	}
}

type Moves struct {
	Moves []Move `json:"moves"`
}

type ServerMove struct {
	Move Moves `json:"move"`
}

func findServer() (conn net.Conn, err error) {
	p := *port
	serverAddress := fmt.Sprintf("%s:%d", *server, p)
	log.Printf("Trying %s", serverAddress)
	conn, err = net.Dial("tcp", serverAddress)
	if err == nil {
		return
	}
	log.Fatal()
	return
}

func send(conn io.Writer, d interface{}) (err error) {
	var b []byte
	buf := bytes.NewBuffer(nil)
	err = json.NewEncoder(buf).Encode(d)
	if err != nil {
		return
	}
	b = buf.Bytes()
	msg := fmt.Sprintf("%d:%s", len(b), b)
	// log.Printf("Sending: %s", msg)
	_, err = conn.Write([]byte(msg))
	if err != nil {
		return
	}
	return err
}

func receive(conn io.Reader, d interface{}) (err error) {
	var i int
	_, err = fmt.Fscanf(conn, "%d:", &i)
	if err != nil {
		return
	}
	// log.Printf("Reading %d bytes", i)
	b1 := make([]byte, i)
	offset := 0
	for offset < i {
		var n int
		n, err = conn.Read(b1[offset:])
		if err != nil {
			return
		}
		offset += n
	}
	// log.Printf("Bytes: %d %v", len(b1), b1)
	// listen for reply
	err = json.Unmarshal(b1, d)
	return err
}

func handshake(conn io.ReadWriter) (err error) {
	handshakeRequest := HandshakeRequest{*name}
	err = send(conn, &handshakeRequest)
	if err != nil {
		return
	}

	// log.Printf("Waiting for reply")
	// listen for reply
	var handshakeResponse HandshakeResponse
	err = receive(conn, &handshakeResponse)
	if err != nil {
		return
	}
	// fmt.Printf("response %v\n", handshakeResponse)
	return
}

func setup(conn io.ReadWriter) (setupRequest SetupRequest, err error) {
	err = receive(conn, &setupRequest)
	if err != nil {
		return
	}
	setupResponse := SetupResponse{setupRequest.Punter}
	err = send(conn, &setupResponse)
	return
}

func processServerMove(setupRequest SetupRequest, serverMove ServerMove) (err error) {
	for _, move := range serverMove.Move.Moves {
		if move.Claim != nil {
			for riverIndex, river := range setupRequest.Map.Rivers {
				if river.Source == move.Claim.Source &&
					river.Target == move.Claim.Target {
					river.Claimed = true
					river.Owner = move.Claim.Punter
					setupRequest.Map.Rivers[riverIndex] = river
					break
				}
			}
		}
	}
	return
}

func main() {
	flag.Parse()
	err := onlineMode()
	if err != nil {
		log.Fatal(err)
	}
}

func onlineMode() (err error) {
	conn, err := findServer()
	if err != nil {
		return
	}
	log.Printf("connected")
	err = handshake(conn)
	if err != nil {
		return
	}

	setupRequest, err := setup(conn)
	if err != nil {
		return
	}

	for {
		log.Printf("Setup %+v", setupRequest)
		var serverMove ServerMove
		err = receive(conn, &serverMove)
		if err != nil {
			return
		}
		err = processServerMove(setupRequest, serverMove)
		move := Move{nil, &Pass{setupRequest.Punter}}
		err = send(conn, move)
		if err != nil {
			return
		}
	}
	return
}
<commit_msg>Pick first claimed river<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"net"
)

var server = flag.String("server", "punter.inf.ed.ac.uk", "server ip")
var port = flag.Int("port", 9001, "server port")
var name = flag.String("name", "iris punter", "bot name")

type HandshakeRequest struct {
	Me string `json:"me"`
}

type HandshakeResponse struct {
	You string `json:"you"`
}

type PunterID int
type SiteID int

type Site struct {
	ID SiteID `json:"id"`
}

type River struct {
	Source  SiteID   `json:"source"`
	Target  SiteID   `json:"target"`
	Claimed bool     `json:"claimed",omitempty`
	Owner   PunterID `json:"owner",omitempty`
}

type Map struct {
	Sites  []Site   `json:"sites"`
	Rivers []River  `json:"rivers"`
	Mines  []SiteID `json:"mines"`
}

type SetupRequest struct {
	Punter  PunterID `json:"punter"`
	Punters int      `json:"punters"`
	Map     Map      `json:"map"`
}

type SetupResponse struct {
	Ready PunterID `json:"ready"`
}

type Claim struct {
	Punter PunterID `json:"punter"`
	Source SiteID   `json:"source"`
	Target SiteID   `json:"target"`
}

type Pass struct {
	Punter PunterID `json:"punter"`
}

// Poor man's union type. Only one of Claim or Pass is non-null
type Move struct {
	Claim *Claim `json:"claim",omitempty`
	Pass  *Pass  `json:"pass",omitempty`
}

func (m Move) String() string {
	if m.Claim != nil {
		return fmt.Sprintf("claim:%+v", m.Claim)
	} else if m.Pass != nil {
		return fmt.Sprintf("pass:%+v", m.Pass)
	} else {
		return "empty"
	}
}

type Moves struct {
	Moves []Move `json:"moves"`
}

type ServerMove struct {
	Move Moves `json:"move"`
}

func findServer() (conn net.Conn, err error) {
	p := *port
	serverAddress := fmt.Sprintf("%s:%d", *server, p)
	log.Printf("Trying %s", serverAddress)
	conn, err = net.Dial("tcp", serverAddress)
	if err == nil {
		return
	}
	log.Fatal()
	return
}

func send(conn io.Writer, d interface{}) (err error) {
	var b []byte
	buf := bytes.NewBuffer(nil)
	err = json.NewEncoder(buf).Encode(d)
	if err != nil {
		return
	}
	b = buf.Bytes()
	msg := fmt.Sprintf("%d:%s", len(b), b)
	// log.Printf("Sending: %s", msg)
	_, err = conn.Write([]byte(msg))
	if err != nil {
		return
	}
	return err
}

func receive(conn io.Reader, d interface{}) (err error) {
	var i int
	_, err = fmt.Fscanf(conn, "%d:", &i)
	if err != nil {
		return
	}
	// log.Printf("Reading %d bytes", i)
	b1 := make([]byte, i)
	offset := 0
	for offset < i {
		var n int
		n, err = conn.Read(b1[offset:])
		if err != nil {
			return
		}
		offset += n
	}
	// log.Printf("Bytes: %d %v", len(b1), b1)
	// listen for reply
	err = json.Unmarshal(b1, d)
	return err
}

func handshake(conn io.ReadWriter) (err error) {
	handshakeRequest := HandshakeRequest{*name}
	err = send(conn, &handshakeRequest)
	if err != nil {
		return
	}

	// log.Printf("Waiting for reply")
	// listen for reply
	var handshakeResponse HandshakeResponse
	err = receive(conn, &handshakeResponse)
	if err != nil {
		return
	}
	// fmt.Printf("response %v\n", handshakeResponse)
	return
}

func setup(conn io.ReadWriter) (setupRequest SetupRequest, err error) {
	err = receive(conn, &setupRequest)
	if err != nil {
		return
	}
	setupResponse := SetupResponse{setupRequest.Punter}
	err = send(conn, &setupResponse)
	return
}

func processServerMove(setupRequest SetupRequest, serverMove ServerMove) (err error) {
	for _, move := range serverMove.Move.Moves {
		if move.Claim != nil {
			for riverIndex, river := range setupRequest.Map.Rivers {
				if river.Source == move.Claim.Source &&
					river.Target == move.Claim.Target {
					river.Claimed = true
					river.Owner = move.Claim.Punter
					setupRequest.Map.Rivers[riverIndex] = river
					break
				}
			}
		}
	}
	return
}

func pickPass(setupRequest SetupRequest) (move Move, err error) {
	move = Move{nil, &Pass{setupRequest.Punter}}
	return
}

func pickFirstUnclaimed(setupRequest SetupRequest) (move Move, err error) {
	for _, river := range setupRequest.Map.Rivers {
		if river.Claimed == false {
			move.Claim = &Claim{setupRequest.Punter, river.Source, river.Target}
			return
		}
	}
	return pickPass(setupRequest)
}

func main() {
	flag.Parse()
	err := onlineMode()
	if err != nil {
		log.Fatal(err)
	}
}

func onlineMode() (err error) {
	conn, err := findServer()
	if err != nil {
		return
	}
	log.Printf("connected")
	err = handshake(conn)
	if err != nil {
		return
	}

	setupRequest, err := setup(conn)
	if err != nil {
		return
	}

	for {
		log.Printf("Setup %+v", setupRequest)
		var serverMove ServerMove
		err = receive(conn, &serverMove)
		if err != nil {
			return
		}
		err = processServerMove(setupRequest, serverMove)
		if err != nil {
			return
		}
		var move Move
		move, err = pickFirstUnclaimed(setupRequest)
		if err != nil {
			return
		}
		log.Printf("Move: %v", move)
		err = send(conn, move)
		if err != nil {
			return
		}
	}
	return
}
<|endoftext|>


<commit_before>package lumberjack

import (
	"crypto/tls"
	"math/rand"
	"net"
	"time"

	"github.com/elastic/libbeat/logp"
)

// TransportClient interfaces adds (re-)connect support to net.Conn.
type TransportClient interface {
	net.Conn
	Connect(timeout time.Duration) error
	IsConnected() bool
}

type tcpClient struct {
	hostport  string
	connected bool
	net.Conn
}

type tlsClient struct {
	tcpClient
	tls tls.Config
}

func newTCPClient(host string) (*tcpClient, error) {
	return &tcpClient{hostport: host}, nil
}

func (c *tcpClient) Connect(timeout time.Duration) error {
	if c.IsConnected() {
		_ = c.Close()
	}

	host, port, err := net.SplitHostPort(c.hostport)
	if err != nil {
		return err
	}

	// TODO: address lookup copied from logstash-forwarded. Really required?
	addresses, err := net.LookupHost(host)
	c.Conn = nil
	if err != nil {
		logp.Warn("DNS lookup failure \"%s\": %s", host, err)
		return err
	}

	// connect to random address
	// Use randomization on DNS reported addresses combined with timeout and ACKs
	// to spread potential load when starting up large number of beats using
	// lumberjack.
	//
	// RFCs discussing reasons for ignoring order of DNS records:
	// http://www.ietf.org/rfc/rfc3484.txt
	// > is specific to locality-based address selection for multiple dns
	// > records, but exists as prior art in "Choose some different ordering for
	// > the dns records" done by a client
	//
	// https://tools.ietf.org/html/rfc1794
	// > "Clients, of course, may reorder this information" - with respect to
	// > handling order of dns records in a response. address :=
	address := addresses[rand.Int()%len(addresses)]
	addressport := net.JoinHostPort(address, port)
	conn, err := net.DialTimeout("tcp", addressport, timeout)
	if err != nil {
		return err
	}

	c.Conn = conn
	c.connected = true
	return nil
}

func (c *tcpClient) IsConnected() bool {
	return c.connected
}

func (c *tcpClient) Close() error {
	err := c.Conn.Close()
	c.connected = false
	return err
}

func newTLSClient(host string, tls *tls.Config) (*tlsClient, error) {
	c := tlsClient{}
	c.hostport = host
	c.tls = *tls
	return &c, nil
}

func (c *tlsClient) Connect(timeout time.Duration) error {
	host, _, err := net.SplitHostPort(c.hostport)
	if err != nil {
		return err
	}

	var tlsconfig tls.Config
	tlsconfig.MinVersion = c.tls.MinVersion
	tlsconfig.RootCAs = c.tls.RootCAs
	tlsconfig.Certificates = c.tls.Certificates
	tlsconfig.ServerName = host

	if err := c.tcpClient.Connect(timeout); err != nil {
		return err
	}

	socket := tls.Client(c.Conn, &tlsconfig)
	if err := socket.SetDeadline(time.Now().Add(timeout)); err != nil {
		_ = socket.Close()
		return c.onFail(err)
	}
	if err := socket.Handshake(); err != nil {
		_ = socket.Close()
		return c.onFail(err)
	}

	c.Conn = socket
	c.connected = true
	return nil
}

func (c *tlsClient) onFail(err error) error {
	c.Conn = nil
	c.connected = false
	return err
}
<commit_msg>log tls connection failures with error log level<commit_after>package lumberjack

import (
	"crypto/tls"
	"math/rand"
	"net"
	"time"

	"github.com/elastic/libbeat/logp"
)

// TransportClient interfaces adds (re-)connect support to net.Conn.
type TransportClient interface {
	net.Conn
	Connect(timeout time.Duration) error
	IsConnected() bool
}

type tcpClient struct {
	hostport  string
	connected bool
	net.Conn
}

type tlsClient struct {
	tcpClient
	tls tls.Config
}

func newTCPClient(host string) (*tcpClient, error) {
	return &tcpClient{hostport: host}, nil
}

func (c *tcpClient) Connect(timeout time.Duration) error {
	if c.IsConnected() {
		_ = c.Close()
	}

	host, port, err := net.SplitHostPort(c.hostport)
	if err != nil {
		return err
	}

	// TODO: address lookup copied from logstash-forwarded. Really required?
	addresses, err := net.LookupHost(host)
	c.Conn = nil
	if err != nil {
		logp.Warn("DNS lookup failure \"%s\": %s", host, err)
		return err
	}

	// connect to random address
	// Use randomization on DNS reported addresses combined with timeout and ACKs
	// to spread potential load when starting up large number of beats using
	// lumberjack.
	//
	// RFCs discussing reasons for ignoring order of DNS records:
	// http://www.ietf.org/rfc/rfc3484.txt
	// > is specific to locality-based address selection for multiple dns
	// > records, but exists as prior art in "Choose some different ordering for
	// > the dns records" done by a client
	//
	// https://tools.ietf.org/html/rfc1794
	// > "Clients, of course, may reorder this information" - with respect to
	// > handling order of dns records in a response. address :=
	address := addresses[rand.Int()%len(addresses)]
	addressport := net.JoinHostPort(address, port)
	conn, err := net.DialTimeout("tcp", addressport, timeout)
	if err != nil {
		return err
	}

	c.Conn = conn
	c.connected = true
	return nil
}

func (c *tcpClient) IsConnected() bool {
	return c.connected
}

func (c *tcpClient) Close() error {
	err := c.Conn.Close()
	c.connected = false
	return err
}

func newTLSClient(host string, tls *tls.Config) (*tlsClient, error) {
	c := tlsClient{}
	c.hostport = host
	c.tls = *tls
	return &c, nil
}

func (c *tlsClient) Connect(timeout time.Duration) error {
	host, _, err := net.SplitHostPort(c.hostport)
	if err != nil {
		return err
	}

	var tlsconfig tls.Config
	tlsconfig.MinVersion = c.tls.MinVersion
	tlsconfig.RootCAs = c.tls.RootCAs
	tlsconfig.Certificates = c.tls.Certificates
	tlsconfig.ServerName = host

	if err := c.tcpClient.Connect(timeout); err != nil {
		return c.onFail(err)
	}

	socket := tls.Client(c.Conn, &tlsconfig)
	if err := socket.SetDeadline(time.Now().Add(timeout)); err != nil {
		_ = socket.Close()
		return c.onFail(err)
	}
	if err := socket.Handshake(); err != nil {
		_ = socket.Close()
		return c.onFail(err)
	}

	c.Conn = socket
	c.connected = true
	return nil
}

func (c *tlsClient) onFail(err error) error {
	logp.Err("SSL client failed to connect with: %v", err)
	c.Conn = nil
	c.connected = false
	return err
}
<|endoftext|>


<commit_before>package weed_server

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/go/glog"
	"github.com/chrislusf/seaweedfs/go/operation"
	"github.com/chrislusf/seaweedfs/go/storage"
	"github.com/chrislusf/seaweedfs/go/topology"
	"github.com/chrislusf/seaweedfs/go/util"
	"github.com/golang/protobuf/proto"
)

func (ms *MasterServer) collectionDeleteHandler(w http.ResponseWriter, r *http.Request) {
	collection, ok := ms.Topo.GetCollection(r.FormValue("collection"))
	if !ok {
		writeJsonError(w, r, http.StatusBadRequest, fmt.Errorf("collection %s does not exist", r.FormValue("collection")))
		return
	}
	for _, server := range collection.ListVolumeServers() {
		_, err := util.Get("http://" + server.Ip + ":" + strconv.Itoa(server.Port) + "/admin/delete_collection?collection=" + r.FormValue("collection"))
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
			return
		}
	}
	ms.Topo.DeleteCollection(r.FormValue("collection"))
}

func (ms *MasterServer) dirJoinHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeJsonError(w, r, http.StatusBadRequest, err)
		return
	}
	joinMessage := &operation.JoinMessage{}
	if err = proto.Unmarshal(body, joinMessage); err != nil {
		writeJsonError(w, r, http.StatusBadRequest, err)
		return
	}
	if *joinMessage.Ip == "" {
		*joinMessage.Ip = r.RemoteAddr[0:strings.Index(r.RemoteAddr, ":")]
	}
	if glog.V(4) {
		if jsonData, jsonError := json.Marshal(joinMessage); jsonError != nil {
			glog.V(0).Infoln("json marshaling error: ", jsonError)
			writeJsonError(w, r, http.StatusBadRequest, jsonError)
			return
		} else {
			glog.V(4).Infoln("Proto size", len(body), "json size", len(jsonData), string(jsonData))
		}
	}

	ms.Topo.ProcessJoinMessage(joinMessage)
	writeJsonQuiet(w, r, http.StatusOK, operation.JoinResult{
		VolumeSizeLimit: uint64(ms.volumeSizeLimitMB) * 1024 * 1024,
		SecretKey:       string(ms.guard.SecretKey),
	})
}

func (ms *MasterServer) dirStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Topology"] = ms.Topo.ToMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) volumeVacuumHandler(w http.ResponseWriter, r *http.Request) {
	gcThreshold := r.FormValue("garbageThreshold")
	if gcThreshold == "" {
		gcThreshold = ms.garbageThreshold
	}
	debug("garbageThreshold =", gcThreshold)
	ms.Topo.Vacuum(gcThreshold)
	ms.dirStatusHandler(w, r)
}

func (ms *MasterServer) volumeGrowHandler(w http.ResponseWriter, r *http.Request) {
	count := 0
	option, err := ms.getVolumeGrowOption(r)
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
		return
	}
	if err == nil {
		if count, err = strconv.Atoi(r.FormValue("count")); err == nil {
			if ms.Topo.FreeSpace() < count*option.ReplicaPlacement.GetCopyCount() {
				err = errors.New("Only " + strconv.Itoa(ms.Topo.FreeSpace()) + " volumes left! Not enough for " + strconv.Itoa(count*option.ReplicaPlacement.GetCopyCount()))
			} else {
				count, err = ms.vg.GrowByCountAndType(count, option, ms.Topo)
			}
		} else {
			err = errors.New("parameter count is not found")
		}
	}
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
	} else {
		writeJsonQuiet(w, r, http.StatusOK, map[string]interface{}{"count": count})
	}
}

func (ms *MasterServer) volumeStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Volumes"] = ms.Topo.ToVolumeMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) redirectHandler(w http.ResponseWriter, r *http.Request) {
	vid, _, _, _, _ := parseURLPath(r.URL.Path)
	volumeId, err := storage.NewVolumeId(vid)
	if err != nil {
		debug("parsing error:", err, r.URL.Path)
		return
	}
	machines := ms.Topo.Lookup("", volumeId)
	if machines != nil && len(machines) > 0 {
		http.Redirect(w, r, util.NormalizeUrl(machines[rand.Intn(len(machines))].PublicUrl)+r.URL.Path, http.StatusMovedPermanently)
	} else {
		writeJsonError(w, r, http.StatusNotFound, fmt.Errorf("volume id %d not found", volumeId))
	}
}

func (ms *MasterServer) submitFromMasterServerHandler(w http.ResponseWriter, r *http.Request) {
	if ms.Topo.IsLeader() {
		submitForClientHandler(w, r, "localhost:"+strconv.Itoa(ms.port))
	} else {
		submitForClientHandler(w, r, ms.Topo.RaftServer.Leader())
	}
}

func (ms *MasterServer) deleteFromMasterServerHandler(w http.ResponseWriter, r *http.Request) {
	if ms.Topo.IsLeader() {
		deleteForClientHandler(w, r, "localhost:"+strconv.Itoa(ms.port))
	} else {
		deleteForClientHandler(w, r, ms.Topo.RaftServer.Leader())
	}
}

func (ms *MasterServer) HasWritableVolume(option *topology.VolumeGrowOption) bool {
	vl := ms.Topo.GetVolumeLayout(option.Collection, option.ReplicaPlacement, option.Ttl)
	return vl.GetActiveVolumeCount(option) > 0
}

func (ms *MasterServer) getVolumeGrowOption(r *http.Request) (*topology.VolumeGrowOption, error) {
	replicationString := r.FormValue("replication")
	if replicationString == "" {
		replicationString = ms.defaultReplicaPlacement
	}
	replicaPlacement, err := storage.NewReplicaPlacementFromString(replicationString)
	if err != nil {
		return nil, err
	}
	ttl, err := storage.ReadTTL(r.FormValue("ttl"))
	if err != nil {
		return nil, err
	}
	volumeGrowOption := &topology.VolumeGrowOption{
		Collection:       r.FormValue("collection"),
		ReplicaPlacement: replicaPlacement,
		Ttl:              ttl,
		DataCenter:       r.FormValue("dataCenter"),
		Rack:             r.FormValue("rack"),
		DataNode:         r.FormValue("dataNode"),
	}
	return volumeGrowOption, nil
}
<commit_msg>FIXED: When RaftServer cannot find a leader, Return a more readable error. Before: curl -F "file=1234" "http://127.0.0.1:9333/submit" {"error":"Post http:///dir/assign: http: no Host in request URL"} After: curl -F "file=1234" "http://127.0.0.1:9333/submit" {"error":"Raft Server not initialized!"}<commit_after>package weed_server

import (
	"encoding/json"
	"errors"
	"fmt"
	"io/ioutil"
	"math/rand"
	"net/http"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/go/glog"
	"github.com/chrislusf/seaweedfs/go/operation"
	"github.com/chrislusf/seaweedfs/go/storage"
	"github.com/chrislusf/seaweedfs/go/topology"
	"github.com/chrislusf/seaweedfs/go/util"
	"github.com/golang/protobuf/proto"
)

func (ms *MasterServer) collectionDeleteHandler(w http.ResponseWriter, r *http.Request) {
	collection, ok := ms.Topo.GetCollection(r.FormValue("collection"))
	if !ok {
		writeJsonError(w, r, http.StatusBadRequest, fmt.Errorf("collection %s does not exist", r.FormValue("collection")))
		return
	}
	for _, server := range collection.ListVolumeServers() {
		_, err := util.Get("http://" + server.Ip + ":" + strconv.Itoa(server.Port) + "/admin/delete_collection?collection=" + r.FormValue("collection"))
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
			return
		}
	}
	ms.Topo.DeleteCollection(r.FormValue("collection"))
}

func (ms *MasterServer) dirJoinHandler(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		writeJsonError(w, r, http.StatusBadRequest, err)
		return
	}
	joinMessage := &operation.JoinMessage{}
	if err = proto.Unmarshal(body, joinMessage); err != nil {
		writeJsonError(w, r, http.StatusBadRequest, err)
		return
	}
	if *joinMessage.Ip == "" {
		*joinMessage.Ip = r.RemoteAddr[0:strings.Index(r.RemoteAddr, ":")]
	}
	if glog.V(4) {
		if jsonData, jsonError := json.Marshal(joinMessage); jsonError != nil {
			glog.V(0).Infoln("json marshaling error: ", jsonError)
			writeJsonError(w, r, http.StatusBadRequest, jsonError)
			return
		} else {
			glog.V(4).Infoln("Proto size", len(body), "json size", len(jsonData), string(jsonData))
		}
	}

	ms.Topo.ProcessJoinMessage(joinMessage)
	writeJsonQuiet(w, r, http.StatusOK, operation.JoinResult{
		VolumeSizeLimit: uint64(ms.volumeSizeLimitMB) * 1024 * 1024,
		SecretKey:       string(ms.guard.SecretKey),
	})
}

func (ms *MasterServer) dirStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Topology"] = ms.Topo.ToMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) volumeVacuumHandler(w http.ResponseWriter, r *http.Request) {
	gcThreshold := r.FormValue("garbageThreshold")
	if gcThreshold == "" {
		gcThreshold = ms.garbageThreshold
	}
	debug("garbageThreshold =", gcThreshold)
	ms.Topo.Vacuum(gcThreshold)
	ms.dirStatusHandler(w, r)
}

func (ms *MasterServer) volumeGrowHandler(w http.ResponseWriter, r *http.Request) {
	count := 0
	option, err := ms.getVolumeGrowOption(r)
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
		return
	}
	if err == nil {
		if count, err = strconv.Atoi(r.FormValue("count")); err == nil {
			if ms.Topo.FreeSpace() < count*option.ReplicaPlacement.GetCopyCount() {
				err = errors.New("Only " + strconv.Itoa(ms.Topo.FreeSpace()) + " volumes left! Not enough for " + strconv.Itoa(count*option.ReplicaPlacement.GetCopyCount()))
			} else {
				count, err = ms.vg.GrowByCountAndType(count, option, ms.Topo)
			}
		} else {
			err = errors.New("parameter count is not found")
		}
	}
	if err != nil {
		writeJsonError(w, r, http.StatusNotAcceptable, err)
	} else {
		writeJsonQuiet(w, r, http.StatusOK, map[string]interface{}{"count": count})
	}
}

func (ms *MasterServer) volumeStatusHandler(w http.ResponseWriter, r *http.Request) {
	m := make(map[string]interface{})
	m["Version"] = util.VERSION
	m["Volumes"] = ms.Topo.ToVolumeMap()
	writeJsonQuiet(w, r, http.StatusOK, m)
}

func (ms *MasterServer) redirectHandler(w http.ResponseWriter, r *http.Request) {
	vid, _, _, _, _ := parseURLPath(r.URL.Path)
	volumeId, err := storage.NewVolumeId(vid)
	if err != nil {
		debug("parsing error:", err, r.URL.Path)
		return
	}
	machines := ms.Topo.Lookup("", volumeId)
	if machines != nil && len(machines) > 0 {
		http.Redirect(w, r, util.NormalizeUrl(machines[rand.Intn(len(machines))].PublicUrl)+r.URL.Path, http.StatusMovedPermanently)
	} else {
		writeJsonError(w, r, http.StatusNotFound, fmt.Errorf("volume id %d not found", volumeId))
	}
}

func (ms *MasterServer) submitFromMasterServerHandler(w http.ResponseWriter, r *http.Request) {
	if ms.Topo.IsLeader() {
		submitForClientHandler(w, r, "localhost:"+strconv.Itoa(ms.port))
	} else {
		masterUrl, err := ms.Topo.Leader()
		if err != nil {
			writeJsonError(w, r, http.StatusInternalServerError, err)
		} else {
			submitForClientHandler(w, r, masterUrl)
		}
	}
}

func (ms *MasterServer) deleteFromMasterServerHandler(w http.ResponseWriter, r *http.Request) {
	if ms.Topo.IsLeader() {
		deleteForClientHandler(w, r, "localhost:"+strconv.Itoa(ms.port))
	} else {
		deleteForClientHandler(w, r, ms.Topo.RaftServer.Leader())
	}
}

func (ms *MasterServer) HasWritableVolume(option *topology.VolumeGrowOption) bool {
	vl := ms.Topo.GetVolumeLayout(option.Collection, option.ReplicaPlacement, option.Ttl)
	return vl.GetActiveVolumeCount(option) > 0
}

func (ms *MasterServer) getVolumeGrowOption(r *http.Request) (*topology.VolumeGrowOption, error) {
	replicationString := r.FormValue("replication")
	if replicationString == "" {
		replicationString = ms.defaultReplicaPlacement
	}
	replicaPlacement, err := storage.NewReplicaPlacementFromString(replicationString)
	if err != nil {
		return nil, err
	}
	ttl, err := storage.ReadTTL(r.FormValue("ttl"))
	if err != nil {
		return nil, err
	}
	volumeGrowOption := &topology.VolumeGrowOption{
		Collection:       r.FormValue("collection"),
		ReplicaPlacement: replicaPlacement,
		Ttl:              ttl,
		DataCenter:       r.FormValue("dataCenter"),
		Rack:             r.FormValue("rack"),
		DataNode:         r.FormValue("dataNode"),
	}
	return volumeGrowOption, nil
}
<|endoftext|>


<commit_before>package peerstore

import (
	"fmt"
	"io"
	"sync"

	peer "github.com/libp2p/go-libp2p-peer"
)

var _ Peerstore = (*peerstore)(nil)

const maxInternedProtocols = 512
const maxInternedProtocolSize = 256

type peerstore struct {
	Metrics

	KeyBook
	AddrBook
	PeerMetadata

	// lock for protocol information, separate from datastore lock
	protolock         sync.Mutex
	internedProtocols map[string]string
}

// NewPeerstore creates a data structure that stores peer data, backed by the
// supplied implementations of KeyBook, AddrBook and PeerMetadata.
func NewPeerstore(kb KeyBook, ab AddrBook, md PeerMetadata) Peerstore {
	return &peerstore{
		KeyBook:           kb,
		AddrBook:          ab,
		PeerMetadata:      md,
		Metrics:           NewMetrics(),
		internedProtocols: make(map[string]string),
	}
}

func (ps *peerstore) Close() (err error) {
	var errs []error
	weakClose := func(name string, c interface{}) {
		if cl, ok := c.(io.Closer); ok {
			if err = cl.Close(); err != nil {
				errs = append(errs, fmt.Errorf("%s error: %s", name, err))
			}
		}
	}

	weakClose("keybook", ps.KeyBook)
	weakClose("addressbook", ps.AddrBook)
	weakClose("peermetadata", ps.PeerMetadata)

	if len(errs) > 0 {
		return fmt.Errorf("failed while closing peerstore; err(s): %q", errs)
	}
	return nil
}

func (ps *peerstore) Peers() peer.IDSlice {
	set := map[peer.ID]struct{}{}
	for _, p := range ps.PeersWithKeys() {
		set[p] = struct{}{}
	}
	for _, p := range ps.PeersWithAddrs() {
		set[p] = struct{}{}
	}

	pps := make(peer.IDSlice, 0, len(set))
	for p := range set {
		pps = append(pps, p)
	}
	return pps
}

func (ps *peerstore) PeerInfo(p peer.ID) PeerInfo {
	return PeerInfo{
		ID:    p,
		Addrs: ps.AddrBook.Addrs(p),
	}
}

func (ps *peerstore) internProtocol(s string) string {
	if len(s) > maxInternedProtocolSize {
		return s
	}

	if interned, ok := ps.internedProtocols[s]; ok {
		return interned
	}

	if len(ps.internedProtocols) >= maxInternedProtocols {
		ps.internedProtocols = make(map[string]string, maxInternedProtocols)
	}

	ps.internedProtocols[s] = s
	return s
}

func (ps *peerstore) SetProtocols(p peer.ID, protos ...string) error {
	ps.protolock.Lock()
	defer ps.protolock.Unlock()

	protomap := make(map[string]struct{}, len(protos))
	for _, proto := range protos {
		protomap[ps.internProtocol(proto)] = struct{}{}
	}

	return ps.Put(p, "protocols", protomap)
}

func (ps *peerstore) AddProtocols(p peer.ID, protos ...string) error {
	ps.protolock.Lock()
	defer ps.protolock.Unlock()
	protomap, err := ps.getProtocolMap(p)
	if err != nil {
		return err
	}

	for _, proto := range protos {
		protomap[ps.internProtocol(proto)] = struct{}{}
	}

	return ps.Put(p, "protocols", protomap)
}

func (ps *peerstore) getProtocolMap(p peer.ID) (map[string]struct{}, error) {
	iprotomap, err := ps.Get(p, "protocols")
	switch err {
	default:
		return nil, err
	case ErrNotFound:
		return make(map[string]struct{}), nil
	case nil:
		cast, ok := iprotomap.(map[string]struct{})
		if !ok {
			return nil, fmt.Errorf("stored protocol set was not a map")
		}

		return cast, nil
	}
}

func (ps *peerstore) GetProtocols(p peer.ID) ([]string, error) {
	ps.protolock.Lock()
	defer ps.protolock.Unlock()
	pmap, err := ps.getProtocolMap(p)
	if err != nil {
		return nil, err
	}

	out := make([]string, 0, len(pmap))
	for k := range pmap {
		out = append(out, k)
	}

	return out, nil
}

func (ps *peerstore) SupportsProtocols(p peer.ID, protos ...string) ([]string, error) {
	ps.protolock.Lock()
	defer ps.protolock.Unlock()
	pmap, err := ps.getProtocolMap(p)
	if err != nil {
		return nil, err
	}

	out := make([]string, 0, len(protos))
	for _, proto := range protos {
		if _, ok := pmap[proto]; ok {
			out = append(out, proto)
		}
	}

	return out, nil
}

func PeerInfos(ps Peerstore, peers peer.IDSlice) []PeerInfo {
	pi := make([]PeerInfo, len(peers))
	for i, p := range peers {
		pi[i] = ps.PeerInfo(p)
	}
	return pi
}

func PeerInfoIDs(pis []PeerInfo) peer.IDSlice {
	ps := make(peer.IDSlice, len(pis))
	for i, pi := range pis {
		ps[i] = pi.ID
	}
	return ps
}
<commit_msg>read/write protocol lock<commit_after>package peerstore

import (
	"fmt"
	"io"
	"sync"

	peer "github.com/libp2p/go-libp2p-peer"
)

var _ Peerstore = (*peerstore)(nil)

const maxInternedProtocols = 512
const maxInternedProtocolSize = 256

type peerstore struct {
	Metrics

	KeyBook
	AddrBook
	PeerMetadata

	// lock for protocol information, separate from datastore lock
	protolock         sync.RWMutex
	internedProtocols map[string]string
}

// NewPeerstore creates a data structure that stores peer data, backed by the
// supplied implementations of KeyBook, AddrBook and PeerMetadata.
func NewPeerstore(kb KeyBook, ab AddrBook, md PeerMetadata) Peerstore {
	return &peerstore{
		KeyBook:           kb,
		AddrBook:          ab,
		PeerMetadata:      md,
		Metrics:           NewMetrics(),
		internedProtocols: make(map[string]string),
	}
}

func (ps *peerstore) Close() (err error) {
	var errs []error
	weakClose := func(name string, c interface{}) {
		if cl, ok := c.(io.Closer); ok {
			if err = cl.Close(); err != nil {
				errs = append(errs, fmt.Errorf("%s error: %s", name, err))
			}
		}
	}

	weakClose("keybook", ps.KeyBook)
	weakClose("addressbook", ps.AddrBook)
	weakClose("peermetadata", ps.PeerMetadata)

	if len(errs) > 0 {
		return fmt.Errorf("failed while closing peerstore; err(s): %q", errs)
	}
	return nil
}

func (ps *peerstore) Peers() peer.IDSlice {
	set := map[peer.ID]struct{}{}
	for _, p := range ps.PeersWithKeys() {
		set[p] = struct{}{}
	}
	for _, p := range ps.PeersWithAddrs() {
		set[p] = struct{}{}
	}

	pps := make(peer.IDSlice, 0, len(set))
	for p := range set {
		pps = append(pps, p)
	}
	return pps
}

func (ps *peerstore) PeerInfo(p peer.ID) PeerInfo {
	return PeerInfo{
		ID:    p,
		Addrs: ps.AddrBook.Addrs(p),
	}
}

func (ps *peerstore) internProtocol(s string) string {
	if len(s) > maxInternedProtocolSize {
		return s
	}

	if interned, ok := ps.internedProtocols[s]; ok {
		return interned
	}

	if len(ps.internedProtocols) >= maxInternedProtocols {
		ps.internedProtocols = make(map[string]string, maxInternedProtocols)
	}

	ps.internedProtocols[s] = s
	return s
}

func (ps *peerstore) SetProtocols(p peer.ID, protos ...string) error {
	ps.protolock.Lock()
	defer ps.protolock.Unlock()

	protomap := make(map[string]struct{}, len(protos))
	for _, proto := range protos {
		protomap[ps.internProtocol(proto)] = struct{}{}
	}

	return ps.Put(p, "protocols", protomap)
}

func (ps *peerstore) AddProtocols(p peer.ID, protos ...string) error {
	ps.protolock.Lock()
	defer ps.protolock.Unlock()
	protomap, err := ps.getProtocolMap(p)
	if err != nil {
		return err
	}

	for _, proto := range protos {
		protomap[ps.internProtocol(proto)] = struct{}{}
	}

	return ps.Put(p, "protocols", protomap)
}

func (ps *peerstore) getProtocolMap(p peer.ID) (map[string]struct{}, error) {
	iprotomap, err := ps.Get(p, "protocols")
	switch err {
	default:
		return nil, err
	case ErrNotFound:
		return make(map[string]struct{}), nil
	case nil:
		cast, ok := iprotomap.(map[string]struct{})
		if !ok {
			return nil, fmt.Errorf("stored protocol set was not a map")
		}

		return cast, nil
	}
}

func (ps *peerstore) GetProtocols(p peer.ID) ([]string, error) {
	ps.protolock.RLock()
	defer ps.protolock.RUnlock()
	pmap, err := ps.getProtocolMap(p)
	if err != nil {
		return nil, err
	}

	out := make([]string, 0, len(pmap))
	for k := range pmap {
		out = append(out, k)
	}

	return out, nil
}

func (ps *peerstore) SupportsProtocols(p peer.ID, protos ...string) ([]string, error) {
	ps.protolock.RLock()
	defer ps.protolock.RUnlock()
	pmap, err := ps.getProtocolMap(p)
	if err != nil {
		return nil, err
	}

	out := make([]string, 0, len(protos))
	for _, proto := range protos {
		if _, ok := pmap[proto]; ok {
			out = append(out, proto)
		}
	}

	return out, nil
}

func PeerInfos(ps Peerstore, peers peer.IDSlice) []PeerInfo {
	pi := make([]PeerInfo, len(peers))
	for i, p := range peers {
		pi[i] = ps.PeerInfo(p)
	}
	return pi
}

func PeerInfoIDs(pis []PeerInfo) peer.IDSlice {
	ps := make(peer.IDSlice, len(pis))
	for i, pi := range pis {
		ps[i] = pi.ID
	}
	return ps
}
<|endoftext|>