<commit_before>package main

import (
    "bufio"
    "bytes"
    "encoding/json"
    "errors"
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "metrics"
    "net"
    "net/http"
    "net/url"
    "os"
    "os/signal"
    "strconv"
    "strings"
    "sync"
    "syscall"
    "time"
    "tokenbucket"
)

const _VERSION string = "0.3.0"
var _logger *log.Logger = log.New(os.Stdout, "", 0)

type Route struct {
    Bandwidth int64 `json:"bandwidth"`    // Bits per second
    Buffersize uint64 `json:"buffersize"` // Bytes
    Inspect bool `json:"inspect"`         // True = proxy, false = reverse proxy
    Guide string `json:"guide"`           // HTTP(S) probe to query a load balancer for backend addresses, response field name containing IP address, and static destination port
    Src string `json:"src"`
    Dst []string `json:"dst"`
}

type Itinerary struct {
    Map map[string]Route
}

func sigHandler(ch *chan os.Signal) {
    sig := <-*ch
    fmt.Println("Captured sig", sig)
    os.Exit(3)
}

func main() {
    sigs := make(chan os.Signal, 1)
    signal.Notify(sigs,
                  syscall.SIGHUP,
                  syscall.SIGINT,
                  syscall.SIGQUIT,
                  syscall.SIGABRT,
                  syscall.SIGKILL,
                  syscall.SIGSEGV,
                  syscall.SIGTERM)

    go sigHandler(&sigs)

    log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

    itineraryFile := flag.String("itinerary", "itinerary.json", "file containing source and destination routes")
    flag.Usage = func() {
        fmt.Fprintf(os.Stderr, "version %s\n", _VERSION)
        fmt.Fprintln(os.Stderr, "usage:")
        flag.PrintDefaults()
    }
    flag.Parse()

    itinerary := loadItinerary(itineraryFile)

    var wg sync.WaitGroup
    wg.Add(len(itinerary.Map))

    for _, m := range itinerary.Map {
        go intercept(&wg, m)
    }

    wg.Wait()
}

func loadItinerary(fileName *string) Itinerary {

    file, _ := os.Open(*fileName)
    decoder := json.NewDecoder(file)
    itinerary := Itinerary{}
    err := decoder.Decode(&itinerary)
    if err != nil {
        log.Println(err.Error())
    }

    log.Println("Asking guides for directions")

    var wg sync.WaitGroup
    wg.Add(len(itinerary.Map))

    for i, m := range itinerary.Map {
        if len(m.Guide) > 0 {
            var uri, key string
            var port int
            args, _ := fmt.Sscanf(m.Guide, "%s %s %d", &uri, &key, &port)

            if args >= 2 {
                findDestinations(&wg, &m, uri, key, port)
                itinerary.Map[i] = m
            } else {
                wg.Done()
            }
        } else {
            wg.Done()
        }
    }

    wg.Wait()

    log.Printf("Finished asking guides for directions")

    return itinerary
}

func findDestinations(wg *sync.WaitGroup, route *Route, guide, key string, port int) {
    const TimeToLive = 10
    ask := TimeToLive
    count := 0

    for ask > 0 {
        if dst, err := askGuide(guide, key); err == nil {
            found := false
            for _, v := range route.Dst {
                if v == (dst + ":" + strconv.Itoa(port)) {
                    found = true
                    ask = ask - 1
                    break
                }
            }

            if !found {
                route.Dst = append(route.Dst, dst + ":" + strconv.Itoa(port))
                count = count + 1
                ask = TimeToLive
                log.Println(guide + ": found " + dst)
            } else {
                time.Sleep(250 * time.Millisecond)
            }
        } else {
            ask = ask - 1
        }
    }

    log.Printf("%s: found %d destinations\n", guide, count)
    wg.Done()
}

func askGuide(guide, key string) (string, error) {
    var dir string = ""
    tp := &http.Transport{ DisableKeepAlives: true, }
    client := &http.Client{ Transport: tp, Timeout: 4 * time.Second }
    r, err := client.Get(guide)

    if err == nil {
        if r.StatusCode == 200 {
            var dat map[string]interface{}
            body, _ := ioutil.ReadAll(r.Body)
            json.Unmarshal([]byte(string(body)), &dat)

            if val, ok := dat[key]; ok {
                if dst, ok := val.(string); ok {
                    dir = dst
                }
            }
        }

        r.Body.Close()
    }

    if len(dir) == 0 {
        err = errors.New("'" + key + "' key does not exist")
    }

    return dir, err
}

func intercept(wg *sync.WaitGroup, route Route) {
    listener, err := net.Listen("tcp", route.Src)
    if err == nil {
        log.Println("Listening on", route.Src)
        i := 0
        count := 0

        for {
            if con, err := listener.Accept(); err == nil {
                count = count + 1
                m := metrics.New(_logger, 1000000000, strconv.Itoa(count) + ",tx")
                if route.Inspect { // Proxy mode
                    go findHttpRoute(count, con, route.Bandwidth / 8, route.Buffersize, m)
                } else if len(route.Dst) > 0 { // Load balancer mode
                    go findRoute(count, con, route.Dst[i], route.Bandwidth / 8, route.Buffersize, m)
                    i = (i + 1) % len(route.Dst) // Round-robin for now
                }
            } else {
                log.Println(err.Error())
            }
        }
        listener.Close()
    } else {
        log.Println(err.Error())
    }

    wg.Done()
}

func parseHttpRequestUri(header *http.Request) string {
    var uri string
    if header.Method == "CONNECT" {
        uri = header.RequestURI
    } else {
        url, err := url.Parse(header.RequestURI)

        if err == nil {
            if strings.ContainsAny(url.Host, ":") {
                uri = url.Host
            } else {
                switch url.Scheme {
                    case "http":
                        uri = url.Host + ":80"
                    case "https":
                        uri = url.Host + ":443"
                    default:
                        uri = url.Host
                }
            }
        } else {
            log.Println(err.Error())
        }
    }

    return uri
}

func findHttpRoute(id int, src net.Conn, bandwidth int64, bufferSize uint64, m *metrics.Metrics) {
    buf := make([]byte, 65536)
    size, err := src.Read(buf) // Assume one read will include entire HTTP header

    if err == nil {
        reader := bufio.NewReader(strings.NewReader(string(buf[0:size])))
        header, _ := http.ReadRequest(reader)

        dstAddr := parseHttpRequestUri(header)
        dst, err := net.Dial("tcp", dstAddr)

        if err == nil {
            if header.Method == "CONNECT" {
                body := ""
                rsp := &http.Response {
                    Status:        "200 OK",
                    StatusCode:    200,
                    Proto:         "HTTP/1.1",
                    ProtoMajor:    1,
                    ProtoMinor:    1,
                    Body:          ioutil.NopCloser(bytes.NewBufferString(body)),
                    ContentLength: int64(len(body)),
                    Request:       nil,
                    Header:        make(http.Header, 0),
                }
                rspBuf := bytes.NewBuffer(nil)
                rsp.Write(rspBuf)
                src.Write(rspBuf.Bytes())
            } else {
                dst.Write(buf[0:size])
            }
            startDetour(id, src, dst, bandwidth, bufferSize, m)
        } else {
            src.Close()
            log.Println(err.Error())
        }
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func findRoute(id int, src net.Conn, dstAddr string, bandwidth int64, bufferSize uint64, m *metrics.Metrics) {
    dst, err := net.Dial("tcp", dstAddr)

    if err == nil {
        startDetour(id, src, dst, bandwidth, bufferSize, m)
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func startDetour(id int, src net.Conn, dst net.Conn, bandwidth int64, bufferSize uint64, m *metrics.Metrics) {
    log.Println("Opening route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())

    var wg sync.WaitGroup
    wg.Add(1)
    go reroute(&wg, src, dst, bandwidth, bufferSize, m)
    reroute(nil, dst, src, bandwidth, bufferSize, metrics.New(_logger, 1000000000, strconv.Itoa(id) + ",rx"))
    wg.Wait()

    log.Println("Closing route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())
}

func reroute(wg *sync.WaitGroup, src net.Conn, dst net.Conn, bandwidth int64, bufferSize uint64, m *metrics.Metrics) {
    tb := tokenbucket.New(uint64(bandwidth), 10 * uint64(bandwidth))
    buf := make([]byte, bufferSize)
    defer src.Close()
    defer dst.Close()

    for {
        bytes := tb.Remove(bufferSize)
        if bytes < bufferSize {
            tb.Return(bytes)
            time.Sleep(1 * time.Millisecond)
            continue
        }

        size, err := src.Read(buf)
        if err == nil {
            m.Add(int64(size))
            dst.Write(buf[0:size])
            if size < int(bufferSize) {
                tb.Return(bufferSize - uint64(size))
            }
        } else {
            break
        }
    }

    m.Dump()

    if wg != nil {
        wg.Done()
    }
}
<commit_msg>Track HTTP requests and responses<commit_after>package main

import (
    "bufio"
    "bytes"
    "encoding/json"
    "errors"
    "flag"
    "fmt"
    "io/ioutil"
    "log"
    "metrics"
    "net"
    "net/http"
    "net/url"
    "os"
    "os/signal"
    "strconv"
    "strings"
    "sync"
    "syscall"
    "time"
    "tokenbucket"
)

const _VERSION string = "0.3.0"
var _logger *log.Logger = log.New(os.Stdout, "", 0)

type Route struct {
    Bandwidth int64 `json:"bandwidth"`    // Bits per second
    Buffersize uint64 `json:"buffersize"` // Bytes
    Inspect bool `json:"inspect"`         // True = proxy, false = reverse proxy
    Guide string `json:"guide"`           // HTTP(S) probe to query a load balancer for backend addresses, response field name containing IP address, and static destination port
    Src string `json:"src"`
    Dst []string `json:"dst"`
}

type Itinerary struct {
    Map map[string]Route
}

func sigHandler(ch *chan os.Signal) {
    sig := <-*ch
    fmt.Println("Captured sig", sig)
    os.Exit(3)
}

func main() {
    sigs := make(chan os.Signal, 1)
    signal.Notify(sigs,
                  syscall.SIGHUP,
                  syscall.SIGINT,
                  syscall.SIGQUIT,
                  syscall.SIGABRT,
                  syscall.SIGKILL,
                  syscall.SIGSEGV,
                  syscall.SIGTERM)

    go sigHandler(&sigs)

    log.SetFlags(log.Ldate | log.Ltime | log.Lshortfile)

    itineraryFile := flag.String("itinerary", "itinerary.json", "file containing source and destination routes")
    flag.Usage = func() {
        fmt.Fprintf(os.Stderr, "version %s\n", _VERSION)
        fmt.Fprintln(os.Stderr, "usage:")
        flag.PrintDefaults()
    }
    flag.Parse()

    itinerary := loadItinerary(itineraryFile)

    var wg sync.WaitGroup
    wg.Add(len(itinerary.Map))

    for _, m := range itinerary.Map {
        go intercept(&wg, m)
    }

    wg.Wait()
}

func loadItinerary(fileName *string) Itinerary {

    file, _ := os.Open(*fileName)
    decoder := json.NewDecoder(file)
    itinerary := Itinerary{}
    err := decoder.Decode(&itinerary)
    if err != nil {
        log.Println(err.Error())
    }

    log.Println("Asking guides for directions")

    var wg sync.WaitGroup
    wg.Add(len(itinerary.Map))

    for i, m := range itinerary.Map {
        if len(m.Guide) > 0 {
            var uri, key string
            var port int
            args, _ := fmt.Sscanf(m.Guide, "%s %s %d", &uri, &key, &port)

            if args >= 2 {
                findDestinations(&wg, &m, uri, key, port)
                itinerary.Map[i] = m
            } else {
                wg.Done()
            }
        } else {
            wg.Done()
        }
    }

    wg.Wait()

    log.Printf("Finished asking guides for directions")

    return itinerary
}

func findDestinations(wg *sync.WaitGroup, route *Route, guide, key string, port int) {
    const TimeToLive = 10
    ask := TimeToLive
    count := 0

    for ask > 0 {
        if dst, err := askGuide(guide, key); err == nil {
            found := false
            for _, v := range route.Dst {
                if v == (dst + ":" + strconv.Itoa(port)) {
                    found = true
                    ask = ask - 1
                    break
                }
            }

            if !found {
                route.Dst = append(route.Dst, dst + ":" + strconv.Itoa(port))
                count = count + 1
                ask = TimeToLive
                log.Println(guide + ": found " + dst)
            } else {
                time.Sleep(250 * time.Millisecond)
            }
        } else {
            ask = ask - 1
        }
    }

    log.Printf("%s: found %d destinations\n", guide, count)
    wg.Done()
}

func askGuide(guide, key string) (string, error) {
    var dir string = ""
    tp := &http.Transport{ DisableKeepAlives: true, }
    client := &http.Client{ Transport: tp, Timeout: 4 * time.Second }
    r, err := client.Get(guide)

    if err == nil {
        if r.StatusCode == 200 {
            var dat map[string]interface{}
            body, _ := ioutil.ReadAll(r.Body)
            json.Unmarshal([]byte(string(body)), &dat)

            if val, ok := dat[key]; ok {
                if dst, ok := val.(string); ok {
                    dir = dst
                }
            }
        }

        r.Body.Close()
    }

    if len(dir) == 0 {
        err = errors.New("'" + key + "' key does not exist")
    }

    return dir, err
}

func intercept(wg *sync.WaitGroup, route Route) {
    listener, err := net.Listen("tcp", route.Src)
    if err == nil {
        log.Println("Listening on", route.Src)
        i := 0
        count := 0

        for {
            if con, err := listener.Accept(); err == nil {
                count = count + 1
                m := metrics.New(_logger, 1000000000, strconv.Itoa(count) + ",tx")
                if route.Inspect { // Proxy mode
                    go findHttpRoute(count, con, &route, m)
                } else if len(route.Dst) > 0 { // Load balancer mode
                    go findRoute(count, con, route.Dst[i], &route, m)
                    i = (i + 1) % len(route.Dst) // Round-robin for now
                }
            } else {
                log.Println(err.Error())
            }
        }
        listener.Close()
    } else {
        log.Println(err.Error())
    }

    wg.Done()
}

func parseHttpRequestUri(header *http.Request) string {
    var uri string
    if header.Method == "CONNECT" {
        uri = header.RequestURI
    } else {
        url, err := url.Parse(header.RequestURI)

        if err == nil {
            if strings.ContainsAny(url.Host, ":") {
                uri = url.Host
            } else {
                switch url.Scheme {
                    case "http":
                        uri = url.Host + ":80"
                    case "https":
                        uri = url.Host + ":443"
                    default:
                        uri = url.Host
                }
            }
        } else {
            log.Println(err.Error())
        }
    }

    return uri
}

func findHttpRoute(id int, src net.Conn, route *Route, m *metrics.Metrics) {
    buf := make([]byte, 65536)
    size, err := src.Read(buf) // Assume one read will include entire HTTP header

    if err == nil {
        reader := bufio.NewReader(strings.NewReader(string(buf[0:size])))
        request, _ := http.ReadRequest(reader)

        dstAddr := parseHttpRequestUri(request)
        dst, err := net.Dial("tcp", dstAddr)

        log.Println(request)
        log.Printf("[REQUEST] Content Length: %d\n", request.ContentLength)

        if err == nil {
            if request.Method == "CONNECT" {
                body := ""
                rsp := &http.Response {
                    Status:        "200 OK",
                    StatusCode:    200,
                    Proto:         "HTTP/1.1",
                    ProtoMajor:    1,
                    ProtoMinor:    1,
                    Body:          ioutil.NopCloser(bytes.NewBufferString(body)),
                    ContentLength: int64(len(body)),
                    Request:       nil,
                    Header:        make(http.Header, 0),
                }
                rspBuf := bytes.NewBuffer(nil)
                rsp.Write(rspBuf)
                src.Write(rspBuf.Bytes())
            } else {
                dst.Write(buf[0:size])
            }
            startDetour(id, src, dst, route, m)
        } else {
            src.Close()
            log.Println(err.Error())
        }
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func findRoute(id int, src net.Conn, dstAddr string, route *Route, m *metrics.Metrics) {
    dst, err := net.Dial("tcp", dstAddr)

    if err == nil {
        startDetour(id, src, dst, route, m)
    } else {
        src.Close()
        log.Println(err.Error())
    }
}

func startDetour(id int, src net.Conn, dst net.Conn, route *Route, m *metrics.Metrics) {
    log.Println("Opening route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())

    var wg sync.WaitGroup
    wg.Add(1)
    go reroute(&wg, src, dst, route, m)
    reroute(nil, dst, src, route, metrics.New(_logger, 1000000000, strconv.Itoa(id) + ",rx"))
    wg.Wait()

    log.Println("Closing route", id, ":", src.RemoteAddr().String(), "to", dst.RemoteAddr().String())
}

func reroute(wg *sync.WaitGroup, src net.Conn, dst net.Conn, route *Route, m *metrics.Metrics) {
    bandwidth := route.Bandwidth / 8
    bufferSize := route.Buffersize
    var contentLength int64 = 0
    direction := ""

    tb := tokenbucket.New(uint64(bandwidth), 10 * uint64(bandwidth))
    buf := make([]byte, bufferSize)
    defer src.Close()
    defer dst.Close()

    for {
        bytes := tb.Remove(bufferSize)
        if bytes < bufferSize {
            tb.Return(bytes)
            time.Sleep(1 * time.Millisecond)
            continue
        }

        size, err := src.Read(buf)

        if err == nil {
            if route.Inspect == true  {
                if contentLength == 0 {
                    if r := getHttpRequest(string(buf[0:size])); r != nil {
                        contentLength = r.ContentLength
                        direction = "[REQUEST]"
                        log.Println(r)
                        log.Printf("%s Content Length: %d\n", direction, contentLength)
                    } else if r := getHttpResponse(string(buf[0:size])); r != nil {
                        contentLength = r.ContentLength
                        direction = "[RESPONSE]"
                        log.Println(r)
                        log.Printf("%s Content Length: %d\n", direction, contentLength)
                    }
                } else {
                    contentLength = contentLength - int64(size)
                    //log.Println("%s Content Length: %d (bytes read %d)\n", direction, contentLength, size)
                }
            }

            //m.Add(int64(size))
            dst.Write(buf[0:size])
            if size < int(bufferSize) {
                tb.Return(bufferSize - uint64(size))
            }
        } else {
            break
        }
    }

    //m.Dump()

    if wg != nil {
        wg.Done()
    }
}

func getHttpRequest(request string) *http.Request {
    var req *http.Request = nil

    reader := bufio.NewReader(strings.NewReader(request))

    if r, err := http.ReadRequest(reader); err == nil {
        req = r
    }

    return req
}

func getHttpResponse(response string) *http.Response {
    var resp *http.Response = nil

    reader := bufio.NewReader(strings.NewReader(response))

    if r, err := http.ReadResponse(reader, nil); err == nil {
        resp = r
    }

    return resp
}
<|endoftext|>


<commit_before>package wemo

import (
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"regexp"
	"strconv"
)

type Device struct {
	Host string
}

type DeviceInfo struct {
	DeviceType      string `xml:"deviceType" json:"device-type"`
	FriendlyName    string `xml:"friendlyName" json:"friendly-name"`
	MacAddress      string `xml:"macAddress" json:"mac-address"`
	FirmwareVersion string `xml:"firmwareVersion" json:"firmware-version"`
	SerialNumber    string `xml:"serialNumber" json:"serial-number"`
}

type BelkinResponse struct {
	Device DeviceInfo `xml:"device"`
}

func (self *Device) FetchDeviceInfo() (*DeviceInfo, error) {
	uri := fmt.Sprintf("http://%s/setup.xml", self.Host)
	response, err := client.Get(uri)
	if err != nil {
		return nil, err
	}

	defer response.Body.Close()
	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	resp := new(BelkinResponse)
	err = xml.Unmarshal(data, resp)
	if err != nil {
		return nil, err
	}

	fmt.Printf("%+v\n", resp.Device)

	return &resp.Device, nil
}

func (self *Device) FetchBinaryState() int {
	message := newGetBinaryStateMessage()
	response, err := post(self.Host, "GetBinaryState", message)
	if err != nil {
		log.Printf("unable to fetch BinaryState => %s\n", err)
		return -1
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		log.Printf("GetBinaryState returned status code => %d\n", response.StatusCode)
		return -1
	}

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Printf("unable to read data => %s\n", err)
		return -1
	}

	re := regexp.MustCompile(`.*<BinaryState>(\d+)</BinaryState>.*`)
	matches := re.FindStringSubmatch(string(data))
	if len(matches) != 2 {
		log.Printf("unable to find BinaryState response in message => %s\n", string(data))
		return -1
	}

	result, _ := strconv.Atoi(matches[1])
	return result
}

func (self *Device) Off() {
	self.changeState(false)
}

func (self *Device) On() {
	self.changeState(true)
}

func (self *Device) Toggle() {
	if binaryState := self.FetchBinaryState(); binaryState == 0 {
		self.On()
	} else {
		self.Off()
	}
}

func (self *Device) changeState(newState bool) error {
	fmt.Printf("changeState(%v)\n", newState)
	message := newSetBinaryStateMessage(newState)
	response, err := post(self.Host, "SetBinaryState", message)
	if err != nil {
		log.Println("unable to SetBinaryState")
		return err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		data, err := ioutil.ReadAll(response.Body)
		if err != nil {
			log.Println("couldn't read body from message => " + err.Error())
			return err
		}
		content := string(data)

		gripe := fmt.Sprintf("changeState(%v) => %s", newState, content)
		log.Println(gripe)
		return errors.New(gripe)
	}

	return nil
}
<commit_msg>- renamed FetchBinaryState to GetBinaryState<commit_after>package wemo

import (
	"encoding/xml"
	"errors"
	"fmt"
	"io/ioutil"
	"log"
	"regexp"
	"strconv"
)

type Device struct {
	Host string
}

type DeviceInfo struct {
	DeviceType      string `xml:"deviceType" json:"device-type"`
	FriendlyName    string `xml:"friendlyName" json:"friendly-name"`
	MacAddress      string `xml:"macAddress" json:"mac-address"`
	FirmwareVersion string `xml:"firmwareVersion" json:"firmware-version"`
	SerialNumber    string `xml:"serialNumber" json:"serial-number"`
}

type BelkinResponse struct {
	Device DeviceInfo `xml:"device"`
}

func (self *Device) FetchDeviceInfo() (*DeviceInfo, error) {
	uri := fmt.Sprintf("http://%s/setup.xml", self.Host)
	response, err := client.Get(uri)
	if err != nil {
		return nil, err
	}

	defer response.Body.Close()
	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		return nil, err
	}

	resp := new(BelkinResponse)
	err = xml.Unmarshal(data, resp)
	if err != nil {
		return nil, err
	}

	fmt.Printf("%+v\n", resp.Device)

	return &resp.Device, nil
}

func (self *Device) GetBinaryState() int {
	message := newGetBinaryStateMessage()
	response, err := post(self.Host, "GetBinaryState", message)
	if err != nil {
		log.Printf("unable to fetch BinaryState => %s\n", err)
		return -1
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		log.Printf("GetBinaryState returned status code => %d\n", response.StatusCode)
		return -1
	}

	data, err := ioutil.ReadAll(response.Body)
	if err != nil {
		log.Printf("unable to read data => %s\n", err)
		return -1
	}

	re := regexp.MustCompile(`.*<BinaryState>(\d+)</BinaryState>.*`)
	matches := re.FindStringSubmatch(string(data))
	if len(matches) != 2 {
		log.Printf("unable to find BinaryState response in message => %s\n", string(data))
		return -1
	}

	result, _ := strconv.Atoi(matches[1])
	return result
}

func (self *Device) Off() {
	self.changeState(false)
}

func (self *Device) On() {
	self.changeState(true)
}

func (self *Device) Toggle() {
	if binaryState := self.GetBinaryState(); binaryState == 0 {
		self.On()
	} else {
		self.Off()
	}
}

func (self *Device) changeState(newState bool) error {
	fmt.Printf("changeState(%v)\n", newState)
	message := newSetBinaryStateMessage(newState)
	response, err := post(self.Host, "SetBinaryState", message)
	if err != nil {
		log.Println("unable to SetBinaryState")
		return err
	}
	defer response.Body.Close()

	if response.StatusCode != 200 {
		data, err := ioutil.ReadAll(response.Body)
		if err != nil {
			log.Println("couldn't read body from message => " + err.Error())
			return err
		}
		content := string(data)

		gripe := fmt.Sprintf("changeState(%v) => %s", newState, content)
		log.Println(gripe)
		return errors.New(gripe)
	}

	return nil
}
<|endoftext|>


<commit_before>package qshell

import (
	"bufio"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"github.com/qiniu/api.v6/auth/digest"
	"github.com/qiniu/api.v6/conf"
	fio "github.com/qiniu/api.v6/io"
	rio "github.com/qiniu/api.v6/resumable/io"
	"github.com/qiniu/api.v6/rs"
	"github.com/qiniu/log"
	"github.com/qiniu/rpc"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
)

/*
Config file like:

{
	"up_host"		:	"http://upload.qiniu.com",
	"src_dir" 		:	"/Users/jemy/Photos",
	"access_key" 	:	"<Your AccessKey>",
	"secret_key"	:	"<Your SecretKey>",
	"bucket"		:	"test-bucket",
	"ignore_dir"	:	false,
	"key_prefix"	:	"2014/12/01/",
	"overwrite"		:	false,
	"check_exists"	:	true
}

or without up_host and key_prefix and ignore_dir and check_exists

{
	"src_dir" 		:	"/Users/jemy/Photos",
	"access_key" 	:	"<Your AccessKey>",
	"secret_key"	:	"<Your SecretKey>",
	"bucket"		:	"test-bucket",
}
*/

const (
	PUT_THRESHOLD           int64 = 10 * 1 << 20
	MIN_UPLOAD_THREAD_COUNT int64 = 1
	MAX_UPLOAD_THREAD_COUNT int64 = 100
)

type UploadConfig struct {
	SrcDir       string `json:"src_dir"`
	AccessKey    string `json:"access_key"`
	SecretKey    string `json:"secret_key"`
	Bucket       string `json:"bucket"`
	UpHost       string `json:"up_host,omitempty"`
	KeyPrefix    string `json:"key_prefix,omitempty"`
	IgnoreDir    bool   `json:"ignore_dir,omitempty"`
	Overwrite    bool   `json:"overwrite,omitempty"`
	CheckExists  bool   `json:"check_exists,omitempty"`
	SkipPrefixes string `json:"skip_prefixes,omitempty"`
	SkipSuffixes string `json:"skip_suffixes,omitempty"`
}

var upSettings = rio.Settings{
	ChunkSize: 1 * 1024 * 1024,
	TryTimes:  5,
}

func QiniuUpload(threadCount int, uploadConfigFile string) {
	fp, err := os.Open(uploadConfigFile)
	if err != nil {
		log.Error(fmt.Sprintf("Open upload config file `%s' error due to `%s'", uploadConfigFile, err))
		return
	}
	defer fp.Close()
	configData, err := ioutil.ReadAll(fp)
	if err != nil {
		log.Error(fmt.Sprintf("Read upload config file `%s' error due to `%s'", uploadConfigFile, err))
		return
	}
	var uploadConfig UploadConfig
	err = json.Unmarshal(configData, &uploadConfig)
	if err != nil {
		log.Error(fmt.Sprintf("Parse upload config file `%s' errror due to `%s'", uploadConfigFile, err))
		return
	}
	if _, err := os.Stat(uploadConfig.SrcDir); err != nil {
		log.Error("Upload config error for parameter `SrcDir`,", err)
		return
	}
	dirCache := DirCache{}
	currentUser, err := user.Current()
	if err != nil {
		log.Error("Failed to get current user", err)
		return
	}

	pathSep := string(os.PathSeparator)
	//create job id
	md5Hasher := md5.New()
	md5Hasher.Write([]byte(strings.TrimSuffix(uploadConfig.SrcDir, pathSep) + ":" + uploadConfig.Bucket))
	jobId := fmt.Sprintf("%x", md5Hasher.Sum(nil))

	//local storage path
	storePath := filepath.Join(currentUser.HomeDir, ".qshell", "qupload", jobId)
	err = os.MkdirAll(storePath, 0775)
	if err != nil {
		log.Error(fmt.Sprintf("Failed to mkdir `%s' due to `%s'", storePath, err))
		return
	}

	//cache file
	cacheFileName := filepath.Join(storePath, jobId+".cache")
	//leveldb folder
	leveldbFileName := filepath.Join(storePath, jobId+".ldb")

	totalFileCount := dirCache.Cache(uploadConfig.SrcDir, cacheFileName)
	ldb, err := leveldb.OpenFile(leveldbFileName, nil)
	if err != nil {
		log.Error(fmt.Sprintf("Open leveldb `%s' failed due to `%s'", leveldbFileName, err))
		return
	}
	defer ldb.Close()
	//sync
	ufp, err := os.Open(cacheFileName)
	if err != nil {
		log.Error(fmt.Sprintf("Open cache file `%s' failed due to `%s'", cacheFileName, err))
		return
	}
	defer ufp.Close()
	bScanner := bufio.NewScanner(ufp)
	bScanner.Split(bufio.ScanLines)

	var currentFileCount int64 = 0
	var successFileCount int64 = 0
	var failureFileCount int64 = 0
	var skippedFileCount int64 = 0

	ldbWOpt := opt.WriteOptions{
		Sync: true,
	}

	upWorkGroup := sync.WaitGroup{}
	upCounter := 0
	threadThreshold := threadCount + 1

	//use host if not empty
	if uploadConfig.UpHost != "" {
		conf.UP_HOST = uploadConfig.UpHost
	}
	//set resume upload settings
	rio.SetSettings(&upSettings)
	mac := digest.Mac{uploadConfig.AccessKey, []byte(uploadConfig.SecretKey)}

	for bScanner.Scan() {
		line := strings.TrimSpace(bScanner.Text())
		items := strings.Split(line, "\t")
		if len(items) != 3 {
			log.Error(fmt.Sprintf("Invalid cache line `%s'", line))
			continue
		}

		localFname := items[0]
		currentFileCount += 1

		skip := false
		//check skip local file or folder
		if uploadConfig.SkipPrefixes != "" {
			//unpack skip prefix
			skipPrefixes := strings.Split(uploadConfig.SkipPrefixes, ",")
			for _, prefix := range skipPrefixes {
				if strings.HasPrefix(localFname, strings.TrimSpace(prefix)) {
					log.Debug(fmt.Sprintf("Skip by prefix '%s' for local file %s",
						strings.TrimSpace(prefix), localFname))
					skip = true
					skippedFileCount += 1
					break
				}
			}

			if skip {

				continue
			}
		}

		if uploadConfig.SkipSuffixes != "" {
			skipSuffixes := strings.Split(uploadConfig.SkipSuffixes, ",")
			for _, suffix := range skipSuffixes {
				if strings.HasSuffix(localFname, strings.TrimSpace(suffix)) {
					log.Debug(fmt.Sprintf("Skip by suffix '%s' for local file %s",
						strings.TrimSpace(suffix), localFname))
					skip = true
					skippedFileCount += 1
					break
				}
			}

			if skip {
				continue
			}
		}

		//pack the upload file key
		localFlmd, _ := strconv.Atoi(items[2])
		uploadFileKey := localFname

		if uploadConfig.IgnoreDir {
			if i := strings.LastIndex(uploadFileKey, pathSep); i != -1 {
				uploadFileKey = uploadFileKey[i+1:]
			}
		}
		if uploadConfig.KeyPrefix != "" {
			uploadFileKey = strings.Join([]string{uploadConfig.KeyPrefix, uploadFileKey}, "")
		}
		//convert \ to / under windows
		if runtime.GOOS == "windows" {
			uploadFileKey = strings.Replace(uploadFileKey, "\\", "/", -1)
		}

		localFilePath := filepath.Join(uploadConfig.SrcDir, localFname)
		fstat, err := os.Stat(localFilePath)
		if err != nil {
			log.Error(fmt.Sprintf("Error stat local file `%s' due to `%s'", localFilePath, err))
			return
		}

		fsize := fstat.Size()
		ldbKey := fmt.Sprintf("%s => %s", localFilePath, uploadFileKey)

		log.Info(fmt.Sprintf("Uploading %s (%d/%d, %.1f%%) ...", ldbKey, currentFileCount, totalFileCount,
			float32(currentFileCount)*100/float32(totalFileCount)))

		rsClient := rs.New(&mac)
		log.Debug(fmt.Sprintf("Checking %s ...", ldbKey))

		//check exists
		if uploadConfig.CheckExists {
			rsEntry, checkErr := rsClient.Stat(nil, uploadConfig.Bucket, uploadFileKey)
			if checkErr == nil {
				//compare hash
				localEtag, cErr := GetEtag(localFilePath)
				if cErr != nil {
					atomic.AddInt64(&failureFileCount, 1)
					log.Error("Calc local file hash failed,", cErr)
					continue
				}
				if rsEntry.Hash == localEtag {
					atomic.AddInt64(&skippedFileCount, 1)
					log.Debug(fmt.Sprintf("File %s already exists in bucket, ignore this upload", uploadFileKey))
					continue
				}
			} else {
				if _, ok := checkErr.(*rpc.ErrorInfo); !ok {
					//not logic error, should be network error
					atomic.AddInt64(&failureFileCount, 1)
					continue
				}
			}
		} else {
			//check leveldb
			ldbFlmd, err := ldb.Get([]byte(ldbKey), nil)
			flmd, _ := strconv.Atoi(string(ldbFlmd))
			//not exist, return ErrNotFound
			//check last modified

			if err == nil && localFlmd == flmd {
				log.Debug("Skip by local log for file", localFname)
				atomic.AddInt64(&skippedFileCount, 1)
				continue
			}
		}

		//worker
		upCounter += 1
		if upCounter%threadThreshold == 0 {
			upWorkGroup.Wait()
		}
		upWorkGroup.Add(1)

		//start to upload
		go func() {
			defer upWorkGroup.Done()

			policy := rs.PutPolicy{}
			policy.Scope = uploadConfig.Bucket
			if uploadConfig.Overwrite {
				policy.Scope = uploadConfig.Bucket + ":" + uploadFileKey
				policy.InsertOnly = 0
			}
			policy.Expires = 24 * 3600
			uptoken := policy.Token(&mac)
			if fsize > PUT_THRESHOLD {
				putRet := rio.PutRet{}
				err := rio.PutFile(nil, &putRet, uptoken, uploadFileKey, localFilePath, nil)
				if err != nil {
					atomic.AddInt64(&failureFileCount, 1)
					if pErr, ok := err.(*rpc.ErrorInfo); ok {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, pErr.Err))
					} else {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, err))
					}
				} else {
					atomic.AddInt64(&successFileCount, 1)
					perr := ldb.Put([]byte(ldbKey), []byte("Y"), &ldbWOpt)
					if perr != nil {
						log.Error(fmt.Sprintf("Put key `%s' into leveldb error due to `%s'", ldbKey, perr))
					}
				}
			} else {
				putRet := fio.PutRet{}
				err := fio.PutFile(nil, &putRet, uptoken, uploadFileKey, localFilePath, nil)
				if err != nil {
					atomic.AddInt64(&failureFileCount, 1)
					if pErr, ok := err.(*rpc.ErrorInfo); ok {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, pErr.Err))
					} else {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, err))
					}
				} else {
					atomic.AddInt64(&successFileCount, 1)
					perr := ldb.Put([]byte(ldbKey), []byte(strconv.Itoa(localFlmd)), &ldbWOpt)
					if perr != nil {
						log.Error(fmt.Sprintf("Put key `%s' into leveldb error due to `%s'", ldbKey, perr))
					}
				}
			}
		}()

	}
	upWorkGroup.Wait()

	log.Info("-------Upload Done-------")
	log.Info("Total:\t", currentFileCount)
	log.Info("Success:\t", successFileCount)
	log.Info("Failure:\t", failureFileCount)
	log.Info("Skipped:\t", skippedFileCount)
	log.Info("-------------------------")

}
<commit_msg>remove unused log<commit_after>package qshell

import (
	"bufio"
	"crypto/md5"
	"encoding/json"
	"fmt"
	"github.com/qiniu/api.v6/auth/digest"
	"github.com/qiniu/api.v6/conf"
	fio "github.com/qiniu/api.v6/io"
	rio "github.com/qiniu/api.v6/resumable/io"
	"github.com/qiniu/api.v6/rs"
	"github.com/qiniu/log"
	"github.com/qiniu/rpc"
	"github.com/syndtr/goleveldb/leveldb"
	"github.com/syndtr/goleveldb/leveldb/opt"
	"io/ioutil"
	"os"
	"os/user"
	"path/filepath"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"sync/atomic"
)

/*
Config file like:

{
	"up_host"		:	"http://upload.qiniu.com",
	"src_dir" 		:	"/Users/jemy/Photos",
	"access_key" 	:	"<Your AccessKey>",
	"secret_key"	:	"<Your SecretKey>",
	"bucket"		:	"test-bucket",
	"ignore_dir"	:	false,
	"key_prefix"	:	"2014/12/01/",
	"overwrite"		:	false,
	"check_exists"	:	true
}

or without up_host and key_prefix and ignore_dir and check_exists

{
	"src_dir" 		:	"/Users/jemy/Photos",
	"access_key" 	:	"<Your AccessKey>",
	"secret_key"	:	"<Your SecretKey>",
	"bucket"		:	"test-bucket",
}
*/

const (
	PUT_THRESHOLD           int64 = 10 * 1 << 20
	MIN_UPLOAD_THREAD_COUNT int64 = 1
	MAX_UPLOAD_THREAD_COUNT int64 = 100
)

type UploadConfig struct {
	SrcDir       string `json:"src_dir"`
	AccessKey    string `json:"access_key"`
	SecretKey    string `json:"secret_key"`
	Bucket       string `json:"bucket"`
	UpHost       string `json:"up_host,omitempty"`
	KeyPrefix    string `json:"key_prefix,omitempty"`
	IgnoreDir    bool   `json:"ignore_dir,omitempty"`
	Overwrite    bool   `json:"overwrite,omitempty"`
	CheckExists  bool   `json:"check_exists,omitempty"`
	SkipPrefixes string `json:"skip_prefixes,omitempty"`
	SkipSuffixes string `json:"skip_suffixes,omitempty"`
}

var upSettings = rio.Settings{
	ChunkSize: 1 * 1024 * 1024,
	TryTimes:  5,
}

func QiniuUpload(threadCount int, uploadConfigFile string) {
	fp, err := os.Open(uploadConfigFile)
	if err != nil {
		log.Error(fmt.Sprintf("Open upload config file `%s' error due to `%s'", uploadConfigFile, err))
		return
	}
	defer fp.Close()
	configData, err := ioutil.ReadAll(fp)
	if err != nil {
		log.Error(fmt.Sprintf("Read upload config file `%s' error due to `%s'", uploadConfigFile, err))
		return
	}
	var uploadConfig UploadConfig
	err = json.Unmarshal(configData, &uploadConfig)
	if err != nil {
		log.Error(fmt.Sprintf("Parse upload config file `%s' errror due to `%s'", uploadConfigFile, err))
		return
	}
	if _, err := os.Stat(uploadConfig.SrcDir); err != nil {
		log.Error("Upload config error for parameter `SrcDir`,", err)
		return
	}
	dirCache := DirCache{}
	currentUser, err := user.Current()
	if err != nil {
		log.Error("Failed to get current user", err)
		return
	}

	pathSep := string(os.PathSeparator)
	//create job id
	md5Hasher := md5.New()
	md5Hasher.Write([]byte(strings.TrimSuffix(uploadConfig.SrcDir, pathSep) + ":" + uploadConfig.Bucket))
	jobId := fmt.Sprintf("%x", md5Hasher.Sum(nil))

	//local storage path
	storePath := filepath.Join(currentUser.HomeDir, ".qshell", "qupload", jobId)
	err = os.MkdirAll(storePath, 0775)
	if err != nil {
		log.Error(fmt.Sprintf("Failed to mkdir `%s' due to `%s'", storePath, err))
		return
	}

	//cache file
	cacheFileName := filepath.Join(storePath, jobId+".cache")
	//leveldb folder
	leveldbFileName := filepath.Join(storePath, jobId+".ldb")

	totalFileCount := dirCache.Cache(uploadConfig.SrcDir, cacheFileName)
	ldb, err := leveldb.OpenFile(leveldbFileName, nil)
	if err != nil {
		log.Error(fmt.Sprintf("Open leveldb `%s' failed due to `%s'", leveldbFileName, err))
		return
	}
	defer ldb.Close()
	//sync
	ufp, err := os.Open(cacheFileName)
	if err != nil {
		log.Error(fmt.Sprintf("Open cache file `%s' failed due to `%s'", cacheFileName, err))
		return
	}
	defer ufp.Close()
	bScanner := bufio.NewScanner(ufp)
	bScanner.Split(bufio.ScanLines)

	var currentFileCount int64 = 0
	var successFileCount int64 = 0
	var failureFileCount int64 = 0
	var skippedFileCount int64 = 0

	ldbWOpt := opt.WriteOptions{
		Sync: true,
	}

	upWorkGroup := sync.WaitGroup{}
	upCounter := 0
	threadThreshold := threadCount + 1

	//use host if not empty
	if uploadConfig.UpHost != "" {
		conf.UP_HOST = uploadConfig.UpHost
	}
	//set resume upload settings
	rio.SetSettings(&upSettings)
	mac := digest.Mac{uploadConfig.AccessKey, []byte(uploadConfig.SecretKey)}

	for bScanner.Scan() {
		line := strings.TrimSpace(bScanner.Text())
		items := strings.Split(line, "\t")
		if len(items) != 3 {
			log.Error(fmt.Sprintf("Invalid cache line `%s'", line))
			continue
		}

		localFname := items[0]
		currentFileCount += 1

		skip := false
		//check skip local file or folder
		if uploadConfig.SkipPrefixes != "" {
			//unpack skip prefix
			skipPrefixes := strings.Split(uploadConfig.SkipPrefixes, ",")
			for _, prefix := range skipPrefixes {
				if strings.HasPrefix(localFname, strings.TrimSpace(prefix)) {
					log.Debug(fmt.Sprintf("Skip by prefix '%s' for local file %s",
						strings.TrimSpace(prefix), localFname))
					skip = true
					skippedFileCount += 1
					break
				}
			}

			if skip {

				continue
			}
		}

		if uploadConfig.SkipSuffixes != "" {
			skipSuffixes := strings.Split(uploadConfig.SkipSuffixes, ",")
			for _, suffix := range skipSuffixes {
				if strings.HasSuffix(localFname, strings.TrimSpace(suffix)) {
					log.Debug(fmt.Sprintf("Skip by suffix '%s' for local file %s",
						strings.TrimSpace(suffix), localFname))
					skip = true
					skippedFileCount += 1
					break
				}
			}

			if skip {
				continue
			}
		}

		//pack the upload file key
		localFlmd, _ := strconv.Atoi(items[2])
		uploadFileKey := localFname

		if uploadConfig.IgnoreDir {
			if i := strings.LastIndex(uploadFileKey, pathSep); i != -1 {
				uploadFileKey = uploadFileKey[i+1:]
			}
		}
		if uploadConfig.KeyPrefix != "" {
			uploadFileKey = strings.Join([]string{uploadConfig.KeyPrefix, uploadFileKey}, "")
		}
		//convert \ to / under windows
		if runtime.GOOS == "windows" {
			uploadFileKey = strings.Replace(uploadFileKey, "\\", "/", -1)
		}

		localFilePath := filepath.Join(uploadConfig.SrcDir, localFname)
		fstat, err := os.Stat(localFilePath)
		if err != nil {
			log.Error(fmt.Sprintf("Error stat local file `%s' due to `%s'", localFilePath, err))
			return
		}

		fsize := fstat.Size()
		ldbKey := fmt.Sprintf("%s => %s", localFilePath, uploadFileKey)

		log.Info(fmt.Sprintf("Uploading %s (%d/%d, %.1f%%) ...", ldbKey, currentFileCount, totalFileCount,
			float32(currentFileCount)*100/float32(totalFileCount)))

		rsClient := rs.New(&mac)

		//check exists
		if uploadConfig.CheckExists {
			rsEntry, checkErr := rsClient.Stat(nil, uploadConfig.Bucket, uploadFileKey)
			if checkErr == nil {
				//compare hash
				localEtag, cErr := GetEtag(localFilePath)
				if cErr != nil {
					atomic.AddInt64(&failureFileCount, 1)
					log.Error("Calc local file hash failed,", cErr)
					continue
				}
				if rsEntry.Hash == localEtag {
					atomic.AddInt64(&skippedFileCount, 1)
					log.Debug(fmt.Sprintf("File %s already exists in bucket, ignore this upload", uploadFileKey))
					continue
				}
			} else {
				if _, ok := checkErr.(*rpc.ErrorInfo); !ok {
					//not logic error, should be network error
					atomic.AddInt64(&failureFileCount, 1)
					continue
				}
			}
		} else {
			//check leveldb
			ldbFlmd, err := ldb.Get([]byte(ldbKey), nil)
			flmd, _ := strconv.Atoi(string(ldbFlmd))
			//not exist, return ErrNotFound
			//check last modified

			if err == nil && localFlmd == flmd {
				log.Debug("Skip by local log for file", localFname)
				atomic.AddInt64(&skippedFileCount, 1)
				continue
			}
		}

		//worker
		upCounter += 1
		if upCounter%threadThreshold == 0 {
			upWorkGroup.Wait()
		}
		upWorkGroup.Add(1)

		//start to upload
		go func() {
			defer upWorkGroup.Done()

			policy := rs.PutPolicy{}
			policy.Scope = uploadConfig.Bucket
			if uploadConfig.Overwrite {
				policy.Scope = uploadConfig.Bucket + ":" + uploadFileKey
				policy.InsertOnly = 0
			}
			policy.Expires = 24 * 3600
			uptoken := policy.Token(&mac)
			if fsize > PUT_THRESHOLD {
				putRet := rio.PutRet{}
				err := rio.PutFile(nil, &putRet, uptoken, uploadFileKey, localFilePath, nil)
				if err != nil {
					atomic.AddInt64(&failureFileCount, 1)
					if pErr, ok := err.(*rpc.ErrorInfo); ok {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, pErr.Err))
					} else {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, err))
					}
				} else {
					atomic.AddInt64(&successFileCount, 1)
					perr := ldb.Put([]byte(ldbKey), []byte("Y"), &ldbWOpt)
					if perr != nil {
						log.Error(fmt.Sprintf("Put key `%s' into leveldb error due to `%s'", ldbKey, perr))
					}
				}
			} else {
				putRet := fio.PutRet{}
				err := fio.PutFile(nil, &putRet, uptoken, uploadFileKey, localFilePath, nil)
				if err != nil {
					atomic.AddInt64(&failureFileCount, 1)
					if pErr, ok := err.(*rpc.ErrorInfo); ok {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, pErr.Err))
					} else {
						log.Error(fmt.Sprintf("Put file `%s' => `%s' failed due to `%s'", localFilePath, uploadFileKey, err))
					}
				} else {
					atomic.AddInt64(&successFileCount, 1)
					perr := ldb.Put([]byte(ldbKey), []byte(strconv.Itoa(localFlmd)), &ldbWOpt)
					if perr != nil {
						log.Error(fmt.Sprintf("Put key `%s' into leveldb error due to `%s'", ldbKey, perr))
					}
				}
			}
		}()

	}
	upWorkGroup.Wait()

	log.Info("-------Upload Done-------")
	log.Info("Total:\t", currentFileCount)
	log.Info("Success:\t", successFileCount)
	log.Info("Failure:\t", failureFileCount)
	log.Info("Skipped:\t", skippedFileCount)
	log.Info("-------------------------")

}
<|endoftext|>


<commit_before>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package sshstorage

import (
	"bytes"
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	stdtesting "testing"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs/storage"
	coreerrors "launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/utils"
)

type storageSuite struct {
	testing.LoggingSuite
	restoreEnv func()
}

var _ = gc.Suite(&storageSuite{})

func Test(t *stdtesting.T) {
	gc.TestingT(t)
}

var sshCommandOrig = sshCommand

func sshCommandTesting(host string, tty bool, command string) *exec.Cmd {
	cmd := exec.Command("bash", "-c", command)
	uid := fmt.Sprint(os.Getuid())
	gid := fmt.Sprint(os.Getgid())
	defer testing.PatchEnvironment("SUDO_UID", uid)()
	defer testing.PatchEnvironment("SUDO_GID", gid)()
	cmd.Env = os.Environ()
	return cmd
}

// flockBin is the path to the original "flock" binary.
var flockBin string

func (s *storageSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)

	var err error
	flockBin, err = exec.LookPath("flock")
	c.Assert(err, gc.IsNil)

	bin := c.MkDir()
	s.restoreEnv = testing.PatchEnvironment("PATH", bin+":"+os.Getenv("PATH"))

	// Create a "sudo" command which just executes its args.
	c.Assert(os.Symlink("/usr/bin/env", filepath.Join(bin, "sudo")), gc.IsNil)
	sshCommand = sshCommandTesting

	// Create a new "flock" which calls the original, but in non-blocking mode.
	data := []byte(fmt.Sprintf("#!/bin/sh\nexec %s --nonblock \"$@\"", flockBin))
	c.Assert(ioutil.WriteFile(filepath.Join(bin, "flock"), data, 0755), gc.IsNil)
}

func (s *storageSuite) TearDownSuite(c *gc.C) {
	sshCommand = sshCommandOrig
	s.restoreEnv()
	s.LoggingSuite.TearDownSuite(c)
}

func (s *storageSuite) TestNewSSHStorage(c *gc.C) {
	storageDir := c.MkDir()
	for i := 0; i < 2; i++ {
		stor, err := NewSSHStorage("example.com", storageDir)
		c.Assert(err, gc.IsNil)
		c.Assert(stor, gc.NotNil)
		c.Assert(stor.Close(), gc.IsNil)
	}
	c.Assert(os.RemoveAll(storageDir), gc.IsNil)

	// You must have permissions to create the directory.
	storageDir = c.MkDir()
	c.Assert(os.Chmod(storageDir, 0555), gc.IsNil)
	_, err := NewSSHStorage("example.com", filepath.Join(storageDir))
	c.Assert(err, gc.ErrorMatches, "(.|\n)*cannot change owner and permissions of(.|\n)*")
}

func (s *storageSuite) TestPathValidity(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()

	c.Assert(os.Mkdir(filepath.Join(storageDir, contentdir, "a"), 0755), gc.IsNil)
	f, err := os.Create(filepath.Join(storageDir, contentdir, "a", "b"))
	c.Assert(err, gc.IsNil)
	c.Assert(f.Close(), gc.IsNil)

	for _, prefix := range []string{"..", "a/../.."} {
		c.Logf("prefix: %q", prefix)
		_, err := storage.List(stor, prefix)
		c.Check(err, gc.ErrorMatches, regexp.QuoteMeta(fmt.Sprintf("%q escapes storage directory", prefix)))
	}

	// Paths are always relative, so a leading "/" may as well not be there.
	names, err := storage.List(stor, "/")
	c.Assert(err, gc.IsNil)
	c.Assert(names, gc.DeepEquals, []string{"a/b"})

	// Paths will be canonicalised.
	names, err = storage.List(stor, "a/..")
	c.Assert(err, gc.IsNil)
	c.Assert(names, gc.DeepEquals, []string{"a/b"})
}

func (s *storageSuite) TestGet(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	data := []byte("abc\000def")
	c.Assert(os.Mkdir(filepath.Join(storageDir, contentdir, "a"), 0755), gc.IsNil)
	for _, name := range []string{"b", filepath.Join("a", "b")} {
		err = ioutil.WriteFile(filepath.Join(storageDir, contentdir, name), data, 0644)
		c.Assert(err, gc.IsNil)
		r, err := storage.Get(stor, name)
		c.Assert(err, gc.IsNil)
		out, err := ioutil.ReadAll(r)
		c.Assert(err, gc.IsNil)
		c.Assert(out, gc.DeepEquals, data)
	}

	_, err = storage.Get(stor, "notthere")
	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)
}

func (s *storageSuite) TestPut(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	data := []byte("abc\000def")
	for _, name := range []string{"b", filepath.Join("a", "b")} {
		err = stor.Put(name, bytes.NewBuffer(data), int64(len(data)))
		c.Assert(err, gc.IsNil)
		out, err := ioutil.ReadFile(filepath.Join(storageDir, contentdir, name))
		c.Assert(err, gc.IsNil)
		c.Assert(out, gc.DeepEquals, data)
	}
}

func (s *storageSuite) assertList(c *gc.C, stor storage.StorageReader, prefix string, expected []string) {
	c.Logf("List: %v", prefix)
	names, err := storage.List(stor, prefix)
	c.Assert(err, gc.IsNil)
	c.Assert(names, gc.DeepEquals, expected)
}

func (s *storageSuite) TestList(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	s.assertList(c, stor, "", nil)

	// Directories don't show up in List.
	contentDir := filepath.Join(storageDir, contentdir)
	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
	s.assertList(c, stor, "", nil)
	s.assertList(c, stor, "a", nil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "b"), nil, 0), gc.IsNil)
	s.assertList(c, stor, "", []string{"a/b1", "a/b2", "b"})
	s.assertList(c, stor, "a", []string{"a/b1", "a/b2"})
	s.assertList(c, stor, "a/b", []string{"a/b1", "a/b2"})
	s.assertList(c, stor, "a/b1", []string{"a/b1"})
	s.assertList(c, stor, "a/b3", nil)
	s.assertList(c, stor, "a/b/c", nil)
	s.assertList(c, stor, "b", []string{"b"})
}

func (s *storageSuite) TestRemove(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()

	contentDir := filepath.Join(storageDir, contentdir)
	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
	c.Assert(stor.Remove("a"), gc.ErrorMatches, "rm: cannot remove.*Is a directory")
	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
	c.Assert(stor.Remove("a/b"), gc.IsNil) // doesn't exist; not an error
	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
	c.Assert(stor.Remove("a/b2"), gc.IsNil)
	s.assertList(c, stor, "", []string{"a/b1"})
	c.Assert(stor.Remove("a/b1"), gc.IsNil)
	s.assertList(c, stor, "", nil)
}

func (s *storageSuite) TestRemoveAll(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()

	contentDir := filepath.Join(storageDir, contentdir)
	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
	c.Assert(stor.RemoveAll(), gc.IsNil)
	s.assertList(c, stor, "", nil)

	// RemoveAll does not remove the base storage directory.
	_, err = os.Stat(contentDir)
	c.Assert(err, gc.IsNil)
}

func (s *storageSuite) TestURL(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	url, err := stor.URL("a/b")
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "sftp://example.com/"+path.Join(storageDir, contentdir, "a/b"))
}

func (s *storageSuite) TestConsistencyStrategy(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	c.Assert(stor.DefaultConsistencyStrategy(), gc.Equals, utils.AttemptStrategy{})
}

// flock is a test helper that flocks a file,
// executes "sleep" with the specified duration,
// and returns the *Cmd so it can be early terminated.
func (s *storageSuite) flock(c *gc.C, mode flockmode, lockfile string, duration time.Duration) *os.Process {
	sleepcmd := fmt.Sprintf("sleep %vs", duration.Seconds())
	cmd := exec.Command(flockBin, "--nonblock", "--close", string(mode), lockfile, "-c", sleepcmd)
	c.Assert(cmd.Start(), gc.IsNil)
	return cmd.Process
}

const defaultFlockTimeout = 5 * time.Second

func (s *storageSuite) TestSynchronisation(c *gc.C) {
	storageDir := c.MkDir()
	proc := s.flock(c, flockShared, storageDir, defaultFlockTimeout)
	defer proc.Wait()
	defer proc.Kill()

	// Creating storage requires an exclusive lock initially.
	//
	// flock exits with exit code 1 if it can't acquire the
	// lock immediately in non-blocking mode (which the tests force).
	_, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.ErrorMatches, "exit code 1")

	proc.Kill()
	proc.Wait()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)

	// Get and List should be able to proceed with a shared lock.
	// All other methods should fail.
	data := []byte("abc\000def")
	c.Assert(ioutil.WriteFile(filepath.Join(storageDir, contentdir, "a"), data, 0644), gc.IsNil)

	proc = s.flock(c, flockShared, storageDir, defaultFlockTimeout)
	_, err = storage.Get(stor, "a")
	c.Assert(err, gc.IsNil)
	_, err = storage.List(stor, "")
	c.Assert(err, gc.IsNil)
	c.Assert(stor.Put("a", bytes.NewBuffer(nil), 0), gc.NotNil)
	c.Assert(stor.Remove("a"), gc.NotNil)
	c.Assert(stor.RemoveAll(), gc.NotNil)
	proc.Kill()
	proc.Wait()

	// None of the methods (apart from URL) should be able to do anything
	// while an exclusive lock is held.
	proc = s.flock(c, flockExclusive, storageDir, defaultFlockTimeout)
	_, err = stor.URL("a")
	c.Assert(err, gc.IsNil)
	c.Assert(stor.Put("a", bytes.NewBuffer(nil), 0), gc.NotNil)
	c.Assert(stor.Remove("a"), gc.NotNil)
	c.Assert(stor.RemoveAll(), gc.NotNil)
	_, err = storage.Get(stor, "a")
	c.Assert(err, gc.NotNil)
	_, err = storage.List(stor, "")
	c.Assert(err, gc.NotNil)
}
<commit_msg>[r=thumper] Fix race condition in SSHStorage test<commit_after>// Copyright 2013 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package sshstorage

import (
	"bytes"
	"fmt"
	"io"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"path/filepath"
	"regexp"
	stdtesting "testing"
	"time"

	gc "launchpad.net/gocheck"

	"launchpad.net/juju-core/environs/storage"
	coreerrors "launchpad.net/juju-core/errors"
	"launchpad.net/juju-core/testing"
	jc "launchpad.net/juju-core/testing/checkers"
	"launchpad.net/juju-core/utils"
)

type storageSuite struct {
	testing.LoggingSuite
	restoreEnv func()
}

var _ = gc.Suite(&storageSuite{})

func Test(t *stdtesting.T) {
	gc.TestingT(t)
}

var sshCommandOrig = sshCommand

func sshCommandTesting(host string, tty bool, command string) *exec.Cmd {
	cmd := exec.Command("bash", "-c", command)
	uid := fmt.Sprint(os.Getuid())
	gid := fmt.Sprint(os.Getgid())
	defer testing.PatchEnvironment("SUDO_UID", uid)()
	defer testing.PatchEnvironment("SUDO_GID", gid)()
	cmd.Env = os.Environ()
	return cmd
}

// flockBin is the path to the original "flock" binary.
var flockBin string

func (s *storageSuite) SetUpSuite(c *gc.C) {
	s.LoggingSuite.SetUpSuite(c)

	var err error
	flockBin, err = exec.LookPath("flock")
	c.Assert(err, gc.IsNil)

	bin := c.MkDir()
	s.restoreEnv = testing.PatchEnvironment("PATH", bin+":"+os.Getenv("PATH"))

	// Create a "sudo" command which just executes its args.
	c.Assert(os.Symlink("/usr/bin/env", filepath.Join(bin, "sudo")), gc.IsNil)
	sshCommand = sshCommandTesting

	// Create a new "flock" which calls the original, but in non-blocking mode.
	data := []byte(fmt.Sprintf("#!/bin/sh\nexec %s --nonblock \"$@\"", flockBin))
	c.Assert(ioutil.WriteFile(filepath.Join(bin, "flock"), data, 0755), gc.IsNil)
}

func (s *storageSuite) TearDownSuite(c *gc.C) {
	sshCommand = sshCommandOrig
	s.restoreEnv()
	s.LoggingSuite.TearDownSuite(c)
}

func (s *storageSuite) TestNewSSHStorage(c *gc.C) {
	storageDir := c.MkDir()
	for i := 0; i < 2; i++ {
		stor, err := NewSSHStorage("example.com", storageDir)
		c.Assert(err, gc.IsNil)
		c.Assert(stor, gc.NotNil)
		c.Assert(stor.Close(), gc.IsNil)
	}
	c.Assert(os.RemoveAll(storageDir), gc.IsNil)

	// You must have permissions to create the directory.
	storageDir = c.MkDir()
	c.Assert(os.Chmod(storageDir, 0555), gc.IsNil)
	_, err := NewSSHStorage("example.com", filepath.Join(storageDir))
	c.Assert(err, gc.ErrorMatches, "(.|\n)*cannot change owner and permissions of(.|\n)*")
}

func (s *storageSuite) TestPathValidity(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()

	c.Assert(os.Mkdir(filepath.Join(storageDir, contentdir, "a"), 0755), gc.IsNil)
	f, err := os.Create(filepath.Join(storageDir, contentdir, "a", "b"))
	c.Assert(err, gc.IsNil)
	c.Assert(f.Close(), gc.IsNil)

	for _, prefix := range []string{"..", "a/../.."} {
		c.Logf("prefix: %q", prefix)
		_, err := storage.List(stor, prefix)
		c.Check(err, gc.ErrorMatches, regexp.QuoteMeta(fmt.Sprintf("%q escapes storage directory", prefix)))
	}

	// Paths are always relative, so a leading "/" may as well not be there.
	names, err := storage.List(stor, "/")
	c.Assert(err, gc.IsNil)
	c.Assert(names, gc.DeepEquals, []string{"a/b"})

	// Paths will be canonicalised.
	names, err = storage.List(stor, "a/..")
	c.Assert(err, gc.IsNil)
	c.Assert(names, gc.DeepEquals, []string{"a/b"})
}

func (s *storageSuite) TestGet(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	data := []byte("abc\000def")
	c.Assert(os.Mkdir(filepath.Join(storageDir, contentdir, "a"), 0755), gc.IsNil)
	for _, name := range []string{"b", filepath.Join("a", "b")} {
		err = ioutil.WriteFile(filepath.Join(storageDir, contentdir, name), data, 0644)
		c.Assert(err, gc.IsNil)
		r, err := storage.Get(stor, name)
		c.Assert(err, gc.IsNil)
		out, err := ioutil.ReadAll(r)
		c.Assert(err, gc.IsNil)
		c.Assert(out, gc.DeepEquals, data)
	}

	_, err = storage.Get(stor, "notthere")
	c.Assert(err, jc.Satisfies, coreerrors.IsNotFoundError)
}

func (s *storageSuite) TestPut(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	data := []byte("abc\000def")
	for _, name := range []string{"b", filepath.Join("a", "b")} {
		err = stor.Put(name, bytes.NewBuffer(data), int64(len(data)))
		c.Assert(err, gc.IsNil)
		out, err := ioutil.ReadFile(filepath.Join(storageDir, contentdir, name))
		c.Assert(err, gc.IsNil)
		c.Assert(out, gc.DeepEquals, data)
	}
}

func (s *storageSuite) assertList(c *gc.C, stor storage.StorageReader, prefix string, expected []string) {
	c.Logf("List: %v", prefix)
	names, err := storage.List(stor, prefix)
	c.Assert(err, gc.IsNil)
	c.Assert(names, gc.DeepEquals, expected)
}

func (s *storageSuite) TestList(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	s.assertList(c, stor, "", nil)

	// Directories don't show up in List.
	contentDir := filepath.Join(storageDir, contentdir)
	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
	s.assertList(c, stor, "", nil)
	s.assertList(c, stor, "a", nil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "b"), nil, 0), gc.IsNil)
	s.assertList(c, stor, "", []string{"a/b1", "a/b2", "b"})
	s.assertList(c, stor, "a", []string{"a/b1", "a/b2"})
	s.assertList(c, stor, "a/b", []string{"a/b1", "a/b2"})
	s.assertList(c, stor, "a/b1", []string{"a/b1"})
	s.assertList(c, stor, "a/b3", nil)
	s.assertList(c, stor, "a/b/c", nil)
	s.assertList(c, stor, "b", []string{"b"})
}

func (s *storageSuite) TestRemove(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()

	contentDir := filepath.Join(storageDir, contentdir)
	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
	c.Assert(stor.Remove("a"), gc.ErrorMatches, "rm: cannot remove.*Is a directory")
	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
	c.Assert(stor.Remove("a/b"), gc.IsNil) // doesn't exist; not an error
	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
	c.Assert(stor.Remove("a/b2"), gc.IsNil)
	s.assertList(c, stor, "", []string{"a/b1"})
	c.Assert(stor.Remove("a/b1"), gc.IsNil)
	s.assertList(c, stor, "", nil)
}

func (s *storageSuite) TestRemoveAll(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()

	contentDir := filepath.Join(storageDir, contentdir)
	c.Assert(os.Mkdir(filepath.Join(contentDir, "a"), 0755), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b1"), nil, 0), gc.IsNil)
	c.Assert(ioutil.WriteFile(filepath.Join(contentDir, "a", "b2"), nil, 0), gc.IsNil)
	s.assertList(c, stor, "", []string{"a/b1", "a/b2"})
	c.Assert(stor.RemoveAll(), gc.IsNil)
	s.assertList(c, stor, "", nil)

	// RemoveAll does not remove the base storage directory.
	_, err = os.Stat(contentDir)
	c.Assert(err, gc.IsNil)
}

func (s *storageSuite) TestURL(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	url, err := stor.URL("a/b")
	c.Assert(err, gc.IsNil)
	c.Assert(url, gc.Equals, "sftp://example.com/"+path.Join(storageDir, contentdir, "a/b"))
}

func (s *storageSuite) TestConsistencyStrategy(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	defer stor.Close()
	c.Assert(stor.DefaultConsistencyStrategy(), gc.Equals, utils.AttemptStrategy{})
}

const defaultFlockTimeout = 5 * time.Second

// flock is a test helper that flocks a file, executes "sleep" with the
// specified duration, the command is terminated in the test tear down.
func (s *storageSuite) flock(c *gc.C, mode flockmode, lockfile string) {
	sleepcmd := fmt.Sprintf("echo started && sleep %vs", defaultFlockTimeout.Seconds())
	cmd := exec.Command(flockBin, "--nonblock", "--close", string(mode), lockfile, "-c", sleepcmd)
	stdout, err := cmd.StdoutPipe()
	c.Assert(err, gc.IsNil)
	c.Assert(cmd.Start(), gc.IsNil)
	// Make sure the flock has been taken before returning by reading stdout waiting for "started"
	_, err = io.ReadFull(stdout, make([]byte, len("started")))
	c.Assert(err, gc.IsNil)
	s.AddCleanup(func(*gc.C) {
		cmd.Process.Kill()
		cmd.Process.Wait()
	})
}

func (s *storageSuite) TestCreateFailsIfFlockNotAvailable(c *gc.C) {
	storageDir := c.MkDir()
	s.flock(c, flockShared, storageDir)
	// Creating storage requires an exclusive lock initially.
	//
	// flock exits with exit code 1 if it can't acquire the
	// lock immediately in non-blocking mode (which the tests force).
	_, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.ErrorMatches, "exit code 1")
}

func (s *storageSuite) TestWithSharedLocks(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)

	// Get and List should be able to proceed with a shared lock.
	// All other methods should fail.
	data := []byte("abc\000def")
	c.Assert(ioutil.WriteFile(filepath.Join(storageDir, contentdir, "a"), data, 0644), gc.IsNil)

	s.flock(c, flockShared, storageDir)
	_, err = storage.Get(stor, "a")
	c.Assert(err, gc.IsNil)
	_, err = storage.List(stor, "")
	c.Assert(err, gc.IsNil)
	c.Assert(stor.Put("a", bytes.NewBuffer(nil), 0), gc.NotNil)
	c.Assert(stor.Remove("a"), gc.NotNil)
	c.Assert(stor.RemoveAll(), gc.NotNil)
}

func (s *storageSuite) TestWithExclusiveLocks(c *gc.C) {
	storageDir := c.MkDir()
	stor, err := NewSSHStorage("example.com", storageDir)
	c.Assert(err, gc.IsNil)
	// None of the methods (apart from URL) should be able to do anything
	// while an exclusive lock is held.
	s.flock(c, flockExclusive, storageDir)
	_, err = stor.URL("a")
	c.Assert(err, gc.IsNil)
	c.Assert(stor.Put("a", bytes.NewBuffer(nil), 0), gc.NotNil)
	c.Assert(stor.Remove("a"), gc.NotNil)
	c.Assert(stor.RemoveAll(), gc.NotNil)
	_, err = storage.Get(stor, "a")
	c.Assert(err, gc.NotNil)
	_, err = storage.List(stor, "")
	c.Assert(err, gc.NotNil)
}
<|endoftext|>


<commit_before>package nexus

import (
	"io"
	"net"
	"time"

	"github.com/gammazero/nexus/stdlog"
	"github.com/gammazero/nexus/transport"
)

// RawSocketServer handles socket connections.
type RawSocketServer struct {
	router Router

	log       stdlog.StdLog
	recvLimit int
	keepalive time.Duration
}

// NewRawSocketServer takes a router instance and creates a new socket server.
func NewRawSocketServer(r Router, recvLimit int, keepalive time.Duration) *RawSocketServer {
	return &RawSocketServer{
		router:    r,
		log:       r.Logger(),
		recvLimit: recvLimit,
		keepalive: keepalive,
	}
}

// ListenAndServe listens on the specified endpoint and starts a goroutine that
// accept new client connections until the returned io.closer is closed.
func (s *RawSocketServer) ListenAndServe(network, address string) (io.Closer, error) {
	l, err := net.Listen(network, address)
	if err != nil {
		s.log.Print(err)
		return nil, err
	}
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				// Error normal when listener closed, do not log.
				l.Close()
				return
			}
			if tcpConn, ok := conn.(*net.TCPConn); ok && s.keepalive != 0 {
				tcpConn.SetKeepAlive(true)
				tcpConn.SetKeepAlivePeriod(s.keepalive)
			} else {
				tcpConn.SetKeepAlive(false)
			}
			go s.handleRawSocket(conn)
		}
	}()

	return l, nil
}

func (s *RawSocketServer) handleRawSocket(conn net.Conn) {
	peer, err := transport.AcceptRawSocket(conn, s.log, s.recvLimit)
	if err != nil {
		s.log.Println("Error accepting rawsocket client:", err)
		return
	}

	if err := s.router.Attach(peer); err != nil {
		s.log.Println("Error attaching to router:", err)
	}
}
<commit_msg>fix set keepalive<commit_after>package nexus

import (
	"io"
	"net"
	"time"

	"github.com/gammazero/nexus/stdlog"
	"github.com/gammazero/nexus/transport"
)

// RawSocketServer handles socket connections.
type RawSocketServer struct {
	router Router

	log       stdlog.StdLog
	recvLimit int
	keepalive time.Duration
}

// NewRawSocketServer takes a router instance and creates a new socket server.
func NewRawSocketServer(r Router, recvLimit int, keepalive time.Duration) *RawSocketServer {
	return &RawSocketServer{
		router:    r,
		log:       r.Logger(),
		recvLimit: recvLimit,
		keepalive: keepalive,
	}
}

// ListenAndServe listens on the specified endpoint and starts a goroutine that
// accept new client connections until the returned io.closer is closed.
func (s *RawSocketServer) ListenAndServe(network, address string) (io.Closer, error) {
	l, err := net.Listen(network, address)
	if err != nil {
		s.log.Print(err)
		return nil, err
	}
	go func() {
		for {
			conn, err := l.Accept()
			if err != nil {
				// Error normal when listener closed, do not log.
				l.Close()
				return
			}
			if tcpConn, ok := conn.(*net.TCPConn); ok {
				if s.keepalive != 0 {
					tcpConn.SetKeepAlive(true)
					tcpConn.SetKeepAlivePeriod(s.keepalive)
				} else {
					tcpConn.SetKeepAlive(false)
				}
			}
			go s.handleRawSocket(conn)
		}
	}()

	return l, nil
}

func (s *RawSocketServer) handleRawSocket(conn net.Conn) {
	peer, err := transport.AcceptRawSocket(conn, s.log, s.recvLimit)
	if err != nil {
		s.log.Println("Error accepting rawsocket client:", err)
		return
	}

	if err := s.router.Attach(peer); err != nil {
		s.log.Println("Error attaching to router:", err)
	}
}
<|endoftext|>


<commit_before>// Copyright (c) 2015, Jose Luis Aracil Gomez (pepe@diselpro.com)
// Based on Nick Patavalis (npat@efault.net) poller package.
// All rights reserved.
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE.txt file.

package poll

import (
	"io"
	"sync"
	"syscall"
	"time"
)

const (
	O_RDONLY   int = syscall.O_RDONLY   // open the file read-only.
	O_WRONLY   int = syscall.O_WRONLY   // open the file write-only.
	O_RDWR     int = syscall.O_RDWR     // open the file read-write.
	O_NONBLOCK int = syscall.O_NONBLOCK // open in non block mode.
)

// fdCtl keeps control fields (locks, timers, etc) for a single
// direction. For every File there is one fdCtl for Read operations and
// another for Write operations.
type fdCtl struct {
	m        sync.Mutex
	cond     *sync.Cond
	deadline time.Time
	timer    *time.Timer
	timeout  bool
}

// OsFile interface with *os.File methods used in NewFromFile
type OsFile interface {
	Close() error
	Fd() uintptr
	Name() string
}

// File is an *os.File like object who adds polling capabilities
type File struct {
	closed bool // Set by Close(), never cleared
	fd     int
	name   string
	closeF func() error
	// Must hold respective lock to access
	r fdCtl // Control fields for Read operations
	w fdCtl // Control fields for Write operations
}

// NewFile returns a new File with the given file descriptor and name.
func NewFile(fd uintptr, name string) (*File, error) {
	err := syscall.SetNonblock(int(fd), true)
	if err != nil {
		return nil, err
	}
	file := &File{fd: int(fd), name: name}
	file.r.cond = sync.NewCond(&sync.Mutex{})
	file.w.cond = sync.NewCond(&sync.Mutex{})
	err = register(file)
	if err != nil {
		return nil, err
	}
	return file, nil
}

// NewFromFile returns a new *poll.File based on the given *os.File.
// You don't need to worry about closing the *os.File, *poll.File already does it.
func NewFromFile(of OsFile) (*File, error) {
	f, err := NewFile(of.Fd(), of.Name())
	if err != nil {
		return nil, err
	}
	f.closeF = of.Close
	return f, nil
}

// Name returns the name of the file as presented to Open.
func (f *File) Name() string {
	return f.name
}

// Fd returns the integer Unix file descriptor referencing the open file.
func (f *File) Fd() uintptr {
	return uintptr(f.fd)
}

// WriteString is like Write, but writes the contents of string s rather than a slice of bytes.
func (f *File) WriteString(s string) (int, error) {
	return f.Write([]byte(s))
}

// Read reads up to len(b) bytes from the File.
// It returns the number of bytes read and an error, if any.
func (f *File) Read(p []byte) (n int, err error) {
	f.r.m.Lock()
	n, err = f.sysrw(false, p)
	f.r.m.Unlock()
	return
}

// Write writes len(b) bytes to the File.
// It returns the number of bytes written and an error, if any.
// Write returns a non-nil error when n != len(b).
func (f *File) Write(p []byte) (n int, err error) {
	f.w.m.Lock()
	for n != len(p) {
		var nn int
		nn, err = f.sysrw(true, p[n:])
		n += nn
		if err != nil {
			break
		}
	}
	f.w.m.Unlock()
	return
}

func (f *File) sysrw(write bool, p []byte) (n int, err error) {
	var fdc *fdCtl
	var rwfun func(int, []byte) (int, error)
	var errEOF error

	if !write {
		// Prepare things for Read.
		fdc = &f.r
		rwfun = syscall.Read
		errEOF = io.EOF
	} else {
		// Prepare things for Write.
		fdc = &f.w
		rwfun = syscall.Write
		errEOF = io.ErrUnexpectedEOF
	}
	// Read & Write are identical
	fdc.cond.L.Lock()
	defer fdc.cond.L.Unlock()
	for {
		if f.closed {
			return 0, ErrClosed
		}
		if fdc.timeout {
			return 0, ErrTimeout
		}
		n, err = rwfun(f.fd, p)
		if err != nil {
			n = 0
			if err != syscall.EAGAIN {
				break
			}
			// EAGAIN
			startTrack(f.fd, write)
			fdc.cond.Wait()
			if f.closed || fdc.timeout {
				stopTrack(f.fd, write)
			}
			continue
		}
		if n == 0 && len(p) != 0 {
			err = errEOF
			break
		}
		break
	}
	return n, err
}

//Close closes the File, rendering it unusable for I/O. It returns an error, if any.
func (f *File) Close() error {
	if err := f.Lock(); err != nil {
		return err
	}
	defer f.Unlock()
	if f.closed {
		return ErrClosed
	}
	f.closed = true
	unregister(f)
	if f.r.timer != nil {
		f.r.timer.Stop()
	}
	if f.w.timer != nil {
		f.w.timer.Stop()
	}
	// Wake up everybody waiting on File.
	f.r.cond.Broadcast()
	f.w.cond.Broadcast()
	if f.closeF != nil {
		return f.closeF()
	}
	return syscall.Close(f.fd)
}

// SetDeadline sets the deadline for Read and write operations on File.
func (f *File) SetDeadline(t time.Time) error {
	if err := f.SetReadDeadline(t); err != nil {
		return err
	}
	return f.SetWriteDeadline(t)
}

// SetReadDeadline sets the deadline for Read operations on File.
func (f *File) SetReadDeadline(t time.Time) error {
	return f.setDeadline(false, t)
}

// SetWriteDeadline sets the deadline for Write operations on File.
func (f *File) SetWriteDeadline(t time.Time) error {
	return f.setDeadline(true, t)
}

func (f *File) setDeadline(write bool, t time.Time) error {
	var fdc *fdCtl

	if !write {
		// Setting read deadline
		fdc = &f.r
	} else {
		// Setting write deadline
		fdc = &f.w
	}
	// R & W deadlines are handled identically
	fdc.cond.L.Lock()
	if f.closed {
		fdc.cond.L.Unlock()
		return ErrClosed
	}
	fdc.deadline = t
	fdc.timeout = false
	if t.IsZero() {
		if fdc.timer != nil {
			fdc.timer.Stop()
		}
	} else {
		d := t.Sub(time.Now())
		if fdc.timer == nil {
			fdc.timer = time.AfterFunc(d,
				func() { f.timerEvent(write) })
		} else {
			fdc.timer.Stop()
			fdc.timer.Reset(d)
		}
	}
	fdc.cond.L.Unlock()
	return nil
}

// Lock locks the file. It must be called before perfoming
// miscellaneous operations (e.g. ioctls) on the underlying system
// file descriptor.
func (f *File) Lock() error {
	f.r.cond.L.Lock()
	f.w.cond.L.Lock()
	if f.closed {
		f.w.cond.L.Unlock()
		f.r.cond.L.Unlock()
		return ErrClosed
	}
	return nil
}

// Unlock unlocks the file.
func (f *File) Unlock() {
	f.w.cond.L.Unlock()
	f.r.cond.L.Unlock()
}

func (f *File) timerEvent(write bool) {
	var fdc *fdCtl

	if !write {
		// A Read timeout
		fdc = &f.r
	} else {
		// A Write timeout
		fdc = &f.w
	}
	fdc.cond.L.Lock()
	if !f.closed && !fdc.timeout &&
		!fdc.deadline.IsZero() && !fdc.deadline.After(time.Now()) {
		fdc.timeout = true
		fdc.cond.Broadcast()
	}
	fdc.cond.L.Unlock()
}
<commit_msg>Add Open function<commit_after>// Copyright (c) 2015, Jose Luis Aracil Gomez (pepe@diselpro.com)
// Based on Nick Patavalis (npat@efault.net) poller package.
// All rights reserved.
// Use of this source code is governed by a BSD-style license that can
// be found in the LICENSE.txt file.

package poll

import (
	"io"
	"sync"
	"syscall"
	"time"
)

const (
	O_RDONLY   int = syscall.O_RDONLY   // open the file read-only.
	O_WRONLY   int = syscall.O_WRONLY   // open the file write-only.
	O_RDWR     int = syscall.O_RDWR     // open the file read-write.
	O_NONBLOCK int = syscall.O_NONBLOCK // open in non block mode.
)

// fdCtl keeps control fields (locks, timers, etc) for a single
// direction. For every File there is one fdCtl for Read operations and
// another for Write operations.
type fdCtl struct {
	m        sync.Mutex
	cond     *sync.Cond
	deadline time.Time
	timer    *time.Timer
	timeout  bool
}

// OsFile interface with *os.File methods used in NewFromFile
type OsFile interface {
	Close() error
	Fd() uintptr
	Name() string
}

// File is an *os.File like object who adds polling capabilities
type File struct {
	closed bool // Set by Close(), never cleared
	fd     int
	name   string
	closeF func() error
	// Must hold respective lock to access
	r fdCtl // Control fields for Read operations
	w fdCtl // Control fields for Write operations
}

// NewFile returns a new File with the given file descriptor and name.
func NewFile(fd uintptr, name string) (*File, error) {
	err := syscall.SetNonblock(int(fd), true)
	if err != nil {
		return nil, err
	}
	file := &File{fd: int(fd), name: name}
	file.r.cond = sync.NewCond(&sync.Mutex{})
	file.w.cond = sync.NewCond(&sync.Mutex{})
	err = register(file)
	if err != nil {
		return nil, err
	}
	return file, nil
}

// Open the named path for reading, writing or both, depnding on the
// flags argument.
func Open(name string, flags int) (*File, error) {
	fd, err := syscall.Open(name, flags|syscall.O_CLOEXEC|syscall.O_NONBLOCK, 0666)
	if err != nil {
		return nil, err
	}
	return NewFile(uintptr(fd), name)
}

// NewFromFile returns a new *poll.File based on the given *os.File.
// You don't need to worry about closing the *os.File, *poll.File already does it.
func NewFromFile(of OsFile) (*File, error) {
	f, err := NewFile(of.Fd(), of.Name())
	if err != nil {
		return nil, err
	}
	f.closeF = of.Close
	return f, nil
}

// Name returns the name of the file as presented to Open.
func (f *File) Name() string {
	return f.name
}

// Fd returns the integer Unix file descriptor referencing the open file.
func (f *File) Fd() uintptr {
	return uintptr(f.fd)
}

// WriteString is like Write, but writes the contents of string s rather than a slice of bytes.
func (f *File) WriteString(s string) (int, error) {
	return f.Write([]byte(s))
}

// Read reads up to len(b) bytes from the File.
// It returns the number of bytes read and an error, if any.
func (f *File) Read(p []byte) (n int, err error) {
	f.r.m.Lock()
	n, err = f.sysrw(false, p)
	f.r.m.Unlock()
	return
}

// Write writes len(b) bytes to the File.
// It returns the number of bytes written and an error, if any.
// Write returns a non-nil error when n != len(b).
func (f *File) Write(p []byte) (n int, err error) {
	f.w.m.Lock()
	for n != len(p) {
		var nn int
		nn, err = f.sysrw(true, p[n:])
		n += nn
		if err != nil {
			break
		}
	}
	f.w.m.Unlock()
	return
}

func (f *File) sysrw(write bool, p []byte) (n int, err error) {
	var fdc *fdCtl
	var rwfun func(int, []byte) (int, error)
	var errEOF error

	if !write {
		// Prepare things for Read.
		fdc = &f.r
		rwfun = syscall.Read
		errEOF = io.EOF
	} else {
		// Prepare things for Write.
		fdc = &f.w
		rwfun = syscall.Write
		errEOF = io.ErrUnexpectedEOF
	}
	// Read & Write are identical
	fdc.cond.L.Lock()
	defer fdc.cond.L.Unlock()
	for {
		if f.closed {
			return 0, ErrClosed
		}
		if fdc.timeout {
			return 0, ErrTimeout
		}
		n, err = rwfun(f.fd, p)
		if err != nil {
			n = 0
			if err != syscall.EAGAIN {
				break
			}
			// EAGAIN
			startTrack(f.fd, write)
			fdc.cond.Wait()
			if f.closed || fdc.timeout {
				stopTrack(f.fd, write)
			}
			continue
		}
		if n == 0 && len(p) != 0 {
			err = errEOF
			break
		}
		break
	}
	return n, err
}

//Close closes the File, rendering it unusable for I/O. It returns an error, if any.
func (f *File) Close() error {
	if err := f.Lock(); err != nil {
		return err
	}
	defer f.Unlock()
	if f.closed {
		return ErrClosed
	}
	f.closed = true
	unregister(f)
	if f.r.timer != nil {
		f.r.timer.Stop()
	}
	if f.w.timer != nil {
		f.w.timer.Stop()
	}
	// Wake up everybody waiting on File.
	f.r.cond.Broadcast()
	f.w.cond.Broadcast()
	if f.closeF != nil {
		return f.closeF()
	}
	return syscall.Close(f.fd)
}

// SetDeadline sets the deadline for Read and write operations on File.
func (f *File) SetDeadline(t time.Time) error {
	if err := f.SetReadDeadline(t); err != nil {
		return err
	}
	return f.SetWriteDeadline(t)
}

// SetReadDeadline sets the deadline for Read operations on File.
func (f *File) SetReadDeadline(t time.Time) error {
	return f.setDeadline(false, t)
}

// SetWriteDeadline sets the deadline for Write operations on File.
func (f *File) SetWriteDeadline(t time.Time) error {
	return f.setDeadline(true, t)
}

func (f *File) setDeadline(write bool, t time.Time) error {
	var fdc *fdCtl

	if !write {
		// Setting read deadline
		fdc = &f.r
	} else {
		// Setting write deadline
		fdc = &f.w
	}
	// R & W deadlines are handled identically
	fdc.cond.L.Lock()
	if f.closed {
		fdc.cond.L.Unlock()
		return ErrClosed
	}
	fdc.deadline = t
	fdc.timeout = false
	if t.IsZero() {
		if fdc.timer != nil {
			fdc.timer.Stop()
		}
	} else {
		d := t.Sub(time.Now())
		if fdc.timer == nil {
			fdc.timer = time.AfterFunc(d,
				func() { f.timerEvent(write) })
		} else {
			fdc.timer.Stop()
			fdc.timer.Reset(d)
		}
	}
	fdc.cond.L.Unlock()
	return nil
}

// Lock locks the file. It must be called before perfoming
// miscellaneous operations (e.g. ioctls) on the underlying system
// file descriptor.
func (f *File) Lock() error {
	f.r.cond.L.Lock()
	f.w.cond.L.Lock()
	if f.closed {
		f.w.cond.L.Unlock()
		f.r.cond.L.Unlock()
		return ErrClosed
	}
	return nil
}

// Unlock unlocks the file.
func (f *File) Unlock() {
	f.w.cond.L.Unlock()
	f.r.cond.L.Unlock()
}

func (f *File) timerEvent(write bool) {
	var fdc *fdCtl

	if !write {
		// A Read timeout
		fdc = &f.r
	} else {
		// A Write timeout
		fdc = &f.w
	}
	fdc.cond.L.Lock()
	if !f.closed && !fdc.timeout &&
		!fdc.deadline.IsZero() && !fdc.deadline.After(time.Now()) {
		fdc.timeout = true
		fdc.cond.Broadcast()
	}
	fdc.cond.L.Unlock()
}
<|endoftext|>


<commit_before>package memalpha

import (
	"bufio"
	"bytes"
	"strconv"
	"strings"
	"testing"

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

func fakedClient(response []byte) *Client {
	request := bytes.NewBuffer([]byte{})

	serverReadWriter := bufio.NewReadWriter(
		bufio.NewReader(bytes.NewReader(response)),
		bufio.NewWriter(request),
	)

	return &Client{rw: serverReadWriter}
}

func TestServerError(t *testing.T) {
	errorMessage := "test fake"
	c := fakedClient([]byte("SERVER_ERROR " + errorMessage))

	err := c.Set("foo", []byte("bar"), false)
	e, ok := err.(ServerError)
	if ok && strings.Contains(e.Error(), "server error: "+errorMessage) {
		return
	}

	t.Fatalf("set(foo): Error = %v, want ServerError: test fake", err)
}

func TestClientError(t *testing.T) {
	errorMessage := "test fake"
	c := fakedClient([]byte("CLIENT_ERROR " + errorMessage))

	err := c.Set("foo", []byte("bar"), false)
	e, ok := err.(ClientError)
	if ok && strings.Contains(e.Error(), "client error: "+errorMessage) {
		return
	}

	t.Fatalf("set(foo): Error = %v, want ClientError: test fake", err)
}

func TestReplyError(t *testing.T) {
	c := fakedClient([]byte("ERROR"))

	err := c.Set("foo", []byte("bar"), false)
	assert.Equal(t, ErrReplyError, err)
}

func TestStatsProtocolError(t *testing.T) {
	c := fakedClient([]byte("foobar"))

	_, err := c.Stats()
	assert.Equal(t, ProtocolError("malformed stats response"), err)
}

func TestIncrValueError(t *testing.T) {
	c := fakedClient([]byte("foobar"))

	err := c.Set("foo", []byte("42"), true)
	_, err = c.Increment("foo", 1, false)
	assert.IsType(t, &strconv.NumError{}, err)
}

func TestIllegalVersionResponse(t *testing.T) {
	c := fakedClient([]byte("Illegal Ver 1"))

	_, err := c.Version()
	assert.IsType(t, ProtocolError(""), err)
}
<commit_msg>Refactor mem_textproto_test.go<commit_after>package memalpha

import (
	"bufio"
	"bytes"
	"io"
	"io/ioutil"
	"strconv"
	"strings"
	"testing"

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

func newFakedClient(response string, requestWriter io.Writer) *Client {
	return &Client{rw: bufio.NewReadWriter(
		bufio.NewReader(bytes.NewReader([]byte(response))),
		bufio.NewWriter(requestWriter),
	)}
}

func TestServerError(t *testing.T) {
	errorMessage := "test fake"
	c := newFakedClient("SERVER_ERROR "+errorMessage, ioutil.Discard)

	err := c.Set("foo", []byte("bar"), false)
	e, ok := err.(ServerError)
	if ok && strings.Contains(e.Error(), "server error: "+errorMessage) {
		return
	}

	t.Fatalf("set(foo): Error = %v, want ServerError: test fake", err)
}

func TestClientError(t *testing.T) {
	errorMessage := "test fake"
	c := newFakedClient("CLIENT_ERROR "+errorMessage, ioutil.Discard)

	err := c.Set("foo", []byte("bar"), false)
	e, ok := err.(ClientError)
	if ok && strings.Contains(e.Error(), "client error: "+errorMessage) {
		return
	}

	t.Fatalf("set(foo): Error = %v, want ClientError: test fake", err)
}

func TestReplyError(t *testing.T) {
	c := newFakedClient("ERROR", ioutil.Discard)
	err := c.Set("foo", []byte("bar"), false)
	assert.Equal(t, ErrReplyError, err)
}

func TestMalformedStatsResponse(t *testing.T) {
	{
		c := newFakedClient("foobar", ioutil.Discard)
		_, err := c.Stats()
		assert.Equal(t, ProtocolError("malformed stats response"), err)
	}
}

func TestIncrValueError(t *testing.T) {
	c := newFakedClient("foobar", ioutil.Discard)

	err := c.Set("foo", []byte("42"), true)
	_, err = c.Increment("foo", 1, false)
	assert.IsType(t, &strconv.NumError{}, err)
}

func TestMalformedVersionResponse(t *testing.T) {
	{
		c := newFakedClient("Malformed Ver 1", ioutil.Discard)
		_, err := c.Version()
		assert.IsType(t, ProtocolError(""), err)
	}

	{
		c := newFakedClient("ERROR", ioutil.Discard)
		_, err := c.Version()
		assert.Equal(t, ErrReplyError, err)
	}
}
<|endoftext|>


<commit_before>// Package mysql implements MySQL migrations.
package mysql

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/blue-jay/core/storage"
	driver "github.com/blue-jay/core/storage/driver/mysql"
	"github.com/blue-jay/core/storage/migration"
	"github.com/jmoiron/sqlx"
)

// *****************************************************************************
// Thread-Safe Configuration
// *****************************************************************************
/*
var (
	info      database.Info
	infoMutex sync.RWMutex
)

// SetConfig stores the config.
func SetConfig(i database.Info) {
	infoMutex.Lock()
	info = i
	infoMutex.Unlock()
}

// ResetConfig removes the config.
func ResetConfig() {
	infoMutex.Lock()
	info = database.Info{}
	infoMutex.Unlock()
}

// Config returns the config.
func Config() database.Info {
	infoMutex.RLock()
	defer infoMutex.RUnlock()
	return info
}*/

// Configuration defines the shared configuration interface.
type Configuration struct {
	driver.Info
}

/*
// Shared returns the global configuration information.
func Shared() Configuration {
	return Configuration{
		Config(),
	}
}*/

// *****************************************************************************
// Migration Creation
// *****************************************************************************

// New creates a migration connection to the database.
func (c Configuration) New() (*migration.Info, error) {
	var mig *migration.Info

	// Load the config
	i := c.Info

	// Update the config
	i.Parameter = "parseTime=true&multiStatements=true"

	// Create MySQL entity
	mi := &Entity{}

	// Connect to the database
	con, err := i.Connect(true)

	// If the database doesn't exist or can't connect
	if err != nil {
		// Close the open connection (since 'unknown database' is still an
		// active connection)
		con.Close()

		// Connect to database without a database
		con, err = i.Connect(false)
		if err != nil {
			return mig, err
		}

		// Create the database
		err = i.Create(con)
		if err != nil {
			return mig, err
		}

		// Close connection
		con.Close()

		// Reconnect to the database
		con, err = i.Connect(true)
		if err != nil {
			return mig, err
		}
	}

	// Store the connection in the entity
	mi.sql = con

	// Store the migration table name
	mi.table = c.Migration.Table

	if len(mi.table) == 0 {
		return mig, errors.New("MySQL.Migration.Table key is missing in config file.")
	}

	// Setup logic was here
	return migration.New(mi, mi.table, c.Migration.Folder)
}

// *****************************************************************************
// Interface
// *****************************************************************************

// Item defines the migration table.
type Item struct {
	ID        uint32    `db:"id"`
	Name      string    `db:"name"`
	CreatedAt time.Time `db:"created_at"`
}

// Entity defines fulfills the migration interface.
type Entity struct {
	table string
	sql   *sqlx.DB
}

// Extension returns the file extension with a period
func (t *Entity) Extension() string {
	return ".sql"
}

// TableExist returns true if the migration table exists
func (t *Entity) TableExist() error {
	_, err := t.sql.Exec(fmt.Sprintf("SELECT 1 FROM %v LIMIT 1;", t.table))
	if err != nil {
		return err
	}

	return err
}

// CreateTable returns true if the migration was created
func (t *Entity) CreateTable() error {
	_, err := t.sql.Exec(fmt.Sprintf(`CREATE TABLE %v (
		id INT UNSIGNED NOT NULL AUTO_INCREMENT,
  		name VARCHAR(191) NOT NULL,
		created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
		UNIQUE KEY (name),
  		PRIMARY KEY (id)
		);`, t.table))

	if err != nil {
		return err
	}

	return err
}

// Status returns last migration name
func (t *Entity) Status() (string, error) {
	result := &Item{}
	err := t.sql.Get(result, fmt.Sprintf("SELECT * FROM %v ORDER BY id DESC LIMIT 1;", t.table))

	// If no rows, then set to nil
	if err == sql.ErrNoRows {
		err = nil
	}

	return result.Name, err
}

// statusID returns last migration ID
func (t *Entity) statusID() (uint32, error) {
	result := &Item{}
	err := t.sql.Get(result, fmt.Sprintf("SELECT * FROM %v ORDER BY id DESC LIMIT 1;", t.table))
	return result.ID, err
}

// Migrate runs a query and returns error
func (t *Entity) Migrate(qry string) error {
	_, err := t.sql.Exec(qry)
	return err
}

// RecordUp adds a record to the database
func (t *Entity) RecordUp(name string) error {
	_, err := t.sql.Exec(fmt.Sprintf("INSERT INTO %v (name) VALUES (?);", t.table), name)
	return err
}

// RecordDown removes a record from the database and updates the AUTO_INCREMENT value
func (t *Entity) RecordDown(name string) error {
	_, err := t.sql.Exec(fmt.Sprintf("DELETE FROM %v WHERE name = ? LIMIT 1;", t.table), name)

	// If the record was removed successfully
	if err == nil {
		var ID uint32
		var nextID uint32 = 1

		// Get the last migration record now
		ID, err = t.statusID()

		// If there are no more migrations in the table
		if err == sql.ErrNoRows {
			// Leave ID at 1
		} else if err != nil {
			return err
		} else {
			nextID = ID
		}

		_, err = t.sql.Exec(fmt.Sprintf("ALTER TABLE %v AUTO_INCREMENT = %v;", t.table, nextID))
	}
	return err
}

// *****************************************************************************
// Test Helpers
// *****************************************************************************

// SetUp is a function for unit tests on a separate database.
func SetUp(envPath string, dbName string) (*migration.Info, Configuration) {
	// Get the environment variable
	if len(os.Getenv("JAYCONFIG")) == 0 {
		// Attempt to find env.json
		p, err := filepath.Abs(envPath)
		if err != nil {
			log.Fatalf("%v", err)
		}

		// Set the environment variable
		os.Setenv("JAYCONFIG", p)
	}

	// Get the config file path
	configFile := os.Getenv("JAYCONFIG")

	// Load the config
	config, err := storage.LoadConfig(configFile)
	if err != nil {
		log.Fatalf("%v", err)
	}

	// Set the database name
	config.MySQL.Database = dbName

	// Set the migration folder to the absolute path
	config.MySQL.Migration.Folder = filepath.Join(filepath.Dir(configFile), config.MySQL.Migration.Folder)

	// Create the migration configuration
	conf := Configuration{
		config.MySQL,
	}

	// Create the migration
	mig, err := conf.New()
	if err != nil {
		log.Fatalf("%v", err)
	}

	// Refresh the data with verbose logging
	err = mig.DownAll()
	if err != nil {
		log.Println(err)
	}

	err = mig.UpAll()
	if err != nil {
		log.Println(err)
	}

	return mig, conf
}

// TearDown removes the unit test database.
func TearDown(db *sqlx.DB, dbName string) error {
	// Drop the database
	_, err := db.Exec(fmt.Sprintf(`DROP DATABASE %v;`, dbName))
	return err
}
<commit_msg>Update mysql.go<commit_after>// Package mysql implements MySQL migrations.
package mysql

import (
	"database/sql"
	"errors"
	"fmt"
	"log"
	"os"
	"path/filepath"
	"time"

	"github.com/blue-jay/core/storage"
	driver "github.com/blue-jay/core/storage/driver/mysql"
	"github.com/blue-jay/core/storage/migration"
	"github.com/jmoiron/sqlx"
)

// *****************************************************************************
// Thread-Safe Configuration
// *****************************************************************************
/*
var (
	info      database.Info
	infoMutex sync.RWMutex
)

// SetConfig stores the config.
func SetConfig(i database.Info) {
	infoMutex.Lock()
	info = i
	infoMutex.Unlock()
}

// ResetConfig removes the config.
func ResetConfig() {
	infoMutex.Lock()
	info = database.Info{}
	infoMutex.Unlock()
}

// Config returns the config.
func Config() database.Info {
	infoMutex.RLock()
	defer infoMutex.RUnlock()
	return info
}*/

// Configuration defines the shared configuration interface.
type Configuration struct {
	driver.Info
}

/*
// Shared returns the global configuration information.
func Shared() Configuration {
	return Configuration{
		Config(),
	}
}*/

// *****************************************************************************
// Migration Creation
// *****************************************************************************

// New creates a migration connection to the database.
func (c Configuration) New() (*migration.Info, error) {
	var mig *migration.Info

	// Load the config
	i := c.Info

	// Update the config
	i.Parameter = "parseTime=true&multiStatements=true"

	// Create MySQL entity
	mi := &Entity{}

	// Connect to the database
	con, err := i.Connect(true)

	// If the database doesn't exist or can't connect
	if err != nil {
		// Close the open connection (since 'unknown database' is still an
		// active connection)
		if con != nil {
			con.Close()
		}

		// Connect to database without a database
		con, err = i.Connect(false)
		if err != nil {
			return mig, err
		}

		// Create the database
		err = i.Create(con)
		if err != nil {
			return mig, err
		}

		// Close connection
		con.Close()

		// Reconnect to the database
		con, err = i.Connect(true)
		if err != nil {
			return mig, err
		}
	}

	// Store the connection in the entity
	mi.sql = con

	// Store the migration table name
	mi.table = c.Migration.Table

	if len(mi.table) == 0 {
		return mig, errors.New("MySQL.Migration.Table key is missing in config file.")
	}

	// Setup logic was here
	return migration.New(mi, mi.table, c.Migration.Folder)
}

// *****************************************************************************
// Interface
// *****************************************************************************

// Item defines the migration table.
type Item struct {
	ID        uint32    `db:"id"`
	Name      string    `db:"name"`
	CreatedAt time.Time `db:"created_at"`
}

// Entity defines fulfills the migration interface.
type Entity struct {
	table string
	sql   *sqlx.DB
}

// Extension returns the file extension with a period
func (t *Entity) Extension() string {
	return ".sql"
}

// TableExist returns true if the migration table exists
func (t *Entity) TableExist() error {
	_, err := t.sql.Exec(fmt.Sprintf("SELECT 1 FROM %v LIMIT 1;", t.table))
	if err != nil {
		return err
	}

	return err
}

// CreateTable returns true if the migration was created
func (t *Entity) CreateTable() error {
	_, err := t.sql.Exec(fmt.Sprintf(`CREATE TABLE %v (
		id INT UNSIGNED NOT NULL AUTO_INCREMENT,
  		name VARCHAR(191) NOT NULL,
		created_at TIMESTAMP DEFAULT CURRENT_TIMESTAMP,
		UNIQUE KEY (name),
  		PRIMARY KEY (id)
		);`, t.table))

	if err != nil {
		return err
	}

	return err
}

// Status returns last migration name
func (t *Entity) Status() (string, error) {
	result := &Item{}
	err := t.sql.Get(result, fmt.Sprintf("SELECT * FROM %v ORDER BY id DESC LIMIT 1;", t.table))

	// If no rows, then set to nil
	if err == sql.ErrNoRows {
		err = nil
	}

	return result.Name, err
}

// statusID returns last migration ID
func (t *Entity) statusID() (uint32, error) {
	result := &Item{}
	err := t.sql.Get(result, fmt.Sprintf("SELECT * FROM %v ORDER BY id DESC LIMIT 1;", t.table))
	return result.ID, err
}

// Migrate runs a query and returns error
func (t *Entity) Migrate(qry string) error {
	_, err := t.sql.Exec(qry)
	return err
}

// RecordUp adds a record to the database
func (t *Entity) RecordUp(name string) error {
	_, err := t.sql.Exec(fmt.Sprintf("INSERT INTO %v (name) VALUES (?);", t.table), name)
	return err
}

// RecordDown removes a record from the database and updates the AUTO_INCREMENT value
func (t *Entity) RecordDown(name string) error {
	_, err := t.sql.Exec(fmt.Sprintf("DELETE FROM %v WHERE name = ? LIMIT 1;", t.table), name)

	// If the record was removed successfully
	if err == nil {
		var ID uint32
		var nextID uint32 = 1

		// Get the last migration record now
		ID, err = t.statusID()

		// If there are no more migrations in the table
		if err == sql.ErrNoRows {
			// Leave ID at 1
		} else if err != nil {
			return err
		} else {
			nextID = ID
		}

		_, err = t.sql.Exec(fmt.Sprintf("ALTER TABLE %v AUTO_INCREMENT = %v;", t.table, nextID))
	}
	return err
}

// *****************************************************************************
// Test Helpers
// *****************************************************************************

// SetUp is a function for unit tests on a separate database.
func SetUp(envPath string, dbName string) (*migration.Info, Configuration) {
	// Get the environment variable
	if len(os.Getenv("JAYCONFIG")) == 0 {
		// Attempt to find env.json
		p, err := filepath.Abs(envPath)
		if err != nil {
			log.Fatalf("%v", err)
		}

		// Set the environment variable
		os.Setenv("JAYCONFIG", p)
	}

	// Get the config file path
	configFile := os.Getenv("JAYCONFIG")

	// Load the config
	config, err := storage.LoadConfig(configFile)
	if err != nil {
		log.Fatalf("%v", err)
	}

	// Set the database name
	config.MySQL.Database = dbName

	// Set the migration folder to the absolute path
	config.MySQL.Migration.Folder = filepath.Join(filepath.Dir(configFile), config.MySQL.Migration.Folder)

	// Create the migration configuration
	conf := Configuration{
		config.MySQL,
	}

	// Create the migration
	mig, err := conf.New()
	if err != nil {
		log.Fatalf("%v", err)
	}

	// Refresh the data with verbose logging
	err = mig.DownAll()
	if err != nil {
		log.Println(err)
	}

	err = mig.UpAll()
	if err != nil {
		log.Println(err)
	}

	return mig, conf
}

// TearDown removes the unit test database.
func TearDown(db *sqlx.DB, dbName string) error {
	// Drop the database
	_, err := db.Exec(fmt.Sprintf(`DROP DATABASE %v;`, dbName))
	return err
}
<|endoftext|>