<commit_before>package utils

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"runtime/debug"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor"

	"strings"
)

// HumanizeString Humanize separates string based on capitalizd letters
// e.g. "OrderItem" -> "Order Item"
func HumanizeString(str string) string {
	var human []rune
	for i, l := range str {
		if i > 0 && isUppercase(byte(l)) {
			if i > 0 && !isUppercase(str[i-1]) || i+1 < len(str) && !isUppercase(str[i+1]) {
				human = append(human, rune(' '))
			}
		}
		human = append(human, l)
	}
	return strings.Title(string(human))
}

func isUppercase(char byte) bool {
	return 'A' <= char && char <= 'Z'
}

var upcaseRegexp = regexp.MustCompile("[A-Z]{3,}[a-z]")

// ToParamString replaces spaces and separates words (by uppercase letters) with
// underscores in a string, also downcase it
// e.g. ToParamString -> to_param_string, To ParamString -> to_param_string
func ToParamString(str string) string {
	return gorm.ToDBName(strings.Replace(str, " ", "_", -1))
}

// PatchURL updates the query part of the current request url. You can
// access it in template by `patch_url`.
//     patch_url "google.com" "key" "value"
func PatchURL(originalURL string, params ...interface{}) (patchedURL string, err error) {
	url, err := url.Parse(originalURL)
	if err != nil {
		return
	}

	query := url.Query()
	for i := 0; i < len(params)/2; i++ {
		// Check if params is key&value pair
		key := fmt.Sprintf("%v", params[i*2])
		value := fmt.Sprintf("%v", params[i*2+1])

		if value == "" {
			query.Del(key)
		} else {
			query.Set(key, value)
		}
	}

	url.RawQuery = query.Encode()
	patchedURL = url.String()
	return
}

// GetLocale get locale from request, cookie, after get the locale, will write the locale to the cookie if possible
func GetLocale(context *qor.Context) string {
	if locale := context.Request.Header.Get("Locale"); locale != "" {
		return locale
	}

	if locale := context.Request.URL.Query().Get("locale"); locale != "" {
		if context.Writer != nil {
			context.Request.Header.Set("Locale", locale)
			c := http.Cookie{Name: "locale", Value: locale, Expires: time.Now().AddDate(1, 0, 0), Path: "/"}
			http.SetCookie(context.Writer, &c)
		}
		return locale
	}

	if locale, err := context.Request.Cookie("locale"); err == nil {
		return locale.Value
	}

	return ""
}

// Stringify stringify any data, if it is a struct, will try to use its Name, Title, Code field, else will use its primary key
func Stringify(object interface{}) string {
	if obj, ok := object.(interface {
		Stringify() string
	}); ok {
		return obj.Stringify()
	}

	scope := gorm.Scope{Value: object}
	for _, column := range []string{"Name", "Title", "Code"} {
		if field, ok := scope.FieldByName(column); ok {
			return fmt.Sprintf("%v", field.Field.Interface())
		}
	}

	if scope.PrimaryField() != nil {
		if scope.PrimaryKeyZero() {
			return ""
		}
		return fmt.Sprintf("%v#%v", scope.GetModelStruct().ModelType.Name(), scope.PrimaryKeyValue())
	}

	return fmt.Sprint(reflect.Indirect(reflect.ValueOf(object)).Interface())
}

// ModelType get value's model type
func ModelType(value interface{}) reflect.Type {
	reflectType := reflect.Indirect(reflect.ValueOf(value)).Type()

	for reflectType.Kind() == reflect.Ptr || reflectType.Kind() == reflect.Slice {
		reflectType = reflectType.Elem()
	}

	return reflectType
}

// ParseTagOption parse tag options to hash
func ParseTagOption(str string) map[string]string {
	tags := strings.Split(str, ";")
	setting := map[string]string{}
	for _, value := range tags {
		v := strings.Split(value, ":")
		k := strings.TrimSpace(strings.ToUpper(v[0]))
		if len(v) == 2 {
			setting[k] = v[1]
		} else {
			setting[k] = k
		}
	}
	return setting
}

// ExitWithMsg debug error messages and print stack
func ExitWithMsg(msg interface{}, value ...interface{}) {
	fmt.Printf("\n"+filenameWithLineNum()+"\n%v\n", append([]interface{}{msg}, value...)...)
	debug.PrintStack()
}

func filenameWithLineNum() string {
	var total = 10
	var results []string
	for i := 2; i < 15; i++ {
		if _, file, line, ok := runtime.Caller(i); ok {
			total--
			results = append(results[:0],
				append(
					[]string{fmt.Sprintf("%v:%v", strings.TrimPrefix(file, os.Getenv("GOPATH")+"src/"), line)},
					results[0:]...)...)

			if total == 0 {
				return strings.Join(results, "\n")
			}
		}
	}
	return ""
}
<commit_msg>HttpOnly cookie<commit_after>package utils

import (
	"fmt"
	"net/http"
	"net/url"
	"os"
	"reflect"
	"regexp"
	"runtime"
	"runtime/debug"
	"time"

	"github.com/jinzhu/gorm"
	"github.com/qor/qor"

	"strings"
)

// HumanizeString Humanize separates string based on capitalizd letters
// e.g. "OrderItem" -> "Order Item"
func HumanizeString(str string) string {
	var human []rune
	for i, l := range str {
		if i > 0 && isUppercase(byte(l)) {
			if i > 0 && !isUppercase(str[i-1]) || i+1 < len(str) && !isUppercase(str[i+1]) {
				human = append(human, rune(' '))
			}
		}
		human = append(human, l)
	}
	return strings.Title(string(human))
}

func isUppercase(char byte) bool {
	return 'A' <= char && char <= 'Z'
}

var upcaseRegexp = regexp.MustCompile("[A-Z]{3,}[a-z]")

// ToParamString replaces spaces and separates words (by uppercase letters) with
// underscores in a string, also downcase it
// e.g. ToParamString -> to_param_string, To ParamString -> to_param_string
func ToParamString(str string) string {
	return gorm.ToDBName(strings.Replace(str, " ", "_", -1))
}

// PatchURL updates the query part of the current request url. You can
// access it in template by `patch_url`.
//     patch_url "google.com" "key" "value"
func PatchURL(originalURL string, params ...interface{}) (patchedURL string, err error) {
	url, err := url.Parse(originalURL)
	if err != nil {
		return
	}

	query := url.Query()
	for i := 0; i < len(params)/2; i++ {
		// Check if params is key&value pair
		key := fmt.Sprintf("%v", params[i*2])
		value := fmt.Sprintf("%v", params[i*2+1])

		if value == "" {
			query.Del(key)
		} else {
			query.Set(key, value)
		}
	}

	url.RawQuery = query.Encode()
	patchedURL = url.String()
	return
}

// GetLocale get locale from request, cookie, after get the locale, will write the locale to the cookie if possible
func GetLocale(context *qor.Context) string {
	if locale := context.Request.Header.Get("Locale"); locale != "" {
		return locale
	}

	if locale := context.Request.URL.Query().Get("locale"); locale != "" {
		if context.Writer != nil {
			context.Request.Header.Set("Locale", locale)
			c := http.Cookie{Name: "locale", Value: locale, Expires: time.Now().AddDate(1, 0, 0), Path: "/", HttpOnly: true}
			http.SetCookie(context.Writer, &c)
		}
		return locale
	}

	if locale, err := context.Request.Cookie("locale"); err == nil {
		return locale.Value
	}

	return ""
}

// Stringify stringify any data, if it is a struct, will try to use its Name, Title, Code field, else will use its primary key
func Stringify(object interface{}) string {
	if obj, ok := object.(interface {
		Stringify() string
	}); ok {
		return obj.Stringify()
	}

	scope := gorm.Scope{Value: object}
	for _, column := range []string{"Name", "Title", "Code"} {
		if field, ok := scope.FieldByName(column); ok {
			return fmt.Sprintf("%v", field.Field.Interface())
		}
	}

	if scope.PrimaryField() != nil {
		if scope.PrimaryKeyZero() {
			return ""
		}
		return fmt.Sprintf("%v#%v", scope.GetModelStruct().ModelType.Name(), scope.PrimaryKeyValue())
	}

	return fmt.Sprint(reflect.Indirect(reflect.ValueOf(object)).Interface())
}

// ModelType get value's model type
func ModelType(value interface{}) reflect.Type {
	reflectType := reflect.Indirect(reflect.ValueOf(value)).Type()

	for reflectType.Kind() == reflect.Ptr || reflectType.Kind() == reflect.Slice {
		reflectType = reflectType.Elem()
	}

	return reflectType
}

// ParseTagOption parse tag options to hash
func ParseTagOption(str string) map[string]string {
	tags := strings.Split(str, ";")
	setting := map[string]string{}
	for _, value := range tags {
		v := strings.Split(value, ":")
		k := strings.TrimSpace(strings.ToUpper(v[0]))
		if len(v) == 2 {
			setting[k] = v[1]
		} else {
			setting[k] = k
		}
	}
	return setting
}

// ExitWithMsg debug error messages and print stack
func ExitWithMsg(msg interface{}, value ...interface{}) {
	fmt.Printf("\n"+filenameWithLineNum()+"\n%v\n", append([]interface{}{msg}, value...)...)
	debug.PrintStack()
}

func filenameWithLineNum() string {
	var total = 10
	var results []string
	for i := 2; i < 15; i++ {
		if _, file, line, ok := runtime.Caller(i); ok {
			total--
			results = append(results[:0],
				append(
					[]string{fmt.Sprintf("%v:%v", strings.TrimPrefix(file, os.Getenv("GOPATH")+"src/"), line)},
					results[0:]...)...)

			if total == 0 {
				return strings.Join(results, "\n")
			}
		}
	}
	return ""
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"

	"github.com/gorilla/mux"

	"github.com/lxc/lxd/shared"
)

// Helper functions
func networkGetInterfaces(d *Daemon) ([]string, error) {
	networks, err := dbNetworks(d.db)
	if err != nil {
		return nil, err
	}

	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, iface := range ifaces {
		if !shared.StringInSlice(iface.Name, networks) {
			networks = append(networks, iface.Name)
		}
	}

	return networks, nil
}

func networkIsInUse(c container, name string) bool {
	for _, d := range c.ExpandedDevices() {
		if d["type"] != "nic" {
			continue
		}

		if !shared.StringInSlice(d["nictype"], []string{"bridged", "macvlan"}) {
			continue
		}

		if d["parent"] == "" {
			continue
		}

		if d["parent"] == name {
			return true
		}
	}

	return false
}

// API endpoints
func networksGet(d *Daemon, r *http.Request) Response {
	recursionStr := r.FormValue("recursion")
	recursion, err := strconv.Atoi(recursionStr)
	if err != nil {
		recursion = 0
	}

	ifs, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	resultString := []string{}
	resultMap := []shared.NetworkConfig{}
	for _, iface := range ifs {
		if recursion == 0 {
			resultString = append(resultString, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, iface))
		} else {
			net, err := doNetworkGet(d, iface)
			if err != nil {
				continue
			}
			resultMap = append(resultMap, net)
		}
	}

	if recursion == 0 {
		return SyncResponse(true, resultString)
	}

	return SyncResponse(true, resultMap)
}

func networksPost(d *Daemon, r *http.Request) Response {
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Sanity checks
	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	if req.Type != "" && req.Type != "bridge" {
		return BadRequest(fmt.Errorf("Only 'bridge' type networks can be created"))
	}

	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return BadRequest(fmt.Errorf("The network already exists"))
	}

	// Create the database entry
	_, err = dbNetworkCreate(d.db, req.Name, req.Config)
	if err != nil {
		return InternalError(
			fmt.Errorf("Error inserting %s into database: %s", req.Name, err))
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

var networksCmd = Command{name: "networks", get: networksGet, post: networksPost}

func networkGet(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	n, err := doNetworkGet(d, name)
	if err != nil {
		return SmartError(err)
	}

	return SyncResponse(true, &n)
}

func doNetworkGet(d *Daemon, name string) (shared.NetworkConfig, error) {
	// Get some information
	osInfo, _ := net.InterfaceByName(name)
	_, dbInfo, _ := dbNetworkGet(d.db, name)

	// Sanity check
	if osInfo == nil && dbInfo == nil {
		return shared.NetworkConfig{}, os.ErrNotExist
	}

	// Prepare the response
	n := shared.NetworkConfig{}
	n.Name = name
	n.UsedBy = []string{}
	n.Config = map[string]string{}

	// Look for containers using the interface
	cts, err := dbContainersList(d.db, cTypeRegular)
	if err != nil {
		return shared.NetworkConfig{}, err
	}

	for _, ct := range cts {
		c, err := containerLoadByName(d, ct)
		if err != nil {
			return shared.NetworkConfig{}, err
		}

		if networkIsInUse(c, n.Name) {
			n.UsedBy = append(n.UsedBy, fmt.Sprintf("/%s/containers/%s", shared.APIVersion, ct))
		}
	}

	// Set the device type as needed
	if osInfo != nil && shared.IsLoopback(osInfo) {
		n.Type = "loopback"
	} else if dbInfo != nil || shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bridge", n.Name)) {
		if dbInfo != nil {
			n.Managed = true
			n.Config = dbInfo.Config
		}

		n.Type = "bridge"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/device", n.Name)) {
		n.Type = "physical"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bonding", n.Name)) {
		n.Type = "bond"
	} else {
		_, err := exec.Command("ovs-vsctl", "br-exists", n.Name).CombinedOutput()
		if err == nil {
			n.Type = "bridge"
		} else {
			n.Type = "unknown"
		}
	}

	return n, nil
}

func networkDelete(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Sanity checks
	if len(dbInfo.UsedBy) != 0 {
		return BadRequest(fmt.Errorf("Network is currently in use)"))
	}

	// Remove the network
	err := dbNetworkDelete(d.db, name)
	if err != nil {
		return SmartError(err)
	}

	return EmptySyncResponse
}

func networkPost(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Sanity checks
	if len(dbInfo.UsedBy) != 0 {
		return BadRequest(fmt.Errorf("Network is currently in use)"))
	}

	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	// Check that the name isn't already in use
	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return Conflict
	}

	// Rename the database entry
	err = dbNetworkRename(d.db, name, req.Name)
	if err != nil {
		return SmartError(err)
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

var networkCmd = Command{name: "networks/{name}", get: networkGet, delete: networkDelete, post: networkPost}
<commit_msg>network: Add ETag support<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"

	"github.com/gorilla/mux"

	"github.com/lxc/lxd/shared"
)

// Helper functions
func networkGetInterfaces(d *Daemon) ([]string, error) {
	networks, err := dbNetworks(d.db)
	if err != nil {
		return nil, err
	}

	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, iface := range ifaces {
		if !shared.StringInSlice(iface.Name, networks) {
			networks = append(networks, iface.Name)
		}
	}

	return networks, nil
}

func networkIsInUse(c container, name string) bool {
	for _, d := range c.ExpandedDevices() {
		if d["type"] != "nic" {
			continue
		}

		if !shared.StringInSlice(d["nictype"], []string{"bridged", "macvlan"}) {
			continue
		}

		if d["parent"] == "" {
			continue
		}

		if d["parent"] == name {
			return true
		}
	}

	return false
}

// API endpoints
func networksGet(d *Daemon, r *http.Request) Response {
	recursionStr := r.FormValue("recursion")
	recursion, err := strconv.Atoi(recursionStr)
	if err != nil {
		recursion = 0
	}

	ifs, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	resultString := []string{}
	resultMap := []shared.NetworkConfig{}
	for _, iface := range ifs {
		if recursion == 0 {
			resultString = append(resultString, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, iface))
		} else {
			net, err := doNetworkGet(d, iface)
			if err != nil {
				continue
			}
			resultMap = append(resultMap, net)
		}
	}

	if recursion == 0 {
		return SyncResponse(true, resultString)
	}

	return SyncResponse(true, resultMap)
}

func networksPost(d *Daemon, r *http.Request) Response {
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Sanity checks
	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	if req.Type != "" && req.Type != "bridge" {
		return BadRequest(fmt.Errorf("Only 'bridge' type networks can be created"))
	}

	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return BadRequest(fmt.Errorf("The network already exists"))
	}

	// Create the database entry
	_, err = dbNetworkCreate(d.db, req.Name, req.Config)
	if err != nil {
		return InternalError(
			fmt.Errorf("Error inserting %s into database: %s", req.Name, err))
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

var networksCmd = Command{name: "networks", get: networksGet, post: networksPost}

func networkGet(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	n, err := doNetworkGet(d, name)
	if err != nil {
		return SmartError(err)
	}

	etag := []interface{}{n.Name, n.Managed, n.Type, n.Config}

	return SyncResponseETag(true, &n, etag)
}

func doNetworkGet(d *Daemon, name string) (shared.NetworkConfig, error) {
	// Get some information
	osInfo, _ := net.InterfaceByName(name)
	_, dbInfo, _ := dbNetworkGet(d.db, name)

	// Sanity check
	if osInfo == nil && dbInfo == nil {
		return shared.NetworkConfig{}, os.ErrNotExist
	}

	// Prepare the response
	n := shared.NetworkConfig{}
	n.Name = name
	n.UsedBy = []string{}
	n.Config = map[string]string{}

	// Look for containers using the interface
	cts, err := dbContainersList(d.db, cTypeRegular)
	if err != nil {
		return shared.NetworkConfig{}, err
	}

	for _, ct := range cts {
		c, err := containerLoadByName(d, ct)
		if err != nil {
			return shared.NetworkConfig{}, err
		}

		if networkIsInUse(c, n.Name) {
			n.UsedBy = append(n.UsedBy, fmt.Sprintf("/%s/containers/%s", shared.APIVersion, ct))
		}
	}

	// Set the device type as needed
	if osInfo != nil && shared.IsLoopback(osInfo) {
		n.Type = "loopback"
	} else if dbInfo != nil || shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bridge", n.Name)) {
		if dbInfo != nil {
			n.Managed = true
			n.Config = dbInfo.Config
		}

		n.Type = "bridge"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/device", n.Name)) {
		n.Type = "physical"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bonding", n.Name)) {
		n.Type = "bond"
	} else {
		_, err := exec.Command("ovs-vsctl", "br-exists", n.Name).CombinedOutput()
		if err == nil {
			n.Type = "bridge"
		} else {
			n.Type = "unknown"
		}
	}

	return n, nil
}

func networkDelete(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Sanity checks
	if len(dbInfo.UsedBy) != 0 {
		return BadRequest(fmt.Errorf("Network is currently in use)"))
	}

	// Remove the network
	err := dbNetworkDelete(d.db, name)
	if err != nil {
		return SmartError(err)
	}

	return EmptySyncResponse
}

func networkPost(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Sanity checks
	if len(dbInfo.UsedBy) != 0 {
		return BadRequest(fmt.Errorf("Network is currently in use)"))
	}

	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	// Check that the name isn't already in use
	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return Conflict
	}

	// Rename the database entry
	err = dbNetworkRename(d.db, name, req.Name)
	if err != nil {
		return SmartError(err)
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

var networkCmd = Command{name: "networks/{name}", get: networkGet, delete: networkDelete, post: networkPost}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"

	"github.com/gorilla/mux"

	"github.com/lxc/lxd/shared"
)

// Helper functions
func networkGetInterfaces(d *Daemon) ([]string, error) {
	networks, err := dbNetworks(d.db)
	if err != nil {
		return nil, err
	}

	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, iface := range ifaces {
		if !shared.StringInSlice(iface.Name, networks) {
			networks = append(networks, iface.Name)
		}
	}

	return networks, nil
}

func networkIsInUse(c container, name string) bool {
	for _, d := range c.ExpandedDevices() {
		if d["type"] != "nic" {
			continue
		}

		if !shared.StringInSlice(d["nictype"], []string{"bridged", "macvlan"}) {
			continue
		}

		if d["parent"] == "" {
			continue
		}

		if d["parent"] == name {
			return true
		}
	}

	return false
}

// API endpoints
func networksGet(d *Daemon, r *http.Request) Response {
	recursionStr := r.FormValue("recursion")
	recursion, err := strconv.Atoi(recursionStr)
	if err != nil {
		recursion = 0
	}

	ifs, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	resultString := []string{}
	resultMap := []shared.NetworkConfig{}
	for _, iface := range ifs {
		if recursion == 0 {
			resultString = append(resultString, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, iface))
		} else {
			net, err := doNetworkGet(d, iface)
			if err != nil {
				continue
			}
			resultMap = append(resultMap, net)
		}
	}

	if recursion == 0 {
		return SyncResponse(true, resultString)
	}

	return SyncResponse(true, resultMap)
}

func networksPost(d *Daemon, r *http.Request) Response {
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Sanity checks
	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	err = networkValidName(req.Name)
	if err != nil {
		return BadRequest(err)
	}

	if req.Type != "" && req.Type != "bridge" {
		return BadRequest(fmt.Errorf("Only 'bridge' type networks can be created"))
	}

	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return BadRequest(fmt.Errorf("The network already exists"))
	}

	err = networkValidateConfig(req.Config)
	if err != nil {
		return BadRequest(err)
	}

	// Set some default values where needed
	if req.Config["bridge.mode"] == "fan" {
		if req.Config["fan.underlay_subnet"] == "" {
			req.Config["fan.underlay_subnet"] = "auto"
		}
	} else {
		if req.Config["ipv4.address"] == "" {
			req.Config["ipv4.address"] = "auto"
			if req.Config["ipv4.nat"] == "" {
				req.Config["ipv4.nat"] = "true"
			}
		}

		if req.Config["ipv6.address"] == "" {
			content, err := ioutil.ReadFile("/proc/sys/net/ipv6/conf/default/disable_ipv6")
			if err == nil && string(content) == "0\n" {
				req.Config["ipv6.address"] = "auto"
				if req.Config["ipv6.nat"] == "" {
					req.Config["ipv6.nat"] = "true"
				}
			}
		}
	}

	// Replace "auto" by actual values
	err = networkFillAuto(req.Config)
	if err != nil {
		return InternalError(err)
	}

	// Create the database entry
	_, err = dbNetworkCreate(d.db, req.Name, req.Config)
	if err != nil {
		return InternalError(
			fmt.Errorf("Error inserting %s into database: %s", req.Name, err))
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

var networksCmd = Command{name: "networks", get: networksGet, post: networksPost}

func networkGet(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	n, err := doNetworkGet(d, name)
	if err != nil {
		return SmartError(err)
	}

	etag := []interface{}{n.Name, n.Managed, n.Type, n.Config}

	return SyncResponseETag(true, &n, etag)
}

func doNetworkGet(d *Daemon, name string) (shared.NetworkConfig, error) {
	// Get some information
	osInfo, _ := net.InterfaceByName(name)
	_, dbInfo, _ := dbNetworkGet(d.db, name)

	// Sanity check
	if osInfo == nil && dbInfo == nil {
		return shared.NetworkConfig{}, os.ErrNotExist
	}

	// Prepare the response
	n := shared.NetworkConfig{}
	n.Name = name
	n.UsedBy = []string{}
	n.Config = map[string]string{}

	// Look for containers using the interface
	cts, err := dbContainersList(d.db, cTypeRegular)
	if err != nil {
		return shared.NetworkConfig{}, err
	}

	for _, ct := range cts {
		c, err := containerLoadByName(d, ct)
		if err != nil {
			return shared.NetworkConfig{}, err
		}

		if networkIsInUse(c, n.Name) {
			n.UsedBy = append(n.UsedBy, fmt.Sprintf("/%s/containers/%s", shared.APIVersion, ct))
		}
	}

	// Set the device type as needed
	if osInfo != nil && shared.IsLoopback(osInfo) {
		n.Type = "loopback"
	} else if dbInfo != nil || shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bridge", n.Name)) {
		if dbInfo != nil {
			n.Managed = true
			n.Config = dbInfo.Config
		}

		n.Type = "bridge"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/device", n.Name)) {
		n.Type = "physical"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bonding", n.Name)) {
		n.Type = "bond"
	} else {
		_, err := exec.Command("ovs-vsctl", "br-exists", n.Name).CombinedOutput()
		if err == nil {
			n.Type = "bridge"
		} else {
			n.Type = "unknown"
		}
	}

	return n, nil
}

func networkDelete(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Sanity checks
	if len(dbInfo.UsedBy) != 0 {
		return BadRequest(fmt.Errorf("Network is currently in use)"))
	}

	// Remove the network
	err := dbNetworkDelete(d.db, name)
	if err != nil {
		return SmartError(err)
	}

	return EmptySyncResponse
}

func networkPost(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Sanity checks
	if len(dbInfo.UsedBy) != 0 {
		return BadRequest(fmt.Errorf("Network is currently in use)"))
	}

	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	err = networkValidName(req.Name)
	if err != nil {
		return BadRequest(err)
	}

	// Check that the name isn't already in use
	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return Conflict
	}

	// Rename the database entry
	err = dbNetworkRename(d.db, name, req.Name)
	if err != nil {
		return SmartError(err)
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

func networkPut(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Validate the ETag
	etag := []interface{}{dbInfo.Name, dbInfo.Managed, dbInfo.Type, dbInfo.Config}

	err := etagCheck(r, etag)
	if err != nil {
		return PreconditionFailed(err)
	}

	req := shared.NetworkConfig{}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return BadRequest(err)
	}

	return doNetworkUpdate(d, name, dbInfo.Config, req.Config)
}

func networkPatch(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Validate the ETag
	etag := []interface{}{dbInfo.Name, dbInfo.Managed, dbInfo.Type, dbInfo.Config}

	err := etagCheck(r, etag)
	if err != nil {
		return PreconditionFailed(err)
	}

	req := shared.NetworkConfig{}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return BadRequest(err)
	}

	// Config stacking
	if req.Config == nil {
		req.Config = map[string]string{}
	}

	for k, v := range dbInfo.Config {
		_, ok := req.Config[k]
		if !ok {
			req.Config[k] = v
		}
	}

	return doNetworkUpdate(d, name, dbInfo.Config, req.Config)
}

func doNetworkUpdate(d *Daemon, name string, oldConfig map[string]string, newConfig map[string]string) Response {
	// Validate the configuration
	err := networkValidateConfig(newConfig)
	if err != nil {
		return BadRequest(err)
	}

	if newConfig["bridge.mode"] == "fan" {
		if newConfig["fan.underlay_subnet"] == "" {
			newConfig["fan.underlay_subnet"] = "auto"
		}
	}

	// Replace "auto" by actual values
	err = networkFillAuto(newConfig)
	if err != nil {
		return InternalError(err)
	}

	err = dbNetworkUpdate(d.db, name, newConfig)
	if err != nil {
		return InternalError(err)
	}

	return EmptySyncResponse
}

var networkCmd = Command{name: "networks/{name}", get: networkGet, delete: networkDelete, post: networkPost, put: networkPut, patch: networkPatch}
<commit_msg>network: Implement initial network struct<commit_after>package main

import (
	"encoding/json"
	"fmt"
	"io/ioutil"
	"net"
	"net/http"
	"os"
	"os/exec"
	"strconv"

	"github.com/gorilla/mux"

	"github.com/lxc/lxd/shared"
)

// Helper functions
func networkGetInterfaces(d *Daemon) ([]string, error) {
	networks, err := dbNetworks(d.db)
	if err != nil {
		return nil, err
	}

	ifaces, err := net.Interfaces()
	if err != nil {
		return nil, err
	}

	for _, iface := range ifaces {
		if !shared.StringInSlice(iface.Name, networks) {
			networks = append(networks, iface.Name)
		}
	}

	return networks, nil
}

func networkIsInUse(c container, name string) bool {
	for _, d := range c.ExpandedDevices() {
		if d["type"] != "nic" {
			continue
		}

		if !shared.StringInSlice(d["nictype"], []string{"bridged", "macvlan"}) {
			continue
		}

		if d["parent"] == "" {
			continue
		}

		if d["parent"] == name {
			return true
		}
	}

	return false
}

// API endpoints
func networksGet(d *Daemon, r *http.Request) Response {
	recursionStr := r.FormValue("recursion")
	recursion, err := strconv.Atoi(recursionStr)
	if err != nil {
		recursion = 0
	}

	ifs, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	resultString := []string{}
	resultMap := []shared.NetworkConfig{}
	for _, iface := range ifs {
		if recursion == 0 {
			resultString = append(resultString, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, iface))
		} else {
			net, err := doNetworkGet(d, iface)
			if err != nil {
				continue
			}
			resultMap = append(resultMap, net)
		}
	}

	if recursion == 0 {
		return SyncResponse(true, resultString)
	}

	return SyncResponse(true, resultMap)
}

func networksPost(d *Daemon, r *http.Request) Response {
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Sanity checks
	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	err = networkValidName(req.Name)
	if err != nil {
		return BadRequest(err)
	}

	if req.Type != "" && req.Type != "bridge" {
		return BadRequest(fmt.Errorf("Only 'bridge' type networks can be created"))
	}

	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return BadRequest(fmt.Errorf("The network already exists"))
	}

	err = networkValidateConfig(req.Config)
	if err != nil {
		return BadRequest(err)
	}

	// Set some default values where needed
	if req.Config["bridge.mode"] == "fan" {
		if req.Config["fan.underlay_subnet"] == "" {
			req.Config["fan.underlay_subnet"] = "auto"
		}
	} else {
		if req.Config["ipv4.address"] == "" {
			req.Config["ipv4.address"] = "auto"
			if req.Config["ipv4.nat"] == "" {
				req.Config["ipv4.nat"] = "true"
			}
		}

		if req.Config["ipv6.address"] == "" {
			content, err := ioutil.ReadFile("/proc/sys/net/ipv6/conf/default/disable_ipv6")
			if err == nil && string(content) == "0\n" {
				req.Config["ipv6.address"] = "auto"
				if req.Config["ipv6.nat"] == "" {
					req.Config["ipv6.nat"] = "true"
				}
			}
		}
	}

	// Replace "auto" by actual values
	err = networkFillAuto(req.Config)
	if err != nil {
		return InternalError(err)
	}

	// Create the database entry
	_, err = dbNetworkCreate(d.db, req.Name, req.Config)
	if err != nil {
		return InternalError(
			fmt.Errorf("Error inserting %s into database: %s", req.Name, err))
	}

	// Start the network
	n, err := networkLoadByName(d, req.Name)
	if err != nil {
		return InternalError(err)
	}

	err = n.Start()
	if err != nil {
		n.Delete()
		return InternalError(err)
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

var networksCmd = Command{name: "networks", get: networksGet, post: networksPost}

func networkGet(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	n, err := doNetworkGet(d, name)
	if err != nil {
		return SmartError(err)
	}

	etag := []interface{}{n.Name, n.Managed, n.Type, n.Config}

	return SyncResponseETag(true, &n, etag)
}

func doNetworkGet(d *Daemon, name string) (shared.NetworkConfig, error) {
	// Get some information
	osInfo, _ := net.InterfaceByName(name)
	_, dbInfo, _ := dbNetworkGet(d.db, name)

	// Sanity check
	if osInfo == nil && dbInfo == nil {
		return shared.NetworkConfig{}, os.ErrNotExist
	}

	// Prepare the response
	n := shared.NetworkConfig{}
	n.Name = name
	n.UsedBy = []string{}
	n.Config = map[string]string{}

	// Look for containers using the interface
	cts, err := dbContainersList(d.db, cTypeRegular)
	if err != nil {
		return shared.NetworkConfig{}, err
	}

	for _, ct := range cts {
		c, err := containerLoadByName(d, ct)
		if err != nil {
			return shared.NetworkConfig{}, err
		}

		if networkIsInUse(c, n.Name) {
			n.UsedBy = append(n.UsedBy, fmt.Sprintf("/%s/containers/%s", shared.APIVersion, ct))
		}
	}

	// Set the device type as needed
	if osInfo != nil && shared.IsLoopback(osInfo) {
		n.Type = "loopback"
	} else if dbInfo != nil || shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bridge", n.Name)) {
		if dbInfo != nil {
			n.Managed = true
			n.Config = dbInfo.Config
		}

		n.Type = "bridge"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/device", n.Name)) {
		n.Type = "physical"
	} else if shared.PathExists(fmt.Sprintf("/sys/class/net/%s/bonding", n.Name)) {
		n.Type = "bond"
	} else {
		_, err := exec.Command("ovs-vsctl", "br-exists", n.Name).CombinedOutput()
		if err == nil {
			n.Type = "bridge"
		} else {
			n.Type = "unknown"
		}
	}

	return n, nil
}

func networkDelete(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	n, err := networkLoadByName(d, name)
	if err != nil {
		return NotFound
	}

	// Attempt to delete the network
	err = n.Delete()
	if err != nil {
		return SmartError(err)
	}

	return EmptySyncResponse
}

func networkPost(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]
	req := shared.NetworkConfig{}

	// Parse the request
	err := json.NewDecoder(r.Body).Decode(&req)
	if err != nil {
		return BadRequest(err)
	}

	// Get the existing network
	n, err := networkLoadByName(d, name)
	if err != nil {
		return NotFound
	}

	// Sanity checks
	if req.Name == "" {
		return BadRequest(fmt.Errorf("No name provided"))
	}

	err = networkValidName(req.Name)
	if err != nil {
		return BadRequest(err)
	}

	// Check that the name isn't already in use
	networks, err := networkGetInterfaces(d)
	if err != nil {
		return InternalError(err)
	}

	if shared.StringInSlice(req.Name, networks) {
		return Conflict
	}

	// Rename it
	err = n.Rename(req.Name)
	if err != nil {
		return SmartError(err)
	}

	return SyncResponseLocation(true, nil, fmt.Sprintf("/%s/networks/%s", shared.APIVersion, req.Name))
}

func networkPut(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Validate the ETag
	etag := []interface{}{dbInfo.Name, dbInfo.Managed, dbInfo.Type, dbInfo.Config}

	err := etagCheck(r, etag)
	if err != nil {
		return PreconditionFailed(err)
	}

	req := shared.NetworkConfig{}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return BadRequest(err)
	}

	return doNetworkUpdate(d, name, dbInfo.Config, req.Config)
}

func networkPatch(d *Daemon, r *http.Request) Response {
	name := mux.Vars(r)["name"]

	// Get the existing network
	_, dbInfo, _ := dbNetworkGet(d.db, name)
	if dbInfo == nil {
		return NotFound
	}

	// Validate the ETag
	etag := []interface{}{dbInfo.Name, dbInfo.Managed, dbInfo.Type, dbInfo.Config}

	err := etagCheck(r, etag)
	if err != nil {
		return PreconditionFailed(err)
	}

	req := shared.NetworkConfig{}
	if err := json.NewDecoder(r.Body).Decode(&req); err != nil {
		return BadRequest(err)
	}

	// Config stacking
	if req.Config == nil {
		req.Config = map[string]string{}
	}

	for k, v := range dbInfo.Config {
		_, ok := req.Config[k]
		if !ok {
			req.Config[k] = v
		}
	}

	return doNetworkUpdate(d, name, dbInfo.Config, req.Config)
}

func doNetworkUpdate(d *Daemon, name string, oldConfig map[string]string, newConfig map[string]string) Response {
	// Validate the configuration
	err := networkValidateConfig(newConfig)
	if err != nil {
		return BadRequest(err)
	}

	// When switching to a fan bridge, auto-detect the underlay
	if newConfig["bridge.mode"] == "fan" {
		if newConfig["fan.underlay_subnet"] == "" {
			newConfig["fan.underlay_subnet"] = "auto"
		}
	}

	// Load the network
	n, err := networkLoadByName(d, name)
	if err != nil {
		return NotFound
	}

	err = n.Update(shared.NetworkConfig{Config: newConfig})
	if err != nil {
		return SmartError(err)
	}

	return EmptySyncResponse
}

var networkCmd = Command{name: "networks/{name}", get: networkGet, delete: networkDelete, post: networkPost, put: networkPut, patch: networkPatch}

// The network structs and functions
func networkLoadByName(d *Daemon, name string) (*network, error) {
	id, dbInfo, err := dbNetworkGet(d.db, name)
	if err != nil {
		return nil, err
	}

	n := network{d: d, id: id, name: name, config: dbInfo.Config}

	return &n, nil
}

type network struct {
	// Properties
	d    *Daemon
	id   int64
	name string

	// config
	config map[string]string
}

func (n *network) Config() map[string]string {
	return n.config
}

func (n *network) IsRunning() bool {
	return shared.PathExists(fmt.Sprintf("/sys/class/net/%s", n.name))
}

func (n *network) IsUsed() bool {
	// Look for containers using the interface
	cts, err := dbContainersList(n.d.db, cTypeRegular)
	if err != nil {
		return true
	}

	for _, ct := range cts {
		c, err := containerLoadByName(n.d, ct)
		if err != nil {
			return true
		}

		if networkIsInUse(c, n.name) {
			return true
		}
	}

	return false
}

func (n *network) Delete() error {
	// Sanity checks
	if n.IsUsed() {
		return fmt.Errorf("The network is currently in use")
	}

	// Bring the network down
	if n.IsRunning() {
		err := n.Stop()
		if err != nil {
			return err
		}
	}

	// Remove the network from the database
	err := dbNetworkDelete(n.d.db, n.name)
	if err != nil {
		return err
	}

	return nil
}

func (n *network) Rename(name string) error {
	// Sanity checks
	if n.IsUsed() {
		return fmt.Errorf("The network is currently in use")
	}

	// Bring the network down
	if n.IsRunning() {
		err := n.Stop()
		if err != nil {
			return err
		}
	}

	// Rename the database entry
	err := dbNetworkRename(n.d.db, n.name, name)
	if err != nil {
		return err
	}

	// Bring the network up
	err = n.Start()
	if err != nil {
		return err
	}

	return nil
}

func (n *network) Start() error {
	if n.IsRunning() {
		return fmt.Errorf("The network is already running")
	}

	return nil
}

func (n *network) Stop() error {
	if !n.IsRunning() {
		return fmt.Errorf("The network is already stopped")
	}

	return nil
}

func (n *network) Update(newNetwork shared.NetworkConfig) error {
	err := networkFillAuto(newNetwork.Config)
	if err != nil {
		return err
	}
	newConfig := newNetwork.Config

	// Backup the current state
	oldConfig := map[string]string{}
	err = shared.DeepCopy(&n.config, &oldConfig)
	if err != nil {
		return err
	}

	// Define a function which reverts everything.  Defer this function
	// so that it doesn't need to be explicitly called in every failing
	// return path.  Track whether or not we want to undo the changes
	// using a closure.
	undoChanges := true
	defer func() {
		if undoChanges {
			n.config = oldConfig
		}
	}()

	// Diff the configurations
	changedConfig := []string{}
	userOnly := true
	for key, _ := range oldConfig {
		if oldConfig[key] != newConfig[key] {
			if !strings.HasPrefix(key, "user.") {
				userOnly = false
			}

			if !shared.StringInSlice(key, changedConfig) {
				changedConfig = append(changedConfig, key)
			}
		}
	}

	for key, _ := range newConfig {
		if oldConfig[key] != newConfig[key] {
			if !strings.HasPrefix(key, "user.") {
				userOnly = false
			}

			if !shared.StringInSlice(key, changedConfig) {
				changedConfig = append(changedConfig, key)
			}
		}
	}

	// Skip on no change
	if len(changedConfig) == 0 {
		return nil
	}

	// Update the network
	if !userOnly {
		if shared.StringInSlice("bridge.driver", changedConfig) && n.IsRunning() {
			err = n.Stop()
			if err != nil {
				return err
			}
		}
	}

	// Apply the new configuration
	n.config = newConfig

	// Update the database
	err = dbNetworkUpdate(n.d.db, n.name, n.config)
	if err != nil {
		return err
	}

	// Restart the network
	if !userOnly {
		err = n.Start()
		if err != nil {
			return err
		}
	}

	// Success, update the closure to mark that the changes should be kept.
	undoChanges = false

	return nil
}
<|endoftext|>


<commit_before>// Copyright (c) 2010, Andrei Vieru. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	//"compress/lzma"
	"code.google.com/p/lzma"
	"flag"
	"fmt"
	"io"
	"log"
	"os"
	"path"
	"runtime"
	"strings"
)

var (
	stdout     = flag.Bool("c", false, "write on standard output, keep original files unchanged")
	decompress = flag.Bool("d", false, "decompress; see also -c and -k")
	force      = flag.Bool("f", false, "force overwrite of output file")
	help       = flag.Bool("h", false, "print this help message")
	keep       = flag.Bool("k", false, "keep original files unchaned")
	suffix     = flag.String("s", "lzma", "use provided suffix on compressed files")
	level      = flag.Int("l", 5, "compression level [1 ... 9]")
	cores      = flag.Int("cores", 1, "number of cores to use for parallelization")

	stdin bool
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTION]... [FILE]\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "Compress or uncompress FILE (by default, compress FILE in-place).\n\n")
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nWith no FILE, or when FILE is -, read standard input.\n")
}

func exit(msg string) {
	usage()
	fmt.Fprintln(os.Stderr)
	log.Fatalf("%s: check args: %s\n\n", os.Args[0], msg)
}

func setByUser(name string) (isSet bool) {
	flag.Visit(func(f *flag.Flag) {
		if f.Name == name {
			isSet = true
		}
	})
	return
}

func main() {
	flag.Parse()
	if *help == true {
		usage()
		log.Fatal(0)
	}
	//if *stdout == true && *suffix != "lzma" {
	if *stdout == true && setByUser("s") == true {
		exit("stdout set, suffix not used")
	}
	if *stdout == true && *force == true {
		exit("stdout set, force not used")
	}
	if *stdout == true && *keep == true {
		exit("stdout set, keep is redundant")
	}
	if flag.NArg() > 1 {
		exit("too many file, provide at most one file at a time or check order of flags")
	}
	//if *decompress == true && *level != 5 {
	if *decompress == true && setByUser("l") == true {
		exit("compression level is used for compression only")
	}
	if *level < 1 || *level > 9 {
		exit("compression level out of range")
	}
	if *cores < 1 || *cores > 32 {
		exit("invalid number of cores")
	}

	runtime.GOMAXPROCS(*cores)

	var inFilePath string
	var outFilePath string
	if flag.NArg() == 0 || flag.NArg() == 1 && flag.Args()[0] == "-" { // parse args: read from stdin
		if *stdout != true {
			exit("reading from stdin, can write only to stdout")
		}
		//if *suffix != "lzma" {
		if setByUser("s") == true {
			exit("reading from stdin, suffix not needed")
		}
		stdin = true

	} else if flag.NArg() == 1 { // parse args: read from file
		inFilePath = flag.Args()[0]
		f, err := os.Lstat(inFilePath)
		if err != nil {
			log.Fatal(err.Error())
		}
		if f == nil {
			exit(fmt.Sprintf("file %s not found", inFilePath))
		}
		if !!f.IsDir() {
			exit(fmt.Sprintf("%s is not a regular file", inFilePath))
		}

		if *stdout == false { // parse args: write to file
			if *suffix == "" {
				exit("suffix can't be an empty string")
			}

			if *decompress == true {
				outFileDir, outFileName := path.Split(inFilePath)
				if strings.HasSuffix(outFileName, "."+*suffix) {
					if len(outFileName) > len("."+*suffix) {
						nstr := strings.SplitN(outFileName, ".", len(outFileName))
						estr := strings.Join(nstr[0:len(nstr)-1], ".")
						outFilePath = outFileDir + estr
					} else {
						log.Fatalf("error: can't strip suffix .%s from file %s", *suffix, inFilePath)
					}
				} else {
					exit(fmt.Sprintf("file %s doesn't have suffix .%s", inFilePath, *suffix))
				}

			} else {
				outFilePath = inFilePath + "." + *suffix
			}

			f, err = os.Lstat(outFilePath)
			if err != nil && f != nil { // should be: ||| if err != nil && err != "file not found" ||| but i can't find the error's id
				log.Fatal(err.Error())
			}
			if f != nil && !f.IsDir() {
				if *force == true {
					err = os.Remove(outFilePath)
					if err != nil {
						log.Fatal(err.Error())
					}
				} else {
					exit(fmt.Sprintf("outFile %s exists. use force to overwrite", outFilePath))
				}
			} else if f != nil {
				exit(fmt.Sprintf("outFile %s exists and is not a regular file", outFilePath))
			}
		}
	}

	pr, pw := io.Pipe()
	//defer pr.Close()
	//defer pw.Close()

	if *decompress {
		// read from inFile into pw
		go func() {
			defer pw.Close()
			var inFile *os.File
			var err error
			if stdin == true {
				inFile = os.Stdin
			} else {
				inFile, err = os.Open(inFilePath)
			}
			defer inFile.Close()
			if err != nil {
				log.Fatal(err.Error())
			}

			_, err = io.Copy(pw, inFile)
			if err != nil {
				log.Fatal(err.Error())
			}

		}()

		// write into outFile from z
		defer pr.Close()
		z := lzma.NewReader(pr)
		defer z.Close()
		var outFile *os.File
		var err error
		if *stdout == true {
			outFile = os.Stdout
		} else {
			outFile, err = os.Create(outFilePath)
		}
		defer outFile.Close()
		if err != nil {
			log.Fatal(err.Error())
		}

		_, err = io.Copy(outFile, z)
		if err != nil {
			log.Fatal(err.Error())
		}

	} else {
		// read from inFile into z
		go func() {
			defer pw.Close()
			var z io.WriteCloser
			var inFile *os.File
			var err error
			if stdin == true {
				inFile = os.Stdin
				defer inFile.Close()
				z = lzma.NewWriterLevel(pw, *level)
				defer z.Close()
			} else {
				inFile, err = os.Open(inFilePath)
				defer inFile.Close()
				if err != nil {
					log.Fatal(err.Error())
				}
				f, err := os.Lstat(inFilePath)
				if err != nil {
					log.Fatal(err.Error())
				}
				z = lzma.NewWriterSizeLevel(pw, int64(f.Size()), *level)
				defer z.Close()
			}

			_, err = io.Copy(z, inFile)
			if err != nil {
				log.Fatal(err.Error())
			}
		}()

		// write into outFile from pr
		defer pr.Close()
		var outFile *os.File
		var err error
		if *stdout == true {
			outFile = os.Stdout
		} else {
			outFile, err = os.Create(outFilePath)
		}
		defer outFile.Close()
		if err != nil {
			log.Fatal(err.Error())
		}

		_, err = io.Copy(outFile, pr)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	if *stdout == false && *keep == false {
		err := os.Remove(inFilePath)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

}
<commit_msg>Fix the package import.<commit_after>// Copyright (c) 2010, Andrei Vieru. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package main

import (
	//"compress/lzma"
	"flag"
	"fmt"
	"github.com/linuxerwang/lzma"
	"io"
	"log"
	"os"
	"path"
	"runtime"
	"strings"
)

var (
	stdout     = flag.Bool("c", false, "write on standard output, keep original files unchanged")
	decompress = flag.Bool("d", false, "decompress; see also -c and -k")
	force      = flag.Bool("f", false, "force overwrite of output file")
	help       = flag.Bool("h", false, "print this help message")
	keep       = flag.Bool("k", false, "keep original files unchaned")
	suffix     = flag.String("s", "lzma", "use provided suffix on compressed files")
	level      = flag.Int("l", 5, "compression level [1 ... 9]")
	cores      = flag.Int("cores", 1, "number of cores to use for parallelization")

	stdin bool
)

func usage() {
	fmt.Fprintf(os.Stderr, "Usage: %s [OPTION]... [FILE]\n", os.Args[0])
	fmt.Fprintf(os.Stderr, "Compress or uncompress FILE (by default, compress FILE in-place).\n\n")
	flag.PrintDefaults()
	fmt.Fprintf(os.Stderr, "\nWith no FILE, or when FILE is -, read standard input.\n")
}

func exit(msg string) {
	usage()
	fmt.Fprintln(os.Stderr)
	log.Fatalf("%s: check args: %s\n\n", os.Args[0], msg)
}

func setByUser(name string) (isSet bool) {
	flag.Visit(func(f *flag.Flag) {
		if f.Name == name {
			isSet = true
		}
	})
	return
}

func main() {
	flag.Parse()
	if *help == true {
		usage()
		log.Fatal(0)
	}
	//if *stdout == true && *suffix != "lzma" {
	if *stdout == true && setByUser("s") == true {
		exit("stdout set, suffix not used")
	}
	if *stdout == true && *force == true {
		exit("stdout set, force not used")
	}
	if *stdout == true && *keep == true {
		exit("stdout set, keep is redundant")
	}
	if flag.NArg() > 1 {
		exit("too many file, provide at most one file at a time or check order of flags")
	}
	//if *decompress == true && *level != 5 {
	if *decompress == true && setByUser("l") == true {
		exit("compression level is used for compression only")
	}
	if *level < 1 || *level > 9 {
		exit("compression level out of range")
	}
	if *cores < 1 || *cores > 32 {
		exit("invalid number of cores")
	}

	runtime.GOMAXPROCS(*cores)

	var inFilePath string
	var outFilePath string
	if flag.NArg() == 0 || flag.NArg() == 1 && flag.Args()[0] == "-" { // parse args: read from stdin
		if *stdout != true {
			exit("reading from stdin, can write only to stdout")
		}
		//if *suffix != "lzma" {
		if setByUser("s") == true {
			exit("reading from stdin, suffix not needed")
		}
		stdin = true

	} else if flag.NArg() == 1 { // parse args: read from file
		inFilePath = flag.Args()[0]
		f, err := os.Lstat(inFilePath)
		if err != nil {
			log.Fatal(err.Error())
		}
		if f == nil {
			exit(fmt.Sprintf("file %s not found", inFilePath))
		}
		if !!f.IsDir() {
			exit(fmt.Sprintf("%s is not a regular file", inFilePath))
		}

		if *stdout == false { // parse args: write to file
			if *suffix == "" {
				exit("suffix can't be an empty string")
			}

			if *decompress == true {
				outFileDir, outFileName := path.Split(inFilePath)
				if strings.HasSuffix(outFileName, "."+*suffix) {
					if len(outFileName) > len("."+*suffix) {
						nstr := strings.SplitN(outFileName, ".", len(outFileName))
						estr := strings.Join(nstr[0:len(nstr)-1], ".")
						outFilePath = outFileDir + estr
					} else {
						log.Fatalf("error: can't strip suffix .%s from file %s", *suffix, inFilePath)
					}
				} else {
					exit(fmt.Sprintf("file %s doesn't have suffix .%s", inFilePath, *suffix))
				}

			} else {
				outFilePath = inFilePath + "." + *suffix
			}

			f, err = os.Lstat(outFilePath)
			if err != nil && f != nil { // should be: ||| if err != nil && err != "file not found" ||| but i can't find the error's id
				log.Fatal(err.Error())
			}
			if f != nil && !f.IsDir() {
				if *force == true {
					err = os.Remove(outFilePath)
					if err != nil {
						log.Fatal(err.Error())
					}
				} else {
					exit(fmt.Sprintf("outFile %s exists. use force to overwrite", outFilePath))
				}
			} else if f != nil {
				exit(fmt.Sprintf("outFile %s exists and is not a regular file", outFilePath))
			}
		}
	}

	pr, pw := io.Pipe()
	//defer pr.Close()
	//defer pw.Close()

	if *decompress {
		// read from inFile into pw
		go func() {
			defer pw.Close()
			var inFile *os.File
			var err error
			if stdin == true {
				inFile = os.Stdin
			} else {
				inFile, err = os.Open(inFilePath)
			}
			defer inFile.Close()
			if err != nil {
				log.Fatal(err.Error())
			}

			_, err = io.Copy(pw, inFile)
			if err != nil {
				log.Fatal(err.Error())
			}

		}()

		// write into outFile from z
		defer pr.Close()
		z := lzma.NewReader(pr)
		defer z.Close()
		var outFile *os.File
		var err error
		if *stdout == true {
			outFile = os.Stdout
		} else {
			outFile, err = os.Create(outFilePath)
		}
		defer outFile.Close()
		if err != nil {
			log.Fatal(err.Error())
		}

		_, err = io.Copy(outFile, z)
		if err != nil {
			log.Fatal(err.Error())
		}

	} else {
		// read from inFile into z
		go func() {
			defer pw.Close()
			var z io.WriteCloser
			var inFile *os.File
			var err error
			if stdin == true {
				inFile = os.Stdin
				defer inFile.Close()
				z = lzma.NewWriterLevel(pw, *level)
				defer z.Close()
			} else {
				inFile, err = os.Open(inFilePath)
				defer inFile.Close()
				if err != nil {
					log.Fatal(err.Error())
				}
				f, err := os.Lstat(inFilePath)
				if err != nil {
					log.Fatal(err.Error())
				}
				z = lzma.NewWriterSizeLevel(pw, int64(f.Size()), *level)
				defer z.Close()
			}

			_, err = io.Copy(z, inFile)
			if err != nil {
				log.Fatal(err.Error())
			}
		}()

		// write into outFile from pr
		defer pr.Close()
		var outFile *os.File
		var err error
		if *stdout == true {
			outFile = os.Stdout
		} else {
			outFile, err = os.Create(outFilePath)
		}
		defer outFile.Close()
		if err != nil {
			log.Fatal(err.Error())
		}

		_, err = io.Copy(outFile, pr)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

	if *stdout == false && *keep == false {
		err := os.Remove(inFilePath)
		if err != nil {
			log.Fatal(err.Error())
		}
	}

}
<|endoftext|>


<commit_before>// Copyright 2016 Afshin Darian. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package sleuth

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/ursiform/logger"
)

func init() {
	// Tests should run using a different group than production.
	group = "SLEUTH-vT"
}

// testCodes compares the error codes in an error with a list of wanted codes.
func testCodes(t *testing.T, err error, want []int) {
	codes := err.(*Error).Codes
	if len(codes) != len(want) {
		t.Errorf("expected codes length %d to be %d", len(codes), len(want))
		return
	}
	for i, code := range codes {
		if code != want[i] {
			t.Errorf("expected code [%d] to be [%d]", code, want[i])
			return
		}
	}
}

// Test client.go units.

func TestClientAddBadMember(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.add(group, "foo", "bar", "", "")
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
	}
	testCodes(t, err, []int{errAdd})
}

func TestClientDispatchBadAction(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte(group + "FAIL"))
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
	}
	testCodes(t, err, []int{errDispatchAction})
}

func TestClientDispatchEmpty(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte{})
	if err == nil {
		t.Errorf("expected client dispatch to fail on empty payload")
	}
	testCodes(t, err, []int{errDispatchHeader})
}

func TestClientDoUnknownScheme(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "foo://bar/baz", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown scheme")
	}
	testCodes(t, err, []int{errScheme})
}

func TestClientDoUnknownService(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "sleuth://foo/bar", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown service")
	}
	testCodes(t, err, []int{errUnknownService})
}

func TestClientReceiveBadHandle(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	res := &response{Handle: "foo"}
	err := c.receive(marshalResponse(res)[len(group)+len(recv):])
	if err == nil {
		t.Errorf("expected client receive to fail on bad handle")
	}
	testCodes(t, err, []int{errRECV})
}

func TestClientReceiveBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.receive([]byte(""))
	if err == nil {
		t.Errorf("expected client receive to fail on bad payload")
	}
	testCodes(t, err, []int{errUnzip, errResUnmarshal, errRECV})
}

func TestClientReplyBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.reply([]byte(""))
	if err == nil {
		t.Errorf("expected client reply to fail on bad payload")
	}
	testCodes(t, err, []int{errUnzip, errReqUnmarshal, errREPL})
}

// Test error.go units.

func TestError(t *testing.T) {
	code := 1
	message := "test"
	want := "sleuth: test [1]"
	err := newError(code, message)
	if err.Error() != want {
		t.Errorf("expected error to be formatted as: %s", want)
	}
}

// Test request.go units.

func TestUnmarshalRequestBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalRequest(payload)
	if err == nil {
		t.Errorf("expected unmarshalRequest to fail on bad json")
	}
	testCodes(t, err, []int{errReqUnmarshalJSON})
}

// Test response.go units.

func TestUnmarshalResponseBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalResponse(payload)
	if err == nil {
		t.Errorf("expected unmarshalResponse to fail on bad json")
	}
	testCodes(t, err, []int{errResUnmarshalJSON})
}

// Test sleuth.go units.

func TestNewBadInterface(t *testing.T) {
	_, err := New(&Config{Interface: "foo"})
	if err == nil {
		t.Errorf("expected New to fail on start with bad interface")
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestNewBadLogLevel(t *testing.T) {
	c, _ := New(&Config{LogLevel: "foo"})
	if c.log.Level() != logger.Debug {
		t.Errorf("expected log level 'foo' to be coerced to 'debug'")
		return
	}
}

func TestNewBadPort(t *testing.T) {
	_, err := New(&Config{Port: 1})
	if err == nil {
		t.Errorf("expected New to fail on start with bad port")
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestNewBadService(t *testing.T) {
	_, err := New(&Config{Handler: http.FileServer(http.Dir("."))})
	if err == nil {
		t.Errorf("expected New to fail with bad service")
	}
	testCodes(t, err, []int{errService})
}

// Test workers.go units.

func TestWorkersAddDuplicate(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if n := w.add(p); n != w.add(p) {
		t.Error("expected duplicate addition to be ignored")
	}
}

func TestWorkersAvailable(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if w.add(p); !w.available() {
		t.Error("expected workers to be available")
	}
}

func TestWorkersNext(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	w.add(&peer{name: "qux", node: "quux", service: "corge"})
	if w.next().name != "foo" {
		t.Error("expected next to return first added item")
	}
	if w.next().name != "qux" {
		t.Error("expected next to return second added item")
	}
	if w.next().name != "foo" {
		t.Error("expected next to loop back to first added item")
	}
}

func TestWorkersNextNonexistent(t *testing.T) {
	w := newWorkers()
	if p := w.next(); p != nil {
		t.Error("expected nonexistent worker to be nil")
	}
}

func TestWorkersRemove(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	if _, p := w.remove("foo"); p == nil || p.name != "foo" || w.available() {
		t.Error("expected worker to be removed")
	}
}

func TestWorkersRemoveNonexistent(t *testing.T) {
	w := newWorkers()
	if _, p := w.remove("foo"); p != nil {
		t.Error("expected nonexistent worker removal to return nil")
	}
}

// Test writer.go units.

type goodWhisperer struct{}

func (g *goodWhisperer) Whisper(addr string, payload []byte) error {
	return nil
}

type badWhisperer struct{}

func (b *badWhisperer) Whisper(addr string, payload []byte) error {
	return errors.New("bad whisperer error")
}

func TestWriterWrite(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(goodWhisperer), &destination{node: "qux", handle: "quux"})
	if n, err := w.Write(data); err != nil {
		t.Errorf("expected write to succeed: %s", err.Error())
	} else if n <= 0 {
		t.Errorf("expected (%d) to be greater than 0", n)
	}
}

func TestWriterErrWhisper(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(badWhisperer), &destination{node: "qux", handle: "quux"})
	if _, err := w.Write(data); err == nil {
		t.Errorf("expected writer to fail")
	}
}

// Test zip.go units.

func TestZipUnzip(t *testing.T) {
	in := []byte("a value that should be zipped")
	zipped := zip(in)
	if out, err := unzip(zipped); err != nil {
		t.Errorf("unzip failed: %s", err.Error())
	} else if string(out) != string(in) {
		t.Errorf("zip failed")
	}
}

func TestUnzipErr(t *testing.T) {
	in := []byte("a value that cannot be unzipped")
	_, err := unzip(in)
	testCodes(t, err, []int{errUnzip})
}

// Test integrated package.

type echoHandler struct{}

func (h *echoHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	body, _ := ioutil.ReadAll(req.Body)
	res.Write(body)
}

type silentHandler struct{}

func (*silentHandler) ServeHTTP(http.ResponseWriter, *http.Request) {
	// Silent handler does nothing.
}

func TestRequestResponseCycle(t *testing.T) {
	// Create client.
	client, err := New(nil)
	if err != nil {
		t.Errorf("client instantiation failed: %s", err.Error())
		return
	}
	defer func(client *Client, t *testing.T) {
		if err := client.Close(); err != nil {
			t.Errorf("client close failed: %s", err.Error())
		}
	}(client, t)
	// Create server.
	addr := "sleuth-test-server-one"
	server, err := New(&Config{
		Handler: new(echoHandler), Service: addr})
	if err != nil {
		t.Errorf("server instantiation failed: %s", err.Error())
		return
	}
	defer func(server *Client, t *testing.T) {
		if err := server.Close(); err != nil {
			t.Errorf("server close failed: %s", err.Error())
		}
	}(server, t)
	// Wait until the server has been added to the client pool.
	client.WaitFor(addr)
	// Set timeout to 10 seconds to accommodate slow test spin-up.
	client.Timeout = time.Second * 10
	if client.block(addr) {
		t.Errorf("call to block should have returned immediately")
	}
	body := "foo bar baz"
	buffer := bytes.NewBuffer([]byte(body))
	request, err := http.NewRequest("GET", scheme+"://"+addr+"/", buffer)
	if err != nil {
		t.Errorf("request instantiation failed: %s", err.Error())
		return
	}
	response, err := client.Do(request)
	if err != nil {
		t.Errorf("client.Do failed: %s", err.Error())
		return
	}
	if response.StatusCode != http.StatusOK {
		t.Errorf("client.Do expected %d got %d", http.StatusOK, response.StatusCode)
		return
	}
	output, _ := ioutil.ReadAll(response.Body)
	// It's not really necessary to close a sleuth response.
	response.Body.Close()
	if string(output) != body {
		t.Errorf("client.Do expected %s to equal %s", string(output), body)
		return
	}
}

func TestTimeout(t *testing.T) {
	// Create server.
	addr := "sleuth-test-server-two"
	server, err := New(&Config{Handler: new(silentHandler), Service: addr})
	if err != nil {
		t.Errorf("server instantiation failed: %s", err.Error())
		return
	}
	defer func(server *Client, t *testing.T) {
		if err := server.Close(); err != nil {
			t.Errorf("server close failed: %s", err.Error())
		}
	}(server, t)
	// Create client.
	client, err := New(nil)
	if err != nil {
		t.Errorf("client instantiation failed: %s", err.Error())
		return
	}
	defer func(client *Client, t *testing.T) {
		if err := client.Close(); err != nil {
			t.Errorf("client close failed: %s", err.Error())
		}
	}(client, t)
	// Wait long enough that the server should be ready (but not guaranteed).
	// If the server is ready, then WaitFor will not block.
	<-time.After(2000 * time.Millisecond)
	// Wait until the server has been added to the client pool.
	client.WaitFor(addr)
	request, err := http.NewRequest("GET", scheme+"://"+addr+"/", nil)
	if err != nil {
		t.Errorf("request instantiation failed: %s", err.Error())
		return
	}
	_, err = client.Do(request)
	if err == nil {
		t.Errorf("client request should have timed out")
		return
	}
}
<commit_msg>TestClientDoTimeout<commit_after>// Copyright 2016 Afshin Darian. All rights reserved.
// Use of this source code is governed by The MIT License
// that can be found in the LICENSE file.

package sleuth

import (
	"bytes"
	"errors"
	"io/ioutil"
	"net/http"
	"testing"
	"time"

	"github.com/ursiform/logger"
)

func init() {
	// Tests should run using a different group than production.
	group = "SLEUTH-vT"
}

// testCodes compares the error codes in an error with a list of wanted codes.
func testCodes(t *testing.T, err error, want []int) {
	codes := err.(*Error).Codes
	if len(codes) != len(want) {
		t.Errorf("expected codes length %d to be %d", len(codes), len(want))
		return
	}
	for i, code := range codes {
		if code != want[i] {
			t.Errorf("expected code [%d] to be [%d]", code, want[i])
			return
		}
	}
}

// Test client.go

func TestClientAddBadMember(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.add(group, "foo", "bar", "", "")
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
		return
	}
	testCodes(t, err, []int{errAdd})
}

func TestClientDispatchBadAction(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte(group + "FAIL"))
	if err == nil {
		t.Errorf("expected client dispatch to fail on bad action")
		return
	}
	testCodes(t, err, []int{errDispatchAction})
}

func TestClientDispatchEmpty(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.dispatch([]byte{})
	if err == nil {
		t.Errorf("expected client dispatch to fail on empty payload")
		return
	}
	testCodes(t, err, []int{errDispatchHeader})
}

func TestClientDoUnknownScheme(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "foo://bar/baz", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown scheme")
		return
	}
	testCodes(t, err, []int{errScheme})
}

func TestClientDoTimeout(t *testing.T) {
	c, _ := New(nil)
	defer c.Close()
	service := "foo"
	c.add(group, "bar", "baz", service, "")
	req, _ := http.NewRequest("POST", "sleuth://"+service+"/", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail by timing out")
		return
	}
	testCodes(t, err, []int{errTimeout})
}

func TestClientDoUnknownService(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	req, _ := http.NewRequest("POST", "sleuth://foo/bar", nil)
	_, err := c.Do(req)
	if err == nil {
		t.Errorf("expected client Do to fail on unknown service")
		return
	}
	testCodes(t, err, []int{errUnknownService})
}

func TestClientReceiveBadHandle(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	res := &response{Handle: "foo"}
	err := c.receive(marshalResponse(res)[len(group)+len(recv):])
	if err == nil {
		t.Errorf("expected client receive to fail on bad handle")
		return
	}
	testCodes(t, err, []int{errRECV})
}

func TestClientReceiveBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.receive([]byte(""))
	if err == nil {
		t.Errorf("expected client receive to fail on bad payload")
		return
	}
	testCodes(t, err, []int{errUnzip, errResUnmarshal, errRECV})
}

func TestClientReplyBadPayload(t *testing.T) {
	log, _ := logger.New(logger.Silent)
	c := newClient(nil, log)
	err := c.reply([]byte(""))
	if err == nil {
		t.Errorf("expected client reply to fail on bad payload")
		return
	}
	testCodes(t, err, []int{errUnzip, errReqUnmarshal, errREPL})
}

// Test error.go

func TestError(t *testing.T) {
	code := 1
	message := "test"
	want := "sleuth: test [1]"
	err := newError(code, message)
	if err.Error() != want {
		t.Errorf("expected error to be formatted as: %s", want)
	}
}

// Test request.go

func TestUnmarshalRequestBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalRequest(payload)
	if err == nil {
		t.Errorf("expected unmarshalRequest to fail on bad json")
		return
	}
	testCodes(t, err, []int{errReqUnmarshalJSON})
}

// Test response.go

func TestUnmarshalResponseBadJSON(t *testing.T) {
	payload := zip([]byte("{bad json}"))
	_, _, err := unmarshalResponse(payload)
	if err == nil {
		t.Errorf("expected unmarshalResponse to fail on bad json")
		return
	}
	testCodes(t, err, []int{errResUnmarshalJSON})
}

// Test sleuth.go

func TestNewBadInterface(t *testing.T) {
	_, err := New(&Config{Interface: "foo"})
	if err == nil {
		t.Errorf("expected New to fail on start with bad interface")
		return
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestNewBadLogLevel(t *testing.T) {
	c, _ := New(&Config{LogLevel: "foo"})
	if c.log.Level() != logger.Debug {
		t.Errorf("expected log level 'foo' to be coerced to 'debug'")
		return
	}
}

func TestNewBadPort(t *testing.T) {
	_, err := New(&Config{Port: 1})
	if err == nil {
		t.Errorf("expected New to fail on start with bad port")
		return
	}
	testCodes(t, err, []int{errStart, errCreate, errNew})
}

func TestNewBadService(t *testing.T) {
	_, err := New(&Config{Handler: http.FileServer(http.Dir("."))})
	if err == nil {
		t.Errorf("expected New to fail with bad service")
		return
	}
	testCodes(t, err, []int{errService})
}

// Test workers.go

func TestWorkersAddDuplicate(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if n := w.add(p); n != w.add(p) {
		t.Error("expected duplicate addition to be ignored")
	}
}

func TestWorkersAvailable(t *testing.T) {
	w := newWorkers()
	p := &peer{name: "foo", node: "bar", service: "baz"}
	if w.add(p); !w.available() {
		t.Error("expected workers to be available")
	}
}

func TestWorkersNext(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	w.add(&peer{name: "qux", node: "quux", service: "corge"})
	if w.next().name != "foo" {
		t.Error("expected next to return first added item")
	}
	if w.next().name != "qux" {
		t.Error("expected next to return second added item")
	}
	if w.next().name != "foo" {
		t.Error("expected next to loop back to first added item")
	}
}

func TestWorkersNextNonexistent(t *testing.T) {
	w := newWorkers()
	if p := w.next(); p != nil {
		t.Error("expected nonexistent worker to be nil")
	}
}

func TestWorkersRemove(t *testing.T) {
	w := newWorkers()
	w.add(&peer{name: "foo", node: "bar", service: "baz"})
	if _, p := w.remove("foo"); p == nil || p.name != "foo" || w.available() {
		t.Error("expected worker to be removed")
	}
}

func TestWorkersRemoveNonexistent(t *testing.T) {
	w := newWorkers()
	if _, p := w.remove("foo"); p != nil {
		t.Error("expected nonexistent worker removal to return nil")
	}
}

// Test writer.go

type goodWhisperer struct{}

func (g *goodWhisperer) Whisper(addr string, payload []byte) error {
	return nil
}

type badWhisperer struct{}

func (b *badWhisperer) Whisper(addr string, payload []byte) error {
	return errors.New("bad whisperer error")
}

func TestWriterWrite(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(goodWhisperer), &destination{node: "qux", handle: "quux"})
	if n, err := w.Write(data); err != nil {
		t.Errorf("expected write to succeed: %s", err.Error())
	} else if n <= 0 {
		t.Errorf("expected (%d) to be greater than 0", n)
	}
}

func TestWriterErrWhisper(t *testing.T) {
	data := []byte("foo bar baz")
	w := newWriter(new(badWhisperer), &destination{node: "qux", handle: "quux"})
	if _, err := w.Write(data); err == nil {
		t.Errorf("expected writer to fail")
	}
}

// Test zip.go

func TestZipUnzip(t *testing.T) {
	in := []byte("a value that should be zipped")
	zipped := zip(in)
	if out, err := unzip(zipped); err != nil {
		t.Errorf("unzip failed: %s", err.Error())
	} else if string(out) != string(in) {
		t.Errorf("zip failed")
	}
}

func TestUnzipErr(t *testing.T) {
	in := []byte("a value that cannot be unzipped")
	_, err := unzip(in)
	testCodes(t, err, []int{errUnzip})
}

// Test integrated package.

type echoHandler struct{}

func (h *echoHandler) ServeHTTP(res http.ResponseWriter, req *http.Request) {
	body, _ := ioutil.ReadAll(req.Body)
	res.Write(body)
}

type silentHandler struct{}

func (*silentHandler) ServeHTTP(http.ResponseWriter, *http.Request) {
	// Silent handler does nothing.
}

func TestRequestResponseCycle(t *testing.T) {
	// Create client.
	client, err := New(nil)
	if err != nil {
		t.Errorf("client instantiation failed: %s", err.Error())
		return
	}
	defer func(client *Client, t *testing.T) {
		if err := client.Close(); err != nil {
			t.Errorf("client close failed: %s", err.Error())
		}
	}(client, t)
	// Create server.
	addr := "sleuth-test-server-one"
	server, err := New(&Config{
		Handler: new(echoHandler), Service: addr})
	if err != nil {
		t.Errorf("server instantiation failed: %s", err.Error())
		return
	}
	defer func(server *Client, t *testing.T) {
		if err := server.Close(); err != nil {
			t.Errorf("server close failed: %s", err.Error())
		}
	}(server, t)
	// Wait until the server has been added to the client pool.
	client.WaitFor(addr)
	// Set timeout to 10 seconds to accommodate slow test spin-up.
	client.Timeout = time.Second * 10
	if client.block(addr) {
		t.Errorf("call to block should have returned immediately")
	}
	body := "foo bar baz"
	buffer := bytes.NewBuffer([]byte(body))
	request, err := http.NewRequest("GET", scheme+"://"+addr+"/", buffer)
	if err != nil {
		t.Errorf("request instantiation failed: %s", err.Error())
		return
	}
	response, err := client.Do(request)
	if err != nil {
		t.Errorf("client.Do failed: %s", err.Error())
		return
	}
	if response.StatusCode != http.StatusOK {
		t.Errorf("client.Do expected %d got %d", http.StatusOK, response.StatusCode)
		return
	}
	output, _ := ioutil.ReadAll(response.Body)
	// It's not really necessary to close a sleuth response.
	response.Body.Close()
	if string(output) != body {
		t.Errorf("client.Do expected %s to equal %s", string(output), body)
		return
	}
}
<|endoftext|>


<commit_before>package meta

import (
	"fmt"

	"github.com/influxdb/influxdb/influxql"
)

// StatementExecutor translates InfluxQL queries to meta store methods.
type StatementExecutor struct {
	Store interface {
		Nodes() ([]NodeInfo, error)

		Database(name string) (*DatabaseInfo, error)
		Databases() ([]DatabaseInfo, error)
		CreateDatabase(name string) (*DatabaseInfo, error)
		DropDatabase(name string) error

		DefaultRetentionPolicy(database string) (*RetentionPolicyInfo, error)
		CreateRetentionPolicy(database string, rpi *RetentionPolicyInfo) (*RetentionPolicyInfo, error)
		UpdateRetentionPolicy(database, name string, rpu *RetentionPolicyUpdate) error
		SetDefaultRetentionPolicy(database, name string) error
		DropRetentionPolicy(database, name string) error

		Users() ([]UserInfo, error)
		CreateUser(name, password string, admin bool) (*UserInfo, error)
		UpdateUser(name, password string) error
		DropUser(name string) error
		SetPrivilege(username, database string, p influxql.Privilege) error
		UserPrivileges(username string) (map[string]influxql.Privilege, error)

		CreateContinuousQuery(database, name, query string) error
		DropContinuousQuery(database, name string) error
	}
}

// ExecuteStatement executes stmt against the meta store as user.
func (e *StatementExecutor) ExecuteStatement(stmt influxql.Statement) *influxql.Result {
	switch stmt := stmt.(type) {
	case *influxql.CreateDatabaseStatement:
		return e.executeCreateDatabaseStatement(stmt)
	case *influxql.DropDatabaseStatement:
		return e.executeDropDatabaseStatement(stmt)
	case *influxql.ShowDatabasesStatement:
		return e.executeShowDatabasesStatement(stmt)
	case *influxql.ShowGrantsForUserStatement:
		return e.executeShowGrantsForUserStatement(stmt)
	case *influxql.ShowServersStatement:
		return e.executeShowServersStatement(stmt)
	case *influxql.CreateUserStatement:
		return e.executeCreateUserStatement(stmt)
	case *influxql.SetPasswordUserStatement:
		return e.executeSetPasswordUserStatement(stmt)
	case *influxql.DropUserStatement:
		return e.executeDropUserStatement(stmt)
	case *influxql.ShowUsersStatement:
		return e.executeShowUsersStatement(stmt)
	case *influxql.GrantStatement:
		return e.executeGrantStatement(stmt)
	case *influxql.RevokeStatement:
		return e.executeRevokeStatement(stmt)
	case *influxql.CreateRetentionPolicyStatement:
		return e.executeCreateRetentionPolicyStatement(stmt)
	case *influxql.AlterRetentionPolicyStatement:
		return e.executeAlterRetentionPolicyStatement(stmt)
	case *influxql.DropRetentionPolicyStatement:
		return e.executeDropRetentionPolicyStatement(stmt)
	case *influxql.ShowRetentionPoliciesStatement:
		return e.executeShowRetentionPoliciesStatement(stmt)
	case *influxql.CreateContinuousQueryStatement:
		return e.executeCreateContinuousQueryStatement(stmt)
	case *influxql.DropContinuousQueryStatement:
		return e.executeDropContinuousQueryStatement(stmt)
	case *influxql.ShowContinuousQueriesStatement:
		return e.executeShowContinuousQueriesStatement(stmt)
	default:
		panic(fmt.Sprintf("unsupported statement type: %T", stmt))
	}
}

func (e *StatementExecutor) executeCreateDatabaseStatement(q *influxql.CreateDatabaseStatement) *influxql.Result {
	_, err := e.Store.CreateDatabase(q.Name)
	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeDropDatabaseStatement(q *influxql.DropDatabaseStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.DropDatabase(q.Name)}
}

func (e *StatementExecutor) executeShowDatabasesStatement(q *influxql.ShowDatabasesStatement) *influxql.Result {
	dis, err := e.Store.Databases()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Name: "databases", Columns: []string{"name"}}
	for _, di := range dis {
		row.Values = append(row.Values, []interface{}{di.Name})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeShowGrantsForUserStatement(q *influxql.ShowGrantsForUserStatement) *influxql.Result {
	priv, err := e.Store.UserPrivileges(q.Name)
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Columns: []string{"database", "privilege"}}
	for d, p := range priv {
		row.Values = append(row.Values, []interface{}{d, p.String()})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeShowServersStatement(q *influxql.ShowServersStatement) *influxql.Result {
	nis, err := e.Store.Nodes()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Columns: []string{"id", "url"}}
	for _, ni := range nis {
		row.Values = append(row.Values, []interface{}{ni.ID, "http://" + ni.Host})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeCreateUserStatement(q *influxql.CreateUserStatement) *influxql.Result {
	admin := false
	if q.Privilege != nil {
		admin = (*q.Privilege == influxql.AllPrivileges)
	}

	_, err := e.Store.CreateUser(q.Name, q.Password, admin)
	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeSetPasswordUserStatement(q *influxql.SetPasswordUserStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.UpdateUser(q.Name, q.Password)}
}

func (e *StatementExecutor) executeDropUserStatement(q *influxql.DropUserStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.DropUser(q.Name)}
}

func (e *StatementExecutor) executeShowUsersStatement(q *influxql.ShowUsersStatement) *influxql.Result {
	uis, err := e.Store.Users()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Columns: []string{"user", "admin"}}
	for _, ui := range uis {
		row.Values = append(row.Values, []interface{}{ui.Name, ui.Admin})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeGrantStatement(stmt *influxql.GrantStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.SetPrivilege(stmt.User, stmt.On, stmt.Privilege)}
}

func (e *StatementExecutor) executeRevokeStatement(stmt *influxql.RevokeStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.SetPrivilege(stmt.User, stmt.On, influxql.NoPrivileges)}
}

func (e *StatementExecutor) executeCreateRetentionPolicyStatement(stmt *influxql.CreateRetentionPolicyStatement) *influxql.Result {
	rpi := NewRetentionPolicyInfo(stmt.Name)
	rpi.Duration = stmt.Duration
	rpi.ReplicaN = stmt.Replication

	// Create new retention policy.
	_, err := e.Store.CreateRetentionPolicy(stmt.Database, rpi)
	if err != nil {
		return &influxql.Result{Err: err}
	}

	// If requested, set new policy as the default.
	if stmt.Default {
		err = e.Store.SetDefaultRetentionPolicy(stmt.Database, stmt.Name)
	}

	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeAlterRetentionPolicyStatement(stmt *influxql.AlterRetentionPolicyStatement) *influxql.Result {
	rpu := &RetentionPolicyUpdate{
		Duration: stmt.Duration,
		ReplicaN: stmt.Replication,
	}

	// Update the retention policy.
	err := e.Store.UpdateRetentionPolicy(stmt.Database, stmt.Name, rpu)
	if err != nil {
		return &influxql.Result{Err: err}
	}

	// If requested, set as default retention policy.
	if stmt.Default {
		err = e.Store.SetDefaultRetentionPolicy(stmt.Database, stmt.Name)
	}

	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeDropRetentionPolicyStatement(q *influxql.DropRetentionPolicyStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.DropRetentionPolicy(q.Database, q.Name)}
}

func (e *StatementExecutor) executeShowRetentionPoliciesStatement(q *influxql.ShowRetentionPoliciesStatement) *influxql.Result {
	di, err := e.Store.Database(q.Database)
	if err != nil {
		return &influxql.Result{Err: err}
	} else if di == nil {
		return &influxql.Result{Err: ErrDatabaseNotFound}
	}

	row := &influxql.Row{Columns: []string{"name", "duration", "replicaN", "default"}}
	for _, rpi := range di.RetentionPolicies {
		row.Values = append(row.Values, []interface{}{rpi.Name, rpi.Duration.String(), rpi.ReplicaN, di.DefaultRetentionPolicy == rpi.Name})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeCreateContinuousQueryStatement(q *influxql.CreateContinuousQueryStatement) *influxql.Result {
	return &influxql.Result{
		Err: e.Store.CreateContinuousQuery(q.Database, q.Name, q.String()),
	}
}

func (e *StatementExecutor) executeDropContinuousQueryStatement(q *influxql.DropContinuousQueryStatement) *influxql.Result {
	return &influxql.Result{
		Err: e.Store.DropContinuousQuery(q.Database, q.Name),
	}
}

func (e *StatementExecutor) executeShowContinuousQueriesStatement(stmt *influxql.ShowContinuousQueriesStatement) *influxql.Result {
	dis, err := e.Store.Databases()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	rows := []*influxql.Row{}
	for _, di := range dis {
		row := &influxql.Row{Columns: []string{"name", "query"}, Name: di.Name}
		for _, cqi := range di.ContinuousQueries {
			row.Values = append(row.Values, []interface{}{cqi.Name, cqi.Query})
		}
		rows = append(rows, row)
	}
	return &influxql.Result{Series: rows}
}
<commit_msg>Make show stats return error instead of panic<commit_after>package meta

import (
	"fmt"

	"github.com/influxdb/influxdb/influxql"
)

// StatementExecutor translates InfluxQL queries to meta store methods.
type StatementExecutor struct {
	Store interface {
		Nodes() ([]NodeInfo, error)

		Database(name string) (*DatabaseInfo, error)
		Databases() ([]DatabaseInfo, error)
		CreateDatabase(name string) (*DatabaseInfo, error)
		DropDatabase(name string) error

		DefaultRetentionPolicy(database string) (*RetentionPolicyInfo, error)
		CreateRetentionPolicy(database string, rpi *RetentionPolicyInfo) (*RetentionPolicyInfo, error)
		UpdateRetentionPolicy(database, name string, rpu *RetentionPolicyUpdate) error
		SetDefaultRetentionPolicy(database, name string) error
		DropRetentionPolicy(database, name string) error

		Users() ([]UserInfo, error)
		CreateUser(name, password string, admin bool) (*UserInfo, error)
		UpdateUser(name, password string) error
		DropUser(name string) error
		SetPrivilege(username, database string, p influxql.Privilege) error
		UserPrivileges(username string) (map[string]influxql.Privilege, error)

		CreateContinuousQuery(database, name, query string) error
		DropContinuousQuery(database, name string) error
	}
}

// ExecuteStatement executes stmt against the meta store as user.
func (e *StatementExecutor) ExecuteStatement(stmt influxql.Statement) *influxql.Result {
	switch stmt := stmt.(type) {
	case *influxql.CreateDatabaseStatement:
		return e.executeCreateDatabaseStatement(stmt)
	case *influxql.DropDatabaseStatement:
		return e.executeDropDatabaseStatement(stmt)
	case *influxql.ShowDatabasesStatement:
		return e.executeShowDatabasesStatement(stmt)
	case *influxql.ShowGrantsForUserStatement:
		return e.executeShowGrantsForUserStatement(stmt)
	case *influxql.ShowServersStatement:
		return e.executeShowServersStatement(stmt)
	case *influxql.CreateUserStatement:
		return e.executeCreateUserStatement(stmt)
	case *influxql.SetPasswordUserStatement:
		return e.executeSetPasswordUserStatement(stmt)
	case *influxql.DropUserStatement:
		return e.executeDropUserStatement(stmt)
	case *influxql.ShowUsersStatement:
		return e.executeShowUsersStatement(stmt)
	case *influxql.GrantStatement:
		return e.executeGrantStatement(stmt)
	case *influxql.RevokeStatement:
		return e.executeRevokeStatement(stmt)
	case *influxql.CreateRetentionPolicyStatement:
		return e.executeCreateRetentionPolicyStatement(stmt)
	case *influxql.AlterRetentionPolicyStatement:
		return e.executeAlterRetentionPolicyStatement(stmt)
	case *influxql.DropRetentionPolicyStatement:
		return e.executeDropRetentionPolicyStatement(stmt)
	case *influxql.ShowRetentionPoliciesStatement:
		return e.executeShowRetentionPoliciesStatement(stmt)
	case *influxql.CreateContinuousQueryStatement:
		return e.executeCreateContinuousQueryStatement(stmt)
	case *influxql.DropContinuousQueryStatement:
		return e.executeDropContinuousQueryStatement(stmt)
	case *influxql.ShowContinuousQueriesStatement:
		return e.executeShowContinuousQueriesStatement(stmt)
	case *influxql.ShowStatsStatement:
		return e.executeShowStatsStatement(stmt)
	default:
		panic(fmt.Sprintf("unsupported statement type: %T", stmt))
	}
}

func (e *StatementExecutor) executeCreateDatabaseStatement(q *influxql.CreateDatabaseStatement) *influxql.Result {
	_, err := e.Store.CreateDatabase(q.Name)
	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeDropDatabaseStatement(q *influxql.DropDatabaseStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.DropDatabase(q.Name)}
}

func (e *StatementExecutor) executeShowDatabasesStatement(q *influxql.ShowDatabasesStatement) *influxql.Result {
	dis, err := e.Store.Databases()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Name: "databases", Columns: []string{"name"}}
	for _, di := range dis {
		row.Values = append(row.Values, []interface{}{di.Name})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeShowGrantsForUserStatement(q *influxql.ShowGrantsForUserStatement) *influxql.Result {
	priv, err := e.Store.UserPrivileges(q.Name)
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Columns: []string{"database", "privilege"}}
	for d, p := range priv {
		row.Values = append(row.Values, []interface{}{d, p.String()})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeShowServersStatement(q *influxql.ShowServersStatement) *influxql.Result {
	nis, err := e.Store.Nodes()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Columns: []string{"id", "url"}}
	for _, ni := range nis {
		row.Values = append(row.Values, []interface{}{ni.ID, "http://" + ni.Host})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeCreateUserStatement(q *influxql.CreateUserStatement) *influxql.Result {
	admin := false
	if q.Privilege != nil {
		admin = (*q.Privilege == influxql.AllPrivileges)
	}

	_, err := e.Store.CreateUser(q.Name, q.Password, admin)
	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeSetPasswordUserStatement(q *influxql.SetPasswordUserStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.UpdateUser(q.Name, q.Password)}
}

func (e *StatementExecutor) executeDropUserStatement(q *influxql.DropUserStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.DropUser(q.Name)}
}

func (e *StatementExecutor) executeShowUsersStatement(q *influxql.ShowUsersStatement) *influxql.Result {
	uis, err := e.Store.Users()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	row := &influxql.Row{Columns: []string{"user", "admin"}}
	for _, ui := range uis {
		row.Values = append(row.Values, []interface{}{ui.Name, ui.Admin})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeGrantStatement(stmt *influxql.GrantStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.SetPrivilege(stmt.User, stmt.On, stmt.Privilege)}
}

func (e *StatementExecutor) executeRevokeStatement(stmt *influxql.RevokeStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.SetPrivilege(stmt.User, stmt.On, influxql.NoPrivileges)}
}

func (e *StatementExecutor) executeCreateRetentionPolicyStatement(stmt *influxql.CreateRetentionPolicyStatement) *influxql.Result {
	rpi := NewRetentionPolicyInfo(stmt.Name)
	rpi.Duration = stmt.Duration
	rpi.ReplicaN = stmt.Replication

	// Create new retention policy.
	_, err := e.Store.CreateRetentionPolicy(stmt.Database, rpi)
	if err != nil {
		return &influxql.Result{Err: err}
	}

	// If requested, set new policy as the default.
	if stmt.Default {
		err = e.Store.SetDefaultRetentionPolicy(stmt.Database, stmt.Name)
	}

	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeAlterRetentionPolicyStatement(stmt *influxql.AlterRetentionPolicyStatement) *influxql.Result {
	rpu := &RetentionPolicyUpdate{
		Duration: stmt.Duration,
		ReplicaN: stmt.Replication,
	}

	// Update the retention policy.
	err := e.Store.UpdateRetentionPolicy(stmt.Database, stmt.Name, rpu)
	if err != nil {
		return &influxql.Result{Err: err}
	}

	// If requested, set as default retention policy.
	if stmt.Default {
		err = e.Store.SetDefaultRetentionPolicy(stmt.Database, stmt.Name)
	}

	return &influxql.Result{Err: err}
}

func (e *StatementExecutor) executeDropRetentionPolicyStatement(q *influxql.DropRetentionPolicyStatement) *influxql.Result {
	return &influxql.Result{Err: e.Store.DropRetentionPolicy(q.Database, q.Name)}
}

func (e *StatementExecutor) executeShowRetentionPoliciesStatement(q *influxql.ShowRetentionPoliciesStatement) *influxql.Result {
	di, err := e.Store.Database(q.Database)
	if err != nil {
		return &influxql.Result{Err: err}
	} else if di == nil {
		return &influxql.Result{Err: ErrDatabaseNotFound}
	}

	row := &influxql.Row{Columns: []string{"name", "duration", "replicaN", "default"}}
	for _, rpi := range di.RetentionPolicies {
		row.Values = append(row.Values, []interface{}{rpi.Name, rpi.Duration.String(), rpi.ReplicaN, di.DefaultRetentionPolicy == rpi.Name})
	}
	return &influxql.Result{Series: []*influxql.Row{row}}
}

func (e *StatementExecutor) executeCreateContinuousQueryStatement(q *influxql.CreateContinuousQueryStatement) *influxql.Result {
	return &influxql.Result{
		Err: e.Store.CreateContinuousQuery(q.Database, q.Name, q.String()),
	}
}

func (e *StatementExecutor) executeDropContinuousQueryStatement(q *influxql.DropContinuousQueryStatement) *influxql.Result {
	return &influxql.Result{
		Err: e.Store.DropContinuousQuery(q.Database, q.Name),
	}
}

func (e *StatementExecutor) executeShowContinuousQueriesStatement(stmt *influxql.ShowContinuousQueriesStatement) *influxql.Result {
	dis, err := e.Store.Databases()
	if err != nil {
		return &influxql.Result{Err: err}
	}

	rows := []*influxql.Row{}
	for _, di := range dis {
		row := &influxql.Row{Columns: []string{"name", "query"}, Name: di.Name}
		for _, cqi := range di.ContinuousQueries {
			row.Values = append(row.Values, []interface{}{cqi.Name, cqi.Query})
		}
		rows = append(rows, row)
	}
	return &influxql.Result{Series: rows}
}

func (e *StatementExecutor) executeShowStatsStatement(stmt *influxql.ShowStatsStatement) *influxql.Result {
	return &influxql.Result{Err: fmt.Errorf("SHOW STATS is not implemented yet")}
}
<|endoftext|>