<commit_before>package amadeus

import (
	cc "github.com/tmconsulting/amadeus-ws-go/reqstructs/command_cryptic"
	diit "github.com/tmconsulting/amadeus-ws-go/reqstructs/docissuance_issueticket"
	fmptbs "github.com/tmconsulting/amadeus-ws-go/reqstructs/fare_masterpricertravelboardsearch"
	fppwbc "github.com/tmconsulting/amadeus-ws-go/reqstructs/fare_pricepnrwithbookingclass"
	pame "github.com/tmconsulting/amadeus-ws-go/reqstructs/pnr_addmultielements"
	pc "github.com/tmconsulting/amadeus-ws-go/reqstructs/pnr_cancel"
	pr "github.com/tmconsulting/amadeus-ws-go/reqstructs/pnr_retrieve"
	ql "github.com/tmconsulting/amadeus-ws-go/reqstructs/queue_list"
	qpp "github.com/tmconsulting/amadeus-ws-go/reqstructs/queue_placepnr"
	qri "github.com/tmconsulting/amadeus-ws-go/reqstructs/queue_removeitem"
	sa "github.com/tmconsulting/amadeus-ws-go/reqstructs/security_authenticate"
	ss "github.com/tmconsulting/amadeus-ws-go/reqstructs/security_signout"
	tctfp "github.com/tmconsulting/amadeus-ws-go/reqstructs/ticket_createtstfrompricing"
	tdt "github.com/tmconsulting/amadeus-ws-go/reqstructs/ticket_display_tst"

	ccr "github.com/tmconsulting/amadeus-ws-go/respstructs/command_cryptic_reply"
	diitr "github.com/tmconsulting/amadeus-ws-go/respstructs/docissuance_issueticket_reply"
	fmptbsr "github.com/tmconsulting/amadeus-ws-go/respstructs/fare_masterpricertravelboardsearch_reply"
	fppwbcr "github.com/tmconsulting/amadeus-ws-go/respstructs/fare_pricepnrwithbookingclass_reply"
	pl "github.com/tmconsulting/amadeus-ws-go/respstructs/pnr_list"
	prepl "github.com/tmconsulting/amadeus-ws-go/respstructs/pnr_reply"
	qlr "github.com/tmconsulting/amadeus-ws-go/respstructs/queue_list_reply"
	qppr "github.com/tmconsulting/amadeus-ws-go/respstructs/queue_placepnr_reply"
	qrir "github.com/tmconsulting/amadeus-ws-go/respstructs/queue_removeitem_reply"
	sar "github.com/tmconsulting/amadeus-ws-go/respstructs/security_authenticate_reply"
	ssr "github.com/tmconsulting/amadeus-ws-go/respstructs/security_signout_reply"
	tctfpr "github.com/tmconsulting/amadeus-ws-go/respstructs/ticket_createtstfrompricing_reply"
	tdtr "github.com/tmconsulting/amadeus-ws-go/respstructs/ticket_display_tst_reply"
)

var url string = "http://webservices.amadeus.com"

func (service *WebServicesPT) SecurityAuthenticate(request *sa.SecurityAuthenticate) (*sar.SecurityAuthenticateReply, *Session, error) {
	soapAction := "VLSSLQ_06_1_1A"

	response := new(sar.SecurityAuthenticateReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, nil, err
	}

	return response, session, nil
}

func (service *WebServicesPT) SecuritySignOut(request *ss.SecuritySignOut) (*ssr.SecuritySignOutReply, error) {
	soapAction := "VLSSOQ_04_1_1A"

	response := new(ssr.SecuritySignOutReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) CommandCryptic(request *cc.CommandCryptic) (*ccr.CommandCrypticReply, *Session, error) {
	soapAction := "HSFREQ_07_3_1A"

	response := new(ccr.CommandCrypticReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, nil, err
	}

	return response, session, nil
}

func (service *WebServicesPT) QueueList(request *ql.QueueList) (*qlr.QueueListReply, error) {
	soapAction := "QDQLRQ_11_1_1A"

	response := new(qlr.QueueListReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) TicketDisplayTST(request *tdt.TicketDisplayTST) (*tdtr.TicketDisplayTSTReply, error) {
	soapAction := "TTSTRQ_07_1_1A"

	response := new(tdtr.TicketDisplayTSTReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) DocIssuanceIssueTicket(request *diit.DocIssuanceIssueTicket) (*diitr.DocIssuanceIssueTicketReply, error) {
	soapAction := "TTKTIQ_09_1_1A"

	response := new(diitr.DocIssuanceIssueTicketReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) FarePricePNRWithBookingClass(request *fppwbc.FarePricePNRWithBookingClass) (*fppwbcr.FarePricePNRWithBookingClassReply, error) {
	soapAction := "TPCBRQ_12_4_1A"

	response := new(fppwbcr.FarePricePNRWithBookingClassReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRRetrieve2(request *pr.PNRRetrieve) (*pl.PNRList, error) {
	soapAction := "PNRRET_11_3_1A"

	response := new(pl.PNRList)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRRetrieve(request *pr.PNRRetrieve) (*prepl.PNRReply, error) {
	soapAction := "PNRRET_11_3_1A"

	response := new(prepl.PNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRCancel(request *pc.PNRCancel) (*prepl.PNRReply, error) {
	soapAction := "PNRXCL_11_3_1A"

	response := new(prepl.PNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) TicketCreateTSTFromPricing(request *tctfp.TicketCreateTSTFromPricing) (*tctfpr.TicketCreateTSTFromPricingReply, error) {
	soapAction := "TAUTCQ_04_1_1A"

	response := new(tctfpr.TicketCreateTSTFromPricingReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) QueueRemoveItem(request *qri.QueueRemoveItem) (*qrir.QueueRemoveItemReply, error) {
	soapAction := "QUQMDQ_03_1_1A"

	response := new(qrir.QueueRemoveItemReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) QueuePlacePNR(request *qpp.QueuePlacePNR) (*qppr.QueuePlacePNRReply, error) {
	soapAction := "QUQPCQ_03_1_1A"

	response := new(qppr.QueuePlacePNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRAddMultiElements(request *pame.PNRAddMultiElements) (*prepl.PNRReply, error) {
	soapAction := "PNRADD_11_3_1A"

	response := new(prepl.PNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) FareMasterPricerTravelBoardSearch(request *fmptbs.FareMasterPricerTravelBoardSearch) (*fmptbsr.FareMasterPricerTravelBoardSearchReply, error) {
	soapAction := "FMPTBQ_14_3_1A"

	response := new(fmptbsr.FareMasterPricerTravelBoardSearchReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}
<commit_msg>added FareInformativePricingWithoutPNR<commit_after>package amadeus

import (
	cc "github.com/tmconsulting/amadeus-ws-go/reqstructs/command_cryptic"
	diit "github.com/tmconsulting/amadeus-ws-go/reqstructs/docissuance_issueticket"
	fipwpnr "github.com/tmconsulting/amadeus-ws-go/reqstructs/fare_informativepricingwithoutpnr"
	fmptbs "github.com/tmconsulting/amadeus-ws-go/reqstructs/fare_masterpricertravelboardsearch"
	fppwbc "github.com/tmconsulting/amadeus-ws-go/reqstructs/fare_pricepnrwithbookingclass"
	pame "github.com/tmconsulting/amadeus-ws-go/reqstructs/pnr_addmultielements"
	pc "github.com/tmconsulting/amadeus-ws-go/reqstructs/pnr_cancel"
	pr "github.com/tmconsulting/amadeus-ws-go/reqstructs/pnr_retrieve"
	ql "github.com/tmconsulting/amadeus-ws-go/reqstructs/queue_list"
	qpp "github.com/tmconsulting/amadeus-ws-go/reqstructs/queue_placepnr"
	qri "github.com/tmconsulting/amadeus-ws-go/reqstructs/queue_removeitem"
	sa "github.com/tmconsulting/amadeus-ws-go/reqstructs/security_authenticate"
	ss "github.com/tmconsulting/amadeus-ws-go/reqstructs/security_signout"
	tctfp "github.com/tmconsulting/amadeus-ws-go/reqstructs/ticket_createtstfrompricing"
	tdt "github.com/tmconsulting/amadeus-ws-go/reqstructs/ticket_display_tst"

	ccr "github.com/tmconsulting/amadeus-ws-go/respstructs/command_cryptic_reply"
	diitr "github.com/tmconsulting/amadeus-ws-go/respstructs/docissuance_issueticket_reply"
	fipwpnrr "github.com/tmconsulting/amadeus-ws-go/respstructs/fare_informativepricingwithoutpnr_reply"
	fmptbsr "github.com/tmconsulting/amadeus-ws-go/respstructs/fare_masterpricertravelboardsearch_reply"
	fppwbcr "github.com/tmconsulting/amadeus-ws-go/respstructs/fare_pricepnrwithbookingclass_reply"
	pl "github.com/tmconsulting/amadeus-ws-go/respstructs/pnr_list"
	prepl "github.com/tmconsulting/amadeus-ws-go/respstructs/pnr_reply"
	qlr "github.com/tmconsulting/amadeus-ws-go/respstructs/queue_list_reply"
	qppr "github.com/tmconsulting/amadeus-ws-go/respstructs/queue_placepnr_reply"
	qrir "github.com/tmconsulting/amadeus-ws-go/respstructs/queue_removeitem_reply"
	sar "github.com/tmconsulting/amadeus-ws-go/respstructs/security_authenticate_reply"
	ssr "github.com/tmconsulting/amadeus-ws-go/respstructs/security_signout_reply"
	tctfpr "github.com/tmconsulting/amadeus-ws-go/respstructs/ticket_createtstfrompricing_reply"
	tdtr "github.com/tmconsulting/amadeus-ws-go/respstructs/ticket_display_tst_reply"
)

var url string = "http://webservices.amadeus.com"

func (service *WebServicesPT) SecurityAuthenticate(request *sa.SecurityAuthenticate) (*sar.SecurityAuthenticateReply, *Session, error) {
	soapAction := "VLSSLQ_06_1_1A"

	response := new(sar.SecurityAuthenticateReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, nil, err
	}

	return response, session, nil
}

func (service *WebServicesPT) SecuritySignOut(request *ss.SecuritySignOut) (*ssr.SecuritySignOutReply, error) {
	soapAction := "VLSSOQ_04_1_1A"

	response := new(ssr.SecuritySignOutReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) CommandCryptic(request *cc.CommandCryptic) (*ccr.CommandCrypticReply, *Session, error) {
	soapAction := "HSFREQ_07_3_1A"

	response := new(ccr.CommandCrypticReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, nil, err
	}

	return response, session, nil
}

func (service *WebServicesPT) QueueList(request *ql.QueueList) (*qlr.QueueListReply, error) {
	soapAction := "QDQLRQ_11_1_1A"

	response := new(qlr.QueueListReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) TicketDisplayTST(request *tdt.TicketDisplayTST) (*tdtr.TicketDisplayTSTReply, error) {
	soapAction := "TTSTRQ_07_1_1A"

	response := new(tdtr.TicketDisplayTSTReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) DocIssuanceIssueTicket(request *diit.DocIssuanceIssueTicket) (*diitr.DocIssuanceIssueTicketReply, error) {
	soapAction := "TTKTIQ_09_1_1A"

	response := new(diitr.DocIssuanceIssueTicketReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) FarePricePNRWithBookingClass(request *fppwbc.FarePricePNRWithBookingClass) (*fppwbcr.FarePricePNRWithBookingClassReply, error) {
	soapAction := "TPCBRQ_12_4_1A"

	response := new(fppwbcr.FarePricePNRWithBookingClassReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRRetrieve2(request *pr.PNRRetrieve) (*pl.PNRList, error) {
	soapAction := "PNRRET_11_3_1A"

	response := new(pl.PNRList)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRRetrieve(request *pr.PNRRetrieve) (*prepl.PNRReply, error) {
	soapAction := "PNRRET_11_3_1A"

	response := new(prepl.PNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRCancel(request *pc.PNRCancel) (*prepl.PNRReply, error) {
	soapAction := "PNRXCL_11_3_1A"

	response := new(prepl.PNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) TicketCreateTSTFromPricing(request *tctfp.TicketCreateTSTFromPricing) (*tctfpr.TicketCreateTSTFromPricingReply, error) {
	soapAction := "TAUTCQ_04_1_1A"

	response := new(tctfpr.TicketCreateTSTFromPricingReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) QueueRemoveItem(request *qri.QueueRemoveItem) (*qrir.QueueRemoveItemReply, error) {
	soapAction := "QUQMDQ_03_1_1A"

	response := new(qrir.QueueRemoveItemReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) QueuePlacePNR(request *qpp.QueuePlacePNR) (*qppr.QueuePlacePNRReply, error) {
	soapAction := "QUQPCQ_03_1_1A"

	response := new(qppr.QueuePlacePNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) PNRAddMultiElements(request *pame.PNRAddMultiElements) (*prepl.PNRReply, error) {
	soapAction := "PNRADD_11_3_1A"

	response := new(prepl.PNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) FareMasterPricerTravelBoardSearch(request *fmptbs.FareMasterPricerTravelBoardSearch) (*fmptbsr.FareMasterPricerTravelBoardSearchReply, error) {
	soapAction := "FMPTBQ_14_3_1A"

	response := new(fmptbsr.FareMasterPricerTravelBoardSearchReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}

func (service *WebServicesPT) FareInformativePricingWithoutPNR(request *fipwpnr.FareInformativePricingWithoutPNR) (*fipwpnrr.FareInformativePricingWithoutPNRReply, error) {
	soapAction := "TIPNRQ_12_4_1A"

	response := new(fipwpnrr.FareInformativePricingWithoutPNRReply)
	session := new(Session)
	err := service.client.Call(url+"/"+service.wsap+"/"+soapAction, request, response, session)
	if err != nil {
		return nil, err
	}

	return response, nil
}
<|endoftext|>


<commit_before>package pget

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/net/context/ctxhttp"
	"golang.org/x/sync/errgroup"
)

// Range struct for range access
type Range struct {
	low    uint
	high   uint
	worker uint
}

func isNotLastURL(url, purl string) bool {
	return url != purl && url != ""
}

func isLastProc(i, procs uint) bool {
	return i == procs-1
}

// Checking is check to can request
func (p *Pget) Checking() error {

	ctx, cancelAll := context.WithTimeout(context.Background(), time.Duration(p.timeout)*time.Second)

	ch := MakeCh()
	defer ch.Close()

	for _, url := range p.URLs {
		fmt.Fprintf(os.Stdout, "Checking now %s\n", url)
		go p.CheckMirrors(ctx, url, ch)
	}

	// listen for error or size channel
	size, err := ch.CheckingListen(ctx, cancelAll, len(p.URLs))
	if err != nil {
		return err
	}

	// did already get filename from -o option
	filename := p.Utils.FileName()
	if filename == "" {
		filename = p.Utils.URLFileName(p.TargetDir, p.TargetURLs[0])
	}
	p.SetFileName(filename)
	p.SetFullFileName(p.TargetDir, filename)
	p.Utils.SetDirName(p.TargetDir, filename, p.Procs)

	p.SetFileSize(size)

	return nil
}

// CheckMirrors method check be able to range access. also get redirected url.
func (p *Pget) CheckMirrors(ctx context.Context, url string, ch *Ch) {

	res, err := ctxhttp.Head(ctx, http.DefaultClient, url)
	if err != nil {
		ch.Err <- errors.Wrap(err, "failed to head request: "+url)
		return
	}

	if res.Header.Get("Accept-Ranges") != "bytes" {
		ch.Err <- errors.Errorf("not supported range access: %s", url)
		return
	}

	// To perform with the correct "range access"
	// get the last url in the redirect
	_url := res.Request.URL.String()
	if isNotLastURL(_url, url) {
		p.TargetURLs = append(p.TargetURLs, _url)
	} else {
		p.TargetURLs = append(p.TargetURLs, url)
	}

	// get of ContentLength
	size := uint(res.ContentLength)
	if size <= 0 {
		ch.Err <- errors.New("invalid content length")
	} else {
		ch.Size <- size
	}
}

// Download method distributes the task to each goroutine for each URL
func (p *Pget) Download() error {

	procs := uint(p.Procs)

	filesize := p.FileSize()
	dirname := p.DirName()

	// create download location
	if err := os.MkdirAll(dirname, 0755); err != nil {
		return errors.Wrap(err, "failed to mkdir for download location")
	}

	// calculate split file size
	split := filesize / procs

	if err := p.Utils.IsFree(split); err != nil {
		return err
	}

	grp, ctx := errgroup.WithContext(context.Background())

	// on an assignment for request
	p.Assignment(grp, procs, split)

	if err := p.Utils.ProgressBar(ctx); err != nil {
		return err
	}

	// wait for Assignment method
	if err := grp.Wait(); err != nil {
		return err
	}

	return nil
}

// Assignment method that to each goroutine gives the task
func (p Pget) Assignment(grp *errgroup.Group, procs, split uint) {
	filename := p.FileName()
	dirname := p.DirName()

	assignment := uint(p.Procs / len(p.TargetURLs))

	var lasturl string
	totalActiveProcs := uint(0)
	for i := uint(0); i < procs; i++ {
		partName := fmt.Sprintf("%s/%s.%d.%d", dirname, filename, procs, i)

		// make range
		r := p.Utils.MakeRange(i, split, procs)

		if info, err := os.Stat(partName); err == nil {
			infosize := uint(info.Size())
			// check if the part is fully downloaded
			if isLastProc(i, procs) {
				if infosize == r.high-r.low {
					continue
				}
			} else if infosize == split {
				// skip as the part is already downloaded
				continue
			}

			// make low range from this next byte
			r.low += infosize
		}

		totalActiveProcs++

		url := p.TargetURLs[0]

		// give efficiency and equality work
		if totalActiveProcs%assignment == 0 {
			// Like shift method
			if len(p.TargetURLs) > 1 {
				p.TargetURLs = p.TargetURLs[1:]
			}

			// check whether to output the message
			if lasturl != url {
				fmt.Fprintf(os.Stdout, "Download start from %s\n", url)
				lasturl = url
			}
		}

		// execute get request
		grp.Go(func() error {
			return p.Requests(r, filename, dirname, url)
		})
	}
}

// Requests method will download the file
func (p Pget) Requests(r Range, filename, dirname, url string) error {

	res, err := p.MakeResponse(r, url) // ctxhttp
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to split get requests: %d", r.worker))
	}
	defer res.Body.Close()

	partName := fmt.Sprintf("%s/%s.%d.%d", dirname, filename, p.Procs, r.worker)

	output, err := os.OpenFile(partName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to create %s in %s", filename, dirname))
	}
	defer output.Close()

	io.Copy(output, res.Body)

	return nil
}

// MakeResponse return *http.Response include context and range header
func (p Pget) MakeResponse(r Range, url string) (*http.Response, error) {
	// create get request
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("failed to split NewRequest for get: %d", r.worker))
	}

	// set download ranges
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", r.low, r.high))

	// set useragent
	if p.useragent != "" {
		req.Header.Set("User-Agent", p.useragent)
	}

	// set referer
	if p.referer != "" {
		req.Header.Set("Referer", p.referer)
	}

	return http.DefaultClient.Do(req)
}
<commit_msg>Evaluate ContentLength as int<commit_after>package pget

import (
	"context"
	"fmt"
	"io"
	"net/http"
	"os"
	"time"

	"github.com/pkg/errors"
	"golang.org/x/net/context/ctxhttp"
	"golang.org/x/sync/errgroup"
)

// Range struct for range access
type Range struct {
	low    uint
	high   uint
	worker uint
}

func isNotLastURL(url, purl string) bool {
	return url != purl && url != ""
}

func isLastProc(i, procs uint) bool {
	return i == procs-1
}

// Checking is check to can request
func (p *Pget) Checking() error {

	ctx, cancelAll := context.WithTimeout(context.Background(), time.Duration(p.timeout)*time.Second)

	ch := MakeCh()
	defer ch.Close()

	for _, url := range p.URLs {
		fmt.Fprintf(os.Stdout, "Checking now %s\n", url)
		go p.CheckMirrors(ctx, url, ch)
	}

	// listen for error or size channel
	size, err := ch.CheckingListen(ctx, cancelAll, len(p.URLs))
	if err != nil {
		return err
	}

	// did already get filename from -o option
	filename := p.Utils.FileName()
	if filename == "" {
		filename = p.Utils.URLFileName(p.TargetDir, p.TargetURLs[0])
	}
	p.SetFileName(filename)
	p.SetFullFileName(p.TargetDir, filename)
	p.Utils.SetDirName(p.TargetDir, filename, p.Procs)

	p.SetFileSize(size)

	return nil
}

// CheckMirrors method check be able to range access. also get redirected url.
func (p *Pget) CheckMirrors(ctx context.Context, url string, ch *Ch) {

	res, err := ctxhttp.Head(ctx, http.DefaultClient, url)
	if err != nil {
		ch.Err <- errors.Wrap(err, "failed to head request: "+url)
		return
	}

	if res.Header.Get("Accept-Ranges") != "bytes" {
		ch.Err <- errors.Errorf("not supported range access: %s", url)
		return
	}

	// To perform with the correct "range access"
	// get the last url in the redirect
	_url := res.Request.URL.String()
	if isNotLastURL(_url, url) {
		p.TargetURLs = append(p.TargetURLs, _url)
	} else {
		p.TargetURLs = append(p.TargetURLs, url)
	}

	// get of ContentLength
	if res.ContentLength <= 0 {
		ch.Err <- errors.New("invalid content length")
		return
	}
	ch.Size <- uint(res.ContentLength)
}

// Download method distributes the task to each goroutine for each URL
func (p *Pget) Download() error {

	procs := uint(p.Procs)

	filesize := p.FileSize()
	dirname := p.DirName()

	// create download location
	if err := os.MkdirAll(dirname, 0755); err != nil {
		return errors.Wrap(err, "failed to mkdir for download location")
	}

	// calculate split file size
	split := filesize / procs

	if err := p.Utils.IsFree(split); err != nil {
		return err
	}

	grp, ctx := errgroup.WithContext(context.Background())

	// on an assignment for request
	p.Assignment(grp, procs, split)

	if err := p.Utils.ProgressBar(ctx); err != nil {
		return err
	}

	// wait for Assignment method
	if err := grp.Wait(); err != nil {
		return err
	}

	return nil
}

// Assignment method that to each goroutine gives the task
func (p Pget) Assignment(grp *errgroup.Group, procs, split uint) {
	filename := p.FileName()
	dirname := p.DirName()

	assignment := uint(p.Procs / len(p.TargetURLs))

	var lasturl string
	totalActiveProcs := uint(0)
	for i := uint(0); i < procs; i++ {
		partName := fmt.Sprintf("%s/%s.%d.%d", dirname, filename, procs, i)

		// make range
		r := p.Utils.MakeRange(i, split, procs)

		if info, err := os.Stat(partName); err == nil {
			infosize := uint(info.Size())
			// check if the part is fully downloaded
			if isLastProc(i, procs) {
				if infosize == r.high-r.low {
					continue
				}
			} else if infosize == split {
				// skip as the part is already downloaded
				continue
			}

			// make low range from this next byte
			r.low += infosize
		}

		totalActiveProcs++

		url := p.TargetURLs[0]

		// give efficiency and equality work
		if totalActiveProcs%assignment == 0 {
			// Like shift method
			if len(p.TargetURLs) > 1 {
				p.TargetURLs = p.TargetURLs[1:]
			}

			// check whether to output the message
			if lasturl != url {
				fmt.Fprintf(os.Stdout, "Download start from %s\n", url)
				lasturl = url
			}
		}

		// execute get request
		grp.Go(func() error {
			return p.Requests(r, filename, dirname, url)
		})
	}
}

// Requests method will download the file
func (p Pget) Requests(r Range, filename, dirname, url string) error {

	res, err := p.MakeResponse(r, url) // ctxhttp
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to split get requests: %d", r.worker))
	}
	defer res.Body.Close()

	partName := fmt.Sprintf("%s/%s.%d.%d", dirname, filename, p.Procs, r.worker)

	output, err := os.OpenFile(partName, os.O_WRONLY|os.O_CREATE|os.O_APPEND, 0666)
	if err != nil {
		return errors.Wrap(err, fmt.Sprintf("failed to create %s in %s", filename, dirname))
	}
	defer output.Close()

	io.Copy(output, res.Body)

	return nil
}

// MakeResponse return *http.Response include context and range header
func (p Pget) MakeResponse(r Range, url string) (*http.Response, error) {
	// create get request
	req, err := http.NewRequest("GET", url, nil)
	if err != nil {
		return nil, errors.Wrap(err, fmt.Sprintf("failed to split NewRequest for get: %d", r.worker))
	}

	// set download ranges
	req.Header.Set("Range", fmt.Sprintf("bytes=%d-%d", r.low, r.high))

	// set useragent
	if p.useragent != "" {
		req.Header.Set("User-Agent", p.useragent)
	}

	// set referer
	if p.referer != "" {
		req.Header.Set("Referer", p.referer)
	}

	return http.DefaultClient.Do(req)
}
<|endoftext|>


<commit_before>package readers

import (
	"errors"
)

func NewGoStruct(name string) (IReaderWriter, error) {
	var structInstance IReaderWriter

	if name == "CpuInfo" {
		structInstance = NewCpuInfo()
	}
	if name == "Df" {
		structInstance = NewDf()
	}
	if name == "Du" {
		structInstance = NewDu()
	}
	if name == "DiskPartitions" {
		structInstance = NewDiskPartitions()
	}
	if name == "DockerContainersMemory" {
		structInstance = NewDockerContainersMemory()
	}
	if name == "DockerContainersCpu" {
		structInstance = NewDockerContainersCpu()
	}
	if name == "Free" {
		structInstance = NewFree()
	}
	if name == "LoadAvg" {
		structInstance = NewLoadAvg()
	}
	if name == "Ps" {
		structInstance = NewPs()
	}
	if name == "Meminfo" {
		structInstance = NewMeminfo()
	}
	if name == "NetworkInterfaces" {
		structInstance = NewNetworkInterfaces()
	}
	if name == "Uptime" {
		structInstance = NewUptime()
	}

	if structInstance == nil {
		return nil, errors.New("GoStruct is undefined.")
	}

	return structInstance, nil
}

type IReaderWriter interface {
	Run() error
	ToJson() ([]byte, error)
}
<commit_msg>enable DiskIO<commit_after>package readers

import (
	"errors"
)

func NewGoStruct(name string) (IReaderWriter, error) {
	var structInstance IReaderWriter

	if name == "CpuInfo" {
		structInstance = NewCpuInfo()
	}
	if name == "Df" {
		structInstance = NewDf()
	}
	if name == "Du" {
		structInstance = NewDu()
	}
	if name == "DiskPartitions" {
		structInstance = NewDiskPartitions()
	}
	if name == "DiskIO" {
		structInstance = NewDiskIO()
	}
	if name == "DockerContainersMemory" {
		structInstance = NewDockerContainersMemory()
	}
	if name == "DockerContainersCpu" {
		structInstance = NewDockerContainersCpu()
	}
	if name == "Free" {
		structInstance = NewFree()
	}
	if name == "LoadAvg" {
		structInstance = NewLoadAvg()
	}
	if name == "Ps" {
		structInstance = NewPs()
	}
	if name == "Meminfo" {
		structInstance = NewMeminfo()
	}
	if name == "NetworkInterfaces" {
		structInstance = NewNetworkInterfaces()
	}
	if name == "Uptime" {
		structInstance = NewUptime()
	}

	if structInstance == nil {
		return nil, errors.New("GoStruct is undefined.")
	}

	return structInstance, nil
}

type IReaderWriter interface {
	Run() error
	ToJson() ([]byte, error)
}
<|endoftext|>


<commit_before>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.
//
// Windows utility for silently starting console processes
// without showing a console.
// Must be built with -ldflags "-H windowsgui"

// +build windows

package main

import (
	"fmt"
	"log"
	"os"
	"strings"
	"syscall"
	"time"
)

const flagCreateNewConsole = 0x00000010

func main() {
	argsIndex := 1 // 0 is the name of this program, 1 is either the one to launch or the "wait" option

	if len(os.Args) < 2 {
		log.Fatal("ERROR: no arguments. Use [-wait] programname [arg arg arg]\n")
	}

	// Do this awkward thing so we can pass along the rest of the command line as-is

	doWait := false
	doHide := true
	for i := 1; i < 3 && (i+1) < len(os.Args); i++ {
		if strings.EqualFold(os.Args[argsIndex], "-wait") {
			argsIndex++
			doWait = true
		} else if strings.EqualFold(os.Args[argsIndex], "-show") {
			argsIndex++
			doHide = false
		}
	}

	attr := &syscall.ProcAttr{
		Files: []uintptr{uintptr(syscall.Stdin), uintptr(syscall.Stdout), uintptr(syscall.Stderr)},
		Env:   syscall.Environ(),
		Sys: &syscall.SysProcAttr{
			HideWindow:    doHide,
			CreationFlags: flagCreateNewConsole,
		},
	}
	fmt.Printf("Launching %s with args %v\n", os.Args[argsIndex], os.Args[argsIndex:])
	pid, _, err := syscall.StartProcess(os.Args[argsIndex], os.Args[argsIndex:], attr)
	if err != nil {
		fmt.Printf("StartProcess error: %s\n", err.Error())
	} else if doWait {
		p, err := os.FindProcess(pid)
		if err != nil {
			fmt.Printf("Launcher can't find %d\n", pid)
		}

		timeout := make(chan time.Time, 1)

		go func() {
			pstate, err := p.Wait()

			if err == nil && pstate.Success() {
				time.Sleep(100 * time.Millisecond)
			} else {
				fmt.Printf("Unsuccessful wait: Error %v, pstate %v\n", err, *pstate)
			}
			timeout <- time.Now()
		}()

		// Only wait 15 seconds because an erroring command was shown to hang
		// up an installer on Win7
		select {
		case _ = <-timeout:
			// success
		case <-time.After(15 * time.Second):
			fmt.Println("timed out")
		}
	}
}
<commit_msg>Adapt windows launcher utility to run programs as user despite elevation (#13951)<commit_after>// Copyright 2015 Keybase, Inc. All rights reserved. Use of
// this source code is governed by the included BSD license.
//
// Windows utility for silently starting console processes
// without showing a console.
// Must be built with -ldflags "-H windowsgui"

// +build windows

package main

import (
	"errors"
	"fmt"
	"log"
	"os"
	"strings"
	"syscall"
	"time"
	"unsafe"

	"github.com/keybase/go-winio"
	"golang.org/x/sys/windows"
)

const flagCreateNewConsole = 0x00000010

var (
	modadvapi32 *windows.LazyDLL = windows.NewLazySystemDLL("advapi32.dll")
	moduser32   *windows.LazyDLL = windows.NewLazySystemDLL("user32.dll")

	procDuplicateTokenEx         *windows.LazyProc = modadvapi32.NewProc("DuplicateTokenEx")
	procGetShellWindow           *windows.LazyProc = moduser32.NewProc("GetShellWindow")
	procGetWindowThreadProcessId *windows.LazyProc = moduser32.NewProc("GetWindowThreadProcessId")
	procCreateProcessWithTokenW  *windows.LazyProc = modadvapi32.NewProc("CreateProcessWithTokenW")
)

func GetWindowThreadProcessId(hwnd syscall.Handle) int {
	var processID int
	_, _, _ = procGetWindowThreadProcessId.Call(
		uintptr(hwnd),
		uintptr(unsafe.Pointer(&processID)))

	return processID
}

const SecurityImpersonation = 2

type TokenType uint32

const (
	TokenPrimary       TokenType = 1
	TokenImpersonation TokenType = 2
)

func DuplicateTokenEx(hExistingToken windows.Token, dwDesiredAccess uint32, lpTokenAttributes *syscall.SecurityAttributes, impersonationLevel uint32, tokenType TokenType, phNewToken *windows.Token) (err error) {
	r1, _, e1 := syscall.Syscall6(procDuplicateTokenEx.Addr(), 6, uintptr(hExistingToken), uintptr(dwDesiredAccess), uintptr(unsafe.Pointer(lpTokenAttributes)), uintptr(impersonationLevel), uintptr(tokenType), uintptr(unsafe.Pointer(phNewToken)))
	if r1 == 0 {
		if e1 != 0 {
			err = syscall.Errno(e1)
		} else {
			err = syscall.EINVAL
		}
	}
	return
}

// makeCmdLine builds a command line out of args by escaping "special"
// characters and joining the arguments with spaces.
func makeCmdLine(args []string) string {
	var s string
	for _, v := range args {
		if s != "" {
			s += " "
		}
		s += syscall.EscapeArg(v)
	}
	return s
}

func CreateProcessWithTokenW(hToken syscall.Token, argv []string, attr *syscall.ProcAttr) (pid int, handle uintptr, err error) {
	var sys = attr.Sys

	var cmdline string
	// Windows CreateProcess takes the command line as a single string:
	// use attr.CmdLine if set, else build the command line by escaping
	// and joining each argument with spaces
	if sys.CmdLine != "" {
		cmdline = sys.CmdLine
	} else {
		cmdline = makeCmdLine(argv)
	}

	var argvp *uint16
	if len(cmdline) != 0 {
		argvp, err = syscall.UTF16PtrFromString(cmdline)
		if err != nil {
			return 0, 0, err
		}
	}

	si := new(syscall.StartupInfo)
	si.Cb = uint32(unsafe.Sizeof(*si))
	if sys.HideWindow {
		si.Flags |= syscall.STARTF_USESHOWWINDOW
		si.ShowWindow = syscall.SW_HIDE
	}

	pi := new(syscall.ProcessInformation)

	flags := sys.CreationFlags | syscall.CREATE_UNICODE_ENVIRONMENT
	r1, _, e1 := procCreateProcessWithTokenW.Call(
		uintptr(hToken),                // HANDLE                hToken,
		0,                              // DWORD                 dwLogonFlags,
		uintptr(0),                     // LPCWSTR               lpApplicationName,
		uintptr(unsafe.Pointer(argvp)), // LPWSTR                lpCommandLine,
		uintptr(flags),                 // DWORD                 dwCreationFlags,
		uintptr(0),                     // LPVOID                lpEnvironment,
		uintptr(0),                     // LPCWSTR               lpCurrentDirectory,
		uintptr(unsafe.Pointer(si)),    // LPSTARTUPINFOW        lpStartupInfo,
		uintptr(unsafe.Pointer(pi)),    // LPPROCESS_INFORMATION lpProcessInformation
	)

	if r1 != 0 {
		e1 = nil
	}

	return int(pi.ProcessId), uintptr(pi.Process), e1
}

// Protection against running elevated: attempt to run as regular user instead.
// Assume an error means we weren't elevated, which should be the normal case.
// see https://blogs.msdn.microsoft.com/aaron_margosis/2009/06/06/faq-how-do-i-start-a-program-as-the-desktop-user-from-an-elevated-app/
func getUserToken() (syscall.Token, error) {
	var err error
	var shellWindow uintptr
	if shellWindow, _, err = procGetShellWindow.Call(); shellWindow == 0 {
		return 0, fmt.Errorf("call native GetShellWindow: %s", err)
	}

	processID := GetWindowThreadProcessId(syscall.Handle(shellWindow))

	if processID == 0 {
		return 0, errors.New("can't get desktop window proc ID")
	}

	h, e := syscall.OpenProcess(syscall.PROCESS_QUERY_INFORMATION, false, uint32(processID))
	if e != nil {
		return 0, fmt.Errorf("OpenProcess error: %s", e)
	}
	defer syscall.CloseHandle(h)

	var token, dupToken windows.Token
	err = windows.OpenProcessToken(windows.Handle(h), windows.TOKEN_DUPLICATE, &token)
	if err != nil {
		return 0, err
	}
	const TOKEN_ADJUST_SESSIONID = 256
	dwTokenRights := uint32(syscall.TOKEN_QUERY | syscall.TOKEN_ASSIGN_PRIMARY | syscall.TOKEN_DUPLICATE | syscall.TOKEN_ADJUST_DEFAULT | TOKEN_ADJUST_SESSIONID)
	err = DuplicateTokenEx(token, dwTokenRights, nil, SecurityImpersonation, TokenPrimary, &dupToken)
	if err != nil || dupToken == 0 {
		return 0, fmt.Errorf("DuplicateTokenEx error: %s", err.Error())
	}

	return syscall.Token(dupToken), nil
}

// elevated means we try running as user
func doRun(elevated bool) error {
	argsIndex := 1 // 0 is the name of this program, 1 is either the one to launch or the "wait" option

	if len(os.Args) < 2 {
		log.Fatal("ERROR: no arguments. Use [-wait] programname [arg arg arg]\n")
	}

	// Do this awkward thing so we can pass along the rest of the command line as-is

	doWait := false
	doHide := true
	for i := 1; i < 3 && (i+1) < len(os.Args); i++ {
		if strings.EqualFold(os.Args[argsIndex], "-wait") {
			argsIndex++
			doWait = true
		} else if strings.EqualFold(os.Args[argsIndex], "-show") {
			argsIndex++
			doHide = false
		}
	}
	attr := &syscall.ProcAttr{
		Files: []uintptr{uintptr(syscall.Stdin), uintptr(syscall.Stdout), uintptr(syscall.Stderr)},
		Env:   syscall.Environ(),
		Sys: &syscall.SysProcAttr{
			HideWindow:    doHide,
			CreationFlags: flagCreateNewConsole,
		},
	}
	fmt.Printf("Launching %s with args %v\n", os.Args[argsIndex], os.Args[argsIndex:])

	var err error
	var pid int

	if elevated {
		token, _ := getUserToken()

		pid, _, err = CreateProcessWithTokenW(token, os.Args[argsIndex:], attr)
		defer syscall.CloseHandle(syscall.Handle(token))
		if err != nil {
			fmt.Printf("CreateProcessWithTokenW error: %s\n", err.Error())
			return err
		}
	} else {
		pid, _, err = syscall.StartProcess(os.Args[argsIndex], os.Args[argsIndex:], attr)
	}

	if err != nil {
		fmt.Printf("StartProcess error: %s\n", err.Error())
		return err
	} else if doWait {
		p, err := os.FindProcess(pid)
		if err != nil {
			fmt.Printf("Launcher can't find %d\n", pid)
			return err
		}

		timeout := make(chan time.Time, 1)

		go func() {
			pstate, err := p.Wait()

			if err == nil && pstate.Success() {
				time.Sleep(100 * time.Millisecond)
			} else {
				fmt.Printf("Unsuccessful wait: Error %v, pstate %v\n", err, *pstate)
			}
			timeout <- time.Now()
		}()

		// Only wait 15 seconds because an erroring command was shown to hang
		// up an installer on Win7
		select {
		case _ = <-timeout:
			// success
		case <-time.After(15 * time.Second):
			fmt.Println("timed out")
		}
	}
	return nil
}

func main() {
	// RunWithPrivilege enables a single privilege for a function call.
	// SeIncreaseQuotaPrivilege will only work when elevated
	err := winio.RunWithPrivilege("SeIncreaseQuotaPrivilege", func() error {
		result := doRun(true)
		if result != nil {
			// Print this failure but not the RunWithPrivilege result, since
			// RunWithPrivilege failure is the normal case
			fmt.Printf("De-elevation failure: %v\n", result)
		}
		return result
	})
	if err != nil {
		// This means we weren't elevated, or de-elevation failed. Run without.
		doRun(false)
	}
}
<|endoftext|>


<commit_before>package main

import (
	"io"
	"log"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/buger/gor/requestfiles"
)

type DiffReporter struct {
	total int64
	diff  int64

	settings *Settings

	stats     *Stats
	statNames StatNames

	// If writing out diffs, need a queue to serialize to a single writer.
	outQueue       chan []byte
	requestsWriter io.Writer
}

// Compute these once at startup to avoid allocating them every time
type StatNames struct {
	diff  string
	match string
	total string

	errA string
	errB string

	rttA string
	rttB string
}

func NewDiffReporter(s *Settings, stats *Stats) (d *DiffReporter) {
	r := new(DiffReporter)

	r.settings = s

	r.stats = stats

	r.statNames = StatNames{
		total: "diffing.total",
		match: "diffing.match",
		diff:  "diffing.diff",
		errA:  "diffing.err." + s.nameA,
		errB:  "diffing.err." + s.nameB,
		rttA:  "diffing.rtt." + s.nameA,
		rttB:  "diffing.rtt." + s.nameB,
	}

	if s.requestsFile != "" {
		r.outQueue = make(chan []byte, 100)
		r.requestsWriter = requestfiles.NewFileOutput(s.requestsFile)
		go r.writeDiffs()
	}

	return r
}

func (d *DiffReporter) writeDiffs() {
	for {
		req := <-d.outQueue
		d.requestsWriter.Write(req)
	}
}

func isErr(err error, status int) bool {
	return err != nil || status/100 == 5
}

func (d *DiffReporter) Compare(req *http.Request, raw []byte, statusA, statusB int, respA, respB string, rawErrA, rawErrB error, rttA, rttB time.Duration) {
	atomic.AddInt64(&d.total, 1)
	d.stats.Inc(d.statNames.total)

	errA := isErr(rawErrA, statusA)
	errB := isErr(rawErrB, statusB)

	if errA {
		d.stats.Inc(d.statNames.errA)
	} else {
		d.stats.Timing(d.statNames.rttA, rttA)
	}

	if errB {
		d.stats.Inc(d.statNames.errB)
	} else {
		d.stats.Timing(d.statNames.rttB, rttB)
	}

	if (errA && errB) || (d.settings.ignoreErrors && (errA || errB)) {
		return
	}

	if !errA && !errB && respA == respB {
		d.stats.Inc(d.statNames.match)
		return
	}

	atomic.AddInt64(&d.diff, 1)
	d.stats.Inc(d.statNames.diff)

	sizeA := len(respA)
	sizeB := len(respB)

	limit := sizeA
	if sizeA > sizeB {
		limit = sizeB
	}

	i := 0
	for i < limit {
		if respA[i] != respB[i] {
			break
		}
		i++
	}

	start := 0
	if i > 100 {
		start = i - 100
	}

	end := i + 100
	if end > limit {
		end = limit
	}

	log.Printf(
		`[DIFF %d/%d] %s %s [status: %d v %d size: %d v %d (%d) time: %dms vs %dms (%d)]
		bytes %d - %d
		######## %s ########
		%s
		######## %s ########
		%s
		####################
		`,
		atomic.LoadInt64(&d.diff),
		atomic.LoadInt64(&d.total),
		req.Method,
		req.RequestURI,
		statusA, statusB,
		sizeA, sizeB, sizeA-sizeB,
		ms(rttA), ms(rttB), ms(rttA-rttB),
		start,
		end,
		d.settings.nameA,
		string(respA[start:end]),
		d.settings.nameB,
		string(respB[start:end]),
	)

	if d.requestsWriter != nil {
		d.outQueue <- raw
	}

}

func ms(d time.Duration) time.Duration {
	return d / time.Millisecond
}
<commit_msg>use mirror of gor's requestfile format for now<commit_after>package main

import (
	"io"
	"log"
	"net/http"
	"sync/atomic"
	"time"

	"github.com/dt/gor_request_files/requestfiles"
)

type DiffReporter struct {
	total int64
	diff  int64

	settings *Settings

	stats     *Stats
	statNames StatNames

	// If writing out diffs, need a queue to serialize to a single writer.
	outQueue       chan []byte
	requestsWriter io.Writer
}

// Compute these once at startup to avoid allocating them every time
type StatNames struct {
	diff  string
	match string
	total string

	errA string
	errB string

	rttA string
	rttB string
}

func NewDiffReporter(s *Settings, stats *Stats) (d *DiffReporter) {
	r := new(DiffReporter)

	r.settings = s

	r.stats = stats

	r.statNames = StatNames{
		total: "diffing.total",
		match: "diffing.match",
		diff:  "diffing.diff",
		errA:  "diffing.err." + s.nameA,
		errB:  "diffing.err." + s.nameB,
		rttA:  "diffing.rtt." + s.nameA,
		rttB:  "diffing.rtt." + s.nameB,
	}

	if s.requestsFile != "" {
		r.outQueue = make(chan []byte, 100)
		r.requestsWriter = requestfiles.NewFileOutput(s.requestsFile)
		go r.writeDiffs()
	}

	return r
}

func (d *DiffReporter) writeDiffs() {
	for {
		req := <-d.outQueue
		d.requestsWriter.Write(req)
	}
}

func isErr(err error, status int) bool {
	return err != nil || status/100 == 5
}

func (d *DiffReporter) Compare(req *http.Request, raw []byte, statusA, statusB int, respA, respB string, rawErrA, rawErrB error, rttA, rttB time.Duration) {
	atomic.AddInt64(&d.total, 1)
	d.stats.Inc(d.statNames.total)

	errA := isErr(rawErrA, statusA)
	errB := isErr(rawErrB, statusB)

	if errA {
		d.stats.Inc(d.statNames.errA)
	} else {
		d.stats.Timing(d.statNames.rttA, rttA)
	}

	if errB {
		d.stats.Inc(d.statNames.errB)
	} else {
		d.stats.Timing(d.statNames.rttB, rttB)
	}

	if (errA && errB) || (d.settings.ignoreErrors && (errA || errB)) {
		return
	}

	if !errA && !errB && respA == respB {
		d.stats.Inc(d.statNames.match)
		return
	}

	atomic.AddInt64(&d.diff, 1)
	d.stats.Inc(d.statNames.diff)

	sizeA := len(respA)
	sizeB := len(respB)

	limit := sizeA
	if sizeA > sizeB {
		limit = sizeB
	}

	i := 0
	for i < limit {
		if respA[i] != respB[i] {
			break
		}
		i++
	}

	start := 0
	if i > 100 {
		start = i - 100
	}

	end := i + 100
	if end > limit {
		end = limit
	}

	log.Printf(
		`[DIFF %d/%d] %s %s [status: %d v %d size: %d v %d (%d) time: %dms vs %dms (%d)]
		bytes %d - %d
		######## %s ########
		%s
		######## %s ########
		%s
		####################
		`,
		atomic.LoadInt64(&d.diff),
		atomic.LoadInt64(&d.total),
		req.Method,
		req.RequestURI,
		statusA, statusB,
		sizeA, sizeB, sizeA-sizeB,
		ms(rttA), ms(rttB), ms(rttA-rttB),
		start,
		end,
		d.settings.nameA,
		string(respA[start:end]),
		d.settings.nameB,
		string(respB[start:end]),
	)

	if d.requestsWriter != nil {
		d.outQueue <- raw
	}

}

func ms(d time.Duration) time.Duration {
	return d / time.Millisecond
}
<|endoftext|>


<commit_before>package parser

//go:generate goyacc -o parser.go parser.y

import (
	"bytes"
	"fmt"
	"os"
	"unicode/utf8"

	"github.com/elpinal/coco3/extra/ast"
	"github.com/elpinal/coco3/extra/token"
	"github.com/elpinal/coco3/extra/types"
)

const eof = 0

type exprLexer struct {
	src   []byte // source
	r     rune   // current character
	errCh chan *ParseError

	// result
	expr *ast.Command

	// information for error messages
	off    uint // start at 0
	line   uint // start at 1
	column uint // start at 1

	// information for current token
	tokLine   uint
	tokColumn uint
}

func newLexer(src []byte) *exprLexer {
	l := &exprLexer{
		src:   src,
		line:  1,
		errCh: make(chan *ParseError),
	}
	l.next()
	return l
}

func isAlphabet(c rune) bool {
	return 'A' <= c && c <= 'Z' || 'a' <= c && c <= 'z'
}

func isIdent(c rune) bool {
	return isAlphabet(c) || isNumber(c) || c == '-'
}

func isNumber(c rune) bool {
	return '0' <= c && c <= '9'
}

func isQuote(c rune) bool {
	return c == '\''
}

func (x *exprLexer) Lex(yylval *yySymType) int {
	for {
		x.tokLine = x.line
		x.tokColumn = x.column
		c := x.r
		switch c {
		case eof:
			return eof
		case ' ', '\n':
			x.next()
		case '\'':
			x.next()
			return x.str(yylval)
		case '[':
			x.next()
			return LBRACK
		case ']':
			x.next()
			return RBRACK
		case ':':
			x.next()
			return COLON
		case ',':
			x.next()
			return COMMA
		default:
			if isAlphabet(c) {
				return x.ident(yylval)
			}
			if isNumber(c) {
				return x.num(yylval)
			}
			fmt.Fprintf(os.Stderr, "%d:%d: invalid character: %[3]U %[3]q\n", x.line, x.column, c)
			return ILLEGAL
		}
	}
}

func (x *exprLexer) ident(yylval *yySymType) int {
	x.takeWhile(types.Ident, isIdent, yylval)
	return IDENT
}

func (x *exprLexer) str(yylval *yySymType) int {
	add := func(b *bytes.Buffer, c rune) {
		if _, err := b.WriteRune(c); err != nil {
			x.errCh <- &ParseError{
				Line:   x.line,
				Column: x.column,
				Msg:    fmt.Sprintf("WriteRune: %s", err),
			}
		}
	}
	var b bytes.Buffer
	for !isQuote(x.r) {
		if x.r == eof {
			x.errCh <- &ParseError{
				Line:   x.tokLine,
				Column: x.tokColumn,
				Msg:    "string literal not terminated: unexpected EOF",
			}
			return STRING
		}
		if x.r == '\\' {
			line := x.line
			column := x.column
			x.next()
			switch x.r {
			case '\'', '\\':
				add(&b, x.r)
				x.next()
			case eof:
				x.errCh <- &ParseError{
					Line:   line,
					Column: column,
					Msg:    "string literal not terminated: unexpected EOF",
				}
				return STRING
			default:
				x.errCh <- &ParseError{
					Line:   line,
					Column: column,
					Msg:    fmt.Sprintf("unknown escape sequence: \\%c", x.r),
				}
				x.next()
				return STRING
			}
			continue
		}
		add(&b, x.r)
		x.next()
	}
	yylval.token = token.Token{
		Kind:   types.String,
		Lit:    b.String(),
		Line:   x.tokLine,
		Column: x.tokColumn,
	}
	x.next()
	return STRING
}

func (x *exprLexer) num(yylval *yySymType) int {
	x.takeWhile(types.Int, isNumber, yylval)
	return NUM
}

func (x *exprLexer) takeWhile(kind types.Type, f func(rune) bool, yylval *yySymType) {
	add := func(b *bytes.Buffer, c rune) {
		if _, err := b.WriteRune(c); err != nil {
			x.errCh <- &ParseError{
				Line:   x.line,
				Column: x.column,
				Msg:    fmt.Sprintf("WriteRune: %s", err),
			}
		}
	}
	var b bytes.Buffer
	for f(x.r) && x.r != eof {
		add(&b, x.r)
		x.next()
	}
	yylval.token = token.Token{
		Kind:   kind,
		Lit:    b.String(),
		Line:   x.tokLine,
		Column: x.tokColumn,
	}
}

func (x *exprLexer) next() {
	if len(x.src) == 0 {
		x.r = eof
		return
	}
	c, size := utf8.DecodeRune(x.src)
	x.src = x.src[size:]
	x.off++
	if c == '\n' {
		x.line++
		x.column = 0
	} else {
		x.column++
	}
	if c == utf8.RuneError && size == 1 {
		x.errCh <- &ParseError{
			Line:   x.line,
			Column: x.column,
			Msg:    "next: invalid utf8",
		}
		x.next()
		return
	}
	x.r = c
}

func (x *exprLexer) Error(s string) {
	x.errCh <- &ParseError{
		Line:   x.tokLine,
		Column: x.tokColumn,
		Msg:    s,
	}
}

func (x *exprLexer) run() <-chan struct{} {
	ch := make(chan struct{})
	go func() {
		yyParse(x)
		ch <- struct{}{}
	}()
	return ch
}

func Parse(src []byte) (*ast.Command, error) {
	l := newLexer(src)
	yyErrorVerbose = true
	ch := l.run()
	select {
	case err := <-l.errCh:
		err.Src = string(src)
		return nil, err
	case <-ch:
	}
	return l.expr, nil
}
<commit_msg>Give the channel better name<commit_after>package parser

//go:generate goyacc -o parser.go parser.y

import (
	"bytes"
	"fmt"
	"os"
	"unicode/utf8"

	"github.com/elpinal/coco3/extra/ast"
	"github.com/elpinal/coco3/extra/token"
	"github.com/elpinal/coco3/extra/types"
)

const eof = 0

type exprLexer struct {
	src   []byte // source
	r     rune   // current character
	errCh chan *ParseError

	// result
	expr *ast.Command

	// information for error messages
	off    uint // start at 0
	line   uint // start at 1
	column uint // start at 1

	// information for current token
	tokLine   uint
	tokColumn uint
}

func newLexer(src []byte) *exprLexer {
	l := &exprLexer{
		src:   src,
		line:  1,
		errCh: make(chan *ParseError),
	}
	l.next()
	return l
}

func isAlphabet(c rune) bool {
	return 'A' <= c && c <= 'Z' || 'a' <= c && c <= 'z'
}

func isIdent(c rune) bool {
	return isAlphabet(c) || isNumber(c) || c == '-'
}

func isNumber(c rune) bool {
	return '0' <= c && c <= '9'
}

func isQuote(c rune) bool {
	return c == '\''
}

func (x *exprLexer) Lex(yylval *yySymType) int {
	for {
		x.tokLine = x.line
		x.tokColumn = x.column
		c := x.r
		switch c {
		case eof:
			return eof
		case ' ', '\n':
			x.next()
		case '\'':
			x.next()
			return x.str(yylval)
		case '[':
			x.next()
			return LBRACK
		case ']':
			x.next()
			return RBRACK
		case ':':
			x.next()
			return COLON
		case ',':
			x.next()
			return COMMA
		default:
			if isAlphabet(c) {
				return x.ident(yylval)
			}
			if isNumber(c) {
				return x.num(yylval)
			}
			fmt.Fprintf(os.Stderr, "%d:%d: invalid character: %[3]U %[3]q\n", x.line, x.column, c)
			return ILLEGAL
		}
	}
}

func (x *exprLexer) ident(yylval *yySymType) int {
	x.takeWhile(types.Ident, isIdent, yylval)
	return IDENT
}

func (x *exprLexer) str(yylval *yySymType) int {
	add := func(b *bytes.Buffer, c rune) {
		if _, err := b.WriteRune(c); err != nil {
			x.errCh <- &ParseError{
				Line:   x.line,
				Column: x.column,
				Msg:    fmt.Sprintf("WriteRune: %s", err),
			}
		}
	}
	var b bytes.Buffer
	for !isQuote(x.r) {
		if x.r == eof {
			x.errCh <- &ParseError{
				Line:   x.tokLine,
				Column: x.tokColumn,
				Msg:    "string literal not terminated: unexpected EOF",
			}
			return STRING
		}
		if x.r == '\\' {
			line := x.line
			column := x.column
			x.next()
			switch x.r {
			case '\'', '\\':
				add(&b, x.r)
				x.next()
			case eof:
				x.errCh <- &ParseError{
					Line:   line,
					Column: column,
					Msg:    "string literal not terminated: unexpected EOF",
				}
				return STRING
			default:
				x.errCh <- &ParseError{
					Line:   line,
					Column: column,
					Msg:    fmt.Sprintf("unknown escape sequence: \\%c", x.r),
				}
				x.next()
				return STRING
			}
			continue
		}
		add(&b, x.r)
		x.next()
	}
	yylval.token = token.Token{
		Kind:   types.String,
		Lit:    b.String(),
		Line:   x.tokLine,
		Column: x.tokColumn,
	}
	x.next()
	return STRING
}

func (x *exprLexer) num(yylval *yySymType) int {
	x.takeWhile(types.Int, isNumber, yylval)
	return NUM
}

func (x *exprLexer) takeWhile(kind types.Type, f func(rune) bool, yylval *yySymType) {
	add := func(b *bytes.Buffer, c rune) {
		if _, err := b.WriteRune(c); err != nil {
			x.errCh <- &ParseError{
				Line:   x.line,
				Column: x.column,
				Msg:    fmt.Sprintf("WriteRune: %s", err),
			}
		}
	}
	var b bytes.Buffer
	for f(x.r) && x.r != eof {
		add(&b, x.r)
		x.next()
	}
	yylval.token = token.Token{
		Kind:   kind,
		Lit:    b.String(),
		Line:   x.tokLine,
		Column: x.tokColumn,
	}
}

func (x *exprLexer) next() {
	if len(x.src) == 0 {
		x.r = eof
		return
	}
	c, size := utf8.DecodeRune(x.src)
	x.src = x.src[size:]
	x.off++
	if c == '\n' {
		x.line++
		x.column = 0
	} else {
		x.column++
	}
	if c == utf8.RuneError && size == 1 {
		x.errCh <- &ParseError{
			Line:   x.line,
			Column: x.column,
			Msg:    "next: invalid utf8",
		}
		x.next()
		return
	}
	x.r = c
}

func (x *exprLexer) Error(s string) {
	x.errCh <- &ParseError{
		Line:   x.tokLine,
		Column: x.tokColumn,
		Msg:    s,
	}
}

func (x *exprLexer) run() <-chan struct{} {
	done := make(chan struct{})
	go func() {
		yyParse(x)
		done <- struct{}{}
	}()
	return done
}

func Parse(src []byte) (*ast.Command, error) {
	l := newLexer(src)
	yyErrorVerbose = true
	done := l.run()
	select {
	case err := <-l.errCh:
		err.Src = string(src)
		return nil, err
	case <-done:
	}
	return l.expr, nil
}
<|endoftext|>


<commit_before>package flag

import (
	"fmt"
	"os"
	"reflect"
)

var envParser = os.Getenv

type resolver struct {
	LastSet *FlagSet
}

func (r *resolver) fromDefault(f *Flag) []string {
	if !isSlicePtr(f.Ptr) {
		return []string{fmt.Sprint(f.Default)}
	}

	refval := reflect.ValueOf(f.Default)
	vals := make([]string, 0, refval.Len())
	for i, l := 0, refval.Len(); i < l; i++ {
		val := fmt.Sprint(refval.Index(i).Interface())
		if val != "" {
			vals = append(vals, val)
		}
	}
	return vals
}

func (r *resolver) fromEnv(f *Flag) []string {
	val := envParser(f.Env)
	if val == "" {
		return nil
	}

	var vals []string
	if isSlicePtr(f.Ptr) {
		vals = splitAndTrimSpace(val, f.ValSep)
	} else {
		vals = []string{val}
	}
	return vals
}

func (r *resolver) applyVals(f *Flag, vals ...string) error {
	for _, val := range vals {
		err := applyValToPtr(f.Names, f.Ptr, val, f.Selects)
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *resolver) applyEnvAndDefault(f *FlagSet, applied map[*Flag]bool) error {
	for i := range f.flags {
		flag := &f.flags[i]
		if applied[flag] {
			continue
		}
		applied[flag] = true

		var vals []string
		if flag.Env != "" {
			vals = r.fromEnv(flag)
		}
		if len(vals) == 0 && flag.Default != nil {
			vals = r.fromDefault(flag)
		}
		err := r.applyVals(flag, vals...)
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *resolver) resolveFlags(f *FlagSet, context []string, args []argument) error {
	var positional []*Flag
	for i := range f.flags {
		if f.flags[i].Names == flagNamePositional {
			positional = append(positional, &f.flags[i])
		}
	}
	var (
		applied = make(map[*Flag]bool)
		flag    *Flag
		err     error

		positionalIndex int
		applyValue      = func(flag *Flag, val string) error {
			applied[flag] = true
			return r.applyVals(flag, val)
		}
		applyLastFlag = func() error {
			if flag == nil {
				return nil
			}
			return newErrorf(errFlagValueNotProvided, "flag value is not provided: %v.%s", context, flag.Names)
		}
		hasFlag = func(args []argument) bool {
			for i := range args {
				if args[i].Type == argumentFlag {
					return true
				}
			}
			return false
		}
		appendNonFlagArg = func(arg argument, args []argument) error {
			if f.self.ArgsPtr == nil || (!f.self.ArgsAnywhere && hasFlag(args[1:])) {
				return newErrorf(errNonFlagValue, "unexpected non-flag value: %v %s", context, arg.Value)
			}
			if positionalIndex < len(positional) {
				err = applyValue(positional[positionalIndex], arg.Value)
				if err != nil {
					return err
				}
				positionalIndex++
				return nil
			}

			slice := *f.self.ArgsPtr
			slice = append(slice, arg.Value)
			*f.self.ArgsPtr = slice
			return nil
		}
	)

	for i, arg := range args {
		switch arg.Type {
		case argumentFlag:
			err = applyLastFlag()
			if err != nil {
				return err
			}

			flag = f.searchFlag(arg.Value)
			if flag == nil {
				return newErrorf(errFlagNotFound, "unsupported flag: %v.%s", context, arg.Value)
			}
			if applied[flag] && !isSlicePtr(flag.Ptr) {
				return newErrorf(errDuplicateFlagParsed, "duplicated flag: %v.%s", context, flag.Names)
			}

			if arg.AttachValid {
				// directly consume flag attached value
				err = applyValue(flag, arg.Attached)
				if err != nil {
					return err
				}
				flag = nil
			} else if isBoolPtr(flag.Ptr) {
				// bool flag should not consume next value to not affect positional or non flag parsing
				err = applyValue(flag, "true")
				if err != nil {
					return err
				}
				flag = nil
			}
		case argumentValue:
			if flag == nil {
				err = appendNonFlagArg(args[i], args[i:])
				if err != nil {
					return err
				}
			} else {
				err = applyValue(flag, arg.Value)
				if err != nil {
					return err
				}
				flag = nil
			}
		default:
			panic("unreachable")
		}
	}
	err = applyLastFlag()
	if err != nil {
		return err
	}
	//if positionalIndex < len(positional) {
	//	var names []string
	//	for i := positionalIndex; i < len(positional); i++ {
	//		names = append(names, positional[i].Arglist)
	//	}
	//	return newErrorf(errPositionalFlagNotProvided, "flag not provided: %v.%v", context, names)
	//}

	return r.applyEnvAndDefault(f, applied)
}

func (r *resolver) resolveSet(f *FlagSet, context []string, args *scanArgs) (lastSubset *FlagSet, err error) {
	context = append(context, f.self.Names)
	err = r.resolveFlags(f, context, args.Flags[1:])
	if err != nil {
		return nil, err
	}
	for sub, subArgs := range args.Sets {
		set := &f.subsets[f.subsetIndexes[sub]]
		err = r.applyVals(&set.self, "true")
		if err != nil {
			return nil, err
		}

		last, err := r.resolveSet(set, context, subArgs)
		if err != nil {
			return nil, err
		}
		if sub == args.FirstSubset {
			lastSubset = last
			if lastSubset == nil {
				lastSubset = set
			}
		}
	}

	if lastSubset == nil {
		lastSubset = f
	}
	return lastSubset, nil
}

func (r *resolver) resolve(f *FlagSet, args *scanArgs) error {
	var err error
	r.LastSet, err = r.resolveSet(f, nil, args)
	return err
}

func (r *resolver) reset(f *FlagSet) {
	if f.self.ArgsPtr != nil {
		*f.self.ArgsPtr = nil
	}
	resetPtrVal(f.self.Ptr)
	for i := range f.flags {
		resetPtrVal(f.flags[i].Ptr)
	}
	for i := range f.subsets {
		r.reset(&f.subsets[i])
	}
}
<commit_msg>fix non-flag argument handling.<commit_after>package flag

import (
	"fmt"
	"os"
	"reflect"
)

var envParser = os.Getenv

type resolver struct {
	LastSet *FlagSet
}

func (r *resolver) fromDefault(f *Flag) []string {
	if !isSlicePtr(f.Ptr) {
		return []string{fmt.Sprint(f.Default)}
	}

	refval := reflect.ValueOf(f.Default)
	vals := make([]string, 0, refval.Len())
	for i, l := 0, refval.Len(); i < l; i++ {
		val := fmt.Sprint(refval.Index(i).Interface())
		if val != "" {
			vals = append(vals, val)
		}
	}
	return vals
}

func (r *resolver) fromEnv(f *Flag) []string {
	val := envParser(f.Env)
	if val == "" {
		return nil
	}

	var vals []string
	if isSlicePtr(f.Ptr) {
		vals = splitAndTrimSpace(val, f.ValSep)
	} else {
		vals = []string{val}
	}
	return vals
}

func (r *resolver) applyVals(f *Flag, vals ...string) error {
	for _, val := range vals {
		err := applyValToPtr(f.Names, f.Ptr, val, f.Selects)
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *resolver) applyEnvAndDefault(f *FlagSet, applied map[*Flag]bool) error {
	for i := range f.flags {
		flag := &f.flags[i]
		if applied[flag] {
			continue
		}
		applied[flag] = true

		var vals []string
		if flag.Env != "" {
			vals = r.fromEnv(flag)
		}
		if len(vals) == 0 && flag.Default != nil {
			vals = r.fromDefault(flag)
		}
		err := r.applyVals(flag, vals...)
		if err != nil {
			return err
		}
	}
	return nil
}

func (r *resolver) resolveFlags(f *FlagSet, context []string, args []argument) error {
	var positional []*Flag
	for i := range f.flags {
		if f.flags[i].Names == flagNamePositional {
			positional = append(positional, &f.flags[i])
		}
	}
	var (
		applied = make(map[*Flag]bool)
		flag    *Flag
		err     error

		positionalIndex int
		applyValue      = func(flag *Flag, val string) error {
			applied[flag] = true
			return r.applyVals(flag, val)
		}
		applyLastFlag = func() error {
			if flag == nil {
				return nil
			}
			return newErrorf(errFlagValueNotProvided, "flag value is not provided: %v.%s", context, flag.Names)
		}
		hasFlag = func(args []argument) bool {
			for i := range args {
				if args[i].Type == argumentFlag {
					return true
				}
			}
			return false
		}
		appendNonFlagArg = func(arg argument, args []argument) error {
			if (positionalIndex >= len(positional) && f.self.ArgsPtr == nil) ||
				(!f.self.ArgsAnywhere && hasFlag(args[1:])) {
				return newErrorf(errNonFlagValue, "unexpected non-flag value: %v %s", context, arg.Value)
			}
			if positionalIndex < len(positional) {
				err = applyValue(positional[positionalIndex], arg.Value)
				if err != nil {
					return err
				}
				positionalIndex++
				return nil
			}

			slice := *f.self.ArgsPtr
			slice = append(slice, arg.Value)
			*f.self.ArgsPtr = slice
			return nil
		}
	)

	for i, arg := range args {
		switch arg.Type {
		case argumentFlag:
			err = applyLastFlag()
			if err != nil {
				return err
			}

			flag = f.searchFlag(arg.Value)
			if flag == nil {
				return newErrorf(errFlagNotFound, "unsupported flag: %v.%s", context, arg.Value)
			}
			if applied[flag] && !isSlicePtr(flag.Ptr) {
				return newErrorf(errDuplicateFlagParsed, "duplicated flag: %v.%s", context, flag.Names)
			}

			if arg.AttachValid {
				// directly consume flag attached value
				err = applyValue(flag, arg.Attached)
				if err != nil {
					return err
				}
				flag = nil
			} else if isBoolPtr(flag.Ptr) {
				// bool flag should not consume next value to not affect positional or non flag parsing
				err = applyValue(flag, "true")
				if err != nil {
					return err
				}
				flag = nil
			}
		case argumentValue:
			if flag == nil {
				err = appendNonFlagArg(args[i], args[i:])
				if err != nil {
					return err
				}
			} else {
				err = applyValue(flag, arg.Value)
				if err != nil {
					return err
				}
				flag = nil
			}
		default:
			panic("unreachable")
		}
	}
	err = applyLastFlag()
	if err != nil {
		return err
	}
	//if positionalIndex < len(positional) {
	//	var names []string
	//	for i := positionalIndex; i < len(positional); i++ {
	//		names = append(names, positional[i].Arglist)
	//	}
	//	return newErrorf(errPositionalFlagNotProvided, "flag not provided: %v.%v", context, names)
	//}

	return r.applyEnvAndDefault(f, applied)
}

func (r *resolver) resolveSet(f *FlagSet, context []string, args *scanArgs) (lastSubset *FlagSet, err error) {
	context = append(context, f.self.Names)
	err = r.resolveFlags(f, context, args.Flags[1:])
	if err != nil {
		return nil, err
	}
	for sub, subArgs := range args.Sets {
		set := &f.subsets[f.subsetIndexes[sub]]
		err = r.applyVals(&set.self, "true")
		if err != nil {
			return nil, err
		}

		last, err := r.resolveSet(set, context, subArgs)
		if err != nil {
			return nil, err
		}
		if sub == args.FirstSubset {
			lastSubset = last
			if lastSubset == nil {
				lastSubset = set
			}
		}
	}

	if lastSubset == nil {
		lastSubset = f
	}
	return lastSubset, nil
}

func (r *resolver) resolve(f *FlagSet, args *scanArgs) error {
	var err error
	r.LastSet, err = r.resolveSet(f, nil, args)
	return err
}

func (r *resolver) reset(f *FlagSet) {
	if f.self.ArgsPtr != nil {
		*f.self.ArgsPtr = nil
	}
	resetPtrVal(f.self.Ptr)
	for i := range f.flags {
		resetPtrVal(f.flags[i].Ptr)
	}
	for i := range f.subsets {
		r.reset(&f.subsets[i])
	}
}
<|endoftext|>


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

// DNS resolver client: see RFC 1035.

package dns

import (
	"os"
	"net"
	"time"
)

type Resolver struct {
	Servers  []string            // servers to use
	Search   []string            // suffixes to append to local name (not implemented)
	Port     string              // what port to use
	Ndots    int                 // number of dots in name to trigger absolute lookup (not implemented)
	Timeout  int                 // seconds before giving up on packet
	Attempts int                 // lost packets before giving up on server
	Tcp      bool                // use TCP
	Mangle   func([]byte) []byte // mangle the packet
	Rtt      map[string]int64    // Store round trip times
	Rrb      int                 // Last used server (for round robin)
}

// Send a query to the nameserver using the res.
func (res *Resolver) Query(q *Msg) (d *Msg, err os.Error) {
	return res.QueryTsig(q, nil)
}

// Send a query to the nameserver using res, but perform TSIG validation.
func (res *Resolver) QueryTsig(q *Msg, tsig *Tsig) (d *Msg, err os.Error) {
	var inb []byte
	in := new(Msg)
	port, err := check(res, q)
	if err != nil {
		return nil, err
	}

	sending, ok := q.Pack()
	if !ok {
		return nil, ErrPack
	}
	if res.Mangle != nil {
		sending = res.Mangle(sending)
	}

	for i := 0; i < len(res.Servers); i++ {
                var d *Conn
		server := res.Servers[i] + ":" + port
		t := time.Nanoseconds()
		if res.Tcp {
			d, err = Dial("tcp", "", server)
			if err != nil {
				continue
			}
		} else {
			d, err = Dial("udp", "", server)
			if err != nil {
				continue
			}
		}
		d.Tsig = tsig
		inb, err = d.Exchange(sending, false)
		if err != nil {
			continue
		}
		in.Unpack(inb) // Discard error.
		res.Rtt[server] = time.Nanoseconds() - t
                d.Close()
		break
	}
	if err != nil {
		return nil, err
	}
	return in, nil
}

// Perform an incoming Ixfr or Axfr. If the message q's question
// section contains an AXFR type an Axfr is performed. If q's question
// section contains an IXFR type an Ixfr is performed.
func (res *Resolver) Xfr(q *Msg, m chan Xfr) {
	res.XfrTsig(q, nil, m)
}

// Perform an incoming Ixfr or Axfr with Tsig validation. If the message
// q's question section contains an AXFR type an Axfr is performed. If q's question
// section contains an IXFR type an Ixfr is performed.
func (res *Resolver) XfrTsig(q *Msg, t *Tsig, m chan Xfr) {
	port, err := check(res, q)
	if err != nil {
		close(m)
		return
	}
	sending, ok := q.Pack()
	if !ok {
		close(m)
		return
	}
	// No defer close(m) as m is closed in d.XfrRead()
Server:
	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
		c, err := net.Dial("tcp", "", server)
		if err != nil {
			continue Server
		}
		d := new(Conn)
		d.TCP = c.(*net.TCPConn)
		d.Addr = d.TCP.RemoteAddr()
		d.Tsig = t

		_, err = d.Write(sending)
		if err != nil {
			continue Server
		}
		d.XfrRead(q, m) // check
	}
	return
}

func check(res *Resolver, q *Msg) (port string, err os.Error) {
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}
	if res.Rtt == nil {
		res.Rtt = make(map[string]int64)
	}
	if q.Id == 0 {
		q.Id = Id()
	}
	return
}
<commit_msg>Kill net package usage<commit_after>// Copyright 2011 Miek Gieben. All rights reserved.
// Use of this source code is governed by a BSD-style
// license that can be found in the LICENSE file.

// DNS resolver client: see RFC 1035.

package dns

import (
	"os"
	"time"
)

type Resolver struct {
	Servers  []string            // servers to use
	Search   []string            // suffixes to append to local name (not implemented)
	Port     string              // what port to use
	Ndots    int                 // number of dots in name to trigger absolute lookup (not implemented)
	Timeout  int                 // seconds before giving up on packet
	Attempts int                 // lost packets before giving up on server
	Tcp      bool                // use TCP
	Mangle   func([]byte) []byte // mangle the packet
	Rtt      map[string]int64    // Store round trip times
	Rrb      int                 // Last used server (for round robin)
}

// Send a query to the nameserver using the res.
func (res *Resolver) Query(q *Msg) (d *Msg, err os.Error) {
	return res.QueryTsig(q, nil)
}

// Send a query to the nameserver using res, but perform TSIG validation.
func (res *Resolver) QueryTsig(q *Msg, tsig *Tsig) (d *Msg, err os.Error) {
	var inb []byte
	in := new(Msg)
	port, err := check(res, q)
	if err != nil {
		return nil, err
	}

	sending, ok := q.Pack()
	if !ok {
		return nil, ErrPack
	}
	if res.Mangle != nil {
		sending = res.Mangle(sending)
	}

	for i := 0; i < len(res.Servers); i++ {
                var d *Conn
		server := res.Servers[i] + ":" + port
		t := time.Nanoseconds()
		if res.Tcp {
			d, err = Dial("tcp", "", server)
			if err != nil {
				continue
			}
		} else {
			d, err = Dial("udp", "", server)
			if err != nil {
				continue
			}
		}
		d.Tsig = tsig
		inb, err = d.Exchange(sending, false)
		if err != nil {
			continue
		}
		in.Unpack(inb) // Discard error.
		res.Rtt[server] = time.Nanoseconds() - t
                d.Close()
		break
	}
	if err != nil {
		return nil, err
	}
	return in, nil
}

// Perform an incoming Ixfr or Axfr. If the message q's question
// section contains an AXFR type an Axfr is performed. If q's question
// section contains an IXFR type an Ixfr is performed.
func (res *Resolver) Xfr(q *Msg, m chan Xfr) {
	res.XfrTsig(q, nil, m)
}

// Perform an incoming Ixfr or Axfr with Tsig validation. If the message
// q's question section contains an AXFR type an Axfr is performed. If q's question
// section contains an IXFR type an Ixfr is performed.
func (res *Resolver) XfrTsig(q *Msg, t *Tsig, m chan Xfr) {
	port, err := check(res, q)
	if err != nil {
		close(m)
		return
	}
	sending, ok := q.Pack()
	if !ok {
		close(m)
		return
	}
	// No defer close(m) as m is closed in d.XfrRead()
Server:
	for i := 0; i < len(res.Servers); i++ {
		server := res.Servers[i] + ":" + port
                d, err := Dial("tcp", "", server)
		if err != nil {
			continue Server
		}
		d.Tsig = t

		_, err = d.Write(sending)
		if err != nil {
			continue Server
		}
		d.XfrRead(q, m) // check
	}
	return
}

func check(res *Resolver, q *Msg) (port string, err os.Error) {
	if res.Port == "" {
		port = "53"
	} else {
		port = res.Port
	}
	if res.Rtt == nil {
		res.Rtt = make(map[string]int64)
	}
	if q.Id == 0 {
		q.Id = Id()
	}
	return
}
<|endoftext|>


<commit_before>// Package solvere provides an implementation of a recursive, validating,
// DNSSEC aware DNS resolver. It also provides a basic question and answer
// cache implementation.
package solvere

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	mrand "math/rand"
	"net"
	"os"
	"strings"
	"time"

	"github.com/miekg/dns"
	"golang.org/x/net/context"
)

func init() {
	// Initialize math/rand with 8 bytes from crypto/rand so the randomness
	// isn't _super_ terrible
	b := [8]byte{}
	_, err := rand.Read(b[:])
	if err != nil {
		fmt.Fprintf(os.Stderr, "solvere: Failed to read bytes for PSRNG initialization: %s\n", err)
		return
	}
	i, err := binary.ReadVarint(bytes.NewBuffer(b[:]))
	if err != nil {
		fmt.Fprintf(os.Stderr, "solvere: Failed to read bytes for PSRNG initialization: %s\n", err)
		return
	}
	mrand.Seed(i)
}

var (
	// MaxReferrals is the maximum number of referral responses before failing
	MaxReferrals = 10

	dnsPort = "53"

	ErrTooManyReferrals   = errors.New("solvere: Too many referrals")
	ErrNoNSAuthorties     = errors.New("solvere: No NS authority records found")
	ErrNoAuthorityAddress = errors.New("solvere: No A/AAAA records found for the chosen authority")
)

// Question represents a DNS IN question
type Question struct {
	Name string
	Type uint16
}

// QueryLog describes a query to a upstream nameserver
type QueryLog struct {
	Query       *Question
	Rcode       int
	CacheHit    bool `json:",omitempty"`
	DNSSECValid bool
	Latency     time.Duration
	Error       string `json:",omitempty"`
	Truncated   bool   `json:",omitempty"`
	Referral    bool   `json:",omitempty"`

	// Only present if CacheHit == false
	NS *Nameserver `json:",omitempty"`

	Composites []*QueryLog `json:",omitempty"`
}

// LookupLog describes a iterative resolution
type LookupLog struct {
	Query       *Question
	DNSSECValid bool
	Started     time.Time
	Latency     time.Duration
	Rcode       uint16

	Composites []*QueryLog
}

// Answer contains the answer to a iterative resolution performed
// by RecursiveResolver.Lookup
type Answer struct {
	Answer        []dns.RR
	Authority     []dns.RR
	Additional    []dns.RR
	Rcode         int
	Authenticated bool
}

// Nameserver describes a upstream authoritative nameserver
type Nameserver struct {
	Name string
	Addr string
	Zone string
}

// RecursiveResolver defines the parameters for running a recursive resolver
type RecursiveResolver struct {
	useIPv6   bool
	useDNSSEC bool

	c *dns.Client

	cache           QuestionAnswerCache
	rootNameservers []Nameserver
}

// NewRecursiveResolver returns an initialized RecursiveResolver. If cache is nil
// answers won't be cached.
func NewRecursiveResolver(useIPv6 bool, useDNSSEC bool, rootHints []dns.RR, rootKeys []dns.RR, cache QuestionAnswerCache) *RecursiveResolver {
	rr := &RecursiveResolver{
		useIPv6:   useIPv6,
		useDNSSEC: useDNSSEC,
		c:         new(dns.Client),
		cache:     cache,
	}
	// Initialize root nameservers
	addrs := extractRRSet(rootHints, "", dns.TypeA)
	if useIPv6 {
		addrs = append(addrs, extractRRSet(rootHints, "", dns.TypeAAAA)...)
	}
	for _, a := range addrs {
		switch r := a.(type) {
		case *dns.A:
			rr.rootNameservers = append(rr.rootNameservers, Nameserver{a.Header().Name, r.A.String(), "."})
		case *dns.AAAA:
			rr.rootNameservers = append(rr.rootNameservers, Nameserver{a.Header().Name, r.AAAA.String(), "."})
		}
	}
	// Add root DNSSEC keys to cache indefinitely
	if rr.cache != nil {
		rr.cache.Add(&Question{Name: ".", Type: dns.TypeDNSKEY}, &Answer{rootKeys, nil, nil, dns.RcodeSuccess, true}, true)
	}
	return rr
}

func (rr *RecursiveResolver) query(ctx context.Context, q *Question, auth *Nameserver) (*dns.Msg, *QueryLog, error) {
	ql := &QueryLog{Query: q, NS: auth}
	s := time.Now()
	defer func() { ql.Latency = time.Since(s) }()
	m := new(dns.Msg)
	m.SetEdns0(4096, rr.useDNSSEC)
	m.Question = []dns.Question{{Name: q.Name, Qtype: q.Type, Qclass: dns.ClassINET}}
	if rr.cache != nil {
		if answer := rr.cache.Get(q); answer != nil {
			m.Rcode = dns.RcodeSuccess
			m.Answer = answer.Answer
			m.Ns = answer.Authority
			m.Extra = answer.Additional
			ql.CacheHit = true
			ql.NS = nil
			ql.DNSSECValid = answer.Authenticated
			ql.Rcode = dns.RcodeSuccess
			return m, ql, nil
		}
	}
	r, _, err := rr.c.Exchange(m, net.JoinHostPort(auth.Addr, dnsPort))
	if err != nil {
		return nil, ql, err
	}
	ql.Rcode = r.Rcode

	// check all returned records are in-bailiwick, ignore extra section?
	for _, section := range [][]dns.RR{r.Answer, r.Ns} {
		for _, record := range section {
			if record.Header().Rrtype != dns.TypeOPT && !strings.HasSuffix(record.Header().Name, auth.Zone) {
				return nil, ql, errors.New("Out of bailiwick record in message") // or just strip invalid records...?
			}
		}
	}
	return r, ql, nil
}

func (rr *RecursiveResolver) lookupNS(ctx context.Context, name string) (*Nameserver, error) {
	// BUG(roland): LookupLog for lookupNS calls isn't included in parent chain
	// BUG(roland): There is no maximum depth to Lookup -> lookupNS -> Lookup calls, looping is possible
	// BUG(roland): I'm not sure how the lookup of a NS addr should be taken into account in terms of the
	//              dnssec chain
	r, _, err := rr.Lookup(ctx, Question{Name: name, Type: dns.TypeA})
	if err != nil {
		return nil, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("Authority lookup failed for %s: %s", name, dns.RcodeToString[r.Rcode])
	}
	if len(r.Answer) == 0 {
		return nil, ErrNoAuthorityAddress
	}
	addresses := extractRRSet(r.Answer, name, dns.TypeA)
	if len(addresses) == 0 {
		return nil, ErrNoAuthorityAddress
	}
	return &Nameserver{Name: name, Addr: addresses[mrand.Intn(len(addresses))].(*dns.A).A.String()}, nil // ewwww
}

func splitAuthsByZone(auths []dns.RR, extras []dns.RR, useIPv6 bool) (map[string][]string, map[string]int, map[string]string) {
	zones := make(map[string][]string)
	minTTLs := make(map[string]int)
	nsToZone := make(map[string]string)

	for _, rr := range auths {
		if rr.Header().Rrtype == dns.TypeNS {
			ns := rr.(*dns.NS)
			nsToZone[ns.Ns] = rr.Header().Name
		}
	}

	for _, rr := range extras {
		zone, present := nsToZone[rr.Header().Name]
		if present && (rr.Header().Rrtype == dns.TypeA || (useIPv6 && rr.Header().Rrtype == dns.TypeAAAA)) {
			switch a := rr.(type) {
			case *dns.A:
				zones[zone] = append(zones[zone], a.A.String())
			case *dns.AAAA:
				if useIPv6 {
					zones[zone] = append(zones[zone], a.AAAA.String())
				}
			}
			if minTTLs[zone] == 0 || int(rr.Header().Ttl) < minTTLs[zone] {
				minTTLs[zone] = int(rr.Header().Ttl)
			}
		}
	}

	return zones, minTTLs, nsToZone
}

func (rr *RecursiveResolver) pickAuthority(ctx context.Context, auths []dns.RR, extras []dns.RR) (*Nameserver, error) {
	zones, _, nsToZone := splitAuthsByZone(auths, extras, rr.useIPv6)
	if len(zones) == 0 {
		if len(nsToZone) == 0 {
			return nil, ErrNoNSAuthorties
		}
		var ns, z string
		// abuse how ranging over maps works to select a 'random' element
		for ns, z = range nsToZone {
			break
		}
		a, err := rr.lookupNS(ctx, ns)
		if err != nil {
			return nil, err
		}
		a.Zone = z
		return a, nil
	}
	// abuse how ranging over maps works to select a 'random' element
	for ns, z := range nsToZone {
		if len(zones[z]) > 0 {
			return &Nameserver{ns, zones[z][mrand.Intn(len(zones[z]))], z}, nil
		}
	}
	return nil, ErrNoNSAuthorties
}

func extractAnswer(m *dns.Msg, authenticated bool) *Answer {
	return &Answer{
		Answer:        m.Answer,
		Authority:     m.Ns,
		Additional:    m.Extra,
		Rcode:         m.Rcode,
		Authenticated: authenticated,
	}
}

// Lookup a Question iteratively. All upstream responses are validated
// and a DNSSEC chain is built if the RecursiveResolver was initialized to do so.
// If responses are found in the underlying cache they will be used instead of
// sending messages to remote nameservers.
func (rr *RecursiveResolver) Lookup(ctx context.Context, q Question) (*Answer, *LookupLog, error) {
	ll := &LookupLog{Query: &q, Started: time.Now()}

	authority := &rr.rootNameservers[mrand.Intn(len(rr.rootNameservers))]

	defer func() {
		ll.Latency = time.Since(ll.Started)
	}()

	var parentDSSet []dns.RR
	for i := 0; i < MaxReferrals; i++ {
		r, log, err := rr.query(ctx, &q, authority)
		ll.Composites = append(ll.Composites, log)
		if err != nil && err != dns.ErrTruncated { // if truncated still try...
			log.Error = err.Error()
			return nil, ll, err
		} else if err == dns.ErrTruncated {
			log.Truncated = true
		}

		// validate
		validated := false
		if log.CacheHit {
			validated = log.DNSSECValid
		}
		if (i == 0 || len(parentDSSet) > 0) && !log.CacheHit {
			dkLog, err := rr.checkDNSKEY(ctx, r, authority, parentDSSet)
			log.Composites = append(log.Composites, dkLog)
			if err != nil {
				log.Error = err.Error()
				return nil, ll, err
			}
			validated = true
		}
		log.DNSSECValid = validated
		ll.DNSSECValid = validated

		if r.Rcode != dns.RcodeSuccess {
			// BUG(roland): This should be moved into checkDNSKEY (which needs a better name...)
			if r.Rcode == dns.RcodeNameError {
				nsecSet := extractRRSet(r.Ns, "", dns.TypeNSEC3)
				if len(nsecSet) != 0 {
					err = verifyNameError(&q, nsecSet)
					if err != nil {
						log.Error = err.Error()
						return nil, ll, err
					}
				}
			}
			return extractAnswer(r, validated), ll, nil
		}

		// good response
		if len(r.Answer) > 0 {
			// BUG(roland): if after stripping dnssec records the only remaining record is a CNAME and the
			//              question type wasn't CNAME then the alias should be chased
			if !log.CacheHit {
				if rr.cache != nil {
					go rr.cache.Add(&q, &Answer{r.Answer, r.Ns, r.Extra, r.Rcode, validated}, false)
				}
			}

			// check for alias and chase or do that in RecursiveResolver.query?
			return extractAnswer(r, validated), ll, nil
		}

		// NODATA validation
		// BUG(roland): This catches referrals :/
		if len(r.Answer) == 0 {
			// BUG(roland): This should be moved into checkDNSKEY maybe?
			if len(r.Ns) > 0 {
				nsecSet := extractRRSet(r.Ns, "", dns.TypeNSEC3)
				if len(nsecSet) != 0 {
					err = verifyNODATA(&q, nsecSet)
					if err != nil {
						log.Error = err.Error()
						return nil, ll, err
					}
				}
			}
		}

		// referral response
		if len(r.Ns) > 0 {
			log.Referral = true
			// BUG(roland): NSEC# DS delegation denials aren't checked
			authority, err = rr.pickAuthority(ctx, r.Ns, r.Extra)
			if err != nil {
				log.Error = err.Error()
				return nil, ll, err
			}
			if i == 0 || len(parentDSSet) > 0 {
				parentDSSet = extractRRSet(r.Ns, authority.Zone, dns.TypeDS)
			}
			continue
		}

		// useless response...
		return nil, ll, errors.New("No authority or additional records! IDK") // ???
	}
	return nil, ll, ErrTooManyReferrals
}

func extractRRSet(in []dns.RR, name string, t ...uint16) []dns.RR {
	out := []dns.RR{}
	tMap := make(map[uint16]struct{}, len(t))
	for _, t := range t {
		tMap[t] = struct{}{}
	}
	for _, r := range in {
		if _, present := tMap[r.Header().Rrtype]; present {
			if name != "" && name != r.Header().Name {
				continue
			}
			out = append(out, r)
		}
	}
	return out
}

func extractAndMapRRSet(in []dns.RR, name string, t ...uint16) map[uint16][]dns.RR {
	out := make(map[uint16][]dns.RR, len(t))
	for _, rt := range t {
		out[rt] = []dns.RR{}
	}
	for _, r := range in {
		rt := r.Header().Rrtype
		if _, present := out[rt]; !present {
			continue
		}
		if name != "" && name != r.Header().Name {
			continue
		}
		out[rt] = append(out[rt], r)
	}
	return out
}
<commit_msg>Juggle things around a bit<commit_after>// Package solvere provides an implementation of a recursive, validating,
// DNSSEC aware DNS resolver. It also provides a basic question and answer
// cache implementation.
package solvere

import (
	"bytes"
	"crypto/rand"
	"encoding/binary"
	"errors"
	"fmt"
	mrand "math/rand"
	"net"
	"os"
	"strings"
	"time"

	"github.com/miekg/dns"
	"golang.org/x/net/context"
)

func init() {
	// Initialize math/rand with 8 bytes from crypto/rand so the randomness
	// isn't _super_ terrible
	b := [8]byte{}
	_, err := rand.Read(b[:])
	if err != nil {
		fmt.Fprintf(os.Stderr, "solvere: Failed to read bytes for PSRNG initialization: %s\n", err)
		return
	}
	i, err := binary.ReadVarint(bytes.NewBuffer(b[:]))
	if err != nil {
		fmt.Fprintf(os.Stderr, "solvere: Failed to read bytes for PSRNG initialization: %s\n", err)
		return
	}
	mrand.Seed(i)
}

var (
	// MaxReferrals is the maximum number of referral responses before failing
	MaxReferrals = 10

	dnsPort = "53"

	ErrTooManyReferrals   = errors.New("solvere: Too many referrals")
	ErrNoNSAuthorties     = errors.New("solvere: No NS authority records found")
	ErrNoAuthorityAddress = errors.New("solvere: No A/AAAA records found for the chosen authority")
)

// Question represents a DNS IN question
type Question struct {
	Name string
	Type uint16
}

// QueryLog describes a query to a upstream nameserver
type QueryLog struct {
	Query       *Question
	Rcode       int
	CacheHit    bool `json:",omitempty"`
	DNSSECValid bool
	Latency     time.Duration
	Error       string `json:",omitempty"`
	Truncated   bool   `json:",omitempty"`
	Referral    bool   `json:",omitempty"`

	// Only present if CacheHit == false
	NS *Nameserver `json:",omitempty"`

	Composites []*QueryLog `json:",omitempty"`
}

// LookupLog describes a iterative resolution
type LookupLog struct {
	Query       *Question
	DNSSECValid bool
	Started     time.Time
	Latency     time.Duration
	Rcode       uint16

	Composites []*QueryLog
}

// Answer contains the answer to a iterative resolution performed
// by RecursiveResolver.Lookup
type Answer struct {
	Answer        []dns.RR
	Authority     []dns.RR
	Additional    []dns.RR
	Rcode         int
	Authenticated bool
}

// Nameserver describes a upstream authoritative nameserver
type Nameserver struct {
	Name string
	Addr string
	Zone string
}

// RecursiveResolver defines the parameters for running a recursive resolver
type RecursiveResolver struct {
	useIPv6   bool
	useDNSSEC bool

	c *dns.Client

	cache           QuestionAnswerCache
	rootNameservers []Nameserver
}

// NewRecursiveResolver returns an initialized RecursiveResolver. If cache is nil
// answers won't be cached.
func NewRecursiveResolver(useIPv6 bool, useDNSSEC bool, rootHints []dns.RR, rootKeys []dns.RR, cache QuestionAnswerCache) *RecursiveResolver {
	rr := &RecursiveResolver{
		useIPv6:   useIPv6,
		useDNSSEC: useDNSSEC,
		c:         new(dns.Client),
		cache:     cache,
	}
	// Initialize root nameservers
	addrs := extractRRSet(rootHints, "", dns.TypeA)
	if useIPv6 {
		addrs = append(addrs, extractRRSet(rootHints, "", dns.TypeAAAA)...)
	}
	for _, a := range addrs {
		switch r := a.(type) {
		case *dns.A:
			rr.rootNameservers = append(rr.rootNameservers, Nameserver{a.Header().Name, r.A.String(), "."})
		case *dns.AAAA:
			rr.rootNameservers = append(rr.rootNameservers, Nameserver{a.Header().Name, r.AAAA.String(), "."})
		}
	}
	// Add root DNSSEC keys to cache indefinitely
	if rr.cache != nil {
		rr.cache.Add(&Question{Name: ".", Type: dns.TypeDNSKEY}, &Answer{rootKeys, nil, nil, dns.RcodeSuccess, true}, true)
	}
	return rr
}

func (rr *RecursiveResolver) query(ctx context.Context, q *Question, auth *Nameserver) (*dns.Msg, *QueryLog, error) {
	ql := &QueryLog{Query: q, NS: auth}
	s := time.Now()
	defer func() { ql.Latency = time.Since(s) }()
	m := new(dns.Msg)
	m.SetEdns0(4096, rr.useDNSSEC)
	m.Question = []dns.Question{{Name: q.Name, Qtype: q.Type, Qclass: dns.ClassINET}}
	if rr.cache != nil {
		if answer := rr.cache.Get(q); answer != nil {
			m.Rcode = dns.RcodeSuccess
			m.Answer = answer.Answer
			m.Ns = answer.Authority
			m.Extra = answer.Additional
			ql.CacheHit = true
			ql.NS = nil
			ql.DNSSECValid = answer.Authenticated
			ql.Rcode = dns.RcodeSuccess
			return m, ql, nil
		}
	}
	r, _, err := rr.c.Exchange(m, net.JoinHostPort(auth.Addr, dnsPort))
	if err != nil {
		return nil, ql, err
	}
	ql.Rcode = r.Rcode

	// check all returned records are in-bailiwick, ignore extra section?
	for _, section := range [][]dns.RR{r.Answer, r.Ns} {
		for _, record := range section {
			if record.Header().Rrtype != dns.TypeOPT && !strings.HasSuffix(record.Header().Name, auth.Zone) {
				return nil, ql, errors.New("Out of bailiwick record in message") // or just strip invalid records...?
			}
		}
	}
	return r, ql, nil
}

func (rr *RecursiveResolver) lookupNS(ctx context.Context, name string) (*Nameserver, error) {
	// BUG(roland): LookupLog for lookupNS calls isn't included in parent chain
	// BUG(roland): There is no maximum depth to Lookup -> lookupNS -> Lookup calls, looping is possible
	// BUG(roland): I'm not sure how the lookup of a NS addr should be taken into account in terms of the
	//              dnssec chain
	r, _, err := rr.Lookup(ctx, Question{Name: name, Type: dns.TypeA})
	if err != nil {
		return nil, err
	}
	if r.Rcode != dns.RcodeSuccess {
		return nil, fmt.Errorf("Authority lookup failed for %s: %s", name, dns.RcodeToString[r.Rcode])
	}
	if len(r.Answer) == 0 {
		return nil, ErrNoAuthorityAddress
	}
	addresses := extractRRSet(r.Answer, name, dns.TypeA)
	if len(addresses) == 0 {
		return nil, ErrNoAuthorityAddress
	}
	return &Nameserver{Name: name, Addr: addresses[mrand.Intn(len(addresses))].(*dns.A).A.String()}, nil // ewwww
}

func splitAuthsByZone(auths []dns.RR, extras []dns.RR, useIPv6 bool) (map[string][]string, map[string]int, map[string]string) {
	zones := make(map[string][]string)
	minTTLs := make(map[string]int)
	nsToZone := make(map[string]string)

	for _, rr := range auths {
		if rr.Header().Rrtype == dns.TypeNS {
			ns := rr.(*dns.NS)
			nsToZone[ns.Ns] = rr.Header().Name
		}
	}

	for _, rr := range extras {
		zone, present := nsToZone[rr.Header().Name]
		if present && (rr.Header().Rrtype == dns.TypeA || (useIPv6 && rr.Header().Rrtype == dns.TypeAAAA)) {
			switch a := rr.(type) {
			case *dns.A:
				zones[zone] = append(zones[zone], a.A.String())
			case *dns.AAAA:
				if useIPv6 {
					zones[zone] = append(zones[zone], a.AAAA.String())
				}
			}
			if minTTLs[zone] == 0 || int(rr.Header().Ttl) < minTTLs[zone] {
				minTTLs[zone] = int(rr.Header().Ttl)
			}
		}
	}

	return zones, minTTLs, nsToZone
}

func (rr *RecursiveResolver) pickAuthority(ctx context.Context, auths []dns.RR, extras []dns.RR) (*Nameserver, error) {
	zones, _, nsToZone := splitAuthsByZone(auths, extras, rr.useIPv6)
	if len(zones) == 0 {
		if len(nsToZone) == 0 {
			return nil, ErrNoNSAuthorties
		}
		var ns, z string
		// abuse how ranging over maps works to select a 'random' element
		for ns, z = range nsToZone {
			break
		}
		a, err := rr.lookupNS(ctx, ns)
		if err != nil {
			return nil, err
		}
		a.Zone = z
		return a, nil
	}
	// abuse how ranging over maps works to select a 'random' element
	for ns, z := range nsToZone {
		if len(zones[z]) > 0 {
			return &Nameserver{ns, zones[z][mrand.Intn(len(zones[z]))], z}, nil
		}
	}
	return nil, ErrNoNSAuthorties
}

func extractAnswer(m *dns.Msg, authenticated bool) *Answer {
	return &Answer{
		Answer:        m.Answer,
		Authority:     m.Ns,
		Additional:    m.Extra,
		Rcode:         m.Rcode,
		Authenticated: authenticated,
	}
}

// Lookup a Question iteratively. All upstream responses are validated
// and a DNSSEC chain is built if the RecursiveResolver was initialized to do so.
// If responses are found in the underlying cache they will be used instead of
// sending messages to remote nameservers.
func (rr *RecursiveResolver) Lookup(ctx context.Context, q Question) (*Answer, *LookupLog, error) {
	ll := &LookupLog{Query: &q, Started: time.Now()}

	authority := &rr.rootNameservers[mrand.Intn(len(rr.rootNameservers))]

	defer func() {
		ll.Latency = time.Since(ll.Started)
	}()

	var parentDSSet []dns.RR
	for i := 0; i < MaxReferrals; i++ {
		r, log, err := rr.query(ctx, &q, authority)
		ll.Composites = append(ll.Composites, log)
		if err != nil && err != dns.ErrTruncated { // if truncated still try...
			log.Error = err.Error()
			return nil, ll, err
		} else if err == dns.ErrTruncated {
			log.Truncated = true
		}

		// validate
		validated := false
		if log.CacheHit {
			validated = log.DNSSECValid
		}
		if (i == 0 || len(parentDSSet) > 0) && !log.CacheHit {
			dkLog, err := rr.checkDNSKEY(ctx, r, authority, parentDSSet)
			log.Composites = append(log.Composites, dkLog)
			if err != nil {
				log.Error = err.Error()
				return nil, ll, err
			}
			validated = true
		}
		log.DNSSECValid = validated
		ll.DNSSECValid = validated

		if r.Rcode != dns.RcodeSuccess {
			// BUG(roland): This should be moved into checkDNSKEY (which needs a better name...)
			if r.Rcode == dns.RcodeNameError {
				nsecSet := extractRRSet(r.Ns, "", dns.TypeNSEC3)
				if len(nsecSet) != 0 {
					err = verifyNameError(&q, nsecSet)
					if err != nil {
						log.Error = err.Error()
						return nil, ll, err
					}
				}
			}
			return extractAnswer(r, validated), ll, nil
		}

		// good response
		if len(r.Answer) > 0 {
			// BUG(roland): if after stripping dnssec records the only remaining record is a CNAME and the
			//              question type wasn't CNAME then the alias should be chased
			if !log.CacheHit {
				if rr.cache != nil {
					go rr.cache.Add(&q, &Answer{r.Answer, r.Ns, r.Extra, r.Rcode, validated}, false)
				}
			}

			// check for alias and chase or do that in RecursiveResolver.query?
			return extractAnswer(r, validated), ll, nil
		}

		// NODATA or referral response
		nsecSet := extractRRSet(r.Ns, "", dns.TypeNSEC3)
		if len(r.Ns) == 0 {
			// NODATA response
			if len(nsecSet) != 0 { // if the zone is signed and this is missing its a failure...
				err = verifyNODATA(&q, nsecSet)
				if err != nil {
					log.Error = err.Error()
					log.DNSSECValid = false
					ll.DNSSECValid = false
					return nil, ll, err
				}
			}
			return &Answer{Rcode: dns.RcodeSuccess, Authenticated: validated}, ll, nil // ignore anything in additional section
		}

		// Referral response
		log.Referral = true
		// BUG(roland): NSEC# DS delegation denials aren't checked
		authority, err = rr.pickAuthority(ctx, r.Ns, r.Extra)
		if err != nil {
			log.Error = err.Error()
			return nil, ll, err
		}
		if len(nsecSet) != 0 { // if the zone is signed and this is missing its a failure...
			err = verifyDelegation(authority.Zone, nsecSet)
			if err != nil {
				log.Error = err.Error()
				log.DNSSECValid = false
				ll.DNSSECValid = false
				return nil, ll, err
			}
		}
		if i == 0 || len(parentDSSet) > 0 {
			parentDSSet = extractRRSet(r.Ns, authority.Zone, dns.TypeDS)
		}
	}
	return nil, ll, ErrTooManyReferrals
}

func extractRRSet(in []dns.RR, name string, t ...uint16) []dns.RR {
	out := []dns.RR{}
	tMap := make(map[uint16]struct{}, len(t))
	for _, t := range t {
		tMap[t] = struct{}{}
	}
	for _, r := range in {
		if _, present := tMap[r.Header().Rrtype]; present {
			if name != "" && name != r.Header().Name {
				continue
			}
			out = append(out, r)
		}
	}
	return out
}

func extractAndMapRRSet(in []dns.RR, name string, t ...uint16) map[uint16][]dns.RR {
	out := make(map[uint16][]dns.RR, len(t))
	for _, rt := range t {
		out[rt] = []dns.RR{}
	}
	for _, r := range in {
		rt := r.Header().Rrtype
		if _, present := out[rt]; !present {
			continue
		}
		if name != "" && name != r.Header().Name {
			continue
		}
		out[rt] = append(out[rt], r)
	}
	return out
}
<|endoftext|>


<commit_before>package milter

/* Response represents a response structure returned by callback
   handlers to indicate how the milter server should proceed */
type Response interface {
	Response() *Message
	Continue() bool
}

/* SimpleResponse interface to define list of pre-defined responses */
type SimpleResponse byte

/* Response returns a Message object reference */
func (r SimpleResponse) Response() *Message {
	return &Message{byte(r), nil}
}

/* Continue to process milter messages only if current code is Continue */
func (r SimpleResponse) Continue() bool {
	return byte(r) == Continue
}

/* Define standard responses with no data */
const (
	RespAccept   = SimpleResponse(Accept)
	RespContinue = SimpleResponse(Continue)
	RespDiscard  = SimpleResponse(Discard)
	RespReject   = SimpleResponse(Reject)
	RespTempFail = SimpleResponse(TempFail)
)

/* CustomResponse is a response object used by callback handlers to indicate how the milter should continue processing of current message */
type CustomResponse struct {
	Code byte
	Data []byte
}

/* Response returns message object with data */
func (c *CustomResponse) Response() *Message {
	return &Message{c.Code, c.Data}
}

/* Continue returns false if milter chain should be stopped, true otherwise */
func (c *CustomResponse) Continue() bool {
	for _, q := range []byte{Accept, Discard, Reject, TempFail} {
		if c.Code == q {
			return false
		}
	}
	return true
}

/* NewResponse generates a new CustomRespanse suitable for WritePacket */
func NewResponse(code byte, data []byte) *CustomResponse {
	return &CustomResponse{code, data}
}

/* NewResponseStr generates a new CustomResponse with string payload */
func NewResponseStr(code byte, data string) *CustomResponse {
	return NewResponse(code, EncodeCStrings([]string{data}))
}
<commit_msg>directly convert golang string to a null terminated string<commit_after>package milter

/* Response represents a response structure returned by callback
   handlers to indicate how the milter server should proceed */
type Response interface {
	Response() *Message
	Continue() bool
}

/* SimpleResponse interface to define list of pre-defined responses */
type SimpleResponse byte

/* Response returns a Message object reference */
func (r SimpleResponse) Response() *Message {
	return &Message{byte(r), nil}
}

/* Continue to process milter messages only if current code is Continue */
func (r SimpleResponse) Continue() bool {
	return byte(r) == Continue
}

/* Define standard responses with no data */
const (
	RespAccept   = SimpleResponse(Accept)
	RespContinue = SimpleResponse(Continue)
	RespDiscard  = SimpleResponse(Discard)
	RespReject   = SimpleResponse(Reject)
	RespTempFail = SimpleResponse(TempFail)
)

/* CustomResponse is a response object used by callback handlers to indicate how the milter should continue processing of current message */
type CustomResponse struct {
	Code byte
	Data []byte
}

/* Response returns message object with data */
func (c *CustomResponse) Response() *Message {
	return &Message{c.Code, c.Data}
}

/* Continue returns false if milter chain should be stopped, true otherwise */
func (c *CustomResponse) Continue() bool {
	for _, q := range []byte{Accept, Discard, Reject, TempFail} {
		if c.Code == q {
			return false
		}
	}
	return true
}

/* NewResponse generates a new CustomRespanse suitable for WritePacket */
func NewResponse(code byte, data []byte) *CustomResponse {
	return &CustomResponse{code, data}
}

/* NewResponseStr generates a new CustomResponse with string payload */
func NewResponseStr(code byte, data string) *CustomResponse {
	return NewResponse(code, []byte(data+NULL))
}
<|endoftext|>