<commit_before>package srvclient

// At the moment go's dns resolver which is built into the net package doesn't
// properly handle the case of a response being too big. Which leads us to
// having to manually parse /etc/resolv.conf and manually make the SRV requests.

import (
	"errors"
	"math/rand"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"

	"net"
	"sort"

	"github.com/miekg/dns"
)

// sortableSRV implements sort.Interface for []*dns.SRV based on
// the Priority and Weight fields
type sortableSRV []*dns.SRV

func (a sortableSRV) Len() int      { return len(a) }
func (a sortableSRV) Swap(i, j int) { a[i], a[j] = a[j], a[i] }
func (a sortableSRV) Less(i, j int) bool {
	if a[i].Priority == a[j].Priority {
		return a[i].Weight > a[j].Weight
	}
	return a[i].Priority < a[j].Priority
}

func init() {
	go dnsConfigLoop()
}

// SRVClient is a holder for methods related to SRV lookups. Use new(SRVClient)
// to initialize one.
type SRVClient struct {
	cacheLast  map[string]*dns.Msg
	cacheLastL sync.RWMutex

	client        *dns.Client
	lastConfig    dns.ClientConfig
	clientConfigL sync.RWMutex
	UDPSize       uint16

	// A list of addresses ("ip:port") which should be used as the resolver
	// list. If none are set then the resolver settings in /etc/resolv.conf are
	// used
	ResolverAddrs []string

	// If non-nill, will be called on messages returned from dns servers prior
	// to them being processed (i.e. before they are cached, sorted,
	// ip-replaced, etc...)
	Preprocess func(*dns.Msg)
}

// EnableCacheLast is used to make SRVClient cache the last successful SRV
// response for each domain requested, and if the next request results in some
// kind of error it will use that last response instead.
func (sc *SRVClient) EnableCacheLast() {
	sc.cacheLastL.Lock()
	if sc.cacheLast == nil {
		sc.cacheLast = map[string]*dns.Msg{}
	}
	sc.cacheLastL.Unlock()
}

// DefaultSRVClient is an instance of SRVClient with all zero'd values, used as
// the default client for all global methods. It can be overwritten prior to any
// of the methods being used in order to modify their behavior
var DefaultSRVClient = new(SRVClient)

func replaceSRVTarget(r *dns.SRV, extra []dns.RR) *dns.SRV {
	for _, e := range extra {
		if eA, ok := e.(*dns.A); ok && eA.Hdr.Name == r.Target {
			r.Target = eA.A.String()
		} else if eAAAA, ok := e.(*dns.AAAA); ok && eAAAA.Hdr.Name == r.Target {
			r.Target = eAAAA.AAAA.String()
		}
	}
	return r
}

func (sc *SRVClient) doCacheLast(hostname string, res *dns.Msg) *dns.Msg {
	if sc.cacheLast == nil {
		return res
	}

	if res == nil || len(res.Answer) == 0 {
		sc.cacheLastL.RLock()
		defer sc.cacheLastL.RUnlock()
		if cres, ok := sc.cacheLast[hostname]; ok {
			res = cres
		}
		return res
	}

	sc.cacheLastL.Lock()
	defer sc.cacheLastL.Unlock()
	sc.cacheLast[hostname] = res
	return res
}

func (sc *SRVClient) newClient(cfg dns.ClientConfig) *dns.Client {
	c := new(dns.Client)
	if sc.UDPSize != 0 {
		c.UDPSize = sc.UDPSize
	} else {
		c.UDPSize = dns.DefaultMsgSize
	}
	c.SingleInflight = true
	if cfg.Timeout > 0 {
		timeout := time.Duration(cfg.Timeout) * time.Second
		c.DialTimeout = timeout
		c.ReadTimeout = timeout
		c.WriteTimeout = timeout
	}
	return c
}

func (sc *SRVClient) clientConfig() (*dns.Client, dns.ClientConfig, error) {
	cfg, err := dnsGetConfig()
	if err != nil {
		return nil, cfg, err
	} else if len(sc.ResolverAddrs) > 0 {
		cfg.Servers = sc.ResolverAddrs
	}

	sc.clientConfigL.Lock()
	defer sc.clientConfigL.Unlock()
	if sc.client == nil || !reflect.DeepEqual(sc.lastConfig, cfg) {
		sc.client = sc.newClient(cfg)
		sc.lastConfig = cfg
	}

	return sc.client, sc.lastConfig, nil
}

func (sc *SRVClient) doExchange(c *dns.Client, fqdn, server string) (*dns.Msg, error) {
	m := new(dns.Msg)
	m.SetQuestion(fqdn, dns.TypeSRV)
	if c.UDPSize != 0 {
		m.SetEdns0(c.UDPSize, false)
	}

	res, _, err := c.Exchange(m, server)
	if err != nil {
		// return the res also in case it was a truncation error
		return res, err
	}
	if res.Rcode != dns.RcodeFormatError || c.UDPSize == 0 {
		return res, nil
	}

	// At this point we got a response, but it was just to tell us that
	// edns0 isn't supported, so we try again without it
	m2 := new(dns.Msg)
	m2.SetQuestion(fqdn, dns.TypeSRV)
	res, _, err = c.Exchange(m2, server)
	return res, err
}

func answersFromMsg(m *dns.Msg, replaceWithIPs bool) []*dns.SRV {
	ans := make([]*dns.SRV, 0, len(m.Answer))
	for i := range m.Answer {
		if ansSRV, ok := m.Answer[i].(*dns.SRV); ok {
			if replaceWithIPs {
				// attempt to replace SRV's Target with the actual IP
				ansSRV = replaceSRVTarget(ansSRV, m.Extra)
			}
			ans = append(ans, ansSRV)
		}
	}
	return ans
}

func (sc *SRVClient) lookupSRV(hostname string, replaceWithIPs bool) ([]*dns.SRV, error) {
	c, cfg, err := sc.clientConfig()
	if err != nil {
		return nil, err
	}

	fqdn := dns.Fqdn(hostname)
	var res *dns.Msg
	var tres *dns.Msg
	for _, server := range cfg.Servers {
		res, err = sc.doExchange(c, fqdn, server)
		if err == dns.ErrTruncated && res != nil {
			tres = res
			continue
		}
		if err != nil || res == nil {
			continue
		}
		// no error so stop
		break
	}

	if sc.Preprocess != nil {
		// preprocess both since we don't know which one we'll use yet
		if res != nil {
			sc.Preprocess(res)
		}
		if tres != nil {
			sc.Preprocess(tres)
		}
	}

	// Handles caching this response if it's a successful one, or replacing res
	// with the last response if not. Does nothing if sc.cacheLast is false.
	res = sc.doCacheLast(hostname, res)

	// if we got a truncated error from a server but it was a success, use it
	// we check this AFTER the cache in case we have a better one in the cache
	if res != nil && res.Rcode != dns.RcodeSuccess && tres != nil && tres.Rcode == dns.RcodeSuccess {
		res = tres
		err = dns.ErrTruncated
	}

	if res == nil {
		if err == nil {
			err = errors.New("no available nameservers")
		}
		return nil, err
	}

	ans := answersFromMsg(res, replaceWithIPs)
	if len(ans) == 0 {
		return nil, &ErrNotFound{hostname}
	}

	return ans, err
}

func srvToStr(srv *dns.SRV, port string) string {
	if port == "" {
		port = strconv.Itoa(int(srv.Port))
	}
	return net.JoinHostPort(srv.Target, port)
}

// SRV calls the SRV method on the DefaultSRVClient
func SRV(hostname string) (string, error) {
	return DefaultSRVClient.SRV(hostname)
}

func (sc *SRVClient) srv(hostname string, replaceWithIPs bool) (string, error) {
	var portStr string
	if h, p, _ := net.SplitHostPort(hostname); p != "" && h != "" {
		// check for host being an IP and if so, just return what they sent
		if ip := net.ParseIP(h); ip != nil {
			return hostname, nil
		}
		hostname = h
		portStr = p
	}

	ans, err := sc.lookupSRV(hostname, replaceWithIPs)
	// ignore truncated errors unless there were no answers
	if err != nil && (err != dns.ErrTruncated || len(ans) == 0) {
		return "", err
	}

	srv := pickSRV(ans)

	return srvToStr(srv, portStr), err
}

// SRV will perform a SRV request on the given hostname, and then choose one of
// the returned entries randomly based on the priority and weight fields it
// sees. It will return the address ("host:port") of the winning entry, or an
// error if the query couldn't be made or it returned no entries. If the DNS
// server provided the A records for the hosts, then the result will have the
// target replaced with its respective IP.
//
// If the given hostname already has a ":port" appended to it, only the ip will
// be looked up from the SRV request, but the port given will be returned
//
// If the given hostname is "ip:port", it'll just immediately return what you
// sent.
func (sc *SRVClient) SRV(hostname string) (string, error) {
	return sc.srv(hostname, true)
}

// SRVNoTranslate calls the SRVNoTranslate method on the DefaultSRVClient
func SRVNoTranslate(hostname string) (string, error) {
	return DefaultSRVClient.SRVNoTranslate(hostname)
}

// SRVNoTranslate is exactly like SRV except it won't translate names to their
// respective IPs
func (sc *SRVClient) SRVNoTranslate(hostname string) (string, error) {
	return sc.srv(hostname, false)
}

// SRVNoPort calls the SRVNoPort method on the DefaultSRVClient
func SRVNoPort(hostname string) (string, error) {
	return DefaultSRVClient.SRVNoPort(hostname)
}

// SRVNoPort behaves the same as SRV, but the returned address string will not
// contain the port
func (sc *SRVClient) SRVNoPort(hostname string) (string, error) {
	addr, err := sc.SRV(hostname)
	if err != nil {
		return "", err
	}

	host, _, err := net.SplitHostPort(addr)
	return host, err
}

// AllSRV calls the AllSRV method on the DefaultSRVClient
func AllSRV(hostname string) ([]string, error) {
	return DefaultSRVClient.AllSRV(hostname)
}

// AllSRV returns the list of all hostnames and ports for the SRV lookup
// The results are sorted by priority and then weight. Like SRV, if hostname
// contained a port then the port on all results will be replaced with the
// originally-passed port
// AllSRV will NOT replace hostnames with their respective IPs
func (sc *SRVClient) AllSRV(hostname string) ([]string, error) {
	var ogPort string
	if parts := strings.Split(hostname, ":"); len(parts) == 2 {
		hostname = parts[0]
		ogPort = parts[1]
	}

	ans, err := sc.lookupSRV(hostname, false)
	// ignore truncated errors unless there were no answers
	if err != nil && (err != dns.ErrTruncated || len(ans) == 0) {
		return nil, err
	}

	sort.Sort(sortableSRV(ans))

	res := make([]string, len(ans))
	for i := range ans {
		res[i] = srvToStr(ans[i], ogPort)
	}
	return res, err
}

// MaybeSRV calls the MaybeSRV method on the DefaultSRVClient
func MaybeSRV(host string) string {
	return DefaultSRVClient.MaybeSRV(host)
}

// MaybeSRV attempts a SRV lookup if the host doesn't contain a port and if the
// SRV lookup succeeds it'll rewrite the host and return it with the lookup
// result. If it fails it'll just return the host originally sent
func (sc *SRVClient) MaybeSRV(host string) string {
	if _, p, _ := net.SplitHostPort(host); p == "" {
		if addr, err := sc.SRV(host); err == nil {
			host = addr
		}
	}
	return host
}

func pickSRV(srvs []*dns.SRV) *dns.SRV {
	randSrc := rand.NewSource(time.Now().UnixNano())
	rand := rand.New(randSrc)

	lowPrio := srvs[0].Priority
	picks := make([]*dns.SRV, 0, len(srvs))
	weights := make([]int, 0, len(srvs))

	for i := range srvs {
		if srvs[i].Priority < lowPrio {
			picks = picks[:0]
			weights = weights[:0]
			lowPrio = srvs[i].Priority
		}

		if srvs[i].Priority == lowPrio {
			picks = append(picks, srvs[i])
			weights = append(weights, int(srvs[i].Weight))
		}
	}

	sum := 0
	for i := range weights {
		sum += weights[i]
	}

	if sum == 0 {
		return picks[rand.Intn(len(picks))]
	}

	r := rand.Intn(sum)
	for i := range weights {
		r -= weights[i]
		if r < 0 {
			return picks[i]
		}
	}

	// We should never get here, just return the first pick
	return picks[0]
}

// MaybeSRVURL calls the MaybeSRVURL method on the DefaultSRVClient
func MaybeSRVURL(host string) string {
	return DefaultSRVClient.MaybeSRVURL(host)
}

// MaybeSRVURL calls MaybeSRV and also prepends http:// if no scheme was sent
func (sc *SRVClient) MaybeSRVURL(host string) string {
	host = sc.MaybeSRV(host)
	if !strings.Contains(host, "://") {
		return "http://" + host
	}
	return host
}
<commit_msg>Switch to stable sort so server order is respected<commit_after>package srvclient

// At the moment go's dns resolver which is built into the net package doesn't
// properly handle the case of a response being too big. Which leads us to
// having to manually parse /etc/resolv.conf and manually make the SRV requests.

import (
	"errors"
	"math/rand"
	"reflect"
	"strconv"
	"strings"
	"sync"
	"time"

	"net"
	"sort"

	"github.com/miekg/dns"
)

func init() {
	go dnsConfigLoop()
}

// SRVClient is a holder for methods related to SRV lookups. Use new(SRVClient)
// to initialize one.
type SRVClient struct {
	cacheLast  map[string]*dns.Msg
	cacheLastL sync.RWMutex

	client        *dns.Client
	lastConfig    dns.ClientConfig
	clientConfigL sync.RWMutex
	UDPSize       uint16

	// A list of addresses ("ip:port") which should be used as the resolver
	// list. If none are set then the resolver settings in /etc/resolv.conf are
	// used
	ResolverAddrs []string

	// If non-nill, will be called on messages returned from dns servers prior
	// to them being processed (i.e. before they are cached, sorted,
	// ip-replaced, etc...)
	Preprocess func(*dns.Msg)
}

// EnableCacheLast is used to make SRVClient cache the last successful SRV
// response for each domain requested, and if the next request results in some
// kind of error it will use that last response instead.
func (sc *SRVClient) EnableCacheLast() {
	sc.cacheLastL.Lock()
	if sc.cacheLast == nil {
		sc.cacheLast = map[string]*dns.Msg{}
	}
	sc.cacheLastL.Unlock()
}

// DefaultSRVClient is an instance of SRVClient with all zero'd values, used as
// the default client for all global methods. It can be overwritten prior to any
// of the methods being used in order to modify their behavior
var DefaultSRVClient = new(SRVClient)

func replaceSRVTarget(r *dns.SRV, extra []dns.RR) *dns.SRV {
	for _, e := range extra {
		if eA, ok := e.(*dns.A); ok && eA.Hdr.Name == r.Target {
			r.Target = eA.A.String()
		} else if eAAAA, ok := e.(*dns.AAAA); ok && eAAAA.Hdr.Name == r.Target {
			r.Target = eAAAA.AAAA.String()
		}
	}
	return r
}

func (sc *SRVClient) doCacheLast(hostname string, res *dns.Msg) *dns.Msg {
	if sc.cacheLast == nil {
		return res
	}

	if res == nil || len(res.Answer) == 0 {
		sc.cacheLastL.RLock()
		defer sc.cacheLastL.RUnlock()
		if cres, ok := sc.cacheLast[hostname]; ok {
			res = cres
		}
		return res
	}

	sc.cacheLastL.Lock()
	defer sc.cacheLastL.Unlock()
	sc.cacheLast[hostname] = res
	return res
}

func (sc *SRVClient) newClient(cfg dns.ClientConfig) *dns.Client {
	c := new(dns.Client)
	if sc.UDPSize != 0 {
		c.UDPSize = sc.UDPSize
	} else {
		c.UDPSize = dns.DefaultMsgSize
	}
	c.SingleInflight = true
	if cfg.Timeout > 0 {
		timeout := time.Duration(cfg.Timeout) * time.Second
		c.DialTimeout = timeout
		c.ReadTimeout = timeout
		c.WriteTimeout = timeout
	}
	return c
}

func (sc *SRVClient) clientConfig() (*dns.Client, dns.ClientConfig, error) {
	cfg, err := dnsGetConfig()
	if err != nil {
		return nil, cfg, err
	} else if len(sc.ResolverAddrs) > 0 {
		cfg.Servers = sc.ResolverAddrs
	}

	sc.clientConfigL.Lock()
	defer sc.clientConfigL.Unlock()
	if sc.client == nil || !reflect.DeepEqual(sc.lastConfig, cfg) {
		sc.client = sc.newClient(cfg)
		sc.lastConfig = cfg
	}

	return sc.client, sc.lastConfig, nil
}

func (sc *SRVClient) doExchange(c *dns.Client, fqdn, server string) (*dns.Msg, error) {
	m := new(dns.Msg)
	m.SetQuestion(fqdn, dns.TypeSRV)
	if c.UDPSize != 0 {
		m.SetEdns0(c.UDPSize, false)
	}

	res, _, err := c.Exchange(m, server)
	if err != nil {
		// return the res also in case it was a truncation error
		return res, err
	}
	if res.Rcode != dns.RcodeFormatError || c.UDPSize == 0 {
		return res, nil
	}

	// At this point we got a response, but it was just to tell us that
	// edns0 isn't supported, so we try again without it
	m2 := new(dns.Msg)
	m2.SetQuestion(fqdn, dns.TypeSRV)
	res, _, err = c.Exchange(m2, server)
	return res, err
}

func answersFromMsg(m *dns.Msg, replaceWithIPs bool) []*dns.SRV {
	ans := make([]*dns.SRV, 0, len(m.Answer))
	for i := range m.Answer {
		if ansSRV, ok := m.Answer[i].(*dns.SRV); ok {
			if replaceWithIPs {
				// attempt to replace SRV's Target with the actual IP
				ansSRV = replaceSRVTarget(ansSRV, m.Extra)
			}
			ans = append(ans, ansSRV)
		}
	}
	return ans
}

func (sc *SRVClient) lookupSRV(hostname string, replaceWithIPs bool) ([]*dns.SRV, error) {
	c, cfg, err := sc.clientConfig()
	if err != nil {
		return nil, err
	}

	fqdn := dns.Fqdn(hostname)
	var res *dns.Msg
	var tres *dns.Msg
	for _, server := range cfg.Servers {
		res, err = sc.doExchange(c, fqdn, server)
		if err == dns.ErrTruncated && res != nil {
			tres = res
			continue
		}
		if err != nil || res == nil {
			continue
		}
		// no error so stop
		break
	}

	if sc.Preprocess != nil {
		// preprocess both since we don't know which one we'll use yet
		if res != nil {
			sc.Preprocess(res)
		}
		if tres != nil {
			sc.Preprocess(tres)
		}
	}

	// Handles caching this response if it's a successful one, or replacing res
	// with the last response if not. Does nothing if sc.cacheLast is false.
	res = sc.doCacheLast(hostname, res)

	// if we got a truncated error from a server but it was a success, use it
	// we check this AFTER the cache in case we have a better one in the cache
	if res != nil && res.Rcode != dns.RcodeSuccess && tres != nil && tres.Rcode == dns.RcodeSuccess {
		res = tres
		err = dns.ErrTruncated
	}

	if res == nil {
		if err == nil {
			err = errors.New("no available nameservers")
		}
		return nil, err
	}

	ans := answersFromMsg(res, replaceWithIPs)
	if len(ans) == 0 {
		return nil, &ErrNotFound{hostname}
	}

	return ans, err
}

func srvToStr(srv *dns.SRV, port string) string {
	if port == "" {
		port = strconv.Itoa(int(srv.Port))
	}
	return net.JoinHostPort(srv.Target, port)
}

// SRV calls the SRV method on the DefaultSRVClient
func SRV(hostname string) (string, error) {
	return DefaultSRVClient.SRV(hostname)
}

func (sc *SRVClient) srv(hostname string, replaceWithIPs bool) (string, error) {
	var portStr string
	if h, p, _ := net.SplitHostPort(hostname); p != "" && h != "" {
		// check for host being an IP and if so, just return what they sent
		if ip := net.ParseIP(h); ip != nil {
			return hostname, nil
		}
		hostname = h
		portStr = p
	}

	ans, err := sc.lookupSRV(hostname, replaceWithIPs)
	// ignore truncated errors unless there were no answers
	if err != nil && (err != dns.ErrTruncated || len(ans) == 0) {
		return "", err
	}

	srv := pickSRV(ans)

	return srvToStr(srv, portStr), err
}

// SRV will perform a SRV request on the given hostname, and then choose one of
// the returned entries randomly based on the priority and weight fields it
// sees. It will return the address ("host:port") of the winning entry, or an
// error if the query couldn't be made or it returned no entries. If the DNS
// server provided the A records for the hosts, then the result will have the
// target replaced with its respective IP.
//
// If the given hostname already has a ":port" appended to it, only the ip will
// be looked up from the SRV request, but the port given will be returned
//
// If the given hostname is "ip:port", it'll just immediately return what you
// sent.
func (sc *SRVClient) SRV(hostname string) (string, error) {
	return sc.srv(hostname, true)
}

// SRVNoTranslate calls the SRVNoTranslate method on the DefaultSRVClient
func SRVNoTranslate(hostname string) (string, error) {
	return DefaultSRVClient.SRVNoTranslate(hostname)
}

// SRVNoTranslate is exactly like SRV except it won't translate names to their
// respective IPs
func (sc *SRVClient) SRVNoTranslate(hostname string) (string, error) {
	return sc.srv(hostname, false)
}

// SRVNoPort calls the SRVNoPort method on the DefaultSRVClient
func SRVNoPort(hostname string) (string, error) {
	return DefaultSRVClient.SRVNoPort(hostname)
}

// SRVNoPort behaves the same as SRV, but the returned address string will not
// contain the port
func (sc *SRVClient) SRVNoPort(hostname string) (string, error) {
	addr, err := sc.SRV(hostname)
	if err != nil {
		return "", err
	}

	host, _, err := net.SplitHostPort(addr)
	return host, err
}

// AllSRV calls the AllSRV method on the DefaultSRVClient
func AllSRV(hostname string) ([]string, error) {
	return DefaultSRVClient.AllSRV(hostname)
}

// AllSRV returns the list of all hostnames and ports for the SRV lookup
// The results are sorted by priority and then weight. Like SRV, if hostname
// contained a port then the port on all results will be replaced with the
// originally-passed port
// AllSRV will NOT replace hostnames with their respective IPs
func (sc *SRVClient) AllSRV(hostname string) ([]string, error) {
	var ogPort string
	if parts := strings.Split(hostname, ":"); len(parts) == 2 {
		hostname = parts[0]
		ogPort = parts[1]
	}

	ans, err := sc.lookupSRV(hostname, false)
	// ignore truncated errors unless there were no answers
	if err != nil && (err != dns.ErrTruncated || len(ans) == 0) {
		return nil, err
	}

	// sort the lowest priority to the front and if priorities match
	// sort the highest weights to the front
	// use a stable sort in case the server's order is meaningful
	sort.SliceStable(ans, func(i, j int) bool {
		if ans[i].Priority == ans[j].Priority {
			return ans[i].Weight > ans[j].Weight
		}
		return ans[i].Priority < ans[j].Priority
	})

	res := make([]string, len(ans))
	for i := range ans {
		res[i] = srvToStr(ans[i], ogPort)
	}
	return res, err
}

// MaybeSRV calls the MaybeSRV method on the DefaultSRVClient
func MaybeSRV(host string) string {
	return DefaultSRVClient.MaybeSRV(host)
}

// MaybeSRV attempts a SRV lookup if the host doesn't contain a port and if the
// SRV lookup succeeds it'll rewrite the host and return it with the lookup
// result. If it fails it'll just return the host originally sent
func (sc *SRVClient) MaybeSRV(host string) string {
	if _, p, _ := net.SplitHostPort(host); p == "" {
		if addr, err := sc.SRV(host); err == nil {
			host = addr
		}
	}
	return host
}

func pickSRV(srvs []*dns.SRV) *dns.SRV {
	randSrc := rand.NewSource(time.Now().UnixNano())
	rand := rand.New(randSrc)

	lowPrio := srvs[0].Priority
	picks := make([]*dns.SRV, 0, len(srvs))
	weights := make([]int, 0, len(srvs))

	for i := range srvs {
		if srvs[i].Priority < lowPrio {
			picks = picks[:0]
			weights = weights[:0]
			lowPrio = srvs[i].Priority
		}

		if srvs[i].Priority == lowPrio {
			picks = append(picks, srvs[i])
			weights = append(weights, int(srvs[i].Weight))
		}
	}

	sum := 0
	for i := range weights {
		sum += weights[i]
	}

	if sum == 0 {
		return picks[rand.Intn(len(picks))]
	}

	r := rand.Intn(sum)
	for i := range weights {
		r -= weights[i]
		if r < 0 {
			return picks[i]
		}
	}

	// We should never get here, just return the first pick
	return picks[0]
}

// MaybeSRVURL calls the MaybeSRVURL method on the DefaultSRVClient
func MaybeSRVURL(host string) string {
	return DefaultSRVClient.MaybeSRVURL(host)
}

// MaybeSRVURL calls MaybeSRV and also prepends http:// if no scheme was sent
func (sc *SRVClient) MaybeSRVURL(host string) string {
	host = sc.MaybeSRV(host)
	if !strings.Contains(host, "://") {
		return "http://" + host
	}
	return host
}
<|endoftext|>


<commit_before>package main

import "github.com/dmulholl/janus-go/janus"

import (
    "fmt"
    "os"
    "path/filepath"
)

const version = "2.2.0.dev1"

var helptext = fmt.Sprintf(`
Usage: %s [FLAGS] [COMMAND]

  A utility for creating and managing encrypted password databases.

Flags:
  -h, --help        Print the application's help text.
  -v, --version     Print the application's version number.

Basic Commands:
  add               Add a new entry to a password database.
  edit              Edit an existing database entry.
  gen               Generate a new random password.
  init              Initialize a new password database.
  list              List database entries.
  pass              Copy a password to the clipboard.
  retire            Mark one or more entries as inactive.
  url               Copy a url to the clipboard.
  user              Copy a username to the clipboard.

Additional Commands:
  config            Set or print a configuration option.
  decrypt           Decrypt a file.
  dump              Dump a database's internal JSON data store.
  encrypt           Encrypt a file.
  export            Export entries from a database.
  import            Import entries into a database.
  purge             Purge inactive entries from a database.
  restore           Restore inactive entries to active status.
  setcachepass      Change a database's cache password.
  setmasterpass     Change a database's master password.
  tags              List database tags.

Aliases:
  new               Alias for 'add'.
  show              Alias for 'list --verbose'.

Command Help:
  help <command>    Print the specified command's help text and exit.
`, filepath.Base(os.Args[0]))

func main() {

    // Instantiate an argument parser.
    parser := janus.NewParser()
    parser.Helptext = helptext
    parser.Version = version

    // Register commands.
    registerAddCmd(parser)
    registerCacheCmd(parser)
    registerConfigCmd(parser)
    registerDecryptCmd(parser)
    registerDumpCmd(parser)
    registerEditCmd(parser)
    registerEncryptCmd(parser)
    registerExportCmd(parser)
    registerGenCmd(parser)
    registerImportCmd(parser)
    registerInitCmd(parser)
    registerListCmd(parser)
    registerPassCmd(parser)
    registerPurgeCmd(parser)
    registerRestoreCmd(parser)
    registerRetireCmd(parser)
    registerSetMasterPassCmd(parser)
    registerSetCachePassCmd(parser)
    registerTagsCmd(parser)
    registerUrlCmd(parser)
    registerUserCmd(parser)

    // Parse the command line arguments.
    parser.Parse()
    if !parser.HasCmd() {
        parser.ExitHelp()
    }
}
<commit_msg>Bump version to 2.2.0<commit_after>package main

import "github.com/dmulholl/janus-go/janus"

import (
    "fmt"
    "os"
    "path/filepath"
)

const version = "2.2.0"

var helptext = fmt.Sprintf(`
Usage: %s [FLAGS] [COMMAND]

  A utility for creating and managing encrypted password databases.

Flags:
  -h, --help        Print the application's help text.
  -v, --version     Print the application's version number.

Basic Commands:
  add               Add a new entry to a password database.
  edit              Edit an existing database entry.
  gen               Generate a new random password.
  init              Initialize a new password database.
  list              List database entries.
  pass              Copy a password to the clipboard.
  retire            Mark one or more entries as inactive.
  url               Copy a url to the clipboard.
  user              Copy a username to the clipboard.

Additional Commands:
  config            Set or print a configuration option.
  decrypt           Decrypt a file.
  dump              Dump a database's internal JSON data store.
  encrypt           Encrypt a file.
  export            Export entries from a database.
  import            Import entries into a database.
  purge             Purge inactive entries from a database.
  restore           Restore inactive entries to active status.
  setcachepass      Change a database's cache password.
  setmasterpass     Change a database's master password.
  tags              List database tags.

Aliases:
  new               Alias for 'add'.
  show              Alias for 'list --verbose'.

Command Help:
  help <command>    Print the specified command's help text and exit.
`, filepath.Base(os.Args[0]))

func main() {

    // Instantiate an argument parser.
    parser := janus.NewParser()
    parser.Helptext = helptext
    parser.Version = version

    // Register commands.
    registerAddCmd(parser)
    registerCacheCmd(parser)
    registerConfigCmd(parser)
    registerDecryptCmd(parser)
    registerDumpCmd(parser)
    registerEditCmd(parser)
    registerEncryptCmd(parser)
    registerExportCmd(parser)
    registerGenCmd(parser)
    registerImportCmd(parser)
    registerInitCmd(parser)
    registerListCmd(parser)
    registerPassCmd(parser)
    registerPurgeCmd(parser)
    registerRestoreCmd(parser)
    registerRetireCmd(parser)
    registerSetMasterPassCmd(parser)
    registerSetCachePassCmd(parser)
    registerTagsCmd(parser)
    registerUrlCmd(parser)
    registerUserCmd(parser)

    // Parse the command line arguments.
    parser.Parse()
    if !parser.HasCmd() {
        parser.ExitHelp()
    }
}
<|endoftext|>


<commit_before>/*
Copyright 2019 The Jetstack cert-manager 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 gen

import (
	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1alpha2"
	cmmeta "github.com/jetstack/cert-manager/pkg/apis/meta/v1"
)

type OrderModifier func(*cmacme.Order)

func Order(name string, mods ...OrderModifier) *cmacme.Order {
	c := &cmacme.Order{
		ObjectMeta: ObjectMeta(name),
	}
	for _, mod := range mods {
		mod(c)
	}
	return c
}

func OrderFrom(crt *cmacme.Order, mods ...OrderModifier) *cmacme.Order {
	crt = crt.DeepCopy()
	for _, mod := range mods {
		mod(crt)
	}
	return crt
}

// SetIssuer sets the Order.spec.issuerRef field
func SetOrderIssuer(o cmmeta.ObjectReference) OrderModifier {
	return func(c *cmacme.Order) {
		c.Spec.IssuerRef = o
	}
}

func SetOrderDNSNames(dnsNames ...string) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.Spec.DNSNames = dnsNames
	}
}

func SetOrderURL(url string) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.Status.URL = url
	}
}

func SetOrderState(s cmacme.State) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.Status.State = s
	}
}

func SetOrderReason(reason string) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.Status.Reason = reason
	}
}

func SetOrderStatus(s cmacme.OrderStatus) OrderModifier {
	return func(o *cmacme.Order) {
		o.Status = s
	}
}

func SetOrderCertificate(d []byte) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.Status.Certificate = d
	}
}

func SetOrderCommonName(commonName string) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.Spec.CommonName = commonName
	}
}

func SetOrderNamespace(namespace string) OrderModifier {
	return func(crt *cmacme.Order) {
		crt.ObjectMeta.Namespace = namespace
	}
}
<commit_msg>User order as variable name in gen/order.go<commit_after>/*
Copyright 2019 The Jetstack cert-manager 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 gen

import (
	cmacme "github.com/jetstack/cert-manager/pkg/apis/acme/v1alpha2"
	cmmeta "github.com/jetstack/cert-manager/pkg/apis/meta/v1"
)

type OrderModifier func(*cmacme.Order)

func Order(name string, mods ...OrderModifier) *cmacme.Order {
	order := &cmacme.Order{
		ObjectMeta: ObjectMeta(name),
	}
	for _, mod := range mods {
		mod(order)
	}
	return order
}

func OrderFrom(order *cmacme.Order, mods ...OrderModifier) *cmacme.Order {
	order = order.DeepCopy()
	for _, mod := range mods {
		mod(order)
	}
	return order
}

// SetIssuer sets the Order.spec.issuerRef field
func SetOrderIssuer(o cmmeta.ObjectReference) OrderModifier {
	return func(order *cmacme.Order) {
		order.Spec.IssuerRef = o
	}
}

func SetOrderDNSNames(dnsNames ...string) OrderModifier {
	return func(order *cmacme.Order) {
		order.Spec.DNSNames = dnsNames
	}
}

func SetOrderURL(url string) OrderModifier {
	return func(order *cmacme.Order) {
		order.Status.URL = url
	}
}

func SetOrderState(s cmacme.State) OrderModifier {
	return func(order *cmacme.Order) {
		order.Status.State = s
	}
}

func SetOrderReason(reason string) OrderModifier {
	return func(order *cmacme.Order) {
		order.Status.Reason = reason
	}
}

func SetOrderStatus(s cmacme.OrderStatus) OrderModifier {
	return func(order *cmacme.Order) {
		order.Status = s
	}
}

func SetOrderCertificate(d []byte) OrderModifier {
	return func(order *cmacme.Order) {
		order.Status.Certificate = d
	}
}

func SetOrderCommonName(commonName string) OrderModifier {
	return func(order *cmacme.Order) {
		order.Spec.CommonName = commonName
	}
}

func SetOrderNamespace(namespace string) OrderModifier {
	return func(order *cmacme.Order) {
		order.ObjectMeta.Namespace = namespace
	}
}
<|endoftext|>


<commit_before>package streams

import (
	"context"
	"sync"
	"time"

	"github.com/matrix-org/dendrite/syncapi/types"
	userapi "github.com/matrix-org/dendrite/userapi/api"
	"github.com/matrix-org/gomatrixserverlib"
	"go.uber.org/atomic"
)

// The max number of per-room goroutines to have running.
// Too high and this will consume lots of CPU, too low and complete
// sync responses will take longer to process.
const PDU_STREAM_WORKERS = 256

// The maximum number of tasks that can be queued in total before
// backpressure will build up and the rests will start to block.
const PDU_STREAM_QUEUESIZE = PDU_STREAM_WORKERS * 8

type PDUStreamProvider struct {
	StreamProvider

	tasks   chan func()
	workers atomic.Int32
}

func (p *PDUStreamProvider) worker() {
	defer p.workers.Dec()
	for {
		select {
		case f := <-p.tasks:
			f()
		case <-time.After(time.Second * 10):
			return
		}
	}
}

func (p *PDUStreamProvider) queue(f func()) {
	if p.workers.Load() < PDU_STREAM_WORKERS {
		p.workers.Inc()
		go p.worker()
	}
	p.tasks <- f
}

func (p *PDUStreamProvider) Setup() {
	p.StreamProvider.Setup()
	p.tasks = make(chan func(), PDU_STREAM_QUEUESIZE)

	p.latestMutex.Lock()
	defer p.latestMutex.Unlock()

	id, err := p.DB.MaxStreamPositionForPDUs(context.Background())
	if err != nil {
		panic(err)
	}
	p.latest = id
}

func (p *PDUStreamProvider) CompleteSync(
	ctx context.Context,
	req *types.SyncRequest,
) types.StreamPosition {
	from := types.StreamPosition(0)
	to := p.LatestPosition(ctx)

	// Get the current sync position which we will base the sync response on.
	// For complete syncs, we want to start at the most recent events and work
	// backwards, so that we show the most recent events in the room.
	r := types.Range{
		From:      to,
		To:        0,
		Backwards: true,
	}

	// Extract room state and recent events for all rooms the user is joined to.
	joinedRoomIDs, err := p.DB.RoomIDsWithMembership(ctx, req.Device.UserID, gomatrixserverlib.Join)
	if err != nil {
		req.Log.WithError(err).Error("p.DB.RoomIDsWithMembership failed")
		return from
	}

	stateFilter := req.Filter.Room.State
	eventFilter := req.Filter.Room.Timeline

	// Build up a /sync response. Add joined rooms.
	var reqMutex sync.Mutex
	var reqWaitGroup sync.WaitGroup
	reqWaitGroup.Add(len(joinedRoomIDs))
	for _, room := range joinedRoomIDs {
		roomID := room
		p.queue(func() {
			defer reqWaitGroup.Done()

			var jr *types.JoinResponse
			jr, err = p.getJoinResponseForCompleteSync(
				ctx, roomID, r, &stateFilter, &eventFilter, req.WantFullState, req.Device,
			)
			if err != nil {
				req.Log.WithError(err).Error("p.getJoinResponseForCompleteSync failed")
				return
			}

			reqMutex.Lock()
			defer reqMutex.Unlock()
			req.Response.Rooms.Join[roomID] = *jr
			req.Rooms[roomID] = gomatrixserverlib.Join
		})
	}

	reqWaitGroup.Wait()

	// Add peeked rooms.
	peeks, err := p.DB.PeeksInRange(ctx, req.Device.UserID, req.Device.ID, r)
	if err != nil {
		req.Log.WithError(err).Error("p.DB.PeeksInRange failed")
		return from
	}
	for _, peek := range peeks {
		if !peek.Deleted {
			var jr *types.JoinResponse
			jr, err = p.getJoinResponseForCompleteSync(
				ctx, peek.RoomID, r, &stateFilter, &eventFilter, req.WantFullState, req.Device,
			)
			if err != nil {
				req.Log.WithError(err).Error("p.getJoinResponseForCompleteSync failed")
				return from
			}
			req.Response.Rooms.Peek[peek.RoomID] = *jr
		}
	}

	return to
}

func (p *PDUStreamProvider) IncrementalSync(
	ctx context.Context,
	req *types.SyncRequest,
	from, to types.StreamPosition,
) (newPos types.StreamPosition) {
	r := types.Range{
		From:      from,
		To:        to,
		Backwards: from > to,
	}
	newPos = to

	var err error
	var stateDeltas []types.StateDelta
	var joinedRooms []string

	stateFilter := req.Filter.Room.State
	eventFilter := req.Filter.Room.Timeline

	if req.WantFullState {
		if stateDeltas, joinedRooms, err = p.DB.GetStateDeltasForFullStateSync(ctx, req.Device, r, req.Device.UserID, &stateFilter); err != nil {
			req.Log.WithError(err).Error("p.DB.GetStateDeltasForFullStateSync failed")
			return
		}
	} else {
		if stateDeltas, joinedRooms, err = p.DB.GetStateDeltas(ctx, req.Device, r, req.Device.UserID, &stateFilter); err != nil {
			req.Log.WithError(err).Error("p.DB.GetStateDeltas failed")
			return
		}
	}

	for _, roomID := range joinedRooms {
		req.Rooms[roomID] = gomatrixserverlib.Join
	}

	for _, delta := range stateDeltas {
		if err = p.addRoomDeltaToResponse(ctx, req.Device, r, delta, &eventFilter, req.Response); err != nil {
			req.Log.WithError(err).Error("d.addRoomDeltaToResponse failed")
			return newPos
		}
	}

	return r.To
}

func (p *PDUStreamProvider) addRoomDeltaToResponse(
	ctx context.Context,
	device *userapi.Device,
	r types.Range,
	delta types.StateDelta,
	eventFilter *gomatrixserverlib.RoomEventFilter,
	res *types.Response,
) error {
	if delta.MembershipPos > 0 && delta.Membership == gomatrixserverlib.Leave {
		// make sure we don't leak recent events after the leave event.
		// TODO: History visibility makes this somewhat complex to handle correctly. For example:
		// TODO: This doesn't work for join -> leave in a single /sync request (see events prior to join).
		// TODO: This will fail on join -> leave -> sensitive msg -> join -> leave
		//       in a single /sync request
		// This is all "okay" assuming history_visibility == "shared" which it is by default.
		r.To = delta.MembershipPos
	}
	recentStreamEvents, limited, err := p.DB.RecentEvents(
		ctx, delta.RoomID, r,
		eventFilter, true, true,
	)
	if err != nil {
		return err
	}
	recentEvents := p.DB.StreamEventsToEvents(device, recentStreamEvents)
	delta.StateEvents = removeDuplicates(delta.StateEvents, recentEvents) // roll back
	prevBatch, err := p.DB.GetBackwardTopologyPos(ctx, recentStreamEvents)
	if err != nil {
		return err
	}

	// XXX: should we ever get this far if we have no recent events or state in this room?
	// in practice we do for peeks, but possibly not joins?
	if len(recentEvents) == 0 && len(delta.StateEvents) == 0 {
		return nil
	}

	switch delta.Membership {
	case gomatrixserverlib.Join:
		jr := types.NewJoinResponse()
		jr.Timeline.PrevBatch = &prevBatch
		jr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
		jr.Timeline.Limited = limited
		jr.State.Events = gomatrixserverlib.HeaderedToClientEvents(delta.StateEvents, gomatrixserverlib.FormatSync)
		res.Rooms.Join[delta.RoomID] = *jr

	case gomatrixserverlib.Peek:
		jr := types.NewJoinResponse()
		jr.Timeline.PrevBatch = &prevBatch
		jr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
		jr.Timeline.Limited = limited
		jr.State.Events = gomatrixserverlib.HeaderedToClientEvents(delta.StateEvents, gomatrixserverlib.FormatSync)
		res.Rooms.Peek[delta.RoomID] = *jr

	case gomatrixserverlib.Leave:
		fallthrough // transitions to leave are the same as ban

	case gomatrixserverlib.Ban:
		// TODO: recentEvents may contain events that this user is not allowed to see because they are
		//       no longer in the room.
		lr := types.NewLeaveResponse()
		lr.Timeline.PrevBatch = &prevBatch
		lr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
		lr.Timeline.Limited = false // TODO: if len(events) >= numRecents + 1 and then set limited:true
		lr.State.Events = gomatrixserverlib.HeaderedToClientEvents(delta.StateEvents, gomatrixserverlib.FormatSync)
		res.Rooms.Leave[delta.RoomID] = *lr
	}

	return nil
}

func (p *PDUStreamProvider) getJoinResponseForCompleteSync(
	ctx context.Context,
	roomID string,
	r types.Range,
	stateFilter *gomatrixserverlib.StateFilter,
	eventFilter *gomatrixserverlib.RoomEventFilter,
	wantFullState bool,
	device *userapi.Device,
) (jr *types.JoinResponse, err error) {
	// TODO: When filters are added, we may need to call this multiple times to get enough events.
	//       See: https://github.com/matrix-org/synapse/blob/v0.19.3/synapse/handlers/sync.py#L316
	recentStreamEvents, limited, err := p.DB.RecentEvents(
		ctx, roomID, r, eventFilter, true, true,
	)
	if err != nil {
		return
	}

	// Get the event IDs of the stream events we fetched. There's no point in us
	var excludingEventIDs []string
	if !wantFullState {
		excludingEventIDs = make([]string, 0, len(recentStreamEvents))
		for _, event := range recentStreamEvents {
			if event.StateKey() != nil {
				excludingEventIDs = append(excludingEventIDs, event.EventID())
			}
		}
	}

	stateEvents, err := p.DB.CurrentState(ctx, roomID, stateFilter, excludingEventIDs)
	if err != nil {
		return
	}

	// TODO FIXME: We don't fully implement history visibility yet. To avoid leaking events which the
	// user shouldn't see, we check the recent events and remove any prior to the join event of the user
	// which is equiv to history_visibility: joined
	joinEventIndex := -1
	for i := len(recentStreamEvents) - 1; i >= 0; i-- {
		ev := recentStreamEvents[i]
		if ev.Type() == gomatrixserverlib.MRoomMember && ev.StateKeyEquals(device.UserID) {
			membership, _ := ev.Membership()
			if membership == "join" {
				joinEventIndex = i
				if i > 0 {
					// the create event happens before the first join, so we should cut it at that point instead
					if recentStreamEvents[i-1].Type() == gomatrixserverlib.MRoomCreate && recentStreamEvents[i-1].StateKeyEquals("") {
						joinEventIndex = i - 1
						break
					}
				}
				break
			}
		}
	}
	if joinEventIndex != -1 {
		// cut all events earlier than the join (but not the join itself)
		recentStreamEvents = recentStreamEvents[joinEventIndex:]
		limited = false // so clients know not to try to backpaginate
	}

	// Retrieve the backward topology position, i.e. the position of the
	// oldest event in the room's topology.
	var prevBatch *types.TopologyToken
	if len(recentStreamEvents) > 0 {
		var backwardTopologyPos, backwardStreamPos types.StreamPosition
		backwardTopologyPos, backwardStreamPos, err = p.DB.PositionInTopology(ctx, recentStreamEvents[0].EventID())
		if err != nil {
			return
		}
		prevBatch = &types.TopologyToken{
			Depth:       backwardTopologyPos,
			PDUPosition: backwardStreamPos,
		}
		prevBatch.Decrement()
	}

	// We don't include a device here as we don't need to send down
	// transaction IDs for complete syncs, but we do it anyway because Sytest demands it for:
	// "Can sync a room with a message with a transaction id" - which does a complete sync to check.
	recentEvents := p.DB.StreamEventsToEvents(device, recentStreamEvents)
	stateEvents = removeDuplicates(stateEvents, recentEvents)
	jr = types.NewJoinResponse()
	jr.Timeline.PrevBatch = prevBatch
	jr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
	jr.Timeline.Limited = limited
	jr.State.Events = gomatrixserverlib.HeaderedToClientEvents(stateEvents, gomatrixserverlib.FormatSync)
	return jr, nil
}

func removeDuplicates(stateEvents, recentEvents []*gomatrixserverlib.HeaderedEvent) []*gomatrixserverlib.HeaderedEvent {
	for _, recentEv := range recentEvents {
		if recentEv.StateKey() == nil {
			continue // not a state event
		}
		// TODO: This is a linear scan over all the current state events in this room. This will
		//       be slow for big rooms. We should instead sort the state events by event ID  (ORDER BY)
		//       then do a binary search to find matching events, similar to what roomserver does.
		for j := 0; j < len(stateEvents); j++ {
			if stateEvents[j].EventID() == recentEv.EventID() {
				// overwrite the element to remove with the last element then pop the last element.
				// This is orders of magnitude faster than re-slicing, but doesn't preserve ordering
				// (we don't care about the order of stateEvents)
				stateEvents[j] = stateEvents[len(stateEvents)-1]
				stateEvents = stateEvents[:len(stateEvents)-1]
				break // there shouldn't be multiple events with the same event ID
			}
		}
	}
	return stateEvents
}
<commit_msg>Truncate `recentStreamEvents` before working out which event IDs to exclude from `stateEvents` (#2281)<commit_after>package streams

import (
	"context"
	"sync"
	"time"

	"github.com/matrix-org/dendrite/syncapi/types"
	userapi "github.com/matrix-org/dendrite/userapi/api"
	"github.com/matrix-org/gomatrixserverlib"
	"go.uber.org/atomic"
)

// The max number of per-room goroutines to have running.
// Too high and this will consume lots of CPU, too low and complete
// sync responses will take longer to process.
const PDU_STREAM_WORKERS = 256

// The maximum number of tasks that can be queued in total before
// backpressure will build up and the rests will start to block.
const PDU_STREAM_QUEUESIZE = PDU_STREAM_WORKERS * 8

type PDUStreamProvider struct {
	StreamProvider

	tasks   chan func()
	workers atomic.Int32
}

func (p *PDUStreamProvider) worker() {
	defer p.workers.Dec()
	for {
		select {
		case f := <-p.tasks:
			f()
		case <-time.After(time.Second * 10):
			return
		}
	}
}

func (p *PDUStreamProvider) queue(f func()) {
	if p.workers.Load() < PDU_STREAM_WORKERS {
		p.workers.Inc()
		go p.worker()
	}
	p.tasks <- f
}

func (p *PDUStreamProvider) Setup() {
	p.StreamProvider.Setup()
	p.tasks = make(chan func(), PDU_STREAM_QUEUESIZE)

	p.latestMutex.Lock()
	defer p.latestMutex.Unlock()

	id, err := p.DB.MaxStreamPositionForPDUs(context.Background())
	if err != nil {
		panic(err)
	}
	p.latest = id
}

func (p *PDUStreamProvider) CompleteSync(
	ctx context.Context,
	req *types.SyncRequest,
) types.StreamPosition {
	from := types.StreamPosition(0)
	to := p.LatestPosition(ctx)

	// Get the current sync position which we will base the sync response on.
	// For complete syncs, we want to start at the most recent events and work
	// backwards, so that we show the most recent events in the room.
	r := types.Range{
		From:      to,
		To:        0,
		Backwards: true,
	}

	// Extract room state and recent events for all rooms the user is joined to.
	joinedRoomIDs, err := p.DB.RoomIDsWithMembership(ctx, req.Device.UserID, gomatrixserverlib.Join)
	if err != nil {
		req.Log.WithError(err).Error("p.DB.RoomIDsWithMembership failed")
		return from
	}

	stateFilter := req.Filter.Room.State
	eventFilter := req.Filter.Room.Timeline

	// Build up a /sync response. Add joined rooms.
	var reqMutex sync.Mutex
	var reqWaitGroup sync.WaitGroup
	reqWaitGroup.Add(len(joinedRoomIDs))
	for _, room := range joinedRoomIDs {
		roomID := room
		p.queue(func() {
			defer reqWaitGroup.Done()

			var jr *types.JoinResponse
			jr, err = p.getJoinResponseForCompleteSync(
				ctx, roomID, r, &stateFilter, &eventFilter, req.WantFullState, req.Device,
			)
			if err != nil {
				req.Log.WithError(err).Error("p.getJoinResponseForCompleteSync failed")
				return
			}

			reqMutex.Lock()
			defer reqMutex.Unlock()
			req.Response.Rooms.Join[roomID] = *jr
			req.Rooms[roomID] = gomatrixserverlib.Join
		})
	}

	reqWaitGroup.Wait()

	// Add peeked rooms.
	peeks, err := p.DB.PeeksInRange(ctx, req.Device.UserID, req.Device.ID, r)
	if err != nil {
		req.Log.WithError(err).Error("p.DB.PeeksInRange failed")
		return from
	}
	for _, peek := range peeks {
		if !peek.Deleted {
			var jr *types.JoinResponse
			jr, err = p.getJoinResponseForCompleteSync(
				ctx, peek.RoomID, r, &stateFilter, &eventFilter, req.WantFullState, req.Device,
			)
			if err != nil {
				req.Log.WithError(err).Error("p.getJoinResponseForCompleteSync failed")
				return from
			}
			req.Response.Rooms.Peek[peek.RoomID] = *jr
		}
	}

	return to
}

func (p *PDUStreamProvider) IncrementalSync(
	ctx context.Context,
	req *types.SyncRequest,
	from, to types.StreamPosition,
) (newPos types.StreamPosition) {
	r := types.Range{
		From:      from,
		To:        to,
		Backwards: from > to,
	}
	newPos = to

	var err error
	var stateDeltas []types.StateDelta
	var joinedRooms []string

	stateFilter := req.Filter.Room.State
	eventFilter := req.Filter.Room.Timeline

	if req.WantFullState {
		if stateDeltas, joinedRooms, err = p.DB.GetStateDeltasForFullStateSync(ctx, req.Device, r, req.Device.UserID, &stateFilter); err != nil {
			req.Log.WithError(err).Error("p.DB.GetStateDeltasForFullStateSync failed")
			return
		}
	} else {
		if stateDeltas, joinedRooms, err = p.DB.GetStateDeltas(ctx, req.Device, r, req.Device.UserID, &stateFilter); err != nil {
			req.Log.WithError(err).Error("p.DB.GetStateDeltas failed")
			return
		}
	}

	for _, roomID := range joinedRooms {
		req.Rooms[roomID] = gomatrixserverlib.Join
	}

	for _, delta := range stateDeltas {
		if err = p.addRoomDeltaToResponse(ctx, req.Device, r, delta, &eventFilter, req.Response); err != nil {
			req.Log.WithError(err).Error("d.addRoomDeltaToResponse failed")
			return newPos
		}
	}

	return r.To
}

func (p *PDUStreamProvider) addRoomDeltaToResponse(
	ctx context.Context,
	device *userapi.Device,
	r types.Range,
	delta types.StateDelta,
	eventFilter *gomatrixserverlib.RoomEventFilter,
	res *types.Response,
) error {
	if delta.MembershipPos > 0 && delta.Membership == gomatrixserverlib.Leave {
		// make sure we don't leak recent events after the leave event.
		// TODO: History visibility makes this somewhat complex to handle correctly. For example:
		// TODO: This doesn't work for join -> leave in a single /sync request (see events prior to join).
		// TODO: This will fail on join -> leave -> sensitive msg -> join -> leave
		//       in a single /sync request
		// This is all "okay" assuming history_visibility == "shared" which it is by default.
		r.To = delta.MembershipPos
	}
	recentStreamEvents, limited, err := p.DB.RecentEvents(
		ctx, delta.RoomID, r,
		eventFilter, true, true,
	)
	if err != nil {
		return err
	}
	recentEvents := p.DB.StreamEventsToEvents(device, recentStreamEvents)
	delta.StateEvents = removeDuplicates(delta.StateEvents, recentEvents) // roll back
	prevBatch, err := p.DB.GetBackwardTopologyPos(ctx, recentStreamEvents)
	if err != nil {
		return err
	}

	// XXX: should we ever get this far if we have no recent events or state in this room?
	// in practice we do for peeks, but possibly not joins?
	if len(recentEvents) == 0 && len(delta.StateEvents) == 0 {
		return nil
	}

	switch delta.Membership {
	case gomatrixserverlib.Join:
		jr := types.NewJoinResponse()
		jr.Timeline.PrevBatch = &prevBatch
		jr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
		jr.Timeline.Limited = limited
		jr.State.Events = gomatrixserverlib.HeaderedToClientEvents(delta.StateEvents, gomatrixserverlib.FormatSync)
		res.Rooms.Join[delta.RoomID] = *jr

	case gomatrixserverlib.Peek:
		jr := types.NewJoinResponse()
		jr.Timeline.PrevBatch = &prevBatch
		jr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
		jr.Timeline.Limited = limited
		jr.State.Events = gomatrixserverlib.HeaderedToClientEvents(delta.StateEvents, gomatrixserverlib.FormatSync)
		res.Rooms.Peek[delta.RoomID] = *jr

	case gomatrixserverlib.Leave:
		fallthrough // transitions to leave are the same as ban

	case gomatrixserverlib.Ban:
		// TODO: recentEvents may contain events that this user is not allowed to see because they are
		//       no longer in the room.
		lr := types.NewLeaveResponse()
		lr.Timeline.PrevBatch = &prevBatch
		lr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
		lr.Timeline.Limited = false // TODO: if len(events) >= numRecents + 1 and then set limited:true
		lr.State.Events = gomatrixserverlib.HeaderedToClientEvents(delta.StateEvents, gomatrixserverlib.FormatSync)
		res.Rooms.Leave[delta.RoomID] = *lr
	}

	return nil
}

func (p *PDUStreamProvider) getJoinResponseForCompleteSync(
	ctx context.Context,
	roomID string,
	r types.Range,
	stateFilter *gomatrixserverlib.StateFilter,
	eventFilter *gomatrixserverlib.RoomEventFilter,
	wantFullState bool,
	device *userapi.Device,
) (jr *types.JoinResponse, err error) {
	// TODO: When filters are added, we may need to call this multiple times to get enough events.
	//       See: https://github.com/matrix-org/synapse/blob/v0.19.3/synapse/handlers/sync.py#L316
	recentStreamEvents, limited, err := p.DB.RecentEvents(
		ctx, roomID, r, eventFilter, true, true,
	)
	if err != nil {
		return
	}

	// TODO FIXME: We don't fully implement history visibility yet. To avoid leaking events which the
	// user shouldn't see, we check the recent events and remove any prior to the join event of the user
	// which is equiv to history_visibility: joined
	joinEventIndex := -1
	for i := len(recentStreamEvents) - 1; i >= 0; i-- {
		ev := recentStreamEvents[i]
		if ev.Type() == gomatrixserverlib.MRoomMember && ev.StateKeyEquals(device.UserID) {
			membership, _ := ev.Membership()
			if membership == "join" {
				joinEventIndex = i
				if i > 0 {
					// the create event happens before the first join, so we should cut it at that point instead
					if recentStreamEvents[i-1].Type() == gomatrixserverlib.MRoomCreate && recentStreamEvents[i-1].StateKeyEquals("") {
						joinEventIndex = i - 1
						break
					}
				}
				break
			}
		}
	}
	if joinEventIndex != -1 {
		// cut all events earlier than the join (but not the join itself)
		recentStreamEvents = recentStreamEvents[joinEventIndex:]
		limited = false // so clients know not to try to backpaginate
	}

	// Work our way through the timeline events and pick out the event IDs
	// of any state events that appear in the timeline. We'll specifically
	// exclude them at the next step, so that we don't get duplicate state
	// events in both `recentStreamEvents` and `stateEvents`.
	var excludingEventIDs []string
	if !wantFullState {
		excludingEventIDs = make([]string, 0, len(recentStreamEvents))
		for _, event := range recentStreamEvents {
			if event.StateKey() != nil {
				excludingEventIDs = append(excludingEventIDs, event.EventID())
			}
		}
	}

	stateEvents, err := p.DB.CurrentState(ctx, roomID, stateFilter, excludingEventIDs)
	if err != nil {
		return
	}

	// Retrieve the backward topology position, i.e. the position of the
	// oldest event in the room's topology.
	var prevBatch *types.TopologyToken
	if len(recentStreamEvents) > 0 {
		var backwardTopologyPos, backwardStreamPos types.StreamPosition
		backwardTopologyPos, backwardStreamPos, err = p.DB.PositionInTopology(ctx, recentStreamEvents[0].EventID())
		if err != nil {
			return
		}
		prevBatch = &types.TopologyToken{
			Depth:       backwardTopologyPos,
			PDUPosition: backwardStreamPos,
		}
		prevBatch.Decrement()
	}

	// We don't include a device here as we don't need to send down
	// transaction IDs for complete syncs, but we do it anyway because Sytest demands it for:
	// "Can sync a room with a message with a transaction id" - which does a complete sync to check.
	recentEvents := p.DB.StreamEventsToEvents(device, recentStreamEvents)
	stateEvents = removeDuplicates(stateEvents, recentEvents)
	jr = types.NewJoinResponse()
	jr.Timeline.PrevBatch = prevBatch
	jr.Timeline.Events = gomatrixserverlib.HeaderedToClientEvents(recentEvents, gomatrixserverlib.FormatSync)
	jr.Timeline.Limited = limited
	jr.State.Events = gomatrixserverlib.HeaderedToClientEvents(stateEvents, gomatrixserverlib.FormatSync)
	return jr, nil
}

func removeDuplicates(stateEvents, recentEvents []*gomatrixserverlib.HeaderedEvent) []*gomatrixserverlib.HeaderedEvent {
	for _, recentEv := range recentEvents {
		if recentEv.StateKey() == nil {
			continue // not a state event
		}
		// TODO: This is a linear scan over all the current state events in this room. This will
		//       be slow for big rooms. We should instead sort the state events by event ID  (ORDER BY)
		//       then do a binary search to find matching events, similar to what roomserver does.
		for j := 0; j < len(stateEvents); j++ {
			if stateEvents[j].EventID() == recentEv.EventID() {
				// overwrite the element to remove with the last element then pop the last element.
				// This is orders of magnitude faster than re-slicing, but doesn't preserve ordering
				// (we don't care about the order of stateEvents)
				stateEvents[j] = stateEvents[len(stateEvents)-1]
				stateEvents = stateEvents[:len(stateEvents)-1]
				break // there shouldn't be multiple events with the same event ID
			}
		}
	}
	return stateEvents
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"

	"github.com/dougfort/gocards/standard"

	"github.com/dougfort/arachne/game"
)

func displayTableauStrings(g gameData) {
	var row int
	image := make([]string, game.TableauWidth)

ROW_LOOP:
	for {
		var found bool
		for col := 0; col < game.TableauWidth; col++ {
			if row < g.remote.Tableau[col].HiddenCount {
				image[col] = standard.StringBack
				found = true
			} else {
				visibleRow := row - g.remote.Tableau[col].HiddenCount
				if visibleRow < len(g.remote.Tableau[col].Cards) {
					image[col] = standard.Strings[g.remote.Tableau[col].Cards[visibleRow]]
					found = true
				} else {
					image[col] = ""
				}
			}
		}
		if !found {
			break ROW_LOOP
		}
		fmt.Printf("%-17s %-17s %-17s %-17s %-17s %-17s %-17s %-17s %-17s %-17s\n",
			image[0], image[1], image[2], image[3], image[4],
			image[5], image[6], image[7], image[8], image[9],
		)

		row++
	}
}
<commit_msg>tighten te display size to fit inside an 80 col screen<commit_after>package main

import (
	"fmt"

	"github.com/dougfort/gocards/standard"

	"github.com/dougfort/arachne/game"
)

func displayTableauStrings(g gameData) {
	var row int
	image := make([]string, game.TableauWidth)

ROW_LOOP:
	for {
		var found bool
		for col := 0; col < game.TableauWidth; col++ {
			if row < g.remote.Tableau[col].HiddenCount {
				image[col] = standard.StringBack
				found = true
			} else {
				visibleRow := row - g.remote.Tableau[col].HiddenCount
				if visibleRow < len(g.remote.Tableau[col].Cards) {
					image[col] = standard.Strings[g.remote.Tableau[col].Cards[visibleRow]]
					found = true
				} else {
					image[col] = ""
				}
			}
		}
		if !found {
			break ROW_LOOP
		}
		fmt.Printf("%-7s %-7s %-7s %-7s %-7s %-7s %-7s %-7s %-7s %-7s\n",
			image[0], image[1], image[2], image[3], image[4],
			image[5], image[6], image[7], image[8], image[9],
		)

		row++
	}
}
<|endoftext|>


<commit_before>package weed_server

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (fs *FilerServer) GetOrHeadHandler(w http.ResponseWriter, r *http.Request, isGetMethod bool) {

	path := r.URL.Path
	isForDirectory := strings.HasSuffix(path, "/")
	if isForDirectory && len(path) > 1 {
		path = path[:len(path)-1]
	}

	entry, err := fs.filer.FindEntry(context.Background(), filer2.FullPath(path))
	if err != nil {
		if path == "/" {
			fs.listDirectoryHandler(w, r)
			return
		}
		if err == filer_pb.ErrNotFound {
			glog.V(1).Infof("Not found %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.notfound").Inc()
			w.WriteHeader(http.StatusNotFound)
		} else {
			glog.V(0).Infof("Internal %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.internalerror").Inc()
			w.WriteHeader(http.StatusInternalServerError)
		}
		return
	}

	if entry.IsDirectory() {
		if fs.option.DisableDirListing {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		fs.listDirectoryHandler(w, r)
		return
	}

	if isForDirectory {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	if len(entry.Chunks) == 0 {
		glog.V(1).Infof("no file chunks for %s, attr=%+v", path, entry.Attr)
		stats.FilerRequestCounter.WithLabelValues("read.nocontent").Inc()
		w.WriteHeader(http.StatusNoContent)
		return
	}

	w.Header().Set("Accept-Ranges", "bytes")
	if r.Method == "HEAD" {
		w.Header().Set("Content-Length", strconv.FormatInt(int64(filer2.TotalSize(entry.Chunks)), 10))
		w.Header().Set("Last-Modified", entry.Attr.Mtime.Format(http.TimeFormat))
		if entry.Attr.Mime != "" {
			w.Header().Set("Content-Type", entry.Attr.Mime)
		}
		setEtag(w, filer2.ETag(entry.Chunks))
		return
	}

	if len(entry.Chunks) == 1 {
		fs.handleSingleChunk(w, r, entry)
		return
	}

	fs.handleMultipleChunks(w, r, entry)

}

func (fs *FilerServer) handleSingleChunk(w http.ResponseWriter, r *http.Request, entry *filer2.Entry) {

	fileId := entry.Chunks[0].GetFileIdString()

	urlString, err := fs.filer.MasterClient.LookupFileId(fileId)
	if err != nil {
		glog.V(1).Infof("operation LookupFileId %s failed, err: %v", fileId, err)
		w.WriteHeader(http.StatusNotFound)
		return
	}

	if fs.option.RedirectOnRead && entry.Chunks[0].CipherKey == nil {
		stats.FilerRequestCounter.WithLabelValues("redirect").Inc()
		http.Redirect(w, r, urlString, http.StatusFound)
		return
	}

	u, _ := url.Parse(urlString)
	q := u.Query()
	for key, values := range r.URL.Query() {
		for _, value := range values {
			q.Add(key, value)
		}
	}
	u.RawQuery = q.Encode()
	request := &http.Request{
		Method:        r.Method,
		URL:           u,
		Proto:         r.Proto,
		ProtoMajor:    r.ProtoMajor,
		ProtoMinor:    r.ProtoMinor,
		Header:        r.Header,
		Body:          r.Body,
		Host:          r.Host,
		ContentLength: r.ContentLength,
	}
	glog.V(3).Infoln("retrieving from", u)
	resp, do_err := util.Do(request)
	if do_err != nil {
		glog.V(0).Infoln("failing to connect to volume server", do_err.Error())
		writeJsonError(w, r, http.StatusInternalServerError, do_err)
		return
	}
	defer func() {
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}()
	for k, v := range resp.Header {
		w.Header()[k] = v
	}
	if entry.Attr.Mime != "" {
		w.Header().Set("Content-Type", entry.Attr.Mime)
	}
	if entry.Chunks[0].CipherKey == nil {
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	} else {
		fs.writeEncryptedChunk(w, resp, entry)
	}
}

func (fs *FilerServer) writeEncryptedChunk(w http.ResponseWriter, resp *http.Response, entry *filer2.Entry) {
	chunk := entry.Chunks[0]
	encryptedData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		glog.V(1).Infof("read encrypted %s failed, err: %v", chunk.FileId, err)
		w.WriteHeader(http.StatusNotFound)
		return
	}
	decryptedData, err := util.Decrypt(encryptedData, util.CipherKey(chunk.CipherKey))
	if err != nil {
		glog.V(1).Infof("decrypt %s failed, err: %v", chunk.FileId, err)
		w.WriteHeader(http.StatusNotFound)
		return
	}
	w.Header().Set("Content-Length", fmt.Sprintf("%d", chunk.Size))
	w.WriteHeader(resp.StatusCode)
	w.Write(decryptedData)
}

func (fs *FilerServer) handleMultipleChunks(w http.ResponseWriter, r *http.Request, entry *filer2.Entry) {

	mimeType := entry.Attr.Mime
	if mimeType == "" {
		if ext := path.Ext(entry.Name()); ext != "" {
			mimeType = mime.TypeByExtension(ext)
		}
	}
	if mimeType != "" {
		w.Header().Set("Content-Type", mimeType)
	}
	setEtag(w, filer2.ETag(entry.Chunks))

	totalSize := int64(filer2.TotalSize(entry.Chunks))

	processRangeRequst(r, w, totalSize, mimeType, func(writer io.Writer, offset int64, size int64) error {
		return fs.writeContent(writer, entry, offset, int(size))
	})

}

func processRangeRequst(r *http.Request, w http.ResponseWriter, totalSize int64, mimeType string, writeFn func(writer io.Writer, offset int64, size int64) error) {
	rangeReq := r.Header.Get("Range")

	if rangeReq == "" {
		w.Header().Set("Content-Length", strconv.FormatInt(totalSize, 10))
		if err := writeFn(w, 0, totalSize); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		return
	}

	//the rest is dealing with partial content request
	//mostly copy from src/pkg/net/http/fs.go
	ranges, err := parseRange(rangeReq, totalSize)
	if err != nil {
		http.Error(w, err.Error(), http.StatusRequestedRangeNotSatisfiable)
		return
	}
	if sumRangesSize(ranges) > totalSize {
		// The total number of bytes in all the ranges
		// is larger than the size of the file by
		// itself, so this is probably an attack, or a
		// dumb client.  Ignore the range request.
		return
	}
	if len(ranges) == 0 {
		return
	}
	if len(ranges) == 1 {
		// RFC 2616, Section 14.16:
		// "When an HTTP message includes the content of a single
		// range (for example, a response to a request for a
		// single range, or to a request for a set of ranges
		// that overlap without any holes), this content is
		// transmitted with a Content-Range header, and a
		// Content-Length header showing the number of bytes
		// actually transferred.
		// ...
		// A response to a request for a single range MUST NOT
		// be sent using the multipart/byteranges media type."
		ra := ranges[0]
		w.Header().Set("Content-Length", strconv.FormatInt(ra.length, 10))
		w.Header().Set("Content-Range", ra.contentRange(totalSize))
		w.WriteHeader(http.StatusPartialContent)

		err = writeFn(w, ra.start, ra.length)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		return
	}

	// process multiple ranges
	for _, ra := range ranges {
		if ra.start > totalSize {
			http.Error(w, "Out of Range", http.StatusRequestedRangeNotSatisfiable)
			return
		}
	}
	sendSize := rangesMIMESize(ranges, mimeType, totalSize)
	pr, pw := io.Pipe()
	mw := multipart.NewWriter(pw)
	w.Header().Set("Content-Type", "multipart/byteranges; boundary="+mw.Boundary())
	sendContent := pr
	defer pr.Close() // cause writing goroutine to fail and exit if CopyN doesn't finish.
	go func() {
		for _, ra := range ranges {
			part, e := mw.CreatePart(ra.mimeHeader(mimeType, totalSize))
			if e != nil {
				pw.CloseWithError(e)
				return
			}
			if e = writeFn(part, ra.start, ra.length); e != nil {
				pw.CloseWithError(e)
				return
			}
		}
		mw.Close()
		pw.Close()
	}()
	if w.Header().Get("Content-Encoding") == "" {
		w.Header().Set("Content-Length", strconv.FormatInt(sendSize, 10))
	}
	w.WriteHeader(http.StatusPartialContent)
	if _, err := io.CopyN(w, sendContent, sendSize); err != nil {
		http.Error(w, "Internal Error", http.StatusInternalServerError)
		return
	}
}

func (fs *FilerServer) writeContent(w io.Writer, entry *filer2.Entry, offset int64, size int) error {

	return filer2.StreamContent(fs.filer.MasterClient, w, entry.Chunks, offset, size)

}
<commit_msg>refactoring<commit_after>package weed_server

import (
	"context"
	"fmt"
	"io"
	"io/ioutil"
	"mime"
	"mime/multipart"
	"net/http"
	"net/url"
	"path"
	"strconv"
	"strings"

	"github.com/chrislusf/seaweedfs/weed/filer2"
	"github.com/chrislusf/seaweedfs/weed/glog"
	"github.com/chrislusf/seaweedfs/weed/pb/filer_pb"
	"github.com/chrislusf/seaweedfs/weed/stats"
	"github.com/chrislusf/seaweedfs/weed/util"
)

func (fs *FilerServer) GetOrHeadHandler(w http.ResponseWriter, r *http.Request, isGetMethod bool) {

	path := r.URL.Path
	isForDirectory := strings.HasSuffix(path, "/")
	if isForDirectory && len(path) > 1 {
		path = path[:len(path)-1]
	}

	entry, err := fs.filer.FindEntry(context.Background(), filer2.FullPath(path))
	if err != nil {
		if path == "/" {
			fs.listDirectoryHandler(w, r)
			return
		}
		if err == filer_pb.ErrNotFound {
			glog.V(1).Infof("Not found %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.notfound").Inc()
			w.WriteHeader(http.StatusNotFound)
		} else {
			glog.V(0).Infof("Internal %s: %v", path, err)
			stats.FilerRequestCounter.WithLabelValues("read.internalerror").Inc()
			w.WriteHeader(http.StatusInternalServerError)
		}
		return
	}

	if entry.IsDirectory() {
		if fs.option.DisableDirListing {
			w.WriteHeader(http.StatusMethodNotAllowed)
			return
		}
		fs.listDirectoryHandler(w, r)
		return
	}

	if isForDirectory {
		w.WriteHeader(http.StatusNotFound)
		return
	}

	if len(entry.Chunks) == 0 {
		glog.V(1).Infof("no file chunks for %s, attr=%+v", path, entry.Attr)
		stats.FilerRequestCounter.WithLabelValues("read.nocontent").Inc()
		w.WriteHeader(http.StatusNoContent)
		return
	}

	w.Header().Set("Accept-Ranges", "bytes")
	if r.Method == "HEAD" {
		w.Header().Set("Content-Length", strconv.FormatInt(int64(filer2.TotalSize(entry.Chunks)), 10))
		w.Header().Set("Last-Modified", entry.Attr.Mtime.Format(http.TimeFormat))
		if entry.Attr.Mime != "" {
			w.Header().Set("Content-Type", entry.Attr.Mime)
		}
		setEtag(w, filer2.ETag(entry.Chunks))
		return
	}

	if len(entry.Chunks) == 1 {
		fs.handleSingleChunk(w, r, entry)
		return
	}

	fs.handleMultipleChunks(w, r, entry)

}

func (fs *FilerServer) handleSingleChunk(w http.ResponseWriter, r *http.Request, entry *filer2.Entry) {

	fileId := entry.Chunks[0].GetFileIdString()

	urlString, err := fs.filer.MasterClient.LookupFileId(fileId)
	if err != nil {
		glog.V(1).Infof("operation LookupFileId %s failed, err: %v", fileId, err)
		w.WriteHeader(http.StatusNotFound)
		return
	}

	if fs.option.RedirectOnRead && entry.Chunks[0].CipherKey == nil {
		stats.FilerRequestCounter.WithLabelValues("redirect").Inc()
		http.Redirect(w, r, urlString, http.StatusFound)
		return
	}

	u, _ := url.Parse(urlString)
	q := u.Query()
	for key, values := range r.URL.Query() {
		for _, value := range values {
			q.Add(key, value)
		}
	}
	u.RawQuery = q.Encode()
	request := &http.Request{
		Method:        r.Method,
		URL:           u,
		Proto:         r.Proto,
		ProtoMajor:    r.ProtoMajor,
		ProtoMinor:    r.ProtoMinor,
		Header:        r.Header,
		Body:          r.Body,
		Host:          r.Host,
		ContentLength: r.ContentLength,
	}
	glog.V(3).Infoln("retrieving from", u)
	resp, do_err := util.Do(request)
	if do_err != nil {
		glog.V(0).Infoln("failing to connect to volume server", do_err.Error())
		writeJsonError(w, r, http.StatusInternalServerError, do_err)
		return
	}
	defer func() {
		io.Copy(ioutil.Discard, resp.Body)
		resp.Body.Close()
	}()
	for k, v := range resp.Header {
		w.Header()[k] = v
	}
	if entry.Attr.Mime != "" {
		w.Header().Set("Content-Type", entry.Attr.Mime)
	}
	if entry.Chunks[0].CipherKey == nil {
		w.WriteHeader(resp.StatusCode)
		io.Copy(w, resp.Body)
	} else {
		fs.writeEncryptedChunk(w, resp, entry)
	}
}

func (fs *FilerServer) writeEncryptedChunk(w http.ResponseWriter, resp *http.Response, entry *filer2.Entry) {
	chunk := entry.Chunks[0]
	encryptedData, err := ioutil.ReadAll(resp.Body)
	if err != nil {
		glog.V(1).Infof("read encrypted %s failed, err: %v", chunk.FileId, err)
		w.WriteHeader(http.StatusNotFound)
		return
	}
	decryptedData, err := util.Decrypt(encryptedData, util.CipherKey(chunk.CipherKey))
	if err != nil {
		glog.V(1).Infof("decrypt %s failed, err: %v", chunk.FileId, err)
		w.WriteHeader(http.StatusNotFound)
		return
	}
	w.Header().Set("Content-Length", fmt.Sprintf("%d", chunk.Size))
	w.WriteHeader(resp.StatusCode)
	w.Write(decryptedData)
}

func (fs *FilerServer) handleMultipleChunks(w http.ResponseWriter, r *http.Request, entry *filer2.Entry) {

	mimeType := entry.Attr.Mime
	if mimeType == "" {
		if ext := path.Ext(entry.Name()); ext != "" {
			mimeType = mime.TypeByExtension(ext)
		}
	}
	if mimeType != "" {
		w.Header().Set("Content-Type", mimeType)
	}
	setEtag(w, filer2.ETag(entry.Chunks))

	totalSize := int64(filer2.TotalSize(entry.Chunks))

	processRangeRequst(r, w, totalSize, mimeType, func(writer io.Writer, offset int64, size int64) error {
		return filer2.StreamContent(fs.filer.MasterClient, w, entry.Chunks, offset, int(size))
	})

}

func processRangeRequst(r *http.Request, w http.ResponseWriter, totalSize int64, mimeType string, writeFn func(writer io.Writer, offset int64, size int64) error) {
	rangeReq := r.Header.Get("Range")

	if rangeReq == "" {
		w.Header().Set("Content-Length", strconv.FormatInt(totalSize, 10))
		if err := writeFn(w, 0, totalSize); err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		return
	}

	//the rest is dealing with partial content request
	//mostly copy from src/pkg/net/http/fs.go
	ranges, err := parseRange(rangeReq, totalSize)
	if err != nil {
		http.Error(w, err.Error(), http.StatusRequestedRangeNotSatisfiable)
		return
	}
	if sumRangesSize(ranges) > totalSize {
		// The total number of bytes in all the ranges
		// is larger than the size of the file by
		// itself, so this is probably an attack, or a
		// dumb client.  Ignore the range request.
		return
	}
	if len(ranges) == 0 {
		return
	}
	if len(ranges) == 1 {
		// RFC 2616, Section 14.16:
		// "When an HTTP message includes the content of a single
		// range (for example, a response to a request for a
		// single range, or to a request for a set of ranges
		// that overlap without any holes), this content is
		// transmitted with a Content-Range header, and a
		// Content-Length header showing the number of bytes
		// actually transferred.
		// ...
		// A response to a request for a single range MUST NOT
		// be sent using the multipart/byteranges media type."
		ra := ranges[0]
		w.Header().Set("Content-Length", strconv.FormatInt(ra.length, 10))
		w.Header().Set("Content-Range", ra.contentRange(totalSize))
		w.WriteHeader(http.StatusPartialContent)

		err = writeFn(w, ra.start, ra.length)
		if err != nil {
			http.Error(w, err.Error(), http.StatusInternalServerError)
			return
		}
		return
	}

	// process multiple ranges
	for _, ra := range ranges {
		if ra.start > totalSize {
			http.Error(w, "Out of Range", http.StatusRequestedRangeNotSatisfiable)
			return
		}
	}
	sendSize := rangesMIMESize(ranges, mimeType, totalSize)
	pr, pw := io.Pipe()
	mw := multipart.NewWriter(pw)
	w.Header().Set("Content-Type", "multipart/byteranges; boundary="+mw.Boundary())
	sendContent := pr
	defer pr.Close() // cause writing goroutine to fail and exit if CopyN doesn't finish.
	go func() {
		for _, ra := range ranges {
			part, e := mw.CreatePart(ra.mimeHeader(mimeType, totalSize))
			if e != nil {
				pw.CloseWithError(e)
				return
			}
			if e = writeFn(part, ra.start, ra.length); e != nil {
				pw.CloseWithError(e)
				return
			}
		}
		mw.Close()
		pw.Close()
	}()
	if w.Header().Get("Content-Encoding") == "" {
		w.Header().Set("Content-Length", strconv.FormatInt(sendSize, 10))
	}
	w.WriteHeader(http.StatusPartialContent)
	if _, err := io.CopyN(w, sendContent, sendSize); err != nil {
		http.Error(w, "Internal Error", http.StatusInternalServerError)
		return
	}
}
<|endoftext|>


<commit_before>package redisence

import (
	"errors"
	"fmt"
	"testing"
	"time"
)

func initRedisence(t *testing.T) *Session {
	ses, err := New("localhost:6379", 10, time.Second*1)
	if err != nil {
		t.Fatal(err)
	}
	conn := ses.redis.Pool().Get()
	conn.Do("CONFIG", "SET", "notify-keyspace-events Ex$")
	if err := conn.Close(); err != nil {
		t.Fatal(err)
	}

	return ses
}

func TestInitialization(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()
}

func TestSinglePing(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id"); err != nil {
		t.Fatal(err)
	}
}

func TestMultiPing(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id", "id2"); err != nil {
		t.Fatal(err)
	}
}

func TestOnlineStatus(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	id := "id3"
	if err := s.Online(id); err != nil {
		t.Fatal(err)
	}

	status, err := s.Status(id)
	if err != nil {
		t.Fatal(err)
	}

	if status.Status != Online {
		t.Fatal(errors.New("User should be active"))
	}
}

func TestOfflineStatus(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	id := "id4"
	if err := s.Online(id); err != nil {
		t.Fatal(err)
	}

	status, err := s.Status("id5")
	if err != nil {
		t.Fatal(err)
	}

	if status.Status != Offline {
		t.Fatal(errors.New("User should be offline"))
	}
}

func TestMultiStatusAllOnline(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id6", "id7"); err != nil {
		t.Fatal(err)
	}

	status, err := s.MultipleStatus([]string{"id6", "id7"})
	if err != nil {
		t.Fatal(err)
	}
	for _, st := range status {
		if st.Status != Online {
			t.Fatal(errors.New("User should be active"))
		}
	}
}

func TestMultiStatusAllOffline(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id8", "id9"); err != nil {
		t.Fatal(err)
	}

	status, err := s.MultipleStatus([]string{"id10", "id11"})
	if err != nil {
		t.Fatal(err)
	}

	for _, st := range status {
		if st.Status != Offline {
			t.Fatal(errors.New("User should be offline"))
		}
	}
}

func TestStatusWithTimeout(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	id := "12"
	if err := s.Online(id); err != nil {
		t.Fatal(err)
	}
	time.Sleep(time.Second * 2)
	status, err := s.Status(id)
	if err != nil {
		t.Fatal(err)
	}
	if status.Status == Online {
		t.Fatal(errors.New("User should not be active"))
	}
}

func TestSubscriptions(t *testing.T) {
	t.Skip("Skipped to travis")
	s := initRedisence(t)
	defer s.Close()

	// wait for all keys to expire
	time.Sleep(time.Second * 1)

	events := make(chan Event)

	id1 := "13"
	id2 := "14"
	id3 := "15"

	go s.ListenStatusChanges(events)

	time.AfterFunc(time.Second*5, func() {
		err := s.Close()
		if err != nil {
			t.Fatal(err)
		}
	})

	time.AfterFunc(time.Second*1, func() {
		err := s.Online(id1, id2, id3)
		if err != nil {
			t.Fatal(err)
		}
		// err = s.Offline(id1, id2, id3)
		// if err != nil {
		// 	t.Fatal(err)
		// }
	})

	onlineCount := 0
	offlineCount := 0
	closedCount := 0
	for event := range events {
		switch event.Status {
		case Online:
			onlineCount++
		case Offline:
			offlineCount++
		case Closed:
			closedCount++
			close(events)
			// return
		default:
		}
	}

	if onlineCount != 3 {
		t.Fatal(
			errors.New(
				fmt.Sprintf("online count should be 3 it is %d", onlineCount),
			),
		)
	}

	if offlineCount != 3 {
		t.Fatal(
			errors.New(
				fmt.Sprintf("offline count should be 3 it is %d", offlineCount),
			),
		)
	}

	if closedCount != 1 {
		t.Fatal(
			errors.New(
				fmt.Sprintf("closedCount count should be 3 it is %d", closedCount),
			),
		)
	}

}

func TestJustMultiOffline(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Offline("id16", "id17"); err != nil {
		t.Fatal(err)
	}
}

func TestMultiOnlineAndOfflineTogether(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id18", "id19"); err != nil {
		t.Fatal(err)
	}
	if err := s.Offline("id18", "id19"); err != nil {
		t.Fatal(err)
	}
}

func TestMultiOfflineWithMultiStatus(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id20", "id21"); err != nil {
		t.Fatal(err)
	}
	if err := s.Offline("id20", "id21"); err != nil {
		t.Fatal(err)
	}
	status, err := s.MultipleStatus([]string{"id20", "id21"})
	if err != nil {
		t.Fatal(err)
	}

	for _, st := range status {
		if st.Status != Offline {
			t.Fatal(errors.New("User should be offline"))
		}
	}
}
<commit_msg>Redisence:  update tests according to new listenStatusChange system<commit_after>package redisence

import (
	"errors"
	"fmt"
	"testing"
	"time"
)

func initRedisence(t *testing.T) *Session {
	ses, err := New("localhost:6379", 10, time.Second*1)
	if err != nil {
		t.Fatal(err)
	}
	conn := ses.redis.Pool().Get()
	conn.Do("CONFIG", "SET", "notify-keyspace-events Ex$")
	if err := conn.Close(); err != nil {
		t.Fatal(err)
	}

	return ses
}

func TestInitialization(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()
}

func TestSinglePing(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id"); err != nil {
		t.Fatal(err)
	}
}

func TestMultiPing(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id", "id2"); err != nil {
		t.Fatal(err)
	}
}

func TestOnlineStatus(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	id := "id3"
	if err := s.Online(id); err != nil {
		t.Fatal(err)
	}

	status, err := s.Status(id)
	if err != nil {
		t.Fatal(err)
	}

	if status.Status != Online {
		t.Fatal(errors.New("User should be active"))
	}
}

func TestOfflineStatus(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	id := "id4"
	if err := s.Online(id); err != nil {
		t.Fatal(err)
	}

	status, err := s.Status("id5")
	if err != nil {
		t.Fatal(err)
	}

	if status.Status != Offline {
		t.Fatal(errors.New("User should be offline"))
	}
}

func TestMultiStatusAllOnline(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id6", "id7"); err != nil {
		t.Fatal(err)
	}

	status, err := s.MultipleStatus([]string{"id6", "id7"})
	if err != nil {
		t.Fatal(err)
	}
	for _, st := range status {
		if st.Status != Online {
			t.Fatal(errors.New("User should be active"))
		}
	}
}

func TestMultiStatusAllOffline(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id8", "id9"); err != nil {
		t.Fatal(err)
	}

	status, err := s.MultipleStatus([]string{"id10", "id11"})
	if err != nil {
		t.Fatal(err)
	}

	for _, st := range status {
		if st.Status != Offline {
			t.Fatal(errors.New("User should be offline"))
		}
	}
}

func TestStatusWithTimeout(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	id := "12"
	if err := s.Online(id); err != nil {
		t.Fatal(err)
	}
	time.Sleep(time.Second * 2)
	status, err := s.Status(id)
	if err != nil {
		t.Fatal(err)
	}
	if status.Status == Online {
		t.Fatal(errors.New("User should not be active"))
	}
}

func TestSubscriptions(t *testing.T) {
	t.Skip("Skipped to travis")
	s := initRedisence(t)

	// wait for all keys to expire
	time.Sleep(time.Second * 1)

	id1 := "13"
	id2 := "14"
	id3 := "15"

	time.AfterFunc(time.Second*5, func() {
		err := s.Close()
		if err != nil {
			t.Fatal(err)
		}
	})

	time.AfterFunc(time.Second*1, func() {
		err := s.Online(id1, id2, id3)
		if err != nil {
			t.Fatal(err)
		}
		// err = s.Offline(id1, id2, id3)
		// if err != nil {
		// 	t.Fatal(err)
		// }
	})

	onlineCount := 0
	offlineCount := 0
	for event := range s.ListenStatusChanges() {
		switch event.Status {
		case Online:
			onlineCount++
		case Offline:
			offlineCount++
		}
	}

	if onlineCount != 3 {
		t.Fatal(
			errors.New(
				fmt.Sprintf("online count should be 3 it is %d", onlineCount),
			),
		)
	}

	if offlineCount != 3 {
		t.Fatal(
			errors.New(
				fmt.Sprintf("offline count should be 3 it is %d", offlineCount),
			),
		)
	}
}

func TestJustMultiOffline(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Offline("id16", "id17"); err != nil {
		t.Fatal(err)
	}
}

func TestMultiOnlineAndOfflineTogether(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id18", "id19"); err != nil {
		t.Fatal(err)
	}
	if err := s.Offline("id18", "id19"); err != nil {
		t.Fatal(err)
	}
}

func TestMultiOfflineWithMultiStatus(t *testing.T) {
	s := initRedisence(t)
	defer s.Close()

	if err := s.Online("id20", "id21"); err != nil {
		t.Fatal(err)
	}
	if err := s.Offline("id20", "id21"); err != nil {
		t.Fatal(err)
	}
	status, err := s.MultipleStatus([]string{"id20", "id21"})
	if err != nil {
		t.Fatal(err)
	}

	for _, st := range status {
		if st.Status != Offline {
			t.Fatal(errors.New("User should be offline"))
		}
	}
}
<|endoftext|>