<commit_before>package serf

import (
	"bytes"
	"github.com/ugorji/go/codec"
)

// messageType are the types of gossip messages Serf will send along
// memberlist.
type messageType uint8

const (
	messageLeaveType messageType = iota
	messageRemoveFailedType
)

// messageLeave is the message broadcasted to signal the intentional to
// leave.
type messageLeave struct {
	Node string
}

// messageRemoveFailed is the message broadcasted to force remove
// a failed node from the failed node list and the member list.
type messageRemoveFailed struct {
	Node string
}

func decodeMessage(buf []byte, out interface{}) error {
	var handle codec.MsgpackHandle
	return codec.NewDecoder(bytes.NewBuffer(buf), &handle).Decode(out)
}

func encodeMessage(t messageType, msg interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	buf.WriteByte(uint8(t))

	handle := codec.MsgpackHandle{}
	encoder := codec.NewEncoder(buf, &handle)
	err := encoder.Encode(msg)
	return buf.Bytes(), err
}
<commit_msg>serf: adding LamportTime field to messages<commit_after>package serf

import (
	"bytes"
	"github.com/ugorji/go/codec"
)

// messageType are the types of gossip messages Serf will send along
// memberlist.
type messageType uint8

const (
	messageLeaveType messageType = iota
	messageRemoveFailedType
)

// messageLeave is the message broadcasted to signal the intentional to
// leave.
type messageLeave struct {
	LTime LamportTime
	Node  string
}

// messageRemoveFailed is the message broadcasted to force remove
// a failed node from the failed node list and the member list.
type messageRemoveFailed struct {
	LTime LamportTime
	Node  string
}

func decodeMessage(buf []byte, out interface{}) error {
	var handle codec.MsgpackHandle
	return codec.NewDecoder(bytes.NewBuffer(buf), &handle).Decode(out)
}

func encodeMessage(t messageType, msg interface{}) ([]byte, error) {
	buf := bytes.NewBuffer(nil)
	buf.WriteByte(uint8(t))

	handle := codec.MsgpackHandle{}
	encoder := codec.NewEncoder(buf, &handle)
	err := encoder.Encode(msg)
	return buf.Bytes(), err
}
<|endoftext|>


<commit_before>package workday

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

// DataDir is where the root directory
// where the data is stored
var DataDir = "data"

// Day is a single workday
type Day struct {
	Start time.Time
	End   time.Time
	Tasks []*Task
}

func pathForTime(t time.Time) string {
	year, month, day := t.Date()
	return filepath.Join(
		DataDir,
		"tasks",
		strconv.Itoa(year),
		month.String(),
		strconv.Itoa(day)+".json",
	)
}

func readDayFromFile(fpath string) (*Day, error) {
	f, err := os.Open(fpath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	d := new(Day)
	if err := json.NewDecoder(f).Decode(d); err != nil {
		return nil, err
	}
	return d, nil
}

// LoadDay reads the Day corresponding to the supplied time.
// If one is not found, a new one is created.
func LoadDay(t time.Time) (*Day, error) {
	fpath := pathForTime(t)
	if exists, err := fileExists(fpath); err != nil {
		return nil, err
	} else if exists {
		return readDayFromFile(fpath)
	}
	return &Day{t, t, make([]*Task, 0)}, nil
}

// Save writes the Day to the data directory as json
// and creates any required parent directories
func (d *Day) Save() error {
	d.MustBeSane()
	fpath := pathForTime(d.Start)
	if err := createParentDirs(fpath); err != nil {
		return err
	}
	f, err := os.Create(fpath)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := json.NewEncoder(f).Encode(d); err != nil {
		return err
	}
	return nil
}

// Ping updates the current Day's End time.
// A Day is created if it does not already exist.
func Ping() error {
	t := time.Now()
	d, err := LoadDay(t)
	if err != nil {
		return err
	}
	d.End = t
	return d.Save()
}

// MustBeSane panics if there are Day's data is inconsistent
func (d *Day) MustBeSane() {
	sYear, sMonth, sDay := d.Start.Date()
	eYear, eMonth, eDay := d.End.Date()
	if sYear != eYear {
		panic("years do not match")
	}
	if sMonth != eMonth {
		panic("months do not match")
	}
	if sDay != eDay {
		panic("days do not match")
	}
	if d.End.Before(d.Start) {
		panic("day ends before it starts")
	}
}

func roundDuration(d time.Duration) time.Duration {
	return time.Duration(d.Seconds()) * time.Second
}

func (d *Day) String() string {
	var buf bytes.Buffer

	fmt.Fprintf(&buf,
		"Date: %s\nTime: %s - %s (%s)\n\n",
		d.Start.Format("Mon Jan _2 2006"),
		d.Start.Format(time.Kitchen),
		d.End.Format(time.Kitchen),
		roundDuration(d.End.Sub(d.Start)),
	)
	for _, t := range d.Tasks {
		fmt.Fprintf(&buf, "\t- %s\n", t.Desc)
	}

	return buf.String()
}
<commit_msg>don't indent tasks<commit_after>package workday

import (
	"bytes"
	"encoding/json"
	"fmt"
	"os"
	"path/filepath"
	"strconv"
	"time"
)

// DataDir is where the root directory
// where the data is stored
var DataDir = "data"

// Day is a single workday
type Day struct {
	Start time.Time
	End   time.Time
	Tasks []*Task
}

func pathForTime(t time.Time) string {
	year, month, day := t.Date()
	return filepath.Join(
		DataDir,
		"tasks",
		strconv.Itoa(year),
		month.String(),
		strconv.Itoa(day)+".json",
	)
}

func readDayFromFile(fpath string) (*Day, error) {
	f, err := os.Open(fpath)
	if err != nil {
		return nil, err
	}
	defer f.Close()
	d := new(Day)
	if err := json.NewDecoder(f).Decode(d); err != nil {
		return nil, err
	}
	return d, nil
}

// LoadDay reads the Day corresponding to the supplied time.
// If one is not found, a new one is created.
func LoadDay(t time.Time) (*Day, error) {
	fpath := pathForTime(t)
	if exists, err := fileExists(fpath); err != nil {
		return nil, err
	} else if exists {
		return readDayFromFile(fpath)
	}
	return &Day{t, t, make([]*Task, 0)}, nil
}

// Save writes the Day to the data directory as json
// and creates any required parent directories
func (d *Day) Save() error {
	d.MustBeSane()
	fpath := pathForTime(d.Start)
	if err := createParentDirs(fpath); err != nil {
		return err
	}
	f, err := os.Create(fpath)
	if err != nil {
		return err
	}
	defer f.Close()
	if err := json.NewEncoder(f).Encode(d); err != nil {
		return err
	}
	return nil
}

// Ping updates the current Day's End time.
// A Day is created if it does not already exist.
func Ping() error {
	t := time.Now()
	d, err := LoadDay(t)
	if err != nil {
		return err
	}
	d.End = t
	return d.Save()
}

// MustBeSane panics if there are Day's data is inconsistent
func (d *Day) MustBeSane() {
	sYear, sMonth, sDay := d.Start.Date()
	eYear, eMonth, eDay := d.End.Date()
	if sYear != eYear {
		panic("years do not match")
	}
	if sMonth != eMonth {
		panic("months do not match")
	}
	if sDay != eDay {
		panic("days do not match")
	}
	if d.End.Before(d.Start) {
		panic("day ends before it starts")
	}
}

func roundDuration(d time.Duration) time.Duration {
	return time.Duration(d.Seconds()) * time.Second
}

func (d *Day) String() string {
	var buf bytes.Buffer

	fmt.Fprintf(&buf,
		"Date: %s\nTime: %s - %s (%s)\n\n",
		d.Start.Format("Mon Jan _2 2006"),
		d.Start.Format(time.Kitchen),
		d.End.Format(time.Kitchen),
		roundDuration(d.End.Sub(d.Start)),
	)
	for _, t := range d.Tasks {
		fmt.Fprintf(&buf, "%s\n", t.Desc)
	}

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


<commit_before>package server

import (
	"context"
	"database/sql"
	"encoding/binary"
	"io"
	"time"

	"github.com/AsynkronIT/protoactor-go/actor"
	"github.com/AsynkronIT/protoactor-go/router"
	"github.com/Nyarum/noterius/core"
	"github.com/Nyarum/noterius/entities"
	"github.com/Nyarum/noterius/network"
	"github.com/Nyarum/noterius/network/out"
	"github.com/valyala/bytebufferpool"
	"go.uber.org/zap"

	"net"
)

type Server struct {
	config   core.Config
	database *sql.DB
	logger   *zap.SugaredLogger
}

func NewServer(config core.Config, database *sql.DB, logger *zap.SugaredLogger) *Server {
	return &Server{
		config:   config,
		database: database,
		logger:   logger,
	}
}

func (s *Server) Run(ctx context.Context) error {
	defer func() {
		if r := recover(); r != nil {
			s.logger.Errorw("Recovered main server", "details", r)
		}
	}()

	world := actor.Spawn(actor.FromInstance(entities.NewWorld(s.database)))
	defer world.Stop()

	listen, err := net.Listen("tcp", s.config.Common.Host)
	if err != nil {
		return err
	}

	s.logger.Infow("Started server", "host", s.config.Common.Host)

	tick := time.Tick(1 * time.Second / 2)

	// Graceful shutdown and global ticker
	go func() {
		for {
			select {
			case tm := <-tick:
				world.Tell(entities.GlobalTick{
					Now: tm,
				})
			case <-ctx.Done():
				listen.Close()
			}
		}
	}()

	for {
		client, err := listen.Accept()
		if err != nil {
			if ctx.Err() == context.Canceled {
				return nil
			}

			s.logger.Errorw("Error accept connection", "err", err)
			continue
		}

		s.logger.Infow("New connect", "ip", client.RemoteAddr())

		go s.acceptConnect(ctx, client, world)
	}
}

func (s *Server) acceptConnect(ctx context.Context, client net.Conn, world *actor.PID) {
	var (
		packetSender = actor.Spawn(router.NewRoundRobinPool(5).WithInstance(&entities.PacketSender{
			Client:  client,
			Network: network.NewNetwork(),
			Logger:  s.logger,
		}))
		player = actor.Spawn(actor.FromInstance(&entities.Player{
			DB:           s.database,
			World:        world,
			PacketSender: packetSender,
			Logger:       s.logger,
		}))
		packetReader = actor.Spawn(router.NewRoundRobinPool(5).WithInstance(&entities.PacketReader{
			Client:  client,
			Player:  player,
			Network: network.NewNetwork(),
			Logger:  s.logger,
		}))
	)
	defer func() {
		s.logger.Debugw("Shutdown connection actors", "ip", client.RemoteAddr())

		packetReader.Stop()
		packetSender.Stop()
		player.Stop()
	}()

	player.Tell(entities.RecordTime{
		Time: (&out.Date{}).GetCurrentTime(),
	})

	var (
		lenPacket int
		bb        *bytebufferpool.ByteBuffer
	)
	for {
		if ctx.Err() == context.Canceled {
			break
		}

		if lenPacket == 0 {
			bb = bytebufferpool.Get()
		}

		bufTemp := make([]byte, 4096)
		ln, err := client.Read(bufTemp)
		if err != nil {
			if val, ok := err.(net.Error); ok && val.Timeout() {
				s.logger.Errorw("Client is timeout", "error", err)
			}

			if err == io.EOF {
				s.logger.Errorw("Client is disconnected", "error", err)
			}

			break
		}

		bb.Write(bufTemp[:ln])

		// Func to receive many sub packets in an one main packet
		var manyDataFunc func() bool
		manyDataFunc = func() bool {
			if lenPacket == 0 {
				lenPacket = int(binary.BigEndian.Uint16(bb.Bytes()[0:2]))
			}

			if lenPacket < int(ln) {
				return false
			}

			packetReader.Tell(entities.ReadPacket{
				Len: lenPacket,
				Buf: bb.Bytes(),
			})

			bb.Set(bb.Bytes()[lenPacket:])
			lenPacket = 0

			if bb.Len() != 0 {
				return manyDataFunc()
			}

			return false
		}

		if !manyDataFunc() {
			continue
		}

		bytebufferpool.Put(bb)
	}
}
<commit_msg>Just small fix<commit_after>package server

import (
	"context"
	"database/sql"
	"encoding/binary"
	"io"
	"time"

	"github.com/AsynkronIT/protoactor-go/actor"
	"github.com/AsynkronIT/protoactor-go/router"
	"github.com/Nyarum/noterius/core"
	"github.com/Nyarum/noterius/entities"
	"github.com/Nyarum/noterius/network"
	"github.com/Nyarum/noterius/network/out"
	"github.com/valyala/bytebufferpool"
	"go.uber.org/zap"

	"net"
)

type Server struct {
	config   core.Config
	database *sql.DB
	logger   *zap.SugaredLogger
}

func NewServer(config core.Config, database *sql.DB, logger *zap.SugaredLogger) *Server {
	return &Server{
		config:   config,
		database: database,
		logger:   logger,
	}
}

func (s *Server) Run(ctx context.Context) error {
	defer func() {
		if r := recover(); r != nil {
			s.logger.Errorw("Recovered main server", "details", r)
		}
	}()

	world := actor.Spawn(actor.FromInstance(entities.NewWorld(s.database)))
	defer world.Stop()

	listen, err := net.Listen("tcp", s.config.Common.Host)
	if err != nil {
		return err
	}

	s.logger.Infow("Started server", "host", s.config.Common.Host)

	// Graceful shutdown and global ticker
	tick := time.Tick(1 * time.Second / 2)

	go func() {
		for {
			select {
			case tm := <-tick:
				world.Tell(entities.GlobalTick{
					Now: tm,
				})
			case <-ctx.Done():
				listen.Close()
			}
		}
	}()

	for {
		client, err := listen.Accept()
		if err != nil {
			if ctx.Err() == context.Canceled {
				return nil
			}

			s.logger.Errorw("Error accept connection", "err", err)
			continue
		}

		s.logger.Infow("New connect", "ip", client.RemoteAddr())

		go s.acceptConnect(ctx, client, world)
	}
}

func (s *Server) acceptConnect(ctx context.Context, client net.Conn, world *actor.PID) {
	var (
		packetSender = actor.Spawn(router.NewRoundRobinPool(5).WithInstance(&entities.PacketSender{
			Client:  client,
			Network: network.NewNetwork(),
			Logger:  s.logger,
		}))
		player = actor.Spawn(actor.FromInstance(&entities.Player{
			DB:           s.database,
			World:        world,
			PacketSender: packetSender,
			Logger:       s.logger,
		}))
		packetReader = actor.Spawn(router.NewRoundRobinPool(5).WithInstance(&entities.PacketReader{
			Client:  client,
			Player:  player,
			Network: network.NewNetwork(),
			Logger:  s.logger,
		}))
	)
	defer func() {
		s.logger.Debugw("Shutdown connection actors", "ip", client.RemoteAddr())

		packetReader.Stop()
		packetSender.Stop()
		player.Stop()
	}()

	player.Tell(entities.RecordTime{
		Time: (&out.Date{}).GetCurrentTime(),
	})

	var (
		lenPacket int
		bb        *bytebufferpool.ByteBuffer
	)
	for {
		if ctx.Err() == context.Canceled {
			break
		}

		if lenPacket == 0 {
			bb = bytebufferpool.Get()
		}

		bufTemp := make([]byte, 4096)
		ln, err := client.Read(bufTemp)
		if err != nil {
			if val, ok := err.(net.Error); ok && val.Timeout() {
				s.logger.Errorw("Client is timeout", "error", err)
			}

			if err == io.EOF {
				s.logger.Errorw("Client is disconnected", "error", err)
			}

			break
		}

		bb.Write(bufTemp[:ln])

		// Func to receive many sub packets in an one main packet
		var manyDataFunc func() bool
		manyDataFunc = func() bool {
			if lenPacket == 0 {
				lenPacket = int(binary.BigEndian.Uint16(bb.Bytes()[0:2]))
			}

			if lenPacket < int(ln) {
				return false
			}

			packetReader.Tell(entities.ReadPacket{
				Len: lenPacket,
				Buf: bb.Bytes(),
			})

			bb.Set(bb.Bytes()[lenPacket:])
			lenPacket = 0

			if bb.Len() != 0 {
				return manyDataFunc()
			}

			return false
		}

		if !manyDataFunc() {
			continue
		}

		bytebufferpool.Put(bb)
	}
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/subtle"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/braintree/manners"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/meatballhat/negroni-logrus"
	"github.com/pkg/errors"
	"github.com/travis-ci/jupiter-brain"
	"github.com/travis-ci/jupiter-brain/metrics"
	"github.com/travis-ci/jupiter-brain/server/jsonapi"
	"github.com/travis-ci/jupiter-brain/server/negroniraven"
	"golang.org/x/net/context"
)

type server struct {
	addr, authToken, sentryDSN string

	log *logrus.Logger

	i jupiterbrain.InstanceManager

	n *negroni.Negroni
	r *mux.Router
	s *manners.GracefulServer

	db       database
	bootTime time.Time
}

func newServer(cfg *Config) (*server, error) {
	log := logrus.New()
	if cfg.Debug {
		log.Level = logrus.DebugLevel
	}

	log.Formatter = &logrus.TextFormatter{DisableColors: true}

	u, err := url.Parse(cfg.VSphereURL)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse vsphere url")
	}

	if !u.IsAbs() {
		return nil, errors.Errorf("vSphere API URL must be absolute")
	}

	db, err := newPGDatabase(cfg.DatabaseURL)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create postgres database")
	}

	paths := jupiterbrain.VSpherePaths{
		BasePath:    cfg.VSphereBasePath,
		VMPath:      cfg.VSphereVMPath,
		ClusterPath: cfg.VSphereClusterPath,
	}

	srv := &server{
		addr:      cfg.Addr,
		authToken: cfg.AuthToken,
		sentryDSN: cfg.SentryDSN,

		log: log,

		i: jupiterbrain.NewVSphereInstanceManager(log, u, paths),

		n: negroni.New(),
		r: mux.NewRouter(),
		s: manners.NewServer(),

		db:       db,
		bootTime: time.Now().UTC(),
	}

	return srv, nil
}

func (srv *server) Setup() {
	srv.setupRoutes()
	srv.setupMiddleware()
	go srv.signalHandler()
}

func (srv *server) Run() {
	srv.log.WithField("addr", srv.addr).Info("Listening")
	srv.s.Addr = srv.addr
	srv.s.Handler = srv.n
	err := srv.s.ListenAndServe()
	if err != nil {
		srv.log.WithField("err", err).Error("ListenAndServe failed")
	}
}

func (srv *server) setupRoutes() {
	srv.r.HandleFunc(`/instances`, srv.handleInstancesList).Methods("GET").Name("instances-list")
	srv.r.HandleFunc(`/instances`, srv.handleInstancesCreate).Methods("POST").Name("instances-create")
	srv.r.HandleFunc(`/instances/{id}`, srv.handleInstanceByIDFetch).Methods("GET").Name("instance-by-id")
	srv.r.HandleFunc(`/instances/{id}`, srv.handleInstanceByIDTerminate).Methods("DELETE").Name("instance-by-id-terminate")
	srv.r.HandleFunc(`/instance-syncs`, srv.handleInstanceSync).Methods("POST").Name("instance-syncs-create")
}

func (srv *server) setupMiddleware() {
	srv.n.Use(negroni.NewRecovery())
	srv.n.Use(negronilogrus.NewCustomMiddleware(srv.log.Level, srv.log.Formatter, "web"))
	srv.n.Use(negroni.HandlerFunc(srv.authMiddleware))
	nr, err := negroniraven.NewMiddleware(srv.sentryDSN)
	if err != nil {
		panic(err)
	}
	srv.n.Use(nr)
	srv.n.UseHandler(srv.r)
}

func (srv *server) authMiddleware(w http.ResponseWriter, req *http.Request, f http.HandlerFunc) {
	authHeader := req.Header.Get("Authorization")
	srv.log.WithField("authorization", authHeader).Debug("raw authorization header")

	if authHeader == "" {
		w.Header().Set("WWW-Authenticate", "token")
		srv.log.WithField("request_id", req.Header.Get("X-Request-ID")).Debug("responding 401 due to empty Authorization header")

		jsonapi.Error(w, errors.New("token is required"), http.StatusUnauthorized)
		return
	}

	if strings.HasPrefix(authHeader, "token ") && subtle.ConstantTimeCompare([]byte("token "+srv.authToken), []byte(authHeader)) == 1 {
		f(w, req)
		return
	} else if strings.HasPrefix(authHeader, "token=") && subtle.ConstantTimeCompare([]byte("token="+srv.authToken), []byte(authHeader)) == 1 {
		f(w, req)
		return
	}

	jsonapi.Error(w, errors.New("incorrect token"), http.StatusUnauthorized)
}

func (srv *server) handleInstancesList(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instances-list", time.Now())

	instances, err := srv.i.List(context.TODO())
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	dbInstanceIDs := []string{}
	dbInstanceIDCreatedMap := map[string]time.Time{}
	applyDBFilter := false

	if req.FormValue("min_age") != "" {
		dur, err := time.ParseDuration(req.FormValue("min_age"))
		if err != nil {
			jsonapi.Error(w, err, http.StatusBadRequest)
			return
		}

		res, err := srv.db.FetchInstances(&databaseQuery{MinAge: dur})
		if err != nil {
			jsonapi.Error(w, err, http.StatusBadRequest)
			return
		}

		srv.log.WithFields(logrus.Fields{
			"n": len(res),
		}).Debug("retrieved instances from database")

		for _, r := range res {
			dbInstanceIDCreatedMap[r.ID] = r.CreatedAt
			dbInstanceIDs = append(dbInstanceIDs, r.ID)
		}

		applyDBFilter = true
	}

	response := map[string][]interface{}{
		"data": make([]interface{}, 0),
	}

	if applyDBFilter {
		keptInstances := []*jupiterbrain.Instance{}
		for _, instance := range instances {
			for _, instID := range dbInstanceIDs {
				if instID == instance.ID {
					instance.CreatedAt = dbInstanceIDCreatedMap[instID]
					keptInstances = append(keptInstances, instance)
				}
			}
		}

		srv.log.WithFields(logrus.Fields{
			"pre_filter":  len(instances),
			"post_filter": len(keptInstances),
		}).Debug("applying known instance filter")

		instances = keptInstances
	}

	for _, instance := range instances {
		response["data"] = append(response["data"], MarshalInstance(instance))
	}

	b, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/vnd.api+json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, string(b)+"\n")
}

func (srv *server) handleInstancesCreate(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instances-create", time.Now())

	var requestBody map[string]map[string]string

	err := json.NewDecoder(req.Body).Decode(&requestBody)
	if err != nil {
		jsonapi.Error(w, err, http.StatusBadRequest)
		return
	}

	if requestBody["data"] == nil {
		jsonapi.Error(w, &jsonapi.JSONError{Status: "422", Code: "missing-field", Title: "root object must have data field"}, 422)
		return
	}

	if requestBody["data"]["type"] != "instances" {
		jsonapi.Error(w, &jsonapi.JSONError{Status: "409", Code: "incorrect-type", Title: "data must be of type instances"}, http.StatusConflict)
		return
	}

	if requestBody["data"]["base-image"] == "" {
		jsonapi.Error(w, &jsonapi.JSONError{Status: "422", Code: "missing-field", Title: "instance must have base-image field"}, 422)
		return
	}

	instance, err := srv.i.Start(context.TODO(), requestBody["data"]["base-image"])
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	recoverDelete := false
	defer func() {
		if recoverDelete && instance != nil {
			go func() { _ = srv.i.Terminate(context.TODO(), instance.ID) }()
		}
	}()

	instance.CreatedAt = time.Now().UTC()
	err = srv.db.SaveInstance(instance)
	if err != nil {
		recoverDelete = true
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	response := map[string][]interface{}{
		"data": {MarshalInstance(instance)},
	}

	b, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		recoverDelete = true
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/vnd.api+json")
	w.Header().Set("Location", fmt.Sprintf("/instances/%s", instance.ID))
	w.WriteHeader(http.StatusCreated)
	fmt.Fprintf(w, string(b)+"\n")
}

func (srv *server) handleInstanceByIDFetch(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instance-by-id-fetch", time.Now())

	vars := mux.Vars(req)
	instance, err := srv.i.Fetch(context.TODO(), vars["id"])
	if err != nil {
		switch err.(type) {
		case jupiterbrain.VirtualMachineNotFoundError:
			jsonapi.Error(w, err, http.StatusNotFound)
			return
		default:
			srv.log.WithFields(logrus.Fields{
				"err": err,
				"id":  vars["id"],
			}).Error("failed to fetch instance")
			jsonapi.Error(w, err, http.StatusInternalServerError)
			return
		}
	}

	response := map[string][]interface{}{
		"data": {MarshalInstance(instance)},
	}

	b, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/vnd.api+json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, string(b)+"\n")
}

func (srv *server) handleInstanceByIDTerminate(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instance-by-id-terminate", time.Now())

	vars := mux.Vars(req)
	err := srv.i.Terminate(context.TODO(), vars["id"])
	if err != nil {
		switch err.(type) {
		case jupiterbrain.VirtualMachineNotFoundError:
			jsonapi.Error(w, err, http.StatusNotFound)
			return
		default:
			srv.log.WithFields(logrus.Fields{
				"err": err,
				"id":  vars["id"],
			}).Error("failed to terminate instance")
			jsonapi.Error(w, err, http.StatusInternalServerError)
			return
		}
	}

	err = srv.db.DestroyInstance(vars["id"])
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (srv *server) handleInstanceSync(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instance-sync", time.Now())
	instances, err := srv.i.List(context.TODO())
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	for _, instance := range instances {
		instance.CreatedAt = time.Now().UTC()
		err = srv.db.SaveInstance(instance)
		if err != nil {
			srv.log.WithFields(logrus.Fields{
				"err": err,
				"id":  instance.ID,
			}).Warn("failed to save instance")
			continue
		}

		srv.log.WithField("id", instance.ID).Debug("synced instance")
	}

	w.WriteHeader(http.StatusNoContent)
}

func (srv *server) signalHandler() {
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1)
	for {
		select {
		case sig := <-signalChan:
			switch sig {
			case syscall.SIGTERM:
				srv.log.Info("Received SIGTERM, shutting down now.")
				srv.s.Close()
			case syscall.SIGINT:
				srv.log.Info("Received SIGINT, shutting down now.")
				srv.s.Close()
			case syscall.SIGUSR1:
				srv.log.WithFields(logrus.Fields{
					"version":   os.Getenv("VERSION"),
					"revision":  os.Getenv("REVISION"),
					"boot_time": srv.bootTime,
					"uptime":    time.Since(srv.bootTime),
				}).Info("Received SIGUSR1.")
			default:
				log.Print("ignoring unknown signal")
			}
		default:
			time.Sleep(time.Second)
		}
	}
}
<commit_msg>server: block shutdown until instance deletion completes<commit_after>package server

import (
	"crypto/subtle"
	"encoding/json"
	"fmt"
	"log"
	"net/http"
	"net/url"
	"os"
	"os/signal"
	"strings"
	"syscall"
	"time"

	"github.com/Sirupsen/logrus"
	"github.com/braintree/manners"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/meatballhat/negroni-logrus"
	"github.com/pkg/errors"
	"github.com/travis-ci/jupiter-brain"
	"github.com/travis-ci/jupiter-brain/metrics"
	"github.com/travis-ci/jupiter-brain/server/jsonapi"
	"github.com/travis-ci/jupiter-brain/server/negroniraven"
	"golang.org/x/net/context"
)

type server struct {
	addr, authToken, sentryDSN string

	log *logrus.Logger

	i jupiterbrain.InstanceManager

	n *negroni.Negroni
	r *mux.Router
	s *manners.GracefulServer

	db       database
	bootTime time.Time
}

func newServer(cfg *Config) (*server, error) {
	log := logrus.New()
	if cfg.Debug {
		log.Level = logrus.DebugLevel
	}

	log.Formatter = &logrus.TextFormatter{DisableColors: true}

	u, err := url.Parse(cfg.VSphereURL)
	if err != nil {
		return nil, errors.Wrap(err, "failed to parse vsphere url")
	}

	if !u.IsAbs() {
		return nil, errors.Errorf("vSphere API URL must be absolute")
	}

	db, err := newPGDatabase(cfg.DatabaseURL)
	if err != nil {
		return nil, errors.Wrap(err, "failed to create postgres database")
	}

	paths := jupiterbrain.VSpherePaths{
		BasePath:    cfg.VSphereBasePath,
		VMPath:      cfg.VSphereVMPath,
		ClusterPath: cfg.VSphereClusterPath,
	}

	srv := &server{
		addr:      cfg.Addr,
		authToken: cfg.AuthToken,
		sentryDSN: cfg.SentryDSN,

		log: log,

		i: jupiterbrain.NewVSphereInstanceManager(log, u, paths),

		n: negroni.New(),
		r: mux.NewRouter(),
		s: manners.NewServer(),

		db:       db,
		bootTime: time.Now().UTC(),
	}

	return srv, nil
}

func (srv *server) Setup() {
	srv.setupRoutes()
	srv.setupMiddleware()
	go srv.signalHandler()
}

func (srv *server) Run() {
	srv.log.WithField("addr", srv.addr).Info("Listening")
	srv.s.Addr = srv.addr
	srv.s.Handler = srv.n
	err := srv.s.ListenAndServe()
	if err != nil {
		srv.log.WithField("err", err).Error("ListenAndServe failed")
	}
}

func (srv *server) setupRoutes() {
	srv.r.HandleFunc(`/instances`, srv.handleInstancesList).Methods("GET").Name("instances-list")
	srv.r.HandleFunc(`/instances`, srv.handleInstancesCreate).Methods("POST").Name("instances-create")
	srv.r.HandleFunc(`/instances/{id}`, srv.handleInstanceByIDFetch).Methods("GET").Name("instance-by-id")
	srv.r.HandleFunc(`/instances/{id}`, srv.handleInstanceByIDTerminate).Methods("DELETE").Name("instance-by-id-terminate")
	srv.r.HandleFunc(`/instance-syncs`, srv.handleInstanceSync).Methods("POST").Name("instance-syncs-create")
}

func (srv *server) setupMiddleware() {
	srv.n.Use(negroni.NewRecovery())
	srv.n.Use(negronilogrus.NewCustomMiddleware(srv.log.Level, srv.log.Formatter, "web"))
	srv.n.Use(negroni.HandlerFunc(srv.authMiddleware))
	nr, err := negroniraven.NewMiddleware(srv.sentryDSN)
	if err != nil {
		panic(err)
	}
	srv.n.Use(nr)
	srv.n.UseHandler(srv.r)
}

func (srv *server) authMiddleware(w http.ResponseWriter, req *http.Request, f http.HandlerFunc) {
	authHeader := req.Header.Get("Authorization")
	srv.log.WithField("authorization", authHeader).Debug("raw authorization header")

	if authHeader == "" {
		w.Header().Set("WWW-Authenticate", "token")
		srv.log.WithField("request_id", req.Header.Get("X-Request-ID")).Debug("responding 401 due to empty Authorization header")

		jsonapi.Error(w, errors.New("token is required"), http.StatusUnauthorized)
		return
	}

	if strings.HasPrefix(authHeader, "token ") && subtle.ConstantTimeCompare([]byte("token "+srv.authToken), []byte(authHeader)) == 1 {
		f(w, req)
		return
	} else if strings.HasPrefix(authHeader, "token=") && subtle.ConstantTimeCompare([]byte("token="+srv.authToken), []byte(authHeader)) == 1 {
		f(w, req)
		return
	}

	jsonapi.Error(w, errors.New("incorrect token"), http.StatusUnauthorized)
}

func (srv *server) handleInstancesList(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instances-list", time.Now())

	instances, err := srv.i.List(context.TODO())
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	dbInstanceIDs := []string{}
	dbInstanceIDCreatedMap := map[string]time.Time{}
	applyDBFilter := false

	if req.FormValue("min_age") != "" {
		dur, err := time.ParseDuration(req.FormValue("min_age"))
		if err != nil {
			jsonapi.Error(w, err, http.StatusBadRequest)
			return
		}

		res, err := srv.db.FetchInstances(&databaseQuery{MinAge: dur})
		if err != nil {
			jsonapi.Error(w, err, http.StatusBadRequest)
			return
		}

		srv.log.WithFields(logrus.Fields{
			"n": len(res),
		}).Debug("retrieved instances from database")

		for _, r := range res {
			dbInstanceIDCreatedMap[r.ID] = r.CreatedAt
			dbInstanceIDs = append(dbInstanceIDs, r.ID)
		}

		applyDBFilter = true
	}

	response := map[string][]interface{}{
		"data": make([]interface{}, 0),
	}

	if applyDBFilter {
		keptInstances := []*jupiterbrain.Instance{}
		for _, instance := range instances {
			for _, instID := range dbInstanceIDs {
				if instID == instance.ID {
					instance.CreatedAt = dbInstanceIDCreatedMap[instID]
					keptInstances = append(keptInstances, instance)
				}
			}
		}

		srv.log.WithFields(logrus.Fields{
			"pre_filter":  len(instances),
			"post_filter": len(keptInstances),
		}).Debug("applying known instance filter")

		instances = keptInstances
	}

	for _, instance := range instances {
		response["data"] = append(response["data"], MarshalInstance(instance))
	}

	b, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/vnd.api+json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, string(b)+"\n")
}

func (srv *server) handleInstancesCreate(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instances-create", time.Now())

	var requestBody map[string]map[string]string

	err := json.NewDecoder(req.Body).Decode(&requestBody)
	if err != nil {
		jsonapi.Error(w, err, http.StatusBadRequest)
		return
	}

	if requestBody["data"] == nil {
		jsonapi.Error(w, &jsonapi.JSONError{Status: "422", Code: "missing-field", Title: "root object must have data field"}, 422)
		return
	}

	if requestBody["data"]["type"] != "instances" {
		jsonapi.Error(w, &jsonapi.JSONError{Status: "409", Code: "incorrect-type", Title: "data must be of type instances"}, http.StatusConflict)
		return
	}

	if requestBody["data"]["base-image"] == "" {
		jsonapi.Error(w, &jsonapi.JSONError{Status: "422", Code: "missing-field", Title: "instance must have base-image field"}, 422)
		return
	}

	instance, err := srv.i.Start(context.TODO(), requestBody["data"]["base-image"])
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	recoverDelete := false
	defer func() {
		if recoverDelete && instance != nil {
			go func() {
				srv.s.StartRoutine()
				defer srv.s.FinishRoutine()
				srv.i.Terminate(context.TODO(), instance.ID)
			}()
		}
	}()

	instance.CreatedAt = time.Now().UTC()
	err = srv.db.SaveInstance(instance)
	if err != nil {
		recoverDelete = true
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	response := map[string][]interface{}{
		"data": {MarshalInstance(instance)},
	}

	b, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		recoverDelete = true
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/vnd.api+json")
	w.Header().Set("Location", fmt.Sprintf("/instances/%s", instance.ID))
	w.WriteHeader(http.StatusCreated)
	fmt.Fprintf(w, string(b)+"\n")
}

func (srv *server) handleInstanceByIDFetch(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instance-by-id-fetch", time.Now())

	vars := mux.Vars(req)
	instance, err := srv.i.Fetch(context.TODO(), vars["id"])
	if err != nil {
		switch err.(type) {
		case jupiterbrain.VirtualMachineNotFoundError:
			jsonapi.Error(w, err, http.StatusNotFound)
			return
		default:
			srv.log.WithFields(logrus.Fields{
				"err": err,
				"id":  vars["id"],
			}).Error("failed to fetch instance")
			jsonapi.Error(w, err, http.StatusInternalServerError)
			return
		}
	}

	response := map[string][]interface{}{
		"data": {MarshalInstance(instance)},
	}

	b, err := json.MarshalIndent(response, "", "  ")
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.Header().Set("Content-Type", "application/vnd.api+json")
	w.WriteHeader(http.StatusOK)
	fmt.Fprintf(w, string(b)+"\n")
}

func (srv *server) handleInstanceByIDTerminate(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instance-by-id-terminate", time.Now())

	vars := mux.Vars(req)
	err := srv.i.Terminate(context.TODO(), vars["id"])
	if err != nil {
		switch err.(type) {
		case jupiterbrain.VirtualMachineNotFoundError:
			jsonapi.Error(w, err, http.StatusNotFound)
			return
		default:
			srv.log.WithFields(logrus.Fields{
				"err": err,
				"id":  vars["id"],
			}).Error("failed to terminate instance")
			jsonapi.Error(w, err, http.StatusInternalServerError)
			return
		}
	}

	err = srv.db.DestroyInstance(vars["id"])
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	w.WriteHeader(http.StatusNoContent)
}

func (srv *server) handleInstanceSync(w http.ResponseWriter, req *http.Request) {
	defer metrics.TimeSince("travis.jupiter-brain.endpoints.instance-sync", time.Now())
	instances, err := srv.i.List(context.TODO())
	if err != nil {
		jsonapi.Error(w, err, http.StatusInternalServerError)
		return
	}

	for _, instance := range instances {
		instance.CreatedAt = time.Now().UTC()
		err = srv.db.SaveInstance(instance)
		if err != nil {
			srv.log.WithFields(logrus.Fields{
				"err": err,
				"id":  instance.ID,
			}).Warn("failed to save instance")
			continue
		}

		srv.log.WithField("id", instance.ID).Debug("synced instance")
	}

	w.WriteHeader(http.StatusNoContent)
}

func (srv *server) signalHandler() {
	signalChan := make(chan os.Signal, 1)
	signal.Notify(signalChan, syscall.SIGTERM, syscall.SIGINT, syscall.SIGUSR1)
	for {
		select {
		case sig := <-signalChan:
			switch sig {
			case syscall.SIGTERM:
				srv.log.Info("Received SIGTERM, shutting down now.")
				srv.s.Close()
			case syscall.SIGINT:
				srv.log.Info("Received SIGINT, shutting down now.")
				srv.s.Close()
			case syscall.SIGUSR1:
				srv.log.WithFields(logrus.Fields{
					"version":   os.Getenv("VERSION"),
					"revision":  os.Getenv("REVISION"),
					"boot_time": srv.bootTime,
					"uptime":    time.Since(srv.bootTime),
				}).Info("Received SIGUSR1.")
			default:
				log.Print("ignoring unknown signal")
			}
		default:
			time.Sleep(time.Second)
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/meatballhat/negroni-logrus"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/phyber/negroni-gzip/gzip"
	"github.com/rs/cors"
	"github.com/unrolled/render"
	"html/template"
	"net/http"
	_ "net/http/pprof"
	"strings"
)

type Server struct {
	ts          TokenStore
	middlewares []negroni.HandlerFunc
	ver         *versioning
	*mux.Router
}
type TokenStore interface {
	Validate(token auth.Token) bool
}

func NewServer(ts TokenStore, middlewares []negroni.HandlerFunc) *Server {
	return &Server{
		ts,
		middlewares,
		newVersioning(),
		mux.NewRouter(),
	}
}

func (s *Server) AddRest(path string, rests []interface{}) {
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IndentJSON:    appgo.Conf.DevMode,
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range rests {
		h := newHandler(api, HandlerTypeJson, s.ts, renderer)
		s.Handle(path+h.path, h).Methods(h.supports...)
	}
}

func (s *Server) AddHtml(path, layout string, htmls []interface{}, funcs template.FuncMap) {
	// add "static" template function
	static := func(path string) string {
		return s.ver.getStatic(path)
	}
	if funcs == nil {
		funcs = template.FuncMap{}
	}
	funcs["static"] = static

	renderer := render.New(render.Options{
		Directory:     appgo.Conf.TemplatePath,
		Layout:        layout,
		Funcs:         []template.FuncMap{funcs},
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range htmls {
		h := newHandler(api, HandlerTypeHtml, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddFeed(path string, feeds []interface{}) {
	// renderer is only for rendering error
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range feeds {
		h := newHandler(api, HandlerTypeFeed, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddProxy(path string, handler http.Handler) {
	s.PathPrefix(path).Handler(http.StripPrefix(path, handler))
}

func (s *Server) AddStatic(path, fileDir string) {
	s.ver.addMap(path, fileDir)
	s.AddProxy(path, http.FileServer(http.Dir(fileDir)))
}

func (s *Server) AddAppleAppSiteAsso(content []byte) {
	f := func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
	s.HandleFunc("/apple-app-site-association", f)
}

func (s *Server) Serve() {
	if appgo.Conf.Pprof.Enable {
		go func() {
			log.Infoln(http.ListenAndServe(":"+appgo.Conf.Pprof.Port, nil))
		}()
	}

	n := negroni.New()
	n.Use(negroni.NewRecovery())
	n.Use(negronilogrus.NewCustomMiddleware(
		appgo.Conf.LogLevel, &log.TextFormatter{}, "appgo"))
	n.Use(cors.New(corsOptions()))
	if appgo.Conf.Negroni.GZip {
		n.Use(gzip.Gzip(gzip.DefaultCompression))
	}
	for _, mw := range s.middlewares {
		n.Use(negroni.HandlerFunc(mw))
	}
	n.UseHandler(s)
	n.Run(appgo.Conf.Negroni.Port)
}

func corsOptions() cors.Options {
	origins := strings.Split(appgo.Conf.Cors.AllowedOrigins, ",")
	methods := strings.Split(appgo.Conf.Cors.AllowedMethods, ",")
	headers := strings.Split(appgo.Conf.Cors.AllowedHeaders, ",")
	return cors.Options{
		AllowedOrigins:     origins,
		AllowedMethods:     methods,
		AllowedHeaders:     headers,
		OptionsPassthrough: appgo.Conf.Cors.OptionsPassthrough,
		Debug:              appgo.Conf.Cors.Debug,
	}
}
<commit_msg>compression level<commit_after>package server

import (
	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/meatballhat/negroni-logrus"
	"github.com/oxfeeefeee/appgo"
	"github.com/oxfeeefeee/appgo/auth"
	"github.com/phyber/negroni-gzip/gzip"
	"github.com/rs/cors"
	"github.com/unrolled/render"
	"html/template"
	"net/http"
	_ "net/http/pprof"
	"strings"
)

type Server struct {
	ts          TokenStore
	middlewares []negroni.HandlerFunc
	ver         *versioning
	*mux.Router
}
type TokenStore interface {
	Validate(token auth.Token) bool
}

func NewServer(ts TokenStore, middlewares []negroni.HandlerFunc) *Server {
	return &Server{
		ts,
		middlewares,
		newVersioning(),
		mux.NewRouter(),
	}
}

func (s *Server) AddRest(path string, rests []interface{}) {
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IndentJSON:    appgo.Conf.DevMode,
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range rests {
		h := newHandler(api, HandlerTypeJson, s.ts, renderer)
		s.Handle(path+h.path, h).Methods(h.supports...)
	}
}

func (s *Server) AddHtml(path, layout string, htmls []interface{}, funcs template.FuncMap) {
	// add "static" template function
	static := func(path string) string {
		return s.ver.getStatic(path)
	}
	if funcs == nil {
		funcs = template.FuncMap{}
	}
	funcs["static"] = static

	renderer := render.New(render.Options{
		Directory:     appgo.Conf.TemplatePath,
		Layout:        layout,
		Funcs:         []template.FuncMap{funcs},
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range htmls {
		h := newHandler(api, HandlerTypeHtml, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddFeed(path string, feeds []interface{}) {
	// renderer is only for rendering error
	renderer := render.New(render.Options{
		Directory:     "N/A",
		IsDevelopment: appgo.Conf.DevMode,
	})
	for _, api := range feeds {
		h := newHandler(api, HandlerTypeFeed, s.ts, renderer)
		s.Handle(path+h.path, h).Methods("GET")
	}
}

func (s *Server) AddProxy(path string, handler http.Handler) {
	s.PathPrefix(path).Handler(http.StripPrefix(path, handler))
}

func (s *Server) AddStatic(path, fileDir string) {
	s.ver.addMap(path, fileDir)
	s.AddProxy(path, http.FileServer(http.Dir(fileDir)))
}

func (s *Server) AddAppleAppSiteAsso(content []byte) {
	f := func(w http.ResponseWriter, req *http.Request) {
		w.Header().Set("Content-Type", "application/json")
		w.Write(content)
	}
	s.HandleFunc("/apple-app-site-association", f)
}

func (s *Server) Serve() {
	if appgo.Conf.Pprof.Enable {
		go func() {
			log.Infoln(http.ListenAndServe(":"+appgo.Conf.Pprof.Port, nil))
		}()
	}

	n := negroni.New()
	n.Use(negroni.NewRecovery())
	n.Use(negronilogrus.NewCustomMiddleware(
		appgo.Conf.LogLevel, &log.TextFormatter{}, "appgo"))
	n.Use(cors.New(corsOptions()))
	if appgo.Conf.Negroni.GZip {
		n.Use(gzip.Gzip(gzip.BestSpeed))
	}
	for _, mw := range s.middlewares {
		n.Use(negroni.HandlerFunc(mw))
	}
	n.UseHandler(s)
	n.Run(appgo.Conf.Negroni.Port)
}

func corsOptions() cors.Options {
	origins := strings.Split(appgo.Conf.Cors.AllowedOrigins, ",")
	methods := strings.Split(appgo.Conf.Cors.AllowedMethods, ",")
	headers := strings.Split(appgo.Conf.Cors.AllowedHeaders, ",")
	return cors.Options{
		AllowedOrigins:     origins,
		AllowedMethods:     methods,
		AllowedHeaders:     headers,
		OptionsPassthrough: appgo.Conf.Cors.OptionsPassthrough,
		Debug:              appgo.Conf.Cors.Debug,
	}
}
<|endoftext|>


<commit_before>package server

import (
	"crypto/rand"
	"encoding/hex"
	"net"
	"net/http"
	"net/http/pprof"
	"net/url"
	"strings"

	"github.com/tchap/steemwatch/config"
	"github.com/tchap/steemwatch/server/auth"
	"github.com/tchap/steemwatch/server/auth/facebook"
	"github.com/tchap/steemwatch/server/auth/github"
	"github.com/tchap/steemwatch/server/auth/google"
	"github.com/tchap/steemwatch/server/auth/reddit"
	"github.com/tchap/steemwatch/server/context"
	"github.com/tchap/steemwatch/server/db"
	"github.com/tchap/steemwatch/server/routes/api/eventstream"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/discord"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/slack"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/steemitchat"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/telegram"
	"github.com/tchap/steemwatch/server/routes/api/profile"
	"github.com/tchap/steemwatch/server/routes/api/v1/info"
	"github.com/tchap/steemwatch/server/routes/home"
	"github.com/tchap/steemwatch/server/routes/logout"
	"github.com/tchap/steemwatch/server/sessions"
	"github.com/tchap/steemwatch/server/users/stores/mongodb"
	"github.com/tchap/steemwatch/server/views"

	"github.com/bwmarrin/discordgo"
	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	gorillaSessions "github.com/gorilla/sessions"
	"github.com/labstack/echo"
	"github.com/labstack/echo-contrib/session"
	"github.com/labstack/echo/middleware"
	"github.com/pkg/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/tomb.v2"
)

type Context struct {
	EventStreamManager *eventstream.Manager

	listener net.Listener

	discordSession *discordgo.Session

	t tomb.Tomb
}

func Run(mongo *mgo.Database, cfg *config.Config) (*Context, *discordgo.Session, error) {
	serverCtx := &context.Context{}

	// Environment.
	switch cfg.Env {
	case "development":
		serverCtx.Env = context.EnvironmentDevelopment
	case "production":
		serverCtx.Env = context.EnvironmentProduction
		serverCtx.SSLEnabled = true
	default:
		return nil, nil, errors.New("invalid environment: " + cfg.Env)
	}

	// Database.
	serverCtx.DB = mongo

	// User store.
	userStore := mongodb.NewUserStore(mongo.C("users"))

	// Session manager.
	hashKey, blockKey, err := getSecureCookieKeys(mongo)
	if err != nil {
		return nil, nil, err
	}
	sessionManager, err := sessions.NewSessionManager(userStore)
	if err != nil {
		return nil, nil, err
	}

	serverCtx.SessionManager = sessionManager

	// Server context.
	canonicalURL, err := url.Parse(cfg.CanonicalURL)
	if err != nil {
		return nil, nil, errors.Wrap(err, "invalid canonical URL")
	}

	serverCtx.CanonicalURL = canonicalURL

	// Echo.
	e := echo.New()

	// Templates.
	renderer, err := views.NewRenderer("./server/views/*.html")
	if err != nil {
		return nil, nil, err
	}
	e.Renderer = renderer

	// Assets.
	e.Static("/app", "server/app")
	e.Static("/modules", "server/app/node_modules")
	e.Static("/assets/css", "server/assets/css")
	e.Static("/assets/js", "server/assets/js")
	e.Static("/assets/img", "server/assets/img")
	e.Static("/assets/fonts", "server/assets/fonts")
	e.Static("/assets/bootstrap", "server/app/node_modules/bootstrap/dist")

	// Middleware
	e.Pre(middleware.AddTrailingSlash())
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	e.Use(middleware.Secure())
	e.Use(session.Middleware(gorillaSessions.NewCookieStore(hashKey, blockKey)))

	// A temporary fix. We need to encode the CSRF header value.
	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			req := c.Request()
			token := req.Header.Get(echo.HeaderXCSRFToken)
			req.Header.Set(echo.HeaderXCSRFToken, url.QueryEscape(token))
			return next(c)
		}
	})

	csrfConfig := middleware.DefaultCSRFConfig
	csrfConfig.CookieName = "csrf"
	csrfConfig.CookiePath = "/"
	csrf := middleware.CSRFWithConfig(csrfConfig)

	// Debug
	e.GET("/debug/pprof/cmdline", echo.WrapHandler(http.HandlerFunc(pprof.Cmdline)))
	e.GET("/debug/pprof/profile", echo.WrapHandler(http.HandlerFunc(pprof.Profile)))
	e.GET("/debug/pprof/symbol", echo.WrapHandler(http.HandlerFunc(pprof.Symbol)))
	e.GET("/debug/pprof/trace", echo.WrapHandler(http.HandlerFunc(pprof.Trace)))
	e.GET(
		"/debug/pprof",
		echo.WrapHandler(http.HandlerFunc(pprof.Index)),
		middleware.AddTrailingSlashWithConfig(middleware.TrailingSlashConfig{
			RedirectCode: http.StatusPermanentRedirect,
		}),
	)
	e.GET("/debug/pprof/*", echo.WrapHandler(http.HandlerFunc(pprof.Index)))

	// Web
	homeHandler := home.NewHandlerFunc(serverCtx)

	e.GET("/", homeHandler, csrf)
	e.GET("/logout/", logout.NewHandlerFunc(serverCtx), csrf)

	e.GET("/home/", homeHandler, csrf)
	e.GET("/events/", homeHandler, csrf)
	e.GET("/eventstream/", homeHandler, csrf)
	e.GET("/notifications/", homeHandler, csrf)
	e.GET("/profile/", homeHandler, csrf)

	facebookCallbackPath, _ := url.Parse("/auth/facebook/callback")
	facebookCallback := serverCtx.CanonicalURL.ResolveReference(facebookCallbackPath).String()
	facebookAuth := facebook.NewAuthenticator(
		cfg.FacebookClientId, cfg.FacebookClientSecret, facebookCallback)
	auth.Bind(serverCtx, e.Group("/auth/facebook", csrf), facebookAuth)

	redditCallbackPath, _ := url.Parse("/auth/reddit/callback")
	redditCallback := serverCtx.CanonicalURL.ResolveReference(redditCallbackPath).String()
	redditAuth := reddit.NewAuthenticator(
		cfg.RedditClientId, cfg.RedditClientSecret, redditCallback, serverCtx.SSLEnabled)
	auth.Bind(serverCtx, e.Group("/auth/reddit", csrf), redditAuth)

	googleCallbackPath, _ := url.Parse("/auth/google/callback")
	googleCallback := serverCtx.CanonicalURL.ResolveReference(googleCallbackPath).String()
	googleAuth := google.NewAuthenticator(
		cfg.GoogleClientId, cfg.GoogleClientSecret, googleCallback)
	auth.Bind(serverCtx, e.Group("/auth/google", csrf), googleAuth)

	githubCallbackPath, _ := url.Parse("/auth/github/callback")
	githubCallback := serverCtx.CanonicalURL.ResolveReference(githubCallbackPath).String()
	githubAuth := github.NewAuthenticator(
		cfg.GitHubClientId, cfg.GitHubClientSecret, githubCallback)
	auth.Bind(serverCtx, e.Group("/auth/github", csrf), githubAuth)

	// Public API
	info.Bind(serverCtx, e.Group("/api/v1/info"))

	// API
	api := e.Group("/api", csrf, auth.Required(serverCtx))

	// API - Events
	db.BindList(serverCtx, api.Group("/events/:kind/:list"))

	// API - Event Stream
	manager := eventstream.NewManager()
	manager.Bind(serverCtx, api.Group("/eventstream"))

	// API - Notifiers
	slack.Bind(serverCtx, api.Group("/notifiers/slack"))
	steemitchat.Bind(serverCtx, api.Group("/notifiers/steemit-chat"))

	// Telegram
	botSecret := make([]byte, 256/8)
	if _, err := rand.Read(botSecret); err != nil {
		return nil, nil, errors.Wrap(err, "failed to generate Telegram bot secret")
	}
	botSecretHex := hex.EncodeToString(botSecret)

	botURL, _ := url.Parse("/bots/telegram/both-" + botSecretHex)
	botWebhookURL := serverCtx.CanonicalURL.ResolveReference(botURL)

	bot, err := tgbotapi.NewBotAPI(cfg.TelegramBotToken)
	if err != nil {
		return nil, nil, errors.Wrap(err, "failed to initialize Telegram bot")
	}

	if !strings.Contains(cfg.CanonicalURL, "localhost") {
		if _, err := bot.SetWebhook(tgbotapi.NewWebhook(botWebhookURL.String())); err != nil {
			return nil, nil, errors.Wrap(err, "failed to register Telegram webhook URL")
		}
	}

	telegram.BindWebhook(serverCtx, e.Group(botURL.String()))
	telegram.BindAPI(serverCtx, api.Group("/notifiers/telegram"))

	// API - Profile
	profile.Bind(serverCtx, api.Group("/profile"))

	// Start server
	listener, err := net.Listen("tcp", cfg.ListenAddress)
	if err != nil {
		return nil, nil, errors.Wrap(err, "failed to start the web server")
	}

	ctx := &Context{
		EventStreamManager: manager,
		listener:           listener,
	}

	// Discord
	dg, err := discord.InitBot(&ctx.t, cfg.DiscordBotToken, serverCtx)
	if err != nil {
		return nil, nil, err
	}

	discord.BindAPI(serverCtx, api.Group("/notifiers/discord"))

	// Start listening.
	ctx.t.Go(func() error {
		http.Serve(listener, e)
		return nil
	})

	go func() {
		<-ctx.t.Dying()
		listener.Close()
	}()

	return ctx, dg, nil
}

func (ctx *Context) Interrupt() {
	ctx.t.Kill(nil)
}

func (ctx *Context) Wait() error {
	return ctx.t.Wait()
}
<commit_msg>server: Fix debug endpoint routing<commit_after>package server

import (
	"crypto/rand"
	"encoding/hex"
	"net"
	"net/http"
	"net/http/pprof"
	"net/url"
	"strings"

	"github.com/tchap/steemwatch/config"
	"github.com/tchap/steemwatch/server/auth"
	"github.com/tchap/steemwatch/server/auth/facebook"
	"github.com/tchap/steemwatch/server/auth/github"
	"github.com/tchap/steemwatch/server/auth/google"
	"github.com/tchap/steemwatch/server/auth/reddit"
	"github.com/tchap/steemwatch/server/context"
	"github.com/tchap/steemwatch/server/db"
	"github.com/tchap/steemwatch/server/routes/api/eventstream"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/discord"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/slack"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/steemitchat"
	"github.com/tchap/steemwatch/server/routes/api/notifiers/telegram"
	"github.com/tchap/steemwatch/server/routes/api/profile"
	"github.com/tchap/steemwatch/server/routes/api/v1/info"
	"github.com/tchap/steemwatch/server/routes/home"
	"github.com/tchap/steemwatch/server/routes/logout"
	"github.com/tchap/steemwatch/server/sessions"
	"github.com/tchap/steemwatch/server/users/stores/mongodb"
	"github.com/tchap/steemwatch/server/views"

	"github.com/bwmarrin/discordgo"
	tgbotapi "github.com/go-telegram-bot-api/telegram-bot-api"
	gorillaSessions "github.com/gorilla/sessions"
	"github.com/labstack/echo"
	"github.com/labstack/echo-contrib/session"
	"github.com/labstack/echo/middleware"
	"github.com/pkg/errors"
	"gopkg.in/mgo.v2"
	"gopkg.in/tomb.v2"
)

type Context struct {
	EventStreamManager *eventstream.Manager

	listener net.Listener

	discordSession *discordgo.Session

	t tomb.Tomb
}

func Run(mongo *mgo.Database, cfg *config.Config) (*Context, *discordgo.Session, error) {
	serverCtx := &context.Context{}

	// Environment.
	switch cfg.Env {
	case "development":
		serverCtx.Env = context.EnvironmentDevelopment
	case "production":
		serverCtx.Env = context.EnvironmentProduction
		serverCtx.SSLEnabled = true
	default:
		return nil, nil, errors.New("invalid environment: " + cfg.Env)
	}

	// Database.
	serverCtx.DB = mongo

	// User store.
	userStore := mongodb.NewUserStore(mongo.C("users"))

	// Session manager.
	hashKey, blockKey, err := getSecureCookieKeys(mongo)
	if err != nil {
		return nil, nil, err
	}
	sessionManager, err := sessions.NewSessionManager(userStore)
	if err != nil {
		return nil, nil, err
	}

	serverCtx.SessionManager = sessionManager

	// Server context.
	canonicalURL, err := url.Parse(cfg.CanonicalURL)
	if err != nil {
		return nil, nil, errors.Wrap(err, "invalid canonical URL")
	}

	serverCtx.CanonicalURL = canonicalURL

	// Echo.
	e := echo.New()

	// Templates.
	renderer, err := views.NewRenderer("./server/views/*.html")
	if err != nil {
		return nil, nil, err
	}
	e.Renderer = renderer

	// Assets.
	e.Static("/app", "server/app")
	e.Static("/modules", "server/app/node_modules")
	e.Static("/assets/css", "server/assets/css")
	e.Static("/assets/js", "server/assets/js")
	e.Static("/assets/img", "server/assets/img")
	e.Static("/assets/fonts", "server/assets/fonts")
	e.Static("/assets/bootstrap", "server/app/node_modules/bootstrap/dist")

	// Middleware
	e.Pre(middleware.AddTrailingSlash())
	e.Use(middleware.Logger())
	e.Use(middleware.Recover())
	e.Use(middleware.Secure())
	e.Use(session.Middleware(gorillaSessions.NewCookieStore(hashKey, blockKey)))

	// A temporary fix. We need to encode the CSRF header value.
	e.Use(func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			req := c.Request()
			token := req.Header.Get(echo.HeaderXCSRFToken)
			req.Header.Set(echo.HeaderXCSRFToken, url.QueryEscape(token))
			return next(c)
		}
	})

	csrfConfig := middleware.DefaultCSRFConfig
	csrfConfig.CookieName = "csrf"
	csrfConfig.CookiePath = "/"
	csrf := middleware.CSRFWithConfig(csrfConfig)

	// Debug
	e.GET("/debug/pprof/cmdline/", echo.WrapHandler(http.HandlerFunc(pprof.Cmdline)))
	e.GET("/debug/pprof/profile/", echo.WrapHandler(http.HandlerFunc(pprof.Profile)))
	e.GET("/debug/pprof/symbol/", echo.WrapHandler(http.HandlerFunc(pprof.Symbol)))
	e.GET("/debug/pprof/trace/", echo.WrapHandler(http.HandlerFunc(pprof.Trace)))
	e.GET("/debug/pprof/", echo.WrapHandler(http.HandlerFunc(pprof.Index)))
	e.GET(
		"/debug/pprof/*",
		echo.WrapHandler(http.HandlerFunc(pprof.Index)),
		middleware.RemoveTrailingSlash(),
	)

	// Web
	homeHandler := home.NewHandlerFunc(serverCtx)

	e.GET("/", homeHandler, csrf)
	e.GET("/logout/", logout.NewHandlerFunc(serverCtx), csrf)

	e.GET("/home/", homeHandler, csrf)
	e.GET("/events/", homeHandler, csrf)
	e.GET("/eventstream/", homeHandler, csrf)
	e.GET("/notifications/", homeHandler, csrf)
	e.GET("/profile/", homeHandler, csrf)

	facebookCallbackPath, _ := url.Parse("/auth/facebook/callback")
	facebookCallback := serverCtx.CanonicalURL.ResolveReference(facebookCallbackPath).String()
	facebookAuth := facebook.NewAuthenticator(
		cfg.FacebookClientId, cfg.FacebookClientSecret, facebookCallback)
	auth.Bind(serverCtx, e.Group("/auth/facebook", csrf), facebookAuth)

	redditCallbackPath, _ := url.Parse("/auth/reddit/callback")
	redditCallback := serverCtx.CanonicalURL.ResolveReference(redditCallbackPath).String()
	redditAuth := reddit.NewAuthenticator(
		cfg.RedditClientId, cfg.RedditClientSecret, redditCallback, serverCtx.SSLEnabled)
	auth.Bind(serverCtx, e.Group("/auth/reddit", csrf), redditAuth)

	googleCallbackPath, _ := url.Parse("/auth/google/callback")
	googleCallback := serverCtx.CanonicalURL.ResolveReference(googleCallbackPath).String()
	googleAuth := google.NewAuthenticator(
		cfg.GoogleClientId, cfg.GoogleClientSecret, googleCallback)
	auth.Bind(serverCtx, e.Group("/auth/google", csrf), googleAuth)

	githubCallbackPath, _ := url.Parse("/auth/github/callback")
	githubCallback := serverCtx.CanonicalURL.ResolveReference(githubCallbackPath).String()
	githubAuth := github.NewAuthenticator(
		cfg.GitHubClientId, cfg.GitHubClientSecret, githubCallback)
	auth.Bind(serverCtx, e.Group("/auth/github", csrf), githubAuth)

	// Public API
	info.Bind(serverCtx, e.Group("/api/v1/info"))

	// API
	api := e.Group("/api", csrf, auth.Required(serverCtx))

	// API - Events
	db.BindList(serverCtx, api.Group("/events/:kind/:list"))

	// API - Event Stream
	manager := eventstream.NewManager()
	manager.Bind(serverCtx, api.Group("/eventstream"))

	// API - Notifiers
	slack.Bind(serverCtx, api.Group("/notifiers/slack"))
	steemitchat.Bind(serverCtx, api.Group("/notifiers/steemit-chat"))

	// Telegram
	botSecret := make([]byte, 256/8)
	if _, err := rand.Read(botSecret); err != nil {
		return nil, nil, errors.Wrap(err, "failed to generate Telegram bot secret")
	}
	botSecretHex := hex.EncodeToString(botSecret)

	botURL, _ := url.Parse("/bots/telegram/both-" + botSecretHex)
	botWebhookURL := serverCtx.CanonicalURL.ResolveReference(botURL)

	bot, err := tgbotapi.NewBotAPI(cfg.TelegramBotToken)
	if err != nil {
		return nil, nil, errors.Wrap(err, "failed to initialize Telegram bot")
	}

	if !strings.Contains(cfg.CanonicalURL, "localhost") {
		if _, err := bot.SetWebhook(tgbotapi.NewWebhook(botWebhookURL.String())); err != nil {
			return nil, nil, errors.Wrap(err, "failed to register Telegram webhook URL")
		}
	}

	telegram.BindWebhook(serverCtx, e.Group(botURL.String()))
	telegram.BindAPI(serverCtx, api.Group("/notifiers/telegram"))

	// API - Profile
	profile.Bind(serverCtx, api.Group("/profile"))

	// Start server
	listener, err := net.Listen("tcp", cfg.ListenAddress)
	if err != nil {
		return nil, nil, errors.Wrap(err, "failed to start the web server")
	}

	ctx := &Context{
		EventStreamManager: manager,
		listener:           listener,
	}

	// Discord
	dg, err := discord.InitBot(&ctx.t, cfg.DiscordBotToken, serverCtx)
	if err != nil {
		return nil, nil, err
	}

	discord.BindAPI(serverCtx, api.Group("/notifiers/discord"))

	// Start listening.
	ctx.t.Go(func() error {
		http.Serve(listener, e)
		return nil
	})

	go func() {
		<-ctx.t.Dying()
		listener.Close()
	}()

	return ctx, dg, nil
}

func (ctx *Context) Interrupt() {
	ctx.t.Kill(nil)
}

func (ctx *Context) Wait() error {
	return ctx.t.Wait()
}
<|endoftext|>


<commit_before>package server

import (
	"math/rand"
	"net"
	"net/http"
	"runtime"
	"strconv"
	"time"

	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/bolt"
	"github.com/influxdata/chronograf/canned"
	"github.com/influxdata/chronograf/influx"
	"github.com/influxdata/chronograf/layouts"
	clog "github.com/influxdata/chronograf/log"
	"github.com/influxdata/chronograf/uuid"
	client "github.com/influxdata/usage-client/v1"
	"github.com/tylerb/graceful"
)

// Server for the chronograf API
type Server struct {
	Host string `long:"host" description:"the IP to listen on" default:"0.0.0.0" env:"HOST"`
	Port int    `long:"port" description:"the port to listen on for insecure connections, defaults to a random value" default:"8888" env:"PORT"`

	/* TODO: add in support for TLS
	TLSHost           string         `long:"tls-host" description:"the IP to listen on for tls, when not specified it's the same as --host" env:"TLS_HOST"`
	TLSPort           int            `long:"tls-port" description:"the port to listen on for secure connections, defaults to a random value" env:"TLS_PORT"`
	TLSCertificate    flags.Filename `long:"tls-certificate" description:"the certificate to use for secure connections" env:"TLS_CERTIFICATE"`
	TLSCertificateKey flags.Filename `long:"tls-key" description:"the private key to use for secure conections" env:"TLS_PRIVATE_KEY"`
	*/

	Develop            bool   `short:"d" long:"develop" description:"Run server in develop mode."`
	BoltPath           string `short:"b" long:"bolt-path" description:"Full path to boltDB file (/var/lib/chronograf/chronograf-v1.db)" env:"BOLT_PATH" default:"chronograf-v1.db"`
	CannedPath         string `short:"c" long:"canned-path" description:"Path to directory of pre-canned application layouts (/usr/share/chronograf/canned)" env:"CANNED_PATH" default:"canned"`
	TokenSecret        string `short:"t" long:"token-secret" description:"Secret to sign tokens" env:"TOKEN_SECRET"`
	GithubClientID     string `short:"i" long:"github-client-id" description:"Github Client ID for OAuth 2 support" env:"GH_CLIENT_ID"`
	GithubClientSecret string `short:"s" long:"github-client-secret" description:"Github Client Secret for OAuth 2 support" env:"GH_CLIENT_SECRET"`
	ReportingDisabled  bool   `short:"r" long:"reporting-disabled" description:"Disable reporting of usage stats (os,arch,version,cluster_id) once every 24hr" env:"REPORTING_DISABLED"`
	LogLevel           string `short:"l" long:"log-level" value-name:"choice" choice:"debug" choice:"info" choice:"warn" choice:"error" choice:"fatal" choice:"panic" default:"info" description:"Set the logging level" env:"LOG_LEVEL"`

	ShowVersion bool `short:"v" long:"version" description:"Show Chronograf version info"`
	BuildInfo   BuildInfo

	Listener net.Listener
	handler  http.Handler
}

type BuildInfo struct {
	Version string
	Commit  string
}

func (s *Server) useAuth() bool {
	return s.TokenSecret != "" && s.GithubClientID != "" && s.GithubClientSecret != ""
}

// Serve starts and runs the chronograf server
func (s *Server) Serve() error {
	logger := clog.New(clog.ParseLevel(s.LogLevel))
	service := openService(s.BoltPath, s.CannedPath, logger)
	s.handler = NewMux(MuxOpts{
		Develop:            s.Develop,
		TokenSecret:        s.TokenSecret,
		GithubClientID:     s.GithubClientID,
		GithubClientSecret: s.GithubClientSecret,
		Logger:             logger,
		UseAuth:            s.useAuth(),
	}, service)

	var err error
	s.Listener, err = net.Listen("tcp", net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
	if err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	httpServer := &graceful.Server{Server: new(http.Server)}
	httpServer.SetKeepAlivesEnabled(true)
	httpServer.TCPKeepAlive = 1 * time.Minute
	httpServer.Handler = s.handler

	if !s.ReportingDisabled {
		go reportUsageStats(s.BuildInfo, logger)
	}

	logger.
		WithField("component", "server").
		Info("Serving chronograf at http://", s.Listener.Addr())

	if err := httpServer.Serve(s.Listener); err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	logger.
		WithField("component", "server").
		Info("Stopped serving chronograf at http://", s.Listener.Addr())

	return nil
}

func openService(boltPath, cannedPath string, logger chronograf.Logger) Service {
	db := bolt.NewClient()
	db.Path = boltPath
	if err := db.Open(); err != nil {
		logger.
			WithField("component", "boltstore").
			Fatal("Unable to open boltdb; is there a chronograf already running?  ", err)
	}

	// These apps are those handled from a directory
	apps := canned.NewApps(cannedPath, &uuid.V4{}, logger)
	// These apps are statically compiled into chronograf
	binApps := &canned.BinLayoutStore{
		Logger: logger,
	}

	// Acts as a front-end to both the bolt layouts, filesystem layouts and binary statically compiled layouts.
	// The idea here is that these stores form a hierarchy in which each is tried sequentially until
	// the operation has success.  So, the database is preferred over filesystem over binary data.
	layouts := &layouts.MultiLayoutStore{
		Stores: []chronograf.LayoutStore{
			db.LayoutStore,
			apps,
			binApps,
		},
	}

	return Service{
		ExplorationStore: db.ExplorationStore,
		SourcesStore:     db.SourcesStore,
		ServersStore:     db.ServersStore,
		TimeSeries: &influx.Client{
			Logger: logger,
		},
		LayoutStore:     layouts,
		AlertRulesStore: db.AlertsStore,
	}
}

// reportUsageStats starts periodic server reporting.
func reportUsageStats(bi BuildInfo, logger chronograf.Logger) {
	rand.Seed(time.Now().UTC().UnixNano())
	serverID := strconv.FormatUint(uint64(rand.Int63()), 10)
	reporter := client.New("")
	u := &client.Usage{
		Product: "chronograf",
		Data: []client.UsageData{
			{
				Values: client.Values{
					"os":         runtime.GOOS,
					"arch":       runtime.GOARCH,
					"version":    bi.Version,
					"cluster_id": serverID,
				},
			},
		},
	}
	l := logger.WithField("component", "usage").
		WithField("reporting_addr", reporter.URL).
		WithField("freq", "24h").
		WithField("stats", "os,arch,version,cluster_id")
	l.Info("Reporting usage stats")
	reporter.Save(u)

	ticker := time.NewTicker(24 * time.Hour)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			l.Debug("Reporting usage stats")
			go reporter.Save(u)
		}
	}
}
<commit_msg>Update reporting tag to use chronograf-ng<commit_after>package server

import (
	"math/rand"
	"net"
	"net/http"
	"runtime"
	"strconv"
	"time"

	"github.com/influxdata/chronograf"
	"github.com/influxdata/chronograf/bolt"
	"github.com/influxdata/chronograf/canned"
	"github.com/influxdata/chronograf/influx"
	"github.com/influxdata/chronograf/layouts"
	clog "github.com/influxdata/chronograf/log"
	"github.com/influxdata/chronograf/uuid"
	client "github.com/influxdata/usage-client/v1"
	"github.com/tylerb/graceful"
)

// Server for the chronograf API
type Server struct {
	Host string `long:"host" description:"the IP to listen on" default:"0.0.0.0" env:"HOST"`
	Port int    `long:"port" description:"the port to listen on for insecure connections, defaults to a random value" default:"8888" env:"PORT"`

	/* TODO: add in support for TLS
	TLSHost           string         `long:"tls-host" description:"the IP to listen on for tls, when not specified it's the same as --host" env:"TLS_HOST"`
	TLSPort           int            `long:"tls-port" description:"the port to listen on for secure connections, defaults to a random value" env:"TLS_PORT"`
	TLSCertificate    flags.Filename `long:"tls-certificate" description:"the certificate to use for secure connections" env:"TLS_CERTIFICATE"`
	TLSCertificateKey flags.Filename `long:"tls-key" description:"the private key to use for secure conections" env:"TLS_PRIVATE_KEY"`
	*/

	Develop            bool   `short:"d" long:"develop" description:"Run server in develop mode."`
	BoltPath           string `short:"b" long:"bolt-path" description:"Full path to boltDB file (/var/lib/chronograf/chronograf-v1.db)" env:"BOLT_PATH" default:"chronograf-v1.db"`
	CannedPath         string `short:"c" long:"canned-path" description:"Path to directory of pre-canned application layouts (/usr/share/chronograf/canned)" env:"CANNED_PATH" default:"canned"`
	TokenSecret        string `short:"t" long:"token-secret" description:"Secret to sign tokens" env:"TOKEN_SECRET"`
	GithubClientID     string `short:"i" long:"github-client-id" description:"Github Client ID for OAuth 2 support" env:"GH_CLIENT_ID"`
	GithubClientSecret string `short:"s" long:"github-client-secret" description:"Github Client Secret for OAuth 2 support" env:"GH_CLIENT_SECRET"`
	ReportingDisabled  bool   `short:"r" long:"reporting-disabled" description:"Disable reporting of usage stats (os,arch,version,cluster_id) once every 24hr" env:"REPORTING_DISABLED"`
	LogLevel           string `short:"l" long:"log-level" value-name:"choice" choice:"debug" choice:"info" choice:"warn" choice:"error" choice:"fatal" choice:"panic" default:"info" description:"Set the logging level" env:"LOG_LEVEL"`

	ShowVersion bool `short:"v" long:"version" description:"Show Chronograf version info"`
	BuildInfo   BuildInfo

	Listener net.Listener
	handler  http.Handler
}

type BuildInfo struct {
	Version string
	Commit  string
}

func (s *Server) useAuth() bool {
	return s.TokenSecret != "" && s.GithubClientID != "" && s.GithubClientSecret != ""
}

// Serve starts and runs the chronograf server
func (s *Server) Serve() error {
	logger := clog.New(clog.ParseLevel(s.LogLevel))
	service := openService(s.BoltPath, s.CannedPath, logger)
	s.handler = NewMux(MuxOpts{
		Develop:            s.Develop,
		TokenSecret:        s.TokenSecret,
		GithubClientID:     s.GithubClientID,
		GithubClientSecret: s.GithubClientSecret,
		Logger:             logger,
		UseAuth:            s.useAuth(),
	}, service)

	var err error
	s.Listener, err = net.Listen("tcp", net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
	if err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	httpServer := &graceful.Server{Server: new(http.Server)}
	httpServer.SetKeepAlivesEnabled(true)
	httpServer.TCPKeepAlive = 1 * time.Minute
	httpServer.Handler = s.handler

	if !s.ReportingDisabled {
		go reportUsageStats(s.BuildInfo, logger)
	}

	logger.
		WithField("component", "server").
		Info("Serving chronograf at http://", s.Listener.Addr())

	if err := httpServer.Serve(s.Listener); err != nil {
		logger.
			WithField("component", "server").
			Error(err)
		return err
	}

	logger.
		WithField("component", "server").
		Info("Stopped serving chronograf at http://", s.Listener.Addr())

	return nil
}

func openService(boltPath, cannedPath string, logger chronograf.Logger) Service {
	db := bolt.NewClient()
	db.Path = boltPath
	if err := db.Open(); err != nil {
		logger.
			WithField("component", "boltstore").
			Fatal("Unable to open boltdb; is there a chronograf already running?  ", err)
	}

	// These apps are those handled from a directory
	apps := canned.NewApps(cannedPath, &uuid.V4{}, logger)
	// These apps are statically compiled into chronograf
	binApps := &canned.BinLayoutStore{
		Logger: logger,
	}

	// Acts as a front-end to both the bolt layouts, filesystem layouts and binary statically compiled layouts.
	// The idea here is that these stores form a hierarchy in which each is tried sequentially until
	// the operation has success.  So, the database is preferred over filesystem over binary data.
	layouts := &layouts.MultiLayoutStore{
		Stores: []chronograf.LayoutStore{
			db.LayoutStore,
			apps,
			binApps,
		},
	}

	return Service{
		ExplorationStore: db.ExplorationStore,
		SourcesStore:     db.SourcesStore,
		ServersStore:     db.ServersStore,
		TimeSeries: &influx.Client{
			Logger: logger,
		},
		LayoutStore:     layouts,
		AlertRulesStore: db.AlertsStore,
	}
}

// reportUsageStats starts periodic server reporting.
func reportUsageStats(bi BuildInfo, logger chronograf.Logger) {
	rand.Seed(time.Now().UTC().UnixNano())
	serverID := strconv.FormatUint(uint64(rand.Int63()), 10)
	reporter := client.New("")
	u := &client.Usage{
		Product: "chronograf-ng",
		Data: []client.UsageData{
			{
				Values: client.Values{
					"os":         runtime.GOOS,
					"arch":       runtime.GOARCH,
					"version":    bi.Version,
					"cluster_id": serverID,
				},
			},
		},
	}
	l := logger.WithField("component", "usage").
		WithField("reporting_addr", reporter.URL).
		WithField("freq", "24h").
		WithField("stats", "os,arch,version,cluster_id")
	l.Info("Reporting usage stats")
	reporter.Save(u)

	ticker := time.NewTicker(24 * time.Hour)
	defer ticker.Stop()
	for {
		select {
		case <-ticker.C:
			l.Debug("Reporting usage stats")
			go reporter.Save(u)
		}
	}
}
<|endoftext|>


<commit_before>/*This package represents varuous backend endpoints
processing requests
*/
package endpoint

import (
	"fmt"
	"github.com/mailgun/vulcan/netutils"
	"net/url"
)

type Endpoint interface {
	GetId() string
	GetUrl() *url.URL
}

type HttpEndpoint struct {
	url *url.URL
	id  string
}

func ParseUrl(in string) (Endpoint, error) {
	url, err := netutils.ParseUrl(in)
	if err != nil {
		return nil, err
	}
	return &HttpEndpoint{url: url, id: fmt.Sprintf("%s://%s", url.Scheme, url.Host)}, nil
}

func MustParseUrl(in string) Endpoint {
	u, err := ParseUrl(in)
	if err != nil {
		panic(err)
	}
	return u
}

func NewHttpEndpoint(in *url.URL) (*HttpEndpoint, error) {
	if in == nil {
		return nil, fmt.Errorf("Provide url")
	}
	return &HttpEndpoint{
		url: netutils.CopyUrl(in),
		id:  fmt.Sprintf("%s://%s", in.Scheme, in.Host)}, nil
}

func (e *HttpEndpoint) GetId() string {
	return e.id
}

func (e *HttpEndpoint) GetUrl() *url.URL {
	return e.url
}
<commit_msg>Remove some unused code<commit_after><|endoftext|>


<commit_before>package main

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

	"github.com/gorilla/websocket"
)

type Server struct {
	sync.Mutex

	config Config

	// the 10 minute timer to clean sending grants
	lastSweepTime     time.Time
}

var upgrader = websocket.Upgrader{}
func clientWebConnect(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		_, _, err := c.ReadMessage()
		if err != nil {
			log.Println("read err:", err)
			break
		}
		break
	}
}

func NewServer(conf Config) *Server {
	var addr = fmt.Sprintf("0.0.0.0:%d", conf.port)

	http.HandleFunc("/", clientWebConnect)
	http.ListenAndServe(addr, nil)

	return &Server{
		config:     conf,
	}
}
<commit_msg>runs web-(socket)-server<commit_after>package main

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

	"github.com/gorilla/websocket"
)

type Server struct {
	sync.Mutex
	webServer http.Server
	config Config

	// the 10 minute timer to clean sending grants
	lastSweepTime     time.Time
}

var upgrader = websocket.Upgrader{}
func clientWebConnect(w http.ResponseWriter, r *http.Request) {
	c, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Print("upgrade:", err)
		return
	}
	defer c.Close()
	for {
		_, _, err := c.ReadMessage()
		if err != nil {
			log.Println("read err:", err)
			break
		}
		break
	}
}

func NewServer(conf Config) *Server {
	addr := fmt.Sprintf("0.0.0.0:%d", conf.port)
	mux := http.NewServeMux()
	mux.HandleFunc("/", clientWebConnect)
	webServer := &http.Server{Addr: addr, Handler: mux}
	webServer.ListenAndServe()

	return &Server{
		config:     conf,
		webServer:		  *webServer,
	}
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"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 moveFile(src, dest string) (err error) {
	srcfile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcfile.Close()
	defer os.Remove(src)

	destfile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destfile.Close()

	if _, err := io.Copy(destfile, srcfile); err != nil {
		return err
	}

	srcinfo, err := os.Stat(src)
	if err != nil {
		return err
	}

	return os.Chmod(dest, srcinfo.Mode())
}

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 private 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 cli.NewExitError(fmt.Sprintf("Error creating index: %s", err.Error()), 1)
		}

		// 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
					}
					if c.GlobalString("clair") != "" {
						wg.Wait()
					}
				}
			}
		}()

		// 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

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

			var createdDate string
			for _, h := range m1.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)*3) * time.Second))
					<-throttle

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

					if err := createVulnStaticPage(r, staticDir, clairURI, repo, tag, m1); 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 := moveFile(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, m schema1.SignedManifest) error {
	// 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
}
<commit_msg>do not hard fail on manifest fail<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"html/template"
	"io"
	"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 moveFile(src, dest string) (err error) {
	srcfile, err := os.Open(src)
	if err != nil {
		return err
	}
	defer srcfile.Close()
	defer os.Remove(src)

	destfile, err := os.Create(dest)
	if err != nil {
		return err
	}
	defer destfile.Close()

	if _, err := io.Copy(destfile, srcfile); err != nil {
		return err
	}

	srcinfo, err := os.Stat(src)
	if err != nil {
		return err
	}

	return os.Chmod(dest, srcinfo.Mode())
}

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 private 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 cli.NewExitError(fmt.Sprintf("Error creating index: %s", err.Error()), 1)
		}

		// 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
					}
					if c.GlobalString("clair") != "" {
						wg.Wait()
					}
				}
			}
		}()

		// 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

			m1, err := r.ManifestV1(repo, tag)
			if err != nil {
				logrus.Warnf("getting v1 manifest for %s:%s failed: %v", repo, tag, err)
			}

			var createdDate string
			for _, h := range m1.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)*3) * time.Second))
					<-throttle

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

					if err := createVulnStaticPage(r, staticDir, clairURI, repo, tag, m1); 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 := moveFile(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, m schema1.SignedManifest) error {
	// 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|>