<commit_before>package server

import (
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/auth"
	"github.com/smancke/guble/store"

	log "github.com/Sirupsen/logrus"

	"fmt"
	"github.com/smancke/guble/server/cluster"
	"runtime"
	"strings"
	"sync"
	"time"
)

const (
	overloadedHandleChannelRatio = 0.9
	handleChannelCapacity        = 500
	subscribeChannelCapacity     = 10
	unsubscribeChannelCapacity   = 10
)

// Router interface provides a mechanism for PubSub messaging
type Router interface {
	AccessManager() (auth.AccessManager, error)
	MessageStore() (store.MessageStore, error)
	KVStore() (store.KVStore, error)
	Cluster() *cluster.Cluster

	Subscribe(r *Route) (*Route, error)
	Unsubscribe(r *Route)
	HandleMessage(message *protocol.Message) error
}

// Helper struct to pass `Route` to subscription channel and provide a notification channel.
type subRequest struct {
	route *Route
	doneC chan bool
}

type router struct {
	routes       map[protocol.Path][]*Route // mapping the path to the route slice
	handleC      chan *protocol.Message
	subscribeC   chan subRequest
	unsubscribeC chan subRequest
	stopC        chan bool      // Channel that signals stop of the router
	stopping     bool           // Flag: the router is in stopping process and no incoming messages are accepted
	wg           sync.WaitGroup // Add any operation that we need to wait upon here

	accessManager auth.AccessManager
	messageStore  store.MessageStore
	kvStore       store.KVStore
	cluster       *cluster.Cluster
}

// NewRouter returns a pointer to Router
func NewRouter(accessManager auth.AccessManager, messageStore store.MessageStore, kvStore store.KVStore, cluster *cluster.Cluster) Router {
	return &router{
		routes: make(map[protocol.Path][]*Route),

		handleC:      make(chan *protocol.Message, handleChannelCapacity),
		subscribeC:   make(chan subRequest, subscribeChannelCapacity),
		unsubscribeC: make(chan subRequest, unsubscribeChannelCapacity),
		stopC:        make(chan bool, 1),

		accessManager: accessManager,
		messageStore:  messageStore,
		kvStore:       kvStore,
		cluster:       cluster,
	}
}

func (router *router) Start() error {
	router.panicIfInternalDependenciesAreNil()
	resetRouterMetrics()

	go func() {
		router.wg.Add(1)
		for {
			if router.stopping && router.channelsAreEmpty() {
				router.closeRoutes()
				router.wg.Done()
				return
			}

			func() {
				defer protocol.PanicLogger()

				select {
				case message := <-router.handleC:
					router.routeMessage(message)
					runtime.Gosched()
				case subscriber := <-router.subscribeC:
					router.subscribe(subscriber.route)
					subscriber.doneC <- true
				case unsubscriber := <-router.unsubscribeC:
					router.unsubscribe(unsubscriber.route)
					unsubscriber.doneC <- true
				case <-router.stopC:
					router.stopping = true
				}
			}()
		}
	}()

	return nil
}

// Stop stops the router by closing the stop channel, and waiting on the WaitGroup
func (router *router) Stop() error {
	logger.Debug("Stopping router")

	router.stopC <- true
	router.wg.Wait()
	return nil
}

func (router *router) Check() error {
	if router.accessManager == nil || router.messageStore == nil || router.kvStore == nil {
		logger.WithFields(log.Fields{
			"err": ErrServiceNotProvided,
		}).Error("Some services are not provided")
		return ErrServiceNotProvided
	}

	err := router.messageStore.Check()
	if err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("MessageStore check failed")
		return err
	}

	err = router.kvStore.Check()
	if err != nil {
		log.WithFields(log.Fields{
			"module": "router",
			"err":    err,
		}).Error("KvStore check failed")
		return err
	}

	return nil
}

// HandleMessage assigns the new message id, stores the message and passes it to the handleC channel, and asynchronously to the cluster (if available).
func (router *router) HandleMessage(message *protocol.Message) error {
	logger.WithFields(log.Fields{
		"userID": message.UserID,
		"path":   message.Path,
	}).Debug("HandleMessage")

	mTotalMessagesIncoming.Add(1)
	if err := router.isStopping(); err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("Router is stopping")
		return err
	}

	if !router.accessManager.IsAllowed(auth.WRITE, message.UserID, message.Path) {
		return &PermissionDeniedError{message.UserID, auth.WRITE, message.Path}
	}

	txCallback := func(msgId uint64) []byte {
		message.ID = msgId
		message.Time = time.Now().Unix()
		return message.Bytes()
	}
	lenMsg := int64(len(message.Bytes()))
	mTotalMessagesIncomingBytes.Add(lenMsg)

	if err := router.messageStore.StoreTx(message.Path.Partition(), txCallback); err != nil {
		logger.WithFields(log.Fields{
			"err":          err,
			"msgPartition": message.Path.Partition(),
		}).Error("Error storing message in partition")
		mTotalMessageStoreErrors.Add(1)
		return err
	}
	mTotalMessagesStoredBytes.Add(lenMsg)

	if float32(len(router.handleC))/float32(cap(router.handleC)) > overloadedHandleChannelRatio {
		logger.WithFields(log.Fields{
			"currentLength": len(router.handleC),
			"maxCapacity":   cap(router.handleC),
		}).Warn("storeAndChannelMessage: handleC channel is almost full")
		mTotalOverloadedHandleChannel.Add(1)
	}

	router.handleC <- message

	if router.cluster != nil {
		go router.cluster.BroadcastMessage(message)
	}

	return nil
}

// Subscribe adds a route to the subscribers.
// If there is already a route with same Application Id and Path, it will be replaced.
func (router *router) Subscribe(r *Route) (*Route, error) {

	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Subscribe")

	if err := router.isStopping(); err != nil {
		return nil, err
	}

	accessAllowed := router.accessManager.IsAllowed(auth.READ, r.UserID, r.Path)
	if !accessAllowed {
		return r, &PermissionDeniedError{UserID: r.UserID, AccessType: auth.READ, Path: r.Path}
	}
	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.subscribeC <- req
	<-req.doneC
	return r, nil
}

func (router *router) Unsubscribe(r *Route) {
	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Unsubscribe")

	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.unsubscribeC <- req
	<-req.doneC
}

func (router *router) subscribe(r *Route) {
	logger.WithFields(log.Fields{
		"userID": r.UserID,
		"path":   r.Path,
	}).Debug("Internal subscribe")
	mTotalSubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	var removed bool
	if present {
		// Try to remove, to avoid double subscriptions of the same app
		slice, removed = removeIfMatching(slice, r)
	} else {
		// Path not present yet. Initialize the slice
		slice = make([]*Route, 0, 1)
		router.routes[r.Path] = slice
		mCurrentRoutes.Add(1)
	}
	router.routes[r.Path] = append(slice, r)
	if removed {
		mTotalDuplicateSubscriptionsAttempts.Add(1)
	} else {
		mTotalSubscriptions.Add(1)
		mCurrentSubscriptions.Add(1)
	}
}

func (router *router) unsubscribe(r *Route) {

	logger.WithFields(log.Fields{
		"userID": r.UserID,
		"path":   r.Path,
	}).Debug("Internal unsubscribe")
	mTotalUnsubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	if !present {
		mTotalInvalidTopicOnUnsubscriptionAttempts.Add(1)
		return
	}
	var removed bool
	router.routes[r.Path], removed = removeIfMatching(slice, r)
	if removed {
		mTotalUnsubscriptions.Add(1)
		mCurrentSubscriptions.Add(-1)
	} else {
		mTotalInvalidUnsubscriptionAttempts.Add(1)
	}
	if len(router.routes[r.Path]) == 0 {
		delete(router.routes, r.Path)
		mCurrentRoutes.Add(-1)
	}
}

func (router *router) panicIfInternalDependenciesAreNil() {
	if router.accessManager == nil || router.kvStore == nil || router.messageStore == nil {
		panic(fmt.Sprintf("router: the internal dependencies marked with `true` are not set: AccessManager=%v, KVStore=%v, MessageStore=%v",
			router.accessManager == nil, router.kvStore == nil, router.messageStore == nil))
	}
}

func (router *router) channelsAreEmpty() bool {
	return len(router.handleC) == 0 && len(router.subscribeC) == 0 && len(router.unsubscribeC) == 0
}

func (router *router) isStopping() error {
	if router.stopping {
		return &ModuleStoppingError{"Router"}
	}
	return nil
}

func (router *router) routeMessage(message *protocol.Message) {

	logger.WithFields(log.Fields{
		"msgMetadata": message.Metadata(),
	}).Debug("Called routeMessage")
	mTotalMessagesRouted.Add(1)

	for path, list := range router.routes {
		if matchesTopic(message.Path, path) {
			for _, route := range list {
				router.deliverMessage(route, message)
			}
		} else {
			mTotalMessagesNotMatchingTopic.Add(1)
		}
	}
}

func (router *router) deliverMessage(route *Route, message *protocol.Message) {
	defer protocol.PanicLogger()

	select {
	case route.MessagesChannel() <- &MessageForRoute{Message: message, Route: route}:
	// fine, we could send the message
	default:
		logger.WithFields(log.Fields{
			"route": route.String(),
		}).Warn("deliverMessage: queue was full, unsubscribing and closing delivery channel for route")
		router.unsubscribe(route)
		route.Close()
		mTotalDeliverMessageErrors.Add(1)
	}
}

func (router *router) closeRoutes() {
	logger.Debug("Called closeRoutes")

	for _, currentRouteList := range router.routes {
		for _, route := range currentRouteList {
			router.unsubscribe(route)
			log.WithFields(log.Fields{
				"module": "router",
				"route":  route.String(),
			}).Debug("Closing route for ")
			route.Close()
		}
	}
}

func copyOf(message []byte) []byte {
	messageCopy := make([]byte, len(message))
	copy(messageCopy, message)
	return messageCopy
}

// matchesTopic checks whether the supplied routePath matches the message topic
func matchesTopic(messagePath, routePath protocol.Path) bool {
	messagePathLen := len(string(messagePath))
	routePathLen := len(string(routePath))
	return strings.HasPrefix(string(messagePath), string(routePath)) &&
		(messagePathLen == routePathLen ||
			(messagePathLen > routePathLen && string(messagePath)[routePathLen] == '/'))
}

// removeIfMatching removes a route from the supplied list, based on same ApplicationID id and same path (if existing)
// returns: the (possibly updated) slide, and a boolean value (true if route was removed, false otherwise)
func removeIfMatching(slice []*Route, route *Route) ([]*Route, bool) {
	position := -1
	for p, r := range slice {
		if r.ApplicationID == route.ApplicationID && r.Path == route.Path {
			position = p
		}
	}
	if position == -1 {
		return slice, false
	}
	return append(slice[:position], slice[position+1:]...), true
}

// AccessManager returns the `accessManager` provided for the router
func (router *router) AccessManager() (auth.AccessManager, error) {
	if router.accessManager == nil {
		return nil, ErrServiceNotProvided
	}
	return router.accessManager, nil
}

// MessageStore returns the `messageStore` provided for the router
func (router *router) MessageStore() (store.MessageStore, error) {
	if router.messageStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.messageStore, nil
}

// KVStore returns the `kvStore` provided for the router
func (router *router) KVStore() (store.KVStore, error) {
	if router.kvStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.kvStore, nil
}

// Cluster returns the `cluster` provided for the router, or nil if no cluster was set-up
func (router *router) Cluster() *cluster.Cluster {
	return router.cluster
}
<commit_msg>removing unused copyOf func<commit_after>package server

import (
	"github.com/smancke/guble/protocol"
	"github.com/smancke/guble/server/auth"
	"github.com/smancke/guble/server/cluster"
	"github.com/smancke/guble/store"

	log "github.com/Sirupsen/logrus"

	"fmt"
	"runtime"
	"strings"
	"sync"
	"time"
)

const (
	overloadedHandleChannelRatio = 0.9
	handleChannelCapacity        = 500
	subscribeChannelCapacity     = 10
	unsubscribeChannelCapacity   = 10
)

// Router interface provides a mechanism for PubSub messaging
type Router interface {
	AccessManager() (auth.AccessManager, error)
	MessageStore() (store.MessageStore, error)
	KVStore() (store.KVStore, error)
	Cluster() *cluster.Cluster

	Subscribe(r *Route) (*Route, error)
	Unsubscribe(r *Route)
	HandleMessage(message *protocol.Message) error
}

// Helper struct to pass `Route` to subscription channel and provide a notification channel.
type subRequest struct {
	route *Route
	doneC chan bool
}

type router struct {
	routes       map[protocol.Path][]*Route // mapping the path to the route slice
	handleC      chan *protocol.Message
	subscribeC   chan subRequest
	unsubscribeC chan subRequest
	stopC        chan bool      // Channel that signals stop of the router
	stopping     bool           // Flag: the router is in stopping process and no incoming messages are accepted
	wg           sync.WaitGroup // Add any operation that we need to wait upon here

	accessManager auth.AccessManager
	messageStore  store.MessageStore
	kvStore       store.KVStore
	cluster       *cluster.Cluster
}

// NewRouter returns a pointer to Router
func NewRouter(accessManager auth.AccessManager, messageStore store.MessageStore, kvStore store.KVStore, cluster *cluster.Cluster) Router {
	return &router{
		routes: make(map[protocol.Path][]*Route),

		handleC:      make(chan *protocol.Message, handleChannelCapacity),
		subscribeC:   make(chan subRequest, subscribeChannelCapacity),
		unsubscribeC: make(chan subRequest, unsubscribeChannelCapacity),
		stopC:        make(chan bool, 1),

		accessManager: accessManager,
		messageStore:  messageStore,
		kvStore:       kvStore,
		cluster:       cluster,
	}
}

func (router *router) Start() error {
	router.panicIfInternalDependenciesAreNil()
	resetRouterMetrics()

	go func() {
		router.wg.Add(1)
		for {
			if router.stopping && router.channelsAreEmpty() {
				router.closeRoutes()
				router.wg.Done()
				return
			}

			func() {
				defer protocol.PanicLogger()

				select {
				case message := <-router.handleC:
					router.routeMessage(message)
					runtime.Gosched()
				case subscriber := <-router.subscribeC:
					router.subscribe(subscriber.route)
					subscriber.doneC <- true
				case unsubscriber := <-router.unsubscribeC:
					router.unsubscribe(unsubscriber.route)
					unsubscriber.doneC <- true
				case <-router.stopC:
					router.stopping = true
				}
			}()
		}
	}()

	return nil
}

// Stop stops the router by closing the stop channel, and waiting on the WaitGroup
func (router *router) Stop() error {
	logger.Debug("Stopping router")

	router.stopC <- true
	router.wg.Wait()
	return nil
}

func (router *router) Check() error {
	if router.accessManager == nil || router.messageStore == nil || router.kvStore == nil {
		logger.WithFields(log.Fields{
			"err": ErrServiceNotProvided,
		}).Error("Some mandatory services are not provided")
		return ErrServiceNotProvided
	}

	err := router.messageStore.Check()
	if err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("MessageStore check failed")
		return err
	}

	err = router.kvStore.Check()
	if err != nil {
		log.WithFields(log.Fields{
			"module": "router",
			"err":    err,
		}).Error("KVStore check failed")
		return err
	}

	return nil
}

// HandleMessage assigns the new message id, stores the message and passes it to: the internal channel, and asynchronously to the cluster (if available).
func (router *router) HandleMessage(message *protocol.Message) error {
	logger.WithFields(log.Fields{
		"userID": message.UserID,
		"path":   message.Path,
	}).Debug("HandleMessage")

	mTotalMessagesIncoming.Add(1)
	if err := router.isStopping(); err != nil {
		logger.WithFields(log.Fields{
			"err": err,
		}).Error("Router is stopping")
		return err
	}

	if !router.accessManager.IsAllowed(auth.WRITE, message.UserID, message.Path) {
		return &PermissionDeniedError{message.UserID, auth.WRITE, message.Path}
	}

	txCallback := func(msgId uint64) []byte {
		message.ID = msgId
		message.Time = time.Now().Unix()
		return message.Bytes()
	}
	lenMsg := int64(len(message.Bytes()))
	mTotalMessagesIncomingBytes.Add(lenMsg)

	if err := router.messageStore.StoreTx(message.Path.Partition(), txCallback); err != nil {
		logger.WithFields(log.Fields{
			"err":          err,
			"msgPartition": message.Path.Partition(),
		}).Error("Error storing message in partition")
		mTotalMessageStoreErrors.Add(1)
		return err
	}
	mTotalMessagesStoredBytes.Add(lenMsg)

	if float32(len(router.handleC))/float32(cap(router.handleC)) > overloadedHandleChannelRatio {
		logger.WithFields(log.Fields{
			"currentLength": len(router.handleC),
			"maxCapacity":   cap(router.handleC),
		}).Warn("handleC channel is almost full")
		mTotalOverloadedHandleChannel.Add(1)
	}

	router.handleC <- message

	if router.cluster != nil {
		go router.cluster.BroadcastMessage(message)
	}

	return nil
}

// Subscribe adds a route to the subscribers. If there is already a route with same Application Id and Path, it will be replaced.
func (router *router) Subscribe(r *Route) (*Route, error) {

	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Subscribe")

	if err := router.isStopping(); err != nil {
		return nil, err
	}

	accessAllowed := router.accessManager.IsAllowed(auth.READ, r.UserID, r.Path)
	if !accessAllowed {
		return r, &PermissionDeniedError{UserID: r.UserID, AccessType: auth.READ, Path: r.Path}
	}
	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.subscribeC <- req
	<-req.doneC
	return r, nil
}

func (router *router) Unsubscribe(r *Route) {
	logger.WithFields(log.Fields{
		"accessManager": router.accessManager,
		"userID":        r.UserID,
		"path":          r.Path,
	}).Debug("Unsubscribe")

	req := subRequest{
		route: r,
		doneC: make(chan bool),
	}
	router.unsubscribeC <- req
	<-req.doneC
}

func (router *router) subscribe(r *Route) {
	logger.WithFields(log.Fields{
		"userID": r.UserID,
		"path":   r.Path,
	}).Debug("Internal subscribe")
	mTotalSubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	var removed bool
	if present {
		// Try to remove, to avoid double subscriptions of the same app
		slice, removed = removeIfMatching(slice, r)
	} else {
		// Path not present yet. Initialize the slice
		slice = make([]*Route, 0, 1)
		router.routes[r.Path] = slice
		mCurrentRoutes.Add(1)
	}
	router.routes[r.Path] = append(slice, r)
	if removed {
		mTotalDuplicateSubscriptionsAttempts.Add(1)
	} else {
		mTotalSubscriptions.Add(1)
		mCurrentSubscriptions.Add(1)
	}
}

func (router *router) unsubscribe(r *Route) {

	logger.WithFields(log.Fields{
		"userID": r.UserID,
		"path":   r.Path,
	}).Debug("Internal unsubscribe")
	mTotalUnsubscriptionAttempts.Add(1)

	slice, present := router.routes[r.Path]
	if !present {
		mTotalInvalidTopicOnUnsubscriptionAttempts.Add(1)
		return
	}
	var removed bool
	router.routes[r.Path], removed = removeIfMatching(slice, r)
	if removed {
		mTotalUnsubscriptions.Add(1)
		mCurrentSubscriptions.Add(-1)
	} else {
		mTotalInvalidUnsubscriptionAttempts.Add(1)
	}
	if len(router.routes[r.Path]) == 0 {
		delete(router.routes, r.Path)
		mCurrentRoutes.Add(-1)
	}
}

func (router *router) panicIfInternalDependenciesAreNil() {
	if router.accessManager == nil || router.kvStore == nil || router.messageStore == nil {
		panic(fmt.Sprintf("router: the internal dependencies marked with `true` are not set: AccessManager=%v, KVStore=%v, MessageStore=%v",
			router.accessManager == nil, router.kvStore == nil, router.messageStore == nil))
	}
}

func (router *router) channelsAreEmpty() bool {
	return len(router.handleC) == 0 && len(router.subscribeC) == 0 && len(router.unsubscribeC) == 0
}

func (router *router) isStopping() error {
	if router.stopping {
		return &ModuleStoppingError{"Router"}
	}
	return nil
}

func (router *router) routeMessage(message *protocol.Message) {

	logger.WithFields(log.Fields{
		"msgMetadata": message.Metadata(),
	}).Debug("Called routeMessage")
	mTotalMessagesRouted.Add(1)

	for path, list := range router.routes {
		if matchesTopic(message.Path, path) {
			for _, route := range list {
				router.deliverMessage(route, message)
			}
		} else {
			mTotalMessagesNotMatchingTopic.Add(1)
		}
	}
}

func (router *router) deliverMessage(route *Route, message *protocol.Message) {
	defer protocol.PanicLogger()

	select {
	case route.MessagesChannel() <- &MessageForRoute{Message: message, Route: route}:
	// fine, we could send the message
	default:
		logger.WithFields(log.Fields{
			"route": route.String(),
		}).Warn("deliverMessage: queue was full, unsubscribing and closing delivery channel for route")
		router.unsubscribe(route)
		route.Close()
		mTotalDeliverMessageErrors.Add(1)
	}
}

func (router *router) closeRoutes() {
	logger.Debug("Called closeRoutes")

	for _, currentRouteList := range router.routes {
		for _, route := range currentRouteList {
			router.unsubscribe(route)
			log.WithFields(log.Fields{
				"module": "router",
				"route":  route.String(),
			}).Debug("Closing route for ")
			route.Close()
		}
	}
}

// matchesTopic checks whether the supplied routePath matches the message topic
func matchesTopic(messagePath, routePath protocol.Path) bool {
	messagePathLen := len(string(messagePath))
	routePathLen := len(string(routePath))
	return strings.HasPrefix(string(messagePath), string(routePath)) &&
		(messagePathLen == routePathLen ||
			(messagePathLen > routePathLen && string(messagePath)[routePathLen] == '/'))
}

// removeIfMatching removes a route from the supplied list, based on same ApplicationID id and same path (if existing)
// returns: the (possibly updated) slide, and a boolean value (true if route was removed, false otherwise)
func removeIfMatching(slice []*Route, route *Route) ([]*Route, bool) {
	position := -1
	for p, r := range slice {
		if r.ApplicationID == route.ApplicationID && r.Path == route.Path {
			position = p
		}
	}
	if position == -1 {
		return slice, false
	}
	return append(slice[:position], slice[position+1:]...), true
}

// AccessManager returns the `accessManager` provided for the router
func (router *router) AccessManager() (auth.AccessManager, error) {
	if router.accessManager == nil {
		return nil, ErrServiceNotProvided
	}
	return router.accessManager, nil
}

// MessageStore returns the `messageStore` provided for the router
func (router *router) MessageStore() (store.MessageStore, error) {
	if router.messageStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.messageStore, nil
}

// KVStore returns the `kvStore` provided for the router
func (router *router) KVStore() (store.KVStore, error) {
	if router.kvStore == nil {
		return nil, ErrServiceNotProvided
	}
	return router.kvStore, nil
}

// Cluster returns the `cluster` provided for the router, or nil if no cluster was set-up
func (router *router) Cluster() *cluster.Cluster {
	return router.cluster
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/manifest/schema1"
	humanize "github.com/dustin/go-humanize"
	"github.com/jessfraz/reg/clair"
	"github.com/jessfraz/reg/registry"
	"github.com/jessfraz/reg/utils"
	"github.com/urfave/cli"
)

const (
	// VERSION is the binary version.
	VERSION = "v0.1.0"

	dockerConfigPath = ".docker/config.json"
)

var (
	updating = false
	wg       sync.WaitGroup
)

// preload initializes any global options and configuration
// before the main or sub commands are run.
func preload(c *cli.Context) (err error) {
	if c.GlobalBool("debug") {
		logrus.SetLevel(logrus.DebugLevel)
	}

	return nil
}

func main() {
	app := cli.NewApp()
	app.Name = "reg-server"
	app.Version = VERSION
	app.Author = "@jessfraz"
	app.Email = "no-reply@butts.com"
	app.Usage = "Docker registry v2 static UI server."
	app.Before = preload
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "run in debug mode",
		},
		cli.StringFlag{
			Name:  "username, u",
			Usage: "username for the registry",
		},
		cli.StringFlag{
			Name:  "password, p",
			Usage: "password for the registry",
		},
		cli.StringFlag{
			Name:  "registry, r",
			Usage: "URL to the provate registry (ex. r.j3ss.co)",
		},
		cli.StringFlag{
			Name:  "port",
			Value: "8080",
			Usage: "port for server to run on",
		},
		cli.StringFlag{
			Name:  "cert",
			Usage: "path to ssl cert",
		},
		cli.StringFlag{
			Name:  "key",
			Usage: "path to ssl key",
		},
		cli.StringFlag{
			Name:  "interval",
			Value: "5m",
			Usage: "interval to generate new index.html's at",
		},
		cli.StringFlag{
			Name:  "clair",
			Usage: "url to clair instance",
		},
	}
	app.Action = func(c *cli.Context) error {
		auth, err := utils.GetAuthConfig(c)
		if err != nil {
			return err
		}

		// create the registry client
		r, err := registry.New(auth, c.GlobalBool("debug"))
		if err != nil {
			return err
		}

		// get the path to the static directory
		wd, err := os.Getwd()
		if err != nil {
			return err
		}
		staticDir := filepath.Join(wd, "static")

		// create the initial index
		if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
			return err
		}

		// parse the duration
		dur, err := time.ParseDuration(c.String("interval"))
		if err != nil {
			logrus.Fatalf("parsing %s as duration failed: %v", c.String("interval"), err)
		}
		ticker := time.NewTicker(dur)

		go func() {
			// create more indexes every X minutes based off interval
			for range ticker.C {
				if !updating {
					if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
						logrus.Warnf("creating static index failed: %v", err)
						wg.Wait()
						updating = false
					}
				}
			}
		}()

		// create mux server
		mux := http.NewServeMux()

		// static files handler
		staticHandler := http.FileServer(http.Dir(staticDir))
		mux.Handle("/", staticHandler)

		// set up the server
		port := c.String("port")
		server := &http.Server{
			Addr:    ":" + port,
			Handler: mux,
		}
		logrus.Infof("Starting server on port %q", port)
		if c.String("cert") != "" && c.String("key") != "" {
			logrus.Fatal(server.ListenAndServeTLS(c.String("cert"), c.String("key")))
		} else {
			logrus.Fatal(server.ListenAndServe())
		}

		return nil
	}

	app.Run(os.Args)
}

type data struct {
	RegistryURL string
	LastUpdated string
	Repos       []repository
}

type repository struct {
	Name        string
	Tag         string
	RepoURI     string
	CreatedDate string
	VulnURI     string
}

type v1Compatibility struct {
	ID      string    `json:"id"`
	Created time.Time `json:"created"`
}

func createStaticIndex(r *registry.Registry, staticDir, clairURI string) error {
	updating = true
	logrus.Info("fetching catalog")
	repoList, err := r.Catalog("")
	if err != nil {
		return fmt.Errorf("getting catalog failed: %v", err)
	}

	logrus.Info("fetching tags")
	var repos []repository
	for i, repo := range repoList {
		// get the tags
		tags, err := r.Tags(repo)
		if err != nil {
			return fmt.Errorf("getting tags for %s failed: %v", repo, err)
		}
		for j, tag := range tags {
			// get the manifest

			manifest, err := r.Manifest(repo, tag)
			if err != nil {
				return fmt.Errorf("getting tags for %s:%s failed: %v", repo, tag, err)
			}

			var createdDate string
			if m1, ok := manifest.(schema1.SignedManifest); ok {
				history := m1.History
				for _, h := range history {
					var comp v1Compatibility
					if err := json.Unmarshal([]byte(h.V1Compatibility), &comp); err != nil {
						return fmt.Errorf("unmarshal v1compatibility failed: %v", err)
					}
					createdDate = humanize.Time(comp.Created)
					break
				}
			}

			repoURI := fmt.Sprintf("%s/%s", r.Domain, repo)
			if tag != "latest" {
				repoURI += ":" + tag
			}

			newrepo := repository{
				Name:        repo,
				Tag:         tag,
				RepoURI:     repoURI,
				CreatedDate: createdDate,
			}

			if clairURI != "" {
				wg.Add(1)

				go func(repo, tag string, i, j int) {
					defer wg.Done()

					throttle := time.Tick(time.Duration(time.Duration((i+1)*(j+1)) * time.Second))
					<-throttle

					logrus.Infof("creating vulns.txt for %s:%s", repo, tag)

					if err := createVulnStaticPage(r, staticDir, clairURI, repo, tag); err != nil {
						// return fmt.Errorf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
						logrus.Warnf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
					}
				}(repo, tag, i, j)

				newrepo.VulnURI = filepath.Join(repo, tag, "vulns.txt")
			}
			repos = append(repos, newrepo)
		}
	}

	// create temporoary file to save template to
	logrus.Info("creating temporary file for template")
	f, err := ioutil.TempFile("", "reg-server")
	if err != nil {
		return fmt.Errorf("creating temp file failed: %v", err)
	}
	defer f.Close()
	defer os.Remove(f.Name())

	// parse & execute the template
	logrus.Info("parsing and executing the template")
	templateDir := filepath.Join(staticDir, "../templates")
	lp := filepath.Join(templateDir, "layout.html")

	d := data{
		RegistryURL: r.Domain,
		Repos:       repos,
		LastUpdated: time.Now().Local().Format(time.RFC1123),
	}
	tmpl := template.Must(template.New("").ParseFiles(lp))
	if err := tmpl.ExecuteTemplate(f, "layout", d); err != nil {
		return fmt.Errorf("execute template failed: %v", err)
	}
	f.Close()

	index := filepath.Join(staticDir, "index.html")
	logrus.Infof("renaming the temporary file %s to %s", f.Name(), index)
	if err := os.Rename(f.Name(), index); err != nil {
		return fmt.Errorf("renaming result from %s to %s failed: %v", f.Name(), index, err)
	}
	updating = false
	return nil
}

func createVulnStaticPage(r *registry.Registry, staticDir, clairURI, repo, tag string) error {
	// get the manifest
	m, err := r.ManifestV1(repo, tag)
	if err != nil {
		return err
	}

	// filter out the empty layers
	var filteredLayers []schema1.FSLayer
	for _, layer := range m.FSLayers {
		if layer.BlobSum != clair.EmptyLayerBlobSum {
			filteredLayers = append(filteredLayers, layer)
		}
	}
	m.FSLayers = filteredLayers
	if len(m.FSLayers) == 0 {
		fmt.Printf("No need to analyse image %s:%s as there is no non-emtpy layer", repo, tag)
		return nil
	}

	// initialize clair
	cr, err := clair.New(clairURI, false)
	if err != nil {
		return err
	}

	for i := len(m.FSLayers) - 1; i >= 0; i-- {
		// form the clair layer
		l, err := utils.NewClairLayer(r, repo, m.FSLayers, i)
		if err != nil {
			return err
		}

		// post the layer
		if _, err := cr.PostLayer(l); err != nil {
			return err
		}
	}

	vl, err := cr.GetLayer(m.FSLayers[0].BlobSum.String(), false, true)
	if err != nil {
		return err
	}

	// get the vulns
	var vulns []clair.Vulnerability
	for _, f := range vl.Features {
		for _, v := range f.Vulnerabilities {
			vulns = append(vulns, v)
		}
	}

	path := filepath.Join(staticDir, repo, tag, "vulns.txt")
	if err := os.MkdirAll(filepath.Dir(path), 0644); err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	fmt.Fprintf(file, "Found %d vulnerabilities \n", len(vulns))

	vulnsBy := func(sev string, store map[string][]clair.Vulnerability) []clair.Vulnerability {
		items, found := store[sev]
		if !found {
			items = make([]clair.Vulnerability, 0)
			store[sev] = items
		}
		return items
	}

	// group by severity
	store := make(map[string][]clair.Vulnerability)
	for _, v := range vulns {
		sevRow := vulnsBy(v.Severity, store)
		store[v.Severity] = append(sevRow, v)
	}

	// iterate over the priorities list
	iteratePriorities := func(f func(sev string)) {
		for _, sev := range clair.Priorities {
			if len(store[sev]) != 0 {
				f(sev)
			}
		}
	}

	iteratePriorities(func(sev string) {
		fmt.Fprintf(file, "%s: %d\n", sev, len(store[sev]))
	})
	fmt.Fprintln(file, "")

	// return an error if there are more than 10 bad vulns
	lenBadVulns := len(store["High"]) + len(store["Critical"]) + len(store["Defcon1"])
	if lenBadVulns > 10 {
		fmt.Fprintln(file, "--------------- ALERT ---------------")
		fmt.Fprintf(file, "%d bad vunerabilities found", lenBadVulns)
	}
	fmt.Fprintln(file, "")

	iteratePriorities(func(sev string) {
		for _, v := range store[sev] {
			fmt.Fprintf(file, "%s: [%s] \n%s\n%s\n", v.Name, v.Severity, v.Description, v.Link)
			fmt.Fprintln(file, "-----------------------------------------")
		}
	})

	return nil
}
<commit_msg>return<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io/ioutil"
	"net/http"
	"os"
	"path/filepath"
	"sync"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/docker/distribution/manifest/schema1"
	humanize "github.com/dustin/go-humanize"
	"github.com/jessfraz/reg/clair"
	"github.com/jessfraz/reg/registry"
	"github.com/jessfraz/reg/utils"
	"github.com/urfave/cli"
)

const (
	// VERSION is the binary version.
	VERSION = "v0.1.0"

	dockerConfigPath = ".docker/config.json"
)

var (
	updating = false
	wg       sync.WaitGroup
)

// preload initializes any global options and configuration
// before the main or sub commands are run.
func preload(c *cli.Context) (err error) {
	if c.GlobalBool("debug") {
		logrus.SetLevel(logrus.DebugLevel)
	}

	return nil
}

func main() {
	app := cli.NewApp()
	app.Name = "reg-server"
	app.Version = VERSION
	app.Author = "@jessfraz"
	app.Email = "no-reply@butts.com"
	app.Usage = "Docker registry v2 static UI server."
	app.Before = preload
	app.Flags = []cli.Flag{
		cli.BoolFlag{
			Name:  "debug, d",
			Usage: "run in debug mode",
		},
		cli.StringFlag{
			Name:  "username, u",
			Usage: "username for the registry",
		},
		cli.StringFlag{
			Name:  "password, p",
			Usage: "password for the registry",
		},
		cli.StringFlag{
			Name:  "registry, r",
			Usage: "URL to the provate registry (ex. r.j3ss.co)",
		},
		cli.StringFlag{
			Name:  "port",
			Value: "8080",
			Usage: "port for server to run on",
		},
		cli.StringFlag{
			Name:  "cert",
			Usage: "path to ssl cert",
		},
		cli.StringFlag{
			Name:  "key",
			Usage: "path to ssl key",
		},
		cli.StringFlag{
			Name:  "interval",
			Value: "5m",
			Usage: "interval to generate new index.html's at",
		},
		cli.StringFlag{
			Name:  "clair",
			Usage: "url to clair instance",
		},
	}
	app.Action = func(c *cli.Context) error {
		auth, err := utils.GetAuthConfig(c)
		if err != nil {
			return err
		}

		// create the registry client
		r, err := registry.New(auth, c.GlobalBool("debug"))
		if err != nil {
			return err
		}

		// get the path to the static directory
		wd, err := os.Getwd()
		if err != nil {
			return err
		}
		staticDir := filepath.Join(wd, "static")

		// create the initial index
		if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
			return err
		}

		// parse the duration
		dur, err := time.ParseDuration(c.String("interval"))
		if err != nil {
			logrus.Fatalf("parsing %s as duration failed: %v", c.String("interval"), err)
		}
		ticker := time.NewTicker(dur)

		go func() {
			// create more indexes every X minutes based off interval
			for range ticker.C {
				if !updating {
					if err := createStaticIndex(r, staticDir, c.GlobalString("clair")); err != nil {
						logrus.Warnf("creating static index failed: %v", err)
						wg.Wait()
						updating = false
					}
				}
			}
		}()

		// create mux server
		mux := http.NewServeMux()

		// static files handler
		staticHandler := http.FileServer(http.Dir(staticDir))
		mux.Handle("/", staticHandler)

		// set up the server
		port := c.String("port")
		server := &http.Server{
			Addr:    ":" + port,
			Handler: mux,
		}
		logrus.Infof("Starting server on port %q", port)
		if c.String("cert") != "" && c.String("key") != "" {
			logrus.Fatal(server.ListenAndServeTLS(c.String("cert"), c.String("key")))
		} else {
			logrus.Fatal(server.ListenAndServe())
		}

		return nil
	}

	app.Run(os.Args)
}

type data struct {
	RegistryURL string
	LastUpdated string
	Repos       []repository
}

type repository struct {
	Name        string
	Tag         string
	RepoURI     string
	CreatedDate string
	VulnURI     string
}

type v1Compatibility struct {
	ID      string    `json:"id"`
	Created time.Time `json:"created"`
}

func createStaticIndex(r *registry.Registry, staticDir, clairURI string) error {
	updating = true
	logrus.Info("fetching catalog")
	repoList, err := r.Catalog("")
	if err != nil {
		return fmt.Errorf("getting catalog failed: %v", err)
	}

	logrus.Info("fetching tags")
	var repos []repository
	for i, repo := range repoList {
		// get the tags
		tags, err := r.Tags(repo)
		if err != nil {
			return fmt.Errorf("getting tags for %s failed: %v", repo, err)
		}
		for j, tag := range tags {
			// get the manifest

			manifest, err := r.Manifest(repo, tag)
			if err != nil {
				return fmt.Errorf("getting tags for %s:%s failed: %v", repo, tag, err)
			}

			var createdDate string
			if m1, ok := manifest.(schema1.SignedManifest); ok {
				history := m1.History
				for _, h := range history {
					var comp v1Compatibility
					if err := json.Unmarshal([]byte(h.V1Compatibility), &comp); err != nil {
						return fmt.Errorf("unmarshal v1compatibility failed: %v", err)
					}
					createdDate = humanize.Time(comp.Created)
					break
				}
			}

			repoURI := fmt.Sprintf("%s/%s", r.Domain, repo)
			if tag != "latest" {
				repoURI += ":" + tag
			}

			newrepo := repository{
				Name:        repo,
				Tag:         tag,
				RepoURI:     repoURI,
				CreatedDate: createdDate,
			}

			if clairURI != "" {
				wg.Add(1)

				go func(repo, tag string, i, j int) {
					defer wg.Done()

					throttle := time.Tick(time.Duration(time.Duration((i+1)*(j+1)) * time.Second))
					<-throttle

					logrus.Infof("creating vulns.txt for %s:%s", repo, tag)

					if err := createVulnStaticPage(r, staticDir, clairURI, repo, tag); err != nil {
						// return fmt.Errorf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
						logrus.Warnf("creating vuln static page for %s:%s failed: %v", repo, tag, err)
					}
				}(repo, tag, i, j)

				newrepo.VulnURI = filepath.Join(repo, tag, "vulns.txt")
			}
			repos = append(repos, newrepo)
		}
	}

	// create temporoary file to save template to
	logrus.Info("creating temporary file for template")
	f, err := ioutil.TempFile("", "reg-server")
	if err != nil {
		return fmt.Errorf("creating temp file failed: %v", err)
	}
	defer f.Close()
	defer os.Remove(f.Name())

	// parse & execute the template
	logrus.Info("parsing and executing the template")
	templateDir := filepath.Join(staticDir, "../templates")
	lp := filepath.Join(templateDir, "layout.html")

	d := data{
		RegistryURL: r.Domain,
		Repos:       repos,
		LastUpdated: time.Now().Local().Format(time.RFC1123),
	}
	tmpl := template.Must(template.New("").ParseFiles(lp))
	if err := tmpl.ExecuteTemplate(f, "layout", d); err != nil {
		return fmt.Errorf("execute template failed: %v", err)
	}
	f.Close()

	index := filepath.Join(staticDir, "index.html")
	logrus.Infof("renaming the temporary file %s to %s", f.Name(), index)
	if err := os.Rename(f.Name(), index); err != nil {
		return fmt.Errorf("renaming result from %s to %s failed: %v", f.Name(), index, err)
	}
	updating = false
	return nil
}

func createVulnStaticPage(r *registry.Registry, staticDir, clairURI, repo, tag string) error {
	// get the manifest
	m, err := r.ManifestV1(repo, tag)
	if err != nil {
		return err
	}

	// filter out the empty layers
	var filteredLayers []schema1.FSLayer
	for _, layer := range m.FSLayers {
		if layer.BlobSum != clair.EmptyLayerBlobSum {
			filteredLayers = append(filteredLayers, layer)
		}
	}
	m.FSLayers = filteredLayers
	if len(m.FSLayers) == 0 {
		fmt.Printf("No need to analyse image %s:%s as there is no non-emtpy layer", repo, tag)
		return nil
	}

	// initialize clair
	cr, err := clair.New(clairURI, false)
	if err != nil {
		return err
	}

	for i := len(m.FSLayers) - 1; i >= 0; i-- {
		// form the clair layer
		l, err := utils.NewClairLayer(r, repo, m.FSLayers, i)
		if err != nil {
			return err
		}

		// post the layer
		if _, err := cr.PostLayer(l); err != nil {
			return err
		}
	}

	vl, err := cr.GetLayer(m.FSLayers[0].BlobSum.String(), false, true)
	if err != nil {
		return err
	}

	// get the vulns
	var vulns []clair.Vulnerability
	for _, f := range vl.Features {
		for _, v := range f.Vulnerabilities {
			vulns = append(vulns, v)
		}
	}

	path := filepath.Join(staticDir, repo, tag, "vulns.txt")
	if err := os.MkdirAll(filepath.Dir(path), 0644); err != nil {
		return err
	}
	file, err := os.Create(path)
	if err != nil {
		return err
	}
	defer file.Close()

	fmt.Fprintf(file, "Found %d vulnerabilities \n", len(vulns))
	fmt.Fprintln(file, "")

	vulnsBy := func(sev string, store map[string][]clair.Vulnerability) []clair.Vulnerability {
		items, found := store[sev]
		if !found {
			items = make([]clair.Vulnerability, 0)
			store[sev] = items
		}
		return items
	}

	// group by severity
	store := make(map[string][]clair.Vulnerability)
	for _, v := range vulns {
		sevRow := vulnsBy(v.Severity, store)
		store[v.Severity] = append(sevRow, v)
	}

	// iterate over the priorities list
	iteratePriorities := func(f func(sev string)) {
		for _, sev := range clair.Priorities {
			if len(store[sev]) != 0 {
				f(sev)
			}
		}
	}

	iteratePriorities(func(sev string) {
		fmt.Fprintf(file, "%s: %d\n", sev, len(store[sev]))
	})
	fmt.Fprintln(file, "")

	// return an error if there are more than 10 bad vulns
	lenBadVulns := len(store["High"]) + len(store["Critical"]) + len(store["Defcon1"])
	if lenBadVulns > 10 {
		fmt.Fprintln(file, "--------------- ALERT ---------------")
		fmt.Fprintf(file, "%d bad vunerabilities found", lenBadVulns)
	}
	fmt.Fprintln(file, "")

	iteratePriorities(func(sev string) {
		for _, v := range store[sev] {
			fmt.Fprintf(file, "%s: [%s] \n%s\n%s\n", v.Name, v.Severity, v.Description, v.Link)
			fmt.Fprintln(file, "-----------------------------------------")
		}
	})

	return nil
}
<|endoftext|>


<commit_before>package server

import (
	"bytes"
	"fmt"
	"github.com/GrappigPanda/notorious/bencode"
	"gopkg.in/redis.v3"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"strings"
)

const (
	STARTED = iota
	COMPLETED
	STOPPED
)

type TorrentEvent struct {
	started   int
	completed int
	stopped   int
}

type TorrentRequestData struct {
	info_hash  string //20 byte sha1 hash
	peer_id    string //max len 20
	ip         string //optional
	port       string // port number the peer is listening on
	uploaded   int    // base10 ascii amount uploaded so far
	downloaded int    // base10 ascii amount downloaded so far
	left       int    // # of bytes left to download (base 10 ascii)
	event      int    // TorrentEvent
	numwant    int    // Number of peers requested by client.
	compact    bool   // Bep23 peer list compression decision: True -> compress bep23
}

var ANNOUNCE_URL = "/announce"

var FIELDS = []string{"port", "uploaded", "downloaded", "left", "event", "compact"}

func parseTorrentGetRequestURI(request string, ip string) map[string]interface{} {
	torrentdata := make(map[string]interface{})
	querydata := decodeQueryURL(request)

	lookupvalues := []string{
		"info_hash",
		"peer_id",
		// ip intentially left out
		"port",
		"uploaded",
		"downloaded",
		"left",
		"event",
		"numwant",
		"compact",
	}

	lookupints := []string{
		"uploaded",
		"downloaded",
		"left",
		"event",
		"numwant",
	}

	for i := range lookupvalues {
		index := lookupvalues[i]

		// TODO(ian): Delegate responsibility to another function.
		x := querydata[index]
		if len(x) <= 0 {
			continue
		}

		if index == "info_hash" {
			torrentdata[index] = ParseInfoHash(x[0])
		} else if index == "compact" {
			if querydata[index][0] == "1" {
				torrentdata[index] = true
			} else {
				torrentdata[index] = false
			}
		} else {
			torrentdata[index] = querydata[index][0]
		}

		// TODO(ian): Delegate responsibility here, as well.
		for j := range lookupints {
			if lookupints[j] == index {
				val, _ := strconv.Atoi(index)
				torrentdata[index] = val
			}
		}
	}

	// TODO(ian): Add a generic size lookup function and assert > 0
	x := strings.Split(ip, ":")
	if len(x) > 0 {
		torrentdata["ip"] = x[0]
	}

	fmt.Println(torrentdata)
	return torrentdata
}

func ParseInfoHash(s string) string {
	return fmt.Sprintf("%x", s)
}

func fillEmptyMapValues(torrentMap map[string]interface{}) *TorrentRequestData {
	_, ok := torrentMap["uploaded"]
	if !ok {
		torrentMap["uploaded"] = 0
	}

	_, ok = torrentMap["downloaded"]
	if !ok {
		torrentMap["downloaded"] = 0
	}

	_, ok = torrentMap["left"]
	if !ok {
		torrentMap["left"] = 0
	}

	_, ok = torrentMap["event"]
	if !ok {
		torrentMap["event"] = STOPPED
	}

	_, ok = torrentMap["compact"]
	if !ok {
		torrentMap["compact"] = true
	}

	_, ok = torrentMap["numwant"]
	if !ok {
		torrentMap["numwant"] = 30
	}

	x := TorrentRequestData{
		torrentMap["info_hash"].(string),
		torrentMap["peer_id"].(string),
		torrentMap["ip"].(string),
		torrentMap["port"].(string),
		torrentMap["uploaded"].(int),
		torrentMap["downloaded"].(int),
		torrentMap["left"].(int),
		torrentMap["event"].(int),
		torrentMap["numwant"].(int),
		torrentMap["compact"].(bool),
	}

	return &x
}

func worker(client *redis.Client, torrdata *TorrentRequestData) []string {
	if RedisGetBoolKeyVal(client, torrdata.info_hash, torrdata) {
		x := RedisGetKeyVal(client, torrdata.info_hash, torrdata)

		RedisSetKeyVal(client,
			concatenateKeyMember(torrdata.info_hash, "ip"),
			createIpPortPair(torrdata))

		return x

	} else {
		CreateNewTorrentKey(client, torrdata.info_hash, torrdata)
		return worker(client, torrdata)
	}
}
func compactIPPort(ip string, port string) []byte {
    res := bytes.NewBuffer(make([]byte, 0))

    intPort, err := strconv.Atoi(port)
    if err != nil {
        panic("failure1")
    }

    if err := binary.Write(res, binary.BigEndian, binary.BigEndian.Uint32(net.ParseIP(ip).To4())); err != nil {
        panic("failure0")
    }

    err = binary.Write(res, binary.BigEndian, uint16(intPort));
        if err != nil {
        panic("failure2")
    }

    return res.Bytes()
}

func CompactAllPeers(ipport []string) []byte {
    ret := bytes.NewBuffer(make([]byte, 0))
    for i := range ipport {
        sz := strings.Split(ipport[i], ":")
        ip := sz[0]
        port := sz[1]

        ret += compactIPPort(ip, port)
    }

    return ret.Bytes()
}

func formatResponseData(ips []string, torrentdata *TorrentRequestData) string {
	if torrentdata.compact {
		ips = compactaAllPeers(ips[i])
		return EncodeResponse(ips, torrentdata)
	} else {
		return EncodeResponse(ips, torrentdata)
	}
}

func decodeQueryURL(s string) url.Values {
	u, err := url.Parse(s)
	if err != nil {
		panic(err)
	}

	m, _ := url.ParseQuery(u.RawQuery)
	return m
}

func encodeKV(key string, value string) string {
	return fmt.Sprintf("%s%s", bencode.EncodeByteString(key), bencode.EncodeByteString(value))
}

func EncodeResponse(ipport []string, torrentdata *TorrentRequestData) string {
	ret := ""
	ret += encodeKV("interval", "30")
	ret += encodeKV("complete", "1")
	ipstr := ""

	if torrentdata.compact {
		count := 0

		for i := range ipport {
			ipstr += fmt.Sprintf("%s", ipport[i])
			count += 1
		}

		ret += encodeKV("incomplete", fmt.Sprintf("%d", count))
		ret += bencode.EncodeDictionary("peers", ipstr)
	} else {

		for i := range ipport {
			data := strings.Split(ipport[i], ":")

			ipstr += encodeKV("peer_id", "1")
			ipstr += encodeKV("ip", data[0])
			ipstr += encodeKV("port", data[1])
		}

		ret += bencode.EncodeDictionary("peers", ipstr)
	}

	return bencode.EncodeDictionary("", ret)
}

func requestHandler(w http.ResponseWriter, req *http.Request) {
	client := OpenClient()

	torrentdata := parseTorrentGetRequestURI(req.RequestURI, req.RemoteAddr)
	if torrentdata != nil {
		data := fillEmptyMapValues(torrentdata)

		worker(client, data)
		x := RedisGetKeyVal(client, data.info_hash, data)
		fmt.Println(x)

		response := formatResponseData(x, data)
		fmt.Println(response)

		w.Write([]byte(response))
	}
}

func RunServer() {
	mux := http.NewServeMux()

	mux.HandleFunc("/announce", requestHandler)
	http.ListenAndServe(":3000", mux)
}

func OpenClient() *redis.Client {
	client := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "",
		DB:       0,
	})

	return client
}

func CreateNewTorrentKey(client *redis.Client, key string, value *TorrentRequestData) {
	// CreateNewTorrentKey creates a new key. By default, it adds a member
	// ":ip". I don't think this ought to ever be generalized, as I just want
	// Redis to function in one specific way in notorious.

	// TODO(ian): You might want to set this explicitly in parameters
	// value := *TorrentRequestData
	client.SAdd(key, "ip")
}

func createIpPortPair(value *TorrentRequestData) string {
	// createIpPortPair creates a string formatted ("%s:%s", value.ip,
	// value.port) looking like so: "127.0.0.1:6886" and returns this value.
	return fmt.Sprintf("%s:%s", value.ip, value.port)
}

func RedisSetKeyVal(client *redis.Client, keymember string, value string) {
	// RedisSetKeyVal sets a key:member's value to value. Returns nothing as of
	// yet.
	client.SAdd(keymember, value)
}

func RedisGetKeyVal(client *redis.Client, key string, value *TorrentRequestData) []string {
	// RedisGetKeyVal retrieves a value from the Redis store by looking up the
	// provided key. If the key does not yet exist, we create the key in the KV
	// storage or if the value is empty, we add the current requester to the
	// list.
	keymember := concatenateKeyMember(key, "ip")

	val, err := client.SMembers(keymember).Result()
	if err != nil {
		// Fail because the key doesn't exist in the KV storage.
		CreateNewTorrentKey(client, keymember, value)
	}

	// If no keys yet exist in the KV storage.
	if len(val) == 0 {
		RedisSetKeyVal(client, keymember, createIpPortPair(value))
	}

	return val
}

func RedisGetBoolKeyVal(client *redis.Client, key string, value interface{}) bool {
	_, err := client.Get(key).Result()

	return err != nil
}

func concatenateKeyMember(key string, member string) string {
	var buffer bytes.Buffer

	buffer.WriteString(key)
	buffer.WriteString(":")
	buffer.WriteString(member)

	return buffer.String()
}
<commit_msg>Update server.go<commit_after>package server

import (
	"bytes"
	"fmt"
	"github.com/GrappigPanda/notorious/bencode"
	"gopkg.in/redis.v3"
	"net"
	"net/http"
	"net/url"
	"strconv"
	"binary"
	"strings"
)

const (
	STARTED = iota
	COMPLETED
	STOPPED
)

type TorrentEvent struct {
	started   int
	completed int
	stopped   int
}

type TorrentRequestData struct {
	info_hash  string //20 byte sha1 hash
	peer_id    string //max len 20
	ip         string //optional
	port       string // port number the peer is listening on
	uploaded   int    // base10 ascii amount uploaded so far
	downloaded int    // base10 ascii amount downloaded so far
	left       int    // # of bytes left to download (base 10 ascii)
	event      int    // TorrentEvent
	numwant    int    // Number of peers requested by client.
	compact    bool   // Bep23 peer list compression decision: True -> compress bep23
}

var ANNOUNCE_URL = "/announce"

var FIELDS = []string{"port", "uploaded", "downloaded", "left", "event", "compact"}

func parseTorrentGetRequestURI(request string, ip string) map[string]interface{} {
	torrentdata := make(map[string]interface{})
	querydata := decodeQueryURL(request)

	lookupvalues := []string{
		"info_hash",
		"peer_id",
		// ip intentially left out
		"port",
		"uploaded",
		"downloaded",
		"left",
		"event",
		"numwant",
		"compact",
	}

	lookupints := []string{
		"uploaded",
		"downloaded",
		"left",
		"event",
		"numwant",
	}

	for i := range lookupvalues {
		index := lookupvalues[i]

		// TODO(ian): Delegate responsibility to another function.
		x := querydata[index]
		if len(x) <= 0 {
			continue
		}

		if index == "info_hash" {
			torrentdata[index] = ParseInfoHash(x[0])
		} else if index == "compact" {
			if querydata[index][0] == "1" {
				torrentdata[index] = true
			} else {
				torrentdata[index] = false
			}
		} else {
			torrentdata[index] = querydata[index][0]
		}

		// TODO(ian): Delegate responsibility here, as well.
		for j := range lookupints {
			if lookupints[j] == index {
				val, _ := strconv.Atoi(index)
				torrentdata[index] = val
			}
		}
	}

	// TODO(ian): Add a generic size lookup function and assert > 0
	x := strings.Split(ip, ":")
	if len(x) > 0 {
		torrentdata["ip"] = x[0]
	}

	fmt.Println(torrentdata)
	return torrentdata
}

func ParseInfoHash(s string) string {
	return fmt.Sprintf("%x", s)
}

func fillEmptyMapValues(torrentMap map[string]interface{}) *TorrentRequestData {
	_, ok := torrentMap["uploaded"]
	if !ok {
		torrentMap["uploaded"] = 0
	}

	_, ok = torrentMap["downloaded"]
	if !ok {
		torrentMap["downloaded"] = 0
	}

	_, ok = torrentMap["left"]
	if !ok {
		torrentMap["left"] = 0
	}

	_, ok = torrentMap["event"]
	if !ok {
		torrentMap["event"] = STOPPED
	}

	_, ok = torrentMap["compact"]
	if !ok {
		torrentMap["compact"] = true
	}

	_, ok = torrentMap["numwant"]
	if !ok {
		torrentMap["numwant"] = 30
	}

	x := TorrentRequestData{
		torrentMap["info_hash"].(string),
		torrentMap["peer_id"].(string),
		torrentMap["ip"].(string),
		torrentMap["port"].(string),
		torrentMap["uploaded"].(int),
		torrentMap["downloaded"].(int),
		torrentMap["left"].(int),
		torrentMap["event"].(int),
		torrentMap["numwant"].(int),
		torrentMap["compact"].(bool),
	}

	return &x
}

func worker(client *redis.Client, torrdata *TorrentRequestData) []string {
	if RedisGetBoolKeyVal(client, torrdata.info_hash, torrdata) {
		x := RedisGetKeyVal(client, torrdata.info_hash, torrdata)

		RedisSetKeyVal(client,
			concatenateKeyMember(torrdata.info_hash, "ip"),
			createIpPortPair(torrdata))

		return x

	} else {
		CreateNewTorrentKey(client, torrdata.info_hash, torrdata)
		return worker(client, torrdata)
	}
}
func compactIPPort(ip string, port string) []byte {
    res := bytes.NewBuffer(make([]byte, 0))

    intPort, err := strconv.Atoi(port)
    if err != nil {
        panic("failure1")
    }

    if err := binary.Write(res, binary.BigEndian, binary.BigEndian.Uint32(net.ParseIP(ip).To4())); err != nil {
        panic("failure0")
    }

    err = binary.Write(res, binary.BigEndian, uint16(intPort));
        if err != nil {
        panic("failure2")
    }

    return res.Bytes()
}

func CompactAllPeers(ipport []string) []byte {
    ret := bytes.NewBuffer(make([]byte, 0))
    for i := range ipport {
        sz := strings.Split(ipport[i], ":")
        ip := sz[0]
        port := sz[1]

        ret += compactIPPort(ip, port)
    }

    return ret.Bytes()
}

func formatResponseData(ips []string, torrentdata *TorrentRequestData) string {
	if torrentdata.compact {
		ips = compactaAllPeers(ips)
		return EncodeResponse(ips, torrentdata)
	} else {
		return EncodeResponse(ips, torrentdata)
	}
}

func decodeQueryURL(s string) url.Values {
	u, err := url.Parse(s)
	if err != nil {
		panic(err)
	}

	m, _ := url.ParseQuery(u.RawQuery)
	return m
}

func encodeKV(key string, value string) string {
	return fmt.Sprintf("%s%s", bencode.EncodeByteString(key), bencode.EncodeByteString(value))
}

func EncodeResponse(ipport []string, torrentdata *TorrentRequestData) string {
	ret := ""
	ret += encodeKV("interval", "30")
	ret += encodeKV("complete", "1")
	ipstr := ""

	if torrentdata.compact {
		count := 0

		for i := range ipport {
			ipstr += fmt.Sprintf("%s", ipport[i])
			count += 1
		}

		ret += encodeKV("incomplete", fmt.Sprintf("%d", count))
		ret += bencode.EncodeDictionary("peers", ipstr)
	} else {

		for i := range ipport {
			data := strings.Split(ipport[i], ":")

			ipstr += encodeKV("peer_id", "1")
			ipstr += encodeKV("ip", data[0])
			ipstr += encodeKV("port", data[1])
		}

		ret += bencode.EncodeDictionary("peers", ipstr)
	}

	return bencode.EncodeDictionary("", ret)
}

func requestHandler(w http.ResponseWriter, req *http.Request) {
	client := OpenClient()

	torrentdata := parseTorrentGetRequestURI(req.RequestURI, req.RemoteAddr)
	if torrentdata != nil {
		data := fillEmptyMapValues(torrentdata)

		worker(client, data)
		x := RedisGetKeyVal(client, data.info_hash, data)
		fmt.Println(x)

		response := formatResponseData(x, data)
		fmt.Println(response)

		w.Write([]byte(response))
	}
}

func RunServer() {
	mux := http.NewServeMux()

	mux.HandleFunc("/announce", requestHandler)
	http.ListenAndServe(":3000", mux)
}

func OpenClient() *redis.Client {
	client := redis.NewClient(&redis.Options{
		Addr:     "localhost:6379",
		Password: "",
		DB:       0,
	})

	return client
}

func CreateNewTorrentKey(client *redis.Client, key string, value *TorrentRequestData) {
	// CreateNewTorrentKey creates a new key. By default, it adds a member
	// ":ip". I don't think this ought to ever be generalized, as I just want
	// Redis to function in one specific way in notorious.

	// TODO(ian): You might want to set this explicitly in parameters
	// value := *TorrentRequestData
	client.SAdd(key, "ip")
}

func createIpPortPair(value *TorrentRequestData) string {
	// createIpPortPair creates a string formatted ("%s:%s", value.ip,
	// value.port) looking like so: "127.0.0.1:6886" and returns this value.
	return fmt.Sprintf("%s:%s", value.ip, value.port)
}

func RedisSetKeyVal(client *redis.Client, keymember string, value string) {
	// RedisSetKeyVal sets a key:member's value to value. Returns nothing as of
	// yet.
	client.SAdd(keymember, value)
}

func RedisGetKeyVal(client *redis.Client, key string, value *TorrentRequestData) []string {
	// RedisGetKeyVal retrieves a value from the Redis store by looking up the
	// provided key. If the key does not yet exist, we create the key in the KV
	// storage or if the value is empty, we add the current requester to the
	// list.
	keymember := concatenateKeyMember(key, "ip")

	val, err := client.SMembers(keymember).Result()
	if err != nil {
		// Fail because the key doesn't exist in the KV storage.
		CreateNewTorrentKey(client, keymember, value)
	}

	// If no keys yet exist in the KV storage.
	if len(val) == 0 {
		RedisSetKeyVal(client, keymember, createIpPortPair(value))
	}

	return val
}

func RedisGetBoolKeyVal(client *redis.Client, key string, value interface{}) bool {
	_, err := client.Get(key).Result()

	return err != nil
}

func concatenateKeyMember(key string, member string) string {
	var buffer bytes.Buffer

	buffer.WriteString(key)
	buffer.WriteString(":")
	buffer.WriteString(member)

	return buffer.String()
}
<|endoftext|>


<commit_before>// Package server provides a server implementation to connect network transport
// protocols and service business logic by defining server endpoints.
package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	kitendpoint "github.com/go-kit/kit/endpoint"
	kithttp "github.com/go-kit/kit/transport/http"
	"github.com/gorilla/mux"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tylerb/graceful"
	"golang.org/x/net/context"

	microerror "github.com/giantswarm/microkit/error"
	"github.com/giantswarm/microkit/logger"
)

// Config represents the configuration used to create a new server object.
type Config struct {
	// Dependencies.
	Endpoints    []Endpoint
	ErrorEncoder kithttp.ErrorEncoder
	Logger       logger.Logger
	RequestFuncs []kithttp.RequestFunc
	Router       *mux.Router

	// Settings.
	ListenAddress string
	ServiceName   string
}

// DefaultConfig provides a default configuration to create a new server object
// by best effort.
func DefaultConfig() Config {
	return Config{
		// Dependencies.
		Endpoints:    nil,
		ErrorEncoder: nil,
		Logger:       nil,
		RequestFuncs: nil,
		Router:       mux.NewRouter(),

		// Settings.
		ListenAddress: "http://127.0.0.1:8000",
		ServiceName:   "microkit",
	}
}

// New creates a new configured server object.
func New(config Config) (Server, error) {
	// Dependencies.
	if config.Endpoints == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "endpoints must not be empty")
	}
	if config.ErrorEncoder == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "error encoder must not be empty")
	}
	if config.Logger == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "logger must not be empty")
	}
	if config.RequestFuncs == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "request funcs must not be empty")
	}
	if config.Router == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "router must not be empty")
	}

	// Settings.
	if config.ListenAddress == "" {
		return nil, microerror.MaskAnyf(invalidConfigError, "listen address must not be empty")
	}
	if config.ServiceName == "" {
		return nil, microerror.MaskAnyf(invalidConfigError, "service name must not be empty")
	}

	listenURL, err := url.Parse(config.ListenAddress)
	if err != nil {
		return nil, microerror.MaskAnyf(invalidConfigError, err.Error())
	}

	newServer := &server{
		bootOnce:     sync.Once{},
		endpoints:    config.Endpoints,
		errorEncoder: config.ErrorEncoder,
		httpServer:   nil,
		listenURL:    listenURL,
		logger:       config.Logger,
		requestFuncs: config.RequestFuncs,
		router:       config.Router,
		serviceName:  config.ServiceName,
		shutdownOnce: sync.Once{},
	}

	return newServer, nil
}

// server manages the transport logic and endpoint registration.
type server struct {
	bootOnce     sync.Once
	endpoints    []Endpoint
	errorEncoder kithttp.ErrorEncoder
	httpServer   *graceful.Server
	listenURL    *url.URL
	logger       logger.Logger
	requestFuncs []kithttp.RequestFunc
	router       *mux.Router
	serviceName  string
	shutdownOnce sync.Once
}

func (s *server) Boot() {
	s.bootOnce.Do(func() {
		// Define our custom not found handler. Here we take care about logging,
		// metrics and a proper response.
		s.router.NotFoundHandler = http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Log the error and its message. This is really useful for debugging.
			errDomain := errorDomain(nil)
			errMessage := fmt.Sprintf("not found: %s %s", r.Method, r.URL.Path)
			errTrace := ""
			s.logger.Log("error", map[string]string{"domain": errDomain, "message": errMessage, "trace": errTrace})

			// This defered callback will be executed at the very end of the request.
			defer func(t time.Time) {
				endpointCode := strconv.Itoa(http.StatusNotFound)
				endpointMethod := strings.ToLower(r.Method)
				endpointName := "notfound"

				endpointTotal.WithLabelValues(endpointCode, endpointMethod, endpointName).Inc()
				endpointTime.WithLabelValues(endpointCode, endpointMethod, endpointName).Set(float64(time.Since(t) / time.Millisecond))

				errorTotal.WithLabelValues(errDomain).Inc()
			}(time.Now())

			// Write the actual response body.
			w.WriteHeader(http.StatusNotFound)
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"error": errMessage,
				"from":  s.ServiceName(),
			})
		}))

		// We go through all endpoints this server defines and register them to the
		// router.
		for _, e := range s.Endpoints() {
			decoder := e.Decoder()
			encoder := e.Encoder()
			endpoint := e.Endpoint()
			method := e.Method()
			middlewares := e.Middlewares()
			name := e.Name()
			path := e.Path()

			// Prepare the actual endpoint depending on the provided middlewares of the
			// endpoint implementation. There might be cases in which there are none or
			// only one middleware. The go-kit interface is not that nice so we need to
			// make it fit here.
			{
				if len(middlewares) == 1 {
					endpoint = kitendpoint.Chain(middlewares[0])(endpoint)
				}
				if len(middlewares) > 1 {
					endpoint = kitendpoint.Chain(middlewares[0], middlewares[1:]...)(endpoint)
				}
			}

			// Combine all options this server defines.
			options := []kithttp.ServerOption{
				kithttp.ServerBefore(s.RequestFuncs()...),
				kithttp.ServerErrorEncoder(s.ErrorEncoder()),
			}

			// Register all endpoints to the router depending on their HTTP methods and
			// request paths. The registered http.Handler is instrumented using
			// prometheus. We track counts of execution and duration it took to complete
			// the http.Handler.
			s.router.Methods(method).Path(path).Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctx := context.Background()

				// Here we create a new wrapper for the http.ResponseWriter of the current
				// request. We inject it into the called http.Handler so it can track the
				// status code we are interested in. It will help us gathering the
				// response status code after it was written by the underlying
				// http.ResponseWriter.
				responseWriterConfig := DefaultResponseWriterConfig()
				responseWriterConfig.ResponseWriter = w
				responseWriter, err := NewResponseWriter(responseWriterConfig)
				if err != nil {
					panic(err)
				}
				w = responseWriter

				// Here we define the metrics relevant labels. These will be used to
				// instrument the current request.
				endpointName := strings.Replace(name, "/", "_", -1)
				endpointMethod := strings.ToLower(method)
				endpointCode := responseWriter.StatusCode()

				// This defered callback will be executed at the very end of the request.
				// When it is executed we know all necessary information to instrument the
				// complete request, including its response status code.
				defer func(t time.Time) {
					s.logger.Log("code", endpointCode, "endpoint", name, "method", endpointMethod, "path", r.URL.Path)

					// At the time this code is executed the status code is properly set. So
					// we can use it for our instrumentation.
					endpointCode := strconv.Itoa(endpointCode)
					endpointTotal.WithLabelValues(endpointCode, endpointMethod, endpointName).Inc()
					endpointTime.WithLabelValues(endpointCode, endpointMethod, endpointName).Set(float64(time.Since(t) / time.Millisecond))
				}(time.Now())

				// Now we execute the actual endpoint handler.
				kithttp.NewServer(
					ctx,
					endpoint,
					decoder,
					encoder,
					options...,
				).ServeHTTP(w, r)

				// Here we now the status code.
				endpointCode = responseWriter.StatusCode()
			}))
		}

		// Register prometheus metrics endpoint.
		s.router.Path("/metrics").Handler(promhttp.Handler())

		// Register the router which has all of the configured custom endpoints
		// registered.
		s.httpServer = &graceful.Server{
			NoSignalHandling: true,
			Server: &http.Server{
				Addr:    s.listenURL.Host,
				Handler: s.router,
			},
			Timeout: 3 * time.Second,
		}

		go func() {
			err := s.httpServer.ListenAndServe()
			if err != nil {
				panic(err)
			}
		}()
	})
}

func (s *server) Endpoints() []Endpoint {
	return s.endpoints
}

func (s *server) ErrorEncoder() kithttp.ErrorEncoder {
	return func(ctx context.Context, err error, w http.ResponseWriter) {
		// At first we have to set the content type of the actual error response. If
		// we would set it at the end we would set a trailing header that would not
		// be recognized by most of the clients out there. This is because in the
		// next call to the errorEncoder below the client's implementation of the
		// errorEncoder probably writes the status code header, which marks the
		// beginning of trailing headers in HTTP.
		w.Header().Set("Content-Type", "application/json; charset=utf-8")

		// Create the microkit specific response error, which acts as error wrapper
		// within the client's error encoder. It is used to propagate response codes
		// and messages, so we can use them below.
		var responseError ResponseError
		{
			responseConfig := DefaultResponseErrorConfig()
			responseConfig.Underlying = err
			responseError, err = NewResponseError(responseConfig)
			if err != nil {
				panic(err)
			}
		}

		// Run the custom error encoder. This is used to let the implementing
		// microservice do something with errors occured during runtime. Things like
		// writing specific HTTP status codes to the given response writer can be
		// done.
		s.errorEncoder(ctx, responseError, w)

		// Log the error and its errgo trace. This is really useful for debugging.
		errDomain := errorDomain(err)
		errMessage := errorMessage(err)
		errTrace := errorTrace(err)
		s.logger.Log("error", map[string]string{"domain": errDomain, "message": errMessage, "trace": errTrace})

		// Emit metrics about the occured errors. That way we can feed our
		// instrumentation stack to have nice dashboards to get a picture about the
		// general system health.
		errorTotal.WithLabelValues(errDomain).Inc()

		// Write the actual response body.
		json.NewEncoder(w).Encode(map[string]interface{}{
			"code":  responseError.Code(),
			"error": responseError.Message(),
			"from":  s.ServiceName(),
		})
	}
}

func (s *server) RequestFuncs() []kithttp.RequestFunc {
	return s.requestFuncs
}

// NewRouter returns a HTTP handler for the server. Here we register all
// endpoints listed in the endpoint collection.
func (s *server) Router() *mux.Router {
	return s.router
}

func (s *server) ServiceName() string {
	return s.serviceName
}

func (s *server) Shutdown() {
	s.shutdownOnce.Do(func() {
		var wg sync.WaitGroup

		wg.Add(1)
		go func() {
			// Stop the HTTP server gracefully and wait some time for open connections
			// to be closed. Then force it to be stopped.
			s.httpServer.Stop(s.httpServer.Timeout)
			<-s.httpServer.StopChan()
			wg.Done()
		}()

		wg.Wait()
	})
}
<commit_msg>added response code to not found handler (#41)<commit_after>// Package server provides a server implementation to connect network transport
// protocols and service business logic by defining server endpoints.
package server

import (
	"encoding/json"
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"strings"
	"sync"
	"time"

	kitendpoint "github.com/go-kit/kit/endpoint"
	kithttp "github.com/go-kit/kit/transport/http"
	"github.com/gorilla/mux"
	"github.com/prometheus/client_golang/prometheus/promhttp"
	"github.com/tylerb/graceful"
	"golang.org/x/net/context"

	microerror "github.com/giantswarm/microkit/error"
	"github.com/giantswarm/microkit/logger"
)

// Config represents the configuration used to create a new server object.
type Config struct {
	// Dependencies.
	Endpoints    []Endpoint
	ErrorEncoder kithttp.ErrorEncoder
	Logger       logger.Logger
	RequestFuncs []kithttp.RequestFunc
	Router       *mux.Router

	// Settings.
	ListenAddress string
	ServiceName   string
}

// DefaultConfig provides a default configuration to create a new server object
// by best effort.
func DefaultConfig() Config {
	return Config{
		// Dependencies.
		Endpoints:    nil,
		ErrorEncoder: nil,
		Logger:       nil,
		RequestFuncs: nil,
		Router:       mux.NewRouter(),

		// Settings.
		ListenAddress: "http://127.0.0.1:8000",
		ServiceName:   "microkit",
	}
}

// New creates a new configured server object.
func New(config Config) (Server, error) {
	// Dependencies.
	if config.Endpoints == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "endpoints must not be empty")
	}
	if config.ErrorEncoder == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "error encoder must not be empty")
	}
	if config.Logger == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "logger must not be empty")
	}
	if config.RequestFuncs == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "request funcs must not be empty")
	}
	if config.Router == nil {
		return nil, microerror.MaskAnyf(invalidConfigError, "router must not be empty")
	}

	// Settings.
	if config.ListenAddress == "" {
		return nil, microerror.MaskAnyf(invalidConfigError, "listen address must not be empty")
	}
	if config.ServiceName == "" {
		return nil, microerror.MaskAnyf(invalidConfigError, "service name must not be empty")
	}

	listenURL, err := url.Parse(config.ListenAddress)
	if err != nil {
		return nil, microerror.MaskAnyf(invalidConfigError, err.Error())
	}

	newServer := &server{
		bootOnce:     sync.Once{},
		endpoints:    config.Endpoints,
		errorEncoder: config.ErrorEncoder,
		httpServer:   nil,
		listenURL:    listenURL,
		logger:       config.Logger,
		requestFuncs: config.RequestFuncs,
		router:       config.Router,
		serviceName:  config.ServiceName,
		shutdownOnce: sync.Once{},
	}

	return newServer, nil
}

// server manages the transport logic and endpoint registration.
type server struct {
	bootOnce     sync.Once
	endpoints    []Endpoint
	errorEncoder kithttp.ErrorEncoder
	httpServer   *graceful.Server
	listenURL    *url.URL
	logger       logger.Logger
	requestFuncs []kithttp.RequestFunc
	router       *mux.Router
	serviceName  string
	shutdownOnce sync.Once
}

func (s *server) Boot() {
	s.bootOnce.Do(func() {
		// Define our custom not found handler. Here we take care about logging,
		// metrics and a proper response.
		s.router.NotFoundHandler = http.Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			// Log the error and its message. This is really useful for debugging.
			errDomain := errorDomain(nil)
			errMessage := fmt.Sprintf("not found: %s %s", r.Method, r.URL.Path)
			errTrace := ""
			s.logger.Log("error", map[string]string{"domain": errDomain, "message": errMessage, "trace": errTrace})

			// This defered callback will be executed at the very end of the request.
			defer func(t time.Time) {
				endpointCode := strconv.Itoa(http.StatusNotFound)
				endpointMethod := strings.ToLower(r.Method)
				endpointName := "notfound"

				endpointTotal.WithLabelValues(endpointCode, endpointMethod, endpointName).Inc()
				endpointTime.WithLabelValues(endpointCode, endpointMethod, endpointName).Set(float64(time.Since(t) / time.Millisecond))

				errorTotal.WithLabelValues(errDomain).Inc()
			}(time.Now())

			// Write the actual response body.
			w.WriteHeader(http.StatusNotFound)
			w.Header().Set("Content-Type", "application/json; charset=utf-8")
			json.NewEncoder(w).Encode(map[string]interface{}{
				"code":  CodeResourceNotFound,
				"error": errMessage,
				"from":  s.ServiceName(),
			})
		}))

		// We go through all endpoints this server defines and register them to the
		// router.
		for _, e := range s.Endpoints() {
			decoder := e.Decoder()
			encoder := e.Encoder()
			endpoint := e.Endpoint()
			method := e.Method()
			middlewares := e.Middlewares()
			name := e.Name()
			path := e.Path()

			// Prepare the actual endpoint depending on the provided middlewares of the
			// endpoint implementation. There might be cases in which there are none or
			// only one middleware. The go-kit interface is not that nice so we need to
			// make it fit here.
			{
				if len(middlewares) == 1 {
					endpoint = kitendpoint.Chain(middlewares[0])(endpoint)
				}
				if len(middlewares) > 1 {
					endpoint = kitendpoint.Chain(middlewares[0], middlewares[1:]...)(endpoint)
				}
			}

			// Combine all options this server defines.
			options := []kithttp.ServerOption{
				kithttp.ServerBefore(s.RequestFuncs()...),
				kithttp.ServerErrorEncoder(s.ErrorEncoder()),
			}

			// Register all endpoints to the router depending on their HTTP methods and
			// request paths. The registered http.Handler is instrumented using
			// prometheus. We track counts of execution and duration it took to complete
			// the http.Handler.
			s.router.Methods(method).Path(path).Handler(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
				ctx := context.Background()

				// Here we create a new wrapper for the http.ResponseWriter of the current
				// request. We inject it into the called http.Handler so it can track the
				// status code we are interested in. It will help us gathering the
				// response status code after it was written by the underlying
				// http.ResponseWriter.
				responseWriterConfig := DefaultResponseWriterConfig()
				responseWriterConfig.ResponseWriter = w
				responseWriter, err := NewResponseWriter(responseWriterConfig)
				if err != nil {
					panic(err)
				}
				w = responseWriter

				// Here we define the metrics relevant labels. These will be used to
				// instrument the current request.
				endpointName := strings.Replace(name, "/", "_", -1)
				endpointMethod := strings.ToLower(method)
				endpointCode := responseWriter.StatusCode()

				// This defered callback will be executed at the very end of the request.
				// When it is executed we know all necessary information to instrument the
				// complete request, including its response status code.
				defer func(t time.Time) {
					s.logger.Log("code", endpointCode, "endpoint", name, "method", endpointMethod, "path", r.URL.Path)

					// At the time this code is executed the status code is properly set. So
					// we can use it for our instrumentation.
					endpointCode := strconv.Itoa(endpointCode)
					endpointTotal.WithLabelValues(endpointCode, endpointMethod, endpointName).Inc()
					endpointTime.WithLabelValues(endpointCode, endpointMethod, endpointName).Set(float64(time.Since(t) / time.Millisecond))
				}(time.Now())

				// Now we execute the actual endpoint handler.
				kithttp.NewServer(
					ctx,
					endpoint,
					decoder,
					encoder,
					options...,
				).ServeHTTP(w, r)

				// Here we now the status code.
				endpointCode = responseWriter.StatusCode()
			}))
		}

		// Register prometheus metrics endpoint.
		s.router.Path("/metrics").Handler(promhttp.Handler())

		// Register the router which has all of the configured custom endpoints
		// registered.
		s.httpServer = &graceful.Server{
			NoSignalHandling: true,
			Server: &http.Server{
				Addr:    s.listenURL.Host,
				Handler: s.router,
			},
			Timeout: 3 * time.Second,
		}

		go func() {
			err := s.httpServer.ListenAndServe()
			if err != nil {
				panic(err)
			}
		}()
	})
}

func (s *server) Endpoints() []Endpoint {
	return s.endpoints
}

func (s *server) ErrorEncoder() kithttp.ErrorEncoder {
	return func(ctx context.Context, err error, w http.ResponseWriter) {
		// At first we have to set the content type of the actual error response. If
		// we would set it at the end we would set a trailing header that would not
		// be recognized by most of the clients out there. This is because in the
		// next call to the errorEncoder below the client's implementation of the
		// errorEncoder probably writes the status code header, which marks the
		// beginning of trailing headers in HTTP.
		w.Header().Set("Content-Type", "application/json; charset=utf-8")

		// Create the microkit specific response error, which acts as error wrapper
		// within the client's error encoder. It is used to propagate response codes
		// and messages, so we can use them below.
		var responseError ResponseError
		{
			responseConfig := DefaultResponseErrorConfig()
			responseConfig.Underlying = err
			responseError, err = NewResponseError(responseConfig)
			if err != nil {
				panic(err)
			}
		}

		// Run the custom error encoder. This is used to let the implementing
		// microservice do something with errors occured during runtime. Things like
		// writing specific HTTP status codes to the given response writer can be
		// done.
		s.errorEncoder(ctx, responseError, w)

		// Log the error and its errgo trace. This is really useful for debugging.
		errDomain := errorDomain(err)
		errMessage := errorMessage(err)
		errTrace := errorTrace(err)
		s.logger.Log("error", map[string]string{"domain": errDomain, "message": errMessage, "trace": errTrace})

		// Emit metrics about the occured errors. That way we can feed our
		// instrumentation stack to have nice dashboards to get a picture about the
		// general system health.
		errorTotal.WithLabelValues(errDomain).Inc()

		// Write the actual response body.
		json.NewEncoder(w).Encode(map[string]interface{}{
			"code":  responseError.Code(),
			"error": responseError.Message(),
			"from":  s.ServiceName(),
		})
	}
}

func (s *server) RequestFuncs() []kithttp.RequestFunc {
	return s.requestFuncs
}

// NewRouter returns a HTTP handler for the server. Here we register all
// endpoints listed in the endpoint collection.
func (s *server) Router() *mux.Router {
	return s.router
}

func (s *server) ServiceName() string {
	return s.serviceName
}

func (s *server) Shutdown() {
	s.shutdownOnce.Do(func() {
		var wg sync.WaitGroup

		wg.Add(1)
		go func() {
			// Stop the HTTP server gracefully and wait some time for open connections
			// to be closed. Then force it to be stopped.
			s.httpServer.Stop(s.httpServer.Timeout)
			<-s.httpServer.StopChan()
			wg.Done()
		}()

		wg.Wait()
	})
}
<|endoftext|>