<commit_before>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"

	// auth
	"github.com/SpectoLabs/hoverfly/authentication"
	"github.com/SpectoLabs/hoverfly/authentication/controllers"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	go func() {
		// starting admin interface
		mux := getBoneRouter(*d)
		n := negroni.Classic()

		logLevel := log.ErrorLevel

		if d.Cfg.Verbose {
			logLevel = log.DebugLevel
		}

		n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
		n.UseHandler(mux)

		// admin interface starting message
		log.WithFields(log.Fields{
			"AdminPort": d.Cfg.AdminPort,
		}).Info("Admin interface is starting...")

		n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
	}()
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// getting auth controllers and middleware
	ac := controllers.GetNewAuthenticationController(d.AB, d.Cfg.SecretKey, d.Cfg.JWTExpirationDelta)
	am := authentication.GetNewAuthenticationMiddleware(d.AB,
		d.Cfg.SecretKey,
		d.Cfg.JWTExpirationDelta,
		d.Cfg.AuthEnabled)

	mux.Post("/token-auth", http.HandlerFunc(ac.Login))
	mux.Get("/refresh-token-auth", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.RefreshToken),
	))
	mux.Get("/logout", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.Logout),
	))

	mux.Get("/users", http.HandlerFunc(ac.GetAllUsersHandler))
	// TODO: add users delete/add functionality

	mux.Get("/records", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(d.AllRecordsHandler),
	))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		// since hoverfly is not started from cmd/hoverfly/hoverfly
		// we have to target to that directory
		log.Warn("Hoverfly is serving files from /static/dist instead of statik binary!")
		mux.Handle("/*", http.FileServer(http.Dir("../../static/dist")))
	} else {
		// preparing static assets for embedded admin
		statikFS, err := fs.New()

		if err != nil {
			log.WithFields(log.Fields{
				"Error": err.Error(),
			}).Error("Failed to load statikFS, admin UI might not work :(")
		}

		mux.Handle("/*", http.FileServer(statikFS))
	}
	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Debug("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	err = d.ImportPayloads(requests.Data)

	if err != nil {
		response.Message = err.Error()
		w.WriteHeader(400)
	} else {
		response.Message = fmt.Sprintf("%d payloads import complete.", len(requests.Data))
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteData()

	var en Entry
	en.ActionType = ActionTypeWipeDB
	en.Message = "wipe"
	en.Time = time.Now()

	if err := d.Hooks.Fire(ActionTypeWipeDB, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeWipeDB,
		}).Error("failed to fire hook")
	}

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var en Entry
	en.ActionType = ActionTypeConfigurationChanged
	en.Message = "changed"
	en.Time = time.Now()
	en.Data = []byte("sr.Mode")

	if err := d.Hooks.Fire(ActionTypeConfigurationChanged, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeConfigurationChanged,
		}).Error("failed to fire hook")
	}

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<commit_msg>users now require authentication<commit_after>package hoverfly

import (
	"bytes"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net/http"
	"time"

	// static assets
	_ "github.com/SpectoLabs/hoverfly/statik"
	"github.com/rakyll/statik/fs"

	log "github.com/Sirupsen/logrus"
	"github.com/codegangsta/negroni"
	"github.com/go-zoo/bone"
	"github.com/gorilla/websocket"
	"github.com/meatballhat/negroni-logrus"

	// auth
	"github.com/SpectoLabs/hoverfly/authentication"
	"github.com/SpectoLabs/hoverfly/authentication/controllers"
)

// recordedRequests struct encapsulates payload data
type recordedRequests struct {
	Data []Payload `json:"data"`
}

type recordsCount struct {
	Count int `json:"count"`
}

type statsResponse struct {
	Stats        Stats `json:"stats"`
	RecordsCount int   `json:"recordsCount"`
}

type stateRequest struct {
	Mode        string `json:"mode"`
	Destination string `json:"destination"`
}

type messageResponse struct {
	Message string `json:"message"`
}

// StartAdminInterface - starts admin interface web server
func (d *DBClient) StartAdminInterface() {
	go func() {
		// starting admin interface
		mux := getBoneRouter(*d)
		n := negroni.Classic()

		logLevel := log.ErrorLevel

		if d.Cfg.Verbose {
			logLevel = log.DebugLevel
		}

		n.Use(negronilogrus.NewCustomMiddleware(logLevel, &log.JSONFormatter{}, "admin"))
		n.UseHandler(mux)

		// admin interface starting message
		log.WithFields(log.Fields{
			"AdminPort": d.Cfg.AdminPort,
		}).Info("Admin interface is starting...")

		n.Run(fmt.Sprintf(":%s", d.Cfg.AdminPort))
	}()
}

// getBoneRouter returns mux for admin interface
func getBoneRouter(d DBClient) *bone.Mux {
	mux := bone.New()

	// getting auth controllers and middleware
	ac := controllers.GetNewAuthenticationController(d.AB, d.Cfg.SecretKey, d.Cfg.JWTExpirationDelta)
	am := authentication.GetNewAuthenticationMiddleware(d.AB,
		d.Cfg.SecretKey,
		d.Cfg.JWTExpirationDelta,
		d.Cfg.AuthEnabled)

	mux.Post("/token-auth", http.HandlerFunc(ac.Login))
	mux.Get("/refresh-token-auth", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.RefreshToken),
	))
	mux.Get("/logout", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.Logout),
	))

	mux.Get("/users", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(ac.GetAllUsersHandler),
	))

	mux.Get("/records", negroni.New(
		negroni.HandlerFunc(am.RequireTokenAuthentication),
		negroni.HandlerFunc(d.AllRecordsHandler),
	))
	mux.Delete("/records", http.HandlerFunc(d.DeleteAllRecordsHandler))
	mux.Post("/records", http.HandlerFunc(d.ImportRecordsHandler))

	mux.Get("/count", http.HandlerFunc(d.RecordsCount))
	mux.Get("/stats", http.HandlerFunc(d.StatsHandler))
	mux.Get("/statsws", http.HandlerFunc(d.StatsWSHandler))

	mux.Get("/state", http.HandlerFunc(d.CurrentStateHandler))
	mux.Post("/state", http.HandlerFunc(d.StateHandler))

	if d.Cfg.Development {
		// since hoverfly is not started from cmd/hoverfly/hoverfly
		// we have to target to that directory
		log.Warn("Hoverfly is serving files from /static/dist instead of statik binary!")
		mux.Handle("/*", http.FileServer(http.Dir("../../static/dist")))
	} else {
		// preparing static assets for embedded admin
		statikFS, err := fs.New()

		if err != nil {
			log.WithFields(log.Fields{
				"Error": err.Error(),
			}).Error("Failed to load statikFS, admin UI might not work :(")
		}

		mux.Handle("/*", http.FileServer(statikFS))
	}
	return mux
}

// AllRecordsHandler returns JSON content type http response
func (d *DBClient) AllRecordsHandler(w http.ResponseWriter, req *http.Request, next http.HandlerFunc) {
	records, err := d.Cache.GetAllRequests()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordedRequests
		response.Data = records
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// RecordsCount returns number of captured requests as a JSON payload
func (d *DBClient) RecordsCount(w http.ResponseWriter, req *http.Request) {
	count, err := d.Cache.RecordsCount()

	if err == nil {

		w.Header().Set("Content-Type", "application/json")

		var response recordsCount
		response.Count = count
		b, err := json.Marshal(response)

		if err != nil {
			log.Error(err)
			http.Error(w, err.Error(), http.StatusInternalServerError)
		} else {
			w.Write(b)
			return
		}
	} else {
		log.WithFields(log.Fields{
			"Error": err.Error(),
		}).Error("Failed to get data from cache!")

		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(500) // can't process this entity
		return
	}
}

// StatsHandler - returns current stats about Hoverfly (request counts, record count)
func (d *DBClient) StatsHandler(w http.ResponseWriter, req *http.Request) {
	stats := d.Counter.Flush()

	count, err := d.Cache.RecordsCount()

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	}

	var sr statsResponse
	sr.Stats = stats
	sr.RecordsCount = count

	w.Header().Set("Content-Type", "application/json")

	b, err := json.Marshal(sr)

	if err != nil {
		log.Error(err)
		http.Error(w, err.Error(), http.StatusInternalServerError)
	} else {
		w.Write(b)
		return
	}

}

var upgrader = websocket.Upgrader{
	ReadBufferSize:  1024,
	WriteBufferSize: 1024,
	CheckOrigin: func(r *http.Request) bool {
		return true
	},
}

// StatsWSHandler - returns current stats about Hoverfly (request counts, record count) through the websocket
func (d *DBClient) StatsWSHandler(w http.ResponseWriter, r *http.Request) {
	conn, err := upgrader.Upgrade(w, r, nil)
	if err != nil {
		log.Println(err)
		return
	}

	for {
		messageType, p, err := conn.ReadMessage()
		if err != nil {
			return
		}
		log.WithFields(log.Fields{
			"message": string(p),
		}).Info("Got message...")

		for _ = range time.Tick(1 * time.Second) {

			count, err := d.Cache.RecordsCount()

			if err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Error("got error while trying to get records count")
				return
			}

			stats := d.Counter.Flush()

			var sr statsResponse
			sr.Stats = stats
			sr.RecordsCount = count

			b, err := json.Marshal(sr)

			if err = conn.WriteMessage(messageType, b); err != nil {
				log.WithFields(log.Fields{
					"message": p,
					"error":   err.Error(),
				}).Debug("Got error when writing message...")
				return
			}
		}

	}

}

// ImportRecordsHandler - accepts JSON payload and saves it to cache
func (d *DBClient) ImportRecordsHandler(w http.ResponseWriter, req *http.Request) {

	var requests recordedRequests

	defer req.Body.Close()
	body, err := ioutil.ReadAll(req.Body)

	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	var response messageResponse

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		response.Message = "Bad request. Nothing to import!"
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &requests)

	if err != nil {
		w.WriteHeader(422) // can't process this entity
		return
	}

	err = d.ImportPayloads(requests.Data)

	if err != nil {
		response.Message = err.Error()
		w.WriteHeader(400)
	} else {
		response.Message = fmt.Sprintf("%d payloads import complete.", len(requests.Data))
	}

	b, err := json.Marshal(response)
	w.Write(b)

}

// DeleteAllRecordsHandler - deletes all captured requests
func (d *DBClient) DeleteAllRecordsHandler(w http.ResponseWriter, req *http.Request) {
	err := d.Cache.DeleteData()

	var en Entry
	en.ActionType = ActionTypeWipeDB
	en.Message = "wipe"
	en.Time = time.Now()

	if err := d.Hooks.Fire(ActionTypeWipeDB, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeWipeDB,
		}).Error("failed to fire hook")
	}

	w.Header().Set("Content-Type", "application/json")

	var response messageResponse
	if err != nil {
		if err.Error() == "bucket not found" {
			response.Message = fmt.Sprintf("No records found")
			w.WriteHeader(200)
		} else {
			response.Message = fmt.Sprintf("Something went wrong: %s", err.Error())
			w.WriteHeader(500)
		}
	} else {
		response.Message = "Proxy cache deleted successfuly"
		w.WriteHeader(200)
	}
	b, err := json.Marshal(response)

	w.Write(b)
	return
}

// CurrentStateHandler returns current state
func (d *DBClient) CurrentStateHandler(w http.ResponseWriter, req *http.Request) {
	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination

	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)
}

// StateHandler handles current proxy state
func (d *DBClient) StateHandler(w http.ResponseWriter, r *http.Request) {
	var sr stateRequest

	// this is mainly for testing, since when you create
	if r.Body == nil {
		r.Body = ioutil.NopCloser(bytes.NewBuffer([]byte("")))
	}

	defer r.Body.Close()
	body, err := ioutil.ReadAll(r.Body)

	if err != nil {
		// failed to read response body
		log.WithFields(log.Fields{
			"error": err.Error(),
		}).Error("Could not read response body!")
		http.Error(w, "Failed to read request body.", 400)
		return
	}

	err = json.Unmarshal(body, &sr)

	if err != nil {
		w.Header().Set("Content-Type", "application/json; charset=UTF-8")
		w.WriteHeader(400) // can't process this entity
		return
	}

	availableModes := map[string]bool{
		"virtualize": true,
		"capture":    true,
		"modify":     true,
		"synthesize": true,
	}

	if !availableModes[sr.Mode] {
		log.WithFields(log.Fields{
			"suppliedMode": sr.Mode,
		}).Error("Wrong mode found, can't change state")
		http.Error(w, "Bad mode supplied, available modes: virtualize, capture, modify, synthesize.", 400)
		return
	}

	log.WithFields(log.Fields{
		"newState": sr.Mode,
		"body":     string(body),
	}).Info("Handling state change request!")

	// setting new state
	d.Cfg.SetMode(sr.Mode)

	var en Entry
	en.ActionType = ActionTypeConfigurationChanged
	en.Message = "changed"
	en.Time = time.Now()
	en.Data = []byte("sr.Mode")

	if err := d.Hooks.Fire(ActionTypeConfigurationChanged, &en); err != nil {
		log.WithFields(log.Fields{
			"error":      err.Error(),
			"message":    en.Message,
			"actionType": ActionTypeConfigurationChanged,
		}).Error("failed to fire hook")
	}

	var resp stateRequest
	resp.Mode = d.Cfg.GetMode()
	resp.Destination = d.Cfg.Destination
	b, _ := json.Marshal(resp)
	w.Header().Set("Content-Type", "application/json; charset=UTF-8")
	w.Write(b)

}
<|endoftext|>


<commit_before>package gorelic

import (
//	"encoding/json"
//	"errors"
//	"fmt"
//	"io/ioutil"
//	"net/http"
//	"net/url"
//	"strings"
)

type AgentSettings struct {
}

type environmentAttribute []interface{}
type AgentEnvironment []environmentAttribute

func NewAgentEnvironment() *AgentEnvironment {
    //TODO:  ["Plugin List", []]

    env := &AgentEnvironment{
        environmentAttribute{"Agent Version", "1.10.2.38"},
        environmentAttribute{"Arch", "x86_64"},
        environmentAttribute{"OS", "Linux"},
        environmentAttribute{"OS version", "3.2.0-24-generic"},
        environmentAttribute{"CPU Count", "1"},
        environmentAttribute{"System Memory", "2003.6328125"},
        environmentAttribute{"Python Program Name", "/usr/local/bin/newrelic-admin"},
        environmentAttribute{"Python Executable", "/usr/bin/python"},
        environmentAttribute{"Python Home", ""},
        environmentAttribute{"Python Path", ""},
        environmentAttribute{"Python Prefix", "/usr"},
        environmentAttribute{"Python Exec Prefix", "/usr"},
        environmentAttribute{"Python Version", "2.7.3 (default, Apr 20 2012, 22:39:59) \n[GCC 4.6.3]"},
        environmentAttribute{"Python Platform", "linux2"},
        environmentAttribute{"Python Max Unicode", "1114111"},
        environmentAttribute{"Compiled Extensions", ""},
    }
    return env
}

type Agent struct {
    AppName []string `json:"app_name"`
    Language string `json:"language"`
    Settings *AgentSettings `json:"settings"`
    Pid      int `json:"pid"`
    Environment *AgentEnvironment `json:"environment"`
    Host  string `json:"host"`
    Identifier  string `json:"identifier"`
    AgentVersion string `json:"agent_version"`
}

func NewAgent() *Agent {
    a := &Agent{
        AppName: []string{"Python Agent Test"},
        Language: "python",
        Identifier: "Python Agent Test",
        AgentVersion: "1.10.2.38",
        Environment: NewAgentEnvironment(),
    }
    return a
}

type Settings struct {
}

<commit_msg>Continue working on agent settings<commit_after>package gorelic

import (
//	"encoding/json"
//	"errors"
//	"fmt"
//	"io/ioutil"
//	"net/http"
//	"net/url"
//	"strings"
)

type AgentSettings struct {
}

type environmentAttribute []interface{}
type AgentEnvironment []environmentAttribute

func NewAgentEnvironment() *AgentEnvironment {
    //TODO:  ["Plugin List", []]

    env := &AgentEnvironment{
        environmentAttribute{"Agent Version", "1.10.2.38"},
        environmentAttribute{"Arch", "x86_64"},
        environmentAttribute{"OS", "Linux"},
        environmentAttribute{"OS version", "3.2.0-24-generic"},
        environmentAttribute{"CPU Count", "1"},
        environmentAttribute{"System Memory", "2003.6328125"},
        environmentAttribute{"Python Program Name", "/usr/local/bin/newrelic-admin"},
        environmentAttribute{"Python Executable", "/usr/bin/python"},
        environmentAttribute{"Python Home", ""},
        environmentAttribute{"Python Path", ""},
        environmentAttribute{"Python Prefix", "/usr"},
        environmentAttribute{"Python Exec Prefix", "/usr"},
        environmentAttribute{"Python Version", "2.7.3 (default, Apr 20 2012, 22:39:59) \n[GCC 4.6.3]"},
        environmentAttribute{"Python Platform", "linux2"},
        environmentAttribute{"Python Max Unicode", "1114111"},
        environmentAttribute{"Compiled Extensions", ""},
    }
    return env
}

type Agent struct {
    AppName []string `json:"app_name"`
    Language string `json:"language"`
    Settings *AgentSettings `json:"settings"`
    Pid      int `json:"pid"`
    Environment *AgentEnvironment `json:"environment"`
    Host  string `json:"host"`
    Identifier  string `json:"identifier"`
    AgentVersion string `json:"agent_version"`
}

func NewAgent() *Agent {
    a := &Agent{
        AppName: []string{"Python Agent Test"},
        Language: "python",
        Identifier: "Python Agent Test",
        AgentVersion: "1.10.2.38",
        Environment: NewAgentEnvironment(),
    }
    return a
}

type AgentSettings struct {
    StartupTimeout float `json:"startup_timeout"`
    DebugLogDataCollectorCalls bool `json:"debug.log_data_collector_calls"`
    EncodingKey string `json:"encoding_key"`
    ApplicationId string `json:"application_id"`
    ThreadProfilerEnabled bool `json:"thread_profiler.enabled"`
    ErrorCollectorCaptureSource bool `json:"error_collector.capture_source"`
    CaptureParams bool `json:"capture_params"`
    AgentLimitsSqlQueryLengthMaximum int `json:"agent_limits.sql_query_length_maximum"`
    ProxyPort int `json:"proxy_port"`
    IncludeEnviron []string `json:"include_environ"`
    TransactionNameLimit int `json:"transaction_name.limit"`
    BrowserKey string `json:"browser_key"`
    DebugLogTransactionTracePayload bool `json:"debug.log_transaction_trace_payload"`
    ShutdownTimeout float `json:"shutdown_timeout"`
    TrustedAccountIds []int `json:"trusted_account_ids"`
    WebTransactionsApdex interface{} `json:"web_transactions_apdex"`
    Port int `json:"port"`
    AppName string `json:"app_name"`
    TransactionNameRules []string `json:"transaction_name_rules"`
    AgentLimitsTransactionTracesNodes int `json:"agent_limits.transaction_traces_nodes"`
    TransactionTracerEnabled bool `json:"transaction_tracer.enabled"`
    LogLevel int `json:"log_level"`
    ProxyHost string `json:"proxy_host"`
    IgnoredParams []string `json:"ignored_params"
    AgentLimitsSqlExplainPlans int `json:"agent_limits.sql_explain_plans"`
}

func NewAgentSettings() *AgentSettings {
    s := &AgentSettings{
        StartupTimeout: 0.0,
        DebugLogDataCollectorCalls: true,
        ThreadProfilerEnabled: true,
        ErrorCollectorCaptureSource: false,
        CaptureParams: true,
        AgentLimitsSqlQueryLengthMaximum: 16384,
        ProxyPort: 0,
        IncludeEnviron: []string{"REQUEST_METHOD", "HTTP_USER_AGENT", "HTTP_REFERER", "CONTENT_TYPE", "CONTENT_LENGTH"},
        TransactionNameLimit: 0,
        DebugLogTransactionTracePayload: false,
        ShutdownTimeout: 30.0,
        TrustedAccountIds: []int{},
        WebTransactionsApdex: map[string]string{},
        Port: 0,
        AppName: "Python Agent Test",
        TransactionNameRules: []string{},
        AgentLimitsTransactionTracesNodes: 2000,
        TransactionTracerEnabled: true,
        LogLevel: 10,
        IgnoredParams: []string{},
        AgentLimitsSqlExplainPlans: 30,
    }
    return s
}

<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/armon/consul-api"
)

const (
	MANIPT_KEY = "service/manipt/leader"
)

var (
	AgentWaitTime = time.Duration(10) * time.Second
)

type Agent struct {
	Client *consulapi.Client
	Leader *consulapi.Node

	Bind string
	Port int

	PubListener    *net.TCPListener
	WebAppListener *net.TCPListener
	Server         *http.Server

	Session string

	leaderchan chan *consulapi.Node
	connchan   chan *net.TCPConn
	proxychan  chan *net.TCPConn
}

func NewAgent(c *consulapi.Client, bind string, port int, server *http.Server) *Agent {
	return &Agent{
		Client:     c,
		Leader:     nil,
		Bind:       bind,
		Port:       port,
		Server:     server,
		leaderchan: make(chan *consulapi.Node),
		connchan:   make(chan *net.TCPConn, 1024),
		proxychan:  make(chan *net.TCPConn, 1024),
	}
}

func (a *Agent) proxyConns() {
	for {
		c, err := a.PubListener.AcceptTCP()
		if err != nil {
			log.Printf("[err] failed to accept conn: %s", err)
			continue
		}
		a.connchan <- c
	}
}

func (a *Agent) setup() {
	publisten, err := net.ListenTCP("tcp", a.Bind+":"+strconv.Itoa(a.Port))
	if err != nil {
		log.Fatalf("[fatal] failed to setup public listener")
	}

	a.PubListener = publisten

	webapp, err := net.ListenTCP("tcp", "127.0.0.1:0")
	if err != nil {
		log.Fatalf("[fatal] failed to setup webapp listener")
	}

	a.WebAppListener = webapp

	go a.Server.Serve(a.WebAppListener)
}

func (a *Agent) Run() {
	// set our Listener
	a.setup()

	// get the current status first
	a.Leader = <-a.leaderchan

	// start proxying connections
	go a.proxyConns()

	for {
		select {
		case node := <-a.leaderchan:
			log.Printf("[info] received leader node: %s", node)

			if node == nil && a.Leader != nil {
				log.Printf("[info] becoming leader")
			} else if node != nil && a.Leader == nil {
				log.Printf("[info] becoming proxy")
			} else {
				log.Printf("[info] no change")
			}

			a.Leader = node
		case incoming := <-a.connchan:
			if a.Leader != nil {
				log.Printf("[info] proxying %s", incoming.RemoteAddr())
				a.proxyTo(a.Leader, incoming)
			} else {
				a.proxyWebApp(incoming)
			}
		}
	}
}

func (a *Agent) proxyWebApp(conn *net.TCPConn) {
	waddr := a.WebAppListener.Addr()
	proxy, err := net.Dial(waddr.Network(), waddr.String())
	if err != nil {
		log.Printf("[err] failed to send to webapp, dropping connection: %s", err)
		conn.Close()
		return
	}

	go a.proxyStream(conn, proxy)
	go a.proxyStream(proxy, conn)
}

func (a *Agent) proxyTo(to *consulapi.Node, conn *net.TCPConn) {
	proxy, err := net.Dial("tcp", to.Address+":"+strconv.Itoa(a.Port))
	if err != nil {
		log.Printf("[err] failed to proxy to leader, will drop connection: %s", err)
		conn.Close()
		return
	}

	go a.proxyStream(conn, proxy)
	go a.proxyStream(proxy, conn)
}

func (a *Agent) proxyStream(from, to *net.TCPConn) {
	for {
		buf := make([]byte, 1024)
		_, err := from.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("[err] error reading from %s to %s: %s", from, to, err)
			}
			from.CloseRead()
			to.CloseWrite()
			return
		}

		_, err = to.Write(buf)
		if err != nil {
			log.Printf("[err] error writing from %s to %s: %s", from, to, err)
			return
		}
	}
}

func (a *Agent) LeaderUpdater() {

	nodename, err := a.Client.Agent().NodeName()
	if err != nil {
		log.Fatalf("[fatal] can't get our own node name: %s", err)
	}

	if a.Session == "" {
		str, _, err := a.Client.Session().Create(nil, nil)
		if err != nil {
			log.Fatalf("[fatal] failed to create session in consul: %s", err)
		}
		a.Session = str
	}

	var lastindex uint64

	kv := a.Client.KV()
	for {
		res, _, err := kv.Acquire(&consulapi.KVPair{
			Key:     MANIPT_KEY + "/" + strconv.Itoa(a.Port),
			Value:   []byte(nodename),
			Session: a.Session,
		}, nil)

		if err != nil {
			log.Fatalf("[fatal] couldn't contact consul for acquiring lock: %s", err)
		}

		if res {
			// we're leader!
			a.leaderchan <- nil
		}

		for {
			// who's leader? let's check
			pair, meta, err := kv.Get(MANIPT_KEY+"/"+strconv.Itoa(a.Port), &consulapi.QueryOptions{
				WaitIndex: lastindex,
				WaitTime:  AgentWaitTime,
			})

			if err != nil {
				log.Printf("[err] failed to get kv for current leader node: %s", err)
				break
			}

			if pair.Session == "" {
				// whoever was no longer is
				// handle until we get a new leader
				a.leaderchan <- nil
				// set the last index
				lastindex = meta.LastIndex
				// wait for a moment (lockdelays)
				time.Sleep(10 * time.Second)
				// so we're going to try getting it
				break
			}

			if string(pair.Value) != nodename { // non-leaders
				// find this person
				catalog := a.Client.Catalog()
				node, _, err := catalog.Node(string(pair.Value), nil)
				if err != nil {
					log.Printf("[err] failed to get current leader node: %s", err)
					continue
				}
				a.leaderchan <- node.Node // send new leader off
			}

			lastindex = meta.LastIndex
		}
	}

}
<commit_msg>Minor bugfixes<commit_after>package main

import (
	"io"
	"log"
	"net"
	"net/http"
	"strconv"
	"time"

	"github.com/armon/consul-api"
)

const (
	MANIPT_KEY = "service/manipt/leader"
)

var (
	AgentWaitTime = time.Duration(10) * time.Second
)

type Agent struct {
	Client *consulapi.Client
	Leader *consulapi.Node

	Bind *net.TCPAddr

	PubListener    *net.TCPListener
	WebAppListener *net.TCPListener
	Server         *http.Server

	Session string

	leaderchan chan *consulapi.Node
	connchan   chan *net.TCPConn
	proxychan  chan *net.TCPConn
}

func NewAgent(c *consulapi.Client, bind string, port int, server *http.Server) *Agent {
	tcpaddr, err := net.ResolveTCPAddr("tcp", bind+":"+strconv.Itoa(port))
	if err != nil {
		panic(err)
	}

	return &Agent{
		Client:     c,
		Leader:     nil,
		Bind:       tcpaddr,
		Server:     server,
		leaderchan: make(chan *consulapi.Node),
		connchan:   make(chan *net.TCPConn, 1024),
		proxychan:  make(chan *net.TCPConn, 1024),
	}
}

func (a *Agent) proxyConns() {
	for {
		c, err := a.PubListener.AcceptTCP()
		if err != nil {
			log.Printf("[err] failed to accept conn: %s", err)
			continue
		}
		a.connchan <- c
	}
}

func (a *Agent) setup() {
	publisten, err := net.ListenTCP("tcp", a.Bind)
	if err != nil {
		log.Fatalf("[fatal] failed to setup public listener")
	}

	a.PubListener = publisten

	webapp, err := net.ListenTCP("tcp", &net.TCPAddr{
		IP:   net.IPv4(127, 0, 0, 1),
		Port: 0,
	})
	if err != nil {
		log.Fatalf("[fatal] failed to setup webapp listener")
	}

	a.WebAppListener = webapp

	go a.Server.Serve(a.WebAppListener)
}

func (a *Agent) Run() {
	// set our Listener
	a.setup()

	// get the current status first
	a.Leader = <-a.leaderchan

	if a.Leader == nil {
		log.Printf("[info] starting as leader")
	} else {
		log.Printf("[info] starting as proxy to %s @ %s", a.Leader.Node, a.Leader.Address)
	}

	// start proxying connections
	go a.proxyConns()

	for {
		select {
		case node := <-a.leaderchan:
			log.Printf("[info] received leader node: %s", node)

			if node == nil && a.Leader != nil {
				log.Printf("[info] becoming leader")
			} else if node != nil {
				log.Printf("[info] becoming proxy to %s @ %s", node.Node, node.Address)
			}

			a.Leader = node
		case incoming := <-a.connchan:
			if a.Leader != nil {
				log.Printf("[info] proxying %s", incoming.RemoteAddr())
				a.proxyTo(a.Leader, incoming)
			} else {
				a.proxyWebApp(incoming)
			}
		}
	}
}

func (a *Agent) proxyWebApp(conn *net.TCPConn) {
	waddr := a.WebAppListener.Addr().(*net.TCPAddr)
	proxy, err := net.DialTCP("tcp", nil, waddr)
	if err != nil {
		log.Printf("[err] failed to send to webapp, dropping connection: %s", err)
		conn.Close()
		return
	}

	go a.proxyStream(conn, proxy)
	go a.proxyStream(proxy, conn)
}

func (a *Agent) proxyTo(to *consulapi.Node, conn *net.TCPConn) {
	proxy, err := net.Dial("tcp", to.Address+":"+strconv.Itoa(a.Bind.Port))
	if err != nil {
		log.Printf("[err] failed to proxy to leader, will drop connection: %s", err)
		conn.Close()
		return
	}

	tcpproxy, ok := proxy.(*net.TCPConn)
	if !ok {
		log.Fatalf("[fatal] assert failed - tcpproxy should be a *net.TCPConn")
	}

	go a.proxyStream(conn, tcpproxy)
	go a.proxyStream(tcpproxy, conn)
}

func (a *Agent) proxyStream(from, to *net.TCPConn) {
	for {
		buf := make([]byte, 1024)
		_, err := from.Read(buf)
		if err != nil {
			if err != io.EOF {
				log.Printf("[err] error reading from %s to %s: %s", from, to, err)
			}
			from.CloseRead()
			to.CloseWrite()
			return
		}

		_, err = to.Write(buf)
		if err != nil {
			log.Printf("[err] error writing from %s to %s: %s", from, to, err)
			return
		}
	}
}

func (a *Agent) LeaderUpdater() {

	nodename, err := a.Client.Agent().NodeName()
	if err != nil {
		log.Fatalf("[fatal] can't get our own node name: %s", err)
	}

	if a.Session == "" {
		str, _, err := a.Client.Session().Create(nil, nil)
		if err != nil {
			log.Fatalf("[fatal] failed to create session in consul: %s", err)
		}
		a.Session = str
	}

	var lastindex uint64

	kv := a.Client.KV()
	for {
		res, _, err := kv.Acquire(&consulapi.KVPair{
			Key:     MANIPT_KEY + "/" + strconv.Itoa(a.Bind.Port),
			Value:   []byte(nodename),
			Session: a.Session,
		}, nil)

		if err != nil {
			log.Fatalf("[fatal] couldn't contact consul for acquiring lock: %s", err)
		}

		if res {
			// we're leader!
			a.leaderchan <- nil
		}

		for {
			// who's leader? let's check
			pair, meta, err := kv.Get(MANIPT_KEY+"/"+strconv.Itoa(a.Bind.Port), &consulapi.QueryOptions{
				WaitIndex: lastindex,
				WaitTime:  AgentWaitTime,
			})

			if err != nil {
				log.Printf("[err] failed to get kv for current leader node: %s", err)
				break
			}

			if pair.Session == "" {
				// whoever was no longer is
				// handle until we get a new leader
				a.leaderchan <- nil
				// set the last index
				lastindex = meta.LastIndex
				// wait for a moment (lockdelays)
				time.Sleep(10 * time.Second)
				// so we're going to try getting it
				break
			}

			if string(pair.Value) != nodename { // non-leaders
				// find this person
				catalog := a.Client.Catalog()
				node, _, err := catalog.Node(string(pair.Value), nil)
				if err != nil {
					log.Printf("[err] failed to get current leader node: %s", err)
					continue
				}
				a.leaderchan <- node.Node // send new leader off
			}

			lastindex = meta.LastIndex
		}
	}

}
<|endoftext|>


<commit_before>package downloads

import (
	"context"
	"time"

	"github.com/go-errors/errors"

	"github.com/itchio/butler/buse"
	"github.com/itchio/butler/buse/messages"
	"github.com/itchio/butler/cmd/operate"
	"github.com/itchio/butler/cmd/wipe"
	"github.com/itchio/butler/database/models"
)

var downloadsDriveCancelID = "Downloads.Drive"

func DownloadsDrive(rc *buse.RequestContext, params *buse.DownloadsDriveParams) (*buse.DownloadsDriveResult, error) {
	consumer := rc.Consumer
	consumer.Infof("Now driving downloads...")

	parentCtx := rc.Ctx
	ctx, cancelFunc := context.WithCancel(parentCtx)

	rc.CancelFuncs.Add(downloadsDriveCancelID, cancelFunc)
	defer rc.CancelFuncs.Remove(downloadsDriveCancelID)

poll:
	for {
		select {
		case <-ctx.Done():
			consumer.Infof("Drive cancelled, bye!")
			break poll
		default:
			// let's keep going
		}

		err := cleanDiscarded(rc)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		err = performOne(ctx, rc)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		time.Sleep(1 * time.Second)
	}

	res := &buse.DownloadsDriveResult{}
	return res, nil
}

func cleanDiscarded(rc *buse.RequestContext) error {
	consumer := rc.Consumer

	var discardedDownloads []*models.Download
	err := rc.DB().Where(`discarded`).Find(&discardedDownloads).Error
	if err != nil {
		return errors.Wrap(err, 0)
	}

	models.PreloadDownloads(rc.DB(), discardedDownloads)
	for _, download := range discardedDownloads {
		consumer.Opf("Cleaning up download for %s", operate.GameToString(download.Game))

		if download.StagingFolder == "" {
			consumer.Warnf("No staging folder specified, can't wipe")
		} else {
			consumer.Opf("Wiping staging folder...")
			err := wipe.Do(consumer, download.StagingFolder)
			if err != nil {
				consumer.Warnf("While wiping staging folder: %s", err.Error())
			}
		}

		if download.Fresh {
			if download.StagingFolder == "" {
				consumer.Warnf("No (fresh) install folder specified, can't wipe")
			} else {
				consumer.Opf("Wiping (fresh) install folder...")
				err := wipe.Do(consumer, download.InstallFolder)
				if err != nil {
					consumer.Warnf("While wiping (fresh) install folder: %s", err.Error())
				}
			}
		}

		err := rc.DB().Delete(download).Error
		if err != nil {
			return errors.Wrap(err, 0)
		}
	}
	return nil
}

func performOne(parentCtx context.Context, rc *buse.RequestContext) error {
	consumer := rc.Consumer

	var pendingDownloads []*models.Download
	err := rc.DB().Where(`finished_at IS NULL AND NOT discarded`).Order(`position ASC`).Find(&pendingDownloads).Error
	if err != nil {
		return errors.Wrap(err, 0)
	}

	if len(pendingDownloads) == 0 {
		return nil
	}

	download := pendingDownloads[0]
	download.Preload(rc.DB())
	consumer.Infof("%d pending downloads, performing for %s", len(pendingDownloads), operate.GameToString(download.Game))

	ctx, cancelFunc := context.WithCancel(parentCtx)
	defer cancelFunc()
	goGadgetoDiscardWatcher := func() {
		for {
			select {
			case <-time.After(5 * time.Second):
				var row = struct {
					Discarded bool
				}{}
				err := rc.DB().Raw(`SELECT discarded FROM downloads WHERE id = ?`, download.ID).Scan(&row).Error
				if err != nil {
					consumer.Warnf("Could not check whether download is discarded: %s", err.Error())
				}

				if row.Discarded {
					consumer.Infof("Download was cancelled from under us, bailing out!")
					cancelFunc()
				}
			case <-ctx.Done():
				return
			}
		}
	}
	go goGadgetoDiscardWatcher()

	var stage = "prepare"
	var progress, eta, bps float64

	sendProgress := func() error {
		// TODO: send BPS history in here too
		return messages.DownloadsDriveProgress.Notify(rc, &buse.DownloadsDriveProgressNotification{
			Download: formatDownload(download),
			Progress: &buse.DownloadProgress{
				Stage:    stage,
				Progress: progress,
				ETA:      eta,
				BPS:      bps,
			},
		})
	}

	defer rc.StopInterceptingNotification(messages.Progress.Method())
	rc.InterceptNotification(messages.Progress.Method(), func(method string, paramsIn interface{}) error {
		params := paramsIn.(*buse.ProgressNotification)
		progress = params.Progress
		eta = params.ETA
		bps = params.BPS
		return sendProgress()
	})

	defer rc.StopInterceptingNotification(messages.TaskStarted.Method())
	rc.InterceptNotification(messages.TaskStarted.Method(), func(method string, paramsIn interface{}) error {
		params := paramsIn.(*buse.TaskStartedNotification)
		stage = string(params.Type)
		return sendProgress()
	})

	err = func() (err error) {
		defer func() {
			if r := recover(); r != nil {
				consumer.Warnf("Recovered from panic!")
				if rErr, ok := r.(error); ok {
					err = errors.Wrap(rErr, 0)
				} else {
					err = errors.New(r)
				}
			}
		}()

		err = operate.InstallPerform(ctx, rc, &buse.InstallPerformParams{
			ID:            download.ID,
			StagingFolder: download.StagingFolder,
		})
		return
	}()
	if err != nil {
		if be, ok := buse.AsBuseError(err); ok {
			switch buse.Code(be.AsJsonRpc2().Code) {
			case buse.CodeOperationCancelled:
				// the whole drive was probably cancelled?
				return nil
			case buse.CodeOperationAborted:
				consumer.Warnf("Download aborted, cleaning it out.")
				err := rc.DB().Delete(download).Error
				if err != nil {
					return errors.Wrap(err, 0)
				}
				return nil
			}
		}

		var errString = err.Error()
		if se, ok := err.(*errors.Error); ok {
			errString = se.ErrorStack()
		}

		consumer.Warnf("Download failed: %s", errString)
		download.Error = &errString
		finishedAt := time.Now().UTC()
		download.FinishedAt = &finishedAt
		download.Save(rc.DB())
		return nil
	}

	consumer.Infof("Download finished!")
	finishedAt := time.Now().UTC()
	download.FinishedAt = &finishedAt
	download.Save(rc.DB())

	messages.DownloadsDriveFinished.Notify(rc, &buse.DownloadsDriveFinishedNotification{
		Download: formatDownload(download),
	})

	return nil
}
<commit_msg>Fix deprioritize detection, and send DownloadProgress at most once per second<commit_after>package downloads

import (
	"context"
	"time"

	"github.com/go-errors/errors"

	"github.com/itchio/butler/buse"
	"github.com/itchio/butler/buse/messages"
	"github.com/itchio/butler/cmd/operate"
	"github.com/itchio/butler/cmd/wipe"
	"github.com/itchio/butler/database/models"
)

var downloadsDriveCancelID = "Downloads.Drive"

func DownloadsDrive(rc *buse.RequestContext, params *buse.DownloadsDriveParams) (*buse.DownloadsDriveResult, error) {
	consumer := rc.Consumer
	consumer.Infof("Now driving downloads...")

	parentCtx := rc.Ctx
	ctx, cancelFunc := context.WithCancel(parentCtx)

	rc.CancelFuncs.Add(downloadsDriveCancelID, cancelFunc)
	defer rc.CancelFuncs.Remove(downloadsDriveCancelID)

poll:
	for {
		select {
		case <-ctx.Done():
			consumer.Infof("Drive cancelled, bye!")
			break poll
		default:
			// let's keep going
		}

		err := cleanDiscarded(rc)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		err = performOne(ctx, rc)
		if err != nil {
			return nil, errors.Wrap(err, 0)
		}

		time.Sleep(1 * time.Second)
	}

	res := &buse.DownloadsDriveResult{}
	return res, nil
}

func cleanDiscarded(rc *buse.RequestContext) error {
	consumer := rc.Consumer

	var discardedDownloads []*models.Download
	err := rc.DB().Where(`discarded`).Find(&discardedDownloads).Error
	if err != nil {
		return errors.Wrap(err, 0)
	}

	models.PreloadDownloads(rc.DB(), discardedDownloads)
	for _, download := range discardedDownloads {
		consumer.Opf("Cleaning up download for %s", operate.GameToString(download.Game))

		if download.StagingFolder == "" {
			consumer.Warnf("No staging folder specified, can't wipe")
		} else {
			consumer.Opf("Wiping staging folder...")
			err := wipe.Do(consumer, download.StagingFolder)
			if err != nil {
				consumer.Warnf("While wiping staging folder: %s", err.Error())
			}
		}

		if download.Fresh {
			if download.StagingFolder == "" {
				consumer.Warnf("No (fresh) install folder specified, can't wipe")
			} else {
				consumer.Opf("Wiping (fresh) install folder...")
				err := wipe.Do(consumer, download.InstallFolder)
				if err != nil {
					consumer.Warnf("While wiping (fresh) install folder: %s", err.Error())
				}
			}
		}

		err := rc.DB().Delete(download).Error
		if err != nil {
			return errors.Wrap(err, 0)
		}
	}
	return nil
}

func performOne(parentCtx context.Context, rc *buse.RequestContext) error {
	consumer := rc.Consumer

	var pendingDownloads []*models.Download
	err := rc.DB().Where(`finished_at IS NULL AND NOT discarded`).Order(`position ASC`).Find(&pendingDownloads).Error
	if err != nil {
		return errors.Wrap(err, 0)
	}

	if len(pendingDownloads) == 0 {
		return nil
	}

	download := pendingDownloads[0]
	download.Preload(rc.DB())
	consumer.Infof("%d pending downloads, performing for %s", len(pendingDownloads), operate.GameToString(download.Game))

	ctx, cancelFunc := context.WithCancel(parentCtx)
	defer cancelFunc()
	goGadgetoDiscardWatcher := func() {
		for {
			select {
			case <-time.After(5 * time.Second):
				// have we been discarded?
				{
					var row = struct {
						Discarded bool
					}{}
					err := rc.DB().Raw(`SELECT discarded FROM downloads WHERE id = ?`, download.ID).Scan(&row).Error
					if err != nil {
						consumer.Warnf("Could not check whether download is discarded: %s", err.Error())
					}

					if row.Discarded {
						consumer.Infof("Download was cancelled from under us, bailing out!")
						cancelFunc()
					}
				}

				// has something else been prioritized?
				{
					var row = struct {
						ID string
					}{}
					err := rc.DB().Raw(`SELECT id FROM downloads WHERE finished_at IS NULL AND NOT discarded ORDER BY position ASC LIMIT 1`).Scan(&row).Error
					if err != nil {
						consumer.Warnf("Could not check whether download is discarded: %s", err.Error())
					}

					if row.ID != download.ID {
						consumer.Infof("%s deprioritized (for %s), bailing out!", download.ID, row.ID)
						cancelFunc()
					}
				}
			case <-ctx.Done():
				return
			}
		}
	}
	go goGadgetoDiscardWatcher()

	var stage = "prepare"
	var progress, eta, bps float64

	lastProgress := time.Now()

	sendProgress := func() error {
		if time.Since(lastProgress).Seconds() < 1 {
			return nil
		}
		lastProgress = time.Now()

		// TODO: send BPS history in here too
		return messages.DownloadsDriveProgress.Notify(rc, &buse.DownloadsDriveProgressNotification{
			Download: formatDownload(download),
			Progress: &buse.DownloadProgress{
				Stage:    stage,
				Progress: progress,
				ETA:      eta,
				BPS:      bps,
			},
		})
	}

	defer rc.StopInterceptingNotification(messages.Progress.Method())
	rc.InterceptNotification(messages.Progress.Method(), func(method string, paramsIn interface{}) error {
		params := paramsIn.(*buse.ProgressNotification)
		progress = params.Progress
		eta = params.ETA
		bps = params.BPS
		return sendProgress()
	})

	defer rc.StopInterceptingNotification(messages.TaskStarted.Method())
	rc.InterceptNotification(messages.TaskStarted.Method(), func(method string, paramsIn interface{}) error {
		params := paramsIn.(*buse.TaskStartedNotification)
		stage = string(params.Type)
		return sendProgress()
	})

	err = func() (err error) {
		defer func() {
			if r := recover(); r != nil {
				consumer.Warnf("Recovered from panic!")
				if rErr, ok := r.(error); ok {
					err = errors.Wrap(rErr, 0)
				} else {
					err = errors.New(r)
				}
			}
		}()

		err = operate.InstallPerform(ctx, rc, &buse.InstallPerformParams{
			ID:            download.ID,
			StagingFolder: download.StagingFolder,
		})
		return
	}()
	if err != nil {
		if be, ok := buse.AsBuseError(err); ok {
			switch buse.Code(be.AsJsonRpc2().Code) {
			case buse.CodeOperationCancelled:
				// the whole drive was probably cancelled?
				return nil
			case buse.CodeOperationAborted:
				consumer.Warnf("Download aborted, cleaning it out.")
				err := rc.DB().Delete(download).Error
				if err != nil {
					return errors.Wrap(err, 0)
				}
				return nil
			}
		}

		var errString = err.Error()
		if se, ok := err.(*errors.Error); ok {
			errString = se.ErrorStack()
		}

		consumer.Warnf("Download failed: %s", errString)
		download.Error = &errString
		finishedAt := time.Now().UTC()
		download.FinishedAt = &finishedAt
		download.Save(rc.DB())
		return nil
	}

	consumer.Infof("Download finished!")
	finishedAt := time.Now().UTC()
	download.FinishedAt = &finishedAt
	download.Save(rc.DB())

	messages.DownloadsDriveFinished.Notify(rc, &buse.DownloadsDriveFinishedNotification{
		Download: formatDownload(download),
	})

	return nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"time"

	"github.com/uber/tchannel/golang"
	"github.com/uber/tchannel/golang/hyperbahn"
	"golang.org/x/net/context"
)

func main() {
	tchan, err := tchannel.NewChannel("go-echo-server", nil)
	if err != nil {
		log.Fatalf("Failed to create channel: %v", err)
	}

	l, err := net.Listen("tcp", ":61543")
	if err != nil {
		log.Fatalf("Could not listen: %v", err)
	}
	log.Printf("Listening on %v", l.Addr())

	tchan.Register(handler{}, "echo")
	tchan.Serve(l)

	// register service with Hyperbahn.
	client := hyperbahn.NewClient(tchan, os.Args[1:], &hyperbahn.ClientOptions{
		Handler: eventHandler{},
		Timeout: time.Second,
	})
	if err := client.Register(); err != nil {
		fmt.Println("Register threw error:", err)
	}

	// Server will keep running till Ctrl-C.
	select {}
}

type eventHandler struct{}

func (eventHandler) On(event hyperbahn.Event) {
	fmt.Printf("On(%v)\n", event)
}

func (eventHandler) OnError(err error) {
	fmt.Printf("OnError(%v)\n", err)
}

type handler struct{}

func (handler) Handle(ctx context.Context, call *tchannel.InboundCall) {
	var arg2 []byte
	if err := tchannel.NewArgReader(call.Arg2Reader()).Read(&arg2); err != nil {
		log.Printf("Read arg2 failed: %v\n", err)
	}

	var arg3 []byte
	if err := tchannel.NewArgReader(call.Arg3Reader()).Read(&arg3); err != nil {
		log.Printf("Read arg2 failed: %v\n", err)
	}

	resp := call.Response()
	if err := tchannel.NewArgWriter(resp.Arg2Writer()).Write(arg2); err != nil {
		log.Printf("Write arg2 failed: %v", arg2)
	}

	if err := tchannel.NewArgWriter(resp.Arg3Writer()).Write(arg3); err != nil {
		log.Printf("Write arg3 failed: %v", arg3)
	}

}
<commit_msg>Improve example hyperbahn server<commit_after>package main

import (
	"fmt"
	"log"
	"net"
	"os"
	"time"

	"github.com/uber/tchannel/golang"
	"github.com/uber/tchannel/golang/hyperbahn"
	"golang.org/x/net/context"
)

func main() {
	tchan, err := tchannel.NewChannel("go-echo-server", nil)
	if err != nil {
		log.Fatalf("Failed to create channel: %v", err)
	}

	l, err := net.Listen("tcp", "127.0.0.1:61543")
	if err != nil {
		log.Fatalf("Could not listen: %v", err)
	}
	log.Printf("Listening on %v", l.Addr())

	tchan.Register(handler{}, "echo")
	go tchan.Serve(l)

	time.Sleep(100 * time.Millisecond)

	if len(os.Args[1:]) == 0 {
		log.Fatalf("You must provide Hyperbahn nodes as arguments")
	}

	// register service with Hyperbahn.
	client := hyperbahn.NewClient(tchan, os.Args[1:], &hyperbahn.ClientOptions{
		Handler: eventHandler{},
		Timeout: time.Second,
	})
	if err := client.Register(); err != nil {
		log.Fatalf("Register threw error: %v", err)
	}

	// Server will keep running till Ctrl-C.
	select {}
}

type eventHandler struct{}

func (eventHandler) On(event hyperbahn.Event) {
	fmt.Printf("On(%v)\n", event)
}

func (eventHandler) OnError(err error) {
	fmt.Printf("OnError(%v)\n", err)
}

type handler struct{}

func (handler) Handle(ctx context.Context, call *tchannel.InboundCall) {
	var arg2 []byte
	if err := tchannel.NewArgReader(call.Arg2Reader()).Read(&arg2); err != nil {
		log.Printf("Read arg2 failed: %v\n", err)
	}

	var arg3 []byte
	if err := tchannel.NewArgReader(call.Arg3Reader()).Read(&arg3); err != nil {
		log.Printf("Read arg2 failed: %v\n", err)
	}

	resp := call.Response()
	if err := tchannel.NewArgWriter(resp.Arg2Writer()).Write(arg2); err != nil {
		log.Printf("Write arg2 failed: %v", arg2)
	}

	if err := tchannel.NewArgWriter(resp.Arg3Writer()).Write(arg3); err != nil {
		log.Printf("Write arg3 failed: %v", arg3)
	}

}
<|endoftext|>


<commit_before>/*

	test is a package that is used to run a boardgame/server.StorageManager
	implementation through its paces and verify it does everything correctly.

*/
package test

import (
	"encoding/json"
	"github.com/jkomoros/boardgame"
	"github.com/jkomoros/boardgame/examples/blackjack"
	"github.com/jkomoros/boardgame/examples/tictactoe"
	"github.com/jkomoros/boardgame/server/api/users"
	"github.com/workfit/tester/assert"
	"reflect"
	"testing"
)

type StorageManager interface {
	boardgame.StorageManager

	//CleanUp will be called when a given manager is done and can be dispoed of.
	CleanUp()

	//The methods past this point are the same ones that are included in Server.StorageManager
	Name() string

	Connect(config string) error

	Close()

	ListGames(max int) []*boardgame.GameStorageRecord

	UserIdsForGame(gameId string) []string

	SetPlayerForGame(gameId string, playerIndex boardgame.PlayerIndex, userId string) error

	UpdateUser(user *users.StorageRecord) error

	GetUserById(uid string) *users.StorageRecord

	GetUserByCookie(cookie string) *users.StorageRecord

	ConnectCookieToUser(cookie string, user *users.StorageRecord) error
}

type managerMap map[string]*boardgame.GameManager

func (m managerMap) Get(name string) *boardgame.GameManager {
	return m[name]
}

type StorageManagerFactory func() StorageManager

func Test(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {

	BasicTest(factory, testName, connectConfig, t)
	UsersTest(factory, testName, connectConfig, t)
	AgentsTest(factory, testName, connectConfig, t)

}

func BasicTest(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {
	storage := factory()

	defer storage.Close()

	defer storage.CleanUp()

	if err := storage.Connect(connectConfig); err != nil {
		t.Fatal("Unexpected error connecting: ", err.Error())
	}

	assert.For(t).ThatActual(storage.Name()).Equals(testName)

	managers := make(managerMap)

	tictactoeManager := tictactoe.NewManager(storage)

	managers[tictactoeManager.Delegate().Name()] = tictactoeManager

	blackjackManager := blackjack.NewManager(storage)

	managers[blackjackManager.Delegate().Name()] = blackjackManager

	tictactoeGame := boardgame.NewGame(tictactoeManager)

	tictactoeGame.SetUp(0, nil)

	move := tictactoeGame.PlayerMoveByName("Place Token")

	if move == nil {
		t.Fatal(testName, "Couldn't find a move")
	}

	if err := <-tictactoeGame.ProposeMove(move, boardgame.AdminPlayerIndex); err != nil {
		t.Fatal(testName, "Couldn't make move", err)
	}

	refriedMove, err := tictactoeGame.Move(1)

	assert.For(t).ThatActual(err).IsNil()
	assert.For(t).ThatActual(refriedMove).Equals(move)

	//OK, now test that the manager and SetUp and everyone did the right thing.

	localGame, err := storage.Game(tictactoeGame.Id())

	if err != nil {
		t.Error(testName, "Unexpected error", err)
	}

	if localGame == nil {
		t.Fatal(testName, "Couldn't get game copy out")
	}

	assert.For(t).ThatActual(tictactoeGame.SecretSalt()).Equals(localGame.SecretSalt)

	blob, err := json.MarshalIndent(tictactoeGame.StorageRecord(), "", "  ")

	if err != nil {
		t.Fatal(testName, "couldn't marshal game", err)
	}

	localBlob, err := json.MarshalIndent(localGame, "", "  ")

	if err != nil {
		t.Fatal(testName, "Couldn't marshal localGame", err)
	}

	compareJSONObjects(blob, localBlob, testName+"Comparing game and local game", t)

	//Verify that if the game is stored with wrong name that doesn't match manager it won't load up.

	blackjackGame := boardgame.NewGame(blackjackManager)

	blackjackGame.SetUp(0, nil)

	games := storage.ListGames(10)

	if games == nil {
		t.Error(testName, "ListGames gave back nothing")
	}

	if len(games) != 2 {
		t.Error(testName, "We called listgames with a tictactoe game and a blackjack game, but got", len(games), "back.")
	}

	//TODO: figure out how to test that name is matched when retrieving from store.

}

func UsersTest(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {
	storage := factory()

	defer storage.Close()

	defer storage.CleanUp()

	if err := storage.Connect(connectConfig); err != nil {
		t.Fatal("Err connecting to storage: ", err)
	}

	manager := tictactoe.NewManager(storage)

	game := boardgame.NewGame(manager)

	game.SetUp(2, nil)

	var nilIds []string

	ids := storage.UserIdsForGame("DEADBEEF")

	assert.For(t).ThatActual(ids).Equals(nilIds)

	ids = storage.UserIdsForGame(game.Id())

	assert.For(t).ThatActual(ids).Equals([]string{"", ""})

	userId := "THISISAVERYLONGUSERIDTOTESTTHATWEDONTCLIPSHORTUSERIDSTOOAGGRESSIVELY"

	cookie := "MYCOOKIE"

	fetchedUser := storage.GetUserById(userId)

	var nilUser *users.StorageRecord

	assert.For(t).ThatActual(fetchedUser).Equals(nilUser)

	user := &users.StorageRecord{Id: userId}

	err := storage.UpdateUser(user)

	assert.For(t).ThatActual(err).IsNil()

	fetchedUser = storage.GetUserById(userId)

	assert.For(t).ThatActual(fetchedUser).Equals(user)

	fetchedUser = storage.GetUserByCookie(cookie)

	assert.For(t).ThatActual(fetchedUser).Equals(nilUser)

	err = storage.ConnectCookieToUser(cookie, user)

	assert.For(t).ThatActual(err).IsNil()

	fetchedUser = storage.GetUserByCookie(cookie)

	assert.For(t).ThatActual(fetchedUser).Equals(user)

	err = storage.SetPlayerForGame(game.Id(), 0, userId)

	assert.For(t).ThatActual(err).IsNil()

	ids = storage.UserIdsForGame(game.Id())

	assert.For(t).ThatActual(ids).Equals([]string{userId, ""})

	err = storage.SetPlayerForGame(game.Id(), 0, userId)

	assert.For(t).ThatActual(err).IsNotNil()
}

func AgentsTest(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {

	storage := factory()

	defer storage.Close()
	defer storage.CleanUp()

	if err := storage.Connect(connectConfig); err != nil {
		t.Fatal("Err connecting to storage: ", err)
	}

	manager := tictactoe.NewManager(storage)

	game := boardgame.NewGame(manager)

	err := game.SetUp(2, []string{"", "ai"})

	assert.For(t).ThatActual(err).IsNil()

	refriedGame := manager.Game(game.Id())

	assert.For(t).ThatActual(refriedGame.Agents()).Equals(game.Agents())

	refriedBlob, err := storage.AgentState(game.Id(), 0)

	assert.For(t).ThatActual(err).IsNil()

	var nilBlob []byte

	assert.For(t).ThatActual(refriedBlob).Equals(nilBlob)

	blob := []byte("ThisIsABlob")

	err = storage.SaveAgentState(game.Id(), 0, blob)

	assert.For(t).ThatActual(err).IsNil()

	refriedBlob, err = storage.AgentState(game.Id(), 0)

	assert.For(t).ThatActual(err).IsNil()

	assert.For(t).ThatActual(refriedBlob).Equals(blob)

	newBlob := []byte("ThisIsANewBlob")

	err = storage.SaveAgentState(game.Id(), 0, newBlob)

	assert.For(t).ThatActual(err).IsNil()

	refriedBlob, err = storage.AgentState(game.Id(), 0)

	assert.For(t).ThatActual(err).IsNil()

	assert.For(t).ThatActual(refriedBlob).Equals(newBlob)

}

func compareJSONObjects(in []byte, golden []byte, message string, t *testing.T) {

	//recreated in boardgame/state_test.go

	var deserializedIn interface{}
	var deserializedGolden interface{}

	json.Unmarshal(in, &deserializedIn)
	json.Unmarshal(golden, &deserializedGolden)

	if deserializedIn == nil {
		t.Error("In didn't deserialize", message)
	}

	if deserializedGolden == nil {
		t.Error("Golden didn't deserialize", message)
	}

	if !reflect.DeepEqual(deserializedIn, deserializedGolden) {
		t.Error("Got wrong json.", message, "Got", string(in), "wanted", string(golden))
	}
}
<commit_msg>StorageManager test suite expects a manager that implements ExtendedGame() and UpdateExtendedGame. Part of #436. TESTS FAIL for memory and bolt.<commit_after>/*

	test is a package that is used to run a boardgame/server.StorageManager
	implementation through its paces and verify it does everything correctly.

*/
package test

import (
	"encoding/json"
	"github.com/jkomoros/boardgame"
	"github.com/jkomoros/boardgame/examples/blackjack"
	"github.com/jkomoros/boardgame/examples/tictactoe"
	"github.com/jkomoros/boardgame/server/api/extendedgame"
	"github.com/jkomoros/boardgame/server/api/users"
	"github.com/workfit/tester/assert"
	"reflect"
	"testing"
)

type StorageManager interface {
	boardgame.StorageManager

	//CleanUp will be called when a given manager is done and can be dispoed of.
	CleanUp()

	//The methods past this point are the same ones that are included in Server.StorageManager
	Name() string

	Connect(config string) error

	ExtendedGame(id string) (*extendedgame.CombinedStorageRecord, error)

	UpdateExtendedGame(id string, eGame *extendedgame.StorageRecord) error

	Close()
	ListGames(max int) []*extendedgame.CombinedStorageRecord

	UserIdsForGame(gameId string) []string

	SetPlayerForGame(gameId string, playerIndex boardgame.PlayerIndex, userId string) error

	UpdateUser(user *users.StorageRecord) error

	GetUserById(uid string) *users.StorageRecord

	GetUserByCookie(cookie string) *users.StorageRecord

	ConnectCookieToUser(cookie string, user *users.StorageRecord) error
}

type managerMap map[string]*boardgame.GameManager

func (m managerMap) Get(name string) *boardgame.GameManager {
	return m[name]
}

type StorageManagerFactory func() StorageManager

func Test(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {

	BasicTest(factory, testName, connectConfig, t)
	UsersTest(factory, testName, connectConfig, t)
	AgentsTest(factory, testName, connectConfig, t)

}

func BasicTest(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {
	storage := factory()

	defer storage.Close()

	defer storage.CleanUp()

	if err := storage.Connect(connectConfig); err != nil {
		t.Fatal("Unexpected error connecting: ", err.Error())
	}

	assert.For(t).ThatActual(storage.Name()).Equals(testName)

	managers := make(managerMap)

	tictactoeManager := tictactoe.NewManager(storage)

	managers[tictactoeManager.Delegate().Name()] = tictactoeManager

	blackjackManager := blackjack.NewManager(storage)

	managers[blackjackManager.Delegate().Name()] = blackjackManager

	tictactoeGame := boardgame.NewGame(tictactoeManager)

	tictactoeGame.SetUp(0, nil)

	move := tictactoeGame.PlayerMoveByName("Place Token")

	if move == nil {
		t.Fatal(testName, "Couldn't find a move")
	}

	if err := <-tictactoeGame.ProposeMove(move, boardgame.AdminPlayerIndex); err != nil {
		t.Fatal(testName, "Couldn't make move", err)
	}

	refriedMove, err := tictactoeGame.Move(1)

	assert.For(t).ThatActual(err).IsNil()
	assert.For(t).ThatActual(refriedMove).Equals(move)

	//OK, now test that the manager and SetUp and everyone did the right thing.

	localGame, err := storage.Game(tictactoeGame.Id())

	if err != nil {
		t.Error(testName, "Unexpected error", err)
	}

	if localGame == nil {
		t.Fatal(testName, "Couldn't get game copy out")
	}

	assert.For(t).ThatActual(tictactoeGame.SecretSalt()).Equals(localGame.SecretSalt)

	blob, err := json.MarshalIndent(tictactoeGame.StorageRecord(), "", "  ")

	if err != nil {
		t.Fatal(testName, "couldn't marshal game", err)
	}

	localBlob, err := json.MarshalIndent(localGame, "", "  ")

	if err != nil {
		t.Fatal(testName, "Couldn't marshal localGame", err)
	}

	compareJSONObjects(blob, localBlob, testName+"Comparing game and local game", t)

	//Verify that if the game is stored with wrong name that doesn't match manager it won't load up.

	blackjackGame := boardgame.NewGame(blackjackManager)

	blackjackGame.SetUp(0, nil)

	games := storage.ListGames(10)

	if games == nil {
		t.Error(testName, "ListGames gave back nothing")
	}

	if len(games) != 2 {
		t.Error(testName, "We called listgames with a tictactoe game and a blackjack game, but got", len(games), "back.")
	}

	//TODO: figure out how to test that name is matched when retrieving from store.

}

func UsersTest(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {
	storage := factory()

	defer storage.Close()

	defer storage.CleanUp()

	if err := storage.Connect(connectConfig); err != nil {
		t.Fatal("Err connecting to storage: ", err)
	}

	manager := tictactoe.NewManager(storage)

	game := boardgame.NewGame(manager)

	game.SetUp(2, nil)

	var nilIds []string

	ids := storage.UserIdsForGame("DEADBEEF")

	assert.For(t).ThatActual(ids).Equals(nilIds)

	ids = storage.UserIdsForGame(game.Id())

	assert.For(t).ThatActual(ids).Equals([]string{"", ""})

	userId := "THISISAVERYLONGUSERIDTOTESTTHATWEDONTCLIPSHORTUSERIDSTOOAGGRESSIVELY"

	cookie := "MYCOOKIE"

	fetchedUser := storage.GetUserById(userId)

	var nilUser *users.StorageRecord

	assert.For(t).ThatActual(fetchedUser).Equals(nilUser)

	user := &users.StorageRecord{Id: userId}

	err := storage.UpdateUser(user)

	assert.For(t).ThatActual(err).IsNil()

	fetchedUser = storage.GetUserById(userId)

	assert.For(t).ThatActual(fetchedUser).Equals(user)

	fetchedUser = storage.GetUserByCookie(cookie)

	assert.For(t).ThatActual(fetchedUser).Equals(nilUser)

	err = storage.ConnectCookieToUser(cookie, user)

	assert.For(t).ThatActual(err).IsNil()

	fetchedUser = storage.GetUserByCookie(cookie)

	assert.For(t).ThatActual(fetchedUser).Equals(user)

	err = storage.SetPlayerForGame(game.Id(), 0, userId)

	assert.For(t).ThatActual(err).IsNil()

	ids = storage.UserIdsForGame(game.Id())

	assert.For(t).ThatActual(ids).Equals([]string{userId, ""})

	err = storage.SetPlayerForGame(game.Id(), 0, userId)

	assert.For(t).ThatActual(err).IsNotNil()
}

func AgentsTest(factory StorageManagerFactory, testName string, connectConfig string, t *testing.T) {

	storage := factory()

	defer storage.Close()
	defer storage.CleanUp()

	if err := storage.Connect(connectConfig); err != nil {
		t.Fatal("Err connecting to storage: ", err)
	}

	manager := tictactoe.NewManager(storage)

	game := boardgame.NewGame(manager)

	err := game.SetUp(2, []string{"", "ai"})

	assert.For(t).ThatActual(err).IsNil()

	refriedGame := manager.Game(game.Id())

	assert.For(t).ThatActual(refriedGame.Agents()).Equals(game.Agents())

	refriedBlob, err := storage.AgentState(game.Id(), 0)

	assert.For(t).ThatActual(err).IsNil()

	var nilBlob []byte

	assert.For(t).ThatActual(refriedBlob).Equals(nilBlob)

	blob := []byte("ThisIsABlob")

	err = storage.SaveAgentState(game.Id(), 0, blob)

	assert.For(t).ThatActual(err).IsNil()

	refriedBlob, err = storage.AgentState(game.Id(), 0)

	assert.For(t).ThatActual(err).IsNil()

	assert.For(t).ThatActual(refriedBlob).Equals(blob)

	newBlob := []byte("ThisIsANewBlob")

	err = storage.SaveAgentState(game.Id(), 0, newBlob)

	assert.For(t).ThatActual(err).IsNil()

	refriedBlob, err = storage.AgentState(game.Id(), 0)

	assert.For(t).ThatActual(err).IsNil()

	assert.For(t).ThatActual(refriedBlob).Equals(newBlob)

}

func compareJSONObjects(in []byte, golden []byte, message string, t *testing.T) {

	//recreated in boardgame/state_test.go

	var deserializedIn interface{}
	var deserializedGolden interface{}

	json.Unmarshal(in, &deserializedIn)
	json.Unmarshal(golden, &deserializedGolden)

	if deserializedIn == nil {
		t.Error("In didn't deserialize", message)
	}

	if deserializedGolden == nil {
		t.Error("Golden didn't deserialize", message)
	}

	if !reflect.DeepEqual(deserializedIn, deserializedGolden) {
		t.Error("Got wrong json.", message, "Got", string(in), "wanted", string(golden))
	}
}
<|endoftext|>


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

// Package https implements the HTTPS transport protocol for upspin.Store.
package https

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"io"

	"upspin.io/bind"
	"upspin.io/upspin"
)

// Store is an implementation of upspin.Store that interfaces
// with an HTTP server for serving data.
type Store struct {
	upspin.NoConfiguration
	serverURL  string
	httpClient HTTPClient
}

// Guarantee we implement the interface
var _ upspin.Store = (*Store)(nil)

// maxBytesLimit is the maximum number of bytes to retrieve in one request.
const maxBytesLimit = 1 << 30 // 1GB

// Some error messages.
const (
	invalidRefError = "invalid reference"
	notHTTPError    = "not an HTTP(S) reference"
	httpClientError = "HTTP client error: %v"

	Get = "GET" // HTTP Get method
)

// HTTPClient is a minimal HTTP client interface. An instance of
// http.Client implements this interface.
type HTTPClient interface {
	Do(req *http.Request) (resp *http.Response, err error)
}

// New returns a concrete implementation of Store, pointing to a
// server at a given URL (including the port), for performing Get and
// Put requests on blocks of data. Use this only for testing.
func New(serverURL string, httpClient HTTPClient) *Store {
	return &Store{
		serverURL:  serverURL,
		httpClient: httpClient,
	}
}

// IsServerReachable reports whether the server at an URL can be reached.
func IsServerReachable(serverURL string) bool {
	_, err := http.Head(serverURL)
	return err == nil
}

// Dial implements Dialer.
func (s *Store) Dial(context *upspin.Context, endpoint upspin.Endpoint) (upspin.Service, error) {
	const op = "Dial"
	if context == nil {
		return nil, newStoreError(op, "nil context", "")
	}
	serverURL, err := url.Parse(string(endpoint.NetAddr))
	if err != nil {
		return nil, newStoreError(op, fmt.Sprintf("invalid HTTP address for endpoint: %v", err), "")
	}
	s.serverURL = serverURL.String()
	if !IsServerReachable(s.serverURL) {
		return nil, newStoreError(op, "HTTPS store server unreachable", "")
	}
	return s, nil
}

// Ping implements Service.
func (s *Store) Ping() bool {
	return IsServerReachable(s.serverURL)
}

// ServerUserName implements Service.
func (s *Store) ServerUserName() string {
	return ""
}

// Get implements Store.
func (s *Store) Get(ref upspin.Reference) ([]byte, []upspin.Location, error) {
	const op = "Get"
	if ref == "" {
		return nil, nil, newStoreError(op, invalidRefError, "")
	}
	url := string(ref)
	if !strings.HasPrefix(string(ref), "http://") && !strings.HasPrefix(string(ref), "https://") {
		return nil, nil, newStoreError(op, notHTTPError, ref)
	}
	httpReq, err := http.NewRequest(Get, url, nil)
	if err != nil {
		return nil, nil, err
	}
	body, err := s.requestAndReadResponseBody(op, ref, httpReq)
	if err != nil {
		return nil, nil, err
	}
	return body, nil, nil
}

// Put implements Store.
func (s *Store) Put(data []byte) (upspin.Reference, error) {
	return "", errors.New("https: Put: not implemented")
}

// Delete implements Store.
func (s *Store) Delete(ref upspin.Reference) error {
	return errors.New("https: Delete: not implemented")
}

// requestAndReadResponseBody is an internal helper function that
// sends a given request over the HTTP client and parses the body of
// the reply, using op and key to build an error if one is
// encountered along the way.
func (s *Store) requestAndReadResponseBody(op string, ref upspin.Reference, req *http.Request) ([]byte, error) {
	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, newStoreError(op, err.Error(), ref)
	}

	// Check the response
	if resp.StatusCode != http.StatusOK {
		return nil, newStoreError(op, fmt.Sprintf(httpClientError, resp.StatusCode), ref)
	}

	// Read the body of the response
	defer resp.Body.Close()
	respBody, err := BufferResponse(resp, maxBytesLimit)
	if err != nil {
		return nil, newStoreError(op, err.Error(), ref)
	}
	return respBody, nil
}

// Endpoint implements upspin.Service.
func (s *Store) Endpoint() upspin.Endpoint {
	return upspin.Endpoint{
		Transport: upspin.HTTPS,
		NetAddr:   upspin.NetAddr(s.serverURL),
	}
}

// Close implements upspin.Service.
func (s *Store) Close() {
	// Nothing to do.
}

// Authenticate implements upspin.Service.
func (s *Store) Authenticate(*upspin.Context) error {
	return nil
}

type storeError struct {
	op    string
	error string
	ref   upspin.Reference
}

// Error implements error
func (s storeError) Error() string {
	if s.ref != "" {
		return fmt.Sprintf("https: store error: %s: %s: %s", s.op, s.ref, s.error)
	}
	return fmt.Sprintf("https: store error: %s: %s", s.op, s.error)
}

func newStoreError(op string, error string, ref upspin.Reference) *storeError {
	return &storeError{
		op:    op,
		error: error,
		ref:   ref,
	}
}

// errTooLong is returned when a BufferResponse would not fit in the buffer budget.
var errTooLong = errors.New("response body too long")

// BufferResponse reads the body of an HTTP response up to maxBufLen bytes. It closes the response body.
// If the response is larger than maxBufLen, it returns ErrTooLong.
func BufferResponse(resp *http.Response, maxBufLen int64) ([]byte, error) {
	var buf []byte
	defer resp.Body.Close()
	if resp.ContentLength > 0 {
		if resp.ContentLength <= maxBufLen {
			buf = make([]byte, resp.ContentLength)
		} else {
			// Return an error
			return nil, errTooLong
		}
	} else {
		buf = make([]byte, maxBufLen)
	}
	n, err := resp.Body.Read(buf)
	if err != nil && err != io.EOF {
		if err == io.ErrShortBuffer {
			return nil, errTooLong
		}
		return nil, err
	}
	return buf[:n], nil
}

func init() {
	// By default, set up only the HTTP client. The server URL gets bound at Dial time.
	bind.RegisterStore(upspin.HTTPS, New("", &http.Client{}))
}
<commit_msg>https: fix from deleted netutil CL<commit_after>// Copyright 2016 The Upspin Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// Package https implements the HTTPS transport protocol for upspin.Store.
package https

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"strings"

	"io"

	"upspin.io/bind"
	"upspin.io/upspin"
)

// Store is an implementation of upspin.Store that interfaces
// with an HTTP server for serving data.
type Store struct {
	upspin.NoConfiguration
	serverURL  string
	httpClient HTTPClient
}

// Guarantee we implement the interface
var _ upspin.Store = (*Store)(nil)

// maxBytesLimit is the maximum number of bytes to retrieve in one request.
const maxBytesLimit = 1 << 30 // 1GB

// Some error messages.
const (
	invalidRefError = "invalid reference"
	notHTTPError    = "not an HTTP(S) reference"
	httpClientError = "HTTP client error: %v"

	Get = "GET" // HTTP Get method
)

// HTTPClient is a minimal HTTP client interface. An instance of
// http.Client implements this interface.
type HTTPClient interface {
	Do(req *http.Request) (resp *http.Response, err error)
}

// New returns a concrete implementation of Store, pointing to a
// server at a given URL (including the port), for performing Get and
// Put requests on blocks of data. Use this only for testing.
func New(serverURL string, httpClient HTTPClient) *Store {
	return &Store{
		serverURL:  serverURL,
		httpClient: httpClient,
	}
}

// IsServerReachable reports whether the server at an URL can be reached.
func IsServerReachable(serverURL string) bool {
	_, err := http.Head(serverURL)
	return err == nil
}

// Dial implements Dialer.
func (s *Store) Dial(context *upspin.Context, endpoint upspin.Endpoint) (upspin.Service, error) {
	const op = "Dial"
	if context == nil {
		return nil, newStoreError(op, "nil context", "")
	}
	serverURL, err := url.Parse(string(endpoint.NetAddr))
	if err != nil {
		return nil, newStoreError(op, fmt.Sprintf("invalid HTTP address for endpoint: %v", err), "")
	}
	s.serverURL = serverURL.String()
	if !IsServerReachable(s.serverURL) {
		return nil, newStoreError(op, "HTTPS store server unreachable", "")
	}
	return s, nil
}

// Ping implements Service.
func (s *Store) Ping() bool {
	return IsServerReachable(s.serverURL)
}

// ServerUserName implements Service.
func (s *Store) ServerUserName() string {
	return ""
}

// Get implements Store.
func (s *Store) Get(ref upspin.Reference) ([]byte, []upspin.Location, error) {
	const op = "Get"
	if ref == "" {
		return nil, nil, newStoreError(op, invalidRefError, "")
	}
	url := string(ref)
	if !strings.HasPrefix(string(ref), "http://") && !strings.HasPrefix(string(ref), "https://") {
		return nil, nil, newStoreError(op, notHTTPError, ref)
	}
	httpReq, err := http.NewRequest(Get, url, nil)
	if err != nil {
		return nil, nil, err
	}
	body, err := s.requestAndReadResponseBody(op, ref, httpReq)
	if err != nil {
		return nil, nil, err
	}
	return body, nil, nil
}

// Put implements Store.
func (s *Store) Put(data []byte) (upspin.Reference, error) {
	return "", errors.New("https: Put: not implemented")
}

// Delete implements Store.
func (s *Store) Delete(ref upspin.Reference) error {
	return errors.New("https: Delete: not implemented")
}

// requestAndReadResponseBody is an internal helper function that
// sends a given request over the HTTP client and parses the body of
// the reply, using op and key to build an error if one is
// encountered along the way.
func (s *Store) requestAndReadResponseBody(op string, ref upspin.Reference, req *http.Request) ([]byte, error) {
	resp, err := s.httpClient.Do(req)
	if err != nil {
		return nil, newStoreError(op, err.Error(), ref)
	}

	// Check the response
	if resp.StatusCode != http.StatusOK {
		return nil, newStoreError(op, fmt.Sprintf(httpClientError, resp.StatusCode), ref)
	}

	// Read the body of the response
	defer resp.Body.Close()
	respBody, err := BufferResponse(resp, maxBytesLimit)
	if err != nil {
		return nil, newStoreError(op, err.Error(), ref)
	}
	return respBody, nil
}

// Endpoint implements upspin.Service.
func (s *Store) Endpoint() upspin.Endpoint {
	return upspin.Endpoint{
		Transport: upspin.HTTPS,
		NetAddr:   upspin.NetAddr(s.serverURL),
	}
}

// Close implements upspin.Service.
func (s *Store) Close() {
	// Nothing to do.
}

// Authenticate implements upspin.Service.
func (s *Store) Authenticate(*upspin.Context) error {
	return nil
}

type storeError struct {
	op    string
	error string
	ref   upspin.Reference
}

// Error implements error
func (s storeError) Error() string {
	if s.ref != "" {
		return fmt.Sprintf("https: store error: %s: %s: %s", s.op, s.ref, s.error)
	}
	return fmt.Sprintf("https: store error: %s: %s", s.op, s.error)
}

func newStoreError(op string, error string, ref upspin.Reference) *storeError {
	return &storeError{
		op:    op,
		error: error,
		ref:   ref,
	}
}

// errTooLong is returned when a BufferResponse would not fit in the buffer budget.
var errTooLong = errors.New("response body too long")

// BufferResponse reads the body of an HTTP response up to maxBufLen bytes. It closes the response body.
// If the response is larger than maxBufLen, it returns ErrTooLong.
func BufferResponse(resp *http.Response, maxBufLen int64) ([]byte, error) {
	var buf []byte
	defer resp.Body.Close()
	if resp.ContentLength > 0 {
		if resp.ContentLength <= maxBufLen {
			buf = make([]byte, resp.ContentLength)
		} else {
			// Return an error
			return nil, errTooLong
		}
	} else {
		buf = make([]byte, maxBufLen)
	}
	n, err := io.ReadFull(resp.Body, buf)
	if err != nil {
		if err == io.ErrUnexpectedEOF {
			return buf[:n], nil
		}
		return nil, err
	}
	return buf, nil
}

func init() {
	// By default, set up only the HTTP client. The server URL gets bound at Dial time.
	bind.RegisterStore(upspin.HTTPS, New("", &http.Client{}))
}
<|endoftext|>


<commit_before>package checklogfile

import (
	"io"
	"io/ioutil"
	"os"
	"testing"
)

var positions = [...]struct {
	description string
	offset      int64
	want        int64
	whence      int
}{
	{
		description: "Getting position",
		whence:      1, offset: 0, want: 20,
	},
	{
		description: "Seeking 100 bytes forward",
		whence:      1, offset: 100, want: 120,
	},
	{
		description: "Seeking to 200 bytes",
		whence:      0, offset: 200, want: 200,
	},
	{
		description: "Rewind to beginning",
		whence:      0, offset: 0, want: 0,
	},
}

func TestGzipSeeker(t *testing.T) {
	fp, err := os.Open("testdata/unattended-upgrades.log.gz")
	if err != nil {
		t.Fatal("testdata missing. Error: ", err)
		return
	}
	// This also ensures we have the right interface
	var cw ReadSeekCloser
	cw = NewCompressorSeekWrapper(fp, "gz")
	if cw == nil {
		t.Fatal("cannot open compressing seeker")
		return
	}
	// do a bit io
	_, err = io.CopyN(ioutil.Discard, cw, 20)
	if err != nil {
		t.Fatal("testdata too small (need at least 100 bytes) Error ", err)
		return
	}
	for i, p := range positions {
		t.Logf("%d:%s", i, p.description)
		if pos, err := cw.Seek(p.offset, p.whence); err != nil {
			t.Errorf("%d:err: %s", i, err)
		} else if pos != p.want {
			t.Errorf("%d: got %d, want %d", i, pos, p.want)
		} else {
			t.Logf("%d: ok, got %d", i, pos)
		}
	}
}

func TestBzip2Seeker(t *testing.T) {
	fp, err := os.Open("testdata/unattended-upgrades.log.bz2")
	if err != nil {
		t.Fatal("testdata missing. Error: ", err)
		return
	}
	// This also ensures we have the right interface
	var cw ReadSeekCloser
	cw = NewCompressorSeekWrapper(fp, "bz2")
	if cw == nil {
		t.Fatal("cannot open compressing seeker")
		return
	}
	// do a bit io
	_, err = io.CopyN(ioutil.Discard, cw, 20)
	if err != nil {
		t.Fatal("testdata too small (need at least 100 bytes) Error ", err)
		return
	}
	for i, p := range positions {
		t.Logf("%d:%s", i, p.description)
		if pos, err := cw.Seek(p.offset, p.whence); err != nil {
			t.Errorf("%d:err: %s", i, err)
		} else if pos != p.want {
			t.Errorf("%d: got %d, want %d", i, pos, p.want)
		} else {
			t.Logf("%d: ok, got %d", i, pos)
		}
	}
}

func BenchmarkGzipWrapper(b *testing.B) {
	fp, err := os.Open("testdata/unattended-upgrades.log.gz")
	if err != nil {
		b.Fatal("testdata missing. Error: ", err)
		return
	}
	// This also ensures we have the right interface
	var cw ReadSeekCloser
	cw = NewCompressorSeekWrapper(fp, "gz")
	if cw == nil {
		b.Fatal("cannot open compressing seeker")
		return
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		b.StartTimer()
		n, _ := io.Copy(ioutil.Discard, cw)
		b.StopTimer()
		if err != nil {
			b.Fatal("invalid testdata: ", err)
			return
		}
		b.SetBytes(n)
		cw.Seek(0, 0)
	}
}
<commit_msg>Better testing<commit_after>package checklogfile

import (
	"io"
	"io/ioutil"
	"os"
	"testing"
)

var positions = [...]struct {
	description string
	offset      int64
	want        int64
	whence      int
}{
	{
		description: "Getting position",
		whence:      1, offset: 0, want: 20,
	},
	{
		description: "Seeking 100 bytes forward",
		whence:      1, offset: 100, want: 120,
	},
	{
		description: "Seeking to 200 bytes",
		whence:      0, offset: 200, want: 200,
	},
	{
		description: "Rewind to beginning",
		whence:      0, offset: 0, want: 0,
	},
}

var files = map[string]string{
	"bz2": "testdata/unattended-upgrades.log.bz2",
	"gz":  "testdata/unattended-upgrades.log.gz",
	"":    "testdata/unattended-upgrades.log",
}

func TestCompressorSeeker(t *testing.T) {
	want_line := []byte("2013-02-20 11:05:29,")
	if n := len(want_line); n != 20 {
		t.Fatalf("testsetup broken. want_line should be %d bytes, got %d bytes", 20, n)
		return
	}
	for ext, f := range files {
		fp, err := os.Open(f)
		if err != nil {
			t.Fatal("testdata missing. Error: ", err)
			return
		}
		// This also ensures we have the right interface
		var cw ReadSeekCloser
		cw = NewCompressorSeekWrapper(fp, ext)
		if cw == nil {
			t.Errorf("%s:cannot open compressing seeker", ext)
			continue
		}

		got_line := make([]byte, len(want_line))

		n, err := cw.Read(got_line)
		if err != nil {
			t.Error(ext, ": testdata too small (need at least 200 bytes) Error ", err)
			continue
		} else if n < 20 {
			t.Errorf("%s: testdata too small (need at least 200 bytes, got %v) Error %v", ext, n, err)
			continue
		}

		if string(want_line) != string(got_line) {
			t.Errorf("%s: decompressor b0rken: want %s, got %s", ext, want_line, got_line)
			continue
		}
		for i, p := range positions {
			t.Logf("%d:%s", i, p.description)
			if pos, err := cw.Seek(p.offset, p.whence); err != nil {
				t.Errorf("%s:%d:err: %s", ext, i, err)
			} else if pos != p.want {
				t.Errorf("%s:%d: got %d, want %d", ext, i, pos, p.want)
			} else {
				t.Logf("%s:%d: ok, got %d", ext, i, pos)
			}
		}
	}
}

func BenchmarkGzipWrapper(b *testing.B) {
	fp, err := os.Open("testdata/unattended-upgrades.log.gz")
	if err != nil {
		b.Fatal("testdata missing. Error: ", err)
		return
	}
	// This also ensures we have the right interface
	var cw ReadSeekCloser
	cw = NewCompressorSeekWrapper(fp, "gz")
	if cw == nil {
		b.Fatal("cannot open compressing seeker")
		return
	}
	b.ResetTimer()
	for i := 0; i < b.N; i++ {
		b.StopTimer()
		b.StartTimer()
		n, _ := io.Copy(ioutil.Discard, cw)
		b.StopTimer()
		if err != nil {
			b.Fatal("invalid testdata: ", err)
			return
		}
		b.SetBytes(n)
		cw.Seek(0, 0)
	}
}
<|endoftext|>