<commit_before>// Copyright Banrai LLC. All rights reserved. Use of this source code is
// governed by the license that can be found in the LICENSE file.

package ui

import (
	"bytes"
	"database/sql"
	"fmt"
	"github.com/Banrai/TeamWork.io/server/database"
	"io"
	"net/http"
	"strings"
)

const NO_SUCH_MESSAGE = "There is no such message among the list of available posts"

// Lookup and stream the given message back to the client
func DownloadMessage(w http.ResponseWriter, r *http.Request, db database.DBConnection, opts ...interface{}) {
	var (
		m *database.MESSAGE
		s *database.SESSION
		p *database.PERSON
		d []*database.MESSAGE_DIGEST
	)
	alert := new(Alert)
	messageFound := false

	messageId := r.URL.Query().Get("message")

	// preserve the session data, if any
	if "POST" == r.Method {
		r.ParseForm()
		sessionCode, sessionCodeExists := r.PostForm["session"]
		personCode, personCodeExists := r.PostForm["person"]
		if sessionCodeExists && personCodeExists {
			sessionId := strings.Join(sessionCode, "")
			personId := strings.Join(personCode, "")
			if len(sessionId) > 0 && len(personId) > 0 {

				fn := func(stmt map[string]*sql.Stmt) {
					session, sessionErr := ConfirmSessionCode(sessionId, stmt[database.SESSION_CLEANUP], stmt[database.SESSION_LOOKUP_BY_ID])
					if sessionErr != nil {
						alert.AsError(OTHER_ERROR)
						return
					}

					if len(session.Id) == 0 {
						alert.AsError(INVALID_SESSION)
						return
					}

					if !session.Verified {
						session.Verified = true
						if session.Update(stmt[database.SESSION_UPDATE]) != nil {
							alert.AsError(OTHER_ERROR)
							return
						}
					}

					// attempt to find the person for this session
					person, personErr := database.LookupPerson(stmt[database.PERSON_LOOKUP_BY_ID], session.PersonId)
					if personErr != nil {
						alert.AsError(OTHER_ERROR)
						return
					}

					if len(person.Id) == 0 {
						alert.AsError(UNKNOWN)
						return
					}

					if !person.Enabled {
						alert.AsError(DISABLED)
						return
					}

					if !person.Verified {
						person.Verified = true
						if person.Update(stmt[database.PERSON_UPDATE]) != nil {
							alert.AsError(OTHER_ERROR)
							return
						}
					}

					// make sure the session matches the person from the form
					if person.Id != personId {
						alert.AsError(INVALID_SESSION)
						return
					}

					// found a valid session and person
					s = session
					p = person
				}
				database.WithDatabase(db, fn)
			}
		}
	}

	if len(messageId) > 0 {
		// attempt to find the specific message
		fn := func(stmt map[string]*sql.Stmt) {
			messages, messagesErr := database.RetrieveMessages(stmt[database.MESSAGE_BY_ID], messageId, 1, 0)
			if messagesErr != nil {
				alert.AsError(OTHER_ERROR)
				return
			}
			if len(messages) > 0 {
				m = messages[0]
				messageFound = true
			}
		}
		database.WithDatabase(db, fn)
	}

	if messageFound {
		// send its contents to the client
		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s.txt", messageId))
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(m.Message)))
		io.Copy(w, bytes.NewBufferString(m.Message))
	} else {
		// go back to the list of posts
		alert.AsError(NO_SUCH_MESSAGE)

		if s == nil && p == nil {
			// retrieve the latest digests, without session/person
			fn := func(stmt map[string]*sql.Stmt) {
				messages, _ := database.RetrieveMessages(stmt[database.LATEST_MESSAGES], "", POSTS_PER_PAGE, 0)
				digests, _ := database.GetMessageDigests(stmt[database.PERSON_LOOKUP_BY_ID], stmt[database.RECIPIENTS_BY_MESSAGE], messages, "")
				d = digests
			}
			database.WithDatabase(db, fn)

			// define these as empty, so the session template renders properly
			s = new(database.SESSION)
			p = new(database.PERSON)
		} else {
			// use the session data
			fn := func(stmt map[string]*sql.Stmt) {
				messages, _ := p.LookupLatestMessages(stmt[database.LATEST_MESSAGES], POSTS_PER_PAGE, 0)
				digests, _ := database.GetMessageDigests(stmt[database.PERSON_LOOKUP_BY_ID], stmt[database.RECIPIENTS_BY_MESSAGE], messages, p.Id)
				d = digests
			}
			database.WithDatabase(db, fn)
		}

		posts := &DisplayPostsPage{Title: "Latest Posts", Alert: alert, Session: s, Person: p, Posts: d}
		ALL_POSTS_TEMPLATE.Execute(w, posts)
	}
}
<commit_msg>Armored text file extension is '.asc' not '.txt'<commit_after>// Copyright Banrai LLC. All rights reserved. Use of this source code is
// governed by the license that can be found in the LICENSE file.

package ui

import (
	"bytes"
	"database/sql"
	"fmt"
	"github.com/Banrai/TeamWork.io/server/database"
	"io"
	"net/http"
	"strings"
)

const NO_SUCH_MESSAGE = "There is no such message among the list of available posts"

// Lookup and stream the given message back to the client
func DownloadMessage(w http.ResponseWriter, r *http.Request, db database.DBConnection, opts ...interface{}) {
	var (
		m *database.MESSAGE
		s *database.SESSION
		p *database.PERSON
		d []*database.MESSAGE_DIGEST
	)
	alert := new(Alert)
	messageFound := false

	messageId := r.URL.Query().Get("message")

	// preserve the session data, if any
	if "POST" == r.Method {
		r.ParseForm()
		sessionCode, sessionCodeExists := r.PostForm["session"]
		personCode, personCodeExists := r.PostForm["person"]
		if sessionCodeExists && personCodeExists {
			sessionId := strings.Join(sessionCode, "")
			personId := strings.Join(personCode, "")
			if len(sessionId) > 0 && len(personId) > 0 {

				fn := func(stmt map[string]*sql.Stmt) {
					session, sessionErr := ConfirmSessionCode(sessionId, stmt[database.SESSION_CLEANUP], stmt[database.SESSION_LOOKUP_BY_ID])
					if sessionErr != nil {
						alert.AsError(OTHER_ERROR)
						return
					}

					if len(session.Id) == 0 {
						alert.AsError(INVALID_SESSION)
						return
					}

					if !session.Verified {
						session.Verified = true
						if session.Update(stmt[database.SESSION_UPDATE]) != nil {
							alert.AsError(OTHER_ERROR)
							return
						}
					}

					// attempt to find the person for this session
					person, personErr := database.LookupPerson(stmt[database.PERSON_LOOKUP_BY_ID], session.PersonId)
					if personErr != nil {
						alert.AsError(OTHER_ERROR)
						return
					}

					if len(person.Id) == 0 {
						alert.AsError(UNKNOWN)
						return
					}

					if !person.Enabled {
						alert.AsError(DISABLED)
						return
					}

					if !person.Verified {
						person.Verified = true
						if person.Update(stmt[database.PERSON_UPDATE]) != nil {
							alert.AsError(OTHER_ERROR)
							return
						}
					}

					// make sure the session matches the person from the form
					if person.Id != personId {
						alert.AsError(INVALID_SESSION)
						return
					}

					// found a valid session and person
					s = session
					p = person
				}
				database.WithDatabase(db, fn)
			}
		}
	}

	if len(messageId) > 0 {
		// attempt to find the specific message
		fn := func(stmt map[string]*sql.Stmt) {
			messages, messagesErr := database.RetrieveMessages(stmt[database.MESSAGE_BY_ID], messageId, 1, 0)
			if messagesErr != nil {
				alert.AsError(OTHER_ERROR)
				return
			}
			if len(messages) > 0 {
				m = messages[0]
				messageFound = true
			}
		}
		database.WithDatabase(db, fn)
	}

	if messageFound {
		// send its contents to the client
		w.Header().Set("Content-Disposition", fmt.Sprintf("attachment; filename=%s.asc", messageId))
		w.Header().Set("Content-Type", "text/plain")
		w.Header().Set("Content-Length", fmt.Sprintf("%d", len(m.Message)))
		io.Copy(w, bytes.NewBufferString(m.Message))
	} else {
		// go back to the list of posts
		alert.AsError(NO_SUCH_MESSAGE)

		if s == nil && p == nil {
			// retrieve the latest digests, without session/person
			fn := func(stmt map[string]*sql.Stmt) {
				messages, _ := database.RetrieveMessages(stmt[database.LATEST_MESSAGES], "", POSTS_PER_PAGE, 0)
				digests, _ := database.GetMessageDigests(stmt[database.PERSON_LOOKUP_BY_ID], stmt[database.RECIPIENTS_BY_MESSAGE], messages, "")
				d = digests
			}
			database.WithDatabase(db, fn)

			// define these as empty, so the session template renders properly
			s = new(database.SESSION)
			p = new(database.PERSON)
		} else {
			// use the session data
			fn := func(stmt map[string]*sql.Stmt) {
				messages, _ := p.LookupLatestMessages(stmt[database.LATEST_MESSAGES], POSTS_PER_PAGE, 0)
				digests, _ := database.GetMessageDigests(stmt[database.PERSON_LOOKUP_BY_ID], stmt[database.RECIPIENTS_BY_MESSAGE], messages, p.Id)
				d = digests
			}
			database.WithDatabase(db, fn)
		}

		posts := &DisplayPostsPage{Title: "Latest Posts", Alert: alert, Session: s, Person: p, Posts: d}
		ALL_POSTS_TEMPLATE.Execute(w, posts)
	}
}
<|endoftext|>


<commit_before>package myrouter

import (
	"fmt"
	"strings"
	"testing"
)

func TestGenerateEmptyPath(t *testing.T) {
	var path, err = generatePath("", make(map[string]string))
	if err != nil {
		t.Fail()
	}
	if path != "" {
		t.Fail()
	}
}

func TestGeneratePathWithMissingParameter(t *testing.T) {
	var path, err = generatePath("/{id}", make(map[string]string))
	if err == nil {
		t.Fail()
	}
	if path != "" {
		t.Fail()
	}
}

func TestGeneratePathWithParameter(t *testing.T) {
	var path, err = generatePath("/{id}", map[string]string{"id": "test"})
	if err != nil {
		t.Fail()
	}
	if path != "/test" {
		t.Fail()
	}
}

func TestGeneratePathWithExtraParameter(t *testing.T) {
	var path, err = generatePath("/{id}", map[string]string{"id": "test", "slug": "poltorak-dariusz"})
	if err != nil {
		t.Fail()
	}
	if path != "/test" {
		t.Fail()
	}
}

func TestGeneratePathWithTwoParameters(t *testing.T) {
	var path, err = generatePath("/{id}/{slug}", map[string]string{"id": "test", "slug": "poltorak-dariusz"})
	if err != nil {
		t.Fail()
	}
	if path != "/test/poltorak-dariusz" {
		t.Fail()
	}
}

func TestGenerateRegExpFromPathNoArgs(t *testing.T) {
	var path = "/api/user"
	var url = "/api/user"
	var expect = "/api/user"

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	if regexp.String() != expect {
		fmt.Print(regexp.String())
		fmt.Print("\n")
		t.Fail()
	}
	if !regexp.MatchString(url) {
		t.Fail()
	}
}

func TestGenerateRegExpFromPathSingleArg(t *testing.T) {
	var path = "/api/user/{id}"
	var url = "/api/user/5"
	var expect = "/api/user/([^/]+)"

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	if regexp.String() != expect {
		fmt.Print(regexp.String())
		fmt.Print("\n")
		t.Fail()
	}
	if !regexp.MatchString(url) {
		t.Fail()
	}
}

func TestGenerateRegExpFromPathSingleCopiedArg(t *testing.T) {
	var path = "/api/user/{id}/user-{id}"
	var url = "/api/user/5/user-5"
	var expect = "/api/user/([^/]+)/user-([^/]+)"

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	if regexp.String() != expect {
		fmt.Print(regexp.String())
		fmt.Print("\n")
		t.Fail()
	}
	if !regexp.MatchString(url) {
		t.Fail()
	}
}

func TestGenerateRegExpFromPathTwoArgs(t *testing.T) {
	var path = "/api/user/{id}/client-{slug}"
	var url = "/api/user/5/client-poltorak-dariusz"
	var expect = "/api/user/([^/]+)/client-([^/]+)"

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	if regexp.String() != expect {
		fmt.Print(regexp.String())
		fmt.Print("\n")
		t.Fail()
	}
	if !regexp.MatchString(url) {
		t.Fail()
	}
}

func TestGenerateRegExpFromPathNoValues(t *testing.T) {
	var path = "/api/user/5/client-poltorak-dariusz"
	var url = "/api/user/5/client-poltorak-dariusz"
	var expect = []string{}

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	var result = regexp.FindAllStringSubmatch(url, -1)
	var parameters []string
	for i := 1; i < len(result[0]); i++ {
		parameters = append(parameters, result[0][i])
	}

	if !regexp.MatchString(url) {
		t.Fail()
	}

	if !arrayCompareString(expect, parameters) {
		fmt.Print(strings.Join(parameters, ", "))
		fmt.Print("\n")
		t.Fail()
	}
}

func TestGenerateRegExpFromPathOneValue(t *testing.T) {
	var path = "/api/user/{id}"
	var url = "/api/user/5"
	var expect = []string{"5"}

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	var result = regexp.FindAllStringSubmatch(url, -1)
	var parameters []string
	for i := 1; i < len(result[0]); i++ {
		parameters = append(parameters, result[0][i])
	}

	if !regexp.MatchString(url) {
		t.Fail()
	}

	if !arrayCompareString(expect, parameters) {
		fmt.Print(strings.Join(parameters, ", "))
		fmt.Print("\n")
		t.Fail()
	}
}

func TestGenerateRegExpFromPathTwoValues(t *testing.T) {
	var path = "/api/user/{id}/client-{slug}"
	var url = "/api/user/5/client-poltorak-dariusz"
	var expect = []string{"5", "poltorak-dariusz"}

	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	var result = regexp.FindAllStringSubmatch(url, -1)
	var parameters []string
	for i := 1; i < len(result[0]); i++ {
		parameters = append(parameters, result[0][i])
	}

	if !regexp.MatchString(url) {
		t.Fail()
	}

	if !arrayCompareString(expect, parameters) {
		fmt.Print(strings.Join(parameters, ", "))
		fmt.Print("\n")
		t.Fail()
	}
}
<commit_msg>more tests<commit_after>package myrouter

import (
	"fmt"
	"strings"
	"testing"
)

func generateRegExpFromPathVerification(t *testing.T, path string, url string, expect string) {
	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	if regexp.String() != expect {
		fmt.Print(regexp.String())
		fmt.Print("\n")
		t.Fail()
	}
	if !regexp.MatchString(url) {
		t.Fail()
	}
}

func generateRegExpFromPathValueVerification(t *testing.T, path string, url string, expect []string) {
	var regexp = GenerateRegExpFromPath(path, map[string]string{})
	var result = regexp.FindAllStringSubmatch(url, -1)
	var parameters []string
	for i := 1; i < len(result[0]); i++ {
		parameters = append(parameters, result[0][i])
	}

	if !regexp.MatchString(url) {
		t.Fail()
	}

	if !arrayCompareString(expect, parameters) {
		fmt.Print(strings.Join(parameters, ", "))
		fmt.Print("\n")
		t.Fail()
	}
}

func TestGenerateEmptyPath(t *testing.T) {
	var path, err = generatePath("", make(map[string]string))
	if err != nil {
		t.Fail()
	}
	if path != "" {
		t.Fail()
	}
}

func TestGeneratePathWithMissingParameter(t *testing.T) {
	var path, err = generatePath("/{id}", make(map[string]string))
	if err == nil {
		t.Fail()
	}
	if path != "" {
		t.Fail()
	}
}

func TestGeneratePathWithParameter(t *testing.T) {
	var path, err = generatePath("/{id}", map[string]string{"id": "test"})
	if err != nil {
		t.Fail()
	}
	if path != "/test" {
		t.Fail()
	}
}

func TestGeneratePathWithExtraParameter(t *testing.T) {
	var path, err = generatePath("/{id}", map[string]string{"id": "test", "slug": "poltorak-dariusz"})
	if err != nil {
		t.Fail()
	}
	if path != "/test" {
		t.Fail()
	}
}

func TestGeneratePathWithTwoParameters(t *testing.T) {
	var path, err = generatePath("/{id}/{slug}", map[string]string{"id": "test", "slug": "poltorak-dariusz"})
	if err != nil {
		t.Fail()
	}
	if path != "/test/poltorak-dariusz" {
		t.Fail()
	}
}

func TestGenerateRegExpFromPathNoArgs(t *testing.T) {
	var path = "/api/user"
	var url = "/api/user"
	var expect = "/api/user"

	generateRegExpFromPathVerification(t, path, url, expect)
}

func TestGenerateRegExpFromPathSingleArg(t *testing.T) {
	var path = "/api/user/{id}"
	var url = "/api/user/5"
	var expect = "/api/user/([^/]+)"

	generateRegExpFromPathVerification(t, path, url, expect)
}

func TestGenerateRegExpFromPathSingleCopiedArg(t *testing.T) {
	var path = "/api/user/{id}/user-{id}"
	var url = "/api/user/5/user-5"
	var expect = "/api/user/([^/]+)/user-([^/]+)"

	generateRegExpFromPathVerification(t, path, url, expect)
}

func TestGenerateRegExpFromPathTwoArgs(t *testing.T) {
	var path = "/api/user/{id}/client-{slug}"
	var url = "/api/user/5/client-poltorak-dariusz"
	var expect = "/api/user/([^/]+)/client-([^/]+)"

	generateRegExpFromPathVerification(t, path, url, expect)
}

func TestGenerateRegExpFromPathNoValues(t *testing.T) {
	var path = "/api/user/5/client-poltorak-dariusz"
	var url = "/api/user/5/client-poltorak-dariusz"
	var expect = []string{}

	generateRegExpFromPathValueVerification(t, path, url, expect)
}

func TestGenerateRegExpFromPathOneValue(t *testing.T) {
	var path = "/api/user/{id}"
	var url = "/api/user/5"
	var expect = []string{"5"}

	generateRegExpFromPathValueVerification(t, path, url, expect)
}

func TestGenerateRegExpFromPathTwoValues(t *testing.T) {
	var path = "/api/user/{id}/client-{slug}"
	var url = "/api/user/5/client-poltorak-dariusz"
	var expect = []string{"5", "poltorak-dariusz"}

	generateRegExpFromPathValueVerification(t, path, url, expect)
}
<|endoftext|>


<commit_before>/*-
 * Copyright (c) 2012 Caoimhe Chaos <caoimhechaos@protonmail.com>,
 *                    Ancient Solutions. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions  of source code must retain  the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions  in   binary  form  must   reproduce  the  above
 *    copyright  notice, this  list  of conditions  and the  following
 *    disclaimer in the  documentation and/or other materials provided
 *    with the distribution.
 *
 * THIS  SOFTWARE IS  PROVIDED BY  ANCIENT SOLUTIONS  AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO,  THE IMPLIED WARRANTIES OF  MERCHANTABILITY AND FITNESS
 * FOR A  PARTICULAR PURPOSE  ARE DISCLAIMED.  IN  NO EVENT  SHALL THE
 * FOUNDATION  OR CONTRIBUTORS  BE  LIABLE FOR  ANY DIRECT,  INDIRECT,
 * INCIDENTAL,   SPECIAL,    EXEMPLARY,   OR   CONSEQUENTIAL   DAMAGES
 * (INCLUDING, BUT NOT LIMITED  TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE,  DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT  LIABILITY,  OR  TORT  (INCLUDING NEGLIGENCE  OR  OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package geocolo

import (
	"database/sql"
	"errors"
	"fmt"
	"net"
	"strings"

	_ "github.com/lib/pq"
	libgeo "github.com/nranchev/go-libGeoIP"
)

var (
	addr192 = net.IPv4(192, 168, 0, 0)
	addr172 = net.IPv4(172, 16, 0, 0)
	addr10  = net.IPv4(10, 0, 0, 0)

	mask8  = net.IPv4Mask(255, 0, 0, 0)
	mask16 = net.IPv4Mask(255, 255, 0, 0)

	net192 = net.IPNet{
		IP:   addr192,
		Mask: mask16,
	}
	net172 = net.IPNet{
		IP:   addr172,
		Mask: mask16,
	}
	net10 = net.IPNet{
		IP:   addr10,
		Mask: mask8,
	}
)

type GeoProximityService struct {
	conn            *sql.DB
	gi              *libgeo.GeoIP
	rfc1918_country string
}

func isRFC1918(addr string) bool {
	var ip net.IP = net.ParseIP(addr)

	if ip.To4() == nil {
		return false
	}
	return net192.Contains(ip) || net172.Contains(ip) || net10.Contains(ip)
}

func NewGeoProximityService(
	config *GeoProximityServiceConfig) (*GeoProximityService, error) {
	var gi *libgeo.GeoIP
	var dsn string
	var rfc1918 string
	var c *sql.DB
	var err error

	dsn = fmt.Sprintf("user=%s dbname=%s host=%s port=%d",
		*config.User, *config.Dbname, *config.Host, *config.Port)

	if config.Password != nil {
		dsn += fmt.Sprintf(" password=%s", *config.Password)
	}

	if config.Sslmode != nil {
		dsn += fmt.Sprintf(" sslmode=%s", *config.Sslmode)
	}

	c, err = sql.Open("postgres", dsn)
	if err != nil {
		return nil, err
	}

	if config.GeoipPath != nil {
		gi, err = libgeo.Load(*config.GeoipPath)
		if err != nil {
			return nil, err
		}
	}

	if config.Rfc1918Country != nil {
		rfc1918 = strings.ToUpper(*config.Rfc1918Country)
	}

	if rfc1918 == "" {
		rfc1918 = "CH"
	}

	return &GeoProximityService{
		conn:            c,
		gi:              gi,
		rfc1918_country: rfc1918,
	}, nil
}

func (self *GeoProximityService) GetProximity(req GeoProximityRequest,
	res *GeoProximityResponse) error {
	var rows *sql.Rows
	var err error

	if req.Origin == nil {
		return errors.New("No origin specified")
	}

	if len(req.Candidates) > 0 {
		var fullsql string
		var valuecollection []string
		var value string

		for _, value = range req.Candidates {
			value = strings.ToUpper(value)
			if len(value) == 2 && value[0] >= 'A' &&
				value[0] <= 'Z' && value[1] >= 'A' &&
				value[1] <= 'Z' {
				valuecollection = append(valuecollection,
					"'"+value+"'")
			}
		}

		fullsql = strings.Join(valuecollection, ",")

		rows, err = self.conn.Query("SELECT s.iso_a2, distance("+
			"s.the_geom, (SELECT g.the_geom FROM geoborders g "+
			"WHERE g.iso_a2 = $1 ) ) AS dist FROM geoborders s "+
			"WHERE s.iso_a2 IN ( "+fullsql+" ) ORDER BY "+
			"dist ASC", strings.ToUpper(*req.Origin))
	} else {
		rows, err = self.conn.Query("SELECT s.iso_a2, distance("+
			"s.the_geom, (SELECT g.the_geom FROM geoborders g "+
			"WHERE g.iso_a2 = $1 ) ) AS dist FROM geoborders s "+
			"ORDER BY dist ASC", strings.ToUpper(*req.Origin))
	}
	if err != nil {
		return err
	}

	for rows.Next() {
		var detail *GeoProximityDetail = new(GeoProximityDetail)
		detail.Country = new(string)
		detail.Distance = new(float64)

		err = rows.Scan(detail.Country, detail.Distance)
		if err != nil {
			return err
		}

		// Go RPC hates 0 values :S
		if *detail.Distance == 0 {
			*detail.Distance -= 0.001
		}

		if res.Closest == nil {
			res.Closest = detail.Country
		}

		if *req.DetailedResponse {
			res.FullMap = append(res.FullMap, detail)
		}
	}

	return nil
}

// Request which list of destination IPs are closest to a given source
// IP. Takes a radius around which results are returned but we always
// return at least one.
func (self *GeoProximityService) GetProximityByIP(req GeoProximityByIPRequest,
	res *GeoProximityByIPResponse) error {
	var addrlocations = make(map[string][]*GeoProximityByIPDetail)
	var locdata []string = make([]string, 0)
	var maxdistance float64
	var loc *libgeo.Location
	var origin string

	var fullsql string
	var rows *sql.Rows
	var initclosest int
	var err error

	if self.gi == nil {
		return errors.New("GeoIP not loaded")
	}

	if req.Origin == nil {
		return errors.New("No origin specified")
	}

	// First, determine the geodata for all of the given IPs and
	// sort them into a map.
	for _, addr := range req.Candidates {
		loc = self.gi.GetLocationByIP(addr)

		if loc == nil {
			// No country data? Always return…
			res.Closest = append(res.Closest, addr)
		} else {
			var cc string = strings.ToUpper(
				loc.CountryCode)
			var detail *GeoProximityByIPDetail = new(GeoProximityByIPDetail)
			var ok bool

			detail.Ip = new(string)
			*detail.Ip = addr

			_, ok = addrlocations[cc]
			if !ok {
				addrlocations[cc] =
					make([]*GeoProximityByIPDetail,
						0)
			}

			addrlocations[cc] =
				append(addrlocations[cc], detail)
		}
	}

	initclosest = len(res.Closest)

	for cc, _ := range addrlocations {
		locdata = append(locdata, "'"+cc+"'")
	}

	fullsql = strings.Join(locdata, ",")

	if isRFC1918(*req.Origin) {
		origin = self.rfc1918_country
	} else {
		// Now let's figure out where the request came from.
		// TODO(caoimhe): Handle RFC1918 IPs.
		loc = self.gi.GetLocationByIP(*req.Origin)

		if loc == nil {
			// Fail open: return all IPs.
			// TODO(caoimhe): Filter out RFC1918 IPs.
			res.Closest = req.Candidates
			return nil
		}
		origin = strings.ToUpper(loc.CountryCode)
	}

	rows, err = self.conn.Query("SELECT s.iso_a2, distance("+
		"s.the_geom, (SELECT g.the_geom FROM geoborders g "+
		"WHERE g.iso_a2 = $1 ) ) AS dist FROM geoborders s "+
		"WHERE s.iso_a2 IN ( "+fullsql+" ) ORDER BY "+
		"dist ASC", origin)
	if err != nil {
		return err
	}

	for rows.Next() {
		var country string
		var distance float64
		var detail *GeoProximityByIPDetail

		err = rows.Scan(&country, &distance)
		if err != nil {
			return err
		}

		if country != origin {
			distance += 0.01
		}

		// Go RPC hates 0 values :S
		if distance == 0 {
			distance += 0.001
		}

		for _, detail = range addrlocations[country] {
			detail.Distance = new(float64)
			*detail.Distance = distance

			if len(res.Closest) == initclosest {
				maxdistance = distance

				if req.MaxDistance != nil {
					maxdistance +=
						*req.MaxDistance
				}
			}

			if distance <= maxdistance {
				res.Closest = append(res.Closest,
					*detail.Ip)
			}

			if *req.DetailedResponse {
				res.FullMap = append(res.FullMap, detail)
			}
		}
	}

	return nil
}
<commit_msg>Check whether DetailedResponse exists before dereferencing it.<commit_after>/*-
 * Copyright (c) 2012 Caoimhe Chaos <caoimhechaos@protonmail.com>,
 *                    Ancient Solutions. All rights reserved.
 *
 * Redistribution and use in source and binary forms, with or without
 * modification, are permitted provided that the following conditions
 * are met:
 *
 * 1. Redistributions  of source code must retain  the above copyright
 *    notice, this list of conditions and the following disclaimer.
 * 2. Redistributions  in   binary  form  must   reproduce  the  above
 *    copyright  notice, this  list  of conditions  and the  following
 *    disclaimer in the  documentation and/or other materials provided
 *    with the distribution.
 *
 * THIS  SOFTWARE IS  PROVIDED BY  ANCIENT SOLUTIONS  AND CONTRIBUTORS
 * ``AS IS'' AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
 * LIMITED TO,  THE IMPLIED WARRANTIES OF  MERCHANTABILITY AND FITNESS
 * FOR A  PARTICULAR PURPOSE  ARE DISCLAIMED.  IN  NO EVENT  SHALL THE
 * FOUNDATION  OR CONTRIBUTORS  BE  LIABLE FOR  ANY DIRECT,  INDIRECT,
 * INCIDENTAL,   SPECIAL,    EXEMPLARY,   OR   CONSEQUENTIAL   DAMAGES
 * (INCLUDING, BUT NOT LIMITED  TO, PROCUREMENT OF SUBSTITUTE GOODS OR
 * SERVICES; LOSS OF USE,  DATA, OR PROFITS; OR BUSINESS INTERRUPTION)
 * HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT,
 * STRICT  LIABILITY,  OR  TORT  (INCLUDING NEGLIGENCE  OR  OTHERWISE)
 * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED
 * OF THE POSSIBILITY OF SUCH DAMAGE.
 */

package geocolo

import (
	"database/sql"
	"errors"
	"fmt"
	"net"
	"strings"

	_ "github.com/lib/pq"
	libgeo "github.com/nranchev/go-libGeoIP"
)

var (
	addr192 = net.IPv4(192, 168, 0, 0)
	addr172 = net.IPv4(172, 16, 0, 0)
	addr10  = net.IPv4(10, 0, 0, 0)

	mask8  = net.IPv4Mask(255, 0, 0, 0)
	mask16 = net.IPv4Mask(255, 255, 0, 0)

	net192 = net.IPNet{
		IP:   addr192,
		Mask: mask16,
	}
	net172 = net.IPNet{
		IP:   addr172,
		Mask: mask16,
	}
	net10 = net.IPNet{
		IP:   addr10,
		Mask: mask8,
	}
)

type GeoProximityService struct {
	conn            *sql.DB
	gi              *libgeo.GeoIP
	rfc1918_country string
}

func isRFC1918(addr string) bool {
	var ip net.IP = net.ParseIP(addr)

	if ip.To4() == nil {
		return false
	}
	return net192.Contains(ip) || net172.Contains(ip) || net10.Contains(ip)
}

func NewGeoProximityService(
	config *GeoProximityServiceConfig) (*GeoProximityService, error) {
	var gi *libgeo.GeoIP
	var dsn string
	var rfc1918 string
	var c *sql.DB
	var err error

	dsn = fmt.Sprintf("user=%s dbname=%s host=%s port=%d",
		*config.User, *config.Dbname, *config.Host, *config.Port)

	if config.Password != nil {
		dsn += fmt.Sprintf(" password=%s", *config.Password)
	}

	if config.Sslmode != nil {
		dsn += fmt.Sprintf(" sslmode=%s", *config.Sslmode)
	}

	c, err = sql.Open("postgres", dsn)
	if err != nil {
		return nil, err
	}

	if config.GeoipPath != nil {
		gi, err = libgeo.Load(*config.GeoipPath)
		if err != nil {
			return nil, err
		}
	}

	if config.Rfc1918Country != nil {
		rfc1918 = strings.ToUpper(*config.Rfc1918Country)
	}

	if rfc1918 == "" {
		rfc1918 = "CH"
	}

	return &GeoProximityService{
		conn:            c,
		gi:              gi,
		rfc1918_country: rfc1918,
	}, nil
}

func (self *GeoProximityService) GetProximity(req GeoProximityRequest,
	res *GeoProximityResponse) error {
	var rows *sql.Rows
	var err error

	if req.Origin == nil {
		return errors.New("No origin specified")
	}

	if len(req.Candidates) > 0 {
		var fullsql string
		var valuecollection []string
		var value string

		for _, value = range req.Candidates {
			value = strings.ToUpper(value)
			if len(value) == 2 && value[0] >= 'A' &&
				value[0] <= 'Z' && value[1] >= 'A' &&
				value[1] <= 'Z' {
				valuecollection = append(valuecollection,
					"'"+value+"'")
			}
		}

		fullsql = strings.Join(valuecollection, ",")

		rows, err = self.conn.Query("SELECT s.iso_a2, distance("+
			"s.the_geom, (SELECT g.the_geom FROM geoborders g "+
			"WHERE g.iso_a2 = $1 ) ) AS dist FROM geoborders s "+
			"WHERE s.iso_a2 IN ( "+fullsql+" ) ORDER BY "+
			"dist ASC", strings.ToUpper(*req.Origin))
	} else {
		rows, err = self.conn.Query("SELECT s.iso_a2, distance("+
			"s.the_geom, (SELECT g.the_geom FROM geoborders g "+
			"WHERE g.iso_a2 = $1 ) ) AS dist FROM geoborders s "+
			"ORDER BY dist ASC", strings.ToUpper(*req.Origin))
	}
	if err != nil {
		return err
	}

	for rows.Next() {
		var detail *GeoProximityDetail = new(GeoProximityDetail)
		detail.Country = new(string)
		detail.Distance = new(float64)

		err = rows.Scan(detail.Country, detail.Distance)
		if err != nil {
			return err
		}

		// Go RPC hates 0 values :S
		if *detail.Distance == 0 {
			*detail.Distance -= 0.001
		}

		if res.Closest == nil {
			res.Closest = detail.Country
		}

		if req.DetailedResponse != nil && *req.DetailedResponse {
			res.FullMap = append(res.FullMap, detail)
		}
	}

	return nil
}

// Request which list of destination IPs are closest to a given source
// IP. Takes a radius around which results are returned but we always
// return at least one.
func (self *GeoProximityService) GetProximityByIP(req GeoProximityByIPRequest,
	res *GeoProximityByIPResponse) error {
	var addrlocations = make(map[string][]*GeoProximityByIPDetail)
	var locdata []string = make([]string, 0)
	var maxdistance float64
	var loc *libgeo.Location
	var origin string

	var fullsql string
	var rows *sql.Rows
	var initclosest int
	var err error

	if self.gi == nil {
		return errors.New("GeoIP not loaded")
	}

	if req.Origin == nil {
		return errors.New("No origin specified")
	}

	// First, determine the geodata for all of the given IPs and
	// sort them into a map.
	for _, addr := range req.Candidates {
		loc = self.gi.GetLocationByIP(addr)

		if loc == nil {
			// No country data? Always return…
			res.Closest = append(res.Closest, addr)
		} else {
			var cc string = strings.ToUpper(
				loc.CountryCode)
			var detail *GeoProximityByIPDetail = new(GeoProximityByIPDetail)
			var ok bool

			detail.Ip = new(string)
			*detail.Ip = addr

			_, ok = addrlocations[cc]
			if !ok {
				addrlocations[cc] =
					make([]*GeoProximityByIPDetail,
						0)
			}

			addrlocations[cc] =
				append(addrlocations[cc], detail)
		}
	}

	initclosest = len(res.Closest)

	for cc, _ := range addrlocations {
		locdata = append(locdata, "'"+cc+"'")
	}

	fullsql = strings.Join(locdata, ",")

	if isRFC1918(*req.Origin) {
		origin = self.rfc1918_country
	} else {
		// Now let's figure out where the request came from.
		// TODO(caoimhe): Handle RFC1918 IPs.
		loc = self.gi.GetLocationByIP(*req.Origin)

		if loc == nil {
			// Fail open: return all IPs.
			// TODO(caoimhe): Filter out RFC1918 IPs.
			res.Closest = req.Candidates
			return nil
		}
		origin = strings.ToUpper(loc.CountryCode)
	}

	rows, err = self.conn.Query("SELECT s.iso_a2, distance("+
		"s.the_geom, (SELECT g.the_geom FROM geoborders g "+
		"WHERE g.iso_a2 = $1 ) ) AS dist FROM geoborders s "+
		"WHERE s.iso_a2 IN ( "+fullsql+" ) ORDER BY "+
		"dist ASC", origin)
	if err != nil {
		return err
	}

	for rows.Next() {
		var country string
		var distance float64
		var detail *GeoProximityByIPDetail

		err = rows.Scan(&country, &distance)
		if err != nil {
			return err
		}

		if country != origin {
			distance += 0.01
		}

		// Go RPC hates 0 values :S
		if distance == 0 {
			distance += 0.001
		}

		for _, detail = range addrlocations[country] {
			detail.Distance = new(float64)
			*detail.Distance = distance

			if len(res.Closest) == initclosest {
				maxdistance = distance

				if req.MaxDistance != nil {
					maxdistance +=
						*req.MaxDistance
				}
			}

			if distance <= maxdistance {
				res.Closest = append(res.Closest,
					*detail.Ip)
			}

			if req.DetailedResponse != nil && *req.DetailedResponse {
				res.FullMap = append(res.FullMap, detail)
			}
		}
	}

	return nil
}
<|endoftext|>


<commit_before>package pixel_test

import (
	"testing"

	"github.com/faiface/pixel"
)

func TestResizeRect(t *testing.T) {
	testCases := []pixel.Rect{
		pixel.R(-10, -10, 10, 10),
		pixel.R(10, 10, 30, 30),
	}

	resizeAroundCenterAnswers := []pixel.Rect{
		pixel.R(-5, -5, 5, 5),
		pixel.R(15, 15, 25, 25),
	}

	for i, rect := range testCases {
		answer := resizeAroundCenterAnswers[i]

		// resize rectangle by 50% anchored at it's current center point
		resizedRect := rect.Resized(rect.Center(), rect.Size().Scaled(0.5))

		if resizedRect != answer {
			t.Errorf("Rectangle resize was incorrect, got %v, want: %v.", resizedRect, answer)
		}
	}

	resizeAroundMinAnswers := []pixel.Rect{
		pixel.R(-10, -10, 0, 0),
		pixel.R(10, 10, 20, 20),
	}

	for i, rect := range testCases {
		answer := resizeAroundMinAnswers[i]

		// resize rectangle by 50% anchored at it's Min coordinate
		resizedRect := rect.Resized(rect.Min, rect.Size().Scaled(0.5))

		if resizedRect != answer {
			t.Errorf("Rectangle resize was incorrect, got %v, want: %v.", resizedRect, answer)
		}
	}

	resizeAroundMaxAnswers := []pixel.Rect{
		pixel.R(0, 0, 10, 10),
		pixel.R(20, 20, 30, 30),
	}

	for i, rect := range testCases {
		answer := resizeAroundMaxAnswers[i]

		// resize rectangle by 50% anchored at it's Max coordinate
		resizedRect := rect.Resized(rect.Max, rect.Size().Scaled(0.5))

		if resizedRect != answer {
			t.Errorf("Rectangle resize was incorrect, got %v, want: %v.", resizedRect, answer)
		}
	}

	resizeAroundMiddleOfLeftSideAnswers := []pixel.Rect{
		pixel.R(-10, -5, 0, 5),
		pixel.R(10, 15, 20, 25),
	}

	for i, rect := range testCases {
		answer := resizeAroundMiddleOfLeftSideAnswers[i]

		// resize rectangle by 50% anchored at the middle of it's left side
		resizedRect := rect.Resized(pixel.V(rect.Min.X, rect.Center().Y), rect.Size().Scaled(0.5))

		if resizedRect != answer {
			t.Errorf("Rectangle resize was incorrect, got %v, want: %v.", resizedRect, answer)
		}
	}

	resizeAroundOriginAnswers := []pixel.Rect{
		pixel.R(-5, -5, 5, 5),
		pixel.R(5, 5, 15, 15),
	}

	for i, rect := range testCases {
		answer := resizeAroundOriginAnswers[i]

		// resize rectangle by 50% anchored at the origin
		resizedRect := rect.Resized(pixel.ZV, rect.Size().Scaled(0.5))

		if resizedRect != answer {
			t.Errorf("Rectangle resize was incorrect, got %v, want: %v.", resizedRect, answer)
		}
	}
}
<commit_msg>refactor and add test cases<commit_after>package pixel_test

import (
	"testing"

	"github.com/faiface/pixel"
)

type rectTestInput struct {
	name string
	rect pixel.Rect
}

type rectTestTransform struct {
	name string
	f    func(pixel.Rect) pixel.Rect
}

func TestResizeRect(t *testing.T) {

	// rectangles
	squareAroundOrigin := rectTestInput{"square around origin", pixel.R(-10, -10, 10, 10)}
	squareAround2020 := rectTestInput{"square around 20, 20", pixel.R(10, 10, 30, 30)}
	rectangleAroundOrigin := rectTestInput{"rectangle around origin", pixel.R(-20, -10, 20, 10)}
	rectangleAround2020 := rectTestInput{"rectangle around 20, 20", pixel.R(0, 10, 40, 30)}

	// resize transformations
	resizeByHalfAroundCenter := rectTestTransform{"by half around center", func(rect pixel.Rect) pixel.Rect {
		return rect.Resized(rect.Center(), rect.Size().Scaled(0.5))
	}}
	resizeByHalfAroundMin := rectTestTransform{"by half around Min", func(rect pixel.Rect) pixel.Rect {
		return rect.Resized(rect.Min, rect.Size().Scaled(0.5))
	}}
	resizeByHalfAroundMax := rectTestTransform{"by half around Max", func(rect pixel.Rect) pixel.Rect {
		return rect.Resized(rect.Max, rect.Size().Scaled(0.5))
	}}
	resizeByHalfAroundMiddleOfLeftSide := rectTestTransform{"by half around middle of left side", func(rect pixel.Rect) pixel.Rect {
		return rect.Resized(pixel.V(rect.Min.X, rect.Center().Y), rect.Size().Scaled(0.5))
	}}
	resizeByHalfAroundOrigin := rectTestTransform{"by half around the origin", func(rect pixel.Rect) pixel.Rect {
		return rect.Resized(pixel.ZV, rect.Size().Scaled(0.5))
	}}

	testCases := []struct {
		input     rectTestInput
		transform rectTestTransform
		answer    pixel.Rect
	}{
		{squareAroundOrigin, resizeByHalfAroundCenter, pixel.R(-5, -5, 5, 5)},
		{squareAround2020, resizeByHalfAroundCenter, pixel.R(15, 15, 25, 25)},
		{rectangleAroundOrigin, resizeByHalfAroundCenter, pixel.R(-10, -5, 10, 5)},
		{rectangleAround2020, resizeByHalfAroundCenter, pixel.R(10, 15, 30, 25)},

		{squareAroundOrigin, resizeByHalfAroundMin, pixel.R(-10, -10, 0, 0)},
		{squareAround2020, resizeByHalfAroundMin, pixel.R(10, 10, 20, 20)},
		{rectangleAroundOrigin, resizeByHalfAroundMin, pixel.R(-20, -10, 0, 0)},
		{rectangleAround2020, resizeByHalfAroundMin, pixel.R(0, 10, 20, 20)},

		{squareAroundOrigin, resizeByHalfAroundMax, pixel.R(0, 0, 10, 10)},
		{squareAround2020, resizeByHalfAroundMax, pixel.R(20, 20, 30, 30)},
		{rectangleAroundOrigin, resizeByHalfAroundMax, pixel.R(0, 0, 20, 10)},
		{rectangleAround2020, resizeByHalfAroundMax, pixel.R(20, 20, 40, 30)},

		{squareAroundOrigin, resizeByHalfAroundMiddleOfLeftSide, pixel.R(-10, -5, 0, 5)},
		{squareAround2020, resizeByHalfAroundMiddleOfLeftSide, pixel.R(10, 15, 20, 25)},
		{rectangleAroundOrigin, resizeByHalfAroundMiddleOfLeftSide, pixel.R(-20, -5, 0, 5)},
		{rectangleAround2020, resizeByHalfAroundMiddleOfLeftSide, pixel.R(0, 15, 20, 25)},

		{squareAroundOrigin, resizeByHalfAroundOrigin, pixel.R(-5, -5, 5, 5)},
		{squareAround2020, resizeByHalfAroundOrigin, pixel.R(5, 5, 15, 15)},
		{rectangleAroundOrigin, resizeByHalfAroundOrigin, pixel.R(-10, -5, 10, 5)},
		{rectangleAround2020, resizeByHalfAroundOrigin, pixel.R(0, 5, 20, 15)},
	}

	for _, testCase := range testCases {
		testResult := testCase.transform.f(testCase.input.rect)
		if testResult != testCase.answer {
			t.Errorf("Resizing %s %s failed, got: %v, wanted: %v\n", testCase.input.name, testCase.transform.name, testResult, testCase.answer)
		}
	}
}
<|endoftext|>


<commit_before>package main

var directDomainList = []string{
	// 视频
	"xunlei.com", // 迅雷
	"kankan.com",
	"kanimg.com",

	"tdimg.com", // 土豆
	"tudou.com",
	"tudouui.com",

	"soku.com", // 优酷
	"youku.com",
	"ykimg.com",

	"ku6.cn", // 酷六
	"ku6.com",
	"ku6cdn.com",
	"ku6img.com",

	// 电商
	"z.cn",
	"amazon.cn",

	"360buy.com",
	"360buyimg.com",

	"51buy.com",
	"icson.com",

	"dangdang.com",
	"ddimg.cn",

	"yihaodian.com",
	"yihaodianimg.com",

	"paipai.com",
	"paipaiimg.com",

	"tmall.com",
	"taobao.com",
	"taobaocdn.com",
	"tbcdn.cn",
	"etao.com",

	"aicdn.com",
	"alicdn.com",
	"alimama.cn",
	"alimama.com",
	"alipay.com",
	"alipayobjects.com",

	// 银行
	"bankcomm.com",
	"bankofchina.com",
	"95559.com.cn",
	"abchina.com",
	"95599.cn",
	"boc.cn",
	"ccb.com",
	"cmbchina.com",
	"icbc.com.cn",
	"spdb.com.cn",

	// 社交
	"douban.com",
	"t.cn",
	"weibo.com",
	"zhihu.com",
	"v2ex.com",
	"kaixin001.com",
	"qq.com",

	"renren.com",
	"rrimg.com",
	"xiaonei.com",
	"xnimg.cn",
	"xnpic.com",

	"dianping.com", // 点评
	"dpfile.com",

	"huaban.com", // 又拍云的几个
	"yupoo.com",
	"upyun.com",
	"upaiyun.com",

	// 新闻门户
	"ifanr.cn",
	"ifanr.com",
	"163.com",
	"hexun.com",
	"sina.com.cn",
	"sinaapp.com",
	"sinaimg.cn",
	"sinajs.cn",
	"sohu.com",
	"solidot.org",

	// 搜索
	"bing.com",
	"bing.com.cn",
	"baidu.com",
	"bdstatic.com",
	"bdimg.com",
	"youdao.com",
	"sogou.com",

	// Apple
	"apple.com",
	"apple.com.cn",
	"icloud.com",

	// 微软
	"hotmail.com",
	"live.com",
	"outlook.com",
	"microsoft.com",

	// 其他
	"macromedia.com",
	"mmcdn.cn",
	"12306.cn",
}
<commit_msg>Add jd.com to direct list.<commit_after>package main

var directDomainList = []string{
	// 视频
	"xunlei.com", // 迅雷
	"kankan.com",
	"kanimg.com",

	"tdimg.com", // 土豆
	"tudou.com",
	"tudouui.com",

	"soku.com", // 优酷
	"youku.com",
	"ykimg.com",

	"ku6.cn", // 酷六
	"ku6.com",
	"ku6cdn.com",
	"ku6img.com",

	// 电商
	"z.cn",
	"amazon.cn",

	"360buy.com",
	"360buyimg.com",
	"jd.com",

	"51buy.com",
	"icson.com",

	"dangdang.com",
	"ddimg.cn",

	"yihaodian.com",
	"yihaodianimg.com",

	"paipai.com",
	"paipaiimg.com",

	"tmall.com",
	"taobao.com",
	"taobaocdn.com",
	"tbcdn.cn",
	"etao.com",

	"aicdn.com",
	"alicdn.com",
	"alimama.cn",
	"alimama.com",
	"alipay.com",
	"alipayobjects.com",

	// 银行
	"bankcomm.com",
	"bankofchina.com",
	"95559.com.cn",
	"abchina.com",
	"95599.cn",
	"boc.cn",
	"ccb.com",
	"cmbchina.com",
	"icbc.com.cn",
	"spdb.com.cn",

	// 社交
	"douban.com",
	"t.cn",
	"weibo.com",
	"zhihu.com",
	"v2ex.com",
	"kaixin001.com",
	"qq.com",

	"renren.com",
	"rrimg.com",
	"xiaonei.com",
	"xnimg.cn",
	"xnpic.com",

	"dianping.com", // 点评
	"dpfile.com",

	"huaban.com", // 又拍云的几个
	"yupoo.com",
	"upyun.com",
	"upaiyun.com",

	// 新闻门户
	"ifanr.cn",
	"ifanr.com",
	"163.com",
	"hexun.com",
	"sina.com.cn",
	"sinaapp.com",
	"sinaimg.cn",
	"sinajs.cn",
	"sohu.com",
	"solidot.org",

	// 搜索
	"bing.com",
	"bing.com.cn",
	"baidu.com",
	"bdstatic.com",
	"bdimg.com",
	"youdao.com",
	"sogou.com",

	// Apple
	"apple.com",
	"apple.com.cn",
	"icloud.com",

	// 微软
	"hotmail.com",
	"live.com",
	"outlook.com",
	"microsoft.com",

	// 其他
	"macromedia.com",
	"mmcdn.cn",
	"12306.cn",
}
<|endoftext|>


<commit_before>// Copyright 2018 Jigsaw Operations LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package shadowsocks

import (
	"bytes"
	"errors"
	"net"
	"sync"
	"testing"
	"time"

	logging "github.com/op/go-logging"
	"github.com/shadowsocks/go-shadowsocks2/core"
	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
)

const timeout = 5 * time.Minute

var clientAddr = net.UDPAddr{IP: []byte{192, 0, 2, 1}, Port: 12345}
var targetAddr = net.UDPAddr{IP: []byte{192, 0, 2, 2}, Port: 54321}
var dnsAddr = net.UDPAddr{IP: []byte{192, 0, 2, 3}, Port: 53}
var natCipher shadowaead.Cipher

func init() {
	coreCipher, _ := core.PickCipher(testCipher, nil, "test password")
	natCipher = coreCipher.(shadowaead.Cipher)
}

type packet struct {
	addr    net.Addr
	payload []byte
	err     error
}

type fakePacketConn struct {
	net.PacketConn
	send     chan packet
	recv     chan packet
	deadline time.Time
}

func makePacketConn() *fakePacketConn {
	return &fakePacketConn{
		send: make(chan packet, 1),
		recv: make(chan packet),
	}
}

func (conn *fakePacketConn) SetReadDeadline(deadline time.Time) error {
	conn.deadline = deadline
	return nil
}

func (conn *fakePacketConn) WriteTo(payload []byte, addr net.Addr) (int, error) {
	conn.send <- packet{addr, payload, nil}
	return len(payload), nil
}

func (conn *fakePacketConn) ReadFrom(buffer []byte) (int, net.Addr, error) {
	pkt, ok := <-conn.recv
	if !ok {
		return 0, nil, errors.New("Receive closed")
	}
	n := copy(buffer, pkt.payload)
	if n < len(pkt.payload) {
		return n, pkt.addr, errors.New("Buffer was too short")
	}
	return n, pkt.addr, pkt.err
}

func (conn *fakePacketConn) Close() error {
	close(conn.send)
	close(conn.recv)
	return nil
}

func assertAlmostEqual(t *testing.T, a, b time.Time) {
	delta := a.Sub(b)
	limit := 100 * time.Millisecond
	if delta > limit || -delta > limit {
		t.Errorf("Times are not close: %v, %v", a, b)
	}
}

func TestNATEmpty(t *testing.T) {
	nat := newNATmap(timeout, &probeTestMetrics{}, &sync.WaitGroup{})
	if nat.Get("foo") != nil {
		t.Error("Expected nil value from empty NAT map")
	}
}

func setup() (*fakePacketConn, *fakePacketConn, *natconn) {
	nat := newNATmap(timeout, &probeTestMetrics{}, &sync.WaitGroup{})
	clientConn := makePacketConn()
	targetConn := makePacketConn()
	nat.Add(&clientAddr, clientConn, natCipher, targetConn, "ZZ", "key id")
	entry := nat.Get(clientAddr.String())
	return clientConn, targetConn, entry
}

func TestNATGet(t *testing.T) {
	_, targetConn, entry := setup()
	if entry == nil {
		t.Fatal("Failed to find target conn")
	}
	if entry.PacketConn != targetConn {
		t.Error("Mismatched connection returned")
	}
}

func TestNATWrite(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate one generic packet being sent
	buf := []byte{1}
	entry.WriteTo([]byte{1}, &targetAddr)
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(timeout))
	sent := <-targetConn.send
	if !bytes.Equal(sent.payload, buf) {
		t.Errorf("Mismatched payload: %v != %v", sent.payload, buf)
	}
	if sent.addr != &targetAddr {
		t.Errorf("Mismatched address: %v != %v", sent.addr, &targetAddr)
	}
}

func TestNATWriteDNS(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate one DNS query being sent.
	buf := []byte{1}
	entry.WriteTo(buf, &dnsAddr)
	// DNS-only connections have a fixed timeout of 17 seconds.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(17*time.Second))
	sent := <-targetConn.send
	if !bytes.Equal(sent.payload, buf) {
		t.Errorf("Mismatched payload: %v != %v", sent.payload, buf)
	}
	if sent.addr != &dnsAddr {
		t.Errorf("Mismatched address: %v != %v", sent.addr, &targetAddr)
	}
}

func TestNATWriteDNSMultiple(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate three DNS queries being sent.
	buf := []byte{1}
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	// DNS-only connections have a fixed timeout of 17 seconds.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(17*time.Second))
}

func TestNATWriteMixed(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate both non-DNS and DNS packets being sent.
	buf := []byte{1}
	entry.WriteTo(buf, &targetAddr)
	<-targetConn.send
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	// Mixed DNS and non-DNS connections should have the user-specified timeout.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(timeout))
}

func TestNATFastClose(t *testing.T) {
	clientConn, targetConn, entry := setup()

	// Send one DNS query.
	query := []byte{1}
	entry.WriteTo(query, &dnsAddr)
	sent := <-targetConn.send
	// Send the response.
	response := []byte{1, 2, 3, 4, 5}
	received := packet{addr: &dnsAddr, payload: response}
	targetConn.recv <- received
	sent, ok := <-clientConn.send
	if !ok {
		t.Error("clientConn was closed")
	}
	if len(sent.payload) <= len(response) {
		t.Error("Packet is too short to be shadowsocks-AEAD")
	}
	if sent.addr != &clientAddr {
		t.Errorf("Address mismatch: %v != %v", sent.addr, clientAddr)
	}

	// targetConn should be scheduled to close immediately.
	assertAlmostEqual(t, targetConn.deadline, time.Now())
}

func TestNATNoFastClose_NotDNS(t *testing.T) {
	clientConn, targetConn, entry := setup()

	// Send one non-DNS packet.
	query := []byte{1}
	entry.WriteTo(query, &targetAddr)
	sent := <-targetConn.send
	// Send the response.
	response := []byte{1, 2, 3, 4, 5}
	received := packet{addr: &targetAddr, payload: response}
	targetConn.recv <- received
	sent, ok := <-clientConn.send
	if !ok {
		t.Error("clientConn was closed")
	}
	if len(sent.payload) <= len(response) {
		t.Error("Packet is too short to be shadowsocks-AEAD")
	}
	if sent.addr != &clientAddr {
		t.Errorf("Address mismatch: %v != %v", sent.addr, clientAddr)
	}
	// targetConn should be scheduled to close after the full timeout.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(timeout))
}

func TestNATNoFastClose_MultipleDNS(t *testing.T) {
	clientConn, targetConn, entry := setup()

	// Send two DNS packets.
	query1 := []byte{1}
	entry.WriteTo(query1, &dnsAddr)
	<-targetConn.send
	query2 := []byte{2}
	entry.WriteTo(query2, &dnsAddr)
	<-targetConn.send

	// Send a response.
	response := []byte{1, 2, 3, 4, 5}
	received := packet{addr: &dnsAddr, payload: response}
	targetConn.recv <- received
	<-clientConn.send

	// targetConn should be scheduled to close after the DNS timeout.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(17*time.Second))
}

// Implements net.Error
type fakeTimeoutError struct {
	error
}

func (e *fakeTimeoutError) Timeout() bool {
	return true
}

func (e *fakeTimeoutError) Temporary() bool {
	return false
}

func TestNATTimeout(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate a non-DNS initial packet.
	entry.WriteTo([]byte{1}, &targetAddr)
	<-targetConn.send
	// Simulate a read timeout.
	received := packet{err: &fakeTimeoutError{}}
	before := time.Now()
	targetConn.recv <- received
	// Wait for targetConn to close.
	if _, ok := <-targetConn.send; ok {
		t.Error("targetConn should be closed due to read timeout")
	}
	// targetConn should be closed as soon as the timeout error is received.
	assertAlmostEqual(t, before, time.Now())
}

// Simulates receiving invalid UDP packets on a server with 100 ciphers.
func BenchmarkUDPUnpackFail(b *testing.B) {
	logging.SetLevel(logging.INFO, "")

	cipherList, err := MakeTestCiphers(MakeTestSecrets(100))
	if err != nil {
		b.Fatal(err)
	}
	testPayload := MakeTestPayload(50)
	textBuf := make([]byte, udpBufSize)
	testIP := net.ParseIP("192.0.2.1")
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		unpack(testIP, textBuf, testPayload, cipherList)
	}
}

// Simulates receiving valid UDP packets from 100 different users, each with
// their own cipher and IP address.
func BenchmarkUDPUnpackRepeat(b *testing.B) {
	logging.SetLevel(logging.INFO, "")

	const numCiphers = 100 // Must be <256
	cipherList, err := MakeTestCiphers(MakeTestSecrets(numCiphers))
	if err != nil {
		b.Fatal(err)
	}
	testBuf := make([]byte, udpBufSize)
	packets := [numCiphers][]byte{}
	ips := [numCiphers]net.IP{}
	_, snapshot := cipherList.SnapshotForClientIP(nil)
	for i, element := range snapshot {
		packets[i] = make([]byte, 0, udpBufSize)
		plaintext := MakeTestPayload(50)
		packets[i], err = shadowaead.Pack(make([]byte, udpBufSize), plaintext, element.Value.(*CipherEntry).Cipher)
		if err != nil {
			b.Error(err)
		}
		ips[i] = net.IPv4(192, 0, 2, byte(i))
	}
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		cipherNumber := n % numCiphers
		ip := ips[cipherNumber]
		packet := packets[cipherNumber]
		_, _, _, err := unpack(ip, testBuf, packet, cipherList)
		if err != nil {
			b.Error(err)
		}
	}
}

// Simulates receiving valid UDP packets from 100 different IP addresses,
// all using the same cipher.
func BenchmarkUDPUnpackSharedKey(b *testing.B) {
	logging.SetLevel(logging.INFO, "")

	cipherList, err := MakeTestCiphers(MakeTestSecrets(1)) // One widely shared key
	if err != nil {
		b.Fatal(err)
	}
	testBuf := make([]byte, udpBufSize)
	plaintext := MakeTestPayload(50)
	_, snapshot := cipherList.SnapshotForClientIP(nil)
	cipher := snapshot[0].Value.(*CipherEntry).Cipher
	packet, err := shadowaead.Pack(make([]byte, udpBufSize), plaintext, cipher)

	const numIPs = 100 // Must be <256
	ips := [numIPs]net.IP{}
	for i := 0; i < numIPs; i++ {
		ips[i] = net.IPv4(192, 0, 2, byte(i))
	}
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		ip := ips[n%numIPs]
		_, _, _, err := unpack(ip, testBuf, packet, cipherList)
		if err != nil {
			b.Error(err)
		}
	}
}

func TestUDPDoubleServe(t *testing.T) {
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1))
	if err != nil {
		t.Fatal(err)
	}
	testMetrics := &probeTestMetrics{}
	const testTimeout = 200 * time.Millisecond
	s := NewUDPService(testTimeout, cipherList, testMetrics)

	c := make(chan error)
	for i := 0; i < 2; i++ {
		clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
		if err != nil {
			t.Fatalf("ListenUDP failed: %v", err)
		}
		go func() {
			err := s.Serve(clientConn)
			if err != nil {
				c <- err
				close(c)
			}
		}()
	}

	err = <-c
	if err == nil {
		t.Error("Expected an error from one of the two Serve calls")
	}

	if err := s.Stop(); err != nil {
		t.Error(err)
	}
}

func TestUDPEarlyStop(t *testing.T) {
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1))
	if err != nil {
		t.Fatal(err)
	}
	testMetrics := &probeTestMetrics{}
	const testTimeout = 200 * time.Millisecond
	s := NewUDPService(testTimeout, cipherList, testMetrics)

	if err := s.Stop(); err != nil {
		t.Error(err)
	}
	clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenUDP failed: %v", err)
	}
	if err := s.Serve(clientConn); err != nil {
		t.Error(err)
	}
}
<commit_msg>Remove data race in the benchmarks due to logging<commit_after>// Copyright 2018 Jigsaw Operations LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     https://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package shadowsocks

import (
	"bytes"
	"errors"
	"net"
	"sync"
	"testing"
	"time"

	logging "github.com/op/go-logging"
	"github.com/shadowsocks/go-shadowsocks2/core"
	"github.com/shadowsocks/go-shadowsocks2/shadowaead"
)

const timeout = 5 * time.Minute

var clientAddr = net.UDPAddr{IP: []byte{192, 0, 2, 1}, Port: 12345}
var targetAddr = net.UDPAddr{IP: []byte{192, 0, 2, 2}, Port: 54321}
var dnsAddr = net.UDPAddr{IP: []byte{192, 0, 2, 3}, Port: 53}
var natCipher shadowaead.Cipher

func init() {
	logging.SetLevel(logging.INFO, "")
	coreCipher, _ := core.PickCipher(testCipher, nil, "test password")
	natCipher = coreCipher.(shadowaead.Cipher)
}

type packet struct {
	addr    net.Addr
	payload []byte
	err     error
}

type fakePacketConn struct {
	net.PacketConn
	send     chan packet
	recv     chan packet
	deadline time.Time
}

func makePacketConn() *fakePacketConn {
	return &fakePacketConn{
		send: make(chan packet, 1),
		recv: make(chan packet),
	}
}

func (conn *fakePacketConn) SetReadDeadline(deadline time.Time) error {
	conn.deadline = deadline
	return nil
}

func (conn *fakePacketConn) WriteTo(payload []byte, addr net.Addr) (int, error) {
	conn.send <- packet{addr, payload, nil}
	return len(payload), nil
}

func (conn *fakePacketConn) ReadFrom(buffer []byte) (int, net.Addr, error) {
	pkt, ok := <-conn.recv
	if !ok {
		return 0, nil, errors.New("Receive closed")
	}
	n := copy(buffer, pkt.payload)
	if n < len(pkt.payload) {
		return n, pkt.addr, errors.New("Buffer was too short")
	}
	return n, pkt.addr, pkt.err
}

func (conn *fakePacketConn) Close() error {
	close(conn.send)
	close(conn.recv)
	return nil
}

func assertAlmostEqual(t *testing.T, a, b time.Time) {
	delta := a.Sub(b)
	limit := 100 * time.Millisecond
	if delta > limit || -delta > limit {
		t.Errorf("Times are not close: %v, %v", a, b)
	}
}

func TestNATEmpty(t *testing.T) {
	nat := newNATmap(timeout, &probeTestMetrics{}, &sync.WaitGroup{})
	if nat.Get("foo") != nil {
		t.Error("Expected nil value from empty NAT map")
	}
}

func setup() (*fakePacketConn, *fakePacketConn, *natconn) {
	nat := newNATmap(timeout, &probeTestMetrics{}, &sync.WaitGroup{})
	clientConn := makePacketConn()
	targetConn := makePacketConn()
	nat.Add(&clientAddr, clientConn, natCipher, targetConn, "ZZ", "key id")
	entry := nat.Get(clientAddr.String())
	return clientConn, targetConn, entry
}

func TestNATGet(t *testing.T) {
	_, targetConn, entry := setup()
	if entry == nil {
		t.Fatal("Failed to find target conn")
	}
	if entry.PacketConn != targetConn {
		t.Error("Mismatched connection returned")
	}
}

func TestNATWrite(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate one generic packet being sent
	buf := []byte{1}
	entry.WriteTo([]byte{1}, &targetAddr)
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(timeout))
	sent := <-targetConn.send
	if !bytes.Equal(sent.payload, buf) {
		t.Errorf("Mismatched payload: %v != %v", sent.payload, buf)
	}
	if sent.addr != &targetAddr {
		t.Errorf("Mismatched address: %v != %v", sent.addr, &targetAddr)
	}
}

func TestNATWriteDNS(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate one DNS query being sent.
	buf := []byte{1}
	entry.WriteTo(buf, &dnsAddr)
	// DNS-only connections have a fixed timeout of 17 seconds.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(17*time.Second))
	sent := <-targetConn.send
	if !bytes.Equal(sent.payload, buf) {
		t.Errorf("Mismatched payload: %v != %v", sent.payload, buf)
	}
	if sent.addr != &dnsAddr {
		t.Errorf("Mismatched address: %v != %v", sent.addr, &targetAddr)
	}
}

func TestNATWriteDNSMultiple(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate three DNS queries being sent.
	buf := []byte{1}
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	// DNS-only connections have a fixed timeout of 17 seconds.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(17*time.Second))
}

func TestNATWriteMixed(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate both non-DNS and DNS packets being sent.
	buf := []byte{1}
	entry.WriteTo(buf, &targetAddr)
	<-targetConn.send
	entry.WriteTo(buf, &dnsAddr)
	<-targetConn.send
	// Mixed DNS and non-DNS connections should have the user-specified timeout.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(timeout))
}

func TestNATFastClose(t *testing.T) {
	clientConn, targetConn, entry := setup()

	// Send one DNS query.
	query := []byte{1}
	entry.WriteTo(query, &dnsAddr)
	sent := <-targetConn.send
	// Send the response.
	response := []byte{1, 2, 3, 4, 5}
	received := packet{addr: &dnsAddr, payload: response}
	targetConn.recv <- received
	sent, ok := <-clientConn.send
	if !ok {
		t.Error("clientConn was closed")
	}
	if len(sent.payload) <= len(response) {
		t.Error("Packet is too short to be shadowsocks-AEAD")
	}
	if sent.addr != &clientAddr {
		t.Errorf("Address mismatch: %v != %v", sent.addr, clientAddr)
	}

	// targetConn should be scheduled to close immediately.
	assertAlmostEqual(t, targetConn.deadline, time.Now())
}

func TestNATNoFastClose_NotDNS(t *testing.T) {
	clientConn, targetConn, entry := setup()

	// Send one non-DNS packet.
	query := []byte{1}
	entry.WriteTo(query, &targetAddr)
	sent := <-targetConn.send
	// Send the response.
	response := []byte{1, 2, 3, 4, 5}
	received := packet{addr: &targetAddr, payload: response}
	targetConn.recv <- received
	sent, ok := <-clientConn.send
	if !ok {
		t.Error("clientConn was closed")
	}
	if len(sent.payload) <= len(response) {
		t.Error("Packet is too short to be shadowsocks-AEAD")
	}
	if sent.addr != &clientAddr {
		t.Errorf("Address mismatch: %v != %v", sent.addr, clientAddr)
	}
	// targetConn should be scheduled to close after the full timeout.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(timeout))
}

func TestNATNoFastClose_MultipleDNS(t *testing.T) {
	clientConn, targetConn, entry := setup()

	// Send two DNS packets.
	query1 := []byte{1}
	entry.WriteTo(query1, &dnsAddr)
	<-targetConn.send
	query2 := []byte{2}
	entry.WriteTo(query2, &dnsAddr)
	<-targetConn.send

	// Send a response.
	response := []byte{1, 2, 3, 4, 5}
	received := packet{addr: &dnsAddr, payload: response}
	targetConn.recv <- received
	<-clientConn.send

	// targetConn should be scheduled to close after the DNS timeout.
	assertAlmostEqual(t, targetConn.deadline, time.Now().Add(17*time.Second))
}

// Implements net.Error
type fakeTimeoutError struct {
	error
}

func (e *fakeTimeoutError) Timeout() bool {
	return true
}

func (e *fakeTimeoutError) Temporary() bool {
	return false
}

func TestNATTimeout(t *testing.T) {
	_, targetConn, entry := setup()

	// Simulate a non-DNS initial packet.
	entry.WriteTo([]byte{1}, &targetAddr)
	<-targetConn.send
	// Simulate a read timeout.
	received := packet{err: &fakeTimeoutError{}}
	before := time.Now()
	targetConn.recv <- received
	// Wait for targetConn to close.
	if _, ok := <-targetConn.send; ok {
		t.Error("targetConn should be closed due to read timeout")
	}
	// targetConn should be closed as soon as the timeout error is received.
	assertAlmostEqual(t, before, time.Now())
}

// Simulates receiving invalid UDP packets on a server with 100 ciphers.
func BenchmarkUDPUnpackFail(b *testing.B) {
	cipherList, err := MakeTestCiphers(MakeTestSecrets(100))
	if err != nil {
		b.Fatal(err)
	}
	testPayload := MakeTestPayload(50)
	textBuf := make([]byte, udpBufSize)
	testIP := net.ParseIP("192.0.2.1")
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		unpack(testIP, textBuf, testPayload, cipherList)
	}
}

// Simulates receiving valid UDP packets from 100 different users, each with
// their own cipher and IP address.
func BenchmarkUDPUnpackRepeat(b *testing.B) {
	const numCiphers = 100 // Must be <256
	cipherList, err := MakeTestCiphers(MakeTestSecrets(numCiphers))
	if err != nil {
		b.Fatal(err)
	}
	testBuf := make([]byte, udpBufSize)
	packets := [numCiphers][]byte{}
	ips := [numCiphers]net.IP{}
	_, snapshot := cipherList.SnapshotForClientIP(nil)
	for i, element := range snapshot {
		packets[i] = make([]byte, 0, udpBufSize)
		plaintext := MakeTestPayload(50)
		packets[i], err = shadowaead.Pack(make([]byte, udpBufSize), plaintext, element.Value.(*CipherEntry).Cipher)
		if err != nil {
			b.Error(err)
		}
		ips[i] = net.IPv4(192, 0, 2, byte(i))
	}
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		cipherNumber := n % numCiphers
		ip := ips[cipherNumber]
		packet := packets[cipherNumber]
		_, _, _, err := unpack(ip, testBuf, packet, cipherList)
		if err != nil {
			b.Error(err)
		}
	}
}

// Simulates receiving valid UDP packets from 100 different IP addresses,
// all using the same cipher.
func BenchmarkUDPUnpackSharedKey(b *testing.B) {
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1)) // One widely shared key
	if err != nil {
		b.Fatal(err)
	}
	testBuf := make([]byte, udpBufSize)
	plaintext := MakeTestPayload(50)
	_, snapshot := cipherList.SnapshotForClientIP(nil)
	cipher := snapshot[0].Value.(*CipherEntry).Cipher
	packet, err := shadowaead.Pack(make([]byte, udpBufSize), plaintext, cipher)

	const numIPs = 100 // Must be <256
	ips := [numIPs]net.IP{}
	for i := 0; i < numIPs; i++ {
		ips[i] = net.IPv4(192, 0, 2, byte(i))
	}
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		ip := ips[n%numIPs]
		_, _, _, err := unpack(ip, testBuf, packet, cipherList)
		if err != nil {
			b.Error(err)
		}
	}
}

func TestUDPDoubleServe(t *testing.T) {
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1))
	if err != nil {
		t.Fatal(err)
	}
	testMetrics := &probeTestMetrics{}
	const testTimeout = 200 * time.Millisecond
	s := NewUDPService(testTimeout, cipherList, testMetrics)

	c := make(chan error)
	for i := 0; i < 2; i++ {
		clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
		if err != nil {
			t.Fatalf("ListenUDP failed: %v", err)
		}
		go func() {
			err := s.Serve(clientConn)
			if err != nil {
				c <- err
				close(c)
			}
		}()
	}

	err = <-c
	if err == nil {
		t.Error("Expected an error from one of the two Serve calls")
	}

	if err := s.Stop(); err != nil {
		t.Error(err)
	}
}

func TestUDPEarlyStop(t *testing.T) {
	cipherList, err := MakeTestCiphers(MakeTestSecrets(1))
	if err != nil {
		t.Fatal(err)
	}
	testMetrics := &probeTestMetrics{}
	const testTimeout = 200 * time.Millisecond
	s := NewUDPService(testTimeout, cipherList, testMetrics)

	if err := s.Stop(); err != nil {
		t.Error(err)
	}
	clientConn, err := net.ListenUDP("udp", &net.UDPAddr{IP: net.ParseIP("127.0.0.1"), Port: 0})
	if err != nil {
		t.Fatalf("ListenUDP failed: %v", err)
	}
	if err := s.Serve(clientConn); err != nil {
		t.Error(err)
	}
}
<|endoftext|>


<commit_before>package mqb

import (
	"errors"
	"fmt"
	"math"
	"net/http"
	"reflect"
	"strconv"
	"strings"

	"github.com/ansel1/merry"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	DefaultPageSize uint = 20 // DefaultPageSize defines how many elements a page contains per default.
)

// Page the paging information.
type Page struct {
	Size    uint `json:"size"`    // Size defines how many elements a page contains.
	Items   uint `json:"items"`   // Items defines the total number of items the corresponding query returns.
	Last    uint `json:"last"`    // Last represents total number of pages a query generates (depends on the page size and the total number of elements returned by the query).
	Current uint `json:"current"` // Current is the current page nuber for the query.
}

// Response contains the result of the query, including the Page information.
type Response struct {
	Content interface{} `json:"content,omitempty"`
	Page    Page        `json:"page"`
}

// MongoQuery can be used to to create mgo.Query from http request parameters.
type MongoQuery struct {
	endPointStruct               interface{}
	dataBase                     *mgo.Database
	supportedParameters          map[string]reflect.Kind
	additionalSupportedParamters map[string]reflect.Kind
	disabledParameters           []string
	page                         Page
}

// NewMongoQuery returns a new MongoQuery.
func NewMongoQuery(endPointStruct interface{}, database *mgo.Database) *MongoQuery {
	return &MongoQuery{
		dataBase:                     database,
		supportedParameters:          createValidParametersMap(endPointStruct),
		disabledParameters:           []string{},
		additionalSupportedParamters: make(map[string]reflect.Kind),
		endPointStruct:               endPointStruct,
	}
}

// CreateQuery creates a mgo.Query from a HTTP Request for a collection represented by endpointStruct.
//
// Examples:
//     mq := NewMongoQuery(People{}, db)
//     q, _ := mq.CreateQuery(req) // creates a query from the request for the people collection
//
//     mq.DisableParameters("name", "sort")
//     q, _ := mq.CreateQuery(req) // creates a query from the request for the people collection with the parameters "name" and "sort" disabled.
//
func (mq *MongoQuery) CreateQuery(req *http.Request) (*mgo.Query, error) {
	filterMap, err := mq.createQueryFilter(req)
	if err != nil {
		return nil, err
	}
	q := mq.dataBase.C(structName(mq.endPointStruct)).Find(filterMap)

	selectFields, err := mq.createFieldsMap(req)
	q.Select(selectFields)

	sortFields, err := mq.createSortFields(req)
	if err != nil {
		return nil, err
	}
	q.Sort(sortFields...)

	mq.page.Size, err = getUint(req, "limit", DefaultPageSize)
	if err != nil {
		return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
	}
	mq.page.Current, err = getUint(req, "page", 1)
	if err != nil {
		return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
	}
	if mq.page.Current == 0 {
		return nil, merry.Wrap(errors.New("page cannot be 0")).WithHTTPCode(http.StatusBadRequest)
	}
	q = q.Limit(int(mq.page.Size))
	q = q.Skip(int((mq.page.Current - 1) * mq.page.Size))
	return q, nil
}

// Run runs the query on the database and returns a *Response.
func (mq *MongoQuery) Run(req *http.Request) (*Response, error) {
	q, err := mq.CreateQuery(req)
	if err != nil {
		return nil, err
	}

	// copy query and reset limit and skip values to count total items
	// that would be returned for a query
	countQuery := &mgo.Query{}
	*countQuery = *q
	countQuery.Limit(0)
	countQuery.Skip(0)
	items, err := countQuery.Count()
	if err != nil {
		return nil, merry.Wrap(err).WithHTTPCode(http.StatusInternalServerError)
	}

	response := &Response{
		Page: mq.page,
	}
	response.Page.Items = uint(items)
	response.Page.calculateLastPage()

	// create a pointer to an empty slice with same type as enpointStruct to store the
	// result of the query
	slice := reflect.MakeSlice(reflect.SliceOf(reflect.TypeOf(mq.endPointStruct)), 0, 0)
	content := reflect.New(slice.Type()).Interface()
	err = q.All(content)
	if err != nil {
		return nil, merry.Wrap(err).WithHTTPCode(http.StatusInternalServerError)
	}
	response.Content = content
	return response, nil
}

// DisableParameters disables paramters. If a URL query contains any
// of those paramters, an error is returned.
func (mq *MongoQuery) DisableParameters(paramters ...string) {
	for _, p := range paramters {
		if !contains(mq.disabledParameters, p) {
			mq.disabledParameters = append(mq.disabledParameters, p)
		}
	}
	mq.supportedParameters = createValidParametersMap(mq.endPointStruct, mq.disabledParameters...)
	for k, v := range mq.additionalSupportedParamters {
		mq.supportedParameters[k] = v
	}
}

// AddOrOverwriteValidParameter adds or overwrites a valid parmeter with name and reflect.Kind.
func (mq *MongoQuery) AddOrOverwriteValidParameter(name string, value reflect.Kind) {
	mq.additionalSupportedParamters[name] = value
	for k, v := range mq.additionalSupportedParamters {
		mq.supportedParameters[k] = v
	}
}

func (mq *MongoQuery) createQueryFilter(req *http.Request) (map[string]interface{}, error) {
	filter := make(map[string]interface{})

	for parameterName, parameterValues := range req.URL.Query() {
		s := []interface{}{}
		if kind, ok := mq.supportedParameters[parameterName]; ok {
			// meta parameters are not filters
			if _, ok := validMetaParameters[parameterName]; ok {
				continue
			}
			switch kind {
			case reflect.Bool:
				for _, v := range parameterValues {
					b, err := strconv.ParseBool(v)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, b)
				}
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				for _, v := range parameterValues {
					i, err := strconv.Atoi(v)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, i)
				}
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				for _, v := range parameterValues {
					i, err := strconv.ParseUint(v, 10, 0)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, uint(i))
				}
			case reflect.Float32, reflect.Float64:
				for _, v := range parameterValues {
					f, err := strconv.ParseFloat(v, 64)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, f)
				}
			case reflect.String:
				if len(parameterValues) == 1 {
					if bson.IsObjectIdHex(parameterValues[0]) {
						s = []interface{}{bson.ObjectIdHex(parameterValues[0])}
					} else {
						s = []interface{}{bson.RegEx{Pattern: parameterValues[0], Options: ""}}
					}
				} else {
					for _, v := range parameterValues {
						if bson.IsObjectIdHex(v) {
							s = append(s, bson.ObjectIdHex(v))
						} else {
							s = append(s, v)
						}
					}
				}
			default:
				return nil, merry.Wrap(fmt.Errorf("reflection kind '%s' is not supported", kind)).WithHTTPCode(http.StatusBadRequest)
			}
		} else {
			return nil, merry.Wrap(fmt.Errorf("parameter '%s' is not supported", parameterName)).WithHTTPCode(http.StatusBadRequest)
		}
		if len(s) == 1 {
			filter[parameterName] = s[0]
		} else {
			filter[parameterName] = map[string]interface{}{
				"$in": s,
			}
		}
	}
	return filter, nil
}

func (mq *MongoQuery) createFieldsMap(req *http.Request) (map[string]interface{}, error) {
	fields := make(map[string]interface{})
	if _field, ok := req.URL.Query()["field"]; ok {
		for _, v := range _field {
			if _, ok := mq.supportedParameters[v]; !ok {
				return nil, fmt.Errorf("unsupported field value: %s", v)
			}
			fields[v] = 1
		}
	}
	return fields, nil
}

func (mq *MongoQuery) createSortFields(req *http.Request) ([]string, error) {
	sortFields := []string{}
	if _sortField, ok := req.URL.Query()["sort"]; ok {
		for _, v := range _sortField {
			if _, ok := mq.supportedParameters[strings.Trim(v, "-")]; !ok {
				return nil, merry.Wrap(fmt.Errorf("unsupported field value: %s", v)).WithHTTPCode(http.StatusBadRequest)
			}
			sortFields = append(sortFields, v)
		}
	}
	return sortFields, nil
}

func (p *Page) calculateLastPage() {
	p.Last = uint(math.Ceil(float64(p.Items) / float64(p.Size)))
}
<commit_msg>add more detail for internal server errors<commit_after>package mqb

import (
	"errors"
	"fmt"
	"math"
	"net/http"
	"reflect"
	"strconv"
	"strings"

	"github.com/ansel1/merry"
	"gopkg.in/mgo.v2"
	"gopkg.in/mgo.v2/bson"
)

var (
	DefaultPageSize uint = 20 // DefaultPageSize defines how many elements a page contains per default.
)

// Page the paging information.
type Page struct {
	Size    uint `json:"size"`    // Size defines how many elements a page contains.
	Items   uint `json:"items"`   // Items defines the total number of items the corresponding query returns.
	Last    uint `json:"last"`    // Last represents total number of pages a query generates (depends on the page size and the total number of elements returned by the query).
	Current uint `json:"current"` // Current is the current page nuber for the query.
}

// Response contains the result of the query, including the Page information.
type Response struct {
	Content interface{} `json:"content,omitempty"`
	Page    Page        `json:"page"`
}

// MongoQuery can be used to to create mgo.Query from http request parameters.
type MongoQuery struct {
	endPointStruct               interface{}
	dataBase                     *mgo.Database
	supportedParameters          map[string]reflect.Kind
	additionalSupportedParamters map[string]reflect.Kind
	disabledParameters           []string
	page                         Page
}

// NewMongoQuery returns a new MongoQuery.
func NewMongoQuery(endPointStruct interface{}, database *mgo.Database) *MongoQuery {
	return &MongoQuery{
		dataBase:                     database,
		supportedParameters:          createValidParametersMap(endPointStruct),
		disabledParameters:           []string{},
		additionalSupportedParamters: make(map[string]reflect.Kind),
		endPointStruct:               endPointStruct,
	}
}

// CreateQuery creates a mgo.Query from a HTTP Request for a collection represented by endpointStruct.
//
// Examples:
//     mq := NewMongoQuery(People{}, db)
//     q, _ := mq.CreateQuery(req) // creates a query from the request for the people collection
//
//     mq.DisableParameters("name", "sort")
//     q, _ := mq.CreateQuery(req) // creates a query from the request for the people collection with the parameters "name" and "sort" disabled.
//
func (mq *MongoQuery) CreateQuery(req *http.Request) (*mgo.Query, error) {
	filterMap, err := mq.createQueryFilter(req)
	if err != nil {
		return nil, err
	}
	q := mq.dataBase.C(structName(mq.endPointStruct)).Find(filterMap)

	selectFields, err := mq.createFieldsMap(req)
	q.Select(selectFields)

	sortFields, err := mq.createSortFields(req)
	if err != nil {
		return nil, err
	}
	q.Sort(sortFields...)

	mq.page.Size, err = getUint(req, "limit", DefaultPageSize)
	if err != nil {
		return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
	}
	mq.page.Current, err = getUint(req, "page", 1)
	if err != nil {
		return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
	}
	if mq.page.Current == 0 {
		return nil, merry.Wrap(errors.New("page cannot be 0")).WithHTTPCode(http.StatusBadRequest)
	}
	q = q.Limit(int(mq.page.Size))
	q = q.Skip(int((mq.page.Current - 1) * mq.page.Size))
	return q, nil
}

// Run runs the query on the database and returns a *Response.
func (mq *MongoQuery) Run(req *http.Request) (*Response, error) {
	q, err := mq.CreateQuery(req)
	if err != nil {
		return nil, err
	}

	// copy query and reset limit and skip values to count total items
	// that would be returned for a query
	countQuery := &mgo.Query{}
	*countQuery = *q
	countQuery.Limit(0)
	countQuery.Skip(0)
	items, err := countQuery.Count()
	if err != nil {
		return nil, merry.New("could not create count query").Append(err.Error()).WithHTTPCode(http.StatusInternalServerError)
	}

	response := &Response{
		Page: mq.page,
	}
	response.Page.Items = uint(items)
	response.Page.calculateLastPage()

	// create a pointer to an empty slice with same type as enpointStruct to store the
	// result of the query
	slice := reflect.MakeSlice(reflect.SliceOf(reflect.TypeOf(mq.endPointStruct)), 0, 0)
	content := reflect.New(slice.Type()).Interface()
	err = q.All(content)
	if err != nil {
		return nil, merry.New("could not execute q.All()").Append(err.Error()).WithHTTPCode(http.StatusInternalServerError)
	}
	response.Content = content
	return response, nil
}

// DisableParameters disables paramters. If a URL query contains any
// of those paramters, an error is returned.
func (mq *MongoQuery) DisableParameters(paramters ...string) {
	for _, p := range paramters {
		if !contains(mq.disabledParameters, p) {
			mq.disabledParameters = append(mq.disabledParameters, p)
		}
	}
	mq.supportedParameters = createValidParametersMap(mq.endPointStruct, mq.disabledParameters...)
	for k, v := range mq.additionalSupportedParamters {
		mq.supportedParameters[k] = v
	}
}

// AddOrOverwriteValidParameter adds or overwrites a valid parmeter with name and reflect.Kind.
func (mq *MongoQuery) AddOrOverwriteValidParameter(name string, value reflect.Kind) {
	mq.additionalSupportedParamters[name] = value
	for k, v := range mq.additionalSupportedParamters {
		mq.supportedParameters[k] = v
	}
}

func (mq *MongoQuery) createQueryFilter(req *http.Request) (map[string]interface{}, error) {
	filter := make(map[string]interface{})

	for parameterName, parameterValues := range req.URL.Query() {
		s := []interface{}{}
		if kind, ok := mq.supportedParameters[parameterName]; ok {
			// meta parameters are not filters
			if _, ok := validMetaParameters[parameterName]; ok {
				continue
			}
			switch kind {
			case reflect.Bool:
				for _, v := range parameterValues {
					b, err := strconv.ParseBool(v)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, b)
				}
			case reflect.Int, reflect.Int8, reflect.Int16, reflect.Int32, reflect.Int64:
				for _, v := range parameterValues {
					i, err := strconv.Atoi(v)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, i)
				}
			case reflect.Uint, reflect.Uint8, reflect.Uint16, reflect.Uint32, reflect.Uint64:
				for _, v := range parameterValues {
					i, err := strconv.ParseUint(v, 10, 0)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, uint(i))
				}
			case reflect.Float32, reflect.Float64:
				for _, v := range parameterValues {
					f, err := strconv.ParseFloat(v, 64)
					if err != nil {
						return nil, merry.Wrap(err).WithHTTPCode(http.StatusBadRequest)
					}
					s = append(s, f)
				}
			case reflect.String:
				if len(parameterValues) == 1 {
					if bson.IsObjectIdHex(parameterValues[0]) {
						s = []interface{}{bson.ObjectIdHex(parameterValues[0])}
					} else {
						s = []interface{}{bson.RegEx{Pattern: parameterValues[0], Options: ""}}
					}
				} else {
					for _, v := range parameterValues {
						if bson.IsObjectIdHex(v) {
							s = append(s, bson.ObjectIdHex(v))
						} else {
							s = append(s, v)
						}
					}
				}
			default:
				return nil, merry.Wrap(fmt.Errorf("reflection kind '%s' is not supported", kind)).WithHTTPCode(http.StatusBadRequest)
			}
		} else {
			return nil, merry.Wrap(fmt.Errorf("parameter '%s' is not supported", parameterName)).WithHTTPCode(http.StatusBadRequest)
		}
		if len(s) == 1 {
			filter[parameterName] = s[0]
		} else {
			filter[parameterName] = map[string]interface{}{
				"$in": s,
			}
		}
	}
	return filter, nil
}

func (mq *MongoQuery) createFieldsMap(req *http.Request) (map[string]interface{}, error) {
	fields := make(map[string]interface{})
	if _field, ok := req.URL.Query()["field"]; ok {
		for _, v := range _field {
			if _, ok := mq.supportedParameters[v]; !ok {
				return nil, fmt.Errorf("unsupported field value: %s", v)
			}
			fields[v] = 1
		}
	}
	return fields, nil
}

func (mq *MongoQuery) createSortFields(req *http.Request) ([]string, error) {
	sortFields := []string{}
	if _sortField, ok := req.URL.Query()["sort"]; ok {
		for _, v := range _sortField {
			if _, ok := mq.supportedParameters[strings.Trim(v, "-")]; !ok {
				return nil, merry.Wrap(fmt.Errorf("unsupported field value: %s", v)).WithHTTPCode(http.StatusBadRequest)
			}
			sortFields = append(sortFields, v)
		}
	}
	return sortFields, nil
}

func (p *Page) calculateLastPage() {
	p.Last = uint(math.Ceil(float64(p.Items) / float64(p.Size)))
}
<|endoftext|>


<commit_before>package source

import (
	"errors"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path"
	"time"
)

var (
	ErrFileType = errors.New("not a Go file")
)

type File struct {
	Options []*Option // Options found in file.
	Name    string    // Name of command.
	Doc     string    // Comment at top of file.
	Time    time.Time // Modification time.

	// Unexported fields
	fileSet  *token.FileSet
	file     *ast.File
	comments map[int]*ast.CommentGroup // map ending line -> comment group
}

func NewFile(filePath string) (*File, error) {
	result := new(File)

	// Get modification time and base name of source file.
	fileInfo, err := os.Stat(filePath)
	if err != nil {
		return nil, err
	}
	result.Time = fileInfo.ModTime()
	result.Name = fileInfo.Name()

	// Check if file has a .go extension.
	if ext := path.Ext(result.Name); ext != ".go" {
		return nil, ErrFileType
	} else {
		result.Name = result.Name[:len(result.Name)-3] // remove ".go"
	}

	fileSet := token.NewFileSet()
	file, err := parser.ParseFile(fileSet, filePath, nil, parser.ParseComments)
	if err != nil {
		return nil, err
	}
	result.fileSet = fileSet
	result.file = file

	// Load comment groups and map them to their ending line number.
	// We assume a comment belongs to a command line flag declaration if it
	// ends on the previous line of the flag declaration.
	result.comments = make(map[int]*ast.CommentGroup)
	for _, group := range file.Comments {
		pos := fileSet.Position(group.Pos())
		end := fileSet.Position(group.End())

		if pos.Line == 1 {
			result.Doc = group.Text()
		}
		result.comments[end.Line] = group
	}

	result.parseOptions()

	return result, nil
}

func (f *File) parseOptions() {
	options := make([]*Option, 0)

	// Memorize options with variable names
	bound := make(map[string]*Option)

	// Load all options in source file. This means, detect and parse
	// all flag.Bool, flag.Duration, ... calls.
	ast.Inspect(f.file, func(node ast.Node) bool {
		if call, ok := node.(*ast.CallExpr); ok {
			if opt, err := NewOptionFromCallExpr(f.fileSet, call); err == nil {
				// Check if we have a comment that belongs to option
				if comment, ok := f.comments[opt.Line-1]; ok {
					opt.Doc = comment.Text()
				}

				// Check if we already encountered an option bound to the
				// variable.
				if opt.Variable != "" {
					if reg, ok := bound[opt.Variable]; ok {
						// Merge currrent option with the one we already found
						reg.merge(opt)
						// Don't add the current option to the list, since the list
						// already contains the struct stored in the map.
						return true
					} else {
						// Register variable and the proceed to add option
						// struct to the options list
						bound[opt.Variable] = opt
					}
				}
				options = append(options, opt)
			}
		}
		return true
	})

	f.Options = options
}
<commit_msg>Close issue #5 submitted by mewmew<commit_after>package source

import (
	"errors"
	"go/ast"
	"go/parser"
	"go/token"
	"os"
	"path"
	"time"
)

var (
	ErrFileType = errors.New("not a Go file")
)

type File struct {
	Options []*Option // Options found in file.
	Name    string    // Name of command.
	Doc     string    // Comment preceding the "package" keyword.
	Time    time.Time // Modification time.

	// Unexported fields
	fileSet  *token.FileSet
	file     *ast.File
	comments map[int]*ast.CommentGroup // map ending line -> comment group
}

func NewFile(filePath string) (*File, error) {
	result := new(File)

	// Get modification time and base name of source file.
	fileInfo, err := os.Stat(filePath)
	if err != nil {
		return nil, err
	}
	result.Time = fileInfo.ModTime()
	result.Name = fileInfo.Name()

	// Check if file has a .go extension.
	if ext := path.Ext(result.Name); ext != ".go" {
		return nil, ErrFileType
	} else {
		result.Name = result.Name[:len(result.Name)-3] // remove ".go"
	}

	fileSet := token.NewFileSet()
	file, err := parser.ParseFile(fileSet, filePath, nil, parser.ParseComments)
	if err != nil {
		return nil, err
	}
	result.fileSet = fileSet
	result.file = file

	// The last comment group before a package declaration must contain the
	// command description.
	packageLine := 2
	if packagePos := fileSet.Position(file.Package); packagePos.IsValid() {
		packageLine = packagePos.Line
	}

	// Load comment groups and map them to their ending line number.
	// We assume a comment belongs to a command line flag declaration if it
	// ends on the previous line of the flag declaration.
	result.comments = make(map[int]*ast.CommentGroup)
	for _, group := range file.Comments {
		pos := fileSet.Position(group.Pos())
		end := fileSet.Position(group.End())

		if pos.Line < packageLine {
			result.Doc = group.Text()
		}
		result.comments[end.Line] = group
	}

	result.parseOptions()

	return result, nil
}

func (f *File) parseOptions() {
	options := make([]*Option, 0)

	// Memorize options with variable names
	bound := make(map[string]*Option)

	// Load all options in source file. This means, detect and parse
	// all flag.Bool, flag.Duration, ... calls.
	ast.Inspect(f.file, func(node ast.Node) bool {
		if call, ok := node.(*ast.CallExpr); ok {
			if opt, err := NewOptionFromCallExpr(f.fileSet, call); err == nil {
				// Check if we have a comment that belongs to option
				if comment, ok := f.comments[opt.Line-1]; ok {
					opt.Doc = comment.Text()
				}

				// Check if we already encountered an option bound to the
				// variable.
				if opt.Variable != "" {
					if reg, ok := bound[opt.Variable]; ok {
						// Merge currrent option with the one we already found
						reg.merge(opt)
						// Don't add the current option to the list, since the list
						// already contains the struct stored in the map.
						return true
					} else {
						// Register variable and the proceed to add option
						// struct to the options list
						bound[opt.Variable] = opt
					}
				}
				options = append(options, opt)
			}
		}
		return true
	})

	f.Options = options
}
<|endoftext|>


<commit_before>package main

import (
	. "fmt"
	"github.com/go-vgo/robotgo"
	// "go-vgo/robotgo"
)

func main() {
	//Control the keyboard
	robotgo.TypeString("Hello World") //importing "Hello World"
	robotgo.KeyTap("enter")           //Press "enter"
	robotgo.KeyTap("a", "control")
	robotgo.KeyTap("h", "command") //Hide window

	robotgo.KeyTap("i", "alt", "command")
	//Press "i", "alt", "command" Key combination
	arr := []string{"alt", "command"}
	robotgo.KeyTap("i", arr)

	robotgo.KeyTap("w", "command") //close window
	robotgo.KeyTap("m", "command") //minimize window
	robotgo.KeyTap("f1", "control")
	robotgo.KeyTap("a", "control")
	robotgo.KeyToggle("a", "down")
	robotgo.KeyToggle("a", "down", "alt")
	robotgo.KeyToggle("a", "down", "alt", "command")
	robotgo.KeyToggle("enter", "down")
	robotgo.TypeString("en")

	//Control the mouse
	robotgo.MoveMouse(100, 200)          // Move the mouse to 100, 200
	robotgo.MouseClick()                 //Click the left mouse button
	robotgo.MouseClick("right", false)   //Click the right mouse button
	robotgo.MouseClick("left", true)     //double click the left mouse button
	robotgo.ScrollMouse(10, "up")        //Scrolls the mouse either up
	robotgo.MouseToggle("down", "right") //Toggles right mouse button
	robotgo.MoveMouseSmooth(100, 200)    //Smooth move the mouse to 100, 200
	robotgo.MoveMouseSmooth(100, 200, 1.0, 100.0)
	x, y := robotgo.GetMousePos() //Gets the mouse coordinates
	Println("pos:", x, y)
	if x == 456 && y == 586 {
		Println("mouse...", "586")
	}

	robotgo.MouseToggle("up")
	robotgo.MoveMouse(x, y)
	robotgo.MoveMouse(100, 200)

	for i := 0; i < 1080; i += 1000 {
		Println(i)
		robotgo.MoveMouse(800, i)
	}

	//read the screen
	// robotgo.CaptureScreen()
	// bit_map := robotgo.CaptureScreen()
	// Println("CaptureScreen...", bit_map)
	gbit_map := robotgo.Capture_Screen()
	Println("Capture_Screen...", gbit_map.Width)

	sx, sy := robotgo.GetScreenSize()
	//Gets the screen width and height
	Println("...", sx, sy)

	color := robotgo.GetPixelColor(100, 200)
	//Gets the pixel color at 100, 200.
	Println("color----", color, "-----------------")

	color2 := robotgo.GetPixelColor(10, 20)
	//Gets the pixel color at 10, 20.
	Println("color---", color2)

	// Bitmap
	abit_map := robotgo.CaptureScreen()
	//Gets all of the screen
	Println("a...", abit_map)

	bit_map := robotgo.CaptureScreen(100, 200, 30, 40)
	//Gets part of the screen
	Println("CaptureScreen...", bit_map)
	// Println("...", bit_map.Width, bit_map.BytesPerPixel)

	fx, fy := robotgo.FindBitmap(bit_map)
	//Searches for needle in bitmap
	Println("FindBitmap------", fx, fy)

	bit_pos := robotgo.GetPortion(bit_map, 10, 10, 11, 10)
	//Returns new bitmap object created from a portion of another
	Println(bit_pos)

	bit_str := robotgo.TostringBitmap(bit_map)
	//Creates bitmap from string by bit_map
	Println("bit_str...", bit_str)

	// sbit_map := robotgo.BitmapFromstring(bit_str, 2)
	// Println("...", sbit_map)

	robotgo.SaveBitmap(bit_map, "test.png")
	//Saves image to absolute filepath in the given format
	robotgo.SaveBitmap(bit_map, "test31.tif", 1)
	robotgo.Convert("test.png", "test.tif")
	//Convert image

	open_bit := robotgo.OpenBitmap("test.tif")
	// open image bitmap
	Println("open...", open_bit)

	//global event listener
	Println("---please press v---")
	eve := robotgo.AddEvent("v")

	if eve == 0 {
		Println("---you press v---", "v")
	}

	Println("---please press k---")
	keve := robotgo.AddEvent("k")
	if keve == 0 {
		Println("---you press k---", "k")
	}

	Println("---please press left mouse button---")
	mleft := robotgo.AddEvent("mleft")
	if mleft == 0 {
		Println("---you press left mouse button---", "mleft")
	}

	// mright := robotgo.AddEvent("mright")
	// if mright == 0 {
	// 	Println("---you press right mouse button---", "mright")
	// }

	// robotgo.LStop()

	//Window Handle
	abool := robotgo.ShowAlert("hello", "robotgo") //Show Alert Window
	if abool == 0 {
		Println("ok@@@", "ok")
	}
	robotgo.ShowAlert("hello", "robotgo", "Ok", "Cancel")
	// robotgo.GetPID()
	mdata := robotgo.GetActive() //Get current Window Active
	hwnd := robotgo.GetHandle()  //Get current Window Handle
	Println("hwnd---", hwnd)
	title := robotgo.GetTitle() //Get current Window title
	Println("title-----", title)
	robotgo.CloseWindow()    //close current Window
	robotgo.SetActive(mdata) //set Window Active
}
<commit_msg>Update golint<commit_after>package main

import (
	. "fmt"

	"github.com/go-vgo/robotgo"
	// "go-vgo/robotgo"
)

func main() {
	//Control the keyboard
	robotgo.TypeString("Hello World") //importing "Hello World"
	robotgo.KeyTap("enter")           //Press "enter"
	robotgo.KeyTap("a", "control")
	robotgo.KeyTap("h", "command") //Hide window

	robotgo.KeyTap("i", "alt", "command")
	//Press "i", "alt", "command" Key combination
	arr := []string{"alt", "command"}
	robotgo.KeyTap("i", arr)

	robotgo.KeyTap("w", "command") //close window
	robotgo.KeyTap("m", "command") //minimize window
	robotgo.KeyTap("f1", "control")
	robotgo.KeyTap("a", "control")
	robotgo.KeyToggle("a", "down")
	robotgo.KeyToggle("a", "down", "alt")
	robotgo.KeyToggle("a", "down", "alt", "command")
	robotgo.KeyToggle("enter", "down")
	robotgo.TypeString("en")

	//Control the mouse
	robotgo.MoveMouse(100, 200)          // Move the mouse to 100, 200
	robotgo.MouseClick()                 //Click the left mouse button
	robotgo.MouseClick("right", false)   //Click the right mouse button
	robotgo.MouseClick("left", true)     //double click the left mouse button
	robotgo.ScrollMouse(10, "up")        //Scrolls the mouse either up
	robotgo.MouseToggle("down", "right") //Toggles right mouse button
	robotgo.MoveMouseSmooth(100, 200)    //Smooth move the mouse to 100, 200
	robotgo.MoveMouseSmooth(100, 200, 1.0, 100.0)
	x, y := robotgo.GetMousePos() //Gets the mouse coordinates
	Println("pos:", x, y)
	if x == 456 && y == 586 {
		Println("mouse...", "586")
	}

	robotgo.MouseToggle("up")
	robotgo.MoveMouse(x, y)
	robotgo.MoveMouse(100, 200)

	for i := 0; i < 1080; i += 1000 {
		Println(i)
		robotgo.MoveMouse(800, i)
	}

	//read the screen
	// robotgo.CaptureScreen()
	// bit_map := robotgo.CaptureScreen()
	// Println("CaptureScreen...", bit_map)
	// gbit_map := robotgo.Capture_Screen()
	gbitMap := robotgo.BCaptureScreen()
	Println("Capture_Screen...", gbitMap.Width)

	sx, sy := robotgo.GetScreenSize()
	//Gets the screen width and height
	Println("...", sx, sy)

	color := robotgo.GetPixelColor(100, 200)
	//Gets the pixel color at 100, 200.
	Println("color----", color, "-----------------")

	color2 := robotgo.GetPixelColor(10, 20)
	//Gets the pixel color at 10, 20.
	Println("color---", color2)

	// Bitmap
	abitMap := robotgo.CaptureScreen()
	//Gets all of the screen
	Println("a...", abitMap)

	bitmap := robotgo.CaptureScreen(100, 200, 30, 40)
	//Gets part of the screen
	Println("CaptureScreen...", bitmap)
	// Println("...", bit_map.Width, bit_map.BytesPerPixel)

	fx, fy := robotgo.FindBitmap(bitmap)
	//Searches for needle in bitmap
	Println("FindBitmap------", fx, fy)

	bitpos := robotgo.GetPortion(bitmap, 10, 10, 11, 10)
	//Returns new bitmap object created from a portion of another
	Println(bitpos)

	bitstr := robotgo.TostringBitmap(bitmap)
	//Creates bitmap from string by bit_map
	Println("bit_str...", bitstr)

	// sbit_map := robotgo.BitmapFromstring(bit_str, 2)
	// Println("...", sbit_map)

	robotgo.SaveBitmap(bitmap, "test.png")
	//Saves image to absolute filepath in the given format
	robotgo.SaveBitmap(bitmap, "test31.tif", 1)
	robotgo.Convert("test.png", "test.tif")
	//Convert image

	// open_bit := robotgo.OpenBitmap("test.tif")
	openbit := robotgo.OpenBitmap("test.tif")
	// open image bitmap
	Println("open...", openbit)

	//global event listener
	Println("---please press v---")
	eve := robotgo.AddEvent("v")

	if eve == 0 {
		Println("---you press v---", "v")
	}

	Println("---please press k---")
	keve := robotgo.AddEvent("k")
	if keve == 0 {
		Println("---you press k---", "k")
	}

	Println("---please press left mouse button---")
	mleft := robotgo.AddEvent("mleft")
	if mleft == 0 {
		Println("---you press left mouse button---", "mleft")
	}

	// mright := robotgo.AddEvent("mright")
	// if mright == 0 {
	// 	Println("---you press right mouse button---", "mright")
	// }

	// robotgo.LStop()

	//Window Handle
	abool := robotgo.ShowAlert("hello", "robotgo") //Show Alert Window
	if abool == 0 {
		Println("ok@@@", "ok")
	}
	robotgo.ShowAlert("hello", "robotgo", "Ok", "Cancel")
	// robotgo.GetPID()
	mdata := robotgo.GetActive() //Get current Window Active
	hwnd := robotgo.GetHandle()  //Get current Window Handle
	Println("hwnd---", hwnd)
	title := robotgo.GetTitle() //Get current Window title
	Println("title-----", title)
	robotgo.CloseWindow()    //close current Window
	robotgo.SetActive(mdata) //set Window Active
}
<|endoftext|>


<commit_before>package main

import (
	"fmt"
	"log"
	"../"
	"github.com/timob/javabind"
	"text/tabwriter"
	"os"
	"strings"
)

type proxy struct {
	keyspace string
	bean *nodeprobe.DbColumnFamilyStoreMBean
}

type stat struct {
	title string
	format string
	f func(p *proxy) interface{}
}

func main() {
	err := javabind.SetupJVM(
		"/files/data/apache-cassandra-2.0.8/lib/apache-cassandra-2.0.8.jar",
		"/files/data/apache-cassandra-2.0.8/lib/guava-15.0.jar",
	)
	if err != nil {
		log.Fatal(err)
	}

	nodeProbe, err := nodeprobe.NewToolsNodeProbe3("localhost")
	if err != nil {
		log.Fatal(err)
	}

	writer := tabwriter.NewWriter(os.Stdout, 2, 4, 2, ' ', 0)

	stats := make([]stat, 0)
	stats = append(stats, stat{"Keyspace", "%s", func(p *proxy) interface{} { return p.keyspace }})
	stats = append(stats, stat{"CF name", "%s", func(p *proxy) interface{} { return p.bean.GetColumnFamilyName() }})
	stats = append(stats, stat{"DiskSpace", "%d", func(p *proxy) interface{} { return p.bean.GetTotalDiskSpaceUsed() }})
	stats = append(stats, stat{"ReadCount", "%d", func(p *proxy) interface{} { return p.bean.GetReadCount() }})
	stats = append(stats, stat{"WriteCount", "%d", func(p *proxy) interface{} { return p.bean.GetWriteCount() }})
	stats = append(stats, stat{"LiveSSTableCount", "%d", func(p *proxy) interface{} { return p.bean.GetLiveSSTableCount() }})

//	titles := forall(stats, func(s *stat) string { return s.title}).([]string)
	titles := make([]string, 0)
	formats := make([]string, 0)
	for _, s := range stats {
		titles = append(titles, s.title)
		formats = append(formats, s.format)
	}
	fmt.Fprintf(writer, strings.Join(titles, "\t") + "\n")
	format := strings.Join(formats, "\t") + "\n"
	results := make([]interface{}, len(stats))
	for _, v := range nodeProbe.GetColumnFamilyStoreMBeanProxies() {
		proxy := &proxy{v.Key, v.Value}
		for i, s := range stats {
			results[i] = s.f(proxy)
		}
		fmt.Fprintf(writer, format, results...)
	}

	if err = nodeProbe.Close(); err != nil {
		log.Fatal(err)
	}
	writer.Flush()
}

<commit_msg>update example<commit_after>package main

import (
	"fmt"
	"log"
	"../"
	"github.com/timob/javabind"
	"text/tabwriter"
	"os"
	"strings"
)

type proxy struct {
	keyspace string
	bean *nodeprobe.DbColumnFamilyStoreMBean
}

type stat struct {
	title string
	format string
	f func(p *proxy) interface{}
}

func main() {
	err := javabind.SetupJVM(os.Getenv("CLASSPATH"))
	if err != nil {
		log.Fatal(err)
	}

	nodeProbe, err := nodeprobe.NewToolsNodeProbe3("localhost")
	if err != nil {
		log.Fatal(err)
	}

	writer := tabwriter.NewWriter(os.Stdout, 2, 4, 2, ' ', 0)

	stats := make([]stat, 0)
	stats = append(stats, stat{"Keyspace", "%s", func(p *proxy) interface{} { return p.keyspace }})
	stats = append(stats, stat{"CF name", "%s", func(p *proxy) interface{} { return p.bean.GetColumnFamilyName() }})
	stats = append(stats, stat{"DiskSpace", "%d", func(p *proxy) interface{} { return p.bean.GetTotalDiskSpaceUsed() }})
	stats = append(stats, stat{"ReadCount", "%d", func(p *proxy) interface{} { return p.bean.GetReadCount() }})
	stats = append(stats, stat{"WriteCount", "%d", func(p *proxy) interface{} { return p.bean.GetWriteCount() }})
	stats = append(stats, stat{"LiveSSTableCount", "%d", func(p *proxy) interface{} { return p.bean.GetLiveSSTableCount() }})

//	titles := forall(stats, func(s *stat) string { return s.title}).([]string)
	titles := make([]string, 0)
	formats := make([]string, 0)
	for _, s := range stats {
		titles = append(titles, s.title)
		formats = append(formats, s.format)
	}
	fmt.Fprintf(writer, strings.Join(titles, "\t") + "\n")
	format := strings.Join(formats, "\t") + "\n"
	results := make([]interface{}, len(stats))
	for _, v := range nodeProbe.GetColumnFamilyStoreMBeanProxies() {
		proxy := &proxy{v.Key, v.Value}
		for i, s := range stats {
			results[i] = s.f(proxy)
		}
		fmt.Fprintf(writer, format, results...)
	}

	if err = nodeProbe.Close(); err != nil {
		log.Fatal(err)
	}
	writer.Flush()
}

<|endoftext|>