<commit_before>package transforms

import (
	. "connectordb/streamdb/datastream"
	"errors"
	"testing"

	"github.com/connectordb/duck"
	"github.com/stretchr/testify/require"
)

func TestPipelineGenerator(t *testing.T) {
	testcases := []struct {
		Pipeline       string
		HasSyntaxError bool
		Haserror2      bool
		Input          *Datapoint
		Output         *Datapoint
	}{
		// Identity functions
		{"true", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"45.555", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 45.555}},
		{"\"string\"", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string"}},
		{"\"❤ ☀ ☆ ☂ ☻ ♞ ☯ ☭ ☢ €\"", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "❤ ☀ ☆ ☂ ☻ ♞ ☯ ☭ ☢ €"}},

		// Literal identity
		{"$", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},

		// Basic Testing
		{"4 < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},

		// Logical tests
		{"true or false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"false or false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"true and false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"true and (false or true)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"true and true", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"true and not false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},

		// Logical filter tests
		{"if true", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"if true | 42", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 42}},
		{"if false", false, false, &Datapoint{Data: 4}, nil},
		{"if $ < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},

		// Comparison
		{"$ > 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"$ > 3", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ >= 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ < 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"$ < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ <= 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ != 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"$ != 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ == 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ == 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},

		// Logical pipelines
		{"if $ < 5 and $ > 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"if $ < 5 | if $ > 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"if $ < 5 | if $ > 33", false, false, &Datapoint{Data: 4}, nil},
		{"if $ < 5 | $ > 33", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"has(\"test\") | $ < 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"if has(\"test\")| $ < 1", false, false, &Datapoint{Data: 4}, nil},
		{"if has(\"test\")| $[\"test\"] < 1", false, false, &Datapoint{Data: map[string]interface{}{"test": 25}}, &Datapoint{Data: false}},
		{"if has(\"tst\")| $[\"test\"] < 1", false, false, &Datapoint{Data: map[string]interface{}{"test": 25}}, nil},
		{"if has(\"test\")| $[\"test\"] > 1", false, false, &Datapoint{Data: map[string]interface{}{"test": 25}}, &Datapoint{Data: true}},

		// Invalid
		{"if has(\"test\"", true, false, nil, nil},
		{"$[\"test\"]", false, true, &Datapoint{Data: 4}, nil},

		// Multiple stage pipeline
		{"$ | false | 42", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 42}},

		// implicit logicals
		{"gt(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"gt(3)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"gte(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"lt(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"lt(5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"lte(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"ne(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"ne(5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"eq(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"eq(5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},

		// Test custom functions
		{"identity()", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"passthrough($ > 5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"passthrough($ > 5 | eq(false))", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"fortyTwo()", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 42}},
		{"doesnotexist()", true, false, &Datapoint{Data: 4}, nil},

		// wrong number of args on generation
		{"passthrough($ > 5 | eq(false), $)", true, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},

		// setting values
		{"set($, 4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"set($, true)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"set($, \"foo\")", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "foo"}},
		{"set($[\"bar\"], \"foo\")", false, true, &Datapoint{Data: 4}, &Datapoint{Data: "foo"}},

		// maths
		{"1 + 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 2}},
		{"$ + 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 5}},
		{"$ + \"4\"", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 8}},
		{"$ * 2", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 8}},
		{"$ / 2", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 2}},
		{"1 + 2 * 3 + 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 11}},
		{"1 + 2 * (3 + 4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 15}},
		{"-1 + 2", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 1}},
		{"-(1 + 2)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: -3}},
	}

	// function that should nilt out
	identityFunc := func(name string, children ...TransformFunc) (TransformFunc, error) {
		return func(dp *Datapoint) (tdp *Datapoint, err error) {
			return dp, nil
		}, nil
	}
	RegisterCustomFunction("identity", identityFunc)

	// passthrough
	passthroughFunc := func(name string, children ...TransformFunc) (TransformFunc, error) {
		if len(children) != 1 {
			return pipelineGeneratorIdentity(), errors.New("passthrough error")
		}
		return func(dp *Datapoint) (tdp *Datapoint, err error) {
			return children[0](dp)
		}, nil
	}
	RegisterCustomFunction("passthrough", passthroughFunc)

	fortyTwo := func(name string, children ...TransformFunc) (TransformFunc, error) {
		return func(dp *Datapoint) (tdp *Datapoint, err error) {
			dp.Data = 42
			return dp, nil
		}, nil
	}
	RegisterCustomFunction("fortyTwo", fortyTwo)

	for _, c := range testcases {

		result, err := ParseTransform(c.Pipeline)

		if c.HasSyntaxError {
			require.Error(t, err)
			continue
		}

		require.NoError(t, err, duck.JSONString(c))

		dp, err := result(c.Input)
		if c.Haserror2 {
			require.Error(t, err, duck.JSONString(c))
		} else {
			require.NoError(t, err, duck.JSONString(c))
			if c.Output != nil {
				require.Equal(t, c.Output.String(), dp.String(), duck.JSONString(c))
			} else {
				require.Nil(t, dp, duck.JSONString(c))
			}
		}
	}
}

func TestParseTransform(t *testing.T) {
	// Valid pipeline
	{
		transform, err := ParseTransform("42")
		require.Nil(t, err)
		require.NotNil(t, transform)
	}

	// invalid pipeline
	{
		transform, err := ParseTransform("(")
		require.NotNil(t, err)
		require.Nil(t, transform)
	}
}
<commit_msg>added tests for escape sequences<commit_after>package transforms

import (
	. "connectordb/streamdb/datastream"
	"errors"
	"testing"

	"github.com/connectordb/duck"
	"github.com/stretchr/testify/require"
)

func TestPipelineGenerator(t *testing.T) {
	testcases := []struct {
		Pipeline       string
		HasSyntaxError bool
		Haserror2      bool
		Input          *Datapoint
		Output         *Datapoint
	}{
		// Identity functions
		{"true", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"45.555", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 45.555}},

		// String testing -- escaping, unicode and pipes
		{"\"string\"", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string"}},
		{"'string'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string"}},
		{"'string\\n'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string\n"}},
		{"'string\\t'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string\t"}},
		{"'string\\\\'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string\\"}},
		{"'string\\r'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string\r"}},
		{"'string\"'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "string\""}},
		{"'|'", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "|"}},
		{"\"❤ ☀ ☆ ☂ ☻ ♞ ☯ ☭ ☢ €\"", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "❤ ☀ ☆ ☂ ☻ ♞ ☯ ☭ ☢ €"}},

		// Literal identity
		{"$", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},

		// Basic Testing
		{"4 < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},

		// Logical tests
		{"true or false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"false or false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"true and false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"true and (false or true)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"true and true", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"true and not false", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},

		// Logical filter tests
		{"if true", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"if true | 42", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 42}},
		{"if false", false, false, &Datapoint{Data: 4}, nil},
		{"if $ < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},

		// Comparison
		{"$ > 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"$ > 3", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ >= 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ < 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"$ < 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ <= 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ != 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"$ != 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ == 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"$ == 5", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},

		// Logical pipelines
		{"if $ < 5 and $ > 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"if $ < 5 | if $ > 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"if $ < 5 | if $ > 33", false, false, &Datapoint{Data: 4}, nil},
		{"if $ < 5 | $ > 33", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"has(\"test\") | $ < 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"if has(\"test\")| $ < 1", false, false, &Datapoint{Data: 4}, nil},
		{"if has(\"test\")| $[\"test\"] < 1", false, false, &Datapoint{Data: map[string]interface{}{"test": 25}}, &Datapoint{Data: false}},
		{"if has(\"tst\")| $[\"test\"] < 1", false, false, &Datapoint{Data: map[string]interface{}{"test": 25}}, nil},
		{"if has(\"test\")| $[\"test\"] > 1", false, false, &Datapoint{Data: map[string]interface{}{"test": 25}}, &Datapoint{Data: true}},

		// Invalid
		{"if has(\"test\"", true, false, nil, nil},
		{"$[\"test\"]", false, true, &Datapoint{Data: 4}, nil},

		// Multiple stage pipeline
		{"$ | false | 42", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 42}},

		// implicit logicals
		{"gt(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"gt(3)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"gte(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"lt(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"lt(5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"lte(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"ne(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"ne(5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"eq(4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"eq(5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},

		// Test custom functions
		{"identity()", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"passthrough($ > 5)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: false}},
		{"passthrough($ > 5 | eq(false))", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"fortyTwo()", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 42}},
		{"doesnotexist()", true, false, &Datapoint{Data: 4}, nil},

		// wrong number of args on generation
		{"passthrough($ > 5 | eq(false), $)", true, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},

		// setting values
		{"set($, 4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 4}},
		{"set($, true)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: true}},
		{"set($, \"foo\")", false, false, &Datapoint{Data: 4}, &Datapoint{Data: "foo"}},
		{"set($[\"bar\"], \"foo\")", false, true, &Datapoint{Data: 4}, &Datapoint{Data: "foo"}},

		// maths
		{"1 + 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 2}},
		{"$ + 1", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 5}},
		{"$ + \"4\"", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 8}},
		{"$ * 2", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 8}},
		{"$ / 2", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 2}},
		{"1 + 2 * 3 + 4", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 11}},
		{"1 + 2 * (3 + 4)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 15}},
		{"-1 + 2", false, false, &Datapoint{Data: 4}, &Datapoint{Data: 1}},
		{"-(1 + 2)", false, false, &Datapoint{Data: 4}, &Datapoint{Data: -3}},
	}

	// function that should nilt out
	identityFunc := func(name string, children ...TransformFunc) (TransformFunc, error) {
		return func(dp *Datapoint) (tdp *Datapoint, err error) {
			return dp, nil
		}, nil
	}
	RegisterCustomFunction("identity", identityFunc)

	// passthrough
	passthroughFunc := func(name string, children ...TransformFunc) (TransformFunc, error) {
		if len(children) != 1 {
			return pipelineGeneratorIdentity(), errors.New("passthrough error")
		}
		return func(dp *Datapoint) (tdp *Datapoint, err error) {
			return children[0](dp)
		}, nil
	}
	RegisterCustomFunction("passthrough", passthroughFunc)

	fortyTwo := func(name string, children ...TransformFunc) (TransformFunc, error) {
		return func(dp *Datapoint) (tdp *Datapoint, err error) {
			dp.Data = 42
			return dp, nil
		}, nil
	}
	RegisterCustomFunction("fortyTwo", fortyTwo)

	for _, c := range testcases {

		result, err := ParseTransform(c.Pipeline)

		if c.HasSyntaxError {
			require.Error(t, err)
			continue
		}

		require.NoError(t, err, duck.JSONString(c))

		dp, err := result(c.Input)
		if c.Haserror2 {
			require.Error(t, err, duck.JSONString(c))
		} else {
			require.NoError(t, err, duck.JSONString(c))
			if c.Output != nil {
				require.Equal(t, c.Output.String(), dp.String(), duck.JSONString(c))
			} else {
				require.Nil(t, dp, duck.JSONString(c))
			}
		}
	}
}

func TestParseTransform(t *testing.T) {
	// Valid pipeline
	{
		transform, err := ParseTransform("42")
		require.Nil(t, err)
		require.NotNil(t, transform)
	}

	// invalid pipeline
	{
		transform, err := ParseTransform("(")
		require.NotNil(t, err)
		require.Nil(t, transform)
	}
}
<|endoftext|>


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

package main

import (
	"flag"
	"net/http"
	"os"

	"github.com/gorilla/mux"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/accounts"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/devices"
	"github.com/matrix-org/dendrite/clientapi/producers"
	"github.com/matrix-org/dendrite/common"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/common/keydb"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/naffka"

	mediaapi_routing "github.com/matrix-org/dendrite/mediaapi/routing"
	mediaapi_storage "github.com/matrix-org/dendrite/mediaapi/storage"

	roomserver_alias "github.com/matrix-org/dendrite/roomserver/alias"
	roomserver_input "github.com/matrix-org/dendrite/roomserver/input"
	roomserver_query "github.com/matrix-org/dendrite/roomserver/query"
	roomserver_storage "github.com/matrix-org/dendrite/roomserver/storage"

	clientapi_consumers "github.com/matrix-org/dendrite/clientapi/consumers"
	clientapi_routing "github.com/matrix-org/dendrite/clientapi/routing"

	syncapi_consumers "github.com/matrix-org/dendrite/syncapi/consumers"
	syncapi_routing "github.com/matrix-org/dendrite/syncapi/routing"
	syncapi_storage "github.com/matrix-org/dendrite/syncapi/storage"
	syncapi_sync "github.com/matrix-org/dendrite/syncapi/sync"
	syncapi_types "github.com/matrix-org/dendrite/syncapi/types"

	federationapi_routing "github.com/matrix-org/dendrite/federationapi/routing"

	federationsender_consumers "github.com/matrix-org/dendrite/federationsender/consumers"
	"github.com/matrix-org/dendrite/federationsender/queue"
	federationsender_storage "github.com/matrix-org/dendrite/federationsender/storage"

	log "github.com/Sirupsen/logrus"
	sarama "gopkg.in/Shopify/sarama.v1"
)

var (
	logDir        = os.Getenv("LOG_DIR")
	configPath    = flag.String("config", "dendrite.yaml", "The path to the config file. For more information, see the config file in this repository.")
	httpBindAddr  = flag.String("http-bind-address", ":8008", "The HTTP listening port for the server")
	httpsBindAddr = flag.String("https-bind-address", ":8448", "The HTTPS listening port for the server")
	certFile      = flag.String("tls-cert", "", "The PEM formatted X509 certificate to use for TLS")
	keyFile       = flag.String("tls-key", "", "The PEM private key to use for TLS")
)

func main() {
	common.SetupLogging(logDir)

	flag.Parse()

	if *configPath == "" {
		log.Fatal("--config must be supplied")
	}
	cfg, err := config.LoadMonolithic(*configPath)
	if err != nil {
		log.Fatalf("Invalid config file: %s", err)
	}

	m := newMonolith(cfg)
	m.setupDatabases()
	m.setupFederation()
	m.setupKafka()
	m.setupRoomServer()
	m.setupProducers()
	m.setupNotifiers()
	m.setupConsumers()
	m.setupAPIs()

	// Expose the matrix APIs directly rather than putting them under a /api path.
	go func() {
		log.Info("Listening on ", *httpBindAddr)
		log.Fatal(http.ListenAndServe(*httpBindAddr, m.api))
	}()
	// Handle HTTPS if certificate and key are provided
	go func() {
		if *certFile != "" && *keyFile != "" {
			log.Info("Listening on ", *httpsBindAddr)
			log.Fatal(http.ListenAndServeTLS(*httpsBindAddr, *certFile, *keyFile, m.api))
		}
	}()

	// We want to block forever to let the HTTP and HTTPS handler serve the APIs
	select {}
}

// A monolith contains all the dendrite components.
// Some of the setup functions depend on previous setup functions, so they must
// be called in the same order as they are defined in the file.
type monolith struct {
	cfg *config.Dendrite
	api *mux.Router

	roomServerDB       *roomserver_storage.Database
	accountDB          *accounts.Database
	deviceDB           *devices.Database
	keyDB              *keydb.Database
	mediaAPIDB         *mediaapi_storage.Database
	syncAPIDB          *syncapi_storage.SyncServerDatabase
	federationSenderDB *federationsender_storage.Database

	federation *gomatrixserverlib.FederationClient
	keyRing    gomatrixserverlib.KeyRing

	inputAPI *roomserver_input.RoomserverInputAPI
	queryAPI *roomserver_query.RoomserverQueryAPI
	aliasAPI *roomserver_alias.RoomserverAliasAPI

	kafkaConsumer sarama.Consumer
	kafkaProducer sarama.SyncProducer

	roomServerProducer *producers.RoomserverProducer
	userUpdateProducer *producers.UserUpdateProducer
	syncProducer       *producers.SyncAPIProducer

	syncAPINotifier *syncapi_sync.Notifier
}

func newMonolith(cfg *config.Dendrite) *monolith {
	return &monolith{cfg: cfg, api: mux.NewRouter()}
}

func (m *monolith) setupDatabases() {
	var err error
	m.roomServerDB, err = roomserver_storage.Open(string(m.cfg.Database.RoomServer))
	if err != nil {
		panic(err)
	}
	m.accountDB, err = accounts.NewDatabase(string(m.cfg.Database.Account), m.cfg.Matrix.ServerName)
	if err != nil {
		log.Panicf("Failed to setup account database(%q): %s", m.cfg.Database.Account, err.Error())
	}
	m.deviceDB, err = devices.NewDatabase(string(m.cfg.Database.Device), m.cfg.Matrix.ServerName)
	if err != nil {
		log.Panicf("Failed to setup device database(%q): %s", m.cfg.Database.Device, err.Error())
	}
	m.keyDB, err = keydb.NewDatabase(string(m.cfg.Database.ServerKey))
	if err != nil {
		log.Panicf("Failed to setup key database(%q): %s", m.cfg.Database.ServerKey, err.Error())
	}
	m.mediaAPIDB, err = mediaapi_storage.Open(string(m.cfg.Database.MediaAPI))
	if err != nil {
		log.Panicf("Failed to setup sync api database(%q): %s", m.cfg.Database.MediaAPI, err.Error())
	}
	m.syncAPIDB, err = syncapi_storage.NewSyncServerDatabase(string(m.cfg.Database.SyncAPI))
	if err != nil {
		log.Panicf("Failed to setup sync api database(%q): %s", m.cfg.Database.SyncAPI, err.Error())
	}
	m.federationSenderDB, err = federationsender_storage.NewDatabase(string(m.cfg.Database.FederationSender))
	if err != nil {
		log.Panicf("startup: failed to create federation sender database with data source %s : %s", m.cfg.Database.FederationSender, err)
	}
}

func (m *monolith) setupFederation() {
	m.federation = gomatrixserverlib.NewFederationClient(
		m.cfg.Matrix.ServerName, m.cfg.Matrix.KeyID, m.cfg.Matrix.PrivateKey,
	)

	m.keyRing = gomatrixserverlib.KeyRing{
		KeyFetchers: []gomatrixserverlib.KeyFetcher{
			// TODO: Use perspective key fetchers for production.
			&gomatrixserverlib.DirectKeyFetcher{Client: m.federation.Client},
		},
		KeyDatabase: m.keyDB,
	}
}

func (m *monolith) setupKafka() {
	var err error
	if m.cfg.Kafka.UseNaffka {
		naff, err := naffka.New(&naffka.MemoryDatabase{})
		if err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
			}).Panic("Failed to setup naffka")
		}
		m.kafkaConsumer = naff
		m.kafkaProducer = naff
	} else {
		m.kafkaConsumer, err = sarama.NewConsumer(m.cfg.Kafka.Addresses, nil)
		if err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
				"addresses":  m.cfg.Kafka.Addresses,
			}).Panic("Failed to setup kafka consumers")
		}
		m.kafkaProducer, err = sarama.NewSyncProducer(m.cfg.Kafka.Addresses, nil)
		if err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
				"addresses":  m.cfg.Kafka.Addresses,
			}).Panic("Failed to setup kafka producers")
		}
	}
}

func (m *monolith) setupRoomServer() {
	m.inputAPI = &roomserver_input.RoomserverInputAPI{
		DB:                   m.roomServerDB,
		Producer:             m.kafkaProducer,
		OutputRoomEventTopic: string(m.cfg.Kafka.Topics.OutputRoomEvent),
	}

	m.queryAPI = &roomserver_query.RoomserverQueryAPI{
		DB: m.roomServerDB,
	}

	m.aliasAPI = &roomserver_alias.RoomserverAliasAPI{
		DB:       m.roomServerDB,
		Cfg:      m.cfg,
		InputAPI: m.inputAPI,
		QueryAPI: m.queryAPI,
	}
}

func (m *monolith) setupProducers() {
	m.roomServerProducer = producers.NewRoomserverProducer(m.inputAPI)
	m.userUpdateProducer = &producers.UserUpdateProducer{
		Producer: m.kafkaProducer,
		Topic:    string(m.cfg.Kafka.Topics.UserUpdates),
	}
	m.syncProducer = &producers.SyncAPIProducer{
		Producer: m.kafkaProducer,
		Topic:    string(m.cfg.Kafka.Topics.OutputClientData),
	}
}

func (m *monolith) setupNotifiers() {
	pos, err := m.syncAPIDB.SyncStreamPosition()
	if err != nil {
		log.Panicf("startup: failed to get latest sync stream position : %s", err)
	}

	m.syncAPINotifier = syncapi_sync.NewNotifier(syncapi_types.StreamPosition(pos))
	if err = m.syncAPINotifier.Load(m.syncAPIDB); err != nil {
		log.Panicf("startup: failed to set up notifier: %s", err)
	}
}

func (m *monolith) setupConsumers() {
	var err error

	clientAPIConsumer := clientapi_consumers.NewOutputRoomEvent(
		m.cfg, m.kafkaConsumer, m.accountDB, m.queryAPI,
	)
	if err = clientAPIConsumer.Start(); err != nil {
		log.Panicf("startup: failed to start room server consumer")
	}

	syncAPIRoomConsumer := syncapi_consumers.NewOutputRoomEvent(
		m.cfg, m.kafkaConsumer, m.syncAPINotifier, m.syncAPIDB, m.queryAPI,
	)
	if err = syncAPIRoomConsumer.Start(); err != nil {
		log.Panicf("startup: failed to start room server consumer: %s", err)
	}

	syncAPIClientConsumer := syncapi_consumers.NewOutputClientData(
		m.cfg, m.kafkaConsumer, m.syncAPINotifier, m.syncAPIDB,
	)
	if err = syncAPIClientConsumer.Start(); err != nil {
		log.Panicf("startup: failed to start client API server consumer: %s", err)
	}

	federationSenderQueues := queue.NewOutgoingQueues(m.cfg.Matrix.ServerName, m.federation)

	federationSenderRoomConsumer := federationsender_consumers.NewOutputRoomEvent(
		m.cfg, m.kafkaConsumer, federationSenderQueues, m.federationSenderDB, m.queryAPI,
	)
	if err = federationSenderRoomConsumer.Start(); err != nil {
		log.WithError(err).Panicf("startup: failed to start room server consumer")
	}
}

func (m *monolith) setupAPIs() {
	clientapi_routing.Setup(
		m.api, http.DefaultClient, *m.cfg, m.roomServerProducer,
		m.queryAPI, m.aliasAPI, m.accountDB, m.deviceDB, m.federation, m.keyRing,
		m.userUpdateProducer, m.syncProducer,
	)

	mediaapi_routing.Setup(
		m.api, http.DefaultClient, m.cfg, m.mediaAPIDB,
	)

	syncapi_routing.Setup(m.api, syncapi_sync.NewRequestPool(
		m.syncAPIDB, m.syncAPINotifier, m.accountDB,
	), m.deviceDB)

	federationapi_routing.Setup(
		m.api, *m.cfg, m.queryAPI, m.roomServerProducer, m.keyRing, m.federation,
	)
}
<commit_msg>Fix kafka consumer setup in monolith. (#184)<commit_after>// Copyright 2017 Vector Creations Ltd
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package main

import (
	"flag"
	"net/http"
	"os"

	"github.com/gorilla/mux"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/accounts"
	"github.com/matrix-org/dendrite/clientapi/auth/storage/devices"
	"github.com/matrix-org/dendrite/clientapi/producers"
	"github.com/matrix-org/dendrite/common"
	"github.com/matrix-org/dendrite/common/config"
	"github.com/matrix-org/dendrite/common/keydb"
	"github.com/matrix-org/gomatrixserverlib"
	"github.com/matrix-org/naffka"

	mediaapi_routing "github.com/matrix-org/dendrite/mediaapi/routing"
	mediaapi_storage "github.com/matrix-org/dendrite/mediaapi/storage"

	roomserver_alias "github.com/matrix-org/dendrite/roomserver/alias"
	roomserver_input "github.com/matrix-org/dendrite/roomserver/input"
	roomserver_query "github.com/matrix-org/dendrite/roomserver/query"
	roomserver_storage "github.com/matrix-org/dendrite/roomserver/storage"

	clientapi_consumers "github.com/matrix-org/dendrite/clientapi/consumers"
	clientapi_routing "github.com/matrix-org/dendrite/clientapi/routing"

	syncapi_consumers "github.com/matrix-org/dendrite/syncapi/consumers"
	syncapi_routing "github.com/matrix-org/dendrite/syncapi/routing"
	syncapi_storage "github.com/matrix-org/dendrite/syncapi/storage"
	syncapi_sync "github.com/matrix-org/dendrite/syncapi/sync"
	syncapi_types "github.com/matrix-org/dendrite/syncapi/types"

	federationapi_routing "github.com/matrix-org/dendrite/federationapi/routing"

	federationsender_consumers "github.com/matrix-org/dendrite/federationsender/consumers"
	"github.com/matrix-org/dendrite/federationsender/queue"
	federationsender_storage "github.com/matrix-org/dendrite/federationsender/storage"

	log "github.com/Sirupsen/logrus"
	sarama "gopkg.in/Shopify/sarama.v1"
)

var (
	logDir        = os.Getenv("LOG_DIR")
	configPath    = flag.String("config", "dendrite.yaml", "The path to the config file. For more information, see the config file in this repository.")
	httpBindAddr  = flag.String("http-bind-address", ":8008", "The HTTP listening port for the server")
	httpsBindAddr = flag.String("https-bind-address", ":8448", "The HTTPS listening port for the server")
	certFile      = flag.String("tls-cert", "", "The PEM formatted X509 certificate to use for TLS")
	keyFile       = flag.String("tls-key", "", "The PEM private key to use for TLS")
)

func main() {
	common.SetupLogging(logDir)

	flag.Parse()

	if *configPath == "" {
		log.Fatal("--config must be supplied")
	}
	cfg, err := config.LoadMonolithic(*configPath)
	if err != nil {
		log.Fatalf("Invalid config file: %s", err)
	}

	m := newMonolith(cfg)
	m.setupDatabases()
	m.setupFederation()
	m.setupKafka()
	m.setupRoomServer()
	m.setupProducers()
	m.setupNotifiers()
	m.setupConsumers()
	m.setupAPIs()

	// Expose the matrix APIs directly rather than putting them under a /api path.
	go func() {
		log.Info("Listening on ", *httpBindAddr)
		log.Fatal(http.ListenAndServe(*httpBindAddr, m.api))
	}()
	// Handle HTTPS if certificate and key are provided
	go func() {
		if *certFile != "" && *keyFile != "" {
			log.Info("Listening on ", *httpsBindAddr)
			log.Fatal(http.ListenAndServeTLS(*httpsBindAddr, *certFile, *keyFile, m.api))
		}
	}()

	// We want to block forever to let the HTTP and HTTPS handler serve the APIs
	select {}
}

// A monolith contains all the dendrite components.
// Some of the setup functions depend on previous setup functions, so they must
// be called in the same order as they are defined in the file.
type monolith struct {
	cfg *config.Dendrite
	api *mux.Router

	roomServerDB       *roomserver_storage.Database
	accountDB          *accounts.Database
	deviceDB           *devices.Database
	keyDB              *keydb.Database
	mediaAPIDB         *mediaapi_storage.Database
	syncAPIDB          *syncapi_storage.SyncServerDatabase
	federationSenderDB *federationsender_storage.Database

	federation *gomatrixserverlib.FederationClient
	keyRing    gomatrixserverlib.KeyRing

	inputAPI *roomserver_input.RoomserverInputAPI
	queryAPI *roomserver_query.RoomserverQueryAPI
	aliasAPI *roomserver_alias.RoomserverAliasAPI

	naffka        *naffka.Naffka
	kafkaProducer sarama.SyncProducer

	roomServerProducer *producers.RoomserverProducer
	userUpdateProducer *producers.UserUpdateProducer
	syncProducer       *producers.SyncAPIProducer

	syncAPINotifier *syncapi_sync.Notifier
}

func newMonolith(cfg *config.Dendrite) *monolith {
	return &monolith{cfg: cfg, api: mux.NewRouter()}
}

func (m *monolith) setupDatabases() {
	var err error
	m.roomServerDB, err = roomserver_storage.Open(string(m.cfg.Database.RoomServer))
	if err != nil {
		panic(err)
	}
	m.accountDB, err = accounts.NewDatabase(string(m.cfg.Database.Account), m.cfg.Matrix.ServerName)
	if err != nil {
		log.Panicf("Failed to setup account database(%q): %s", m.cfg.Database.Account, err.Error())
	}
	m.deviceDB, err = devices.NewDatabase(string(m.cfg.Database.Device), m.cfg.Matrix.ServerName)
	if err != nil {
		log.Panicf("Failed to setup device database(%q): %s", m.cfg.Database.Device, err.Error())
	}
	m.keyDB, err = keydb.NewDatabase(string(m.cfg.Database.ServerKey))
	if err != nil {
		log.Panicf("Failed to setup key database(%q): %s", m.cfg.Database.ServerKey, err.Error())
	}
	m.mediaAPIDB, err = mediaapi_storage.Open(string(m.cfg.Database.MediaAPI))
	if err != nil {
		log.Panicf("Failed to setup sync api database(%q): %s", m.cfg.Database.MediaAPI, err.Error())
	}
	m.syncAPIDB, err = syncapi_storage.NewSyncServerDatabase(string(m.cfg.Database.SyncAPI))
	if err != nil {
		log.Panicf("Failed to setup sync api database(%q): %s", m.cfg.Database.SyncAPI, err.Error())
	}
	m.federationSenderDB, err = federationsender_storage.NewDatabase(string(m.cfg.Database.FederationSender))
	if err != nil {
		log.Panicf("startup: failed to create federation sender database with data source %s : %s", m.cfg.Database.FederationSender, err)
	}
}

func (m *monolith) setupFederation() {
	m.federation = gomatrixserverlib.NewFederationClient(
		m.cfg.Matrix.ServerName, m.cfg.Matrix.KeyID, m.cfg.Matrix.PrivateKey,
	)

	m.keyRing = gomatrixserverlib.KeyRing{
		KeyFetchers: []gomatrixserverlib.KeyFetcher{
			// TODO: Use perspective key fetchers for production.
			&gomatrixserverlib.DirectKeyFetcher{Client: m.federation.Client},
		},
		KeyDatabase: m.keyDB,
	}
}

func (m *monolith) setupKafka() {
	var err error
	if m.cfg.Kafka.UseNaffka {
		naff, err := naffka.New(&naffka.MemoryDatabase{})
		if err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
			}).Panic("Failed to setup naffka")
		}
		m.naffka = naff
		m.kafkaProducer = naff
	} else {
		m.kafkaProducer, err = sarama.NewSyncProducer(m.cfg.Kafka.Addresses, nil)
		if err != nil {
			log.WithFields(log.Fields{
				log.ErrorKey: err,
				"addresses":  m.cfg.Kafka.Addresses,
			}).Panic("Failed to setup kafka producers")
		}
	}
}

func (m *monolith) kafkaConsumer() sarama.Consumer {
	if m.cfg.Kafka.UseNaffka {
		return m.naffka
	}
	consumer, err := sarama.NewConsumer(m.cfg.Kafka.Addresses, nil)
	if err != nil {
		log.WithFields(log.Fields{
			log.ErrorKey: err,
			"addresses":  m.cfg.Kafka.Addresses,
		}).Panic("Failed to setup kafka consumers")
	}
	return consumer
}

func (m *monolith) setupRoomServer() {
	m.inputAPI = &roomserver_input.RoomserverInputAPI{
		DB:                   m.roomServerDB,
		Producer:             m.kafkaProducer,
		OutputRoomEventTopic: string(m.cfg.Kafka.Topics.OutputRoomEvent),
	}

	m.queryAPI = &roomserver_query.RoomserverQueryAPI{
		DB: m.roomServerDB,
	}

	m.aliasAPI = &roomserver_alias.RoomserverAliasAPI{
		DB:       m.roomServerDB,
		Cfg:      m.cfg,
		InputAPI: m.inputAPI,
		QueryAPI: m.queryAPI,
	}
}

func (m *monolith) setupProducers() {
	m.roomServerProducer = producers.NewRoomserverProducer(m.inputAPI)
	m.userUpdateProducer = &producers.UserUpdateProducer{
		Producer: m.kafkaProducer,
		Topic:    string(m.cfg.Kafka.Topics.UserUpdates),
	}
	m.syncProducer = &producers.SyncAPIProducer{
		Producer: m.kafkaProducer,
		Topic:    string(m.cfg.Kafka.Topics.OutputClientData),
	}
}

func (m *monolith) setupNotifiers() {
	pos, err := m.syncAPIDB.SyncStreamPosition()
	if err != nil {
		log.Panicf("startup: failed to get latest sync stream position : %s", err)
	}

	m.syncAPINotifier = syncapi_sync.NewNotifier(syncapi_types.StreamPosition(pos))
	if err = m.syncAPINotifier.Load(m.syncAPIDB); err != nil {
		log.Panicf("startup: failed to set up notifier: %s", err)
	}
}

func (m *monolith) setupConsumers() {
	var err error

	clientAPIConsumer := clientapi_consumers.NewOutputRoomEvent(
		m.cfg, m.kafkaConsumer(), m.accountDB, m.queryAPI,
	)
	if err = clientAPIConsumer.Start(); err != nil {
		log.Panicf("startup: failed to start room server consumer")
	}

	syncAPIRoomConsumer := syncapi_consumers.NewOutputRoomEvent(
		m.cfg, m.kafkaConsumer(), m.syncAPINotifier, m.syncAPIDB, m.queryAPI,
	)
	if err = syncAPIRoomConsumer.Start(); err != nil {
		log.Panicf("startup: failed to start room server consumer: %s", err)
	}

	syncAPIClientConsumer := syncapi_consumers.NewOutputClientData(
		m.cfg, m.kafkaConsumer(), m.syncAPINotifier, m.syncAPIDB,
	)
	if err = syncAPIClientConsumer.Start(); err != nil {
		log.Panicf("startup: failed to start client API server consumer: %s", err)
	}

	federationSenderQueues := queue.NewOutgoingQueues(m.cfg.Matrix.ServerName, m.federation)

	federationSenderRoomConsumer := federationsender_consumers.NewOutputRoomEvent(
		m.cfg, m.kafkaConsumer(), federationSenderQueues, m.federationSenderDB, m.queryAPI,
	)
	if err = federationSenderRoomConsumer.Start(); err != nil {
		log.WithError(err).Panicf("startup: failed to start room server consumer")
	}
}

func (m *monolith) setupAPIs() {
	clientapi_routing.Setup(
		m.api, http.DefaultClient, *m.cfg, m.roomServerProducer,
		m.queryAPI, m.aliasAPI, m.accountDB, m.deviceDB, m.federation, m.keyRing,
		m.userUpdateProducer, m.syncProducer,
	)

	mediaapi_routing.Setup(
		m.api, http.DefaultClient, m.cfg, m.mediaAPIDB,
	)

	syncapi_routing.Setup(m.api, syncapi_sync.NewRequestPool(
		m.syncAPIDB, m.syncAPINotifier, m.accountDB,
	), m.deviceDB)

	federationapi_routing.Setup(
		m.api, *m.cfg, m.queryAPI, m.roomServerProducer, m.keyRing, m.federation,
	)
}
<|endoftext|>


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

package clientv3

import (
	"context"
	"fmt"
	"io"

	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

type (
	DefragmentResponse pb.DefragmentResponse
	AlarmResponse      pb.AlarmResponse
	AlarmMember        pb.AlarmMember
	StatusResponse     pb.StatusResponse
	HashKVResponse     pb.HashKVResponse
	MoveLeaderResponse pb.MoveLeaderResponse
)

type Maintenance interface {
	// AlarmList gets all active alarms.
	AlarmList(ctx context.Context) (*AlarmResponse, error)

	// AlarmDisarm disarms a given alarm.
	AlarmDisarm(ctx context.Context, m *AlarmMember) (*AlarmResponse, error)

	// Defragment releases wasted space from internal fragmentation on a given etcd member.
	// Defragment is only needed when deleting a large number of keys and want to reclaim
	// the resources.
	// Defragment is an expensive operation. User should avoid defragmenting multiple members
	// at the same time.
	// To defragment multiple members in the cluster, user need to call defragment multiple
	// times with different endpoints.
	Defragment(ctx context.Context, endpoint string) (*DefragmentResponse, error)

	// Status gets the status of the endpoint.
	Status(ctx context.Context, endpoint string) (*StatusResponse, error)

	// HashKV returns a hash of the KV state at the time of the RPC.
	// If revision is zero, the hash is computed on all keys. If the revision
	// is non-zero, the hash is computed on all keys at or below the given revision.
	HashKV(ctx context.Context, endpoint string, rev int64) (*HashKVResponse, error)

	// Snapshot provides a reader for a point-in-time snapshot of etcd.
	// If the context "ctx" is canceled or timed out, reading from returned
	// "io.ReadCloser" would error out (e.g. context.Canceled, context.DeadlineExceeded).
	Snapshot(ctx context.Context) (io.ReadCloser, error)

	// MoveLeader requests current leader to transfer its leadership to the transferee.
	// Request must be made to the leader.
	MoveLeader(ctx context.Context, transfereeID uint64) (*MoveLeaderResponse, error)
}

type maintenance struct {
	lg       *zap.Logger
	dial     func(endpoint string) (pb.MaintenanceClient, func(), error)
	remote   pb.MaintenanceClient
	callOpts []grpc.CallOption
}

func NewMaintenance(c *Client) Maintenance {
	api := &maintenance{
		lg: c.lg,
		dial: func(endpoint string) (pb.MaintenanceClient, func(), error) {
			conn, err := c.Dial(endpoint)
			if err != nil {
				return nil, nil, fmt.Errorf("failed to dial endpoint %s with maintenance client: %v", endpoint, err)
			}

			//get token with established connection
			dctx := c.ctx
			cancel := func() {}
			if c.cfg.DialTimeout > 0 {
				dctx, cancel = context.WithTimeout(c.ctx, c.cfg.DialTimeout)
			}
			err = c.getToken(dctx)
			cancel()
			if err != nil {
				return nil, nil, fmt.Errorf("failed to getToken from endpoint %s with maintenance client: %v", endpoint, err)
			}
			cancel = func() { conn.Close() }
			return RetryMaintenanceClient(c, conn), cancel, nil
		},
		remote: RetryMaintenanceClient(c, c.conn),
	}
	if c != nil {
		api.callOpts = c.callOpts
	}
	return api
}

func NewMaintenanceFromMaintenanceClient(remote pb.MaintenanceClient, c *Client) Maintenance {
	api := &maintenance{
		lg: c.lg,
		dial: func(string) (pb.MaintenanceClient, func(), error) {
			return remote, func() {}, nil
		},
		remote: remote,
	}
	if c != nil {
		api.callOpts = c.callOpts
	}
	return api
}

func (m *maintenance) AlarmList(ctx context.Context) (*AlarmResponse, error) {
	req := &pb.AlarmRequest{
		Action:   pb.AlarmRequest_GET,
		MemberID: 0,                 // all
		Alarm:    pb.AlarmType_NONE, // all
	}
	resp, err := m.remote.Alarm(ctx, req, m.callOpts...)
	if err == nil {
		return (*AlarmResponse)(resp), nil
	}
	return nil, toErr(ctx, err)
}

func (m *maintenance) AlarmDisarm(ctx context.Context, am *AlarmMember) (*AlarmResponse, error) {
	req := &pb.AlarmRequest{
		Action:   pb.AlarmRequest_DEACTIVATE,
		MemberID: am.MemberID,
		Alarm:    am.Alarm,
	}

	if req.MemberID == 0 && req.Alarm == pb.AlarmType_NONE {
		ar, err := m.AlarmList(ctx)
		if err != nil {
			return nil, toErr(ctx, err)
		}
		ret := AlarmResponse{}
		for _, am := range ar.Alarms {
			dresp, derr := m.AlarmDisarm(ctx, (*AlarmMember)(am))
			if derr != nil {
				return nil, toErr(ctx, derr)
			}
			ret.Alarms = append(ret.Alarms, dresp.Alarms...)
		}
		return &ret, nil
	}

	resp, err := m.remote.Alarm(ctx, req, m.callOpts...)
	if err == nil {
		return (*AlarmResponse)(resp), nil
	}
	return nil, toErr(ctx, err)
}

func (m *maintenance) Defragment(ctx context.Context, endpoint string) (*DefragmentResponse, error) {
	remote, cancel, err := m.dial(endpoint)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	defer cancel()
	resp, err := remote.Defragment(ctx, &pb.DefragmentRequest{}, m.callOpts...)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	return (*DefragmentResponse)(resp), nil
}

func (m *maintenance) Status(ctx context.Context, endpoint string) (*StatusResponse, error) {
	remote, cancel, err := m.dial(endpoint)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	defer cancel()
	resp, err := remote.Status(ctx, &pb.StatusRequest{}, m.callOpts...)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	return (*StatusResponse)(resp), nil
}

func (m *maintenance) HashKV(ctx context.Context, endpoint string, rev int64) (*HashKVResponse, error) {
	remote, cancel, err := m.dial(endpoint)
	if err != nil {

		return nil, toErr(ctx, err)
	}
	defer cancel()
	resp, err := remote.HashKV(ctx, &pb.HashKVRequest{Revision: rev}, m.callOpts...)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	return (*HashKVResponse)(resp), nil
}

func (m *maintenance) Snapshot(ctx context.Context) (io.ReadCloser, error) {
	ss, err := m.remote.Snapshot(ctx, &pb.SnapshotRequest{}, append(m.callOpts, withMax(defaultStreamMaxRetries))...)
	if err != nil {
		return nil, toErr(ctx, err)
	}

	m.lg.Info("opened snapshot stream; downloading")
	pr, pw := io.Pipe()
	go func() {
		for {
			resp, err := ss.Recv()
			if err != nil {
				switch err {
				case io.EOF:
					m.lg.Info("completed snapshot read; closing")
				default:
					m.lg.Warn("failed to receive from snapshot stream; closing", zap.Error(err))
				}
				pw.CloseWithError(err)
				return
			}

			// can "resp == nil && err == nil"
			// before we receive snapshot SHA digest?
			// No, server sends EOF with an empty response
			// after it sends SHA digest at the end

			if _, werr := pw.Write(resp.Blob); werr != nil {
				pw.CloseWithError(werr)
				return
			}
		}
	}()
	return &snapshotReadCloser{ctx: ctx, ReadCloser: pr}, nil
}

type snapshotReadCloser struct {
	ctx context.Context
	io.ReadCloser
}

func (rc *snapshotReadCloser) Read(p []byte) (n int, err error) {
	n, err = rc.ReadCloser.Read(p)
	return n, toErr(rc.ctx, err)
}

func (m *maintenance) MoveLeader(ctx context.Context, transfereeID uint64) (*MoveLeaderResponse, error) {
	resp, err := m.remote.MoveLeader(ctx, &pb.MoveLeaderRequest{TargetID: transfereeID}, m.callOpts...)
	return (*MoveLeaderResponse)(resp), toErr(ctx, err)
}
<commit_msg>client/v3/maintenance.go: Add Downgrade support to client<commit_after>// Copyright 2016 The etcd Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package clientv3

import (
	"context"
	"errors"
	"fmt"
	"io"

	pb "go.etcd.io/etcd/api/v3/etcdserverpb"
	"go.uber.org/zap"
	"google.golang.org/grpc"
)

type (
	DefragmentResponse pb.DefragmentResponse
	AlarmResponse      pb.AlarmResponse
	AlarmMember        pb.AlarmMember
	StatusResponse     pb.StatusResponse
	HashKVResponse     pb.HashKVResponse
	MoveLeaderResponse pb.MoveLeaderResponse
	DowngradeResponse  pb.DowngradeResponse
)

type Maintenance interface {
	// AlarmList gets all active alarms.
	AlarmList(ctx context.Context) (*AlarmResponse, error)

	// AlarmDisarm disarms a given alarm.
	AlarmDisarm(ctx context.Context, m *AlarmMember) (*AlarmResponse, error)

	// Defragment releases wasted space from internal fragmentation on a given etcd member.
	// Defragment is only needed when deleting a large number of keys and want to reclaim
	// the resources.
	// Defragment is an expensive operation. User should avoid defragmenting multiple members
	// at the same time.
	// To defragment multiple members in the cluster, user need to call defragment multiple
	// times with different endpoints.
	Defragment(ctx context.Context, endpoint string) (*DefragmentResponse, error)

	// Status gets the status of the endpoint.
	Status(ctx context.Context, endpoint string) (*StatusResponse, error)

	// HashKV returns a hash of the KV state at the time of the RPC.
	// If revision is zero, the hash is computed on all keys. If the revision
	// is non-zero, the hash is computed on all keys at or below the given revision.
	HashKV(ctx context.Context, endpoint string, rev int64) (*HashKVResponse, error)

	// Snapshot provides a reader for a point-in-time snapshot of etcd.
	// If the context "ctx" is canceled or timed out, reading from returned
	// "io.ReadCloser" would error out (e.g. context.Canceled, context.DeadlineExceeded).
	Snapshot(ctx context.Context) (io.ReadCloser, error)

	// MoveLeader requests current leader to transfer its leadership to the transferee.
	// Request must be made to the leader.
	MoveLeader(ctx context.Context, transfereeID uint64) (*MoveLeaderResponse, error)

	// Downgrade requests downgrades, verifies feasibility or cancels downgrade
	// on the cluster version.
	// action is one of the following:
	// VALIDATE = 0;
	// ENABLE = 1;
	// CANCEL = 2;
	// Supported since etcd 3.5.
	Downgrade(ctx context.Context, action int32, version string) (*DowngradeResponse, error)
}

type maintenance struct {
	lg       *zap.Logger
	dial     func(endpoint string) (pb.MaintenanceClient, func(), error)
	remote   pb.MaintenanceClient
	callOpts []grpc.CallOption
}

func NewMaintenance(c *Client) Maintenance {
	api := &maintenance{
		lg: c.lg,
		dial: func(endpoint string) (pb.MaintenanceClient, func(), error) {
			conn, err := c.Dial(endpoint)
			if err != nil {
				return nil, nil, fmt.Errorf("failed to dial endpoint %s with maintenance client: %v", endpoint, err)
			}

			//get token with established connection
			dctx := c.ctx
			cancel := func() {}
			if c.cfg.DialTimeout > 0 {
				dctx, cancel = context.WithTimeout(c.ctx, c.cfg.DialTimeout)
			}
			err = c.getToken(dctx)
			cancel()
			if err != nil {
				return nil, nil, fmt.Errorf("failed to getToken from endpoint %s with maintenance client: %v", endpoint, err)
			}
			cancel = func() { conn.Close() }
			return RetryMaintenanceClient(c, conn), cancel, nil
		},
		remote: RetryMaintenanceClient(c, c.conn),
	}
	if c != nil {
		api.callOpts = c.callOpts
	}
	return api
}

func NewMaintenanceFromMaintenanceClient(remote pb.MaintenanceClient, c *Client) Maintenance {
	api := &maintenance{
		lg: c.lg,
		dial: func(string) (pb.MaintenanceClient, func(), error) {
			return remote, func() {}, nil
		},
		remote: remote,
	}
	if c != nil {
		api.callOpts = c.callOpts
	}
	return api
}

func (m *maintenance) AlarmList(ctx context.Context) (*AlarmResponse, error) {
	req := &pb.AlarmRequest{
		Action:   pb.AlarmRequest_GET,
		MemberID: 0,                 // all
		Alarm:    pb.AlarmType_NONE, // all
	}
	resp, err := m.remote.Alarm(ctx, req, m.callOpts...)
	if err == nil {
		return (*AlarmResponse)(resp), nil
	}
	return nil, toErr(ctx, err)
}

func (m *maintenance) AlarmDisarm(ctx context.Context, am *AlarmMember) (*AlarmResponse, error) {
	req := &pb.AlarmRequest{
		Action:   pb.AlarmRequest_DEACTIVATE,
		MemberID: am.MemberID,
		Alarm:    am.Alarm,
	}

	if req.MemberID == 0 && req.Alarm == pb.AlarmType_NONE {
		ar, err := m.AlarmList(ctx)
		if err != nil {
			return nil, toErr(ctx, err)
		}
		ret := AlarmResponse{}
		for _, am := range ar.Alarms {
			dresp, derr := m.AlarmDisarm(ctx, (*AlarmMember)(am))
			if derr != nil {
				return nil, toErr(ctx, derr)
			}
			ret.Alarms = append(ret.Alarms, dresp.Alarms...)
		}
		return &ret, nil
	}

	resp, err := m.remote.Alarm(ctx, req, m.callOpts...)
	if err == nil {
		return (*AlarmResponse)(resp), nil
	}
	return nil, toErr(ctx, err)
}

func (m *maintenance) Defragment(ctx context.Context, endpoint string) (*DefragmentResponse, error) {
	remote, cancel, err := m.dial(endpoint)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	defer cancel()
	resp, err := remote.Defragment(ctx, &pb.DefragmentRequest{}, m.callOpts...)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	return (*DefragmentResponse)(resp), nil
}

func (m *maintenance) Status(ctx context.Context, endpoint string) (*StatusResponse, error) {
	remote, cancel, err := m.dial(endpoint)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	defer cancel()
	resp, err := remote.Status(ctx, &pb.StatusRequest{}, m.callOpts...)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	return (*StatusResponse)(resp), nil
}

func (m *maintenance) HashKV(ctx context.Context, endpoint string, rev int64) (*HashKVResponse, error) {
	remote, cancel, err := m.dial(endpoint)
	if err != nil {

		return nil, toErr(ctx, err)
	}
	defer cancel()
	resp, err := remote.HashKV(ctx, &pb.HashKVRequest{Revision: rev}, m.callOpts...)
	if err != nil {
		return nil, toErr(ctx, err)
	}
	return (*HashKVResponse)(resp), nil
}

func (m *maintenance) Snapshot(ctx context.Context) (io.ReadCloser, error) {
	ss, err := m.remote.Snapshot(ctx, &pb.SnapshotRequest{}, append(m.callOpts, withMax(defaultStreamMaxRetries))...)
	if err != nil {
		return nil, toErr(ctx, err)
	}

	m.lg.Info("opened snapshot stream; downloading")
	pr, pw := io.Pipe()
	go func() {
		for {
			resp, err := ss.Recv()
			if err != nil {
				switch err {
				case io.EOF:
					m.lg.Info("completed snapshot read; closing")
				default:
					m.lg.Warn("failed to receive from snapshot stream; closing", zap.Error(err))
				}
				pw.CloseWithError(err)
				return
			}

			// can "resp == nil && err == nil"
			// before we receive snapshot SHA digest?
			// No, server sends EOF with an empty response
			// after it sends SHA digest at the end

			if _, werr := pw.Write(resp.Blob); werr != nil {
				pw.CloseWithError(werr)
				return
			}
		}
	}()
	return &snapshotReadCloser{ctx: ctx, ReadCloser: pr}, nil
}

type snapshotReadCloser struct {
	ctx context.Context
	io.ReadCloser
}

func (rc *snapshotReadCloser) Read(p []byte) (n int, err error) {
	n, err = rc.ReadCloser.Read(p)
	return n, toErr(rc.ctx, err)
}

func (m *maintenance) MoveLeader(ctx context.Context, transfereeID uint64) (*MoveLeaderResponse, error) {
	resp, err := m.remote.MoveLeader(ctx, &pb.MoveLeaderRequest{TargetID: transfereeID}, m.callOpts...)
	return (*MoveLeaderResponse)(resp), toErr(ctx, err)
}

func (m *maintenance) Downgrade(ctx context.Context, action int32, version string) (*DowngradeResponse, error) {
	actionType := pb.DowngradeRequest_VALIDATE
	switch action {
	case 0:
		actionType = pb.DowngradeRequest_VALIDATE
	case 1:
		actionType = pb.DowngradeRequest_ENABLE
	case 2:
		actionType = pb.DowngradeRequest_CANCEL
	default:
		return nil, errors.New("etcdclient: unknown downgrade action")
	}
	resp, err := m.remote.Downgrade(ctx, &pb.DowngradeRequest{Action: actionType, Version: version}, m.callOpts...)
	return (*DowngradeResponse)(resp), toErr(ctx, err)
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"log"
	"os"
	"time"

	"github.com/oinume/lekcije/server/bootstrap"
	"github.com/oinume/lekcije/server/errors"
	"github.com/oinume/lekcije/server/fetcher"
	"github.com/oinume/lekcije/server/logger"
	"github.com/oinume/lekcije/server/model"
	"github.com/oinume/lekcije/server/notifier"
	"github.com/pkg/profile"
	"github.com/uber-go/zap"
)

var (
	dryRun       = flag.Bool("dry-run", false, "Don't update database with fetched lessons")
	sendEmail    = flag.Bool("send-email", true, "flag to send email")
	concurrency  = flag.Int("concurrency", 1, "concurrency of fetcher")
	fetcherCache = flag.Bool("fetcher-cache", false, "Cache teacher and lesson data in Fetcher")
	logLevel     = flag.String("log-level", "info", "Log level")
	profileMode  = flag.String("profile-mode", "", "block|cpu|mem|trace")
)

func main() {
	flag.Parse()
	if err := run(); err != nil {
		log.Fatalf("err = %v", err) // TODO: Error handling
	}
	os.Exit(0)
}

func run() error {
	switch *profileMode {
	case "block":
		defer profile.Start(profile.ProfilePath("."), profile.BlockProfile).Stop()
	case "cpu":
		defer profile.Start(profile.ProfilePath("."), profile.CPUProfile).Stop()
	case "mem":
		defer profile.Start(profile.ProfilePath("."), profile.MemProfile).Stop()
	case "trace":
		defer profile.Start(profile.ProfilePath("."), profile.TraceProfile).Stop()
	}

	bootstrap.CheckCLIEnvVars()
	startedAt := time.Now().UTC()
	if *logLevel != "" {
		logger.App.SetLevel(logger.NewLevel(*logLevel))
	}
	logger.App.Info("notifier started")
	defer func() {
		elapsed := time.Now().UTC().Sub(startedAt) / time.Millisecond
		logger.App.Info("notifier finished", zap.Int("elapsed", int(elapsed)))
	}()

	// TODO: Wrap up as function
	var dbLogging bool
	//dbLogging := !config.IsProductionEnv()
	if *logLevel == "debug" {
		dbLogging = true
	} else {
		dbLogging = false
	}
	db, err := model.OpenDB(bootstrap.CLIEnvVars.DBURL, 1, dbLogging)
	if err != nil {
		return err
	}
	defer db.Close()

	users, err := model.NewUserService(db).FindAllEmailVerifiedIsTrue()
	if err != nil {
		return err
	}
	mCountries, err := model.NewMCountryService(db).LoadAll()
	if err != nil {
		return errors.InternalWrapf(err, "Failed to load all MCountries")
	}
	fetcher := fetcher.NewTeacherLessonFetcher(nil, *concurrency, *fetcherCache, mCountries, logger.App)
	notifier := notifier.NewNotifier(db, fetcher, *dryRun, *sendEmail)
	defer notifier.Close()
	for _, user := range users {
		if err := notifier.SendNotification(user); err != nil {
			return err
		}
	}

	return nil
}
<commit_msg>Fix problem of staticcheck<commit_after>package main

import (
	"flag"
	"log"
	"os"
	"time"

	"github.com/oinume/lekcije/server/bootstrap"
	"github.com/oinume/lekcije/server/errors"
	"github.com/oinume/lekcije/server/fetcher"
	"github.com/oinume/lekcije/server/logger"
	"github.com/oinume/lekcije/server/model"
	"github.com/oinume/lekcije/server/notifier"
	"github.com/pkg/profile"
	"github.com/uber-go/zap"
)

var (
	dryRun       = flag.Bool("dry-run", false, "Don't update database with fetched lessons")
	sendEmail    = flag.Bool("send-email", true, "flag to send email")
	concurrency  = flag.Int("concurrency", 1, "concurrency of fetcher")
	fetcherCache = flag.Bool("fetcher-cache", false, "Cache teacher and lesson data in Fetcher")
	logLevel     = flag.String("log-level", "info", "Log level")
	profileMode  = flag.String("profile-mode", "", "block|cpu|mem|trace")
)

func main() {
	flag.Parse()
	if err := run(); err != nil {
		log.Fatalf("err = %v", err) // TODO: Error handling
	}
	os.Exit(0)
}

func run() error {
	switch *profileMode {
	case "block":
		defer profile.Start(profile.ProfilePath("."), profile.BlockProfile).Stop()
	case "cpu":
		defer profile.Start(profile.ProfilePath("."), profile.CPUProfile).Stop()
	case "mem":
		defer profile.Start(profile.ProfilePath("."), profile.MemProfile).Stop()
	case "trace":
		defer profile.Start(profile.ProfilePath("."), profile.TraceProfile).Stop()
	}

	bootstrap.CheckCLIEnvVars()
	startedAt := time.Now().UTC()
	if *logLevel != "" {
		logger.App.SetLevel(logger.NewLevel(*logLevel))
	}
	logger.App.Info("notifier started")
	defer func() {
		elapsed := time.Now().UTC().Sub(startedAt) / time.Millisecond
		logger.App.Info("notifier finished", zap.Int("elapsed", int(elapsed)))
	}()

	// TODO: Wrap up as function
	dbLogging := false
	// TODO: something wrong with staticcheck? this value of dbLogging is never used (SA4006)
	//dbLogging := !config.IsProductionEnv()x
	if *logLevel == "debug" {
		dbLogging = true
	} else {
		dbLogging = false
	}
	db, err := model.OpenDB(bootstrap.CLIEnvVars.DBURL, 1, dbLogging)
	if err != nil {
		return err
	}
	defer db.Close()

	users, err := model.NewUserService(db).FindAllEmailVerifiedIsTrue()
	if err != nil {
		return err
	}
	mCountries, err := model.NewMCountryService(db).LoadAll()
	if err != nil {
		return errors.InternalWrapf(err, "Failed to load all MCountries")
	}
	fetcher := fetcher.NewTeacherLessonFetcher(nil, *concurrency, *fetcherCache, mCountries, logger.App)
	notifier := notifier.NewNotifier(db, fetcher, *dryRun, *sendEmail)
	defer notifier.Close()
	for _, user := range users {
		if err := notifier.SendNotification(user); err != nil {
			return err
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package henchman

import (
	"log"

	"code.google.com/p/go-uuid/uuid"
	"github.com/flosch/pongo2"

	"github.com/sudharsh/henchman/ansi"
)

var statuses = map[string]string{
	"reset":   ansi.ColorCode("reset"),
	"success": ansi.ColorCode("green"),
	"ignored": ansi.ColorCode("yellow"),
	"failure": ansi.ColorCode("red"),
}

// Task is the unit of work in henchman.
type Task struct {
	Id string

	Name         string
	Action       string
	IgnoreErrors bool `yaml:"ignore_errors"`
}

func prepareTemplate(data string, vars *TaskVars, machine *Machine) (string, error) {
	tmpl, err := pongo2.FromString(data)
	if err != nil {
		panic(err)
	}
	return tmpl.Execute(&pongo2.Context{"vars": vars, "machine": machine})
}

// Renders the template parts in the task field.
// Also assigns a new UUID to the task uniquely identifying it.
func (task *Task) prepare(vars *TaskVars, machine *Machine) {
	var err error
	task.Id = uuid.New()
	task.Name, err = prepareTemplate(task.Name, vars, machine)
	if err != nil {
		panic(err)
	}
	task.Action, err = prepareTemplate(task.Action, vars, machine)
	if err != nil {
		panic(err)
	}
}

// Runs the task on the machine. The task might mutate `vars` so that other
// tasks down the `plan` can see any additions/updates.
func (task *Task) Run(machine *Machine, vars *TaskVars) string {
	task.prepare(vars, machine)
	log.Printf("%s: %s '%s'\n", task.Id, machine.Hostname, task.Name)
	out, err := machine.Exec(task.Action)
	var taskStatus string = "success"
	if err != nil {
		if task.IgnoreErrors {
			taskStatus = "ignored"
		} else {
			taskStatus = "failure"
		}
	}
	escapeCode := statuses[taskStatus]
	var reset string = statuses["reset"]
	log.Printf("%s: %s [%s] - %s", task.Id, escapeCode, taskStatus, out.String()+reset)
	return taskStatus
}
<commit_msg>Fix build<commit_after>package henchman

import (
	"log"

	"code.google.com/p/go-uuid/uuid"
	"github.com/flosch/pongo2"

	"github.com/sudharsh/henchman/ansi"
)

var statuses = map[string]string{
	"reset":   ansi.ColorCode("reset"),
	"success": ansi.ColorCode("green"),
	"ignored": ansi.ColorCode("yellow"),
	"failure": ansi.ColorCode("red"),
}

// Task is the unit of work in henchman.
type Task struct {
	Id string

	Name         string
	Action       string
	IgnoreErrors bool `yaml:"ignore_errors"`
}

func prepareTemplate(data string, vars *TaskVars, machine *Machine) (string, error) {
	tmpl, err := pongo2.FromString(data)
	if err != nil {
		panic(err)
	}
	ctxt := pongo2.Context{"vars": vars, "machine": machine}
	return tmpl.Execute(&ctxt)
}

// Renders the template parts in the task field.
// Also assigns a new UUID to the task uniquely identifying it.
func (task *Task) prepare(vars *TaskVars, machine *Machine) {
	var err error
	task.Id = uuid.New()
	task.Name, err = prepareTemplate(task.Name, vars, machine)
	if err != nil {
		panic(err)
	}
	task.Action, err = prepareTemplate(task.Action, vars, machine)
	if err != nil {
		panic(err)
	}
}

// Runs the task on the machine. The task might mutate `vars` so that other
// tasks down the `plan` can see any additions/updates.
func (task *Task) Run(machine *Machine, vars *TaskVars) string {
	task.prepare(vars, machine)
	log.Printf("%s: %s '%s'\n", task.Id, machine.Hostname, task.Name)
	out, err := machine.Exec(task.Action)
	var taskStatus string = "success"
	if err != nil {
		if task.IgnoreErrors {
			taskStatus = "ignored"
		} else {
			taskStatus = "failure"
		}
	}
	escapeCode := statuses[taskStatus]
	var reset string = statuses["reset"]
	log.Printf("%s: %s [%s] - %s", task.Id, escapeCode, taskStatus, out.String()+reset)
	return taskStatus
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"github.com/GeertJohan/go.rice/embedded"
	"github.com/akavel/rsrc/coff"
	"go/build"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"
	"time"
)

type sizedBytes []byte

func (s sizedBytes) Size() int64 {
	return int64(len(s))
}

var tmplEmbeddedSysoHelper *template.Template

func init() {
	var err error
	tmplEmbeddedSysoHelper, err = template.New("embeddedSysoHelper").Parse(`package {{.Package}}

// extern char _bricebox_{{.Symname}}[], _ericebox_{{.Symname}};
// int get_{{.Symname}}_length() {
// 	return &_ericebox_{{.Symname}} - _bricebox_{{.Symname}};
// }
import "C"
import (
	"bytes"
	"encoding/gob"
	"github.com/GeertJohan/go.rice/embedded"
	"unsafe"
)

func init() {
	ptr := unsafe.Pointer(&C._bricebox_{{.Symname}})
	bts := C.GoBytes(ptr, C.get_{{.Symname}}_length())
	embeddedBox := &embedded.EmbeddedBox{}
	err := gob.NewDecoder(bytes.NewReader(bts)).Decode(embeddedBox)
	if err != nil {
		panic("error decoding embedded box: "+err.Error())
	}
	embeddedBox.Link()
	embedded.RegisterEmbeddedBox(embeddedBox.Name, embeddedBox)
}`)
	if err != nil {
		panic("could not parse template embeddedSysoHelper: " + err.Error())
	}
}

type embeddedSysoHelperData struct {
	Package string
	Symname string
}

func operationEmbedSyso(pkg *build.Package) {

	regexpSynameReplacer := regexp.MustCompile(`[^a-z0-9_]`)

	boxMap := findBoxes(pkg)

	// notify user when no calls to rice.FindBox are made (is this an error and therefore os.Exit(1) ?
	if len(boxMap) == 0 {
		fmt.Println("no calls to rice.FindBox() found")
		return
	}

	verbosef("\n")

	for boxname := range boxMap {
		// find path and filename for this box
		boxPath := filepath.Join(pkg.Dir, boxname)
		boxFilename := strings.Replace(boxname, "/", "-", -1)
		boxFilename = strings.Replace(boxFilename, "..", "back", -1)

		// verbose info
		verbosef("embedding box '%s'\n", boxname)
		verbosef("\tto file %s\n", boxFilename)

		// create box datastructure (used by template)
		box := &embedded.EmbeddedBox{
			Name:      boxname,
			Time:      time.Now(),
			EmbedType: embedded.EmbedTypeSyso,
			Files:     make(map[string]*embedded.EmbeddedFile),
			Dirs:      make(map[string]*embedded.EmbeddedDir),
		}

		// fill box datastructure with file data
		filepath.Walk(boxPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf("error walking box: %s\n", err)
				os.Exit(1)
			}

			filename := strings.TrimPrefix(path, boxPath)
			filename = strings.Replace(filename, "\\", "/", -1)
			filename = strings.TrimPrefix(filename, "/")
			if info.IsDir() {
				embeddedDir := &embedded.EmbeddedDir{
					Filename:   filename,
					DirModTime: info.ModTime(),
				}
				verbosef("\tincludes dir: '%s'\n", embeddedDir.Filename)
				box.Dirs[embeddedDir.Filename] = embeddedDir

				// add tree entry (skip for root, it'll create a recursion)
				if embeddedDir.Filename != "" {
					pathParts := strings.Split(embeddedDir.Filename, "/")
					parentDir := box.Dirs[strings.Join(pathParts[:len(pathParts)-1], "/")]
					parentDir.ChildDirs = append(parentDir.ChildDirs, embeddedDir)
				}
			} else {
				embeddedFile := &embedded.EmbeddedFile{
					Filename:    filename,
					FileModTime: info.ModTime(),
					Content:     "",
				}
				verbosef("\tincludes file: '%s'\n", embeddedFile.Filename)
				contentBytes, err := ioutil.ReadFile(path)
				if err != nil {
					fmt.Printf("error reading file content while walking box: %s\n", err)
					os.Exit(1)
				}
				embeddedFile.Content = string(contentBytes)
				box.Files[embeddedFile.Filename] = embeddedFile
			}
			return nil
		})

		// encode embedded box to gob file
		boxGobBuf := &bytes.Buffer{}
		err := gob.NewEncoder(boxGobBuf).Encode(box)
		if err != nil {
			fmt.Printf("error encoding box to gob: %v\n", err)
			os.Exit(1)
		}

		// write coff
		symname := regexpSynameReplacer.ReplaceAllString(boxname, "_")
		createCoffSyso(boxname, symname, "386", boxGobBuf.Bytes())
		createCoffSyso(boxname, symname, "amd64", boxGobBuf.Bytes())

		// write go
		sysoHelperData := embeddedSysoHelperData{
			Package: pkg.Name,
			Symname: symname,
		}
		fileSysoHelper, err := os.Create(boxFilename + ".rice-box.go")
		if err != nil {
			fmt.Printf("error creating syso helper: %v\n", err)
			os.Exit(1)
		}
		err = tmplEmbeddedSysoHelper.Execute(fileSysoHelper, sysoHelperData)
		if err != nil {
			fmt.Printf("error executing tmplEmbeddedSysoHelper: %v\n", err)
			os.Exit(1)
		}
	}
}

func createCoffSyso(boxFilename string, symname string, arch string, data []byte) {
	boxCoff := coff.NewRDATA()
	switch arch {
	case "386":
	case "amd64":
		boxCoff.FileHeader.Machine = 0x8664
	default:
		panic("invalid arch")
	}
	boxCoff.AddData("_bricebox_"+symname, sizedBytes(data))
	boxCoff.AddData("_ericebox_"+symname, io.NewSectionReader(strings.NewReader("\000\000"), 0, 2)) // TODO: why? copied from rsrc, which copied it from as-generated
	boxCoff.Freeze()
	err := writeCoff(boxCoff, boxFilename+"_"+arch+".rice-box.syso")
	if err != nil {
		fmt.Printf("error writing %s coff/.syso: %v\n", arch, err)
		os.Exit(1)
	}
}
<commit_msg>Add comment for generated code<commit_after>package main

import (
	"bytes"
	"encoding/gob"
	"fmt"
	"github.com/GeertJohan/go.rice/embedded"
	"github.com/akavel/rsrc/coff"
	"go/build"
	"io"
	"io/ioutil"
	"os"
	"path/filepath"
	"regexp"
	"strings"
	"text/template"
	"time"
)

type sizedBytes []byte

func (s sizedBytes) Size() int64 {
	return int64(len(s))
}

var tmplEmbeddedSysoHelper *template.Template

func init() {
	var err error
	tmplEmbeddedSysoHelper, err = template.New("embeddedSysoHelper").Parse(`package {{.Package}}
// ############# GENERATED CODE #####################
// ## This file was generated by the rice tool.
// ## Do not edit unless you know what you're doing.
// ##################################################

// extern char _bricebox_{{.Symname}}[], _ericebox_{{.Symname}};
// int get_{{.Symname}}_length() {
// 	return &_ericebox_{{.Symname}} - _bricebox_{{.Symname}};
// }
import "C"
import (
	"bytes"
	"encoding/gob"
	"github.com/GeertJohan/go.rice/embedded"
	"unsafe"
)

func init() {
	ptr := unsafe.Pointer(&C._bricebox_{{.Symname}})
	bts := C.GoBytes(ptr, C.get_{{.Symname}}_length())
	embeddedBox := &embedded.EmbeddedBox{}
	err := gob.NewDecoder(bytes.NewReader(bts)).Decode(embeddedBox)
	if err != nil {
		panic("error decoding embedded box: "+err.Error())
	}
	embeddedBox.Link()
	embedded.RegisterEmbeddedBox(embeddedBox.Name, embeddedBox)
}`)
	if err != nil {
		panic("could not parse template embeddedSysoHelper: " + err.Error())
	}
}

type embeddedSysoHelperData struct {
	Package string
	Symname string
}

func operationEmbedSyso(pkg *build.Package) {

	regexpSynameReplacer := regexp.MustCompile(`[^a-z0-9_]`)

	boxMap := findBoxes(pkg)

	// notify user when no calls to rice.FindBox are made (is this an error and therefore os.Exit(1) ?
	if len(boxMap) == 0 {
		fmt.Println("no calls to rice.FindBox() found")
		return
	}

	verbosef("\n")

	for boxname := range boxMap {
		// find path and filename for this box
		boxPath := filepath.Join(pkg.Dir, boxname)
		boxFilename := strings.Replace(boxname, "/", "-", -1)
		boxFilename = strings.Replace(boxFilename, "..", "back", -1)

		// verbose info
		verbosef("embedding box '%s'\n", boxname)
		verbosef("\tto file %s\n", boxFilename)

		// create box datastructure (used by template)
		box := &embedded.EmbeddedBox{
			Name:      boxname,
			Time:      time.Now(),
			EmbedType: embedded.EmbedTypeSyso,
			Files:     make(map[string]*embedded.EmbeddedFile),
			Dirs:      make(map[string]*embedded.EmbeddedDir),
		}

		// fill box datastructure with file data
		filepath.Walk(boxPath, func(path string, info os.FileInfo, err error) error {
			if err != nil {
				fmt.Printf("error walking box: %s\n", err)
				os.Exit(1)
			}

			filename := strings.TrimPrefix(path, boxPath)
			filename = strings.Replace(filename, "\\", "/", -1)
			filename = strings.TrimPrefix(filename, "/")
			if info.IsDir() {
				embeddedDir := &embedded.EmbeddedDir{
					Filename:   filename,
					DirModTime: info.ModTime(),
				}
				verbosef("\tincludes dir: '%s'\n", embeddedDir.Filename)
				box.Dirs[embeddedDir.Filename] = embeddedDir

				// add tree entry (skip for root, it'll create a recursion)
				if embeddedDir.Filename != "" {
					pathParts := strings.Split(embeddedDir.Filename, "/")
					parentDir := box.Dirs[strings.Join(pathParts[:len(pathParts)-1], "/")]
					parentDir.ChildDirs = append(parentDir.ChildDirs, embeddedDir)
				}
			} else {
				embeddedFile := &embedded.EmbeddedFile{
					Filename:    filename,
					FileModTime: info.ModTime(),
					Content:     "",
				}
				verbosef("\tincludes file: '%s'\n", embeddedFile.Filename)
				contentBytes, err := ioutil.ReadFile(path)
				if err != nil {
					fmt.Printf("error reading file content while walking box: %s\n", err)
					os.Exit(1)
				}
				embeddedFile.Content = string(contentBytes)
				box.Files[embeddedFile.Filename] = embeddedFile
			}
			return nil
		})

		// encode embedded box to gob file
		boxGobBuf := &bytes.Buffer{}
		err := gob.NewEncoder(boxGobBuf).Encode(box)
		if err != nil {
			fmt.Printf("error encoding box to gob: %v\n", err)
			os.Exit(1)
		}

		// write coff
		symname := regexpSynameReplacer.ReplaceAllString(boxname, "_")
		createCoffSyso(boxname, symname, "386", boxGobBuf.Bytes())
		createCoffSyso(boxname, symname, "amd64", boxGobBuf.Bytes())

		// write go
		sysoHelperData := embeddedSysoHelperData{
			Package: pkg.Name,
			Symname: symname,
		}
		fileSysoHelper, err := os.Create(boxFilename + ".rice-box.go")
		if err != nil {
			fmt.Printf("error creating syso helper: %v\n", err)
			os.Exit(1)
		}
		err = tmplEmbeddedSysoHelper.Execute(fileSysoHelper, sysoHelperData)
		if err != nil {
			fmt.Printf("error executing tmplEmbeddedSysoHelper: %v\n", err)
			os.Exit(1)
		}
	}
}

func createCoffSyso(boxFilename string, symname string, arch string, data []byte) {
	boxCoff := coff.NewRDATA()
	switch arch {
	case "386":
	case "amd64":
		boxCoff.FileHeader.Machine = 0x8664
	default:
		panic("invalid arch")
	}
	boxCoff.AddData("_bricebox_"+symname, sizedBytes(data))
	boxCoff.AddData("_ericebox_"+symname, io.NewSectionReader(strings.NewReader("\000\000"), 0, 2)) // TODO: why? copied from rsrc, which copied it from as-generated
	boxCoff.Freeze()
	err := writeCoff(boxCoff, boxFilename+"_"+arch+".rice-box.syso")
	if err != nil {
		fmt.Printf("error writing %s coff/.syso: %v\n", arch, err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"time"

	"github.com/hprose/hprose-golang/rpc"
)

// Args ...
type Args struct {
	A, B int
}

// Quotient ...
type Quotient struct {
	Quo, Rem int
}

// Stub ...
type Stub struct {
	// Synchronous call
	Multiply func(args *Args) int
	// Asynchronous call
	Divide func(func(*Quotient, error), *Args)
}

func main() {
	client := rpc.NewClient("http://127.0.0.1:8080")
	var stub *Stub
	client.UseService(&stub)
	fmt.Println(stub.Multiply(&Args{8, 7}))
	stub.Divide(func(result *Quotient, err error) {
		if err != nil {
			log.Fatal("arith error:", err)
		} else {
			fmt.Println(result.Quo, result.Rem)
		}
	}, &Args{8, 7})
	time.Sleep(1 * time.Second)
}
<commit_msg>Improved example<commit_after>package main

import (
	"fmt"
	"log"

	"github.com/hprose/hprose-golang/rpc"
)

// Args ...
type Args struct {
	A, B int
}

// Quotient ...
type Quotient struct {
	Quo, Rem int
}

// Stub ...
type Stub struct {
	// Synchronous call
	Multiply func(args *Args) int
	// Asynchronous call
	Divide func(func(*Quotient, error), *Args)
}

func main() {
	client := rpc.NewClient("http://127.0.0.1:8080")
	var stub *Stub
	client.UseService(&stub)
	fmt.Println(stub.Multiply(&Args{8, 7}))
	done := make(chan struct{})
	stub.Divide(func(result *Quotient, err error) {
		if err != nil {
			log.Fatal("arith error:", err)
		} else {
			fmt.Println(result.Quo, result.Rem)
		}
		done <- struct{}{}
	}, &Args{8, 7})
	<-done
}
<|endoftext|>


<commit_before>package irmaserver

import (
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/go-chi/chi"
	"github.com/go-errors/errors"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/server"
	"github.com/privacybydesign/irmago/server/irmarequestor"
)

type Configuration struct {
	*server.Configuration
	Port int
}

var s *http.Server

// Start the server. If successful then it will not return until Stop() is called.
func Start(conf *Configuration) error {
	if err := irmarequestor.Initialize(conf.Configuration); err != nil {
		return err
	}

	router := chi.NewRouter()

	// Mount server for irmaclient
	router.Mount("/irma/", irmarequestor.HttpHandlerFunc("/irma/"))

	// Server routes
	router.Post("/create", handleCreate)
	router.Get("/status/{token}", handleStatus)
	router.Get("/result/{token}", handleResult)

	// Start server
	s = &http.Server{Addr: fmt.Sprintf(":%d", conf.Port), Handler: router}
	err := s.ListenAndServe()
	if err == http.ErrServerClosed {
		return nil // Server was closed normally
	}
	return err
}

func Stop() {
	s.Close()
}

func handleCreate(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		server.WriteError(w, server.ErrorInvalidRequest, err.Error())
		return
	}
	request, err := parseRequest(body)
	if err != nil {
		server.WriteError(w, server.ErrorInvalidRequest, err.Error())
		return
	}

	qr, _, err := irmarequestor.StartSession(request, nil)
	if err != nil {
		server.WriteError(w, server.ErrorInvalidRequest, err.Error())
		return
	}

	server.WriteJson(w, qr)
}

func handleStatus(w http.ResponseWriter, r *http.Request) {
	res := irmarequestor.GetSessionResult(chi.URLParam(r, "token"))
	if res == nil {
		server.WriteError(w, server.ErrorSessionUnknown, "")
		return
	}
	server.WriteJson(w, res.Status)
}

func handleResult(w http.ResponseWriter, r *http.Request) {
	res := irmarequestor.GetSessionResult(chi.URLParam(r, "token"))
	if res == nil {
		server.WriteError(w, server.ErrorSessionUnknown, "")
		return
	}
	server.WriteJson(w, res)
}

func parseRequest(bts []byte) (request irma.SessionRequest, err error) {
	request = &irma.DisclosureRequest{}
	if err = irma.UnmarshalValidate(bts, request); err == nil {
		return request, nil
	}
	request = &irma.SignatureRequest{}
	if err = irma.UnmarshalValidate(bts, request); err == nil {
		return request, nil
	}
	request = &irma.IssuanceRequest{}
	if err = irma.UnmarshalValidate(bts, request); err == nil {
		return request, nil
	}
	return nil, errors.New("Invalid session type")
}
<commit_msg>Allow irmaserver to be used as library<commit_after>package irmaserver

import (
	"fmt"
	"io/ioutil"
	"net/http"

	"github.com/go-chi/chi"
	"github.com/go-errors/errors"
	"github.com/privacybydesign/irmago"
	"github.com/privacybydesign/irmago/server"
	"github.com/privacybydesign/irmago/server/irmarequestor"
)

type Configuration struct {
	*server.Configuration
	Port int
}

var s *http.Server

// Start the server. If successful then it will not return until Stop() is called.
func Start(conf *Configuration) error {
	handler, err := Handler(conf.Configuration)
	if err != nil {
		return err
	}

	// Start server
	s = &http.Server{Addr: fmt.Sprintf(":%d", conf.Port), Handler: handler}
	err = s.ListenAndServe()
	if err == http.ErrServerClosed {
		return nil // Server was closed normally
	}

	return err
}

func Stop() {
	s.Close()
}

func Handler(conf *server.Configuration) (http.Handler, error) {
	if err := irmarequestor.Initialize(conf); err != nil {
		return nil, err
	}

	router := chi.NewRouter()

	// Mount server for irmaclient
	router.Mount("/irma/", irmarequestor.HttpHandlerFunc("/irma/"))

	// Server routes
	router.Post("/create", handleCreate)
	router.Get("/status/{token}", handleStatus)
	router.Get("/result/{token}", handleResult)

	return router, nil
}

func handleCreate(w http.ResponseWriter, r *http.Request) {
	body, err := ioutil.ReadAll(r.Body)
	if err != nil {
		server.WriteError(w, server.ErrorInvalidRequest, err.Error())
		return
	}
	request, err := parseRequest(body)
	if err != nil {
		server.WriteError(w, server.ErrorInvalidRequest, err.Error())
		return
	}

	qr, _, err := irmarequestor.StartSession(request, nil)
	if err != nil {
		server.WriteError(w, server.ErrorInvalidRequest, err.Error())
		return
	}

	server.WriteJson(w, qr)
}

func handleStatus(w http.ResponseWriter, r *http.Request) {
	res := irmarequestor.GetSessionResult(chi.URLParam(r, "token"))
	if res == nil {
		server.WriteError(w, server.ErrorSessionUnknown, "")
		return
	}
	server.WriteJson(w, res.Status)
}

func handleResult(w http.ResponseWriter, r *http.Request) {
	res := irmarequestor.GetSessionResult(chi.URLParam(r, "token"))
	if res == nil {
		server.WriteError(w, server.ErrorSessionUnknown, "")
		return
	}
	server.WriteJson(w, res)
}

func parseRequest(bts []byte) (request irma.SessionRequest, err error) {
	request = &irma.DisclosureRequest{}
	if err = irma.UnmarshalValidate(bts, request); err == nil {
		return request, nil
	}
	request = &irma.SignatureRequest{}
	if err = irma.UnmarshalValidate(bts, request); err == nil {
		return request, nil
	}
	request = &irma.IssuanceRequest{}
	if err = irma.UnmarshalValidate(bts, request); err == nil {
		return request, nil
	}
	return nil, errors.New("Invalid session type")
}
<|endoftext|>


<commit_before>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

// This binary provides sample code for using the gopacket TCP assembler and TCP
// stream reader.  It reads packets off the wire and reconstructs HTTP requests
// it sees, logging them.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/examples/util"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/tcpassembly"
	"github.com/google/gopacket/tcpassembly/tcpreader"
)

var iface = flag.String("i", "eth0", "Interface to get packets from")
var fname = flag.String("r", "", "Filename to read from, overrides -i")
var snaplen = flag.Int("s", 1600, "SnapLen for pcap packet capture")
var filter = flag.String("f", "tcp and dst port 80", "BPF filter for pcap")
var logAllPackets = flag.Bool("v", false, "Logs every packet in great detail")

// Build a simple HTTP request parser using tcpassembly.StreamFactory and tcpassembly.Stream interfaces

// httpStreamFactory implements tcpassembly.StreamFactory
type httpStreamFactory struct{}

// httpStream will handle the actual decoding of http requests.
type httpStream struct {
	net, transport gopacket.Flow
	r              tcpreader.ReaderStream
}

func (h *httpStreamFactory) New(net, transport gopacket.Flow) tcpassembly.Stream {
	hstream := &httpStream{
		net:       net,
		transport: transport,
		r:         tcpreader.NewReaderStream(),
	}
	go hstream.run() // Important... we must guarantee that data from the reader stream is read.

	// ReaderStream implements tcpassembly.Stream, so we can return a pointer to it.
	return &hstream.r
}

func (h *httpStream) run() {
	buf := bufio.NewReader(&h.r)
	for {
		resp, err := http.ReadResponse(buf, nil)
		if err == io.EOF || err == io.ErrUnexpectedEOF {
			// We must read until we see an EOF... very important!
			return
		} else if err != nil {
			//log.Println("ERROR IN RESPONSE:", h.net, ":", err)
		} else {

			contentType := resp.Header["Content-Type"]
			//contentEnc := resp.Header["Content-Encoding"]

			//log.Println("ENCODING:", resp.TransferEncoding, ":", contentEnc, ":", resp.Uncompressed)

			if len(contentType) != 0 {

				reader := resp.Body
				/*
					if len(contentEnc) != 0 {
						if contentEnc[0] == "gzip" {
							r, qerr := gzip.NewReader(resp.Body)
							if qerr != nil {
								log.Println("ERROR GZIP:", qerr)
							}
							reader = r
						}
					}
				*/

				switch contentType[0] {
				// TODO: ASCII, ANSI (Windows-1252)
				case "text/html", "text/html; charset=utf-8", "text/html; charset=UTF-8":
					// Default charset for HTML5
					//fmt.Println("FOUND ONE:", contentType)

					log.Print("MATCHED:", contentType[0])

					b, err := ioutil.ReadAll(reader)
					if err != nil {
						log.Println(err)
					}
					fmt.Println(b)
					/*
						body, perr := html.Parse(resp.Body)
						if perr != nil {
							log.Println("PARSE ERROR:", perr)
							break
						} else {
							doc := goquery.NewDocumentFromNode(body)
							fmt.Println("DOC:", doc.Find("h1").Text())
						}
					*/
				case "text/html; charset=iso-8859-1", "text/html; charset=ISO-8859-1":
					// Default charset before HTML5
					// TODO: Do something with it, e.g. convert with iconv.
				default:
					//log.Println("UNUSED TYPE:", contentType)
				}
			}

			for {
				bytes, err := tcpreader.DiscardBytesToFirstError(resp.Body)

				if err == io.EOF || err == io.ErrUnexpectedEOF {
					break

				} else if err != nil && bytes == 0 {
					log.Println("ERROR BUT ZERO:", h.net, ":", err, ":", bytes)
					break

				} else if err != nil {
					log.Println("ERROR NOT ZERO:", h.net, ":", err, ":", bytes)
				}
			}
			resp.Body.Close()
			//log.Println(h.net, ":", contentType, resp.Status)
		}
	}
}

func main() {
	defer util.Run()()
	var handle *pcap.Handle
	var err error

	// Set up pcap packet capture
	if *fname != "" {
		log.Printf("Reading from pcap dump %q", *fname)
		handle, err = pcap.OpenOffline(*fname)
	} else {
		log.Printf("Starting capture on interface %q", *iface)
		handle, err = pcap.OpenLive(*iface, int32(*snaplen), true, pcap.BlockForever)
	}
	if err != nil {
		log.Fatal(err)
	}

	if err := handle.SetBPFFilter(*filter); err != nil {
		log.Fatal(err)
	}

	// Set up assembly
	streamFactory := &httpStreamFactory{}
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	assembler := tcpassembly.NewAssembler(streamPool)

	log.Println("reading in packets")
	// Read in packets, pass to assembler.
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets := packetSource.Packets()
	ticker := time.Tick(time.Minute)
	for {
		select {
		case packet := <-packets:
			// A nil packet indicates the end of a pcap file.
			if packet == nil {
				return
			}
			if *logAllPackets {
				log.Println(packet)
			}
			if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
				log.Println("Unusable packet")
				continue
			}
			tcp := packet.TransportLayer().(*layers.TCP)
			assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)

		case <-ticker:
			// Every minute, flush connections that haven't seen activity in the past 2 minutes.
			assembler.FlushOlderThan(time.Now().Add(time.Minute * -2))
		}
	}
}
<commit_msg>Prune back to try and trace the end of file issues.<commit_after>// Copyright 2012 Google, Inc. All rights reserved.
//
// Use of this source code is governed by a BSD-style license
// that can be found in the LICENSE file in the root of the source
// tree.

// This binary provides sample code for using the gopacket TCP assembler and TCP
// stream reader.  It reads packets off the wire and reconstructs HTTP requests
// it sees, logging them.
package main

import (
	"bufio"
	"flag"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net/http"
	"time"

	"github.com/google/gopacket"
	"github.com/google/gopacket/examples/util"
	"github.com/google/gopacket/layers"
	"github.com/google/gopacket/pcap"
	"github.com/google/gopacket/tcpassembly"
	"github.com/google/gopacket/tcpassembly/tcpreader"
)

var iface = flag.String("i", "eth0", "Interface to get packets from")
var fname = flag.String("r", "", "Filename to read from, overrides -i")
var snaplen = flag.Int("s", 1600, "SnapLen for pcap packet capture")
var filter = flag.String("f", "tcp and dst port 80", "BPF filter for pcap")
var logAllPackets = flag.Bool("v", false, "Logs every packet in great detail")

// Build a simple HTTP request parser using tcpassembly.StreamFactory and tcpassembly.Stream interfaces

// httpStreamFactory implements tcpassembly.StreamFactory
type httpStreamFactory struct{}

// httpStream will handle the actual decoding of http requests.
type httpStream struct {
	net, transport gopacket.Flow
	r              tcpreader.ReaderStream
}

func (h *httpStreamFactory) New(net, transport gopacket.Flow) tcpassembly.Stream {
	hstream := &httpStream{
		net:       net,
		transport: transport,
		r:         tcpreader.NewReaderStream(),
	}
	go hstream.run() // Important... we must guarantee that data from the reader stream is read.

	// ReaderStream implements tcpassembly.Stream, so we can return a pointer to it.
	return &hstream.r
}

func (h *httpStream) run() {
	buf := bufio.NewReader(&h.r)
	for {
		resp, err := http.ReadResponse(buf, nil)
		if err == io.EOF {
			// We must read until we see an EOF... very important!
			return
		} else if err == io.ErrUnexpectedEOF {
			// TODO: need to establish if we get these in the header.
			//log.Println("UEOF IN RESP:", h.net, ":",, err)
			return
		} else if err != nil {
			log.Println("ERROR IN RESP:", h.net, ":", err)
			// TODO: What else?
		} else {

			contentType := resp.Header["Content-Type"]
			//contentEnc := resp.Header["Content-Encoding"]

			//log.Println("ENCODING:", resp.TransferEncoding, ":", contentEnc, ":", resp.Uncompressed)

			if len(contentType) != 0 {

				reader := resp.Body
				/*
					if len(contentEnc) != 0 {
						if contentEnc[0] == "gzip" {
							r, qerr := gzip.NewReader(resp.Body)
							if qerr != nil {
								log.Println("ERROR GZIP:", qerr)
							}
							reader = r
						}
					}
				*/

				switch contentType[0] {
				// TODO: ASCII, ANSI (Windows-1252)
				case "text/html", "text/html; charset=utf-8", "text/html; charset=UTF-8":
					// Default charset for HTML5
					log.Print("MATCHED:", contentType[0])

					b, err := ioutil.ReadAll(reader)
					if err != nil {
						log.Println(err)
					}
					fmt.Println(string(b))
					/*
						body, perr := html.Parse(resp.Body)
						if perr != nil {
							log.Println("PARSE ERROR:", perr)
							break
						} else {
							doc := goquery.NewDocumentFromNode(body)
							fmt.Println("DOC:", doc.Find("h1").Text())
						}
					*/
				case "text/html; charset=iso-8859-1", "text/html; charset=ISO-8859-1":
					// Default charset for HTML 2 to 4
					// TODO: Do something with it, e.g. convert with iconv.
					log.Print("MATCHED:", contentType[0])
					fallthrough
				default:
					//log.Println("UNUSED TYPE:", contentTyp
				}
			}

			for {
				bytes, err := tcpreader.DiscardBytesToFirstError(resp.Body)

				if err == io.EOF || err == io.ErrUnexpectedEOF {
					break

				} else if err != nil && bytes == 0 {
					log.Println("ERROR BUT ZERO:", h.net, ":", err, ":", bytes)
					break

				} else if err != nil {
					log.Println("ERROR NOT ZERO:", h.net, ":", err, ":", bytes)
				}
			}
			resp.Body.Close()
			//log.Println(h.net, ":", contentType, resp.Status)
		}
	}
}

func main() {
	defer util.Run()()
	var handle *pcap.Handle
	var err error

	// Set up pcap packet capture
	if *fname != "" {
		log.Printf("Reading from pcap dump %q", *fname)
		handle, err = pcap.OpenOffline(*fname)
	} else {
		log.Printf("Starting capture on interface %q", *iface)
		// TODO: Not sure about BlockForever.
		handle, err = pcap.OpenLive(*iface, int32(*snaplen), true, pcap.BlockForever)
	}
	if err != nil {
		log.Fatal(err)
	}

	if err := handle.SetBPFFilter(*filter); err != nil {
		log.Fatal(err)
	}

	// Set up assembly
	streamFactory := &httpStreamFactory{}
	streamPool := tcpassembly.NewStreamPool(streamFactory)
	assembler := tcpassembly.NewAssembler(streamPool)

	log.Println("reading in packets")
	// Read in packets, pass to assembler.
	packetSource := gopacket.NewPacketSource(handle, handle.LinkType())
	packets := packetSource.Packets()
	ticker := time.Tick(time.Second * 10)
	for {
		select {
		case packet := <-packets:
			// A nil packet indicates the end of a pcap file.
			if packet == nil {
				return
			}
			if *logAllPackets {
				log.Println(packet)
			}
			if packet.NetworkLayer() == nil || packet.TransportLayer() == nil || packet.TransportLayer().LayerType() != layers.LayerTypeTCP {
				log.Println("Unusable packet")
				continue
			}
			tcp := packet.TransportLayer().(*layers.TCP)
			assembler.AssembleWithTimestamp(packet.NetworkLayer().NetworkFlow(), tcp, packet.Metadata().Timestamp)

		case <-ticker:
			// Was: Every minute, flush connections that haven't seen activity in the past 2 minutes.
			assembler.FlushOlderThan(time.Now().Add(time.Second * -20))
		}
	}
}
<|endoftext|>


<commit_before>package routing

import (
	"log"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/victorspringer/trapAdvisor/authenticating"
	"github.com/victorspringer/trapAdvisor/handling"
	"github.com/victorspringer/trapAdvisor/persistence"
)

type route struct {
	Method      string
	Pattern     string
	Name        string
	HandlerFunc http.HandlerFunc
}

// Router initializer.
func Router() *mux.Router {
	router := mux.NewRouter().StrictSlash(true)

	aSvc := authenticating.NewService()

	hSvc := handling.NewService(
		persistence.NewTravellerRepository(),
		persistence.NewFriendshipRepository(),
		persistence.NewTripRepository(),
		persistence.NewTouristAttractionRepository(),
	)

	routes := []route{
		route{"GET", "/health", "Health", hSvc.Health},
		route{"GET", "/login", "Login", aSvc.HandleFacebookLogin},
		route{"GET", "/auth_callback", "AuthCallback", aSvc.HandleFacebookCallback},
		route{"GET", "/logout", "Logout", aSvc.HandleFacebookLogout},

		route{"POST", "/v1/trip/store", "StoreTrip", aSvc.AuthMiddleware(hSvc.StoreTrip)},
		route{"POST", "/v1/ta/store", "StoreTouristAttraction", aSvc.AuthMiddleware(hSvc.StoreTouristAttraction)},

		route{"GET", "/v1/traveller/find/{id}", "FindTraveller", aSvc.AuthMiddleware(hSvc.FindTraveller)},
		route{"GET", "/v1/friendship/find/traveller/{id}", "FindFriendshipByTravellerID", aSvc.AuthMiddleware(hSvc.FindFriendshipByTravellerID)},
		route{"GET", "/v1/trip/find/{id}", "FindTrip", aSvc.AuthMiddleware(hSvc.FindTrip)},
		route{"GET", "/v1/trip/find/traveller/{id}", "FindTripByTravellerID", aSvc.AuthMiddleware(hSvc.FindTripByTravellerID)},
		route{"GET", "/v1/ta/find/{id}", "FindTouristAttraction", aSvc.AuthMiddleware(hSvc.FindTouristAttraction)},
		route{"GET", "/v1/ta/find/trip/{id}", "FindTouristAttractionByTripID", aSvc.AuthMiddleware(hSvc.FindTouristAttractionByTripID)},
		route{"GET", "/v1/ta/find/name_part/{namePart}", "FindTouristAttractionByNamePart", aSvc.AuthMiddleware(hSvc.FindTouristAttractionByNamePart)},
		route{"GET", "/v1/ta/most_visited", "FindMostVisitedTouristAttractions", aSvc.AuthMiddleware(hSvc.FindMostVisitedTouristAttractions)},
		route{"GET", "/v1/ta/best_rated", "FindBestRatedTouristAttractions", aSvc.AuthMiddleware(hSvc.FindBestRatedTouristAttractions)},
	}

	for _, route := range routes {
		var handler http.Handler

		handler = route.HandlerFunc
		handler = logger(handler, route.Name)

		router.Methods(route.Method).Path(route.Pattern).Name(route.Name).Handler(handler)
	}

	return router
}

func logger(inner http.Handler, name string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		inner.ServeHTTP(w, r)
		log.Printf("%v\t%v\t%v\t%v", r.Method, r.RequestURI, name, time.Since(start))
	})
}
<commit_msg>allows cors<commit_after>package routing

import (
	"log"
	"net/http"
	"time"

	"github.com/gorilla/mux"
	"github.com/rs/cors"
	"github.com/victorspringer/trapAdvisor/authenticating"
	"github.com/victorspringer/trapAdvisor/handling"
	"github.com/victorspringer/trapAdvisor/persistence"
)

type route struct {
	Method      string
	Pattern     string
	Name        string
	HandlerFunc http.HandlerFunc
}

// Router initializer.
func Router() *mux.Router {
	router := mux.NewRouter().StrictSlash(true)

	aSvc := authenticating.NewService()

	hSvc := handling.NewService(
		persistence.NewTravellerRepository(),
		persistence.NewFriendshipRepository(),
		persistence.NewTripRepository(),
		persistence.NewTouristAttractionRepository(),
	)

	routes := []route{
		route{"GET", "/health", "Health", hSvc.Health},
		route{"GET", "/login", "Login", aSvc.HandleFacebookLogin},
		route{"GET", "/auth_callback", "AuthCallback", aSvc.HandleFacebookCallback},
		route{"GET", "/logout", "Logout", aSvc.HandleFacebookLogout},

		route{"POST", "/v1/trip/store", "StoreTrip", aSvc.AuthMiddleware(hSvc.StoreTrip)},
		route{"POST", "/v1/ta/store", "StoreTouristAttraction", aSvc.AuthMiddleware(hSvc.StoreTouristAttraction)},

		route{"GET", "/v1/traveller/find/{id}", "FindTraveller", aSvc.AuthMiddleware(hSvc.FindTraveller)},
		route{"GET", "/v1/friendship/find/traveller/{id}", "FindFriendshipByTravellerID", aSvc.AuthMiddleware(hSvc.FindFriendshipByTravellerID)},
		route{"GET", "/v1/trip/find/{id}", "FindTrip", aSvc.AuthMiddleware(hSvc.FindTrip)},
		route{"GET", "/v1/trip/find/traveller/{id}", "FindTripByTravellerID", aSvc.AuthMiddleware(hSvc.FindTripByTravellerID)},
		route{"GET", "/v1/ta/find/{id}", "FindTouristAttraction", aSvc.AuthMiddleware(hSvc.FindTouristAttraction)},
		route{"GET", "/v1/ta/find/trip/{id}", "FindTouristAttractionByTripID", aSvc.AuthMiddleware(hSvc.FindTouristAttractionByTripID)},
		route{"GET", "/v1/ta/find/name_part/{namePart}", "FindTouristAttractionByNamePart", aSvc.AuthMiddleware(hSvc.FindTouristAttractionByNamePart)},
		route{"GET", "/v1/ta/most_visited", "FindMostVisitedTouristAttractions", aSvc.AuthMiddleware(hSvc.FindMostVisitedTouristAttractions)},
		route{"GET", "/v1/ta/best_rated", "FindBestRatedTouristAttractions", aSvc.AuthMiddleware(hSvc.FindBestRatedTouristAttractions)},
	}

	for _, route := range routes {
		var handler http.Handler

		handler = route.HandlerFunc
		handler = logger(handler, route.Name)
		handler = cors.Default().Handler(handler)

		router.Methods(route.Method).Path(route.Pattern).Name(route.Name).Handler(handler)
	}

	return router
}

func logger(inner http.Handler, name string) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		start := time.Now()
		inner.ServeHTTP(w, r)
		log.Printf("%v\t%v\t%v\t%v", r.Method, r.RequestURI, name, time.Since(start))
	})
}
<|endoftext|>