<commit_before>// main simply contains the primary web serving code that allows peers to
// register and unregister as give mode peers running within the Lantern
// network
package main

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"github.com/getlantern/keyman"
	"github.com/getlantern/tlsdefaults"
)

const (
	PKFile   = "pk.pem"
	CertFile = "cert.pem"
)

func startHttp() {
	http.HandleFunc("/register", register)
	http.HandleFunc("/unregister", unregister)
	laddr := fmt.Sprintf(":%d", *port)

	tlsConfig := tlsdefaults.Server()
	_, _, err := keyman.StoredPKAndCert(PKFile, CertFile, "Lantern", "localhost")
	if err != nil {
		log.Fatalf("Unable to initialize private key and certificate: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(CertFile, PKFile)
	if err != nil {
		log.Fatalf("Unable to load certificate and key from %s and %s: %s", CertFile, PKFile, err)
	}
	tlsConfig.Certificates = []tls.Certificate{cert}

	log.Debugf("About to listen at %v", laddr)
	l, err := tls.Listen("tcp", laddr, tlsConfig)
	if err != nil {
		log.Fatalf("Unable to listen for tls connections at %s: %s", laddr, err)
	}

	log.Debug("About to serve")
	err = http.Serve(l, nil)
	if err != nil {
		log.Fatalf("Unable to serve: %s", err)
	}
}

// register is the entry point for peers registering themselves with the service.
// If peers are successfully vetted, they'll be added to the DNS round robin.
func register(resp http.ResponseWriter, req *http.Request) {
	name, ip, port, err := getHostInfo(req)
	if err == nil && port != 443 {
		err = fmt.Errorf("Port %d not supported, only port 443 is supported", port)
	}
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getOrCreateHost(name, ip)
	online, connectionRefused, timedOut := h.status()
	if timedOut {
		log.Debugf("%v timed out waiting for status, returning 500 error", h)
		resp.WriteHeader(500)
		fmt.Fprintf(resp, "Timed out waiting for status")
		return
	}

	if online {
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Connectivity to proxy confirmed")
		return
	}

	// Note this may not work across platforms, but the intent
	// is to tell the client if the connection was flat out
	// refused as opposed to timed out in order to allow them
	// to configure their router if possible.
	if connectionRefused {
		// 417 response code.
		resp.WriteHeader(http.StatusExpectationFailed)
		fmt.Fprintln(resp, "No connectivity to proxy - connection refused")
	} else {
		// 408 response code.
		resp.WriteHeader(http.StatusRequestTimeout)
		fmt.Fprintln(resp, "No connectivity to proxy - test request timed out")
	}
}

// unregister is the HTTP endpoint for removing peers from DNS. Peers are
// unregistered based on their ip (not their name).
func unregister(resp http.ResponseWriter, req *http.Request) {
	_, ip, _, err := getHostInfo(req)
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getHostByIp(ip)
	msg := "Host not registered"
	if h != nil {
		h.unregister()
		msg = "Host unregistered"
	}
	resp.WriteHeader(200)
	fmt.Fprintln(resp, msg)
}

func getHostInfo(req *http.Request) (name string, ip string, port int, err error) {
	name = req.FormValue("name")
	if name == "" {
		err = fmt.Errorf("Please specify a name")
		return
	}
	ip = clientIpFor(req, name)
	if ip == "" {
		err = fmt.Errorf("Unable to determine IP address")
		return
	}
	portString := req.FormValue("port")

	if portString != "" {
		port, err = strconv.Atoi(portString)
		if err != nil {
			err = fmt.Errorf("Received invalid port for %v - %v: %v", name, ip, portString)
		}
	}

	return
}

func clientIpFor(req *http.Request, name string) string {
	// Client requested their info
	clientIp := req.Header.Get("X-Peerscanner-Forwarded-For")
	if clientIp == "" {
		clientIp = req.Header.Get("X-Forwarded-For")
	}
	if clientIp == "" && isFallback(name) {
		// Use direct IP for fallbacks
		clientIp = strings.Split(req.RemoteAddr, ":")[0]
	}
	// clientIp may contain multiple ips, use the first
	ips := strings.Split(clientIp, ",")
	return strings.TrimSpace(ips[0])
}
<commit_msg>Always using direct IP for fallbacks<commit_after>// main simply contains the primary web serving code that allows peers to
// register and unregister as give mode peers running within the Lantern
// network
package main

import (
	"crypto/tls"
	"fmt"
	"net/http"
	"strconv"
	"strings"

	"github.com/getlantern/keyman"
	"github.com/getlantern/tlsdefaults"
)

const (
	PKFile   = "pk.pem"
	CertFile = "cert.pem"
)

func startHttp() {
	http.HandleFunc("/register", register)
	http.HandleFunc("/unregister", unregister)
	laddr := fmt.Sprintf(":%d", *port)

	tlsConfig := tlsdefaults.Server()
	_, _, err := keyman.StoredPKAndCert(PKFile, CertFile, "Lantern", "localhost")
	if err != nil {
		log.Fatalf("Unable to initialize private key and certificate: %v", err)
	}
	cert, err := tls.LoadX509KeyPair(CertFile, PKFile)
	if err != nil {
		log.Fatalf("Unable to load certificate and key from %s and %s: %s", CertFile, PKFile, err)
	}
	tlsConfig.Certificates = []tls.Certificate{cert}

	log.Debugf("About to listen at %v", laddr)
	l, err := tls.Listen("tcp", laddr, tlsConfig)
	if err != nil {
		log.Fatalf("Unable to listen for tls connections at %s: %s", laddr, err)
	}

	log.Debug("About to serve")
	err = http.Serve(l, nil)
	if err != nil {
		log.Fatalf("Unable to serve: %s", err)
	}
}

// register is the entry point for peers registering themselves with the service.
// If peers are successfully vetted, they'll be added to the DNS round robin.
func register(resp http.ResponseWriter, req *http.Request) {
	name, ip, port, err := getHostInfo(req)
	if err == nil && port != 443 {
		err = fmt.Errorf("Port %d not supported, only port 443 is supported", port)
	}
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getOrCreateHost(name, ip)
	online, connectionRefused, timedOut := h.status()
	if timedOut {
		log.Debugf("%v timed out waiting for status, returning 500 error", h)
		resp.WriteHeader(500)
		fmt.Fprintf(resp, "Timed out waiting for status")
		return
	}

	if online {
		resp.WriteHeader(200)
		fmt.Fprintln(resp, "Connectivity to proxy confirmed")
		return
	}

	// Note this may not work across platforms, but the intent
	// is to tell the client if the connection was flat out
	// refused as opposed to timed out in order to allow them
	// to configure their router if possible.
	if connectionRefused {
		// 417 response code.
		resp.WriteHeader(http.StatusExpectationFailed)
		fmt.Fprintln(resp, "No connectivity to proxy - connection refused")
	} else {
		// 408 response code.
		resp.WriteHeader(http.StatusRequestTimeout)
		fmt.Fprintln(resp, "No connectivity to proxy - test request timed out")
	}
}

// unregister is the HTTP endpoint for removing peers from DNS. Peers are
// unregistered based on their ip (not their name).
func unregister(resp http.ResponseWriter, req *http.Request) {
	_, ip, _, err := getHostInfo(req)
	if err != nil {
		resp.WriteHeader(http.StatusBadRequest)
		fmt.Fprintln(resp, err.Error())
		return
	}

	h := getHostByIp(ip)
	msg := "Host not registered"
	if h != nil {
		h.unregister()
		msg = "Host unregistered"
	}
	resp.WriteHeader(200)
	fmt.Fprintln(resp, msg)
}

func getHostInfo(req *http.Request) (name string, ip string, port int, err error) {
	name = req.FormValue("name")
	if name == "" {
		err = fmt.Errorf("Please specify a name")
		return
	}
	ip = clientIpFor(req, name)
	if ip == "" {
		err = fmt.Errorf("Unable to determine IP address")
		return
	}
	portString := req.FormValue("port")

	if portString != "" {
		port, err = strconv.Atoi(portString)
		if err != nil {
			err = fmt.Errorf("Received invalid port for %v - %v: %v", name, ip, portString)
		}
	}

	return
}

func clientIpFor(req *http.Request, name string) string {
	// Client requested their info
	clientIp := req.Header.Get("X-Peerscanner-Forwarded-For")
	if clientIp == "" {
		clientIp = req.Header.Get("X-Forwarded-For")
	}
	if isFallback(name) {
		// Always use direct IP for fallbacks
		clientIp = strings.Split(req.RemoteAddr, ":")[0]
	}
	// clientIp may contain multiple ips, use the first
	ips := strings.Split(clientIp, ",")
	return strings.TrimSpace(ips[0])
}
<|endoftext|>


<commit_before>package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

var vowels = [...]string{"a", "e", "i", "o", "u"}
var bads = [...]string{"ab", "cd", "pq", "xy"}

func hasVowels(s string) bool {
	var cnt int
	for _, v := range vowels {
		cnt += strings.Count(s, v)

		if cnt >= 3 {
			return true
		}
	}

	return false
}

func hasRepeated(s string) bool {
	var prev rune
	for i, c := range s {
		if i > 0 && prev == c {
			return true
		}
		prev = c
	}

	return false
}

func hasNoBad(s string) bool {
	for _, b := range bads {
		if strings.Count(s, b) > 0 {
			return false
		}
	}

	return true
}

func hasRepeatedPair(s string) bool {
	return false
}

func hasRepeatWithGap(s string) bool {
	return false
}

func checkString(s string) bool {
	return hasVowels(s) && hasRepeated(s) && hasNoBad(s)
}

func checkString2(s string) bool {
	return hasRepeatedPair(s) && hasRepeatWithGap(s)
}

func run() int {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "%s filename\n", os.Args[0])
		return 1
	}

	f, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}
	defer f.Close()

	nice := 0
	nice2 := 0
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		s := scanner.Text()

		if checkString(s) {
			nice++
		}

		if checkString2(s) {
			nice2++
		}
	}

	fmt.Printf("nice: %d\nnice2: %d\n", nice, nice2)
	return 0
}

func main() {
	os.Exit(run())
}
<commit_msg>day5 part 1 changes<commit_after>package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

func hasVowels(s string) bool {
	var cnt int
	for _, c := range s {
		switch c {
		case 'a', 'e', 'i', 'o', 'u':
			cnt++
			if cnt >= 3 {
				return true
			}
		}
	}

	return false
}

func hasRepeated(s string) bool {
	r := []rune(s)
	for i := 0; i < len(r)-1; i++ {
		if r[i+1] == r[i] {
			return true
		}
	}

	return false
}

func hasNoBad(s string) bool {
	r := []rune(s)
	for i := 0; i < len(r)-1; i++ {
		switch string(r[i : i+2]) {
		case "ab", "cd", "pq", "xy":
			return false
		}
	}

	return true
}

func hasRepeatedPair(s string) bool {
	return false
}

func hasRepeatWithGap(s string) bool {
	return false
}

func checkString(s string) bool {
	return hasVowels(s) && hasRepeated(s) && hasNoBad(s)
}

func checkString2(s string) bool {
	return hasRepeatedPair(s) && hasRepeatWithGap(s)
}

func run() int {
	if len(os.Args) != 2 {
		fmt.Fprintf(os.Stderr, "%s filename\n", os.Args[0])
		return 1
	}

	f, err := os.Open(os.Args[1])
	if err != nil {
		fmt.Fprintln(os.Stderr, err)
		return 1
	}
	defer f.Close()

	nice := 0
	nice2 := 0
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		s := scanner.Text()

		if checkString(s) {
			nice++
		}

		if checkString2(s) {
			nice2++
		}
	}

	fmt.Printf("nice: %d\nnice2: %d\n", nice, nice2)
	return 0
}

func main() {
	os.Exit(run())
}
<|endoftext|>


<commit_before>package factorlib

import (
	"fmt"
	"math"
	"math/rand"

	"github.com/randall77/factorlib/big"
)

// width of window to sieve at once.  TODO: make configurable?
const sieverange = 1 << 20

// use an array of this size to do the sieving
const window = 1 << 14

// check to see if anything returned from the sieve isn't actually usable
const checkFalsePositive = false

// check to see if anything usable isn't returned by the sieve
const checkFalseNegative = false

// TODO: do this more systematically
var falsepos int

// Records f(x) == product(factors)*remainder
// The values in the factors slice are indexes into the factor base
type sieveResult struct {
	x         big.Int
	factors   []uint
	remainder int64
}

// Find values of x for which f(x) = a x^2 + b x + c factors (within one bigprime) over the primes in fb.
// Returns each x found, togegther with the factorization of f(x) into factor base primes and a
// possible bigprime remainder.
// Searches in the window [x0, x0+sieverange).
// requires: a > 0
func sievesmooth(a, b, c big.Int, fb []int64, x0 big.Int, rnd *rand.Rand) []sieveResult {
	var result []sieveResult

	maxp := fb[len(fb)-1]

	// Compute scaling factor for logarithms
	// Note: we assume that the maximum f(x) is achieved at one of the endpoints of the sieve range.
	y0 := a.Mul(x0).Add(b).Mul(x0).Add(c).Abs()
	x1 := x0.Add64(sieverange-1)
	y1 := a.Mul(x1).Add(b).Mul(x1).Add(c).Abs()
	maxf := y0
	if y1.Cmp(maxf) > 0 {
		maxf = y1
	}
	scale := 255/maxf.Log()

	// find starting points
	si := makeSieveInfo(a, b, c, fb, x0, scale, rnd)

	// compute thresholds for each window
	var thresholds [sieverange / window]byte
	for i := 0; i < sieverange; i += window {
		x := x0.Add64(int64(i))
		y0 = a.Mul(x).Add(b).Mul(x).Add(c).Abs()
		x = x.Add64(window)
		y1 = a.Mul(x).Add(b).Mul(x).Add(c).Abs()
		m := y0
		if y1.Cmp(m) > 0 {
			m = y1
		}
		thresholds[i/window] = byte(scale * (m.Log() - 2 * math.Log(float64(maxp))))
	}

	// sieve to find any potential smooth f(x)
	res := sieveinner(si, thresholds[:])

	// results buffer
	var factors []uint
	s := &big.Scratch{}

	// check potential results using trial factorization
	for _, i := range res {
		// compute y=f(x)
		x := x0.Add64(int64(i))
		y := a.Mul(x).Add(b).Mul(x).Add(c)

		// trial divide y by the factor base
		// accumulate factor base indexes of factors
		factors = factors[:0]
		for k, p := range fb {
			if p == -1 {
				if y.Sign() < 0 {
					y = y.Neg()
					factors = append(factors, uint(k))
				}
				continue
			}
			for y.Mod64s(p, s) == 0 {
				y = y.Div64(p)
				factors = append(factors, uint(k))
			}
		}

		// if remainder > B^2, it's too big, might not be prime.
		if y.Cmp64(maxp*maxp) > 0 {
			if checkFalsePositive {
				var f []int64
				for _, j := range factors {
					f = append(f, fb[j])
				}
				fmt.Printf("  false positive x=%d f(x)=%d f=%v·%d\n", x, a.Mul(x).Add(b).Mul(x).Add(c), f, y)
			}
			falsepos++
			continue
		}

		result = append(result, sieveResult{x, dup(factors), y.Int64()})
		if len(result) > 2*len(fb) {
			// early out when we're factoring small n
			// TODO: include in spec for sievesmooth function?
			return result
		}
	}
	if checkFalseNegative {
		// This is expensive, we have to trial divide all the numbers in the
		// sieve range.  Oh well, testing.
	checkloop:
		for i := 0; i < sieverange; i++ {
			for _, r := range res {
				if r == i {
					continue checkloop
				}
			}
			x := x0.Add64(int64(i))
			y := a.Mul(x).Add(b).Mul(x).Add(c)
			var f []int64
			for _, p := range fb {
				if p == -1 {
					if y.Sign() < 0 {
						y = y.Neg()
						f = append(f, p)
					}
					continue
				}
				for y.Mod64s(p, s) == 0 {
					y = y.Div64(p)
					f = append(f, p)
				}
			}

			// if remainder <= B^2, leftover is a prime
			if y.Cmp64(maxp*maxp) <= 0 {
				fmt.Printf("  false negative x=%d f(x)=%d f=%v·%d\n", x, a.Mul(x).Add(b).Mul(x).Add(c), f, y)
			}
		}
	}
	return result
}

// sieveinfo
type sieveinfo_inner struct {
	pk   int32  // p^k, the amount we step by through the sieve array
	off  uint16 // offset within the window which holds the next muliple of p^k
	lg_p uint8  // log_2(p), the amount we add to each sieve array slot
}

// a block of sieveinfo_inners
type block struct {
	next *block
	num  int
	data [64 - 2]sieveinfo_inner
}

func sieveinner(si []sieveinfo, thresholds []byte) []int {
	if window >= 1<<16 {
		panic("window too big")
	}

	// compute stats about our sieve work
	maxpk := int32(0)
	nsmall := 0
	for _, s := range si {
		pk := s.pk
		if pk < window {
			nsmall++
		}
		if pk > maxpk {
			maxpk = pk
		}
	}
	nlarge := len(si) - nsmall

	// Slice to store sieve work with pk < window
	small := make([]sieveinfo_inner, 0, nsmall)

	// For pk >= window, arrange sieve work in buckets of size w by offset.
	// The offsets stored in this array are relative to the start
	// of the bucket.  So all entries s in large[3] have an effective
	// offset of 3*window+s.off
	nw := 2 + int(maxpk/window)
	large := make([]*block, nw)

	// Allocate enough free blocks to hold everything.
	// There will be at most one partially empty block for each window.
	// Compute as the maximum # of full blocks + maximum # of nonfull blocks.
	n := nlarge/len(block{}.data) + nw
	blockstore := make([]block, n)
	// start each window with one empty block
	for i := 0; i < nw; i++ {
		large[i] = &blockstore[i]
	}
	// put the rest in a free list
	var freeblocks *block
	if nw < n {
		for i := nw; i < n-1; i++ {
			blockstore[i].next = &blockstore[i+1]
		}
		freeblocks = &blockstore[nw]
	}

	// put each sieve entry into the appropriate bucket
	for _, s := range si {
		pk := s.pk
		off := s.off
		lg_p := s.lg_p
		if pk < window {
			small = append(small, sieveinfo_inner{pk, uint16(off), lg_p})
			continue
		}
		i := off / window
		b := large[i]
		if b.num == len(b.data) {
			c := freeblocks
			freeblocks = c.next
			c.next = b
			large[i] = c
			b = c
		}
		b.data[b.num] = sieveinfo_inner{pk, uint16(off % window), lg_p}
		b.num++
	}
	/*
		log.Printf("len(si)=%d, len(small)=%d", len(si), len(small))
		for i := 0; i < nw; i++ {
			n := 0
			for b := large[i]; b != nil; b = b.next {
				n += b.num
			}
			log.Printf("len(large[%d])=%d", i, n)
		}
	*/

	var r []int
	for i := 0; i < sieverange; i += window {
		// start with a clear sieve
		var sieve [window]byte

		// increment sieve entries for f(x) that are divisible
		// by each factor base prime.

		// first, the small primes
		for j, s := range small {
			pk := int(s.pk)
			off := int(s.off)
			lg_p := s.lg_p
			for ; off < window; off += pk {
				sieve[off] += lg_p
			}
			small[j] = sieveinfo_inner{int32(pk), uint16(off - window), lg_p}
		}

		// second, the large primes
		for b := large[0]; b != nil; {
			for _, s := range b.data[:b.num] {
				pk := int(s.pk)
				off := int(s.off)
				lg_p := s.lg_p
				sieve[off] += lg_p
				off += pk
				j := off / window
				c := large[j]
				if c.num == len(c.data) {
					d := freeblocks
					freeblocks = d.next
					d.next = c
					large[j] = d
					c = d
				}
				c.data[c.num] = sieveinfo_inner{int32(pk), uint16(off % window), lg_p}
				c.num++
			}
			c := b
			b = b.next
			c.num = 0
			c.next = freeblocks
			freeblocks = c
		}

		// shift large buckets down by one
		copy(large, large[1:])
		// allocate a new empty bucket for the last one
		b := freeblocks
		freeblocks = b.next
		b.next = nil
		large[len(large)-1] = b

		// check for smooth numbers
		threshold := thresholds[i/window]
		for j := 0; j < window; j++ {
			if sieve[j] >= threshold {
				r = append(r, i+j)
			}
		}
	}
	return r
}

type sieveinfo struct {
	pk   int32 // p^k for this factor base entry
	lg_p uint8 // ~log_2(p)
	off  int32 // working offset in sieve array
}

func makeSieveInfo(a, b, c big.Int, fb []int64, x0 big.Int, scale float64, rnd *rand.Rand) []sieveinfo {
	var si []sieveinfo
	s := &big.Scratch{}
	maxp := fb[len(fb)-1]

	for _, p := range fb[1:] {
		if a.Mod64(p) == 0 {
			// This can happen in mpqs if p is one of the factors we used to construct a.
			// Ignore this prime - we might miss a few smooth f(x), but such is life.
			continue
		}
		lg_p := byte(scale * math.Log(float64(p))) // note: round down to avoid overflow in sieve buckets
		pk := p
		for k := uint(1); ; k++ {
			if pk > maxp {
				// Kind of arbitrary, but use powers of p as long as p^k is
				// smaller than than the maximum factor base prime.
				break
			}
			st := x0.Mod64s(pk, s)
			for _, r := range quadraticModPK(a.Mod64s(pk, s), b.Mod64s(pk, s), c.Mod64s(pk, s), p, k, pk, rnd) {
				// find first pk*i+r which is >= x0
				off := (r - st + pk) % pk
				si = append(si, sieveinfo{int32(pk), lg_p, int32(off)})
			}
			pk *= p
		}
	}
	return si
}
<commit_msg>Use larger sieve range.<commit_after>package factorlib

import (
	"fmt"
	"math"
	"math/rand"

	"github.com/randall77/factorlib/big"
)

// width of window to sieve at once.  TODO: make configurable?
const sieverange = 1 << 24

// use an array of this size to do the sieving
const window = 1 << 14

// check to see if anything returned from the sieve isn't actually usable
const checkFalsePositive = false

// check to see if anything usable isn't returned by the sieve
const checkFalseNegative = false

// TODO: do this more systematically
var falsepos int

// Records f(x) == product(factors)*remainder
// The values in the factors slice are indexes into the factor base
type sieveResult struct {
	x         big.Int
	factors   []uint
	remainder int64
}

// Find values of x for which f(x) = a x^2 + b x + c factors (within one bigprime) over the primes in fb.
// Returns each x found, togegther with the factorization of f(x) into factor base primes and a
// possible bigprime remainder.
// Searches in the window [x0, x0+sieverange).
// requires: a > 0
func sievesmooth(a, b, c big.Int, fb []int64, x0 big.Int, rnd *rand.Rand) []sieveResult {
	var result []sieveResult

	maxp := fb[len(fb)-1]

	// Compute scaling factor for logarithms
	// Note: we assume that the maximum f(x) is achieved at one of the endpoints of the sieve range.
	y0 := a.Mul(x0).Add(b).Mul(x0).Add(c).Abs()
	x1 := x0.Add64(sieverange-1)
	y1 := a.Mul(x1).Add(b).Mul(x1).Add(c).Abs()
	maxf := y0
	if y1.Cmp(maxf) > 0 {
		maxf = y1
	}
	scale := 255/maxf.Log()

	// find starting points
	si := makeSieveInfo(a, b, c, fb, x0, scale, rnd)

	// compute thresholds for each window
	var thresholds [sieverange / window]byte
	for i := 0; i < sieverange; i += window {
		x := x0.Add64(int64(i))
		y0 = a.Mul(x).Add(b).Mul(x).Add(c).Abs()
		x = x.Add64(window)
		y1 = a.Mul(x).Add(b).Mul(x).Add(c).Abs()
		m := y0
		if y1.Cmp(m) > 0 {
			m = y1
		}
		thresholds[i/window] = byte(scale * (m.Log() - 2 * math.Log(float64(maxp))))
	}

	// sieve to find any potential smooth f(x)
	res := sieveinner(si, thresholds[:])

	// results buffer
	var factors []uint
	s := &big.Scratch{}

	// check potential results using trial factorization
	for _, i := range res {
		// compute y=f(x)
		x := x0.Add64(int64(i))
		y := a.Mul(x).Add(b).Mul(x).Add(c)

		// trial divide y by the factor base
		// accumulate factor base indexes of factors
		factors = factors[:0]
		for k, p := range fb {
			if p == -1 {
				if y.Sign() < 0 {
					y = y.Neg()
					factors = append(factors, uint(k))
				}
				continue
			}
			for y.Mod64s(p, s) == 0 {
				y = y.Div64(p)
				factors = append(factors, uint(k))
			}
		}

		// if remainder > B^2, it's too big, might not be prime.
		if y.Cmp64(maxp*maxp) > 0 {
			if checkFalsePositive {
				var f []int64
				for _, j := range factors {
					f = append(f, fb[j])
				}
				fmt.Printf("  false positive x=%d f(x)=%d f=%v·%d\n", x, a.Mul(x).Add(b).Mul(x).Add(c), f, y)
			}
			falsepos++
			continue
		}

		result = append(result, sieveResult{x, dup(factors), y.Int64()})
		if len(result) > 2*len(fb) {
			// early out when we're factoring small n
			// TODO: include in spec for sievesmooth function?
			return result
		}
	}
	if checkFalseNegative {
		// This is expensive, we have to trial divide all the numbers in the
		// sieve range.  Oh well, testing.
	checkloop:
		for i := 0; i < sieverange; i++ {
			for _, r := range res {
				if r == i {
					continue checkloop
				}
			}
			x := x0.Add64(int64(i))
			y := a.Mul(x).Add(b).Mul(x).Add(c)
			var f []int64
			for _, p := range fb {
				if p == -1 {
					if y.Sign() < 0 {
						y = y.Neg()
						f = append(f, p)
					}
					continue
				}
				for y.Mod64s(p, s) == 0 {
					y = y.Div64(p)
					f = append(f, p)
				}
			}

			// if remainder <= B^2, leftover is a prime
			if y.Cmp64(maxp*maxp) <= 0 {
				fmt.Printf("  false negative x=%d f(x)=%d f=%v·%d\n", x, a.Mul(x).Add(b).Mul(x).Add(c), f, y)
			}
		}
	}
	return result
}

// sieveinfo
type sieveinfo_inner struct {
	pk   int32  // p^k, the amount we step by through the sieve array
	off  uint16 // offset within the window which holds the next muliple of p^k
	lg_p uint8  // log_2(p), the amount we add to each sieve array slot
}

// a block of sieveinfo_inners
type block struct {
	next *block
	num  int
	data [64 - 2]sieveinfo_inner
}

func sieveinner(si []sieveinfo, thresholds []byte) []int {
	if window >= 1<<16 {
		panic("window too big")
	}

	// compute stats about our sieve work
	maxpk := int32(0)
	nsmall := 0
	for _, s := range si {
		pk := s.pk
		if pk < window {
			nsmall++
		}
		if pk > maxpk {
			maxpk = pk
		}
	}
	nlarge := len(si) - nsmall

	// Slice to store sieve work with pk < window
	small := make([]sieveinfo_inner, 0, nsmall)

	// For pk >= window, arrange sieve work in buckets of size w by offset.
	// The offsets stored in this array are relative to the start
	// of the bucket.  So all entries s in large[3] have an effective
	// offset of 3*window+s.off
	nw := 2 + int(maxpk/window)
	large := make([]*block, nw)

	// Allocate enough free blocks to hold everything.
	// There will be at most one partially empty block for each window.
	// Compute as the maximum # of full blocks + maximum # of nonfull blocks.
	n := nlarge/len(block{}.data) + nw
	blockstore := make([]block, n)
	// start each window with one empty block
	for i := 0; i < nw; i++ {
		large[i] = &blockstore[i]
	}
	// put the rest in a free list
	var freeblocks *block
	if nw < n {
		for i := nw; i < n-1; i++ {
			blockstore[i].next = &blockstore[i+1]
		}
		freeblocks = &blockstore[nw]
	}

	// put each sieve entry into the appropriate bucket
	for _, s := range si {
		pk := s.pk
		off := s.off
		lg_p := s.lg_p
		if pk < window {
			small = append(small, sieveinfo_inner{pk, uint16(off), lg_p})
			continue
		}
		i := off / window
		b := large[i]
		if b.num == len(b.data) {
			c := freeblocks
			freeblocks = c.next
			c.next = b
			large[i] = c
			b = c
		}
		b.data[b.num] = sieveinfo_inner{pk, uint16(off % window), lg_p}
		b.num++
	}
	/*
		log.Printf("len(si)=%d, len(small)=%d", len(si), len(small))
		for i := 0; i < nw; i++ {
			n := 0
			for b := large[i]; b != nil; b = b.next {
				n += b.num
			}
			log.Printf("len(large[%d])=%d", i, n)
		}
	*/

	var r []int
	for i := 0; i < sieverange; i += window {
		// start with a clear sieve
		var sieve [window]byte

		// increment sieve entries for f(x) that are divisible
		// by each factor base prime.

		// first, the small primes
		for j, s := range small {
			pk := int(s.pk)
			off := int(s.off)
			lg_p := s.lg_p
			for ; off < window; off += pk {
				sieve[off] += lg_p
			}
			small[j] = sieveinfo_inner{int32(pk), uint16(off - window), lg_p}
		}

		// second, the large primes
		for b := large[0]; b != nil; {
			for _, s := range b.data[:b.num] {
				pk := int(s.pk)
				off := int(s.off)
				lg_p := s.lg_p
				sieve[off] += lg_p
				off += pk
				j := off / window
				c := large[j]
				if c.num == len(c.data) {
					d := freeblocks
					freeblocks = d.next
					d.next = c
					large[j] = d
					c = d
				}
				c.data[c.num] = sieveinfo_inner{int32(pk), uint16(off % window), lg_p}
				c.num++
			}
			c := b
			b = b.next
			c.num = 0
			c.next = freeblocks
			freeblocks = c
		}

		// shift large buckets down by one
		copy(large, large[1:])
		// allocate a new empty bucket for the last one
		b := freeblocks
		freeblocks = b.next
		b.next = nil
		large[len(large)-1] = b

		// check for smooth numbers
		threshold := thresholds[i/window]
		for j := 0; j < window; j++ {
			if sieve[j] >= threshold {
				r = append(r, i+j)
			}
		}
	}
	return r
}

type sieveinfo struct {
	pk   int32 // p^k for this factor base entry
	lg_p uint8 // ~log_2(p)
	off  int32 // working offset in sieve array
}

func makeSieveInfo(a, b, c big.Int, fb []int64, x0 big.Int, scale float64, rnd *rand.Rand) []sieveinfo {
	var si []sieveinfo
	s := &big.Scratch{}
	maxp := fb[len(fb)-1]

	for _, p := range fb[1:] {
		if a.Mod64(p) == 0 {
			// This can happen in mpqs if p is one of the factors we used to construct a.
			// Ignore this prime - we might miss a few smooth f(x), but such is life.
			continue
		}
		lg_p := byte(scale * math.Log(float64(p))) // note: round down to avoid overflow in sieve buckets
		pk := p
		for k := uint(1); ; k++ {
			if pk > maxp {
				// Kind of arbitrary, but use powers of p as long as p^k is
				// smaller than than the maximum factor base prime.
				break
			}
			st := x0.Mod64s(pk, s)
			for _, r := range quadraticModPK(a.Mod64s(pk, s), b.Mod64s(pk, s), c.Mod64s(pk, s), p, k, pk, rnd) {
				// find first pk*i+r which is >= x0
				off := (r - st + pk) % pk
				si = append(si, sieveinfo{int32(pk), lg_p, int32(off)})
			}
			pk *= p
		}
	}
	return si
}
<|endoftext|>


<commit_before>// Licensed under the MIT license, see LICENCE file for details.

package quicktest_test

import (
	"errors"
	"os"
	"path/filepath"
	"testing"

	qt "github.com/frankban/quicktest"
)

func TestPatchSetInt(t *testing.T) {
	c := qt.New(t)
	i := 99
	testDefer(c, func(c *qt.C) {
		c.Patch(&i, 88)
		c.Assert(i, qt.Equals, 88)
	})
	c.Assert(i, qt.Equals, 99)
}

func TestPatchSetError(t *testing.T) {
	c := qt.New(t)
	oldErr := errors.New("foo")
	newErr := errors.New("bar")
	err := oldErr
	testDefer(c, func(c *qt.C) {
		c.Patch(&err, newErr)
		c.Assert(err, qt.Equals, newErr)
	})
	c.Assert(err, qt.Equals, oldErr)
}

func TestPatchSetErrorToNil(t *testing.T) {
	c := qt.New(t)
	oldErr := errors.New("foo")
	err := oldErr
	testDefer(c, func(c *qt.C) {
		c.Patch(&err, nil)
		c.Assert(err, qt.Equals, nil)
	})
	c.Assert(err, qt.Equals, oldErr)
}

func TestPatchSetMapToNil(t *testing.T) {
	c := qt.New(t)
	oldMap := map[string]int{"foo": 1234}
	m := oldMap
	testDefer(c, func(c *qt.C) {
		c.Patch(&m, nil)
		c.Assert(m, qt.IsNil)
	})
	c.Assert(m, qt.DeepEquals, oldMap)
}

func TestSetPatchPanicsWhenNotAssignable(t *testing.T) {
	c := qt.New(t)
	i := 99
	type otherInt int
	c.Assert(func() { c.Patch(&i, otherInt(88)) }, qt.PanicMatches, `reflect\.Set: value of type quicktest_test\.otherInt is not assignable to type int`)
}

func TestSetenv(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Setenv(envName, "initial")
	testDefer(c, func(c *qt.C) {
		c.Setenv(envName, "new value")
		c.Check(os.Getenv(envName), qt.Equals, "new value")
	})
	c.Check(os.Getenv(envName), qt.Equals, "initial")
}

func TestSetenvWithUnsetVariable(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Unsetenv(envName)
	testDefer(c, func(c *qt.C) {
		c.Setenv(envName, "new value")
		c.Check(os.Getenv(envName), qt.Equals, "new value")
	})
	_, ok := os.LookupEnv(envName)
	c.Assert(ok, qt.Equals, false)
}

func TestUnsetenv(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Setenv(envName, "initial")
	testDefer(c, func(c *qt.C) {
		c.Unsetenv(envName)
		_, ok := os.LookupEnv(envName)
		c.Assert(ok, qt.Equals, false)
	})
	c.Check(os.Getenv(envName), qt.Equals, "initial")
}

func TestUnsetenvWithUnsetVariable(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Unsetenv(envName)
	testDefer(c, func(c *qt.C) {
		c.Unsetenv(envName)
		_, ok := os.LookupEnv(envName)
		c.Assert(ok, qt.Equals, false)
	})
	_, ok := os.LookupEnv(envName)
	c.Assert(ok, qt.Equals, false)
}

func TestMkdir(t *testing.T) {
	c := qt.New(t)
	var dir string
	testDefer(c, func(c *qt.C) {
		dir = c.Mkdir()
		c.Assert(dir, qt.Not(qt.Equals), "")
		info, err := os.Stat(dir)
		c.Assert(err, qt.Equals, nil)
		c.Assert(info.IsDir(), qt.Equals, true)
		f, err := os.Create(filepath.Join(dir, "hello"))
		c.Assert(err, qt.Equals, nil)
		f.Close()
	})
	_, err := os.Stat(dir)
	c.Assert(err, qt.Not(qt.IsNil))
}
<commit_msg>Minor fixes to tests<commit_after>// Licensed under the MIT license, see LICENCE file for details.

package quicktest_test

import (
	"errors"
	"os"
	"path/filepath"
	"testing"

	qt "github.com/frankban/quicktest"
)

func TestPatchSetInt(t *testing.T) {
	c := qt.New(t)
	i := 99
	testDefer(c, func(c *qt.C) {
		c.Patch(&i, 88)
		c.Assert(i, qt.Equals, 88)
	})
	c.Assert(i, qt.Equals, 99)
}

func TestPatchSetError(t *testing.T) {
	c := qt.New(t)
	oldErr := errors.New("foo")
	newErr := errors.New("bar")
	err := oldErr
	testDefer(c, func(c *qt.C) {
		c.Patch(&err, newErr)
		c.Assert(err, qt.Equals, newErr)
	})
	c.Assert(err, qt.Equals, oldErr)
}

func TestPatchSetErrorToNil(t *testing.T) {
	c := qt.New(t)
	oldErr := errors.New("foo")
	err := oldErr
	testDefer(c, func(c *qt.C) {
		c.Patch(&err, nil)
		c.Assert(err, qt.IsNil)
	})
	c.Assert(err, qt.Equals, oldErr)
}

func TestPatchSetMapToNil(t *testing.T) {
	c := qt.New(t)
	oldMap := map[string]int{"foo": 1234}
	m := oldMap
	testDefer(c, func(c *qt.C) {
		c.Patch(&m, nil)
		c.Assert(m, qt.IsNil)
	})
	c.Assert(m, qt.DeepEquals, oldMap)
}

func TestSetPatchPanicsWhenNotAssignable(t *testing.T) {
	c := qt.New(t)
	i := 99
	type otherInt int
	c.Assert(func() { c.Patch(&i, otherInt(88)) }, qt.PanicMatches, `reflect\.Set: value of type quicktest_test\.otherInt is not assignable to type int`)
}

func TestSetenv(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Setenv(envName, "initial")
	testDefer(c, func(c *qt.C) {
		c.Setenv(envName, "new value")
		c.Check(os.Getenv(envName), qt.Equals, "new value")
	})
	c.Check(os.Getenv(envName), qt.Equals, "initial")
}

func TestSetenvWithUnsetVariable(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Unsetenv(envName)
	testDefer(c, func(c *qt.C) {
		c.Setenv(envName, "new value")
		c.Check(os.Getenv(envName), qt.Equals, "new value")
	})
	_, ok := os.LookupEnv(envName)
	c.Assert(ok, qt.IsFalse)
}

func TestUnsetenv(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Setenv(envName, "initial")
	testDefer(c, func(c *qt.C) {
		c.Unsetenv(envName)
		_, ok := os.LookupEnv(envName)
		c.Assert(ok, qt.IsFalse)
	})
	c.Check(os.Getenv(envName), qt.Equals, "initial")
}

func TestUnsetenvWithUnsetVariable(t *testing.T) {
	c := qt.New(t)
	const envName = "SOME_VAR"
	os.Unsetenv(envName)
	testDefer(c, func(c *qt.C) {
		c.Unsetenv(envName)
		_, ok := os.LookupEnv(envName)
		c.Assert(ok, qt.IsFalse)
	})
	_, ok := os.LookupEnv(envName)
	c.Assert(ok, qt.IsFalse)
}

func TestMkdir(t *testing.T) {
	c := qt.New(t)
	var dir string
	testDefer(c, func(c *qt.C) {
		dir = c.Mkdir()
		c.Assert(dir, qt.Not(qt.Equals), "")
		info, err := os.Stat(dir)
		c.Assert(err, qt.IsNil)
		c.Assert(info.IsDir(), qt.IsTrue)
		f, err := os.Create(filepath.Join(dir, "hello"))
		c.Assert(err, qt.IsNil)
		f.Close()
	})
	_, err := os.Stat(dir)
	c.Assert(err, qt.Not(qt.IsNil))
}
<|endoftext|>


<commit_before>package jsonpatch

import (
	"github.com/stretchr/testify/assert"
	"testing"
)

func TestMakePatch(t *testing.T) {
	// Test map
	docA := getMapDoc(`{"this":{"is":"my", "document":"sir"}}`)
	docB := getMapDoc(`{"this":{"document":"my", "is":"sir", "now":{"go":"away!"}}}`)
	patch, err := MakePatch(docA, docB)
	assert.Nil(t, err)
	assert.Equal(t, 3, len(patch.Operations))
	// TODO: This test must depend on the ordering of a map iteration, because I'm getting
	// occasional failures occasionally with the ordering of the following two lines.
	assert.Equal(t, PatchOperation{Op: "replace", Path: "/this/is", Value: "sir"}, patch.Operations[0])
	assert.Equal(t, PatchOperation{Op: "replace", Path: "/this/document", Value: "my"}, patch.Operations[1])
	assert.Equal(t, PatchOperation{Op: "add", Path: "/this/now", Value: map[string]interface{}{"go": "away!"}}, patch.Operations[2])

	// Test array
	docA = getMapDoc(`{"a":[0, 1, 2, 3]}`)
	docB = getMapDoc(`{"a":[1, 2, 4, "hi"]}`)
	patch, err = MakePatch(docA, docB)
	assert.Nil(t, err)
	assert.Equal(t, 3, len(patch.Operations))
	assert.Equal(t, PatchOperation{Op: "remove", Path: "/a/0"}, patch.Operations[0])
	assert.Equal(t, PatchOperation{Op: "replace", Path: "/a/2", Value: 4}, patch.Operations[1])
	assert.Equal(t, PatchOperation{Op: "add", Path: "/a/3", Value: "hi"}, patch.Operations[2])
}

func TestApplyPatchFromString(t *testing.T) {
	doc := getMapDoc(`{"foo": "bar"}`)

	patchOp, err := FromString(`[{"op": "add", "path": "/baz", "value": "qux"}]`)
	assert.Nil(t, err)
	patchOp.Apply(&doc)
	val, found := doc["baz"]
	assert.True(t, found)
	assert.Equal(t, "qux", val.(string))
}

func TestLcs(t *testing.T) {
	pairA, pairB := longestCommonSubseq(slice(1, 2, 3, 4), slice(0, 1, 2, 3, 5))
	assert.Equal(t, intPair{0, 3}, *pairA)
	assert.Equal(t, intPair{1, 4}, *pairB)

	pairA, pairB = longestCommonSubseq(slice(1, 3, 5), slice(0, 1, 2, 3, 4, 5, 6))
	assert.Equal(t, intPair{2, 3}, *pairA)
	assert.Equal(t, intPair{5, 6}, *pairB)
}

func TestSplitByCommonSeq(t *testing.T) {
	node := splitByCommonSeq(slice(0, 1, 2, 3), slice(1, 2, 4, 5), &intPair{0, -1}, &intPair{0, -1})
	assert.Nil(t, node.left)
	assert.Nil(t, node.right)

	// Left subtree
	assert.NotNil(t, node.leftPtr)
	assert.Equal(t, intPair{0, 1}, *node.leftPtr.left)
	assert.Nil(t, node.leftPtr.leftPtr)
	assert.Nil(t, node.leftPtr.rightPtr)
	// Right subtree
	assert.NotNil(t, node.rightPtr)
	assert.Equal(t, intPair{3, 4}, *node.rightPtr.left)
	assert.Equal(t, intPair{2, 4}, *node.rightPtr.right)
	assert.Nil(t, node.rightPtr.rightPtr)
	assert.Nil(t, node.rightPtr.leftPtr)
}

func slice(args ...interface{}) []interface{} {
	s := []interface{}{}
	s = append(s, args...)
	return s
}
<commit_msg>TestMakePatch tests patch.Apply() result and not individual patch ops<commit_after>package jsonpatch

import (
	"fmt"
	"testing"

	"github.com/stretchr/testify/assert"
)

// TestMakePatch_nop ensures that MakePatch returns an empty patch when the
// input equals the output.
func TestMakePatch_nop(t *testing.T) {
	for i, test := range []struct {
		src string
	}{
		{ // simple map
			`{"a":1,"b":2}`,
		},
		{ // nested map
			`{"a":{"b":"c", "d":"e"}}`,
		},
		{ // array
			`{"a":[0, 1, 2, 3]}`,
		},
	} {
		index := fmt.Sprintf("test %d", i)
		docA := getMapDoc(test.src)
		docB := getMapDoc(test.src)
		patch, err := MakePatch(docA, docB)
		assert.Nil(t, err, index)
		assert.Equal(t, len(patch.Operations), 0, index)
	}
}

// TestMakePatch ensures that patches cleanly apply to the original documents
// and produce a document equivalent to the target.
func TestMakePatch(t *testing.T) {
	for i, test := range []struct {
		src, dst string
	}{
		{ // simple map operations.
			`{"a":1,"b":2}`,
			`{"a":2,"b":1}`,
		},
		{ // nested map operations
			`{"this":{"is":"my", "document":"sir"}}`,
			`{"this":{"document":"my", "is":"sir", "now":{"go":"away!"}}}`,
		},
		{ // array operations
			`{"a":[0, 1, 2, 3]}`,
			`{"a":[1, 2, 4, "hi"]}`,
		},
		/* BUG: this test panics
		{ // nested array operations
			`{"a":[0, [1, 2, 3], 5, 2]}`,
			`{"a":[1, [1, 0, 2, 3, 5], "x", 2]}`,
		},
		*/
	} {
		index := fmt.Sprintf("test %d", i)
		docA := getMapDoc(test.src)
		docB := getMapDoc(test.dst)
		patch, err := MakePatch(docA, docB)
		assert.Nil(t, err, index)
		err = patch.Apply(&docA)
		assert.Nil(t, err, index)
		assert.Equal(t, docA, docB, index)
	}
}

func TestApplyPatchFromString(t *testing.T) {
	doc := getMapDoc(`{"foo": "bar"}`)

	patchOp, err := FromString(`[{"op": "add", "path": "/baz", "value": "qux"}]`)
	assert.Nil(t, err)
	patchOp.Apply(&doc)
	val, found := doc["baz"]
	assert.True(t, found)
	assert.Equal(t, "qux", val.(string))
}

func TestLcs(t *testing.T) {
	pairA, pairB := longestCommonSubseq(slice(1, 2, 3, 4), slice(0, 1, 2, 3, 5))
	assert.Equal(t, intPair{0, 3}, *pairA)
	assert.Equal(t, intPair{1, 4}, *pairB)

	pairA, pairB = longestCommonSubseq(slice(1, 3, 5), slice(0, 1, 2, 3, 4, 5, 6))
	assert.Equal(t, intPair{2, 3}, *pairA)
	assert.Equal(t, intPair{5, 6}, *pairB)
}

func TestSplitByCommonSeq(t *testing.T) {
	node := splitByCommonSeq(slice(0, 1, 2, 3), slice(1, 2, 4, 5), &intPair{0, -1}, &intPair{0, -1})
	assert.Nil(t, node.left)
	assert.Nil(t, node.right)

	// Left subtree
	assert.NotNil(t, node.leftPtr)
	assert.Equal(t, intPair{0, 1}, *node.leftPtr.left)
	assert.Nil(t, node.leftPtr.leftPtr)
	assert.Nil(t, node.leftPtr.rightPtr)
	// Right subtree
	assert.NotNil(t, node.rightPtr)
	assert.Equal(t, intPair{3, 4}, *node.rightPtr.left)
	assert.Equal(t, intPair{2, 4}, *node.rightPtr.right)
	assert.Nil(t, node.rightPtr.rightPtr)
	assert.Nil(t, node.rightPtr.leftPtr)
}

func slice(args ...interface{}) []interface{} {
	s := []interface{}{}
	s = append(s, args...)
	return s
}
<|endoftext|>


<commit_before>package parse

func Term(name string) *R {
	return &R{Name: name}
}

func Rule(name string, rules ...*R) *R {
	return Con(rules...).initRecursiveRule().As(name)
}
func (r *R) initRecursiveRule() *R {
	r.traverseAlt(make(map[*R]bool), func(a *Alt) {
		if a.Parent == Self {
			a.Parent = r
		}
		for i := range a.Rules {
			if a.Rules[i] == Self {
				a.Rules[i] = r
			}
		}
	})
	return r
}

func Con(rules ...*R) *R {
	if len(rules) == 1 && rules[0].Name == "" {
		return rules[0]
	}
	r := &R{}
	r.Alts = Alts{{r, rules}}
	return r
}

func Or(rules ...*R) *R {
	if len(rules) == 1 && rules[0].Name == "" {
		return rules[0]
	}
	r := &R{Alts: make(Alts, len(rules))}
	for i := range rules {
		r.Alts[i] = rules[i].toAlt(r)
	}
	return r
}
func (r *R) toAlt(parent *R) *Alt {
	if len(r.Alts) == 1 {
		return &Alt{parent, r.Alts[0].Rules}
	}
	return &Alt{parent, Rules{r}}
}

func (r *R) As(name string) *R {
	r.Name = name
	return r
}

func (r *R) ZeroOrOne() *R {
	return Or(r, Null)
}

func (r *R) ZeroOrMore() *R {
	x := &R{}
	x.Alts = Con(x, r).ZeroOrOne().toAlts(x)
	return x
}
func (r *R) toAlts(parent *R) Alts {
	r.eachAlt(func(a *Alt) {
		a.Parent = parent
	})
	return r.Alts
}
<commit_msg>implement OneOrMore.<commit_after>package parse

func Term(name string) *R {
	return &R{Name: name}
}

func Rule(name string, rules ...*R) *R {
	return Con(rules...).initRecursiveRule().As(name)
}
func (r *R) initRecursiveRule() *R {
	r.traverseAlt(make(map[*R]bool), func(a *Alt) {
		if a.Parent == Self {
			a.Parent = r
		}
		for i := range a.Rules {
			if a.Rules[i] == Self {
				a.Rules[i] = r
			}
		}
	})
	return r
}

func Con(rules ...*R) *R {
	if len(rules) == 1 && rules[0].Name == "" {
		return rules[0]
	}
	r := &R{}
	r.Alts = Alts{{r, rules}}
	return r
}

func Or(rules ...*R) *R {
	if len(rules) == 1 && rules[0].Name == "" {
		return rules[0]
	}
	r := &R{Alts: make(Alts, len(rules))}
	for i := range rules {
		r.Alts[i] = rules[i].toAlt(r)
	}
	return r
}
func (r *R) toAlt(parent *R) *Alt {
	if len(r.Alts) == 1 {
		return &Alt{parent, r.Alts[0].Rules}
	}
	return &Alt{parent, Rules{r}}
}

func (r *R) As(name string) *R {
	r.Name = name
	return r
}

func (r *R) ZeroOrOne() *R {
	return Or(r, Null)
}

func (r *R) OneOrMore() *R {
	return Con(r.ZeroOrMore(), r)
}

func (r *R) ZeroOrMore() *R {
	x := &R{}
	x.Alts = Con(x, r).ZeroOrOne().toAlts(x)
	return x
}
func (r *R) toAlts(parent *R) Alts {
	r.eachAlt(func(a *Alt) {
		a.Parent = parent
	})
	return r.Alts
}
<|endoftext|>


<commit_before>package httputils

import (
	"errors"
	"strconv"
	"strings"

	"github.com/koofr/go-ioutils"
)

var ErrInvalidRange = errors.New("invalid range")

func ParseRange(s string, size int64) (spans []ioutils.FileSpan, hasEnd bool, err error) {
	if s == "" {
		return nil, false, nil // header not present
	}

	const b = "bytes="
	if !strings.HasPrefix(s, b) {
		return nil, false, ErrInvalidRange
	}

	hasEnd = true

	for _, ra := range strings.Split(s[len(b):], ",") {
		ra = strings.TrimSpace(ra)
		if ra == "" {
			continue
		}
		i := strings.Index(ra, "-")
		if i < 0 {
			return nil, false, ErrInvalidRange
		}
		start, end := strings.TrimSpace(ra[:i]), strings.TrimSpace(ra[i+1:])
		var s ioutils.FileSpan
		if start == "" {
			// If no start is specified, end specifies the
			// range start relative to the end of the file.
			i, err := strconv.ParseInt(end, 10, 64)
			if err != nil {
				return nil, false, ErrInvalidRange
			}
			if i > size {
				i = size
			}
			s.Start = size - i
			s.End = size - 1
		} else {
			i, err := strconv.ParseInt(start, 10, 64)
			if err != nil || i >= size || i < 0 {
				return nil, false, ErrInvalidRange
			}
			s.Start = i
			if end == "" {
				// If no end is specified, range extends to end of the file.
				s.End = size - 1
				hasEnd = false
			} else {
				i, err := strconv.ParseInt(end, 10, 64)
				if err != nil || s.Start > i {
					return nil, false, ErrInvalidRange
				}
				if i >= size {
					i = size - 1
				}
				s.End = i
			}
		}
		spans = append(spans, s)
	}

	return spans, hasEnd, nil
}
<commit_msg>Use textproto.TrimString<commit_after>package httputils

import (
	"errors"
	"net/textproto"
	"strconv"
	"strings"

	"github.com/koofr/go-ioutils"
)

var ErrInvalidRange = errors.New("invalid range")

func ParseRange(s string, size int64) (spans []ioutils.FileSpan, hasEnd bool, err error) {
	if s == "" {
		return nil, false, nil // header not present
	}

	const b = "bytes="
	if !strings.HasPrefix(s, b) {
		return nil, false, ErrInvalidRange
	}

	hasEnd = true

	for _, ra := range strings.Split(s[len(b):], ",") {
		ra = textproto.TrimString(ra)
		if ra == "" {
			continue
		}
		i := strings.Index(ra, "-")
		if i < 0 {
			return nil, false, ErrInvalidRange
		}
		start, end := textproto.TrimString(ra[:i]), textproto.TrimString(ra[i+1:])
		var s ioutils.FileSpan
		if start == "" {
			// If no start is specified, end specifies the
			// range start relative to the end of the file.
			i, err := strconv.ParseInt(end, 10, 64)
			if err != nil {
				return nil, false, ErrInvalidRange
			}
			if i > size {
				i = size
			}
			s.Start = size - i
			s.End = size - 1
		} else {
			i, err := strconv.ParseInt(start, 10, 64)
			if err != nil || i >= size || i < 0 {
				return nil, false, ErrInvalidRange
			}
			s.Start = i
			if end == "" {
				// If no end is specified, range extends to end of the file.
				s.End = size - 1
				hasEnd = false
			} else {
				i, err := strconv.ParseInt(end, 10, 64)
				if err != nil || s.Start > i {
					return nil, false, ErrInvalidRange
				}
				if i >= size {
					i = size - 1
				}
				s.End = i
			}
		}
		spans = append(spans, s)
	}

	return spans, hasEnd, nil
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"
)

// TODO testable

// partitions represents a list of partitions for a single version and their
// mapping to nodes, synced from zookeeper. It's also responsible for
// advertising the partitions we have locally.
type partitions struct {
	peers     *peers
	zkWatcher *zkWatcher

	db      string
	version string
	zkPath  string

	total   int
	missing int
	local   map[int]bool
	remote  map[int][]string

	lock  sync.RWMutex
	ready chan bool
}

func watchPartitions(zkWatcher *zkWatcher, peers *peers, db, version string, numPartitions, replication int) *partitions {
	p := &partitions{
		peers:     peers,
		zkWatcher: zkWatcher,
		db:        db,
		version:   version,
		zkPath:    path.Join("partitions", db, version),
		total:     numPartitions,
		remote:    make(map[int][]string),
		ready:     make(chan bool),
	}

	p.pickLocalPartitions(replication)

	// Create the partitions path we're going to watch, in case no one has done
	// that yet.
	p.zkPath = path.Join("partitions", db, version)
	zkWatcher.createPath(p.zkPath)

	updates, _ := zkWatcher.watchChildren(p.zkPath)
	go p.sync(updates)
	return p
}

// pickLocalPartitions selects which partitions are local by iterating through
// them all, and checking the hashring to see if this peer is one of the
// replicas.
func (p *partitions) pickLocalPartitions(replication int) {
	local := make(map[int]bool)
	disp := make([]int, 0)

	for i := 0; i < p.total; i++ {
		partitionId := p.partitionId(i)

		replicas := p.peers.pick(partitionId, replication)
		for _, replica := range replicas {
			if replica == peerSelf {
				local[i] = true
				disp = append(disp, i)
			}
		}
	}

	log.Printf("Selected partitions for %s version %s: %v", p.db, p.version, disp)
	p.local = local
	p.missing = p.total - len(local)
}

// sync syncs the remote partitions from zoolander whenever they change.
func (p *partitions) sync(updates chan []string) {
	for {
		nodes := <-updates
		p.updatePartitions(nodes)
	}
}

func (p *partitions) updatePartitions(nodes []string) {
	p.lock.Lock()
	defer p.lock.Unlock()

	remote := make(map[int][]string)
	for _, node := range nodes {
		parts := strings.SplitN(node, "@", 2)
		partition, _ := strconv.Atoi(parts[0])
		host := parts[1]
		if host != p.peers.address {
			remote[partition] = append(remote[partition], host)
		}
	}

	// Check for each partition. If every one is available on at least one node,
	// then we're ready to rumble.
	p.remote = remote
	missing := 0
	for i := 0; i < p.total; i++ {
		if _, ok := p.local[i]; ok {
			continue
		}

		if _, ok := p.remote[i]; ok {
			continue
		}

		missing += 1
	}

	p.missing = missing
	if missing == 0 {
		select {
		case p.ready <- true:
		default:
		}
	}
}

// advertiseAndWait advertises the partitions we have locally, and waits until
// it sees at least one peer for every remote partition. It returns false only
// if it was closed before that happens.
func (p *partitions) advertiseAndWait() bool {
	// Advertise that our local partitions are ready.
	p.advertisePartitions()

	for {
		p.lock.RLock()
		missing := p.missing
		p.lock.RUnlock()
		if missing == 0 {
			break
		}

		log.Printf("Waiting for all partitions of %s version %s to be available (missing %d)",
			p.db, p.version, missing)

		t := time.NewTimer(10 * time.Second)
		select {
		case <-t.C:
		case success := <-p.ready:
			// If success is false, it's because the close() was called before we
			// finished waiting on peers
			return success
		}
	}

	return true
}

// advertisePartitions creates an ephemeral node for each partition this local
// peer is responsible for.
// TODO: this should maybe be a zk multi op?
func (p *partitions) advertisePartitions() {
	for partition := range p.local {
		node := fmt.Sprintf("%05d@%s", partition, p.peers.address)
		p.zkWatcher.createEphemeral(path.Join(p.zkPath, node))
	}
}

// getPeers returns the list of peers who have the given partition available.
func (p *partitions) getPeers(partition int) []string {
	p.lock.RLock()
	defer p.lock.RUnlock()

	peers := make([]string, len(p.remote[partition]))
	copy(peers, p.remote[partition])
	return peers
}

// partitionId returns a string id for the given partition, to be used for the
// consistent hashing ring. It's not really meant to be unique, but it should be
// different for different versions with the same number of partitions, so that
// they don't shard identically.
func (p *partitions) partitionId(partition int) string {
	return fmt.Sprintf("%s:%05d", p.zkPath, partition)
}

func (p *partitions) close() {
	// Delete ephemeral keys
	close(p.ready)
}
<commit_msg>Don't close partitions.ready, since we'll panic when it updates<commit_after>package main

import (
	"fmt"
	"log"
	"path"
	"strconv"
	"strings"
	"sync"
	"time"
)

// TODO testable

// partitions represents a list of partitions for a single version and their
// mapping to nodes, synced from zookeeper. It's also responsible for
// advertising the partitions we have locally.
type partitions struct {
	peers     *peers
	zkWatcher *zkWatcher

	db      string
	version string
	zkPath  string

	total   int
	missing int
	local   map[int]bool
	remote  map[int][]string

	lock  sync.RWMutex
	ready chan bool
}

func watchPartitions(zkWatcher *zkWatcher, peers *peers, db, version string, numPartitions, replication int) *partitions {
	p := &partitions{
		peers:     peers,
		zkWatcher: zkWatcher,
		db:        db,
		version:   version,
		zkPath:    path.Join("partitions", db, version),
		total:     numPartitions,
		remote:    make(map[int][]string),
		ready:     make(chan bool),
	}

	p.pickLocalPartitions(replication)

	// Create the partitions path we're going to watch, in case no one has done
	// that yet.
	p.zkPath = path.Join("partitions", db, version)
	zkWatcher.createPath(p.zkPath)

	updates, _ := zkWatcher.watchChildren(p.zkPath)
	go p.sync(updates)
	return p
}

// pickLocalPartitions selects which partitions are local by iterating through
// them all, and checking the hashring to see if this peer is one of the
// replicas.
func (p *partitions) pickLocalPartitions(replication int) {
	local := make(map[int]bool)
	disp := make([]int, 0)

	for i := 0; i < p.total; i++ {
		partitionId := p.partitionId(i)

		replicas := p.peers.pick(partitionId, replication)
		for _, replica := range replicas {
			if replica == peerSelf {
				local[i] = true
				disp = append(disp, i)
			}
		}
	}

	log.Printf("Selected partitions for %s version %s: %v", p.db, p.version, disp)
	p.local = local
	p.missing = p.total - len(local)
}

// sync syncs the remote partitions from zoolander whenever they change.
func (p *partitions) sync(updates chan []string) {
	for {
		nodes := <-updates
		p.updatePartitions(nodes)
	}
}

func (p *partitions) updatePartitions(nodes []string) {
	p.lock.Lock()
	defer p.lock.Unlock()

	remote := make(map[int][]string)
	for _, node := range nodes {
		parts := strings.SplitN(node, "@", 2)
		partition, _ := strconv.Atoi(parts[0])
		host := parts[1]
		if host != p.peers.address {
			remote[partition] = append(remote[partition], host)
		}
	}

	// Check for each partition. If every one is available on at least one node,
	// then we're ready to rumble.
	p.remote = remote
	missing := 0
	for i := 0; i < p.total; i++ {
		if _, ok := p.local[i]; ok {
			continue
		}

		if _, ok := p.remote[i]; ok {
			continue
		}

		missing += 1
	}

	p.missing = missing
	if missing == 0 {
		select {
		case p.ready <- true:
		default:
		}
	}
}

// advertiseAndWait advertises the partitions we have locally, and waits until
// it sees at least one peer for every remote partition. It returns false only
// if it was closed before that happens.
func (p *partitions) advertiseAndWait() bool {
	// Advertise that our local partitions are ready.
	p.advertisePartitions()

	for {
		p.lock.RLock()
		missing := p.missing
		p.lock.RUnlock()
		if missing == 0 {
			break
		}

		log.Printf("Waiting for all partitions of %s version %s to be available (missing %d)",
			p.db, p.version, missing)

		t := time.NewTimer(10 * time.Second)
		select {
		case <-t.C:
		case success := <-p.ready:
			// If success is false, it's because the close() was called before we
			// finished waiting on peers
			return success
		}
	}

	return true
}

// advertisePartitions creates an ephemeral node for each partition this local
// peer is responsible for.
// TODO: this should maybe be a zk multi op?
func (p *partitions) advertisePartitions() {
	for partition := range p.local {
		node := fmt.Sprintf("%05d@%s", partition, p.peers.address)
		p.zkWatcher.createEphemeral(path.Join(p.zkPath, node))
	}
}

// getPeers returns the list of peers who have the given partition available.
func (p *partitions) getPeers(partition int) []string {
	p.lock.RLock()
	defer p.lock.RUnlock()

	peers := make([]string, len(p.remote[partition]))
	copy(peers, p.remote[partition])
	return peers
}

// partitionId returns a string id for the given partition, to be used for the
// consistent hashing ring. It's not really meant to be unique, but it should be
// different for different versions with the same number of partitions, so that
// they don't shard identically.
func (p *partitions) partitionId(partition int) string {
	return fmt.Sprintf("%s:%05d", p.zkPath, partition)
}

func (p *partitions) close() {
	// Delete ephemeral keys
	// Remove watch
	select {
	case p.ready <- false:
	default:
	}
}
<|endoftext|>


<commit_before>package tusd_test

import (
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"

	. "github.com/tus/tusd"
)

func TestPatch(t *testing.T) {
	SubTest(t, "UploadChunk", func(t *testing.T, store *MockFullDataStore) {
		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				ID:     "yes",
				Offset: 5,
				Size:   10,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(5), NewReaderMatcher("hello")).Return(int64(5), nil),
			store.EXPECT().FinishUpload("yes"),
		)

		handler, _ := NewHandler(Config{
			DataStore:             store,
			NotifyCompleteUploads: true,
		})

		c := make(chan FileInfo, 1)
		handler.CompleteUploads = c

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			ReqBody: strings.NewReader("hello"),
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "10",
			},
		}).Run(handler, t)

		a := assert.New(t)
		info := <-c
		a.Equal("yes", info.ID)
		a.EqualValues(int64(10), info.Size)
		a.Equal(int64(10), info.Offset)
	})

	SubTest(t, "MethodOverriding", func(t *testing.T, store *MockFullDataStore) {
		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				ID:     "yes",
				Offset: 5,
				Size:   20,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(5), NewReaderMatcher("hello")).Return(int64(5), nil),
		)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "POST",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable":          "1.0.0",
				"Upload-Offset":          "5",
				"Content-Type":           "application/offset+octet-stream",
				"X-HTTP-Method-Override": "PATCH",
			},
			ReqBody: strings.NewReader("hello"),
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "10",
			},
		}).Run(handler, t)
	})

	SubTest(t, "UploadChunkToFinished", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("yes").Return(FileInfo{
			Offset: 20,
			Size:   20,
		}, nil)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "20",
			},
			ReqBody: strings.NewReader(""),
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "20",
			},
		}).Run(handler, t)
	})

	SubTest(t, "UploadNotFoundFail", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("no").Return(FileInfo{}, os.ErrNotExist)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "no",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			Code: http.StatusNotFound,
		}).Run(handler, t)
	})

	SubTest(t, "MissmatchingOffsetFail", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("yes").Return(FileInfo{
			Offset: 5,
		}, nil)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "4",
			},
			Code: http.StatusConflict,
		}).Run(handler, t)
	})

	SubTest(t, "ExceedingMaxSizeFail", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("yes").Return(FileInfo{
			Offset: 5,
			Size:   10,
		}, nil)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			ReqBody: strings.NewReader("hellothisismorethan15bytes"),
			Code:    http.StatusRequestEntityTooLarge,
		}).Run(handler, t)
	})

	SubTest(t, "InvalidContentTypeFail", func(t *testing.T, store *MockFullDataStore) {
		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/fail",
				"Upload-Offset": "5",
			},
			ReqBody: strings.NewReader("hellothisismorethan15bytes"),
			Code:    http.StatusBadRequest,
		}).Run(handler, t)
	})

	SubTest(t, "InvalidOffsetFail", func(t *testing.T, store *MockFullDataStore) {
		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "-5",
			},
			ReqBody: strings.NewReader("hellothisismorethan15bytes"),
			Code:    http.StatusBadRequest,
		}).Run(handler, t)
	})

	SubTest(t, "OverflowWithoutLength", func(t *testing.T, store *MockFullDataStore) {
		// In this test we attempt to upload more than 15 bytes to an upload
		// which has only space for 15 bytes (offset of 5 and size of 20).
		// The request does not contain the Content-Length header and the handler
		// therefore does not know the chunk's size before. The wanted behavior
		// is that even if the uploader supplies more than 15 bytes, we only
		// pass 15 bytes to the data store and ignore the rest.

		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				Offset: 5,
				Size:   20,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(5), NewReaderMatcher("hellothisismore")).Return(int64(15), nil),
			store.EXPECT().FinishUpload("yes"),
		)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		// Wrap the string.Reader in a NopCloser to hide its type. else
		// http.NewRequest() will detect the we supply a strings.Reader as body
		// and use this information to set the Content-Length header which we
		// explicitly do not want (see comment above for reason).
		body := ioutil.NopCloser(strings.NewReader("hellothisismorethan15bytes"))

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			ReqBody: body,
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "20",
			},
		}).Run(handler, t)
	})

	SubTest(t, "Locker", func(t *testing.T, store *MockFullDataStore) {
		ctrl := gomock.NewController(t)
		defer ctrl.Finish()
		locker := NewMockLocker(ctrl)

		gomock.InOrder(
			locker.EXPECT().LockUpload("yes").Return(nil),
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				Offset: 0,
				Size:   20,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(0), NewReaderMatcher("hello")).Return(int64(5), nil),
			locker.EXPECT().UnlockUpload("yes").Return(nil),
		)

		composer := NewStoreComposer()
		composer.UseCore(store)
		composer.UseLocker(locker)

		handler, _ := NewHandler(Config{
			StoreComposer: composer,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "0",
			},
			ReqBody: strings.NewReader("hello"),
			Code:    http.StatusNoContent,
		}).Run(handler, t)
	})

	SubTest(t, "NotifyUploadProgress", func(t *testing.T, store *MockFullDataStore) {
		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				ID:     "yes",
				Offset: 0,
				Size:   100,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(0), NewReaderMatcher("first second third")).Return(int64(18), nil),
		)

		handler, _ := NewHandler(Config{
			DataStore:            store,
			NotifyUploadProgress: true,
		})

		c := make(chan FileInfo, 10)
		handler.UploadProgress = c

		reader, writer := io.Pipe()

		go func() {
			writer.Write([]byte("first "))
			<-time.After(time.Second)
			writer.Write([]byte("second "))
			writer.Write([]byte("third"))
			writer.Close()
		}()

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "0",
			},
			ReqBody: reader,
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "18",
			},
		}).Run(handler, t)

		// Wait a short time after the request has been handled before closing the
		// channel because another goroutine may still write to the channel.
		<-time.After(10 * time.Millisecond)
		close(handler.UploadProgress)

		a := assert.New(t)

		info := <-c
		a.Equal("yes", info.ID)
		a.Equal(int64(100), info.Size)
		a.Equal(int64(6), info.Offset)

		info = <-c
		a.Equal("yes", info.ID)
		a.Equal(int64(100), info.Size)
		a.Equal(int64(18), info.Offset)

		_, more := <-c
		a.False(more)
	})
}
<commit_msg>Include channel testing in goroutine fore NotifyUploadProgress test<commit_after>package tusd_test

import (
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"strings"
	"testing"
	"time"

	"github.com/golang/mock/gomock"
	"github.com/stretchr/testify/assert"

	. "github.com/tus/tusd"
)

func TestPatch(t *testing.T) {
	SubTest(t, "UploadChunk", func(t *testing.T, store *MockFullDataStore) {
		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				ID:     "yes",
				Offset: 5,
				Size:   10,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(5), NewReaderMatcher("hello")).Return(int64(5), nil),
			store.EXPECT().FinishUpload("yes"),
		)

		handler, _ := NewHandler(Config{
			DataStore:             store,
			NotifyCompleteUploads: true,
		})

		c := make(chan FileInfo, 1)
		handler.CompleteUploads = c

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			ReqBody: strings.NewReader("hello"),
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "10",
			},
		}).Run(handler, t)

		a := assert.New(t)
		info := <-c
		a.Equal("yes", info.ID)
		a.EqualValues(int64(10), info.Size)
		a.Equal(int64(10), info.Offset)
	})

	SubTest(t, "MethodOverriding", func(t *testing.T, store *MockFullDataStore) {
		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				ID:     "yes",
				Offset: 5,
				Size:   20,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(5), NewReaderMatcher("hello")).Return(int64(5), nil),
		)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "POST",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable":          "1.0.0",
				"Upload-Offset":          "5",
				"Content-Type":           "application/offset+octet-stream",
				"X-HTTP-Method-Override": "PATCH",
			},
			ReqBody: strings.NewReader("hello"),
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "10",
			},
		}).Run(handler, t)
	})

	SubTest(t, "UploadChunkToFinished", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("yes").Return(FileInfo{
			Offset: 20,
			Size:   20,
		}, nil)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "20",
			},
			ReqBody: strings.NewReader(""),
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "20",
			},
		}).Run(handler, t)
	})

	SubTest(t, "UploadNotFoundFail", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("no").Return(FileInfo{}, os.ErrNotExist)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "no",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			Code: http.StatusNotFound,
		}).Run(handler, t)
	})

	SubTest(t, "MissmatchingOffsetFail", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("yes").Return(FileInfo{
			Offset: 5,
		}, nil)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "4",
			},
			Code: http.StatusConflict,
		}).Run(handler, t)
	})

	SubTest(t, "ExceedingMaxSizeFail", func(t *testing.T, store *MockFullDataStore) {
		store.EXPECT().GetInfo("yes").Return(FileInfo{
			Offset: 5,
			Size:   10,
		}, nil)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			ReqBody: strings.NewReader("hellothisismorethan15bytes"),
			Code:    http.StatusRequestEntityTooLarge,
		}).Run(handler, t)
	})

	SubTest(t, "InvalidContentTypeFail", func(t *testing.T, store *MockFullDataStore) {
		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/fail",
				"Upload-Offset": "5",
			},
			ReqBody: strings.NewReader("hellothisismorethan15bytes"),
			Code:    http.StatusBadRequest,
		}).Run(handler, t)
	})

	SubTest(t, "InvalidOffsetFail", func(t *testing.T, store *MockFullDataStore) {
		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "-5",
			},
			ReqBody: strings.NewReader("hellothisismorethan15bytes"),
			Code:    http.StatusBadRequest,
		}).Run(handler, t)
	})

	SubTest(t, "OverflowWithoutLength", func(t *testing.T, store *MockFullDataStore) {
		// In this test we attempt to upload more than 15 bytes to an upload
		// which has only space for 15 bytes (offset of 5 and size of 20).
		// The request does not contain the Content-Length header and the handler
		// therefore does not know the chunk's size before. The wanted behavior
		// is that even if the uploader supplies more than 15 bytes, we only
		// pass 15 bytes to the data store and ignore the rest.

		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				Offset: 5,
				Size:   20,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(5), NewReaderMatcher("hellothisismore")).Return(int64(15), nil),
			store.EXPECT().FinishUpload("yes"),
		)

		handler, _ := NewHandler(Config{
			DataStore: store,
		})

		// Wrap the string.Reader in a NopCloser to hide its type. else
		// http.NewRequest() will detect the we supply a strings.Reader as body
		// and use this information to set the Content-Length header which we
		// explicitly do not want (see comment above for reason).
		body := ioutil.NopCloser(strings.NewReader("hellothisismorethan15bytes"))

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "5",
			},
			ReqBody: body,
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "20",
			},
		}).Run(handler, t)
	})

	SubTest(t, "Locker", func(t *testing.T, store *MockFullDataStore) {
		ctrl := gomock.NewController(t)
		defer ctrl.Finish()
		locker := NewMockLocker(ctrl)

		gomock.InOrder(
			locker.EXPECT().LockUpload("yes").Return(nil),
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				Offset: 0,
				Size:   20,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(0), NewReaderMatcher("hello")).Return(int64(5), nil),
			locker.EXPECT().UnlockUpload("yes").Return(nil),
		)

		composer := NewStoreComposer()
		composer.UseCore(store)
		composer.UseLocker(locker)

		handler, _ := NewHandler(Config{
			StoreComposer: composer,
		})

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "0",
			},
			ReqBody: strings.NewReader("hello"),
			Code:    http.StatusNoContent,
		}).Run(handler, t)
	})

	SubTest(t, "NotifyUploadProgress", func(t *testing.T, store *MockFullDataStore) {
		gomock.InOrder(
			store.EXPECT().GetInfo("yes").Return(FileInfo{
				ID:     "yes",
				Offset: 0,
				Size:   100,
			}, nil),
			store.EXPECT().WriteChunk("yes", int64(0), NewReaderMatcher("first second third")).Return(int64(18), nil),
		)

		handler, _ := NewHandler(Config{
			DataStore:            store,
			NotifyUploadProgress: true,
		})

		c := make(chan FileInfo)
		handler.UploadProgress = c

		reader, writer := io.Pipe()
		a := assert.New(t)

		go func() {
			writer.Write([]byte("first "))

			info := <-c
			a.Equal("yes", info.ID)
			a.Equal(int64(100), info.Size)
			a.Equal(int64(6), info.Offset)

			writer.Write([]byte("second "))
			writer.Write([]byte("third"))

			info = <-c
			a.Equal("yes", info.ID)
			a.Equal(int64(100), info.Size)
			a.Equal(int64(18), info.Offset)

			writer.Close()

			info = <-c
			a.Equal("yes", info.ID)
			a.Equal(int64(100), info.Size)
			a.Equal(int64(18), info.Offset)
		}()

		(&httpTest{
			Method: "PATCH",
			URL:    "yes",
			ReqHeader: map[string]string{
				"Tus-Resumable": "1.0.0",
				"Content-Type":  "application/offset+octet-stream",
				"Upload-Offset": "0",
			},
			ReqBody: reader,
			Code:    http.StatusNoContent,
			ResHeader: map[string]string{
				"Upload-Offset": "18",
			},
		}).Run(handler, t)

		// Wait a short time after the request has been handled before closing the
		// channel because another goroutine may still write to the channel.
		<-time.After(10 * time.Millisecond)
		close(handler.UploadProgress)

		_, more := <-c
		a.False(more)
	})
}
<|endoftext|>


<commit_before>package main

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"golang.org/x/net/html"
)

// Results -- struct returned by Crawl() to represent the entire crawl process
type Results struct {
	// Inherit the Resource struct
	Resource

	// Body represents a string implementation of the byte array returned by
	// http.Response
	Body string

	// Slice of Resource structs containing the needed resources for the given URL
	Resources []*Resource

	// ResourceTime shows how long it took to fetch all resources
	ResourceTime *TimerResult

	// TotalTime represents the time it took to crawl the site
	TotalTime *TimerResult
}

// Resource represents a single entity of many within a given crawl. These should
// only be of type css, js, jpg, png, etc (static resources).
type Resource struct {
	// connURL is the initial URL received by input
	connURL string

	// connIP is the initial IP address received by input
	connIP string

	// connHostname represents the original requested hostname for the resource
	connHostname string

	// URL represents the resulting static URL derived by the original result page
	URL string

	// Hostname represents the resulting hostname derived by the original returned
	// resource
	Hostname string

	// Remote represents if the resulting resource is remote to the original domain
	Remote bool

	// Error represents any errors that may have occurred when fetching the resource
	Error error

	// Code represents the numeric HTTP based status code
	Code int

	// Proto represents the end protocol used to fetch the page. For example, HTTP/2.0
	Proto string

	// ContentLength represents the number of bytes in the body of the response
	ContentLength int64

	// TLS represents the SSL/TLS handshake/session if the resource was loaded over
	// SSL.
	TLS *tls.ConnectionState

	// Time represents the time it took to complete the request
	Time *TimerResult
}

var resourcePool sync.WaitGroup

// getSrc crawls the body of the Results page, yielding all img/script/link resources
// so they can later be fetched.
func getSrc(b io.ReadCloser, req *http.Request) (urls []string) {
	urls = []string{}

	z := html.NewTokenizer(b)

	for {
		// loop through all tokens in the html body response
		tt := z.Next()

		switch {
		case tt == html.ErrorToken:
			// this assumes that there are no further tokens -- end of document
			return
		case tt == html.StartTagToken:
			t := z.Token()

			// the tokens that we are pulling resources from, and the attribute we are
			// pulling from
			allowed := map[string]string{
				"link":   "href",
				"script": "src",
				"img":    "src",
			}
			var isInAllowed bool
			var checkType string
			var src string

			// loop through all allowed elements, and see if the current element is
			// allowed
			for key := range allowed {
				if t.Data == key {
					isInAllowed = true
					checkType = allowed[key]
					break
				}
			}

			if !isInAllowed {
				continue
			}

			for _, a := range t.Attr {
				if a.Key == checkType {
					src = a.Val
					break
				}
			}

			// this assumes that the resource is something along the lines of:
			//   http://something.com/ -- which we don't care about
			if len(src) == 0 || strings.HasSuffix(src, "/") {
				continue
			}

			// add trailing slash to the end of the path
			if len(req.URL.Path) == 0 {
				req.URL.Path = "/"
			}

			// site was developed using relative paths. E.g:
			//  - url: http://domain.com/sub/path and resource: ./something/main.js
			//    would equal http://domain.com/sub/path/something/main.js
			if strings.HasPrefix(src, "./") {
				src = req.URL.Scheme + "://" + req.URL.Host + req.URL.Path + strings.SplitN(src, "./", 2)[1]
			}

			// site is loading resources from a remote location that supports both
			// http and https. browsers should natively tack on the current sites
			// protocol to the url. E.g:
			//  - url: http://domain.com/ and resource: //other.com/some-resource.js
			//    generates: http://other.com/some-resource.js
			//  - url: https://domain.com/ and resource: //other.com/some-resource.js
			//    generates: https://other.com/some-resource.js
			if strings.HasPrefix(src, "//") {
				src = req.URL.Scheme + ":" + src
			}

			// non-host-absolute resource. E.g. resource is loaded based on the docroot
			// of the domain. E.g:
			//  - url: http://domain.com/ and resource: /some-resource.js
			//    generates: http://domain.com/some-resource.js
			//  - url: https://domain.com/sub/resource and resource: /some-resource.js
			//    generates: https://domain.com/some-resource.js
			if strings.HasPrefix(src, "/") {
				src = req.URL.Scheme + "://" + req.URL.Host + src
			}

			// ignore anything else that isn't http based. E.g. ftp://, and other svg-like
			// data urls, as we really can't fetch those.
			if !strings.HasPrefix(src, "http") {
				continue
			}

			urls = append(urls, src)
		}
	}
}

func connHostname(URL string) (host string, err error) {
	tmp, err := url.Parse(URL)

	if err != nil {
		return
	}

	host = tmp.Host
	return
}

// FetchResource fetches a singular resource from a page, returning a *Resource struct.
// As we don't care much about the body of the resource, that can safely be ignored. We
// must still close the body object, however.
func (rsrc *Resource) FetchResource() {
	var err error

	defer resourcePool.Done()

	// calculate the time it takes to fetch the request
	timer := NewTimer()
	resp, err := Get(rsrc.connURL, rsrc.connIP)
	rsrc.Time = timer.End()
	resp.Body.Close()

	if err != nil {
		rsrc.Error = err
		return
	}

	rsrc.connHostname, err = connHostname(rsrc.connURL)
	if err != nil {
		rsrc.Error = err
		return
	}

	rsrc.Hostname = resp.Request.Host
	rsrc.URL = resp.Request.URL.String()
	rsrc.Code = resp.StatusCode
	rsrc.Proto = resp.Proto
	rsrc.ContentLength = resp.ContentLength
	rsrc.TLS = resp.TLS

	if rsrc.Hostname != rsrc.connHostname {
		rsrc.Remote = true
	}

	fmt.Printf("[%d] [%s] %s\n", rsrc.Code, rsrc.Proto, rsrc.URL)

	return
}

// Crawl manages the fetching of the main resource, as well as all child resources,
// providing a Results struct containing the entire crawl data needed
func Crawl(URL string, IP string) (res *Results) {
	res = &Results{}

	crawlTimer := NewTimer()
	reqTimer := NewTimer()

	// actually fetch the request
	resp, err := Get(URL, IP)
	defer resp.Body.Close()

	res.Time = reqTimer.End()

	if err != nil {
		res.Error = err
		return
	}

	res.connHostname, err = connHostname(URL)
	if err != nil {
		res.Error = err
		return
	}

	res.connURL = URL
	res.connIP = IP
	res.Hostname = resp.Request.Host
	res.URL = URL
	res.Code = resp.StatusCode
	res.Proto = resp.Proto
	res.ContentLength = resp.ContentLength
	res.TLS = resp.TLS

	if res.Hostname != res.connHostname {
		res.Remote = true
	}

	buf, _ := ioutil.ReadAll(resp.Body)
	b := ioutil.NopCloser(bytes.NewReader(buf))
	defer b.Close()

	bbytes, err := ioutil.ReadAll(bytes.NewBuffer(buf))
	if len(bbytes) != 0 {
		res.Body = string(bbytes[:])
	}

	urls := getSrc(b, resp.Request)

	fmt.Printf("[%d] [%s] %s\n", res.Code, res.Proto, res.URL)

	resourceTime := NewTimer()

	for i := range urls {
		resourcePool.Add(1)

		rsrc := &Resource{connURL: urls[i], connIP: ""}
		res.Resources = append(res.Resources, rsrc)
		go res.Resources[i].FetchResource()
	}

	resourcePool.Wait()

	res.ResourceTime = resourceTime.End()
	res.TotalTime = crawlTimer.End()

	return
}
<commit_msg>add scheme<commit_after>package main

import (
	"bytes"
	"crypto/tls"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"net/url"
	"strings"
	"sync"

	"golang.org/x/net/html"
)

// Results -- struct returned by Crawl() to represent the entire crawl process
type Results struct {
	// Inherit the Resource struct
	Resource

	// Body represents a string implementation of the byte array returned by
	// http.Response
	Body string

	// Slice of Resource structs containing the needed resources for the given URL
	Resources []*Resource

	// ResourceTime shows how long it took to fetch all resources
	ResourceTime *TimerResult

	// TotalTime represents the time it took to crawl the site
	TotalTime *TimerResult
}

// Resource represents a single entity of many within a given crawl. These should
// only be of type css, js, jpg, png, etc (static resources).
type Resource struct {
	// connURL is the initial URL received by input
	connURL string

	// connIP is the initial IP address received by input
	connIP string

	// connHostname represents the original requested hostname for the resource
	connHostname string

	// URL represents the resulting static URL derived by the original result page
	URL string

	// Hostname represents the resulting hostname derived by the original returned
	// resource
	Hostname string

	// Remote represents if the resulting resource is remote to the original domain
	Remote bool

	// Error represents any errors that may have occurred when fetching the resource
	Error error

	// Code represents the numeric HTTP based status code
	Code int

	// Proto represents the end protocol used to fetch the page. For example, HTTP/2.0
	Proto string

	// Scheme represents the end scheme used to fetch the page. For example, https
	Scheme string

	// ContentLength represents the number of bytes in the body of the response
	ContentLength int64

	// TLS represents the SSL/TLS handshake/session if the resource was loaded over
	// SSL.
	TLS *tls.ConnectionState

	// Time represents the time it took to complete the request
	Time *TimerResult
}

var resourcePool sync.WaitGroup

// getSrc crawls the body of the Results page, yielding all img/script/link resources
// so they can later be fetched.
func getSrc(b io.ReadCloser, req *http.Request) (urls []string) {
	urls = []string{}

	z := html.NewTokenizer(b)

	for {
		// loop through all tokens in the html body response
		tt := z.Next()

		switch {
		case tt == html.ErrorToken:
			// this assumes that there are no further tokens -- end of document
			return
		case tt == html.StartTagToken:
			t := z.Token()

			// the tokens that we are pulling resources from, and the attribute we are
			// pulling from
			allowed := map[string]string{
				"link":   "href",
				"script": "src",
				"img":    "src",
			}
			var isInAllowed bool
			var checkType string
			var src string

			// loop through all allowed elements, and see if the current element is
			// allowed
			for key := range allowed {
				if t.Data == key {
					isInAllowed = true
					checkType = allowed[key]
					break
				}
			}

			if !isInAllowed {
				continue
			}

			for _, a := range t.Attr {
				if a.Key == checkType {
					src = a.Val
					break
				}
			}

			// this assumes that the resource is something along the lines of:
			//   http://something.com/ -- which we don't care about
			if len(src) == 0 || strings.HasSuffix(src, "/") {
				continue
			}

			// add trailing slash to the end of the path
			if len(req.URL.Path) == 0 {
				req.URL.Path = "/"
			}

			// site was developed using relative paths. E.g:
			//  - url: http://domain.com/sub/path and resource: ./something/main.js
			//    would equal http://domain.com/sub/path/something/main.js
			if strings.HasPrefix(src, "./") {
				src = req.URL.Scheme + "://" + req.URL.Host + req.URL.Path + strings.SplitN(src, "./", 2)[1]
			}

			// site is loading resources from a remote location that supports both
			// http and https. browsers should natively tack on the current sites
			// protocol to the url. E.g:
			//  - url: http://domain.com/ and resource: //other.com/some-resource.js
			//    generates: http://other.com/some-resource.js
			//  - url: https://domain.com/ and resource: //other.com/some-resource.js
			//    generates: https://other.com/some-resource.js
			if strings.HasPrefix(src, "//") {
				src = req.URL.Scheme + ":" + src
			}

			// non-host-absolute resource. E.g. resource is loaded based on the docroot
			// of the domain. E.g:
			//  - url: http://domain.com/ and resource: /some-resource.js
			//    generates: http://domain.com/some-resource.js
			//  - url: https://domain.com/sub/resource and resource: /some-resource.js
			//    generates: https://domain.com/some-resource.js
			if strings.HasPrefix(src, "/") {
				src = req.URL.Scheme + "://" + req.URL.Host + src
			}

			// ignore anything else that isn't http based. E.g. ftp://, and other svg-like
			// data urls, as we really can't fetch those.
			if !strings.HasPrefix(src, "http") {
				continue
			}

			urls = append(urls, src)
		}
	}
}

func connHostname(URL string) (host string, err error) {
	tmp, err := url.Parse(URL)

	if err != nil {
		return
	}

	host = tmp.Host
	return
}

// FetchResource fetches a singular resource from a page, returning a *Resource struct.
// As we don't care much about the body of the resource, that can safely be ignored. We
// must still close the body object, however.
func (rsrc *Resource) FetchResource() {
	var err error

	defer resourcePool.Done()

	// calculate the time it takes to fetch the request
	timer := NewTimer()
	resp, err := Get(rsrc.connURL, rsrc.connIP)
	rsrc.Time = timer.End()
	resp.Body.Close()

	if err != nil {
		rsrc.Error = err
		return
	}

	rsrc.connHostname, err = connHostname(rsrc.connURL)
	if err != nil {
		rsrc.Error = err
		return
	}

	rsrc.Hostname = resp.Request.Host
	rsrc.URL = resp.Request.URL.String()
	rsrc.Code = resp.StatusCode
	rsrc.Proto = resp.Proto
	rsrc.Scheme = resp.Request.URL.Scheme
	rsrc.ContentLength = resp.ContentLength
	rsrc.TLS = resp.TLS

	if rsrc.Hostname != rsrc.connHostname {
		rsrc.Remote = true
	}

	fmt.Printf("[%d] [%s] %s\n", rsrc.Code, rsrc.Proto, rsrc.URL)

	return
}

// Crawl manages the fetching of the main resource, as well as all child resources,
// providing a Results struct containing the entire crawl data needed
func Crawl(URL string, IP string) (res *Results) {
	res = &Results{}

	crawlTimer := NewTimer()
	reqTimer := NewTimer()

	// actually fetch the request
	resp, err := Get(URL, IP)
	defer resp.Body.Close()

	res.Time = reqTimer.End()

	if err != nil {
		res.Error = err
		return
	}

	res.connHostname, err = connHostname(URL)
	if err != nil {
		res.Error = err
		return
	}

	res.connURL = URL
	res.connIP = IP
	res.Hostname = resp.Request.Host
	res.URL = URL
	res.Code = resp.StatusCode
	res.Proto = resp.Proto
	res.Scheme = resp.Request.URL.Scheme
	res.ContentLength = resp.ContentLength
	res.TLS = resp.TLS

	if res.Hostname != res.connHostname {
		res.Remote = true
	}

	buf, _ := ioutil.ReadAll(resp.Body)
	b := ioutil.NopCloser(bytes.NewReader(buf))
	defer b.Close()

	bbytes, err := ioutil.ReadAll(bytes.NewBuffer(buf))
	if len(bbytes) != 0 {
		res.Body = string(bbytes[:])
	}

	urls := getSrc(b, resp.Request)

	fmt.Printf("[%d] [%s] %s\n", res.Code, res.Proto, res.URL)

	resourceTime := NewTimer()

	for i := range urls {
		resourcePool.Add(1)

		rsrc := &Resource{connURL: urls[i], connIP: ""}
		res.Resources = append(res.Resources, rsrc)
		go res.Resources[i].FetchResource()
	}

	resourcePool.Wait()

	res.ResourceTime = resourceTime.End()
	res.TotalTime = crawlTimer.End()

	return
}
<|endoftext|>


<commit_before>package db

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
	"time"
)

const (
	//tokensExpireAfter time.Duration = time.Duration(time.Second * 60 * 60 * 24 * 7) // 1 week
	// tokensExpireAfter controls how long authentication tokens are valid.
	tokensExpireAfter time.Duration = time.Duration(time.Second * 60 * 60) // 1h
)

type (
	// Token holds information associated with a given authentication token
	Token struct {
		ID        bson.ObjectId `bson:"_id"`
		CreatedAt time.Time     `bson:"createdAt"`
		UserID    bson.ObjectId `bson:"userID"`
	}
)

// tokensC returns the tokens collection.
func tokensC() *mgo.Collection {
	return adminSession.Copy().DB("").C(tokensCName)
}

func tokensSetup() {
	c := tokensC()
	index := mgo.Index{
		Key:         []string{"createdAt"},
		Unique:      false,
		DropDups:    false,
		Background:  false,
		Sparse:      false,
		ExpireAfter: tokensExpireAfter,
	}
	err := c.EnsureIndex(index)
	if err != nil {
		log.Fatalf("MongoDB ensureIndex createdAt on tokens failed: %s.\n", err)
	}
	usrIndex := mgo.Index{
		Key:        []string{"userID"},
		Unique:     false,
		DropDups:   false,
		Background: false,
		Sparse:     false,
	}
	err = c.EnsureIndex(usrIndex)
	if err != nil {
		log.Fatalf("MongoDB ensureIndex userID on tokens failed: %s.\n", err)
	}
}

// tokenNew creates new token for user usrID.
// Returns token string and error message.
func tokenNew(usrID bson.ObjectId) (string, error) {
	var token Token
	token.ID = bson.NewObjectId()
	token.CreatedAt = time.Now()
	token.UserID = usrID
	err := tokensC().Insert(token)
	return token.ID.Hex(), err
}

// TokenCheck check if the given token is valid.
// Returns authenticated user and error.
func TokenCheck(tknID bson.ObjectId) (bson.ObjectId, error) {
	var token Token
	err := tokensC().Find(bson.M{"_id": tknID}).One(&token)
	return token.UserID, err
}
<commit_msg>tokens: small cleanup<commit_after>package db

import (
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
	"log"
	"time"
)

const (
	//tokensExpireAfter time.Duration = time.Duration(time.Second * 60 * 60 * 24 * 7) // 1 week
	// tokensExpireAfter controls how long authentication tokens are valid.
	tokensExpireAfter time.Duration = time.Duration(time.Second * 60 * 60) // 1h
)

type (
	// Token holds information associated with a given authentication token
	Token struct {
		ID        bson.ObjectId `bson:"_id"`
		CreatedAt time.Time     `bson:"createdAt"`
		UserID    bson.ObjectId `bson:"userID"`
	}
)

// tokensC returns the tokens collection.
func tokensC() *mgo.Collection {
	return adminSession.Copy().DB("").C(tokensCName)
}

func tokensSetup() {
	c := tokensC()
	index := mgo.Index{
		Key:         []string{"createdAt"},
		Unique:      false,
		DropDups:    false,
		Background:  false,
		Sparse:      false,
		ExpireAfter: tokensExpireAfter,
	}
	if err := c.EnsureIndex(index); err != nil {
		log.Fatalf("MongoDB ensureIndex createdAt on tokens failed: %s.\n", err)
	}
	usrIndex := mgo.Index{
		Key:        []string{"userID"},
		Unique:     false,
		DropDups:   false,
		Background: false,
		Sparse:     false,
	}
	if err := c.EnsureIndex(usrIndex); err != nil {
		log.Fatalf("MongoDB ensureIndex userID on tokens failed: %s.\n", err)
	}
}

// tokenNew creates new token for user usrID.
// Returns token string and error message. FIXME: 'encrypt/randomize' the string
func tokenNew(usrID bson.ObjectId) (string, error) {
	var token Token
	token.ID = bson.NewObjectId()
	token.CreatedAt = time.Now()
	token.UserID = usrID
	err := tokensC().Insert(token)
	return token.ID.Hex(), err
}

// TokenCheck check if the given token is valid.
// Returns authenticated user and error.
func TokenCheck(tknID bson.ObjectId) (bson.ObjectId, error) {
	var token Token
	err := tokensC().Find(bson.M{"_id": tknID}).One(&token)
	return token.UserID, err
}
<|endoftext|>