<commit_before>// Package server implements a configurable, general-purpose web server.
// It relies on configurations obtained from the adjacent config package
// and can execute middleware as defined by the adjacent middleware package.
package server

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"

	"github.com/bradfitz/http2"
)

// Server represents an instance of a server, which serves
// static content at a particular address (host and port).
type Server struct {
	HTTP2   bool                   // temporary while http2 is not in std lib (TODO: remove flag when part of std lib)
	address string                 // the actual address for net.Listen to listen on
	tls     bool                   // whether this server is serving all HTTPS hosts or not
	vhosts  map[string]virtualHost // virtual hosts keyed by their address
}

// New creates a new Server which will bind to addr and serve
// the sites/hosts configured in configs. This function does
// not start serving.
func New(addr string, configs []Config) (*Server, error) {
	var tls bool
	if len(configs) > 0 {
		tls = configs[0].TLS.Enabled
	}

	s := &Server{
		address: addr,
		tls:     tls,
		vhosts:  make(map[string]virtualHost),
	}

	for _, conf := range configs {
		if _, exists := s.vhosts[conf.Host]; exists {
			return nil, fmt.Errorf("cannot serve %s - host already defined for address %s", conf.Address(), s.address)
		}

		vh := virtualHost{config: conf}

		// Build middleware stack
		err := vh.buildStack()
		if err != nil {
			return nil, err
		}

		s.vhosts[conf.Host] = vh
	}

	return s, nil
}

// Serve starts the server. It blocks until the server quits.
func (s *Server) Serve() error {
	server := &http.Server{
		Addr:    s.address,
		Handler: s,
	}

	if s.HTTP2 {
		// TODO: This call may not be necessary after HTTP/2 is merged into std lib
		http2.ConfigureServer(server, nil)
	}

	for _, vh := range s.vhosts {
		// Execute startup functions now
		for _, start := range vh.config.Startup {
			err := start()
			if err != nil {
				return err
			}
		}

		// Execute shutdown commands on exit
		if len(vh.config.Shutdown) > 0 {
			go func(vh virtualHost) {
				// Wait for signal
				interrupt := make(chan os.Signal, 1)
				signal.Notify(interrupt, os.Interrupt, os.Kill)
				<-interrupt

				// Run callbacks
				exitCode := 0
				for _, shutdownFunc := range vh.config.Shutdown {
					err := shutdownFunc()
					if err != nil {
						exitCode = 1
						log.Println(err)
					}
				}
				os.Exit(exitCode) // BUG: Other shutdown goroutines might be running; use sync.WaitGroup
			}(vh)
		}
	}

	if s.tls {
		var tlsConfigs []TLSConfig
		for _, vh := range s.vhosts {
			tlsConfigs = append(tlsConfigs, vh.config.TLS)
		}
		return ListenAndServeTLSWithSNI(server, tlsConfigs)
	}
	return server.ListenAndServe()
}

// ListenAndServeTLSWithSNI serves TLS with Server Name Indication (SNI) support, which allows
// multiple sites (different hostnames) to be served from the same address. This method is
// adapted directly from the std lib's net/http ListenAndServeTLS function, which was
// written by the Go Authors. It has been modified to support multiple certificate/key pairs.
func ListenAndServeTLSWithSNI(srv *http.Server, tlsConfigs []TLSConfig) error {
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	config := new(tls.Config)
	if srv.TLSConfig != nil {
		*config = *srv.TLSConfig
	}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	// Here we diverge from the stdlib a bit by loading multiple certs/key pairs
	// then we map the server names to their certs
	var err error
	config.Certificates = make([]tls.Certificate, len(tlsConfigs))
	for i, tlsConfig := range tlsConfigs {
		config.Certificates[i], err = tls.LoadX509KeyPair(tlsConfig.Certificate, tlsConfig.Key)
		if err != nil {
			return err
		}
	}
	config.BuildNameToCertificate()

	// Customize our TLS configuration
	config.MinVersion = tlsConfigs[0].ProtocolMinVersion
	config.MaxVersion = tlsConfigs[0].ProtocolMaxVersion
	config.CipherSuites = tlsConfigs[0].Ciphers
	config.PreferServerCipherSuites = tlsConfigs[0].PreferServerCipherSuites

	// TLS client authentication, if user enabled it
	err = setupClientAuth(tlsConfigs, config)
	if err != nil {
		return err
	}

	// Create listener and we're on our way
	conn, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	tlsListener := tls.NewListener(conn, config)

	return srv.Serve(tlsListener)
}

// setupClientAuth sets up TLS client authentication only if
// any of the TLS configs specified at least one cert file.
func setupClientAuth(tlsConfigs []TLSConfig, config *tls.Config) error {
	var clientAuth bool
	for _, cfg := range tlsConfigs {
		if len(cfg.ClientCerts) > 0 {
			clientAuth = true
			break
		}
	}

	if clientAuth {
		pool := x509.NewCertPool()
		for _, cfg := range tlsConfigs {
			for _, caFile := range cfg.ClientCerts {
				caCrt, err := ioutil.ReadFile(caFile) // Anyone that gets a cert from Matt Holt can connect
				if err != nil {
					return err
				}
				if !pool.AppendCertsFromPEM(caCrt) {
					return fmt.Errorf("error loading client certificate '%s': no certificates were successfully parsed", caFile)
				}
			}
		}
		config.ClientCAs = pool
		config.ClientAuth = tls.RequireAndVerifyClientCert
	}

	return nil
}

// ServeHTTP is the entry point for every request to the address that s
// is bound to. It acts as a multiplexer for the requests hostname as
// defined in the Host header so that the correct virtualhost
// (configuration and middleware stack) will handle the request.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer func() {
		// In case the user doesn't enable error middleware, we still
		// need to make sure that we stay alive up here
		if rec := recover(); rec != nil {
			http.Error(w, http.StatusText(http.StatusInternalServerError),
				http.StatusInternalServerError)
		}
	}()

	host, _, err := net.SplitHostPort(r.Host)
	if err != nil {
		host = r.Host // oh well
	}

	// Try the host as given, or try falling back to 0.0.0.0 (wildcard)
	if _, ok := s.vhosts[host]; !ok {
		if _, ok2 := s.vhosts["0.0.0.0"]; ok2 {
			host = "0.0.0.0"
		} else if _, ok2 := s.vhosts[""]; ok2 {
			host = ""
		}
	}

	if vh, ok := s.vhosts[host]; ok {
		w.Header().Set("Server", "Caddy")

		status, _ := vh.stack.ServeHTTP(w, r)

		// Fallback error response in case error handling wasn't chained in
		if status >= 400 {
			DefaultErrorFunc(w, r, status)
		}
	} else {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "No such host at %s", s.address)
	}
}

func DefaultErrorFunc(w http.ResponseWriter, r *http.Request, status int) {
	w.WriteHeader(status)
	fmt.Fprintf(w, "%d %s", status, http.StatusText(status))
}
<commit_msg>Use func from stdlib to clone *tls.Config for calming vet<commit_after>// Package server implements a configurable, general-purpose web server.
// It relies on configurations obtained from the adjacent config package
// and can execute middleware as defined by the adjacent middleware package.
package server

import (
	"crypto/tls"
	"crypto/x509"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"

	"github.com/bradfitz/http2"
)

// Server represents an instance of a server, which serves
// static content at a particular address (host and port).
type Server struct {
	HTTP2   bool                   // temporary while http2 is not in std lib (TODO: remove flag when part of std lib)
	address string                 // the actual address for net.Listen to listen on
	tls     bool                   // whether this server is serving all HTTPS hosts or not
	vhosts  map[string]virtualHost // virtual hosts keyed by their address
}

// New creates a new Server which will bind to addr and serve
// the sites/hosts configured in configs. This function does
// not start serving.
func New(addr string, configs []Config) (*Server, error) {
	var tls bool
	if len(configs) > 0 {
		tls = configs[0].TLS.Enabled
	}

	s := &Server{
		address: addr,
		tls:     tls,
		vhosts:  make(map[string]virtualHost),
	}

	for _, conf := range configs {
		if _, exists := s.vhosts[conf.Host]; exists {
			return nil, fmt.Errorf("cannot serve %s - host already defined for address %s", conf.Address(), s.address)
		}

		vh := virtualHost{config: conf}

		// Build middleware stack
		err := vh.buildStack()
		if err != nil {
			return nil, err
		}

		s.vhosts[conf.Host] = vh
	}

	return s, nil
}

// Serve starts the server. It blocks until the server quits.
func (s *Server) Serve() error {
	server := &http.Server{
		Addr:    s.address,
		Handler: s,
	}

	if s.HTTP2 {
		// TODO: This call may not be necessary after HTTP/2 is merged into std lib
		http2.ConfigureServer(server, nil)
	}

	for _, vh := range s.vhosts {
		// Execute startup functions now
		for _, start := range vh.config.Startup {
			err := start()
			if err != nil {
				return err
			}
		}

		// Execute shutdown commands on exit
		if len(vh.config.Shutdown) > 0 {
			go func(vh virtualHost) {
				// Wait for signal
				interrupt := make(chan os.Signal, 1)
				signal.Notify(interrupt, os.Interrupt, os.Kill)
				<-interrupt

				// Run callbacks
				exitCode := 0
				for _, shutdownFunc := range vh.config.Shutdown {
					err := shutdownFunc()
					if err != nil {
						exitCode = 1
						log.Println(err)
					}
				}
				os.Exit(exitCode) // BUG: Other shutdown goroutines might be running; use sync.WaitGroup
			}(vh)
		}
	}

	if s.tls {
		var tlsConfigs []TLSConfig
		for _, vh := range s.vhosts {
			tlsConfigs = append(tlsConfigs, vh.config.TLS)
		}
		return ListenAndServeTLSWithSNI(server, tlsConfigs)
	}
	return server.ListenAndServe()
}

// copy from net/http/transport.go
func cloneTLSConfig(cfg *tls.Config) *tls.Config {
	if cfg == nil {
		return &tls.Config{}
	}
	return &tls.Config{
		Rand:                     cfg.Rand,
		Time:                     cfg.Time,
		Certificates:             cfg.Certificates,
		NameToCertificate:        cfg.NameToCertificate,
		GetCertificate:           cfg.GetCertificate,
		RootCAs:                  cfg.RootCAs,
		NextProtos:               cfg.NextProtos,
		ServerName:               cfg.ServerName,
		ClientAuth:               cfg.ClientAuth,
		ClientCAs:                cfg.ClientCAs,
		InsecureSkipVerify:       cfg.InsecureSkipVerify,
		CipherSuites:             cfg.CipherSuites,
		PreferServerCipherSuites: cfg.PreferServerCipherSuites,
		SessionTicketsDisabled:   cfg.SessionTicketsDisabled,
		SessionTicketKey:         cfg.SessionTicketKey,
		ClientSessionCache:       cfg.ClientSessionCache,
		MinVersion:               cfg.MinVersion,
		MaxVersion:               cfg.MaxVersion,
		CurvePreferences:         cfg.CurvePreferences,
	}
}

// ListenAndServeTLSWithSNI serves TLS with Server Name Indication (SNI) support, which allows
// multiple sites (different hostnames) to be served from the same address. This method is
// adapted directly from the std lib's net/http ListenAndServeTLS function, which was
// written by the Go Authors. It has been modified to support multiple certificate/key pairs.
func ListenAndServeTLSWithSNI(srv *http.Server, tlsConfigs []TLSConfig) error {
	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	config := cloneTLSConfig(srv.TLSConfig)
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	// Here we diverge from the stdlib a bit by loading multiple certs/key pairs
	// then we map the server names to their certs
	var err error
	config.Certificates = make([]tls.Certificate, len(tlsConfigs))
	for i, tlsConfig := range tlsConfigs {
		config.Certificates[i], err = tls.LoadX509KeyPair(tlsConfig.Certificate, tlsConfig.Key)
		if err != nil {
			return err
		}
	}
	config.BuildNameToCertificate()

	// Customize our TLS configuration
	config.MinVersion = tlsConfigs[0].ProtocolMinVersion
	config.MaxVersion = tlsConfigs[0].ProtocolMaxVersion
	config.CipherSuites = tlsConfigs[0].Ciphers
	config.PreferServerCipherSuites = tlsConfigs[0].PreferServerCipherSuites

	// TLS client authentication, if user enabled it
	err = setupClientAuth(tlsConfigs, config)
	if err != nil {
		return err
	}

	// Create listener and we're on our way
	conn, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	tlsListener := tls.NewListener(conn, config)

	return srv.Serve(tlsListener)
}

// setupClientAuth sets up TLS client authentication only if
// any of the TLS configs specified at least one cert file.
func setupClientAuth(tlsConfigs []TLSConfig, config *tls.Config) error {
	var clientAuth bool
	for _, cfg := range tlsConfigs {
		if len(cfg.ClientCerts) > 0 {
			clientAuth = true
			break
		}
	}

	if clientAuth {
		pool := x509.NewCertPool()
		for _, cfg := range tlsConfigs {
			for _, caFile := range cfg.ClientCerts {
				caCrt, err := ioutil.ReadFile(caFile) // Anyone that gets a cert from Matt Holt can connect
				if err != nil {
					return err
				}
				if !pool.AppendCertsFromPEM(caCrt) {
					return fmt.Errorf("error loading client certificate '%s': no certificates were successfully parsed", caFile)
				}
			}
		}
		config.ClientCAs = pool
		config.ClientAuth = tls.RequireAndVerifyClientCert
	}

	return nil
}

// ServeHTTP is the entry point for every request to the address that s
// is bound to. It acts as a multiplexer for the requests hostname as
// defined in the Host header so that the correct virtualhost
// (configuration and middleware stack) will handle the request.
func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	defer func() {
		// In case the user doesn't enable error middleware, we still
		// need to make sure that we stay alive up here
		if rec := recover(); rec != nil {
			http.Error(w, http.StatusText(http.StatusInternalServerError),
				http.StatusInternalServerError)
		}
	}()

	host, _, err := net.SplitHostPort(r.Host)
	if err != nil {
		host = r.Host // oh well
	}

	// Try the host as given, or try falling back to 0.0.0.0 (wildcard)
	if _, ok := s.vhosts[host]; !ok {
		if _, ok2 := s.vhosts["0.0.0.0"]; ok2 {
			host = "0.0.0.0"
		} else if _, ok2 := s.vhosts[""]; ok2 {
			host = ""
		}
	}

	if vh, ok := s.vhosts[host]; ok {
		w.Header().Set("Server", "Caddy")

		status, _ := vh.stack.ServeHTTP(w, r)

		// Fallback error response in case error handling wasn't chained in
		if status >= 400 {
			DefaultErrorFunc(w, r, status)
		}
	} else {
		w.WriteHeader(http.StatusNotFound)
		fmt.Fprintf(w, "No such host at %s", s.address)
	}
}

func DefaultErrorFunc(w http.ResponseWriter, r *http.Request, status int) {
	w.WriteHeader(status)
	fmt.Fprintf(w, "%d %s", status, http.StatusText(status))
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"meowtrics/model"
	"os"
	"strconv"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/spf13/viper"
	"github.com/stretchr/graceful"
)

var (
	file            *os.File
	meowtricsLogger = log.New()
	logFileName     = "log-meowtrics.log"
	configFileName  = "meowtricsConfig"
	router          *mux.Router
	postSubrouter   *mux.Router
	getSubrouter    *mux.Router
	eventMap        map[string]model.ClientEventData
)

func init() {
	initApp()
	initRouter()
}

func initApp() {

	err := InitializeLogger(file, logFileName, meowtricsLogger, new(log.JSONFormatter))
	if err != nil {
		fmt.Println("Error initializing logger, defaulting to stdout. Error: " + err.Error())
	}

	err = LoadAppProperties(configFileName, DeploymentConfigDefaultPath, meowtricsLogger)
	if err != nil {
		meowtricsLogger.Panicln("Error reading app properties:" + err.Error())
	}

	eventMap = make(map[string]model.ClientEventData)
}

func initRouter() {

	router = mux.NewRouter()
	router.StrictSlash(true)

	v1Route := router.PathPrefix("/v1/")

	postSubrouter = v1Route.Methods("POST").Subrouter()
	postSubrouter.Handle("/events", CreateEventHandler())

	getSubrouter = v1Route.Methods("GET").Subrouter()
	getSubrouter.Handle("/events/{id:[0-9]+}", RetrieveEventHandler())

	router.Handle("/heartbeat", HeartBeatHandler())
	router.NotFoundHandler = NotFoundHandler()
}

func main() {

	n := negroni.Classic()
	n.UseHandler(router)

	appGracefulShutdownTimeinSeconds, err := strconv.Atoi(viper.GetString("appGracefulShutdownTimeinSeconds"))
	if err != nil {
		meowtricsLogger.Errorln("Error reading config: " + err.Error())
	}

	graceful.Run(":"+viper.GetString("appPort"), time.Duration(appGracefulShutdownTimeinSeconds)*time.Second, n)

	defer cleanup(file, meowtricsLogger)
}

func cleanup(file *os.File, logger *log.Logger) {
	logger.Println("Starting clean up")

	logger.Println("Closing file stream")
	file.Close()
}
<commit_msg>Refactored routes<commit_after>package main

import (
	"fmt"
	"meowtrics/model"
	"os"
	"strconv"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/gorilla/mux"
	"github.com/spf13/viper"
	"github.com/stretchr/graceful"
)

var (
	file            *os.File
	meowtricsLogger = log.New()
	logFileName     = "log-meowtrics.log"
	configFileName  = "meowtricsConfig"
	router          *mux.Router
	postSubrouter   *mux.Router
	getSubrouter    *mux.Router
	eventMap        map[string]model.ClientEventData
)

func init() {
	initApp()
	initRouter()
}

func initApp() {

	err := InitializeLogger(file, logFileName, meowtricsLogger, new(log.JSONFormatter))
	if err != nil {
		fmt.Println("Error initializing logger, defaulting to stdout. Error: " + err.Error())
	}

	err = LoadAppProperties(configFileName, DeploymentConfigDefaultPath, meowtricsLogger)
	if err != nil {
		meowtricsLogger.Panicln("Error reading app properties:" + err.Error())
	}

	eventMap = make(map[string]model.ClientEventData)
}

func initRouter() {

	router = mux.NewRouter()
	router.StrictSlash(true)

	postSubrouter = router.PathPrefix("/v1/").Methods("POST").Subrouter()
	postSubrouter.Handle("/events", CreateEventHandler())

	getSubrouter = router.PathPrefix("/v1/").Methods("GET").Subrouter()
	getSubrouter.Handle("/events/{id:[0-9]+}", RetrieveEventHandler())

	router.Handle("/heartbeat", HeartBeatHandler())
	router.NotFoundHandler = NotFoundHandler()
}

func main() {

	n := negroni.Classic()
	n.UseHandler(router)

	appGracefulShutdownTimeinSeconds, err := strconv.Atoi(viper.GetString("appGracefulShutdownTimeinSeconds"))
	if err != nil {
		meowtricsLogger.Errorln("Error reading config: " + err.Error())
	}

	graceful.Run(":"+viper.GetString("appPort"), time.Duration(appGracefulShutdownTimeinSeconds)*time.Second, n)

	defer cleanup(file, meowtricsLogger)
}

func cleanup(file *os.File, logger *log.Logger) {
	logger.Println("Starting clean up")

	logger.Println("Closing file stream")
	file.Close()
}
<|endoftext|>


<commit_before>package server

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/input"
	"github.com/packetzoom/logslammer/output"
)

type Server struct {
	Config *Config
	Buffer *buffer.Buffer

	mtx     sync.Mutex
	inputs  map[string]input.Input
	outputs map[string]output.Output
}

func signalCatcher() chan os.Signal {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	return c
}

func New(configFile string) (*Server, error) {
	config, err := LoadConfig(configFile)
	if err != nil {
		return nil, err
	}

	return &Server{
		Config:  config,
		Buffer:  buffer.New(),
		inputs:  make(map[string]input.Input),
		outputs: make(map[string]output.Output),
	}, nil
}

func (s *Server) Start() {
	log.Println("Starting server")

	// Start buffer
	log.Println("Starting buffer")
	go s.Buffer.Start()

	s.mtx.Lock()

	// Start inputs
	for name, config := range s.Config.Inputs {
		in, err := input.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := in.Init(config, s.Buffer); err != nil {
			log.Printf("Failed to init %s input: %v", name, err)
			continue
		}

		go func(name string, in input.Input) {
			log.Printf("Starting input %s", name)
			if err := in.Start(); err != nil {
				log.Printf("Error starting input %s: %v", name, err)
			}
		}(name, in)

		s.inputs[name] = in
	}

	// Start ouputs
	for name, config := range s.Config.Outputs {
		out, err := output.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := out.Init(config, s.Buffer); err != nil {
			log.Printf("Failed to init %s output: %v", name, err)
			continue
		}

		go func(name string, out output.Output) {
			log.Printf("Starting output %s", name)
			if err := out.Start(); err != nil {
				log.Printf("Error starting output %s: %v", name, err)
			}
		}(name, out)

		s.outputs[name] = out
	}

	s.mtx.Unlock()

	// Wait for kill signal
	<-signalCatcher()
	log.Printf("Received quit signal")

	// Stop Server
	s.Stop()
}

func (s *Server) Stop() {
	log.Println("Stopping server")

	s.mtx.Lock()

	// stop inputs
	for name, in := range s.inputs {
		log.Printf("Stopping input %s", name)
		if err := in.Stop(); err != nil {
			log.Printf("Error stopping %s input: %v", name, err)
		}
	}

	// stop ouputs
	for name, out := range s.outputs {
		log.Printf("Stopping output %s", name)
		if err := out.Stop(); err != nil {
			log.Printf("Error stopping %s output: %v", name, err)
		}
	}

	s.mtx.Unlock()

	log.Println("Stopping buffer")
	s.Buffer.Stop()
}
<commit_msg>Fail fast if input configure fails<commit_after>package server

import (
	"log"
	"os"
	"os/signal"
	"sync"
	"syscall"

	"github.com/packetzoom/logslammer/buffer"
	"github.com/packetzoom/logslammer/input"
	"github.com/packetzoom/logslammer/output"
)

type Server struct {
	Config *Config
	Buffer *buffer.Buffer

	mtx     sync.Mutex
	inputs  map[string]input.Input
	outputs map[string]output.Output
}

func signalCatcher() chan os.Signal {
	c := make(chan os.Signal, 1)
	signal.Notify(c, syscall.SIGINT, syscall.SIGTERM, syscall.SIGQUIT)
	return c
}

func New(configFile string) (*Server, error) {
	config, err := LoadConfig(configFile)
	if err != nil {
		return nil, err
	}

	return &Server{
		Config:  config,
		Buffer:  buffer.New(),
		inputs:  make(map[string]input.Input),
		outputs: make(map[string]output.Output),
	}, nil
}

func (s *Server) Start() {
	log.Println("Starting server")

	// Start buffer
	log.Println("Starting buffer")
	go s.Buffer.Start()

	s.mtx.Lock()

	// Start inputs
	for name, config := range s.Config.Inputs {
		in, err := input.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := in.Init(config, s.Buffer); err != nil {
			log.Fatalf("Failed to init %s input: %v", name, err)
		}

		go func(name string, in input.Input) {
			log.Printf("Starting input %s", name)
			if err := in.Start(); err != nil {
				log.Printf("Error starting input %s: %v", name, err)
			}
		}(name, in)

		s.inputs[name] = in
	}

	// Start ouputs
	for name, config := range s.Config.Outputs {
		out, err := output.Load(name)
		if err != nil {
			log.Println(err.Error)
			continue
		}

		if err := out.Init(config, s.Buffer); err != nil {
			log.Printf("Failed to init %s output: %v", name, err)
			continue
		}

		go func(name string, out output.Output) {
			log.Printf("Starting output %s", name)
			if err := out.Start(); err != nil {
				log.Printf("Error starting output %s: %v", name, err)
			}
		}(name, out)

		s.outputs[name] = out
	}

	s.mtx.Unlock()

	// Wait for kill signal
	<-signalCatcher()
	log.Printf("Received quit signal")

	// Stop Server
	s.Stop()
}

func (s *Server) Stop() {
	log.Println("Stopping server")

	s.mtx.Lock()

	// stop inputs
	for name, in := range s.inputs {
		log.Printf("Stopping input %s", name)
		if err := in.Stop(); err != nil {
			log.Printf("Error stopping %s input: %v", name, err)
		}
	}

	// stop ouputs
	for name, out := range s.outputs {
		log.Printf("Stopping output %s", name)
		if err := out.Stop(); err != nil {
			log.Printf("Error stopping %s output: %v", name, err)
		}
	}

	s.mtx.Unlock()

	log.Println("Stopping buffer")
	s.Buffer.Stop()
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/gophergala2016/3wordgame/validation"
	"net"
)

// Client struct
type Client struct {
	incoming chan string
	outgoing chan string
	reader   *bufio.Reader
	writer   *bufio.Writer
}

// Read line by line into the client.incoming
func (client *Client) Read() {
	for {
		line, _ := client.reader.ReadString('\n')
		client.incoming <- line
	}
}

// Write client outgoing data to the client writer
func (client *Client) Write() {
	for data := range client.outgoing {
		client.writer.WriteString(data)
		client.writer.Flush()
	}
}

// Listen for reads and writes on the client
func (client *Client) Listen() {
	go client.Read()
	go client.Write()
}

// NewClient returns new instance of client.
func NewClient(connection net.Conn) *Client {
	writer := bufio.NewWriter(connection)
	reader := bufio.NewReader(connection)

	client := &Client{
		incoming: make(chan string),
		outgoing: make(chan string),
		reader:   reader,
		writer:   writer,
	}

	client.Listen()

	return client
}

// ChatRoom struct
type ChatRoom struct {
	clients  []*Client
	joins    chan net.Conn
	incoming chan string
	outgoing chan string
}

// Broadcast data to all connected chatRoom.clients
func (chatRoom *ChatRoom) Broadcast(data string) {
	for _, client := range chatRoom.clients {
		client.outgoing <- data
	}
}

// Join attaches a new client to the chatRoom clients
func (chatRoom *ChatRoom) Join(connection net.Conn) {
	client := NewClient(connection)
	chatRoom.clients = append(chatRoom.clients, client)
	go func() {
		for {
			chatRoom.incoming <- <-client.incoming
		}
	}()
}

// Listen to all incoming messages for the chatRoom
func (chatRoom *ChatRoom) Listen() {
	go func() {
		for {
			select {
			case data := <-chatRoom.incoming:
				msg, err := validation.ValidateMsg(data)
				if err == nil {
					chatRoom.Broadcast(msg)
				}
			case conn := <-chatRoom.joins:
				chatRoom.Join(conn)
			}
		}
	}()
}

// NewChatRoom factories a ChatRoom instance
func NewChatRoom() *ChatRoom {
	chatRoom := &ChatRoom{
		clients:  make([]*Client, 0),
		joins:    make(chan net.Conn),
		incoming: make(chan string),
		outgoing: make(chan string),
	}

	chatRoom.Listen()

	return chatRoom
}

func main() {
	var server string
	var port int

	flag.StringVar(&server, "server", "127.0.0.1", "Server host")
	flag.IntVar(&port, "port", 6666, "Server port")
	flag.Parse()

	listener, _ := net.Listen("tcp", fmt.Sprintf("%s:%d", server, port))

	chatRoom := NewChatRoom()

	for {
		conn, _ := listener.Accept()
		chatRoom.joins <- conn
	}
}
<commit_msg>Ensure consecutive messages are from unique users<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/gophergala2016/3wordgame/validation"
	"net"
)

// Client struct
type Client struct {
	incoming chan Message
	outgoing chan string
	reader *bufio.Reader
	writer *bufio.Writer
	address string
}

type Message struct {
	text string
	address string
}

// Read line by line into the client.incoming
func (client *Client) Read() {
	for {
		line, _ := client.reader.ReadString('\n')
		client.incoming <- Message{
				text: line,
				address: client.address,
			}
	}
}

// Write client outgoing data to the client writer
func (client *Client) Write() {
	for data := range client.outgoing {
		client.writer.WriteString(data)
		client.writer.Flush()
	}
}

// Listen for reads and writes on the client
func (client *Client) Listen() {
	go client.Read()
	go client.Write()
}

// NewClient returns new instance of client.
func NewClient(connection net.Conn) *Client {
	writer := bufio.NewWriter(connection)
	reader := bufio.NewReader(connection)

	address := connection.RemoteAddr().String()

	client := &Client{
		incoming: make(chan Message),
		outgoing: make(chan string),
		reader:   reader,
		writer:   writer,
		address: address,
	}

	client.Listen()

	return client
}

// ChatRoom struct
type ChatRoom struct {
	clients  []*Client
	joins    chan net.Conn
	incoming chan Message
	outgoing chan string
	last_msg_user_address string
}

// Broadcast data to all connected chatRoom.clients
func (chatRoom *ChatRoom) Broadcast(data string) {
	for _, client := range chatRoom.clients {
		client.outgoing <- data
	}
}

// Join attaches a new client to the chatRoom clients
func (chatRoom *ChatRoom) Join(connection net.Conn) {
	client := NewClient(connection)
	chatRoom.clients = append(chatRoom.clients, client)
	go func() {
		for {
			chatRoom.incoming <- <-client.incoming
		}
	}()
}

// Listen to all incoming messages for the chatRoom
func (chatRoom *ChatRoom) Listen() {
	go func() {
		for {
			select {
			case data := <-chatRoom.incoming:
				msg, err := validation.ValidateMsg(data.text)
				if err == nil && chatRoom.last_msg_user_address != data.address {
					chatRoom.Broadcast(msg)
					chatRoom.last_msg_user_address = data.address
				}
			case conn := <-chatRoom.joins:
				chatRoom.Join(conn)
			}
		}
	}()
}

// NewChatRoom factories a ChatRoom instance
func NewChatRoom() *ChatRoom {
	chatRoom := &ChatRoom{
		clients:  make([]*Client, 0),
		joins:    make(chan net.Conn),
		incoming: make(chan Message),
		outgoing: make(chan string),
	}

	chatRoom.Listen()

	return chatRoom
}

func main() {
	var server string
	var port int

	flag.StringVar(&server, "server", "127.0.0.1", "Server host")
	flag.IntVar(&port, "port", 6666, "Server port")
	flag.Parse()

	listener, _ := net.Listen("tcp", fmt.Sprintf("%s:%d", server, port))

	chatRoom := NewChatRoom()

	for {
		conn, _ := listener.Accept()
		chatRoom.joins <- conn
	}
}
<|endoftext|>


<commit_before>package server

import (
	"context"
	"io/ioutil"
	"net"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/golog"
	"github.com/getlantern/ops"
	"github.com/getlantern/proxy"
	"github.com/getlantern/proxy/filters"
	"github.com/getlantern/tlsdefaults"

	"github.com/getlantern/http-proxy/buffers"
	"github.com/getlantern/http-proxy/listeners"
)

var (
	testingLocal = false
	log          = golog.LoggerFor("server")
)

type listenerGenerator func(net.Listener) net.Listener

// Opts are used to configure a Server
type Opts struct {
	IdleTimeout time.Duration
	Filter      filters.Filter
	Dial        proxy.DialFunc

	// OKDoesNotWaitForUpstream can be set to true in order to immediately return
	// OK to CONNECT requests.
	OKDoesNotWaitForUpstream bool

	// On Error provides a callback that's invoked if the proxy encounters an
	// error while proxying for the given client connection.
	OnError func(conn net.Conn, err error)
}

// Server is an HTTP proxy server.
type Server struct {
	// Allow is a function that determines whether or not to allow connections
	// from the given IP address. If unspecified, all connections are allowed.
	Allow              func(string) bool
	proxy              proxy.Proxy
	listenerGenerators []listenerGenerator
	onError            func(conn net.Conn, err error)
}

// New constructs a new HTTP proxy server using the given options
func New(opts *Opts) *Server {
	p, _ := proxy.New(&proxy.Opts{
		IdleTimeout:         opts.IdleTimeout,
		Dial:                opts.Dial,
		Filter:              opts.Filter,
		BufferSource:        buffers.Pool(),
		OKWaitsForUpstream:  !opts.OKDoesNotWaitForUpstream,
		OKSendsServerTiming: true,
		OnError: func(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
			status := http.StatusBadGateway
			if read {
				status = http.StatusBadRequest
			}
			return &http.Response{
				Request:    req,
				StatusCode: status,
				Body:       ioutil.NopCloser(strings.NewReader(err.Error())),
			}
		},
	})

	if opts.OnError == nil {
		opts.OnError = func(conn net.Conn, err error) {}
	}
	return &Server{
		proxy:   p,
		onError: opts.OnError,
	}
}

func (s *Server) AddListenerWrappers(listenerGens ...listenerGenerator) {
	for _, g := range listenerGens {
		s.listenerGenerators = append(s.listenerGenerators, g)
	}
}

func (s *Server) ListenAndServeHTTP(addr string, readyCb func(addr string)) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.Debugf("Listen http on %s", addr)
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) ListenAndServeHTTPS(addr, keyfile, certfile string, readyCb func(addr string)) error {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	listener, err := tlsdefaults.NewListener(s.wrapListenerIfNecessary(l), keyfile, certfile)
	if err != nil {
		return err
	}
	log.Debugf("Listen https on %s", addr)
	return s.serve(listener, readyCb)
}

func (s *Server) Serve(listener net.Listener, readyCb func(addr string)) error {
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) serve(listener net.Listener, readyCb func(addr string)) error {
	l := listeners.NewDefaultListener(listener)

	for _, wrap := range s.listenerGenerators {
		l = wrap(l)
	}

	if readyCb != nil {
		readyCb(l.Addr().String())
	}

	var tempDelay time.Duration // how long to sleep on accept failure
	for {
		conn, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				// delay code based on net/http.Server
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Errorf("http: Accept error: %v; retrying in %v", err, tempDelay)
				time.Sleep(tempDelay)
				continue
			}
			return errors.New("Error accepting: %v", err)
		}
		tempDelay = 0
		s.handle(conn)
	}
}

func (s *Server) handle(conn net.Conn) {
	wrapConn, isWrapConn := conn.(listeners.WrapConn)
	if isWrapConn {
		wrapConn.OnState(http.StateNew)
	}
	go s.doHandle(conn, isWrapConn, wrapConn)
}

func (s *Server) doHandle(conn net.Conn, isWrapConn bool, wrapConn listeners.WrapConn) {
	clientIP := ""
	remoteAddr := conn.RemoteAddr()
	if remoteAddr != nil {
		clientIP, _, _ = net.SplitHostPort(remoteAddr.String())
	}
	op := ops.Begin("http_proxy_handle").Set("client_ip", clientIP)
	defer op.End()

	defer func() {
		p := recover()
		if p != nil {
			err := log.Errorf("Caught panic handling connection from %v: %v", conn.RemoteAddr(), p)
			if op != nil {
				op.FailIf(err)
			}
			safeClose(conn)
		}
	}()

	err := s.proxy.Handle(context.Background(), conn, conn)
	if err != nil {
		op.FailIf(errors.New("Error handling connection from %v: %v", conn.RemoteAddr(), err))
		s.onError(conn, err)
	}
	if isWrapConn {
		wrapConn.OnState(http.StateClosed)
	}
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection from %v: %v", conn.RemoteAddr(), p)
		}
	}()

	conn.Close()
}

func (s *Server) wrapListenerIfNecessary(l net.Listener) net.Listener {
	if s.Allow != nil {
		log.Debug("Wrapping listener with Allow")
		return &allowinglistener{l, s.Allow}
	}
	return l
}

type allowinglistener struct {
	wrapped net.Listener
	allow   func(string) bool
}

func (l *allowinglistener) Accept() (net.Conn, error) {
	conn, err := l.wrapped.Accept()
	if err != nil {
		return conn, err
	}

	ip := ""
	remoteAddr := conn.RemoteAddr()
	switch addr := remoteAddr.(type) {
	case *net.TCPAddr:
		ip = addr.IP.String()
	case *net.UDPAddr:
		ip = addr.IP.String()
	default:
		log.Errorf("Remote addr %v is of unknown type %v, unable to determine IP", remoteAddr, reflect.TypeOf(remoteAddr))
		return conn, err
	}
	if !l.allow(ip) {
		conn.Close()
		// Note - we don't return an error, because that causes http.Server to stop
		// serving.
	}

	return conn, err
}

func (l *allowinglistener) Close() error {
	return l.wrapped.Close()
}

func (l *allowinglistener) Addr() net.Addr {
	return l.wrapped.Addr()
}
<commit_msg>Allow callers to interpret accept errors<commit_after>package server

import (
	"context"
	"io/ioutil"
	"net"
	"net/http"
	"reflect"
	"strings"
	"time"

	"github.com/getlantern/errors"
	"github.com/getlantern/golog"
	"github.com/getlantern/ops"
	"github.com/getlantern/proxy"
	"github.com/getlantern/proxy/filters"
	"github.com/getlantern/tlsdefaults"

	"github.com/getlantern/http-proxy/buffers"
	"github.com/getlantern/http-proxy/listeners"
)

var (
	testingLocal = false
	log          = golog.LoggerFor("server")
)

type listenerGenerator func(net.Listener) net.Listener

// Opts are used to configure a Server
type Opts struct {
	IdleTimeout time.Duration
	Filter      filters.Filter
	Dial        proxy.DialFunc

	// OKDoesNotWaitForUpstream can be set to true in order to immediately return
	// OK to CONNECT requests.
	OKDoesNotWaitForUpstream bool

	// OnError provides a callback that's invoked if the proxy encounters an
	// error while proxying for the given client connection.
	OnError func(conn net.Conn, err error)

	// OnAcceptError is called when the server fails to accept a connection.
	// If the error is fatal and should halt server operations, this callback
	// should return an error. That error will be returned by functions like
	// Serve, ListenAndServeHTTP, etc. If this callback returns nil, the
	// server will carry on.
	//
	// Temporary network errors (errors of type net.Error for which Temporary()
	// returns true) will not trigger this callback.
	OnAcceptError func(err error) (fatalErr error)
}

// Server is an HTTP proxy server.
type Server struct {
	// Allow is a function that determines whether or not to allow connections
	// from the given IP address. If unspecified, all connections are allowed.
	Allow              func(string) bool
	proxy              proxy.Proxy
	listenerGenerators []listenerGenerator
	onError            func(conn net.Conn, err error)
	onAcceptError      func(err error) (fatalErr error)
}

// New constructs a new HTTP proxy server using the given options
func New(opts *Opts) *Server {
	p, _ := proxy.New(&proxy.Opts{
		IdleTimeout:         opts.IdleTimeout,
		Dial:                opts.Dial,
		Filter:              opts.Filter,
		BufferSource:        buffers.Pool(),
		OKWaitsForUpstream:  !opts.OKDoesNotWaitForUpstream,
		OKSendsServerTiming: true,
		OnError: func(ctx filters.Context, req *http.Request, read bool, err error) *http.Response {
			status := http.StatusBadGateway
			if read {
				status = http.StatusBadRequest
			}
			return &http.Response{
				Request:    req,
				StatusCode: status,
				Body:       ioutil.NopCloser(strings.NewReader(err.Error())),
			}
		},
	})

	if opts.OnError == nil {
		opts.OnError = func(conn net.Conn, err error) {}
	}
	if opts.OnAcceptError == nil {
		opts.OnAcceptError = func(err error) (fatalErr error) { return err }
	}
	return &Server{
		proxy:         p,
		onError:       opts.OnError,
		onAcceptError: opts.OnAcceptError,
	}
}

func (s *Server) AddListenerWrappers(listenerGens ...listenerGenerator) {
	for _, g := range listenerGens {
		s.listenerGenerators = append(s.listenerGenerators, g)
	}
}

func (s *Server) ListenAndServeHTTP(addr string, readyCb func(addr string)) error {
	listener, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}
	log.Debugf("Listen http on %s", addr)
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) ListenAndServeHTTPS(addr, keyfile, certfile string, readyCb func(addr string)) error {
	l, err := net.Listen("tcp", addr)
	if err != nil {
		return err
	}

	listener, err := tlsdefaults.NewListener(s.wrapListenerIfNecessary(l), keyfile, certfile)
	if err != nil {
		return err
	}
	log.Debugf("Listen https on %s", addr)
	return s.serve(listener, readyCb)
}

func (s *Server) Serve(listener net.Listener, readyCb func(addr string)) error {
	return s.serve(s.wrapListenerIfNecessary(listener), readyCb)
}

func (s *Server) serve(listener net.Listener, readyCb func(addr string)) error {
	l := listeners.NewDefaultListener(listener)

	for _, wrap := range s.listenerGenerators {
		l = wrap(l)
	}

	if readyCb != nil {
		readyCb(l.Addr().String())
	}

	var tempDelay time.Duration // how long to sleep on accept failure
	for {
		conn, err := l.Accept()
		if err != nil {
			if ne, ok := err.(net.Error); ok && ne.Temporary() {
				// delay code based on net/http.Server
				if tempDelay == 0 {
					tempDelay = 5 * time.Millisecond
				} else {
					tempDelay *= 2
				}
				if max := 1 * time.Second; tempDelay > max {
					tempDelay = max
				}
				log.Errorf("http: Accept error: %v; retrying in %v", err, tempDelay)
				time.Sleep(tempDelay)
			} else if fatalErr := s.onAcceptError(err); fatalErr != nil {
				return fatalErr
			}
			continue
		}
		tempDelay = 0
		s.handle(conn)
	}
}

func (s *Server) handle(conn net.Conn) {
	wrapConn, isWrapConn := conn.(listeners.WrapConn)
	if isWrapConn {
		wrapConn.OnState(http.StateNew)
	}
	go s.doHandle(conn, isWrapConn, wrapConn)
}

func (s *Server) doHandle(conn net.Conn, isWrapConn bool, wrapConn listeners.WrapConn) {
	clientIP := ""
	remoteAddr := conn.RemoteAddr()
	if remoteAddr != nil {
		clientIP, _, _ = net.SplitHostPort(remoteAddr.String())
	}
	op := ops.Begin("http_proxy_handle").Set("client_ip", clientIP)
	defer op.End()

	defer func() {
		p := recover()
		if p != nil {
			err := log.Errorf("Caught panic handling connection from %v: %v", conn.RemoteAddr(), p)
			if op != nil {
				op.FailIf(err)
			}
			safeClose(conn)
		}
	}()

	err := s.proxy.Handle(context.Background(), conn, conn)
	if err != nil {
		op.FailIf(errors.New("Error handling connection from %v: %v", conn.RemoteAddr(), err))
		s.onError(conn, err)
	}
	if isWrapConn {
		wrapConn.OnState(http.StateClosed)
	}
}

func safeClose(conn net.Conn) {
	defer func() {
		p := recover()
		if p != nil {
			log.Errorf("Panic on closing connection from %v: %v", conn.RemoteAddr(), p)
		}
	}()

	conn.Close()
}

func (s *Server) wrapListenerIfNecessary(l net.Listener) net.Listener {
	if s.Allow != nil {
		log.Debug("Wrapping listener with Allow")
		return &allowinglistener{l, s.Allow}
	}
	return l
}

type allowinglistener struct {
	wrapped net.Listener
	allow   func(string) bool
}

func (l *allowinglistener) Accept() (net.Conn, error) {
	conn, err := l.wrapped.Accept()
	if err != nil {
		return conn, err
	}

	ip := ""
	remoteAddr := conn.RemoteAddr()
	switch addr := remoteAddr.(type) {
	case *net.TCPAddr:
		ip = addr.IP.String()
	case *net.UDPAddr:
		ip = addr.IP.String()
	default:
		log.Errorf("Remote addr %v is of unknown type %v, unable to determine IP", remoteAddr, reflect.TypeOf(remoteAddr))
		return conn, err
	}
	if !l.allow(ip) {
		conn.Close()
		// Note - we don't return an error, because that causes http.Server to stop
		// serving.
	}

	return conn, err
}

func (l *allowinglistener) Close() error {
	return l.wrapped.Close()
}

func (l *allowinglistener) Addr() net.Addr {
	return l.wrapped.Addr()
}
<|endoftext|>


<commit_before>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"os"

	"strings"
	"time"

	env "github.com/segmentio/go-env"

	redshift "github.com/Clever/s3-to-redshift/redshift"
	s3filepath "github.com/Clever/s3-to-redshift/s3filepath"
)

var (
	// things we are likely to change when running the worker normally are flags
	inputSchemaName = flag.String("schema", "mongo", "what target schema we load into")
	inputTables     = flag.String("tables", "", "target tables to run on, comma separated")
	inputBucket     = flag.String("bucket", "metrics", "bucket to load from, not including s3:// protocol")
	truncate        = flag.Bool("truncate", false, "do we truncate the table before inserting")
	force           = flag.Bool("force", false, "do we refresh the data even if it's already handled?")
	dataDate        = flag.String("date", "", "data date we should process, must be full RFC3339")
	configFile      = flag.String("config", "", "schema & table config to use in YAML format")
	gzip            = flag.Bool("gzip", true, "whether target files are gzipped, defaults to true")
	delimiter       = flag.String("delimiter", "", "delimiter for CSV files, usually pipe character")
	timeGranularity = flag.String("granularity", "day", "how often we expect to append new data")
	// things which will would strongly suggest launching as a second worker are env vars
	// also the secrets ... shhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh
	host               = os.Getenv("REDSHIFT_HOST")
	port               = os.Getenv("REDSHIFT_PORT")
	dbName             = env.MustGet("REDSHIFT_DB")
	user               = env.MustGet("REDSHIFT_USER")
	pwd                = env.MustGet("REDSHIFT_PASSWORD")
	awsRegion          = env.MustGet("AWS_REGION")
	awsAccessKeyID     = env.MustGet("AWS_ACCESS_KEY_ID")
	awsSecretAccessKey = env.MustGet("AWS_SECRET_ACCESS_KEY")
)

func fatalIfErr(err error, msg string) {
	if err != nil {
		panic(fmt.Sprintf("%s: %s", msg, err)) // TODO: kayvee
	}
}

// helper function used for verifying inputs
func getMapKeys(m map[string]bool) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

// Rounds down a dateTime to a granularity
// For instance, 11:50AM will be truncated to 11:00AM
// if given a granularity of an hour
func truncateDate(date time.Time, granularity string) time.Time {
	switch granularity {
	case "hour":
		return date.Truncate(time.Hour)
	default:
		// Round down to day granularity by default
		return date.Truncate(24 * time.Hour)
	}
}

// in a transaction, truncate, create or update, and then copy from the s3 dagta file or manifest
// yell loudly if there is anything different in the target table compared to config (different distkey, etc)
func runCopy(db *redshift.Redshift, inputConf s3filepath.S3File, inputTable redshift.Table, targetTable *redshift.Table, truncate, gzip bool, delimiter, timeGranularity string) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	// TRUNCATE for dimension tables, but not fact tables
	if truncate && targetTable != nil {
		log.Println("truncating table!")
		if err = db.Truncate(tx, inputConf.Schema, inputTable.Name); err != nil {
			return fmt.Errorf("err running truncate table: %s", err)
		}
	}
	if targetTable == nil {
		if err = db.CreateTable(tx, inputTable); err != nil {
			return fmt.Errorf("err running create table: %s", err)
		}
	} else {
		// To prevent duplicates, clear away any existing data within a certain time range as the data date
		// (that is, sharing the same data date up to a certain time granularity)
		if err = db.TruncateInTimeRange(tx, inputConf.Schema, inputTable.Name, inputConf.DataDate, timeGranularity, inputTable.Meta.DataDateColumn); err != nil {
			return fmt.Errorf("err truncating data for data refresh: %s", err)
		}

		if err = db.UpdateTable(tx, *targetTable, inputTable); err != nil {
			return fmt.Errorf("err running update table: %s", err)
		}
	}

	// COPY direct into it, ok to do since we're in a transaction
	// can't switch on file ending as manifest files b/c
	// manifest files obscure the underlying file types
	// instead just pass the delimiter along even if it's null
	if err = db.Copy(tx, inputConf, delimiter, true, gzip); err != nil {
		return fmt.Errorf("err running copy: %s", err)
	}

	if err = tx.Commit(); err != nil {
		return fmt.Errorf("err committing transaction: %s", err)
	}
	return nil
}

// This worker finds the latest file in s3 and uploads it to redshift
// If the destination table does not exist, the worker creates it
// If the destination table lacks columns, the worker creates those as well
// The worker also uses a column in the data to figure out whether the s3 data is
// newer than what already exists.
func main() {
	flag.Parse()

	// verify that timeGranularity is a supported value. for convenience,
	// we use the convention that granularities must be valid PostgreSQL dateparts
	// (see: http://www.postgresql.org/docs/8.1/static/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC)
	supportedGranularities := map[string]bool{"hour": true, "day": true}
	if !supportedGranularities[*timeGranularity] {
		panic(fmt.Sprintf("Unsupported granularity, must be one of %v", getMapKeys(supportedGranularities)))
	}

	// use an custom bucket type for testablitity
	bucket := s3filepath.S3Bucket{*inputBucket, awsRegion, awsAccessKeyID, awsSecretAccessKey}

	timeout := 60 // can parameterize later if this is an issue
	if host == "" {
		host = "localhost"
	}
	if port == "" {
		port = "5439"
	}
	db, err := redshift.NewRedshift(host, port, dbName, user, pwd, timeout)
	fatalIfErr(err, "error getting redshift instance")

	// for each table passed in - likely we could goroutine this out
	for _, t := range strings.Split(*inputTables, ",") {
		log.Printf("attempting to run on schema: %s table: %s", *inputSchemaName, t)
		// override most recent data file
		if *dataDate == "" {
			panic("No date provided")
		}
		parsedDate, err := time.Parse(time.RFC3339, *dataDate)
		fatalIfErr(err, fmt.Sprintf("issue parsing date: %s", *dataDate))
		inputConf, err := s3filepath.CreateS3File(s3filepath.S3PathChecker{}, bucket, *inputSchemaName, t, *configFile, parsedDate)
		fatalIfErr(err, "Issue getting data file from s3")
		inputTable, err := db.GetTableFromConf(*inputConf) // allow passing explicit config later
		fatalIfErr(err, "Issue getting table from input")

		// figure out what the current state of the table is to determine if the table is already up to date
		targetTable, lastTargetData, err := db.GetTableMetadata(inputConf.Schema, inputConf.Table, inputTable.Meta.DataDateColumn)
		if err != nil && err != sql.ErrNoRows { // ErrNoRows is fine, just means the table doesn't exist
			fatalIfErr(err, "Error getting existing latest table metadata") // use fatalIfErr to stay the same
		}

		// unless --force, don't update unless input data is new
		// Since lastTargetData comes from the columns, and
		// inputConf.DataDate comes from the filename, we round to
		// compare at the time granularity level
		if lastTargetData != nil && (truncateDate(*lastTargetData, *timeGranularity)).After(truncateDate(inputConf.DataDate, *timeGranularity)) {
			if *force == false {
				log.Printf("Recent data already exists in db: %s", *lastTargetData)
				return
			}
			log.Printf("Forcing update of inputTable: %s", inputConf.Table)
		}

		fatalIfErr(runCopy(db, *inputConf, *inputTable, targetTable, *truncate, *gzip, *delimiter, *timeGranularity), "Issue running copy")
		// DON'T NEED TO CREATE VIEWS - will be handled by the refresh script
		log.Printf("done with table: %s.%s", inputConf.Schema, t)
	}
	log.Println("done with full run")
}
<commit_msg>better scope error vars since its easy here, thx Christian :)<commit_after>package main

import (
	"database/sql"
	"flag"
	"fmt"
	"log"
	"os"

	"strings"
	"time"

	env "github.com/segmentio/go-env"

	redshift "github.com/Clever/s3-to-redshift/redshift"
	s3filepath "github.com/Clever/s3-to-redshift/s3filepath"
)

var (
	// things we are likely to change when running the worker normally are flags
	inputSchemaName = flag.String("schema", "mongo", "what target schema we load into")
	inputTables     = flag.String("tables", "", "target tables to run on, comma separated")
	inputBucket     = flag.String("bucket", "metrics", "bucket to load from, not including s3:// protocol")
	truncate        = flag.Bool("truncate", false, "do we truncate the table before inserting")
	force           = flag.Bool("force", false, "do we refresh the data even if it's already handled?")
	dataDate        = flag.String("date", "", "data date we should process, must be full RFC3339")
	configFile      = flag.String("config", "", "schema & table config to use in YAML format")
	gzip            = flag.Bool("gzip", true, "whether target files are gzipped, defaults to true")
	delimiter       = flag.String("delimiter", "", "delimiter for CSV files, usually pipe character")
	timeGranularity = flag.String("granularity", "day", "how often we expect to append new data")
	// things which will would strongly suggest launching as a second worker are env vars
	// also the secrets ... shhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhhh
	host               = os.Getenv("REDSHIFT_HOST")
	port               = os.Getenv("REDSHIFT_PORT")
	dbName             = env.MustGet("REDSHIFT_DB")
	user               = env.MustGet("REDSHIFT_USER")
	pwd                = env.MustGet("REDSHIFT_PASSWORD")
	awsRegion          = env.MustGet("AWS_REGION")
	awsAccessKeyID     = env.MustGet("AWS_ACCESS_KEY_ID")
	awsSecretAccessKey = env.MustGet("AWS_SECRET_ACCESS_KEY")
)

func fatalIfErr(err error, msg string) {
	if err != nil {
		panic(fmt.Sprintf("%s: %s", msg, err)) // TODO: kayvee
	}
}

// helper function used for verifying inputs
func getMapKeys(m map[string]bool) []string {
	keys := make([]string, 0, len(m))
	for k := range m {
		keys = append(keys, k)
	}
	return keys
}

// Rounds down a dateTime to a granularity
// For instance, 11:50AM will be truncated to 11:00AM
// if given a granularity of an hour
func truncateDate(date time.Time, granularity string) time.Time {
	switch granularity {
	case "hour":
		return date.Truncate(time.Hour)
	default:
		// Round down to day granularity by default
		return date.Truncate(24 * time.Hour)
	}
}

// in a transaction, truncate, create or update, and then copy from the s3 dagta file or manifest
// yell loudly if there is anything different in the target table compared to config (different distkey, etc)
func runCopy(db *redshift.Redshift, inputConf s3filepath.S3File, inputTable redshift.Table, targetTable *redshift.Table, truncate, gzip bool, delimiter, timeGranularity string) error {
	tx, err := db.Begin()
	if err != nil {
		return err
	}

	// TRUNCATE for dimension tables, but not fact tables
	if truncate && targetTable != nil {
		log.Println("truncating table!")
		if err := db.Truncate(tx, inputConf.Schema, inputTable.Name); err != nil {
			return fmt.Errorf("err running truncate table: %s", err)
		}
	}
	if targetTable == nil {
		if err := db.CreateTable(tx, inputTable); err != nil {
			return fmt.Errorf("err running create table: %s", err)
		}
	} else {
		// To prevent duplicates, clear away any existing data within a certain time range as the data date
		// (that is, sharing the same data date up to a certain time granularity)
		if err := db.TruncateInTimeRange(tx, inputConf.Schema, inputTable.Name, inputConf.DataDate, timeGranularity, inputTable.Meta.DataDateColumn); err != nil {
			return fmt.Errorf("err truncating data for data refresh: %s", err)
		}

		if err := db.UpdateTable(tx, *targetTable, inputTable); err != nil {
			return fmt.Errorf("err running update table: %s", err)
		}
	}

	// COPY direct into it, ok to do since we're in a transaction
	// can't switch on file ending as manifest files b/c
	// manifest files obscure the underlying file types
	// instead just pass the delimiter along even if it's null
	if err := db.Copy(tx, inputConf, delimiter, true, gzip); err != nil {
		return fmt.Errorf("err running copy: %s", err)
	}

	if err := tx.Commit(); err != nil {
		return fmt.Errorf("err committing transaction: %s", err)
	}
	return nil
}

// This worker finds the latest file in s3 and uploads it to redshift
// If the destination table does not exist, the worker creates it
// If the destination table lacks columns, the worker creates those as well
// The worker also uses a column in the data to figure out whether the s3 data is
// newer than what already exists.
func main() {
	flag.Parse()

	// verify that timeGranularity is a supported value. for convenience,
	// we use the convention that granularities must be valid PostgreSQL dateparts
	// (see: http://www.postgresql.org/docs/8.1/static/functions-datetime.html#FUNCTIONS-DATETIME-TRUNC)
	supportedGranularities := map[string]bool{"hour": true, "day": true}
	if !supportedGranularities[*timeGranularity] {
		panic(fmt.Sprintf("Unsupported granularity, must be one of %v", getMapKeys(supportedGranularities)))
	}

	// use an custom bucket type for testablitity
	bucket := s3filepath.S3Bucket{*inputBucket, awsRegion, awsAccessKeyID, awsSecretAccessKey}

	timeout := 60 // can parameterize later if this is an issue
	if host == "" {
		host = "localhost"
	}
	if port == "" {
		port = "5439"
	}
	db, err := redshift.NewRedshift(host, port, dbName, user, pwd, timeout)
	fatalIfErr(err, "error getting redshift instance")

	// for each table passed in - likely we could goroutine this out
	for _, t := range strings.Split(*inputTables, ",") {
		log.Printf("attempting to run on schema: %s table: %s", *inputSchemaName, t)
		// override most recent data file
		if *dataDate == "" {
			panic("No date provided")
		}
		parsedDate, err := time.Parse(time.RFC3339, *dataDate)
		fatalIfErr(err, fmt.Sprintf("issue parsing date: %s", *dataDate))
		inputConf, err := s3filepath.CreateS3File(s3filepath.S3PathChecker{}, bucket, *inputSchemaName, t, *configFile, parsedDate)
		fatalIfErr(err, "Issue getting data file from s3")
		inputTable, err := db.GetTableFromConf(*inputConf) // allow passing explicit config later
		fatalIfErr(err, "Issue getting table from input")

		// figure out what the current state of the table is to determine if the table is already up to date
		targetTable, lastTargetData, err := db.GetTableMetadata(inputConf.Schema, inputConf.Table, inputTable.Meta.DataDateColumn)
		if err != nil && err != sql.ErrNoRows { // ErrNoRows is fine, just means the table doesn't exist
			fatalIfErr(err, "Error getting existing latest table metadata") // use fatalIfErr to stay the same
		}

		// unless --force, don't update unless input data is new
		// Since lastTargetData comes from the columns, and
		// inputConf.DataDate comes from the filename, we round to
		// compare at the time granularity level
		if lastTargetData != nil && (truncateDate(*lastTargetData, *timeGranularity)).After(truncateDate(inputConf.DataDate, *timeGranularity)) {
			if *force == false {
				log.Printf("Recent data already exists in db: %s", *lastTargetData)
				return
			}
			log.Printf("Forcing update of inputTable: %s", inputConf.Table)
		}

		fatalIfErr(runCopy(db, *inputConf, *inputTable, targetTable, *truncate, *gzip, *delimiter, *timeGranularity), "Issue running copy")
		// DON'T NEED TO CREATE VIEWS - will be handled by the refresh script
		log.Printf("done with table: %s.%s", inputConf.Schema, t)
	}
	log.Println("done with full run")
}
<|endoftext|>


<commit_before>package server

import (
	"github.com/gin-gonic/gin"
)

// serveStaticFiles simply servs static files
func serveStaticFiles(r *gin.Engine) {
	r.StaticFile("/", "./web_static/build/index.html")
	r.StaticFile("/icons/json", "./icons/allIcons.json")
	r.StaticFile("/icons/png", "./icons/allIcons.png")
	r.Static("/js", "./web_static/build/js")
}
<commit_msg>Added the webp location to the icon<commit_after>package server

import (
	"github.com/gin-gonic/gin"
)

// serveStaticFiles simply servs static files
func serveStaticFiles(r *gin.Engine) {
	r.StaticFile("/", "./web_static/build/index.html")
	r.StaticFile("/icons/json", "./icons/allIcons.json")
	r.StaticFile("/icons/png", "./icons/allIcons.png")
	r.StaticFile("/icons/webp", "./icons/allIcons.webp")
	r.Static("/js", "./web_static/build/js")
}
<|endoftext|>


<commit_before>package quotas_test

import(
	"github.com/cloudfoundry/cf-acceptance-tests/helpers"
	"testing"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var context helpers.SuiteContext

func TestQuotas(t *testing.T) {
	RegisterFailHandler(Fail)

	config := helpers.LoadConfig()
	context = helpers.NewContext(config)
	environment := helpers.NewEnvironment(context)

	BeforeSuite(func() {
		environment.Setup()
	})

	AfterSuite(func() {
		environment.Teardown()
	})

	RunSpecs(t, "Quotas")
}
<commit_msg>Rejibber import lines<commit_after>package quotas_test

import (
	"testing"

	"github.com/cloudfoundry/cf-acceptance-tests/helpers"

	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var context helpers.SuiteContext

func TestQuotas(t *testing.T) {
	RegisterFailHandler(Fail)

	config := helpers.LoadConfig()
	context = helpers.NewContext(config)
	environment := helpers.NewEnvironment(context)

	BeforeSuite(func() {
		environment.Setup()
	})

	AfterSuite(func() {
		environment.Teardown()
	})

	RunSpecs(t, "Quotas")
}
<|endoftext|>


<commit_before>// Copyright 2012 Dorival de Moraes Pedroso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cec09

import (
	"testing"

	"github.com/cpmech/gosl/io"
	"github.com/cpmech/gosl/plt"
	"github.com/cpmech/gosl/utl"
	"github.com/cpmech/gosl/vtk"
)

func Test_data2d(tst *testing.T) {
	prob := "CF4"
	dat := PFdata(prob)
	X := utl.DblsGetColumn(0, dat)
	Y := utl.DblsGetColumn(1, dat)
	plt.SetForEps(1.0, 250, nil)
	plt.Plot(X, Y, nil)
	plt.Gll("$f_1$", "$f_2$", nil)
	plt.SaveD("/tmp/goga", io.Sf("cec09-%s.eps", prob))
}

func Test_data3d(tst *testing.T) {

	// data
	prob := "CF9"
	dat := PFdata(prob)
	X := utl.DblsGetColumn(0, dat)
	Y := utl.DblsGetColumn(1, dat)
	Z := utl.DblsGetColumn(2, dat)

	// figure
	plt.SetForEps(1.0, 400, nil)
	plt.Plot3dPoints(X, Y, Z, true, nil)
	plt.AxisRange3d(0, 1, 0, 1, 0, 1)
	plt.Camera(10, -135, nil)
	//plt.Camera(10, 45, nil)
	plt.SaveD("/tmp/goga", io.Sf("cec09-%s.eps", prob))

	// interactive
	if false {
		r := 0.005
		scn := vtk.NewScene()
		P := vtk.Spheres{X: X, Y: Y, Z: Z, R: utl.DblVals(len(X), r), Color: []float64{1, 0, 0, 1}}
		P.AddTo(scn)
		scn.Run()
	}
}
<commit_msg>Fix example after update on plt<commit_after>// Copyright 2012 Dorival de Moraes Pedroso. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package cec09

import (
	"testing"

	"github.com/cpmech/gosl/io"
	"github.com/cpmech/gosl/plt"
	"github.com/cpmech/gosl/utl"
	"github.com/cpmech/gosl/vtk"
)

func Test_data2d(tst *testing.T) {
	prob := "CF4"
	dat := PFdata(prob)
	X := utl.DblsGetColumn(0, dat)
	Y := utl.DblsGetColumn(1, dat)
	plt.Reset(false, nil)
	plt.Plot(X, Y, nil)
	plt.Gll("$f_1$", "$f_2$", nil)
	plt.Save("/tmp/goga", io.Sf("cec09-%s", prob))
}

func Test_data3d(tst *testing.T) {

	// data
	prob := "CF9"
	dat := PFdata(prob)
	X := utl.DblsGetColumn(0, dat)
	Y := utl.DblsGetColumn(1, dat)
	Z := utl.DblsGetColumn(2, dat)

	// figure
	plt.Reset(false, nil)
	plt.Plot3dPoints(X, Y, Z, true, nil)
	plt.AxisRange3d(0, 1, 0, 1, 0, 1)
	plt.Camera(10, -135, nil)
	//plt.Camera(10, 45, nil)
	plt.Save("/tmp/goga", io.Sf("cec09-%s", prob))

	// interactive
	if false {
		r := 0.005
		scn := vtk.NewScene()
		P := vtk.Spheres{X: X, Y: Y, Z: Z, R: utl.DblVals(len(X), r), Color: []float64{1, 0, 0, 1}}
		P.AddTo(scn)
		scn.Run()
	}
}
<|endoftext|>


<commit_before>package iamy

import (
	"fmt"
	"log"
	"regexp"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/pkg/errors"
)

var cfnResourceRegexp = regexp.MustCompile(`-[A-Z0-9]{10,20}$`)

// AwsFetcher fetches account data from AWS
type AwsFetcher struct {
	// As Policy and Role descriptions are immutable, we can skip fetching them
	// when pushing to AWS
	SkipFetchingPolicyAndRoleDescriptions bool

	Debug *log.Logger

	iam     *iamClient
	s3      *s3Client
	account *Account
	data    AccountData

	descriptionFetchWaitGroup sync.WaitGroup
	descriptionFetchError     error
}

func (a *AwsFetcher) init() error {
	var err error

	s := awsSession()
	a.iam = newIamClient(s)
	a.s3 = newS3Client(s)
	if a.account, err = a.getAccount(); err != nil {
		return err
	}
	a.data = AccountData{
		Account: a.account,
	}

	return nil
}

// Fetch queries AWS for account data
func (a *AwsFetcher) Fetch() (*AccountData, error) {
	if err := a.init(); err != nil {
		return nil, errors.Wrap(err, "Error in init")
	}

	var wg sync.WaitGroup
	var iamErr, s3Err error

	log.Println("Fetching IAM data")
	wg.Add(1)
	go func() {
		defer wg.Done()
		iamErr = a.fetchIamData()
	}()

	log.Println("Fetching S3 data")
	wg.Add(1)
	go func() {
		defer wg.Done()
		s3Err = a.fetchS3Data()
	}()

	wg.Wait()

	if iamErr != nil {
		return nil, errors.Wrap(iamErr, "Error fetching IAM data")
	}
	if s3Err != nil {
		return nil, errors.Wrap(s3Err, "Error fetching S3 data")
	}

	return &a.data, nil
}

func (a *AwsFetcher) fetchS3Data() error {
	buckets, err := a.s3.listAllBuckets()
	if err != nil {
		return errors.Wrap(err, "Error listing buckets")
	}
	for _, b := range buckets {
		if b.policyJson == "" {
			continue
		}

		policyDoc, err := NewPolicyDocumentFromEncodedJson(b.policyJson)
		if err != nil {
			return errors.Wrap(err, "Error creating Policy document")
		}

		bp := BucketPolicy{
			BucketName: b.name,
			Policy:     policyDoc,
		}

		a.data.BucketPolicies = append(a.data.BucketPolicies, &bp)
	}

	return nil
}
func (a *AwsFetcher) fetchIamData() error {
	var populateIamDataErr error
	var populateInstanceProfileErr error
	err := a.iam.GetAccountAuthorizationDetailsPages(
		&iam.GetAccountAuthorizationDetailsInput{
			Filter: aws.StringSlice([]string{
				iam.EntityTypeUser,
				iam.EntityTypeGroup,
				iam.EntityTypeRole,
				iam.EntityTypeLocalManagedPolicy,
			}),
		},
		func(resp *iam.GetAccountAuthorizationDetailsOutput, lastPage bool) bool {
			populateIamDataErr = a.populateIamData(resp)
			if populateIamDataErr != nil {
				return false
			}
			return true
		},
	)
	if populateIamDataErr != nil {
		return err
	}
	if err != nil {
		return err
	}
	// Fetch instance profiles
	err = a.iam.ListInstanceProfilesPages(&iam.ListInstanceProfilesInput{},
		func(resp *iam.ListInstanceProfilesOutput, lastPage bool) bool {
			populateInstanceProfileErr = a.populateInstanceProfileData(resp)
			if populateInstanceProfileErr != nil {
				return false
			}
			return true
		})
	if populateInstanceProfileErr != nil {
		return err
	}
	if err != nil {
		return err
	}
	return nil
}

func (a *AwsFetcher) populateInlinePolicies(source []*iam.PolicyDetail, target *[]InlinePolicy) error {
	for _, ip := range source {
		doc, err := NewPolicyDocumentFromEncodedJson(*ip.PolicyDocument)
		if err != nil {
			return err
		}
		*target = append(*target, InlinePolicy{
			Name:   *ip.PolicyName,
			Policy: doc,
		})
	}

	return nil
}

func (a *AwsFetcher) marshalPolicyDescriptionAsync(policyArn string, target *string) {
	a.descriptionFetchWaitGroup.Add(1)
	go func() {
		defer a.descriptionFetchWaitGroup.Done()
		log.Println("Fetching policy description for", policyArn)

		var err error
		*target, err = a.iam.getPolicyDescription(policyArn)
		if err != nil {
			a.descriptionFetchError = err
		}
	}()
}

func (a *AwsFetcher) marshalRoleDescriptionAsync(roleName string, target *string) {
	a.descriptionFetchWaitGroup.Add(1)
	go func() {
		defer a.descriptionFetchWaitGroup.Done()
		log.Println("Fetching role description for", roleName)

		var err error
		*target, err = a.iam.getRoleDescription(roleName)
		if err != nil {
			a.descriptionFetchError = err
		}
	}()
}

func (a *AwsFetcher) populateInstanceProfileData(resp *iam.ListInstanceProfilesOutput) error {
	for _, profileResp := range resp.InstanceProfiles {
		if ok, err := isSkippableManagedResource(*profileResp.InstanceProfileName); ok {
			log.Printf(err)
			continue
		}

		profile := InstanceProfile{iamService: iamService{
			Name: *profileResp.InstanceProfileName,
			Path: *profileResp.Path,
		}}
		for _, roleResp := range profileResp.Roles {
			role := *(roleResp.RoleName)
			profile.Roles = append(profile.Roles, role)
		}
		a.data.InstanceProfiles = append(a.data.InstanceProfiles, &profile)
	}
	return nil
}

func (a *AwsFetcher) populateIamData(resp *iam.GetAccountAuthorizationDetailsOutput) error {
	for _, userResp := range resp.UserDetailList {
		if ok, err := isSkippableManagedResource(*userResp.UserName); ok {
			log.Printf(err)
			continue
		}

		user := User{
			iamService: iamService{
				Name: *userResp.UserName,
				Path: *userResp.Path,
			},
			Tags: make(map[string]string),
		}

		for _, g := range userResp.GroupList {
			user.Groups = append(user.Groups, *g)
		}
		for _, p := range userResp.AttachedManagedPolicies {
			user.Policies = append(user.Policies, a.account.normalisePolicyArn(*p.PolicyArn))
		}
		if err := a.populateInlinePolicies(userResp.UserPolicyList, &user.InlinePolicies); err != nil {
			return err
		}
		for _, t := range userResp.Tags {
			user.Tags[*t.Key] = *t.Value
		}

		a.data.Users = append(a.data.Users, &user)
	}

	for _, groupResp := range resp.GroupDetailList {
		if ok, err := isSkippableManagedResource(*groupResp.GroupName); ok {
			log.Printf(err)
			continue
		}

		group := Group{iamService: iamService{
			Name: *groupResp.GroupName,
			Path: *groupResp.Path,
		}}

		for _, p := range groupResp.AttachedManagedPolicies {
			group.Policies = append(group.Policies, a.account.normalisePolicyArn(*p.PolicyArn))
		}
		if err := a.populateInlinePolicies(groupResp.GroupPolicyList, &group.InlinePolicies); err != nil {
			return err
		}

		a.data.Groups = append(a.data.Groups, &group)
	}

	for _, roleResp := range resp.RoleDetailList {
		if ok, err := isSkippableManagedResource(*roleResp.RoleName); ok {
			log.Printf(err)
			continue
		}

		role := Role{iamService: iamService{
			Name: *roleResp.RoleName,
			Path: *roleResp.Path,
		}}

		if !a.SkipFetchingPolicyAndRoleDescriptions {
			a.marshalRoleDescriptionAsync(*roleResp.RoleName, &role.Description)
		}

		var err error
		role.AssumeRolePolicyDocument, err = NewPolicyDocumentFromEncodedJson(*roleResp.AssumeRolePolicyDocument)
		if err != nil {
			return err
		}
		for _, p := range roleResp.AttachedManagedPolicies {
			role.Policies = append(role.Policies, a.account.normalisePolicyArn(*p.PolicyArn))
		}
		if err := a.populateInlinePolicies(roleResp.RolePolicyList, &role.InlinePolicies); err != nil {
			return err
		}

		a.data.addRole(&role)
	}

	for _, policyResp := range resp.Policies {
		if ok, err := isSkippableManagedResource(*policyResp.PolicyName); ok {
			log.Printf(err)
			continue
		}

		defaultPolicyVersion := findDefaultPolicyVersion(policyResp.PolicyVersionList)
		doc, err := NewPolicyDocumentFromEncodedJson(*defaultPolicyVersion.Document)
		if err != nil {
			return err
		}

		p := Policy{
			iamService: iamService{
				Name: *policyResp.PolicyName,
				Path: *policyResp.Path,
			},
			oldestVersionId:      findOldestPolicyVersionId(policyResp.PolicyVersionList),
			numberOfVersions:     len(policyResp.PolicyVersionList),
			nondefaultVersionIds: findNonDefaultPolicyVersionIds(policyResp.PolicyVersionList),
			Policy:               doc,
		}

		if !a.SkipFetchingPolicyAndRoleDescriptions {
			a.marshalPolicyDescriptionAsync(*policyResp.Arn, &p.Description)
		}

		a.data.addPolicy(&p)
	}

	a.descriptionFetchWaitGroup.Wait()

	return a.descriptionFetchError
}

func findDefaultPolicyVersion(versions []*iam.PolicyVersion) *iam.PolicyVersion {
	for _, version := range versions {
		if *version.IsDefaultVersion {
			return version
		}
	}
	panic("Expected a default policy version")
}

func findNonDefaultPolicyVersionIds(versions []*iam.PolicyVersion) []string {
	ss := []string{}
	for _, version := range versions {
		if !*version.IsDefaultVersion {
			ss = append(ss, *version.VersionId)
		}
	}
	return ss
}

func findOldestPolicyVersionId(versions []*iam.PolicyVersion) string {
	oldest := versions[0]
	for _, version := range versions[1:] {
		if version.CreateDate.Before(*oldest.CreateDate) {
			oldest = version
		}
	}
	return *oldest.VersionId
}

func (a *AwsFetcher) getAccount() (*Account, error) {
	var err error
	acct := Account{}

	acct.Id, err = GetAwsAccountId(awsSession(), a.Debug)
	if err == aws.ErrMissingRegion {
		return nil, errors.New("Error determining the AWS account id - check the AWS_REGION environment variable is set")
	}
	if err != nil {
		return nil, err
	}

	aliasResp, err := a.iam.ListAccountAliases(&iam.ListAccountAliasesInput{})
	if err != nil {
		return nil, err
	}
	if len(aliasResp.AccountAliases) > 0 {
		acct.Alias = *aliasResp.AccountAliases[0]
	}

	return &acct, nil
}

// isSkippableResource takes the resource identifier as a string and
// checks it against known resources that we shouldn't need to manage as
// it will already be managed by another process (such as Cloudformation
// roles).
//
// Returns a boolean of whether it can be skipped and a string of the
// reasoning why it was skipped.
func isSkippableManagedResource(resourceIdentifier string) (bool, string) {
	if cfnResourceRegexp.MatchString(resourceIdentifier) {
		return true, fmt.Sprintf("CloudFormation generated resource %s", resourceIdentifier)
	}

	if strings.Contains(resourceIdentifier, "AWSServiceRole") || strings.Contains(resourceIdentifier, "aws-service-role") {
		return true, fmt.Sprintf("AWS Service role generated resource %s", resourceIdentifier)
	}

	return false, ""
}
<commit_msg>Print errors<commit_after>package iamy

import (
	"fmt"
	"log"
	"regexp"
	"strings"
	"sync"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/service/iam"
	"github.com/pkg/errors"
)

var cfnResourceRegexp = regexp.MustCompile(`-[A-Z0-9]{10,20}$`)

// AwsFetcher fetches account data from AWS
type AwsFetcher struct {
	// As Policy and Role descriptions are immutable, we can skip fetching them
	// when pushing to AWS
	SkipFetchingPolicyAndRoleDescriptions bool

	Debug *log.Logger

	iam     *iamClient
	s3      *s3Client
	account *Account
	data    AccountData

	descriptionFetchWaitGroup sync.WaitGroup
	descriptionFetchError     error
}

func (a *AwsFetcher) init() error {
	var err error

	s := awsSession()
	a.iam = newIamClient(s)
	a.s3 = newS3Client(s)
	if a.account, err = a.getAccount(); err != nil {
		return err
	}
	a.data = AccountData{
		Account: a.account,
	}

	return nil
}

// Fetch queries AWS for account data
func (a *AwsFetcher) Fetch() (*AccountData, error) {
	if err := a.init(); err != nil {
		return nil, errors.Wrap(err, "Error in init")
	}

	var wg sync.WaitGroup
	var iamErr, s3Err error

	log.Println("Fetching IAM data")
	wg.Add(1)
	go func() {
		defer wg.Done()
		iamErr = a.fetchIamData()
	}()

	log.Println("Fetching S3 data")
	wg.Add(1)
	go func() {
		defer wg.Done()
		s3Err = a.fetchS3Data()
	}()

	wg.Wait()

	if iamErr != nil {
		return nil, errors.Wrap(iamErr, "Error fetching IAM data")
	}
	if s3Err != nil {
		return nil, errors.Wrap(s3Err, "Error fetching S3 data")
	}

	return &a.data, nil
}

func (a *AwsFetcher) fetchS3Data() error {
	buckets, err := a.s3.listAllBuckets()
	if err != nil {
		return errors.Wrap(err, "Error listing buckets")
	}
	for _, b := range buckets {
		if b.policyJson == "" {
			continue
		}

		policyDoc, err := NewPolicyDocumentFromEncodedJson(b.policyJson)
		if err != nil {
			return errors.Wrap(err, "Error creating Policy document")
		}

		bp := BucketPolicy{
			BucketName: b.name,
			Policy:     policyDoc,
		}

		a.data.BucketPolicies = append(a.data.BucketPolicies, &bp)
	}

	return nil
}
func (a *AwsFetcher) fetchIamData() error {
	var populateIamDataErr error
	var populateInstanceProfileErr error
	err := a.iam.GetAccountAuthorizationDetailsPages(
		&iam.GetAccountAuthorizationDetailsInput{
			Filter: aws.StringSlice([]string{
				iam.EntityTypeUser,
				iam.EntityTypeGroup,
				iam.EntityTypeRole,
				iam.EntityTypeLocalManagedPolicy,
			}),
		},
		func(resp *iam.GetAccountAuthorizationDetailsOutput, lastPage bool) bool {
			// There's a bug here. This doesn't set the variable outside this function scope
			// So IAMY just swallows any errors returned from populateIamData()
			populateIamDataErr = a.populateIamData(resp)
			if populateIamDataErr != nil {
				log.Println("Error Fetching IAM Data")
				return false
			}
			return true
		},
	)
	if populateIamDataErr != nil {
		return err
	}
	if err != nil {
		return err
	}
	// Fetch instance profiles
	err = a.iam.ListInstanceProfilesPages(&iam.ListInstanceProfilesInput{},
		func(resp *iam.ListInstanceProfilesOutput, lastPage bool) bool {
			populateInstanceProfileErr = a.populateInstanceProfileData(resp)
			if populateInstanceProfileErr != nil {
				return false
			}
			return true
		})
	if populateInstanceProfileErr != nil {
		return err
	}
	if err != nil {
		return err
	}
	return nil
}

func (a *AwsFetcher) populateInlinePolicies(source []*iam.PolicyDetail, target *[]InlinePolicy) error {
	for _, ip := range source {
		doc, err := NewPolicyDocumentFromEncodedJson(*ip.PolicyDocument)
		if err != nil {
			return err
		}
		*target = append(*target, InlinePolicy{
			Name:   *ip.PolicyName,
			Policy: doc,
		})
	}

	return nil
}

func (a *AwsFetcher) marshalPolicyDescriptionAsync(policyArn string, target *string) {
	a.descriptionFetchWaitGroup.Add(1)
	go func() {
		defer a.descriptionFetchWaitGroup.Done()
		log.Println("Fetching policy description for", policyArn)

		var err error
		*target, err = a.iam.getPolicyDescription(policyArn)
		if err != nil {
			a.descriptionFetchError = err
		}
	}()
}

func (a *AwsFetcher) marshalRoleDescriptionAsync(roleName string, target *string) {
	a.descriptionFetchWaitGroup.Add(1)
	go func() {
		defer a.descriptionFetchWaitGroup.Done()
		log.Println("Fetching role description for", roleName)

		var err error
		*target, err = a.iam.getRoleDescription(roleName)
		if err != nil {
			a.descriptionFetchError = err
		}
	}()
}

func (a *AwsFetcher) populateInstanceProfileData(resp *iam.ListInstanceProfilesOutput) error {
	for _, profileResp := range resp.InstanceProfiles {
		if ok, err := isSkippableManagedResource(*profileResp.InstanceProfileName); ok {
			log.Printf(err)
			continue
		}

		profile := InstanceProfile{iamService: iamService{
			Name: *profileResp.InstanceProfileName,
			Path: *profileResp.Path,
		}}
		for _, roleResp := range profileResp.Roles {
			role := *(roleResp.RoleName)
			profile.Roles = append(profile.Roles, role)
		}
		a.data.InstanceProfiles = append(a.data.InstanceProfiles, &profile)
	}
	return nil
}

func (a *AwsFetcher) populateIamData(resp *iam.GetAccountAuthorizationDetailsOutput) error {
	for _, userResp := range resp.UserDetailList {
		if ok, err := isSkippableManagedResource(*userResp.UserName); ok {
			log.Printf(err)
			continue
		}

		user := User{
			iamService: iamService{
				Name: *userResp.UserName,
				Path: *userResp.Path,
			},
			Tags: make(map[string]string),
		}

		for _, g := range userResp.GroupList {
			user.Groups = append(user.Groups, *g)
		}
		for _, p := range userResp.AttachedManagedPolicies {
			user.Policies = append(user.Policies, a.account.normalisePolicyArn(*p.PolicyArn))
		}
		if err := a.populateInlinePolicies(userResp.UserPolicyList, &user.InlinePolicies); err != nil {
			return err
		}
		for _, t := range userResp.Tags {
			user.Tags[*t.Key] = *t.Value
		}

		a.data.Users = append(a.data.Users, &user)
	}

	for _, groupResp := range resp.GroupDetailList {
		if ok, err := isSkippableManagedResource(*groupResp.GroupName); ok {
			log.Printf(err)
			continue
		}

		group := Group{iamService: iamService{
			Name: *groupResp.GroupName,
			Path: *groupResp.Path,
		}}

		for _, p := range groupResp.AttachedManagedPolicies {
			group.Policies = append(group.Policies, a.account.normalisePolicyArn(*p.PolicyArn))
		}
		if err := a.populateInlinePolicies(groupResp.GroupPolicyList, &group.InlinePolicies); err != nil {
			return err
		}

		a.data.Groups = append(a.data.Groups, &group)
	}

	for _, roleResp := range resp.RoleDetailList {
		if ok, err := isSkippableManagedResource(*roleResp.RoleName); ok {
			log.Printf(err)
			continue
		}

		role := Role{iamService: iamService{
			Name: *roleResp.RoleName,
			Path: *roleResp.Path,
		}}

		if !a.SkipFetchingPolicyAndRoleDescriptions {
			a.marshalRoleDescriptionAsync(*roleResp.RoleName, &role.Description)
		}

		var err error
		role.AssumeRolePolicyDocument, err = NewPolicyDocumentFromEncodedJson(*roleResp.AssumeRolePolicyDocument)
		if err != nil {
			return err
		}
		for _, p := range roleResp.AttachedManagedPolicies {
			role.Policies = append(role.Policies, a.account.normalisePolicyArn(*p.PolicyArn))
		}
		if err := a.populateInlinePolicies(roleResp.RolePolicyList, &role.InlinePolicies); err != nil {
			return err
		}

		a.data.addRole(&role)
	}

	for _, policyResp := range resp.Policies {
		if ok, err := isSkippableManagedResource(*policyResp.PolicyName); ok {
			log.Printf(err)
			continue
		}

		defaultPolicyVersion := findDefaultPolicyVersion(policyResp.PolicyVersionList)
		doc, err := NewPolicyDocumentFromEncodedJson(*defaultPolicyVersion.Document)
		if err != nil {
			return err
		}

		p := Policy{
			iamService: iamService{
				Name: *policyResp.PolicyName,
				Path: *policyResp.Path,
			},
			oldestVersionId:      findOldestPolicyVersionId(policyResp.PolicyVersionList),
			numberOfVersions:     len(policyResp.PolicyVersionList),
			nondefaultVersionIds: findNonDefaultPolicyVersionIds(policyResp.PolicyVersionList),
			Policy:               doc,
		}

		if !a.SkipFetchingPolicyAndRoleDescriptions {
			a.marshalPolicyDescriptionAsync(*policyResp.Arn, &p.Description)
		}

		a.data.addPolicy(&p)
	}

	a.descriptionFetchWaitGroup.Wait()

	return a.descriptionFetchError
}

func findDefaultPolicyVersion(versions []*iam.PolicyVersion) *iam.PolicyVersion {
	for _, version := range versions {
		if *version.IsDefaultVersion {
			return version
		}
	}
	panic("Expected a default policy version")
}

func findNonDefaultPolicyVersionIds(versions []*iam.PolicyVersion) []string {
	ss := []string{}
	for _, version := range versions {
		if !*version.IsDefaultVersion {
			ss = append(ss, *version.VersionId)
		}
	}
	return ss
}

func findOldestPolicyVersionId(versions []*iam.PolicyVersion) string {
	oldest := versions[0]
	for _, version := range versions[1:] {
		if version.CreateDate.Before(*oldest.CreateDate) {
			oldest = version
		}
	}
	return *oldest.VersionId
}

func (a *AwsFetcher) getAccount() (*Account, error) {
	var err error
	acct := Account{}

	acct.Id, err = GetAwsAccountId(awsSession(), a.Debug)
	if err == aws.ErrMissingRegion {
		return nil, errors.New("Error determining the AWS account id - check the AWS_REGION environment variable is set")
	}
	if err != nil {
		return nil, err
	}

	aliasResp, err := a.iam.ListAccountAliases(&iam.ListAccountAliasesInput{})
	if err != nil {
		return nil, err
	}
	if len(aliasResp.AccountAliases) > 0 {
		acct.Alias = *aliasResp.AccountAliases[0]
	}

	return &acct, nil
}

// isSkippableResource takes the resource identifier as a string and
// checks it against known resources that we shouldn't need to manage as
// it will already be managed by another process (such as Cloudformation
// roles).
//
// Returns a boolean of whether it can be skipped and a string of the
// reasoning why it was skipped.
func isSkippableManagedResource(resourceIdentifier string) (bool, string) {
	if cfnResourceRegexp.MatchString(resourceIdentifier) {
		return true, fmt.Sprintf("CloudFormation generated resource %s", resourceIdentifier)
	}

	if strings.Contains(resourceIdentifier, "AWSServiceRole") || strings.Contains(resourceIdentifier, "aws-service-role") {
		return true, fmt.Sprintf("AWS Service role generated resource %s", resourceIdentifier)
	}

	return false, ""
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"os"
	"testing"
)

func TestDiff(t *testing.T) {
	os.Args = append(os.Args, "diff", "bash")
	main()
}

func TestGroup(t *testing.T) {
	os.Args = append([]string{}, "build", "devel")
	fmt.Println(os.Args)
	main()
}
<commit_msg>test additional arguments<commit_after>package main

import (
	"os"
	"testing"
)

var (
	args = os.Args
)

func TestDiff(t *testing.T) {
	os.Args = append(args, "diff", "bash")
	main()
}

func TestGroup(t *testing.T) {
	os.Args = append(args, "build", "sed")
	main()
}

func TestList(t *testing.T) {
	os.Args = append(args, "list", "glibc")
	main()
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

type report struct {
	Container topology
	Plugins   []pluginSpec
}

type topology struct {
	Nodes             map[string]node             `json:"nodes"`
	Controls          map[string]control          `json:"controls"`
	MetadataTemplates map[string]metadataTemplate `json:"metadata_templates,omitempty"`
}

type metadataTemplate struct {
	ID       string  `json:"id"`
	Label    string  `json:"label,omitempty"`    // Human-readable descriptor for this row
	Truncate int     `json:"truncate,omitempty"` // If > 0, truncate the value to this length.
	Datatype string  `json:"dataType,omitempty"`
	Priority float64 `json:"priority,omitempty"`
	From     string  `json:"from,omitempty"` // Defines how to get the value from a report node
}

type node struct {
	LatestControls map[string]controlEntry `json:"latestControls,omitempty"`
	Latest         map[string]stringEntry  `json:"latest,omitempty"`
}

type controlEntry struct {
	Timestamp time.Time   `json:"timestamp"`
	Value     controlData `json:"value"`
}

type controlData struct {
	Dead bool `json:"dead"`
}

type control struct {
	ID    string `json:"id"`
	Human string `json:"human"`
	Icon  string `json:"icon"`
	Rank  int    `json:"rank"`
}

type stringEntry struct {
	Timestamp time.Time `json:"timestamp"`
	Value     string    `json:"value"`
}

type pluginSpec struct {
	ID          string   `json:"id"`
	Label       string   `json:"label"`
	Description string   `json:"description,omitempty"`
	Interfaces  []string `json:"interfaces"`
	APIVersion  string   `json:"api_version,omitempty"`
}

type Reporter struct {
	store *Store
}

func NewReporter(store *Store) *Reporter {
	return &Reporter{
		store: store,
	}
}

func (r *Reporter) RawReport() ([]byte, error) {
	rpt := &report{
		Container: topology{
			Nodes:             r.getContainerNodes(),
			Controls:          getTrafficControls(),
			MetadataTemplates: getMetadataTemplate(),
		},
		Plugins: []pluginSpec{
			{
				ID:          "traffic-control",
				Label:       "Traffic control",
				Description: "Adds traffic controls to the running Docker containers",
				Interfaces:  []string{"reporter", "controller"},
				APIVersion:  "1",
			},
		},
	}
	raw, err := json.Marshal(rpt)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal the report: %v", err)
	}
	return raw, nil
}

func (r *Reporter) GetHandler(nodeID, controlID string) (func() error, error) {
	containerID, err := nodeIDToContainerID(nodeID)
	if err != nil {
		return nil, fmt.Errorf("failed to get container ID from node ID %q: %v", nodeID)
	}
	container, found := r.store.Container(containerID)
	if !found {
		return nil, fmt.Errorf("container %s not found")
	}
	var handler func(pid int) error
	for _, c := range getControls() {
		if c.control.ID == controlID {
			handler = c.handler
			break
		}
	}
	if handler == nil {
		return nil, fmt.Errorf("unknown control ID %q for node ID %q", controlID, nodeID)
	}
	return func() error {
		return handler(container.PID)
	}, nil
}

// states:
// created, destroyed - don't create any node
// running, not running - create node with controls

func (r *Reporter) getContainerNodes() map[string]node {
	nodes := map[string]node{}
	timestamp := time.Now()
	r.store.ForEach(func(containerID string, container Container) {
		dead := false
		switch container.State {
		case Created, Destroyed:
			// do nothing, to prevent adding a stale node
			// to a report
		case Stopped:
			dead = true
			fallthrough
		case Running:
			nodeID := containerIDToNodeID(containerID)
			latency, _ := getLatency(container.PID)
			nodes[nodeID] = node{
				LatestControls: getTrafficNodeControls(timestamp, dead),
				Latest: map[string]stringEntry{
					"traffic-control-latency": {
						Timestamp: timestamp,
						Value:     latency,
					},
				},
			}
		}
	})
	return nodes
}

func getMetadataTemplate() map[string]metadataTemplate {
	return map[string]metadataTemplate{
		"traffic-control-latency": metadataTemplate{
			ID:       "traffic-control-latency",
			Label:    "Latency",
			Truncate: 0,
			Datatype: "",
			Priority: 13.5,
			From:     "latest",
		},
	}
}

func getTrafficNodeControls(timestamp time.Time, dead bool) map[string]controlEntry {
	controls := map[string]controlEntry{}
	entry := controlEntry{
		Timestamp: timestamp,
		Value: controlData{
			Dead: dead,
		},
	}
	for _, c := range getControls() {
		controls[c.control.ID] = entry
	}
	return controls
}

func getTrafficControls() map[string]control {
	controls := map[string]control{}
	for _, c := range getControls() {
		controls[c.control.ID] = c.control
	}
	return controls
}

type extControl struct {
	control control
	handler func(pid int) error
}

func getControls() []extControl {
	return []extControl{
		{
			control: control{
				ID:    "slow",
				Human: "Traffic speed: slow",
				Icon:  "fa-hourglass-1",
				Rank:  20,
			},
			handler: func(pid int) error {
				return DoTrafficControl(pid, "2000ms")
			},
		},
		{
			control: control{
				ID:    "medium",
				Human: "Traffic speed: medium",
				Icon:  "fa-hourglass-2",
				Rank:  21,
			},
			handler: func(pid int) error {
				return DoTrafficControl(pid, "300ms")
			},
		},
		{
			control: control{
				ID:    "fast",
				Human: "Traffic speed: fast",
				Icon:  "fa-hourglass-3",
				Rank:  22,
			},
			handler: func(pid int) error {
				return DoTrafficControl(pid, "1ms")
			},
		},
	}
}

const nodeSuffix = ";<container>"

func containerIDToNodeID(containerID string) string {
	return fmt.Sprintf("%s%s", containerID, nodeSuffix)
}

func nodeIDToContainerID(nodeID string) (string, error) {
	if !strings.HasSuffix(nodeID, nodeSuffix) {
		return "", fmt.Errorf("no suffix %q in node ID %q", nodeSuffix, nodeID)
	}
	return strings.TrimSuffix(nodeID, nodeSuffix), nil
}
<commit_msg>plugins/traffic-control: create Traffic control table<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"strings"
	"time"
)

type report struct {
	Container topology
	Plugins   []pluginSpec
}

type topology struct {
	Nodes             map[string]node             `json:"nodes"`
	Controls          map[string]control          `json:"controls"`
	MetadataTemplates map[string]metadataTemplate `json:"metadata_templates,omitempty"`
	TableTemplates    map[string]tableTemplate    `json:"table_templates,omitempty"`
}

type tableTemplate struct {
	ID     string `json:"id"`
	Label  string `json:"label"`
	Prefix string `json:"prefix"`
}

type metadataTemplate struct {
	ID       string  `json:"id"`
	Label    string  `json:"label,omitempty"`    // Human-readable descriptor for this row
	Truncate int     `json:"truncate,omitempty"` // If > 0, truncate the value to this length.
	Datatype string  `json:"dataType,omitempty"`
	Priority float64 `json:"priority,omitempty"`
	From     string  `json:"from,omitempty"` // Defines how to get the value from a report node
}

type node struct {
	LatestControls map[string]controlEntry `json:"latestControls,omitempty"`
	Latest         map[string]stringEntry  `json:"latest,omitempty"`
}

type controlEntry struct {
	Timestamp time.Time   `json:"timestamp"`
	Value     controlData `json:"value"`
}

type controlData struct {
	Dead bool `json:"dead"`
}

type control struct {
	ID    string `json:"id"`
	Human string `json:"human"`
	Icon  string `json:"icon"`
	Rank  int    `json:"rank"`
}

type stringEntry struct {
	Timestamp time.Time `json:"timestamp"`
	Value     string    `json:"value"`
}

type pluginSpec struct {
	ID          string   `json:"id"`
	Label       string   `json:"label"`
	Description string   `json:"description,omitempty"`
	Interfaces  []string `json:"interfaces"`
	APIVersion  string   `json:"api_version,omitempty"`
}

type Reporter struct {
	store *Store
}

func NewReporter(store *Store) *Reporter {
	return &Reporter{
		store: store,
	}
}

func (r *Reporter) RawReport() ([]byte, error) {
	rpt := &report{
		Container: topology{
			Nodes:             r.getContainerNodes(),
			Controls:          getTrafficControls(),
			MetadataTemplates: getMetadataTemplate(),
			TableTemplates:    getTableTemplate(),
		},
		Plugins: []pluginSpec{
			{
				ID:          "traffic-control",
				Label:       "Traffic control",
				Description: "Adds traffic controls to the running Docker containers",
				Interfaces:  []string{"reporter", "controller"},
				APIVersion:  "1",
			},
		},
	}
	raw, err := json.Marshal(rpt)
	if err != nil {
		return nil, fmt.Errorf("failed to marshal the report: %v", err)
	}
	return raw, nil
}

func (r *Reporter) GetHandler(nodeID, controlID string) (func() error, error) {
	containerID, err := nodeIDToContainerID(nodeID)
	if err != nil {
		return nil, fmt.Errorf("failed to get container ID from node ID %q: %v", nodeID)
	}
	container, found := r.store.Container(containerID)
	if !found {
		return nil, fmt.Errorf("container %s not found")
	}
	var handler func(pid int) error
	for _, c := range getControls() {
		if c.control.ID == controlID {
			handler = c.handler
			break
		}
	}
	if handler == nil {
		return nil, fmt.Errorf("unknown control ID %q for node ID %q", controlID, nodeID)
	}
	return func() error {
		return handler(container.PID)
	}, nil
}

// states:
// created, destroyed - don't create any node
// running, not running - create node with controls

func (r *Reporter) getContainerNodes() map[string]node {
	nodes := map[string]node{}
	timestamp := time.Now()
	r.store.ForEach(func(containerID string, container Container) {
		dead := false
		switch container.State {
		case Created, Destroyed:
			// do nothing, to prevent adding a stale node
			// to a report
		case Stopped:
			dead = true
			fallthrough
		case Running:
			nodeID := containerIDToNodeID(containerID)
			latency, _ := getLatency(container.PID)
			nodes[nodeID] = node{
				LatestControls: getTrafficNodeControls(timestamp, dead),
				Latest: map[string]stringEntry{
					"traffic-control-latency": {
						Timestamp: timestamp,
						Value:     latency,
					},
					"traffic-control-table-latency": {
						Timestamp: timestamp,
						Value:     latency,
					},
				},
			}
		}
	})
	return nodes
}

func getMetadataTemplate() map[string]metadataTemplate {
	return map[string]metadataTemplate{
		"traffic-control-latency": metadataTemplate{
			ID:       "traffic-control-latency",
			Label:    "Latency",
			Truncate: 0,
			Datatype: "",
			Priority: 13.5,
			From:     "latest",
		},
	}
}

func getTableTemplate() map[string]tableTemplate {
	return map[string]tableTemplate{
		"traffic-control-table": tableTemplate{
			ID:     "traffic-control-table",
			Label:  "Traffic Control",
			Prefix: "traffic-control-table-",
		},
	}
}

func getTrafficNodeControls(timestamp time.Time, dead bool) map[string]controlEntry {
	controls := map[string]controlEntry{}
	entry := controlEntry{
		Timestamp: timestamp,
		Value: controlData{
			Dead: dead,
		},
	}
	for _, c := range getControls() {
		controls[c.control.ID] = entry
	}
	return controls
}

func getTrafficControls() map[string]control {
	controls := map[string]control{}
	for _, c := range getControls() {
		controls[c.control.ID] = c.control
	}
	return controls
}

type extControl struct {
	control control
	handler func(pid int) error
}

func getControls() []extControl {
	return []extControl{
		{
			control: control{
				ID:    "slow",
				Human: "Traffic speed: slow",
				Icon:  "fa-hourglass-1",
				Rank:  20,
			},
			handler: func(pid int) error {
				return DoTrafficControl(pid, "2000ms")
			},
		},
		{
			control: control{
				ID:    "medium",
				Human: "Traffic speed: medium",
				Icon:  "fa-hourglass-2",
				Rank:  21,
			},
			handler: func(pid int) error {
				return DoTrafficControl(pid, "300ms")
			},
		},
		{
			control: control{
				ID:    "fast",
				Human: "Traffic speed: fast",
				Icon:  "fa-hourglass-3",
				Rank:  22,
			},
			handler: func(pid int) error {
				return DoTrafficControl(pid, "1ms")
			},
		},
	}
}

const nodeSuffix = ";<container>"

func containerIDToNodeID(containerID string) string {
	return fmt.Sprintf("%s%s", containerID, nodeSuffix)
}

func nodeIDToContainerID(nodeID string) (string, error) {
	if !strings.HasSuffix(nodeID, nodeSuffix) {
		return "", fmt.Errorf("no suffix %q in node ID %q", nodeSuffix, nodeID)
	}
	return strings.TrimSuffix(nodeID, nodeSuffix), nil
}
<|endoftext|>