<commit_before>package integration_test

import (
	"testing"

	"github.com/ojongerius/go-datadog-api"
)

/* TODO: Add tests for:
* ToplistWidget
* EventStreamWidget
* ImageWidget
 */

func TestFreeTextWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.FreeTextWidget

	expected.X = 1
	expected.Y = 1
	expected.Height = 10
	expected.Width = 10
	expected.Text = "Test"
	expected.FontSize = "16"
	expected.TextAlign = "center"

	w := datadog.Widget{FreeTextWidget: expected}

	board.Widgets = append(board.Widgets, w)

	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].FreeTextWidget

	assertEquals(t, "font-size", actualWidget.FontSize, expected.FontSize)
	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "text", actualWidget.Text, expected.Text)
	assertEquals(t, "text-align", actualWidget.TextAlign, expected.TextAlign)
	assertEquals(t, "type", actualWidget.Type, expected.Type)

	cleanUpScreenboard(t, board.Id)
}

func TestTimeseriesWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.TimeseriesWidget
	expected.X = 1
	expected.Y = 1
	expected.Width = 20
	expected.Height = 30
	expected.Title = true
	expected.TitleAlign = "centre"
	expected.TitleSize = datadog.TextSize{Size: 16}
	expected.TitleText = "Test"
	expected.Timeframe = "1m"

	w := datadog.Widget{TimeseriesWidget: expected}

	board.Widgets = append(board.Widgets, w)
	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].TimeseriesWidget

	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "title", actualWidget.Title, expected.Title)
	assertEquals(t, "title-align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "title-size.size", actualWidget.TitleSize.Size, expected.TitleSize.Size)
	assertEquals(t, "title-size.auto", actualWidget.TitleSize.Auto, expected.TitleSize.Auto)
	assertEquals(t, "title-text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "type", actualWidget.Type, expected.Type)
	assertEquals(t, "timeframe", actualWidget.Timeframe, expected.Timeframe)
	assertEquals(t, "legend", actualWidget.Legend, expected.Legend)
	assertTileDefEquals(t, actualWidget.TileDef, expected.TileDef)

	cleanUpScreenboard(t, board.Id)
}

func TestQueryValueWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.QueryValueWidget
	expected.X = 1
	expected.Y = 1
	expected.Width = 20
	expected.Height = 30
	expected.Title = true
	expected.TitleAlign = "centre"
	expected.TitleSize = datadog.TextSize{Size: 16}
	expected.TitleText = "Test"
	expected.Timeframe = "1m"
	expected.TimeframeAggregator = "sum"
	expected.Aggregator = "min"
	expected.Query = "docker.containers.running"

	w := datadog.Widget{QueryValueWidget: expected}

	board.Widgets = append(board.Widgets, w)
	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].QueryValueWidget

	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "title", actualWidget.Title, expected.Title)
	assertEquals(t, "title-align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "title-size.size", actualWidget.TitleSize.Size, expected.TitleSize.Size)
	assertEquals(t, "title-size.auto", actualWidget.TitleSize.Auto, expected.TitleSize.Auto)
	assertEquals(t, "title-text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "type", actualWidget.Type, expected.Type)
	assertEquals(t, "timeframe", actualWidget.Timeframe, expected.Timeframe)
	assertEquals(t, "timeframe-aggregator", actualWidget.TimeframeAggregator, expected.TimeframeAggregator)
	assertEquals(t, "aggregator", actualWidget.Aggregator, expected.Aggregator)
	assertEquals(t, "query", actualWidget.Query, expected.Query)

	cleanUpScreenboard(t, board.Id)
}

func assertTileDefEquals(t *testing.T, actual datadog.TileDef, expected datadog.TileDef) {
	assertEquals(t, "num-events", len(actual.Events), len(expected.Events))
	assertEquals(t, "num-requests", len(actual.Requests), len(expected.Requests))
	assertEquals(t, "viz", actual.Viz, expected.Viz)

	for i, event := range actual.Events {
		assertEquals(t, "event-query", event.Query, expected.Events[i].Query)
	}

	for i, request := range actual.Requests {
		assertEquals(t, "request-query", request.Query, expected.Requests[i].Query)
		assertEquals(t, "request-type", request.Type, expected.Requests[i].Type)
	}
}

func assertEquals(t *testing.T, attribute string, a, b interface{}) {
	if a != b {
		t.Errorf("The two %s values '%v' and '%v' are not equal", attribute, a, b)
	}
}
<commit_msg>Add ToplistWidget, EventStreamWidget and ImageWidget tests.<commit_after>package integration_test

import (
	"testing"

	"github.com/ojongerius/go-datadog-api"
)

func TestToplistWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.ToplistWidget
	expected.X = 1
	expected.Y = 1
	expected.Width = 5
	expected.Height = 5
	expected.Type = "toplist"
	expected.TitleText = "foo"
	expected.TitleSize.Auto = false
	expected.TitleSize.Size = 5
	expected.TitleAlign = "center"
	expected.Title = false
	expected.Timeframe = "5m"
	expected.Legend = false
	expected.LegendSize = "5"

	w := datadog.Widget{ToplistWidget: expected}

	board.Widgets = append(board.Widgets, w)

	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].ToplistWidget

	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "title_text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "title_size", actualWidget.TitleSize, expected.TitleSize)
	assertEquals(t, "title_align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "legend", actualWidget.Legend, expected.Legend)
	assertEquals(t, "legend_size", actualWidget.LegendSize, expected.LegendSize)

	cleanUpScreenboard(t, board.Id)
}

func TestEventSteamWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.EventStreamWidget
	expected.EventSize = "1"
	expected.Width = 1
	expected.Height = 1
	expected.X = 1
	expected.Y = 1
	expected.Query = "foo"
	expected.Timeframe = "5w"
	expected.Title = false
	expected.TitleAlign = "center"
	expected.TitleSize.Auto = false
	expected.TitleSize.Size = 5
	expected.TitleText = "bar"
	expected.Type = "baz"

	w := datadog.Widget{EventStreamWidget: expected}

	board.Widgets = append(board.Widgets, w)

	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].EventStreamWidget

	assertEquals(t, "event_size", actualWidget.EventSize, expected.EventSize)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "query", actualWidget.Query, expected.Query)
	assertEquals(t, "timeframe", actualWidget.Timeframe, expected.Timeframe)
	assertEquals(t, "title", actualWidget.Title, expected.Title)
	assertEquals(t, "title_align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "title_size", actualWidget.TitleSize, expected.TitleSize)
	assertEquals(t, "title_text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "type", actualWidget.Type, expected.Type)

	cleanUpScreenboard(t, board.Id)
}

func TestImageWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.ImageWidget

	expected.Width = 1
	expected.Height = 1
	expected.X = 1
	expected.Y = 1
	expected.Title = false
	expected.TitleAlign = "center"
	expected.TitleSize.Auto = false
	expected.TitleSize.Size = 5
	expected.TitleText = "bar"
	expected.Type = "baz"
	expected.Url = "qux"
	expected.Sizing = "quuz"

	w := datadog.Widget{ImageWidget: expected}

	board.Widgets = append(board.Widgets, w)

	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].ImageWidget

	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "title", actualWidget.Title, expected.Title)
	assertEquals(t, "title_align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "title_size", actualWidget.TitleSize, expected.TitleSize)
	assertEquals(t, "title_text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "type", actualWidget.Type, expected.Type)
	assertEquals(t, "url", actualWidget.Url, expected.Url)
	assertEquals(t, "sizing", actualWidget.Sizing, expected.Sizing)

	cleanUpScreenboard(t, board.Id)
}

func TestFreeTextWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.FreeTextWidget

	expected.X = 1
	expected.Y = 1
	expected.Height = 10
	expected.Width = 10
	expected.Text = "Test"
	expected.FontSize = "16"
	expected.TextAlign = "center"

	w := datadog.Widget{FreeTextWidget: expected}

	board.Widgets = append(board.Widgets, w)

	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].FreeTextWidget

	assertEquals(t, "font-size", actualWidget.FontSize, expected.FontSize)
	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "text", actualWidget.Text, expected.Text)
	assertEquals(t, "text-align", actualWidget.TextAlign, expected.TextAlign)
	assertEquals(t, "type", actualWidget.Type, expected.Type)

	cleanUpScreenboard(t, board.Id)
}

func TestTimeseriesWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.TimeseriesWidget
	expected.X = 1
	expected.Y = 1
	expected.Width = 20
	expected.Height = 30
	expected.Title = true
	expected.TitleAlign = "centre"
	expected.TitleSize = datadog.TextSize{Size: 16}
	expected.TitleText = "Test"
	expected.Timeframe = "1m"

	w := datadog.Widget{TimeseriesWidget: expected}

	board.Widgets = append(board.Widgets, w)
	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].TimeseriesWidget

	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "title", actualWidget.Title, expected.Title)
	assertEquals(t, "title-align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "title-size.size", actualWidget.TitleSize.Size, expected.TitleSize.Size)
	assertEquals(t, "title-size.auto", actualWidget.TitleSize.Auto, expected.TitleSize.Auto)
	assertEquals(t, "title-text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "type", actualWidget.Type, expected.Type)
	assertEquals(t, "timeframe", actualWidget.Timeframe, expected.Timeframe)
	assertEquals(t, "legend", actualWidget.Legend, expected.Legend)
	assertTileDefEquals(t, actualWidget.TileDef, expected.TileDef)

	cleanUpScreenboard(t, board.Id)
}

func TestQueryValueWidget(t *testing.T) {
	board := createTestScreenboard(t)

	expected := datadog.Widget{}.QueryValueWidget
	expected.X = 1
	expected.Y = 1
	expected.Width = 20
	expected.Height = 30
	expected.Title = true
	expected.TitleAlign = "centre"
	expected.TitleSize = datadog.TextSize{Size: 16}
	expected.TitleText = "Test"
	expected.Timeframe = "1m"
	expected.TimeframeAggregator = "sum"
	expected.Aggregator = "min"
	expected.Query = "docker.containers.running"

	w := datadog.Widget{QueryValueWidget: expected}

	board.Widgets = append(board.Widgets, w)
	if err := client.UpdateScreenboard(board); err != nil {
		t.Fatalf("Updating a screenboard failed: %s", err)
	}

	actual, err := client.GetScreenboard(board.Id)
	if err != nil {
		t.Fatalf("Retreiving a screenboard failed: %s", err)
	}

	actualWidget := actual.Widgets[0].QueryValueWidget

	assertEquals(t, "height", actualWidget.Height, expected.Height)
	assertEquals(t, "width", actualWidget.Width, expected.Width)
	assertEquals(t, "x", actualWidget.X, expected.X)
	assertEquals(t, "y", actualWidget.Y, expected.Y)
	assertEquals(t, "title", actualWidget.Title, expected.Title)
	assertEquals(t, "title-align", actualWidget.TitleAlign, expected.TitleAlign)
	assertEquals(t, "title-size.size", actualWidget.TitleSize.Size, expected.TitleSize.Size)
	assertEquals(t, "title-size.auto", actualWidget.TitleSize.Auto, expected.TitleSize.Auto)
	assertEquals(t, "title-text", actualWidget.TitleText, expected.TitleText)
	assertEquals(t, "type", actualWidget.Type, expected.Type)
	assertEquals(t, "timeframe", actualWidget.Timeframe, expected.Timeframe)
	assertEquals(t, "timeframe-aggregator", actualWidget.TimeframeAggregator, expected.TimeframeAggregator)
	assertEquals(t, "aggregator", actualWidget.Aggregator, expected.Aggregator)
	assertEquals(t, "query", actualWidget.Query, expected.Query)

	cleanUpScreenboard(t, board.Id)
}

func assertTileDefEquals(t *testing.T, actual datadog.TileDef, expected datadog.TileDef) {
	assertEquals(t, "num-events", len(actual.Events), len(expected.Events))
	assertEquals(t, "num-requests", len(actual.Requests), len(expected.Requests))
	assertEquals(t, "viz", actual.Viz, expected.Viz)

	for i, event := range actual.Events {
		assertEquals(t, "event-query", event.Query, expected.Events[i].Query)
	}

	for i, request := range actual.Requests {
		assertEquals(t, "request-query", request.Query, expected.Requests[i].Query)
		assertEquals(t, "request-type", request.Type, expected.Requests[i].Type)
	}
}

func assertEquals(t *testing.T, attribute string, a, b interface{}) {
	if a != b {
		t.Errorf("The two %s values '%v' and '%v' are not equal", attribute, a, b)
	}
}
<|endoftext|>


<commit_before>package regtest

import (
	"testing"
)

const unformattedProgram = `
-- main.go --
package main
import "fmt"
func main(  ) {
	fmt.Println("Hello World.")
}
-- main.go.golden --
package main

import "fmt"

func main() {
	fmt.Println("Hello World.")
}
`

func TestFormatting(t *testing.T) {
	runner.Run(t, unformattedProgram, func(t *testing.T, env *Env) {
		env.OpenFile("main.go")
		env.FormatBuffer("main.go")
		got := env.E.BufferText("main.go")
		want := env.ReadWorkspaceFile("main.go.golden")
		if got != want {
			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
		}
	})
}

const disorganizedProgram = `
-- main.go --
package main

import (
	"fmt"
	"errors"
)
func main(  ) {
	fmt.Println(errors.New("bad"))
}
-- main.go.organized --
package main

import (
	"errors"
	"fmt"
)
func main(  ) {
	fmt.Println(errors.New("bad"))
}
-- main.go.formatted --
package main

import (
	"errors"
	"fmt"
)

func main() {
	fmt.Println(errors.New("bad"))
}
`

func TestOrganizeImports(t *testing.T) {
	runner.Run(t, disorganizedProgram, func(t *testing.T, env *Env) {
		env.OpenFile("main.go")
		env.OrganizeImports("main.go")
		got := env.E.BufferText("main.go")
		want := env.ReadWorkspaceFile("main.go.organized")
		if got != want {
			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
		}
	})
}

func TestFormattingOnSave(t *testing.T) {
	runner.Run(t, disorganizedProgram, func(t *testing.T, env *Env) {
		env.OpenFile("main.go")
		env.SaveBuffer("main.go")
		got := env.E.BufferText("main.go")
		want := env.ReadWorkspaceFile("main.go.formatted")
		if got != want {
			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
		}
	})
}
<commit_msg>internal/lsp: add a regtest for formatting one-line files<commit_after>package regtest

import (
	"testing"
)

const unformattedProgram = `
-- main.go --
package main
import "fmt"
func main(  ) {
	fmt.Println("Hello World.")
}
-- main.go.golden --
package main

import "fmt"

func main() {
	fmt.Println("Hello World.")
}
`

func TestFormatting(t *testing.T) {
	runner.Run(t, unformattedProgram, func(t *testing.T, env *Env) {
		env.OpenFile("main.go")
		env.FormatBuffer("main.go")
		got := env.E.BufferText("main.go")
		want := env.ReadWorkspaceFile("main.go.golden")
		if got != want {
			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
		}
	})
}

// this is the fixed case from #36824
const onelineProgram = `
-- a.go --
package main; func f() {}

-- a.go.formatted --
package main

func f() {}
`

func TestFormattingOneLine36824(t *testing.T) {
	runner.Run(t, onelineProgram, func(t *testing.T, env *Env) {
		env.OpenFile("a.go")
		env.FormatBuffer("a.go")
		got := env.E.BufferText("a.go")
		want := env.ReadWorkspaceFile("a.go.formatted")
		if got != want {
			t.Errorf("got\n%q wanted\n%q", got, want)
		}
	})
}

const onelineProgramA = `
-- a.go --
package x; func f() {fmt.Println()}

-- a.go.imported --
package x

import "fmt"

func f() { fmt.Println() }
`

// this is the example from #36824 done properly
// but gopls does not reformat before fixing the imports
func TestFormattingOneLineImports36824(t *testing.T) {
	runner.Run(t, onelineProgramA, func(t *testing.T, env *Env) {
		env.OpenFile("a.go")
		env.FormatBuffer("a.go")
		env.OrganizeImports("a.go")
		got := env.E.BufferText("a.go")
		want := env.ReadWorkspaceFile("a.go.imported")
		if got != want {
			t.Errorf("OneLineImports go\n%q wanted\n%q", got, want)
		}
	})
}

const disorganizedProgram = `
-- main.go --
package main

import (
	"fmt"
	"errors"
)
func main(  ) {
	fmt.Println(errors.New("bad"))
}
-- main.go.organized --
package main

import (
	"errors"
	"fmt"
)
func main(  ) {
	fmt.Println(errors.New("bad"))
}
-- main.go.formatted --
package main

import (
	"errors"
	"fmt"
)

func main() {
	fmt.Println(errors.New("bad"))
}
`

func TestOrganizeImports(t *testing.T) {
	runner.Run(t, disorganizedProgram, func(t *testing.T, env *Env) {
		env.OpenFile("main.go")
		env.OrganizeImports("main.go")
		got := env.E.BufferText("main.go")
		want := env.ReadWorkspaceFile("main.go.organized")
		if got != want {
			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
		}
	})
}

func TestFormattingOnSave(t *testing.T) {
	runner.Run(t, disorganizedProgram, func(t *testing.T, env *Env) {
		env.OpenFile("main.go")
		env.SaveBuffer("main.go")
		got := env.E.BufferText("main.go")
		want := env.ReadWorkspaceFile("main.go.formatted")
		if got != want {
			t.Errorf("\n## got formatted file:\n%s\n## want:\n%s", got, want)
		}
	})
}
<|endoftext|>


<commit_before>// Package jwriter contains a JSON writer.
package jwriter

import (
	"io"
	"strconv"
	"unicode/utf8"

	"github.com/mailru/easyjson/buffer"
)

// Writer is a JSON writer.
type Writer struct {
	Error  error
	Buffer buffer.Buffer
}

// Size returns the size of the data that was written out.
func (w *Writer) Size() int {
	return w.Buffer.Size()
}

// DumpTo outputs the data to given io.Writer, resetting the buffer.
func (w *Writer) DumpTo(out io.Writer) (written int, err error) {
	return w.Buffer.DumpTo(out)
}

// BuildBytes returns writer data as a single byte slice.
func (w *Writer) BuildBytes() ([]byte, error) {
	if w.Error != nil {
		return nil, w.Error
	}

	return w.Buffer.BuildBytes(), nil
}

// RawByte appends raw binary data to the buffer.
func (w *Writer) RawByte(c byte) {
	w.Buffer.AppendByte(c)
}

// RawByte appends raw binary data to the buffer.
func (w *Writer) RawString(s string) {
	w.Buffer.AppendString(s)
}

// RawByte appends raw binary data to the buffer or sets the error if it is given. Useful for
// calling with results of MarshalJSON-like functions.
func (w *Writer) Raw(data []byte, err error) {
	switch {
	case w.Error != nil:
		return
	case err != nil:
		w.Error = err
	case len(data) > 0:
		w.Buffer.AppendBytes(data)
	default:
		w.RawString("null")
	}
}

func (w *Writer) Uint8(n uint8) {
	w.Buffer.EnsureSpace(3)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint16(n uint16) {
	w.Buffer.EnsureSpace(5)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint32(n uint32) {
	w.Buffer.EnsureSpace(10)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint(n uint) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint64(n uint64) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
}

func (w *Writer) Int8(n int8) {
	w.Buffer.EnsureSpace(4)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int16(n int16) {
	w.Buffer.EnsureSpace(6)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int32(n int32) {
	w.Buffer.EnsureSpace(11)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int(n int) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int64(n int64) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
}

func (w *Writer) Uint8Str(n uint8) {
	w.Buffer.EnsureSpace(3)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Uint16Str(n uint16) {
	w.Buffer.EnsureSpace(5)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Uint32Str(n uint32) {
	w.Buffer.EnsureSpace(10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) UintStr(n uint) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Uint64Str(n uint64) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int8Str(n int8) {
	w.Buffer.EnsureSpace(4)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int16Str(n int16) {
	w.Buffer.EnsureSpace(6)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int32Str(n int32) {
	w.Buffer.EnsureSpace(11)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) IntStr(n int) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int64Str(n int64) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Float32(n float32) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 32)
}

func (w *Writer) Float64(n float64) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, n, 'g', -1, 64)
}

func (w *Writer) Bool(v bool) {
	w.Buffer.EnsureSpace(5)
	if v {
		w.Buffer.Buf = append(w.Buffer.Buf, "true"...)
	} else {
		w.Buffer.Buf = append(w.Buffer.Buf, "false"...)
	}
}

const chars = "0123456789abcdef"

func (w *Writer) String(s string) {
	w.Buffer.AppendByte('"')

	// Portions of the string that contain no escapes are appended as
	// byte slices.

	p := 0 // last non-escape symbol

	for i := 0; i < len(s); {
		// single-with character
		if c := s[i]; c < utf8.RuneSelf {
			var escape byte
			switch c {
			case '\t':
				escape = 't'
			case '\r':
				escape = 'r'
			case '\n':
				escape = 'n'
			case '\\':
				escape = '\\'
			case '"':
				escape = '"'
			case '<', '>':
				// do nothing
			default:
				if c >= 0x20 {
					// no escaping is required
					i++
					continue
				}
			}
			if escape != 0 {
				w.Buffer.AppendString(s[p:i])
				w.Buffer.AppendByte('\\')
				w.Buffer.AppendByte(escape)
			} else {
				w.Buffer.AppendString(s[p:i])
				w.Buffer.AppendString(`\u00`)
				w.Buffer.AppendByte(chars[c>>4])
				w.Buffer.AppendByte(chars[c&0xf])
			}
			i++
			p = i
			continue
		}

		// broken utf
		runeValue, runeWidth := utf8.DecodeRuneInString(s[i:])
		if runeValue == utf8.RuneError && runeWidth == 1 {
			w.Buffer.AppendString(s[p:i])
			w.Buffer.AppendString(`\ufffd`)
			i++
			p = i
			continue
		}

		// jsonp stuff - tab separator and line separator
		if runeValue == '\u2028' || runeValue == '\u2029' {
			w.Buffer.AppendString(s[p:i])
			w.Buffer.AppendString(`\u202`)
			w.Buffer.AppendByte(chars[runeValue&0xf])
			i += runeWidth
			p = i
			continue
		}
		i += runeWidth
	}
	w.Buffer.AppendString(s[p:])
	w.Buffer.AppendByte('"')
}
<commit_msg>Faster string marshalling.<commit_after>// Package jwriter contains a JSON writer.
package jwriter

import (
	"io"
	"strconv"
	"unicode/utf8"

	"github.com/mailru/easyjson/buffer"
)

// Writer is a JSON writer.
type Writer struct {
	Error  error
	Buffer buffer.Buffer
}

// Size returns the size of the data that was written out.
func (w *Writer) Size() int {
	return w.Buffer.Size()
}

// DumpTo outputs the data to given io.Writer, resetting the buffer.
func (w *Writer) DumpTo(out io.Writer) (written int, err error) {
	return w.Buffer.DumpTo(out)
}

// BuildBytes returns writer data as a single byte slice.
func (w *Writer) BuildBytes() ([]byte, error) {
	if w.Error != nil {
		return nil, w.Error
	}

	return w.Buffer.BuildBytes(), nil
}

// RawByte appends raw binary data to the buffer.
func (w *Writer) RawByte(c byte) {
	w.Buffer.AppendByte(c)
}

// RawByte appends raw binary data to the buffer.
func (w *Writer) RawString(s string) {
	w.Buffer.AppendString(s)
}

// RawByte appends raw binary data to the buffer or sets the error if it is given. Useful for
// calling with results of MarshalJSON-like functions.
func (w *Writer) Raw(data []byte, err error) {
	switch {
	case w.Error != nil:
		return
	case err != nil:
		w.Error = err
	case len(data) > 0:
		w.Buffer.AppendBytes(data)
	default:
		w.RawString("null")
	}
}

func (w *Writer) Uint8(n uint8) {
	w.Buffer.EnsureSpace(3)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint16(n uint16) {
	w.Buffer.EnsureSpace(5)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint32(n uint32) {
	w.Buffer.EnsureSpace(10)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint(n uint) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
}

func (w *Writer) Uint64(n uint64) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
}

func (w *Writer) Int8(n int8) {
	w.Buffer.EnsureSpace(4)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int16(n int16) {
	w.Buffer.EnsureSpace(6)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int32(n int32) {
	w.Buffer.EnsureSpace(11)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int(n int) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
}

func (w *Writer) Int64(n int64) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
}

func (w *Writer) Uint8Str(n uint8) {
	w.Buffer.EnsureSpace(3)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Uint16Str(n uint16) {
	w.Buffer.EnsureSpace(5)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Uint32Str(n uint32) {
	w.Buffer.EnsureSpace(10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) UintStr(n uint) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, uint64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Uint64Str(n uint64) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendUint(w.Buffer.Buf, n, 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int8Str(n int8) {
	w.Buffer.EnsureSpace(4)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int16Str(n int16) {
	w.Buffer.EnsureSpace(6)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int32Str(n int32) {
	w.Buffer.EnsureSpace(11)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) IntStr(n int) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, int64(n), 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Int64Str(n int64) {
	w.Buffer.EnsureSpace(21)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
	w.Buffer.Buf = strconv.AppendInt(w.Buffer.Buf, n, 10)
	w.Buffer.Buf = append(w.Buffer.Buf, '"')
}

func (w *Writer) Float32(n float32) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, float64(n), 'g', -1, 32)
}

func (w *Writer) Float64(n float64) {
	w.Buffer.EnsureSpace(20)
	w.Buffer.Buf = strconv.AppendFloat(w.Buffer.Buf, n, 'g', -1, 64)
}

func (w *Writer) Bool(v bool) {
	w.Buffer.EnsureSpace(5)
	if v {
		w.Buffer.Buf = append(w.Buffer.Buf, "true"...)
	} else {
		w.Buffer.Buf = append(w.Buffer.Buf, "false"...)
	}
}

const chars = "0123456789abcdef"

func isNotEscapedSingleChar(c byte) bool {
	// Note: might make sense to use a table if there are more chars to escape. With 4 chars
	// it benchmarks the same.
	return c != '<' && c != '\\' && c != '"' && c != '>' && c >= 0x20 && c < utf8.RuneSelf
}

func (w *Writer) String(s string) {
	w.Buffer.AppendByte('"')

	// Portions of the string that contain no escapes are appended as
	// byte slices.

	p := 0 // last non-escape symbol

	for i := 0; i < len(s); {
		c := s[i]

		if isNotEscapedSingleChar(c) {
			// single-width character, no escaping is required
			i++
			continue
		} else if c < utf8.RuneSelf {
			// single-with character, need to escape
			w.Buffer.AppendString(s[p:i])
			switch c {
			case '\t':
				w.Buffer.AppendString(`\t`)
			case '\r':
				w.Buffer.AppendString(`\r`)
			case '\n':
				w.Buffer.AppendString(`\n`)
			case '\\':
				w.Buffer.AppendString(`\\`)
			case '"':
				w.Buffer.AppendString(`\"`)
			default:
				w.Buffer.AppendString(`\u00`)
				w.Buffer.AppendByte(chars[c>>4])
				w.Buffer.AppendByte(chars[c&0xf])
			}

			i++
			p = i
			continue
		}

		// broken utf
		runeValue, runeWidth := utf8.DecodeRuneInString(s[i:])
		if runeValue == utf8.RuneError && runeWidth == 1 {
			w.Buffer.AppendString(s[p:i])
			w.Buffer.AppendString(`\ufffd`)
			i++
			p = i
			continue
		}

		// jsonp stuff - tab separator and line separator
		if runeValue == '\u2028' || runeValue == '\u2029' {
			w.Buffer.AppendString(s[p:i])
			w.Buffer.AppendString(`\u202`)
			w.Buffer.AppendByte(chars[runeValue&0xf])
			i += runeWidth
			p = i
			continue
		}
		i += runeWidth
	}
	w.Buffer.AppendString(s[p:])
	w.Buffer.AppendByte('"')
}
<|endoftext|>


<commit_before>package kafka

import (
	"hash"

	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	metrics "github.com/rcrowley/go-metrics"
)

// ConsumerBuilder creates a Kafka consumer.
type ConsumerBuilder func(brokers []string, group, clientID string) (Consumer, error)

// DefaultConsumerBuilder creates a Kafka consumer using the Sarama library.
func DefaultConsumerBuilder(brokers []string, group, clientID string) (Consumer, error) {
	config := CreateDefaultSaramaConfig(clientID, nil, metrics.DefaultRegistry)
	return NewSaramaConsumer(brokers, group, config)
}

// ConsumerBuilderWithConfig creates a Kafka consumer using the Sarama library.
func ConsumerBuilderWithConfig(config *cluster.Config) ConsumerBuilder {
	return func(brokers []string, group, clientID string) (Consumer, error) {
		config.ClientID = clientID
		return NewSaramaConsumer(brokers, group, config)
	}
}

// ProducerBuilder create a Kafka producer.
type ProducerBuilder func(brokers []string, clientID string, hasher func() hash.Hash32) (Producer, error)

// DefaultProducerBuilder creates a Kafka producer using the Sarama library.
func DefaultProducerBuilder(brokers []string, clientID string, hasher func() hash.Hash32) (Producer, error) {
	partitioner := sarama.NewCustomHashPartitioner(hasher)
	config := CreateDefaultSaramaConfig(clientID, partitioner, metrics.DefaultRegistry)
	return NewProducer(brokers, &config.Config)
}

// TopicManagerBuilder creates a TopicManager to check partition counts and
// create tables.
type TopicManagerBuilder func(brokers []string) (TopicManager, error)

// DefaultTopicManagerBuilder creates TopicManager using the Sarama library.
// This topic manager cannot create topics.
func DefaultTopicManagerBuilder(brokers []string) (TopicManager, error) {
	return NewSaramaTopicManager(brokers)
}
<commit_msg>add ProducerBuilderWithConfig<commit_after>package kafka

import (
	"hash"

	"github.com/Shopify/sarama"
	cluster "github.com/bsm/sarama-cluster"
	metrics "github.com/rcrowley/go-metrics"
)

// ConsumerBuilder creates a Kafka consumer.
type ConsumerBuilder func(brokers []string, group, clientID string) (Consumer, error)

// DefaultConsumerBuilder creates a Kafka consumer using the Sarama library.
func DefaultConsumerBuilder(brokers []string, group, clientID string) (Consumer, error) {
	config := CreateDefaultSaramaConfig(clientID, nil, metrics.DefaultRegistry)
	return NewSaramaConsumer(brokers, group, config)
}

// ConsumerBuilderWithConfig creates a Kafka consumer using the Sarama library.
func ConsumerBuilderWithConfig(config *cluster.Config) ConsumerBuilder {
	return func(brokers []string, group, clientID string) (Consumer, error) {
		config.ClientID = clientID
		return NewSaramaConsumer(brokers, group, config)
	}
}

// ProducerBuilder create a Kafka producer.
type ProducerBuilder func(brokers []string, clientID string, hasher func() hash.Hash32) (Producer, error)

// DefaultProducerBuilder creates a Kafka producer using the Sarama library.
func DefaultProducerBuilder(brokers []string, clientID string, hasher func() hash.Hash32) (Producer, error) {
	partitioner := sarama.NewCustomHashPartitioner(hasher)
	config := CreateDefaultSaramaConfig(clientID, partitioner, metrics.DefaultRegistry)
	return NewProducer(brokers, &config.Config)
}

// ProducerBuilderWithConfig creates a Kafka consumer using the Sarama library.
func ProducerBuilderWithConfig(config *cluster.Config) ProducerBuilder {
	return func(brokers []string, clientID string, hasher func() hash.Hash32) (Producer, error) {
		partitioner := sarama.NewCustomHashPartitioner(hasher)
		config.ClientID = clientID
		config.Producer.Partitioner = partitioner
		return NewProducer(brokers, &config.Config)
	}
}

// TopicManagerBuilder creates a TopicManager to check partition counts and
// create tables.
type TopicManagerBuilder func(brokers []string) (TopicManager, error)

// DefaultTopicManagerBuilder creates TopicManager using the Sarama library.
// This topic manager cannot create topics.
func DefaultTopicManagerBuilder(brokers []string) (TopicManager, error) {
	return NewSaramaTopicManager(brokers)
}
<|endoftext|>


<commit_before>package store

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

	etcd "github.com/coreos/go-etcd/etcd"
)

// Etcd embeds the client
type Etcd struct {
	client       *etcd.Client
	ephemeralTTL time.Duration
}

type etcdLock struct {
	client   *etcd.Client
	stopLock chan struct{}
	key      string
	value    string
	ttl      uint64
}

const (
	defaultLockTTL    = 20 * time.Second
	defaultUpdateTime = 5 * time.Second
)

// InitializeEtcd creates a new Etcd client given
// a list of endpoints and optional tls config
func InitializeEtcd(addrs []string, options *Config) (Store, error) {
	s := &Etcd{}

	entries := createEndpoints(addrs, "http")
	s.client = etcd.NewClient(entries)

	// Set options
	if options != nil {
		if options.TLS != nil {
			s.setTLS(options.TLS)
		}
		if options.ConnectionTimeout != 0 {
			s.setTimeout(options.ConnectionTimeout)
		}
		if options.EphemeralTTL != 0 {
			s.setEphemeralTTL(options.EphemeralTTL)
		}
	}

	return s, nil
}

// SetTLS sets the tls configuration given the path
// of certificate files
func (s *Etcd) setTLS(tls *tls.Config) {
	// Change to https scheme
	var addrs []string
	entries := s.client.GetCluster()
	for _, entry := range entries {
		addrs = append(addrs, strings.Replace(entry, "http", "https", -1))
	}
	s.client.SetCluster(addrs)

	// Set transport
	t := http.Transport{
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second, // default timeout
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig:     tls,
	}
	s.client.SetTransport(&t)
}

// SetTimeout sets the timeout used for connecting to the store
func (s *Etcd) setTimeout(time time.Duration) {
	s.client.SetDialTimeout(time)
}

// SetHeartbeat sets the heartbeat value to notify we are alive
func (s *Etcd) setEphemeralTTL(time time.Duration) {
	s.ephemeralTTL = time
}

// Create the entire path for a directory that does not exist
func (s *Etcd) createDirectory(path string) error {
	if _, err := s.client.CreateDir(normalize(path), 10); err != nil {
		if etcdError, ok := err.(*etcd.EtcdError); ok {
			if etcdError.ErrorCode != 105 { // Skip key already exists
				return err
			}
		} else {
			return err
		}
	}
	return nil
}

// Get the value at "key", returns the last modified index
// to use in conjunction to CAS calls
func (s *Etcd) Get(key string) (*KVPair, error) {
	result, err := s.client.Get(normalize(key), false, false)
	if err != nil {
		if etcdError, ok := err.(*etcd.EtcdError); ok {
			// Not a Directory or Not a file
			if etcdError.ErrorCode == 102 || etcdError.ErrorCode == 104 {
				return nil, ErrKeyNotFound
			}
		}
		return nil, err
	}
	return &KVPair{result.Node.Key, []byte(result.Node.Value), result.Node.ModifiedIndex}, nil
}

// Put a value at "key"
func (s *Etcd) Put(key string, value []byte, opts *WriteOptions) error {

	// Default TTL = 0 means no expiration
	var ttl uint64
	if opts != nil && opts.Ephemeral {
		ttl = uint64(s.ephemeralTTL.Seconds())
	}

	if _, err := s.client.Set(key, string(value), ttl); err != nil {
		if etcdError, ok := err.(*etcd.EtcdError); ok {
			if etcdError.ErrorCode == 104 { // Not a directory
				// Remove the last element (the actual key) and set the prefix as a dir
				err = s.createDirectory(getDirectory(key))
				if _, err := s.client.Set(key, string(value), ttl); err != nil {
					return err
				}
			}
		}
		return err
	}
	return nil
}

// Delete a value at "key"
func (s *Etcd) Delete(key string) error {
	if _, err := s.client.Delete(normalize(key), false); err != nil {
		return err
	}
	return nil
}

// Exists checks if the key exists inside the store
func (s *Etcd) Exists(key string) (bool, error) {
	entry, err := s.Get(key)
	if err != nil {
		if err == ErrKeyNotFound || entry.Value == nil {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// Watch changes on a key.
// Returns a channel that will receive changes or an error.
// Upon creating a watch, the current value will be sent to the channel.
// Providing a non-nil stopCh can be used to stop watching.
func (s *Etcd) Watch(key string, stopCh <-chan struct{}) (<-chan *KVPair, error) {
	key = normalize(key)

	// Get the current value
	current, err := s.Get(key)
	if err != nil {
		return nil, err
	}

	// Start an etcd watch.
	// Note: etcd will send the current value through the channel.
	etcdWatchCh := make(chan *etcd.Response)
	etcdStopCh := make(chan bool)
	go s.client.Watch(key, 0, false, etcdWatchCh, etcdStopCh)

	// Adapter goroutine: The goal here is to convert wathever format etcd is
	// using into our interface.
	watchCh := make(chan *KVPair)
	go func() {
		defer close(watchCh)

		// Push the current value through the channel.
		watchCh <- current

		for {
			select {
			case result := <-etcdWatchCh:
				watchCh <- &KVPair{
					result.Node.Key,
					[]byte(result.Node.Value),
					result.Node.ModifiedIndex,
				}
			case <-stopCh:
				etcdStopCh <- true
				return
			}
		}
	}()
	return watchCh, nil
}

// WatchTree watches changes on a "directory"
// Returns a channel that will receive changes or an error.
// Upon creating a watch, the current value will be sent to the channel.
// Providing a non-nil stopCh can be used to stop watching.
func (s *Etcd) WatchTree(prefix string, stopCh <-chan struct{}) (<-chan []*KVPair, error) {
	prefix = normalize(prefix)

	// Get the current value
	current, err := s.List(prefix)
	if err != nil {
		return nil, err
	}

	// Start an etcd watch.
	etcdWatchCh := make(chan *etcd.Response)
	etcdStopCh := make(chan bool)
	go s.client.Watch(prefix, 0, true, etcdWatchCh, etcdStopCh)

	// Adapter goroutine: The goal here is to convert wathever format etcd is
	// using into our interface.
	watchCh := make(chan []*KVPair)
	go func() {
		defer close(watchCh)

		// Push the current value through the channel.
		watchCh <- current

		for {
			select {
			case <-etcdWatchCh:
				// FIXME: We should probably use the value pushed by the channel.
				// However, .Node.Nodes seems to be empty.
				if list, err := s.List(prefix); err == nil {
					watchCh <- list
				}
			case <-stopCh:
				etcdStopCh <- true
				return
			}
		}
	}()
	return watchCh, nil
}

// AtomicPut put a value at "key" if the key has not been
// modified in the meantime, throws an error if this is the case
func (s *Etcd) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, error) {
	_, err := s.client.CompareAndSwap(normalize(key), string(value), 0, "", previous.LastIndex)
	if err != nil {
		return false, err
	}
	return true, nil
}

// AtomicDelete deletes a value at "key" if the key has not
// been modified in the meantime, throws an error if this is the case
func (s *Etcd) AtomicDelete(key string, previous *KVPair) (bool, error) {
	_, err := s.client.CompareAndDelete(normalize(key), "", previous.LastIndex)
	if err != nil {
		return false, err
	}
	return true, nil
}

// List the content of a given prefix
func (s *Etcd) List(prefix string) ([]*KVPair, error) {
	resp, err := s.client.Get(normalize(prefix), true, true)
	if err != nil {
		return nil, err
	}
	kv := []*KVPair{}
	for _, n := range resp.Node.Nodes {
		kv = append(kv, &KVPair{n.Key, []byte(n.Value), n.ModifiedIndex})
	}
	return kv, nil
}

// DeleteTree deletes a range of keys based on prefix
func (s *Etcd) DeleteTree(prefix string) error {
	if _, err := s.client.Delete(normalize(prefix), true); err != nil {
		return err
	}
	return nil
}

// NewLock returns a handle to a lock struct which can be used to acquire and
// release the mutex.
func (s *Etcd) NewLock(key string, options *LockOptions) (Locker, error) {
	var value string
	ttl := uint64(time.Duration(defaultLockTTL).Seconds())

	// Apply options
	if options != nil {
		if options.Value != nil {
			value = string(options.Value)
		}
		if options.TTL != 0 {
			ttl = uint64(options.TTL.Seconds())
		}
	}

	// Create lock object
	lock := &etcdLock{
		client: s.client,
		key:    key,
		value:  value,
		ttl:    ttl,
	}

	return lock, nil
}

// Lock attempts to acquire the lock and blocks while doing so.
// Returns a channel that is closed if our lock is lost or an error.
func (l *etcdLock) Lock() (<-chan struct{}, error) {

	key := normalize(l.key)
	var lastIndex uint64

	// Lock holder channels
	lockHeld := make(chan struct{})
	stopLocking := make(chan struct{})

	for {
		resp, err := l.client.Create(key, l.value, l.ttl)
		if err != nil {
			if etcdError, ok := err.(*etcd.EtcdError); ok {
				// Key already exists
				if etcdError.ErrorCode != 105 {
					lastIndex = ^uint64(0)
				}
			}
		} else {
			lastIndex = resp.Node.ModifiedIndex
		}

		_, err = l.client.CompareAndSwap(key, l.value, l.ttl, "", lastIndex)

		if err == nil {
			// Leader section
			l.stopLock = stopLocking
			go l.holdLock(key, lockHeld, stopLocking)
			break
		} else {
			// Seeker section
			chW := make(chan *etcd.Response)
			chWStop := make(chan bool)
			l.waitLock(key, chW, chWStop)

			// Delete or Expire event occured
			// Retry
		}
	}

	return lockHeld, nil
}

// Hold the lock as long as we can
// Updates the key ttl periodically until we receive
// an explicit stop signal from the Unlock method
func (l *etcdLock) holdLock(key string, lockHeld chan struct{}, stopLocking chan struct{}) {
	defer close(lockHeld)

	update := time.NewTicker(defaultUpdateTime)

	for {
		select {
		case <-update.C:
			_, err := l.client.Update(key, l.value, l.ttl)
			if err != nil {
				return
			}

		case <-stopLocking:
			update.Stop()
			return
		}
	}
}

// Simply wait for the key to be available
func (l *etcdLock) waitLock(key string, eventCh chan *etcd.Response, stopWatchCh chan bool) {
	go l.client.Watch(key, 0, false, eventCh, stopWatchCh)
	for event := range eventCh {
		if event.Action == "delete" || event.Action == "expire" {
			return
		}
	}
}

// Unlock released the lock. It is an error to call this
// if the lock is not currently held.
func (l *etcdLock) Unlock() error {
	if l.stopLock != nil {
		close(l.stopLock)
	}
	_, err := l.client.Delete(normalize(l.key), false)
	if err != nil {
		return err
	}
	return nil
}
<commit_msg>Fix etcd Unlock that allowed other processes (non lock-holder) to unlock the key<commit_after>package store

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

	etcd "github.com/coreos/go-etcd/etcd"
)

// Etcd embeds the client
type Etcd struct {
	client       *etcd.Client
	ephemeralTTL time.Duration
}

type etcdLock struct {
	client   *etcd.Client
	stopLock chan struct{}
	key      string
	value    string
	last     *etcd.Response
	ttl      uint64
}

const (
	defaultLockTTL    = 20 * time.Second
	defaultUpdateTime = 5 * time.Second
)

// InitializeEtcd creates a new Etcd client given
// a list of endpoints and optional tls config
func InitializeEtcd(addrs []string, options *Config) (Store, error) {
	s := &Etcd{}

	entries := createEndpoints(addrs, "http")
	s.client = etcd.NewClient(entries)

	// Set options
	if options != nil {
		if options.TLS != nil {
			s.setTLS(options.TLS)
		}
		if options.ConnectionTimeout != 0 {
			s.setTimeout(options.ConnectionTimeout)
		}
		if options.EphemeralTTL != 0 {
			s.setEphemeralTTL(options.EphemeralTTL)
		}
	}

	return s, nil
}

// SetTLS sets the tls configuration given the path
// of certificate files
func (s *Etcd) setTLS(tls *tls.Config) {
	// Change to https scheme
	var addrs []string
	entries := s.client.GetCluster()
	for _, entry := range entries {
		addrs = append(addrs, strings.Replace(entry, "http", "https", -1))
	}
	s.client.SetCluster(addrs)

	// Set transport
	t := http.Transport{
		Dial: (&net.Dialer{
			Timeout:   30 * time.Second, // default timeout
			KeepAlive: 30 * time.Second,
		}).Dial,
		TLSHandshakeTimeout: 10 * time.Second,
		TLSClientConfig:     tls,
	}
	s.client.SetTransport(&t)
}

// SetTimeout sets the timeout used for connecting to the store
func (s *Etcd) setTimeout(time time.Duration) {
	s.client.SetDialTimeout(time)
}

// SetHeartbeat sets the heartbeat value to notify we are alive
func (s *Etcd) setEphemeralTTL(time time.Duration) {
	s.ephemeralTTL = time
}

// Create the entire path for a directory that does not exist
func (s *Etcd) createDirectory(path string) error {
	if _, err := s.client.CreateDir(normalize(path), 10); err != nil {
		if etcdError, ok := err.(*etcd.EtcdError); ok {
			if etcdError.ErrorCode != 105 { // Skip key already exists
				return err
			}
		} else {
			return err
		}
	}
	return nil
}

// Get the value at "key", returns the last modified index
// to use in conjunction to CAS calls
func (s *Etcd) Get(key string) (*KVPair, error) {
	result, err := s.client.Get(normalize(key), false, false)
	if err != nil {
		if etcdError, ok := err.(*etcd.EtcdError); ok {
			// Not a Directory or Not a file
			if etcdError.ErrorCode == 102 || etcdError.ErrorCode == 104 {
				return nil, ErrKeyNotFound
			}
		}
		return nil, err
	}
	return &KVPair{result.Node.Key, []byte(result.Node.Value), result.Node.ModifiedIndex}, nil
}

// Put a value at "key"
func (s *Etcd) Put(key string, value []byte, opts *WriteOptions) error {

	// Default TTL = 0 means no expiration
	var ttl uint64
	if opts != nil && opts.Ephemeral {
		ttl = uint64(s.ephemeralTTL.Seconds())
	}

	if _, err := s.client.Set(key, string(value), ttl); err != nil {
		if etcdError, ok := err.(*etcd.EtcdError); ok {
			if etcdError.ErrorCode == 104 { // Not a directory
				// Remove the last element (the actual key) and set the prefix as a dir
				err = s.createDirectory(getDirectory(key))
				if _, err := s.client.Set(key, string(value), ttl); err != nil {
					return err
				}
			}
		}
		return err
	}
	return nil
}

// Delete a value at "key"
func (s *Etcd) Delete(key string) error {
	if _, err := s.client.Delete(normalize(key), false); err != nil {
		return err
	}
	return nil
}

// Exists checks if the key exists inside the store
func (s *Etcd) Exists(key string) (bool, error) {
	entry, err := s.Get(key)
	if err != nil {
		if err == ErrKeyNotFound || entry.Value == nil {
			return false, nil
		}
		return false, err
	}
	return true, nil
}

// Watch changes on a key.
// Returns a channel that will receive changes or an error.
// Upon creating a watch, the current value will be sent to the channel.
// Providing a non-nil stopCh can be used to stop watching.
func (s *Etcd) Watch(key string, stopCh <-chan struct{}) (<-chan *KVPair, error) {
	key = normalize(key)

	// Get the current value
	current, err := s.Get(key)
	if err != nil {
		return nil, err
	}

	// Start an etcd watch.
	// Note: etcd will send the current value through the channel.
	etcdWatchCh := make(chan *etcd.Response)
	etcdStopCh := make(chan bool)
	go s.client.Watch(key, 0, false, etcdWatchCh, etcdStopCh)

	// Adapter goroutine: The goal here is to convert wathever format etcd is
	// using into our interface.
	watchCh := make(chan *KVPair)
	go func() {
		defer close(watchCh)

		// Push the current value through the channel.
		watchCh <- current

		for {
			select {
			case result := <-etcdWatchCh:
				watchCh <- &KVPair{
					result.Node.Key,
					[]byte(result.Node.Value),
					result.Node.ModifiedIndex,
				}
			case <-stopCh:
				etcdStopCh <- true
				return
			}
		}
	}()
	return watchCh, nil
}

// WatchTree watches changes on a "directory"
// Returns a channel that will receive changes or an error.
// Upon creating a watch, the current value will be sent to the channel.
// Providing a non-nil stopCh can be used to stop watching.
func (s *Etcd) WatchTree(prefix string, stopCh <-chan struct{}) (<-chan []*KVPair, error) {
	prefix = normalize(prefix)

	// Get the current value
	current, err := s.List(prefix)
	if err != nil {
		return nil, err
	}

	// Start an etcd watch.
	etcdWatchCh := make(chan *etcd.Response)
	etcdStopCh := make(chan bool)
	go s.client.Watch(prefix, 0, true, etcdWatchCh, etcdStopCh)

	// Adapter goroutine: The goal here is to convert wathever format etcd is
	// using into our interface.
	watchCh := make(chan []*KVPair)
	go func() {
		defer close(watchCh)

		// Push the current value through the channel.
		watchCh <- current

		for {
			select {
			case <-etcdWatchCh:
				// FIXME: We should probably use the value pushed by the channel.
				// However, .Node.Nodes seems to be empty.
				if list, err := s.List(prefix); err == nil {
					watchCh <- list
				}
			case <-stopCh:
				etcdStopCh <- true
				return
			}
		}
	}()
	return watchCh, nil
}

// AtomicPut put a value at "key" if the key has not been
// modified in the meantime, throws an error if this is the case
func (s *Etcd) AtomicPut(key string, value []byte, previous *KVPair, options *WriteOptions) (bool, error) {
	_, err := s.client.CompareAndSwap(normalize(key), string(value), 0, "", previous.LastIndex)
	if err != nil {
		return false, err
	}
	return true, nil
}

// AtomicDelete deletes a value at "key" if the key has not
// been modified in the meantime, throws an error if this is the case
func (s *Etcd) AtomicDelete(key string, previous *KVPair) (bool, error) {
	_, err := s.client.CompareAndDelete(normalize(key), "", previous.LastIndex)
	if err != nil {
		return false, err
	}
	return true, nil
}

// List the content of a given prefix
func (s *Etcd) List(prefix string) ([]*KVPair, error) {
	resp, err := s.client.Get(normalize(prefix), true, true)
	if err != nil {
		return nil, err
	}
	kv := []*KVPair{}
	for _, n := range resp.Node.Nodes {
		kv = append(kv, &KVPair{n.Key, []byte(n.Value), n.ModifiedIndex})
	}
	return kv, nil
}

// DeleteTree deletes a range of keys based on prefix
func (s *Etcd) DeleteTree(prefix string) error {
	if _, err := s.client.Delete(normalize(prefix), true); err != nil {
		return err
	}
	return nil
}

// NewLock returns a handle to a lock struct which can be used to acquire and
// release the mutex.
func (s *Etcd) NewLock(key string, options *LockOptions) (Locker, error) {
	var value string
	ttl := uint64(time.Duration(defaultLockTTL).Seconds())

	// Apply options
	if options != nil {
		if options.Value != nil {
			value = string(options.Value)
		}
		if options.TTL != 0 {
			ttl = uint64(options.TTL.Seconds())
		}
	}

	// Create lock object
	lock := &etcdLock{
		client: s.client,
		key:    key,
		value:  value,
		ttl:    ttl,
	}

	return lock, nil
}

// Lock attempts to acquire the lock and blocks while doing so.
// Returns a channel that is closed if our lock is lost or an error.
func (l *etcdLock) Lock() (<-chan struct{}, error) {

	key := normalize(l.key)

	// Lock holder channels
	lockHeld := make(chan struct{})
	stopLocking := make(chan struct{})

	var lastIndex uint64

	for {
		resp, err := l.client.Create(key, l.value, l.ttl)
		if err != nil {
			if etcdError, ok := err.(*etcd.EtcdError); ok {
				// Key already exists
				if etcdError.ErrorCode != 105 {
					lastIndex = ^uint64(0)
				}
			}
		} else {
			lastIndex = resp.Node.ModifiedIndex
		}

		_, err = l.client.CompareAndSwap(key, l.value, l.ttl, "", lastIndex)

		if err == nil {
			// Leader section
			l.stopLock = stopLocking
			go l.holdLock(key, lockHeld, stopLocking)
			break
		} else {
			// Seeker section
			chW := make(chan *etcd.Response)
			chWStop := make(chan bool)
			l.waitLock(key, chW, chWStop)

			// Delete or Expire event occured
			// Retry
		}
	}

	return lockHeld, nil
}

// Hold the lock as long as we can
// Updates the key ttl periodically until we receive
// an explicit stop signal from the Unlock method
func (l *etcdLock) holdLock(key string, lockHeld chan struct{}, stopLocking chan struct{}) {
	defer close(lockHeld)

	update := time.NewTicker(defaultUpdateTime)
	defer update.Stop()

	var err error

	for {
		select {
		case <-update.C:
			l.last, err = l.client.Update(key, l.value, l.ttl)
			if err != nil {
				return
			}

		case <-stopLocking:
			return
		}
	}
}

// WaitLock simply waits for the key to be available for creation
func (l *etcdLock) waitLock(key string, eventCh chan *etcd.Response, stopWatchCh chan bool) {
	go l.client.Watch(key, 0, false, eventCh, stopWatchCh)
	for event := range eventCh {
		if event.Action == "delete" || event.Action == "expire" {
			return
		}
	}
}

// Unlock released the lock. It is an error to call this
// if the lock is not currently held.
func (l *etcdLock) Unlock() error {
	if l.stopLock != nil {
		l.stopLock <- struct{}{}
	}
	if l.last != nil {
		_, err := l.client.CompareAndDelete(normalize(l.key), l.value, l.last.Node.ModifiedIndex)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package client

import (
	"fmt"
	"sync"
	"time"

	"gopkg.in/jcmturner/gokrb5.v6/iana/nametype"
	"gopkg.in/jcmturner/gokrb5.v6/krberror"
	"gopkg.in/jcmturner/gokrb5.v6/messages"
	"gopkg.in/jcmturner/gokrb5.v6/types"
)

// sessions hold TGTs and are keyed on the realm name
type sessions struct {
	Entries map[string]*session
	mux     sync.RWMutex
}

// destroy erases all sessions
func (s *sessions) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	for k, e := range s.Entries {
		e.destroy()
		delete(s.Entries, k)
	}
}

// update replaces a session with the one provided or adds it as a new one
func (s *sessions) update(sess *session) {
	s.mux.Lock()
	defer s.mux.Unlock()
	// if a session already exists for this, cancel its auto renew.
	if i, ok := s.Entries[sess.realm]; ok {
		if i != sess {
			// Session in the sessions cache is not the same as one provided.
			// Cancel the one in the cache and add this one.
			i.mux.Lock()
			i.cancel <- true
			i.mux.Unlock()
			s.Entries[sess.realm] = sess
			return
		}
	}
	// No session for this realm was found so just add it
	s.Entries[sess.realm] = sess
}

// get returns the session for the realm specified
func (s *sessions) get(realm string) (*session, bool) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	sess, ok := s.Entries[realm]
	return sess, ok
}

// session holds the TGT details for a realm
type session struct {
	realm                string
	authTime             time.Time
	endTime              time.Time
	renewTill            time.Time
	tgt                  messages.Ticket
	sessionKey           types.EncryptionKey
	sessionKeyExpiration time.Time
	cancel               chan bool
	mux                  sync.RWMutex
}

// AddSession adds a session for a realm with a TGT to the client's session cache.
// A goroutine is started to automatically renew the TGT before expiry.
func (cl *Client) AddSession(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	realm := cl.spnRealm(tgt.SName)
	s := &session{
		realm:                realm,
		authTime:             dep.AuthTime,
		endTime:              dep.EndTime,
		renewTill:            dep.RenewTill,
		tgt:                  tgt,
		sessionKey:           dep.Key,
		sessionKeyExpiration: dep.KeyExpiration,
	}
	cl.sessions.update(s)
	cl.enableAutoSessionRenewal(s)
}

// update overwrites the session details with those from the TGT and decrypted encPart
func (s *session) update(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	s.mux.Lock()
	defer s.mux.Unlock()
	s.authTime = dep.AuthTime
	s.endTime = dep.EndTime
	s.renewTill = dep.RenewTill
	s.tgt = tgt
	s.sessionKey = dep.Key
	s.sessionKeyExpiration = dep.KeyExpiration
}

// destroy will cancel any auto renewal of the session and set the expiration times to the current time
func (s *session) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	if s.cancel != nil {
		s.cancel <- true
	}
	s.endTime = time.Now().UTC()
	s.renewTill = s.endTime
	s.sessionKeyExpiration = s.endTime
}

// valid informs if the TGT is still within the valid time window
func (s *session) valid() bool {
	s.mux.RLock()
	defer s.mux.RUnlock()
	t := time.Now().UTC()
	if t.Before(s.endTime) && s.authTime.Before(t) {
		return true
	}
	return false
}

// tgtDetails is a thread safe way to get the session's realm, TGT and session key values
func (s *session) tgtDetails() (string, messages.Ticket, types.EncryptionKey) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.tgt, s.sessionKey
}

// timeDetails is a thread safe way to get the session's validity time values
func (s *session) timeDetails() (string, time.Time, time.Time, time.Time, time.Time) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.authTime, s.endTime, s.renewTill, s.sessionKeyExpiration
}

// enableAutoSessionRenewal turns on the automatic renewal for the client's TGT session.
func (cl *Client) enableAutoSessionRenewal(s *session) {
	var timer *time.Timer
	s.cancel = make(chan bool, 1)
	go func(s *session) {
		for {
			s.mux.RLock()
			w := (s.endTime.Sub(time.Now().UTC()) * 5) / 6
			s.mux.RUnlock()
			if w < 0 {
				return
			}
			timer = time.NewTimer(w)
			select {
			case <-timer.C:
				renewal, err := cl.refreshSession(s)
				if !renewal && err == nil {
					// end this goroutine as there will have been a new login and new auto renewal goroutine created.
					return
				}
			case <-s.cancel:
				// cancel has been called. Stop the timer and exit.
				timer.Stop()
				return
			}
		}
	}(s)
}

// renewTGT renews the client's TGT session.
func (cl *Client) renewTGT(s *session) error {
	realm, tgt, skey := s.tgtDetails()
	spn := types.PrincipalName{
		NameType:   nametype.KRB_NT_SRV_INST,
		NameString: []string{"krbtgt", realm},
	}
	_, tgsRep, err := cl.TGSExchange(spn, cl.Credentials.Realm, tgt, skey, true, 0)
	if err != nil {
		return krberror.Errorf(err, krberror.KRBMsgError, "error renewing TGT")
	}
	s.update(tgsRep.Ticket, tgsRep.DecryptedEncPart)
	cl.sessions.update(s)
	return nil
}

// updateSession updates either through renewal or creating a new login.
// The boolean indicates if the update was a renewal.
func (cl *Client) refreshSession(s *session) (bool, error) {
	s.mux.RLock()
	realm := s.realm
	renewTill := s.renewTill
	s.mux.RUnlock()
	if time.Now().UTC().Before(renewTill) {
		err := cl.renewTGT(s)
		return true, err
	}
	if realm != cl.Credentials.Realm {
		// session is not for the client's own realm
		_, err := cl.remoteRealmSession(realm)
		if err != nil {
			return false, err
		}
	}
	err := cl.Login()
	return false, err
}

// ensureValidSession makes sure there is a valid session for the realm
func (cl *Client) ensureValidSession(realm string) error {
	s, ok := cl.sessions.get(realm)
	if ok {
		s.mux.RLock()
		defer s.mux.RUnlock()
		d := s.endTime.Sub(s.authTime) / 6
		if s.endTime.Sub(time.Now().UTC()) > d {
			return nil
		}
		_, err := cl.refreshSession(s)
		return err
	}
	if realm != cl.Credentials.Realm {
		// not for the client's own realm
		_, err := cl.remoteRealmSession(realm)
		return err
	}
	return cl.Login()
}

// remoteRealmSession returns the session for a realm that the client is not a member of but for which there is a trust
func (cl *Client) remoteRealmSession(realm string) (*session, error) {
	s, ok := cl.sessions.get(cl.Credentials.Realm)
	if !ok || !s.valid() {
		err := cl.Login()
		if err != nil {
			return nil, fmt.Errorf("client was unable to login: %v", err)
		}
	}

	spn := types.PrincipalName{
		NameType:   nametype.KRB_NT_SRV_INST,
		NameString: []string{"krbtgt", realm},
	}

	_, tgsRep, err := cl.TGSExchange(spn, cl.Credentials.Realm, s.tgt, s.sessionKey, false, 0)
	if err != nil {
		return nil, err
	}
	cl.AddSession(tgsRep.Ticket, tgsRep.DecryptedEncPart)

	cl.sessions.mux.RLock()
	defer cl.sessions.mux.RUnlock()
	return cl.sessions.Entries[realm], nil
}

// realmSession returns the session for the realm provided.
func (cl *Client) realmSession(realm string) (*session, error) {
	s, ok := cl.sessions.get(realm)
	var err error
	if !ok {
		// Try to request TGT from trusted remote Realm
		s, err = cl.remoteRealmSession(realm)
		if err != nil {
			return s, err
		}
	}
	return s, nil
}

// sessionTGTDetails is a thread safe way to get the TGT and session key values for a realm
func (cl *Client) sessionTGTDetails(realm string) (tgt messages.Ticket, sessionKey types.EncryptionKey, err error) {
	var s *session
	s, err = cl.realmSession(realm)
	if err != nil {
		return
	}
	realm, tgt, sessionKey = s.tgtDetails()
	return
}

// spnRealm resolves the realm name of a service principal name
func (cl *Client) spnRealm(spn types.PrincipalName) string {
	return cl.Config.ResolveRealm(spn.NameString[len(spn.NameString)-1])
}
<commit_msg>race condition fix<commit_after>package client

import (
	"fmt"
	"sync"
	"time"

	"gopkg.in/jcmturner/gokrb5.v6/iana/nametype"
	"gopkg.in/jcmturner/gokrb5.v6/krberror"
	"gopkg.in/jcmturner/gokrb5.v6/messages"
	"gopkg.in/jcmturner/gokrb5.v6/types"
)

// sessions hold TGTs and are keyed on the realm name
type sessions struct {
	Entries map[string]*session
	mux     sync.RWMutex
}

// destroy erases all sessions
func (s *sessions) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	for k, e := range s.Entries {
		e.destroy()
		delete(s.Entries, k)
	}
}

// update replaces a session with the one provided or adds it as a new one
func (s *sessions) update(sess *session) {
	s.mux.Lock()
	defer s.mux.Unlock()
	// if a session already exists for this, cancel its auto renew.
	if i, ok := s.Entries[sess.realm]; ok {
		if i != sess {
			// Session in the sessions cache is not the same as one provided.
			// Cancel the one in the cache and add this one.
			i.mux.Lock()
			defer i.mux.Unlock()
			i.cancel <- true
			s.Entries[sess.realm] = sess
			return
		}
	}
	// No session for this realm was found so just add it
	s.Entries[sess.realm] = sess
}

// get returns the session for the realm specified
func (s *sessions) get(realm string) (*session, bool) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	sess, ok := s.Entries[realm]
	return sess, ok
}

// session holds the TGT details for a realm
type session struct {
	realm                string
	authTime             time.Time
	endTime              time.Time
	renewTill            time.Time
	tgt                  messages.Ticket
	sessionKey           types.EncryptionKey
	sessionKeyExpiration time.Time
	cancel               chan bool
	mux                  sync.RWMutex
}

// AddSession adds a session for a realm with a TGT to the client's session cache.
// A goroutine is started to automatically renew the TGT before expiry.
func (cl *Client) AddSession(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	realm := cl.spnRealm(tgt.SName)
	s := &session{
		realm:                realm,
		authTime:             dep.AuthTime,
		endTime:              dep.EndTime,
		renewTill:            dep.RenewTill,
		tgt:                  tgt,
		sessionKey:           dep.Key,
		sessionKeyExpiration: dep.KeyExpiration,
	}
	cl.sessions.update(s)
	cl.enableAutoSessionRenewal(s)
}

// update overwrites the session details with those from the TGT and decrypted encPart
func (s *session) update(tgt messages.Ticket, dep messages.EncKDCRepPart) {
	s.mux.Lock()
	defer s.mux.Unlock()
	s.authTime = dep.AuthTime
	s.endTime = dep.EndTime
	s.renewTill = dep.RenewTill
	s.tgt = tgt
	s.sessionKey = dep.Key
	s.sessionKeyExpiration = dep.KeyExpiration
}

// destroy will cancel any auto renewal of the session and set the expiration times to the current time
func (s *session) destroy() {
	s.mux.Lock()
	defer s.mux.Unlock()
	if s.cancel != nil {
		s.cancel <- true
	}
	s.endTime = time.Now().UTC()
	s.renewTill = s.endTime
	s.sessionKeyExpiration = s.endTime
}

// valid informs if the TGT is still within the valid time window
func (s *session) valid() bool {
	s.mux.RLock()
	defer s.mux.RUnlock()
	t := time.Now().UTC()
	if t.Before(s.endTime) && s.authTime.Before(t) {
		return true
	}
	return false
}

// tgtDetails is a thread safe way to get the session's realm, TGT and session key values
func (s *session) tgtDetails() (string, messages.Ticket, types.EncryptionKey) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.tgt, s.sessionKey
}

// timeDetails is a thread safe way to get the session's validity time values
func (s *session) timeDetails() (string, time.Time, time.Time, time.Time, time.Time) {
	s.mux.RLock()
	defer s.mux.RUnlock()
	return s.realm, s.authTime, s.endTime, s.renewTill, s.sessionKeyExpiration
}

// enableAutoSessionRenewal turns on the automatic renewal for the client's TGT session.
func (cl *Client) enableAutoSessionRenewal(s *session) {
	var timer *time.Timer
	s.cancel = make(chan bool, 1)
	go func(s *session) {
		for {
			s.mux.RLock()
			w := (s.endTime.Sub(time.Now().UTC()) * 5) / 6
			s.mux.RUnlock()
			if w < 0 {
				return
			}
			timer = time.NewTimer(w)
			select {
			case <-timer.C:
				renewal, err := cl.refreshSession(s)
				if !renewal && err == nil {
					// end this goroutine as there will have been a new login and new auto renewal goroutine created.
					return
				}
			case <-s.cancel:
				// cancel has been called. Stop the timer and exit.
				timer.Stop()
				return
			}
		}
	}(s)
}

// renewTGT renews the client's TGT session.
func (cl *Client) renewTGT(s *session) error {
	realm, tgt, skey := s.tgtDetails()
	spn := types.PrincipalName{
		NameType:   nametype.KRB_NT_SRV_INST,
		NameString: []string{"krbtgt", realm},
	}
	_, tgsRep, err := cl.TGSExchange(spn, cl.Credentials.Realm, tgt, skey, true, 0)
	if err != nil {
		return krberror.Errorf(err, krberror.KRBMsgError, "error renewing TGT")
	}
	s.update(tgsRep.Ticket, tgsRep.DecryptedEncPart)
	cl.sessions.update(s)
	return nil
}

// updateSession updates either through renewal or creating a new login.
// The boolean indicates if the update was a renewal.
func (cl *Client) refreshSession(s *session) (bool, error) {
	s.mux.RLock()
	realm := s.realm
	renewTill := s.renewTill
	s.mux.RUnlock()
	if time.Now().UTC().Before(renewTill) {
		err := cl.renewTGT(s)
		return true, err
	}
	if realm != cl.Credentials.Realm {
		// session is not for the client's own realm
		_, err := cl.remoteRealmSession(realm)
		if err != nil {
			return false, err
		}
	}
	err := cl.Login()
	return false, err
}

// ensureValidSession makes sure there is a valid session for the realm
func (cl *Client) ensureValidSession(realm string) error {
	s, ok := cl.sessions.get(realm)
	if ok {
		s.mux.RLock()
		defer s.mux.RUnlock()
		d := s.endTime.Sub(s.authTime) / 6
		if s.endTime.Sub(time.Now().UTC()) > d {
			return nil
		}
		_, err := cl.refreshSession(s)
		return err
	}
	if realm != cl.Credentials.Realm {
		// not for the client's own realm
		_, err := cl.remoteRealmSession(realm)
		return err
	}
	return cl.Login()
}

// remoteRealmSession returns the session for a realm that the client is not a member of but for which there is a trust
func (cl *Client) remoteRealmSession(realm string) (*session, error) {
	s, ok := cl.sessions.get(cl.Credentials.Realm)
	if !ok || !s.valid() {
		err := cl.Login()
		if err != nil {
			return nil, fmt.Errorf("client was unable to login: %v", err)
		}
	}

	spn := types.PrincipalName{
		NameType:   nametype.KRB_NT_SRV_INST,
		NameString: []string{"krbtgt", realm},
	}

	_, tgsRep, err := cl.TGSExchange(spn, cl.Credentials.Realm, s.tgt, s.sessionKey, false, 0)
	if err != nil {
		return nil, err
	}
	cl.AddSession(tgsRep.Ticket, tgsRep.DecryptedEncPart)

	cl.sessions.mux.RLock()
	defer cl.sessions.mux.RUnlock()
	return cl.sessions.Entries[realm], nil
}

// realmSession returns the session for the realm provided.
func (cl *Client) realmSession(realm string) (*session, error) {
	s, ok := cl.sessions.get(realm)
	var err error
	if !ok {
		// Try to request TGT from trusted remote Realm
		s, err = cl.remoteRealmSession(realm)
		if err != nil {
			return s, err
		}
	}
	return s, nil
}

// sessionTGTDetails is a thread safe way to get the TGT and session key values for a realm
func (cl *Client) sessionTGTDetails(realm string) (tgt messages.Ticket, sessionKey types.EncryptionKey, err error) {
	var s *session
	s, err = cl.realmSession(realm)
	if err != nil {
		return
	}
	realm, tgt, sessionKey = s.tgtDetails()
	return
}

// spnRealm resolves the realm name of a service principal name
func (cl *Client) spnRealm(spn types.PrincipalName) string {
	return cl.Config.ResolveRealm(spn.NameString[len(spn.NameString)-1])
}
<|endoftext|>


<commit_before>package glog

import (
	"fmt"

	"github.com/cayleygraph/cayley/clog"
	"github.com/golang/glog"
)

func init() {
	clog.SetLogger(Logger{})
}

type Logger struct{}

func (Logger) Infof(format string, args ...interface{}) {
	glog.InfoDepth(3, fmt.Sprintf(format, args...))
}
func (Logger) Warningf(format string, args ...interface{}) {
	glog.WarningDepth(3, fmt.Sprintf(format, args...))
}
func (Logger) Errorf(format string, args ...interface{}) {
	glog.ErrorDepth(3, fmt.Sprintf(format, args...))
}
func (Logger) Fatalf(format string, args ...interface{}) {
	glog.FatalDepth(3, fmt.Sprintf(format, args...))
}

func (Logger) V(level int) bool {
	return bool(glog.V(glog.Level(level)))
}

func (Logger) SetV(int) {
	panic("unexposed in glog")
}
<commit_msg>print warning on set verbosity for glog; fix #599<commit_after>package glog

import (
	"fmt"

	"github.com/cayleygraph/cayley/clog"
	"github.com/golang/glog"
)

func init() {
	clog.SetLogger(Logger{})
}

type Logger struct{}

func (Logger) Infof(format string, args ...interface{}) {
	glog.InfoDepth(3, fmt.Sprintf(format, args...))
}
func (Logger) Warningf(format string, args ...interface{}) {
	glog.WarningDepth(3, fmt.Sprintf(format, args...))
}
func (Logger) Errorf(format string, args ...interface{}) {
	glog.ErrorDepth(3, fmt.Sprintf(format, args...))
}
func (Logger) Fatalf(format string, args ...interface{}) {
	glog.FatalDepth(3, fmt.Sprintf(format, args...))
}

func (Logger) V(level int) bool {
	return bool(glog.V(glog.Level(level)))
}

func (Logger) SetV(v int) {
	glog.Warningf("changing log level is not supported; run command with '-v %d' flag", v)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"gopkg.in/yaml.v2"

	"github.com/kusabashira/acgen"
	"github.com/ogier/pflag"
)

var (
	name    = "acgen"
	version = "0.1.0"

	flag       = pflag.NewFlagSet(name, pflag.ContinueOnError)
	outputType = flag.StringP("type", "t", "", "")
	isHelp     = flag.BoolP("help", "h", false, "")
	isVersion  = flag.BoolP("version", "v", false, "")
)

func printUsage() {
	fmt.Fprintf(os.Stderr, `
Usage: %s [OPTION]... SOURCE
Generate auto-completions for specifyed TYPE
by SOURCE written in YAML.

Options:
  -t, --type=TYPE   output auto-completion for specified TYPE
                      TYPE=bash|zsh|fish|yash|tcsh
  -h, --help        display this help text and exit
  -v, --version     output version information and exit
`[1:], name)
}

func printVersion() {
	fmt.Fprintln(os.Stderr, version)
}

func printErr(err interface{}) {
	fmt.Fprintf(os.Stderr, "%s: %s\n", name, err)
}

func guideToHelp() {
	fmt.Fprintf(os.Stderr, "Try '%s --help' for more information.\n", name)
}

func main() {
	flag.SetOutput(ioutil.Discard)
	if err := flag.Parse(os.Args[1:]); err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	switch {
	case *isHelp:
		printUsage()
		os.Exit(0)
	case *isVersion:
		printVersion()
		os.Exit(0)
	}
	switch {
	case flag.NArg() < 1:
		printErr("no input file")
		guideToHelp()
		os.Exit(2)
	case *outputType == "":
		printErr("no specify TYPE")
		guideToHelp()
		os.Exit(2)
	}

	generator, err := acgen.LookGenerator(*outputType)
	if err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	file := flag.Arg(0)
	conf, err := ioutil.ReadFile(file)
	if err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	command := &acgen.Command{}
	if err = yaml.Unmarshal(conf, command); err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}

	if err = generator(os.Stdout, command); err != nil {
		printErr(err)
		os.Exit(1)
	}
}
<commit_msg>Refactor cmd/acgen/acgen.go around variables name<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"

	"gopkg.in/yaml.v2"

	"github.com/kusabashira/acgen"
	"github.com/ogier/pflag"
)

var (
	name    = "acgen"
	version = "0.1.0"

	flagset    = pflag.NewFlagSet(name, pflag.ContinueOnError)
	outputType = flagset.StringP("type", "t", "", "")
	isHelp     = flagset.BoolP("help", "h", false, "")
	isVersion  = flagset.BoolP("version", "v", false, "")
)

func printUsage() {
	fmt.Fprintf(os.Stderr, `
Usage: %s [OPTION]... SOURCE
Generate auto-completions for specifyed TYPE
by SOURCE written in YAML.

Options:
  -t, --type=TYPE   output auto-completion for specified TYPE
                      TYPE=bash|zsh|fish|yash|tcsh
  -h, --help        display this help text and exit
  -v, --version     output version information and exit
`[1:], name)
}

func printVersion() {
	fmt.Fprintln(os.Stderr, version)
}

func printErr(err interface{}) {
	fmt.Fprintf(os.Stderr, "%s: %s\n", name, err)
}

func guideToHelp() {
	fmt.Fprintf(os.Stderr, "Try '%s --help' for more information.\n", name)
}

func main() {
	flagset.SetOutput(ioutil.Discard)
	if err := flagset.Parse(os.Args[1:]); err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	switch {
	case *isHelp:
		printUsage()
		os.Exit(0)
	case *isVersion:
		printVersion()
		os.Exit(0)
	}
	switch {
	case flagset.NArg() < 1:
		printErr("no input file")
		guideToHelp()
		os.Exit(2)
	case *outputType == "":
		printErr("no specify TYPE")
		guideToHelp()
		os.Exit(2)
	}

	generate, err := acgen.LookGenerator(*outputType)
	if err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}

	source, err := ioutil.ReadFile(flagset.Arg(0))
	if err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}
	command := &acgen.Command{}
	if err = yaml.Unmarshal(source, command); err != nil {
		printErr(err)
		guideToHelp()
		os.Exit(2)
	}

	if err = generate(os.Stdout, command); err != nil {
		printErr(err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"image"
	"image/draw"
	_ "image/jpeg"
	"image/png"
	"math/rand"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/urfave/cli"
)

var (
	version = "master"
	app     *cli.App
	args    arguments
)

type arguments struct {
	input  string
	output string
	min    int
	max    int
}

func init() {
	app = &cli.App{
		Name:    "asami",
		Usage:   "pixel sorter using simple brute sorting",
		Version: version,
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:        "input",
				Usage:       "the input file `path` to use",
				Aliases:     []string{"i"},
				Destination: &args.input,
			},
			&cli.StringFlag{
				Name:        "output",
				Usage:       "the output file `path` to use, must be a png file",
				Aliases:     []string{"o"},
				Destination: &args.output,
			},
			&cli.IntFlag{
				Name:        "min",
				Usage:       "the minimum chunk `size`",
				Value:       32,
				Destination: &args.min,
			},
			&cli.IntFlag{
				Name:        "max",
				Usage:       "the maximum chunk `size`",
				Value:       64,
				Destination: &args.max,
			},
		},
		Action: func(c *cli.Context) error {
			if err := checkInput(c); err != nil {
				return err
			}

			return nil
		},
	}
}

func checkInput(c *cli.Context) error {
	errors := []string{}
	if args.input == "" {
		errors = append(errors, "input file path must be set")
	}
	if args.output == "" {
		errors = append(errors, "output file path must be set")
	}
	if args.min < 1 {
		errors = append(errors, "minimum chunk size must be at least 1")
	}
	if args.max < args.min {
		errors = append(errors, "maximum chunk size must be at least as big as the minimum chunk size")
	}

	if len(errors) > 0 {
		errorMessage := strings.Join(errors, "\n")
		return cli.Exit(errorMessage, 1)
	}

	rgba, err := loadImage(args.input)
	if err != nil {
		return err
	}

	sortRGBA(rgba)

	err = saveImage(args.output, rgba)
	if err != nil {
		return err
	}

	return nil
}

func loadImage(filename string) (*image.RGBA, error) {
	inFile, err := os.Open(filename)
	if err != nil {
		return nil, cli.Exit("failed reading input file", 2)
	}
	defer inFile.Close()

	img, _, err := image.Decode(inFile)
	if err != nil {
		return nil, cli.Exit("failed decoding input file", 3)
	}

	rgba := image.NewRGBA(img.Bounds())
	draw.Draw(rgba, rgba.Bounds(), img, image.Point{0, 0}, draw.Src)

	return rgba, nil
}

func saveImage(filename string, rgba *image.RGBA) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return cli.Exit("failed writing output image to file", 4)
	}
	defer outFile.Close()

	png.Encode(outFile, rgba)

	return nil
}

type chunk struct {
	begin int
	size  int
}

func sortRGBA(rgba *image.RGBA) {
	chunks := calculateChunks(rgba.Bounds())

	for _, chunk := range chunks {
		chunk.sort(rgba)
	}
}

func calculateChunks(bounds image.Rectangle) []chunk {
	chunks := []chunk{}

	size := bounds.Max.Sub(bounds.Min)
	width := size.X
	height := size.Y

	rand.Seed(time.Now().UnixNano())

	for y := 0; y < height; y++ {
		x := 0
		for x < width {
			size := chunkSize()
			if size > width-x {
				size = width - x
			}
			begin := y*width + x
			chunks = append(chunks, chunk{begin, size})
			x += size
		}
	}

	return chunks
}

func chunkSize() int {
	return rand.Intn(args.max-args.min) + args.min
}

type rgbaPixel [4]uint8
type rgbaPixels []rgbaPixel

type byRed rgbaPixels

func (p byRed) Len() int           { return len(p) }
func (p byRed) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p byRed) Less(i, j int) bool { return p[i][0] < p[j][0] }

type byGreen rgbaPixels

func (p byGreen) Len() int           { return len(p) }
func (p byGreen) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p byGreen) Less(i, j int) bool { return p[i][1] < p[j][1] }

type byBlue rgbaPixels

func (p byBlue) Len() int           { return len(p) }
func (p byBlue) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p byBlue) Less(i, j int) bool { return p[i][2] < p[j][2] }

type bySum rgbaPixels

func (p bySum) Len() int      { return len(p) }
func (p bySum) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p bySum) Less(i, j int) bool {
	return p[i].sum() < p[j].sum()
}

func (p *rgbaPixel) sum() uint32 {
	return uint32(p[0]) + uint32(p[1]) + uint32(p[2])
}

type byGrayscale rgbaPixels

func (p byGrayscale) Len() int      { return len(p) }
func (p byGrayscale) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p byGrayscale) Less(i, j int) bool {
	return p[i].gray() < p[j].gray()
}

func (p *rgbaPixel) gray() uint8 {
	return p[0]>>2 + p[1]>>1 + p[1]>>3 + p[2]>>3
}

func (chunk *chunk) sort(rgba *image.RGBA) {
	beginIndex := chunk.begin * 4
	endIndex := beginIndex + chunk.size*4
	data := rgba.Pix[beginIndex:endIndex]

	pixels := rgbaPixels{}

	for i := 0; i < len(data); i += 4 {
		pixel := rgbaPixel{
			data[i+0],
			data[i+1],
			data[i+2],
			data[i+3],
		}
		pixels = append(pixels, pixel)
	}

	sort.Sort(byBlue(pixels))

	for i, pixel := range pixels {
		data[i*4+0] = pixel[0]
		data[i*4+1] = pixel[1]
		data[i*4+2] = pixel[2]
		data[i*4+3] = pixel[3]
	}
}

func main() {
	app.Run(os.Args)
}
<commit_msg>sorting code is now inside the Action (and not inside the input checking code)<commit_after>package main

import (
	"image"
	"image/draw"
	_ "image/jpeg" // to be able to import jpeg images
	"image/png"
	"math/rand"
	"os"
	"sort"
	"strings"
	"time"

	"github.com/urfave/cli"
)

var (
	version = "master"
	app     *cli.App
	args    arguments
)

type arguments struct {
	input  string
	output string
	min    int
	max    int
}

func init() {
	app = &cli.App{
		Name:    "asami",
		Usage:   "pixel sorter using simple brute sorting",
		Version: version,
		Flags: []cli.Flag{
			&cli.StringFlag{
				Name:        "input",
				Usage:       "the input file `path` to use",
				Aliases:     []string{"i"},
				Destination: &args.input,
			},
			&cli.StringFlag{
				Name:        "output",
				Usage:       "the output file `path` to use, must be a png file",
				Aliases:     []string{"o"},
				Destination: &args.output,
			},
			&cli.IntFlag{
				Name:        "min",
				Usage:       "the minimum chunk `size`",
				Value:       32,
				Destination: &args.min,
			},
			&cli.IntFlag{
				Name:        "max",
				Usage:       "the maximum chunk `size`",
				Value:       64,
				Destination: &args.max,
			},
		},
		Action: func(c *cli.Context) error {
			if err := checkInput(c); err != nil {
				return err
			}

			rgba, err := loadImage(args.input)
			if err != nil {
				return err
			}

			sortRGBA(rgba)

			err = saveImage(args.output, rgba)
			if err != nil {
				return err
			}

			return nil
		},
	}
}

func checkInput(c *cli.Context) error {
	errors := []string{}
	if args.input == "" {
		errors = append(errors, "input file path must be set")
	}
	if args.output == "" {
		errors = append(errors, "output file path must be set")
	}
	if args.min < 1 {
		errors = append(errors, "minimum chunk size must be at least 1")
	}
	if args.max < args.min {
		errors = append(errors, "maximum chunk size must be at least as big as the minimum chunk size")
	}

	if len(errors) > 0 {
		errorMessage := strings.Join(errors, "\n")
		return cli.Exit(errorMessage, 1)
	}

	return nil
}

func loadImage(filename string) (*image.RGBA, error) {
	inFile, err := os.Open(filename)
	if err != nil {
		return nil, cli.Exit("failed reading input file", 2)
	}
	defer inFile.Close()

	img, _, err := image.Decode(inFile)
	if err != nil {
		return nil, cli.Exit("failed decoding input file", 3)
	}

	rgba := image.NewRGBA(img.Bounds())
	draw.Draw(rgba, rgba.Bounds(), img, image.Point{0, 0}, draw.Src)

	return rgba, nil
}

func saveImage(filename string, rgba *image.RGBA) error {
	outFile, err := os.Create(filename)
	if err != nil {
		return cli.Exit("failed writing output image to file", 4)
	}
	defer outFile.Close()

	png.Encode(outFile, rgba)

	return nil
}

type chunk struct {
	begin int
	size  int
}

func sortRGBA(rgba *image.RGBA) {
	chunks := calculateChunks(rgba.Bounds())

	for _, chunk := range chunks {
		chunk.sort(rgba)
	}
}

func calculateChunks(bounds image.Rectangle) []chunk {
	chunks := []chunk{}

	size := bounds.Max.Sub(bounds.Min)
	width := size.X
	height := size.Y

	rand.Seed(time.Now().UnixNano())

	for y := 0; y < height; y++ {
		x := 0
		for x < width {
			size := chunkSize()
			if size > width-x {
				size = width - x
			}
			begin := y*width + x
			chunks = append(chunks, chunk{begin, size})
			x += size
		}
	}

	return chunks
}

func chunkSize() int {
	return rand.Intn(args.max-args.min) + args.min
}

type rgbaPixel [4]uint8
type rgbaPixels []rgbaPixel

type byRed rgbaPixels

func (p byRed) Len() int           { return len(p) }
func (p byRed) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p byRed) Less(i, j int) bool { return p[i][0] < p[j][0] }

type byGreen rgbaPixels

func (p byGreen) Len() int           { return len(p) }
func (p byGreen) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p byGreen) Less(i, j int) bool { return p[i][1] < p[j][1] }

type byBlue rgbaPixels

func (p byBlue) Len() int           { return len(p) }
func (p byBlue) Swap(i, j int)      { p[i], p[j] = p[j], p[i] }
func (p byBlue) Less(i, j int) bool { return p[i][2] < p[j][2] }

type bySum rgbaPixels

func (p bySum) Len() int      { return len(p) }
func (p bySum) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p bySum) Less(i, j int) bool {
	return p[i].sum() < p[j].sum()
}

func (p *rgbaPixel) sum() uint32 {
	return uint32(p[0]) + uint32(p[1]) + uint32(p[2])
}

type byGrayscale rgbaPixels

func (p byGrayscale) Len() int      { return len(p) }
func (p byGrayscale) Swap(i, j int) { p[i], p[j] = p[j], p[i] }
func (p byGrayscale) Less(i, j int) bool {
	return p[i].gray() < p[j].gray()
}

func (p *rgbaPixel) gray() uint8 {
	return p[0]>>2 + p[1]>>1 + p[1]>>3 + p[2]>>3
}

func (chunk *chunk) sort(rgba *image.RGBA) {
	beginIndex := chunk.begin * 4
	endIndex := beginIndex + chunk.size*4
	data := rgba.Pix[beginIndex:endIndex]

	pixels := rgbaPixels{}

	for i := 0; i < len(data); i += 4 {
		pixel := rgbaPixel{
			data[i+0],
			data[i+1],
			data[i+2],
			data[i+3],
		}
		pixels = append(pixels, pixel)
	}

	sort.Sort(byBlue(pixels))

	for i, pixel := range pixels {
		data[i*4+0] = pixel[0]
		data[i*4+1] = pixel[1]
		data[i*4+2] = pixel[2]
		data[i*4+3] = pixel[3]
	}
}

func main() {
	app.Run(os.Args)
}
<|endoftext|>