<commit_before>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/testutil"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func Test_Cluster_Subscribe_To_Random_Node(t *testing.T) {
	testutil.SkipIfShort(t)
	a := assert.New(t)

	node1 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: "localhost:8090",
		NodeID:     1,
		NodePort:   11000,
		Remotes:    "localhost:11000",
	})
	a.NotNil(node1)
	defer node1.cleanup(true)

	node2 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: "localhost:8091",
		NodeID:     2,
		NodePort:   11001,
		Remotes:    "localhost:11000",
	})
	a.NotNil(node2)
	defer node2.cleanup(true)

	client1, err := node1.client("user1", 10, true)
	a.NoError(err)

	err = client1.Subscribe("/foo/bar")
	a.NoError(err, "Subscribe to first node should work")

	client1.Close()

	time.Sleep(50 * time.Millisecond)

	client1, err = node2.client("user1", 10, true)
	a.NoError(err, "Connection to second node should return no error")

	err = client1.Subscribe("/foo/bar")
	a.NoError(err, "Subscribe to second node should work")
	client1.Close()
}

func Test_Cluster_Integration(t *testing.T) {
	testutil.SkipIfShort(t)
	defer testutil.ResetDefaultRegistryHealthCheck()

	a := assert.New(t)

	node1 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: ":8092",
		NodeID:     1,
		NodePort:   11002,
		Remotes:    "localhost:11002",
	})
	a.NotNil(node1)
	defer node1.cleanup(true)

	node2 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: ":8093",
		NodeID:     2,
		NodePort:   11002,
		Remotes:    "localhost:11003",
	})
	a.NotNil(node2)
	defer node2.cleanup(true)

	client1, err := node1.client("user1", 10, false)
	a.NoError(err)

	client2, err := node2.client("user2", 10, false)
	a.NoError(err)

	err = client2.Subscribe("/testTopic/m")
	a.NoError(err)

	client3, err := node1.client("user3", 10, false)
	a.NoError(err)

	numSent := 3
	for i := 0; i < numSent; i++ {
		err := client1.Send("/testTopic/m", "body", "{jsonHeader:1}")
		a.NoError(err)

		err = client3.Send("/testTopic/m", "body", "{jsonHeader:4}")
		a.NoError(err)
	}

	breakTimer := time.After(3 * time.Second)
	numReceived := 0
	idReceived := make(map[uint64]bool)

	// see if the correct number of messages arrived at the other client, before timeout is reached
WAIT:
	for {
		select {
		case incomingMessage := <-client2.Messages():
			numReceived++
			logger.WithFields(log.Fields{
				"nodeID":            incomingMessage.NodeID,
				"path":              incomingMessage.Path,
				"incomingMsgUserId": incomingMessage.UserID,
				"headerJson":        incomingMessage.HeaderJSON,
				"body":              incomingMessage.BodyAsString(),
				"numReceived":       numReceived,
			}).Info("Client2 received a message")

			a.Equal(protocol.Path("/testTopic/m"), incomingMessage.Path)
			a.Equal("body", incomingMessage.BodyAsString())
			a.True(incomingMessage.ID > 0)
			idReceived[incomingMessage.ID] = true

			if 2*numReceived == numSent {
				break WAIT
			}

		case <-breakTimer:
			break WAIT
		}
	}
}

var syncTopic = "/syncTopic"

// Test synchronizing messages when a new node is
func TestSynchronizerIntegration(t *testing.T) {
	testutil.SkipIfShort(t)
	defer testutil.EnableDebugForMethod()()

	a := assert.New(t)

	node1 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: "localhost:8094",
		NodeID:     1,
		NodePort:   11004,
		Remotes:    "localhost:11004",
	})
	a.NotNil(node1)
	defer node1.cleanup(true)

	time.Sleep(2 * time.Second)

	client1, err := node1.client("client1", 10, true)
	a.NoError(err)

	client1.Send(syncTopic, "nobody", "")
	client1.Send(syncTopic, "nobody", "")
	client1.Send(syncTopic, "nobody", "")

	time.Sleep(2 * time.Second)

	node2 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: ":8095",
		NodeID:     2,
		NodePort:   11005,
		Remotes:    "localhost:11004",
	})
	a.NotNil(node2)
	defer node2.cleanup(true)

	client2, err := node2.client("client2", 10, true)
	a.NoError(err)

	cmd := &protocol.Cmd{
		Name: protocol.CmdReceive,
		Arg:  syncTopic + " -3",
	}
	doneC := make(chan struct{})
	go func() {
		for {
			select {
			case m := <-client2.Messages():
				log.WithField("m", m).Error("Message received from first cluster")
			case e := <-client2.Errors():
				log.WithField("clientError", e).Error("Client error")
			case status := <-client2.StatusMessages():
				log.WithField("status", status).Error("Client status messasge")
			case <-doneC:
				return
			}
		}
	}()
	log.Error(string(cmd.Bytes()))
	client2.WriteRawMessage(cmd.Bytes())
	time.Sleep(10 * time.Second)
	close(doneC)
}
<commit_msg>fix ports in IT<commit_after>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/testutil"
	"github.com/stretchr/testify/assert"
	"testing"
	"time"
)

func Test_Cluster_Subscribe_To_Random_Node(t *testing.T) {
	testutil.SkipIfShort(t)
	a := assert.New(t)

	node1 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: "localhost:8090",
		NodeID:     1,
		NodePort:   11000,
		Remotes:    "localhost:11000",
	})
	a.NotNil(node1)
	defer node1.cleanup(true)

	node2 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: "localhost:8091",
		NodeID:     2,
		NodePort:   11001,
		Remotes:    "localhost:11000",
	})
	a.NotNil(node2)
	defer node2.cleanup(true)

	client1, err := node1.client("user1", 10, true)
	a.NoError(err)

	err = client1.Subscribe("/foo/bar")
	a.NoError(err, "Subscribe to first node should work")

	client1.Close()

	time.Sleep(50 * time.Millisecond)

	client1, err = node2.client("user1", 10, true)
	a.NoError(err, "Connection to second node should return no error")

	err = client1.Subscribe("/foo/bar")
	a.NoError(err, "Subscribe to second node should work")
	client1.Close()
}

func Test_Cluster_Integration(t *testing.T) {
	testutil.SkipIfShort(t)
	defer testutil.ResetDefaultRegistryHealthCheck()

	a := assert.New(t)

	node1 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: ":8092",
		NodeID:     1,
		NodePort:   11002,
		Remotes:    "localhost:11002",
	})
	a.NotNil(node1)
	defer node1.cleanup(true)

	node2 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: ":8093",
		NodeID:     2,
		NodePort:   11003,
		Remotes:    "localhost:11002",
	})
	a.NotNil(node2)
	defer node2.cleanup(true)

	client1, err := node1.client("user1", 10, false)
	a.NoError(err)

	client2, err := node2.client("user2", 10, false)
	a.NoError(err)

	err = client2.Subscribe("/testTopic/m")
	a.NoError(err)

	client3, err := node1.client("user3", 10, false)
	a.NoError(err)

	numSent := 3
	for i := 0; i < numSent; i++ {
		err := client1.Send("/testTopic/m", "body", "{jsonHeader:1}")
		a.NoError(err)

		err = client3.Send("/testTopic/m", "body", "{jsonHeader:4}")
		a.NoError(err)
	}

	breakTimer := time.After(3 * time.Second)
	numReceived := 0
	idReceived := make(map[uint64]bool)

	// see if the correct number of messages arrived at the other client, before timeout is reached
WAIT:
	for {
		select {
		case incomingMessage := <-client2.Messages():
			numReceived++
			logger.WithFields(log.Fields{
				"nodeID":            incomingMessage.NodeID,
				"path":              incomingMessage.Path,
				"incomingMsgUserId": incomingMessage.UserID,
				"headerJson":        incomingMessage.HeaderJSON,
				"body":              incomingMessage.BodyAsString(),
				"numReceived":       numReceived,
			}).Info("Client2 received a message")

			a.Equal(protocol.Path("/testTopic/m"), incomingMessage.Path)
			a.Equal("body", incomingMessage.BodyAsString())
			a.True(incomingMessage.ID > 0)
			idReceived[incomingMessage.ID] = true

			if 2*numReceived == numSent {
				break WAIT
			}

		case <-breakTimer:
			break WAIT
		}
	}
}

var syncTopic = "/syncTopic"

// Test synchronizing messages when a new node is
func TestSynchronizerIntegration(t *testing.T) {
	testutil.SkipIfShort(t)
	defer testutil.EnableDebugForMethod()()

	a := assert.New(t)

	node1 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: "localhost:8094",
		NodeID:     1,
		NodePort:   11004,
		Remotes:    "localhost:11004",
	})
	a.NotNil(node1)
	defer node1.cleanup(true)

	time.Sleep(2 * time.Second)

	client1, err := node1.client("client1", 10, true)
	a.NoError(err)

	client1.Send(syncTopic, "nobody", "")
	client1.Send(syncTopic, "nobody", "")
	client1.Send(syncTopic, "nobody", "")

	time.Sleep(2 * time.Second)

	node2 := newTestClusterNode(t, testClusterNodeConfig{
		HttpListen: ":8095",
		NodeID:     2,
		NodePort:   11005,
		Remotes:    "localhost:11004",
	})
	a.NotNil(node2)
	defer node2.cleanup(true)

	client2, err := node2.client("client2", 10, true)
	a.NoError(err)

	cmd := &protocol.Cmd{
		Name: protocol.CmdReceive,
		Arg:  syncTopic + " -3",
	}
	doneC := make(chan struct{})
	go func() {
		for {
			select {
			case m := <-client2.Messages():
				log.WithField("m", m).Error("Message received from first cluster")
			case e := <-client2.Errors():
				log.WithField("clientError", e).Error("Client error")
			case status := <-client2.StatusMessages():
				log.WithField("status", status).Error("Client status messasge")
			case <-doneC:
				return
			}
		}
	}()
	log.Error(string(cmd.Bytes()))
	client2.WriteRawMessage(cmd.Bytes())
	time.Sleep(10 * time.Second)
	close(doneC)
}
<|endoftext|>


<commit_before>package connector

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/router"
	"github.com/smancke/guble/testutil"
	"github.com/stretchr/testify/assert"
)

type connectorMocks struct {
	router  *MockRouter
	sender  *MockSender
	queue   *MockQueue
	manager *MockManager
	kvstore *MockKVStore
}

// Ensure the subscription is started when posting
func TestConnector_PostSubscription(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	mocks.manager.EXPECT().Load().Return(nil)
	mocks.manager.EXPECT().List().Return(make([]Subscriber, 0))
	err := conn.Start()
	a.NoError(err)
	defer conn.Stop()

	subscriber := NewMockSubscriber(testutil.MockCtrl)
	mocks.manager.EXPECT().Create(gomock.Eq(protocol.Path("/topic1")), gomock.Eq(router.RouteParams{
		"device_token": "device1",
		"user_id":      "user1",
	})).Return(subscriber, nil)

	subscriber.EXPECT().Loop(gomock.Any(), gomock.Any())
	r := router.NewRoute(router.RouteConfig{
		Path: protocol.Path("topic1"),
		RouteParams: router.RouteParams{
			"device_token": "device1",
			"user_id":      "user1",
		},
	})
	subscriber.EXPECT().Route().Return(r)
	mocks.router.EXPECT().Subscribe(gomock.Eq(r)).Return(r, nil)

	req, err := http.NewRequest(http.MethodPost, "/connector/device1/user1/topic1", strings.NewReader(""))
	a.NoError(err)
	conn.ServeHTTP(recorder, req)
	a.Equal(`{"subscribed":"/topic1"}`, recorder.Body.String())
	time.Sleep(100 * time.Millisecond)
}

func TestConnector_PostSubscriptionNoMocks(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, false, false)

	entriesC := make(chan [2]string)
	mocks.kvstore.EXPECT().Iterate(gomock.Eq("test"), gomock.Eq("")).Return(entriesC)
	close(entriesC)

	mocks.kvstore.EXPECT().Put(gomock.Eq("test"), gomock.Eq(GenerateKey("/topic1", map[string]string{
		"device_token": "device1",
		"user_id":      "user1",
	})), gomock.Any())

	mocks.router.EXPECT().Subscribe(gomock.Any())

	err := conn.Start()
	a.NoError(err)
	defer conn.Stop()

	req, err := http.NewRequest(http.MethodPost, "/connector/device1/user1/topic1", strings.NewReader(""))
	a.NoError(err)
	conn.ServeHTTP(recorder, req)
	a.Equal(`{"subscribed":"/topic1"}`, recorder.Body.String())
	time.Sleep(100 * time.Millisecond)
}

func TestConnector_DeleteSubscription(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	subscriber := NewMockSubscriber(testutil.MockCtrl)
	mocks.manager.EXPECT().Find(gomock.Eq(GenerateKey("topic1", map[string]string{
		"device_token": "device1",
		"user_id":      "user1",
	}))).Return(subscriber)
	mocks.manager.EXPECT().Remove(subscriber).Return(nil)

	req, err := http.NewRequest(http.MethodDelete, "/connector/device1/user1/topic1", strings.NewReader(""))
	a.NoError(err)
	conn.ServeHTTP(recorder, req)
	a.Equal(`{"unsubscribed":"topic1"}`, recorder.Body.String())
	time.Sleep(200 * time.Millisecond)
}

func TestConnector_GetList_And_Getters(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	subscriber1 := NewMockSubscriber(testutil.MockCtrl)
	subscriber1.EXPECT().Route().Return(router.NewRoute(router.RouteConfig{
		Path: "topic1",
	}))
	subscriber2 := NewMockSubscriber(testutil.MockCtrl)
	subscriber2.EXPECT().Route().Return(router.NewRoute(router.RouteConfig{
		Path: "topic2",
	}))
	mocks.manager.EXPECT().Filter(gomock.Any()).Return([]Subscriber{subscriber1, subscriber2})

	req, err := http.NewRequest(http.MethodGet, "/connector/", strings.NewReader(""))
	a.NoError(err)

	conn.ServeHTTP(recorder, req)
	expectedJSON := `["topic1","topic2"]`
	a.JSONEq(expectedJSON, recorder.Body.String())

	a.Equal("/connector/", conn.GetPrefix())
	a.Equal(mocks.manager, conn.Manager())
	a.Equal(nil, conn.ResponseHandler())
}

func TestConnector_GetListWithFilters(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	mocks.manager.EXPECT().Filter(gomock.Eq(map[string]string{
		"filter1": "value1",
		"filter2": "value2",
	})).Return([]Subscriber{})

	req, err := http.NewRequest(
		http.MethodGet,
		"/connector/?filter1=value1&filter2=value2",
		strings.NewReader(""))
	a.NoError(err)

	conn.ServeHTTP(recorder, req)
}

func TestConnector_StartWithSubscriptions(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, false, false)

	entriesC := make(chan [2]string)
	mocks.kvstore.EXPECT().Iterate(gomock.Eq("test"), gomock.Eq("")).Return(entriesC)
	close(entriesC)
	mocks.kvstore.EXPECT().Put(gomock.Any(), gomock.Any(), gomock.Any()).Times(4)

	err := conn.Start()
	a.NoError(err)

	routes := make([]*router.Route, 0, 4)
	mocks.router.EXPECT().Subscribe(gomock.Any()).Do(func(r *router.Route) (*router.Route, error) {
		routes = append(routes, r)
		return r, nil
	}).Times(4)

	// create subscriptions
	createSubscriptions(t, conn, 4)
	time.Sleep(100 * time.Millisecond)

	mocks.sender.EXPECT().Send(gomock.Any()).Return(nil, nil).Times(4)

	// send message in route channel
	for i, r := range routes {
		r.Deliver(&protocol.Message{
			ID:   uint64(i),
			Path: protocol.Path("/topic"),
			Body: []byte("test body"),
		})
	}

	time.Sleep(100 * time.Millisecond)

	err = conn.Stop()
	a.NoError(err)
}

func createSubscriptions(t *testing.T, conn Connector, count int) {
	a := assert.New(t)
	for i := 1; i <= count; i++ {
		recorder := httptest.NewRecorder()
		r, err := http.NewRequest(
			http.MethodPost,
			fmt.Sprintf("/connector/device%d/user%d/topic", i, i),
			strings.NewReader(""))
		a.NoError(err)
		conn.ServeHTTP(recorder, r)
		a.Equal(200, recorder.Code)
		a.Equal(`{"subscribed":"topic"}`, recorder.Body.String())
	}
}

func TestConnector_StartAndStopWithoutSubscribers(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()
	a := assert.New(t)

	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, true)
	mocks.manager.EXPECT().Load().Return(nil)
	mocks.manager.EXPECT().List().Return(nil)
	mocks.queue.EXPECT().Start().Return(nil)
	mocks.queue.EXPECT().Stop().Return(nil)

	err := conn.Start()
	a.NoError(err)

	err = conn.Stop()
	a.NoError(err)
}

func getTestConnector(t *testing.T, config Config, mockManager bool, mockQueue bool) (Connector, *connectorMocks) {
	a := assert.New(t)

	var (
		mManager *MockManager
		mQueue   *MockQueue
	)

	mKVS := NewMockKVStore(testutil.MockCtrl)
	mRouter := NewMockRouter(testutil.MockCtrl)
	mRouter.EXPECT().KVStore().Return(mKVS, nil).AnyTimes()
	mSender := NewMockSender(testutil.MockCtrl)

	conn, err := NewConnector(mRouter, mSender, config)
	a.NoError(err)

	if mockManager {
		mManager = NewMockManager(testutil.MockCtrl)
		conn.(*connector).manager = mManager
	}
	if mockQueue {
		mQueue = NewMockQueue(testutil.MockCtrl)
		conn.(*connector).queue = mQueue
	}

	return conn, &connectorMocks{
		mRouter,
		mSender,
		mQueue,
		mManager,
		mKVS,
	}
}
<commit_msg>Fix expected value topic with slash<commit_after>package connector

import (
	"fmt"
	"net/http"
	"net/http/httptest"
	"strings"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/router"
	"github.com/smancke/guble/testutil"
	"github.com/stretchr/testify/assert"
)

type connectorMocks struct {
	router  *MockRouter
	sender  *MockSender
	queue   *MockQueue
	manager *MockManager
	kvstore *MockKVStore
}

// Ensure the subscription is started when posting
func TestConnector_PostSubscription(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	mocks.manager.EXPECT().Load().Return(nil)
	mocks.manager.EXPECT().List().Return(make([]Subscriber, 0))
	err := conn.Start()
	a.NoError(err)
	defer conn.Stop()

	subscriber := NewMockSubscriber(testutil.MockCtrl)
	mocks.manager.EXPECT().Create(gomock.Eq(protocol.Path("/topic1")), gomock.Eq(router.RouteParams{
		"device_token": "device1",
		"user_id":      "user1",
	})).Return(subscriber, nil)

	subscriber.EXPECT().Loop(gomock.Any(), gomock.Any())
	r := router.NewRoute(router.RouteConfig{
		Path: protocol.Path("topic1"),
		RouteParams: router.RouteParams{
			"device_token": "device1",
			"user_id":      "user1",
		},
	})
	subscriber.EXPECT().Route().Return(r)
	mocks.router.EXPECT().Subscribe(gomock.Eq(r)).Return(r, nil)

	req, err := http.NewRequest(http.MethodPost, "/connector/device1/user1/topic1", strings.NewReader(""))
	a.NoError(err)
	conn.ServeHTTP(recorder, req)
	a.Equal(`{"subscribed":"/topic1"}`, recorder.Body.String())
	time.Sleep(100 * time.Millisecond)
}

func TestConnector_PostSubscriptionNoMocks(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, false, false)

	entriesC := make(chan [2]string)
	mocks.kvstore.EXPECT().Iterate(gomock.Eq("test"), gomock.Eq("")).Return(entriesC)
	close(entriesC)

	mocks.kvstore.EXPECT().Put(gomock.Eq("test"), gomock.Eq(GenerateKey("/topic1", map[string]string{
		"device_token": "device1",
		"user_id":      "user1",
	})), gomock.Any())

	mocks.router.EXPECT().Subscribe(gomock.Any())

	err := conn.Start()
	a.NoError(err)
	defer conn.Stop()

	req, err := http.NewRequest(http.MethodPost, "/connector/device1/user1/topic1", strings.NewReader(""))
	a.NoError(err)
	conn.ServeHTTP(recorder, req)
	a.Equal(`{"subscribed":"/topic1"}`, recorder.Body.String())
	time.Sleep(100 * time.Millisecond)
}

func TestConnector_DeleteSubscription(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	subscriber := NewMockSubscriber(testutil.MockCtrl)
	mocks.manager.EXPECT().Find(gomock.Eq(GenerateKey("topic1", map[string]string{
		"device_token": "device1",
		"user_id":      "user1",
	}))).Return(subscriber)
	mocks.manager.EXPECT().Remove(subscriber).Return(nil)

	req, err := http.NewRequest(http.MethodDelete, "/connector/device1/user1/topic1", strings.NewReader(""))
	a.NoError(err)
	conn.ServeHTTP(recorder, req)
	a.Equal(`{"unsubscribed":"topic1"}`, recorder.Body.String())
	time.Sleep(200 * time.Millisecond)
}

func TestConnector_GetList_And_Getters(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	subscriber1 := NewMockSubscriber(testutil.MockCtrl)
	subscriber1.EXPECT().Route().Return(router.NewRoute(router.RouteConfig{
		Path: "topic1",
	}))
	subscriber2 := NewMockSubscriber(testutil.MockCtrl)
	subscriber2.EXPECT().Route().Return(router.NewRoute(router.RouteConfig{
		Path: "topic2",
	}))
	mocks.manager.EXPECT().Filter(gomock.Any()).Return([]Subscriber{subscriber1, subscriber2})

	req, err := http.NewRequest(http.MethodGet, "/connector/", strings.NewReader(""))
	a.NoError(err)

	conn.ServeHTTP(recorder, req)
	expectedJSON := `["topic1","topic2"]`
	a.JSONEq(expectedJSON, recorder.Body.String())

	a.Equal("/connector/", conn.GetPrefix())
	a.Equal(mocks.manager, conn.Manager())
	a.Equal(nil, conn.ResponseHandler())
}

func TestConnector_GetListWithFilters(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)

	recorder := httptest.NewRecorder()
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, false)

	mocks.manager.EXPECT().Filter(gomock.Eq(map[string]string{
		"filter1": "value1",
		"filter2": "value2",
	})).Return([]Subscriber{})

	req, err := http.NewRequest(
		http.MethodGet,
		"/connector/?filter1=value1&filter2=value2",
		strings.NewReader(""))
	a.NoError(err)

	conn.ServeHTTP(recorder, req)
}

func TestConnector_StartWithSubscriptions(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()

	a := assert.New(t)
	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, false, false)

	entriesC := make(chan [2]string)
	mocks.kvstore.EXPECT().Iterate(gomock.Eq("test"), gomock.Eq("")).Return(entriesC)
	close(entriesC)
	mocks.kvstore.EXPECT().Put(gomock.Any(), gomock.Any(), gomock.Any()).Times(4)

	err := conn.Start()
	a.NoError(err)

	routes := make([]*router.Route, 0, 4)
	mocks.router.EXPECT().Subscribe(gomock.Any()).Do(func(r *router.Route) (*router.Route, error) {
		routes = append(routes, r)
		return r, nil
	}).Times(4)

	// create subscriptions
	createSubscriptions(t, conn, 4)
	time.Sleep(100 * time.Millisecond)

	mocks.sender.EXPECT().Send(gomock.Any()).Return(nil, nil).Times(4)

	// send message in route channel
	for i, r := range routes {
		r.Deliver(&protocol.Message{
			ID:   uint64(i),
			Path: protocol.Path("/topic"),
			Body: []byte("test body"),
		})
	}

	time.Sleep(100 * time.Millisecond)

	err = conn.Stop()
	a.NoError(err)
}

func createSubscriptions(t *testing.T, conn Connector, count int) {
	a := assert.New(t)
	for i := 1; i <= count; i++ {
		recorder := httptest.NewRecorder()
		r, err := http.NewRequest(
			http.MethodPost,
			fmt.Sprintf("/connector/device%d/user%d/topic", i, i),
			strings.NewReader(""))
		a.NoError(err)
		conn.ServeHTTP(recorder, r)
		a.Equal(200, recorder.Code)
		a.Equal(`{"subscribed":"/topic"}`, recorder.Body.String())
	}
}

func TestConnector_StartAndStopWithoutSubscribers(t *testing.T) {
	_, finish := testutil.NewMockCtrl(t)
	defer finish()
	a := assert.New(t)

	conn, mocks := getTestConnector(t, Config{
		Name:       "test",
		Schema:     "test",
		Prefix:     "/connector/",
		URLPattern: "/{device_token}/{user_id}/{topic:.*}",
	}, true, true)
	mocks.manager.EXPECT().Load().Return(nil)
	mocks.manager.EXPECT().List().Return(nil)
	mocks.queue.EXPECT().Start().Return(nil)
	mocks.queue.EXPECT().Stop().Return(nil)

	err := conn.Start()
	a.NoError(err)

	err = conn.Stop()
	a.NoError(err)
}

func getTestConnector(t *testing.T, config Config, mockManager bool, mockQueue bool) (Connector, *connectorMocks) {
	a := assert.New(t)

	var (
		mManager *MockManager
		mQueue   *MockQueue
	)

	mKVS := NewMockKVStore(testutil.MockCtrl)
	mRouter := NewMockRouter(testutil.MockCtrl)
	mRouter.EXPECT().KVStore().Return(mKVS, nil).AnyTimes()
	mSender := NewMockSender(testutil.MockCtrl)

	conn, err := NewConnector(mRouter, mSender, config)
	a.NoError(err)

	if mockManager {
		mManager = NewMockManager(testutil.MockCtrl)
		conn.(*connector).manager = mManager
	}
	if mockQueue {
		mQueue = NewMockQueue(testutil.MockCtrl)
		conn.(*connector).queue = mQueue
	}

	return conn, &connectorMocks{
		mRouter,
		mSender,
		mQueue,
		mManager,
		mKVS,
	}
}
<|endoftext|>


<commit_before>package tracer

import (
	"fmt"
	"log"
	"time"

	"github.com/tracer/tracer/pb"

	"github.com/golang/protobuf/ptypes"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// GRPC is a gRPC-based transport for sending spans to a server.
type GRPC struct {
	client        pb.StorerClient
	queue         []RawSpan
	ch            chan RawSpan
	flushInterval time.Duration

	stored  prometheus.Counter
	dropped prometheus.Counter
}

// GRPCOptions are options for the GRPC storer.
type GRPCOptions struct {
	// How many spans to queue before sending them to the server.
	// Additionally, a buffer the size of 2*QueueSize will be used to
	// process new spans. If this buffer runs full, new spans will be
	// dropped.
	QueueSize int
	// How often to flush spans, even if the queue isn't full yet.
	FlushInterval time.Duration
}

// NewGRPC returns a new Storer that sends spans via gRPC to a server.
func NewGRPC(address string, grpcOpts *GRPCOptions, opts ...grpc.DialOption) (Storer, error) {
	if grpcOpts == nil {
		grpcOpts = &GRPCOptions{1024, 1 * time.Second}
	}
	conn, err := grpc.Dial(address, opts...)
	if err != nil {
		return nil, err
	}
	client := pb.NewStorerClient(conn)
	g := &GRPC{
		client:        client,
		queue:         make([]RawSpan, 0, grpcOpts.QueueSize),
		ch:            make(chan RawSpan, grpcOpts.QueueSize*2),
		flushInterval: grpcOpts.FlushInterval,

		stored: prometheus.NewCounter(prometheus.CounterOpts{
			Name: "tracer_stored_spans_total",
			Help: "Number of stored spans",
		}),
		dropped: prometheus.NewCounter(prometheus.CounterOpts{
			Name: "tracer_dropped_spans_total",
			Help: "Number of dropped spans",
		}),
	}
	err = prometheus.Register(g.dropped)
	if err != nil {
		log.Println("couldn't register prometheus counter:", err)
	}
	err = prometheus.Register(g.stored)
	if err != nil {
		log.Println("couldn't register prometheus counter:", err)
	}
	go g.loop()
	return g, nil
}

func (g *GRPC) loop() {
	t := time.NewTicker(g.flushInterval)
	for {
		select {
		case sp := <-g.ch:
			g.queue = append(g.queue, sp)
			if len(g.queue) == cap(g.queue) {
				g.flush()
			}
		case <-t.C:
			g.flush()
		}
	}
}

func (g *GRPC) flush() {
	var pbs []*pb.Span
	for _, sp := range g.queue {
		pst, err := ptypes.TimestampProto(sp.StartTime)
		if err != nil {
			return // XXX
		}
		pft, err := ptypes.TimestampProto(sp.FinishTime)
		if err != nil {
			return // XXX
		}
		var tags []*pb.Tag
		for k, v := range sp.Tags {
			vs := fmt.Sprintf("%v", v) // XXX
			tags = append(tags, &pb.Tag{
				Key:   k,
				Value: vs,
			})
		}
		for _, log := range sp.Logs {
			t, err := ptypes.TimestampProto(log.Timestamp)
			if err != nil {
				return // XXX
			}
			ps := fmt.Sprintf("%v", log.Payload) // XXX
			tags = append(tags, &pb.Tag{
				Key:   log.Event,
				Value: ps,
				Time:  t,
			})
		}
		psp := &pb.Span{
			SpanId:        sp.SpanID,
			ParentId:      sp.ParentID,
			TraceId:       sp.TraceID,
			ServiceName:   sp.ServiceName,
			OperationName: sp.OperationName,
			StartTime:     pst,
			FinishTime:    pft,
			Flags:         sp.Flags,
			Tags:          tags,
		}
		pbs = append(pbs, psp)
	}
	if _, err := g.client.Store(context.Background(), &pb.StoreRequest{Spans: pbs}); err != nil {
		return // XXX
	}
	g.queue = g.queue[0:0]
}

// Store implements the tracer.Storer interface.
func (g *GRPC) Store(sp RawSpan) error {
	select {
	case g.ch <- sp:
		g.stored.Inc()
	default:
		g.dropped.Inc()
	}
	return nil
}
<commit_msg>Skip broken spans/logs and log errors<commit_after>package tracer

import (
	"fmt"
	"log"
	"time"

	"github.com/tracer/tracer/pb"

	"github.com/golang/protobuf/ptypes"
	"github.com/prometheus/client_golang/prometheus"
	"golang.org/x/net/context"
	"google.golang.org/grpc"
)

// GRPC is a gRPC-based transport for sending spans to a server.
type GRPC struct {
	client        pb.StorerClient
	queue         []RawSpan
	ch            chan RawSpan
	flushInterval time.Duration

	stored  prometheus.Counter
	dropped prometheus.Counter
}

// GRPCOptions are options for the GRPC storer.
type GRPCOptions struct {
	// How many spans to queue before sending them to the server.
	// Additionally, a buffer the size of 2*QueueSize will be used to
	// process new spans. If this buffer runs full, new spans will be
	// dropped.
	QueueSize int
	// How often to flush spans, even if the queue isn't full yet.
	FlushInterval time.Duration
}

// NewGRPC returns a new Storer that sends spans via gRPC to a server.
func NewGRPC(address string, grpcOpts *GRPCOptions, opts ...grpc.DialOption) (Storer, error) {
	if grpcOpts == nil {
		grpcOpts = &GRPCOptions{1024, 1 * time.Second}
	}
	conn, err := grpc.Dial(address, opts...)
	if err != nil {
		return nil, err
	}
	client := pb.NewStorerClient(conn)
	g := &GRPC{
		client:        client,
		queue:         make([]RawSpan, 0, grpcOpts.QueueSize),
		ch:            make(chan RawSpan, grpcOpts.QueueSize*2),
		flushInterval: grpcOpts.FlushInterval,

		stored: prometheus.NewCounter(prometheus.CounterOpts{
			Name: "tracer_stored_spans_total",
			Help: "Number of stored spans",
		}),
		dropped: prometheus.NewCounter(prometheus.CounterOpts{
			Name: "tracer_dropped_spans_total",
			Help: "Number of dropped spans",
		}),
	}
	err = prometheus.Register(g.dropped)
	if err != nil {
		log.Println("couldn't register prometheus counter:", err)
	}
	err = prometheus.Register(g.stored)
	if err != nil {
		log.Println("couldn't register prometheus counter:", err)
	}
	go g.loop()
	return g, nil
}

func (g *GRPC) loop() {
	t := time.NewTicker(g.flushInterval)
	for {
		select {
		case sp := <-g.ch:
			g.queue = append(g.queue, sp)
			if len(g.queue) == cap(g.queue) {
				if err := g.flush(); err != nil {
					log.Println("couldn't flush spans:", err)
				}
			}
		case <-t.C:
			if err := g.flush(); err != nil {
				log.Println("couldn't flush spans:", err)
			}
		}
	}
}

func (g *GRPC) flush() error {
	var pbs []*pb.Span
	for _, sp := range g.queue {
		pst, err := ptypes.TimestampProto(sp.StartTime)
		if err != nil {
			log.Println("dropping span because of error:", err)
			continue
		}
		pft, err := ptypes.TimestampProto(sp.FinishTime)
		if err != nil {
			log.Println("dropping span because of error:", err)
			continue
		}
		var tags []*pb.Tag
		for k, v := range sp.Tags {
			vs := fmt.Sprintf("%v", v) // XXX
			tags = append(tags, &pb.Tag{
				Key:   k,
				Value: vs,
			})
		}
		for _, l := range sp.Logs {
			t, err := ptypes.TimestampProto(l.Timestamp)
			if err != nil {
				log.Println("dropping log entry because of error:", err)
				continue
			}
			ps := fmt.Sprintf("%v", l.Payload) // XXX
			tags = append(tags, &pb.Tag{
				Key:   l.Event,
				Value: ps,
				Time:  t,
			})
		}
		psp := &pb.Span{
			SpanId:        sp.SpanID,
			ParentId:      sp.ParentID,
			TraceId:       sp.TraceID,
			ServiceName:   sp.ServiceName,
			OperationName: sp.OperationName,
			StartTime:     pst,
			FinishTime:    pft,
			Flags:         sp.Flags,
			Tags:          tags,
		}
		pbs = append(pbs, psp)
	}
	g.queue = g.queue[0:0]
	if _, err := g.client.Store(context.Background(), &pb.StoreRequest{Spans: pbs}); err != nil {
		return err
	}
	return nil
}

// Store implements the tracer.Storer interface.
func (g *GRPC) Store(sp RawSpan) error {
	select {
	case g.ch <- sp:
		g.stored.Inc()
	default:
		g.dropped.Inc()
	}
	return nil
}
<|endoftext|>


<commit_before>// Copyright 2013 Martini Authors
// Copyright 2015 The Macaron 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 gzip

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"strings"

	gz "github.com/klauspost/compress/gzip"
	"gopkg.in/macaron.v1"
)

const (
	_HEADER_ACCEPT_ENCODING  = "Accept-Encoding"
	_HEADER_CONTENT_ENCODING = "Content-Encoding"
	_HEADER_CONTENT_LENGTH   = "Content-Length"
	_HEADER_CONTENT_TYPE     = "Content-Type"
	_HEADER_VARY             = "Vary"
)

// Options represents a struct for specifying configuration options for the GZip middleware.
type Options struct {
	// Compression level. Can be DefaultCompression(-1), ConstantCompression(-2)
	// or any integer value between BestSpeed(1) and BestCompression(9) inclusive.
	CompressionLevel int
}

func isCompressionLevelValid(level int) bool {
	return level == gz.DefaultCompression ||
		level == gz.ConstantCompression ||
		(level >= gz.BestSpeed && level <= gz.BestCompression)
}

func prepareOptions(options []Options) Options {
	var opt Options
	if len(options) > 0 {
		opt = options[0]
	}

	if !isCompressionLevelValid(opt.CompressionLevel) {
		// For web content, level 4 seems to be a sweet spot.
		opt.CompressionLevel = 4
	}
	return opt
}

// Gziper returns a Handler that adds gzip compression to all requests.
// Make sure to include the Gzip middleware above other middleware
// that alter the response body (like the render middleware).
func Gziper(options ...Options) macaron.Handler {
	opt := prepareOptions(options)

	return func(ctx *macaron.Context) {
		if !strings.Contains(ctx.Req.Header.Get(_HEADER_ACCEPT_ENCODING), "gzip") {
			return
		}

		headers := ctx.Resp.Header()
		headers.Set(_HEADER_CONTENT_ENCODING, "gzip")
		headers.Set(_HEADER_VARY, _HEADER_ACCEPT_ENCODING)

		// We've made sure compression level is valid in prepareGzipOptions,
		// no need to check same error again.
		gz, err := gz.NewWriterLevel(ctx.Resp, opt.CompressionLevel)
		if err != nil {
			panic(err.Error())
		}
		defer gz.Close()

		gzw := gzipResponseWriter{gz, ctx.Resp}
		ctx.Resp = gzw
		ctx.MapTo(gzw, (*http.ResponseWriter)(nil))
		if ctx.Render != nil {
			ctx.Render.SetResponseWriter(gzw)
		}

		ctx.Next()

		// delete content length after we know we have been written to
		gzw.Header().Del("Content-Length")
	}
}

type gzipResponseWriter struct {
	w *gz.Writer
	macaron.ResponseWriter
}

func (grw gzipResponseWriter) Write(p []byte) (int, error) {
	if len(grw.Header().Get(_HEADER_CONTENT_TYPE)) == 0 {
		grw.Header().Set(_HEADER_CONTENT_TYPE, http.DetectContentType(p))
	}
	return grw.w.Write(p)
}

func (grw gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := grw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}
<commit_msg>fix import<commit_after>// Copyright 2013 Martini Authors
// Copyright 2015 The Macaron 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 gzip

import (
	"bufio"
	"fmt"
	"net"
	"net/http"
	"strings"

	"github.com/klauspost/compress/gzip"
	"gopkg.in/macaron.v1"
)

const (
	_HEADER_ACCEPT_ENCODING  = "Accept-Encoding"
	_HEADER_CONTENT_ENCODING = "Content-Encoding"
	_HEADER_CONTENT_LENGTH   = "Content-Length"
	_HEADER_CONTENT_TYPE     = "Content-Type"
	_HEADER_VARY             = "Vary"
)

// Options represents a struct for specifying configuration options for the GZip middleware.
type Options struct {
	// Compression level. Can be DefaultCompression(-1), ConstantCompression(-2)
	// or any integer value between BestSpeed(1) and BestCompression(9) inclusive.
	CompressionLevel int
}

func isCompressionLevelValid(level int) bool {
	return level == gzip.DefaultCompression ||
		level == gzip.ConstantCompression ||
		(level >= gzip.BestSpeed && level <= gzip.BestCompression)
}

func prepareOptions(options []Options) Options {
	var opt Options
	if len(options) > 0 {
		opt = options[0]
	}

	if !isCompressionLevelValid(opt.CompressionLevel) {
		// For web content, level 4 seems to be a sweet spot.
		opt.CompressionLevel = 4
	}
	return opt
}

// Gziper returns a Handler that adds gzip compression to all requests.
// Make sure to include the Gzip middleware above other middleware
// that alter the response body (like the render middleware).
func Gziper(options ...Options) macaron.Handler {
	opt := prepareOptions(options)

	return func(ctx *macaron.Context) {
		if !strings.Contains(ctx.Req.Header.Get(_HEADER_ACCEPT_ENCODING), "gzip") {
			return
		}

		headers := ctx.Resp.Header()
		headers.Set(_HEADER_CONTENT_ENCODING, "gzip")
		headers.Set(_HEADER_VARY, _HEADER_ACCEPT_ENCODING)

		// We've made sure compression level is valid in prepareGzipOptions,
		// no need to check same error again.
		gz, err := gzip.NewWriterLevel(ctx.Resp, opt.CompressionLevel)
		if err != nil {
			panic(err.Error())
		}
		defer gz.Close()

		gzw := gzipResponseWriter{gz, ctx.Resp}
		ctx.Resp = gzw
		ctx.MapTo(gzw, (*http.ResponseWriter)(nil))
		if ctx.Render != nil {
			ctx.Render.SetResponseWriter(gzw)
		}

		ctx.Next()

		// delete content length after we know we have been written to
		gzw.Header().Del("Content-Length")
	}
}

type gzipResponseWriter struct {
	w *gzip.Writer
	macaron.ResponseWriter
}

func (grw gzipResponseWriter) Write(p []byte) (int, error) {
	if len(grw.Header().Get(_HEADER_CONTENT_TYPE)) == 0 {
		grw.Header().Set(_HEADER_CONTENT_TYPE, http.DetectContentType(p))
	}
	return grw.w.Write(p)
}

func (grw gzipResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	hijacker, ok := grw.ResponseWriter.(http.Hijacker)
	if !ok {
		return nil, nil, fmt.Errorf("the ResponseWriter doesn't support the Hijacker interface")
	}
	return hijacker.Hijack()
}
<|endoftext|>


<commit_before>package hawk

import (
	"crypto/sha256"
	"errors"
	"net/http"
	"time"

	"github.com/dchest/uniuri"
	"github.com/gin-gonic/gin"
	hawk "github.com/tent/hawk-go"
)

const (
	AuthKey = "hawk_auth"
	UserKey = "hawk_user"
)

// ErrNotFound is set in context.Err if the GetCredentialFunc
// returns nil
var ErrNotFound = errors.New("Credentials not found")

// Credentials is used to store a key string and a User object.
// It is returned by a function of type GetCredentialFunc.
type Credentials struct {
	Key  string
	User interface{}
}

// GetCredentialFunc is a function that returns a *Credentials by id.
// If nothing is found the result should be nil and it's
// an authentication error (set in context).
// If an error occured (an external problem like db connection),
// return the error and it will be set as the context error.
type GetCredentialFunc func(id string) (*Credentials, error)

// SetNonceFunc is a function that returns false if nonce with the same
// associated id and time already exists. Otherwise true is returned
// an the nonce should be save to avoid replay problems.
type SetNonceFunc func(id string, nonce string, t time.Time) (bool, error)

type AbortHandlerFunc func(*gin.Context, error)

// Middleware is the middleware object.
// GetCredentials is the GetCredentialFunc
// SetNonce is the SetNonceFunc
// UserParam if set will set the user in the context with a matching key
// Ext add an "ext" header in the request
type Middleware struct {
	GetCredentials GetCredentialFunc
	SetNonce       SetNonceFunc
	AbortHandler   AbortHandlerFunc
	UserParam      string
	Ext            string
}

// NewMiddleware creates a new Middleware with the GetCredentials
// and SetNonce params set. UserParam is set to "user" by default.
func NewMiddleware(gcf GetCredentialFunc, snf SetNonceFunc) *Middleware {
	return &Middleware{
		GetCredentials: gcf,
		SetNonce:       snf,
	}
}

func ISHawkError(err error) bool {
	switch err {
	case ErrNotFound,
		hawk.ErrBewitExpired,
		hawk.ErrInvalidBewitMethod,
		hawk.ErrInvalidMAC,
		hawk.ErrMissingServerAuth,
		hawk.ErrNoAuth,
		hawk.ErrReplay,
		hawk.ErrTimestampSkew:
		return true
	}
	return false
}

// Abortequest aborts the request and set the context error and status.
// When possible it will attempt to send a "Server-Authorization" header.
func (hm *Middleware) Abortequest(c *gin.Context, err error, auth *hawk.Auth) {
	isHawk := ISHawkError(err)
	if isHawk && auth != nil {
		c.Header("Server-Authorization", auth.ResponseHeader(hm.Ext))
	}
	if hm.AbortHandler != nil {
		hm.AbortHandler(c, err)
		c.Abort()
	} else if isHawk {
		c.AbortWithError(http.StatusUnauthorized, err)
	} else {
		c.AbortWithError(http.StatusInternalServerError, err)
	}
}

// Filter is the middleware function that validate the hawk authentication.
func (hm *Middleware) Filter(c *gin.Context) {
	res := &Request{
		Hawk: hm,
	}

	auth, err := hawk.NewAuthFromRequest(c.Request, res.CredentialsLookup, res.NonceCheck)
	if res.Error != nil {
		hm.Abortequest(c, res.Error, nil)
	} else if err != nil {
		hm.Abortequest(c, err, auth)
	} else if err := auth.Valid(); err != nil {
		hm.Abortequest(c, err, auth)
	} else {
		c.Header("Server-Authorization", auth.ResponseHeader(hm.Ext))
		c.Set(AuthKey, auth)
		c.Set(UserKey, res.User)
		c.Next()
	}
}

// Request represent the state of a request.
type Request struct {
	Hawk  *Middleware
	ID    string
	User  interface{}
	Ok    bool
	Error error
}

// CredentialsLookup lookup the credantial for hawk-go from the user
// provided GetCredentialFunc.
func (hr *Request) CredentialsLookup(creds *hawk.Credentials) error {

	id := creds.ID
	if res, err := hr.Hawk.GetCredentials(id); err != nil {
		hr.Error = err
		return err
	} else if res == nil {
		return ErrNotFound
	} else {
		creds.Key = res.Key
		hr.User = res.User
		creds.Hash = sha256.New
		hr.Ok = true
		return nil
	}
}

// NonceCheck call the SetNonceFunc on behalf of hawk-go.
func (hr *Request) NonceCheck(nonce string, t time.Time, creds *hawk.Credentials) bool {
	if hr.Error != nil || !hr.Ok || hr.Hawk.SetNonce == nil {
		return false
	}

	ok, err := hr.Hawk.SetNonce(creds.ID, nonce, t)
	if err != nil {
		hr.Error = err
		return false
	}
	return ok
}

// GenIDKey generates a random id and key.
func GenIDKey() (string, string) {
	return uniuri.NewLen(12), uniuri.NewLen(24)
}
<commit_msg>add convenience methods GetUser and GetAuth<commit_after>package hawk

import (
	"crypto/sha256"
	"errors"
	"net/http"
	"time"

	"github.com/dchest/uniuri"
	"github.com/gin-gonic/gin"
	hawk "github.com/tent/hawk-go"
)

const (
	AuthKey = "hawk_auth"
	UserKey = "hawk_user"
)

// ErrNotFound is set in context.Err if the GetCredentialFunc
// returns nil
var ErrNotFound = errors.New("Credentials not found")

// Credentials is used to store a key string and a User object.
// It is returned by a function of type GetCredentialFunc.
type Credentials struct {
	Key  string
	User interface{}
}

// GetCredentialFunc is a function that returns a *Credentials by id.
// If nothing is found the result should be nil and it's
// an authentication error (set in context).
// If an error occured (an external problem like db connection),
// return the error and it will be set as the context error.
type GetCredentialFunc func(id string) (*Credentials, error)

// SetNonceFunc is a function that returns false if nonce with the same
// associated id and time already exists. Otherwise true is returned
// an the nonce should be save to avoid replay problems.
type SetNonceFunc func(id string, nonce string, t time.Time) (bool, error)

type AbortHandlerFunc func(*gin.Context, error)

// Middleware is the middleware object.
// GetCredentials is the GetCredentialFunc
// SetNonce is the SetNonceFunc
// UserParam if set will set the user in the context with a matching key
// Ext add an "ext" header in the request
type Middleware struct {
	GetCredentials GetCredentialFunc
	SetNonce       SetNonceFunc
	AbortHandler   AbortHandlerFunc
	UserParam      string
	Ext            string
}

// NewMiddleware creates a new Middleware with the GetCredentials
// and SetNonce params set. UserParam is set to "user" by default.
func NewMiddleware(gcf GetCredentialFunc, snf SetNonceFunc) *Middleware {
	return &Middleware{
		GetCredentials: gcf,
		SetNonce:       snf,
	}
}

func ISHawkError(err error) bool {
	switch err {
	case ErrNotFound,
		hawk.ErrBewitExpired,
		hawk.ErrInvalidBewitMethod,
		hawk.ErrInvalidMAC,
		hawk.ErrMissingServerAuth,
		hawk.ErrNoAuth,
		hawk.ErrReplay,
		hawk.ErrTimestampSkew:
		return true
	}
	return false
}

// Abortequest aborts the request and set the context error and status.
// When possible it will attempt to send a "Server-Authorization" header.
func (hm *Middleware) Abortequest(c *gin.Context, err error, auth *hawk.Auth) {
	isHawk := ISHawkError(err)
	if isHawk && auth != nil {
		c.Header("Server-Authorization", auth.ResponseHeader(hm.Ext))
	}
	if hm.AbortHandler != nil {
		hm.AbortHandler(c, err)
		c.Abort()
	} else if isHawk {
		c.AbortWithError(http.StatusUnauthorized, err)
	} else {
		c.AbortWithError(http.StatusInternalServerError, err)
	}
}

// Filter is the middleware function that validate the hawk authentication.
func (hm *Middleware) Filter(c *gin.Context) {
	res := &Request{
		Hawk: hm,
	}

	auth, err := hawk.NewAuthFromRequest(c.Request, res.CredentialsLookup, res.NonceCheck)
	if res.Error != nil {
		hm.Abortequest(c, res.Error, nil)
	} else if err != nil {
		hm.Abortequest(c, err, auth)
	} else if err := auth.Valid(); err != nil {
		hm.Abortequest(c, err, auth)
	} else {
		c.Header("Server-Authorization", auth.ResponseHeader(hm.Ext))
		c.Set(AuthKey, auth)
		c.Set(UserKey, res.User)
		c.Next()
	}
}

// Request represent the state of a request.
type Request struct {
	Hawk  *Middleware
	ID    string
	User  interface{}
	Ok    bool
	Error error
}

// CredentialsLookup lookup the credantial for hawk-go from the user
// provided GetCredentialFunc.
func (hr *Request) CredentialsLookup(creds *hawk.Credentials) error {

	id := creds.ID
	if res, err := hr.Hawk.GetCredentials(id); err != nil {
		hr.Error = err
		return err
	} else if res == nil {
		return ErrNotFound
	} else {
		creds.Key = res.Key
		hr.User = res.User
		creds.Hash = sha256.New
		hr.Ok = true
		return nil
	}
}

// NonceCheck call the SetNonceFunc on behalf of hawk-go.
func (hr *Request) NonceCheck(nonce string, t time.Time, creds *hawk.Credentials) bool {
	if hr.Error != nil || !hr.Ok || hr.Hawk.SetNonce == nil {
		return false
	}

	ok, err := hr.Hawk.SetNonce(creds.ID, nonce, t)
	if err != nil {
		hr.Error = err
		return false
	}
	return ok
}

// GenIDKey generates a random id and key.
func GenIDKey() (string, string) {
	return uniuri.NewLen(12), uniuri.NewLen(24)
}

// GetAuth returns the *hawk.Auth from the context.
// Will panic if not set (i.e. when the filter fail or has not happend yet)
func GetAuth(c *gin.Context) *hawk.Auth {
	return c.MustGet(AuthKey).(*hawk.Auth)
}

// GetHawk returns the user object (obtain with the GetCredentialFunc) from
// the context. Will panic if not set (i.e. when the filter fail or
// has not happend yet)
func GetUser(c *gin.Context) interface{} {
	return c.MustGet(UserKey)
}
<|endoftext|>


<commit_before>package cli

import (
	"fmt"
	"io"
	"strings"
	"text/tabwriter"
	"text/template"
)

// The text template for the Default help topic.
// cli.go uses text/template to render templates. You can
// render custom help text by setting this variable.
var AppHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   {{.Name}} {{if .Flags}}[global options] {{end}}command{{if .Flags}} [command options]{{end}} [arguments...]{{if len .Version}}

VERSION:
   {{.Version}}{{end}}{{if len .Authors}}

AUTHOR(S):
   {{range .Authors}}{{ . }}{{end}}{{end}}{{if .Commands}}

COMMANDS:
   {{range .Commands}}{{join .Names ", "}}{{ "\t" }}{{.Usage}}
   {{end}}{{end}}{{if .Flags}}

GLOBAL OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}{{if .Copyright }}

COPYRIGHT:
   {{.Copyright}}{{end}}
`

// The text template for the command help topic.
// cli.go uses text/template to render templates. You can
// render custom help text by setting this variable.
var CommandHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   command {{.Name}}{{if .Flags}} [command options]{{end}} [arguments...]{{if .Description}}

DESCRIPTION:
   {{.Description}}{{end}}{{if .Flags}}

OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{ end }}
`

// The text template for the subcommand help topic.
// cli.go uses text/template to render templates. You can
// render custom help text by setting this variable.
var SubcommandHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   {{.Name}} command{{if .Flags}} [command options]{{end}} [arguments...]

COMMANDS:
   {{range .Commands}}{{join .Names ", "}}{{ "\t" }}{{.Usage}}
   {{end}}{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`

var helpCommand = Command{
	Name:    "help",
	Aliases: []string{"h"},
	Usage:   "Shows a list of commands or help for one command",
	Action: func(c *Context) {
		args := c.Args()
		if args.Present() {
			ShowCommandHelp(c, args.First())
		} else {
			ShowAppHelp(c)
		}
	},
}

var helpSubcommand = Command{
	Name:    "help",
	Aliases: []string{"h"},
	Usage:   "Shows a list of commands or help for one command",
	Action: func(c *Context) {
		args := c.Args()
		if args.Present() {
			ShowCommandHelp(c, args.First())
		} else {
			ShowSubcommandHelp(c)
		}
	},
}

// Prints help for the App or Command
type helpPrinter func(w io.Writer, templ string, data interface{})

var HelpPrinter helpPrinter = printHelp

// Prints version for the App
var VersionPrinter = printVersion

func ShowAppHelp(c *Context) {
	HelpPrinter(c.App.Writer, AppHelpTemplate, c.App)
}

// Prints the list of subcommands as the default app completion method
func DefaultAppComplete(c *Context) {
	for _, command := range c.App.Commands {
		for _, name := range command.Names() {
			fmt.Fprintln(c.App.Writer, name)
		}
	}
}

// Prints help for the given command
func ShowCommandHelp(ctx *Context, command string) {
	// show the subcommand help for a command with subcommands
	if command == "" {
		HelpPrinter(ctx.App.Writer, SubcommandHelpTemplate, ctx.App)
		return
	}

	for _, c := range ctx.App.Commands {
		if c.HasName(command) {
			HelpPrinter(ctx.App.Writer, CommandHelpTemplate, c)
			return
		}
	}

	if ctx.App.CommandNotFound != nil {
		ctx.App.CommandNotFound(ctx, command)
	} else {
		fmt.Fprintf(ctx.App.Writer, "No help topic for '%v'\n", command)
	}
}

// Prints help for the given subcommand
func ShowSubcommandHelp(c *Context) {
	ShowCommandHelp(c, c.Command.Name)
}

// Prints the version number of the App
func ShowVersion(c *Context) {
	VersionPrinter(c)
}

func printVersion(c *Context) {
	fmt.Fprintf(c.App.Writer, "%v version %v\n", c.App.Name, c.App.Version)
}

// Prints the lists of commands within a given context
func ShowCompletions(c *Context) {
	a := c.App
	if a != nil && a.BashComplete != nil {
		a.BashComplete(c)
	}
}

// Prints the custom completions for a given command
func ShowCommandCompletions(ctx *Context, command string) {
	c := ctx.App.Command(command)
	if c != nil && c.BashComplete != nil {
		c.BashComplete(ctx)
	}
}

func printHelp(out io.Writer, templ string, data interface{}) {
	funcMap := template.FuncMap{
		"join": strings.Join,
	}

	w := tabwriter.NewWriter(out, 0, 8, 1, '\t', 0)
	t := template.Must(template.New("help").Funcs(funcMap).Parse(templ))
	err := t.Execute(w, data)
	if err != nil {
		panic(err)
	}
	w.Flush()
}

func checkVersion(c *Context) bool {
	if c.GlobalBool("version") || c.GlobalBool("v") || c.Bool("version") || c.Bool("v") {
		ShowVersion(c)
		return true
	}

	return false
}

func checkHelp(c *Context) bool {
	if c.GlobalBool("h") || c.GlobalBool("help") || c.Bool("h") || c.Bool("help") {
		ShowAppHelp(c)
		return true
	}

	return false
}

func checkCommandHelp(c *Context, name string) bool {
	if c.Bool("h") || c.Bool("help") {
		ShowCommandHelp(c, name)
		return true
	}

	return false
}

func checkSubcommandHelp(c *Context) bool {
	if c.GlobalBool("h") || c.GlobalBool("help") {
		ShowSubcommandHelp(c)
		return true
	}

	return false
}

func checkCompletions(c *Context) bool {
	if (c.GlobalBool(BashCompletionFlag.Name) || c.Bool(BashCompletionFlag.Name)) && c.App.EnableBashCompletion {
		ShowCompletions(c)
		return true
	}

	return false
}

func checkCommandCompletions(c *Context, name string) bool {
	if c.Bool(BashCompletionFlag.Name) && c.App.EnableBashCompletion {
		ShowCommandCompletions(c, name)
		return true
	}

	return false
}
<commit_msg>Fixing more formatting<commit_after>package cli

import (
	"fmt"
	"io"
	"strings"
	"text/tabwriter"
	"text/template"
)

// The text template for the Default help topic.
// cli.go uses text/template to render templates. You can
// render custom help text by setting this variable.
var AppHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   {{.Name}} {{if .Flags}}[global options]{{end}}{{if .Commands}} command [command options]{{end}} [arguments...]
   {{if .Version}}
VERSION:
   {{.Version}}
   {{end}}{{if len .Authors}}
AUTHOR(S):
   {{range .Authors}}{{ . }}{{end}}
   {{end}}{{if .Commands}}
COMMANDS:
   {{range .Commands}}{{join .Names ", "}}{{ "\t" }}{{.Usage}}
   {{end}}{{end}}{{if .Flags}}
GLOBAL OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}{{if .Copyright }}
COPYRIGHT:
   {{.Copyright}}
   {{end}}
`

// The text template for the command help topic.
// cli.go uses text/template to render templates. You can
// render custom help text by setting this variable.
var CommandHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   command {{.Name}}{{if .Flags}} [command options]{{end}} [arguments...]{{if .Description}}

DESCRIPTION:
   {{.Description}}{{end}}{{if .Flags}}

OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{ end }}
`

// The text template for the subcommand help topic.
// cli.go uses text/template to render templates. You can
// render custom help text by setting this variable.
var SubcommandHelpTemplate = `NAME:
   {{.Name}} - {{.Usage}}

USAGE:
   {{.Name}} command{{if .Flags}} [command options]{{end}} [arguments...]

COMMANDS:
   {{range .Commands}}{{join .Names ", "}}{{ "\t" }}{{.Usage}}
   {{end}}{{if .Flags}}
OPTIONS:
   {{range .Flags}}{{.}}
   {{end}}{{end}}
`

var helpCommand = Command{
	Name:    "help",
	Aliases: []string{"h"},
	Usage:   "Shows a list of commands or help for one command",
	Action: func(c *Context) {
		args := c.Args()
		if args.Present() {
			ShowCommandHelp(c, args.First())
		} else {
			ShowAppHelp(c)
		}
	},
}

var helpSubcommand = Command{
	Name:    "help",
	Aliases: []string{"h"},
	Usage:   "Shows a list of commands or help for one command",
	Action: func(c *Context) {
		args := c.Args()
		if args.Present() {
			ShowCommandHelp(c, args.First())
		} else {
			ShowSubcommandHelp(c)
		}
	},
}

// Prints help for the App or Command
type helpPrinter func(w io.Writer, templ string, data interface{})

var HelpPrinter helpPrinter = printHelp

// Prints version for the App
var VersionPrinter = printVersion

func ShowAppHelp(c *Context) {
	HelpPrinter(c.App.Writer, AppHelpTemplate, c.App)
}

// Prints the list of subcommands as the default app completion method
func DefaultAppComplete(c *Context) {
	for _, command := range c.App.Commands {
		for _, name := range command.Names() {
			fmt.Fprintln(c.App.Writer, name)
		}
	}
}

// Prints help for the given command
func ShowCommandHelp(ctx *Context, command string) {
	// show the subcommand help for a command with subcommands
	if command == "" {
		HelpPrinter(ctx.App.Writer, SubcommandHelpTemplate, ctx.App)
		return
	}

	for _, c := range ctx.App.Commands {
		if c.HasName(command) {
			HelpPrinter(ctx.App.Writer, CommandHelpTemplate, c)
			return
		}
	}

	if ctx.App.CommandNotFound != nil {
		ctx.App.CommandNotFound(ctx, command)
	} else {
		fmt.Fprintf(ctx.App.Writer, "No help topic for '%v'\n", command)
	}
}

// Prints help for the given subcommand
func ShowSubcommandHelp(c *Context) {
	ShowCommandHelp(c, c.Command.Name)
}

// Prints the version number of the App
func ShowVersion(c *Context) {
	VersionPrinter(c)
}

func printVersion(c *Context) {
	fmt.Fprintf(c.App.Writer, "%v version %v\n", c.App.Name, c.App.Version)
}

// Prints the lists of commands within a given context
func ShowCompletions(c *Context) {
	a := c.App
	if a != nil && a.BashComplete != nil {
		a.BashComplete(c)
	}
}

// Prints the custom completions for a given command
func ShowCommandCompletions(ctx *Context, command string) {
	c := ctx.App.Command(command)
	if c != nil && c.BashComplete != nil {
		c.BashComplete(ctx)
	}
}

func printHelp(out io.Writer, templ string, data interface{}) {
	funcMap := template.FuncMap{
		"join": strings.Join,
	}

	w := tabwriter.NewWriter(out, 0, 8, 1, '\t', 0)
	t := template.Must(template.New("help").Funcs(funcMap).Parse(templ))
	err := t.Execute(w, data)
	if err != nil {
		panic(err)
	}
	w.Flush()
}

func checkVersion(c *Context) bool {
	if c.GlobalBool("version") || c.GlobalBool("v") || c.Bool("version") || c.Bool("v") {
		ShowVersion(c)
		return true
	}

	return false
}

func checkHelp(c *Context) bool {
	if c.GlobalBool("h") || c.GlobalBool("help") || c.Bool("h") || c.Bool("help") {
		ShowAppHelp(c)
		return true
	}

	return false
}

func checkCommandHelp(c *Context, name string) bool {
	if c.Bool("h") || c.Bool("help") {
		ShowCommandHelp(c, name)
		return true
	}

	return false
}

func checkSubcommandHelp(c *Context) bool {
	if c.GlobalBool("h") || c.GlobalBool("help") {
		ShowSubcommandHelp(c)
		return true
	}

	return false
}

func checkCompletions(c *Context) bool {
	if (c.GlobalBool(BashCompletionFlag.Name) || c.Bool(BashCompletionFlag.Name)) && c.App.EnableBashCompletion {
		ShowCompletions(c)
		return true
	}

	return false
}

func checkCommandCompletions(c *Context, name string) bool {
	if c.Bool(BashCompletionFlag.Name) && c.App.EnableBashCompletion {
		ShowCommandCompletions(c, name)
		return true
	}

	return false
}
<|endoftext|>


<commit_before>package main

import (
	"golang.org/x/crypto/bcrypt"
	"net/http"
)

func ownerResetPassword(resetHex string, password string) error {
	if resetHex == "" || password == "" {
		return errorMissingField
	}

	passwordHash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		logger.Errorf("cannot generate hash from password: %v\n", err)
		return errorInternal
	}

	statement := `
		UPDATE owners SET passwordHash=$1
		WHERE email IN (
			SELECT email FROM ownerResetHexes
			WHERE resetHex=$2
		);
	`
	res, err := db.Exec(statement, string(passwordHash), resetHex)
	if err != nil {
		logger.Errorf("cannot change user's password: %v\n", err)
		return errorInternal
	}

	count, err := res.RowsAffected()
	if err != nil {
		logger.Errorf("cannot count rows affected: %v\n", err)
		return errorInternal
	}

	if count == 0 {
		return errorNoSuchResetToken
	}

	statement = `
		DELETE FROM ownerResetHexes
    WHERE resetHex=$1;
	`
	_, err = db.Exec(statement, resetHex)
	if err != nil {
		logger.Warningf("cannot remove reset token: %v\n", err)
	}

	return nil
}

func ownerResetPasswordHandler(w http.ResponseWriter, r *http.Request) {
	type request struct {
		ResetHex *string `json:"resetHex"`
		Password *string `json:"password"`
	}

	var x request
	if err := bodyUnmarshal(r, &x); err != nil {
		bodyMarshal(w, response{"success": false, "message": err.Error()})
		return
	}

	if err := ownerResetPassword(*x.ResetHex, *x.Password); err != nil {
		bodyMarshal(w, response{"success": false, "message": err.Error()})
		return
	}

	bodyMarshal(w, response{"success": true})
}
<commit_msg>owner_reset_password.go: use ownerHex in SELECT<commit_after>package main

import (
	"golang.org/x/crypto/bcrypt"
	"net/http"
)

func ownerResetPassword(resetHex string, password string) error {
	if resetHex == "" || password == "" {
		return errorMissingField
	}

	passwordHash, err := bcrypt.GenerateFromPassword([]byte(password), bcrypt.DefaultCost)
	if err != nil {
		logger.Errorf("cannot generate hash from password: %v\n", err)
		return errorInternal
	}

	statement := `
		UPDATE owners SET passwordHash=$1
		WHERE ownerHex = (
			SELECT ownerHex
			FROM ownerResetHexes
			WHERE resetHex=$2
		);
	`
	res, err := db.Exec(statement, string(passwordHash), resetHex)
	if err != nil {
		logger.Errorf("cannot change user's password: %v\n", err)
		return errorInternal
	}

	count, err := res.RowsAffected()
	if err != nil {
		logger.Errorf("cannot count rows affected: %v\n", err)
		return errorInternal
	}

	if count == 0 {
		return errorNoSuchResetToken
	}

	statement = `
		DELETE FROM ownerResetHexes
    WHERE resetHex=$1;
	`
	_, err = db.Exec(statement, resetHex)
	if err != nil {
		logger.Warningf("cannot remove reset token: %v\n", err)
	}

	return nil
}

func ownerResetPasswordHandler(w http.ResponseWriter, r *http.Request) {
	type request struct {
		ResetHex *string `json:"resetHex"`
		Password *string `json:"password"`
	}

	var x request
	if err := bodyUnmarshal(r, &x); err != nil {
		bodyMarshal(w, response{"success": false, "message": err.Error()})
		return
	}

	if err := ownerResetPassword(*x.ResetHex, *x.Password); err != nil {
		bodyMarshal(w, response{"success": false, "message": err.Error()})
		return
	}

	bodyMarshal(w, response{"success": true})
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package backups

import (
	"github.com/juju/errors"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/state"
	"github.com/juju/juju/state/backups/db"
)

// Create is the API method that requests juju to create a new backup
// of its state.  It returns the metadata for that backup.
func (a *API) Create(args params.BackupsCreateArgs) (p params.BackupsMetadataResult, err error) {
	backups, closer := newBackups(a.st)
	defer closer.Close()

	mgoInfo := a.st.MongoConnectionInfo()
	dbInfo := db.NewMongoConnInfo(mgoInfo)

	// TODO(ericsnow) The machine ID needs to be introspected from state.
	machine := "0"
	origin := state.NewBackupsOrigin(a.st, machine)

	meta, err := backups.Create(a.paths, *dbInfo, *origin, args.Notes)
	if err != nil {
		return p, errors.Trace(err)
	}

	p.UpdateFromMetadata(meta)

	return p, nil
}
<commit_msg>Clarify a TODO comment.<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package backups

import (
	"github.com/juju/errors"

	"github.com/juju/juju/apiserver/params"
	"github.com/juju/juju/state"
	"github.com/juju/juju/state/backups/db"
)

// Create is the API method that requests juju to create a new backup
// of its state.  It returns the metadata for that backup.
func (a *API) Create(args params.BackupsCreateArgs) (p params.BackupsMetadataResult, err error) {
	backups, closer := newBackups(a.st)
	defer closer.Close()

	mgoInfo := a.st.MongoConnectionInfo()
	dbInfo := db.NewMongoConnInfo(mgoInfo)

	// TODO(ericsnow) The machine ID needs to be introspected from the
	// API server, likely through a Resource.
	machine := "0"
	origin := state.NewBackupsOrigin(a.st, machine)

	meta, err := backups.Create(a.paths, *dbInfo, *origin, args.Notes)
	if err != nil {
		return p, errors.Trace(err)
	}

	p.UpdateFromMetadata(meta)

	return p, nil
}
<|endoftext|>


<commit_before>package htcat

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"sync"
)

const (
	_        = iota
	kB int64 = 1 << (10 * iota)
	mB
	gB
	tB
	pB
	eB
)

type HtCat struct {
	io.WriterTo
	d  defrag
	u  *url.URL
	cl *http.Client

	// Protect httpFragGen with a Mutex.
	httpFragGenMu sync.Mutex
	hfg           httpFragGen
}

type HttpStatusError struct {
	error
	Status string
}

func (cat *HtCat) startup(parallelism int) {
	req := http.Request{
		Method:     "GET",
		URL:        cat.u,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Body:       nil,
		Host:       cat.u.Host,
	}

	resp, err := cat.cl.Do(&req)
	if err != nil {
		go cat.d.cancel(err)
		return
	}

	// Check for non-200 OK response codes from the startup-GET.
	if resp.Status != "200 OK" {
		err = HttpStatusError{
			error: fmt.Errorf(
				"Expected HTTP Status 200, received: %q",
				resp.Status),
			Status: resp.Status}
		go cat.d.cancel(err)
		return
	}

	l := resp.Header.Get("Content-Length")

	// Some kinds of small or indeterminate-length files will
	// receive no parallelism.  This procedure helps prepare the
	// HtCat value for a one-HTTP-Request GET.
	noParallel := func(wtc writerToCloser) {
		f := cat.d.nextFragment()
		cat.d.setLast(cat.d.lastAllocated())
		f.contents = wtc
		cat.d.register(f)
	}

	if l == "" {
		// No Content-Length, stream without parallelism nor
		// assumptions about the length of the stream.
		go noParallel(struct {
			io.WriterTo
			io.Closer
		}{
			WriterTo: bufio.NewReader(resp.Body),
			Closer:   resp.Body,
		})
		return
	}

	length, err := strconv.ParseInt(l, 10, 64)
	if err != nil {
		// Invalid integer for Content-Length, defer reporting
		// the error until a WriteTo call is made.
		go cat.d.cancel(err)
		return
	}

	// Set up httpFrag generator state.
	cat.hfg.totalSize = length
	cat.hfg.targetFragSize = 1 + ((length - 1) / int64(parallelism))
	if cat.hfg.targetFragSize > 20*mB {
		cat.hfg.targetFragSize = 20 * mB
	}

	// Very small fragments are probably not worthwhile to start
	// up new requests for, but it in this case it was possible to
	// ascertain the size, so take advantage of that to start
	// reading in the background as eagerly as possible.
	if cat.hfg.targetFragSize < 1*mB {
		cat.hfg.curPos = cat.hfg.totalSize
		er := newEagerReader(resp.Body, cat.hfg.totalSize)
		go noParallel(er)
		go er.WaitClosed()
		return
	}

	// None of the other special short-circuit cases have been
	// triggered, so begin preparation for full-blown parallel
	// GET.  One GET worker is started here to take advantage of
	// the already pending response (which has no determinate
	// length, so it must be limited).
	hf := cat.nextFragment()
	go func() {
		er := newEagerReader(
			struct {
				io.Reader
				io.Closer
			}{
				Reader: io.LimitReader(resp.Body, hf.size),
				Closer: resp.Body,
			},
			hf.size)

		hf.fragment.contents = er
		cat.d.register(hf.fragment)
		er.WaitClosed()

		// Chain into being a regular worker, having finished
		// the special start-up segment.
		cat.get()
	}()
}

func New(client *http.Client, u *url.URL, parallelism int) *HtCat {
	cat := HtCat{
		u:  u,
		cl: client,
	}

	cat.d.initDefrag()
	cat.WriterTo = &cat.d
	cat.startup(parallelism)

	if cat.hfg.curPos == cat.hfg.totalSize {
		return &cat
	}

	// Start background workers.
	//
	// "startup" starts one worker that is specially constructed
	// to deal with the first request, so back off by one to
	// prevent performing with too much parallelism.
	for i := 1; i < parallelism; i += 1 {
		go cat.get()
	}

	return &cat
}

func (cat *HtCat) nextFragment() *httpFrag {
	cat.httpFragGenMu.Lock()
	defer cat.httpFragGenMu.Unlock()

	var hf *httpFrag

	if cat.hfg.hasNext() {
		f := cat.d.nextFragment()
		hf = cat.hfg.nextFragment(f)
	} else {
		cat.d.setLast(cat.d.lastAllocated())
	}

	return hf
}

func (cat *HtCat) get() {
	for {
		hf := cat.nextFragment()
		if hf == nil {
			return
		}

		req := http.Request{
			Method:     "GET",
			URL:        cat.u,
			Proto:      "HTTP/1.1",
			ProtoMajor: 1,
			ProtoMinor: 1,
			Header:     hf.header,
			Body:       nil,
			Host:       cat.u.Host,
		}

		resp, err := cat.cl.Do(&req)
		if err != nil {
			cat.d.cancel(err)
			return
		}

		// Check for an acceptable HTTP status code.
		if !(resp.Status == "206 Partial Content" ||
			resp.Status == "200 OK") {
			err = HttpStatusError{
				error: fmt.Errorf("Expected HTTP Status "+
					"206 or 200, received: %q",
					resp.Status),
				Status: resp.Status}
			go cat.d.cancel(err)
			return
		}

		er := newEagerReader(resp.Body, hf.size)
		hf.fragment.contents = er
		cat.d.register(hf.fragment)
		er.WaitClosed()
	}
}
<commit_msg>Use Response.StatusCode not .Status<commit_after>package htcat

import (
	"bufio"
	"fmt"
	"io"
	"net/http"
	"net/url"
	"strconv"
	"sync"
)

const (
	_        = iota
	kB int64 = 1 << (10 * iota)
	mB
	gB
	tB
	pB
	eB
)

type HtCat struct {
	io.WriterTo
	d  defrag
	u  *url.URL
	cl *http.Client

	// Protect httpFragGen with a Mutex.
	httpFragGenMu sync.Mutex
	hfg           httpFragGen
}

type HttpStatusError struct {
	error
	Status string
}

func (cat *HtCat) startup(parallelism int) {
	req := http.Request{
		Method:     "GET",
		URL:        cat.u,
		Proto:      "HTTP/1.1",
		ProtoMajor: 1,
		ProtoMinor: 1,
		Body:       nil,
		Host:       cat.u.Host,
	}

	resp, err := cat.cl.Do(&req)
	if err != nil {
		go cat.d.cancel(err)
		return
	}

	// Check for non-200 OK response codes from the startup-GET.
	if resp.StatusCode != 200 {
		err = HttpStatusError{
			error: fmt.Errorf(
				"Expected HTTP Status 200, received: %q",
				resp.Status),
			Status: resp.Status}
		go cat.d.cancel(err)
		return
	}

	l := resp.Header.Get("Content-Length")

	// Some kinds of small or indeterminate-length files will
	// receive no parallelism.  This procedure helps prepare the
	// HtCat value for a one-HTTP-Request GET.
	noParallel := func(wtc writerToCloser) {
		f := cat.d.nextFragment()
		cat.d.setLast(cat.d.lastAllocated())
		f.contents = wtc
		cat.d.register(f)
	}

	if l == "" {
		// No Content-Length, stream without parallelism nor
		// assumptions about the length of the stream.
		go noParallel(struct {
			io.WriterTo
			io.Closer
		}{
			WriterTo: bufio.NewReader(resp.Body),
			Closer:   resp.Body,
		})
		return
	}

	length, err := strconv.ParseInt(l, 10, 64)
	if err != nil {
		// Invalid integer for Content-Length, defer reporting
		// the error until a WriteTo call is made.
		go cat.d.cancel(err)
		return
	}

	// Set up httpFrag generator state.
	cat.hfg.totalSize = length
	cat.hfg.targetFragSize = 1 + ((length - 1) / int64(parallelism))
	if cat.hfg.targetFragSize > 20*mB {
		cat.hfg.targetFragSize = 20 * mB
	}

	// Very small fragments are probably not worthwhile to start
	// up new requests for, but it in this case it was possible to
	// ascertain the size, so take advantage of that to start
	// reading in the background as eagerly as possible.
	if cat.hfg.targetFragSize < 1*mB {
		cat.hfg.curPos = cat.hfg.totalSize
		er := newEagerReader(resp.Body, cat.hfg.totalSize)
		go noParallel(er)
		go er.WaitClosed()
		return
	}

	// None of the other special short-circuit cases have been
	// triggered, so begin preparation for full-blown parallel
	// GET.  One GET worker is started here to take advantage of
	// the already pending response (which has no determinate
	// length, so it must be limited).
	hf := cat.nextFragment()
	go func() {
		er := newEagerReader(
			struct {
				io.Reader
				io.Closer
			}{
				Reader: io.LimitReader(resp.Body, hf.size),
				Closer: resp.Body,
			},
			hf.size)

		hf.fragment.contents = er
		cat.d.register(hf.fragment)
		er.WaitClosed()

		// Chain into being a regular worker, having finished
		// the special start-up segment.
		cat.get()
	}()
}

func New(client *http.Client, u *url.URL, parallelism int) *HtCat {
	cat := HtCat{
		u:  u,
		cl: client,
	}

	cat.d.initDefrag()
	cat.WriterTo = &cat.d
	cat.startup(parallelism)

	if cat.hfg.curPos == cat.hfg.totalSize {
		return &cat
	}

	// Start background workers.
	//
	// "startup" starts one worker that is specially constructed
	// to deal with the first request, so back off by one to
	// prevent performing with too much parallelism.
	for i := 1; i < parallelism; i += 1 {
		go cat.get()
	}

	return &cat
}

func (cat *HtCat) nextFragment() *httpFrag {
	cat.httpFragGenMu.Lock()
	defer cat.httpFragGenMu.Unlock()

	var hf *httpFrag

	if cat.hfg.hasNext() {
		f := cat.d.nextFragment()
		hf = cat.hfg.nextFragment(f)
	} else {
		cat.d.setLast(cat.d.lastAllocated())
	}

	return hf
}

func (cat *HtCat) get() {
	for {
		hf := cat.nextFragment()
		if hf == nil {
			return
		}

		req := http.Request{
			Method:     "GET",
			URL:        cat.u,
			Proto:      "HTTP/1.1",
			ProtoMajor: 1,
			ProtoMinor: 1,
			Header:     hf.header,
			Body:       nil,
			Host:       cat.u.Host,
		}

		resp, err := cat.cl.Do(&req)
		if err != nil {
			cat.d.cancel(err)
			return
		}

		// Check for an acceptable HTTP status code.
		if !(resp.Status == "206 Partial Content" ||
			resp.Status == "200 OK") {
			err = HttpStatusError{
				error: fmt.Errorf("Expected HTTP Status "+
					"206 or 200, received: %q",
					resp.Status),
				Status: resp.Status}
			go cat.d.cancel(err)
			return
		}

		er := newEagerReader(resp.Body, hf.size)
		hf.fragment.contents = er
		cat.d.register(hf.fragment)
		er.WaitClosed()
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/litl/galaxy/utils"

	"github.com/litl/galaxy/log"
	gotoolslog "github.com/mailgun/gotools-log"
	"github.com/mailgun/vulcan"
	"github.com/mailgun/vulcan/endpoint"
	"github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/location/httploc"
	"github.com/mailgun/vulcan/request"
	"github.com/mailgun/vulcan/route"
	"github.com/mailgun/vulcan/route/hostroute"
)

var (
	httpRouter *HTTPRouter
)

type RequestLogger struct{}

type HTTPRouter struct {
	router    *hostroute.HostRouter
	balancers map[string]*roundrobin.RoundRobin
}

func (r *RequestLogger) ObserveRequest(req request.Request) {}

func (r *RequestLogger) ObserveResponse(req request.Request, a request.Attempt) {
	err := ""
	statusCode := ""
	if a.GetError() != nil {
		err = " err=" + a.GetError().Error()
	}

	if a.GetResponse() != nil {
		statusCode = " status=" + strconv.FormatInt(int64(a.GetResponse().StatusCode), 10)
	}

	log.Printf("id=%d method=%s clientIp=%s url=%s backend=%s%s duration=%s%s",
		req.GetId(),
		req.GetHttpRequest().Method,
		req.GetHttpRequest().RemoteAddr,
		req.GetHttpRequest().Host+req.GetHttpRequest().RequestURI,
		a.GetEndpoint(),
		statusCode, a.GetDuration(), err)
}

type SSLRedirect struct{}

func (s *SSLRedirect) ProcessRequest(r request.Request) (*http.Response, error) {
	if sslOnly && r.GetHttpRequest().Header.Get("X-Forwarded-Proto") != "https" {

		resp := &http.Response{
			Status:        "301 Moved Permanently",
			StatusCode:    301,
			Proto:         r.GetHttpRequest().Proto,
			ProtoMajor:    r.GetHttpRequest().ProtoMajor,
			ProtoMinor:    r.GetHttpRequest().ProtoMinor,
			Body:          ioutil.NopCloser(bytes.NewBufferString("")),
			ContentLength: 0,
			Request:       r.GetHttpRequest(),
			Header:        http.Header{},
		}
		resp.Header.Set("Location", "https://"+r.GetHttpRequest().Host+r.GetHttpRequest().RequestURI)
		return resp, nil
	}

	return nil, nil
}

func (s *SSLRedirect) ProcessResponse(r request.Request, a request.Attempt) {
}

func NewHTTPRouter() *HTTPRouter {
	return &HTTPRouter{
		balancers: make(map[string]*roundrobin.RoundRobin),
	}
}

func (s *HTTPRouter) AddBackend(name, vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	var err error
	balancer := s.balancers[vhost]

	if balancer == nil {
		// Create a round robin load balancer with some endpoints
		balancer, err = roundrobin.NewRoundRobin()
		if err != nil {
			return err
		}

		// Create a http location with the load balancer we've just added
		loc, err := httploc.NewLocationWithOptions(name, balancer,
			httploc.Options{
				TrustForwardHeader: true,
			})
		if err != nil {
			return err
		}
		loc.GetObserverChain().Add("logger", &RequestLogger{})
		loc.GetMiddlewareChain().Add("ssl", 0, &SSLRedirect{})

		s.router.SetRouter(vhost, &route.ConstRouter{Location: loc})
		log.Printf("Starting HTTP listener for %s", vhost)
		s.balancers[vhost] = balancer
	}

	// Already registered?
	if balancer.FindEndpointByUrl(url) != nil {
		return nil
	}
	endpoint := endpoint.MustParseUrl(url)
	log.Printf("Adding HTTP endpoint %s to %s", endpoint.GetUrl(), vhost)
	err = balancer.AddEndpoint(endpoint)
	if err != nil {
		return err
	}
	return nil
}

func (s *HTTPRouter) RemoveBackend(vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	balancer := s.balancers[vhost]
	if balancer == nil {
		return nil
	}

	endpoint := balancer.FindEndpointByUrl(url)
	if endpoint == nil {
		return nil
	}
	log.Printf("Removing HTTP endpoint %s from %s ", endpoint.GetUrl(), vhost)
	balancer.RemoveEndpoint(endpoint)

	endpoints := balancer.GetEndpoints()
	if len(endpoints) == 0 {
		s.RemoveRouter(vhost)
	}
	return nil
}

// Remove all backends for vhost that are not listed in addrs
func (s *HTTPRouter) RemoveBackends(vhost string, addrs []string) {

	if vhost == "" {
		return
	}

	// Remove backends that are no longer registered

	balancer := s.balancers[vhost]
	if balancer == nil {
		return
	}

	endpoints := balancer.GetEndpoints()
	for _, endpoint := range endpoints {
		if !utils.StringInSlice(endpoint.GetUrl().String(), addrs) {
			s.RemoveBackend(vhost, endpoint.GetUrl().String())
		}
	}
}

// Removes a virtual host router
func (s *HTTPRouter) RemoveRouter(vhost string) {
	if vhost == "" {
		return
	}

	log.Printf("Removing balancer for %s", vhost)
	delete(s.balancers, vhost)
	s.router.RemoveRouter(vhost)
}

func (s *HTTPRouter) adminHandler(w http.ResponseWriter, r *http.Request) {
	if len(s.balancers) == 0 {
		w.WriteHeader(503)
		return
	}

	keys := make([]string, 0, len(s.balancers))
	for key := range s.balancers {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, k := range keys {
		balancer := s.balancers[k]
		endpoints := balancer.GetEndpoints()
		fmt.Fprintf(w, "%s\n", k)
		for _, endpoint := range endpoints {
			fmt.Fprintf(w, "  %s\t%d\t%d\t%0.2f\n", endpoint.GetUrl(), endpoint.GetOriginalWeight(), endpoint.GetEffectiveWeight(), endpoint.GetMeter().GetRate())
		}
	}
}

func (s *HTTPRouter) statusHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		host := r.Host
		if strings.Contains(host, ":") {
			host, _, err = net.SplitHostPort(r.Host)
			if err != nil {
				log.Warnf("%s", err)
				h.ServeHTTP(w, r)
				return
			}
		}

		if _, exists := s.balancers[host]; !exists {
			s.adminHandler(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}

func (s *HTTPRouter) Start() {

	// init the vulcan logging
	gotoolslog.Init([]*gotoolslog.LogConfig{
		&gotoolslog.LogConfig{Name: "console"},
	})

	log.Printf("HTTP server listening at %s", listenAddr)

	s.router = hostroute.NewHostRouter()

	proxy, err := vulcan.NewProxy(s.router)
	if err != nil {
		log.Fatalf("ERROR: %s", err)
	}

	// Proxy acts as http handler:
	server := &http.Server{
		Addr:           listenAddr,
		Handler:        s.statusHandler(proxy),
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Errorf("%s", server.ListenAndServe())
}

func startHTTPServer() {
	defer wg.Done()
	httpRouter = NewHTTPRouter()
	httpRouter.Start()
}
<commit_msg>Disable vulcan logger unless debug is enabled<commit_after>package main

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"github.com/litl/galaxy/utils"

	"github.com/litl/galaxy/log"
	gotoolslog "github.com/mailgun/gotools-log"
	"github.com/mailgun/vulcan"
	"github.com/mailgun/vulcan/endpoint"
	"github.com/mailgun/vulcan/loadbalance/roundrobin"
	"github.com/mailgun/vulcan/location/httploc"
	"github.com/mailgun/vulcan/request"
	"github.com/mailgun/vulcan/route"
	"github.com/mailgun/vulcan/route/hostroute"
)

var (
	httpRouter *HTTPRouter
)

type RequestLogger struct{}

type HTTPRouter struct {
	router    *hostroute.HostRouter
	balancers map[string]*roundrobin.RoundRobin
}

func (r *RequestLogger) ObserveRequest(req request.Request) {}

func (r *RequestLogger) ObserveResponse(req request.Request, a request.Attempt) {
	err := ""
	statusCode := ""
	if a.GetError() != nil {
		err = " err=" + a.GetError().Error()
	}

	if a.GetResponse() != nil {
		statusCode = " status=" + strconv.FormatInt(int64(a.GetResponse().StatusCode), 10)
	}

	log.Printf("id=%d method=%s clientIp=%s url=%s backend=%s%s duration=%s%s",
		req.GetId(),
		req.GetHttpRequest().Method,
		req.GetHttpRequest().RemoteAddr,
		req.GetHttpRequest().Host+req.GetHttpRequest().RequestURI,
		a.GetEndpoint(),
		statusCode, a.GetDuration(), err)
}

type SSLRedirect struct{}

func (s *SSLRedirect) ProcessRequest(r request.Request) (*http.Response, error) {
	if sslOnly && r.GetHttpRequest().Header.Get("X-Forwarded-Proto") != "https" {

		resp := &http.Response{
			Status:        "301 Moved Permanently",
			StatusCode:    301,
			Proto:         r.GetHttpRequest().Proto,
			ProtoMajor:    r.GetHttpRequest().ProtoMajor,
			ProtoMinor:    r.GetHttpRequest().ProtoMinor,
			Body:          ioutil.NopCloser(bytes.NewBufferString("")),
			ContentLength: 0,
			Request:       r.GetHttpRequest(),
			Header:        http.Header{},
		}
		resp.Header.Set("Location", "https://"+r.GetHttpRequest().Host+r.GetHttpRequest().RequestURI)
		return resp, nil
	}

	return nil, nil
}

func (s *SSLRedirect) ProcessResponse(r request.Request, a request.Attempt) {
}

func NewHTTPRouter() *HTTPRouter {
	return &HTTPRouter{
		balancers: make(map[string]*roundrobin.RoundRobin),
	}
}

func (s *HTTPRouter) AddBackend(name, vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	var err error
	balancer := s.balancers[vhost]

	if balancer == nil {
		// Create a round robin load balancer with some endpoints
		balancer, err = roundrobin.NewRoundRobin()
		if err != nil {
			return err
		}

		// Create a http location with the load balancer we've just added
		loc, err := httploc.NewLocationWithOptions(name, balancer,
			httploc.Options{
				TrustForwardHeader: true,
			})
		if err != nil {
			return err
		}
		loc.GetObserverChain().Add("logger", &RequestLogger{})
		loc.GetMiddlewareChain().Add("ssl", 0, &SSLRedirect{})

		s.router.SetRouter(vhost, &route.ConstRouter{Location: loc})
		log.Printf("Starting HTTP listener for %s", vhost)
		s.balancers[vhost] = balancer
	}

	// Already registered?
	if balancer.FindEndpointByUrl(url) != nil {
		return nil
	}
	endpoint := endpoint.MustParseUrl(url)
	log.Printf("Adding HTTP endpoint %s to %s", endpoint.GetUrl(), vhost)
	err = balancer.AddEndpoint(endpoint)
	if err != nil {
		return err
	}
	return nil
}

func (s *HTTPRouter) RemoveBackend(vhost, url string) error {

	if vhost == "" || url == "" {
		return nil
	}

	balancer := s.balancers[vhost]
	if balancer == nil {
		return nil
	}

	endpoint := balancer.FindEndpointByUrl(url)
	if endpoint == nil {
		return nil
	}
	log.Printf("Removing HTTP endpoint %s from %s ", endpoint.GetUrl(), vhost)
	balancer.RemoveEndpoint(endpoint)

	endpoints := balancer.GetEndpoints()
	if len(endpoints) == 0 {
		s.RemoveRouter(vhost)
	}
	return nil
}

// Remove all backends for vhost that are not listed in addrs
func (s *HTTPRouter) RemoveBackends(vhost string, addrs []string) {

	if vhost == "" {
		return
	}

	// Remove backends that are no longer registered

	balancer := s.balancers[vhost]
	if balancer == nil {
		return
	}

	endpoints := balancer.GetEndpoints()
	for _, endpoint := range endpoints {
		if !utils.StringInSlice(endpoint.GetUrl().String(), addrs) {
			s.RemoveBackend(vhost, endpoint.GetUrl().String())
		}
	}
}

// Removes a virtual host router
func (s *HTTPRouter) RemoveRouter(vhost string) {
	if vhost == "" {
		return
	}

	log.Printf("Removing balancer for %s", vhost)
	delete(s.balancers, vhost)
	s.router.RemoveRouter(vhost)
}

func (s *HTTPRouter) adminHandler(w http.ResponseWriter, r *http.Request) {
	if len(s.balancers) == 0 {
		w.WriteHeader(503)
		return
	}

	keys := make([]string, 0, len(s.balancers))
	for key := range s.balancers {
		keys = append(keys, key)
	}
	sort.Strings(keys)

	for _, k := range keys {
		balancer := s.balancers[k]
		endpoints := balancer.GetEndpoints()
		fmt.Fprintf(w, "%s\n", k)
		for _, endpoint := range endpoints {
			fmt.Fprintf(w, "  %s\t%d\t%d\t%0.2f\n", endpoint.GetUrl(), endpoint.GetOriginalWeight(), endpoint.GetEffectiveWeight(), endpoint.GetMeter().GetRate())
		}
	}
}

func (s *HTTPRouter) statusHandler(h http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		var err error
		host := r.Host
		if strings.Contains(host, ":") {
			host, _, err = net.SplitHostPort(r.Host)
			if err != nil {
				log.Warnf("%s", err)
				h.ServeHTTP(w, r)
				return
			}
		}

		if _, exists := s.balancers[host]; !exists {
			s.adminHandler(w, r)
			return
		}
		h.ServeHTTP(w, r)
	})
}

func (s *HTTPRouter) Start() {

	if debug {
		// init the vulcan logging
		gotoolslog.Init([]*gotoolslog.LogConfig{
			&gotoolslog.LogConfig{Name: "console"},
		})
	}

	log.Printf("HTTP server listening at %s", listenAddr)

	s.router = hostroute.NewHostRouter()

	proxy, err := vulcan.NewProxy(s.router)
	if err != nil {
		log.Fatalf("ERROR: %s", err)
	}

	// Proxy acts as http handler:
	server := &http.Server{
		Addr:           listenAddr,
		Handler:        s.statusHandler(proxy),
		ReadTimeout:    60 * time.Second,
		WriteTimeout:   60 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}

	log.Errorf("%s", server.ListenAndServe())
}

func startHTTPServer() {
	defer wg.Done()
	httpRouter = NewHTTPRouter()
	httpRouter.Start()
}
<|endoftext|>