<commit_before>package main

import (
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"gopkg.in/gavv/httpexpect.v2"
)

// TODO(dotslash): Try out a more modern test framework, that displays test
// exection summaries better. As of now og output looks like this -
// https://gist.github.com/dotslash/56bed5506ff96c2f7740c54c451d5986

type Outcome bool

// Failure outcome.
var FAILURE = Outcome(false)

// Bookmarks web server address used to generate short urls.
var testBookMarksServer = "https://dotslash.com"

// https://stackoverflow.com/q/22892120
func RandStringBytesMaskImprSrc(n int) string {
	const letterBytes = "abcdefghijklmnopqrstuvwxyz0123456789"
	const (
		letterIdxBits = 6                    // 6 bits to represent a letter index
		letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
		letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
	)
	src := rand.NewSource(time.Now().UnixNano())
	b := make([]byte, n)
	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
	for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
		if remain == 0 {
			cache, remain = src.Int63(), letterIdxMax
		}
		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
			b[i] = letterBytes[idx]
			i--
		}
		cache >>= letterIdxBits
		remain--
	}
	return string(b)
}

// Without this:
//   x, err := foo()
//   if err != nil {
//	   panic(err)
//   }
// With this:
//   x := PanicIfErr1(foo()).(string)
func PanicIfErr1(inp interface{}, err error) interface{} {
	if err != nil {
		panic(err)
	}
	return inp
}

// Without this:
//   err := foo()
//   if err != nil {
//	   panic(err)
//   }
// With this:
//   PanicIfErr(foo())
func PanicIfErr(err error) {
	if err != nil {
		panic(err)
	}
}

// An implementation of http.Handler interface, that returns the text "redirected"
// in the response and saves the input request to `requests` slice.
type SaveRequestsHandler struct {
	requests []*http.Request
}

func (h *SaveRequestsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.requests = append(h.requests, r)
	w.Write([]byte("redirected"))
	w.WriteHeader(http.StatusOK)
}

// Creates sqlite db with the correct schema into a random file in /tmp/ and
// returns the path of the sqlite file.
func CreateTmpSqlite() string {
	tmpSqliteFile := "/tmp/bookmarks_test_" + RandStringBytesMaskImprSrc(20)
	sourceFile := "testdata/test_db"
	input := PanicIfErr1(ioutil.ReadFile(sourceFile)).([]byte)
	PanicIfErr(ioutil.WriteFile(tmpSqliteFile, input, 0644))
	return tmpSqliteFile
}

func NewTestHelper(t *testing.T) TestHelper {
	sqliteFile := CreateTmpSqlite()
	saveReqHandler := &SaveRequestsHandler{}
	server := httptest.NewServer(
		NewRouter(testBookMarksServer, sqliteFile))
	return TestHelper{
		server:         server,
		sqliteFile:     sqliteFile,
		redirectReqs:   &saveReqHandler.requests,
		redirectServer: httptest.NewServer(saveReqHandler),
		httpexpect: httpexpect.WithConfig(httpexpect.Config{
			BaseURL:  server.URL,
			Reporter: httpexpect.NewRequireReporter(t),
			Printers: []httpexpect.Printer{
				httpexpect.NewDebugPrinter(t, true),
			},
		}),
		require: require.New(t),
	}
}

type TestHelper struct {
	// bookmarks HTTP server.
	server *httptest.Server
	// A server that will be used as target for the redirects.
	// This saves the incoming requests to `redirectReqs` field.
	redirectServer *httptest.Server
	redirectReqs   *[]*http.Request
	// Path of sqlite file used for test.
	sqliteFile string
	// httpexpect and require are test utilites for checking expectations
	// and actual values returned.
	httpexpect *httpexpect.Expect
	require    *require.Assertions
}

func (t *TestHelper) cleanup() {
	t.redirectServer.Close()
	t.server.Close()
}

func AddAlias(th *TestHelper, short, url, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]
	res := th.httpexpect.POST("/actions/add").
		WithFormField("short", short).
		WithFormField("url", url).
		WithFormField("secret", secret).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}
}

func UpdateOrigURL(th *TestHelper, alias, oldLong, newLong, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]

	res := th.httpexpect.POST("/actions/update").
		WithFormField("id", alias).
		WithFormField("newvalue", newLong).
		WithFormField("oldvalue", oldLong).
		WithFormField("colname", "orig").
		WithFormField("secret", secret).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}

}

func UpdateAlias(th *TestHelper, oldAlias, newAlias, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]

	res := th.httpexpect.POST("/actions/update").
		WithFormField("id", oldAlias).
		WithFormField("newvalue", newAlias).
		WithFormField("oldvalue", oldAlias).
		WithFormField("colname", "alias").
		WithFormField("secret", secret).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}

}

func RemoveAlias(th *TestHelper, short, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]
	res := th.httpexpect.POST("/actions/delete").
		WithFormField("short", short).WithFormField("secret", secret).Expect()

	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}
}

func Lookup(th *TestHelper, short, long string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]
	res := th.httpexpect.POST("/actions/lookup").
		WithFormField("short", short).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Body().Equal(long)
	} else {
		res.Status(http.StatusNotFound)
	}
}

func RevLookup(th *TestHelper, long string, short []interface{}, secret string) {
	th.httpexpect.POST("/actions/revlookup").
		WithFormField("long", long).WithFormField("secret", secret).
		Expect().Status(http.StatusOK).
		JSON().Path("$.shorturls").Array().ContainsOnly(short...)
}
func RedirectIs404(th *TestHelper, short string) {
	prev := len(*th.redirectReqs)
	th.httpexpect.GET(short).Expect().StatusRange(httpexpect.Status4xx)
	th.require.Equal(prev, len(*th.redirectReqs))
}

func RedirectWorks(th *TestHelper, short, urlPath string) {
	prev := len(*th.redirectReqs)
	th.httpexpect.GET("/r/" + short).Expect().Text().Equal("redirected")
	th.require.Equal(prev+1, len(*th.redirectReqs))
	apath := (*th.redirectReqs)[prev].URL.Path
	th.require.Equal(urlPath, apath)
}

func ViewBookmarks(th *TestHelper, aliases []AliasInfo, secret string) {
	expRespStruct := CreateViewResponse(aliases, testBookMarksServer)
	actRespStr := th.httpexpect.POST("/actions/view").
		WithFormField("secret", secret).Expect().
		Status(http.StatusOK).Body().Raw()
	actRespStruct := ViewResponse{}
	th.require.Nil(json.Unmarshal([]byte(actRespStr), &actRespStruct))
	// This is slightly fragile. This fails if the order of rows returned in
	// the response does not match the order of `alises` array. Using insertion
	// order seems to work.
	// If this is getting problematic, sort the rows in the struct before
	// comparing
	th.require.Equal(expRespStruct, actRespStruct)
}

func TestMutationsAndRedirects(t *testing.T) {
	th := NewTestHelper(t)
	defer th.cleanup()

	// Begin: s1, _s1 redirects dont exist.
	RedirectIs404(&th, "s1")
	RedirectIs404(&th, "_s1")

	// Add redirects for s1, _s1.
	AddAlias(&th, "s1", th.redirectServer.URL+"/long1", "strong-secret")
	AddAlias(&th, "_s1", th.redirectServer.URL+"/_long1", "strong-secret")
	// - Check redirects.
	// NOTE: Eventhough _s1 is a private alias, it will still be redirected
	//       without needing the secret. Secret is needed for lookups involving
	//       private alises and all mutations.
	RedirectWorks(&th, "s1", "/long1")
	RedirectWorks(&th, "_s1", "/_long1")

	// Update long urls for s1, _s1.
	UpdateOrigURL(&th,
		"s1",                           // alias
		th.redirectServer.URL+"/long1", // oldLong
		th.redirectServer.URL+"/long2", // newLong
		"strong-secret",                // secret
	)
	UpdateOrigURL(&th,
		"_s1",                           // alias
		th.redirectServer.URL+"/_long1", // oldLong
		th.redirectServer.URL+"/_long2", // newLong
		"strong-secret",                 // secret
	)
	// - Check that redirect urls are updated.
	RedirectWorks(&th, "s1", "/long2")
	RedirectWorks(&th, "_s1", "/_long2")

	// Update alises; s1->s2, _s1->_s2.
	UpdateAlias(&th, "s1", "s2", "strong-secret")
	UpdateAlias(&th, "_s1", "_s2", "strong-secret")
	// - Check redirects for _s2, s2
	RedirectWorks(&th, "s2", "/long2")
	RedirectWorks(&th, "_s2", "/_long2")
	// - Check no redirects for _s1, s1
	RedirectIs404(&th, "s1")
	RedirectIs404(&th, "_s1")

	// Remove aliases for _s2, s2
	RemoveAlias(&th, "s2", "strong-secret")
	RemoveAlias(&th, "_s2", "strong-secret")
	// - Check no redirects for _s2, s2
	RedirectIs404(&th, "_s2")
	RedirectIs404(&th, "s2")
}

func TestLookupsForPublicBookmarks(t *testing.T) {
	th := NewTestHelper(t)
	defer th.cleanup()

	// Add aliases for _s, s1, s2.
	// _s is a private alias.
	AddAlias(&th, "_s", th.redirectServer.URL+"/long", "strong-secret")
	AddAlias(&th, "s1", th.redirectServer.URL+"/long", "strong-secret")
	AddAlias(&th, "s2", th.redirectServer.URL+"/long", "strong-secret")
	// Alias to long url should work for all the aliases. (even the private one)
	Lookup(&th, "_s", th.redirectServer.URL+"/long")
	Lookup(&th, "s1", th.redirectServer.URL+"/long")
	Lookup(&th, "s2", th.redirectServer.URL+"/long")
	// RevLookup with no secret => no private alises returned.
	RevLookup(&th, th.redirectServer.URL+"/long",
		[]interface{}{"s1", "s2"}, "")
	// RevLookup with secret => private alises returned.
	RevLookup(&th, th.redirectServer.URL+"/long",
		[]interface{}{"s1", "s2", "_s"}, "strong-secret")
	// RevLookup with random url => nothing returned.
	RevLookup(&th, th.redirectServer.URL+"/does_not_exist",
		[]interface{}{}, "strong-secret")
	// View bookmarks with correct secret => all aliases returned.
	ViewBookmarks(&th,
		[]AliasInfo{
			{"_s", th.redirectServer.URL + "/long", "_s"},
			{"s1", th.redirectServer.URL + "/long", "s1"},
			{"s2", th.redirectServer.URL + "/long", "s2"},
		},
		"strong-secret")
	// View bookmarks with wrong secret => only public aliases returned.
	ViewBookmarks(&th,
		[]AliasInfo{
			{"s1", th.redirectServer.URL + "/long", "s1"},
			{"s2", th.redirectServer.URL + "/long", "s2"},
		},
		"wrong-secret")
}
func TestWrongPasswordFailsMutations(t *testing.T) {
	th := NewTestHelper(t)
	defer th.cleanup()

	// Setup: Add alias and check it works.
	AddAlias(&th, "s1", th.redirectServer.URL+"/long1", "strong-secret")
	RedirectWorks(&th, "s1", "/long1")

	// Try to add alias request for s2 with wrong secret.
	// Check that redirect fails for s2.
	AddAlias(&th, "s2", th.redirectServer.URL+"/long1", "wrong-secret", FAILURE)
	RedirectIs404(&th, "s2")

	// Try to update long url for s1 with wrong secret.
	// Verify the update was not applied by checking redirect.
	UpdateOrigURL(&th,
		"s1",                           // alias
		th.redirectServer.URL+"/long1", // oldLong
		th.redirectServer.URL+"/long2", // newLong
		"wrong-secret",                 // secret
		FAILURE,
	)
	RedirectWorks(&th, "s1", "/long1")

	// Try to update alias for s1 to s2 with wrong secret.
	// Verify the update was not applied by checking redirect.
	UpdateAlias(&th, "s1", "s2", "wrong-secret", FAILURE)
	RedirectWorks(&th, "s1", "/long1")
	RedirectIs404(&th, "s2")

	// Try to remove alias with wrong secret.
	// Verify the update was not applied by checking redirect.
	RemoveAlias(&th, "s1", "wrong-secret", FAILURE)
	RedirectWorks(&th, "s1", "/long1")
}
<commit_msg>Minor reorg in tests<commit_after>package main

import (
	"encoding/json"
	"io/ioutil"
	"math/rand"
	"net/http"
	"net/http/httptest"
	"testing"
	"time"

	"github.com/stretchr/testify/require"
	"gopkg.in/gavv/httpexpect.v2"
)

// TODO(dotslash): Try out a more modern test framework, that displays test
// exection summaries better. As of now og output looks like this -
// https://gist.github.com/dotslash/56bed5506ff96c2f7740c54c451d5986

type Outcome bool

// Failure outcome.
var FAILURE = Outcome(false)

// Bookmarks web server address used to generate short urls.
var testBookMarksServer = "https://dotslash.com"

// https://stackoverflow.com/q/22892120
func RandStringBytesMaskImprSrc(n int) string {
	const letterBytes = "abcdefghijklmnopqrstuvwxyz0123456789"
	const (
		letterIdxBits = 6                    // 6 bits to represent a letter index
		letterIdxMask = 1<<letterIdxBits - 1 // All 1-bits, as many as letterIdxBits
		letterIdxMax  = 63 / letterIdxBits   // # of letter indices fitting in 63 bits
	)
	src := rand.NewSource(time.Now().UnixNano())
	b := make([]byte, n)
	// A src.Int63() generates 63 random bits, enough for letterIdxMax characters!
	for i, cache, remain := n-1, src.Int63(), letterIdxMax; i >= 0; {
		if remain == 0 {
			cache, remain = src.Int63(), letterIdxMax
		}
		if idx := int(cache & letterIdxMask); idx < len(letterBytes) {
			b[i] = letterBytes[idx]
			i--
		}
		cache >>= letterIdxBits
		remain--
	}
	return string(b)
}

// Without this:
//   x, err := foo()
//   if err != nil {
//	   panic(err)
//   }
// With this:
//   x := PanicIfErr1(foo()).(string)
func PanicIfErr1(inp interface{}, err error) interface{} {
	if err != nil {
		panic(err)
	}
	return inp
}

// Without this:
//   err := foo()
//   if err != nil {
//	   panic(err)
//   }
// With this:
//   PanicIfErr(foo())
func PanicIfErr(err error) {
	if err != nil {
		panic(err)
	}
}

// An implementation of http.Handler interface, that returns the text "redirected"
// in the response and saves the input request to `requests` slice.
type SaveRequestsHandler struct {
	requests []*http.Request
}

func (h *SaveRequestsHandler) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	h.requests = append(h.requests, r)
	w.Write([]byte("redirected"))
	w.WriteHeader(http.StatusOK)
}

// Creates sqlite db with the correct schema into a random file in /tmp/ and
// returns the path of the sqlite file.
func CreateTmpSqlite() string {
	tmpSqliteFile := "/tmp/bookmarks_test_" + RandStringBytesMaskImprSrc(20)
	sourceFile := "testdata/test_db"
	input := PanicIfErr1(ioutil.ReadFile(sourceFile)).([]byte)
	PanicIfErr(ioutil.WriteFile(tmpSqliteFile, input, 0644))
	return tmpSqliteFile
}

func NewTestHelper(t *testing.T) TestHelper {
	sqliteFile := CreateTmpSqlite()
	saveReqHandler := &SaveRequestsHandler{}
	server := httptest.NewServer(
		NewRouter(testBookMarksServer, sqliteFile))
	return TestHelper{
		server:         server,
		sqliteFile:     sqliteFile,
		redirectReqs:   &saveReqHandler.requests,
		redirectServer: httptest.NewServer(saveReqHandler),
		httpexpect: httpexpect.WithConfig(httpexpect.Config{
			BaseURL:  server.URL,
			Reporter: httpexpect.NewRequireReporter(t),
			Printers: []httpexpect.Printer{
				httpexpect.NewDebugPrinter(t, true),
			},
		}),
		require: require.New(t),
	}
}

type TestHelper struct {
	// bookmarks HTTP server.
	server *httptest.Server
	// A server that will be used as target for the redirects.
	// This saves the incoming requests to `redirectReqs` field.
	redirectServer *httptest.Server
	redirectReqs   *[]*http.Request
	// Path of sqlite file used for test.
	sqliteFile string
	// httpexpect and require are test utilites for checking expectations
	// and actual values returned.
	// httpexpect is tied to requests for `server`.
	httpexpect *httpexpect.Expect
	require    *require.Assertions
}

func (th *TestHelper) cleanup() {
	th.redirectServer.Close()
	th.server.Close()
}

func (th *TestHelper) AddAlias(short, url, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]
	res := th.httpexpect.POST("/actions/add").
		WithFormField("short", short).
		WithFormField("url", url).
		WithFormField("secret", secret).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}
}

func (th *TestHelper) UpdateOrigURL(alias, oldLong, newLong, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]

	res := th.httpexpect.POST("/actions/update").
		WithFormField("id", alias).
		WithFormField("newvalue", newLong).
		WithFormField("oldvalue", oldLong).
		WithFormField("colname", "orig").
		WithFormField("secret", secret).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}

}

func (th *TestHelper) UpdateAlias(oldAlias, newAlias, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]

	res := th.httpexpect.POST("/actions/update").
		WithFormField("id", oldAlias).
		WithFormField("newvalue", newAlias).
		WithFormField("oldvalue", oldAlias).
		WithFormField("colname", "alias").
		WithFormField("secret", secret).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}

}

func (th *TestHelper) RemoveAlias(short, secret string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]
	res := th.httpexpect.POST("/actions/delete").
		WithFormField("short", short).WithFormField("secret", secret).Expect()

	if expectSuccess {
		res.Status(http.StatusOK).Text().Equal("ok")
	} else {
		res.Status(http.StatusOK).Text().NotEqual("ok")
	}
}

func (th *TestHelper) Lookup(short, long string, ok ...Outcome) {
	th.require.LessOrEqual(len(ok), 1)
	expectSuccess := len(ok) == 0 || ok[0]
	res := th.httpexpect.POST("/actions/lookup").
		WithFormField("short", short).Expect()
	if expectSuccess {
		res.Status(http.StatusOK).Body().Equal(long)
	} else {
		res.Status(http.StatusNotFound)
	}
}

func (th *TestHelper) RevLookup(long string, short []interface{}, secret string) {
	th.httpexpect.POST("/actions/revlookup").
		WithFormField("long", long).WithFormField("secret", secret).
		Expect().Status(http.StatusOK).
		JSON().Path("$.shorturls").Array().ContainsOnly(short...)
}
func (th *TestHelper) RedirectIs404(short string) {
	prev := len(*th.redirectReqs)
	th.httpexpect.GET(short).Expect().StatusRange(httpexpect.Status4xx)
	th.require.Equal(prev, len(*th.redirectReqs))
}

func (th *TestHelper) RedirectWorks(short, urlPath string) {
	prev := len(*th.redirectReqs)
	th.httpexpect.GET("/r/" + short).Expect().Text().Equal("redirected")
	th.require.Equal(prev+1, len(*th.redirectReqs))
	apath := (*th.redirectReqs)[prev].URL.Path
	th.require.Equal(urlPath, apath)
}

func (th *TestHelper) ViewBookmarks(aliases []AliasInfo, secret string) {
	expRespStruct := CreateViewResponse(aliases, testBookMarksServer)
	actRespStr := th.httpexpect.POST("/actions/view").
		WithFormField("secret", secret).Expect().
		Status(http.StatusOK).Body().Raw()
	actRespStruct := ViewResponse{}
	th.require.Nil(json.Unmarshal([]byte(actRespStr), &actRespStruct))
	// This is slightly fragile. This fails if the order of rows returned in
	// the response does not match the order of `alises` array. Using insertion
	// order seems to work.
	// If this is getting problematic, sort the rows in the struct before
	// comparing
	th.require.Equal(expRespStruct, actRespStruct)
}

func TestMutationsAndRedirects(t *testing.T) {
	th := NewTestHelper(t)
	defer th.cleanup()

	// Begin: s1, _s1 redirects dont exist.
	th.RedirectIs404("s1")
	th.RedirectIs404("_s1")

	// Add redirects for s1, _s1.
	th.AddAlias("s1", th.redirectServer.URL+"/long1", "strong-secret")
	th.AddAlias("_s1", th.redirectServer.URL+"/_long1", "strong-secret")
	// - Check redirects.
	// NOTE: Eventhough _s1 is a private alias, it will still be redirected
	//       without needing the secret. Secret is needed for lookups involving
	//       private alises and all mutations.
	th.RedirectWorks("s1", "/long1")
	th.RedirectWorks("_s1", "/_long1")

	// Update long urls for s1, _s1.
	th.UpdateOrigURL(
		"s1",                           // alias
		th.redirectServer.URL+"/long1", // oldLong
		th.redirectServer.URL+"/long2", // newLong
		"strong-secret",                // secret
	)
	th.UpdateOrigURL(
		"_s1",                           // alias
		th.redirectServer.URL+"/_long1", // oldLong
		th.redirectServer.URL+"/_long2", // newLong
		"strong-secret",                 // secret
	)
	// - Check that redirect urls are updated.
	th.RedirectWorks("s1", "/long2")
	th.RedirectWorks("_s1", "/_long2")

	// Update alises; s1->s2, _s1->_s2.
	th.UpdateAlias("s1", "s2", "strong-secret")
	th.UpdateAlias("_s1", "_s2", "strong-secret")
	// - Check redirects for _s2, s2
	th.RedirectWorks("s2", "/long2")
	th.RedirectWorks("_s2", "/_long2")
	// - Check no redirects for _s1, s1
	th.RedirectIs404("s1")
	th.RedirectIs404("_s1")

	// Remove aliases for _s2, s2
	th.RemoveAlias("s2", "strong-secret")
	th.RemoveAlias("_s2", "strong-secret")
	// - Check no redirects for _s2, s2
	th.RedirectIs404("_s2")
	th.RedirectIs404("s2")
}

func TestLookupsForPublicBookmarks(t *testing.T) {
	th := NewTestHelper(t)
	defer th.cleanup()

	// Add aliases for _s, s1, s2.
	// _s is a private alias.
	th.AddAlias("_s", th.redirectServer.URL+"/long", "strong-secret")
	th.AddAlias("s1", th.redirectServer.URL+"/long", "strong-secret")
	th.AddAlias("s2", th.redirectServer.URL+"/long", "strong-secret")
	// Alias to long url should work for all the aliases. (even the private one)
	th.Lookup("_s", th.redirectServer.URL+"/long")
	th.Lookup("s1", th.redirectServer.URL+"/long")
	th.Lookup("s2", th.redirectServer.URL+"/long")
	// RevLookup with no secret => no private alises returned.
	th.RevLookup(th.redirectServer.URL+"/long",
		[]interface{}{"s1", "s2"}, "")
	// RevLookup with secret => private alises returned.
	th.RevLookup(th.redirectServer.URL+"/long",
		[]interface{}{"s1", "s2", "_s"}, "strong-secret")
	// RevLookup with random url => nothing returned.
	th.RevLookup(th.redirectServer.URL+"/does_not_exist",
		[]interface{}{}, "strong-secret")
	// View bookmarks with correct secret => all aliases returned.
	th.ViewBookmarks(
		[]AliasInfo{
			{"_s", th.redirectServer.URL + "/long", "_s"},
			{"s1", th.redirectServer.URL + "/long", "s1"},
			{"s2", th.redirectServer.URL + "/long", "s2"},
		},
		"strong-secret")
	// View bookmarks with wrong secret => only public aliases returned.
	th.ViewBookmarks(
		[]AliasInfo{
			{"s1", th.redirectServer.URL + "/long", "s1"},
			{"s2", th.redirectServer.URL + "/long", "s2"},
		},
		"wrong-secret")
}
func TestWrongPasswordFailsMutations(t *testing.T) {
	th := NewTestHelper(t)
	defer th.cleanup()

	// Setup: Add alias and check it works.
	th.AddAlias("s1", th.redirectServer.URL+"/long1", "strong-secret")
	th.RedirectWorks("s1", "/long1")

	// Try to add alias request for s2 with wrong secret.
	// Check that redirect fails for s2.
	th.AddAlias("s2", th.redirectServer.URL+"/long1", "wrong-secret", FAILURE)
	th.RedirectIs404("s2")

	// Try to update long url for s1 with wrong secret.
	// Verify the update was not applied by checking redirect.
	th.UpdateOrigURL(
		"s1",                           // alias
		th.redirectServer.URL+"/long1", // oldLong
		th.redirectServer.URL+"/long2", // newLong
		"wrong-secret",                 // secret
		FAILURE,
	)
	th.RedirectWorks("s1", "/long1")

	// Try to update alias for s1 to s2 with wrong secret.
	// Verify the update was not applied by checking redirect.
	th.UpdateAlias("s1", "s2", "wrong-secret", FAILURE)
	th.RedirectWorks("s1", "/long1")
	th.RedirectIs404("s2")

	// Try to remove alias with wrong secret.
	// Verify the update was not applied by checking redirect.
	th.RemoveAlias("s1", "wrong-secret", FAILURE)
	th.RedirectWorks("s1", "/long1")
}
<|endoftext|>


<commit_before>package mqstore

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"os/exec"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/TeaMeow/KitSvc/module/logger"
	"github.com/TeaMeow/KitSvc/shared/mqutil"
	nsq "github.com/bitly/go-nsq"
	"github.com/parnurzeal/gorequest"
)

var (
	// AllConnected returns true when the message queue were all connected.
	AllConnected = false
	// SentTotal returns the total of the sent message.
	SentTotal = 0
	// RecvTotal returns the total of the received message.
	RecvTotal = 0
	// QueueTotal returns the total of the message are still in the queue.
	QueueTotal = 0
)

type mqstore struct {
	*nsq.Producer
	config      *nsq.Config
	isConnected bool
	queue       []message
}

// message represents a message.
type message struct {
	topic string
	body  []byte
}

// lg is a logger for NSQ.
type lg struct {
}

// NewProducer creates a new NSQ producer, and start to capturing the incoming messages.
func NewProducer(url, producer, prodHTTP string, lookupds []string, m *mqutil.Engine, deployed <-chan bool) *mqstore {
	// Ping the Event Store to make sure it's alive.
	if err := pingMQ(prodHTTP); err != nil {
		logger.Fatal(err)
	}
	config := nsq.NewConfig()
	// Create the producer.
	prod, err := nsq.NewProducer(producer, config)
	// Set the logger for the producer.
	prod.SetLogger(&lg{}, nsq.LogLevelDebug)
	if err != nil {
		logger.FatalFields("Error occurred while creating the NSQ producer.", logrus.Fields{
			"err":      err,
			"url":      url,
			"producer": producer,
			"http":     prodHTTP,
			"lookupds": lookupds,
		})
	}
	ms := &mqstore{
		Producer:    prod,
		config:      config,
		isConnected: true,
	}
	// Capturing the messages when the router was ready in the goroutine.
	go ms.capture(url, prodHTTP, lookupds, m, deployed)
	// Pushing the messages which are in the local queue to the remote message queue.
	go ms.push(prodHTTP)

	return ms
}

// Output the error to the global logger.
func (l *lg) Output(calldepth int, s string) error {
	typ := s[0:3]
	switch typ {
	case "DBG":
		logger.Debug(s[9:len(s)])
	case "INF":
		// Use debug so it won't be shown in the terminal.
		logger.Debug(s[9:len(s)])
	case "WRN":
		logger.Warning(s[9:len(s)])
	case "ERR":
		logger.Error(s[9:len(s)])
	}
	return nil
}

// pingMQ pings the NSQ with backoff to ensure
// a connection can be established before we proceed with the
// message queue setup and migration.
func pingMQ(addr string) error {
	for i := 0; i < 30; i++ {
		_, err := http.Get("http://" + addr)
		if err == nil {
			return nil
		}
		logger.Info("Cannot connect to NSQ producer, retry in 1 second.")
		time.Sleep(time.Second * 1)
	}

	return errors.New("Cannot connect to NSQ producer.")
}

// sendToRouter sends the received event data to self router.
func sendToRouter(method string, url string, json []byte) {
	// Send the request via the HTTP client.
	resp, _, err := gorequest.
		New().
		CustomMethod(method, url).
		Send(string(json)).
		End()
	if err != nil {
		logger.ErrorFields("Error occurred while sending the message to self router.", logrus.Fields{
			"err":    err,
			"method": method,
			"url":    url,
			"body":   json,
		})
	}
	if resp.StatusCode != 200 {
		logger.InfoFields("The message has been recevied by the router, but the status code wasn't 200.", logrus.Fields{
			"status": resp.StatusCode,
			"method": method,
			"url":    url,
			"body":   json,
		})
	}
}

// createTopic creates the new topic in the remote message queue,
// so we can subscribe to it.
func createTopic(httpProducer, topic string) {
	cmd := exec.Command("curl", "-X", "POST", fmt.Sprintf("http://%s/topic/create?topic=%s", httpProducer, topic))
	cmd.Start()
	cmd.Wait()
}

// capture the incoming events.
func (mq *mqstore) capture(url string, prodHTTP string, lookupds []string, m *mqutil.Engine, deployed <-chan bool) {
	// Continue if the router was ready.
	<-deployed
	// Each of the topic listener.
	for _, v := range m.Listeners {
		// Create the consumer.
		c, err := nsq.NewConsumer(v.Topic, v.Channel, nsq.NewConfig())
		// Set the logger for the consumer.
		c.SetLogger(&lg{}, nsq.LogLevelDebug)
		if err != nil {
			logger.FatalFields("Cannot create the NSQ consumer.", logrus.Fields{
				"channel": v.Channel,
				"topic":   v.Topic,
				"path":    v.Path,
				"method":  v.Method,
			})
		}

		// Create the topic to make sure it does exist before we subscribe to it.
		createTopic(prodHTTP, v.Topic)
		// Add the topic handler.
		c.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
			RecvTotal++
			// Send the received message to the self router,
			// so we can process it with Gin.
			sendToRouter(v.Method, url+v.Path, msg.Body)
			return nil
		}))

		// Connect to the NSQLookupds instead of a single NSQ node.
		if err := c.ConnectToNSQLookupds(lookupds); err != nil {
			logger.FatalFields("Cannot connect to the NSQ lookupds.", logrus.Fields{
				"err":      err,
				"lookupds": lookupds,
				"channel":  v.Channel,
				"topic":    v.Topic,
			})
		}
	}
}

// push the message which are in the queue to the remote message queue.
func (mq *mqstore) push(prodHTTP string) {
	for {
		<-time.After(time.Millisecond * 10)

		// Ping the NSQ Producer to see if it's back online or not.
		if !mq.isConnected {
			if err := pingMQ(prodHTTP); err == nil {
				mq.isConnected = true
				AllConnected = true

				logger.InfoFields("NSQ producer is back online, there are %d unsent messages that will begin to send.", logrus.Fields{
					"unsent": len(mq.queue),
				})
			}
			continue
		}

		// Skip if there's nothing in the queue.
		if len(mq.queue) == 0 {
			continue
		}

		// A downward loop for the queue.
		for i := len(mq.queue) - 1; i >= 0; i-- {
			m := mq.queue[i]
			// Wait a little bit for another event.
			<-time.After(time.Millisecond * 2)
			// Append the message in the topic.
			err := mq.Publish(m.topic, m.body)
			if err != nil {
				continue
			}
			QueueTotal--
			// Remove the message from the queue since it has been sent.
			mq.queue = append(mq.queue[:i], mq.queue[i+1:]...)
		}
	}
}

// send the event to the specified stream.
func (mq *mqstore) send(topic string, data interface{}) {
	body, err := json.Marshal(data)
	if err != nil {
		logger.ErrorFields("Error occurred while converting the data to json.", logrus.Fields{
			"err":   err,
			"topic": topic,
			"data":  data,
		})
	}
	// Counter.
	SentTotal++

	// Send the message to the remote message queue.
	if err := mq.Publish(topic, body); err != nil {
		switch t := err.(type) {
		case *net.OpError:
			// Push the message to the local queue if connecting refused.
			if t.Op == "dial" {
				// Mark the connection as lost.
				mq.isConnected = false
				AllConnected = false
				// Push the message to local queue.
				mq.queue = append([]message{message{topic, body}}, mq.queue...)
				// Counter.
				QueueTotal++
				logger.WarningFields("Message will be sent when the NSQ producer is back online.", logrus.Fields{
					"topic":  topic,
					"unsent": len(mq.queue),
				})
			}
		}
	}
}
<commit_msg>Fixed %d in unsent warning<commit_after>package mqstore

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"os/exec"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/TeaMeow/KitSvc/module/logger"
	"github.com/TeaMeow/KitSvc/shared/mqutil"
	nsq "github.com/bitly/go-nsq"
	"github.com/parnurzeal/gorequest"
)

var (
	// AllConnected returns true when the message queue were all connected.
	AllConnected = false
	// SentTotal returns the total of the sent message.
	SentTotal = 0
	// RecvTotal returns the total of the received message.
	RecvTotal = 0
	// QueueTotal returns the total of the message are still in the queue.
	QueueTotal = 0
)

type mqstore struct {
	*nsq.Producer
	config      *nsq.Config
	isConnected bool
	queue       []message
}

// message represents a message.
type message struct {
	topic string
	body  []byte
}

// lg is a logger for NSQ.
type lg struct {
}

// NewProducer creates a new NSQ producer, and start to capturing the incoming messages.
func NewProducer(url, producer, prodHTTP string, lookupds []string, m *mqutil.Engine, deployed <-chan bool) *mqstore {
	// Ping the Event Store to make sure it's alive.
	if err := pingMQ(prodHTTP); err != nil {
		logger.Fatal(err)
	}
	config := nsq.NewConfig()
	// Create the producer.
	prod, err := nsq.NewProducer(producer, config)
	// Set the logger for the producer.
	prod.SetLogger(&lg{}, nsq.LogLevelDebug)
	if err != nil {
		logger.FatalFields("Error occurred while creating the NSQ producer.", logrus.Fields{
			"err":      err,
			"url":      url,
			"producer": producer,
			"http":     prodHTTP,
			"lookupds": lookupds,
		})
	}
	ms := &mqstore{
		Producer:    prod,
		config:      config,
		isConnected: true,
	}
	// Capturing the messages when the router was ready in the goroutine.
	go ms.capture(url, prodHTTP, lookupds, m, deployed)
	// Pushing the messages which are in the local queue to the remote message queue.
	go ms.push(prodHTTP)

	return ms
}

// Output the error to the global logger.
func (l *lg) Output(calldepth int, s string) error {
	typ := s[0:3]
	switch typ {
	case "DBG":
		logger.Debug(s[9:len(s)])
	case "INF":
		// Use debug so it won't be shown in the terminal.
		logger.Debug(s[9:len(s)])
	case "WRN":
		logger.Warning(s[9:len(s)])
	case "ERR":
		logger.Error(s[9:len(s)])
	}
	return nil
}

// pingMQ pings the NSQ with backoff to ensure
// a connection can be established before we proceed with the
// message queue setup and migration.
func pingMQ(addr string) error {
	for i := 0; i < 30; i++ {
		_, err := http.Get("http://" + addr)
		if err == nil {
			return nil
		}
		logger.Info("Cannot connect to NSQ producer, retry in 1 second.")
		time.Sleep(time.Second * 1)
	}

	return errors.New("Cannot connect to NSQ producer.")
}

// sendToRouter sends the received event data to self router.
func sendToRouter(method string, url string, json []byte) {
	// Send the request via the HTTP client.
	resp, _, err := gorequest.
		New().
		CustomMethod(method, url).
		Send(string(json)).
		End()
	if err != nil {
		logger.ErrorFields("Error occurred while sending the message to self router.", logrus.Fields{
			"err":    err,
			"method": method,
			"url":    url,
			"body":   json,
		})
	}
	if resp.StatusCode != 200 {
		logger.InfoFields("The message has been recevied by the router, but the status code wasn't 200.", logrus.Fields{
			"status": resp.StatusCode,
			"method": method,
			"url":    url,
			"body":   json,
		})
	}
}

// createTopic creates the new topic in the remote message queue,
// so we can subscribe to it.
func createTopic(httpProducer, topic string) {
	cmd := exec.Command("curl", "-X", "POST", fmt.Sprintf("http://%s/topic/create?topic=%s", httpProducer, topic))
	cmd.Start()
	cmd.Wait()
}

// capture the incoming events.
func (mq *mqstore) capture(url string, prodHTTP string, lookupds []string, m *mqutil.Engine, deployed <-chan bool) {
	// Continue if the router was ready.
	<-deployed
	// Each of the topic listener.
	for _, v := range m.Listeners {
		// Create the consumer.
		c, err := nsq.NewConsumer(v.Topic, v.Channel, nsq.NewConfig())
		// Set the logger for the consumer.
		c.SetLogger(&lg{}, nsq.LogLevelDebug)
		if err != nil {
			logger.FatalFields("Cannot create the NSQ consumer.", logrus.Fields{
				"channel": v.Channel,
				"topic":   v.Topic,
				"path":    v.Path,
				"method":  v.Method,
			})
		}

		// Create the topic to make sure it does exist before we subscribe to it.
		createTopic(prodHTTP, v.Topic)
		// Add the topic handler.
		c.AddHandler(nsq.HandlerFunc(func(msg *nsq.Message) error {
			RecvTotal++
			// Send the received message to the self router,
			// so we can process it with Gin.
			sendToRouter(v.Method, url+v.Path, msg.Body)
			return nil
		}))

		// Connect to the NSQLookupds instead of a single NSQ node.
		if err := c.ConnectToNSQLookupds(lookupds); err != nil {
			logger.FatalFields("Cannot connect to the NSQ lookupds.", logrus.Fields{
				"err":      err,
				"lookupds": lookupds,
				"channel":  v.Channel,
				"topic":    v.Topic,
			})
		}
	}
}

// push the message which are in the queue to the remote message queue.
func (mq *mqstore) push(prodHTTP string) {
	for {
		<-time.After(time.Millisecond * 10)

		// Ping the NSQ Producer to see if it's back online or not.
		if !mq.isConnected {
			if err := pingMQ(prodHTTP); err == nil {
				mq.isConnected = true
				AllConnected = true

				logger.InfoFields("NSQ producer is back online, the unsent messages that will begin to send.", logrus.Fields{
					"unsent": len(mq.queue),
				})
			}
			continue
		}

		// Skip if there's nothing in the queue.
		if len(mq.queue) == 0 {
			continue
		}

		// A downward loop for the queue.
		for i := len(mq.queue) - 1; i >= 0; i-- {
			m := mq.queue[i]
			// Wait a little bit for another event.
			<-time.After(time.Millisecond * 2)
			// Append the message in the topic.
			err := mq.Publish(m.topic, m.body)
			if err != nil {
				continue
			}
			QueueTotal--
			// Remove the message from the queue since it has been sent.
			mq.queue = append(mq.queue[:i], mq.queue[i+1:]...)
		}
	}
}

// send the event to the specified stream.
func (mq *mqstore) send(topic string, data interface{}) {
	body, err := json.Marshal(data)
	if err != nil {
		logger.ErrorFields("Error occurred while converting the data to json.", logrus.Fields{
			"err":   err,
			"topic": topic,
			"data":  data,
		})
	}
	// Counter.
	SentTotal++

	// Send the message to the remote message queue.
	if err := mq.Publish(topic, body); err != nil {
		switch t := err.(type) {
		case *net.OpError:
			// Push the message to the local queue if connecting refused.
			if t.Op == "dial" {
				// Mark the connection as lost.
				mq.isConnected = false
				AllConnected = false
				// Push the message to local queue.
				mq.queue = append([]message{message{topic, body}}, mq.queue...)
				// Counter.
				QueueTotal++
				logger.WarningFields("Message will be sent when the NSQ producer is back online.", logrus.Fields{
					"topic":  topic,
					"unsent": len(mq.queue),
				})
			}
		}
	}
}
<|endoftext|>


<commit_before>package udp

import (
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
)

// Payload represents a single UDP payload.
type Payload struct {
	payload      *buf.Buffer
	source       net.Destination
	originalDest net.Destination
}

// PayloadHandler is function to handle Payload.
type PayloadHandler func(payload *buf.Buffer, source net.Destination, originalDest net.Destination)

// PayloadQueue is a queue of Payload.
type PayloadQueue struct {
	queue    []chan Payload
	callback PayloadHandler
}

// NewPayloadQueue returns a new PayloadQueue.
func NewPayloadQueue(option ListenOption) *PayloadQueue {
	queue := &PayloadQueue{
		callback: option.Callback,
		queue:    make([]chan Payload, option.Concurrency),
	}
	for i := range queue.queue {
		queue.queue[i] = make(chan Payload, 64)
		go queue.Dequeue(queue.queue[i])
	}
	return queue
}

// Enqueue adds the payload to the end of this queue.
func (q *PayloadQueue) Enqueue(payload Payload) {
	size := len(q.queue)
	idx := 0
	if size > 1 {
		idx = dice.Roll(size)
	}
	for i := 0; i < size; i++ {
		select {
		case q.queue[idx%size] <- payload:
			return
		default:
			idx++
		}
	}
}

func (q *PayloadQueue) Dequeue(queue <-chan Payload) {
	for payload := range queue {
		q.callback(payload.payload, payload.source, payload.originalDest)
	}
}

func (q *PayloadQueue) Close() error {
	for _, queue := range q.queue {
		close(queue)
	}
	return nil
}

type ListenOption struct {
	Callback            PayloadHandler
	ReceiveOriginalDest bool
	Concurrency         int
}

type Hub struct {
	conn   *net.UDPConn
	queue  *PayloadQueue
	option ListenOption
}

func ListenUDP(address net.Address, port net.Port, option ListenOption) (*Hub, error) {
	if option.Concurrency < 1 {
		option.Concurrency = 1
	}
	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
		IP:   address.IP(),
		Port: int(port),
	})
	if err != nil {
		return nil, err
	}
	newError("listening UDP on ", address, ":", port).WriteToLog()
	if option.ReceiveOriginalDest {
		rawConn, err := udpConn.SyscallConn()
		if err != nil {
			return nil, newError("failed to get fd").Base(err)
		}
		err = rawConn.Control(func(fd uintptr) {
			if err := SetOriginalDestOptions(int(fd)); err != nil {
				newError("failed to set socket options").Base(err).WriteToLog()
			}
		})
		if err != nil {
			return nil, newError("failed to control socket").Base(err)
		}
	}
	hub := &Hub{
		conn:   udpConn,
		queue:  NewPayloadQueue(option),
		option: option,
	}
	go hub.start()
	return hub, nil
}

func (h *Hub) Close() error {
	h.conn.Close()
	return nil
}

func (h *Hub) WriteTo(payload []byte, dest net.Destination) (int, error) {
	return h.conn.WriteToUDP(payload, &net.UDPAddr{
		IP:   dest.Address.IP(),
		Port: int(dest.Port),
	})
}

func (h *Hub) start() {
	oobBytes := make([]byte, 256)

	for {
		buffer := buf.New()
		var noob int
		var addr *net.UDPAddr
		err := buffer.AppendSupplier(func(b []byte) (int, error) {
			n, nb, _, a, e := ReadUDPMsg(h.conn, b, oobBytes)
			noob = nb
			addr = a
			return n, e
		})

		if err != nil {
			newError("failed to read UDP msg").Base(err).WriteToLog()
			buffer.Release()
			break
		}

		payload := Payload{
			payload: buffer,
		}
		payload.source = net.UDPDestination(net.IPAddress(addr.IP), net.Port(addr.Port))
		if h.option.ReceiveOriginalDest && noob > 0 {
			payload.originalDest = RetrieveOriginalDest(oobBytes[:noob])
			if payload.originalDest.IsValid() {
				newError("UDP original destination: ", payload.originalDest).AtDebug().WriteToLog()
			} else {
				newError("failed to read UDP original destination").WriteToLog()
			}
		}
		h.queue.Enqueue(payload)
	}
	h.queue.Close()
}

func (h *Hub) Addr() net.Addr {
	return h.conn.LocalAddr()
}
<commit_msg>comments<commit_after>package udp

import (
	"v2ray.com/core/common/buf"
	"v2ray.com/core/common/dice"
	"v2ray.com/core/common/net"
)

// Payload represents a single UDP payload.
type Payload struct {
	payload      *buf.Buffer
	source       net.Destination
	originalDest net.Destination
}

// PayloadHandler is function to handle Payload.
type PayloadHandler func(payload *buf.Buffer, source net.Destination, originalDest net.Destination)

// PayloadQueue is a queue of Payload.
type PayloadQueue struct {
	queue    []chan Payload
	callback PayloadHandler
}

// NewPayloadQueue returns a new PayloadQueue.
func NewPayloadQueue(option ListenOption) *PayloadQueue {
	queue := &PayloadQueue{
		callback: option.Callback,
		queue:    make([]chan Payload, option.Concurrency),
	}
	for i := range queue.queue {
		queue.queue[i] = make(chan Payload, 64)
		go queue.Dequeue(queue.queue[i])
	}
	return queue
}

// Enqueue adds the payload to the end of this queue.
func (q *PayloadQueue) Enqueue(payload Payload) {
	size := len(q.queue)
	idx := 0
	if size > 1 {
		idx = dice.Roll(size)
	}
	for i := 0; i < size; i++ {
		select {
		case q.queue[idx%size] <- payload:
			return
		default:
			idx++
		}
	}
}

func (q *PayloadQueue) Dequeue(queue <-chan Payload) {
	for payload := range queue {
		q.callback(payload.payload, payload.source, payload.originalDest)
	}
}

func (q *PayloadQueue) Close() error {
	for _, queue := range q.queue {
		close(queue)
	}
	return nil
}

type ListenOption struct {
	Callback            PayloadHandler
	ReceiveOriginalDest bool
	Concurrency         int
}

type Hub struct {
	conn   *net.UDPConn
	queue  *PayloadQueue
	option ListenOption
}

func ListenUDP(address net.Address, port net.Port, option ListenOption) (*Hub, error) {
	if option.Concurrency < 1 {
		option.Concurrency = 1
	}
	udpConn, err := net.ListenUDP("udp", &net.UDPAddr{
		IP:   address.IP(),
		Port: int(port),
	})
	if err != nil {
		return nil, err
	}
	newError("listening UDP on ", address, ":", port).WriteToLog()
	if option.ReceiveOriginalDest {
		rawConn, err := udpConn.SyscallConn()
		if err != nil {
			return nil, newError("failed to get fd").Base(err)
		}
		err = rawConn.Control(func(fd uintptr) {
			if err := SetOriginalDestOptions(int(fd)); err != nil {
				newError("failed to set socket options").Base(err).WriteToLog()
			}
		})
		if err != nil {
			return nil, newError("failed to control socket").Base(err)
		}
	}
	hub := &Hub{
		conn:   udpConn,
		queue:  NewPayloadQueue(option),
		option: option,
	}
	go hub.start()
	return hub, nil
}

// Close implements net.Listener.
func (h *Hub) Close() error {
	h.conn.Close()
	return nil
}

func (h *Hub) WriteTo(payload []byte, dest net.Destination) (int, error) {
	return h.conn.WriteToUDP(payload, &net.UDPAddr{
		IP:   dest.Address.IP(),
		Port: int(dest.Port),
	})
}

func (h *Hub) start() {
	oobBytes := make([]byte, 256)

	for {
		buffer := buf.New()
		var noob int
		var addr *net.UDPAddr
		err := buffer.AppendSupplier(func(b []byte) (int, error) {
			n, nb, _, a, e := ReadUDPMsg(h.conn, b, oobBytes)
			noob = nb
			addr = a
			return n, e
		})

		if err != nil {
			newError("failed to read UDP msg").Base(err).WriteToLog()
			buffer.Release()
			break
		}

		payload := Payload{
			payload: buffer,
		}
		payload.source = net.UDPDestination(net.IPAddress(addr.IP), net.Port(addr.Port))
		if h.option.ReceiveOriginalDest && noob > 0 {
			payload.originalDest = RetrieveOriginalDest(oobBytes[:noob])
			if payload.originalDest.IsValid() {
				newError("UDP original destination: ", payload.originalDest).AtDebug().WriteToLog()
			} else {
				newError("failed to read UDP original destination").WriteToLog()
			}
		}
		h.queue.Enqueue(payload)
	}
	h.queue.Close()
}

// Addr implements net.Listener.
func (h *Hub) Addr() net.Addr {
	return h.conn.LocalAddr()
}
<|endoftext|>


<commit_before>package json

import (
	"bytes"
	"encoding/json"
	"io"
	"testing"
)

func TestTokens(t *testing.T) {
	j := []byte(`
	{
		"foo": 1,
		"bar": 2,
		"test": "Hello, world!",
		"baz": 123.1,
		"array": [
			{"foo": 1},
			{"bar": 2},
			{"baz": 3}
		],
		"subobj": {
			"foo": 1,
			"subarray": [1,2,3],
			"subsubobj": {
				"bar": 2,
				"baz": 3,
				"array": ["hello", "world"]
			}
		},
		"bool": true
	}`)

	expPaths := []JsonPath{
		{""},
		{"foo"}, {"foo"},
		{"bar"}, {"bar"},
		{"test"}, {"test"},
		{"baz"}, {"baz"},
		{"array"}, {"array", -1},
		{"array", 0, ""}, {"array", 0, "foo"}, {"array", 0, "foo"}, {"array", 0},
		{"array", 1, ""}, {"array", 1, "bar"}, {"array", 1, "bar"}, {"array", 1},
		{"array", 2, ""}, {"array", 2, "baz"}, {"array", 2, "baz"}, {"array", 2},
		{"array"},
		{"subobj"},
		{"subobj", ""},
		{"subobj", "foo"}, {"subobj", "foo"},
		{"subobj", "subarray"}, {"subobj", "subarray", -1},
		{"subobj", "subarray", 0}, {"subobj", "subarray", 1}, {"subobj", "subarray", 2},
		{"subobj", "subarray"},
		{"subobj", "subsubobj"}, {"subobj", "subsubobj", ""},
		{"subobj", "subsubobj", "bar"}, {"subobj", "subsubobj", "bar"},
		{"subobj", "subsubobj", "baz"}, {"subobj", "subsubobj", "baz"},
		{"subobj", "subsubobj", "array"}, {"subobj", "subsubobj", "array", -1},
		{"subobj", "subsubobj", "array", 0}, {"subobj", "subsubobj", "array", 1},
		{"subobj", "subsubobj", "array"},
		{"subobj", "subsubobj"}, {"subobj"},
		{"bool"}, {"bool"},
		{},
	}

	expTokens := []json.Token{
		json.Delim('{'),
		KeyString("foo"), float64(1),
		KeyString("bar"), float64(2),
		KeyString("test"), "Hello, world!",
		KeyString("baz"), float64(123.1),
		KeyString("array"), json.Delim('['),
		json.Delim('{'), KeyString("foo"), float64(1), json.Delim('}'),
		json.Delim('{'), KeyString("bar"), float64(2), json.Delim('}'),
		json.Delim('{'), KeyString("baz"), float64(3), json.Delim('}'),
		json.Delim(']'),
		KeyString("subobj"),
		json.Delim('{'),
		KeyString("foo"), float64(1),
		KeyString("subarray"), json.Delim('['),
		float64(1), float64(2), float64(3),
		json.Delim(']'),
		KeyString("subsubobj"), json.Delim('{'),
		KeyString("bar"), float64(2),
		KeyString("baz"), float64(3),
		KeyString("array"), json.Delim('['),
		"hello", "world",
		json.Delim(']'),
		json.Delim('}'), json.Delim('}'),
		KeyString("bool"), true,
		json.Delim('}'),
	}
	outTokens := []json.Token{}
	outPaths := []JsonPath{}

	d := NewDecoder(bytes.NewBuffer(j))

	for {
		st, err := d.Token()
		if err == io.EOF {
			break
		} else if err != nil {
			t.Error(err)
		}
		outTokens = append(outTokens, st)
		outPaths = append(outPaths, d.Path())
	}

	// Check tokens
	if len(outTokens) != len(expTokens) {
		t.Errorf("Out has %v elements, expected %v", len(outTokens), len(expTokens))
		return
	}
	for i, v := range expTokens {
		if v != outTokens[i] {
			t.Errorf("@%v exp: %T:%v but was: %T:%v", i, v, v, outTokens[i], outTokens[i])
		}
	}

	// Check paths
	if len(outPaths) != len(expPaths) {
		t.Errorf("Outpaths has %v elements, expected %v", len(outPaths), len(expPaths))
	}
	for i, v := range expPaths {
		if !v.Equal(outPaths[i]) {
			t.Errorf("@%v exp: %T:%v but was: %T:%v", i, v, v, outPaths[i], outPaths[i].String())
		}
	}
}
<commit_msg>Another array test.<commit_after>package json

import (
	"bytes"
	"encoding/json"
	"io"
	"testing"
)

func TestTokens(t *testing.T) {
	j := []byte(`
	{
		"foo": 1,
		"bar": 2,
		"test": "Hello, world!",
		"baz": 123.1,
		"array": [
			{"foo": 1},
			{"bar": 2},
			{"baz": 3}
		],
		"subobj": {
			"foo": 1,
			"subarray": [1,2,3],
			"subsubobj": {
				"bar": 2,
				"baz": 3,
				"array": ["hello", "world"]
			}
		},
		"bool": true,
		"a": [[0],[[1]]]
	}`)

	expPaths := []JsonPath{
		{""},
		{"foo"}, {"foo"},
		{"bar"}, {"bar"},
		{"test"}, {"test"},
		{"baz"}, {"baz"},
		{"array"}, {"array", -1},
		{"array", 0, ""}, {"array", 0, "foo"}, {"array", 0, "foo"}, {"array", 0},
		{"array", 1, ""}, {"array", 1, "bar"}, {"array", 1, "bar"}, {"array", 1},
		{"array", 2, ""}, {"array", 2, "baz"}, {"array", 2, "baz"}, {"array", 2},
		{"array"},
		{"subobj"},
		{"subobj", ""},
		{"subobj", "foo"}, {"subobj", "foo"},
		{"subobj", "subarray"}, {"subobj", "subarray", -1},
		{"subobj", "subarray", 0}, {"subobj", "subarray", 1}, {"subobj", "subarray", 2},
		{"subobj", "subarray"},
		{"subobj", "subsubobj"}, {"subobj", "subsubobj", ""},
		{"subobj", "subsubobj", "bar"}, {"subobj", "subsubobj", "bar"},
		{"subobj", "subsubobj", "baz"}, {"subobj", "subsubobj", "baz"},
		{"subobj", "subsubobj", "array"}, {"subobj", "subsubobj", "array", -1},
		{"subobj", "subsubobj", "array", 0}, {"subobj", "subsubobj", "array", 1},
		{"subobj", "subsubobj", "array"},
		{"subobj", "subsubobj"}, {"subobj"},
		{"bool"}, {"bool"},
		{"a"}, {"a", -1}, {"a", 0, -1}, {"a", 0, 0}, {"a", 0}, {"a", 1, -1},
		{"a", 1, 0, -1}, {"a", 1, 0, 0}, {"a", 1, 0}, {"a", 1}, {"a"},
		{},
	}

	expTokens := []json.Token{
		json.Delim('{'),
		KeyString("foo"), float64(1),
		KeyString("bar"), float64(2),
		KeyString("test"), "Hello, world!",
		KeyString("baz"), float64(123.1),
		KeyString("array"), json.Delim('['),
		json.Delim('{'), KeyString("foo"), float64(1), json.Delim('}'),
		json.Delim('{'), KeyString("bar"), float64(2), json.Delim('}'),
		json.Delim('{'), KeyString("baz"), float64(3), json.Delim('}'),
		json.Delim(']'),
		KeyString("subobj"),
		json.Delim('{'),
		KeyString("foo"), float64(1),
		KeyString("subarray"), json.Delim('['),
		float64(1), float64(2), float64(3),
		json.Delim(']'),
		KeyString("subsubobj"), json.Delim('{'),
		KeyString("bar"), float64(2),
		KeyString("baz"), float64(3),
		KeyString("array"), json.Delim('['),
		"hello", "world",
		json.Delim(']'),
		json.Delim('}'), json.Delim('}'),
		KeyString("bool"), true,
		KeyString("a"),
		json.Delim('['), json.Delim('['), float64(0), json.Delim(']'),
		json.Delim('['), json.Delim('['), float64(1), json.Delim(']'), json.Delim(']'),
		json.Delim(']'),
		json.Delim('}'),
	}
	outTokens := []json.Token{}
	outPaths := []JsonPath{}

	d := NewDecoder(bytes.NewBuffer(j))

	for {
		st, err := d.Token()
		if err == io.EOF {
			break
		} else if err != nil {
			t.Error(err)
			break
		}
		outTokens = append(outTokens, st)
		outPaths = append(outPaths, d.Path())
	}

	// Check tokens
	if len(outTokens) != len(expTokens) {
		t.Errorf("Out has %v elements, expected %v", len(outTokens), len(expTokens))
	}
	for i, v := range expTokens {
		if v != outTokens[i] {
			t.Errorf("@%v exp: %T:%v but was: %T:%v", i, v, v, outTokens[i], outTokens[i])
		}
	}

	// Check paths
	if len(outPaths) != len(expPaths) {
		t.Errorf("Outpaths has %v elements, expected %v", len(outPaths), len(expPaths))
	}
	for i, v := range expPaths {
		if !v.Equal(outPaths[i]) {
			t.Errorf("@%v exp: %T:%v but was: %T:%v", i, v, v, outPaths[i], outPaths[i].String())
		}
	}
}
<|endoftext|>


<commit_before>package host

import (
// "testing"
)

/*
// TestSaveLoad tests that saving and loading a Host restores its data.
func TestSaveLoad(t *testing.T) {
	ht := CreateHostTester("TestSaveLoad", t)
	ht.testObligation()
	err := ht.save()
	if err != nil {
		ht.Fatal(err)
	}
	err = ht.load()
	if err != nil {
		ht.Fatal(err)
	}
}
*/
<commit_msg>restore host persist test<commit_after>package host

import (
	"testing"
)

// TestSaveLoad tests that saving and loading a Host restores its data.
func TestSaveLoad(t *testing.T) {
	ht := CreateHostTester("TestSaveLoad", t)
	ht.testAllocation()
	err := ht.host.save()
	if err != nil {
		ht.t.Fatal(err)
	}
	err = ht.host.load()
	if err != nil {
		ht.t.Fatal(err)
	}
}
<|endoftext|>


<commit_before>package db

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
	"io/ioutil"
	"kubernetes/Godeps/_workspace/src/github.com/ghodss/yaml"
	"log"
	"os"
)

type Configure struct {
	Mysqldbhost     string `json:"mysqldbhost"`
	Mysqldbport     string `json:"mysqldbport"`
	Mysqldbname     string `json:"mysqldbname"`
	Mysqldbusername string `json:"mysqldbusername"`
	Mysqldbpassword string `json:"mysqldbpassword"`
}

func Connect() *sql.DB {
	dir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	configureFileData, err := ioutil.ReadFile(dir + "/conf/app.conf.yml")
	if err != nil {
		log.Fatal(err)
	}
	var configure Configure
	err = yaml.Unmarshal([]byte(configureFileData), &configure)
	if err != nil {
		log.Fatal(err)
	}

	dsn := configure.Mysqldbusername + ":" + configure.Mysqldbpassword + "@" + "tcp(" + configure.Mysqldbhost +
		":" + configure.Mysqldbport + ")" + "/" + configure.Mysqldbname + "?charset=utf8"
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		log.Fatal(err)
	}
	return db
}
<commit_msg>update use yaml package<commit_after>package db

import (
	"database/sql"
	_ "github.com/go-sql-driver/mysql"
	"gopkg.in/yaml.v2"
	"io/ioutil"
	"log"
	"os"
)

type Configure struct {
	Mysqldbhost     string `json:"mysqldbhost"`
	Mysqldbport     string `json:"mysqldbport"`
	Mysqldbname     string `json:"mysqldbname"`
	Mysqldbusername string `json:"mysqldbusername"`
	Mysqldbpassword string `json:"mysqldbpassword"`
}

func Connect() *sql.DB {
	dir, err := os.Getwd()
	if err != nil {
		log.Fatal(err)
	}
	configureFileData, err := ioutil.ReadFile(dir + "/conf/app.conf.yml")
	if err != nil {
		log.Fatal(err)
	}
	var configure Configure
	err = yaml.Unmarshal([]byte(configureFileData), &configure)
	if err != nil {
		log.Fatal(err)
	}

	dsn := configure.Mysqldbusername + ":" + configure.Mysqldbpassword + "@" + "tcp(" + configure.Mysqldbhost +
		":" + configure.Mysqldbport + ")" + "/" + configure.Mysqldbname + "?charset=utf8"
	db, err := sql.Open("mysql", dsn)
	if err != nil {
		log.Fatal(err)
	}
	return db
}
<|endoftext|>


<commit_before>package eg_test

import (
	"bytes"
	"flag"
	"go/parser"
	"go/token"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"code.google.com/p/go.tools/go/exact"
	"code.google.com/p/go.tools/go/loader"
	"code.google.com/p/go.tools/go/types"
	"code.google.com/p/go.tools/refactor/eg"
)

// TODO(adonovan): more tests:
// - of command-line tool
// - of all parts of syntax
// - of applying a template to a package it imports:
//   the replacement syntax should use unqualified names for its objects.

var (
	updateFlag  = flag.Bool("update", false, "update the golden files")
	verboseFlag = flag.Bool("verbose", false, "show matcher information")
)

func Test(t *testing.T) {
	switch runtime.GOOS {
	case "windows":
		t.Skipf("skipping test on %q (no /usr/bin/diff)", runtime.GOOS)
	}

	conf := loader.Config{
		Fset:          token.NewFileSet(),
		ParserMode:    parser.ParseComments,
		SourceImports: true,
	}

	// Each entry is a single-file package.
	// (Multi-file packages aren't interesting for this test.)
	// Order matters: each non-template package is processed using
	// the preceding template package.
	for _, filename := range []string{
		"testdata/A.template",
		"testdata/A1.go",
		"testdata/A2.go",

		"testdata/B.template",
		"testdata/B1.go",

		"testdata/C.template",
		"testdata/C1.go",

		"testdata/D.template",
		"testdata/D1.go",

		"testdata/E.template",
		"testdata/E1.go",

		"testdata/bad_type.template",
		"testdata/no_before.template",
		"testdata/no_after_return.template",
		"testdata/type_mismatch.template",
		"testdata/expr_type_mismatch.template",
	} {
		pkgname := strings.TrimSuffix(filepath.Base(filename), ".go")
		if err := conf.CreateFromFilenames(pkgname, filename); err != nil {
			t.Fatal(err)
		}
	}
	iprog, err := conf.Load()
	if err != nil {
		t.Fatal(err)
	}

	var xform *eg.Transformer
	for _, info := range iprog.Created {
		file := info.Files[0]
		filename := iprog.Fset.File(file.Pos()).Name() // foo.go

		if strings.HasSuffix(filename, "template") {
			// a new template
			shouldFail, _ := info.Pkg.Scope().Lookup("shouldFail").(*types.Const)
			xform, err = eg.NewTransformer(iprog.Fset, info, *verboseFlag)
			if err != nil {
				if shouldFail == nil {
					t.Errorf("NewTransformer(%s): %s", filename, err)
				} else if want := exact.StringVal(shouldFail.Val()); !strings.Contains(err.Error(), want) {
					t.Errorf("NewTransformer(%s): got error %q, want error %q", filename, err, want)
				}
			} else if shouldFail != nil {
				t.Errorf("NewTransformer(%s) succeeded unexpectedly; want error %q",
					filename, shouldFail.Val())
			}
			continue
		}

		if xform == nil {
			t.Errorf("%s: no previous template", filename)
			continue
		}

		// apply previous template to this package
		n := xform.Transform(&info.Info, info.Pkg, file)
		if n == 0 {
			t.Errorf("%s: no matches", filename)
			continue
		}

		got := filename + "t"       // foo.got
		golden := filename + "lden" // foo.golden

		// Write actual output to foo.got.
		if err := eg.WriteAST(iprog.Fset, got, file); err != nil {
			t.Error(err)
		}
		defer os.Remove(got)

		// Compare foo.got with foo.golden.
		var cmd *exec.Cmd
		switch runtime.GOOS {
		case "plan9":
			cmd = exec.Command("/bin/diff", "-c", golden, got)
		default:
			cmd = exec.Command("/usr/bin/diff", "-u", "-N", golden, got)
		}
		buf := new(bytes.Buffer)
		cmd.Stdout = buf
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			t.Errorf("eg tests for %s failed: %s.\n%s\n", filename, err, buf)

			if *updateFlag {
				t.Logf("Updating %s...", golden)
				if err := exec.Command("/bin/cp", got, golden).Run(); err != nil {
					t.Errorf("Update failed: %s", err)
				}
			}
		}
	}
}
<commit_msg>go.tools/refactor/eg: avoid nonstandard 'diff -N' option.<commit_after>package eg_test

import (
	"bytes"
	"flag"
	"go/parser"
	"go/token"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"
	"testing"

	"code.google.com/p/go.tools/go/exact"
	"code.google.com/p/go.tools/go/loader"
	"code.google.com/p/go.tools/go/types"
	"code.google.com/p/go.tools/refactor/eg"
)

// TODO(adonovan): more tests:
// - of command-line tool
// - of all parts of syntax
// - of applying a template to a package it imports:
//   the replacement syntax should use unqualified names for its objects.

var (
	updateFlag  = flag.Bool("update", false, "update the golden files")
	verboseFlag = flag.Bool("verbose", false, "show matcher information")
)

func Test(t *testing.T) {
	switch runtime.GOOS {
	case "windows":
		t.Skipf("skipping test on %q (no /usr/bin/diff)", runtime.GOOS)
	}

	conf := loader.Config{
		Fset:          token.NewFileSet(),
		ParserMode:    parser.ParseComments,
		SourceImports: true,
	}

	// Each entry is a single-file package.
	// (Multi-file packages aren't interesting for this test.)
	// Order matters: each non-template package is processed using
	// the preceding template package.
	for _, filename := range []string{
		"testdata/A.template",
		"testdata/A1.go",
		"testdata/A2.go",

		"testdata/B.template",
		"testdata/B1.go",

		"testdata/C.template",
		"testdata/C1.go",

		"testdata/D.template",
		"testdata/D1.go",

		"testdata/E.template",
		"testdata/E1.go",

		"testdata/bad_type.template",
		"testdata/no_before.template",
		"testdata/no_after_return.template",
		"testdata/type_mismatch.template",
		"testdata/expr_type_mismatch.template",
	} {
		pkgname := strings.TrimSuffix(filepath.Base(filename), ".go")
		if err := conf.CreateFromFilenames(pkgname, filename); err != nil {
			t.Fatal(err)
		}
	}
	iprog, err := conf.Load()
	if err != nil {
		t.Fatal(err)
	}

	var xform *eg.Transformer
	for _, info := range iprog.Created {
		file := info.Files[0]
		filename := iprog.Fset.File(file.Pos()).Name() // foo.go

		if strings.HasSuffix(filename, "template") {
			// a new template
			shouldFail, _ := info.Pkg.Scope().Lookup("shouldFail").(*types.Const)
			xform, err = eg.NewTransformer(iprog.Fset, info, *verboseFlag)
			if err != nil {
				if shouldFail == nil {
					t.Errorf("NewTransformer(%s): %s", filename, err)
				} else if want := exact.StringVal(shouldFail.Val()); !strings.Contains(err.Error(), want) {
					t.Errorf("NewTransformer(%s): got error %q, want error %q", filename, err, want)
				}
			} else if shouldFail != nil {
				t.Errorf("NewTransformer(%s) succeeded unexpectedly; want error %q",
					filename, shouldFail.Val())
			}
			continue
		}

		if xform == nil {
			t.Errorf("%s: no previous template", filename)
			continue
		}

		// apply previous template to this package
		n := xform.Transform(&info.Info, info.Pkg, file)
		if n == 0 {
			t.Errorf("%s: no matches", filename)
			continue
		}

		got := filename + "t"       // foo.got
		golden := filename + "lden" // foo.golden

		// Write actual output to foo.got.
		if err := eg.WriteAST(iprog.Fset, got, file); err != nil {
			t.Error(err)
		}
		defer os.Remove(got)

		// Compare foo.got with foo.golden.
		var cmd *exec.Cmd
		switch runtime.GOOS {
		case "plan9":
			cmd = exec.Command("/bin/diff", "-c", golden, got)
		default:
			cmd = exec.Command("/usr/bin/diff", "-u", golden, got)
		}
		buf := new(bytes.Buffer)
		cmd.Stdout = buf
		cmd.Stderr = os.Stderr
		if err := cmd.Run(); err != nil {
			t.Errorf("eg tests for %s failed: %s.\n%s\n", filename, err, buf)

			if *updateFlag {
				t.Logf("Updating %s...", golden)
				if err := exec.Command("/bin/cp", got, golden).Run(); err != nil {
					t.Errorf("Update failed: %s", err)
				}
			}
		}
	}
}
<|endoftext|>


<commit_before>// Copyright 2017 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package topdown

import (
	"encoding/json"
	"fmt"
	"math/big"
	"strconv"
	"time"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/topdown/builtins"
)

type nowKeyID string

var nowKey = nowKeyID("time.now_ns")
var tzCache map[string]*time.Location

func builtinTimeNowNanos(bctx BuiltinContext, _ []*ast.Term, iter func(*ast.Term) error) error {

	exist, ok := bctx.Cache.Get(nowKey)
	var now *ast.Term

	if !ok {
		curr := time.Now()
		now = ast.NewTerm(ast.Number(int64ToJSONNumber(curr.UnixNano())))
		bctx.Cache.Put(nowKey, now)
	} else {
		now = exist.(*ast.Term)
	}

	return iter(now)
}

func builtinTimeParseNanos(a, b ast.Value) (ast.Value, error) {

	format, err := builtins.StringOperand(a, 1)
	if err != nil {
		return nil, err
	}

	value, err := builtins.StringOperand(b, 2)
	if err != nil {
		return nil, err
	}

	result, err := time.Parse(string(format), string(value))
	if err != nil {
		return nil, err
	}

	return ast.Number(int64ToJSONNumber(result.UnixNano())), nil
}

func builtinTimeParseRFC3339Nanos(a ast.Value) (ast.Value, error) {

	value, err := builtins.StringOperand(a, 1)
	if err != nil {
		return nil, err
	}

	result, err := time.Parse(time.RFC3339, string(value))
	if err != nil {
		return nil, err
	}

	return ast.Number(int64ToJSONNumber(result.UnixNano())), nil
}
func builtinParseDurationNanos(a ast.Value) (ast.Value, error) {

	duration, err := builtins.StringOperand(a, 1)
	if err != nil {
		return nil, err
	}
	value, err := time.ParseDuration(string(duration))
	if err != nil {
		return nil, err
	}
	return ast.Number(int64ToJSONNumber(int64(value))), nil
}

func builtinDate(a ast.Value) (ast.Value, error) {
	t, err := tzTime(a)
	if err != nil {
		return nil, err
	}
	year, month, day := t.Date()
	result := ast.Array{ast.IntNumberTerm(year), ast.IntNumberTerm(int(month)), ast.IntNumberTerm(day)}
	return result, nil
}

func builtinClock(a ast.Value) (ast.Value, error) {
	t, err := tzTime(a)
	if err != nil {
		return nil, err
	}
	hour, minute, second := t.Clock()
	result := ast.Array{ast.IntNumberTerm(hour), ast.IntNumberTerm(minute), ast.IntNumberTerm(second)}
	return result, nil
}

func builtinWeekday(a ast.Value) (ast.Value, error) {
	t, err := tzTime(a)
	if err != nil {
		return nil, err
	}
	weekday := t.Weekday().String()
	return ast.String(weekday), nil
}

func tzTime(a ast.Value) (t time.Time, err error) {
	var nVal ast.Value
	loc := time.UTC

	switch va := a.(type) {
	case ast.Array:

		// not sure what error to return if passed an empty array, but if this isn't handled it panics if passed an empty array
		// if len(va) == 0 {
		// 	return time.Time{}, builtins.NewOperandElementErr(<something>)
		// }

		nVal, err = builtins.NumberOperand(va[0].Value, 1)
		if err != nil {
			return time.Time{}, err
		}

		if len(va) > 1 {
			tzVal, err := builtins.StringOperand(va[1].Value, 1)
			if err != nil {
				return time.Time{}, err
			}

			tzName := string(tzVal)

			switch tzName {
			case "", "UTC":
				// loc is already UTC
			case "Local":
				loc = time.Local

			default:
				var ok bool
				loc, ok = tzCache[tzName]

				if !ok {
					loc, err = time.LoadLocation(tzName)
					if err != nil {
						return time.Time{}, err
					}
					tzCache[tzName] = loc
				}
			}
		}

	case ast.Number:
		nVal = a
	}

	value, err := builtins.NumberOperand(nVal, 1)
	if err != nil {
		return time.Time{}, err
	}

	f := builtins.NumberToFloat(value)
	i64, acc := f.Int64()
	if acc != big.Exact {
		return time.Time{}, fmt.Errorf("timestamp too big")
	}

	t = time.Unix(0, i64).In(loc)

	return t, nil
}

func int64ToJSONNumber(i int64) json.Number {
	return json.Number(strconv.FormatInt(i, 10))
}

func init() {
	RegisterBuiltinFunc(ast.NowNanos.Name, builtinTimeNowNanos)
	RegisterFunctionalBuiltin1(ast.ParseRFC3339Nanos.Name, builtinTimeParseRFC3339Nanos)
	RegisterFunctionalBuiltin2(ast.ParseNanos.Name, builtinTimeParseNanos)
	RegisterFunctionalBuiltin1(ast.ParseDurationNanos.Name, builtinParseDurationNanos)
	RegisterFunctionalBuiltin1(ast.Date.Name, builtinDate)
	RegisterFunctionalBuiltin1(ast.Clock.Name, builtinClock)
	RegisterFunctionalBuiltin1(ast.Weekday.Name, builtinWeekday)
	tzCache = make(map[string]*time.Location)
}
<commit_msg>Add additional error handling<commit_after>// Copyright 2017 The OPA Authors.  All rights reserved.
// Use of this source code is governed by an Apache2
// license that can be found in the LICENSE file.

package topdown

import (
	"encoding/json"
	"fmt"
	"math/big"
	"strconv"
	"time"

	"github.com/open-policy-agent/opa/ast"
	"github.com/open-policy-agent/opa/topdown/builtins"
)

type nowKeyID string

var nowKey = nowKeyID("time.now_ns")
var tzCache map[string]*time.Location

func builtinTimeNowNanos(bctx BuiltinContext, _ []*ast.Term, iter func(*ast.Term) error) error {

	exist, ok := bctx.Cache.Get(nowKey)
	var now *ast.Term

	if !ok {
		curr := time.Now()
		now = ast.NewTerm(ast.Number(int64ToJSONNumber(curr.UnixNano())))
		bctx.Cache.Put(nowKey, now)
	} else {
		now = exist.(*ast.Term)
	}

	return iter(now)
}

func builtinTimeParseNanos(a, b ast.Value) (ast.Value, error) {

	format, err := builtins.StringOperand(a, 1)
	if err != nil {
		return nil, err
	}

	value, err := builtins.StringOperand(b, 2)
	if err != nil {
		return nil, err
	}

	result, err := time.Parse(string(format), string(value))
	if err != nil {
		return nil, err
	}

	return ast.Number(int64ToJSONNumber(result.UnixNano())), nil
}

func builtinTimeParseRFC3339Nanos(a ast.Value) (ast.Value, error) {

	value, err := builtins.StringOperand(a, 1)
	if err != nil {
		return nil, err
	}

	result, err := time.Parse(time.RFC3339, string(value))
	if err != nil {
		return nil, err
	}

	return ast.Number(int64ToJSONNumber(result.UnixNano())), nil
}
func builtinParseDurationNanos(a ast.Value) (ast.Value, error) {

	duration, err := builtins.StringOperand(a, 1)
	if err != nil {
		return nil, err
	}
	value, err := time.ParseDuration(string(duration))
	if err != nil {
		return nil, err
	}
	return ast.Number(int64ToJSONNumber(int64(value))), nil
}

func builtinDate(a ast.Value) (ast.Value, error) {
	t, err := tzTime(a)
	if err != nil {
		return nil, err
	}
	year, month, day := t.Date()
	result := ast.Array{ast.IntNumberTerm(year), ast.IntNumberTerm(int(month)), ast.IntNumberTerm(day)}
	return result, nil
}

func builtinClock(a ast.Value) (ast.Value, error) {
	t, err := tzTime(a)
	if err != nil {
		return nil, err
	}
	hour, minute, second := t.Clock()
	result := ast.Array{ast.IntNumberTerm(hour), ast.IntNumberTerm(minute), ast.IntNumberTerm(second)}
	return result, nil
}

func builtinWeekday(a ast.Value) (ast.Value, error) {
	t, err := tzTime(a)
	if err != nil {
		return nil, err
	}
	weekday := t.Weekday().String()
	return ast.String(weekday), nil
}

func tzTime(a ast.Value) (t time.Time, err error) {
	var nVal ast.Value
	loc := time.UTC

	switch va := a.(type) {
	case ast.Array:

		if len(va) == 0 {
			return time.Time{}, builtins.NewOperandTypeErr(1, a, "either number (ns) or [number (ns), string (tz)]")
		}

		nVal, err = builtins.NumberOperand(va[0].Value, 1)
		if err != nil {
			return time.Time{}, err
		}

		if len(va) > 1 {
			tzVal, err := builtins.StringOperand(va[1].Value, 1)
			if err != nil {
				return time.Time{}, err
			}

			tzName := string(tzVal)

			switch tzName {
			case "", "UTC":
				// loc is already UTC
			case "Local":
				loc = time.Local

			default:
				var ok bool
				loc, ok = tzCache[tzName]

				if !ok {
					loc, err = time.LoadLocation(tzName)
					if err != nil {
						return time.Time{}, err
					}
					tzCache[tzName] = loc
				}
			}
		}

	case ast.Number:
		nVal = a

	default:
		return time.Time{}, builtins.NewOperandTypeErr(1, a, "either number (ns) or [number (ns), string (tz)]")
	}

	value, err := builtins.NumberOperand(nVal, 1)
	if err != nil {
		return time.Time{}, err
	}

	f := builtins.NumberToFloat(value)
	i64, acc := f.Int64()
	if acc != big.Exact {
		return time.Time{}, fmt.Errorf("timestamp too big")
	}

	t = time.Unix(0, i64).In(loc)

	return t, nil
}

func int64ToJSONNumber(i int64) json.Number {
	return json.Number(strconv.FormatInt(i, 10))
}

func init() {
	RegisterBuiltinFunc(ast.NowNanos.Name, builtinTimeNowNanos)
	RegisterFunctionalBuiltin1(ast.ParseRFC3339Nanos.Name, builtinTimeParseRFC3339Nanos)
	RegisterFunctionalBuiltin2(ast.ParseNanos.Name, builtinTimeParseNanos)
	RegisterFunctionalBuiltin1(ast.ParseDurationNanos.Name, builtinParseDurationNanos)
	RegisterFunctionalBuiltin1(ast.Date.Name, builtinDate)
	RegisterFunctionalBuiltin1(ast.Clock.Name, builtinClock)
	RegisterFunctionalBuiltin1(ast.Weekday.Name, builtinWeekday)
	tzCache = make(map[string]*time.Location)
}
<|endoftext|>


<commit_before>//+build !local_tests

package myirmaserver

import (
	"path/filepath"
	"testing"

	"github.com/privacybydesign/irmago/internal/test"
	"github.com/stretchr/testify/assert"
)

func validConfWithEmail(t *testing.T) *Configuration {
	testdataPath := test.FindTestdataFolder(t)
	conf := validConf(t)
	conf.EmailServer = "localhost:1025"
	conf.DefaultLanguage = "en"
	conf.LoginEmailFiles = map[string]string{"en": filepath.Join(testdataPath, "emailtemplate.html")}
	conf.LoginEmailSubject = map[string]string{"en": "testsubject"}
	conf.LoginEmailBaseURL = map[string]string{"en": "localhost:8000/test/"}
	conf.DeleteEmailFiles = map[string]string{"en": filepath.Join(testdataPath, "emailtemplate.html")}
	conf.DeleteEmailSubject = map[string]string{"en": "testsubject"}
	conf.DeleteAccountFiles = map[string]string{"en": filepath.Join(testdataPath, "emailtemplate.html")}
	conf.DeleteAccountSubject = map[string]string{"en": "testsubject"}
	return conf
}

func TestConfEmailValidation(t *testing.T) {
	testdataPath := test.FindTestdataFolder(t)

	conf := validConf(t)
	conf.DBType = DatabaseTypePostgres
	conf.DBConnstring = "postgresql://localhost:5432/test"
	_, err := New(conf)
	assert.NoError(t, err)

	_, err = New(validConfWithEmail(t))
	assert.NoError(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailFiles = map[string]string{"en": filepath.Join(testdataPath, "invalidemailtemplate.html")}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailFiles = map[string]string{}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailSubject = map[string]string{}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailBaseURL = map[string]string{}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.DeleteEmailFiles = map[string]string{"de": filepath.Join(testdataPath, "emailtemplate.html")}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.DeleteEmailSubject = map[string]string{"de": "testsubject"}
	_, err = New(conf)
	assert.Error(t, err)
}
<commit_msg>test: invalid email server and postgres URLs in myirmaserver configuration<commit_after>//+build !local_tests

package myirmaserver

import (
	"path/filepath"
	"testing"

	"github.com/privacybydesign/irmago/internal/test"
	"github.com/stretchr/testify/assert"
)

func validConfWithEmail(t *testing.T) *Configuration {
	testdataPath := test.FindTestdataFolder(t)
	conf := validConf(t)
	conf.EmailServer = "localhost:1025"
	conf.DefaultLanguage = "en"
	conf.LoginEmailFiles = map[string]string{"en": filepath.Join(testdataPath, "emailtemplate.html")}
	conf.LoginEmailSubject = map[string]string{"en": "testsubject"}
	conf.LoginEmailBaseURL = map[string]string{"en": "localhost:8000/test/"}
	conf.DeleteEmailFiles = map[string]string{"en": filepath.Join(testdataPath, "emailtemplate.html")}
	conf.DeleteEmailSubject = map[string]string{"en": "testsubject"}
	conf.DeleteAccountFiles = map[string]string{"en": filepath.Join(testdataPath, "emailtemplate.html")}
	conf.DeleteAccountSubject = map[string]string{"en": "testsubject"}
	return conf
}

func TestConfEmailValidation(t *testing.T) {
	testdataPath := test.FindTestdataFolder(t)

	conf := validConf(t)
	conf.DBType = DatabaseTypePostgres
	conf.DBConnstring = "postgresql://localhost:5432/test"
	_, err := New(conf)
	assert.NoError(t, err)

	_, err = New(validConfWithEmail(t))
	assert.NoError(t, err)

	conf = validConfWithEmail(t)
	conf.DBType = DatabaseTypePostgres
	conf.DBConnstring = "postgresql://localhost:54321/test"
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.EmailServer = "http://localhost:1025"
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailFiles = map[string]string{"en": filepath.Join(testdataPath, "invalidemailtemplate.html")}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailFiles = map[string]string{}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailSubject = map[string]string{}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.LoginEmailBaseURL = map[string]string{}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.DeleteEmailFiles = map[string]string{"de": filepath.Join(testdataPath, "emailtemplate.html")}
	_, err = New(conf)
	assert.Error(t, err)

	conf = validConfWithEmail(t)
	conf.DeleteEmailSubject = map[string]string{"de": "testsubject"}
	_, err = New(conf)
	assert.Error(t, err)
}
<|endoftext|>


<commit_before>package rabbitsmpp

import (
	"log"
	"sync"

	"github.com/cenk/backoff"
	"github.com/streadway/amqp"
)

var conn struct {
	conn *amqp.Connection
	url  string
	sync.Mutex
}

func initConn(url string) error {
	conn.Lock()
	if conn.conn != nil {
		conn.Unlock()
		return nil
	}
	conn.Unlock()
	return startConn(url)
}

func startConn(url string) error {
	c, err := amqp.Dial(url)
	if err != nil {
		return err
	}
	conn.Lock()
	conn.conn = c
	conn.url = url
	conn.Unlock()
	rebindOnClose()
	return nil
}

func rebindOnClose() {
	conn.Lock()
	errors := make(chan *amqp.Error)
	errChan := conn.conn.NotifyClose(errors)
	conn.Unlock()
	go func() {
		closeNotification := <-errChan
		log.Println("connection was closed:", closeNotification, "rebinding...")

		ticker := backoff.NewTicker(backoff.NewExponentialBackOff())
		var err error
		for _ = range ticker.C {
			err = startConn(conn.url)
			if err != nil {
				continue
			}

			// stop the ticker
			ticker.Stop()
			return
		}
		log.Fatal("Failed to rebind connection:", err)
	}()
}

func closeConn() error {
	conn.Lock()
	defer conn.Unlock()
	if conn.conn == nil {
		return nil
	}
	err := conn.conn.Close()
	conn.conn = nil
	return err
}

func getConn() *amqp.Connection {
	conn.Lock()
	defer conn.Unlock()
	return conn.conn
}

type Config struct {
	URL       string
	QueueName string
}

type Closer interface {
	// Close terminates the connection.
	Close() error
}

type Channel interface {
	ExchangeDeclare(name, kind string, durable, autoDelete, internal, noWait bool, args amqp.Table) error
	QueueBind(name, key, exchange string, noWait bool, args amqp.Table) error
	QueueDeclare(string, bool, bool, bool, bool, amqp.Table) (amqp.Queue, error)
	QueueInspect(string) (amqp.Queue, error)
	Consume(string, string, bool, bool, bool, bool, amqp.Table) (<-chan amqp.Delivery, error)
	Publish(string, string, bool, bool, amqp.Publishing) error
	Qos(int, int, bool) error
	Closer
}

type Client interface {
	Channel() (Channel, error)
	Config() Config
	GetCloseChan() chan *amqp.Error
	Closer
}

type ClientFactory func() (Client, error)

func clientFactory(conf Config) ClientFactory {
	return func() (Client, error) {
		return NewClient(conf)
	}
}

var defaultClientFactory = clientFactory

type client struct {
	config Config
}

func NewClient(conf Config) (Client, error) {
	err := initConn(conf.URL)
	if err != nil {
		return nil, err
	}
	return &client{config: conf}, nil
}

func (c *client) Channel() (Channel, error) {
	conn.Lock()
	defer conn.Unlock()
	return conn.conn.Channel()
}

func (c *client) Config() Config {
	return c.config
}

func (c *client) Close() error {
	return getConn().Close()
}

func (c *client) GetCloseChan() chan *amqp.Error {
	conn.Lock()
	defer conn.Unlock()
	if conn.conn == nil {
		return nil
	}
	ch := make(chan *amqp.Error)
	return conn.conn.NotifyClose(ch)
}
<commit_msg>More logging<commit_after>package rabbitsmpp

import (
	"log"
	"sync"

	"github.com/cenk/backoff"
	"github.com/streadway/amqp"
)

var conn struct {
	conn *amqp.Connection
	url  string
	sync.Mutex
}

func initConn(url string) error {
	conn.Lock()
	if conn.conn != nil {
		conn.Unlock()
		return nil
	}
	conn.Unlock()
	return startConn(url)
}

func startConn(url string) error {
	c, err := amqp.Dial(url)
	if err != nil {
		return err
	}
	conn.Lock()
	conn.conn = c
	conn.url = url
	conn.Unlock()
	rebindOnClose()
	return nil
}

func rebindOnClose() {
	conn.Lock()
	errors := make(chan *amqp.Error)
	errChan := conn.conn.NotifyClose(errors)
	conn.Unlock()
	go func() {
		closeNotification := <-errChan
		log.Println("connection was closed:", closeNotification, "rebinding...")

		ticker := backoff.NewTicker(backoff.NewExponentialBackOff())
		var err error
		for _ = range ticker.C {
			err = startConn(conn.url)
			if err != nil {
				log.Println("failed to restore connection:", err)
				continue
			}
			log.Println("connection was restored")

			// stop the ticker
			ticker.Stop()
			return
		}
		log.Fatal("Failed to rebind connection:", err)
	}()
}

func closeConn() error {
	conn.Lock()
	defer conn.Unlock()
	if conn.conn == nil {
		return nil
	}
	err := conn.conn.Close()
	conn.conn = nil
	return err
}

func getConn() *amqp.Connection {
	conn.Lock()
	defer conn.Unlock()
	return conn.conn
}

type Config struct {
	URL       string
	QueueName string
}

type Closer interface {
	// Close terminates the connection.
	Close() error
}

type Channel interface {
	ExchangeDeclare(name, kind string, durable, autoDelete, internal, noWait bool, args amqp.Table) error
	QueueBind(name, key, exchange string, noWait bool, args amqp.Table) error
	QueueDeclare(string, bool, bool, bool, bool, amqp.Table) (amqp.Queue, error)
	QueueInspect(string) (amqp.Queue, error)
	Consume(string, string, bool, bool, bool, bool, amqp.Table) (<-chan amqp.Delivery, error)
	Publish(string, string, bool, bool, amqp.Publishing) error
	Qos(int, int, bool) error
	Closer
}

type Client interface {
	Channel() (Channel, error)
	Config() Config
	GetCloseChan() chan *amqp.Error
	Closer
}

type ClientFactory func() (Client, error)

func clientFactory(conf Config) ClientFactory {
	return func() (Client, error) {
		return NewClient(conf)
	}
}

var defaultClientFactory = clientFactory

type client struct {
	config Config
}

func NewClient(conf Config) (Client, error) {
	err := initConn(conf.URL)
	if err != nil {
		return nil, err
	}
	return &client{config: conf}, nil
}

func (c *client) Channel() (Channel, error) {
	conn.Lock()
	defer conn.Unlock()
	return conn.conn.Channel()
}

func (c *client) Config() Config {
	return c.config
}

func (c *client) Close() error {
	return getConn().Close()
}

func (c *client) GetCloseChan() chan *amqp.Error {
	conn.Lock()
	defer conn.Unlock()
	if conn.conn == nil {
		return nil
	}
	ch := make(chan *amqp.Error)
	return conn.conn.NotifyClose(ch)
}
<|endoftext|>


<commit_before>package sshprovider

import (
	"context"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/session/sshforward"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// AgentConfig is the config for a single exposed SSH agent
type AgentConfig struct {
	ID    string
	Paths []string
}

// NewSSHAgentProvider creates a session provider that allows access to ssh agent
func NewSSHAgentProvider(confs []AgentConfig) (session.Attachable, error) {
	m := map[string]source{}
	for _, conf := range confs {
		if len(conf.Paths) == 0 || len(conf.Paths) == 1 && conf.Paths[0] == "" {
			conf.Paths = []string{os.Getenv("SSH_AUTH_SOCK")}
		}

		if conf.Paths[0] == "" {
			return nil, errors.Errorf("invalid empty ssh agent socket, make sure SSH_AUTH_SOCK is set")
		}

		src, err := toAgentSource(conf.Paths)
		if err != nil {
			return nil, err
		}
		if conf.ID == "" {
			conf.ID = sshforward.DefaultID
		}
		if _, ok := m[conf.ID]; ok {
			return nil, errors.Errorf("invalid duplicate ID %s", conf.ID)
		}
		m[conf.ID] = src
	}

	return &socketProvider{m: m}, nil
}

type source struct {
	agent  agent.Agent
	socket string
}

type socketProvider struct {
	m map[string]source
}

func (sp *socketProvider) Register(server *grpc.Server) {
	sshforward.RegisterSSHServer(server, sp)
}

func (sp *socketProvider) CheckAgent(ctx context.Context, req *sshforward.CheckAgentRequest) (*sshforward.CheckAgentResponse, error) {
	id := sshforward.DefaultID
	if req.ID != "" {
		id = req.ID
	}
	if _, ok := sp.m[id]; !ok {
		return &sshforward.CheckAgentResponse{}, errors.Errorf("unset ssh forward key %s", id)
	}
	return &sshforward.CheckAgentResponse{}, nil
}

func (sp *socketProvider) ForwardAgent(stream sshforward.SSH_ForwardAgentServer) error {
	id := sshforward.DefaultID

	opts, _ := metadata.FromIncomingContext(stream.Context()) // if no metadata continue with empty object

	if v, ok := opts[sshforward.KeySSHID]; ok && len(v) > 0 && v[0] != "" {
		id = v[0]
	}

	src, ok := sp.m[id]
	if !ok {
		return errors.Errorf("unset ssh forward key %s", id)
	}

	var a agent.Agent

	if src.socket != "" {
		conn, err := net.DialTimeout("unix", src.socket, time.Second)
		if err != nil {
			return errors.Wrapf(err, "failed to connect to %s", src.socket)
		}

		a = &readOnlyAgent{agent.NewClient(conn)}
		defer conn.Close()
	} else {
		a = src.agent
	}

	s1, s2 := sockPair()

	eg, ctx := errgroup.WithContext(context.TODO())

	eg.Go(func() error {
		return agent.ServeAgent(a, s1)
	})

	eg.Go(func() error {
		defer s1.Close()
		return sshforward.Copy(ctx, s2, stream, nil)
	})

	return eg.Wait()
}

func toAgentSource(paths []string) (source, error) {
	var keys bool
	var socket string
	a := agent.NewKeyring()
	for _, p := range paths {
		if socket != "" {
			return source{}, errors.New("only single socket allowed")
		}
		fi, err := os.Stat(p)
		if err != nil {
			return source{}, errors.WithStack(err)
		}
		if fi.Mode()&os.ModeSocket > 0 {
			if keys {
				return source{}, errors.Errorf("invalid combination of keys and sockets")
			}
			socket = p
			continue
		}
		keys = true
		f, err := os.Open(p)
		if err != nil {
			return source{}, errors.Wrapf(err, "failed to open %s", p)
		}
		dt, err := ioutil.ReadAll(&io.LimitedReader{R: f, N: 100 * 1024})
		if err != nil {
			return source{}, errors.Wrapf(err, "failed to read %s", p)
		}

		k, err := ssh.ParseRawPrivateKey(dt)
		if err != nil {
			return source{}, errors.Wrapf(err, "failed to parse %s", p) // TODO: prompt passphrase?
		}
		if err := a.Add(agent.AddedKey{PrivateKey: k}); err != nil {
			return source{}, errors.Wrapf(err, "failed to add %s to agent", p)
		}
	}

	if socket != "" {
		return source{socket: socket}, nil
	}

	return source{agent: a}, nil
}

func sockPair() (io.ReadWriteCloser, io.ReadWriteCloser) {
	pr1, pw1 := io.Pipe()
	pr2, pw2 := io.Pipe()
	return &sock{pr1, pw2, pw1}, &sock{pr2, pw1, pw2}
}

type sock struct {
	io.Reader
	io.Writer
	io.Closer
}

type readOnlyAgent struct {
	agent.ExtendedAgent
}

func (a *readOnlyAgent) Add(_ agent.AddedKey) error {
	return errors.Errorf("adding new keys not allowed by buildkit")
}

func (a *readOnlyAgent) Remove(_ ssh.PublicKey) error {
	return errors.Errorf("removing keys not allowed by buildkit")
}

func (a *readOnlyAgent) RemoveAll() error {
	return errors.Errorf("removing keys not allowed by buildkit")
}

func (a *readOnlyAgent) Lock(_ []byte) error {
	return errors.Errorf("locking agent not allowed by buildkit")
}
<commit_msg>Disallow Extensions for now<commit_after>package sshprovider

import (
	"context"
	"io"
	"io/ioutil"
	"net"
	"os"
	"time"

	"github.com/moby/buildkit/session"
	"github.com/moby/buildkit/session/sshforward"
	"github.com/pkg/errors"
	"golang.org/x/crypto/ssh"
	"golang.org/x/crypto/ssh/agent"
	"golang.org/x/sync/errgroup"
	"google.golang.org/grpc"
	"google.golang.org/grpc/metadata"
)

// AgentConfig is the config for a single exposed SSH agent
type AgentConfig struct {
	ID    string
	Paths []string
}

// NewSSHAgentProvider creates a session provider that allows access to ssh agent
func NewSSHAgentProvider(confs []AgentConfig) (session.Attachable, error) {
	m := map[string]source{}
	for _, conf := range confs {
		if len(conf.Paths) == 0 || len(conf.Paths) == 1 && conf.Paths[0] == "" {
			conf.Paths = []string{os.Getenv("SSH_AUTH_SOCK")}
		}

		if conf.Paths[0] == "" {
			return nil, errors.Errorf("invalid empty ssh agent socket, make sure SSH_AUTH_SOCK is set")
		}

		src, err := toAgentSource(conf.Paths)
		if err != nil {
			return nil, err
		}
		if conf.ID == "" {
			conf.ID = sshforward.DefaultID
		}
		if _, ok := m[conf.ID]; ok {
			return nil, errors.Errorf("invalid duplicate ID %s", conf.ID)
		}
		m[conf.ID] = src
	}

	return &socketProvider{m: m}, nil
}

type source struct {
	agent  agent.Agent
	socket string
}

type socketProvider struct {
	m map[string]source
}

func (sp *socketProvider) Register(server *grpc.Server) {
	sshforward.RegisterSSHServer(server, sp)
}

func (sp *socketProvider) CheckAgent(ctx context.Context, req *sshforward.CheckAgentRequest) (*sshforward.CheckAgentResponse, error) {
	id := sshforward.DefaultID
	if req.ID != "" {
		id = req.ID
	}
	if _, ok := sp.m[id]; !ok {
		return &sshforward.CheckAgentResponse{}, errors.Errorf("unset ssh forward key %s", id)
	}
	return &sshforward.CheckAgentResponse{}, nil
}

func (sp *socketProvider) ForwardAgent(stream sshforward.SSH_ForwardAgentServer) error {
	id := sshforward.DefaultID

	opts, _ := metadata.FromIncomingContext(stream.Context()) // if no metadata continue with empty object

	if v, ok := opts[sshforward.KeySSHID]; ok && len(v) > 0 && v[0] != "" {
		id = v[0]
	}

	src, ok := sp.m[id]
	if !ok {
		return errors.Errorf("unset ssh forward key %s", id)
	}

	var a agent.Agent

	if src.socket != "" {
		conn, err := net.DialTimeout("unix", src.socket, time.Second)
		if err != nil {
			return errors.Wrapf(err, "failed to connect to %s", src.socket)
		}

		a = &readOnlyAgent{agent.NewClient(conn)}
		defer conn.Close()
	} else {
		a = src.agent
	}

	s1, s2 := sockPair()

	eg, ctx := errgroup.WithContext(context.TODO())

	eg.Go(func() error {
		return agent.ServeAgent(a, s1)
	})

	eg.Go(func() error {
		defer s1.Close()
		return sshforward.Copy(ctx, s2, stream, nil)
	})

	return eg.Wait()
}

func toAgentSource(paths []string) (source, error) {
	var keys bool
	var socket string
	a := agent.NewKeyring()
	for _, p := range paths {
		if socket != "" {
			return source{}, errors.New("only single socket allowed")
		}
		fi, err := os.Stat(p)
		if err != nil {
			return source{}, errors.WithStack(err)
		}
		if fi.Mode()&os.ModeSocket > 0 {
			if keys {
				return source{}, errors.Errorf("invalid combination of keys and sockets")
			}
			socket = p
			continue
		}
		keys = true
		f, err := os.Open(p)
		if err != nil {
			return source{}, errors.Wrapf(err, "failed to open %s", p)
		}
		dt, err := ioutil.ReadAll(&io.LimitedReader{R: f, N: 100 * 1024})
		if err != nil {
			return source{}, errors.Wrapf(err, "failed to read %s", p)
		}

		k, err := ssh.ParseRawPrivateKey(dt)
		if err != nil {
			return source{}, errors.Wrapf(err, "failed to parse %s", p) // TODO: prompt passphrase?
		}
		if err := a.Add(agent.AddedKey{PrivateKey: k}); err != nil {
			return source{}, errors.Wrapf(err, "failed to add %s to agent", p)
		}
	}

	if socket != "" {
		return source{socket: socket}, nil
	}

	return source{agent: a}, nil
}

func sockPair() (io.ReadWriteCloser, io.ReadWriteCloser) {
	pr1, pw1 := io.Pipe()
	pr2, pw2 := io.Pipe()
	return &sock{pr1, pw2, pw1}, &sock{pr2, pw1, pw2}
}

type sock struct {
	io.Reader
	io.Writer
	io.Closer
}

type readOnlyAgent struct {
	agent.ExtendedAgent
}

func (a *readOnlyAgent) Add(_ agent.AddedKey) error {
	return errors.Errorf("adding new keys not allowed by buildkit")
}

func (a *readOnlyAgent) Remove(_ ssh.PublicKey) error {
	return errors.Errorf("removing keys not allowed by buildkit")
}

func (a *readOnlyAgent) RemoveAll() error {
	return errors.Errorf("removing keys not allowed by buildkit")
}

func (a *readOnlyAgent) Lock(_ []byte) error {
	return errors.Errorf("locking agent not allowed by buildkit")
}

func (a *readOnlyAgent) Extension(_ string, _ []byte) ([]byte, error) {
	return nil, errors.Errorf("extensions not allowed by buildkit")
}
<|endoftext|>


<commit_before>package daemon

import (
	// Importing packages here only to make sure their init gets called and
	// therefore they register themselves to the logdriver factory.
	_ "github.com/docker/docker/daemon/logger/jsonfilelog"
	// _ "github.com/docker/docker/daemon/logger/kafkalog"
	_ "github.com/docker/docker/daemon/logger/loghub"
	_ "github.com/docker/docker/daemon/logger/fifo"
	// _ "github.com/docker/docker/daemon/logger/redislog"
)
<commit_msg>Enable all drivers<commit_after>package daemon

import (
	// Importing packages here only to make sure their init gets called and
	// therefore they register themselves to the logdriver factory.
	_ "github.com/docker/docker/daemon/logger/fifo"
	_ "github.com/docker/docker/daemon/logger/jsonfilelog"
	_ "github.com/docker/docker/daemon/logger/kafkalog"
	_ "github.com/docker/docker/daemon/logger/loghub"
	_ "github.com/docker/docker/daemon/logger/redislog"
)
<|endoftext|>