<commit_before>package main

import (
	"bytes"
	"fmt"
	"github.com/garyburd/redigo/redis"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"text/template"
	"time"
)

// A struct to hold the list of VMs
type VMList struct {
	mux sync.Mutex
	Vms map[int]VMInformation
}

type VMInformation struct {
	Status    string
	Id        int
	OpenPorts map[string]string
}

var configLock sync.Mutex

func main() {
	os.Exit(run())
}

func run() int {
	// Create our datastructures
	v := sync.Mutex{}
	configLock = sync.Mutex{}

	// Connect here rather than in the handler so that we can ensure we can connect and exit if need be
	{
		// Scope our variable to ensure no one accidently uses the connection
		redisCon, err := redis.Dial("tcp", "10.0.5.20:6379")
		if err != nil {
			fmt.Printf("Could not connect to redis database: %v", err)
			return 1
		}
		go redisPubSubHandle(redisCon)
		defer redisCon.Close()

	}

	http.HandleFunc("/create/", func(w http.ResponseWriter, r *http.Request) {
		createHandler(w, r, v)
	})

	http.HandleFunc("/view/", func(w http.ResponseWriter, r *http.Request) {
		viewHandler(w, r, v)
	})

	// IMPORTANT OMG
	// Only ever bind to one address!
	http.ListenAndServe("10.0.0.5:80", nil)

	return 0
}

func redisPubSubHandle(redisCon redis.Conn) {
	psc := redis.PubSubConn{Conn: redisCon}
	psc.Subscribe("openport")
	psc.Subscribe("deletevm")

	for {
		// TODO Can we use an if/continue instead of a switch?
		switch v := psc.Receive().(type) {
		case redis.Message:
			switch v.Channel {
			case "openport":
				// Lets tell Tor to regenerate its configuration (which happens without the state of the previous message)
				// TODO: We should just add the port to the config, not regenerate it from scratch, probably
				go rewriteConfig()
			case "deletevm":
				// Parse out the ID and if required, do the deed
				vmId, err := strconv.Atoi(string(v.Data))
				// Check whether the ID is valid
				if err != nil {
					continue
				}
				if vmId < 50 || vmId > 255 {
					return
				}

				go deleteVm(vmId)
			}
		}
		fmt.Println(fmt.Sprintf("[%v] Got a PUBSUB message", time.Now()))
	}
}

func deleteVm(vmId int) {
	if vmId < 50 || vmId > 255 {
		return
	}

	// delete the networking file
	err := os.Remove(fmt.Sprintf("/etc/network/interfaces.d/vlan%v", vmId))
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not delete networking interface (vmId: %v): %v", vmId, err)
	}

	// regenerate configuration
	rewriteConfig()

	// now clean up the extra files in /var/lib/tor
	err = os.Remove(fmt.Sprintf("/var/lib/tor/guest-%v", vmId))
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not delete tor datadir (vmId: %v): %v", vmId, err)
	}
}

func rewriteConfig() {
	configLock.Lock()
	defer configLock.Unlock()

	// To generate the new torrc we need the list of every VM, which we can get by looking in /etc/network/interfaces.d
	vmsfiles, err := filepath.Glob("/etc/network/interfaces.d/vlan*")
	if err != nil {
		fmt.Fprintln(os.Stderr, "error globbing for VMs: %v", err)
		return
	}

	vms := VMList{Vms: make(map[int]VMInformation)}

	// A variable for whether we should bother talking to redis
	doRedis := true
	redisCon, err := redis.Dial("tcp", "10.0.5.20:6379")
	if err != nil {
		fmt.Fprintln(os.Stderr, "error connecting to redis: %v", err)
		doRedis = false
	}

	for _, f := range vmsfiles {
		fid := strings.Split(f, "vlan")
		i, err := strconv.Atoi(fid[1])
		if err == nil {
			// We have a valid ID we should build a VMInformation for
			var vminfo VMInformation
			vminfo = VMInformation{Id: i, Status: "complete"}
			// Check with redis to see if we need to have an open port
			if doRedis {
				vminfo.OpenPorts = make(map[string]string)
				inter, err := redis.Strings(redisCon.Do("SMEMBERS", fmt.Sprintf("vm:%v:hostedports", i)))
				for _, f := range inter {
					vminfo.OpenPorts[f] = f
				}
				if err != nil {
					fmt.Fprintln(os.Stderr, "error talking to redis: %v", err)
				}
				// We can ignore an error since the map is still intialized
				vms.Vms[i] = vminfo
			}
			// Add it to our VMList
		} // else it probably was something we can ignore
	}

	// Configure ip tables
	// TODO rewrite so this doesn't require a reload
	t, err := template.ParseFiles("assets/iptables")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error creating iptables template for new VM: %v", err)
		return
	}

	var iptables bytes.Buffer
	err = t.Execute(&iptables, vms)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error executing iptables template for new VM: %v", err)
		return
	}

	// Write new iptables file
	err = ioutil.WriteFile("/etc/iptables", iptables.Bytes(), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing iptables template for new VM: %v", err)
		return
	}

	// Reload iptables
	err = exec.Command("iptables-restore", "/etc/iptables").Run()
	if err != nil {
		// TODO More graceful handling of this. If iptables is down, HOLY SHIT FIRE, like shutdown -h now
		fmt.Fprintln(os.Stderr, "error executing iptables restore for new VM: %v", err)
		return
	}

	// Generate new torrc
	t, err = template.ParseFiles("assets/torrc")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error creating torrc template for new VM: %v", err)
		return
	}

	var torrc bytes.Buffer
	err = t.Execute(&torrc, vms)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error executing torrc template for new VM: %v", err)
		return
	}

	// Write new torrc
	err = ioutil.WriteFile("/etc/tor/torrc", torrc.Bytes(), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing torrc template for new VM: %v", err)
		return
	}

	// Send SIGHUP to Tor
	// From docs: The signal instructs Tor to reload its configuration (including closing and reopening logs), and kill and restart its helper processes if applicable.
	// TODO Look into using control port to make changes without a reload required
	err = exec.Command("pkill", "-SIGHUP", "tor$").Run()
	if err != nil {
		// TODO More graceful handling of this. If tor is down, HOLY SHIT FIRE
		fmt.Fprintln(os.Stderr, "error executing tor restart for new VM: %v", err)
		return
	}

}

func viewHandler(w http.ResponseWriter, r *http.Request, v sync.Mutex) {
	fmt.Println(fmt.Sprintf("[%v] %v", time.Now(), r.URL.Path))
	// Get the ID of the new VM
	vmIdStr := r.URL.Path[len("/view/"):]
	vmId, err := strconv.Atoi(vmIdStr)
	// Check whether the ID is valid
	if err != nil {
		fmt.Fprintf(w, "invalid")
		return
	}

	if vmId < 50 || vmId > 255 {
		fmt.Fprintf(w, "invalid")
		return
	}

	// There's a chance it doesn't exist yet, but meh, we can't do much
	buf, err := ioutil.ReadFile(fmt.Sprintf("/var/lib/tor/guest-%v/hostname", vmId))
	if err != nil {
		fmt.Fprintf(w, "unknown")
		fmt.Fprintln(os.Stderr, "error fetching hostname for hidden service", err)
		return
	}

	// TODO We should really return success + the hostname (json n shit yo)
	fmt.Fprintf(w, string(buf))
}

func createHandler(w http.ResponseWriter, r *http.Request, v sync.Mutex) {
	fmt.Println(fmt.Sprintf("[%v] %v", time.Now(), r.URL.Path))
	// TODO In the future we will allow more than just sshd port to be a hidden service
	// Lock to be safe (unlock in the actual create goroutine)
	v.Lock()

	// Get the ID of the new VM
	vmIdStr := r.URL.Path[len("/create/"):]
	vmId, err := strconv.Atoi(vmIdStr)
	// Check whether the ID is valid
	if err != nil {
		fmt.Fprintf(w, "invalid")
		return
	}

	// TODO More validation for if it already exists rather than just error out later

	// create in background
	go create(vmId, v)
	fmt.Fprintf(w, "creating")
}

func create(vmId int, v sync.Mutex) {
	// remove the lock when we're done
	defer v.Unlock()

	// Add network configuration
	t, err := template.ParseFiles("assets/networks-vlan")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error creating template for new VM: %v", err)
		return
	}
	var net bytes.Buffer
	err = t.Execute(&net, VMInformation{Id: vmId})
	if err != nil {
		fmt.Fprintf(os.Stderr, "error executing template for new VM: %v", err)
		return
	}
	// Write net file
	err = ioutil.WriteFile(fmt.Sprintf("/etc/network/interfaces.d/vlan%v", vmId), net.Bytes(), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing net template for new VM: %v", err)
		return
	}

	// Instead of restarting network, lets just run the commands to bring up that one vlan
	err = exec.Command("vconfig", "add", "eth0", fmt.Sprintf("%v", vmId)).Run()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error running vconfig", err)
		return
	}
	err = exec.Command("ifconfig", fmt.Sprintf("eth0.%v", vmId), "up").Run()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error running ifconfig up", err)
		return
	}
	err = exec.Command("ifconfig", fmt.Sprintf("eth0.%v", vmId), fmt.Sprintf("10.0.%v.5", vmId), "netmask", "255.255.255.0").Run()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error running ifconfig configuration", err)
		return
	}

	// Now we just create it
	rewriteConfig()
}
<commit_msg>Update torcontrol-daemon to fix fast deletevm<commit_after>package main

import (
	"bytes"
	"fmt"
	"github.com/garyburd/redigo/redis"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"sync"
	"text/template"
	"time"
)

// A struct to hold the list of VMs
type VMList struct {
	mux sync.Mutex
	Vms map[int]VMInformation
}

type VMInformation struct {
	Status    string
	Id        int
	OpenPorts map[string]string
}

var configLock sync.Mutex

func main() {
	os.Exit(run())
}

func run() int {
	// Create our datastructures
	v := sync.Mutex{}
	configLock = sync.Mutex{}

	// Connect here rather than in the handler so that we can ensure we can connect and exit if need be
	{
		// Scope our variable to ensure no one accidently uses the connection
		redisCon, err := redis.Dial("tcp", "10.0.5.20:6379")
		if err != nil {
			fmt.Printf("Could not connect to redis database: %v", err)
			return 1
		}
		go redisPubSubHandle(redisCon)
		defer redisCon.Close()

	}

	http.HandleFunc("/create/", func(w http.ResponseWriter, r *http.Request) {
		createHandler(w, r, v)
	})

	http.HandleFunc("/view/", func(w http.ResponseWriter, r *http.Request) {
		viewHandler(w, r, v)
	})

	// IMPORTANT OMG
	// Only ever bind to one address!
	http.ListenAndServe("10.0.0.5:80", nil)

	return 0
}

func redisPubSubHandle(redisCon redis.Conn) {
	psc := redis.PubSubConn{Conn: redisCon}
	psc.Subscribe("openport")
	psc.Subscribe("deletevm")

	for {
		// TODO Can we use an if/continue instead of a switch?
		switch v := psc.Receive().(type) {
		case redis.Message:
			switch v.Channel {
			case "openport":
				// Lets tell Tor to regenerate its configuration (which happens without the state of the previous message)
				// TODO: We should just add the port to the config, not regenerate it from scratch, probably
				go rewriteConfig()
			case "deletevm":
				// Parse out the ID and if required, do the deed
				vmId, err := strconv.Atoi(string(v.Data))
				// Check whether the ID is valid
				if err != nil {
					continue
				}
				if vmId < 50 || vmId > 255 {
					return
				}

				go deleteVm(vmId)
			}
		}
		fmt.Println(fmt.Sprintf("[%v] Got a PUBSUB message", time.Now()))
	}
}

func deleteVm(vmId int) {
	if vmId < 50 || vmId > 255 {
		return
	}

	// delete the networking file
	err := os.Remove(fmt.Sprintf("/etc/network/interfaces.d/vlan%v", vmId))
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not delete networking interface (vmId: %v): %v", vmId, err)
	}

	// regenerate configuration
	rewriteConfig()

	// now clean up the extra files in /var/lib/tor
	err = os.RemoveAll(fmt.Sprintf("/var/lib/tor/guest-%v", vmId))
	if err != nil {
		fmt.Fprintln(os.Stderr, "could not delete tor datadir (vmId: %v): %v", vmId, err)
	}
}

func rewriteConfig() {
	configLock.Lock()
	defer configLock.Unlock()

	// To generate the new torrc we need the list of every VM, which we can get by looking in /etc/network/interfaces.d
	vmsfiles, err := filepath.Glob("/etc/network/interfaces.d/vlan*")
	if err != nil {
		fmt.Fprintln(os.Stderr, "error globbing for VMs: %v", err)
		return
	}

	vms := VMList{Vms: make(map[int]VMInformation)}

	// A variable for whether we should bother talking to redis
	doRedis := true
	redisCon, err := redis.Dial("tcp", "10.0.5.20:6379")
	if err != nil {
		fmt.Fprintln(os.Stderr, "error connecting to redis: %v", err)
		doRedis = false
	}

	for _, f := range vmsfiles {
		fid := strings.Split(f, "vlan")
		i, err := strconv.Atoi(fid[1])
		if err == nil {
			// We have a valid ID we should build a VMInformation for
			var vminfo VMInformation
			vminfo = VMInformation{Id: i, Status: "complete"}
			// Check with redis to see if we need to have an open port
			if doRedis {
				vminfo.OpenPorts = make(map[string]string)
				inter, err := redis.Strings(redisCon.Do("SMEMBERS", fmt.Sprintf("vm:%v:hostedports", i)))
				for _, f := range inter {
					vminfo.OpenPorts[f] = f
				}
				if err != nil {
					fmt.Fprintln(os.Stderr, "error talking to redis: %v", err)
				}
				// We can ignore an error since the map is still intialized
				vms.Vms[i] = vminfo
			}
			// Add it to our VMList
		} // else it probably was something we can ignore
	}

	// Configure ip tables
	// TODO rewrite so this doesn't require a reload
	t, err := template.ParseFiles("assets/iptables")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error creating iptables template for new VM: %v", err)
		return
	}

	var iptables bytes.Buffer
	err = t.Execute(&iptables, vms)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error executing iptables template for new VM: %v", err)
		return
	}

	// Write new iptables file
	err = ioutil.WriteFile("/etc/iptables", iptables.Bytes(), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing iptables template for new VM: %v", err)
		return
	}

	// Reload iptables
	err = exec.Command("iptables-restore", "/etc/iptables").Run()
	if err != nil {
		// TODO More graceful handling of this. If iptables is down, HOLY SHIT FIRE, like shutdown -h now
		fmt.Fprintln(os.Stderr, "error executing iptables restore for new VM: %v", err)
		return
	}

	// Generate new torrc
	t, err = template.ParseFiles("assets/torrc")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error creating torrc template for new VM: %v", err)
		return
	}

	var torrc bytes.Buffer
	err = t.Execute(&torrc, vms)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error executing torrc template for new VM: %v", err)
		return
	}

	// Write new torrc
	err = ioutil.WriteFile("/etc/tor/torrc", torrc.Bytes(), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing torrc template for new VM: %v", err)
		return
	}

	// Send SIGHUP to Tor
	// From docs: The signal instructs Tor to reload its configuration (including closing and reopening logs), and kill and restart its helper processes if applicable.
	// TODO Look into using control port to make changes without a reload required
	err = exec.Command("pkill", "-SIGHUP", "tor$").Run()
	if err != nil {
		// TODO More graceful handling of this. If tor is down, HOLY SHIT FIRE
		fmt.Fprintln(os.Stderr, "error executing tor restart for new VM: %v", err)
		return
	}

}

func viewHandler(w http.ResponseWriter, r *http.Request, v sync.Mutex) {
	fmt.Println(fmt.Sprintf("[%v] %v", time.Now(), r.URL.Path))
	// Get the ID of the new VM
	vmIdStr := r.URL.Path[len("/view/"):]
	vmId, err := strconv.Atoi(vmIdStr)
	// Check whether the ID is valid
	if err != nil {
		fmt.Fprintf(w, "invalid")
		return
	}

	if vmId < 50 || vmId > 255 {
		fmt.Fprintf(w, "invalid")
		return
	}

	// There's a chance it doesn't exist yet, but meh, we can't do much
	buf, err := ioutil.ReadFile(fmt.Sprintf("/var/lib/tor/guest-%v/hostname", vmId))
	if err != nil {
		fmt.Fprintf(w, "unknown")
		fmt.Fprintln(os.Stderr, "error fetching hostname for hidden service", err)
		return
	}

	// TODO We should really return success + the hostname (json n shit yo)
	fmt.Fprintf(w, string(buf))
}

func createHandler(w http.ResponseWriter, r *http.Request, v sync.Mutex) {
	fmt.Println(fmt.Sprintf("[%v] %v", time.Now(), r.URL.Path))
	// TODO In the future we will allow more than just sshd port to be a hidden service
	// Lock to be safe (unlock in the actual create goroutine)
	v.Lock()

	// Get the ID of the new VM
	vmIdStr := r.URL.Path[len("/create/"):]
	vmId, err := strconv.Atoi(vmIdStr)
	// Check whether the ID is valid
	if err != nil {
		fmt.Fprintf(w, "invalid")
		return
	}

	// TODO More validation for if it already exists rather than just error out later

	// create in background
	go create(vmId, v)
	fmt.Fprintf(w, "creating")
}

func create(vmId int, v sync.Mutex) {
	// remove the lock when we're done
	defer v.Unlock()

	// Add network configuration
	t, err := template.ParseFiles("assets/networks-vlan")
	if err != nil {
		fmt.Fprintf(os.Stderr, "error creating template for new VM: %v", err)
		return
	}
	var net bytes.Buffer
	err = t.Execute(&net, VMInformation{Id: vmId})
	if err != nil {
		fmt.Fprintf(os.Stderr, "error executing template for new VM: %v", err)
		return
	}
	// Write net file
	err = ioutil.WriteFile(fmt.Sprintf("/etc/network/interfaces.d/vlan%v", vmId), net.Bytes(), 0644)
	if err != nil {
		fmt.Fprintf(os.Stderr, "error writing net template for new VM: %v", err)
		return
	}

	// Instead of restarting network, lets just run the commands to bring up that one vlan
	err = exec.Command("vconfig", "add", "eth0", fmt.Sprintf("%v", vmId)).Run()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error running vconfig", err)
		return
	}
	err = exec.Command("ifconfig", fmt.Sprintf("eth0.%v", vmId), "up").Run()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error running ifconfig up", err)
		return
	}
	err = exec.Command("ifconfig", fmt.Sprintf("eth0.%v", vmId), fmt.Sprintf("10.0.%v.5", vmId), "netmask", "255.255.255.0").Run()
	if err != nil {
		fmt.Fprintln(os.Stderr, "error running ifconfig configuration", err)
		return
	}

	// Now we just create it
	rewriteConfig()
}
<|endoftext|>


<commit_before>package conn

import (
	"errors"
	"fmt"
	"regexp"
	"strings"

	"github.com/jordwest/imap-server/mailstore"
	"github.com/jordwest/imap-server/types"
	"github.com/jordwest/imap-server/util"
)

const (
	fetchArgUID    int = 0
	fetchArgRange  int = 1
	fetchArgParams int = 2
)

var registeredFetchParams []fetchParamDefinition
var peekRE *regexp.Regexp

// ErrUnrecognisedParameter indicates that the parameter requested in a FETCH
// command is unrecognised or not implemented in this IMAP server
var ErrUnrecognisedParameter = errors.New("Unrecognised Parameter")

type fetchParamDefinition struct {
	re      *regexp.Regexp
	handler func([]string, *Conn, mailstore.Message, bool) string
}

func cmdFetch(args commandArgs, c *Conn) {
	// Fetch the messages

	seqSet, err := types.InterpretSequenceSet(args.Arg(fetchArgRange))
	if err != nil {
		c.writeResponse(args.ID(), "NO "+err.Error())
		return
	}

	searchByUID := strings.ToUpper(args.Arg(fetchArgUID)) == "UID "

	var msgs []mailstore.Message
	if searchByUID {
		msgs = c.SelectedMailbox.MessageSetByUID(seqSet)
	} else {
		msgs = c.SelectedMailbox.MessageSetBySequenceNumber(seqSet)
	}

	fetchParamString := args.Arg(fetchArgParams)
	if searchByUID && !strings.Contains(fetchParamString, "UID") {
		fetchParamString += " UID"
	}

	for _, msg := range msgs {
		fetchParams, err := fetch(fetchParamString, c, msg)
		if err != nil {
			if err == ErrUnrecognisedParameter {
				c.writeResponse(args.ID(), "BAD Unrecognised Parameter")
				return
			}

			c.writeResponse(args.ID(), "BAD")
			return
		}

		fullReply := fmt.Sprintf("%d FETCH (%s)",
			msg.SequenceNumber(),
			fetchParams)

		c.writeResponse("", fullReply)
	}

	if searchByUID {
		c.writeResponse(args.ID(), "OK UID FETCH Completed")
	} else {
		c.writeResponse(args.ID(), "OK FETCH Completed")
	}
}

// Fetch requested params from a given message
// eg fetch("UID BODY[TEXT] RFC822.SIZE", c, message)
func fetch(params string, c *Conn, m mailstore.Message) (string, error) {
	paramList := util.SplitParams(params)

	// Prepare the list of responses
	responseParams := make([]string, 0, len(paramList))

	for _, param := range paramList {
		paramResponse, err := fetchParam(param, c, m)
		if err != nil {
			return "", err
		}
		responseParams = append(responseParams, paramResponse)
	}
	return strings.Join(responseParams, " "), nil
}

// Match a single fetch parameter and return the data
func fetchParam(param string, c *Conn, m mailstore.Message) (string, error) {
	peek := false
	if peekRE.MatchString(param) {
		peek = true
	}
	// Search through the parameter list until a parameter handler is found
	for _, element := range registeredFetchParams {
		if element.re.MatchString(param) {
			return element.handler(element.re.FindStringSubmatch(param), c, m, peek), nil
		}
	}
	return "", ErrUnrecognisedParameter
}

// Register all supported fetch parameters
func init() {
	peekRE = regexp.MustCompile("\\.PEEK")
	registeredFetchParams = make([]fetchParamDefinition, 0)
	registerFetchParam("UID", fetchUID)
	registerFetchParam("FLAGS", fetchFlags)
	registerFetchParam("RFC822\\.SIZE", fetchRfcSize)
	registerFetchParam("INTERNALDATE", fetchInternalDate)
	registerFetchParam("BODY(?:\\.PEEK)?\\[HEADER\\]", fetchHeaders)
	registerFetchParam("BODY(?:\\.PEEK)?"+
		"\\[HEADER\\.FIELDS \\(([A-z\\s-]+)\\)\\]", fetchHeaderSpecificFields)
	registerFetchParam("BODY(?:\\.PEEK)?\\[TEXT\\]", fetchBody)
	registerFetchParam("BODY(?:\\.PEEK)?\\[\\]", fetchFullText)
}

func registerFetchParam(regex string, handler func([]string, *Conn, mailstore.Message, bool) string) {
	newParam := fetchParamDefinition{
		re:      regexp.MustCompile(regex),
		handler: handler,
	}
	registeredFetchParams = append(registeredFetchParams, newParam)
}

// Fetch the UID of the mail message
func fetchUID(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	return fmt.Sprintf("UID %d", m.UID())
}

func fetchFlags(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	flags := append(m.Flags().Strings(), m.Keywords()...)
	flagList := strings.Join(flags, " ")
	return fmt.Sprintf("FLAGS (%s)", flagList)
}

func fetchRfcSize(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	return fmt.Sprintf("RFC822.SIZE %d", m.Size())
}

func fetchInternalDate(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	dateStr := m.InternalDate().Format(util.InternalDate)
	return fmt.Sprintf("INTERNALDATE \"%s\"", dateStr)
}

func fetchHeaders(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	hdr := fmt.Sprintf("%s\r\n\r\n", m.Header())
	hdrLen := len(hdr)

	peekStr := ""
	if peekOnly {
		peekStr = ".PEEK"
	}

	return fmt.Sprintf("BODY%s[HEADER] {%d}\r\n%s", peekStr, hdrLen, hdr)
}

func fetchHeaderSpecificFields(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	if !peekOnly {
		fmt.Printf("TODO: Peek not requested, mark all as non-recent\n")
	}
	fields := strings.Split(args[1], " ")
	hdrs := m.Header()
	requestedHeaders := make(types.MIMEHeader)
	replyFieldList := make([]string, len(fields))
	for i, key := range fields {
		replyFieldList[i] = "\"" + key + "\""
		// If the key exists in the headers, copy it over
		if k, v, ok := hdrs.FindKey(key); ok {
			requestedHeaders[k] = v
		}
	}
	hdr := fmt.Sprintf("%s\r\n\r\n", requestedHeaders)
	hdrLen := len(hdr)

	return fmt.Sprintf("BODY[HEADER.FIELDS (%s)] {%d}\r\n%s",
		strings.Join(replyFieldList, " "),
		hdrLen,
		hdr)

}

func fetchBody(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	body := fmt.Sprintf("%s\r\n", m.Body())
	bodyLen := len(body)

	return fmt.Sprintf("BODY[TEXT] {%d}\r\n%s",
		bodyLen, body)
}

func fetchFullText(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	mail := fmt.Sprintf("%s\r\n\r\n%s\r\n", m.Header(), m.Body())
	mailLen := len(mail)

	return fmt.Sprintf("BODY[] {%d}\r\n%s",
		mailLen, mail)
}
<commit_msg>Mark messages as non-recent after being fetched. Closes #15<commit_after>package conn

import (
	"errors"
	"fmt"
	"regexp"
	"strings"

	"github.com/jordwest/imap-server/mailstore"
	"github.com/jordwest/imap-server/types"
	"github.com/jordwest/imap-server/util"
)

const (
	fetchArgUID    int = 0
	fetchArgRange  int = 1
	fetchArgParams int = 2
)

var registeredFetchParams []fetchParamDefinition
var peekRE *regexp.Regexp

// ErrUnrecognisedParameter indicates that the parameter requested in a FETCH
// command is unrecognised or not implemented in this IMAP server
var ErrUnrecognisedParameter = errors.New("Unrecognised Parameter")

type fetchParamDefinition struct {
	re      *regexp.Regexp
	handler func([]string, *Conn, mailstore.Message, bool) string
}

func cmdFetch(args commandArgs, c *Conn) {
	// Fetch the messages

	seqSet, err := types.InterpretSequenceSet(args.Arg(fetchArgRange))
	if err != nil {
		c.writeResponse(args.ID(), "NO "+err.Error())
		return
	}

	searchByUID := strings.ToUpper(args.Arg(fetchArgUID)) == "UID "

	var msgs []mailstore.Message
	if searchByUID {
		msgs = c.SelectedMailbox.MessageSetByUID(seqSet)
	} else {
		msgs = c.SelectedMailbox.MessageSetBySequenceNumber(seqSet)
	}

	fetchParamString := args.Arg(fetchArgParams)
	if searchByUID && !strings.Contains(fetchParamString, "UID") {
		fetchParamString += " UID"
	}

	for _, msg := range msgs {
		fetchParams, err := fetch(fetchParamString, c, msg)
		if err != nil {
			if err == ErrUnrecognisedParameter {
				c.writeResponse(args.ID(), "BAD Unrecognised Parameter")
				return
			}

			c.writeResponse(args.ID(), "BAD")
			return
		}

		if c.mailboxWritable == ReadWrite {
			msg = msg.RemoveFlags(types.FlagRecent)
			err = msg.Save()
			if err != nil {
				// TODO: this error is not fatal, but should still be logged
			}
		}

		fullReply := fmt.Sprintf("%d FETCH (%s)",
			msg.SequenceNumber(),
			fetchParams)

		c.writeResponse("", fullReply)
	}

	if searchByUID {
		c.writeResponse(args.ID(), "OK UID FETCH Completed")
	} else {
		c.writeResponse(args.ID(), "OK FETCH Completed")
	}
}

// Fetch requested params from a given message
// eg fetch("UID BODY[TEXT] RFC822.SIZE", c, message)
func fetch(params string, c *Conn, m mailstore.Message) (string, error) {
	paramList := util.SplitParams(params)

	// Prepare the list of responses
	responseParams := make([]string, 0, len(paramList))

	for _, param := range paramList {
		paramResponse, err := fetchParam(param, c, m)
		if err != nil {
			return "", err
		}
		responseParams = append(responseParams, paramResponse)
	}
	return strings.Join(responseParams, " "), nil
}

// Match a single fetch parameter and return the data
func fetchParam(param string, c *Conn, m mailstore.Message) (string, error) {
	peek := false
	if peekRE.MatchString(param) {
		peek = true
	}
	// Search through the parameter list until a parameter handler is found
	for _, element := range registeredFetchParams {
		if element.re.MatchString(param) {
			return element.handler(element.re.FindStringSubmatch(param), c, m, peek), nil
		}
	}
	return "", ErrUnrecognisedParameter
}

// Register all supported fetch parameters
func init() {
	peekRE = regexp.MustCompile("\\.PEEK")
	registeredFetchParams = make([]fetchParamDefinition, 0)
	registerFetchParam("UID", fetchUID)
	registerFetchParam("FLAGS", fetchFlags)
	registerFetchParam("RFC822\\.SIZE", fetchRfcSize)
	registerFetchParam("INTERNALDATE", fetchInternalDate)
	registerFetchParam("BODY(?:\\.PEEK)?\\[HEADER\\]", fetchHeaders)
	registerFetchParam("BODY(?:\\.PEEK)?"+
		"\\[HEADER\\.FIELDS \\(([A-z\\s-]+)\\)\\]", fetchHeaderSpecificFields)
	registerFetchParam("BODY(?:\\.PEEK)?\\[TEXT\\]", fetchBody)
	registerFetchParam("BODY(?:\\.PEEK)?\\[\\]", fetchFullText)
}

func registerFetchParam(regex string, handler func([]string, *Conn, mailstore.Message, bool) string) {
	newParam := fetchParamDefinition{
		re:      regexp.MustCompile(regex),
		handler: handler,
	}
	registeredFetchParams = append(registeredFetchParams, newParam)
}

// Fetch the UID of the mail message
func fetchUID(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	return fmt.Sprintf("UID %d", m.UID())
}

func fetchFlags(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	flags := append(m.Flags().Strings(), m.Keywords()...)
	flagList := strings.Join(flags, " ")
	return fmt.Sprintf("FLAGS (%s)", flagList)
}

func fetchRfcSize(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	return fmt.Sprintf("RFC822.SIZE %d", m.Size())
}

func fetchInternalDate(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	dateStr := m.InternalDate().Format(util.InternalDate)
	return fmt.Sprintf("INTERNALDATE \"%s\"", dateStr)
}

func fetchHeaders(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	hdr := fmt.Sprintf("%s\r\n\r\n", m.Header())
	hdrLen := len(hdr)

	peekStr := ""
	if peekOnly {
		peekStr = ".PEEK"
	}

	return fmt.Sprintf("BODY%s[HEADER] {%d}\r\n%s", peekStr, hdrLen, hdr)
}

func fetchHeaderSpecificFields(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	if !peekOnly {
		fmt.Printf("TODO: Peek not requested, mark all as non-recent\n")
	}
	fields := strings.Split(args[1], " ")
	hdrs := m.Header()
	requestedHeaders := make(types.MIMEHeader)
	replyFieldList := make([]string, len(fields))
	for i, key := range fields {
		replyFieldList[i] = "\"" + key + "\""
		// If the key exists in the headers, copy it over
		if k, v, ok := hdrs.FindKey(key); ok {
			requestedHeaders[k] = v
		}
	}
	hdr := fmt.Sprintf("%s\r\n\r\n", requestedHeaders)
	hdrLen := len(hdr)

	return fmt.Sprintf("BODY[HEADER.FIELDS (%s)] {%d}\r\n%s",
		strings.Join(replyFieldList, " "),
		hdrLen,
		hdr)

}

func fetchBody(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	body := fmt.Sprintf("%s\r\n", m.Body())
	bodyLen := len(body)

	return fmt.Sprintf("BODY[TEXT] {%d}\r\n%s",
		bodyLen, body)
}

func fetchFullText(args []string, c *Conn, m mailstore.Message, peekOnly bool) string {
	mail := fmt.Sprintf("%s\r\n\r\n%s\r\n", m.Header(), m.Body())
	mailLen := len(mail)

	return fmt.Sprintf("BODY[] {%d}\r\n%s",
		mailLen, mail)
}
<|endoftext|>


<commit_before>package file_storage

import (
	"fmt"
	"github.com/photoshelf/photoshelf-storage/domain/model"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

var storage *FileStorage
var testdata []byte

func TestMain(m *testing.M) {
	testdataPath := path.Join(os.Getenv("GOPATH"), "src/github.com/photoshelf/photoshelf-storage", "testdata")
	body, _ := os.Open(path.Join(testdataPath, "e3158990bdee63f8594c260cd51a011d"))
	testdata, _ = ioutil.ReadAll(body)

	dataPath := path.Join(os.TempDir(), "file")
	os.RemoveAll(dataPath)
	os.MkdirAll(dataPath, 0700)

	storage = NewFileStorage(dataPath)

	os.Exit(m.Run())
}

func TestEmptyDirectory(t *testing.T) {
	os.RemoveAll(storage.baseDir)
	os.MkdirAll(storage.baseDir, 0700)

	t.Run("same data between src and dst", func(t *testing.T) {
		photo := model.PhotoOf(*model.IdentifierOf("testdata"), testdata)
		_, err := storage.Save(*photo)

		if assert.NoError(t, err) {
			file, _ := os.Open(path.Join(storage.baseDir, "testdata"))
			actual, _ := ioutil.ReadAll(file)

			assert.EqualValues(t, testdata, actual)

			files, _ := ioutil.ReadDir(storage.baseDir)
			assert.EqualValues(t, 1, len(files))
		}
	})
}

func TestExistData(t *testing.T) {
	os.RemoveAll(storage.baseDir)
	os.MkdirAll(storage.baseDir, 0700)
	err := ioutil.WriteFile(path.Join(storage.baseDir, "testdata"), testdata, 0700)
	assert.NoError(t, err, "failure testdata setting.")

	t.Run("same data between src and read", func(t *testing.T) {
		file, _ := os.Open(path.Join(storage.baseDir, "testdata"))
		actual, _ := ioutil.ReadAll(file)

		assert.EqualValues(t, testdata, actual)
	})

	t.Run("deleted data", func(t *testing.T) {
		err := storage.Delete(*model.IdentifierOf("testdata"))
		if assert.NoError(t, err) {
			files, _ := ioutil.ReadDir(storage.baseDir)
			assert.EqualValues(t, 0, len(files))
		}
	})
}

func BenchmarkWithEmptyData(b *testing.B) {
	os.RemoveAll(storage.baseDir)
	os.MkdirAll(storage.baseDir, 0700)
	err := ioutil.WriteFile(path.Join(storage.baseDir, "testdata"), testdata, 0700)
	assert.NoError(b, err, "failure testdata setting.")

	b.Run("write override", func(b *testing.B) {
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			photo := model.PhotoOf(*model.IdentifierOf("testdata"), testdata)
			storage.Save(*photo)
		}
	})

	b.Run("write new", func(b *testing.B) {
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			photo := model.PhotoOf(*model.IdentifierOf(fmt.Sprintf("testdata-%d", i)), testdata)
			storage.Save(*photo)
		}
	})
}

func BenchmarkWithData(b *testing.B) {
	os.RemoveAll(storage.baseDir)
	os.MkdirAll(storage.baseDir, 0700)
	err := ioutil.WriteFile(path.Join(storage.baseDir, "testdata"), testdata, 0700)
	assert.NoError(b, err, "failure testdata setting.")

	for i := 0; i < 100; i++ {
		key := fmt.Sprintf("testdata-%d", i)
		if err := ioutil.WriteFile(path.Join(storage.baseDir, key), testdata, 0700); err != nil {
			assert.NoError(b, err, "failure testdata setting.")
		}
	}

	b.Run("read same data", func(b *testing.B) {
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			storage.Read(*model.IdentifierOf("testdata"))
		}
	})

	b.Run("read different data", func(b *testing.B) {
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			storage.Read(*model.IdentifierOf(fmt.Sprintf("testdata-%d", i)))
		}
	})
}
<commit_msg>Fix and Clean test cases in file storage test<commit_after>package file_storage

import (
	"fmt"
	"github.com/photoshelf/photoshelf-storage/domain/model"
	"github.com/stretchr/testify/assert"
	"io/ioutil"
	"os"
	"path"
	"testing"
)

func TestNew(t *testing.T) {
	t.Run("with correct directory", func(t *testing.T) {
		instance := NewFileStorage(path.Join(os.TempDir(), "file_storage"))
		assert.NotNil(t, instance)
	})
}

func TestFileStorage_Save(t *testing.T) {
	t.Run("save without identifier, generate new identifier", func(t *testing.T) {
		instance := createInstance(t)
		photo := model.NewPhoto(readTestData(t))

		identifier, err := instance.Save(*photo)
		if assert.NoError(t, err) {
			assert.NotNil(t, identifier)
		}
	})

	t.Run("save with identifier", func(t *testing.T) {
		instance := createInstance(t)
		photo := model.PhotoOf(*model.IdentifierOf("testdata"), readTestData(t))

		identifier, err := instance.Save(*photo)

		if assert.NoError(t, err) {
			t.Run("returns identifier has same value", func(t *testing.T) {
				actual := photo.Id()
				assert.EqualValues(t, actual.Value(), identifier.Value())
			})

			t.Run("stored same binary", func(t *testing.T) {
				file, err := os.Open(path.Join(instance.baseDir, "testdata"))
				if err != nil {
					t.Fatal(err)
				}
				actual, err := ioutil.ReadAll(file)
				if err != nil {
					t.Fatal(err)
				}

				assert.EqualValues(t, readTestData(t), actual)
			})
		}
	})
}

func TestFileStorage_Read(t *testing.T) {
	instance := createInstance(t)
	if err := ioutil.WriteFile(path.Join(instance.baseDir, "testdata"), readTestData(t), 0700); err != nil {
		t.Fatal(err)
	}

	t.Run("with no key, returns err", func(t *testing.T) {
		_, err := instance.Read(*model.IdentifierOf("noKey"))
		assert.Error(t, err)
	})

	t.Run("returns same data with source", func(t *testing.T) {
		photo, err := instance.Read(*model.IdentifierOf("testdata"))
		if assert.NoError(t, err) {
			assert.EqualValues(t, readTestData(t), photo.Image())
		}
	})
}

func TestFileStorage_Delete(t *testing.T) {
	instance := createInstance(t)
	if err := ioutil.WriteFile(path.Join(instance.baseDir, "testdata"), readTestData(t), 0700); err != nil {
		t.Fatal(err)
	}

	t.Run("when delete existing key, returns no error", func(t *testing.T) {
		err := instance.Delete(*model.IdentifierOf("testdata"))
		if assert.NoError(t, err) {
			files, _ := ioutil.ReadDir(instance.baseDir)
			assert.EqualValues(t, 0, len(files))
		}
	})
}

func BenchmarkFileStorage_Save(b *testing.B) {
	data := readTestData(b)

	b.Run("override", func(b *testing.B) {
		instance := createInstance(b)

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("testdata-%d", 0)
			photo := *model.PhotoOf(*model.IdentifierOf(key), data)
			instance.Save(photo)
		}
		b.StopTimer()
	})

	b.Run("with new key", func(b *testing.B) {
		instance := createInstance(b)

		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("testdata-%d", i)
			photo := *model.PhotoOf(*model.IdentifierOf(key), data)
			instance.Save(photo)
		}
		b.StopTimer()
	})
}

func BenchmarkFileStorage_Read(b *testing.B) {
	data := readTestData(b)
	instance := createInstance(b)
	for i := 0; i < 100; i++ {
		key := fmt.Sprintf("testdata-%d", i)
		if err := ioutil.WriteFile(path.Join(instance.baseDir, key), data, 0700); err != nil {
			b.Fatal(err)
		}
	}

	b.Run("same data", func(b *testing.B) {
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("testdata-%d", 0)
			instance.Read(*model.IdentifierOf(key))
		}
		b.StopTimer()
	})

	b.Run("sequential", func(b *testing.B) {
		b.ResetTimer()
		for i := 0; i < b.N; i++ {
			key := fmt.Sprintf("testdata-%d", i)
			instance.Read(*model.IdentifierOf(key))
		}
		b.StopTimer()
	})
}

func readTestData(tb testing.TB) []byte {
	tb.Helper()

	testdataPath := path.Join(os.Getenv("GOPATH"), "src/github.com/photoshelf/photoshelf-storage", "testdata")
	body, err := os.Open(path.Join(testdataPath, "e3158990bdee63f8594c260cd51a011d"))
	if err != nil {
		tb.Fatal(err)
	}
	bytea, err := ioutil.ReadAll(body)
	if err != nil {
		tb.Fatal(err)
	}
	return bytea
}

func createInstance(tb testing.TB) *FileStorage {
	tb.Helper()

	dataPath := path.Join(os.TempDir(), "file_storage")
	if err := os.RemoveAll(dataPath); err != nil {
		tb.Fatal(err)
	}
	if err := os.MkdirAll(dataPath, 0700); err != nil {
		tb.Fatal(err)
	}
	return NewFileStorage(dataPath)
}
<|endoftext|>


<commit_before>package datastore

import (
	"errors"
	"time"

	query "github.com/ipfs/go-datastore/query"
)

/*
Datastore represents storage for any key-value pair.

Datastores are general enough to be backed by all kinds of different storage:
in-memory caches, databases, a remote datastore, flat files on disk, etc.

The general idea is to wrap a more complicated storage facility in a simple,
uniform interface, keeping the freedom of using the right tools for the job.
In particular, a Datastore can aggregate other datastores in interesting ways,
like sharded (to distribute load) or tiered access (caches before databases).

While Datastores should be written general enough to accept all sorts of
values, some implementations will undoubtedly have to be specific (e.g. SQL
databases where fields should be decomposed into columns), particularly to
support queries efficiently. Moreover, certain datastores may enforce certain
types of values (e.g. requiring an io.Reader, a specific struct, etc) or
serialization formats (JSON, Protobufs, etc).

IMPORTANT: No Datastore should ever Panic! This is a cross-module interface,
and thus it should behave predictably and handle exceptional conditions with
proper error reporting. Thus, all Datastore calls may return errors, which
should be checked by callers.
*/
type Datastore interface {
	// Put stores the object `value` named by `key`.
	//
	// The generalized Datastore interface does not impose a value type,
	// allowing various datastore middleware implementations (which do not
	// handle the values directly) to be composed together.
	//
	// Ultimately, the lowest-level datastore will need to do some value checking
	// or risk getting incorrect values. It may also be useful to expose a more
	// type-safe interface to your application, and do the checking up-front.
	Put(key Key, value []byte) error

	// Get retrieves the object `value` named by `key`.
	// Get will return ErrNotFound if the key is not mapped to a value.
	Get(key Key) (value []byte, err error)

	// Has returns whether the `key` is mapped to a `value`.
	// In some contexts, it may be much cheaper only to check for existence of
	// a value, rather than retrieving the value itself. (e.g. HTTP HEAD).
	// The default implementation is found in `GetBackedHas`.
	Has(key Key) (exists bool, err error)

	// Delete removes the value for given `key`.
	Delete(key Key) error

	// Query searches the datastore and returns a query result. This function
	// may return before the query actually runs. To wait for the query:
	//
	//   result, _ := ds.Query(q)
	//
	//   // use the channel interface; result may come in at different times
	//   for entry := range result.Next() { ... }
	//
	//   // or wait for the query to be completely done
	//   entries, _ := result.Rest()
	//   for entry := range entries { ... }
	//
	Query(q query.Query) (query.Results, error)
}

type Batching interface {
	Datastore

	Batch() (Batch, error)
}

var ErrBatchUnsupported = errors.New("this datastore does not support batching")

// ThreadSafeDatastore is an interface that all threadsafe datastore should
// implement to leverage type safety checks.
type ThreadSafeDatastore interface {
	Datastore

	IsThreadSafe()
}

// CheckedDatastore is an interface that should be implemented by datastores
// which may need checking on-disk data integrity.
type CheckedDatastore interface {
	Datastore

	Check() error
}

// CheckedDatastore is an interface that should be implemented by datastores
// which want to provide a mechanism to check data integrity and/or
// error correction.
type ScrubbedDatastore interface {
	Datastore

	Scrub() error
}

// GCDatastore is an interface that should be implemented by datastores which
// don't free disk space by just removing data from them.
type GCDatastore interface {
	Datastore

	CollectGarbage() error
}

// PersistentDatastore is an interface that should be implemented by datastores
// which can report disk usage.
type PersistentDatastore interface {
	Datastore

	// DiskUsage returns the space used by a datastore, in bytes.
	DiskUsage() (uint64, error)
}

// DiskUsage checks if a Datastore is a
// PersistentDatastore and returns its DiskUsage(),
// otherwise returns 0.
func DiskUsage(d Datastore) (uint64, error) {
	persDs, ok := d.(PersistentDatastore)
	if !ok {
		return 0, nil
	}
	return persDs.DiskUsage()
}

// TTLDatastore is an interface that should be implemented by datastores that
// support expiring entries.
type TTLDatastore interface {
	Datastore

	PutWithTTL(key Key, value interface{}, ttl time.Duration) error
	SetTTL(key Key, ttl time.Duration) error
}

// Txn is an interface for transactions that can be committed or rolled back.
type Txn interface {
	Datastore

	Commit() error
	Rollback()
}

// TxDatastore is an interface that should be implemented by datastores that
// support transactions.
type TxnDatastore interface {
	Datastore

	NewTransaction(readOnly bool) Txn
}

// Errors

// ErrNotFound is returned by Get, Has, and Delete when a datastore does not
// map the given key to a value.
var ErrNotFound = errors.New("datastore: key not found")

// ErrInvalidType is returned by Put when a given value is incopatible with
// the type the datastore supports. This means a conversion (or serialization)
// is needed beforehand.
var ErrInvalidType = errors.New("datastore: invalid type error")

// GetBackedHas provides a default Datastore.Has implementation.
// It exists so Datastore.Has implementations can use it, like so:
//
// func (*d SomeDatastore) Has(key Key) (exists bool, err error) {
//   return GetBackedHas(d, key)
// }
func GetBackedHas(ds Datastore, key Key) (bool, error) {
	_, err := ds.Get(key)
	switch err {
	case nil:
		return true, nil
	case ErrNotFound:
		return false, nil
	default:
		return false, err
	}
}

type Batch interface {
	Put(key Key, val []byte) error

	Delete(key Key) error

	Commit() error
}
<commit_msg>Rename Rollback to Discard, document Txn<commit_after>package datastore

import (
	"errors"
	"time"

	query "github.com/ipfs/go-datastore/query"
)

/*
Datastore represents storage for any key-value pair.

Datastores are general enough to be backed by all kinds of different storage:
in-memory caches, databases, a remote datastore, flat files on disk, etc.

The general idea is to wrap a more complicated storage facility in a simple,
uniform interface, keeping the freedom of using the right tools for the job.
In particular, a Datastore can aggregate other datastores in interesting ways,
like sharded (to distribute load) or tiered access (caches before databases).

While Datastores should be written general enough to accept all sorts of
values, some implementations will undoubtedly have to be specific (e.g. SQL
databases where fields should be decomposed into columns), particularly to
support queries efficiently. Moreover, certain datastores may enforce certain
types of values (e.g. requiring an io.Reader, a specific struct, etc) or
serialization formats (JSON, Protobufs, etc).

IMPORTANT: No Datastore should ever Panic! This is a cross-module interface,
and thus it should behave predictably and handle exceptional conditions with
proper error reporting. Thus, all Datastore calls may return errors, which
should be checked by callers.
*/
type Datastore interface {
	// Put stores the object `value` named by `key`.
	//
	// The generalized Datastore interface does not impose a value type,
	// allowing various datastore middleware implementations (which do not
	// handle the values directly) to be composed together.
	//
	// Ultimately, the lowest-level datastore will need to do some value checking
	// or risk getting incorrect values. It may also be useful to expose a more
	// type-safe interface to your application, and do the checking up-front.
	Put(key Key, value []byte) error

	// Get retrieves the object `value` named by `key`.
	// Get will return ErrNotFound if the key is not mapped to a value.
	Get(key Key) (value []byte, err error)

	// Has returns whether the `key` is mapped to a `value`.
	// In some contexts, it may be much cheaper only to check for existence of
	// a value, rather than retrieving the value itself. (e.g. HTTP HEAD).
	// The default implementation is found in `GetBackedHas`.
	Has(key Key) (exists bool, err error)

	// Delete removes the value for given `key`.
	Delete(key Key) error

	// Query searches the datastore and returns a query result. This function
	// may return before the query actually runs. To wait for the query:
	//
	//   result, _ := ds.Query(q)
	//
	//   // use the channel interface; result may come in at different times
	//   for entry := range result.Next() { ... }
	//
	//   // or wait for the query to be completely done
	//   entries, _ := result.Rest()
	//   for entry := range entries { ... }
	//
	Query(q query.Query) (query.Results, error)
}

type Batching interface {
	Datastore

	Batch() (Batch, error)
}

var ErrBatchUnsupported = errors.New("this datastore does not support batching")

// ThreadSafeDatastore is an interface that all threadsafe datastore should
// implement to leverage type safety checks.
type ThreadSafeDatastore interface {
	Datastore

	IsThreadSafe()
}

// CheckedDatastore is an interface that should be implemented by datastores
// which may need checking on-disk data integrity.
type CheckedDatastore interface {
	Datastore

	Check() error
}

// CheckedDatastore is an interface that should be implemented by datastores
// which want to provide a mechanism to check data integrity and/or
// error correction.
type ScrubbedDatastore interface {
	Datastore

	Scrub() error
}

// GCDatastore is an interface that should be implemented by datastores which
// don't free disk space by just removing data from them.
type GCDatastore interface {
	Datastore

	CollectGarbage() error
}

// PersistentDatastore is an interface that should be implemented by datastores
// which can report disk usage.
type PersistentDatastore interface {
	Datastore

	// DiskUsage returns the space used by a datastore, in bytes.
	DiskUsage() (uint64, error)
}

// DiskUsage checks if a Datastore is a
// PersistentDatastore and returns its DiskUsage(),
// otherwise returns 0.
func DiskUsage(d Datastore) (uint64, error) {
	persDs, ok := d.(PersistentDatastore)
	if !ok {
		return 0, nil
	}
	return persDs.DiskUsage()
}

// TTLDatastore is an interface that should be implemented by datastores that
// support expiring entries.
type TTLDatastore interface {
	Datastore

	PutWithTTL(key Key, value interface{}, ttl time.Duration) error
	SetTTL(key Key, ttl time.Duration) error
}

// Txn is an interface for transactions that can be committed or discarded.
type Txn interface {
	// Txn extends the Datastore type. Txns allow users to batch queries and
	// mutations to the Datastore into atomic groups, or transactions. Actions
	// performed on a transaction will not take hold until a successful call to
	// Commit has been made. Likewise, transactions can be aborted by calling
	// Discard before a successful Commit has been made.
	Datastore

	// Commit finalizes a transaction, attempting to commit it to the Datastore.
	// May return an error if the transaction has gone stale. The presence of an
	// error is an indication that the data was not committed to the Datastore.
	Commit() error
	// Discard throws away changes recorded in a transaction without committing
	// them to the underlying Datastore. Any calls made to Discard after Commit
	// has been successfully called will have no effect on the transaction and
	// state of the Datastore, making it safe to defer.
	Discard()
}

// TxDatastore is an interface that should be implemented by datastores that
// support transactions.
type TxnDatastore interface {
	Datastore

	NewTransaction(readOnly bool) Txn
}

// Errors

// ErrNotFound is returned by Get, Has, and Delete when a datastore does not
// map the given key to a value.
var ErrNotFound = errors.New("datastore: key not found")

// ErrInvalidType is returned by Put when a given value is incopatible with
// the type the datastore supports. This means a conversion (or serialization)
// is needed beforehand.
var ErrInvalidType = errors.New("datastore: invalid type error")

// GetBackedHas provides a default Datastore.Has implementation.
// It exists so Datastore.Has implementations can use it, like so:
//
// func (*d SomeDatastore) Has(key Key) (exists bool, err error) {
//   return GetBackedHas(d, key)
// }
func GetBackedHas(ds Datastore, key Key) (bool, error) {
	_, err := ds.Get(key)
	switch err {
	case nil:
		return true, nil
	case ErrNotFound:
		return false, nil
	default:
		return false, err
	}
}

type Batch interface {
	Put(key Key, val []byte) error

	Delete(key Key) error

	Commit() error
}
<|endoftext|>


<commit_before>package content

import (
	"bytes"
	"compress/zlib"
	"database/sql/driver"
	"encoding/gob"
	"errors"
	"fmt"
	"io/ioutil"
)

// Compressed - field compressed by zlib
type Compressed struct {
	Data []byte
}

// CompressedHeaders - field with http headers compressed by zlib
type CompressedHeaders struct {
	Compressed
}

// Compress - compress c.Data value
func (c *Compressed) Compress() ([]byte, error) {
	if len(c.Data) == 0 {
		return nil, nil
	}

	var zContent bytes.Buffer
	w, _ := zlib.NewWriterLevelDict(&zContent, 6, nil)
	_, err := w.Write(c.Data)
	w.Close()
	if err != nil {
		return nil, fmt.Errorf("zlib write error: %s", err)
	}

	return zContent.Bytes(), nil
}

// Value - prepare value for save to DB
func (c Compressed) Value() (driver.Value, error) {
	return c.Compress()
}

// Scan - load data from DB to value
func (c *Compressed) Scan(value interface{}) error {
	if value == nil {
		var nilResult []byte
		c.Data = nilResult
		return nil
	}

	obj, ok := value.([]byte)
	if !ok {
		return errors.New("Scan source was not string")
	}

	r, err := zlib.NewReader(bytes.NewReader(obj))
	if err != nil {
		return fmt.Errorf("zlib new reader error: %s", err)
	}
	result, err := ioutil.ReadAll(r)
	r.Close()
	if err != nil {
		return fmt.Errorf("zlib read all error: %s", err)
	}
	c.Data = result

	return nil
}

// IsNull - check is null
func (c *Compressed) IsNull() bool {
	return len(c.Data) == 0
}

// Equals - check is equals
func (c *Compressed) Equals(other Compressed) bool {
	return bytes.Equal([]byte(c.Data), []byte(other.Data))
}

// Set - convert headers to bytes
func (c *CompressedHeaders) Set(headers map[string]string) error {
	var buf bytes.Buffer
	encoder := gob.NewEncoder(&buf)
	err := encoder.Encode(headers)
	if err != nil {
		return fmt.Errorf("serialize headers to bytes, error: %s", err)
	}
	c.Data = buf.Bytes()

	return nil
}

// Get - convert bytes to headers
func (c *CompressedHeaders) Get() (map[string]string, error) {
	decoder := gob.NewDecoder(bytes.NewReader(c.Data))
	var headers map[string]string
	err := decoder.Decode(&headers)
	if err != nil {
		return headers, fmt.Errorf("deserialize bytes to headers, error: %s", err)
	}

	return headers, nil
}
<commit_msg>удалил неиспользуемый код для код для сжатия заголовков<commit_after>package content

import (
	"bytes"
	"compress/zlib"
	"database/sql/driver"
	"errors"
	"fmt"
	"io/ioutil"
)

// Compressed - field compressed by zlib
type Compressed struct {
	Data []byte
}

// Compress - compress c.Data value
func (c *Compressed) Compress() ([]byte, error) {
	if len(c.Data) == 0 {
		return nil, nil
	}

	var zContent bytes.Buffer
	w, _ := zlib.NewWriterLevelDict(&zContent, 6, nil)
	_, err := w.Write(c.Data)
	w.Close()
	if err != nil {
		return nil, fmt.Errorf("zlib write error: %s", err)
	}

	return zContent.Bytes(), nil
}

// Value - prepare value for save to DB
func (c Compressed) Value() (driver.Value, error) {
	return c.Compress()
}

// Scan - load data from DB to value
func (c *Compressed) Scan(value interface{}) error {
	if value == nil {
		var nilResult []byte
		c.Data = nilResult
		return nil
	}

	obj, ok := value.([]byte)
	if !ok {
		return errors.New("Scan source was not string")
	}

	r, err := zlib.NewReader(bytes.NewReader(obj))
	if err != nil {
		return fmt.Errorf("zlib new reader error: %s", err)
	}
	result, err := ioutil.ReadAll(r)
	r.Close()
	if err != nil {
		return fmt.Errorf("zlib read all error: %s", err)
	}
	c.Data = result

	return nil
}

// IsNull - check is null
func (c *Compressed) IsNull() bool {
	return len(c.Data) == 0
}

// Equals - check is equals
func (c *Compressed) Equals(other Compressed) bool {
	return bytes.Equal([]byte(c.Data), []byte(other.Data))
}
<|endoftext|>


<commit_before>package sso

import (
	"encoding/json"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/skygeario/skygear-server/pkg/auth/dependency/sso"
	"github.com/skygeario/skygear-server/pkg/core/skyerr"

	"github.com/skygeario/skygear-server/pkg/auth"
	coreAuth "github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz/policy"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/handler"
	"github.com/skygeario/skygear-server/pkg/core/inject"
	"github.com/skygeario/skygear-server/pkg/core/server"
)

func AttachAuthURLHandler(
	server *server.Server,
	authDependency auth.DependencyMap,
) *server.Server {
	server.Handle("/sso/{provider}/login_auth_url", &AuthURLHandlerFactory{
		Dependency: authDependency,
		Action:     "login",
	}).Methods("OPTIONS", "POST")
	server.Handle("/sso/{provider}/link_auth_url", &AuthURLHandlerFactory{
		Dependency: authDependency,
		Action:     "link",
	}).Methods("OPTIONS", "POST")
	return server
}

type AuthURLHandlerFactory struct {
	Dependency auth.DependencyMap
	Action     string
}

func (f AuthURLHandlerFactory) NewHandler(request *http.Request) http.Handler {
	h := &AuthURLHandler{}
	inject.DefaultRequestInject(h, f.Dependency, request)
	vars := mux.Vars(request)
	h.ProviderName = vars["provider"]
	h.Provider = h.ProviderFactory.NewProvider(h.ProviderName)
	h.Action = f.Action
	return handler.APIHandlerToHandler(h, h.TxContext)
}

func (f AuthURLHandlerFactory) ProvideAuthzPolicy() authz.Policy {
	return authz.PolicyFunc(policy.DenyNoAccessKey)
}

// AuthURLRequestPayload login handler request payload
type AuthURLRequestPayload struct {
	Options     map[string]interface{} `json:"options"`
	CallbackURL string                 `json:"callback_url"`
	UXMode      sso.UXMode             `json:"ux_mode"`
}

// Validate request payload
func (p AuthURLRequestPayload) Validate() error {
	if p.CallbackURL == "" {
		return skyerr.NewInvalidArgument("Callback url is required", []string{"callback_url"})
	}

	if !sso.IsValidUXMode(p.UXMode) {
		return skyerr.NewInvalidArgument("Invalid UX mode", []string{"ux_mode"})
	}

	return nil
}

// AuthURLHandler returns the SSO auth url by provider.
//
// curl \
//   -X POST \
//   -H "Content-Type: application/json" \
//   -H "X-Skygear-Api-Key: API_KEY" \
//   -d @- \
//   http://localhost:3000/sso/<provider>/login_auth_url \
// <<EOF
// {
//     "options": {
//       "prompt": "select_account"
//     },
//     callback_url: <url>,
//     ux_mode: <ux_mode>
// }
// EOF
//
// {
//     "result": "<auth_url>"
// }
//
// curl \
//   -X POST \
//   -H "Content-Type: application/json" \
//   -H "X-Skygear-Api-Key: API_KEY" \
//   -d @- \
//   http://localhost:3000/sso/<provider>/link_auth_url \
// <<EOF
// {
//     "options": {
//       "prompt": "select_account"
//     },
//     callback_url: <url>,
//     ux_mode: <ux_mode>
// }
// EOF
//
// {
//     "result": "<auth_url>"
// }
type AuthURLHandler struct {
	TxContext       db.TxContext           `dependency:"TxContext"`
	AuthContext     coreAuth.ContextGetter `dependency:"AuthContextGetter"`
	ProviderFactory *sso.ProviderFactory   `dependency:"SSOProviderFactory"`
	Provider        sso.Provider
	ProviderName    string
	Action          string
}

func (h AuthURLHandler) WithTx() bool {
	return true
}

func (h AuthURLHandler) DecodeRequest(request *http.Request) (handler.RequestPayload, error) {
	payload := AuthURLRequestPayload{
		// avoid nil pointer
		Options: make(sso.Options),
	}
	err := json.NewDecoder(request.Body).Decode(&payload)

	return payload, err
}

func (h AuthURLHandler) Handle(req interface{}) (resp interface{}, err error) {
	if h.Provider == nil {
		err = skyerr.NewInvalidArgument("Provider is not supported", []string{h.ProviderName})
		return
	}
	payload := req.(AuthURLRequestPayload)
	params := sso.GetURLParams{
		Options:     payload.Options,
		CallbackURL: payload.CallbackURL,
		UXMode:      payload.UXMode,
		Action:      h.Action,
	}
	if h.AuthContext.AuthInfo() != nil {
		params.UserID = h.AuthContext.AuthInfo().ID
	}
	url, err := h.Provider.GetAuthURL(params)
	if err != nil {
		return
	}
	resp = url
	return
}
<commit_msg>Tidy up /sso/<provider>/{login,link}_auth_url<commit_after>package sso

import (
	"encoding/json"
	"net/http"

	"github.com/gorilla/mux"
	"github.com/skygeario/skygear-server/pkg/auth/dependency/sso"
	"github.com/skygeario/skygear-server/pkg/core/skyerr"

	"github.com/skygeario/skygear-server/pkg/auth"
	coreAuth "github.com/skygeario/skygear-server/pkg/core/auth"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz"
	"github.com/skygeario/skygear-server/pkg/core/auth/authz/policy"
	"github.com/skygeario/skygear-server/pkg/core/db"
	"github.com/skygeario/skygear-server/pkg/core/handler"
	"github.com/skygeario/skygear-server/pkg/core/inject"
	"github.com/skygeario/skygear-server/pkg/core/server"
)

func AttachAuthURLHandler(
	server *server.Server,
	authDependency auth.DependencyMap,
) *server.Server {
	server.Handle("/sso/{provider}/login_auth_url", &AuthURLHandlerFactory{
		Dependency: authDependency,
		Action:     "login",
	}).Methods("OPTIONS", "POST")
	server.Handle("/sso/{provider}/link_auth_url", &AuthURLHandlerFactory{
		Dependency: authDependency,
		Action:     "link",
	}).Methods("OPTIONS", "POST")
	return server
}

type AuthURLHandlerFactory struct {
	Dependency auth.DependencyMap
	Action     string
}

func (f AuthURLHandlerFactory) NewHandler(request *http.Request) http.Handler {
	h := &AuthURLHandler{}
	inject.DefaultRequestInject(h, f.Dependency, request)
	vars := mux.Vars(request)
	h.ProviderID = vars["provider"]
	h.Provider = h.ProviderFactory.NewProvider(h.ProviderID)
	h.Action = f.Action
	return handler.APIHandlerToHandler(h, h.TxContext)
}

func (f AuthURLHandlerFactory) ProvideAuthzPolicy() authz.Policy {
	return authz.PolicyFunc(policy.DenyNoAccessKey)
}

// AuthURLRequestPayload login handler request payload
type AuthURLRequestPayload struct {
	Options     map[string]interface{} `json:"options"`
	CallbackURL string                 `json:"callback_url"`
	UXMode      sso.UXMode             `json:"ux_mode"`
}

// Validate request payload
func (p AuthURLRequestPayload) Validate() error {
	if p.CallbackURL == "" {
		return skyerr.NewInvalidArgument("Callback url is required", []string{"callback_url"})
	}

	if !sso.IsValidUXMode(p.UXMode) {
		return skyerr.NewInvalidArgument("Invalid UX mode", []string{"ux_mode"})
	}

	return nil
}

// AuthURLHandler returns the SSO auth url by provider.
//
// curl \
//   -X POST \
//   -H "Content-Type: application/json" \
//   -H "X-Skygear-Api-Key: API_KEY" \
//   -d @- \
//   http://localhost:3000/sso/<provider>/login_auth_url \
// <<EOF
// {
//     "options": {
//       "prompt": "select_account"
//     },
//     callback_url: <url>,
//     ux_mode: <ux_mode>
// }
// EOF
//
// {
//     "result": "<auth_url>"
// }
//
// curl \
//   -X POST \
//   -H "Content-Type: application/json" \
//   -H "X-Skygear-Api-Key: API_KEY" \
//   -d @- \
//   http://localhost:3000/sso/<provider>/link_auth_url \
// <<EOF
// {
//     "options": {
//       "prompt": "select_account"
//     },
//     callback_url: <url>,
//     ux_mode: <ux_mode>
// }
// EOF
//
// {
//     "result": "<auth_url>"
// }
type AuthURLHandler struct {
	TxContext       db.TxContext           `dependency:"TxContext"`
	AuthContext     coreAuth.ContextGetter `dependency:"AuthContextGetter"`
	ProviderFactory *sso.ProviderFactory   `dependency:"SSOProviderFactory"`
	Provider        sso.Provider
	ProviderID      string
	Action          string
}

func (h AuthURLHandler) WithTx() bool {
	return true
}

func (h AuthURLHandler) DecodeRequest(request *http.Request) (handler.RequestPayload, error) {
	payload := AuthURLRequestPayload{
		// avoid nil pointer
		Options: make(sso.Options),
	}
	err := json.NewDecoder(request.Body).Decode(&payload)

	return payload, err
}

func (h AuthURLHandler) Handle(req interface{}) (resp interface{}, err error) {
	if h.Provider == nil {
		err = skyerr.NewInvalidArgument("Provider is not supported", []string{h.ProviderID})
		return
	}
	payload := req.(AuthURLRequestPayload)
	params := sso.GetURLParams{
		Options:     payload.Options,
		CallbackURL: payload.CallbackURL,
		UXMode:      payload.UXMode,
		Action:      h.Action,
	}
	if h.AuthContext.AuthInfo() != nil {
		params.UserID = h.AuthContext.AuthInfo().ID
	}
	url, err := h.Provider.GetAuthURL(params)
	if err != nil {
		return
	}
	resp = url
	return
}
<|endoftext|>


<commit_before>package buildkite

import "fmt"
import "strings"

// You can overridden buildVersion at compile time by using:
//
//  go run -ldflags "-X github.com/buildkite/agent/buildkite.buildVersion abc" *.go --version
//
// On CI, the binaries are always build with the buildVersion variable set.

var baseVersion string = "1.0-beta.33"
var buildVersion string = ""

func Version() string {
	// Only output the build version if a pre-release
	if strings.Contains(baseVersion, "beta") || strings.Contains(baseVersion, "alpha") {
		// Use a default buildVersion if one doesn't exist
		actualBuildVersion := buildVersion
		if actualBuildVersion == "" {
			actualBuildVersion = "x"
		}

		return fmt.Sprintf("%s.%s", baseVersion, actualBuildVersion)
	} else {
		return baseVersion
	}
}
<commit_msg>Bumping to 1.0-beta.34<commit_after>package buildkite

import "fmt"
import "strings"

// You can overridden buildVersion at compile time by using:
//
//  go run -ldflags "-X github.com/buildkite/agent/buildkite.buildVersion abc" *.go --version
//
// On CI, the binaries are always build with the buildVersion variable set.

var baseVersion string = "1.0-beta.34"
var buildVersion string = ""

func Version() string {
	// Only output the build version if a pre-release
	if strings.Contains(baseVersion, "beta") || strings.Contains(baseVersion, "alpha") {
		// Use a default buildVersion if one doesn't exist
		actualBuildVersion := buildVersion
		if actualBuildVersion == "" {
			actualBuildVersion = "x"
		}

		return fmt.Sprintf("%s.%s", baseVersion, actualBuildVersion)
	} else {
		return baseVersion
	}
}
<|endoftext|>


<commit_before>package httphelper

import (
	"net"
	"net/http"
	"strings"
	"time"

	log "github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/pkg/ctxhelper"
)

func NewRequestLogger(handler http.Handler) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		rw := w.(*ResponseWriter)

		reqID, _ := ctxhelper.RequestIDFromContext(rw.Context())
		componentName, _ := ctxhelper.ComponentNameFromContext(rw.Context())
		logger := log.New(log.Ctx{"component": componentName, "req_id": reqID})
		rw.ctx = ctxhelper.NewContextLogger(rw.Context(), logger)

		start := time.Now()
		var clientIP string
		clientIPs := strings.Split(req.Header.Get("X-Forwarded-For"), ",")
		if len(clientIPs) > 0 {
			clientIP = strings.TrimSpace(clientIPs[len(clientIPs)-1])
		}
		if clientIP == "" {
			clientIP, _, _ = net.SplitHostPort(req.RemoteAddr)
		}

		logger.Info("request started", "method", req.Method, "path", req.URL.Path, "client_ip", clientIP)

		handler.ServeHTTP(rw, req)

		logger.Info("request completed", "status", rw.Status(), "duration", time.Since(start))
	})
}
<commit_msg>pkg/httphelper: Make request logger customisable<commit_after>package httphelper

import (
	"net"
	"net/http"
	"strings"
	"time"

	log "github.com/flynn/flynn/Godeps/_workspace/src/gopkg.in/inconshreveable/log15.v2"
	"github.com/flynn/flynn/pkg/ctxhelper"
)

type RequestLoggerFn func(handler http.Handler, logger log.Logger, clientIP string, rw *ResponseWriter, req *http.Request)

func NewRequestLogger(handler http.Handler) http.Handler {
	return newRequestLogger(handler, defaultLoggerFn)
}

func NewRequestLoggerCustom(handler http.Handler, loggerFn RequestLoggerFn) http.Handler {
	return newRequestLogger(handler, loggerFn)
}

func defaultLoggerFn(handler http.Handler, logger log.Logger, clientIP string, rw *ResponseWriter, req *http.Request) {
	start := time.Now()
	logger.Info("request started", "method", req.Method, "path", req.URL.Path, "client_ip", clientIP)
	handler.ServeHTTP(rw, req)
	logger.Info("request completed", "status", rw.Status(), "duration", time.Since(start))
}

func newRequestLogger(handler http.Handler, loggerFn RequestLoggerFn) http.Handler {
	return http.HandlerFunc(func(w http.ResponseWriter, req *http.Request) {
		rw := w.(*ResponseWriter)

		reqID, _ := ctxhelper.RequestIDFromContext(rw.Context())
		componentName, _ := ctxhelper.ComponentNameFromContext(rw.Context())
		logger := log.New(log.Ctx{"component": componentName, "req_id": reqID})
		rw.ctx = ctxhelper.NewContextLogger(rw.Context(), logger)

		var clientIP string
		clientIPs := strings.Split(req.Header.Get("X-Forwarded-For"), ",")
		if len(clientIPs) > 0 {
			clientIP = strings.TrimSpace(clientIPs[len(clientIPs)-1])
		}
		if clientIP == "" {
			clientIP, _, _ = net.SplitHostPort(req.RemoteAddr)
		}

		loggerFn(handler, logger, clientIP, rw, req)
	})
}
<|endoftext|>


<commit_before>/*
Copyright 2018 the Velero contributors.

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

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

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

package kuberesource

import (
	"k8s.io/apimachinery/pkg/runtime/schema"
)

var (
	ClusterRoleBindings       = schema.GroupResource{Group: "rbac.authorization.k8s.io", Resource: "clusterrolebindings"}
	ClusterRoles              = schema.GroupResource{Group: "rbac.authorization.k8s.io", Resource: "clusterroles"}
	CustomResourceDefinitions = schema.GroupResource{Group: "apiextensions.k8s.io", Resource: "customresourcedefinitions"}
	Jobs                      = schema.GroupResource{Group: "batch", Resource: "jobs"}
	Namespaces                = schema.GroupResource{Group: "", Resource: "namespaces"}
	PersistentVolumeClaims    = schema.GroupResource{Group: "", Resource: "persistentvolumeclaims"}
	PersistentVolumes         = schema.GroupResource{Group: "", Resource: "persistentvolumes"}
	Pods                      = schema.GroupResource{Group: "", Resource: "pods"}
	ServiceAccounts           = schema.GroupResource{Group: "", Resource: "serviceaccounts"}
	VolumeSnapshotClasses     = schema.GroupResource{Group: "snapshot.storage.k8s.io", Resource: "volumesnapshotclasses"}
	VolumeSnapshots           = schema.GroupResource{Group: "snapshot.storage.k8s.io", Resource: "volumesnapshots"}
	VolumeSnapshotContents    = schema.GroupResource{Group: "snapshot.storage.k8s.io", Resource: "volumesnapshotcontents"}
)
<commit_msg>add secrets to kuberesource (#2369)<commit_after>/*
Copyright 2018 the Velero contributors.

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

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

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

package kuberesource

import (
	"k8s.io/apimachinery/pkg/runtime/schema"
)

var (
	ClusterRoleBindings       = schema.GroupResource{Group: "rbac.authorization.k8s.io", Resource: "clusterrolebindings"}
	ClusterRoles              = schema.GroupResource{Group: "rbac.authorization.k8s.io", Resource: "clusterroles"}
	CustomResourceDefinitions = schema.GroupResource{Group: "apiextensions.k8s.io", Resource: "customresourcedefinitions"}
	Jobs                      = schema.GroupResource{Group: "batch", Resource: "jobs"}
	Namespaces                = schema.GroupResource{Group: "", Resource: "namespaces"}
	PersistentVolumeClaims    = schema.GroupResource{Group: "", Resource: "persistentvolumeclaims"}
	PersistentVolumes         = schema.GroupResource{Group: "", Resource: "persistentvolumes"}
	Pods                      = schema.GroupResource{Group: "", Resource: "pods"}
	ServiceAccounts           = schema.GroupResource{Group: "", Resource: "serviceaccounts"}
	Secrets                   = schema.GroupResource{Group: "", Resource: "secrets"}
	VolumeSnapshotClasses     = schema.GroupResource{Group: "snapshot.storage.k8s.io", Resource: "volumesnapshotclasses"}
	VolumeSnapshots           = schema.GroupResource{Group: "snapshot.storage.k8s.io", Resource: "volumesnapshots"}
	VolumeSnapshotContents    = schema.GroupResource{Group: "snapshot.storage.k8s.io", Resource: "volumesnapshotcontents"}
)
<|endoftext|>


<commit_before>package scheduler

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/go-redis/redis"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// luaPoll returns the lua script used for polling triggers in redis
// TODO we should poll sched too
const luaPoll = `
local res = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #res > 0 then
  redis.call("ZREM", "` + TriggersKey + `", res[1])
  redis.call("ZADD", "` + SchedKey + `", res[2], res[1])
end
return res`

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	Broker jobs.Broker
	client *redis.Client
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client *redis.Client) *RedisScheduler {
	return &RedisScheduler{
		client: client,
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + ":" + infos.TID
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.Broker = b
	go func() {
		for range time.Tick(pollInterval) {
			if err := s.Poll(); err != nil {
				log.Warnf("[Scheduler] Failed to poll redis: %s", err)
			}
		}
	}()
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll() error {
	now := strconv.FormatInt(time.Now().UTC().Unix(), 10)
	for {
		res, err := s.client.Eval(luaPoll, []string{now}).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), ":", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			s.client.ZRem(SchedKey, results[0])
			return err
		}
		switch t := t.(type) {
		case *AtTrigger:
			job := t.Trigger()
			if _, _, err = s.Broker.PushJob(job); err != nil {
				return err
			}
			if err := s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Trigger()
			if _, _, err = s.Broker.PushJob(job); err != nil {
				return err
			}
			if err := s.addToRedis(t); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t)
}

func (s *RedisScheduler) addToRedis(t Trigger) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(time.Now())
	case *EventTrigger:
		// TODO implement this (we ignore it because of the thumbnails trigger)
		return nil
	default:
		return errors.New("Not implemented yet")
	}
	return s.client.ZAddNX(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	pipe := s.client.Pipeline()
	pipe.ZRem(TriggersKey, t.ID())
	pipe.ZRem(SchedKey, t.ID())
	_, err := pipe.Exec()
	return err
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	// TODO(pagination): use a sort of couchdb.WalkDocs function when available.
	req := &couchdb.AllDocsRequest{Limit: 1000}
	err := couchdb.GetAllDocs(db, consts.Triggers, req, &infos)
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

var _ Scheduler = &RedisScheduler{}
<commit_msg>Try again triggers that have failed<commit_after>package scheduler

import (
	"errors"
	"fmt"
	"strconv"
	"strings"
	"time"

	log "github.com/Sirupsen/logrus"
	"github.com/cozy/cozy-stack/pkg/consts"
	"github.com/cozy/cozy-stack/pkg/couchdb"
	"github.com/cozy/cozy-stack/pkg/jobs"
	"github.com/go-redis/redis"
)

// TriggersKey is the the key of the sorted set in redis used for triggers
// waiting to be activated
const TriggersKey = "triggers"

// SchedKey is the the key of the sorted set in redis used for triggers
// currently being executed
const SchedKey = "scheduling"

// pollInterval is the time interval between 2 redis polling
const pollInterval = 1 * time.Second

// luaPoll returns the lua script used for polling triggers in redis.
// If a trigger is in the scheduling key for more than 10 seconds, it is
// an error and we can try again to schedule it.
// TODO add tests
const luaPoll = `
local w = KEYS[1] - 10
local s = redis.call("ZRANGEBYSCORE", "` + SchedKey + `", 0, w, "WITHSCORES", "LIMIT", 0, 1)
if #s > 0 then
  redis.call("ZADD", "` + SchedKey + `", KEYS[1], t[1])
  return s
end
local t = redis.call("ZRANGEBYSCORE", "` + TriggersKey + `", 0, KEYS[1], "WITHSCORES", "LIMIT", 0, 1)
if #t > 0 then
  redis.call("ZREM", "` + TriggersKey + `", t[1])
  redis.call("ZADD", "` + SchedKey + `", t[2], t[1])
end
return t`

// RedisScheduler is a centralized scheduler of many triggers. It starts all of
// them and schedules jobs accordingly.
type RedisScheduler struct {
	Broker jobs.Broker
	client *redis.Client
}

// NewRedisScheduler creates a new scheduler that use redis to synchronize with
// other cozy-stack processes to schedule jobs.
func NewRedisScheduler(client *redis.Client) *RedisScheduler {
	return &RedisScheduler{
		client: client,
	}
}

func redisKey(infos *TriggerInfos) string {
	return infos.Domain + ":" + infos.TID
}

// Start a goroutine that will fetch triggers in redis to schedule their jobs
func (s *RedisScheduler) Start(b jobs.Broker) error {
	s.Broker = b
	go func() {
		for range time.Tick(pollInterval) {
			now := time.Now().UTC().Unix()
			if err := s.Poll(now); err != nil {
				log.Warnf("[Scheduler] Failed to poll redis: %s", err)
			}
		}
	}()
	return nil
}

// Poll redis to see if there are some triggers ready
func (s *RedisScheduler) Poll(now int64) error {
	keys := []string{strconv.FormatInt(now, 10)}
	for {
		res, err := s.client.Eval(luaPoll, keys).Result()
		if err != nil || res == nil {
			return err
		}
		results, ok := res.([]interface{})
		if !ok {
			return errors.New("Unexpected response from redis")
		}
		if len(results) < 2 {
			return nil
		}
		parts := strings.SplitN(results[0].(string), ":", 2)
		if len(parts) != 2 {
			s.client.ZRem(SchedKey, results[0])
			return fmt.Errorf("Invalid key %s", res)
		}
		t, err := s.Get(parts[0], parts[1])
		if err != nil {
			s.client.ZRem(SchedKey, results[0])
			return err
		}
		switch t := t.(type) {
		case *AtTrigger:
			job := t.Trigger()
			if _, _, err = s.Broker.PushJob(job); err != nil {
				return err
			}
			if err := s.deleteTrigger(t); err != nil {
				return err
			}
		case *CronTrigger:
			job := t.Trigger()
			if _, _, err = s.Broker.PushJob(job); err != nil {
				return err
			}
			if err := s.addToRedis(t); err != nil {
				return err
			}
		default:
			return errors.New("Not implemented yet")
		}
	}
}

// Add a trigger to the system, by persisting it and using redis for scheduling
// its jobs
func (s *RedisScheduler) Add(t Trigger) error {
	infos := t.Infos()
	db := couchdb.SimpleDatabasePrefix(infos.Domain)
	if err := couchdb.CreateDoc(db, infos); err != nil {
		return err
	}
	return s.addToRedis(t)
}

func (s *RedisScheduler) addToRedis(t Trigger) error {
	var timestamp time.Time
	switch t := t.(type) {
	case *AtTrigger:
		timestamp = t.at
	case *CronTrigger:
		timestamp = t.NextExecution(time.Now())
	case *EventTrigger:
		// TODO implement this (we ignore it because of the thumbnails trigger)
		return nil
	default:
		return errors.New("Not implemented yet")
	}
	return s.client.ZAddNX(TriggersKey, redis.Z{
		Score:  float64(timestamp.UTC().Unix()),
		Member: redisKey(t.Infos()),
	}).Err()
}

// Get returns the trigger with the specified ID.
func (s *RedisScheduler) Get(domain, id string) (Trigger, error) {
	var infos TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	if err := couchdb.GetDoc(db, consts.Triggers, id, &infos); err != nil {
		if couchdb.IsNotFoundError(err) {
			return nil, ErrNotFoundTrigger
		}
		return nil, err
	}
	return NewTrigger(&infos)
}

// Delete removes the trigger with the specified ID. The trigger is unscheduled
// and remove from the storage.
func (s *RedisScheduler) Delete(domain, id string) error {
	t, err := s.Get(domain, id)
	if err != nil {
		return err
	}
	return s.deleteTrigger(t)
}

func (s *RedisScheduler) deleteTrigger(t Trigger) error {
	db := couchdb.SimpleDatabasePrefix(t.Infos().Domain)
	if err := couchdb.DeleteDoc(db, t.Infos()); err != nil {
		return err
	}
	pipe := s.client.Pipeline()
	pipe.ZRem(TriggersKey, t.ID())
	pipe.ZRem(SchedKey, t.ID())
	_, err := pipe.Exec()
	return err
}

// GetAll returns all the triggers for a domain, from couch.
func (s *RedisScheduler) GetAll(domain string) ([]Trigger, error) {
	var infos []*TriggerInfos
	db := couchdb.SimpleDatabasePrefix(domain)
	// TODO(pagination): use a sort of couchdb.WalkDocs function when available.
	req := &couchdb.AllDocsRequest{Limit: 1000}
	err := couchdb.GetAllDocs(db, consts.Triggers, req, &infos)
	if err != nil {
		if couchdb.IsNoDatabaseError(err) {
			return nil, nil
		}
		return nil, err
	}
	v := make([]Trigger, 0, len(infos))
	for _, info := range infos {
		t, err := NewTrigger(info)
		if err != nil {
			return nil, err
		}
		v = append(v, t)
	}
	return v, nil
}

var _ Scheduler = &RedisScheduler{}
<|endoftext|>


<commit_before>// bulk_load_es loads an ElasticSearch daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
)

// Program option vars:
var (
	daemonUrl         string
	refreshEachBatch  bool
	workers           int
	batchSize         int
	indexTemplateName string
	doLoad            bool
)

// Global vars
var (
	bufPool      sync.Pool
	batchChan    chan *bytes.Buffer
	inputDone    chan struct{}
	workersGroup sync.WaitGroup
)

// Args parsing vars
var (
	indexTemplateChoices = map[string][]byte{
		"default": defaultTemplate,
		"lossy":   lossyAggregationTemplate,
	}
)

var defaultTemplate = []byte(`
{
  "template": "*",
  "settings": {
    "index": {
      "refresh_interval": "5s"
    }
  },
  "mappings": {
    "point": {
      "_all":            { "enabled": false },
      "_source":         { "enabled": true },
      "properties": {
        "timestamp":    { "type": "date", "doc_values": true }
      }
    }
  }
}
`)

var lossyAggregationTemplate = []byte(`
{
  "template": "*",
  "settings": {
    "index": {
      "refresh_interval": "5s"
    }
  },
  "mappings": {
    "_default_": {
      "dynamic_templates": [
        {
          "strings_are_stored_exactly_for_filtering": {
            "match": "*",
            "match_mapping_type": "string",
            "mapping": { "type": "string",  "doc_values": true, "index": "not_analyzed" }
          }
        },
	{
	  "all_other_types_are_only_stored_in_column_index": {
            "match": "*",
            "mapping": { "doc_values": true, "index": "no" }
          }
        }
      ],
      "_all": { "enabled": false },
      "_source": { "enabled": false },
      "properties": {
        "timestamp": { "type": "date", "doc_values": true, "format": "epoch_millis" }
      }
    }
  }
}

`)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "http://localhost:9200", "ElasticSearch URL.")
	flag.BoolVar(&refreshEachBatch, "refresh", true, "Whether each batch is immediately indexed.")

	flag.IntVar(&batchSize, "batch-size", 5000, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")

	flag.StringVar(&indexTemplateName, "index-template", "default", "ElasticSearch index template to use (choices: default, lossy).")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.Parse()

	if _, ok := indexTemplateChoices[indexTemplateName]; !ok {
		log.Fatalf("invalid index template type")
	}
}

func main() {
	if doLoad {
		// check that there are no pre-existing index templates:
		existingIndexTemplates, err := listIndexTemplates(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		if len(existingIndexTemplates) > 0 {
			log.Fatal("There are index templates already in the data store. If you know what you are doing, clear them first with a command like: curl -XDELETE 'http://localhost:9200/_template/*'")
		}

		// check that there are no pre-existing indices:
		existingIndices, err := listIndices(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		if len(existingIndices) > 0 {
			log.Fatal("There are indices already in the data store. If you know what you are doing, clear them first with a command like: curl -XDELETE 'http://localhost:9200/_all'")
		}

		// create the index template:
		indexTemplate := indexTemplateChoices[indexTemplateName]
		err = createESTemplate(daemonUrl, "measurements_template", indexTemplate)
		if err != nil {
			log.Fatal(err)
		}
	}
	bufPool = sync.Pool{
		New: func() interface{} {
			return bytes.NewBuffer(make([]byte, 0, 4*1024*1024))
		},
	}

	batchChan = make(chan *bytes.Buffer, workers)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		cfg := HTTPWriterConfig{
			Host: daemonUrl,
			//Database: dbName,
		}
		go processBatches(NewHTTPWriter(cfg, refreshEachBatch))
	}

	scan(batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
}

// scan reads lines from stdin. It expects input in the ElasticSearch bulk
// format: two line pairs, the first line being an 'action' and the second line
// being the payload.
func scan(itemsPerBatch int) {
	buf := bufPool.Get().(*bytes.Buffer)

	var n int
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		buf.Write(scanner.Bytes())
		buf.Write([]byte("\n"))

		n++
		if n%2 == 0 && (n/2) >= itemsPerBatch {
			batchChan <- buf
			buf = bufPool.Get().(*bytes.Buffer)
			n = 0
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading input: %s", err.Error())
	}

	// Finished reading input, make sure last batch goes out.
	if n > 0 {
		batchChan <- buf
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)
}

// processBatches reads byte buffers from batchChan and writes them to the target server, while tracking stats on the write.
func processBatches(w LineProtocolWriter) {
	for batch := range batchChan {
		if !doLoad {
			continue
		}

		// Write the batch.
		_, err := w.WriteLineProtocol(batch.Bytes())
		if err != nil {
			log.Fatalf("Error writing: %s\n", err.Error())
		}

		// Return the batch buffer to the pool.
		batch.Reset()
		bufPool.Put(batch)
	}
	workersGroup.Done()
}

func createESTemplate(daemonUrl, templateName string, templateBody []byte) error {
	u, err := url.Parse(daemonUrl)
	if err != nil {
		return err
	}

	u.Path = fmt.Sprintf("_template/%s", templateName)

	req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(templateBody))
	if err != nil {
		return err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// does the body need to be read into the void?

	if resp.StatusCode != 200 {
		return fmt.Errorf("bad mapping create")
	}
	return nil
}

func createDb(daemon_url, dbname string) error {
	u, err := url.Parse(daemon_url)
	if err != nil {
		return err
	}

	// serialize params the right way:
	u.Path = "query"
	v := u.Query()
	v.Set("q", fmt.Sprintf("CREATE DATABASE %s", dbname))
	u.RawQuery = v.Encode()

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// does the body need to be read into the void?

	if resp.StatusCode != 200 {
		return fmt.Errorf("bad db create")
	}
	return nil
}

// listIndexTemplates lists the existing index templates in ElasticSearch.
func listIndexTemplates(daemonUrl string) (map[string]interface{}, error) {
	u := fmt.Sprintf("%s/_template", daemonUrl)
	resp, err := http.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var listing map[string]interface{}
	err = json.Unmarshal(body, &listing)
	if err != nil {
		return nil, err
	}

	return listing, nil
}

// listIndices lists the existing indices in ElasticSearch.
func listIndices(daemonUrl string) (map[string]interface{}, error) {
	u := fmt.Sprintf("%s/*", daemonUrl)
	resp, err := http.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var listing map[string]interface{}
	err = json.Unmarshal(body, &listing)
	if err != nil {
		return nil, err
	}

	return listing, nil
}
<commit_msg>Add Jon's schema changes.<commit_after>// bulk_load_es loads an ElasticSearch daemon with data from stdin.
//
// The caller is responsible for assuring that the database is empty before
// bulk load.
package main

import (
	"bufio"
	"bytes"
	"encoding/json"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"net/http"
	"net/url"
	"os"
	"sync"
)

// Program option vars:
var (
	daemonUrl         string
	refreshEachBatch  bool
	workers           int
	batchSize         int
	indexTemplateName string
	doLoad            bool
)

// Global vars
var (
	bufPool      sync.Pool
	batchChan    chan *bytes.Buffer
	inputDone    chan struct{}
	workersGroup sync.WaitGroup
)

// Args parsing vars
var (
	indexTemplateChoices = map[string][]byte{
		"default": defaultTemplate,
		"lossy":   lossyAggregationTemplate,
	}
)

var defaultTemplate = []byte(`
{
  "template": "*",
  "settings": {
    "index": {
      "refresh_interval": "5s"
    }
  },
  "mappings": {
    "point": {
      "_all":            { "enabled": false },
      "_source":         { "enabled": true },
      "properties": {
        "timestamp":    { "type": "date", "doc_values": true }
      }
    }
  }
}
`)

var lossyAggregationTemplate = []byte(`
{
  "template": "*",
  "settings": {
    "index": {
      "refresh_interval": "5s"
    }
  },
  "mappings": {
    "_default_": {
      "dynamic_templates": [
        {
          "strings_are_indexed_exactly_for_filtering": {
            "match": "*",
            "match_mapping_type": "string",
            "mapping": { "type": "string",  "doc_values": true, "index": "not_analyzed" }
          }
        },
	{
	  "all_other_types_are_in_column_index_but_not_searchable": {
            "match": "*",
            "mapping": { "doc_values": true, "index": "no" }
          }
        }
      ],
      "_all": { "enabled": false },
      "_source": { "enabled": false },
      "properties": {
        "timestamp": { "type": "date", "doc_values": true }
      }
    }
  }
}

`)

// Parse args:
func init() {
	flag.StringVar(&daemonUrl, "url", "http://localhost:9200", "ElasticSearch URL.")
	flag.BoolVar(&refreshEachBatch, "refresh", true, "Whether each batch is immediately indexed.")

	flag.IntVar(&batchSize, "batch-size", 5000, "Batch size (input items).")
	flag.IntVar(&workers, "workers", 1, "Number of parallel requests to make.")

	flag.StringVar(&indexTemplateName, "index-template", "default", "ElasticSearch index template to use (choices: default, lossy).")

	flag.BoolVar(&doLoad, "do-load", true, "Whether to write data. Set this flag to false to check input read speed.")

	flag.Parse()

	if _, ok := indexTemplateChoices[indexTemplateName]; !ok {
		log.Fatalf("invalid index template type")
	}
}

func main() {
	if doLoad {
		// check that there are no pre-existing index templates:
		existingIndexTemplates, err := listIndexTemplates(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		if len(existingIndexTemplates) > 0 {
			log.Fatal("There are index templates already in the data store. If you know what you are doing, clear them first with a command like: curl -XDELETE 'http://localhost:9200/_template/*'")
		}

		// check that there are no pre-existing indices:
		existingIndices, err := listIndices(daemonUrl)
		if err != nil {
			log.Fatal(err)
		}

		if len(existingIndices) > 0 {
			log.Fatal("There are indices already in the data store. If you know what you are doing, clear them first with a command like: curl -XDELETE 'http://localhost:9200/_all'")
		}

		// create the index template:
		indexTemplate := indexTemplateChoices[indexTemplateName]
		err = createESTemplate(daemonUrl, "measurements_template", indexTemplate)
		if err != nil {
			log.Fatal(err)
		}
	}
	bufPool = sync.Pool{
		New: func() interface{} {
			return bytes.NewBuffer(make([]byte, 0, 4*1024*1024))
		},
	}

	batchChan = make(chan *bytes.Buffer, workers)
	inputDone = make(chan struct{})

	for i := 0; i < workers; i++ {
		workersGroup.Add(1)
		cfg := HTTPWriterConfig{
			Host: daemonUrl,
			//Database: dbName,
		}
		go processBatches(NewHTTPWriter(cfg, refreshEachBatch))
	}

	scan(batchSize)

	<-inputDone
	close(batchChan)
	workersGroup.Wait()
}

// scan reads lines from stdin. It expects input in the ElasticSearch bulk
// format: two line pairs, the first line being an 'action' and the second line
// being the payload.
func scan(itemsPerBatch int) {
	buf := bufPool.Get().(*bytes.Buffer)

	var n int
	scanner := bufio.NewScanner(os.Stdin)
	for scanner.Scan() {
		buf.Write(scanner.Bytes())
		buf.Write([]byte("\n"))

		n++
		if n%2 == 0 && (n/2) >= itemsPerBatch {
			batchChan <- buf
			buf = bufPool.Get().(*bytes.Buffer)
			n = 0
		}
	}

	if err := scanner.Err(); err != nil {
		log.Fatalf("Error reading input: %s", err.Error())
	}

	// Finished reading input, make sure last batch goes out.
	if n > 0 {
		batchChan <- buf
	}

	// Closing inputDone signals to the application that we've read everything and can now shut down.
	close(inputDone)
}

// processBatches reads byte buffers from batchChan and writes them to the target server, while tracking stats on the write.
func processBatches(w LineProtocolWriter) {
	for batch := range batchChan {
		if !doLoad {
			continue
		}

		// Write the batch.
		_, err := w.WriteLineProtocol(batch.Bytes())
		if err != nil {
			log.Fatalf("Error writing: %s\n", err.Error())
		}

		// Return the batch buffer to the pool.
		batch.Reset()
		bufPool.Put(batch)
	}
	workersGroup.Done()
}

func createESTemplate(daemonUrl, templateName string, templateBody []byte) error {
	u, err := url.Parse(daemonUrl)
	if err != nil {
		return err
	}

	u.Path = fmt.Sprintf("_template/%s", templateName)

	req, err := http.NewRequest("PUT", u.String(), bytes.NewReader(templateBody))
	if err != nil {
		return err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// does the body need to be read into the void?

	if resp.StatusCode != 200 {
		return fmt.Errorf("bad mapping create")
	}
	return nil
}

func createDb(daemon_url, dbname string) error {
	u, err := url.Parse(daemon_url)
	if err != nil {
		return err
	}

	// serialize params the right way:
	u.Path = "query"
	v := u.Query()
	v.Set("q", fmt.Sprintf("CREATE DATABASE %s", dbname))
	u.RawQuery = v.Encode()

	req, err := http.NewRequest("GET", u.String(), nil)
	if err != nil {
		return err
	}

	client := &http.Client{}
	resp, err := client.Do(req)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	// does the body need to be read into the void?

	if resp.StatusCode != 200 {
		return fmt.Errorf("bad db create")
	}
	return nil
}

// listIndexTemplates lists the existing index templates in ElasticSearch.
func listIndexTemplates(daemonUrl string) (map[string]interface{}, error) {
	u := fmt.Sprintf("%s/_template", daemonUrl)
	resp, err := http.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var listing map[string]interface{}
	err = json.Unmarshal(body, &listing)
	if err != nil {
		return nil, err
	}

	return listing, nil
}

// listIndices lists the existing indices in ElasticSearch.
func listIndices(daemonUrl string) (map[string]interface{}, error) {
	u := fmt.Sprintf("%s/*", daemonUrl)
	resp, err := http.Get(u)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()

	body, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		return nil, err
	}

	var listing map[string]interface{}
	err = json.Unmarshal(body, &listing)
	if err != nil {
		return nil, err
	}

	return listing, nil
}
<|endoftext|>