<commit_before>// +build go1.8

package sqlx

import (
	"context"
	"database/sql"
)

// A union interface of contextPreparer and binder, required to be able to
// prepare named statements with context (as the bindtype must be determined).
type namedPreparerContext interface {
	PreparerContext
	binder
}

func prepareNamedContext(ctx context.Context, p namedPreparerContext, query string) (*NamedStmt, error) {
	bindType := BindType(p.DriverName())
	q, args, err := compileNamedQuery([]byte(query), bindType)
	if err != nil {
		return nil, err
	}
	stmt, err := PreparexContext(ctx, p, q)
	if err != nil {
		return nil, err
	}
	return &NamedStmt{
		QueryString: q,
		Params:      args,
		Stmt:        stmt,
	}, nil
}

// ExecContext executes a named statement using the struct passed.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) ExecContext(ctx context.Context, arg interface{}) (sql.Result, error) {
	args, err := bindAnyArgs(n.Params, arg, n.Stmt.Mapper)
	if err != nil {
		return *new(sql.Result), err
	}
	return n.Stmt.ExecContext(ctx, args...)
}

// QueryContext executes a named statement using the struct argument, returning rows.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryContext(ctx context.Context, arg interface{}) (*sql.Rows, error) {
	args, err := bindAnyArgs(n.Params, arg, n.Stmt.Mapper)
	if err != nil {
		return nil, err
	}
	return n.Stmt.QueryContext(ctx, args...)
}

// QueryRowContext executes a named statement against the database.  Because sqlx cannot
// create a *sql.Row with an error condition pre-set for binding errors, sqlx
// returns a *sqlx.Row instead.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryRowContext(ctx context.Context, arg interface{}) *Row {
	args, err := bindAnyArgs(n.Params, arg, n.Stmt.Mapper)
	if err != nil {
		return &Row{err: err}
	}
	return n.Stmt.QueryRowxContext(ctx, args...)
}

// MustExecContext execs a NamedStmt, panicing on error
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) MustExecContext(ctx context.Context, arg interface{}) sql.Result {
	res, err := n.ExecContext(ctx, arg)
	if err != nil {
		panic(err)
	}
	return res
}

// QueryxContext using this NamedStmt
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryxContext(ctx context.Context, arg interface{}) (*Rows, error) {
	r, err := n.QueryContext(ctx, arg)
	if err != nil {
		return nil, err
	}
	return &Rows{Rows: r, Mapper: n.Stmt.Mapper, unsafe: isUnsafe(n)}, err
}

// QueryRowxContext this NamedStmt.  Because of limitations with QueryRow, this is
// an alias for QueryRow.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryRowxContext(ctx context.Context, arg interface{}) *Row {
	return n.QueryRowContext(ctx, arg)
}

// SelectContext using this NamedStmt
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) SelectContext(ctx context.Context, dest interface{}, arg interface{}) error {
	rows, err := n.QueryxContext(ctx, arg)
	if err != nil {
		return err
	}
	// if something happens here, we want to make sure the rows are Closed
	defer rows.Close()
	return scanAll(rows, dest, false)
}

// GetContext using this NamedStmt
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) GetContext(ctx context.Context, dest interface{}, arg interface{}) error {
	r := n.QueryRowxContext(ctx, arg)
	return r.scanAny(dest, false)
}

// NamedQueryContext binds a named query and then runs Query on the result using the
// provided Ext (sqlx.Tx, sqlx.Db).  It works with both structs and with
// map[string]interface{} types.
func NamedQueryContext(ctx context.Context, e ExtContext, query string, arg interface{}) (*Rows, error) {
	q, args, err := bindNamedMapper(BindType(e.DriverName()), query, arg, mapperFor(e))
	if err != nil {
		return nil, err
	}
	return e.QueryxContext(ctx, q, args...)
}

// NamedExecContext uses BindStruct to get a query executable by the driver and
// then runs Exec on the result.  Returns an error from the binding
// or the query excution itself.
func NamedExecContext(ctx context.Context, e ExtContext, query string, arg interface{}) (sql.Result, error) {
	q, args, err := bindNamedMapper(BindType(e.DriverName()), query, arg, mapperFor(e))
	if err != nil {
		return nil, err
	}
	return e.ExecContext(ctx, q, args...)
}
<commit_msg>fix typo<commit_after>// +build go1.8

package sqlx

import (
	"context"
	"database/sql"
)

// A union interface of contextPreparer and binder, required to be able to
// prepare named statements with context (as the bindtype must be determined).
type namedPreparerContext interface {
	PreparerContext
	binder
}

func prepareNamedContext(ctx context.Context, p namedPreparerContext, query string) (*NamedStmt, error) {
	bindType := BindType(p.DriverName())
	q, args, err := compileNamedQuery([]byte(query), bindType)
	if err != nil {
		return nil, err
	}
	stmt, err := PreparexContext(ctx, p, q)
	if err != nil {
		return nil, err
	}
	return &NamedStmt{
		QueryString: q,
		Params:      args,
		Stmt:        stmt,
	}, nil
}

// ExecContext executes a named statement using the struct passed.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) ExecContext(ctx context.Context, arg interface{}) (sql.Result, error) {
	args, err := bindAnyArgs(n.Params, arg, n.Stmt.Mapper)
	if err != nil {
		return *new(sql.Result), err
	}
	return n.Stmt.ExecContext(ctx, args...)
}

// QueryContext executes a named statement using the struct argument, returning rows.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryContext(ctx context.Context, arg interface{}) (*sql.Rows, error) {
	args, err := bindAnyArgs(n.Params, arg, n.Stmt.Mapper)
	if err != nil {
		return nil, err
	}
	return n.Stmt.QueryContext(ctx, args...)
}

// QueryRowContext executes a named statement against the database.  Because sqlx cannot
// create a *sql.Row with an error condition pre-set for binding errors, sqlx
// returns a *sqlx.Row instead.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryRowContext(ctx context.Context, arg interface{}) *Row {
	args, err := bindAnyArgs(n.Params, arg, n.Stmt.Mapper)
	if err != nil {
		return &Row{err: err}
	}
	return n.Stmt.QueryRowxContext(ctx, args...)
}

// MustExecContext execs a NamedStmt, panicing on error
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) MustExecContext(ctx context.Context, arg interface{}) sql.Result {
	res, err := n.ExecContext(ctx, arg)
	if err != nil {
		panic(err)
	}
	return res
}

// QueryxContext using this NamedStmt
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryxContext(ctx context.Context, arg interface{}) (*Rows, error) {
	r, err := n.QueryContext(ctx, arg)
	if err != nil {
		return nil, err
	}
	return &Rows{Rows: r, Mapper: n.Stmt.Mapper, unsafe: isUnsafe(n)}, err
}

// QueryRowxContext this NamedStmt.  Because of limitations with QueryRow, this is
// an alias for QueryRow.
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) QueryRowxContext(ctx context.Context, arg interface{}) *Row {
	return n.QueryRowContext(ctx, arg)
}

// SelectContext using this NamedStmt
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) SelectContext(ctx context.Context, dest interface{}, arg interface{}) error {
	rows, err := n.QueryxContext(ctx, arg)
	if err != nil {
		return err
	}
	// if something happens here, we want to make sure the rows are Closed
	defer rows.Close()
	return scanAll(rows, dest, false)
}

// GetContext using this NamedStmt
// Any named placeholder parameters are replaced with fields from arg.
func (n *NamedStmt) GetContext(ctx context.Context, dest interface{}, arg interface{}) error {
	r := n.QueryRowxContext(ctx, arg)
	return r.scanAny(dest, false)
}

// NamedQueryContext binds a named query and then runs Query on the result using the
// provided Ext (sqlx.Tx, sqlx.Db).  It works with both structs and with
// map[string]interface{} types.
func NamedQueryContext(ctx context.Context, e ExtContext, query string, arg interface{}) (*Rows, error) {
	q, args, err := bindNamedMapper(BindType(e.DriverName()), query, arg, mapperFor(e))
	if err != nil {
		return nil, err
	}
	return e.QueryxContext(ctx, q, args...)
}

// NamedExecContext uses BindStruct to get a query executable by the driver and
// then runs Exec on the result.  Returns an error from the binding
// or the query execution itself.
func NamedExecContext(ctx context.Context, e ExtContext, query string, arg interface{}) (sql.Result, error) {
	q, args, err := bindNamedMapper(BindType(e.DriverName()), query, arg, mapperFor(e))
	if err != nil {
		return nil, err
	}
	return e.ExecContext(ctx, q, args...)
}
<|endoftext|>


<commit_before>// Copyright 2017-2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/filter"
	amassnet "github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/PuerkitoBio/goquery"
	"github.com/caffix/stringset"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

const (
	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	httpTimeout      = 30 * time.Second
	handshakeTimeout = 5 * time.Second
)

var (
  // UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent       string
	subRE           = dns.AnySubdomainRegex()
	crawlRE         = regexp.MustCompile(`\.\w{2,6}($|\?|#)`)
  crawlFileEnds   = []string{"html", "do", "action", "cgi"}
	crawlFileStarts = []string{"js", "htm", "as", "php", "inc"}
  nameStripRE     = regexp.MustCompile(`^u[0-9a-f]{4}|20|22|25|2b|2f|3d|3a|40`)
)

// DefaultClient is the same HTTP client used by the package methods.
var DefaultClient *http.Client

// BasicAuth contains the data used for HTTP basic authentication.
type BasicAuth struct {
	Username string
	Password string
}

func init() {
	jar, _ := cookiejar.New(nil)
	DefaultClient = &http.Client{
		Timeout: httpTimeout,
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           amassnet.DialContext,
			MaxIdleConns:          200,
			MaxConnsPerHost:       50,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   handshakeTimeout,
			ExpectContinueTimeout: 10 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}

	switch runtime.GOOS {
	case "windows":
		UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	case "darwin":
		UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 11_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	default:
		UserAgent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	DefaultClient.Jar.SetCookies(destURL, DefaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range DefaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for the provided URL when successful.
func RequestWebPage(ctx context.Context, u string, body io.Reader, hvals map[string]string, auth *BasicAuth) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequestWithContext(ctx, method, u, body)
	if err != nil {
		return "", err
	}
	if auth != nil && auth.Username != "" && auth.Password != "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)

	for k, v := range hvals {
		req.Header.Set(k, v)
	}

	resp, err := DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		err = errors.New(resp.Status)
	}
	return string(in), err
}

// Crawl will spider the web page at the URL argument looking for DNS names within the scope argument.
func Crawl(ctx context.Context, u string, scope []string, max int, f filter.Filter) ([]string, error) {
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("The context expired")
	default:
	}

	newScope := append([]string{}, scope...)

	target := subRE.FindString(u)
	if target != "" {
		var found bool
		for _, domain := range newScope {
			if target == domain {
				found = true
				break
			}
		}
		if !found {
			newScope = append(newScope, target)
		}
	}

	if f == nil {
		f = filter.NewStringFilter()
	}

	var count int
	var m sync.Mutex
	results := stringset.New()
	g := geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:        newScope,
		StartURLs:             []string{u},
		Timeout:               5 * time.Minute,
		RobotsTxtDisabled:     true,
		UserAgent:             UserAgent,
		LogDisabled:           true,
		ConcurrentRequests:    5,
		RequestDelay:          750 * time.Millisecond,
		RequestDelayRandomize: true,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			resp, err := httputil.DumpResponse(interface{}(r).(*http.Response), true)
			if err == nil {
				for _, n := range subRE.FindAllString(string(resp), -1) {
					if name := CleanName(n); whichDomain(name, scope) != "" {
						m.Lock()
						results.Insert(name)
						m.Unlock()
					}
				}
			}

			processURL := func(u *url.URL) {
				// Attempt to save the name in our results
				m.Lock()
				results.Insert(u.Hostname())
				m.Unlock()

				if s := crawlFilterURLs(u, f); s != "" {
					// Be sure the crawl has not exceeded the maximum links to be followed
					m.Lock()
					count++
					current := count
					m.Unlock()
					if max <= 0 || current < max {
						g.Get(s, g.Opt.ParseFunc)
					}
				}
			}

			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if u, err := r.JoinURL(href); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})

			r.HTMLDoc.Find("script").Each(func(i int, s *goquery.Selection) {
				if src, ok := s.Attr("src"); ok {
					if u, err := r.JoinURL(src); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})
		},
	})
	options := &client.Options{
		MaxBodySize:    100 * 1024 * 1024, // 100MB
		RetryTimes:     2,
		RetryHTTPCodes: []int{408, 500, 502, 503, 504, 522, 524},
	}
	g.Client = client.NewClient(options)
	g.Client.Client = http.DefaultClient

	done := make(chan struct{}, 2)
	go func() {
		g.Start()
		done <- struct{}{}
	}()

	var err error
	select {
	case <-ctx.Done():
		err = fmt.Errorf("The context expired during the crawl of %s", u)
	case <-done:
		if len(results.Slice()) == 0 {
			err = fmt.Errorf("No DNS names were discovered during the crawl of %s", u)
		}
	}

	return results.Slice(), err
}

func crawlFilterURLs(p *url.URL, f filter.Filter) string {
	// Check that the URL has an appropriate scheme for scraping
	if !p.IsAbs() || (p.Scheme != "http" && p.Scheme != "https") {
		return ""
	}
	// If the URL path has a file extension, check that it's of interest
	if ext := crawlRE.FindString(p.Path); ext != "" {
		ext = strings.TrimRight(ext, "?#")

		var found bool
		for _, s := range crawlFileStarts {
			if strings.HasPrefix(ext, "." + s) {
				found = true
				break
			}
		}
		for _, e := range crawlFileEnds {
			if strings.HasSuffix(ext, e) {
				found = true
				break
			}
		}
		if !found {
			return ""
		}
	}
	// Remove fragments and check if we've seen this URL before
	p.Fragment = ""
	p.RawFragment = ""
	if f.Duplicate(p.String()) {
		return ""
	}
	return p.String()
}

func whichDomain(name string, scope []string) string {
	n := strings.TrimSpace(name)

	for _, d := range scope {
		if strings.HasSuffix(n, d) {
			// fork made me do it :>
			nlen := len(n)
			dlen := len(d)
			// Check for exact match first to guard against out of bound index
			if nlen == dlen || n[nlen-dlen-1] == '.' {
				return d
			}
		}
	}
	return ""
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(ctx context.Context, addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		select {
		case <-ctx.Done():
			return names
		default:
		}

		// Set the maximum time allowed for making the connection
		tCtx, cancel := context.WithTimeout(ctx, handshakeTimeout)
		defer cancel()
		// Obtain the connection
		conn, err := amassnet.DialContext(tCtx, "tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		go func() {
			errChan <- c.Handshake()
		}()

		t := time.NewTimer(handshakeTimeout)
		select {
		case <-t.C:
			err = errors.New("Handshake timeout")
		case e := <-errChan:
			err = e
		}
		t.Stop()

		if err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}

	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode(ctx context.Context) string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage(ctx, "https://ipapi.co/json", nil, headers, nil)
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	if err := json.Unmarshal([]byte(page), &ipinfo); err != nil {
		return ""
	}
	return strings.ToLower(ipinfo.CountryCode)
}

// CleanName will clean up the names scraped from the web.
func CleanName(name string) string {
	var err error

	name, err = strconv.Unquote("\"" + strings.TrimSpace(name) + "\"")
	if err == nil {
		name = subRE.FindString(name)
	}

	name = strings.ToLower(name)
	for {
		name = strings.Trim(name, "-.")

		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	return name
}
<commit_msg>Indent fix<commit_after>// Copyright 2017-2021 Jeff Foley. All rights reserved.
// Use of this source code is governed by Apache 2 LICENSE that can be found in the LICENSE file.

package http

import (
	"context"
	"crypto/tls"
	"crypto/x509"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	"net/http/cookiejar"
	"net/http/httputil"
	"net/url"
	"regexp"
	"runtime"
	"strconv"
	"strings"
	"sync"
	"time"

	"github.com/OWASP/Amass/v3/filter"
	amassnet "github.com/OWASP/Amass/v3/net"
	"github.com/OWASP/Amass/v3/net/dns"
	"github.com/PuerkitoBio/goquery"
	"github.com/caffix/stringset"
	"github.com/geziyor/geziyor"
	"github.com/geziyor/geziyor/client"
)

const (
	// Accept is the default HTTP Accept header value used by Amass.
	Accept = "text/html,application/xhtml+xml,application/xml;q=0.9,image/webp,*/*;q=0.8"

	// AcceptLang is the default HTTP Accept-Language header value used by Amass.
	AcceptLang = "en-US,en;q=0.8"

	httpTimeout      = 30 * time.Second
	handshakeTimeout = 5 * time.Second
)

var (
  // UserAgent is the default user agent used by Amass during HTTP requests.
	UserAgent       string
	subRE           = dns.AnySubdomainRegex()
	crawlRE         = regexp.MustCompile(`\.\w{2,6}($|\?|#)`)
	crawlFileEnds   = []string{"html", "do", "action", "cgi"}
	crawlFileStarts = []string{"js", "htm", "as", "php", "inc"}
	nameStripRE     = regexp.MustCompile(`^u[0-9a-f]{4}|20|22|25|2b|2f|3d|3a|40`)
)

// DefaultClient is the same HTTP client used by the package methods.
var DefaultClient *http.Client

// BasicAuth contains the data used for HTTP basic authentication.
type BasicAuth struct {
	Username string
	Password string
}

func init() {
	jar, _ := cookiejar.New(nil)
	DefaultClient = &http.Client{
		Timeout: httpTimeout,
		Transport: &http.Transport{
			Proxy:                 http.ProxyFromEnvironment,
			DialContext:           amassnet.DialContext,
			MaxIdleConns:          200,
			MaxConnsPerHost:       50,
			IdleConnTimeout:       90 * time.Second,
			TLSHandshakeTimeout:   handshakeTimeout,
			ExpectContinueTimeout: 10 * time.Second,
			TLSClientConfig:       &tls.Config{InsecureSkipVerify: true},
		},
		Jar: jar,
	}

	switch runtime.GOOS {
	case "windows":
		UserAgent = "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	case "darwin":
		UserAgent = "Mozilla/5.0 (Macintosh; Intel Mac OS X 11_4) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	default:
		UserAgent = "Mozilla/5.0 (X11; Linux x86_64) AppleWebKit/537.36 (KHTML, like Gecko) Chrome/91.0.4472.164 Safari/537.36"
	}
}

// CopyCookies copies cookies from one domain to another. Some of our data
// sources rely on shared auth tokens and this avoids sending extra requests
// to have the site reissue cookies for the other domains.
func CopyCookies(src string, dest string) {
	srcURL, _ := url.Parse(src)
	destURL, _ := url.Parse(dest)
	DefaultClient.Jar.SetCookies(destURL, DefaultClient.Jar.Cookies(srcURL))
}

// CheckCookie checks if a cookie exists in the cookie jar for a given host
func CheckCookie(urlString string, cookieName string) bool {
	cookieURL, _ := url.Parse(urlString)
	found := false
	for _, cookie := range DefaultClient.Jar.Cookies(cookieURL) {
		if cookie.Name == cookieName {
			found = true
			break
		}
	}
	return found
}

// RequestWebPage returns a string containing the entire response for the provided URL when successful.
func RequestWebPage(ctx context.Context, u string, body io.Reader, hvals map[string]string, auth *BasicAuth) (string, error) {
	method := "GET"
	if body != nil {
		method = "POST"
	}
	req, err := http.NewRequestWithContext(ctx, method, u, body)
	if err != nil {
		return "", err
	}
	if auth != nil && auth.Username != "" && auth.Password != "" {
		req.SetBasicAuth(auth.Username, auth.Password)
	}
	req.Header.Set("User-Agent", UserAgent)
	req.Header.Set("Accept", Accept)
	req.Header.Set("Accept-Language", AcceptLang)

	for k, v := range hvals {
		req.Header.Set(k, v)
	}

	resp, err := DefaultClient.Do(req)
	if err != nil {
		return "", err
	}

	in, err := ioutil.ReadAll(resp.Body)
	resp.Body.Close()

	if resp.StatusCode < 200 || resp.StatusCode >= 400 {
		err = errors.New(resp.Status)
	}
	return string(in), err
}

// Crawl will spider the web page at the URL argument looking for DNS names within the scope argument.
func Crawl(ctx context.Context, u string, scope []string, max int, f filter.Filter) ([]string, error) {
	select {
	case <-ctx.Done():
		return nil, fmt.Errorf("The context expired")
	default:
	}

	newScope := append([]string{}, scope...)

	target := subRE.FindString(u)
	if target != "" {
		var found bool
		for _, domain := range newScope {
			if target == domain {
				found = true
				break
			}
		}
		if !found {
			newScope = append(newScope, target)
		}
	}

	if f == nil {
		f = filter.NewStringFilter()
	}

	var count int
	var m sync.Mutex
	results := stringset.New()
	g := geziyor.NewGeziyor(&geziyor.Options{
		AllowedDomains:        newScope,
		StartURLs:             []string{u},
		Timeout:               5 * time.Minute,
		RobotsTxtDisabled:     true,
		UserAgent:             UserAgent,
		LogDisabled:           true,
		ConcurrentRequests:    5,
		RequestDelay:          750 * time.Millisecond,
		RequestDelayRandomize: true,
		ParseFunc: func(g *geziyor.Geziyor, r *client.Response) {
			resp, err := httputil.DumpResponse(interface{}(r).(*http.Response), true)
			if err == nil {
				for _, n := range subRE.FindAllString(string(resp), -1) {
					if name := CleanName(n); whichDomain(name, scope) != "" {
						m.Lock()
						results.Insert(name)
						m.Unlock()
					}
				}
			}

			processURL := func(u *url.URL) {
				// Attempt to save the name in our results
				m.Lock()
				results.Insert(u.Hostname())
				m.Unlock()

				if s := crawlFilterURLs(u, f); s != "" {
					// Be sure the crawl has not exceeded the maximum links to be followed
					m.Lock()
					count++
					current := count
					m.Unlock()
					if max <= 0 || current < max {
						g.Get(s, g.Opt.ParseFunc)
					}
				}
			}

			r.HTMLDoc.Find("a").Each(func(i int, s *goquery.Selection) {
				if href, ok := s.Attr("href"); ok {
					if u, err := r.JoinURL(href); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})

			r.HTMLDoc.Find("script").Each(func(i int, s *goquery.Selection) {
				if src, ok := s.Attr("src"); ok {
					if u, err := r.JoinURL(src); err == nil && whichDomain(u.Hostname(), newScope) != "" {
						processURL(u)
					}
				}
			})
		},
	})
	options := &client.Options{
		MaxBodySize:    100 * 1024 * 1024, // 100MB
		RetryTimes:     2,
		RetryHTTPCodes: []int{408, 500, 502, 503, 504, 522, 524},
	}
	g.Client = client.NewClient(options)
	g.Client.Client = http.DefaultClient

	done := make(chan struct{}, 2)
	go func() {
		g.Start()
		done <- struct{}{}
	}()

	var err error
	select {
	case <-ctx.Done():
		err = fmt.Errorf("The context expired during the crawl of %s", u)
	case <-done:
		if len(results.Slice()) == 0 {
			err = fmt.Errorf("No DNS names were discovered during the crawl of %s", u)
		}
	}

	return results.Slice(), err
}

func crawlFilterURLs(p *url.URL, f filter.Filter) string {
	// Check that the URL has an appropriate scheme for scraping
	if !p.IsAbs() || (p.Scheme != "http" && p.Scheme != "https") {
		return ""
	}
	// If the URL path has a file extension, check that it's of interest
	if ext := crawlRE.FindString(p.Path); ext != "" {
		ext = strings.TrimRight(ext, "?#")

		var found bool
		for _, s := range crawlFileStarts {
			if strings.HasPrefix(ext, "." + s) {
				found = true
				break
			}
		}
		for _, e := range crawlFileEnds {
			if strings.HasSuffix(ext, e) {
				found = true
				break
			}
		}
		if !found {
			return ""
		}
	}
	// Remove fragments and check if we've seen this URL before
	p.Fragment = ""
	p.RawFragment = ""
	if f.Duplicate(p.String()) {
		return ""
	}
	return p.String()
}

func whichDomain(name string, scope []string) string {
	n := strings.TrimSpace(name)

	for _, d := range scope {
		if strings.HasSuffix(n, d) {
			// fork made me do it :>
			nlen := len(n)
			dlen := len(d)
			// Check for exact match first to guard against out of bound index
			if nlen == dlen || n[nlen-dlen-1] == '.' {
				return d
			}
		}
	}
	return ""
}

// PullCertificateNames attempts to pull a cert from one or more ports on an IP.
func PullCertificateNames(ctx context.Context, addr string, ports []int) []string {
	var names []string

	// Check hosts for certificates that contain subdomain names
	for _, port := range ports {
		select {
		case <-ctx.Done():
			return names
		default:
		}

		// Set the maximum time allowed for making the connection
		tCtx, cancel := context.WithTimeout(ctx, handshakeTimeout)
		defer cancel()
		// Obtain the connection
		conn, err := amassnet.DialContext(tCtx, "tcp", net.JoinHostPort(addr, strconv.Itoa(port)))
		if err != nil {
			continue
		}
		defer conn.Close()

		c := tls.Client(conn, &tls.Config{InsecureSkipVerify: true})
		// Attempt to acquire the certificate chain
		errChan := make(chan error, 2)
		go func() {
			errChan <- c.Handshake()
		}()

		t := time.NewTimer(handshakeTimeout)
		select {
		case <-t.C:
			err = errors.New("Handshake timeout")
		case e := <-errChan:
			err = e
		}
		t.Stop()

		if err != nil {
			continue
		}
		// Get the correct certificate in the chain
		certChain := c.ConnectionState().PeerCertificates
		cert := certChain[0]
		// Create the new requests from names found within the cert
		names = append(names, namesFromCert(cert)...)
	}

	return names
}

func namesFromCert(cert *x509.Certificate) []string {
	var cn string

	for _, name := range cert.Subject.Names {
		oid := name.Type
		if len(oid) == 4 && oid[0] == 2 && oid[1] == 5 && oid[2] == 4 {
			if oid[3] == 3 {
				cn = fmt.Sprintf("%s", name.Value)
				break
			}
		}
	}

	subdomains := stringset.New()
	// Add the subject common name to the list of subdomain names
	commonName := dns.RemoveAsteriskLabel(cn)
	if commonName != "" {
		subdomains.Insert(commonName)
	}
	// Add the cert DNS names to the list of subdomain names
	for _, name := range cert.DNSNames {
		n := dns.RemoveAsteriskLabel(name)
		if n != "" {
			subdomains.Insert(n)
		}
	}
	return subdomains.Slice()
}

// ClientCountryCode returns the country code for the public-facing IP address for the host of the process.
func ClientCountryCode(ctx context.Context) string {
	headers := map[string]string{"Content-Type": "application/json"}

	page, err := RequestWebPage(ctx, "https://ipapi.co/json", nil, headers, nil)
	if err != nil {
		return ""
	}

	// Extract the country code from the REST API results
	var ipinfo struct {
		CountryCode string `json:"country"`
	}

	if err := json.Unmarshal([]byte(page), &ipinfo); err != nil {
		return ""
	}
	return strings.ToLower(ipinfo.CountryCode)
}

// CleanName will clean up the names scraped from the web.
func CleanName(name string) string {
	var err error

	name, err = strconv.Unquote("\"" + strings.TrimSpace(name) + "\"")
	if err == nil {
		name = subRE.FindString(name)
	}

	name = strings.ToLower(name)
	for {
		name = strings.Trim(name, "-.")

		if i := nameStripRE.FindStringIndex(name); i != nil {
			name = name[i[1]:]
		} else {
			break
		}
	}

	return name
}
<|endoftext|>


<commit_before>package node

import (
	"DNA/common/log"
	. "DNA/config"
	. "DNA/net/message"
	. "DNA/net/protocol"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

type link struct {
	//Todo Add lock here
	addr  string    // The address of the node
	conn  net.Conn  // Connect socket with the peer node
	port  uint16    // The server port of the node
	time  time.Time // The latest time the node activity
	rxBuf struct {  // The RX buffer of this node to solve mutliple packets problem
		p   []byte
		len int
	}
	connCnt uint64 // The connection count
}

// Shrinking the buf to the exactly reading in byte length
//@Return @1 the start header of next message, the left length of the next message
func unpackNodeBuf(node *node, buf []byte) {
	var msgLen int
	var msgBuf []byte
	if node.rxBuf.p == nil {
		if len(buf) < MSGHDRLEN {
			log.Warn("Unexpected size of received message")
			errors.New("Unexpected size of received message")
			return
		}
		// FIXME Check the payload < 0 error case
		msgLen = PayloadLen(buf) + MSGHDRLEN
	} else {
		msgLen = node.rxBuf.len
	}

	if len(buf) == msgLen {
		msgBuf = append(node.rxBuf.p, buf[:]...)
		go HandleNodeMsg(node, msgBuf, len(msgBuf))
		node.rxBuf.p = nil
		node.rxBuf.len = 0
	} else if len(buf) < msgLen {
		node.rxBuf.p = append(node.rxBuf.p, buf[:]...)
		node.rxBuf.len = msgLen - len(buf)
	} else {
		msgBuf = append(node.rxBuf.p, buf[0:msgLen]...)
		go HandleNodeMsg(node, msgBuf, len(msgBuf))
		node.rxBuf.p = nil
		node.rxBuf.len = 0

		unpackNodeBuf(node, buf[msgLen:])
	}

	// TODO we need reset the node.rxBuf.p pointer and length if CheckSUM error happened?
}

func (node *node) rx() error {
	conn := node.getConn()
	from := conn.RemoteAddr().String()
	buf := make([]byte, MAXBUFLEN)
	for {
		len, err := conn.Read(buf[0:(MAXBUFLEN - 1)])
		buf[MAXBUFLEN-1] = 0 //Prevent overflow
		switch err {
		case nil:
			unpackNodeBuf(node, buf[0:len])
			//go handleNodeMsg(node, buf, len)
			break
		case io.EOF:
			//log.Debug("Reading EOF of network conn")
			break
		default:
			log.Error("Read connetion error ", err)
			goto disconnect
		}
	}

disconnect:
	err := conn.Close()
	node.SetState(INACTIVITY)
	log.Debug("Close connection ", from)
	return err
}

func printIPAddr() {
	host, _ := os.Hostname()
	addrs, _ := net.LookupIP(host)
	for _, addr := range addrs {
		if ipv4 := addr.To4(); ipv4 != nil {
			log.Info("IPv4: ", ipv4)
		}
	}
}

func (link link) CloseConn() {
	link.conn.Close()
}

// Init the server port, should be run in another thread
func (n *node) initConnection() {
	isTls := Parameters.IsTLS
	var listener net.Listener
	var err error
	if isTls {
		listener, err = initTlsListen()
		if err != nil {
			log.Error("TLS listen failed")
			return
		}
	} else {
		listener, err = initNonTlsListen()
		if err != nil {
			log.Error("non TLS listen failed")
			return
		}
	}
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Error("Error accepting ", err.Error())
			return
		}
		log.Info("Remote node connect with ", conn.RemoteAddr(), conn.LocalAddr())

		n.link.connCnt++

		node := NewNode()
		node.addr, err = parseIPaddr(conn.RemoteAddr().String())
		node.local = n
		node.conn = conn
		go node.rx()
	}
	//TODO When to free the net listen resouce?
}

func initNonTlsListen() (net.Listener, error) {
	log.Debug()
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(Parameters.NodePort))
	if err != nil {
		log.Error("Error listening\n", err.Error())
		return nil, err
	}
	return listener, nil
}

func initTlsListen() (net.Listener, error) {
	CertPath := Parameters.CertPath
	KeyPath := Parameters.KeyPath
	CAPath := Parameters.CAPath

	// load cert
	cert, err := tls.LoadX509KeyPair(CertPath, KeyPath)
	if err != nil {
		log.Error("load keys fail", err)
		return nil, err
	}
	// load root ca
	caData, err := ioutil.ReadFile(CAPath)
	if err != nil {
		log.Error("read ca fail", err)
		return nil, err
	}
	pool := x509.NewCertPool()
	ret := pool.AppendCertsFromPEM(caData)
	if !ret {
		return nil, errors.New("failed to parse root certificate")
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      pool,
		ClientAuth:   tls.RequireAndVerifyClientCert,
		ClientCAs:    pool,
	}

	log.Info("TLS listen port is ", strconv.Itoa(Parameters.NodePort))
	listener, err := tls.Listen("tcp", ":"+strconv.Itoa(Parameters.NodePort), tlsConfig)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	return listener, nil
}

func parseIPaddr(s string) (string, error) {
	i := strings.Index(s, ":")
	if i < 0 {
		log.Warn("Split IP address&port error")
		return s, errors.New("Split IP address&port error")
	}
	return s[:i], nil
}

func (node *node) Connect(nodeAddr string) error {
	log.Debug()
	isTls := Parameters.IsTLS
	var conn net.Conn
	var err error
	if isTls {
		conn, err = TLSDial(nodeAddr)
		if err != nil {
			log.Error("TLS connect failed: ", err)
			return nil
		}
	} else {
		conn, err = NonTLSDial(nodeAddr)
		if err != nil {
			log.Error("non TLS connect failed:", err)
			return nil
		}
	}
	node.link.connCnt++

	n := NewNode()
	n.conn = conn
	n.addr, err = parseIPaddr(conn.RemoteAddr().String())
	n.local = node

	log.Info(fmt.Sprintf("Connect node %s connect with %s with %s",
		conn.LocalAddr().String(), conn.RemoteAddr().String(),
		conn.RemoteAddr().Network()))
	go n.rx()

	n.SetState(HAND)
	buf, _ := NewVersion(node)
	n.Tx(buf)

	return nil
}

func NonTLSDial(nodeAddr string) (net.Conn, error) {
	log.Debug()
	conn, err := net.Dial("tcp", nodeAddr)
	if err != nil {
		log.Error("Error dialing\n", err.Error())
		return nil, err
	}
	return conn, nil
}

func TLSDial(nodeAddr string) (net.Conn, error) {
	CertPath := Parameters.CertPath
	KeyPath := Parameters.KeyPath
	CAPath := Parameters.CAPath

	clientCertPool := x509.NewCertPool()

	cacert, err := ioutil.ReadFile(CAPath)
	cert, err := tls.LoadX509KeyPair(CertPath, KeyPath)
	if err != nil {
		log.Error("ReadFile err: ", err)
		return nil, err
	}

	ret := clientCertPool.AppendCertsFromPEM(cacert)
	if !ret {
		return nil, errors.New("failed to parse root certificate")
	}

	conf := &tls.Config{
		RootCAs:      clientCertPool,
		Certificates: []tls.Certificate{cert},
	}

	conn, err := tls.Dial("tcp", nodeAddr, conf)
	if err != nil {
		log.Error("Dial failed: ", err)
		return nil, err
	}
	return conn, nil
}

func (node node) Tx(buf []byte) {
	log.Debug()
	str := hex.EncodeToString(buf)
	log.Debug(fmt.Sprintf("TX buf length: %d\n%s", len(buf), str))

	_, err := node.conn.Write(buf)
	if err != nil {
		log.Error("Error sending messge to peer node ", err.Error())
	}
}
<commit_msg>Set the node inactive status at the period update process<commit_after>package node

import (
	"DNA/common/log"
	. "DNA/config"
	. "DNA/net/message"
	. "DNA/net/protocol"
	"crypto/tls"
	"crypto/x509"
	"encoding/hex"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net"
	"os"
	"strconv"
	"strings"
	"time"
)

type link struct {
	//Todo Add lock here
	addr  string    // The address of the node
	conn  net.Conn  // Connect socket with the peer node
	port  uint16    // The server port of the node
	time  time.Time // The latest time the node activity
	rxBuf struct {  // The RX buffer of this node to solve mutliple packets problem
		p   []byte
		len int
	}
	connCnt uint64 // The connection count
}

// Shrinking the buf to the exactly reading in byte length
//@Return @1 the start header of next message, the left length of the next message
func unpackNodeBuf(node *node, buf []byte) {
	var msgLen int
	var msgBuf []byte
	if node.rxBuf.p == nil {
		if len(buf) < MSGHDRLEN {
			log.Warn("Unexpected size of received message")
			errors.New("Unexpected size of received message")
			return
		}
		// FIXME Check the payload < 0 error case
		msgLen = PayloadLen(buf) + MSGHDRLEN
	} else {
		msgLen = node.rxBuf.len
	}

	if len(buf) == msgLen {
		msgBuf = append(node.rxBuf.p, buf[:]...)
		go HandleNodeMsg(node, msgBuf, len(msgBuf))
		node.rxBuf.p = nil
		node.rxBuf.len = 0
	} else if len(buf) < msgLen {
		node.rxBuf.p = append(node.rxBuf.p, buf[:]...)
		node.rxBuf.len = msgLen - len(buf)
	} else {
		msgBuf = append(node.rxBuf.p, buf[0:msgLen]...)
		go HandleNodeMsg(node, msgBuf, len(msgBuf))
		node.rxBuf.p = nil
		node.rxBuf.len = 0

		unpackNodeBuf(node, buf[msgLen:])
	}
}

func (node *node) rx() error {
	conn := node.getConn()
	buf := make([]byte, MAXBUFLEN)
	for {
		len, err := conn.Read(buf[0:(MAXBUFLEN - 1)])
		buf[MAXBUFLEN-1] = 0 //Prevent overflow
		switch err {
		case nil:
			unpackNodeBuf(node, buf[0:len])
			break
		case io.EOF:
			break
		default:
			log.Error("Read connetion error ", err)
			goto disconnect
		}
	}

disconnect:
	err := conn.Close()
	return err
}

func printIPAddr() {
	host, _ := os.Hostname()
	addrs, _ := net.LookupIP(host)
	for _, addr := range addrs {
		if ipv4 := addr.To4(); ipv4 != nil {
			log.Info("IPv4: ", ipv4)
		}
	}
}

func (link link) CloseConn() {
	link.conn.Close()
}

func (n *node) initConnection() {
	isTls := Parameters.IsTLS
	var listener net.Listener
	var err error
	if isTls {
		listener, err = initTlsListen()
		if err != nil {
			log.Error("TLS listen failed")
			return
		}
	} else {
		listener, err = initNonTlsListen()
		if err != nil {
			log.Error("non TLS listen failed")
			return
		}
	}
	for {
		conn, err := listener.Accept()
		if err != nil {
			log.Error("Error accepting ", err.Error())
			return
		}
		log.Info("Remote node connect with ", conn.RemoteAddr(), conn.LocalAddr())

		n.link.connCnt++

		node := NewNode()
		node.addr, err = parseIPaddr(conn.RemoteAddr().String())
		node.local = n
		node.conn = conn
		go node.rx()
	}
	//TODO Release the net listen resouce
}

func initNonTlsListen() (net.Listener, error) {
	log.Debug()
	listener, err := net.Listen("tcp", ":"+strconv.Itoa(Parameters.NodePort))
	if err != nil {
		log.Error("Error listening\n", err.Error())
		return nil, err
	}
	return listener, nil
}

func initTlsListen() (net.Listener, error) {
	CertPath := Parameters.CertPath
	KeyPath := Parameters.KeyPath
	CAPath := Parameters.CAPath

	// load cert
	cert, err := tls.LoadX509KeyPair(CertPath, KeyPath)
	if err != nil {
		log.Error("load keys fail", err)
		return nil, err
	}
	// load root ca
	caData, err := ioutil.ReadFile(CAPath)
	if err != nil {
		log.Error("read ca fail", err)
		return nil, err
	}
	pool := x509.NewCertPool()
	ret := pool.AppendCertsFromPEM(caData)
	if !ret {
		return nil, errors.New("failed to parse root certificate")
	}

	tlsConfig := &tls.Config{
		Certificates: []tls.Certificate{cert},
		RootCAs:      pool,
		ClientAuth:   tls.RequireAndVerifyClientCert,
		ClientCAs:    pool,
	}

	log.Info("TLS listen port is ", strconv.Itoa(Parameters.NodePort))
	listener, err := tls.Listen("tcp", ":"+strconv.Itoa(Parameters.NodePort), tlsConfig)
	if err != nil {
		log.Error(err)
		return nil, err
	}
	return listener, nil
}

func parseIPaddr(s string) (string, error) {
	i := strings.Index(s, ":")
	if i < 0 {
		log.Warn("Split IP address&port error")
		return s, errors.New("Split IP address&port error")
	}
	return s[:i], nil
}

func (node *node) Connect(nodeAddr string) error {
	log.Debug()
	isTls := Parameters.IsTLS
	var conn net.Conn
	var err error
	if isTls {
		conn, err = TLSDial(nodeAddr)
		if err != nil {
			log.Error("TLS connect failed: ", err)
			return nil
		}
	} else {
		conn, err = NonTLSDial(nodeAddr)
		if err != nil {
			log.Error("non TLS connect failed:", err)
			return nil
		}
	}
	node.link.connCnt++

	n := NewNode()
	n.conn = conn
	n.addr, err = parseIPaddr(conn.RemoteAddr().String())
	n.local = node

	log.Info(fmt.Sprintf("Connect node %s connect with %s with %s",
		conn.LocalAddr().String(), conn.RemoteAddr().String(),
		conn.RemoteAddr().Network()))
	go n.rx()

	n.SetState(HAND)
	buf, _ := NewVersion(node)
	n.Tx(buf)

	return nil
}

func NonTLSDial(nodeAddr string) (net.Conn, error) {
	log.Debug()
	conn, err := net.Dial("tcp", nodeAddr)
	if err != nil {
		log.Error("Error dialing\n", err.Error())
		return nil, err
	}
	return conn, nil
}

func TLSDial(nodeAddr string) (net.Conn, error) {
	CertPath := Parameters.CertPath
	KeyPath := Parameters.KeyPath
	CAPath := Parameters.CAPath

	clientCertPool := x509.NewCertPool()

	cacert, err := ioutil.ReadFile(CAPath)
	cert, err := tls.LoadX509KeyPair(CertPath, KeyPath)
	if err != nil {
		log.Error("ReadFile err: ", err)
		return nil, err
	}

	ret := clientCertPool.AppendCertsFromPEM(cacert)
	if !ret {
		return nil, errors.New("failed to parse root certificate")
	}

	conf := &tls.Config{
		RootCAs:      clientCertPool,
		Certificates: []tls.Certificate{cert},
	}

	conn, err := tls.Dial("tcp", nodeAddr, conf)
	if err != nil {
		log.Error("Dial failed: ", err)
		return nil, err
	}
	return conn, nil
}

func (node node) Tx(buf []byte) {
	log.Debug()
	str := hex.EncodeToString(buf)
	log.Debug(fmt.Sprintf("TX buf length: %d\n%s", len(buf), str))

	_, err := node.conn.Write(buf)
	if err != nil {
		log.Error("Error sending messge to peer node ", err.Error())
	}
}
<|endoftext|>


<commit_before>package net

import (
	"io"

	_ "github.com/v2ray/v2ray-core/log"
)

const (
	bufferSize = 32 * 1024
)

func ReaderToChan(stream chan<- []byte, reader io.Reader) error {
	for {
		buffer := make([]byte, bufferSize)
		nBytes, err := reader.Read(buffer)
		if err != nil {
			return err
		}
		stream <- buffer[:nBytes]
	}
	return nil
}

func ChanToWriter(writer io.Writer, stream <-chan []byte) error {
	for buffer := range stream {
		_, err := writer.Write(buffer)
		if err != nil {
			return err
		}
	}
	return nil
}
<commit_msg>Transport buffer even when there is an error<commit_after>package net

import (
	"io"

	_ "github.com/v2ray/v2ray-core/log"
)

const (
	bufferSize = 32 * 1024
)

func ReaderToChan(stream chan<- []byte, reader io.Reader) error {
	for {
		buffer := make([]byte, bufferSize)
		nBytes, err := reader.Read(buffer)
    if nBytes > 0 {
      stream <- buffer[:nBytes]
    }
		if err != nil {
			return err
		}
	}
	return nil
}

func ChanToWriter(writer io.Writer, stream <-chan []byte) error {
	for buffer := range stream {
		_, err := writer.Write(buffer)
		if err != nil {
			return err
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package raft

import (
	"bufio"
	"fmt"
	"github.com/ugorji/go/codec"
	"io"
	"log"
	"net"
	"sync"
	"time"
)

const (
	rpcAppendEntries uint8 = iota
	rpcRequestVote
	rpcInstallSnapshot
	DefaultTimeoutScale = 1024 * 1024 // 1MB
)

var (
	TransportShutdown = fmt.Errorf("transport shutdown")
)

/*

NetworkTransport provides a network based transport that can be
used to communicate with Raft on remote machines. It requires
an underlying layer to provide a stream abstraction, which can
be simple TCP, TLS, etc. Underlying addresses must be castable to TCPAddr

This transport is very simple and lightweight. Each RPC request is
framed by sending a byte that indicates the message type, followed
by the MsgPack encoded request.

The response is an error string followed by the response object,
both are encoded using MsgPack.

InstallSnapshot is special, in that after the RPC request we stream
the entire state. That socket is not re-used as the connection state
is not known if there is an error.

*/
type NetworkTransport struct {
	connPool     map[string][]*netConn
	connPoolLock sync.Mutex

	consumeCh chan RPC

	dialer   net.Dialer
	listener net.Listener
	maxPool  int

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex

	timeout      time.Duration
	timeoutScale int
}

type netConn struct {
	target net.Addr
	conn   net.Conn
	r      *bufio.Reader
	w      *bufio.Writer
	dec    *codec.Decoder
	enc    *codec.Encoder
}

func (n *netConn) Release() error {
	return n.conn.Close()
}

// Creates a new network transport with the given dailer and listener
// The timeout is used to apply I/O deadlines. For InstallSnapshot, we multiply
// the timeout by (SnapshotSize / TimeoutScale)
func NewNetworkTransport(dialer net.Dialer, listener net.Listener, timeout time.Duration) *NetworkTransport {
	trans := &NetworkTransport{
		connPool:     make(map[string][]*netConn),
		consumeCh:    make(chan RPC),
		dialer:       dialer,
		listener:     listener,
		maxPool:      2,
		shutdownCh:   make(chan struct{}),
		timeout:      timeout,
		timeoutScale: DefaultTimeoutScale,
	}
	go trans.listen()
	return trans
}

// SetMaxPool is used to set the maximum number of pooled connections per host
func (n *NetworkTransport) SetMaxPool(maxPool int) {
	reduced := maxPool < n.maxPool
	n.maxPool = maxPool

	// If we reduced the pool size, we need to close any open connections
	if reduced {
		n.shutdownLock.Lock()
		defer n.shutdownLock.Lock()

		for key := range n.connPool {
			conns := n.connPool[key]
			if len(conns) > maxPool {
				for i := maxPool; i < len(conns); i++ {
					conns[i].Release()
					conns[i] = nil
				}
				n.connPool[key] = conns[:maxPool]
			}
		}
	}
}

// SetTimeoutScale is used to change the default timeout scale
func (n *NetworkTransport) SetTimeoutScale(scale int) {
	n.timeoutScale = scale
}

// Close is used to stop the network transport
func (n *NetworkTransport) Close() error {
	n.shutdownLock.Lock()
	defer n.shutdownLock.Lock()

	if !n.shutdown {
		close(n.shutdownCh)
		n.listener.Close()
		n.shutdown = true
		n.SetMaxPool(0)
	}
	return nil
}

func (n *NetworkTransport) Consumer() <-chan RPC {
	return n.consumeCh
}

func (n *NetworkTransport) LocalAddr() net.Addr {
	return n.listener.Addr()
}

// getExistingConn is used to grab a pooled connection
func (n *NetworkTransport) getPooledConn(target net.Addr) *netConn {
	n.connPoolLock.Lock()
	defer n.connPoolLock.Unlock()

	key := target.String()
	conns, ok := n.connPool[key]
	if !ok || len(conns) == 0 {
		return nil
	}

	var conn *netConn
	num := len(conns)
	conn, conns[num-1] = conns[num-1], nil
	n.connPool[key] = conns[:num-1]
	return conn
}

// getConn is used to get a connection from the pool
func (n *NetworkTransport) getConn(target net.Addr) (*netConn, error) {
	// Check for a pooled conn
	if conn := n.getPooledConn(target); conn != nil {
		return conn, nil
	}

	// Dial a new connection
	conn, err := n.dialer.Dial(target.Network(), target.String())
	if err != nil {
		return nil, err
	}

	// Wrap the conn
	netConn := &netConn{
		target: target,
		conn:   conn,
		r:      bufio.NewReader(conn),
		w:      bufio.NewWriter(conn),
	}

	// Setup encoder/decoders
	netConn.dec = codec.NewDecoder(netConn.r, &codec.MsgpackHandle{})
	netConn.enc = codec.NewEncoder(netConn.w, &codec.MsgpackHandle{})

	// Done
	return netConn, nil
}

// returnConn returns a connection back to the pool
func (n *NetworkTransport) returnConn(conn *netConn) {
	n.connPoolLock.Lock()
	defer n.connPoolLock.Unlock()

	key := conn.target.String()
	conns, _ := n.connPool[key]

	if len(conns) < n.maxPool {
		n.connPool[key] = append(conns, conn)
	} else {
		conn.Release()
	}
}

func (n *NetworkTransport) AppendEntries(target net.Addr, args *AppendEntriesRequest, resp *AppendEntriesResponse) error {
	return n.genericRPC(target, rpcAppendEntries, args, resp)
}

func (n *NetworkTransport) RequestVote(target net.Addr, args *RequestVoteRequest, resp *RequestVoteResponse) error {
	return n.genericRPC(target, rpcRequestVote, args, resp)
}

// genericRPC handles a simple request/response RPC
func (n *NetworkTransport) genericRPC(target net.Addr, rpcType uint8, args interface{}, resp interface{}) error {
	// Get a conn
	conn, err := n.getConn(target)
	if err != nil {
		return err
	}

	// Set a deadline
	if n.timeout > 0 {
		conn.conn.SetDeadline(time.Now().Add(n.timeout))
	}

	// Send the RPC
	if err := sendRPC(conn, rpcType, args); err != nil {
		return err
	}

	// Decode the response
	return n.decodeResponse(conn, resp, true)
}

func (n *NetworkTransport) InstallSnapshot(target net.Addr, args *InstallSnapshotRequest, resp *InstallSnapshotResponse, data io.ReadCloser) error {
	// Make sure the state file is closed
	defer data.Close()

	// Get a conn, always close for InstallSnapshot
	conn, err := n.getConn(target)
	if err != nil {
		return err
	}
	defer conn.Release()

	// Set a deadline, scaled by request size
	if n.timeout > 0 {
		timeout := n.timeout * time.Duration(args.Size/int64(n.timeoutScale))
		conn.conn.SetDeadline(time.Now().Add(timeout))
	}

	// Send the RPC
	if err := sendRPC(conn, rpcInstallSnapshot, args); err != nil {
		return err
	}

	// Stream the state
	if _, err := io.Copy(conn.w, data); err != nil {
		return err
	}

	// Flush
	if err := conn.w.Flush(); err != nil {
		return err
	}

	// Decode the response, do not return conn
	return n.decodeResponse(conn, resp, false)
}

func (n *NetworkTransport) EncodePeer(p net.Addr) []byte {
	return []byte(p.String())
}

func (n *NetworkTransport) DecodePeer(buf []byte) net.Addr {
	addr, err := net.ResolveTCPAddr("tcp", string(buf))
	if err != nil {
		panic(fmt.Errorf("Failed to parse network address: %s", buf))
	}
	return addr
}

// listen is used to handling incoming connections
func (n *NetworkTransport) listen() {
	for {
		// Accept incoming connections
		conn, err := n.listener.Accept()
		if err != nil {
			if n.shutdown {
				return
			}
			log.Printf("[ERR] Failed to accept connection: %v", err)
			continue
		}

		// Handle the connection in dedicated routine
		go n.handleConn(conn)
	}
}

// handleConn is used to handle an inbound connection for its lifespan
func (n *NetworkTransport) handleConn(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)
	w := bufio.NewWriter(conn)
	dec := codec.NewDecoder(r, &codec.MsgpackHandle{})
	enc := codec.NewEncoder(w, &codec.MsgpackHandle{})

	for {
		if err := n.handleCommand(r, dec, enc); err != nil {
			if err != io.EOF {
				log.Printf("[ERR] Failed to decode incoming command: %v", err)
			}
			return
		}
		if err := w.Flush(); err != nil {
			log.Printf("[ERR] Failed to flush response: %v", err)
			return
		}
	}
}

// handleCommand is used to decode and dispatch a single command
func (n *NetworkTransport) handleCommand(r *bufio.Reader, dec *codec.Decoder, enc *codec.Encoder) error {
	// Get the rpc type
	rpcType, err := r.ReadByte()
	if err != nil {
		return err
	}

	// Create the RPC object
	respCh := make(chan RPCResponse)
	rpc := RPC{
		RespChan: respCh,
	}

	// Decode the command
	switch rpcType {
	case rpcAppendEntries:
		var req AppendEntriesRequest
		if err := dec.Decode(&req); err != nil {
			return err
		}
		rpc.Command = &req

	case rpcRequestVote:
		var req RequestVoteRequest
		if err := dec.Decode(&req); err != nil {
			return err
		}
		rpc.Command = &req

	case rpcInstallSnapshot:
		var req InstallSnapshotRequest
		if err := dec.Decode(&req); err != nil {
			return err
		}
		rpc.Command = &req
		rpc.Reader = io.LimitReader(r, req.Size)

	default:
		return fmt.Errorf("unknown rpc type %d", rpcType)
	}

	// Dispatch the RPC
	select {
	case n.consumeCh <- rpc:
	case <-n.shutdownCh:
		return TransportShutdown
	}

	// Wait for response
	select {
	case resp := <-respCh:
		// Send the error first
		respErr := ""
		if resp.Error != nil {
			respErr = resp.Error.Error()
		}
		if err := enc.Encode(respErr); err != nil {
			return err
		}

		// Send the response
		if err := enc.Encode(resp.Response); err != nil {
			return err
		}
	case <-n.shutdownCh:
		return TransportShutdown
	}
	return nil
}

// decodeResponse is used to decode an RPC response and return the conn
func (n *NetworkTransport) decodeResponse(conn *netConn, resp interface{}, retConn bool) error {
	// Decode the error if any
	var rpcError string
	if err := conn.dec.Decode(&rpcError); err != nil {
		conn.Release()
		return err
	}

	// Decode the response
	if err := conn.dec.Decode(resp); err != nil {
		conn.Release()
		return err
	}

	// Return the conn
	if retConn {
		n.returnConn(conn)
	}

	// Format an error if any
	if rpcError != "" {
		return fmt.Errorf(rpcError)
	}
	return nil
}

// sendRPC is used to encode and send the RPC
func sendRPC(conn *netConn, rpcType uint8, args interface{}) error {
	// Write the request type
	if err := conn.w.WriteByte(rpcType); err != nil {
		conn.Release()
		return err
	}

	// Send the request
	if err := conn.enc.Encode(args); err != nil {
		conn.Release()
		return err
	}

	// Flush
	if err := conn.w.Flush(); err != nil {
		conn.Release()
		return err
	}
	return nil
}
<commit_msg>Adding StreamLayer abstraction<commit_after>package raft

import (
	"bufio"
	"fmt"
	"github.com/ugorji/go/codec"
	"io"
	"log"
	"net"
	"sync"
	"time"
)

const (
	rpcAppendEntries uint8 = iota
	rpcRequestVote
	rpcInstallSnapshot
	DefaultTimeoutScale = 1024 * 1024 // 1MB
)

var (
	TransportShutdown = fmt.Errorf("transport shutdown")
)

/*

NetworkTransport provides a network based transport that can be
used to communicate with Raft on remote machines. It requires
an underlying stream layer to provide a stream abstraction, which can
be simple TCP, TLS, etc. Underlying addresses must be castable to TCPAddr

This transport is very simple and lightweight. Each RPC request is
framed by sending a byte that indicates the message type, followed
by the MsgPack encoded request.

The response is an error string followed by the response object,
both are encoded using MsgPack.

InstallSnapshot is special, in that after the RPC request we stream
the entire state. That socket is not re-used as the connection state
is not known if there is an error.

*/
type NetworkTransport struct {
	connPool     map[string][]*netConn
	connPoolLock sync.Mutex

	consumeCh chan RPC

	maxPool int

	shutdown     bool
	shutdownCh   chan struct{}
	shutdownLock sync.Mutex

	stream StreamLayer

	timeout      time.Duration
	timeoutScale int
}

// StreamLayer is used with the NetworkTransport to provide
// the low level stream abstraction
type StreamLayer interface {
	net.Listener

	// Dial is used to create a new outgoing connection
	Dial(network, address string, timeout time.Duration) (net.Conn, error)
}

type netConn struct {
	target net.Addr
	conn   net.Conn
	r      *bufio.Reader
	w      *bufio.Writer
	dec    *codec.Decoder
	enc    *codec.Encoder
}

func (n *netConn) Release() error {
	return n.conn.Close()
}

// Creates a new network transport with the given dailer and listener
// The timeout is used to apply I/O deadlines. For InstallSnapshot, we multiply
// the timeout by (SnapshotSize / TimeoutScale)
func NewNetworkTransport(stream StreamLayer, timeout time.Duration) *NetworkTransport {
	trans := &NetworkTransport{
		connPool:     make(map[string][]*netConn),
		consumeCh:    make(chan RPC),
		maxPool:      2,
		shutdownCh:   make(chan struct{}),
		stream:       stream,
		timeout:      timeout,
		timeoutScale: DefaultTimeoutScale,
	}
	go trans.listen()
	return trans
}

// SetMaxPool is used to set the maximum number of pooled connections per host
func (n *NetworkTransport) SetMaxPool(maxPool int) {
	reduced := maxPool < n.maxPool
	n.maxPool = maxPool

	// If we reduced the pool size, we need to close any open connections
	if reduced {
		n.shutdownLock.Lock()
		defer n.shutdownLock.Lock()

		for key := range n.connPool {
			conns := n.connPool[key]
			if len(conns) > maxPool {
				for i := maxPool; i < len(conns); i++ {
					conns[i].Release()
					conns[i] = nil
				}
				n.connPool[key] = conns[:maxPool]
			}
		}
	}
}

// SetTimeoutScale is used to change the default timeout scale
func (n *NetworkTransport) SetTimeoutScale(scale int) {
	n.timeoutScale = scale
}

// Close is used to stop the network transport
func (n *NetworkTransport) Close() error {
	n.shutdownLock.Lock()
	defer n.shutdownLock.Lock()

	if !n.shutdown {
		close(n.shutdownCh)
		n.stream.Close()
		n.shutdown = true
		n.SetMaxPool(0)
	}
	return nil
}

func (n *NetworkTransport) Consumer() <-chan RPC {
	return n.consumeCh
}

func (n *NetworkTransport) LocalAddr() net.Addr {
	return n.stream.Addr()
}

// getExistingConn is used to grab a pooled connection
func (n *NetworkTransport) getPooledConn(target net.Addr) *netConn {
	n.connPoolLock.Lock()
	defer n.connPoolLock.Unlock()

	key := target.String()
	conns, ok := n.connPool[key]
	if !ok || len(conns) == 0 {
		return nil
	}

	var conn *netConn
	num := len(conns)
	conn, conns[num-1] = conns[num-1], nil
	n.connPool[key] = conns[:num-1]
	return conn
}

// getConn is used to get a connection from the pool
func (n *NetworkTransport) getConn(target net.Addr) (*netConn, error) {
	// Check for a pooled conn
	if conn := n.getPooledConn(target); conn != nil {
		return conn, nil
	}

	// Dial a new connection
	conn, err := n.stream.Dial(target.Network(), target.String(), n.timeout)
	if err != nil {
		return nil, err
	}

	// Wrap the conn
	netConn := &netConn{
		target: target,
		conn:   conn,
		r:      bufio.NewReader(conn),
		w:      bufio.NewWriter(conn),
	}

	// Setup encoder/decoders
	netConn.dec = codec.NewDecoder(netConn.r, &codec.MsgpackHandle{})
	netConn.enc = codec.NewEncoder(netConn.w, &codec.MsgpackHandle{})

	// Done
	return netConn, nil
}

// returnConn returns a connection back to the pool
func (n *NetworkTransport) returnConn(conn *netConn) {
	n.connPoolLock.Lock()
	defer n.connPoolLock.Unlock()

	key := conn.target.String()
	conns, _ := n.connPool[key]

	if len(conns) < n.maxPool {
		n.connPool[key] = append(conns, conn)
	} else {
		conn.Release()
	}
}

func (n *NetworkTransport) AppendEntries(target net.Addr, args *AppendEntriesRequest, resp *AppendEntriesResponse) error {
	return n.genericRPC(target, rpcAppendEntries, args, resp)
}

func (n *NetworkTransport) RequestVote(target net.Addr, args *RequestVoteRequest, resp *RequestVoteResponse) error {
	return n.genericRPC(target, rpcRequestVote, args, resp)
}

// genericRPC handles a simple request/response RPC
func (n *NetworkTransport) genericRPC(target net.Addr, rpcType uint8, args interface{}, resp interface{}) error {
	// Get a conn
	conn, err := n.getConn(target)
	if err != nil {
		return err
	}

	// Set a deadline
	if n.timeout > 0 {
		conn.conn.SetDeadline(time.Now().Add(n.timeout))
	}

	// Send the RPC
	if err := sendRPC(conn, rpcType, args); err != nil {
		return err
	}

	// Decode the response
	return n.decodeResponse(conn, resp, true)
}

func (n *NetworkTransport) InstallSnapshot(target net.Addr, args *InstallSnapshotRequest, resp *InstallSnapshotResponse, data io.ReadCloser) error {
	// Make sure the state file is closed
	defer data.Close()

	// Get a conn, always close for InstallSnapshot
	conn, err := n.getConn(target)
	if err != nil {
		return err
	}
	defer conn.Release()

	// Set a deadline, scaled by request size
	if n.timeout > 0 {
		timeout := n.timeout * time.Duration(args.Size/int64(n.timeoutScale))
		conn.conn.SetDeadline(time.Now().Add(timeout))
	}

	// Send the RPC
	if err := sendRPC(conn, rpcInstallSnapshot, args); err != nil {
		return err
	}

	// Stream the state
	if _, err := io.Copy(conn.w, data); err != nil {
		return err
	}

	// Flush
	if err := conn.w.Flush(); err != nil {
		return err
	}

	// Decode the response, do not return conn
	return n.decodeResponse(conn, resp, false)
}

func (n *NetworkTransport) EncodePeer(p net.Addr) []byte {
	return []byte(p.String())
}

func (n *NetworkTransport) DecodePeer(buf []byte) net.Addr {
	addr, err := net.ResolveTCPAddr("tcp", string(buf))
	if err != nil {
		panic(fmt.Errorf("Failed to parse network address: %s", buf))
	}
	return addr
}

// listen is used to handling incoming connections
func (n *NetworkTransport) listen() {
	for {
		// Accept incoming connections
		conn, err := n.stream.Accept()
		if err != nil {
			if n.shutdown {
				return
			}
			log.Printf("[ERR] Failed to accept connection: %v", err)
			continue
		}

		// Handle the connection in dedicated routine
		go n.handleConn(conn)
	}
}

// handleConn is used to handle an inbound connection for its lifespan
func (n *NetworkTransport) handleConn(conn net.Conn) {
	defer conn.Close()
	r := bufio.NewReader(conn)
	w := bufio.NewWriter(conn)
	dec := codec.NewDecoder(r, &codec.MsgpackHandle{})
	enc := codec.NewEncoder(w, &codec.MsgpackHandle{})

	for {
		if err := n.handleCommand(r, dec, enc); err != nil {
			if err != io.EOF {
				log.Printf("[ERR] Failed to decode incoming command: %v", err)
			}
			return
		}
		if err := w.Flush(); err != nil {
			log.Printf("[ERR] Failed to flush response: %v", err)
			return
		}
	}
}

// handleCommand is used to decode and dispatch a single command
func (n *NetworkTransport) handleCommand(r *bufio.Reader, dec *codec.Decoder, enc *codec.Encoder) error {
	// Get the rpc type
	rpcType, err := r.ReadByte()
	if err != nil {
		return err
	}

	// Create the RPC object
	respCh := make(chan RPCResponse)
	rpc := RPC{
		RespChan: respCh,
	}

	// Decode the command
	switch rpcType {
	case rpcAppendEntries:
		var req AppendEntriesRequest
		if err := dec.Decode(&req); err != nil {
			return err
		}
		rpc.Command = &req

	case rpcRequestVote:
		var req RequestVoteRequest
		if err := dec.Decode(&req); err != nil {
			return err
		}
		rpc.Command = &req

	case rpcInstallSnapshot:
		var req InstallSnapshotRequest
		if err := dec.Decode(&req); err != nil {
			return err
		}
		rpc.Command = &req
		rpc.Reader = io.LimitReader(r, req.Size)

	default:
		return fmt.Errorf("unknown rpc type %d", rpcType)
	}

	// Dispatch the RPC
	select {
	case n.consumeCh <- rpc:
	case <-n.shutdownCh:
		return TransportShutdown
	}

	// Wait for response
	select {
	case resp := <-respCh:
		// Send the error first
		respErr := ""
		if resp.Error != nil {
			respErr = resp.Error.Error()
		}
		if err := enc.Encode(respErr); err != nil {
			return err
		}

		// Send the response
		if err := enc.Encode(resp.Response); err != nil {
			return err
		}
	case <-n.shutdownCh:
		return TransportShutdown
	}
	return nil
}

// decodeResponse is used to decode an RPC response and return the conn
func (n *NetworkTransport) decodeResponse(conn *netConn, resp interface{}, retConn bool) error {
	// Decode the error if any
	var rpcError string
	if err := conn.dec.Decode(&rpcError); err != nil {
		conn.Release()
		return err
	}

	// Decode the response
	if err := conn.dec.Decode(resp); err != nil {
		conn.Release()
		return err
	}

	// Return the conn
	if retConn {
		n.returnConn(conn)
	}

	// Format an error if any
	if rpcError != "" {
		return fmt.Errorf(rpcError)
	}
	return nil
}

// sendRPC is used to encode and send the RPC
func sendRPC(conn *netConn, rpcType uint8, args interface{}) error {
	// Write the request type
	if err := conn.w.WriteByte(rpcType); err != nil {
		conn.Release()
		return err
	}

	// Send the request
	if err := conn.enc.Encode(args); err != nil {
		conn.Release()
		return err
	}

	// Flush
	if err := conn.w.Flush(); err != nil {
		conn.Release()
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package neterr

import (
	"io"
	"net"
	"net/url"
)

// IsNetworkError returns true if the error's cause is: io.ErrUnexpectedEOF,
// any *net.OpError, any *url.Error, any URL that implements `Temporary()`
// (and returns true)
func IsNetworkError(err error) bool {
	if err == io.ErrUnexpectedEOF {
		return true
	}

	if causer, ok := err.(causer); ok {
		return IsNetworkError(causer.Cause())
	}

	if urlError, ok := err.(*url.Error); ok {
		return IsNetworkError(urlError.Err)
	}

	if _, ok := err.(*net.OpError); ok {
		return true
	}

	if te, ok := err.(temporary); ok {
		return te.Temporary()
	}

	return false
}

type temporary interface {
	Temporary() bool
}

type causer interface {
	Cause() error
}
<commit_msg>An idled connection is a network error<commit_after>package neterr

import (
	"io"
	"net"
	"net/url"

	"github.com/getlantern/idletiming"
)

// IsNetworkError returns true if the error's cause is: io.ErrUnexpectedEOF,
// any *net.OpError, any *url.Error, any URL that implements `Temporary()`
// (and returns true)
func IsNetworkError(err error) bool {
	if err == io.ErrUnexpectedEOF {
		return true
	}

	if causer, ok := err.(causer); ok {
		return IsNetworkError(causer.Cause())
	}

	if urlError, ok := err.(*url.Error); ok {
		return IsNetworkError(urlError.Err)
	}

	if _, ok := err.(*net.OpError); ok {
		return true
	}

	if err == idletiming.ErrIdled {
		return true
	}

	if te, ok := err.(temporary); ok {
		return te.Temporary()
	}

	return false
}

type temporary interface {
	Temporary() bool
}

type causer interface {
	Cause() error
}
<|endoftext|>


<commit_before>// Config is key, value map for system level and component configuration.
// Key is a string and represents a config parameter, and corresponding
// value is an interface{} that can be consumed using accessor methods
// based on the context of config-value.
//
// Config maps are immutable and newer versions can be created using accessor
// methods.
//
// Shape of config-parameter, the key string, is sequence of alpha-numeric
// characters separated by one or more '.' , eg,
//      "projector.adminport.readtimeout"

package common

import "encoding/json"
import "strings"

// Config is a key, value map with key always being a string
// represents a config-parameter.
type Config map[string]ConfigValue

// ConfigValue for each parameter.
type ConfigValue struct {
	Value      interface{}
	Help       string
	DefaultVal interface{}
}

// SystemConfig is default configuration for system and components.
// configuration parameters follow flat namespacing like,
//      "maxVbuckets"  for system-level config parameter
//      "projector.xxx" for projector component.
//      "projector.adminport.xxx" for adminport under projector component.
// etc...
var SystemConfig = Config{
	// system parameters
	"maxVbuckets": ConfigValue{
		1024,
		"number of vbuckets configured in KV",
		1024,
	},
	// log parameters
	// TODO: add configuration for log file-name and other types of writer.
	"log.ignore": ConfigValue{
		false,
		"ignores all logging, irrespective of the log-level",
		false,
	},
	"log.level": ConfigValue{
		"info",
		"logging level for the system",
		"info",
	},
	// projector parameters
	"projector.name": ConfigValue{
		"projector",
		"human readable name for this projector",
		"projector",
	},
	"projector.clusterAddr": ConfigValue{
		"localhost:9000",
		"KV cluster's address to be used by projector",
		"localhost:9000",
	},
	"projector.kvAddrs": ConfigValue{
		"127.0.0.1:9000",
		"Comma separated list of KV-address to read mutations, this need to " +
			"exactly match with KV-node's configured address",
		"127.0.0.1:9000",
	},
	"projector.colocate": ConfigValue{
		true,
		"Whether projector will be colocated with KV. In which case " +
			"`kvaddrs` specified above will be discarded",
		true,
	},
	"projector.routerEndpointFactory": ConfigValue{
		RouterEndpointFactory(nil),
		"RouterEndpointFactory callback to generate endpoint instances " +
			"to push data to downstream",
		RouterEndpointFactory(nil),
	},
	"projector.feedWaitStreamReqTimeout": ConfigValue{
		10 * 1000,
		"timeout, in milliseconds, to await a response for StreamRequest",
		10 * 1000,
	},
	"projector.feedWaitStreamEndTimeout": ConfigValue{
		10 * 1000,
		"timeout, in milliseconds, to await a response for StreamEnd",
		10 * 1000,
	},
	"projector.mutationChanSize": ConfigValue{
		10000,
		"channel size of projector's data path routine",
		10000,
	},
	"projector.feedChanSize": ConfigValue{
		100,
		"channel size for feed's control path and back path.",
		100,
	},
	"projector.vbucketSyncTimeout": ConfigValue{
		500,
		"timeout, in milliseconds, for sending periodic Sync messages.",
		500,
	},
	// projector adminport parameters
	"projector.adminport.name": ConfigValue{
		"projector.adminport",
		"human readable name for this adminport, must be supplied",
		"projector.adminport",
	},
	"projector.adminport.listenAddr": ConfigValue{
		"",
		"projector's adminport address listen for request.",
		"",
	},
	"projector.adminport.urlPrefix": ConfigValue{
		"/adminport/",
		"url prefix (script-path) for adminport used by projector",
		"/adminport/",
	},
	"projector.adminport.readTimeout": ConfigValue{
		0,
		"timeout in milliseconds, is read timeout for adminport http server " +
			"used by projector",
		0,
	},
	"projector.adminport.writeTimeout": ConfigValue{
		0,
		"timeout in milliseconds, is write timeout for adminport http server " +
			"used by projector",
		0,
	},
	"projector.adminport.maxHeaderBytes": ConfigValue{
		1 << 20, // 1 MegaByte
		"in bytes, is max. length of adminport http header " +
			"used by projector",
		1 << 20, // 1 MegaByte
	},
	// projector's adminport client
	"projector.client.retryInterval": ConfigValue{
		16,
		"retryInterval, in milliseconds, when connection refused by server",
		16,
	},
	"projector.client.maxRetries": ConfigValue{
		5,
		"maximum number of timest to retry",
		5,
	},
	"projector.client.exponentialBackoff": ConfigValue{
		2,
		"multiplying factor on retryInterval for every attempt with server",
		2,
	},
	// TODO: This configuration param is same as the above.
	"projector.client.urlPrefix": ConfigValue{
		"/adminport/",
		"url prefix (script-path) for adminport used by projector",
		"/adminport/",
	},
	// projector dataport client parameters
	// TODO: this configuration option should be tunnable for each feed.
	"endpoint.dataport.remoteBlock": ConfigValue{
		false,
		"should dataport endpoint block when remote is slow ?",
		false,
	},
	"endpoint.dataport.keyChanSize": ConfigValue{
		10000,
		"channel size of dataport endpoints data input",
		10000,
	},
	"endpoint.dataport.bufferSize": ConfigValue{
		100,
		"number of entries to buffer before flushing it, where each entry " +
			"is for a vbucket's set of mutations that was flushed by the endpoint.",
		100,
	},
	"endpoint.dataport.bufferTimeout": ConfigValue{
		1,
		"timeout in milliseconds, to flush vbucket-mutations from endpoint",
		1, // 1ms
	},
	"endpoint.dataport.harakiriTimeout": ConfigValue{
		10 * 1000,
		"timeout in milliseconds, after which endpoint will commit harakiri " +
			"if not activity",
		10 * 1000, //10s
	},
	"endpoint.dataport.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload length, in bytes, for transmission data from " +
			"router to downstream client",
		1000 * 1024, // bytes
	},
	// indexer dataport parameters
	"projector.dataport.indexer.genServerChanSize": ConfigValue{
		64,
		"request channel size of indexer dataport's gen-server routine",
		64,
	},
	"projector.dataport.indexer.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload length, in bytes, for receiving data from router",
		1000 * 1024, // bytes
	},
	"projector.dataport.indexer.tcpReadDeadline": ConfigValue{
		10 * 1000,
		"timeout, in milliseconds, while reading from socket",
		10 * 1000, // 10s
	},
	// indexer queryport configuration
	"queryport.indexer.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload, in bytes, for receiving data from client",
		1000 * 1024,
	},
	"queryport.indexer.readDeadline": ConfigValue{
		4000,
		"timeout, in milliseconds, is timeout while reading from socket",
		4000,
	},
	"queryport.indexer.writeDeadline": ConfigValue{
		4000,
		"timeout, in milliseconds, is timeout while writing to socket",
		4000,
	},
	"queryport.indexer.pageSize": ConfigValue{
		1,
		"number of index-entries that shall be returned as single payload",
		1,
	},
	"queryport.indexer.streamChanSize": ConfigValue{
		16,
		"size of the buffered channels used to stream request and response.",
		16,
	},
	// queryport client configuration
	"queryport.client.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload, in bytes, for receiving data from server",
		1000 * 1024,
	},
	"queryport.client.readDeadline": ConfigValue{
		300000,
		"timeout, in milliseconds, is timeout while reading from socket",
		300000,
	},
	"queryport.client.writeDeadline": ConfigValue{
		4000,
		"timeout, in milliseconds, is timeout while writing to socket",
		4000,
	},
	"queryport.client.poolSize": ConfigValue{
		2,
		"number simultaneous active connections connections in a pool",
		2,
	},
	"queryport.client.poolOverflow": ConfigValue{
		4,
		"maximum number of connections in a pool",
		4,
	},
	"queryport.client.connPoolTimeout": ConfigValue{
		1000,
		"timeout, in milliseconds, is timeout for retrieving a connection " +
			"from the pool",
		1000,
	},
	"queryport.client.connPoolAvailWaitTimeout": ConfigValue{
		1,
		"timeout, in milliseconds, to wait for an existing connection " +
			"from the pool before considering the creation of a new one",
		1,
	},
	"indexer.scanTimeout": ConfigValue{
		12000,
		"timeout, in milliseconds, timeout for index scan processing",
		12000,
	},
	"indexer.adminPort": ConfigValue{
		"9100",
		"port for index ddl and status operations",
		"9100",
	},
	"indexer.scanPort": ConfigValue{
		"9101",
		"port for index scan operations",
		"9101",
	},
	"indexer.streamInitPort": ConfigValue{
		"9102",
		"port for inital build stream",
		"9102",
	},
	"indexer.streamCatchupPort": ConfigValue{
		"9103",
		"port for catchup stream",
		"9103",
	},
	"indexer.streamMaintPort": ConfigValue{
		"9104",
		"port for maintenance stream",
		"9104",
	},
	"indexer.clusterAddr": ConfigValue{
		"127.0.0.1:8091",
		"Local cluster manager address",
		"127.0.0.1:8091",
	},
	"indexer.numVbuckets": ConfigValue{
		1024,
		"Number of vbuckets",
		1024,
	},
	"indexer.enableManager": ConfigValue{
		false,
		"Enable index manager",
		false,
	},
	"indexer.storage_dir": ConfigValue{
		"./",
		"Index file storage directory",
		"./",
	},
	"indexer.compaction.interval": ConfigValue{
		60,
		"Compaction poll interval in seconds",
		60,
	},
	"indexer.compaction.minFrag": ConfigValue{
		30,
		"Compaction fragmentation threshold percentage",
		30,
	},
	"indexer.compaction.minSize": ConfigValue{
		uint64(1024 * 1024),
		"Compaction min file size",
		uint64(1024 * 1024),
	},
}

// NewConfig from another
// Config object or from map[string]interface{} object
// or from []byte slice, a byte-slice of JSON string.
func NewConfig(data interface{}) (Config, error) {
	config := SystemConfig.Clone()
	switch v := data.(type) {
	case Config: // Clone
		for key, value := range v {
			config.Set(key, value)
		}

	case map[string]interface{}: // transform
		for key, value := range v {
			config.SetValue(key, value)
		}

	case []byte: // parse JSON
		m := make(map[string]interface{})
		if err := json.Unmarshal(v, m); err != nil {
			return nil, err
		}
		for key, value := range m {
			config.SetValue(key, value)
		}

	default:
		return nil, nil
	}

	return config, nil
}

// Clone a new config object.
func (config Config) Clone() Config {
	clone := make(Config)
	for key, value := range config {
		clone[key] = value
	}
	return clone
}

// Override will clone `config` object and update parameters with
// values from `others` instance.
func (config Config) Override(others ...Config) Config {
	newconfig := config.Clone()
	for _, other := range others {
		for key, cv := range other {
			ocv, ok := newconfig[key]
			if !ok {
				ocv = cv
			} else {
				ocv.Value = cv.Value
			}
			config[key] = ocv
		}
	}
	return config
}

// SectionConfig will create a new config object with parameters
// starting with `prefix`. If `trim` is true, then config
// parameter will be trimmed with the prefix string.
func (config Config) SectionConfig(prefix string, trim bool) Config {
	section := make(Config)
	for key, value := range config {
		if strings.HasPrefix(key, prefix) {
			if trim {
				section[strings.TrimPrefix(key, prefix)] = value
			} else {
				section[key] = value
			}
		}
	}
	return section
}

// Set ConfigValue for parameter. Mutates the config object.
func (config Config) Set(key string, cv ConfigValue) Config {
	config[key] = cv
	return config
}

// SetValue config parameter with value. Mutates the config object.
func (config Config) SetValue(key string, value interface{}) Config {
	cv := config[key]
	cv.Value = value
	config[key] = cv
	return config
}

// Int assumes config value is an integer and returns the same.
func (cv ConfigValue) Int() int {
	return cv.Value.(int)
}

// Uint64 assumes config value is 64-bit integer and returns the same.
func (cv ConfigValue) Uint64() uint64 {
	return cv.Value.(uint64)
}

// String assumes config value is a string and returns the same.
func (cv ConfigValue) String() string {
	return cv.Value.(string)
}

// Strings assumes config value is comma separated string items.
func (cv ConfigValue) Strings() []string {
	ss := make([]string, 0)
	for _, s := range strings.Split(cv.Value.(string), ",") {
		s = strings.Trim(s, " \t\r\n")
		if len(s) > 0 {
			ss = append(ss, s)
		}
	}
	return ss
}

// Bool assumes config value is a Bool and returns the same.
func (cv ConfigValue) Bool() bool {
	return cv.Value.(bool)
}
<commit_msg>Increase ScanTimeout to 120 secs for large scale tests<commit_after>// Config is key, value map for system level and component configuration.
// Key is a string and represents a config parameter, and corresponding
// value is an interface{} that can be consumed using accessor methods
// based on the context of config-value.
//
// Config maps are immutable and newer versions can be created using accessor
// methods.
//
// Shape of config-parameter, the key string, is sequence of alpha-numeric
// characters separated by one or more '.' , eg,
//      "projector.adminport.readtimeout"

package common

import "encoding/json"
import "strings"

// Config is a key, value map with key always being a string
// represents a config-parameter.
type Config map[string]ConfigValue

// ConfigValue for each parameter.
type ConfigValue struct {
	Value      interface{}
	Help       string
	DefaultVal interface{}
}

// SystemConfig is default configuration for system and components.
// configuration parameters follow flat namespacing like,
//      "maxVbuckets"  for system-level config parameter
//      "projector.xxx" for projector component.
//      "projector.adminport.xxx" for adminport under projector component.
// etc...
var SystemConfig = Config{
	// system parameters
	"maxVbuckets": ConfigValue{
		1024,
		"number of vbuckets configured in KV",
		1024,
	},
	// log parameters
	// TODO: add configuration for log file-name and other types of writer.
	"log.ignore": ConfigValue{
		false,
		"ignores all logging, irrespective of the log-level",
		false,
	},
	"log.level": ConfigValue{
		"info",
		"logging level for the system",
		"info",
	},
	// projector parameters
	"projector.name": ConfigValue{
		"projector",
		"human readable name for this projector",
		"projector",
	},
	"projector.clusterAddr": ConfigValue{
		"localhost:9000",
		"KV cluster's address to be used by projector",
		"localhost:9000",
	},
	"projector.kvAddrs": ConfigValue{
		"127.0.0.1:9000",
		"Comma separated list of KV-address to read mutations, this need to " +
			"exactly match with KV-node's configured address",
		"127.0.0.1:9000",
	},
	"projector.colocate": ConfigValue{
		true,
		"Whether projector will be colocated with KV. In which case " +
			"`kvaddrs` specified above will be discarded",
		true,
	},
	"projector.routerEndpointFactory": ConfigValue{
		RouterEndpointFactory(nil),
		"RouterEndpointFactory callback to generate endpoint instances " +
			"to push data to downstream",
		RouterEndpointFactory(nil),
	},
	"projector.feedWaitStreamReqTimeout": ConfigValue{
		10 * 1000,
		"timeout, in milliseconds, to await a response for StreamRequest",
		10 * 1000,
	},
	"projector.feedWaitStreamEndTimeout": ConfigValue{
		10 * 1000,
		"timeout, in milliseconds, to await a response for StreamEnd",
		10 * 1000,
	},
	"projector.mutationChanSize": ConfigValue{
		10000,
		"channel size of projector's data path routine",
		10000,
	},
	"projector.feedChanSize": ConfigValue{
		100,
		"channel size for feed's control path and back path.",
		100,
	},
	"projector.vbucketSyncTimeout": ConfigValue{
		500,
		"timeout, in milliseconds, for sending periodic Sync messages.",
		500,
	},
	// projector adminport parameters
	"projector.adminport.name": ConfigValue{
		"projector.adminport",
		"human readable name for this adminport, must be supplied",
		"projector.adminport",
	},
	"projector.adminport.listenAddr": ConfigValue{
		"",
		"projector's adminport address listen for request.",
		"",
	},
	"projector.adminport.urlPrefix": ConfigValue{
		"/adminport/",
		"url prefix (script-path) for adminport used by projector",
		"/adminport/",
	},
	"projector.adminport.readTimeout": ConfigValue{
		0,
		"timeout in milliseconds, is read timeout for adminport http server " +
			"used by projector",
		0,
	},
	"projector.adminport.writeTimeout": ConfigValue{
		0,
		"timeout in milliseconds, is write timeout for adminport http server " +
			"used by projector",
		0,
	},
	"projector.adminport.maxHeaderBytes": ConfigValue{
		1 << 20, // 1 MegaByte
		"in bytes, is max. length of adminport http header " +
			"used by projector",
		1 << 20, // 1 MegaByte
	},
	// projector's adminport client
	"projector.client.retryInterval": ConfigValue{
		16,
		"retryInterval, in milliseconds, when connection refused by server",
		16,
	},
	"projector.client.maxRetries": ConfigValue{
		5,
		"maximum number of timest to retry",
		5,
	},
	"projector.client.exponentialBackoff": ConfigValue{
		2,
		"multiplying factor on retryInterval for every attempt with server",
		2,
	},
	// TODO: This configuration param is same as the above.
	"projector.client.urlPrefix": ConfigValue{
		"/adminport/",
		"url prefix (script-path) for adminport used by projector",
		"/adminport/",
	},
	// projector dataport client parameters
	// TODO: this configuration option should be tunnable for each feed.
	"endpoint.dataport.remoteBlock": ConfigValue{
		false,
		"should dataport endpoint block when remote is slow ?",
		false,
	},
	"endpoint.dataport.keyChanSize": ConfigValue{
		10000,
		"channel size of dataport endpoints data input",
		10000,
	},
	"endpoint.dataport.bufferSize": ConfigValue{
		100,
		"number of entries to buffer before flushing it, where each entry " +
			"is for a vbucket's set of mutations that was flushed by the endpoint.",
		100,
	},
	"endpoint.dataport.bufferTimeout": ConfigValue{
		1,
		"timeout in milliseconds, to flush vbucket-mutations from endpoint",
		1, // 1ms
	},
	"endpoint.dataport.harakiriTimeout": ConfigValue{
		10 * 1000,
		"timeout in milliseconds, after which endpoint will commit harakiri " +
			"if not activity",
		10 * 1000, //10s
	},
	"endpoint.dataport.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload length, in bytes, for transmission data from " +
			"router to downstream client",
		1000 * 1024, // bytes
	},
	// indexer dataport parameters
	"projector.dataport.indexer.genServerChanSize": ConfigValue{
		64,
		"request channel size of indexer dataport's gen-server routine",
		64,
	},
	"projector.dataport.indexer.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload length, in bytes, for receiving data from router",
		1000 * 1024, // bytes
	},
	"projector.dataport.indexer.tcpReadDeadline": ConfigValue{
		10 * 1000,
		"timeout, in milliseconds, while reading from socket",
		10 * 1000, // 10s
	},
	// indexer queryport configuration
	"queryport.indexer.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload, in bytes, for receiving data from client",
		1000 * 1024,
	},
	"queryport.indexer.readDeadline": ConfigValue{
		4000,
		"timeout, in milliseconds, is timeout while reading from socket",
		4000,
	},
	"queryport.indexer.writeDeadline": ConfigValue{
		4000,
		"timeout, in milliseconds, is timeout while writing to socket",
		4000,
	},
	"queryport.indexer.pageSize": ConfigValue{
		1,
		"number of index-entries that shall be returned as single payload",
		1,
	},
	"queryport.indexer.streamChanSize": ConfigValue{
		16,
		"size of the buffered channels used to stream request and response.",
		16,
	},
	// queryport client configuration
	"queryport.client.maxPayload": ConfigValue{
		1000 * 1024,
		"maximum payload, in bytes, for receiving data from server",
		1000 * 1024,
	},
	"queryport.client.readDeadline": ConfigValue{
		300000,
		"timeout, in milliseconds, is timeout while reading from socket",
		300000,
	},
	"queryport.client.writeDeadline": ConfigValue{
		4000,
		"timeout, in milliseconds, is timeout while writing to socket",
		4000,
	},
	"queryport.client.poolSize": ConfigValue{
		2,
		"number simultaneous active connections connections in a pool",
		2,
	},
	"queryport.client.poolOverflow": ConfigValue{
		4,
		"maximum number of connections in a pool",
		4,
	},
	"queryport.client.connPoolTimeout": ConfigValue{
		1000,
		"timeout, in milliseconds, is timeout for retrieving a connection " +
			"from the pool",
		1000,
	},
	"queryport.client.connPoolAvailWaitTimeout": ConfigValue{
		1,
		"timeout, in milliseconds, to wait for an existing connection " +
			"from the pool before considering the creation of a new one",
		1,
	},
	"indexer.scanTimeout": ConfigValue{
		120000,
		"timeout, in milliseconds, timeout for index scan processing",
		120000,
	},
	"indexer.adminPort": ConfigValue{
		"9100",
		"port for index ddl and status operations",
		"9100",
	},
	"indexer.scanPort": ConfigValue{
		"9101",
		"port for index scan operations",
		"9101",
	},
	"indexer.streamInitPort": ConfigValue{
		"9102",
		"port for inital build stream",
		"9102",
	},
	"indexer.streamCatchupPort": ConfigValue{
		"9103",
		"port for catchup stream",
		"9103",
	},
	"indexer.streamMaintPort": ConfigValue{
		"9104",
		"port for maintenance stream",
		"9104",
	},
	"indexer.clusterAddr": ConfigValue{
		"127.0.0.1:8091",
		"Local cluster manager address",
		"127.0.0.1:8091",
	},
	"indexer.numVbuckets": ConfigValue{
		1024,
		"Number of vbuckets",
		1024,
	},
	"indexer.enableManager": ConfigValue{
		false,
		"Enable index manager",
		false,
	},
	"indexer.storage_dir": ConfigValue{
		"./",
		"Index file storage directory",
		"./",
	},
	"indexer.compaction.interval": ConfigValue{
		60,
		"Compaction poll interval in seconds",
		60,
	},
	"indexer.compaction.minFrag": ConfigValue{
		30,
		"Compaction fragmentation threshold percentage",
		30,
	},
	"indexer.compaction.minSize": ConfigValue{
		uint64(1024 * 1024),
		"Compaction min file size",
		uint64(1024 * 1024),
	},
}

// NewConfig from another
// Config object or from map[string]interface{} object
// or from []byte slice, a byte-slice of JSON string.
func NewConfig(data interface{}) (Config, error) {
	config := SystemConfig.Clone()
	switch v := data.(type) {
	case Config: // Clone
		for key, value := range v {
			config.Set(key, value)
		}

	case map[string]interface{}: // transform
		for key, value := range v {
			config.SetValue(key, value)
		}

	case []byte: // parse JSON
		m := make(map[string]interface{})
		if err := json.Unmarshal(v, m); err != nil {
			return nil, err
		}
		for key, value := range m {
			config.SetValue(key, value)
		}

	default:
		return nil, nil
	}

	return config, nil
}

// Clone a new config object.
func (config Config) Clone() Config {
	clone := make(Config)
	for key, value := range config {
		clone[key] = value
	}
	return clone
}

// Override will clone `config` object and update parameters with
// values from `others` instance.
func (config Config) Override(others ...Config) Config {
	newconfig := config.Clone()
	for _, other := range others {
		for key, cv := range other {
			ocv, ok := newconfig[key]
			if !ok {
				ocv = cv
			} else {
				ocv.Value = cv.Value
			}
			config[key] = ocv
		}
	}
	return config
}

// SectionConfig will create a new config object with parameters
// starting with `prefix`. If `trim` is true, then config
// parameter will be trimmed with the prefix string.
func (config Config) SectionConfig(prefix string, trim bool) Config {
	section := make(Config)
	for key, value := range config {
		if strings.HasPrefix(key, prefix) {
			if trim {
				section[strings.TrimPrefix(key, prefix)] = value
			} else {
				section[key] = value
			}
		}
	}
	return section
}

// Set ConfigValue for parameter. Mutates the config object.
func (config Config) Set(key string, cv ConfigValue) Config {
	config[key] = cv
	return config
}

// SetValue config parameter with value. Mutates the config object.
func (config Config) SetValue(key string, value interface{}) Config {
	cv := config[key]
	cv.Value = value
	config[key] = cv
	return config
}

// Int assumes config value is an integer and returns the same.
func (cv ConfigValue) Int() int {
	return cv.Value.(int)
}

// Uint64 assumes config value is 64-bit integer and returns the same.
func (cv ConfigValue) Uint64() uint64 {
	return cv.Value.(uint64)
}

// String assumes config value is a string and returns the same.
func (cv ConfigValue) String() string {
	return cv.Value.(string)
}

// Strings assumes config value is comma separated string items.
func (cv ConfigValue) Strings() []string {
	ss := make([]string, 0)
	for _, s := range strings.Split(cv.Value.(string), ",") {
		s = strings.Trim(s, " \t\r\n")
		if len(s) > 0 {
			ss = append(ss, s)
		}
	}
	return ss
}

// Bool assumes config value is a Bool and returns the same.
func (cv ConfigValue) Bool() bool {
	return cv.Value.(bool)
}
<|endoftext|>


<commit_before>package main

import (
	"encoding/json"
	"log"
	"net/http"
	"strconv"

	"github.com/guregu/kami"
)

func main() {
	kami.Get("/contacts", getContacts)
	kami.Serve()
}

func getContacts(
	w http.ResponseWriter,
	r *http.Request,
) {
	defer func() {
		if err := recover(); err != nil {
			log.Print(err)
		}
	}()

	page, err := strconv.Atoi(r.FormValue("page"))
	if err != nil {
		page = 1
	}

	perPage, err := strconv.Atoi(r.FormValue("per_page"))
	if err != nil {
		perPage = 100
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	err = json.NewEncoder(w).Encode(
		NewContactQuery(page, perPage).All())

	if err != nil {
		log.Print(err)
	}
}
<commit_msg>[kami] Adopt handler for use one of the query objects<commit_after>package main

import (
	"bytes"
	"encoding/json"
	"log"
	"net/http"
	"strconv"
	"sync"

	"github.com/guregu/kami"
)

var bufPool = sync.Pool{
	New: func() interface{} {
		return new(bytes.Buffer)
	},
}

func init() {
	DBConn()
	// PgxDBConn()
	// PgDBConn()
}

func main() {
	kami.Get("/contacts", getContacts)
	kami.Serve()
}

func getContacts(
	w http.ResponseWriter,
	r *http.Request,
) {
	defer func() {
		if err := recover(); err != nil {
			log.Print(err)
		}
	}()

	page, err := strconv.Atoi(r.FormValue("page"))
	if err != nil {
		page = 1
	}

	perPage, err := strconv.Atoi(r.FormValue("per_page"))
	if err != nil {
		perPage = 100
	}

	w.Header().Set("Content-Type", "application/json; charset=utf-8")

	contacts := NewContactQuery(page, perPage).All()
	// contacts := NewPGXContactQuery(page, perPage).All()
	// contacts := NewPGContactQuery(page, perPage).All()
	buf := bufPool.Get().(*bytes.Buffer)
	err = json.NewEncoder(buf).Encode(contacts)
	w.Write(buf.Bytes())
	buf.Reset()
	bufPool.Put(buf)

	if err != nil {
		log.Print(err)
	}
}
<|endoftext|>


<commit_before>package jsonproperty

import (
	"appengine/datastore"
	"encoding/json"
	"reflect"
)

type JsonProperty map[string]interface{}

var jsonPropertyType = reflect.TypeOf(JsonProperty{})

// entity must be a pointer to a struct
func LoadJsonProperties(entity interface{}, c <-chan datastore.Property) (<-chan datastore.Property, error) {
	jsonProperties := map[string]reflect.Value{} // {name: value}

	// Builds jsonProperties
	value := reflect.ValueOf(entity).Elem()
	for i := 0; i < value.NumField(); i++ {
		value2 := value.Field(i)
		if value2.Type() == jsonPropertyType {
			jsonProperties[nameFromField(&value.Type().Field(i))] = value2
		}
	}

	// Builds return channel
	c2 := make(chan datastore.Property, value.NumField()-len(jsonProperties))
	defer close(c2)
	for property := range c {
		if jsonValue, ok := jsonProperties[property.Name]; ok {
			bytes := []byte(property.Value.(string))
			if err := json.Unmarshal(bytes, jsonValue.Addr().Interface()); err != nil {
				return c2, err
			}
		} else {
			c2 <- property
		}
	}

	return c2, nil
}

// entity must be a pointer to a struct
func SaveJsonProperties(entity interface{}, c chan<- datastore.Property) (chan<- datastore.Property, error) {
	value := reflect.ValueOf(entity).Elem()
	for i := 0; i < value.NumField(); i++ {
		value2 := value.Field(i)
		if value2.Type() == jsonPropertyType {
			bytes, err := json.Marshal(value2.Interface())
			if err != nil {
				return c, err
			}
			c <- datastore.Property{
				Name:  nameFromField(&value.Type().Field(i)),
				Value: string(bytes),
			}
		}
	}

	return c, nil
}

func nameFromField(f *reflect.StructField) string {
	if name := f.Tag.Get("jsonproperty"); name != "" {
		return name
	}
	return f.Name
}
<commit_msg>Fixing error<commit_after>package jsonproperty

import (
	"appengine/datastore"
	"encoding/json"
	"reflect"
)

type JsonProperty map[string]interface{}

var jsonPropertyType = reflect.TypeOf(JsonProperty{})

// entity must be a pointer to a struct
func LoadJsonProperties(entity interface{}, c <-chan datastore.Property) (<-chan datastore.Property, error) {
	jsonProperties := map[string]reflect.Value{} // {name: value}

	// Builds jsonProperties
	value := reflect.ValueOf(entity).Elem()
	for i := 0; i < value.NumField(); i++ {
		value2 := value.Field(i)
		if value2.Type() == jsonPropertyType {
			jsonProperties[nameFromField(value.Type().Field(i))] = value2
		}
	}

	// Builds return channel
	c2 := make(chan datastore.Property, value.NumField()-len(jsonProperties))
	defer close(c2)
	for property := range c {
		if jsonValue, ok := jsonProperties[property.Name]; ok {
			bytes := []byte(property.Value.(string))
			if err := json.Unmarshal(bytes, jsonValue.Addr().Interface()); err != nil {
				return c2, err
			}
		} else {
			c2 <- property
		}
	}

	return c2, nil
}

// entity must be a pointer to a struct
func SaveJsonProperties(entity interface{}, c chan<- datastore.Property) (chan<- datastore.Property, error) {
	value := reflect.ValueOf(entity).Elem()
	for i := 0; i < value.NumField(); i++ {
		value2 := value.Field(i)
		if value2.Type() == jsonPropertyType {
			bytes, err := json.Marshal(value2.Interface())
			if err != nil {
				return c, err
			}
			c <- datastore.Property{
				Name:  nameFromField(value.Type().Field(i)),
				Value: string(bytes),
			}
		}
	}

	return c, nil
}

func nameFromField(f reflect.StructField) string {
	if name := f.Tag.Get("jsonproperty"); name != "" {
		return name
	}
	return f.Name
}
<|endoftext|>