<commit_before>package main

import(
	"fmt"
	"net/http"
	"log"
	"strings"
	"strconv"
	"html/template"

	"github.com/kyledayton/requesthub/templates"
)

var config *Config

func foundationCSShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_foundation_css_bytes()
	w.Header().Set("Content-Type", "text/css")
	w.Write(f)
}

func foundationJShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_foundation_js_bytes()
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(f)
}

func jqueryJShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_jquery_js_bytes()
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(f)
}

func modernizrJShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_modernizr_js_bytes()
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(f)
}

func authFailed(r *http.Request) bool {
	reqUser, reqPass, ok := r.BasicAuth()

	return config.AuthEnabled() && (!ok || reqUser != config.Username || reqPass != config.Password)
}

func requireAuth(w http.ResponseWriter) {
	w.Header().Add("WWW-Authenticate", "Basic realm=\"requesthub\"")
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte("Authorization Failed."))
}

func Start() {
	config = NewConfig()

	log.Printf("Initializing hub database (maxReq=%d)\n", config.MaxRequests)
	db := newHubDatabase(config.MaxRequests)

	if config.HasYAMLConfig() {
		err := config.ApplyYAMLConfig(db)
		if err != nil {
			log.Printf("Error loading config file (%s): %s", config.YamlConfigFile, err)
		}
	}

	if config.AuthEnabled() {
		log.Printf("Using HTTP Basic Auth\n")
	}

	viewPage := template.Must(template.New("show").Parse(templates.SHOW_HUB))
	indexPage := template.Must(template.New("index").Parse(templates.INDEX))

	forwardClient := new(http.Client)

	router := MakeRouter()

	router.HandleFunc(`/assets/foundation.css`, foundationCSShandler)
	router.HandleFunc(`/assets/foundation.js`, foundationJShandler)
	router.HandleFunc(`/assets/jquery.js`, jqueryJShandler)
	router.HandleFunc(`/assets/modernizr.js`, modernizrJShandler)

	router.HandleFunc(`/show/([\d\w\-_]+)`, func(w http.ResponseWriter, r *http.Request) {
		if authFailed(r) {
			requireAuth(w)
			return
		}

		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[2]
		hub := db.Get(hubName)

		if hub != nil {
			w.Header().Add("Content-Type", "text/html")

			if authFailed(r) {
				requireAuth(w)
				return
			}

			viewPage.Execute(w, hub)
		}
	})

	router.HandleFunc(`/([\w\d\-_]+)/forward`, func(w http.ResponseWriter, r *http.Request) {
		if authFailed(r) {
			requireAuth(w)
			return
		}

		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]

		if hubName != "" {
			hub := db.Get(hubName)
			dest := strings.TrimSpace( r.FormValue("url") )

			if hub != nil {
				hub.ForwardURL = dest
			}
		}
	})

	router.HandleFunc(`/([\w\d\-_]+)/latest`, func(w http.ResponseWriter, r *http.Request) {
		if authFailed(r) {
			requireAuth(w)
			return
		}

		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]

		if hubName != "" {
			hub := db.Get(hubName)

			if hub != nil {
				w.Write([]byte(strconv.Itoa(int(hub.Requests.Count))))
			}
		}
	})

	router.HandleFunc(`/([\w\d\-_]+)/delete`, func(w http.ResponseWriter, r *http.Request) {
		if authFailed(r) {
			requireAuth(w)
			return
		}

		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]

		if hubName != "" {
			hub := db.Get(hubName)

			if hub != nil {
				db.Delete(hub.Id)
				http.Redirect(w, r, "/", 302)
			}
		}
	})

	router.HandleFunc(`/([\w\d\-_]+)/requests`, func(w http.ResponseWriter, r *http.Request) {
		if authFailed(r) {
			requireAuth(w)
			return
		}

		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]

		if hubName != "" {
			hub := db.Get(hubName)

			if hub != nil {
				json, err := db.Get(hubName).Requests.ToJson()

				if err != nil {
					log.Panic(err)
				}

				w.Header().Add("Content-Type", "application/json")
				w.Write(json)
			}
		}
	})

	router.HandleFunc(`/([\d\w\-_]+)/clear`, func(w http.ResponseWriter, r *http.Request) {
		if authFailed(r) {
			requireAuth(w)
			return
		}

		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]
		hub := db.Get(hubName)

		if hub != nil {
			hub.Requests.Clear()
		}
	})

	router.HandleFunc(`/([\d\w\-_]+)`, func(w http.ResponseWriter, r *http.Request) {
		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]

		hub := db.Get(hubName)

		if hub != nil {
			req := hub.Requests.Insert(r)

			if hub.ForwardURL != "" {
				go req.Forward(forwardClient, hub.ForwardURL)
			}
		}
	})

	router.HandleFunc(`/`, func(w http.ResponseWriter, r *http.Request) {

		if authFailed(r) {
			requireAuth(w)
			return
		}

		if r.Method == "GET" {
			w.Header().Add("Content-Type", "text/html")
			indexPage.Execute(w, db.hubs)
		} else if r.Method == "POST" {
			hubName := strings.TrimSpace( r.FormValue("hub_name") )

			if hubName == "" {
				http.Redirect(w, r, "/", 302)
				return
			}

			hub, err := db.Create(hubName)

			if err != nil {
				http.Redirect(w, r, `/`, 302)
			} else {
				http.Redirect(w, r, fmt.Sprintf("/show/%s", hub.Id), 302)
			}
		}
	})

	log.Printf("Listening on port %d\n", config.Port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", config.Port), router))
}
<commit_msg>Support disabling the web UI<commit_after>package main

import(
	"fmt"
	"net/http"
	"log"
	"strings"
	"strconv"
	"html/template"

	"github.com/kyledayton/requesthub/templates"
)

var config *Config

func foundationCSShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_foundation_css_bytes()
	w.Header().Set("Content-Type", "text/css")
	w.Write(f)
}

func foundationJShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_foundation_js_bytes()
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(f)
}

func jqueryJShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_jquery_js_bytes()
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(f)
}

func modernizrJShandler(w http.ResponseWriter, r *http.Request) {
	f, _ := assets_modernizr_js_bytes()
	w.Header().Set("Content-Type", "application/javascript")
	w.Write(f)
}

func authFailed(r *http.Request) bool {
	reqUser, reqPass, ok := r.BasicAuth()

	return config.AuthEnabled() && (!ok || reqUser != config.Username || reqPass != config.Password)
}

func requireAuth(w http.ResponseWriter) {
	w.Header().Add("WWW-Authenticate", "Basic realm=\"requesthub\"")
	w.WriteHeader(http.StatusUnauthorized)
	w.Write([]byte("Authorization Failed."))
}

func Start() {
	config = NewConfig()

	log.Printf("Initializing hub database (maxReq=%d)\n", config.MaxRequests)
	db := newHubDatabase(config.MaxRequests)

	if config.HasYAMLConfig() {
		err := config.ApplyYAMLConfig(db)
		if err != nil {
			log.Printf("Error loading config file (%s): %s", config.YamlConfigFile, err)
		}
	}

	if config.AuthEnabled() {
		log.Printf("Using HTTP Basic Auth\n")
	}

	viewPage := template.Must(template.New("show").Parse(templates.SHOW_HUB))
	indexPage := template.Must(template.New("index").Parse(templates.INDEX))

	forwardClient := new(http.Client)

	router := MakeRouter()

	router.HandleFunc(`/([\d\w\-_]+)`, func(w http.ResponseWriter, r *http.Request) {
		parts := strings.Split(r.URL.Path, "/")
		hubName := parts[1]

		hub := db.Get(hubName)

		if hub != nil {
			req := hub.Requests.Insert(r)

			if hub.ForwardURL != "" {
				go req.Forward(forwardClient, hub.ForwardURL)
			}
		}
	})

	if !config.UIEnabled {
		log.Println("Starting without web UI")
	}

	if config.UIEnabled {
		router.HandleFunc(`/assets/foundation.css`, foundationCSShandler)
		router.HandleFunc(`/assets/foundation.js`, foundationJShandler)
		router.HandleFunc(`/assets/jquery.js`, jqueryJShandler)
		router.HandleFunc(`/assets/modernizr.js`, modernizrJShandler)

		router.HandleFunc(`/show/([\d\w\-_]+)`, func(w http.ResponseWriter, r *http.Request) {
			if authFailed(r) {
				requireAuth(w)
				return
			}

			parts := strings.Split(r.URL.Path, "/")
			hubName := parts[2]
			hub := db.Get(hubName)

			if hub != nil {
				w.Header().Add("Content-Type", "text/html")

				if authFailed(r) {
					requireAuth(w)
					return
				}

				viewPage.Execute(w, hub)
			}
		})

		router.HandleFunc(`/([\w\d\-_]+)/forward`, func(w http.ResponseWriter, r *http.Request) {
			if authFailed(r) {
				requireAuth(w)
				return
			}

			parts := strings.Split(r.URL.Path, "/")
			hubName := parts[1]

			if hubName != "" {
				hub := db.Get(hubName)
				dest := strings.TrimSpace( r.FormValue("url") )

				if hub != nil {
					hub.ForwardURL = dest
				}
			}
		})

		router.HandleFunc(`/([\w\d\-_]+)/latest`, func(w http.ResponseWriter, r *http.Request) {
			if authFailed(r) {
				requireAuth(w)
				return
			}

			parts := strings.Split(r.URL.Path, "/")
			hubName := parts[1]

			if hubName != "" {
				hub := db.Get(hubName)

				if hub != nil {
					w.Write([]byte(strconv.Itoa(int(hub.Requests.Count))))
				}
			}
		})

		router.HandleFunc(`/([\w\d\-_]+)/delete`, func(w http.ResponseWriter, r *http.Request) {
			if authFailed(r) {
				requireAuth(w)
				return
			}

			parts := strings.Split(r.URL.Path, "/")
			hubName := parts[1]

			if hubName != "" {
				hub := db.Get(hubName)

				if hub != nil {
					db.Delete(hub.Id)
					http.Redirect(w, r, "/", 302)
				}
			}
		})

		router.HandleFunc(`/([\w\d\-_]+)/requests`, func(w http.ResponseWriter, r *http.Request) {
			if authFailed(r) {
				requireAuth(w)
				return
			}

			parts := strings.Split(r.URL.Path, "/")
			hubName := parts[1]

			if hubName != "" {
				hub := db.Get(hubName)

				if hub != nil {
					json, err := db.Get(hubName).Requests.ToJson()

					if err != nil {
						log.Panic(err)
					}

					w.Header().Add("Content-Type", "application/json")
					w.Write(json)
				}
			}
		})

		router.HandleFunc(`/([\d\w\-_]+)/clear`, func(w http.ResponseWriter, r *http.Request) {
			if authFailed(r) {
				requireAuth(w)
				return
			}

			parts := strings.Split(r.URL.Path, "/")
			hubName := parts[1]
			hub := db.Get(hubName)

			if hub != nil {
				hub.Requests.Clear()
			}
		})

		router.HandleFunc(`/`, func(w http.ResponseWriter, r *http.Request) {

			if authFailed(r) {
				requireAuth(w)
				return
			}

			if r.Method == "GET" {
				w.Header().Add("Content-Type", "text/html")
				indexPage.Execute(w, db.hubs)
			} else if r.Method == "POST" {
				hubName := strings.TrimSpace( r.FormValue("hub_name") )

				if hubName == "" {
					http.Redirect(w, r, "/", 302)
					return
				}

				hub, err := db.Create(hubName)

				if err != nil {
					http.Redirect(w, r, `/`, 302)
				} else {
					http.Redirect(w, r, fmt.Sprintf("/show/%s", hub.Id), 302)
				}
			}
		})
	}

	log.Printf("Listening on port %d\n", config.Port)
	log.Fatal(http.ListenAndServe(fmt.Sprintf(":%d", config.Port), router))
}
<|endoftext|>


<commit_before>package lazycache

import (
       "fmt"
       "net/url"
       "strings"
       kitlog "github.com/go-kit/kit/log"
)

func AddMirror(serverAddr string) {

	url, err := url.Parse(serverAddr)
	fs, err := OpenHttpFS(*url)

	if err != nil {
		panic(fmt.Sprintf("Error opening HTTP FS Source: %s", err.Error()))
	}

	//serverAddr := fmt.Sprintf("%s:%d", config.ServerIp, config.ServerPort)

	// Reverse hostname
	splitHN := MungeHostname(fs.Uri.Host)
	root := fmt.Sprintf("/%s%s", strings.Join(splitHN, "/"), fs.Uri.Path)
	MakeRootNode(fs, root)

	RootMap[serverAddr] = root

	// fmt.Printf("Starting http handler at http://%s/\n", serverAddr)
	// fmt.Printf("Fs at http://%s%s\n", serverAddr, root )

	//http.ListenAndServe(serverAddr, nil)
}

func ConfigureImageStore(store_type string, bucket string, logger kitlog.Logger) {
	switch strings.ToLower(store_type) {
	case "", "none":
		 DefaultImageStore = NullImageStore{}
	case "google":
	   DefaultImageStore = CreateGoogleStore(bucket, logger)
	}
}
<commit_msg>Added version to path information.  Breaking change.<commit_after>package lazycache

import (
       "fmt"
       "net/url"
       "strings"
       kitlog "github.com/go-kit/kit/log"
)

const ApiVersion = "v1"

func AddMirror(serverAddr string) {

	url, err := url.Parse(serverAddr)
	fs, err := OpenHttpFS(*url)

	if err != nil {
		panic(fmt.Sprintf("Error opening HTTP FS Source: %s", err.Error()))
	}

	//serverAddr := fmt.Sprintf("%s:%d", config.ServerIp, config.ServerPort)

	// Reverse hostname
	splitHN := MungeHostname(fs.Uri.Host)
	root := fmt.Sprintf("/%s/%s%s", ApiVersion, strings.Join(splitHN, "/"), fs.Uri.Path)
	MakeRootNode(fs, root)

	RootMap[serverAddr] = root

	// fmt.Printf("Starting http handler at http://%s/\n", serverAddr)
	// fmt.Printf("Fs at http://%s%s\n", serverAddr, root )

	//http.ListenAndServe(serverAddr, nil)
}

func ConfigureImageStore(store_type string, bucket string, logger kitlog.Logger) {
	switch strings.ToLower(store_type) {
	case "", "none":
		 DefaultImageStore = NullImageStore{}
	case "google":
	   DefaultImageStore = CreateGoogleStore(bucket, logger)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"github.com/dimfeld/glog"
	"github.com/dimfeld/httptreemux"
	"net"
	"net/http"
	"strings"
)

type HookHandler func(http.ResponseWriter, *http.Request, map[string]string, *Hook)

func hookHandler(w http.ResponseWriter, r *http.Request, params map[string]string, hook *Hook) {
	githubEventType := r.Header.Get("X-GitHub-Event")

	if r.ContentLength > 16384 {
		// We should never get a request this large.
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		return
	}

	buffer := bytes.Buffer{}
	buffer.ReadFrom(r.Body)
	r.Body.Close()

	if glog.V(2) {
		niceBuffer := &bytes.Buffer{}
		json.Indent(niceBuffer, buffer.Bytes(), "", "  ")
		glog.Infof("Hook %s received data %s\n",
			r.URL.Path, string(niceBuffer.Bytes()))
	}

	if hook.Secret != "" {
		secret := r.Header.Get("X-Hub-Signature")
		if !strings.HasPrefix(secret, "sha1=") {
			glog.Warningf("Request with no secret for hook %s from %s\n",
				r.URL.Path, r.RemoteAddr)
			w.WriteHeader(http.StatusForbidden)
			return
		}

		hash := hmac.New(sha1.New, []byte(hook.Secret))
		hash.Write(buffer.Bytes())
		expected := hash.Sum(nil)
		seen, err := hex.DecodeString(secret[5:])
		if err != nil || !hmac.Equal(expected, seen) {
			glog.Warningf("Request with bad secret for hook %s from %s\nExpected %s, saw %s",
				r.URL.Path, r.RemoteAddr, hex.EncodeToString(expected), secret)
			w.WriteHeader(http.StatusForbidden)
			return
		}
	}

	event, err := NewEvent(buffer.Bytes(), githubEventType)
	if err != nil {
		glog.Errorf("Error parinsg JSON for %s: %s", r.URL.Path, err)
	}
	event["urlparams"] = params
	go hook.Execute(event)
}

func handlerWrapper(handler HookHandler, hook *Hook) httptreemux.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, params map[string]string) {
		glog.Infoln("Called", r.URL.Path)
		handler(w, r, params, hook)
	}
}

func SetupServer(config *Config) (net.Listener, http.Handler) {
	var listener net.Listener = nil

	listener, err := net.Listen("tcp", config.ListenAddress)
	if err != nil {
		glog.Fatal("Could not listen on", config.ListenAddress)
	}

	if len(config.AcceptIps) != 0 {
		listenFilter := NewListenFilter(listener, WhiteList)
		for _, a := range config.AcceptIps {
			glog.Infoln("Adding IP filter", a)
			listenFilter.AddString(a)
		}
		listener = listenFilter
	}

	router := httptreemux.New()

	for _, hook := range config.Hook {
		router.POST(hook.Url, handlerWrapper(hookHandler, hook))
	}

	return listener, router
}

func RunServer(config *Config) {
	listener, router := SetupServer(config)
	glog.Fatal(http.Serve(listener, router))
}
<commit_msg>Don't panic or call the hook when JSON parsing fails<commit_after>package main

import (
	"bytes"
	"crypto/hmac"
	"crypto/sha1"
	"encoding/hex"
	"encoding/json"
	"github.com/dimfeld/glog"
	"github.com/dimfeld/httptreemux"
	"net"
	"net/http"
	"strings"
)

type HookHandler func(http.ResponseWriter, *http.Request, map[string]string, *Hook)

func hookHandler(w http.ResponseWriter, r *http.Request, params map[string]string, hook *Hook) {
	githubEventType := r.Header.Get("X-GitHub-Event")

	if r.ContentLength > 16384 {
		// We should never get a request this large.
		w.WriteHeader(http.StatusRequestEntityTooLarge)
		return
	}

	buffer := bytes.Buffer{}
	buffer.ReadFrom(r.Body)
	r.Body.Close()

	if glog.V(2) {
		niceBuffer := &bytes.Buffer{}
		json.Indent(niceBuffer, buffer.Bytes(), "", "  ")
		glog.Infof("Hook %s received data %s\n",
			r.URL.Path, string(niceBuffer.Bytes()))
	}

	if hook.Secret != "" {
		secret := r.Header.Get("X-Hub-Signature")
		if !strings.HasPrefix(secret, "sha1=") {
			glog.Warningf("Request with no secret for hook %s from %s\n",
				r.URL.Path, r.RemoteAddr)
			w.WriteHeader(http.StatusForbidden)
			return
		}

		hash := hmac.New(sha1.New, []byte(hook.Secret))
		hash.Write(buffer.Bytes())
		expected := hash.Sum(nil)
		seen, err := hex.DecodeString(secret[5:])
		if err != nil || !hmac.Equal(expected, seen) {
			glog.Warningf("Request with bad secret for hook %s from %s\nExpected %s, saw %s",
				r.URL.Path, r.RemoteAddr, hex.EncodeToString(expected), secret)
			w.WriteHeader(http.StatusForbidden)
			return
		}
	}

	event, err := NewEvent(buffer.Bytes(), githubEventType)
	if err != nil {
		glog.Errorf("Error parinsg JSON for %s: %s", r.URL.Path, err)
		return
	}
	event["urlparams"] = params
	go hook.Execute(event)
}

func handlerWrapper(handler HookHandler, hook *Hook) httptreemux.HandlerFunc {
	return func(w http.ResponseWriter, r *http.Request, params map[string]string) {
		glog.Infoln("Called", r.URL.Path)
		handler(w, r, params, hook)
	}
}

func SetupServer(config *Config) (net.Listener, http.Handler) {
	var listener net.Listener = nil

	listener, err := net.Listen("tcp", config.ListenAddress)
	if err != nil {
		glog.Fatal("Could not listen on", config.ListenAddress)
	}

	if len(config.AcceptIps) != 0 {
		listenFilter := NewListenFilter(listener, WhiteList)
		for _, a := range config.AcceptIps {
			glog.Infoln("Adding IP filter", a)
			listenFilter.AddString(a)
		}
		listener = listenFilter
	}

	router := httptreemux.New()

	for _, hook := range config.Hook {
		router.POST(hook.Url, handlerWrapper(hookHandler, hook))
	}

	return listener, router
}

func RunServer(config *Config) {
	listener, router := SetupServer(config)
	glog.Fatal(http.Serve(listener, router))
}
<|endoftext|>


<commit_before>package pixel

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"log/syslog"
	"net/http"
	"time"
)

type Server struct {
	listenAddress  string
	httpServer     *http.Server
	syslogAddress  string
	syslogPriority syslog.Priority
	logger         *log.Logger
	now            func() time.Time
}

// Every pixel GET or JSON POST is an event.
type Event struct {
	Time string `json:"t"`

	// Params are client-supplied event attributes.
	// Pixel events have params of form map[string]string.
	// JSON POST events have arbitrary JSON.
	Params interface{} `json:"params,omitempty"`

	UserAgent string `json:"ua,omitempty"`
	IP        string `json:"ip,omitempty"`
	Proto     string `json:"proto,omitempty"`

	// Error is a server-defined string reporting any error that
	// occurred while constructing the event, such as parse failures.
	Error string `json:"error,omitempty"`
}

const ISO8601Format = "2006-01-02T15:04:05Z"

var Transparent1PxGIF = []byte("\x47\x49\x46\x38\x39\x61\x01\x00" +
	"\x01\x00\x80\x00\x00\xff\xff\xff\x00\x00\x00\x2c\x00\x00\x00\x00" +
	"\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b")

const UDPMaxBytes = 65507

var BadRequest = http.StatusText(http.StatusBadRequest)

const PostBodyMaxLen = 2048

var levelMap = map[string]syslog.Priority{
	"LOG_EMERG":   syslog.LOG_EMERG,
	"LOG_ALERT":   syslog.LOG_ALERT,
	"LOG_CRIT":    syslog.LOG_CRIT,
	"LOG_ERR":     syslog.LOG_ERR,
	"LOG_WARNING": syslog.LOG_WARNING,
	"LOG_NOTICE":  syslog.LOG_NOTICE,
	"LOG_INFO":    syslog.LOG_INFO,
	"LOG_DEBUG":   syslog.LOG_DEBUG,
}

var facilityMap = map[string]syslog.Priority{
	"LOG_KERN":     syslog.LOG_KERN,
	"LOG_USER":     syslog.LOG_USER,
	"LOG_MAIL":     syslog.LOG_MAIL,
	"LOG_DAEMON":   syslog.LOG_DAEMON,
	"LOG_AUTH":     syslog.LOG_AUTH,
	"LOG_SYSLOG":   syslog.LOG_SYSLOG,
	"LOG_LPR":      syslog.LOG_LPR,
	"LOG_NEWS":     syslog.LOG_NEWS,
	"LOG_UUCP":     syslog.LOG_UUCP,
	"LOG_CRON":     syslog.LOG_CRON,
	"LOG_AUTHPRIV": syslog.LOG_AUTHPRIV,
	"LOG_FTP":      syslog.LOG_FTP,
	"LOG_LOCAL0":   syslog.LOG_LOCAL0,
	"LOG_LOCAL1":   syslog.LOG_LOCAL1,
	"LOG_LOCAL2":   syslog.LOG_LOCAL2,
	"LOG_LOCAL3":   syslog.LOG_LOCAL3,
	"LOG_LOCAL4":   syslog.LOG_LOCAL4,
	"LOG_LOCAL5":   syslog.LOG_LOCAL5,
	"LOG_LOCAL6":   syslog.LOG_LOCAL6,
}

func NewSyslogPriority(level string, facility string) syslog.Priority {
	p, ok := levelMap[level]
	if !ok {
		p = syslog.LOG_INFO
	}

	f, ok := facilityMap[facility]
	if !ok {
		f = syslog.LOG_LOCAL7
	}

	return p | f
}

func NewServer(syslogAddress string, syslogPriority syslog.Priority) (*Server, error) {
	s := new(Server)
	s.syslogAddress = syslogAddress
	s.syslogPriority = syslogPriority

	writer, err := syslog.Dial("udp", s.syslogAddress, s.syslogPriority, "pixel")
	if err != nil {
		return nil, err
	}

	s.logger = log.New(writer, "", 0)
	s.now = time.Now
	return s, err
}

// Returns true if an interface is a flat mapping of strings to strings
// or numbers.
func isFlatJSON(i interface{}) bool {
	m, ok := i.(map[string]interface{})
	if !ok {
		return false
	}
	for _, v := range m {
		switch v.(type) {
		case string:
		case int:
		case float64:
			continue
		default:
			return false
		}
	}
	return true
}

// Returns a new *Event. *http.Request can be either a pixel or a
// JSON POST request.
func NewEvent(t time.Time, r *http.Request) (*Event, error) {
	event := &Event{
		Time:      t.UTC().Format(ISO8601Format),
		UserAgent: r.Header.Get("User-Agent"),
		IP:        r.Header.Get("X-Forwarded-For"),
		Proto:     r.Header.Get("X-Forwarded-Proto"),
	}
	var err error

	switch r.Method {
	case "GET":
		err = r.ParseForm()
		if err != nil {
			err = fmt.Errorf("Malformed query string: %s", err)
			break
		}

		params := make(map[string]string)
		for key, values := range r.Form {
			params[key] = values[0]
		}
		event.Params = params

	case "POST":
		var body []byte
		var n int
		if r.ContentLength > 0 {
			if r.ContentLength > PostBodyMaxLen {
				err = fmt.Errorf("POST body exceeds max length: %d", r.ContentLength)
				break
			}
			body = make([]byte, r.ContentLength)
		} else if r.ContentLength == -1 {
			body = make([]byte, PostBodyMaxLen)
		} else {
			err = fmt.Errorf("Empty POST body")
			break
		}

		n, err = r.Body.Read(body)
		if err != nil && err != io.EOF {
			break
		}

		var params interface{}
		err = json.Unmarshal(body[:n], &params)
		if err != nil {
			break
		}
		event.Params = &params

	default:
		err = fmt.Errorf("Invalid method %s", r.Method)
	}

	if err != nil {
		event.Error = err.Error()
	}
	return event, err
}

func (s *Server) parseRequest(w http.ResponseWriter, r *http.Request) (*Event, error) {
	event, err := NewEvent(s.now(), r)
	if err != nil {
		log.Printf("%s", err)
		http.Error(w, BadRequest, http.StatusBadRequest)
	}
	return event, err
}

func (s *Server) logEvent(event *Event) {
	jsondata, err := json.Marshal(event)
	if err != nil {
		log.Printf("json encode error: %s", err)
	} else {
		s.logger.Printf("%s", jsondata)
	}
}

func (s *Server) trackPixel(w http.ResponseWriter, r *http.Request) {
	event, err := s.parseRequest(w, r)
	s.logEvent(event)
	if err != nil {
		return
	}
	w.Header().Set("Content-Type", "image/gif")
	w.Write(Transparent1PxGIF)
}

func (s *Server) trackJSON(w http.ResponseWriter, r *http.Request) {
	event, err := s.parseRequest(w, r)
	s.logEvent(event)
	if err != nil {
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte("{}"))
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		s.trackPixel(w, r)
		return

	case "POST":
		s.trackJSON(w, r)

	default:
		http.Error(w, BadRequest, http.StatusBadRequest)
	}
}

func (s *Server) ListenAndServe(address string) {
	s.httpServer = &http.Server{
		Addr:           address,
		Handler:        s,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}
	log.Fatal(s.httpServer.ListenAndServe())
}
<commit_msg>simplify NewEvent() so it makes same-sized slice<commit_after>package pixel

import (
	"encoding/json"
	"fmt"
	"io"
	"log"
	"log/syslog"
	"net/http"
	"time"
)

type Server struct {
	listenAddress  string
	httpServer     *http.Server
	syslogAddress  string
	syslogPriority syslog.Priority
	logger         *log.Logger
	now            func() time.Time
}

// Every pixel GET or JSON POST is an event.
type Event struct {
	Time string `json:"t"`

	// Params are client-supplied event attributes.
	// Pixel events have params of form map[string]string.
	// JSON POST events have arbitrary JSON.
	Params interface{} `json:"params,omitempty"`

	UserAgent string `json:"ua,omitempty"`
	IP        string `json:"ip,omitempty"`
	Proto     string `json:"proto,omitempty"`

	// Error is a server-defined string reporting any error that
	// occurred while constructing the event, such as parse failures.
	Error string `json:"error,omitempty"`
}

const ISO8601Format = "2006-01-02T15:04:05Z"

var Transparent1PxGIF = []byte("\x47\x49\x46\x38\x39\x61\x01\x00" +
	"\x01\x00\x80\x00\x00\xff\xff\xff\x00\x00\x00\x2c\x00\x00\x00\x00" +
	"\x01\x00\x01\x00\x00\x02\x02\x44\x01\x00\x3b")

const UDPMaxBytes = 65507

var BadRequest = http.StatusText(http.StatusBadRequest)

const PostBodyMaxLen = 2048

var levelMap = map[string]syslog.Priority{
	"LOG_EMERG":   syslog.LOG_EMERG,
	"LOG_ALERT":   syslog.LOG_ALERT,
	"LOG_CRIT":    syslog.LOG_CRIT,
	"LOG_ERR":     syslog.LOG_ERR,
	"LOG_WARNING": syslog.LOG_WARNING,
	"LOG_NOTICE":  syslog.LOG_NOTICE,
	"LOG_INFO":    syslog.LOG_INFO,
	"LOG_DEBUG":   syslog.LOG_DEBUG,
}

var facilityMap = map[string]syslog.Priority{
	"LOG_KERN":     syslog.LOG_KERN,
	"LOG_USER":     syslog.LOG_USER,
	"LOG_MAIL":     syslog.LOG_MAIL,
	"LOG_DAEMON":   syslog.LOG_DAEMON,
	"LOG_AUTH":     syslog.LOG_AUTH,
	"LOG_SYSLOG":   syslog.LOG_SYSLOG,
	"LOG_LPR":      syslog.LOG_LPR,
	"LOG_NEWS":     syslog.LOG_NEWS,
	"LOG_UUCP":     syslog.LOG_UUCP,
	"LOG_CRON":     syslog.LOG_CRON,
	"LOG_AUTHPRIV": syslog.LOG_AUTHPRIV,
	"LOG_FTP":      syslog.LOG_FTP,
	"LOG_LOCAL0":   syslog.LOG_LOCAL0,
	"LOG_LOCAL1":   syslog.LOG_LOCAL1,
	"LOG_LOCAL2":   syslog.LOG_LOCAL2,
	"LOG_LOCAL3":   syslog.LOG_LOCAL3,
	"LOG_LOCAL4":   syslog.LOG_LOCAL4,
	"LOG_LOCAL5":   syslog.LOG_LOCAL5,
	"LOG_LOCAL6":   syslog.LOG_LOCAL6,
}

func NewSyslogPriority(level string, facility string) syslog.Priority {
	p, ok := levelMap[level]
	if !ok {
		p = syslog.LOG_INFO
	}

	f, ok := facilityMap[facility]
	if !ok {
		f = syslog.LOG_LOCAL7
	}

	return p | f
}

func NewServer(syslogAddress string, syslogPriority syslog.Priority) (*Server, error) {
	s := new(Server)
	s.syslogAddress = syslogAddress
	s.syslogPriority = syslogPriority

	writer, err := syslog.Dial("udp", s.syslogAddress, s.syslogPriority, "pixel")
	if err != nil {
		return nil, err
	}

	s.logger = log.New(writer, "", 0)
	s.now = time.Now
	return s, err
}

// Returns true if an interface is a flat mapping of strings to strings
// or numbers.
func isFlatJSON(i interface{}) bool {
	m, ok := i.(map[string]interface{})
	if !ok {
		return false
	}
	for _, v := range m {
		switch v.(type) {
		case string:
		case int:
		case float64:
			continue
		default:
			return false
		}
	}
	return true
}

// Returns a new *Event. *http.Request can be either a pixel or a
// JSON POST request.
func NewEvent(t time.Time, r *http.Request) (*Event, error) {
	event := &Event{
		Time:      t.UTC().Format(ISO8601Format),
		UserAgent: r.Header.Get("User-Agent"),
		IP:        r.Header.Get("X-Forwarded-For"),
		Proto:     r.Header.Get("X-Forwarded-Proto"),
	}
	var err error

	switch r.Method {
	case "GET":
		err = r.ParseForm()
		if err != nil {
			err = fmt.Errorf("Malformed query string: %s", err)
			break
		}

		params := make(map[string]string)
		for key, values := range r.Form {
			params[key] = values[0]
		}
		event.Params = params

	case "POST":
		if r.ContentLength == 0 {
			err = fmt.Errorf("Empty POST body")
			break
		}
		if r.ContentLength > PostBodyMaxLen {
			err = fmt.Errorf("POST body exceeds max length: %d", r.ContentLength)
			break
		}

		body := make([]byte, PostBodyMaxLen)

		var n int
		n, err = r.Body.Read(body)
		if err != nil && err != io.EOF {
			break
		}

		var params interface{}
		err = json.Unmarshal(body[:n], &params)
		if err != nil {
			break
		}
		event.Params = &params

	default:
		err = fmt.Errorf("Invalid method %s", r.Method)
	}

	if err != nil {
		event.Error = err.Error()
	}
	return event, err
}

func (s *Server) parseRequest(w http.ResponseWriter, r *http.Request) (*Event, error) {
	event, err := NewEvent(s.now(), r)
	if err != nil {
		log.Printf("%s", err)
		http.Error(w, BadRequest, http.StatusBadRequest)
	}
	return event, err
}

func (s *Server) logEvent(event *Event) {
	jsondata, err := json.Marshal(event)
	if err != nil {
		log.Printf("json encode error: %s", err)
	} else {
		s.logger.Printf("%s", jsondata)
	}
}

func (s *Server) trackPixel(w http.ResponseWriter, r *http.Request) {
	event, err := s.parseRequest(w, r)
	s.logEvent(event)
	if err != nil {
		return
	}
	w.Header().Set("Content-Type", "image/gif")
	w.Write(Transparent1PxGIF)
}

func (s *Server) trackJSON(w http.ResponseWriter, r *http.Request) {
	event, err := s.parseRequest(w, r)
	s.logEvent(event)
	if err != nil {
		return
	}
	w.Header().Set("Content-Type", "application/json")
	w.Write([]byte("{}"))
}

func (s *Server) ServeHTTP(w http.ResponseWriter, r *http.Request) {
	switch r.Method {
	case "GET":
		s.trackPixel(w, r)
		return

	case "POST":
		s.trackJSON(w, r)

	default:
		http.Error(w, BadRequest, http.StatusBadRequest)
	}
}

func (s *Server) ListenAndServe(address string) {
	s.httpServer = &http.Server{
		Addr:           address,
		Handler:        s,
		ReadTimeout:    10 * time.Second,
		WriteTimeout:   10 * time.Second,
		MaxHeaderBytes: 1 << 20,
	}
	log.Fatal(s.httpServer.ListenAndServe())
}
<|endoftext|>


<commit_before>// Copyright 2014 struktur AG. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package httputils

import (
	"crypto/tls"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

// Provide a HTTP server implementation which can listen on TPC
// and Unix Domain sockets.
type Server struct {
	http.Server
	*log.Logger
}

// ListenAndServe binds sockets according to the configuration of srv and blocks
// until the socket closes or an exit signal is received.
func (srv *Server) ListenAndServe() error {

	addr := srv.Addr
	if addr == "" {
		addr = ":http"
	}

	var closing = false
	var err error
	var l net.Listener
	if l, err = srv.socketListen(addr); err != nil {
		return err
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() {
		s := <-sig
		msg := "Received exit signal %d - Closing ..."
		if srv.Logger != nil {
			srv.Logger.Printf(msg, s)
		} else {
			log.Printf(msg, s)
		}
		closing = true
		l.Close()
	}()

	err = srv.Serve(l)
	if err != nil {
		if closing {
			return nil
		}
	}
	return err

}

// ListenAndServeTLS binds sockets according to the configuration of srv and blocks
// until the socket closes or an exit signal is received.
func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {

	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

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

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	var closing = false
	var l net.Listener
	if l, err = srv.socketListen(addr); err != nil {
		return err
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() {
		s := <-sig
		msg := "Received exit signal %d - Closing ..."
		if srv.Logger != nil {
			srv.Logger.Printf(msg, s)
		} else {
			log.Printf(msg, s)
		}
		closing = true
		l.Close()
	}()

	tlsListener := tls.NewListener(l, config)

	err = srv.Serve(tlsListener)
	if err != nil {
		if closing {
			return nil
		}
	}
	return err

}

func (srv *Server) socketListen(addr string) (net.Listener, error) {

	var err error
	var l net.Listener

	if strings.HasPrefix(addr, "/") {
		var laddr *net.UnixAddr
		if laddr, err = net.ResolveUnixAddr("unix", addr); err != nil {
			return nil, err
		}
		if l, err = createUnixSocket(laddr); err != nil {
			// Unix-domain-socket already exists, try to connect to it to
			// see if it still is usedb by another process
			if _, err = net.Dial("unix", addr); err != nil {
				if err = os.Remove(addr); err != nil {
					return nil, err
				}
				if l, err = createUnixSocket(laddr); err != nil {
					return nil, err
				}
			} else {
				return nil, fmt.Errorf("another process seems to be listening on %s already", addr)
			}
		}
	} else {
		var laddr *net.TCPAddr
		if laddr, err = net.ResolveTCPAddr("tcp", addr); err != nil {
			return nil, err
		}
		if l, err = net.ListenTCP("tcp", laddr); err != nil {
			return nil, err
		}
	}

	return l, nil

}

func createUnixSocket(addr *net.UnixAddr) (l net.Listener, err error) {
	l, err = net.ListenUnix("unix", addr)

	if err != nil {
		return
	}

	// TODO(lcooper): Not sure if this sequence is completely safe.
	// It would be better if we could get the underlying FD of the socket
	// and stat() + chmod() that instead.
	fi, err := os.Stat(addr.String())
	if err != nil {
		return
	}

	// NOTE(lcooper): This only ensures g+w other then on Linux,
	// BSD systems only use parent directory permissions.
	// See http://stackoverflow.com/questions/5977556 .
	err = os.Chmod(addr.String(), fi.Mode()|0060)
	return
}
<commit_msg>Added advanced tls listen function.<commit_after>// Copyright 2014 struktur AG. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package httputils

import (
	"crypto/tls"
	"errors"
	"fmt"
	"log"
	"net"
	"net/http"
	"os"
	"os/signal"
	"strings"
	"syscall"
)

// Provide a HTTP server implementation which can listen on TPC
// and Unix Domain sockets.
type Server struct {
	http.Server
	*log.Logger
}

// ListenAndServe binds sockets according to the configuration of srv and blocks
// until the socket closes or an exit signal is received.
func (srv *Server) ListenAndServe() error {

	addr := srv.Addr
	if addr == "" {
		addr = ":http"
	}

	var closing = false
	var err error
	var l net.Listener
	if l, err = srv.socketListen(addr); err != nil {
		return err
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() {
		s := <-sig
		msg := "Received exit signal %d - Closing ..."
		if srv.Logger != nil {
			srv.Logger.Printf(msg, s)
		} else {
			log.Printf(msg, s)
		}
		closing = true
		l.Close()
	}()

	err = srv.Serve(l)
	if err != nil {
		if closing {
			return nil
		}
	}
	return err

}

// ListenAndServeTLS binds sockets according to the configuration of srv and blocks
// until the socket closes or an exit signal is received.
func (srv *Server) ListenAndServeTLS(certFile, keyFile string) error {

	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

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

	var err error
	config.Certificates = make([]tls.Certificate, 1)
	config.Certificates[0], err = tls.LoadX509KeyPair(certFile, keyFile)
	if err != nil {
		return err
	}

	var closing = false
	var l net.Listener
	if l, err = srv.socketListen(addr); err != nil {
		return err
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() {
		s := <-sig
		msg := "Received exit signal %d - Closing ..."
		if srv.Logger != nil {
			srv.Logger.Printf(msg, s)
		} else {
			log.Printf(msg, s)
		}
		closing = true
		l.Close()
	}()

	tlsListener := tls.NewListener(l, config)

	err = srv.Serve(tlsListener)
	if err != nil {
		if closing {
			return nil
		}
	}
	return err

}

// ListenAndServeTLSAdvanced binds sockets according to the configuration
// of srv and blocks until the socket closes or an exit signal is received. A
// TLSConfig needs to be available at the server.
func (srv *Server) ListenAndServeTLSAdvanced() error {

	addr := srv.Addr
	if addr == "" {
		addr = ":https"
	}

	config := srv.TLSConfig

	if config == nil {
		return errors.New("TLSConfig required")
	}
	if config.NextProtos == nil {
		config.NextProtos = []string{"http/1.1"}
	}

	if len(config.Certificates) == 0 {
		return errors.New("TLSConfig has no certificate")
	}

	var err error
	var closing = false
	var l net.Listener
	if l, err = srv.socketListen(addr); err != nil {
		return err
	}

	sig := make(chan os.Signal, 1)
	signal.Notify(sig, os.Interrupt, syscall.SIGTERM)
	go func() {
		s := <-sig
		msg := "Received exit signal %d - Closing ..."
		if srv.Logger != nil {
			srv.Logger.Printf(msg, s)
		} else {
			log.Printf(msg, s)
		}
		closing = true
		l.Close()
	}()

	tlsListener := tls.NewListener(l, config)

	err = srv.Serve(tlsListener)
	if err != nil {
		if closing {
			return nil
		}
	}
	return err

}

func (srv *Server) socketListen(addr string) (net.Listener, error) {

	var err error
	var l net.Listener

	if strings.HasPrefix(addr, "/") {
		var laddr *net.UnixAddr
		if laddr, err = net.ResolveUnixAddr("unix", addr); err != nil {
			return nil, err
		}
		if l, err = createUnixSocket(laddr); err != nil {
			// Unix-domain-socket already exists, try to connect to it to
			// see if it still is usedb by another process
			if _, err = net.Dial("unix", addr); err != nil {
				if err = os.Remove(addr); err != nil {
					return nil, err
				}
				if l, err = createUnixSocket(laddr); err != nil {
					return nil, err
				}
			} else {
				return nil, fmt.Errorf("another process seems to be listening on %s already", addr)
			}
		}
	} else {
		var laddr *net.TCPAddr
		if laddr, err = net.ResolveTCPAddr("tcp", addr); err != nil {
			return nil, err
		}
		if l, err = net.ListenTCP("tcp", laddr); err != nil {
			return nil, err
		}
	}

	return l, nil

}

func createUnixSocket(addr *net.UnixAddr) (l net.Listener, err error) {
	l, err = net.ListenUnix("unix", addr)

	if err != nil {
		return
	}

	// TODO(lcooper): Not sure if this sequence is completely safe.
	// It would be better if we could get the underlying FD of the socket
	// and stat() + chmod() that instead.
	fi, err := os.Stat(addr.String())
	if err != nil {
		return
	}

	// NOTE(lcooper): This only ensures g+w other then on Linux,
	// BSD systems only use parent directory permissions.
	// See http://stackoverflow.com/questions/5977556 .
	err = os.Chmod(addr.String(), fi.Mode()|0060)
	return
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"github.com/RobinUS2/golang-jresp"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Server methods (you probably only need one or two in HA failover mode)

type Server struct {
	clientsMux sync.RWMutex
	clients    map[string]*RegisteredClient
}

// Register client
func (s *Server) RegisterClient(hostname string) {
	s.clientsMux.Lock()
	if s.clients[hostname] == nil {
		s.clients[hostname] = newRegisteredClient(hostname)
		log.Printf("Client %s registered", hostname)
	}
	s.clients[hostname].mux.Lock()
	s.clients[hostname].LastPing = time.Now()
	s.clients[hostname].mux.Unlock()
	s.clientsMux.Unlock()
}

func (s *Server) GetClient(hostname string) *RegisteredClient {
	s.clientsMux.Lock()
	defer s.clientsMux.Unlock()
	return s.clients[hostname]
}

// Scan for old clients
func (s *Server) CleanupClients() {
	s.clientsMux.Lock()
	for k, client := range s.clients {
		if time.Now().Sub(client.LastPing).Seconds() > float64(CLIENT_PING_INTERVAL*5) {
			// Disconnect
			log.Printf("Client %s disconnected", client.Hostname)
			delete(s.clients, k)
		}
	}
	s.clientsMux.Unlock()
}

type RegisteredClient struct {
	mux      sync.RWMutex
	Hostname string
	LastPing time.Time
	Cmds     map[string]*Cmd
	CmdChan  chan bool
}

// Generate keys
func (s *Server) _prepareTlsKeys() {
	if _, err := os.Stat("./private_key"); os.IsNotExist(err) {
		// No keys, generate
		log.Println("Auto-generating keys for server")
		cmd := newCmd("./generate_key.sh", 60)
		cmd.Execute()
		log.Println("Finished generating keys for server")
	}
}

// Start server
func (s *Server) Start() bool {
	log.Printf("Starting server at https://localhost:%d/", serverPort)

	// Start webserver
	go func() {
		router := httprouter.New()
		router.GET("/", Home)
		router.GET("/ping", Ping)
		router.GET("/client/:hostname/ping", ClientPing)
		router.GET("/client/:hostname/cmds", ClientCmds)
		router.POST("/client/:hostname/cmd", PostClientCmd)
		router.ServeFiles("/console/*filepath", http.Dir("console"))

		// Auto generate key
		s._prepareTlsKeys()

		// Start server
		log.Fatal(http.ListenAndServeTLS(fmt.Sprintf(":%d", serverPort), "./public_key", "./private_key", router))
	}()

	// Minutely cleanups etc
	go func() {
		c := time.Tick(1 * time.Minute)
		for _ = range c {
			server.CleanupClients()
		}
	}()

	return true
}

// Submit client command
func PostClientCmd(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {

	jr := jresp.NewJsonResp()
	if !auth(r) {
		jr.Error("Not authorized")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// Get client
	registeredClient := server.GetClient(ps.ByName("hostname"))
	if registeredClient == nil {
		jr.Error("Client not registered")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// Timeout
	timeoutStr := r.URL.Query().Get("timeout")
	var timeout int = DEFAULT_COMMAND_TIMEOUT
	if len(strings.TrimSpace(timeoutStr)) > 0 {
		timeoutI, timeoutE := strconv.ParseInt(timeoutStr, 10, 0)
		if timeoutE != nil || timeoutI < 1 {
			jr.Error("Invalid timeout value")
			fmt.Fprint(w, jr.ToString(debug))
			return
		}
		timeout = int(timeoutI)
	}

	// Create command
	command := r.URL.Query().Get("cmd")
	if len(strings.TrimSpace(command)) < 1 {
		jr.Error("Provide a command")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}
	cmd := newCmd(command, timeout) // @todo dynamic

	// Add to list
	registeredClient.mux.Lock()
	registeredClient.Cmds[cmd.Id] = cmd
	registeredClient.CmdChan <- true // Signal for work
	registeredClient.mux.Unlock()

	jr.Set("ack", true)
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Commands
func ClientCmds(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	jr := jresp.NewJsonResp()
	if !auth(r) {
		jr.Error("Not authorized")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// Get client
	registeredClient := server.GetClient(ps.ByName("hostname"))
	if registeredClient == nil {
		jr.Error("Client not registered")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// @todo Read from channel flag and dispatch before timeout
	select {
	case <-registeredClient.CmdChan:
		cmds := make([]*Cmd, 0)
		registeredClient.mux.Lock()
		for _, cmd := range registeredClient.Cmds {
			if cmd.Pending {
				cmds = append(cmds, cmd)
				cmd.Pending = false
			}
		}
		registeredClient.mux.Unlock()
		jr.Set("cmds", cmds)
	case <-time.After(time.Second * LONG_POLL_TIMEOUT):
		// No commands
		jr.Set("cmds", make([]string, 0))
	}
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Ping
func ClientPing(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	jr := jresp.NewJsonResp()
	if !auth(r) {
		jr.Error("Not authorized")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}
	tags := strings.Split(r.URL.Query().Get("tags"), ",")
	log.Printf("Tags %v", tags)
	server.RegisterClient(ps.ByName("hostname"))
	jr.Set("ack", true)
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Home
func Home(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	// Redirect to console
	http.Redirect(w, r, r.URL.String()+"console/", 301)
}

// Ping
func Ping(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	jr := jresp.NewJsonResp()
	jr.Set("ping", "pong")
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Auth
func auth(r *http.Request) bool {
	// Signed token
	uri := r.URL.String()
	hasher := sha256.New()
	hasher.Write([]byte(uri))
	signedToken := base64.URLEncoding.EncodeToString(hasher.Sum(nil))

	// Validate
	if r.Header.Get("X-Auth") != signedToken {
		return false
	}
	return true
}

// Create new server
func newServer() *Server {
	return &Server{
		clients: make(map[string]*RegisteredClient),
	}
}

// New registered client
func newRegisteredClient(hostname string) *RegisteredClient {
	return &RegisteredClient{
		Hostname: hostname,
		Cmds:     make(map[string]*Cmd),
		CmdChan:  make(chan bool),
	}
}
<commit_msg>Server track tags<commit_after>package main

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"github.com/RobinUS2/golang-jresp"
	"github.com/julienschmidt/httprouter"
	"net/http"
	"os"
	"strconv"
	"strings"
	"sync"
	"time"
)

// Server methods (you probably only need one or two in HA failover mode)

type Server struct {
	clientsMux sync.RWMutex
	clients    map[string]*RegisteredClient
}

// Register client
func (s *Server) RegisterClient(clientId string, tags []string) {
	s.clientsMux.Lock()
	if s.clients[clientId] == nil {
		s.clients[clientId] = newRegisteredClient(clientId)
		log.Printf("Client %s registered with tags %s", clientId, tags)
	}
	s.clients[clientId].mux.Lock()
	s.clients[clientId].LastPing = time.Now()
	s.clients[clientId].Tags = tags
	s.clients[clientId].mux.Unlock()
	s.clientsMux.Unlock()
}

func (s *Server) GetClient(clientId string) *RegisteredClient {
	s.clientsMux.Lock()
	defer s.clientsMux.Unlock()
	return s.clients[clientId]
}

// Scan for old clients
func (s *Server) CleanupClients() {
	s.clientsMux.Lock()
	for k, client := range s.clients {
		if time.Now().Sub(client.LastPing).Seconds() > float64(CLIENT_PING_INTERVAL*5) {
			// Disconnect
			log.Printf("Client %s disconnected", client.clientId)
			delete(s.clients, k)
		}
	}
	s.clientsMux.Unlock()
}

type RegisteredClient struct {
	mux      sync.RWMutex
	clientId string
	LastPing time.Time
	Tags     []string
	Cmds     map[string]*Cmd
	CmdChan  chan bool
}

// Generate keys
func (s *Server) _prepareTlsKeys() {
	if _, err := os.Stat("./private_key"); os.IsNotExist(err) {
		// No keys, generate
		log.Println("Auto-generating keys for server")
		cmd := newCmd("./generate_key.sh", 60)
		cmd.Execute()
		log.Println("Finished generating keys for server")
	}
}

// Start server
func (s *Server) Start() bool {
	log.Printf("Starting server at https://localhost:%d/", serverPort)

	// Start webserver
	go func() {
		router := httprouter.New()
		router.GET("/", Home)
		router.GET("/ping", Ping)
		router.GET("/client/:clientId/ping", ClientPing)
		router.GET("/client/:clientId/cmds", ClientCmds)
		router.POST("/client/:clientId/cmd", PostClientCmd)
		router.ServeFiles("/console/*filepath", http.Dir("console"))

		// Auto generate key
		s._prepareTlsKeys()

		// Start server
		log.Fatal(http.ListenAndServeTLS(fmt.Sprintf(":%d", serverPort), "./public_key", "./private_key", router))
	}()

	// Minutely cleanups etc
	go func() {
		c := time.Tick(1 * time.Minute)
		for _ = range c {
			server.CleanupClients()
		}
	}()

	return true
}

// Submit client command
func PostClientCmd(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {

	jr := jresp.NewJsonResp()
	if !auth(r) {
		jr.Error("Not authorized")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// Get client
	registeredClient := server.GetClient(ps.ByName("clientId"))
	if registeredClient == nil {
		jr.Error("Client not registered")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// Timeout
	timeoutStr := r.URL.Query().Get("timeout")
	var timeout int = DEFAULT_COMMAND_TIMEOUT
	if len(strings.TrimSpace(timeoutStr)) > 0 {
		timeoutI, timeoutE := strconv.ParseInt(timeoutStr, 10, 0)
		if timeoutE != nil || timeoutI < 1 {
			jr.Error("Invalid timeout value")
			fmt.Fprint(w, jr.ToString(debug))
			return
		}
		timeout = int(timeoutI)
	}

	// Create command
	command := r.URL.Query().Get("cmd")
	if len(strings.TrimSpace(command)) < 1 {
		jr.Error("Provide a command")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}
	cmd := newCmd(command, timeout) // @todo dynamic

	// Add to list
	registeredClient.mux.Lock()
	registeredClient.Cmds[cmd.Id] = cmd
	registeredClient.CmdChan <- true // Signal for work
	registeredClient.mux.Unlock()

	jr.Set("ack", true)
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Commands
func ClientCmds(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	jr := jresp.NewJsonResp()
	if !auth(r) {
		jr.Error("Not authorized")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// Get client
	registeredClient := server.GetClient(ps.ByName("clientId"))
	if registeredClient == nil {
		jr.Error("Client not registered")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}

	// @todo Read from channel flag and dispatch before timeout
	select {
	case <-registeredClient.CmdChan:
		cmds := make([]*Cmd, 0)
		registeredClient.mux.Lock()
		for _, cmd := range registeredClient.Cmds {
			if cmd.Pending {
				cmds = append(cmds, cmd)
				cmd.Pending = false
			}
		}
		registeredClient.mux.Unlock()
		jr.Set("cmds", cmds)
	case <-time.After(time.Second * LONG_POLL_TIMEOUT):
		// No commands
		jr.Set("cmds", make([]string, 0))
	}
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Ping
func ClientPing(w http.ResponseWriter, r *http.Request, ps httprouter.Params) {
	jr := jresp.NewJsonResp()
	if !auth(r) {
		jr.Error("Not authorized")
		fmt.Fprint(w, jr.ToString(debug))
		return
	}
	tags := strings.Split(r.URL.Query().Get("tags"), ",")
	server.RegisterClient(ps.ByName("clientId"), tags)
	jr.Set("ack", true)
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Home
func Home(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	// Redirect to console
	http.Redirect(w, r, r.URL.String()+"console/", 301)
}

// Ping
func Ping(w http.ResponseWriter, r *http.Request, _ httprouter.Params) {
	jr := jresp.NewJsonResp()
	jr.Set("ping", "pong")
	jr.OK()
	fmt.Fprint(w, jr.ToString(debug))
}

// Auth
func auth(r *http.Request) bool {
	// Signed token
	uri := r.URL.String()
	hasher := sha256.New()
	hasher.Write([]byte(uri))
	signedToken := base64.URLEncoding.EncodeToString(hasher.Sum(nil))

	// Validate
	if r.Header.Get("X-Auth") != signedToken {
		return false
	}
	return true
}

// Create new server
func newServer() *Server {
	return &Server{
		clients: make(map[string]*RegisteredClient),
	}
}

// New registered client
func newRegisteredClient(clientId string) *RegisteredClient {
	return &RegisteredClient{
		clientId: clientId,
		Cmds:     make(map[string]*Cmd),
		CmdChan:  make(chan bool),
	}
}
<|endoftext|>


<commit_before>package main

import (
	"crypto/md5"
	"fmt"
	"net"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"
)

const (
	maxNameLength        = 32
	historyLength        = 20
	systemMessageFormat  = "\033[1;90m"
	privateMessageFormat = "\033[1m"
	highlightFormat      = Bold + "\033[48;5;11m\033[38;5;16m"
	beep                 = "\007"
)

var (
	reStripText = regexp.MustCompile("[^0-9A-Za-z_.-]")
)

// Clients is a map of clients
type Clients map[string]*Client

// Server holds all the fields used by a server
type Server struct {
	sshConfig *ssh.ServerConfig
	done      chan struct{}
	clients   Clients
	count     int
	history   *History
	motd      string
	whitelist map[string]struct{}   // fingerprint lookup
	admins    map[string]struct{}   // fingerprint lookup
	bannedPK  map[string]*time.Time // fingerprint lookup
	started   time.Time
	sync.Mutex
}

// NewServer constructs a new server
func NewServer(privateKey []byte) (*Server, error) {
	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}

	server := Server{
		done:      make(chan struct{}),
		clients:   Clients{},
		count:     0,
		history:   NewHistory(historyLength),
		motd:      "",
		whitelist: map[string]struct{}{},
		admins:    map[string]struct{}{},
		bannedPK:  map[string]*time.Time{},
		started:   time.Now(),
	}

	config := ssh.ServerConfig{
		NoClientAuth: false,
		// Auth-related things should be constant-time to avoid timing attacks.
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			fingerprint := Fingerprint(key)
			if server.IsBanned(fingerprint) {
				return nil, fmt.Errorf("Banned.")
			}
			if !server.IsWhitelisted(fingerprint) {
				return nil, fmt.Errorf("Not Whitelisted.")
			}
			perm := &ssh.Permissions{Extensions: map[string]string{"fingerprint": fingerprint}}
			return perm, nil
		},
	}
	config.AddHostKey(signer)

	server.sshConfig = &config

	return &server, nil
}

// Len returns the number of clients
func (s *Server) Len() int {
	return len(s.clients)
}

// SysMsg broadcasts the given message to everyone
func (s *Server) SysMsg(msg string, args ...interface{}) {
	s.Broadcast(ContinuousFormat(systemMessageFormat, " * "+fmt.Sprintf(msg, args...)), nil)
}

// Broadcast broadcasts the given message to everyone except for the given client
func (s *Server) Broadcast(msg string, except *Client) {
	logger.Debugf("Broadcast to %d: %s", s.Len(), msg)
	s.history.Add(msg)

	for _, client := range s.clients {
		if except != nil && client == except {
			continue
		}

		if strings.Contains(msg, client.Name) {
			// Turn message red if client's name is mentioned, and send BEL if they have enabled beeping
			personalMsg := strings.Replace(msg, client.Name, highlightFormat+client.Name+Reset, -1)
			if client.beepMe {
				personalMsg += beep
			}
			client.Send(personalMsg)
		} else {
			client.Send(msg)
		}
	}
}

// Privmsg sends a message to a particular nick, if it exists
func (s *Server) Privmsg(nick, message string, sender *Client) error {
	// Get the recipient
	target, ok := s.clients[strings.ToLower(nick)]
	if !ok {
		return fmt.Errorf("no client with that nick")
	}
	// Send the message
	target.Msg <- fmt.Sprintf(beep+"[PM from %v] %s%v%s", sender.ColoredName(), privateMessageFormat, message, Reset)
	logger.Debugf("PM from %v to %v: %v", sender.Name, nick, message)
	return nil
}

// SetMotd sets the Message of the Day (MOTD)
func (s *Server) SetMotd(motd string) {
	s.Lock()
	s.motd = motd
	s.Unlock()
}

// MotdUnicast sends the MOTD as a SysMsg
func (s *Server) MotdUnicast(client *Client) {
	if s.motd == "" {
		return
	}
	client.SysMsg(s.motd)
}

// MotdBroadcast broadcasts the MOTD
func (s *Server) MotdBroadcast(client *Client) {
	if s.motd == "" {
		return
	}
	s.Broadcast(ContinuousFormat(systemMessageFormat, fmt.Sprintf(" * New MOTD set by %s.", client.ColoredName())), client)
	s.Broadcast(s.motd, client)
}

// Add adds the client to the list of clients
func (s *Server) Add(client *Client) {
	go func() {
		s.MotdUnicast(client)
		client.SendLines(s.history.Get(10))
	}()

	s.Lock()
	s.count++

	newName, err := s.proposeName(client.Name)
	if err != nil {
		client.SysMsg("Your name '%s' is not available, renamed to '%s'. Use /nick <name> to change it.", client.Name, ColorString(client.Color, newName))
	}

	client.Rename(newName)
	s.clients[strings.ToLower(client.Name)] = client
	num := len(s.clients)
	s.Unlock()

	s.Broadcast(ContinuousFormat(systemMessageFormat, fmt.Sprintf(" * %s joined. (Total connected: %d)", client.Name, num)), client)
}

// Remove removes the given client from the list of clients
func (s *Server) Remove(client *Client) {
	s.Lock()
	delete(s.clients, strings.ToLower(client.Name))
	s.Unlock()

	s.SysMsg("%s left.", client.Name)
}

func (s *Server) proposeName(name string) (string, error) {
	// Assumes caller holds lock.
	var err error
	name = reStripText.ReplaceAllString(name, "")

	if len(name) > maxNameLength {
		name = name[:maxNameLength]
	} else if len(name) == 0 {
		name = fmt.Sprintf("Guest%d", s.count)
	}

	_, collision := s.clients[strings.ToLower(name)]
	if collision {
		err = fmt.Errorf("%s is not available", name)
		name = fmt.Sprintf("Guest%d", s.count)
	}

	return name, err
}

// Rename renames the given client (user)
func (s *Server) Rename(client *Client, newName string) {
	s.Lock()
	var oldName string
	if strings.ToLower(newName) == strings.ToLower(client.Name) {
		oldName = client.Name
		client.Rename(newName)
	} else {
		newName, err := s.proposeName(newName)
		if err != nil {
			client.SysMsg("%s", err)
			s.Unlock()
			return
		}

		// TODO: Use a channel/goroutine for adding clients, rather than locks?
		delete(s.clients, strings.ToLower(client.Name))
		oldName = client.Name
		client.Rename(newName)
		s.clients[strings.ToLower(client.Name)] = client
		s.Unlock()
	}
	s.SysMsg("%s is now known as %s.", ColorString(client.Color, oldName), ColorString(client.Color, newName))
}

// List lists the clients with the given prefix
func (s *Server) List(prefix *string) []string {
	r := []string{}

	for name := range s.clients {
		if prefix != nil && !strings.HasPrefix(name, *prefix) {
			continue
		}
		r = append(r, name)
	}

	return r
}

// Who returns the client with a given name
func (s *Server) Who(name string) *Client {
	return s.clients[strings.ToLower(name)]
}

// Op adds the given fingerprint to the list of admins
func (s *Server) Op(fingerprint string) {
	logger.Infof("Adding admin: %s", fingerprint)
	s.Lock()
	s.admins[fingerprint] = struct{}{}
	s.Unlock()
}

// Whitelist adds the given fingerprint to the whitelist
func (s *Server) Whitelist(fingerprint string) {
	logger.Infof("Adding whitelist: %s", fingerprint)
	s.Lock()
	s.whitelist[fingerprint] = struct{}{}
	s.Unlock()
}

// Uptime returns the time since the server was started
func (s *Server) Uptime() string {
	return time.Now().Sub(s.started).String()
}

// IsOp checks if the given client is Op
func (s *Server) IsOp(client *Client) bool {
	_, r := s.admins[client.Fingerprint()]
	return r
}

// IsWhitelisted checks if the given fingerprint is whitelisted
func (s *Server) IsWhitelisted(fingerprint string) bool {
	/* if no whitelist, anyone is welcome */
	if len(s.whitelist) == 0 {
		return true
	}

	/* otherwise, check for whitelist presence */
	_, r := s.whitelist[fingerprint]
	return r
}

// IsBanned checks if the given fingerprint is banned
func (s *Server) IsBanned(fingerprint string) bool {
	ban, hasBan := s.bannedPK[fingerprint]
	if !hasBan {
		return false
	}
	if ban == nil {
		return true
	}
	if ban.Before(time.Now()) {
		s.Unban(fingerprint)
		return false
	}
	return true
}

// Ban bans a fingerprint for the given duration
func (s *Server) Ban(fingerprint string, duration *time.Duration) {
	var until *time.Time
	s.Lock()
	if duration != nil {
		when := time.Now().Add(*duration)
		until = &when
	}
	s.bannedPK[fingerprint] = until
	s.Unlock()
}

// Unban unbans a banned fingerprint
func (s *Server) Unban(fingerprint string) {
	s.Lock()
	delete(s.bannedPK, fingerprint)
	s.Unlock()
}

// Start starts the server
func (s *Server) Start(laddr string) error {
	// Once a ServerConfig has been configured, connections can be
	// accepted.
	socket, err := net.Listen("tcp", laddr)
	if err != nil {
		return err
	}

	logger.Infof("Listening on %s", laddr)

	go func() {
		defer socket.Close()
		for {
			conn, err := socket.Accept()

			if err != nil {
				logger.Errorf("Failed to accept connection: %v", err)
				if err == syscall.EINVAL {
					// TODO: Handle shutdown more gracefully?
					return
				}
			}

			// Goroutineify to resume accepting sockets early.
			go func() {
				// From a standard TCP connection to an encrypted SSH connection
				sshConn, channels, requests, err := ssh.NewServerConn(conn, s.sshConfig)
				if err != nil {
					logger.Errorf("Failed to handshake: %v", err)
					return
				}

				version := reStripText.ReplaceAllString(string(sshConn.ClientVersion()), "")
				if len(version) > 100 {
					version = "Evil Jerk with a superlong string"
				}
				logger.Infof("Connection #%d from: %s, %s, %s", s.count+1, sshConn.RemoteAddr(), sshConn.User(), version)

				go ssh.DiscardRequests(requests)

				client := NewClient(s, sshConn)
				go client.handleChannels(channels)
			}()
		}
	}()

	go func() {
		<-s.done
		socket.Close()
	}()

	return nil
}

// AutoCompleteFunction handles auto completion of nicks
func (s *Server) AutoCompleteFunction(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
	if key == 9 {
		shortLine := strings.Split(line[:pos], " ")
		partialNick := shortLine[len(shortLine)-1]

		nicks := s.List(&partialNick)
		if len(nicks) > 0 {
			nick := nicks[len(nicks)-1]
			posPartialNick := pos - len(partialNick)
			if len(shortLine) < 2 {
				nick += ": "
			} else {
				nick += " "
			}
			newLine = strings.Replace(line[posPartialNick:],
				partialNick, nick, 1)
			newLine = line[:posPartialNick] + newLine
			newPos = pos + (len(nick) - len(partialNick))
			ok = true
		}
	} else {
		ok = false
	}
	return
}

// Stop stops the server
func (s *Server) Stop() {
	for _, client := range s.clients {
		client.Conn.Close()
	}

	close(s.done)
}

// Fingerprint returns the fingerprint based on a public key
func Fingerprint(k ssh.PublicKey) string {
	hash := md5.Sum(k.Marshal())
	r := fmt.Sprintf("% x", hash)
	return strings.Replace(r, " ", ":", -1)
}
<commit_msg>lowercase prefix<commit_after>package main

import (
	"crypto/md5"
	"fmt"
	"net"
	"regexp"
	"strings"
	"sync"
	"syscall"
	"time"

	"golang.org/x/crypto/ssh"
)

const (
	maxNameLength        = 32
	historyLength        = 20
	systemMessageFormat  = "\033[1;90m"
	privateMessageFormat = "\033[1m"
	highlightFormat      = Bold + "\033[48;5;11m\033[38;5;16m"
	beep                 = "\007"
)

var (
	reStripText = regexp.MustCompile("[^0-9A-Za-z_.-]")
)

// Clients is a map of clients
type Clients map[string]*Client

// Server holds all the fields used by a server
type Server struct {
	sshConfig *ssh.ServerConfig
	done      chan struct{}
	clients   Clients
	count     int
	history   *History
	motd      string
	whitelist map[string]struct{}   // fingerprint lookup
	admins    map[string]struct{}   // fingerprint lookup
	bannedPK  map[string]*time.Time // fingerprint lookup
	started   time.Time
	sync.Mutex
}

// NewServer constructs a new server
func NewServer(privateKey []byte) (*Server, error) {
	signer, err := ssh.ParsePrivateKey(privateKey)
	if err != nil {
		return nil, err
	}

	server := Server{
		done:      make(chan struct{}),
		clients:   Clients{},
		count:     0,
		history:   NewHistory(historyLength),
		motd:      "",
		whitelist: map[string]struct{}{},
		admins:    map[string]struct{}{},
		bannedPK:  map[string]*time.Time{},
		started:   time.Now(),
	}

	config := ssh.ServerConfig{
		NoClientAuth: false,
		// Auth-related things should be constant-time to avoid timing attacks.
		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
			fingerprint := Fingerprint(key)
			if server.IsBanned(fingerprint) {
				return nil, fmt.Errorf("Banned.")
			}
			if !server.IsWhitelisted(fingerprint) {
				return nil, fmt.Errorf("Not Whitelisted.")
			}
			perm := &ssh.Permissions{Extensions: map[string]string{"fingerprint": fingerprint}}
			return perm, nil
		},
	}
	config.AddHostKey(signer)

	server.sshConfig = &config

	return &server, nil
}

// Len returns the number of clients
func (s *Server) Len() int {
	return len(s.clients)
}

// SysMsg broadcasts the given message to everyone
func (s *Server) SysMsg(msg string, args ...interface{}) {
	s.Broadcast(ContinuousFormat(systemMessageFormat, " * "+fmt.Sprintf(msg, args...)), nil)
}

// Broadcast broadcasts the given message to everyone except for the given client
func (s *Server) Broadcast(msg string, except *Client) {
	logger.Debugf("Broadcast to %d: %s", s.Len(), msg)
	s.history.Add(msg)

	for _, client := range s.clients {
		if except != nil && client == except {
			continue
		}

		if strings.Contains(msg, client.Name) {
			// Turn message red if client's name is mentioned, and send BEL if they have enabled beeping
			personalMsg := strings.Replace(msg, client.Name, highlightFormat+client.Name+Reset, -1)
			if client.beepMe {
				personalMsg += beep
			}
			client.Send(personalMsg)
		} else {
			client.Send(msg)
		}
	}
}

// Privmsg sends a message to a particular nick, if it exists
func (s *Server) Privmsg(nick, message string, sender *Client) error {
	// Get the recipient
	target, ok := s.clients[strings.ToLower(nick)]
	if !ok {
		return fmt.Errorf("no client with that nick")
	}
	// Send the message
	target.Msg <- fmt.Sprintf(beep+"[PM from %v] %s%v%s", sender.ColoredName(), privateMessageFormat, message, Reset)
	logger.Debugf("PM from %v to %v: %v", sender.Name, nick, message)
	return nil
}

// SetMotd sets the Message of the Day (MOTD)
func (s *Server) SetMotd(motd string) {
	s.Lock()
	s.motd = motd
	s.Unlock()
}

// MotdUnicast sends the MOTD as a SysMsg
func (s *Server) MotdUnicast(client *Client) {
	if s.motd == "" {
		return
	}
	client.SysMsg(s.motd)
}

// MotdBroadcast broadcasts the MOTD
func (s *Server) MotdBroadcast(client *Client) {
	if s.motd == "" {
		return
	}
	s.Broadcast(ContinuousFormat(systemMessageFormat, fmt.Sprintf(" * New MOTD set by %s.", client.ColoredName())), client)
	s.Broadcast(s.motd, client)
}

// Add adds the client to the list of clients
func (s *Server) Add(client *Client) {
	go func() {
		s.MotdUnicast(client)
		client.SendLines(s.history.Get(10))
	}()

	s.Lock()
	s.count++

	newName, err := s.proposeName(client.Name)
	if err != nil {
		client.SysMsg("Your name '%s' is not available, renamed to '%s'. Use /nick <name> to change it.", client.Name, ColorString(client.Color, newName))
	}

	client.Rename(newName)
	s.clients[strings.ToLower(client.Name)] = client
	num := len(s.clients)
	s.Unlock()

	s.Broadcast(ContinuousFormat(systemMessageFormat, fmt.Sprintf(" * %s joined. (Total connected: %d)", client.Name, num)), client)
}

// Remove removes the given client from the list of clients
func (s *Server) Remove(client *Client) {
	s.Lock()
	delete(s.clients, strings.ToLower(client.Name))
	s.Unlock()

	s.SysMsg("%s left.", client.Name)
}

func (s *Server) proposeName(name string) (string, error) {
	// Assumes caller holds lock.
	var err error
	name = reStripText.ReplaceAllString(name, "")

	if len(name) > maxNameLength {
		name = name[:maxNameLength]
	} else if len(name) == 0 {
		name = fmt.Sprintf("Guest%d", s.count)
	}

	_, collision := s.clients[strings.ToLower(name)]
	if collision {
		err = fmt.Errorf("%s is not available", name)
		name = fmt.Sprintf("Guest%d", s.count)
	}

	return name, err
}

// Rename renames the given client (user)
func (s *Server) Rename(client *Client, newName string) {
	s.Lock()
	var oldName string
	if strings.ToLower(newName) == strings.ToLower(client.Name) {
		oldName = client.Name
		client.Rename(newName)
	} else {
		newName, err := s.proposeName(newName)
		if err != nil {
			client.SysMsg("%s", err)
			s.Unlock()
			return
		}

		// TODO: Use a channel/goroutine for adding clients, rather than locks?
		delete(s.clients, strings.ToLower(client.Name))
		oldName = client.Name
		client.Rename(newName)
		s.clients[strings.ToLower(client.Name)] = client
		s.Unlock()
	}
	s.SysMsg("%s is now known as %s.", ColorString(client.Color, oldName), ColorString(client.Color, newName))
}

// List lists the clients with the given prefix
func (s *Server) List(prefix *string) []string {
	r := []string{}

	for name := range s.clients {
		if prefix != nil && !strings.HasPrefix(name, strings.ToLower(*prefix)) {
			continue
		}
		r = append(r, name)
	}

	return r
}

// Who returns the client with a given name
func (s *Server) Who(name string) *Client {
	return s.clients[strings.ToLower(name)]
}

// Op adds the given fingerprint to the list of admins
func (s *Server) Op(fingerprint string) {
	logger.Infof("Adding admin: %s", fingerprint)
	s.Lock()
	s.admins[fingerprint] = struct{}{}
	s.Unlock()
}

// Whitelist adds the given fingerprint to the whitelist
func (s *Server) Whitelist(fingerprint string) {
	logger.Infof("Adding whitelist: %s", fingerprint)
	s.Lock()
	s.whitelist[fingerprint] = struct{}{}
	s.Unlock()
}

// Uptime returns the time since the server was started
func (s *Server) Uptime() string {
	return time.Now().Sub(s.started).String()
}

// IsOp checks if the given client is Op
func (s *Server) IsOp(client *Client) bool {
	_, r := s.admins[client.Fingerprint()]
	return r
}

// IsWhitelisted checks if the given fingerprint is whitelisted
func (s *Server) IsWhitelisted(fingerprint string) bool {
	/* if no whitelist, anyone is welcome */
	if len(s.whitelist) == 0 {
		return true
	}

	/* otherwise, check for whitelist presence */
	_, r := s.whitelist[fingerprint]
	return r
}

// IsBanned checks if the given fingerprint is banned
func (s *Server) IsBanned(fingerprint string) bool {
	ban, hasBan := s.bannedPK[fingerprint]
	if !hasBan {
		return false
	}
	if ban == nil {
		return true
	}
	if ban.Before(time.Now()) {
		s.Unban(fingerprint)
		return false
	}
	return true
}

// Ban bans a fingerprint for the given duration
func (s *Server) Ban(fingerprint string, duration *time.Duration) {
	var until *time.Time
	s.Lock()
	if duration != nil {
		when := time.Now().Add(*duration)
		until = &when
	}
	s.bannedPK[fingerprint] = until
	s.Unlock()
}

// Unban unbans a banned fingerprint
func (s *Server) Unban(fingerprint string) {
	s.Lock()
	delete(s.bannedPK, fingerprint)
	s.Unlock()
}

// Start starts the server
func (s *Server) Start(laddr string) error {
	// Once a ServerConfig has been configured, connections can be
	// accepted.
	socket, err := net.Listen("tcp", laddr)
	if err != nil {
		return err
	}

	logger.Infof("Listening on %s", laddr)

	go func() {
		defer socket.Close()
		for {
			conn, err := socket.Accept()

			if err != nil {
				logger.Errorf("Failed to accept connection: %v", err)
				if err == syscall.EINVAL {
					// TODO: Handle shutdown more gracefully?
					return
				}
			}

			// Goroutineify to resume accepting sockets early.
			go func() {
				// From a standard TCP connection to an encrypted SSH connection
				sshConn, channels, requests, err := ssh.NewServerConn(conn, s.sshConfig)
				if err != nil {
					logger.Errorf("Failed to handshake: %v", err)
					return
				}

				version := reStripText.ReplaceAllString(string(sshConn.ClientVersion()), "")
				if len(version) > 100 {
					version = "Evil Jerk with a superlong string"
				}
				logger.Infof("Connection #%d from: %s, %s, %s", s.count+1, sshConn.RemoteAddr(), sshConn.User(), version)

				go ssh.DiscardRequests(requests)

				client := NewClient(s, sshConn)
				go client.handleChannels(channels)
			}()
		}
	}()

	go func() {
		<-s.done
		socket.Close()
	}()

	return nil
}

// AutoCompleteFunction handles auto completion of nicks
func (s *Server) AutoCompleteFunction(line string, pos int, key rune) (newLine string, newPos int, ok bool) {
	if key == 9 {
		shortLine := strings.Split(line[:pos], " ")
		partialNick := shortLine[len(shortLine)-1]

		nicks := s.List(&partialNick)
		if len(nicks) > 0 {
			nick := nicks[len(nicks)-1]
			posPartialNick := pos - len(partialNick)
			if len(shortLine) < 2 {
				nick += ": "
			} else {
				nick += " "
			}
			newLine = strings.Replace(line[posPartialNick:],
				partialNick, nick, 1)
			newLine = line[:posPartialNick] + newLine
			newPos = pos + (len(nick) - len(partialNick))
			ok = true
		}
	} else {
		ok = false
	}
	return
}

// Stop stops the server
func (s *Server) Stop() {
	for _, client := range s.clients {
		client.Conn.Close()
	}

	close(s.done)
}

// Fingerprint returns the fingerprint based on a public key
func Fingerprint(k ssh.PublicKey) string {
	hash := md5.Sum(k.Marshal())
	r := fmt.Sprintf("% x", hash)
	return strings.Replace(r, " ", ":", -1)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/cenkalti/backoff"
	rfc3164 "github.com/jeromer/syslogparser/rfc3164"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	cfg       *config
	logger    = log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lshortfile)
	port      = flag.Int("port", 514, "port on which to listen")
	verbose   = flag.Bool("verbose", false, "print all messages to stdout")
	publish   = flag.Bool("publish", false, "publish messages to kafka")
	topic     = flag.String("topic", "syslog", "kafka topic to publish on")
	zkstring  = flag.String("zkstring", "localhost:2181", "ZooKeeper broker connection string")
	flushTime = flag.String("flushTime", "10s", "max time before flushing messages to kafka, e.g. 1s, 2m")
)

type config struct {
	port      int
	verbose   bool
	publish   bool
	topic     string
	zkstring  string
	flushTime time.Duration
}

func (c *config) String() string {
	return fmt.Sprintf("port: %d, verbose: %t, publish: %t, topic: %s, zkstring: %s, flushTime: %s",
		c.port, c.verbose, c.publish, c.topic, c.zkstring, c.flushTime)
}

type openConnection struct {
	connection net.Conn
	done       chan bool
}

type server struct {
	listener    net.Listener
	connections []*openConnection
	client      *sarama.Client
	producer    *sarama.Producer
	shutdown    chan bool
}

func (s *server) process(line []byte) {
	p := rfc3164.NewParser(line)
	if err := p.Parse(); err != nil {
		logger.Println("failed to parse:", err)
		return
	}

	parts := p.Dump()
	content := parts["content"].(string)

	if cfg.publish {
		if cfg.verbose {
			logger.Println("enqueuing", content)
		}

		put := func() error {
			err := s.producer.QueueMessage(cfg.topic, nil, sarama.StringEncoder(content))
			if err != nil {
				logger.Println("error queueing message, will retry", err)
			}
			return err
		}

		policy := backoff.NewExponentialBackOff()
		policy.MaxElapsedTime = time.Minute * 5
		err := backoff.Retry(put, policy)
		if err != nil {
			// retrying a bunch of times failed...
			logger.Println("failed sending message", err)
		}

	}

	if cfg.verbose {
		logger.Println(content)
	}
}

func (s *server) handleConnection(conn *openConnection) {
	defer conn.connection.Close()

	logger.Println("got connection from:", conn.connection.RemoteAddr())
	scanner := bufio.NewScanner(conn.connection)
	for scanner.Scan() {
		b := []byte(scanner.Text())
		logger.Printf("received %d bytes\n", len(b))
		s.process(b)
	}

	if err := scanner.Err(); err != nil {
		logger.Println("error reading from connection:", err)
	}

	logger.Println("exiting connection handler")
	conn.done <- true
}

func (s *server) stop() {
	s.shutdown <- true
	s.listener.Close()

	for _, conn := range s.connections {
		logger.Println("closing connection to", conn.connection.RemoteAddr())
		conn.connection.Close()
	}

	logger.Println("waiting for", len(s.connections), "connections to close")
	for _, conn := range s.connections {
		<-conn.done
	}
	logger.Println("all connections closed")

	if cfg.publish {
		s.producer.Close()
		s.client.Close()
		logger.Println("finished stopping producer")
	}
}

func (s *server) start() error {
	logger.Printf("starting to listen on %d\n", cfg.port)

	if cfg.verbose {
		logger.Println("verbose is enabled, all messages will be printed to stderr")
	}

	// listen for inbound syslog messages over tcp
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", cfg.port))
	if err != nil {
		return err
	}
	s.listener = listener

	if cfg.publish {
		// client for kafka
		brokers, err := LookupBrokers(cfg.zkstring)
		if err != nil {
			return err
		}

		brokerStr := make([]string, len(brokers))
		for i, b := range brokers {
			brokerStr[i] = fmt.Sprintf("%s:%d", b.Host, b.Port)
		}

		logger.Println("connecting to Kafka, using brokers from ZooKeeper:", brokerStr)
		client, err := sarama.NewClient("syslog", brokerStr, sarama.NewClientConfig())
		if err != nil {
			return err
		}
		s.client = client

		// producer for kafka
		producerConfig := sarama.NewProducerConfig()
		producerConfig.MaxBufferTime = cfg.flushTime

		producer, err := sarama.NewProducer(client, producerConfig)
		if err != nil {
			return err
		}
		s.producer = producer

		go func() {
			for err := range s.producer.Errors() {
				if err != nil {
					logger.Println(err)
				}
			}
		}()
	}

	connections := make(chan *openConnection)
	errors := make(chan error)

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				errors <- err
			} else {
				connection := &openConnection{
					connection: conn,
					done:       make(chan bool),
				}
				connections <- connection
			}
		}
	}()

	logger.Println("waiting for syslog connection")
	for {
		select {
		case conn := <-connections:
			s.connections = append(s.connections, conn)
			go s.handleConnection(conn)
		case err := <-errors:
			logger.Println("failed to accept connection:", err)
		case <-s.shutdown:
			goto exit
		}
	}

exit:
	logger.Println("exiting listen loop")
	return nil
}

func handleInterrupt(s *server) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGUSR1)

	go func() {
		for sig := range signals {
			switch sig {
			case syscall.SIGINT:
				logger.Println("got kill signal, closing", len(s.connections), "connections")
				s.stop()
				os.Exit(0)
				break
			case syscall.SIGUSR1:
				// TODO: refresh/reload things?
				logger.Println("refreshing kafka endpoints?")
				break
			}
		}
	}()
}

func main() {
	flag.Parse()

	ft, err := time.ParseDuration(*flushTime)
	if err != nil {
		logger.Fatalln(err)
	}

	cfg = &config{
		port:      *port,
		verbose:   *verbose,
		publish:   *publish,
		topic:     *topic,
		zkstring:  *zkstring,
		flushTime: ft,
	}

	logger.Println("starting with config:", cfg)

	server := &server{
		connections: []*openConnection{},
		shutdown:    make(chan bool),
	}

	handleInterrupt(server)

	if err := server.start(); err != nil {
		logger.Println("failed to listen:", err)
		os.Exit(1)
	}
}
<commit_msg>use synchronous send with backoff, will retry sending messages and not drop stuff<commit_after>package main

import (
	"bufio"
	"flag"
	"fmt"
	"github.com/Shopify/sarama"
	"github.com/cenkalti/backoff"
	rfc3164 "github.com/jeromer/syslogparser/rfc3164"
	"log"
	"net"
	"os"
	"os/signal"
	"syscall"
	"time"
)

var (
	cfg       *config
	logger    = log.New(os.Stderr, "", log.Ldate|log.Ltime|log.Lshortfile)
	port      = flag.Int("port", 514, "port on which to listen")
	verbose   = flag.Bool("verbose", false, "print all messages to stdout")
	publish   = flag.Bool("publish", false, "publish messages to kafka")
	topic     = flag.String("topic", "syslog", "kafka topic to publish on")
	zkstring  = flag.String("zkstring", "localhost:2181", "ZooKeeper broker connection string")
	flushTime = flag.String("flushTime", "10s", "max time before flushing messages to kafka, e.g. 1s, 2m")
)

type config struct {
	port      int
	verbose   bool
	publish   bool
	topic     string
	zkstring  string
	flushTime time.Duration
}

func (c *config) String() string {
	return fmt.Sprintf("port: %d, verbose: %t, publish: %t, topic: %s, zkstring: %s, flushTime: %s",
		c.port, c.verbose, c.publish, c.topic, c.zkstring, c.flushTime)
}

type openConnection struct {
	connection net.Conn
	done       chan bool
}

type server struct {
	listener    net.Listener
	connections []*openConnection
	client      *sarama.Client
	producer    *sarama.Producer
	shutdown    chan bool
}

func (s *server) process(line []byte) {
	p := rfc3164.NewParser(line)
	if err := p.Parse(); err != nil {
		logger.Println("failed to parse:", err)
		return
	}

	parts := p.Dump()
	content := parts["content"].(string)

	if cfg.publish {
		if cfg.verbose {
			logger.Println("enqueuing", content)
		}

		put := func() error {
			err := s.producer.SendMessage(cfg.topic, nil, sarama.StringEncoder(content))
			if err != nil {
				logger.Println("error queueing message, will retry.", err)
			}
			return err
		}

		policy := backoff.NewExponentialBackOff()
		policy.MaxElapsedTime = time.Minute * 5
		err := backoff.Retry(put, policy)
		if err != nil {
			// retrying a bunch of times failed...
			logger.Println("failed sending message.", err)
		}

	}

	if cfg.verbose {
		logger.Println(content)
	}
}

func (s *server) handleConnection(conn *openConnection) {
	defer conn.connection.Close()

	logger.Println("got connection from:", conn.connection.RemoteAddr())
	scanner := bufio.NewScanner(conn.connection)
	for scanner.Scan() {
		b := []byte(scanner.Text())
		s.process(b)
	}

	if err := scanner.Err(); err != nil {
		logger.Println("error reading from connection:", err)
	}

	logger.Println("exiting connection handler")
	conn.done <- true
}

func (s *server) stop() {
	s.shutdown <- true
	s.listener.Close()

	for _, conn := range s.connections {
		logger.Println("closing connection to", conn.connection.RemoteAddr())
		conn.connection.Close()
	}

	logger.Println("waiting for", len(s.connections), "connections to close")
	for _, conn := range s.connections {
		<-conn.done
	}
	logger.Println("all connections closed")

	if cfg.publish {
		s.producer.Close()
		s.client.Close()
		logger.Println("finished stopping producer")
	}
}

func (s *server) start() error {
	logger.Printf("starting to listen on %d\n", cfg.port)

	if cfg.verbose {
		logger.Println("verbose is enabled, all messages will be printed to stderr")
	}

	// listen for inbound syslog messages over tcp
	listener, err := net.Listen("tcp", fmt.Sprintf(":%d", cfg.port))
	if err != nil {
		return err
	}
	s.listener = listener

	if cfg.publish {
		// client for kafka
		brokers, err := LookupBrokers(cfg.zkstring)
		if err != nil {
			return err
		}

		brokerStr := make([]string, len(brokers))
		for i, b := range brokers {
			brokerStr[i] = fmt.Sprintf("%s:%d", b.Host, b.Port)
		}

		logger.Println("connecting to Kafka, using brokers from ZooKeeper:", brokerStr)
		client, err := sarama.NewClient("syslog", brokerStr, sarama.NewClientConfig())
		if err != nil {
			return err
		}
		s.client = client

		// producer for kafka
		producerConfig := sarama.NewProducerConfig()
		producerConfig.MaxBufferTime = cfg.flushTime

		producer, err := sarama.NewProducer(client, producerConfig)
		if err != nil {
			return err
		}
		s.producer = producer
	}

	connections := make(chan *openConnection)
	errors := make(chan error)

	go func() {
		for {
			conn, err := listener.Accept()
			if err != nil {
				errors <- err
			} else {
				connection := &openConnection{
					connection: conn,
					done:       make(chan bool),
				}
				connections <- connection
			}
		}
	}()

	logger.Println("waiting for syslog connection")
	for {
		select {
		case conn := <-connections:
			s.connections = append(s.connections, conn)
			go s.handleConnection(conn)
		case err := <-errors:
			logger.Println("failed to accept connection:", err)
		case <-s.shutdown:
			goto exit
		}
	}

exit:
	logger.Println("exiting listen loop")
	return nil
}

func handleInterrupt(s *server) {
	signals := make(chan os.Signal, 1)
	signal.Notify(signals, syscall.SIGINT, syscall.SIGUSR1)

	go func() {
		for sig := range signals {
			switch sig {
			case syscall.SIGINT:
				logger.Println("got kill signal, closing", len(s.connections), "connections")
				s.stop()
				os.Exit(0)
				break
			case syscall.SIGUSR1:
				// TODO: refresh/reload things?
				logger.Println("refreshing kafka endpoints?")
				break
			}
		}
	}()
}

func main() {
	flag.Parse()

	ft, err := time.ParseDuration(*flushTime)
	if err != nil {
		logger.Fatalln(err)
	}

	cfg = &config{
		port:      *port,
		verbose:   *verbose,
		publish:   *publish,
		topic:     *topic,
		zkstring:  *zkstring,
		flushTime: ft,
	}

	logger.Println("starting with config:", cfg)

	server := &server{
		connections: []*openConnection{},
		shutdown:    make(chan bool),
	}

	handleInterrupt(server)

	if err := server.start(); err != nil {
		logger.Println("failed to listen:", err)
		os.Exit(1)
	}
}
<|endoftext|>


<commit_before>package inigo_test

import (
	"fmt"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	"strings"

	"github.com/cloudfoundry-incubator/inigo/inigolistener"
	"github.com/cloudfoundry-incubator/inigo/stager_runner"
	"github.com/cloudfoundry-incubator/inigo/zipper"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Stager", func() {
	var otherStagerRunner *stager_runner.StagerRunner

	BeforeEach(func() {
		otherStagerRunner = stager_runner.New(
			stagerPath,
			etcdRunner.NodeURLS(),
			[]string{fmt.Sprintf("127.0.0.1:%d", natsPort)},
		)
	})

	Context("when unable to find an appropriate compiler", func() {
		BeforeEach(func() {
			executorRunner.Start()
			stagerRunner.Start()
		})

		It("returns an error", func() {
			errorMessages := 0
			natsRunner.MessageBus.Subscribe("compiler-stagers-test", func(message *yagnats.Message) {
				if strings.Contains(string(message.Payload), "error") {
					errorMessages++
				}
			})

			natsRunner.MessageBus.PublishWithReplyTo(
				"diego.staging.start",
				"compiler-stagers-test",
				[]byte(`{"app_id": "some-app-guid", "task_id": "some-task-id", "stack": "no-compiler"}`))

			Eventually(func() int {
				return errorMessages
			}, 2.0).Should(Equal(1))

			Consistently(func() int {
				return errorMessages
			}, 2.0).Should(Equal(1))
		})
	})

	Describe("Staging", func() {
		var compilerGuid string
		var appGuid string
		var adminBuildpackGuid string
		var outputGuid string
		var stagingMessage []byte

		BeforeEach(func() {
			fileServerRunner.Start()
			executorRunner.Start()

			compilerGuid = factories.GenerateGuid()
			appGuid = factories.GenerateGuid()
			adminBuildpackGuid = factories.GenerateGuid()
			outputGuid = factories.GenerateGuid()

			//make and provide a compiler via the file-server
			var compilerFiles = []zipper.ZipFile{
				{"run", fmt.Sprintf(`echo $FOO && %s && $APP_DIR/run && $BUILDPACKS_DIR/test/run`,
					inigolistener.CurlCommand(compilerGuid),
				)},
			}
			zipper.CreateZipFile("/tmp/compiler.zip", compilerFiles)
			fileServerRunner.ServeFile("compiler.zip", "/tmp/compiler.zip")

			//make and upload an app
			var appFiles = []zipper.ZipFile{
				{"run", inigolistener.CurlCommand(appGuid)},
			}
			zipper.CreateZipFile("/tmp/app.zip", appFiles)
			inigolistener.UploadFile("app.zip", "/tmp/app.zip")

			//make and upload a buildpack
			var adminBuildpackFiles = []zipper.ZipFile{
				{"run", inigolistener.CurlCommand(adminBuildpackGuid)},
			}
			zipper.CreateZipFile("/tmp/admin_buildpack.zip", adminBuildpackFiles)
			inigolistener.UploadFile("admin_buildpack.zip", "/tmp/admin_buildpack.zip")

			stagingMessage = []byte(
				fmt.Sprintf(
					`{
						"app_id": "some-app-guid",
						"task_id": "some-task-id",
						"stack": "default",
						"download_uri": "%s",
						"admin_buildpacks" : [{ "key": "test", "url": "%s" }],
						"environment": [["FOO", "%s"]]
					}`,
					inigolistener.DownloadUrl("app.zip"),
					inigolistener.DownloadUrl("admin_buildpack.zip"),
					outputGuid,
				),
			)
		})

		Context("with one stager running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"compiler.zip"}`)
			})

			It("runs the compiler on executor with the correct environment variables, bits and log tag", func(done Done) {
				messages, stop := loggredile.StreamMessages(
					loggregatorRunner.Config.OutgoingPort,
					"/tail/?app=some-app-guid",
				)

				err := natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"stager-test",
					stagingMessage,
				)
				Ω(err).ShouldNot(HaveOccurred())
				Eventually(inigolistener.ReportingGuids, 10.0).Should(ContainElement(compilerGuid))
				Eventually(inigolistener.ReportingGuids, 5.0).Should(ContainElement(appGuid))
				Eventually(inigolistener.ReportingGuids, 5.0).Should(ContainElement(adminBuildpackGuid))

				message := <-messages
				Ω(message.GetSourceName()).To(Equal("STG"))
				Ω(string(message.GetMessage())).To(Equal(outputGuid))

				close(stop)
				close(done)
			}, 20)
		})

		Context("with two stagers running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"compiler.zip"}`)
				otherStagerRunner.Start("--compilers", `{"default":"compiler.zip"}`)
			})

			AfterEach(func() {
				otherStagerRunner.Stop()
			})

			It("only one returns a staging completed response", func() {
				messages := 0
				natsRunner.MessageBus.Subscribe("two-stagers-test", func(message *yagnats.Message) {
					messages++
				})

				natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"two-stagers-test",
					stagingMessage,
				)

				Eventually(func() int {
					return messages
				}, 10.0).Should(Equal(1))

				Consistently(func() int {
					return messages
				}, 2.0).Should(Equal(1))
			})
		})
	})
})
<commit_msg>add test for stager response, add timeouts<commit_after>package inigo_test

import (
	"fmt"
	"github.com/cloudfoundry-incubator/inigo/loggredile"
	"strings"

	"github.com/cloudfoundry-incubator/inigo/inigolistener"
	"github.com/cloudfoundry-incubator/inigo/stager_runner"
	"github.com/cloudfoundry-incubator/inigo/zipper"
	"github.com/cloudfoundry-incubator/runtime-schema/models/factories"
	"github.com/cloudfoundry/yagnats"
	. "github.com/onsi/ginkgo"
	. "github.com/onsi/gomega"
)

var _ = Describe("Stager", func() {
	var otherStagerRunner *stager_runner.StagerRunner

	BeforeEach(func() {
		otherStagerRunner = stager_runner.New(
			stagerPath,
			etcdRunner.NodeURLS(),
			[]string{fmt.Sprintf("127.0.0.1:%d", natsPort)},
		)
	})

	Context("when unable to find an appropriate compiler", func() {
		BeforeEach(func() {
			executorRunner.Start()
			stagerRunner.Start()
		})

		It("returns an error", func(done Done) {
			errorMessages := 0
			natsRunner.MessageBus.Subscribe("compiler-stagers-test", func(message *yagnats.Message) {
				if strings.Contains(string(message.Payload), "error") {
					errorMessages++
				}
			})

			natsRunner.MessageBus.PublishWithReplyTo(
				"diego.staging.start",
				"compiler-stagers-test",
				[]byte(`{"app_id": "some-app-guid", "task_id": "some-task-id", "stack": "no-compiler"}`))

			Eventually(func() int {
				return errorMessages
			}, 2.0).Should(Equal(1))

			Consistently(func() int {
				return errorMessages
			}, 2.0).Should(Equal(1))
		}, 10.0)
	})

	Describe("Staging", func() {
		var compilerGuid string
		var appGuid string
		var adminBuildpackGuid string
		var outputGuid string
		var stagingMessage []byte
		var compilerExitStatus int

		BeforeEach(func() {
			fileServerRunner.Start()
			executorRunner.Start()

			compilerGuid = factories.GenerateGuid()
			appGuid = factories.GenerateGuid()
			adminBuildpackGuid = factories.GenerateGuid()
			outputGuid = factories.GenerateGuid()

			//make and provide a compiler via the file-server
			var compilerFiles = []zipper.ZipFile{
				{"run", fmt.Sprintf(`#!/bin/bash

set -e

# use an environment variable provided for staging
echo $SOME_STAGING_ENV

# tell the listener we're here
%s

# tell the listener the app's bits are here
$APP_DIR/run

# tell the listener the buildpacks are here
$BUILDPACKS_DIR/test/run

# write a staging result
mkdir -p $RESULT_DIR
echo '%s' > $RESULT_DIR/result.json

# inject success/failure
exit $COMPILER_EXIT_STATUS
`,
					inigolistener.CurlCommand(compilerGuid),
					`{"detected_buildpack":"Test Buildpack"}`,
				)},
			}
			zipper.CreateZipFile("/tmp/compiler.zip", compilerFiles)
			fileServerRunner.ServeFile("compiler.zip", "/tmp/compiler.zip")

			//make and upload an app
			var appFiles = []zipper.ZipFile{
				{"run", inigolistener.CurlCommand(appGuid)},
			}
			zipper.CreateZipFile("/tmp/app.zip", appFiles)
			inigolistener.UploadFile("app.zip", "/tmp/app.zip")

			//make and upload a buildpack
			var adminBuildpackFiles = []zipper.ZipFile{
				{"run", inigolistener.CurlCommand(adminBuildpackGuid)},
			}
			zipper.CreateZipFile("/tmp/admin_buildpack.zip", adminBuildpackFiles)
			inigolistener.UploadFile("admin_buildpack.zip", "/tmp/admin_buildpack.zip")

		})

		JustBeforeEach(func() {
			stagingMessage = []byte(
				fmt.Sprintf(
					`{
						"app_id": "some-app-guid",
						"task_id": "some-task-id",
						"stack": "default",
						"download_uri": "%s",
						"admin_buildpacks" : [{ "key": "test", "url": "%s" }],
						"environment": [["SOME_STAGING_ENV", "%s"], ["COMPILER_EXIT_STATUS", "%d"]]
					}`,
					inigolistener.DownloadUrl("app.zip"),
					inigolistener.DownloadUrl("admin_buildpack.zip"),
					outputGuid,
					compilerExitStatus,
				),
			)
		})

		Context("with one stager running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"compiler.zip"}`)
			})

			It("runs the compiler on the executor with the correct environment variables, bits and log tag, and responds with the detected buildpack", func(done Done) {
				defer close(done)

				payloads := make(chan []byte)

				natsRunner.MessageBus.Subscribe("stager-test", func(msg *yagnats.Message) {
					payloads <- msg.Payload
				})

				messages, stop := loggredile.StreamMessages(
					loggregatorRunner.Config.OutgoingPort,
					"/tail/?app=some-app-guid",
				)
				defer close(stop)

				err := natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"stager-test",
					stagingMessage,
				)
				Ω(err).ShouldNot(HaveOccurred())
				Eventually(inigolistener.ReportingGuids, 10.0).Should(ContainElement(compilerGuid))
				Eventually(inigolistener.ReportingGuids, 5.0).Should(ContainElement(appGuid))
				Eventually(inigolistener.ReportingGuids, 5.0).Should(ContainElement(adminBuildpackGuid))

				message := <-messages
				Ω(message.GetSourceName()).To(Equal("STG"))
				Ω(string(message.GetMessage())).To(Equal(outputGuid))

				payload := <-payloads
				Ω(string(payload)).Should(Equal(`{"detected_buildpack":"Test Buildpack"}`))
			}, 20.0)

			Context("when compilation fails", func() {
				BeforeEach(func() {
					compilerExitStatus = 42
				})

				It("responds with the error, and no detected buildpack present", func(done Done) {
					defer close(done)

					payloads := make(chan []byte)

					natsRunner.MessageBus.Subscribe("stager-test", func(msg *yagnats.Message) {
						payloads <- msg.Payload
					})

					messages, stop := loggredile.StreamMessages(
						loggregatorRunner.Config.OutgoingPort,
						"/tail/?app=some-app-guid",
					)
					defer close(stop)

					err := natsRunner.MessageBus.PublishWithReplyTo(
						"diego.staging.start",
						"stager-test",
						stagingMessage,
					)
					Ω(err).ShouldNot(HaveOccurred())
					Eventually(inigolistener.ReportingGuids, 10.0).Should(ContainElement(compilerGuid))
					Eventually(inigolistener.ReportingGuids, 5.0).Should(ContainElement(appGuid))
					Eventually(inigolistener.ReportingGuids, 5.0).Should(ContainElement(adminBuildpackGuid))

					message := <-messages
					Ω(message.GetSourceName()).To(Equal("STG"))
					Ω(string(message.GetMessage())).To(Equal(outputGuid))

					payload := <-payloads
					Ω(string(payload)).Should(Equal(`{"error":"Process returned with exit value: 42"}`))
				}, 20.0)
			})
		})

		Context("with two stagers running", func() {
			BeforeEach(func() {
				stagerRunner.Start("--compilers", `{"default":"compiler.zip"}`)
				otherStagerRunner.Start("--compilers", `{"default":"compiler.zip"}`)
			})

			AfterEach(func() {
				otherStagerRunner.Stop()
			})

			It("only one returns a staging completed response", func(done Done) {
				defer close(done)

				messages := 0
				natsRunner.MessageBus.Subscribe("two-stagers-test", func(message *yagnats.Message) {
					messages++
				})

				natsRunner.MessageBus.PublishWithReplyTo(
					"diego.staging.start",
					"two-stagers-test",
					stagingMessage,
				)

				Eventually(func() int {
					return messages
				}, 10.0).Should(Equal(1))

				Consistently(func() int {
					return messages
				}, 2.0).Should(Equal(1))
			}, 10.0)
		})
	})
})
<|endoftext|>