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

// Package server contains the `pilosa server` subcommand which runs Pilosa
// itself. The purpose of this package is to define an easily tested Command
// object which handles interpreting configuration and setting up all the
// objects that Pilosa needs.
package server

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/httpbroadcast"
	"github.com/pilosa/pilosa/statsd"
)

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

const (
	// DefaultDataDir is the default data directory.
	DefaultDataDir = "~/.pilosa"
)

// Command represents the state of the pilosa server command.
type Command struct {
	Server *pilosa.Server

	// Configuration.
	Config *pilosa.Config

	// Profiling options.
	CPUProfile string
	CPUTime    time.Duration

	// Standard input/output
	*pilosa.CmdIO

	// running will be closed once Command.Run is finished.
	Started chan struct{}
	// Done will be closed when Command.Close() is called
	Done chan struct{}
}

// NewCommand returns a new instance of Main.
func NewCommand(stdin io.Reader, stdout, stderr io.Writer) *Command {
	return &Command{
		Server: pilosa.NewServer(),
		Config: pilosa.NewConfig(),

		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),

		Started: make(chan struct{}),
		Done:    make(chan struct{}),
	}
}

// Run executes the pilosa server.
func (m *Command) Run(args ...string) (err error) {
	defer close(m.Started)
	prefix := "~" + string(filepath.Separator)
	if strings.HasPrefix(m.Config.DataDir, prefix) {
		HomeDir := os.Getenv("HOME")
		if HomeDir == "" {
			return errors.New("data directory not specified and no home dir available")
		}
		m.Config.DataDir = filepath.Join(HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
	}

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return err
	}

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return fmt.Errorf("server.Open: %v", err)
	}
	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", m.Server.Host)
	return nil
}

// SetupServer use the cluster configuration to setup this server
func (m *Command) SetupServer() error {
	var err error
	cluster := pilosa.NewCluster()
	cluster.ReplicaN = m.Config.Cluster.ReplicaN

	for _, hostport := range m.Config.Cluster.Hosts {
		cluster.Nodes = append(cluster.Nodes, &pilosa.Node{Host: hostport})
	}
	// TODO: if InternalHosts is not provided then pilosa.Node.InternalHost is empty.
	// This will throw an error when trying to Broadcast messages over HTTP.
	// One option may be to fall back to using host from hostport + config.InternalPort.
	for i, internalhostport := range m.Config.Cluster.InternalHosts {
		cluster.Nodes[i].InternalHost = internalhostport
	}
	m.Server.Cluster = cluster

	// Setup logging output.
	if m.Config.LogPath == "" {
		m.Server.LogOutput = m.Stderr
	} else {
		logFile, err := os.OpenFile(m.Config.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
		if err != nil {
			return err
		}
		m.Server.LogOutput = logFile
	}

	// Configure holder.
	fmt.Fprintf(m.Stderr, "Using data from: %s\n", m.Config.DataDir)
	m.Server.Holder.Path = m.Config.DataDir
	m.Server.MetricInterval = time.Duration(m.Config.Metric.PollingInterval)
	m.Server.Holder.Stats, err = NewStatsClient(m.Config.Metric.Service, m.Config.Metric.Host)
	if err != nil {
		return err
	}

	m.Server.Holder.Stats.SetLogger(m.Server.LogOutput)

	// Copy configuration flags.
	m.Server.MaxWritesPerRequest = m.Config.MaxWritesPerRequest

	m.Server.Host, err = normalizeHost(m.Config.Host)
	if err != nil {
		return err
	}

	// Set internal port (string).
	internalPortStr := pilosa.DefaultInternalPort
	if m.Config.Cluster.InternalPort != "" {
		internalPortStr = m.Config.Cluster.InternalPort
	}

	switch m.Config.Cluster.Type {
	case "http":
		m.Server.Broadcaster = httpbroadcast.NewHTTPBroadcaster(m.Server, internalPortStr)
		m.Server.BroadcastReceiver = httpbroadcast.NewHTTPBroadcastReceiver(internalPortStr, m.Stderr)
		m.Server.Cluster.NodeSet = httpbroadcast.NewHTTPNodeSet()
		err := m.Server.Cluster.NodeSet.(*httpbroadcast.HTTPNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	case "gossip":
		gossipPort, err := strconv.Atoi(internalPortStr)
		if err != nil {
			return err
		}
		gossipSeed := pilosa.DefaultHost
		if m.Config.Cluster.GossipSeed != "" {
			gossipSeed = m.Config.Cluster.GossipSeed
		}
		// get the host portion of addr to use for binding
		gossipHost, _, err := net.SplitHostPort(m.Config.Host)
		if err != nil {
			gossipHost = m.Config.Host
		}
		gossipNodeSet := gossip.NewGossipNodeSet(m.Config.Host, gossipHost, gossipPort, gossipSeed, m.Server)
		m.Server.Cluster.NodeSet = gossipNodeSet
		m.Server.Broadcaster = gossipNodeSet
		m.Server.BroadcastReceiver = gossipNodeSet
	case "static", "":
		m.Server.Broadcaster = pilosa.NopBroadcaster
		m.Server.Cluster.NodeSet = pilosa.NewStaticNodeSet()
		m.Server.BroadcastReceiver = pilosa.NopBroadcastReceiver
		err := m.Server.Cluster.NodeSet.(*pilosa.StaticNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("'%v' is not a supported value for broadcaster type", m.Config.Cluster.Type)
	}

	// Set configuration options.
	m.Server.AntiEntropyInterval = time.Duration(m.Config.AntiEntropy.Interval)
	return nil
}

func normalizeHost(host string) (string, error) {
	if !strings.Contains(host, ":") {
		host = host + ":"
	} else if strings.Contains(host, "://") {
		if strings.HasPrefix(host, "http://") {
			host = host[7:]
		} else {
			return "", fmt.Errorf("invalid scheme or host: '%s'. use the format [http://]<host>:<port>", host)
		}
	}
	return host, nil
}

// Close shuts down the server.
func (m *Command) Close() error {
	var logErr error
	serveErr := m.Server.Close()
	logOutput := m.Server.LogOutput
	if closer, ok := logOutput.(io.Closer); ok {
		logErr = closer.Close()
	}
	close(m.Done)
	if serveErr != nil && logErr != nil {
		return fmt.Errorf("closing server: '%v', closing logs: '%v'", serveErr, logErr)
	} else if logErr != nil {
		return logErr
	}
	return serveErr
}

// NewStatsClient creates a stats client from the config
func NewStatsClient(name string, host string) (pilosa.StatsClient, error) {
	switch name {
	case "expvar":
		return pilosa.NewExpvarStatsClient(), nil
	case "statsd":
		return statsd.NewStatsClient(host)
	default:
		return pilosa.NopStatsClient, nil
	}
}
<commit_msg>set long query time parameter at startup<commit_after>// Copyright 2017 Pilosa Corp.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

// Package server contains the `pilosa server` subcommand which runs Pilosa
// itself. The purpose of this package is to define an easily tested Command
// object which handles interpreting configuration and setting up all the
// objects that Pilosa needs.
package server

import (
	"errors"
	"fmt"
	"io"
	"math/rand"
	"net"
	"os"
	"path/filepath"
	"strconv"
	"strings"
	"time"

	"github.com/pilosa/pilosa"
	"github.com/pilosa/pilosa/gossip"
	"github.com/pilosa/pilosa/httpbroadcast"
	"github.com/pilosa/pilosa/statsd"
)

func init() {
	rand.Seed(time.Now().UTC().UnixNano())
}

const (
	// DefaultDataDir is the default data directory.
	DefaultDataDir = "~/.pilosa"
)

// Command represents the state of the pilosa server command.
type Command struct {
	Server *pilosa.Server

	// Configuration.
	Config *pilosa.Config

	// Profiling options.
	CPUProfile string
	CPUTime    time.Duration

	// Standard input/output
	*pilosa.CmdIO

	// running will be closed once Command.Run is finished.
	Started chan struct{}
	// Done will be closed when Command.Close() is called
	Done chan struct{}
}

// NewCommand returns a new instance of Main.
func NewCommand(stdin io.Reader, stdout, stderr io.Writer) *Command {
	return &Command{
		Server: pilosa.NewServer(),
		Config: pilosa.NewConfig(),

		CmdIO: pilosa.NewCmdIO(stdin, stdout, stderr),

		Started: make(chan struct{}),
		Done:    make(chan struct{}),
	}
}

// Run executes the pilosa server.
func (m *Command) Run(args ...string) (err error) {
	defer close(m.Started)
	prefix := "~" + string(filepath.Separator)
	if strings.HasPrefix(m.Config.DataDir, prefix) {
		HomeDir := os.Getenv("HOME")
		if HomeDir == "" {
			return errors.New("data directory not specified and no home dir available")
		}
		m.Config.DataDir = filepath.Join(HomeDir, strings.TrimPrefix(m.Config.DataDir, prefix))
	}

	// SetupServer
	err = m.SetupServer()
	if err != nil {
		return err
	}

	// Initialize server.
	if err = m.Server.Open(); err != nil {
		return fmt.Errorf("server.Open: %v", err)
	}
	fmt.Fprintf(m.Stderr, "Listening as http://%s\n", m.Server.Host)
	return nil
}

// SetupServer use the cluster configuration to setup this server
func (m *Command) SetupServer() error {
	var err error
	cluster := pilosa.NewCluster()
	cluster.ReplicaN = m.Config.Cluster.ReplicaN

	for _, hostport := range m.Config.Cluster.Hosts {
		cluster.Nodes = append(cluster.Nodes, &pilosa.Node{Host: hostport})
	}
	// TODO: if InternalHosts is not provided then pilosa.Node.InternalHost is empty.
	// This will throw an error when trying to Broadcast messages over HTTP.
	// One option may be to fall back to using host from hostport + config.InternalPort.
	for i, internalhostport := range m.Config.Cluster.InternalHosts {
		cluster.Nodes[i].InternalHost = internalhostport
	}
	m.Server.Cluster = cluster

	// Setup logging output.
	if m.Config.LogPath == "" {
		m.Server.LogOutput = m.Stderr
	} else {
		logFile, err := os.OpenFile(m.Config.LogPath, os.O_RDWR|os.O_CREATE|os.O_APPEND, 0600)
		if err != nil {
			return err
		}
		m.Server.LogOutput = logFile
	}

	// Configure holder.
	fmt.Fprintf(m.Stderr, "Using data from: %s\n", m.Config.DataDir)
	m.Server.Holder.Path = m.Config.DataDir
	m.Server.MetricInterval = time.Duration(m.Config.Metric.PollingInterval)
	m.Server.Holder.Stats, err = NewStatsClient(m.Config.Metric.Service, m.Config.Metric.Host)
	if err != nil {
		return err
	}

	m.Server.Holder.Stats.SetLogger(m.Server.LogOutput)

	// Copy configuration flags.
	m.Server.MaxWritesPerRequest = m.Config.MaxWritesPerRequest

	m.Server.Host, err = normalizeHost(m.Config.Host)
	if err != nil {
		return err
	}

	// Set internal port (string).
	internalPortStr := pilosa.DefaultInternalPort
	if m.Config.Cluster.InternalPort != "" {
		internalPortStr = m.Config.Cluster.InternalPort
	}

	switch m.Config.Cluster.Type {
	case "http":
		m.Server.Broadcaster = httpbroadcast.NewHTTPBroadcaster(m.Server, internalPortStr)
		m.Server.BroadcastReceiver = httpbroadcast.NewHTTPBroadcastReceiver(internalPortStr, m.Stderr)
		m.Server.Cluster.NodeSet = httpbroadcast.NewHTTPNodeSet()
		err := m.Server.Cluster.NodeSet.(*httpbroadcast.HTTPNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	case "gossip":
		gossipPort, err := strconv.Atoi(internalPortStr)
		if err != nil {
			return err
		}
		gossipSeed := pilosa.DefaultHost
		if m.Config.Cluster.GossipSeed != "" {
			gossipSeed = m.Config.Cluster.GossipSeed
		}
		// get the host portion of addr to use for binding
		gossipHost, _, err := net.SplitHostPort(m.Config.Host)
		if err != nil {
			gossipHost = m.Config.Host
		}
		gossipNodeSet := gossip.NewGossipNodeSet(m.Config.Host, gossipHost, gossipPort, gossipSeed, m.Server)
		m.Server.Cluster.NodeSet = gossipNodeSet
		m.Server.Broadcaster = gossipNodeSet
		m.Server.BroadcastReceiver = gossipNodeSet
	case "static", "":
		m.Server.Broadcaster = pilosa.NopBroadcaster
		m.Server.Cluster.NodeSet = pilosa.NewStaticNodeSet()
		m.Server.BroadcastReceiver = pilosa.NopBroadcastReceiver
		err := m.Server.Cluster.NodeSet.(*pilosa.StaticNodeSet).Join(m.Server.Cluster.Nodes)
		if err != nil {
			return err
		}
	default:
		return fmt.Errorf("'%v' is not a supported value for broadcaster type", m.Config.Cluster.Type)
	}

	// Set configuration options.
	m.Server.AntiEntropyInterval = time.Duration(m.Config.AntiEntropy.Interval)
	m.Server.Cluster.LongQueryTime = time.Duration(m.Config.Cluster.LongQueryTime)
	return nil
}

func normalizeHost(host string) (string, error) {
	if !strings.Contains(host, ":") {
		host = host + ":"
	} else if strings.Contains(host, "://") {
		if strings.HasPrefix(host, "http://") {
			host = host[7:]
		} else {
			return "", fmt.Errorf("invalid scheme or host: '%s'. use the format [http://]<host>:<port>", host)
		}
	}
	return host, nil
}

// Close shuts down the server.
func (m *Command) Close() error {
	var logErr error
	serveErr := m.Server.Close()
	logOutput := m.Server.LogOutput
	if closer, ok := logOutput.(io.Closer); ok {
		logErr = closer.Close()
	}
	close(m.Done)
	if serveErr != nil && logErr != nil {
		return fmt.Errorf("closing server: '%v', closing logs: '%v'", serveErr, logErr)
	} else if logErr != nil {
		return logErr
	}
	return serveErr
}

// NewStatsClient creates a stats client from the config
func NewStatsClient(name string, host string) (pilosa.StatsClient, error) {
	switch name {
	case "expvar":
		return pilosa.NewExpvarStatsClient(), nil
	case "statsd":
		return statsd.NewStatsClient(host)
	default:
		return pilosa.NopStatsClient, nil
	}
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/marcusolsson/goddd/booking"
	"github.com/marcusolsson/goddd/cargo"
	"github.com/marcusolsson/goddd/location"
	"github.com/marcusolsson/goddd/routing"

	"github.com/go-martini/martini"
	"github.com/martini-contrib/binding"
	"github.com/martini-contrib/cors"
	"github.com/martini-contrib/render"
)

type eventDTO struct {
	Description string `json:"description"`
	Expected    bool   `json:"expected"`
}

type locationDTO struct {
	UNLocode string `json:"locode"`
	Name     string `json:"name"`
}

type cargoDTO struct {
	TrackingId           string     `json:"trackingId"`
	StatusText           string     `json:"statusText"`
	Origin               string     `json:"origin"`
	Destination          string     `json:"destination"`
	ETA                  string     `json:"eta"`
	NextExpectedActivity string     `json:"nextExpectedActivity"`
	Misrouted            bool       `json:"misrouted"`
	Routed               bool       `json:"routed"`
	ArrivalDeadline      string     `json:"arrivalDeadline"`
	Events               []eventDTO `json:"events"`
	Legs                 []legDTO   `json:"legs"`
}

type legDTO struct {
	VoyageNumber string `json:"voyage"`
	From         string `json:"from"`
	To           string `json:"to"`
	LoadTime     string `json:"loadTime"`
	UnloadTime   string `json:"unloadTime"`
}

type routeCandidate struct {
	Legs []legDTO `json:"legs"`
}

// Assemble converts the Cargo domain object to a serializable DTO.
func Assemble(c cargo.Cargo) cargoDTO {
	eta := time.Date(2009, time.March, 12, 12, 0, 0, 0, time.UTC)
	dto := cargoDTO{
		TrackingId:           string(c.TrackingId),
		StatusText:           fmt.Sprintf("%s %s", cargo.InPort, c.Origin.Name),
		Origin:               c.Origin.Name,
		Destination:          c.RouteSpecification.Destination.Name,
		ETA:                  eta.Format(time.RFC3339),
		NextExpectedActivity: "Next expected activity is to load cargo onto voyage 0200T in New York",
		Misrouted:            c.Delivery.RoutingStatus == cargo.Misrouted,
		Routed:               !c.Itinerary.IsEmpty(),
		ArrivalDeadline:      c.ArrivalDeadline.Format(time.RFC3339),
	}

	legs := make([]legDTO, 0)
	for _, l := range c.Itinerary.Legs {
		legs = append(legs, legDTO{
			From: string(l.LoadLocation.UNLocode),
			To:   string(l.UnloadLocation.UNLocode),
		})
	}
	dto.Legs = legs

	dto.Events = make([]eventDTO, 3)
	dto.Events[0] = eventDTO{Description: "Received in Hongkong, at 3/1/09 12:00 AM.", Expected: true}
	dto.Events[1] = eventDTO{Description: "Loaded onto voyage 0100S in Hongkong, at 3/2/09 12:00 AM."}
	dto.Events[2] = eventDTO{Description: "Unloaded off voyage 0100S in New York, at 3/5/09 12:00 AM."}

	return dto
}

type JSONObject map[string]interface{}

var (
	ResourceNotFound              = JSONObject{"error": "The specified resource does not exist."}
	MissingRequiredQueryParameter = JSONObject{"error": "A required query parameter was not specified for this request."}
	InvalidInput                  = JSONObject{"error": "One of the request inputs is not valid."}
)

// TODO: Globals are bad!
var (
	cargoRepository    = cargo.NewCargoRepository()
	locationRepository = location.NewLocationRepository()
	routingService     = routing.NewRoutingService(locationRepository)
	bookingService     = booking.NewBookingService(cargoRepository, locationRepository, routingService)
)

func RegisterHandlers() {

	// Store some sample cargos.
	storeTestData(cargoRepository)

	m := martini.Classic()

	m.Use(cors.Allow(&cors.Options{
		AllowOrigins: []string{"*"},
		AllowMethods: []string{"GET", "POST", "OPTIONS"},
		AllowHeaders: []string{"Origin"},
	}))

	m.Use(martini.Static("app"))
	m.Use(render.Renderer(render.Options{
		IndentJSON: true,
	}))

	// GET /cargos
	// Returns an array of all booked cargos.
	m.Get("/cargos", func(r render.Render) {
		cargos := cargoRepository.FindAll()
		dtos := make([]cargoDTO, len(cargos))

		for i, c := range cargos {
			dtos[i] = Assemble(c)
		}

		r.JSON(200, dtos)
	})

	// GET /cargos/:id
	// Finds and returns a cargo with a specified tracking id.
	m.Get("/cargos/:id", func(params martini.Params, r render.Render) {
		trackingId := cargo.TrackingId(params["id"])
		c, err := cargoRepository.Find(trackingId)

		if err != nil {
			r.JSON(404, ResourceNotFound)
		} else {
			r.JSON(200, Assemble(c))
		}
	})

	// POST /cargos/:id/change_destination
	// Updates the route specification of a cargo with a new destination.
	m.Post("/cargos/:id/change_destination", func(req *http.Request, params martini.Params, r render.Render) {
		v := QueryParams(req.URL.Query())
		found, missing := v.validateQueryParams("destination")

		if len(missing) > 0 {
			e := MissingRequiredQueryParameter
			e["missing"] = missing
			r.JSON(400, e)
			return
		}

		var (
			trackingId  = cargo.TrackingId(params["id"])
			destination = location.UNLocode(fmt.Sprintf("%s", found["destination"]))
		)

		if err := bookingService.ChangeDestination(trackingId, destination); err != nil {
			r.JSON(400, InvalidInput)
			return
		}

		r.JSON(200, JSONObject{})
	})

	// POST /cargos/:id/assign_to_route
	// Assigns the cargo to a route.
	m.Post("/cargos/:id/assign_to_route", binding.Bind(routeCandidate{}), func(rc routeCandidate, params martini.Params, r render.Render) {
		trackingId := cargo.TrackingId(params["id"])

		legs := make([]cargo.Leg, 0)
		for _, l := range rc.Legs {

			var (
				loadLocation   = locationRepository.Find(location.UNLocode(l.From))
				unloadLocation = locationRepository.Find(location.UNLocode(l.To))
			)

			legs = append(legs, cargo.Leg{
				LoadLocation:   loadLocation,
				UnloadLocation: unloadLocation,
			})
		}

		itinerary := cargo.Itinerary{Legs: legs}

		if err := bookingService.AssignCargoToRoute(itinerary, trackingId); err != nil {
			r.JSON(400, InvalidInput)
			return
		}

		r.JSON(200, itinerary)
	})

	// GET /cargos/:id/request_routes
	// Requests the possible routes for a booked cargo.
	m.Get("/cargos/:id/request_routes", func(params martini.Params, r render.Render) {
		trackingId := cargo.TrackingId(params["id"])
		itineraries := bookingService.RequestPossibleRoutesForCargo(trackingId)

		candidates := make([]routeCandidate, 0)
		for _, itin := range itineraries {
			legs := make([]legDTO, 0)
			for _, leg := range itin.Legs {
				legs = append(legs, legDTO{
					VoyageNumber: "S0001",
					From:         string(leg.LoadLocation.UNLocode),
					To:           string(leg.UnloadLocation.UNLocode),
					LoadTime:     "N/A",
					UnloadTime:   "N/A",
				})
			}
			candidates = append(candidates, routeCandidate{Legs: legs})
		}

		r.JSON(200, candidates)
	})

	// POST /cargos
	// Books a cargo from an origin to a destination within a specified arrival deadline.
	m.Post("/cargos", func(req *http.Request, r render.Render) {
		v := QueryParams(req.URL.Query())
		found, missing := v.validateQueryParams("origin", "destination", "arrivalDeadline")

		if len(missing) > 0 {
			e := MissingRequiredQueryParameter
			e["missing"] = missing
			r.JSON(400, e)
			return
		}

		var (
			origin      = location.UNLocode(fmt.Sprintf("%s", found["origin"]))
			destination = location.UNLocode(fmt.Sprintf("%s", found["destination"]))
		)

		millis, _ := strconv.ParseInt(fmt.Sprintf("%s", found["arrivalDeadline"]), 10, 64)
		arrivalDeadline := time.Unix(millis/1000, 0)

		trackingId, err := bookingService.BookNewCargo(origin, destination, arrivalDeadline)

		if err != nil {
			r.JSON(400, InvalidInput)
			return
		}

		c, err := cargoRepository.Find(trackingId)

		if err != nil {
			r.JSON(404, ResourceNotFound)
			return
		}

		r.JSON(200, Assemble(c))
	})

	// GET /locations
	// Returns an array of known locations.
	m.Get("/locations", func(r render.Render) {
		locationRepository := location.NewLocationRepository()
		locations := locationRepository.FindAll()

		dtos := make([]locationDTO, len(locations))
		for i, loc := range locations {
			dtos[i] = locationDTO{
				UNLocode: string(loc.UNLocode),
				Name:     loc.Name,
			}
		}

		r.JSON(200, dtos)
	})

	http.Handle("/", m)
}

func storeTestData(r cargo.CargoRepository) {
	test1 := cargo.NewCargo("FTL456", cargo.RouteSpecification{
		Origin:      location.Melbourne,
		Destination: location.Stockholm,
	})
	r.Store(*test1)

	test2 := cargo.NewCargo("ABC123", cargo.RouteSpecification{
		Origin:      location.Stockholm,
		Destination: location.Hongkong,
	})
	r.Store(*test2)
}

type QueryParams url.Values

func (p QueryParams) validateQueryParams(params ...string) (found JSONObject, missing []string) {
	found = make(JSONObject)
	missing = make([]string, 0)

	for _, param := range params {
		s := url.Values(p).Get(param)
		if len(s) > 0 {
			found[param] = s
		} else {
			missing = append(missing, param)
		}
	}
	return found, missing
}
<commit_msg>Add Content-Type as allowed header.<commit_after>package server

import (
	"fmt"
	"net/http"
	"net/url"
	"strconv"
	"time"

	"github.com/marcusolsson/goddd/booking"
	"github.com/marcusolsson/goddd/cargo"
	"github.com/marcusolsson/goddd/location"
	"github.com/marcusolsson/goddd/routing"

	"github.com/go-martini/martini"
	"github.com/martini-contrib/binding"
	"github.com/martini-contrib/cors"
	"github.com/martini-contrib/render"
)

type eventDTO struct {
	Description string `json:"description"`
	Expected    bool   `json:"expected"`
}

type locationDTO struct {
	UNLocode string `json:"locode"`
	Name     string `json:"name"`
}

type cargoDTO struct {
	TrackingId           string     `json:"trackingId"`
	StatusText           string     `json:"statusText"`
	Origin               string     `json:"origin"`
	Destination          string     `json:"destination"`
	ETA                  string     `json:"eta"`
	NextExpectedActivity string     `json:"nextExpectedActivity"`
	Misrouted            bool       `json:"misrouted"`
	Routed               bool       `json:"routed"`
	ArrivalDeadline      string     `json:"arrivalDeadline"`
	Events               []eventDTO `json:"events"`
	Legs                 []legDTO   `json:"legs"`
}

type legDTO struct {
	VoyageNumber string `json:"voyage"`
	From         string `json:"from"`
	To           string `json:"to"`
	LoadTime     string `json:"loadTime"`
	UnloadTime   string `json:"unloadTime"`
}

type routeCandidate struct {
	Legs []legDTO `json:"legs"`
}

// Assemble converts the Cargo domain object to a serializable DTO.
func Assemble(c cargo.Cargo) cargoDTO {
	eta := time.Date(2009, time.March, 12, 12, 0, 0, 0, time.UTC)
	dto := cargoDTO{
		TrackingId:           string(c.TrackingId),
		StatusText:           fmt.Sprintf("%s %s", cargo.InPort, c.Origin.Name),
		Origin:               c.Origin.Name,
		Destination:          c.RouteSpecification.Destination.Name,
		ETA:                  eta.Format(time.RFC3339),
		NextExpectedActivity: "Next expected activity is to load cargo onto voyage 0200T in New York",
		Misrouted:            c.Delivery.RoutingStatus == cargo.Misrouted,
		Routed:               !c.Itinerary.IsEmpty(),
		ArrivalDeadline:      c.ArrivalDeadline.Format(time.RFC3339),
	}

	legs := make([]legDTO, 0)
	for _, l := range c.Itinerary.Legs {
		legs = append(legs, legDTO{
			From: string(l.LoadLocation.UNLocode),
			To:   string(l.UnloadLocation.UNLocode),
		})
	}
	dto.Legs = legs

	dto.Events = make([]eventDTO, 3)
	dto.Events[0] = eventDTO{Description: "Received in Hongkong, at 3/1/09 12:00 AM.", Expected: true}
	dto.Events[1] = eventDTO{Description: "Loaded onto voyage 0100S in Hongkong, at 3/2/09 12:00 AM."}
	dto.Events[2] = eventDTO{Description: "Unloaded off voyage 0100S in New York, at 3/5/09 12:00 AM."}

	return dto
}

type JSONObject map[string]interface{}

var (
	ResourceNotFound              = JSONObject{"error": "The specified resource does not exist."}
	MissingRequiredQueryParameter = JSONObject{"error": "A required query parameter was not specified for this request."}
	InvalidInput                  = JSONObject{"error": "One of the request inputs is not valid."}
)

// TODO: Globals are bad!
var (
	cargoRepository    = cargo.NewCargoRepository()
	locationRepository = location.NewLocationRepository()
	routingService     = routing.NewRoutingService(locationRepository)
	bookingService     = booking.NewBookingService(cargoRepository, locationRepository, routingService)
)

func RegisterHandlers() {

	// Store some sample cargos.
	storeTestData(cargoRepository)

	m := martini.Classic()

	m.Use(cors.Allow(&cors.Options{
		AllowOrigins: []string{"*"},
		AllowMethods: []string{"GET", "POST", "OPTIONS"},
		AllowHeaders: []string{"Origin, Content-Type"},
	}))

	m.Use(martini.Static("app"))
	m.Use(render.Renderer(render.Options{
		IndentJSON: true,
	}))

	// GET /cargos
	// Returns an array of all booked cargos.
	m.Get("/cargos", func(r render.Render) {
		cargos := cargoRepository.FindAll()
		dtos := make([]cargoDTO, len(cargos))

		for i, c := range cargos {
			dtos[i] = Assemble(c)
		}

		r.JSON(200, dtos)
	})

	// GET /cargos/:id
	// Finds and returns a cargo with a specified tracking id.
	m.Get("/cargos/:id", func(params martini.Params, r render.Render) {
		trackingId := cargo.TrackingId(params["id"])
		c, err := cargoRepository.Find(trackingId)

		if err != nil {
			r.JSON(404, ResourceNotFound)
		} else {
			r.JSON(200, Assemble(c))
		}
	})

	// POST /cargos/:id/change_destination
	// Updates the route specification of a cargo with a new destination.
	m.Post("/cargos/:id/change_destination", func(req *http.Request, params martini.Params, r render.Render) {
		v := QueryParams(req.URL.Query())
		found, missing := v.validateQueryParams("destination")

		if len(missing) > 0 {
			e := MissingRequiredQueryParameter
			e["missing"] = missing
			r.JSON(400, e)
			return
		}

		var (
			trackingId  = cargo.TrackingId(params["id"])
			destination = location.UNLocode(fmt.Sprintf("%s", found["destination"]))
		)

		if err := bookingService.ChangeDestination(trackingId, destination); err != nil {
			r.JSON(400, InvalidInput)
			return
		}

		r.JSON(200, JSONObject{})
	})

	// POST /cargos/:id/assign_to_route
	// Assigns the cargo to a route.
	m.Post("/cargos/:id/assign_to_route", binding.Bind(routeCandidate{}), func(rc routeCandidate, params martini.Params, r render.Render) {
		trackingId := cargo.TrackingId(params["id"])

		legs := make([]cargo.Leg, 0)
		for _, l := range rc.Legs {

			var (
				loadLocation   = locationRepository.Find(location.UNLocode(l.From))
				unloadLocation = locationRepository.Find(location.UNLocode(l.To))
			)

			legs = append(legs, cargo.Leg{
				LoadLocation:   loadLocation,
				UnloadLocation: unloadLocation,
			})
		}

		itinerary := cargo.Itinerary{Legs: legs}

		if err := bookingService.AssignCargoToRoute(itinerary, trackingId); err != nil {
			r.JSON(400, InvalidInput)
			return
		}

		r.JSON(200, itinerary)
	})

	// GET /cargos/:id/request_routes
	// Requests the possible routes for a booked cargo.
	m.Get("/cargos/:id/request_routes", func(params martini.Params, r render.Render) {
		trackingId := cargo.TrackingId(params["id"])
		itineraries := bookingService.RequestPossibleRoutesForCargo(trackingId)

		candidates := make([]routeCandidate, 0)
		for _, itin := range itineraries {
			legs := make([]legDTO, 0)
			for _, leg := range itin.Legs {
				legs = append(legs, legDTO{
					VoyageNumber: "S0001",
					From:         string(leg.LoadLocation.UNLocode),
					To:           string(leg.UnloadLocation.UNLocode),
					LoadTime:     "N/A",
					UnloadTime:   "N/A",
				})
			}
			candidates = append(candidates, routeCandidate{Legs: legs})
		}

		r.JSON(200, candidates)
	})

	// POST /cargos
	// Books a cargo from an origin to a destination within a specified arrival deadline.
	m.Post("/cargos", func(req *http.Request, r render.Render) {
		v := QueryParams(req.URL.Query())
		found, missing := v.validateQueryParams("origin", "destination", "arrivalDeadline")

		if len(missing) > 0 {
			e := MissingRequiredQueryParameter
			e["missing"] = missing
			r.JSON(400, e)
			return
		}

		var (
			origin      = location.UNLocode(fmt.Sprintf("%s", found["origin"]))
			destination = location.UNLocode(fmt.Sprintf("%s", found["destination"]))
		)

		millis, _ := strconv.ParseInt(fmt.Sprintf("%s", found["arrivalDeadline"]), 10, 64)
		arrivalDeadline := time.Unix(millis/1000, 0)

		trackingId, err := bookingService.BookNewCargo(origin, destination, arrivalDeadline)

		if err != nil {
			r.JSON(400, InvalidInput)
			return
		}

		c, err := cargoRepository.Find(trackingId)

		if err != nil {
			r.JSON(404, ResourceNotFound)
			return
		}

		r.JSON(200, Assemble(c))
	})

	// GET /locations
	// Returns an array of known locations.
	m.Get("/locations", func(r render.Render) {
		locationRepository := location.NewLocationRepository()
		locations := locationRepository.FindAll()

		dtos := make([]locationDTO, len(locations))
		for i, loc := range locations {
			dtos[i] = locationDTO{
				UNLocode: string(loc.UNLocode),
				Name:     loc.Name,
			}
		}

		r.JSON(200, dtos)
	})

	http.Handle("/", m)
}

func storeTestData(r cargo.CargoRepository) {
	test1 := cargo.NewCargo("FTL456", cargo.RouteSpecification{
		Origin:      location.Melbourne,
		Destination: location.Stockholm,
	})
	r.Store(*test1)

	test2 := cargo.NewCargo("ABC123", cargo.RouteSpecification{
		Origin:      location.Stockholm,
		Destination: location.Hongkong,
	})
	r.Store(*test2)
}

type QueryParams url.Values

func (p QueryParams) validateQueryParams(params ...string) (found JSONObject, missing []string) {
	found = make(JSONObject)
	missing = make([]string, 0)

	for _, param := range params {
		s := url.Values(p).Get(param)
		if len(s) > 0 {
			found[param] = s
		} else {
			missing = append(missing, param)
		}
	}
	return found, missing
}
<|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/rs/cors"
	"github.com/unrolled/render"
	"html/template"
	"net/http"
	"strings"
)

type Server struct {
	ts TokenStore
	*mux.Router
}
type TokenStore interface {
	Validate(token auth.Token) bool
}

func NewServer(ts TokenStore) *Server {
	return &Server{
		ts,
		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 string, htmls []interface{}, funcs template.FuncMap) {
	renderer := render.New(render.Options{
		Directory:     appgo.Conf.TemplatePath,
		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) AddProxy(path string, handler http.Handler) {
	s.PathPrefix(path).Handler(http.StripPrefix(path, handler))
}

func (s *Server) AddStatic(path, fileDir string) {
	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", "text/plain; charset=utf-8")
		w.Write(content)
	}
	s.HandleFunc("/.well-known/apple-app-site-association", f)
}

func (s *Server) Serve() {
	n := negroni.New()
	n.Use(negroni.NewRecovery())
	n.Use(negronilogrus.NewCustomMiddleware(
		appgo.Conf.LogLevel, &log.TextFormatter{}, "appgo"))
	n.Use(cors.New(corsOptions()))
	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>change apple-app-site-association url<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/rs/cors"
	"github.com/unrolled/render"
	"html/template"
	"net/http"
	"strings"
)

type Server struct {
	ts TokenStore
	*mux.Router
}
type TokenStore interface {
	Validate(token auth.Token) bool
}

func NewServer(ts TokenStore) *Server {
	return &Server{
		ts,
		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 string, htmls []interface{}, funcs template.FuncMap) {
	renderer := render.New(render.Options{
		Directory:     appgo.Conf.TemplatePath,
		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) AddProxy(path string, handler http.Handler) {
	s.PathPrefix(path).Handler(http.StripPrefix(path, handler))
}

func (s *Server) AddStatic(path, fileDir string) {
	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; charset=utf-8")
		w.Write(content)
	}
	s.HandleFunc("/apple-app-site-association", f)
}

func (s *Server) Serve() {
	n := negroni.New()
	n.Use(negroni.NewRecovery())
	n.Use(negronilogrus.NewCustomMiddleware(
		appgo.Conf.LogLevel, &log.TextFormatter{}, "appgo"))
	n.Use(cors.New(corsOptions()))
	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 (
	"expvar"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"context"

	afex "github.com/afex/hystrix-go/hystrix"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/justinas/alice"
	gsh "github.com/mchudgins/go-service-helper/handlers"
	"github.com/mchudgins/playground/pkg/healthz"
	"github.com/mwitkow/go-grpc-middleware"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"go.uber.org/zap"
	xcontext "golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type Config struct {
	Insecure          bool
	CertFilename      string
	KeyFilename       string
	HTTPListenPort    int
	RPCListenPort     int
	MetricsListenPort int
	RPCRegister       RPCRegistration
	logger            *zap.Logger
	rpcServer         *grpc.Server
	httpServer        *http.Server
	metricsServer     *http.Server
}

type Option func(*Config) error

type RPCRegistration func(*grpc.Server) error

func WithRPCServer(fn RPCRegistration) Option {
	return func(cfg *Config) error {
		/*
			echoServer, err := NewServer(p.logger)
			if err != nil {
				cfg.logger.Panic("while creating new EchoServer", zap.Error(err))
			}
			rpc.RegisterEchoServiceServer(s, echoServer)
		*/
		cfg.RPCRegister = fn

		return nil
	}
}

func WithCertificate(certFilename, keyFilename string) Option {
	return func(cfg *Config) error {
		cfg.CertFilename = certFilename
		cfg.KeyFilename = keyFilename
		cfg.Insecure = false
		return nil
	}
}

func WithHTTPListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.HTTPListenPort = port
		return nil
	}
}

func WithRPCListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.RPCListenPort = port
		return nil
	}
}

func WithMetricsListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.MetricsListenPort = port
		return nil
	}
}

func WithLogger(l *zap.Logger) Option {
	return func(cfg *Config) error {
		cfg.logger = l
		return nil
	}
}

type sourcetype int

const (
	interrupt sourcetype = iota
	httpServer
	metricsServer
	rpcServer
)

type errorSource struct {
	source sourcetype
	err    error
}

func grpcEndpointLog(logger *zap.Logger, s string) grpc.UnaryServerInterceptor {
	return func(ctx xcontext.Context,
		req interface{},
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (interface{}, error) {
		logger.Info("grpcEndpointLog+", zap.String("", s))
		defer func() {
			logger.Info("grpcEndpointLog-", zap.String("", s))
			logger.Sync()
		}()

		return handler(ctx, req)
	}
}

func Run(opts ...Option) {
	cfg := &Config{
		Insecure:          true,
		HTTPListenPort:    8443,
		MetricsListenPort: 8080,
		RPCListenPort:     50050,
	}
	for _, o := range opts {
		o(cfg)
	}

	// make a channel to listen on events,
	// then launch the servers.

	errc := make(chan errorSource)

	// interrupt handler
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
		errc <- errorSource{
			source: interrupt,
			err:    fmt.Errorf("%s", <-c),
		}
	}()

	// gRPC server
	if cfg.RPCRegister != nil {
		go func() {
			rpcListenPort := ":" + strconv.Itoa(cfg.RPCListenPort)
			lis, err := net.Listen("tcp", rpcListenPort)
			if err != nil {
				errc <- errorSource{
					err:    err,
					source: rpcServer,
				}
				return
			}

			if cfg.Insecure {
				cfg.rpcServer = grpc.NewServer(
					grpc_middleware.WithUnaryServerChain(
						grpc_prometheus.UnaryServerInterceptor,
						grpcEndpointLog(cfg.logger, "certMgr")))
			} else {
				tlsCreds, err := credentials.NewServerTLSFromFile(cfg.CertFilename, cfg.KeyFilename)
				if err != nil {
					cfg.logger.Fatal("Failed to generate grpc TLS credentials", zap.Error(err))
				}
				cfg.rpcServer = grpc.NewServer(
					grpc.Creds(tlsCreds),
					grpc.RPCCompressor(grpc.NewGZIPCompressor()),
					grpc.RPCDecompressor(grpc.NewGZIPDecompressor()),
					grpc_middleware.WithUnaryServerChain(
						grpc_prometheus.UnaryServerInterceptor,
						grpcEndpointLog(cfg.logger, "Echo RPC server")))
			}

			cfg.RPCRegister(cfg.rpcServer)

			if cfg.Insecure {
				log.Warnf("gRPC service listening insecurely on %s", rpcListenPort)
			} else {
				log.Infof("gRPC service listening on %s", rpcListenPort)
			}
			errc <- errorSource{
				err:    cfg.rpcServer.Serve(lis),
				source: rpcServer,
			}
		}()
	}

	// health & metrics via https
	go func() {
		rootMux := mux.NewRouter() //actuator.NewActuatorMux("")

		hc, err := healthz.NewConfig()
		healthzHandler, err := healthz.Handler(hc)
		if err != nil {
			cfg.logger.Panic("Constructing healthz.Handler", zap.Error(err))
		}

		// set up handlers for THIS instance
		// (these are not expected to be proxied)
		rootMux.Handle("/debug/vars", expvar.Handler())
		rootMux.Handle("/healthz", healthzHandler)
		rootMux.Handle("/metrics", prometheus.Handler())

		canonical := handlers.CanonicalHost("http://fubar.local.dstcorp.io:7070", http.StatusPermanentRedirect)
		var tracer func(http.Handler) http.Handler
		tracer = gsh.TracerFromHTTPRequest(gsh.NewTracer("commandName"), "proxy")

		//rootMux.PathPrefix("/").Handler(p)

		chain := alice.New(tracer,
			gsh.HTTPMetricsCollector,
			gsh.HTTPLogrusLogger,
			canonical,
			handlers.CompressHandler)

		//errc <- http.ListenAndServe(p.address, chain.Then(rootMux))
		httpListenAddress := ":" + strconv.Itoa(cfg.HTTPListenPort)
		cfg.httpServer = &http.Server{
			Addr:              httpListenAddress,
			Handler:           chain.Then(rootMux),
			ReadTimeout:       time.Duration(5) * time.Second,
			ReadHeaderTimeout: time.Duration(2) * time.Second,
		}

		errc <- errorSource{
			err:    cfg.httpServer.ListenAndServeTLS(cfg.CertFilename, cfg.KeyFilename),
			source: httpServer,
		}
	}()

	// start the hystrix stream provider
	go func() {
		hystrixStreamHandler := afex.NewStreamHandler()
		hystrixStreamHandler.Start()
		listenPort := ":" + strconv.Itoa(cfg.MetricsListenPort)
		cfg.metricsServer = &http.Server{
			Addr:    listenPort,
			Handler: hystrixStreamHandler,
		}
		errc <- errorSource{
			err:    cfg.metricsServer.ListenAndServe(),
			source: metricsServer,
		}
	}()

	// wait for somthin'
	cfg.logger.Info("Echo Server",
		zap.Int("http port", cfg.HTTPListenPort),
		zap.Int("metrics port", cfg.MetricsListenPort),
		zap.Int("RPC port", cfg.RPCListenPort))
	rc := <-errc

	// somethin happened, now shut everything down gracefully
	cfg.logger.Info("exit", zap.Error(rc.err), zap.Int("source", int(rc.source)))
	waitDuration := time.Duration(5) * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), waitDuration)
	defer cancel()

	if rc.source != httpServer && cfg.httpServer != nil {
		err = cfg.httpServer.Shutdown(ctx)
	}
	if rc.source != rpcServer && cfg.rpcServer != nil {
		cfg.rpcServer.GracefulStop()
	}
	if rc.source != metricsServer && cfg.metricsServer != nil {
		cfg.metricsServer.Shutdown(ctx)
	}

	// wait for shutdown to complete or time to expire
	select {
	case <-time.After(time.Duration(4) * time.Second):
		cfg.logger.Info("server shutdown complete")

	case <-ctx.Done():
		cfg.logger.Warn("shutdown time expired -- performing hard shutdown", zap.Error(ctx.Err()))
	}
}
<commit_msg>better shutdown handing<commit_after>package server

import (
	"expvar"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strconv"
	"syscall"
	"time"

	"context"

	afex "github.com/afex/hystrix-go/hystrix"
	"github.com/gorilla/handlers"
	"github.com/gorilla/mux"
	"github.com/grpc-ecosystem/go-grpc-prometheus"
	"github.com/justinas/alice"
	gsh "github.com/mchudgins/go-service-helper/handlers"
	"github.com/mchudgins/playground/pkg/healthz"
	"github.com/mwitkow/go-grpc-middleware"
	"github.com/prometheus/client_golang/prometheus"
	"github.com/prometheus/common/log"
	"go.uber.org/zap"
	xcontext "golang.org/x/net/context"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
)

type Config struct {
	Insecure          bool
	CertFilename      string
	KeyFilename       string
	HTTPListenPort    int
	RPCListenPort     int
	MetricsListenPort int
	RPCRegister       RPCRegistration
	logger            *zap.Logger
	rpcServer         *grpc.Server
	httpServer        *http.Server
	metricsServer     *http.Server
}

type Option func(*Config) error

type RPCRegistration func(*grpc.Server) error

func WithRPCServer(fn RPCRegistration) Option {
	return func(cfg *Config) error {
		/*
			echoServer, err := NewServer(p.logger)
			if err != nil {
				cfg.logger.Panic("while creating new EchoServer", zap.Error(err))
			}
			rpc.RegisterEchoServiceServer(s, echoServer)
		*/
		cfg.RPCRegister = fn

		return nil
	}
}

func WithCertificate(certFilename, keyFilename string) Option {
	return func(cfg *Config) error {
		cfg.CertFilename = certFilename
		cfg.KeyFilename = keyFilename
		cfg.Insecure = false
		return nil
	}
}

func WithHTTPListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.HTTPListenPort = port
		return nil
	}
}

func WithRPCListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.RPCListenPort = port
		return nil
	}
}

func WithMetricsListenPort(port int) Option {
	return func(cfg *Config) error {
		cfg.MetricsListenPort = port
		return nil
	}
}

func WithLogger(l *zap.Logger) Option {
	return func(cfg *Config) error {
		cfg.logger = l
		return nil
	}
}

type sourcetype int

const (
	interrupt sourcetype = iota
	httpServer
	metricsServer
	rpcServer
)

type eventSource struct {
	source sourcetype
	err    error
}

func grpcEndpointLog(logger *zap.Logger, s string) grpc.UnaryServerInterceptor {
	return func(ctx xcontext.Context,
		req interface{},
		info *grpc.UnaryServerInfo,
		handler grpc.UnaryHandler) (interface{}, error) {
		logger.Info("grpcEndpointLog+", zap.String("", s))
		defer func() {
			logger.Info("grpcEndpointLog-", zap.String("", s))
			logger.Sync()
		}()

		return handler(ctx, req)
	}
}

func Run(opts ...Option) {
	cfg := &Config{
		Insecure:          true,
		HTTPListenPort:    8443,
		MetricsListenPort: 8080,
		RPCListenPort:     50050,
	}

	for _, o := range opts {
		o(cfg)
	}

	// make a channel to listen on events,
	// then launch the servers.

	errc := make(chan eventSource)

	// interrupt handler
	go func() {
		c := make(chan os.Signal, 1)
		signal.Notify(c, syscall.SIGINT, syscall.SIGTERM)
		errc <- eventSource{
			source: interrupt,
			err:    fmt.Errorf("%s", <-c),
		}
	}()

	// gRPC server
	if cfg.RPCRegister != nil {
		go func() {
			rpcListenPort := ":" + strconv.Itoa(cfg.RPCListenPort)
			lis, err := net.Listen("tcp", rpcListenPort)
			if err != nil {
				errc <- eventSource{
					err:    err,
					source: rpcServer,
				}
				return
			}

			if cfg.Insecure {
				cfg.rpcServer = grpc.NewServer(
					grpc_middleware.WithUnaryServerChain(
						grpc_prometheus.UnaryServerInterceptor,
						grpcEndpointLog(cfg.logger, "certMgr")))
			} else {
				tlsCreds, err := credentials.NewServerTLSFromFile(cfg.CertFilename, cfg.KeyFilename)
				if err != nil {
					cfg.logger.Fatal("Failed to generate grpc TLS credentials", zap.Error(err))
				}
				cfg.rpcServer = grpc.NewServer(
					grpc.Creds(tlsCreds),
					grpc.RPCCompressor(grpc.NewGZIPCompressor()),
					grpc.RPCDecompressor(grpc.NewGZIPDecompressor()),
					grpc_middleware.WithUnaryServerChain(
						grpc_prometheus.UnaryServerInterceptor,
						grpcEndpointLog(cfg.logger, "Echo RPC server")))
			}

			cfg.RPCRegister(cfg.rpcServer)

			if cfg.Insecure {
				log.Warnf("gRPC service listening insecurely on %s", rpcListenPort)
			} else {
				log.Infof("gRPC service listening on %s", rpcListenPort)
			}
			errc <- eventSource{
				err:    cfg.rpcServer.Serve(lis),
				source: rpcServer,
			}
		}()
	}

	// health & metrics via https
	go func() {
		rootMux := mux.NewRouter() //actuator.NewActuatorMux("")

		hc, err := healthz.NewConfig()
		healthzHandler, err := healthz.Handler(hc)
		if err != nil {
			cfg.logger.Panic("Constructing healthz.Handler", zap.Error(err))
		}

		// set up handlers for THIS instance
		// (these are not expected to be proxied)
		rootMux.Handle("/debug/vars", expvar.Handler())
		rootMux.Handle("/healthz", healthzHandler)
		rootMux.Handle("/metrics", prometheus.Handler())

		canonical := handlers.CanonicalHost("http://fubar.local.dstcorp.io:7070", http.StatusPermanentRedirect)
		var tracer func(http.Handler) http.Handler
		tracer = gsh.TracerFromHTTPRequest(gsh.NewTracer("commandName"), "proxy")

		//rootMux.PathPrefix("/").Handler(p)

		chain := alice.New(tracer,
			gsh.HTTPMetricsCollector,
			gsh.HTTPLogrusLogger,
			canonical,
			handlers.CompressHandler)

		//errc <- http.ListenAndServe(p.address, chain.Then(rootMux))
		httpListenAddress := ":" + strconv.Itoa(cfg.HTTPListenPort)
		cfg.httpServer = &http.Server{
			Addr:              httpListenAddress,
			Handler:           chain.Then(rootMux),
			ReadTimeout:       time.Duration(5) * time.Second,
			ReadHeaderTimeout: time.Duration(2) * time.Second,
		}

		errc <- eventSource{
			err:    cfg.httpServer.ListenAndServeTLS(cfg.CertFilename, cfg.KeyFilename),
			source: httpServer,
		}
	}()

	// start the hystrix stream provider
	go func() {
		hystrixStreamHandler := afex.NewStreamHandler()
		hystrixStreamHandler.Start()
		listenPort := ":" + strconv.Itoa(cfg.MetricsListenPort)
		cfg.metricsServer = &http.Server{
			Addr:    listenPort,
			Handler: hystrixStreamHandler,
		}
		errc <- eventSource{
			err:    cfg.metricsServer.ListenAndServe(),
			source: metricsServer,
		}
	}()

	// wait for somthin'
	cfg.logger.Info("Echo Server",
		zap.Int("http port", cfg.HTTPListenPort),
		zap.Int("metrics port", cfg.MetricsListenPort),
		zap.Int("RPC port", cfg.RPCListenPort))
	rc := <-errc

	// somethin happened, now shut everything down gracefully
	cfg.logger.Info("exit", zap.Error(rc.err), zap.Int("source", int(rc.source)))
	waitDuration := time.Duration(5) * time.Second
	ctx, cancel := context.WithTimeout(context.Background(), waitDuration)
	defer cancel()

	waitEvents := 0
	evtc := make(chan eventSource)

	if rc.source != httpServer && cfg.httpServer != nil {
		waitEvents++
		go func() {
			evtc <- eventSource{
				err:    cfg.httpServer.Shutdown(ctx),
				source: httpServer,
			}
		}()
	}
	if rc.source != rpcServer && cfg.rpcServer != nil {
		waitEvents++
		go func() {
			cfg.rpcServer.GracefulStop()
			evtc <- eventSource{source: rpcServer}
		}()
	}
	if rc.source != metricsServer && cfg.metricsServer != nil {
		waitEvents++
		go func() {
			evtc <- eventSource{
				err:    cfg.metricsServer.Shutdown(ctx),
				source: metricsServer,
			}
		}()
	}

	// wait for shutdown to complete or time to expire
	for {
		select {
		case <-time.After(time.Duration(4) * time.Second):
			cfg.logger.Info("server shutdown complete")
			os.Exit(1)

		case <-ctx.Done():
			cfg.logger.Warn("shutdown time expired -- performing hard shutdown", zap.Error(ctx.Err()))
			os.Exit(2)

		case evt := <-evtc:
			waitEvents--
			cfg.logger.Info("waitEvent recv'ed", zap.Error(evt.err), zap.Int("eventSource", int(evt.source)))
			if waitEvents == 0 {
				os.Exit(0)
			}
		}
	}
}
<|endoftext|>


<commit_before>package server

import (
	"fmt"
	"github.com/godfried/cabanga/db"
	"github.com/godfried/cabanga/submission"
	"github.com/godfried/cabanga/user"
	"github.com/godfried/cabanga/util"
	"labix.org/v2/mgo/bson"
	"net"
)

const (
	OK     = "ok"
	SEND   = "send"
	LOGIN  = "begin"
	LOGOUT = "end"
	REQ    = "req"
)

//Run listens for new connections and creates a new goroutine for each connection.
func RunFileReceiver(port string, subChan chan *submission.Submission, fileChan chan *submission.File) {
	netListen, err := net.Listen("tcp", ":"+port)
	if err != nil {
		util.Log(fmt.Errorf("Encountered error %q when listening on port %q", err, port))
		return
	}
	defer netListen.Close()
	for {
		conn, err := netListen.Accept()
		if err != nil {
			util.Log(fmt.Errorf("Encountered error %q when accepting connection", err))
		} else {
			go ConnHandler(conn, subChan, fileChan)
		}
	}
}

func RunTestReceiver(port string){
	netListen, err := net.Listen("tcp", ":"+port)
	if err != nil {
		util.Log(fmt.Errorf("Encountered error %q when listening on port %q", err, port))
		return
	}
	defer netListen.Close()
	for {
		conn, err := netListen.Accept()
		if err != nil {
			util.Log(fmt.Errorf("Encountered error %q when accepting connection", err))
		} else {
			fmt.Println("received connection")
			go ReceiveTests(conn)
		}
	}
}

//ConnHandler manages an incoming connection request.
//It authenticates the request and processes files sent on the connection.
func ReceiveTests(conn net.Conn) {
	testInfo, err := util.ReadJSON(conn)
	fmt.Println("received info", testInfo)
	if err != nil {
		EndSession(conn, err)
		return
	}
	project, err := util.GetString(testInfo, submission.PROJECT)
	if err != nil {
		EndSession(conn, err)
		return
	}
	lang, err := util.GetString(testInfo, submission.LANG)
	if err != nil {
		EndSession(conn, err)
		return
	}
	names, err := util.GetStrings(testInfo, submission.NAMES)
	if err != nil {
		EndSession(conn, err)
		return
	}
	conn.Write([]byte(OK))
	testFiles, err := util.ReadData(conn)
	if err != nil {
		EndSession(conn, err)
		return
	}
	fmt.Println("received tests", testFiles)
	conn.Write([]byte(OK))
	dataFiles, err := util.ReadData(conn)
	if err != nil {
		EndSession(conn, err)
		return
	}
	fmt.Println("received tests", dataFiles)
	conn.Write([]byte(OK))
	test := submission.NewTest(project, lang, names, testFiles, dataFiles)
	err = db.AddTest(test)
	EndSession(conn, err)
}

//ConnHandler manages an incoming connection request.
//It authenticates the request and processes files sent on the connection.
func ConnHandler(conn net.Conn, subChan chan *submission.Submission, fileChan chan *submission.File) {
	jobj, err := util.ReadJSON(conn)
	if err != nil {
		EndSession(conn, err)
		return
	}
	sub, err := Login(jobj, conn)
	if err != nil {
		EndSession(conn, err)
		return
	}
	subChan <- sub
	util.Log("Created submission: ", sub)
	receiving := true
	for receiving && err == nil {
		jobj, err = util.ReadJSON(conn)
		if err != nil {
			EndSession(conn, err)
			return
		}
		req, err := util.GetString(jobj, REQ)
		if err != nil {
			EndSession(conn, err)
			return
		}
		if req == SEND {
			delete(jobj, REQ)
			err = ProcessFile(sub.Id, jobj, conn, fileChan)
		} else if req == LOGOUT {
			receiving = false
		} else {
			err = fmt.Errorf("Unknown request %q", req)
		}
	}
	EndSession(conn, err)
}

//processFile reads file data from connection and stores it in the db.
//The file data is then sent on fileChan for further processing.
func ProcessFile(subId bson.ObjectId, finfo map[string]interface{}, conn net.Conn, fileChan chan *submission.File) error {
	conn.Write([]byte(OK))
	buffer, err := util.ReadData(conn)
	if err != nil {
		return err
	}
	conn.Write([]byte(OK))
	f := submission.NewFile(subId, finfo, buffer)
	err = db.AddFile(f)
	if err != nil {
		return err
	}
	fileChan <- f
	return nil
}

//Login creates a new submission if the login request is valid.
func Login(jobj map[string]interface{}, conn net.Conn) (*submission.Submission, error) {
	sub, err := CreateSubmission(jobj)
	if err != nil {
		return nil, err
	}
	err = db.AddSubmission(sub)
	if err != nil {
		return nil, err
	}
	conn.Write([]byte(OK))
	return sub, nil
}

//EndSession ends a session and reports any errors to the client.
func EndSession(conn net.Conn, err error) {
	var msg string
	if err != nil {
		msg = "ERROR: " + err.Error()
		util.Log(err)
	} else {
		msg = OK
	}
	conn.Write([]byte(msg))
	conn.Close()
}

//CreateSubmission validates a login request.
//It reads submission values from a json object and checks privilege level and password.
func CreateSubmission(jobj map[string]interface{}) (*submission.Submission, error) {
	username, err := util.GetString(jobj, submission.USER)
	if err != nil {
		return nil, err
	}
	pword, err := util.GetString(jobj, user.PWORD)
	if err != nil {
		return nil, err
	}
	project, err := util.GetString(jobj, submission.PROJECT)
	if err != nil {
		return nil, err
	}
	mode, err := util.GetString(jobj, submission.MODE)
	if err != nil {
		return nil, err
	}
	lang, err := util.GetString(jobj, submission.LANG)
	if err != nil {
		return nil, err
	}
	u, err := db.GetUserById(username)
	if err != nil {
		return nil, err
	}
	if !u.CheckSubmit(mode) {
		return nil, fmt.Errorf("User %q has insufficient permissions for %q", username, mode)
	}
	if !util.Validate(u.Password, u.Salt, pword) {
		return nil, fmt.Errorf("User %q attempted to login with an invalid username or password", username)
	}
	return submission.NewSubmission(project, username, mode, lang), nil
}
<commit_msg>added authentication for test submissions<commit_after>package server

import (
	"fmt"
	"github.com/godfried/cabanga/db"
	"github.com/godfried/cabanga/submission"
	"github.com/godfried/cabanga/user"
	"github.com/godfried/cabanga/util"
	"labix.org/v2/mgo/bson"
	"net"
)

const (
	OK     = "ok"
	SEND   = "send"
	LOGIN  = "begin"
	LOGOUT = "end"
	REQ    = "req"
)

//Run listens for new connections and creates a new goroutine for each connection.
func RunFileReceiver(port string, subChan chan *submission.Submission, fileChan chan *submission.File) {
	netListen, err := net.Listen("tcp", ":"+port)
	if err != nil {
		util.Log(fmt.Errorf("Encountered error %q when listening on port %q", err, port))
		return
	}
	defer netListen.Close()
	for {
		conn, err := netListen.Accept()
		if err != nil {
			util.Log(fmt.Errorf("Encountered error %q when accepting connection", err))
		} else {
			go ConnHandler(conn, subChan, fileChan)
		}
	}
}

func RunTestReceiver(port string){
	netListen, err := net.Listen("tcp", ":"+port)
	if err != nil {
		util.Log(fmt.Errorf("Encountered error %q when listening on port %q", err, port))
		return
	}
	defer netListen.Close()
	for {
		conn, err := netListen.Accept()
		if err != nil {
			util.Log(fmt.Errorf("Encountered error %q when accepting connection", err))
		} else {
			fmt.Println("received connection")
			go ReceiveTests(conn)
		}
	}
}

//ConnHandler manages an incoming connection request.
//It authenticates the request and processes files sent on the connection.
func ReceiveTests(conn net.Conn) {
	err := TestLogin(conn)
	if err != nil {
		EndSession(conn, err)
		return
	}
	test, err := ReadTest(conn)
	if err != nil {
		EndSession(conn, err)
		return
	}
	err = db.AddTest(test)
	EndSession(conn, err)
}

//Login creates a new submission if the login request is valid.
func TestLogin(conn net.Conn) error {
	loginInfo, err := util.ReadJSON(conn)
	if err != nil {
		return err
	}
	username, err := util.GetString(loginInfo, submission.USER)
	if err != nil {
		return err
	}
	pword, err := util.GetString(loginInfo, user.PWORD)
	if err != nil {
		return err
	}
	u, err := db.GetUserById(username)
	if err != nil {
		return err
	}
	if !u.CheckSubmit(submission.TEST_MODE) {
		return fmt.Errorf("User %q has insufficient permissions for %q", username, submission.TEST_MODE)
	}
	if !util.Validate(u.Password, u.Salt, pword) {
		return fmt.Errorf("User %q attempted to login with an invalid username or password", username)
	}
	conn.Write([]byte(OK))
	return nil
}

func ReadTest(conn net.Conn) (*submission.Test, error){
	testInfo, err := util.ReadJSON(conn)
	if err != nil {
		return nil, err
	}
	project, err := util.GetString(testInfo, submission.PROJECT)
	if err != nil {
		return nil, err
	}
	lang, err := util.GetString(testInfo, submission.LANG)
	if err != nil {
		return nil, err
	}
	names, err := util.GetStrings(testInfo, submission.NAMES)
	if err != nil {
		return nil, err
	}
	conn.Write([]byte(OK))
	testFiles, err := util.ReadData(conn)
	if err != nil {
		return nil, err
	}
	conn.Write([]byte(OK))
	dataFiles, err := util.ReadData(conn)
	if err != nil {
		return nil, err
	}
	conn.Write([]byte(OK))
	return submission.NewTest(project, lang, names, testFiles, dataFiles), nil
}

//ConnHandler manages an incoming connection request.
//It authenticates the request and processes files sent on the connection.
func ConnHandler(conn net.Conn, subChan chan *submission.Submission, fileChan chan *submission.File) {
	err := ReceiveFiles(conn, subChan, fileChan)
	EndSession(conn, err)
}

func ReceiveFiles(conn net.Conn, subChan chan *submission.Submission, fileChan chan *submission.File)error{
	sub, err := Login(conn)
	if err != nil {
		return err
	}
	subChan <- sub
	for {
		requestInfo, err := util.ReadJSON(conn)
		if err != nil {
			return err
		}
		req, err := util.GetString(requestInfo, REQ)
		if err != nil {
			return err
		}
		if req == SEND {
			err = ProcessFile(sub.Id, requestInfo, conn, fileChan)
			if err != nil {
				return err
			}
		} else if req == LOGOUT {
			return nil
		} else{
			return fmt.Errorf("Unknown request %q", req)
		}
	}
}

//processFile reads file data from connection and stores it in the db.
//The file data is then sent on fileChan for further processing.
func ProcessFile(subId bson.ObjectId, finfo map[string]interface{}, conn net.Conn, fileChan chan *submission.File) error {
	conn.Write([]byte(OK))
	buffer, err := util.ReadData(conn)
	if err != nil {
		return err
	}
	conn.Write([]byte(OK))
	f := submission.NewFile(subId, finfo, buffer)
	err = db.AddFile(f)
	if err != nil {
		return err
	}
	fileChan <- f
	return nil
}

//Login creates a new submission if the login request is valid.
func Login(conn net.Conn) (*submission.Submission, error) {
	fileInfo, err := util.ReadJSON(conn)
	if err != nil {
		return nil, err
	}
	sub, err := CreateSubmission(fileInfo)
	if err != nil {
		return nil, err
	}
	err = db.AddSubmission(sub)
	if err != nil {
		return nil, err
	}
	conn.Write([]byte(OK))
	return sub, nil
}

//EndSession ends a session and reports any errors to the client.
func EndSession(conn net.Conn, err error) {
	var msg string
	if err != nil {
		msg = "ERROR: " + err.Error()
		util.Log(err)
	} else {
		msg = OK
	}
	conn.Write([]byte(msg))
	conn.Close()
}

//CreateSubmission validates a login request.
//It reads submission values from a json object and checks privilege level and password.
func CreateSubmission(loginInfo map[string]interface{}) (*submission.Submission, error) {
	username, err := util.GetString(loginInfo, submission.USER)
	if err != nil {
		return nil, err
	}
	pword, err := util.GetString(loginInfo, user.PWORD)
	if err != nil {
		return nil, err
	}
	project, err := util.GetString(loginInfo, submission.PROJECT)
	if err != nil {
		return nil, err
	}
	mode, err := util.GetString(loginInfo, submission.MODE)
	if err != nil {
		return nil, err
	}
	lang, err := util.GetString(loginInfo, submission.LANG)
	if err != nil {
		return nil, err
	}
	u, err := db.GetUserById(username)
	if err != nil {
		return nil, err
	}
	if !u.CheckSubmit(mode) {
		return nil, fmt.Errorf("User %q has insufficient permissions for %q", username, mode)
	}
	if !util.Validate(u.Password, u.Salt, pword) {
		return nil, fmt.Errorf("User %q attempted to login with an invalid username or password", username)
	}
	return submission.NewSubmission(project, username, mode, lang), nil
}
<|endoftext|>


<commit_before>/*
Copyright 2018 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Command server launches a stand-alone inverting proxy.
//
// Example usage:
//   go build -o ~/bin/inverting-proxy ./server/server.go
//   ~/bin/inverting-proxy --port 8081
package main

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/google/inverting-proxy/agent/utils"
)

var (
	port = flag.Int("port", 0, "Port on which to listen")
)

// pendingRequest represents a frontend request
type pendingRequest struct {
	startTime time.Time
	req       *http.Request
	respChan  chan *http.Response
}

func newPendingRequest(r *http.Request) *pendingRequest {
	return &pendingRequest{
		startTime: time.Now(),
		req:       r,
		respChan:  make(chan *http.Response),
	}
}

type proxy struct {
	requestIDs    chan string
	randGenerator *rand.Rand

	// protects the map below
	sync.Mutex
	requests map[string]*pendingRequest
}

func newProxy() *proxy {
	return &proxy{
		requestIDs:    make(chan string),
		randGenerator: rand.New(rand.NewSource(time.Now().UnixNano())),
		requests:      make(map[string]*pendingRequest),
	}
}

func (p *proxy) handleAgentPostResponse(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	resp, err := http.ReadResponse(bufio.NewReader(r.Body), pending.req)
	if err != nil {
		log.Printf("Could not parse response to request %q: %v", requestID, err)
		http.Error(w, "Failure parsing request body", http.StatusBadRequest)
		return
	}
	// We want to track whether or not the body has finished being read so that we can
	// make sure that this method does not return until after that. However, we do not
	// want to block the sending of the response to the client while it is being read.
	//
	// To accommodate both goals, we replace the response body with a pipereader, start
	// forwarding the response immediately, and then copy the original body to the
	// corresponding pipewriter.
	respBody := resp.Body
	defer respBody.Close()

	pr, pw := io.Pipe()
	defer pw.Close()

	resp.Body = pr
	select {
	case <-r.Context().Done():
		return
	case pending.respChan <- resp:
	}
	if _, err := io.Copy(pw, respBody); err != nil {
		log.Printf("Could not read response to request %q: %v", requestID, err)
		http.Error(w, "Failure reading request body", http.StatusInternalServerError)
	}
}

func (p *proxy) handleAgentGetRequest(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	log.Printf("Returning pending request: %q", requestID)
	w.Header().Set(utils.HeaderRequestStartTime, pending.startTime.Format(time.RFC3339Nano))
	w.WriteHeader(http.StatusOK)
	pending.req.Write(w)
}

// waitForRequestIDs blocks until at least one request ID is available, and then returns
// a slice of all of the IDs available at that time.
//
// Note that any IDs returned by this method will never be returned again.
func (p *proxy) waitForRequestIDs(ctx context.Context) []string {
	var requestIDs []string
	select {
	case <-ctx.Done():
		return nil
	case <-time.After(30 * time.Second):
		return nil
	case id := <-p.requestIDs:
		requestIDs = append(requestIDs, id)
	}
	for {
		select {
		case id := <-p.requestIDs:
			requestIDs = append(requestIDs, id)
		default:
			return requestIDs
		}
	}
}

func (p *proxy) handleAgentListRequests(w http.ResponseWriter, r *http.Request) {
	requestIDs := p.waitForRequestIDs(r.Context())
	respJSON, err := json.Marshal(requestIDs)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failure serializing the request IDs: %v", err), http.StatusInternalServerError)
		return
	}
	log.Printf("Reporting pending requests: %s", respJSON)
	w.WriteHeader(http.StatusOK)
	w.Write(respJSON)
}

func (p *proxy) handleAgentRequest(w http.ResponseWriter, r *http.Request, backendID string) {
	requestID := r.Header.Get(utils.HeaderRequestID)
	if requestID == "" {
		log.Printf("Received new backend list request from %q", backendID)
		p.handleAgentListRequests(w, r)
		return
	}
	if r.Method == http.MethodPost {
		log.Printf("Received new backend post request from %q", backendID)
		p.handleAgentPostResponse(w, r, requestID)
		return
	}
	log.Printf("Received new backend get request from %q", backendID)
	p.handleAgentGetRequest(w, r, requestID)
}

func (p *proxy) newID() string {
	sum := sha256.Sum256([]byte(fmt.Sprintf("%d", p.randGenerator.Int63())))
	return fmt.Sprintf("%x", sum)
}

// isHopByHopHeader determines whether or not the given header name represents
// a header that is specific to a single network hop and thus should not be
// retransmitted by a proxy.
//
// See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers#hbh
func isHopByHopHeader(name string) bool {
	switch n := strings.ToLower(name); n {
	case "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade":
		return true
	default:
		return false
	}
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if backendID := r.Header.Get(utils.HeaderBackendID); backendID != "" {
		p.handleAgentRequest(w, r, backendID)
		return
	}
	id := p.newID()
	log.Printf("Received new frontend request %q", id)
	pending := newPendingRequest(r)
	p.Lock()
	p.requests[id] = pending
	p.Unlock()

	// Enqueue the request
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting to enqueue the request ID for %q", id)
		return
	case p.requestIDs <- id:
	}
	log.Printf("Request %q enqueued after %s", id, time.Since(pending.startTime))

	// Pull out and copy the response
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting for the response to %q", id)
		return
	case resp := <-pending.respChan:
		defer resp.Body.Close()
		// Copy all of the non-hop-by-hop headers to the proxied response
		for key, vals := range resp.Header {
			if !isHopByHopHeader(key) {
				w.Header()[key] = vals
			}
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
		return
	}
	log.Printf("Response for %q received after %s", id, time.Since(pending.startTime))
}

func main() {
	flag.Parse()
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
	if err != nil {
		log.Fatalf("Failed to create the TCP listener for port %d: %v", *port, err)
	}
	log.Printf("Listening on %s", listener.Addr())
	log.Fatal(http.Serve(listener, newProxy()))
}
<commit_msg>Make proxied request headers compliant with RFC 2616<commit_after>/*
Copyright 2018 Google Inc. All rights reserved.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

// Command server launches a stand-alone inverting proxy.
//
// Example usage:
//   go build -o ~/bin/inverting-proxy ./server/server.go
//   ~/bin/inverting-proxy --port 8081
package main

import (
	"bufio"
	"context"
	"crypto/sha256"
	"encoding/json"
	"flag"
	"fmt"
	"io"
	"log"
	"math/rand"
	"net"
	"net/http"
	"strings"
	"sync"
	"time"

	"github.com/google/inverting-proxy/agent/utils"
)

var (
	port = flag.Int("port", 0, "Port on which to listen")
)

// pendingRequest represents a frontend request
type pendingRequest struct {
	startTime time.Time
	req       *http.Request
	respChan  chan *http.Response
}

func newPendingRequest(r *http.Request) *pendingRequest {
	return &pendingRequest{
		startTime: time.Now(),
		req:       r,
		respChan:  make(chan *http.Response),
	}
}

type proxy struct {
	requestIDs    chan string
	randGenerator *rand.Rand

	// protects the map below
	sync.Mutex
	requests map[string]*pendingRequest
}

func newProxy() *proxy {
	return &proxy{
		requestIDs:    make(chan string),
		randGenerator: rand.New(rand.NewSource(time.Now().UnixNano())),
		requests:      make(map[string]*pendingRequest),
	}
}

func (p *proxy) handleAgentPostResponse(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	resp, err := http.ReadResponse(bufio.NewReader(r.Body), pending.req)
	if err != nil {
		log.Printf("Could not parse response to request %q: %v", requestID, err)
		http.Error(w, "Failure parsing request body", http.StatusBadRequest)
		return
	}
	// We want to track whether or not the body has finished being read so that we can
	// make sure that this method does not return until after that. However, we do not
	// want to block the sending of the response to the client while it is being read.
	//
	// To accommodate both goals, we replace the response body with a pipereader, start
	// forwarding the response immediately, and then copy the original body to the
	// corresponding pipewriter.
	respBody := resp.Body
	defer respBody.Close()

	pr, pw := io.Pipe()
	defer pw.Close()

	resp.Body = pr
	select {
	case <-r.Context().Done():
		return
	case pending.respChan <- resp:
	}
	if _, err := io.Copy(pw, respBody); err != nil {
		log.Printf("Could not read response to request %q: %v", requestID, err)
		http.Error(w, "Failure reading request body", http.StatusInternalServerError)
	}
}

func (p *proxy) handleAgentGetRequest(w http.ResponseWriter, r *http.Request, requestID string) {
	p.Lock()
	pending, ok := p.requests[requestID]
	p.Unlock()
	if !ok {
		log.Printf("Could not find pending request: %q", requestID)
		http.NotFound(w, r)
		return
	}
	log.Printf("Returning pending request: %q", requestID)
	w.Header().Set(utils.HeaderRequestStartTime, pending.startTime.Format(time.RFC3339Nano))
	w.WriteHeader(http.StatusOK)
	pending.req.Write(w)
}

// waitForRequestIDs blocks until at least one request ID is available, and then returns
// a slice of all of the IDs available at that time.
//
// Note that any IDs returned by this method will never be returned again.
func (p *proxy) waitForRequestIDs(ctx context.Context) []string {
	var requestIDs []string
	select {
	case <-ctx.Done():
		return nil
	case <-time.After(30 * time.Second):
		return nil
	case id := <-p.requestIDs:
		requestIDs = append(requestIDs, id)
	}
	for {
		select {
		case id := <-p.requestIDs:
			requestIDs = append(requestIDs, id)
		default:
			return requestIDs
		}
	}
}

func (p *proxy) handleAgentListRequests(w http.ResponseWriter, r *http.Request) {
	requestIDs := p.waitForRequestIDs(r.Context())
	respJSON, err := json.Marshal(requestIDs)
	if err != nil {
		http.Error(w, fmt.Sprintf("Failure serializing the request IDs: %v", err), http.StatusInternalServerError)
		return
	}
	log.Printf("Reporting pending requests: %s", respJSON)
	w.WriteHeader(http.StatusOK)
	w.Write(respJSON)
}

func (p *proxy) handleAgentRequest(w http.ResponseWriter, r *http.Request, backendID string) {
	requestID := r.Header.Get(utils.HeaderRequestID)
	if requestID == "" {
		log.Printf("Received new backend list request from %q", backendID)
		p.handleAgentListRequests(w, r)
		return
	}
	if r.Method == http.MethodPost {
		log.Printf("Received new backend post request from %q", backendID)
		p.handleAgentPostResponse(w, r, requestID)
		return
	}
	log.Printf("Received new backend get request from %q", backendID)
	p.handleAgentGetRequest(w, r, requestID)
}

func (p *proxy) newID() string {
	sum := sha256.Sum256([]byte(fmt.Sprintf("%d", p.randGenerator.Int63())))
	return fmt.Sprintf("%x", sum)
}

// isHopByHopHeader determines whether or not the given header name represents
// a header that is specific to a single network hop and thus should not be
// retransmitted by a proxy.
//
// See: https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers#hbh
func isHopByHopHeader(name string) bool {
	switch n := strings.ToLower(name); n {
	case "connection", "keep-alive", "proxy-authenticate", "proxy-authorization", "te", "trailer", "transfer-encoding", "upgrade":
		return true
	default:
		return false
	}
}

func (p *proxy) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	if backendID := r.Header.Get(utils.HeaderBackendID); backendID != "" {
		p.handleAgentRequest(w, r, backendID)
		return
	}
	id := p.newID()
	log.Printf("Received new frontend request %q", id)
	// Filter out hop-by-hop headers from the request
	for name := range r.Header {
		if isHopByHopHeader(name) {
			r.Header.Del(name)
		}
	}
	pending := newPendingRequest(r)
	p.Lock()
	p.requests[id] = pending
	p.Unlock()

	// Enqueue the request
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting to enqueue the request ID for %q", id)
		return
	case p.requestIDs <- id:
	}
	log.Printf("Request %q enqueued after %s", id, time.Since(pending.startTime))

	// Pull out and copy the response
	select {
	case <-r.Context().Done():
		// The client request was cancelled
		log.Printf("Timeout waiting for the response to %q", id)
		return
	case resp := <-pending.respChan:
		defer resp.Body.Close()
		// Copy all of the non-hop-by-hop headers to the proxied response
		for name, vals := range resp.Header {
			if !isHopByHopHeader(name) {
				w.Header()[name] = vals
			}
		}
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
		return
	}
	log.Printf("Response for %q received after %s", id, time.Since(pending.startTime))
}

func main() {
	flag.Parse()
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", *port))
	if err != nil {
		log.Fatalf("Failed to create the TCP listener for port %d: %v", *port, err)
	}
	log.Printf("Listening on %s", listener.Addr())
	log.Fatal(http.Serve(listener, newProxy()))
}
<|endoftext|>


<commit_before>package notifications

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"strconv"
	"strings"
	"time"

	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/lager"
)

//go:generate counterfeiter . Notifier

type Notifier interface {
	SendNotification(lager.Logger, Notification) error
	SendBatchNotification(lager.Logger, []Notification) error
}

type Notification struct {
	Owner      string
	Repository string
	Private    bool

	SHA string

	Path       string
	LineNumber int
}

func (n Notification) FullName() string {
	return fmt.Sprintf("%s/%s", n.Owner, n.Repository)
}

func (n Notification) ShortSHA() string {
	return n.SHA[:7]
}

type slackNotifier struct {
	webhookURL string
	client     *http.Client
	clock      clock.Clock
}

type slackMessage struct {
	Attachments []slackAttachment `json:"attachments"`
}

type slackAttachment struct {
	Fallback string `json:"fallback"`
	Color    string `json:"color"`
	Title    string `json:"title"`
	Text     string `json:"text"`
}

func NewSlackNotifier(webhookURL string, clock clock.Clock) Notifier {
	if webhookURL == "" {
		return &nullSlackNotifier{}
	}

	return &slackNotifier{
		webhookURL: webhookURL,
		clock:      clock,
		client: &http.Client{
			Timeout: 3 * time.Second,
			Transport: &http.Transport{
				DisableKeepAlives: true,
			},
		},
	}
}

const maxRetries = 3

func (n *slackNotifier) SendNotification(logger lager.Logger, notification Notification) error {
	logger = logger.Session("send-notification")
	logger.Debug("starting")

	message := n.formatSlackMessage(notification)

	body, err := json.Marshal(message)
	if err != nil {
		logger.Error("unmarshal-faiiled", err)
		return err
	}

	return n.send(logger, body)
}

func (n *slackNotifier) SendBatchNotification(logger lager.Logger, batch []Notification) error {
	logger = logger.Session("send-batch-notification", lager.Data{"batch-size": len(batch)})

	logger.Debug("starting")

	if len(batch) == 0 {
		logger.Debug("done")
		return nil
	}

	messages := n.formatBatchSlackMessages(batch)

	for _, message := range messages {
		body, err := json.Marshal(message)
		if err != nil {
			logger.Error("unmarshal-faiiled", err)
			return err
		}

		err = n.send(logger, body)
		if err != nil {
			return err
		}
	}

	return nil
}

func (n *slackNotifier) send(logger lager.Logger, body []byte) error {
	for numReq := 0; numReq < maxRetries; numReq++ {
		req, err := http.NewRequest("POST", n.webhookURL, bytes.NewBuffer(body))
		if err != nil {
			logger.Error("request-failed", err)
			return err
		}

		req.Header.Set("Content-type", "application/json")

		resp, err := n.client.Do(req)
		if err != nil {
			logger.Error("response-error", err)
			return err
		}

		switch resp.StatusCode {
		case http.StatusOK:
			logger.Debug("done")
			return nil
		case http.StatusTooManyRequests:
			lastLoop := (numReq == maxRetries-1)
			if lastLoop {
				break
			}

			afterStr := resp.Header.Get("Retry-After")
			logger.Info("told-to-wait", lager.Data{"after": afterStr})
			after, err := strconv.Atoi(afterStr)
			if err != nil {
				logger.Error("failed", err)
				return err
			}

			wait := after + 1 // +1 for luck

			n.clock.Sleep(time.Duration(wait) * time.Second)
			continue
		default:
			err = fmt.Errorf("bad response (!200): %d", resp.StatusCode)
			logger.Error("bad-response", err)
			return err
		}
	}

	err := errors.New("retried too many times")
	logger.Error("failed", err)

	return err
}

func (n *slackNotifier) formatSlackMessage(not Notification) slackMessage {
	link := fmt.Sprintf("https://github.com/%s/%s/blob/%s/%s#L%d", not.Owner, not.Repository, not.SHA, not.Path, not.LineNumber)

	color := "danger"
	if not.Private {
		color = "warning"
	}

	return slackMessage{
		Attachments: []slackAttachment{
			{
				Fallback: link,
				Color:    color,
				Title:    fmt.Sprintf("Credential detected in %s!", not.FullName()),
				Text:     fmt.Sprintf("<%s|%s:%d>", link, not.Path, not.LineNumber),
			},
		},
	}
}

type slackLink struct {
	Text string
	Href string
}

func (l slackLink) String() string {
	return fmt.Sprintf("<%s|%s>", l.Href, l.Text)
}

type slackBatchRepo struct {
	Owner      string
	Repository string
	SHA        string
	Private    bool
}

func (r slackBatchRepo) FullName() string {
	return fmt.Sprintf("%s/%s", r.Owner, r.Repository)
}

func (r slackBatchRepo) ShortSHA() string {
	return r.SHA[:7]
}

func (n *slackNotifier) formatBatchSlackMessages(batch []Notification) []slackMessage {
	messages := []slackMessage{}

	messageMap := make(map[slackBatchRepo]map[string][]Notification)

	for _, not := range batch {
		repo := slackBatchRepo{
			Owner:      not.Owner,
			Repository: not.Repository,
			SHA:        not.SHA,
			Private:    not.Private,
		}

		_, found := messageMap[repo]
		if !found {
			messageMap[repo] = make(map[string][]Notification)
		}

		messageMap[repo][not.Path] = append(messageMap[repo][not.Path], not)
	}

	for repo, files := range messageMap {
		commitLink := fmt.Sprintf("https://github.com/%s/%s/commit/%s", repo.Owner, repo.Repository, repo.SHA)
		title := fmt.Sprintf("Possible credentials found in %s!", slackLink{
			Text: fmt.Sprintf("%s / %s", repo.FullName(), repo.ShortSHA()),
			Href: commitLink,
		})
		fallback := fmt.Sprintf("Possible credentials found in %s!", commitLink)

		color := "danger"
		if repo.Private {
			color = "warning"
		}

		fileLines := []string{}

		for path, nots := range files {
			fileLink := fmt.Sprintf("https://github.com/%s/%s/blob/%s/%s", repo.Owner, repo.Repository, repo.SHA, path)

			lineLinks := []string{}

			for _, not := range nots {
				lineLink := fmt.Sprintf("%s#L%d", fileLink, not.LineNumber)

				lineLinks = append(lineLinks, slackLink{
					Text: strconv.Itoa(not.LineNumber),
					Href: lineLink,
				}.String())
			}

			plurality := "line"
			if len(lineLinks) > 1 {
				plurality = "lines"
			}

			text := fmt.Sprintf("• %s on %s %s", slackLink{
				Text: path,
				Href: fileLink,
			}, plurality, humanizeList(lineLinks))

			fileLines = append(fileLines, text)
		}

		messages = append(messages, slackMessage{
			Attachments: []slackAttachment{
				{
					Title:    title,
					Text:     strings.Join(fileLines, "\n"),
					Color:    color,
					Fallback: fallback,
				},
			},
		})
	}

	return messages
}

func humanizeList(list []string) string {
	joinedLines := &bytes.Buffer{}

	if len(list) <= 1 {
		joinedLines.WriteString(list[0])
	} else if len(list) == 2 {
		joinedLines.WriteString(list[0])
		joinedLines.WriteString(" and ")
		joinedLines.WriteString(list[1])
	} else {
		for _, line := range list[:len(list)-1] {
			joinedLines.WriteString(line)
			joinedLines.WriteString(", ")
		}

		joinedLines.WriteString("and ")
		joinedLines.WriteString(list[len(list)-1])
	}

	return joinedLines.String()
}

type nullSlackNotifier struct{}

func (n *nullSlackNotifier) SendNotification(logger lager.Logger, notification Notification) error {
	logger.Session("send-notification").Debug("done")

	return nil
}

func (n *nullSlackNotifier) SendBatchNotification(logger lager.Logger, batch []Notification) error {
	logger.Session("send-batch-notification").Debug("done")

	return nil
}
<commit_msg>ensure map iteration order<commit_after>package notifications

import (
	"bytes"
	"encoding/json"
	"errors"
	"fmt"
	"net/http"
	"sort"
	"strconv"
	"strings"
	"time"

	"code.cloudfoundry.org/clock"
	"code.cloudfoundry.org/lager"
)

//go:generate counterfeiter . Notifier

type Notifier interface {
	SendNotification(lager.Logger, Notification) error
	SendBatchNotification(lager.Logger, []Notification) error
}

type Notification struct {
	Owner      string
	Repository string
	Private    bool

	SHA string

	Path       string
	LineNumber int
}

func (n Notification) FullName() string {
	return fmt.Sprintf("%s/%s", n.Owner, n.Repository)
}

func (n Notification) ShortSHA() string {
	return n.SHA[:7]
}

type slackNotifier struct {
	webhookURL string
	client     *http.Client
	clock      clock.Clock
}

type slackMessage struct {
	Attachments []slackAttachment `json:"attachments"`
}

type slackAttachment struct {
	Fallback string `json:"fallback"`
	Color    string `json:"color"`
	Title    string `json:"title"`
	Text     string `json:"text"`
}

func NewSlackNotifier(webhookURL string, clock clock.Clock) Notifier {
	if webhookURL == "" {
		return &nullSlackNotifier{}
	}

	return &slackNotifier{
		webhookURL: webhookURL,
		clock:      clock,
		client: &http.Client{
			Timeout: 3 * time.Second,
			Transport: &http.Transport{
				DisableKeepAlives: true,
			},
		},
	}
}

const maxRetries = 3

func (n *slackNotifier) SendNotification(logger lager.Logger, notification Notification) error {
	logger = logger.Session("send-notification")
	logger.Debug("starting")

	message := n.formatSlackMessage(notification)

	body, err := json.Marshal(message)
	if err != nil {
		logger.Error("unmarshal-faiiled", err)
		return err
	}

	return n.send(logger, body)
}

func (n *slackNotifier) SendBatchNotification(logger lager.Logger, batch []Notification) error {
	logger = logger.Session("send-batch-notification", lager.Data{"batch-size": len(batch)})

	logger.Debug("starting")

	if len(batch) == 0 {
		logger.Debug("done")
		return nil
	}

	messages := n.formatBatchSlackMessages(batch)

	for _, message := range messages {
		body, err := json.Marshal(message)
		if err != nil {
			logger.Error("unmarshal-faiiled", err)
			return err
		}

		err = n.send(logger, body)
		if err != nil {
			return err
		}
	}

	return nil
}

func (n *slackNotifier) send(logger lager.Logger, body []byte) error {
	for numReq := 0; numReq < maxRetries; numReq++ {
		req, err := http.NewRequest("POST", n.webhookURL, bytes.NewBuffer(body))
		if err != nil {
			logger.Error("request-failed", err)
			return err
		}

		req.Header.Set("Content-type", "application/json")

		resp, err := n.client.Do(req)
		if err != nil {
			logger.Error("response-error", err)
			return err
		}

		switch resp.StatusCode {
		case http.StatusOK:
			logger.Debug("done")
			return nil
		case http.StatusTooManyRequests:
			lastLoop := (numReq == maxRetries-1)
			if lastLoop {
				break
			}

			afterStr := resp.Header.Get("Retry-After")
			logger.Info("told-to-wait", lager.Data{"after": afterStr})
			after, err := strconv.Atoi(afterStr)
			if err != nil {
				logger.Error("failed", err)
				return err
			}

			wait := after + 1 // +1 for luck

			n.clock.Sleep(time.Duration(wait) * time.Second)
			continue
		default:
			err = fmt.Errorf("bad response (!200): %d", resp.StatusCode)
			logger.Error("bad-response", err)
			return err
		}
	}

	err := errors.New("retried too many times")
	logger.Error("failed", err)

	return err
}

func (n *slackNotifier) formatSlackMessage(not Notification) slackMessage {
	link := fmt.Sprintf("https://github.com/%s/%s/blob/%s/%s#L%d", not.Owner, not.Repository, not.SHA, not.Path, not.LineNumber)

	color := "danger"
	if not.Private {
		color = "warning"
	}

	return slackMessage{
		Attachments: []slackAttachment{
			{
				Fallback: link,
				Color:    color,
				Title:    fmt.Sprintf("Credential detected in %s!", not.FullName()),
				Text:     fmt.Sprintf("<%s|%s:%d>", link, not.Path, not.LineNumber),
			},
		},
	}
}

type slackLink struct {
	Text string
	Href string
}

func (l slackLink) String() string {
	return fmt.Sprintf("<%s|%s>", l.Href, l.Text)
}

type slackBatchRepo struct {
	Owner      string
	Repository string
	SHA        string
	Private    bool
}

func (r slackBatchRepo) FullName() string {
	return fmt.Sprintf("%s/%s", r.Owner, r.Repository)
}

func (r slackBatchRepo) ShortSHA() string {
	return r.SHA[:7]
}

func (n *slackNotifier) formatBatchSlackMessages(batch []Notification) []slackMessage {
	messages := []slackMessage{}

	messageMap := make(map[slackBatchRepo]map[string][]Notification)

	for _, not := range batch {
		repo := slackBatchRepo{
			Owner:      not.Owner,
			Repository: not.Repository,
			SHA:        not.SHA,
			Private:    not.Private,
		}

		_, found := messageMap[repo]
		if !found {
			messageMap[repo] = make(map[string][]Notification)
		}

		messageMap[repo][not.Path] = append(messageMap[repo][not.Path], not)
	}

	for repo, files := range messageMap {
		commitLink := fmt.Sprintf("https://github.com/%s/%s/commit/%s", repo.Owner, repo.Repository, repo.SHA)
		title := fmt.Sprintf("Possible credentials found in %s!", slackLink{
			Text: fmt.Sprintf("%s / %s", repo.FullName(), repo.ShortSHA()),
			Href: commitLink,
		})
		fallback := fmt.Sprintf("Possible credentials found in %s!", commitLink)

		color := "danger"
		if repo.Private {
			color = "warning"
		}

		// Make sure we get a consistent map iteration order.
		fileNames := []string{}
		for path, _ := range files {
			fileNames = append(fileNames, path)
		}
		sort.Strings(fileNames)

		fileLines := []string{}

		for _, path := range fileNames {
			nots := files[path]
			fileLink := fmt.Sprintf("https://github.com/%s/%s/blob/%s/%s", repo.Owner, repo.Repository, repo.SHA, path)

			lineLinks := []string{}

			for _, not := range nots {
				lineLink := fmt.Sprintf("%s#L%d", fileLink, not.LineNumber)

				lineLinks = append(lineLinks, slackLink{
					Text: strconv.Itoa(not.LineNumber),
					Href: lineLink,
				}.String())
			}

			plurality := "line"
			if len(lineLinks) > 1 {
				plurality = "lines"
			}

			text := fmt.Sprintf("• %s on %s %s", slackLink{
				Text: path,
				Href: fileLink,
			}, plurality, humanizeList(lineLinks))

			fileLines = append(fileLines, text)
		}

		messages = append(messages, slackMessage{
			Attachments: []slackAttachment{
				{
					Title:    title,
					Text:     strings.Join(fileLines, "\n"),
					Color:    color,
					Fallback: fallback,
				},
			},
		})
	}

	return messages
}

func humanizeList(list []string) string {
	joinedLines := &bytes.Buffer{}

	if len(list) <= 1 {
		joinedLines.WriteString(list[0])
	} else if len(list) == 2 {
		joinedLines.WriteString(list[0])
		joinedLines.WriteString(" and ")
		joinedLines.WriteString(list[1])
	} else {
		for _, line := range list[:len(list)-1] {
			joinedLines.WriteString(line)
			joinedLines.WriteString(", ")
		}

		joinedLines.WriteString("and ")
		joinedLines.WriteString(list[len(list)-1])
	}

	return joinedLines.String()
}

type nullSlackNotifier struct{}

func (n *nullSlackNotifier) SendNotification(logger lager.Logger, notification Notification) error {
	logger.Session("send-notification").Debug("done")

	return nil
}

func (n *nullSlackNotifier) SendBatchNotification(logger lager.Logger, batch []Notification) error {
	logger.Session("send-batch-notification").Debug("done")

	return nil
}
<|endoftext|>


<commit_before>package file

import (
	"testing"

	"github.com/pierrre/imageserver"
	imageserver_source "github.com/pierrre/imageserver/source"
	"github.com/pierrre/imageserver/testdata"
)

func BenchmarkServerGetSmall(b *testing.B) {
	benchmarkServerGet(b, testdata.SmallFileName)
}

func BenchmarkServerGetMedium(b *testing.B) {
	benchmarkServerGet(b, testdata.MediumFileName)
}

func BenchmarkServerGetLarge(b *testing.B) {
	benchmarkServerGet(b, testdata.LargeFileName)
}

func BenchmarkServerGetHuge(b *testing.B) {
	benchmarkServerGet(b, testdata.HugeFileName)
}

func benchmarkServerGet(b *testing.B, filename string) {
	srv := &Server{
		Root: testdata.Dir,
	}
	params := imageserver.Params{
		imageserver_source.Param: filename,
	}
	var bs int
	b.ResetTimer()
	b.RunParallel(func(pb *testing.PB) {
		for pb.Next() {
			im, err := srv.Get(params)
			if err != nil {
				b.Fatal(err)
			}
			bs = len(im.Data)
		}
	})
	b.SetBytes(int64(bs))
}
<commit_msg>source/file: remove parallelism in benchmarks, it fixes a data race, and is not really useful<commit_after>package file

import (
	"testing"

	"github.com/pierrre/imageserver"
	imageserver_source "github.com/pierrre/imageserver/source"
	"github.com/pierrre/imageserver/testdata"
)

func BenchmarkServerGetSmall(b *testing.B) {
	benchmarkServerGet(b, testdata.SmallFileName)
}

func BenchmarkServerGetMedium(b *testing.B) {
	benchmarkServerGet(b, testdata.MediumFileName)
}

func BenchmarkServerGetLarge(b *testing.B) {
	benchmarkServerGet(b, testdata.LargeFileName)
}

func BenchmarkServerGetHuge(b *testing.B) {
	benchmarkServerGet(b, testdata.HugeFileName)
}

func benchmarkServerGet(b *testing.B, filename string) {
	srv := &Server{
		Root: testdata.Dir,
	}
	params := imageserver.Params{
		imageserver_source.Param: filename,
	}
	var bs int
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		im, err := srv.Get(params)
		if err != nil {
			b.Fatal(err)
		}
		bs = len(im.Data)
	}
	b.SetBytes(int64(bs))
}
<|endoftext|>