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

package verify_test

import (
	"errors"
	"fmt"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/comeback/verify"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
)

func TestVisitor(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const namePrefix = "blobs/"

type superTest struct {
	ctx       context.Context
	blobStore mock_blob.MockStore
	visitor   graph.Visitor
}

func (t *superTest) setUp(
	ti *TestInfo,
	readFiles bool,
	allScores []blob.Score) {
	t.ctx = ti.Ctx
	t.blobStore = mock_blob.NewMockStore(ti.MockController, "blobStore")
	t.visitor = verify.NewVisitor(readFiles, allScores, t.blobStore)
}

////////////////////////////////////////////////////////////////////////
// Common
////////////////////////////////////////////////////////////////////////

type CommonTest struct {
	superTest
}

var _ SetUpInterface = &CommonTest{}

func init() { RegisterTestSuite(&CommonTest{}) }

func (t *CommonTest) SetUp(ti *TestInfo) {
	t.superTest.setUp(ti, false, nil)
}

func (t *CommonTest) InvalidNodeName() {
	_, err := t.visitor.Visit(t.ctx, "taco")

	ExpectThat(err, Error(HasSubstr("ParseNodeName")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

////////////////////////////////////////////////////////////////////////
// Directories
////////////////////////////////////////////////////////////////////////

type DirsTest struct {
	superTest

	listing  []*fs.DirectoryEntry
	contents []byte
	score    blob.Score
	node     string
}

var _ SetUpInterface = &DirsTest{}

func init() { RegisterTestSuite(&DirsTest{}) }

func (t *DirsTest) SetUp(ti *TestInfo) {
	t.superTest.setUp(ti, false, nil)

	// Set up canned data for a valid listing.
	t.listing = []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type: fs.TypeFile,
			Name: "foo",
			Scores: []blob.Score{
				blob.ComputeScore([]byte("0")),
				blob.ComputeScore([]byte("1")),
			},
		},

		&fs.DirectoryEntry{
			Type: fs.TypeDirectory,
			Name: "bar",
			Scores: []blob.Score{
				blob.ComputeScore([]byte("2")),
			},
		},
	}

	var err error
	t.contents, err = repr.MarshalDir(t.listing)
	AssertEq(nil, err)

	t.score = blob.ComputeScore(t.contents)
	t.node = verify.FormatNodeName(true, t.score)
}

func (t *DirsTest) CallsBlobStore() {
	// Load
	ExpectCall(t.blobStore, "Load")(t.score).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.visitor.Visit(t.ctx, t.node)
}

func (t *DirsTest) BlobStoreReturnsError() {
	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	_, err := t.visitor.Visit(t.ctx, t.node)

	ExpectThat(err, Error(HasSubstr("Load")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *DirsTest) BlobIsJunk() {
	// Set up junk contents.
	t.contents = append(t.contents, 'a')
	t.score = blob.ComputeScore(t.contents)
	t.node = verify.FormatNodeName(true, t.score)

	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(t.contents, nil))

	// Call
	_, err := t.visitor.Visit(t.ctx, t.node)

	ExpectThat(err, Error(HasSubstr(t.score.Hex())))
	ExpectThat(err, Error(HasSubstr("UnmarshalDir")))
}

func (t *DirsTest) UnknownEntryType() {
	// Set up a listing with a junk entry type.
	t.listing = []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type: fs.TypeSymlink,
			Name: "foo",
			Scores: []blob.Score{
				blob.ComputeScore([]byte("0")),
			},
		},
	}

	var err error
	t.contents, err = repr.MarshalDir(t.listing)
	AssertEq(nil, err)

	t.score = blob.ComputeScore(t.contents)
	t.node = verify.FormatNodeName(true, t.score)

	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(t.contents, nil))

	// Call
	_, err = t.visitor.Visit(t.ctx, t.node)

	ExpectThat(err, Error(HasSubstr("entry type")))
	ExpectThat(err, Error(HasSubstr(fmt.Sprintf("%d", fs.TypeSymlink))))
}

func (t *DirsTest) ReturnsAppropriateAdjacentNodes() {
	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(t.contents, nil))

	// Call
	adjacent, err := t.visitor.Visit(t.ctx, t.node)
	AssertEq(nil, err)

	var expected []interface{}
	for _, entry := range t.listing {
		dir := entry.Type == fs.TypeDirectory
		for _, score := range entry.Scores {
			expected = append(expected, verify.FormatNodeName(dir, score))
		}
	}

	ExpectThat(adjacent, ElementsAre(expected...))
}

////////////////////////////////////////////////////////////////////////
// Files, no read
////////////////////////////////////////////////////////////////////////

type FilesLiteTest struct {
	superTest

	contents []byte

	knownScore blob.Score
	knownNode  string

	unknownScore blob.Score
	unknownNode  string
}

var _ SetUpInterface = &FilesLiteTest{}

func init() { RegisterTestSuite(&FilesLiteTest{}) }

func (t *FilesLiteTest) setUp(ti *TestInfo, readFiles bool) {
	t.contents = []byte("foobarbaz")

	t.knownScore = blob.ComputeScore(t.contents)
	t.knownNode = verify.FormatNodeName(false, t.knownScore)

	t.unknownScore = blob.ComputeScore(append(t.contents, 'a'))
	t.unknownNode = verify.FormatNodeName(false, t.unknownScore)

	t.superTest.setUp(ti, readFiles, []blob.Score{t.knownScore})
}

func (t *FilesLiteTest) SetUp(ti *TestInfo) {
	t.setUp(ti, false)
}

func (t *FilesLiteTest) ScoreNotInList() {
	// Call
	_, err := t.visitor.Visit(t.ctx, t.unknownNode)

	ExpectThat(err, Error(HasSubstr("Unknown score")))
	ExpectThat(err, Error(HasSubstr(t.unknownScore.Hex())))
}

func (t *FilesLiteTest) ScoreIsInList() {
	// Call
	adjacent, err := t.visitor.Visit(t.ctx, t.knownNode)

	AssertEq(nil, err)
	ExpectThat(adjacent, ElementsAre())
}

////////////////////////////////////////////////////////////////////////
// Files, read
////////////////////////////////////////////////////////////////////////

type FilesFullTest struct {
	FilesLiteTest
}

var _ SetUpInterface = &FilesFullTest{}

func init() { RegisterTestSuite(&FilesFullTest{}) }

func (t *FilesFullTest) SetUp(ti *TestInfo) {
	t.FilesLiteTest.setUp(ti, true)
}

func (t *FilesFullTest) CallsBlobStore() {
	AssertFalse(true, "TODO")
}

func (t *FilesFullTest) BlobStoreReturnsError() {
	AssertFalse(true, "TODO")
}

func (t *FilesFullTest) BlobStoreSucceeds() {
	AssertFalse(true, "TODO")
}
<commit_msg>Fixed a test bug.<commit_after>// Copyright 2015 Aaron Jacobs. All Rights Reserved.
// Author: aaronjjacobs@gmail.com (Aaron Jacobs)
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package verify_test

import (
	"errors"
	"fmt"
	"testing"

	"golang.org/x/net/context"

	"github.com/jacobsa/comeback/blob"
	"github.com/jacobsa/comeback/blob/mock"
	"github.com/jacobsa/comeback/fs"
	"github.com/jacobsa/comeback/graph"
	"github.com/jacobsa/comeback/repr"
	"github.com/jacobsa/comeback/verify"
	. "github.com/jacobsa/oglematchers"
	. "github.com/jacobsa/oglemock"
	. "github.com/jacobsa/ogletest"
)

func TestVisitor(t *testing.T) { RunTests(t) }

////////////////////////////////////////////////////////////////////////
// Boilerplate
////////////////////////////////////////////////////////////////////////

const namePrefix = "blobs/"

type superTest struct {
	ctx       context.Context
	blobStore mock_blob.MockStore
	visitor   graph.Visitor
}

func (t *superTest) setUp(
	ti *TestInfo,
	readFiles bool,
	allScores []blob.Score) {
	t.ctx = ti.Ctx
	t.blobStore = mock_blob.NewMockStore(ti.MockController, "blobStore")
	t.visitor = verify.NewVisitor(readFiles, allScores, t.blobStore)
}

////////////////////////////////////////////////////////////////////////
// Common
////////////////////////////////////////////////////////////////////////

type CommonTest struct {
	superTest
}

var _ SetUpInterface = &CommonTest{}

func init() { RegisterTestSuite(&CommonTest{}) }

func (t *CommonTest) SetUp(ti *TestInfo) {
	t.superTest.setUp(ti, false, nil)
}

func (t *CommonTest) InvalidNodeName() {
	_, err := t.visitor.Visit(t.ctx, "taco")

	ExpectThat(err, Error(HasSubstr("ParseNodeName")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

////////////////////////////////////////////////////////////////////////
// Directories
////////////////////////////////////////////////////////////////////////

type DirsTest struct {
	superTest

	listing  []*fs.DirectoryEntry
	contents []byte
	score    blob.Score
	node     string
}

var _ SetUpInterface = &DirsTest{}

func init() { RegisterTestSuite(&DirsTest{}) }

func (t *DirsTest) SetUp(ti *TestInfo) {
	t.superTest.setUp(ti, false, nil)

	// Set up canned data for a valid listing.
	t.listing = []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type: fs.TypeFile,
			Name: "foo",
			Scores: []blob.Score{
				blob.ComputeScore([]byte("0")),
				blob.ComputeScore([]byte("1")),
			},
		},

		&fs.DirectoryEntry{
			Type: fs.TypeDirectory,
			Name: "bar",
			Scores: []blob.Score{
				blob.ComputeScore([]byte("2")),
			},
		},
	}

	var err error
	t.contents, err = repr.MarshalDir(t.listing)
	AssertEq(nil, err)

	t.score = blob.ComputeScore(t.contents)
	t.node = verify.FormatNodeName(true, t.score)
}

func (t *DirsTest) CallsBlobStore() {
	// Load
	ExpectCall(t.blobStore, "Load")(t.score).
		WillOnce(Return(nil, errors.New("")))

	// Call
	t.visitor.Visit(t.ctx, t.node)
}

func (t *DirsTest) BlobStoreReturnsError() {
	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(nil, errors.New("taco")))

	// Call
	_, err := t.visitor.Visit(t.ctx, t.node)

	ExpectThat(err, Error(HasSubstr("Load")))
	ExpectThat(err, Error(HasSubstr("taco")))
}

func (t *DirsTest) BlobIsJunk() {
	// Set up junk contents.
	t.contents = append(t.contents, 'a')
	t.score = blob.ComputeScore(t.contents)
	t.node = verify.FormatNodeName(true, t.score)

	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(t.contents, nil))

	// Call
	_, err := t.visitor.Visit(t.ctx, t.node)

	ExpectThat(err, Error(HasSubstr(t.score.Hex())))
	ExpectThat(err, Error(HasSubstr("UnmarshalDir")))
}

func (t *DirsTest) UnknownEntryType() {
	// Set up a listing with a junk entry type.
	t.listing = []*fs.DirectoryEntry{
		&fs.DirectoryEntry{
			Type: fs.TypeSymlink,
			Name: "foo",
			Scores: []blob.Score{
				blob.ComputeScore([]byte("0")),
			},
		},
	}

	var err error
	t.contents, err = repr.MarshalDir(t.listing)
	AssertEq(nil, err)

	t.score = blob.ComputeScore(t.contents)
	t.node = verify.FormatNodeName(true, t.score)

	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(t.contents, nil))

	// Call
	_, err = t.visitor.Visit(t.ctx, t.node)

	ExpectThat(err, Error(HasSubstr("entry type")))
	ExpectThat(err, Error(HasSubstr(fmt.Sprintf("%d", fs.TypeSymlink))))
}

func (t *DirsTest) ReturnsAppropriateAdjacentNodes() {
	// Load
	ExpectCall(t.blobStore, "Load")(Any()).
		WillOnce(Return(t.contents, nil))

	// Call
	adjacent, err := t.visitor.Visit(t.ctx, t.node)
	AssertEq(nil, err)

	var expected []interface{}
	for _, entry := range t.listing {
		dir := entry.Type == fs.TypeDirectory
		for _, score := range entry.Scores {
			expected = append(expected, verify.FormatNodeName(dir, score))
		}
	}

	ExpectThat(adjacent, ElementsAre(expected...))
}

////////////////////////////////////////////////////////////////////////
// Files, no read
////////////////////////////////////////////////////////////////////////

type filesCommonTest struct {
	superTest

	contents []byte

	knownScore blob.Score
	knownNode  string

	unknownScore blob.Score
	unknownNode  string
}

func (t *filesCommonTest) setUp(ti *TestInfo, readFiles bool) {
	t.contents = []byte("foobarbaz")

	t.knownScore = blob.ComputeScore(t.contents)
	t.knownNode = verify.FormatNodeName(false, t.knownScore)

	t.unknownScore = blob.ComputeScore(append(t.contents, 'a'))
	t.unknownNode = verify.FormatNodeName(false, t.unknownScore)

	t.superTest.setUp(ti, readFiles, []blob.Score{t.knownScore})
}

type FilesLiteTest struct {
	filesCommonTest

	contents []byte

	knownScore blob.Score
	knownNode  string

	unknownScore blob.Score
	unknownNode  string
}

var _ SetUpInterface = &FilesLiteTest{}

func init() { RegisterTestSuite(&FilesLiteTest{}) }

func (t *FilesLiteTest) SetUp(ti *TestInfo) {
	t.filesCommonTest.setUp(ti, false)
}

func (t *FilesLiteTest) ScoreNotInList() {
	// Call
	_, err := t.visitor.Visit(t.ctx, t.unknownNode)

	ExpectThat(err, Error(HasSubstr("Unknown score")))
	ExpectThat(err, Error(HasSubstr(t.unknownScore.Hex())))
}

func (t *FilesLiteTest) ScoreIsInList() {
	// Call
	adjacent, err := t.visitor.Visit(t.ctx, t.knownNode)

	AssertEq(nil, err)
	ExpectThat(adjacent, ElementsAre())
}

////////////////////////////////////////////////////////////////////////
// Files, read
////////////////////////////////////////////////////////////////////////

type FilesFullTest struct {
	filesCommonTest
}

var _ SetUpInterface = &FilesFullTest{}

func init() { RegisterTestSuite(&FilesFullTest{}) }

func (t *FilesFullTest) SetUp(ti *TestInfo) {
	t.filesCommonTest.setUp(ti, true)
}

func (t *FilesFullTest) CallsBlobStore() {
	AssertFalse(true, "TODO")
}

func (t *FilesFullTest) BlobStoreReturnsError() {
	AssertFalse(true, "TODO")
}

func (t *FilesFullTest) BlobStoreSucceeds() {
	AssertFalse(true, "TODO")
}
<|endoftext|>


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

package valuelabel

import (
	"bytes"
	"encoding/json"
	"math"
	"sort"

	"github.com/corestoreio/csfw/util"
	"github.com/juju/errgo"
)

// Slice type is returned by the SourceModel.Options() interface
type Slice []Pair

// Options satisfies the SourceModeller interface
func (s Slice) Options() Slice { return s }

// SortByLabel sorts by label in asc or desc direction
func (s Slice) SortByLabel(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByLabel{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByValue sorts by value in asc or desc direction. The underlying value
// will be converted to a string. You might expect strange results when sorting
// integers or other non-strings.
func (s Slice) SortByValue(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByValue{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByInt sorts by field Int in asc or desc direction
func (s Slice) SortByInt(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByInt{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByFloat64 sorts by field Float64 in asc or desc direction
func (s Slice) SortByFloat64(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByFloat64{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByBool sorts by field Bool in asc or desc direction
func (s Slice) SortByBool(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByBool{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// Len returns the length of the slice
func (s Slice) Len() int { return len(s) }

// Swap swaps elements. Will panic when slice index does not exists.
func (s Slice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// ToJSON returns a JSON string, convenience function.
func (s Slice) ToJSON() (string, error) {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(s); err != nil {
		if PkgLog.IsDebug() {
			PkgLog.Debug("config.ValueLabelSlice.ToJSON.Encode", "err", err, "slice", s)
		}
		return "", errgo.Mask(err)
	}
	return buf.String(), nil
}

// ContainsKeyString checks if k has an entry as a key.
func (s Slice) ContainsKeyString(k string) bool {
	for _, p := range s {
		if p.String == k {
			return true
		}
	}
	return false
}

// ContainsKeyInt checks if k has an entry as a key.
func (s Slice) ContainsKeyInt(k int) bool {
	for _, p := range s {
		if p.Int == k {
			return true
		}
	}
	return false
}

// ContainsKeyFloat64 checks if k has an entry as a key.
func (s Slice) ContainsKeyFloat64(k float64) bool {
	for _, p := range s {
		abs := math.Abs(p.Float64 - k)
		if abs >= 0 && abs < 0.000001 {
			return true
		}
	}
	return false
}

// ContainsLabel checks if k has an entry as a label.
func (s Slice) ContainsLabel(l string) bool {
	for _, p := range s {
		if p.Label() == l {
			return true
		}
	}
	return false
}

type (
	vlSortByLabel struct {
		Slice
	}
	vlSortByValue struct {
		Slice
	}
	vlSortByInt struct {
		Slice
	}
	vlSortByFloat64 struct {
		Slice
	}
	vlSortByBool struct {
		Slice
	}
)

func (v vlSortByLabel) Less(i, j int) bool {
	return v.Slice[i].Label() < v.Slice[j].Label()
}

func (v vlSortByValue) Less(i, j int) bool {
	return v.Slice[i].Value() < v.Slice[j].Value()
}

func (v vlSortByInt) Less(i, j int) bool {
	return v.Slice[i].Int < v.Slice[j].Int
}

func (v vlSortByFloat64) Less(i, j int) bool {
	return v.Slice[i].Float64 < v.Slice[j].Float64
}

func (v vlSortByBool) Less(i, j int) bool {
	if !v.Slice[i].Bool {
		return v.Slice[i].Label() < v.Slice[j].Label()
	}
	return v.Slice[i].Bool && v.Slice[i].Label() < v.Slice[j].Label()
}
<commit_msg>config/valuelabel: Remove Options() from Slice<commit_after>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package valuelabel

import (
	"bytes"
	"encoding/json"
	"math"
	"sort"

	"github.com/corestoreio/csfw/util"
	"github.com/juju/errgo"
)

// Slice type is returned by the SourceModel.Options() interface
type Slice []Pair

// SortByLabel sorts by label in asc or desc direction
func (s Slice) SortByLabel(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByLabel{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByValue sorts by value in asc or desc direction. The underlying value
// will be converted to a string. You might expect strange results when sorting
// integers or other non-strings.
func (s Slice) SortByValue(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByValue{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByInt sorts by field Int in asc or desc direction
func (s Slice) SortByInt(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByInt{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByFloat64 sorts by field Float64 in asc or desc direction
func (s Slice) SortByFloat64(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByFloat64{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// SortByBool sorts by field Bool in asc or desc direction
func (s Slice) SortByBool(d util.SortDirection) Slice {
	var si sort.Interface
	si = vlSortByBool{s}
	if d == util.SortDesc {
		si = sort.Reverse(si)
	}
	sort.Sort(si)
	return s
}

// Len returns the length of the slice
func (s Slice) Len() int { return len(s) }

// Swap swaps elements. Will panic when slice index does not exists.
func (s Slice) Swap(i, j int) { s[i], s[j] = s[j], s[i] }

// ToJSON returns a JSON string, convenience function.
func (s Slice) ToJSON() (string, error) {
	var buf bytes.Buffer
	if err := json.NewEncoder(&buf).Encode(s); err != nil {
		if PkgLog.IsDebug() {
			PkgLog.Debug("config.ValueLabelSlice.ToJSON.Encode", "err", err, "slice", s)
		}
		return "", errgo.Mask(err)
	}
	return buf.String(), nil
}

// ContainsKeyString checks if k has an entry as a key.
func (s Slice) ContainsKeyString(k string) bool {
	for _, p := range s {
		if p.String == k {
			return true
		}
	}
	return false
}

// ContainsKeyInt checks if k has an entry as a key.
func (s Slice) ContainsKeyInt(k int) bool {
	for _, p := range s {
		if p.Int == k {
			return true
		}
	}
	return false
}

// ContainsKeyFloat64 checks if k has an entry as a key.
func (s Slice) ContainsKeyFloat64(k float64) bool {
	for _, p := range s {
		abs := math.Abs(p.Float64 - k)
		if abs >= 0 && abs < 0.000001 {
			return true
		}
	}
	return false
}

// ContainsLabel checks if k has an entry as a label.
func (s Slice) ContainsLabel(l string) bool {
	for _, p := range s {
		if p.Label() == l {
			return true
		}
	}
	return false
}

type (
	vlSortByLabel struct {
		Slice
	}
	vlSortByValue struct {
		Slice
	}
	vlSortByInt struct {
		Slice
	}
	vlSortByFloat64 struct {
		Slice
	}
	vlSortByBool struct {
		Slice
	}
)

func (v vlSortByLabel) Less(i, j int) bool {
	return v.Slice[i].Label() < v.Slice[j].Label()
}

func (v vlSortByValue) Less(i, j int) bool {
	return v.Slice[i].Value() < v.Slice[j].Value()
}

func (v vlSortByInt) Less(i, j int) bool {
	return v.Slice[i].Int < v.Slice[j].Int
}

func (v vlSortByFloat64) Less(i, j int) bool {
	return v.Slice[i].Float64 < v.Slice[j].Float64
}

func (v vlSortByBool) Less(i, j int) bool {
	if !v.Slice[i].Bool {
		return v.Slice[i].Label() < v.Slice[j].Label()
	}
	return v.Slice[i].Bool && v.Slice[i].Label() < v.Slice[j].Label()
}
<|endoftext|>


<commit_before>package pinger

import (
	"encoding/binary"
	"errors"
	"net"
	"os"
	"sync"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

const (
	protocolICMP = 1
)

type (
	PingResult struct {
		Times          int
		ReplyRatio     float64
		LatencyAverage float64
	}

	bySeq struct {
		ip      net.IP
		ipIdx   int       // set when start echo
		echoAt  time.Time //
		latency float64   // updated at each reply
	}

	byIP struct {
		times      int     // updated at each reply
		latencySum float64 //
	}
)

var (
	ErrParseReply   = errors.New("error: Parse ICMP message expect ipv4.ICMPTypeEchoReply")
	ErrDisconnected = errors.New("error: No active connection")
	ErrSendEcho     = errors.New("error: Failed to send ICMP echo message")

	conn struct {
		sync.Mutex
		*icmp.PacketConn
		onReply map[int]func(seq int) // Map seq to callback function
		seq     int
	}
)

func Start() error {
	syncConnDestroy()
	return syncConnCreate()
}

func Stop() {
	syncConnDestroy()
}

func bgReply() error {
	//if err := conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
	//	stop()
	//	return err
	//}

	for {
		if seq, e := syncReplyRead(); e == nil {
			// Notes:
			// So far the only identified error is the "destination unreachable".
			// So it is safe to ignore the error.
			syncReplyCall(seq)
		}
	}
}

func Ping(pingList []net.IP, times int, durBetween, durLast time.Duration) ([]*PingResult, error) {
	rst := make([]*PingResult, len(pingList))
	rstByIP := make([]*byIP, len(pingList))
	rstBySeq := make(map[int]*bySeq)

	for ipIdx := 0; ipIdx < len(pingList); ipIdx++ {
		rst[ipIdx] = new(PingResult)
		rstByIP[ipIdx] = new(byIP)
	}

	done := make(chan bool, 1)

	for loop := 0; loop < times; loop++ {
		for ipIdx, ip := range pingList {
			err := syncPingEcho(ip,
				func(seq int, ip net.IP) { // On echo
					rstBySeq[seq] = &bySeq{ip: ip, ipIdx: ipIdx, echoAt: time.Now()}
				},
				func(seq int) { // On reply
					rs, ok := rstBySeq[seq]
					if !ok {
						return
					}
					//fmt.Println("Reply:", rs.ip)
					rs.latency = time.Since(rs.echoAt).Seconds()

					rip := rstByIP[rs.ipIdx]
					rip.times++
					rip.latencySum += rs.latency

					delete(rstBySeq, seq)
					if len(rstBySeq) == 0 {
						select {
						case done <- true:
						default:
							return
						}
					}
				})
			if err != nil {
				return nil, err
			}
		}
		time.Sleep(durBetween)
	}

	go func() {
		time.Sleep(durLast)
		syncPingClearSeq(rstBySeq)
		done <- true
	}()

	<-done
	return syncPingSummarize(rst, rstByIP, times)
}

func syncConnCreate() error {
	conn.Lock()
	defer conn.Unlock()
	c, err := icmp.ListenPacket("ip4:icmp", "0.0.0.0")
	if err != nil {
		return err
	}

	conn.PacketConn = c
	conn.onReply = make(map[int]func(int))
	go bgReply()

	return nil
}

func syncConnDestroy() {
	conn.Lock()
	defer conn.Unlock()
	if conn.PacketConn == nil {
		return // Already disconnected
	}
	conn.Close()

	conn.PacketConn = nil
	conn.onReply = nil
}

func syncReplyRead() (seq int, err error) {
	conn.Lock()
	c := conn.PacketConn
	if c == nil {
		conn.Unlock()
		return 0, ErrDisconnected
	}
	conn.Unlock() // no need to defer because read should not concurrent

	readBuffer := make([]byte, 1500)
	n, peer, err := c.ReadFrom(readBuffer)
	if err != nil {
		//switch err.(type) {
		//case *net.OpError:
		//	// Timeout error may happen
		//	if strings.Contains(err.Error(), "i/o timeout") {
		//		return ErrReadTimeout
		//	}
		//default:
		//}
		return 0, err
	}

	msg, err := icmp.ParseMessage(protocolICMP, readBuffer[:n])
	if err != nil {
		return 0, err
	}

	switch msg.Type {
	case ipv4.ICMPTypeEchoReply:
	default:
		// Notes:
		// So far the only identified error is the "destination unreachable".
		// So it is safe to ignore the error.

		//fmt.Println("error: Got", msg, "from", peer)
		_ = peer
		return 0, ErrParseReply
	}

	b, err := msg.Body.Marshal(protocolICMP)
	if err != nil {
		return 0, err
	}

	return int(binary.BigEndian.Uint16(b[2:4])), nil
}

func syncReplyCall(seq int) {
	conn.Lock()
	defer conn.Unlock()
	if onReply, ok := conn.onReply[seq]; ok {
		onReply(seq)
		delete(conn.onReply, seq)
	}
}

func syncPingEcho(ip net.IP, onEcho func(seq int, ip net.IP), onReply func(seq int)) error {
	conn.Lock()
	defer conn.Unlock()

	onEcho(conn.seq, ip)
	conn.onReply[conn.seq] = onReply

	msg := icmp.Message{
		Type: ipv4.ICMPTypeEcho, Code: 0,
		Body: &icmp.Echo{
			ID:   os.Getpid() & 0xffff, Seq: conn.seq,
			Data: []byte("HELLO-R-U-THERE"),
		},
	}
	conn.seq++

	b, err := msg.Marshal(nil)
	if err != nil {
		return err
	}

	n, err := conn.WriteTo(b, &net.IPAddr{IP: ip})
	if err != nil {
		return err
	}
	if n != len(b) {
		return ErrSendEcho
	}

	return nil
}

func syncPingClearSeq(rstBySeq map[int]*bySeq) {
	conn.Lock()
	defer conn.Unlock()
	for seq := range rstBySeq {
		delete(rstBySeq, seq)
	}
}

func syncPingSummarize(rst []*PingResult, rstByIP []*byIP, times int) ([]*PingResult, error) {
	conn.Lock()
	defer conn.Unlock()
	for ipIdx, r := range rst {
		rip := rstByIP[ipIdx]
		if rip.times == 0 {
			continue
		}
		r.Times = rip.times
		r.ReplyRatio = float64(rip.times) / float64(times)
		r.LatencyAverage = rip.latencySum / float64(rip.times)
	}
	return rst, nil
}
<commit_msg>Remove bgReply() returns because it is not necessary. Improve Ping() in how it handle message read.<commit_after>package pinger

import (
	"encoding/binary"
	"errors"
	"fmt"
	"net"
	"os"
	"sync"
	"time"

	"golang.org/x/net/icmp"
	"golang.org/x/net/ipv4"
)

const (
	protocolICMP = 1
)

type (
	PingResult struct {
		Times          int
		ReplyRatio     float64
		LatencyAverage float64
	}

	bySeq struct {
		ip      net.IP
		ipIdx   int       // set when start echo
		echoAt  time.Time //
		latency float64   // updated at each reply
	}

	byIP struct {
		times      int     // updated at each reply
		latencySum float64 //
	}
)

var (
	ErrParseReply   = errors.New("error: Parse ICMP message expect ipv4.ICMPTypeEchoReply")
	ErrDisconnected = errors.New("error: No active connection")
	ErrSendEcho     = errors.New("error: Failed to send ICMP echo message")

	readBuffer = make([]byte, 1500)

	conn struct {
		sync.Mutex
		*icmp.PacketConn
		onReply map[int]func(seq int) // Map seq to callback function
		seq     int
	}
)

func Start() error {
	syncConnDestroy()
	return syncConnCreate()
}

func Stop() {
	syncConnDestroy()
}

func bgReply() {
	//if err := conn.SetReadDeadline(time.Now().Add(time.Second)); err != nil {
	//	stop()
	//	return err
	//}

	for {
		seq, e := syncReplyRead()
		if e != nil {
			fmt.Println(e)
			continue
		}
		syncReplyCall(seq)
	}
}

func Ping(pingList []net.IP, times int, durBetween, durLast time.Duration) ([]*PingResult, error) {
	rst := make([]*PingResult, len(pingList))
	rstByIP := make([]*byIP, len(pingList))
	rstBySeq := make(map[int]*bySeq)

	for ipIdx := 0; ipIdx < len(pingList); ipIdx++ {
		rst[ipIdx] = new(PingResult)
		rstByIP[ipIdx] = new(byIP)
	}

	done := make(chan bool, 1)

	for loop := 0; loop < times; loop++ {
		for ipIdx, ip := range pingList {
			err := syncPingEcho(ip,
				func(seq int, ip net.IP) { // On echo
					rstBySeq[seq] = &bySeq{ip: ip, ipIdx: ipIdx, echoAt: time.Now()}
				},
				func(seq int) { // On reply
					rs, ok := rstBySeq[seq]
					if !ok {
						return
					}
					//fmt.Println("Reply:", rs.ip)
					rs.latency = time.Since(rs.echoAt).Seconds()

					rip := rstByIP[rs.ipIdx]
					rip.times++
					rip.latencySum += rs.latency

					delete(rstBySeq, seq)
					if len(rstBySeq) == 0 {
						select {
						case done <- true:
						default:
							return
						}
					}
				})
			if err != nil {
				return nil, err
			}
		}
		time.Sleep(durBetween)
	}

	go func() {
		time.Sleep(durLast)
		syncPingClearSeq(rstBySeq)
		done <- true
	}()

	<-done
	return syncPingSummarize(rst, rstByIP, times)
}

func syncConnCreate() error {
	conn.Lock()
	defer conn.Unlock()
	c, err := icmp.ListenPacket("ip4:icmp", "0.0.0.0")
	if err != nil {
		return err
	}

	conn.PacketConn = c
	conn.onReply = make(map[int]func(int))
	go bgReply()

	return nil
}

func syncConnDestroy() {
	conn.Lock()
	defer conn.Unlock()
	if conn.PacketConn == nil {
		return // Already disconnected
	}
	conn.Close()

	conn.PacketConn = nil
	conn.onReply = nil
}

func syncReplyRead() (seq int, err error) {
	conn.Lock()
	c := conn.PacketConn
	if c == nil {
		conn.Unlock()
		return 0, ErrDisconnected
	}
	conn.Unlock() // no need to defer because read should not concurrent

	n, _, err := c.ReadFrom(readBuffer)
	if err != nil {
		//switch err.(type) {
		//case *net.OpError:
		//	// Timeout error may happen
		//	if strings.Contains(err.Error(), "i/o timeout") {
		//		return ErrReadTimeout
		//	}
		//default:
		//}
		return 0, err
	}

	msg, err := icmp.ParseMessage(protocolICMP, readBuffer[:n])
	if err != nil {
		return 0, err
	}

	if msg.Type != ipv4.ICMPTypeEchoReply {
		// Not interested in other than ICMPTypeEchoReply message
		return 0, ErrParseReply
	}

	b, err := msg.Body.Marshal(protocolICMP)
	if err != nil {
		return 0, err
	}

	return int(binary.BigEndian.Uint16(b[2:4])), nil
}

func syncReplyCall(seq int) {
	conn.Lock()
	defer conn.Unlock()
	if onReply, ok := conn.onReply[seq]; ok {
		onReply(seq)
		delete(conn.onReply, seq)
	}
}

func syncPingEcho(ip net.IP, onEcho func(seq int, ip net.IP), onReply func(seq int)) error {
	conn.Lock()
	defer conn.Unlock()

	onEcho(conn.seq, ip)
	conn.onReply[conn.seq] = onReply

	msg := icmp.Message{
		Type: ipv4.ICMPTypeEcho, Code: 0,
		Body: &icmp.Echo{
			ID:   os.Getpid() & 0xffff, Seq: conn.seq,
			Data: []byte("HELLO-R-U-THERE"),
		},
	}
	conn.seq++

	b, err := msg.Marshal(nil)
	if err != nil {
		return err
	}

	n, err := conn.WriteTo(b, &net.IPAddr{IP: ip})
	if err != nil {
		return err
	}
	if n != len(b) {
		return ErrSendEcho
	}

	return nil
}

func syncPingClearSeq(rstBySeq map[int]*bySeq) {
	conn.Lock()
	defer conn.Unlock()
	for seq := range rstBySeq {
		delete(rstBySeq, seq)
	}
}

func syncPingSummarize(rst []*PingResult, rstByIP []*byIP, times int) ([]*PingResult, error) {
	conn.Lock()
	defer conn.Unlock()
	for ipIdx, r := range rst {
		rip := rstByIP[ipIdx]
		if rip.times == 0 {
			continue
		}
		r.Times = rip.times
		r.ReplyRatio = float64(rip.times) / float64(times)
		r.LatencyAverage = rip.latencySum / float64(rip.times)
	}
	return rst, nil
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	// "io"
	"flag"
	"io/ioutil"
	"strings"
	"time"
)

func get(hostname string, port int, path string, auth string, urls bool, verbose bool, timeout int) (rv bool, err error) {

	defer func() {
		if err := recover(); err != nil {
			return
		}
	}()

	rv = true

	if verbose {
		fmt.Fprintf(os.Stderr, "fetching:hostname:%s:\n", hostname)
	}

	res := &http.Response{}

	if urls {

		url := hostname
		res, err = http.Head(url)
		defer res.Body.Close()

		if err != nil {
			fmt.Println(err.Error())
			rv = false
			return
		}

	} else {

		client := &http.Client{Timeout: time.Duration(timeout) * time.Second}

		// had to allocate this or the SetBasicAuth will panic
		headers := make(map[string][]string)
		hostPort := fmt.Sprintf("%s:%d", hostname, port)

		if verbose {

			fmt.Fprintf(os.Stderr, "adding hostPort:%s:%d:path:%s:\n", hostname, port, path)

		}
		req := &http.Request{
			Method: "HEAD",
			// Host:  hostPort,
			URL: &url.URL{
				Host:   hostPort,
				Scheme: "http",
				Opaque: path,
			},
			Header: headers,
		}

		if auth != "" {

			up := strings.SplitN(auth, ":", 2)

			if verbose {

				fmt.Fprintf(os.Stderr, "Doing auth with:username:%s:password:%s:", up[0], up[1])

			}
			req.SetBasicAuth(up[0], up[1])

		}

		if verbose {

			dump, _ := httputil.DumpRequestOut(req, true)
			fmt.Fprintf(os.Stderr, "%s", dump)

		}

		res, err = client.Do(req)
		if err != nil {
			fmt.Println(err.Error())
			rv = false
			return
		}

		defer res.Body.Close()
		_, err = ioutil.ReadAll(res.Body)

	}

	if verbose {

		fmt.Println(res.Status)
		for k, v := range res.Header {
			fmt.Println(k+":", v)
		}

	}

	if res.StatusCode != http.StatusOK {
		rv = false
	}

	return
}

func main() {

	status := "OK"
	rv := 0
	name := "Bulk HTTP Check"
	bad := 0
	total := 0

	// this needs improvement. the number of spaces here has to equal the number of chars in the badHosts append line suffix
	badHosts := []byte("  ")

	verbose := flag.Bool("v", false, "verbose output")
	warn := flag.Int("w", 10, "warning level - number of non-200s or percentage of non-200s (default is numeric not percentage)")
	crit := flag.Int("c", 20, "critical level - number of non-200s or percentage of non-200s (default is numeric not percentage)")
	timeout := flag.Int("t", 2, "timeout in seconds - don't wait.  Do Head requests and don't wait.")
	pct := flag.Bool("pct", false, "interpret warming and critical levels are percentages")
	path := flag.String("path", "", "optional path to append to the input lines including the leading slash - these will not be urlencoded. This is ignored is the urls option is given.")
	file := flag.String("file", "", "input data source: a filename or '-' for STDIN.")
	port := flag.Int("port", 80, "optional port for the http request - ignored if urls is specified")
	urls := flag.Bool("urls", false, "Assume the input data is full urls - its normally a list of hostnames")
	auth := flag.String("auth", "", "Do basic auth with this username:passwd - ignored if urls is specified - make this use .netrc instead")
	checkName := flag.String("name", "", "a name to be included in the check output to distinguish the check output")

	flag.Usage = func() {

		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, `
	Read hostnames from a file or STDIN and do a single nagios check over
	them all.  Just check for 200s.  Warning and Critical are either
	percentages of the total, or a regular numeric thresholds.

	The output contains the hostname of any non-200 reporting hosts.

	Skip input lines that are commented out with shell style comments
	like /^#/.

	Do Head requests since we don't care about the content.  Make this
	optional some day.

	The -path is appended to the hostnames to make full URLs for the checks.

	If the -urls option is specified, then the input is assumed a complete URLs, like http://$hostname:$port/$path.

	Examples:

	./someCommand |  ./check_http_bulk  -w 1 -c 2 -path '/api/aliveness-test/%%2F/' -port 15672 -file - -auth zup:nuch

	./check_http_bulk -urls -file urls.txt

`)

		flag.PrintDefaults()
	}

	flag.Parse()

	if len(flag.Args()) > 0 {

		flag.Usage()
		os.Exit(3)

	}

	// it urls is specified, the input is full urls to be used enmasse and to be url encoded
	if *urls {
		*path = ""
	}

	if *checkName != "" {
		name = *checkName
	}

	defer func() {
		if err := recover(); err != nil {
			fmt.Println(name+" Unknown: ", err)
			os.Exit(3)
		}
	}()

	if file == nil || *file == "" {
		flag.Usage()
		os.Exit(3)
	}

	inputSource := os.Stdin

	if (*file)[0] != "-"[0] {

		var err error

		inputSource, err = os.Open(*file)

		if err != nil {

			fmt.Printf("Couldn't open the specified input file:%s:error:%v:\n\n", name, err)
			flag.Usage()
			os.Exit(3)

		}

	}

	scanner := bufio.NewScanner(inputSource)
	for scanner.Scan() {

		total++

		hostname := scanner.Text()

		if hostname[0] == "#"[0] {

			if *verbose {

				fmt.Printf("skipping:%s:\n", hostname)

			}

			continue
		}

		if *verbose {

			fmt.Printf("working on:%s:\n", hostname)

		}

		goodCheck, err := get(hostname, *port, *path, *auth, *urls, *verbose, *timeout)
		if err != nil {

			fmt.Printf("%s get error: %T %s %#v\n", name, err, err, err)

			continue

		}

		if !goodCheck {
			badHosts = append(badHosts, hostname...)
			badHosts = append(badHosts, ", "...)
			bad++
		}

	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
		status = "Unknown"
		rv = 3
	}

	if *pct {

		ratio := int(float64(bad)/float64(total)*100)

		if *verbose {

			fmt.Fprintf(os.Stderr, "ratio:%d:\n", ratio)

		}

		if ratio >= *crit {
			status = "Critical"
			rv = 1
		} else if ratio >= *warn {
			status = "Warning"
			rv = 2
		}

	} else {

		if bad >= *crit {
			status = "Critical"
			rv = 1
		} else if bad >= *warn {
			status = "Warning"
			rv = 2
		}

	}

	fmt.Printf("%s %s: %d of %d |%s\n", name, status, bad, total, badHosts[:len(badHosts)-2])
	os.Exit(rv)
}
<commit_msg>fix bug missed conx failures in error count<commit_after>package main

import (
	"bufio"
	"fmt"
	"net/http"
	"net/http/httputil"
	"net/url"
	"os"
	// "io"
	"flag"
	"io/ioutil"
	"strings"
	"time"
)

func get(hostname string, port int, path string, auth string, urls bool, verbose bool, timeout int) (rv bool, err error) {

	defer func() {
		if err := recover(); err != nil {
			return
		}
	}()

	rv = true

	if verbose {
		fmt.Fprintf(os.Stderr, "fetching:hostname:%s:\n", hostname)
	}

	res := &http.Response{}

	if urls {

		url := hostname
		res, err = http.Head(url)
		defer res.Body.Close()

		if err != nil {
			fmt.Println(err.Error())
			rv = false
			return
		}

	} else {

		client := &http.Client{Timeout: time.Duration(timeout) * time.Second}

		// had to allocate this or the SetBasicAuth will panic
		headers := make(map[string][]string)
		hostPort := fmt.Sprintf("%s:%d", hostname, port)

		if verbose {

			fmt.Fprintf(os.Stderr, "adding hostPort:%s:%d:path:%s:\n", hostname, port, path)

		}
		req := &http.Request{
			Method: "HEAD",
			// Host:  hostPort,
			URL: &url.URL{
				Host:   hostPort,
				Scheme: "http",
				Opaque: path,
			},
			Header: headers,
		}

		if auth != "" {

			up := strings.SplitN(auth, ":", 2)

			if verbose {

				fmt.Fprintf(os.Stderr, "Doing auth with:username:%s:password:%s:", up[0], up[1])

			}
			req.SetBasicAuth(up[0], up[1])

		}

		if verbose {

			dump, _ := httputil.DumpRequestOut(req, true)
			fmt.Fprintf(os.Stderr, "%s", dump)

		}

		res, err = client.Do(req)
		if err != nil {
			fmt.Println(err.Error())
			rv = false
			return
		}

		defer res.Body.Close()
		_, err = ioutil.ReadAll(res.Body)

	}

	if verbose {

		fmt.Println(res.Status)
		for k, v := range res.Header {
			fmt.Println(k+":", v)
		}

	}

	if res.StatusCode != http.StatusOK {
		rv = false
	}

	return
}

func main() {

	status := "OK"
	rv := 0
	name := "Bulk HTTP Check"
	bad := 0
	total := 0

	// this needs improvement. the number of spaces here has to equal the number of chars in the badHosts append line suffix
	badHosts := []byte("  ")

	verbose := flag.Bool("v", false, "verbose output")
	warn := flag.Int("w", 10, "warning level - number of non-200s or percentage of non-200s (default is numeric not percentage)")
	crit := flag.Int("c", 20, "critical level - number of non-200s or percentage of non-200s (default is numeric not percentage)")
	timeout := flag.Int("t", 2, "timeout in seconds - don't wait.  Do Head requests and don't wait.")
	pct := flag.Bool("pct", false, "interpret warming and critical levels are percentages")
	path := flag.String("path", "", "optional path to append to the input lines including the leading slash - these will not be urlencoded. This is ignored is the urls option is given.")
	file := flag.String("file", "", "input data source: a filename or '-' for STDIN.")
	port := flag.Int("port", 80, "optional port for the http request - ignored if urls is specified")
	urls := flag.Bool("urls", false, "Assume the input data is full urls - its normally a list of hostnames")
	auth := flag.String("auth", "", "Do basic auth with this username:passwd - ignored if urls is specified - make this use .netrc instead")
	checkName := flag.String("name", "", "a name to be included in the check output to distinguish the check output")

	flag.Usage = func() {

		fmt.Fprintf(os.Stderr, "Usage of %s:\n", os.Args[0])
		fmt.Fprintf(os.Stderr, `
	Read hostnames from a file or STDIN and do a single nagios check over
	them all.  Just check for 200s.  Warning and Critical are either
	percentages of the total, or a regular numeric thresholds.

	The output contains the hostname of any non-200 reporting hosts.

	Skip input lines that are commented out with shell style comments
	like /^#/.

	Do Head requests since we don't care about the content.  Make this
	optional some day.

	The -path is appended to the hostnames to make full URLs for the checks.

	If the -urls option is specified, then the input is assumed a complete URLs, like http://$hostname:$port/$path.

	Examples:

	./someCommand |  ./check_http_bulk  -w 1 -c 2 -path '/api/aliveness-test/%%2F/' -port 15672 -file - -auth zup:nuch

	./check_http_bulk -urls -file urls.txt

`)

		flag.PrintDefaults()
	}

	flag.Parse()

	if len(flag.Args()) > 0 {

		flag.Usage()
		os.Exit(3)

	}

	// it urls is specified, the input is full urls to be used enmasse and to be url encoded
	if *urls {
		*path = ""
	}

	if *checkName != "" {
		name = *checkName
	}

	defer func() {
		if err := recover(); err != nil {
			fmt.Println(name+" Unknown: ", err)
			os.Exit(3)
		}
	}()

	if file == nil || *file == "" {
		flag.Usage()
		os.Exit(3)
	}

	inputSource := os.Stdin

	if (*file)[0] != "-"[0] {

		var err error

		inputSource, err = os.Open(*file)

		if err != nil {

			fmt.Printf("Couldn't open the specified input file:%s:error:%v:\n\n", name, err)
			flag.Usage()
			os.Exit(3)

		}

	}

	scanner := bufio.NewScanner(inputSource)
	for scanner.Scan() {

		total++

		hostname := scanner.Text()

		if hostname[0] == "#"[0] {

			if *verbose {

				fmt.Printf("skipping:%s:\n", hostname)

			}

			continue
		}

		if *verbose {

			fmt.Printf("working on:%s:\n", hostname)

		}

		goodCheck, err := get(hostname, *port, *path, *auth, *urls, *verbose, *timeout)
		if err != nil {

			fmt.Printf("%s get error: %T %s %#v\n", name, err, err, err)
			badHosts = append(badHosts, hostname...)
			badHosts = append(badHosts, ", "...)
			bad++

			continue

		}

		if !goodCheck {
			badHosts = append(badHosts, hostname...)
			badHosts = append(badHosts, ", "...)
			bad++
		}

	}
	if err := scanner.Err(); err != nil {
		fmt.Fprintln(os.Stderr, "reading standard input:", err)
		status = "Unknown"
		rv = 3
	}

	if *pct {

		ratio := int(float64(bad)/float64(total)*100)

		if *verbose {

			fmt.Fprintf(os.Stderr, "ratio:%d:\n", ratio)

		}

		if ratio >= *crit {
			status = "Critical"
			rv = 1
		} else if ratio >= *warn {
			status = "Warning"
			rv = 2
		}

	} else {

		if bad >= *crit {
			status = "Critical"
			rv = 1
		} else if bad >= *warn {
			status = "Warning"
			rv = 2
		}

	}

	fmt.Printf("%s %s: %d of %d |%s\n", name, status, bad, total, badHosts[:len(badHosts)-2])
	os.Exit(rv)
}
<|endoftext|>


<commit_before>package http

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"github.com/viant/endly"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/data"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

//ServiceID represents http runner service id.
const ServiceID = "http/runner"

//PreviousTripStateKey keys to store previous request details for multi trip HTTP Send request in context state
const PreviousTripStateKey = "previous"

type service struct {
	*endly.AbstractService
}

func (s *service) processResponse(context *endly.Context, sendRequest *SendRequest, sendHTTPRequest *Request, response *Response, httpResponse *http.Response, isBase64Encoded bool, extracted map[string]string) (string, error) {
	response.Header = make(map[string][]string)
	copyHeaders(httpResponse.Header, response.Header)
	readBody(httpResponse, response, isBase64Encoded)
	if sendHTTPRequest.ResponseUdf != "" {
		transformed, err := endly.TransformWithUDF(context, sendHTTPRequest.ResponseUdf, sendHTTPRequest.URL, response.Body)
		if err != nil {
			return "", err
		}
		if toolbox.IsMap(transformed) {
			response.Body, _ = toolbox.AsJSONText(transformed)
		} else {
			response.Body = toolbox.AsString(transformed)
		}
	}

	var responseBody = replaceResponseBodyIfNeeded(sendHTTPRequest, response.Body)
	return responseBody, nil
}

func (s *service) sendRequest(context *endly.Context, client *http.Client, sendHTTPRequest *Request, sessionCookies *Cookies, sendGroupRequest *SendRequest, sendGroupResponse *SendResponse) error {
	var err error
	var state = context.State()
	cookies := state.GetMap("cookies")
	var reader io.Reader
	var isBase64Encoded = false
	sendHTTPRequest = sendHTTPRequest.Expand(context)
	var body []byte
	var ok bool
	if len(sendHTTPRequest.Body) > 0 {
		body = []byte(sendHTTPRequest.Body)
		if sendHTTPRequest.RequestUdf != "" {
			transformed, err := endly.TransformWithUDF(context, sendHTTPRequest.RequestUdf, sendHTTPRequest.URL, string(body))
			if err != nil {
				return err
			}
			if body, ok = transformed.([]byte); !ok {
				body = []byte(toolbox.AsString(transformed))
			}
		}
		isBase64Encoded = strings.HasPrefix(string(body), "base64:")
		body, err = endly.FromPayload(string(body))
		if err != nil {
			return err
		}
		reader = bytes.NewReader(body)
	}

	httpRequest, err := http.NewRequest(strings.ToUpper(sendHTTPRequest.Method), sendHTTPRequest.URL, reader)
	if err != nil {
		return err
	}

	copyHeaders(sendHTTPRequest.Header, httpRequest.Header)
	sessionCookies.SetHeader(sendHTTPRequest.Header)
	sendHTTPRequest.Cookies.SetHeader(httpRequest.Header)

	response := &Response{}
	sendGroupResponse.Responses = append(sendGroupResponse.Responses, response)
	startEvent := s.Begin(context, sendHTTPRequest)
	repeatable := sendHTTPRequest.Repeatable.Get()

	var httpResponse *http.Response
	var responseBody string
	var bodyCache []byte
	var useCachedBody = repeatable.Repeat > 1 && httpRequest.ContentLength > 0
	if useCachedBody {
		bodyCache, err = ioutil.ReadAll(httpRequest.Body)
		if err != nil {
			return err
		}
	}

	handler := func() (interface{}, error) {
		if useCachedBody {
			httpRequest.Body = ioutil.NopCloser(bytes.NewReader(bodyCache))
		}

		httpResponse, err = client.Do(httpRequest)
		if err != nil {
			return nil, err
		}
		responseBody, err = s.processResponse(context, sendGroupRequest, sendHTTPRequest, response, httpResponse, isBase64Encoded, sendGroupResponse.Extracted)
		return responseBody, err
	}

	err = repeatable.Run(s.AbstractService, "HTTPRunner", context, handler, sendGroupResponse.Extracted)
	if err != nil {
		return err
	}
	var responseCookies Cookies = httpResponse.Cookies()
	response.Cookies = responseCookies.IndexByName()
	for k, cookie := range response.Cookies {
		cookies.Put(k, cookie.Value)
	}
	sessionCookies.AddCookies(responseCookies...)

	endEvent := s.End(context)(startEvent, toolbox.Pairs("response", response))

	var previous = state.GetMap(PreviousTripStateKey)
	if previous == nil {
		previous = data.NewMap()
	}
	response.Code = httpResponse.StatusCode
	response.TimeTakenMs = int(startEvent.Timestamp.Sub(endEvent.Timestamp) / time.Millisecond)

	if toolbox.IsCompleteJSON(responseBody) {
		response.JSONBody, err = toolbox.JSONToMap(responseBody)
		if err == nil && sendHTTPRequest.Repeatable != nil {
			_ = sendHTTPRequest.Variables.Apply(data.Map(response.JSONBody), previous)
		}
	}

	for k, v := range sendGroupResponse.Extracted {
		var expanded = previous.Expand(v)
		previous[k] = state.Expand(expanded)
	}

	err = repeatable.Variables.Apply(previous, previous)
	if err != nil {
		return err
	}
	if len(previous) > 0 {
		state.Put(PreviousTripStateKey, previous)
	}
	if sendHTTPRequest.MatchBody != "" {
		return nil
	}

	for _, candidate := range sendGroupRequest.Requests {
		if candidate.MatchBody != "" && strings.Contains(response.Body, candidate.MatchBody) {
			err = s.sendRequest(context, client, candidate, sessionCookies, sendGroupRequest, sendGroupResponse)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func replaceResponseBodyIfNeeded(sendHTTPRequest *Request, responseBody string) string {
	if len(sendHTTPRequest.Replace) > 0 {
		for k, v := range sendHTTPRequest.Replace {
			responseBody = strings.Replace(responseBody, k, v, len(responseBody))
		}
	}
	return responseBody
}

func (s *service) send(context *endly.Context, request *SendRequest) (*SendResponse, error) {
	client, err := toolbox.NewHttpClient(request.Options...)
	if err != nil {
		return nil, fmt.Errorf("failed to send req: %v", err)
	}
	defer s.resetContext(context, request)
	var result = &SendResponse{
		Responses: make([]*Response, 0),
		Extracted: make(map[string]string),
	}
	var sessionCookies Cookies = make([]*http.Cookie, 0)
	var state = context.State()
	if !state.Has("cookies") {
		state.Put("cookies", data.NewMap())
	}
	for _, req := range request.Requests {
		if req.MatchBody != "" {
			continue
		}
		err = s.sendRequest(context, client, req, &sessionCookies, request, result)
		if err != nil {
			return nil, err
		}
	}
	return result, nil

}
func readBody(httpResponse *http.Response, response *Response, expectBased64Encoded bool) {

	body, err := ioutil.ReadAll(httpResponse.Body)
	if err != nil {
		response.Error = fmt.Sprintf("%v", err)
		return
	}
	httpResponse.Body.Close()
	if expectBased64Encoded {
		buf := new(bytes.Buffer)
		encoder := base64.NewEncoder(base64.StdEncoding, buf)
		_, _ = encoder.Write(body)
		_ = encoder.Close()
		response.Body = "base64:" + string(buf.Bytes())

	} else {
		response.Body = string(body)
	}
}

func copyHeaders(source http.Header, target http.Header) {
	for key, values := range source {
		if _, has := target[key]; !has {
			target[key] = make([]string, 0)
		}
		if len(values) == 1 {
			target.Set(key, values[0])
		} else {

			for _, value := range values {
				target.Add(key, value)
			}
		}
	}
}

func copyExpandedHeaders(source http.Header, target http.Header, context *endly.Context) {
	for key, values := range source {
		if _, has := target[key]; !has {
			target[key] = make([]string, 0)
		}
		if len(values) == 1 {
			target.Set(key, context.Expand(values[0]))
		} else {
			for _, value := range values {
				target.Add(key, context.Expand(value))
			}
		}
	}
}

//resetContext resets context for variables with Reset flag set, and removes PreviousTripStateKey
func (s *service) resetContext(context *endly.Context, request *SendRequest) {
	state := context.State()
	state.Delete(PreviousTripStateKey)
	for _, request := range request.Requests {
		if request.Repeatable != nil && len(request.Extraction) > 0 {
			request.Extraction.Reset(state)
		}
	}
}

const httpRunnerSendRequestExample = `{
  "Requests": [
    {
      "Method": "GET",
      "URL": "http://127.0.0.1:8777/event1/?k1=v1\u0026k2=v2"
    },
    {
      "Method": "GET",
      "URL": "http://127.0.0.1:8777/event1/?k10=v1\u0026k2=v2"
    },
    {
      "MatchBody": "",
      "Method": "POST",
      "URL": "http://127.0.0.1:8777/event4/",
      "Body": "Lorem Ipsum is simply dummy text of the printing and typesetting industry."
    }
  ]
}`

func (s *service) registerRoutes() {
	s.Register(&endly.ServiceActionRoute{
		Action: "send",
		RequestInfo: &endly.ActionInfo{
			Description: "send http request(s)",
			Examples: []*endly.ExampleUseCase{
				{
					UseCase: "send",
					Data:    httpRunnerSendRequestExample,
				},
			},
		},
		RequestProvider: func() interface{} {
			return &SendRequest{}
		},
		ResponseProvider: func() interface{} {
			return &SendResponse{}
		},
		Handler: func(context *endly.Context, request interface{}) (interface{}, error) {
			if req, ok := request.(*SendRequest); ok {

				return s.send(context, req)
			}
			return nil, fmt.Errorf("unsupported request type: %T", request)
		},
	})
}

//New creates a new http runner service
func New() endly.Service {
	var result = &service{
		AbstractService: endly.NewAbstractService(ServiceID),
	}
	result.AbstractService.Service = result
	result.registerRoutes()
	return result
}
<commit_msg>patched http response logging<commit_after>package http

import (
	"bytes"
	"encoding/base64"
	"fmt"
	"github.com/viant/endly"
	"github.com/viant/toolbox"
	"github.com/viant/toolbox/data"
	"io"
	"io/ioutil"
	"net/http"
	"strings"
	"time"
)

//ServiceID represents http runner service id.
const ServiceID = "http/runner"

//PreviousTripStateKey keys to store previous request details for multi trip HTTP Send request in context state
const PreviousTripStateKey = "previous"

type service struct {
	*endly.AbstractService
}

func (s *service) processResponse(context *endly.Context, sendRequest *SendRequest, sendHTTPRequest *Request, response *Response, httpResponse *http.Response, isBase64Encoded bool, extracted map[string]string) (string, error) {
	response.Header = make(map[string][]string)
	copyHeaders(httpResponse.Header, response.Header)
	readBody(httpResponse, response, isBase64Encoded)
	if sendHTTPRequest.ResponseUdf != "" {
		transformed, err := endly.TransformWithUDF(context, sendHTTPRequest.ResponseUdf, sendHTTPRequest.URL, response.Body)
		if err != nil {
			return "", err
		}
		if toolbox.IsMap(transformed) {
			response.Body, _ = toolbox.AsJSONText(transformed)
		} else {
			response.Body = toolbox.AsString(transformed)
		}
	}

	var responseBody = replaceResponseBodyIfNeeded(sendHTTPRequest, response.Body)
	return responseBody, nil
}

func (s *service) sendRequest(context *endly.Context, client *http.Client, HTTPRequest *Request, sessionCookies *Cookies, sendGroupRequest *SendRequest, sendGroupResponse *SendResponse) error {
	var err error
	var state = context.State()
	cookies := state.GetMap("cookies")
	var reader io.Reader
	var isBase64Encoded = false
	HTTPRequest = HTTPRequest.Expand(context)
	var body []byte
	var ok bool
	if len(HTTPRequest.Body) > 0 {
		body = []byte(HTTPRequest.Body)
		if HTTPRequest.RequestUdf != "" {
			transformed, err := endly.TransformWithUDF(context, HTTPRequest.RequestUdf, HTTPRequest.URL, string(body))
			if err != nil {
				return err
			}
			if body, ok = transformed.([]byte); !ok {
				body = []byte(toolbox.AsString(transformed))
			}
		}
		isBase64Encoded = strings.HasPrefix(string(body), "base64:")
		body, err = endly.FromPayload(string(body))
		if err != nil {
			return err
		}
		reader = bytes.NewReader(body)
	}

	httpRequest, err := http.NewRequest(strings.ToUpper(HTTPRequest.Method), HTTPRequest.URL, reader)
	if err != nil {
		return err
	}

	copyHeaders(HTTPRequest.Header, httpRequest.Header)
	sessionCookies.SetHeader(HTTPRequest.Header)
	HTTPRequest.Cookies.SetHeader(httpRequest.Header)

	response := &Response{}
	sendGroupResponse.Responses = append(sendGroupResponse.Responses, response)
	startEvent := s.Begin(context, HTTPRequest)
	repeatable := HTTPRequest.Repeatable.Get()

	var HTTPResponse *http.Response
	var responseBody string
	var bodyCache []byte
	var useCachedBody = repeatable.Repeat > 1 && httpRequest.ContentLength > 0
	if useCachedBody {
		bodyCache, err = ioutil.ReadAll(httpRequest.Body)
		if err != nil {
			return err
		}
	}

	handler := func() (interface{}, error) {
		if useCachedBody {
			httpRequest.Body = ioutil.NopCloser(bytes.NewReader(bodyCache))
		}

		HTTPResponse, err = client.Do(httpRequest)
		if err != nil {
			return nil, err
		}
		responseBody, err = s.processResponse(context, sendGroupRequest, HTTPRequest, response, HTTPResponse, isBase64Encoded, sendGroupResponse.Extracted)
		return responseBody, err
	}

	err = repeatable.Run(s.AbstractService, "HTTPRunner", context, handler, sendGroupResponse.Extracted)
	if err != nil {
		return err
	}
	var responseCookies Cookies = HTTPResponse.Cookies()
	response.Cookies = responseCookies.IndexByName()
	for k, cookie := range response.Cookies {
		cookies.Put(k, cookie.Value)
	}
	sessionCookies.AddCookies(responseCookies...)
	endEvent := s.End(context)(startEvent, response)

	var previous = state.GetMap(PreviousTripStateKey)
	if previous == nil {
		previous = data.NewMap()
	}
	response.Code = HTTPResponse.StatusCode
	response.TimeTakenMs = int(startEvent.Timestamp.Sub(endEvent.Timestamp) / time.Millisecond)

	if toolbox.IsCompleteJSON(responseBody) {
		response.JSONBody, err = toolbox.JSONToMap(responseBody)
		if err == nil && HTTPRequest.Repeatable != nil {
			_ = HTTPRequest.Variables.Apply(data.Map(response.JSONBody), previous)
		}
	}

	for k, v := range sendGroupResponse.Extracted {
		var expanded = previous.Expand(v)
		previous[k] = state.Expand(expanded)
	}

	err = repeatable.Variables.Apply(previous, previous)
	if err != nil {
		return err
	}
	if len(previous) > 0 {
		state.Put(PreviousTripStateKey, previous)
	}
	if HTTPRequest.MatchBody != "" {
		return nil
	}

	for _, candidate := range sendGroupRequest.Requests {
		if candidate.MatchBody != "" && strings.Contains(response.Body, candidate.MatchBody) {
			err = s.sendRequest(context, client, candidate, sessionCookies, sendGroupRequest, sendGroupResponse)
			if err != nil {
				return err
			}
		}
	}
	return nil
}

func replaceResponseBodyIfNeeded(sendHTTPRequest *Request, responseBody string) string {
	if len(sendHTTPRequest.Replace) > 0 {
		for k, v := range sendHTTPRequest.Replace {
			responseBody = strings.Replace(responseBody, k, v, len(responseBody))
		}
	}
	return responseBody
}

func (s *service) send(context *endly.Context, request *SendRequest) (*SendResponse, error) {
	client, err := toolbox.NewHttpClient(request.Options...)
	if err != nil {
		return nil, fmt.Errorf("failed to send req: %v", err)
	}
	defer s.resetContext(context, request)
	var result = &SendResponse{
		Responses: make([]*Response, 0),
		Extracted: make(map[string]string),
	}
	var sessionCookies Cookies = make([]*http.Cookie, 0)
	var state = context.State()
	if !state.Has("cookies") {
		state.Put("cookies", data.NewMap())
	}
	for _, req := range request.Requests {
		if req.MatchBody != "" {
			continue
		}
		err = s.sendRequest(context, client, req, &sessionCookies, request, result)
		if err != nil {
			return nil, err
		}
	}
	return result, nil

}
func readBody(httpResponse *http.Response, response *Response, expectBased64Encoded bool) {

	body, err := ioutil.ReadAll(httpResponse.Body)
	if err != nil {
		response.Error = fmt.Sprintf("%v", err)
		return
	}
	httpResponse.Body.Close()
	if expectBased64Encoded {
		buf := new(bytes.Buffer)
		encoder := base64.NewEncoder(base64.StdEncoding, buf)
		_, _ = encoder.Write(body)
		_ = encoder.Close()
		response.Body = "base64:" + string(buf.Bytes())

	} else {
		response.Body = string(body)
	}
}

func copyHeaders(source http.Header, target http.Header) {
	for key, values := range source {
		if _, has := target[key]; !has {
			target[key] = make([]string, 0)
		}
		if len(values) == 1 {
			target.Set(key, values[0])
		} else {

			for _, value := range values {
				target.Add(key, value)
			}
		}
	}
}

func copyExpandedHeaders(source http.Header, target http.Header, context *endly.Context) {
	for key, values := range source {
		if _, has := target[key]; !has {
			target[key] = make([]string, 0)
		}
		if len(values) == 1 {
			target.Set(key, context.Expand(values[0]))
		} else {
			for _, value := range values {
				target.Add(key, context.Expand(value))
			}
		}
	}
}

//resetContext resets context for variables with Reset flag set, and removes PreviousTripStateKey
func (s *service) resetContext(context *endly.Context, request *SendRequest) {
	state := context.State()
	state.Delete(PreviousTripStateKey)
	for _, request := range request.Requests {
		if request.Repeatable != nil && len(request.Extraction) > 0 {
			request.Extraction.Reset(state)
		}
	}
}

const httpRunnerSendRequestExample = `{
  "Requests": [
    {
      "Method": "GET",
      "URL": "http://127.0.0.1:8777/event1/?k1=v1\u0026k2=v2"
    },
    {
      "Method": "GET",
      "URL": "http://127.0.0.1:8777/event1/?k10=v1\u0026k2=v2"
    },
    {
      "MatchBody": "",
      "Method": "POST",
      "URL": "http://127.0.0.1:8777/event4/",
      "Body": "Lorem Ipsum is simply dummy text of the printing and typesetting industry."
    }
  ]
}`

func (s *service) registerRoutes() {
	s.Register(&endly.ServiceActionRoute{
		Action: "send",
		RequestInfo: &endly.ActionInfo{
			Description: "send http request(s)",
			Examples: []*endly.ExampleUseCase{
				{
					UseCase: "send",
					Data:    httpRunnerSendRequestExample,
				},
			},
		},
		RequestProvider: func() interface{} {
			return &SendRequest{}
		},
		ResponseProvider: func() interface{} {
			return &SendResponse{}
		},
		Handler: func(context *endly.Context, request interface{}) (interface{}, error) {
			if req, ok := request.(*SendRequest); ok {

				return s.send(context, req)
			}
			return nil, fmt.Errorf("unsupported request type: %T", request)
		},
	})
}

//New creates a new http runner service
func New() endly.Service {
	var result = &service{
		AbstractService: endly.NewAbstractService(ServiceID),
	}
	result.AbstractService.Service = result
	result.registerRoutes()
	return result
}
<|endoftext|>


<commit_before>package checker

import (
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*CheckJob, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	Chan      chan bool
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*result.Result, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	finishedResources := 0
	errorCount := 0
	c.CheckURL(URL)
	for finishedResources != len(c.Resources) && finishedResources < RESOURCE_LIMIT {
		select {
		case ret := <-c.Chan:
			if ret == false {
				errorCount += 1
			}
			finishedResources += 1
		}
	}
	if finishedResources == 0 || (errorCount > 0 && errorCount == finishedResources) {
		return c.Result, errors.New("Could not download host")
	}
	return c.Result, nil
}

func newCheckJob(URL string) *CheckJob {
	return &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		Chan:      make(chan bool, RESOURCE_LIMIT),
	}
}

func (c *CheckJob) CheckURL(URL string) {
	// URL already added
	if _, found := c.Resources[URL]; found {
		return
	}
	var p *PageInfo
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	go func() {
		r, err := fetchURL(URL)
		if err != nil {
			c.Result.AddError(err)
			c.Chan <- false
			return
		}
		defer r.Body.Close()
		body, err := ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
		u, _ := url.Parse(URL)
		p = &PageInfo{
			body,
			r.Header.Get("Content-Type"),
			r.StatusCode,
			r.Request.URL,
			u,
			r.Cookies(),
			utils.CropSubdomains(r.Request.URL.Host),
			&r.Header,
		}
		if err != nil {
			c.Result.AddError(err)
		}
		for _, ch := range checkers {
			ch.Check(c, p)
		}
		c.Chan <- true
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<commit_msg>[mod] do not fetch response body on non-text content-type<commit_after>package checker

import (
	"errors"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/asciimoo/privacyscore/result"
	"github.com/asciimoo/privacyscore/utils"
)

const RESOURCE_LIMIT = 64
const TIMEOUT = 5
const USER_AGENT = "Mozilla/5.0 (compatible) PrivacyScore Checker v0.1.0"
const maxResponseBodySize = 1024 * 1024 * 5

var mutex = &sync.Mutex{}

type Checker interface {
	Check(*CheckJob, *PageInfo)
}

type CheckJob struct {
	sync.RWMutex
	Result    *result.Result
	Resources map[string]*PageInfo
	Chan      chan bool
}

type PageInfo struct {
	ResponseBody   []byte
	ContentType    string
	StatusCode     int
	URL            *url.URL
	OriginalURL    *url.URL
	Cookies        []*http.Cookie
	Domain         string
	ResponseHeader *http.Header
}

var checkers []Checker = []Checker{
	&CookieChecker{},
	&HTMLChecker{},
	&HTTPSChecker{},
	&SecureHeaderChecker{},
	&CSSChecker{},
}

func Run(URL string) (*result.Result, error) {
	if !strings.HasPrefix(URL, "http://") && !strings.HasPrefix(URL, "https://") {
		URL = "http://" + URL
	}
	c := newCheckJob(URL)
	finishedResources := 0
	errorCount := 0
	c.CheckURL(URL)
	for finishedResources != len(c.Resources) && finishedResources < RESOURCE_LIMIT {
		select {
		case ret := <-c.Chan:
			if ret == false {
				errorCount += 1
			}
			finishedResources += 1
		}
	}
	if finishedResources == 0 || (errorCount > 0 && errorCount == finishedResources) {
		return c.Result, errors.New("Could not download host")
	}
	return c.Result, nil
}

func newCheckJob(URL string) *CheckJob {
	return &CheckJob{
		Result:    result.New(URL),
		Resources: make(map[string]*PageInfo),
		Chan:      make(chan bool, RESOURCE_LIMIT),
	}
}

func (c *CheckJob) CheckURL(URL string) {
	// URL already added
	if _, found := c.Resources[URL]; found {
		return
	}
	var p *PageInfo
	c.Lock()
	c.Resources[URL] = p
	c.Unlock()
	go func() {
		r, err := fetchURL(URL)
		if err != nil {
			c.Result.AddError(err)
			c.Chan <- false
			return
		}
		var body []byte
		contentType := r.Header.Get("Content-Type")
		if strings.Contains(contentType, "text") && r.StatusCode == 200 {
			body, err = ioutil.ReadAll(io.LimitReader(r.Body, maxResponseBodySize))
			if err != nil {
				c.Result.AddError(err)
			}
		} else {
			body = []byte{}
		}
		r.Body.Close()

		u, _ := url.Parse(URL)
		p = &PageInfo{
			body,
			r.Header.Get("Content-Type"),
			r.StatusCode,
			r.Request.URL,
			u,
			r.Cookies(),
			utils.CropSubdomains(r.Request.URL.Host),
			&r.Header,
		}
		for _, ch := range checkers {
			ch.Check(c, p)
		}
		c.Chan <- true
	}()
	return
}

func fetchURL(URL string) (*http.Response, error) {
	client := http.Client{
		Timeout: time.Duration(TIMEOUT * time.Second),
		Transport: &http.Transport{
			DisableKeepAlives: true,
		},
	}
	req, err := http.NewRequest("GET", URL, nil)
	if err != nil {
		return nil, err
	}
	req.Header.Set("User-Agent", USER_AGENT)
	response, err := client.Do(req)
	return response, err
}
<|endoftext|>


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

package gochimp

const (
	chimp_templates_list_endpoint string = "/templates/list.json"
)

func (a *ChimpAPI) TemplatesList(req TemplatesList) (TemplatesListResponse, error) {
	req.ApiKey = a.Key
	var response TemplatesListResponse
	err := parseChimpJson(a, chimp_templates_list_endpoint, req, &response)
	return response, err
}

type TemplateListType struct {
	User    bool `json:"user"`
	Gallery bool `json:"gallery"`
	Base    bool `json:"base"`
}

type TemplateListFilter struct {
	Category           string `json:"category"`
	FolderId           string `json:"folder_id"`
	IncludeInactive    bool   `json:"include_inactive"`
	InactiveOnly       bool   `json:"inactive_only"`
	IncludeDragAndDrop bool   `json:"include_drag_and_drop"`
}

type UserTemplate struct {
	Id           int    `json:"id"`
	Name         string `json:"name"`
	Layout       string `json:"layout"`
	Category     string `json:"category"`
	PreviewImage string `json:"preview_image"`
	DateCreated  string `json:"date_created"`
	Active       bool   `json:"active"`
	EditSource   bool   `json:"edit_source"`
	FolderId     bool   `json:"folder_id"`
}

type GalleryTemplate struct {
	Id           int    `json:"id"`
	Name         string `json:"name"`
	Layout       string `json:"layout"`
	Category     string `json:"category"`
	PreviewImage string `json:"preview_image"`
	DateCreated  string `json:"date_created"`
	Active       bool   `json:"active"`
	EditSource   bool   `json:"edit_source"`
}

type TemplatesList struct {
	ApiKey  string             `json:"apikey"`
	Types   TemplateListType   `json:"types"`
	Filters TemplateListFilter `json:"filters"`
}

type TemplatesListResponse struct {
	User    []UserTemplate    `json:"user"`
	Gallery []GalleryTemplate `json:"gallery"`
}

const (
	chimp_template_info_endpoint string = "/templates/info.json"
)

func (a *ChimpAPI) TemplatesInfo(req TemplateInfo) (TemplateInfoResponse, error) {
	req.ApiKey = a.Key
	var response TemplateInfoResponse
	err := parseChimpJson(a, chimp_template_info_endpoint, req, &response)
	return response, err
}

type TemplateInfo struct {
	ApiKey     string `json:"apikey"`
	TemplateID int    `json:"template_id"`
	Type       string `json:"type"`
}

type TemplateInfoResponse struct {
	DefaultContent interface{}
	Sections       interface{}
	Source         string
	Preview        string
}
<commit_msg>Add TemplateAdd and TemplateUpdate support<commit_after>// Copyright 2013 Matthew Baird
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//     http://www.apache.org/licenses/LICENSE-2.0
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package gochimp

const (
	chimp_templates_list_endpoint string = "/templates/list.json"
)

func (a *ChimpAPI) TemplatesList(req TemplatesList) (TemplatesListResponse, error) {
	req.ApiKey = a.Key
	var response TemplatesListResponse
	err := parseChimpJson(a, chimp_templates_list_endpoint, req, &response)
	return response, err
}

type TemplateListType struct {
	User    bool `json:"user"`
	Gallery bool `json:"gallery"`
	Base    bool `json:"base"`
}

type TemplateListFilter struct {
	Category           string `json:"category"`
	FolderId           string `json:"folder_id"`
	IncludeInactive    bool   `json:"include_inactive"`
	InactiveOnly       bool   `json:"inactive_only"`
	IncludeDragAndDrop bool   `json:"include_drag_and_drop"`
}

type UserTemplate struct {
	Id           int    `json:"id"`
	Name         string `json:"name"`
	Layout       string `json:"layout"`
	Category     string `json:"category"`
	PreviewImage string `json:"preview_image"`
	DateCreated  string `json:"date_created"`
	Active       bool   `json:"active"`
	EditSource   bool   `json:"edit_source"`
	FolderId     bool   `json:"folder_id"`
}

type GalleryTemplate struct {
	Id           int    `json:"id"`
	Name         string `json:"name"`
	Layout       string `json:"layout"`
	Category     string `json:"category"`
	PreviewImage string `json:"preview_image"`
	DateCreated  string `json:"date_created"`
	Active       bool   `json:"active"`
	EditSource   bool   `json:"edit_source"`
}

type TemplatesList struct {
	ApiKey  string             `json:"apikey"`
	Types   TemplateListType   `json:"types"`
	Filters TemplateListFilter `json:"filters"`
}

type TemplatesListResponse struct {
	User    []UserTemplate    `json:"user"`
	Gallery []GalleryTemplate `json:"gallery"`
}

const (
	chimp_template_info_endpoint   string = "/templates/info.json"
	chimp_template_add_endpoint    string = "/templates/add.json"
	chimp_template_update_endpoint string = "/templates/update.json"
)

func (a *ChimpAPI) TemplatesInfo(req TemplateInfo) (TemplateInfoResponse, error) {
	req.ApiKey = a.Key
	var response TemplateInfoResponse
	err := parseChimpJson(a, chimp_template_info_endpoint, req, &response)
	return response, err
}

type TemplateInfo struct {
	ApiKey     string `json:"apikey"`
	TemplateID int    `json:"template_id"`
	Type       string `json:"type"`
}

type TemplateInfoResponse struct {
	DefaultContent interface{}
	Sections       interface{}
	Source         string
	Preview        string
}

func (a *ChimpAPI) TemplatesAdd(req TemplatesAdd) (TemplatesAddResponse, error) {
	req.ApiKey = a.Key
	var response TemplatesAddResponse
	err := parseChimpJson(a, chimp_template_add_endpoint, req, &response)
	return response, err
}

type TemplatesAdd struct {
	ApiKey   string `json:"apikey"`
	Name     string `json:"name"`
	HTML     string `json:"html"`
	FolderID int    `json:"folder_id,omitempty"`
}

type TemplatesAddResponse struct {
	TemplateID int `json:"template_id"`
}

func (a *ChimpAPI) TemplatesUpdate(req TemplatesUpdate) (TemplatesUpdateResponse, error) {
	req.ApiKey = a.Key
	var response TemplatesUpdateResponse
	err := parseChimpJson(a, chimp_template_update_endpoint, req, &response)
	return response, err
}

type TemplatesUpdate struct {
	ApiKey     string                `json:"apikey"`
	TemplateID int                   `json:"template_id"`
	Values     TemplatesUpdateValues `json:"values"`
}

type TemplatesUpdateValues struct {
	Name     string `json:"name"`
	HTML     string `json:"html"`
	FolderID int    `json:"folder_id,omitempty"`
}

type TemplatesUpdateResponse struct {
	Complete bool `json:"complete"`
}
<|endoftext|>


<commit_before>/*
 *
 * Copyright 2019 gRPC 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
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

// Binary grpclb_fallback is an interop test client for grpclb fallback.
package main

import (
	"context"
	"flag"
	"log"
	"net"
	"os"
	"os/exec"
	"syscall"
	"time"

	"golang.org/x/sys/unix"
	"google.golang.org/grpc"
	_ "google.golang.org/grpc/balancer/grpclb"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/credentials/google"
	_ "google.golang.org/grpc/xds/googledirectpath"

	testgrpc "google.golang.org/grpc/interop/grpc_testing"
	testpb "google.golang.org/grpc/interop/grpc_testing"
)

var (
	customCredentialsType         = flag.String("custom_credentials_type", "", "Client creds to use")
	serverURI                     = flag.String("server_uri", "dns:///staging-grpc-directpath-fallback-test.googleapis.com:443", "The server host name")
	unrouteLBAndBackendAddrsCmd   = flag.String("unroute_lb_and_backend_addrs_cmd", "", "Command to make LB and backend address unroutable")
	blackholeLBAndBackendAddrsCmd = flag.String("blackhole_lb_and_backend_addrs_cmd", "", "Command to make LB and backend addresses blackholed")
	testCase                      = flag.String("test_case", "",
		`Configure different test cases. Valid options are:
        fast_fallback_before_startup : LB/backend connections fail fast before RPC's have been made;
        fast_fallback_after_startup : LB/backend connections fail fast after RPC's have been made;
        slow_fallback_before_startup : LB/backend connections black hole before RPC's have been made;
        slow_fallback_after_startup : LB/backend connections black hole after RPC's have been made;`)
	infoLog  = log.New(os.Stderr, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
	errorLog = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
)

func doRPCAndGetPath(client testgrpc.TestServiceClient, timeout time.Duration) testpb.GrpclbRouteType {
	infoLog.Printf("doRPCAndGetPath timeout:%v\n", timeout)
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	req := &testpb.SimpleRequest{
		FillGrpclbRouteType: true,
	}
	reply, err := client.UnaryCall(ctx, req)
	if err != nil {
		infoLog.Printf("doRPCAndGetPath error:%v\n", err)
		return testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_UNKNOWN
	}
	g := reply.GetGrpclbRouteType()
	infoLog.Printf("doRPCAndGetPath got grpclb route type: %v\n", g)
	if g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_FALLBACK && g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
		errorLog.Fatalf("Expected grpclb route type to be either backend or fallback; got: %d", g)
	}
	return g
}

func dialTCPUserTimeout(ctx context.Context, addr string) (net.Conn, error) {
	control := func(network, address string, c syscall.RawConn) error {
		var syscallErr error
		controlErr := c.Control(func(fd uintptr) {
			syscallErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, unix.TCP_USER_TIMEOUT, 20000)
		})
		if syscallErr != nil {
			errorLog.Fatalf("syscall error setting sockopt TCP_USER_TIMEOUT: %v", syscallErr)
		}
		if controlErr != nil {
			errorLog.Fatalf("control error setting sockopt TCP_USER_TIMEOUT: %v", syscallErr)
		}
		return nil
	}
	d := &net.Dialer{
		Control: control,
	}
	return d.DialContext(ctx, "tcp", addr)
}

func createTestConn() *grpc.ClientConn {
	opts := []grpc.DialOption{
		grpc.WithContextDialer(dialTCPUserTimeout),
		grpc.WithBlock(),
	}
	switch *customCredentialsType {
	case "tls":
		creds := credentials.NewClientTLSFromCert(nil, "")
		opts = append(opts, grpc.WithTransportCredentials(creds))
	case "alts":
		creds := alts.NewClientCreds(alts.DefaultClientOptions())
		opts = append(opts, grpc.WithTransportCredentials(creds))
	case "google_default_credentials":
		opts = append(opts, grpc.WithCredentialsBundle(google.NewDefaultCredentials()))
	case "compute_engine_channel_creds":
		opts = append(opts, grpc.WithCredentialsBundle(google.NewComputeEngineCredentials()))
	default:
		errorLog.Fatalf("Invalid --custom_credentials_type:%v", *customCredentialsType)
	}
	conn, err := grpc.Dial(*serverURI, opts...)
	if err != nil {
		errorLog.Fatalf("Fail to dial: %v", err)
	}
	return conn
}

func runCmd(command string) {
	infoLog.Printf("Running cmd:|%v|\n", command)
	if err := exec.Command("bash", "-c", command).Run(); err != nil {
		errorLog.Fatalf("error running cmd:|%v| : %v", command, err)
	}
}

func waitForFallbackAndDoRPCs(client testgrpc.TestServiceClient, fallbackDeadline time.Time) {
	fallbackRetryCount := 0
	fellBack := false
	for time.Now().Before(fallbackDeadline) {
		g := doRPCAndGetPath(client, 1*time.Second)
		if g == testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_FALLBACK {
			infoLog.Println("Made one successul RPC to a fallback. Now expect the same for the rest.")
			fellBack = true
			break
		} else if g == testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
			errorLog.Fatalf("Got RPC type backend. This suggests an error in test implementation")
		} else {
			infoLog.Println("Retryable RPC failure on iteration:", fallbackRetryCount)
		}
		fallbackRetryCount++
	}
	if !fellBack {
		infoLog.Fatalf("Didn't fall back before deadline: %v\n", fallbackDeadline)
	}
	for i := 0; i < 30; i++ {
		if g := doRPCAndGetPath(client, 20*time.Second); g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_FALLBACK {
			errorLog.Fatalf("Expected RPC to take grpclb route type FALLBACK. Got: %v", g)
		}
		time.Sleep(time.Second)
	}
}

func doFastFallbackBeforeStartup() {
	runCmd(*unrouteLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(5 * time.Second)
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func doSlowFallbackBeforeStartup() {
	runCmd(*blackholeLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(20 * time.Second)
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func doFastFallbackAfterStartup() {
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	if g := doRPCAndGetPath(client, 20*time.Second); g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
		errorLog.Fatalf("Expected RPC to take grpclb route type BACKEND. Got: %v", g)
	}
	runCmd(*unrouteLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(40 * time.Second)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func doSlowFallbackAfterStartup() {
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	if g := doRPCAndGetPath(client, 20*time.Second); g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
		errorLog.Fatalf("Expected RPC to take grpclb route type BACKEND. Got: %v", g)
	}
	runCmd(*blackholeLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(40 * time.Second)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func main() {
	flag.Parse()
	if len(*unrouteLBAndBackendAddrsCmd) == 0 {
		errorLog.Fatalf("--unroute_lb_and_backend_addrs_cmd unset")
	}
	if len(*blackholeLBAndBackendAddrsCmd) == 0 {
		errorLog.Fatalf("--blackhole_lb_and_backend_addrs_cmd unset")
	}
	switch *testCase {
	case "fast_fallback_before_startup":
		doFastFallbackBeforeStartup()
		log.Printf("FastFallbackBeforeStartup done!\n")
	case "fast_fallback_after_startup":
		doFastFallbackAfterStartup()
		log.Printf("FastFallbackAfterStartup done!\n")
	case "slow_fallback_before_startup":
		doSlowFallbackBeforeStartup()
		log.Printf("SlowFallbackBeforeStartup done!\n")
	case "slow_fallback_after_startup":
		doSlowFallbackAfterStartup()
		log.Printf("SlowFallbackAfterStartup done!\n")
	default:
		errorLog.Fatalf("Unsupported test case: %v", *testCase)
	}
}
<commit_msg>xds: Remove WithBlock option from the fallback test client (#5066)<commit_after>/*
 *
 * Copyright 2019 gRPC 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
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 *
 */

// Binary grpclb_fallback is an interop test client for grpclb fallback.
package main

import (
	"context"
	"flag"
	"log"
	"net"
	"os"
	"os/exec"
	"syscall"
	"time"

	"golang.org/x/sys/unix"
	"google.golang.org/grpc"
	_ "google.golang.org/grpc/balancer/grpclb"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/alts"
	"google.golang.org/grpc/credentials/google"
	_ "google.golang.org/grpc/xds/googledirectpath"

	testgrpc "google.golang.org/grpc/interop/grpc_testing"
	testpb "google.golang.org/grpc/interop/grpc_testing"
)

var (
	customCredentialsType         = flag.String("custom_credentials_type", "", "Client creds to use")
	serverURI                     = flag.String("server_uri", "dns:///staging-grpc-directpath-fallback-test.googleapis.com:443", "The server host name")
	unrouteLBAndBackendAddrsCmd   = flag.String("unroute_lb_and_backend_addrs_cmd", "", "Command to make LB and backend address unroutable")
	blackholeLBAndBackendAddrsCmd = flag.String("blackhole_lb_and_backend_addrs_cmd", "", "Command to make LB and backend addresses blackholed")
	testCase                      = flag.String("test_case", "",
		`Configure different test cases. Valid options are:
        fast_fallback_before_startup : LB/backend connections fail fast before RPC's have been made;
        fast_fallback_after_startup : LB/backend connections fail fast after RPC's have been made;
        slow_fallback_before_startup : LB/backend connections black hole before RPC's have been made;
        slow_fallback_after_startup : LB/backend connections black hole after RPC's have been made;`)
	infoLog  = log.New(os.Stderr, "INFO: ", log.Ldate|log.Ltime|log.Lshortfile)
	errorLog = log.New(os.Stderr, "ERROR: ", log.Ldate|log.Ltime|log.Lshortfile)
)

func doRPCAndGetPath(client testgrpc.TestServiceClient, timeout time.Duration) testpb.GrpclbRouteType {
	infoLog.Printf("doRPCAndGetPath timeout:%v\n", timeout)
	ctx, cancel := context.WithTimeout(context.Background(), timeout)
	defer cancel()
	req := &testpb.SimpleRequest{
		FillGrpclbRouteType: true,
	}
	reply, err := client.UnaryCall(ctx, req)
	if err != nil {
		infoLog.Printf("doRPCAndGetPath error:%v\n", err)
		return testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_UNKNOWN
	}
	g := reply.GetGrpclbRouteType()
	infoLog.Printf("doRPCAndGetPath got grpclb route type: %v\n", g)
	if g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_FALLBACK && g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
		errorLog.Fatalf("Expected grpclb route type to be either backend or fallback; got: %d", g)
	}
	return g
}

func dialTCPUserTimeout(ctx context.Context, addr string) (net.Conn, error) {
	control := func(network, address string, c syscall.RawConn) error {
		var syscallErr error
		controlErr := c.Control(func(fd uintptr) {
			syscallErr = syscall.SetsockoptInt(int(fd), syscall.IPPROTO_TCP, unix.TCP_USER_TIMEOUT, 20000)
		})
		if syscallErr != nil {
			errorLog.Fatalf("syscall error setting sockopt TCP_USER_TIMEOUT: %v", syscallErr)
		}
		if controlErr != nil {
			errorLog.Fatalf("control error setting sockopt TCP_USER_TIMEOUT: %v", syscallErr)
		}
		return nil
	}
	d := &net.Dialer{
		Control: control,
	}
	return d.DialContext(ctx, "tcp", addr)
}

func createTestConn() *grpc.ClientConn {
	opts := []grpc.DialOption{
		grpc.WithContextDialer(dialTCPUserTimeout),
	}
	switch *customCredentialsType {
	case "tls":
		creds := credentials.NewClientTLSFromCert(nil, "")
		opts = append(opts, grpc.WithTransportCredentials(creds))
	case "alts":
		creds := alts.NewClientCreds(alts.DefaultClientOptions())
		opts = append(opts, grpc.WithTransportCredentials(creds))
	case "google_default_credentials":
		opts = append(opts, grpc.WithCredentialsBundle(google.NewDefaultCredentials()))
	case "compute_engine_channel_creds":
		opts = append(opts, grpc.WithCredentialsBundle(google.NewComputeEngineCredentials()))
	default:
		errorLog.Fatalf("Invalid --custom_credentials_type:%v", *customCredentialsType)
	}
	conn, err := grpc.Dial(*serverURI, opts...)
	if err != nil {
		errorLog.Fatalf("Fail to dial: %v", err)
	}
	return conn
}

func runCmd(command string) {
	infoLog.Printf("Running cmd:|%v|\n", command)
	if err := exec.Command("bash", "-c", command).Run(); err != nil {
		errorLog.Fatalf("error running cmd:|%v| : %v", command, err)
	}
}

func waitForFallbackAndDoRPCs(client testgrpc.TestServiceClient, fallbackDeadline time.Time) {
	fallbackRetryCount := 0
	fellBack := false
	for time.Now().Before(fallbackDeadline) {
		g := doRPCAndGetPath(client, 1*time.Second)
		if g == testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_FALLBACK {
			infoLog.Println("Made one successul RPC to a fallback. Now expect the same for the rest.")
			fellBack = true
			break
		} else if g == testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
			errorLog.Fatalf("Got RPC type backend. This suggests an error in test implementation")
		} else {
			infoLog.Println("Retryable RPC failure on iteration:", fallbackRetryCount)
		}
		fallbackRetryCount++
	}
	if !fellBack {
		infoLog.Fatalf("Didn't fall back before deadline: %v\n", fallbackDeadline)
	}
	for i := 0; i < 30; i++ {
		if g := doRPCAndGetPath(client, 20*time.Second); g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_FALLBACK {
			errorLog.Fatalf("Expected RPC to take grpclb route type FALLBACK. Got: %v", g)
		}
		time.Sleep(time.Second)
	}
}

func doFastFallbackBeforeStartup() {
	runCmd(*unrouteLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(5 * time.Second)
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func doSlowFallbackBeforeStartup() {
	runCmd(*blackholeLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(20 * time.Second)
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func doFastFallbackAfterStartup() {
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	if g := doRPCAndGetPath(client, 20*time.Second); g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
		errorLog.Fatalf("Expected RPC to take grpclb route type BACKEND. Got: %v", g)
	}
	runCmd(*unrouteLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(40 * time.Second)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func doSlowFallbackAfterStartup() {
	conn := createTestConn()
	defer conn.Close()
	client := testgrpc.NewTestServiceClient(conn)
	if g := doRPCAndGetPath(client, 20*time.Second); g != testpb.GrpclbRouteType_GRPCLB_ROUTE_TYPE_BACKEND {
		errorLog.Fatalf("Expected RPC to take grpclb route type BACKEND. Got: %v", g)
	}
	runCmd(*blackholeLBAndBackendAddrsCmd)
	fallbackDeadline := time.Now().Add(40 * time.Second)
	waitForFallbackAndDoRPCs(client, fallbackDeadline)
}

func main() {
	flag.Parse()
	if len(*unrouteLBAndBackendAddrsCmd) == 0 {
		errorLog.Fatalf("--unroute_lb_and_backend_addrs_cmd unset")
	}
	if len(*blackholeLBAndBackendAddrsCmd) == 0 {
		errorLog.Fatalf("--blackhole_lb_and_backend_addrs_cmd unset")
	}
	switch *testCase {
	case "fast_fallback_before_startup":
		doFastFallbackBeforeStartup()
		log.Printf("FastFallbackBeforeStartup done!\n")
	case "fast_fallback_after_startup":
		doFastFallbackAfterStartup()
		log.Printf("FastFallbackAfterStartup done!\n")
	case "slow_fallback_before_startup":
		doSlowFallbackBeforeStartup()
		log.Printf("SlowFallbackBeforeStartup done!\n")
	case "slow_fallback_after_startup":
		doSlowFallbackAfterStartup()
		log.Printf("SlowFallbackAfterStartup done!\n")
	default:
		errorLog.Fatalf("Unsupported test case: %v", *testCase)
	}
}
<|endoftext|>


<commit_before>package signage

import (
	"net/http"
	"strings"
	"time"

	"github.com/DeedleFake/signage/errors"

	"golang.org/x/net/html"
)

const (
	base = "https://www.whitehouse.gov"

	signed  = base + "/briefing-room/signed-legislation"
	vetoed  = base + "/briefing-room/vetoed-legislation"
	pending = base + "/briefing-room/pending-legislation"
)

// GetSigned fetches a list of signed bills.
func GetSigned() ([]Bill, error) {
	return scrape(signed)
}

// GetPending fetches a list of pending bills.
func GetPending() ([]Bill, error) {
	return scrape(pending)
}

// GetVetoed fetches a list of vetoed bills.
func GetVetoed() ([]Bill, error) {
	return scrape(vetoed)
}

// scrape pulls a list of bills from a URL.
//
// TODO: Handle scraping multiple pages.
func scrape(url string) ([]Bill, error) {
	const (
		DateFormat = "2006-01-02T15:04:05-07:00"
	)

	rsp, err := http.Get(url)
	if err != nil {
		return nil, errors.Err(err)
	}
	defer rsp.Body.Close()

	root, err := html.Parse(rsp.Body)
	if err != nil {
		return nil, errors.Err(err)
	}

	content := findNode(root, func(n *html.Node) bool {
		return (n.Type == html.ElementNode) && (n.Data == "div") && (getAttr(n.Attr, "class") == "view-content")
	})

	var entries []Bill
	for cur := content.FirstChild; cur != nil; cur = cur.NextSibling {
		if (cur.Type != html.ElementNode) || (cur.Data != "div") {
			continue
		}

		found := findNode(cur, func(n *html.Node) bool {
			return (n.Type == html.ElementNode) && (n.Data == "span") && (getAttr(n.Attr, "datatype") == "xsd:dateTime")
		})
		date, err := time.Parse(DateFormat, getAttr(found.Attr, "content"))
		if err != nil {
			return nil, errors.Err(err)
		}

		found = findNode(cur, func(n *html.Node) bool {
			return (n.Type == html.ElementNode) && (n.Data == "h3") && (getAttr(n.Attr, "class") == "field-content")
		})
		title := strings.TrimSpace(found.FirstChild.FirstChild.Data)
		url := getAttr(found.FirstChild.Attr, "href")

		entries = append(entries, Bill{
			Date:  date,
			Title: title,
			URL:   base + url,
		})
	}

	return entries, nil
}

// findNode recursively searches an HTML node tree until it finds one
// on which match returns true, at which point it returns that node.
// If no nodes match, it returns nil.
func findNode(root *html.Node, match func(*html.Node) bool) *html.Node {
	if (root == nil) || match(root) {
		return root
	}

	if found := findNode(root.FirstChild, match); found != nil {
		return found
	}

	return findNode(root.NextSibling, match)
}

func getAttr(attrs []html.Attribute, key string) (val string) {
	for _, attr := range attrs {
		if attr.Key == key {
			return attr.Val
		}
	}

	return
}

// Bill contains information about a specific entry on the White House
// site.
type Bill struct {
	Date  time.Time
	Title string
	URL   string
}
<commit_msg>signage: Don't crash if there are no bills in a list.<commit_after>package signage

import (
	"net/http"
	"strings"
	"time"

	"github.com/DeedleFake/signage/errors"

	"golang.org/x/net/html"
)

const (
	base = "https://www.whitehouse.gov"

	signed  = base + "/briefing-room/signed-legislation"
	vetoed  = base + "/briefing-room/vetoed-legislation"
	pending = base + "/briefing-room/pending-legislation"
)

// GetSigned fetches a list of signed bills.
func GetSigned() ([]Bill, error) {
	return scrape(signed)
}

// GetPending fetches a list of pending bills.
func GetPending() ([]Bill, error) {
	return scrape(pending)
}

// GetVetoed fetches a list of vetoed bills.
func GetVetoed() ([]Bill, error) {
	return scrape(vetoed)
}

// scrape pulls a list of bills from a URL.
//
// TODO: Handle scraping multiple pages.
func scrape(url string) ([]Bill, error) {
	const (
		DateFormat = "2006-01-02T15:04:05-07:00"
	)

	rsp, err := http.Get(url)
	if err != nil {
		return nil, errors.Err(err)
	}
	defer rsp.Body.Close()

	root, err := html.Parse(rsp.Body)
	if err != nil {
		return nil, errors.Err(err)
	}

	content := findNode(root, func(n *html.Node) bool {
		return (n.Type == html.ElementNode) && (n.Data == "div") && (getAttr(n.Attr, "class") == "view-content")
	})
	if content == nil {
		return nil, nil
	}

	var entries []Bill
	for cur := content.FirstChild; cur != nil; cur = cur.NextSibling {
		if (cur.Type != html.ElementNode) || (cur.Data != "div") {
			continue
		}

		found := findNode(cur, func(n *html.Node) bool {
			return (n.Type == html.ElementNode) && (n.Data == "span") && (getAttr(n.Attr, "datatype") == "xsd:dateTime")
		})
		date, err := time.Parse(DateFormat, getAttr(found.Attr, "content"))
		if err != nil {
			return nil, errors.Err(err)
		}

		found = findNode(cur, func(n *html.Node) bool {
			return (n.Type == html.ElementNode) && (n.Data == "h3") && (getAttr(n.Attr, "class") == "field-content")
		})
		title := strings.TrimSpace(found.FirstChild.FirstChild.Data)
		url := getAttr(found.FirstChild.Attr, "href")

		entries = append(entries, Bill{
			Date:  date,
			Title: title,
			URL:   base + url,
		})
	}

	return entries, nil
}

// findNode recursively searches an HTML node tree until it finds one
// on which match returns true, at which point it returns that node.
// If no nodes match, it returns nil.
func findNode(root *html.Node, match func(*html.Node) bool) *html.Node {
	if (root == nil) || match(root) {
		return root
	}

	if found := findNode(root.FirstChild, match); found != nil {
		return found
	}

	return findNode(root.NextSibling, match)
}

func getAttr(attrs []html.Attribute, key string) (val string) {
	for _, attr := range attrs {
		if attr.Key == key {
			return attr.Val
		}
	}

	return
}

// Bill contains information about a specific entry on the White House
// site.
type Bill struct {
	Date  time.Time
	Title string
	URL   string
}
<|endoftext|>