<commit_before>// Copyright 2016 e-Xpert Solutions SA. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ltm

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

	"github.com/e-XpertSolutions/f5-rest-client/f5"
)

// VirtualServerList holds a list of virtual server uration.
type VirtualServerList struct {
	Items    []VirtualServer `json:"items,omitempty"`
	Kind     string          `json:"kind,omitempty" pretty:",expanded"`
	SelfLink string          `json:"selfLink,omitempty" pretty:",expanded"`
}

// VirtualServer holds the uration of a single virtual server.
type VirtualServer struct {
	AddressStatus     string `json:"addressStatus,omitempty"`
	AutoLasthop       string `json:"autoLasthop,omitempty"`
	CmpEnabled        string `json:"cmpEnabled,omitempty"`
	ConnectionLimit   int64  `json:"connectionLimit,omitempty"`
	Description       string `json:"description,omitempty"`
	Destination       string `json:"destination,omitempty"`
	Disabled          bool   `json:"disabled,omitempty"`
	Enabled           bool   `json:"enabled,omitempty"`
	FullPath          string `json:"fullPath,omitempty" pretty:",expanded"`
	Generation        int64  `json:"generation,omitempty" pretty:",expanded"`
	GtmScore          int64  `json:"gtmScore,omitempty" pretty:",expanded"`
	IPProtocol        string `json:"ipProtocol,omitempty"`
	Kind              string `json:"kind,omitempty" pretty:",expanded"`
	Mask              string `json:"mask,omitempty"`
	Mirror            string `json:"mirror,omitempty"`
	MobileAppTunnel   string `json:"mobileAppTunnel,omitempty" pretty:",expanded"`
	Name              string `json:"name,omitempty"`
	Nat64             string `json:"nat64,omitempty" pretty:",expanded"`
	Partition         string `json:"partition,omitempty"`
	PoliciesReference struct {
		IsSubcollection bool   `json:"isSubcollection,omitempty"`
		Link            string `json:"link,omitempty"`
	} `json:"policiesReference,omitempty"`
	Pool              string `json:"pool,omitempty"`
	ProfilesReference struct {
		IsSubcollection bool      `json:"isSubcollection,omitempty"`
		Link            string    `json:"link,omitempty"`
		Profiles        []Profile `json:"items,omitempty"`
	} `json:"profilesReference,omitempty"`
	RateLimit                string                   `json:"rateLimit,omitempty" pretty:",expanded"`
	RateLimitDstMask         int64                    `json:"rateLimitDstMask,omitempty" pretty:",expanded"`
	RateLimitMode            string                   `json:"rateLimitMode,omitempty" pretty:",expanded"`
	RateLimitSrcMask         int64                    `json:"rateLimitSrcMask,omitempty" pretty:",expanded"`
	Rules                    []string                 `json:"rules,omitempty"`
	SelfLink                 string                   `json:"selfLink,omitempty" pretty:",expanded"`
	Source                   string                   `json:"source,omitempty"`
	SourceAddressTranslation SourceAddressTranslation `json:"sourceAddressTranslation,omitempty"`
	SourcePort               string                   `json:"sourcePort,omitempty"`
	SynCookieStatus          string                   `json:"synCookieStatus,omitempty"`
	TranslateAddress         string                   `json:"translateAddress,omitempty"`
	TranslatePort            string                   `json:"translatePort,omitempty"`
	Vlans                    []string                 `json:"vlans,omitempty"`
	VlansDisabled            bool                     `json:"vlansDisabled,omitempty"`
	VlansEnabled             bool                     `json:"vlansEnabled,omitempty"`
	VsIndex                  int64                    `json:"vsIndex,omitempty" pretty:",expanded"`
}

type SourceAddressTranslation struct {
	Type string `json:"type,omitempty"`
}

type Profile struct {
	Name    string `json:"name,omitempty"`
	Context string `json:"context,omitempty"`
}

// VirtualEndpoint represents the REST resource for managing virtual server.
const VirtualEndpoint = "/virtual"

// VirtualResponse provide a simple mechanism to read paginated results.
//
// TODO(gilliek): use VirtualResponse object where pagination is needed.
type VirtualResponse struct {
}

// VirtualResource provides an API to manage virtual server urations.
type VirtualResource struct {
	c *f5.Client
}

// ListAll lists all the virtual server urations.
func (vr *VirtualResource) ListAll() (*VirtualServerList, error) {
	resp, err := vr.doRequest("GET", "", nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return nil, err
	}
	var vsc VirtualServerList
	dec := json.NewDecoder(resp.Body)
	if err := dec.Decode(&vsc); err != nil {
		return nil, err
	}
	return &vsc, nil
}

// Get a single virtual server uration identified by id.
func (vr *VirtualResource) Get(id string) (*VirtualServer, error) {
	resp, err := vr.doRequest("GET", id, nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return nil, err
	}
	var vsci VirtualServer
	dec := json.NewDecoder(resp.Body)
	if err := dec.Decode(&vsci); err != nil {
		return nil, err
	}
	return &vsci, nil
}

// Create a new virtual server uration.
func (vr *VirtualResource) Create(item VirtualServer) error {
	if err := vr.c.ModQuery("POST", BasePath+VirtualEndpoint, item); err != nil {
		return err
	}
	return nil
}

// Edit a virtual server uration identified by id.
func (vr *VirtualResource) Edit(id string, item VirtualServer) error {
	resp, err := vr.doRequest("PUT", id, item)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Enabling a virtual server item identified by id.
func (vr *VirtualResource) Enable(id string) error {
	item := VirtualServer{Enabled: true}
	resp, err := vr.doRequest("PATCH", id, item)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Disabling a virtual server item identified by id.
func (vr *VirtualResource) Disable(id string) error {
	item := VirtualServer{Disabled: true}
	resp, err := vr.doRequest("PATCH", id, item)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Delete a single server uration identified by id.
func (vr *VirtualResource) Delete(id string) error {
	resp, err := vr.doRequest("DELETE", id, nil)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Rules gets the iRules uration for a virtual server identified by id.
func (vr *VirtualResource) Rules(id string) ([]Rule, error) {
	resp, err := vr.doRequest("GET", id+"/rule", nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return nil, err
	}
	var rules []Rule
	dec := json.NewDecoder(resp.Body)
	if err := dec.Decode(&rules); err != nil {
		return nil, err
	}
	return rules, nil
}

// AddRule adds an iRule to the virtual server identified by id.
func (vr *VirtualResource) AddRule(id string, rule Rule) error {
	resp, err := vr.doRequest("POST", id+"/rule", rule)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// RemoveRule removes a single  iRule from the virtual server identified by id.
func (vr *VirtualResource) RemoveRule(vsID, ruleID string) error {
	resp, err := vr.doRequest("DELETE", vsID+"/rule/"+ruleID, nil)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// doRequest creates and send HTTP request using the F5 client.
//
// TODO(gilliek): decorate errors
func (vr *VirtualResource) doRequest(method, id string, data interface{}) (*http.Response, error) {
	req, err := vr.c.MakeRequest(method, BasePath+VirtualEndpoint+"/"+id, data)
	if err != nil {
		return nil, err
	}
	resp, err := vr.c.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// readError reads request error object from a HTTP response.
//
// TODO(gilliek): move this function into F5 package.
func (vr *VirtualResource) readError(resp *http.Response) error {
	if resp.StatusCode >= 400 {
		errResp, err := f5.NewRequestError(resp.Body)
		if err != nil {
			return err
		}
		return errResp
	}
	return nil
}

func (vr *VirtualResource) ShowStats(id string) (*VirtualStatsList, error) {
	var item VirtualStatsList
	if err := vr.c.ReadQuery(BasePath+VirtualEndpoint+"/"+id+"/stats", &item); err != nil {
		return nil, err
	}
	return &item, nil
}
<commit_msg>f5/ltm: add support for profiles linking in VirtualServer resource<commit_after>// Copyright 2016 e-Xpert Solutions SA. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

package ltm

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

	"github.com/e-XpertSolutions/f5-rest-client/f5"
)

// VirtualServerList holds a list of virtual server uration.
type VirtualServerList struct {
	Items    []VirtualServer `json:"items,omitempty"`
	Kind     string          `json:"kind,omitempty" pretty:",expanded"`
	SelfLink string          `json:"selfLink,omitempty" pretty:",expanded"`
}

// VirtualServer holds the uration of a single virtual server.
type VirtualServer struct {
	AddressStatus     string `json:"addressStatus,omitempty"`
	AutoLasthop       string `json:"autoLasthop,omitempty"`
	CmpEnabled        string `json:"cmpEnabled,omitempty"`
	ConnectionLimit   int64  `json:"connectionLimit,omitempty"`
	Description       string `json:"description,omitempty"`
	Destination       string `json:"destination,omitempty"`
	Disabled          bool   `json:"disabled,omitempty"`
	Enabled           bool   `json:"enabled,omitempty"`
	FullPath          string `json:"fullPath,omitempty" pretty:",expanded"`
	Generation        int64  `json:"generation,omitempty" pretty:",expanded"`
	GtmScore          int64  `json:"gtmScore,omitempty" pretty:",expanded"`
	IPProtocol        string `json:"ipProtocol,omitempty"`
	Kind              string `json:"kind,omitempty" pretty:",expanded"`
	Mask              string `json:"mask,omitempty"`
	Mirror            string `json:"mirror,omitempty"`
	MobileAppTunnel   string `json:"mobileAppTunnel,omitempty" pretty:",expanded"`
	Name              string `json:"name,omitempty"`
	Nat64             string `json:"nat64,omitempty" pretty:",expanded"`
	Partition         string `json:"partition,omitempty"`
	PoliciesReference struct {
		IsSubcollection bool   `json:"isSubcollection,omitempty"`
		Link            string `json:"link,omitempty"`
	} `json:"policiesReference,omitempty"`
	Pool              string   `json:"pool,omitempty"`
	Profiles          []string `json:"profiles,omitempty"` // only used to link existing profiles a creation or update
	ProfilesReference struct {
		IsSubcollection bool      `json:"isSubcollection,omitempty"`
		Link            string    `json:"link,omitempty"`
		Profiles        []Profile `json:"items,omitempty"`
	} `json:"profilesReference,omitempty"`
	RateLimit                string                   `json:"rateLimit,omitempty" pretty:",expanded"`
	RateLimitDstMask         int64                    `json:"rateLimitDstMask,omitempty" pretty:",expanded"`
	RateLimitMode            string                   `json:"rateLimitMode,omitempty" pretty:",expanded"`
	RateLimitSrcMask         int64                    `json:"rateLimitSrcMask,omitempty" pretty:",expanded"`
	Rules                    []string                 `json:"rules,omitempty"`
	SelfLink                 string                   `json:"selfLink,omitempty" pretty:",expanded"`
	Source                   string                   `json:"source,omitempty"`
	SourceAddressTranslation SourceAddressTranslation `json:"sourceAddressTranslation,omitempty"`
	SourcePort               string                   `json:"sourcePort,omitempty"`
	SynCookieStatus          string                   `json:"synCookieStatus,omitempty"`
	TranslateAddress         string                   `json:"translateAddress,omitempty"`
	TranslatePort            string                   `json:"translatePort,omitempty"`
	Vlans                    []string                 `json:"vlans,omitempty"`
	VlansDisabled            bool                     `json:"vlansDisabled,omitempty"`
	VlansEnabled             bool                     `json:"vlansEnabled,omitempty"`
	VsIndex                  int64                    `json:"vsIndex,omitempty" pretty:",expanded"`
}

type SourceAddressTranslation struct {
	Type string `json:"type,omitempty"`
}

type Profile struct {
	Name    string `json:"name,omitempty"`
	Context string `json:"context,omitempty"`
}

// VirtualEndpoint represents the REST resource for managing virtual server.
const VirtualEndpoint = "/virtual"

// VirtualResponse provide a simple mechanism to read paginated results.
//
// TODO(gilliek): use VirtualResponse object where pagination is needed.
type VirtualResponse struct {
}

// VirtualResource provides an API to manage virtual server urations.
type VirtualResource struct {
	c *f5.Client
}

// ListAll lists all the virtual server urations.
func (vr *VirtualResource) ListAll() (*VirtualServerList, error) {
	resp, err := vr.doRequest("GET", "", nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return nil, err
	}
	var vsc VirtualServerList
	dec := json.NewDecoder(resp.Body)
	if err := dec.Decode(&vsc); err != nil {
		return nil, err
	}
	return &vsc, nil
}

// Get a single virtual server uration identified by id.
func (vr *VirtualResource) Get(id string) (*VirtualServer, error) {
	resp, err := vr.doRequest("GET", id, nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return nil, err
	}
	var vsci VirtualServer
	dec := json.NewDecoder(resp.Body)
	if err := dec.Decode(&vsci); err != nil {
		return nil, err
	}
	return &vsci, nil
}

// Create a new virtual server uration.
func (vr *VirtualResource) Create(item VirtualServer) error {
	if err := vr.c.ModQuery("POST", BasePath+VirtualEndpoint, item); err != nil {
		return err
	}
	return nil
}

// Edit a virtual server uration identified by id.
func (vr *VirtualResource) Edit(id string, item VirtualServer) error {
	resp, err := vr.doRequest("PUT", id, item)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Enabling a virtual server item identified by id.
func (vr *VirtualResource) Enable(id string) error {
	item := VirtualServer{Enabled: true}
	resp, err := vr.doRequest("PATCH", id, item)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Disabling a virtual server item identified by id.
func (vr *VirtualResource) Disable(id string) error {
	item := VirtualServer{Disabled: true}
	resp, err := vr.doRequest("PATCH", id, item)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Delete a single server uration identified by id.
func (vr *VirtualResource) Delete(id string) error {
	resp, err := vr.doRequest("DELETE", id, nil)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// Rules gets the iRules uration for a virtual server identified by id.
func (vr *VirtualResource) Rules(id string) ([]Rule, error) {
	resp, err := vr.doRequest("GET", id+"/rule", nil)
	if err != nil {
		return nil, err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return nil, err
	}
	var rules []Rule
	dec := json.NewDecoder(resp.Body)
	if err := dec.Decode(&rules); err != nil {
		return nil, err
	}
	return rules, nil
}

// AddRule adds an iRule to the virtual server identified by id.
func (vr *VirtualResource) AddRule(id string, rule Rule) error {
	resp, err := vr.doRequest("POST", id+"/rule", rule)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// RemoveRule removes a single  iRule from the virtual server identified by id.
func (vr *VirtualResource) RemoveRule(vsID, ruleID string) error {
	resp, err := vr.doRequest("DELETE", vsID+"/rule/"+ruleID, nil)
	if err != nil {
		return err
	}
	defer resp.Body.Close()
	if err := vr.readError(resp); err != nil {
		return err
	}
	return nil
}

// doRequest creates and send HTTP request using the F5 client.
//
// TODO(gilliek): decorate errors
func (vr *VirtualResource) doRequest(method, id string, data interface{}) (*http.Response, error) {
	req, err := vr.c.MakeRequest(method, BasePath+VirtualEndpoint+"/"+id, data)
	if err != nil {
		return nil, err
	}
	resp, err := vr.c.Do(req)
	if err != nil {
		return nil, err
	}
	return resp, nil
}

// readError reads request error object from a HTTP response.
//
// TODO(gilliek): move this function into F5 package.
func (vr *VirtualResource) readError(resp *http.Response) error {
	if resp.StatusCode >= 400 {
		errResp, err := f5.NewRequestError(resp.Body)
		if err != nil {
			return err
		}
		return errResp
	}
	return nil
}

func (vr *VirtualResource) ShowStats(id string) (*VirtualStatsList, error) {
	var item VirtualStatsList
	if err := vr.c.ReadQuery(BasePath+VirtualEndpoint+"/"+id+"/stats", &item); err != nil {
		return nil, err
	}
	return &item, nil
}
<|endoftext|>


<commit_before>package main

import (
	"io/ioutil"
	"os"
	"path"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/itchio/butler/comm"
	"github.com/itchio/wharf/counter"
	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/sync"
	"github.com/itchio/wharf/tlc"
	"github.com/itchio/wharf/wire"
)

func diff(target string, source string, patch string, compression pwr.CompressionSettings) {
	must(doDiff(target, source, patch, compression))
}

func doDiff(target string, source string, patch string, compression pwr.CompressionSettings) error {
	startTime := time.Now()

	var targetSignature []sync.BlockHash
	var targetContainer *tlc.Container

	if target == "/dev/null" {
		targetContainer = &tlc.Container{}
	} else {
		targetInfo, err := os.Lstat(target)
		if err != nil {
			return err
		}

		if targetInfo.IsDir() {
			comm.Opf("Hashing %s", target)
			targetContainer, err = tlc.Walk(target, filterPaths)
			if err != nil {
				return err
			}

			comm.StartProgress()
			targetSignature, err = pwr.ComputeSignature(targetContainer, target, comm.NewStateConsumer())
			comm.EndProgress()
			if err != nil {
				return err
			}

			{
				prettySize := humanize.Bytes(uint64(targetContainer.Size))
				perSecond := humanize.Bytes(uint64(float64(targetContainer.Size) / time.Since(startTime).Seconds()))
				comm.Statf("%s (%s) @ %s/s\n", prettySize, targetContainer.Stats(), perSecond)
			}
		} else {
			comm.Opf("Reading signature from %s", target)
			signatureReader, err := os.Open(target)
			if err != nil {
				return err
			}
			targetContainer, targetSignature, err = pwr.ReadSignature(signatureReader)
			if err != nil {
				return err
			}
			err = signatureReader.Close()
			if err != nil {
				return err
			}
		}

	}

	startTime = time.Now()

	var sourceContainer *tlc.Container
	if source == "/dev/null" {
		sourceContainer = &tlc.Container{}
	} else {
		var err error
		sourceContainer, err = tlc.Walk(source, filterPaths)
		if err != nil {
			return err
		}
	}

	patchWriter, err := os.Create(patch)
	if err != nil {
		return err
	}
	defer patchWriter.Close()

	signaturePath := patch + ".sig"
	signatureWriter, err := os.Create(signaturePath)
	if err != nil {
		return err
	}
	defer signatureWriter.Close()

	patchCounter := counter.NewWriter(patchWriter)
	signatureCounter := counter.NewWriter(signatureWriter)

	dctx := &pwr.DiffContext{
		SourceContainer: sourceContainer,
		FilePool:        sourceContainer.NewFilePool(source),

		TargetContainer: targetContainer,
		TargetSignature: targetSignature,

		Consumer:    comm.NewStateConsumer(),
		Compression: &compression,
	}

	comm.Opf("Diffing %s", source)
	comm.StartProgress()
	err = dctx.WritePatch(patchCounter, signatureCounter)
	if err != nil {
		return err
	}
	comm.EndProgress()

	{
		prettySize := humanize.Bytes(uint64(sourceContainer.Size))
		perSecond := humanize.Bytes(uint64(float64(sourceContainer.Size) / time.Since(startTime).Seconds()))
		comm.Statf("%s (%s) @ %s/s\n", prettySize, sourceContainer.Stats(), perSecond)
	}

	if *diffArgs.verify {
		tmpDir, err := ioutil.TempDir("", "pwr")
		if err != nil {
			return err
		}
		defer os.RemoveAll(tmpDir)

		apply(patch, target, tmpDir, false, signaturePath)
	}

	{
		prettyPatchSize := humanize.Bytes(uint64(patchCounter.Count()))
		percReused := 100.0 * float64(dctx.ReusedBytes) / float64(dctx.FreshBytes+dctx.ReusedBytes)
		relToNew := 100.0 * float64(patchCounter.Count()) / float64(sourceContainer.Size)
		prettyFreshSize := humanize.Bytes(uint64(dctx.FreshBytes))

		comm.Statf("Re-used %.2f%% of old, added %s fresh data", percReused, prettyFreshSize)
		comm.Statf("%s patch (%.2f%% of the full size)", prettyPatchSize, relToNew)
	}

	return nil
}

func apply(patch string, target string, output string, inplace bool, sigpath string) {
	must(doApply(patch, target, output, inplace, sigpath))
}

func doApply(patch string, target string, output string, inplace bool, sigpath string) error {
	if output == "" {
		output = target
	}

	target = path.Clean(target)
	output = path.Clean(output)
	if output == target {
		if !inplace {
			comm.Dief("Refusing to destructively patch %s without --inplace", output)
		}
	}

	comm.Opf("Patching %s", output)
	startTime := time.Now()

	patchReader, err := os.Open(patch)
	if err != nil {
		return err
	}

	actx := &pwr.ApplyContext{
		TargetPath:        target,
		OutputPath:        output,
		InPlace:           inplace,
		SignatureFilePath: sigpath,

		Consumer: comm.NewStateConsumer(),
	}

	comm.StartProgress()
	err = actx.ApplyPatch(patchReader)
	if err != nil {
		return err
	}
	comm.EndProgress()

	container := actx.SourceContainer
	prettySize := humanize.Bytes(uint64(container.Size))
	perSecond := humanize.Bytes(uint64(float64(container.Size) / time.Since(startTime).Seconds()))

	if actx.InPlace {
		comm.Statf("patched %d, kept %d, deleted %d (%s stage)", actx.TouchedFiles, actx.NoopFiles, actx.DeletedFiles, humanize.Bytes(uint64(actx.StageSize)))
	}
	comm.Statf("%s (%s) @ %s/s\n", prettySize, container.Stats(), perSecond)

	return nil
}

func sign(output string, signature string, compression pwr.CompressionSettings) {
	must(doSign(output, signature, compression))
}

func doSign(output string, signature string, compression pwr.CompressionSettings) error {
	comm.Opf("Creating signature for %s", output)
	startTime := time.Now()

	container, err := tlc.Walk(output, filterPaths)
	if err != nil {
		return err
	}

	signatureWriter, err := os.Create(signature)
	if err != nil {
		return err
	}

	rawSigWire := wire.NewWriteContext(signatureWriter)
	rawSigWire.WriteMagic(pwr.SignatureMagic)

	rawSigWire.WriteMessage(&pwr.SignatureHeader{
		Compression: &compression,
	})

	sigWire, err := pwr.CompressWire(rawSigWire, &compression)
	if err != nil {
		return err
	}
	sigWire.WriteMessage(container)

	comm.StartProgress()
	err = pwr.ComputeSignatureToWriter(container, output, comm.NewStateConsumer(), func(hash sync.BlockHash) error {
		return sigWire.WriteMessage(&pwr.BlockHash{
			WeakHash:   hash.WeakHash,
			StrongHash: hash.StrongHash,
		})
	})
	comm.EndProgress()
	if err != nil {
		return err
	}

	err = sigWire.Close()
	if err != nil {
		return err
	}

	prettySize := humanize.Bytes(uint64(container.Size))
	perSecond := humanize.Bytes(uint64(float64(container.Size) / time.Since(startTime).Seconds()))
	comm.Statf("%s (%s) @ %s/s\n", prettySize, container.Stats(), perSecond)

	return nil
}

func verify(signature string, output string) {
	must(doVerify(signature, output))
}

func doVerify(signature string, output string) error {
	comm.Opf("Verifying %s", output)
	startTime := time.Now()

	signatureReader, err := os.Open(signature)
	if err != nil {
		return err
	}
	defer signatureReader.Close()

	refContainer, refHashes, err := pwr.ReadSignature(signatureReader)
	if err != nil {
		return err
	}

	comm.StartProgress()
	hashes, err := pwr.ComputeSignature(refContainer, output, comm.NewStateConsumer())
	comm.EndProgress()
	if err != nil {
		return err
	}

	err = pwr.CompareHashes(refHashes, hashes, refContainer)
	if err != nil {
		comm.Logf(err.Error())
		comm.Dief("Some checks failed after checking %d block.", len(refHashes))
	}

	prettySize := humanize.Bytes(uint64(refContainer.Size))
	perSecond := humanize.Bytes(uint64(float64(refContainer.Size) / time.Since(startTime).Seconds()))
	comm.Statf("%s (%s) @ %s/s\n", prettySize, refContainer.Stats(), perSecond)

	return nil
}
<commit_msg>Preliminary zip signature support<commit_after>package main

import (
	"archive/zip"
	"io/ioutil"
	"os"
	"path"
	"time"

	"github.com/dustin/go-humanize"
	"github.com/itchio/butler/comm"
	"github.com/itchio/wharf/counter"
	"github.com/itchio/wharf/pwr"
	"github.com/itchio/wharf/sync"
	"github.com/itchio/wharf/tlc"
	"github.com/itchio/wharf/wire"
)

func diff(target string, source string, patch string, compression pwr.CompressionSettings) {
	must(doDiff(target, source, patch, compression))
}

func doDiff(target string, source string, patch string, compression pwr.CompressionSettings) error {
	startTime := time.Now()

	var targetSignature []sync.BlockHash
	var targetContainer *tlc.Container

	if target == "/dev/null" {
		targetContainer = &tlc.Container{}
	} else {
		targetInfo, err := os.Lstat(target)
		if err != nil {
			return err
		}

		if targetInfo.IsDir() {
			comm.Opf("Hashing %s", target)
			targetContainer, err = tlc.Walk(target, filterPaths)
			if err != nil {
				return err
			}

			comm.StartProgress()
			targetSignature, err = pwr.ComputeSignature(targetContainer, targetContainer.NewFilePool(target), comm.NewStateConsumer())
			comm.EndProgress()
			if err != nil {
				return err
			}

			{
				prettySize := humanize.Bytes(uint64(targetContainer.Size))
				perSecond := humanize.Bytes(uint64(float64(targetContainer.Size) / time.Since(startTime).Seconds()))
				comm.Statf("%s (%s) @ %s/s\n", prettySize, targetContainer.Stats(), perSecond)
			}
		} else {
			signatureReader, err := os.Open(target)
			if err != nil {
				return err
			}

			targetContainer, targetSignature, err = pwr.ReadSignature(signatureReader)
			if err != nil {
				if err != wire.ErrFormat {
					return err
				}

				_, err = signatureReader.Seek(0, os.SEEK_SET)
				if err != nil {
					return err
				}

				stats, err := os.Lstat(target)
				if err != nil {
					return err
				}

				zr, err := zip.NewReader(signatureReader, stats.Size())
				if err != nil {
					return err
				}

				targetContainer, err = tlc.WalkZip(zr, filterPaths)
				if err != nil {
					return err
				}
				comm.Opf("Walking archive (%s)", targetContainer.Stats())

				// targetContainer.NewZipFilePool(zr)

				comm.StartProgress()
				targetSignature, err = pwr.ComputeSignature(targetContainer, targetContainer.NewFilePool(target), comm.NewStateConsumer())
				comm.EndProgress()
				if err != nil {
					return err
				}

				{
					prettySize := humanize.Bytes(uint64(targetContainer.Size))
					perSecond := humanize.Bytes(uint64(float64(targetContainer.Size) / time.Since(startTime).Seconds()))
					comm.Statf("%s (%s) @ %s/s\n", prettySize, targetContainer.Stats(), perSecond)
				}
				comm.Opf("Read signature from %s", target)
			} else {
				comm.Opf("Read signature from %s", target)
			}

			err = signatureReader.Close()
			if err != nil {
				return err
			}
		}

	}

	startTime = time.Now()

	var sourceContainer *tlc.Container
	if source == "/dev/null" {
		sourceContainer = &tlc.Container{}
	} else {
		var err error
		sourceContainer, err = tlc.Walk(source, filterPaths)
		if err != nil {
			return err
		}
	}

	patchWriter, err := os.Create(patch)
	if err != nil {
		return err
	}
	defer patchWriter.Close()

	signaturePath := patch + ".sig"
	signatureWriter, err := os.Create(signaturePath)
	if err != nil {
		return err
	}
	defer signatureWriter.Close()

	patchCounter := counter.NewWriter(patchWriter)
	signatureCounter := counter.NewWriter(signatureWriter)

	dctx := &pwr.DiffContext{
		SourceContainer: sourceContainer,
		FilePool:        sourceContainer.NewFilePool(source),

		TargetContainer: targetContainer,
		TargetSignature: targetSignature,

		Consumer:    comm.NewStateConsumer(),
		Compression: &compression,
	}

	comm.Opf("Diffing %s", source)
	comm.StartProgress()
	err = dctx.WritePatch(patchCounter, signatureCounter)
	if err != nil {
		return err
	}
	comm.EndProgress()

	{
		prettySize := humanize.Bytes(uint64(sourceContainer.Size))
		perSecond := humanize.Bytes(uint64(float64(sourceContainer.Size) / time.Since(startTime).Seconds()))
		comm.Statf("%s (%s) @ %s/s\n", prettySize, sourceContainer.Stats(), perSecond)
	}

	if *diffArgs.verify {
		tmpDir, err := ioutil.TempDir("", "pwr")
		if err != nil {
			return err
		}
		defer os.RemoveAll(tmpDir)

		apply(patch, target, tmpDir, false, signaturePath)
	}

	{
		prettyPatchSize := humanize.Bytes(uint64(patchCounter.Count()))
		percReused := 100.0 * float64(dctx.ReusedBytes) / float64(dctx.FreshBytes+dctx.ReusedBytes)
		relToNew := 100.0 * float64(patchCounter.Count()) / float64(sourceContainer.Size)
		prettyFreshSize := humanize.Bytes(uint64(dctx.FreshBytes))

		comm.Statf("Re-used %.2f%% of old, added %s fresh data", percReused, prettyFreshSize)
		comm.Statf("%s patch (%.2f%% of the full size)", prettyPatchSize, relToNew)
	}

	return nil
}

func apply(patch string, target string, output string, inplace bool, sigpath string) {
	must(doApply(patch, target, output, inplace, sigpath))
}

func doApply(patch string, target string, output string, inplace bool, sigpath string) error {
	if output == "" {
		output = target
	}

	target = path.Clean(target)
	output = path.Clean(output)
	if output == target {
		if !inplace {
			comm.Dief("Refusing to destructively patch %s without --inplace", output)
		}
	}

	comm.Opf("Patching %s", output)
	startTime := time.Now()

	patchReader, err := os.Open(patch)
	if err != nil {
		return err
	}

	actx := &pwr.ApplyContext{
		TargetPath:        target,
		OutputPath:        output,
		InPlace:           inplace,
		SignatureFilePath: sigpath,

		Consumer: comm.NewStateConsumer(),
	}

	comm.StartProgress()
	err = actx.ApplyPatch(patchReader)
	if err != nil {
		return err
	}
	comm.EndProgress()

	container := actx.SourceContainer
	prettySize := humanize.Bytes(uint64(container.Size))
	perSecond := humanize.Bytes(uint64(float64(container.Size) / time.Since(startTime).Seconds()))

	if actx.InPlace {
		comm.Statf("patched %d, kept %d, deleted %d (%s stage)", actx.TouchedFiles, actx.NoopFiles, actx.DeletedFiles, humanize.Bytes(uint64(actx.StageSize)))
	}
	comm.Statf("%s (%s) @ %s/s\n", prettySize, container.Stats(), perSecond)

	return nil
}

func sign(output string, signature string, compression pwr.CompressionSettings) {
	must(doSign(output, signature, compression))
}

func doSign(output string, signature string, compression pwr.CompressionSettings) error {
	comm.Opf("Creating signature for %s", output)
	startTime := time.Now()

	container, err := tlc.Walk(output, filterPaths)
	if err != nil {
		return err
	}

	signatureWriter, err := os.Create(signature)
	if err != nil {
		return err
	}

	rawSigWire := wire.NewWriteContext(signatureWriter)
	rawSigWire.WriteMagic(pwr.SignatureMagic)

	rawSigWire.WriteMessage(&pwr.SignatureHeader{
		Compression: &compression,
	})

	sigWire, err := pwr.CompressWire(rawSigWire, &compression)
	if err != nil {
		return err
	}
	sigWire.WriteMessage(container)

	comm.StartProgress()
	err = pwr.ComputeSignatureToWriter(container, container.NewFilePool(output), comm.NewStateConsumer(), func(hash sync.BlockHash) error {
		return sigWire.WriteMessage(&pwr.BlockHash{
			WeakHash:   hash.WeakHash,
			StrongHash: hash.StrongHash,
		})
	})
	comm.EndProgress()
	if err != nil {
		return err
	}

	err = sigWire.Close()
	if err != nil {
		return err
	}

	prettySize := humanize.Bytes(uint64(container.Size))
	perSecond := humanize.Bytes(uint64(float64(container.Size) / time.Since(startTime).Seconds()))
	comm.Statf("%s (%s) @ %s/s\n", prettySize, container.Stats(), perSecond)

	return nil
}

func verify(signature string, output string) {
	must(doVerify(signature, output))
}

func doVerify(signature string, output string) error {
	comm.Opf("Verifying %s", output)
	startTime := time.Now()

	signatureReader, err := os.Open(signature)
	if err != nil {
		return err
	}
	defer signatureReader.Close()

	refContainer, refHashes, err := pwr.ReadSignature(signatureReader)
	if err != nil {
		return err
	}

	comm.StartProgress()
	hashes, err := pwr.ComputeSignature(refContainer, refContainer.NewFilePool(output), comm.NewStateConsumer())
	comm.EndProgress()
	if err != nil {
		return err
	}

	err = pwr.CompareHashes(refHashes, hashes, refContainer)
	if err != nil {
		comm.Logf(err.Error())
		comm.Dief("Some checks failed after checking %d block.", len(refHashes))
	}

	prettySize := humanize.Bytes(uint64(refContainer.Size))
	perSecond := humanize.Bytes(uint64(float64(refContainer.Size) / time.Since(startTime).Seconds()))
	comm.Statf("%s (%s) @ %s/s\n", prettySize, refContainer.Stats(), perSecond)

	return nil
}
<|endoftext|>


<commit_before>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package network_test

import (
	"io/ioutil"
	"net"
	"path/filepath"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/network"
	"github.com/juju/juju/testing"
)

type InterfaceInfoSuite struct {
	info []network.InterfaceInfo
}

var _ = gc.Suite(&InterfaceInfoSuite{})

func (s *InterfaceInfoSuite) SetUpTest(c *gc.C) {
	s.info = []network.InterfaceInfo{
		{VLANTag: 1, DeviceIndex: 0, InterfaceName: "eth0"},
		{VLANTag: 0, DeviceIndex: 1, InterfaceName: "eth1"},
		{VLANTag: 42, DeviceIndex: 2, InterfaceName: "br2"},
		{ConfigType: network.ConfigDHCP, NoAutoStart: true},
		{Address: network.NewAddress("0.1.2.3")},
		{DNSServers: network.NewAddresses("1.1.1.1", "2.2.2.2")},
		{GatewayAddress: network.NewAddress("4.3.2.1")},
		{ExtraConfig: map[string]string{
			"foo": "bar",
			"baz": "nonsense",
		}},
	}
}

func (s *InterfaceInfoSuite) TestActualInterfaceName(c *gc.C) {
	c.Check(s.info[0].ActualInterfaceName(), gc.Equals, "eth0.1")
	c.Check(s.info[1].ActualInterfaceName(), gc.Equals, "eth1")
	c.Check(s.info[2].ActualInterfaceName(), gc.Equals, "br2.42")
}

func (s *InterfaceInfoSuite) TestIsVirtual(c *gc.C) {
	c.Check(s.info[0].IsVirtual(), jc.IsTrue)
	c.Check(s.info[1].IsVirtual(), jc.IsFalse)
	c.Check(s.info[2].IsVirtual(), jc.IsTrue)
}

func (s *InterfaceInfoSuite) TestIsVLAN(c *gc.C) {
	c.Check(s.info[0].IsVLAN(), jc.IsTrue)
	c.Check(s.info[1].IsVLAN(), jc.IsFalse)
	c.Check(s.info[2].IsVLAN(), jc.IsTrue)
}

func (s *InterfaceInfoSuite) TestAdditionalFields(c *gc.C) {
	c.Check(s.info[3].ConfigType, gc.Equals, network.ConfigDHCP)
	c.Check(s.info[3].NoAutoStart, jc.IsTrue)
	c.Check(s.info[4].Address, jc.DeepEquals, network.NewAddress("0.1.2.3"))
	c.Check(s.info[5].DNSServers, jc.DeepEquals, network.NewAddresses("1.1.1.1", "2.2.2.2"))
	c.Check(s.info[6].GatewayAddress, jc.DeepEquals, network.NewAddress("4.3.2.1"))
	c.Check(s.info[7].ExtraConfig, jc.DeepEquals, map[string]string{
		"foo": "bar",
		"baz": "nonsense",
	})
}

func (s *InterfaceInfoSuite) TestSortInterfaceInfo(c *gc.C) {
	info := []network.InterfaceInfo{
		{VLANTag: 42, DeviceIndex: 2, InterfaceName: "br2"},
		{VLANTag: 0, DeviceIndex: 1, InterfaceName: "eth1"},
		{VLANTag: 1, DeviceIndex: 0, InterfaceName: "eth0"},
	}
	expectedInfo := []network.InterfaceInfo{
		{VLANTag: 1, DeviceIndex: 0, InterfaceName: "eth0"},
		{VLANTag: 0, DeviceIndex: 1, InterfaceName: "eth1"},
		{VLANTag: 42, DeviceIndex: 2, InterfaceName: "br2"},
	}
	network.SortInterfaceInfo(info)
	c.Assert(info, jc.DeepEquals, expectedInfo)
}

type NetworkSuite struct {
	testing.BaseSuite
}

var _ = gc.Suite(&NetworkSuite{})

func (*NetworkSuite) TestInitializeFromConfig(c *gc.C) {
	c.Check(network.GetPreferIPv6(), jc.IsFalse)

	envConfig := testing.CustomEnvironConfig(c, testing.Attrs{
		"prefer-ipv6": true,
	})
	network.InitializeFromConfig(envConfig)
	c.Check(network.GetPreferIPv6(), jc.IsTrue)

	envConfig = testing.CustomEnvironConfig(c, testing.Attrs{
		"prefer-ipv6": false,
	})
	network.InitializeFromConfig(envConfig)
	c.Check(network.GetPreferIPv6(), jc.IsFalse)
}

func (s *NetworkSuite) TestFilterLXCAddresses(c *gc.C) {
	lxcFakeNetConfig := filepath.Join(c.MkDir(), "lxc-net")
	netConf := []byte(`
  # comments ignored
LXC_BR= ignored
LXC_ADDR = "fooo"
 LXC_BRIDGE = " foobar " # detected, spaces stripped
anything else ignored
LXC_BRIDGE="ignored"`[1:])
	err := ioutil.WriteFile(lxcFakeNetConfig, netConf, 0644)
	c.Assert(err, jc.ErrorIsNil)
	s.PatchValue(&network.InterfaceByNameAddrs, func(name string) ([]net.Addr, error) {
		c.Assert(name, gc.Equals, "foobar")
		return []net.Addr{
			&net.IPAddr{IP: net.IPv4(10, 0, 3, 1)},
			&net.IPAddr{IP: net.IPv4(10, 0, 3, 4)},
			// Try a CIDR 10.0.3.5/24 as well.
			&net.IPNet{IP: net.IPv4(10, 0, 3, 5), Mask: net.IPv4Mask(255, 255, 255, 0)},
		}, nil
	})
	s.PatchValue(&network.LXCNetDefaultConfig, lxcFakeNetConfig)

	inputAddresses := network.NewAddresses(
		"127.0.0.1",
		"2001:db8::1",
		"10.0.0.1",
		"10.0.3.1", // filtered (directly as IP)
		"10.0.3.3", // filtered (by the 10.0.3.5/24 CIDR)
		"10.0.3.5", // filtered (directly)
		"10.0.3.4", // filtered (directly)
		"192.168.123.42",
	)
	filteredAddresses := network.NewAddresses(
		"127.0.0.1",
		"2001:db8::1",
		"10.0.0.1",
		"192.168.123.42",
	)
	c.Assert(network.FilterLXCAddresses(inputAddresses), jc.DeepEquals, filteredAddresses)
}
<commit_msg>Add test for InterfaceInfo<commit_after>// Copyright 2014 Canonical Ltd.
// Licensed under the AGPLv3, see LICENCE file for details.

package network_test

import (
	"io/ioutil"
	"net"
	"path/filepath"

	jc "github.com/juju/testing/checkers"
	gc "gopkg.in/check.v1"

	"github.com/juju/juju/network"
	"github.com/juju/juju/testing"
)

type InterfaceInfoSuite struct {
	info []network.InterfaceInfo
}

var _ = gc.Suite(&InterfaceInfoSuite{})

func (s *InterfaceInfoSuite) SetUpTest(c *gc.C) {
	s.info = []network.InterfaceInfo{
		{VLANTag: 1, DeviceIndex: 0, InterfaceName: "eth0"},
		{VLANTag: 0, DeviceIndex: 1, InterfaceName: "eth1"},
		{VLANTag: 42, DeviceIndex: 2, InterfaceName: "br2"},
		{ConfigType: network.ConfigDHCP, NoAutoStart: true},
		{Address: network.NewAddress("0.1.2.3")},
		{DNSServers: network.NewAddresses("1.1.1.1", "2.2.2.2")},
		{GatewayAddress: network.NewAddress("4.3.2.1")},
		{ExtraConfig: map[string]string{
			"foo": "bar",
			"baz": "nonsense",
		}},
		{AvailabilityZones: []string{"foo", "bar"}},
	}
}

func (s *InterfaceInfoSuite) TestActualInterfaceName(c *gc.C) {
	c.Check(s.info[0].ActualInterfaceName(), gc.Equals, "eth0.1")
	c.Check(s.info[1].ActualInterfaceName(), gc.Equals, "eth1")
	c.Check(s.info[2].ActualInterfaceName(), gc.Equals, "br2.42")
}

func (s *InterfaceInfoSuite) TestIsVirtual(c *gc.C) {
	c.Check(s.info[0].IsVirtual(), jc.IsTrue)
	c.Check(s.info[1].IsVirtual(), jc.IsFalse)
	c.Check(s.info[2].IsVirtual(), jc.IsTrue)
}

func (s *InterfaceInfoSuite) TestIsVLAN(c *gc.C) {
	c.Check(s.info[0].IsVLAN(), jc.IsTrue)
	c.Check(s.info[1].IsVLAN(), jc.IsFalse)
	c.Check(s.info[2].IsVLAN(), jc.IsTrue)
}

func (s *InterfaceInfoSuite) TestAdditionalFields(c *gc.C) {
	c.Check(s.info[3].ConfigType, gc.Equals, network.ConfigDHCP)
	c.Check(s.info[3].NoAutoStart, jc.IsTrue)
	c.Check(s.info[4].Address, jc.DeepEquals, network.NewAddress("0.1.2.3"))
	c.Check(s.info[5].DNSServers, jc.DeepEquals, network.NewAddresses("1.1.1.1", "2.2.2.2"))
	c.Check(s.info[6].GatewayAddress, jc.DeepEquals, network.NewAddress("4.3.2.1"))
	c.Check(s.info[7].ExtraConfig, jc.DeepEquals, map[string]string{
		"foo": "bar",
		"baz": "nonsense",
	})
	c.Check(s.info[8].AvailabilityZones, jc.DeepEquals, []string{"foo", "bar"})
}

func (s *InterfaceInfoSuite) TestSortInterfaceInfo(c *gc.C) {
	info := []network.InterfaceInfo{
		{VLANTag: 42, DeviceIndex: 2, InterfaceName: "br2"},
		{VLANTag: 0, DeviceIndex: 1, InterfaceName: "eth1"},
		{VLANTag: 1, DeviceIndex: 0, InterfaceName: "eth0"},
	}
	expectedInfo := []network.InterfaceInfo{
		{VLANTag: 1, DeviceIndex: 0, InterfaceName: "eth0"},
		{VLANTag: 0, DeviceIndex: 1, InterfaceName: "eth1"},
		{VLANTag: 42, DeviceIndex: 2, InterfaceName: "br2"},
	}
	network.SortInterfaceInfo(info)
	c.Assert(info, jc.DeepEquals, expectedInfo)
}

type NetworkSuite struct {
	testing.BaseSuite
}

var _ = gc.Suite(&NetworkSuite{})

func (*NetworkSuite) TestInitializeFromConfig(c *gc.C) {
	c.Check(network.GetPreferIPv6(), jc.IsFalse)

	envConfig := testing.CustomEnvironConfig(c, testing.Attrs{
		"prefer-ipv6": true,
	})
	network.InitializeFromConfig(envConfig)
	c.Check(network.GetPreferIPv6(), jc.IsTrue)

	envConfig = testing.CustomEnvironConfig(c, testing.Attrs{
		"prefer-ipv6": false,
	})
	network.InitializeFromConfig(envConfig)
	c.Check(network.GetPreferIPv6(), jc.IsFalse)
}

func (s *NetworkSuite) TestFilterLXCAddresses(c *gc.C) {
	lxcFakeNetConfig := filepath.Join(c.MkDir(), "lxc-net")
	netConf := []byte(`
  # comments ignored
LXC_BR= ignored
LXC_ADDR = "fooo"
 LXC_BRIDGE = " foobar " # detected, spaces stripped
anything else ignored
LXC_BRIDGE="ignored"`[1:])
	err := ioutil.WriteFile(lxcFakeNetConfig, netConf, 0644)
	c.Assert(err, jc.ErrorIsNil)
	s.PatchValue(&network.InterfaceByNameAddrs, func(name string) ([]net.Addr, error) {
		c.Assert(name, gc.Equals, "foobar")
		return []net.Addr{
			&net.IPAddr{IP: net.IPv4(10, 0, 3, 1)},
			&net.IPAddr{IP: net.IPv4(10, 0, 3, 4)},
			// Try a CIDR 10.0.3.5/24 as well.
			&net.IPNet{IP: net.IPv4(10, 0, 3, 5), Mask: net.IPv4Mask(255, 255, 255, 0)},
		}, nil
	})
	s.PatchValue(&network.LXCNetDefaultConfig, lxcFakeNetConfig)

	inputAddresses := network.NewAddresses(
		"127.0.0.1",
		"2001:db8::1",
		"10.0.0.1",
		"10.0.3.1", // filtered (directly as IP)
		"10.0.3.3", // filtered (by the 10.0.3.5/24 CIDR)
		"10.0.3.5", // filtered (directly)
		"10.0.3.4", // filtered (directly)
		"192.168.123.42",
	)
	filteredAddresses := network.NewAddresses(
		"127.0.0.1",
		"2001:db8::1",
		"10.0.0.1",
		"192.168.123.42",
	)
	c.Assert(network.FilterLXCAddresses(inputAddresses), jc.DeepEquals, filteredAddresses)
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"net/textproto"
	"strings"
)

var botlist []*Bot

// TCPServer simple tcp server for commands
func TCPServer() (ret int) {
	ln, err := net.Listen("tcp", ":"+TCPPort)
	if err != nil {
		log.Println("[control] error listening:", err.Error())
		return 1
	}
	log.Printf("[control] listening to port %s for connections...", TCPPort)
	defer ln.Close()

	for {
		conn, err := ln.Accept()
		bot := NewBot()
		bot.inconn = conn
		botlist = append(botlist, bot)
		if err != nil {
			fmt.Println("[control] error accepting: ", err.Error())
			return 1
		}
		go handleRequest(conn, bot)
	}
}

// CloseBot closes all bots
func CloseBot(bot *Bot) {
	// Iterate over the list of bots
	for i := range botlist {
		if bot == botlist[i] {
			// Remove the closed bot from the list
			botlist = append(botlist[:i], botlist[i+1:]...)
			return
		}
	}

	// Let the bot clean itself up
	bot.Close()
}

func handleRequest(conn net.Conn, bot *Bot) {
	defer CloseBot(bot)
	defer bot.Close()

	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()

		if err != nil {
			fmt.Println("[control] read error:", err)
			fmt.Fprintf(conn, "[control] read error: %s", err)
			conn.Close()
			return
		}
		if !strings.HasPrefix(line, "PASS ") {
			log.Println("[control] " + line)
		}
		handleMessage(line, bot)
	}
}

// Handle an IRC received from a bot
func handleMessage(message string, bot *Bot) {
	remoteAddr := bot.inconn.RemoteAddr().String()
	remoteAddrIP := strings.Split(remoteAddr, ":")

	if strings.HasPrefix(message, "JOIN ") {
		joinComm := strings.Split(message, "JOIN ")
		go bot.Join(joinComm[1])
	} else if strings.HasPrefix(message, "PART ") {
		partComm := strings.Split(message, "PART ")
		go bot.Part(partComm[1])
	} else if strings.HasPrefix(message, "PASS ") {
		passComm := strings.Split(message, "PASS ")
		passwordParts := strings.Split(passComm[1], ";")
		if passwordParts[0] == TCPPass {
			bot.oauth = passwordParts[1]
			log.Printf("[control] authenticated! %s\n", remoteAddrIP)
		} else {
			log.Printf("[control] invalid broker pass! %s\n", remoteAddrIP)
			bot.inconn.Close()
			return
		}
	} else if strings.HasPrefix(message, "NICK ") || strings.HasPrefix(message, "USER ") {
		if strings.HasPrefix(message, "NICK ") {
			nickComm := strings.Split(message, "NICK ")
			bot.nick = nickComm[1]
		} else if strings.HasPrefix(message, "USER ") {
			nickComm := strings.Split(message, "USER ")
			bot.nick = nickComm[1]
		} else {
			bot.nick = "justinfan123321"
		}
		if bot.oauth != "" {
			bot.CreateConnection()
			go bot.CreateConnection()
			go bot.CreateConnection()
			go bot.CreateConnection()
			go bot.CreateConnection()
		}
	} else if strings.HasPrefix(message, "PRIVMSG #jtv :/w ") {
		privmsgComm := strings.Split(message, "PRIVMSG #jtv :")
		go bot.Whisper(privmsgComm[1])
	} else {
		bot.Message(message)
	}
}
<commit_msg>removed unreachable code<commit_after>package main

import (
	"bufio"
	"fmt"
	"log"
	"net"
	"net/textproto"
	"strings"
)

var botlist []*Bot

// TCPServer simple tcp server for commands
func TCPServer() (ret int) {
	ln, err := net.Listen("tcp", ":"+TCPPort)
	if err != nil {
		log.Println("[control] error listening:", err.Error())
		return 1
	}
	log.Printf("[control] listening to port %s for connections...", TCPPort)
	defer ln.Close()

	for {
		conn, err := ln.Accept()
		bot := NewBot()
		bot.inconn = conn
		botlist = append(botlist, bot)
		if err != nil {
			fmt.Println("[control] error accepting: ", err.Error())
			return 1
		}
		go handleRequest(conn, bot)
	}
}

// CloseBot closes all bots
func CloseBot(bot *Bot) {
	// Iterate over the list of bots
	for i := range botlist {
		if bot == botlist[i] {
			// Remove the closed bot from the list
			botlist = append(botlist[:i], botlist[i+1:]...)
			return
		}
	}

	// Let the bot clean itself up
	bot.Close()
}

func handleRequest(conn net.Conn, bot *Bot) {
	defer CloseBot(bot)
	defer bot.Close()

	reader := bufio.NewReader(conn)
	tp := textproto.NewReader(reader)
	for {
		line, err := tp.ReadLine()

		if err != nil {
			fmt.Println("[control] read error:", err)
			fmt.Fprintf(conn, "[control] read error: %s", err)
			conn.Close()
			return
		}
		if !strings.HasPrefix(line, "PASS ") {
			log.Println("[control] " + line)
		}
		handleMessage(line, bot)
	}
}

// Handle an IRC received from a bot
func handleMessage(message string, bot *Bot) {
	remoteAddr := bot.inconn.RemoteAddr().String()
	remoteAddrIP := strings.Split(remoteAddr, ":")

	if strings.HasPrefix(message, "JOIN ") {
		joinComm := strings.Split(message, "JOIN ")
		go bot.Join(joinComm[1])
	} else if strings.HasPrefix(message, "PART ") {
		partComm := strings.Split(message, "PART ")
		go bot.Part(partComm[1])
	} else if strings.HasPrefix(message, "PASS ") {
		passComm := strings.Split(message, "PASS ")
		passwordParts := strings.Split(passComm[1], ";")
		if passwordParts[0] == TCPPass {
			bot.oauth = passwordParts[1]
			log.Printf("[control] authenticated! %s\n", remoteAddrIP)
		} else {
			log.Printf("[control] invalid broker pass! %s\n", remoteAddrIP)
			bot.inconn.Close()
			return
		}
	} else if strings.HasPrefix(message, "NICK ") || strings.HasPrefix(message, "USER ") {
		if strings.HasPrefix(message, "NICK ") {
			nickComm := strings.Split(message, "NICK ")
			bot.nick = nickComm[1]
		} else if strings.HasPrefix(message, "USER ") {
			nickComm := strings.Split(message, "USER ")
			bot.nick = nickComm[1]
		}

		if bot.oauth != "" {
			bot.CreateConnection()
			go bot.CreateConnection()
			go bot.CreateConnection()
			go bot.CreateConnection()
			go bot.CreateConnection()
		}
	} else if strings.HasPrefix(message, "PRIVMSG #jtv :/w ") {
		privmsgComm := strings.Split(message, "PRIVMSG #jtv :")
		go bot.Whisper(privmsgComm[1])
	} else {
		bot.Message(message)
	}
}
<|endoftext|>


<commit_before>package gocql

import (
	crand "crypto/rand"
	"errors"
	"fmt"
	"golang.org/x/net/context"
	"log"
	"math/rand"
	"net"
	"strconv"
	"sync/atomic"
	"time"
)

var (
	randr *rand.Rand
)

func init() {
	b := make([]byte, 4)
	if _, err := crand.Read(b); err != nil {
		panic(fmt.Sprintf("unable to seed random number generator: %v", err))
	}

	randr = rand.New(rand.NewSource(int64(readInt(b))))
}

// Ensure that the atomic variable is aligned to a 64bit boundary
// so that atomic operations can be applied on 32bit architectures.
type controlConn struct {
	session *Session
	conn    atomic.Value

	retry RetryPolicy

	started int32
	quit    chan struct{}
}

func createControlConn(session *Session) *controlConn {
	control := &controlConn{
		session: session,
		quit:    make(chan struct{}),
		retry:   &SimpleRetryPolicy{NumRetries: 3},
	}

	control.conn.Store((*Conn)(nil))

	return control
}

func (c *controlConn) heartBeat() {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return
	}

	sleepTime := 1 * time.Second

	for {
		select {
		case <-c.quit:
			return
		case <-time.After(sleepTime):
		}

		resp, err := c.writeFrame(&writeOptionsFrame{})
		if err != nil {
			goto reconn
		}

		switch resp.(type) {
		case *supportedFrame:
			// Everything ok
			sleepTime = 5 * time.Second
			continue
		case error:
			goto reconn
		default:
			panic(fmt.Sprintf("gocql: unknown frame in response to options: %T", resp))
		}

	reconn:
		// try to connect a bit faster
		sleepTime = 1 * time.Second
		c.reconnect(true)
		// time.Sleep(5 * time.Second)
		continue
	}
}

func (c *controlConn) shuffleDial(endpoints []string) (conn *Conn, err error) {
	perm := randr.Perm(len(endpoints))
	shuffled := make([]string, len(endpoints))

	for i, endpoint := range endpoints {
		shuffled[perm[i]] = endpoint
	}

	// shuffle endpoints so not all drivers will connect to the same initial
	// node.
	for _, addr := range shuffled {
		if addr == "" {
			return nil, fmt.Errorf("control: invalid address: %q", addr)
		}

		port := c.session.cfg.Port
		addr = JoinHostPort(addr, port)
		host, portStr, err := net.SplitHostPort(addr)
		if err != nil {
			host = addr
			port = c.session.cfg.Port
			err = nil
		} else {
			port, err = strconv.Atoi(portStr)
			if err != nil {
				return nil, err
			}
		}

		hostInfo, _ := c.session.ring.addHostIfMissing(&HostInfo{peer: host, port: port})
		conn, err = c.session.connect(addr, c, hostInfo)
		if err == nil {
			return conn, err
		}

		log.Printf("gocql: unable to dial control conn %v: %v\n", addr, err)
	}

	return
}

func (c *controlConn) connect(endpoints []string) error {
	if len(endpoints) == 0 {
		return errors.New("control: no endpoints specified")
	}

	conn, err := c.shuffleDial(endpoints)
	if err != nil {
		return fmt.Errorf("control: unable to connect: %v", err)
	} else if conn == nil {
		return errors.New("control: unable to connect to initial endpoints")
	}

	if err := c.setupConn(conn); err != nil {
		conn.Close()
		return fmt.Errorf("control: unable to setup connection: %v", err)
	}

	// we could fetch the initial ring here and update initial host data. So that
	// when we return from here we have a ring topology ready to go.

	go c.heartBeat()

	return nil
}

func (c *controlConn) setupConn(conn *Conn) error {
	if err := c.registerEvents(conn); err != nil {
		conn.Close()
		return err
	}

	c.conn.Store(conn)

	host, portstr, err := net.SplitHostPort(conn.conn.RemoteAddr().String())
	if err != nil {
		return err
	}
	port, err := strconv.Atoi(portstr)
	if err != nil {
		return err
	}

	c.session.handleNodeUp(net.ParseIP(host), port, false)

	return nil
}

func (c *controlConn) registerEvents(conn *Conn) error {
	var events []string

	if !c.session.cfg.Events.DisableTopologyEvents {
		events = append(events, "TOPOLOGY_CHANGE")
	}
	if !c.session.cfg.Events.DisableNodeStatusEvents {
		events = append(events, "STATUS_CHANGE")
	}
	if !c.session.cfg.Events.DisableSchemaEvents {
		events = append(events, "SCHEMA_CHANGE")
	}

	if len(events) == 0 {
		return nil
	}

	framer, err := conn.exec(context.Background(),
		&writeRegisterFrame{
			events: events,
		}, nil)
	if err != nil {
		return err
	}

	frame, err := framer.parseFrame()
	if err != nil {
		return err
	} else if _, ok := frame.(*readyFrame); !ok {
		return fmt.Errorf("unexpected frame in response to register: got %T: %v\n", frame, frame)
	}

	return nil
}

func (c *controlConn) reconnect(refreshring bool) {
	// TODO: simplify this function, use session.ring to get hosts instead of the
	// connection pool

	addr := c.addr()
	oldConn := c.conn.Load().(*Conn)
	if oldConn != nil {
		oldConn.Close()
	}

	var newConn *Conn
	if addr != "" {
		// try to connect to the old host
		conn, err := c.session.connect(addr, c, oldConn.host)
		if err != nil {
			// host is dead
			// TODO: this is replicated in a few places
			ip, portStr, _ := net.SplitHostPort(addr)
			port, _ := strconv.Atoi(portStr)
			c.session.handleNodeDown(net.ParseIP(ip), port)
		} else {
			newConn = conn
		}
	}

	// TODO: should have our own roundrobbin for hosts so that we can try each
	// in succession and guantee that we get a different host each time.
	if newConn == nil {
		host := c.session.ring.rrHost()
		if host == nil {
			c.connect(c.session.ring.endpoints)
			return
		}

		var err error
		newConn, err = c.session.connect(host.Peer(), c, host)
		if err != nil {
			// TODO: add log handler for things like this
			return
		}
	}

	if err := c.setupConn(newConn); err != nil {
		newConn.Close()
		log.Printf("gocql: control unable to register events: %v\n", err)
		return
	}

	if refreshring {
		c.session.hostSource.refreshRing()
	}
}

func (c *controlConn) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		return
	}

	oldConn := c.conn.Load().(*Conn)
	if oldConn != conn {
		return
	}

	c.reconnect(true)
}

func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
	conn := c.conn.Load().(*Conn)
	if conn == nil {
		return nil, errNoControl
	}

	framer, err := conn.exec(context.Background(), w, nil)
	if err != nil {
		return nil, err
	}

	return framer.parseFrame()
}

func (c *controlConn) withConn(fn func(*Conn) *Iter) *Iter {
	const maxConnectAttempts = 5
	connectAttempts := 0

	for i := 0; i < maxConnectAttempts; i++ {
		conn := c.conn.Load().(*Conn)
		if conn == nil {
			if connectAttempts > maxConnectAttempts {
				break
			}

			connectAttempts++

			c.reconnect(false)
			continue
		}

		return fn(conn)
	}

	return &Iter{err: errNoControl}
}

// query will return nil if the connection is closed or nil
func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter) {
	q := c.session.Query(statement, values...).Consistency(One).RoutingKey([]byte{})

	for {
		iter = c.withConn(func(conn *Conn) *Iter {
			return conn.executeQuery(q)
		})

		if gocqlDebug && iter.err != nil {
			log.Printf("control: error executing %q: %v\n", statement, iter.err)
		}

		q.attempts++
		if iter.err == nil || !c.retry.Attempt(q) {
			break
		}
	}

	return
}

func (c *controlConn) fetchHostInfo(addr net.IP, port int) (*HostInfo, error) {
	// TODO(zariel): we should probably move this into host_source or atleast
	// share code with it.
	hostname, _, err := net.SplitHostPort(c.addr())
	if err != nil {
		return nil, fmt.Errorf("unable to fetch host info, invalid conn addr: %q: %v", c.addr(), err)
	}

	isLocal := hostname == addr.String()

	var fn func(*HostInfo) error

	if isLocal {
		fn = func(host *HostInfo) error {
			// TODO(zariel): should we fetch rpc_address from here?
			iter := c.query("SELECT data_center, rack, host_id, tokens, release_version FROM system.local WHERE key='local'")
			iter.Scan(&host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version)
			return iter.Close()
		}
	} else {
		fn = func(host *HostInfo) error {
			// TODO(zariel): should we fetch rpc_address from here?
			iter := c.query("SELECT data_center, rack, host_id, tokens, release_version FROM system.peers WHERE peer=?", addr)
			iter.Scan(&host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version)
			return iter.Close()
		}
	}

	host := &HostInfo{
		port: port,
	}

	if err := fn(host); err != nil {
		return nil, err
	}
	host.peer = addr.String()

	return host, nil
}

func (c *controlConn) awaitSchemaAgreement() error {
	return c.withConn(func(conn *Conn) *Iter {
		return &Iter{err: conn.awaitSchemaAgreement()}
	}).err
}

func (c *controlConn) addr() string {
	conn := c.conn.Load().(*Conn)
	if conn == nil {
		return ""
	}
	return conn.addr
}

func (c *controlConn) close() {
	if atomic.CompareAndSwapInt32(&c.started, 1, -1) {
		c.quit <- struct{}{}
	}
	conn := c.conn.Load().(*Conn)
	if conn != nil {
		conn.Close()
	}
}

var errNoControl = errors.New("gocql: no control connection available")
<commit_msg>control: fix shaddowing error in shuffle dial<commit_after>package gocql

import (
	crand "crypto/rand"
	"errors"
	"fmt"
	"golang.org/x/net/context"
	"log"
	"math/rand"
	"net"
	"strconv"
	"sync/atomic"
	"time"
)

var (
	randr *rand.Rand
)

func init() {
	b := make([]byte, 4)
	if _, err := crand.Read(b); err != nil {
		panic(fmt.Sprintf("unable to seed random number generator: %v", err))
	}

	randr = rand.New(rand.NewSource(int64(readInt(b))))
}

// Ensure that the atomic variable is aligned to a 64bit boundary
// so that atomic operations can be applied on 32bit architectures.
type controlConn struct {
	session *Session
	conn    atomic.Value

	retry RetryPolicy

	started int32
	quit    chan struct{}
}

func createControlConn(session *Session) *controlConn {
	control := &controlConn{
		session: session,
		quit:    make(chan struct{}),
		retry:   &SimpleRetryPolicy{NumRetries: 3},
	}

	control.conn.Store((*Conn)(nil))

	return control
}

func (c *controlConn) heartBeat() {
	if !atomic.CompareAndSwapInt32(&c.started, 0, 1) {
		return
	}

	sleepTime := 1 * time.Second

	for {
		select {
		case <-c.quit:
			return
		case <-time.After(sleepTime):
		}

		resp, err := c.writeFrame(&writeOptionsFrame{})
		if err != nil {
			goto reconn
		}

		switch resp.(type) {
		case *supportedFrame:
			// Everything ok
			sleepTime = 5 * time.Second
			continue
		case error:
			goto reconn
		default:
			panic(fmt.Sprintf("gocql: unknown frame in response to options: %T", resp))
		}

	reconn:
		// try to connect a bit faster
		sleepTime = 1 * time.Second
		c.reconnect(true)
		// time.Sleep(5 * time.Second)
		continue
	}
}

func hostInfo(addr string, defaultPort int) (*HostInfo, error) {
	var port int
	host, portStr, err := net.SplitHostPort(addr)
	if err != nil {
		host = addr
		port = defaultPort
	} else {
		port, err = strconv.Atoi(portStr)
		if err != nil {
			return nil, err
		}
	}

	return &HostInfo{peer: host, port: port}, nil
}

func (c *controlConn) shuffleDial(endpoints []string) (conn *Conn, err error) {
	perm := randr.Perm(len(endpoints))
	shuffled := make([]string, len(endpoints))

	for i, endpoint := range endpoints {
		shuffled[perm[i]] = endpoint
	}

	// shuffle endpoints so not all drivers will connect to the same initial
	// node.
	for _, addr := range shuffled {
		if addr == "" {
			return nil, fmt.Errorf("invalid address: %q", addr)
		}

		port := c.session.cfg.Port
		addr = JoinHostPort(addr, port)

		var host *HostInfo
		host, err = hostInfo(addr, port)
		if err != nil {
			return nil, fmt.Errorf("invalid address: %q: %v", addr, err)
		}

		hostInfo, _ := c.session.ring.addHostIfMissing(host)
		conn, err = c.session.connect(addr, c, hostInfo)
		if err == nil {
			return conn, err
		}

		log.Printf("gocql: unable to dial control conn %v: %v\n", addr, err)
	}

	if err != nil {
		return nil, err
	}

	return conn, nil
}

func (c *controlConn) connect(endpoints []string) error {
	if len(endpoints) == 0 {
		return errors.New("control: no endpoints specified")
	}

	conn, err := c.shuffleDial(endpoints)
	if err != nil {
		return fmt.Errorf("control: unable to connect to initial hosts: %v", err)
	}

	if err := c.setupConn(conn); err != nil {
		conn.Close()
		return fmt.Errorf("control: unable to setup connection: %v", err)
	}

	// we could fetch the initial ring here and update initial host data. So that
	// when we return from here we have a ring topology ready to go.

	go c.heartBeat()

	return nil
}

func (c *controlConn) setupConn(conn *Conn) error {
	if err := c.registerEvents(conn); err != nil {
		conn.Close()
		return err
	}

	c.conn.Store(conn)

	host, portstr, err := net.SplitHostPort(conn.conn.RemoteAddr().String())
	if err != nil {
		return err
	}
	port, err := strconv.Atoi(portstr)
	if err != nil {
		return err
	}

	c.session.handleNodeUp(net.ParseIP(host), port, false)

	return nil
}

func (c *controlConn) registerEvents(conn *Conn) error {
	var events []string

	if !c.session.cfg.Events.DisableTopologyEvents {
		events = append(events, "TOPOLOGY_CHANGE")
	}
	if !c.session.cfg.Events.DisableNodeStatusEvents {
		events = append(events, "STATUS_CHANGE")
	}
	if !c.session.cfg.Events.DisableSchemaEvents {
		events = append(events, "SCHEMA_CHANGE")
	}

	if len(events) == 0 {
		return nil
	}

	framer, err := conn.exec(context.Background(),
		&writeRegisterFrame{
			events: events,
		}, nil)
	if err != nil {
		return err
	}

	frame, err := framer.parseFrame()
	if err != nil {
		return err
	} else if _, ok := frame.(*readyFrame); !ok {
		return fmt.Errorf("unexpected frame in response to register: got %T: %v\n", frame, frame)
	}

	return nil
}

func (c *controlConn) reconnect(refreshring bool) {
	// TODO: simplify this function, use session.ring to get hosts instead of the
	// connection pool

	addr := c.addr()
	oldConn := c.conn.Load().(*Conn)
	if oldConn != nil {
		oldConn.Close()
	}

	var newConn *Conn
	if addr != "" {
		// try to connect to the old host
		conn, err := c.session.connect(addr, c, oldConn.host)
		if err != nil {
			// host is dead
			// TODO: this is replicated in a few places
			ip, portStr, _ := net.SplitHostPort(addr)
			port, _ := strconv.Atoi(portStr)
			c.session.handleNodeDown(net.ParseIP(ip), port)
		} else {
			newConn = conn
		}
	}

	// TODO: should have our own roundrobbin for hosts so that we can try each
	// in succession and guantee that we get a different host each time.
	if newConn == nil {
		host := c.session.ring.rrHost()
		if host == nil {
			c.connect(c.session.ring.endpoints)
			return
		}

		var err error
		newConn, err = c.session.connect(host.Peer(), c, host)
		if err != nil {
			// TODO: add log handler for things like this
			return
		}
	}

	if err := c.setupConn(newConn); err != nil {
		newConn.Close()
		log.Printf("gocql: control unable to register events: %v\n", err)
		return
	}

	if refreshring {
		c.session.hostSource.refreshRing()
	}
}

func (c *controlConn) HandleError(conn *Conn, err error, closed bool) {
	if !closed {
		return
	}

	oldConn := c.conn.Load().(*Conn)
	if oldConn != conn {
		return
	}

	c.reconnect(true)
}

func (c *controlConn) writeFrame(w frameWriter) (frame, error) {
	conn := c.conn.Load().(*Conn)
	if conn == nil {
		return nil, errNoControl
	}

	framer, err := conn.exec(context.Background(), w, nil)
	if err != nil {
		return nil, err
	}

	return framer.parseFrame()
}

func (c *controlConn) withConn(fn func(*Conn) *Iter) *Iter {
	const maxConnectAttempts = 5
	connectAttempts := 0

	for i := 0; i < maxConnectAttempts; i++ {
		conn := c.conn.Load().(*Conn)
		if conn == nil {
			if connectAttempts > maxConnectAttempts {
				break
			}

			connectAttempts++

			c.reconnect(false)
			continue
		}

		return fn(conn)
	}

	return &Iter{err: errNoControl}
}

// query will return nil if the connection is closed or nil
func (c *controlConn) query(statement string, values ...interface{}) (iter *Iter) {
	q := c.session.Query(statement, values...).Consistency(One).RoutingKey([]byte{})

	for {
		iter = c.withConn(func(conn *Conn) *Iter {
			return conn.executeQuery(q)
		})

		if gocqlDebug && iter.err != nil {
			log.Printf("control: error executing %q: %v\n", statement, iter.err)
		}

		q.attempts++
		if iter.err == nil || !c.retry.Attempt(q) {
			break
		}
	}

	return
}

func (c *controlConn) fetchHostInfo(addr net.IP, port int) (*HostInfo, error) {
	// TODO(zariel): we should probably move this into host_source or atleast
	// share code with it.
	hostname, _, err := net.SplitHostPort(c.addr())
	if err != nil {
		return nil, fmt.Errorf("unable to fetch host info, invalid conn addr: %q: %v", c.addr(), err)
	}

	isLocal := hostname == addr.String()

	var fn func(*HostInfo) error

	if isLocal {
		fn = func(host *HostInfo) error {
			// TODO(zariel): should we fetch rpc_address from here?
			iter := c.query("SELECT data_center, rack, host_id, tokens, release_version FROM system.local WHERE key='local'")
			iter.Scan(&host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version)
			return iter.Close()
		}
	} else {
		fn = func(host *HostInfo) error {
			// TODO(zariel): should we fetch rpc_address from here?
			iter := c.query("SELECT data_center, rack, host_id, tokens, release_version FROM system.peers WHERE peer=?", addr)
			iter.Scan(&host.dataCenter, &host.rack, &host.hostId, &host.tokens, &host.version)
			return iter.Close()
		}
	}

	host := &HostInfo{
		port: port,
	}

	if err := fn(host); err != nil {
		return nil, err
	}
	host.peer = addr.String()

	return host, nil
}

func (c *controlConn) awaitSchemaAgreement() error {
	return c.withConn(func(conn *Conn) *Iter {
		return &Iter{err: conn.awaitSchemaAgreement()}
	}).err
}

func (c *controlConn) addr() string {
	conn := c.conn.Load().(*Conn)
	if conn == nil {
		return ""
	}
	return conn.addr
}

func (c *controlConn) close() {
	if atomic.CompareAndSwapInt32(&c.started, 1, -1) {
		c.quit <- struct{}{}
	}
	conn := c.conn.Load().(*Conn)
	if conn != nil {
		conn.Close()
	}
}

var errNoControl = errors.New("gocql: no control connection available")
<|endoftext|>


<commit_before>package main

import (
	"container/list"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/divan/gotrace/trace"
)

func ConvertEvents(events []*trace.Event) (Commands, error) {
	var c Commands

	sends := list.New()

	debug := os.Getenv("GOTRACE_DEBUG") == "1"
	var lastG uint64
	for _, ev := range events {
		switch ev.Type {
		case trace.EvGoStart:
			fmt.Println(" ---> GoStart:", ev.G, "from", lastG)
			if ev.G == 1 && lastG == 0 {
				c.StartGoroutine(ev.Ts, ev.Stk[0].Fn, ev.G, lastG)
			}
			lastG = ev.G
		case trace.EvGoCreate:
			if len(ev.Stk) > 0 {
				if strings.HasPrefix(ev.Stk[0].Fn, "runtime") {
					if ev.Stk[0].Fn != "runtime.main" {
						break
					}
				}
				c.StartGoroutine(ev.Ts, ev.Stk[0].Fn, ev.Args[0], ev.G)
				if debug {
					fmt.Println(" ---> Create:", ev.Args[0], "from", ev.G)
				}
			}
		case trace.EvGoUnblock:
			lastG = ev.Args[0]
			fmt.Println("Go: Unblock: set lastG to", lastG)
		case trace.EvGoEnd:
			c.StopGoroutine(ev.Ts, "", ev.G)
			lastG = ev.G
			if debug {
				fmt.Println("End:", ev.G)
			}
		case trace.EvGoSend:
			sends.PushBack(ev)
			if debug {
				fmt.Printf("[DD] %d, Send: G:%d, CH: %d, EvID: %d, Val:%d\n", ev.Ts, ev.G, ev.Args[1], ev.Args[0], ev.Args[2])
			}
		case trace.EvGoRecv:
			if debug {
				fmt.Printf("[DD] %d, Recv: G:%d, CH: %d, EvID: %d, Val:%d - %v\n", ev.Ts, ev.G, ev.Args[1], ev.Args[0], ev.Args[2], ev)
			}
			send := findSource(sends, ev)
			if send == nil {
				// it's either channel close() or error in findSource
				continue
			}
			if debug {
				fmt.Printf("[DD] %d, Recv->Send: FromG:%d, ToG: %d, CH: %d, EvID: %d, Val:%d\n", send.Ts, send.G, ev.G, ev.Args[1], ev.Args[0], ev.Args[2])
			}
			c.ChanSend(send.Ts, ev.Args[1], send.G, ev.G, send.Args[2])
		case trace.EvGCStart, trace.EvGCDone, trace.EvGCScanStart, trace.EvGCScanDone:
			lastG = 1
			/*
				case trace.EvGoSched, trace.EvGoPreempt,
					trace.EvGoSleep, trace.EvGoBlock, trace.EvGoBlockSelect, trace.EvGoBlockSend, trace.EvGoBlockRecv,
					trace.EvGoBlockSync, trace.EvGoBlockCond, trace.EvGoBlockNet,
					trace.EvGoSysBlock:
					fmt.Println("Ev:", ev.Type, ev.G, ev.Args)
					//lastG = ev.G
			*/
		case trace.EvGoStop:
			if debug {
				fmt.Println("Stop:", ev.G)
			}
			lastG = 1
		default:
			if debug {
				fmt.Println("Ev:", ev.Type, ev.G, ev.Args)
			}
		}
	}

	// sort events
	sort.Sort(ByTimestamp(c))

	// insert stop main
	// TODO: figure out why it's not in the trace
	lastTs := c[len(c)-1].Time
	c.StopGoroutine(lastTs+1000, "", 1)

	return c, nil
}

// findSource tries to find corresponding Send event to ev.
func findSource(sends *list.List, ev *trace.Event) *trace.Event {
	for e := sends.Back(); e != nil; e = e.Prev() {
		send := e.Value.(*trace.Event)
		if send.Args[1] == ev.Args[1] && send.Args[0] == ev.Args[0] {
			sends.Remove(e)
			return send
		}
	}
	return nil
}
<commit_msg>Improve logging<commit_after>package main

import (
	"container/list"
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/divan/gotrace/trace"
)

func ConvertEvents(events []*trace.Event) (Commands, error) {
	var c Commands

	sends := list.New()

	debug := os.Getenv("GOTRACE_DEBUG") == "1"
	var lastG uint64
	for _, ev := range events {
		switch ev.Type {
		case trace.EvGoStart:
			if debug {
				fmt.Println(ev.Ts, "GoStart:", ev.G, "from", lastG, ev.Args)
			}
			if ev.G == 1 && lastG == 0 {
				c.StartGoroutine(ev.Ts, ev.Stk[0].Fn, ev.G, lastG)
			}
			lastG = ev.G
		case trace.EvGoCreate:
			if len(ev.Stk) > 0 {
				if strings.HasPrefix(ev.Stk[0].Fn, "runtime") {
					if ev.Stk[0].Fn != "runtime.main" {
						break
					}
				}
				c.StartGoroutine(ev.Ts, ev.Stk[0].Fn, ev.Args[0], ev.G)
				if debug {
					fmt.Println(ev.Ts, "GoCreate:", ev.Args[0], "from", ev.G)
				}
			}
		case trace.EvGoUnblock:
			lastG = ev.Args[0]
			if debug {
				fmt.Println(ev.Ts, "GoUnblock: set lastG to", lastG, ev.Args)
			}
		case trace.EvGoEnd:
			if debug {
				fmt.Println(ev.Ts, "GoEnd:", ev.G)
			}
			c.StopGoroutine(ev.Ts, "", ev.G)
			lastG = ev.G
		case trace.EvGoSend:
			if debug {
				fmt.Printf("[DD] %d, Send: G:%d, CH: %d, EvID: %d, Val:%d\n", ev.Ts, ev.G, ev.Args[1], ev.Args[0], ev.Args[2])
			}
			sends.PushBack(ev)
		case trace.EvGoRecv:
			if debug {
				fmt.Printf("[DD] %d, Recv: G:%d, CH: %d, EvID: %d, Val:%d - %v\n", ev.Ts, ev.G, ev.Args[1], ev.Args[0], ev.Args[2], ev)
			}
			send := findSource(sends, ev)
			if send == nil {
				// it's either channel close() or error in findSource
				continue
			}
			if debug {
				fmt.Printf("[DD] %d, Recv->Send: FromG:%d, ToG: %d, CH: %d, EvID: %d, Val:%d\n", send.Ts, send.G, ev.G, ev.Args[1], ev.Args[0], ev.Args[2])
			}
			c.ChanSend(send.Ts, ev.Args[1], send.G, ev.G, send.Args[2])
		case trace.EvGCStart, trace.EvGCDone, trace.EvGCScanStart, trace.EvGCScanDone:
			if debug {
				fmt.Println(ev.Ts, "GoGC...", ev.Type, ev.Args)
			}
			lastG = 1
		case trace.EvGoBlockSelect:
			if debug {
				fmt.Println("[DD] BlockSelect:", ev.Ts, ev.G, ev.Args)
			}
		case trace.EvGoBlockSend:
			if debug {
				fmt.Println("[DD] BlockSend:", ev.Ts, ev.G, ev.Args)
			}
		case trace.EvGoBlockRecv:
			if debug {
				fmt.Println("[DD] BlockRecv:", ev.Ts, ev.G, ev.Args)
			}
			/*
				case trace.EvGoSched, trace.EvGoPreempt,
					trace.EvGoSleep, trace.EvGoBlock, trace.EvGoBlockSelect, trace.EvGoBlockSend, trace.EvGoBlockRecv,
					trace.EvGoBlockSync, trace.EvGoBlockCond, trace.EvGoBlockNet,
					trace.EvGoSysBlock:
					fmt.Println("Ev:", ev.Type, ev.G, ev.Args)
					//lastG = ev.G
			*/
		case trace.EvGoStop:
			if debug {
				fmt.Println(ev.Ts, "GoStop:", ev.G)
			}
			lastG = 1
		default:
			if debug {
				fmt.Println(ev.Ts, "Ev:", ev.Type, ev.G, ev.Args)
			}
		}
	}

	// sort events
	sort.Sort(ByTimestamp(c))

	// insert stop main
	// TODO: figure out why it's not in the trace
	lastTs := c[len(c)-1].Time
	c.StopGoroutine(lastTs+1000, "", 1)

	return c, nil
}

// findSource tries to find corresponding Send event to ev.
func findSource(sends *list.List, ev *trace.Event) *trace.Event {
	for e := sends.Back(); e != nil; e = e.Prev() {
		send := e.Value.(*trace.Event)
		if send.Args[1] == ev.Args[1] && send.Args[0] == ev.Args[0] {
			sends.Remove(e)
			return send
		}
	}
	return nil
}
<|endoftext|>


<commit_before>package gocrawl

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"path"

	"github.com/PuerkitoBio/goquery"
	"github.com/temoto/robotstxt"
	"golang.org/x/net/html"
)

// The worker is dedicated to fetching and visiting a given host, respecting
// this host's robots.txt crawling policies.
type worker struct {
	// Worker identification
	host  string
	index int

	// Communication channels and sync
	push    chan<- *workerResponse
	pop     popChannel
	stop    chan struct{}
	enqueue chan<- interface{}
	wg      *sync.WaitGroup

	// Robots validation
	robotsGroup *robotstxt.Group

	// Logging
	logFunc func(LogFlags, string, ...interface{})

	// Implementation fields
	wait           <-chan time.Time
	lastFetch      *FetchInfo
	lastCrawlDelay time.Duration
	opts           *Options
}

// Start crawling the host.
func (w *worker) run() {
	defer func() {
		w.logFunc(LogInfo, "worker done.")
		w.wg.Done()
	}()

	// Enter loop to process URLs until stop signal is received
	for {
		var idleChan <-chan time.Time

		w.logFunc(LogInfo, "waiting for pop...")

		// Initialize the idle timeout channel, if required
		if w.opts.WorkerIdleTTL > 0 {
			idleChan = time.After(w.opts.WorkerIdleTTL)
		}

		select {
		case <-w.stop:
			w.logFunc(LogInfo, "stop signal received.")
			return

		case <-idleChan:
			w.logFunc(LogInfo, "idle timeout received.")
			w.sendResponse(nil, false, nil, true)
			return

		case batch := <-w.pop:

			// Got a batch of urls to crawl, loop and check at each iteration if a stop
			// is received.
			for _, ctx := range batch {
				w.logFunc(LogInfo, "popped: %s", ctx.url)

				if ctx.IsRobotsURL() {
					w.requestRobotsTxt(ctx)
				} else if w.isAllowedPerRobotsPolicies(ctx.url) {
					w.requestURL(ctx, ctx.HeadBeforeGet)
				} else {
					// Must still notify Crawler that this URL was processed, although not visited
					w.opts.Extender.Disallowed(ctx)
					w.sendResponse(ctx, false, nil, false)
				}

				// No need to check for idle timeout here, no idling while looping through
				// a batch of URLs.
				select {
				case <-w.stop:
					w.logFunc(LogInfo, "stop signal received.")
					return
				default:
					// Nothing, just continue...
				}
			}
		}
	}
}

// Checks if the given URL can be fetched based on robots.txt policies.
func (w *worker) isAllowedPerRobotsPolicies(u *url.URL) bool {
	if w.robotsGroup != nil {
		// Is this URL allowed per robots.txt policy?
		ok := w.robotsGroup.Test(u.Path)
		if !ok {
			w.logFunc(LogIgnored, "ignored on robots.txt policy: %s", u.String())
		}
		return ok
	}

	// No robots.txt = everything is allowed
	return true
}

// Process the specified URL.
func (w *worker) requestURL(ctx *URLContext, headRequest bool) {
	if res, ok := w.fetchURL(ctx, w.opts.UserAgent, headRequest); ok {
		var harvested interface{}
		var visited bool

		// Close the body on function end
		defer res.Body.Close()

		// Any 2xx status code is good to go
		if res.StatusCode >= 200 && res.StatusCode < 300 {
			// Success, visit the URL
			harvested = w.visitURL(ctx, res)
			visited = true
		} else {
			// Error based on status code received
			w.opts.Extender.Error(newCrawlErrorMessage(ctx, res.Status, CekHttpStatusCode))
			w.logFunc(LogError, "ERROR status code for %s: %s", ctx.url, res.Status)
		}
		w.sendResponse(ctx, visited, harvested, false)
	}
}

// Process the robots.txt URL.
func (w *worker) requestRobotsTxt(ctx *URLContext) {
	// Ask if it should be fetched
	if robData, reqRob := w.opts.Extender.RequestRobots(ctx, w.opts.RobotUserAgent); !reqRob {
		w.logFunc(LogInfo, "using robots.txt from cache")
		w.robotsGroup = w.getRobotsTxtGroup(ctx, robData, nil)

	} else if res, ok := w.fetchURL(ctx, w.opts.UserAgent, false); ok {
		// Close the body on function end
		defer res.Body.Close()
		w.robotsGroup = w.getRobotsTxtGroup(ctx, nil, res)
	}
}

// Get the robots.txt group for this crawler.
func (w *worker) getRobotsTxtGroup(ctx *URLContext, b []byte, res *http.Response) (g *robotstxt.Group) {
	var data *robotstxt.RobotsData
	var e error

	if res != nil {
		var buf bytes.Buffer
		io.Copy(&buf, res.Body)
		res.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
		data, e = robotstxt.FromResponse(res)
		// Rewind the res.Body (by re-creating it from the bytes)
		res.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
		// Error or not, the robots.txt has been fetched, so notify
		w.opts.Extender.FetchedRobots(ctx, res)
	} else {
		data, e = robotstxt.FromBytes(b)
	}

	// If robots data cannot be parsed, will return nil, which will allow access by default.
	// Reasonable, since by default no robots.txt means full access, so invalid
	// robots.txt is similar behavior.
	if e != nil {
		w.opts.Extender.Error(newCrawlError(nil, e, CekParseRobots))
		w.logFunc(LogError, "ERROR parsing robots.txt for host %s: %s", w.host, e)
	} else {
		g = data.FindGroup(w.opts.RobotUserAgent)
	}
	return g
}

// Set the crawl delay between this request and the next.
func (w *worker) setCrawlDelay() {
	var robDelay time.Duration

	if w.robotsGroup != nil {
		robDelay = w.robotsGroup.CrawlDelay
	}
	w.lastCrawlDelay = w.opts.Extender.ComputeDelay(w.host,
		&DelayInfo{
			w.opts.CrawlDelay,
			robDelay,
			w.lastCrawlDelay,
		},
		w.lastFetch)
	w.logFunc(LogInfo, "using crawl-delay: %v", w.lastCrawlDelay)
}

// Request the specified URL and return the response.
func (w *worker) fetchURL(ctx *URLContext, agent string, headRequest bool) (res *http.Response, ok bool) {
	var e error
	var silent bool

	for {
		// Wait for crawl delay, if one is pending.
		w.logFunc(LogTrace, "waiting for crawl delay")
		if w.wait != nil {
			<-w.wait
			w.wait = nil
		}

		// Compute the next delay
		w.setCrawlDelay()

		// Compute the fetch duration
		now := time.Now()

		// Request the URL
		if res, e = w.opts.Extender.Fetch(ctx, agent, headRequest); e != nil {
			// Check if this is an ErrEnqueueRedirect, in which case we will enqueue
			// the redirect-to URL.
			if ue, ok := e.(*url.Error); ok {
				// We have a *url.Error, check if it was returned because of an ErrEnqueueRedirect
				if ue.Err == ErrEnqueueRedirect {
					// Do not notify this error outside of this if block, this is not a
					// "real" error. We either enqueue the new URL, or fail to parse it,
					// and then stop processing the current URL.
					silent = true
					// Parse the URL in the context of the original URL (so that relative URLs are ok).
					// Absolute URLs that point to another host are ok too.
					if ur, e := ctx.url.Parse(ue.URL); e != nil {
						// Notify error
						w.opts.Extender.Error(newCrawlError(nil, e, CekParseRedirectURL))
						w.logFunc(LogError, "ERROR parsing redirect URL %s: %s", ue.URL, e)
					} else {
						w.logFunc(LogTrace, "redirect to %s from %s, linked from %s", ur, ctx.URL(), ctx.SourceURL())
						// Enqueue the redirect-to URL with the original source
						rCtx := ctx.cloneForRedirect(ur, w.opts.URLNormalizationFlags)
						w.enqueue <- rCtx
					}
				}
			}

			// No fetch, so set to nil
			w.lastFetch = nil

			if !silent {
				// Notify error
				w.opts.Extender.Error(newCrawlError(ctx, e, CekFetch))
				w.logFunc(LogError, "ERROR fetching %s: %s", ctx.url, e)
			}

			// Return from this URL crawl
			w.sendResponse(ctx, false, nil, false)
			return nil, false

		}
		// Get the fetch duration
		fetchDuration := time.Now().Sub(now)
		// Crawl delay starts now.
		w.wait = time.After(w.lastCrawlDelay)

		// Keep trace of this last fetch info
		w.lastFetch = &FetchInfo{
			ctx,
			fetchDuration,
			res.StatusCode,
			headRequest,
		}

		if headRequest {
			// Close the HEAD request's body
			defer res.Body.Close()
			// Next up is GET request, maybe
			headRequest = false
			// Ask caller if we should proceed with a GET
			if !w.opts.Extender.RequestGet(ctx, res) {
				w.logFunc(LogIgnored, "ignored on HEAD filter policy: %s", ctx.url)
				w.sendResponse(ctx, false, nil, false)
				ok = false
				break
			}
		} else {
			ok = true
			break
		}
	}
	return
}

// Send a response to the crawler.
func (w *worker) sendResponse(ctx *URLContext, visited bool, harvested interface{}, idleDeath bool) {
	// Push harvested urls back to crawler, even if empty (uses the channel communication
	// to decrement reference count of pending URLs)
	if ctx == nil || !isRobotsURL(ctx.url) {
		// If a stop signal has been received, ignore the response, since the push
		// channel may be full and could block indefinitely.
		select {
		case <-w.stop:
			w.logFunc(LogInfo, "ignoring send response, will stop.")
			return
		default:
			// Nothing, just continue...
		}

		// No stop signal, send the response
		res := &workerResponse{
			ctx,
			visited,
			harvested,
			w.host,
			idleDeath,
		}
		w.push <- res
	}
}

// Process the response for a URL.
func (w *worker) visitURL(ctx *URLContext, res *http.Response) interface{} {
	var doc *goquery.Document
	var harvested interface{}
	var doLinks bool

	// Load a goquery document and call the visitor function
	if bd, e := ioutil.ReadAll(res.Body); e != nil {
		w.opts.Extender.Error(newCrawlError(ctx, e, CekReadBody))
		w.logFunc(LogError, "ERROR reading body %s: %s", ctx.url, e)
	} else {
		if node, e := html.Parse(bytes.NewBuffer(bd)); e != nil {
			w.opts.Extender.Error(newCrawlError(ctx, e, CekParseBody))
			w.logFunc(LogError, "ERROR parsing %s: %s", ctx.url, e)
		} else {
			doc = goquery.NewDocumentFromNode(node)
			doc.Url = res.Request.URL
		}
		// Re-assign the body so it can be consumed by the visitor function
		res.Body = ioutil.NopCloser(bytes.NewBuffer(bd))
	}

	// Visit the document (with nil goquery doc if failed to load)
	if harvested, doLinks = w.opts.Extender.Visit(ctx, res, doc); doLinks {
		// Links were not processed by the visitor, so process links
		if doc != nil {
			harvested = w.processLinks(doc)
		} else {
			w.opts.Extender.Error(newCrawlErrorMessage(ctx, "No goquery document to process links.", CekProcessLinks))
			w.logFunc(LogError, "ERROR processing links %s", ctx.url)
		}
	}
	// Notify that this URL has been visited
	w.opts.Extender.Visited(ctx, harvested)

	return harvested
}

func handleBaseTag(root *url.URL, baseHref string, aHref string) string {
	resolvedBase, err := root.Parse(baseHref)
	if err != nil {
		return ""
	}

	parsedURL, err := url.Parse(aHref)
	if err != nil {
		return ""
	}
	// If a[href] starts with a /, it overrides the base[href]
	if parsedURL.Host == "" && !strings.HasPrefix(aHref, "/") {
		aHref = path.Join(resolvedBase.Path, aHref)
	}

	resolvedURL, err := resolvedBase.Parse(aHref)
	if err != nil {
		return ""
	}
	return resolvedURL.String()
}

// Scrape the document's content to gather all links
func (w *worker) processLinks(doc *goquery.Document) (result []*url.URL) {
	baseURL, _ := doc.Find("base[href]").Attr("href")
	urls := doc.Find("a[href]").Map(func(_ int, s *goquery.Selection) string {
		val, _ := s.Attr("href")
		if baseURL != "" {
			val = handleBaseTag(doc.Url, baseURL, val)
		}
		return val
	})
	for _, s := range urls {
		// If href starts with "#", then it points to this same exact URL, ignore (will fail to parse anyway)
		if len(s) > 0 && !strings.HasPrefix(s, "#") {
			if parsed, e := url.Parse(s); e == nil {
				parsed = doc.Url.ResolveReference(parsed)
				result = append(result, parsed)
			} else {
				w.logFunc(LogIgnored, "ignore on unparsable policy %s: %s", s, e.Error())
			}
		}
	}
	return
}
<commit_msg>Added precompiled cascadia matchers (#68)<commit_after>package gocrawl

import (
	"bytes"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"
	"time"

	"path"

	"github.com/PuerkitoBio/goquery"
	"github.com/andybalholm/cascadia"
	"github.com/temoto/robotstxt"
	"golang.org/x/net/html"
)

// The worker is dedicated to fetching and visiting a given host, respecting
// this host's robots.txt crawling policies.
type worker struct {
	// Worker identification
	host  string
	index int

	// Communication channels and sync
	push    chan<- *workerResponse
	pop     popChannel
	stop    chan struct{}
	enqueue chan<- interface{}
	wg      *sync.WaitGroup

	// Robots validation
	robotsGroup *robotstxt.Group

	// Logging
	logFunc func(LogFlags, string, ...interface{})

	// Implementation fields
	wait           <-chan time.Time
	lastFetch      *FetchInfo
	lastCrawlDelay time.Duration
	opts           *Options
}

// Start crawling the host.
func (w *worker) run() {
	defer func() {
		w.logFunc(LogInfo, "worker done.")
		w.wg.Done()
	}()

	// Enter loop to process URLs until stop signal is received
	for {
		var idleChan <-chan time.Time

		w.logFunc(LogInfo, "waiting for pop...")

		// Initialize the idle timeout channel, if required
		if w.opts.WorkerIdleTTL > 0 {
			idleChan = time.After(w.opts.WorkerIdleTTL)
		}

		select {
		case <-w.stop:
			w.logFunc(LogInfo, "stop signal received.")
			return

		case <-idleChan:
			w.logFunc(LogInfo, "idle timeout received.")
			w.sendResponse(nil, false, nil, true)
			return

		case batch := <-w.pop:

			// Got a batch of urls to crawl, loop and check at each iteration if a stop
			// is received.
			for _, ctx := range batch {
				w.logFunc(LogInfo, "popped: %s", ctx.url)

				if ctx.IsRobotsURL() {
					w.requestRobotsTxt(ctx)
				} else if w.isAllowedPerRobotsPolicies(ctx.url) {
					w.requestURL(ctx, ctx.HeadBeforeGet)
				} else {
					// Must still notify Crawler that this URL was processed, although not visited
					w.opts.Extender.Disallowed(ctx)
					w.sendResponse(ctx, false, nil, false)
				}

				// No need to check for idle timeout here, no idling while looping through
				// a batch of URLs.
				select {
				case <-w.stop:
					w.logFunc(LogInfo, "stop signal received.")
					return
				default:
					// Nothing, just continue...
				}
			}
		}
	}
}

// Checks if the given URL can be fetched based on robots.txt policies.
func (w *worker) isAllowedPerRobotsPolicies(u *url.URL) bool {
	if w.robotsGroup != nil {
		// Is this URL allowed per robots.txt policy?
		ok := w.robotsGroup.Test(u.Path)
		if !ok {
			w.logFunc(LogIgnored, "ignored on robots.txt policy: %s", u.String())
		}
		return ok
	}

	// No robots.txt = everything is allowed
	return true
}

// Process the specified URL.
func (w *worker) requestURL(ctx *URLContext, headRequest bool) {
	if res, ok := w.fetchURL(ctx, w.opts.UserAgent, headRequest); ok {
		var harvested interface{}
		var visited bool

		// Close the body on function end
		defer res.Body.Close()

		// Any 2xx status code is good to go
		if res.StatusCode >= 200 && res.StatusCode < 300 {
			// Success, visit the URL
			harvested = w.visitURL(ctx, res)
			visited = true
		} else {
			// Error based on status code received
			w.opts.Extender.Error(newCrawlErrorMessage(ctx, res.Status, CekHttpStatusCode))
			w.logFunc(LogError, "ERROR status code for %s: %s", ctx.url, res.Status)
		}
		w.sendResponse(ctx, visited, harvested, false)
	}
}

// Process the robots.txt URL.
func (w *worker) requestRobotsTxt(ctx *URLContext) {
	// Ask if it should be fetched
	if robData, reqRob := w.opts.Extender.RequestRobots(ctx, w.opts.RobotUserAgent); !reqRob {
		w.logFunc(LogInfo, "using robots.txt from cache")
		w.robotsGroup = w.getRobotsTxtGroup(ctx, robData, nil)

	} else if res, ok := w.fetchURL(ctx, w.opts.UserAgent, false); ok {
		// Close the body on function end
		defer res.Body.Close()
		w.robotsGroup = w.getRobotsTxtGroup(ctx, nil, res)
	}
}

// Get the robots.txt group for this crawler.
func (w *worker) getRobotsTxtGroup(ctx *URLContext, b []byte, res *http.Response) (g *robotstxt.Group) {
	var data *robotstxt.RobotsData
	var e error

	if res != nil {
		var buf bytes.Buffer
		io.Copy(&buf, res.Body)
		res.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
		data, e = robotstxt.FromResponse(res)
		// Rewind the res.Body (by re-creating it from the bytes)
		res.Body = ioutil.NopCloser(bytes.NewReader(buf.Bytes()))
		// Error or not, the robots.txt has been fetched, so notify
		w.opts.Extender.FetchedRobots(ctx, res)
	} else {
		data, e = robotstxt.FromBytes(b)
	}

	// If robots data cannot be parsed, will return nil, which will allow access by default.
	// Reasonable, since by default no robots.txt means full access, so invalid
	// robots.txt is similar behavior.
	if e != nil {
		w.opts.Extender.Error(newCrawlError(nil, e, CekParseRobots))
		w.logFunc(LogError, "ERROR parsing robots.txt for host %s: %s", w.host, e)
	} else {
		g = data.FindGroup(w.opts.RobotUserAgent)
	}
	return g
}

// Set the crawl delay between this request and the next.
func (w *worker) setCrawlDelay() {
	var robDelay time.Duration

	if w.robotsGroup != nil {
		robDelay = w.robotsGroup.CrawlDelay
	}
	w.lastCrawlDelay = w.opts.Extender.ComputeDelay(w.host,
		&DelayInfo{
			w.opts.CrawlDelay,
			robDelay,
			w.lastCrawlDelay,
		},
		w.lastFetch)
	w.logFunc(LogInfo, "using crawl-delay: %v", w.lastCrawlDelay)
}

// Request the specified URL and return the response.
func (w *worker) fetchURL(ctx *URLContext, agent string, headRequest bool) (res *http.Response, ok bool) {
	var e error
	var silent bool

	for {
		// Wait for crawl delay, if one is pending.
		w.logFunc(LogTrace, "waiting for crawl delay")
		if w.wait != nil {
			<-w.wait
			w.wait = nil
		}

		// Compute the next delay
		w.setCrawlDelay()

		// Compute the fetch duration
		now := time.Now()

		// Request the URL
		if res, e = w.opts.Extender.Fetch(ctx, agent, headRequest); e != nil {
			// Check if this is an ErrEnqueueRedirect, in which case we will enqueue
			// the redirect-to URL.
			if ue, ok := e.(*url.Error); ok {
				// We have a *url.Error, check if it was returned because of an ErrEnqueueRedirect
				if ue.Err == ErrEnqueueRedirect {
					// Do not notify this error outside of this if block, this is not a
					// "real" error. We either enqueue the new URL, or fail to parse it,
					// and then stop processing the current URL.
					silent = true
					// Parse the URL in the context of the original URL (so that relative URLs are ok).
					// Absolute URLs that point to another host are ok too.
					if ur, e := ctx.url.Parse(ue.URL); e != nil {
						// Notify error
						w.opts.Extender.Error(newCrawlError(nil, e, CekParseRedirectURL))
						w.logFunc(LogError, "ERROR parsing redirect URL %s: %s", ue.URL, e)
					} else {
						w.logFunc(LogTrace, "redirect to %s from %s, linked from %s", ur, ctx.URL(), ctx.SourceURL())
						// Enqueue the redirect-to URL with the original source
						rCtx := ctx.cloneForRedirect(ur, w.opts.URLNormalizationFlags)
						w.enqueue <- rCtx
					}
				}
			}

			// No fetch, so set to nil
			w.lastFetch = nil

			if !silent {
				// Notify error
				w.opts.Extender.Error(newCrawlError(ctx, e, CekFetch))
				w.logFunc(LogError, "ERROR fetching %s: %s", ctx.url, e)
			}

			// Return from this URL crawl
			w.sendResponse(ctx, false, nil, false)
			return nil, false

		}
		// Get the fetch duration
		fetchDuration := time.Now().Sub(now)
		// Crawl delay starts now.
		w.wait = time.After(w.lastCrawlDelay)

		// Keep trace of this last fetch info
		w.lastFetch = &FetchInfo{
			ctx,
			fetchDuration,
			res.StatusCode,
			headRequest,
		}

		if headRequest {
			// Close the HEAD request's body
			defer res.Body.Close()
			// Next up is GET request, maybe
			headRequest = false
			// Ask caller if we should proceed with a GET
			if !w.opts.Extender.RequestGet(ctx, res) {
				w.logFunc(LogIgnored, "ignored on HEAD filter policy: %s", ctx.url)
				w.sendResponse(ctx, false, nil, false)
				ok = false
				break
			}
		} else {
			ok = true
			break
		}
	}
	return
}

// Send a response to the crawler.
func (w *worker) sendResponse(ctx *URLContext, visited bool, harvested interface{}, idleDeath bool) {
	// Push harvested urls back to crawler, even if empty (uses the channel communication
	// to decrement reference count of pending URLs)
	if ctx == nil || !isRobotsURL(ctx.url) {
		// If a stop signal has been received, ignore the response, since the push
		// channel may be full and could block indefinitely.
		select {
		case <-w.stop:
			w.logFunc(LogInfo, "ignoring send response, will stop.")
			return
		default:
			// Nothing, just continue...
		}

		// No stop signal, send the response
		res := &workerResponse{
			ctx,
			visited,
			harvested,
			w.host,
			idleDeath,
		}
		w.push <- res
	}
}

// Process the response for a URL.
func (w *worker) visitURL(ctx *URLContext, res *http.Response) interface{} {
	var doc *goquery.Document
	var harvested interface{}
	var doLinks bool

	// Load a goquery document and call the visitor function
	if bd, e := ioutil.ReadAll(res.Body); e != nil {
		w.opts.Extender.Error(newCrawlError(ctx, e, CekReadBody))
		w.logFunc(LogError, "ERROR reading body %s: %s", ctx.url, e)
	} else {
		if node, e := html.Parse(bytes.NewBuffer(bd)); e != nil {
			w.opts.Extender.Error(newCrawlError(ctx, e, CekParseBody))
			w.logFunc(LogError, "ERROR parsing %s: %s", ctx.url, e)
		} else {
			doc = goquery.NewDocumentFromNode(node)
			doc.Url = res.Request.URL
		}
		// Re-assign the body so it can be consumed by the visitor function
		res.Body = ioutil.NopCloser(bytes.NewBuffer(bd))
	}

	// Visit the document (with nil goquery doc if failed to load)
	if harvested, doLinks = w.opts.Extender.Visit(ctx, res, doc); doLinks {
		// Links were not processed by the visitor, so process links
		if doc != nil {
			harvested = w.processLinks(doc)
		} else {
			w.opts.Extender.Error(newCrawlErrorMessage(ctx, "No goquery document to process links.", CekProcessLinks))
			w.logFunc(LogError, "ERROR processing links %s", ctx.url)
		}
	}
	// Notify that this URL has been visited
	w.opts.Extender.Visited(ctx, harvested)

	return harvested
}

func handleBaseTag(root *url.URL, baseHref string, aHref string) string {
	resolvedBase, err := root.Parse(baseHref)
	if err != nil {
		return ""
	}

	parsedURL, err := url.Parse(aHref)
	if err != nil {
		return ""
	}
	// If a[href] starts with a /, it overrides the base[href]
	if parsedURL.Host == "" && !strings.HasPrefix(aHref, "/") {
		aHref = path.Join(resolvedBase.Path, aHref)
	}

	resolvedURL, err := resolvedBase.Parse(aHref)
	if err != nil {
		return ""
	}
	return resolvedURL.String()
}

var (
	aHrefMatcher    = cascadia.MustCompile("a[href]")
	baseHrefMatcher = cascadia.MustCompile("base[href]")
)

// Scrape the document's content to gather all links
func (w *worker) processLinks(doc *goquery.Document) (result []*url.URL) {
	baseURL, _ := doc.FindMatcher(baseHrefMatcher).Attr("href")
	urls := doc.FindMatcher(aHrefMatcher).Map(func(_ int, s *goquery.Selection) string {
		val, _ := s.Attr("href")
		if baseURL != "" {
			val = handleBaseTag(doc.Url, baseURL, val)
		}
		return val
	})
	for _, s := range urls {
		// If href starts with "#", then it points to this same exact URL, ignore (will fail to parse anyway)
		if len(s) > 0 && !strings.HasPrefix(s, "#") {
			if parsed, e := url.Parse(s); e == nil {
				parsed = doc.Url.ResolveReference(parsed)
				result = append(result, parsed)
			} else {
				w.logFunc(LogIgnored, "ignore on unparsable policy %s: %s", s, e.Error())
			}
		}
	}
	return
}
<|endoftext|>


<commit_before>package goridge

import (
	"bytes"
	"encoding/gob"

	"github.com/spiral/errors"
	"github.com/vmihailenco/msgpack"
)

func decodeJSON(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode json")
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	return json.Unmarshal(frame.Payload()[opts[1]:], out)
}

func decodeGob(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode json")
	buf := new(bytes.Buffer)
	dec := gob.NewDecoder(buf)
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	payload := frame.Payload()[opts[1]:]
	buf.Write(payload)

	return dec.Decode(out)
}

func decodeRaw(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode raw")
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	payload := frame.Payload()[opts[1]:]

	if raw, ok := out.(*[]byte); ok {
		*raw = append(*raw, payload...)
		return nil
	}

	return nil
}

func decodeMsgPack(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode msgpack")
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	payload := frame.Payload()[opts[1]:]
	return msgpack.Unmarshal(payload, out)
}
<commit_msg>Correct error string in decoders<commit_after>package goridge

import (
	"bytes"
	"encoding/gob"

	"github.com/spiral/errors"
	"github.com/vmihailenco/msgpack"
)

func decodeJSON(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode json")
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	return json.Unmarshal(frame.Payload()[opts[1]:], out)
}

func decodeGob(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode GOB")
	buf := new(bytes.Buffer)
	dec := gob.NewDecoder(buf)
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	payload := frame.Payload()[opts[1]:]
	buf.Write(payload)

	return dec.Decode(out)
}

func decodeRaw(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode raw")
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	payload := frame.Payload()[opts[1]:]

	if raw, ok := out.(*[]byte); ok {
		*raw = append(*raw, payload...)
		return nil
	}

	return nil
}

func decodeMsgPack(out interface{}, frame *Frame) error {
	const op = errors.Op("codec: decode msgpack")
	opts := frame.ReadOptions()
	if len(opts) != 2 {
		return errors.E(op, errors.Str("should be 2 options. SEQ_ID and METHOD_LEN"))
	}
	payload := frame.Payload()[opts[1]:]
	return msgpack.Unmarshal(payload, out)
}
<|endoftext|>