<commit_before>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/cyfdecyf/bufio"
	"io/ioutil"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"
)

func init() {
	rand.Seed(time.Now().Unix())
}

// VisitCnt and SiteStat are used to track how many times a site is visited.
// With this information: COW knows which sites are frequently visited, and
// judging whether a site is blocked or not is more reliable.

const (
	directDelta  = 15
	blockedDelta = 10
	maxCnt       = 100 // no protect to update visit cnt, smaller value is unlikely to overflow
	userCnt      = -1  // this represents user specified host or domain
)

type siteVisitMethod int

// minus operation on visit count may get negative value, so use signed int
type vcntint int8

type Date time.Time

const dateLayout = "2006-01-02"

func (d Date) MarshalJSON() ([]byte, error) {
	return []byte("\"" + time.Time(d).Format(dateLayout) + "\""), nil
}

func (d *Date) UnmarshalJSON(input []byte) error {
	if len(input) != len(dateLayout)+2 {
		return errors.New(fmt.Sprintf("unmarshaling date: invalid input %s", string(input)))
	}
	input = input[1 : len(dateLayout)+1]
	t, err := time.Parse(dateLayout, string(input))
	*d = Date(t)
	return err
}

// COW don't need very accurate visit count, so update to visit count value is
// not protected.
type VisitCnt struct {
	Direct    vcntint   `json:"direct"`
	Blocked   vcntint   `json:"block"`
	Recent    Date      `json:"recent"`
	rUpdated  bool      // whether Recent is updated, we only need date precision
	blockedOn time.Time // when is the site last blocked
}

func newVisitCnt(direct, blocked vcntint) *VisitCnt {
	return &VisitCnt{direct, blocked, Date(time.Now()), true, zeroTime}
}

func newVisitCntWithTime(direct, blocked vcntint, t time.Time) *VisitCnt {
	return &VisitCnt{direct, blocked, Date(t), true, zeroTime}
}

func (vc *VisitCnt) userSpecified() bool {
	return vc.Blocked == userCnt || vc.Direct == userCnt
}

const siteStaleThreshold = 10 * 24 * time.Hour

func (vc *VisitCnt) isStale() bool {
	return time.Now().Sub(time.Time(vc.Recent)) > siteStaleThreshold
}

// shouldNotSave returns true if the a VisitCnt is not visited for a long time
// (several days) or is specified by user.
func (vc *VisitCnt) shouldNotSave() bool {
	return vc.userSpecified() || vc.isStale() || (vc.Blocked == 0 && vc.Direct == 0)
}

const tmpBlockedTimeout = 2 * time.Minute

func (vc *VisitCnt) AsTempBlocked() bool {
	return time.Now().Sub(vc.blockedOn) < tmpBlockedTimeout
}

func (vc *VisitCnt) AsDirect() bool {
	return (vc.Direct == userCnt) || (vc.Direct-vc.Blocked >= directDelta && vc.Blocked == 0)
}

func (vc *VisitCnt) AsBlocked() bool {
	if vc.Blocked == userCnt || vc.AsTempBlocked() {
		return true
	}
	// add some randomness to fix mistake
	delta := vc.Blocked - vc.Direct
	return delta >= blockedDelta && rand.Intn(int(delta)) != 0
}

func (vc *VisitCnt) AlwaysDirect() bool {
	return vc.Direct == userCnt
}

func (vc *VisitCnt) AlwaysBlocked() bool {
	return vc.Blocked == userCnt
}

func (vc *VisitCnt) OnceBlocked() bool {
	return vc.Blocked > 0 || vc.AlwaysBlocked() || vc.AsTempBlocked()
}

func (vc *VisitCnt) tempBlocked() {
	vc.blockedOn = time.Now()
}

// time.Time is composed of 3 fields, so need lock to protect update. As
// update of last visit is not frequent (at most once for each domain), use a
// global lock to avoid associating a lock to each VisitCnt.
var visitLock sync.Mutex

// visit updates visit cnt
func (vc *VisitCnt) visit(inc *vcntint) {
	if *inc < maxCnt {
		*inc++
	}
	// Because of concurrent update, possible for *inc to overflow and become
	// negative, but very unlikely.
	if *inc > maxCnt || *inc < 0 {
		*inc = maxCnt
	}

	if !vc.rUpdated {
		vc.rUpdated = true
		visitLock.Lock()
		vc.Recent = Date(time.Now())
		visitLock.Unlock()
	}
}

func (vc *VisitCnt) DirectVisit() {
	if networkBad() || vc.userSpecified() {
		return
	}
	// one successful direct visit probably means the site is not actually
	// blocked
	vc.visit(&vc.Direct)
	vc.Blocked = 0
}

func (vc *VisitCnt) BlockedVisit() {
	if networkBad() || vc.userSpecified() {
		return
	}
	// When a site changes from direct to blocked by GFW, COW should learn
	// this quickly and remove it from the PAC ASAP. So change direct to 0
	// once there's a single blocked visit, this ensures the site is removed
	// upon the next PAC update.
	vc.visit(&vc.Blocked)
	vc.Direct = 0
}

type SiteStat struct {
	Update Date                 `json:"update"`
	Vcnt   map[string]*VisitCnt `json:"site_info"` // Vcnt uses host as key
	vcLock sync.RWMutex

	// Whether a domain has blocked host. Used to avoid considering a domain as
	// direct though it has blocked hosts.
	hasBlockedHost map[string]bool
	hbhLock        sync.RWMutex
}

func newSiteStat() *SiteStat {
	return &SiteStat{
		Vcnt:           map[string]*VisitCnt{},
		hasBlockedHost: map[string]bool{},
	}
}

func (ss *SiteStat) get(s string) *VisitCnt {
	ss.vcLock.RLock()
	Vcnt, ok := ss.Vcnt[s]
	ss.vcLock.RUnlock()
	if ok {
		return Vcnt
	}
	return nil
}

func (ss *SiteStat) create(s string) (vcnt *VisitCnt) {
	vcnt = newVisitCnt(0, 0)
	ss.vcLock.Lock()
	ss.Vcnt[s] = vcnt
	ss.vcLock.Unlock()
	return
}

// Caller should guarantee that always direct url does not attempt
// blocked visit.
func (ss *SiteStat) TempBlocked(url *URL) {
	debug.Printf("%s temp blocked\n", url.Host)

	vcnt := ss.get(url.Host)
	if vcnt == nil {
		panic("TempBlocked should always get existing visitCnt")
	}
	vcnt.tempBlocked()

	// Mistakenly consider a partial blocked domain as direct will make that
	// domain into PAC and never have a chance to correct the error.
	// Once using blocked visit, a host is considered to maybe blocked even if
	// it's block visit count decrease to 0. As hasBlockedHost is not saved,
	// upon next start up of COW, the information will reflect the current
	// status of that host.
	ss.hbhLock.RLock()
	t := ss.hasBlockedHost[url.Domain]
	ss.hbhLock.RUnlock()
	if !t {
		ss.hbhLock.Lock()
		ss.hasBlockedHost[url.Domain] = true
		ss.hbhLock.Unlock()
	}
}

var alwaysDirectVisitCnt = newVisitCnt(userCnt, 0)

func (ss *SiteStat) GetVisitCnt(url *URL) (vcnt *VisitCnt) {
	if !hasParentProxy() { // no way to retry, so always visit directly
		return alwaysDirectVisitCnt
	}
	if url.Domain == "" { // simple host or private ip
		return alwaysDirectVisitCnt
	}
	if vcnt = ss.get(url.Host); vcnt != nil {
		return
	}
	if len(url.Domain) != len(url.Host) {
		if dmcnt := ss.get(url.Domain); dmcnt != nil && dmcnt.userSpecified() {
			// if the domain is not specified by user, should create a new host
			// visitCnt
			return dmcnt
		}
	}
	return ss.create(url.Host)
}

func (ss *SiteStat) store(file string) (err error) {
	if err = mkConfigDir(); err != nil {
		return
	}

	now := time.Now()
	var savedSS *SiteStat
	if ss.Update == Date(zeroTime) {
		ss.Update = Date(time.Now())
	}
	if now.Sub(time.Time(ss.Update)) > siteStaleThreshold {
		// Not updated for a long time, don't drop any record
		savedSS = ss
		// Changing update time too fast will also drop useful record
		savedSS.Update = Date(time.Time(ss.Update).Add(siteStaleThreshold / 2))
		if time.Time(savedSS.Update).After(now) {
			savedSS.Update = Date(now)
		}
	} else {
		savedSS = newSiteStat()
		savedSS.Update = Date(now)
		ss.vcLock.RLock()
		for site, vcnt := range ss.Vcnt {
			if vcnt.shouldNotSave() {
				continue
			}
			savedSS.Vcnt[site] = vcnt
		}
		ss.vcLock.RUnlock()
	}

	b, err := json.MarshalIndent(savedSS, "", "\t")
	if err != nil {
		errl.Println("Error marshalling site stat:", err)
		panic("internal error: error marshalling site")
	}

	f, err := os.Create(file)
	if err != nil {
		errl.Println("Can't create stat file:", err)
		return
	}
	defer f.Close()
	if _, err = f.Write(b); err != nil {
		errl.Println("Error writing stat file:", err)
		return
	}
	return
}

func (ss *SiteStat) loadList(lst []string, direct, blocked vcntint) {
	for _, d := range lst {
		ss.Vcnt[d] = newVisitCntWithTime(direct, blocked, zeroTime)
	}
}

func (ss *SiteStat) loadBuiltinList() {
	ss.loadList(blockedDomainList, 0, userCnt)
	ss.loadList(directDomainList, userCnt, 0)
}

func (ss *SiteStat) loadUserList() {
	if directList, err := loadSiteList(dsFile.alwaysDirect); err == nil {
		ss.loadList(directList, userCnt, 0)
	}
	if blockedList, err := loadSiteList(dsFile.alwaysBlocked); err == nil {
		ss.loadList(blockedList, 0, userCnt)
	}
}

// Filter sites covered by user specified domains, also filter out stale
// sites.
func (ss *SiteStat) filterSites() {
	// It's not safe to remove element while iterating over a map.
	var removeSites []string

	// find what to remove first
	ss.vcLock.RLock()
	for site, vcnt := range ss.Vcnt {
		if vcnt.userSpecified() {
			continue
		}
		if vcnt.isStale() {
			removeSites = append(removeSites, site)
			continue
		}
		var dmcnt *VisitCnt
		domain := host2Domain(site)
		if domain != site {
			dmcnt = ss.get(domain)
		}
		if dmcnt != nil && dmcnt.userSpecified() {
			removeSites = append(removeSites, site)
		}
	}
	ss.vcLock.RUnlock()

	// do remove
	ss.vcLock.Lock()
	for _, site := range removeSites {
		delete(ss.Vcnt, site)
	}
	ss.vcLock.Unlock()
}

func (ss *SiteStat) load(file string) (err error) {
	defer func() {
		// load builtin list first, so user list can override builtin
		ss.loadBuiltinList()
		ss.loadUserList()
		ss.filterSites()
		for host, vcnt := range ss.Vcnt {
			if vcnt.OnceBlocked() {
				ss.hasBlockedHost[host2Domain(host)] = true
			}
		}
	}()
	var exists bool
	if exists, err = isFileExists(file); err != nil {
		fmt.Println("Error loading stat:", err)
		return
	}
	if !exists {
		return
	}
	var f *os.File
	if f, err = os.Open(file); err != nil {
		fmt.Printf("Error opening site stat %s: %v\n", file, err)
		return
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		fmt.Println("Error reading site stat:", err)
		return
	}
	if err = json.Unmarshal(b, ss); err != nil {
		fmt.Println("Error decoding site stat:", err)
		return
	}
	return
}

func (ss *SiteStat) GetDirectList() []string {
	lst := make([]string, 0)
	// anyway to do more fine grained locking?
	ss.vcLock.RLock()
	for site, vc := range ss.Vcnt {
		if ss.hasBlockedHost[host2Domain(site)] {
			continue
		}
		if vc.AsDirect() {
			lst = append(lst, site)
		}
	}
	ss.vcLock.RUnlock()
	return lst
}

var siteStat = newSiteStat()

func initSiteStat() {
	if err := siteStat.load(dsFile.stat); err != nil {
		os.Exit(1)
	}
	// dump site stat once every hour, so we don't always need to close cow to
	// get updated stat
	go func() {
		for {
			time.Sleep(time.Hour)
			storeSiteStat()
		}
	}()
}

func storeSiteStat() {
	siteStat.store(dsFile.stat)
}

func loadSiteList(fpath string) (lst []string, err error) {
	var exists bool
	if exists, err = isFileExists(fpath); err != nil {
		errl.Printf("Error loading domaint list: %v\n", err)
	}
	if !exists {
		return
	}
	f, err := os.Open(fpath)
	if err != nil {
		errl.Println("Error opening domain list:", err)
		return
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	lst = make([]string, 0)
	for scanner.Scan() {
		site := strings.TrimSpace(scanner.Text())
		if site == "" {
			continue
		}
		lst = append(lst, site)
	}
	if scanner.Err() != nil {
		errl.Printf("Error reading domain list %s: %v\n", fpath, scanner.Err())
	}
	return lst, scanner.Err()
}
<commit_msg>Fix AsDirect returns false for new site.<commit_after>package main

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/cyfdecyf/bufio"
	"io/ioutil"
	"math/rand"
	"os"
	"strings"
	"sync"
	"time"
)

func init() {
	rand.Seed(time.Now().Unix())
}

// VisitCnt and SiteStat are used to track how many times a site is visited.
// With this information: COW knows which sites are frequently visited, and
// judging whether a site is blocked or not is more reliable.

const (
	directDelta  = 15
	blockedDelta = 10
	maxCnt       = 100 // no protect to update visit cnt, smaller value is unlikely to overflow
	userCnt      = -1  // this represents user specified host or domain
)

type siteVisitMethod int

// minus operation on visit count may get negative value, so use signed int
type vcntint int8

type Date time.Time

const dateLayout = "2006-01-02"

func (d Date) MarshalJSON() ([]byte, error) {
	return []byte("\"" + time.Time(d).Format(dateLayout) + "\""), nil
}

func (d *Date) UnmarshalJSON(input []byte) error {
	if len(input) != len(dateLayout)+2 {
		return errors.New(fmt.Sprintf("unmarshaling date: invalid input %s", string(input)))
	}
	input = input[1 : len(dateLayout)+1]
	t, err := time.Parse(dateLayout, string(input))
	*d = Date(t)
	return err
}

// COW don't need very accurate visit count, so update to visit count value is
// not protected.
type VisitCnt struct {
	Direct    vcntint   `json:"direct"`
	Blocked   vcntint   `json:"block"`
	Recent    Date      `json:"recent"`
	rUpdated  bool      // whether Recent is updated, we only need date precision
	blockedOn time.Time // when is the site last blocked
}

func newVisitCnt(direct, blocked vcntint) *VisitCnt {
	return &VisitCnt{direct, blocked, Date(time.Now()), true, zeroTime}
}

func newVisitCntWithTime(direct, blocked vcntint, t time.Time) *VisitCnt {
	return &VisitCnt{direct, blocked, Date(t), true, zeroTime}
}

func (vc *VisitCnt) userSpecified() bool {
	return vc.Blocked == userCnt || vc.Direct == userCnt
}

const siteStaleThreshold = 10 * 24 * time.Hour

func (vc *VisitCnt) isStale() bool {
	return time.Now().Sub(time.Time(vc.Recent)) > siteStaleThreshold
}

// shouldNotSave returns true if the a VisitCnt is not visited for a long time
// (several days) or is specified by user.
func (vc *VisitCnt) shouldNotSave() bool {
	return vc.userSpecified() || vc.isStale() || (vc.Blocked == 0 && vc.Direct == 0)
}

const tmpBlockedTimeout = 2 * time.Minute

func (vc *VisitCnt) AsTempBlocked() bool {
	return time.Now().Sub(vc.blockedOn) < tmpBlockedTimeout
}

func (vc *VisitCnt) AsDirect() bool {
	return (vc.Blocked == 0) || (vc.Direct-vc.Blocked >= directDelta) || vc.AlwaysDirect()
}

func (vc *VisitCnt) AsBlocked() bool {
	if vc.Blocked == userCnt || vc.AsTempBlocked() {
		return true
	}
	// add some randomness to fix mistake
	delta := vc.Blocked - vc.Direct
	return delta >= blockedDelta && rand.Intn(int(delta)) != 0
}

func (vc *VisitCnt) AlwaysDirect() bool {
	return vc.Direct == userCnt
}

func (vc *VisitCnt) AlwaysBlocked() bool {
	return vc.Blocked == userCnt
}

func (vc *VisitCnt) OnceBlocked() bool {
	return vc.Blocked > 0 || vc.AlwaysBlocked() || vc.AsTempBlocked()
}

func (vc *VisitCnt) tempBlocked() {
	vc.blockedOn = time.Now()
}

// time.Time is composed of 3 fields, so need lock to protect update. As
// update of last visit is not frequent (at most once for each domain), use a
// global lock to avoid associating a lock to each VisitCnt.
var visitLock sync.Mutex

// visit updates visit cnt
func (vc *VisitCnt) visit(inc *vcntint) {
	if *inc < maxCnt {
		*inc++
	}
	// Because of concurrent update, possible for *inc to overflow and become
	// negative, but very unlikely.
	if *inc > maxCnt || *inc < 0 {
		*inc = maxCnt
	}

	if !vc.rUpdated {
		vc.rUpdated = true
		visitLock.Lock()
		vc.Recent = Date(time.Now())
		visitLock.Unlock()
	}
}

func (vc *VisitCnt) DirectVisit() {
	if networkBad() || vc.userSpecified() {
		return
	}
	// one successful direct visit probably means the site is not actually
	// blocked
	vc.visit(&vc.Direct)
	vc.Blocked = 0
}

func (vc *VisitCnt) BlockedVisit() {
	if networkBad() || vc.userSpecified() {
		return
	}
	// When a site changes from direct to blocked by GFW, COW should learn
	// this quickly and remove it from the PAC ASAP. So change direct to 0
	// once there's a single blocked visit, this ensures the site is removed
	// upon the next PAC update.
	vc.visit(&vc.Blocked)
	vc.Direct = 0
}

type SiteStat struct {
	Update Date                 `json:"update"`
	Vcnt   map[string]*VisitCnt `json:"site_info"` // Vcnt uses host as key
	vcLock sync.RWMutex

	// Whether a domain has blocked host. Used to avoid considering a domain as
	// direct though it has blocked hosts.
	hasBlockedHost map[string]bool
	hbhLock        sync.RWMutex
}

func newSiteStat() *SiteStat {
	return &SiteStat{
		Vcnt:           map[string]*VisitCnt{},
		hasBlockedHost: map[string]bool{},
	}
}

func (ss *SiteStat) get(s string) *VisitCnt {
	ss.vcLock.RLock()
	Vcnt, ok := ss.Vcnt[s]
	ss.vcLock.RUnlock()
	if ok {
		return Vcnt
	}
	return nil
}

func (ss *SiteStat) create(s string) (vcnt *VisitCnt) {
	vcnt = newVisitCnt(0, 0)
	ss.vcLock.Lock()
	ss.Vcnt[s] = vcnt
	ss.vcLock.Unlock()
	return
}

// Caller should guarantee that always direct url does not attempt
// blocked visit.
func (ss *SiteStat) TempBlocked(url *URL) {
	debug.Printf("%s temp blocked\n", url.Host)

	vcnt := ss.get(url.Host)
	if vcnt == nil {
		panic("TempBlocked should always get existing visitCnt")
	}
	vcnt.tempBlocked()

	// Mistakenly consider a partial blocked domain as direct will make that
	// domain into PAC and never have a chance to correct the error.
	// Once using blocked visit, a host is considered to maybe blocked even if
	// it's block visit count decrease to 0. As hasBlockedHost is not saved,
	// upon next start up of COW, the information will reflect the current
	// status of that host.
	ss.hbhLock.RLock()
	t := ss.hasBlockedHost[url.Domain]
	ss.hbhLock.RUnlock()
	if !t {
		ss.hbhLock.Lock()
		ss.hasBlockedHost[url.Domain] = true
		ss.hbhLock.Unlock()
	}
}

var alwaysDirectVisitCnt = newVisitCnt(userCnt, 0)

func (ss *SiteStat) GetVisitCnt(url *URL) (vcnt *VisitCnt) {
	if !hasParentProxy() { // no way to retry, so always visit directly
		return alwaysDirectVisitCnt
	}
	if url.Domain == "" { // simple host or private ip
		return alwaysDirectVisitCnt
	}
	if vcnt = ss.get(url.Host); vcnt != nil {
		return
	}
	if len(url.Domain) != len(url.Host) {
		if dmcnt := ss.get(url.Domain); dmcnt != nil && dmcnt.userSpecified() {
			// if the domain is not specified by user, should create a new host
			// visitCnt
			return dmcnt
		}
	}
	return ss.create(url.Host)
}

func (ss *SiteStat) store(file string) (err error) {
	if err = mkConfigDir(); err != nil {
		return
	}

	now := time.Now()
	var savedSS *SiteStat
	if ss.Update == Date(zeroTime) {
		ss.Update = Date(time.Now())
	}
	if now.Sub(time.Time(ss.Update)) > siteStaleThreshold {
		// Not updated for a long time, don't drop any record
		savedSS = ss
		// Changing update time too fast will also drop useful record
		savedSS.Update = Date(time.Time(ss.Update).Add(siteStaleThreshold / 2))
		if time.Time(savedSS.Update).After(now) {
			savedSS.Update = Date(now)
		}
	} else {
		savedSS = newSiteStat()
		savedSS.Update = Date(now)
		ss.vcLock.RLock()
		for site, vcnt := range ss.Vcnt {
			if vcnt.shouldNotSave() {
				continue
			}
			savedSS.Vcnt[site] = vcnt
		}
		ss.vcLock.RUnlock()
	}

	b, err := json.MarshalIndent(savedSS, "", "\t")
	if err != nil {
		errl.Println("Error marshalling site stat:", err)
		panic("internal error: error marshalling site")
	}

	f, err := os.Create(file)
	if err != nil {
		errl.Println("Can't create stat file:", err)
		return
	}
	defer f.Close()
	if _, err = f.Write(b); err != nil {
		errl.Println("Error writing stat file:", err)
		return
	}
	return
}

func (ss *SiteStat) loadList(lst []string, direct, blocked vcntint) {
	for _, d := range lst {
		ss.Vcnt[d] = newVisitCntWithTime(direct, blocked, zeroTime)
	}
}

func (ss *SiteStat) loadBuiltinList() {
	ss.loadList(blockedDomainList, 0, userCnt)
	ss.loadList(directDomainList, userCnt, 0)
}

func (ss *SiteStat) loadUserList() {
	if directList, err := loadSiteList(dsFile.alwaysDirect); err == nil {
		ss.loadList(directList, userCnt, 0)
	}
	if blockedList, err := loadSiteList(dsFile.alwaysBlocked); err == nil {
		ss.loadList(blockedList, 0, userCnt)
	}
}

// Filter sites covered by user specified domains, also filter out stale
// sites.
func (ss *SiteStat) filterSites() {
	// It's not safe to remove element while iterating over a map.
	var removeSites []string

	// find what to remove first
	ss.vcLock.RLock()
	for site, vcnt := range ss.Vcnt {
		if vcnt.userSpecified() {
			continue
		}
		if vcnt.isStale() {
			removeSites = append(removeSites, site)
			continue
		}
		var dmcnt *VisitCnt
		domain := host2Domain(site)
		if domain != site {
			dmcnt = ss.get(domain)
		}
		if dmcnt != nil && dmcnt.userSpecified() {
			removeSites = append(removeSites, site)
		}
	}
	ss.vcLock.RUnlock()

	// do remove
	ss.vcLock.Lock()
	for _, site := range removeSites {
		delete(ss.Vcnt, site)
	}
	ss.vcLock.Unlock()
}

func (ss *SiteStat) load(file string) (err error) {
	defer func() {
		// load builtin list first, so user list can override builtin
		ss.loadBuiltinList()
		ss.loadUserList()
		ss.filterSites()
		for host, vcnt := range ss.Vcnt {
			if vcnt.OnceBlocked() {
				ss.hasBlockedHost[host2Domain(host)] = true
			}
		}
	}()
	var exists bool
	if exists, err = isFileExists(file); err != nil {
		fmt.Println("Error loading stat:", err)
		return
	}
	if !exists {
		return
	}
	var f *os.File
	if f, err = os.Open(file); err != nil {
		fmt.Printf("Error opening site stat %s: %v\n", file, err)
		return
	}
	b, err := ioutil.ReadAll(f)
	if err != nil {
		fmt.Println("Error reading site stat:", err)
		return
	}
	if err = json.Unmarshal(b, ss); err != nil {
		fmt.Println("Error decoding site stat:", err)
		return
	}
	return
}

func (ss *SiteStat) GetDirectList() []string {
	lst := make([]string, 0)
	// anyway to do more fine grained locking?
	ss.vcLock.RLock()
	for site, vc := range ss.Vcnt {
		if ss.hasBlockedHost[host2Domain(site)] {
			continue
		}
		if vc.AsDirect() {
			lst = append(lst, site)
		}
	}
	ss.vcLock.RUnlock()
	return lst
}

var siteStat = newSiteStat()

func initSiteStat() {
	if err := siteStat.load(dsFile.stat); err != nil {
		os.Exit(1)
	}
	// dump site stat once every hour, so we don't always need to close cow to
	// get updated stat
	go func() {
		for {
			time.Sleep(time.Hour)
			storeSiteStat()
		}
	}()
}

func storeSiteStat() {
	siteStat.store(dsFile.stat)
}

func loadSiteList(fpath string) (lst []string, err error) {
	var exists bool
	if exists, err = isFileExists(fpath); err != nil {
		errl.Printf("Error loading domaint list: %v\n", err)
	}
	if !exists {
		return
	}
	f, err := os.Open(fpath)
	if err != nil {
		errl.Println("Error opening domain list:", err)
		return
	}
	defer f.Close()

	scanner := bufio.NewScanner(f)
	lst = make([]string, 0)
	for scanner.Scan() {
		site := strings.TrimSpace(scanner.Text())
		if site == "" {
			continue
		}
		lst = append(lst, site)
	}
	if scanner.Err() != nil {
		errl.Printf("Error reading domain list %s: %v\n", fpath, scanner.Err())
	}
	return lst, scanner.Err()
}
<|endoftext|>


<commit_before>package libp2p

// This file contains all libp2p configuration options (except the defaults,
// those are in defaults.go).

import (
	"fmt"
	"net"

	"github.com/libp2p/go-libp2p-core/connmgr"
	"github.com/libp2p/go-libp2p-core/crypto"
	"github.com/libp2p/go-libp2p-core/metrics"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/peerstore"
	"github.com/libp2p/go-libp2p-core/pnet"

	circuit "github.com/libp2p/go-libp2p-circuit"
	config "github.com/libp2p/go-libp2p/config"
	bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
	autorelay "github.com/libp2p/go-libp2p/p2p/host/relay"

	filter "github.com/libp2p/go-maddr-filter"
	ma "github.com/multiformats/go-multiaddr"
)

// ListenAddrStrings configures libp2p to listen on the given (unparsed)
// addresses.
func ListenAddrStrings(s ...string) Option {
	return func(cfg *Config) error {
		for _, addrstr := range s {
			a, err := ma.NewMultiaddr(addrstr)
			if err != nil {
				return err
			}
			cfg.ListenAddrs = append(cfg.ListenAddrs, a)
		}
		return nil
	}
}

// ListenAddrs configures libp2p to listen on the given addresses.
func ListenAddrs(addrs ...ma.Multiaddr) Option {
	return func(cfg *Config) error {
		cfg.ListenAddrs = append(cfg.ListenAddrs, addrs...)
		return nil
	}
}

// Security configures libp2p to use the given security transport (or transport
// constructor).
//
// Name is the protocol name.
//
// The transport can be a constructed security.Transport or a function taking
// any subset of this libp2p node's:
// * Public key
// * Private key
// * Peer ID
// * Host
// * Network
// * Peerstore
func Security(name string, tpt interface{}) Option {
	stpt, err := config.SecurityConstructor(tpt)
	err = traceError(err, 1)
	return func(cfg *Config) error {
		if err != nil {
			return err
		}
		if cfg.Insecure {
			return fmt.Errorf("cannot use security transports with an insecure libp2p configuration")
		}
		cfg.SecurityTransports = append(cfg.SecurityTransports, config.MsSecC{SecC: stpt, ID: name})
		return nil
	}
}

// NoSecurity is an option that completely disables all transport security.
// It's incompatible with all other transport security protocols.
var NoSecurity Option = func(cfg *Config) error {
	if len(cfg.SecurityTransports) > 0 {
		return fmt.Errorf("cannot use security transports with an insecure libp2p configuration")
	}
	cfg.Insecure = true
	return nil
}

// Muxer configures libp2p to use the given stream multiplexer (or stream
// multiplexer constructor).
//
// Name is the protocol name.
//
// The transport can be a constructed mux.Transport or a function taking any
// subset of this libp2p node's:
// * Peer ID
// * Host
// * Network
// * Peerstore
func Muxer(name string, tpt interface{}) Option {
	mtpt, err := config.MuxerConstructor(tpt)
	err = traceError(err, 1)
	return func(cfg *Config) error {
		if err != nil {
			return err
		}
		cfg.Muxers = append(cfg.Muxers, config.MsMuxC{MuxC: mtpt, ID: name})
		return nil
	}
}

// Transport configures libp2p to use the given transport (or transport
// constructor).
//
// The transport can be a constructed transport.Transport or a function taking
// any subset of this libp2p node's:
// * Transport Upgrader (*tptu.Upgrader)
// * Host
// * Stream muxer (muxer.Transport)
// * Security transport (security.Transport)
// * Private network protector (pnet.Protector)
// * Peer ID
// * Private Key
// * Public Key
// * Address filter (filter.Filter)
// * Peerstore
func Transport(tpt interface{}) Option {
	tptc, err := config.TransportConstructor(tpt)
	err = traceError(err, 1)
	return func(cfg *Config) error {
		if err != nil {
			return err
		}
		cfg.Transports = append(cfg.Transports, tptc)
		return nil
	}
}

// Peerstore configures libp2p to use the given peerstore.
func Peerstore(ps peerstore.Peerstore) Option {
	return func(cfg *Config) error {
		if cfg.Peerstore != nil {
			return fmt.Errorf("cannot specify multiple peerstore options")
		}

		cfg.Peerstore = ps
		return nil
	}
}

// PrivateNetwork configures libp2p to use the given private network protector.
func PrivateNetwork(psk pnet.PSK) Option {
	return func(cfg *Config) error {
		if cfg.PSK != nil {
			return fmt.Errorf("cannot specify multiple private network options")
		}

		cfg.PSK = psk
		return nil
	}
}

// BandwidthReporter configures libp2p to use the given bandwidth reporter.
func BandwidthReporter(rep metrics.Reporter) Option {
	return func(cfg *Config) error {
		if cfg.Reporter != nil {
			return fmt.Errorf("cannot specify multiple bandwidth reporter options")
		}

		cfg.Reporter = rep
		return nil
	}
}

// Identity configures libp2p to use the given private key to identify itself.
func Identity(sk crypto.PrivKey) Option {
	return func(cfg *Config) error {
		if cfg.PeerKey != nil {
			return fmt.Errorf("cannot specify multiple identities")
		}

		cfg.PeerKey = sk
		return nil
	}
}

// ConnectionManager configures libp2p to use the given connection manager.
func ConnectionManager(connman connmgr.ConnManager) Option {
	return func(cfg *Config) error {
		if cfg.ConnManager != nil {
			return fmt.Errorf("cannot specify multiple connection managers")
		}
		cfg.ConnManager = connman
		return nil
	}
}

// AddrsFactory configures libp2p to use the given address factory.
func AddrsFactory(factory config.AddrsFactory) Option {
	return func(cfg *Config) error {
		if cfg.AddrsFactory != nil {
			return fmt.Errorf("cannot specify multiple address factories")
		}
		cfg.AddrsFactory = factory
		return nil
	}
}

// EnableRelay configures libp2p to enable the relay transport with
// configuration options. By default, this option only configures libp2p to
// accept inbound connections from relays and make outbound connections
// _through_ relays when requested by the remote peer. (default: enabled)
//
// To _act_ as a relay, pass the circuit.OptHop option.
func EnableRelay(options ...circuit.RelayOpt) Option {
	return func(cfg *Config) error {
		cfg.RelayCustom = true
		cfg.Relay = true
		cfg.RelayOpts = options
		return nil
	}
}

// DisableRelay configures libp2p to disable the relay transport.
func DisableRelay() Option {
	return func(cfg *Config) error {
		cfg.RelayCustom = true
		cfg.Relay = false
		return nil
	}
}

// EnableAutoRelay configures libp2p to enable the AutoRelay subsystem.
//
// Dependencies:
//  * Relay (enabled by default)
//  * Routing (to find relays), or StaticRelays/DefaultStaticRelays.
//
// This subsystem performs two functions:
//
// 1. When this libp2p node is configured to act as a relay "hop"
//    (circuit.OptHop is passed to EnableRelay), this node will advertise itself
//    as a public relay using the provided routing system.
// 2. When this libp2p node is _not_ configured as a relay "hop", it will
//    automatically detect if it is unreachable (e.g., behind a NAT). If so, it will
//    find, configure, and announce a set of public relays.
func EnableAutoRelay() Option {
	return func(cfg *Config) error {
		cfg.EnableAutoRelay = true
		return nil
	}
}

// StaticRelays configures known relays for autorelay; when this option is enabled
// then the system will use the configured relays instead of querying the DHT to
// discover relays.
func StaticRelays(relays []peer.AddrInfo) Option {
	return func(cfg *Config) error {
		cfg.StaticRelays = append(cfg.StaticRelays, relays...)
		return nil
	}
}

// DefaultStaticRelays configures the static relays to use the known PL-operated relays.
func DefaultStaticRelays() Option {
	return func(cfg *Config) error {
		for _, addr := range autorelay.DefaultRelays {
			a, err := ma.NewMultiaddr(addr)
			if err != nil {
				return err
			}
			pi, err := peer.AddrInfoFromP2pAddr(a)
			if err != nil {
				return err
			}
			cfg.StaticRelays = append(cfg.StaticRelays, *pi)
		}

		return nil
	}
}

// WithReachability overrides automatic reachability detection to force the local node
// to believe it is either unreachable or reachable externally.
func WithReachability(reachable bool) Option {
	return func(cfg *Config) error {
		if reachable {
			cfg.Reachability = network.ReachabilityPublic
		} else {
			cfg.Reachability = network.ReachabilityPrivate
		}
		return nil
	}
}

// EnableNATService configures libp2p to provide a service to peers for determining
// their reachability status. When enabled, the host will attempt to dial back
// to peers, and then tell them if it was successful in making such connections.
func EnableNATService() Option {
	return func(cfg *Config) error {
		cfg.AutoNATService = true
		return nil
	}
}

// FilterAddresses configures libp2p to never dial nor accept connections from
// the given addresses. FilterAddresses should be used for cases where the
// addresses you want to deny are known ahead of time.
func FilterAddresses(addrs ...*net.IPNet) Option {
	return func(cfg *Config) error {
		if cfg.Filters == nil {
			cfg.Filters = filter.NewFilters()
		}
		for _, addr := range addrs {
			cfg.Filters.AddDialFilter(addr)
		}
		return nil
	}
}

// Filters configures libp2p to use the given filters for accepting/denying
// certain addresses. Filters offers more control and should be used when the
// addresses you want to accept/deny are not known ahead of time and can
// dynamically change.
func Filters(filters *filter.Filters) Option {
	return func(cfg *Config) error {
		cfg.Filters = filters
		return nil
	}
}

// NATPortMap configures libp2p to use the default NATManager. The default
// NATManager will attempt to open a port in your network's firewall using UPnP.
func NATPortMap() Option {
	return NATManager(bhost.NewNATManager)
}

// NATManager will configure libp2p to use the requested NATManager. This
// function should be passed a NATManager *constructor* that takes a libp2p Network.
func NATManager(nm config.NATManagerC) Option {
	return func(cfg *Config) error {
		if cfg.NATManager != nil {
			return fmt.Errorf("cannot specify multiple NATManagers")
		}
		cfg.NATManager = nm
		return nil
	}
}

// Ping will configure libp2p to support the ping service; enable by default.
func Ping(enable bool) Option {
	return func(cfg *Config) error {
		cfg.DisablePing = !enable
		return nil
	}
}

// Routing will configure libp2p to use routing.
func Routing(rt config.RoutingC) Option {
	return func(cfg *Config) error {
		if cfg.Routing != nil {
			return fmt.Errorf("cannot specify multiple routing options")
		}
		cfg.Routing = rt
		return nil
	}
}

// NoListenAddrs will configure libp2p to not listen by default.
//
// This will both clear any configured listen addrs and prevent libp2p from
// applying the default listen address option. It also disables relay, unless the
// user explicitly specifies with an option, as the transport creates an implicit
// listen address that would make the node dialable through any relay it was connected to.
var NoListenAddrs = func(cfg *Config) error {
	cfg.ListenAddrs = []ma.Multiaddr{}
	if !cfg.RelayCustom {
		cfg.RelayCustom = true
		cfg.Relay = false
	}
	return nil
}

// NoTransports will configure libp2p to not enable any transports.
//
// This will both clear any configured transports (specified in prior libp2p
// options) and prevent libp2p from applying the default transports.
var NoTransports = func(cfg *Config) error {
	cfg.Transports = []config.TptC{}
	return nil
}

// UserAgent sets the libp2p user-agent sent along with the identify protocol
func UserAgent(userAgent string) Option {
	return func(cfg *Config) error {
		cfg.UserAgent = userAgent
		return nil
	}
}
<commit_msg>Clearer naming of nat override options<commit_after>package libp2p

// This file contains all libp2p configuration options (except the defaults,
// those are in defaults.go).

import (
	"fmt"
	"net"

	"github.com/libp2p/go-libp2p-core/connmgr"
	"github.com/libp2p/go-libp2p-core/crypto"
	"github.com/libp2p/go-libp2p-core/metrics"
	"github.com/libp2p/go-libp2p-core/network"
	"github.com/libp2p/go-libp2p-core/peer"
	"github.com/libp2p/go-libp2p-core/peerstore"
	"github.com/libp2p/go-libp2p-core/pnet"

	circuit "github.com/libp2p/go-libp2p-circuit"
	config "github.com/libp2p/go-libp2p/config"
	bhost "github.com/libp2p/go-libp2p/p2p/host/basic"
	autorelay "github.com/libp2p/go-libp2p/p2p/host/relay"

	filter "github.com/libp2p/go-maddr-filter"
	ma "github.com/multiformats/go-multiaddr"
)

// ListenAddrStrings configures libp2p to listen on the given (unparsed)
// addresses.
func ListenAddrStrings(s ...string) Option {
	return func(cfg *Config) error {
		for _, addrstr := range s {
			a, err := ma.NewMultiaddr(addrstr)
			if err != nil {
				return err
			}
			cfg.ListenAddrs = append(cfg.ListenAddrs, a)
		}
		return nil
	}
}

// ListenAddrs configures libp2p to listen on the given addresses.
func ListenAddrs(addrs ...ma.Multiaddr) Option {
	return func(cfg *Config) error {
		cfg.ListenAddrs = append(cfg.ListenAddrs, addrs...)
		return nil
	}
}

// Security configures libp2p to use the given security transport (or transport
// constructor).
//
// Name is the protocol name.
//
// The transport can be a constructed security.Transport or a function taking
// any subset of this libp2p node's:
// * Public key
// * Private key
// * Peer ID
// * Host
// * Network
// * Peerstore
func Security(name string, tpt interface{}) Option {
	stpt, err := config.SecurityConstructor(tpt)
	err = traceError(err, 1)
	return func(cfg *Config) error {
		if err != nil {
			return err
		}
		if cfg.Insecure {
			return fmt.Errorf("cannot use security transports with an insecure libp2p configuration")
		}
		cfg.SecurityTransports = append(cfg.SecurityTransports, config.MsSecC{SecC: stpt, ID: name})
		return nil
	}
}

// NoSecurity is an option that completely disables all transport security.
// It's incompatible with all other transport security protocols.
var NoSecurity Option = func(cfg *Config) error {
	if len(cfg.SecurityTransports) > 0 {
		return fmt.Errorf("cannot use security transports with an insecure libp2p configuration")
	}
	cfg.Insecure = true
	return nil
}

// Muxer configures libp2p to use the given stream multiplexer (or stream
// multiplexer constructor).
//
// Name is the protocol name.
//
// The transport can be a constructed mux.Transport or a function taking any
// subset of this libp2p node's:
// * Peer ID
// * Host
// * Network
// * Peerstore
func Muxer(name string, tpt interface{}) Option {
	mtpt, err := config.MuxerConstructor(tpt)
	err = traceError(err, 1)
	return func(cfg *Config) error {
		if err != nil {
			return err
		}
		cfg.Muxers = append(cfg.Muxers, config.MsMuxC{MuxC: mtpt, ID: name})
		return nil
	}
}

// Transport configures libp2p to use the given transport (or transport
// constructor).
//
// The transport can be a constructed transport.Transport or a function taking
// any subset of this libp2p node's:
// * Transport Upgrader (*tptu.Upgrader)
// * Host
// * Stream muxer (muxer.Transport)
// * Security transport (security.Transport)
// * Private network protector (pnet.Protector)
// * Peer ID
// * Private Key
// * Public Key
// * Address filter (filter.Filter)
// * Peerstore
func Transport(tpt interface{}) Option {
	tptc, err := config.TransportConstructor(tpt)
	err = traceError(err, 1)
	return func(cfg *Config) error {
		if err != nil {
			return err
		}
		cfg.Transports = append(cfg.Transports, tptc)
		return nil
	}
}

// Peerstore configures libp2p to use the given peerstore.
func Peerstore(ps peerstore.Peerstore) Option {
	return func(cfg *Config) error {
		if cfg.Peerstore != nil {
			return fmt.Errorf("cannot specify multiple peerstore options")
		}

		cfg.Peerstore = ps
		return nil
	}
}

// PrivateNetwork configures libp2p to use the given private network protector.
func PrivateNetwork(psk pnet.PSK) Option {
	return func(cfg *Config) error {
		if cfg.PSK != nil {
			return fmt.Errorf("cannot specify multiple private network options")
		}

		cfg.PSK = psk
		return nil
	}
}

// BandwidthReporter configures libp2p to use the given bandwidth reporter.
func BandwidthReporter(rep metrics.Reporter) Option {
	return func(cfg *Config) error {
		if cfg.Reporter != nil {
			return fmt.Errorf("cannot specify multiple bandwidth reporter options")
		}

		cfg.Reporter = rep
		return nil
	}
}

// Identity configures libp2p to use the given private key to identify itself.
func Identity(sk crypto.PrivKey) Option {
	return func(cfg *Config) error {
		if cfg.PeerKey != nil {
			return fmt.Errorf("cannot specify multiple identities")
		}

		cfg.PeerKey = sk
		return nil
	}
}

// ConnectionManager configures libp2p to use the given connection manager.
func ConnectionManager(connman connmgr.ConnManager) Option {
	return func(cfg *Config) error {
		if cfg.ConnManager != nil {
			return fmt.Errorf("cannot specify multiple connection managers")
		}
		cfg.ConnManager = connman
		return nil
	}
}

// AddrsFactory configures libp2p to use the given address factory.
func AddrsFactory(factory config.AddrsFactory) Option {
	return func(cfg *Config) error {
		if cfg.AddrsFactory != nil {
			return fmt.Errorf("cannot specify multiple address factories")
		}
		cfg.AddrsFactory = factory
		return nil
	}
}

// EnableRelay configures libp2p to enable the relay transport with
// configuration options. By default, this option only configures libp2p to
// accept inbound connections from relays and make outbound connections
// _through_ relays when requested by the remote peer. (default: enabled)
//
// To _act_ as a relay, pass the circuit.OptHop option.
func EnableRelay(options ...circuit.RelayOpt) Option {
	return func(cfg *Config) error {
		cfg.RelayCustom = true
		cfg.Relay = true
		cfg.RelayOpts = options
		return nil
	}
}

// DisableRelay configures libp2p to disable the relay transport.
func DisableRelay() Option {
	return func(cfg *Config) error {
		cfg.RelayCustom = true
		cfg.Relay = false
		return nil
	}
}

// EnableAutoRelay configures libp2p to enable the AutoRelay subsystem.
//
// Dependencies:
//  * Relay (enabled by default)
//  * Routing (to find relays), or StaticRelays/DefaultStaticRelays.
//
// This subsystem performs two functions:
//
// 1. When this libp2p node is configured to act as a relay "hop"
//    (circuit.OptHop is passed to EnableRelay), this node will advertise itself
//    as a public relay using the provided routing system.
// 2. When this libp2p node is _not_ configured as a relay "hop", it will
//    automatically detect if it is unreachable (e.g., behind a NAT). If so, it will
//    find, configure, and announce a set of public relays.
func EnableAutoRelay() Option {
	return func(cfg *Config) error {
		cfg.EnableAutoRelay = true
		return nil
	}
}

// StaticRelays configures known relays for autorelay; when this option is enabled
// then the system will use the configured relays instead of querying the DHT to
// discover relays.
func StaticRelays(relays []peer.AddrInfo) Option {
	return func(cfg *Config) error {
		cfg.StaticRelays = append(cfg.StaticRelays, relays...)
		return nil
	}
}

// DefaultStaticRelays configures the static relays to use the known PL-operated relays.
func DefaultStaticRelays() Option {
	return func(cfg *Config) error {
		for _, addr := range autorelay.DefaultRelays {
			a, err := ma.NewMultiaddr(addr)
			if err != nil {
				return err
			}
			pi, err := peer.AddrInfoFromP2pAddr(a)
			if err != nil {
				return err
			}
			cfg.StaticRelays = append(cfg.StaticRelays, *pi)
		}

		return nil
	}
}

// ForceReachabilityPublic overrides automatic reachability detection in the AutoNAT subsystem,
// forcing the local node to believe it is reachable externally.
func ForceReachabilityPublic() Option {
	return func(cfg *Config) error {
		cfg.Reachability = network.ReachabilityPublic
		return nil
	}
}

// ForceReachabilityPrivate overrides automatic reachability detection in the AutoNAT subsystem,
// forceing the local node to believe it is behind a NAT and not reachable externally.
func ForceReachabilityPrivate() Option {
	return func(cfg *Config) error {
		cfg.Reachability = network.ReachabilityPrivate
		return nil
	}
}

// EnableNATService configures libp2p to provide a service to peers for determining
// their reachability status. When enabled, the host will attempt to dial back
// to peers, and then tell them if it was successful in making such connections.
func EnableNATService() Option {
	return func(cfg *Config) error {
		cfg.AutoNATService = true
		return nil
	}
}

// FilterAddresses configures libp2p to never dial nor accept connections from
// the given addresses. FilterAddresses should be used for cases where the
// addresses you want to deny are known ahead of time.
func FilterAddresses(addrs ...*net.IPNet) Option {
	return func(cfg *Config) error {
		if cfg.Filters == nil {
			cfg.Filters = filter.NewFilters()
		}
		for _, addr := range addrs {
			cfg.Filters.AddDialFilter(addr)
		}
		return nil
	}
}

// Filters configures libp2p to use the given filters for accepting/denying
// certain addresses. Filters offers more control and should be used when the
// addresses you want to accept/deny are not known ahead of time and can
// dynamically change.
func Filters(filters *filter.Filters) Option {
	return func(cfg *Config) error {
		cfg.Filters = filters
		return nil
	}
}

// NATPortMap configures libp2p to use the default NATManager. The default
// NATManager will attempt to open a port in your network's firewall using UPnP.
func NATPortMap() Option {
	return NATManager(bhost.NewNATManager)
}

// NATManager will configure libp2p to use the requested NATManager. This
// function should be passed a NATManager *constructor* that takes a libp2p Network.
func NATManager(nm config.NATManagerC) Option {
	return func(cfg *Config) error {
		if cfg.NATManager != nil {
			return fmt.Errorf("cannot specify multiple NATManagers")
		}
		cfg.NATManager = nm
		return nil
	}
}

// Ping will configure libp2p to support the ping service; enable by default.
func Ping(enable bool) Option {
	return func(cfg *Config) error {
		cfg.DisablePing = !enable
		return nil
	}
}

// Routing will configure libp2p to use routing.
func Routing(rt config.RoutingC) Option {
	return func(cfg *Config) error {
		if cfg.Routing != nil {
			return fmt.Errorf("cannot specify multiple routing options")
		}
		cfg.Routing = rt
		return nil
	}
}

// NoListenAddrs will configure libp2p to not listen by default.
//
// This will both clear any configured listen addrs and prevent libp2p from
// applying the default listen address option. It also disables relay, unless the
// user explicitly specifies with an option, as the transport creates an implicit
// listen address that would make the node dialable through any relay it was connected to.
var NoListenAddrs = func(cfg *Config) error {
	cfg.ListenAddrs = []ma.Multiaddr{}
	if !cfg.RelayCustom {
		cfg.RelayCustom = true
		cfg.Relay = false
	}
	return nil
}

// NoTransports will configure libp2p to not enable any transports.
//
// This will both clear any configured transports (specified in prior libp2p
// options) and prevent libp2p from applying the default transports.
var NoTransports = func(cfg *Config) error {
	cfg.Transports = []config.TptC{}
	return nil
}

// UserAgent sets the libp2p user-agent sent along with the identify protocol
func UserAgent(userAgent string) Option {
	return func(cfg *Config) error {
		cfg.UserAgent = userAgent
		return nil
	}
}
<|endoftext|>


<commit_before>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package present

import (
	"fmt"
	"net/url"
	"strings"
)

func init() {
	Register("link", parseLink)
}

type Link struct {
	URL   *url.URL
	Label string
}

func (l Link) TemplateName() string { return "link" }

func parseLink(ctx *Context, fileName string, lineno int, text string) (Elem, error) {
	args := strings.Fields(text)
	url, err := url.Parse(args[1])
	if err != nil {
		return nil, err
	}
	label := ""
	if len(args) > 2 {
		label = strings.Join(args[2:], " ")
	} else {
		scheme := url.Scheme + "://"
		if url.Scheme == "mailto" {
			scheme = "mailto:"
		}
		label = strings.Replace(url.String(), scheme, "", 1)
	}
	return Link{url, label}, nil
}

func renderLink(url, text string) string {
	text = font(text)
	if text == "" {
		text = url
	}
	return fmt.Sprintf(`<a href="%s" target="_blank">%s</a>`, url, text)
}

// parseInlineLink parses an inline link at the start of s, and returns
// a rendered HTML link and the total length of the raw inline link.
// If no inline link is present, it returns all zeroes.
func parseInlineLink(s string) (link string, length int) {
	if !strings.HasPrefix(s, "[[") {
		return
	}
	end := strings.Index(s, "]]")
	if end == -1 {
		return
	}
	urlEnd := strings.Index(s, "]")
	rawURL := s[2:urlEnd]
	const badURLChars = `<>"{}|\^[] ` + "`" // per RFC2396 section 2.4.3
	if strings.ContainsAny(rawURL, badURLChars) {
		return
	}
	if urlEnd == end {
		simpleUrl := ""
		url, err := url.Parse(rawURL)
		if err == nil {
			// If the URL is http://foo.com, drop the http://
			// In other words, render [[http://golang.org]] as:
			//   golang.org
			if strings.HasPrefix(rawURL, url.Scheme+"://") {
				simpleUrl = strings.TrimPrefix(rawURL, url.Scheme+"://")
			} else if strings.HasPrefix(rawURL, url.Scheme+":") {
				simpleUrl = strings.TrimPrefix(rawURL, url.Scheme+":")
			}
		}
		return renderLink(rawURL, simpleUrl), end + 2
	}
	if s[urlEnd:urlEnd+2] != "][" {
		return
	}
	text := s[urlEnd+2 : end]
	return renderLink(rawURL, text), end + 2
}
<commit_msg>go.tools/present: Open links with local (non absolute) URLs in the same window.<commit_after>// Copyright 2012 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package present

import (
	"fmt"
	"net/url"
	"strings"
)

func init() {
	Register("link", parseLink)
}

type Link struct {
	URL   *url.URL
	Label string
}

func (l Link) TemplateName() string { return "link" }

func parseLink(ctx *Context, fileName string, lineno int, text string) (Elem, error) {
	args := strings.Fields(text)
	url, err := url.Parse(args[1])
	if err != nil {
		return nil, err
	}
	label := ""
	if len(args) > 2 {
		label = strings.Join(args[2:], " ")
	} else {
		scheme := url.Scheme + "://"
		if url.Scheme == "mailto" {
			scheme = "mailto:"
		}
		label = strings.Replace(url.String(), scheme, "", 1)
	}
	return Link{url, label}, nil
}

func renderLink(url, text string) string {
	text = font(text)
	if text == "" {
		text = url
	}
	// Open links in new window only when their url is absolute.
	target := "_blank"
	if url[0] == '/' {
		target = "_self"
	}

	return fmt.Sprintf(`<a href="%s" target="%s">%s</a>`, url, target, text)
}

// parseInlineLink parses an inline link at the start of s, and returns
// a rendered HTML link and the total length of the raw inline link.
// If no inline link is present, it returns all zeroes.
func parseInlineLink(s string) (link string, length int) {
	if !strings.HasPrefix(s, "[[") {
		return
	}
	end := strings.Index(s, "]]")
	if end == -1 {
		return
	}
	urlEnd := strings.Index(s, "]")
	rawURL := s[2:urlEnd]
	const badURLChars = `<>"{}|\^[] ` + "`" // per RFC2396 section 2.4.3
	if strings.ContainsAny(rawURL, badURLChars) {
		return
	}
	if urlEnd == end {
		simpleUrl := ""
		url, err := url.Parse(rawURL)
		if err == nil {
			// If the URL is http://foo.com, drop the http://
			// In other words, render [[http://golang.org]] as:
			//   golang.org
			if strings.HasPrefix(rawURL, url.Scheme+"://") {
				simpleUrl = strings.TrimPrefix(rawURL, url.Scheme+"://")
			} else if strings.HasPrefix(rawURL, url.Scheme+":") {
				simpleUrl = strings.TrimPrefix(rawURL, url.Scheme+":")
			}
		}
		return renderLink(rawURL, simpleUrl), end + 2
	}
	if s[urlEnd:urlEnd+2] != "][" {
		return
	}
	text := s[urlEnd+2 : end]
	return renderLink(rawURL, text), end + 2
}
<|endoftext|>


<commit_before>package printer

import (
	"bytes"
	"fmt"
	"io"
	"regexp"
	"sort"

	"github.com/mibk/dupl/syntax"
)

type html struct {
	iota int
	w    io.Writer
	ReadFile
}

func NewHTML(w io.Writer, fread ReadFile) Printer {
	return &html{w: w, ReadFile: fread}
}

func (p *html) PrintHeader() error {
	_, err := fmt.Fprint(p.w, `<!DOCTYPE html>
<meta charset="utf-8"/>
<title>Duplicates</title>

`)
	return err
}

func (p *html) PrintClones(dups [][]*syntax.Node) error {
	p.iota++
	fmt.Fprintf(p.w, "
#%d found %d clones
\n", p.iota, len(dups))

	clones := make([]clone, len(dups))
	for i, dup := range dups {
		cnt := len(dup)
		if cnt == 0 {
			panic("zero length dup")
		}
		nstart := dup[0]
		nend := dup[cnt-1]

		file, err := p.ReadFile(nstart.Filename)
		if err != nil {
			return err
		}

		lineStart, _ := blockLines(file, nstart.Pos, nend.End)
		cl := clone{filename: nstart.Filename, lineStart: lineStart}
		start := findLineBeg(file, nstart.Pos)
		content := append(toWhitespace(file[start:nstart.Pos]), file[nstart.Pos:nend.End]...)
		cl.fragment = deindent(content)
		clones[i] = cl
	}

	sort.Sort(byNameAndLine(clones))
	for _, cl := range clones {
		fmt.Fprintf(p.w, "
%s:%d
\n
%s
\n", cl.filename, cl.lineStart, cl.fragment)
	}
	return nil
}

func (*html) PrintFooter() error { return nil }

func findLineBeg(file []byte, index int) int {
	for i := index; i >= 0; i-- {
		if file[i] == '\n' {
			return i + 1
		}
	}
	return 0
}

func toWhitespace(str []byte) []byte {
	var out []byte
	for _, c := range bytes.Runes(str) {
		if c == '\t' {
			out = append(out, '\t')
		} else {
			out = append(out, ' ')
		}
	}
	return out
}

func deindent(block []byte) []byte {
	const maxVal = 99
	min := maxVal
	re := regexp.MustCompile(`(^|\n)(\t*)\S`)
	for _, line := range re.FindAllSubmatch(block, -1) {
		indent := line[2]
		if len(indent) < min {
			min = len(indent)
		}
	}
	if min == 0 || min == maxVal {
		return block
	}
	block = block[min:]
Loop:
	for i := 0; i < len(block); i++ {
		if block[i] == '\n' && i != len(block)-1 {
			for j := 0; j < min; j++ {
				if block[i+j+1] != '\t' {
					continue Loop
				}
			}
			block = append(block[:i+1], block[i+1+min:]...)
		}
	}
	return block
}
<commit_msg>printer: fix HTML escaping<commit_after>package printer

import (
	"bytes"
	"fmt"
	"html"
	"io"
	"regexp"
	"sort"

	"github.com/mibk/dupl/syntax"
)

type htmlprinter struct {
	iota int
	w    io.Writer
	ReadFile
}

func NewHTML(w io.Writer, fread ReadFile) Printer {
	return &htmlprinter{w: w, ReadFile: fread}
}

func (p *htmlprinter) PrintHeader() error {
	_, err := fmt.Fprint(p.w, `<!DOCTYPE html>
<meta charset="utf-8"/>
<title>Duplicates</title>

`)
	return err
}

func (p *htmlprinter) PrintClones(dups [][]*syntax.Node) error {
	p.iota++
	fmt.Fprintf(p.w, "
#%d found %d clones
\n", p.iota, len(dups))

	clones := make([]clone, len(dups))
	for i, dup := range dups {
		cnt := len(dup)
		if cnt == 0 {
			panic("zero length dup")
		}
		nstart := dup[0]
		nend := dup[cnt-1]

		file, err := p.ReadFile(nstart.Filename)
		if err != nil {
			return err
		}

		lineStart, _ := blockLines(file, nstart.Pos, nend.End)
		cl := clone{filename: nstart.Filename, lineStart: lineStart}
		start := findLineBeg(file, nstart.Pos)
		content := append(toWhitespace(file[start:nstart.Pos]), file[nstart.Pos:nend.End]...)
		cl.fragment = deindent(content)
		clones[i] = cl
	}

	sort.Sort(byNameAndLine(clones))
	for _, cl := range clones {
		fmt.Fprintf(p.w, "
%s:%d
\n
%s
\n", cl.filename, cl.lineStart,
			html.EscapeString(string(cl.fragment)))
	}
	return nil
}

func (*htmlprinter) PrintFooter() error { return nil }

func findLineBeg(file []byte, index int) int {
	for i := index; i >= 0; i-- {
		if file[i] == '\n' {
			return i + 1
		}
	}
	return 0
}

func toWhitespace(str []byte) []byte {
	var out []byte
	for _, c := range bytes.Runes(str) {
		if c == '\t' {
			out = append(out, '\t')
		} else {
			out = append(out, ' ')
		}
	}
	return out
}

func deindent(block []byte) []byte {
	const maxVal = 99
	min := maxVal
	re := regexp.MustCompile(`(^|\n)(\t*)\S`)
	for _, line := range re.FindAllSubmatch(block, -1) {
		indent := line[2]
		if len(indent) < min {
			min = len(indent)
		}
	}
	if min == 0 || min == maxVal {
		return block
	}
	block = block[min:]
Loop:
	for i := 0; i < len(block); i++ {
		if block[i] == '\n' && i != len(block)-1 {
			for j := 0; j < min; j++ {
				if block[i+j+1] != '\t' {
					continue Loop
				}
			}
			block = append(block[:i+1], block[i+1+min:]...)
		}
	}
	return block
}
<|endoftext|>


<commit_before>package gopsutil

import (
	"fmt"
	"os"
	"runtime"
	"testing"
)

func test_getProcess() Process {
	checkPid := os.Getpid()
	if runtime.GOOS == "windows" {
		checkPid = 7960
	}
	ret, _ := NewProcess(int32(checkPid))
	return *ret
}

func Test_Pids(t *testing.T) {
	ret, err := Pids()
	if err != nil {
		t.Errorf("error %v", err)
	}
	if len(ret) == 0 {
		t.Errorf("could not get pids %v", ret)
	}
}

func Test_Pid_exists(t *testing.T) {
	checkPid := 1
	if runtime.GOOS == "windows" {
		checkPid = 0
	}

	ret, err := PidExists(int32(checkPid))
	if err != nil {
		t.Errorf("error %v", err)
	}

	if ret == false {
		t.Errorf("could not get init process %v", ret)
	}
}

func Test_NewProcess(t *testing.T) {
	checkPid := 1
	if runtime.GOOS == "windows" {
		checkPid = 0
	}

	ret, err := NewProcess(int32(checkPid))
	if err != nil {
		t.Errorf("error %v", err)
	}

	fmt.Println(ret)
}

func Test_Process_memory_maps(t *testing.T) {
	checkPid := os.Getpid()
	if runtime.GOOS == "windows" {
		checkPid = 0
	}
	return
	ret, err := NewProcess(int32(checkPid))

	mmaps, err := ret.MemoryMaps(false)
	if err != nil {
		t.Errorf("memory map get error %v", err)
	}
	for _, m := range *mmaps {
		fmt.Println(m)
	}

}

func Test_Process_Ppid(t *testing.T) {
	p := test_getProcess()

	v, err := p.Ppid()
	if err != nil {
		t.Errorf("geting ppid error %v", err)
	}
	if v == 0 {
		t.Errorf("return value is 0 %v", v)
	}

}

func Test_Process_IOCounters(t *testing.T) {
	p := test_getProcess()

	v, err := p.IOCounters()
	if err != nil {
		t.Errorf("geting ppid error %v", err)
	}
	fmt.Println(v)
	if v.ReadCount == 0 {
		t.Errorf("return value is 0 %v", v)
	}

}
<commit_msg>fix if disk_io_counters failed.<commit_after>package gopsutil

import (
	"fmt"
	"os"
	"runtime"
	"testing"
)

func test_getProcess() Process {
	checkPid := os.Getpid()
	if runtime.GOOS == "windows" {
		checkPid = 7960
	}
	ret, _ := NewProcess(int32(checkPid))
	return *ret
}

func Test_Pids(t *testing.T) {
	ret, err := Pids()
	if err != nil {
		t.Errorf("error %v", err)
	}
	if len(ret) == 0 {
		t.Errorf("could not get pids %v", ret)
	}
}

func Test_Pid_exists(t *testing.T) {
	checkPid := 1
	if runtime.GOOS == "windows" {
		checkPid = 0
	}

	ret, err := PidExists(int32(checkPid))
	if err != nil {
		t.Errorf("error %v", err)
	}

	if ret == false {
		t.Errorf("could not get init process %v", ret)
	}
}

func Test_NewProcess(t *testing.T) {
	checkPid := 1
	if runtime.GOOS == "windows" {
		checkPid = 0
	}

	ret, err := NewProcess(int32(checkPid))
	if err != nil {
		t.Errorf("error %v", err)
	}

	fmt.Println(ret)
}

func Test_Process_memory_maps(t *testing.T) {
	checkPid := os.Getpid()
	if runtime.GOOS == "windows" {
		checkPid = 0
	}
	return
	ret, err := NewProcess(int32(checkPid))

	mmaps, err := ret.MemoryMaps(false)
	if err != nil {
		t.Errorf("memory map get error %v", err)
	}
	for _, m := range *mmaps {
		fmt.Println(m)
	}

}

func Test_Process_Ppid(t *testing.T) {
	p := test_getProcess()

	v, err := p.Ppid()
	if err != nil {
		t.Errorf("geting ppid error %v", err)
	}
	if v == 0 {
		t.Errorf("return value is 0 %v", v)
	}

}

func Test_Process_IOCounters(t *testing.T) {
	p := test_getProcess()

	v, err := p.IOCounters()
	if err != nil {
		t.Errorf("geting ppid error %v", err)
		return
	}
	if v.ReadCount == 0 {
		t.Errorf("return value is 0 %v", v)
	}
	fmt.Println(v)

}
<|endoftext|>


<commit_before>package procker

import (
	"bytes"
	"testing"
)

func TestProcessStart(t *testing.T) {
	stdOut := &bytes.Buffer{}
	stdErr := &bytes.Buffer{}
	var env []string

	p := NewProcess("simple", "echo -n procker")
	err := p.Start(env, stdOut, stdErr)

	if err != nil {
		t.Fatal("process failed")
	}

	err = p.Wait()
	if err != nil {
		t.Fatal("process failed")
	}

	if stdOut.String() != "procker" {
		t.Error("bad output")
	}

	if stdErr.String() != "" {
		t.Error("bad output")
	}
}

func TestProcessStartUsingEnv(t *testing.T) {
	stdOut := &bytes.Buffer{}
	stdErr := &bytes.Buffer{}
	var env []string = []string{"PROCKER_MSG=hello", "PROCKER_MSG2=world"}

	p := NewProcess("simple", "echo -n $PROCKER_MSG $PROCKER_MSG2")
	err := p.Start(env, stdOut, stdErr)

	if err != nil {
		t.Fatal("process failed")
	}

	err = p.Wait()
	if err != nil {
		t.Fatal("process failed")
	}

	if stdOut.String() != "hello world" {
		t.Error("bad output")
	}

	if stdErr.String() != "" {
		t.Error("bad output")
	}
}
<commit_msg>add assert test helper function<commit_after>package procker

import (
	"bytes"
	"reflect"
	"testing"
)

func assert(t *testing.T, expected, actual interface{}) {
	if !reflect.DeepEqual(expected, actual) {
		t.Errorf("expected: %+v, actual: %+v", expected, actual)
	}
}

func TestProcessStart(t *testing.T) {
	stdOut := &bytes.Buffer{}
	stdErr := &bytes.Buffer{}
	var env []string

	p := NewProcess("simple", "echo -n procker")
	err := p.Start(env, stdOut, stdErr)

	if err != nil {
		t.Fatal("process failed")
	}

	err = p.Wait()
	if err != nil {
		t.Fatal("process failed")
	}

	assert(t, "procker", stdOut.String())
	assert(t, "", stdErr.String())
}

func TestProcessStartUsingEnv(t *testing.T) {
	stdOut := &bytes.Buffer{}
	stdErr := &bytes.Buffer{}
	var env []string = []string{"PROCKER_MSG=hello", "PROCKER_MSG2=world"}

	p := NewProcess("simple", "echo -n $PROCKER_MSG $PROCKER_MSG2")
	err := p.Start(env, stdOut, stdErr)

	if err != nil {
		t.Fatal("process failed")
	}

	err = p.Wait()
	if err != nil {
		t.Fatal("process failed")
	}

	assert(t, "hello world", stdOut.String())
	assert(t, "", stdErr.String())
}
<|endoftext|>


<commit_before>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style.
// license that can be found in the LICENSE file.

package proto

import (
	"google.golang.org/protobuf/internal/encoding/messageset"
	"google.golang.org/protobuf/internal/encoding/wire"
	"google.golang.org/protobuf/internal/errors"
	"google.golang.org/protobuf/internal/flags"
	"google.golang.org/protobuf/internal/pragma"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"
	"google.golang.org/protobuf/runtime/protoiface"
)

// UnmarshalOptions configures the unmarshaler.
//
// Example usage:
//   err := UnmarshalOptions{DiscardUnknown: true}.Unmarshal(b, m)
type UnmarshalOptions struct {
	pragma.NoUnkeyedLiterals

	// Merge merges the input into the destination message.
	// The default behavior is to always reset the message before unmarshaling,
	// unless Merge is specified.
	Merge bool

	// AllowPartial accepts input for messages that will result in missing
	// required fields. If AllowPartial is false (the default), Unmarshal will
	// return an error if there are any missing required fields.
	AllowPartial bool

	// If DiscardUnknown is set, unknown fields are ignored.
	DiscardUnknown bool

	// Resolver is used for looking up types when unmarshaling extension fields.
	// If nil, this defaults to using protoregistry.GlobalTypes.
	Resolver interface {
		FindExtensionByName(field protoreflect.FullName) (protoreflect.ExtensionType, error)
		FindExtensionByNumber(message protoreflect.FullName, field protoreflect.FieldNumber) (protoreflect.ExtensionType, error)
	}
}

// Unmarshal parses the wire-format message in b and places the result in m.
func Unmarshal(b []byte, m Message) error {
	_, err := UnmarshalOptions{}.unmarshal(b, m)
	return err
}

// Unmarshal parses the wire-format message in b and places the result in m.
func (o UnmarshalOptions) Unmarshal(b []byte, m Message) error {
	_, err := o.unmarshal(b, m)
	return err
}

// UnmarshalState parses a wire-format message and places the result in m.
//
// This method permits fine-grained control over the unmarshaler.
// Most users should use Unmarshal instead.
func (o UnmarshalOptions) UnmarshalState(m Message, in protoiface.UnmarshalInput) (protoiface.UnmarshalOutput, error) {
	return o.unmarshal(in.Buf, m)
}

func (o UnmarshalOptions) unmarshal(b []byte, message Message) (out protoiface.UnmarshalOutput, err error) {
	defaultResolver := false
	if o.Resolver == nil {
		o.Resolver = protoregistry.GlobalTypes
		defaultResolver = true
	} else if o.Resolver == protoregistry.GlobalTypes {
		defaultResolver = true
	}
	if !o.Merge {
		Reset(message)
	}
	allowPartial := o.AllowPartial
	o.Merge = true
	o.AllowPartial = true
	m := message.ProtoReflect()
	methods := protoMethods(m)
	if methods != nil && methods.Unmarshal != nil &&
		!(o.DiscardUnknown && methods.Flags&protoiface.SupportUnmarshalDiscardUnknown == 0) {
		opts := protoiface.UnmarshalOptions{
			Resolver: o.Resolver,
		}
		if o.DiscardUnknown {
			opts.Flags |= protoiface.UnmarshalDiscardUnknown
		}
		if defaultResolver {
			opts.Flags |= protoiface.UnmarshalDefaultResolver
		}
		out, err = methods.Unmarshal(m, protoiface.UnmarshalInput{
			Buf: b,
		}, opts)
	} else {
		err = o.unmarshalMessageSlow(b, m)
	}
	if err != nil {
		return out, err
	}
	if allowPartial {
		return out, nil
	}
	return out, isInitialized(m)
}

func (o UnmarshalOptions) unmarshalMessage(b []byte, m protoreflect.Message) error {
	_, err := o.unmarshal(b, m.Interface())
	return err
}

func (o UnmarshalOptions) unmarshalMessageSlow(b []byte, m protoreflect.Message) error {
	md := m.Descriptor()
	if messageset.IsMessageSet(md) {
		return unmarshalMessageSet(b, m, o)
	}
	fields := md.Fields()
	for len(b) > 0 {
		// Parse the tag (field number and wire type).
		num, wtyp, tagLen := wire.ConsumeTag(b)
		if tagLen < 0 {
			return wire.ParseError(tagLen)
		}
		if num > wire.MaxValidNumber {
			return errors.New("invalid field number")
		}

		// Find the field descriptor for this field number.
		fd := fields.ByNumber(num)
		if fd == nil && md.ExtensionRanges().Has(num) {
			extType, err := o.Resolver.FindExtensionByNumber(md.FullName(), num)
			if err != nil && err != protoregistry.NotFound {
				return err
			}
			if extType != nil {
				fd = extType.TypeDescriptor()
			}
		}
		var err error
		if fd == nil {
			err = errUnknown
		} else if flags.ProtoLegacy {
			if fd.IsWeak() && fd.Message().IsPlaceholder() {
				err = errUnknown // weak referent is not linked in
			}
		}

		// Parse the field value.
		var valLen int
		switch {
		case err != nil:
		case fd.IsList():
			valLen, err = o.unmarshalList(b[tagLen:], wtyp, m.Mutable(fd).List(), fd)
		case fd.IsMap():
			valLen, err = o.unmarshalMap(b[tagLen:], wtyp, m.Mutable(fd).Map(), fd)
		default:
			valLen, err = o.unmarshalSingular(b[tagLen:], wtyp, m, fd)
		}
		if err != nil {
			if err != errUnknown {
				return err
			}
			valLen = wire.ConsumeFieldValue(num, wtyp, b[tagLen:])
			if valLen < 0 {
				return wire.ParseError(valLen)
			}
			if !o.DiscardUnknown {
				m.SetUnknown(append(m.GetUnknown(), b[:tagLen+valLen]...))
			}
		}
		b = b[tagLen+valLen:]
	}
	return nil
}

func (o UnmarshalOptions) unmarshalSingular(b []byte, wtyp wire.Type, m protoreflect.Message, fd protoreflect.FieldDescriptor) (n int, err error) {
	v, n, err := o.unmarshalScalar(b, wtyp, fd)
	if err != nil {
		return 0, err
	}
	switch fd.Kind() {
	case protoreflect.GroupKind, protoreflect.MessageKind:
		m2 := m.Mutable(fd).Message()
		if err := o.unmarshalMessage(v.Bytes(), m2); err != nil {
			return n, err
		}
	default:
		// Non-message scalars replace the previous value.
		m.Set(fd, v)
	}
	return n, nil
}

func (o UnmarshalOptions) unmarshalMap(b []byte, wtyp wire.Type, mapv protoreflect.Map, fd protoreflect.FieldDescriptor) (n int, err error) {
	if wtyp != wire.BytesType {
		return 0, errUnknown
	}
	b, n = wire.ConsumeBytes(b)
	if n < 0 {
		return 0, wire.ParseError(n)
	}
	var (
		keyField = fd.MapKey()
		valField = fd.MapValue()
		key      protoreflect.Value
		val      protoreflect.Value
		haveKey  bool
		haveVal  bool
	)
	switch valField.Kind() {
	case protoreflect.GroupKind, protoreflect.MessageKind:
		val = mapv.NewValue()
	}
	// Map entries are represented as a two-element message with fields
	// containing the key and value.
	for len(b) > 0 {
		num, wtyp, n := wire.ConsumeTag(b)
		if n < 0 {
			return 0, wire.ParseError(n)
		}
		if num > wire.MaxValidNumber {
			return 0, errors.New("invalid field number")
		}
		b = b[n:]
		err = errUnknown
		switch num {
		case 1:
			key, n, err = o.unmarshalScalar(b, wtyp, keyField)
			if err != nil {
				break
			}
			haveKey = true
		case 2:
			var v protoreflect.Value
			v, n, err = o.unmarshalScalar(b, wtyp, valField)
			if err != nil {
				break
			}
			switch valField.Kind() {
			case protoreflect.GroupKind, protoreflect.MessageKind:
				if err := o.unmarshalMessage(v.Bytes(), val.Message()); err != nil {
					return 0, err
				}
			default:
				val = v
			}
			haveVal = true
		}
		if err == errUnknown {
			n = wire.ConsumeFieldValue(num, wtyp, b)
			if n < 0 {
				return 0, wire.ParseError(n)
			}
		} else if err != nil {
			return 0, err
		}
		b = b[n:]
	}
	// Every map entry should have entries for key and value, but this is not strictly required.
	if !haveKey {
		key = keyField.Default()
	}
	if !haveVal {
		switch valField.Kind() {
		case protoreflect.GroupKind, protoreflect.MessageKind:
		default:
			val = valField.Default()
		}
	}
	mapv.Set(key.MapKey(), val)
	return n, nil
}

// errUnknown is used internally to indicate fields which should be added
// to the unknown field set of a message. It is never returned from an exported
// function.
var errUnknown = errors.New("BUG: internal error (unknown)")
<commit_msg>proto: make use of fast-path initialization checks<commit_after>// Copyright 2018 The Go Authors. All rights reserved.
// Use of this source code is governed by a BSD-style.
// license that can be found in the LICENSE file.

package proto

import (
	"google.golang.org/protobuf/internal/encoding/messageset"
	"google.golang.org/protobuf/internal/encoding/wire"
	"google.golang.org/protobuf/internal/errors"
	"google.golang.org/protobuf/internal/flags"
	"google.golang.org/protobuf/internal/pragma"
	"google.golang.org/protobuf/reflect/protoreflect"
	"google.golang.org/protobuf/reflect/protoregistry"
	"google.golang.org/protobuf/runtime/protoiface"
)

// UnmarshalOptions configures the unmarshaler.
//
// Example usage:
//   err := UnmarshalOptions{DiscardUnknown: true}.Unmarshal(b, m)
type UnmarshalOptions struct {
	pragma.NoUnkeyedLiterals

	// Merge merges the input into the destination message.
	// The default behavior is to always reset the message before unmarshaling,
	// unless Merge is specified.
	Merge bool

	// AllowPartial accepts input for messages that will result in missing
	// required fields. If AllowPartial is false (the default), Unmarshal will
	// return an error if there are any missing required fields.
	AllowPartial bool

	// If DiscardUnknown is set, unknown fields are ignored.
	DiscardUnknown bool

	// Resolver is used for looking up types when unmarshaling extension fields.
	// If nil, this defaults to using protoregistry.GlobalTypes.
	Resolver interface {
		FindExtensionByName(field protoreflect.FullName) (protoreflect.ExtensionType, error)
		FindExtensionByNumber(message protoreflect.FullName, field protoreflect.FieldNumber) (protoreflect.ExtensionType, error)
	}
}

// Unmarshal parses the wire-format message in b and places the result in m.
func Unmarshal(b []byte, m Message) error {
	_, err := UnmarshalOptions{}.unmarshal(b, m)
	return err
}

// Unmarshal parses the wire-format message in b and places the result in m.
func (o UnmarshalOptions) Unmarshal(b []byte, m Message) error {
	_, err := o.unmarshal(b, m)
	return err
}

// UnmarshalState parses a wire-format message and places the result in m.
//
// This method permits fine-grained control over the unmarshaler.
// Most users should use Unmarshal instead.
func (o UnmarshalOptions) UnmarshalState(m Message, in protoiface.UnmarshalInput) (protoiface.UnmarshalOutput, error) {
	return o.unmarshal(in.Buf, m)
}

func (o UnmarshalOptions) unmarshal(b []byte, message Message) (out protoiface.UnmarshalOutput, err error) {
	defaultResolver := false
	if o.Resolver == nil {
		o.Resolver = protoregistry.GlobalTypes
		defaultResolver = true
	} else if o.Resolver == protoregistry.GlobalTypes {
		defaultResolver = true
	}
	if !o.Merge {
		Reset(message)
	}
	allowPartial := o.AllowPartial
	o.Merge = true
	o.AllowPartial = true
	m := message.ProtoReflect()
	methods := protoMethods(m)
	if methods != nil && methods.Unmarshal != nil &&
		!(o.DiscardUnknown && methods.Flags&protoiface.SupportUnmarshalDiscardUnknown == 0) {
		opts := protoiface.UnmarshalOptions{
			Resolver: o.Resolver,
		}
		if o.DiscardUnknown {
			opts.Flags |= protoiface.UnmarshalDiscardUnknown
		}
		if defaultResolver {
			opts.Flags |= protoiface.UnmarshalDefaultResolver
		}
		out, err = methods.Unmarshal(m, protoiface.UnmarshalInput{
			Buf: b,
		}, opts)
	} else {
		err = o.unmarshalMessageSlow(b, m)
	}
	if err != nil {
		return out, err
	}
	if allowPartial || out.Initialized {
		return out, nil
	}
	return out, isInitialized(m)
}

func (o UnmarshalOptions) unmarshalMessage(b []byte, m protoreflect.Message) error {
	_, err := o.unmarshal(b, m.Interface())
	return err
}

func (o UnmarshalOptions) unmarshalMessageSlow(b []byte, m protoreflect.Message) error {
	md := m.Descriptor()
	if messageset.IsMessageSet(md) {
		return unmarshalMessageSet(b, m, o)
	}
	fields := md.Fields()
	for len(b) > 0 {
		// Parse the tag (field number and wire type).
		num, wtyp, tagLen := wire.ConsumeTag(b)
		if tagLen < 0 {
			return wire.ParseError(tagLen)
		}
		if num > wire.MaxValidNumber {
			return errors.New("invalid field number")
		}

		// Find the field descriptor for this field number.
		fd := fields.ByNumber(num)
		if fd == nil && md.ExtensionRanges().Has(num) {
			extType, err := o.Resolver.FindExtensionByNumber(md.FullName(), num)
			if err != nil && err != protoregistry.NotFound {
				return err
			}
			if extType != nil {
				fd = extType.TypeDescriptor()
			}
		}
		var err error
		if fd == nil {
			err = errUnknown
		} else if flags.ProtoLegacy {
			if fd.IsWeak() && fd.Message().IsPlaceholder() {
				err = errUnknown // weak referent is not linked in
			}
		}

		// Parse the field value.
		var valLen int
		switch {
		case err != nil:
		case fd.IsList():
			valLen, err = o.unmarshalList(b[tagLen:], wtyp, m.Mutable(fd).List(), fd)
		case fd.IsMap():
			valLen, err = o.unmarshalMap(b[tagLen:], wtyp, m.Mutable(fd).Map(), fd)
		default:
			valLen, err = o.unmarshalSingular(b[tagLen:], wtyp, m, fd)
		}
		if err != nil {
			if err != errUnknown {
				return err
			}
			valLen = wire.ConsumeFieldValue(num, wtyp, b[tagLen:])
			if valLen < 0 {
				return wire.ParseError(valLen)
			}
			if !o.DiscardUnknown {
				m.SetUnknown(append(m.GetUnknown(), b[:tagLen+valLen]...))
			}
		}
		b = b[tagLen+valLen:]
	}
	return nil
}

func (o UnmarshalOptions) unmarshalSingular(b []byte, wtyp wire.Type, m protoreflect.Message, fd protoreflect.FieldDescriptor) (n int, err error) {
	v, n, err := o.unmarshalScalar(b, wtyp, fd)
	if err != nil {
		return 0, err
	}
	switch fd.Kind() {
	case protoreflect.GroupKind, protoreflect.MessageKind:
		m2 := m.Mutable(fd).Message()
		if err := o.unmarshalMessage(v.Bytes(), m2); err != nil {
			return n, err
		}
	default:
		// Non-message scalars replace the previous value.
		m.Set(fd, v)
	}
	return n, nil
}

func (o UnmarshalOptions) unmarshalMap(b []byte, wtyp wire.Type, mapv protoreflect.Map, fd protoreflect.FieldDescriptor) (n int, err error) {
	if wtyp != wire.BytesType {
		return 0, errUnknown
	}
	b, n = wire.ConsumeBytes(b)
	if n < 0 {
		return 0, wire.ParseError(n)
	}
	var (
		keyField = fd.MapKey()
		valField = fd.MapValue()
		key      protoreflect.Value
		val      protoreflect.Value
		haveKey  bool
		haveVal  bool
	)
	switch valField.Kind() {
	case protoreflect.GroupKind, protoreflect.MessageKind:
		val = mapv.NewValue()
	}
	// Map entries are represented as a two-element message with fields
	// containing the key and value.
	for len(b) > 0 {
		num, wtyp, n := wire.ConsumeTag(b)
		if n < 0 {
			return 0, wire.ParseError(n)
		}
		if num > wire.MaxValidNumber {
			return 0, errors.New("invalid field number")
		}
		b = b[n:]
		err = errUnknown
		switch num {
		case 1:
			key, n, err = o.unmarshalScalar(b, wtyp, keyField)
			if err != nil {
				break
			}
			haveKey = true
		case 2:
			var v protoreflect.Value
			v, n, err = o.unmarshalScalar(b, wtyp, valField)
			if err != nil {
				break
			}
			switch valField.Kind() {
			case protoreflect.GroupKind, protoreflect.MessageKind:
				if err := o.unmarshalMessage(v.Bytes(), val.Message()); err != nil {
					return 0, err
				}
			default:
				val = v
			}
			haveVal = true
		}
		if err == errUnknown {
			n = wire.ConsumeFieldValue(num, wtyp, b)
			if n < 0 {
				return 0, wire.ParseError(n)
			}
		} else if err != nil {
			return 0, err
		}
		b = b[n:]
	}
	// Every map entry should have entries for key and value, but this is not strictly required.
	if !haveKey {
		key = keyField.Default()
	}
	if !haveVal {
		switch valField.Kind() {
		case protoreflect.GroupKind, protoreflect.MessageKind:
		default:
			val = valField.Default()
		}
	}
	mapv.Set(key.MapKey(), val)
	return n, nil
}

// errUnknown is used internally to indicate fields which should be added
// to the unknown field set of a message. It is never returned from an exported
// function.
var errUnknown = errors.New("BUG: internal error (unknown)")
<|endoftext|>


<commit_before>package main

import (
	"debug/dwarf"
	"debug/elf"
	"errors"
	"fmt"
	"log"
	"os"
	"sort"
)

const attrMIPSLinkageName dwarf.Attr = 0x2007

type costInfo struct {
	count uint64 // Number of times the function was inlined.
	bytes uint64 // Total bytes inlined for the function.
}

type subprogramMap map[dwarf.Offset]*dwarf.Entry

// Attempts to extract a function name from the DIE at the provided offset. Unfortunately, since
// it's C++ and DWARF, it's not just a simple matter of getting name attribute and returning it.
func nameForEntry(subprograms subprogramMap, offset dwarf.Offset) (string, error) {
	subprogram, ok := subprograms[offset]
	if !ok {
		return "", errors.New("couldn't find subprogram")
	}

	// TODO(dcheng): Ideally, demangle this name.
	if name, ok := subprogram.Val(attrMIPSLinkageName).(string); ok {
		return name, nil
	}

	// Some DIEs chain to another DIE with the specification attribute.
	// TODO(dcheng): Document when this happens.
	if specOffset, ok := subprogram.Val(dwarf.AttrSpecification).(dwarf.Offset); ok {
		return nameForEntry(subprograms, specOffset)
	}

	// Otherwise, fall back to the name attribute. This is probably a plain C function.
	if name, ok := subprogram.Val(dwarf.AttrName).(string); ok {
		return name, nil
	}

	return "", fmt.Errorf("%v missing name, linkage name, and spec", subprogram)
}

func analyze(data *dwarf.Data) (map[string]*costInfo, error) {
	// DIEs may refer to a DIE with a greater offset, so collect all interesting DIEs first.
	reader := data.Reader()
	subprograms := make(subprogramMap)
	inlinedSubroutines := make([]*dwarf.Entry, 0, 1024)
	for i := 0; ; i++ {
		if i % 1000000 == 0 {
			log.Printf("read %d DIEs...", i)
		}
		entry, err := reader.Next()
		if err != nil {
			return nil, err
		}
		if entry == nil {
			break
		}
		switch entry.Tag {
		case dwarf.TagSubprogram:
			subprograms[entry.Offset] = entry
		case dwarf.TagInlinedSubroutine:
			inlinedSubroutines = append(inlinedSubroutines, entry)
		}
	}

	log.Printf("got %d inlined subroutines", len(inlinedSubroutines))
	results := make(map[string]*costInfo)
	for _, entry := range inlinedSubroutines {
		subprogramOffset, ok := entry.Val(dwarf.AttrAbstractOrigin).(dwarf.Offset)
		if !ok {
			log.Printf("error: %v missing abstract origin", entry)
			continue
		}
		name, err := nameForEntry(subprograms, subprogramOffset)
		if err != nil {
			log.Printf("error: couldn't extract name for %v: %v", entry, err)
		}

		// Despite the name, this is relative to Lowpc, so it's actually a byte count.
		bytes, ok := entry.Val(dwarf.AttrHighpc).(int64)
		if !ok {
			// TODO(dcheng): Handle AttrRanges.
			continue
			log.Printf("error: couldn't extract AttrHighpc for %v", entry)
		}
		if bytes < 0 {
			log.Printf("error: %v has a negative AttrHighpc", entry)
			continue
		}

		info, ok := results[name]
		if !ok {
			info = &costInfo{}
			results[name] = info
		}
		info.count++
		info.bytes += uint64(bytes)
	}
	return results, nil
}

type resultSorter struct {
	// TOOO(dcheng): Bad names are bad.
	keys    []string
	results map[string]*costInfo
}

func (s *resultSorter) Len() int {
	return len(s.keys)
}

func (s *resultSorter) Swap(i, j int) {
	s.keys[i], s.keys[j] = s.keys[j], s.keys[i]
}

func (s *resultSorter) Less(i, j int) bool {
	return s.results[s.keys[i]].bytes > s.results[s.keys[j]].bytes
}

func sortAndPrintTop100(results map[string]*costInfo) {
	keys := make([]string, 0, len(results))
	for k := range results {
		keys = append(keys, k)
	}
	sort.Sort(&resultSorter{keys, results})
	fmt.Printf("     Count      Bytes   Name\n")
	fmt.Printf("  --------  ---------   ---------------------------------\n")
	for i, k := range keys {
		if i > 100 {
			break
		}
		fmt.Printf("%10d %10d   %s\n", results[k].count, results[k].bytes, k)
	}
}

func main() {
	for _, f := range os.Args[1:] {
		log.Printf("analyzing %s...", f)
		elf, err := elf.Open(f)
		if err != nil {
			log.Printf("error: couldn't open %s: %v", f, err)
			continue
		}
		debugData, err := elf.DWARF()
		if err != nil {
			log.Printf("error: couldn't load debug data for %s: %v", f, err)
			continue
		}
		results, err := analyze(debugData)
		if err != nil {
			log.Printf("error: couldn't analyze debug data for %s: %v", f, err)
			continue
		}
		sortAndPrintTop100(results)
	}
}
<commit_msg>Minor renames and preparation for parsing .debug_ranges<commit_after>package main

import (
	"debug/dwarf"
	"debug/elf"
	"errors"
	"fmt"
	"log"
	"os"
	"sort"
)

const attrMIPSLinkageName dwarf.Attr = 0x2007

type subprogramMap map[dwarf.Offset]*dwarf.Entry

// Attempts to extract a function name from the DIE at the provided offset. Unfortunately, since
// it's C++ and DWARF, it's not just a simple matter of getting name attribute and returning it.
func nameForEntry(subprograms subprogramMap, offset dwarf.Offset) (string, error) {
	subprogram, ok := subprograms[offset]
	if !ok {
		return "", errors.New("couldn't find subprogram")
	}

	// TODO(dcheng): Ideally, demangle this name.
	if name, ok := subprogram.Val(attrMIPSLinkageName).(string); ok {
		return name, nil
	}

	// Some DIEs chain to another DIE with the specification attribute.
	// TODO(dcheng): Document when this happens.
	if specOffset, ok := subprogram.Val(dwarf.AttrSpecification).(dwarf.Offset); ok {
		return nameForEntry(subprograms, specOffset)
	}

	// Otherwise, fall back to the name attribute. This is probably a plain C function.
	if name, ok := subprogram.Val(dwarf.AttrName).(string); ok {
		return name, nil
	}

	return "", fmt.Errorf("%v missing name, linkage name, and spec", subprogram)
}

type inlineStats struct {
	count uint64 // Number of times the function was inlined.
	bytes uint64 // Total bytes inlined for the function.
}

func analyze(elf *elf.File) (map[string]*inlineStats, error) {
	// Strictly speaking, dwarf.Data should have other debug sections too, but in practice,
	// only .debug_info is exposed.
	debugInfo, err := elf.DWARF()
	if err != nil {
		return nil, err
	}

	// DIEs may refer to a DIE with a greater offset, so collect all interesting DIEs first.
	reader := debugInfo.Reader()
	subprograms := make(subprogramMap)
	inlinedSubroutines := make([]*dwarf.Entry, 0, 1024)
	for i := 0; ; i++ {
		if i % 1000000 == 0 {
			log.Printf("read %d DIEs...", i)
		}
		entry, err := reader.Next()
		if err != nil {
			return nil, err
		}
		if entry == nil {
			break
		}
		switch entry.Tag {
		case dwarf.TagSubprogram:
			subprograms[entry.Offset] = entry
		case dwarf.TagInlinedSubroutine:
			inlinedSubroutines = append(inlinedSubroutines, entry)
		}
	}

	log.Printf("got %d inlined subroutines", len(inlinedSubroutines))
	results := make(map[string]*inlineStats)
	for _, entry := range inlinedSubroutines {
		subprogramOffset, ok := entry.Val(dwarf.AttrAbstractOrigin).(dwarf.Offset)
		if !ok {
			log.Printf("error: %v missing abstract origin", entry)
			continue
		}
		name, err := nameForEntry(subprograms, subprogramOffset)
		if err != nil {
			log.Printf("error: couldn't extract name for %v: %v", entry, err)
		}

		// Despite the name, this is relative to Lowpc, so it's actually a byte count.
		bytes, ok := entry.Val(dwarf.AttrHighpc).(int64)
		if !ok {
			// TODO(dcheng): Handle AttrRanges.
			continue
			log.Printf("error: couldn't extract AttrHighpc for %v", entry)
		}
		if bytes < 0 {
			log.Printf("error: %v has a negative AttrHighpc", entry)
			continue
		}

		info, ok := results[name]
		if !ok {
			info = &inlineStats{}
			results[name] = info
		}
		info.count++
		info.bytes += uint64(bytes)
	}
	return results, nil
}

type resultSorter struct {
	// TOOO(dcheng): Bad names are bad.
	keys    []string
	results map[string]*inlineStats
}

func (s *resultSorter) Len() int {
	return len(s.keys)
}

func (s *resultSorter) Swap(i, j int) {
	s.keys[i], s.keys[j] = s.keys[j], s.keys[i]
}

func (s *resultSorter) Less(i, j int) bool {
	return s.results[s.keys[i]].bytes > s.results[s.keys[j]].bytes
}

func sortAndPrintTop100(results map[string]*inlineStats) {
	keys := make([]string, 0, len(results))
	for k := range results {
		keys = append(keys, k)
	}
	sort.Sort(&resultSorter{keys, results})
	fmt.Printf("     Count      Bytes   Name\n")
	fmt.Printf("  --------  ---------   ---------------------------------\n")
	for i, k := range keys {
		if i > 100 {
			break
		}
		fmt.Printf("%10d %10d   %s\n", results[k].count, results[k].bytes, k)
	}
}

func main() {
	for _, f := range os.Args[1:] {
		log.Printf("analyzing %s...", f)
		elf, err := elf.Open(f)
		if err != nil {
			log.Printf("error: couldn't open %s: %v", f, err)
			continue
		}
		results, err := analyze(elf)
		if err != nil {
			log.Printf("error: couldn't analyze debug data for %s: %v", f, err)
			continue
		}
		sortAndPrintTop100(results)
	}
}
<|endoftext|>