<commit_before>package blockchain

import (
	"bytes"
	crand "crypto/rand"
	"testing"

	log "github.com/Sirupsen/logrus"
)

// TestMain sets logging levels for tests.
func TestMain(t *testing.T) {
	log.SetLevel(log.DebugLevel)
}

func TestValidTransactionNotInBlock(t *testing.T) {
	tr := newTransactionValue(5, Address{})

	inputTransactions := make([]*Transaction, 0)
	outputTransactions := make([]*Transaction, 0)

	inputBlock := newInputBlock(inputTransactions)
	outputBlock := newOutputBlock(outputTransactions, inputBlock)

	bc := BlockChain{
		Blocks: []*Block{inputBlock, outputBlock},
		Head:   newHash(),
	}

	if bc.ValidTransaction(tr) {
		t.Fail()
	}
}

func TestValidTransactionInputsFail(t *testing.T) {
	// 2 + 2 = 5 ?
	trA := newTransactionValue(2, newWallet().Public())
	trB := newTransactionValue(2, newWallet().Public())
	trC := newTransactionValue(5, newWallet().Public())

	inputTransactions := []*Transaction{trA, trB}
	outputTransactions := []*Transaction{trC}

	inputBlock := newInputBlock(inputTransactions)
	outputBlock := newOutputBlock(outputTransactions, inputBlock)

	bc := BlockChain{
		Blocks: []*Block{inputBlock, outputBlock},
		Head:   newHash(),
	}

	if bc.ValidTransaction(trC) {
		t.Fail()
	}
}

func TestValidTransactionSignatureFail(t *testing.T) {
	tr := newTransactionValue(2, newWallet().Public())

	// Fake sender puts in a silly signature.
	fakeSender := newWallet()
	fakeTransaction := newTransaction()
	digest := HashSum(fakeTransaction.TxBody)
	tr.Sig, _ = fakeSender.Sign(digest, crand.Reader)

	transactions := []*Transaction{tr}
	b := newInputBlock(transactions)

	bc := BlockChain{
		Blocks: []*Block{b},
		Head:   newHash(),
	}

	if bc.ValidTransaction(tr) {
		t.Fail()
	}

}

// TestValidBlock tests the three cases in which a block can fail to be valid.
func TestValidBlock(t *testing.T) {
	// b conains invalid transaction.
	// Block number for b is not one greater than the last block.
	// The hash of the last block is incorrect.
}

func TestEncodeDecodeBlockChain(t *testing.T) {
	b1 := newBlockChain()

	buf := bytes.NewBuffer(make([]byte, 0, b1.Len()))

	b1.Encode(buf)
	b2 := DecodeBlockChain(buf)

	if HashSum(b1) != HashSum(b2) {
		t.Fail()
	}
}
<commit_msg>remove empty space<commit_after>package blockchain

import (
	"bytes"
	crand "crypto/rand"
	"testing"

	log "github.com/Sirupsen/logrus"
)

// TestMain sets logging levels for tests.
func TestMain(t *testing.T) {
	log.SetLevel(log.DebugLevel)
}

func TestValidTransactionNotInBlock(t *testing.T) {
	tr := newTransactionValue(5, Address{})

	inputTransactions := make([]*Transaction, 0)
	outputTransactions := make([]*Transaction, 0)

	inputBlock := newInputBlock(inputTransactions)
	outputBlock := newOutputBlock(outputTransactions, inputBlock)

	bc := BlockChain{
		Blocks: []*Block{inputBlock, outputBlock},
		Head:   newHash(),
	}

	if bc.ValidTransaction(tr) {
		t.Fail()
	}
}

func TestValidTransactionInputsFail(t *testing.T) {
	// 2 + 2 = 5 ?
	trA := newTransactionValue(2, newWallet().Public())
	trB := newTransactionValue(2, newWallet().Public())
	trC := newTransactionValue(5, newWallet().Public())

	inputTransactions := []*Transaction{trA, trB}
	outputTransactions := []*Transaction{trC}

	inputBlock := newInputBlock(inputTransactions)
	outputBlock := newOutputBlock(outputTransactions, inputBlock)

	bc := BlockChain{
		Blocks: []*Block{inputBlock, outputBlock},
		Head:   newHash(),
	}

	if bc.ValidTransaction(trC) {
		t.Fail()
	}
}

func TestValidTransactionSignatureFail(t *testing.T) {
	tr := newTransactionValue(2, newWallet().Public())

	// Fake sender puts in a silly signature.
	fakeSender := newWallet()
	fakeTransaction := newTransaction()
	digest := HashSum(fakeTransaction.TxBody)
	tr.Sig, _ = fakeSender.Sign(digest, crand.Reader)

	transactions := []*Transaction{tr}
	b := newInputBlock(transactions)

	bc := BlockChain{
		Blocks: []*Block{b},
		Head:   newHash(),
	}

	if bc.ValidTransaction(tr) {
		t.Fail()
	}
}

// TestValidBlock tests the three cases in which a block can fail to be valid.
func TestValidBlock(t *testing.T) {
	// b conains invalid transaction.
	// Block number for b is not one greater than the last block.
	// The hash of the last block is incorrect.
}

func TestEncodeDecodeBlockChain(t *testing.T) {
	b1 := newBlockChain()

	buf := bytes.NewBuffer(make([]byte, 0, b1.Len()))

	b1.Encode(buf)
	b2 := DecodeBlockChain(buf)

	if HashSum(b1) != HashSum(b2) {
		t.Fail()
	}
}
<|endoftext|>


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

package merkle

import (
	"bytes"
	"errors"
	"fmt"
	"math/bits"
)

// RootMismatchError occurs when an inclusion proof fails.
type RootMismatchError struct {
	ExpectedRoot   []byte
	CalculatedRoot []byte
}

func (e RootMismatchError) Error() string {
	return fmt.Sprintf("calculated root:\n%v\n does not match expected root:\n%v", e.CalculatedRoot, e.ExpectedRoot)
}

func verifyMatch(calculated, expected []byte) error {
	if !bytes.Equal(calculated, expected) {
		return RootMismatchError{ExpectedRoot: expected, CalculatedRoot: calculated}
	}
	return nil
}

// LogHasher provides the hash functions needed to compute dense merkle trees.
type LogHasher interface {
	// EmptyRoot supports returning a special case for the root of an empty tree.
	EmptyRoot() []byte
	// HashLeaf computes the hash of a leaf that exists.
	HashLeaf(leaf []byte) []byte
	// HashChildren computes interior nodes.
	HashChildren(l, r []byte) []byte
	// Size returns the number of bytes the Hash* functions will return.
	Size() int
}

// LogVerifier verifies inclusion and consistency proofs for append only logs.
type LogVerifier struct {
	hasher LogHasher
}

// NewLogVerifier returns a new LogVerifier for a tree.
func NewLogVerifier(hasher LogHasher) LogVerifier {
	return LogVerifier{hasher}
}

// VerifyInclusion verifies the correctness of the inclusion proof for the leaf
// with the specified hash and index, relatively to the tree of the given size
// and root hash. Requires 0 <= index < size.
func (v LogVerifier) VerifyInclusion(index, size uint64, leafHash []byte, proof [][]byte, root []byte) error {
	calcRoot, err := v.RootFromInclusionProof(index, size, leafHash, proof)
	if err != nil {
		return err
	}
	return verifyMatch(calcRoot, root)
}

// RootFromInclusionProof calculates the expected root hash for a tree of the
// given size, provided a leaf index and hash with the corresponding inclusion
// proof. Requires 0 <= index < size.
func (v LogVerifier) RootFromInclusionProof(index, size uint64, leafHash []byte, proof [][]byte) ([]byte, error) {
	if index >= size {
		return nil, fmt.Errorf("index is beyond size: %d >= %d", index, size)
	}
	if got, want := len(leafHash), v.hasher.Size(); got != want {
		return nil, fmt.Errorf("leafHash has unexpected size %d, want %d", got, want)
	}

	inner, border := decompInclProof(index, size)
	if got, want := len(proof), inner+border; got != want {
		return nil, fmt.Errorf("wrong proof size %d, want %d", got, want)
	}

	res := chainInner(v.hasher, leafHash, proof[:inner], index)
	res = chainBorderRight(v.hasher, res, proof[inner:])
	return res, nil
}

// VerifyConsistency checks that the passed-in consistency proof is valid
// between the passed in tree sizes, with respect to the corresponding root
// hashes. Requires 0 <= size1 <= size2.
func (v LogVerifier) VerifyConsistency(size1, size2 uint64, root1, root2 []byte, proof [][]byte) error {
	switch {
	case size2 < size1:
		return fmt.Errorf("size2 (%d) < size1 (%d)", size1, size2)
	case size1 == size2:
		if !bytes.Equal(root1, root2) {
			return RootMismatchError{
				CalculatedRoot: root1,
				ExpectedRoot:   root2,
			}
		} else if len(proof) > 0 {
			return errors.New("root1 and root2 match, but proof is non-empty")
		}
		return nil // Proof OK.
	case size1 == 0:
		// Any size greater than 0 is consistent with size 0.
		if len(proof) > 0 {
			return fmt.Errorf("expected empty proof, but got %d components", len(proof))
		}
		return nil // Proof OK.
	case len(proof) == 0:
		return errors.New("empty proof")
	}

	inner, border := decompInclProof(size1-1, size2)
	shift := bits.TrailingZeros64(size1)
	inner -= shift // Note: shift < inner if size1 < size2.

	// The proof includes the root hash for the sub-tree of size 2^shift.
	seed, start := proof[0], 1
	if size1 == 1<<uint(shift) { // Unless size1 is that very 2^shift.
		seed, start = root1, 0
	}
	if got, want := len(proof), start+inner+border; got != want {
		return fmt.Errorf("wrong proof size %d, want %d", got, want)
	}
	proof = proof[start:]
	// Now len(proof) == inner+border, and proof is effectively a suffix of
	// inclusion proof for entry |size1-1| in a tree of size |size2|.

	// Verify the first root.
	mask := (size1 - 1) >> uint(shift) // Start chaining from level |shift|.
	hash1 := chainInnerRight(v.hasher, seed, proof[:inner], mask)
	hash1 = chainBorderRight(v.hasher, hash1, proof[inner:])
	if err := verifyMatch(hash1, root1); err != nil {
		return err
	}

	// Verify the second root.
	hash2 := chainInner(v.hasher, seed, proof[:inner], mask)
	hash2 = chainBorderRight(v.hasher, hash2, proof[inner:])
	return verifyMatch(hash2, root2)
}

// decompInclProof breaks down inclusion proof for a leaf at the specified
// |index| in a tree of the specified |size| into 2 components. The splitting
// point between them is where paths to leaves |index| and |size-1| diverge.
// Returns lengths of the bottom and upper proof parts correspondingly. The sum
// of the two determines the correct length of the inclusion proof.
func decompInclProof(index, size uint64) (int, int) {
	inner := innerProofSize(index, size)
	border := bits.OnesCount64(index >> uint(inner))
	return inner, border
}

func innerProofSize(index, size uint64) int {
	return bits.Len64(index ^ (size - 1))
}

// chainInner computes a subtree hash for a node on or below the tree's right
// border. Assumes |proof| hashes are ordered from lower levels to upper, and
// |seed| is the initial subtree/leaf hash on the path located at the specified
// |index| on its level.
func chainInner(hasher LogHasher, seed []byte, proof [][]byte, index uint64) []byte {
	for i, h := range proof {
		if (index>>uint(i))&1 == 0 {
			seed = hasher.HashChildren(seed, h)
		} else {
			seed = hasher.HashChildren(h, seed)
		}
	}
	return seed
}

// chainInnerRight computes a subtree hash like chainInner, but only takes
// hashes to the left from the path into consideration, which effectively means
// the result is a hash of the corresponding earlier version of this subtree.
func chainInnerRight(hasher LogHasher, seed []byte, proof [][]byte, index uint64) []byte {
	for i, h := range proof {
		if (index>>uint(i))&1 == 1 {
			seed = hasher.HashChildren(h, seed)
		}
	}
	return seed
}

// chainBorderRight chains proof hashes along tree borders. This differs from
// inner chaining because |proof| contains only left-side subtree hashes.
func chainBorderRight(hasher LogHasher, seed []byte, proof [][]byte) []byte {
	for _, h := range proof {
		seed = hasher.HashChildren(h, seed)
	}
	return seed
}
<commit_msg>Replace hash comparison with verifyMatch<commit_after>// Copyright 2017 Google LLC. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package merkle

import (
	"bytes"
	"errors"
	"fmt"
	"math/bits"
)

// RootMismatchError occurs when an inclusion proof fails.
type RootMismatchError struct {
	ExpectedRoot   []byte
	CalculatedRoot []byte
}

func (e RootMismatchError) Error() string {
	return fmt.Sprintf("calculated root:\n%v\n does not match expected root:\n%v", e.CalculatedRoot, e.ExpectedRoot)
}

func verifyMatch(calculated, expected []byte) error {
	if !bytes.Equal(calculated, expected) {
		return RootMismatchError{ExpectedRoot: expected, CalculatedRoot: calculated}
	}
	return nil
}

// LogHasher provides the hash functions needed to compute dense merkle trees.
type LogHasher interface {
	// EmptyRoot supports returning a special case for the root of an empty tree.
	EmptyRoot() []byte
	// HashLeaf computes the hash of a leaf that exists.
	HashLeaf(leaf []byte) []byte
	// HashChildren computes interior nodes.
	HashChildren(l, r []byte) []byte
	// Size returns the number of bytes the Hash* functions will return.
	Size() int
}

// LogVerifier verifies inclusion and consistency proofs for append only logs.
type LogVerifier struct {
	hasher LogHasher
}

// NewLogVerifier returns a new LogVerifier for a tree.
func NewLogVerifier(hasher LogHasher) LogVerifier {
	return LogVerifier{hasher}
}

// VerifyInclusion verifies the correctness of the inclusion proof for the leaf
// with the specified hash and index, relatively to the tree of the given size
// and root hash. Requires 0 <= index < size.
func (v LogVerifier) VerifyInclusion(index, size uint64, leafHash []byte, proof [][]byte, root []byte) error {
	calcRoot, err := v.RootFromInclusionProof(index, size, leafHash, proof)
	if err != nil {
		return err
	}
	return verifyMatch(calcRoot, root)
}

// RootFromInclusionProof calculates the expected root hash for a tree of the
// given size, provided a leaf index and hash with the corresponding inclusion
// proof. Requires 0 <= index < size.
func (v LogVerifier) RootFromInclusionProof(index, size uint64, leafHash []byte, proof [][]byte) ([]byte, error) {
	if index >= size {
		return nil, fmt.Errorf("index is beyond size: %d >= %d", index, size)
	}
	if got, want := len(leafHash), v.hasher.Size(); got != want {
		return nil, fmt.Errorf("leafHash has unexpected size %d, want %d", got, want)
	}

	inner, border := decompInclProof(index, size)
	if got, want := len(proof), inner+border; got != want {
		return nil, fmt.Errorf("wrong proof size %d, want %d", got, want)
	}

	res := chainInner(v.hasher, leafHash, proof[:inner], index)
	res = chainBorderRight(v.hasher, res, proof[inner:])
	return res, nil
}

// VerifyConsistency checks that the passed-in consistency proof is valid
// between the passed in tree sizes, with respect to the corresponding root
// hashes. Requires 0 <= size1 <= size2.
func (v LogVerifier) VerifyConsistency(size1, size2 uint64, root1, root2 []byte, proof [][]byte) error {
	switch {
	case size2 < size1:
		return fmt.Errorf("size2 (%d) < size1 (%d)", size1, size2)
	case size1 == size2:
		if len(proof) > 0 {
			return errors.New("size1=size2, but proof is not empty")
		}
		return verifyMatch(root1, root2)
	case size1 == 0:
		// Any size greater than 0 is consistent with size 0.
		if len(proof) > 0 {
			return fmt.Errorf("expected empty proof, but got %d components", len(proof))
		}
		return nil // Proof OK.
	case len(proof) == 0:
		return errors.New("empty proof")
	}

	inner, border := decompInclProof(size1-1, size2)
	shift := bits.TrailingZeros64(size1)
	inner -= shift // Note: shift < inner if size1 < size2.

	// The proof includes the root hash for the sub-tree of size 2^shift.
	seed, start := proof[0], 1
	if size1 == 1<<uint(shift) { // Unless size1 is that very 2^shift.
		seed, start = root1, 0
	}
	if got, want := len(proof), start+inner+border; got != want {
		return fmt.Errorf("wrong proof size %d, want %d", got, want)
	}
	proof = proof[start:]
	// Now len(proof) == inner+border, and proof is effectively a suffix of
	// inclusion proof for entry |size1-1| in a tree of size |size2|.

	// Verify the first root.
	mask := (size1 - 1) >> uint(shift) // Start chaining from level |shift|.
	hash1 := chainInnerRight(v.hasher, seed, proof[:inner], mask)
	hash1 = chainBorderRight(v.hasher, hash1, proof[inner:])
	if err := verifyMatch(hash1, root1); err != nil {
		return err
	}

	// Verify the second root.
	hash2 := chainInner(v.hasher, seed, proof[:inner], mask)
	hash2 = chainBorderRight(v.hasher, hash2, proof[inner:])
	return verifyMatch(hash2, root2)
}

// decompInclProof breaks down inclusion proof for a leaf at the specified
// |index| in a tree of the specified |size| into 2 components. The splitting
// point between them is where paths to leaves |index| and |size-1| diverge.
// Returns lengths of the bottom and upper proof parts correspondingly. The sum
// of the two determines the correct length of the inclusion proof.
func decompInclProof(index, size uint64) (int, int) {
	inner := innerProofSize(index, size)
	border := bits.OnesCount64(index >> uint(inner))
	return inner, border
}

func innerProofSize(index, size uint64) int {
	return bits.Len64(index ^ (size - 1))
}

// chainInner computes a subtree hash for a node on or below the tree's right
// border. Assumes |proof| hashes are ordered from lower levels to upper, and
// |seed| is the initial subtree/leaf hash on the path located at the specified
// |index| on its level.
func chainInner(hasher LogHasher, seed []byte, proof [][]byte, index uint64) []byte {
	for i, h := range proof {
		if (index>>uint(i))&1 == 0 {
			seed = hasher.HashChildren(seed, h)
		} else {
			seed = hasher.HashChildren(h, seed)
		}
	}
	return seed
}

// chainInnerRight computes a subtree hash like chainInner, but only takes
// hashes to the left from the path into consideration, which effectively means
// the result is a hash of the corresponding earlier version of this subtree.
func chainInnerRight(hasher LogHasher, seed []byte, proof [][]byte, index uint64) []byte {
	for i, h := range proof {
		if (index>>uint(i))&1 == 1 {
			seed = hasher.HashChildren(h, seed)
		}
	}
	return seed
}

// chainBorderRight chains proof hashes along tree borders. This differs from
// inner chaining because |proof| contains only left-side subtree hashes.
func chainBorderRight(hasher LogHasher, seed []byte, proof [][]byte) []byte {
	for _, h := range proof {
		seed = hasher.HashChildren(h, seed)
	}
	return seed
}
<|endoftext|>


<commit_before>package cmd

import (
	"os/exec"

	"github.com/markbates/buffalo/buffalo/cmd/generate"
	"github.com/markbates/gentronics"
)

func newAppGenerator() *gentronics.Generator {
	g := gentronics.New()
	g.Add(gentronics.NewFile("main.go", nMain))
	g.Add(gentronics.NewFile("Procfile", nProcfile))
	g.Add(gentronics.NewFile("Procfile.development", nProcfileDev))
	g.Add(gentronics.NewFile(".buffalo.dev.yml", nRefresh))
	g.Add(gentronics.NewFile("actions/app.go", nApp))
	g.Add(gentronics.NewFile("actions/home.go", nHomeHandler))
	g.Add(gentronics.NewFile("actions/home_test.go", nHomeHandlerTest))
	g.Add(gentronics.NewFile("actions/render.go", nRender))
	g.Add(gentronics.NewFile("grifts/routes.go", nGriftRoutes))
	g.Add(gentronics.NewFile("templates/index.html", nIndexHTML))
	g.Add(gentronics.NewFile("templates/application.html", nApplicationHTML))
	g.Add(gentronics.NewFile("assets/application.js", ""))
	g.Add(gentronics.NewFile("assets/application.css", nApplicationCSS))
	g.Add(gentronics.NewFile(".gitignore", nGitignore))
	g.Add(gentronics.NewCommand(goGet("github.com/markbates/refresh/...")))
	g.Add(gentronics.NewCommand(goInstall("github.com/markbates/refresh")))
	g.Add(gentronics.NewCommand(goGet("github.com/markbates/grift/...")))
	g.Add(gentronics.NewCommand(goInstall("github.com/markbates/grift")))
	g.Add(generate.NewJQueryGenerator())
	g.Add(generate.NewBootstrapGenerator())
	g.Add(newSodaGenerator())
	g.Add(gentronics.NewCommand(appGoGet()))
	g.Add(generate.Fmt)
	return g
}

func appGoGet() *exec.Cmd {
	appArgs := []string{"get", "-t"}
	if verbose {
		appArgs = append(appArgs, "-v")
	}
	appArgs = append(appArgs, "./...")
	return exec.Command("go", appArgs...)
}

const nMain = `package main

import (
	"log"
	"net/http"

	"{{.actionsPath}}"
)

func main() {
	port := defaults.String(os.Getenv("PORT"), "3000")
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%s", port), actions.App()))
}

`
const nApp = `package actions

import (
	"net/http"

	"github.com/markbates/buffalo"
	{{if .withPop -}}
	"github.com/markbates/buffalo/middleware"
	"{{.modelsPath}}"
	{{end -}}
)

// App is where all routes and middleware for buffalo
// should be defined. This is the nerve center of your
// application.
func App() http.Handler {
	a := buffalo.Automatic(buffalo.Options{
		Env: "development",
	})

	{{if .withPop -}}
	a.Use(middleware.PopTransaction(models.DB))
	{{end -}}

	a.ServeFiles("/assets", assetsPath())
	a.GET("/", HomeHandler)

	return a
}
`

const nRender = `package actions

import (
	"net/http"
	"path"
	"runtime"

	"github.com/markbates/buffalo/render"
)

var r *render.Engine

func init() {
	r = render.New(render.Options{
		TemplatesPath: fromHere("../templates"),
		HTMLLayout:    "application.html",
	})
}

func assetsPath() http.Dir {
	return http.Dir(fromHere("../assets"))
}

func fromHere(p string) string {
	_, filename, _, _ := runtime.Caller(1)
	return path.Join(path.Dir(filename), p)
}
`

const nHomeHandler = `package actions

import "github.com/markbates/buffalo"

// HomeHandler is a default handler to serve up
// a home page.
func HomeHandler(c buffalo.Context) error {
	return c.Render(200, r.HTML("index.html"))
}
`

const nHomeHandlerTest = `package actions_test

import (
	"testing"

	"{{.actionsPath}}"
	"github.com/markbates/willie"
	"github.com/stretchr/testify/require"
)

func Test_HomeHandler(t *testing.T) {
	r := require.New(t)

	w := willie.New(actions.App())
	res := w.Request("/").Get()

	r.Equal(200, res.Code)
	r.Contains(res.Body.String(), "Welcome to Buffalo!")
}
`

const nIndexHTML = `<h1>Welcome to Buffalo!</h1>`

const nApplicationHTML = `<html>
<head>
  <meta charset="utf-8">
  <title>Buffalo - {{ .titleName }}</title>
  {{if .withBootstrap -}}
  <link rel="stylesheet" href="/assets/bootstrap.css" type="text/css" media="all" />
  {{end -}}
  <link rel="stylesheet" href="/assets/application.css" type="text/css" media="all" />
</head>
<body>

  {{"{{"}} yield {{"}}"}}

  {{if .withJQuery -}}
  <script src="/assets/jquery.js" type="text/javascript" charset="utf-8"></script>
  {{end -}}
  {{if .withBootstrap -}}
  <script src="/assets/bootstrap.js" type="text/javascript" charset="utf-8"></script>
  {{end -}}
  <script src="/assets/application.js" type="text/javascript" charset="utf-8"></script>
</body>
</html>
`

const nApplicationCSS = `body {
  font-family: helvetica;
}
`

const nGitignore = `vendor/
**/*.log
**/*.sqlite
bin/
node_modules/
{{ .name }}
`

const nGriftRoutes = `package grifts

import (
	"os"

	"github.com/markbates/buffalo"
	. "github.com/markbates/grift/grift"
	"{{.actionsPath}}"
	"github.com/olekukonko/tablewriter"
)

var _ = Add("routes", func(c *Context) error {
	a := actions.App().(*buffalo.App)
	routes := a.Routes()

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Method", "Path", "Handler"})
	for _, r := range routes {
		table.Append([]string{r.Method, r.Path, r.HandlerName})
	}
	table.SetCenterSeparator("|")
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()
	return nil
})`

const nRefresh = `app_root: .
ignored_folders:
- vendor
- log
- tmp
included_extensions:
- .go
- .html
build_path: /tmp
build_delay: 200ns
binary_name: {{.name}}-build
command_flags: []
enable_colors: true
log_name: buffalo
`

const nProcfile = `web: {{.name}}`
const nProcfileDev = `web: buffalo dev`
<commit_msg>define ENV at the top of actions/app.go closes #46<commit_after>package cmd

import (
	"os/exec"

	"github.com/markbates/buffalo/buffalo/cmd/generate"
	"github.com/markbates/gentronics"
)

func newAppGenerator() *gentronics.Generator {
	g := gentronics.New()
	g.Add(gentronics.NewFile("main.go", nMain))
	g.Add(gentronics.NewFile("Procfile", nProcfile))
	g.Add(gentronics.NewFile("Procfile.development", nProcfileDev))
	g.Add(gentronics.NewFile(".buffalo.dev.yml", nRefresh))
	g.Add(gentronics.NewFile("actions/app.go", nApp))
	g.Add(gentronics.NewFile("actions/home.go", nHomeHandler))
	g.Add(gentronics.NewFile("actions/home_test.go", nHomeHandlerTest))
	g.Add(gentronics.NewFile("actions/render.go", nRender))
	g.Add(gentronics.NewFile("grifts/routes.go", nGriftRoutes))
	g.Add(gentronics.NewFile("templates/index.html", nIndexHTML))
	g.Add(gentronics.NewFile("templates/application.html", nApplicationHTML))
	g.Add(gentronics.NewFile("assets/application.js", ""))
	g.Add(gentronics.NewFile("assets/application.css", nApplicationCSS))
	g.Add(gentronics.NewFile(".gitignore", nGitignore))
	g.Add(gentronics.NewCommand(goGet("github.com/markbates/refresh/...")))
	g.Add(gentronics.NewCommand(goInstall("github.com/markbates/refresh")))
	g.Add(gentronics.NewCommand(goGet("github.com/markbates/grift/...")))
	g.Add(gentronics.NewCommand(goInstall("github.com/markbates/grift")))
	g.Add(generate.NewJQueryGenerator())
	g.Add(generate.NewBootstrapGenerator())
	g.Add(newSodaGenerator())
	g.Add(gentronics.NewCommand(appGoGet()))
	g.Add(generate.Fmt)
	return g
}

func appGoGet() *exec.Cmd {
	appArgs := []string{"get", "-t"}
	if verbose {
		appArgs = append(appArgs, "-v")
	}
	appArgs = append(appArgs, "./...")
	return exec.Command("go", appArgs...)
}

const nMain = `package main

import (
	"log"
	"net/http"

	"{{.actionsPath}}"
)

func main() {
	port := defaults.String(os.Getenv("PORT"), "3000")
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%s", port), actions.App()))
}

`
const nApp = `package actions

import (
	"net/http"

	"github.com/markbates/buffalo"
	{{if .withPop -}}
	"github.com/markbates/buffalo/middleware"
	"{{.modelsPath}}"
	{{end -}}
)

var ENV = defaults.String(os.Getenv("GO_ENV"), "development")

// App is where all routes and middleware for buffalo
// should be defined. This is the nerve center of your
// application.
func App() http.Handler {
	a := buffalo.Automatic(buffalo.Options{
		Env: ENV,
	})

	{{if .withPop -}}
	a.Use(middleware.PopTransaction(models.DB))
	{{end -}}

	a.ServeFiles("/assets", assetsPath())
	a.GET("/", HomeHandler)

	return a
}
`

const nRender = `package actions

import (
	"net/http"
	"path"
	"runtime"

	"github.com/markbates/buffalo/render"
)

var r *render.Engine

func init() {
	r = render.New(render.Options{
		TemplatesPath: fromHere("../templates"),
		HTMLLayout:    "application.html",
	})
}

func assetsPath() http.Dir {
	return http.Dir(fromHere("../assets"))
}

func fromHere(p string) string {
	_, filename, _, _ := runtime.Caller(1)
	return path.Join(path.Dir(filename), p)
}
`

const nHomeHandler = `package actions

import "github.com/markbates/buffalo"

// HomeHandler is a default handler to serve up
// a home page.
func HomeHandler(c buffalo.Context) error {
	return c.Render(200, r.HTML("index.html"))
}
`

const nHomeHandlerTest = `package actions_test

import (
	"testing"

	"{{.actionsPath}}"
	"github.com/markbates/willie"
	"github.com/stretchr/testify/require"
)

func Test_HomeHandler(t *testing.T) {
	r := require.New(t)

	w := willie.New(actions.App())
	res := w.Request("/").Get()

	r.Equal(200, res.Code)
	r.Contains(res.Body.String(), "Welcome to Buffalo!")
}
`

const nIndexHTML = `<h1>Welcome to Buffalo!</h1>`

const nApplicationHTML = `<html>
<head>
  <meta charset="utf-8">
  <title>Buffalo - {{ .titleName }}</title>
  {{if .withBootstrap -}}
  <link rel="stylesheet" href="/assets/bootstrap.css" type="text/css" media="all" />
  {{end -}}
  <link rel="stylesheet" href="/assets/application.css" type="text/css" media="all" />
</head>
<body>

  {{"{{"}} yield {{"}}"}}

  {{if .withJQuery -}}
  <script src="/assets/jquery.js" type="text/javascript" charset="utf-8"></script>
  {{end -}}
  {{if .withBootstrap -}}
  <script src="/assets/bootstrap.js" type="text/javascript" charset="utf-8"></script>
  {{end -}}
  <script src="/assets/application.js" type="text/javascript" charset="utf-8"></script>
</body>
</html>
`

const nApplicationCSS = `body {
  font-family: helvetica;
}
`

const nGitignore = `vendor/
**/*.log
**/*.sqlite
bin/
node_modules/
{{ .name }}
`

const nGriftRoutes = `package grifts

import (
	"os"

	"github.com/markbates/buffalo"
	. "github.com/markbates/grift/grift"
	"{{.actionsPath}}"
	"github.com/olekukonko/tablewriter"
)

var _ = Add("routes", func(c *Context) error {
	a := actions.App().(*buffalo.App)
	routes := a.Routes()

	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Method", "Path", "Handler"})
	for _, r := range routes {
		table.Append([]string{r.Method, r.Path, r.HandlerName})
	}
	table.SetCenterSeparator("|")
	table.SetAlignment(tablewriter.ALIGN_LEFT)
	table.Render()
	return nil
})`

const nRefresh = `app_root: .
ignored_folders:
- vendor
- log
- tmp
included_extensions:
- .go
- .html
build_path: /tmp
build_delay: 200ns
binary_name: {{.name}}-build
command_flags: []
enable_colors: true
log_name: buffalo
`

const nProcfile = `web: {{.name}}`
const nProcfileDev = `web: buffalo dev`
<|endoftext|>


<commit_before>package quiclatest

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"

	"github.com/ami-GS/gQUIC/latest/qtype"
)

type Packet interface {
	// doesn't need genWire
	GetWire() ([]byte, error)
	SetWire(wire []byte)
	GetHeader() PacketHeader
	String() string
	SetHeader(ph PacketHeader)
	GetFrames() []Frame
	SetFrames(fs []Frame)
}

func ParsePacket(data []byte) (packet Packet, idx int, err error) {
	version := binary.BigEndian.Uint32(data[1:5])
	if version == 0 {
		// VersionNegotiationPacket use all UDP datagram (doesn't have frame)
		return ParseVersionNegotiationPacket(data)
	}
	header, idx, err := PacketHeaderParserMap[PacketHeaderType(data[0]&0x80)](data) // ParseHeader
	if err != nil {
		return nil, 0, err
	}
	//TODO: want to know the number of padding
	fs, idxTmp, err := ParseFrames(data[idx:])
	if err != nil {
		return nil, 0, err
	}
	packet, err = newPacket(header, fs)
	packet.SetWire(data[idx:])
	if err != nil {
		return nil, 0, err
	}
	return packet, idx + idxTmp, err
}

type BasePacket struct {
	Header     PacketHeader
	Frames     []Frame
	PaddingNum int
	payload    []byte
}

func (bp *BasePacket) GetHeader() PacketHeader {
	return bp.Header
}
func (bp *BasePacket) SetHeader(h PacketHeader) {
	bp.Header = h
}

func (bp *BasePacket) SetFrames(fs []Frame) {
	bp.Frames = fs
}

func (bp *BasePacket) GetFrames() []Frame {
	return bp.Frames
}

func (bp *BasePacket) String() string {
	frameStr := ""
	for _, frame := range bp.Frames {
		frameStr += fmt.Sprintf("\n\t%s", frame)
	}
	return fmt.Sprintf("%s\nPaddingLen:%d\n\t{%s\n\t}", bp.Header.String(), bp.PaddingNum, frameStr)
}

func (bp *BasePacket) SetWire(wire []byte) {
	bp.payload = wire
}

func (bp *BasePacket) GetPayloadLen() int {
	if bp.payload != nil {
		return len(bp.payload)
	}

	length := 0
	for _, frame := range bp.Frames {
		length += frame.GetWireSize()
	}
	return length
}

//GetWire of BasePacket assembles all wires, from header wire to frame wires
func (bp *BasePacket) GetWire() (wire []byte, err error) {
	if bp.payload != nil {
		// bp.wire is filled after parsing
		return append(bp.Header.GetWire(), bp.payload...), nil
	}
	hWire := bp.Header.GetWire()
	if err != nil {
		return nil, err
	}
	bp.payload, err = GetFrameWires(bp.Frames)
	if err != nil {
		return nil, err
	}
	if bp.PaddingNum != 0 {
		bp.payload = append(bp.payload, make([]byte, bp.PaddingNum)...)
	}
	// TODO: protect for short header?
	return append(hWire, bp.payload...), nil
}

func newPacket(ph PacketHeader, fs []Frame) (Packet, error) {
	// TODO: needs frame type validation for each packet type
	if lh, ok := ph.(*LongHeader); ok {
		payloadLen := 0
		for _, frame := range fs {
			payloadLen += frame.GetWireSize()
		}

		switch lh.PacketType {
		case InitialPacketType:
			// TODO: check whether it is stream frame or not
			return &InitialPacket{
				BasePacket: &BasePacket{
					Header:     ph,
					Frames:     fs,
					PaddingNum: InitialPacketMinimumPayloadSize - payloadLen,
				},
			}, nil
		case RetryPacketType:
			return &RetryPacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case HandshakePacketType:
			return &HandshakePacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case ZeroRTTProtectedPacketType:
			return &ProtectedPacket{
				BasePacket: &BasePacket{
					Header: ph,
					Frames: fs,
				},
				RTT: 0,
			}, nil
		default:
			// error type is not defined
			return nil, qtype.ProtocolViolation
		}
	} else if _, ok := ph.(*ShortHeader); ok {
		return &ProtectedPacket{
			BasePacket: &BasePacket{
				Header: ph,
				Frames: fs,
			},
			RTT: 1,
		}, nil
	}
	// error type is not defined
	return nil, qtype.ProtocolViolation
}

// long header with type of 0x7F
type InitialPacket struct {
	*BasePacket
}

const InitialPacketMinimumPayloadSize = 1200

func NewInitialPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, sFrame *StreamFrame) *InitialPacket {
	sFrameLen := sFrame.GetWireSize()
	var frames []Frame
	var lh *LongHeader
	paddingNum := 0
	if InitialPacketMinimumPayloadSize <= sFrameLen {
		// TODO: need to check when sFrameLen is over MTUIPv4
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(sFrameLen))
		frames = []Frame{sFrame}
	} else {
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, InitialPacketMinimumPayloadSize)
		frames = []Frame{sFrame}
		paddingNum = InitialPacketMinimumPayloadSize - sFrameLen
	}
	p := &InitialPacket{
		BasePacket: &BasePacket{
			Header:     lh,
			Frames:     frames,
			PaddingNum: paddingNum,
		},
	}
	return p
}

// long header with type of 0x7E
type RetryPacket struct {
	*BasePacket
}

func NewRetryPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *RetryPacket {
	if len(frames) < 2 {
		return nil
	}
	if _, ok := frames[0].(*StreamFrame); !ok {
		return nil
	}
	if _, ok := frames[1].(*AckFrame); !ok {
		return nil
	}
	for i := 2; i < len(frames); i++ {
		if _, ok := frames[i].(*PaddingFrame); !ok {
			return nil
		}
	}

	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		payloadLen += frames[i].GetWireSize()
	}

	return &RetryPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(RetryPacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with type of 0x7D
type HandshakePacket struct {
	*BasePacket
}

func NewHandshakePacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *HandshakePacket {
	minimumReq := false
	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		switch frames[i].(type) {
		case *StreamFrame:
			minimumReq = true
			payloadLen += frames[i].GetWireSize()
		case *AckFrame, *PathChallengeFrame /*or*/, *PathResponseFrame, *ConnectionCloseFrame:
			payloadLen += frames[i].GetWireSize()
		default:
			return nil
		}
	}
	if !minimumReq {
		return nil
	}

	return &HandshakePacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(HandshakePacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with 0-RTT (type:0x7C)
// short header with 1-RTT
type ProtectedPacket struct {
	*BasePacket
	RTT byte
}

func NewProtectedPacket(version qtype.Version, key bool, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, rtt byte, frames []Frame) *ProtectedPacket {
	var header PacketHeader
	if rtt == 0 {
		payloadLen := 0
		for _, frame := range frames {
			payloadLen += frame.GetWireSize()
		}
		header = NewLongHeader(ZeroRTTProtectedPacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(payloadLen))
	} else if rtt == 1 {
		header = NewShortHeader(key, destConnID, packetNumber)
	} else {
		// error
	}

	return &ProtectedPacket{
		BasePacket: &BasePacket{
			Header: header,
			Frames: frames,
		},
		RTT: rtt,
	}
}

func (p ProtectedPacket) String() string {
	return fmt.Sprintf("%s, RTT:%d", p.BasePacket, p.RTT)
}

/*
    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
   +-+-+-+-+-+-+-+-+
   |1|  Unused (7) |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                          Version (32)                         |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |DCIL(4)|SCIL(4)|
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |               Destination Connection ID (0/32..144)         ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                 Source Connection ID (0/32..144)            ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                    Supported Version 1 (32)                 ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version 2 (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
                                  ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version N (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/

// Version negotiation doesn't use long header, but have similar form
type VersionNegotiationPacket struct {
	*BasePacketHeader
	wire              []byte
	Version           qtype.Version
	DCIL              byte
	SCIL              byte
	SupportedVersions []qtype.Version
}

func NewVersionNegotiationPacket(destConnID, srcConnID qtype.ConnectionID, supportedVersions []qtype.Version) *VersionNegotiationPacket {
	dcil := 0
	if destConnID != nil {
		dcil = len(destConnID) - 3
	}
	scil := 0
	if srcConnID != nil {
		scil = len(srcConnID) - 3
	}

	p := &VersionNegotiationPacket{
		BasePacketHeader: &BasePacketHeader{
			DestConnID:   destConnID,
			SrcConnID:    srcConnID,
			PacketNumber: 0,
		},
		Version:           0,
		DCIL:              byte(dcil),
		SCIL:              byte(scil),
		SupportedVersions: supportedVersions,
	}
	var err error
	p.wire, err = p.genWire()
	if err != nil {
		return nil
	}
	return p
}

func ParseVersionNegotiationPacket(data []byte) (Packet, int, error) {
	idx := 0
	packet := NewVersionNegotiationPacket(nil, nil, nil)
	if data[0]&0x80 != 0x80 {
		//TODO: error
	}
	idx++
	packet.Version = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
	if packet.Version != 0 {
		// must be zero
		// TODO: error
	}
	idx += 4
	packet.DCIL = data[idx] >> 4
	packet.SCIL = data[idx] & 0x0f
	idx++
	if packet.DCIL != 0 {
		dcil := int(packet.DCIL + 3)
		packet.BasePacketHeader.DestConnID, _ = qtype.ReadConnectionID(data[idx:], dcil)
		idx += dcil
	}
	if packet.SCIL != 0 {
		scil := int(packet.SCIL + 3)
		packet.BasePacketHeader.SrcConnID, _ = qtype.ReadConnectionID(data[idx:], scil)
		idx += scil
	}
	numVersions := (len(data) - idx) / 4
	packet.SupportedVersions = make([]qtype.Version, numVersions)
	for i := 0; i < numVersions; i++ {
		packet.SupportedVersions[i] = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
		idx += 4
	}
	packet.wire = data
	return packet, idx, nil
}

func (p VersionNegotiationPacket) genWire() (wire []byte, err error) {
	wireLen := 6 + len(p.SupportedVersions)*4
	if p.DCIL != 0 {
		wireLen += int(p.DCIL + 3)
	}
	if p.SCIL != 0 {
		wireLen += int(p.SCIL + 3)
	}
	//wireLen = qtype.MTUIPv4
	wire = make([]byte, wireLen)
	if _, err := rand.Read(wire[0:1]); err != nil {
		return nil, err
	}
	wire[0] |= 0x80
	binary.BigEndian.PutUint32(wire[1:], uint32(p.Version))
	wire[5] = (p.DCIL << 4) | p.SCIL
	idx := 6
	if p.DCIL != 0 {
		for i := 0; i < int(p.DCIL+3); i++ {
			wire[idx+i] = p.DestConnID[i]
		}
		idx += int(p.DCIL + 3)
	}
	if p.SCIL != 0 {
		for i := 0; i < int(p.SCIL+3); i++ {
			wire[idx+i] = p.SrcConnID[i]
		}
		idx += int(p.SCIL + 3)
	}

	for i, version := range p.SupportedVersions {
		binary.BigEndian.PutUint32(wire[idx+i*4:], uint32(version))
	}
	// TODO: VersionNegotiationPacket fills MTU
	return
}

func (p VersionNegotiationPacket) SetWire(wire []byte) {
	// actual set is done in ParseVersionNegotiationPacket()
}

func (p VersionNegotiationPacket) GetWire() ([]byte, error) {
	return p.wire, nil
}

func (p VersionNegotiationPacket) GetHeader() PacketHeader {
	return nil
}
func (p VersionNegotiationPacket) SetHeader(h PacketHeader) {
	// no op?
}
func (p VersionNegotiationPacket) SetFrames(fs []Frame) {
	// no op?
}
func (p VersionNegotiationPacket) GetFrames() []Frame {
	return nil
}
func (p VersionNegotiationPacket) String() string {
	return fmt.Sprintf("NoHeader:VersionNegotiationPacket\tVer:%d\nDCIL:%d,SCIL:%d\n%s\nSupported Versions:%v", p.Version, p.DCIL, p.SCIL, p.BasePacketHeader, p.SupportedVersions)
}
<commit_msg>detail handshake packet<commit_after>package quiclatest

import (
	"crypto/rand"
	"encoding/binary"
	"fmt"

	"github.com/ami-GS/gQUIC/latest/qtype"
)

type Packet interface {
	// doesn't need genWire
	GetWire() ([]byte, error)
	SetWire(wire []byte)
	GetHeader() PacketHeader
	String() string
	SetHeader(ph PacketHeader)
	GetFrames() []Frame
	SetFrames(fs []Frame)
}

func ParsePacket(data []byte) (packet Packet, idx int, err error) {
	version := binary.BigEndian.Uint32(data[1:5])
	if version == 0 {
		// VersionNegotiationPacket use all UDP datagram (doesn't have frame)
		return ParseVersionNegotiationPacket(data)
	}
	header, idx, err := PacketHeaderParserMap[PacketHeaderType(data[0]&0x80)](data) // ParseHeader
	if err != nil {
		return nil, 0, err
	}
	//TODO: want to know the number of padding
	fs, idxTmp, err := ParseFrames(data[idx:])
	if err != nil {
		return nil, 0, err
	}
	packet, err = newPacket(header, fs)
	packet.SetWire(data[idx:])
	if err != nil {
		return nil, 0, err
	}
	return packet, idx + idxTmp, err
}

type BasePacket struct {
	Header     PacketHeader
	Frames     []Frame
	PaddingNum int
	payload    []byte
}

func (bp *BasePacket) GetHeader() PacketHeader {
	return bp.Header
}
func (bp *BasePacket) SetHeader(h PacketHeader) {
	bp.Header = h
}

func (bp *BasePacket) SetFrames(fs []Frame) {
	bp.Frames = fs
}

func (bp *BasePacket) GetFrames() []Frame {
	return bp.Frames
}

func (bp *BasePacket) String() string {
	frameStr := ""
	for _, frame := range bp.Frames {
		frameStr += fmt.Sprintf("\n\t%s", frame)
	}
	return fmt.Sprintf("%s\nPaddingLen:%d\n\t{%s\n\t}", bp.Header.String(), bp.PaddingNum, frameStr)
}

func (bp *BasePacket) SetWire(wire []byte) {
	bp.payload = wire
}

func (bp *BasePacket) GetPayloadLen() int {
	if bp.payload != nil {
		return len(bp.payload)
	}

	length := 0
	for _, frame := range bp.Frames {
		length += frame.GetWireSize()
	}
	return length
}

//GetWire of BasePacket assembles all wires, from header wire to frame wires
func (bp *BasePacket) GetWire() (wire []byte, err error) {
	if bp.payload != nil {
		// bp.wire is filled after parsing
		return append(bp.Header.GetWire(), bp.payload...), nil
	}
	hWire := bp.Header.GetWire()
	if err != nil {
		return nil, err
	}
	bp.payload, err = GetFrameWires(bp.Frames)
	if err != nil {
		return nil, err
	}
	if bp.PaddingNum != 0 {
		bp.payload = append(bp.payload, make([]byte, bp.PaddingNum)...)
	}
	// TODO: protect for short header?
	return append(hWire, bp.payload...), nil
}

func newPacket(ph PacketHeader, fs []Frame) (Packet, error) {
	// TODO: needs frame type validation for each packet type
	if lh, ok := ph.(*LongHeader); ok {
		payloadLen := 0
		for _, frame := range fs {
			payloadLen += frame.GetWireSize()
		}

		switch lh.PacketType {
		case InitialPacketType:
			// TODO: check whether it is stream frame or not
			return &InitialPacket{
				BasePacket: &BasePacket{
					Header:     ph,
					Frames:     fs,
					PaddingNum: InitialPacketMinimumPayloadSize - payloadLen,
				},
			}, nil
		case RetryPacketType:
			return &RetryPacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case HandshakePacketType:
			return &HandshakePacket{
				&BasePacket{
					Header: ph,
					Frames: fs,
				},
			}, nil
		case ZeroRTTProtectedPacketType:
			return &ProtectedPacket{
				BasePacket: &BasePacket{
					Header: ph,
					Frames: fs,
				},
				RTT: 0,
			}, nil
		default:
			// error type is not defined
			return nil, qtype.ProtocolViolation
		}
	} else if _, ok := ph.(*ShortHeader); ok {
		return &ProtectedPacket{
			BasePacket: &BasePacket{
				Header: ph,
				Frames: fs,
			},
			RTT: 1,
		}, nil
	}
	// error type is not defined
	return nil, qtype.ProtocolViolation
}

// long header with type of 0x7F
type InitialPacket struct {
	*BasePacket
}

const InitialPacketMinimumPayloadSize = 1200

func NewInitialPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, sFrame *StreamFrame) *InitialPacket {
	sFrameLen := sFrame.GetWireSize()
	var frames []Frame
	var lh *LongHeader
	paddingNum := 0
	if InitialPacketMinimumPayloadSize <= sFrameLen {
		// TODO: need to check when sFrameLen is over MTUIPv4
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(sFrameLen))
		frames = []Frame{sFrame}
	} else {
		lh = NewLongHeader(InitialPacketType, version, destConnID, srcConnID, packetNumber, InitialPacketMinimumPayloadSize)
		frames = []Frame{sFrame}
		paddingNum = InitialPacketMinimumPayloadSize - sFrameLen
	}
	p := &InitialPacket{
		BasePacket: &BasePacket{
			Header:     lh,
			Frames:     frames,
			PaddingNum: paddingNum,
		},
	}
	return p
}

// long header with type of 0x7E
type RetryPacket struct {
	*BasePacket
}

func NewRetryPacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *RetryPacket {
	if len(frames) < 2 {
		return nil
	}
	if _, ok := frames[0].(*StreamFrame); !ok {
		return nil
	}
	if _, ok := frames[1].(*AckFrame); !ok {
		return nil
	}
	for i := 2; i < len(frames); i++ {
		if _, ok := frames[i].(*PaddingFrame); !ok {
			return nil
		}
	}

	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		payloadLen += frames[i].GetWireSize()
	}

	return &RetryPacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(RetryPacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with type of 0x7D
type HandshakePacket struct {
	*BasePacket
}

func NewHandshakePacket(version qtype.Version, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, frames []Frame) *HandshakePacket {
	minimumReq := false
	payloadLen := 0
	for i := 0; i < len(frames); i++ {
		switch frames[i].(type) {
		case *ConnectionCloseFrame:
			// This stands for the handshake is unsaccessful
			minimumReq = true
		case *StreamFrame:
			minimumReq = true
		case *AckFrame, *PathChallengeFrame /*or*/, *PathResponseFrame:
		default:
			// TODO: error, handshake packet cannot have these frames
			return nil
		}
		payloadLen += frames[i].GetWireSize()
	}
	if !minimumReq {
		return nil
	}

	return &HandshakePacket{
		BasePacket: &BasePacket{
			Header: NewLongHeader(HandshakePacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(payloadLen)),
			Frames: frames,
		},
	}
}

// long header with 0-RTT (type:0x7C)
// short header with 1-RTT
type ProtectedPacket struct {
	*BasePacket
	RTT byte
}

func NewProtectedPacket(version qtype.Version, key bool, destConnID, srcConnID qtype.ConnectionID, packetNumber qtype.PacketNumber, rtt byte, frames []Frame) *ProtectedPacket {
	var header PacketHeader
	if rtt == 0 {
		payloadLen := 0
		for _, frame := range frames {
			payloadLen += frame.GetWireSize()
		}
		header = NewLongHeader(ZeroRTTProtectedPacketType, version, destConnID, srcConnID, packetNumber, qtype.QuicInt(payloadLen))
	} else if rtt == 1 {
		header = NewShortHeader(key, destConnID, packetNumber)
	} else {
		// error
	}

	return &ProtectedPacket{
		BasePacket: &BasePacket{
			Header: header,
			Frames: frames,
		},
		RTT: rtt,
	}
}

func (p ProtectedPacket) String() string {
	return fmt.Sprintf("%s, RTT:%d", p.BasePacket, p.RTT)
}

/*
    0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1 2 3 4 5 6 7 8 9 0 1
   +-+-+-+-+-+-+-+-+
   |1|  Unused (7) |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                          Version (32)                         |
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |DCIL(4)|SCIL(4)|
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |               Destination Connection ID (0/32..144)         ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                 Source Connection ID (0/32..144)            ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                    Supported Version 1 (32)                 ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version 2 (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
                                  ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
   |                   [Supported Version N (32)]                ...
   +-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+-+
*/

// Version negotiation doesn't use long header, but have similar form
type VersionNegotiationPacket struct {
	*BasePacketHeader
	wire              []byte
	Version           qtype.Version
	DCIL              byte
	SCIL              byte
	SupportedVersions []qtype.Version
}

func NewVersionNegotiationPacket(destConnID, srcConnID qtype.ConnectionID, supportedVersions []qtype.Version) *VersionNegotiationPacket {
	dcil := 0
	if destConnID != nil {
		dcil = len(destConnID) - 3
	}
	scil := 0
	if srcConnID != nil {
		scil = len(srcConnID) - 3
	}

	p := &VersionNegotiationPacket{
		BasePacketHeader: &BasePacketHeader{
			DestConnID:   destConnID,
			SrcConnID:    srcConnID,
			PacketNumber: 0,
		},
		Version:           0,
		DCIL:              byte(dcil),
		SCIL:              byte(scil),
		SupportedVersions: supportedVersions,
	}
	var err error
	p.wire, err = p.genWire()
	if err != nil {
		return nil
	}
	return p
}

func ParseVersionNegotiationPacket(data []byte) (Packet, int, error) {
	idx := 0
	packet := NewVersionNegotiationPacket(nil, nil, nil)
	if data[0]&0x80 != 0x80 {
		//TODO: error
	}
	idx++
	packet.Version = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
	if packet.Version != 0 {
		// must be zero
		// TODO: error
	}
	idx += 4
	packet.DCIL = data[idx] >> 4
	packet.SCIL = data[idx] & 0x0f
	idx++
	if packet.DCIL != 0 {
		dcil := int(packet.DCIL + 3)
		packet.BasePacketHeader.DestConnID, _ = qtype.ReadConnectionID(data[idx:], dcil)
		idx += dcil
	}
	if packet.SCIL != 0 {
		scil := int(packet.SCIL + 3)
		packet.BasePacketHeader.SrcConnID, _ = qtype.ReadConnectionID(data[idx:], scil)
		idx += scil
	}
	numVersions := (len(data) - idx) / 4
	packet.SupportedVersions = make([]qtype.Version, numVersions)
	for i := 0; i < numVersions; i++ {
		packet.SupportedVersions[i] = qtype.Version(binary.BigEndian.Uint32(data[idx:]))
		idx += 4
	}
	packet.wire = data
	return packet, idx, nil
}

func (p VersionNegotiationPacket) genWire() (wire []byte, err error) {
	wireLen := 6 + len(p.SupportedVersions)*4
	if p.DCIL != 0 {
		wireLen += int(p.DCIL + 3)
	}
	if p.SCIL != 0 {
		wireLen += int(p.SCIL + 3)
	}
	//wireLen = qtype.MTUIPv4
	wire = make([]byte, wireLen)
	if _, err := rand.Read(wire[0:1]); err != nil {
		return nil, err
	}
	wire[0] |= 0x80
	binary.BigEndian.PutUint32(wire[1:], uint32(p.Version))
	wire[5] = (p.DCIL << 4) | p.SCIL
	idx := 6
	if p.DCIL != 0 {
		for i := 0; i < int(p.DCIL+3); i++ {
			wire[idx+i] = p.DestConnID[i]
		}
		idx += int(p.DCIL + 3)
	}
	if p.SCIL != 0 {
		for i := 0; i < int(p.SCIL+3); i++ {
			wire[idx+i] = p.SrcConnID[i]
		}
		idx += int(p.SCIL + 3)
	}

	for i, version := range p.SupportedVersions {
		binary.BigEndian.PutUint32(wire[idx+i*4:], uint32(version))
	}
	// TODO: VersionNegotiationPacket fills MTU
	return
}

func (p VersionNegotiationPacket) SetWire(wire []byte) {
	// actual set is done in ParseVersionNegotiationPacket()
}

func (p VersionNegotiationPacket) GetWire() ([]byte, error) {
	return p.wire, nil
}

func (p VersionNegotiationPacket) GetHeader() PacketHeader {
	return nil
}
func (p VersionNegotiationPacket) SetHeader(h PacketHeader) {
	// no op?
}
func (p VersionNegotiationPacket) SetFrames(fs []Frame) {
	// no op?
}
func (p VersionNegotiationPacket) GetFrames() []Frame {
	return nil
}
func (p VersionNegotiationPacket) String() string {
	return fmt.Sprintf("NoHeader:VersionNegotiationPacket\tVer:%d\nDCIL:%d,SCIL:%d\n%s\nSupported Versions:%v", p.Version, p.DCIL, p.SCIL, p.BasePacketHeader, p.SupportedVersions)
}
<|endoftext|>


<commit_before>/*
Package measured wraps a dialer to measure the delay, throughput and errors of the connection made.
A list of reporters can be plugged in to distribute the results to different target.
*/
package measured

import (
	"net"
	"strings"

	"github.com/getlantern/golog"
)

// Stats encapsulates the statistics to report
type Stats struct {
	Type   string
	Tags   map[string]string
	Fields map[string]interface{}
}

// Reporter encapsulates different ways to report statistics
type Reporter interface {
	Submit(*Stats) error
}

var (
	reporters   []Reporter
	defaultTags map[string]string
	log         = golog.LoggerFor("measured")
	chStats     = make(chan *Stats)
	chStop      = make(chan interface{})
)

// DialFunc is the type of function measured can wrap
type DialFunc func(net, addr string) (net.Conn, error)

// Reset resets the measured package
func Reset() {
	reporters = []Reporter{}
}

// AddReporter add a new way to report statistics
func AddReporter(r Reporter) {
	reporters = append(reporters, r)
}

// SetDefaults set a few default tags sending every time
func SetDefaults(defaults map[string]string) {
	defaultTags = defaults
}

// Start runs the measured loop
func Start() {
	go run()
}

// Stop stops the measured loop
func Stop() {
	log.Debug("Stopping measured loop...")
	select {
	case chStop <- nil:
	default:
		log.Error("Failed to send stop signal")
	}
}

// Dialer wraps a dial function to measure various statistics
func Dialer(d DialFunc, via string) DialFunc {
	return func(net, addr string) (net.Conn, error) {
		c, err := d(net, addr)
		if err != nil {
			reportError(via, err, "dial")
		}
		return measuredConn{c, via}, err
	}
}

func run() {
	log.Debug("Measured loop started")
	for {
		select {
		case s := <-chStats:
			for _, r := range reporters {
				for k, v := range defaultTags {
					s.Tags[k] = v
				}
				if err := r.Submit(s); err != nil {
					log.Errorf("report error to influxdb failed: %s", err)
				} else {
					log.Tracef("submitted error to influxdb: %v", s)
				}
			}
		case <-chStop:
			log.Debug("Measured loop stopped")
			return
		}
	}
}

func reportError(addr string, err error, phase string) {
	splitted := strings.Split(err.Error(), ":")
	lastIndex := len(splitted) - 1
	if lastIndex < 0 {
		lastIndex = 0
	}
	e := strings.Trim(splitted[lastIndex], " ")
	select {
	case chStats <- &Stats{
		Type: "errors",
		Tags: map[string]string{
			"server": addr,
			"error":  e,
			"phase":  phase,
		},
		Fields: map[string]interface{}{"value": 1},
	}:
	default:
	}
}

type measuredConn struct {
	net.Conn
	addr string
}

// Read() implements the function from net.Conn
func (mc measuredConn) Read(b []byte) (n int, err error) {
	n, err = mc.Conn.Read(b)
	if err != nil {
		reportError(mc.addr, err, "read")
	}
	return
}

// Write() implements the function from net.Conn
func (mc measuredConn) Write(b []byte) (n int, err error) {
	n, err = mc.Conn.Write(b)
	if err != nil {
		reportError(mc.addr, err, "write")
	}
	return
}

// Close() implements the function from net.Conn
func (mc measuredConn) Close() (err error) {
	err = mc.Conn.Close()
	if err != nil {
		reportError(mc.addr, err, "close")
	}
	return
}
<commit_msg>get country from pubsub<commit_after>/*
Package measured wraps a dialer to measure the delay, throughput and errors of the connection made.
A list of reporters can be plugged in to distribute the results to different target.
*/
package measured

import (
	"net"
	"strings"
	"sync/atomic"

	"github.com/getlantern/golog"
)

// Stats encapsulates the statistics to report
type Stats struct {
	Type   string
	Tags   map[string]string
	Fields map[string]interface{}
}

// Reporter encapsulates different ways to report statistics
type Reporter interface {
	Submit(*Stats) error
}

var (
	reporters   []Reporter
	defaultTags atomic.Value
	running     uint32
	log         = golog.LoggerFor("measured")
	chStats     = make(chan *Stats)
	chStop      = make(chan interface{})
)

func init() {
	defaultTags.Store(map[string]string{})
}

// DialFunc is the type of function measured can wrap
type DialFunc func(net, addr string) (net.Conn, error)

// Reset resets the measured package
func Reset() {
	reporters = []Reporter{}
}

// AddReporter add a new way to report statistics
func AddReporter(r Reporter) {
	reporters = append(reporters, r)
}

// SetDefaults set a few default tags sending every time
func SetDefaults(defaults map[string]string) {
	defaultTags.Store(defaults)
}

// Start runs the measured loop
func Start() {
	go run()
}

// Stop stops the measured loop
func Stop() {
	if atomic.LoadUint32(&running) == 0 {
		return
	}
	log.Debug("Stopping measured loop...")
	select {
	case chStop <- nil:
	default:
		log.Error("Failed to send stop signal")
	}
}

// Dialer wraps a dial function to measure various statistics
func Dialer(d DialFunc, via string) DialFunc {
	return func(net, addr string) (net.Conn, error) {
		c, err := d(net, addr)
		if err != nil {
			reportError(via, err, "dial")
		}
		return measuredConn{c, via}, err
	}
}

func run() {
	log.Debug("Measured loop started")
	atomic.StoreUint32(&running, 1)
	for {
		select {
		case s := <-chStats:
			defaults := defaultTags.Load().(map[string]string)
			for _, r := range reporters {
				for k, v := range defaults {
					s.Tags[k] = v
				}
				if err := r.Submit(s); err != nil {
					log.Errorf("Failed to report error to influxdb: %s", err)
				} else {
					log.Tracef("Submitted error to influxdb: %v", s)
				}
			}
		case <-chStop:
			log.Debug("Measured loop stopped")
			atomic.StoreUint32(&running, 0)
			return
		}
	}
}

func reportError(addr string, err error, phase string) {
	splitted := strings.Split(err.Error(), ":")
	lastIndex := len(splitted) - 1
	if lastIndex < 0 {
		lastIndex = 0
	}
	e := strings.Trim(splitted[lastIndex], " ")
	select {
	case chStats <- &Stats{
		Type: "errors",
		Tags: map[string]string{
			"server": addr,
			"error":  e,
			"phase":  phase,
		},
		Fields: map[string]interface{}{"value": 1},
	}:
	default:
	}
}

type measuredConn struct {
	net.Conn
	addr string
}

// Read() implements the function from net.Conn
func (mc measuredConn) Read(b []byte) (n int, err error) {
	n, err = mc.Conn.Read(b)
	if err != nil {
		reportError(mc.addr, err, "read")
	}
	return
}

// Write() implements the function from net.Conn
func (mc measuredConn) Write(b []byte) (n int, err error) {
	n, err = mc.Conn.Write(b)
	if err != nil {
		reportError(mc.addr, err, "write")
	}
	return
}

// Close() implements the function from net.Conn
func (mc measuredConn) Close() (err error) {
	err = mc.Conn.Close()
	if err != nil {
		reportError(mc.addr, err, "close")
	}
	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	mailgun "github.com/mailgun/mailgun-go"
)

const (
	defaultCIDR = "0.0.0.0/0"

	arinAPIEndpoint = "http://whois.arin.net/rest/ip/%s"

	emailSender = "k8scan@jessfraz.com"
)

var (
	timeoutPing time.Duration
	timeoutGet  time.Duration

	cidr string

	ports = []int{80, 443, 9001, 8001}

	mailgunDomain  string
	mailgunAPIKey  string
	emailRecipient string

	debug bool
)

func init() {
	flag.DurationVar(&timeoutPing, "timeout-ping", 2*time.Second, "Timeout for checking that the port is open")
	flag.DurationVar(&timeoutGet, "timeout-get", 10*time.Second, "Timeout for getting the contents of the URL")

	flag.StringVar(&cidr, "cidr", defaultCIDR, "IP CIDR to scan")

	flag.StringVar(&mailgunAPIKey, "mailgun-api-key", "", "Mailgun API Key to use for sending email (optional)")
	flag.StringVar(&mailgunDomain, "mailgun-domain", "", "Mailgun Domain to use for sending email (optional)")
	flag.StringVar(&emailRecipient, "email-recipient", "", "Recipient for email notifications (optional)")

	flag.BoolVar(&debug, "d", false, "Run in debug mode")

	flag.Usage = func() {
		flag.PrintDefaults()
	}

	flag.Parse()

	// set log level
	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}
}

func main() {
	// On ^C, or SIGTERM handle exit.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	go func() {
		for sig := range c {
			logrus.Infof("Received %s, exiting.", sig.String())
			os.Exit(0)
		}
	}()

	// Set the logger to nil so we ignore messages from the Dial that don't matter.
	// See: https://github.com/golang/go/issues/19895#issuecomment-292793756
	log.SetFlags(0)
	log.SetOutput(ioutil.Discard)

	logrus.Infof("Scanning for Kubernetes Dashboards and API Servers on %s over port range %#v", cidr, ports)
	if len(mailgunDomain) > 0 && len(mailgunAPIKey) > 0 && len(emailRecipient) > 0 {
		logrus.Infof("Using Mailgun Domain %s, API Key %s to send emails to %s", mailgunDomain, mailgunAPIKey, emailRecipient)
	}
	logrus.Infof("This may take a bit...")

	startTime := time.Now()

	ip, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		logrus.Fatal(err)
	}

	var wg sync.WaitGroup
	for ip := ip.Mask(ipnet.Mask); ipnet.Contains(ip); inc(ip) {
		for _, port := range ports {
			wg.Add(1)
			go func(ip string, port int) {
				defer wg.Done()

				scanIP(ip, port)

			}(ip.String(), port)
		}
	}

	wg.Wait()

	since := time.Since(startTime)
	logrus.Infof("Scan took: %s", since.String())
}

func scanIP(ip string, port int) {
	// Check if the port is open.
	ok := portOpen(ip, port)
	if !ok {
		return
	}

	// Check if it's a kubernetes dashboard.
	ok, uri := isKubernetesDashboard(ip, port)
	if !ok {
		return
	}

	// Get the info for the ip address.
	info, err := getIPInfo(ip)
	if err != nil {
		logrus.Warnf("ip info err: %v", err)
	}

	fmt.Printf("%s:%d\t%s\t%s\t%s\n",
		ip, port,
		info.Net.Organization.Handle, info.Net.Organization.Name, info.Net.Organization.Reference)

	// send an email
	if len(mailgunDomain) > 0 && len(mailgunAPIKey) > 0 && len(emailRecipient) > 0 {
		if err := sendEmail(uri, ip, port, info); err != nil {
			logrus.Warn(err)
		}
	}
}

func portOpen(ip string, port int) bool {
	c, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), timeoutPing)
	if err != nil {
		// logrus.Warnf("listen at %s:%s failed: %v", ip, port, err)
		return false
	}
	defer c.Close()

	return true
}

func isKubernetesDashboard(ip string, port int) (bool, string) {
	client := &http.Client{
		Timeout: timeoutGet,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		},
	}

	tryAddrs := []string{
		fmt.Sprintf("http://%s:%d", ip, port),
		fmt.Sprintf("https://%s:%d", ip, port),
		fmt.Sprintf("http://%s:%d/api/", ip, port),
		fmt.Sprintf("https://%s:%d/api/", ip, port),
	}

	var (
		resp *http.Response
		err  = errors.New("not yet run")
		uri  string
	)

	for i := 0; i < len(tryAddrs) && err != nil; i++ {
		uri = tryAddrs[i]
		resp, err = client.Get(uri)
	}
	if err != nil {
		//logrus.Warnf("getting %s:%s failed: %v", ip, port, err)
		return false, ""
	}
	defer resp.Body.Close()

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return false, ""
	}

	body := strings.ToLower(string(b))
	if (strings.Contains(body, "kubernetes") && strings.Contains(body, "dashboard")) ||
		(strings.Contains(body, `"versions"`) && strings.Contains(body, `"serverAddress`)) ||
		(strings.Contains(body, `"paths"`) && strings.Contains(body, `"/api"`)) {
		return true, uri
	}

	return false, ""
}

type ARINResponse struct {
	Net NetJSON `json:"net,omitempty"`
}

type NetJSON struct {
	Organization OrganizationJSON `json:"orgRef,omitempty"`
}

type OrganizationJSON struct {
	Handle    string `json:"@handle,omitempty"`
	Name      string `json:"@name,omitempty"`
	Reference string `json:"$,omitempty"`
}

func getIPInfo(ip string) (b ARINResponse, err error) {
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf(arinAPIEndpoint, ip), nil)
	if err != nil {
		return b, err
	}
	req.Header.Set("Accept", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return b, err
	}
	defer resp.Body.Close()

	if err := json.NewDecoder(resp.Body).Decode(&b); err != nil {
		return b, err
	}

	return b, nil
}

func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] > 0 {
			break
		}
	}
}

func sendEmail(uri, ip string, port int, arinInfo ARINResponse) error {
	mailgunClient := mailgun.NewMailgun(mailgunDomain, mailgunAPIKey, "")

	msg, _, err := mailgunClient.Send(mailgunClient.NewMessage(
		/* From */ fmt.Sprintf("%s <%s>", emailSender, emailSender),
		/* Subject */ fmt.Sprintf("[k8scan]: found dashboard %s", uri),
		/* Body */ fmt.Sprintf(`Time: %s

IP: %s:%d
URL: %s

ARIN: %s
	  %s
	  %s
`,
			time.Now().Format(time.UnixDate),
			arinInfo.Net.Organization.Handle,
			arinInfo.Net.Organization.Name,
			arinInfo.Net.Organization.Reference,
		),
		/* To */ emailRecipient,
	))
	if err != nil {
		return fmt.Errorf("sending Mailgun message failed: response: %#v error: %v", msg, err)
	}

	return nil
}
<commit_msg>fix format<commit_after>package main

import (
	"crypto/tls"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"sync"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	mailgun "github.com/mailgun/mailgun-go"
)

const (
	defaultCIDR = "0.0.0.0/0"

	arinAPIEndpoint = "http://whois.arin.net/rest/ip/%s"

	emailSender = "k8scan@jessfraz.com"
)

var (
	timeoutPing time.Duration
	timeoutGet  time.Duration

	cidr string

	ports = []int{80, 443, 9001, 8001}

	mailgunDomain  string
	mailgunAPIKey  string
	emailRecipient string

	debug bool
)

func init() {
	flag.DurationVar(&timeoutPing, "timeout-ping", 2*time.Second, "Timeout for checking that the port is open")
	flag.DurationVar(&timeoutGet, "timeout-get", 10*time.Second, "Timeout for getting the contents of the URL")

	flag.StringVar(&cidr, "cidr", defaultCIDR, "IP CIDR to scan")

	flag.StringVar(&mailgunAPIKey, "mailgun-api-key", "", "Mailgun API Key to use for sending email (optional)")
	flag.StringVar(&mailgunDomain, "mailgun-domain", "", "Mailgun Domain to use for sending email (optional)")
	flag.StringVar(&emailRecipient, "email-recipient", "", "Recipient for email notifications (optional)")

	flag.BoolVar(&debug, "d", false, "Run in debug mode")

	flag.Usage = func() {
		flag.PrintDefaults()
	}

	flag.Parse()

	// set log level
	if debug {
		logrus.SetLevel(logrus.DebugLevel)
	}
}

func main() {
	// On ^C, or SIGTERM handle exit.
	c := make(chan os.Signal, 1)
	signal.Notify(c, os.Interrupt)
	signal.Notify(c, syscall.SIGTERM)
	go func() {
		for sig := range c {
			logrus.Infof("Received %s, exiting.", sig.String())
			os.Exit(0)
		}
	}()

	// Set the logger to nil so we ignore messages from the Dial that don't matter.
	// See: https://github.com/golang/go/issues/19895#issuecomment-292793756
	log.SetFlags(0)
	log.SetOutput(ioutil.Discard)

	logrus.Infof("Scanning for Kubernetes Dashboards and API Servers on %s over port range %#v", cidr, ports)
	if len(mailgunDomain) > 0 && len(mailgunAPIKey) > 0 && len(emailRecipient) > 0 {
		logrus.Infof("Using Mailgun Domain %s, API Key %s to send emails to %s", mailgunDomain, mailgunAPIKey, emailRecipient)
	}
	logrus.Infof("This may take a bit...")

	startTime := time.Now()

	ip, ipnet, err := net.ParseCIDR(cidr)
	if err != nil {
		logrus.Fatal(err)
	}

	var wg sync.WaitGroup
	for ip := ip.Mask(ipnet.Mask); ipnet.Contains(ip); inc(ip) {
		for _, port := range ports {
			wg.Add(1)
			go func(ip string, port int) {
				defer wg.Done()

				scanIP(ip, port)

			}(ip.String(), port)
		}
	}

	wg.Wait()

	since := time.Since(startTime)
	logrus.Infof("Scan took: %s", since.String())
}

func scanIP(ip string, port int) {
	// Check if the port is open.
	ok := portOpen(ip, port)
	if !ok {
		return
	}

	// Check if it's a kubernetes dashboard.
	ok, uri := isKubernetesDashboard(ip, port)
	if !ok {
		return
	}

	// Get the info for the ip address.
	info, err := getIPInfo(ip)
	if err != nil {
		logrus.Warnf("ip info err: %v", err)
	}

	fmt.Printf("%s:%d\t%s\t%s\t%s\n",
		ip, port,
		info.Net.Organization.Handle, info.Net.Organization.Name, info.Net.Organization.Reference)

	// send an email
	if len(mailgunDomain) > 0 && len(mailgunAPIKey) > 0 && len(emailRecipient) > 0 {
		if err := sendEmail(uri, ip, port, info); err != nil {
			logrus.Warn(err)
		}
	}
}

func portOpen(ip string, port int) bool {
	c, err := net.DialTimeout("tcp", fmt.Sprintf("%s:%d", ip, port), timeoutPing)
	if err != nil {
		// logrus.Warnf("listen at %s:%s failed: %v", ip, port, err)
		return false
	}
	defer c.Close()

	return true
}

func isKubernetesDashboard(ip string, port int) (bool, string) {
	client := &http.Client{
		Timeout: timeoutGet,
		Transport: &http.Transport{
			TLSClientConfig: &tls.Config{
				InsecureSkipVerify: true,
			},
		},
	}

	tryAddrs := []string{
		fmt.Sprintf("http://%s:%d", ip, port),
		fmt.Sprintf("https://%s:%d", ip, port),
		fmt.Sprintf("http://%s:%d/api/", ip, port),
		fmt.Sprintf("https://%s:%d/api/", ip, port),
	}

	var (
		resp *http.Response
		err  = errors.New("not yet run")
		uri  string
	)

	for i := 0; i < len(tryAddrs) && err != nil; i++ {
		uri = tryAddrs[i]
		resp, err = client.Get(uri)
	}
	if err != nil {
		//logrus.Warnf("getting %s:%s failed: %v", ip, port, err)
		return false, ""
	}
	defer resp.Body.Close()

	b, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return false, ""
	}

	body := strings.ToLower(string(b))
	if (strings.Contains(body, "kubernetes") && strings.Contains(body, "dashboard")) ||
		(strings.Contains(body, `"versions"`) && strings.Contains(body, `"serverAddress`)) ||
		(strings.Contains(body, `"paths"`) && strings.Contains(body, `"/api"`)) {
		return true, uri
	}

	return false, ""
}

type ARINResponse struct {
	Net NetJSON `json:"net,omitempty"`
}

type NetJSON struct {
	Organization OrganizationJSON `json:"orgRef,omitempty"`
}

type OrganizationJSON struct {
	Handle    string `json:"@handle,omitempty"`
	Name      string `json:"@name,omitempty"`
	Reference string `json:"$,omitempty"`
}

func getIPInfo(ip string) (b ARINResponse, err error) {
	req, err := http.NewRequest(http.MethodGet, fmt.Sprintf(arinAPIEndpoint, ip), nil)
	if err != nil {
		return b, err
	}
	req.Header.Set("Accept", "application/json")

	resp, err := http.DefaultClient.Do(req)
	if err != nil {
		return b, err
	}
	defer resp.Body.Close()

	if err := json.NewDecoder(resp.Body).Decode(&b); err != nil {
		return b, err
	}

	return b, nil
}

func inc(ip net.IP) {
	for j := len(ip) - 1; j >= 0; j-- {
		ip[j]++
		if ip[j] > 0 {
			break
		}
	}
}

func sendEmail(uri, ip string, port int, arinInfo ARINResponse) error {
	mailgunClient := mailgun.NewMailgun(mailgunDomain, mailgunAPIKey, "")

	msg, _, err := mailgunClient.Send(mailgunClient.NewMessage(
		/* From */ fmt.Sprintf("%s <%s>", emailSender, emailSender),
		/* Subject */ fmt.Sprintf("[k8scan]: found dashboard %s", uri),
		/* Body */ fmt.Sprintf(`Time: %s

IP: %s:%d
URL: %s

ARIN: %s
	  %s
	  %s
`,
			time.Now().Format(time.UnixDate),
			ip,
			port,
			uri,
			arinInfo.Net.Organization.Handle,
			arinInfo.Net.Organization.Name,
			arinInfo.Net.Organization.Reference,
		),
		/* To */ emailRecipient,
	))
	if err != nil {
		return fmt.Errorf("sending Mailgun message failed: response: %#v error: %v", msg, err)
	}

	return nil
}
<|endoftext|>


<commit_before>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 kafka

import (
	"strings"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core/ctypes"

	"gopkg.in/Shopify/sarama.v1"
)

const (
	PluginName    = "kafka"
	PluginVersion = 6
	PluginType    = plugin.PublisherPluginType
)

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(PluginName, PluginVersion, PluginType, []string{plugin.SnapGOBContentType}, []string{plugin.SnapGOBContentType})
}

type kafkaPublisher struct{}

func NewKafkaPublisher() *kafkaPublisher {
	return &kafkaPublisher{}
}

// Publish sends data to a Kafka server
func (k *kafkaPublisher) Publish(contentType string, content []byte, config map[string]ctypes.ConfigValue) error {
	topic := config["topic"].(ctypes.ConfigValueStr).Value
	brokers := parseBrokerString(config["brokers"].(ctypes.ConfigValueStr).Value)
	//
	err := k.publish(topic, brokers, content)
	return err
}

func (k *kafkaPublisher) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	config := cpolicy.NewPolicyNode()

	r1, err := cpolicy.NewStringRule("topic", true)
	handleErr(err)
	r1.Description = "Kafka topic for publishing"

	r2, _ := cpolicy.NewStringRule("brokers", true)
	handleErr(err)
	r2.Description = "List of brokers in the format: broker-ip:port;broker-ip:port (ex: 192.168.1.1:9092;172.16.9.99:9092"

	config.Add(r1, r2)
	cp.Add([]string{""}, config)
	return cp, nil
}

// Internal method after data has been converted to serialized bytes to send
func (k *kafkaPublisher) publish(topic string, brokers []string, content []byte) error {
	producer, err := sarama.NewSyncProducer(brokers, nil)
	if err != nil {
		return err
	}

	_, _, err = producer.SendMessage(&sarama.ProducerMessage{
		Topic: topic,
		Value: sarama.ByteEncoder(content),
	})
	return err
}

func parseBrokerString(brokerStr string) []string {
	return strings.Split(brokerStr, ";")
}

func handleErr(e error) {
	if e != nil {
		panic(e)
	}
}
<commit_msg>Solved issue #14: The plugin now sends metrics to Kafka in JSON type, instead of raw byte[] from processor.<commit_after>/*
http://www.apache.org/licenses/LICENSE-2.0.txt

Copyright 2015 Intel Corporation

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 kafka

import (
//	"bufio"
	"bytes"
	"encoding/gob"
	"encoding/json"
	"fmt"
//	"os"

	"strings"

	log "github.com/Sirupsen/logrus"

	"github.com/intelsdi-x/snap/control/plugin"
	"github.com/intelsdi-x/snap/control/plugin/cpolicy"
	"github.com/intelsdi-x/snap/core/ctypes"

	"gopkg.in/Shopify/sarama.v1"
)

const (
	PluginName    = "kafka"
	PluginVersion = 6
	PluginType    = plugin.PublisherPluginType
)

func Meta() *plugin.PluginMeta {
	return plugin.NewPluginMeta(PluginName, PluginVersion, PluginType, []string{plugin.SnapGOBContentType}, []string{plugin.SnapGOBContentType})
}

type kafkaPublisher struct{}

func NewKafkaPublisher() *kafkaPublisher {
	return &kafkaPublisher{}
}

// Publish sends data to a Kafka server
func (k *kafkaPublisher) Publish(contentType string, content []byte, config map[string]ctypes.ConfigValue) error {

	// Inserted and modified codes from intelsdi-x/snap-plugin-publisher-file/file/file.go

	logger := log.New()
	logger.Println("Publishing started")
	var metrics []plugin.MetricType

	switch contentType {
	case plugin.SnapGOBContentType:
		dec := gob.NewDecoder(bytes.NewBuffer(content))
		if err := dec.Decode(&metrics); err != nil {
			logger.Printf("Error decoding: error=%v content=%v", err, content)
			return err
		}
	default:
		logger.Printf("Error unknown content type '%v'", contentType)
		return fmt.Errorf("Unknown content type '%s'", contentType)
	}

	logger.Printf("publishing %v metrics to %v", len(metrics), config)

	jsonOut, err := json.Marshal(metrics)
	if err != nil {
		return fmt.Errorf("Error while marshalling metrics to JSON: %v", err)
	}

	// Inserted codes end

	topic := config["topic"].(ctypes.ConfigValueStr).Value
	brokers := parseBrokerString(config["brokers"].(ctypes.ConfigValueStr).Value)
	//
//	err := k.publish(topic, brokers, content)
	err = k.publish(topic, brokers, []byte(jsonOut))
	return err
}

func (k *kafkaPublisher) GetConfigPolicy() (*cpolicy.ConfigPolicy, error) {
	cp := cpolicy.New()
	config := cpolicy.NewPolicyNode()

	r1, err := cpolicy.NewStringRule("topic", true)
	handleErr(err)
	r1.Description = "Kafka topic for publishing"

	r2, _ := cpolicy.NewStringRule("brokers", true)
	handleErr(err)
	r2.Description = "List of brokers in the format: broker-ip:port;broker-ip:port (ex: 192.168.1.1:9092;172.16.9.99:9092"

	config.Add(r1, r2)
	cp.Add([]string{""}, config)
	return cp, nil
}

// Internal method after data has been converted to serialized bytes to send
func (k *kafkaPublisher) publish(topic string, brokers []string, content []byte) error {
	producer, err := sarama.NewSyncProducer(brokers, nil)
	if err != nil {
		return err
	}

	_, _, err = producer.SendMessage(&sarama.ProducerMessage{
		Topic: topic,
		Value: sarama.ByteEncoder(content),
	})
	return err
}

func parseBrokerString(brokerStr string) []string {
	return strings.Split(brokerStr, ";")
}

func handleErr(e error) {
	if e != nil {
		panic(e)
	}
}
<|endoftext|>


<commit_before>package db

import "code.cloudfoundry.org/lager"

var CheckDeleteBatchSize = 500

//counterfeiter:generate . CheckLifecycle
type CheckLifecycle interface {
	DeleteCompletedChecks(logger lager.Logger) error
}

type checkLifecycle struct {
	conn Conn
}

func NewCheckLifecycle(conn Conn) CheckLifecycle {
	return &checkLifecycle{
		conn: conn,
	}
}

func (cl *checkLifecycle) DeleteCompletedChecks(logger lager.Logger) error {
	var counter int
	var err1 error
	for {
		var numChecksDeleted int
		err1 = cl.conn.QueryRow(`
      WITH resource_builds AS (
        SELECT distinct(last_check_build_id) as build_id
        FROM resource_config_scopes
        WHERE last_check_build_id IS NOT NULL
      ),
      deleted_builds AS (
        DELETE FROM builds USING (
          (SELECT id
          FROM builds b
          WHERE completed AND resource_id IS NOT NULL
          AND NOT EXISTS ( SELECT 1 FROM resource_builds WHERE build_id = b.id )
					LIMIT $1)
            UNION ALL
          SELECT id
          FROM builds b
          WHERE completed AND resource_type_id IS NOT NULL
          AND EXISTS (SELECT * FROM builds b2 WHERE b.resource_type_id = b2.resource_type_id AND b.id < b2.id)
    ) AS deletable_builds WHERE builds.id = deletable_builds.id
      RETURNING builds.id
      ), deleted_events AS (
        DELETE FROM check_build_events USING deleted_builds WHERE build_id = deleted_builds.id
      )
      SELECT COUNT(*) FROM deleted_builds
    `, CheckDeleteBatchSize).Scan(&numChecksDeleted)
		if err1 != nil {
			break
		}
		logger.Debug("deleted-check-builds", lager.Data{"count": numChecksDeleted, "batch": counter})

		if numChecksDeleted < CheckDeleteBatchSize {
			break
		}
		counter++
	}

	_, err2 := cl.conn.Exec(`
      WITH resource_builds AS (
        SELECT distinct(last_check_build_id) as build_id
        FROM resource_config_scopes
        WHERE last_check_build_id IS NOT NULL
		  UNION ALL
        SELECT distinct(in_memory_build_id) as build_id
        FROM resources
        WHERE in_memory_build_id IS NOT NULL
      )
      DELETE FROM check_build_events WHERE build_id NOT IN (SELECT build_id FROM resource_builds)
    `)

	if err1 != nil {
		return err1
	}

	if err2 != nil {
		return err2
	}

	return nil
}
<commit_msg>fix a bad SQL for check gc.<commit_after>package db

import "code.cloudfoundry.org/lager"

var CheckDeleteBatchSize = 500

//counterfeiter:generate . CheckLifecycle
type CheckLifecycle interface {
	DeleteCompletedChecks(logger lager.Logger) error
}

type checkLifecycle struct {
	conn Conn
}

func NewCheckLifecycle(conn Conn) CheckLifecycle {
	return &checkLifecycle{
		conn: conn,
	}
}

func (cl *checkLifecycle) DeleteCompletedChecks(logger lager.Logger) error {
	var counter int
	var err1 error
	for {
		var numChecksDeleted int
		err1 = cl.conn.QueryRow(`
      WITH resource_builds AS (
        SELECT distinct(last_check_build_id) as build_id
        FROM resource_config_scopes
        WHERE last_check_build_id IS NOT NULL
      ),
      deleted_builds AS (
        DELETE FROM builds USING (
          (SELECT id
          FROM builds b
          WHERE completed AND resource_id IS NOT NULL
          AND NOT EXISTS ( SELECT 1 FROM resource_builds WHERE build_id = b.id )
					LIMIT $1)
            UNION ALL
          SELECT id
          FROM builds b
          WHERE completed AND resource_type_id IS NOT NULL
          AND EXISTS (SELECT * FROM builds b2 WHERE b.resource_type_id = b2.resource_type_id AND b.id < b2.id)
    ) AS deletable_builds WHERE builds.id = deletable_builds.id
      RETURNING builds.id
      ), deleted_events AS (
        DELETE FROM check_build_events USING deleted_builds WHERE build_id = deleted_builds.id
      )
      SELECT COUNT(*) FROM deleted_builds
    `, CheckDeleteBatchSize).Scan(&numChecksDeleted)
		if err1 != nil {
			break
		}
		logger.Debug("deleted-check-builds", lager.Data{"count": numChecksDeleted, "batch": counter})

		if numChecksDeleted < CheckDeleteBatchSize {
			break
		}
		counter++
	}

	_, err2 := cl.conn.Exec(`
      WITH expired_imb_ids AS (
	    SELECT distinct(build_id) AS build_id FROM check_build_events cbe
		WHERE NOT EXISTS (SELECT 1 FROM resources WHERE in_memory_build_id = cbe.build_id) AND
		      NOT EXISTS (SELECT 1 FROM builds WHERE id = cbe.build_id)
      )
      DELETE FROM check_build_events cbe2 USING expired_imb_ids WHERE cbe2.build_id = expired_imb_ids.build_id;
    `)

	if err1 != nil {
		return err1
	}

	if err2 != nil {
		return err2
	}

	return nil
}
<|endoftext|>


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

// +build go1.8

package spanner

import (
	"fmt"

	ocgrpc "go.opencensus.io/plugin/grpc"
	"go.opencensus.io/trace"
	"golang.org/x/net/context"
	"google.golang.org/api/option"
	"google.golang.org/grpc"
)

func openCensusOptions() []option.ClientOption {
	return []option.ClientOption{
		option.WithGRPCDialOption(grpc.WithStatsHandler(ocgrpc.NewClientStatsHandler())),
	}
}

func traceStartSpan(ctx context.Context, name string) context.Context {
	return trace.StartSpan(ctx, name)
}

func traceEndSpan(ctx context.Context, err error) {
	span := trace.FromContext(ctx)
	if err != nil {
		span.SetStatus(trace.Status{Message: err.Error()})
	}
	trace.EndSpan(ctx)
}

func tracePrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
	var attrs []trace.Attribute
	for k, v := range attrMap {
		var a trace.Attribute
		switch v := v.(type) {
		case string:
			a = trace.StringAttribute{k, v}
		case bool:
			a = trace.BoolAttribute{k, v}
		case int:
			a = trace.Int64Attribute{k, int64(v)}
		case int64:
			a = trace.Int64Attribute{k, v}
		default:
			a = trace.StringAttribute{k, fmt.Sprintf("%#v", v)}
		}
		attrs = append(attrs, a)
	}
	trace.FromContext(ctx).Annotatef(attrs, format, args...)
}
<commit_msg>Switch to the new OpenCensus trace API<commit_after>// Copyright 2017 Google Inc. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//      http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// +build go1.8

package spanner

import (
	"fmt"

	ocgrpc "go.opencensus.io/plugin/grpc"
	"go.opencensus.io/trace"
	"golang.org/x/net/context"
	"google.golang.org/api/option"
	"google.golang.org/grpc"
)

func openCensusOptions() []option.ClientOption {
	return []option.ClientOption{
		option.WithGRPCDialOption(grpc.WithStatsHandler(ocgrpc.NewClientStatsHandler())),
	}
}

func traceStartSpan(ctx context.Context, name string) context.Context {
	ctx, _ = trace.StartSpan(ctx, name)
	return ctx
}

func traceEndSpan(ctx context.Context, err error) {
	span := trace.FromContext(ctx)
	if err != nil {
		// TODO(jba): Add error code to the status.
		span.SetStatus(trace.Status{Message: err.Error()})
	}
	span.End()
}

func tracePrintf(ctx context.Context, attrMap map[string]interface{}, format string, args ...interface{}) {
	var attrs []trace.Attribute
	for k, v := range attrMap {
		var a trace.Attribute
		switch v := v.(type) {
		case string:
			a = trace.StringAttribute{k, v}
		case bool:
			a = trace.BoolAttribute{k, v}
		case int:
			a = trace.Int64Attribute{k, int64(v)}
		case int64:
			a = trace.Int64Attribute{k, v}
		default:
			a = trace.StringAttribute{k, fmt.Sprintf("%#v", v)}
		}
		attrs = append(attrs, a)
	}
	trace.FromContext(ctx).Annotatef(attrs, format, args...)
}
<|endoftext|>