<commit_before>// Copyright 2014 gandalf 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 repository

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/tsuru/config"
	"github.com/tsuru/gandalf/db"
	"github.com/tsuru/gandalf/fs"
	"github.com/tsuru/tsuru/log"
	"io/ioutil"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"os/exec"
	"regexp"
	"strings"
)

// Repository represents a Git repository. A Git repository is a record in the
// database and a directory in the filesystem (the bare repository).
type Repository struct {
	Name     string `bson:"_id"`
	Users    []string
	IsPublic bool
}

// MarshalJSON marshals the Repository in json format.
func (r *Repository) MarshalJSON() ([]byte, error) {
	data := map[string]interface{}{
		"name":    r.Name,
		"public":  r.IsPublic,
		"ssh_url": r.ReadWriteURL(),
		"git_url": r.ReadOnlyURL(),
	}
	return json.Marshal(&data)
}

// New creates a representation of a git repository. It creates a Git
// repository using the "bare-dir" setting and saves repository's meta data in
// the database.
func New(name string, users []string, isPublic bool) (*Repository, error) {
	log.Debugf("Creating repository %q", name)
	r := &Repository{Name: name, Users: users, IsPublic: isPublic}
	if v, err := r.isValid(); !v {
		log.Errorf("repository.New: Invalid repository %q: %s", name, err)
		return r, err
	}
	if err := newBare(name); err != nil {
		log.Errorf("repository.New: Error creating bare repository for %q: %s", name, err)
		return r, err
	}
	barePath := barePath(name)
	if barePath != "" && isPublic {
		ioutil.WriteFile(barePath+"/git-daemon-export-ok", []byte(""), 0644)
		if f, err := fs.Filesystem().Create(barePath + "/git-daemon-export-ok"); err == nil {
			f.Close()
		}
	}
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	err = conn.Repository().Insert(&r)
	if mgo.IsDup(err) {
		log.Errorf("repository.New: Duplicate repository %q", name)
		return r, fmt.Errorf("A repository with this name already exists.")
	}
	return r, err
}

// Get find a repository by name.
func Get(name string) (Repository, error) {
	var r Repository
	conn, err := db.Conn()
	if err != nil {
		return r, err
	}
	defer conn.Close()
	err = conn.Repository().FindId(name).One(&r)
	return r, err
}

// Remove deletes the repository from the database and removes it's bare Git
// repository.
func Remove(name string) error {
	log.Debugf("Removing repository %q", name)
	if err := removeBare(name); err != nil {
		log.Errorf("repository.Remove: Error removing bare repository %q: %s", name, err)
		return err
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	if err := conn.Repository().RemoveId(name); err != nil {
		log.Errorf("repository.Remove: Error removing repository %q from db: %s", name, err)
		return fmt.Errorf("Could not remove repository: %s", err)
	}
	return nil
}

// Rename renames a repository.
func Rename(oldName, newName string) error {
	log.Debugf("Renaming repository %q to %q", oldName, newName)
	repo, err := Get(oldName)
	if err != nil {
		log.Errorf("repository.Rename: Repository %q not found: %s", oldName, err)
		return err
	}
	newRepo := repo
	newRepo.Name = newName
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	err = conn.Repository().Insert(newRepo)
	if err != nil {
		log.Errorf("repository.Rename: Error adding new repository %q: %s", newName, err)
		return err
	}
	err = conn.Repository().RemoveId(oldName)
	if err != nil {
		log.Errorf("repository.Rename: Error removing old repository %q: %s", oldName, err)
		return err
	}
	return fs.Filesystem().Rename(barePath(oldName), barePath(newName))
}

// ReadWriteURL formats the git ssh url and return it. If no remote is configured in
// gandalf.conf, this method panics.
func (r *Repository) ReadWriteURL() string {
	uid, err := config.GetString("uid")
	if err != nil {
		panic(err.Error())
	}
	remote := uid + "@%s:%s.git"
	if useSSH, _ := config.GetBool("git:ssh:use"); useSSH {
		port, err := config.GetString("git:ssh:port")
		if err == nil {
			remote = "ssh://" + uid + "@%s:" + port + "/%s.git"
		} else {
			remote = "ssh://" + uid + "@%s/%s.git"
		}
	}
	host, err := config.GetString("host")
	if err != nil {
		panic(err.Error())
	}
	return fmt.Sprintf(remote, host, r.Name)
}

// ReadOnly formats the git url and return it. If no host is configured in
// gandalf.conf, this method panics.
func (r *Repository) ReadOnlyURL() string {
	remote := "git://%s/%s.git"
	if useSSH, _ := config.GetBool("git:ssh:use"); useSSH {
		uid, err := config.GetString("uid")
		if err != nil {
			panic(err.Error())
		}
		port, err := config.GetString("git:ssh:port")
		if err == nil {
			remote = "ssh://" + uid + "@%s:" + port + "/%s.git"
		} else {
			remote = "ssh://" + uid + "@%s/%s.git"
		}
	}
	host, err := config.GetString("readonly-host")
	if err != nil {
		host, err = config.GetString("host")
		if err != nil {
			panic(err)
		}
	}
	return fmt.Sprintf(remote, host, r.Name)
}

// Validates a repository
// A valid repository must have:
//  - a name without any special chars only alphanumeric and underlines are allowed.
//  - at least one user in users array
func (r *Repository) isValid() (bool, error) {
	m, e := regexp.Match(`^[\w-]+$`, []byte(r.Name))
	if e != nil {
		panic(e)
	}
	if !m {
		return false, errors.New("Validation Error: repository name is not valid")
	}
	if len(r.Users) == 0 {
		return false, errors.New("Validation Error: repository should have at least one user")
	}
	return true, nil
}

// GrantAccess gives write permission for users in all specified repositories.
// If any of the repositories/users do not exists, GrantAccess just skips it.
func GrantAccess(rNames, uNames []string) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	_, err = conn.Repository().UpdateAll(bson.M{"_id": bson.M{"$in": rNames}}, bson.M{"$addToSet": bson.M{"users": bson.M{"$each": uNames}}})
	return err
}

// RevokeAccess revokes write permission from users in all specified
// repositories.
func RevokeAccess(rNames, uNames []string) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	_, err = conn.Repository().UpdateAll(bson.M{"_id": bson.M{"$in": rNames}}, bson.M{"$pullAll": bson.M{"users": uNames}})
	return err
}

type ArchiveFormat int

const (
	Zip ArchiveFormat = iota
	Tar
	TarGz
)

type ContentRetriever interface {
	GetContents(repo, ref, path string) ([]byte, error)
	GetArchive(repo, ref string, format ArchiveFormat) ([]byte, error)
	GetTree(repo, ref, path string) ([]map[string]string, error)
}

var Retriever ContentRetriever

type GitContentRetriever struct{}

func (*GitContentRetriever) GetContents(repo, ref, path string) ([]byte, error) {
	gitPath, err := exec.LookPath("git")
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain file %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	cwd := barePath(repo)
	cmd := exec.Command(gitPath, "show", fmt.Sprintf("%s:%s", ref, path))
	cmd.Dir = cwd
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain file %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	return out, nil
}

func (*GitContentRetriever) GetArchive(repo, ref string, format ArchiveFormat) ([]byte, error) {
	gitPath, err := exec.LookPath("git")
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain archive for ref %s of repository %s (%s).", ref, repo, err)
	}
	var archiveFormat string
	switch format {
	case Tar:
		archiveFormat = "--format=tar"
	case TarGz:
		archiveFormat = "--format=tar.gz"
	default:
		archiveFormat = "--format=zip"
	}
	prefix := fmt.Sprintf("--prefix=%s-%s/", repo, ref)
	cwd := barePath(repo)
	cmd := exec.Command(gitPath, "archive", ref, prefix, archiveFormat)
	cmd.Dir = cwd
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain archive for ref %s of repository %s (%s).", ref, repo, err)
	}
	return out, nil
}

func (*GitContentRetriever) GetTree(repo, ref, path string) ([]map[string]string, error) {
	gitPath, err := exec.LookPath("git")
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain file %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	cwd := barePath(repo)
	cmd := exec.Command(gitPath, "ls-tree", "-r", ref, path)
	cmd.Dir = cwd
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain tree %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	lines := strings.Split(string(out), "\n")
	objectCount := 0
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		objectCount++
	}
	objects := make([]map[string]string, len(lines)-1)
	objectCount = 0
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		tabbed := strings.Split(line, "\t")
		meta, filepath := tabbed[0], tabbed[1]
		meta_parts := strings.Split(meta, " ")
		permission, filetype, hash := meta_parts[0], meta_parts[1], meta_parts[2]
		object := make(map[string]string)
		object["permission"] = permission
		object["filetype"] = filetype
		object["hash"] = hash
		object["path"] = strings.TrimSpace(strings.Trim(filepath, "\""))
		object["rawPath"] = filepath
		objects[objectCount] = object
		objectCount++
	}
	return objects, nil
}

func retriever() ContentRetriever {
	if Retriever == nil {
		Retriever = &GitContentRetriever{}
	}
	return Retriever
}

// GetFileContents returns the contents for a given file
// in a given ref for the specified repository
func GetFileContents(repo, ref, path string) ([]byte, error) {
	return retriever().GetContents(repo, ref, path)
}

// GetArchive returns the contents for a given file
// in a given ref for the specified repository
func GetArchive(repo, ref string, format ArchiveFormat) ([]byte, error) {
	return retriever().GetArchive(repo, ref, format)
}

func GetTree(repo, ref, path string) ([]map[string]string, error) {
	return retriever().GetTree(repo, ref, path)
}
<commit_msg>using proper count<commit_after>// Copyright 2014 gandalf 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 repository

import (
	"encoding/json"
	"errors"
	"fmt"
	"github.com/tsuru/config"
	"github.com/tsuru/gandalf/db"
	"github.com/tsuru/gandalf/fs"
	"github.com/tsuru/tsuru/log"
	"io/ioutil"
	"labix.org/v2/mgo"
	"labix.org/v2/mgo/bson"
	"os/exec"
	"regexp"
	"strings"
)

// Repository represents a Git repository. A Git repository is a record in the
// database and a directory in the filesystem (the bare repository).
type Repository struct {
	Name     string `bson:"_id"`
	Users    []string
	IsPublic bool
}

// MarshalJSON marshals the Repository in json format.
func (r *Repository) MarshalJSON() ([]byte, error) {
	data := map[string]interface{}{
		"name":    r.Name,
		"public":  r.IsPublic,
		"ssh_url": r.ReadWriteURL(),
		"git_url": r.ReadOnlyURL(),
	}
	return json.Marshal(&data)
}

// New creates a representation of a git repository. It creates a Git
// repository using the "bare-dir" setting and saves repository's meta data in
// the database.
func New(name string, users []string, isPublic bool) (*Repository, error) {
	log.Debugf("Creating repository %q", name)
	r := &Repository{Name: name, Users: users, IsPublic: isPublic}
	if v, err := r.isValid(); !v {
		log.Errorf("repository.New: Invalid repository %q: %s", name, err)
		return r, err
	}
	if err := newBare(name); err != nil {
		log.Errorf("repository.New: Error creating bare repository for %q: %s", name, err)
		return r, err
	}
	barePath := barePath(name)
	if barePath != "" && isPublic {
		ioutil.WriteFile(barePath+"/git-daemon-export-ok", []byte(""), 0644)
		if f, err := fs.Filesystem().Create(barePath + "/git-daemon-export-ok"); err == nil {
			f.Close()
		}
	}
	conn, err := db.Conn()
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	err = conn.Repository().Insert(&r)
	if mgo.IsDup(err) {
		log.Errorf("repository.New: Duplicate repository %q", name)
		return r, fmt.Errorf("A repository with this name already exists.")
	}
	return r, err
}

// Get find a repository by name.
func Get(name string) (Repository, error) {
	var r Repository
	conn, err := db.Conn()
	if err != nil {
		return r, err
	}
	defer conn.Close()
	err = conn.Repository().FindId(name).One(&r)
	return r, err
}

// Remove deletes the repository from the database and removes it's bare Git
// repository.
func Remove(name string) error {
	log.Debugf("Removing repository %q", name)
	if err := removeBare(name); err != nil {
		log.Errorf("repository.Remove: Error removing bare repository %q: %s", name, err)
		return err
	}
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	if err := conn.Repository().RemoveId(name); err != nil {
		log.Errorf("repository.Remove: Error removing repository %q from db: %s", name, err)
		return fmt.Errorf("Could not remove repository: %s", err)
	}
	return nil
}

// Rename renames a repository.
func Rename(oldName, newName string) error {
	log.Debugf("Renaming repository %q to %q", oldName, newName)
	repo, err := Get(oldName)
	if err != nil {
		log.Errorf("repository.Rename: Repository %q not found: %s", oldName, err)
		return err
	}
	newRepo := repo
	newRepo.Name = newName
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	err = conn.Repository().Insert(newRepo)
	if err != nil {
		log.Errorf("repository.Rename: Error adding new repository %q: %s", newName, err)
		return err
	}
	err = conn.Repository().RemoveId(oldName)
	if err != nil {
		log.Errorf("repository.Rename: Error removing old repository %q: %s", oldName, err)
		return err
	}
	return fs.Filesystem().Rename(barePath(oldName), barePath(newName))
}

// ReadWriteURL formats the git ssh url and return it. If no remote is configured in
// gandalf.conf, this method panics.
func (r *Repository) ReadWriteURL() string {
	uid, err := config.GetString("uid")
	if err != nil {
		panic(err.Error())
	}
	remote := uid + "@%s:%s.git"
	if useSSH, _ := config.GetBool("git:ssh:use"); useSSH {
		port, err := config.GetString("git:ssh:port")
		if err == nil {
			remote = "ssh://" + uid + "@%s:" + port + "/%s.git"
		} else {
			remote = "ssh://" + uid + "@%s/%s.git"
		}
	}
	host, err := config.GetString("host")
	if err != nil {
		panic(err.Error())
	}
	return fmt.Sprintf(remote, host, r.Name)
}

// ReadOnly formats the git url and return it. If no host is configured in
// gandalf.conf, this method panics.
func (r *Repository) ReadOnlyURL() string {
	remote := "git://%s/%s.git"
	if useSSH, _ := config.GetBool("git:ssh:use"); useSSH {
		uid, err := config.GetString("uid")
		if err != nil {
			panic(err.Error())
		}
		port, err := config.GetString("git:ssh:port")
		if err == nil {
			remote = "ssh://" + uid + "@%s:" + port + "/%s.git"
		} else {
			remote = "ssh://" + uid + "@%s/%s.git"
		}
	}
	host, err := config.GetString("readonly-host")
	if err != nil {
		host, err = config.GetString("host")
		if err != nil {
			panic(err)
		}
	}
	return fmt.Sprintf(remote, host, r.Name)
}

// Validates a repository
// A valid repository must have:
//  - a name without any special chars only alphanumeric and underlines are allowed.
//  - at least one user in users array
func (r *Repository) isValid() (bool, error) {
	m, e := regexp.Match(`^[\w-]+$`, []byte(r.Name))
	if e != nil {
		panic(e)
	}
	if !m {
		return false, errors.New("Validation Error: repository name is not valid")
	}
	if len(r.Users) == 0 {
		return false, errors.New("Validation Error: repository should have at least one user")
	}
	return true, nil
}

// GrantAccess gives write permission for users in all specified repositories.
// If any of the repositories/users do not exists, GrantAccess just skips it.
func GrantAccess(rNames, uNames []string) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	_, err = conn.Repository().UpdateAll(bson.M{"_id": bson.M{"$in": rNames}}, bson.M{"$addToSet": bson.M{"users": bson.M{"$each": uNames}}})
	return err
}

// RevokeAccess revokes write permission from users in all specified
// repositories.
func RevokeAccess(rNames, uNames []string) error {
	conn, err := db.Conn()
	if err != nil {
		return err
	}
	defer conn.Close()
	_, err = conn.Repository().UpdateAll(bson.M{"_id": bson.M{"$in": rNames}}, bson.M{"$pullAll": bson.M{"users": uNames}})
	return err
}

type ArchiveFormat int

const (
	Zip ArchiveFormat = iota
	Tar
	TarGz
)

type ContentRetriever interface {
	GetContents(repo, ref, path string) ([]byte, error)
	GetArchive(repo, ref string, format ArchiveFormat) ([]byte, error)
	GetTree(repo, ref, path string) ([]map[string]string, error)
}

var Retriever ContentRetriever

type GitContentRetriever struct{}

func (*GitContentRetriever) GetContents(repo, ref, path string) ([]byte, error) {
	gitPath, err := exec.LookPath("git")
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain file %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	cwd := barePath(repo)
	cmd := exec.Command(gitPath, "show", fmt.Sprintf("%s:%s", ref, path))
	cmd.Dir = cwd
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain file %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	return out, nil
}

func (*GitContentRetriever) GetArchive(repo, ref string, format ArchiveFormat) ([]byte, error) {
	gitPath, err := exec.LookPath("git")
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain archive for ref %s of repository %s (%s).", ref, repo, err)
	}
	var archiveFormat string
	switch format {
	case Tar:
		archiveFormat = "--format=tar"
	case TarGz:
		archiveFormat = "--format=tar.gz"
	default:
		archiveFormat = "--format=zip"
	}
	prefix := fmt.Sprintf("--prefix=%s-%s/", repo, ref)
	cwd := barePath(repo)
	cmd := exec.Command(gitPath, "archive", ref, prefix, archiveFormat)
	cmd.Dir = cwd
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain archive for ref %s of repository %s (%s).", ref, repo, err)
	}
	return out, nil
}

func (*GitContentRetriever) GetTree(repo, ref, path string) ([]map[string]string, error) {
	gitPath, err := exec.LookPath("git")
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain file %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	cwd := barePath(repo)
	cmd := exec.Command(gitPath, "ls-tree", "-r", ref, path)
	cmd.Dir = cwd
	out, err := cmd.Output()
	if err != nil {
		return nil, fmt.Errorf("Error when trying to obtain tree %s on ref %s of repository %s (%s).", path, ref, repo, err)
	}
	lines := strings.Split(string(out), "\n")
	objectCount := 0
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		objectCount++
	}
	objects := make([]map[string]string, objectCount)
	objectCount = 0
	for _, line := range lines {
		if strings.TrimSpace(line) == "" {
			continue
		}
		tabbed := strings.Split(line, "\t")
		meta, filepath := tabbed[0], tabbed[1]
		meta_parts := strings.Split(meta, " ")
		permission, filetype, hash := meta_parts[0], meta_parts[1], meta_parts[2]
		object := make(map[string]string)
		object["permission"] = permission
		object["filetype"] = filetype
		object["hash"] = hash
		object["path"] = strings.TrimSpace(strings.Trim(filepath, "\""))
		object["rawPath"] = filepath
		objects[objectCount] = object
		objectCount++
	}
	return objects, nil
}

func retriever() ContentRetriever {
	if Retriever == nil {
		Retriever = &GitContentRetriever{}
	}
	return Retriever
}

// GetFileContents returns the contents for a given file
// in a given ref for the specified repository
func GetFileContents(repo, ref, path string) ([]byte, error) {
	return retriever().GetContents(repo, ref, path)
}

// GetArchive returns the contents for a given file
// in a given ref for the specified repository
func GetArchive(repo, ref string, format ArchiveFormat) ([]byte, error) {
	return retriever().GetArchive(repo, ref, format)
}

func GetTree(repo, ref, path string) ([]map[string]string, error) {
	return retriever().GetTree(repo, ref, path)
}
<|endoftext|>


<commit_before>package scm

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
)

func GitRootPath(path ...string) (string, error) {
	cmd := exec.Command("git", "rev-parse", "--show-toplevel")
	if len(path) > 0 {
		cmd.Dir = path[0]
	}

	b, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("Unable to parse repository path, %s %s", string(b), err)
	}

	s := strings.TrimSpace(string(b))
	if s == "" {
		return "", fmt.Errorf("Unable to parse repository path, %s", err)
	}

	return s, nil
}

func GitBranch(path ...string) (string, error) {
	cmd := exec.Command("git", "rev-parse", "--abbrev-ref", "HEAD")
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return "", fmt.Errorf("Unable to parse branch name, %s %s", string(b), err)
	}
	return strings.TrimSpace(string(b)), nil

}

func GitEmail(path ...string) (string, error) {
	cmd := exec.Command("git", "config", "--get", "user.email")
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return "", fmt.Errorf("Unable to get user email, %s %s", string(b), err)
	}
	return strings.TrimSpace(string(b)), nil
}

func GitCommitMsg(path ...string) (string, error) {
	cmd := exec.Command("git", "log", "-1", "--oneline", "--raw")
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		// if there are no git commits yet it will fail
		// ignoring this error
		return "", nil
	}
	return string(b), err
}

func GitParseMessage(m string) (uuid, msg string, files []string) {
	l := strings.Split(m, "\n")
	files = make([]string, 0)
	for i, v := range l {
		if i == 0 {
			s := strings.SplitN(v, " ", 2)
			uuid = s[0]
			msg = s[1]
		} else {
			if strings.TrimSpace(v) != "" {
				s := strings.Split(v, "\t")
				files = append(files, s[1])
			}
		}
	}
	return
}

func GitAddNote(n string, nameSpace string, path ...string) error {
	cmd := exec.Command("git", "notes", fmt.Sprintf("--ref=%s", nameSpace), "add", "-f", "-m", n)
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	if b, err := cmd.Output(); err != nil {
		return fmt.Errorf("Unable to add git note, %s %s", string(b), err)
	}
	return nil
}

func GitSetRewriteRef(ref string, path ...string) error {
	cmd := exec.Command("git", "config", "-l")
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return fmt.Errorf("Unable to run git config -l notes.rewriteref, %s %s", string(b), err)
	}
	if !strings.Contains(string(b), ref+"\n") {
		cmd := exec.Command("git", "config", "--add", "notes.rewriteref", ref)
		if len(path) > 0 {
			cmd.Dir = path[0]
		}
		if b, err := cmd.Output(); err != nil {
			return fmt.Errorf("Unable to run git config --add notes.rewriteref %s, %s %s", ref, string(b), err)
		}
	}
	return nil
}

func GitTracked(f string, path ...string) (bool, error) {
	cmd := exec.Command("git", "ls-files", f)
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return false, fmt.Errorf("Unable to determine git tracked status for %s, %s %s", f, string(b), err)
	}
	return strings.TrimSpace(string(b)) != "", nil
}

func GitModified(f string, path ...string) (bool, error) {
	cmd := exec.Command("git", "ls-files", "-m", f)
	if len(path) > 0 {
		cmd.Dir = path[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return false, fmt.Errorf("Unable to determine git modified status for %s, %s %s", f, string(b), err)
	}
	return strings.TrimSpace(string(b)) != "", nil
}

func GitInitHook(hook, command string, wd ...string) error {
	var (
		p   string
		err error
	)

	if len(wd) > 0 {
		p = wd[0]
	} else {
		p, err = os.Getwd()
		if err != nil {
			return err
		}
	}
	fp := path.Join(p, ".git", "hooks", hook)

	var output string
	if _, err := os.Stat(fp); !os.IsNotExist(err) {
		b, err := ioutil.ReadFile(fp)
		if err != nil {
			return err
		}
		output = string(b)

		if strings.Contains(output, command+"\n") {
			// if file already exists this will make sure it's executable
			if err := os.Chmod(fp, 0755); err != nil {
				return err
			}
			return nil
		}
	}

	if err = ioutil.WriteFile(
		fp, []byte(fmt.Sprintf("%s\n%s\n", output, command)), 0755); err != nil {
		return err
	}
	// if file already exists this will make sure it's executable
	if err := os.Chmod(fp, 0755); err != nil {
		return err
	}

	return nil
}

func GitIgnore(ignore string, wd ...string) error {
	var (
		p   string
		err error
	)

	if len(wd) > 0 {
		p = wd[0]
	} else {
		p, err = os.Getwd()
		if err != nil {
			return err
		}
	}
	fp := path.Join(p, ".gitignore")

	var output string
	if _, err := os.Stat(fp); !os.IsNotExist(err) {
		b, err := ioutil.ReadFile(fp)
		if err != nil {
			return err
		}
		output = string(b)

		if strings.Contains(output, ignore+"\n") {
			return nil
		}
	}

	if err = ioutil.WriteFile(
		fp, []byte(fmt.Sprintf("%s\n%s\n", output, ignore)), 0644); err != nil {
		return err
	}
	return nil
}
<commit_msg>Use wd instead of path for argument name<commit_after>package scm

import (
	"fmt"
	"io/ioutil"
	"os"
	"os/exec"
	"path"
	"strings"
)

func GitRootPath(path ...string) (string, error) {
	cmd := exec.Command("git", "rev-parse", "--show-toplevel")
	if len(path) > 0 {
		cmd.Dir = path[0]
	}

	b, err := cmd.Output()
	if err != nil {
		return "", fmt.Errorf("Unable to parse repository path, %s %s", string(b), err)
	}

	s := strings.TrimSpace(string(b))
	if s == "" {
		return "", fmt.Errorf("Unable to parse repository path, %s", err)
	}

	return s, nil
}

func GitBranch(wd ...string) (string, error) {
	cmd := exec.Command("git", "rev-parse", "--abbrev-ref", "HEAD")
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return "", fmt.Errorf("Unable to parse branch name, %s %s", string(b), err)
	}
	return strings.TrimSpace(string(b)), nil

}

func GitEmail(wd ...string) (string, error) {
	cmd := exec.Command("git", "config", "--get", "user.email")
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return "", fmt.Errorf("Unable to get user email, %s %s", string(b), err)
	}
	return strings.TrimSpace(string(b)), nil
}

func GitCommitMsg(wd ...string) (string, error) {
	cmd := exec.Command("git", "log", "-1", "--oneline", "--raw")
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		// if there are no git commits yet it will fail
		// ignoring this error
		return "", nil
	}
	return string(b), err
}

func GitParseMessage(m string) (uuid, msg string, files []string) {
	l := strings.Split(m, "\n")
	files = make([]string, 0)
	for i, v := range l {
		if i == 0 {
			s := strings.SplitN(v, " ", 2)
			uuid = s[0]
			msg = s[1]
		} else {
			if strings.TrimSpace(v) != "" {
				s := strings.Split(v, "\t")
				files = append(files, s[1])
			}
		}
	}
	return
}

func GitAddNote(n string, nameSpace string, wd ...string) error {
	cmd := exec.Command("git", "notes", fmt.Sprintf("--ref=%s", nameSpace), "add", "-f", "-m", n)
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	if b, err := cmd.Output(); err != nil {
		return fmt.Errorf("Unable to add git note, %s %s", string(b), err)
	}
	return nil
}

func GitSetRewriteRef(ref string, wd ...string) error {
	cmd := exec.Command("git", "config", "-l")
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return fmt.Errorf("Unable to run git config -l notes.rewriteref, %s %s", string(b), err)
	}
	if !strings.Contains(string(b), ref+"\n") {
		cmd := exec.Command("git", "config", "--add", "notes.rewriteref", ref)
		if len(wd) > 0 {
			cmd.Dir = wd[0]
		}
		if b, err := cmd.Output(); err != nil {
			return fmt.Errorf("Unable to run git config --add notes.rewriteref %s, %s %s", ref, string(b), err)
		}
	}
	return nil
}

func GitTracked(f string, wd ...string) (bool, error) {
	cmd := exec.Command("git", "ls-files", f)
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return false, fmt.Errorf("Unable to determine git tracked status for %s, %s %s", f, string(b), err)
	}
	return strings.TrimSpace(string(b)) != "", nil
}

func GitModified(f string, wd ...string) (bool, error) {
	cmd := exec.Command("git", "ls-files", "-m", f)
	if len(wd) > 0 {
		cmd.Dir = wd[0]
	}
	var (
		b   []byte
		err error
	)
	if b, err = cmd.Output(); err != nil {
		return false, fmt.Errorf("Unable to determine git modified status for %s, %s %s", f, string(b), err)
	}
	return strings.TrimSpace(string(b)) != "", nil
}

func GitInitHook(hook, command string, wd ...string) error {
	var (
		p   string
		err error
	)

	if len(wd) > 0 {
		p = wd[0]
	} else {
		p, err = os.Getwd()
		if err != nil {
			return err
		}
	}
	fp := path.Join(p, ".git", "hooks", hook)

	var output string
	if _, err := os.Stat(fp); !os.IsNotExist(err) {
		b, err := ioutil.ReadFile(fp)
		if err != nil {
			return err
		}
		output = string(b)

		if strings.Contains(output, command+"\n") {
			// if file already exists this will make sure it's executable
			if err := os.Chmod(fp, 0755); err != nil {
				return err
			}
			return nil
		}
	}

	if err = ioutil.WriteFile(
		fp, []byte(fmt.Sprintf("%s\n%s\n", output, command)), 0755); err != nil {
		return err
	}
	// if file already exists this will make sure it's executable
	if err := os.Chmod(fp, 0755); err != nil {
		return err
	}

	return nil
}

func GitIgnore(ignore string, wd ...string) error {
	var (
		p   string
		err error
	)

	if len(wd) > 0 {
		p = wd[0]
	} else {
		p, err = os.Getwd()
		if err != nil {
			return err
		}
	}
	fp := path.Join(p, ".gitignore")

	var output string
	if _, err := os.Stat(fp); !os.IsNotExist(err) {
		b, err := ioutil.ReadFile(fp)
		if err != nil {
			return err
		}
		output = string(b)

		if strings.Contains(output, ignore+"\n") {
			return nil
		}
	}

	if err = ioutil.WriteFile(
		fp, []byte(fmt.Sprintf("%s\n%s\n", output, ignore)), 0644); err != nil {
		return err
	}
	return nil
}
<|endoftext|>


<commit_before>package gotana

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	URL "net/url"
	"strings"
	"sync"
	"time"
)

const (
	EVENT_SCRAPER_OPENED     = "SCRAPER_OPENED"
	EVENT_SCRAPER_CLOSED     = "SCRAPER_CLOSED"
	EVENT_SAVEABLE_EXTRACTED = "SAVEABLE_EXTRACTED"
	TIMEOUT_DIALER           = time.Duration(time.Second * 30)
	TIMEOUT_REQUEST          = time.Duration(time.Second * 30)
	TIMEOUT_TLS              = time.Duration(time.Second * 10)
)

type SaveableItem interface {
	Scraper() *Scraper
	Validate() bool
	RecordData() []string
}

type ScraperMixin struct {
	Proxy ScrapedItem
}

func (s *ScraperMixin) SetProxy(proxy ScrapedItem) *ScraperMixin {
	s.Proxy = proxy
	return s
}

func (item ScraperMixin) Scraper() *Scraper {
	return item.Proxy.scraper
}

type recordWriter interface {
	Write(record []string) error
	Flush()
}

type ScrapingHandlerFunc func(ScrapedItem, chan<- SaveableItem)

func GetHref(t html.Token) (ok bool, href string) {
	for _, a := range t.Attr {
		if a.Key == "href" {
			href = a.Val
			ok = true
		}
	}

	return
}

type extensionParameters struct {
	scraper *Scraper
	item    SaveableItem
}

type Extractable interface {
	Extract(io.ReadCloser, func(string))
}

type LinkExtractor struct {
	Extractable
}

func (extractor *LinkExtractor) Extract(r io.ReadCloser, callback func(string)) {
	z := html.NewTokenizer(r)
	defer r.Close()

	for {
		tt := z.Next()

		switch {
		case tt == html.ErrorToken:
			return
		case tt == html.StartTagToken:
			t := z.Token()

			isAnchor := t.Data == "a"
			if !isAnchor {
				continue
			}

			ok, url := GetHref(t)
			if !ok {
				continue
			}
			callback(url)
		}
	}
}

type ScraperConfig struct {
	Project     string `required:"true"`
	TcpAddress  string
	OutFileName string
	Scrapers    []struct {
		RequestLimit int `required:"true"`
		Extractor    string
		Name         string `required:"true"`
		Url          string `required:"true"`
	}
}

type ScrapedItem struct {
	Url       string
	FinalUrl  string
	scraper   *Scraper
	BodyBytes []byte
}

func (proxy ScrapedItem) String() (result string) {
	result = fmt.Sprintf("Result of scraping: %s", proxy.Url)
	return
}

func (proxy ScrapedItem) CheckIfRedirected() bool {
	return proxy.Url != proxy.FinalUrl
}

func (proxy ScrapedItem) FinalResponseBody() (io.ReadCloser, error) {
	if proxy.CheckIfRedirected() {
		client := NewHTTPClient()
		response, err := client.Get(proxy.FinalUrl)
		if err != nil {
			return nil, err
		}
		bodyBytes, _ := ioutil.ReadAll(response.Body)
		proxy.BodyBytes = bodyBytes
	}
	return ioutil.NopCloser(bytes.NewBuffer(proxy.BodyBytes)), nil
}

func (proxy ScrapedItem) HTMLDocument() (document *goquery.Document, err error) {
	responseBody, err := proxy.FinalResponseBody()
	if err == nil {
		document, err = goquery.NewDocumentFromReader(responseBody)
	}

	return
}

type Scraper struct {
	crawled      int
	successful   int
	failed       int
	handler      ScrapingHandlerFunc
	fetchMutex   *sync.Mutex
	crawledMutex *sync.Mutex
	Name         string
	Domain       string
	BaseUrl      string
	CurrentUrl   string
	fetchedUrls  map[string]bool
	engine       *Engine
	extractor    Extractable
	chDone       chan struct{}
	chRequestUrl chan string
	requestLimit int
}

func (scraper *Scraper) MarkAsFetched(url string) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	scraper.CurrentUrl = url
	scraper.fetchedUrls[url] = true
}

func (scraper *Scraper) CheckIfShouldStop() (ok bool) {
	scraper.crawledMutex.Lock()
	defer scraper.crawledMutex.Unlock()
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]

	if stats.crawled == scraper.engine.limitCrawl {
		Logger().Warningf("Crawl limit exceeded: %s", scraper)
		ok = true
	} else if stats.failed == scraper.engine.limitFail {
		Logger().Warningf("Fail limit exceeeded: %s", scraper)
		ok = true
	} else if stats.failed == 1 && scraper.crawled == 1 {
		Logger().Warningf("Base URL is corrupted: %s", scraper)
		ok = true
	}
	return
}

func (scraper *Scraper) CheckIfFetched(url string) (ok bool) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	_, ok = scraper.fetchedUrls[url]
	return
}

func (scraper *Scraper) CheckUrl(sourceUrl string) (ok bool, url string) {
	if strings.Contains(sourceUrl, scraper.Domain) && strings.Index(sourceUrl, "http") == 0 {
		url = sourceUrl
		ok = true
	} else if strings.Index(sourceUrl, "/") == 0 {
		url = scraper.BaseUrl + sourceUrl
		ok = true
	}
	return
}

func (scraper *Scraper) RunExtractor(resp *http.Response) {
	defer SilentRecover("EXTRACTOR")

	scraper.extractor.Extract(resp.Body, func(url string) {
		ok, url := scraper.CheckUrl(url)

		if ok {
			scraper.chRequestUrl <- url
		}
	})
}

func (scraper *Scraper) Stop() {
	Logger().Warningf("Stopping %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_CLOSED,
		extensionParameters{scraper: scraper})

	scraper.chDone <- struct{}{}
	scraper.engine.wg.Done()
}

func (scraper *Scraper) Start() {
	scraper.engine.wg.Add(1)
	Logger().Infof("Starting: %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_OPENED,
		extensionParameters{scraper: scraper})

	scraper.chRequestUrl <- scraper.BaseUrl
	duration := time.Duration(scraper.requestLimit)

	if scraper.requestLimit == 0 {
		duration = defaultRequestLimit()
	}

	limiter := time.Tick(time.Millisecond * duration)

	for {
		select {
		case url := <-scraper.chRequestUrl:
			<-limiter
			go scraper.Fetch(url)
		case <-scraper.chDone:
			Logger().Warningf("Stopped %s", scraper)
			scraper.engine.IncrFinishedCounter()
			scraper.engine.chDone <- struct{}{}
			return
		}
	}
	return
}

func (scraper *Scraper) Notify(url string, resp *http.Response) {
	scraper.engine.Meta.IncrScraped(scraper)
	scraper.engine.chScraped <- NewScrapedItem(url, scraper, resp)
}

func (scraper *Scraper) Fetch(url string) (resp *http.Response, err error) {
	if ok := scraper.CheckIfFetched(url); ok {
		return
	}
	scraper.MarkAsFetched(url)

	Logger().Infof("Fetching: %s", url)
	tic := time.Now()

	req, _ := http.NewRequest("GET", url, nil)
	req = scraper.engine.PrepareRequest(req)

	resp, err = NewHTTPClient().Do(req)

	statusCode := 0
	if err == nil {
		statusCode = resp.StatusCode
	}

	Logger().Debugf("[%d]Request to %s took: %s", statusCode, url, time.Since(tic))

	isSuccessful := (err == nil)

	scraper.engine.Meta.UpdateRequestStats(scraper, isSuccessful, req, resp)

	if err == nil {
		scraper.Notify(url, resp)
		scraper.RunExtractor(resp)
	} else {
		Logger().Warningf("Failed to crawl %s", url)
		Logger().Warning(err)
	}

	if scraper.CheckIfShouldStop() {
		scraper.Stop()
	}
	return
}

func (scraper *Scraper) SetHandler(handler ScrapingHandlerFunc) *Scraper {
	scraper.handler = handler
	return scraper
}

func (scraper *Scraper) String() (result string) {
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]
	result = fmt.Sprintf("<Scraper: %s>. Crawled: %d, successful: %d, failed: %d. Items scraped: %d, saved: %d",
		scraper.Domain, stats.crawled, stats.successful, stats.failed, stats.scraped, stats.saved)
	return
}

func NewScraper(name string, sourceUrl string, requestLimit int, extractor Extractable) (s *Scraper) {
	parsed, err := URL.Parse(sourceUrl)
	if err != nil {
		Logger().Infof("Inappropriate URL: %s", sourceUrl)
		return
	}

	if extractor == nil {
		Logger().Warning("Switching to default extractor")
		extractor = defaultExtractor()
	}

	s = &Scraper{
		Name:         name,
		Domain:       parsed.Host,
		BaseUrl:      sourceUrl,
		fetchedUrls:  make(map[string]bool),
		crawledMutex: &sync.Mutex{},
		fetchMutex:   &sync.Mutex{},
		extractor:    extractor,
		chDone:       make(chan struct{}),
		chRequestUrl: make(chan string, 5),
		requestLimit: requestLimit,
	}
	return
}

func NewScrapedItem(url string, scraper *Scraper, resp *http.Response) ScrapedItem {
	bodyBytes, _ := ioutil.ReadAll(resp.Body)
	resp.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))

	return ScrapedItem{
		BodyBytes: bodyBytes,
		FinalUrl:  resp.Request.URL.String(),
		Url:       url,
		scraper:   scraper,
	}
}

func NewHTTPClient() (client *http.Client) {
	client = &http.Client{
		Timeout: TIMEOUT_REQUEST,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: TIMEOUT_DIALER,
			}).Dial,
			TLSHandshakeTimeout: TIMEOUT_TLS,
		},
	}
	return
}

func defaultExtractor() Extractable {
	return &LinkExtractor{}
}

func defaultRequestLimit() time.Duration {
	return time.Duration(1)
}

func SaveItem(item SaveableItem, writer recordWriter) {
	if writer == nil {
		return
	}

	if !item.Validate() {
		Logger().Warning("Item is not valid. Skipping...")
		return
	}

	defer writer.Flush()
	writer.Write(item.RecordData())
}

func NewSpiderConfig(file string) (config *ScraperConfig) {
	config = &ScraperConfig{}
	ProcessFile(config, file)
	return
}
<commit_msg>Simplified extractor<commit_after>package gotana

import (
	"bytes"
	"fmt"
	"github.com/PuerkitoBio/goquery"
	"golang.org/x/net/html"
	"io"
	"io/ioutil"
	"net"
	"net/http"
	URL "net/url"
	"strings"
	"sync"
	"time"
)

const (
	EVENT_SCRAPER_OPENED     = "SCRAPER_OPENED"
	EVENT_SCRAPER_CLOSED     = "SCRAPER_CLOSED"
	EVENT_SAVEABLE_EXTRACTED = "SAVEABLE_EXTRACTED"
	TIMEOUT_DIALER           = time.Duration(time.Second * 30)
	TIMEOUT_REQUEST          = time.Duration(time.Second * 30)
	TIMEOUT_TLS              = time.Duration(time.Second * 10)
)

type SaveableItem interface {
	Scraper() *Scraper
	Validate() bool
	RecordData() []string
}

type ScraperMixin struct {
	Proxy ScrapedItem
}

func (s *ScraperMixin) SetProxy(proxy ScrapedItem) *ScraperMixin {
	s.Proxy = proxy
	return s
}

func (item ScraperMixin) Scraper() *Scraper {
	return item.Proxy.scraper
}

type recordWriter interface {
	Write(record []string) error
	Flush()
}

type ScrapingHandlerFunc func(ScrapedItem, chan<- SaveableItem)

func GetHref(t html.Token) (ok bool, href string) {
	for _, a := range t.Attr {
		if a.Key == "href" {
			href = a.Val
			ok = true
		}
	}

	return
}

type extensionParameters struct {
	scraper *Scraper
	item    SaveableItem
}

type Extractable interface {
	Extract(io.ReadCloser, func(string))
}

type LinkExtractor struct {
	Extractable
}

func (extractor *LinkExtractor) Extract(r io.ReadCloser, callback func(string)) {
	page := html.NewTokenizer(r)
	defer r.Close()

	for {
		tokenType := page.Next()
		if tokenType == html.ErrorToken {
			return
		}
		token := page.Token()
		if tokenType == html.StartTagToken && token.DataAtom.String() == "a" {
			ok, url := GetHref(token)
			if !ok {
				continue
			}
			callback(url)
		}
	}
}

type ScraperConfig struct {
	Project     string `required:"true"`
	TcpAddress  string
	OutFileName string
	Scrapers    []struct {
		RequestLimit int `required:"true"`
		Extractor    string
		Name         string `required:"true"`
		Url          string `required:"true"`
	}
}

type ScrapedItem struct {
	Url       string
	FinalUrl  string
	scraper   *Scraper
	BodyBytes []byte
}

func (proxy ScrapedItem) String() (result string) {
	result = fmt.Sprintf("Result of scraping: %s", proxy.Url)
	return
}

func (proxy ScrapedItem) CheckIfRedirected() bool {
	return proxy.Url != proxy.FinalUrl
}

func (proxy ScrapedItem) FinalResponseBody() (io.ReadCloser, error) {
	if proxy.CheckIfRedirected() {
		client := NewHTTPClient()
		response, err := client.Get(proxy.FinalUrl)
		if err != nil {
			return nil, err
		}
		bodyBytes, _ := ioutil.ReadAll(response.Body)
		proxy.BodyBytes = bodyBytes
	}
	return ioutil.NopCloser(bytes.NewBuffer(proxy.BodyBytes)), nil
}

func (proxy ScrapedItem) HTMLDocument() (document *goquery.Document, err error) {
	responseBody, err := proxy.FinalResponseBody()
	if err == nil {
		document, err = goquery.NewDocumentFromReader(responseBody)
	}

	return
}

type Scraper struct {
	crawled      int
	successful   int
	failed       int
	handler      ScrapingHandlerFunc
	fetchMutex   *sync.Mutex
	crawledMutex *sync.Mutex
	Name         string
	Domain       string
	BaseUrl      string
	CurrentUrl   string
	fetchedUrls  map[string]bool
	engine       *Engine
	extractor    Extractable
	chDone       chan struct{}
	chRequestUrl chan string
	requestLimit int
}

func (scraper *Scraper) MarkAsFetched(url string) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	scraper.CurrentUrl = url
	scraper.fetchedUrls[url] = true
}

func (scraper *Scraper) CheckIfShouldStop() (ok bool) {
	scraper.crawledMutex.Lock()
	defer scraper.crawledMutex.Unlock()
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]

	if stats.crawled == scraper.engine.limitCrawl {
		Logger().Warningf("Crawl limit exceeded: %s", scraper)
		ok = true
	} else if stats.failed == scraper.engine.limitFail {
		Logger().Warningf("Fail limit exceeeded: %s", scraper)
		ok = true
	} else if stats.failed == 1 && scraper.crawled == 1 {
		Logger().Warningf("Base URL is corrupted: %s", scraper)
		ok = true
	}
	return
}

func (scraper *Scraper) CheckIfFetched(url string) (ok bool) {
	scraper.fetchMutex.Lock()
	defer scraper.fetchMutex.Unlock()

	_, ok = scraper.fetchedUrls[url]
	return
}

func (scraper *Scraper) CheckUrl(sourceUrl string) (ok bool, url string) {
	if strings.Contains(sourceUrl, scraper.Domain) && strings.Index(sourceUrl, "http") == 0 {
		url = sourceUrl
		ok = true
	} else if strings.Index(sourceUrl, "/") == 0 {
		url = scraper.BaseUrl + sourceUrl
		ok = true
	}
	return
}

func (scraper *Scraper) RunExtractor(resp *http.Response) {
	defer SilentRecover("EXTRACTOR")

	scraper.extractor.Extract(resp.Body, func(url string) {
		ok, url := scraper.CheckUrl(url)

		if ok {
			scraper.chRequestUrl <- url
		}
	})
}

func (scraper *Scraper) Stop() {
	Logger().Warningf("Stopping %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_CLOSED,
		extensionParameters{scraper: scraper})

	scraper.chDone <- struct{}{}
	scraper.engine.wg.Done()
}

func (scraper *Scraper) Start() {
	scraper.engine.wg.Add(1)
	Logger().Infof("Starting: %s", scraper)
	scraper.engine.notifyExtensions(EVENT_SCRAPER_OPENED,
		extensionParameters{scraper: scraper})

	scraper.chRequestUrl <- scraper.BaseUrl
	duration := time.Duration(scraper.requestLimit)

	if scraper.requestLimit == 0 {
		duration = defaultRequestLimit()
	}

	limiter := time.Tick(time.Millisecond * duration)

	for {
		select {
		case url := <-scraper.chRequestUrl:
			<-limiter
			go scraper.Fetch(url)
		case <-scraper.chDone:
			Logger().Warningf("Stopped %s", scraper)
			scraper.engine.IncrFinishedCounter()
			scraper.engine.chDone <- struct{}{}
			return
		}
	}
	return
}

func (scraper *Scraper) Notify(url string, resp *http.Response) {
	scraper.engine.Meta.IncrScraped(scraper)
	scraper.engine.chScraped <- NewScrapedItem(url, scraper, resp)
}

func (scraper *Scraper) Fetch(url string) (resp *http.Response, err error) {
	if ok := scraper.CheckIfFetched(url); ok {
		return
	}
	scraper.MarkAsFetched(url)

	Logger().Infof("Fetching: %s", url)
	tic := time.Now()

	req, _ := http.NewRequest("GET", url, nil)
	req = scraper.engine.PrepareRequest(req)

	resp, err = NewHTTPClient().Do(req)

	statusCode := 0
	if err == nil {
		statusCode = resp.StatusCode
	}

	Logger().Debugf("[%d]Request to %s took: %s", statusCode, url, time.Since(tic))

	isSuccessful := (err == nil)

	scraper.engine.Meta.UpdateRequestStats(scraper, isSuccessful, req, resp)

	if err == nil {
		scraper.Notify(url, resp)
		scraper.RunExtractor(resp)
	} else {
		Logger().Warningf("Failed to crawl %s", url)
		Logger().Warning(err)
	}

	if scraper.CheckIfShouldStop() {
		scraper.Stop()
	}
	return
}

func (scraper *Scraper) SetHandler(handler ScrapingHandlerFunc) *Scraper {
	scraper.handler = handler
	return scraper
}

func (scraper *Scraper) String() (result string) {
	stats := scraper.engine.Meta.ScraperStats[scraper.Name]
	result = fmt.Sprintf("<Scraper: %s>. Crawled: %d, successful: %d, failed: %d. Items scraped: %d, saved: %d",
		scraper.Domain, stats.crawled, stats.successful, stats.failed, stats.scraped, stats.saved)
	return
}

func NewScraper(name string, sourceUrl string, requestLimit int, extractor Extractable) (s *Scraper) {
	parsed, err := URL.Parse(sourceUrl)
	if err != nil {
		Logger().Infof("Inappropriate URL: %s", sourceUrl)
		return
	}

	if extractor == nil {
		Logger().Warning("Switching to default extractor")
		extractor = defaultExtractor()
	}

	s = &Scraper{
		Name:         name,
		Domain:       parsed.Host,
		BaseUrl:      sourceUrl,
		fetchedUrls:  make(map[string]bool),
		crawledMutex: &sync.Mutex{},
		fetchMutex:   &sync.Mutex{},
		extractor:    extractor,
		chDone:       make(chan struct{}),
		chRequestUrl: make(chan string, 5),
		requestLimit: requestLimit,
	}
	return
}

func NewScrapedItem(url string, scraper *Scraper, resp *http.Response) ScrapedItem {
	bodyBytes, _ := ioutil.ReadAll(resp.Body)
	resp.Body = ioutil.NopCloser(bytes.NewBuffer(bodyBytes))

	return ScrapedItem{
		BodyBytes: bodyBytes,
		FinalUrl:  resp.Request.URL.String(),
		Url:       url,
		scraper:   scraper,
	}
}

func NewHTTPClient() (client *http.Client) {
	client = &http.Client{
		Timeout: TIMEOUT_REQUEST,
		Transport: &http.Transport{
			Dial: (&net.Dialer{
				Timeout: TIMEOUT_DIALER,
			}).Dial,
			TLSHandshakeTimeout: TIMEOUT_TLS,
		},
	}
	return
}

func defaultExtractor() Extractable {
	return &LinkExtractor{}
}

func defaultRequestLimit() time.Duration {
	return time.Duration(1)
}

func SaveItem(item SaveableItem, writer recordWriter) {
	if writer == nil {
		return
	}

	if !item.Validate() {
		Logger().Warning("Item is not valid. Skipping...")
		return
	}

	defer writer.Flush()
	writer.Write(item.RecordData())
}

func NewSpiderConfig(file string) (config *ScraperConfig) {
	config = &ScraperConfig{}
	ProcessFile(config, file)
	return
}
<|endoftext|>


<commit_before>/*
Package detour provides a net.Conn interface which detects blockage
of a site automatically and access it through alternative connection.

Basically, if a site is not whitelisted, following steps will be taken:
1. Dial proxied connection (detour) a small delay after dialed directly
2. Return to caller when any connection is established
3. Read/write through all open connections in parallel
4. Check for blockage on direct connection and closes it if it happens
5. If possible, replay operations on detour connection. [1]
6. After sucessfully read from a connection, stick with it and close others.
7. Add those sites failed on direct connection but succeeded on detour ones
   to proxied list, so above steps can be skipped next time. The list can be
   exported and persisted if required.

Blockage can happen at several stages of a connection, what detour can detect are:
1. Connection attempt is blocked (IP blocking / DNS hijack).
   Symptoms can be connection time out / TCP RST / connection refused.
2. Connection made but real data get blocked (DPI).
3. Successfully exchanged a few packets, while follow up packets are blocked. [2]
4. Connection made but get fake response or HTTP redirect to a fixed URL.

[1] Detour will not replay nonidempotent plain HTTP requests, but will add it to
    proxied list to be detoured next time.
[2] Detour can only handle exact 1 successful read followed by failed read,
    which covers most cases in reality.
*/
package detour

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/getlantern/golog"
)

// If no any connection made after this period, stop dialing and fail
var TimeoutToConnect = 30 * time.Second

// To avoid unnecessarily proxy not-blocked url, detour will dial detour connection
// after this small delay. Set to zero to dial in parallel to not introducing any delay.
var DelayBeforeDetour = 0 * time.Millisecond

// If DirectAddrCh is set, when a direct connection is closed without any error,
// the connection's remote address (in host:port format) will be send to it
var DirectAddrCh chan string = make(chan string)

var (
	log = golog.LoggerFor("detour")
)

// Conn implements an net.Conn interface by utilizing underlie direct and
// detour connections.
type Conn struct {
	// Keeps track of the total bytes read from this connection, atomic
	// Due to https://golang.org/pkg/sync/atomic/#pkg-note-BUG it requires
	// manual alignment. For this, it is best to keep it as the first field
	readBytes uint64

	// The underlie connections, uses buffered channel as ring queue to avoid
	// locking. We have at most 2 connetions so a length of 2 is enough.
	conns chan conn

	// The chan to notify dialer to dial detour immediately
	chDialDetourNow chan bool
	// The channel to notify read/write that a detour connection is available
	chDetourConn chan conn

	// The chan to receive result of any read operation
	chRead chan ioResult
	// The chan to receive result of any write operation
	chWrite chan ioResult

	addr string

	muWriteBuffer sync.RWMutex
	// Keeps written bytes through direct connection to replay it if required.
	writeBuffer *bytes.Buffer
	// Is it a plain HTTP request or not, atomic
	nonidempotentHTTPRequest uint32
}

// The data structure to pass result of io operation back from underlie connection
type ioResult struct {
	// Number of bytes read/wrote
	n int
	// IO error, if any
	err error
	// The underlie connection itself
	conn conn
}

type connType int

const (
	connTypeDirect connType = iota
	connTypeDetour connType = iota
)

type conn interface {
	ConnType() connType
	FirstRead(b []byte, ch chan ioResult)
	FollowupRead(b []byte, ch chan ioResult)
	Write(b []byte, ch chan ioResult)
	Close() error
	Closed() bool
}

func typeOf(c conn) string {
	var connTypeDesc = []string{"direct", "detour"}
	return connTypeDesc[c.ConnType()]
}

type dialFunc func(network, addr string) (net.Conn, error)

// Dialer returns a function with same signature of net.Dialer.Dial().
func Dialer(detourDialer dialFunc) func(network, addr string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		dc := &Conn{
			addr:            addr,
			writeBuffer:     new(bytes.Buffer),
			conns:           make(chan conn, 2),
			chDetourConn:    make(chan conn),
			chRead:          make(chan ioResult),
			chWrite:         make(chan ioResult),
			chDialDetourNow: make(chan bool),
		}
		// use buffered channel as we may send twice to it but only receive once
		chAnyConn := make(chan bool, 1)
		ch := make(chan conn)

		// dialing sequence
		if whitelisted(addr) {
			dialDetour(network, addr, detourDialer, ch)
		} else {
			go func() {
				dialDirect(network, addr, ch)
				dt := time.NewTimer(DelayBeforeDetour)
				select {
				case <-dt.C:
				case <-dc.chDialDetourNow:
				}
				if dc.anyDataReceived() {
					ch <- nil
					return
				}
				dialDetour(network, addr, detourDialer, ch)
			}()
		}

		// handle dialing result
		go func() {
			t := time.NewTimer(TimeoutToConnect)
			defer t.Stop()
			// At most 2 connections will be made
			for i := 0; i < 2; i++ {
				log.Tracef("Waiting for connection to %s, round %d", dc.addr, i)
				select {
				case c := <-ch:
					if c == nil {
						log.Tracef("No new connection to %s remaining, return", dc.addr)
						return
					}
					// first connection made, pass it back to caller
					if i == 0 {
						dc.conns <- c
						chAnyConn <- true
					} else {
						if c.ConnType() == connTypeDirect {
							// Could happen if direct route is much slower.
							log.Debugf("Direct connection to %s established too late, close it", dc.addr)
							if err := c.Close(); err != nil {
								log.Debugf("Error closing direct connection to %s: %s", dc.addr, err)
							}
							return
						}
						log.Tracef("Feed detour connection to %s to read/write op", dc.addr)
						dc.chDetourConn <- c
						return
					}
				case <-t.C:
					// still no connection made
					chAnyConn <- false
					return
				}
			}
		}()
		// return to caller if any connection available
		if anyConn := <-chAnyConn; anyConn {
			return dc, nil
		}
		return nil, fmt.Errorf("Timeout dialing any connection to %s", addr)
	}
}

func (dc *Conn) anyDataReceived() bool {
	return atomic.LoadUint64(&dc.readBytes) > 0
}

func (dc *Conn) incReadBytes(n int) {
	atomic.AddUint64(&dc.readBytes, uint64(n))
}

// Read() implements the function from net.Conn
func (dc *Conn) Read(b []byte) (n int, err error) {
	if dc.anyDataReceived() {
		return dc.followupRead(b)
	}
	// At initial stage, we only have one connection,
	// but detour connection can be available at anytime.
	if !dc.withValidConn(func(c conn) { c.FirstRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	for count := 1; count > 0; count-- {
		select {
		case newConn := <-dc.chDetourConn:
			if atomic.LoadUint32(&dc.nonidempotentHTTPRequest) == 1 {
				log.Tracef("Not replay nonidempotent request to %s, only add to whitelist", dc.addr)
				AddToWl(dc.addr, false)
				if err := newConn.Close(); err != nil {
					log.Debugf("Error closing detour connection to %s: %s", dc.addr, err)
				}
				return
			}
			log.Tracef("Got detour connection to %s, replay previous op on it", dc.addr)
			dc.muWriteBuffer.RLock()
			sentBytes := dc.writeBuffer.Bytes()
			dc.muWriteBuffer.RUnlock()
			newConn.Write(sentBytes, dc.chWrite)
			newConn.FirstRead(b, dc.chRead)
			count++
			// add new connection to connections
			dc.conns <- newConn
		case result := <-dc.chRead:
			conn, n, err := result.conn, result.n, result.err
			if err != nil {
				log.Tracef("Read from %s connection to %s failed, closing: %s", typeOf(conn), dc.addr, err)
				if err := conn.Close(); err != nil {
					log.Debugf("Error closing %s connection to %s: %s", typeOf(conn), dc.addr, err)
				}
				// skip failed connection as we have more
				if count > 1 {
					continue
				}
				switch conn.ConnType() {
				case connTypeDirect:
					// if we haven't dial detour yet, do so now
					select {
					case dc.chDialDetourNow <- true:
						count++
					default:
					}
					continue
				case connTypeDetour:
					log.Tracef("Detour connection to %s failed, removing from whitelist", dc.addr)
					RemoveFromWl(dc.addr)
					// no more connections, return directly to avoid dead lock
					return n, err
				}
			}
			log.Tracef("Read %d bytes from %s connection to %s", n, typeOf(conn), dc.addr)
			dc.incReadBytes(n)
			return n, err
		}
	}
	return
}

// followUpRead is called by Read() if a connection's state already settled
func (dc *Conn) followupRead(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.FollowupRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chRead
	dc.incReadBytes(result.n)
	return result.n, result.err
}

// Write() implements the function from net.Conn
func (dc *Conn) Write(b []byte) (n int, err error) {
	if dc.anyDataReceived() {
		return dc.followupWrite(b)
	}
	if isNonidempotentHTTPRequest(b) {
		atomic.StoreUint32(&dc.nonidempotentHTTPRequest, 1)
	} else {
		dc.muWriteBuffer.Lock()
		_, _ = dc.writeBuffer.Write(b)
		dc.muWriteBuffer.Unlock()
	}
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}

	result := <-dc.chWrite
	if n, err = result.n, result.err; err != nil {
		log.Tracef("Error writing %s connection to %s: %s", typeOf(result.conn), dc.addr, err)
		if err := result.conn.Close(); err != nil {
			log.Debugf("Error closing %s connection to %s: %s", typeOf(result.conn), dc.addr, err)
		}
		return
	}
	log.Tracef("Wrote %d bytes to %s connection to %s", n, typeOf(result.conn), dc.addr)
	return
}

// followupWrite is called by Write() if a connection's state already settled
func (dc *Conn) followupWrite(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chWrite
	return result.n, result.err
}

// Close implements the function from net.Conn
func (dc *Conn) Close() error {
	log.Tracef("Closing connection to %s", dc.addr)
	for len(dc.conns) > 0 {
		conn := <-dc.conns
		if err := conn.Close(); err != nil {
			log.Debugf("Error closing %s connection to %s: %s", typeOf(conn), dc.addr, err)
		}
	}
	return nil
}

// LocalAddr implements the function from net.Conn
func (dc *Conn) LocalAddr() net.Addr {
	log.Trace("LocalAddr not implemented")
	return nil
}

// RemoteAddr implements the function from net.Conn
func (dc *Conn) RemoteAddr() net.Addr {
	log.Trace("RemoteAddr not implemented")
	return nil
}

// SetDeadline implements the function from net.Conn
func (dc *Conn) SetDeadline(t time.Time) error {
	return fmt.Errorf("SetDeadline not implemented")
}

// SetReadDeadline implements the function from net.Conn
func (dc *Conn) SetReadDeadline(t time.Time) error {
	return fmt.Errorf("SetReadDeadline not implemented")
}

// SetWriteDeadline implements the function from net.Conn
func (dc *Conn) SetWriteDeadline(t time.Time) error {
	return fmt.Errorf("SetWriteDeadline not implemented")
}

func (dc *Conn) withValidConn(f func(conn)) bool {
	for i := 0; i < len(dc.conns); i++ {
		select {
		case c := <-dc.conns:
			if c.Closed() {
				log.Tracef("Drain closed %s connection to %s", typeOf(c), dc.addr)
				continue
			}
			f(c)
			dc.conns <- c
			return true
		default:
			break
		}
	}
	return false
}

var nonidempotentMethods = [][]byte{
	[]byte("PUT "),
	[]byte("POST "),
	[]byte("PATCH "),
}

// Ref section 9.1.2 of https://www.ietf.org/rfc/rfc2616.txt.
// We consider the https handshake phase to be idemponent.
func isNonidempotentHTTPRequest(b []byte) bool {
	if len(b) > 4 {
		for _, m := range nonidempotentMethods {
			if bytes.HasPrefix(b, m) {
				return true
			}
		}
	}
	return false
}
<commit_msg>detour supports LocalAddr() and RemoteAddr()<commit_after>/*
Package detour provides a net.Conn interface which detects blockage
of a site automatically and access it through alternative connection.

Basically, if a site is not whitelisted, following steps will be taken:
1. Dial proxied connection (detour) a small delay after dialed directly
2. Return to caller when any connection is established
3. Read/write through all open connections in parallel
4. Check for blockage on direct connection and closes it if it happens
5. If possible, replay operations on detour connection. [1]
6. After sucessfully read from a connection, stick with it and close others.
7. Add those sites failed on direct connection but succeeded on detour ones
   to proxied list, so above steps can be skipped next time. The list can be
   exported and persisted if required.

Blockage can happen at several stages of a connection, what detour can detect are:
1. Connection attempt is blocked (IP blocking / DNS hijack).
   Symptoms can be connection time out / TCP RST / connection refused.
2. Connection made but real data get blocked (DPI).
3. Successfully exchanged a few packets, while follow up packets are blocked. [2]
4. Connection made but get fake response or HTTP redirect to a fixed URL.

[1] Detour will not replay nonidempotent plain HTTP requests, but will add it to
    proxied list to be detoured next time.
[2] Detour can only handle exact 1 successful read followed by failed read,
    which covers most cases in reality.
*/
package detour

import (
	"bytes"
	"fmt"
	"net"
	"sync"
	"sync/atomic"
	"time"

	"github.com/getlantern/golog"
)

// If no any connection made after this period, stop dialing and fail
var TimeoutToConnect = 30 * time.Second

// To avoid unnecessarily proxy not-blocked url, detour will dial detour connection
// after this small delay. Set to zero to dial in parallel to not introducing any delay.
var DelayBeforeDetour = 0 * time.Millisecond

// If DirectAddrCh is set, when a direct connection is closed without any error,
// the connection's remote address (in host:port format) will be send to it
var DirectAddrCh = make(chan string)

var (
	log = golog.LoggerFor("detour")
)

// Conn implements an net.Conn interface by utilizing underlie direct and
// detour connections.
type Conn struct {
	// Keeps track of the total bytes read from this connection, atomic
	// Due to https://golang.org/pkg/sync/atomic/#pkg-note-BUG it requires
	// manual alignment. For this, it is best to keep it as the first field
	readBytes uint64

	// The underlie connections, uses buffered channel as ring queue to avoid
	// locking. We have at most 2 connetions so a length of 2 is enough.
	conns chan conn

	// The chan to notify dialer to dial detour immediately
	chDialDetourNow chan bool
	// The channel to notify read/write that a detour connection is available
	chDetourConn chan conn

	// The chan to receive result of any read operation
	chRead chan ioResult
	// The chan to receive result of any write operation
	chWrite chan ioResult

	addr string

	muWriteBuffer sync.RWMutex
	// Keeps written bytes through direct connection to replay it if required.
	writeBuffer *bytes.Buffer
	// Is it a plain HTTP request or not, atomic
	nonidempotentHTTPRequest uint32
}

// The data structure to pass result of io operation back from underlie connection
type ioResult struct {
	// Number of bytes read/wrote
	n int
	// IO error, if any
	err error
	// The underlie connection itself
	conn conn
}

type connType int

const (
	connTypeDirect connType = iota
	connTypeDetour connType = iota
)

type conn interface {
	ConnType() connType
	FirstRead(b []byte, ch chan ioResult)
	FollowupRead(b []byte, ch chan ioResult)
	Write(b []byte, ch chan ioResult)
	Close() error
	Closed() bool
	LocalAddr() net.Addr
	RemoteAddr() net.Addr
}

func typeOf(c conn) string {
	var connTypeDesc = []string{"direct", "detour"}
	return connTypeDesc[c.ConnType()]
}

type dialFunc func(network, addr string) (net.Conn, error)

// Dialer returns a function with same signature of net.Dialer.Dial().
func Dialer(detourDialer dialFunc) func(network, addr string) (net.Conn, error) {
	return func(network, addr string) (net.Conn, error) {
		dc := &Conn{
			addr:            addr,
			writeBuffer:     new(bytes.Buffer),
			conns:           make(chan conn, 2),
			chDetourConn:    make(chan conn),
			chRead:          make(chan ioResult),
			chWrite:         make(chan ioResult),
			chDialDetourNow: make(chan bool),
		}
		// use buffered channel as we may send twice to it but only receive once
		chAnyConn := make(chan bool, 1)
		ch := make(chan conn)

		// dialing sequence
		if whitelisted(addr) {
			dialDetour(network, addr, detourDialer, ch)
		} else {
			go func() {
				dialDirect(network, addr, ch)
				dt := time.NewTimer(DelayBeforeDetour)
				select {
				case <-dt.C:
				case <-dc.chDialDetourNow:
				}
				if dc.anyDataReceived() {
					ch <- nil
					return
				}
				dialDetour(network, addr, detourDialer, ch)
			}()
		}

		// handle dialing result
		go func() {
			t := time.NewTimer(TimeoutToConnect)
			defer t.Stop()
			// At most 2 connections will be made
			for i := 0; i < 2; i++ {
				log.Tracef("Waiting for connection to %s, round %d", dc.addr, i)
				select {
				case c := <-ch:
					if c == nil {
						log.Tracef("No new connection to %s remaining, return", dc.addr)
						return
					}
					// first connection made, pass it back to caller
					if i == 0 {
						dc.conns <- c
						chAnyConn <- true
					} else {
						if c.ConnType() == connTypeDirect {
							// Could happen if direct route is much slower.
							log.Debugf("Direct connection to %s established too late, close it", dc.addr)
							if err := c.Close(); err != nil {
								log.Debugf("Error closing direct connection to %s: %s", dc.addr, err)
							}
							return
						}
						log.Tracef("Feed detour connection to %s to read/write op", dc.addr)
						dc.chDetourConn <- c
						return
					}
				case <-t.C:
					// still no connection made
					chAnyConn <- false
					return
				}
			}
		}()
		// return to caller if any connection available
		if anyConn := <-chAnyConn; anyConn {
			return dc, nil
		}
		return nil, fmt.Errorf("Timeout dialing any connection to %s", addr)
	}
}

func (dc *Conn) anyDataReceived() bool {
	return atomic.LoadUint64(&dc.readBytes) > 0
}

func (dc *Conn) incReadBytes(n int) {
	atomic.AddUint64(&dc.readBytes, uint64(n))
}

// Read() implements the function from net.Conn
func (dc *Conn) Read(b []byte) (n int, err error) {
	if dc.anyDataReceived() {
		return dc.followupRead(b)
	}
	// At initial stage, we only have one connection,
	// but detour connection can be available at anytime.
	if !dc.withValidConn(func(c conn) { c.FirstRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	for count := 1; count > 0; count-- {
		select {
		case newConn := <-dc.chDetourConn:
			if atomic.LoadUint32(&dc.nonidempotentHTTPRequest) == 1 {
				log.Tracef("Not replay nonidempotent request to %s, only add to whitelist", dc.addr)
				AddToWl(dc.addr, false)
				if err := newConn.Close(); err != nil {
					log.Debugf("Error closing detour connection to %s: %s", dc.addr, err)
				}
				return
			}
			log.Tracef("Got detour connection to %s, replay previous op on it", dc.addr)
			dc.muWriteBuffer.RLock()
			sentBytes := dc.writeBuffer.Bytes()
			dc.muWriteBuffer.RUnlock()
			newConn.Write(sentBytes, dc.chWrite)
			newConn.FirstRead(b, dc.chRead)
			count++
			// add new connection to connections
			dc.conns <- newConn
		case result := <-dc.chRead:
			conn, n, err := result.conn, result.n, result.err
			if err != nil {
				log.Tracef("Read from %s connection to %s failed, closing: %s", typeOf(conn), dc.addr, err)
				if err := conn.Close(); err != nil {
					log.Debugf("Error closing %s connection to %s: %s", typeOf(conn), dc.addr, err)
				}
				// skip failed connection as we have more
				if count > 1 {
					continue
				}
				switch conn.ConnType() {
				case connTypeDirect:
					// if we haven't dial detour yet, do so now
					select {
					case dc.chDialDetourNow <- true:
						count++
					default:
					}
					continue
				case connTypeDetour:
					log.Tracef("Detour connection to %s failed, removing from whitelist", dc.addr)
					RemoveFromWl(dc.addr)
					// no more connections, return directly to avoid dead lock
					return n, err
				}
			}
			log.Tracef("Read %d bytes from %s connection to %s", n, typeOf(conn), dc.addr)
			dc.incReadBytes(n)
			return n, err
		}
	}
	return
}

// followUpRead is called by Read() if a connection's state already settled
func (dc *Conn) followupRead(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.FollowupRead(b, dc.chRead) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chRead
	dc.incReadBytes(result.n)
	return result.n, result.err
}

// Write() implements the function from net.Conn
func (dc *Conn) Write(b []byte) (n int, err error) {
	if dc.anyDataReceived() {
		return dc.followupWrite(b)
	}
	if isNonidempotentHTTPRequest(b) {
		atomic.StoreUint32(&dc.nonidempotentHTTPRequest, 1)
	} else {
		dc.muWriteBuffer.Lock()
		_, _ = dc.writeBuffer.Write(b)
		dc.muWriteBuffer.Unlock()
	}
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}

	result := <-dc.chWrite
	if n, err = result.n, result.err; err != nil {
		log.Tracef("Error writing %s connection to %s: %s", typeOf(result.conn), dc.addr, err)
		if err := result.conn.Close(); err != nil {
			log.Debugf("Error closing %s connection to %s: %s", typeOf(result.conn), dc.addr, err)
		}
		return
	}
	log.Tracef("Wrote %d bytes to %s connection to %s", n, typeOf(result.conn), dc.addr)
	return
}

// followupWrite is called by Write() if a connection's state already settled
func (dc *Conn) followupWrite(b []byte) (n int, err error) {
	if !dc.withValidConn(func(c conn) { c.Write(b, dc.chWrite) }) {
		return 0, fmt.Errorf("no connection available to %s", dc.addr)
	}
	result := <-dc.chWrite
	return result.n, result.err
}

// Close implements the function from net.Conn
func (dc *Conn) Close() error {
	log.Tracef("Closing connection to %s", dc.addr)
	for len(dc.conns) > 0 {
		conn := <-dc.conns
		if err := conn.Close(); err != nil {
			log.Debugf("Error closing %s connection to %s: %s", typeOf(conn), dc.addr, err)
		}
	}
	return nil
}

// LocalAddr implements the function from net.Conn
func (dc *Conn) LocalAddr() (addr net.Addr) {
	if !dc.withValidConn(func(c conn) { addr = c.LocalAddr() }) {
		panic("no valid connection to call LocalAddr()")
	}
	return
}

// RemoteAddr implements the function from net.Conn
func (dc *Conn) RemoteAddr() (addr net.Addr) {
	if !dc.withValidConn(func(c conn) { addr = c.RemoteAddr() }) {
		panic("no valid connection to call RemoteAddr()")
	}
	return
}

// SetDeadline implements the function from net.Conn
func (dc *Conn) SetDeadline(t time.Time) error {
	return fmt.Errorf("SetDeadline not implemented")
}

// SetReadDeadline implements the function from net.Conn
func (dc *Conn) SetReadDeadline(t time.Time) error {
	return fmt.Errorf("SetReadDeadline not implemented")
}

// SetWriteDeadline implements the function from net.Conn
func (dc *Conn) SetWriteDeadline(t time.Time) error {
	return fmt.Errorf("SetWriteDeadline not implemented")
}

func (dc *Conn) withValidConn(f func(conn)) bool {
	for i := 0; i < len(dc.conns); i++ {
		select {
		case c := <-dc.conns:
			if c.Closed() {
				log.Tracef("Drain closed %s connection to %s", typeOf(c), dc.addr)
				continue
			}
			f(c)
			dc.conns <- c
			return true
		default:
			break
		}
	}
	return false
}

var nonidempotentMethods = [][]byte{
	[]byte("PUT "),
	[]byte("POST "),
	[]byte("PATCH "),
}

// Ref section 9.1.2 of https://www.ietf.org/rfc/rfc2616.txt.
// We consider the https handshake phase to be idemponent.
func isNonidempotentHTTPRequest(b []byte) bool {
	if len(b) > 4 {
		for _, m := range nonidempotentMethods {
			if bytes.HasPrefix(b, m) {
				return true
			}
		}
	}
	return false
}
<|endoftext|>


<commit_before>package detour

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/getlantern/golog"
)

var (
	log = golog.LoggerFor("detour")
	// if dial or read exceeded this timeout, we consider switch to detour
	timeoutToDetour = 1 * time.Second

	muWhitelist sync.RWMutex
	whitelist   = make(map[string]bool)
)

type dialFunc func(network, addr string) (net.Conn, error)

type detourConn struct {
	muConn sync.RWMutex
	// the actual connection, will change so protect it
	// can't user atomic.Value as the concrete type may vary
	conn net.Conn

	// don't access directly, use inState() and setState() instead
	state uint32

	// the function to dial detour if the site to connect seems blocked
	dialDetour dialFunc

	muBuf sync.Mutex
	// keep track of bytes sent through normal connection
	// so we can resend them when detour
	buf bytes.Buffer

	network, addr string
	readDeadline  time.Time
	writeDeadline time.Time
}

const (
	stateInitial = iota
	stateDirect
	stateDetour
	stateWhitelistCandidate
	stateWhitelist
)

var statesDesc = []string{
	"INITIALLY",
	"DIRECTLY",
	"DETOURED",
	"WHITELIST CANDIDATE",
	"WHITELISTED",
}

// SetTimeout sets the timeout so if dial or read exceeds this timeout, we consider switch to detour
// The value depends on OS and browser and defaults to 1s
// For Windows XP, find TcpMaxConnectRetransmissions in http://support2.microsoft.com/default.aspx?scid=kb;en-us;314053
func SetTimeout(t time.Duration) {
	timeoutToDetour = t
}

func Dialer(dialer dialFunc) dialFunc {
	return func(network, addr string) (conn net.Conn, err error) {
		dc := &detourConn{dialDetour: dialer, network: network, addr: addr}
		if !whitelisted(addr) {
			dc.setState(stateInitial)
			dc.conn, err = net.DialTimeout(network, addr, timeoutToDetour)
			if err == nil {
				log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
				return dc, nil
			}
			log.Debugf("Dial %s to %s failed, try detour: %s", dc.stateDesc(), addr, err)
		}
		dc.setState(stateDetour)
		dc.conn, err = dc.dialDetour(network, addr)
		if err != nil {
			log.Errorf("Dial %s to %s failed", dc.stateDesc(), addr)
			return nil, err
		}
		log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
		return dc, err
	}
}

// Read() implements the function from net.Conn
func (dc *detourConn) Read(b []byte) (n int, err error) {
	conn := dc.getConn()
	if !dc.inState(stateInitial) {
		if n, err = conn.Read(b); err != nil && err != io.EOF {
			log.Tracef("Read from %s %s failed: %s", dc.addr, dc.stateDesc(), err)
			if dc.inState(stateDirect) && blocked(err) {
				// direct route is not reliable even the first read succeeded
				// try again through detour in next dial
				log.Tracef("Seems %s still blocked, add to whitelist so will try detour next time", dc.addr)
				addToWl(dc.addr, false)
			} else if wlTemporarily(dc.addr) {
				log.Tracef("Detoured route is still not reliable for %s, not whitelist it", dc.addr)
				removeFromWl(dc.addr)
			}
			/*if dc.inState(stateDetour) {
				log.Debugf("Add %s to white list temporarily", dc.addr)
				addToWhitelist(dc.addr, false)
			}*/
			return
		}
		log.Tracef("Read %d bytes from %s %s", n, dc.addr, dc.stateDesc())
		return n, err
	}
	/*if in, _ := inWhitelist(dc.addr); in {
		log.Tracef("%s in white list, detour", dc.addr)
		return dc.detour(b)
	}*/
	// state will always be settled after first read, safe to clear buffer at end of it
	defer dc.resetBuffer()
	now := time.Now()
	dl := now.Add(timeoutToDetour)
	if !dc.readDeadline.IsZero() && dc.readDeadline.Sub(now) < 2*timeoutToDetour {
		/*dc.setState(stateDirect)
		n, err = conn.Read(b)
		log.Tracef("No time left to detour, read %d bytes from %s directly, err=%s", n, dc.addr, err)
		return*/
		// if no enough room, reduce timeout to be half before read dead line
		dl = now.Add(dc.readDeadline.Sub(now) / 2)
	}
	conn.SetReadDeadline(dl)

	n, err = conn.Read(b)
	conn.SetReadDeadline(dc.readDeadline)
	if err != nil && err != io.EOF {
		ne := fmt.Errorf("Error while read from %s %s: %s", dc.addr, dc.stateDesc(), err)
		if blocked(err) {
			dc.detour(b)
		}
		return n, ne
	}
	log.Tracef("Read %d bytes from %s %s", n, dc.addr, dc.stateDesc())
	/*dc.setState(stateDirect)
	log.Tracef("Read %d bytes from %s directly, set state to %s", n, dc.addr, dc.stateDesc())*/
	return n, err
}

// Write() implements the function from net.Conn
func (dc *detourConn) Write(b []byte) (n int, err error) {
	if dc.inState(stateInitial) {
		if n, err = dc.writeToBuffer(b); err != nil {
			return n, fmt.Errorf("Unable to write to local buffer: %s", err)
		}
	}
	if n, err = dc.getConn().Write(b); err != nil {
		log.Debugf("Write %d bytes to %s %s failed: %s", len(b), dc.addr, dc.stateDesc(), err)
		/*if !dc.inState(stateDetour) {
			log.Debugf("Add %s to white list temporarily", dc.addr)
			addToWhitelist(dc.addr, false)
		}*/
	}
	return
}

// Close() implements the function from net.Conn
func (dc *detourConn) Close() error {
	log.Tracef("Closing %s connection to %s", dc.stateDesc(), dc.addr)
	if wlTemporarily(dc.addr) {
		log.Tracef("no error found till closing, add %s to permanent whitelist", dc.addr)
		addToWl(dc.addr, true)
	}
	return dc.getConn().Close()
}

func (dc *detourConn) LocalAddr() net.Addr {
	return dc.getConn().LocalAddr()
}

func (dc *detourConn) RemoteAddr() net.Addr {
	return dc.getConn().RemoteAddr()
}

func (dc *detourConn) SetDeadline(t time.Time) error {
	dc.SetReadDeadline(t)
	dc.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) SetReadDeadline(t time.Time) error {
	dc.readDeadline = t
	dc.conn.SetReadDeadline(t)
	return nil
}

func (dc *detourConn) SetWriteDeadline(t time.Time) error {
	dc.writeDeadline = t
	dc.conn.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) writeToBuffer(b []byte) (n int, err error) {
	dc.muBuf.Lock()
	n, err = dc.buf.Write(b)
	dc.muBuf.Unlock()
	return
}

func (dc *detourConn) resetBuffer() {
	dc.muBuf.Lock()
	dc.buf.Reset()
	dc.muBuf.Unlock()
}

func (dc *detourConn) detour(b []byte) (n int, err error) {
	if err = dc.setupDetour(); err != nil {
		log.Errorf("Error to setup detour: %s", err)
		return
	}
	if _, err = dc.resend(); err != nil {
		err = fmt.Errorf("Error resend buffer to %s: %s", dc.addr, err)
		log.Error(err)
		return
	}
	// should getConn() again as it has changed
	if n, err = dc.getConn().Read(b); err != nil {
		log.Debugf("Read from %s %s still failed: %s", dc.addr, dc.stateDesc(), err)
		return
	}
	dc.setState(stateDetour)
	addToWl(dc.addr, false)
	log.Tracef("Read %d bytes from %s through detour, set state to %s", n, dc.addr, dc.stateDesc())
	return
}

func (dc *detourConn) resend() (int, error) {
	dc.muBuf.Lock()
	b := dc.buf.Bytes()
	dc.muBuf.Unlock()
	if len(b) > 0 {
		n, err := dc.getConn().Write(b)
		log.Tracef("Resend %d buffered bytes to %s, %d sent", len(b), dc.addr, n)
		return n, err
	}
	return 0, nil
}

func (dc *detourConn) setupDetour() error {
	c, err := dc.dialDetour("tcp", dc.addr)
	if err != nil {
		return err
	}
	log.Tracef("Dialed a new detour connection to %s", dc.addr)
	dc.setConn(c)
	return nil
}

func (dc *detourConn) getConn() (c net.Conn) {
	dc.muConn.RLock()
	defer dc.muConn.RUnlock()
	return dc.conn
}

func (dc *detourConn) setConn(c net.Conn) {
	dc.muConn.Lock()
	oldConn := dc.conn
	dc.conn = c
	dc.muConn.Unlock()
	dc.conn.SetReadDeadline(dc.readDeadline)
	dc.conn.SetWriteDeadline(dc.writeDeadline)
	log.Tracef("Replaced connection to %s from direct to detour and closing old one", dc.addr)
	oldConn.Close()
}

func (dc *detourConn) stateDesc() string {
	return statesDesc[atomic.LoadUint32(&dc.state)]
}

func (dc *detourConn) inState(s uint32) bool {
	return atomic.LoadUint32(&dc.state) == s
}

func (dc *detourConn) setState(s uint32) {
	atomic.StoreUint32(&dc.state, s)
}

func blocked(err error) bool {
	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		return true
	}
	if oe, ok := err.(*net.OpError); ok && (oe.Err == syscall.EPIPE || oe.Err == syscall.ECONNRESET) {
		return true
	}
	return false
}

func whitelisted(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	_, in := whitelist[addr]
	return in
}

func wlTemporarily(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	return whitelist[addr]
}

func addToWl(addr string, permanent bool) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	whitelist[addr] = permanent
}

func removeFromWl(addr string) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	delete(whitelist, addr)
}
<commit_msg>âclean up again<commit_after>package detour

import (
	"bytes"
	"fmt"
	"io"
	"net"
	"sync"
	"sync/atomic"
	"syscall"
	"time"

	"github.com/getlantern/golog"
)

var (
	log = golog.LoggerFor("detour")
	// if dial or read exceeded this timeout, we consider switch to detour
	timeoutToDetour = 1 * time.Second

	muWhitelist sync.RWMutex
	whitelist   = make(map[string]bool)
)

type dialFunc func(network, addr string) (net.Conn, error)

type detourConn struct {
	muConn sync.RWMutex
	// the actual connection, will change so protect it
	// can't user atomic.Value as the concrete type may vary
	conn net.Conn

	// don't access directly, use inState() and setState() instead
	state uint32

	// the function to dial detour if the site to connect seems blocked
	dialDetour dialFunc

	muBuf sync.Mutex
	// keep track of bytes sent through normal connection
	// so we can resend them when detour
	buf bytes.Buffer

	network, addr string
	readDeadline  time.Time
	writeDeadline time.Time
}

const (
	stateInitial = iota
	stateDirect
	stateDetour
	stateWhitelistCandidate
	stateWhitelist
)

var statesDesc = []string{
	"INITIALLY",
	"DIRECTLY",
	"DETOURED",
	"WHITELIST CANDIDATE",
	"WHITELISTED",
}

// SetTimeout sets the timeout so if dial or read exceeds this timeout, we consider switch to detour
// The value depends on OS and browser and defaults to 1s
// For Windows XP, find TcpMaxConnectRetransmissions in http://support2.microsoft.com/default.aspx?scid=kb;en-us;314053
func SetTimeout(t time.Duration) {
	timeoutToDetour = t
}

func Dialer(dialer dialFunc) dialFunc {
	return func(network, addr string) (conn net.Conn, err error) {
		dc := &detourConn{dialDetour: dialer, network: network, addr: addr}
		if !whitelisted(addr) {
			dc.setState(stateInitial)
			dc.conn, err = net.DialTimeout(network, addr, timeoutToDetour)
			if err == nil {
				log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
				return dc, nil
			}
			log.Debugf("Dial %s to %s failed, try detour: %s", dc.stateDesc(), addr, err)
		}
		dc.setState(stateDetour)
		dc.conn, err = dc.dialDetour(network, addr)
		if err != nil {
			log.Errorf("Dial %s to %s failed", dc.stateDesc(), addr)
			return nil, err
		}
		log.Tracef("Dial %s to %s succeeded", dc.stateDesc(), addr)
		return dc, err
	}
}

// Read() implements the function from net.Conn
func (dc *detourConn) Read(b []byte) (n int, err error) {
	conn := dc.getConn()
	if !dc.inState(stateInitial) {
		if n, err = conn.Read(b); err != nil && err != io.EOF {
			log.Tracef("Read from %s %s failed: %s", dc.addr, dc.stateDesc(), err)
			if dc.inState(stateDirect) && blocked(err) {
				// direct route is not reliable even the first read succeeded
				// try again through detour in next dial
				log.Tracef("Seems %s still blocked, add to whitelist so will try detour next time", dc.addr)
				addToWl(dc.addr, false)
			} else if wlTemporarily(dc.addr) {
				log.Tracef("Detoured route is still not reliable for %s, not whitelist it", dc.addr)
				removeFromWl(dc.addr)
			}
			return
		}
		log.Tracef("Read %d bytes from %s %s", n, dc.addr, dc.stateDesc())
		return n, err
	}
	// state will always be settled after first read, safe to clear buffer at end of it
	defer dc.resetBuffer()
	now := time.Now()
	dl := now.Add(timeoutToDetour)
	if !dc.readDeadline.IsZero() && dc.readDeadline.Sub(now) < 2*timeoutToDetour {
		// if no enough room, reduce timeout to be half before read dead line
		dl = now.Add(dc.readDeadline.Sub(now) / 2)
	}
	conn.SetReadDeadline(dl)

	n, err = conn.Read(b)
	conn.SetReadDeadline(dc.readDeadline)
	if err != nil && err != io.EOF {
		ne := fmt.Errorf("Error while read from %s %s: %s", dc.addr, dc.stateDesc(), err)
		if blocked(err) {
			dc.detour(b)
		}
		return n, ne
	}
	log.Tracef("Read %d bytes from %s %s", n, dc.addr, dc.stateDesc())
	return n, err
}

// Write() implements the function from net.Conn
func (dc *detourConn) Write(b []byte) (n int, err error) {
	if dc.inState(stateInitial) {
		if n, err = dc.writeToBuffer(b); err != nil {
			return n, fmt.Errorf("Unable to write to local buffer: %s", err)
		}
	}
	if n, err = dc.getConn().Write(b); err != nil {
		log.Debugf("Write %d bytes to %s %s failed: %s", len(b), dc.addr, dc.stateDesc(), err)
		return
	}
	log.Debugf("Writed %d bytes to %s %s", len(b), dc.addr, dc.stateDesc())
	return
}

// Close() implements the function from net.Conn
func (dc *detourConn) Close() error {
	log.Tracef("Closing %s connection to %s", dc.stateDesc(), dc.addr)
	if wlTemporarily(dc.addr) {
		log.Tracef("no error found till closing, add %s to permanent whitelist", dc.addr)
		addToWl(dc.addr, true)
	}
	return dc.getConn().Close()
}

func (dc *detourConn) LocalAddr() net.Addr {
	return dc.getConn().LocalAddr()
}

func (dc *detourConn) RemoteAddr() net.Addr {
	return dc.getConn().RemoteAddr()
}

func (dc *detourConn) SetDeadline(t time.Time) error {
	dc.SetReadDeadline(t)
	dc.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) SetReadDeadline(t time.Time) error {
	dc.readDeadline = t
	dc.conn.SetReadDeadline(t)
	return nil
}

func (dc *detourConn) SetWriteDeadline(t time.Time) error {
	dc.writeDeadline = t
	dc.conn.SetWriteDeadline(t)
	return nil
}

func (dc *detourConn) writeToBuffer(b []byte) (n int, err error) {
	dc.muBuf.Lock()
	n, err = dc.buf.Write(b)
	dc.muBuf.Unlock()
	return
}

func (dc *detourConn) resetBuffer() {
	dc.muBuf.Lock()
	dc.buf.Reset()
	dc.muBuf.Unlock()
}

func (dc *detourConn) detour(b []byte) (n int, err error) {
	if err = dc.setupDetour(); err != nil {
		log.Errorf("Error to setup detour: %s", err)
		return
	}
	if _, err = dc.resend(); err != nil {
		err = fmt.Errorf("Error resend buffer to %s: %s", dc.addr, err)
		log.Error(err)
		return
	}
	// should getConn() again as it has changed
	if n, err = dc.getConn().Read(b); err != nil {
		log.Debugf("Read from %s %s still failed: %s", dc.addr, dc.stateDesc(), err)
		return
	}
	dc.setState(stateDetour)
	addToWl(dc.addr, false)
	log.Tracef("Read %d bytes from %s through detour, set state to %s", n, dc.addr, dc.stateDesc())
	return
}

func (dc *detourConn) resend() (int, error) {
	dc.muBuf.Lock()
	b := dc.buf.Bytes()
	dc.muBuf.Unlock()
	if len(b) > 0 {
		n, err := dc.getConn().Write(b)
		log.Tracef("Resend %d buffered bytes to %s, %d sent", len(b), dc.addr, n)
		return n, err
	}
	return 0, nil
}

func (dc *detourConn) setupDetour() error {
	c, err := dc.dialDetour("tcp", dc.addr)
	if err != nil {
		return err
	}
	log.Tracef("Dialed a new detour connection to %s", dc.addr)
	dc.setConn(c)
	return nil
}

func (dc *detourConn) getConn() (c net.Conn) {
	dc.muConn.RLock()
	defer dc.muConn.RUnlock()
	return dc.conn
}

func (dc *detourConn) setConn(c net.Conn) {
	dc.muConn.Lock()
	oldConn := dc.conn
	dc.conn = c
	dc.muConn.Unlock()
	dc.conn.SetReadDeadline(dc.readDeadline)
	dc.conn.SetWriteDeadline(dc.writeDeadline)
	log.Tracef("Replaced connection to %s from direct to detour and closing old one", dc.addr)
	oldConn.Close()
}

func (dc *detourConn) stateDesc() string {
	return statesDesc[atomic.LoadUint32(&dc.state)]
}

func (dc *detourConn) inState(s uint32) bool {
	return atomic.LoadUint32(&dc.state) == s
}

func (dc *detourConn) setState(s uint32) {
	atomic.StoreUint32(&dc.state, s)
}

func blocked(err error) bool {
	if ne, ok := err.(net.Error); ok && ne.Timeout() {
		return true
	}
	if oe, ok := err.(*net.OpError); ok && (oe.Err == syscall.EPIPE || oe.Err == syscall.ECONNRESET) {
		return true
	}
	return false
}

func whitelisted(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	_, in := whitelist[addr]
	return in
}

func wlTemporarily(addr string) bool {
	muWhitelist.RLock()
	defer muWhitelist.RUnlock()
	return whitelist[addr]
}

func addToWl(addr string, permanent bool) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	whitelist[addr] = permanent
}

func removeFromWl(addr string) {
	muWhitelist.Lock()
	defer muWhitelist.Unlock()
	delete(whitelist, addr)
}
<|endoftext|>


<commit_before>// expression.go
package engine

import (
	"fmt"

	"iotrules/mylog"
)

type Expression struct {
	Reference string
	Text      string
	Number    float64
}

func (e Expression) getNumber(n *Notif) (number float64, err error) {
	if mylog.Debugging {
		mylog.Debugf("enter Expression.getNumber %+v, %+v", e, n)
		defer func() { mylog.Debugf("exit Expression.getNumber %+v, %+v", number, err) }()
	}

	if e.Reference == "" {
		return e.Number, nil
	} else {
		return n.GetNumber(e.Reference)
	}
}

func (e Expression) getString(n *Notif) (str string, err error) {
	if mylog.Debugging {
		mylog.Debugf("enter Expression.getString %+v, %+v", e, n)
		defer func() { mylog.Debugf("exit Expression.getString %+v, %+v", str, err) }()
	}

	if e.Reference == "" {
		return e.Text, nil
	} else {
		return n.GetString(e.Reference)
	}
}

func makeExpressionFromJSON(i interface{}, isNumber bool) (exp Expression, err error) {
	if mylog.Debugging {
		mylog.Debugf("enter makeExpressionFromJSON %+v %+v", i, isNumber)
		defer func() { mylog.Debugf("exit makeExpressionFromJSON %+v  %+v", exp, err) }()
	}

	switch i := i.(type) {
	case string:
		if i[0] == '$' {
			exp.Reference = i[1:]
		} else {
			if !isNumber {
				exp.Text = i
			} else {
				return exp, fmt.Errorf("not numerical value %q in numerical condition")
			}
		}
	case int:
		if isNumber {
			exp.Number = float64(i)
		} else {
			return exp, fmt.Errorf("numerical value %v in not numerical condition")
		}
	case float64:
		if isNumber {
			exp.Number = i
		} else {
			return exp, fmt.Errorf("numerical value %v in not numerical condition")
		}
	default:
		return exp, fmt.Errorf("invalid type for expression %T", i)
	}
	return exp, nil
}

func (e Expression) ToRuleJSON() interface{} {
	if e.Reference != "" {
		return "$" + e.Reference
	} else if e.Text != "" {
		return e.Text
	} else {
		return e.Number
	}
}
<commit_msg>Fix some error messages<commit_after>// expression.go
package engine

import (
	"fmt"

	"iotrules/mylog"
)

type Expression struct {
	Reference string
	Text      string
	Number    float64
}

func (e Expression) getNumber(n *Notif) (number float64, err error) {
	if mylog.Debugging {
		mylog.Debugf("enter Expression.getNumber %+v, %+v", e, n)
		defer func() { mylog.Debugf("exit Expression.getNumber %+v, %+v", number, err) }()
	}

	if e.Reference == "" {
		return e.Number, nil
	} else {
		return n.GetNumber(e.Reference)
	}
}

func (e Expression) getString(n *Notif) (str string, err error) {
	if mylog.Debugging {
		mylog.Debugf("enter Expression.getString %+v, %+v", e, n)
		defer func() { mylog.Debugf("exit Expression.getString %+v, %+v", str, err) }()
	}

	if e.Reference == "" {
		return e.Text, nil
	} else {
		return n.GetString(e.Reference)
	}
}

func makeExpressionFromJSON(i interface{}, isNumber bool) (exp Expression, err error) {
	if mylog.Debugging {
		mylog.Debugf("enter makeExpressionFromJSON %+v %+v", i, isNumber)
		defer func() { mylog.Debugf("exit makeExpressionFromJSON %+v  %+v", exp, err) }()
	}

	switch i := i.(type) {
	case string:
		if i[0] == '$' {
			exp.Reference = i[1:]
		} else {
			if !isNumber {
				exp.Text = i
			} else {
				return exp, fmt.Errorf("non numerical value %q in numerical condition", i)
			}
		}
	case int:
		if isNumber {
			exp.Number = float64(i)
		} else {
			return exp, fmt.Errorf("numerical value %v in non numerical condition", i)
		}
	case float64:
		if isNumber {
			exp.Number = i
		} else {
			return exp, fmt.Errorf("numerical value %v in non numerical condition", i)
		}
	default:
		return exp, fmt.Errorf("invalid type for expression %T", i)
	}
	return exp, nil
}

func (e Expression) ToRuleJSON() interface{} {
	if e.Reference != "" {
		return "$" + e.Reference
	} else if e.Text != "" {
		return e.Text
	} else {
		return e.Number
	}
}
<|endoftext|>


<commit_before>package publicsuffix

import (
	"fmt"
	"strings"

	"github.com/globalsign/certlint/certdata"
	"github.com/globalsign/certlint/checks"
	"github.com/globalsign/certlint/errors"

	psl "golang.org/x/net/publicsuffix"
)

const checkName = "Public Suffix (xTLD) Check"

func init() {
	filter := &checks.Filter{
		Type: []string{"DV", "OV", "IV", "EV"},
	}
	checks.RegisterCertificateCheck(checkName, filter, Check)
}

// Check performs a strict verification on the extension according to the standard(s)
func Check(d *certdata.Data) *errors.Errors {
	var e = errors.New(nil)

	if len(d.Cert.Subject.CommonName) > 0 {
		suffix, official := psl.PublicSuffix(strings.ToLower(d.Cert.Subject.CommonName))
		if official && (fmt.Sprintf("*.%s", suffix) == d.Cert.Subject.CommonName || suffix == d.Cert.Subject.CommonName) {
			e.Err("Certificate CommonName '%s' equals '%s' from the public suffix list", d.Cert.Subject.CommonName, suffix)
		}
	}

	for _, n := range d.Cert.DNSNames {
		suffix, official := psl.PublicSuffix(strings.ToLower(n))
		if official && (fmt.Sprintf("*.%s", suffix) == n || suffix == n) {
			e.Err("Certificate subjectAltName '%s' equals '%s' from the public suffix list", n, suffix)
		}
	}

	return e
}
<commit_msg>Don't ignore an unofficial suffix<commit_after>package publicsuffix

import (
	"fmt"
	"strings"

	"github.com/globalsign/certlint/certdata"
	"github.com/globalsign/certlint/checks"
	"github.com/globalsign/certlint/errors"

	psl "golang.org/x/net/publicsuffix"
)

const checkName = "Public Suffix (xTLD) Check"

func init() {
	filter := &checks.Filter{
		Type: []string{"DV", "OV", "IV", "EV"},
	}
	checks.RegisterCertificateCheck(checkName, filter, Check)
}

// Check performs a strict verification on the extension according to the standard(s)
func Check(d *certdata.Data) *errors.Errors {
	var e = errors.New(nil)

	if len(d.Cert.Subject.CommonName) > 0 {
		suffix, _ := psl.PublicSuffix(strings.ToLower(d.Cert.Subject.CommonName))
		if fmt.Sprintf("*.%s", suffix) == d.Cert.Subject.CommonName || suffix == d.Cert.Subject.CommonName {
			e.Err("Certificate CommonName %q equals %q from the public suffix list", d.Cert.Subject.CommonName, suffix)
		}
	}

	for _, n := range d.Cert.DNSNames {
		suffix, _ := psl.PublicSuffix(strings.ToLower(n))
		if fmt.Sprintf("*.%s", suffix) == n || suffix == n {
			e.Err("Certificate subjectAltName %q equals %q from the public suffix list", n, suffix)
		}
	}

	return e
}
<|endoftext|>