<commit_before>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"encoding/base64"
	"encoding/json"

	"github.com/abbot/go-http-auth"
	"github.com/dchest/uniuri"
	"github.com/zenazn/goji/bind"
	"goji.io"
	"goji.io/pat"
)

const cookieName string = "_htpasswd_auth"
const sessionFormat string = time.RFC3339
const slack time.Duration = time.Duration(10) * time.Second

var sessionDir string
var domain string
var htpasswd string
var realm string
var cookieLifetime int

var validChars map[rune]bool

func init() {
	validChars = map[rune]bool{}
	for _, c := range string(uniuri.StdChars) {
		validChars[c] = true
	}
}

type session struct {
	CreatedRaw string `json:"created"`
	Domain     string

	Name    string    `json:"-"`
	Created time.Time `json:"-"`
}

func (s *session) Valid() bool {
	sessionExpiry := s.Created.Add(time.Duration(cookieLifetime) * time.Second)
	return sessionExpiry.After(time.Now())
}
func validateCookie(cookie string) (*session, error) {
	for _, c := range cookie {
		if _, ok := validChars[c]; !ok {
			return nil, errors.New("Cookie value contains invalid characters.")
		}
	}
	sessionPath := filepath.Join(sessionDir, cookie)
	base, err := filepath.Rel(sessionDir, sessionPath)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(base, ".") || strings.HasPrefix(base, "/") {
		return nil, errors.New("Cookie value would traverse directories")
	}

	// From here on, the session pathname is assumed innocent & handle-able.
	data, err := ioutil.ReadFile(sessionPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Printf("Request for a session that doesn't exist: %q", sessionPath)
			return nil, err
		} else {
			log.Printf("Checking session %q failed: %s", sessionPath, err)
			return nil, err
		}
	}
	session := session{Name: cookie}
	json.Unmarshal(data, &session)
	sessionStart, err := time.Parse(sessionFormat, string(session.CreatedRaw))
	if err != nil {
		os.Remove(sessionPath) // clean up the session
		return nil, err
	}
	session.Created = sessionStart
	if !session.Valid() {
		return nil, errors.New("Session is no longer valid.")
	}
	return &session, nil
}

func invalidateCookie() *http.Cookie {
	return &http.Cookie{Name: cookieName, Value: "nope", Domain: domain, Path: "/", MaxAge: -1, Secure: true, HttpOnly: true}
}

func checkSession(w http.ResponseWriter, r *http.Request) {
	cookie, err := r.Cookie(cookieName)
	newCookie := invalidateCookie()
	success := false

	// Cleanup and ensure we always send a decent response:
	defer func() {
		if !success {
			http.SetCookie(w, newCookie)
			http.Error(w, "Session is not/no longer valid", http.StatusUnauthorized)
			return
		}
		w.WriteHeader(200)
		w.Write([]byte{})
	}()
	if err != nil {
		if err == http.ErrNoCookie {
			return
		}
	}
	session, err := validateCookie(cookie.Value)
	if err != nil {
		log.Printf("Error validating session: %s", err)
		return
	}

	if session.Domain != r.Host {
		log.Printf("Session %q was made for another domain: %s; wanted %q", cookie.Value, session.Domain, r.Host)
		return
	}

	// We have a valid session that hasn't yet expired. Let's call it a success.
	success = true
	return
}

func newSession(domain string) (string, error) {
	name := uniuri.NewLen(90)
	sessionPath := filepath.Join(sessionDir, name)
	data, err := json.Marshal(session{CreatedRaw: time.Now().Add(slack).Format(sessionFormat), Domain: domain})
	if err != nil {
		return "", err
	}
	err = ioutil.WriteFile(sessionPath, data, 0400)
	if err != nil {
		return "", err
	}
	return name, nil
}

func login(w http.ResponseWriter, r *http.Request) {
	newCookie := invalidateCookie()
	success := false
	defer func() {
		http.SetCookie(w, newCookie)
		if !success {
			http.Error(w, "Nope", http.StatusUnauthorized)
			return
		}
	}()
	user := r.PostFormValue("login")
	if len(user) == 0 {
		return
	}

	password := r.PostFormValue("password")
	if len(password) == 0 {
		return
	}

	log.Printf("Validating login request by user %q on %q", user, r.Host)

	// This is very awful, but sadly go-http-auth has no better interface:
	authenticator := auth.NewBasicAuthenticator(realm, auth.HtpasswdFileProvider(htpasswd))
	r.Header["Authorization"] = []string{
		fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", user, password)))),
	}
	if user != authenticator.CheckAuth(r) {
		log.Printf("User %q couldn't be authenticated", user)
		return
	}

	// OK, we're authenticated. Let's get a session started:
	session, err := newSession(r.Host)
	if err != nil {
		return
	}
	newCookie.Value = session
	newCookie.MaxAge = cookieLifetime
	if r.PostFormValue("ephemeral") == "" {
		newCookie.Expires = time.Now().Add(time.Duration(cookieLifetime) * time.Second)
	}
	success = true
}

func main() {
	bind.DefaultBind = "127.0.0.1:8000"
	flag.StringVar(&sessionDir, "sessions", "/var/db/http-auth/cookies", "Directory in which htpasswd-login-form places sessions")
	flag.StringVar(&domain, "domain", "", "Domain to set on all cookies")
	flag.StringVar(&htpasswd, "htpasswd", "/etc/nginx/.htpasswd", "htpasswd file to use for authentication")
	flag.StringVar(&realm, "realm", "example.com", "HTTP Basic auth realm to pretend we run in")
	flag.IntVar(&cookieLifetime, "lifetime", 86400, "Maximum cookie lifetime in seconds")
	bind.WithFlag()
	flag.Parse()

	mux := goji.NewMux()
	mux.HandleFunc(pat.Get("/auth"), checkSession)
	mux.HandleFunc(pat.Post("/login"), login)

	http.Serve(bind.Default(), mux)

}
<commit_msg>Add a session cleanup switch<commit_after>package main

import (
	"errors"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"os"
	"path/filepath"
	"strings"
	"time"

	"encoding/base64"
	"encoding/json"

	"github.com/abbot/go-http-auth"
	"github.com/dchest/uniuri"
	"github.com/zenazn/goji/bind"
	"goji.io"
	"goji.io/pat"
)

const cookieName string = "_htpasswd_auth"
const sessionFormat string = time.RFC3339
const slack time.Duration = time.Duration(10) * time.Second

var sessionDir string
var domain string
var htpasswd string
var realm string
var cookieLifetime int

var validChars map[rune]bool

func init() {
	validChars = map[rune]bool{}
	for _, c := range string(uniuri.StdChars) {
		validChars[c] = true
	}
}

type session struct {
	CreatedRaw string `json:"created"`
	Domain     string

	Name    string    `json:"-"`
	Created time.Time `json:"-"`
}

func (s *session) fill(name string) error {
	t, err := time.Parse(sessionFormat, s.CreatedRaw)
	if err != nil {
		return err
	}
	s.Created = t
	s.Name = name
	return nil
}

func (s *session) Valid() bool {
	sessionExpiry := s.Created.Add(time.Duration(cookieLifetime) * time.Second)
	return sessionExpiry.After(time.Now())
}
func validateCookie(cookie string) (*session, error) {
	for _, c := range cookie {
		if _, ok := validChars[c]; !ok {
			return nil, errors.New("Cookie value contains invalid characters.")
		}
	}
	sessionPath := filepath.Join(sessionDir, cookie)
	base, err := filepath.Rel(sessionDir, sessionPath)
	if err != nil {
		return nil, err
	}
	if strings.HasPrefix(base, ".") || strings.HasPrefix(base, "/") {
		return nil, errors.New("Cookie value would traverse directories")
	}

	// From here on, the session pathname is assumed innocent & handle-able.
	data, err := ioutil.ReadFile(sessionPath)
	if err != nil {
		if os.IsNotExist(err) {
			log.Printf("Request for a session that doesn't exist: %q", sessionPath)
			return nil, err
		} else {
			log.Printf("Checking session %q failed: %s", sessionPath, err)
			return nil, err
		}
	}
	session := session{}
	json.Unmarshal(data, &session)
	err = session.fill(cookie)
	if err != nil {
		os.Remove(sessionPath) // clean up the session
		return nil, err
	}
	if !session.Valid() {
		return nil, errors.New("Session is no longer valid.")
	}
	return &session, nil
}

func invalidateCookie() *http.Cookie {
	return &http.Cookie{Name: cookieName, Value: "nope", Domain: domain, Path: "/", MaxAge: -1, Secure: true, HttpOnly: true}
}

func checkSession(w http.ResponseWriter, r *http.Request) {
	cookie, err := r.Cookie(cookieName)
	newCookie := invalidateCookie()
	success := false

	// Cleanup and ensure we always send a decent response:
	defer func() {
		if !success {
			http.SetCookie(w, newCookie)
			http.Error(w, "Session is not/no longer valid", http.StatusUnauthorized)
			return
		}
		w.WriteHeader(200)
		w.Write([]byte{})
	}()
	if err != nil {
		if err == http.ErrNoCookie {
			return
		}
	}
	session, err := validateCookie(cookie.Value)
	if err != nil {
		log.Printf("Error validating session: %s", err)
		return
	}

	if session.Domain != r.Host {
		log.Printf("Session %q was made for another domain: %s; wanted %q", cookie.Value, session.Domain, r.Host)
		return
	}

	// We have a valid session that hasn't yet expired. Let's call it a success.
	success = true
	return
}

func newSession(domain string) (string, error) {
	name := uniuri.NewLen(90)
	sessionPath := filepath.Join(sessionDir, name)
	data, err := json.Marshal(session{CreatedRaw: time.Now().Add(slack).Format(sessionFormat), Domain: domain})
	if err != nil {
		return "", err
	}
	err = ioutil.WriteFile(sessionPath, data, 0400)
	if err != nil {
		return "", err
	}
	return name, nil
}

func login(w http.ResponseWriter, r *http.Request) {
	newCookie := invalidateCookie()
	success := false
	defer func() {
		http.SetCookie(w, newCookie)
		if !success {
			http.Error(w, "Nope", http.StatusUnauthorized)
			return
		}
	}()
	user := r.PostFormValue("login")
	if len(user) == 0 {
		return
	}

	password := r.PostFormValue("password")
	if len(password) == 0 {
		return
	}

	log.Printf("Validating login request by user %q on %q", user, r.Host)

	// This is very awful, but sadly go-http-auth has no better interface:
	authenticator := auth.NewBasicAuthenticator(realm, auth.HtpasswdFileProvider(htpasswd))
	r.Header["Authorization"] = []string{
		fmt.Sprintf("Basic %s", base64.StdEncoding.EncodeToString([]byte(fmt.Sprintf("%s:%s", user, password)))),
	}
	if user != authenticator.CheckAuth(r) {
		log.Printf("User %q couldn't be authenticated", user)
		return
	}

	// OK, we're authenticated. Let's get a session started:
	session, err := newSession(r.Host)
	if err != nil {
		return
	}
	newCookie.Value = session
	newCookie.MaxAge = cookieLifetime
	if r.PostFormValue("ephemeral") == "" {
		newCookie.Expires = time.Now().Add(time.Duration(cookieLifetime) * time.Second)
	}
	success = true
}

func runCleanup() {
	deleted := 0
	err := filepath.Walk(sessionDir, func(path string, info os.FileInfo, err error) error {
		if err != nil {
			return nil
		}
		if info.IsDir() {
			if path != sessionDir {
				return filepath.SkipDir
			} else {
				return nil
			}
		}
		data, err := ioutil.ReadFile(path)
		if err != nil {
			return err
		}
		var sess session
		err = json.Unmarshal(data, &sess)
		if err != nil {
			fmt.Printf("Session %q was invalid JSON (%s), leaving it", path, err)
			return nil
		}
		if err = sess.fill(filepath.Base(path)); err != nil {
			fmt.Printf("Session %q has an invalid date (%s), deleting it", path, err)
			return os.Remove(path)
		}
		if !sess.Valid() {
			deleted += 1
			return os.Remove(path)
		}
		return nil
	})
	if err != nil {
		log.Printf("Couldn't delete sessions: %s", err)
		os.Exit(2)
	}
	if deleted > 0 {
		log.Printf("Cleaned out %d sessions from %q.", deleted, sessionDir)
	}
}

func main() {
	var cleanup bool

	bind.DefaultBind = "127.0.0.1:8000"
	bind.WithFlag()
	flag.StringVar(&sessionDir, "sessions", "/var/db/http-auth/cookies", "Directory in which htpasswd-login-form places sessions")
	flag.StringVar(&domain, "domain", "", "Domain to set on all cookies")
	flag.StringVar(&htpasswd, "htpasswd", "/etc/nginx/.htpasswd", "htpasswd file to use for authentication")
	flag.StringVar(&realm, "realm", "example.com", "HTTP Basic auth realm to pretend we run in")
	flag.IntVar(&cookieLifetime, "lifetime", 86400, "Maximum cookie lifetime in seconds")
	flag.BoolVar(&cleanup, "cleanup", false, "Perform once-in-a-while cleanup actions")
	flag.Parse()

	if cleanup {
		runCleanup()
		return
	}

	mux := goji.NewMux()
	mux.HandleFunc(pat.Get("/auth"), checkSession)
	mux.HandleFunc(pat.Post("/login"), login)

	http.Serve(bind.Default(), mux)
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/drone/drone-plugin-go/plugin"
)

type Ansible struct {
	Inventory 	string   `json:"inventory"`
	Playbook 	string   `json:"playbook"`
}

func main() {
	workspace := plugin.Workspace{}
	build := plugin.Build{}
	vargs := Ansible{}

	plugin.Param("workspace", &workspace)
	plugin.Param("build", &build)
	plugin.Param("vargs", &vargs)
	plugin.MustParse()

	// Set the Inventory file
	if len(vargs.Inventory) == 0 {
		vargs.Inventory = "hosts"
	}
	// Set the Playbook
	if len(vargs.Playbook) == 0 {
		vargs.Playbook = "playbook.yml"
	}

	// write the rsa private key
	if err := writeKey(workspace); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// Docker environment info
	cmd := exec.Command("ls")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	trace(cmd)
	err := cmd.Run()
	if err != nil {
		os.Exit(1)
	}

	// Docker environment info
	cmd = exec.Command("/usr/bin/ansible-playbook", "-i", vargs.Inventory, vargs.Playbook)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	trace(cmd)
	err = cmd.Run()
	if err != nil {
		os.Exit(1)
	}
}

// Trace writes each command to standard error (preceded by a ‘$ ’) before it
// is executed. Used for debugging your build.
func trace(cmd *exec.Cmd) {
	fmt.Println("$", strings.Join(cmd.Args, " "))
}

// Writes the RSA private key
func writeKey(in plugin.Workspace) error {
	if in.Keys == nil || len(in.Keys.Private) == 0 {
		return nil
	}
	home := "/root"
	u, err := user.Current()
	if err == nil {
		home = u.HomeDir
	}
	sshpath := filepath.Join(home, ".ssh")
	if err := os.MkdirAll(sshpath, 0700); err != nil {
		return err
	}
	confpath := filepath.Join(sshpath, "config")
	privpath := filepath.Join(sshpath, "id_rsa")
	ioutil.WriteFile(confpath, []byte("StrictHostKeyChecking no\n"), 0700)
	return ioutil.WriteFile(privpath, []byte(in.Keys.Private), 0600)
}<commit_msg>add drone/src<commit_after>package main

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"os/user"
	"path/filepath"
	"strings"

	"github.com/drone/drone-plugin-go/plugin"
)

type Ansible struct {
	Inventory 	string   `json:"inventory"`
	Playbook 	string   `json:"playbook"`
}

func main() {
	workspace := plugin.Workspace{}
	build := plugin.Build{}
	vargs := Ansible{}

	plugin.Param("workspace", &workspace)
	plugin.Param("build", &build)
	plugin.Param("vargs", &vargs)
	plugin.MustParse()

	// Set the Inventory file
	if len(vargs.Inventory) == 0 {
		vargs.Inventory = "hosts"
	}
	// Set the Playbook
	if len(vargs.Playbook) == 0 {
		vargs.Playbook = "playbook.yml"
	}

	// write the rsa private key
	if err := writeKey(workspace); err != nil {
		fmt.Println(err)
		os.Exit(1)
	}

	// Docker environment info
	cmd := exec.Command("cd /drone/src && ls")
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	trace(cmd)
	err := cmd.Run()
	if err != nil {
		os.Exit(1)
	}

	inventory_path := fmt.Sprintf("/drone/src/%s", vargs.Inventory)
	playbook_path := fmt.Sprintf("/drone/src/%s", vargs.Playbook)
	// Docker environment info
	cmd = exec.Command("/usr/bin/ansible-playbook", "-i", inventory_path, playbook_path)
	cmd.Stdout = os.Stdout
	cmd.Stderr = os.Stderr
	trace(cmd)
	err = cmd.Run()
	if err != nil {
		os.Exit(1)
	}
}

// Trace writes each command to standard error (preceded by a ‘$ ’) before it
// is executed. Used for debugging your build.
func trace(cmd *exec.Cmd) {
	fmt.Println("$", strings.Join(cmd.Args, " "))
}

// Writes the RSA private key
func writeKey(in plugin.Workspace) error {
	if in.Keys == nil || len(in.Keys.Private) == 0 {
		return nil
	}
	home := "/root"
	u, err := user.Current()
	if err == nil {
		home = u.HomeDir
	}
	sshpath := filepath.Join(home, ".ssh")
	if err := os.MkdirAll(sshpath, 0700); err != nil {
		return err
	}
	confpath := filepath.Join(sshpath, "config")
	privpath := filepath.Join(sshpath, "id_rsa")
	ioutil.WriteFile(confpath, []byte("StrictHostKeyChecking no\n"), 0700)
	return ioutil.WriteFile(privpath, []byte(in.Keys.Private), 0600)
}<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/gorilla/mux"
	"github.com/spf13/viper"

	"github.com/njasm/marionette_client"
)

type Program struct {
	URL      string
	Duration time.Duration
}

// InitializeConfig loads our configuration using Viper package.
func InitializeConfig() {

	viper.SetConfigType("yaml")
	viper.SetConfigName("config")

	viper.AddConfigPath("$HOME/.gotator")
	viper.AddConfigPath(".")

	err := viper.ReadInConfig()
	if err != nil {
		panic(fmt.Errorf("Fatal error config file: %s \n", err))
	}

	viper.SetDefault("debug", false)

	viper.SetEnvPrefix("gorotator") // will be uppercased automatically
	viper.BindEnv("debug")
	viper.BindEnv("firefox_ip")
	viper.BindEnv("firefox_port")
	viper.BindEnv("gotator_port")

	if !viper.IsSet("firefox_ip") || !viper.IsSet("firefox_port") {
		fmt.Fprintln(os.Stderr, "Configuration error.  Both FIREFOX_IP and FIREFOX_PORT must be set via either config or environment.")
		os.Exit(1)
	}
	mode := viper.Get("FIREFOX_MODE")
	ipStr := viper.Get("FIREFOX_IP")
	if mode == 1 {
		log.Println("Using MODE1 (aka. FF Remote Control plugin) -- [DEPRECATED in newer versio of Firefox]")
	}
	if mode == 2 {
		log.Printf("Using MODE2:  Firefox Marionette protocol.  ")
		log.Printf("  IP is %s, but localhost will be used instead.\n", ipStr)
	}

	viper.WatchConfig()
	viper.OnConfigChange(func(e fsnotify.Event) {
		log.Println("\nConfig file changed:", e.Name)
		skip <- struct{}{}
		log.Printf("Content will change immediately.\n\n")

	})

}

// Loads a list of programs.
// A program consists of a list things to display on the rotator along
// with a number of seconds to display each one before moving on.
func loadProgramList(filename string) []Program {

	log.Printf("Loading programs from %s\n", filename)

	var list []Program

	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatal(err)
	}
	webpages := string(bytes)

	r := csv.NewReader(strings.NewReader(webpages))
	r.LazyQuotes = true

	for {
		var p Program
		record, err := r.Read()
		if err == io.EOF {
			log.Println("Finished loading programs.")
			break
		}
		if err != nil {
			log.Printf("Error reading line from program file: %s.  Abandoning attempt to read programs.\n", filename)
			log.Fatal(err)
		}
		p.URL = record[0]
		p.Duration, err = time.ParseDuration(record[1])
		if err != nil {
			log.Println("Program rejected.  Invalid duration.")
		}

		log.Printf("  Loaded program %.50s to show for %s.\n", p.URL, p.Duration)
		list = append(list, p)
	}

	return list
}

func runProgram(program Program) {

	ip := viper.Get("FIREFOX_IP")
	port := viper.GetInt("FIREFOX_PORT")

	constr := fmt.Sprintf("%s:%d", ip, port)

	log.Printf("Running program for %s", program.Duration)
	log.Printf("URL %s", program.URL)

	mode := viper.Get("FIREFOX_MODE")
	if mode == 1 {
		// Connect to FF Remote Control
		conn, err := net.Dial("tcp", constr)
		if err != nil {
			log.Printf("  Error making network connection to: %s\n", constr)
			log.Println("  It is possible Firefox needs to be started or restarted.")
			log.Println("  It is possible FF Remote Control plugin is not installed.")
			log.Println("  Pausing for 30s")
			time.Sleep(30 * time.Second) // wait 30 seconds to slow retries
			return
		}

		// Actual controll of browser starts here
		fmt.Fprintf(conn, "window.location='%s'\n", program.URL)
		status, err := bufio.NewReader(conn).ReadString('\n')
		if err != nil {
			log.Println("ERROR - URL didn't load as desired.")
		}

		var statusParsed interface{}
		err = json.Unmarshal([]byte(status), &statusParsed)

		m := statusParsed.(map[string]interface{})

		if m["result"] == program.URL {
			log.Println("RESULT: OK")
		} else {
			log.Println("RESULT: ERROR - URL didn't load as desired.")
		}
	}
	if mode == 2 {
		// Connect using Marionette
		client := marionette_client.NewClient()

		err := client.Connect("", 0) // this are the default marionette values for hostname, and port
		if err != nil {
			log.Println("Can't connect to firefox.  Sorry.")
			log.Println("It is possible Firefox needs to be started or restarted.")
			log.Println("Pausing for 30s")
			time.Sleep(30 * time.Second) // wait 30 seconds to slow retries
			return
		}
		client.NewSession("", nil) // let marionette generate the Session ID with it's default Capabilities
		client.Navigate(program.URL)

		// experiment with injecting randon script
		// Maybe someday have a CSS overlay to give countdown until rotation?
		// args := []interface{}{}
		// client.ExecuteScript("alert('o hai there!');", args, 1000, false)
	}

	select {
	case <-time.After(program.Duration):
		// Do nothing.
		Unpause()
	case <-skip:
		log.Println("Current program skipped")

		return
	}
}

func Pause() {
	mu.Lock()
	pause = true
	mu.Unlock()
}

func Unpause() {
	mu.Lock()
	pause = false
	mu.Unlock()
}

func IsPaused() bool {
	mu.Lock()
	defer mu.Unlock()

	return pause == true
}

func LoadAndRunLoop() {

	// Load and run the acctive program_file indefinately
	for {
		// We pull filename inside the loop because the
		// configuration can change while our program is running.
		filename := viper.GetString("program_file")

		for IsPaused() {
			time.Sleep(1 * time.Second)
		}

		pl := loadProgramList(filename)

		for _, p := range pl {
			for IsPaused() {
				time.Sleep(1 * time.Second)
			}
			runProgram(p)
		}

		log.Println("Looping back to play program list from beginning")
	}

}

func PlayHandler(w http.ResponseWriter, r *http.Request) {

	r.ParseForm()
	var p Program
	p.URL = r.Form.Get("url")
	log.Printf("URL: %s\n", p.URL)

	d := r.Form.Get("duration")
	log.Printf("Duration: %s\n", d)

	var err error
	p.Duration, err = time.ParseDuration(r.Form.Get("duration"))
	if err != nil {
		w.Write([]byte("Program rejected.  Invalid duration.\n"))
		return
	}

	// Needs validation...

	// Now do something with the program.. play it?

	// Stop normal rotation
	Pause()

	// While paused, go ahead and also skip so that upon resume we move to the next program
	skip <- struct{}{}

	runProgram(p)
	w.Write([]byte("Program accepted\n"))

}

func PauseHandler(w http.ResponseWriter, r *http.Request) {
	Pause()
	log.Println("Paused from web request")
	w.Write([]byte("Ok, paused.\n"))
}

func ResumeHandler(w http.ResponseWriter, r *http.Request) {
	Unpause()
	log.Println("Unpausing from web request")
	w.Write([]byte("Ok, unpaused.\n"))
}
func SkipHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Skippingfrom web request")
	Unpause()
	skip <- struct{}{}

	w.Write([]byte("Skipping current programming and resume program list runner from web request.\n"))
}

func readKeyboardLoop() {
	for {
		os.Stdin.Read(make([]byte, 1)) // read a single byte
		log.Printf(" >> Got keyboard input, that means you want to move to the next program.  Can do! << \n\n")
		Unpause()
		skip <- struct{}{}
	}
}

// Control channel to stop running programs immediately (yes, global)

var skip = make(chan struct{})
var exitprogram = make(chan struct{})
var pause bool
var mu = &sync.Mutex{}
var version = "0.0.5"

func main() {

	if len(os.Args) > 1 {
		if os.Args[1] == "version" {
			log.Println("Gotator version:", version)
			os.Exit(0)
		}
	}
	log.Println("Starting gotator: version", version)
	Unpause()

	InitializeConfig()

	go LoadAndRunLoop()

	if viper.IsSet("interactive") && viper.Get("interactive") == true {
		go readKeyboardLoop()
	}

	if viper.IsSet("apienabled") && viper.Get("apienabled") == true {
		listen_port := ":8080"
		if viper.IsSet("gotator_port") {
			listen_port = ":" + viper.GetString("gotator_port")
		}

		log.Printf("Starting API server on port %s.  Notice:  This allows UNAUTHENTICATED remote control of Firefox. set 'apienabled: false' in config.yaml to disable.\n",
			listen_port)

		r := mux.NewRouter()
		r.HandleFunc("/play", PlayHandler)
		r.HandleFunc("/pause", PauseHandler)
		r.HandleFunc("/resume", ResumeHandler)
		r.HandleFunc("/skip", SkipHandler)

		go log.Fatal(http.ListenAndServe(listen_port, r))
	} else {
		log.Println("notice: rest API not enabled in configuration and will be unavailable.  set 'apienabled: true' in config.yaml if you want to use it.\n")
		// If we aren't doing http.ListenAndServe() we need to block here or else gotator would exit immediately
		<-exitprogram
	}

}
<commit_msg>Typos, comments and logging<commit_after>package main

import (
	"bufio"
	"encoding/csv"
	"encoding/json"
	"fmt"
	"io"
	"io/ioutil"
	"log"
	"net"
	"net/http"
	"os"
	"strings"
	"sync"
	"time"

	"github.com/fsnotify/fsnotify"
	"github.com/gorilla/mux"
	"github.com/spf13/viper"

	"github.com/njasm/marionette_client"
)

type Program struct {
	URL      string
	Duration time.Duration
}

// InitializeConfig loads our configuration using Viper package.
func InitializeConfig() {

	viper.SetConfigType("yaml")
	viper.SetConfigName("config")

	viper.AddConfigPath("$HOME/.gotator")
	viper.AddConfigPath(".")

	err := viper.ReadInConfig()
	if err != nil {
		panic(fmt.Errorf("Fatal error config file: %s \n", err))
	}

	viper.SetDefault("debug", false)

	viper.SetEnvPrefix("gorotator") // will be uppercased automatically
	viper.BindEnv("debug")
	viper.BindEnv("firefox_ip")
	viper.BindEnv("firefox_port")
	viper.BindEnv("gotator_port")

	if !viper.IsSet("firefox_ip") || !viper.IsSet("firefox_port") {
		fmt.Fprintln(os.Stderr, "Configuration error.  Both FIREFOX_IP and FIREFOX_PORT must be set via either config or environment.")
		os.Exit(1)
	}
	mode := viper.Get("FIREFOX_MODE")
	ipStr := viper.Get("FIREFOX_IP")
	if mode == 1 {
		log.Println("Using MODE1 (aka. FF Remote Control plugin) -- [DEPRECATED in newer versio of Firefox]")
	}
	if mode == 2 {
		log.Printf("Using MODE2:  Firefox Marionette protocol.  ")
		log.Printf("  IP is %s, but localhost will be used instead.\n", ipStr)
	}

	viper.WatchConfig()
	viper.OnConfigChange(func(e fsnotify.Event) {
		log.Println("\nConfig file changed:", e.Name)
		skip <- struct{}{}
		log.Printf("Content will change immediately.\n\n")

	})

}

// Loads a list of programs.
// A program consists of a list things to display on the rotator along
// with a number of seconds to display each one before moving on.
func loadProgramList(filename string) []Program {

	log.Printf("Loading programs from %s\n", filename)

	var list []Program

	bytes, err := ioutil.ReadFile(filename)
	if err != nil {
		log.Fatal(err)
	}
	webpages := string(bytes)

	r := csv.NewReader(strings.NewReader(webpages))
	r.LazyQuotes = true

	for {
		var p Program
		record, err := r.Read()
		if err == io.EOF {
			log.Println("Finished loading programs.")
			break
		}
		if err != nil {
			log.Printf("Error reading line from program file: %s.  Abandoning attempt to read programs.\n", filename)
			log.Fatal(err)
		}
		p.URL = record[0]
		p.Duration, err = time.ParseDuration(record[1])
		if err != nil {
			log.Println("Program rejected.  Invalid duration.")
		}

		log.Printf("  Loaded program %.50s to show for %s.\n", p.URL, p.Duration)
		list = append(list, p)
	}

	return list
}

func runProgram(program Program) {

	ip := viper.Get("FIREFOX_IP")
	port := viper.GetInt("FIREFOX_PORT")

	constr := fmt.Sprintf("%s:%d", ip, port)

	log.Printf("Running program for %s", program.Duration)
	log.Printf("URL %s", program.URL)

	mode := viper.Get("FIREFOX_MODE")
	if mode == 1 {
		// Connect to FF Remote Control
		conn, err := net.Dial("tcp", constr)
		if err != nil {
			log.Printf("  Error making network connection to: %s\n", constr)
			log.Println("  It is possible Firefox needs to be started or restarted.")
			log.Println("  It is possible FF Remote Control plugin is not installed.")
			log.Println("  Pausing for 30s")
			time.Sleep(30 * time.Second) // wait 30 seconds to slow retries
			return
		}

		// Actual control of browser starts here
		fmt.Fprintf(conn, "window.location='%s'\n", program.URL)
		status, err := bufio.NewReader(conn).ReadString('\n')
		if err != nil {
			log.Println("ERROR - URL didn't load as desired.")
		}

		var statusParsed interface{}
		err = json.Unmarshal([]byte(status), &statusParsed)

		m := statusParsed.(map[string]interface{})

		if m["result"] == program.URL {
			log.Println("RESULT: OK")
		} else {
			log.Println("RESULT: ERROR - URL didn't load as desired.")
		}
	}
	if mode == 2 {
		// Connect using Marionette
		client := marionette_client.NewClient()

		err := client.Connect("", 0) // this are the default marionette values for hostname, and port
		if err != nil {
			log.Println("Can't connect to firefox.  Sorry.")
			log.Println("It is possible Firefox needs to be started or restarted.")
			log.Println("Pausing for 30s")
			time.Sleep(30 * time.Second) // wait 30 seconds to slow retries
			return
		}
		client.NewSession("", nil) // let marionette generate the Session ID with it's default Capabilities
		client.Navigate(program.URL)

		// experiment with injecting randon script
		// Maybe someday have a CSS overlay to give countdown until rotation?
		// args := []interface{}{}
		// client.ExecuteScript("alert('o hai there!');", args, 1000, false)
	}

	select {
	case <-time.After(program.Duration):
		Unpause()
	case <-skip:
		log.Println("Current program skipped")

		return
	}
}

func Pause() {
	mu.Lock()
	pause = true
	mu.Unlock()
	log.Println("Paused")
}

func Unpause() {
	mu.Lock()
	pause = false
	mu.Unlock()
	log.Println("Unpaused")
}

func IsPaused() bool {
	mu.Lock()
	defer mu.Unlock()

	return pause == true
}

func LoadAndRunLoop() {

	// Load and run the acctive program_file indefinately
	for {
		// We pull filename inside the loop because the
		// configuration can change while our program is running.
		filename := viper.GetString("program_file")

		for IsPaused() {
			fmt.Printf(".")
			time.Sleep(1 * time.Second)
		}

		pl := loadProgramList(filename)

		for _, p := range pl {
			for IsPaused() {
				fmt.Printf("X")
				time.Sleep(1 * time.Second)
			}
			runProgram(p)
		}

		log.Println("Looping back to play program list from beginning")
	}

}

func PlayHandler(w http.ResponseWriter, r *http.Request) {

	r.ParseForm()
	var p Program
	p.URL = r.Form.Get("url")
	log.Printf("URL: %s\n", p.URL)

	d := r.Form.Get("duration")
	log.Printf("Duration: %s\n", d)

	// CAREFUL: There may be bugs here...
	var err error
	p.Duration, err = time.ParseDuration(r.Form.Get("duration"))
	if err != nil {
		w.Write([]byte("Program rejected.  Invalid duration.\n"))
		return
	}

	// Stop normal rotation
	Pause()

	// While paused, go ahead and also skip so that upon resume we move to the next program
	skip <- struct{}{}

	runProgram(p)
	w.Write([]byte("Program accepted\n"))

}

func PauseHandler(w http.ResponseWriter, r *http.Request) {
	Pause()
	log.Println("Paused from web request")
	w.Write([]byte("Ok, paused.\n"))
}

func ResumeHandler(w http.ResponseWriter, r *http.Request) {
	Unpause()
	log.Println("Unpausing from web request")
	w.Write([]byte("Ok, unpaused.\n"))
}
func SkipHandler(w http.ResponseWriter, r *http.Request) {
	log.Println("Skipping from web request")
	Unpause()
	skip <- struct{}{}

	w.Write([]byte("Skipping current programming and resume program list runner from web request.\n"))
}

func readKeyboardLoop() {
	for {
		os.Stdin.Read(make([]byte, 1)) // read a single byte
		log.Printf(" >> Got keyboard input, that means you want to move to the next program.  Can do! << \n\n")
		Unpause()
		skip <- struct{}{}
	}
}

// Control channel to stop running programs immediately (yes, global)

var skip = make(chan struct{})
var exitprogram = make(chan struct{})
var pause bool
var mu = &sync.Mutex{}
var version = "0.0.5"

func main() {

	if len(os.Args) > 1 {
		if os.Args[1] == "version" {
			log.Println("Gotator version:", version)
			os.Exit(0)
		}
	}
	log.Println("Starting gotator: version", version)
	Unpause()

	InitializeConfig()

	go LoadAndRunLoop()

	if viper.IsSet("interactive") && viper.Get("interactive") == true {
		go readKeyboardLoop()
	}

	if viper.IsSet("apienabled") && viper.Get("apienabled") == true {
		listen_port := ":8080"
		if viper.IsSet("gotator_port") {
			listen_port = ":" + viper.GetString("gotator_port")
		}

		log.Printf("Starting API server on port %s.  Notice:  This allows UNAUTHENTICATED remote control of Firefox. set 'apienabled: false' in config.yaml to disable.\n",
			listen_port)

		r := mux.NewRouter()
		r.HandleFunc("/play", PlayHandler)
		r.HandleFunc("/pause", PauseHandler)
		r.HandleFunc("/resume", ResumeHandler)
		r.HandleFunc("/skip", SkipHandler)

		go log.Fatal(http.ListenAndServe(listen_port, r))
	} else {
		log.Println("notice: rest API not enabled in configuration and will be unavailable.  set 'apienabled: true' in config.yaml if you want to use it.\n")
		// If we aren't doing http.ListenAndServe() we need to block here or else gotator would exit immediately
		<-exitprogram
	}

}
<|endoftext|>


<commit_before>package gorm

import (
	"database/sql"
)

type DB struct {
	Value         interface{}
	Error         error
	RowsAffected  int64
	callback      *callback
	db            sqlCommon
	parent        *DB
	search        *search
	logMode       int
	logger        logger
	dialect       Dialect
	tagIdentifier string
	singularTable bool
	source        string
}

func Open(driver, source string) (DB, error) {
	var err error
	db := DB{dialect: NewDialect(driver), tagIdentifier: "sql", logger: defaultLogger, callback: DefaultCallback, source: source}
	db.db, err = sql.Open(driver, source)
	db.parent = &db
	return db, err
}

func (s *DB) Close() error {
	return s.parent.db.(*sql.DB).Close()
}

func (s *DB) DB() *sql.DB {
	return s.db.(*sql.DB)
}

// Return the underlying sql.DB or sql.Tx instance.
// Use of this method is discouraged. It's mainly intended to allow
// coexistence with legacy non-GORM code.
func (s *DB) CommonDB() sqlCommon {
	return s.db
}

func (s *DB) Callback() *callback {
	s.parent.callback = s.parent.callback.clone()
	return s.parent.callback
}

func (s *DB) SetTagIdentifier(str string) {
	s.parent.tagIdentifier = str
}

func (s *DB) SetLogger(l logger) {
	s.parent.logger = l
}

func (s *DB) LogMode(b bool) *DB {
	if b {
		s.logMode = 2
	} else {
		s.logMode = 1
	}
	return s
}

func (s *DB) SingularTable(b bool) {
	s.parent.singularTable = b
}

func (s *DB) Where(query interface{}, args ...interface{}) *DB {
	return s.clone().search.where(query, args...).db
}

func (s *DB) Or(query interface{}, args ...interface{}) *DB {
	return s.clone().search.or(query, args...).db
}

func (s *DB) Not(query interface{}, args ...interface{}) *DB {
	return s.clone().search.not(query, args...).db
}

func (s *DB) Limit(value interface{}) *DB {
	return s.clone().search.limit(value).db
}

func (s *DB) Offset(value interface{}) *DB {
	return s.clone().search.offset(value).db
}

func (s *DB) Order(value string, reorder ...bool) *DB {
	return s.clone().search.order(value, reorder...).db
}

func (s *DB) Select(value interface{}) *DB {
	return s.clone().search.selects(value).db
}

func (s *DB) Group(query string) *DB {
	return s.clone().search.group(query).db
}

func (s *DB) Having(query string, values ...interface{}) *DB {
	return s.clone().search.having(query, values...).db
}

func (s *DB) Joins(query string) *DB {
	return s.clone().search.joins(query).db
}

func (s *DB) Includes(value interface{}) *DB {
	return s.clone().search.includes(value).db
}

func (s *DB) Scopes(funcs ...func(*DB) *DB) *DB {
	c := s
	for _, f := range funcs {
		c = f(c)
	}
	return c
}

func (s *DB) Unscoped() *DB {
	return s.clone().search.unscoped().db
}

func (s *DB) Attrs(attrs ...interface{}) *DB {
	return s.clone().search.attrs(attrs...).db
}

func (s *DB) Assign(attrs ...interface{}) *DB {
	return s.clone().search.assign(attrs...).db
}

func (s *DB) First(out interface{}, where ...interface{}) *DB {
	scope := s.clone().NewScope(out)
	scope.Search = scope.Search.clone().order(scope.TableName() + "." + scope.PrimaryKey()).limit(1)
	return scope.inlineCondition(where...).callCallbacks(s.parent.callback.queries).db
}

func (s *DB) Last(out interface{}, where ...interface{}) *DB {
	scope := s.clone().NewScope(out)
	scope.Search = scope.Search.clone().order(scope.TableName() + "." + scope.PrimaryKey() + " DESC").limit(1)
	return scope.inlineCondition(where...).callCallbacks(s.parent.callback.queries).db
}

func (s *DB) Find(out interface{}, where ...interface{}) *DB {
	return s.clone().NewScope(out).inlineCondition(where...).callCallbacks(s.parent.callback.queries).db
}

func (s *DB) Row() *sql.Row {
	return s.NewScope(s.Value).row()
}

func (s *DB) Rows() (*sql.Rows, error) {
	return s.NewScope(s.Value).rows()
}

func (s *DB) Scan(dest interface{}) *DB {
	scope := s.clone().NewScope(s.Value).Set("gorm:query_destination", dest)
	Query(scope)
	return scope.db
}

func (s *DB) FirstOrInit(out interface{}, where ...interface{}) *DB {
	c := s.clone()
	r := c.First(out, where...)
	if r.Error != nil {
		if !r.RecordNotFound() {
			return r
		}
		c.NewScope(out).inlineCondition(where...).initialize()
	} else {
		c.NewScope(out).updatedAttrsWithValues(convertInterfaceToMap(s.search.AssignAttrs), false)
	}
	return c
}

func (s *DB) FirstOrCreate(out interface{}, where ...interface{}) *DB {
	c := s.clone()
	r := c.First(out, where...)
	if r.Error != nil {
		if !r.RecordNotFound() {
			return r
		}
		c.NewScope(out).inlineCondition(where...).initialize().callCallbacks(s.parent.callback.creates)
	} else if len(s.search.AssignAttrs) > 0 {
		c.NewScope(out).Set("gorm:update_interface", s.search.AssignAttrs).callCallbacks(s.parent.callback.updates)
	}
	return c
}

func (s *DB) Update(attrs ...interface{}) *DB {
	return s.Updates(toSearchableMap(attrs...), true)
}

func (s *DB) Updates(values interface{}, ignoreProtectedAttrs ...bool) *DB {
	return s.clone().NewScope(s.Value).
		Set("gorm:update_interface", values).
		Set("gorm:ignore_protected_attrs", len(ignoreProtectedAttrs) > 0).
		callCallbacks(s.parent.callback.updates).db
}

func (s *DB) UpdateColumn(attrs ...interface{}) *DB {
	return s.UpdateColumns(toSearchableMap(attrs...))
}

func (s *DB) UpdateColumns(values interface{}) *DB {
	return s.clone().NewScope(s.Value).
		Set("gorm:update_interface", values).
		Set("gorm:update_column", true).
		callCallbacks(s.parent.callback.updates).db
}

func (s *DB) Save(value interface{}) *DB {
	scope := s.clone().NewScope(value)
	if scope.PrimaryKeyZero() {
		return scope.callCallbacks(s.parent.callback.creates).db
	} else {
		return scope.callCallbacks(s.parent.callback.updates).db
	}
}

func (s *DB) Create(value interface{}) *DB {
	scope := s.clone().NewScope(value)
	return scope.callCallbacks(s.parent.callback.creates).db
}

func (s *DB) Delete(value interface{}) *DB {
	return s.clone().NewScope(value).callCallbacks(s.parent.callback.deletes).db
}

func (s *DB) Raw(sql string, values ...interface{}) *DB {
	return s.clone().search.raw(true).where(sql, values...).db
}

func (s *DB) Exec(sql string, values ...interface{}) *DB {
	scope := s.clone().NewScope(nil)
	scope.Raw(scope.buildWhereCondition(map[string]interface{}{"query": sql, "args": values}))
	return scope.Exec().db
}

func (s *DB) Model(value interface{}) *DB {
	c := s.clone()
	c.Value = value
	return c
}

func (s *DB) Related(value interface{}, foreignKeys ...string) *DB {
	return s.clone().NewScope(s.Value).related(value, foreignKeys...).db
}

func (s *DB) Pluck(column string, value interface{}) *DB {
	return s.NewScope(s.Value).pluck(column, value).db
}

func (s *DB) Count(value interface{}) *DB {
	return s.NewScope(s.Value).count(value).db
}

func (s *DB) Table(name string) *DB {
	return s.clone().search.table(name).db
}

func (s *DB) Debug() *DB {
	return s.clone().LogMode(true)
}

func (s *DB) Begin() *DB {
	c := s.clone()
	if db, ok := c.db.(sqlDb); ok {
		tx, err := db.Begin()
		c.db = interface{}(tx).(sqlCommon)
		c.err(err)
	} else {
		c.err(CantStartTransaction)
	}
	return c
}

func (s *DB) Commit() *DB {
	if db, ok := s.db.(sqlTx); ok {
		s.err(db.Commit())
	} else {
		s.err(NoValidTransaction)
	}
	return s
}

func (s *DB) Rollback() *DB {
	if db, ok := s.db.(sqlTx); ok {
		s.err(db.Rollback())
	} else {
		s.err(NoValidTransaction)
	}
	return s
}

func (s *DB) NewRecord(value interface{}) bool {
	return s.clone().NewScope(value).PrimaryKeyZero()
}

func (s *DB) RecordNotFound() bool {
	return s.Error == RecordNotFound
}

// Migrations
func (s *DB) CreateTable(value interface{}) *DB {
	return s.clone().NewScope(value).createTable().db
}

func (s *DB) DropTable(value interface{}) *DB {
	return s.clone().NewScope(value).dropTable().db
}

func (s *DB) AutoMigrate(value interface{}) *DB {
	return s.clone().NewScope(value).autoMigrate().db
}

func (s *DB) ModifyColumn(column string, typ string) *DB {
	s.clone().NewScope(s.Value).modifyColumn(column, typ)
	return s
}

func (s *DB) DropColumn(column string) *DB {
	s.clone().NewScope(s.Value).dropColumn(column)
	return s
}

func (s *DB) AddIndex(indexName string, column ...string) *DB {
	s.clone().NewScope(s.Value).addIndex(false, indexName, column...)
	return s
}

func (s *DB) AddUniqueIndex(indexName string, column ...string) *DB {
	s.clone().NewScope(s.Value).addIndex(true, indexName, column...)
	return s
}

func (s *DB) RemoveIndex(column string) *DB {
	s.clone().NewScope(s.Value).removeIndex(column)
	return s
}
<commit_msg>Fix exception for FirstOrCreate when search condition is nil<commit_after>package gorm

import (
	"database/sql"
)

type DB struct {
	Value         interface{}
	Error         error
	RowsAffected  int64
	callback      *callback
	db            sqlCommon
	parent        *DB
	search        *search
	logMode       int
	logger        logger
	dialect       Dialect
	tagIdentifier string
	singularTable bool
	source        string
}

func Open(driver, source string) (DB, error) {
	var err error
	db := DB{dialect: NewDialect(driver), tagIdentifier: "sql", logger: defaultLogger, callback: DefaultCallback, source: source}
	db.db, err = sql.Open(driver, source)
	db.parent = &db
	return db, err
}

func (s *DB) Close() error {
	return s.parent.db.(*sql.DB).Close()
}

func (s *DB) DB() *sql.DB {
	return s.db.(*sql.DB)
}

// Return the underlying sql.DB or sql.Tx instance.
// Use of this method is discouraged. It's mainly intended to allow
// coexistence with legacy non-GORM code.
func (s *DB) CommonDB() sqlCommon {
	return s.db
}

func (s *DB) Callback() *callback {
	s.parent.callback = s.parent.callback.clone()
	return s.parent.callback
}

func (s *DB) SetTagIdentifier(str string) {
	s.parent.tagIdentifier = str
}

func (s *DB) SetLogger(l logger) {
	s.parent.logger = l
}

func (s *DB) LogMode(b bool) *DB {
	if b {
		s.logMode = 2
	} else {
		s.logMode = 1
	}
	return s
}

func (s *DB) SingularTable(b bool) {
	s.parent.singularTable = b
}

func (s *DB) Where(query interface{}, args ...interface{}) *DB {
	return s.clone().search.where(query, args...).db
}

func (s *DB) Or(query interface{}, args ...interface{}) *DB {
	return s.clone().search.or(query, args...).db
}

func (s *DB) Not(query interface{}, args ...interface{}) *DB {
	return s.clone().search.not(query, args...).db
}

func (s *DB) Limit(value interface{}) *DB {
	return s.clone().search.limit(value).db
}

func (s *DB) Offset(value interface{}) *DB {
	return s.clone().search.offset(value).db
}

func (s *DB) Order(value string, reorder ...bool) *DB {
	return s.clone().search.order(value, reorder...).db
}

func (s *DB) Select(value interface{}) *DB {
	return s.clone().search.selects(value).db
}

func (s *DB) Group(query string) *DB {
	return s.clone().search.group(query).db
}

func (s *DB) Having(query string, values ...interface{}) *DB {
	return s.clone().search.having(query, values...).db
}

func (s *DB) Joins(query string) *DB {
	return s.clone().search.joins(query).db
}

func (s *DB) Includes(value interface{}) *DB {
	return s.clone().search.includes(value).db
}

func (s *DB) Scopes(funcs ...func(*DB) *DB) *DB {
	c := s
	for _, f := range funcs {
		c = f(c)
	}
	return c
}

func (s *DB) Unscoped() *DB {
	return s.clone().search.unscoped().db
}

func (s *DB) Attrs(attrs ...interface{}) *DB {
	return s.clone().search.attrs(attrs...).db
}

func (s *DB) Assign(attrs ...interface{}) *DB {
	return s.clone().search.assign(attrs...).db
}

func (s *DB) First(out interface{}, where ...interface{}) *DB {
	scope := s.clone().NewScope(out)
	scope.Search = scope.Search.clone().order(scope.TableName() + "." + scope.PrimaryKey()).limit(1)
	return scope.inlineCondition(where...).callCallbacks(s.parent.callback.queries).db
}

func (s *DB) Last(out interface{}, where ...interface{}) *DB {
	scope := s.clone().NewScope(out)
	scope.Search = scope.Search.clone().order(scope.TableName() + "." + scope.PrimaryKey() + " DESC").limit(1)
	return scope.inlineCondition(where...).callCallbacks(s.parent.callback.queries).db
}

func (s *DB) Find(out interface{}, where ...interface{}) *DB {
	return s.clone().NewScope(out).inlineCondition(where...).callCallbacks(s.parent.callback.queries).db
}

func (s *DB) Row() *sql.Row {
	return s.NewScope(s.Value).row()
}

func (s *DB) Rows() (*sql.Rows, error) {
	return s.NewScope(s.Value).rows()
}

func (s *DB) Scan(dest interface{}) *DB {
	scope := s.clone().NewScope(s.Value).Set("gorm:query_destination", dest)
	Query(scope)
	return scope.db
}

func (s *DB) FirstOrInit(out interface{}, where ...interface{}) *DB {
	c := s.clone()
	r := c.First(out, where...)
	if r.Error != nil {
		if !r.RecordNotFound() {
			return r
		}
		c.NewScope(out).inlineCondition(where...).initialize()
	} else {
		c.NewScope(out).updatedAttrsWithValues(convertInterfaceToMap(s.search.AssignAttrs), false)
	}
	return c
}

func (s *DB) FirstOrCreate(out interface{}, where ...interface{}) *DB {
	c := s.clone()
	r := c.First(out, where...)
	if r.Error != nil {
		if !r.RecordNotFound() {
			return r
		}
		c.NewScope(out).inlineCondition(where...).initialize().callCallbacks(s.parent.callback.creates)
	} else if len(c.search.AssignAttrs) > 0 {
		c.NewScope(out).Set("gorm:update_interface", s.search.AssignAttrs).callCallbacks(s.parent.callback.updates)
	}
	return c
}

func (s *DB) Update(attrs ...interface{}) *DB {
	return s.Updates(toSearchableMap(attrs...), true)
}

func (s *DB) Updates(values interface{}, ignoreProtectedAttrs ...bool) *DB {
	return s.clone().NewScope(s.Value).
		Set("gorm:update_interface", values).
		Set("gorm:ignore_protected_attrs", len(ignoreProtectedAttrs) > 0).
		callCallbacks(s.parent.callback.updates).db
}

func (s *DB) UpdateColumn(attrs ...interface{}) *DB {
	return s.UpdateColumns(toSearchableMap(attrs...))
}

func (s *DB) UpdateColumns(values interface{}) *DB {
	return s.clone().NewScope(s.Value).
		Set("gorm:update_interface", values).
		Set("gorm:update_column", true).
		callCallbacks(s.parent.callback.updates).db
}

func (s *DB) Save(value interface{}) *DB {
	scope := s.clone().NewScope(value)
	if scope.PrimaryKeyZero() {
		return scope.callCallbacks(s.parent.callback.creates).db
	} else {
		return scope.callCallbacks(s.parent.callback.updates).db
	}
}

func (s *DB) Create(value interface{}) *DB {
	scope := s.clone().NewScope(value)
	return scope.callCallbacks(s.parent.callback.creates).db
}

func (s *DB) Delete(value interface{}) *DB {
	return s.clone().NewScope(value).callCallbacks(s.parent.callback.deletes).db
}

func (s *DB) Raw(sql string, values ...interface{}) *DB {
	return s.clone().search.raw(true).where(sql, values...).db
}

func (s *DB) Exec(sql string, values ...interface{}) *DB {
	scope := s.clone().NewScope(nil)
	scope.Raw(scope.buildWhereCondition(map[string]interface{}{"query": sql, "args": values}))
	return scope.Exec().db
}

func (s *DB) Model(value interface{}) *DB {
	c := s.clone()
	c.Value = value
	return c
}

func (s *DB) Related(value interface{}, foreignKeys ...string) *DB {
	return s.clone().NewScope(s.Value).related(value, foreignKeys...).db
}

func (s *DB) Pluck(column string, value interface{}) *DB {
	return s.NewScope(s.Value).pluck(column, value).db
}

func (s *DB) Count(value interface{}) *DB {
	return s.NewScope(s.Value).count(value).db
}

func (s *DB) Table(name string) *DB {
	return s.clone().search.table(name).db
}

func (s *DB) Debug() *DB {
	return s.clone().LogMode(true)
}

func (s *DB) Begin() *DB {
	c := s.clone()
	if db, ok := c.db.(sqlDb); ok {
		tx, err := db.Begin()
		c.db = interface{}(tx).(sqlCommon)
		c.err(err)
	} else {
		c.err(CantStartTransaction)
	}
	return c
}

func (s *DB) Commit() *DB {
	if db, ok := s.db.(sqlTx); ok {
		s.err(db.Commit())
	} else {
		s.err(NoValidTransaction)
	}
	return s
}

func (s *DB) Rollback() *DB {
	if db, ok := s.db.(sqlTx); ok {
		s.err(db.Rollback())
	} else {
		s.err(NoValidTransaction)
	}
	return s
}

func (s *DB) NewRecord(value interface{}) bool {
	return s.clone().NewScope(value).PrimaryKeyZero()
}

func (s *DB) RecordNotFound() bool {
	return s.Error == RecordNotFound
}

// Migrations
func (s *DB) CreateTable(value interface{}) *DB {
	return s.clone().NewScope(value).createTable().db
}

func (s *DB) DropTable(value interface{}) *DB {
	return s.clone().NewScope(value).dropTable().db
}

func (s *DB) AutoMigrate(value interface{}) *DB {
	return s.clone().NewScope(value).autoMigrate().db
}

func (s *DB) ModifyColumn(column string, typ string) *DB {
	s.clone().NewScope(s.Value).modifyColumn(column, typ)
	return s
}

func (s *DB) DropColumn(column string) *DB {
	s.clone().NewScope(s.Value).dropColumn(column)
	return s
}

func (s *DB) AddIndex(indexName string, column ...string) *DB {
	s.clone().NewScope(s.Value).addIndex(false, indexName, column...)
	return s
}

func (s *DB) AddUniqueIndex(indexName string, column ...string) *DB {
	s.clone().NewScope(s.Value).addIndex(true, indexName, column...)
	return s
}

func (s *DB) RemoveIndex(column string) *DB {
	s.clone().NewScope(s.Value).removeIndex(column)
	return s
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/exercism/cli/configuration"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	app := cli.NewApp()
	app.Name = "exercism"
	app.Usage = "A command line tool to interact with http://exercism.io"
	app.Version = VERSION
	app.Commands = []cli.Command{
		{
			Name:      "current",
			ShortName: "c",
			Usage:     "Show the current assignments",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				currentAssignments, err := FetchAssignments(config, FetchEndpoints["current"])
				if err != nil {
					fmt.Println(err)
					return
				}

				fmt.Println("Current Assignments")

				for _, a := range currentAssignments {
					fmt.Printf("%v: %v\n", strings.Title(a.Track), a.Slug)
				}
			},
		},
		{
			Name:      "demo",
			ShortName: "d",
			Usage:     "Fetch first assignment for each language from exercism.io",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					config, err = configuration.Demo()
					if err != nil {
						fmt.Println(err)
						return
					}
				}
				assignments, err := FetchAssignments(config, FetchEndpoints["demo"])
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "fetch",
			ShortName: "f",
			Usage:     "Fetch current assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := FetchAssignments(config,
					FetchEndpoints["current"])
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}

				fmt.Printf("Exercises written to %s\n", config.ExercismDirectory)
			},
		},
		{
			Name:      "login",
			ShortName: "l",
			Usage:     "Save exercism.io api credentials",
			Action: func(c *cli.Context) {
				configuration.ToFile(configuration.HomeDir(), askForConfigInfo())
			},
		},
		{
			Name:      "logout",
			ShortName: "o",
			Usage:     "Clear exercism.io api credentials",
			Action: func(c *cli.Context) {
				Logout(configuration.HomeDir())
			},
		},
		{
			Name:      "peek",
			ShortName: "p",
			Usage:     "Fetch upcoming assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := FetchAssignments(config,
					FetchEndpoints["next"])
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "submit",
			ShortName: "s",
			Usage:     "Submit code to exercism.io on your current assignment",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				if len(c.Args()) == 0 {
					fmt.Println("Please enter a file name")
					return
				}

				filename := c.Args()[0]

				// Make filename relative to config.ExercismDirectory.
				absPath, err := absolutePath(filename)
				if err != nil {
					fmt.Printf("Couldn't find %v: %v\n", filename, err)
					return
				}
				exDir := config.ExercismDirectory + string(filepath.Separator)
				if !strings.HasPrefix(absPath, exDir) {
					fmt.Printf("%v is not under your exercism project path (%v)\n", absPath, exDir)
					return
				}
				filename = absPath[len(exDir):]

				if IsTest(filename) {
					fmt.Println("It looks like this is a test, please enter an example file name.")
					return
				}

				code, err := ioutil.ReadFile(absPath)
				if err != nil {
					fmt.Printf("Error reading %v: %v\n", absPath, err)
					return
				}

				response, err := SubmitAssignment(config, filename, code)
				if err != nil {
					fmt.Printf("There was an issue with your submission: %v\n", err)
					return
				}

				fmt.Printf("For feedback on your submission visit %s%s%s.\n",
					config.Hostname, "/submissions/", response.Id)

			},
		},
		{
			Name:      "unsubmit",
			ShortName: "u",
			Usage:     "Delete the last submission",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				response, err := UnsubmitAssignment(config)
				if err != nil {
					fmt.Println(err)
					return
				}

				if response != "" {
					return
				}

				fmt.Println("The last submission was successfully deleted.")
			},
		},
		{
			Name:      "whoami",
			ShortName: "w",
			Usage:     "Get the github username that you are logged in as",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				fmt.Println(config.GithubUsername)
			},
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		fmt.Errorf("%v", err)
		os.Exit(1)
	}
}

func askForConfigInfo() (c configuration.Config) {
	var un, key, dir string

	currentDir, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	fmt.Print("Your GitHub username: ")
	_, err = fmt.Scanln(&un)
	if err != nil {
		panic(err)
	}

	fmt.Print("Your exercism.io API key: ")
	_, err = fmt.Scanln(&key)
	if err != nil {
		panic(err)
	}

	fmt.Println("What is your exercism exercises project path?")
	fmt.Printf("Press Enter to select the default (%s):\n", currentDir)
	fmt.Print("> ")
	_, err = fmt.Scanln(&dir)
	if err != nil && err.Error() != "unexpected newline" {
		panic(err)
	}
	dir, err = absolutePath(dir)
	if err != nil {
		panic(err)
	}

	if dir == "" {
		dir = currentDir
	}

	return configuration.Config{GithubUsername: un, ApiKey: key, ExercismDirectory: configuration.ReplaceTilde(dir), Hostname: "http://exercism.io"}
}

func absolutePath(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	return filepath.EvalSymlinks(path)
}
<commit_msg>Create exercism directory on login. Fixes #28<commit_after>package main

import (
	"fmt"
	"github.com/codegangsta/cli"
	"github.com/exercism/cli/configuration"
	"io/ioutil"
	"os"
	"path/filepath"
	"strings"
)

func main() {
	app := cli.NewApp()
	app.Name = "exercism"
	app.Usage = "A command line tool to interact with http://exercism.io"
	app.Version = VERSION
	app.Commands = []cli.Command{
		{
			Name:      "current",
			ShortName: "c",
			Usage:     "Show the current assignments",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				currentAssignments, err := FetchAssignments(config, FetchEndpoints["current"])
				if err != nil {
					fmt.Println(err)
					return
				}

				fmt.Println("Current Assignments")

				for _, a := range currentAssignments {
					fmt.Printf("%v: %v\n", strings.Title(a.Track), a.Slug)
				}
			},
		},
		{
			Name:      "demo",
			ShortName: "d",
			Usage:     "Fetch first assignment for each language from exercism.io",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					config, err = configuration.Demo()
					if err != nil {
						fmt.Println(err)
						return
					}
				}
				assignments, err := FetchAssignments(config, FetchEndpoints["demo"])
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "fetch",
			ShortName: "f",
			Usage:     "Fetch current assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := FetchAssignments(config,
					FetchEndpoints["current"])
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}

				fmt.Printf("Exercises written to %s\n", config.ExercismDirectory)
			},
		},
		{
			Name:      "login",
			ShortName: "l",
			Usage:     "Save exercism.io api credentials",
			Action: func(c *cli.Context) {
				configuration.ToFile(configuration.HomeDir(), askForConfigInfo())
			},
		},
		{
			Name:      "logout",
			ShortName: "o",
			Usage:     "Clear exercism.io api credentials",
			Action: func(c *cli.Context) {
				Logout(configuration.HomeDir())
			},
		},
		{
			Name:      "peek",
			ShortName: "p",
			Usage:     "Fetch upcoming assignment from exercism.io",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}
				assignments, err := FetchAssignments(config,
					FetchEndpoints["next"])
				if err != nil {
					fmt.Println(err)
					return
				}

				for _, a := range assignments {
					err := SaveAssignment(config.ExercismDirectory, a)
					if err != nil {
						fmt.Println(err)
					}
				}
			},
		},
		{
			Name:      "submit",
			ShortName: "s",
			Usage:     "Submit code to exercism.io on your current assignment",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				if len(c.Args()) == 0 {
					fmt.Println("Please enter a file name")
					return
				}

				filename := c.Args()[0]

				// Make filename relative to config.ExercismDirectory.
				absPath, err := absolutePath(filename)
				if err != nil {
					fmt.Printf("Couldn't find %v: %v\n", filename, err)
					return
				}
				exDir := config.ExercismDirectory + string(filepath.Separator)
				if !strings.HasPrefix(absPath, exDir) {
					fmt.Printf("%v is not under your exercism project path (%v)\n", absPath, exDir)
					return
				}
				filename = absPath[len(exDir):]

				if IsTest(filename) {
					fmt.Println("It looks like this is a test, please enter an example file name.")
					return
				}

				code, err := ioutil.ReadFile(absPath)
				if err != nil {
					fmt.Printf("Error reading %v: %v\n", absPath, err)
					return
				}

				response, err := SubmitAssignment(config, filename, code)
				if err != nil {
					fmt.Printf("There was an issue with your submission: %v\n", err)
					return
				}

				fmt.Printf("For feedback on your submission visit %s%s%s.\n",
					config.Hostname, "/submissions/", response.Id)

			},
		},
		{
			Name:      "unsubmit",
			ShortName: "u",
			Usage:     "Delete the last submission",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				response, err := UnsubmitAssignment(config)
				if err != nil {
					fmt.Println(err)
					return
				}

				if response != "" {
					return
				}

				fmt.Println("The last submission was successfully deleted.")
			},
		},
		{
			Name:      "whoami",
			ShortName: "w",
			Usage:     "Get the github username that you are logged in as",
			Action: func(c *cli.Context) {
				config, err := configuration.FromFile(configuration.HomeDir())
				if err != nil {
					fmt.Println("Are you sure you are logged in? Please login again.")
					return
				}

				fmt.Println(config.GithubUsername)
			},
		},
	}
	err := app.Run(os.Args)
	if err != nil {
		fmt.Errorf("%v", err)
		os.Exit(1)
	}
}

func askForConfigInfo() (c configuration.Config) {
	var un, key, dir string

	currentDir, err := os.Getwd()
	if err != nil {
		panic(err)
	}

	fmt.Print("Your GitHub username: ")
	_, err = fmt.Scanln(&un)
	if err != nil {
		panic(err)
	}

	fmt.Print("Your exercism.io API key: ")
	_, err = fmt.Scanln(&key)
	if err != nil {
		panic(err)
	}

	fmt.Println("What is your exercism exercises project path?")
	fmt.Printf("Press Enter to select the default (%s):\n", currentDir)
	fmt.Print("> ")
	_, err = fmt.Scanln(&dir)
	if err != nil && err.Error() != "unexpected newline" {
		panic(err)
	}

	dir = configuration.ReplaceTilde(dir)

	err = os.MkdirAll(dir, 0755)
	if err != nil {
		err = fmt.Errorf("Error making directory %v: [%v]", dir, err)
		return
	}

	dir, err = absolutePath(dir)
	if err != nil {
		panic(err)
	}

	if dir == "" {
		dir = currentDir
	}

	return configuration.Config{GithubUsername: un, ApiKey: key, ExercismDirectory: dir, Hostname: "http://exercism.io"}
}

func absolutePath(path string) (string, error) {
	path, err := filepath.Abs(path)
	if err != nil {
		return "", err
	}
	return filepath.EvalSymlinks(path)
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

type config struct {
	awsRegion        string // AWS_REGION
	awsAPIEndpoint   string // AWS_API_ENDPOINT
	s3Bucket         string // AWS_S3_BUCKET
	s3KeyPrefix      string // AWS_S3_KEY_PREFIX
	indexDocument    string // INDEX_DOCUMENT
	directoryListing bool   // DIRECTORY_LISTINGS
	httpCacheControl string // HTTP_CACHE_CONTROL (max-age=86400, no-cache ...)
	httpExpires      string // HTTP_EXPIRES (Thu, 01 Dec 1994 16:00:00 GMT ...)
	basicAuthUser    string // BASIC_AUTH_USER
	basicAuthPass    string // BASIC_AUTH_PASS
	port             string // APP_PORT
	accessLog        bool   // ACCESS_LOG
	sslCert          string // SSL_CERT_PATH
	sslKey           string // SSL_KEY_PATH
	stripPath        string // STRIP_PATH
	contentEncoding  bool   // CONTENT_ENCODING
	corsAllowOrigin  string // CORS_ALLOW_ORIGIN
	corsAllowMethods string // CORS_ALLOW_METHODS
	corsAllowHeaders string // CORS_ALLOW_HEADERS
	corsMaxAge       int64  // CORS_MAX_AGE
}

type symlink struct {
	URL string
}

var (
	version string
	date    string
	c       *config
)

func main() {
	c = configFromEnvironmentVariables()

	http.Handle("/", wrapper(awss3))

	http.HandleFunc("/--version", func(w http.ResponseWriter, r *http.Request) {
		if len(version) > 0 && len(date) > 0 {
			fmt.Fprintf(w, "version: %s (built at %s)", version, date)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	})

	// Listen & Serve
	log.Printf("[service] listening on port %s", c.port)
	if (len(c.sslCert) > 0) && (len(c.sslKey) > 0) {
		log.Fatal(http.ListenAndServeTLS(":"+c.port, c.sslCert, c.sslKey, nil))
	} else {
		log.Fatal(http.ListenAndServe(":"+c.port, nil))
	}
}

func configFromEnvironmentVariables() *config {
	if len(os.Getenv("AWS_ACCESS_KEY_ID")) == 0 {
		log.Print("Not defined environment variable: AWS_ACCESS_KEY_ID")
	}
	if len(os.Getenv("AWS_SECRET_ACCESS_KEY")) == 0 {
		log.Print("Not defined environment variable: AWS_SECRET_ACCESS_KEY")
	}
	if len(os.Getenv("AWS_S3_BUCKET")) == 0 {
		log.Fatal("Missing required environment variable: AWS_S3_BUCKET")
	}
	region := os.Getenv("AWS_REGION")
	if len(region) == 0 {
		region = "us-east-1"
	}
	endpoint := os.Getenv("AWS_API_ENDPOINT")
	if len(endpoint) == 0 {
		endpoint = ""
	}
	port := os.Getenv("APP_PORT")
	if len(port) == 0 {
		port = "80"
	}
	indexDocument := os.Getenv("INDEX_DOCUMENT")
	if len(indexDocument) == 0 {
		indexDocument = "index.html"
	}
	directoryListings := false
	if b, err := strconv.ParseBool(os.Getenv("DIRECTORY_LISTINGS")); err == nil {
		directoryListings = b
	}
	accessLog := false
	if b, err := strconv.ParseBool(os.Getenv("ACCESS_LOG")); err == nil {
		accessLog = b
	}
	contentEncoging := false
	if b, err := strconv.ParseBool(os.Getenv("CONTENT_ENCODING")); err == nil {
		contentEncoging = b
	}
	corsMaxAge := int64(600)
	if i, err := strconv.ParseInt(os.Getenv("CORS_MAX_AGE"), 10, 64); err == nil {
		corsMaxAge = i
	}
	conf := &config{
		awsRegion:        region,
		awsAPIEndpoint:   endpoint,
		s3Bucket:         os.Getenv("AWS_S3_BUCKET"),
		s3KeyPrefix:      os.Getenv("AWS_S3_KEY_PREFIX"),
		indexDocument:    indexDocument,
		directoryListing: directoryListings,
		httpCacheControl: os.Getenv("HTTP_CACHE_CONTROL"),
		httpExpires:      os.Getenv("HTTP_EXPIRES"),
		basicAuthUser:    os.Getenv("BASIC_AUTH_USER"),
		basicAuthPass:    os.Getenv("BASIC_AUTH_PASS"),
		port:             port,
		accessLog:        accessLog,
		sslCert:          os.Getenv("SSL_CERT_PATH"),
		sslKey:           os.Getenv("SSL_KEY_PATH"),
		stripPath:        os.Getenv("STRIP_PATH"),
		contentEncoding:  contentEncoging,
		corsAllowOrigin:  os.Getenv("CORS_ALLOW_ORIGIN"),
		corsAllowMethods: os.Getenv("CORS_ALLOW_METHODS"),
		corsAllowHeaders: os.Getenv("CORS_ALLOW_HEADERS"),
		corsMaxAge:       corsMaxAge,
	}
	// Proxy
	log.Printf("[config] Proxy to %v", conf.s3Bucket)
	log.Printf("[config] AWS Region: %v", conf.awsRegion)

	// TLS pem files
	if (len(conf.sslCert) > 0) && (len(conf.sslKey) > 0) {
		log.Print("[config] TLS enabled.")
	}
	// Basic authentication
	if (len(conf.basicAuthUser) > 0) && (len(conf.basicAuthPass) > 0) {
		log.Printf("[config] Basic authentication: %s", conf.basicAuthUser)
	}
	// CORS
	if (len(conf.corsAllowOrigin) > 0) && (conf.corsMaxAge > 0) {
		log.Printf("[config] CORS enabled: %s", conf.corsAllowOrigin)
	}
	return conf
}

type custom struct {
	io.Writer
	http.ResponseWriter
	status int
}

func (r *custom) Write(b []byte) (int, error) {
	if r.Header().Get("Content-Type") == "" {
		r.Header().Set("Content-Type", http.DetectContentType(b))
	}
	return r.Writer.Write(b)
}

func (r *custom) WriteHeader(status int) {
	r.ResponseWriter.WriteHeader(status)
	r.status = status
}

func wrapper(f func(w http.ResponseWriter, r *http.Request)) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if (len(c.corsAllowOrigin) > 0) && (len(c.corsAllowMethods) > 0) && (len(c.corsAllowHeaders) > 0) && (c.corsMaxAge > 0) {
			w.Header().Set("Access-Control-Allow-Origin", c.corsAllowOrigin)
			w.Header().Set("Access-Control-Allow-Methods", c.corsAllowMethods)
			w.Header().Set("Access-Control-Allow-Headers", c.corsAllowHeaders)
			w.Header().Set("Access-Control-Max-Age", strconv.FormatInt(c.corsMaxAge, 10))
		}
		if (len(c.basicAuthUser) > 0) && (len(c.basicAuthPass) > 0) && !auth(r) {
			w.Header().Set("WWW-Authenticate", `Basic realm="REALM"`)
			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		proc := time.Now()
		addr := r.RemoteAddr
		if ip, found := header(r, "X-Forwarded-For"); found {
			addr = ip
		}
		ioWriter := w.(io.Writer)
		if encodings, found := header(r, "Accept-Encoding"); found && c.contentEncoding {
			for _, encoding := range splitCsvLine(encodings) {
				if encoding == "gzip" {
					w.Header().Set("Content-Encoding", "gzip")
					g := gzip.NewWriter(w)
					defer g.Close()
					ioWriter = g
					break
				}
				if encoding == "deflate" {
					w.Header().Set("Content-Encoding", "deflate")
					z := zlib.NewWriter(w)
					defer z.Close()
					ioWriter = z
					break
				}
			}
		}
		writer := &custom{Writer: ioWriter, ResponseWriter: w, status: http.StatusOK}
		f(writer, r)

		if c.accessLog {
			log.Printf("[%s] %.3f %d %s %s",
				addr, time.Now().Sub(proc).Seconds(),
				writer.status, r.Method, r.URL)
		}
	})
}

func auth(r *http.Request) bool {
	if username, password, ok := r.BasicAuth(); ok {
		return username == c.basicAuthUser &&
			password == c.basicAuthPass
	}
	return false
}

func header(r *http.Request, key string) (string, bool) {
	if r.Header == nil {
		return "", false
	}
	if candidate := r.Header[key]; len(candidate) > 0 {
		return candidate[0], true
	}
	return "", false
}

func splitCsvLine(data string) []string {
	splitted := strings.SplitN(data, ",", -1)
	parsed := make([]string, len(splitted))
	for i, val := range splitted {
		parsed[i] = strings.TrimSpace(val)
	}
	return parsed
}

func awss3(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path

	if len(c.stripPath) > 0 {
		path = strings.Replace(path, c.stripPath, "", 1)
	}
	idx := strings.Index(path, "symlink.json")
	if idx > -1 {
		result, err := s3get(c.s3Bucket, c.s3KeyPrefix+path[:idx+12])
		if err != nil {
			code, message := awsError(err)
			http.Error(w, message, code)
			return
		}
		var link symlink
		buf := new(bytes.Buffer)
		buf.ReadFrom(result.Body)
		err = json.Unmarshal(buf.Bytes(), &link)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		path = link.URL + path[idx+12:]
	}
	if strings.HasSuffix(path, "/") {
		if c.directoryListing {
			s3listFiles(w, r, c.s3Bucket, c.s3KeyPrefix+path)
			return
		}
		path += c.indexDocument
	}
	obj, err := s3get(c.s3Bucket, c.s3KeyPrefix+path)
	if err != nil {
		code, message := awsError(err)
		http.Error(w, message, code)
		return
	}
	if len(c.httpCacheControl) > 0 {
		setStrHeader(w, "Cache-Control", &c.httpCacheControl)
	} else {
		setStrHeader(w, "Cache-Control", obj.CacheControl)
	}
	if len(c.httpExpires) > 0 {
		setStrHeader(w, "Expires", &c.httpExpires)
	} else {
		setStrHeader(w, "Expires", obj.Expires)
	}
	setStrHeader(w, "Content-Disposition", obj.ContentDisposition)
	setStrHeader(w, "Content-Encoding", obj.ContentEncoding)
	setStrHeader(w, "Content-Language", obj.ContentLanguage)
	setIntHeader(w, "Content-Length", obj.ContentLength)
	setStrHeader(w, "Content-Range", obj.ContentRange)
	setStrHeader(w, "Content-Type", obj.ContentType)
	setTimeHeader(w, "Last-Modified", obj.LastModified)

	io.Copy(w, obj.Body)
}

func s3get(backet, key string) (*s3.GetObjectOutput, error) {
	req := &s3.GetObjectInput{
		Bucket: aws.String(backet),
		Key:    aws.String(key),
	}
	return s3.New(awsSession()).GetObject(req)
}

func s3listFiles(w http.ResponseWriter, r *http.Request, backet, key string) {
	if strings.HasSuffix(key, "/") {
		key = key[:len(key)-1]
	}
	req := &s3.ListObjectsInput{
		Bucket: aws.String(backet),
		Prefix: aws.String(key),
	}
	result, err := s3.New(awsSession()).ListObjects(req)
	if err != nil {
		code, message := awsError(err)
		http.Error(w, message, code)
		return
	}
	files := []string{}
	for _, obj := range result.Contents {
		files = append(files, strings.Replace(aws.StringValue(obj.Key), key, "", -1))
	}
	bytes, merr := json.Marshal(files)
	if merr != nil {
		http.Error(w, merr.Error(), http.StatusInternalServerError)
		return
	}
	http.Error(w, string(bytes), http.StatusOK)
}

func awsSession() *session.Session {
	config := &aws.Config{
		Region: aws.String(c.awsRegion),
	}
	if len(c.awsAPIEndpoint) > 0 {
		config.Endpoint = aws.String(c.awsAPIEndpoint)
		config.S3ForcePathStyle = aws.Bool(true)
	}
	return session.Must(session.NewSession(config))
}

func setStrHeader(w http.ResponseWriter, key string, value *string) {
	if value != nil && len(*value) > 0 {
		w.Header().Add(key, *value)
	}
}

func setIntHeader(w http.ResponseWriter, key string, value *int64) {
	if value != nil && *value > 0 {
		w.Header().Add(key, strconv.FormatInt(*value, 10))
	}
}

func setTimeHeader(w http.ResponseWriter, key string, value *time.Time) {
	if value != nil && !reflect.DeepEqual(*value, time.Time{}) {
		w.Header().Add(key, value.UTC().Format(http.TimeFormat))
	}
}

func awsError(err error) (int, string) {
	if aerr, ok := err.(awserr.Error); ok {
		switch aerr.Code() {
		case s3.ErrCodeNoSuchBucket, s3.ErrCodeNoSuchKey:
			return http.StatusNotFound, aerr.Error()
		}
		return http.StatusInternalServerError, aerr.Error()
	}
	return http.StatusInternalServerError, err.Error()
}
<commit_msg>list files just only those prefix matched the specified key<commit_after>package main

import (
	"bytes"
	"compress/gzip"
	"compress/zlib"
	"encoding/json"
	"fmt"
	"io"
	"log"
	"net/http"
	"os"
	"reflect"
	"strconv"
	"strings"
	"time"

	"github.com/aws/aws-sdk-go/aws"
	"github.com/aws/aws-sdk-go/aws/awserr"
	"github.com/aws/aws-sdk-go/aws/session"
	"github.com/aws/aws-sdk-go/service/s3"
)

type config struct {
	awsRegion        string // AWS_REGION
	awsAPIEndpoint   string // AWS_API_ENDPOINT
	s3Bucket         string // AWS_S3_BUCKET
	s3KeyPrefix      string // AWS_S3_KEY_PREFIX
	indexDocument    string // INDEX_DOCUMENT
	directoryListing bool   // DIRECTORY_LISTINGS
	httpCacheControl string // HTTP_CACHE_CONTROL (max-age=86400, no-cache ...)
	httpExpires      string // HTTP_EXPIRES (Thu, 01 Dec 1994 16:00:00 GMT ...)
	basicAuthUser    string // BASIC_AUTH_USER
	basicAuthPass    string // BASIC_AUTH_PASS
	port             string // APP_PORT
	accessLog        bool   // ACCESS_LOG
	sslCert          string // SSL_CERT_PATH
	sslKey           string // SSL_KEY_PATH
	stripPath        string // STRIP_PATH
	contentEncoding  bool   // CONTENT_ENCODING
	corsAllowOrigin  string // CORS_ALLOW_ORIGIN
	corsAllowMethods string // CORS_ALLOW_METHODS
	corsAllowHeaders string // CORS_ALLOW_HEADERS
	corsMaxAge       int64  // CORS_MAX_AGE
}

type symlink struct {
	URL string
}

var (
	version string
	date    string
	c       *config
)

func main() {
	c = configFromEnvironmentVariables()

	http.Handle("/", wrapper(awss3))

	http.HandleFunc("/--version", func(w http.ResponseWriter, r *http.Request) {
		if len(version) > 0 && len(date) > 0 {
			fmt.Fprintf(w, "version: %s (built at %s)", version, date)
		} else {
			w.WriteHeader(http.StatusOK)
		}
	})

	// Listen & Serve
	log.Printf("[service] listening on port %s", c.port)
	if (len(c.sslCert) > 0) && (len(c.sslKey) > 0) {
		log.Fatal(http.ListenAndServeTLS(":"+c.port, c.sslCert, c.sslKey, nil))
	} else {
		log.Fatal(http.ListenAndServe(":"+c.port, nil))
	}
}

func configFromEnvironmentVariables() *config {
	if len(os.Getenv("AWS_ACCESS_KEY_ID")) == 0 {
		log.Print("Not defined environment variable: AWS_ACCESS_KEY_ID")
	}
	if len(os.Getenv("AWS_SECRET_ACCESS_KEY")) == 0 {
		log.Print("Not defined environment variable: AWS_SECRET_ACCESS_KEY")
	}
	if len(os.Getenv("AWS_S3_BUCKET")) == 0 {
		log.Fatal("Missing required environment variable: AWS_S3_BUCKET")
	}
	region := os.Getenv("AWS_REGION")
	if len(region) == 0 {
		region = "us-east-1"
	}
	endpoint := os.Getenv("AWS_API_ENDPOINT")
	if len(endpoint) == 0 {
		endpoint = ""
	}
	port := os.Getenv("APP_PORT")
	if len(port) == 0 {
		port = "80"
	}
	indexDocument := os.Getenv("INDEX_DOCUMENT")
	if len(indexDocument) == 0 {
		indexDocument = "index.html"
	}
	directoryListings := false
	if b, err := strconv.ParseBool(os.Getenv("DIRECTORY_LISTINGS")); err == nil {
		directoryListings = b
	}
	accessLog := false
	if b, err := strconv.ParseBool(os.Getenv("ACCESS_LOG")); err == nil {
		accessLog = b
	}
	contentEncoging := false
	if b, err := strconv.ParseBool(os.Getenv("CONTENT_ENCODING")); err == nil {
		contentEncoging = b
	}
	corsMaxAge := int64(600)
	if i, err := strconv.ParseInt(os.Getenv("CORS_MAX_AGE"), 10, 64); err == nil {
		corsMaxAge = i
	}
	conf := &config{
		awsRegion:        region,
		awsAPIEndpoint:   endpoint,
		s3Bucket:         os.Getenv("AWS_S3_BUCKET"),
		s3KeyPrefix:      os.Getenv("AWS_S3_KEY_PREFIX"),
		indexDocument:    indexDocument,
		directoryListing: directoryListings,
		httpCacheControl: os.Getenv("HTTP_CACHE_CONTROL"),
		httpExpires:      os.Getenv("HTTP_EXPIRES"),
		basicAuthUser:    os.Getenv("BASIC_AUTH_USER"),
		basicAuthPass:    os.Getenv("BASIC_AUTH_PASS"),
		port:             port,
		accessLog:        accessLog,
		sslCert:          os.Getenv("SSL_CERT_PATH"),
		sslKey:           os.Getenv("SSL_KEY_PATH"),
		stripPath:        os.Getenv("STRIP_PATH"),
		contentEncoding:  contentEncoging,
		corsAllowOrigin:  os.Getenv("CORS_ALLOW_ORIGIN"),
		corsAllowMethods: os.Getenv("CORS_ALLOW_METHODS"),
		corsAllowHeaders: os.Getenv("CORS_ALLOW_HEADERS"),
		corsMaxAge:       corsMaxAge,
	}
	// Proxy
	log.Printf("[config] Proxy to %v", conf.s3Bucket)
	log.Printf("[config] AWS Region: %v", conf.awsRegion)

	// TLS pem files
	if (len(conf.sslCert) > 0) && (len(conf.sslKey) > 0) {
		log.Print("[config] TLS enabled.")
	}
	// Basic authentication
	if (len(conf.basicAuthUser) > 0) && (len(conf.basicAuthPass) > 0) {
		log.Printf("[config] Basic authentication: %s", conf.basicAuthUser)
	}
	// CORS
	if (len(conf.corsAllowOrigin) > 0) && (conf.corsMaxAge > 0) {
		log.Printf("[config] CORS enabled: %s", conf.corsAllowOrigin)
	}
	return conf
}

type custom struct {
	io.Writer
	http.ResponseWriter
	status int
}

func (r *custom) Write(b []byte) (int, error) {
	if r.Header().Get("Content-Type") == "" {
		r.Header().Set("Content-Type", http.DetectContentType(b))
	}
	return r.Writer.Write(b)
}

func (r *custom) WriteHeader(status int) {
	r.ResponseWriter.WriteHeader(status)
	r.status = status
}

func wrapper(f func(w http.ResponseWriter, r *http.Request)) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
		if (len(c.corsAllowOrigin) > 0) && (len(c.corsAllowMethods) > 0) && (len(c.corsAllowHeaders) > 0) && (c.corsMaxAge > 0) {
			w.Header().Set("Access-Control-Allow-Origin", c.corsAllowOrigin)
			w.Header().Set("Access-Control-Allow-Methods", c.corsAllowMethods)
			w.Header().Set("Access-Control-Allow-Headers", c.corsAllowHeaders)
			w.Header().Set("Access-Control-Max-Age", strconv.FormatInt(c.corsMaxAge, 10))
		}
		if (len(c.basicAuthUser) > 0) && (len(c.basicAuthPass) > 0) && !auth(r) {
			w.Header().Set("WWW-Authenticate", `Basic realm="REALM"`)
			http.Error(w, http.StatusText(http.StatusUnauthorized), http.StatusUnauthorized)
			return
		}
		proc := time.Now()
		addr := r.RemoteAddr
		if ip, found := header(r, "X-Forwarded-For"); found {
			addr = ip
		}
		ioWriter := w.(io.Writer)
		if encodings, found := header(r, "Accept-Encoding"); found && c.contentEncoding {
			for _, encoding := range splitCsvLine(encodings) {
				if encoding == "gzip" {
					w.Header().Set("Content-Encoding", "gzip")
					g := gzip.NewWriter(w)
					defer g.Close()
					ioWriter = g
					break
				}
				if encoding == "deflate" {
					w.Header().Set("Content-Encoding", "deflate")
					z := zlib.NewWriter(w)
					defer z.Close()
					ioWriter = z
					break
				}
			}
		}
		writer := &custom{Writer: ioWriter, ResponseWriter: w, status: http.StatusOK}
		f(writer, r)

		if c.accessLog {
			log.Printf("[%s] %.3f %d %s %s",
				addr, time.Now().Sub(proc).Seconds(),
				writer.status, r.Method, r.URL)
		}
	})
}

func auth(r *http.Request) bool {
	if username, password, ok := r.BasicAuth(); ok {
		return username == c.basicAuthUser &&
			password == c.basicAuthPass
	}
	return false
}

func header(r *http.Request, key string) (string, bool) {
	if r.Header == nil {
		return "", false
	}
	if candidate := r.Header[key]; len(candidate) > 0 {
		return candidate[0], true
	}
	return "", false
}

func splitCsvLine(data string) []string {
	splitted := strings.SplitN(data, ",", -1)
	parsed := make([]string, len(splitted))
	for i, val := range splitted {
		parsed[i] = strings.TrimSpace(val)
	}
	return parsed
}

func awss3(w http.ResponseWriter, r *http.Request) {
	path := r.URL.Path

	if len(c.stripPath) > 0 {
		path = strings.Replace(path, c.stripPath, "", 1)
	}
	idx := strings.Index(path, "symlink.json")
	if idx > -1 {
		result, err := s3get(c.s3Bucket, c.s3KeyPrefix+path[:idx+12])
		if err != nil {
			code, message := awsError(err)
			http.Error(w, message, code)
			return
		}
		var link symlink
		buf := new(bytes.Buffer)
		buf.ReadFrom(result.Body)
		err = json.Unmarshal(buf.Bytes(), &link)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		path = link.URL + path[idx+12:]
	}
	if strings.HasSuffix(path, "/") {
		if c.directoryListing {
			s3listFiles(w, r, c.s3Bucket, c.s3KeyPrefix+path)
			return
		}
		path += c.indexDocument
	}
	obj, err := s3get(c.s3Bucket, c.s3KeyPrefix+path)
	if err != nil {
		code, message := awsError(err)
		http.Error(w, message, code)
		return
	}
	if len(c.httpCacheControl) > 0 {
		setStrHeader(w, "Cache-Control", &c.httpCacheControl)
	} else {
		setStrHeader(w, "Cache-Control", obj.CacheControl)
	}
	if len(c.httpExpires) > 0 {
		setStrHeader(w, "Expires", &c.httpExpires)
	} else {
		setStrHeader(w, "Expires", obj.Expires)
	}
	setStrHeader(w, "Content-Disposition", obj.ContentDisposition)
	setStrHeader(w, "Content-Encoding", obj.ContentEncoding)
	setStrHeader(w, "Content-Language", obj.ContentLanguage)
	setIntHeader(w, "Content-Length", obj.ContentLength)
	setStrHeader(w, "Content-Range", obj.ContentRange)
	setStrHeader(w, "Content-Type", obj.ContentType)
	setTimeHeader(w, "Last-Modified", obj.LastModified)

	io.Copy(w, obj.Body)
}

func s3get(backet, key string) (*s3.GetObjectOutput, error) {
	req := &s3.GetObjectInput{
		Bucket: aws.String(backet),
		Key:    aws.String(key),
	}
	return s3.New(awsSession()).GetObject(req)
}

func s3listFiles(w http.ResponseWriter, r *http.Request, backet, key string) {
	if strings.HasPrefix(key, "/") {
		key = key[1:]
	}
	req := &s3.ListObjectsInput{
		Bucket: aws.String(backet),
		Prefix: aws.String(key),
	}
	result, err := s3.New(awsSession()).ListObjects(req)
	if err != nil {
		code, message := awsError(err)
		http.Error(w, message, code)
		return
	}
	files := []string{}
	for _, obj := range result.Contents {
		candidate := strings.Replace(aws.StringValue(obj.Key), key, "", -1)
		if len(candidate) == 0 {
			continue
		}
		if strings.Contains(candidate, "/") {
			continue
		}
		files = append(files, candidate)
	}
	bytes, merr := json.Marshal(files)
	if merr != nil {
		http.Error(w, merr.Error(), http.StatusInternalServerError)
		return
	}
	http.Error(w, string(bytes), http.StatusOK)
}

func awsSession() *session.Session {
	config := &aws.Config{
		Region: aws.String(c.awsRegion),
	}
	if len(c.awsAPIEndpoint) > 0 {
		config.Endpoint = aws.String(c.awsAPIEndpoint)
		config.S3ForcePathStyle = aws.Bool(true)
	}
	return session.Must(session.NewSession(config))
}

func setStrHeader(w http.ResponseWriter, key string, value *string) {
	if value != nil && len(*value) > 0 {
		w.Header().Add(key, *value)
	}
}

func setIntHeader(w http.ResponseWriter, key string, value *int64) {
	if value != nil && *value > 0 {
		w.Header().Add(key, strconv.FormatInt(*value, 10))
	}
}

func setTimeHeader(w http.ResponseWriter, key string, value *time.Time) {
	if value != nil && !reflect.DeepEqual(*value, time.Time{}) {
		w.Header().Add(key, value.UTC().Format(http.TimeFormat))
	}
}

func awsError(err error) (int, string) {
	if aerr, ok := err.(awserr.Error); ok {
		switch aerr.Code() {
		case s3.ErrCodeNoSuchBucket, s3.ErrCodeNoSuchKey:
			return http.StatusNotFound, aerr.Error()
		}
		return http.StatusInternalServerError, aerr.Error()
	}
	return http.StatusInternalServerError, err.Error()
}
<|endoftext|>


<commit_before>package users

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestCreateStream(t *testing.T) {
	for _, testdb := range testdatabases {
		_, dev, stream, err := CreateUDS(testdb)
		require.Nil(t, err)

		err = testdb.CreateStream(stream.Name, "{}", dev.DeviceId)
		assert.NotNil(t, err, "Created stream with duplicate name")
	}
}

func TestUpdateStream(t *testing.T) {

	for _, testdb := range testdatabases {
		_, _, stream, err := CreateUDS(testdb)
		require.Nil(t, err)

		stream.Nickname = "true"
		stream.Type = "{a:'string'}"

		err = testdb.UpdateStream(stream)
		assert.Nil(t, err, "Could not update stream %v", err)

		stream2, err := testdb.ReadStreamById(stream.StreamId)
		require.Nil(t, err, "got an error when trying to get a stream that should exist %v", err)

		if !reflect.DeepEqual(stream, stream2) {
			t.Errorf("The original and updated objects don't match orig: %v updated %v", stream, stream2)
		}

		err = testdb.UpdateStream(nil)
		assert.Equal(t, err, InvalidPointerError, "Function safeguards failed")
	}
}

func TestDeleteStream(t *testing.T) {

	for _, testdb := range testdatabases {
		_, _, stream, err := CreateUDS(testdb)
		require.Nil(t, err)

		err = testdb.DeleteStream(stream.StreamId)
		require.Nil(t, err, "Error when attempted delete %v", err)

		_, err = testdb.ReadStreamById(stream.StreamId)
		require.NotNil(t, err, "The stream with the selected ID should have errored out, but it was not")
	}
}

func TestReadStreamByDevice(t *testing.T) {

	for _, testdb := range testdatabases {
		_, dev, _, err := CreateUDS(testdb)
		require.Nil(t, err)

		testdb.CreateStream("TestReadStreamByDevice2", "{}", dev.DeviceId)

		streams, err := testdb.ReadStreamsByDevice(dev.DeviceId)
		require.Nil(t, err)
		require.Len(t, streams, 2, "didn't get enough streams")
	}
}

func TestReadStreamsByUser(t *testing.T) {
	for _, testdb := range testdatabases {
		inserted := map[Stream]bool{}

		user, _, stream, err := CreateUDS(testdb)
		require.Nil(t, err)
		require.NotNil(t, user)
		require.NotNil(t, stream)

		fmt.Printf("User Id: %v\n", user.UserId)

		inserted[*stream] = true

		// create a bunch of devices
		for i := 0; i < 10; i++ {
			device, err := CreateTestDevice(testdb, user)

			require.Nil(t, err)

			fmt.Printf("Device Id: %v\n", device.DeviceId)

			// create a bunch of streams
			for j := 0; j < 10; j++ {
				stream, err := CreateTestStream(testdb, device)
				require.Nil(t, err)
				inserted[*stream] = true
			}
		}

		// Test selecting them
		streams, err := testdb.ReadStreamsByUser(user.UserId)
		require.Nil(t, err, "Retrieved streams was nil")

		// We need to add in the other missing log stream.
		require.Equal(t, len(inserted), len(streams)+1, "Wrong number of streams returned")
	}
}
<commit_msg>Added to wrong side<commit_after>package users

import (
	"fmt"
	"reflect"
	"testing"

	"github.com/stretchr/testify/assert"
	"github.com/stretchr/testify/require"
)

func TestCreateStream(t *testing.T) {
	for _, testdb := range testdatabases {
		_, dev, stream, err := CreateUDS(testdb)
		require.Nil(t, err)

		err = testdb.CreateStream(stream.Name, "{}", dev.DeviceId)
		assert.NotNil(t, err, "Created stream with duplicate name")
	}
}

func TestUpdateStream(t *testing.T) {

	for _, testdb := range testdatabases {
		_, _, stream, err := CreateUDS(testdb)
		require.Nil(t, err)

		stream.Nickname = "true"
		stream.Type = "{a:'string'}"

		err = testdb.UpdateStream(stream)
		assert.Nil(t, err, "Could not update stream %v", err)

		stream2, err := testdb.ReadStreamById(stream.StreamId)
		require.Nil(t, err, "got an error when trying to get a stream that should exist %v", err)

		if !reflect.DeepEqual(stream, stream2) {
			t.Errorf("The original and updated objects don't match orig: %v updated %v", stream, stream2)
		}

		err = testdb.UpdateStream(nil)
		assert.Equal(t, err, InvalidPointerError, "Function safeguards failed")
	}
}

func TestDeleteStream(t *testing.T) {

	for _, testdb := range testdatabases {
		_, _, stream, err := CreateUDS(testdb)
		require.Nil(t, err)

		err = testdb.DeleteStream(stream.StreamId)
		require.Nil(t, err, "Error when attempted delete %v", err)

		_, err = testdb.ReadStreamById(stream.StreamId)
		require.NotNil(t, err, "The stream with the selected ID should have errored out, but it was not")
	}
}

func TestReadStreamByDevice(t *testing.T) {

	for _, testdb := range testdatabases {
		_, dev, _, err := CreateUDS(testdb)
		require.Nil(t, err)

		testdb.CreateStream("TestReadStreamByDevice2", "{}", dev.DeviceId)

		streams, err := testdb.ReadStreamsByDevice(dev.DeviceId)
		require.Nil(t, err)
		require.Len(t, streams, 2, "didn't get enough streams")
	}
}

func TestReadStreamsByUser(t *testing.T) {
	for _, testdb := range testdatabases {
		inserted := map[Stream]bool{}

		user, _, stream, err := CreateUDS(testdb)
		require.Nil(t, err)
		require.NotNil(t, user)
		require.NotNil(t, stream)

		fmt.Printf("User Id: %v\n", user.UserId)

		inserted[*stream] = true

		// create a bunch of devices
		for i := 0; i < 10; i++ {
			device, err := CreateTestDevice(testdb, user)

			require.Nil(t, err)

			fmt.Printf("Device Id: %v\n", device.DeviceId)

			// create a bunch of streams
			for j := 0; j < 10; j++ {
				stream, err := CreateTestStream(testdb, device)
				require.Nil(t, err)
				inserted[*stream] = true
			}
		}

		// Test selecting them
		streams, err := testdb.ReadStreamsByUser(user.UserId)
		require.Nil(t, err, "Retrieved streams was nil")

		// We need to add in the other missing log stream.
		require.Equal(t, len(inserted)+1, len(streams), "Wrong number of streams returned")
	}
}
<|endoftext|>


<commit_before>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/influxdb/influxdb/client"
	"github.com/rcrowley/go-metrics"
)

type configuration struct {
	host            *string
	port            *int
	db              *string
	measurement     *string
	retentionPolicy *string
	batchSize       *int
	rate            *int
	cpus            *int
	duration        *int
}

func main() {

	config := configuration{
		host:            flag.String("h", "localhost", "host"),
		port:            flag.Int("p", 8086, "port number"),
		db:              flag.String("db", "load_test", "database"),
		measurement:     flag.String("m", "load_test", "measurement"),
		retentionPolicy: flag.String("rp", "default", "retention policy"),
		batchSize:       flag.Int("batchSize", 5000, "batch size for requests"),
		rate:            flag.Int("rate", 5, "requests per second"),
		cpus:            flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use"),
		duration:        flag.Int("duration", 5, "time in seconds for test to run"),
	}

	flag.Parse()

	runtime.GOMAXPROCS(*config.cpus)

	u, _ := url.Parse(fmt.Sprintf("http://%s:%d", *config.host, *config.port))
	con, err := client.NewClient(client.Config{URL: *u})

	if err != nil {
		panic(err)
	}

	// Log metrics every second
	go metrics.Log(metrics.DefaultRegistry, time.Second, log.New(os.Stderr, "metrics: ", log.Lmicroseconds))

	startLoadTest(con, config)
}

func startLoadTest(con *client.Client, config configuration) {
	durationCounter := 0

	t := metrics.NewTimer()
	metrics.Register("bang", t)

	for _ = range time.Tick(time.Second) {
		if durationCounter >= *config.duration {
			return
		}

		for i := 0; i < *config.rate; i++ {
			go func() {
				t.Time(func() {
					writePoints(con, config)
				})
			}()
		}
		durationCounter++
	}
}

func writePoints(con *client.Client, config configuration) {
	var (
		hosts     = []string{"host1", "host2", "host3", "host4", "host5", "host6"}
		metrics   = []string{"com.addthis.client.watch.Service.total._red_pjson__.1MinuteRate", "com.addthis.client.watch.Service.total._red_lojson_100eng.json.1MinuteRate", "com.addthis.client.watch.Service.total._red_lojson_300lo.json.1MinuteRate"}
		batchSize = *config.batchSize
		pts       = make([]client.Point, batchSize)
	)

	rand.Seed(42)
	for i := 0; i < batchSize; i++ {
		pts[i] = client.Point{
			Measurement: *config.measurement,
			Tags: map[string]string{
				"host":   strconv.Itoa(rand.Intn(len(hosts))),
				"metric": strconv.Itoa(rand.Intn(len(metrics))),
			},
			Fields: map[string]interface{}{
				"value": rand.Float64(),
			},
			Precision: "n",
		}
	}

	bps := client.BatchPoints{
		Points:          pts,
		Database:        *config.db,
		RetentionPolicy: *config.retentionPolicy,
	}
	_, err := con.Write(bps)
	if err != nil {
		log.Fatal(err)
	}
}
<commit_msg>updated metric names<commit_after>package main

import (
	"flag"
	"fmt"
	"log"
	"math/rand"
	"net/url"
	"os"
	"runtime"
	"strconv"
	"time"

	"github.com/influxdb/influxdb/client"
	"github.com/rcrowley/go-metrics"
)

type configuration struct {
	host            *string
	port            *int
	db              *string
	measurement     *string
	retentionPolicy *string
	batchSize       *int
	rate            *int
	cpus            *int
	duration        *int
}

func main() {

	config := configuration{
		host:            flag.String("h", "localhost", "host"),
		port:            flag.Int("p", 8086, "port number"),
		db:              flag.String("db", "load_test", "database"),
		measurement:     flag.String("m", "load_test", "measurement"),
		retentionPolicy: flag.String("rp", "default", "retention policy"),
		batchSize:       flag.Int("batchSize", 5000, "batch size for requests"),
		rate:            flag.Int("rate", 5, "requests per second"),
		cpus:            flag.Int("cpus", runtime.NumCPU(), "Number of CPUs to use"),
		duration:        flag.Int("duration", 5, "time in seconds for test to run"),
	}

	flag.Parse()

	runtime.GOMAXPROCS(*config.cpus)

	u, _ := url.Parse(fmt.Sprintf("http://%s:%d", *config.host, *config.port))
	con, err := client.NewClient(client.Config{URL: *u})

	if err != nil {
		panic(err)
	}

	// Log metrics every second
	go metrics.Log(metrics.DefaultRegistry, time.Second, log.New(os.Stderr, "metrics: ", log.Lmicroseconds))

	startLoadTest(con, config)
}

func startLoadTest(con *client.Client, config configuration) {
	durationCounter := 0

	t := metrics.NewTimer()
	metrics.Register("bang", t)

	for _ = range time.Tick(time.Second) {
		if durationCounter >= *config.duration {
			return
		}

		for i := 0; i < *config.rate; i++ {
			go func() {
				t.Time(func() {
					writePoints(con, config)
				})
			}()
		}
		durationCounter++
	}
}

func writePoints(con *client.Client, config configuration) {
	var (
		hosts     = []string{"host1", "host2", "host3", "host4", "host5", "host6"}
		metrics   = []string{"com.addthis.Service.total._red_pjson__.1MinuteRate", "com.addthis.Service.total._red_lojson_100eng.json.1MinuteRate", "com.addthis.Service.total._red_lojson_300lo.json.1MinuteRate"}
		batchSize = *config.batchSize
		pts       = make([]client.Point, batchSize)
	)

	rand.Seed(42)
	for i := 0; i < batchSize; i++ {
		pts[i] = client.Point{
			Measurement: *config.measurement,
			Tags: map[string]string{
				"host":   strconv.Itoa(rand.Intn(len(hosts))),
				"metric": strconv.Itoa(rand.Intn(len(metrics))),
			},
			Fields: map[string]interface{}{
				"value": rand.Float64(),
			},
			Precision: "n",
		}
	}

	bps := client.BatchPoints{
		Points:          pts,
		Database:        *config.db,
		RetentionPolicy: *config.retentionPolicy,
	}
	_, err := con.Write(bps)
	if err != nil {
		log.Fatal(err)
	}
}
<|endoftext|>