<commit_before>/* vim: set sw=4 sts=4 et foldmethod=syntax : */

/*
 * Copyright (c) 2011 Alexander Færøy <ahf@0x90.dk>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice, this
 *   list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package main

import (
    "bufio"
    "crypto/md5"
    "crypto/sha1"
    "crypto/sha256"
    "crypto/sha512"
    "fmt"
    "hash"
    "os"
)

func F(h hash.Hash, name, str string) {
    h.Write([]byte(str))
    fmt.Printf("%s\t: '%x'\n", name, h.Sum())
}

func main() {
    fmt.Printf("Password: ");
    stdin := bufio.NewReader(os.Stdin);
    input, error := stdin.ReadString('\n');

    if error != nil {
        os.Exit(1)
    }

    F(md5.New(), "MD5", input)
    F(sha1.New(), "SHA1", input)
    F(sha256.New(), "SHA256", input)
    F(sha512.New(), "SHA512", input)
}
<commit_msg>Trim newlines of input string.<commit_after>/* vim: set sw=4 sts=4 et foldmethod=syntax : */

/*
 * Copyright (c) 2011 Alexander Færøy <ahf@0x90.dk>
 * All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions are met:
 *
 * * Redistributions of source code must retain the above copyright notice, this
 *   list of conditions and the following disclaimer.
 *
 * * Redistributions in binary form must reproduce the above copyright notice,
 *   this list of conditions and the following disclaimer in the documentation
 *   and/or other materials provided with the distribution.
 *
 * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
 * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
 * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
 * DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
 * FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
 * DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
 * CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
 * OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
 * OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package main

import (
    "bufio"
    "crypto/md5"
    "crypto/sha1"
    "crypto/sha256"
    "crypto/sha512"
    "fmt"
    "hash"
    "os"
    "strings"
)

func F(h hash.Hash, name, str string) {
    h.Write([]byte(strings.Trim(str, "\n")))
    fmt.Printf("%s\t: '%x'\n", name, h.Sum())
}

func main() {
    fmt.Printf("Password: ");
    stdin := bufio.NewReader(os.Stdin);
    input, error := stdin.ReadString('\n');

    if error != nil {
        os.Exit(1)
    }

    F(md5.New(), "MD5", input)
    F(sha1.New(), "SHA1", input)
    F(sha256.New(), "SHA256", input)
    F(sha512.New(), "SHA512", input)
}
<|endoftext|>


<commit_before>package utils

import (
	"bytes"
	"fmt"
	"log"
	"os"
)

var (
	loggingBuffer bytes.Buffer
	logger        *log.Logger
	logFileObject *os.File
)

// Creates a new logger which uses a buffer where he collects the messages.
func InitializeLogger(logFile, preface string) {
	if logFile == "path/to/logfile.txt" || logFile == "" {
		logger = log.New(os.Stdout, preface, log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile)
	} else {
		var err error
		if exists := CheckIfFileExists(logFile); exists {
			log.Println(fmt.Sprintf("The file %s already exists, append the logging output.", logFile))
			logFileObject, err = os.OpenFile(logFile, os.O_WRONLY|os.O_APPEND, 0777)
		} else {
			logFileObject, err = os.Create(logFile)
			logFileObject.Chmod(0660)
		}
		if err != nil {
			log.Fatalln(err.Error())
		}
		logFileObject.WriteString("Begin logging...\n")
		logger = log.New(&loggingBuffer, preface, log.Ldate|log.Ltime|log.Lmicroseconds|log.Lshortfile)
	}
}

// This function is used as a wrapper for the logging functinality.
// If the logger is not initialized, it calls the log methods from the
// log package.
func PrintMessage(message interface{}) {
	if logger == nil {
		switch typeValue := message.(type) {
		case fmt.Stringer:
			log.Println(typeValue.String())
		default:
			log.Println(typeValue)
		}
	} else {
		switch typeValue := message.(type) {
		case fmt.Stringer:
			logger.Println(typeValue.String())
		default:
			logger.Println(typeValue)
		}
		printAndClearLoggerContent()
	}
}

// This function prints the content of the logging buffer to the STDOUT.
// TODO: Decide if it should be printed to console or wrote to a file or simmilar.
func printAndClearLoggerContent() {
	if loggingBuffer.Len() != 0 {
		if logFileObject == nil {
			loggingBuffer.WriteTo(os.Stdout)
			//fmt.Println(&loggingBuffer)
			loggingBuffer.Reset()
		} else {
			loggingBuffer.WriteTo(logFileObject)
		}
	}
}
<commit_msg>Removed the Lshortfile option because its everytime the same when using my PrintMessage wrapper.<commit_after>package utils

import (
	"bytes"
	"fmt"
	"log"
	"os"
)

var (
	loggingBuffer bytes.Buffer
	logger        *log.Logger
	logFileObject *os.File
)

// Creates a new logger which uses a buffer where he collects the messages.
func InitializeLogger(logFile, preface string) {
	if logFile == "path/to/logfile.txt" || logFile == "" {
		logger = log.New(os.Stdout, preface, log.Ldate|log.Ltime|log.Lmicroseconds)
	} else {
		var err error
		if exists := CheckIfFileExists(logFile); exists {
			log.Println(fmt.Sprintf("The file %s already exists, append the logging output.", logFile))
			logFileObject, err = os.OpenFile(logFile, os.O_WRONLY|os.O_APPEND, 0777)
		} else {
			logFileObject, err = os.Create(logFile)
			logFileObject.Chmod(0660)
		}
		if err != nil {
			log.Fatalln(err.Error())
		}
		logFileObject.WriteString("Begin logging...\n")
		logger = log.New(&loggingBuffer, preface, log.Ldate|log.Ltime|log.Lmicroseconds)
	}
}

// This function is used as a wrapper for the logging functinality.
// If the logger is not initialized, it calls the log methods from the
// log package.
func PrintMessage(message interface{}) {
	if logger == nil {
		switch typeValue := message.(type) {
		case fmt.Stringer:
			log.Println(typeValue.String())
		default:
			log.Println(typeValue)
		}
	} else {
		switch typeValue := message.(type) {
		case fmt.Stringer:
			logger.Println(typeValue.String())
		default:
			logger.Println(typeValue)
		}
		printAndClearLoggerContent()
	}
}

// This function prints the content of the logging buffer to the STDOUT.
// TODO: Decide if it should be printed to console or wrote to a file or simmilar.
func printAndClearLoggerContent() {
	if loggingBuffer.Len() != 0 {
		if logFileObject == nil {
			loggingBuffer.WriteTo(os.Stdout)
			//fmt.Println(&loggingBuffer)
			loggingBuffer.Reset()
		} else {
			loggingBuffer.WriteTo(logFileObject)
		}
	}
}
<|endoftext|>


<commit_before>package cmds

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"

	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/spf13/cobra"
)

const badJSON1 = `
{
"356weryt

}
`

const badJSON2 = `
{
    "a": 1,
    "b": [23,4,4,64,56,36,7456,7],
    "c": {a,f,g,h,j,j},
    "d": 3452.36456,
}
`

func testJSONSyntaxErrorsReported(inputFile string, inputFileValue string, inputCommand []string) {
	address := "0.0.0.0:30650"

	rootCmd := &cobra.Command{
		Use: os.Args[0],
		Long: `Access the Pachyderm API.

Envronment variables:
  ADDRESS=0.0.0.0:30650, the server to connect to.
`,
	}

	cmds, _ := Cmds(address)
	for _, cmd := range cmds {
		rootCmd.AddCommand(cmd)
	}

	ioutil.WriteFile(inputFile, []byte(inputFileValue), 0644)
	os.Args = inputCommand
	rootCmd.Execute()
}

func testBadJSON(t *testing.T, testName string, inputFile string, inputFileValue string, inputCommand []string, expectedOutput string) {

	if os.Getenv("BE_CRASHER") == "1" {
		testJSONSyntaxErrorsReported(inputFile, inputFileValue, inputCommand)
		return
	}

	cmd := exec.Command(os.Args[0], fmt.Sprintf("-test.run=%v", testName))
	cmd.Env = append(os.Environ(), "BE_CRASHER=1")
	out, err := cmd.CombinedOutput()

	// Do our cleanup here, since we have an exit 1 in the actual run:
	wd, _ := os.Getwd()
	fileName := filepath.Join(wd, inputFile)
	os.Remove(fileName)

	require.YesError(t, err)
	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
		require.Equal(t, expectedOutput, string(out))
		return
	}
	t.Fatalf("process ran with err %v, want exit status 1", err)

}

func TestJSONSyntaxErrorsReportedCreateJob(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 3:

"356weryt

         ^
invalid character '\n' in string literal
`
	cmd := []string{"pachctl", "create-job", "-f", "bad1.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreateJob", "bad1.json", badJSON1, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedCreateJob2(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "create-job", "-f", "bad2.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreateJob2", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedCreatePipeline(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "create-pipeline", "-f", "bad2.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreatePipeline", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedRunPipeline(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "run-pipeline", "-f", "bad2.json", "somePipelineName"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedRunPipeline", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedCreatePipelineFromStdin(t *testing.T) {
	descriptiveOutput := `Reading from stdin.
Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	rawCmd := []string{"pachctl", "create-pipeline"}
	testName := "TestJSONSyntaxErrorsReportedCreatePipelineFromStdin"

	if os.Getenv("BE_CRASHER") == "1" {
		address := "0.0.0.0:30650"
		rootCmd := &cobra.Command{
			Use: os.Args[0],
			Long: `Access the Pachyderm API.

Envronment variables:
  ADDRESS=0.0.0.0:30650, the server to connect to.
`,
		}

		cmds, _ := Cmds(address)
		for _, cmd := range cmds {
			rootCmd.AddCommand(cmd)
		}

		os.Args = rawCmd
		ioutil.WriteFile("bad2.json", []byte(badJSON2), 0644)
		os.Stdin, _ = os.Open("bad2.json")
		rootCmd.Execute()
		return
	}

	cmd := exec.Command(os.Args[0], fmt.Sprintf("-test.run=%v", testName))
	cmd.Env = append(os.Environ(), "BE_CRASHER=1")
	out, err := cmd.CombinedOutput()

	// Put our cleanup here, since we have an exit 1 in the actual run:

	wd, _ := os.Getwd()
	os.Remove(filepath.Join(wd, "bad2.json"))

	require.YesError(t, err)
	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
		require.Equal(t, descriptiveOutput, string(out))
		return
	}
	t.Fatalf("process ran with err %v, want exit status 1", err)

}

func TestJSONSyntaxErrorsReportedUpdatePipeline(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "update-pipeline", "-f", "bad2.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreatePipeline", "bad2.json", badJSON2, cmd, descriptiveOutput)
}
<commit_msg>Adds a test for --push-images.<commit_after>package cmds

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path/filepath"
	"testing"

	"golang.org/x/net/context"

	"github.com/fsouza/go-dockerclient"
	"github.com/pachyderm/pachyderm/src/client/pkg/require"
	"github.com/spf13/cobra"
)

const badJSON1 = `
{
"356weryt

}
`

const badJSON2 = `
{
    "a": 1,
    "b": [23,4,4,64,56,36,7456,7],
    "c": {a,f,g,h,j,j},
    "d": 3452.36456,
}
`

func rootCmd() *cobra.Command {
	rootCmd := &cobra.Command{}
	cmds, _ := Cmds("0.0.0.0:30650")
	for _, cmd := range cmds {
		rootCmd.AddCommand(cmd)
	}
	return rootCmd
}

func testJSONSyntaxErrorsReported(inputFile string, inputFileValue string, inputCommand []string) {
	ioutil.WriteFile(inputFile, []byte(inputFileValue), 0644)
	os.Args = inputCommand
	rootCmd().Execute()
}

func testBadJSON(t *testing.T, testName string, inputFile string, inputFileValue string, inputCommand []string, expectedOutput string) {

	if os.Getenv("BE_CRASHER") == "1" {
		testJSONSyntaxErrorsReported(inputFile, inputFileValue, inputCommand)
		return
	}

	cmd := exec.Command(os.Args[0], fmt.Sprintf("-test.run=%v", testName))
	cmd.Env = append(os.Environ(), "BE_CRASHER=1")
	out, err := cmd.CombinedOutput()

	// Do our cleanup here, since we have an exit 1 in the actual run:
	wd, _ := os.Getwd()
	fileName := filepath.Join(wd, inputFile)
	os.Remove(fileName)

	require.YesError(t, err)
	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
		require.Equal(t, expectedOutput, string(out))
		return
	}
	t.Fatalf("process ran with err %v, want exit status 1", err)

}

func TestJSONSyntaxErrorsReportedCreateJob(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 3:

"356weryt

         ^
invalid character '\n' in string literal
`
	cmd := []string{"pachctl", "create-job", "-f", "bad1.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreateJob", "bad1.json", badJSON1, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedCreateJob2(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "create-job", "-f", "bad2.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreateJob2", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedCreatePipeline(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "create-pipeline", "-f", "bad2.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreatePipeline", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedRunPipeline(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "run-pipeline", "-f", "bad2.json", "somePipelineName"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedRunPipeline", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestJSONSyntaxErrorsReportedCreatePipelineFromStdin(t *testing.T) {
	descriptiveOutput := `Reading from stdin.
Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	rawCmd := []string{"pachctl", "create-pipeline"}
	testName := "TestJSONSyntaxErrorsReportedCreatePipelineFromStdin"

	if os.Getenv("BE_CRASHER") == "1" {
		os.Args = rawCmd
		ioutil.WriteFile("bad2.json", []byte(badJSON2), 0644)
		os.Stdin, _ = os.Open("bad2.json")
		rootCmd().Execute()
		return
	}

	cmd := exec.Command(os.Args[0], fmt.Sprintf("-test.run=%v", testName))
	cmd.Env = append(os.Environ(), "BE_CRASHER=1")
	out, err := cmd.CombinedOutput()

	// Put our cleanup here, since we have an exit 1 in the actual run:

	wd, _ := os.Getwd()
	os.Remove(filepath.Join(wd, "bad2.json"))

	require.YesError(t, err)
	if e, ok := err.(*exec.ExitError); ok && !e.Success() {
		require.Equal(t, descriptiveOutput, string(out))
		return
	}
	t.Fatalf("process ran with err %v, want exit status 1", err)

}

func TestJSONSyntaxErrorsReportedUpdatePipeline(t *testing.T) {
	descriptiveOutput := `Syntax Error on line 5:

    "c": {a
          ^
invalid character 'a' looking for beginning of object key string
`
	cmd := []string{"pachctl", "update-pipeline", "-f", "bad2.json"}
	testBadJSON(t, "TestJSONSyntaxErrorsReportedCreatePipeline", "bad2.json", badJSON2, cmd, descriptiveOutput)
}

func TestPushImages(t *testing.T) {
	client, err := docker.NewClientFromEnv()
	require.NoError(t, err)
	require.NoError(t, client.TagImage("pachyderm/job-shim", docker.TagImageOptions{
		Repo:    "test-job-shim",
		Context: context.Background(),
	}))
	ioutil.WriteFile("test-push-images.json", []byte(`{
  "pipeline": {
    "name": "test_push_images"
  },
  "transform": {
    "cmd": [ "true" ],
	"image": "test-job-shim"
  }
}`), 0644)
	os.Args = []string{"pachctl", "create-pipeline", "--push-images", "-f", "test-push-images.json"}
	require.NoError(t, rootCmd().Execute())
}
<|endoftext|>


<commit_before>package luasrc

const templateText_luasrc = `-- Generated by github.com/davyxu/tabtoy
-- Version: {{.Version}}

return {
	init = function( g )
		{{range $di, $tab := .Datas.AllTables}}
		g.{{$tab.HeaderType}} = { {{range $unusedrow,$row := $tab.DataRowIndex}}{{$headers := $.Types.AllFieldByName $tab.OriginalHeaderType}}
			{ {{range $col, $header := $headers}}{{if IsWrapFieldName $ $tab $headers $row $col}}{{$header.FieldName}} = {{WrapTabValue $ $tab $headers $row $col}}, {{end}}{{end}}}, {{end}}
		}
		{{end}}
		{{range $ii, $idx := GetIndices $}}
		-- {{$idx.Table.HeaderType}}
		g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}} = {}
		for _, rec in pairs(g.{{$idx.Table.HeaderType}}) do
			g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}}[rec.{{$idx.FieldInfo.FieldName}}] = rec
		end
		{{end}}
		{{range $sn, $objName := $.Types.EnumNames}}
		---@enum {{$.PackageName}}.{{$objName}}
		g.{{$objName}} =
		{ {{range $fi,$field := $.Types.AllFieldByName $objName}}
			{{$field.FieldName}} = {{$field.Value}}, -- {{$field.Name}} {{end}} {{range $fi,$field := $.Types.AllFieldByName $objName}}
			[{{$field.Value}}] = "{{$field.FieldName}}", -- {{$field.Name}} {{end}}
		}{{end}}
		return g
	end
}
`

const templateText_luadir = `-- Generated by github.com/davyxu/tabtoy
-- Version: {{$.G.Version}}

return {
	init = function( g )
		g.{{$.Tab.HeaderType}} = { {{range $unusedrow,$row := $.Tab.DataRowIndex}}{{$headers := $.G.Types.AllFieldByName $.Tab.OriginalHeaderType}}
			{ {{range $col, $header := $headers}}{{if IsWrapFieldName $.G $.Tab $headers $row $col}}{{$header.FieldName}} = {{WrapTabValue $.G $.Tab $headers $row $col}}, {{end}}{{end}}}, {{end}}
		}
		{{range $ii, $idx := GetIndicesByTable $.Tab}}
		-- {{$idx.Table.HeaderType}}
		g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}} = {}
		for _, rec in pairs(g.{{$idx.Table.HeaderType}}) do
			g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}}[rec.{{$idx.FieldInfo.FieldName}}] = rec
		end
		{{end}}
		return g
	end
}
`
const templateText_luatype = `-- Generated by github.com/davyxu/tabtoy
-- Version: {{.Version}}

return {
	init = function( g )
		{{range $sn, $objName := $.Types.EnumNames}}
		---@enum {{$.PackageName}}.{{$objName}}
		g.{{$objName}} =
		{ {{range $fi,$field := $.Types.AllFieldByName $objName}}
			{{$field.FieldName}} = {{$field.Value}}, -- {{$field.Name}} {{end}} {{range $fi,$field := $.Types.AllFieldByName $objName}}
			[{{$field.Value}}] = "{{$field.FieldName}}", -- {{$field.Name}} {{end}}
		}{{end}}
		return g
	end
}
`
<commit_msg>修改lua生成格式，换行vscode会识别错误格式<commit_after>package luasrc

const templateText_luasrc = `-- Generated by github.com/davyxu/tabtoy
-- Version: {{.Version}}

return {
	init = function( g )
		{{range $di, $tab := .Datas.AllTables}}
		g.{{$tab.HeaderType}} = { {{range $unusedrow,$row := $tab.DataRowIndex}}{{$headers := $.Types.AllFieldByName $tab.OriginalHeaderType}}
			{ {{range $col, $header := $headers}}{{if IsWrapFieldName $ $tab $headers $row $col}}{{$header.FieldName}} = {{WrapTabValue $ $tab $headers $row $col}}, {{end}}{{end}}}, {{end}}
		}
		{{end}}
		{{range $ii, $idx := GetIndices $}}
		-- {{$idx.Table.HeaderType}}
		g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}} = {}
		for _, rec in pairs(g.{{$idx.Table.HeaderType}}) do
			g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}}[rec.{{$idx.FieldInfo.FieldName}}] = rec
		end
		{{end}}
		{{range $sn, $objName := $.Types.EnumNames}}
		---@enum {{$.PackageName}}.{{$objName}}
		g.{{$objName}} = { {{range $fi,$field := $.Types.AllFieldByName $objName}}
			{{$field.FieldName}} = {{$field.Value}}, -- {{$field.Name}} {{end}} {{range $fi,$field := $.Types.AllFieldByName $objName}}
			[{{$field.Value}}] = "{{$field.FieldName}}", -- {{$field.Name}} {{end}}
		}{{end}}
		return g
	end
}
`

const templateText_luadir = `-- Generated by github.com/davyxu/tabtoy
-- Version: {{$.G.Version}}

return {
	init = function( g )
		g.{{$.Tab.HeaderType}} = { {{range $unusedrow,$row := $.Tab.DataRowIndex}}{{$headers := $.G.Types.AllFieldByName $.Tab.OriginalHeaderType}}
			{ {{range $col, $header := $headers}}{{if IsWrapFieldName $.G $.Tab $headers $row $col}}{{$header.FieldName}} = {{WrapTabValue $.G $.Tab $headers $row $col}}, {{end}}{{end}}}, {{end}}
		}
		{{range $ii, $idx := GetIndicesByTable $.Tab}}
		-- {{$idx.Table.HeaderType}}
		g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}} = {}
		for _, rec in pairs(g.{{$idx.Table.HeaderType}}) do
			g.{{$idx.Table.HeaderType}}By{{$idx.FieldInfo.FieldName}}[rec.{{$idx.FieldInfo.FieldName}}] = rec
		end
		{{end}}
		return g
	end
}
`
const templateText_luatype = `-- Generated by github.com/davyxu/tabtoy
-- Version: {{.Version}}

return {
	init = function( g )
		{{range $sn, $objName := $.Types.EnumNames}}
		---@enum {{$.PackageName}}.{{$objName}}
		g.{{$objName}} = { {{range $fi,$field := $.Types.AllFieldByName $objName}}
			{{$field.FieldName}} = {{$field.Value}}, -- {{$field.Name}} {{end}} {{range $fi,$field := $.Types.AllFieldByName $objName}}
			[{{$field.Value}}] = "{{$field.FieldName}}", -- {{$field.Name}} {{end}}
		}{{end}}
		return g
	end
}
`
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"hash/crc32"
	"sync"
	"time"
)

// Receiver, should only be used for one file
type BinaryTransportFileReceiver struct {
	lastChunkReceived  time.Time
	mux                sync.RWMutex
	fileMeta           *FileMeta
	data               *bytes.Buffer
	targetShardId      []byte
	chunkCount         uint32
	receivedDataChunks map[uint32][]byte // Data
}

// Receive chunk, returns true if done
func (this *BinaryTransportFileReceiver) Add(buf *bytes.Reader) bool {
	// Update last chunk received
	this.mux.Lock()
	this.lastChunkReceived = time.Now()
	this.mux.Unlock()

	// Read chunk number
	var err error
	var chunkNumber uint32
	err = binary.Read(buf, binary.BigEndian, &chunkNumber)
	panicErr(err)

	// Dedup based on chunk number
	this.mux.Lock()
	if this.receivedDataChunks[chunkNumber] != nil {
		this.mux.Unlock()
		log.Debug("Ignoring data file transfer chunk, duplicate message")
		return false
	}
	this.receivedDataChunks[chunkNumber] = make([]byte, 0)
	this.mux.Unlock()

	// Handle this message
	if chunkNumber == 0 {
		this._readFirstChunk(chunkNumber, buf)
	} else {
		this._readContentChunk(chunkNumber, buf)
	}

	// Done?
	this.mux.RLock()
	defer this.mux.RUnlock()
	return uint32(len(this.receivedDataChunks)) == this.chunkCount
}

// Get last chunk received time
func (this *BinaryTransportFileReceiver) GetLastChunkReceived() time.Time {
	this.mux.RLock()
	defer this.mux.RUnlock()
	return this.lastChunkReceived
}

// Read first chunk
func (this *BinaryTransportFileReceiver) _readFirstChunk(chunkNumber uint32, buf *bytes.Reader) {
	var err error

	// Meta already here? Deduplicate payloads (chunk number could have been 0 by malformed byte)
	this.mux.RLock()
	if this.fileMeta != nil {
		this.mux.RUnlock()
		log.Warn("Ignoring metadata file transfer chunk, duplicate message")
		return
	}
	this.mux.RUnlock()

	// Read target shard flag
	targetShardSet, targetShardSetErr := buf.ReadByte()
	panicErr(targetShardSetErr)

	// Target shard id
	targetShardIdBytes := make([]byte, 16)
	targetShardIdBytesRead, _ := buf.Read(targetShardIdBytes)
	if uint32(targetShardIdBytesRead) != 16 {
		panic("Target shard id bytes read mismatch")
	}

	// Set target shard
	if targetShardSet == 1 {
		this.mux.Lock()
		this.targetShardId = targetShardIdBytes
		this.mux.Unlock()
	}

	// Read meta length
	var metaLen uint32
	err = binary.Read(buf, binary.BigEndian, &metaLen)
	panicErr(err)

	// Read meta
	metaBytes := make([]byte, metaLen)
	metaBytesRead, _ := buf.Read(metaBytes)
	if uint32(metaBytesRead) != metaLen {
		panic("Meta bytes read mismatch")
	}
	meta := &FileMeta{}
	meta.FromBytes(metaBytes)

	// Read chunk count
	var chunkCount uint32
	err = binary.Read(buf, binary.BigEndian, &chunkCount)
	panicErr(err)

	// Set meta into receiver
	this.mux.Lock()
	this.fileMeta = meta
	this.chunkCount = chunkCount
	this.mux.Unlock()

	// Read actual content
	this._readContent(chunkNumber, buf)
}

// Read content chunk
func (this *BinaryTransportFileReceiver) _readContentChunk(chunkNumber uint32, buf *bytes.Reader) {
	// Read content
	this._readContent(chunkNumber, buf)
}

// Read content
func (this *BinaryTransportFileReceiver) _readContent(chunkNumber uint32, buf *bytes.Reader) {
	var err error

	// Read content length
	var contentLen uint32
	err = binary.Read(buf, binary.BigEndian, &contentLen)
	panicErr(err)

	// Read content
	contentBytes := make([]byte, contentLen)
	contentBytesRead, _ := buf.Read(contentBytes)
	if uint32(contentBytesRead) != contentLen {
		panic("Data bytes read mismatch")
	}
	this.mux.Lock()
	this.receivedDataChunks[chunkNumber] = contentBytes
	this.mux.Unlock()
}

// To file bytes (ordered, deduplicated, complete)
func (this *BinaryTransportFileReceiver) Bytes() ([]byte, error) {
	// Write all to final buffer
	finalBuf := new(bytes.Buffer)
	for i := 0; i < int(this.chunkCount); i++ {
		// Write to final buffer
		k := uint32(i)
		finalBuf.Write(this.receivedDataChunks[k])
		delete(this.receivedDataChunks, k)

	}

	// To byte array
	b := finalBuf.Bytes()

	// Validate final CRC with metadata
	crc := crc32.Checksum(b, crcTable)
	if crc != this.fileMeta.Checksum {
		return nil, errors.New("Checksum not valid")
	}

	return b, nil
}

// New receiver
func newBinaryTransportFileReceiver() *BinaryTransportFileReceiver {
	return &BinaryTransportFileReceiver{
		fileMeta:           nil,
		receivedDataChunks: make(map[uint32][]byte),
		lastChunkReceived:  time.Now(),
	}
}
<commit_msg>Better CRC error messages<commit_after>package main

import (
	"bytes"
	"encoding/binary"
	"errors"
	"fmt"
	"hash/crc32"
	"sync"
	"time"
)

// Receiver, should only be used for one file
type BinaryTransportFileReceiver struct {
	lastChunkReceived  time.Time
	mux                sync.RWMutex
	fileMeta           *FileMeta
	data               *bytes.Buffer
	targetShardId      []byte
	chunkCount         uint32
	receivedDataChunks map[uint32][]byte // Data
}

// Receive chunk, returns true if done
func (this *BinaryTransportFileReceiver) Add(buf *bytes.Reader) bool {
	// Update last chunk received
	this.mux.Lock()
	this.lastChunkReceived = time.Now()
	this.mux.Unlock()

	// Read chunk number
	var err error
	var chunkNumber uint32
	err = binary.Read(buf, binary.BigEndian, &chunkNumber)
	panicErr(err)

	// Dedup based on chunk number
	this.mux.Lock()
	if this.receivedDataChunks[chunkNumber] != nil {
		this.mux.Unlock()
		log.Debug("Ignoring data file transfer chunk, duplicate message")
		return false
	}
	this.receivedDataChunks[chunkNumber] = make([]byte, 0)
	this.mux.Unlock()

	// Handle this message
	if chunkNumber == 0 {
		this._readFirstChunk(chunkNumber, buf)
	} else {
		this._readContentChunk(chunkNumber, buf)
	}

	// Done?
	this.mux.RLock()
	defer this.mux.RUnlock()
	return uint32(len(this.receivedDataChunks)) == this.chunkCount
}

// Get last chunk received time
func (this *BinaryTransportFileReceiver) GetLastChunkReceived() time.Time {
	this.mux.RLock()
	defer this.mux.RUnlock()
	return this.lastChunkReceived
}

// Read first chunk
func (this *BinaryTransportFileReceiver) _readFirstChunk(chunkNumber uint32, buf *bytes.Reader) {
	var err error

	// Meta already here? Deduplicate payloads (chunk number could have been 0 by malformed byte)
	this.mux.RLock()
	if this.fileMeta != nil {
		this.mux.RUnlock()
		log.Warn("Ignoring metadata file transfer chunk, duplicate message")
		return
	}
	this.mux.RUnlock()

	// Read target shard flag
	targetShardSet, targetShardSetErr := buf.ReadByte()
	panicErr(targetShardSetErr)

	// Target shard id
	targetShardIdBytes := make([]byte, 16)
	targetShardIdBytesRead, _ := buf.Read(targetShardIdBytes)
	if uint32(targetShardIdBytesRead) != 16 {
		panic("Target shard id bytes read mismatch")
	}

	// Set target shard
	if targetShardSet == 1 {
		this.mux.Lock()
		this.targetShardId = targetShardIdBytes
		this.mux.Unlock()
	}

	// Read meta length
	var metaLen uint32
	err = binary.Read(buf, binary.BigEndian, &metaLen)
	panicErr(err)

	// Read meta
	metaBytes := make([]byte, metaLen)
	metaBytesRead, _ := buf.Read(metaBytes)
	if uint32(metaBytesRead) != metaLen {
		panic("Meta bytes read mismatch")
	}
	meta := &FileMeta{}
	meta.FromBytes(metaBytes)

	// Read chunk count
	var chunkCount uint32
	err = binary.Read(buf, binary.BigEndian, &chunkCount)
	panicErr(err)

	// Set meta into receiver
	this.mux.Lock()
	this.fileMeta = meta
	this.chunkCount = chunkCount
	this.mux.Unlock()

	// Read actual content
	this._readContent(chunkNumber, buf)
}

// Read content chunk
func (this *BinaryTransportFileReceiver) _readContentChunk(chunkNumber uint32, buf *bytes.Reader) {
	// Read content
	this._readContent(chunkNumber, buf)
}

// Read content
func (this *BinaryTransportFileReceiver) _readContent(chunkNumber uint32, buf *bytes.Reader) {
	var err error

	// Read content length
	var contentLen uint32
	err = binary.Read(buf, binary.BigEndian, &contentLen)
	panicErr(err)

	// Read content
	contentBytes := make([]byte, contentLen)
	contentBytesRead, _ := buf.Read(contentBytes)
	if uint32(contentBytesRead) != contentLen {
		panic("Data bytes read mismatch")
	}
	this.mux.Lock()
	this.receivedDataChunks[chunkNumber] = contentBytes
	this.mux.Unlock()
}

// To file bytes (ordered, deduplicated, complete)
func (this *BinaryTransportFileReceiver) Bytes() ([]byte, error) {
	// Write all to final buffer
	finalBuf := new(bytes.Buffer)
	for i := 0; i < int(this.chunkCount); i++ {
		// Write to final buffer
		k := uint32(i)
		finalBuf.Write(this.receivedDataChunks[k])
		delete(this.receivedDataChunks, k)

	}

	// To byte array
	b := finalBuf.Bytes()

	// Validate final CRC with metadata
	crc := crc32.Checksum(b, crcTable)
	if crc != this.fileMeta.Checksum {
		return nil, errors.New(fmt.Sprintf("Checksum not valid, expected %d (len %d) found %d (len %d)", this.fileMeta.Checksum, this.fileMeta.Size, crc, len(b)))
	}

	return b, nil
}

// New receiver
func newBinaryTransportFileReceiver() *BinaryTransportFileReceiver {
	return &BinaryTransportFileReceiver{
		fileMeta:           nil,
		receivedDataChunks: make(map[uint32][]byte),
		lastChunkReceived:  time.Now(),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"runtime"
	"time"

	"math/rand"
	"modules/amqpmodule"

	"config"
)

const rxQueue = "scan_ready_queue"
const txQueue = "conn_scan_results_queue"

var workerCount int
var broker *amqpmodule.Broker
var cipherscan string

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
		panic(fmt.Sprintf("%s: %s", msg, err))
	}
}

func panicIf(err error) {
	if err != nil {
		log.Println(fmt.Sprintf("%s", err))
	}
}

func printIntro() {
	fmt.Println(`
	##################################
	#         TLSRetriever          #
	##################################
	`)
}

type ScanInfo struct {
	Target       string        `json:"target"`
	Timestamp    string        `json:"utctimestamp"`
	ServerSide   string        `json:"serverside"`
	CipherSuites []Ciphersuite `json:"ciphersuite"`
}

type Ciphersuite struct {
	Cipher       string   `json:"cipher"`
	Protocols    []string `json:"protocols"`
	PubKey       []string `json:"pubkey"`
	SigAlg       []string `json:"sigalg"`
	Trusted      string   `json:"trusted"`
	TicketHint   string   `json:"ticket_hint"`
	OCSPStapling string   `json:"ocsp_stapling"`
	PFS          string   `json:"pfs"`
}

func worker(msg []byte) {

	domain := string(msg)
	ip := getRandomIP(domain)

	if ip == "" {
		log.Println("Could not resolve ip for: ", domain)
		return
	}

	cmd := cipherscan + " -j -servername " + domain + " " + ip + ":443 "
	fmt.Println(cmd)
	comm := exec.Command("bash", "-c", cmd)
	var out bytes.Buffer
	var stderr bytes.Buffer
	comm.Stdout = &out
	comm.Stderr = &stderr
	err := comm.Run()
	if err != nil {
		log.Println(err)
	}

	info := ScanInfo{}
	err = json.Unmarshal([]byte(out.String()), &info)
	if err != nil {
		log.Println(err)
		//should we requeue the domain???
		return
	}

	info.Target = domain + "--" + ip

	jsonInfo, er := json.MarshalIndent(info, "", "    ")
	panicIf(er)

	broker.Publish(txQueue, []byte(jsonInfo))
}

func getRandomIP(domain string) string {
	ips, err := net.LookupIP(domain)
	if err != nil {
		return ""
	}

	max := len(ips)

	for {
		if max == 0 {
			return ""
		}
		index := rand.Intn(len(ips))

		if ips[index].To4() != nil {
			return ips[index].String()
		} else {
			ips = append(ips[:index], ips[index+1:]...)
		}
		max--
	}
}

func main() {
	var (
		err error
	)
	cores := runtime.NumCPU()
	runtime.GOMAXPROCS(cores * 2)

	printIntro()

	conf := config.RetrieverConfig{}

	var cfgFile string
	flag.StringVar(&cfgFile, "c", "/etc/observer/retriever.cfg", "Input file csv format")
	flag.StringVar(&cipherscan, "b", "etc/observer/cipherscan", "Cipherscan binary to be used")
	flag.Parse()

	_, err = os.Stat(cfgFile)
	failOnError(err, "Missing configuration file from '-c' or /etc/observer/retriever.cfg")

	conf, err = config.RetrieverConfigLoad(cfgFile)
	if err != nil {
		conf = config.GetRetrieverDefaults()
	}

	_, err = os.Stat(cipherscan)
	failOnError(err, "Could not find ciphercan binary.")

	broker, err = amqpmodule.RegisterURL(conf.General.RabbitMQRelay)

	failOnError(err, "Failed to register RabbitMQ")

	msgs, err := broker.Consume(rxQueue)

	for d := range msgs {
		// block until a worker is available
		for {
			if workerCount < conf.General.MaxSimConns {
				break
			}
			time.Sleep(100 * time.Millisecond)
		}
		workerCount++
		go worker(d)
		log.Printf("Domain %s sent to worker. %d workers currently active.", d, workerCount)
	}
}
<commit_msg>changed cipherscan command and output json struct to support new curve functionality<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"log"
	"net"
	"os"
	"os/exec"
	"runtime"
	"time"

	"math/rand"
	"modules/amqpmodule"

	"config"
)

const rxQueue = "scan_ready_queue"
const txQueue = "conn_scan_results_queue"

var workerCount int
var broker *amqpmodule.Broker
var cipherscan string

func failOnError(err error, msg string) {
	if err != nil {
		log.Fatalf("%s: %s", msg, err)
		panic(fmt.Sprintf("%s: %s", msg, err))
	}
}

func panicIf(err error) {
	if err != nil {
		log.Println(fmt.Sprintf("%s", err))
	}
}

func printIntro() {
	fmt.Println(`
	##################################
	#         TLSRetriever          #
	##################################
	`)
}

type ScanInfo struct {
	Target         string        `json:"target"`
	Timestamp      string        `json:"utctimestamp"`
	ServerSide     string        `json:"serverside"`
	CipherSuites   []Ciphersuite `json:"ciphersuite"`
	CurvesFallback string        `json:"curves_fallback"`
}

type Ciphersuite struct {
	Cipher       string   `json:"cipher"`
	Protocols    []string `json:"protocols"`
	PubKey       []string `json:"pubkey"`
	SigAlg       []string `json:"sigalg"`
	Trusted      string   `json:"trusted"`
	TicketHint   string   `json:"ticket_hint"`
	OCSPStapling string   `json:"ocsp_stapling"`
	PFS          string   `json:"pfs"`
	Curves       string   `json:"curves,omitempty"`
}

func worker(msg []byte) {

	domain := string(msg)
	ip := getRandomIP(domain)

	if ip == "" {
		log.Println("Could not resolve ip for: ", domain)
		return
	}

	cmd := cipherscan + " -j --curves -servername " + domain + " " + ip + ":443 "
	fmt.Println(cmd)
	comm := exec.Command("bash", "-c", cmd)
	var out bytes.Buffer
	var stderr bytes.Buffer
	comm.Stdout = &out
	comm.Stderr = &stderr
	err := comm.Run()
	if err != nil {
		log.Println(err)
	}

	info := ScanInfo{}
	err = json.Unmarshal([]byte(out.String()), &info)
	if err != nil {
		log.Println(err)
		//should we requeue the domain???
		return
	}

	info.Target = domain + "--" + ip

	jsonInfo, er := json.MarshalIndent(info, "", "    ")
	panicIf(er)

	broker.Publish(txQueue, []byte(jsonInfo))
}

func getRandomIP(domain string) string {
	ips, err := net.LookupIP(domain)
	if err != nil {
		return ""
	}

	max := len(ips)

	for {
		if max == 0 {
			return ""
		}
		index := rand.Intn(len(ips))

		if ips[index].To4() != nil {
			return ips[index].String()
		} else {
			ips = append(ips[:index], ips[index+1:]...)
		}
		max--
	}
}

func main() {
	var (
		err error
	)
	cores := runtime.NumCPU()
	runtime.GOMAXPROCS(cores * 2)

	printIntro()

	conf := config.RetrieverConfig{}

	var cfgFile string
	flag.StringVar(&cfgFile, "c", "/etc/observer/retriever.cfg", "Input file csv format")
	flag.StringVar(&cipherscan, "b", "etc/observer/cipherscan", "Cipherscan binary to be used")
	flag.Parse()

	_, err = os.Stat(cfgFile)
	failOnError(err, "Missing configuration file from '-c' or /etc/observer/retriever.cfg")

	conf, err = config.RetrieverConfigLoad(cfgFile)
	if err != nil {
		conf = config.GetRetrieverDefaults()
	}

	_, err = os.Stat(cipherscan)
	failOnError(err, "Could not find ciphercan binary.")

	broker, err = amqpmodule.RegisterURL(conf.General.RabbitMQRelay)

	failOnError(err, "Failed to register RabbitMQ")

	msgs, err := broker.Consume(rxQueue)

	for d := range msgs {
		// block until a worker is available
		for {
			if workerCount < conf.General.MaxSimConns {
				break
			}
			time.Sleep(100 * time.Millisecond)
		}
		workerCount++
		go worker(d)
		log.Printf("Domain %s sent to worker. %d workers currently active.", d, workerCount)
	}
}
<|endoftext|>


<commit_before>/*
Copyright 2017 Gravitational, Inc.

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 influxdb

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"
	"text/template"

	"github.com/gravitational/monitoring-app/watcher/lib/constants"
	"github.com/gravitational/monitoring-app/watcher/lib/utils"

	"github.com/gravitational/trace"
)

// Rollup is the rollup configuration
type Rollup struct {
	// Retention is the retention policy for this rollup
	Retention string `json:"retention"`
	// Measurement is the name of the measurement to run rollup on
	Measurement string `json:"measurement,omitempty"`
	// Name is both the name of the rollup query and the name of the
	// new measurement rollup data will be inserted into
	Name string `json:"name"`
	// Functions is a list of functions for rollup calculation
	Functions []Function `json:"functions"`
	// CustomFrom is a custom 'from' clause
	CustomFrom string `json:"custom_from,omitempty"`
	// CustomGroupBy is a custom 'group by' clause
	CustomGroupBy string `json:"custom_group_by,omitempty"`
}

// Check verifies that rollup configuration is correct
func (r Rollup) Check() error {
	var errors []error
	if !utils.OneOf(r.Retention, constants.AllRetentions) {
		errors = append(errors, trace.BadParameter(
			"invalid Retention, must be one of: %v", constants.AllRetentions))
	}
	if r.Measurement == "" && r.CustomFrom == "" {
		errors = append(errors, trace.BadParameter("parameter Measurement or CustomFrom is missing"))
	}
	if r.Name == "" {
		errors = append(errors, trace.BadParameter("parameter Name is missing"))
	}
	if len(r.Functions) == 0 {
		errors = append(errors, trace.BadParameter("parameter Functions is empty"))
	}
	for _, function := range r.Functions {
		if err := function.Check(); err != nil {
			errors = append(errors, trace.Wrap(err))
		}
	}
	return trace.NewAggregate(errors...)
}

// buildCreateQuery returns a string with InfluxDB query to create rollup
// based on the rollup configuration
func (r *Rollup) buildCreateQuery() (string, error) {
	var functions []string
	for _, fn := range r.Functions {
		function, err := fn.buildFunction()
		if err != nil {
			return "", trace.Wrap(err)
		}
		functions = append(functions, function)
	}

	var b bytes.Buffer
	err := createQueryTemplate.Execute(&b, map[string]string{
		"name":             r.Name,
		"database":         constants.InfluxDBDatabase,
		"functions":        strings.Join(functions, ", "),
		"retention_into":   r.Retention,
		"measurement_into": r.Name,
		"retention_from":   constants.InfluxDBRetentionPolicy,
		"measurement_from": r.Measurement,
		"custom_from":      r.CustomFrom,
		"custom_group_by":  r.CustomGroupBy,
		"interval":         constants.RetentionToInterval[r.Retention],
	})
	if err != nil {
		return "", trace.Wrap(err)
	}

	return b.String(), nil
}

// buildDeleteQuery returns a string with InfluxDB query to delete rollup
// based on the rollup configuration
func (r *Rollup) buildDeleteQuery() (string, error) {
	var b bytes.Buffer
	err := deleteQueryTemplate.Execute(&b, map[string]string{
		"name":     r.Name,
		"database": constants.InfluxDBDatabase,
	})
	if err != nil {
		return "", trace.Wrap(err)
	}

	return b.String(), nil
}

// Function defines a single rollup function
type Function struct {
	// Function is the function name (mean, max, etc.)
	Function string `json:"function"`
	// Field is the name of the field to apply the function to
	Field string `json:"field"`
	// Alias is the optional alias for the new field in the rollup table
	Alias string `json:"alias,omitempty"`
}

// Check verifies the function configuration is correct
func (f Function) Check() error {
	var errors []error
	if !utils.OneOf(f.Function, constants.SimpleFunctions) && !f.isComposite() {
		errors = append(errors, trace.BadParameter(
			"invalid Function, must be one of %v, or a composite function starting with one of %v prefixes",
			constants.SimpleFunctions, constants.CompositeFunctions))
	}
	if f.isComposite() {
		funcAndValue := strings.Split(f.Function, "_")
		if len(funcAndValue) != 2 {
			errors = append(errors, trace.BadParameter(
				"percentile function must have format like 'percentile_90', 'top_10', 'bottom_10' or 'sample_1000' "))
		}
	}
	if f.Field == "" {
		errors = append(errors, trace.BadParameter("parameter Field is missing"))
	}
	return trace.NewAggregate(errors...)
}

// buildFunction returns a function string based on the provided function configuration
func (f *Function) buildFunction() (string, error) {
	alias := f.Alias
	if alias == "" {
		alias = f.Field
	}

	// split function name, based on the "_" separator (eg: percentile_99, top_10, ecc)
	err := f.Check()
	if err != nil {
		return "", trace.Wrap(err)
	}

	if !f.isComposite() {
		return fmt.Sprintf(`%v("%v") as %v`, f.Function, f.Field, alias), nil
	}

	funcAndValue := strings.Split(f.Function, "_")
	funcName := funcAndValue[0]
	param := funcAndValue[1]

	err = validateParam(funcName, param)
	if err != nil {
		return "", trace.Wrap(err)
	}
	return fmt.Sprintf(`%v("%v", %v) as %v`, funcName, f.Field, param, alias), nil
}

// isComposite checks if the specified function is composite
func (f *Function) isComposite() bool {
	for _, name := range constants.CompositeFunctions {
		if strings.HasPrefix(f.Function, name) {
			return true
		}
	}
	return false
}

// validateParam checks the function parameter for validity.
func validateParam(funcName, param string) error {
	// convert parameter value as it's always going to be an Integer
	value, err := strconv.Atoi(param)
	if err != nil {
		return trace.Wrap(err)
	}

	switch funcName {
	case constants.FunctionPercentile:
		if value < 0 || value > 100 {
			return trace.BadParameter(
				"percentile value must be between 0 and 100 (inclusive)")
		}
	case constants.FunctionTop, constants.FunctionBottom, constants.FunctionSample:
		if value < 0 {
			return trace.BadParameter(
				"top, bottom and sample value must be greater than or equal to 0")
		}
	}

	return nil
}

var (
	// createQueryTemplate is the template for creating InfluxDB continuous query
	createQueryTemplate = template.Must(template.New("query").Parse(
		`create continuous query "{{.name}}" on {{.database}} begin select {{.functions}} into {{.database}}."{{.retention_into}}"."{{.measurement_into}}" from {{.database}}."{{.retention_from}}"."{{.measurement_from}}" group by {{if .custom_group_by ne ""}}{{.custom_group_by}}{{else}}*, time({{.interval}}){{end}} end`))
	// deleteQueryTemplate is the template for deleting Influx continuous query
	deleteQueryTemplate = template.Must(template.New("query").Parse(
		`drop continuous query "{{.name}}" on {{.database}}`))
)
<commit_msg>Fix equal function for text template<commit_after>/*
Copyright 2017 Gravitational, Inc.

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 influxdb

import (
	"bytes"
	"fmt"
	"strconv"
	"strings"
	"text/template"

	"github.com/gravitational/monitoring-app/watcher/lib/constants"
	"github.com/gravitational/monitoring-app/watcher/lib/utils"

	"github.com/gravitational/trace"
)

// Rollup is the rollup configuration
type Rollup struct {
	// Retention is the retention policy for this rollup
	Retention string `json:"retention"`
	// Measurement is the name of the measurement to run rollup on
	Measurement string `json:"measurement,omitempty"`
	// Name is both the name of the rollup query and the name of the
	// new measurement rollup data will be inserted into
	Name string `json:"name"`
	// Functions is a list of functions for rollup calculation
	Functions []Function `json:"functions"`
	// CustomFrom is a custom 'from' clause
	CustomFrom string `json:"custom_from,omitempty"`
	// CustomGroupBy is a custom 'group by' clause
	CustomGroupBy string `json:"custom_group_by,omitempty"`
}

// Check verifies that rollup configuration is correct
func (r Rollup) Check() error {
	var errors []error
	if !utils.OneOf(r.Retention, constants.AllRetentions) {
		errors = append(errors, trace.BadParameter(
			"invalid Retention, must be one of: %v", constants.AllRetentions))
	}
	if r.Measurement == "" && r.CustomFrom == "" {
		errors = append(errors, trace.BadParameter("parameter Measurement or CustomFrom is missing"))
	}
	if r.Name == "" {
		errors = append(errors, trace.BadParameter("parameter Name is missing"))
	}
	if len(r.Functions) == 0 {
		errors = append(errors, trace.BadParameter("parameter Functions is empty"))
	}
	for _, function := range r.Functions {
		if err := function.Check(); err != nil {
			errors = append(errors, trace.Wrap(err))
		}
	}
	return trace.NewAggregate(errors...)
}

// buildCreateQuery returns a string with InfluxDB query to create rollup
// based on the rollup configuration
func (r *Rollup) buildCreateQuery() (string, error) {
	var functions []string
	for _, fn := range r.Functions {
		function, err := fn.buildFunction()
		if err != nil {
			return "", trace.Wrap(err)
		}
		functions = append(functions, function)
	}

	var b bytes.Buffer
	err := createQueryTemplate.Execute(&b, map[string]string{
		"name":             r.Name,
		"database":         constants.InfluxDBDatabase,
		"functions":        strings.Join(functions, ", "),
		"retention_into":   r.Retention,
		"measurement_into": r.Name,
		"retention_from":   constants.InfluxDBRetentionPolicy,
		"measurement_from": r.Measurement,
		"custom_from":      r.CustomFrom,
		"custom_group_by":  r.CustomGroupBy,
		"interval":         constants.RetentionToInterval[r.Retention],
	})
	if err != nil {
		return "", trace.Wrap(err)
	}

	return b.String(), nil
}

// buildDeleteQuery returns a string with InfluxDB query to delete rollup
// based on the rollup configuration
func (r *Rollup) buildDeleteQuery() (string, error) {
	var b bytes.Buffer
	err := deleteQueryTemplate.Execute(&b, map[string]string{
		"name":     r.Name,
		"database": constants.InfluxDBDatabase,
	})
	if err != nil {
		return "", trace.Wrap(err)
	}

	return b.String(), nil
}

// Function defines a single rollup function
type Function struct {
	// Function is the function name (mean, max, etc.)
	Function string `json:"function"`
	// Field is the name of the field to apply the function to
	Field string `json:"field"`
	// Alias is the optional alias for the new field in the rollup table
	Alias string `json:"alias,omitempty"`
}

// Check verifies the function configuration is correct
func (f Function) Check() error {
	var errors []error
	if !utils.OneOf(f.Function, constants.SimpleFunctions) && !f.isComposite() {
		errors = append(errors, trace.BadParameter(
			"invalid Function, must be one of %v, or a composite function starting with one of %v prefixes",
			constants.SimpleFunctions, constants.CompositeFunctions))
	}
	if f.isComposite() {
		funcAndValue := strings.Split(f.Function, "_")
		if len(funcAndValue) != 2 {
			errors = append(errors, trace.BadParameter(
				"percentile function must have format like 'percentile_90', 'top_10', 'bottom_10' or 'sample_1000' "))
		}
	}
	if f.Field == "" {
		errors = append(errors, trace.BadParameter("parameter Field is missing"))
	}
	return trace.NewAggregate(errors...)
}

// buildFunction returns a function string based on the provided function configuration
func (f *Function) buildFunction() (string, error) {
	alias := f.Alias
	if alias == "" {
		alias = f.Field
	}

	// split function name, based on the "_" separator (eg: percentile_99, top_10, ecc)
	err := f.Check()
	if err != nil {
		return "", trace.Wrap(err)
	}

	if !f.isComposite() {
		return fmt.Sprintf(`%v("%v") as %v`, f.Function, f.Field, alias), nil
	}

	funcAndValue := strings.Split(f.Function, "_")
	funcName := funcAndValue[0]
	param := funcAndValue[1]

	err = validateParam(funcName, param)
	if err != nil {
		return "", trace.Wrap(err)
	}
	return fmt.Sprintf(`%v("%v", %v) as %v`, funcName, f.Field, param, alias), nil
}

// isComposite checks if the specified function is composite
func (f *Function) isComposite() bool {
	for _, name := range constants.CompositeFunctions {
		if strings.HasPrefix(f.Function, name) {
			return true
		}
	}
	return false
}

// validateParam checks the function parameter for validity.
func validateParam(funcName, param string) error {
	// convert parameter value as it's always going to be an Integer
	value, err := strconv.Atoi(param)
	if err != nil {
		return trace.Wrap(err)
	}

	switch funcName {
	case constants.FunctionPercentile:
		if value < 0 || value > 100 {
			return trace.BadParameter(
				"percentile value must be between 0 and 100 (inclusive)")
		}
	case constants.FunctionTop, constants.FunctionBottom, constants.FunctionSample:
		if value < 0 {
			return trace.BadParameter(
				"top, bottom and sample value must be greater than or equal to 0")
		}
	}

	return nil
}

var (
	// createQueryTemplate is the template for creating InfluxDB continuous query
	createQueryTemplate = template.Must(template.New("query").Parse(
		`create continuous query "{{.name}}" on {{.database}} begin select {{.functions}} into {{.database}}."{{.retention_into}}"."{{.measurement_into}}" from {{.database}}."{{.retention_from}}"."{{.measurement_from}}" group by {{if ne .custom_group_by ""}}{{.custom_group_by}}{{else}}*, time({{.interval}}){{end}} end`))
	// deleteQueryTemplate is the template for deleting Influx continuous query
	deleteQueryTemplate = template.Must(template.New("query").Parse(
		`drop continuous query "{{.name}}" on {{.database}}`))
)
<|endoftext|>


<commit_before>/*
Copyright 2019 The Knative 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.
*/

package ingress

import (
	"context"
	"errors"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"testing"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	pkgTest "knative.dev/pkg/test"
	"knative.dev/serving/pkg/apis/networking/v1alpha1"
	"knative.dev/serving/test"
)

// CreateService creates a Kubernetes service that will respond to the protocol
// specified with the given portName.  It returns the service name, the port on
// which the service is listening, and a "cancel" function to clean up the
// created resources.
func CreateService(t *testing.T, clients *test.Clients, portName string) (string, int, context.CancelFunc) {
	t.Helper()
	name := test.ObjectNameForTest(t)

	// Avoid zero, but pick a low port number.
	port := 3 + rand.Intn(97)
	t.Logf("Using port %d", port)

	// Pick a high port number.
	containerPort := 8000 + rand.Intn(100)
	t.Logf("Using containerPort %d", containerPort)

	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: test.ServingNamespace,
			Labels: map[string]string{
				"test-pod": name,
			},
		},
		Spec: corev1.PodSpec{
			Containers: []corev1.Container{{
				Name:  "foo",
				Image: pkgTest.ImagePath("runtime"),
				Ports: []corev1.ContainerPort{{
					Name:          portName,
					ContainerPort: int32(containerPort),
				}},
				// This is needed by the runtime image we are using.
				Env: []corev1.EnvVar{{
					Name:  "PORT",
					Value: strconv.Itoa(containerPort),
				}},
			}},
		},
	}
	pod, err := clients.KubeClient.Kube.CoreV1().Pods(pod.Namespace).Create(pod)
	if err != nil {
		t.Fatalf("Error creating Pod: %v", err)
	}
	cancel := func() {
		err := clients.KubeClient.Kube.CoreV1().Pods(pod.Namespace).Delete(pod.Name, &metav1.DeleteOptions{})
		if err != nil {
			t.Errorf("Error cleaning up Pod %s", pod.Name)
		}
	}

	svc := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: test.ServingNamespace,
			Labels: map[string]string{
				"test-pod": name,
			},
		},
		Spec: corev1.ServiceSpec{
			Type: "ClusterIP",
			Ports: []corev1.ServicePort{{
				Name:       portName,
				Port:       int32(port),
				TargetPort: intstr.FromInt(int(containerPort)),
			}},
			Selector: map[string]string{
				"test-pod": name,
			},
		},
	}
	svc, err = clients.KubeClient.Kube.CoreV1().Services(svc.Namespace).Create(svc)
	if err != nil {
		cancel()
		t.Fatalf("Error creating Service: %v", err)
	}

	return name, port, func() {
		err := clients.KubeClient.Kube.CoreV1().Services(svc.Namespace).Delete(svc.Name, &metav1.DeleteOptions{})
		if err != nil {
			t.Errorf("Error cleaning up Service %s", pod.Name)
		}
		cancel()
	}
}

// CreateDialContext looks up the endpoint information to create a "dialer" for
// the provided Ingress' public ingress loas balancer.  It can be used to
// contact external-visibility services with an HTTP client via:
//
//	client := &http.Client{
//		Transport: &http.Transport{
//			DialContext: CreateDialContext(t, ing, clients),
//		},
//	}
func CreateDialContext(t *testing.T, ing *v1alpha1.Ingress, clients *test.Clients) func(context.Context, string, string) (net.Conn, error) {
	t.Helper()
	if ing.Status.PublicLoadBalancer == nil || len(ing.Status.PublicLoadBalancer.Ingress) < 1 {
		t.Fatal("Ingress does not have a public load balancer assigned.")
	}

	// TODO(mattmoor): I'm open to tricks that would let us cleanly test multiple
	// public load balancers or LBs with multiple ingresses (below), but want to
	// keep our simple tests simple, thus the [0]s...

	// We expect an ingress LB with the form foo.bar.svc.cluster.local (though
	// we aren't strictly sensitive to the suffix, this is just illustrative.
	internalDomain := ing.Status.PublicLoadBalancer.Ingress[0].DomainInternal
	parts := strings.SplitN(internalDomain, ".", 3)
	if len(parts) < 3 {
		t.Fatalf("Too few parts in internal domain: %s", internalDomain)
	}
	name, namespace := parts[0], parts[1]

	svc, err := clients.KubeClient.Kube.CoreV1().Services(namespace).Get(name, metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Unable to retrieve Kubernetes service %s/%s: %v", namespace, name, err)
	}
	if len(svc.Status.LoadBalancer.Ingress) < 1 {
		t.Fatal("Service does not have any ingresses (not type LoadBalancer?).")
	}
	ingress := svc.Status.LoadBalancer.Ingress[0]

	return func(_ context.Context, _ string, address string) (net.Conn, error) {
		_, port, err := net.SplitHostPort(address)
		if err != nil {
			return nil, err
		}
		if ingress.IP != "" {
			return net.Dial("tcp", ingress.IP+":"+port)
		}
		if ingress.Hostname != "" {
			return net.Dial("tcp", ingress.Hostname+":"+port)
		}
		return nil, errors.New("Service ingress does not contain dialing information.")
	}
}
<commit_msg>Add ReadinessProbe and Endpoints check to helper. (#6310)<commit_after>/*
Copyright 2019 The Knative 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.
*/

package ingress

import (
	"context"
	"errors"
	"math/rand"
	"net"
	"strconv"
	"strings"
	"testing"

	corev1 "k8s.io/api/core/v1"
	metav1 "k8s.io/apimachinery/pkg/apis/meta/v1"
	"k8s.io/apimachinery/pkg/util/intstr"
	"k8s.io/apimachinery/pkg/util/wait"
	pkgTest "knative.dev/pkg/test"
	"knative.dev/serving/pkg/apis/networking/v1alpha1"
	"knative.dev/serving/test"
)

// CreateService creates a Kubernetes service that will respond to the protocol
// specified with the given portName.  It returns the service name, the port on
// which the service is listening, and a "cancel" function to clean up the
// created resources.
func CreateService(t *testing.T, clients *test.Clients, portName string) (string, int, context.CancelFunc) {
	t.Helper()
	name := test.ObjectNameForTest(t)

	// Avoid zero, but pick a low port number.
	port := 3 + rand.Intn(97)
	t.Logf("Using port %d", port)

	// Pick a high port number.
	containerPort := 8000 + rand.Intn(100)
	t.Logf("Using containerPort %d", containerPort)

	pod := &corev1.Pod{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: test.ServingNamespace,
			Labels: map[string]string{
				"test-pod": name,
			},
		},
		Spec: corev1.PodSpec{
			Containers: []corev1.Container{{
				Name:  "foo",
				Image: pkgTest.ImagePath("runtime"),
				Ports: []corev1.ContainerPort{{
					Name:          portName,
					ContainerPort: int32(containerPort),
				}},
				// This is needed by the runtime image we are using.
				Env: []corev1.EnvVar{{
					Name:  "PORT",
					Value: strconv.Itoa(containerPort),
				}},
				ReadinessProbe: &corev1.Probe{
					Handler: corev1.Handler{
						HTTPGet: &corev1.HTTPGetAction{
							Path: "/healthz",
							Port: intstr.FromInt(containerPort),
						},
					},
				},
			}},
		},
	}
	test.CleanupOnInterrupt(func() { clients.KubeClient.Kube.CoreV1().Pods(pod.Namespace).Delete(pod.Name, &metav1.DeleteOptions{}) })
	pod, err := clients.KubeClient.Kube.CoreV1().Pods(pod.Namespace).Create(pod)
	if err != nil {
		t.Fatalf("Error creating Pod: %v", err)
	}
	cancel := func() {
		err := clients.KubeClient.Kube.CoreV1().Pods(pod.Namespace).Delete(pod.Name, &metav1.DeleteOptions{})
		if err != nil {
			t.Errorf("Error cleaning up Pod %s", pod.Name)
		}
	}

	svc := &corev1.Service{
		ObjectMeta: metav1.ObjectMeta{
			Name:      name,
			Namespace: test.ServingNamespace,
			Labels: map[string]string{
				"test-pod": name,
			},
		},
		Spec: corev1.ServiceSpec{
			Type: "ClusterIP",
			Ports: []corev1.ServicePort{{
				Name:       portName,
				Port:       int32(port),
				TargetPort: intstr.FromInt(int(containerPort)),
			}},
			Selector: map[string]string{
				"test-pod": name,
			},
		},
	}
	test.CleanupOnInterrupt(func() {
		clients.KubeClient.Kube.CoreV1().Services(svc.Namespace).Delete(svc.Name, &metav1.DeleteOptions{})
	})
	svc, err = clients.KubeClient.Kube.CoreV1().Services(svc.Namespace).Create(svc)
	if err != nil {
		cancel()
		t.Fatalf("Error creating Service: %v", err)
	}

	// Wait for the Pod to show up in the Endpoints resource.
	waitErr := wait.PollImmediate(test.PollInterval, test.PollTimeout, func() (bool, error) {
		ep, err := clients.KubeClient.Kube.CoreV1().Endpoints(svc.Namespace).Get(svc.Name, metav1.GetOptions{})
		if err != nil {
			return true, err
		}
		return len(ep.Subsets) == 1, nil
	})
	if waitErr != nil {
		cancel()
		t.Fatalf("Error waiting for Endpoints to contain a Pod IP: %v", err)
	}

	return name, port, func() {
		err := clients.KubeClient.Kube.CoreV1().Services(svc.Namespace).Delete(svc.Name, &metav1.DeleteOptions{})
		if err != nil {
			t.Errorf("Error cleaning up Service %s", svc.Name)
		}
		cancel()
	}
}

// CreateDialContext looks up the endpoint information to create a "dialer" for
// the provided Ingress' public ingress loas balancer.  It can be used to
// contact external-visibility services with an HTTP client via:
//
//	client := &http.Client{
//		Transport: &http.Transport{
//			DialContext: CreateDialContext(t, ing, clients),
//		},
//	}
func CreateDialContext(t *testing.T, ing *v1alpha1.Ingress, clients *test.Clients) func(context.Context, string, string) (net.Conn, error) {
	t.Helper()
	if ing.Status.PublicLoadBalancer == nil || len(ing.Status.PublicLoadBalancer.Ingress) < 1 {
		t.Fatal("Ingress does not have a public load balancer assigned.")
	}

	// TODO(mattmoor): I'm open to tricks that would let us cleanly test multiple
	// public load balancers or LBs with multiple ingresses (below), but want to
	// keep our simple tests simple, thus the [0]s...

	// We expect an ingress LB with the form foo.bar.svc.cluster.local (though
	// we aren't strictly sensitive to the suffix, this is just illustrative.
	internalDomain := ing.Status.PublicLoadBalancer.Ingress[0].DomainInternal
	parts := strings.SplitN(internalDomain, ".", 3)
	if len(parts) < 3 {
		t.Fatalf("Too few parts in internal domain: %s", internalDomain)
	}
	name, namespace := parts[0], parts[1]

	svc, err := clients.KubeClient.Kube.CoreV1().Services(namespace).Get(name, metav1.GetOptions{})
	if err != nil {
		t.Fatalf("Unable to retrieve Kubernetes service %s/%s: %v", namespace, name, err)
	}
	if len(svc.Status.LoadBalancer.Ingress) < 1 {
		t.Fatal("Service does not have any ingresses (not type LoadBalancer?).")
	}
	ingress := svc.Status.LoadBalancer.Ingress[0]

	return func(_ context.Context, _ string, address string) (net.Conn, error) {
		_, port, err := net.SplitHostPort(address)
		if err != nil {
			return nil, err
		}
		if ingress.IP != "" {
			return net.Dial("tcp", ingress.IP+":"+port)
		}
		if ingress.Hostname != "" {
			return net.Dial("tcp", ingress.Hostname+":"+port)
		}
		return nil, errors.New("Service ingress does not contain dialing information.")
	}
}
<|endoftext|>


<commit_before>package foo

import (
	"os"
)

type Namer interface {
	Name() string
}

type MyWalkFunc func(path string, info os.FileInfo, err error) error
type MyWalkFunc2 func(path string, info os.FileInfo, err error) error

func Impl(path string, info os.FileInfo, err error) error {
	info.Name()
	return nil
}

type MyIface interface {
	FooBar(f *os.File)
}
type MyIface2 interface {
	MyIface
}

type st struct{}

func (s st) FooBar(f *os.File) {}

func FooWrong(s st) {
	s.FooBar(nil)
}
<commit_msg>Improve test coverage<commit_after>package foo

import (
	"os"
)

type Namer interface {
	Name() string
}

type MyWalkFunc func(path string, info os.FileInfo, err error) error
type MyWalkFunc2 func(path string, info os.FileInfo, err error) error

func Impl(path string, info os.FileInfo, err error) error {
	info.Name()
	return nil
}

type MyIface interface {
	FooBar(f *os.File)
}
type MyIface2 interface {
	MyIface
}

type st struct{}

func (s st) FooBar(f *os.File) {}

func FooWrong(s st) {
	s.FooBar(nil)
}

type FooBarFunc func(f *os.File)
<|endoftext|>


<commit_before>package gitsync

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"log"
	"net"
	"strings"
)

var (
	Port             = 9999 // mDNS/Bonjour uses 5353
	IP4MulticastAddr = &net.UDPAddr{
		IP:   net.ParseIP("224.0.0.251"),
		Port: Port,
	}

	IP6MulticastAddr = &net.UDPAddr{
		IP:   net.ParseIP("ff02::fb"),
		Port: Port,
	}
)

func init() {
	gob.Register(GitChange{})
}

func establishConnPair(addr *net.UDPAddr) (recvConn, sendConn *net.UDPConn, err error) {
	if recvConn, err = net.ListenMulticastUDP("udp", nil, addr); err != nil {
		return
	}

	if sendConn, err = net.DialUDP("udp", nil, addr); err != nil {
		return
	}

	return
}

// NetIO shares GitChanges on toNet with the network via a multicast group. It
// will pass on GitChanges from the network via fromNet. It uniques the daemon
// instance by changing the .Name member to be name@<host IP>/<original .Name)
func NetIO(name string, addr *net.UDPAddr, fromNet, toNet chan GitChange) {
	var (
		err                error
		netName            string       // name prefix to identify ourselves with
		recvConn, sendConn *net.UDPConn // UDP connections to allow us to send and	receive change updates
	)

	log.Printf("Joining %v multicast(%t) group", addr, addr.IP.IsMulticast())
	if recvConn, sendConn, err = establishConnPair(addr); err != nil {
		log.Fatalf("Error joining listening: %s\n", addr, err)
		return
	}

	log.Printf("Successfully joined %v multicast(%t) group", addr, addr.IP.IsMulticast())
	defer recvConn.Close()
	defer sendConn.Close()
	netName = fmt.Sprintf("%s@%s", name, sendConn.LocalAddr().(*net.UDPAddr).IP)

	term := false
	defer func() { term = true }()
	rawFromNet := make(chan []byte, 128)
	go func() {
		for !term {
			b := make([]byte, 1024)

			if n, err := recvConn.Read(b); err != nil {
				log.Fatalf("Cannot read socket: %s", err)
				continue
			} else {
				rawFromNet <- b[:n]
			}
		}
	}()

	for {
		select {
		case req, ok := <-toNet:
			if !ok {
				return
			}

			req.Name = fmt.Sprintf("%s:%s", netName, req.Name)

			//log.Printf("Sending %+v", req)
			buf := &bytes.Buffer{}
			enc := gob.NewEncoder(buf)

			if err := enc.Encode(req); err != nil {
				log.Fatalf("%s", err)
			}

			//log.Printf("Sending %+v", buf.Bytes())
			if _, err := sendConn.Write(buf.Bytes()); err != nil {
				log.Fatalf("%s", err)
			}

		case resp := <-rawFromNet:
			var msg GitChange
			dec := gob.NewDecoder(bytes.NewReader(resp))

			if err := dec.Decode(&msg); err != nil {
				log.Fatalf("%s", err)
			} else {
				//log.Printf("received %+v", msg)
			}

			if !strings.HasPrefix(msg.Name, netName) {
				fromNet <- msg
			}
		}
	}
}
<commit_msg>Improved netio handling of critical errors<commit_after>package gitsync

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"log"
	"net"
	"strings"
)

var (
	Port             = 9999 // mDNS/Bonjour uses 5353
	IP4MulticastAddr = &net.UDPAddr{
		IP:   net.ParseIP("224.0.0.251"),
		Port: Port,
	}

	IP6MulticastAddr = &net.UDPAddr{
		IP:   net.ParseIP("ff02::fb"),
		Port: Port,
	}
)

func init() {
	gob.Register(GitChange{})
}

func establishConnPair(addr *net.UDPAddr) (recvConn, sendConn *net.UDPConn, err error) {
	if recvConn, err = net.ListenMulticastUDP("udp", nil, addr); err != nil {
		return
	}

	if sendConn, err = net.DialUDP("udp", nil, addr); err != nil {
		return
	}

	return
}

// NetIO shares GitChanges on toNet with the network via a multicast group. It
// will pass on GitChanges from the network via fromNet. It uniques the daemon
// instance by changing the .Name member to be name@<host IP>/<original .Name)
func NetIO(name string, addr *net.UDPAddr, fromNet, toNet chan GitChange) {
	var (
		err                error
		netName            string       // name prefix to identify ourselves with
		recvConn, sendConn *net.UDPConn // UDP connections to allow us to send and	receive change updates
	)

	log.Printf("Joining %v multicast(%t) group", addr, addr.IP.IsMulticast())
	if recvConn, sendConn, err = establishConnPair(addr); err != nil {
		log.Fatalf("Error joining listening: %s\n", addr, err)
		return
	}

	log.Printf("Successfully joined %v multicast(%t) group", addr, addr.IP.IsMulticast())
	defer recvConn.Close()
	defer sendConn.Close()
	netName = fmt.Sprintf("%s@%s", name, sendConn.LocalAddr().(*net.UDPAddr).IP)

	term := false
	defer func() { term = true }()
	rawFromNet := make(chan []byte, 128)
	go func() {
		for !term {
			b := make([]byte, 1024)

			if n, err := recvConn.Read(b); err != nil {
				log.Fatalf("Cannot read socket: %s", err)
				continue
			} else {
				rawFromNet <- b[:n]
			}
		}
	}()

	for {
		select {
		case req, ok := <-toNet:
			if !ok {
				return
			}

			req.Name = fmt.Sprintf("%s:%s", netName, req.Name)

			//log.Printf("Sending %+v", req)
			buf := &bytes.Buffer{}
			enc := gob.NewEncoder(buf)

			if err := enc.Encode(req); err != nil {
				log.Fatalf("%s", err)
			}

			//log.Printf("Sending %+v", buf.Bytes())
				continue
			if _, err := sendConn.Write(buf.Bytes()); err != nil {
				log.Fatalf("%s", err)
				continue
			}

		case resp := <-rawFromNet:
			var msg GitChange
			dec := gob.NewDecoder(bytes.NewReader(resp))

			if err := dec.Decode(&msg); err != nil {
				log.Fatalf("%s", err)
			} else {
				continue
				//log.Printf("received %+v", msg)
			}

			if !strings.HasPrefix(msg.Name, netName) {
				fromNet <- msg
			}
		}
	}
}
<|endoftext|>


<commit_before>// Package checkpoint persists event log state information to disk so that
// event log monitoring can resume from the last read event in the case of a
// restart or unexpected interruption.
package checkpoint

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"github.com/elastic/beats/libbeat/logp"
	"gopkg.in/yaml.v2"
)

// Checkpoint persists event log state information to disk.
type Checkpoint struct {
	wg            sync.WaitGroup // WaitGroup used to wait on the shutdown of the checkpoint worker.
	done          chan struct{}  // Channel for shutting down the checkpoint worker.
	once          sync.Once      // Used to guarantee shutdown happens once.
	file          string         // File where the state is persisted.
	numUpdates    int            // Number of updates received since last persisting to disk.
	maxUpdates    int            // Maximum number of updates to buffer before persisting to disk.
	flushInterval time.Duration  // Maximum time interval that can pass before persisting to disk.
	sort          []string       // Slice used for sorting states map (store to save on mallocs).

	lock   sync.RWMutex
	states map[string]EventLogState

	save chan EventLogState
}

// PersistedState represents the format of the data persisted to disk.
type PersistedState struct {
	UpdateTime time.Time       `config:"update_time"`
	States     []EventLogState `config:"event_logs"`
}

// EventLogState represents the state of an individual event log.
type EventLogState struct {
	Name         string    `config:"name"`
	RecordNumber uint64    `config:"record_number"`
	Timestamp    time.Time `config:"timestamp"`
}

// NewCheckpoint creates and returns a new Checkpoint. This method loads state
// information from disk if it exists and starts a goroutine for persisting
// state information to disk. Shutdown should be called when finished to
// guarantee any in-memory state information is flushed to disk.
//
// file is the name of the file where event log state is persisted as YAML.
// maxUpdates is the maximum number of updates checkpoint will accept before
// triggering a flush to disk. interval is maximum amount of time that can
// pass since the last flush before triggering a flush to disk (minimum value
// is 1s).
func NewCheckpoint(file string, maxUpdates int, interval time.Duration) (*Checkpoint, error) {
	c := &Checkpoint{
		done:          make(chan struct{}),
		file:          file,
		maxUpdates:    maxUpdates,
		flushInterval: interval,
		sort:          make([]string, 0, 10),
		states:        make(map[string]EventLogState),
		save:          make(chan EventLogState, 1),
	}

	// Minimum batch size.
	if c.maxUpdates < 1 {
		c.maxUpdates = 1
	}

	// Minimum flush interval.
	if c.flushInterval < time.Second {
		c.flushInterval = time.Second
	}

	// Read existing state information:
	ps, err := c.read()
	if err != nil {
		return nil, err
	}

	if ps != nil {
		for _, state := range ps.States {
			c.states[state.Name] = state
		}
	}

	// Write the state file to verify we have have permissions.
	err = c.flush()
	if err != nil {
		return nil, err
	}

	c.wg.Add(1)
	go c.run()
	return c, nil
}

// run is worker loop that reads incoming state information from the save
// channel and persists it when the number of changes reaches maxEvents or
// the amount of time since the last disk write reaches flushInterval.
func (c *Checkpoint) run() {
	defer c.wg.Done()

	flushTimer := time.NewTimer(c.flushInterval)
	defer flushTimer.Stop()
loop:
	for {
		select {
		case <-c.done:
			break loop
		case s := <-c.save:
			c.lock.Lock()
			c.states[s.Name] = s
			c.lock.Unlock()
			c.numUpdates++
			if c.numUpdates < c.maxUpdates {
				continue
			}
		case <-flushTimer.C:
		}

		c.persist()
		flushTimer.Reset(c.flushInterval)
	}

	c.persist()
}

// Shutdown stops the checkpoint worker (which persists any state to disk as
// it stops). This method blocks until the checkpoint worker shutdowns. Calling
// this method more once is safe and has no effect.
func (c *Checkpoint) Shutdown() {
	c.once.Do(func() {
		close(c.done)
		c.wg.Wait()
	})
}

// States returns the current in-memory event log state. This state information
// is bootstrapped with any data found on disk at creation time.
func (c *Checkpoint) States() map[string]EventLogState {
	c.lock.RLock()
	defer c.lock.RUnlock()

	copy := make(map[string]EventLogState)
	for k, v := range c.states {
		copy[k] = v
	}

	return copy
}

// Persist queues the given event log state information to be written to disk.
func (c *Checkpoint) Persist(name string, recordNumber uint64, ts time.Time) {
	c.save <- EventLogState{
		Name:         name,
		RecordNumber: recordNumber,
		Timestamp:    ts,
	}
}

// persist writes the current state to disk if the in-memory state is dirty.
func (c *Checkpoint) persist() bool {
	if c.numUpdates == 0 {
		return false
	}

	err := c.flush()
	if err != nil {
		logp.Err("%v", err)
		return false
	}

	logp.Debug("checkpoint", "Checkpoint saved to disk. numUpdates=%d",
		c.numUpdates)
	c.numUpdates = 0
	return true
}

// flush writes the current state to disk.
func (c *Checkpoint) flush() error {
	tempFile := c.file + ".new"
	file, err := os.Create(tempFile)
	if os.IsNotExist(err) {
		// Try to create directory if it does not exist.
		if createDirErr := c.createDir(); createDirErr == nil {
			file, err = os.Create(tempFile)
		}
	}

	if err != nil {
		return fmt.Errorf("Failed to flush state to disk. %v", err)
	}

	// Sort persisted eventLogs by name.
	c.sort = c.sort[:0]
	for k := range c.states {
		c.sort = append(c.sort, k)
	}
	sort.Strings(c.sort)

	ps := PersistedState{
		UpdateTime: time.Now().UTC(),
		States:     make([]EventLogState, len(c.sort)),
	}
	for i, name := range c.sort {
		ps.States[i] = c.states[name]
	}

	data, err := yaml.Marshal(ps)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not marshal "+
			"data to YAML. %v", err)
	}

	_, err = file.Write(data)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not write to "+
			"%s. %v", tempFile, err)
	}

	file.Close()
	err = os.Rename(tempFile, c.file)
	return err
}

// read loads the persisted state from disk. If the file does not exists then
// the method returns nil and no error.
func (c *Checkpoint) read() (*PersistedState, error) {
	contents, err := ioutil.ReadFile(c.file)
	if err != nil {
		if os.IsNotExist(err) {
			err = nil
		}
		return nil, err
	}

	ps := &PersistedState{}
	err = yaml.Unmarshal(contents, ps)
	if err != nil {
		return nil, err
	}

	return ps, nil
}

// createDir creates the directory in which the state file will reside if the
// directory does not already exist.
func (c *Checkpoint) createDir() error {
	dir := filepath.Dir(c.file)
	logp.Info("Creating %s if it does not exist.", dir)
	return os.MkdirAll(dir, os.FileMode(0750))
}
<commit_msg>Change the Winlogbeat checkpoint tags to yaml from config since its not using ucfg for marshalling<commit_after>// Package checkpoint persists event log state information to disk so that
// event log monitoring can resume from the last read event in the case of a
// restart or unexpected interruption.
package checkpoint

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"github.com/elastic/beats/libbeat/logp"
	"gopkg.in/yaml.v2"
)

// Checkpoint persists event log state information to disk.
type Checkpoint struct {
	wg            sync.WaitGroup // WaitGroup used to wait on the shutdown of the checkpoint worker.
	done          chan struct{}  // Channel for shutting down the checkpoint worker.
	once          sync.Once      // Used to guarantee shutdown happens once.
	file          string         // File where the state is persisted.
	numUpdates    int            // Number of updates received since last persisting to disk.
	maxUpdates    int            // Maximum number of updates to buffer before persisting to disk.
	flushInterval time.Duration  // Maximum time interval that can pass before persisting to disk.
	sort          []string       // Slice used for sorting states map (store to save on mallocs).

	lock   sync.RWMutex
	states map[string]EventLogState

	save chan EventLogState
}

// PersistedState represents the format of the data persisted to disk.
type PersistedState struct {
	UpdateTime time.Time       `yaml:"update_time"`
	States     []EventLogState `yaml:"event_logs"`
}

// EventLogState represents the state of an individual event log.
type EventLogState struct {
	Name         string    `yaml:"name"`
	RecordNumber uint64    `yaml:"record_number"`
	Timestamp    time.Time `yaml:"timestamp"`
}

// NewCheckpoint creates and returns a new Checkpoint. This method loads state
// information from disk if it exists and starts a goroutine for persisting
// state information to disk. Shutdown should be called when finished to
// guarantee any in-memory state information is flushed to disk.
//
// file is the name of the file where event log state is persisted as YAML.
// maxUpdates is the maximum number of updates checkpoint will accept before
// triggering a flush to disk. interval is maximum amount of time that can
// pass since the last flush before triggering a flush to disk (minimum value
// is 1s).
func NewCheckpoint(file string, maxUpdates int, interval time.Duration) (*Checkpoint, error) {
	c := &Checkpoint{
		done:          make(chan struct{}),
		file:          file,
		maxUpdates:    maxUpdates,
		flushInterval: interval,
		sort:          make([]string, 0, 10),
		states:        make(map[string]EventLogState),
		save:          make(chan EventLogState, 1),
	}

	// Minimum batch size.
	if c.maxUpdates < 1 {
		c.maxUpdates = 1
	}

	// Minimum flush interval.
	if c.flushInterval < time.Second {
		c.flushInterval = time.Second
	}

	// Read existing state information:
	ps, err := c.read()
	if err != nil {
		return nil, err
	}

	if ps != nil {
		for _, state := range ps.States {
			c.states[state.Name] = state
		}
	}

	// Write the state file to verify we have have permissions.
	err = c.flush()
	if err != nil {
		return nil, err
	}

	c.wg.Add(1)
	go c.run()
	return c, nil
}

// run is worker loop that reads incoming state information from the save
// channel and persists it when the number of changes reaches maxEvents or
// the amount of time since the last disk write reaches flushInterval.
func (c *Checkpoint) run() {
	defer c.wg.Done()

	flushTimer := time.NewTimer(c.flushInterval)
	defer flushTimer.Stop()
loop:
	for {
		select {
		case <-c.done:
			break loop
		case s := <-c.save:
			c.lock.Lock()
			c.states[s.Name] = s
			c.lock.Unlock()
			c.numUpdates++
			if c.numUpdates < c.maxUpdates {
				continue
			}
		case <-flushTimer.C:
		}

		c.persist()
		flushTimer.Reset(c.flushInterval)
	}

	c.persist()
}

// Shutdown stops the checkpoint worker (which persists any state to disk as
// it stops). This method blocks until the checkpoint worker shutdowns. Calling
// this method more once is safe and has no effect.
func (c *Checkpoint) Shutdown() {
	c.once.Do(func() {
		close(c.done)
		c.wg.Wait()
	})
}

// States returns the current in-memory event log state. This state information
// is bootstrapped with any data found on disk at creation time.
func (c *Checkpoint) States() map[string]EventLogState {
	c.lock.RLock()
	defer c.lock.RUnlock()

	copy := make(map[string]EventLogState)
	for k, v := range c.states {
		copy[k] = v
	}

	return copy
}

// Persist queues the given event log state information to be written to disk.
func (c *Checkpoint) Persist(name string, recordNumber uint64, ts time.Time) {
	c.save <- EventLogState{
		Name:         name,
		RecordNumber: recordNumber,
		Timestamp:    ts,
	}
}

// persist writes the current state to disk if the in-memory state is dirty.
func (c *Checkpoint) persist() bool {
	if c.numUpdates == 0 {
		return false
	}

	err := c.flush()
	if err != nil {
		logp.Err("%v", err)
		return false
	}

	logp.Debug("checkpoint", "Checkpoint saved to disk. numUpdates=%d",
		c.numUpdates)
	c.numUpdates = 0
	return true
}

// flush writes the current state to disk.
func (c *Checkpoint) flush() error {
	tempFile := c.file + ".new"
	file, err := os.Create(tempFile)
	if os.IsNotExist(err) {
		// Try to create directory if it does not exist.
		if createDirErr := c.createDir(); createDirErr == nil {
			file, err = os.Create(tempFile)
		}
	}

	if err != nil {
		return fmt.Errorf("Failed to flush state to disk. %v", err)
	}

	// Sort persisted eventLogs by name.
	c.sort = c.sort[:0]
	for k := range c.states {
		c.sort = append(c.sort, k)
	}
	sort.Strings(c.sort)

	ps := PersistedState{
		UpdateTime: time.Now().UTC(),
		States:     make([]EventLogState, len(c.sort)),
	}
	for i, name := range c.sort {
		ps.States[i] = c.states[name]
	}

	data, err := yaml.Marshal(ps)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not marshal "+
			"data to YAML. %v", err)
	}

	_, err = file.Write(data)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not write to "+
			"%s. %v", tempFile, err)
	}

	file.Close()
	err = os.Rename(tempFile, c.file)
	return err
}

// read loads the persisted state from disk. If the file does not exists then
// the method returns nil and no error.
func (c *Checkpoint) read() (*PersistedState, error) {
	contents, err := ioutil.ReadFile(c.file)
	if err != nil {
		if os.IsNotExist(err) {
			err = nil
		}
		return nil, err
	}

	ps := &PersistedState{}
	err = yaml.Unmarshal(contents, ps)
	if err != nil {
		return nil, err
	}

	return ps, nil
}

// createDir creates the directory in which the state file will reside if the
// directory does not already exist.
func (c *Checkpoint) createDir() error {
	dir := filepath.Dir(c.file)
	logp.Info("Creating %s if it does not exist.", dir)
	return os.MkdirAll(dir, os.FileMode(0750))
}
<|endoftext|>


<commit_before>// Package checkpoint persists event log state information to disk so that
// event log monitoring can resume from the last read event in the case of a
// restart or unexpected interruption.
package checkpoint

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"github.com/elastic/beats/libbeat/logp"
	"gopkg.in/yaml.v2"
)

// Checkpoint persists event log state information to disk.
type Checkpoint struct {
	wg            sync.WaitGroup // WaitGroup used to wait on the shutdown of the checkpoint worker.
	done          chan struct{}  // Channel for shutting down the checkpoint worker.
	once          sync.Once      // Used to guarantee shutdown happens once.
	file          string         // File where the state is persisted.
	numUpdates    int            // Number of updates received since last persisting to disk.
	maxUpdates    int            // Maximum number of updates to buffer before persisting to disk.
	flushInterval time.Duration  // Maximum time interval that can pass before persisting to disk.
	sort          []string       // Slice used for sorting states map (store to save on mallocs).

	lock   sync.RWMutex
	states map[string]EventLogState

	save chan EventLogState
}

// PersistedState represents the format of the data persisted to disk.
type PersistedState struct {
	UpdateTime time.Time       `yaml:"update_time"`
	States     []EventLogState `yaml:"event_logs"`
}

// EventLogState represents the state of an individual event log.
type EventLogState struct {
	Name         string    `yaml:"name"`
	RecordNumber uint32    `yaml:"record_number"`
	Timestamp    time.Time `yaml:"timestamp"`
}

// NewCheckpoint creates and returns a new Checkpoint. This method loads state
// information from disk if it exists and starts a goroutine for persisting
// state information to disk. Shutdown should be called when finished to
// guarantee any in-memory state information is flushed to disk.
//
// file is the name of the file where event log state is persisted as YAML.
// maxUpdates is the maximum number of updates checkpoint will accept before
// triggering a flush to disk. interval is maximum amount of time that can
// pass since the last flush before triggering a flush to disk (minimum value
// is 1s).
func NewCheckpoint(file string, maxUpdates int, interval time.Duration) (*Checkpoint, error) {
	c := &Checkpoint{
		done:          make(chan struct{}),
		file:          file,
		maxUpdates:    maxUpdates,
		flushInterval: interval,
		sort:          make([]string, 0, 10),
		states:        make(map[string]EventLogState),
		save:          make(chan EventLogState, 1),
	}

	// Minimum batch size.
	if c.maxUpdates < 1 {
		c.maxUpdates = 1
	}

	// Minimum flush interval.
	if c.flushInterval < time.Second {
		c.flushInterval = time.Second
	}

	// Read existing state information:
	ps, err := c.read()
	if err != nil {
		return nil, err
	}

	if ps != nil {
		for _, state := range ps.States {
			c.states[state.Name] = state
		}
	}

	// Write the state file to verify we have have permissions.
	err = c.flush()
	if err != nil {
		return nil, err
	}

	c.wg.Add(1)
	go c.run()
	return c, nil
}

// run is worker loop that reads incoming state information from the save
// channel and persists it when the number of changes reaches maxEvents or
// the amount of time since the last disk write reaches flushInterval.
func (c *Checkpoint) run() {
	defer c.wg.Done()

	flushTimer := time.NewTimer(c.flushInterval)
	defer flushTimer.Stop()
loop:
	for {
		select {
		case <-c.done:
			break loop
		case s := <-c.save:
			c.lock.Lock()
			c.states[s.Name] = s
			c.lock.Unlock()
			c.numUpdates++
			if c.numUpdates < c.maxUpdates {
				continue
			}
		case <-flushTimer.C:
		}

		c.persist()
		flushTimer.Reset(c.flushInterval)
	}

	c.persist()
}

// Shutdown stops the checkpoint worker (which persists any state to disk as
// it stops). This method blocks until the checkpoint worker shutdowns. Calling
// this method more once is safe and has no effect.
func (c *Checkpoint) Shutdown() {
	c.once.Do(func() {
		close(c.done)
		c.wg.Wait()
	})
}

// States returns the current in-memory event log state. This state information
// is bootstrapped with any data found on disk at creation time.
func (c *Checkpoint) States() map[string]EventLogState {
	c.lock.RLock()
	defer c.lock.RUnlock()

	copy := make(map[string]EventLogState)
	for k, v := range c.states {
		copy[k] = v
	}

	return copy
}

// Persist queues the given event log state information to be written to disk.
func (c *Checkpoint) Persist(name string, recordNumber uint32, ts time.Time) {
	c.save <- EventLogState{
		Name:         name,
		RecordNumber: recordNumber,
		Timestamp:    ts,
	}
}

// persist writes the current state to disk if the in-memory state is dirty.
func (c *Checkpoint) persist() bool {
	if c.numUpdates == 0 {
		return false
	}

	err := c.flush()
	if err != nil {
		logp.Err("%v", err)
		return false
	}

	logp.Debug("checkpoint", "Checkpoint saved to disk. numUpdates=%d",
		c.numUpdates)
	c.numUpdates = 0
	return true
}

// flush writes the current state to disk.
func (c *Checkpoint) flush() error {
	tempFile := c.file + ".new"
	file, err := os.Create(tempFile)
	if os.IsNotExist(err) {
		// Try to create directory if it does not exist.
		if createDirErr := c.createDir(); createDirErr == nil {
			file, err = os.Create(tempFile)
		}
	}

	if err != nil {
		return fmt.Errorf("Failed to flush state to disk. %v", err)
	}

	// Sort persisted eventLogs by name.
	c.sort = c.sort[:0]
	for k := range c.states {
		c.sort = append(c.sort, k)
	}
	sort.Strings(c.sort)

	ps := PersistedState{
		UpdateTime: time.Now().UTC(),
		States:     make([]EventLogState, len(c.sort)),
	}
	for i, name := range c.sort {
		ps.States[i] = c.states[name]
	}

	data, err := yaml.Marshal(ps)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not marshal "+
			"data to YAML. %v", err)
	}

	_, err = file.Write(data)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not write to "+
			"%s. %v", tempFile, err)
	}

	file.Close()
	err = os.Rename(tempFile, c.file)
	return err
}

// read loads the persisted state from disk. If the file does not exists then
// the method returns nil and no error.
func (c *Checkpoint) read() (*PersistedState, error) {
	contents, err := ioutil.ReadFile(c.file)
	if err != nil {
		if os.IsNotExist(err) {
			err = nil
		}
		return nil, err
	}

	ps := &PersistedState{}
	err = yaml.Unmarshal(contents, ps)
	if err != nil {
		return nil, err
	}

	return ps, nil
}

// createDir creates the directory in which the state file will reside if the
// directory does not already exist.
func (c *Checkpoint) createDir() error {
	dir := filepath.Dir(c.file)
	return os.MkdirAll(dir, os.FileMode(0750))
}
<commit_msg>Log a message when creating a directory.<commit_after>// Package checkpoint persists event log state information to disk so that
// event log monitoring can resume from the last read event in the case of a
// restart or unexpected interruption.
package checkpoint

import (
	"fmt"
	"io/ioutil"
	"os"
	"path/filepath"
	"sort"
	"sync"
	"time"

	"github.com/elastic/beats/libbeat/logp"
	"gopkg.in/yaml.v2"
)

// Checkpoint persists event log state information to disk.
type Checkpoint struct {
	wg            sync.WaitGroup // WaitGroup used to wait on the shutdown of the checkpoint worker.
	done          chan struct{}  // Channel for shutting down the checkpoint worker.
	once          sync.Once      // Used to guarantee shutdown happens once.
	file          string         // File where the state is persisted.
	numUpdates    int            // Number of updates received since last persisting to disk.
	maxUpdates    int            // Maximum number of updates to buffer before persisting to disk.
	flushInterval time.Duration  // Maximum time interval that can pass before persisting to disk.
	sort          []string       // Slice used for sorting states map (store to save on mallocs).

	lock   sync.RWMutex
	states map[string]EventLogState

	save chan EventLogState
}

// PersistedState represents the format of the data persisted to disk.
type PersistedState struct {
	UpdateTime time.Time       `yaml:"update_time"`
	States     []EventLogState `yaml:"event_logs"`
}

// EventLogState represents the state of an individual event log.
type EventLogState struct {
	Name         string    `yaml:"name"`
	RecordNumber uint32    `yaml:"record_number"`
	Timestamp    time.Time `yaml:"timestamp"`
}

// NewCheckpoint creates and returns a new Checkpoint. This method loads state
// information from disk if it exists and starts a goroutine for persisting
// state information to disk. Shutdown should be called when finished to
// guarantee any in-memory state information is flushed to disk.
//
// file is the name of the file where event log state is persisted as YAML.
// maxUpdates is the maximum number of updates checkpoint will accept before
// triggering a flush to disk. interval is maximum amount of time that can
// pass since the last flush before triggering a flush to disk (minimum value
// is 1s).
func NewCheckpoint(file string, maxUpdates int, interval time.Duration) (*Checkpoint, error) {
	c := &Checkpoint{
		done:          make(chan struct{}),
		file:          file,
		maxUpdates:    maxUpdates,
		flushInterval: interval,
		sort:          make([]string, 0, 10),
		states:        make(map[string]EventLogState),
		save:          make(chan EventLogState, 1),
	}

	// Minimum batch size.
	if c.maxUpdates < 1 {
		c.maxUpdates = 1
	}

	// Minimum flush interval.
	if c.flushInterval < time.Second {
		c.flushInterval = time.Second
	}

	// Read existing state information:
	ps, err := c.read()
	if err != nil {
		return nil, err
	}

	if ps != nil {
		for _, state := range ps.States {
			c.states[state.Name] = state
		}
	}

	// Write the state file to verify we have have permissions.
	err = c.flush()
	if err != nil {
		return nil, err
	}

	c.wg.Add(1)
	go c.run()
	return c, nil
}

// run is worker loop that reads incoming state information from the save
// channel and persists it when the number of changes reaches maxEvents or
// the amount of time since the last disk write reaches flushInterval.
func (c *Checkpoint) run() {
	defer c.wg.Done()

	flushTimer := time.NewTimer(c.flushInterval)
	defer flushTimer.Stop()
loop:
	for {
		select {
		case <-c.done:
			break loop
		case s := <-c.save:
			c.lock.Lock()
			c.states[s.Name] = s
			c.lock.Unlock()
			c.numUpdates++
			if c.numUpdates < c.maxUpdates {
				continue
			}
		case <-flushTimer.C:
		}

		c.persist()
		flushTimer.Reset(c.flushInterval)
	}

	c.persist()
}

// Shutdown stops the checkpoint worker (which persists any state to disk as
// it stops). This method blocks until the checkpoint worker shutdowns. Calling
// this method more once is safe and has no effect.
func (c *Checkpoint) Shutdown() {
	c.once.Do(func() {
		close(c.done)
		c.wg.Wait()
	})
}

// States returns the current in-memory event log state. This state information
// is bootstrapped with any data found on disk at creation time.
func (c *Checkpoint) States() map[string]EventLogState {
	c.lock.RLock()
	defer c.lock.RUnlock()

	copy := make(map[string]EventLogState)
	for k, v := range c.states {
		copy[k] = v
	}

	return copy
}

// Persist queues the given event log state information to be written to disk.
func (c *Checkpoint) Persist(name string, recordNumber uint32, ts time.Time) {
	c.save <- EventLogState{
		Name:         name,
		RecordNumber: recordNumber,
		Timestamp:    ts,
	}
}

// persist writes the current state to disk if the in-memory state is dirty.
func (c *Checkpoint) persist() bool {
	if c.numUpdates == 0 {
		return false
	}

	err := c.flush()
	if err != nil {
		logp.Err("%v", err)
		return false
	}

	logp.Debug("checkpoint", "Checkpoint saved to disk. numUpdates=%d",
		c.numUpdates)
	c.numUpdates = 0
	return true
}

// flush writes the current state to disk.
func (c *Checkpoint) flush() error {
	tempFile := c.file + ".new"
	file, err := os.Create(tempFile)
	if os.IsNotExist(err) {
		// Try to create directory if it does not exist.
		if createDirErr := c.createDir(); createDirErr == nil {
			file, err = os.Create(tempFile)
		}
	}

	if err != nil {
		return fmt.Errorf("Failed to flush state to disk. %v", err)
	}

	// Sort persisted eventLogs by name.
	c.sort = c.sort[:0]
	for k := range c.states {
		c.sort = append(c.sort, k)
	}
	sort.Strings(c.sort)

	ps := PersistedState{
		UpdateTime: time.Now().UTC(),
		States:     make([]EventLogState, len(c.sort)),
	}
	for i, name := range c.sort {
		ps.States[i] = c.states[name]
	}

	data, err := yaml.Marshal(ps)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not marshal "+
			"data to YAML. %v", err)
	}

	_, err = file.Write(data)
	if err != nil {
		file.Close()
		return fmt.Errorf("Failed to flush state to disk. Could not write to "+
			"%s. %v", tempFile, err)
	}

	file.Close()
	err = os.Rename(tempFile, c.file)
	return err
}

// read loads the persisted state from disk. If the file does not exists then
// the method returns nil and no error.
func (c *Checkpoint) read() (*PersistedState, error) {
	contents, err := ioutil.ReadFile(c.file)
	if err != nil {
		if os.IsNotExist(err) {
			err = nil
		}
		return nil, err
	}

	ps := &PersistedState{}
	err = yaml.Unmarshal(contents, ps)
	if err != nil {
		return nil, err
	}

	return ps, nil
}

// createDir creates the directory in which the state file will reside if the
// directory does not already exist.
func (c *Checkpoint) createDir() error {
	dir := filepath.Dir(c.file)
	logp.Info("Creating %s if it does not exist.", dir)
	return os.MkdirAll(dir, os.FileMode(0750))
}
<|endoftext|>