<commit_before>package consumer

import (
	"encoding/json"
	"flume-log-sdk/config"
	"flume-log-sdk/consumer/client"
	"fmt"
	"github.com/blackbeans/redigo/redis"
	"log"
	"math/rand"
	_ "os"
	"strconv"
	"sync/atomic"
	"time"
)

type counter struct {
	lastSuccValue int64

	currSuccValue int64

	lastFailValue int64

	currFailValue int64
}

// 用于向flume中作为sink 通过thrift客户端写入日志

type SinkServer struct {
	redisPool       map[string][]*redis.Pool
	flumeClientPool []*flumeClientPool
	isStop          bool
	monitorCount    counter
}

func NewSinkServer(option *config.Option) (server *SinkServer) {

	redisPool := make(map[string][]*redis.Pool, 0)

	//创建redis的消费连接
	for _, v := range option.QueueHostPorts {

		pool := redis.NewPool(func() (conn redis.Conn, err error) {

			conn, err = redis.DialTimeout("tcp", v.Host+":"+strconv.Itoa(v.Port),
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second)

			return
		}, time.Duration(v.Timeout*2)*time.Second, v.Maxconn/2, v.Maxconn)

		pools, ok := redisPool[v.QueueName]
		if !ok {
			pools = make([]*redis.Pool, 0)
			redisPool[v.QueueName] = pools
		}

		redisPool[v.QueueName] = append(pools, pool)

	}

	pools := make([]*flumeClientPool, 0)
	//创建flume的client
	for _, v := range option.FlumeAgents {

		pool := newFlumeClientPool(20, 50, 100, 10*time.Second, func() *client.FlumeClient {
			flumeclient := client.NewFlumeClient(v.Host, v.Port)
			flumeclient.Connect()
			return flumeclient
		})
		pools = append(pools, pool)

		go monitorPool(v.Host+":"+strconv.Itoa(v.Port), pool)
	}

	sinkserver := &SinkServer{redisPool: redisPool, flumeClientPool: pools}
	go sinkserver.monitorFlume()
	return sinkserver
}

func (self *SinkServer) monitorFlume() {
	for {
		time.Sleep(1 * time.Second)
		currSucc := self.monitorCount.currSuccValue
		currFail := self.monitorCount.currFailValue
		log.Printf("succ-send:%d,fail-send:%d",
			(currSucc - self.monitorCount.lastSuccValue),
			(currFail - self.monitorCount.lastFailValue))
		self.monitorCount.currSuccValue = currSucc
		self.monitorCount.currFailValue = currFail
	}
}

func monitorPool(hostport string, pool *flumeClientPool) {
	for {
		time.Sleep(1 * time.Second)

		log.Printf("flume:%s|active:%d,core:%d,max:%d",
			hostport, pool.ActivePoolSize(), pool.CorePoolSize(), pool.maxPoolSize)
	}

}

//启动pop
func (self *SinkServer) Start() {

	self.isStop = false
	ch := make(chan int, 1)
	var count = 0
	for k, v := range self.redisPool {

		log.Println("start redis queueserver succ " + k)
		for _, pool := range v {
			count++
			defer pool.Close()
			go func(queuename string, pool *redis.Pool, end chan int) {
				conn := pool.Get()
				defer pool.Release(conn)
				for !self.isStop {

					// log.Println("pool active count :", strconv.Itoa(pool.ActiveCount()))
					reply, err := conn.Do("LPOP", queuename)
					if nil != err || nil == reply {
						if nil != err {
							log.Printf("LPOP|FAIL|%s", err)
							conn.Close()
							conn = pool.Get()
						} else {
							time.Sleep(100 * time.Millisecond)
						}

						continue
					}

					resp := reply.([]byte)
					var cmd config.Command
					err = json.Unmarshal(resp, &cmd)

					if nil != err {
						log.Printf("command unmarshal fail ! %s | error:%s\n", resp, err.Error())
						continue
					} else if rand.Int()%10 == 0 {
						log.Println("trace|command|%s", cmd)
					}

					//
					momoid := cmd.Params["momoid"].(string)

					businessName := cmd.Params["businessName"].(string)

					action := cmd.Params["type"].(string)

					bodyContent := cmd.Params["body"]

					//将businessName 加入到body中
					bodyMap := bodyContent.(map[string]interface{})
					bodyMap["business_type"] = businessName

					body, err := json.Marshal(bodyContent)

					if nil != err {
						log.Printf("marshal log body fail %s", err.Error())
						continue
					}

					//这里需要优化一下body,需要采用其他的方式定义Body格式，写入

					// log.Printf("%s,%s,%s,%s", momoid, businessName, action, string(body))

					//启动处理任务
					go self.innerSend(momoid, businessName, action, string(body))

				}
				end <- -1
			}(k, pool, ch)
		}
	}

	for {
		count += <-ch
		if count <= 0 {
			log.Printf("redis conn  close %d", count)
			break
		}
	}

}

func (self *SinkServer) innerSend(momoid, businessName, action string, body string) {

	for i := 0; i < 3; i++ {
		pool := self.getFlumeClientPool(businessName, action)
		flumeclient, err := pool.Get(5 * time.Second)
		if nil != err {
			continue
		}
		//拼装头部信息
		header := make(map[string]string, 1)
		header["businessName"] = businessName
		header["type"] = action

		//拼Body
		flumeBody := fmt.Sprintf("%s\t%s\t%s", momoid, action, body)
		err = flumeclient.Append(header, []byte(flumeBody))
		defer func() {
			if err := recover(); nil != err {
				//回收这个坏的连接
				pool.ReleaseBroken(flumeclient)
			}
		}()

		if nil != err {
			atomic.AddInt64(&self.monitorCount.currFailValue, 1)
			log.Printf("send 2 flume fail %s \t err:%s\n", body, err.Error())
		} else {
			atomic.AddInt64(&self.monitorCount.currSuccValue, 1)
			// log.Printf("send 2 flume succ %s\n", body)
			pool.Release(flumeclient)
			break
		}
	}
}

//仅供测试使用推送数据
func (self *SinkServer) testPushLog(queuename, logger string) {

	for _, v := range self.redisPool {
		for _, pool := range v {
			conn := pool.Get()
			defer pool.Release(conn)

			reply, err := conn.Do("RPUSH", queuename, logger)
			log.Printf("%s|err:%s", reply, err)
			break

		}
	}

}

func (self *SinkServer) Stop() {
	self.isStop = true
	for _, v := range self.flumeClientPool {
		v.Destroy()
	}

	for _, v := range self.redisPool {
		for _, p := range v {
			p.Close()
		}
	}
}

func (self *SinkServer) getFlumeClientPool(businessName, action string) *flumeClientPool {

	//使用随机算法直接获得

	idx := rand.Intn(len(self.flumeClientPool))
	return self.flumeClientPool[idx]

}
<commit_msg>	modified:   ../consumer/log_sink.go<commit_after>package consumer

import (
	"encoding/json"
	"flume-log-sdk/config"
	"flume-log-sdk/consumer/client"
	"fmt"
	"github.com/blackbeans/redigo/redis"
	"log"
	"math/rand"
	_ "os"
	"strconv"
	"sync/atomic"
	"time"
)

type counter struct {
	lastSuccValue int64

	currSuccValue int64

	lastFailValue int64

	currFailValue int64
}

// 用于向flume中作为sink 通过thrift客户端写入日志

type SinkServer struct {
	redisPool       map[string][]*redis.Pool
	flumeClientPool []*flumeClientPool
	isStop          bool
	monitorCount    counter
}

func NewSinkServer(option *config.Option) (server *SinkServer) {

	redisPool := make(map[string][]*redis.Pool, 0)

	//创建redis的消费连接
	for _, v := range option.QueueHostPorts {

		pool := redis.NewPool(func() (conn redis.Conn, err error) {

			conn, err = redis.DialTimeout("tcp", v.Host+":"+strconv.Itoa(v.Port),
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second,
				time.Duration(v.Timeout)*time.Second)

			return
		}, time.Duration(v.Timeout*2)*time.Second, v.Maxconn/2, v.Maxconn)

		pools, ok := redisPool[v.QueueName]
		if !ok {
			pools = make([]*redis.Pool, 0)
			redisPool[v.QueueName] = pools
		}

		redisPool[v.QueueName] = append(pools, pool)

	}

	pools := make([]*flumeClientPool, 0)
	//创建flume的client
	for _, v := range option.FlumeAgents {

		pool := newFlumeClientPool(20, 50, 100, 10*time.Second, func() *client.FlumeClient {
			flumeclient := client.NewFlumeClient(v.Host, v.Port)
			flumeclient.Connect()
			return flumeclient
		})
		pools = append(pools, pool)

		go monitorPool(v.Host+":"+strconv.Itoa(v.Port), pool)
	}

	sinkserver := &SinkServer{redisPool: redisPool, flumeClientPool: pools}
	go sinkserver.monitorFlume()
	return sinkserver
}

func (self *SinkServer) monitorFlume() {
	for {
		time.Sleep(1 * time.Second)
		currSucc := self.monitorCount.currSuccValue
		currFail := self.monitorCount.currFailValue
		log.Printf("succ-send:%d,fail-send:%d",
			(currSucc - self.monitorCount.lastSuccValue),
			(currFail - self.monitorCount.lastFailValue))
		self.monitorCount.lastSuccValue = currSucc
		self.monitorCount.lastFailValue = currFail
	}
}

func monitorPool(hostport string, pool *flumeClientPool) {
	for {
		time.Sleep(1 * time.Second)

		log.Printf("flume:%s|active:%d,core:%d,max:%d",
			hostport, pool.ActivePoolSize(), pool.CorePoolSize(), pool.maxPoolSize)
	}

}

//启动pop
func (self *SinkServer) Start() {

	self.isStop = false
	ch := make(chan int, 1)
	var count = 0
	for k, v := range self.redisPool {

		log.Println("start redis queueserver succ " + k)
		for _, pool := range v {
			count++
			defer pool.Close()
			go func(queuename string, pool *redis.Pool, end chan int) {
				conn := pool.Get()
				defer pool.Release(conn)
				for !self.isStop {

					// log.Println("pool active count :", strconv.Itoa(pool.ActiveCount()))
					reply, err := conn.Do("LPOP", queuename)
					if nil != err || nil == reply {
						if nil != err {
							log.Printf("LPOP|FAIL|%s", err)
							conn.Close()
							conn = pool.Get()
						} else {
							time.Sleep(100 * time.Millisecond)
						}

						continue
					}

					resp := reply.([]byte)
					var cmd config.Command
					err = json.Unmarshal(resp, &cmd)

					if nil != err {
						log.Printf("command unmarshal fail ! %s | error:%s\n", resp, err.Error())
						continue
					} else if rand.Int()%10 == 0 {
						log.Println("trace|command|%s", cmd)
					}

					//
					momoid := cmd.Params["momoid"].(string)

					businessName := cmd.Params["businessName"].(string)

					action := cmd.Params["type"].(string)

					bodyContent := cmd.Params["body"]

					//将businessName 加入到body中
					bodyMap := bodyContent.(map[string]interface{})
					bodyMap["business_type"] = businessName

					body, err := json.Marshal(bodyContent)

					if nil != err {
						log.Printf("marshal log body fail %s", err.Error())
						continue
					}

					//这里需要优化一下body,需要采用其他的方式定义Body格式，写入

					// log.Printf("%s,%s,%s,%s", momoid, businessName, action, string(body))

					//启动处理任务
					go self.innerSend(momoid, businessName, action, string(body))

				}
				end <- -1
			}(k, pool, ch)
		}
	}

	for {
		count += <-ch
		if count <= 0 {
			log.Printf("redis conn  close %d", count)
			break
		}
	}

}

func (self *SinkServer) innerSend(momoid, businessName, action string, body string) {

	for i := 0; i < 3; i++ {
		pool := self.getFlumeClientPool(businessName, action)
		flumeclient, err := pool.Get(5 * time.Second)
		if nil != err {
			continue
		}
		//拼装头部信息
		header := make(map[string]string, 1)
		header["businessName"] = businessName
		header["type"] = action

		//拼Body
		flumeBody := fmt.Sprintf("%s\t%s\t%s", momoid, action, body)
		err = flumeclient.Append(header, []byte(flumeBody))
		defer func() {
			if err := recover(); nil != err {
				//回收这个坏的连接
				pool.ReleaseBroken(flumeclient)
			}
		}()

		if nil != err {
			atomic.AddInt64(&self.monitorCount.currFailValue, 1)
			log.Printf("send 2 flume fail %s \t err:%s\n", body, err.Error())
		} else {
			atomic.AddInt64(&self.monitorCount.currSuccValue, 1)
			// log.Printf("send 2 flume succ %s\n", body)
			pool.Release(flumeclient)
			break
		}
	}
}

//仅供测试使用推送数据
func (self *SinkServer) testPushLog(queuename, logger string) {

	for _, v := range self.redisPool {
		for _, pool := range v {
			conn := pool.Get()
			defer pool.Release(conn)

			reply, err := conn.Do("RPUSH", queuename, logger)
			log.Printf("%s|err:%s", reply, err)
			break

		}
	}

}

func (self *SinkServer) Stop() {
	self.isStop = true
	for _, v := range self.flumeClientPool {
		v.Destroy()
	}

	for _, v := range self.redisPool {
		for _, p := range v {
			p.Close()
		}
	}
}

func (self *SinkServer) getFlumeClientPool(businessName, action string) *flumeClientPool {

	//使用随机算法直接获得

	idx := rand.Intn(len(self.flumeClientPool))
	return self.flumeClientPool[idx]

}
<|endoftext|>


<commit_before>// Package ldap implements strategies for authenticating using the LDAP protocol.
package ldap

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"strings"
	"unicode"

	"gopkg.in/ldap.v2"

	"github.com/coreos/dex/connector"
)

// Config holds the configuration parameters for the LDAP connector. The LDAP
// connectors require executing two queries, the first to find the user based on
// the username and password given to the connector. The second to use the user
// entry to search for groups.
//
// An example config:
//
//     type: ldap
//     config:
//       host: ldap.example.com:636
//       # The following field is required if using port 389.
//       # insecureNoSSL: true
//       rootCA: /etc/dex/ldap.ca
//       bindDN: uid=seviceaccount,cn=users,dc=example,dc=com
//       bindPW: password
//       userSearch:
//         # Would translate to the query "(&(objectClass=person)(uid=<username>))"
//         baseDN: cn=users,dc=example,dc=com
//         filter: "(objectClass=person)"
//         username: uid
//         idAttr: uid
//         emailAttr: mail
//         nameAttr: name
//       groupSearch:
//         # Would translate to the query "(&(objectClass=group)(member=<user uid>))"
//         baseDN: cn=groups,dc=example,dc=com
//         filter: "(objectClass=group)"
//         userAttr: uid
//         groupAttr: member
//         nameAttr: name
//
type Config struct {
	// The host and optional port of the LDAP server. If port isn't supplied, it will be
	// guessed based on the TLS configuration. 389 or 636.
	Host string `yaml:"host"`

	// Required if LDAP host does not use TLS.
	InsecureNoSSL bool `yaml:"insecureNoSSL"`

	// Path to a trusted root certificate file.
	RootCA string `yaml:"rootCA"`

	// BindDN and BindPW for an application service account. The connector uses these
	// credentials to search for users and groups.
	BindDN string `yaml:"bindDN"`
	BindPW string `yaml:"bindPW"`

	// User entry search configuration.
	UserSearch struct {
		// BsaeDN to start the search from. For example "cn=users,dc=example,dc=com"
		BaseDN string `yaml:"baseDN"`

		// Optional filter to apply when searching the directory. For example "(objectClass=person)"
		Filter string `yaml:"filter"`

		// Attribute to match against the inputted username. This will be translated and combined
		// with the other filter as "(<attr>=<username>)".
		Username string `yaml:"username"`

		// Can either be:
		// * "sub" - search the whole sub tree
		// * "one" - only search one level
		Scope string `yaml:"scope"`

		// A mapping of attributes on the user entry to claims.
		IDAttr    string `yaml:"idAttr"`    // Defaults to "uid"
		EmailAttr string `yaml:"emailAttr"` // Defaults to "mail"
		NameAttr  string `yaml:"nameAttr"`  // No default.

	} `yaml:"userSearch"`

	// Group search configuration.
	GroupSearch struct {
		// BsaeDN to start the search from. For example "cn=groups,dc=example,dc=com"
		BaseDN string `yaml:"baseDN"`

		// Optional filter to apply when searching the directory. For example "(objectClass=posixGroup)"
		Filter string `yaml:"filter"`

		Scope string `yaml:"scope"` // Defaults to "sub"

		// These two fields are use to match a user to a group.
		//
		// It adds an additional requirement to the filter that an attribute in the group
		// match the user's attribute value. For example that the "members" attribute of
		// a group matches the "uid" of the user. The exact filter being added is:
		//
		//   (<groupAttr>=<userAttr value>)
		//
		UserAttr  string `yaml:"userAttr"`
		GroupAttr string `yaml:"groupAttr"`

		// The attribute of the group that represents its name.
		NameAttr string `yaml:"nameAttr"`
	} `yaml:"groupSearch"`
}

func parseScope(s string) (int, bool) {
	// NOTE(ericchiang): ScopeBaseObject doesn't really make sense for us because we
	// never know the user's or group's DN.
	switch s {
	case "", "sub":
		return ldap.ScopeWholeSubtree, true
	case "one":
		return ldap.ScopeSingleLevel, true
	}
	return 0, false
}

// escapeRune maps a rune to a hex encoded value. For example 'é' would become '\\c3\\a9'
func escapeRune(buff *bytes.Buffer, r rune) {
	// Really inefficient, but it seems correct.
	for _, b := range []byte(string(r)) {
		buff.WriteString("\\")
		buff.WriteString(hex.EncodeToString([]byte{b}))
	}
}

// NOTE(ericchiang): There are no good documents on how to escape an LDAP string.
// This implementation is inspired by an Oracle document, and is purposefully
// extremely restrictive.
//
// See: https://docs.oracle.com/cd/E19424-01/820-4811/gdxpo/index.html
func escapeFilter(s string) string {
	r := strings.NewReader(s)
	buff := new(bytes.Buffer)
	for {
		ru, _, err := r.ReadRune()
		if err != nil {
			// ignore decoding issues
			return buff.String()
		}

		switch {
		case ru > unicode.MaxASCII: // Not ASCII
			escapeRune(buff, ru)
		case !unicode.IsPrint(ru): // Not printable
			escapeRune(buff, ru)
		case strings.ContainsRune(`*\()`, ru): // Reserved characters
			escapeRune(buff, ru)
		default:
			buff.WriteRune(ru)
		}
	}
}

// Open returns an authentication strategy using LDAP.
func (c *Config) Open() (connector.Connector, error) {
	requiredFields := []struct {
		name string
		val  string
	}{
		{"host", c.Host},
		{"userSearch.baseDN", c.UserSearch.BaseDN},
		{"userSearch.username", c.UserSearch.Username},
	}

	for _, field := range requiredFields {
		if field.val == "" {
			return nil, fmt.Errorf("ldap: missing required field %q", field.name)
		}
	}

	var (
		host string
		err  error
	)
	if host, _, err = net.SplitHostPort(c.Host); err != nil {
		host = c.Host
		if c.InsecureNoSSL {
			c.Host = c.Host + ":389"
		} else {
			c.Host = c.Host + ":636"
		}
	}

	tlsConfig := new(tls.Config)
	if c.RootCA != "" {
		data, err := ioutil.ReadFile(c.RootCA)
		if err != nil {
			return nil, fmt.Errorf("ldap: read ca file: %v", err)
		}
		rootCAs := x509.NewCertPool()
		if !rootCAs.AppendCertsFromPEM(data) {
			return nil, fmt.Errorf("ldap: no certs found in ca file")
		}
		tlsConfig.RootCAs = rootCAs
		// NOTE(ericchiang): This was required for our internal LDAP server
		// but might be because of an issue with our root CA.
		tlsConfig.ServerName = host
	}
	userSearchScope, ok := parseScope(c.UserSearch.Scope)
	if !ok {
		return nil, fmt.Errorf("userSearch.Scope unknown value %q", c.UserSearch.Scope)
	}
	groupSearchScope, ok := parseScope(c.GroupSearch.Scope)
	if !ok {
		return nil, fmt.Errorf("userSearch.Scope unknown value %q", c.GroupSearch.Scope)
	}
	return &ldapConnector{*c, userSearchScope, groupSearchScope, tlsConfig}, nil
}

type ldapConnector struct {
	Config

	userSearchScope  int
	groupSearchScope int

	tlsConfig *tls.Config
}

var _ connector.PasswordConnector = (*ldapConnector)(nil)

// do initializes a connection to the LDAP directory and passes it to the
// provided function. It then performs appropriate teardown or reuse before
// returning.
func (c *ldapConnector) do(f func(c *ldap.Conn) error) error {
	var (
		conn *ldap.Conn
		err  error
	)
	if c.InsecureNoSSL {
		conn, err = ldap.Dial("tcp", c.Host)
	} else {
		conn, err = ldap.DialTLS("tcp", c.Host, c.tlsConfig)
	}
	if err != nil {
		return fmt.Errorf("failed to connect: %v", err)
	}
	defer conn.Close()

	// If bindDN and bindPW are empty this will default to an anonymous bind.
	if err := conn.Bind(c.BindDN, c.BindPW); err != nil {
		return fmt.Errorf("ldap: initial bind for user %q failed: %v", c.BindDN, err)
	}

	return f(conn)
}

func getAttr(e ldap.Entry, name string) string {
	for _, a := range e.Attributes {
		if a.Name != name {
			continue
		}
		if len(a.Values) == 0 {
			return ""
		}
		return a.Values[0]
	}
	return ""
}

func (c *ldapConnector) Login(username, password string) (ident connector.Identity, validPass bool, err error) {
	var (
		// We want to return a different error if the user's password is incorrect vs
		// if there was an error.
		incorrectPass = false
		user          ldap.Entry
	)

	filter := fmt.Sprintf("(%s=%s)", c.UserSearch.Username, escapeFilter(username))
	if c.UserSearch.Filter != "" {
		filter = fmt.Sprintf("(&%s%s)", c.UserSearch.Filter, filter)
	}

	// Initial search.
	req := &ldap.SearchRequest{
		BaseDN: c.UserSearch.BaseDN,
		Filter: filter,
		Scope:  c.userSearchScope,
		// We only need to search for these specific requests.
		Attributes: []string{
			c.UserSearch.IDAttr,
			c.UserSearch.EmailAttr,
			c.GroupSearch.UserAttr,
			// TODO(ericchiang): what if this contains duplicate values?
		},
	}

	if c.UserSearch.NameAttr != "" {
		req.Attributes = append(req.Attributes, c.UserSearch.NameAttr)
	}

	err = c.do(func(conn *ldap.Conn) error {
		resp, err := conn.Search(req)
		if err != nil {
			return fmt.Errorf("ldap: search with filter %q failed: %v", req.Filter, err)
		}

		switch n := len(resp.Entries); n {
		case 0:
			return fmt.Errorf("ldap: no results returned for filter: %q", filter)
		case 2:
			return fmt.Errorf("ldap: filter returned multiple (%d) results: %q", n, filter)
		}

		user = *resp.Entries[0]

		// Try to authenticate as the distinguished name.
		if err := conn.Bind(user.DN, password); err != nil {
			// Detect a bad password through the LDAP error code.
			if ldapErr, ok := err.(*ldap.Error); ok {
				if ldapErr.ResultCode == ldap.LDAPResultInvalidCredentials {
					log.Printf("ldap: invalid password for user %q", user.DN)
					incorrectPass = true
					return nil
				}
			}
			return fmt.Errorf("ldap: failed to bind as dn %q: %v", user.DN, err)
		}
		return nil
	})
	if err != nil {
		return connector.Identity{}, false, err
	}

	// Encode entry for follow up requests such as the groups query and
	// refresh attempts.
	if ident.ConnectorData, err = json.Marshal(user); err != nil {
		return connector.Identity{}, false, fmt.Errorf("ldap: marshal entry: %v", err)
	}

	// If we're missing any attributes, such as email or ID, we want to report
	// an error rather than continuing.
	missing := []string{}

	// Fill the identity struct using the attributes from the user entry.
	if ident.UserID = getAttr(user, c.UserSearch.IDAttr); ident.UserID == "" {
		missing = append(missing, c.UserSearch.IDAttr)
	}
	if ident.Email = getAttr(user, c.UserSearch.EmailAttr); ident.Email == "" {
		missing = append(missing, c.UserSearch.EmailAttr)
	}
	if c.UserSearch.NameAttr != "" {
		if ident.Username = getAttr(user, c.UserSearch.NameAttr); ident.Username == "" {
			missing = append(missing, c.UserSearch.NameAttr)
		}
	}

	if len(missing) != 0 {
		err := fmt.Errorf("ldap: entry %q missing following required attribute(s): %q", user.DN, missing)
		return connector.Identity{}, false, err
	}

	return ident, !incorrectPass, nil
}

func (c *ldapConnector) Groups(ident connector.Identity) ([]string, error) {
	// Decode the user entry from the identity.
	var user ldap.Entry
	if err := json.Unmarshal(ident.ConnectorData, &user); err != nil {
		return nil, fmt.Errorf("ldap: failed to unmarshal connector data: %v", err)
	}

	filter := fmt.Sprintf("(%s=%s)", c.GroupSearch.GroupAttr, escapeFilter(getAttr(user, c.GroupSearch.UserAttr)))
	if c.GroupSearch.Filter != "" {
		filter = fmt.Sprintf("(&%s%s)", c.GroupSearch.Filter, filter)
	}

	req := &ldap.SearchRequest{
		BaseDN:     c.GroupSearch.BaseDN,
		Filter:     filter,
		Scope:      c.groupSearchScope,
		Attributes: []string{c.GroupSearch.NameAttr},
	}

	var groups []*ldap.Entry
	if err := c.do(func(conn *ldap.Conn) error {
		resp, err := conn.Search(req)
		if err != nil {
			return fmt.Errorf("ldap: search failed: %v", err)
		}
		groups = resp.Entries
		return nil
	}); err != nil {
		return nil, err
	}
	if len(groups) == 0 {
		// TODO(ericchiang): Is this going to spam the logs?
		log.Printf("ldap: groups search with filter %q returned no groups", filter)
	}

	var groupNames []string

	for _, group := range groups {
		name := getAttr(*group, c.GroupSearch.NameAttr)
		if name == "" {
			// Be obnoxious about missing missing attributes. If the group entry is
			// missing its name attribute, that indicates a misconfiguration.
			//
			// In the future we can add configuration options to just log these errors.
			return nil, fmt.Errorf("ldap: group entity %q missing required attribute %q",
				group.DN, c.GroupSearch.NameAttr)
		}

		groupNames = append(groupNames, name)
	}
	return groupNames, nil
}

func (c *ldapConnector) Close() error {
	return nil
}
<commit_msg>connector/ldap: fix bug in switch statement<commit_after>// Package ldap implements strategies for authenticating using the LDAP protocol.
package ldap

import (
	"bytes"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"encoding/json"
	"fmt"
	"io/ioutil"
	"log"
	"net"
	"strings"
	"unicode"

	"gopkg.in/ldap.v2"

	"github.com/coreos/dex/connector"
)

// Config holds the configuration parameters for the LDAP connector. The LDAP
// connectors require executing two queries, the first to find the user based on
// the username and password given to the connector. The second to use the user
// entry to search for groups.
//
// An example config:
//
//     type: ldap
//     config:
//       host: ldap.example.com:636
//       # The following field is required if using port 389.
//       # insecureNoSSL: true
//       rootCA: /etc/dex/ldap.ca
//       bindDN: uid=seviceaccount,cn=users,dc=example,dc=com
//       bindPW: password
//       userSearch:
//         # Would translate to the query "(&(objectClass=person)(uid=<username>))"
//         baseDN: cn=users,dc=example,dc=com
//         filter: "(objectClass=person)"
//         username: uid
//         idAttr: uid
//         emailAttr: mail
//         nameAttr: name
//       groupSearch:
//         # Would translate to the query "(&(objectClass=group)(member=<user uid>))"
//         baseDN: cn=groups,dc=example,dc=com
//         filter: "(objectClass=group)"
//         userAttr: uid
//         groupAttr: member
//         nameAttr: name
//
type Config struct {
	// The host and optional port of the LDAP server. If port isn't supplied, it will be
	// guessed based on the TLS configuration. 389 or 636.
	Host string `yaml:"host"`

	// Required if LDAP host does not use TLS.
	InsecureNoSSL bool `yaml:"insecureNoSSL"`

	// Path to a trusted root certificate file.
	RootCA string `yaml:"rootCA"`

	// BindDN and BindPW for an application service account. The connector uses these
	// credentials to search for users and groups.
	BindDN string `yaml:"bindDN"`
	BindPW string `yaml:"bindPW"`

	// User entry search configuration.
	UserSearch struct {
		// BsaeDN to start the search from. For example "cn=users,dc=example,dc=com"
		BaseDN string `yaml:"baseDN"`

		// Optional filter to apply when searching the directory. For example "(objectClass=person)"
		Filter string `yaml:"filter"`

		// Attribute to match against the inputted username. This will be translated and combined
		// with the other filter as "(<attr>=<username>)".
		Username string `yaml:"username"`

		// Can either be:
		// * "sub" - search the whole sub tree
		// * "one" - only search one level
		Scope string `yaml:"scope"`

		// A mapping of attributes on the user entry to claims.
		IDAttr    string `yaml:"idAttr"`    // Defaults to "uid"
		EmailAttr string `yaml:"emailAttr"` // Defaults to "mail"
		NameAttr  string `yaml:"nameAttr"`  // No default.

	} `yaml:"userSearch"`

	// Group search configuration.
	GroupSearch struct {
		// BsaeDN to start the search from. For example "cn=groups,dc=example,dc=com"
		BaseDN string `yaml:"baseDN"`

		// Optional filter to apply when searching the directory. For example "(objectClass=posixGroup)"
		Filter string `yaml:"filter"`

		Scope string `yaml:"scope"` // Defaults to "sub"

		// These two fields are use to match a user to a group.
		//
		// It adds an additional requirement to the filter that an attribute in the group
		// match the user's attribute value. For example that the "members" attribute of
		// a group matches the "uid" of the user. The exact filter being added is:
		//
		//   (<groupAttr>=<userAttr value>)
		//
		UserAttr  string `yaml:"userAttr"`
		GroupAttr string `yaml:"groupAttr"`

		// The attribute of the group that represents its name.
		NameAttr string `yaml:"nameAttr"`
	} `yaml:"groupSearch"`
}

func parseScope(s string) (int, bool) {
	// NOTE(ericchiang): ScopeBaseObject doesn't really make sense for us because we
	// never know the user's or group's DN.
	switch s {
	case "", "sub":
		return ldap.ScopeWholeSubtree, true
	case "one":
		return ldap.ScopeSingleLevel, true
	}
	return 0, false
}

// escapeRune maps a rune to a hex encoded value. For example 'é' would become '\\c3\\a9'
func escapeRune(buff *bytes.Buffer, r rune) {
	// Really inefficient, but it seems correct.
	for _, b := range []byte(string(r)) {
		buff.WriteString("\\")
		buff.WriteString(hex.EncodeToString([]byte{b}))
	}
}

// NOTE(ericchiang): There are no good documents on how to escape an LDAP string.
// This implementation is inspired by an Oracle document, and is purposefully
// extremely restrictive.
//
// See: https://docs.oracle.com/cd/E19424-01/820-4811/gdxpo/index.html
func escapeFilter(s string) string {
	r := strings.NewReader(s)
	buff := new(bytes.Buffer)
	for {
		ru, _, err := r.ReadRune()
		if err != nil {
			// ignore decoding issues
			return buff.String()
		}

		switch {
		case ru > unicode.MaxASCII: // Not ASCII
			escapeRune(buff, ru)
		case !unicode.IsPrint(ru): // Not printable
			escapeRune(buff, ru)
		case strings.ContainsRune(`*\()`, ru): // Reserved characters
			escapeRune(buff, ru)
		default:
			buff.WriteRune(ru)
		}
	}
}

// Open returns an authentication strategy using LDAP.
func (c *Config) Open() (connector.Connector, error) {
	requiredFields := []struct {
		name string
		val  string
	}{
		{"host", c.Host},
		{"userSearch.baseDN", c.UserSearch.BaseDN},
		{"userSearch.username", c.UserSearch.Username},
	}

	for _, field := range requiredFields {
		if field.val == "" {
			return nil, fmt.Errorf("ldap: missing required field %q", field.name)
		}
	}

	var (
		host string
		err  error
	)
	if host, _, err = net.SplitHostPort(c.Host); err != nil {
		host = c.Host
		if c.InsecureNoSSL {
			c.Host = c.Host + ":389"
		} else {
			c.Host = c.Host + ":636"
		}
	}

	tlsConfig := new(tls.Config)
	if c.RootCA != "" {
		data, err := ioutil.ReadFile(c.RootCA)
		if err != nil {
			return nil, fmt.Errorf("ldap: read ca file: %v", err)
		}
		rootCAs := x509.NewCertPool()
		if !rootCAs.AppendCertsFromPEM(data) {
			return nil, fmt.Errorf("ldap: no certs found in ca file")
		}
		tlsConfig.RootCAs = rootCAs
		// NOTE(ericchiang): This was required for our internal LDAP server
		// but might be because of an issue with our root CA.
		tlsConfig.ServerName = host
	}
	userSearchScope, ok := parseScope(c.UserSearch.Scope)
	if !ok {
		return nil, fmt.Errorf("userSearch.Scope unknown value %q", c.UserSearch.Scope)
	}
	groupSearchScope, ok := parseScope(c.GroupSearch.Scope)
	if !ok {
		return nil, fmt.Errorf("userSearch.Scope unknown value %q", c.GroupSearch.Scope)
	}
	return &ldapConnector{*c, userSearchScope, groupSearchScope, tlsConfig}, nil
}

type ldapConnector struct {
	Config

	userSearchScope  int
	groupSearchScope int

	tlsConfig *tls.Config
}

var _ connector.PasswordConnector = (*ldapConnector)(nil)

// do initializes a connection to the LDAP directory and passes it to the
// provided function. It then performs appropriate teardown or reuse before
// returning.
func (c *ldapConnector) do(f func(c *ldap.Conn) error) error {
	var (
		conn *ldap.Conn
		err  error
	)
	if c.InsecureNoSSL {
		conn, err = ldap.Dial("tcp", c.Host)
	} else {
		conn, err = ldap.DialTLS("tcp", c.Host, c.tlsConfig)
	}
	if err != nil {
		return fmt.Errorf("failed to connect: %v", err)
	}
	defer conn.Close()

	// If bindDN and bindPW are empty this will default to an anonymous bind.
	if err := conn.Bind(c.BindDN, c.BindPW); err != nil {
		return fmt.Errorf("ldap: initial bind for user %q failed: %v", c.BindDN, err)
	}

	return f(conn)
}

func getAttr(e ldap.Entry, name string) string {
	for _, a := range e.Attributes {
		if a.Name != name {
			continue
		}
		if len(a.Values) == 0 {
			return ""
		}
		return a.Values[0]
	}
	return ""
}

func (c *ldapConnector) Login(username, password string) (ident connector.Identity, validPass bool, err error) {
	var (
		// We want to return a different error if the user's password is incorrect vs
		// if there was an error.
		incorrectPass = false
		user          ldap.Entry
	)

	filter := fmt.Sprintf("(%s=%s)", c.UserSearch.Username, escapeFilter(username))
	if c.UserSearch.Filter != "" {
		filter = fmt.Sprintf("(&%s%s)", c.UserSearch.Filter, filter)
	}

	// Initial search.
	req := &ldap.SearchRequest{
		BaseDN: c.UserSearch.BaseDN,
		Filter: filter,
		Scope:  c.userSearchScope,
		// We only need to search for these specific requests.
		Attributes: []string{
			c.UserSearch.IDAttr,
			c.UserSearch.EmailAttr,
			c.GroupSearch.UserAttr,
			// TODO(ericchiang): what if this contains duplicate values?
		},
	}

	if c.UserSearch.NameAttr != "" {
		req.Attributes = append(req.Attributes, c.UserSearch.NameAttr)
	}

	err = c.do(func(conn *ldap.Conn) error {
		resp, err := conn.Search(req)
		if err != nil {
			return fmt.Errorf("ldap: search with filter %q failed: %v", req.Filter, err)
		}

		switch n := len(resp.Entries); n {
		case 0:
			return fmt.Errorf("ldap: no results returned for filter: %q", filter)
		case 1:
		default:
			return fmt.Errorf("ldap: filter returned multiple (%d) results: %q", n, filter)
		}

		user = *resp.Entries[0]

		// Try to authenticate as the distinguished name.
		if err := conn.Bind(user.DN, password); err != nil {
			// Detect a bad password through the LDAP error code.
			if ldapErr, ok := err.(*ldap.Error); ok {
				if ldapErr.ResultCode == ldap.LDAPResultInvalidCredentials {
					log.Printf("ldap: invalid password for user %q", user.DN)
					incorrectPass = true
					return nil
				}
			}
			return fmt.Errorf("ldap: failed to bind as dn %q: %v", user.DN, err)
		}
		return nil
	})
	if err != nil {
		return connector.Identity{}, false, err
	}

	// Encode entry for follow up requests such as the groups query and
	// refresh attempts.
	if ident.ConnectorData, err = json.Marshal(user); err != nil {
		return connector.Identity{}, false, fmt.Errorf("ldap: marshal entry: %v", err)
	}

	// If we're missing any attributes, such as email or ID, we want to report
	// an error rather than continuing.
	missing := []string{}

	// Fill the identity struct using the attributes from the user entry.
	if ident.UserID = getAttr(user, c.UserSearch.IDAttr); ident.UserID == "" {
		missing = append(missing, c.UserSearch.IDAttr)
	}
	if ident.Email = getAttr(user, c.UserSearch.EmailAttr); ident.Email == "" {
		missing = append(missing, c.UserSearch.EmailAttr)
	}
	if c.UserSearch.NameAttr != "" {
		if ident.Username = getAttr(user, c.UserSearch.NameAttr); ident.Username == "" {
			missing = append(missing, c.UserSearch.NameAttr)
		}
	}

	if len(missing) != 0 {
		err := fmt.Errorf("ldap: entry %q missing following required attribute(s): %q", user.DN, missing)
		return connector.Identity{}, false, err
	}

	return ident, !incorrectPass, nil
}

func (c *ldapConnector) Groups(ident connector.Identity) ([]string, error) {
	// Decode the user entry from the identity.
	var user ldap.Entry
	if err := json.Unmarshal(ident.ConnectorData, &user); err != nil {
		return nil, fmt.Errorf("ldap: failed to unmarshal connector data: %v", err)
	}

	filter := fmt.Sprintf("(%s=%s)", c.GroupSearch.GroupAttr, escapeFilter(getAttr(user, c.GroupSearch.UserAttr)))
	if c.GroupSearch.Filter != "" {
		filter = fmt.Sprintf("(&%s%s)", c.GroupSearch.Filter, filter)
	}

	req := &ldap.SearchRequest{
		BaseDN:     c.GroupSearch.BaseDN,
		Filter:     filter,
		Scope:      c.groupSearchScope,
		Attributes: []string{c.GroupSearch.NameAttr},
	}

	var groups []*ldap.Entry
	if err := c.do(func(conn *ldap.Conn) error {
		resp, err := conn.Search(req)
		if err != nil {
			return fmt.Errorf("ldap: search failed: %v", err)
		}
		groups = resp.Entries
		return nil
	}); err != nil {
		return nil, err
	}
	if len(groups) == 0 {
		// TODO(ericchiang): Is this going to spam the logs?
		log.Printf("ldap: groups search with filter %q returned no groups", filter)
	}

	var groupNames []string

	for _, group := range groups {
		name := getAttr(*group, c.GroupSearch.NameAttr)
		if name == "" {
			// Be obnoxious about missing missing attributes. If the group entry is
			// missing its name attribute, that indicates a misconfiguration.
			//
			// In the future we can add configuration options to just log these errors.
			return nil, fmt.Errorf("ldap: group entity %q missing required attribute %q",
				group.DN, c.GroupSearch.NameAttr)
		}

		groupNames = append(groupNames, name)
	}
	return groupNames, nil
}

func (c *ldapConnector) Close() error {
	return nil
}
<|endoftext|>


<commit_before>package routing

import (
	"math"

	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcutil"
)

const (
	// HopLimit is the maximum number hops that is permissible as a route.
	// Any potential paths found that lie above this limit will be rejected
	// with an error. This value is computed using the current fixed-size
	// packet length of the Sphinx construction.
	HopLimit = 20

	// infinity is used as a starting distance in our shortest path search.
	infinity = math.MaxFloat64
)

// Route represents a path through the channel graph which runs over one or
// more channels in succession. This struct carries all the information
// required to craft the Sphinx onion packet, and send the payment along the
// first hop in the path. A route is only selected as valid if all the channels
// have sufficient capacity to carry the initial payment amount after fees are
// accounted for.
type Route struct {
	// TotalTimeLock is the cumulative (final) time lock across the entire
	// route. This is the CLTV value that should be extended to the first
	// hop in the route. All other hops will decrement the time-lock as
	// advertised, leaving enough time for all hops to wait for or present
	// the payment pre-image to complete the payment.
	TotalTimeLock uint32

	// TotalFees is the sum of the fees paid at each hop within the final
	// route. In the case of a one-hop payment, this value will be zero as
	// we don't need to pay a fee it ourself.
	TotalFees btcutil.Amount

	// TotalAmount is the total amount of funds required to complete a
	// payment over this route. This value includes the cumulative fees at
	// each hop. As a result, the HTLC extended to the first-hop in the
	// route will need to have at least this many satoshis, otherwise the
	// route will fail at an intermediate node due to an insufficient
	// amount of fees.
	TotalAmount btcutil.Amount

	// Hops contains details concerning the specific forwarding details at
	// each hop.
	Hops []*Hop
}

// Hop represents the forwarding details at a particular position within the
// final route. This struct houses the values necessary to create the HTLC
// which will travel along this hop, and also encode the per-hop payload
// included within the Sphinx packet.
type Hop struct {
	// Channels is the active payment channel that this hop will travel
	// along.
	Channel *channeldb.ChannelEdge

	// TimeLockDelta is the delta that this hop will subtract from the HTLC
	// before extending it to the next hop in the route.
	TimeLockDelta uint16

	// AmtToForward is the amount that this hop will forward to the next
	// hop. This value is less than the value that the incoming HTLC
	// carries as a fee will be subtracted by the hop.
	AmtToForward btcutil.Amount

	// Fee is the total fee that this hop will subtract from the incoming
	// payment, this difference nets the hop fees for forwarding the
	// payment.
	Fee btcutil.Amount
}

// computeFee computes the fee to forward an HTLC of `amt` satoshis over the
// passed active payment channel. This value is currently computed as specified
// in BOLT07, but will likely change in the near future.
func computeFee(amt btcutil.Amount, edge *channeldb.ChannelEdge) btcutil.Amount {
	return edge.FeeBaseMSat + (amt*edge.FeeProportionalMillionths)/1000000
}

// newRoute returns a fully valid route between the source and target that's
// capable of supporting a payment of `amtToSend` after fees are fully
// computed. IF the route is too long, or the selected path cannot support the
// fully payment including fees, then a non-nil error is returned. prevHop maps
// a vertex to the channel required to get to it.
func newRoute(amtToSend btcutil.Amount, source, target vertex,
	prevHop map[vertex]edgeWithPrev) (*Route, error) {

	// As an initial sanity check, the potential route is immediately
	// invalidate if it spans more than 20 hops. The current Sphinx (onion
	// routing) implementation can only encode up to 20 hops as the entire
	// packet is fixed size. If this route is more than 20 hops, then it's
	// invalid.
	if len(prevHop) > HopLimit {
		return nil, ErrMaxHopsExceeded
	}

	// If the potential route if below the max hop limit, then we'll use
	// the prevHop map to unravel the path. We end up with a list of edges
	// in the reverse direction which we'll use to properly calculate the
	// timelock and fee values.
	pathEdges := make([]*channeldb.ChannelEdge, 0, len(prevHop))
	prev := target
	for prev != source { // TODO(roasbeef): assumes no cycles
		// Add the current hop to the limit of path edges then walk
		// backwards from this hop via the prev pointer for this hop
		// within the prevHop map.
		pathEdges = append(pathEdges, prevHop[prev].edge)
		prev = newVertex(prevHop[prev].prevNode)
	}

	route := &Route{
		Hops: make([]*Hop, len(pathEdges)),
	}

	// The running amount is the total amount of satoshis required at this
	// point in the route. We start this value at the amount we want to
	// send to the destination. This value will then get successively
	// larger as we compute the fees going backwards.
	runningAmt := amtToSend
	pathLength := len(pathEdges)
	for i, edge := range pathEdges {
		// Now we create the hop struct for this point in the route.
		// The amount to forward is the running amount, and we compute
		// the required fee based on this amount.
		nextHop := &Hop{
			Channel:       edge,
			AmtToForward:  runningAmt,
			Fee:           computeFee(runningAmt, edge),
			TimeLockDelta: edge.Expiry,
		}
		edge.Node.PubKey.Curve = nil

		// As a sanity check, we ensure that the selected channel has
		// enough capacity to forward the required amount which
		// includes the fee dictated at each hop.
		if nextHop.AmtToForward > nextHop.Channel.Capacity {
			return nil, ErrInsufficientCapacity
		}

		// We don't pay any fees to ourselves on the first-hop channel,
		// so we don't tally up the running fee and amount.
		if i != len(pathEdges)-1 {
			// For a node to forward an HTLC, then following
			// inequality most hold true: amt_in - fee >=
			// amt_to_forward. Therefore we add the fee this node
			// consumes in order to calculate the amount that it
			// show be forwarded by the prior node which is the
			// next hop in our loop.
			runningAmt += nextHop.Fee

			// Next we tally the total fees (thus far) in the
			// route, and also accumulate the total timelock in the
			// route by adding the node's time lock delta which is
			// the amount of blocks it'll subtract from the
			// incoming time lock.
			route.TotalFees += nextHop.Fee
		} else {
			nextHop.Fee = 0
		}

		route.TotalTimeLock += uint32(nextHop.TimeLockDelta)

		// Finally, as we're currently talking the route backwards, we
		// reverse the index in order to place this hop at the proper
		// spot in the forward direction of the route.
		route.Hops[pathLength-1-i] = nextHop
	}

	// The total amount required for this route will be the value the
	// source extends to the first hop in the route.
	route.TotalAmount = runningAmt

	return route, nil
}

// vertex is a simple alias for the serialization of a compressed Bitcoin
// public key.
type vertex [33]byte

// newVertex returns a new vertex given a public key.
func newVertex(pub *btcec.PublicKey) vertex {
	var v vertex
	copy(v[:], pub.SerializeCompressed())
	return v
}

// nodeWithDist is a helper struct that couples the distance from the current
// source to a node with a pointer to the node itself.
type nodeWithDist struct {
	dist float64
	node *channeldb.LightningNode
}

// edgeWithPrev is a helper struct used in path finding that couples an
// directional edge with the node's ID in the opposite direction.
type edgeWithPrev struct {
	edge     *channeldb.ChannelEdge
	prevNode *btcec.PublicKey
}

// edgeWeight computes the weight of an edge. This value is used when searching
// for the shortest path within the channel graph between two nodes. Currently
// this is just 1 + the cltv delta value required at this hop, this value
// should be tuned with experimental and empirical data.
//
// TODO(roasbeef): compute robust weight metric
func edgeWeight(e *channeldb.ChannelEdge) float64 {
	return float64(1 + e.Expiry)
}

// findRoute attempts to find a path from the source node within the
// ChannelGraph to the target node that's capable of supporting a payment of
// `amt` value. The current approach is used a multiple pass path finding
// algorithm. First we employ a modified version of Dijkstra's algorithm to
// find a potential set of shortest paths, the distance metric is related to
// the time-lock+fee along the route. Once we have a set of candidate routes,
// we calculate the required fee and time lock values running backwards along
// the route. The route that's selected is the one with the lowest total fee.
//
// TODO(roasbeef): make member, add caching
//  * add k-path
func findRoute(graph *channeldb.ChannelGraph, target *btcec.PublicKey,
	amt btcutil.Amount) (*Route, error) {

	// First initialize empty list of all the node that we've yet to
	// visited.
	// TODO(roasbeef): make into incremental fibonacci heap rather than
	// loading all into memory.
	var unvisited []*channeldb.LightningNode

	// For each node/vertex the graph we create an entry in the distance
	// map for the node set with a distance of "infinity".  We also mark
	// add the node to our set of unvisited nodes.
	distance := make(map[vertex]nodeWithDist)
	if err := graph.ForEachNode(func(node *channeldb.LightningNode) error {
		// TODO(roasbeef): with larger graph can just use disk seeks
		// with a visited map
		distance[newVertex(node.PubKey)] = nodeWithDist{
			dist: infinity,
			node: node,
		}

		unvisited = append(unvisited, node)
		return nil
	}); err != nil {
		return nil, err
	}

	// Next we obtain the source node from the graph, and initialize it
	// with a distance of 0. This indicates our starting point in the graph
	// traversal.
	sourceNode, err := graph.SourceNode()
	if err != nil {
		return nil, err
	}
	sourceVertex := newVertex(sourceNode.PubKey)
	distance[sourceVertex] = nodeWithDist{
		dist: 0,
		node: sourceNode,
	}

	// We'll use this map as a series of "previous" hop pointers. So to get
	// to `vertex` we'll take the edge that it's mapped to within `prev`.
	prev := make(map[vertex]edgeWithPrev)

	for len(unvisited) != 0 {
		var bestNode *channeldb.LightningNode
		smallestDist := infinity

		// First we examine our list of unvisited nodes, for the most
		// optimal vertex to examine next.
		for i, node := range unvisited {
			// The "best" node to visit next is node with the
			// smallest distance from the source of all the
			// unvisited nodes.
			v := newVertex(node.PubKey)
			if nodeInfo := distance[v]; nodeInfo.dist < smallestDist {
				smallestDist = nodeInfo.dist
				bestNode = nodeInfo.node

				// Since we're going to visit this node, we can
				// remove it from the set of unvisited nodes.
				unvisited[i] = unvisited[len(unvisited)-1]
				unvisited[len(unvisited)-1] = nil // Avoid GC leak.
				unvisited = unvisited[:len(unvisited)-1]

				break
			}
		}

		// If we've reached our target, then we're done here and can
		// exit the graph traversal early.
		if bestNode.PubKey.IsEqual(target) {
			break
		}

		// Now that we've found the next potential step to take we'll
		// examine all the outgoing edge (channels) from this node to
		// further our graph traversal.
		pivot := newVertex(bestNode.PubKey)
		err := bestNode.ForEachChannel(nil, func(edge *channeldb.ChannelEdge) error {
			// Compute the tentative distance to this new
			// channel/edge which is the distance to our current
			// pivot node plus the weight of this edge.
			tempDist := distance[pivot].dist + edgeWeight(edge)

			// If this new tentative distance is better than the
			// current best known distance to this node, then we
			// record the new better distance, and also populate
			// our "next hop" map with this edge.
			// TODO(roasbeef): add capacity to relaxation criteria?
			//  * also add min payment?
			v := newVertex(edge.Node.PubKey)
			if tempDist < distance[v].dist {
				// TODO(roasbeef): unconditionally add for all
				// paths
				distance[v] = nodeWithDist{
					dist: tempDist,
					node: edge.Node,
				}
				prev[v] = edgeWithPrev{
					edge:     edge,
					prevNode: bestNode.PubKey,
				}
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
	}

	// If the target node isn't found in the prev hop map, then a path
	// doesn't exist, so we terminate in an error.
	if _, ok := prev[newVertex(target)]; !ok {
		return nil, ErrNoPathFound
	}

	// Otherwise, we construct a new route which calculate the relevant
	// total fees and proper time lock values for each hop.
	targetVerex := newVertex(target)
	return newRoute(amt, sourceVertex, targetVerex, prev)
}
<commit_msg>routing: fix panic in inner loop of path finding<commit_after>package routing

import (
	"math"

	"github.com/lightningnetwork/lnd/channeldb"
	"github.com/roasbeef/btcd/btcec"
	"github.com/roasbeef/btcutil"
)

const (
	// HopLimit is the maximum number hops that is permissible as a route.
	// Any potential paths found that lie above this limit will be rejected
	// with an error. This value is computed using the current fixed-size
	// packet length of the Sphinx construction.
	HopLimit = 20

	// infinity is used as a starting distance in our shortest path search.
	infinity = math.MaxFloat64
)

// Route represents a path through the channel graph which runs over one or
// more channels in succession. This struct carries all the information
// required to craft the Sphinx onion packet, and send the payment along the
// first hop in the path. A route is only selected as valid if all the channels
// have sufficient capacity to carry the initial payment amount after fees are
// accounted for.
type Route struct {
	// TotalTimeLock is the cumulative (final) time lock across the entire
	// route. This is the CLTV value that should be extended to the first
	// hop in the route. All other hops will decrement the time-lock as
	// advertised, leaving enough time for all hops to wait for or present
	// the payment pre-image to complete the payment.
	TotalTimeLock uint32

	// TotalFees is the sum of the fees paid at each hop within the final
	// route. In the case of a one-hop payment, this value will be zero as
	// we don't need to pay a fee it ourself.
	TotalFees btcutil.Amount

	// TotalAmount is the total amount of funds required to complete a
	// payment over this route. This value includes the cumulative fees at
	// each hop. As a result, the HTLC extended to the first-hop in the
	// route will need to have at least this many satoshis, otherwise the
	// route will fail at an intermediate node due to an insufficient
	// amount of fees.
	TotalAmount btcutil.Amount

	// Hops contains details concerning the specific forwarding details at
	// each hop.
	Hops []*Hop
}

// Hop represents the forwarding details at a particular position within the
// final route. This struct houses the values necessary to create the HTLC
// which will travel along this hop, and also encode the per-hop payload
// included within the Sphinx packet.
type Hop struct {
	// Channels is the active payment channel that this hop will travel
	// along.
	Channel *channeldb.ChannelEdge

	// TimeLockDelta is the delta that this hop will subtract from the HTLC
	// before extending it to the next hop in the route.
	TimeLockDelta uint16

	// AmtToForward is the amount that this hop will forward to the next
	// hop. This value is less than the value that the incoming HTLC
	// carries as a fee will be subtracted by the hop.
	AmtToForward btcutil.Amount

	// Fee is the total fee that this hop will subtract from the incoming
	// payment, this difference nets the hop fees for forwarding the
	// payment.
	Fee btcutil.Amount
}

// computeFee computes the fee to forward an HTLC of `amt` satoshis over the
// passed active payment channel. This value is currently computed as specified
// in BOLT07, but will likely change in the near future.
func computeFee(amt btcutil.Amount, edge *channeldb.ChannelEdge) btcutil.Amount {
	return edge.FeeBaseMSat + (amt*edge.FeeProportionalMillionths)/1000000
}

// newRoute returns a fully valid route between the source and target that's
// capable of supporting a payment of `amtToSend` after fees are fully
// computed. IF the route is too long, or the selected path cannot support the
// fully payment including fees, then a non-nil error is returned. prevHop maps
// a vertex to the channel required to get to it.
func newRoute(amtToSend btcutil.Amount, source, target vertex,
	prevHop map[vertex]edgeWithPrev) (*Route, error) {

	// As an initial sanity check, the potential route is immediately
	// invalidate if it spans more than 20 hops. The current Sphinx (onion
	// routing) implementation can only encode up to 20 hops as the entire
	// packet is fixed size. If this route is more than 20 hops, then it's
	// invalid.
	if len(prevHop) > HopLimit {
		return nil, ErrMaxHopsExceeded
	}

	// If the potential route if below the max hop limit, then we'll use
	// the prevHop map to unravel the path. We end up with a list of edges
	// in the reverse direction which we'll use to properly calculate the
	// timelock and fee values.
	pathEdges := make([]*channeldb.ChannelEdge, 0, len(prevHop))
	prev := target
	for prev != source { // TODO(roasbeef): assumes no cycles
		// Add the current hop to the limit of path edges then walk
		// backwards from this hop via the prev pointer for this hop
		// within the prevHop map.
		pathEdges = append(pathEdges, prevHop[prev].edge)
		prev = newVertex(prevHop[prev].prevNode)
	}

	route := &Route{
		Hops: make([]*Hop, len(pathEdges)),
	}

	// The running amount is the total amount of satoshis required at this
	// point in the route. We start this value at the amount we want to
	// send to the destination. This value will then get successively
	// larger as we compute the fees going backwards.
	runningAmt := amtToSend
	pathLength := len(pathEdges)
	for i, edge := range pathEdges {
		// Now we create the hop struct for this point in the route.
		// The amount to forward is the running amount, and we compute
		// the required fee based on this amount.
		nextHop := &Hop{
			Channel:       edge,
			AmtToForward:  runningAmt,
			Fee:           computeFee(runningAmt, edge),
			TimeLockDelta: edge.Expiry,
		}
		edge.Node.PubKey.Curve = nil

		// As a sanity check, we ensure that the selected channel has
		// enough capacity to forward the required amount which
		// includes the fee dictated at each hop.
		if nextHop.AmtToForward > nextHop.Channel.Capacity {
			return nil, ErrInsufficientCapacity
		}

		// We don't pay any fees to ourselves on the first-hop channel,
		// so we don't tally up the running fee and amount.
		if i != len(pathEdges)-1 {
			// For a node to forward an HTLC, then following
			// inequality most hold true: amt_in - fee >=
			// amt_to_forward. Therefore we add the fee this node
			// consumes in order to calculate the amount that it
			// show be forwarded by the prior node which is the
			// next hop in our loop.
			runningAmt += nextHop.Fee

			// Next we tally the total fees (thus far) in the
			// route, and also accumulate the total timelock in the
			// route by adding the node's time lock delta which is
			// the amount of blocks it'll subtract from the
			// incoming time lock.
			route.TotalFees += nextHop.Fee
		} else {
			nextHop.Fee = 0
		}

		route.TotalTimeLock += uint32(nextHop.TimeLockDelta)

		// Finally, as we're currently talking the route backwards, we
		// reverse the index in order to place this hop at the proper
		// spot in the forward direction of the route.
		route.Hops[pathLength-1-i] = nextHop
	}

	// The total amount required for this route will be the value the
	// source extends to the first hop in the route.
	route.TotalAmount = runningAmt

	return route, nil
}

// vertex is a simple alias for the serialization of a compressed Bitcoin
// public key.
type vertex [33]byte

// newVertex returns a new vertex given a public key.
func newVertex(pub *btcec.PublicKey) vertex {
	var v vertex
	copy(v[:], pub.SerializeCompressed())
	return v
}

// nodeWithDist is a helper struct that couples the distance from the current
// source to a node with a pointer to the node itself.
type nodeWithDist struct {
	dist float64
	node *channeldb.LightningNode
}

// edgeWithPrev is a helper struct used in path finding that couples an
// directional edge with the node's ID in the opposite direction.
type edgeWithPrev struct {
	edge     *channeldb.ChannelEdge
	prevNode *btcec.PublicKey
}

// edgeWeight computes the weight of an edge. This value is used when searching
// for the shortest path within the channel graph between two nodes. Currently
// this is just 1 + the cltv delta value required at this hop, this value
// should be tuned with experimental and empirical data.
//
// TODO(roasbeef): compute robust weight metric
func edgeWeight(e *channeldb.ChannelEdge) float64 {
	return float64(1 + e.Expiry)
}

// findRoute attempts to find a path from the source node within the
// ChannelGraph to the target node that's capable of supporting a payment of
// `amt` value. The current approach is used a multiple pass path finding
// algorithm. First we employ a modified version of Dijkstra's algorithm to
// find a potential set of shortest paths, the distance metric is related to
// the time-lock+fee along the route. Once we have a set of candidate routes,
// we calculate the required fee and time lock values running backwards along
// the route. The route that's selected is the one with the lowest total fee.
//
// TODO(roasbeef): make member, add caching
//  * add k-path
func findRoute(graph *channeldb.ChannelGraph, target *btcec.PublicKey,
	amt btcutil.Amount) (*Route, error) {

	// First initialize empty list of all the node that we've yet to
	// visited.
	// TODO(roasbeef): make into incremental fibonacci heap rather than
	// loading all into memory.
	var unvisited []*channeldb.LightningNode

	// For each node/vertex the graph we create an entry in the distance
	// map for the node set with a distance of "infinity".  We also mark
	// add the node to our set of unvisited nodes.
	distance := make(map[vertex]nodeWithDist)
	if err := graph.ForEachNode(func(node *channeldb.LightningNode) error {
		// TODO(roasbeef): with larger graph can just use disk seeks
		// with a visited map
		distance[newVertex(node.PubKey)] = nodeWithDist{
			dist: infinity,
			node: node,
		}

		unvisited = append(unvisited, node)
		return nil
	}); err != nil {
		return nil, err
	}

	// Next we obtain the source node from the graph, and initialize it
	// with a distance of 0. This indicates our starting point in the graph
	// traversal.
	sourceNode, err := graph.SourceNode()
	if err != nil {
		return nil, err
	}
	sourceVertex := newVertex(sourceNode.PubKey)
	distance[sourceVertex] = nodeWithDist{
		dist: 0,
		node: sourceNode,
	}

	// We'll use this map as a series of "previous" hop pointers. So to get
	// to `vertex` we'll take the edge that it's mapped to within `prev`.
	prev := make(map[vertex]edgeWithPrev)

	for len(unvisited) != 0 {
		var bestNode *channeldb.LightningNode
		smallestDist := infinity

		// First we examine our list of unvisited nodes, for the most
		// optimal vertex to examine next.
		for i, node := range unvisited {
			// The "best" node to visit next is node with the
			// smallest distance from the source of all the
			// unvisited nodes.
			v := newVertex(node.PubKey)
			if nodeInfo := distance[v]; nodeInfo.dist < smallestDist {
				smallestDist = nodeInfo.dist
				bestNode = nodeInfo.node

				// Since we're going to visit this node, we can
				// remove it from the set of unvisited nodes.
				copy(unvisited[i:], unvisited[i+1:])
				unvisited[len(unvisited)-1] = nil // Avoid GC leak.
				unvisited = unvisited[:len(unvisited)-1]

				break
			}
		}

		// If we've reached our target, then we're done here and can
		// exit the graph traversal early.
		if bestNode.PubKey.IsEqual(target) {
			break
		}

		// Now that we've found the next potential step to take we'll
		// examine all the outgoing edge (channels) from this node to
		// further our graph traversal.
		pivot := newVertex(bestNode.PubKey)
		err := bestNode.ForEachChannel(nil, func(edge *channeldb.ChannelEdge) error {
			// Compute the tentative distance to this new
			// channel/edge which is the distance to our current
			// pivot node plus the weight of this edge.
			tempDist := distance[pivot].dist + edgeWeight(edge)

			// If this new tentative distance is better than the
			// current best known distance to this node, then we
			// record the new better distance, and also populate
			// our "next hop" map with this edge.
			// TODO(roasbeef): add capacity to relaxation criteria?
			//  * also add min payment?
			v := newVertex(edge.Node.PubKey)
			if tempDist < distance[v].dist {
				// TODO(roasbeef): unconditionally add for all
				// paths
				distance[v] = nodeWithDist{
					dist: tempDist,
					node: edge.Node,
				}
				prev[v] = edgeWithPrev{
					edge:     edge,
					prevNode: bestNode.PubKey,
				}
			}
			return nil
		})
		if err != nil {
			return nil, err
		}
	}

	// If the target node isn't found in the prev hop map, then a path
	// doesn't exist, so we terminate in an error.
	if _, ok := prev[newVertex(target)]; !ok {
		return nil, ErrNoPathFound
	}

	// Otherwise, we construct a new route which calculate the relevant
	// total fees and proper time lock values for each hop.
	targetVerex := newVertex(target)
	return newRoute(amt, sourceVertex, targetVerex, prev)
}
<|endoftext|>


<commit_before>// Copyright 2019 Google Inc.
//
// 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 grid helps to build grid layouts.
package grid

import (
	"fmt"

	"github.com/mum4k/termdash/container"
	"github.com/mum4k/termdash/widgetapi"
)

// Builder builds grid layouts.
type Builder struct {
	elems []Element
}

// New returns a new grid builder.
func New() *Builder {
	return &Builder{}
}

// Add adds the specified elements.
// The subElements can be either a single Widget or any combination of Rows and
// Columns.
// Rows are created using RowHeightPerc() and Columns are created using
// ColWidthPerc().
// Can be called repeatedly, e.g. to add multiple Rows or Columns.
func (b *Builder) Add(subElements ...Element) {
	b.elems = append(b.elems, subElements...)
}

// Build builds the grid layout and returns the corresponding container
// options.
func (b *Builder) Build() ([]container.Option, error) {
	if err := validate(b.elems); err != nil {
		return nil, err
	}
	return build(b.elems, 100, 100), nil
}

// validate recursively validates the elements that were added to the builder.
// Validates the following per each level of Rows or Columns.:
//   The subElements are either exactly one Widget or any number of Rows and
//   Columns.
//   Each individual width or height is in the range 0 < v < 100.
//   The sum of all widths is <= 100.
//   The sum of all heights is <= 100.
func validate(elems []Element) error {
	heightSum := 0
	widthSum := 0
	for _, elem := range elems {
		switch e := elem.(type) {
		case *row:
			if min, max := 0, 100; e.heightPerc <= min || e.heightPerc >= max {
				return fmt.Errorf("invalid row heightPerc(%d), must be a value in the range %d < v < %d", e.heightPerc, min, max)
			}
			heightSum += e.heightPerc
			if err := validate(e.subElem); err != nil {
				return err
			}

		case *col:
			if min, max := 0, 100; e.widthPerc <= min || e.widthPerc >= max {
				return fmt.Errorf("invalid column widthPerc(%d), must be a value in the range %d < v < %d", e.widthPerc, min, max)
			}
			widthSum += e.widthPerc
			if err := validate(e.subElem); err != nil {
				return err
			}

		case *widget:
			if len(elems) > 1 {
				return fmt.Errorf("when adding a widget, it must be the only added element at that level, got: %v", elems)
			}
		}
	}

	if max := 100; heightSum > max || widthSum > max {
		return fmt.Errorf("the sum of all height percentages(%d) and width percentages(%d) at one element level cannot be larger than %d", heightSum, widthSum, max)
	}
	return nil
}

// build recursively builds the container options according to the elements
// that were added to the builder.
// The parentHeightPerc and parentWidthPerc percent indicate the relative size
// of the element we are building now in the parent element. See innerPerc()
// for more details.
func build(elems []Element, parentHeightPerc, parentWidthPerc int) []container.Option {
	if len(elems) == 0 {
		return nil
	}

	elem := elems[0]
	elems = elems[1:]

	switch e := elem.(type) {
	case *row:
		if len(elems) > 0 {
			perc := innerPerc(e.heightPerc, parentHeightPerc)
			childHeightPerc := parentHeightPerc - e.heightPerc
			return []container.Option{
				container.SplitHorizontal(
					container.Top(build(e.subElem, 100, parentWidthPerc)...),
					container.Bottom(build(elems, childHeightPerc, parentWidthPerc)...),
					container.SplitPercent(perc),
				),
			}
		} else {
			return build(e.subElem, 100, parentWidthPerc)
		}

	case *col:
		if len(elems) > 0 {
			perc := innerPerc(e.widthPerc, parentWidthPerc)
			childWidthPerc := parentWidthPerc - e.widthPerc
			return []container.Option{
				container.SplitVertical(
					container.Left(build(e.subElem, parentHeightPerc, 100)...),
					container.Right(build(elems, parentHeightPerc, childWidthPerc)...),
					container.SplitPercent(perc),
				),
			}
		} else {
			return build(e.subElem, parentHeightPerc, 100)
		}

	case *widget:
		opts := e.cOpts
		opts = append(opts, container.PlaceWidget(e.widget))
		return opts
	}
	return nil
}

// innerPerc translates the outer split percentage into the inner one.
// E.g. multiple rows would specify that they want the outer split percentage
// of 25% each, but we are representing them in a tree of containers so the
// inner splits vary:
//     ╭─────────╮
// 25% │   25%   │
//     │╭───────╮│ ---
// 25% ││  33%  ││
//     ││╭─────╮││
// 25% │││ 50% │││
//     ││├─────┤││ 75%
// 25% │││ 50% │││
//     ││╰─────╯││
//     │╰───────╯│
//     ╰─────────╯ ---
//
// Argument outerPerc is the user specified percentage for the split, i.e. the
// 25% in the example above.
// Argument parentPerc is the percentage this container has in the parent, i.e.
// 75% for the first inner container in the example above.
func innerPerc(outerPerc, parentPerc int) int {
	// parentPerc * parentHeightCells = childHeightCells
	// innerPerc * childHeightCells = outerPerc * parentHeightCells
	// innerPerc * parentPerc * parentHeightCells = outerPerc * parentHeightCells
	// innerPerc * parentPerc = outerPerc
	// innerPerc = outerPerc / parentPerc
	return int(float64(outerPerc) / float64(parentPerc) * 100)
}

// Element is an element that can be added to the grid.
type Element interface {
	isElement()
}

// row is a row in the grid.
// row implements Element.
type row struct {
	// heightPerc is the height percentage this row occupies.
	heightPerc int

	// subElem are the sub Rows or Columns or a single widget.
	subElem []Element
}

// isElement implements Element.isElement.
func (row) isElement() {}

// String implements fmt.Stringer.
func (r *row) String() string {
	return fmt.Sprintf("row{height:%d, sub:%v}", r.heightPerc, r.subElem)
}

// col is a column in the grid.
// col implements Element.
type col struct {
	// widthPerc is the width percentage this column occupies.
	widthPerc int

	// subElem are the sub Rows or Columns or a single widget.
	subElem []Element
}

// isElement implements Element.isElement.
func (col) isElement() {}

// String implements fmt.Stringer.
func (c *col) String() string {
	return fmt.Sprintf("col{width:%d, sub:%v}", c.widthPerc, c.subElem)
}

// widget is a widget placed into the grid.
// widget implements Element.
type widget struct {
	// widget is the widget instance.
	widget widgetapi.Widget
	// cOpts are the options for the widget's container.
	cOpts []container.Option
}

// String implements fmt.Stringer.
func (w *widget) String() string {
	return fmt.Sprintf("widget{type:%T}", w.widget)
}

// isElement implements Element.isElement.
func (widget) isElement() {}

// RowHeightPerc creates a row of the specified height.
// The height is supplied as height percentage of the outer container.
// The subElements can be either a single Widget or any combination of Rows and
// Columns.
func RowHeightPerc(heightPerc int, subElements ...Element) Element {
	return &row{
		heightPerc: heightPerc,
		subElem:    subElements,
	}
}

// ColWidthPerc creates a column of the specified width.
// The width is supplied as width percentage of the outer container.
// The subElements can be either a single Widget or any combination of Rows and
// Columns.
func ColWidthPerc(widthPerc int, subElements ...Element) Element {
	return &col{
		widthPerc: widthPerc,
		subElem:   subElements,
	}
}

// Widget adds a widget into the Row or Column.
// The options will be applied to the container that directly holds this
// widget.
func Widget(w widgetapi.Widget, cOpts ...container.Option) Element {
	return &widget{
		widget: w,
		cOpts:  cOpts,
	}
}
<commit_msg>Fixing lint issues.<commit_after>// Package grid helps to build grid layouts.
package grid

import (
	"fmt"

	"github.com/mum4k/termdash/container"
	"github.com/mum4k/termdash/widgetapi"
)

// Builder builds grid layouts.
type Builder struct {
	elems []Element
}

// New returns a new grid builder.
func New() *Builder {
	return &Builder{}
}

// Add adds the specified elements.
// The subElements can be either a single Widget or any combination of Rows and
// Columns.
// Rows are created using RowHeightPerc() and Columns are created using
// ColWidthPerc().
// Can be called repeatedly, e.g. to add multiple Rows or Columns.
func (b *Builder) Add(subElements ...Element) {
	b.elems = append(b.elems, subElements...)
}

// Build builds the grid layout and returns the corresponding container
// options.
func (b *Builder) Build() ([]container.Option, error) {
	if err := validate(b.elems); err != nil {
		return nil, err
	}
	return build(b.elems, 100, 100), nil
}

// validate recursively validates the elements that were added to the builder.
// Validates the following per each level of Rows or Columns.:
//   The subElements are either exactly one Widget or any number of Rows and
//   Columns.
//   Each individual width or height is in the range 0 < v < 100.
//   The sum of all widths is <= 100.
//   The sum of all heights is <= 100.
func validate(elems []Element) error {
	heightSum := 0
	widthSum := 0
	for _, elem := range elems {
		switch e := elem.(type) {
		case *row:
			if min, max := 0, 100; e.heightPerc <= min || e.heightPerc >= max {
				return fmt.Errorf("invalid row heightPerc(%d), must be a value in the range %d < v < %d", e.heightPerc, min, max)
			}
			heightSum += e.heightPerc
			if err := validate(e.subElem); err != nil {
				return err
			}

		case *col:
			if min, max := 0, 100; e.widthPerc <= min || e.widthPerc >= max {
				return fmt.Errorf("invalid column widthPerc(%d), must be a value in the range %d < v < %d", e.widthPerc, min, max)
			}
			widthSum += e.widthPerc
			if err := validate(e.subElem); err != nil {
				return err
			}

		case *widget:
			if len(elems) > 1 {
				return fmt.Errorf("when adding a widget, it must be the only added element at that level, got: %v", elems)
			}
		}
	}

	if max := 100; heightSum > max || widthSum > max {
		return fmt.Errorf("the sum of all height percentages(%d) and width percentages(%d) at one element level cannot be larger than %d", heightSum, widthSum, max)
	}
	return nil
}

// build recursively builds the container options according to the elements
// that were added to the builder.
// The parentHeightPerc and parentWidthPerc percent indicate the relative size
// of the element we are building now in the parent element. See innerPerc()
// for more details.
func build(elems []Element, parentHeightPerc, parentWidthPerc int) []container.Option {
	if len(elems) == 0 {
		return nil
	}

	elem := elems[0]
	elems = elems[1:]

	switch e := elem.(type) {
	case *row:
		if len(elems) > 0 {
			perc := innerPerc(e.heightPerc, parentHeightPerc)
			childHeightPerc := parentHeightPerc - e.heightPerc
			return []container.Option{
				container.SplitHorizontal(
					container.Top(build(e.subElem, 100, parentWidthPerc)...),
					container.Bottom(build(elems, childHeightPerc, parentWidthPerc)...),
					container.SplitPercent(perc),
				),
			}
		}
		return build(e.subElem, 100, parentWidthPerc)

	case *col:
		if len(elems) > 0 {
			perc := innerPerc(e.widthPerc, parentWidthPerc)
			childWidthPerc := parentWidthPerc - e.widthPerc
			return []container.Option{
				container.SplitVertical(
					container.Left(build(e.subElem, parentHeightPerc, 100)...),
					container.Right(build(elems, parentHeightPerc, childWidthPerc)...),
					container.SplitPercent(perc),
				),
			}
		}
		return build(e.subElem, parentHeightPerc, 100)

	case *widget:
		opts := e.cOpts
		opts = append(opts, container.PlaceWidget(e.widget))
		return opts
	}
	return nil
}

// innerPerc translates the outer split percentage into the inner one.
// E.g. multiple rows would specify that they want the outer split percentage
// of 25% each, but we are representing them in a tree of containers so the
// inner splits vary:
//     ╭─────────╮
// 25% │   25%   │
//     │╭───────╮│ ---
// 25% ││  33%  ││
//     ││╭─────╮││
// 25% │││ 50% │││
//     ││├─────┤││ 75%
// 25% │││ 50% │││
//     ││╰─────╯││
//     │╰───────╯│
//     ╰─────────╯ ---
//
// Argument outerPerc is the user specified percentage for the split, i.e. the
// 25% in the example above.
// Argument parentPerc is the percentage this container has in the parent, i.e.
// 75% for the first inner container in the example above.
func innerPerc(outerPerc, parentPerc int) int {
	// parentPerc * parentHeightCells = childHeightCells
	// innerPerc * childHeightCells = outerPerc * parentHeightCells
	// innerPerc * parentPerc * parentHeightCells = outerPerc * parentHeightCells
	// innerPerc * parentPerc = outerPerc
	// innerPerc = outerPerc / parentPerc
	return int(float64(outerPerc) / float64(parentPerc) * 100)
}

// Element is an element that can be added to the grid.
type Element interface {
	isElement()
}

// row is a row in the grid.
// row implements Element.
type row struct {
	// heightPerc is the height percentage this row occupies.
	heightPerc int

	// subElem are the sub Rows or Columns or a single widget.
	subElem []Element
}

// isElement implements Element.isElement.
func (row) isElement() {}

// String implements fmt.Stringer.
func (r *row) String() string {
	return fmt.Sprintf("row{height:%d, sub:%v}", r.heightPerc, r.subElem)
}

// col is a column in the grid.
// col implements Element.
type col struct {
	// widthPerc is the width percentage this column occupies.
	widthPerc int

	// subElem are the sub Rows or Columns or a single widget.
	subElem []Element
}

// isElement implements Element.isElement.
func (col) isElement() {}

// String implements fmt.Stringer.
func (c *col) String() string {
	return fmt.Sprintf("col{width:%d, sub:%v}", c.widthPerc, c.subElem)
}

// widget is a widget placed into the grid.
// widget implements Element.
type widget struct {
	// widget is the widget instance.
	widget widgetapi.Widget
	// cOpts are the options for the widget's container.
	cOpts []container.Option
}

// String implements fmt.Stringer.
func (w *widget) String() string {
	return fmt.Sprintf("widget{type:%T}", w.widget)
}

// isElement implements Element.isElement.
func (widget) isElement() {}

// RowHeightPerc creates a row of the specified height.
// The height is supplied as height percentage of the outer container.
// The subElements can be either a single Widget or any combination of Rows and
// Columns.
func RowHeightPerc(heightPerc int, subElements ...Element) Element {
	return &row{
		heightPerc: heightPerc,
		subElem:    subElements,
	}
}

// ColWidthPerc creates a column of the specified width.
// The width is supplied as width percentage of the outer container.
// The subElements can be either a single Widget or any combination of Rows and
// Columns.
func ColWidthPerc(widthPerc int, subElements ...Element) Element {
	return &col{
		widthPerc: widthPerc,
		subElem:   subElements,
	}
}

// Widget adds a widget into the Row or Column.
// The options will be applied to the container that directly holds this
// widget.
func Widget(w widgetapi.Widget, cOpts ...container.Option) Element {
	return &widget{
		widget: w,
		cOpts:  cOpts,
	}
}
<|endoftext|>


<commit_before>// +build freebsd
// tun_freebsd.go -- tun interface with cgo for linux / bsd
//

package samtun

/*

#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <net/if.h>
#include <net/if_tun.h>
#include <stdio.h>

int tundev_open(char * ifname) {
  if (strlen(ifname) > IFNAMSIZ) {
    return -1;
  }
  char name[IFNAMSIZ];
  sprintf(name, "/dev/%s", ifname);
  int fd = open(name, O_RDWR);
  if (fd > 0) {
    int i = 0;
    ioctl(fd, TUNSLMODE, &i);
    ioctl(fd, TUNSIFHEAD, &i);
  }
  return fd;
}

int tundev_up(char * ifname, char * addr, char * netmask, int mtu) {

  struct ifreq ifr;
  memset(&ifr, 0, sizeof(struct ifreq));
  strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
  int fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
  if ( fd > 0 ) {
    ifr.ifr_mtu = mtu;
    if ( ioctl(fd, SIOCSIFMTU, (void*) &ifr) < 0) {
      close(fd);
      perror("SIOCSIFMTU");
      return -1;
    }

    struct sockaddr_in src;
    memset(&src, 0, sizeof(struct sockaddr_in));
    src.sin_family = AF_INET;
    if ( ! inet_aton(addr, &src.sin_addr) ) {
      printf("invalid srcaddr %s\n", addr);
      close(fd);
      return -1;
    }

    memcpy(&ifr.ifr_addr, &src, sizeof(struct sockaddr_in));
    if ( ioctl(fd, SIOCSIFADDR, (void*)&ifr) < 0 ) {
      close(fd);
      perror("SIOCSIFADDR");
     return -1;
    }

    struct sockaddr_in mask;
    memset(&mask, 0, sizeof(struct sockaddr_in));
    mask.sin_family = AF_INET;
    if ( ! inet_aton(netmask, &mask.sin_addr) ) {
      close(fd);
      printf("invalid netmask %s\n", netmask);
      return -1;
    }

    if ( ioctl(fd, SIOCGIFFLAGS, (void*)&ifr) < 0 ) {
      close(fd);
      perror("SIOCGIFFLAGS");
      return -1;
    }
    ifr.ifr_flags |= IFF_UP ;
    if ( ioctl(fd, SIOCSIFFLAGS, (void*)&ifr) < 0 ) {
      perror("SIOCSIFFLAGS");
      close(fd);
      return -1;
    }

    close(fd);
    return 0;
  }
  return -1;
}

void tundev_close(int fd) {
  close(fd);
}

*/
import "C"

import (
  "errors"
)

type tunDev struct {
  fd C.int
}

func newTun(ifname, addr, dstaddr string, mtu int) (t tunDev, err error) {
  fd := C.tundev_open(C.CString(ifname))

  if fd == -1 {
    err = errors.New("cannot open tun interface")
  } else {
    if C.tundev_up(C.CString(ifname), C.CString(addr), C.CString(dstaddr), C.int(mtu)) < C.int(0) {
      err = errors.New("cannot put up interface")
    } else {
      t = tunDev{fd}
    }
  }
  return
}

// read from the tun device
func (t *tunDev) Read(d []byte) (n int, err error) {
  return fdRead(C.int(t.fd), d)
}

func (t *tunDev) Write(d []byte) (n int, err error) {
  return fdWrite(C.int(t.fd), d)
}

func (t *tunDev) Close() {
  C.tundev_close(C.int(t.fd))
}
<commit_msg>try fixing freebsd<commit_after>// +build freebsd
// tun_freebsd.go -- tun interface with cgo for linux / bsd
//

package samtun

/*

#include <string.h>
#include <unistd.h>
#include <fcntl.h>
#include <netinet/in.h>
#include <netinet/ip.h>
#include <arpa/inet.h>
#include <sys/ioctl.h>
#include <sys/socket.h>
#include <sys/types.h>
#include <net/if.h>
#include <net/if_tun.h>
#include <stdio.h>

int tundev_open(char * ifname) {
  if (strlen(ifname) > IFNAMSIZ) {
    return -1;
  }
  char name[IFNAMSIZ];
  sprintf(name, "/dev/%s", ifname);
  int fd = open(name, O_RDWR);
  if (fd > 0) {
    int i = 0;
    ioctl(fd, TUNSLMODE, &i);
    ioctl(fd, TUNSIFHEAD, &i);
  }
  return fd;
}

int tundev_up(char * ifname, char * addr, char * netmask, int mtu) {

  struct ifreq ifr;
  memset(&ifr, 0, sizeof(struct ifreq));
  strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
  int fd = socket(AF_INET, SOCK_DGRAM, IPPROTO_IP);
  if ( fd > 0 ) {
    ifr.ifr_mtu = mtu;
    if ( ioctl(fd, SIOCSIFMTU, (void*) &ifr) < 0) {
      close(fd);
      perror("SIOCSIFMTU");
      return -1;
    }

    struct sockaddr_in src;
    memset(&src, 0, sizeof(struct sockaddr_in));
    src.sin_family = AF_INET;
    if ( ! inet_aton(addr, &src.sin_addr) ) {
      printf("invalid srcaddr %s\n", addr);
      close(fd);
      return -1;
    }

    memset(&ifr, 0, sizeof(struct ifreq));
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
    memcpy(&ifr.ifr_addr, &src, sizeof(struct sockaddr_in));
    if ( ioctl(fd, SIOCSIFADDR, (void*)&ifr) < 0 ) {
      close(fd);
      perror("SIOCSIFADDR");
     return -1;
    }

    memset(&ifr, 0, sizeof(struct ifreq));
    strncpy(ifr.ifr_name, ifname, IFNAMSIZ);
    if ( ioctl(fd, SIOCGIFFLAGS, (void*)&ifr) < 0 ) {
      close(fd);
      perror("SIOCGIFFLAGS");
      return -1;
    }
    ifr.ifr_flags |= IFF_UP ;
    if ( ioctl(fd, SIOCSIFFLAGS, (void*)&ifr) < 0 ) {
      perror("SIOCSIFFLAGS");
      close(fd);
      return -1;
    }

    close(fd);
    return 0;
  }
  return -1;
}

void tundev_close(int fd) {
  close(fd);
}

*/
import "C"

import (
  "errors"
)

type tunDev struct {
  fd C.int
}

func newTun(ifname, addr, dstaddr string, mtu int) (t tunDev, err error) {
  fd := C.tundev_open(C.CString(ifname))

  if fd == -1 {
    err = errors.New("cannot open tun interface")
  } else {
    if C.tundev_up(C.CString(ifname), C.CString(addr), C.CString(dstaddr), C.int(mtu)) < C.int(0) {
      err = errors.New("cannot put up interface")
    } else {
      t = tunDev{fd}
    }
  }
  return
}

// read from the tun device
func (t *tunDev) Read(d []byte) (n int, err error) {
  return fdRead(C.int(t.fd), d)
}

func (t *tunDev) Write(d []byte) (n int, err error) {
  return fdWrite(C.int(t.fd), d)
}

func (t *tunDev) Close() {
  C.tundev_close(C.int(t.fd))
}
<|endoftext|>


<commit_before>package pgstorage

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"strings"

	"github.com/jacksontj/dataman/datamantype"
	"github.com/jacksontj/dataman/record"
	"github.com/jacksontj/dataman/stream"
	"github.com/jacksontj/dataman/stream/local"
)

func DoQuery(ctx context.Context, db *sql.DB, query string, colAddrs []ColAddr, args ...interface{}) ([]record.Record, error) {
	rows, err := db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("Error running query: Err=%v query=%s ", err, query)
	}

	results := make([]record.Record, 0)

	// Get the list of column names
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	// If there aren't any rows, we return a nil result
	for rows.Next() {
		columns := make([]interface{}, len(cols))
		columnPointers := make([]interface{}, len(cols))
		for i := range columns {
			columnPointers[i] = &columns[i]
		}

		// Scan the result into the column pointers...
		if err := rows.Scan(columnPointers...); err != nil {
			rows.Close()
			return nil, err
		}

		// Create our map, and retrieve the value for each column from the pointers slice,
		// storing it in the map with the name of the column as the key.
		data := make(record.Record)
		skipN := 0
		for i, colName := range cols {
			val := columnPointers[i].(*interface{})
			if colAddrs != nil {
				if colAddrs[i].skipN > 0 {
					if *val != true {
						skipN = colAddrs[i].skipN
					} else {
						skipN = 0
					}
				} else {
					if skipN <= 0 {
						data.Set(colAddrs[i].key, *val)
					} else {
						skipN--
					}
				}
			} else {
				data[colName] = *val
			}
		}
		results = append(results, data)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	return results, nil
}

func DoStreamQuery(ctx context.Context, db *sql.DB, query string, colAddrs []ColAddr, args ...interface{}) (stream.ClientStream, error) {
	rows, err := db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("Error running query: Err=%v query=%s ", err, query)
	}

	resultsChan := make(chan stream.Result, 1)
	errorChan := make(chan error, 1)

	serverStream := local.NewServerStream(ctx, resultsChan, errorChan)
	clientStream := local.NewClientStream(ctx, resultsChan, errorChan)

	// TODO: without goroutine?
	go func() {
		defer serverStream.Close()
		// Get the list of column names
		cols, err := rows.Columns()
		if err != nil {
			serverStream.SendError(err)
			return
		}
		// If there aren't any rows, we return a nil result
		for rows.Next() {
			columns := make([]interface{}, len(cols))
			columnPointers := make([]interface{}, len(cols))
			for i := range columns {
				columnPointers[i] = &columns[i]
			}

			// Scan the result into the column pointers...
			if err := rows.Scan(columnPointers...); err != nil {
				rows.Close()
				serverStream.SendError(err)
				return
			}

			// Create our map, and retrieve the value for each column from the pointers slice,
			// storing it in the map with the name of the column as the key.
			data := make(record.Record)
			skipN := 0
			for i, colName := range cols {
				val := columnPointers[i].(*interface{})
				if colAddrs != nil {
					if colAddrs[i].skipN > 0 {
						// if we didn't find the key in the selector, then we skipN
						// this accounts for nil and false return types
						if *val != true {
							skipN = colAddrs[i].skipN
						} else {
							skipN = 0
						}
					} else {
						if skipN <= 0 {
							data.Set(colAddrs[i].key, *val)
						} else {
							skipN--
						}
					}
				} else {
					data[colName] = *val
				}
			}
			serverStream.SendResult(data)
		}

		if err := rows.Err(); err != nil {
			serverStream.SendError(err)
			return
		}
	}()

	return clientStream, nil
}

// Normalize field names. This takes a string such as "(data ->> 'created'::text)"
// and converts it to "data.created"
func normalizeFieldName(in string) string {
	if in[0] != '(' || in[len(in)-1] != ')' {
		return in
	}
	in = in[1 : len(in)-1]

	var output string

	for _, part := range strings.Split(in, " ") {
		if sepIdx := strings.Index(part, "'::"); sepIdx > -1 {
			part = part[1:sepIdx]
		}
		if part == "->>" {
			output += "."
		} else {
			output += part
		}
	}

	return output
}

// TODO: remove?
func serializeValue(v interface{}) (string, error) {
	switch vTyped := v.(type) {
	case time.Time:
		return fmt.Sprintf("'%v'", vTyped.Format(datamantype.DateTimeFormatStr)), nil
	default:
		return fmt.Sprintf("'%v'", v), nil
	}
}

// Take a path to an object and convert it to postgres json addressing
func collectionFieldToSelector(path []string) string {
	switch len(path) {
	case 1:
		return path[0]
	case 2:
		return path[0] + "->>'" + path[1] + "'"
	default:
		fieldChain := path[1:]
		return path[0] + "->'" + strings.Join(fieldChain[:len(fieldChain)-1], "'->'") + "'->>'" + path[len(path)-1] + "'"
	}
}

// TODO: remove? or consolidate?
// When we want to do existence checks ( top->'level'->'key' ? 'subkey' we can't use the
// ->> selector since it will return "text" (seemingly the actual value) whereas -> returns
// a map-like object with which we can do selection and ? checks on.
func collectionFieldParentToSelector(path []string) string {
	switch len(path) {
	case 1:
		return path[0]
	case 2:
		return path[0] + "->'" + path[1] + "'"
	default:
		fieldChain := path[1:]
		return path[0] + "->'" + strings.Join(fieldChain[:len(fieldChain)-1], "'->'") + "'->'" + path[len(path)-1] + "'"
	}
}

// ColAddr is a list of addresses of columns
type ColAddr struct {
	key []string
	// Number of columns this is a "selector" for. This is used for jsonb columns
	// so we can differentiate between nil meaning the value in the json is null
	// and the field not existing in the JSON
	// is this a `?` selector telling us whether or not to skip the next one
	skipN int
}

// selectFields returns a SELECT string and the corresponding ColAddr
func selectFields(fields []string) (string, []ColAddr) {
	// TODO: remove?
	// If no projection, then just return all
	if fields == nil {
		return "*", nil
	}

	fieldSelectors := make([]string, 0, len(fields))
	cAddrs := make([]ColAddr, 0, len(fields))
	for _, field := range fields {
		fieldParts := strings.Split(field, ".")
		if len(fieldParts) > 1 {
			cAddrs = append(cAddrs, ColAddr{skipN: 1})
			fieldSelectors = append(fieldSelectors, collectionFieldParentToSelector(fieldParts[:len(fieldParts)-1])+" ? '"+fieldParts[len(fieldParts)-1]+"'")
		}
		cAddrs = append(cAddrs, ColAddr{
			key: fieldParts,
		})
		fieldSelectors = append(fieldSelectors, collectionFieldToSelector(fieldParts))

	}

	return strings.Join(fieldSelectors, ","), cAddrs
}
<commit_msg>Increase default channel size<commit_after>package pgstorage

import (
	"context"
	"database/sql"
	"fmt"
	"time"

	"strings"

	"github.com/jacksontj/dataman/datamantype"
	"github.com/jacksontj/dataman/record"
	"github.com/jacksontj/dataman/stream"
	"github.com/jacksontj/dataman/stream/local"
)

func DoQuery(ctx context.Context, db *sql.DB, query string, colAddrs []ColAddr, args ...interface{}) ([]record.Record, error) {
	rows, err := db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("Error running query: Err=%v query=%s ", err, query)
	}

	results := make([]record.Record, 0)

	// Get the list of column names
	cols, err := rows.Columns()
	if err != nil {
		return nil, err
	}
	// If there aren't any rows, we return a nil result
	for rows.Next() {
		columns := make([]interface{}, len(cols))
		columnPointers := make([]interface{}, len(cols))
		for i := range columns {
			columnPointers[i] = &columns[i]
		}

		// Scan the result into the column pointers...
		if err := rows.Scan(columnPointers...); err != nil {
			rows.Close()
			return nil, err
		}

		// Create our map, and retrieve the value for each column from the pointers slice,
		// storing it in the map with the name of the column as the key.
		data := make(record.Record)
		skipN := 0
		for i, colName := range cols {
			val := columnPointers[i].(*interface{})
			if colAddrs != nil {
				if colAddrs[i].skipN > 0 {
					if *val != true {
						skipN = colAddrs[i].skipN
					} else {
						skipN = 0
					}
				} else {
					if skipN <= 0 {
						data.Set(colAddrs[i].key, *val)
					} else {
						skipN--
					}
				}
			} else {
				data[colName] = *val
			}
		}
		results = append(results, data)
	}
	if err := rows.Err(); err != nil {
		return nil, err
	}

	return results, nil
}

func DoStreamQuery(ctx context.Context, db *sql.DB, query string, colAddrs []ColAddr, args ...interface{}) (stream.ClientStream, error) {
	rows, err := db.QueryContext(ctx, query, args...)
	if err != nil {
		return nil, fmt.Errorf("Error running query: Err=%v query=%s ", err, query)
	}

	resultsChan := make(chan stream.Result, 100)
	errorChan := make(chan error, 1)

	serverStream := local.NewServerStream(ctx, resultsChan, errorChan)
	clientStream := local.NewClientStream(ctx, resultsChan, errorChan)

	// TODO: without goroutine?
	go func() {
		defer serverStream.Close()
		// Get the list of column names
		cols, err := rows.Columns()
		if err != nil {
			serverStream.SendError(err)
			return
		}
		// If there aren't any rows, we return a nil result
		for rows.Next() {
			columns := make([]interface{}, len(cols))
			columnPointers := make([]interface{}, len(cols))
			for i := range columns {
				columnPointers[i] = &columns[i]
			}

			// Scan the result into the column pointers...
			if err := rows.Scan(columnPointers...); err != nil {
				rows.Close()
				serverStream.SendError(err)
				return
			}

			// Create our map, and retrieve the value for each column from the pointers slice,
			// storing it in the map with the name of the column as the key.
			data := make(record.Record)
			skipN := 0
			for i, colName := range cols {
				val := columnPointers[i].(*interface{})
				if colAddrs != nil {
					if colAddrs[i].skipN > 0 {
						// if we didn't find the key in the selector, then we skipN
						// this accounts for nil and false return types
						if *val != true {
							skipN = colAddrs[i].skipN
						} else {
							skipN = 0
						}
					} else {
						if skipN <= 0 {
							data.Set(colAddrs[i].key, *val)
						} else {
							skipN--
						}
					}
				} else {
					data[colName] = *val
				}
			}
			serverStream.SendResult(data)
		}

		if err := rows.Err(); err != nil {
			serverStream.SendError(err)
			return
		}
	}()

	return clientStream, nil
}

// Normalize field names. This takes a string such as "(data ->> 'created'::text)"
// and converts it to "data.created"
func normalizeFieldName(in string) string {
	if in[0] != '(' || in[len(in)-1] != ')' {
		return in
	}
	in = in[1 : len(in)-1]

	var output string

	for _, part := range strings.Split(in, " ") {
		if sepIdx := strings.Index(part, "'::"); sepIdx > -1 {
			part = part[1:sepIdx]
		}
		if part == "->>" {
			output += "."
		} else {
			output += part
		}
	}

	return output
}

// TODO: remove?
func serializeValue(v interface{}) (string, error) {
	switch vTyped := v.(type) {
	case time.Time:
		return fmt.Sprintf("'%v'", vTyped.Format(datamantype.DateTimeFormatStr)), nil
	default:
		return fmt.Sprintf("'%v'", v), nil
	}
}

// Take a path to an object and convert it to postgres json addressing
func collectionFieldToSelector(path []string) string {
	switch len(path) {
	case 1:
		return path[0]
	case 2:
		return path[0] + "->>'" + path[1] + "'"
	default:
		fieldChain := path[1:]
		return path[0] + "->'" + strings.Join(fieldChain[:len(fieldChain)-1], "'->'") + "'->>'" + path[len(path)-1] + "'"
	}
}

// TODO: remove? or consolidate?
// When we want to do existence checks ( top->'level'->'key' ? 'subkey' we can't use the
// ->> selector since it will return "text" (seemingly the actual value) whereas -> returns
// a map-like object with which we can do selection and ? checks on.
func collectionFieldParentToSelector(path []string) string {
	switch len(path) {
	case 1:
		return path[0]
	case 2:
		return path[0] + "->'" + path[1] + "'"
	default:
		fieldChain := path[1:]
		return path[0] + "->'" + strings.Join(fieldChain[:len(fieldChain)-1], "'->'") + "'->'" + path[len(path)-1] + "'"
	}
}

// ColAddr is a list of addresses of columns
type ColAddr struct {
	key []string
	// Number of columns this is a "selector" for. This is used for jsonb columns
	// so we can differentiate between nil meaning the value in the json is null
	// and the field not existing in the JSON
	// is this a `?` selector telling us whether or not to skip the next one
	skipN int
}

// selectFields returns a SELECT string and the corresponding ColAddr
func selectFields(fields []string) (string, []ColAddr) {
	// TODO: remove?
	// If no projection, then just return all
	if fields == nil {
		return "*", nil
	}

	fieldSelectors := make([]string, 0, len(fields))
	cAddrs := make([]ColAddr, 0, len(fields))
	for _, field := range fields {
		fieldParts := strings.Split(field, ".")
		if len(fieldParts) > 1 {
			cAddrs = append(cAddrs, ColAddr{skipN: 1})
			fieldSelectors = append(fieldSelectors, collectionFieldParentToSelector(fieldParts[:len(fieldParts)-1])+" ? '"+fieldParts[len(fieldParts)-1]+"'")
		}
		cAddrs = append(cAddrs, ColAddr{
			key: fieldParts,
		})
		fieldSelectors = append(fieldSelectors, collectionFieldToSelector(fieldParts))

	}

	return strings.Join(fieldSelectors, ","), cAddrs
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"

	"github.com/pmezard/adblock/adblock"
)

func check() error {
	verbose := flag.Bool("v", false, "print rejected rules")
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		return fmt.Errorf("one input rule file expected")
	}
	fp, err := os.Open(args[0])
	if err != nil {
		return err
	}
	defer fp.Close()

	ok := true
	rules := adblock.NewMatcher()
	scanner := bufio.NewScanner(fp)
	for scanner.Scan() {
		rule, err := adblock.ParseRule(scanner.Text())
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: could not parse rule:\n  %s\n  %s\n",
				scanner.Text(), err)
			ok = false
			continue
		}
		if rule == nil {
			continue
		}
		err = rules.AddRule(rule, 0)
		if *verbose && err != nil {
			fmt.Fprintf(os.Stderr, "error: could not add rule:\n  %s\n  %s\n",
				scanner.Text(), err)
			ok = false
		}
	}
	if !ok {
		return fmt.Errorf("some rules could not be parsed")
	}
	fmt.Printf("%s\n", rules)
	return nil
}

func main() {
	err := check()
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(1)
	}
}
<commit_msg>abpcheck: document the utility<commit_after>/*
abpcheck is a command line utility to verify AdBlockPlus rules. Running:

	$ abpcheck easylist.txt

prints any rules which cannot be parsed by adblock package.
*/
package main

import (
	"bufio"
	"flag"
	"fmt"
	"os"

	"github.com/pmezard/adblock/adblock"
)

func check() error {
	verbose := flag.Bool("v", false, "print rejected rules")
	flag.Parse()
	args := flag.Args()
	if len(args) != 1 {
		return fmt.Errorf("one input rule file expected")
	}
	fp, err := os.Open(args[0])
	if err != nil {
		return err
	}
	defer fp.Close()

	ok := true
	rules := adblock.NewMatcher()
	scanner := bufio.NewScanner(fp)
	for scanner.Scan() {
		rule, err := adblock.ParseRule(scanner.Text())
		if err != nil {
			fmt.Fprintf(os.Stderr, "error: could not parse rule:\n  %s\n  %s\n",
				scanner.Text(), err)
			ok = false
			continue
		}
		if rule == nil {
			continue
		}
		err = rules.AddRule(rule, 0)
		if *verbose && err != nil {
			fmt.Fprintf(os.Stderr, "error: could not add rule:\n  %s\n  %s\n",
				scanner.Text(), err)
			ok = false
		}
	}
	if !ok {
		return fmt.Errorf("some rules could not be parsed")
	}
	fmt.Printf("%s\n", rules)
	return nil
}

func main() {
	err := check()
	if err != nil {
		fmt.Fprintf(os.Stderr, "error: %s\n", err)
		os.Exit(1)
	}
}
<|endoftext|>