<commit_before>package mesos

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/mesos/mesos-go/api/v0/upid"
	"github.com/nemosupremo/vault-gatekeeper/scheduler"
	"github.com/samuel/go-zookeeper/zk"
	"github.com/sirupsen/logrus"
)

type defaultLogger struct{}

func (d defaultLogger) Printf(s string, a ...interface{}) {
	logrus.Debugf(s, a...)
}

type mesosMaster struct {
	Address struct {
		Hostname string `json:"hostname"`
		Ip       string `json:"ip"`
		Port     int    `json:"port"`
	} `json:"address"`
	Hostname string `json:"hostname"`
	Id       string `json:"id"`
	Ip       int64  `json:"ip"`
	Pid      string `json:"pid"`
	Port     int    `json:"port"`
	Version  string `json:"version"`
}

type mesosTask struct {
	Id        string `json:"id"`
	Name      string `json:"name"`
	State     string `json:"state"`
	SlaveId   string `json:"slave_id"`
	Resources struct {
		Cpus  float64 `json:"cpus"`
		Disk  float64 `json:"disk"`
		Mem   float64 `json:"mem"`
		Ports string  `json:"ports"`
	} `json:"resources"`
	Statuses []struct {
		State     string  `json:"state"`
		Timestamp float64 `json:"timestamp"`
	} `json:"statuses"`
	Container struct {
		Type   string `json:"type"`
		Docker struct {
			Image string `json:"image"`
		} `json:"docker"`
	} `json:"container"`
}

type mesosState struct {
	GitSha string `json:"git_sha"`
	GitTag string `json:"git_tag"`
	Leader string `json:"leader"`
	Pid    string `json:"pid"`
	Slaves []struct {
		Id  string `json:"id"`
		PID string `json:"pid"`
	} `json:"slaves"`
	Frameworks []struct {
		Id     string      `json:"id"`
		Name   string      `json:"name"`
		Active bool        `json:"active"`
		Tasks  []mesosTask `json:"tasks"`
	} `json:"frameworks"`
}

var errMesosNoPath = errors.New("No path specified for mesos zk lookup.")
var errMesosParseError = errors.New("Error parsing mesos master data in zk.")
var errMesosNoMaster = errors.New("Error finding mesos master.")
var errUnknownScheme = errors.New("Unknown mesos scheme.")
var errMesosUnreachable = errors.New("No reachable mesos masters.")

type mesosScheduler struct {
	Master string
}

type task struct {
	id        string
	name      string
	group     string
	image     string
	ip        net.IP
	startTime time.Time
}

func (t task) Id() string {
	return t.id
}

func (t task) Group() string {
	return t.group
}

func (t task) Image() string {
	return t.image
}

func (t task) Name() string {
	return t.name
}

func (t task) IP() net.IP {
	return t.ip
}

func (t task) StartTime() time.Time {
	return t.startTime
}

func NewMesosScheduler(master string) (scheduler.Scheduler, error) {
	m := &mesosScheduler{master}
	if _, _, err := m.getMesosMaster(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *mesosScheduler) LookupTask(taskId string) (scheduler.Task, error) {
	mesosTask, framework, slaveHost, err := m.getMesosTask(taskId)
	for i := time.Duration(0); i < 3 && err == nil && len(mesosTask.Statuses) == 0; i++ {
		time.Sleep((500 + 250*i) * time.Millisecond)
		mesosTask, framework, slaveHost, err = m.getMesosTask(taskId)
	}
	runningTime := time.Unix(0, 0)
	if len(mesosTask.Statuses) > 0 {
		// https://github.com/apache/mesos/blob/a61074586d778d432ba991701c9c4de9459db897/src/webui/master/static/js/controllers.js#L148
		runningTime = time.Unix(0, int64(mesosTask.Statuses[0].Timestamp*1000000000))
	}

	var ip net.IP
	if err == nil {
		if ips, err := net.LookupIP(slaveHost); err == nil {
			ip = ips[0]
		}
	}

	return &task{
		id:        mesosTask.Id,
		name:      mesosTask.Name,
		image:     mesosTask.Container.Docker.Image,
		group:     framework,
		startTime: runningTime,
		ip:        ip,
	}, err
}

func (m *mesosScheduler) getMesosMaster() ([]string, string, error) {
	var masterHosts []string
	protocol := "http"
	if path, err := url.Parse(m.Master); err == nil {
		switch path.Scheme {
		case "zks":
			protocol = "https"
			fallthrough
		case "zk":
			if path.Path == "" || path.Path == "/" {
				return nil, protocol, errMesosNoPath
			}
			zookeeperPath := path.Path
			if zookeeperPath[0] != '/' {
				zookeeperPath = "/" + zookeeperPath
			}

			if zoo, _, err := zk.Connect(zk.FormatServers(strings.Split(path.Host, ",")), 10*time.Second, zk.WithLogger(&defaultLogger{})); err == nil {
				defer zoo.Close()
				if children, _, err := zoo.Children(zookeeperPath); err == nil {
					sort.Strings(children)
					for _, child := range children {
						if strings.HasPrefix(child, "json.info_") {
							if data, _, err := zoo.Get(zookeeperPath + "/" + child); err == nil {
								var masterInfo mesosMaster
								if err := json.Unmarshal(data, &masterInfo); err == nil {
									masterHosts = []string{fmt.Sprintf("%s:%d", masterInfo.Address.Hostname, masterInfo.Address.Port)}
									break
								} else {
									return nil, protocol, errMesosParseError

								}
							}
						}
					}
				} else {
					return nil, protocol, errMesosNoMaster
				}
			}
		case "https":
			protocol = "https"
			fallthrough
		case "http":
			masterHosts = strings.Split(path.Host, ",")
		default:
			return nil, protocol, errUnknownScheme
		}
	} else {
		masterHosts = strings.Split(m.Master, ",")
	}

	if len(masterHosts) == 0 {
		return nil, protocol, errMesosUnreachable
	}
	return masterHosts, protocol, nil
}

func (m *mesosScheduler) getMesosTask(taskId string) (mesosTask, string, string, error) {
	var state mesosState
	var masterErr error
	if masterHosts, protocol, err := m.getMesosMaster(); err == nil {
		for _, host := range masterHosts {
			if resp, err := http.Get(protocol + "://" + host + "/state.json"); err == nil {
				defer resp.Body.Close()
				if err := json.NewDecoder(resp.Body).Decode(&state); err == nil {
					if state.Pid == state.Leader {
						masterErr = nil
						break
					}
				} else {
					masterErr = err
				}
			} else {
				masterErr = err
			}
		}
		if masterErr != nil {
			return mesosTask{}, "", "", masterErr
		}
		if state.Pid != state.Leader {
			return mesosTask{}, "", "", errMesosUnreachable
		}

		slaves := make(map[string]string)
		for _, slave := range state.Slaves {
			slaves[slave.Id] = slave.PID
		}

		for _, framework := range state.Frameworks {
			for _, task := range framework.Tasks {
				if task.Id == taskId {
					slaveHost := ""
					if slavePid, ok := slaves[task.SlaveId]; ok {
						if pid, err := upid.Parse(slavePid); err == nil {
							slaveHost = pid.Host
						} else {
							logrus.Warnf("Mesos: Failed to parse PID %v.", slavePid)
						}
					} else {
						logrus.Warnf("Mesos: Task ID %v was running on Slave %v, but no information about that slave was found.", task.Id, task.SlaveId)
					}
					return task, framework.Name, slaveHost, nil
				}
			}
		}
		return mesosTask{}, "", "", scheduler.ErrTaskNotFound
	} else {
		return mesosTask{}, "", "", err
	}
}
<commit_msg>Mesos: state.json is removed in 1.8+ renamed /state. Fixes #75<commit_after>package mesos

import (
	"encoding/json"
	"errors"
	"fmt"
	"net"
	"net/http"
	"net/url"
	"sort"
	"strings"
	"time"

	"github.com/mesos/mesos-go/api/v0/upid"
	"github.com/nemosupremo/vault-gatekeeper/scheduler"
	"github.com/samuel/go-zookeeper/zk"
	"github.com/sirupsen/logrus"
)

type defaultLogger struct{}

func (d defaultLogger) Printf(s string, a ...interface{}) {
	logrus.Debugf(s, a...)
}

type mesosMaster struct {
	Address struct {
		Hostname string `json:"hostname"`
		Ip       string `json:"ip"`
		Port     int    `json:"port"`
	} `json:"address"`
	Hostname string `json:"hostname"`
	Id       string `json:"id"`
	Ip       int64  `json:"ip"`
	Pid      string `json:"pid"`
	Port     int    `json:"port"`
	Version  string `json:"version"`
}

type mesosTask struct {
	Id        string `json:"id"`
	Name      string `json:"name"`
	State     string `json:"state"`
	SlaveId   string `json:"slave_id"`
	Resources struct {
		Cpus  float64 `json:"cpus"`
		Disk  float64 `json:"disk"`
		Mem   float64 `json:"mem"`
		Ports string  `json:"ports"`
	} `json:"resources"`
	Statuses []struct {
		State     string  `json:"state"`
		Timestamp float64 `json:"timestamp"`
	} `json:"statuses"`
	Container struct {
		Type   string `json:"type"`
		Docker struct {
			Image string `json:"image"`
		} `json:"docker"`
	} `json:"container"`
}

type mesosState struct {
	GitSha string `json:"git_sha"`
	GitTag string `json:"git_tag"`
	Leader string `json:"leader"`
	Pid    string `json:"pid"`
	Slaves []struct {
		Id  string `json:"id"`
		PID string `json:"pid"`
	} `json:"slaves"`
	Frameworks []struct {
		Id     string      `json:"id"`
		Name   string      `json:"name"`
		Active bool        `json:"active"`
		Tasks  []mesosTask `json:"tasks"`
	} `json:"frameworks"`
}

var errMesosNoPath = errors.New("No path specified for mesos zk lookup.")
var errMesosParseError = errors.New("Error parsing mesos master data in zk.")
var errMesosNoMaster = errors.New("Error finding mesos master.")
var errUnknownScheme = errors.New("Unknown mesos scheme.")
var errMesosUnreachable = errors.New("No reachable mesos masters.")

type mesosScheduler struct {
	Master string
}

type task struct {
	id        string
	name      string
	group     string
	image     string
	ip        net.IP
	startTime time.Time
}

func (t task) Id() string {
	return t.id
}

func (t task) Group() string {
	return t.group
}

func (t task) Image() string {
	return t.image
}

func (t task) Name() string {
	return t.name
}

func (t task) IP() net.IP {
	return t.ip
}

func (t task) StartTime() time.Time {
	return t.startTime
}

func NewMesosScheduler(master string) (scheduler.Scheduler, error) {
	m := &mesosScheduler{master}
	if _, _, err := m.getMesosMaster(); err != nil {
		return nil, err
	}
	return m, nil
}

func (m *mesosScheduler) LookupTask(taskId string) (scheduler.Task, error) {
	mesosTask, framework, slaveHost, err := m.getMesosTask(taskId)
	for i := time.Duration(0); i < 3 && err == nil && len(mesosTask.Statuses) == 0; i++ {
		time.Sleep((500 + 250*i) * time.Millisecond)
		mesosTask, framework, slaveHost, err = m.getMesosTask(taskId)
	}
	runningTime := time.Unix(0, 0)
	if len(mesosTask.Statuses) > 0 {
		// https://github.com/apache/mesos/blob/a61074586d778d432ba991701c9c4de9459db897/src/webui/master/static/js/controllers.js#L148
		runningTime = time.Unix(0, int64(mesosTask.Statuses[0].Timestamp*1000000000))
	}

	var ip net.IP
	if err == nil {
		if ips, err := net.LookupIP(slaveHost); err == nil {
			ip = ips[0]
		}
	}

	return &task{
		id:        mesosTask.Id,
		name:      mesosTask.Name,
		image:     mesosTask.Container.Docker.Image,
		group:     framework,
		startTime: runningTime,
		ip:        ip,
	}, err
}

func (m *mesosScheduler) getMesosMaster() ([]string, string, error) {
	var masterHosts []string
	protocol := "http"
	if path, err := url.Parse(m.Master); err == nil {
		switch path.Scheme {
		case "zks":
			protocol = "https"
			fallthrough
		case "zk":
			if path.Path == "" || path.Path == "/" {
				return nil, protocol, errMesosNoPath
			}
			zookeeperPath := path.Path
			if zookeeperPath[0] != '/' {
				zookeeperPath = "/" + zookeeperPath
			}

			if zoo, _, err := zk.Connect(zk.FormatServers(strings.Split(path.Host, ",")), 10*time.Second, zk.WithLogger(&defaultLogger{})); err == nil {
				defer zoo.Close()
				if children, _, err := zoo.Children(zookeeperPath); err == nil {
					sort.Strings(children)
					for _, child := range children {
						if strings.HasPrefix(child, "json.info_") {
							if data, _, err := zoo.Get(zookeeperPath + "/" + child); err == nil {
								var masterInfo mesosMaster
								if err := json.Unmarshal(data, &masterInfo); err == nil {
									masterHosts = []string{fmt.Sprintf("%s:%d", masterInfo.Address.Hostname, masterInfo.Address.Port)}
									break
								} else {
									return nil, protocol, errMesosParseError

								}
							}
						}
					}
				} else {
					return nil, protocol, errMesosNoMaster
				}
			}
		case "https":
			protocol = "https"
			fallthrough
		case "http":
			masterHosts = strings.Split(path.Host, ",")
		default:
			return nil, protocol, errUnknownScheme
		}
	} else {
		masterHosts = strings.Split(m.Master, ",")
	}

	if len(masterHosts) == 0 {
		return nil, protocol, errMesosUnreachable
	}
	return masterHosts, protocol, nil
}

func (m *mesosScheduler) getMesosTask(taskId string) (mesosTask, string, string, error) {
	var state mesosState
	var masterErr error
	if masterHosts, protocol, err := m.getMesosMaster(); err == nil {
		for _, host := range masterHosts {
			if resp, err := http.Get(protocol + "://" + host + "/state"); err == nil {
				defer resp.Body.Close()
				if err := json.NewDecoder(resp.Body).Decode(&state); err == nil {
					if state.Pid == state.Leader {
						masterErr = nil
						break
					}
				} else {
					masterErr = err
				}
			} else {
				masterErr = err
			}
		}
		if masterErr != nil {
			return mesosTask{}, "", "", masterErr
		}
		if state.Pid != state.Leader {
			return mesosTask{}, "", "", errMesosUnreachable
		}

		slaves := make(map[string]string)
		for _, slave := range state.Slaves {
			slaves[slave.Id] = slave.PID
		}

		for _, framework := range state.Frameworks {
			for _, task := range framework.Tasks {
				if task.Id == taskId {
					slaveHost := ""
					if slavePid, ok := slaves[task.SlaveId]; ok {
						if pid, err := upid.Parse(slavePid); err == nil {
							slaveHost = pid.Host
						} else {
							logrus.Warnf("Mesos: Failed to parse PID %v.", slavePid)
						}
					} else {
						logrus.Warnf("Mesos: Task ID %v was running on Slave %v, but no information about that slave was found.", task.Id, task.SlaveId)
					}
					return task, framework.Name, slaveHost, nil
				}
			}
		}
		return mesosTask{}, "", "", scheduler.ErrTaskNotFound
	} else {
		return mesosTask{}, "", "", err
	}
}
<|endoftext|>


<commit_before>/*
adstop is an ad-blocking transparent HTTP/HTTPS proxy.

It was designed to run on low power, low memory ARM devices and serve a couple
of clients, mostly old smartphones which cannot run adblockers themselves.

Before using it, you have to configure your devices and network to make it
accessible as a transparent proxy. One way to achieve this is to install
a VPN on the server side and redirect all HTTP/HTTPS traffic to the proxy
with routing rules. Then make the client browse through the VPN.

HTTPS filtering requires the proxy to intercept the device traffic and decrypt
it. To allow this, you have to generate a certificate and add it to your
device.

You also need to generate a certificate

	$ adstop -http localhost:1080 \
		-https localhost:1081     \
		-cache .adstop			  \
		-max-age 24h			  \
		-ca-cert /path/to/ca.cert \
		-ca-key /path/to/ca.key   \
		https://easylist-downloads.adblockplus.org/easylist.txt \
		some_local_list.txt

starts the proxy and makes it listen on HTTP on port 1080, HTTPS on port 1081,
fetch and load rules from easylist and a local file, cache easylist in an
.adstop/ directory and refresh it every 24 hours.

*/
package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"strings"
	"sync"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/inconshreveable/go-vhost"
	"github.com/pmezard/adblock/adblock"
)

var (
	httpAddr  = flag.String("http", "localhost:1080", "HTTP handler address")
	httpsAddr = flag.String("https", "localhost:1081", "HTTPS handler address")
	httpDebug = flag.String("debug-addr", "", "HTTP debug address")
	logp      = flag.Bool("log", false, "enable logging")
	cacheDir  = flag.String("cache", ".cache", "cache directory")
	maxAgeArg = flag.String("max-age", "24h", "cached entries max age")
	caCert    = flag.String("ca-cert", "", "path to CA certificate")
	caKey     = flag.String("ca-key", "", "path to CA key")
)

type FilteringHandler struct {
	Cache *RuleCache
}

func logRequest(r *http.Request) {
	log.Printf("%s %s %s %s\n", r.Proto, r.Method, r.URL, r.Host)
	buf := &bytes.Buffer{}
	r.Header.Write(buf)
	log.Println(string(buf.Bytes()))
}

func getReferrerDomain(r *http.Request) string {
	ref := r.Header.Get("Referer")
	if len(ref) > 0 {
		u, err := url.Parse(ref)
		if err == nil {
			return u.Host
		}
	}
	return ""
}

type ProxyState struct {
	Duration time.Duration
	URL      string
}

func (h *FilteringHandler) OnRequest(r *http.Request, ctx *goproxy.ProxyCtx) (
	*http.Request, *http.Response) {

	if *logp {
		logRequest(r)
	}

	host := r.URL.Host
	if host == "" {
		host = r.Host
	}
	rq := &adblock.Request{
		URL:          r.URL.String(),
		Domain:       host,
		OriginDomain: getReferrerDomain(r),
	}
	rules := h.Cache.Rules()
	start := time.Now()
	matched, id := rules.Matcher.Match(rq)
	end := time.Now()
	duration := end.Sub(start) / time.Millisecond
	if matched {
		rule := rules.Rules[id]
		log.Printf("rejected in %dms: %s\n", duration, r.URL.String())
		log.Printf("  by %s\n", rule)
		return r, goproxy.NewResponse(r, goproxy.ContentTypeText,
			http.StatusNotFound, "Not Found")
	}
	ctx.UserData = &ProxyState{
		Duration: duration,
		URL:      r.URL.String(),
	}
	return r, nil
}

func (h *FilteringHandler) OnResponse(r *http.Response,
	ctx *goproxy.ProxyCtx) *http.Response {

	if r == nil {
		// Happens if RoundTrip fails
		return r
	}

	state, ok := ctx.UserData.(*ProxyState)
	if !ok {
		// The request was rejected by the previous handler
		return r
	}
	duration2 := time.Duration(0)
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err == nil && len(mediaType) > 0 {
		host := ctx.Req.URL.Host
		if host == "" {
			host = ctx.Req.Host
		}
		rq := &adblock.Request{
			URL:          ctx.Req.URL.String(),
			Domain:       host,
			OriginDomain: getReferrerDomain(ctx.Req),
			ContentType:  mediaType,
		}
		// Second level filtering, based on returned content
		rules := h.Cache.Rules()
		start := time.Now()
		matched, id := rules.Matcher.Match(rq)
		end := time.Now()
		duration2 = end.Sub(start) / time.Millisecond
		if matched {
			r.Body.Close()
			rule := rules.Rules[id]
			log.Printf("rejected in %d/%dms: %s\n", state.Duration, duration2,
				state.URL)
			log.Printf("  by %s\n", rule)
			return goproxy.NewResponse(ctx.Req, goproxy.ContentTypeText,
				http.StatusNotFound, "Not Found")
		}
	}
	log.Printf("accepted in %d/%dms: %s\n", state.Duration, duration2, state.URL)
	return r
}

// CachedConfig holds a TLS configuration. It can be in different states:
// - The config is being generated, Config is nil and the Ready channel is set
// - The config is ready, Config is not nil and Ready is closed.
// This mechanism is used to pool concurrent generations of the same certificate.
type CachedConfig struct {
	Config *tls.Config
	Ready  chan struct{}
}

// TLSConfigCache is a goroutine-safe cache of TLS configurations mapped to hosts.
type TLSConfigCache struct {
	cfgBuilder func(string, *goproxy.ProxyCtx) (*tls.Config, error)
	lock       sync.Mutex
	cache      map[string]CachedConfig
	hit        int
	miss       int
}

func NewTLSConfigCache(ca *tls.Certificate) *TLSConfigCache {
	return &TLSConfigCache{
		cfgBuilder: goproxy.TLSConfigFromCA(ca),
		cache:      map[string]CachedConfig{},
	}
}

func getWildcardHost(host string) string {
	first := strings.Index(host, ".")
	if first <= 0 {
		return host
	}
	last := strings.LastIndex(host, ".")
	if last == first {
		// root domain, no wildcard
		return host
	}
	return "*" + host[first:]
}

func (c *TLSConfigCache) GetConfig(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
	host = getWildcardHost(host)
	c.lock.Lock()
	cached, ok := c.cache[host]
	if !ok {
		// Register a config generation event
		cached = CachedConfig{
			Ready: make(chan struct{}),
		}
		c.cache[host] = cached
	}
	if ok {
		c.hit += 1
	} else {
		c.miss += 1
	}
	hit := c.hit
	miss := c.miss
	c.lock.Unlock()

	ctx.Warnf("signing hit/miss: %d/%d (%.1f%%)", hit, miss,
		100.0*float64(hit)/float64(hit+miss))
	if ok {
		// config is being generated or is ready, grab it
		<-cached.Ready
		cfg := cached.Config
		if cfg == nil {
			return nil, fmt.Errorf("failed to generate TLS config for %s", host)
		}
		return cfg, nil
	}

	// Generate it
	start := time.Now()
	cfg, err := c.cfgBuilder(host, ctx)
	stop := time.Now()
	ctx.Warnf("signing %s in %.0fms", host,
		float64(stop.Sub(start))/float64(time.Millisecond))

	c.lock.Lock()
	if err == nil {
		c.cache[host] = CachedConfig{
			Config: cfg,
			Ready:  cached.Ready,
		}
	} else {
		delete(c.cache, host)
		ctx.Warnf("failed to sign %s: %s", host, err)
	}
	close(cached.Ready)
	c.lock.Unlock()
	return cfg, err
}

// copied/converted from https.go
type dumbResponseWriter struct {
	net.Conn
}

func (dumb dumbResponseWriter) Header() http.Header {
	panic("Header() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Write(buf []byte) (int, error) {
	if bytes.Equal(buf, []byte("HTTP/1.0 200 OK\r\n\r\n")) {
		// throw away the HTTP OK response from the faux CONNECT request
		return len(buf), nil
	}
	return dumb.Conn.Write(buf)
}

func (dumb dumbResponseWriter) WriteHeader(code int) {
	panic("WriteHeader() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return dumb, bufio.NewReadWriter(bufio.NewReader(dumb), bufio.NewWriter(dumb)), nil
}

func makeCertificate(certPath, keyPath string) (*tls.Certificate, error) {
	cert, err := ioutil.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA certificate: %s", err)
	}
	key, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA key: %s", err)
	}
	ca, err := tls.X509KeyPair(cert, key)
	return &ca, err
}

func runProxy() error {
	flag.Parse()
	if *caCert == "" || *caKey == "" {
		return fmt.Errorf("CA certificate and key must be specified")
	}
	ca, err := makeCertificate(*caCert, *caKey)
	if err != nil {
		return err
	}

	maxAge, err := time.ParseDuration(*maxAgeArg)
	if err != nil {
		return fmt.Errorf("invalid max-age: %s", err)
	}
	if maxAge < 0 {
		return fmt.Errorf("invalid negative max-age")
	}
	log.Printf("loading rules")
	cache, err := NewRuleCache(*cacheDir, flag.Args(), maxAge)
	if err != nil {
		return err
	}
	h := &FilteringHandler{
		Cache: cache,
	}

	if *httpDebug != "" {
		log.Printf("starting debug server on %s", *httpDebug)
		go func() {
			log.Println(http.ListenAndServe(*httpDebug, nil))
		}()
	}

	log.Printf("starting servers")
	proxy := goproxy.NewProxyHttpServer()
	proxy.NonproxyHandler = http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			if req.Host == "" {
				log.Printf("Cannot handle requests without Host header, e.g., HTTP 1.0")
				return
			}
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			proxy.ServeHTTP(w, req)
		})

	// Cache MITM certificates
	tlsCache := NewTLSConfigCache(ca)
	MitmConnect := &goproxy.ConnectAction{
		Action: goproxy.ConnectMitm,
		TLSConfig: func(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
			return tlsCache.GetConfig(host, ctx)
		},
	}
	var AlwaysMitm goproxy.FuncHttpsHandler = func(host string, ctx *goproxy.ProxyCtx) (
		*goproxy.ConnectAction, string) {

		return MitmConnect, host
	}
	proxy.OnRequest().HandleConnect(AlwaysMitm)

	proxy.OnRequest().DoFunc(h.OnRequest)
	proxy.OnResponse().DoFunc(h.OnResponse)
	go func() {
		http.ListenAndServe(*httpAddr, proxy)
	}()

	// listen to the TLS ClientHello but make it a CONNECT request instead
	ln, err := net.Listen("tcp", *httpsAddr)
	if err != nil {
		return err
	}
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting new connection - %v", err)
			continue
		}
		go func(c net.Conn) {
			tlsConn, err := vhost.TLS(c)
			if err != nil {
				log.Printf("error accepting new connection - %v", err)
			}
			if tlsConn.Host() == "" {
				log.Printf("cannot support non-SNI enabled clients")
				return
			}
			connectReq := &http.Request{
				Method: "CONNECT",
				URL: &url.URL{
					Opaque: tlsConn.Host(),
					Host:   net.JoinHostPort(tlsConn.Host(), "443"),
				},
				Host:   tlsConn.Host(),
				Header: make(http.Header),
			}
			resp := dumbResponseWriter{tlsConn}
			proxy.ServeHTTP(resp, connectReq)
		}(c)
	}
}

func main() {
	err := runProxy()
	if err != nil {
		log.Fatalf("error: %s\n", err)
	}
}
<commit_msg>adstop: add /trace endpoint to toggle headers logging<commit_after>/*
adstop is an ad-blocking transparent HTTP/HTTPS proxy.

It was designed to run on low power, low memory ARM devices and serve a couple
of clients, mostly old smartphones which cannot run adblockers themselves.

Before using it, you have to configure your devices and network to make it
accessible as a transparent proxy. One way to achieve this is to install
a VPN on the server side and redirect all HTTP/HTTPS traffic to the proxy
with routing rules. Then make the client browse through the VPN.

HTTPS filtering requires the proxy to intercept the device traffic and decrypt
it. To allow this, you have to generate a certificate and add it to your
device.

You also need to generate a certificate

	$ adstop -http localhost:1080 \
		-https localhost:1081     \
		-cache .adstop			  \
		-max-age 24h			  \
		-ca-cert /path/to/ca.cert \
		-ca-key /path/to/ca.key   \
		https://easylist-downloads.adblockplus.org/easylist.txt \
		some_local_list.txt

starts the proxy and makes it listen on HTTP on port 1080, HTTPS on port 1081,
fetch and load rules from easylist and a local file, cache easylist in an
.adstop/ directory and refresh it every 24 hours.

*/
package main

import (
	"bufio"
	"bytes"
	"crypto/tls"
	"flag"
	"fmt"
	"io/ioutil"
	"log"
	"mime"
	"net"
	"net/http"
	_ "net/http/pprof"
	"net/url"
	"strings"
	"sync"
	"sync/atomic"
	"time"

	"github.com/elazarl/goproxy"
	"github.com/inconshreveable/go-vhost"
	"github.com/pmezard/adblock/adblock"
)

var (
	httpAddr    = flag.String("http", "localhost:1080", "HTTP handler address")
	httpsAddr   = flag.String("https", "localhost:1081", "HTTPS handler address")
	httpDebug   = flag.String("debug-addr", "", "HTTP debug address")
	logRequests = flag.Uint64("log", 0, "enable logging")
	cacheDir    = flag.String("cache", ".cache", "cache directory")
	maxAgeArg   = flag.String("max-age", "24h", "cached entries max age")
	caCert      = flag.String("ca-cert", "", "path to CA certificate")
	caKey       = flag.String("ca-key", "", "path to CA key")
)

type FilteringHandler struct {
	Cache *RuleCache
}

func logRequest(r *http.Request) {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, "REQ\n<REQUEST\n%s %s %s %s\n", r.Proto, r.Method, r.URL, r.Host)
	r.Header.Write(buf)
	fmt.Fprintf(buf, "REQUEST>\n")
	log.Println(buf.String())
}

func logResponse(r *http.Response) {
	buf := &bytes.Buffer{}
	fmt.Fprintf(buf, "RSP\n<RESPONSE\n%s %s\n", r.Proto, r.Status)
	r.Header.Write(buf)
	fmt.Fprintf(buf, "RESPONSE>\n")
	log.Println(buf.String())
}

func getReferrerDomain(r *http.Request) string {
	ref := r.Header.Get("Referer")
	if len(ref) > 0 {
		u, err := url.Parse(ref)
		if err == nil {
			return u.Host
		}
	}
	return ""
}

type ProxyState struct {
	Duration time.Duration
	URL      string
}

func (h *FilteringHandler) OnRequest(r *http.Request, ctx *goproxy.ProxyCtx) (
	*http.Request, *http.Response) {

	if atomic.LoadUint64(logRequests)%2 == 1 {
		logRequest(r)
	}

	host := r.URL.Host
	if host == "" {
		host = r.Host
	}
	rq := &adblock.Request{
		URL:          r.URL.String(),
		Domain:       host,
		OriginDomain: getReferrerDomain(r),
	}
	rules := h.Cache.Rules()
	start := time.Now()
	matched, id := rules.Matcher.Match(rq)
	end := time.Now()
	duration := end.Sub(start) / time.Millisecond
	if matched {
		rule := rules.Rules[id]
		log.Printf("rejected in %dms: %s\n", duration, r.URL.String())
		log.Printf("  by %s\n", rule)
		return r, goproxy.NewResponse(r, goproxy.ContentTypeText,
			http.StatusNotFound, "Not Found")
	}
	ctx.UserData = &ProxyState{
		Duration: duration,
		URL:      r.URL.String(),
	}
	return r, nil
}

func (h *FilteringHandler) OnResponse(r *http.Response,
	ctx *goproxy.ProxyCtx) *http.Response {

	if r == nil {
		// Happens if RoundTrip fails
		return r
	}

	state, ok := ctx.UserData.(*ProxyState)
	if !ok {
		// The request was rejected by the previous handler
		return r
	}

	if atomic.LoadUint64(logRequests)%2 == 1 {
		logResponse(r)
	}

	duration2 := time.Duration(0)
	mediaType, _, err := mime.ParseMediaType(r.Header.Get("Content-Type"))
	if err == nil && len(mediaType) > 0 {
		host := ctx.Req.URL.Host
		if host == "" {
			host = ctx.Req.Host
		}
		rq := &adblock.Request{
			URL:          ctx.Req.URL.String(),
			Domain:       host,
			OriginDomain: getReferrerDomain(ctx.Req),
			ContentType:  mediaType,
		}
		// Second level filtering, based on returned content
		rules := h.Cache.Rules()
		start := time.Now()
		matched, id := rules.Matcher.Match(rq)
		end := time.Now()
		duration2 = end.Sub(start) / time.Millisecond
		if matched {
			r.Body.Close()
			rule := rules.Rules[id]
			log.Printf("rejected in %d/%dms: %s\n", state.Duration, duration2,
				state.URL)
			log.Printf("  by %s\n", rule)
			return goproxy.NewResponse(ctx.Req, goproxy.ContentTypeText,
				http.StatusNotFound, "Not Found")
		}
	}
	log.Printf("accepted in %d/%dms: %s\n", state.Duration, duration2, state.URL)
	return r
}

// CachedConfig holds a TLS configuration. It can be in different states:
// - The config is being generated, Config is nil and the Ready channel is set
// - The config is ready, Config is not nil and Ready is closed.
// This mechanism is used to pool concurrent generations of the same certificate.
type CachedConfig struct {
	Config *tls.Config
	Ready  chan struct{}
}

// TLSConfigCache is a goroutine-safe cache of TLS configurations mapped to hosts.
type TLSConfigCache struct {
	cfgBuilder func(string, *goproxy.ProxyCtx) (*tls.Config, error)
	lock       sync.Mutex
	cache      map[string]CachedConfig
	hit        int
	miss       int
}

func NewTLSConfigCache(ca *tls.Certificate) *TLSConfigCache {
	return &TLSConfigCache{
		cfgBuilder: goproxy.TLSConfigFromCA(ca),
		cache:      map[string]CachedConfig{},
	}
}

func getWildcardHost(host string) string {
	first := strings.Index(host, ".")
	if first <= 0 {
		return host
	}
	last := strings.LastIndex(host, ".")
	if last == first {
		// root domain, no wildcard
		return host
	}
	return "*" + host[first:]
}

func (c *TLSConfigCache) GetConfig(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
	host = getWildcardHost(host)
	c.lock.Lock()
	cached, ok := c.cache[host]
	if !ok {
		// Register a config generation event
		cached = CachedConfig{
			Ready: make(chan struct{}),
		}
		c.cache[host] = cached
	}
	if ok {
		c.hit += 1
	} else {
		c.miss += 1
	}
	hit := c.hit
	miss := c.miss
	c.lock.Unlock()

	ctx.Warnf("signing hit/miss: %d/%d (%.1f%%)", hit, miss,
		100.0*float64(hit)/float64(hit+miss))
	if ok {
		// config is being generated or is ready, grab it
		<-cached.Ready
		cfg := cached.Config
		if cfg == nil {
			return nil, fmt.Errorf("failed to generate TLS config for %s", host)
		}
		return cfg, nil
	}

	// Generate it
	start := time.Now()
	cfg, err := c.cfgBuilder(host, ctx)
	stop := time.Now()
	ctx.Warnf("signing %s in %.0fms", host,
		float64(stop.Sub(start))/float64(time.Millisecond))

	c.lock.Lock()
	if err == nil {
		c.cache[host] = CachedConfig{
			Config: cfg,
			Ready:  cached.Ready,
		}
	} else {
		delete(c.cache, host)
		ctx.Warnf("failed to sign %s: %s", host, err)
	}
	close(cached.Ready)
	c.lock.Unlock()
	return cfg, err
}

// copied/converted from https.go
type dumbResponseWriter struct {
	net.Conn
}

func (dumb dumbResponseWriter) Header() http.Header {
	panic("Header() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Write(buf []byte) (int, error) {
	if bytes.Equal(buf, []byte("HTTP/1.0 200 OK\r\n\r\n")) {
		// throw away the HTTP OK response from the faux CONNECT request
		return len(buf), nil
	}
	return dumb.Conn.Write(buf)
}

func (dumb dumbResponseWriter) WriteHeader(code int) {
	panic("WriteHeader() should not be called on this ResponseWriter")
}

func (dumb dumbResponseWriter) Hijack() (net.Conn, *bufio.ReadWriter, error) {
	return dumb, bufio.NewReadWriter(bufio.NewReader(dumb), bufio.NewWriter(dumb)), nil
}

func makeCertificate(certPath, keyPath string) (*tls.Certificate, error) {
	cert, err := ioutil.ReadFile(certPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA certificate: %s", err)
	}
	key, err := ioutil.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("cannot load CA key: %s", err)
	}
	ca, err := tls.X509KeyPair(cert, key)
	return &ca, err
}

func runDebugServer(addr string) error {
	http.HandleFunc("/trace", func(w http.ResponseWriter, r *http.Request) {
		res := atomic.AddUint64(logRequests, 1)
		action := "logging requests"
		if res%2 == 0 {
			action = "ignoring requests"
		}
		fmt.Fprintf(w, "%s\n", action)
	})
	return http.ListenAndServe(addr, nil)
}

func runProxy() error {
	flag.Parse()
	if *caCert == "" || *caKey == "" {
		return fmt.Errorf("CA certificate and key must be specified")
	}
	ca, err := makeCertificate(*caCert, *caKey)
	if err != nil {
		return err
	}

	maxAge, err := time.ParseDuration(*maxAgeArg)
	if err != nil {
		return fmt.Errorf("invalid max-age: %s", err)
	}
	if maxAge < 0 {
		return fmt.Errorf("invalid negative max-age")
	}
	log.Printf("loading rules")
	cache, err := NewRuleCache(*cacheDir, flag.Args(), maxAge)
	if err != nil {
		return err
	}
	h := &FilteringHandler{
		Cache: cache,
	}

	if *httpDebug != "" {
		log.Printf("starting debug server on %s", *httpDebug)
		go func() {
			log.Println(runDebugServer(*httpDebug))
		}()
	}

	log.Printf("starting servers")
	proxy := goproxy.NewProxyHttpServer()
	proxy.NonproxyHandler = http.HandlerFunc(
		func(w http.ResponseWriter, req *http.Request) {
			if req.Host == "" {
				log.Printf("Cannot handle requests without Host header, e.g., HTTP 1.0")
				return
			}
			req.URL.Scheme = "http"
			req.URL.Host = req.Host
			proxy.ServeHTTP(w, req)
		})

	// Cache MITM certificates
	tlsCache := NewTLSConfigCache(ca)
	MitmConnect := &goproxy.ConnectAction{
		Action: goproxy.ConnectMitm,
		TLSConfig: func(host string, ctx *goproxy.ProxyCtx) (*tls.Config, error) {
			return tlsCache.GetConfig(host, ctx)
		},
	}
	var AlwaysMitm goproxy.FuncHttpsHandler = func(host string, ctx *goproxy.ProxyCtx) (
		*goproxy.ConnectAction, string) {

		return MitmConnect, host
	}
	proxy.OnRequest().HandleConnect(AlwaysMitm)

	proxy.OnRequest().DoFunc(h.OnRequest)
	proxy.OnResponse().DoFunc(h.OnResponse)
	go func() {
		http.ListenAndServe(*httpAddr, proxy)
	}()

	// listen to the TLS ClientHello but make it a CONNECT request instead
	ln, err := net.Listen("tcp", *httpsAddr)
	if err != nil {
		return err
	}
	for {
		c, err := ln.Accept()
		if err != nil {
			log.Printf("error accepting new connection - %v", err)
			continue
		}
		go func(c net.Conn) {
			tlsConn, err := vhost.TLS(c)
			if err != nil {
				log.Printf("error accepting new connection - %v", err)
			}
			if tlsConn.Host() == "" {
				log.Printf("cannot support non-SNI enabled clients")
				return
			}
			connectReq := &http.Request{
				Method: "CONNECT",
				URL: &url.URL{
					Opaque: tlsConn.Host(),
					Host:   net.JoinHostPort(tlsConn.Host(), "443"),
				},
				Host:   tlsConn.Host(),
				Header: make(http.Header),
			}
			resp := dumbResponseWriter{tlsConn}
			proxy.ServeHTTP(resp, connectReq)
		}(c)
	}
}

func main() {
	err := runProxy()
	if err != nil {
		log.Fatalf("error: %s\n", err)
	}
}
<|endoftext|>


<commit_before>package agent

import (
	"net"
	"os/exec"

	"github.com/pkg/errors"

	"github.com/havoc-io/mutagen"
	processpkg "github.com/havoc-io/mutagen/process"
	"github.com/havoc-io/mutagen/url"
)

const (
	errorCodeCommandNotFound = 127
)

type sshConn struct {
	net.Conn
	process *exec.Cmd
}

func connectSSH(prompter string, remote *url.SSHURL) (net.Conn, bool, error) {
	// Create an SSH process.
	process, err := ssh(prompter, "Connecting to agent", remote, agentSSHCommand)
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to create SSH command")
	}

	// Create pipes to the process.
	stdin, err := process.StdinPipe()
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to redirect SSH input")
	}
	stdout, err := process.StdoutPipe()
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to redirect SSH output")
	}

	// Start the process.
	if err = process.Start(); err != nil {
		return nil, false, errors.Wrap(err, "unable to start SSH process")
	}

	// Confirm that the process started correctly by performing a version
	// handshake.
	if versionMatch, err := mutagen.ReceiveAndCompareVersion(stdout); err != nil {
		code, codeErr := processpkg.ExitCodeForError(process.Wait())
		if codeErr != nil && code == errorCodeCommandNotFound {
			return nil, true, errors.New("command not found")
		} else {
			return nil, false, errors.Wrap(err, "unable to handshake with SSH process")
		}
	} else if !versionMatch {
		return nil, true, errors.New("version mismatch")
	}

	// Create a connection.
	// HACK: TODO: Document why we don't register stdout as a closer.
	conn, _ := NewIOConn(stdout, stdin, stdin)
	return &sshConn{conn, process}, false, nil
}
<commit_msg>Added small TODO for agent connection portability.<commit_after>package agent

import (
	"net"
	"os/exec"

	"github.com/pkg/errors"

	"github.com/havoc-io/mutagen"
	processpkg "github.com/havoc-io/mutagen/process"
	"github.com/havoc-io/mutagen/url"
)

const (
	errorCodeCommandNotFound = 127
)

type sshConn struct {
	net.Conn
	process *exec.Cmd
}

func connectSSH(prompter string, remote *url.SSHURL) (net.Conn, bool, error) {
	// Create an SSH process.
	process, err := ssh(prompter, "Connecting to agent", remote, agentSSHCommand)
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to create SSH command")
	}

	// Create pipes to the process.
	stdin, err := process.StdinPipe()
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to redirect SSH input")
	}
	stdout, err := process.StdoutPipe()
	if err != nil {
		return nil, false, errors.Wrap(err, "unable to redirect SSH output")
	}

	// Start the process.
	if err = process.Start(); err != nil {
		return nil, false, errors.Wrap(err, "unable to start SSH process")
	}

	// Confirm that the process started correctly by performing a version
	// handshake.
	// TODO: Figure out how to identify "command not found" errors for Windows
	// SSH servers.
	if versionMatch, err := mutagen.ReceiveAndCompareVersion(stdout); err != nil {
		code, codeErr := processpkg.ExitCodeForError(process.Wait())
		if codeErr != nil && code == errorCodeCommandNotFound {
			return nil, true, errors.New("command not found")
		} else {
			return nil, false, errors.Wrap(err, "unable to handshake with SSH process")
		}
	} else if !versionMatch {
		return nil, true, errors.New("version mismatch")
	}

	// Create a connection.
	// HACK: TODO: Document why we don't register stdout as a closer.
	conn, _ := NewIOConn(stdout, stdin, stdin)
	return &sshConn{conn, process}, false, nil
}
<|endoftext|>


<commit_before>package agent

import "runtime"

// You can overridden buildVersion at compile time by using:
//
//  go run -ldflags "-X github.com/buildkite/agent/v3/agent.buildVersion abc" *.go --version
//
// On CI, the binaries are always build with the buildVersion variable set.
//
// Pre-release builds' versions must be in the format `x.y-beta`, `x.y-beta.z` or `x.y-beta.z.a`

var baseVersion string = "3.28.1"
var buildVersion string = ""

func Version() string {
	return baseVersion
}

func BuildVersion() string {
	if buildVersion != "" {
		return buildVersion
	} else {
		return "x"
	}
}

func UserAgent() string {
	return "buildkite-agent/" + Version() + "." + BuildVersion() + " (" + runtime.GOOS + "; " + runtime.GOARCH + ")"
}
<commit_msg>bump version to 3.29.0-beta.1<commit_after>package agent

import "runtime"

// You can overridden buildVersion at compile time by using:
//
//  go run -ldflags "-X github.com/buildkite/agent/v3/agent.buildVersion abc" *.go --version
//
// On CI, the binaries are always build with the buildVersion variable set.
//
// Pre-release builds' versions must be in the format `x.y-beta`, `x.y-beta.z` or `x.y-beta.z.a`

var baseVersion string = "3.29.0-beta.1"
var buildVersion string = ""

func Version() string {
	return baseVersion
}

func BuildVersion() string {
	if buildVersion != "" {
		return buildVersion
	} else {
		return "x"
	}
}

func UserAgent() string {
	return "buildkite-agent/" + Version() + "." + BuildVersion() + " (" + runtime.GOOS + "; " + runtime.GOARCH + ")"
}
<|endoftext|>


<commit_before>package bagman_test

import (
	"github.com/APTrust/bagman/bagman"
	"os"
	"path/filepath"
	"testing"
)

func partnerS3ConfigFile() (string) {
	configFile, _ := bagman.RelativeToAbsPath(
		filepath.Join("testdata", "partner_config_integration_test.conf"))
	return configFile
}

func partnerConfigForTest() (*bagman.PartnerConfig) {
	config := &bagman.PartnerConfig{
		AwsAccessKeyId: "abc",
		AwsSecretAccessKey: "xyz",
		ReceivingBucket: "aptrust.receiving.xyz.edu",
		RestorationBucket: "aptrust.receiving.xyz.edu",
		DownloadDir: "~/tmp",
	}
	config.ExpandFilePaths()
	return config
}

func TestNewPartnerS3ClientFromConfigFile(t *testing.T) {
	// This test will fail if AWS keys are not set in the environment,
	// because they are not set in the config file.
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	_, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
}

func TestNewPartnerS3ClientWithConfig(t *testing.T) {
	partnerConfig := partnerConfigForTest()
 	_, err := bagman.NewPartnerS3ClientWithConfig(partnerConfig, false)
	if err != nil {
		t.Error(err)
	}
}

func TestPartnerS3ClientLoadConfig(t *testing.T) {
	// This test will fail if AWS keys are not set in the environment,
	// because they are not set in the config file.
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
	partnerConfig := partnerConfigForTest()
 	client, err := bagman.NewPartnerS3ClientWithConfig(partnerConfig, false)
	if err != nil {
		t.Error(err)
	}
	// Load a new config from a file
	err = client.LoadConfig(partnerS3ConfigFile())
	if err != nil {
		t.Error(err)
	}
}

func TestPartnerS3ClientUploadListAndDelete(t *testing.T) {
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	client, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
	client.Test = true // turn off output
	tarFile, _ := bagman.RelativeToAbsPath(
		filepath.Join("testdata", "example.edu.sample_good.tar"))
	file, err := os.Open(tarFile)
	if err != nil {
		t.Error(err)
	}

	// Test upload
	md5, err := client.UploadFile(file)
	if err != nil {
		t.Error(err)
	}
	if md5 != "05e68e69767c772d36bd8a2baf693428" {
		t.Errorf("Expected md5 sum '05e68e69767c772d36bd8a2baf693428', got '%s'", md5)
	}

	// Test list. There should be only a handful of files in this test bucket.
	keys, err := client.List(client.PartnerConfig.ReceivingBucket, 500)
	if err != nil {
		t.Errorf("Error trying to list receiving bucket: %v", err)
	}
	foundFile := false
	for _, key := range keys {
		if key.Key == "example.edu.sample_good.tar" {
			foundFile = true
			break
		}
	}
	if !foundFile {
		t.Errorf("Uploaded file 'example.edu.sample_good.tar' not found in test bucket.")
	}

	// Test delete
	err = client.Delete(client.PartnerConfig.ReceivingBucket, "example.edu.sample_good.tar")
	if err != nil {
		t.Errorf("Error deleting example.edu.sample_good.tar from S3: %v", err)
	}
}

func TestPartnerS3ClientUploadFiles(t *testing.T) {
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	client, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
	client.Test = true // turn off output
	file0, _ := bagman.RelativeToAbsPath(filepath.Join("testdata", "example.edu.multipart.b01.of02.tar"))
	file1, _ := bagman.RelativeToAbsPath(filepath.Join("testdata", "example.edu.multipart.b02.of02.tar"))
	files := make([]string, 2)
	files[0] = file0
	files[1] = file1
	succeeded, failed := client.UploadFiles(files)
	if succeeded != 2 {
		t.Errorf("Expected 2 files to have uploaded, but %d actually succeeded", succeeded)
	}
	if failed != 0 {
		t.Errorf("%d files failed to upload")
	}
}

func TestPartnerS3ClientDownloadFile(t *testing.T) {
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	client, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
	client.Test = true // turn off output

	// Download with no checksum
	checksum, err := client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "none")
	if err != nil {
		t.Error(err)
	}
	if checksum != "" {
		t.Errorf("DownloadFile should have returned an empty checksum")
	}

	// Download with md5 - our test file is about 9.5kb
	checksum, err = client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "md5")
	if err != nil {
		t.Error(err)
	}
	if checksum != "3f43304f3f7c8d51111d0846e13cb74e" {
		t.Errorf("Expected md5 '3f43304f3f7c8d51111d0846e13cb74e', got '%s'", checksum)
	}

	// Download with sha256 - our test file is about 9.5kb
	checksum, err = client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "sha256")
	if err != nil {
		t.Error(err)
	}
	if checksum != "6f87f5341df2558967da27b12939d5e88b3af06592104041e57043af150f9309" {
		t.Errorf("Expected sha256 '6f87f5341df2558967da27b12939d5e88b3af06592104041e57043af150f9309', got '%s'", checksum)
	}

	checksum, err = client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "invalid_algo")
	if err == nil {
		t.Error("DownloadFile should have returned error for invalid checksum algorithm")
	}

}
<commit_msg>Updated partners3client_test<commit_after>package bagman_test

import (
	"github.com/APTrust/bagman/bagman"
	"os"
	"path/filepath"
	"testing"
)

func partnerS3ConfigFile() (string) {
	configFile, _ := bagman.RelativeToAbsPath(
		filepath.Join("testdata", "partner_config_integration_test.conf"))
	return configFile
}

func partnerConfigForTest() (*bagman.PartnerConfig) {
	config := &bagman.PartnerConfig{
		AwsAccessKeyId: "abc",
		AwsSecretAccessKey: "xyz",
		ReceivingBucket: "aptrust.receiving.xyz.edu",
		RestorationBucket: "aptrust.receiving.xyz.edu",
		DownloadDir: "~/tmp",
	}
	config.ExpandFilePaths()
	return config
}

func TestNewPartnerS3ClientFromConfigFile(t *testing.T) {
	// This test will fail if AWS keys are not set in the environment,
	// because they are not set in the config file.
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	_, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
}

func TestNewPartnerS3ClientWithConfig(t *testing.T) {
	partnerConfig := partnerConfigForTest()
 	_, err := bagman.NewPartnerS3ClientWithConfig(partnerConfig, false)
	if err != nil {
		t.Error(err)
	}
}

func TestPartnerS3ClientLoadConfig(t *testing.T) {
	// This test will fail if AWS keys are not set in the environment,
	// because they are not set in the config file.
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
	partnerConfig := partnerConfigForTest()
 	client, err := bagman.NewPartnerS3ClientWithConfig(partnerConfig, false)
	if err != nil {
		t.Error(err)
	}
	// Load a new config from a file
	err = client.LoadConfig(partnerS3ConfigFile())
	if err != nil {
		t.Error(err)
	}
}

func TestPartnerS3ClientUploadListAndDelete(t *testing.T) {
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	client, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
	client.Test = true // turn off output
	tarFile, _ := bagman.RelativeToAbsPath(
		filepath.Join("testdata", "example.edu.sample_good.tar"))
	file, err := os.Open(tarFile)
	if err != nil {
		t.Error(err)
	}

	// Test upload
	err = client.UploadFile(file)
	if err != nil {
		t.Error(err)
	}

	// Test list. There should be only a handful of files in this test bucket.
	keys, err := client.List(client.PartnerConfig.ReceivingBucket, 500)
	if err != nil {
		t.Errorf("Error trying to list receiving bucket: %v", err)
	}
	foundFile := false
	for _, key := range keys {
		if key.Key == "example.edu.sample_good.tar" {
			foundFile = true
			break
		}
	}
	if !foundFile {
		t.Errorf("Uploaded file 'example.edu.sample_good.tar' not found in test bucket.")
	}

	// Test delete
	err = client.Delete(client.PartnerConfig.ReceivingBucket, "example.edu.sample_good.tar")
	if err != nil {
		t.Errorf("Error deleting example.edu.sample_good.tar from S3: %v", err)
	}
}

func TestPartnerS3ClientUploadFiles(t *testing.T) {
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	client, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
	client.Test = true // turn off output
	file0, _ := bagman.RelativeToAbsPath(filepath.Join("testdata", "example.edu.multipart.b01.of02.tar"))
	file1, _ := bagman.RelativeToAbsPath(filepath.Join("testdata", "example.edu.multipart.b02.of02.tar"))
	files := make([]string, 2)
	files[0] = file0
	files[1] = file1
	succeeded, failed := client.UploadFiles(files)
	if succeeded != 2 {
		t.Errorf("Expected 2 files to have uploaded, but %d actually succeeded", succeeded)
	}
	if failed != 0 {
		t.Errorf("%d files failed to upload")
	}
}

func TestPartnerS3ClientDownloadFile(t *testing.T) {
	if os.Getenv("AWS_ACCESS_KEY_ID") == "" || os.Getenv("AWS_SECRET_ACCESS_KEY") == "" {
		return
	}
 	client, err := bagman.NewPartnerS3ClientFromConfigFile(partnerS3ConfigFile(), false)
	if err != nil {
		t.Error(err)
	}
	client.Test = true // turn off output

	// Download with no checksum
	checksum, err := client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "none")
	if err != nil {
		t.Error(err)
	}
	if checksum != "" {
		t.Errorf("DownloadFile should have returned an empty checksum")
	}

	// Download with md5 - our test file is about 9.5kb
	checksum, err = client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "md5")
	if err != nil {
		t.Error(err)
	}
	if checksum != "3f43304f3f7c8d51111d0846e13cb74e" {
		t.Errorf("Expected md5 '3f43304f3f7c8d51111d0846e13cb74e', got '%s'", checksum)
	}

	// Download with sha256 - our test file is about 9.5kb
	checksum, err = client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "sha256")
	if err != nil {
		t.Error(err)
	}
	if checksum != "6f87f5341df2558967da27b12939d5e88b3af06592104041e57043af150f9309" {
		t.Errorf("Expected sha256 '6f87f5341df2558967da27b12939d5e88b3af06592104041e57043af150f9309', got '%s'", checksum)
	}

	checksum, err = client.DownloadFile("aptrust.receiving.test.test.edu",
		"virginia.edu.uva-lib_2274765.tar", "invalid_algo")
	if err == nil {
		t.Error("DownloadFile should have returned error for invalid checksum algorithm")
	}

}
<|endoftext|>


<commit_before><commit_msg>md_util: don't fetch unmerged MD range when the BID is null<commit_after><|endoftext|>


<commit_before><commit_msg>Migrating to 0.2.7-alpha<commit_after><|endoftext|>


<commit_before>package downloader

import (
	"github.com/cenkalti/rain/internal/downloader/piecedownloader"
	"github.com/cenkalti/rain/internal/downloader/piecewriter"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/peer"
	"github.com/cenkalti/rain/internal/piece"
	"github.com/cenkalti/rain/internal/torrentdata"
	"github.com/cenkalti/rain/internal/worker"
)

const parallelPieceDownloads = 4

const parallelPieceWrites = 4

type Downloader struct {
	data           *torrentdata.Data
	messages       *peer.Messages
	pieces         []Piece
	connectedPeers map[*peer.Peer]struct{}
	unchokingPeers map[*peer.Peer]struct{}
	downloads      map[*peer.Peer]*piecedownloader.PieceDownloader
	downloadDone   chan *piecedownloader.PieceDownloader
	writeRequests  chan piecewriter.Request
	writeResponses chan piecewriter.Response
	errC           chan error
	log            logger.Logger
	limiter        chan struct{}
	workers        worker.Workers
}

type Piece struct {
	*piece.Piece
	havingPeers    map[*peer.Peer]struct{}
	requestedPeers map[*peer.Peer]*piecedownloader.PieceDownloader
	writing        bool
}

func New(d *torrentdata.Data, m *peer.Messages, errC chan error, l logger.Logger) *Downloader {
	pieces := make([]Piece, len(d.Pieces))
	for i := range d.Pieces {
		pieces[i] = Piece{
			Piece:          &d.Pieces[i],
			havingPeers:    make(map[*peer.Peer]struct{}),
			requestedPeers: make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		}
	}
	return &Downloader{
		data:           d,
		messages:       m,
		pieces:         pieces,
		connectedPeers: make(map[*peer.Peer]struct{}),
		unchokingPeers: make(map[*peer.Peer]struct{}),
		downloads:      make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		downloadDone:   make(chan *piecedownloader.PieceDownloader),
		writeRequests:  make(chan piecewriter.Request, 1),
		writeResponses: make(chan piecewriter.Response),
		errC:           errC,
		log:            l,
		limiter:        make(chan struct{}, parallelPieceDownloads),
	}
}

func (d *Downloader) Run(stopC chan struct{}) {
	defer d.workers.Stop()
	for i := 0; i < parallelPieceWrites; i++ {
		w := piecewriter.New(d.writeRequests, d.writeResponses, d.log)
		d.workers.Start(w)
	}
	waitingDownloader := 0
	for {
		select {
		case d.limiter <- struct{}{}:
			// TODO check status of existing downloads
			pi, pe, ok := d.nextDownload()
			if !ok {
				waitingDownloader++
				continue
			}
			d.log.Debugln("downloading piece", pi.Index, "from", pe.String())
			pd := piecedownloader.New(pi, pe)
			d.downloads[pe] = pd
			d.pieces[pi.Index].requestedPeers[pe] = pd
			d.workers.StartWithOnFinishHandler(pd, func() {
				select {
				case d.downloadDone <- pd:
				case <-stopC:
					return
				}
			})
		case pd := <-d.downloadDone:
			delete(d.downloads, pd.Peer)
			delete(d.pieces[pd.Piece.Index].requestedPeers, pd.Peer)
			<-d.limiter
			select {
			case buf := <-pd.DoneC:
				ok := d.pieces[pd.Piece.Index].Piece.Verify(buf)
				if !ok {
					// TODO handle corrupt piece
					continue
				}
				select {
				case d.writeRequests <- piecewriter.Request{Piece: pd.Piece, Data: buf}:
					d.pieces[pd.Piece.Index].writing = true
				case <-stopC:
					return
				}
			case err := <-pd.ErrC:
				d.log.Errorln("could not download piece:", err)
			case <-stopC:
				return
			}
		case resp := <-d.writeResponses:
			d.pieces[resp.Request.Piece.Index].writing = false
			if resp.Error != nil {
				d.errC <- resp.Error
				continue
			}
			d.data.Bitfield().Set(resp.Request.Piece.Index)
			d.data.CheckCompletion()
			// Tell everyone that we have this piece
			for pe := range d.connectedPeers {
				go pe.SendHave(resp.Request.Piece.Index)
			}
			// TODO update interested state
		case msg := <-d.messages.Have:
			if waitingDownloader > 0 {
				waitingDownloader--
				<-d.limiter
			}
			d.pieces[msg.Piece.Index].havingPeers[msg.Peer] = struct{}{}
			// TODO update interested state
			// go checkInterested(peer, bitfield)
			// peer.writeMessages <- interested{}
		case msg := <-d.messages.Bitfield:
			for i := uint32(0); i < msg.Bitfield.Len(); i++ {
				if msg.Bitfield.Test(i) {
					d.pieces[i].havingPeers[msg.Peer] = struct{}{}
				}
			}
			for waitingDownloader > 0 {
				waitingDownloader--
				<-d.limiter
			}
			// TODO update interested state
		case pe := <-d.messages.Unchoke:
			if waitingDownloader > 0 {
				waitingDownloader--
				<-d.limiter
			}
			d.unchokingPeers[pe] = struct{}{}
			if pd, ok := d.downloads[pe]; ok {
				pd.UnchokeC <- struct{}{}
			}
		case pe := <-d.messages.Choke:
			delete(d.unchokingPeers, pe)
			if pd, ok := d.downloads[pe]; ok {
				pd.ChokeC <- struct{}{}
			}
		case msg := <-d.messages.Piece:
			if pd, ok := d.downloads[msg.Peer]; ok {
				pd.PieceC <- msg
			}
		case msg := <-d.messages.Request:
			// pi := d.data.Pieces[msg.Index]
			buf := make([]byte, msg.Length)
			// TODO read data
			// pi.Data.ReadAt(buf, int64(msg.Begin))
			err := msg.Peer.SendPiece(msg.Piece.Index, msg.Begin, buf)
			if err != nil {
				d.log.Error(err)
				return
			}
		case pe := <-d.messages.Connect:
			d.connectedPeers[pe] = struct{}{}
		case pe := <-d.messages.Disconnect:
			delete(d.connectedPeers, pe)
		case <-stopC:
			return
		}
	}
}

func (d *Downloader) nextDownload() (pi *piece.Piece, pe *peer.Peer, ok bool) {
	// TODO selecting pieces in sequential order, change to rarest first
	for _, p := range d.pieces {
		if d.data.Bitfield().Test(p.Index) {
			continue
		}
		if p.writing {
			continue
		}
		if len(p.havingPeers) == 0 {
			continue
		}
		// TODO selecting first peer having the piece, change to more smart decision
		for pe2 := range p.havingPeers {
			if _, ok2 := d.unchokingPeers[pe2]; !ok2 {
				continue
			}
			if _, ok2 := p.requestedPeers[pe2]; ok2 {
				continue
			}
			if _, ok2 := d.downloads[pe2]; ok2 {
				continue
			}
			pe = pe2
			break
		}
		if pe == nil {
			continue
		}
		pi = p.Piece
		ok = true
		break
	}
	return
}
<commit_msg>send correct data<commit_after>package downloader

import (
	"github.com/cenkalti/rain/internal/downloader/piecedownloader"
	"github.com/cenkalti/rain/internal/downloader/piecewriter"
	"github.com/cenkalti/rain/internal/logger"
	"github.com/cenkalti/rain/internal/peer"
	"github.com/cenkalti/rain/internal/piece"
	"github.com/cenkalti/rain/internal/torrentdata"
	"github.com/cenkalti/rain/internal/worker"
)

const parallelPieceDownloads = 4

const parallelPieceWrites = 4

type Downloader struct {
	data           *torrentdata.Data
	messages       *peer.Messages
	pieces         []Piece
	connectedPeers map[*peer.Peer]struct{}
	unchokingPeers map[*peer.Peer]struct{}
	downloads      map[*peer.Peer]*piecedownloader.PieceDownloader
	downloadDone   chan *piecedownloader.PieceDownloader
	writeRequests  chan piecewriter.Request
	writeResponses chan piecewriter.Response
	errC           chan error
	log            logger.Logger
	limiter        chan struct{}
	workers        worker.Workers
}

type Piece struct {
	*piece.Piece
	havingPeers    map[*peer.Peer]struct{}
	requestedPeers map[*peer.Peer]*piecedownloader.PieceDownloader
	writing        bool
}

func New(d *torrentdata.Data, m *peer.Messages, errC chan error, l logger.Logger) *Downloader {
	pieces := make([]Piece, len(d.Pieces))
	for i := range d.Pieces {
		pieces[i] = Piece{
			Piece:          &d.Pieces[i],
			havingPeers:    make(map[*peer.Peer]struct{}),
			requestedPeers: make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		}
	}
	return &Downloader{
		data:           d,
		messages:       m,
		pieces:         pieces,
		connectedPeers: make(map[*peer.Peer]struct{}),
		unchokingPeers: make(map[*peer.Peer]struct{}),
		downloads:      make(map[*peer.Peer]*piecedownloader.PieceDownloader),
		downloadDone:   make(chan *piecedownloader.PieceDownloader),
		writeRequests:  make(chan piecewriter.Request, 1),
		writeResponses: make(chan piecewriter.Response),
		errC:           errC,
		log:            l,
		limiter:        make(chan struct{}, parallelPieceDownloads),
	}
}

func (d *Downloader) Run(stopC chan struct{}) {
	defer d.workers.Stop()
	for i := 0; i < parallelPieceWrites; i++ {
		w := piecewriter.New(d.writeRequests, d.writeResponses, d.log)
		d.workers.Start(w)
	}
	waitingDownloader := 0
	for {
		select {
		case d.limiter <- struct{}{}:
			// TODO check status of existing downloads
			pi, pe, ok := d.nextDownload()
			if !ok {
				waitingDownloader++
				continue
			}
			d.log.Debugln("downloading piece", pi.Index, "from", pe.String())
			pd := piecedownloader.New(pi, pe)
			d.downloads[pe] = pd
			d.pieces[pi.Index].requestedPeers[pe] = pd
			d.workers.StartWithOnFinishHandler(pd, func() {
				select {
				case d.downloadDone <- pd:
				case <-stopC:
					return
				}
			})
		case pd := <-d.downloadDone:
			delete(d.downloads, pd.Peer)
			delete(d.pieces[pd.Piece.Index].requestedPeers, pd.Peer)
			<-d.limiter
			select {
			case buf := <-pd.DoneC:
				ok := d.pieces[pd.Piece.Index].Piece.Verify(buf)
				if !ok {
					// TODO handle corrupt piece
					continue
				}
				select {
				case d.writeRequests <- piecewriter.Request{Piece: pd.Piece, Data: buf}:
					d.pieces[pd.Piece.Index].writing = true
				case <-stopC:
					return
				}
			case err := <-pd.ErrC:
				d.log.Errorln("could not download piece:", err)
			case <-stopC:
				return
			}
		case resp := <-d.writeResponses:
			d.pieces[resp.Request.Piece.Index].writing = false
			if resp.Error != nil {
				d.errC <- resp.Error
				continue
			}
			d.data.Bitfield().Set(resp.Request.Piece.Index)
			d.data.CheckCompletion()
			// Tell everyone that we have this piece
			for pe := range d.connectedPeers {
				go pe.SendHave(resp.Request.Piece.Index)
			}
			// TODO update interested state
		case msg := <-d.messages.Have:
			if waitingDownloader > 0 {
				waitingDownloader--
				<-d.limiter
			}
			d.pieces[msg.Piece.Index].havingPeers[msg.Peer] = struct{}{}
			// TODO update interested state
			// go checkInterested(peer, bitfield)
			// peer.writeMessages <- interested{}
		case msg := <-d.messages.Bitfield:
			for i := uint32(0); i < msg.Bitfield.Len(); i++ {
				if msg.Bitfield.Test(i) {
					d.pieces[i].havingPeers[msg.Peer] = struct{}{}
				}
			}
			for waitingDownloader > 0 {
				waitingDownloader--
				<-d.limiter
			}
			// TODO update interested state
		case pe := <-d.messages.Unchoke:
			if waitingDownloader > 0 {
				waitingDownloader--
				<-d.limiter
			}
			d.unchokingPeers[pe] = struct{}{}
			if pd, ok := d.downloads[pe]; ok {
				pd.UnchokeC <- struct{}{}
			}
		case pe := <-d.messages.Choke:
			delete(d.unchokingPeers, pe)
			if pd, ok := d.downloads[pe]; ok {
				pd.ChokeC <- struct{}{}
			}
		case msg := <-d.messages.Piece:
			if pd, ok := d.downloads[msg.Peer]; ok {
				pd.PieceC <- msg
			}
		case msg := <-d.messages.Request:
			go func(msg peer.Request) {
				buf := make([]byte, msg.Length)
				msg.Piece.Data.ReadAt(buf, int64(msg.Begin))
				err := msg.Peer.SendPiece(msg.Piece.Index, msg.Begin, buf)
				if err != nil {
					// TODO ignore write errors
					d.log.Error(err)
				}
			}(msg)
		case pe := <-d.messages.Connect:
			d.connectedPeers[pe] = struct{}{}
		case pe := <-d.messages.Disconnect:
			delete(d.connectedPeers, pe)
		case <-stopC:
			return
		}
	}
}

func (d *Downloader) nextDownload() (pi *piece.Piece, pe *peer.Peer, ok bool) {
	// TODO selecting pieces in sequential order, change to rarest first
	for _, p := range d.pieces {
		if d.data.Bitfield().Test(p.Index) {
			continue
		}
		if p.writing {
			continue
		}
		if len(p.havingPeers) == 0 {
			continue
		}
		// TODO selecting first peer having the piece, change to more smart decision
		for pe2 := range p.havingPeers {
			if _, ok2 := d.unchokingPeers[pe2]; !ok2 {
				continue
			}
			if _, ok2 := p.requestedPeers[pe2]; ok2 {
				continue
			}
			if _, ok2 := d.downloads[pe2]; ok2 {
				continue
			}
			pe = pe2
			break
		}
		if pe == nil {
			continue
		}
		pi = p.Piece
		ok = true
		break
	}
	return
}
<|endoftext|>


<commit_before>package storage

import (
	"reflect"
	"testing"
)

func TestKeyIndexGet(t *testing.T) {
	// key:     "foo"
	// rev: 12
	// generations:
	//    {empty}
	//    {8[1], 10[2], 12(t)[3]}
	//    {4[2], 6(t)[3]}
	ki := newTestKeyIndex()
	ki.compact(4, make(map[revision]struct{}))

	tests := []struct {
		rev int64

		wrev int64
		werr error
	}{
		{13, 12, nil},
		{13, 12, nil},

		// get on generation 2
		{12, 12, nil},
		{11, 10, nil},
		{10, 10, nil},
		{9, 8, nil},
		{8, 8, nil},
		{7, 6, nil},

		// get on generation 1
		{6, 6, nil},
		{5, 4, nil},
		{4, 4, nil},
	}

	for i, tt := range tests {
		rev, _, _, err := ki.get(tt.rev)
		if err != tt.werr {
			t.Errorf("#%d: err = %v, want %v", i, err, tt.werr)
		}
		if rev.main != tt.wrev {
			t.Errorf("#%d: rev = %d, want %d", i, rev.main, tt.rev)
		}
	}
}

func TestKeyIndexPut(t *testing.T) {
	ki := &keyIndex{key: []byte("foo")}
	ki.put(5, 0)

	wki := &keyIndex{
		key:         []byte("foo"),
		modified:    revision{5, 0},
		generations: []generation{{created: revision{5, 0}, ver: 1, revs: []revision{{main: 5}}}},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}

	ki.put(7, 0)

	wki = &keyIndex{
		key:         []byte("foo"),
		modified:    revision{7, 0},
		generations: []generation{{created: revision{5, 0}, ver: 2, revs: []revision{{main: 5}, {main: 7}}}},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}
}

func TestKeyIndexTombstone(t *testing.T) {
	ki := &keyIndex{key: []byte("foo")}
	ki.put(5, 0)

	ki.tombstone(7, 0)

	wki := &keyIndex{
		key:         []byte("foo"),
		modified:    revision{7, 0},
		generations: []generation{{created: revision{5, 0}, ver: 2, revs: []revision{{main: 5}, {main: 7}}}, {}},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}

	ki.put(8, 0)
	ki.put(9, 0)
	ki.tombstone(15, 0)

	wki = &keyIndex{
		key:      []byte("foo"),
		modified: revision{15, 0},
		generations: []generation{
			{created: revision{5, 0}, ver: 2, revs: []revision{{main: 5}, {main: 7}}},
			{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 9}, {main: 15}}},
			{},
		},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}
}

func TestKeyIndexCompact(t *testing.T) {
	tests := []struct {
		compact int64

		wki *keyIndex
		wam map[revision]struct{}
	}{
		{
			1,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 2}, {main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{},
		},
		{
			2,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 2}, {main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 2}: {},
			},
		},
		{
			3,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 2}, {main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 2}: {},
			},
		},
		{
			4,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 4}: {},
			},
		},
		{
			5,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 4}: {},
			},
		},
		{
			6,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{},
		},
		{
			7,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{},
		},
		{
			8,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 8}: {},
			},
		},
		{
			9,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 8}: {},
			},
		},
		{
			10,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 10}: {},
			},
		},
		{
			11,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 10}: {},
			},
		},
		{
			12,
			&keyIndex{
				key:         []byte("foo"),
				modified:    revision{12, 0},
				generations: []generation{{}},
			},
			map[revision]struct{}{},
		},
	}

	// Continous Compaction
	ki := newTestKeyIndex()
	for i, tt := range tests {
		am := make(map[revision]struct{})
		ki.compact(tt.compact, am)
		if !reflect.DeepEqual(ki, tt.wki) {
			t.Errorf("#%d: ki = %+v, want %+v", i, ki, tt.wki)
		}
		if !reflect.DeepEqual(am, tt.wam) {
			t.Errorf("#%d: am = %+v, want %+v", i, am, tt.wam)
		}
	}

	// Jump Compaction
	for i, tt := range tests {
		if (i%2 == 0 && i < 6) || (i%2 == 1 && i > 6) {
			am := make(map[revision]struct{})
			ki.compact(tt.compact, am)
			if !reflect.DeepEqual(ki, tt.wki) {
				t.Errorf("#%d: ki = %+v, want %+v", i, ki, tt.wki)
			}
			if !reflect.DeepEqual(am, tt.wam) {
				t.Errorf("#%d: am = %+v, want %+v", i, am, tt.wam)
			}
		}
	}

	// OnceCompaction
	for i, tt := range tests {
		ki := newTestKeyIndex()
		am := make(map[revision]struct{})
		ki.compact(tt.compact, am)
		if !reflect.DeepEqual(ki, tt.wki) {
			t.Errorf("#%d: ki = %+v, want %+v", i, ki, tt.wki)
		}
		if !reflect.DeepEqual(am, tt.wam) {
			t.Errorf("#%d: am = %+v, want %+v", i, am, tt.wam)
		}
	}
}

func newTestKeyIndex() *keyIndex {
	// key:     "foo"
	// rev: 12
	// generations:
	//    {empty}
	//    {8[1], 10[2], 12(t)[3]}
	//    {2[1], 4[2], 6(t)[3]}

	ki := &keyIndex{key: []byte("foo")}
	ki.put(2, 0)
	ki.put(4, 0)
	ki.tombstone(6, 0)
	ki.put(8, 0)
	ki.put(10, 0)
	ki.tombstone(12, 0)
	return ki
}
<commit_msg>storage: fix TestKeyIndexCompact<commit_after>package storage

import (
	"reflect"
	"testing"
)

func TestKeyIndexGet(t *testing.T) {
	// key:     "foo"
	// rev: 12
	// generations:
	//    {empty}
	//    {8[1], 10[2], 12(t)[3]}
	//    {4[2], 6(t)[3]}
	ki := newTestKeyIndex()
	ki.compact(4, make(map[revision]struct{}))

	tests := []struct {
		rev int64

		wrev int64
		werr error
	}{
		{13, 12, nil},
		{13, 12, nil},

		// get on generation 2
		{12, 12, nil},
		{11, 10, nil},
		{10, 10, nil},
		{9, 8, nil},
		{8, 8, nil},
		{7, 6, nil},

		// get on generation 1
		{6, 6, nil},
		{5, 4, nil},
		{4, 4, nil},
	}

	for i, tt := range tests {
		rev, _, _, err := ki.get(tt.rev)
		if err != tt.werr {
			t.Errorf("#%d: err = %v, want %v", i, err, tt.werr)
		}
		if rev.main != tt.wrev {
			t.Errorf("#%d: rev = %d, want %d", i, rev.main, tt.rev)
		}
	}
}

func TestKeyIndexPut(t *testing.T) {
	ki := &keyIndex{key: []byte("foo")}
	ki.put(5, 0)

	wki := &keyIndex{
		key:         []byte("foo"),
		modified:    revision{5, 0},
		generations: []generation{{created: revision{5, 0}, ver: 1, revs: []revision{{main: 5}}}},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}

	ki.put(7, 0)

	wki = &keyIndex{
		key:         []byte("foo"),
		modified:    revision{7, 0},
		generations: []generation{{created: revision{5, 0}, ver: 2, revs: []revision{{main: 5}, {main: 7}}}},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}
}

func TestKeyIndexTombstone(t *testing.T) {
	ki := &keyIndex{key: []byte("foo")}
	ki.put(5, 0)

	ki.tombstone(7, 0)

	wki := &keyIndex{
		key:         []byte("foo"),
		modified:    revision{7, 0},
		generations: []generation{{created: revision{5, 0}, ver: 2, revs: []revision{{main: 5}, {main: 7}}}, {}},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}

	ki.put(8, 0)
	ki.put(9, 0)
	ki.tombstone(15, 0)

	wki = &keyIndex{
		key:      []byte("foo"),
		modified: revision{15, 0},
		generations: []generation{
			{created: revision{5, 0}, ver: 2, revs: []revision{{main: 5}, {main: 7}}},
			{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 9}, {main: 15}}},
			{},
		},
	}
	if !reflect.DeepEqual(ki, wki) {
		t.Errorf("ki = %+v, want %+v", ki, wki)
	}
}

func TestKeyIndexCompact(t *testing.T) {
	tests := []struct {
		compact int64

		wki *keyIndex
		wam map[revision]struct{}
	}{
		{
			1,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 2}, {main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{},
		},
		{
			2,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 2}, {main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 2}: {},
			},
		},
		{
			3,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 2}, {main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 2}: {},
			},
		},
		{
			4,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 4}: {},
			},
		},
		{
			5,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{2, 0}, ver: 3, revs: []revision{{main: 4}, {main: 6}}},
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 4}: {},
			},
		},
		{
			6,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{},
		},
		{
			7,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{},
		},
		{
			8,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 8}: {},
			},
		},
		{
			9,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 8}, {main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 8}: {},
			},
		},
		{
			10,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 10}: {},
			},
		},
		{
			11,
			&keyIndex{
				key:      []byte("foo"),
				modified: revision{12, 0},
				generations: []generation{
					{created: revision{8, 0}, ver: 3, revs: []revision{{main: 10}, {main: 12}}},
					{},
				},
			},
			map[revision]struct{}{
				revision{main: 10}: {},
			},
		},
		{
			12,
			&keyIndex{
				key:         []byte("foo"),
				modified:    revision{12, 0},
				generations: []generation{{}},
			},
			map[revision]struct{}{},
		},
	}

	// Continous Compaction
	ki := newTestKeyIndex()
	for i, tt := range tests {
		am := make(map[revision]struct{})
		ki.compact(tt.compact, am)
		if !reflect.DeepEqual(ki, tt.wki) {
			t.Errorf("#%d: ki = %+v, want %+v", i, ki, tt.wki)
		}
		if !reflect.DeepEqual(am, tt.wam) {
			t.Errorf("#%d: am = %+v, want %+v", i, am, tt.wam)
		}
	}

	ki = newTestKeyIndex()
	// Jump Compaction
	for i, tt := range tests {
		if (i%2 == 0 && i < 6) || (i%2 == 1 && i > 6) {
			am := make(map[revision]struct{})
			ki.compact(tt.compact, am)
			if !reflect.DeepEqual(ki, tt.wki) {
				t.Errorf("#%d: ki = %+v, want %+v", i, ki, tt.wki)
			}
			if !reflect.DeepEqual(am, tt.wam) {
				t.Errorf("#%d: am = %+v, want %+v", i, am, tt.wam)
			}
		}
	}

	// OnceCompaction
	for i, tt := range tests {
		ki := newTestKeyIndex()
		am := make(map[revision]struct{})
		ki.compact(tt.compact, am)
		if !reflect.DeepEqual(ki, tt.wki) {
			t.Errorf("#%d: ki = %+v, want %+v", i, ki, tt.wki)
		}
		if !reflect.DeepEqual(am, tt.wam) {
			t.Errorf("#%d: am = %+v, want %+v", i, am, tt.wam)
		}
	}
}

func newTestKeyIndex() *keyIndex {
	// key:     "foo"
	// rev: 12
	// generations:
	//    {empty}
	//    {8[1], 10[2], 12(t)[3]}
	//    {2[1], 4[2], 6(t)[3]}

	ki := &keyIndex{key: []byte("foo")}
	ki.put(2, 0)
	ki.put(4, 0)
	ki.tombstone(6, 0)
	ki.put(8, 0)
	ki.put(10, 0)
	ki.tombstone(12, 0)
	return ki
}
<|endoftext|>


<commit_before>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package money

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"errors"
	"text/template"
	"unicode"
	"unicode/utf8"

	"github.com/corestoreio/csfw/utils/log"
	"github.com/juju/errgo"
)

var (
	_          json.Unmarshaler = (*Currency)(nil)
	_          json.Marshaler   = (*Currency)(nil)
	_          sql.Scanner      = (*Currency)(nil)
	nullString                  = []byte(`null`)
)

// ErrDecodeMissingColon can be returned on malformed JSON value when decoding a currency.
var ErrDecodeMissingColon = errors.New("No colon found in JSON array")

const (
	// JSONNumber encodes/decodes a currency as a number string to directly use
	// in e.g. JavaScript
	JSONNumber JSONType = 1 << iota
	// JSONLocale encodes/decodes a currency according to its locale format.
	// Decoding: Considers the locale if the currency symbol is valid.
	JSONLocale
	// JSONExtended encodes/decodes a currency into a JSON array:
	// [1234.56, "€", "1.234,56 €"].
	// Decoding: Considers the locale if the currency symbol is valid.
	JSONExtended
)

// JSONType defines the type of the marshaller/unmarshaller
type JSONType uint8

type (
	// JSONMarshaller interface for JSON encoding
	JSONMarshaller interface {
		// MarshalJSON encodes the currency
		MarshalJSON(*Currency) ([]byte, error)
	}
	// JSONUnmarshaller interface for JSON decoding
	JSONUnmarshaller interface {
		// UnmarshalJSON reads the bytes and decodes them into the currency
		UnmarshalJSON(*Currency, []byte) error
	}
)

// MarshalJSON generates JSON output depending on the Marshaller.
func (c Currency) MarshalJSON() ([]byte, error) {
	return c.jm.MarshalJSON(&c)
}

// UnmarshalJSON reads JSON and fills the currency struct depending on the Unmarshaller.
func (c *Currency) UnmarshalJSON(src []byte) error {
	c.applyDefaults()
	if src == nil {
		c.m, c.Valid = 0, false
		return nil
	}
	return c.jum.UnmarshalJSON(c, src)
}

// Scan scans a value into the Currency struct. Returns an error on data loss.
// Errors will be logged. Initial default settings are the guard and precision value.
func (c *Currency) Scan(src interface{}) error {
	c.applyDefaults()

	if src == nil {
		c.m, c.Valid = 0, false
		if log.IsDebug() {
			log.Debug("Currency=Scan", "case", 89, "c", c, "src", src)
		}
		return nil
	}

	if b, ok := src.([]byte); ok {
		return c.ParseFloat(string(b))
	}
	return log.Error("Currency=Scan", "err", errgo.Newf("Unsupported Type for value. Supported: []byte"), "src", src)
}

// NewJSONEncoder creates a new encoder depending on the type.
// Accepts either zero or one argument.
// Default encoder is JSONLocale
func NewJSONEncoder(jts ...JSONType) JSONMarshaller {
	if len(jts) != 1 {
		return JSONLocale
	}
	return jts[0]
}

// NewJSONDecoder creates a new decoder depending on the type.
// Accepts either zero or one argument.
// Default decoder is JSONLocale
func NewJSONDecoder(jts ...JSONType) JSONUnmarshaller {
	if len(jts) != 1 {
		return JSONLocale
	}
	return jts[0]
}

var _ JSONMarshaller = new(JSONType)
var _ JSONUnmarshaller = new(JSONType)

// MarshalJSON encodes a currency to JSON bytes according to the defined JSONType
func (t JSONType) MarshalJSON(c *Currency) ([]byte, error) {
	switch t {
	case JSONNumber:
		return jsonNumberMarshal(c)
	case JSONExtended:
		return jsonExtendedMarshal(c)
	default:
		return jsonLocaleMarshal(c)
	}
}

// UnmarshalJSON decodes three different currency representations into a currency
// struct.
func (t JSONType) UnmarshalJSON(c *Currency, b []byte) error {
	if len(b) < 1 || false == utf8.Valid(b) { // we must have a valid string
		if log.IsDebug() {
			log.Debug("JSONType=UnmarshalJSON", "case", "invalid_bytes", "c", c, "bytes", string(b))
		}
		c.m, c.Valid = 0, false
		return nil
	}

	runes := bytes.Runes(b)
	lenRunes := len(runes)
	var realNumber, isNull, lRunes, posSepComma, posSepDot int
	var isArray bool
	number := make([]rune, 0, lenRunes)
	// atm not needed because currency symbol depends on the formatter
	//symbol := make([]rune, 0, lenRunes)

	// strip quotes
	if lenRunes > 1 && runes[0] == '"' && runes[lenRunes-1] == '"' {
		runes = runes[1 : lenRunes-1]
	}
	lenRunes = len(runes)

	if 0 == lenRunes {
		if log.IsDebug() {
			log.Debug("JSONType=UnmarshalJSON", "case", "lenRunes=0", "c", c, "bytes", string(b))
		}
		c.m, c.Valid = 0, false
		return nil
	}

OuterLoop:
	for i, r := range runes {

		switch {
		case unicode.IsSpace(r):
			continue
		case r == '[':
			isArray = true // [999.0000,"$","$ 999.00"] only until the first comma will be considered.
		case unicode.IsNumber(r): // 1234.56
			number = append(number, r)
			realNumber++
		case r == '.', r == '-': // -1234.56
			number = append(number, r)
			realNumber++
		case r == ',': // -1,234.56 or -1.234,56 or -1 234,56
			if isArray { // we stop after the first colon, because then the 2nd entry starts in the array
				isArray = false
				break OuterLoop
			}
			number = append(number, r)
			//case unicode.IsLetter(r), unicode.IsSymbol(r):
			//	symbol = append(symbol, r)
		}

		if posSepComma == 0 && r == ',' { // check for first occurrence of the comma
			posSepComma = i
		}
		if posSepDot == 0 && r == '.' {
			posSepDot = i
		}

		switch unicode.ToLower(r) {
		case 'n', 'u', 'l':
			isNull++
		}

		if isNull == 4 {
			if log.IsDebug() {
				log.Debug("JSONType=UnmarshalJSON", "case", "isNull", "c", c, "bytes", string(b), "runes", string(runes))
			}
			c.m, c.Valid = 0, false
			return nil
		}

		lRunes++
	}

	if isArray { // now it's an error because no colon found
		c.m, c.Valid = 0, false
		return log.Error("JSONType=UnmarshalJSON", "err", ErrDecodeMissingColon, "bytes", string(b), "number", string(number))
	}

	// I'll keep this redundant IF cases because if the unit tests and coverage checks.
	// Once it's working we can merge them.

	if realNumber == lRunes { // real number e.g. -1234.56 without any other stuff
		return c.ParseFloat(string(runes))
	}
	if posSepComma == 0 && posSepDot == 0 { // no decimals but included any other stripped of character
		return c.ParseFloat(string(number))
	}
	if posSepComma == 0 && posSepDot > 0 { // currency contains only a dot
		return c.ParseFloat(string(number))
	}
	if posSepComma > 0 && posSepDot == 0 { // currency contains only a comma
		for i, r := range number {
			if r == ',' {
				number[i] = '.'
			}
		}
		return c.ParseFloat(string(number))
	}
	if posSepComma > 0 && posSepDot > 0 {

		replaceChar := ','           // number is 12,211,232.45 or 1,234.56
		if posSepDot < posSepComma { // number is 12.211.232,45 or 1.234,56
			replaceChar = '.'
		}

		var i int
		for i < len(number) {
			switch {
			case replaceChar == '.' && number[i] == ',':
				number[i] = '.' // replace decimal comma with a dot to create fractals
			case number[i] == replaceChar:
				number = append(number[:i], number[i+1:]...) // cut comma
				i = 0                                        // restart loop
			}
			i++
		}
		return c.ParseFloat(string(number))
	}

	c.m, c.Valid = 0, false
	return log.Error("JSONType=UnmarshalJSON", "err", errors.New("Invalid bytes"), "bytes", string(b), "number", string(number))
}

// jsonNumberMarshal generates a number formatted currency string
func jsonNumberMarshal(c *Currency) ([]byte, error) {
	if c == nil || c.Valid == false {
		return nullString, nil
	}
	return c.Ftoa(), nil
}

// jsonLocaleMarshal encodes into a locale specific quoted string
func jsonLocaleMarshal(c *Currency) ([]byte, error) {
	if c == nil || c.Valid == false {
		return nullString, nil
	}
	var b bytes.Buffer
	b.WriteString(`"`)
	lb, err := c.Localize()
	if err != nil {
		return nil, log.Error("JSONLocale=MarshalJSON", "err", err, "currency", c, "bytes", lb)
	}
	template.JSEscape(&b, lb)
	b.WriteString(`"`)
	return b.Bytes(), err
}

// jsonExtendedMarshal encodes a currency into a JSON array: [1234.56, "€", "1.234,56 €"]
func jsonExtendedMarshal(c *Currency) ([]byte, error) {
	if c == nil || c.Valid == false {
		return nullString, nil
	}
	var b bytes.Buffer
	b.WriteRune('[')
	b.Write(c.Ftoa())
	b.WriteString(`, "`)
	b.WriteString(template.JSEscapeString(string(c.Symbol())))
	b.WriteString(`", "`)
	lb, err := c.Localize()
	if err != nil {
		return nil, log.Error("JSONLocale=MarshalJSON", "err", err, "currency", c, "bytes", lb)
	}
	template.JSEscape(&b, lb)

	b.WriteRune('"')
	b.WriteRune(']')
	return b.Bytes(), err
}
<commit_msg>Package money: UnmarshalJSON merge IFs into a switch<commit_after>// Copyright 2015, Cyrill @ Schumacher.fm and the CoreStore contributors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package money

import (
	"bytes"
	"database/sql"
	"encoding/json"
	"errors"
	"text/template"
	"unicode"
	"unicode/utf8"

	"github.com/corestoreio/csfw/utils/log"
	"github.com/juju/errgo"
)

var (
	_          json.Unmarshaler = (*Currency)(nil)
	_          json.Marshaler   = (*Currency)(nil)
	_          sql.Scanner      = (*Currency)(nil)
	nullString                  = []byte(`null`)
)

// ErrDecodeMissingColon can be returned on malformed JSON value when decoding a currency.
var ErrDecodeMissingColon = errors.New("No colon found in JSON array")

const (
	// JSONNumber encodes/decodes a currency as a number string to directly use
	// in e.g. JavaScript
	JSONNumber JSONType = 1 << iota
	// JSONLocale encodes/decodes a currency according to its locale format.
	// Decoding: Considers the locale if the currency symbol is valid.
	JSONLocale
	// JSONExtended encodes/decodes a currency into a JSON array:
	// [1234.56, "€", "1.234,56 €"].
	// Decoding: Considers the locale if the currency symbol is valid.
	JSONExtended
)

// JSONType defines the type of the marshaller/unmarshaller
type JSONType uint8

type (
	// JSONMarshaller interface for JSON encoding
	JSONMarshaller interface {
		// MarshalJSON encodes the currency
		MarshalJSON(*Currency) ([]byte, error)
	}
	// JSONUnmarshaller interface for JSON decoding
	JSONUnmarshaller interface {
		// UnmarshalJSON reads the bytes and decodes them into the currency
		UnmarshalJSON(*Currency, []byte) error
	}
)

// MarshalJSON generates JSON output depending on the Marshaller.
func (c Currency) MarshalJSON() ([]byte, error) {
	return c.jm.MarshalJSON(&c)
}

// UnmarshalJSON reads JSON and fills the currency struct depending on the Unmarshaller.
func (c *Currency) UnmarshalJSON(src []byte) error {
	c.applyDefaults()
	if src == nil {
		c.m, c.Valid = 0, false
		return nil
	}
	return c.jum.UnmarshalJSON(c, src)
}

// Scan scans a value into the Currency struct. Returns an error on data loss.
// Errors will be logged. Initial default settings are the guard and precision value.
func (c *Currency) Scan(src interface{}) error {
	c.applyDefaults()

	if src == nil {
		c.m, c.Valid = 0, false
		if log.IsDebug() {
			log.Debug("Currency=Scan", "case", 89, "c", c, "src", src)
		}
		return nil
	}

	if b, ok := src.([]byte); ok {
		return c.ParseFloat(string(b))
	}
	return log.Error("Currency=Scan", "err", errgo.Newf("Unsupported Type for value. Supported: []byte"), "src", src)
}

// NewJSONEncoder creates a new encoder depending on the type.
// Accepts either zero or one argument.
// Default encoder is JSONLocale
func NewJSONEncoder(jts ...JSONType) JSONMarshaller {
	if len(jts) != 1 {
		return JSONLocale
	}
	return jts[0]
}

// NewJSONDecoder creates a new decoder depending on the type.
// Accepts either zero or one argument.
// Default decoder is JSONLocale
func NewJSONDecoder(jts ...JSONType) JSONUnmarshaller {
	if len(jts) != 1 {
		return JSONLocale
	}
	return jts[0]
}

var _ JSONMarshaller = new(JSONType)
var _ JSONUnmarshaller = new(JSONType)

// MarshalJSON encodes a currency to JSON bytes according to the defined JSONType
func (t JSONType) MarshalJSON(c *Currency) ([]byte, error) {
	switch t {
	case JSONNumber:
		return jsonNumberMarshal(c)
	case JSONExtended:
		return jsonExtendedMarshal(c)
	default:
		return jsonLocaleMarshal(c)
	}
}

// UnmarshalJSON decodes three different currency representations into a currency
// struct.
func (t JSONType) UnmarshalJSON(c *Currency, b []byte) error {
	if len(b) < 1 || false == utf8.Valid(b) { // we must have a valid string
		if log.IsDebug() {
			log.Debug("JSONType=UnmarshalJSON", "case", "invalid_bytes", "c", c, "bytes", string(b))
		}
		c.m, c.Valid = 0, false
		return nil
	}

	runes := bytes.Runes(b)
	lenRunes := len(runes)
	var realNumber, isNull, lRunes, posSepComma, posSepDot int
	var isArray bool
	number := make([]rune, 0, lenRunes)
	// atm not needed because currency symbol depends on the formatter
	//symbol := make([]rune, 0, lenRunes)

	// strip quotes
	if lenRunes > 1 && runes[0] == '"' && runes[lenRunes-1] == '"' {
		runes = runes[1 : lenRunes-1]
	}
	lenRunes = len(runes)

	if 0 == lenRunes {
		if log.IsDebug() {
			log.Debug("JSONType=UnmarshalJSON", "case", "lenRunes=0", "c", c, "bytes", string(b))
		}
		c.m, c.Valid = 0, false
		return nil
	}

OuterLoop:
	for i, r := range runes {

		switch {
		case unicode.IsSpace(r):
			continue
		case r == '[':
			isArray = true // [999.0000,"$","$ 999.00"] only until the first comma will be considered.
		case unicode.IsNumber(r): // 1234.56
			number = append(number, r)
			realNumber++
		case r == '.', r == '-': // -1234.56
			number = append(number, r)
			realNumber++
		case r == ',': // -1,234.56 or -1.234,56 or -1 234,56
			if isArray { // we stop after the first colon, because then the 2nd entry starts in the array
				isArray = false
				break OuterLoop
			}
			number = append(number, r)
			//case unicode.IsLetter(r), unicode.IsSymbol(r):
			//	symbol = append(symbol, r)
		}

		if posSepComma == 0 && r == ',' { // check for first occurrence of the comma
			posSepComma = i
		}
		if posSepDot == 0 && r == '.' {
			posSepDot = i
		}

		switch unicode.ToLower(r) {
		case 'n', 'u', 'l':
			isNull++
		}

		if isNull == 4 {
			if log.IsDebug() {
				log.Debug("JSONType=UnmarshalJSON", "case", "isNull", "c", c, "bytes", string(b), "runes", string(runes))
			}
			c.m, c.Valid = 0, false
			return nil
		}

		lRunes++
	}

	if isArray { // now it's an error because no colon found
		c.m, c.Valid = 0, false
		return log.Error("JSONType=UnmarshalJSON", "err", ErrDecodeMissingColon, "bytes", string(b), "number", string(number))
	}

	switch {
	case realNumber == lRunes: // real number e.g. -1234.56 without any other stuff
		return c.ParseFloat(string(runes))

	case posSepComma == 0 && posSepDot == 0, // no decimals but included any other stripped of character
		posSepComma == 0 && posSepDot > 0: // currency contains only a dot
		return c.ParseFloat(string(number))

	case posSepComma > 0 && posSepDot == 0: // currency contains only a comma
		for i, r := range number {
			if r == ',' {
				number[i] = '.'
			}
		}
		return c.ParseFloat(string(number))

	case posSepComma > 0 && posSepDot > 0:
		replaceChar := ','           // number is 12,211,232.45 or 1,234.56
		if posSepDot < posSepComma { // number is 12.211.232,45 or 1.234,56
			replaceChar = '.'
		}

		var i int
		for i < len(number) {
			switch {
			case replaceChar == '.' && number[i] == ',':
				number[i] = '.' // replace decimal comma with a dot to create fractals
			case number[i] == replaceChar:
				number = append(number[:i], number[i+1:]...) // cut comma
				i = 0                                        // restart loop
			}
			i++
		}
		return c.ParseFloat(string(number))
	}

	c.m, c.Valid = 0, false
	return log.Error("JSONType=UnmarshalJSON", "err", errors.New("Invalid bytes"), "bytes", string(b), "number", string(number))
}

// jsonNumberMarshal generates a number formatted currency string
func jsonNumberMarshal(c *Currency) ([]byte, error) {
	if c == nil || c.Valid == false {
		return nullString, nil
	}
	return c.Ftoa(), nil
}

// jsonLocaleMarshal encodes into a locale specific quoted string
func jsonLocaleMarshal(c *Currency) ([]byte, error) {
	if c == nil || c.Valid == false {
		return nullString, nil
	}
	var b bytes.Buffer
	b.WriteString(`"`)
	lb, err := c.Localize()
	if err != nil {
		return nil, log.Error("JSONLocale=MarshalJSON", "err", err, "currency", c, "bytes", lb)
	}
	template.JSEscape(&b, lb)
	b.WriteString(`"`)
	return b.Bytes(), err
}

// jsonExtendedMarshal encodes a currency into a JSON array: [1234.56, "€", "1.234,56 €"]
func jsonExtendedMarshal(c *Currency) ([]byte, error) {
	if c == nil || c.Valid == false {
		return nullString, nil
	}
	var b bytes.Buffer
	b.WriteRune('[')
	b.Write(c.Ftoa())
	b.WriteString(`, "`)
	b.WriteString(template.JSEscapeString(string(c.Symbol())))
	b.WriteString(`", "`)
	lb, err := c.Localize()
	if err != nil {
		return nil, log.Error("JSONLocale=MarshalJSON", "err", err, "currency", c, "bytes", lb)
	}
	template.JSEscape(&b, lb)

	b.WriteRune('"')
	b.WriteRune(']')
	return b.Bytes(), err
}
<|endoftext|>


<commit_before><commit_msg>Pending online state<commit_after><|endoftext|>


<commit_before>package refmt

import (
	"io"

	"github.com/polydawn/refmt/cbor"
	"github.com/polydawn/refmt/json"
	"github.com/polydawn/refmt/obj/atlas"
)

type EncodeOptions interface {
	IsEncodeOptions() // marker method.
}

func Marshal(opts EncodeOptions, v interface{}) ([]byte, error) {
	switch opts.(type) {
	case json.EncodeOptions:
		return json.Marshal(v)
	case cbor.EncodeOptions:
		return cbor.Marshal(v)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}

func MarshalAtlased(opts EncodeOptions, v interface{}, atl atlas.Atlas) ([]byte, error) {
	switch opts.(type) {
	case json.EncodeOptions:
		return json.MarshalAtlased(v, atl)
	case cbor.EncodeOptions:
		return cbor.MarshalAtlased(v, atl)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}

type Marshaller interface {
	Marshal(v interface{}) error
}

func NewMarshaller(opts EncodeOptions, wr io.Writer) Marshaller {
	switch opts.(type) {
	case json.EncodeOptions:
		return json.NewMarshaller(wr)
	case cbor.EncodeOptions:
		return cbor.NewMarshaller(wr)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}

func NewMarshallerAtlased(opts EncodeOptions, wr io.Writer, atl atlas.Atlas) Marshaller {
	switch o2 := opts.(type) {
	case json.EncodeOptions:
		return json.NewMarshallerAtlased(wr, o2, atl)
	case cbor.EncodeOptions:
		return cbor.NewMarshallerAtlased(wr, atl)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}
<commit_msg>Fix wiring of options when building not-atlased marshaller.<commit_after>package refmt

import (
	"io"

	"github.com/polydawn/refmt/cbor"
	"github.com/polydawn/refmt/json"
	"github.com/polydawn/refmt/obj/atlas"
)

type EncodeOptions interface {
	IsEncodeOptions() // marker method.
}

func Marshal(opts EncodeOptions, v interface{}) ([]byte, error) {
	switch opts.(type) {
	case json.EncodeOptions:
		return json.Marshal(v)
	case cbor.EncodeOptions:
		return cbor.Marshal(v)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}

func MarshalAtlased(opts EncodeOptions, v interface{}, atl atlas.Atlas) ([]byte, error) {
	switch opts.(type) {
	case json.EncodeOptions:
		return json.MarshalAtlased(v, atl)
	case cbor.EncodeOptions:
		return cbor.MarshalAtlased(v, atl)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}

type Marshaller interface {
	Marshal(v interface{}) error
}

func NewMarshaller(opts EncodeOptions, wr io.Writer) Marshaller {
	switch o2 := opts.(type) {
	case json.EncodeOptions:
		return json.NewMarshallerAtlased(wr, o2, atlas.MustBuild())
	case cbor.EncodeOptions:
		return cbor.NewMarshaller(wr)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}

func NewMarshallerAtlased(opts EncodeOptions, wr io.Writer, atl atlas.Atlas) Marshaller {
	switch o2 := opts.(type) {
	case json.EncodeOptions:
		return json.NewMarshallerAtlased(wr, o2, atl)
	case cbor.EncodeOptions:
		return cbor.NewMarshallerAtlased(wr, atl)
	default:
		panic("incorrect usage: unknown EncodeOptions type")
	}
}
<|endoftext|>